From b077ca79d7e5ade46da1ca8aca2e04aa7f64b0ad Mon Sep 17 00:00:00 2001 From: Saurabh Garg Date: Wed, 24 Jun 2026 19:17:19 -0700 Subject: [PATCH] Route account override through Responses requests --- codex-rs/Cargo.lock | 1 + .../schema/json/ClientRequest.json | 7 ++ .../codex_app_server_protocol.schemas.json | 7 ++ .../codex_app_server_protocol.v2.schemas.json | 7 ++ .../schema/json/v2/TurnStartParams.json | 7 ++ .../schema/typescript/v2/TurnStartParams.ts | 4 + .../src/protocol/v2/tests.rs | 1 + .../src/protocol/v2/turn.rs | 4 + codex-rs/app-server/README.md | 2 +- .../src/message_processor_tracing_tests.rs | 1 + .../src/request_processors/turn_processor.rs | 1 + .../app-server/tests/suite/v2/turn_start.rs | 2 + codex-rs/core/Cargo.toml | 1 + codex-rs/core/src/client.rs | 83 ++++++++++++++++++- codex-rs/core/src/client_tests.rs | 72 +++++++++++++++- codex-rs/core/src/codex_delegate.rs | 1 + codex-rs/core/src/codex_delegate_tests.rs | 1 + codex-rs/core/src/codex_thread.rs | 8 +- codex-rs/core/src/compact_remote.rs | 1 + codex-rs/core/src/compact_remote_v2.rs | 2 + codex-rs/core/src/session/handlers.rs | 22 ++++- codex-rs/core/src/session/mod.rs | 4 + codex-rs/core/src/session/review.rs | 1 + codex-rs/core/src/session/session.rs | 1 + codex-rs/core/src/session/tests.rs | 5 ++ codex-rs/core/src/session/turn.rs | 1 + codex-rs/core/src/session/turn_context.rs | 10 +++ codex-rs/mcp-server/src/codex_tool_runner.rs | 1 + codex-rs/mcp-server/src/message_processor.rs | 1 + codex-rs/protocol/src/protocol.rs | 3 + 30 files changed, 255 insertions(+), 7 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 8cc9ca818a1e..cc7494d3c5af 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2719,6 +2719,7 @@ dependencies = [ "tracing-subscriber", "tracing-test", "url", + "urlencoding", "uuid", "walkdir", "which 8.0.0", diff --git a/codex-rs/app-server-protocol/schema/json/ClientRequest.json b/codex-rs/app-server-protocol/schema/json/ClientRequest.json index 68c242504f7d..69e037f9ba03 100644 --- a/codex-rs/app-server-protocol/schema/json/ClientRequest.json +++ b/codex-rs/app-server-protocol/schema/json/ClientRequest.json @@ -4450,6 +4450,13 @@ }, "TurnStartParams": { "properties": { + "accountRoutingOverride": { + "description": "Optional account routing override for ChatGPT-authenticated Responses requests in this turn.", + "type": [ + "string", + "null" + ] + }, "approvalPolicy": { "anyOf": [ { diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json index 4d417ffb2f53..2f0452d5516e 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json @@ -19983,6 +19983,13 @@ "TurnStartParams": { "$schema": "http://json-schema.org/draft-07/schema#", "properties": { + "accountRoutingOverride": { + "description": "Optional account routing override for ChatGPT-authenticated Responses requests in this turn.", + "type": [ + "string", + "null" + ] + }, "approvalPolicy": { "anyOf": [ { diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json index 0df923a0f4e6..4f598add49c6 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json @@ -17762,6 +17762,13 @@ "TurnStartParams": { "$schema": "http://json-schema.org/draft-07/schema#", "properties": { + "accountRoutingOverride": { + "description": "Optional account routing override for ChatGPT-authenticated Responses requests in this turn.", + "type": [ + "string", + "null" + ] + }, "approvalPolicy": { "anyOf": [ { diff --git a/codex-rs/app-server-protocol/schema/json/v2/TurnStartParams.json b/codex-rs/app-server-protocol/schema/json/v2/TurnStartParams.json index 36b4d2a15d82..62a3f18182fa 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/TurnStartParams.json +++ b/codex-rs/app-server-protocol/schema/json/v2/TurnStartParams.json @@ -498,6 +498,13 @@ } }, "properties": { + "accountRoutingOverride": { + "description": "Optional account routing override for ChatGPT-authenticated Responses requests in this turn.", + "type": [ + "string", + "null" + ] + }, "approvalPolicy": { "anyOf": [ { diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/TurnStartParams.ts b/codex-rs/app-server-protocol/schema/typescript/v2/TurnStartParams.ts index afe1ac6d9488..5e60c60b9a5d 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/TurnStartParams.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/TurnStartParams.ts @@ -11,6 +11,10 @@ import type { SandboxPolicy } from "./SandboxPolicy"; import type { UserInput } from "./UserInput"; export type TurnStartParams = {threadId: string, clientUserMessageId?: string | null, input: Array, /** + * Optional account routing override for ChatGPT-authenticated Responses requests in this + * turn. + */ +accountRoutingOverride?: string | null, /** * Override the working directory for this turn and subsequent turns. */ cwd?: string | null, /** diff --git a/codex-rs/app-server-protocol/src/protocol/v2/tests.rs b/codex-rs/app-server-protocol/src/protocol/v2/tests.rs index 7c78fee4c83e..b40e76e2a058 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/tests.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/tests.rs @@ -3811,6 +3811,7 @@ fn turn_start_params_preserve_explicit_null_service_tier() { thread_id: "thread_123".to_string(), client_user_message_id: None, input: vec![], + account_routing_override: None, responsesapi_client_metadata: None, additional_context: None, environments: None, diff --git a/codex-rs/app-server-protocol/src/protocol/v2/turn.rs b/codex-rs/app-server-protocol/src/protocol/v2/turn.rs index 1fc35d15152c..0e37b729ed2b 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/turn.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/turn.rs @@ -70,6 +70,10 @@ pub struct TurnStartParams { #[ts(optional = nullable)] pub client_user_message_id: Option, pub input: Vec, + /// Optional account routing override for ChatGPT-authenticated Responses requests in this + /// turn. + #[ts(optional = nullable)] + pub account_routing_override: Option, /// Optional metadata to enrich Codex's ResponsesAPI turn metadata. /// /// Entries are flattened into the JSON string sent as diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 71557bb73034..2777fb663d80 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -168,7 +168,7 @@ Example with notification opt-out: - `thread/backgroundTerminals/list` — list running background terminals for a loaded thread (experimental; requires `capabilities.experimentalApi`); returns `data` with the running terminal ids. - `thread/backgroundTerminals/terminate` — terminate one running background terminal by app-server `processId` (experimental; requires `capabilities.experimentalApi`); returns whether a process was terminated. - `thread/rollback` — drop the last N turns from the agent’s in-memory context and persist a rollback marker in the rollout so future resumes see the pruned history; returns the updated `thread` (with `turns` populated) on success. -- `turn/start` — add user input to a thread and begin Codex generation; responds with the initial `turn` object and streams `turn/started`, `item/*`, and `turn/completed` notifications. `clientUserMessageId` is optional; when supplied, the corresponding `userMessage` item echoes it as `clientId`. Experimental `runtimeWorkspaceRoots` replaces the thread-scoped runtime workspace roots used to materialize `:workspace_roots`; paths must be absolute. Prefer experimental `permissions` profile selection by id for permission overrides; the legacy `sandboxPolicy` field is still accepted but cannot be combined with `permissions`. For `collaborationMode`, `settings.developer_instructions: null` means "use built-in instructions for the selected mode". Experimental `multiAgentMode` accepts `none`, `explicitRequestOnly`, or `proactive`; `none` keeps the tools available without injecting mode instructions, and omission keeps the loaded session's current mode. The requested mode is retained for the loaded session without rejecting unsupported configurations, and eligible multi-agent v2 turns use it directly. +- `turn/start` — add user input to a thread and begin Codex generation; responds with the initial `turn` object and streams `turn/started`, `item/*`, and `turn/completed` notifications. `clientUserMessageId` is optional; when supplied, the corresponding `userMessage` item echoes it as `clientId`. `accountRoutingOverride` optionally routes ChatGPT-authenticated Responses requests for this turn. Experimental `runtimeWorkspaceRoots` replaces the thread-scoped runtime workspace roots used to materialize `:workspace_roots`; paths must be absolute. Prefer experimental `permissions` profile selection by id for permission overrides; the legacy `sandboxPolicy` field is still accepted but cannot be combined with `permissions`. For `collaborationMode`, `settings.developer_instructions: null` means "use built-in instructions for the selected mode". Experimental `multiAgentMode` accepts `none`, `explicitRequestOnly`, or `proactive`; `none` keeps the tools available without injecting mode instructions, and omission keeps the loaded session's current mode. The requested mode is retained for the loaded session without rejecting unsupported configurations, and eligible multi-agent v2 turns use it directly. - `thread/inject_items` — append raw Responses API items to a loaded thread’s model-visible history without starting a user turn; returns `{}` on success. - `turn/steer` — add user input to an already in-flight regular turn without starting a new turn; returns the active `turnId` that accepted the input. `clientUserMessageId` is optional; when supplied, the corresponding `userMessage` item echoes it as `clientId`. Review and manual compaction turns reject `turn/steer`. - `turn/interrupt` — request cancellation of an in-flight turn by `(thread_id, turn_id)`; success is an empty `{}` response and the turn finishes with `status: "interrupted"`. diff --git a/codex-rs/app-server/src/message_processor_tracing_tests.rs b/codex-rs/app-server/src/message_processor_tracing_tests.rs index 7c4ae8c98563..879d1081bf20 100644 --- a/codex-rs/app-server/src/message_processor_tracing_tests.rs +++ b/codex-rs/app-server/src/message_processor_tracing_tests.rs @@ -658,6 +658,7 @@ async fn turn_start_jsonrpc_span_parents_core_turn_spans() -> Result<()> { text: "hello".to_string(), text_elements: Vec::new(), }], + account_routing_override: None, responsesapi_client_metadata: None, additional_context: None, cwd: None, diff --git a/codex-rs/app-server/src/request_processors/turn_processor.rs b/codex-rs/app-server/src/request_processors/turn_processor.rs index e0d79a0c3b15..d9c4c5e7db05 100644 --- a/codex-rs/app-server/src/request_processors/turn_processor.rs +++ b/codex-rs/app-server/src/request_processors/turn_processor.rs @@ -534,6 +534,7 @@ impl TurnRequestProcessor { turn_op, self.request_trace_context(&request_id).await, client_user_message_id, + params.account_routing_override, ) .await .map_err(|err| { diff --git a/codex-rs/app-server/tests/suite/v2/turn_start.rs b/codex-rs/app-server/tests/suite/v2/turn_start.rs index 4ff6417ce9b8..67dd84b1814f 100644 --- a/codex-rs/app-server/tests/suite/v2/turn_start.rs +++ b/codex-rs/app-server/tests/suite/v2/turn_start.rs @@ -2558,6 +2558,7 @@ async fn turn_start_updates_sandbox_and_cwd_between_turns_v2() -> Result<()> { text: "first turn".to_string(), text_elements: Vec::new(), }], + account_routing_override: None, responsesapi_client_metadata: None, additional_context: None, cwd: Some(first_cwd.clone()), @@ -2603,6 +2604,7 @@ async fn turn_start_updates_sandbox_and_cwd_between_turns_v2() -> Result<()> { text: "second turn".to_string(), text_elements: Vec::new(), }], + account_routing_override: None, responsesapi_client_metadata: None, additional_context: None, cwd: Some(second_cwd.clone()), diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index 4435a2eb630c..648b1b336cc6 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -117,6 +117,7 @@ toml = { workspace = true } toml_edit = { workspace = true } tracing = { workspace = true, features = ["log"] } url = { workspace = true } +urlencoding = { workspace = true } uuid = { workspace = true, features = ["serde", "v4", "v5"] } which = { workspace = true } whoami = { workspace = true } diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 649f52e54914..b416b75ab1e1 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -90,6 +90,7 @@ use futures::StreamExt; use http::HeaderMap as ApiHeaderMap; use http::HeaderValue; use http::StatusCode as HttpStatusCode; +use http::header::COOKIE; use reqwest::StatusCode; use std::time::Duration; use std::time::Instant; @@ -149,6 +150,7 @@ const X_OPENAI_INTERNAL_CODEX_RESPONSES_LITE_HEADER: &str = "x-openai-internal-codex-responses-lite"; const RESPONSES_ENDPOINT: &str = "/responses"; const RESPONSES_COMPACT_ENDPOINT: &str = "/responses/compact"; +const ACCOUNT_ROUTING_OVERRIDE_COOKIE: &str = "_account_routing_override"; // `/responses/compact` is unary, so the timeout covers the full response rather than one idle // period between stream events. const COMPACT_REQUEST_TIMEOUT_IDLE_MULTIPLIER: u32 = 4; @@ -163,6 +165,37 @@ pub(crate) struct CompactConversationRequestSettings { pub(crate) service_tier: Option, } +fn add_account_routing_override_cookie( + headers: &mut ApiHeaderMap, + account_routing_override: Option<&str>, + auth_mode: Option, + provider: &ModelProviderInfo, +) { + if !provider.is_openai() + || !provider.requires_openai_auth + || provider.env_key.is_some() + || provider.experimental_bearer_token.is_some() + || provider.auth.is_some() + || !auth_mode.is_some_and(AuthMode::has_chatgpt_account) + { + return; + } + let Some(account_routing_override) = account_routing_override else { + return; + }; + let routing_cookie = format!( + "{ACCOUNT_ROUTING_OVERRIDE_COOKIE}={}", + urlencoding::encode(account_routing_override) + ); + let cookie = match headers.get(COOKIE).and_then(|value| value.to_str().ok()) { + Some(existing_cookie) => format!("{existing_cookie}; {routing_cookie}"), + None => routing_cookie, + }; + if let Ok(cookie) = HeaderValue::from_str(&cookie) { + headers.insert(COOKIE, cookie); + } +} + fn session_telemetry_for_request( session_telemetry: &SessionTelemetry, request: &ResponsesApiRequest, @@ -274,6 +307,7 @@ struct LastResponse { #[derive(Debug, Default)] struct WebsocketSession { connection: Option, + account_routing_override: Option, last_request: Option, last_response_rx: Option>, last_response_from_untraced_warmup: bool, @@ -510,6 +544,7 @@ impl ModelClient { session_telemetry: &SessionTelemetry, compaction_trace: &CompactionTraceContext, responses_metadata: &CodexResponsesMetadata, + account_routing_override: Option<&str>, ) -> Result> { if prompt.input.is_empty() { return Ok(Vec::new()); @@ -578,6 +613,12 @@ impl ModelClient { extra_headers.insert(X_OAI_ATTESTATION_HEADER, header_value); } add_responses_lite_header(&mut extra_headers, model_info.use_responses_lite); + add_account_routing_override_cookie( + &mut extra_headers, + account_routing_override, + client_setup.auth.as_ref().map(CodexAuth::api_auth_mode), + self.state.provider.info(), + ); let compact_request_timeout = client_setup .api_provider .stream_idle_timeout @@ -916,11 +957,15 @@ impl ModelClient { session_telemetry: &SessionTelemetry, api_provider: codex_api::Provider, api_auth: SharedAuthProvider, + auth_mode: Option, + account_routing_override: Option<&str>, responses_metadata: &CodexResponsesMetadata, auth_context: AuthRequestTelemetryContext, request_route_telemetry: RequestRouteTelemetry, ) -> std::result::Result { - let headers = self.build_websocket_headers(responses_metadata).await; + let headers = self + .build_websocket_headers(responses_metadata, auth_mode, account_routing_override) + .await; let websocket_telemetry = ModelClientSession::build_websocket_telemetry( session_telemetry, auth_context, @@ -997,6 +1042,8 @@ impl ModelClient { async fn build_websocket_headers( &self, responses_metadata: &CodexResponsesMetadata, + auth_mode: Option, + account_routing_override: Option<&str>, ) -> ApiHeaderMap { let mut headers = build_responses_headers( self.state.beta_features_header.as_deref(), @@ -1024,6 +1071,12 @@ impl ModelClient { HeaderValue::from_static("true"), ); } + add_account_routing_override_cookie( + &mut headers, + account_routing_override, + auth_mode, + self.state.provider.info(), + ); headers } } @@ -1041,6 +1094,17 @@ impl ModelClientSession { Arc::clone(&self.turn_state) } + pub(crate) fn set_account_routing_override( + &mut self, + account_routing_override: Option, + ) { + if self.websocket_session.account_routing_override == account_routing_override { + return; + } + self.reset_websocket_session(); + self.websocket_session.account_routing_override = account_routing_override; + } + fn reset_websocket_session(&mut self) { self.websocket_session.connection = None; self.websocket_session.last_request = None; @@ -1060,6 +1124,7 @@ impl ModelClientSession { responses_metadata: &CodexResponsesMetadata, compression: Compression, use_responses_lite: bool, + auth_mode: Option, ) -> ApiResponsesOptions { ApiResponsesOptions { session_id: Some(responses_metadata.session_id.to_string()), @@ -1079,6 +1144,12 @@ impl ModelClientSession { headers.insert(X_OAI_ATTESTATION_HEADER, header_value); } add_responses_lite_header(&mut headers, use_responses_lite); + add_account_routing_override_cookie( + &mut headers, + self.websocket_session.account_routing_override.as_deref(), + auth_mode, + self.client.state.provider.info(), + ); headers }, compression, @@ -1196,12 +1267,15 @@ impl ModelClientSession { client_setup.api_auth.as_ref(), PendingUnauthorizedRetry::default(), ); + let auth_mode = client_setup.auth.as_ref().map(CodexAuth::api_auth_mode); let connection = self .client .connect_websocket( session_telemetry, client_setup.api_provider, client_setup.api_auth, + auth_mode, + self.websocket_session.account_routing_override.as_deref(), responses_metadata, auth_context, RequestRouteTelemetry::for_endpoint(RESPONSES_ENDPOINT), @@ -1233,6 +1307,7 @@ impl ModelClientSession { session_telemetry, api_provider, api_auth, + auth_mode, responses_metadata, auth_context, request_route_telemetry, @@ -1252,6 +1327,8 @@ impl ModelClientSession { session_telemetry, api_provider, api_auth, + auth_mode, + self.websocket_session.account_routing_override.as_deref(), responses_metadata, auth_context, request_route_telemetry, @@ -1346,6 +1423,7 @@ impl ModelClientSession { responses_metadata, compression, model_info.use_responses_lite, + client_setup.auth.as_ref().map(CodexAuth::api_auth_mode), ) .await; @@ -1461,6 +1539,7 @@ impl ModelClientSession { client_setup.api_auth.as_ref(), pending_retry, ); + let auth_mode = client_setup.auth.as_ref().map(CodexAuth::api_auth_mode); let request = self.client.build_responses_request( &client_setup.api_provider, prompt, @@ -1498,6 +1577,7 @@ impl ModelClientSession { session_telemetry, api_provider: client_setup.api_provider, api_auth: client_setup.api_auth, + auth_mode, responses_metadata, auth_context: request_auth_context, request_route_telemetry: RequestRouteTelemetry::for_endpoint( @@ -2062,6 +2142,7 @@ struct WebsocketConnectParams<'a> { session_telemetry: &'a SessionTelemetry, api_provider: codex_api::Provider, api_auth: SharedAuthProvider, + auth_mode: Option, responses_metadata: &'a CodexResponsesMetadata, auth_context: AuthRequestTelemetryContext, request_route_telemetry: RequestRouteTelemetry, diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index 78dc324c4302..d5dc29b54c6b 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -7,6 +7,7 @@ use super::X_CODEX_PARENT_THREAD_ID_HEADER; use super::X_CODEX_TURN_METADATA_HEADER; use super::X_CODEX_WINDOW_ID_HEADER; use super::X_OPENAI_SUBAGENT_HEADER; +use super::add_account_routing_override_cookie; use crate::AttestationContext; use crate::AttestationProvider; use crate::GenerateAttestationFuture; @@ -154,6 +155,50 @@ fn test_session_telemetry() -> SessionTelemetry { ) } +#[test] +fn account_routing_override_cookie_is_encoded_and_chatgpt_scoped() { + let provider = + ModelProviderInfo::create_openai_provider(Some(CHATGPT_CODEX_BASE_URL.to_string())); + let chatgpt_auth = CodexAuth::create_dummy_chatgpt_auth_for_testing(); + let mut headers = http::HeaderMap::new(); + headers.insert(http::header::COOKIE, http::HeaderValue::from_static("a=b")); + + add_account_routing_override_cookie( + &mut headers, + Some("us cr;foo=bar"), + Some(chatgpt_auth.api_auth_mode()), + &provider, + ); + + assert_eq!( + headers + .get(http::header::COOKIE) + .and_then(|value| value.to_str().ok()), + Some("a=b; _account_routing_override=us%20cr%3Bfoo%3Dbar"), + ); + + let api_key_auth = CodexAuth::from_api_key("test-api-key"); + let mut api_key_headers = http::HeaderMap::new(); + add_account_routing_override_cookie( + &mut api_key_headers, + Some("us_cr"), + Some(api_key_auth.api_auth_mode()), + &provider, + ); + assert_eq!(api_key_headers.get(http::header::COOKIE), None); + + let mut provider_with_api_key = provider.clone(); + provider_with_api_key.env_key = Some("OPENAI_API_KEY".to_string()); + let mut provider_api_key_headers = http::HeaderMap::new(); + add_account_routing_override_cookie( + &mut provider_api_key_headers, + Some("us_cr"), + Some(chatgpt_auth.api_auth_mode()), + &provider_with_api_key, + ); + assert_eq!(provider_api_key_headers.get(http::header::COOKIE), None); +} + #[derive(Default)] struct TagCollectorVisitor { tags: BTreeMap, @@ -637,7 +682,11 @@ async fn websocket_handshake_includes_attestation_for_chatgpt_codex_responses() ); let headers = model_client - .build_websocket_headers(&responses_metadata) + .build_websocket_headers( + &responses_metadata, + /*auth_mode*/ None, + /*account_routing_override*/ None, + ) .await; assert_eq!( @@ -649,6 +698,27 @@ async fn websocket_handshake_includes_attestation_for_chatgpt_codex_responses() assert_eq!(attestation_calls.load(Ordering::Relaxed), 1); } +#[test] +fn changing_account_routing_override_clears_cached_websocket_state() { + let mut client_session = test_model_client(SessionSource::Cli).new_session(); + client_session.websocket_session.account_routing_override = Some("old".to_string()); + client_session + .websocket_session + .last_response_from_untraced_warmup = true; + + client_session.set_account_routing_override(Some("new".to_string())); + + assert_eq!( + client_session.websocket_session.account_routing_override, + Some("new".to_string()), + ); + assert!( + !client_session + .websocket_session + .last_response_from_untraced_warmup + ); +} + #[tokio::test] async fn non_chatgpt_codex_endpoints_omit_attestation_generation() { let (model_client, attestation_calls) = diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 640e1f8dfa9c..e125ee62816b 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -241,6 +241,7 @@ pub(crate) async fn run_codex_thread_one_shot( id: "shutdown".to_string(), op: Op::Shutdown {}, client_user_message_id: None, + account_routing_override: None, trace: None, }) .await; diff --git a/codex-rs/core/src/codex_delegate_tests.rs b/codex-rs/core/src/codex_delegate_tests.rs index 6d51a0032b24..ab664ffcce03 100644 --- a/codex-rs/core/src/codex_delegate_tests.rs +++ b/codex-rs/core/src/codex_delegate_tests.rs @@ -132,6 +132,7 @@ async fn forward_ops_preserves_submission_trace_context() { id: "sub-1".to_string(), op: Op::Interrupt, client_user_message_id: None, + account_routing_override: None, trace: Some(codex_protocol::protocol::W3cTraceContext { traceparent: Some( "00-1234567890abcdef1234567890abcdef-1234567890abcdef-01".to_string(), diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index dccfd4a76306..89c2cd01abd9 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -257,6 +257,7 @@ impl CodexThread { op: Op, trace: Option, client_user_message_id: Option, + account_routing_override: Option, ) -> CodexResult { self.codex .session @@ -265,7 +266,12 @@ impl CodexThread { .ensure_execution_capacity_for_op(self.session_configured.thread_id, &op) .await?; self.codex - .submit_user_input_with_client_user_message_id(op, trace, client_user_message_id) + .submit_user_input_with_client_user_message_id( + op, + trace, + client_user_message_id, + account_routing_override, + ) .await } diff --git a/codex-rs/core/src/compact_remote.rs b/codex-rs/core/src/compact_remote.rs index 6b251562f798..acffbb8eeb12 100644 --- a/codex-rs/core/src/compact_remote.rs +++ b/codex-rs/core/src/compact_remote.rs @@ -263,6 +263,7 @@ async fn run_remote_compact_task_inner_impl( &turn_context.session_telemetry, &compaction_trace, &responses_metadata, + turn_context.account_routing_override.as_deref(), ) .await?; let (new_window_number, new_window_ids) = sess.advance_auto_compact_window().await; diff --git a/codex-rs/core/src/compact_remote_v2.rs b/codex-rs/core/src/compact_remote_v2.rs index 6e110045bbe4..76c200012e88 100644 --- a/codex-rs/core/src/compact_remote_v2.rs +++ b/codex-rs/core/src/compact_remote_v2.rs @@ -267,6 +267,8 @@ async fn run_remote_compact_task_inner_impl( Some(client_session) => client_session, None => { owned_client_session = sess.services.model_client.new_session(); + owned_client_session + .set_account_routing_override(turn_context.account_routing_override.clone()); &mut owned_client_session } }; diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 39e91d6cf95d..8e5545d6a1d0 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -85,8 +85,16 @@ pub async fn user_input_or_turn( sub_id: String, op: Op, client_user_message_id: Option, + account_routing_override: Option, ) { - user_input_or_turn_inner(sess, sub_id, op, client_user_message_id).await; + user_input_or_turn_inner( + sess, + sub_id, + op, + client_user_message_id, + account_routing_override, + ) + .await; } pub async fn update_thread_settings( @@ -188,6 +196,7 @@ pub(super) async fn user_input_or_turn_inner( sub_id: String, op: Op, client_user_message_id: Option, + account_routing_override: Option, ) { let Op::UserInput { items, @@ -206,6 +215,7 @@ pub(super) async fn user_input_or_turn_inner( SessionSettingsUpdate::default() }; updates.final_output_json_schema = Some(final_output_json_schema); + updates.account_routing_override = account_routing_override; let Ok(current_context) = sess.new_turn_with_sub_id(sub_id.clone(), updates).await else { // new_turn_with_sub_id already emits the error event. @@ -758,8 +768,14 @@ pub(super) async fn submission_loop( false } Op::UserInput { .. } => { - user_input_or_turn(&sess, sub.id.clone(), sub.op, sub.client_user_message_id) - .await; + user_input_or_turn( + &sess, + sub.id.clone(), + sub.op, + sub.client_user_message_id, + sub.account_routing_override, + ) + .await; false } Op::ThreadSettings { thread_settings } => { diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index daf981258127..cfc577a068cc 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -731,6 +731,7 @@ impl Codex { id: id.clone(), op, client_user_message_id: None, + account_routing_override: None, trace, }; self.submit_with_id(sub).await?; @@ -742,6 +743,7 @@ impl Codex { op: Op, trace: Option, client_user_message_id: Option, + account_routing_override: Option, ) -> CodexResult { debug_assert!(matches!(op, Op::UserInput { .. })); let id = new_submission_id(); @@ -749,6 +751,7 @@ impl Codex { id: id.clone(), op, client_user_message_id, + account_routing_override, trace, }; self.submit_with_id(sub).await?; @@ -1198,6 +1201,7 @@ impl Session { thread_settings: Default::default(), }, /*client_user_message_id*/ None, + /*account_routing_override*/ None, ) .await; } diff --git a/codex-rs/core/src/session/review.rs b/codex-rs/core/src/session/review.rs index c55afbf14195..091d7b9f90df 100644 --- a/codex-rs/core/src/session/review.rs +++ b/codex-rs/core/src/session/review.rs @@ -138,6 +138,7 @@ pub(super) async fn spawn_review_thread( #[allow(deprecated)] cwd: parent_turn_context.cwd.clone(), final_output_json_schema: None, + account_routing_override: None, dynamic_tools: parent_turn_context.dynamic_tools.clone(), turn_metadata_state, extension_data, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index af64e0bef5dc..bd57f6077e97 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -434,6 +434,7 @@ pub(crate) struct SessionSettingsUpdate { pub(crate) reasoning_summary: Option, pub(crate) service_tier: Option>, pub(crate) final_output_json_schema: Option>, + pub(crate) account_routing_override: Option, pub(crate) personality: Option, pub(crate) app_server_client_name: Option, pub(crate) app_server_client_version: Option, diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 983ed200942d..0c6a1464d370 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -6437,6 +6437,7 @@ async fn submit_with_id_captures_current_span_trace_context() { id: "sub-1".into(), op: Op::Interrupt, client_user_message_id: None, + account_routing_override: None, trace: None, }) .await @@ -6509,6 +6510,7 @@ fn submission_dispatch_span_prefers_submission_trace_context() { id: "sub-1".into(), op: Op::Interrupt, client_user_message_id: None, + account_routing_override: None, trace: Some(submission_trace), }) }); @@ -6536,6 +6538,7 @@ fn submission_dispatch_span_uses_debug_for_realtime_audio() { }, }), client_user_message_id: None, + account_routing_override: None, trace: None, }); @@ -6602,6 +6605,7 @@ async fn user_turn_updates_approvals_reviewer() { }, }, /*client_user_message_id*/ None, + /*account_routing_override*/ None, ) .await; @@ -6849,6 +6853,7 @@ async fn spawn_task_turn_span_inherits_dispatch_trace_context() { id: "sub-1".into(), op: Op::Interrupt, client_user_message_id: None, + account_routing_override: None, trace: Some(submission_trace.clone()), }); let dispatch_span_id = dispatch_span.context().span().span_context().span_id(); diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 3094bc77a795..db689b0a8a0f 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -149,6 +149,7 @@ pub(crate) async fn run_turn( ) -> CodexResult> { let mut client_session = prewarmed_client_session.unwrap_or_else(|| sess.services.model_client.new_session()); + client_session.set_account_routing_override(turn_context.account_routing_override.clone()); // TODO(ccunningham): Pre-turn compaction runs before context updates and the // new user message are recorded. Estimate pending incoming items (context // diffs/full reinjection + user input) and trigger compaction preemptively diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index da526d49111b..5e5c82770461 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -136,6 +136,7 @@ pub struct TurnContext { pub(crate) available_models: Vec, pub(crate) unified_exec_shell_mode: UnifiedExecShellMode, pub(crate) final_output_json_schema: Option, + pub(crate) account_routing_override: Option, pub(crate) dynamic_tools: Vec, pub(crate) turn_metadata_state: Arc, pub(crate) extension_data: Arc, @@ -286,6 +287,7 @@ impl TurnContext { available_models, unified_exec_shell_mode: self.unified_exec_shell_mode.clone(), final_output_json_schema: self.final_output_json_schema.clone(), + account_routing_override: self.account_routing_override.clone(), dynamic_tools: self.dynamic_tools.clone(), turn_metadata_state: self.turn_metadata_state.clone(), extension_data: Arc::clone(&self.extension_data), @@ -573,6 +575,7 @@ impl Session { available_models, unified_exec_shell_mode, final_output_json_schema: None, + account_routing_override: None, dynamic_tools: session_configuration.dynamic_tools.clone(), turn_metadata_state, extension_data, @@ -649,6 +652,7 @@ impl Session { sub_id, session_configuration, updates.final_output_json_schema, + updates.account_routing_override, ) .await) } @@ -658,11 +662,13 @@ impl Session { sub_id: String, session_configuration: SessionConfiguration, final_output_json_schema: Option>, + account_routing_override: Option, ) -> Arc { self.new_turn_context_from_configuration( sub_id, session_configuration, final_output_json_schema, + account_routing_override, TurnMultiAgentRuntime::ResolveAndStore, ) .await @@ -677,6 +683,7 @@ impl Session { sub_id, session_configuration, /*final_output_json_schema*/ None, + /*account_routing_override*/ None, TurnMultiAgentRuntime::Preview, ) .await @@ -688,6 +695,7 @@ impl Session { sub_id: String, session_configuration: SessionConfiguration, final_output_json_schema: Option>, + account_routing_override: Option, multi_agent_runtime: TurnMultiAgentRuntime, ) -> Arc { let turn_environments = self.services.turn_environments.snapshot().await; @@ -778,6 +786,7 @@ impl Session { if let Some(final_schema) = final_output_json_schema { turn_context.final_output_json_schema = final_schema; } + turn_context.account_routing_override = account_routing_override; let turn_context = Arc::new(turn_context); if turn_context .environments @@ -822,6 +831,7 @@ impl Session { sub_id, session_configuration, /*final_output_json_schema*/ None, + /*account_routing_override*/ None, ) .await } diff --git a/codex-rs/mcp-server/src/codex_tool_runner.rs b/codex-rs/mcp-server/src/codex_tool_runner.rs index cbc113ff79e0..359768311877 100644 --- a/codex-rs/mcp-server/src/codex_tool_runner.rs +++ b/codex-rs/mcp-server/src/codex_tool_runner.rs @@ -114,6 +114,7 @@ pub async fn run_codex_tool_session( thread_settings: Default::default(), }, client_user_message_id: None, + account_routing_override: None, trace: None, }; diff --git a/codex-rs/mcp-server/src/message_processor.rs b/codex-rs/mcp-server/src/message_processor.rs index d211ff0a32fc..5d2c5fcfb7e2 100644 --- a/codex-rs/mcp-server/src/message_processor.rs +++ b/codex-rs/mcp-server/src/message_processor.rs @@ -546,6 +546,7 @@ impl MessageProcessor { id: request_id_string, op: codex_protocol::protocol::Op::Interrupt, client_user_message_id: None, + account_routing_override: None, trace: None, }) .await diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 138b36f10f47..367cc8efe4a5 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -162,6 +162,9 @@ pub struct Submission { pub op: Op, /// Client-provided id for the user message represented by `Op::UserInput`. pub client_user_message_id: Option, + /// Optional account routing override for ChatGPT-authenticated Responses requests in this + /// turn. + pub account_routing_override: Option, /// Optional W3C trace carrier propagated across async submission handoffs. pub trace: Option, }