diff --git a/codex-rs/app-server-protocol/schema/json/ClientRequest.json b/codex-rs/app-server-protocol/schema/json/ClientRequest.json index 84fd8014b0..988dea0dbc 100644 --- a/codex-rs/app-server-protocol/schema/json/ClientRequest.json +++ b/codex-rs/app-server-protocol/schema/json/ClientRequest.json @@ -501,6 +501,10 @@ "description": { "type": "string" }, + "injectIntoContext": { + "default": true, + "type": "boolean" + }, "inputSchema": true, "name": { "type": "string" @@ -2795,6 +2799,33 @@ ], "type": "object" }, + "ToolProviderToolSpec": { + "properties": { + "defaultTimeoutMs": { + "format": "uint64", + "minimum": 0.0, + "type": "integer" + }, + "description": { + "type": "string" + }, + "injectIntoContext": { + "type": "boolean" + }, + "inputSchema": true, + "name": { + "type": "string" + } + }, + "required": [ + "defaultTimeoutMs", + "description", + "injectIntoContext", + "inputSchema", + "name" + ], + "type": "object" + }, "TurnInterruptParams": { "properties": { "threadId": { diff --git a/codex-rs/app-server-protocol/schema/json/EventMsg.json b/codex-rs/app-server-protocol/schema/json/EventMsg.json index dbf9fc8e9f..1ab9f4f553 100644 --- a/codex-rs/app-server-protocol/schema/json/EventMsg.json +++ b/codex-rs/app-server-protocol/schema/json/EventMsg.json @@ -1975,6 +1975,13 @@ "callId": { "type": "string" }, + "providerOwned": { + "default": false, + "type": "boolean" + }, + "threadId": { + "type": "string" + }, "tool": { "type": "string" }, @@ -1992,6 +1999,7 @@ "required": [ "arguments", "callId", + "threadId", "tool", "turnId", "type" @@ -8112,6 +8120,13 @@ "callId": { "type": "string" }, + "providerOwned": { + "default": false, + "type": "boolean" + }, + "threadId": { + "type": "string" + }, "tool": { "type": "string" }, @@ -8129,6 +8144,7 @@ "required": [ "arguments", "callId", + "threadId", "tool", "turnId", "type" diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json index 421423bbc0..47f7c63e74 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json @@ -3338,6 +3338,13 @@ "callId": { "type": "string" }, + "providerOwned": { + "default": false, + "type": "boolean" + }, + "threadId": { + "type": "string" + }, "tool": { "type": "string" }, @@ -3355,6 +3362,7 @@ "required": [ "arguments", "callId", + "threadId", "tool", "turnId", "type" @@ -11032,6 +11040,10 @@ "description": { "type": "string" }, + "injectIntoContext": { + "default": true, + "type": "boolean" + }, "inputSchema": true, "name": { "type": "string" @@ -16997,6 +17009,33 @@ ], "type": "object" }, + "ToolProviderToolSpec": { + "properties": { + "defaultTimeoutMs": { + "format": "uint64", + "minimum": 0.0, + "type": "integer" + }, + "description": { + "type": "string" + }, + "injectIntoContext": { + "type": "boolean" + }, + "inputSchema": true, + "name": { + "type": "string" + } + }, + "required": [ + "defaultTimeoutMs", + "description", + "injectIntoContext", + "inputSchema", + "name" + ], + "type": "object" + }, "ToolsV2": { "properties": { "view_image": { diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json index f7a0dbb478..dc9c75f273 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json @@ -3671,6 +3671,10 @@ "description": { "type": "string" }, + "injectIntoContext": { + "default": true, + "type": "boolean" + }, "inputSchema": true, "name": { "type": "string" @@ -5140,6 +5144,13 @@ "callId": { "type": "string" }, + "providerOwned": { + "default": false, + "type": "boolean" + }, + "threadId": { + "type": "string" + }, "tool": { "type": "string" }, @@ -5157,6 +5168,7 @@ "required": [ "arguments", "callId", + "threadId", "tool", "turnId", "type" @@ -14818,6 +14830,33 @@ ], "type": "object" }, + "ToolProviderToolSpec": { + "properties": { + "defaultTimeoutMs": { + "format": "uint64", + "minimum": 0.0, + "type": "integer" + }, + "description": { + "type": "string" + }, + "injectIntoContext": { + "type": "boolean" + }, + "inputSchema": true, + "name": { + "type": "string" + } + }, + "required": [ + "defaultTimeoutMs", + "description", + "injectIntoContext", + "inputSchema", + "name" + ], + "type": "object" + }, "ToolsV2": { "properties": { "view_image": { diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadStartParams.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadStartParams.json index 630176f8c1..5fd8ee8e54 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadStartParams.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadStartParams.json @@ -56,6 +56,10 @@ "description": { "type": "string" }, + "injectIntoContext": { + "default": true, + "type": "boolean" + }, "inputSchema": true, "name": { "type": "string" diff --git a/codex-rs/app-server-protocol/schema/typescript/DynamicToolCallRequest.ts b/codex-rs/app-server-protocol/schema/typescript/DynamicToolCallRequest.ts index 94b0c65c66..d73baa52d1 100644 --- a/codex-rs/app-server-protocol/schema/typescript/DynamicToolCallRequest.ts +++ b/codex-rs/app-server-protocol/schema/typescript/DynamicToolCallRequest.ts @@ -3,4 +3,4 @@ // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. import type { JsonValue } from "./serde_json/JsonValue"; -export type DynamicToolCallRequest = { callId: string, turnId: string, tool: string, arguments: JsonValue, }; +export type DynamicToolCallRequest = { threadId: string, callId: string, turnId: string, tool: string, arguments: JsonValue, providerOwned: boolean, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/DynamicToolSpec.ts b/codex-rs/app-server-protocol/schema/typescript/v2/DynamicToolSpec.ts index 8b39793f3f..bb9191cf28 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/DynamicToolSpec.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/DynamicToolSpec.ts @@ -3,4 +3,4 @@ // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. import type { JsonValue } from "../serde_json/JsonValue"; -export type DynamicToolSpec = { name: string, description: string, inputSchema: JsonValue, }; +export type DynamicToolSpec = { name: string, description: string, inputSchema: JsonValue, injectIntoContext: boolean, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/ToolProviderToolSpec.ts b/codex-rs/app-server-protocol/schema/typescript/v2/ToolProviderToolSpec.ts new file mode 100644 index 0000000000..d10651a3f8 --- /dev/null +++ b/codex-rs/app-server-protocol/schema/typescript/v2/ToolProviderToolSpec.ts @@ -0,0 +1,6 @@ +// GENERATED CODE! DO NOT MODIFY BY HAND! + +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +import type { JsonValue } from "../serde_json/JsonValue"; + +export type ToolProviderToolSpec = { name: string, description: string, inputSchema: JsonValue, defaultTimeoutMs: bigint, injectIntoContext: boolean, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/index.ts b/codex-rs/app-server-protocol/schema/typescript/v2/index.ts index b57daaac3a..68777ea1cc 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/index.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/index.ts @@ -278,6 +278,7 @@ export type { ThreadUnsubscribeParams } from "./ThreadUnsubscribeParams"; export type { ThreadUnsubscribeResponse } from "./ThreadUnsubscribeResponse"; export type { ThreadUnsubscribeStatus } from "./ThreadUnsubscribeStatus"; export type { TokenUsageBreakdown } from "./TokenUsageBreakdown"; +export type { ToolProviderToolSpec } from "./ToolProviderToolSpec"; export type { ToolRequestUserInputAnswer } from "./ToolRequestUserInputAnswer"; export type { ToolRequestUserInputOption } from "./ToolRequestUserInputOption"; export type { ToolRequestUserInputParams } from "./ToolRequestUserInputParams"; diff --git a/codex-rs/app-server-protocol/src/protocol/common.rs b/codex-rs/app-server-protocol/src/protocol/common.rs index bb3c486ee0..0434d1ab1e 100644 --- a/codex-rs/app-server-protocol/src/protocol/common.rs +++ b/codex-rs/app-server-protocol/src/protocol/common.rs @@ -320,6 +320,16 @@ client_request_definitions! { params: v2::PluginUninstallParams, response: v2::PluginUninstallResponse, }, + #[experimental("toolProvider/register")] + ToolProviderRegister => "toolProvider/register" { + params: v2::ToolProviderRegisterParams, + response: v2::ToolProviderRegisterResponse, + }, + #[experimental("toolProvider/unregister")] + ToolProviderUnregister => "toolProvider/unregister" { + params: v2::ToolProviderUnregisterParams, + response: v2::ToolProviderUnregisterResponse, + }, TurnStart => "turn/start" { params: v2::TurnStartParams, inspect_params: true, diff --git a/codex-rs/app-server-protocol/src/protocol/thread_history.rs b/codex-rs/app-server-protocol/src/protocol/thread_history.rs index 38a8c64968..2a819919be 100644 --- a/codex-rs/app-server-protocol/src/protocol/thread_history.rs +++ b/codex-rs/app-server-protocol/src/protocol/thread_history.rs @@ -1739,10 +1739,12 @@ mod tests { }), EventMsg::DynamicToolCallRequest( codex_protocol::dynamic_tools::DynamicToolCallRequest { + thread_id: "thread-1".into(), call_id: "dyn-1".into(), turn_id: "turn-1".into(), tool: "lookup_ticket".into(), arguments: serde_json::json!({"id":"ABC-123"}), + provider_owned: false, }, ), EventMsg::DynamicToolCallResponse(DynamicToolCallResponseEvent { diff --git a/codex-rs/app-server-protocol/src/protocol/v2.rs b/codex-rs/app-server-protocol/src/protocol/v2.rs index aaf0484a58..cfbb496b09 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2.rs @@ -509,6 +509,47 @@ pub struct DynamicToolSpec { pub name: String, pub description: String, pub input_schema: JsonValue, + #[serde(default = "default_dynamic_tool_inject_into_context")] + pub inject_into_context: bool, +} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct ToolProviderToolSpec { + pub name: String, + pub description: String, + pub input_schema: JsonValue, + pub default_timeout_ms: u64, + pub inject_into_context: bool, +} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct ToolProviderRegisterParams { + pub tool: ToolProviderToolSpec, +} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct ToolProviderRegisterResponse {} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct ToolProviderUnregisterParams { + pub name: String, +} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct ToolProviderUnregisterResponse {} + +fn default_dynamic_tool_inject_into_context() -> bool { + true } #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS, ExperimentalApi)] diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 1fec3a35db..92b8c87417 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -934,6 +934,18 @@ If the session approval policy uses `Reject` with `request_permissions: true`, t `dynamicTools` on `thread/start` and the corresponding `item/tool/call` request/response flow are experimental APIs. To enable them, set `initialize.params.capabilities.experimentalApi = true`. +Clients can also register connection-owned dynamic tools with the experimental `toolProvider/register` and `toolProvider/unregister` requests over either `stdio://` or websocket transports. Registered tool-provider specs include: + +- `name` +- `description` +- `inputSchema` +- `defaultTimeoutMs` +- `injectIntoContext` + +If `injectIntoContext` is `true`, the tool is exposed to the model as a function tool for future turns. If it is `false`, the tool remains callable only through `js_repl` via `codex.tool(...)` and is not injected into model context. + +For provider-owned tools, the app-server routes `item/tool/call` only to the connection that registered the tool and applies that tool's `defaultTimeoutMs` while waiting for the client response. + When a dynamic tool is invoked during a turn, the server sends an `item/tool/call` JSON-RPC request to the client: ```json diff --git a/codex-rs/app-server/src/bespoke_event_handling.rs b/codex-rs/app-server/src/bespoke_event_handling.rs index 620f397af9..460528a01d 100644 --- a/codex-rs/app-server/src/bespoke_event_handling.rs +++ b/codex-rs/app-server/src/bespoke_event_handling.rs @@ -12,6 +12,7 @@ use crate::thread_state::ThreadState; use crate::thread_state::TurnSummary; use crate::thread_status::ThreadWatchActiveGuard; use crate::thread_status::ThreadWatchManager; +use crate::tool_provider::ToolProviderRegistry; use codex_app_server_protocol::AccountRateLimitsUpdatedNotification; use codex_app_server_protocol::AdditionalPermissionProfile as V2AdditionalPermissionProfile; use codex_app_server_protocol::AgentMessageDeltaNotification; @@ -189,6 +190,7 @@ pub(crate) async fn apply_bespoke_event_handling( conversation: Arc, thread_manager: Arc, outgoing: ThreadScopedOutgoingMessageSender, + tool_provider_registry: ToolProviderRegistry, thread_state: Arc>, thread_watch_manager: ThreadWatchManager, api_version: ApiVersion, @@ -742,6 +744,7 @@ pub(crate) async fn apply_bespoke_event_handling( let turn_id = request.turn_id; let tool = request.tool; let arguments = request.arguments; + let provider_owned = request.provider_owned; let item = ThreadItem::DynamicToolCall { id: call_id.clone(), tool: tool.clone(), @@ -766,11 +769,40 @@ pub(crate) async fn apply_bespoke_event_handling( tool: tool.clone(), arguments: arguments.clone(), }; - let (_pending_request_id, rx) = outgoing - .send_request(ServerRequestPayload::DynamicToolCall(params)) - .await; + let (pending_request_id, rx, response_timeout) = if provider_owned { + let Some(target) = tool_provider_registry.lookup_target(&tool).await else { + crate::dynamic_tools::submit_failed_tool_response( + call_id, + "dynamic tool provider is unavailable", + conversation, + ) + .await; + return; + }; + let (pending_request_id, rx) = outgoing + .send_request_to_connection( + target.connection_id, + ServerRequestPayload::DynamicToolCall(params), + ) + .await; + (pending_request_id, rx, Some(target.default_timeout)) + } else { + let (pending_request_id, rx) = outgoing + .send_request(ServerRequestPayload::DynamicToolCall(params)) + .await; + (pending_request_id, rx, None) + }; + let outgoing_for_response = outgoing.clone(); tokio::spawn(async move { - crate::dynamic_tools::on_call_response(call_id, rx, conversation).await; + crate::dynamic_tools::on_call_response( + call_id, + pending_request_id, + rx, + conversation, + outgoing_for_response, + response_timeout, + ) + .await; }); } else { error!( diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index a4909c7745..82f9d0723c 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -160,6 +160,10 @@ use codex_app_server_protocol::ThreadUnarchivedNotification; use codex_app_server_protocol::ThreadUnsubscribeParams; use codex_app_server_protocol::ThreadUnsubscribeResponse; use codex_app_server_protocol::ThreadUnsubscribeStatus; +use codex_app_server_protocol::ToolProviderRegisterParams; +use codex_app_server_protocol::ToolProviderRegisterResponse; +use codex_app_server_protocol::ToolProviderUnregisterParams; +use codex_app_server_protocol::ToolProviderUnregisterResponse; use codex_app_server_protocol::Turn; use codex_app_server_protocol::TurnInterruptParams; use codex_app_server_protocol::TurnStartParams; @@ -310,6 +314,7 @@ use crate::filters::source_kind_matches; use crate::thread_state::ThreadListenerCommand; use crate::thread_state::ThreadState; use crate::thread_state::ThreadStateManager; +use crate::tool_provider::ToolProviderRegistry; const THREAD_LIST_DEFAULT_LIMIT: usize = 25; const THREAD_LIST_MAX_LIMIT: usize = 100; @@ -384,6 +389,7 @@ pub(crate) struct CodexMessageProcessor { thread_state_manager: ThreadStateManager, thread_watch_manager: ThreadWatchManager, command_exec_manager: CommandExecManager, + tool_provider_registry: ToolProviderRegistry, pending_fuzzy_searches: Arc>>>, fuzzy_search_sessions: Arc>>, feedback: CodexFeedback, @@ -404,6 +410,7 @@ struct ListenerTaskContext { thread_state_manager: ThreadStateManager, outgoing: Arc, thread_watch_manager: ThreadWatchManager, + tool_provider_registry: ToolProviderRegistry, fallback_model_provider: String, codex_home: PathBuf, } @@ -424,6 +431,7 @@ pub(crate) struct CodexMessageProcessorArgs { pub(crate) cloud_requirements: Arc>, pub(crate) feedback: CodexFeedback, pub(crate) log_db: Option, + pub(crate) tool_provider_registry: ToolProviderRegistry, } impl CodexMessageProcessor { @@ -484,6 +492,7 @@ impl CodexMessageProcessor { cloud_requirements, feedback, log_db, + tool_provider_registry, } = args; Self { auth_manager, @@ -498,6 +507,7 @@ impl CodexMessageProcessor { thread_state_manager: ThreadStateManager::new(), thread_watch_manager: ThreadWatchManager::new_with_outgoing(outgoing), command_exec_manager: CommandExecManager::default(), + tool_provider_registry, pending_fuzzy_searches: Arc::new(Mutex::new(HashMap::new())), fuzzy_search_sessions: Arc::new(Mutex::new(HashMap::new())), feedback, @@ -730,6 +740,14 @@ impl CodexMessageProcessor { self.plugin_uninstall(to_connection_request_id(request_id), params) .await; } + ClientRequest::ToolProviderRegister { request_id, params } => { + self.tool_provider_register(to_connection_request_id(request_id), params) + .await; + } + ClientRequest::ToolProviderUnregister { request_id, params } => { + self.tool_provider_unregister(to_connection_request_id(request_id), params) + .await; + } ClientRequest::TurnStart { request_id, params } => { self.turn_start( to_connection_request_id(request_id), @@ -1844,6 +1862,7 @@ impl CodexMessageProcessor { thread_state_manager: self.thread_state_manager.clone(), outgoing: Arc::clone(&self.outgoing), thread_watch_manager: self.thread_watch_manager.clone(), + tool_provider_registry: self.tool_provider_registry.clone(), fallback_model_provider: self.config.model_provider_id.clone(), codex_home: self.config.codex_home.clone(), }; @@ -1923,6 +1942,8 @@ impl CodexMessageProcessor { name: tool.name, description: tool.description, input_schema: tool.input_schema, + inject_into_context: tool.inject_into_context, + provider_owned: false, }) .collect() }; @@ -1944,6 +1965,14 @@ impl CodexMessageProcessor { session_configured, .. } = new_conv; + thread + .update_registered_dynamic_tools( + listener_task_context + .tool_provider_registry + .snapshot_dynamic_tools() + .await, + ) + .await; let config_snapshot = thread.config_snapshot().await; let mut thread = build_thread_from_snapshot( thread_id, @@ -3149,6 +3178,23 @@ impl CodexMessageProcessor { self.thread_manager.subscribe_thread_created() } + async fn sync_registered_dynamic_tools_for_thread(&self, thread: &CodexThread) { + thread + .update_registered_dynamic_tools( + self.tool_provider_registry.snapshot_dynamic_tools().await, + ) + .await; + } + + async fn refresh_registered_dynamic_tools_for_loaded_threads(&self) { + let dynamic_tools = self.tool_provider_registry.snapshot_dynamic_tools().await; + for thread in self.thread_manager.loaded_threads().await { + thread + .update_registered_dynamic_tools(dynamic_tools.clone()) + .await; + } + } + pub(crate) async fn connection_initialized(&self, connection_id: ConnectionId) { self.thread_state_manager .connection_initialized(connection_id) @@ -3162,6 +3208,14 @@ impl CodexMessageProcessor { self.thread_state_manager .remove_connection(connection_id) .await; + if self + .tool_provider_registry + .remove_connection(connection_id) + .await + { + self.refresh_registered_dynamic_tools_for_loaded_threads() + .await; + } } pub(crate) fn subscribe_running_assistant_turn_count(&self) -> watch::Receiver { @@ -3175,6 +3229,8 @@ impl CodexMessageProcessor { connection_ids: Vec, ) { if let Ok(thread) = self.thread_manager.get_thread(thread_id).await { + self.sync_registered_dynamic_tools_for_thread(thread.as_ref()) + .await; let config_snapshot = thread.config_snapshot().await; let loaded_thread = build_thread_from_snapshot(thread_id, &config_snapshot, thread.rollout_path()); @@ -3306,6 +3362,8 @@ impl CodexMessageProcessor { thread, session_configured, }) => { + self.sync_registered_dynamic_tools_for_thread(thread.as_ref()) + .await; let SessionConfiguredEvent { rollout_path, .. } = session_configured; let Some(rollout_path) = rollout_path else { self.send_internal_error( @@ -3850,6 +3908,8 @@ impl CodexMessageProcessor { return; } }; + self.sync_registered_dynamic_tools_for_thread(thread.as_ref()) + .await; // Auto-attach a conversation listener when forking a thread. Self::log_listener_attach_result( @@ -5749,6 +5809,52 @@ impl CodexMessageProcessor { } } + async fn tool_provider_register( + &self, + request_id: ConnectionRequestId, + params: ToolProviderRegisterParams, + ) { + match self + .tool_provider_registry + .register(request_id.connection_id, params.tool) + .await + { + Ok(()) => { + self.refresh_registered_dynamic_tools_for_loaded_threads() + .await; + self.outgoing + .send_response(request_id, ToolProviderRegisterResponse {}) + .await; + } + Err(message) => { + self.send_invalid_request_error(request_id, message).await; + } + } + } + + async fn tool_provider_unregister( + &self, + request_id: ConnectionRequestId, + params: ToolProviderUnregisterParams, + ) { + match self + .tool_provider_registry + .unregister(request_id.connection_id, ¶ms.name) + .await + { + Ok(()) => { + self.refresh_registered_dynamic_tools_for_loaded_threads() + .await; + self.outgoing + .send_response(request_id, ToolProviderUnregisterResponse {}) + .await; + } + Err(message) => { + self.send_invalid_request_error(request_id, message).await; + } + } + } + async fn turn_start( &self, request_id: ConnectionRequestId, @@ -5772,6 +5878,8 @@ impl CodexMessageProcessor { self.outgoing.send_error(request_id, error).await; return; } + self.sync_registered_dynamic_tools_for_thread(thread.as_ref()) + .await; let collaboration_modes_config = CollaborationModesConfig { default_mode_request_user_input: thread.enabled(Feature::DefaultModeRequestUserInput), @@ -6367,6 +6475,7 @@ impl CodexMessageProcessor { thread_state_manager: self.thread_state_manager.clone(), outgoing: Arc::clone(&self.outgoing), thread_watch_manager: self.thread_watch_manager.clone(), + tool_provider_registry: self.tool_provider_registry.clone(), fallback_model_provider: self.config.model_provider_id.clone(), codex_home: self.config.codex_home.clone(), }, @@ -6454,6 +6563,7 @@ impl CodexMessageProcessor { thread_state_manager: self.thread_state_manager.clone(), outgoing: Arc::clone(&self.outgoing), thread_watch_manager: self.thread_watch_manager.clone(), + tool_provider_registry: self.tool_provider_registry.clone(), fallback_model_provider: self.config.model_provider_id.clone(), codex_home: self.config.codex_home.clone(), }, @@ -6485,6 +6595,7 @@ impl CodexMessageProcessor { thread_manager, thread_state_manager, thread_watch_manager, + tool_provider_registry, fallback_model_provider, codex_home, } = listener_task_context; @@ -6569,6 +6680,7 @@ impl CodexMessageProcessor { conversation.clone(), thread_manager.clone(), thread_outgoing, + tool_provider_registry.clone(), thread_state.clone(), thread_watch_manager.clone(), api_version, @@ -7948,6 +8060,7 @@ mod tests { name: "my_tool".to_string(), description: "test".to_string(), input_schema: json!({"type": "null"}), + inject_into_context: true, }]; let err = validate_dynamic_tools(&tools).expect_err("invalid schema"); assert!(err.contains("my_tool"), "unexpected error: {err}"); @@ -7960,6 +8073,7 @@ mod tests { description: "test".to_string(), // Missing `type` is common; core sanitizes these to a supported schema. input_schema: json!({"properties": {}}), + inject_into_context: true, }]; validate_dynamic_tools(&tools).expect("valid schema"); } diff --git a/codex-rs/app-server/src/dynamic_tools.rs b/codex-rs/app-server/src/dynamic_tools.rs index c5e7550d9f..e9c177c472 100644 --- a/codex-rs/app-server/src/dynamic_tools.rs +++ b/codex-rs/app-server/src/dynamic_tools.rs @@ -1,22 +1,43 @@ use codex_app_server_protocol::DynamicToolCallOutputContentItem; use codex_app_server_protocol::DynamicToolCallResponse; +use codex_app_server_protocol::RequestId; use codex_core::CodexThread; use codex_protocol::dynamic_tools::DynamicToolCallOutputContentItem as CoreDynamicToolCallOutputContentItem; use codex_protocol::dynamic_tools::DynamicToolResponse as CoreDynamicToolResponse; use codex_protocol::protocol::Op; use std::sync::Arc; +use std::time::Duration; use tokio::sync::oneshot; use tracing::error; use crate::outgoing_message::ClientRequestResult; +use crate::outgoing_message::ThreadScopedOutgoingMessageSender; use crate::server_request_error::is_turn_transition_server_request_error; pub(crate) async fn on_call_response( call_id: String, + pending_request_id: RequestId, receiver: oneshot::Receiver, conversation: Arc, + outgoing: ThreadScopedOutgoingMessageSender, + response_timeout: Option, ) { - let response = receiver.await; + let response = match response_timeout { + Some(response_timeout) => match tokio::time::timeout(response_timeout, receiver).await { + Ok(response) => response, + Err(_) => { + let _canceled = outgoing.cancel_request(&pending_request_id).await; + submit_response( + call_id, + fallback_response("dynamic tool request timed out").0, + conversation, + ) + .await; + return; + } + }, + None => receiver.await, + }; let (response, _error) = match response { Ok(Ok(value)) => decode_response(value), Ok(Err(err)) if is_turn_transition_server_request_error(&err) => return, @@ -30,26 +51,15 @@ pub(crate) async fn on_call_response( } }; - let DynamicToolCallResponse { - content_items, - success, - } = response.clone(); - let core_response = CoreDynamicToolResponse { - content_items: content_items - .into_iter() - .map(CoreDynamicToolCallOutputContentItem::from) - .collect(), - success, - }; - if let Err(err) = conversation - .submit(Op::DynamicToolResponse { - id: call_id.clone(), - response: core_response, - }) - .await - { - error!("failed to submit DynamicToolResponse: {err}"); - } + submit_response(call_id, response, conversation).await; +} + +pub(crate) async fn submit_failed_tool_response( + call_id: String, + message: &str, + conversation: Arc, +) { + submit_response(call_id, fallback_response(message).0, conversation).await; } fn decode_response(value: serde_json::Value) -> (DynamicToolCallResponse, Option) { @@ -73,3 +83,30 @@ fn fallback_response(message: &str) -> (DynamicToolCallResponse, Option) Some(message.to_string()), ) } + +async fn submit_response( + call_id: String, + response: DynamicToolCallResponse, + conversation: Arc, +) { + let DynamicToolCallResponse { + content_items, + success, + } = response; + let core_response = CoreDynamicToolResponse { + content_items: content_items + .into_iter() + .map(CoreDynamicToolCallOutputContentItem::from) + .collect(), + success, + }; + if let Err(err) = conversation + .submit(Op::DynamicToolResponse { + id: call_id, + response: core_response, + }) + .await + { + error!("failed to submit DynamicToolResponse: {err}"); + } +} diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index 887c2c650c..6d95c414e4 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -73,6 +73,7 @@ mod outgoing_message; mod server_request_error; mod thread_state; mod thread_status; +mod tool_provider; mod transport; pub use crate::error_code::INPUT_TOO_LARGE_ERROR_CODE; diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 973a504310..82abd3466f 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -12,6 +12,7 @@ use crate::external_agent_config_api::ExternalAgentConfigApi; use crate::outgoing_message::ConnectionId; use crate::outgoing_message::ConnectionRequestId; use crate::outgoing_message::OutgoingMessageSender; +use crate::tool_provider::ToolProviderRegistry; use crate::transport::AppServerTransport; use async_trait::async_trait; use codex_app_server_protocol::ChatgptAuthTokensRefreshParams; @@ -203,6 +204,7 @@ impl MessageProcessor { .plugins_manager() .maybe_start_curated_repo_sync_for_config(&config); let cloud_requirements = Arc::new(RwLock::new(cloud_requirements)); + let tool_provider_registry = ToolProviderRegistry::default(); let codex_message_processor = CodexMessageProcessor::new(CodexMessageProcessorArgs { auth_manager, thread_manager: Arc::clone(&thread_manager), @@ -213,6 +215,7 @@ impl MessageProcessor { cloud_requirements: cloud_requirements.clone(), feedback, log_db, + tool_provider_registry, }); let config_api = ConfigApi::new( config.codex_home.clone(), diff --git a/codex-rs/app-server/src/outgoing_message.rs b/codex-rs/app-server/src/outgoing_message.rs index d615e51518..86a3a4aefd 100644 --- a/codex-rs/app-server/src/outgoing_message.rs +++ b/codex-rs/app-server/src/outgoing_message.rs @@ -92,6 +92,17 @@ impl ThreadScopedOutgoingMessageSender { .await } + pub(crate) async fn send_request_to_connection( + &self, + connection_id: ConnectionId, + payload: ServerRequestPayload, + ) -> (RequestId, oneshot::Receiver) { + let connection_ids = [connection_id]; + self.outgoing + .send_request_to_connections(Some(&connection_ids), payload, Some(self.thread_id)) + .await + } + pub(crate) async fn send_server_notification(&self, notification: ServerNotification) { if self.connection_ids.is_empty() { return; @@ -134,6 +145,10 @@ impl ThreadScopedOutgoingMessageSender { ) { self.outgoing.send_error(request_id, error).await; } + + pub(crate) async fn cancel_request(&self, request_id: &RequestId) -> bool { + self.outgoing.cancel_request(request_id).await + } } impl OutgoingMessageSender { diff --git a/codex-rs/app-server/src/tool_provider.rs b/codex-rs/app-server/src/tool_provider.rs new file mode 100644 index 0000000000..c8e83ca990 --- /dev/null +++ b/codex-rs/app-server/src/tool_provider.rs @@ -0,0 +1,142 @@ +use crate::outgoing_message::ConnectionId; +use codex_app_server_protocol::ToolProviderToolSpec; +use codex_protocol::dynamic_tools::DynamicToolSpec as CoreDynamicToolSpec; +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Duration; +use tokio::sync::Mutex; + +#[derive(Clone, Default)] +pub(crate) struct ToolProviderRegistry { + state: Arc>, +} + +#[derive(Default)] +struct ToolProviderRegistryState { + tools_by_name: HashMap, +} + +#[derive(Clone)] +struct RegisteredToolProviderTool { + connection_id: ConnectionId, + name: String, + description: String, + input_schema: serde_json::Value, + default_timeout: Duration, + inject_into_context: bool, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct ToolProviderTarget { + pub(crate) connection_id: ConnectionId, + pub(crate) default_timeout: Duration, +} + +impl ToolProviderRegistry { + pub(crate) async fn register( + &self, + connection_id: ConnectionId, + tool: ToolProviderToolSpec, + ) -> Result<(), String> { + validate_tool_provider_tool(&tool)?; + let mut state = self.state.lock().await; + if state.tools_by_name.contains_key(tool.name.as_str()) { + return Err(format!( + "dynamic tool name is already registered: {}", + tool.name + )); + } + state.tools_by_name.insert( + tool.name.clone(), + RegisteredToolProviderTool { + connection_id, + name: tool.name, + description: tool.description, + input_schema: tool.input_schema, + default_timeout: Duration::from_millis(tool.default_timeout_ms), + inject_into_context: tool.inject_into_context, + }, + ); + Ok(()) + } + + pub(crate) async fn unregister( + &self, + connection_id: ConnectionId, + name: &str, + ) -> Result<(), String> { + let mut state = self.state.lock().await; + let Some(existing) = state.tools_by_name.get(name) else { + return Err(format!("tool provider tool not found: {name}")); + }; + if existing.connection_id != connection_id { + return Err(format!("tool provider does not own tool: {name}")); + } + state.tools_by_name.remove(name); + Ok(()) + } + + pub(crate) async fn remove_connection(&self, connection_id: ConnectionId) -> bool { + let mut state = self.state.lock().await; + let before = state.tools_by_name.len(); + state + .tools_by_name + .retain(|_, tool| tool.connection_id != connection_id); + before != state.tools_by_name.len() + } + + pub(crate) async fn snapshot_dynamic_tools(&self) -> Vec { + let state = self.state.lock().await; + let mut tools = state + .tools_by_name + .values() + .map(|tool| CoreDynamicToolSpec { + name: tool.name.clone(), + description: tool.description.clone(), + input_schema: tool.input_schema.clone(), + inject_into_context: tool.inject_into_context, + provider_owned: true, + }) + .collect::>(); + tools.sort_by(|left, right| left.name.cmp(&right.name)); + tools + } + + pub(crate) async fn lookup_target(&self, tool_name: &str) -> Option { + let state = self.state.lock().await; + state + .tools_by_name + .get(tool_name) + .map(|tool| ToolProviderTarget { + connection_id: tool.connection_id, + default_timeout: tool.default_timeout, + }) + } +} + +fn validate_tool_provider_tool(tool: &ToolProviderToolSpec) -> Result<(), String> { + let name = tool.name.trim(); + if name.is_empty() { + return Err("dynamic tool name must not be empty".to_string()); + } + if name != tool.name { + return Err(format!( + "dynamic tool name has leading/trailing whitespace: {}", + tool.name + )); + } + if name == "mcp" || name.starts_with("mcp__") { + return Err(format!("dynamic tool name is reserved: {name}")); + } + if tool.default_timeout_ms == 0 { + return Err(format!( + "tool provider default timeout must be greater than zero for {name}" + )); + } + if let Err(err) = codex_core::parse_tool_input_schema(&tool.input_schema) { + return Err(format!( + "dynamic tool input schema is not supported for {name}: {err}" + )); + } + Ok(()) +} diff --git a/codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs b/codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs index 3a8ae92430..e8647e040b 100644 --- a/codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs +++ b/codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs @@ -3,6 +3,7 @@ use anyhow::Result; use anyhow::bail; use app_test_support::create_mock_responses_server_sequence_unchecked; use codex_app_server_protocol::ClientInfo; +use codex_app_server_protocol::InitializeCapabilities; use codex_app_server_protocol::InitializeParams; use codex_app_server_protocol::JSONRPCError; use codex_app_server_protocol::JSONRPCMessage; @@ -227,7 +228,10 @@ pub(super) async fn send_initialize_request( title: Some("WebSocket Test Client".to_string()), version: "0.1.0".to_string(), }, - capabilities: None, + capabilities: Some(InitializeCapabilities { + experimental_api: true, + ..Default::default() + }), }; send_request( stream, diff --git a/codex-rs/app-server/tests/suite/v2/dynamic_tools.rs b/codex-rs/app-server/tests/suite/v2/dynamic_tools.rs index 338593c465..cb04ecbc9a 100644 --- a/codex-rs/app-server/tests/suite/v2/dynamic_tools.rs +++ b/codex-rs/app-server/tests/suite/v2/dynamic_tools.rs @@ -18,6 +18,7 @@ use codex_app_server_protocol::ServerRequest; use codex_app_server_protocol::ThreadItem; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; +use codex_app_server_protocol::ToolProviderRegisterResponse; use codex_app_server_protocol::TurnStartParams; use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::UserInput as V2UserInput; @@ -61,6 +62,7 @@ async fn thread_start_injects_dynamic_tools_into_model_requests() -> Result<()> name: "demo_tool".to_string(), description: "Demo dynamic tool".to_string(), input_schema: input_schema.clone(), + inject_into_context: true, }; // Thread start injects dynamic tools into the thread's tool registry. @@ -118,6 +120,234 @@ async fn thread_start_injects_dynamic_tools_into_model_requests() -> Result<()> Ok(()) } +#[tokio::test] +async fn tool_provider_registration_routes_stdio_dynamic_tool_calls() -> Result<()> { + let call_id = "dyn-provider-call-1"; + let tool_name = "provider_tool"; + let tool_args = json!({ "city": "Paris" }); + let tool_call_arguments = serde_json::to_string(&tool_args)?; + let responses = vec![ + responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_function_call(call_id, tool_name, &tool_call_arguments), + responses::ev_completed("resp-1"), + ]), + create_final_assistant_message_sse_response("Done")?, + ]; + let server = create_mock_responses_server_sequence_unchecked(responses).await; + + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), &server.uri())?; + + let mut mcp = McpProcess::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let register_request_id = mcp + .send_raw_request( + "toolProvider/register", + Some(json!({ + "tool": { + "name": tool_name, + "description": "Demo provider tool", + "inputSchema": { + "type": "object", + "properties": { + "city": { "type": "string" } + }, + "required": ["city"], + "additionalProperties": false + }, + "defaultTimeoutMs": 1_000, + "injectIntoContext": true + } + })), + ) + .await?; + let register_response: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(register_request_id)), + ) + .await??; + let _: ToolProviderRegisterResponse = to_response(register_response)?; + + let thread_request_id = mcp + .send_thread_start_request(ThreadStartParams::default()) + .await?; + let thread_response: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(thread_request_id)), + ) + .await??; + let ThreadStartResponse { thread, .. } = to_response::(thread_response)?; + let thread_id = thread.id.clone(); + + let turn_request_id = mcp + .send_turn_start_request(TurnStartParams { + thread_id: thread_id.clone(), + input: vec![V2UserInput::Text { + text: "Run the provider tool".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }) + .await?; + let turn_response: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(turn_request_id)), + ) + .await??; + let TurnStartResponse { turn } = to_response::(turn_response)?; + let turn_id = turn.id.clone(); + + let started = wait_for_dynamic_tool_started(&mut mcp, call_id).await?; + assert_eq!(started.thread_id, thread_id.clone()); + assert_eq!(started.turn_id, turn_id.clone()); + + let request = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_request_message(), + ) + .await??; + let (request_id, params) = match request { + ServerRequest::DynamicToolCall { request_id, params } => (request_id, params), + other => panic!("expected DynamicToolCall request, got {other:?}"), + }; + assert_eq!( + params, + DynamicToolCallParams { + thread_id: thread_id.clone(), + turn_id: turn_id.clone(), + call_id: call_id.to_string(), + tool: tool_name.to_string(), + arguments: tool_args.clone(), + } + ); + + mcp.send_response( + request_id, + serde_json::to_value(DynamicToolCallResponse { + content_items: vec![DynamicToolCallOutputContentItem::InputText { + text: "dynamic-ok".to_string(), + }], + success: true, + })?, + ) + .await?; + + let completed = wait_for_dynamic_tool_completed(&mut mcp, call_id).await?; + assert_eq!(completed.thread_id, thread_id); + assert_eq!(completed.turn_id, turn_id); + + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("turn/completed"), + ) + .await??; + + let bodies = responses_bodies(&server).await?; + let tool = bodies + .iter() + .find_map(|body| find_tool(body, tool_name)) + .context("expected provider tool to be injected into request")?; + assert_eq!( + tool.get("description"), + Some(&Value::String("Demo provider tool".to_string())) + ); + + let payload = bodies + .iter() + .find_map(|body| function_call_output_payload(body, call_id)) + .context("expected function_call_output in follow-up request")?; + assert_eq!( + payload, + FunctionCallOutputPayload::from_text("dynamic-ok".to_string()) + ); + + Ok(()) +} + +#[tokio::test] +async fn tool_provider_registration_with_inject_disabled_stays_out_of_model_context() -> Result<()> +{ + let server = create_mock_responses_server_sequence_unchecked(vec![ + create_final_assistant_message_sse_response("Done")?, + ]) + .await; + + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), &server.uri())?; + + let mut mcp = McpProcess::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let register_request_id = mcp + .send_raw_request( + "toolProvider/register", + Some(json!({ + "tool": { + "name": "hidden_provider_tool", + "description": "Hidden provider tool", + "inputSchema": { + "type": "object", + "properties": {}, + "additionalProperties": false + }, + "defaultTimeoutMs": 1_000, + "injectIntoContext": false + } + })), + ) + .await?; + let register_response: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(register_request_id)), + ) + .await??; + let _: ToolProviderRegisterResponse = to_response(register_response)?; + + let thread_request_id = mcp + .send_thread_start_request(ThreadStartParams::default()) + .await?; + let thread_response: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(thread_request_id)), + ) + .await??; + let ThreadStartResponse { thread, .. } = to_response::(thread_response)?; + + let turn_request_id = mcp + .send_turn_start_request(TurnStartParams { + thread_id: thread.id, + input: vec![V2UserInput::Text { + text: "Hello".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }) + .await?; + let turn_response: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(turn_request_id)), + ) + .await??; + let _: TurnStartResponse = to_response::(turn_response)?; + + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("turn/completed"), + ) + .await??; + + let bodies = responses_bodies(&server).await?; + assert!( + bodies + .iter() + .all(|body| find_tool(body, "hidden_provider_tool").is_none()) + ); + + Ok(()) +} + /// Exercises the full dynamic tool call path (server request, client response, model output). #[tokio::test] async fn dynamic_tool_call_round_trip_sends_text_content_items_to_model() -> Result<()> { @@ -154,6 +384,7 @@ async fn dynamic_tool_call_round_trip_sends_text_content_items_to_model() -> Res "required": ["city"], "additionalProperties": false, }), + inject_into_context: true, }; let thread_req = mcp @@ -322,6 +553,7 @@ async fn dynamic_tool_call_round_trip_sends_content_items_to_model() -> Result<( "required": ["city"], "additionalProperties": false, }), + inject_into_context: true, }; let thread_req = mcp diff --git a/codex-rs/app-server/tests/suite/v2/mod.rs b/codex-rs/app-server/tests/suite/v2/mod.rs index 20e9758796..086e157d67 100644 --- a/codex-rs/app-server/tests/suite/v2/mod.rs +++ b/codex-rs/app-server/tests/suite/v2/mod.rs @@ -40,6 +40,7 @@ mod thread_start; mod thread_status; mod thread_unarchive; mod thread_unsubscribe; +mod tool_provider_websocket; mod turn_interrupt; mod turn_start; mod turn_start_zsh_fork; diff --git a/codex-rs/app-server/tests/suite/v2/tool_provider_websocket.rs b/codex-rs/app-server/tests/suite/v2/tool_provider_websocket.rs new file mode 100644 index 0000000000..1e00c642bd --- /dev/null +++ b/codex-rs/app-server/tests/suite/v2/tool_provider_websocket.rs @@ -0,0 +1,202 @@ +use super::connection_handling_websocket::DEFAULT_READ_TIMEOUT; +use super::connection_handling_websocket::connect_websocket; +use super::connection_handling_websocket::create_config_toml; +use super::connection_handling_websocket::read_jsonrpc_message; +use super::connection_handling_websocket::read_response_for_id; +use super::connection_handling_websocket::reserve_local_addr; +use super::connection_handling_websocket::send_initialize_request; +use super::connection_handling_websocket::send_request; +use super::connection_handling_websocket::spawn_websocket_server; +use anyhow::Context; +use anyhow::Result; +use anyhow::bail; +use app_test_support::create_final_assistant_message_sse_response; +use app_test_support::create_mock_responses_server_sequence_unchecked; +use app_test_support::to_response; +use codex_app_server_protocol::DynamicToolCallOutputContentItem; +use codex_app_server_protocol::DynamicToolCallParams; +use codex_app_server_protocol::DynamicToolCallResponse; +use codex_app_server_protocol::JSONRPCMessage; +use codex_app_server_protocol::JSONRPCResponse; +use codex_app_server_protocol::RequestId; +use codex_app_server_protocol::ServerRequest; +use codex_app_server_protocol::ThreadStartParams; +use codex_app_server_protocol::ThreadStartResponse; +use codex_app_server_protocol::ToolProviderRegisterResponse; +use codex_app_server_protocol::TurnStartParams; +use codex_app_server_protocol::TurnStartResponse; +use codex_app_server_protocol::UserInput as V2UserInput; +use core_test_support::responses; +use futures::SinkExt; +use serde_json::json; +use tempfile::TempDir; +use tokio::time::timeout; +use tokio_tungstenite::tungstenite::Message as WebSocketMessage; + +#[tokio::test] +async fn websocket_tool_provider_routes_dynamic_tool_calls_to_owner_connection() -> Result<()> { + let call_id = "ws-provider-call-1"; + let tool_name = "provider_tool"; + let tool_args = json!({ "city": "Paris" }); + let tool_call_arguments = serde_json::to_string(&tool_args)?; + let responses = vec![ + responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_function_call(call_id, tool_name, &tool_call_arguments), + responses::ev_completed("resp-1"), + ]), + create_final_assistant_message_sse_response("Done")?, + ]; + let server = create_mock_responses_server_sequence_unchecked(responses).await; + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), &server.uri(), "never")?; + + let bind_addr = reserve_local_addr()?; + let mut process = spawn_websocket_server(codex_home.path(), bind_addr).await?; + + let test_result = async { + let mut provider = connect_websocket(bind_addr).await?; + let mut thread_client = connect_websocket(bind_addr).await?; + + send_initialize_request(&mut provider, 1, "ws_provider").await?; + let provider_init = read_response_for_id(&mut provider, 1).await?; + assert_eq!(provider_init.id, RequestId::Integer(1)); + + send_initialize_request(&mut thread_client, 2, "ws_thread_client").await?; + let thread_init = read_response_for_id(&mut thread_client, 2).await?; + assert_eq!(thread_init.id, RequestId::Integer(2)); + + send_request( + &mut provider, + "toolProvider/register", + 3, + Some(json!({ + "tool": { + "name": tool_name, + "description": "Provider tool", + "inputSchema": { + "type": "object", + "properties": { + "city": { "type": "string" } + }, + "required": ["city"], + "additionalProperties": false + }, + "defaultTimeoutMs": 1_000, + "injectIntoContext": true + } + })), + ) + .await?; + let register_response = read_response_for_id(&mut provider, 3).await?; + let _: ToolProviderRegisterResponse = to_response(register_response)?; + + send_request( + &mut thread_client, + "thread/start", + 4, + Some(serde_json::to_value(ThreadStartParams::default())?), + ) + .await?; + let thread_start_response = read_response_for_id(&mut thread_client, 4).await?; + let ThreadStartResponse { thread, .. } = to_response(thread_start_response)?; + + send_request( + &mut thread_client, + "turn/start", + 5, + Some(serde_json::to_value(TurnStartParams { + thread_id: thread.id.clone(), + input: vec![V2UserInput::Text { + text: "Run the provider tool".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + })?), + ) + .await?; + let turn_start_response = read_response_for_id(&mut thread_client, 5).await?; + let TurnStartResponse { turn } = to_response::(turn_start_response)?; + + let (request_id, params) = loop { + match timeout(DEFAULT_READ_TIMEOUT, read_jsonrpc_message(&mut provider)).await?? { + JSONRPCMessage::Request(request) => match ServerRequest::try_from(request)? { + ServerRequest::DynamicToolCall { request_id, params } => { + break (request_id, params); + } + other => bail!("expected provider DynamicToolCall request, got {other:?}"), + }, + JSONRPCMessage::Notification(_) | JSONRPCMessage::Response(_) => {} + JSONRPCMessage::Error(error) => { + bail!("unexpected websocket error on provider client: {error:?}"); + } + } + }; + assert_eq!( + params, + DynamicToolCallParams { + thread_id: thread.id.clone(), + turn_id: turn.id, + call_id: call_id.to_string(), + tool: tool_name.to_string(), + arguments: tool_args, + } + ); + + send_response( + &mut provider, + request_id, + serde_json::to_value(DynamicToolCallResponse { + content_items: vec![DynamicToolCallOutputContentItem::InputText { + text: "dynamic-ok".to_string(), + }], + success: true, + })?, + ) + .await?; + + loop { + match timeout( + DEFAULT_READ_TIMEOUT, + read_jsonrpc_message(&mut thread_client), + ) + .await?? + { + JSONRPCMessage::Request(request) => { + bail!("thread client unexpectedly received request: {request:?}"); + } + JSONRPCMessage::Notification(notification) + if notification.method == "turn/completed" => + { + break; + } + JSONRPCMessage::Notification(_) | JSONRPCMessage::Response(_) => {} + JSONRPCMessage::Error(error) => { + bail!("unexpected websocket error on thread client: {error:?}"); + } + } + } + + Ok(()) + } + .await; + + process + .kill() + .await + .context("failed to stop websocket app-server process")?; + test_result +} + +async fn send_response( + stream: &mut super::connection_handling_websocket::WsClient, + id: RequestId, + result: serde_json::Value, +) -> Result<()> { + let payload = serde_json::to_string(&JSONRPCMessage::Response(JSONRPCResponse { id, result }))?; + stream + .send(WebSocketMessage::Text(payload.into())) + .await + .context("failed to send websocket response frame")?; + Ok(()) +} diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index b2bbae4a78..6df63c7069 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -526,6 +526,7 @@ impl Codex { app_server_client_name: None, session_source, dynamic_tools, + registered_dynamic_tools: Vec::new(), persist_extended_history, inherited_shell_snapshot, }; @@ -946,6 +947,7 @@ pub(crate) struct SessionConfiguration { /// Source of the session (cli, vscode, exec, mcp, ...) session_source: SessionSource, dynamic_tools: Vec, + registered_dynamic_tools: Vec, persist_extended_history: bool, inherited_shell_snapshot: Option>, } @@ -955,6 +957,22 @@ impl SessionConfiguration { &self.codex_home } + fn effective_dynamic_tools(&self) -> Vec { + let mut seen_names = HashSet::new(); + let mut tools = + Vec::with_capacity(self.dynamic_tools.len() + self.registered_dynamic_tools.len()); + for tool in self + .dynamic_tools + .iter() + .chain(self.registered_dynamic_tools.iter()) + { + if seen_names.insert(tool.name.clone()) { + tools.push(tool.clone()); + } + } + tools + } + fn thread_config_snapshot(&self) -> ThreadConfigSnapshot { ThreadConfigSnapshot { model: self.collaboration_mode.model().to_string(), @@ -1245,7 +1263,7 @@ impl Session { tool_call_gate: Arc::new(ReadinessFlag::new()), truncation_policy: model_info.truncation_policy.into(), js_repl, - dynamic_tools: session_configuration.dynamic_tools.clone(), + dynamic_tools: session_configuration.effective_dynamic_tools(), turn_metadata_state, turn_skills: TurnSkillsContext::new(skills_outcome), turn_timing_state: Arc::new(TurnTimingState::default()), @@ -2417,6 +2435,14 @@ impl Session { .clone() } + pub(crate) async fn update_registered_dynamic_tools( + &self, + dynamic_tools: Vec, + ) { + let mut state = self.state.lock().await; + state.session_configuration.registered_dynamic_tools = dynamic_tools; + } + pub(crate) async fn provider(&self) -> ModelProviderInfo { let state = self.state.lock().await; state.session_configuration.provider.clone() diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index 69ce86b61b..8e139dabc0 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/codex-rs/core/src/codex_tests.rs @@ -1476,6 +1476,7 @@ async fn set_rate_limits_retains_previous_credits() { app_server_client_name: None, session_source: SessionSource::Exec, dynamic_tools: Vec::new(), + registered_dynamic_tools: Vec::new(), persist_extended_history: false, inherited_shell_snapshot: None, }; @@ -1572,6 +1573,7 @@ async fn set_rate_limits_updates_plan_type_when_present() { app_server_client_name: None, session_source: SessionSource::Exec, dynamic_tools: Vec::new(), + registered_dynamic_tools: Vec::new(), persist_extended_history: false, inherited_shell_snapshot: None, }; @@ -1926,6 +1928,7 @@ pub(crate) async fn make_session_configuration_for_tests() -> SessionConfigurati app_server_client_name: None, session_source: SessionSource::Exec, dynamic_tools: Vec::new(), + registered_dynamic_tools: Vec::new(), persist_extended_history: false, inherited_shell_snapshot: None, } @@ -2079,6 +2082,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_zsh_path() { app_server_client_name: None, session_source: SessionSource::Exec, dynamic_tools: Vec::new(), + registered_dynamic_tools: Vec::new(), persist_extended_history: false, inherited_shell_snapshot: None, }; @@ -2172,6 +2176,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { app_server_client_name: None, session_source: SessionSource::Exec, dynamic_tools: Vec::new(), + registered_dynamic_tools: Vec::new(), persist_extended_history: false, inherited_shell_snapshot: None, }; @@ -2733,6 +2738,7 @@ pub(crate) async fn make_session_and_context_with_dynamic_tools_and_rx( app_server_client_name: None, session_source: SessionSource::Exec, dynamic_tools, + registered_dynamic_tools: Vec::new(), persist_extended_history: false, inherited_shell_snapshot: None, }; diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index b33a66c3ed..e4a1055d22 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -11,6 +11,7 @@ use crate::protocol::Op; use crate::protocol::Submission; use codex_protocol::config_types::Personality; use codex_protocol::config_types::ServiceTier; +use codex_protocol::dynamic_tools::DynamicToolSpec; use codex_protocol::models::ContentItem; use codex_protocol::models::ResponseInputItem; use codex_protocol::models::ResponseItem; @@ -144,6 +145,13 @@ impl CodexThread { self.codex.thread_config_snapshot().await } + pub async fn update_registered_dynamic_tools(&self, dynamic_tools: Vec) { + self.codex + .session + .update_registered_dynamic_tools(dynamic_tools) + .await; + } + pub fn enabled(&self, feature: Feature) -> bool { self.codex.enabled(feature) } diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index cc6de47104..645ebd84fa 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -312,6 +312,10 @@ impl ThreadManager { self.state.get_thread(thread_id).await } + pub async fn loaded_threads(&self) -> Vec> { + self.state.threads.read().await.values().cloned().collect() + } + pub async fn start_thread(&self, config: Config) -> CodexResult { // Box delegated thread-spawn futures so these convenience wrappers do // not inline the full spawn path into every caller's async state. diff --git a/codex-rs/core/src/tools/handlers/dynamic.rs b/codex-rs/core/src/tools/handlers/dynamic.rs index 23285bc907..842ea546af 100644 --- a/codex-rs/core/src/tools/handlers/dynamic.rs +++ b/codex-rs/core/src/tools/handlers/dynamic.rs @@ -97,11 +97,18 @@ async fn request_dynamic_tool( } let started_at = Instant::now(); + let provider_owned = turn_context + .dynamic_tools + .iter() + .find(|spec| spec.name == tool) + .is_some_and(|spec| spec.provider_owned); let event = EventMsg::DynamicToolCallRequest(DynamicToolCallRequest { + thread_id: session.conversation_id.to_string(), call_id: call_id.clone(), turn_id: turn_id.clone(), tool: tool.clone(), arguments: arguments.clone(), + provider_owned, }); session.send_event(turn_context, event).await; let response = rx_response.await.ok(); diff --git a/codex-rs/core/src/tools/js_repl/mod.rs b/codex-rs/core/src/tools/js_repl/mod.rs index d8d043a7d9..38bcf53d04 100644 --- a/codex-rs/core/src/tools/js_repl/mod.rs +++ b/codex-rs/core/src/tools/js_repl/mod.rs @@ -3096,6 +3096,8 @@ await codex.emitImage({ bytes: png, mimeType: "image/png", detail: "ultra" }); "properties": {}, "additionalProperties": false }), + inject_into_context: true, + provider_owned: false, }]) .await; if !turn diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index 7095a38cea..32f1531551 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -183,6 +183,23 @@ impl ToolRouter { err, )); } + if source == ToolCallSource::Direct + && turn + .dynamic_tools + .iter() + .find(|tool| tool.name == tool_name) + .is_some_and(|tool| !tool.inject_into_context) + { + let err = FunctionCallError::RespondToModel( + "direct tool calls are disabled for this dynamic tool; use js_repl and codex.tool(...) instead" + .to_string(), + ); + return Ok(Self::failure_response( + failure_call_id, + payload_outputs_custom, + err, + )); + } let invocation = ToolInvocation { session, @@ -234,8 +251,10 @@ mod tests { use std::sync::Arc; use crate::codex::make_session_and_context; + use crate::codex::make_session_and_context_with_dynamic_tools_and_rx; use crate::tools::context::ToolPayload; use crate::turn_diff_tracker::TurnDiffTracker; + use codex_protocol::dynamic_tools::DynamicToolSpec; use codex_protocol::models::ResponseInputItem; use super::ToolCall; @@ -347,4 +366,55 @@ mod tests { Ok(()) } + + #[tokio::test] + async fn non_injected_dynamic_tools_block_direct_tool_calls() -> anyhow::Result<()> { + let (session, turn, _rx) = + make_session_and_context_with_dynamic_tools_and_rx(vec![DynamicToolSpec { + name: "hidden_dynamic_tool".to_string(), + description: "Only callable from js_repl.".to_string(), + input_schema: serde_json::json!({ + "type": "object", + "properties": {}, + "additionalProperties": false + }), + inject_into_context: false, + provider_owned: false, + }]) + .await; + + let session = Arc::clone(&session); + let turn = Arc::clone(&turn); + let router = ToolRouter::from_config( + &turn.tools_config, + None, + None, + turn.dynamic_tools.as_slice(), + ); + + let call = ToolCall { + tool_name: "hidden_dynamic_tool".to_string(), + call_id: "call-3".to_string(), + payload: ToolPayload::Function { + arguments: "{}".to_string(), + }, + }; + let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); + let response = router + .dispatch_tool_call(session, turn, tracker, call, ToolCallSource::Direct) + .await?; + + match response { + ResponseInputItem::FunctionCallOutput { output, .. } => { + let content = output.text_content().unwrap_or_default(); + assert!( + content.contains("direct tool calls are disabled for this dynamic tool"), + "unexpected tool call message: {content}", + ); + } + other => panic!("expected function call output, got {other:?}"), + } + + Ok(()) + } } diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index cf9eaa32f7..f7dc3f652a 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/codex-rs/core/src/tools/spec.rs @@ -2465,6 +2465,10 @@ pub(crate) fn build_specs( if !dynamic_tools.is_empty() { for tool in dynamic_tools { + builder.register_handler(tool.name.clone(), dynamic_tool_handler.clone()); + if !tool.inject_into_context { + continue; + } match dynamic_tool_to_openai_tool(tool) { Ok(converted_tool) => { push_tool_spec( @@ -2473,7 +2477,6 @@ pub(crate) fn build_specs( false, config.code_mode_enabled, ); - builder.register_handler(tool.name.clone(), dynamic_tool_handler.clone()); } Err(e) => { tracing::error!( diff --git a/codex-rs/core/tests/suite/sqlite_state.rs b/codex-rs/core/tests/suite/sqlite_state.rs index b17219e5f1..99075b321e 100644 --- a/codex-rs/core/tests/suite/sqlite_state.rs +++ b/codex-rs/core/tests/suite/sqlite_state.rs @@ -110,6 +110,8 @@ async fn backfill_scans_existing_rollouts() -> Result<()> { "required": ["city"], "properties": { "city": { "type": "string" } } }), + inject_into_context: true, + provider_owned: false, }, DynamicToolSpec { name: "weather_lookup".to_string(), @@ -119,6 +121,8 @@ async fn backfill_scans_existing_rollouts() -> Result<()> { "required": ["zip"], "properties": { "zip": { "type": "string" } } }), + inject_into_context: true, + provider_owned: false, }, ]; let dynamic_tools_for_hook = dynamic_tools.clone(); diff --git a/codex-rs/protocol/src/dynamic_tools.rs b/codex-rs/protocol/src/dynamic_tools.rs index 8b5405f307..b6a08a36f7 100644 --- a/codex-rs/protocol/src/dynamic_tools.rs +++ b/codex-rs/protocol/src/dynamic_tools.rs @@ -10,15 +10,22 @@ pub struct DynamicToolSpec { pub name: String, pub description: String, pub input_schema: JsonValue, + #[serde(default = "default_inject_into_context")] + pub inject_into_context: bool, + #[serde(default)] + pub provider_owned: bool, } #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, JsonSchema, TS)] #[serde(rename_all = "camelCase")] pub struct DynamicToolCallRequest { + pub thread_id: String, pub call_id: String, pub turn_id: String, pub tool: String, pub arguments: JsonValue, + #[serde(default)] + pub provider_owned: bool, } #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, JsonSchema, TS)] @@ -37,3 +44,7 @@ pub enum DynamicToolCallOutputContentItem { #[serde(rename_all = "camelCase")] InputImage { image_url: String }, } + +fn default_inject_into_context() -> bool { + true +} diff --git a/codex-rs/state/src/runtime/threads.rs b/codex-rs/state/src/runtime/threads.rs index 344a893640..4687e5e9d3 100644 --- a/codex-rs/state/src/runtime/threads.rs +++ b/codex-rs/state/src/runtime/threads.rs @@ -70,6 +70,8 @@ ORDER BY position ASC name: row.try_get("name")?, description: row.try_get("description")?, input_schema, + inject_into_context: true, + provider_owned: false, }); } Ok(Some(tools))