diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 6b5b5afdf6..5a32a8e18d 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3776,7 +3776,6 @@ dependencies = [ "chrono", "codex-file-search", "codex-git-utils", - "codex-login", "codex-otel", "codex-protocol", "codex-state", diff --git a/codex-rs/analytics/src/analytics_client_tests.rs b/codex-rs/analytics/src/analytics_client_tests.rs index 38d9539615..9f8df04b40 100644 --- a/codex-rs/analytics/src/analytics_client_tests.rs +++ b/codex-rs/analytics/src/analytics_client_tests.rs @@ -168,6 +168,17 @@ use std::sync::Arc; use std::sync::Mutex; use tokio::sync::mpsc; +const TEST_PRODUCT_CLIENT_ID: &str = "codex_work_desktop"; + +fn test_tracking_context(thread_id: &str, turn_id: &str) -> TrackEventsContext { + TrackEventsContext { + model_slug: "gpt-5".to_string(), + thread_id: thread_id.to_string(), + turn_id: turn_id.to_string(), + product_client_id: TEST_PRODUCT_CLIENT_ID.to_string(), + } +} + fn sample_thread_with_metadata( thread_id: &str, ephemeral: bool, @@ -1024,11 +1035,7 @@ fn normalize_path_for_skill_id_repo_root_not_in_skill_path_uses_absolute_path() #[test] fn app_mentioned_event_serializes_expected_shape() { - let tracking = TrackEventsContext { - model_slug: "gpt-5".to_string(), - thread_id: "thread-1".to_string(), - turn_id: "turn-1".to_string(), - }; + let tracking = test_tracking_context("thread-1", "turn-1"); let event = TrackEventRequest::AppMentioned(CodexAppMentionedEventRequest { event_type: "codex_app_mentioned", event_params: codex_app_metadata( @@ -1052,7 +1059,7 @@ fn app_mentioned_event_serializes_expected_shape() { "thread_id": "thread-1", "turn_id": "turn-1", "app_name": "Calendar", - "product_client_id": originator().value, + "product_client_id": TEST_PRODUCT_CLIENT_ID, "invoke_type": "explicit", "model_slug": "gpt-5" } @@ -1062,11 +1069,7 @@ fn app_mentioned_event_serializes_expected_shape() { #[test] fn app_used_event_serializes_expected_shape() { - let tracking = TrackEventsContext { - model_slug: "gpt-5".to_string(), - thread_id: "thread-2".to_string(), - turn_id: "turn-2".to_string(), - }; + let tracking = test_tracking_context("thread-2", "turn-2"); let event = TrackEventRequest::AppUsed(CodexAppUsedEventRequest { event_type: "codex_app_used", event_params: codex_app_metadata( @@ -1090,7 +1093,7 @@ fn app_used_event_serializes_expected_shape() { "thread_id": "thread-2", "turn_id": "turn-2", "app_name": "Google Drive", - "product_client_id": originator().value, + "product_client_id": TEST_PRODUCT_CLIENT_ID, "invoke_type": "implicit", "model_slug": "gpt-5" } @@ -1381,16 +1384,8 @@ fn app_used_dedupe_is_keyed_by_turn_and_connector() { invocation_type: Some(InvocationType::Implicit), }; - let turn_1 = TrackEventsContext { - model_slug: "gpt-5".to_string(), - thread_id: "thread-1".to_string(), - turn_id: "turn-1".to_string(), - }; - let turn_2 = TrackEventsContext { - model_slug: "gpt-5".to_string(), - thread_id: "thread-1".to_string(), - turn_id: "turn-2".to_string(), - }; + let turn_1 = test_tracking_context("thread-1", "turn-1"); + let turn_2 = test_tracking_context("thread-1", "turn-2"); assert_eq!(queue.should_enqueue_app_used(&turn_1, &app), true); assert_eq!(queue.should_enqueue_app_used(&turn_1, &app), false); @@ -3009,11 +3004,7 @@ async fn subagent_tool_items_inherit_parent_connection_metadata() { #[test] fn plugin_used_event_serializes_expected_shape() { - let tracking = TrackEventsContext { - model_slug: "gpt-5".to_string(), - thread_id: "thread-3".to_string(), - turn_id: "turn-3".to_string(), - }; + let tracking = test_tracking_context("thread-3", "turn-3"); let event = TrackEventRequest::PluginUsed(CodexPluginUsedEventRequest { event_type: "codex_plugin_used", event_params: codex_plugin_used_metadata(&tracking, sample_plugin_metadata()), @@ -3033,7 +3024,7 @@ fn plugin_used_event_serializes_expected_shape() { "has_skills": true, "mcp_server_count": 2, "connector_ids": ["calendar", "drive"], - "product_client_id": originator().value, + "product_client_id": TEST_PRODUCT_CLIENT_ID, "mcp_server_names": ["mcp-1", "mcp-2"], "thread_id": "thread-3", "turn_id": "turn-3", @@ -3132,11 +3123,7 @@ fn plugin_management_event_keeps_plugin_id_local_when_remote_id_exists() { #[test] fn hook_run_event_serializes_expected_shape() { - let tracking = TrackEventsContext { - model_slug: "gpt-5".to_string(), - thread_id: "thread-3".to_string(), - turn_id: "turn-3".to_string(), - }; + let tracking = test_tracking_context("thread-3", "turn-3"); let event = TrackEventRequest::HookRun(CodexHookRunEventRequest { event_type: "codex_hook_run", event_params: codex_hook_run_metadata( @@ -3158,6 +3145,7 @@ fn hook_run_event_serializes_expected_shape() { "event_params": { "thread_id": "thread-3", "turn_id": "turn-3", + "product_client_id": TEST_PRODUCT_CLIENT_ID, "model_slug": "gpt-5", "hook_name": "PreToolUse", "hook_source": "user", @@ -3169,11 +3157,7 @@ fn hook_run_event_serializes_expected_shape() { #[test] fn hook_run_metadata_maps_sources_and_statuses() { - let tracking = TrackEventsContext { - model_slug: "gpt-5".to_string(), - thread_id: "thread-1".to_string(), - turn_id: "turn-1".to_string(), - }; + let tracking = test_tracking_context("thread-1", "turn-1"); let system = serde_json::to_value(codex_hook_run_metadata( &tracking, @@ -3224,11 +3208,7 @@ fn hook_run_metadata_maps_sources_and_statuses() { #[test] fn hook_run_metadata_maps_stopped_status() { - let tracking = TrackEventsContext { - model_slug: "gpt-5".to_string(), - thread_id: "thread-1".to_string(), - turn_id: "turn-1".to_string(), - }; + let tracking = test_tracking_context("thread-1", "turn-1"); let stopped = serde_json::to_value(codex_hook_run_metadata( &tracking, @@ -3254,16 +3234,8 @@ fn plugin_used_dedupe_is_keyed_by_turn_and_plugin() { }; let plugin = sample_plugin_metadata(); - let turn_1 = TrackEventsContext { - model_slug: "gpt-5".to_string(), - thread_id: "thread-1".to_string(), - turn_id: "turn-1".to_string(), - }; - let turn_2 = TrackEventsContext { - model_slug: "gpt-5".to_string(), - thread_id: "thread-1".to_string(), - turn_id: "turn-2".to_string(), - }; + let turn_1 = test_tracking_context("thread-1", "turn-1"); + let turn_2 = test_tracking_context("thread-1", "turn-2"); assert_eq!(queue.should_enqueue_plugin_used(&turn_1, &plugin), true); assert_eq!(queue.should_enqueue_plugin_used(&turn_1, &plugin), false); @@ -3274,11 +3246,7 @@ fn plugin_used_dedupe_is_keyed_by_turn_and_plugin() { async fn reducer_ingests_skill_invoked_fact() { let mut reducer = AnalyticsReducer::default(); let mut events = Vec::new(); - let tracking = TrackEventsContext { - model_slug: "gpt-5".to_string(), - thread_id: "thread-1".to_string(), - turn_id: "turn-1".to_string(), - }; + let tracking = test_tracking_context("thread-1", "turn-1"); let skill_path = PathBuf::from("/Users/abc/.codex/skills/doc/SKILL.md"); let expected_skill_id = skill_id_for_local_skill( /*repo_url*/ None, @@ -3311,7 +3279,7 @@ async fn reducer_ingests_skill_invoked_fact() { "skill_id": expected_skill_id, "skill_name": "doc", "event_params": { - "product_client_id": originator().value, + "product_client_id": TEST_PRODUCT_CLIENT_ID, "skill_scope": "user", "plugin_id": null, "repo_url": null, @@ -3328,11 +3296,7 @@ async fn reducer_ingests_skill_invoked_fact() { async fn reducer_includes_plugin_id_for_plugin_skill_invocations() { let mut reducer = AnalyticsReducer::default(); let mut events = Vec::new(); - let tracking = TrackEventsContext { - model_slug: "gpt-5".to_string(), - thread_id: "thread-1".to_string(), - turn_id: "turn-1".to_string(), - }; + let tracking = test_tracking_context("thread-1", "turn-1"); let skill_path = PathBuf::from("/Users/abc/.codex/plugins/cache/test/sample/skills/doc/SKILL.md"); @@ -3367,11 +3331,7 @@ async fn reducer_ingests_hook_run_fact() { reducer .ingest( AnalyticsFact::Custom(CustomAnalyticsFact::HookRun(HookRunInput { - tracking: TrackEventsContext { - model_slug: "gpt-5".to_string(), - thread_id: "thread-1".to_string(), - turn_id: "turn-1".to_string(), - }, + tracking: test_tracking_context("thread-1", "turn-1"), hook: HookRunFact { event_name: HookEventName::PostToolUse, hook_source: HookSource::Unknown, @@ -3394,11 +3354,7 @@ async fn reducer_ingests_hook_run_fact() { async fn reducer_ingests_app_and_plugin_facts() { let mut reducer = AnalyticsReducer::default(); let mut events = Vec::new(); - let tracking = TrackEventsContext { - model_slug: "gpt-5".to_string(), - thread_id: "thread-1".to_string(), - turn_id: "turn-1".to_string(), - }; + let tracking = test_tracking_context("thread-1", "turn-1"); reducer .ingest( @@ -3441,6 +3397,18 @@ async fn reducer_ingests_app_and_plugin_facts() { assert_eq!(payload[0]["event_type"], "codex_app_mentioned"); assert_eq!(payload[1]["event_type"], "codex_app_used"); assert_eq!(payload[2]["event_type"], "codex_plugin_used"); + assert_eq!( + payload[0]["event_params"]["product_client_id"], + TEST_PRODUCT_CLIENT_ID + ); + assert_eq!( + payload[1]["event_params"]["product_client_id"], + TEST_PRODUCT_CLIENT_ID + ); + assert_eq!( + payload[2]["event_params"]["product_client_id"], + TEST_PRODUCT_CLIENT_ID + ); } #[tokio::test] diff --git a/codex-rs/analytics/src/events.rs b/codex-rs/analytics/src/events.rs index 2fe465c9bd..1365e54062 100644 --- a/codex-rs/analytics/src/events.rs +++ b/codex-rs/analytics/src/events.rs @@ -773,6 +773,7 @@ pub(crate) struct CodexAppUsedEventRequest { pub(crate) struct CodexHookRunMetadata { pub(crate) thread_id: Option, pub(crate) turn_id: Option, + pub(crate) product_client_id: Option, pub(crate) model_slug: Option, pub(crate) hook_name: Option, pub(crate) hook_source: Option<&'static str>, @@ -1030,13 +1031,20 @@ pub(crate) fn codex_app_metadata( thread_id: Some(tracking.thread_id.clone()), turn_id: Some(tracking.turn_id.clone()), app_name: app.app_name, - product_client_id: Some(originator().value), + product_client_id: Some(tracking.product_client_id.clone()), invoke_type: app.invocation_type, model_slug: Some(tracking.model_slug.clone()), } } pub(crate) fn codex_plugin_metadata(plugin: PluginTelemetryMetadata) -> CodexPluginMetadata { + codex_plugin_metadata_with_product_client_id(plugin, originator().value) +} + +fn codex_plugin_metadata_with_product_client_id( + plugin: PluginTelemetryMetadata, + product_client_id: String, +) -> CodexPluginMetadata { let PluginTelemetryMetadata { plugin_id, remote_plugin_id, @@ -1062,7 +1070,7 @@ pub(crate) fn codex_plugin_metadata(plugin: PluginTelemetryMetadata) -> CodexPlu .map(|connector_id| connector_id.0) .collect() }), - product_client_id: Some(originator().value), + product_client_id: Some(product_client_id), } } @@ -1139,7 +1147,10 @@ pub(crate) fn codex_plugin_used_metadata( .as_ref() .map(|summary| summary.mcp_server_names.clone()); CodexPluginUsedMetadata { - plugin: codex_plugin_metadata(plugin), + plugin: codex_plugin_metadata_with_product_client_id( + plugin, + tracking.product_client_id.clone(), + ), mcp_server_names, thread_id: Some(tracking.thread_id.clone()), turn_id: Some(tracking.turn_id.clone()), @@ -1154,6 +1165,7 @@ pub(crate) fn codex_hook_run_metadata( CodexHookRunMetadata { thread_id: Some(tracking.thread_id.clone()), turn_id: Some(tracking.turn_id.clone()), + product_client_id: Some(tracking.product_client_id.clone()), model_slug: Some(tracking.model_slug.clone()), hook_name: Some(analytics_hook_event_name(hook.event_name).to_owned()), hook_source: Some(analytics_hook_source(hook.hook_source)), diff --git a/codex-rs/analytics/src/facts.rs b/codex-rs/analytics/src/facts.rs index 026c2e8760..4505c0c697 100644 --- a/codex-rs/analytics/src/facts.rs +++ b/codex-rs/analytics/src/facts.rs @@ -41,17 +41,20 @@ pub struct TrackEventsContext { pub model_slug: String, pub thread_id: String, pub turn_id: String, + pub product_client_id: String, } pub fn build_track_events_context( model_slug: String, thread_id: String, turn_id: String, + product_client_id: String, ) -> TrackEventsContext { TrackEventsContext { model_slug, thread_id, turn_id, + product_client_id, } } diff --git a/codex-rs/analytics/src/reducer.rs b/codex-rs/analytics/src/reducer.rs index b647984dff..666e756d3c 100644 --- a/codex-rs/analytics/src/reducer.rs +++ b/codex-rs/analytics/src/reducer.rs @@ -729,7 +729,7 @@ impl AnalyticsReducer { turn_id: Some(tracking.turn_id.clone()), invoke_type: Some(invocation.invocation_type), model_slug: Some(tracking.model_slug.clone()), - product_client_id: Some(originator().value), + product_client_id: Some(tracking.product_client_id.clone()), repo_url, skill_scope: Some(skill_scope.to_string()), plugin_id: invocation.plugin_id, diff --git a/codex-rs/app-server/src/request_processors/external_agent_session_import.rs b/codex-rs/app-server/src/request_processors/external_agent_session_import.rs index dfa2ea94fc..77b9bfcbf3 100644 --- a/codex-rs/app-server/src/request_processors/external_agent_session_import.rs +++ b/codex-rs/app-server/src/request_processors/external_agent_session_import.rs @@ -207,6 +207,7 @@ impl ExternalAgentSessionImporter { parent_thread_id: None, source: source.clone(), thread_source: None, + originator: codex_login::default_client::originator().value, base_instructions: BaseInstructions { text: config .base_instructions diff --git a/codex-rs/app-server/src/request_processors/thread_processor_tests.rs b/codex-rs/app-server/src/request_processors/thread_processor_tests.rs index 1f5ae12f44..3f7cbc9d04 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor_tests.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor_tests.rs @@ -781,6 +781,7 @@ mod thread_processor_behavior_tests { forked_from_thread_id: None, parent_thread_id: None, thread_source: None, + originator: "test_originator".to_string(), }; assert_eq!( diff --git a/codex-rs/app-server/tests/suite/conversation_summary.rs b/codex-rs/app-server/tests/suite/conversation_summary.rs index d4d0cf0615..7f964d3147 100644 --- a/codex-rs/app-server/tests/suite/conversation_summary.rs +++ b/codex-rs/app-server/tests/suite/conversation_summary.rs @@ -128,6 +128,7 @@ async fn get_conversation_summary_by_thread_id_reads_pathless_store_thread() -> parent_thread_id: None, source: SessionSource::Cli, thread_source: None, + originator: "test_originator".to_string(), base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), multi_agent_version: None, diff --git a/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs b/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs index 1a6d1eb4f6..717fe228f4 100644 --- a/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs +++ b/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs @@ -154,6 +154,7 @@ async fn thread_delete_with_non_local_thread_store_does_not_create_local_persist parent_thread_id: None, source: SessionSource::Cli, thread_source: None, + originator: "test_originator".to_string(), base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), multi_agent_version: None, diff --git a/codex-rs/app-server/tests/suite/v2/thread_read.rs b/codex-rs/app-server/tests/suite/v2/thread_read.rs index 0b4f101c09..3e39070a73 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_read.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_read.rs @@ -1364,6 +1364,7 @@ async fn seed_pathless_store_thread( parent_thread_id: None, source: ProtocolSessionSource::Cli, thread_source: None, + originator: "test_originator".to_string(), base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), multi_agent_version: None, diff --git a/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs b/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs index db8b60d162..708d4bdfb7 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs @@ -215,6 +215,7 @@ async fn thread_unarchive_preserves_pathless_store_metadata() -> Result<()> { parent_thread_id: None, source: SessionSource::Cli, thread_source: None, + originator: "test_originator".to_string(), base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), multi_agent_version: None, diff --git a/codex-rs/codex-mcp/src/codex_apps.rs b/codex-rs/codex-mcp/src/codex_apps.rs index 437e750de8..c1b5fd5fa5 100644 --- a/codex-rs/codex-mcp/src/codex_apps.rs +++ b/codex-rs/codex-mcp/src/codex_apps.rs @@ -19,6 +19,7 @@ use serde::Deserialize; use serde::Serialize; use sha1::Digest; use sha1::Sha1; +use tracing::instrument; #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct CodexAppsToolsCacheKey { @@ -170,6 +171,7 @@ pub(crate) fn read_cached_codex_apps_tools( } } +#[instrument(level = "trace", skip_all)] pub(crate) fn load_cached_codex_apps_tools( cache_context: &CodexAppsToolsCacheContext, ) -> CachedCodexAppsToolsLoad { @@ -210,6 +212,7 @@ pub(crate) fn write_cached_codex_apps_tools( let _ = std::fs::write(cache_path, bytes); } +#[instrument(level = "trace", skip_all)] pub(crate) fn load_cached_codex_apps_server_info( cache_context: &CodexAppsToolsCacheContext, ) -> Option { diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index f9cc6a5e81..1e8f5a331d 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -67,6 +67,8 @@ use rmcp::model::JsonObject; use rmcp::model::ProtocolVersion; use rmcp::model::Tool as RmcpTool; use tokio_util::sync::CancellationToken; +use tracing::Instrument; +use tracing::instrument; use tracing::warn; /// MCP server capability indicating that Codex should include [`SandboxState`] @@ -141,6 +143,7 @@ pub(crate) struct AsyncManagedClient { impl AsyncManagedClient { // Keep this constructor flat so the startup inputs remain readable at the // single call site instead of introducing a one-off params wrapper. + #[instrument(level = "trace", skip_all, fields(server_name = %server_name))] #[allow(clippy::too_many_arguments)] pub(crate) fn new( server_name: String, @@ -228,7 +231,7 @@ impl AsyncManagedClient { startup_complete_for_fut.store(true, Ordering::Release); outcome }; - let client = fut.boxed().shared(); + let client = fut.in_current_span().boxed().shared(); if cached_tool_info_snapshot.is_some() { let startup_task = client.clone(); tokio::spawn(async move { @@ -306,6 +309,7 @@ impl From for StartupOutcomeError { } } +#[instrument(level = "trace", skip_all, fields(server_name = %server_name))] pub(crate) async fn list_tools_for_client_uncached( server_name: &str, is_codex_apps_mcp_server: bool, @@ -529,6 +533,7 @@ fn validate_mcp_server_name(server_name: &str) -> Result<()> { Ok(()) } +#[instrument(level = "trace", skip_all, fields(server_name = %server_name))] async fn start_server_task( server_name: String, client: Arc, @@ -658,6 +663,7 @@ struct StartServerTaskParams { supports_openai_form_elicitation: bool, } +#[instrument(level = "trace", skip_all, fields(server_name = %server_name))] async fn make_rmcp_client( server_name: &str, server: EffectiveMcpServer, diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index e37e9b04db..225fd4ba9c 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -9,7 +9,9 @@ use crate::config::ConfigBuilder; use crate::context::ContextualUserFragment; use crate::context::SubagentNotification; use crate::init_state_db; +use crate::thread_manager::StartThreadOptions; use assert_matches::assert_matches; +use codex_extension_api::ExtensionDataInit; use codex_features::Feature; use codex_login::CodexAuth; use codex_protocol::AgentPath; @@ -21,7 +23,9 @@ use codex_protocol::models::ResponseItem; use codex_protocol::protocol::CompactedItem; use codex_protocol::protocol::ErrorEvent; use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::InitialHistory; use codex_protocol::protocol::InterAgentCommunication; +use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::TurnAbortReason; @@ -139,6 +143,33 @@ impl AgentControlHarness { } } +async fn persisted_originator(thread: &CodexThread) -> String { + thread.ensure_rollout_materialized().await; + thread + .flush_rollout() + .await + .expect("thread rollout should flush"); + let stored_thread = thread + .read_thread( + /*include_archived*/ true, /*include_history*/ true, + ) + .await + .expect("thread should be readable"); + let history = stored_thread.history.expect("history should be loaded"); + history + .items + .iter() + .find_map(|item| match item { + RolloutItem::SessionMeta(meta_line) => Some(meta_line.meta.originator.clone()), + RolloutItem::ResponseItem(_) + | RolloutItem::InterAgentCommunication(_) + | RolloutItem::EventMsg(_) + | RolloutItem::Compacted(_) + | RolloutItem::TurnContext(_) => None, + }) + .expect("session metadata should be persisted") +} + fn has_subagent_notification(history_items: &[ResponseItem]) -> bool { history_items.iter().any(|item| { let ResponseItem::Message { role, content, .. } = item else { @@ -2199,6 +2230,83 @@ async fn spawn_thread_subagent_gets_random_nickname_in_session_source() { assert_eq!(agent_role, Some("explorer".to_string())); } +#[tokio::test] +async fn spawn_thread_subagents_persist_parent_originator_across_new_and_truncated_fork() { + let harness = AgentControlHarness::new().await; + let parent = harness + .manager + .start_thread_with_options(StartThreadOptions { + config: harness.config.clone(), + initial_history: InitialHistory::New, + session_source: None, + thread_source: None, + dynamic_tools: Vec::new(), + metrics_service_name: Some("codex_work_desktop".to_string()), + multi_agent_mode: None, + parent_trace: None, + environments: Vec::new(), + thread_extension_init: ExtensionDataInit::default(), + supports_openai_form_elicitation: false, + }) + .await + .expect("parent thread should start"); + let parent_originator = persisted_originator(&parent.thread).await; + assert_eq!(parent_originator, "codex_work_desktop"); + + let child_thread_id = harness + .control + .spawn_agent( + harness.config.clone(), + text_input("hello child"), + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id: parent.thread_id, + depth: 1, + agent_path: None, + agent_nickname: None, + agent_role: Some("explorer".to_string()), + })), + ) + .await + .expect("child spawn should succeed"); + + let child_thread = harness + .manager + .get_thread(child_thread_id) + .await + .expect("child thread should be registered"); + let child_originator = persisted_originator(&child_thread).await; + assert_eq!(child_originator, parent_originator); + + let child = harness + .control + .spawn_agent_with_metadata( + harness.config.clone(), + text_input("hello forked child"), + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id: parent.thread_id, + depth: 1, + agent_path: None, + agent_nickname: None, + agent_role: Some("explorer".to_string()), + })), + SpawnAgentOptions { + fork_parent_spawn_call_id: Some("spawn-call-last-n".to_string()), + fork_mode: Some(SpawnAgentForkMode::LastNTurns(1)), + ..Default::default() + }, + ) + .await + .expect("forked child spawn should succeed"); + + let child_thread = harness + .manager + .get_thread(child.thread_id) + .await + .expect("child thread should be registered"); + let child_originator = persisted_originator(&child_thread).await; + assert_eq!(child_originator, parent_originator); +} + #[tokio::test] async fn spawn_thread_subagent_uses_role_specific_nickname_candidates() { let mut harness = AgentControlHarness::new().await; diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 9e240e968d..67a940dc9e 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -186,6 +186,7 @@ struct ModelClientState { provider: SharedModelProvider, auth_env_telemetry: AuthEnvTelemetry, session_source: SessionSource, + originator: String, model_verbosity: Option, enable_request_compression: bool, include_timing_metrics: bool, @@ -387,6 +388,7 @@ impl ModelClient { thread_id: ThreadId, provider_info: ModelProviderInfo, session_source: SessionSource, + originator: String, model_verbosity: Option, enable_request_compression: bool, include_timing_metrics: bool, @@ -408,6 +410,7 @@ impl ModelClient { provider: model_provider, auth_env_telemetry, session_source, + originator, model_verbosity, enable_request_compression, include_timing_metrics, @@ -565,6 +568,7 @@ impl ModelClient { self.state.beta_features_header.as_deref(), turn_state.as_ref(), )); + add_originator_header(&mut extra_headers, self.state.originator.as_str()); extra_headers.extend(self.build_responses_compatibility_headers(responses_metadata)); extra_headers.extend(build_session_headers( Some(responses_metadata.session_id.to_string()), @@ -676,6 +680,7 @@ impl ModelClient { fn build_subagent_headers(&self) -> ApiHeaderMap { let mut extra_headers = ApiHeaderMap::new(); + add_originator_header(&mut extra_headers, self.state.originator.as_str()); if let Some(subagent) = subagent_header_value(&self.state.session_source) && let Ok(val) = HeaderValue::from_str(&subagent) { @@ -997,6 +1002,7 @@ impl ModelClient { self.state.beta_features_header.as_deref(), /*turn_state*/ None, ); + add_originator_header(&mut headers, self.state.originator.as_str()); if let Ok(header_value) = HeaderValue::from_str(&responses_metadata.thread_id) { headers.insert("x-client-request-id", header_value); } @@ -1064,6 +1070,7 @@ impl ModelClientSession { self.client.state.beta_features_header.as_deref(), Some(&self.turn_state), ); + add_originator_header(&mut headers, self.client.state.originator.as_str()); headers.extend( self.client .build_responses_compatibility_headers(responses_metadata), @@ -1791,6 +1798,22 @@ 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/client_tests.rs b/codex-rs/core/src/client_tests.rs index 9df7d4d01d..8d58fec9bd 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -75,6 +75,7 @@ fn test_model_client(session_source: SessionSource) -> ModelClient { thread_id, provider, session_source, + "test_originator".to_string(), /*model_verbosity*/ None, /*enable_request_compression*/ false, /*include_timing_metrics*/ false, @@ -298,6 +299,10 @@ fn build_subagent_headers_sets_internal_memory_consolidation_label() { .get(X_OPENAI_SUBAGENT_HEADER) .and_then(|value| value.to_str().ok()); assert_eq!(value, Some("memory_consolidation")); + assert_eq!( + headers.get("originator"), + Some(&http::HeaderValue::from_static("test_originator")) + ); } #[test] @@ -606,6 +611,7 @@ fn model_client_with_counting_attestation( ThreadId::new(), provider, SessionSource::Exec, + "test_originator".to_string(), /*model_verbosity*/ None, /*enable_request_compression*/ false, /*include_timing_metrics*/ false, diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 9f5b867f75..640e1f8dfa 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -103,6 +103,7 @@ pub(crate) async fn run_codex_thread_interactive( forked_from_thread_id, parent_thread_id: Some(parent_session.thread_id), thread_source: Some(ThreadSource::Subagent), + originator: parent_ctx.originator.clone(), agent_control: parent_session.services.agent_control.clone(), dynamic_tools: Vec::new(), metrics_service_name: None, diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index a522cd62aa..dccfd4a763 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -77,6 +77,7 @@ pub struct ThreadConfigSnapshot { pub forked_from_thread_id: Option, pub parent_thread_id: Option, pub thread_source: Option, + pub originator: String, } /// Explains why `CodexThread::try_start_turn_if_idle` rejected an automatic diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index 1bed53fcfd..fa86dded55 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -57,6 +57,12 @@ impl fmt::Debug for StartingTurnEnvironment { } } +impl StartingTurnEnvironment { + pub(crate) async fn wait_until_ready(&self) -> Result<(), Arc> { + self.resolution.clone().await.map(|_| ()) + } +} + pub(crate) struct ThreadEnvironments { environment_manager: Arc, local_shell: Shell, diff --git a/codex-rs/core/src/hook_runtime.rs b/codex-rs/core/src/hook_runtime.rs index 12627e2d52..342e4d4b8f 100644 --- a/codex-rs/core/src/hook_runtime.rs +++ b/codex-rs/core/src/hook_runtime.rs @@ -685,6 +685,7 @@ fn hook_run_analytics_payload( .turn_id .clone() .unwrap_or_else(|| turn_context.sub_id.clone()), + turn_context.originator.clone(), ), HookRunFact { event_name: completed.run.event_name, diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index c7187c18f6..085fd496dc 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -971,6 +971,7 @@ async fn maybe_track_codex_app_used( turn_context.model_info.slug.clone(), sess.thread_id.to_string(), turn_context.sub_id.clone(), + turn_context.originator.clone(), ); sess.services.analytics_events_client.track_app_used( tracking, diff --git a/codex-rs/core/src/realtime_conversation.rs b/codex-rs/core/src/realtime_conversation.rs index 667939f105..43ecc3e576 100644 --- a/codex-rs/core/src/realtime_conversation.rs +++ b/codex-rs/core/src/realtime_conversation.rs @@ -1,4 +1,5 @@ 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; @@ -775,6 +776,7 @@ async fn prepare_realtime_start( let session_config = build_realtime_session_config(sess, ¶ms, version, configured_voice).await?; let requested_realtime_session_id = session_config.session_id.clone(); + let originator = sess.originator().await; let extra_headers = match transport { ConversationStartTransport::Websocket => { let realtime_api_key = realtime_api_key(auth.as_ref(), &provider)?; @@ -782,6 +784,7 @@ async fn prepare_realtime_start( requested_realtime_session_id.as_deref(), Some(realtime_api_key.as_str()), version, + originator.as_str(), )? } ConversationStartTransport::Webrtc { .. } => { @@ -789,6 +792,7 @@ async fn prepare_realtime_start( requested_realtime_session_id.as_deref(), /*api_key*/ None, version, + originator.as_str(), )? } }; @@ -1171,6 +1175,7 @@ fn realtime_request_headers( realtime_session_id: Option<&str>, api_key: Option<&str>, version: RealtimeWsVersion, + originator: &str, ) -> CodexResult> { let mut headers = HeaderMap::new(); @@ -1191,6 +1196,8 @@ fn realtime_request_headers( headers.insert(AUTHORIZATION, auth_value); } + add_originator_header(&mut headers, originator); + Ok(Some(headers)) } diff --git a/codex-rs/core/src/realtime_conversation_tests.rs b/codex-rs/core/src/realtime_conversation_tests.rs index 55525774b3..51425e4c3f 100644 --- a/codex-rs/core/src/realtime_conversation_tests.rs +++ b/codex-rs/core/src/realtime_conversation_tests.rs @@ -149,10 +149,14 @@ async fn clears_active_handoff_explicitly() { #[test] fn uses_quicksilver_alpha_header_for_realtime_v1() { - let headers = - realtime_request_headers(Some("session_1"), Some("sk-test"), RealtimeWsVersion::V1) - .expect("headers") - .expect("headers"); + let headers = realtime_request_headers( + Some("session_1"), + Some("sk-test"), + RealtimeWsVersion::V1, + "codex_work_desktop", + ) + .expect("headers") + .expect("headers"); assert_eq!( headers @@ -164,10 +168,39 @@ fn uses_quicksilver_alpha_header_for_realtime_v1() { #[test] fn omits_quicksilver_alpha_header_for_realtime_v2() { - let headers = - realtime_request_headers(Some("session_1"), Some("sk-test"), RealtimeWsVersion::V2) - .expect("headers") - .expect("headers"); + let headers = realtime_request_headers( + Some("session_1"), + Some("sk-test"), + RealtimeWsVersion::V2, + "codex_work_desktop", + ) + .expect("headers") + .expect("headers"); assert!(headers.get("openai-alpha").is_none()); } + +#[test] +fn realtime_headers_include_only_non_default_originator() { + let default_originator = codex_login::default_client::originator(); + for (originator, expected_header) in [ + ("codex_work_desktop", Some("codex_work_desktop")), + (default_originator.value.as_str(), None), + ] { + let headers = realtime_request_headers( + Some("session_1"), + Some("sk-test"), + RealtimeWsVersion::V2, + originator, + ) + .expect("headers") + .expect("headers"); + + assert_eq!( + headers + .get("originator") + .and_then(|value| value.to_str().ok()), + expected_header + ); + } +} diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 0eb60ab0f2..d22a023a31 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -73,7 +73,6 @@ use codex_hooks::HooksConfig; use codex_login::AuthManager; use codex_login::CodexAuth; use codex_login::auth_env_telemetry::collect_auth_env_telemetry; -use codex_login::default_client::originator; use codex_mcp::McpConnectionManager; use codex_mcp::McpResourceClient; use codex_mcp::McpRuntimeContext; @@ -433,6 +432,7 @@ pub(crate) struct CodexSpawnArgs { pub(crate) forked_from_thread_id: Option, pub(crate) parent_thread_id: Option, pub(crate) thread_source: Option, + pub(crate) originator: String, pub(crate) agent_control: AgentControl, pub(crate) dynamic_tools: Vec, pub(crate) metrics_service_name: Option, @@ -522,6 +522,7 @@ impl Codex { forked_from_thread_id, parent_thread_id, thread_source, + originator, agent_control, dynamic_tools, metrics_service_name, @@ -654,6 +655,7 @@ impl Codex { forked_from_thread_id, parent_thread_id, thread_source, + originator, dynamic_tools, user_shell_override, }; @@ -3913,7 +3915,7 @@ pub(crate) fn emit_subagent_session_started( forked_from_thread_id: thread_config .forked_from_thread_id .map(|thread_id| thread_id.to_string()), - product_client_id: client_name.clone(), + product_client_id: thread_config.originator.clone(), client_name, client_version, model: thread_config.model, diff --git a/codex-rs/core/src/session/review.rs b/codex-rs/core/src/session/review.rs index 8986ffd052..c55afbf141 100644 --- a/codex-rs/core/src/session/review.rs +++ b/codex-rs/core/src/session/review.rs @@ -118,6 +118,7 @@ pub(super) async fn spawn_review_thread( reasoning_summary, session_source, parent_thread_id: parent_turn_context.parent_thread_id, + originator: parent_turn_context.originator.clone(), environments: parent_turn_context.environments.clone(), available_models, unified_exec_shell_mode, diff --git a/codex-rs/core/src/session/rollout_reconstruction.rs b/codex-rs/core/src/session/rollout_reconstruction.rs index b3fbeb1038..b2a81392ca 100644 --- a/codex-rs/core/src/session/rollout_reconstruction.rs +++ b/codex-rs/core/src/session/rollout_reconstruction.rs @@ -85,15 +85,26 @@ fn finalize_active_segment<'a>( *window = active_segment.window; } - // `previous_turn_settings` come from the newest surviving user turn that established them. - if previous_turn_settings.is_none() && active_segment.counts_as_user_turn { + let compact_segment_reference_context_item = active_segment.base_replacement_history.is_some() + && matches!( + active_segment.reference_context_item, + TurnReferenceContextItem::Latest(_) + ); + + // `previous_turn_settings` come from the newest surviving user turn that established them, or + // from a compact replacement history that persisted the turn context used to build it. + if previous_turn_settings.is_none() + && (active_segment.counts_as_user_turn || compact_segment_reference_context_item) + { *previous_turn_settings = active_segment.previous_turn_settings; } - // `reference_context_item` comes from the newest surviving user turn baseline, or - // from a surviving compaction that explicitly cleared that baseline. + // `reference_context_item` comes from the newest surviving user turn baseline, a compact + // replacement history that persisted its new baseline, or a surviving compaction that + // explicitly cleared the previous baseline. if matches!(reference_context_item, TurnReferenceContextItem::NeverSet) && (active_segment.counts_as_user_turn + || compact_segment_reference_context_item || matches!( active_segment.reference_context_item, TurnReferenceContextItem::Cleared diff --git a/codex-rs/core/src/session/rollout_reconstruction_tests.rs b/codex-rs/core/src/session/rollout_reconstruction_tests.rs index b875e1f49e..62d9a157a0 100644 --- a/codex-rs/core/src/session/rollout_reconstruction_tests.rs +++ b/codex-rs/core/src/session/rollout_reconstruction_tests.rs @@ -1111,6 +1111,66 @@ async fn reconstruct_history_legacy_compaction_without_replacement_history_clear assert!(reconstructed.reference_context_item.is_none()); } +#[tokio::test] +async fn record_initial_history_resumed_standalone_compact_reestablishes_reference_context_item() { + let (session, turn_context) = make_session_and_context().await; + let compact_context_item = turn_context.to_turn_context_item(); + let compact_turn_id = compact_context_item + .turn_id + .clone() + .expect("turn context should have turn_id"); + let rollout_items = vec![ + RolloutItem::EventMsg(EventMsg::TurnStarted( + codex_protocol::protocol::TurnStartedEvent { + turn_id: compact_turn_id.clone(), + trace_id: None, + started_at: None, + model_context_window: Some(128_000), + collaboration_mode_kind: ModeKind::Default, + }, + )), + RolloutItem::Compacted(CompactedItem { + message: String::new(), + replacement_history: Some(vec![assistant_message("compacted baseline")]), + window_number: Some(1), + first_window_id: None, + previous_window_id: None, + window_id: None, + }), + RolloutItem::TurnContext(compact_context_item.clone()), + RolloutItem::EventMsg(EventMsg::TurnComplete( + codex_protocol::protocol::TurnCompleteEvent { + turn_id: compact_turn_id, + last_agent_message: None, + completed_at: None, + duration_ms: None, + time_to_first_token_ms: None, + }, + )), + ]; + + session + .record_initial_history(InitialHistory::Resumed(ResumedHistory { + conversation_id: ThreadId::default(), + history: Arc::new(rollout_items), + rollout_path: Some(PathBuf::from("/tmp/resume.jsonl")), + })) + .await; + + assert_eq!( + session.previous_turn_settings().await, + Some(PreviousTurnSettings { + model: compact_context_item.model.clone(), + comp_hash: compact_context_item.comp_hash.clone(), + realtime_active: compact_context_item.realtime_active, + }) + ); + assert_eq!( + session.reference_context_item().await, + Some(compact_context_item) + ); +} + #[tokio::test] async fn record_initial_history_resumed_turn_context_after_compaction_reestablishes_reference_context_item() { diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index ee62155a87..af64e0bef5 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -106,6 +106,8 @@ pub(crate) struct SessionConfiguration { pub(super) parent_thread_id: Option, /// Optional analytics source classification for this thread. pub(super) thread_source: Option, + /// Effective originator used for this thread's Responses requests and analytics events. + pub(super) originator: String, pub(super) dynamic_tools: Vec, pub(super) user_shell_override: Option, } @@ -198,6 +200,7 @@ impl SessionConfiguration { forked_from_thread_id: self.forked_from_thread_id, parent_thread_id: self.parent_thread_id, thread_source: self.thread_source.clone(), + originator: self.originator.clone(), } } @@ -473,6 +476,11 @@ impl Session { self.services.agent_control.session_id() } + pub(crate) async fn originator(&self) -> String { + let state = self.state.lock().await; + state.session_configuration.originator.clone() + } + #[instrument(name = "session_init", level = "info", skip_all)] #[allow(clippy::too_many_arguments)] pub(crate) async fn new( @@ -581,6 +589,7 @@ impl Session { parent_thread_id, source: session_source, thread_source: session_configuration.thread_source.clone(), + originator: session_configuration.originator.clone(), base_instructions: BaseInstructions { text: session_configuration.base_instructions.clone(), }, @@ -756,7 +765,7 @@ impl Session { let auth_mode = auth.map(CodexAuth::auth_mode).map(TelemetryAuthMode::from); let account_id = auth.and_then(CodexAuth::get_account_id); let account_email = auth.and_then(CodexAuth::get_account_email); - let originator = originator().value; + let originator = session_configuration.originator.clone(); let terminal_type = user_agent(); let session_model = session_configuration.collaboration_mode.model().to_string(); let auth_env_telemetry = collect_auth_env_telemetry( @@ -1057,6 +1066,7 @@ impl Session { thread_id, session_configuration.provider.clone(), session_configuration.session_source.clone(), + session_configuration.originator.clone(), config.model_verbosity, config.features.enabled(Feature::EnableRequestCompression), config.features.enabled(Feature::RuntimeMetrics), diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 8c6bc8301d..2da1b9a814 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -479,6 +479,7 @@ fn test_model_client_session() -> crate::client::ModelClientSession { thread_id, ModelProviderInfo::create_openai_provider(/* base_url */ /*base_url*/ None), codex_protocol::protocol::SessionSource::Exec, + "test_originator".to_string(), /*model_verbosity*/ None, /*enable_request_compression*/ false, /*include_timing_metrics*/ false, @@ -3554,6 +3555,7 @@ async fn set_rate_limits_retains_previous_credits() { forked_from_thread_id: None, parent_thread_id: None, thread_source: None, + originator: "test_originator".to_string(), dynamic_tools: Vec::new(), user_shell_override: None, }; @@ -3661,6 +3663,7 @@ async fn set_rate_limits_updates_plan_type_when_present() { forked_from_thread_id: None, parent_thread_id: None, thread_source: None, + originator: "test_originator".to_string(), dynamic_tools: Vec::new(), user_shell_override: None, }; @@ -3910,6 +3913,7 @@ async fn attach_thread_persistence(session: &mut Session) -> PathBuf { parent_thread_id: None, source: SessionSource::Exec, thread_source: None, + originator: "test_originator".to_string(), base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), multi_agent_version: None, @@ -4191,13 +4195,14 @@ pub(crate) async fn make_session_configuration_for_tests() -> SessionConfigurati forked_from_thread_id: None, parent_thread_id: None, thread_source: None, + originator: "test_originator".to_string(), dynamic_tools: Vec::new(), user_shell_override: None, } } #[tokio::test] -async fn emit_subagent_session_started_includes_fork_lineage_from_session_configuration() { +async fn emit_subagent_session_started_includes_fork_lineage_and_originator() { use wiremock::Mock; use wiremock::MockServer; use wiremock::ResponseTemplate; @@ -4273,6 +4278,10 @@ async fn emit_subagent_session_started_includes_fork_lineage_from_session_config event["event_params"]["forked_from_thread_id"], forked_from_thread_id.to_string() ); + assert_eq!( + event["event_params"]["app_server_client"]["product_client_id"], + "test_originator" + ); } async fn resolved_environments_for_configuration( @@ -5058,6 +5067,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() { forked_from_thread_id: None, parent_thread_id: None, thread_source: None, + originator: "test_originator".to_string(), dynamic_tools: Vec::new(), user_shell_override: None, }; @@ -5187,6 +5197,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { forked_from_thread_id: None, parent_thread_id: None, thread_source: None, + originator: "test_originator".to_string(), dynamic_tools: Vec::new(), user_shell_override: None, }; @@ -5289,6 +5300,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { thread_id, session_configuration.provider.clone(), session_configuration.session_source.clone(), + session_configuration.originator.clone(), config.model_verbosity, config.features.enabled(Feature::EnableRequestCompression), config.features.enabled(Feature::RuntimeMetrics), @@ -5433,6 +5445,7 @@ async fn make_session_with_config_and_rx( forked_from_thread_id: None, parent_thread_id: None, thread_source: None, + originator: "test_originator".to_string(), dynamic_tools: Vec::new(), user_shell_override: None, }; @@ -5540,6 +5553,7 @@ async fn make_session_with_history_source_and_agent_control_and_rx( forked_from_thread_id: None, parent_thread_id: None, thread_source: None, + originator: "test_originator".to_string(), dynamic_tools: Vec::new(), user_shell_override: None, }; @@ -6782,6 +6796,7 @@ async fn shutdown_complete_does_not_append_to_thread_store_after_shutdown() { parent_thread_id: None, source: SessionSource::Exec, thread_source: None, + originator: "test_originator".to_string(), base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), multi_agent_version: None, @@ -7263,6 +7278,7 @@ where forked_from_thread_id: None, parent_thread_id: None, thread_source: None, + originator: "test_originator".to_string(), dynamic_tools, user_shell_override: None, }; @@ -7364,6 +7380,7 @@ where thread_id, session_configuration.provider.clone(), session_configuration.session_source.clone(), + session_configuration.originator.clone(), config.model_verbosity, config.features.enabled(Feature::EnableRequestCompression), config.features.enabled(Feature::RuntimeMetrics), diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index e61578929f..6f23cb0960 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -733,6 +733,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { forked_from_thread_id: None, parent_thread_id: None, thread_source: None, + originator: "test_originator".to_string(), agent_control: AgentControl::default(), dynamic_tools: Vec::new(), metrics_service_name: None, diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 397dd17b65..21cdec2269 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -526,6 +526,7 @@ async fn build_skills_and_plugins( turn_context.model_info.slug.clone(), sess.thread_id.to_string(), turn_context.sub_id.clone(), + turn_context.originator.clone(), ); let loaded_plugins = sess .services @@ -848,10 +849,15 @@ async fn maybe_run_previous_model_inline_compact( .with_model(previous_turn_settings.model, &sess.services.models_manager) .await, ); + let compaction_turn_context = if turn_context.config.features.enabled(Feature::TokenBudget) { + Arc::clone(turn_context) + } else { + Arc::clone(&previous_model_turn_context) + }; if should_compact_for_comp_hash_change { let step_context = sess - .capture_step_context(Arc::clone(&previous_model_turn_context)) + .capture_step_context(Arc::clone(&compaction_turn_context)) .await; run_auto_compact( sess, @@ -891,7 +897,7 @@ async fn maybe_run_previous_model_inline_compact( && old_context_window > new_context_window; if should_run { let step_context = sess - .capture_step_context(Arc::clone(&previous_model_turn_context)) + .capture_step_context(Arc::clone(&compaction_turn_context)) .await; run_auto_compact( sess, diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index fdd9df5d51..da526d4911 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -113,6 +113,7 @@ pub struct TurnContext { pub(crate) reasoning_summary: ReasoningSummaryConfig, pub(crate) session_source: SessionSource, pub(crate) parent_thread_id: Option, + pub(crate) originator: String, pub(crate) environments: TurnEnvironmentSnapshot, /// The session's absolute working directory. All relative paths provided /// by the model as well as sandbox policies are resolved against this path @@ -265,6 +266,7 @@ impl TurnContext { reasoning_summary: self.reasoning_summary, session_source: self.session_source.clone(), parent_thread_id: self.parent_thread_id, + originator: self.originator.clone(), environments: self.environments.clone(), #[allow(deprecated)] cwd: self.cwd.clone(), @@ -548,6 +550,7 @@ impl Session { reasoning_summary, session_source, parent_thread_id: session_configuration.parent_thread_id, + originator: session_configuration.originator.clone(), environments, #[allow(deprecated)] cwd, diff --git a/codex-rs/core/src/skills.rs b/codex-rs/core/src/skills.rs index 2c1ba8245a..7b5d87b26f 100644 --- a/codex-rs/core/src/skills.rs +++ b/codex-rs/core/src/skills.rs @@ -102,6 +102,7 @@ pub(crate) async fn maybe_emit_implicit_skill_invocation( turn_context.model_info.slug.clone(), sess.thread_id.to_string(), turn_context.sub_id.clone(), + turn_context.originator.clone(), ), vec![invocation], ); diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 816df89c39..75a02d9996 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -29,6 +29,8 @@ use codex_extension_api::empty_extension_registry; use codex_features::Feature; use codex_login::AuthManager; use codex_login::CodexAuth; +use codex_login::default_client::CODEX_INTERNAL_ORIGINATOR_OVERRIDE_ENV_VAR; +use codex_login::default_client::originator; use codex_model_provider::create_model_provider; use codex_model_provider_info::ModelProviderInfo; use codex_model_provider_info::OPENAI_PROVIDER_ID; @@ -192,6 +194,28 @@ pub struct StartThreadOptions { pub supports_openai_form_elicitation: bool, } +fn originator_from_service_name(service_name: Option<&str>) -> Option { + let service_name = service_name?.trim(); + if service_name.eq_ignore_ascii_case("codex_work_desktop") { + return Some("codex_work_desktop".to_string()); + } + None +} + +fn effective_originator_value( + metrics_service_name: Option<&str>, + env_originator: Option, + persisted_originator: Option, + inherited_originator: Option, + default_originator: String, +) -> String { + originator_from_service_name(metrics_service_name) + .or(persisted_originator) + .or(inherited_originator) + .or(env_originator) + .unwrap_or(default_originator) +} + pub(crate) struct ResumeThreadWithHistoryOptions { pub(crate) config: Config, pub(crate) initial_history: InitialHistory, @@ -1220,6 +1244,64 @@ impl ThreadManagerState { } } + async fn inherited_originator_for_parent_thread( + &self, + session_source: &SessionSource, + parent_thread_id: Option, + forked_from_thread_id: Option, + ) -> Option { + let inherited_thread_id = match session_source { + SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, .. + }) => Some(*parent_thread_id), + _ => parent_thread_id.or(forked_from_thread_id), + }; + let thread = self.get_thread(inherited_thread_id?).await.ok()?; + let originator = thread.config_snapshot().await.originator; + (!originator.is_empty()).then_some(originator) + } + + async fn effective_originator( + &self, + initial_history: &InitialHistory, + metrics_service_name: Option<&str>, + session_source: &SessionSource, + parent_thread_id: Option, + forked_from_thread_id: Option, + ) -> String { + let persisted_originator = initial_history.get_session_originator(); + let inherited_originator = match initial_history { + InitialHistory::New | InitialHistory::Cleared => { + self.inherited_originator_for_parent_thread( + session_source, + parent_thread_id, + forked_from_thread_id, + ) + .await + } + InitialHistory::Forked(_) if persisted_originator.is_none() => { + self.inherited_originator_for_parent_thread( + session_source, + parent_thread_id, + forked_from_thread_id, + ) + .await + } + InitialHistory::Resumed(_) | InitialHistory::Forked(_) => None, + }; + + let env_originator = std::env::var(CODEX_INTERNAL_ORIGINATOR_OVERRIDE_ENV_VAR) + .is_ok() + .then(|| originator().value); + effective_originator_value( + metrics_service_name, + env_originator, + persisted_originator, + inherited_originator, + originator().value, + ) + } + /// Spawn a new thread with no history using a provided config. pub(crate) async fn spawn_new_thread( &self, @@ -1466,6 +1548,15 @@ impl ThreadManagerState { forked_from_thread_id, ) .await; + let originator = self + .effective_originator( + &initial_history, + metrics_service_name.as_deref(), + &session_source, + parent_thread_id, + forked_from_thread_id, + ) + .await; let CodexSpawnOk { codex, thread_id, .. } = Box::pin(Codex::spawn(CodexSpawnArgs { @@ -1484,6 +1575,7 @@ impl ThreadManagerState { forked_from_thread_id, parent_thread_id, thread_source, + originator, agent_control, dynamic_tools, metrics_service_name, diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index 990c5ca5d6..e6b222f6bb 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -68,6 +68,41 @@ fn developer_interrupted_marker() -> ResponseItem { .expect("developer interrupted marker should be enabled") } +#[test] +fn effective_originator_prefers_thread_scoped_sources_before_env_originator() { + for (metrics_service_name, persisted_originator, inherited_originator, expected_originator) in [ + ( + Some("codex_work_desktop"), + Some("persisted_originator"), + Some("inherited_originator"), + "codex_work_desktop", + ), + ( + None, + Some("persisted_originator"), + Some("inherited_originator"), + "persisted_originator", + ), + ( + None, + None, + Some("inherited_originator"), + "inherited_originator", + ), + ] { + assert_eq!( + effective_originator_value( + metrics_service_name, + Some("Codex Desktop".to_string()), + persisted_originator.map(str::to_string), + inherited_originator.map(str::to_string), + "codex_cli_rs".to_string(), + ), + expected_originator + ); + } +} + #[test] fn truncates_before_requested_user_message() { let items = [ diff --git a/codex-rs/core/src/tools/handlers/mod.rs b/codex-rs/core/src/tools/handlers/mod.rs index 47c81eabc6..29ad172979 100644 --- a/codex-rs/core/src/tools/handlers/mod.rs +++ b/codex-rs/core/src/tools/handlers/mod.rs @@ -35,6 +35,7 @@ pub(crate) mod tool_search_spec; pub(crate) mod unified_exec; mod view_image; pub(crate) mod view_image_spec; +mod wait_for_environment; use codex_sandboxing::policy_transforms::intersect_permission_profiles; use codex_sandboxing::policy_transforms::merge_permission_profiles; @@ -78,6 +79,7 @@ pub use unified_exec::ExecCommandHandler; pub(crate) use unified_exec::ExecCommandHandlerOptions; pub use unified_exec::WriteStdinHandler; pub use view_image::ViewImageHandler; +pub(crate) use wait_for_environment::WaitForEnvironmentHandler; pub(crate) fn parse_arguments(arguments: &str) -> Result where diff --git a/codex-rs/core/src/tools/handlers/wait_for_environment.rs b/codex-rs/core/src/tools/handlers/wait_for_environment.rs new file mode 100644 index 0000000000..c9103521c1 --- /dev/null +++ b/codex-rs/core/src/tools/handlers/wait_for_environment.rs @@ -0,0 +1,105 @@ +use std::collections::BTreeMap; + +use codex_tools::JsonSchema; +use codex_tools::JsonToolOutput; +use codex_tools::ResponsesApiTool; +use codex_tools::ToolName; +use codex_tools::ToolSpec; +use serde::Deserialize; +use serde_json::json; + +use crate::function_tool::FunctionCallError; +use crate::tools::context::ToolInvocation; +use crate::tools::context::ToolPayload; +use crate::tools::context::boxed_tool_output; +use crate::tools::handlers::parse_arguments; +use crate::tools::registry::CoreToolRuntime; +use crate::tools::registry::ToolExecutor; + +const WAIT_FOR_ENVIRONMENT_TOOL_NAME: &str = "wait_for_environment"; + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct WaitForEnvironmentArgs { + environment_id: String, +} + +pub(crate) struct WaitForEnvironmentHandler; + +impl ToolExecutor for WaitForEnvironmentHandler { + fn tool_name(&self) -> ToolName { + ToolName::plain(WAIT_FOR_ENVIRONMENT_TOOL_NAME) + } + + fn spec(&self) -> ToolSpec { + ToolSpec::Function(ResponsesApiTool { + name: WAIT_FOR_ENVIRONMENT_TOOL_NAME.to_string(), + description: "Wait for a starting environment to become available before continuing." + .to_string(), + strict: false, + defer_loading: None, + parameters: JsonSchema::object( + BTreeMap::from([( + "environment_id".to_string(), + JsonSchema::string(Some( + "The id of an environment currently marked as starting.".to_string(), + )), + )]), + /*required*/ Some(vec!["environment_id".to_string()]), + /*additional_properties*/ Some(false.into()), + ), + output_schema: None, + }) + } + + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(async move { + let ToolInvocation { + payload, + step_context, + .. + } = invocation; + let arguments = match payload { + ToolPayload::Function { arguments } => arguments, + _ => { + return Err(FunctionCallError::Fatal(format!( + "{WAIT_FOR_ENVIRONMENT_TOOL_NAME} handler received unsupported payload" + ))); + } + }; + let args: WaitForEnvironmentArgs = parse_arguments(&arguments)?; + let environment_id = args.environment_id; + let already_ready = step_context + .environments + .turn_environments + .iter() + .any(|environment| environment.environment_id == environment_id); + if !already_ready { + let Some(environment) = step_context + .environments + .starting + .iter() + .find(|environment| environment.selection.environment_id == environment_id) + .cloned() + else { + return Err(FunctionCallError::RespondToModel(format!( + "environment `{environment_id}` is neither ready nor starting" + ))); + }; + + environment.wait_until_ready().await.map_err(|_| { + FunctionCallError::RespondToModel(format!( + "Environment `{environment_id}` failed to start and is unavailable. Continue without it." + )) + })?; + } + + Ok(boxed_tool_output(JsonToolOutput::new(json!({ + "environment_id": environment_id, + "status": "ready", + })))) + }) + } +} + +impl CoreToolRuntime for WaitForEnvironmentHandler {} diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index 1ac255a839..bc4066cd2d 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -29,6 +29,7 @@ use crate::tools::handlers::SleepHandler; use crate::tools::handlers::TestSyncHandler; use crate::tools::handlers::ToolSearchHandlerCache; use crate::tools::handlers::ViewImageHandler; +use crate::tools::handlers::WaitForEnvironmentHandler; use crate::tools::handlers::WriteStdinHandler; use crate::tools::handlers::agent_jobs::ReportAgentJobResultHandler; use crate::tools::handlers::agent_jobs::SpawnAgentsOnCsvHandler; @@ -719,6 +720,10 @@ fn add_core_utility_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut planned_tools.add(PlanHandler); + if features.enabled(Feature::DeferredExecutor) { + planned_tools.add(WaitForEnvironmentHandler); + } + if turn_context.config.experimental_request_user_input_enabled { planned_tools.add_with_exposure( RequestUserInputHandler { diff --git a/codex-rs/core/tests/responses_headers.rs b/codex-rs/core/tests/responses_headers.rs index 2f323e0dd3..7e213991cb 100644 --- a/codex-rs/core/tests/responses_headers.rs +++ b/codex-rs/core/tests/responses_headers.rs @@ -123,6 +123,7 @@ async fn responses_stream_includes_subagent_header_on_review() { thread_id, provider.clone(), session_source.clone(), + "test_originator".to_string(), config.model_verbosity, /*enable_request_compression*/ false, /*include_timing_metrics*/ false, @@ -255,6 +256,7 @@ async fn responses_stream_includes_subagent_header_on_other() { thread_id, provider.clone(), session_source.clone(), + "test_originator".to_string(), config.model_verbosity, /*enable_request_compression*/ false, /*include_timing_metrics*/ false, @@ -373,6 +375,7 @@ async fn responses_respects_model_info_overrides_from_config() { thread_id, provider.clone(), session_source.clone(), + "test_originator".to_string(), config.model_verbosity, /*enable_request_compression*/ false, /*include_timing_metrics*/ false, diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index aa2026be1c..98303d09b7 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -1223,6 +1223,7 @@ async fn send_provider_auth_request(server: &MockServer, auth: ModelProviderAuth thread_id, provider, SessionSource::Exec, + "test_originator".to_string(), config.model_verbosity, /*enable_request_compression*/ false, /*include_timing_metrics*/ false, @@ -2835,6 +2836,7 @@ async fn azure_responses_request_includes_store_and_reasoning_ids() { thread_id, provider.clone(), SessionSource::Exec, + "test_originator".to_string(), config.model_verbosity, /*enable_request_compression*/ false, /*include_timing_metrics*/ false, diff --git a/codex-rs/core/tests/suite/client_websockets.rs b/codex-rs/core/tests/suite/client_websockets.rs index d35714ea45..b1b731a085 100755 --- a/codex-rs/core/tests/suite/client_websockets.rs +++ b/codex-rs/core/tests/suite/client_websockets.rs @@ -2192,6 +2192,7 @@ async fn websocket_harness_with_provider_options( thread_id, provider.clone(), SessionSource::Exec, + "test_originator".to_string(), config.model_verbosity, /*enable_request_compression*/ false, runtime_metrics_enabled, diff --git a/codex-rs/core/tests/suite/remote_env.rs b/codex-rs/core/tests/suite/remote_env.rs index d360bde25f..e665e86a0d 100644 --- a/codex-rs/core/tests/suite/remote_env.rs +++ b/codex-rs/core/tests/suite/remote_env.rs @@ -35,6 +35,7 @@ use codex_utils_path_uri::PathUri; use core_test_support::PathBufExt; use core_test_support::PathExt; use core_test_support::TestTargetOs; +use core_test_support::responses::ResponseMock; use core_test_support::responses::ev_apply_patch_custom_tool_call; use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; @@ -401,32 +402,40 @@ async fn serve_environment_info(listener: TcpListener) { .expect("environment info response"); } +fn tool_names(body: &Value) -> Vec { + body["tools"] + .as_array() + .expect("tools array") + .iter() + .filter_map(|tool| tool.get("name").and_then(Value::as_str).map(str::to_owned)) + .collect() +} + +async fn wait_for_response_request_count(response_mock: &ResponseMock, expected_count: usize) { + timeout(Duration::from_secs(5), async { + while response_mock.requests().len() < expected_count { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("timed out waiting for Responses API request"); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn deferred_executor_updates_context_and_tools_after_startup() -> Result<()> { let listener = TcpListener::bind("127.0.0.1:0").await?; let server = start_mock_server().await; - let user_input_call_id = "wait-for-startup"; + let wait_call_id = "wait-for-startup"; let response_mock = mount_sse_sequence( &server, vec![ sse(vec![ ev_response_created("resp-1"), ev_function_call( - user_input_call_id, - "request_user_input", + wait_call_id, + "wait_for_environment", &json!({ - "questions": [{ - "id": "continue", - "header": "Continue", - "question": "Continue after startup?", - "options": [{ - "label": "Yes (Recommended)", - "description": "Continue the test." - }, { - "label": "No", - "description": "Stop the test." - }] - }] + "environment_id": REMOTE_ENVIRONMENT_ID, }) .to_string(), ), @@ -469,12 +478,6 @@ async fn deferred_executor_updates_context_and_tools_after_startup() -> Result<( .enable(Feature::RequestPermissionsTool) .is_ok() ); - assert!( - config - .features - .enable(Feature::DefaultModeRequestUserInput) - .is_ok() - ); }); let test = timeout(Duration::from_secs(5), builder.build(&server)) .await @@ -492,26 +495,9 @@ async fn deferred_executor_updates_context_and_tools_after_startup() -> Result<( thread_settings: Default::default(), }) .await?; - let request = wait_for_event_match(&test.codex, |event| match event { - EventMsg::RequestUserInput(request) => Some(request.clone()), - _ => None, - }) - .await; - + wait_for_response_request_count(&response_mock, /*expected_count*/ 1).await; + assert_eq!(response_mock.requests().len(), 1); serve_environment_info(listener).await; - test.codex - .submit(Op::UserInputAnswer { - id: request.turn_id, - response: RequestUserInputResponse { - answers: HashMap::from([( - "continue".to_string(), - RequestUserInputAnswer { - answers: vec!["Yes (Recommended)".to_string()], - }, - )]), - }, - }) - .await?; let event = wait_for_event(&test.codex, |event| { matches!( event, @@ -543,16 +529,22 @@ async fn deferred_executor_updates_context_and_tools_after_startup() -> Result<( let requests = response_mock.requests(); assert_eq!(requests.len(), 3); - let tool_names = |request_index: usize| { - requests[request_index].body_json()["tools"] - .as_array() - .expect("tools array") - .iter() - .filter_map(|tool| tool.get("name").and_then(Value::as_str).map(str::to_owned)) - .collect::>() - }; - assert!(!tool_names(0).contains(&"exec_command".to_string())); - assert!(tool_names(1).contains(&"exec_command".to_string())); + let starting_tools = tool_names(&requests[0].body_json()); + let ready_tools = tool_names(&requests[1].body_json()); + assert!(starting_tools.contains(&"wait_for_environment".to_string())); + assert!(!starting_tools.contains(&"exec_command".to_string())); + assert!(ready_tools.contains(&"exec_command".to_string())); + assert!(ready_tools.contains(&"wait_for_environment".to_string())); + let (wait_output, _) = requests[1] + .function_call_output_content_and_success(wait_call_id) + .context("wait_for_environment output should be present")?; + assert_eq!( + serde_json::from_str::(&wait_output.context("wait output should contain text")?)?, + json!({ + "environment_id": REMOTE_ENVIRONMENT_ID, + "status": "ready", + }) + ); assert!( requests[0] .message_input_texts("user") @@ -586,6 +578,93 @@ async fn deferred_executor_updates_context_and_tools_after_startup() -> Result<( Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn deferred_executor_wait_reports_startup_failure() -> Result<()> { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let server = start_mock_server().await; + let wait_call_id = "wait-for-failure"; + let response_mock = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_function_call( + wait_call_id, + "wait_for_environment", + &json!({ + "environment_id": REMOTE_ENVIRONMENT_ID, + }) + .to_string(), + ), + ev_completed("resp-1"), + ]), + sse(vec![ + ev_response_created("resp-2"), + ev_assistant_message("msg-2", "done"), + ev_completed("resp-2"), + ]), + ], + ) + .await; + let mut builder = test_codex() + .with_exec_server_url(format!("ws://{}", listener.local_addr()?)) + .with_config(|config| { + config.use_experimental_unified_exec_tool = true; + assert!(config.features.enable(Feature::DeferredExecutor).is_ok()); + assert!(config.features.enable(Feature::UnifiedExec).is_ok()); + }); + let test = timeout(Duration::from_secs(5), builder.build(&server)) + .await + .context("thread startup should not wait for the remote environment")??; + + test.codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: "wait for the environment".into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: Default::default(), + }) + .await?; + wait_for_response_request_count(&response_mock, /*expected_count*/ 1).await; + assert_eq!(response_mock.requests().len(), 1); + let (stream, _) = timeout(Duration::from_secs(5), listener.accept()) + .await + .context("exec-server connection should arrive")??; + drop(stream); + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + let requests = response_mock.requests(); + assert_eq!(requests.len(), 2); + let starting_tools = tool_names(&requests[0].body_json()); + let failed_tools = tool_names(&requests[1].body_json()); + assert!(starting_tools.contains(&"wait_for_environment".to_string())); + assert!(!starting_tools.contains(&"exec_command".to_string())); + assert!(failed_tools.contains(&"wait_for_environment".to_string())); + assert!(!failed_tools.contains(&"exec_command".to_string())); + let (wait_output, _) = requests[1] + .function_call_output_content_and_success(wait_call_id) + .context("wait_for_environment output should be present")?; + assert_eq!( + wait_output.as_deref(), + Some("Environment `remote` failed to start and is unavailable. Continue without it.") + ); + assert!( + requests[1] + .message_input_texts("user") + .iter() + .any(|text| text.contains("status=\"unavailable\"")) + ); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn deferred_executor_compaction_preserves_then_updates_environment_once() -> Result<()> { let listener = TcpListener::bind("127.0.0.1:0").await?; diff --git a/codex-rs/core/tests/suite/rollout_list_find.rs b/codex-rs/core/tests/suite/rollout_list_find.rs index 1461c1af35..c8b10b24e8 100644 --- a/codex-rs/core/tests/suite/rollout_list_find.rs +++ b/codex-rs/core/tests/suite/rollout_list_find.rs @@ -188,6 +188,7 @@ async fn find_locates_rollout_file_written_by_recorder() -> std::io::Result<()> /*parent_thread_id*/ None, SessionSource::Exec, /*thread_source*/ None, + "test_originator".to_string(), BaseInstructions::default(), Vec::new(), ), diff --git a/codex-rs/core/tests/suite/token_budget.rs b/codex-rs/core/tests/suite/token_budget.rs index a4aac9d93e..63e9f52bb0 100644 --- a/codex-rs/core/tests/suite/token_budget.rs +++ b/codex-rs/core/tests/suite/token_budget.rs @@ -4,8 +4,13 @@ use codex_config::types::McpServerTransportConfig; use codex_core::config::TokenBudgetConfig; use codex_features::Feature; use codex_model_provider_info::built_in_model_providers; +use codex_models_manager::bundled_models_response; use codex_protocol::config_types::AutoCompactTokenLimitScope; use codex_protocol::items::TurnItem; +use codex_protocol::models::PermissionProfile; +use codex_protocol::openai_models::ModelInfo; +use codex_protocol::openai_models::ModelsResponse; +use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::CONTEXT_WINDOW_CLOSE_TAG; use codex_protocol::protocol::CONTEXT_WINDOW_OPEN_TAG; use codex_protocol::protocol::EventMsg; @@ -14,6 +19,7 @@ use codex_protocol::protocol::HookRunStatus; use codex_protocol::protocol::ItemCompletedEvent; use codex_protocol::protocol::ItemStartedEvent; use codex_protocol::protocol::Op; +use codex_protocol::user_input::UserInput; use core_test_support::PathBufExt; use core_test_support::assert_regex_match; use core_test_support::context_snapshot; @@ -31,8 +37,11 @@ use core_test_support::responses::sse; use core_test_support::responses::start_mock_server; use core_test_support::skip_if_no_network; use core_test_support::stdio_server_bin; +use core_test_support::test_codex::TestCodex; use core_test_support::test_codex::local; +use core_test_support::test_codex::local_selections; use core_test_support::test_codex::test_codex; +use core_test_support::test_codex::turn_permission_fields; use core_test_support::wait_for_event; use core_test_support::wait_for_event_match; use core_test_support::wait_for_mcp_server; @@ -41,10 +50,30 @@ use serde_json::Value; use serde_json::json; use std::collections::HashMap; use std::path::Path; +use std::sync::Arc; use std::time::Duration; +use tempfile::TempDir; const CONFIGURED_CONTEXT_WINDOW: i64 = 128_000; +fn model_info_with_context_window_and_comp_hash( + slug: &str, + context_window: i64, + comp_hash: &str, +) -> ModelInfo { + let models_response = bundled_models_response().expect("bundled models.json should parse"); + let mut model_info = models_response + .models + .into_iter() + .find(|model| model.slug == slug) + .expect("model missing from models.json"); + model_info.context_window = Some(context_window); + model_info.max_context_window = None; + model_info.comp_hash = Some(comp_hash.to_string()); + model_info.effective_context_window_percent = 100; + model_info +} + fn token_budget_contexts(request: &ResponsesRequest) -> Vec { let context_window_prefix = format!("{CONTEXT_WINDOW_OPEN_TAG}\nThread id: "); request @@ -137,6 +166,42 @@ fn write_token_budget_compact_hooks(home: &Path) { std::fs::write(home.join("hooks.json"), hooks.to_string()).expect("write hooks.json"); } +async fn submit_turn_with_model(test: &TestCodex, prompt: &str, model: &str) -> Result<()> { + let (sandbox_policy, permission_profile) = + turn_permission_fields(PermissionProfile::Disabled, test.config.cwd.as_path()); + test.codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: prompt.to_string(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: codex_protocol::protocol::ThreadSettingsOverrides { + environments: Some(local_selections(test.config.cwd.clone())), + approval_policy: Some(AskForApproval::Never), + sandbox_policy: Some(sandbox_policy), + permission_profile, + collaboration_mode: Some(codex_protocol::config_types::CollaborationMode { + mode: codex_protocol::config_types::ModeKind::Default, + settings: codex_protocol::config_types::Settings { + model: model.to_string(), + reasoning_effort: None, + developer_instructions: None, + }, + }), + ..Default::default() + }, + }) + .await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + Ok(()) +} + async fn assert_context_compaction_item_lifecycle(codex: &std::sync::Arc) { let mut saw_compaction_started = false; let mut saw_compaction_completed = false; @@ -211,6 +276,95 @@ async fn token_budget_context_is_only_emitted_with_full_context() -> Result<()> Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn token_budget_comp_hash_reset_uses_current_model_context() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let responses = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_completed_with_tokens("resp-1", /*total_tokens*/ 100), + ]), + sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), + ], + ) + .await; + + let codex_home = Arc::new(TempDir::new()?); + let skill_dir = codex_home.path().join("skills/budget-sensitive"); + std::fs::create_dir_all(&skill_dir)?; + std::fs::write( + skill_dir.join("SKILL.md"), + "---\nname: budget-sensitive-skill\ndescription: visible only with the current model budget\n---\n\n# Body\n", + )?; + + let previous_model = "gpt-5.3-codex"; + let current_model = "gpt-5.2"; + let model_catalog = ModelsResponse { + models: vec![ + model_info_with_context_window_and_comp_hash( + previous_model, + /*context_window*/ 100, + "hash-a", + ), + model_info_with_context_window_and_comp_hash( + current_model, + /*context_window*/ 200_000, + "hash-b", + ), + ], + }; + let codex_home_path = codex_home.path().to_path_buf(); + let test = test_codex() + .with_home(codex_home) + .with_model(previous_model) + .with_config(move |config| { + config.cwd = codex_home_path.abs(); + config.model_catalog = Some(model_catalog); + config + .features + .enable(Feature::TokenBudget) + .expect("test config should allow token budget"); + }) + .build(&server) + .await?; + + submit_turn_with_model(&test, "before comp hash switch", previous_model).await?; + submit_turn_with_model(&test, "after comp hash switch", current_model).await?; + + let requests = responses.requests(); + assert_eq!(requests.len(), 2); + assert_eq!( + requests[0].body_json()["model"].as_str(), + Some(previous_model) + ); + assert_eq!( + requests[1].body_json()["model"].as_str(), + Some(current_model) + ); + + let initial_developer_text = requests[0].message_input_texts("developer").join("\n\n"); + assert!( + !initial_developer_text.contains("budget-sensitive-skill"), + "previous model's tiny context budget should omit skill metadata" + ); + let current_developer_text = requests[1].message_input_texts("developer").join("\n\n"); + assert!( + current_developer_text + .contains("budget-sensitive-skill: visible only with the current model budget"), + "pre-turn token-budget reset should rebuild full context with the current model budget, got {current_developer_text:?}" + ); + assert!( + !requests[1].body_contains_text("before comp hash switch"), + "token-budget comp-hash reset should drop prior window history" + ); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn token_budget_context_injects_plain_thread_hint_text() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/memories/write/src/runtime.rs b/codex-rs/memories/write/src/runtime.rs index 0835e15520..bd86d83130 100644 --- a/codex-rs/memories/write/src/runtime.rs +++ b/codex-rs/memories/write/src/runtime.rs @@ -75,6 +75,38 @@ pub(crate) struct MemoryStartupContext { session_telemetry: SessionTelemetry, } +fn build_session_telemetry( + auth_manager: &AuthManager, + thread_id: ThreadId, + config: &Config, + source: SessionSource, + model: &str, + originator: String, +) -> SessionTelemetry { + let auth = auth_manager.auth_cached(); + let auth = auth.as_ref(); + let auth_mode = auth.map(CodexAuth::auth_mode).map(TelemetryAuthMode::from); + let account_id = auth.and_then(CodexAuth::get_account_id); + let account_email = auth.and_then(CodexAuth::get_account_email); + let auth_env_telemetry = collect_auth_env_telemetry( + &config.model_provider, + auth_manager.codex_api_key_env_enabled(), + ); + SessionTelemetry::new( + thread_id, + model, + model, + account_id, + account_email, + auth_mode, + originator, + config.otel.log_user_prompt, + user_agent(), + source, + ) + .with_auth_env(auth_env_telemetry.to_otel_metadata()) +} + impl MemoryStartupContext { pub(crate) fn new( thread_manager: Arc, @@ -129,29 +161,15 @@ impl MemoryStartupContext { source: SessionSource, provider: SharedModelProvider, ) -> Self { - let auth = auth_manager.auth_cached(); - let auth = auth.as_ref(); - let auth_mode = auth.map(CodexAuth::auth_mode).map(TelemetryAuthMode::from); - let account_id = auth.and_then(CodexAuth::get_account_id); - let account_email = auth.and_then(CodexAuth::get_account_email); let model = config.model.as_deref().unwrap_or("unknown"); - let auth_env_telemetry = collect_auth_env_telemetry( - &config.model_provider, - auth_manager.codex_api_key_env_enabled(), - ); - let session_telemetry = SessionTelemetry::new( + let session_telemetry = build_session_telemetry( + &auth_manager, thread_id, - model, - model, - account_id, - account_email, - auth_mode, - originator().value, - config.otel.log_user_prompt, - user_agent(), + config, source, - ) - .with_auth_env(auth_env_telemetry.to_otel_metadata()); + model, + originator().value, + ); Self { thread_id, @@ -205,10 +223,14 @@ impl MemoryStartupContext { StageOneRequestContext { model_info, - session_telemetry: self - .session_telemetry - .clone() - .with_model(model_name, model_name), + session_telemetry: build_session_telemetry( + &self.auth_manager, + self.thread_id, + config, + config_snapshot.session_source, + model_name, + config_snapshot.originator, + ), reasoning_effort: Some(reasoning_effort), reasoning_summary, service_tier: config_snapshot.service_tier, @@ -231,6 +253,7 @@ impl MemoryStartupContext { self.thread_id, config.model_provider.clone(), session_source.clone(), + config_snapshot.originator, config.model_verbosity, config.features.enabled(Feature::EnableRequestCompression), config.features.enabled(Feature::RuntimeMetrics), diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 5a0341103f..444e440752 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -2615,11 +2615,33 @@ impl InitialHistory { .and_then(|meta| meta.thread_source.clone()) } + pub fn get_session_originator(&self) -> Option { + self.get_session_meta() + .map(|meta| meta.originator.clone()) + .filter(|originator| !originator.is_empty()) + } + pub fn get_resumed_parent_thread_id(&self) -> Option { self.get_resumed_session_meta() .and_then(|meta| meta.parent_thread_id) } + fn get_session_meta(&self) -> Option<&SessionMeta> { + match self { + InitialHistory::New | InitialHistory::Cleared => None, + InitialHistory::Resumed(resumed) => { + resumed.history.iter().find_map(|item| match item { + RolloutItem::SessionMeta(meta_line) => Some(&meta_line.meta), + _ => None, + }) + } + InitialHistory::Forked(items) => items.iter().find_map(|item| match item { + RolloutItem::SessionMeta(meta_line) => Some(&meta_line.meta), + _ => None, + }), + } + } + fn get_resumed_session_meta(&self) -> Option<&SessionMeta> { match self { InitialHistory::New | InitialHistory::Cleared | InitialHistory::Forked(_) => None, diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index 003f65cd93..721d801e11 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -58,6 +58,7 @@ use tokio::sync::Mutex; use tokio::sync::Semaphore; use tokio::sync::watch; use tokio::time; +use tracing::instrument; use tracing::warn; use crate::elicitation_client_service::ElicitationClientService; @@ -420,6 +421,7 @@ impl RmcpClient { /// Perform the initialization handshake with the MCP server. /// https://modelcontextprotocol.io/specification/2025-06-18/basic/lifecycle#initialization + #[instrument(level = "trace", skip_all)] pub async fn initialize( &self, params: InitializeRequestParams, @@ -501,6 +503,7 @@ impl RmcpClient { Ok(result) } + #[instrument(level = "trace", skip_all)] pub async fn list_tools_with_connector_ids( &self, params: Option, diff --git a/codex-rs/rollout/Cargo.toml b/codex-rs/rollout/Cargo.toml index edc9059688..838081bb53 100644 --- a/codex-rs/rollout/Cargo.toml +++ b/codex-rs/rollout/Cargo.toml @@ -17,7 +17,6 @@ anyhow = { workspace = true } chrono = { workspace = true, features = ["serde"] } codex-file-search = { workspace = true } codex-git-utils = { workspace = true } -codex-login = { workspace = true } codex-otel = { workspace = true } codex-protocol = { workspace = true } codex-state = { workspace = true } diff --git a/codex-rs/rollout/src/lib.rs b/codex-rs/rollout/src/lib.rs index 804b8ac101..3c103750c1 100644 --- a/codex-rs/rollout/src/lib.rs +++ b/codex-rs/rollout/src/lib.rs @@ -16,10 +16,6 @@ pub(crate) mod session_index; mod sqlite_metrics; pub mod state_db; -pub(crate) mod default_client { - pub use codex_login::default_client::*; -} - pub(crate) use codex_protocol::protocol; pub const SESSIONS_SUBDIR: &str = "sessions"; diff --git a/codex-rs/rollout/src/recorder.rs b/codex-rs/rollout/src/recorder.rs index 53bf115d8d..dad7f8bc32 100644 --- a/codex-rs/rollout/src/recorder.rs +++ b/codex-rs/rollout/src/recorder.rs @@ -46,7 +46,6 @@ use super::list::parse_timestamp_uuid_from_filename; use super::metadata; use super::session_index::find_thread_names_by_ids; use crate::config::RolloutConfigView; -use crate::default_client::originator; use crate::state_db; use crate::state_db::StateDbHandle; use codex_git_utils::collect_git_info; @@ -89,6 +88,7 @@ pub enum RolloutRecorderParams { parent_thread_id: Option, source: Box, thread_source: Option, + originator: String, base_instructions: BaseInstructions, dynamic_tools: Vec, multi_agent_version: Option, @@ -161,12 +161,14 @@ fn clone_io_error(err: &IoError) -> IoError { } impl RolloutRecorderParams { + #[allow(clippy::too_many_arguments)] pub fn new( conversation_id: ThreadId, forked_from_id: Option, parent_thread_id: Option, source: SessionSource, thread_source: Option, + originator: String, base_instructions: BaseInstructions, dynamic_tools: Vec, ) -> Self { @@ -177,6 +179,7 @@ impl RolloutRecorderParams { parent_thread_id, source: Box::new(source), thread_source, + originator, base_instructions, dynamic_tools, multi_agent_version: None, @@ -726,6 +729,7 @@ impl RolloutRecorder { parent_thread_id, source, thread_source, + originator, base_instructions, dynamic_tools, multi_agent_version, @@ -751,7 +755,7 @@ impl RolloutRecorder { parent_thread_id, timestamp, cwd: config.cwd().to_path_buf(), - originator: originator().value, + originator, cli_version: env!("CARGO_PKG_VERSION").to_string(), agent_nickname: source.get_nickname(), agent_role: source.get_agent_role(), diff --git a/codex-rs/rollout/src/recorder_tests.rs b/codex-rs/rollout/src/recorder_tests.rs index 23618cb815..4b72ff7673 100644 --- a/codex-rs/rollout/src/recorder_tests.rs +++ b/codex-rs/rollout/src/recorder_tests.rs @@ -385,6 +385,7 @@ async fn recorder_materializes_on_flush_with_pending_items() -> std::io::Result< /*parent_thread_id*/ None, SessionSource::Exec, /*thread_source*/ None, + "test_originator".to_string(), BaseInstructions::default(), Vec::new(), ) @@ -477,6 +478,7 @@ async fn persist_reports_filesystem_error_and_retries_buffered_items() -> std::i /*parent_thread_id*/ None, SessionSource::Exec, /*thread_source*/ None, + "test_originator".to_string(), BaseInstructions::default(), Vec::new(), ), diff --git a/codex-rs/thread-store/src/in_memory.rs b/codex-rs/thread-store/src/in_memory.rs index d7d85c3cfb..8cd801c9c0 100644 --- a/codex-rs/thread-store/src/in_memory.rs +++ b/codex-rs/thread-store/src/in_memory.rs @@ -118,6 +118,7 @@ mod tests { parent_thread_id, source: SessionSource::Exec, thread_source: None, + originator: "test_originator".to_string(), base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), multi_agent_version: None, @@ -242,6 +243,7 @@ impl InMemoryThreadStore { agent_nickname: params.source.get_nickname(), agent_role: params.source.get_agent_role(), agent_path: params.source.get_agent_path().map(Into::into), + originator: params.originator.clone(), source: params.source.clone(), thread_source: params.thread_source.clone(), model_provider: Some(params.metadata.model_provider.clone()), diff --git a/codex-rs/thread-store/src/local/create_thread.rs b/codex-rs/thread-store/src/local/create_thread.rs index 36c3cdc13a..eba6a5a6a9 100644 --- a/codex-rs/thread-store/src/local/create_thread.rs +++ b/codex-rs/thread-store/src/local/create_thread.rs @@ -33,6 +33,7 @@ pub(super) async fn create_thread( params.parent_thread_id, params.source, params.thread_source, + params.originator, params.base_instructions, params.dynamic_tools, ) diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs index d0a4d65187..51e622fc5b 100644 --- a/codex-rs/thread-store/src/local/mod.rs +++ b/codex-rs/thread-store/src/local/mod.rs @@ -1130,6 +1130,7 @@ mod tests { parent_thread_id: None, source: SessionSource::Exec, thread_source: None, + originator: "test_originator".to_string(), base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), multi_agent_version: None, diff --git a/codex-rs/thread-store/src/types.rs b/codex-rs/thread-store/src/types.rs index c18f398821..feea814f94 100644 --- a/codex-rs/thread-store/src/types.rs +++ b/codex-rs/thread-store/src/types.rs @@ -79,6 +79,8 @@ pub struct CreateThreadParams { pub source: SessionSource, /// Optional analytics source classification for this thread. pub thread_source: Option, + /// Effective originator used for this thread's Responses requests and analytics events. + pub originator: String, /// Base instructions persisted in session metadata. pub base_instructions: BaseInstructions, /// Dynamic tools available to the thread at startup.