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

use super::*;
use crate::session::session::Session;
use crate::session::step_context::StepContext;
use codex_protocol::DEFAULT_FUNCTION_NAMESPACE;
use codex_protocol::models::ResponseItem;
use codex_utils_output_truncation::TruncationPolicy;
use futures::future::BoxFuture;
use pretty_assertions::assert_eq;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;

struct TestHandler {
    tool_name: codex_tools::ToolName,
}

impl ToolExecutor<ToolInvocation> for TestHandler {
    fn tool_name(&self) -> codex_tools::ToolName {
        self.tool_name.clone()
    }

    fn spec(&self) -> codex_tools::ToolSpec {
        test_spec(&self.tool_name)
    }

    fn handle<'a>(&'a self, _invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'a>
    where
        ToolInvocation: 'a,
    {
        Box::pin(async {
            Ok(
                Box::new(crate::tools::context::FunctionToolOutput::from_text(
                    "ok".to_string(),
                    Some(true),
                )) as Box<dyn crate::tools::context::ToolOutput>,
            )
        })
    }
}

impl CoreToolRuntime for TestHandler {}

struct ReadinessTestHandler {
    handler: TestHandler,
    readiness_waits: Arc<AtomicUsize>,
}

impl ToolExecutor<ToolInvocation> for ReadinessTestHandler {
    fn tool_name(&self) -> codex_tools::ToolName {
        self.handler.tool_name()
    }

    fn spec(&self) -> codex_tools::ToolSpec {
        self.handler.spec()
    }

    fn handle<'a>(&'a self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'a>
    where
        ToolInvocation: 'a,
    {
        self.handler.handle(invocation)
    }
}

impl CoreToolRuntime for ReadinessTestHandler {
    fn wait_until_ready<'a>(&'a self, _session: &'a Arc<Session>) -> Option<BoxFuture<'a, ()>> {
        Some(Box::pin(async {
            self.readiness_waits.fetch_add(1, Ordering::Relaxed);
        }))
    }
}

#[derive(Clone)]
enum LifecycleTestResult {
    Ok { success: bool },
    Err,
}

struct LifecycleTestHandler {
    tool_name: codex_tools::ToolName,
    result: LifecycleTestResult,
}

impl ToolExecutor<ToolInvocation> for LifecycleTestHandler {
    fn tool_name(&self) -> codex_tools::ToolName {
        self.tool_name.clone()
    }

    fn spec(&self) -> codex_tools::ToolSpec {
        test_spec(&self.tool_name)
    }

    fn handle<'a>(&'a self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'a>
    where
        ToolInvocation: 'a,
    {
        assert_eq!(
            invocation.tool_name,
            self.tool_name.clone().with_default_namespace()
        );
        Box::pin(self.handle_call())
    }
}

impl LifecycleTestHandler {
    async fn handle_call(
        &self,
    ) -> Result<Box<dyn crate::tools::context::ToolOutput>, FunctionCallError> {
        match self.result.clone() {
            LifecycleTestResult::Ok { success } => Ok(Box::new(
                crate::tools::context::FunctionToolOutput::from_text(
                    "ok".to_string(),
                    Some(success),
                ),
            )
                as Box<dyn crate::tools::context::ToolOutput>),
            LifecycleTestResult::Err => Err(FunctionCallError::RespondToModel(
                "handler failed".to_string(),
            )),
        }
    }
}

impl CoreToolRuntime for LifecycleTestHandler {}

fn test_spec(tool_name: &codex_tools::ToolName) -> codex_tools::ToolSpec {
    codex_tools::ToolSpec::Function(codex_tools::ResponsesApiTool {
        name: tool_name.name.clone(),
        description: "Test tool.".to_string(),
        strict: false,
        defer_loading: None,
        parameters: codex_tools::JsonSchema::default(),
        output_schema: None,
    })
}

#[derive(Debug, PartialEq, Eq)]
enum RecordedToolLifecycle {
    Start {
        call_id: String,
        tool_name: codex_tools::ToolName,
        root_turn_id: Option<String>,
    },
    Finish {
        call_id: String,
        tool_name: codex_tools::ToolName,
        outcome: codex_extension_api::ToolCallOutcome,
    },
}

struct ToolLifecycleRecorder {
    records: Arc<std::sync::Mutex<Vec<RecordedToolLifecycle>>>,
}

impl codex_extension_api::ToolLifecycleContributor for ToolLifecycleRecorder {
    fn on_tool_start<'a>(
        &'a self,
        input: codex_extension_api::ToolStartInput<'a>,
    ) -> codex_extension_api::ToolLifecycleFuture<'a> {
        let records = Arc::clone(&self.records);
        let record = RecordedToolLifecycle::Start {
            call_id: input.call_id.to_string(),
            tool_name: input.tool_name.clone(),
            root_turn_id: input.root_turn_id.map(str::to_owned),
        };
        Box::pin(async move {
            records
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner)
                .push(record);
        })
    }

    fn on_tool_finish<'a>(
        &'a self,
        input: codex_extension_api::ToolFinishInput<'a>,
    ) -> codex_extension_api::ToolLifecycleFuture<'a> {
        let records = Arc::clone(&self.records);
        let record = RecordedToolLifecycle::Finish {
            call_id: input.call_id.to_string(),
            tool_name: input.tool_name.clone(),
            outcome: input.outcome,
        };
        Box::pin(async move {
            records
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner)
                .push(record);
        })
    }
}

#[test]
fn handler_normalizes_only_the_default_namespace() {
    let namespace = "mcp__codex_apps__gmail";
    let tool_name = "gmail_get_recent_emails";
    let plain_name = codex_tools::ToolName::plain(tool_name);
    let namespaced_name = codex_tools::ToolName::namespaced(namespace, tool_name);
    let plain_handler = Arc::new(TestHandler {
        tool_name: plain_name.clone(),
    }) as Arc<dyn CoreToolRuntime>;
    let namespaced_handler = Arc::new(TestHandler {
        tool_name: namespaced_name.clone(),
    }) as Arc<dyn CoreToolRuntime>;
    let registry =
        ToolRegistry::from_tools([Arc::clone(&plain_handler), Arc::clone(&namespaced_handler)]);

    let plain = registry.tool(&plain_name);
    let default_namespaced = registry.tool(&codex_tools::ToolName::namespaced(
        DEFAULT_FUNCTION_NAMESPACE,
        tool_name,
    ));
    let empty_namespaced = registry.tool(&codex_tools::ToolName::namespaced("", tool_name));
    let namespaced = registry.tool(&namespaced_name);
    let missing_namespaced = registry.tool(&codex_tools::ToolName::namespaced(
        "mcp__codex_apps__calendar",
        tool_name,
    ));

    assert_eq!(plain.is_some(), true);
    assert_eq!(namespaced.is_some(), true);
    assert_eq!(missing_namespaced.is_none(), true);
    assert!(
        plain
            .as_ref()
            .is_some_and(|handler| Arc::ptr_eq(handler, &plain_handler))
    );
    assert!(
        default_namespaced
            .as_ref()
            .is_some_and(|handler| Arc::ptr_eq(handler, &plain_handler))
    );
    assert!(
        empty_namespaced
            .as_ref()
            .is_some_and(|handler| Arc::ptr_eq(handler, &plain_handler))
    );
    assert!(
        namespaced
            .as_ref()
            .is_some_and(|handler| Arc::ptr_eq(handler, &namespaced_handler))
    );
}

#[test]
fn registry_rejects_default_namespace_alias_collisions() {
    let plain_name = codex_tools::ToolName::plain("lookup");
    let namespaced_name = codex_tools::ToolName::namespaced(DEFAULT_FUNCTION_NAMESPACE, "lookup");

    for [first_name, duplicate_name] in [
        [plain_name.clone(), namespaced_name.clone()],
        [namespaced_name, plain_name],
    ] {
        let winner = Arc::new(TestHandler {
            tool_name: first_name.clone(),
        }) as Arc<dyn CoreToolRuntime>;
        let mut registry = ToolRegistry::from_tools([Arc::clone(&winner)]);

        assert!(!registry.register_external(Arc::new(TestHandler {
            tool_name: duplicate_name.clone(),
        })));
        assert!(
            registry
                .tool(&duplicate_name)
                .is_some_and(|handler| Arc::ptr_eq(&handler, &winner))
        );
        assert_eq!(
            registry.tool_exposure(&duplicate_name),
            Some(ToolExposure::Direct)
        );
        assert_eq!(
            registry.supports_parallel_tool_calls(&duplicate_name),
            Some(false)
        );
        assert!(
            registry
                .remove(&duplicate_name)
                .is_some_and(|handler| Arc::ptr_eq(&handler, &winner))
        );
        assert!(registry.tool(&first_name).is_none());
    }
}

#[test]
fn registry_preserves_external_winners_and_trusted_synthetic_order() {
    let handler = |tool_name| Arc::new(TestHandler { tool_name }) as Arc<dyn CoreToolRuntime>;
    let [first_name, second_name, synthetic_name] =
        ["first", "second", "synthetic"].map(codex_tools::ToolName::plain);
    let first_handler = handler(first_name.clone());

    let mut registry = ToolRegistry::from_tools([Arc::clone(&first_handler)]);
    assert!(!registry.register_external(handler(first_name.clone())));
    let canonical_first_name = first_name.clone().with_default_namespace();
    assert_eq!(registry.first_collision(), Some(&canonical_first_name));
    assert!(registry.register_external(handler(second_name.clone())));
    registry.prepend_trusted(handler(synthetic_name.clone()));

    assert_eq!(
        registry
            .entries()
            .map(|tool| tool.runtime.tool_name())
            .collect::<Vec<_>>(),
        vec![synthetic_name, first_name.clone(), second_name],
    );
    assert!(
        registry
            .remove(&first_name)
            .is_some_and(|handler| Arc::ptr_eq(&handler, &first_handler))
    );
}

#[test]
fn reserved_command_tools_reject_external_runtimes_without_a_builtin() {
    let handler = |tool_name| Arc::new(TestHandler { tool_name }) as Arc<dyn CoreToolRuntime>;
    let mut registry = ToolRegistry::default();

    for reserved_name in ["exec_command", "shell_command"] {
        let tool_name = codex_tools::ToolName::plain(reserved_name);
        let namespaced_tool_name = codex_tools::ToolName::namespaced("client", reserved_name);

        assert!(!registry.register_external(handler(tool_name.clone())));
        assert!(
            !registry
                .register_external_with_exposure(handler(tool_name.clone()), ToolExposure::Direct)
        );
        assert!(
            !registry.register_external(handler(codex_tools::ToolName::namespaced(
                DEFAULT_FUNCTION_NAMESPACE,
                reserved_name,
            )))
        );
        assert!(registry.tool(&tool_name).is_none());
        assert_eq!(registry.first_collision(), None);

        let namespaced_handler = handler(namespaced_tool_name.clone());
        assert!(registry.register_external(Arc::clone(&namespaced_handler)));
        assert!(
            registry
                .tool(&namespaced_tool_name)
                .is_some_and(|runtime| Arc::ptr_eq(&runtime, &namespaced_handler))
        );
    }
}

#[test]
fn registry_records_reserved_exec_command_when_a_matching_tool_exists() {
    let tool_name = codex_tools::ToolName::plain("exec_command");
    let trusted = Arc::new(TestHandler {
        tool_name: tool_name.clone(),
    }) as Arc<dyn CoreToolRuntime>;
    let external = Arc::new(TestHandler {
        tool_name: tool_name.clone(),
    });
    let mut registry = ToolRegistry::from_tools([trusted]);

    assert!(!registry.register_external(external));
    let canonical_tool_name = tool_name.with_default_namespace();
    assert_eq!(registry.first_collision(), Some(&canonical_tool_name));
}

#[test]
fn registry_allows_identical_names_in_different_namespaces() {
    let handler = |tool_name| Arc::new(TestHandler { tool_name }) as Arc<dyn CoreToolRuntime>;
    let mut registry = ToolRegistry::from_tools([handler(codex_tools::ToolName::namespaced(
        "first", "lookup",
    ))]);

    assert!(
        registry.register_external(handler(codex_tools::ToolName::namespaced(
            "second", "lookup",
        )))
    );
    assert_eq!(registry.first_collision(), None);
}

#[tokio::test]
async fn readiness_selects_exact_tool_with_registry_owned_exposure() {
    let (session, _turn) = crate::session::tests::make_session_and_context().await;
    let session = Arc::new(session);
    let plain_name = codex_tools::ToolName::plain("echo");
    let namespaced_name = codex_tools::ToolName::namespaced("mcp__server__", "echo");
    assert!(
        TestHandler {
            tool_name: plain_name.clone(),
        }
        .wait_until_ready(&session)
        .is_none()
    );
    let plain_readiness_waits = Arc::new(AtomicUsize::new(0));
    let namespaced_readiness_waits = Arc::new(AtomicUsize::new(0));
    let plain_handler = Arc::new(ReadinessTestHandler {
        handler: TestHandler {
            tool_name: plain_name.clone(),
        },
        readiness_waits: Arc::clone(&plain_readiness_waits),
    }) as Arc<dyn CoreToolRuntime>;
    let namespaced_handler = Arc::new(ReadinessTestHandler {
        handler: TestHandler {
            tool_name: namespaced_name.clone(),
        },
        readiness_waits: Arc::clone(&namespaced_readiness_waits),
    });
    let mut registry = ToolRegistry::from_tools([plain_handler]);
    registry.register_trusted_with_exposure(namespaced_handler, ToolExposure::DirectModelOnly);

    registry
        .tool(&plain_name)
        .expect("plain runtime should be registered")
        .wait_until_ready(&session)
        .expect("plain runtime should provide a readiness wait")
        .await;
    assert_eq!(
        [
            plain_readiness_waits.load(Ordering::Relaxed),
            namespaced_readiness_waits.load(Ordering::Relaxed),
        ],
        [1, 0]
    );

    registry
        .tool(&namespaced_name)
        .expect("namespaced runtime should be registered")
        .wait_until_ready(&session)
        .expect("namespaced runtime should forward its readiness wait")
        .await;
    assert_eq!(
        [
            plain_readiness_waits.load(Ordering::Relaxed),
            namespaced_readiness_waits.load(Ordering::Relaxed),
        ],
        [1, 1]
    );

    assert!(
        registry
            .tool(&codex_tools::ToolName::namespaced("mcp__missing__", "echo"))
            .is_none()
    );
    assert_eq!(
        [
            plain_readiness_waits.load(Ordering::Relaxed),
            namespaced_readiness_waits.load(Ordering::Relaxed),
        ],
        [1, 1]
    );
}

#[tokio::test]
async fn function_tools_expose_default_hook_payloads_and_rewrites() -> anyhow::Result<()> {
    let (session, turn) = crate::session::tests::make_session_and_context().await;
    let tool_name = codex_tools::ToolName::namespaced("functions.", "echo");
    let handler = TestHandler {
        tool_name: tool_name.clone(),
    };
    let invocation = ToolInvocation {
        payload: ToolPayload::Function {
            arguments: serde_json::json!({ "message": "hello" }).to_string(),
        },
        ..test_invocation(Arc::new(session), Arc::new(turn), "call-1", tool_name)
    };
    let output =
        crate::tools::context::FunctionToolOutput::from_text("echoed".to_string(), Some(true));

    assert_eq!(
        handler.pre_tool_use_payload(&invocation),
        Some(PreToolUsePayload {
            tool_name: HookToolName::new("functions.echo"),
            tool_input: serde_json::json!({ "message": "hello" }),
        })
    );
    assert_eq!(
        handler.post_tool_use_payload(&invocation, &output),
        Some(PostToolUsePayload {
            tool_name: HookToolName::new("functions.echo"),
            tool_use_id: "call-1".to_string(),
            tool_input: serde_json::json!({ "message": "hello" }),
            tool_response: serde_json::json!("echoed"),
        })
    );

    let invocation = handler
        .with_updated_hook_input(invocation, serde_json::json!({ "message": "rewritten" }))?;
    let ToolPayload::Function { arguments } = invocation.payload else {
        panic!("generic rewritten function payload should remain function-shaped");
    };
    assert_eq!(
        serde_json::from_str::<serde_json::Value>(&arguments)?,
        serde_json::json!({ "message": "rewritten" })
    );

    Ok(())
}

#[tokio::test]
async fn function_hook_input_defaults_empty_arguments_to_object() {
    let (session, turn) = crate::session::tests::make_session_and_context().await;
    let tool_name = codex_tools::ToolName::plain("echo");
    let handler = TestHandler {
        tool_name: tool_name.clone(),
    };
    let invocation = ToolInvocation {
        payload: ToolPayload::Function {
            arguments: "  ".to_string(),
        },
        ..test_invocation(Arc::new(session), Arc::new(turn), "call-1", tool_name)
    };

    assert_eq!(
        handler.pre_tool_use_payload(&invocation),
        Some(PreToolUsePayload {
            tool_name: HookToolName::new("echo"),
            tool_input: serde_json::json!({}),
        })
    );
}

#[tokio::test]
async fn spawn_agent_function_tools_use_agent_matcher_alias() {
    let (session, turn) = crate::session::tests::make_session_and_context().await;
    let session = Arc::new(session);
    let turn = Arc::new(turn);

    let hook_payloads = [
        codex_tools::ToolName::plain("spawn_agent"),
        codex_tools::ToolName::namespaced(DEFAULT_FUNCTION_NAMESPACE, "spawn_agent"),
        codex_tools::ToolName::namespaced(MULTI_AGENT_V1_NAMESPACE, "spawn_agent"),
    ]
    .into_iter()
    .map(|tool_name| {
        let handler = TestHandler {
            tool_name: tool_name.clone(),
        };
        let invocation = ToolInvocation {
            payload: ToolPayload::Function {
                arguments: serde_json::json!({ "message": "inspect this repo" }).to_string(),
            },
            ..test_invocation(Arc::clone(&session), Arc::clone(&turn), "call-1", tool_name)
        };
        handler.pre_tool_use_payload(&invocation)
    })
    .collect::<Vec<_>>();

    assert_eq!(
        hook_payloads,
        vec![
            Some(PreToolUsePayload {
                tool_name: HookToolName::spawn_agent(),
                tool_input: serde_json::json!({ "message": "inspect this repo" }),
            }),
            Some(PreToolUsePayload {
                tool_name: HookToolName::spawn_agent(),
                tool_input: serde_json::json!({ "message": "inspect this repo" }),
            }),
            Some(PreToolUsePayload {
                tool_name: HookToolName::spawn_agent(),
                tool_input: serde_json::json!({ "message": "inspect this repo" }),
            }),
        ]
    );
}

#[tokio::test]
async fn code_mode_wait_does_not_expose_default_hook_payloads() {
    let (session, turn) = crate::session::tests::make_session_and_context().await;
    let output = crate::tools::context::FunctionToolOutput::from_text("ok".to_string(), Some(true));

    let wait = crate::tools::handlers::CodeModeWaitHandler;
    let wait_invocation = test_invocation(
        Arc::new(session),
        Arc::new(turn),
        "wait-call",
        wait.tool_name(),
    );
    assert_eq!(wait.pre_tool_use_payload(&wait_invocation), None);
    assert_eq!(wait.post_tool_use_payload(&wait_invocation, &output), None);
}

#[tokio::test]
async fn write_stdin_does_not_expose_default_pre_tool_use_payload() {
    let (session, turn) = crate::session::tests::make_session_and_context().await;

    let write_stdin = crate::tools::handlers::WriteStdinHandler;
    let invocation = test_invocation(
        Arc::new(session),
        Arc::new(turn),
        "write-stdin-call",
        write_stdin.tool_name(),
    );

    assert_eq!(write_stdin.pre_tool_use_payload(&invocation), None);
}

#[test_case::test_case(TruncationPolicy::Tokens(1), 2; "token budget")]
#[test_case::test_case(TruncationPolicy::Bytes(401), 121; "scale bytes before converting to tokens")]
fn post_tool_use_feedback_output_preserves_fallback_token_limit_override(
    truncation_policy: TruncationPolicy,
    expected_token_limit: usize,
) {
    let result = AnyToolResult {
        call_id: "call-1".to_string(),
        payload: ToolPayload::Function {
            arguments: "{}".to_string(),
        },
        result: Box::new(PostToolUseFeedbackOutput {
            original: Box::new(crate::tools::context::McpToolOutput {
                result: codex_protocol::mcp::CallToolResult {
                    content: Vec::new(),
                    structured_content: None,
                    is_error: None,
                    meta: None,
                },
                tool_input: serde_json::json!({}),
                result_metadata_capture_allowed: false,
                wall_time: Duration::ZERO,
                original_image_detail_supported: false,
                truncation_policy,
            }),
            model_visible: crate::tools::context::FunctionToolOutput::from_text(
                "hook feedback".to_string(),
                /*success*/ None,
            ),
        }),
        post_tool_use_payload: None,
    };

    assert_eq!(
        result.into_response(),
        ResponseItemEnvelope {
            item: ResponseItem::from(ResponseInputItem::FunctionCallOutput {
                call_id: "call-1".to_string(),
                output: FunctionCallOutputPayload::from_text("hook feedback".to_string()),
            }),
            metadata: Some(CodexHarnessMetadata {
                history_truncation_token_limit: Some(expected_token_limit),
                ..Default::default()
            }),
        }
    );
}

#[test]
fn post_tool_use_feedback_output_keeps_code_mode_result_typed() {
    let result = AnyToolResult {
        call_id: "call-1".to_string(),
        payload: ToolPayload::Function {
            arguments: "{}".to_string(),
        },
        result: Box::new(PostToolUseFeedbackOutput {
            original: Box::new(codex_tools::JsonToolOutput::new(
                serde_json::json!({ "typed": true }),
            )),
            model_visible: crate::tools::context::FunctionToolOutput::from_text(
                "hook feedback".to_string(),
                /*success*/ None,
            ),
        }),
        post_tool_use_payload: None,
    };

    assert_eq!(
        result.into_response(),
        ResponseItemEnvelope::new(
            ResponseInputItem::FunctionCallOutput {
                call_id: "call-1".to_string(),
                output: codex_protocol::models::FunctionCallOutputPayload::from_text(
                    "hook feedback".to_string()
                ),
            }
            .into()
        )
    );

    let result = AnyToolResult {
        call_id: "call-1".to_string(),
        payload: ToolPayload::Function {
            arguments: "{}".to_string(),
        },
        result: Box::new(PostToolUseFeedbackOutput {
            original: Box::new(codex_tools::JsonToolOutput::new(
                serde_json::json!({ "typed": true }),
            )),
            model_visible: crate::tools::context::FunctionToolOutput::from_text(
                "hook feedback".to_string(),
                /*success*/ None,
            ),
        }),
        post_tool_use_payload: None,
    };

    assert_eq!(
        result.code_mode_result(),
        serde_json::json!({ "typed": true })
    );
}

#[test_case::test_case(false; "success")]
#[test_case::test_case(true; "tool error")]
fn post_tool_use_feedback_output_preserves_mcp_result_metadata(tool_error: bool) {
    let metadata = serde_json::json!({
        "openai/resource_access": { "resources": [] },
        "provider/custom": { "items": [1, null, { "value": true }] },
    });
    let result = PostToolUseFeedbackOutput {
        original: Box::new(crate::tools::context::McpToolOutput {
            result: codex_protocol::mcp::CallToolResult {
                content: vec![serde_json::json!({
                    "type": "text",
                    "text": "original result",
                })],
                structured_content: None,
                is_error: Some(tool_error),
                meta: Some(metadata.clone()),
            },
            tool_input: serde_json::json!({ "query": "rewritten query" }),
            result_metadata_capture_allowed: true,
            wall_time: std::time::Duration::ZERO,
            original_image_detail_supported: false,
            truncation_policy: codex_utils_output_truncation::TruncationPolicy::Bytes(64),
        }),
        model_visible: FunctionToolOutput::from_text(
            "unrelated hook feedback".to_string(),
            /*success*/ None,
        ),
    };
    let payload = ToolPayload::Function {
        arguments: "{}".to_string(),
    };
    let original_result = result.original.code_mode_result(&payload);
    assert_eq!(result.tool_result_metadata(), Some(&metadata));
    assert_eq!(result.code_mode_result(&payload), original_result);
}

#[tokio::test]
async fn dispatch_uses_canonical_tool_names_for_lifecycle_contributors() -> anyhow::Result<()> {
    let (mut session, turn) = crate::session::tests::make_session_and_context().await;
    turn.turn_metadata_state
        .set_root_turn_id("root-turn".to_string());
    let records = Arc::new(std::sync::Mutex::new(Vec::new()));
    let mut builder = codex_extension_api::ExtensionRegistryBuilder::<crate::config::Config>::new();
    builder.tool_lifecycle_contributor(Arc::new(ToolLifecycleRecorder {
        records: Arc::clone(&records),
    }));
    session.services.extensions = Arc::new(builder.build());

    let ok_tool = codex_tools::ToolName::plain("ok_tool");
    let failing_tool = codex_tools::ToolName::namespaced("extensions", "failing_tool");
    let ok_handler = Arc::new(LifecycleTestHandler {
        tool_name: ok_tool.clone(),
        result: LifecycleTestResult::Ok { success: false },
    }) as Arc<dyn CoreToolRuntime>;
    let failing_handler = Arc::new(LifecycleTestHandler {
        tool_name: failing_tool.clone(),
        result: LifecycleTestResult::Err,
    }) as Arc<dyn CoreToolRuntime>;
    let registry = ToolRegistry::from_tools([ok_handler, failing_handler]);
    let session = Arc::new(session);
    let turn = Arc::new(turn);

    registry
        .dispatch_any_with_terminal_outcome(
            test_invocation(
                Arc::clone(&session),
                Arc::clone(&turn),
                "ok-call",
                codex_tools::ToolName::namespaced(DEFAULT_FUNCTION_NAMESPACE, "ok_tool"),
            ),
            /*terminal_outcome_reached*/ None,
        )
        .await?;
    let err = match registry
        .dispatch_any_with_terminal_outcome(
            test_invocation(
                Arc::clone(&session),
                Arc::clone(&turn),
                "failing-call",
                failing_tool.clone(),
            ),
            /*terminal_outcome_reached*/ None,
        )
        .await
    {
        Ok(_) => panic!("failing handler should return an error"),
        Err(err) => err,
    };
    assert_eq!(err.to_string(), "handler failed");

    let expected = vec![
        RecordedToolLifecycle::Start {
            call_id: "ok-call".to_string(),
            tool_name: ok_tool.clone().with_default_namespace(),
            root_turn_id: Some("root-turn".to_string()),
        },
        RecordedToolLifecycle::Finish {
            call_id: "ok-call".to_string(),
            tool_name: ok_tool.with_default_namespace(),
            outcome: codex_extension_api::ToolCallOutcome::Completed { success: false },
        },
        RecordedToolLifecycle::Start {
            call_id: "failing-call".to_string(),
            tool_name: failing_tool.clone(),
            root_turn_id: Some("root-turn".to_string()),
        },
        RecordedToolLifecycle::Finish {
            call_id: "failing-call".to_string(),
            tool_name: failing_tool,
            outcome: codex_extension_api::ToolCallOutcome::Failed {
                handler_executed: true,
            },
        },
    ];
    let actual = records
        .lock()
        .unwrap_or_else(std::sync::PoisonError::into_inner)
        .drain(..)
        .collect::<Vec<_>>();
    assert_eq!(expected, actual);

    Ok(())
}

fn test_invocation(
    session: Arc<crate::session::session::Session>,
    turn: Arc<crate::session::turn_context::TurnContext>,
    call_id: &str,
    tool_name: codex_tools::ToolName,
) -> ToolInvocation {
    let step_context = StepContext::for_test(Arc::clone(&turn));
    ToolInvocation {
        session,
        step_context,
        turn,
        cancellation_token: tokio_util::sync::CancellationToken::new(),
        tracker: Arc::new(tokio::sync::Mutex::new(
            crate::turn_diff_tracker::TurnDiffTracker::new(),
        )),
        call_id: call_id.to_string(),
        tool_name,
        source: crate::tools::context::ToolCallSource::Direct,
        payload: ToolPayload::Function {
            arguments: "{}".to_string(),
        },
    }
}