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

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

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

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;

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<Box<dyn Future<Output = ()> + 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 {

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;

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(),

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(),

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(),

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<ToolCall> 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<ListedSkill> {

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,

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);
         }
     }

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(),

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;

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<String>,
     pub truncation_policy: TruncationPolicy,
+    pub source: ToolCallSource,
     pub conversation_history: ConversationHistory,
     pub turn_item_emitter: Arc<dyn TurnItemEmitter>,
     pub environments: Vec<ToolEnvironment>,
@@ -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", &"<host 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),