From e4f53e5e990a2fccff310f0967bbd83b9bd4f2cf Mon Sep 17 00:00:00 2001 From: Eric Traut Date: Thu, 19 Mar 2026 09:23:29 -0600 Subject: [PATCH] Preserve CLI originator for codex-tui app-server sessions --- codex-rs/app-server/src/message_processor.rs | 9 ++- .../app-server/tests/suite/v2/initialize.rs | 35 +++++++++ .../app-server/tests/suite/v2/turn_start.rs | 77 +++++++++++++++++++ 3 files changed, 120 insertions(+), 1 deletion(-) diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 3804c4f9b4..0d21d9c867 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -57,6 +57,7 @@ use codex_core::auth::ExternalAuthTokens; use codex_core::config::Config; use codex_core::config_loader::CloudRequirementsLoader; use codex_core::config_loader::LoaderOverrides; +use codex_core::default_client::DEFAULT_ORIGINATOR; use codex_core::default_client::SetOriginatorError; use codex_core::default_client::USER_AGENT_SUFFIX; use codex_core::default_client::get_codex_user_agent; @@ -77,6 +78,7 @@ use toml::Value as TomlValue; use tracing::Instrument; const EXTERNAL_AUTH_REFRESH_TIMEOUT: Duration = Duration::from_secs(10); +const TUI_APP_SERVER_CLIENT_NAME: &str = "codex-tui"; #[derive(Clone)] struct ExternalAuthRefreshBridge { @@ -550,7 +552,12 @@ impl MessageProcessor { } = params.client_info; session.app_server_client_name = Some(name.clone()); session.client_version = Some(version.clone()); - if let Err(error) = set_default_originator(name.clone()) { + let originator = if name.eq_ignore_ascii_case(TUI_APP_SERVER_CLIENT_NAME) { + DEFAULT_ORIGINATOR.to_string() + } else { + name.clone() + }; + if let Err(error) = set_default_originator(originator) { match error { SetOriginatorError::InvalidHeaderValue => { let error = JSONRPCErrorError { diff --git a/codex-rs/app-server/tests/suite/v2/initialize.rs b/codex-rs/app-server/tests/suite/v2/initialize.rs index 9b5f0cacc8..ce20fb7153 100644 --- a/codex-rs/app-server/tests/suite/v2/initialize.rs +++ b/codex-rs/app-server/tests/suite/v2/initialize.rs @@ -14,6 +14,7 @@ use codex_app_server_protocol::ThreadStartResponse; use codex_app_server_protocol::TurnStartParams; use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::UserInput as V2UserInput; +use codex_core::default_client::DEFAULT_ORIGINATOR; use codex_utils_cargo_bin::cargo_bin; use core_test_support::fs_wait; use pretty_assertions::assert_eq; @@ -24,6 +25,7 @@ use tempfile::TempDir; use tokio::time::timeout; const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); +const TUI_APP_SERVER_CLIENT_NAME: &str = "codex-tui"; #[tokio::test] async fn initialize_uses_client_info_name_as_originator() -> Result<()> { @@ -58,6 +60,39 @@ async fn initialize_uses_client_info_name_as_originator() -> Result<()> { Ok(()) } +#[tokio::test] +async fn initialize_keeps_codex_tui_on_cli_originator() -> Result<()> { + let responses = Vec::new(); + let server = create_mock_responses_server_sequence_unchecked(responses).await; + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), &server.uri(), "never")?; + let mut mcp = McpProcess::new(codex_home.path()).await?; + + let message = timeout( + DEFAULT_READ_TIMEOUT, + mcp.initialize_with_client_info(ClientInfo { + name: TUI_APP_SERVER_CLIENT_NAME.to_string(), + title: Some("Codex TUI".to_string()), + version: "0.1.0".to_string(), + }), + ) + .await??; + + let JSONRPCMessage::Response(response) = message else { + anyhow::bail!("expected initialize response, got {message:?}"); + }; + let InitializeResponse { + user_agent, + platform_family, + platform_os, + } = to_response::(response)?; + + assert!(user_agent.starts_with(&format!("{DEFAULT_ORIGINATOR}/"))); + assert_eq!(platform_family, std::env::consts::FAMILY); + assert_eq!(platform_os, std::env::consts::OS); + Ok(()) +} + #[tokio::test] async fn initialize_respects_originator_override_env_var() -> Result<()> { let responses = Vec::new(); diff --git a/codex-rs/app-server/tests/suite/v2/turn_start.rs b/codex-rs/app-server/tests/suite/v2/turn_start.rs index e235ee6b61..5c47906692 100644 --- a/codex-rs/app-server/tests/suite/v2/turn_start.rs +++ b/codex-rs/app-server/tests/suite/v2/turn_start.rs @@ -45,6 +45,7 @@ use codex_app_server_protocol::TurnStartedNotification; use codex_app_server_protocol::TurnStatus; use codex_app_server_protocol::UserInput as V2UserInput; use codex_core::config::ConfigToml; +use codex_core::default_client::DEFAULT_ORIGINATOR; use codex_core::features::FEATURES; use codex_core::features::Feature; use codex_core::personality_migration::PERSONALITY_MIGRATION_FILENAME; @@ -70,6 +71,7 @@ const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs #[cfg(not(windows))] const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); const TEST_ORIGINATOR: &str = "codex_vscode"; +const TUI_APP_SERVER_CLIENT_NAME: &str = "codex-tui"; const LOCAL_PRAGMATIC_TEMPLATE: &str = "You are a deeply pragmatic, effective software engineer."; fn body_contains(req: &wiremock::Request, text: &str) -> bool { @@ -153,6 +155,81 @@ async fn turn_start_sends_originator_header() -> Result<()> { Ok(()) } +#[tokio::test] +async fn turn_start_sends_cli_originator_header_for_codex_tui() -> Result<()> { + let responses = vec![create_final_assistant_message_sse_response("Done")?]; + let server = create_mock_responses_server_sequence_unchecked(responses).await; + + let codex_home = TempDir::new()?; + create_config_toml( + codex_home.path(), + &server.uri(), + "never", + &BTreeMap::from([(Feature::Personality, true)]), + )?; + + let mut mcp = McpProcess::new(codex_home.path()).await?; + timeout( + DEFAULT_READ_TIMEOUT, + mcp.initialize_with_client_info(ClientInfo { + name: TUI_APP_SERVER_CLIENT_NAME.to_string(), + title: Some("Codex TUI".to_string()), + version: "0.1.0".to_string(), + }), + ) + .await??; + + let thread_req = mcp + .send_thread_start_request(ThreadStartParams { + model: Some("mock-model".to_string()), + ..Default::default() + }) + .await?; + let thread_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), + ) + .await??; + let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; + + let turn_req = mcp + .send_turn_start_request(TurnStartParams { + thread_id: thread.id.clone(), + input: vec![V2UserInput::Text { + text: "Hello".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }) + .await?; + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), + ) + .await??; + + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("turn/completed"), + ) + .await??; + + let requests = server + .received_requests() + .await + .expect("failed to fetch received requests"); + assert!(!requests.is_empty()); + for request in requests { + let originator = request + .headers + .get("originator") + .expect("originator header missing"); + assert_eq!(originator.to_str()?, DEFAULT_ORIGINATOR); + } + + Ok(()) +} + #[tokio::test] async fn turn_start_emits_user_message_item_with_text_elements() -> Result<()> { let responses = vec![create_final_assistant_message_sse_response("Done")?];