mirror of
https://github.com/openai/codex.git
synced 2026-11-07 00:33:29 +00:00
Enable remote compaction for Amazon Bedrock (#36981)
## What changed - Add provider-owned remote compaction capabilities for unsupported, v1, and v2 protocols. - Mark Amazon Bedrock as v1-only so manual and automatic compaction use `/v1/responses/compact`, even when the v2 feature is enabled. - Preserve v2 support for OpenAI and Azure Responses providers and local compaction for unsupported providers. ## Testing - Cover Amazon Bedrock manual and automatic remote compaction, including endpoint selection, authentication headers, model selection, and reuse of the compacted summary. GitOrigin-RevId: d58093178605d1f8155b93d9a161408177f6ac72
This commit is contained in:
1 parent
e1f39b5f5b
commit
b9ba969f39
10 files changed
+277
-156
No files matched your search
@@ -10,6 +10,10 @@ use codex_features::Feature;
|
||||
use codex_login::CodexAuth;
|
||||
use codex_login::auth::AgentIdentityAuth;
|
||||
use codex_login::auth::AgentIdentityAuthRecord;
|
||||
use codex_login::auth::BedrockApiKeyAuth;
|
||||
use codex_model_provider_info::AMAZON_BEDROCK_GPT_5_5_MODEL_ID;
|
||||
use codex_model_provider_info::AMAZON_BEDROCK_PROVIDER_ID;
|
||||
use codex_model_provider_info::ModelProviderInfo;
|
||||
use codex_protocol::AgentPath;
|
||||
use codex_protocol::account::PlanType as AccountPlanType;
|
||||
use codex_protocol::config_types::ServiceTier;
|
||||
@@ -346,6 +350,148 @@ async fn wait_for_turn_complete(codex: &codex_core::CodexThread) {
|
||||
.await;
|
||||
}
|
||||
|
||||
fn amazon_bedrock_test_codex() -> TestCodexBuilder {
|
||||
let auth = CodexAuth::BedrockApiKey(BedrockApiKeyAuth {
|
||||
api_key: "bedrock-test-api-key".to_string(),
|
||||
region: "us-east-1".to_string(),
|
||||
});
|
||||
test_codex()
|
||||
.with_auth(auth)
|
||||
.with_model(AMAZON_BEDROCK_GPT_5_5_MODEL_ID)
|
||||
.with_config(|config| {
|
||||
let _ = config.features.enable(Feature::RemoteCompactionV2);
|
||||
config.model_provider = ModelProviderInfo {
|
||||
base_url: config.model_provider.base_url.clone(),
|
||||
..ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None)
|
||||
};
|
||||
config.model_provider_id = AMAZON_BEDROCK_PROVIDER_ID.to_string();
|
||||
})
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn amazon_bedrock_uses_remote_compaction_endpoint() -> Result<()> {
|
||||
skip_if_no_network!(Ok(()));
|
||||
|
||||
let harness = TestCodexHarness::with_auto_env_builder(amazon_bedrock_test_codex()).await?;
|
||||
|
||||
let response_mock = responses::mount_sse_sequence(
|
||||
harness.server(),
|
||||
vec![
|
||||
sse(vec![
|
||||
responses::ev_assistant_message("message-1", "before compaction"),
|
||||
responses::ev_completed("response-1"),
|
||||
]),
|
||||
sse(vec![
|
||||
responses::ev_assistant_message("message-2", "after compaction"),
|
||||
responses::ev_completed("response-2"),
|
||||
]),
|
||||
],
|
||||
)
|
||||
.await;
|
||||
let compact_mock = responses::mount_compact_user_history_with_summary_once(
|
||||
harness.server(),
|
||||
"BEDROCK_REMOTE_COMPACTED_SUMMARY",
|
||||
)
|
||||
.await;
|
||||
|
||||
harness.test().submit_turn("before compact").await?;
|
||||
harness.test().codex.submit(Op::Compact).await?;
|
||||
wait_for_turn_complete(&harness.test().codex).await;
|
||||
harness.test().submit_turn("after compact").await?;
|
||||
|
||||
let compact_request = compact_mock.single_request();
|
||||
assert_eq!(compact_request.path(), "/v1/responses/compact");
|
||||
assert_eq!(
|
||||
compact_request.header("authorization").as_deref(),
|
||||
Some("Bearer bedrock-test-api-key")
|
||||
);
|
||||
assert_eq!(
|
||||
compact_request
|
||||
.header("x-amzn-mantle-client-agent")
|
||||
.as_deref(),
|
||||
Some("codex")
|
||||
);
|
||||
assert_eq!(
|
||||
compact_request.body_json()["model"],
|
||||
AMAZON_BEDROCK_GPT_5_5_MODEL_ID
|
||||
);
|
||||
|
||||
let response_requests = response_mock.requests();
|
||||
assert_eq!(response_requests.len(), 2);
|
||||
assert!(
|
||||
response_requests
|
||||
.iter()
|
||||
.all(|request| request.inputs_of_type("compaction_trigger").is_empty())
|
||||
);
|
||||
assert!(response_requests[1].input().iter().any(|item| {
|
||||
item["type"] == "compaction"
|
||||
&& item["encrypted_content"] == "BEDROCK_REMOTE_COMPACTED_SUMMARY"
|
||||
}));
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn amazon_bedrock_automatic_compaction_uses_v1_endpoint_when_v2_is_enabled() -> Result<()> {
|
||||
skip_if_no_network!(Ok(()));
|
||||
|
||||
let harness = TestCodexHarness::with_auto_env_builder(amazon_bedrock_test_codex().with_config(
|
||||
|config| {
|
||||
config.model_auto_compact_token_limit = Some(200);
|
||||
},
|
||||
))
|
||||
.await?;
|
||||
let response_mock = responses::mount_sse_sequence(
|
||||
harness.server(),
|
||||
vec![
|
||||
sse(vec![
|
||||
responses::ev_assistant_message("message-1", "before automatic compaction"),
|
||||
responses::ev_completed_with_tokens("response-1", /*total_tokens*/ 500),
|
||||
]),
|
||||
sse(vec![
|
||||
responses::ev_assistant_message("message-2", "after automatic compaction"),
|
||||
responses::ev_completed("response-2"),
|
||||
]),
|
||||
],
|
||||
)
|
||||
.await;
|
||||
let compact_mock = responses::mount_compact_user_history_with_summary_once(
|
||||
harness.server(),
|
||||
"BEDROCK_AUTOMATIC_REMOTE_COMPACTED_SUMMARY",
|
||||
)
|
||||
.await;
|
||||
|
||||
harness
|
||||
.test()
|
||||
.submit_turn("before automatic compact")
|
||||
.await?;
|
||||
harness
|
||||
.test()
|
||||
.submit_turn("after automatic compact")
|
||||
.await?;
|
||||
|
||||
let compact_request = compact_mock.single_request();
|
||||
assert_eq!(compact_request.path(), "/v1/responses/compact");
|
||||
assert_eq!(
|
||||
compact_request.header("authorization").as_deref(),
|
||||
Some("Bearer bedrock-test-api-key")
|
||||
);
|
||||
|
||||
let response_requests = response_mock.requests();
|
||||
assert_eq!(response_requests.len(), 2);
|
||||
assert!(
|
||||
response_requests
|
||||
.iter()
|
||||
.all(|request| request.inputs_of_type("compaction_trigger").is_empty())
|
||||
);
|
||||
assert!(response_requests[1].input().iter().any(|item| {
|
||||
item["type"] == "compaction"
|
||||
&& item["encrypted_content"] == "BEDROCK_AUTOMATIC_REMOTE_COMPACTED_SUMMARY"
|
||||
}));
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn remote_compact_replaces_history_for_followups() -> Result<()> {
|
||||
skip_if_no_network!(Ok(()));
|
||||
|
||||
Reference in new issue
Block a user