use std::sync::Arc;use std::sync::Weak;use codex_protocol::items::TurnItem;use codex_tools::ConversationHistory;use codex_tools::ExtensionTurnItem;use codex_tools::ToolCall as ExtensionToolCall;use codex_tools::ToolEnvironment;use codex_tools::ToolName;use codex_tools::ToolSearchInfo;use codex_tools::ToolSpec;use codex_tools::TurnItemEmissionFuture;use codex_tools::TurnItemEmitter;use crate::sandboxing::SandboxPermissions;use crate::session::session::Session;use crate::session::turn_context::TurnContext;use crate::stream_events_utils::TurnItemContributorPolicy;use crate::stream_events_utils::finalize_turn_item;use crate::tools::context::ToolInvocation;use crate::tools::context::ToolPayload;use crate::tools::handlers::apply_granted_turn_permissions;use crate::tools::registry::CoreToolRuntime;use crate::tools::registry::ToolExecutor;pub(crate) struct ExtensionToolAdapter(Arc<dyn codex_tools::ToolExecutor<ExtensionToolCall>>);impl ExtensionToolAdapter {    pub(crate) fn new(executor: Arc<dyn codex_tools::ToolExecutor<ExtensionToolCall>>) -> Self {        Self(executor)    }}impl ToolExecutor<ToolInvocation> for ExtensionToolAdapter {    fn tool_name(&self) -> ToolName {        self.0.tool_name()    }    fn spec(&self) -> ToolSpec {        self.0.spec()    }    fn exposure(&self) -> crate::tools::registry::ToolExposure {        self.0.exposure()    }    fn supports_parallel_tool_calls(&self) -> bool {        self.0.supports_parallel_tool_calls()    }    fn search_info(&self) -> Option<ToolSearchInfo> {        self.0.search_info()    }    fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> {        Box::pin(async move { self.0.handle(to_extension_call(&invocation).await).await })    }}impl CoreToolRuntime for ExtensionToolAdapter {    fn matches_kind(&self, payload: &ToolPayload) -> bool {        matches!(payload, ToolPayload::Function { .. })    }}struct CoreTurnItemEmitter {    session: Weak<Session>,    turn: Weak<TurnContext>,}fn extension_turn_item(item: ExtensionTurnItem) -> TurnItem {    match item {        ExtensionTurnItem::WebSearch(item) => TurnItem::WebSearch(item),        ExtensionTurnItem::ImageGeneration(mut item) => {            item.saved_path = None;            TurnItem::ImageGeneration(item)        }    }}impl TurnItemEmitter for CoreTurnItemEmitter {    fn emit_started<'a>(&'a self, item: ExtensionTurnItem) -> TurnItemEmissionFuture<'a> {        Box::pin(async move {            let (Some(session), Some(turn)) = (self.session.upgrade(), self.turn.upgrade()) else {                return;            };            session                .emit_turn_item_started(turn.as_ref(), &extension_turn_item(item))                .await;        })    }    fn emit_completed<'a>(&'a self, item: ExtensionTurnItem) -> TurnItemEmissionFuture<'a> {        Box::pin(async move {            let (Some(session), Some(turn)) = (self.session.upgrade(), self.turn.upgrade()) else {                return;            };            let mut item = extension_turn_item(item);            finalize_turn_item(                session.as_ref(),                turn.as_ref(),                TurnItemContributorPolicy::Run(turn.extension_data.as_ref()),                &mut item,                turn.collaboration_mode.mode == codex_protocol::config_types::ModeKind::Plan,            )            .await;            session.emit_turn_item_completed(turn.as_ref(), item).await;        })    }}async fn to_extension_call(invocation: &ToolInvocation) -> ExtensionToolCall {    let conversation_history =        ConversationHistory::new(invocation.session.clone_history().await.into_raw_items());    let mut environments = Vec::with_capacity(invocation.turn.environments.turn_environments.len());    for environment in &invocation.turn.environments.turn_environments {        // TODO(anp): Migrate extension ToolEnvironment and granted-permission lookup to PathUri        // so extensions can receive foreign environment cwd values.        let Ok(native_cwd) = environment.cwd().to_abs_path() else {            continue;        };        let additional_permissions = apply_granted_turn_permissions(            invocation.session.as_ref(),            &environment.environment_id,            native_cwd.as_path(),            SandboxPermissions::UseDefault,            /*additional_permissions*/ None,        )        .await        .additional_permissions;        let file_system_sandbox_context = invocation            .turn            .file_system_sandbox_context(additional_permissions, environment.cwd());        environments.push(ToolEnvironment {            environment_id: environment.environment_id.clone(),            cwd: native_cwd,            file_system: environment.environment.get_filesystem(),            file_system_sandbox_context,        });    }    ExtensionToolCall {        turn_id: invocation.turn.sub_id.clone(),        call_id: invocation.call_id.clone(),        tool_name: invocation.tool_name.clone(),        model: invocation.turn.model_info.slug.clone(),        truncation_policy: invocation.turn.truncation_policy,        conversation_history,        turn_item_emitter: Arc::new(CoreTurnItemEmitter {            session: Arc::downgrade(&invocation.session),            turn: Arc::downgrade(&invocation.turn),        }),        environments,        payload: invocation.payload.clone(),    }}#[cfg(test)]mod tests {    use std::sync::Arc;    use codex_extension_api::ExtensionData;    use codex_extension_api::TurnItemContributor;    use codex_protocol::items::TurnItem;    use codex_protocol::items::WebSearchItem;    use codex_protocol::models::ContentItem;    use codex_protocol::models::ResponseItem;    use codex_protocol::models::WebSearchAction;    use codex_protocol::protocol::EventMsg;    use codex_tools::ExtensionTurnItem;    use codex_utils_absolute_path::test_support::PathExt;    use codex_utils_absolute_path::test_support::test_path_buf;    use pretty_assertions::assert_eq;    use serde_json::json;    use tokio::sync::Mutex;    use super::CoreTurnItemEmitter;    use super::ExtensionToolAdapter;    use crate::tools::context::ToolCallSource;    use crate::tools::context::ToolInvocation;    use crate::tools::context::ToolPayload;    use crate::tools::hook_names::HookToolName;    use crate::tools::registry::CoreToolRuntime;    use crate::tools::registry::PostToolUsePayload;    use crate::tools::registry::PreToolUsePayload;    use crate::turn_diff_tracker::TurnDiffTracker;    struct StubExtensionExecutor;    impl codex_extension_api::ToolExecutor<codex_tools::ToolCall> for StubExtensionExecutor {        fn tool_name(&self) -> codex_tools::ToolName {            codex_tools::ToolName::plain("extension_echo")        }        fn spec(&self) -> codex_tools::ToolSpec {            codex_tools::ToolSpec::Function(codex_tools::ResponsesApiTool {                name: "extension_echo".to_string(),                description: "Echoes arguments.".to_string(),                strict: true,                parameters: codex_tools::parse_tool_input_schema(&json!({                    "type": "object",                    "properties": {                        "message": { "type": "string" },                    },                    "required": ["message"],                    "additionalProperties": false,                }))                .expect("extension schema should parse"),                output_schema: None,                defer_loading: None,            })        }        fn handle(&self, _call: codex_tools::ToolCall) -> codex_tools::ToolExecutorFuture<'_> {            Box::pin(async {                Ok(                    Box::new(codex_tools::JsonToolOutput::new(json!({ "ok": true })))                        as Box<dyn codex_tools::ToolOutput>,                )            })        }    }    struct CapturingExtensionExecutor {        captured_call: Arc<Mutex<Option<codex_tools::ToolCall>>>,    }    impl codex_extension_api::ToolExecutor<codex_tools::ToolCall> for CapturingExtensionExecutor {        fn tool_name(&self) -> codex_tools::ToolName {            codex_tools::ToolName::plain("extension_echo")        }        fn spec(&self) -> codex_tools::ToolSpec {            codex_tools::ToolSpec::Function(codex_tools::ResponsesApiTool {                name: "extension_echo".to_string(),                description: "Captures arguments.".to_string(),                strict: false,                parameters: codex_tools::JsonSchema::default(),                output_schema: None,                defer_loading: None,            })        }        fn handle(&self, call: codex_tools::ToolCall) -> codex_tools::ToolExecutorFuture<'_> {            Box::pin(self.handle_call(call))        }    }    impl CapturingExtensionExecutor {        async fn handle_call(            &self,            call: codex_tools::ToolCall,        ) -> Result<Box<dyn codex_tools::ToolOutput>, codex_tools::FunctionCallError> {            let item = ExtensionTurnItem::WebSearch(WebSearchItem {                id: call.call_id.clone(),                query: "rust trait object".to_string(),                action: WebSearchAction::Search {                    query: Some("rust trait object".to_string()),                    queries: None,                },            });            call.turn_item_emitter.emit_started(item.clone()).await;            call.turn_item_emitter.emit_completed(item).await;            *self.captured_call.lock().await = Some(call);            Ok(                Box::new(codex_tools::JsonToolOutput::new(json!({ "ok": true })))                    as Box<dyn codex_tools::ToolOutput>,            )        }    }    #[tokio::test]    async fn exposes_generic_hook_payloads() {        let handler = ExtensionToolAdapter::new(Arc::new(StubExtensionExecutor));        let (session, turn) = crate::session::tests::make_session_and_context().await;        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(),            },        };        let output = codex_tools::JsonToolOutput::new(json!({ "ok": true }));        assert_eq!(            CoreToolRuntime::pre_tool_use_payload(&handler, &invocation),            Some(PreToolUsePayload {                tool_name: HookToolName::new("extension_echo"),                tool_input: json!({ "message": "hello" }),            })        );        assert_eq!(            CoreToolRuntime::post_tool_use_payload(&handler, &invocation, &output),            Some(PostToolUsePayload {                tool_name: HookToolName::new("extension_echo"),                tool_use_id: "call-extension".to_string(),                tool_input: json!({ "message": "hello" }),                tool_response: json!({ "ok": true }),            })        );    }    #[tokio::test]    async fn passes_turn_fields_and_scoped_turn_item_emitter_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, rx) = crate::session::tests::make_session_and_context_with_rx().await;        let weak_session = Arc::downgrade(&session);        let weak_turn = Arc::downgrade(&turn);        let turn_id = turn.sub_id.clone();        let model = turn.model_info.slug.clone();        let truncation_policy = turn.truncation_policy;        let expected_sandbox_cwds = turn            .environments            .turn_environments            .iter()            .map(|environment| Some(environment.cwd().clone()))            .collect::<Vec<_>>();        let history_item = ResponseItem::Message {            id: None,            role: "user".to_string(),            content: vec![ContentItem::InputText {                text: "extension history".to_string(),            }],            phase: None,            metadata: None,        };        session            .record_conversation_items(&turn, std::slice::from_ref(&history_item))            .await;        let raw_history_event = rx.recv().await.expect("history raw response item event");        let EventMsg::RawResponseItem(raw_history_item) = raw_history_event.msg else {            panic!("expected raw response item event");        };        assert_eq!(raw_history_item.item, history_item);        let invocation = ToolInvocation {            session,            turn,            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!(weak_session.upgrade().is_none());        assert!(weak_turn.upgrade().is_none());        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.model, model);        assert_eq!(captured_call.truncation_policy, truncation_policy);        assert_eq!(            captured_call                .environments                .iter()                .map(|environment| environment.file_system_sandbox_context.cwd.clone())                .collect::<Vec<_>>(),            expected_sandbox_cwds        );        assert_eq!(            captured_call.conversation_history.items(),            std::slice::from_ref(&history_item)        );        match captured_call.payload {            ToolPayload::Function { arguments } => {                assert_eq!(arguments, json!({ "message": "hello" }).to_string());            }            payload => panic!("expected function payload, got {payload:?}"),        }        let started = rx.recv().await.expect("item started event");        let EventMsg::ItemStarted(started) = started.msg else {            panic!("expected item started event");        };        let TurnItem::WebSearch(started_item) = started.item else {            panic!("expected web search item");        };        let begin = rx.recv().await.expect("legacy web search begin event");        let EventMsg::WebSearchBegin(begin) = begin.msg else {            panic!("expected legacy web search begin event");        };        let completed = rx.recv().await.expect("item completed event");        let EventMsg::ItemCompleted(completed) = completed.msg else {            panic!("expected item completed event");        };        let TurnItem::WebSearch(completed_item) = completed.item else {            panic!("expected web search item");        };        let end = rx.recv().await.expect("legacy web search end event");        let EventMsg::WebSearchEnd(end) = end.msg else {            panic!("expected legacy web search end event");        };        let expected = WebSearchItem {            id: "call-extension".to_string(),            query: "rust trait object".to_string(),            action: WebSearchAction::Search {                query: Some("rust trait object".to_string()),                queries: None,            },        };        assert_eq!(started_item, expected);        assert_eq!(completed_item, expected);        assert_eq!(begin.call_id, expected.id);        assert_eq!(end.call_id, expected.id);        assert_eq!(end.query, expected.query);        assert_eq!(end.action, expected.action);    }    struct ImageGenerationExtensionExecutor;    #[derive(Debug)]    struct ExtensionTurnItemContributorRan;    struct RecordExtensionTurnItemContributor;    impl TurnItemContributor for RecordExtensionTurnItemContributor {        fn contribute<'a>(            &'a self,            _thread_store: &'a ExtensionData,            turn_store: &'a ExtensionData,            _item: &'a mut TurnItem,        ) -> codex_extension_api::ExtensionFuture<'a, Result<(), String>> {            Box::pin(async move {                turn_store.insert(ExtensionTurnItemContributorRan);                Ok(())            })        }    }    #[tokio::test]    async fn extension_completion_runs_turn_item_contributors() {        let (mut session, turn) = crate::session::tests::make_session_and_context().await;        let mut builder = codex_extension_api::ExtensionRegistryBuilder::new();        builder.turn_item_contributor(Arc::new(RecordExtensionTurnItemContributor));        session.services.extensions = Arc::new(builder.build());        let session = Arc::new(session);        let turn = Arc::new(turn);        let emitter = CoreTurnItemEmitter {            session: Arc::downgrade(&session),            turn: Arc::downgrade(&turn),        };        codex_tools::TurnItemEmitter::emit_completed(            &emitter,            ExtensionTurnItem::WebSearch(WebSearchItem {                id: "search-1".to_string(),                query: "contributors".to_string(),                action: WebSearchAction::Other,            }),        )        .await;        assert!(            turn.extension_data                .get::<ExtensionTurnItemContributorRan>()                .is_some()        );    }    impl codex_extension_api::ToolExecutor<codex_tools::ToolCall> for ImageGenerationExtensionExecutor {        fn tool_name(&self) -> codex_tools::ToolName {            codex_tools::ToolName::namespaced("image_gen", "imagegen")        }        fn spec(&self) -> codex_tools::ToolSpec {            codex_tools::ToolSpec::Function(codex_tools::ResponsesApiTool {                name: "imagegen".to_string(),                description: "Generates an image.".to_string(),                strict: false,                parameters: codex_tools::JsonSchema::default(),                output_schema: None,                defer_loading: None,            })        }        fn handle(&self, call: codex_tools::ToolCall) -> codex_tools::ToolExecutorFuture<'_> {            Box::pin(self.handle_call(call))        }    }    impl ImageGenerationExtensionExecutor {        async fn handle_call(            &self,            call: codex_tools::ToolCall,        ) -> Result<Box<dyn codex_tools::ToolOutput>, codex_tools::FunctionCallError> {            call.turn_item_emitter                .emit_started(ExtensionTurnItem::ImageGeneration(                    codex_protocol::items::ImageGenerationItem {                        id: call.call_id.clone(),                        status: "in_progress".to_string(),                        revised_prompt: None,                        result: String::new(),                        saved_path: None,                    },                ))                .await;            call.turn_item_emitter                .emit_completed(ExtensionTurnItem::ImageGeneration(                    codex_protocol::items::ImageGenerationItem {                        id: call.call_id,                        status: "completed".to_string(),                        revised_prompt: Some("A tiny blue square".to_string()),                        result: "cG5n".to_string(),                        saved_path: Some(test_path_buf("/tmp/extension-claimed.png").abs()),                    },                ))                .await;            Ok(                Box::new(codex_tools::JsonToolOutput::new(json!({ "ok": true })))                    as Box<dyn codex_tools::ToolOutput>,            )        }    }    #[tokio::test]    async fn image_generation_publication_is_finalized_by_core() {        let handler = ExtensionToolAdapter::new(Arc::new(ImageGenerationExtensionExecutor));        let (session, turn, rx) = crate::session::tests::make_session_and_context_with_rx().await;        let expected_path = crate::stream_events_utils::image_generation_artifact_path(            &turn.config.codex_home,            &session.thread_id.to_string(),            "call-image",        );        let invocation = ToolInvocation {            session,            turn,            cancellation_token: tokio_util::sync::CancellationToken::new(),            tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())),            call_id: "call-image".to_string(),            tool_name: codex_tools::ToolName::namespaced("image_gen", "imagegen"),            source: ToolCallSource::Direct,            payload: ToolPayload::Function {                arguments: "{}".to_string(),            },        };        crate::tools::registry::ToolExecutor::handle(&handler, invocation)            .await            .expect("extension call should succeed");        let started = rx.recv().await.expect("item started event");        let EventMsg::ItemStarted(started) = started.msg else {            panic!("expected item started event");        };        let TurnItem::ImageGeneration(started_item) = started.item else {            panic!("expected image generation item");        };        let begin = rx.recv().await.expect("legacy image start event");        assert!(matches!(begin.msg, EventMsg::ImageGenerationBegin(_)));        let completed = rx.recv().await.expect("item completed event");        let EventMsg::ItemCompleted(completed) = completed.msg else {            panic!("expected item completed event");        };        let TurnItem::ImageGeneration(completed_item) = completed.item else {            panic!("expected image generation item");        };        let end = rx.recv().await.expect("legacy image end event");        assert!(matches!(end.msg, EventMsg::ImageGenerationEnd(_)));        assert_eq!(            started_item,            codex_protocol::items::ImageGenerationItem {                id: "call-image".to_string(),                status: "in_progress".to_string(),                revised_prompt: None,                result: String::new(),                saved_path: None,            }        );        assert_eq!(            completed_item,            codex_protocol::items::ImageGenerationItem {                id: "call-image".to_string(),                status: "completed".to_string(),                revised_prompt: Some("A tiny blue square".to_string()),                result: "cG5n".to_string(),                saved_path: Some(expected_path.clone()),            }        );        assert_eq!(            std::fs::read(&expected_path).expect("generated artifact should be saved"),            b"png"        );    }}