openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-15 12:08:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

moving to headers

Browse Source
This commit is contained in:
pash
2026-01-29 18:04:43 -08:00
parent c479bff448
commit b267af326d
5 changed files with 183 additions and 149 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
codex-rs/core/src/client.rs
View File

@@ -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<OnceLock<String>>,
/// Turn-scoped metadata attached to every request in the turn.
turn_metadata_header: Option<HeaderValue>,
}
#[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<String>,
) -> 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<OnceLock<String>>>,
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
}

89
codex-rs/core/src/codex.rs
View File

@@ -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<ReadinessFlag>,
pub(crate) truncation_policy: TruncationPolicy,
pub(crate) dynamic_tools: Vec<DynamicToolSpec>,
/// Per-turn metadata serialized as a header value for outbound model requests.
pub(crate) turn_metadata_header: Option<String>,
}
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<String> {
#[derive(Serialize)]
struct TurnMetadataWorkspace {
#[serde(skip_serializing_if = "Option::is_none")]
associated_remote_urls: Option<BTreeMap<String, String>>,
#[serde(skip_serializing_if = "Option::is_none")]
latest_git_commit_hash: Option<String>,
}
#[derive(Serialize)]
struct TurnMetadata {
workspaces: BTreeMap<String, TurnMetadataWorkspace>,
}
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<ResponseItem> {
fn build_workspace_configuration(
cwd: &Path,
git_info: Option<&codex_protocol::protocol::GitInfo>,
remote_urls: Option<&BTreeMap<String, String>>,
) -> Option<WorkspaceConfiguration> {
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::<ResponseItem>::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;

124
codex-rs/core/src/environment_context.rs
View File

@@ -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<PathBuf>,
pub shell: Shell,
/// Workspace metadata, structured to support future multi-root workspaces.
#[serde(skip_serializing_if = "Option::is_none")]
pub workspace_configuration: Option<WorkspaceConfiguration>,
}
impl EnvironmentContext {
pub fn new(
cwd: Option<PathBuf>,
shell: Shell,
workspace_configuration: Option<WorkspaceConfiguration>,
) -> Self {
Self {
cwd,
shell,
workspace_configuration,
}
pub fn new(cwd: Option<PathBuf>, 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<String, WorkspaceEntry>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub(crate) struct WorkspaceEntry {
#[serde(skip_serializing_if = "Option::is_none")]
pub associated_remote_urls: Option<BTreeMap<String, String>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub latest_git_commit_hash: Option<String>,
}
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>{shell_name}</shell>"));
if let Some(workspace_configuration) = self.workspace_configuration {
lines.push(" <workspace_configuration>".to_string());
for (path, workspace) in workspace_configuration.workspaces {
lines.push(format!(" <workspace path=\"{path}\">"));
if let Some(latest_git_commit_hash) = workspace.latest_git_commit_hash {
lines.push(format!(
" <latest_git_commit_hash>{latest_git_commit_hash}</latest_git_commit_hash>"
));
}
if let Some(associated_remote_urls) = workspace.associated_remote_urls
&& !associated_remote_urls.is_empty()
{
lines.push(" <associated_remote_urls>".to_string());
for (name, url) in associated_remote_urls {
lines.push(format!(" <remote name=\"{name}\">{url}</remote>"));
}
lines.push(" </associated_remote_urls>".to_string());
}
lines.push(" </workspace>".to_string());
}
lines.push(" </workspace_configuration>".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#"<environment_context>
@@ -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#"<environment_context>
<shell>bash</shell>
@@ -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#"<environment_context>
<shell>bash</shell>
@@ -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#"<environment_context>
<shell>bash</shell>
@@ -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#"<environment_context>
<shell>bash</shell>
@@ -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#"<environment_context>
<cwd>{cwd_str}</cwd>
<shell>bash</shell>
<workspace_configuration>
<workspace path="{cwd_str}">
<latest_git_commit_hash>abc123</latest_git_commit_hash>
<associated_remote_urls>
<remote name="origin">https://example.com/repo.git</remote>
</associated_remote_urls>
</workspace>
</workspace_configuration>
</environment_context>"#
);
assert_eq!(context.serialize_to_xml(), expected);
}
}

1
codex-rs/core/src/lib.rs
View File

@@ -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;

95
codex-rs/core/tests/responses_headers.rs
View File

@@ -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!();