Run tool start callbacks after pre-tool hooks (#38568)

## What changed

- Invoke `ToolLifecycleContributor::on_tool_start` only after pre-tool hooks
  have finalized the invocation.
- Pass hook-rewritten arguments and the post-hook conversation snapshot to the
  callback.
- Skip the start callback when a hook denies execution or supplies input that
  cannot be applied.

## Testing

Add lifecycle tests covering rewritten input, hook-added context, denied tool
calls, and invalid rewritten input.

GitOrigin-RevId: 936efaf4a8c35321f9982ff2a33bb8bece1ebf39

codex-rs/core/src/tools/registry.rs

@@ -566,8 +566,6 @@ impl ToolRegistry {
             return Err(err);
         }
 
-        notify_tool_start(&invocation).await;
-
         if let Some(pre_tool_use_payload) = tool.pre_tool_use_payload(&invocation) {
             match run_pre_tool_use_hooks(
                 &invocation.session,
@@ -614,6 +612,8 @@ impl ToolRegistry {
             }
         }
 
+        notify_tool_start(&invocation).await;
+
         if let Some(command) = shell_script_for_invocation(&invocation) {
             let parsed = parse_shell_script(&command);
             let mut categories = parsed.iter().map(|command| match command {

codex-rs/core/tests/suite/tool_lifecycle.rs

@@ -1,6 +1,9 @@
+use std::fs;
+use std::path::Path;
 use std::sync::Arc;
 use std::sync::Mutex;
 
+use anyhow::Context;
 use anyhow::Result;
 use codex_core::config::Config;
 use codex_extension_api::ExtensionRegistryBuilder;
@@ -9,8 +12,10 @@ use codex_extension_api::ToolLifecycleContributor;
 use codex_extension_api::ToolLifecycleFuture;
 use codex_extension_api::ToolStartInput;
 use codex_protocol::models::ContentItem;
+use core_test_support::hooks::trust_discovered_hooks;
 use core_test_support::responses;
 use core_test_support::skip_if_no_network;
+use core_test_support::skip_if_wine_exec;
 use core_test_support::test_codex::test_codex;
 use pretty_assertions::assert_eq;
 use serde_json::json;
@@ -21,6 +26,7 @@ struct RecordedHistory {
     items: Vec<ResponseItem>,
 }
 
+#[derive(Default)]
 struct ConversationHistoryRecorder {
     histories: Mutex<Vec<RecordedHistory>>,
 }
@@ -81,9 +87,7 @@ async fn tool_start_receives_conversation_history() -> Result<()> {
     )
     .await;
 
-    let recorder = Arc::new(ConversationHistoryRecorder {
-        histories: Mutex::new(Vec::new()),
-    });
+    let recorder = Arc::new(ConversationHistoryRecorder::default());
     let mut extensions = ExtensionRegistryBuilder::<Config>::new();
     extensions.tool_lifecycle_contributor(recorder.clone());
     let test = test_codex()
@@ -147,3 +151,182 @@ async fn tool_start_receives_conversation_history() -> Result<()> {
 
     Ok(())
 }
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn tool_start_receives_rewritten_payload_and_post_hook_history() -> Result<()> {
+    skip_if_no_network!(Ok(()));
+    skip_if_wine_exec!(Ok(()), "command hooks require a host-native executor");
+
+    let server = responses::start_mock_server().await;
+    let call_id = "rewritten-plan-call";
+    let original_input = json!({
+        "plan": [{ "step": "Original step", "status": "in_progress" }]
+    });
+    let rewritten_input = json!({
+        "plan": [{ "step": "Rewritten step", "status": "completed" }]
+    });
+    let additional_context = "Only available after the pre-tool hook.";
+    responses::mount_sse_sequence(
+        &server,
+        vec![
+            responses::sse(vec![
+                responses::ev_function_call(call_id, "update_plan", &original_input.to_string()),
+                responses::ev_completed("first-response"),
+            ]),
+            responses::sse(vec![
+                responses::ev_assistant_message("assistant-1", "done"),
+                responses::ev_completed("second-response"),
+            ]),
+        ],
+    )
+    .await;
+
+    let hook_output = json!({
+        "hookSpecificOutput": {
+            "hookEventName": "PreToolUse",
+            "permissionDecision": "allow",
+            "updatedInput": rewritten_input,
+            "additionalContext": additional_context,
+        }
+    });
+    let recorder = Arc::new(ConversationHistoryRecorder::default());
+    let mut extensions = ExtensionRegistryBuilder::<Config>::new();
+    extensions.tool_lifecycle_contributor(recorder.clone());
+    let test = test_codex()
+        .with_extensions(Arc::new(extensions.build()))
+        .with_pre_build_hook(move |home| {
+            write_pre_tool_hook(home, "^update_plan$", &hook_output)
+                .expect("write pre-tool hook fixture");
+        })
+        .with_config(trust_discovered_hooks)
+        .build_with_auto_env(&server)
+        .await?;
+
+    test.submit_text_turn("Update the plan.").await?;
+
+    let histories = recorder
+        .histories
+        .lock()
+        .unwrap_or_else(std::sync::PoisonError::into_inner);
+    let [history] = histories.as_slice() else {
+        panic!("expected one tool start, got {}", histories.len());
+    };
+    assert_eq!(
+        (
+            history.call_id.as_str(),
+            serde_json::from_str::<serde_json::Value>(&history.arguments)?,
+        ),
+        (call_id, rewritten_input)
+    );
+    assert!(history.items.iter().any(|item| matches!(
+        item,
+        ResponseItem::Message { role, content, .. }
+            if role == "developer"
+                && content.iter().any(|content| matches!(
+                    content,
+                    ContentItem::InputText { text } if text == additional_context
+                ))
+    )));
+
+    Ok(())
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn tool_start_is_not_called_when_pre_tool_hook_prevents_execution() -> Result<()> {
+    skip_if_no_network!(Ok(()));
+    skip_if_wine_exec!(Ok(()), "command hooks require a host-native executor");
+
+    for (tool_name, matcher, arguments, hook_output) in [
+        (
+            "update_plan",
+            "^update_plan$",
+            json!({ "plan": [{ "step": "Blocked step", "status": "in_progress" }] }),
+            json!({
+                "hookSpecificOutput": {
+                    "hookEventName": "PreToolUse",
+                    "permissionDecision": "deny",
+                    "permissionDecisionReason": "blocked by lifecycle test",
+                }
+            }),
+        ),
+        (
+            "shell_command",
+            "^Bash$",
+            json!({ "command": "echo original" }),
+            json!({
+                "hookSpecificOutput": {
+                    "hookEventName": "PreToolUse",
+                    "permissionDecision": "allow",
+                    "updatedInput": { "command": 123 },
+                }
+            }),
+        ),
+    ] {
+        let server = responses::start_mock_server().await;
+        let call_id = format!("prevented-{tool_name}-call");
+        responses::mount_sse_sequence(
+            &server,
+            vec![
+                responses::sse(vec![
+                    responses::ev_function_call(&call_id, tool_name, &arguments.to_string()),
+                    responses::ev_completed("first-response"),
+                ]),
+                responses::sse(vec![
+                    responses::ev_assistant_message("assistant-1", "done"),
+                    responses::ev_completed("second-response"),
+                ]),
+            ],
+        )
+        .await;
+
+        let recorder = Arc::new(ConversationHistoryRecorder::default());
+        let mut extensions = ExtensionRegistryBuilder::<Config>::new();
+        extensions.tool_lifecycle_contributor(recorder.clone());
+        let test = test_codex()
+            .with_extensions(Arc::new(extensions.build()))
+            .with_pre_build_hook(move |home| {
+                write_pre_tool_hook(home, matcher, &hook_output)
+                    .expect("write pre-tool hook fixture");
+            })
+            .with_config(trust_discovered_hooks)
+            .build_with_auto_env(&server)
+            .await?;
+
+        test.submit_text_turn("Run the tool.").await?;
+
+        assert!(
+            recorder
+                .histories
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner)
+                .is_empty(),
+            "tool start should not run for {tool_name}"
+        );
+    }
+
+    Ok(())
+}
+
+fn write_pre_tool_hook(home: &Path, matcher: &str, output: &serde_json::Value) -> Result<()> {
+    let script_path = home.join("tool_lifecycle_hook.py");
+    let output_json = serde_json::to_string(output).context("serialize pre-tool hook output")?;
+    fs::write(
+        &script_path,
+        format!("import json\nimport sys\njson.load(sys.stdin)\nprint({output_json:?})\n"),
+    )
+    .context("write pre-tool hook script")?;
+    let hooks = json!({
+        "hooks": {
+            "PreToolUse": [{
+                "matcher": matcher,
+                "hooks": [{
+                    "type": "command",
+                    "command": format!("python3 {}", script_path.display()),
+                }]
+            }]
+        }
+    });
+    fs::write(home.join("hooks.json"), hooks.to_string()).context("write pre-tool hooks.json")?;
+
+    Ok(())
+}

codex-rs/ext/extension-api/src/contributors.rs

@@ -299,12 +299,18 @@ pub trait ToolContributor: Send + Sync {
 /// exposed input without rewriting the invocation. Use `ToolContributor` for
 /// owning a tool implementation and hooks for policy that changes tool payloads.
 pub trait ToolLifecycleContributor: Send + Sync {
-    /// Called once the host has accepted a tool call for execution.
+    /// Called after pre-tool hooks finalize an invocation and before execution.
+    ///
+    /// Calls blocked by hooks, or whose hook-provided input cannot be applied,
+    /// do not reach this callback.
     fn on_tool_start<'a>(&'a self, _input: ToolStartInput<'a>) -> ToolLifecycleFuture<'a> {
         Box::pin(std::future::ready(()))
     }
 
     /// Called after the tool call returns, is blocked, fails, or is cancelled.
+    ///
+    /// A matching start callback does not exist when execution is blocked,
+    /// hook-provided input cannot be applied, or cancellation wins first.
     fn on_tool_finish<'a>(&'a self, _input: ToolFinishInput<'a>) -> ToolLifecycleFuture<'a> {
         Box::pin(std::future::ready(()))
     }

codex-rs/ext/extension-api/src/contributors/tool_lifecycle.rs

@@ -60,11 +60,11 @@ pub struct ToolStartInput<'a> {
     pub call_id: &'a str,
     /// Tool name as routed by the host.
     pub tool_name: &'a ToolName,
-    /// Tool arguments before any pre-tool-use hook rewrites.
+    /// Finalized tool arguments, including any pre-tool-use hook rewrites.
     ///
     /// Payloads can contain sensitive plaintext and must not be logged.
     pub payload: &'a ToolPayload,
-    /// Shared read-only snapshot of the conversation when the tool started.
+    /// Shared read-only snapshot taken after pre-tool hooks have completed.
     pub conversation_history: Arc<dyn ConversationHistorySnapshot>,
     /// Source that issued the tool call.
     pub source: ToolCallSource,