diff --git a/codex-rs/codex-api/src/endpoint/streaming.rs b/codex-rs/codex-api/src/endpoint/streaming.rs index ea7e4c85f3..a2f1106797 100644 --- a/codex-rs/codex-api/src/endpoint/streaming.rs +++ b/codex-rs/codex-api/src/endpoint/streaming.rs @@ -20,6 +20,8 @@ use http::header::CONTENT_TYPE; use serde_json::Value; use std::sync::Arc; use std::time::Duration; +use std::time::Instant; +use tracing::info; use zstd::stream::encode_all; pub(crate) struct StreamingClient { @@ -105,8 +107,16 @@ fn encode_body(body: &Value, compression: RequestCompression) -> Result { let json = serde_json::to_vec(body) .map_err(|err| format!("failed to encode request body as json: {err}"))?; + let started_at = Instant::now(); let compressed = encode_all(json.as_slice(), 0) .map_err(|err| format!("failed to compress request body: {err}"))?; + let elapsed = started_at.elapsed(); + info!( + input_bytes = json.len(), + output_bytes = compressed.len(), + elapsed_ms = elapsed.as_millis(), + "compressed request body" + ); Ok(Body::Bytes(Bytes::from(compressed))) } } diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 11a3c5c65f..11e15f2278 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -158,7 +158,7 @@ impl ModelClient { let auth = auth_manager.as_ref().and_then(|m| m.auth()); let api_provider = self .provider - .to_api_provider(auth.as_ref().map(|a| a.mode))?; + .to_api_provider(auth.as_ref().map(|a| a.mode), &self.config.features)?; let api_auth = auth_provider_from_auth(auth.clone(), &self.provider).await?; let transport = ReqwestTransport::new(build_reqwest_client()); let (request_telemetry, sse_telemetry) = self.build_streaming_telemetry(); @@ -247,7 +247,7 @@ impl ModelClient { let auth = auth_manager.as_ref().and_then(|m| m.auth()); let api_provider = self .provider - .to_api_provider(auth.as_ref().map(|a| a.mode))?; + .to_api_provider(auth.as_ref().map(|a| a.mode), &self.config.features)?; let api_auth = auth_provider_from_auth(auth.clone(), &self.provider).await?; let transport = ReqwestTransport::new(build_reqwest_client()); let (request_telemetry, sse_telemetry) = self.build_streaming_telemetry(); @@ -332,7 +332,7 @@ impl ModelClient { let auth = auth_manager.as_ref().and_then(|m| m.auth()); let api_provider = self .provider - .to_api_provider(auth.as_ref().map(|a| a.mode))?; + .to_api_provider(auth.as_ref().map(|a| a.mode), &self.config.features)?; let api_auth = auth_provider_from_auth(auth.clone(), &self.provider).await?; let transport = ReqwestTransport::new(build_reqwest_client()); let request_telemetry = self.build_request_telemetry(); diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 51a1bd07f7..e72907e4e6 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -1972,7 +1972,10 @@ trust_level = "trusted" fn feature_table_overrides_legacy_flags() -> std::io::Result<()> { let codex_home = TempDir::new()?; let mut entries = BTreeMap::new(); - entries.insert("apply_patch_freeform".to_string(), false); + entries.insert( + "apply_patch_freeform".to_string(), + crate::features::FeatureValue::Bool(false), + ); let cfg = ConfigToml { features: Some(crate::features::FeaturesToml { entries }), ..Default::default() diff --git a/codex-rs/core/src/features.rs b/codex-rs/core/src/features.rs index 3b22bfc3f4..1cb4cf5d81 100644 --- a/codex-rs/core/src/features.rs +++ b/codex-rs/core/src/features.rs @@ -74,6 +74,8 @@ pub enum Feature { ApplyPatchFreeform, /// Allow the model to request web searches. WebSearchRequest, + /// Allow request body compression when using ChatGPT auth. + RequestCompression, /// Gate the execpolicy enforcement for shell/unified exec. ExecPolicy, /// Enable Windows sandbox (restricted token) on Windows. @@ -128,6 +130,7 @@ pub struct LegacyFeatureUsage { pub struct Features { enabled: BTreeSet, legacy_usages: BTreeSet, + request_compression: RequestCompressionFeature, } #[derive(Debug, Clone, Default)] @@ -150,16 +153,17 @@ impl FeatureOverrides { impl Features { /// Starts with built-in defaults. pub fn with_defaults() -> Self { - let mut set = BTreeSet::new(); + let mut features = Self { + enabled: BTreeSet::new(), + legacy_usages: BTreeSet::new(), + request_compression: RequestCompressionFeature::Disabled, + }; for spec in FEATURES { if spec.default_enabled { - set.insert(spec.id); + features.enable(spec.id); } } - Self { - enabled: set, - legacy_usages: BTreeSet::new(), - } + features } pub fn enabled(&self, f: Feature) -> bool { @@ -168,11 +172,17 @@ impl Features { pub fn enable(&mut self, f: Feature) -> &mut Self { self.enabled.insert(f); + if matches!(f, Feature::RequestCompression) { + self.request_compression = RequestCompressionFeature::Zstd; + } self } pub fn disable(&mut self, f: Feature) -> &mut Self { self.enabled.remove(&f); + if matches!(f, Feature::RequestCompression) { + self.request_compression = RequestCompressionFeature::Disabled; + } self } @@ -196,18 +206,62 @@ impl Features { .map(|usage| (usage.alias.as_str(), usage.feature)) } - /// Apply a table of key -> bool toggles (e.g. from TOML). - pub fn apply_map(&mut self, m: &BTreeMap) { + pub fn request_compression(&self) -> RequestCompressionFeature { + self.request_compression + } + + pub fn set_request_compression( + &mut self, + request_compression: RequestCompressionFeature, + ) -> &mut Self { + self.request_compression = request_compression; + if self.request_compression == RequestCompressionFeature::Disabled { + self.enabled.remove(&Feature::RequestCompression); + } else { + self.enabled.insert(Feature::RequestCompression); + } + self + } + + /// Apply a table of key -> value toggles (e.g. from TOML). + pub fn apply_map(&mut self, m: &BTreeMap) { for (k, v) in m { match feature_for_key(k) { Some(feat) => { if k != feat.key() { self.record_legacy_usage(k.as_str(), feat); } - if *v { - self.enable(feat); + if feat == Feature::RequestCompression { + match v { + FeatureValue::Bool(enabled) => { + let request_compression = if *enabled { + RequestCompressionFeature::Zstd + } else { + RequestCompressionFeature::Disabled + }; + self.set_request_compression(request_compression); + } + FeatureValue::String(value) => { + match RequestCompressionFeature::parse(value) { + Some(request_compression) => { + self.set_request_compression(request_compression); + } + None => { + tracing::warn!( + "unknown request_compression feature value in config: {value}" + ); + } + } + } + } + } else if let FeatureValue::Bool(enabled) = v { + if *enabled { + self.enable(feat); + } else { + self.disable(feat); + } } else { - self.disable(feat); + tracing::warn!("feature key expects boolean value: {k}"); } } None => { @@ -280,7 +334,31 @@ pub fn is_known_feature_key(key: &str) -> bool { #[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq)] pub struct FeaturesToml { #[serde(flatten)] - pub entries: BTreeMap, + pub entries: BTreeMap, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(untagged)] +pub enum FeatureValue { + Bool(bool), + String(String), +} + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub enum RequestCompressionFeature { + #[default] + Disabled, + Zstd, +} + +impl RequestCompressionFeature { + pub fn parse(value: &str) -> Option { + match value.to_ascii_lowercase().as_str() { + "none" | "disabled" => Some(Self::Disabled), + "zstd" => Some(Self::Zstd), + _ => None, + } + } } /// Single, easy-to-read registry of all feature definitions. @@ -330,6 +408,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::Stable, default_enabled: false, }, + FeatureSpec { + id: Feature::RequestCompression, + key: "request_compression", + stage: Stage::Experimental, + default_enabled: false, + }, // Beta program. Rendered in the `/experimental` menu for users. FeatureSpec { id: Feature::UnifiedExec, diff --git a/codex-rs/core/src/model_provider_info.rs b/codex-rs/core/src/model_provider_info.rs index 1271dbd587..9aab61b289 100644 --- a/codex-rs/core/src/model_provider_info.rs +++ b/codex-rs/core/src/model_provider_info.rs @@ -19,6 +19,8 @@ use std::env::VarError; use std::time::Duration; use crate::error::EnvVarError; +use crate::features::Features; +use crate::features::RequestCompressionFeature; const DEFAULT_STREAM_IDLE_TIMEOUT_MS: u64 = 300_000; const DEFAULT_STREAM_MAX_RETRIES: u64 = 5; const DEFAULT_REQUEST_MAX_RETRIES: u64 = 4; @@ -143,6 +145,7 @@ impl ModelProviderInfo { pub(crate) fn to_api_provider( &self, auth_mode: Option, + features: &Features, ) -> crate::error::Result { let default_base_url = if matches!(auth_mode, Some(AuthMode::ChatGPT)) { "https://chatgpt.com/backend-api/codex" @@ -173,7 +176,7 @@ impl ModelProviderInfo { }, headers, retry, - request_compression: match self.request_compression { + request_compression: match self.request_compression_for(auth_mode, features) { RequestCompression::None => codex_api::provider::RequestCompression::None, RequestCompression::Zstd => codex_api::provider::RequestCompression::Zstd, }, @@ -271,6 +274,21 @@ impl ModelProviderInfo { pub fn is_openai(&self) -> bool { self.name == OPENAI_PROVIDER_NAME } + + pub fn request_compression_for( + &self, + auth_mode: Option, + features: &Features, + ) -> RequestCompression { + if self.is_openai() && matches!(auth_mode, Some(AuthMode::ChatGPT)) { + match features.request_compression() { + RequestCompressionFeature::Zstd => RequestCompression::Zstd, + RequestCompressionFeature::Disabled => RequestCompression::None, + } + } else { + RequestCompression::None + } + } } pub const DEFAULT_LMSTUDIO_PORT: u16 = 1234; @@ -466,7 +484,9 @@ env_http_headers = { "X-Example-Env-Header" = "EXAMPLE_ENV_VAR" } stream_idle_timeout_ms: None, requires_openai_auth: false, }; - let api = provider.to_api_provider(None).expect("api provider"); + let api = provider + .to_api_provider(None, &Features::with_defaults()) + .expect("api provider"); assert!( api.is_azure_responses_endpoint(), "expected {base_url} to be detected as Azure" @@ -489,7 +509,9 @@ env_http_headers = { "X-Example-Env-Header" = "EXAMPLE_ENV_VAR" } stream_idle_timeout_ms: None, requires_openai_auth: false, }; - let named_api = named_provider.to_api_provider(None).expect("api provider"); + let named_api = named_provider + .to_api_provider(None, &Features::with_defaults()) + .expect("api provider"); assert!(named_api.is_azure_responses_endpoint()); let negative_cases = [ @@ -514,7 +536,9 @@ env_http_headers = { "X-Example-Env-Header" = "EXAMPLE_ENV_VAR" } stream_idle_timeout_ms: None, requires_openai_auth: false, }; - let api = provider.to_api_provider(None).expect("api provider"); + let api = provider + .to_api_provider(None, &Features::with_defaults()) + .expect("api provider"); assert!( !api.is_azure_responses_endpoint(), "expected {base_url} not to be detected as Azure" diff --git a/codex-rs/core/src/models_manager/manager.rs b/codex-rs/core/src/models_manager/manager.rs index da96a4147a..d84c30d208 100644 --- a/codex-rs/core/src/models_manager/manager.rs +++ b/codex-rs/core/src/models_manager/manager.rs @@ -98,7 +98,9 @@ impl ModelsManager { return Ok(()); } let auth = self.auth_manager.auth(); - let api_provider = self.provider.to_api_provider(Some(AuthMode::ChatGPT))?; + let api_provider = self + .provider + .to_api_provider(Some(AuthMode::ChatGPT), &config.features)?; let api_auth = auth_provider_from_auth(auth.clone(), &self.provider).await?; let transport = ReqwestTransport::new(build_reqwest_client()); let client = ModelsClient::new(transport, api_provider, api_auth); diff --git a/codex-rs/core/tests/responses_headers.rs b/codex-rs/core/tests/responses_headers.rs index 28f79217a0..4625286308 100644 --- a/codex-rs/core/tests/responses_headers.rs +++ b/codex-rs/core/tests/responses_headers.rs @@ -12,6 +12,7 @@ use codex_core::ResponseEvent; use codex_core::ResponseItem; use codex_core::WireApi; use codex_core::models_manager::manager::ModelsManager; +use codex_core::features::RequestCompressionFeature; use codex_otel::otel_manager::OtelManager; use codex_protocol::ConversationId; use codex_protocol::config_types::ReasoningSummary; @@ -47,7 +48,7 @@ async fn responses_stream_includes_subagent_header_on_review() { env_key_instructions: None, experimental_bearer_token: None, wire_api: WireApi::Responses, - request_compression: RequestCompression::Zstd, + request_compression: RequestCompression::None, query_params: None, http_headers: None, env_http_headers: None, @@ -142,7 +143,7 @@ async fn responses_stream_includes_subagent_header_on_other() { env_key_instructions: None, experimental_bearer_token: None, wire_api: WireApi::Responses, - request_compression: RequestCompression::Zstd, + request_compression: RequestCompression::None, query_params: None, http_headers: None, env_http_headers: None, @@ -335,13 +336,13 @@ async fn responses_request_body_is_zstd_encoded() { let request_recorder = responses::mount_sse_once(&server, response_body).await; let provider = ModelProviderInfo { - name: "mock".into(), + name: "OpenAI".into(), base_url: Some(format!("{}/v1", server.uri())), env_key: None, env_key_instructions: None, experimental_bearer_token: None, wire_api: WireApi::Responses, - request_compression: RequestCompression::Zstd, + request_compression: RequestCompression::None, query_params: None, http_headers: None, env_http_headers: None, @@ -355,6 +356,9 @@ async fn responses_request_body_is_zstd_encoded() { let mut config = load_default_config_for_test(&codex_home); config.model_provider_id = provider.name.clone(); config.model_provider = provider.clone(); + config + .features + .set_request_compression(RequestCompressionFeature::Zstd); let effort = config.model_reasoning_effort; let summary = config.model_reasoning_summary; let model = ModelsManager::get_model_offline(config.model.as_deref()); @@ -376,9 +380,11 @@ async fn responses_request_body_is_zstd_encoded() { session_source.clone(), ); + let auth_manager = + AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); let client = ModelClient::new( Arc::clone(&config), - None, + Some(auth_manager), model_family, otel_manager, provider, @@ -387,11 +393,6 @@ async fn responses_request_body_is_zstd_encoded() { conversation_id, session_source, ); - assert_eq!( - client.get_provider().request_compression, - RequestCompression::Zstd - ); - let mut prompt = Prompt::default(); prompt.input = vec![ResponseItem::Message { id: None, @@ -435,7 +436,7 @@ async fn responses_request_body_is_uncompressed_when_disabled() { let request_recorder = responses::mount_sse_once(&server, response_body).await; let provider = ModelProviderInfo { - name: "mock".into(), + name: "OpenAI".into(), base_url: Some(format!("{}/v1", server.uri())), env_key: None, env_key_instructions: None, @@ -476,9 +477,11 @@ async fn responses_request_body_is_uncompressed_when_disabled() { session_source.clone(), ); + let auth_manager = + AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); let client = ModelClient::new( Arc::clone(&config), - None, + Some(auth_manager), model_family, otel_manager, provider,