github/codex
4
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-20 00:18:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

Forward thread originators from standalone extensions (#33677)

Browse Source 
## Why

Standalone web search and image requests need to preserve the trusted,
thread-scoped originator used for billing attribution.

## What changed

- Seed extension data with the resolved thread originator and recognize
  `codex_work_cca` and `chatgpt_cca` service names as originators.
- Forward non-default originators from the web search and image generation
  extensions with the shared originator-header helper.

## Testing

- Cover originator resolution, header validation, and forwarding for web
  search, image generation, and image editing.

GitOrigin-RevId: 165dace8e4d24e4bd92f234bff5ec77a21a7ff6e
This commit is contained in:
charleschen-oai
2026-07-16 19:31:26 +00:00
committed by copyberry
parent 71448a29e7
commit 78ba047bda
16 changed files with 248 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server/tests/suite/v2/imagegen_extension.rs
+60 -11
View File
@@ -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::<serde_json::Value>()?)
.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::<serde_json::Value>()?)
}
async fn start_turn(mcp: &mut TestAppServer, input: Vec<V2UserInput>) -> Result<()> {
async fn start_turn(
mcp: &mut TestAppServer,
thread_start_params: ThreadStartParams,
input: Vec<V2UserInput>,
) -> 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,
codex-rs/app-server/tests/suite/v2/web_search.rs
+13 -1
View File
@@ -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::<Value>()
.context("search request body should be JSON")?;
codex-rs/core/src/client.rs
+1 -16
View File
@@ -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(
codex-rs/core/src/realtime_conversation.rs
+1 -1
View File
@@ -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;
codex-rs/core/src/session/session.rs
+3
View File
@@ -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(),
codex-rs/core/src/thread_manager.rs
+7 -1
View File
@@ -204,7 +204,13 @@ pub struct StartThreadOptions {
fn originator_from_service_name(service_name: Option<&str>) -> Option<String> {
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());
}
codex-rs/core/src/thread_manager_tests.rs
+18
View File
@@ -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"),
codex-rs/ext/extension-api/src/contributors.rs
+1
View File
@@ -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;
codex-rs/ext/extension-api/src/contributors/thread_lifecycle.rs
+8
View File
@@ -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.
codex-rs/ext/extension-api/src/lib.rs
+1
View File
@@ -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;
codex-rs/ext/image-generation/src/backend.rs
+17 -4
View File
@@ -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<String>,
}
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<String>) -> Self {
Self {
provider,
originator,
}
}
/// Resolves the provider and auth required for the current image API request.
@@ -44,7 +49,7 @@ impl CodexImagesBackend {
) -> Result<ImageResponse, String> {
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<ImageResponse, String> {
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
}
codex-rs/ext/image-generation/src/extension.rs
+7 -4
View File
@@ -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::<ThreadOriginator>()
.map(|originator| originator.0.clone()),
),
config.save_root.clone(),
thread_store.level_id().to_string(),
))]
codex-rs/ext/web-search/src/extension.rs
+4
View File
@@ -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::<ThreadOriginator>()
.map(|originator| originator.0.clone()),
})]
}
}
codex-rs/ext/web-search/src/tool.rs
+39 -6
View File
@@ -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<String>,
}
impl ToolExecutor<ToolCall> 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<SearchCommands, FunctionCallError> {
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() {
codex-rs/login/src/auth/default_client.rs
+20
View File
@@ -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"
codex-rs/login/src/auth/default_client_tests.rs
+48
View File
@@ -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!();