diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index b6f7505d1b..fcd827fddc 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -3945,6 +3945,12 @@ impl Session { arguments: Option, meta: Option, ) -> anyhow::Result { + if server == CODEX_APPS_MCP_SERVER_NAME + && let Some((turn_context, _)) = self.active_turn_context_and_cancellation_token().await + { + self.sync_mcp_request_headers_for_turn(turn_context.as_ref()) + .await; + } self.services .mcp_connection_manager .read() diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index f767c05f4e..22e83ee216 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/codex-rs/core/src/codex_tests.rs @@ -26,6 +26,8 @@ use codex_protocol::protocol::ReadOnlyAccess; use codex_protocol::protocol::SandboxPolicy; use codex_protocol::request_permissions::PermissionGrantScope; use codex_protocol::request_permissions::RequestPermissionProfile; +use codex_protocol::user_input::UserInput; +use reqwest::header::HeaderValue; use tracing::Span; use crate::protocol::CompactedItem; @@ -56,6 +58,7 @@ use crate::tools::handlers::UnifiedExecHandler; use crate::tools::registry::ToolHandler; use crate::tools::router::ToolCallSource; use crate::turn_diff_tracker::TurnDiffTracker; +use async_trait::async_trait; use codex_app_server_protocol::AppInfo; use codex_execpolicy::Decision; use codex_execpolicy::NetworkRuleProtocol; @@ -78,7 +81,10 @@ use opentelemetry::trace::TracerProvider as _; use opentelemetry_sdk::trace::SdkTracerProvider; use std::path::Path; use std::time::Duration; +use tokio::sync::Notify; use tokio::time::sleep; +use tokio_util::sync::CancellationToken; +use tokio_util::task::AbortOnDropHandle; use tracing_opentelemetry::OpenTelemetrySpanExt; use tracing_subscriber::prelude::*; @@ -2716,6 +2722,99 @@ async fn request_permissions_is_auto_denied_when_granular_policy_blocks_tool_req ); } +struct NoopTask; + +#[async_trait] +impl SessionTask for NoopTask { + fn kind(&self) -> TaskKind { + TaskKind::Regular + } + + fn span_name(&self) -> &'static str { + "noop" + } + + async fn run( + self: Arc, + _ctx: Arc, + _turn_context: Arc, + _input: Vec, + _cancellation_token: CancellationToken, + ) -> Option { + None + } +} + +#[tokio::test] +async fn call_tool_refreshes_codex_apps_request_headers_from_active_turn() { + let (session, turn_context) = make_session_and_context().await; + let session = Arc::new(session); + let turn_context = Arc::new(turn_context); + + session + .services + .mcp_connection_manager + .write() + .await + .register_test_server_for_request_headers(CODEX_APPS_MCP_SERVER_NAME); + + session + .sync_mcp_request_headers_for_turn(turn_context.as_ref()) + .await; + let base_header = turn_context + .turn_metadata_state + .current_header_value() + .expect("base turn metadata header"); + + let updated_header = serde_json::json!({ + "turn_id": turn_context.sub_id, + "sandbox": "test", + "workspaces": { + "/tmp/repo": { + "has_changes": true + } + } + }) + .to_string(); + turn_context + .turn_metadata_state + .set_enriched_header_for_tests(Some(updated_header.clone())); + + let mut active_turn = ActiveTurn::default(); + let handle = tokio::spawn(async {}); + active_turn.add_task(crate::state::RunningTask { + done: Arc::new(Notify::new()), + kind: TaskKind::Regular, + task: Arc::new(NoopTask), + cancellation_token: CancellationToken::new(), + handle: Arc::new(AbortOnDropHandle::new(handle)), + turn_context: Arc::clone(&turn_context), + _timer: None, + }); + *session.active_turn.lock().await = Some(active_turn); + + let _err = session + .call_tool(CODEX_APPS_MCP_SERVER_NAME, "echo", None, None) + .await + .expect_err("test server is not initialized"); + + let headers = session + .services + .mcp_connection_manager + .read() + .await + .request_headers_for_server(CODEX_APPS_MCP_SERVER_NAME) + .expect("request headers should be tracked for codex apps"); + assert_eq!( + headers.get(crate::X_CODEX_TURN_METADATA_HEADER), + Some(&HeaderValue::from_str(&updated_header).expect("valid enriched header")), + ); + assert_ne!( + headers.get(crate::X_CODEX_TURN_METADATA_HEADER), + Some(&HeaderValue::from_str(&base_header).expect("valid base header")), + ); +} + #[tokio::test] async fn submit_with_id_captures_current_span_trace_context() { let (session, _turn_context) = make_session_and_context().await; diff --git a/codex-rs/core/src/mcp_connection_manager.rs b/codex-rs/core/src/mcp_connection_manager.rs index 7c8a343070..cd80220d71 100644 --- a/codex-rs/core/src/mcp_connection_manager.rs +++ b/codex-rs/core/src/mcp_connection_manager.rs @@ -636,6 +636,40 @@ impl McpConnectionManager { Self::new_uninitialized(approval_policy) } + #[cfg(test)] + pub(crate) fn register_test_server_for_request_headers(&mut self, server_name: &str) { + let failed_client = futures::future::ready::>( + Err(StartupOutcomeError::Failed { + error: "test request headers stub".to_string(), + }), + ) + .boxed() + .shared(); + self.clients.insert( + server_name.to_string(), + AsyncManagedClient { + client: failed_client, + request_headers: Arc::new(StdMutex::new(None)), + startup_snapshot: None, + startup_complete: Arc::new(AtomicBool::new(true)), + tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), + }, + ); + } + + #[cfg(test)] + pub(crate) fn request_headers_for_server( + &self, + server_name: &str, + ) -> Option { + let client = self.clients.get(server_name)?; + client + .request_headers + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone() + } + pub(crate) fn has_servers(&self) -> bool { !self.clients.is_empty() } diff --git a/codex-rs/core/src/turn_metadata.rs b/codex-rs/core/src/turn_metadata.rs index c0298c5221..1056752728 100644 --- a/codex-rs/core/src/turn_metadata.rs +++ b/codex-rs/core/src/turn_metadata.rs @@ -217,6 +217,14 @@ impl TurnMetadataState { } } + #[cfg(test)] + pub(crate) fn set_enriched_header_for_tests(&self, header: Option) { + *self + .enriched_header + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner) = header; + } + async fn fetch_workspace_git_metadata(&self) -> WorkspaceGitMetadata { let (latest_git_commit_hash, associated_remote_urls, has_changes) = tokio::join!( get_head_commit_hash(&self.cwd), diff --git a/codex-rs/core/tests/suite/search_tool.rs b/codex-rs/core/tests/suite/search_tool.rs index 118f1bd585..62659aa4de 100644 --- a/codex-rs/core/tests/suite/search_tool.rs +++ b/codex-rs/core/tests/suite/search_tool.rs @@ -30,6 +30,8 @@ use core_test_support::wait_for_event; use pretty_assertions::assert_eq; use serde_json::Value; use serde_json::json; +use std::fs; +use std::process::Command; const SEARCH_TOOL_DESCRIPTION_SNIPPETS: [&str; 2] = [ "You have access to all the tools of the following apps/connectors", @@ -86,6 +88,15 @@ fn tool_search_output_tools(request: &ResponsesRequest, call_id: &str) -> Vec Option { + request + .body_json::() + .ok()? + .get("method") + .and_then(Value::as_str) + .map(str::to_string) +} + fn configure_apps(config: &mut Config, apps_base_url: &str) { config .features @@ -499,5 +510,195 @@ async fn tool_search_returns_deferred_tools_without_follow_up_tool_injection() - "post-tool follow-up should still rely on tool_search_output history, not tool injection: {third_request_tools:?}" ); + let mcp_requests = server + .received_requests() + .await + .expect("failed to fetch recorded requests"); + let tools_list_request = mcp_requests + .iter() + .find(|request| json_rpc_method(request).as_deref() == Some("tools/list")) + .expect("tools/list MCP request"); + assert!( + tools_list_request + .headers + .get("x-codex-turn-metadata") + .is_none(), + "tools/list should not include per-turn MCP headers" + ); + + let tools_call_request = mcp_requests + .iter() + .find(|request| json_rpc_method(request).as_deref() == Some("tools/call")) + .expect("tools/call MCP request"); + let session_id_header = tools_call_request + .headers + .get("session_id") + .expect("tools/call session_id header"); + let request_id_header = tools_call_request + .headers + .get("x-client-request-id") + .expect("tools/call x-client-request-id header"); + let turn_metadata_header = tools_call_request + .headers + .get("x-codex-turn-metadata") + .expect("tools/call turn metadata header"); + assert_eq!( + session_id_header + .to_str() + .expect("session_id header to be utf8"), + request_id_header + .to_str() + .expect("x-client-request-id header to be utf8") + ); + assert!( + turn_metadata_header + .to_str() + .expect("turn metadata header to be utf8") + .contains("\"turn_id\""), + "expected turn metadata header to contain serialized turn metadata" + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn apps_mcp_tool_call_uses_enriched_turn_metadata_header() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let apps_server = AppsTestServer::mount_searchable(&server).await?; + let call_id = "tool-search-git-metadata"; + let mock = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_tool_search_call( + call_id, + &json!({ + "query": "create calendar event", + "limit": 1, + }), + ), + ev_completed("resp-1"), + ]), + sse(vec![ + ev_response_created("resp-2"), + json!({ + "type": "response.output_item.done", + "item": { + "type": "function_call", + "call_id": "calendar-call-git-metadata", + "name": SEARCH_CALENDAR_CREATE_TOOL, + "namespace": SEARCH_CALENDAR_NAMESPACE, + "arguments": serde_json::to_string(&json!({ + "title": "Lunch", + "starts_at": "2026-03-10T12:00:00Z" + })).expect("serialize calendar args") + } + }), + ev_completed("resp-2"), + ]), + sse(vec![ + ev_response_created("resp-3"), + ev_assistant_message("msg-1", "done"), + ev_completed("resp-3"), + ]), + ], + ) + .await; + + let mut builder = configured_builder(apps_server.chatgpt_base_url.clone()); + let test = builder.build(&server).await?; + + let cwd = test.cwd_path().to_path_buf(); + assert!( + Command::new("git") + .arg("init") + .current_dir(&cwd) + .status()? + .success() + ); + assert!( + Command::new("git") + .args(["config", "user.name", "Codex Test"]) + .current_dir(&cwd) + .status()? + .success() + ); + assert!( + Command::new("git") + .args(["config", "user.email", "codex@example.com"]) + .current_dir(&cwd) + .status()? + .success() + ); + assert!( + Command::new("git") + .args(["remote", "add", "origin", "https://example.test/repo.git"]) + .current_dir(&cwd) + .status()? + .success() + ); + for idx in 0..400 { + fs::write( + cwd.join(format!("file-{idx:04}.txt")), + format!("fixture file {idx}\n"), + )?; + } + assert!( + Command::new("git") + .args(["add", "."]) + .current_dir(&cwd) + .status()? + .success() + ); + assert!( + Command::new("git") + .args(["commit", "-m", "init"]) + .current_dir(&cwd) + .status()? + .success() + ); + + test.codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: "Find the calendar create tool".to_string(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + }) + .await?; + + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + let _requests = mock.requests(); + let mcp_requests = server + .received_requests() + .await + .expect("failed to fetch recorded requests"); + let tools_call_request = mcp_requests + .iter() + .find(|request| json_rpc_method(request).as_deref() == Some("tools/call")) + .expect("tools/call MCP request"); + let turn_metadata_header = tools_call_request + .headers + .get("x-codex-turn-metadata") + .expect("tools/call turn metadata header") + .to_str() + .expect("turn metadata header to be utf8"); + let parsed: Value = serde_json::from_str(turn_metadata_header)?; + assert!( + parsed + .get("workspaces") + .and_then(Value::as_object) + .is_some_and(|workspaces| !workspaces.is_empty()), + "expected enriched MCP turn metadata header with workspace git metadata, got {parsed:#?}" + ); + Ok(()) }