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:
rka-oai authored and copyberry committed 2026-08-04 22:07:56 +00:00
1 parent e1f39b5f5b
commit b9ba969f39
10 files changed
+277 -156

No files matched your search

+146
View File
@@ -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(()));