From e3609f2d02a5896c391fa4c07335165c9272b686 Mon Sep 17 00:00:00 2001 From: Tamir Duberstein Date: Mon, 24 Aug 2026 10:53:48 +0000 Subject: [PATCH] Expose invocation sources to extension tools (#40382) ## What changed - Add `ToolCallSource` to extension `ToolCall`s and propagate direct or Code Mode invocation metadata, including the runtime cell and nested tool-call IDs. - Add `ToolCall::response_byte_budget` so extension tools can bound direct-call responses by the host truncation policy while allowing Code Mode responses up to the tool's own limit. - Share the skills extension's response-size limit between its `list` and `read` tools. ## Testing - Verify that Code Mode source metadata is preserved when core constructs an extension tool call. GitOrigin-RevId: b10aef01be30d35b67bee3af24a1a4175c9ee112 --- .../src/tools/handlers/extension_tools.rs | 15 ++++++++- codex-rs/core/src/tools/lifecycle.rs | 2 +- .../ext/extension-api/src/contributors.rs | 1 - .../src/contributors/tool_lifecycle.rs | 15 +-------- codex-rs/ext/extension-api/src/lib.rs | 2 +- .../ext/goal/tests/goal_extension_backend.rs | 1 + .../tests/history_notes_extension.rs | 2 ++ codex-rs/ext/memories/src/tests.rs | 7 ++++ codex-rs/ext/skills/src/tools/list.rs | 6 ++-- codex-rs/ext/skills/src/tools/mod.rs | 1 + codex-rs/ext/skills/src/tools/read.rs | 6 ++-- codex-rs/ext/skills/tests/skills_extension.rs | 2 ++ codex-rs/tools/src/lib.rs | 1 + codex-rs/tools/src/tool_call.rs | 33 +++++++++++++++++++ 14 files changed, 70 insertions(+), 24 deletions(-) diff --git a/codex-rs/core/src/tools/handlers/extension_tools.rs b/codex-rs/core/src/tools/handlers/extension_tools.rs index 57b6e51e95..d31872a525 100644 --- a/codex-rs/core/src/tools/handlers/extension_tools.rs +++ b/codex-rs/core/src/tools/handlers/extension_tools.rs @@ -22,6 +22,7 @@ use crate::session::turn_context::TurnContext; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; use crate::tools::handlers::apply_granted_turn_permissions; +use crate::tools::lifecycle::extension_tool_call_source; use crate::tools::registry::CoreToolRuntime; use crate::tools::registry::ToolExecutor; use crate::turn_metadata::McpTurnMetadataContext; @@ -183,6 +184,7 @@ async fn to_extension_call(invocation: &ToolInvocation) -> ExtensionToolCall { model: invocation.turn.model_info.slug.clone(), codex_turn_metadata, truncation_policy: invocation.turn.model_info.truncation_policy.into(), + source: extension_tool_call_source(invocation.source.clone()), conversation_history, turn_item_emitter: Arc::new(CoreTurnItemEmitter { session: Arc::downgrade(&invocation.session), @@ -207,6 +209,7 @@ mod tests { use codex_protocol::protocol::ImageGenerationBeginEvent; use codex_protocol::protocol::ImageGenerationEndEvent; use codex_tools::ExtensionTurnItem; + use codex_tools::ToolCallSource as ExtensionToolCallSource; use codex_utils_absolute_path::test_support::PathExt; use codex_utils_absolute_path::test_support::test_path_buf; use core_test_support::responses::strip_response_item_id; @@ -417,7 +420,10 @@ mod tests { 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, + source: ToolCallSource::CodeMode { + cell_id: "cell-1".to_string(), + runtime_tool_call_id: "nested-call-1".to_string(), + }, payload: ToolPayload::Function { arguments: json!({ "message": "hello" }).to_string(), }, @@ -438,6 +444,13 @@ mod tests { ); assert_eq!(captured_call.model, model); assert_eq!(captured_call.truncation_policy, truncation_policy); + assert_eq!( + captured_call.source, + ExtensionToolCallSource::CodeMode { + cell_id: "cell-1".to_string(), + runtime_tool_call_id: "nested-call-1".to_string(), + } + ); assert_eq!( captured_call .environments diff --git a/codex-rs/core/src/tools/lifecycle.rs b/codex-rs/core/src/tools/lifecycle.rs index 3e63439aee..b435f26ea5 100644 --- a/codex-rs/core/src/tools/lifecycle.rs +++ b/codex-rs/core/src/tools/lifecycle.rs @@ -94,7 +94,7 @@ async fn notify_tool_finish_parts( } } -fn extension_tool_call_source(source: ToolCallSource) -> ExtensionToolCallSource { +pub(crate) fn extension_tool_call_source(source: ToolCallSource) -> ExtensionToolCallSource { match source { ToolCallSource::Direct | ToolCallSource::DirectPlaintextMessage => { ExtensionToolCallSource::Direct diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index e79b0da8fc..5dfc128de3 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -39,7 +39,6 @@ pub use thread_lifecycle::ThreadResumeInput; pub use thread_lifecycle::ThreadStartInput; pub use thread_lifecycle::ThreadStopInput; pub use tool_lifecycle::ToolCallOutcome; -pub use tool_lifecycle::ToolCallSource; pub use tool_lifecycle::ToolFinishInput; pub use tool_lifecycle::ToolLifecycleFuture; pub use tool_lifecycle::ToolStartInput; diff --git a/codex-rs/ext/extension-api/src/contributors/tool_lifecycle.rs b/codex-rs/ext/extension-api/src/contributors/tool_lifecycle.rs index c0f2f62870..46e0d1f3b6 100644 --- a/codex-rs/ext/extension-api/src/contributors/tool_lifecycle.rs +++ b/codex-rs/ext/extension-api/src/contributors/tool_lifecycle.rs @@ -2,6 +2,7 @@ use std::future::Future; use std::pin::Pin; use std::sync::Arc; +use codex_tools::ToolCallSource; use codex_tools::ToolName; use codex_tools::ToolPayload; @@ -11,20 +12,6 @@ use crate::ExtensionData; /// Future returned by one tool-lifecycle callback. pub type ToolLifecycleFuture<'a> = Pin + Send + 'a>>; -/// Host-visible source for a model tool call. -#[derive(Clone, Debug, Eq, PartialEq)] -pub enum ToolCallSource { - /// The model invoked the tool directly. - Direct, - /// Code mode invoked the tool while executing a runtime cell. - CodeMode { - /// Runtime cell that issued the nested tool request. - cell_id: String, - /// Code-mode's per-cell tool invocation id. - runtime_tool_call_id: String, - }, -} - /// Extension-facing outcome for a finished tool call. #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub enum ToolCallOutcome { diff --git a/codex-rs/ext/extension-api/src/lib.rs b/codex-rs/ext/extension-api/src/lib.rs index 92e359aa25..04d6caf50e 100644 --- a/codex-rs/ext/extension-api/src/lib.rs +++ b/codex-rs/ext/extension-api/src/lib.rs @@ -27,6 +27,7 @@ pub use codex_tools::JsonToolOutput; pub use codex_tools::NoopTurnItemEmitter; pub use codex_tools::ResponsesApiTool; pub use codex_tools::ToolCall; +pub use codex_tools::ToolCallSource; pub use codex_tools::ToolEnvironment; pub use codex_tools::ToolExecutor; pub use codex_tools::ToolExecutorFuture; @@ -64,7 +65,6 @@ pub use contributors::ThreadStartInput; pub use contributors::ThreadStopInput; pub use contributors::TokenUsageContributor; pub use contributors::ToolCallOutcome; -pub use contributors::ToolCallSource; pub use contributors::ToolContributor; pub use contributors::ToolFinishInput; pub use contributors::ToolLifecycleContributor; diff --git a/codex-rs/ext/goal/tests/goal_extension_backend.rs b/codex-rs/ext/goal/tests/goal_extension_backend.rs index 7ac890183e..2e856ffedd 100644 --- a/codex-rs/ext/goal/tests/goal_extension_backend.rs +++ b/codex-rs/ext/goal/tests/goal_extension_backend.rs @@ -1487,6 +1487,7 @@ fn tool_call(tool_name: &str, call_id: &str, arguments: serde_json::Value) -> To model: "gpt-test".to_string(), codex_turn_metadata: None, truncation_policy: TruncationPolicy::Bytes(1024), + source: ToolCallSource::Direct, conversation_history: codex_extension_api::ConversationHistory::default(), turn_item_emitter: Arc::new(NoopTurnItemEmitter), environments: Vec::new(), diff --git a/codex-rs/ext/history-notes/tests/history_notes_extension.rs b/codex-rs/ext/history-notes/tests/history_notes_extension.rs index d5675410e3..c87c9443a2 100644 --- a/codex-rs/ext/history-notes/tests/history_notes_extension.rs +++ b/codex-rs/ext/history-notes/tests/history_notes_extension.rs @@ -11,6 +11,7 @@ use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::NoopTurnItemEmitter; use codex_extension_api::ThreadStartInput; use codex_extension_api::ToolCall; +use codex_extension_api::ToolCallSource; use codex_extension_api::ToolExecutor; use codex_extension_api::ToolName; use codex_extension_api::ToolPayload; @@ -235,6 +236,7 @@ fn tool_call(tool_name: ToolName, arguments: serde_json::Value) -> ToolCall { model: "gpt-test".to_string(), codex_turn_metadata: None, truncation_policy: TruncationPolicy::Bytes(1024), + source: ToolCallSource::Direct, conversation_history: ConversationHistory::default(), turn_item_emitter: Arc::new(NoopTurnItemEmitter), environments: Vec::new(), diff --git a/codex-rs/ext/memories/src/tests.rs b/codex-rs/ext/memories/src/tests.rs index e0ad227891..de581de6e4 100644 --- a/codex-rs/ext/memories/src/tests.rs +++ b/codex-rs/ext/memories/src/tests.rs @@ -7,6 +7,7 @@ use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::NoopTurnItemEmitter; use codex_extension_api::PromptSlot; use codex_extension_api::ToolCall; +use codex_extension_api::ToolCallSource; use codex_extension_api::ToolContributor; use codex_extension_api::ToolExecutor; use codex_extension_api::ToolName; @@ -221,6 +222,7 @@ async fn add_ad_hoc_note_tool_creates_note_file() { model: "gpt-test".to_string(), codex_turn_metadata: None, truncation_policy: TruncationPolicy::Bytes(1024), + source: ToolCallSource::Direct, conversation_history: codex_extension_api::ConversationHistory::default(), turn_item_emitter: Arc::new(NoopTurnItemEmitter), environments: Vec::new(), @@ -266,6 +268,7 @@ async fn add_ad_hoc_note_tool_rejects_paths_as_filenames() { model: "gpt-test".to_string(), codex_turn_metadata: None, truncation_policy: TruncationPolicy::Bytes(1024), + source: ToolCallSource::Direct, conversation_history: codex_extension_api::ConversationHistory::default(), turn_item_emitter: Arc::new(NoopTurnItemEmitter), environments: Vec::new(), @@ -312,6 +315,7 @@ async fn read_tool_reads_memory_file() { model: "gpt-test".to_string(), codex_turn_metadata: None, truncation_policy: TruncationPolicy::Bytes(1024), + source: ToolCallSource::Direct, conversation_history: codex_extension_api::ConversationHistory::default(), turn_item_emitter: Arc::new(NoopTurnItemEmitter), environments: Vec::new(), @@ -428,6 +432,7 @@ async fn search_tool_accepts_multiple_queries() { model: "gpt-test".to_string(), codex_turn_metadata: None, truncation_policy: TruncationPolicy::Bytes(1024), + source: ToolCallSource::Direct, conversation_history: codex_extension_api::ConversationHistory::default(), turn_item_emitter: Arc::new(NoopTurnItemEmitter), environments: Vec::new(), @@ -503,6 +508,7 @@ async fn search_tool_accepts_windowed_all_match_mode() { model: "gpt-test".to_string(), codex_turn_metadata: None, truncation_policy: TruncationPolicy::Bytes(1024), + source: ToolCallSource::Direct, conversation_history: codex_extension_api::ConversationHistory::default(), turn_item_emitter: Arc::new(NoopTurnItemEmitter), environments: Vec::new(), @@ -558,6 +564,7 @@ async fn search_tool_rejects_legacy_single_query() { model: "gpt-test".to_string(), codex_turn_metadata: None, truncation_policy: TruncationPolicy::Bytes(1024), + source: ToolCallSource::Direct, conversation_history: codex_extension_api::ConversationHistory::default(), turn_item_emitter: Arc::new(NoopTurnItemEmitter), environments: Vec::new(), diff --git a/codex-rs/ext/skills/src/tools/list.rs b/codex-rs/ext/skills/src/tools/list.rs index 030bf4249f..3885716621 100644 --- a/codex-rs/ext/skills/src/tools/list.rs +++ b/codex-rs/ext/skills/src/tools/list.rs @@ -15,6 +15,7 @@ use crate::render::truncate_utf8_to_bytes; use crate::warnings::bounded_warnings; use super::MAX_HANDLE_BYTES; +use super::MAX_SKILL_RESPONSE_BYTES; use super::SkillToolAuthority; use super::SkillToolAuthoritySelector; use super::SkillToolContext; @@ -29,7 +30,6 @@ use super::skill_tool_name; const TOOL_NAME: &str = "list"; const MAX_SKILLS_PER_PAGE: usize = 20; -const MAX_LIST_RESPONSE_BYTES: usize = 512 * 1024; const OVERSIZED_ENTRY_WARNING: &str = "Some skills were omitted because their metadata is too large."; @@ -114,7 +114,7 @@ impl ToolExecutor for ListTool { warnings: warnings.clone(), next_cursor: (end < skills.len()).then(|| pagination_cursor(&skills, end)), }; - if serialized_len(&response)? <= MAX_LIST_RESPONSE_BYTES { + if serialized_len(&response)? <= MAX_SKILL_RESPONSE_BYTES { return skill_json_output(&response, args.authority); } if end.saturating_sub(start) > 1 { @@ -137,7 +137,7 @@ fn single_entry_response_is_bounded(skill: &ListedSkill) -> bool { warnings: Vec::new(), next_cursor: Some(pagination_cursor(skill, usize::MAX)), }) - .is_ok_and(|size| size <= MAX_LIST_RESPONSE_BYTES) + .is_ok_and(|size| size <= MAX_SKILL_RESPONSE_BYTES) } fn listed_skill(entry: SkillCatalogEntry) -> Option { diff --git a/codex-rs/ext/skills/src/tools/mod.rs b/codex-rs/ext/skills/src/tools/mod.rs index 62913d230b..4869c5f497 100644 --- a/codex-rs/ext/skills/src/tools/mod.rs +++ b/codex-rs/ext/skills/src/tools/mod.rs @@ -50,6 +50,7 @@ mod schema; const SKILLS_NAMESPACE: &str = "skills"; const MAX_HANDLE_BYTES: usize = 2_048; +const MAX_SKILL_RESPONSE_BYTES: usize = 512 * 1024; pub(crate) fn skill_tools( providers: SkillProviders, diff --git a/codex-rs/ext/skills/src/tools/read.rs b/codex-rs/ext/skills/src/tools/read.rs index fc092008bc..6b5b19ef6c 100644 --- a/codex-rs/ext/skills/src/tools/read.rs +++ b/codex-rs/ext/skills/src/tools/read.rs @@ -14,6 +14,7 @@ use crate::provider::SkillReadRequest; use crate::render::build_alias_plan; use super::MAX_HANDLE_BYTES; +use super::MAX_SKILL_RESPONSE_BYTES; use super::SkillToolAuthority; use super::SkillToolContext; use super::pagination_cursor; @@ -26,7 +27,6 @@ use super::skill_tool_name; use super::validate_handle; const TOOL_NAME: &str = "read"; -const MAX_READ_RESPONSE_BYTES: usize = 512 * 1024; #[derive(Deserialize, JsonSchema)] struct ReadArgs { @@ -237,7 +237,7 @@ fn page_response( next_cursor, }; let complete = response(contents.len(), None); - if serialized_len(&complete)? <= MAX_READ_RESPONSE_BYTES { + if serialized_len(&complete)? <= MAX_SKILL_RESPONSE_BYTES { return Ok(complete); } @@ -248,7 +248,7 @@ fn page_response( end -= 1; } let candidate = response(end, Some(pagination_cursor(contents, end))); - if serialized_len(&candidate)? <= MAX_READ_RESPONSE_BYTES { + if serialized_len(&candidate)? <= MAX_SKILL_RESPONSE_BYTES { return Ok(candidate); } } diff --git a/codex-rs/ext/skills/tests/skills_extension.rs b/codex-rs/ext/skills/tests/skills_extension.rs index dd0a303064..33f8a07f9e 100644 --- a/codex-rs/ext/skills/tests/skills_extension.rs +++ b/codex-rs/ext/skills/tests/skills_extension.rs @@ -24,6 +24,7 @@ use codex_extension_api::SkillInvocationInput; use codex_extension_api::SkillInvocationKind; use codex_extension_api::ThreadStartInput; use codex_extension_api::ToolCall; +use codex_extension_api::ToolCallSource; use codex_extension_api::ToolPayload; use codex_extension_api::TurnInputContext; use codex_extension_api::WorldStateContributionInput; @@ -1604,6 +1605,7 @@ async fn skills_list_only_returns_model_visible_bounded_metadata() -> TestResult model: "gpt-test".to_string(), codex_turn_metadata: None, truncation_policy: TruncationPolicy::Bytes(1_024), + source: ToolCallSource::Direct, conversation_history: ConversationHistory::default(), turn_item_emitter: Arc::new(NoopTurnItemEmitter), environments: Vec::new(), diff --git a/codex-rs/tools/src/lib.rs b/codex-rs/tools/src/lib.rs index 86491196b3..db84442b17 100644 --- a/codex-rs/tools/src/lib.rs +++ b/codex-rs/tools/src/lib.rs @@ -68,6 +68,7 @@ pub use tool_call::ConversationHistory; pub use tool_call::ExtensionTurnItem; pub use tool_call::NoopTurnItemEmitter; pub use tool_call::ToolCall; +pub use tool_call::ToolCallSource; pub use tool_call::ToolEnvironment; pub use tool_call::TurnItemEmissionFuture; pub use tool_call::TurnItemEmitter; diff --git a/codex-rs/tools/src/tool_call.rs b/codex-rs/tools/src/tool_call.rs index f93a0122fc..4de11fad95 100644 --- a/codex-rs/tools/src/tool_call.rs +++ b/codex-rs/tools/src/tool_call.rs @@ -86,6 +86,20 @@ impl TurnItemEmitter for NoopTurnItemEmitter { } } +/// Host-visible source for a model tool call. +#[derive(Clone, Debug, Eq, PartialEq)] +pub enum ToolCallSource { + /// The model invoked the tool directly. + Direct, + /// Code mode invoked the tool while executing a runtime cell. + CodeMode { + /// Runtime cell that issued the nested tool request. + cell_id: String, + /// Code-mode's per-cell tool invocation id. + runtime_tool_call_id: String, + }, +} + #[derive(Clone)] pub struct ToolCall { pub turn_id: String, @@ -94,6 +108,7 @@ pub struct ToolCall { pub model: String, pub codex_turn_metadata: Option, pub truncation_policy: TruncationPolicy, + pub source: ToolCallSource, pub conversation_history: ConversationHistory, pub turn_item_emitter: Arc, pub environments: Vec, @@ -112,6 +127,7 @@ impl std::fmt::Debug for ToolCall { &self.codex_turn_metadata.is_some(), ) .field("truncation_policy", &self.truncation_policy) + .field("source", &self.source) .field("conversation_history", &self.conversation_history) .field("turn_item_emitter", &"") .field("environment_count", &self.environments.len()) @@ -121,6 +137,23 @@ impl std::fmt::Debug for ToolCall { } impl ToolCall { + /// Returns the response-content budget, bounded by the tool's own size limit. + /// + /// Direct calls use the host's effective text-output allowance. Code Mode receives + /// typed results without that truncation, so only the tool's limit applies. + /// Callers must include serialization overhead when fitting a response to this budget. + pub fn response_byte_budget(&self, max_response_bytes: usize) -> usize { + match &self.source { + ToolCallSource::Direct => { + max_response_bytes.min((self.truncation_policy * 1.2).byte_budget()) + } + ToolCallSource::CodeMode { + cell_id: _, + runtime_tool_call_id: _, + } => max_response_bytes, + } + } + pub fn function_arguments(&self) -> Result<&str, FunctionCallError> { match &self.payload { ToolPayload::Function { arguments } => Ok(arguments),