From b267af326de4a1b0a98a1ddbda05cbdc03dcb5ce Mon Sep 17 00:00:00 2001 From: pash Date: Thu, 29 Jan 2026 18:04:43 -0800 Subject: [PATCH] moving to headers --- codex-rs/core/src/client.rs | 23 ++++- codex-rs/core/src/codex.rs | 89 +++++++++------- codex-rs/core/src/environment_context.rs | 124 +++-------------------- codex-rs/core/src/lib.rs | 1 + codex-rs/core/tests/responses_headers.rs | 95 +++++++++++++++++ 5 files changed, 183 insertions(+), 149 deletions(-) diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index f01c145e56..5022ab9de0 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -69,6 +69,7 @@ use crate::transport_manager::TransportManager; pub const WEB_SEARCH_ELIGIBLE_HEADER: &str = "x-oai-web-search-eligible"; pub const X_CODEX_TURN_STATE_HEADER: &str = "x-codex-turn-state"; +pub const X_CODEX_TURN_METADATA_HEADER: &str = "x-codex-turn-metadata"; #[derive(Debug)] struct ModelClientState { @@ -105,6 +106,8 @@ pub struct ModelClientSession { /// keep sending it unchanged between turn requests (e.g., for retries, incremental /// appends, or continuation requests), and must not send it between different turns. turn_state: Arc>, + /// Turn-scoped metadata attached to every request in the turn. + turn_metadata_header: Option, } #[allow(clippy::too_many_arguments)] @@ -138,12 +141,22 @@ impl ModelClient { } pub fn new_session(&self) -> ModelClientSession { + self.new_session_with_turn_metadata(None) + } + + pub fn new_session_with_turn_metadata( + &self, + turn_metadata_header: Option, + ) -> ModelClientSession { + let turn_metadata_header = + turn_metadata_header.and_then(|value| HeaderValue::from_str(&value).ok()); ModelClientSession { state: Arc::clone(&self.state), connection: None, websocket_last_items: Vec::new(), transport_manager: self.state.transport_manager.clone(), turn_state: Arc::new(OnceLock::new()), + turn_metadata_header, } } } @@ -377,7 +390,11 @@ impl ModelClientSession { store_override: None, conversation_id: Some(conversation_id), session_source: Some(self.state.session_source.clone()), - extra_headers: build_responses_headers(&self.state.config, Some(&self.turn_state)), + extra_headers: build_responses_headers( + &self.state.config, + Some(&self.turn_state), + self.turn_metadata_header.as_ref(), + ), compression, turn_state: Some(Arc::clone(&self.turn_state)), } @@ -698,6 +715,7 @@ fn experimental_feature_headers(config: &Config) -> ApiHeaderMap { fn build_responses_headers( config: &Config, turn_state: Option<&Arc>>, + turn_metadata_header: Option<&HeaderValue>, ) -> ApiHeaderMap { let mut headers = experimental_feature_headers(config); headers.insert( @@ -716,6 +734,9 @@ fn build_responses_headers( { headers.insert(X_CODEX_TURN_STATE_HEADER, header_value); } + if let Some(header_value) = turn_metadata_header { + headers.insert(X_CODEX_TURN_METADATA_HEADER, header_value.clone()); + } headers } diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 680a27a199..b400b56ca1 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -79,6 +79,7 @@ use mcp_types::ListResourcesResult; use mcp_types::ReadResourceRequestParams; use mcp_types::ReadResourceResult; use mcp_types::RequestId; +use serde::Serialize; use serde_json; use serde_json::Value; use tokio::sync::Mutex; @@ -112,14 +113,15 @@ use crate::config::types::McpServerConfig; use crate::config::types::ShellEnvironmentPolicy; use crate::context_manager::ContextManager; use crate::environment_context::EnvironmentContext; -use crate::environment_context::WorkspaceConfiguration; -use crate::environment_context::WorkspaceEntry; use crate::error::CodexErr; use crate::error::Result as CodexResult; #[cfg(test)] use crate::exec::StreamOutput; use crate::exec_policy::ExecPolicyUpdateError; use crate::feedback_tags; +use crate::git_info::collect_git_info; +use crate::git_info::get_git_remote_urls; +use crate::git_info::get_git_repo_root; use crate::instructions::UserInstructions; use crate::mcp::CODEX_APPS_MCP_SERVER_NAME; use crate::mcp::auth::compute_auth_statuses; @@ -505,6 +507,8 @@ pub(crate) struct TurnContext { pub(crate) tool_call_gate: Arc, pub(crate) truncation_policy: TruncationPolicy, pub(crate) dynamic_tools: Vec, + /// Per-turn metadata serialized as a header value for outbound model requests. + pub(crate) turn_metadata_header: Option, } impl TurnContext { @@ -706,9 +710,45 @@ impl Session { tool_call_gate: Arc::new(ReadinessFlag::new()), truncation_policy: model_info.truncation_policy.into(), dynamic_tools: session_configuration.dynamic_tools.clone(), + turn_metadata_header: None, } } + async fn build_turn_metadata_header(cwd: &Path) -> Option { + #[derive(Serialize)] + struct TurnMetadataWorkspace { + #[serde(skip_serializing_if = "Option::is_none")] + associated_remote_urls: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + latest_git_commit_hash: Option, + } + + #[derive(Serialize)] + struct TurnMetadata { + workspaces: BTreeMap, + } + + if get_git_repo_root(cwd).is_none() { + return None; + } + let git_info = collect_git_info(cwd).await?; + let latest_git_commit_hash = git_info.commit_hash; + let associated_remote_urls = get_git_remote_urls(cwd).await; + if latest_git_commit_hash.is_none() && associated_remote_urls.is_none() { + return None; + } + + let mut workspaces = BTreeMap::new(); + workspaces.insert( + cwd.to_string_lossy().into_owned(), + TurnMetadataWorkspace { + associated_remote_urls, + latest_git_commit_hash, + }, + ); + serde_json::to_string(&TurnMetadata { workspaces }).ok() + } + #[allow(clippy::too_many_arguments)] async fn new( mut session_configuration: SessionConfiguration, @@ -1245,6 +1285,8 @@ impl Session { if let Some(final_schema) = final_output_json_schema { turn_context.final_output_json_schema = final_schema; } + turn_context.turn_metadata_header = + Self::build_turn_metadata_header(turn_context.cwd.as_path()).await; Arc::new(turn_context) } @@ -1871,28 +1913,6 @@ impl Session { &self, turn_context: &TurnContext, ) -> Vec { - fn build_workspace_configuration( - cwd: &Path, - git_info: Option<&codex_protocol::protocol::GitInfo>, - remote_urls: Option<&BTreeMap>, - ) -> Option { - let latest_git_commit_hash = git_info.and_then(|info| info.commit_hash.clone()); - let associated_remote_urls = remote_urls.cloned(); - if latest_git_commit_hash.is_none() && associated_remote_urls.is_none() { - return None; - } - - let mut workspaces = BTreeMap::new(); - workspaces.insert( - cwd.to_string_lossy().into_owned(), - WorkspaceEntry { - associated_remote_urls, - latest_git_commit_hash, - }, - ); - Some(WorkspaceConfiguration { workspaces }) - } - let mut items = Vec::::with_capacity(4); let shell = self.user_shell(); items.push( @@ -1945,21 +1965,9 @@ impl Session { .into(), ); } - let git_info = crate::git_info::collect_git_info(turn_context.cwd.as_path()).await; - let remote_urls = if git_info.is_some() { - crate::git_info::get_git_remote_urls(turn_context.cwd.as_path()).await - } else { - None - }; - let workspace_configuration = build_workspace_configuration( - turn_context.cwd.as_path(), - git_info.as_ref(), - remote_urls.as_ref(), - ); items.push(ResponseItem::from(EnvironmentContext::new( Some(turn_context.cwd.clone()), shell.as_ref().clone(), - workspace_configuration, ))); items } @@ -3234,6 +3242,8 @@ async fn spawn_review_thread( parent_turn_context.client.transport_manager(), ); + let turn_metadata_header = + Session::build_turn_metadata_header(parent_turn_context.cwd.as_path()).await; let review_turn_context = TurnContext { sub_id: sub_id.to_string(), client, @@ -3254,6 +3264,7 @@ async fn spawn_review_thread( tool_call_gate: Arc::new(ReadinessFlag::new()), dynamic_tools: parent_turn_context.dynamic_tools.clone(), truncation_policy: model_info.truncation_policy.into(), + turn_metadata_header, }; // Seed the child task with the review prompt as the initial user message. @@ -3449,7 +3460,9 @@ pub(crate) async fn run_turn( // many turns, from the perspective of the user, it is a single turn. let turn_diff_tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); - let mut client_session = turn_context.client.new_session(); + let mut client_session = turn_context + .client + .new_session_with_turn_metadata(turn_context.turn_metadata_header.clone()); loop { // Note that pending_input would be something like a message the user @@ -4512,8 +4525,6 @@ pub(super) fn get_last_assistant_message_from_turn(responses: &[ResponseItem]) - #[cfg(test)] pub(crate) use tests::make_session_and_context; - -use crate::git_info::get_git_repo_root; #[cfg(test)] pub(crate) use tests::make_session_and_context_with_rx; diff --git a/codex-rs/core/src/environment_context.rs b/codex-rs/core/src/environment_context.rs index e57f8461aa..d0ad3d6f0d 100644 --- a/codex-rs/core/src/environment_context.rs +++ b/codex-rs/core/src/environment_context.rs @@ -6,7 +6,6 @@ use codex_protocol::protocol::ENVIRONMENT_CONTEXT_CLOSE_TAG; use codex_protocol::protocol::ENVIRONMENT_CONTEXT_OPEN_TAG; use serde::Deserialize; use serde::Serialize; -use std::collections::BTreeMap; use std::path::PathBuf; #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] @@ -14,22 +13,11 @@ use std::path::PathBuf; pub(crate) struct EnvironmentContext { pub cwd: Option, pub shell: Shell, - /// Workspace metadata, structured to support future multi-root workspaces. - #[serde(skip_serializing_if = "Option::is_none")] - pub workspace_configuration: Option, } impl EnvironmentContext { - pub fn new( - cwd: Option, - shell: Shell, - workspace_configuration: Option, - ) -> Self { - Self { - cwd, - shell, - workspace_configuration, - } + pub fn new(cwd: Option, shell: Shell) -> Self { + Self { cwd, shell } } /// Compares two environment contexts, ignoring the shell. Useful when @@ -52,30 +40,14 @@ impl EnvironmentContext { } else { None }; - // Only include workspace configuration on the initial prefix message. - EnvironmentContext::new(cwd, shell.clone(), None) + EnvironmentContext::new(cwd, shell.clone()) } pub fn from_turn_context(turn_context: &TurnContext, shell: &Shell) -> Self { - // Only include workspace configuration on the initial prefix message. - Self::new(Some(turn_context.cwd.clone()), shell.clone(), None) + Self::new(Some(turn_context.cwd.clone()), shell.clone()) } } -/// Multi-root-friendly workspace metadata modeled after Cline's structure. -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] -pub(crate) struct WorkspaceConfiguration { - pub workspaces: BTreeMap, -} - -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] -pub(crate) struct WorkspaceEntry { - #[serde(skip_serializing_if = "Option::is_none")] - pub associated_remote_urls: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - pub latest_git_commit_hash: Option, -} - impl EnvironmentContext { /// Serializes the environment context to XML. Libraries like `quick-xml` /// require custom macros to handle Enums with newtypes, so we just do it @@ -96,32 +68,6 @@ impl EnvironmentContext { let shell_name = self.shell.name(); lines.push(format!(" {shell_name}")); - if let Some(workspace_configuration) = self.workspace_configuration { - lines.push(" ".to_string()); - for (path, workspace) in workspace_configuration.workspaces { - lines.push(format!(" ")); - - if let Some(latest_git_commit_hash) = workspace.latest_git_commit_hash { - lines.push(format!( - " {latest_git_commit_hash}" - )); - } - - if let Some(associated_remote_urls) = workspace.associated_remote_urls - && !associated_remote_urls.is_empty() - { - lines.push(" ".to_string()); - for (name, url) in associated_remote_urls { - lines.push(format!(" {url}")); - } - lines.push(" ".to_string()); - } - - lines.push(" ".to_string()); - } - lines.push(" ".to_string()); - } - lines.push(ENVIRONMENT_CONTEXT_CLOSE_TAG.to_string()); lines.join("\n") } @@ -147,7 +93,6 @@ mod tests { use super::*; use core_test_support::test_path_buf; use pretty_assertions::assert_eq; - use std::collections::BTreeMap; fn fake_shell() -> Shell { Shell { @@ -160,7 +105,7 @@ mod tests { #[test] fn serialize_workspace_write_environment_context() { let cwd = test_path_buf("/repo"); - let context = EnvironmentContext::new(Some(cwd.clone()), fake_shell(), None); + let context = EnvironmentContext::new(Some(cwd.clone()), fake_shell()); let expected = format!( r#" @@ -175,7 +120,7 @@ mod tests { #[test] fn serialize_read_only_environment_context() { - let context = EnvironmentContext::new(None, fake_shell(), None); + let context = EnvironmentContext::new(None, fake_shell()); let expected = r#" bash @@ -186,7 +131,7 @@ mod tests { #[test] fn serialize_external_sandbox_environment_context() { - let context = EnvironmentContext::new(None, fake_shell(), None); + let context = EnvironmentContext::new(None, fake_shell()); let expected = r#" bash @@ -197,7 +142,7 @@ mod tests { #[test] fn serialize_external_sandbox_with_restricted_network_environment_context() { - let context = EnvironmentContext::new(None, fake_shell(), None); + let context = EnvironmentContext::new(None, fake_shell()); let expected = r#" bash @@ -208,7 +153,7 @@ mod tests { #[test] fn serialize_full_access_environment_context() { - let context = EnvironmentContext::new(None, fake_shell(), None); + let context = EnvironmentContext::new(None, fake_shell()); let expected = r#" bash @@ -219,23 +164,23 @@ mod tests { #[test] fn equals_except_shell_compares_cwd() { - let context1 = EnvironmentContext::new(Some(PathBuf::from("/repo")), fake_shell(), None); - let context2 = EnvironmentContext::new(Some(PathBuf::from("/repo")), fake_shell(), None); + let context1 = EnvironmentContext::new(Some(PathBuf::from("/repo")), fake_shell()); + let context2 = EnvironmentContext::new(Some(PathBuf::from("/repo")), fake_shell()); assert!(context1.equals_except_shell(&context2)); } #[test] fn equals_except_shell_ignores_sandbox_policy() { - let context1 = EnvironmentContext::new(Some(PathBuf::from("/repo")), fake_shell(), None); - let context2 = EnvironmentContext::new(Some(PathBuf::from("/repo")), fake_shell(), None); + let context1 = EnvironmentContext::new(Some(PathBuf::from("/repo")), fake_shell()); + let context2 = EnvironmentContext::new(Some(PathBuf::from("/repo")), fake_shell()); assert!(context1.equals_except_shell(&context2)); } #[test] fn equals_except_shell_compares_cwd_differences() { - let context1 = EnvironmentContext::new(Some(PathBuf::from("/repo1")), fake_shell(), None); - let context2 = EnvironmentContext::new(Some(PathBuf::from("/repo2")), fake_shell(), None); + let context1 = EnvironmentContext::new(Some(PathBuf::from("/repo1")), fake_shell()); + let context2 = EnvironmentContext::new(Some(PathBuf::from("/repo2")), fake_shell()); assert!(!context1.equals_except_shell(&context2)); } @@ -249,7 +194,6 @@ mod tests { shell_path: "/bin/bash".into(), shell_snapshot: crate::shell::empty_shell_snapshot_receiver(), }, - None, ); let context2 = EnvironmentContext::new( Some(PathBuf::from("/repo")), @@ -258,46 +202,8 @@ mod tests { shell_path: "/bin/zsh".into(), shell_snapshot: crate::shell::empty_shell_snapshot_receiver(), }, - None, ); assert!(context1.equals_except_shell(&context2)); } - - #[test] - fn serialize_environment_context_with_workspace_configuration() { - let cwd = test_path_buf("/repo"); - let cwd_str = cwd.to_string_lossy().to_string(); - let mut workspaces = BTreeMap::new(); - workspaces.insert( - cwd_str.clone(), - WorkspaceEntry { - associated_remote_urls: Some(BTreeMap::from([( - "origin".to_string(), - "https://example.com/repo.git".to_string(), - )])), - latest_git_commit_hash: Some("abc123".to_string()), - }, - ); - let workspace_configuration = WorkspaceConfiguration { workspaces }; - let context = - EnvironmentContext::new(Some(cwd), fake_shell(), Some(workspace_configuration)); - - let expected = format!( - r#" - {cwd_str} - bash - - - abc123 - - https://example.com/repo.git - - - -"# - ); - - assert_eq!(context.serialize_to_xml(), expected); - } } diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index ba47838e1f..fbb874d3fe 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -128,6 +128,7 @@ pub mod util; pub use apply_patch::CODEX_APPLY_PATCH_ARG1; pub use client::WEB_SEARCH_ELIGIBLE_HEADER; +pub use client::X_CODEX_TURN_METADATA_HEADER; pub use command_safety::is_dangerous_command; pub use command_safety::is_safe_command; pub use exec_policy::ExecPolicyError; diff --git a/codex-rs/core/tests/responses_headers.rs b/codex-rs/core/tests/responses_headers.rs index 86d6e8ce61..4df60c180f 100644 --- a/codex-rs/core/tests/responses_headers.rs +++ b/codex-rs/core/tests/responses_headers.rs @@ -12,6 +12,7 @@ use codex_core::ResponseItem; use codex_core::TransportManager; use codex_core::WEB_SEARCH_ELIGIBLE_HEADER; use codex_core::WireApi; +use codex_core::X_CODEX_TURN_METADATA_HEADER; use codex_core::models_manager::manager::ModelsManager; use codex_otel::OtelManager; use codex_protocol::ThreadId; @@ -23,6 +24,7 @@ use core_test_support::load_default_config_for_test; use core_test_support::responses; use core_test_support::test_codex::test_codex; use futures::StreamExt; +use pretty_assertions::assert_eq; use tempfile::TempDir; use wiremock::matchers::header; @@ -124,6 +126,99 @@ async fn responses_stream_includes_subagent_header_on_review() { ); } +#[tokio::test] +async fn responses_stream_includes_turn_metadata_header_when_set() { + core_test_support::skip_if_no_network!(); + + let server = responses::start_mock_server().await; + let response_body = responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_completed("resp-1"), + ]); + + let request_recorder = responses::mount_sse_once(&server, response_body).await; + + let provider = ModelProviderInfo { + name: "mock".into(), + base_url: Some(format!("{}/v1", server.uri())), + env_key: None, + env_key_instructions: None, + experimental_bearer_token: None, + wire_api: WireApi::Responses, + query_params: None, + http_headers: None, + env_http_headers: None, + request_max_retries: Some(0), + stream_max_retries: Some(0), + stream_idle_timeout_ms: Some(5_000), + requires_openai_auth: false, + }; + + let codex_home = TempDir::new().expect("failed to create TempDir"); + let mut config = load_default_config_for_test(&codex_home).await; + config.model_provider_id = provider.name.clone(); + config.model_provider = provider.clone(); + let effort = config.model_reasoning_effort; + let summary = config.model_reasoning_summary; + let model = ModelsManager::get_model_offline(config.model.as_deref()); + config.model = Some(model.clone()); + let config = Arc::new(config); + + let conversation_id = ThreadId::new(); + let auth_mode = AuthMode::ChatGPT; + let session_source = SessionSource::SubAgent(SubAgentSource::Review); + let model_info = ModelsManager::construct_model_info_offline(model.as_str(), &config); + let otel_manager = OtelManager::new( + conversation_id, + model.as_str(), + model_info.slug.as_str(), + None, + Some("test@test.com".to_string()), + Some(auth_mode), + false, + "test".to_string(), + session_source.clone(), + ); + + let turn_metadata = + r#"{"workspaces":{"/repo":{"latest_git_commit_hash":"abc123"}}}"#.to_string(); + let mut client_session = ModelClient::new( + Arc::clone(&config), + None, + model_info, + otel_manager, + provider, + effort, + summary, + conversation_id, + session_source, + ) + .new_session_with_turn_metadata(Some(turn_metadata.clone())); + + let mut prompt = Prompt::default(); + prompt.input = vec![ResponseItem::Message { + id: None, + role: "user".into(), + content: vec![ContentItem::InputText { + text: "hello".into(), + }], + end_turn: None, + }]; + + let mut stream = client_session.stream(&prompt).await.expect("stream failed"); + while let Some(event) = stream.next().await { + if matches!(event, Ok(ResponseEvent::Completed { .. })) { + break; + } + } + + let request = request_recorder.single_request(); + assert_eq!( + request.header(X_CODEX_TURN_METADATA_HEADER).as_deref(), + Some(turn_metadata.as_str()) + ); +} + #[tokio::test] async fn responses_stream_includes_subagent_header_on_other() { core_test_support::skip_if_no_network!();