diff --git a/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs b/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs index 8f8d670313d6..a4d1ac298d9d 100644 --- a/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs +++ b/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs @@ -95,7 +95,14 @@ async fn standalone_image_generation_returns_saved_path_hint_to_model() -> Resul .build() .await?; timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - start_image_generation_turn(&mut mcp).await?; + start_image_generation_turn( + &mut mcp, + ThreadStartParams { + service_name: Some("chatgpt_cca".to_string()), + ..Default::default() + }, + ) + .await?; let completed = timeout( DEFAULT_READ_TIMEOUT, @@ -123,6 +130,23 @@ async fn standalone_image_generation_returns_saved_path_hint_to_model() -> Resul assert_eq!(result, RESULT); assert_eq!(std::fs::read(&saved_path)?, TINY_PNG_BYTES); + let image_request = server + .received_requests() + .await + .context("failed to fetch received requests")? + .into_iter() + .find(|request| request.url.path() == "/api/codex/images/generations") + .context("image generation request should be sent")?; + assert_eq!( + image_request + .headers + .get("originator") + .context("standalone image generation should include the thread originator")? + .to_str() + .context("standalone image generation originator should be valid ASCII")?, + "chatgpt_cca" + ); + let requests = response_mock.requests(); assert_eq!(requests.len(), 2); let output = requests[1].function_call_output(call_id); @@ -196,7 +220,7 @@ async fn standalone_image_generation_failure_emits_terminal_item() -> Result<()> .build() .await?; timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - start_image_generation_turn(&mut mcp).await?; + start_image_generation_turn(&mut mcp, ThreadStartParams::default()).await?; let completed = timeout( DEFAULT_READ_TIMEOUT, @@ -325,7 +349,7 @@ async fn standalone_image_generation_is_exposed_in_code_mode_only() -> Result<() .build() .await?; timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - start_image_generation_turn(&mut mcp).await?; + start_image_generation_turn(&mut mcp, ThreadStartParams::default()).await?; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), @@ -391,7 +415,7 @@ generatedImage(result); .build() .await?; timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - start_image_generation_turn(&mut mcp).await?; + start_image_generation_turn(&mut mcp, ThreadStartParams::default()).await?; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), @@ -420,9 +444,13 @@ generatedImage(result); Ok(()) } -async fn start_image_generation_turn(mcp: &mut TestAppServer) -> Result<()> { +async fn start_image_generation_turn( + mcp: &mut TestAppServer, + thread_start_params: ThreadStartParams, +) -> Result<()> { start_turn( mcp, + thread_start_params, vec![V2UserInput::Text { text: "Generate an image".to_string(), text_elements: Vec::new(), @@ -474,7 +502,15 @@ async fn run_image_edit_test( .build() .await?; timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - start_turn(&mut mcp, input).await?; + start_turn( + &mut mcp, + ThreadStartParams { + service_name: Some("chatgpt_cca".to_string()), + ..Default::default() + }, + input, + ) + .await?; timeout( DEFAULT_READ_TIMEOUT, wait_for_image_generation_completed(&mut mcp), @@ -491,16 +527,29 @@ async fn run_image_edit_test( .received_requests() .await .context("failed to fetch received requests")?; - Ok(requests + let image_request = requests .iter() .find(|request| request.url.path() == "/api/codex/images/edits") - .context("image edit request should be sent")? - .body_json::()?) + .context("image edit request should be sent")?; + assert_eq!( + image_request + .headers + .get("originator") + .context("standalone image edit should include the thread originator")? + .to_str() + .context("standalone image edit originator should be valid ASCII")?, + "chatgpt_cca" + ); + Ok(image_request.body_json::()?) } -async fn start_turn(mcp: &mut TestAppServer, input: Vec) -> Result<()> { +async fn start_turn( + mcp: &mut TestAppServer, + thread_start_params: ThreadStartParams, + input: Vec, +) -> Result<()> { let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams::default()) + .send_thread_start_request_with_auto_env(thread_start_params) .await?; let thread_resp: JSONRPCResponse = timeout( DEFAULT_READ_TIMEOUT, diff --git a/codex-rs/app-server/tests/suite/v2/web_search.rs b/codex-rs/app-server/tests/suite/v2/web_search.rs index d34f14e57a7f..0327c65d536e 100644 --- a/codex-rs/app-server/tests/suite/v2/web_search.rs +++ b/codex-rs/app-server/tests/suite/v2/web_search.rs @@ -100,7 +100,10 @@ async fn standalone_web_search_round_trips_output() -> Result<()> { timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams::default()) + .send_thread_start_request_with_auto_env(ThreadStartParams { + service_name: Some("chatgpt_cca".to_string()), + ..Default::default() + }) .await?; let thread_resp: JSONRPCResponse = timeout( DEFAULT_READ_TIMEOUT, @@ -159,6 +162,15 @@ async fn standalone_web_search_round_trips_output() -> Result<()> { ); let search_request = search_request(&server).await?; + assert_eq!( + search_request + .headers + .get("originator") + .context("standalone search should include the thread originator")? + .to_str() + .context("standalone search originator should be valid ASCII")?, + "chatgpt_cca" + ); let search_body = search_request .body_json::() .context("search request body should be JSON")?; diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 7b0620f86049..283506f25f40 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -69,6 +69,7 @@ use codex_login::AuthManager; use codex_login::CodexAuth; use codex_login::RefreshTokenError; use codex_login::UnauthorizedRecovery; +use codex_login::default_client::add_originator_header; use codex_login::default_client::build_default_reqwest_client_for_route; use codex_otel::SessionTelemetry; use codex_otel::current_span_w3c_trace_context; @@ -1884,22 +1885,6 @@ fn build_responses_headers( headers } -pub(crate) fn add_originator_header(headers: &mut ApiHeaderMap, originator: &str) { - let default_originator = codex_login::default_client::originator(); - if originator == default_originator.value.as_str() { - return; - } - - match HeaderValue::from_str(originator) { - Ok(header_value) => { - headers.insert("originator", header_value); - } - Err(err) => { - warn!("ignoring invalid thread originator header value: {err}"); - } - } -} - fn add_responses_lite_header(headers: &mut ApiHeaderMap, use_responses_lite: bool) { if use_responses_lite { headers.insert( diff --git a/codex-rs/core/src/realtime_conversation.rs b/codex-rs/core/src/realtime_conversation.rs index b1b86b7b622b..33106868ca1c 100644 --- a/codex-rs/core/src/realtime_conversation.rs +++ b/codex-rs/core/src/realtime_conversation.rs @@ -1,5 +1,4 @@ use crate::client::ModelClient; -use crate::client::add_originator_header; use crate::realtime_context::build_realtime_startup_context; use crate::realtime_context::truncate_realtime_text_to_token_budget; use crate::realtime_prompt::prepare_realtime_backend_prompt; @@ -25,6 +24,7 @@ use codex_api::map_api_error; use codex_config::config_toml::RealtimeWsMode; use codex_config::config_toml::RealtimeWsVersion; use codex_login::CodexAuth; +use codex_login::default_client::add_originator_header; use codex_login::default_client::default_headers; use codex_login::read_openai_api_key_from_env; use codex_model_provider_info::ModelProviderInfo; diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index d931a136c622..02293ffe09fd 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -583,6 +583,9 @@ impl Session { roots } }; + thread_extension_init.insert(codex_extension_api::ThreadOriginator( + session_configuration.originator.clone(), + )); let mcp_thread_init = thread_extension_init.clone(); let thread_extension_data = codex_extension_api::ExtensionData::new_with_init( thread_id.to_string(), diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 7bbbb9efe62d..70599398bbd2 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -204,7 +204,13 @@ pub struct StartThreadOptions { fn originator_from_service_name(service_name: Option<&str>) -> Option { let service_name = service_name?.trim(); - for originator in ["codex_work_desktop", "codex_work_web", "codex_work_mobile"] { + for originator in [ + "codex_work_desktop", + "codex_work_web", + "codex_work_mobile", + "codex_work_cca", + "chatgpt_cca", + ] { if service_name.eq_ignore_ascii_case(originator) { return Some(originator.to_string()); } diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index 3980471611f7..ce3589357cd1 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -135,6 +135,24 @@ fn effective_originator_prefers_thread_scoped_sources_before_env_originator() { Some("inherited_originator"), "codex_work_mobile", ), + ( + Some("codex_work_cca"), + Some("persisted_originator"), + Some("inherited_originator"), + "codex_work_cca", + ), + ( + Some("chatgpt_cca"), + Some("persisted_originator"), + Some("inherited_originator"), + "chatgpt_cca", + ), + ( + Some("chatgpt_cca_extra"), + Some("persisted_originator"), + Some("inherited_originator"), + "persisted_originator", + ), ( None, Some("persisted_originator"), diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index 0ddf5a75d9ea..b8cb2cecc0c8 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -29,6 +29,7 @@ pub use prompt::PromptSlot; pub use skill_invocation::SkillInvocationInput; pub use skill_invocation::SkillInvocationKind; pub use thread_lifecycle::ThreadIdleInput; +pub use thread_lifecycle::ThreadOriginator; pub use thread_lifecycle::ThreadResumeInput; pub use thread_lifecycle::ThreadStartInput; pub use thread_lifecycle::ThreadStopInput; diff --git a/codex-rs/ext/extension-api/src/contributors/thread_lifecycle.rs b/codex-rs/ext/extension-api/src/contributors/thread_lifecycle.rs index 2933fb6b13d5..23852a2df0b4 100644 --- a/codex-rs/ext/extension-api/src/contributors/thread_lifecycle.rs +++ b/codex-rs/ext/extension-api/src/contributors/thread_lifecycle.rs @@ -2,6 +2,14 @@ use crate::ExtensionData; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::TurnEnvironmentSelection; +/// Trusted, host-resolved billing attribution for a thread. +/// +/// Extensions may forward this value to first-party APIs. It is seeded by Core +/// after resolving persisted and host-provided originator state, rather than +/// from model- or tool-controlled input. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct ThreadOriginator(pub String); + /// Input supplied when the host starts a runtime for a thread. pub struct ThreadStartInput<'a, C> { /// Host configuration visible at thread start. diff --git a/codex-rs/ext/extension-api/src/lib.rs b/codex-rs/ext/extension-api/src/lib.rs index 7722c74d6e9f..0c7ddce86df2 100644 --- a/codex-rs/ext/extension-api/src/lib.rs +++ b/codex-rs/ext/extension-api/src/lib.rs @@ -47,6 +47,7 @@ pub use contributors::SkillInvocationInput; pub use contributors::SkillInvocationKind; pub use contributors::ThreadIdleInput; pub use contributors::ThreadLifecycleContributor; +pub use contributors::ThreadOriginator; pub use contributors::ThreadResumeInput; pub use contributors::ThreadStartInput; pub use contributors::ThreadStopInput; diff --git a/codex-rs/ext/image-generation/src/backend.rs b/codex-rs/ext/image-generation/src/backend.rs index 9b837772c7b6..58df602426e7 100644 --- a/codex-rs/ext/image-generation/src/backend.rs +++ b/codex-rs/ext/image-generation/src/backend.rs @@ -3,6 +3,7 @@ use codex_api::ImageGenerationRequest; use codex_api::ImageResponse; use codex_api::ImagesClient; use codex_api::ReqwestTransport; +use codex_login::default_client::add_originator_header; use codex_login::default_client::build_reqwest_client; use codex_model_provider::SharedModelProvider; use http::HeaderMap; @@ -10,12 +11,16 @@ use http::HeaderMap; #[derive(Clone)] pub(crate) struct CodexImagesBackend { provider: SharedModelProvider, + originator: Option, } impl CodexImagesBackend { /// Creates a backend that sends image requests through the active model provider. - pub(crate) fn new(provider: SharedModelProvider) -> Self { - Self { provider } + pub(crate) fn new(provider: SharedModelProvider, originator: Option) -> Self { + Self { + provider, + originator, + } } /// Resolves the provider and auth required for the current image API request. @@ -44,7 +49,7 @@ impl CodexImagesBackend { ) -> Result { self.client() .await? - .generate(&request, HeaderMap::new()) + .generate(&request, image_request_headers(self.originator.as_deref())) .await .map_err(|err| err.to_string()) } @@ -53,8 +58,16 @@ impl CodexImagesBackend { pub(crate) async fn edit(&self, request: ImageEditRequest) -> Result { self.client() .await? - .edit(&request, HeaderMap::new()) + .edit(&request, image_request_headers(self.originator.as_deref())) .await .map_err(|err| err.to_string()) } } + +fn image_request_headers(originator: Option<&str>) -> HeaderMap { + let mut headers = HeaderMap::new(); + if let Some(originator) = originator { + add_originator_header(&mut headers, originator); + } + headers +} diff --git a/codex-rs/ext/image-generation/src/extension.rs b/codex-rs/ext/image-generation/src/extension.rs index 18abc2528900..38a24eed7c63 100644 --- a/codex-rs/ext/image-generation/src/extension.rs +++ b/codex-rs/ext/image-generation/src/extension.rs @@ -6,6 +6,7 @@ use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionFuture; use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::ThreadLifecycleContributor; +use codex_extension_api::ThreadOriginator; use codex_extension_api::ThreadStartInput; use codex_extension_api::ToolCall; use codex_extension_api::ToolContributor; @@ -94,10 +95,12 @@ impl ToolContributor for ImageGenerationExtension { } vec![Arc::new(ImageGenerationTool::new( - CodexImagesBackend::new(create_model_provider( - config.provider.clone(), - Some(self.auth_manager.clone()), - )), + CodexImagesBackend::new( + create_model_provider(config.provider.clone(), Some(self.auth_manager.clone())), + thread_store + .get::() + .map(|originator| originator.0.clone()), + ), config.save_root.clone(), thread_store.level_id().to_string(), ))] diff --git a/codex-rs/ext/web-search/src/extension.rs b/codex-rs/ext/web-search/src/extension.rs index 84e3eeef9d9d..d59a34649dd0 100644 --- a/codex-rs/ext/web-search/src/extension.rs +++ b/codex-rs/ext/web-search/src/extension.rs @@ -14,6 +14,7 @@ use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionFuture; use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::ThreadLifecycleContributor; +use codex_extension_api::ThreadOriginator; use codex_extension_api::ThreadStartInput; use codex_extension_api::ToolContributor; use codex_login::AuthManager; @@ -134,6 +135,9 @@ impl ToolContributor for WebSearchExtension { Some(self.auth_manager.clone()), ), settings: config.settings.clone(), + originator: thread_store + .get::() + .map(|originator| originator.0.clone()), })] } } diff --git a/codex-rs/ext/web-search/src/tool.rs b/codex-rs/ext/web-search/src/tool.rs index 93d18dc58012..127be6854c83 100644 --- a/codex-rs/ext/web-search/src/tool.rs +++ b/codex-rs/ext/web-search/src/tool.rs @@ -18,6 +18,7 @@ use codex_extension_api::parse_tool_input_schema_without_compaction; use codex_extension_items::ExtensionItem; use codex_extension_items::web_search::WebSearchAction; use codex_extension_items::web_search::WebSearchItem; +use codex_login::default_client::add_originator_header; use codex_login::default_client::build_reqwest_client; use codex_model_provider::SharedModelProvider; use codex_protocol::models::WebSearchAction as CoreWebSearchAction; @@ -44,6 +45,7 @@ pub(crate) struct WebSearchTool { pub(crate) session_id: String, pub(crate) provider: SharedModelProvider, pub(crate) settings: SearchSettings, + pub(crate) originator: Option, } impl ToolExecutor for WebSearchTool { @@ -115,12 +117,10 @@ impl WebSearchTool { u64::try_from(call.truncation_policy.token_budget()).unwrap_or(u64::MAX), ), }; - let mut extra_headers = HeaderMap::new(); - if let Some(turn_metadata) = call.codex_turn_metadata.as_deref() - && let Ok(header_value) = HeaderValue::from_str(turn_metadata) - { - extra_headers.insert(X_CODEX_TURN_METADATA_HEADER, header_value); - } + let extra_headers = search_request_headers( + self.originator.as_deref(), + call.codex_turn_metadata.as_deref(), + ); call.turn_item_emitter .emit_started(extension_turn_item( WebSearchItem { @@ -174,6 +174,20 @@ impl WebSearchTool { } } +fn search_request_headers(originator: Option<&str>, turn_metadata: Option<&str>) -> HeaderMap { + let mut headers = HeaderMap::new(); + if let Some(turn_metadata) = turn_metadata + && let Ok(header_value) = HeaderValue::from_str(turn_metadata) + { + headers.insert(X_CODEX_TURN_METADATA_HEADER, header_value); + } + + if let Some(originator) = originator { + add_originator_header(&mut headers, originator); + } + headers +} + fn parse_commands(call: &ToolCall) -> Result { let arguments = call.function_arguments()?; if arguments.trim().is_empty() { @@ -245,6 +259,25 @@ mod tests { use pretty_assertions::assert_eq; use super::command_action; + use super::search_request_headers; + use codex_core::X_CODEX_TURN_METADATA_HEADER; + + #[test] + fn search_request_headers_forward_thread_originator_and_turn_metadata() { + let headers = search_request_headers(Some("chatgpt_cca"), Some("turn-metadata")); + assert_eq!( + headers + .get("originator") + .and_then(|value| value.to_str().ok()), + Some("chatgpt_cca") + ); + assert_eq!( + headers + .get(X_CODEX_TURN_METADATA_HEADER) + .and_then(|value| value.to_str().ok()), + Some("turn-metadata") + ); + } #[test] fn command_action_reports_queries_and_navigation_detail() { diff --git a/codex-rs/login/src/auth/default_client.rs b/codex-rs/login/src/auth/default_client.rs index ac153ea8a4ae..6e1df1002e52 100644 --- a/codex-rs/login/src/auth/default_client.rs +++ b/codex-rs/login/src/auth/default_client.rs @@ -127,6 +127,26 @@ pub fn originator() -> Originator { get_originator_value(/*provided*/ None) } +/// Adds a valid, non-default thread originator override to request headers. +/// +/// The default client already supplies the process originator. Thread-scoped callers should use +/// this helper to override that value only when the thread originator differs. +pub fn add_originator_header(headers: &mut HeaderMap, originator_value: &str) { + let default_originator = originator(); + if originator_value == default_originator.value.as_str() { + return; + } + + match HeaderValue::from_str(originator_value) { + Ok(header_value) => { + headers.insert("originator", header_value); + } + Err(err) => { + tracing::warn!("ignoring invalid thread originator header value: {err}"); + } + } +} + pub fn is_first_party_originator(originator_value: &str) -> bool { originator_value == DEFAULT_ORIGINATOR || originator_value == "codex-tui" diff --git a/codex-rs/login/src/auth/default_client_tests.rs b/codex-rs/login/src/auth/default_client_tests.rs index d873c70e0dc1..23bd54791bc3 100644 --- a/codex-rs/login/src/auth/default_client_tests.rs +++ b/codex-rs/login/src/auth/default_client_tests.rs @@ -67,6 +67,54 @@ fn is_first_party_chat_originator_matches_known_values() { assert_eq!(is_first_party_chat_originator("codex_vscode"), false); } +#[test] +fn add_originator_header_inserts_non_default_originator() { + let default_originator = originator(); + let thread_originator = if default_originator.value == "chatgpt_cca" { + "codex_work_cca" + } else { + "chatgpt_cca" + }; + let mut headers = HeaderMap::new(); + + add_originator_header(&mut headers, thread_originator); + + assert_eq!( + headers + .get("originator") + .and_then(|value| value.to_str().ok()), + Some(thread_originator) + ); +} + +#[test] +fn add_originator_header_preserves_provider_default() { + let default_originator = originator(); + let mut headers = HeaderMap::new(); + headers.insert( + "originator", + HeaderValue::from_static("provider-originator"), + ); + + add_originator_header(&mut headers, &default_originator.value); + + assert_eq!( + headers + .get("originator") + .and_then(|value| value.to_str().ok()), + Some("provider-originator") + ); +} + +#[test] +fn add_originator_header_omits_invalid_originator() { + let mut headers = HeaderMap::new(); + + add_originator_header(&mut headers, "invalid\noriginator"); + + assert!(headers.is_empty()); +} + #[tokio::test] async fn test_create_client_sets_default_headers() { skip_if_no_network!();