diff --git a/codex-rs/core/src/compact.rs b/codex-rs/core/src/compact.rs index ba0b18978d..a599c2bc23 100644 --- a/codex-rs/core/src/compact.rs +++ b/codex-rs/core/src/compact.rs @@ -50,8 +50,6 @@ use codex_utils_output_truncation::truncate_text; use futures::prelude::*; use tracing::error; -use codex_model_provider_info::ModelProviderInfo; - pub use codex_prompts::SUMMARIZATION_PROMPT; pub use codex_prompts::SUMMARY_PREFIX; const COMPACT_USER_MESSAGE_MAX_TOKENS: usize = 20_000; @@ -105,10 +103,6 @@ pub(crate) async fn build_compaction_initial_context( } } -pub(crate) fn should_use_remote_compact_task(provider: &ModelProviderInfo) -> bool { - provider.supports_remote_compaction() -} - pub(crate) async fn run_inline_auto_compact_task( sess: Arc, turn_context: Arc, diff --git a/codex-rs/core/src/compact_tests.rs b/codex-rs/core/src/compact_tests.rs index ac7cc87bee..e818220a4a 100644 --- a/codex-rs/core/src/compact_tests.rs +++ b/codex-rs/core/src/compact_tests.rs @@ -1,6 +1,4 @@ use super::*; -use codex_model_provider_info::ModelProviderInfo; -use codex_model_provider_info::WireApi; use codex_protocol::ResponseItemId; use codex_protocol::models::DEFAULT_IMAGE_DETAIL; use codex_protocol::models::InternalChatMessageMetadataPassthrough; @@ -266,31 +264,6 @@ fn build_compacted_history_preserves_user_message_passthrough_metadata() { assert_eq!(history[1].turn_id(), None); } -#[test] -fn should_use_remote_compact_task_for_azure_provider() { - let provider = ModelProviderInfo { - name: "Azure".into(), - base_url: Some("https://example.com/openai".into()), - env_key: Some("AZURE_OPENAI_API_KEY".into()), - env_key_instructions: None, - experimental_bearer_token: None, - auth: None, - aws: None, - wire_api: WireApi::Responses, - query_params: None, - http_headers: None, - env_http_headers: None, - request_max_retries: None, - stream_max_retries: None, - stream_idle_timeout_ms: None, - websocket_connect_timeout_ms: None, - requires_openai_auth: false, - supports_websockets: false, - supports_standalone_web_search: false, - }; - - assert!(should_use_remote_compact_task(&provider)); -} #[tokio::test] async fn process_compacted_history_replaces_developer_messages() { let compacted_history = vec![ diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index a873582796..e5c658401e 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -11,7 +11,6 @@ use crate::client_common::ResponseEvent; use crate::collect_explicit_skill_mentions; use crate::compact::InitialContextInjection; use crate::compact::run_inline_auto_compact_task; -use crate::compact::should_use_remote_compact_task; use crate::compact_remote::run_inline_remote_auto_compact_task; use crate::compact_remote_v2::run_inline_remote_auto_compact_task as run_inline_remote_auto_compact_task_v2; use crate::connectors; @@ -81,6 +80,7 @@ use codex_features::Feature; use codex_file_system::FindUpErrorPolicy; use codex_file_system::find_nearest_ancestor_with_markers; use codex_login::CodexAuth; +use codex_model_provider::RemoteCompactionSupport; use codex_protocol::ResponseItemId; use codex_protocol::config_types::AutoCompactTokenLimitScope; use codex_protocol::config_types::ModeKind; @@ -1171,11 +1171,12 @@ async fn run_auto_compact( return Ok(()); } - if should_use_remote_compact_task(turn_context.provider.info()) { - if turn_context - .config - .features - .enabled(Feature::RemoteCompactionV2) + match turn_context.provider.capabilities().remote_compaction { + RemoteCompactionSupport::V2 + if turn_context + .config + .features + .enabled(Feature::RemoteCompactionV2) => { emit_compact_metric( &sess.services.session_telemetry, @@ -1192,37 +1193,39 @@ async fn run_auto_compact( phase, ) .await?; - return Ok(()); } - emit_compact_metric( - &sess.services.session_telemetry, - "remote", - /*manual*/ false, - ); - run_inline_remote_auto_compact_task( - Arc::clone(sess), - step_context, - fallback_step_context, - client_session.turn_state(), - initial_context_injection, - reason, - phase, - ) - .await?; - } else { - emit_compact_metric( - &sess.services.session_telemetry, - "local", - /*manual*/ false, - ); - run_inline_auto_compact_task( - Arc::clone(sess), - Arc::clone(turn_context), - initial_context_injection, - reason, - phase, - ) - .await?; + RemoteCompactionSupport::V1 | RemoteCompactionSupport::V2 => { + emit_compact_metric( + &sess.services.session_telemetry, + "remote", + /*manual*/ false, + ); + run_inline_remote_auto_compact_task( + Arc::clone(sess), + step_context, + fallback_step_context, + client_session.turn_state(), + initial_context_injection, + reason, + phase, + ) + .await?; + } + RemoteCompactionSupport::Unsupported => { + emit_compact_metric( + &sess.services.session_telemetry, + "local", + /*manual*/ false, + ); + run_inline_auto_compact_task( + Arc::clone(sess), + Arc::clone(turn_context), + initial_context_injection, + reason, + phase, + ) + .await?; + } } Ok(()) } diff --git a/codex-rs/core/src/tasks/compact.rs b/codex-rs/core/src/tasks/compact.rs index 18c9a27a5b..1c7de737eb 100644 --- a/codex-rs/core/src/tasks/compact.rs +++ b/codex-rs/core/src/tasks/compact.rs @@ -8,6 +8,7 @@ use crate::session::session::Session; use crate::session::turn_context::TurnContext; use crate::state::TaskKind; use codex_features::Feature; +use codex_model_provider::RemoteCompactionSupport; use codex_protocol::error::CodexErrorDetails; use codex_protocol::user_input::UserInput; use tokio_util::sync::CancellationToken; @@ -37,11 +38,9 @@ impl SessionTask for CompactTask { return Ok(None); } - let result = if crate::compact::should_use_remote_compact_task(ctx.provider.info()) { - if ctx - .config - .features - .enabled(codex_features::Feature::RemoteCompactionV2) + let result = match ctx.provider.capabilities().remote_compaction { + RemoteCompactionSupport::V2 + if ctx.config.features.enabled(Feature::RemoteCompactionV2) => { emit_compact_metric( &session.services.session_telemetry, @@ -49,7 +48,8 @@ impl SessionTask for CompactTask { /*manual*/ true, ); crate::compact_remote_v2::run_remote_compact_task(session.clone(), ctx).await - } else { + } + RemoteCompactionSupport::V1 | RemoteCompactionSupport::V2 => { emit_compact_metric( &session.services.session_telemetry, "remote", @@ -57,23 +57,24 @@ impl SessionTask for CompactTask { ); crate::compact_remote::run_remote_compact_task(session.clone(), ctx).await } - } else { - emit_compact_metric( - &session.services.session_telemetry, - "local", - /*manual*/ true, - ); - let input = vec![UserInput::Text { - text: ctx - .config - .compact_prompt - .as_deref() - .unwrap_or(crate::compact::SUMMARIZATION_PROMPT) - .to_string(), - // Compaction prompt is synthesized; no UI element ranges to preserve. - text_elements: Vec::new(), - }]; - crate::compact::run_compact_task(session.clone(), ctx, input).await + RemoteCompactionSupport::Unsupported => { + emit_compact_metric( + &session.services.session_telemetry, + "local", + /*manual*/ true, + ); + let input = vec![UserInput::Text { + text: ctx + .config + .compact_prompt + .as_deref() + .unwrap_or(crate::compact::SUMMARIZATION_PROMPT) + .to_string(), + // Compaction prompt is synthesized; no UI element ranges to preserve. + text_elements: Vec::new(), + }]; + crate::compact::run_compact_task(session.clone(), ctx, input).await + } }; if let Err(err) = result && matches!(err.details(), CodexErrorDetails::TurnAborted) diff --git a/codex-rs/core/tests/suite/compact_remote.rs b/codex-rs/core/tests/suite/compact_remote.rs index 8a1c5ad16c..88f5e5d3a5 100644 --- a/codex-rs/core/tests/suite/compact_remote.rs +++ b/codex-rs/core/tests/suite/compact_remote.rs @@ -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(())); diff --git a/codex-rs/model-provider-info/src/lib.rs b/codex-rs/model-provider-info/src/lib.rs index 82fc5fc2a6..95d9401fa5 100644 --- a/codex-rs/model-provider-info/src/lib.rs +++ b/codex-rs/model-provider-info/src/lib.rs @@ -7,7 +7,6 @@ use codex_api::Provider as ApiProvider; use codex_api::RetryConfig as ApiRetryConfig; -use codex_api::is_azure_responses_provider; use codex_protocol::auth::AuthMode; use codex_protocol::config_types::ModelProviderAuthInfo; use codex_protocol::error::CodexErr; @@ -419,10 +418,6 @@ impl ModelProviderInfo { self.name == AMAZON_BEDROCK_PROVIDER_NAME } - pub fn supports_remote_compaction(&self) -> bool { - self.is_openai() || is_azure_responses_provider(&self.name, self.base_url.as_deref()) - } - pub fn has_command_auth(&self) -> bool { self.auth.is_some() } diff --git a/codex-rs/model-provider-info/src/model_provider_info_tests.rs b/codex-rs/model-provider-info/src/model_provider_info_tests.rs index 354c7fc2b8..39ceba8199 100644 --- a/codex-rs/model-provider-info/src/model_provider_info_tests.rs +++ b/codex-rs/model-provider-info/src/model_provider_info_tests.rs @@ -136,13 +136,6 @@ supports_websockets = true assert_eq!(provider.websocket_connect_timeout_ms, Some(15_000)); } -#[test] -fn test_supports_remote_compaction_for_openai() { - let provider = ModelProviderInfo::create_openai_provider(/*base_url*/ None); - - assert!(provider.supports_remote_compaction()); -} - #[test] fn test_personal_access_token_uses_chatgpt_codex_base_url() { let api_provider = ModelProviderInfo::create_openai_provider(/*base_url*/ None) @@ -161,58 +154,6 @@ fn test_header_auth_uses_chatgpt_codex_base_url() { assert_eq!(api_provider.base_url, CHATGPT_CODEX_BASE_URL); } -#[test] -fn test_supports_remote_compaction_for_azure_name() { - let provider = ModelProviderInfo { - name: "Azure".into(), - base_url: Some("https://example.com/openai".into()), - env_key: Some("AZURE_OPENAI_API_KEY".into()), - env_key_instructions: None, - experimental_bearer_token: None, - auth: None, - aws: None, - wire_api: WireApi::Responses, - query_params: None, - http_headers: None, - env_http_headers: None, - request_max_retries: None, - stream_max_retries: None, - stream_idle_timeout_ms: None, - websocket_connect_timeout_ms: None, - requires_openai_auth: false, - supports_websockets: false, - supports_standalone_web_search: false, - }; - - assert!(provider.supports_remote_compaction()); -} - -#[test] -fn test_supports_remote_compaction_for_non_openai_non_azure_provider() { - let provider = ModelProviderInfo { - name: "Example".into(), - base_url: Some("https://example.com/v1".into()), - env_key: Some("API_KEY".into()), - env_key_instructions: None, - experimental_bearer_token: None, - auth: None, - aws: None, - wire_api: WireApi::Responses, - query_params: None, - http_headers: None, - env_http_headers: None, - request_max_retries: None, - stream_max_retries: None, - stream_idle_timeout_ms: None, - websocket_connect_timeout_ms: None, - requires_openai_auth: false, - supports_websockets: false, - supports_standalone_web_search: false, - }; - - assert!(!provider.supports_remote_compaction()); -} - #[test] fn test_uses_openai_actor_authorization() { let mut provider = ModelProviderInfo { diff --git a/codex-rs/model-provider/src/amazon_bedrock/mod.rs b/codex-rs/model-provider/src/amazon_bedrock/mod.rs index 51bfafdce7..6e3140d0e2 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/mod.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/mod.rs @@ -29,6 +29,7 @@ use crate::provider::ModelProviderFuture; use crate::provider::ProviderAccountResult; use crate::provider::ProviderAccountState; use crate::provider::ProviderCapabilities; +use crate::provider::RemoteCompactionSupport; use auth::resolve_provider_auth as resolve_bedrock_provider_auth; use catalog::normalize_bedrock_catalog; pub(crate) use catalog::static_model_catalog; @@ -126,6 +127,7 @@ impl ModelProvider for AmazonBedrockModelProvider { image_generation: false, web_search: true, external_web_access: false, + remote_compaction: RemoteCompactionSupport::V1, } } @@ -368,6 +370,7 @@ mod tests { image_generation: false, web_search: true, external_web_access: false, + remote_compaction: RemoteCompactionSupport::V1, } ); } diff --git a/codex-rs/model-provider/src/lib.rs b/codex-rs/model-provider/src/lib.rs index 47d477b88b..6da0053d9d 100644 --- a/codex-rs/model-provider/src/lib.rs +++ b/codex-rs/model-provider/src/lib.rs @@ -22,5 +22,6 @@ pub use provider::ProviderAccountError; pub use provider::ProviderAccountResult; pub use provider::ProviderAccountState; pub use provider::ProviderCapabilities; +pub use provider::RemoteCompactionSupport; pub use provider::SharedModelProvider; pub use provider::create_model_provider; diff --git a/codex-rs/model-provider/src/provider.rs b/codex-rs/model-provider/src/provider.rs index 95a873bea6..1964864f2c 100644 --- a/codex-rs/model-provider/src/provider.rs +++ b/codex-rs/model-provider/src/provider.rs @@ -7,6 +7,7 @@ use std::sync::Arc; use codex_api::ApiError; use codex_api::Provider; use codex_api::SharedAuthProvider; +use codex_api::is_azure_responses_provider; use codex_login::AuthManager; use codex_login::CodexAuth; use codex_model_provider_info::ModelProviderInfo; @@ -25,6 +26,17 @@ use crate::auth::resolve_provider_auth; use crate::auth::resolve_provider_auth_for_scope; use crate::models_endpoint::OpenAiModelsEndpoint; +/// Remote context-compaction protocols supported by a model provider. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum RemoteCompactionSupport { + /// The provider does not support remote compaction. + Unsupported, + /// The provider supports only the dedicated `/v1/responses/compact` endpoint. + V1, + /// The provider supports both the dedicated endpoint and `compaction_trigger` items. + V2, +} + /// Optional provider-backed features that Codex may expose at runtime. /// /// These capabilities are a provider-owned upper bound. Callers can disable @@ -36,6 +48,7 @@ pub struct ProviderCapabilities { pub image_generation: bool, pub web_search: bool, pub external_web_access: bool, + pub remote_compaction: RemoteCompactionSupport, } impl Default for ProviderCapabilities { @@ -45,6 +58,7 @@ impl Default for ProviderCapabilities { image_generation: true, web_search: true, external_web_access: true, + remote_compaction: RemoteCompactionSupport::V2, } } } @@ -264,6 +278,21 @@ impl ModelProvider for ConfiguredModelProvider { &self.info } + fn capabilities(&self) -> ProviderCapabilities { + let remote_compaction = if self.info.is_openai() + || is_azure_responses_provider(&self.info.name, self.info.base_url.as_deref()) + { + RemoteCompactionSupport::V2 + } else { + RemoteCompactionSupport::Unsupported + }; + + ProviderCapabilities { + remote_compaction, + ..ProviderCapabilities::default() + } + } + fn auth_manager(&self) -> Option> { self.auth_manager.clone() } @@ -508,6 +537,41 @@ mod tests { assert_eq!(provider.capabilities(), ProviderCapabilities::default()); } + #[test] + fn configured_provider_remote_compaction_matches_provider_support() { + let cases = [ + ( + ModelProviderInfo::create_openai_provider(/*base_url*/ None), + RemoteCompactionSupport::V2, + ), + ( + ModelProviderInfo { + name: "Azure".to_string(), + base_url: Some("https://example.com/openai".to_string()), + ..ModelProviderInfo::default() + }, + RemoteCompactionSupport::V2, + ), + ( + ModelProviderInfo { + name: "Custom".to_string(), + base_url: Some("https://example.openai.azure.com/openai/v1".to_string()), + ..ModelProviderInfo::default() + }, + RemoteCompactionSupport::V2, + ), + ( + provider_for("https://example.test/v1".to_string()), + RemoteCompactionSupport::Unsupported, + ), + ]; + + for (provider_info, expected) in cases { + let provider = create_model_provider(provider_info, /*auth_manager*/ None); + assert_eq!(provider.capabilities().remote_compaction, expected); + } + } + #[test] fn configured_provider_uses_default_approval_review_preferred_model() { let provider = create_model_provider(