From c5bd131567b39d0b70238b1b4290555e5978ddc2 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Wed, 20 May 2026 20:14:41 +0200 Subject: [PATCH] feat: add turn_id and truncation_policy to extension tool calls (#23666) ## Why Extension-owned tools currently receive a stripped `ToolCall` with only `call_id`, `tool_name`, and `payload`. That makes extension work that needs turn-local execution context awkward, especially web-search extension work that needs the active `truncation_policy` at tool invocation time. Reconstructing that value from config or `ExtensionData` would be indirect and could drift from the actual turn context, so the cleaner fix is to pass the needed turn metadata directly on the extension-facing invocation type. ## What changed - added `turn_id` and `truncation_policy` to `codex_tools::ToolCall` - populated those fields when core adapts `ToolInvocation` into an extension tool call - added a focused adapter test that verifies extension executors receive the forwarded turn metadata - updated the memories extension tests to construct the richer `ToolCall` - added the `codex-utils-output-truncation` dependency to `codex-tools` and refreshed lockfiles ## Testing - `cargo test -p codex-tools` - `cargo test -p codex-memories-extension` - `cargo test -p codex-core passes_turn_fields_to_extension_call` - `just bazel-lock-update` - `just bazel-lock-check` --- codex-rs/Cargo.lock | 1 + .../src/tools/handlers/extension_tools.rs | 66 +++++++++++++++++++ codex-rs/ext/memories/src/tests.rs | 9 +++ codex-rs/tools/Cargo.toml | 1 + codex-rs/tools/src/tool_call.rs | 3 + 5 files changed, 80 insertions(+) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 7ea7dce4d..72d395c27 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3765,6 +3765,7 @@ dependencies = [ "codex-features", "codex-protocol", "codex-utils-absolute-path", + "codex-utils-output-truncation", "codex-utils-pty", "codex-utils-string", "pretty_assertions", diff --git a/codex-rs/core/src/tools/handlers/extension_tools.rs b/codex-rs/core/src/tools/handlers/extension_tools.rs index 37ad2d65f..dd2ac42ea 100644 --- a/codex-rs/core/src/tools/handlers/extension_tools.rs +++ b/codex-rs/core/src/tools/handlers/extension_tools.rs @@ -88,8 +88,10 @@ impl CoreToolRuntime for ExtensionToolAdapter { fn to_extension_call(invocation: &ToolInvocation) -> ExtensionToolCall { ExtensionToolCall { + turn_id: invocation.turn.sub_id.clone(), call_id: invocation.call_id.clone(), tool_name: invocation.tool_name.clone(), + truncation_policy: invocation.turn.truncation_policy, payload: invocation.payload.clone(), } } @@ -108,6 +110,7 @@ mod tests { use pretty_assertions::assert_eq; use serde_json::json; + use tokio::sync::Mutex; use super::ExtensionToolAdapter; use crate::tools::context::ToolCallSource; @@ -158,6 +161,27 @@ mod tests { } } + struct CapturingExtensionExecutor { + captured_call: Arc>>, + } + + #[async_trait::async_trait] + impl codex_extension_api::ToolExecutor for CapturingExtensionExecutor { + fn tool_name(&self) -> codex_tools::ToolName { + codex_tools::ToolName::plain("extension_echo") + } + + async fn handle( + &self, + call: codex_tools::ToolCall, + ) -> Result, codex_tools::FunctionCallError> { + *self.captured_call.lock().await = Some(call); + Ok(Box::new(codex_tools::JsonToolOutput::new( + json!({ "ok": true }), + ))) + } + } + #[tokio::test] async fn exposes_generic_hook_payloads() { let handler = ExtensionToolAdapter::new(Arc::new(StubExtensionExecutor)); @@ -193,4 +217,46 @@ mod tests { }) ); } + + #[tokio::test] + async fn passes_turn_fields_to_extension_call() { + let captured_call = Arc::new(Mutex::new(None)); + let handler = ExtensionToolAdapter::new(Arc::new(CapturingExtensionExecutor { + captured_call: Arc::clone(&captured_call), + })); + let (session, turn) = crate::session::tests::make_session_and_context().await; + let turn_id = turn.sub_id.clone(); + let truncation_policy = turn.truncation_policy; + let invocation = ToolInvocation { + session: session.into(), + turn: turn.into(), + cancellation_token: tokio_util::sync::CancellationToken::new(), + tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())), + call_id: "call-extension".to_string(), + tool_name: codex_tools::ToolName::plain("extension_echo"), + source: ToolCallSource::Direct, + payload: ToolPayload::Function { + arguments: json!({ "message": "hello" }).to_string(), + }, + }; + + crate::tools::registry::ToolExecutor::handle(&handler, invocation) + .await + .expect("extension call should succeed"); + + let captured_call = captured_call.lock().await.clone().expect("captured call"); + assert_eq!(captured_call.turn_id, turn_id); + assert_eq!(captured_call.call_id, "call-extension"); + assert_eq!( + captured_call.tool_name, + codex_tools::ToolName::plain("extension_echo") + ); + assert_eq!(captured_call.truncation_policy, truncation_policy); + match captured_call.payload { + ToolPayload::Function { arguments } => { + assert_eq!(arguments, json!({ "message": "hello" }).to_string()); + } + payload => panic!("expected function payload, got {payload:?}"), + } + } } diff --git a/codex-rs/ext/memories/src/tests.rs b/codex-rs/ext/memories/src/tests.rs index b732ca22c..e88d7d8db 100644 --- a/codex-rs/ext/memories/src/tests.rs +++ b/codex-rs/ext/memories/src/tests.rs @@ -13,6 +13,7 @@ use codex_tools::ToolOutput; use codex_utils_absolute_path::test_support::PathBufExt; use codex_utils_absolute_path::test_support::PathExt; use codex_utils_absolute_path::test_support::test_path_buf; +use codex_utils_output_truncation::TruncationPolicy; use pretty_assertions::assert_eq; use serde_json::json; @@ -134,8 +135,10 @@ async fn read_tool_reads_memory_file() { let output = tool .handle(ToolCall { + turn_id: "turn-1".to_string(), call_id: "call-1".to_string(), tool_name: memory_tool_name(crate::READ_TOOL_NAME), + truncation_policy: TruncationPolicy::Bytes(1024), payload: payload.clone(), }) .await @@ -176,8 +179,10 @@ async fn search_tool_accepts_multiple_queries() { let output = tool .handle(ToolCall { + turn_id: "turn-1".to_string(), call_id: "call-1".to_string(), tool_name: memory_tool_name(crate::SEARCH_TOOL_NAME), + truncation_policy: TruncationPolicy::Bytes(1024), payload: payload.clone(), }) .await @@ -244,8 +249,10 @@ async fn search_tool_accepts_windowed_all_match_mode() { let output = tool .handle(ToolCall { + turn_id: "turn-1".to_string(), call_id: "call-1".to_string(), tool_name: memory_tool_name(crate::SEARCH_TOOL_NAME), + truncation_policy: TruncationPolicy::Bytes(1024), payload: payload.clone(), }) .await @@ -292,8 +299,10 @@ async fn search_tool_rejects_legacy_single_query() { let result = tool .handle(ToolCall { + turn_id: "turn-1".to_string(), call_id: "call-1".to_string(), tool_name: memory_tool_name(crate::SEARCH_TOOL_NAME), + truncation_policy: TruncationPolicy::Bytes(1024), payload, }) .await; diff --git a/codex-rs/tools/Cargo.toml b/codex-rs/tools/Cargo.toml index e8d1134a6..334ce7958 100644 --- a/codex-rs/tools/Cargo.toml +++ b/codex-rs/tools/Cargo.toml @@ -14,6 +14,7 @@ codex-code-mode = { workspace = true } codex-features = { workspace = true } codex-protocol = { workspace = true } codex-utils-absolute-path = { workspace = true } +codex-utils-output-truncation = { workspace = true } codex-utils-pty = { workspace = true } codex-utils-string = { workspace = true } rmcp = { workspace = true, default-features = false, features = [ diff --git a/codex-rs/tools/src/tool_call.rs b/codex-rs/tools/src/tool_call.rs index f12a338cf..f92c92f97 100644 --- a/codex-rs/tools/src/tool_call.rs +++ b/codex-rs/tools/src/tool_call.rs @@ -1,12 +1,15 @@ use crate::FunctionCallError; use crate::ToolName; use crate::ToolPayload; +use codex_utils_output_truncation::TruncationPolicy; // TODO: this is temporary and will disappear in the next PR (as we make codex-extension-api generic on Invocation. #[derive(Clone, Debug)] pub struct ToolCall { + pub turn_id: String, pub call_id: String, pub tool_name: ToolName, + pub truncation_policy: TruncationPolicy, pub payload: ToolPayload, }