From 6d8f311b00e823052622eac2b463919ad86a666d Mon Sep 17 00:00:00 2001 From: Alex Kira Date: Mon, 26 Oct 2026 19:59:58 +0000 Subject: [PATCH] Render untrusted additional context as named tool outputs (#57279) ## What changed When `features.untrusted_as_tool_output` is enabled, emit untrusted additional context as `function_call_output` items in the `agent` namespace before user input. Use the explicit context name when supplied; otherwise derive a bounded name from the source key, falling back to `untrusted_content`. Deduplicate these outputs by name and truncated value so name changes are emitted while changes confined to truncated content are ignored. Application context remains in developer messages, and disabling the feature preserves the existing message format. ## Testing Extend core and app-server tests and add request-history snapshots covering both formats, turn start and steering, derived names, truncation, and deduplication. GitOrigin-RevId: cd5e2f4b3121714acb7f4baf5607f3c27d408cae --- .../app-server/tests/suite/v2/turn_start.rs | 74 +++++-- .../app-server/tests/suite/v2/turn_steer.rs | 55 +++++- .../src/additional_context.rs | 49 +++++ codex-rs/context-fragments/src/lib.rs | 1 + codex-rs/core/src/context/mod.rs | 1 + codex-rs/core/src/session/turn_input.rs | 8 +- codex-rs/core/src/state/additional_context.rs | 77 ++++++-- codex-rs/core/src/state/mod.rs | 1 + .../core/tests/suite/additional_context.rs | 182 ++++++++++++++++-- codex-rs/core/tests/suite/scenarios.rs | 127 ++++++++++++ ...s__additional_context_legacy_messages.snap | 29 +++ ...rios__additional_context_tool_outputs.snap | 28 +++ 12 files changed, 589 insertions(+), 43 deletions(-) create mode 100644 codex-rs/core/tests/suite/snapshots/all__suite__scenarios__additional_context_legacy_messages.snap create mode 100644 codex-rs/core/tests/suite/snapshots/all__suite__scenarios__additional_context_tool_outputs.snap 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 748f0a6171..eae342d437 100644 --- a/codex-rs/app-server/tests/suite/v2/turn_start.rs +++ b/codex-rs/app-server/tests/suite/v2/turn_start.rs @@ -790,13 +790,17 @@ async fn turn_start_steers_active_turn_and_returns_active_turn_id() -> Result<() Ok(()) } -#[test_case::test_case(Some("browser"); "named_context")] -#[test_case::test_case(Some("a"); "minimum_length")] -#[test_case::test_case(Some("browser_2"); "digits_and_underscore")] -#[test_case::test_case(Some(&"a".repeat(64)); "maximum_length")] -#[test_case::test_case(None; "legacy_caller")] +#[test_case::test_case(false, Some("browser"); "disabled")] +#[test_case::test_case(true, Some("browser"); "named_output")] +#[test_case::test_case(true, Some("a"); "minimum_length")] +#[test_case::test_case(true, Some("browser_2"); "digits_and_underscore")] +#[test_case::test_case(true, Some(&"a".repeat(64)); "maximum_length")] +#[test_case::test_case(true, None; "legacy_caller")] #[tokio::test] -async fn turn_start_additional_context_flows_to_model_input(name: Option<&str>) -> Result<()> { +async fn turn_start_additional_context_flows_to_model_input( + tool_outputs: bool, + name: Option<&str>, +) -> Result<()> { let responses = vec![create_final_assistant_message_sse_response("Done")?]; let server = create_mock_responses_server_sequence_unchecked(responses).await; @@ -811,6 +815,10 @@ async fn turn_start_additional_context_flows_to_model_input(name: Option<&str>) let ThreadStartResponse { thread, .. } = mcp .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), + config: Some(HashMap::from([( + "features.untrusted_as_tool_output".to_owned(), + json!(tool_outputs), + )])), ..Default::default() }) .await?; @@ -863,6 +871,9 @@ async fn turn_start_additional_context_flows_to_model_input(name: Option<&str>) if name.is_none() { entry.remove("name"); } + // Unknown wire fields cannot override the harness-owned namespace. + entry.insert("namespace".to_owned(), json!("other")); + let mut malformed_params = params.clone(); malformed_params["additionalContext"]["custom_source"]["name"] = json!("agent.browser"); let rejected_request = mcp @@ -882,6 +893,21 @@ async fn turn_start_additional_context_flows_to_model_input(name: Option<&str>) } ); + let capped_name = "A".repeat(128); + let long_key = format!("!{capped_name}z"); + let unnamed_sources: &[(&str, &str)] = if tool_outputs && name.is_none() { + &[ + ("Browser Tab/\té-_2", "Browser_Tab_2"), + ("Browser_Tab_2", "Browser_Tab_2"), + (&long_key, &capped_name), + ("!?", "untrusted_content"), + ] + } else { + &[] + }; + for &(key, _) in unnamed_sources { + params["additionalContext"][key] = json!({"kind": "untrusted", "value": key}); + } let request_id = mcp.send_request("turn/start", Some(params)).await?; let _: TurnStartResponse = mcp.read_response(request_id).await?; timeout( @@ -910,11 +936,37 @@ async fn turn_start_additional_context_flows_to_model_input(name: Option<&str>) .collect::>(); assert!(!texts.contains(&"rejected input")); assert!(!texts.contains(&"rejected context")); - assert!( - body.to_string() - .contains("source value") - ); - assert!(!body.to_string().contains("function_call_output")); + if tool_outputs { + let outputs = input + .iter() + .filter(|item| item["type"] == "function_call_output") + .map(|item| { + ( + item["name"].clone(), + item["namespace"].clone(), + item["output"].clone(), + item.get("call_id").cloned(), + ) + }) + .collect::>(); + assert_eq!(outputs.len(), 1 + unnamed_sources.len()); + for (value, expected_name) in + std::iter::once(("source value", name.unwrap_or("custom_source"))) + .chain(unnamed_sources.iter().copied()) + { + assert!( + outputs.contains(&(json!(expected_name), json!("agent"), json!(value), None)), + "missing agent.{expected_name} output for {value:?}: {outputs:?}" + ); + } + assert!(!body.to_string().contains("")); + } else { + assert!( + body.to_string() + .contains("source value") + ); + assert!(!body.to_string().contains("function_call_output")); + } Ok(()) } diff --git a/codex-rs/app-server/tests/suite/v2/turn_steer.rs b/codex-rs/app-server/tests/suite/v2/turn_steer.rs index c6a677ec97..38f5cace78 100644 --- a/codex-rs/app-server/tests/suite/v2/turn_steer.rs +++ b/codex-rs/app-server/tests/suite/v2/turn_steer.rs @@ -252,8 +252,10 @@ async fn turn_steer_rejects_oversized_text_input() -> Result<()> { Ok(()) } +#[test_case::test_case(false; "legacy_context")] +#[test_case::test_case(true; "named_context")] #[tokio::test] -async fn turn_steer_returns_active_turn_id() -> Result<()> { +async fn turn_steer_returns_active_turn_id(tool_outputs: bool) -> Result<()> { // TODO(anp): Remove after the active-turn fixture can run in the selected remote environment. skip_if_remote!( Ok(()), @@ -301,6 +303,10 @@ async fn turn_steer_returns_active_turn_id() -> Result<()> { let ThreadStartResponse { thread, .. } = mcp .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), + config: Some(HashMap::from([( + "features.untrusted_as_tool_output".to_owned(), + serde_json::json!(tool_outputs), + )])), ..Default::default() }) .await?; @@ -339,7 +345,14 @@ async fn turn_steer_returns_active_turn_id() -> Result<()> { }], responsesapi_client_metadata: None, responsesapi_client_metadata_patch: None, - additional_context: None, + additional_context: Some(HashMap::from([( + "browser_info".to_owned(), + AdditionalContextEntry { + value: "tab one".to_owned(), + kind: AdditionalContextKind::Untrusted, + name: Some("browser".to_owned()), + }, + )])), expected_turn_id: turn.id.clone(), }, }) @@ -393,6 +406,44 @@ async fn turn_steer_returns_active_turn_id() -> Result<()> { ) .await??; + let requests = server + .received_requests() + .await + .context("recorded requests")?; + let body = requests + .iter() + .rfind(|request| request.url.path().ends_with("/responses")) + .context("model request after steering")? + .body_json::()?; + let input = body["input"].as_array().context("input array")?; + let context_index = if tool_outputs { + input.iter().position(|item| { + item["type"] == "function_call_output" + && item["namespace"] == "agent" + && item["name"] == "browser" + && item["output"] == "tab one" + }) + } else { + input.iter().position(|item| { + item["content"].as_array().is_some_and(|content| { + content.iter().any(|part| { + part["text"] == "tab one" + }) + }) + }) + } + .context("steered context")?; + let user_index = input + .iter() + .position(|item| { + item["role"] == "user" + && item["content"] + .as_array() + .is_some_and(|content| content.iter().any(|part| part["text"] == "steer")) + }) + .context("steered user input")?; + assert!(context_index < user_index); + Ok(()) } diff --git a/codex-rs/context-fragments/src/additional_context.rs b/codex-rs/context-fragments/src/additional_context.rs index bf5f3317f9..e409bc9a56 100644 --- a/codex-rs/context-fragments/src/additional_context.rs +++ b/codex-rs/context-fragments/src/additional_context.rs @@ -1,14 +1,49 @@ use codex_protocol::models::ContentItem; use codex_protocol::models::ContentItemKind; +use codex_protocol::models::FunctionCallOutputPayload; +use codex_protocol::models::ResponseItem; use codex_protocol::models::input_text; use codex_utils_string::truncate_middle_with_token_budget; use crate::ContextualUserFragment; const MAX_ADDITIONAL_CONTEXT_VALUE_TOKENS: usize = 1_000; +const MAX_DERIVED_TOOL_NAME_CHARS: usize = 128; +const ADDITIONAL_CONTEXT_TOOL_NAMESPACE: &str = "agent"; const ADDITIONAL_CONTEXT_END_MARKER_SUFFIX: &str = ">"; const ADDITIONAL_CONTEXT_START_MARKER_PREFIX: &str = ", value: &str) -> Self { + let name = name.unwrap_or_else(|| tool_name_from_key(key)); + Self { + name, + value: truncate_middle_with_token_budget(value, MAX_ADDITIONAL_CONTEXT_VALUE_TOKENS).0, + } + } +} + +impl From for ResponseItem { + fn from(output: AdditionalContextToolOutput) -> Self { + Self::FunctionCallOutput { + id: None, + call_id: None, + name: Some(output.name), + namespace: Some(ADDITIONAL_CONTEXT_TOOL_NAMESPACE.to_owned()), + output: FunctionCallOutputPayload::from_text(output.value), + internal_chat_message_metadata_passthrough: None, + } + } +} + #[derive(Debug, Clone, PartialEq, Eq)] pub struct AdditionalContextUserFragment { key: String, @@ -93,6 +128,20 @@ impl ContextualUserFragment for AdditionalContextDeveloperFragment { } } +fn tool_name_from_key(key: &str) -> String { + let derived: String = key + .chars() + .map(|ch| if ch == ' ' { '_' } else { ch }) + .filter(|ch| ch.is_ascii_alphanumeric() || *ch == '_') + .take(MAX_DERIVED_TOOL_NAME_CHARS) + .collect(); + if derived.is_empty() { + "untrusted_content".to_owned() + } else { + derived + } +} + fn additional_context_body(key: &str, value: &str) -> String { let value = truncate_middle_with_token_budget(value, MAX_ADDITIONAL_CONTEXT_VALUE_TOKENS).0; format!("{key}>{value}, ) -> Vec { + let format = if session.enabled(Feature::UntrustedAsToolOutput) { + AdditionalContextFormat::ToolOutputs + } else { + AdditionalContextFormat::Messages + }; let additional_context_input = { let mut state = session.state.lock().await; - state.additional_context.merge(additional_context) + state.additional_context.merge(additional_context, format) }; additional_context_input .into_iter() diff --git a/codex-rs/core/src/state/additional_context.rs b/codex-rs/core/src/state/additional_context.rs index 2ab78d54da..44a15cded6 100644 --- a/codex-rs/core/src/state/additional_context.rs +++ b/codex-rs/core/src/state/additional_context.rs @@ -1,36 +1,87 @@ +//! Deduplicates the active client context using its selected model-visible representation. + use std::collections::BTreeMap; use crate::context::AdditionalContextDeveloperFragment; +use crate::context::AdditionalContextToolOutput; use crate::context::AdditionalContextUserFragment; use crate::context::ContextualUserFragment; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::AdditionalContextEntry; use codex_protocol::protocol::AdditionalContextKind; +pub(crate) enum AdditionalContextFormat { + Messages, + ToolOutputs, +} + +#[derive(Debug, Clone, Eq)] +enum StoredAdditionalContext { + Message(AdditionalContextEntry), + ToolOutput(AdditionalContextToolOutput), +} + +impl PartialEq for StoredAdditionalContext { + fn eq(&self, other: &Self) -> bool { + match (self, other) { + (Self::Message(left), Self::Message(right)) => { + left.kind == right.kind && left.value == right.value + } + (Self::ToolOutput(left), Self::ToolOutput(right)) => left == right, + (Self::Message(_), Self::ToolOutput(_)) | (Self::ToolOutput(_), Self::Message(_)) => { + false + } + } + } +} + #[derive(Debug, Clone, Default, PartialEq, Eq)] pub(crate) struct AdditionalContextStore { - values: BTreeMap, + values: BTreeMap, } impl AdditionalContextStore { pub(crate) fn merge( &mut self, values: BTreeMap, + format: AdditionalContextFormat, ) -> Vec { + let values = + values + .into_iter() + .map(|(key, entry)| { + let value = match (&format, entry.kind) { + ( + AdditionalContextFormat::ToolOutputs, + AdditionalContextKind::Untrusted, + ) => StoredAdditionalContext::ToolOutput(AdditionalContextToolOutput::new( + &key, + entry.name, + &entry.value, + )), + (AdditionalContextFormat::Messages, AdditionalContextKind::Untrusted) + | (AdditionalContextFormat::Messages, AdditionalContextKind::Application) + | ( + AdditionalContextFormat::ToolOutputs, + AdditionalContextKind::Application, + ) => StoredAdditionalContext::Message(entry), + }; + (key, value) + }) + .collect::>(); let fragments = values .iter() - .filter(|(key, entry)| { - self.values.get(*key).is_none_or(|previous| { - previous.kind != entry.kind || previous.value != entry.value - }) - }) - .map(|(key, entry)| match entry.kind { - AdditionalContextKind::Untrusted => ContextualUserFragment::into( - AdditionalContextUserFragment::new(key.clone(), entry.value.clone()), - ), - AdditionalContextKind::Application => ContextualUserFragment::into( - AdditionalContextDeveloperFragment::new(key.clone(), entry.value.clone()), - ), + .filter(|(key, value)| self.values.get(*key) != Some(*value)) + .map(|(key, entry)| match entry { + StoredAdditionalContext::ToolOutput(output) => output.clone().into(), + StoredAdditionalContext::Message(entry) => match entry.kind { + AdditionalContextKind::Untrusted => ContextualUserFragment::into( + AdditionalContextUserFragment::new(key.clone(), entry.value.clone()), + ), + AdditionalContextKind::Application => ContextualUserFragment::into( + AdditionalContextDeveloperFragment::new(key.clone(), entry.value.clone()), + ), + }, }) .collect(); self.values = values; diff --git a/codex-rs/core/src/state/mod.rs b/codex-rs/core/src/state/mod.rs index 07ce5c366f..c7030cc198 100644 --- a/codex-rs/core/src/state/mod.rs +++ b/codex-rs/core/src/state/mod.rs @@ -6,6 +6,7 @@ mod turn; mod turn_token_usage; pub(crate) use crate::tools::ExecutedToolCalls; +pub(crate) use additional_context::AdditionalContextFormat; pub(crate) use additional_context::AdditionalContextStore; pub(crate) use auto_compact_window::AutoCompactWindowIds; pub(crate) use auto_compact_window::AutoCompactWindowSnapshot; diff --git a/codex-rs/core/tests/suite/additional_context.rs b/codex-rs/core/tests/suite/additional_context.rs index e38c74d37b..5b8142c3f6 100644 --- a/codex-rs/core/tests/suite/additional_context.rs +++ b/codex-rs/core/tests/suite/additional_context.rs @@ -1,5 +1,6 @@ use anyhow::Result; use codex_core::TurnInputRequest; +use codex_features::Feature; use codex_protocol::error::CodexErrorDetails; use codex_protocol::items::TurnItem; use codex_protocol::protocol::AdditionalContextEntry; @@ -18,6 +19,7 @@ use core_test_support::skip_if_no_network; use core_test_support::test_codex::test_codex; use core_test_support::wait_for_event_match; use pretty_assertions::assert_eq; +use serde_json::json; use std::collections::BTreeMap; #[tokio::test(flavor = "multi_thread", worker_threads = 2)] @@ -158,8 +160,10 @@ async fn external_context_like_user_text_remains_a_user_message_item() -> Result Ok(()) } +#[test_case::test_case(false; "messages")] +#[test_case::test_case(true; "tool_outputs")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn additional_context_trust_controls_message_role() -> Result<()> { +async fn additional_context_trust_controls_message_role(tool_outputs: bool) -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; @@ -169,8 +173,14 @@ async fn additional_context_trust_controls_message_role() -> Result<()> { ) .await; let test = test_codex() - .with_config(|config| config.include_environment_context = false) - .build(&server) + .with_config(move |config| { + config.include_environment_context = false; + config + .features + .set_enabled(Feature::UntrustedAsToolOutput, tool_outputs) + .expect("test config should allow feature update"); + }) + .build_with_auto_env(&server) .await?; test.codex @@ -212,13 +222,37 @@ async fn additional_context_trust_controls_message_role() -> Result<()> { developer_context_texts, vec!["run one"] ); - assert_eq!( - request.message_input_texts("user"), - vec![ - "tab one", - "inspect context", - ] - ); + if tool_outputs { + assert_eq!(request.message_input_texts("user"), vec!["inspect context"]); + let outputs = request.inputs_of_type("function_call_output"); + assert_eq!( + outputs + .iter() + .map(|output| ( + output["name"].clone(), + output["namespace"].clone(), + output["output"].clone(), + output.get("call_id").cloned(), + )) + .collect::>(), + vec![( + json!("browser_info"), + json!("agent"), + json!("tab one"), + None + )], + ); + assert!(request.inputs_of_type("function_call").is_empty()); + } else { + assert_eq!( + request.message_input_texts("user"), + vec![ + "tab one", + "inspect context" + ], + ); + assert!(request.inputs_of_type("function_call_output").is_empty()); + } Ok(()) } @@ -307,8 +341,12 @@ async fn additional_context_is_deduplicated_between_turns_while_retained( Ok(()) } +#[test_case::test_case(false; "messages")] +#[test_case::test_case(true; "tool_outputs")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn additional_context_removes_one_value_while_adding_another() -> Result<()> { +async fn additional_context_removes_one_value_while_adding_another( + tool_outputs: bool, +) -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; @@ -328,8 +366,14 @@ async fn additional_context_removes_one_value_while_adding_another() -> Result<( ) .await; let test = test_codex() - .with_config(|config| config.include_environment_context = false) - .build(&server) + .with_config(move |config| { + config.include_environment_context = false; + config + .features + .set_enabled(Feature::UntrustedAsToolOutput, tool_outputs) + .expect("test config should allow feature update"); + }) + .build_with_auto_env(&server) .await?; test.codex @@ -426,6 +470,37 @@ async fn additional_context_removes_one_value_while_adding_another() -> Result<( }) .await; + if tool_outputs { + for (request, expected_outputs, expected_users) in [ + ( + first_request.single_request(), + vec!["run one", "tab one"], + vec!["first turn"], + ), + ( + second_request.single_request(), + vec!["run one", "tab one", "pty one"], + vec!["first turn", "second turn"], + ), + ( + third_request.single_request(), + vec!["run one", "tab one", "pty one", "tab one"], + vec!["first turn", "second turn", "third turn"], + ), + ] { + let outputs = request.inputs_of_type("function_call_output"); + assert_eq!( + outputs + .iter() + .map(|output| output["output"].as_str().expect("text output")) + .collect::>(), + expected_outputs, + ); + assert_eq!(request.message_input_texts("user"), expected_users); + } + return Ok(()); + } + assert_eq!( first_request.single_request().message_input_texts("user"), vec![ @@ -460,8 +535,12 @@ async fn additional_context_removes_one_value_while_adding_another() -> Result<( Ok(()) } +#[test_case::test_case(false; "messages")] +#[test_case::test_case(true; "tool_outputs")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn additional_context_values_are_truncated_before_model_input() -> Result<()> { +async fn additional_context_values_are_truncated_before_model_input( + tool_outputs: bool, +) -> Result<()> { skip_if_no_network!(Ok(())); const MAX_EXPECTED_EXTERNAL_CONTEXT_TEXT_BYTES: usize = 5 * 1024; @@ -473,8 +552,14 @@ async fn additional_context_values_are_truncated_before_model_input() -> Result< ) .await; let test = test_codex() - .with_config(|config| config.include_environment_context = false) - .build(&server) + .with_config(move |config| { + config.include_environment_context = false; + config + .features + .set_enabled(Feature::UntrustedAsToolOutput, tool_outputs) + .expect("test config should allow feature update"); + }) + .build_with_auto_env(&server) .await?; let long_browser_value = format!("browser-head-{}browser-tail", "b".repeat(40_000)); let long_automation_value = format!("automation-head-{}automation-tail", "a".repeat(40_000)); @@ -534,6 +619,71 @@ async fn additional_context_values_are_truncated_before_model_input() -> Result< automation_text.len() ); + if tool_outputs { + assert_eq!( + request.message_input_texts("user"), + vec!["summarize context"] + ); + let outputs = request.inputs_of_type("function_call_output"); + assert_eq!(outputs.len(), 1); + assert_eq!(outputs[0]["name"], "browser_info"); + assert_eq!(outputs[0]["namespace"], "agent"); + let body = outputs[0]["output"].as_str().expect("text output"); + assert!(body.starts_with(&format!("browser-head-{}", "b".repeat(1024)))); + assert!(body.ends_with("browser-tail")); + assert!(body.contains("tokens truncated")); + assert!(body.len() <= MAX_EXPECTED_EXTERNAL_CONTEXT_TEXT_BYTES); + let mut changed_middle = long_browser_value.clone(); + changed_middle.replace_range(20_000..20_001, "c"); + let changed_prefix = changed_middle.replacen("browser-head", "browser-view", 1); + for (value, expected_count) in [(changed_middle, 1), (changed_prefix, 2)] { + let next = mount_sse_once(&server, sse(vec![ev_completed("updated-context")])).await; + test.codex + .start_or_steer_turn( + TurnInputRequest::user_input(vec![UserInput::Text { + text: "check updated context".to_owned(), + text_elements: Vec::new(), + }]) + .with_additional_context(BTreeMap::from([ + ( + "automation_info".to_owned(), + AdditionalContextEntry { + name: None, + value: long_automation_value.clone(), + kind: AdditionalContextKind::Application, + }, + ), + ( + "browser_info".to_owned(), + AdditionalContextEntry { + name: None, + value, + kind: AdditionalContextKind::Untrusted, + }, + ), + ])), + ) + .await?; + wait_for_event_match(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)).then_some(()) + }) + .await; + let next = next.single_request(); + let outputs = next.inputs_of_type("function_call_output"); + assert_eq!(outputs.len(), expected_count); + assert_eq!(outputs[0]["output"], body); + if expected_count == 2 { + assert!( + outputs[1]["output"] + .as_str() + .expect("text output") + .starts_with("browser-view") + ); + } + } + return Ok(()); + } + let user_texts = request.message_input_texts("user"); let [external_text, user_text] = user_texts.as_slice() else { panic!("expected external context plus user input, got {user_texts:?}"); diff --git a/codex-rs/core/tests/suite/scenarios.rs b/codex-rs/core/tests/suite/scenarios.rs index deccca2eba..ad8ff10d09 100644 --- a/codex-rs/core/tests/suite/scenarios.rs +++ b/codex-rs/core/tests/suite/scenarios.rs @@ -50,6 +50,8 @@ use codex_protocol::openai_models::CodeModeToolMessages; use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::openai_models::ToolMessage; use codex_protocol::openai_models::ToolMode; +use codex_protocol::protocol::AdditionalContextEntry; +use codex_protocol::protocol::AdditionalContextKind; use codex_protocol::protocol::EnvironmentConfigState; use codex_protocol::protocol::EventMsg; #[cfg(not(target_os = "windows"))] @@ -317,6 +319,131 @@ fn configure_scenario_catalog(config: &mut Config) { config.cloud_skill_enabled = false; } +#[test_case::test_case(false; "messages")] +#[test_case::test_case(true; "tool_outputs")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn additional_context_tool_outputs(tool_outputs: bool) -> Result<()> { + skip_if_no_network!(Ok(())); + let server = start_mock_server().await; + let test = test_codex() + .with_config(move |config| { + configure_scenario_catalog(config); + config.include_environment_context = false; + config + .features + .set_enabled(Feature::UntrustedAsToolOutput, tool_outputs) + .expect("test config should allow feature update"); + }) + .build_with_auto_env(&server) + .await?; + let mock = mount_sse_sequence( + &server, + (0..4) + .map(|index| sse(vec![ev_completed(&format!("response-{index}"))])) + .collect(), + ) + .await; + + for (prompt, name, value) in [ + ( + "Summarize the selected document.", + Some("browser"), + "The launch is on Monday.", + ), + ( + "The same document is now open in the editor.", + Some("editor"), + "The launch is on Monday.", + ), + ( + "Check the date again.", + Some("editor"), + "The launch is on Monday.", + ), + ( + "Summarize the updated document.", + None, + "The launch is on Tuesday.", + ), + ] { + test.codex + .start_or_steer_turn( + TurnInputRequest::user_input(vec![text(prompt)]).with_additional_context( + std::collections::BTreeMap::from([ + ( + "selection".to_owned(), + AdditionalContextEntry { + value: value.to_owned(), + kind: AdditionalContextKind::Untrusted, + name: name.map(str::to_owned), + }, + ), + ( + "application".to_owned(), + AdditionalContextEntry::new( + "Keep the answer brief.".to_owned(), + AdditionalContextKind::Application, + ), + ), + ]), + ), + ) + .await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + } + let requests = mock.requests(); + assert_eq!(requests.len(), 4); + if tool_outputs { + assert_eq!( + requests + .iter() + .map(|request| request.inputs_of_type("function_call_output").len()) + .collect::>(), + vec![1, 2, 2, 3] + ); + assert_eq!(requests[3].message_input_texts("user").len(), 4); + let input = requests[0].input(); + let context_index = input + .iter() + .position(|item| item["type"] == "function_call_output") + .expect("context output"); + let user_index = input + .iter() + .position(|item| item["role"] == "user") + .expect("user message"); + assert!(context_index < user_index); + } else { + assert!( + requests + .iter() + .all(|request| request.inputs_of_type("function_call_output").is_empty()) + ); + assert_eq!( + requests + .iter() + .map(|request| request.message_input_texts("user").len()) + .collect::>(), + vec![2, 3, 4, 6] + ); + } + insta::assert_snapshot!( + if tool_outputs { + "additional_context_tool_outputs" + } else { + "additional_context_legacy_messages" + }, + context_snapshot::format_request_history_snapshot( + "Named context changes are appended before the user; unchanged values are deduplicated.", + &requests, + &ContextSnapshotOptions::default().rewrite_known_segments(), + ) + ); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn realtime_attachment_replacement_preserves_handoff_context() -> Result<()> { let requests = super::realtime_attachment::attachment_replacement_scenario().await?; diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__additional_context_legacy_messages.snap b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__additional_context_legacy_messages.snap new file mode 100644 index 0000000000..112e945915 --- /dev/null +++ b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__additional_context_legacy_messages.snap @@ -0,0 +1,29 @@ +--- +source: core/tests/suite/scenarios.rs +expression: "context_snapshot::format_request_history_snapshot(\"Named context changes are appended before the user; unchanged values are deduplicated.\",\n&requests, &ContextSnapshotOptions::default().rewrite_known_segments(),)" +--- +Scenario: Named context changes are appended before the user; unchanged values are deduplicated. + +## Window 1 +-- request 1 (turn) -- +00:message/developer: + +01:message/developer: + +02:message/developer: + Keep the answer brief. +03:message/user: + The launch is on Monday. +04:message/user: + Summarize the selected document. +-- request 2 (turn) -- +05:message/user: + The same document is now open in the editor. +-- request 3 (turn) -- +06:message/user: + Check the date again. +-- request 4 (turn) -- +07:message/user: + The launch is on Tuesday. +08:message/user: + Summarize the updated document. diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__additional_context_tool_outputs.snap b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__additional_context_tool_outputs.snap new file mode 100644 index 0000000000..3f50e94cf4 --- /dev/null +++ b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__additional_context_tool_outputs.snap @@ -0,0 +1,28 @@ +--- +source: core/tests/suite/scenarios.rs +expression: "context_snapshot::format_request_history_snapshot(\"Named context changes are appended before the user; unchanged values are deduplicated.\",\n&requests, &ContextSnapshotOptions::default().rewrite_known_segments(),)" +--- +Scenario: Named context changes are appended before the user; unchanged values are deduplicated. + +## Window 1 +-- request 1 (turn) -- +00:message/developer: + +01:message/developer: + +02:message/developer: + Keep the answer brief. +03:function_call_output/agent.browser:The launch is on Monday. +04:message/user: + Summarize the selected document. +-- request 2 (turn) -- +05:function_call_output/agent.editor:The launch is on Monday. +06:message/user: + The same document is now open in the editor. +-- request 3 (turn) -- +07:message/user: + Check the date again. +-- request 4 (turn) -- +08:function_call_output/agent.selection:The launch is on Tuesday. +09:message/user: + Summarize the updated document.