From e1e022c9e327870bebefa26e417dc5f482349e10 Mon Sep 17 00:00:00 2001 From: Feng GAO Date: Wed, 15 Jul 2026 09:08:09 +0800 Subject: [PATCH] feat(acp): dispatch dynamic tools through websocket --- crates/aionui-ai-agent/src/factory/acp.rs | 34 +- .../src/factory/acp_assembler.rs | 12 +- crates/aionui-ai-agent/src/factory/mod.rs | 3 + crates/aionui-ai-agent/src/lib.rs | 1 + .../aionui-ai-agent/src/manager/acp/agent.rs | 9 +- .../src/manager/acp/agent_session_flow.rs | 18 + crates/aionui-ai-agent/src/protocol/acp.rs | 47 +- .../src/protocol/custom_agent_probe.rs | 2 +- .../src/protocol/dynamic_tools.rs | 697 ++++++++++++++++++ crates/aionui-ai-agent/src/protocol/mod.rs | 1 + .../tests/factory_provider_integration.rs | 4 +- crates/aionui-api-types/src/dynamic_tools.rs | 144 ++++ crates/aionui-api-types/src/lib.rs | 5 + crates/aionui-app/src/router/state.rs | 6 +- crates/aionui-app/src/services.rs | 11 +- crates/aionui-app/tests/websocket_e2e.rs | 202 ++++- crates/aionui-realtime/src/handler.rs | 1 + crates/aionui-realtime/src/router.rs | 3 + 18 files changed, 1166 insertions(+), 34 deletions(-) create mode 100644 crates/aionui-ai-agent/src/protocol/dynamic_tools.rs create mode 100644 crates/aionui-api-types/src/dynamic_tools.rs diff --git a/crates/aionui-ai-agent/src/factory/acp.rs b/crates/aionui-ai-agent/src/factory/acp.rs index 98745c29e..2154a75db 100644 --- a/crates/aionui-ai-agent/src/factory/acp.rs +++ b/crates/aionui-ai-agent/src/factory/acp.rs @@ -111,23 +111,23 @@ pub(super) async fn build( } } - let params = Arc::new( - assemble_acp_params( - ctx.conversation_id.clone(), - WorkspaceInfo { - path: ctx.workspace, - is_custom: ctx.is_custom_workspace, - }, - meta, - command_spec, - config, - session_mcp_servers, - session_snapshot, - deps.data_dir.clone(), - deps.dump_prompts, - ) - .await, - ); + let mut params = assemble_acp_params( + ctx.conversation_id.clone(), + WorkspaceInfo { + path: ctx.workspace, + is_custom: ctx.is_custom_workspace, + }, + meta, + command_spec, + config, + session_mcp_servers, + session_snapshot, + deps.data_dir.clone(), + deps.dump_prompts, + ) + .await; + params.dynamic_tool_session = deps.dynamic_tool_registry.session_for(&ctx.conversation_id); + let params = Arc::new(params); let skill_mgr = deps.skill_manager.clone(); let catalog_tx = deps.agent_registry.catalog_sender(); diff --git a/crates/aionui-ai-agent/src/factory/acp_assembler.rs b/crates/aionui-ai-agent/src/factory/acp_assembler.rs index 5edaacdce..0330e66fb 100644 --- a/crates/aionui-ai-agent/src/factory/acp_assembler.rs +++ b/crates/aionui-ai-agent/src/factory/acp_assembler.rs @@ -1,3 +1,4 @@ +use crate::protocol::dynamic_tools::DynamicToolSession; use crate::shared_kernel::PersistedSessionState; use agent_client_protocol::schema::{EnvVariable, McpServer, McpServerStdio, NewSessionRequest}; use aionui_api_types::AgentMetadata; @@ -34,16 +35,24 @@ pub struct AcpSessionParams { pub data_dir: PathBuf, /// Whether prompt diagnostics should be dumped under `data_dir/prompt-dumps`. pub dump_prompts: bool, + /// Registration captured before the ACP session starts. The handle binds + /// the CLI-assigned thread id after session/new or session/load. + pub dynamic_tool_session: Option, } impl AcpSessionParams { /// Build a `NewSessionRequest` using the pre-computed MCP servers. pub fn new_session_request(&self) -> NewSessionRequest { let req = NewSessionRequest::new(&self.workspace.path); - if self.mcp_servers.is_empty() { + let req = if self.mcp_servers.is_empty() { req } else { req.mcp_servers(self.mcp_servers.clone()) + }; + if let Some(session) = &self.dynamic_tool_session { + req.meta(session.metadata()) + } else { + req } } } @@ -83,6 +92,7 @@ pub async fn assemble_acp_params( session_snapshot, data_dir, dump_prompts, + dynamic_tool_session: None, } } diff --git a/crates/aionui-ai-agent/src/factory/mod.rs b/crates/aionui-ai-agent/src/factory/mod.rs index d92428455..666f4730e 100644 --- a/crates/aionui-ai-agent/src/factory/mod.rs +++ b/crates/aionui-ai-agent/src/factory/mod.rs @@ -17,6 +17,7 @@ use crate::capability::skill_manager::AcpSkillManager; use crate::error::AgentError; use crate::factory::context::FactoryContext; use crate::persistence::AcpSessionSyncService; +use crate::protocol::dynamic_tools::DynamicToolRegistry; use crate::registry::AgentRegistry; use crate::session_context::AgentSessionKind; use crate::task_manager::AgentFactory; @@ -40,6 +41,8 @@ pub struct AgentFactoryDeps { /// inject enabled servers into `session/new` (ELECTRON-1JG fix). /// `None` for tests/composition paths that do not need MCP injection. pub mcp_server_repo: Option>, + /// WebSocket-owned dynamic tools available to ordinary ACP conversations. + pub dynamic_tool_registry: DynamicToolRegistry, } /// Build a production agent factory that dispatches to concrete agent types. diff --git a/crates/aionui-ai-agent/src/lib.rs b/crates/aionui-ai-agent/src/lib.rs index f9a2b39f3..117d6bfa0 100644 --- a/crates/aionui-ai-agent/src/lib.rs +++ b/crates/aionui-ai-agent/src/lib.rs @@ -37,6 +37,7 @@ pub use error::AgentError; pub use factory::{AgentFactoryDeps, build_agent_factory}; pub use idle_scanner::{IdleCleanupCoordinator, start_idle_scanner, start_idle_scanner_with_coordinator}; pub use persistence::AcpSessionSyncService; +pub use protocol::dynamic_tools::{DynamicToolRegistry, DynamicToolSession}; pub use protocol::error::AcpError; pub use protocol::events::AgentStreamEvent; pub use protocol::send_error::AgentSendError; diff --git a/crates/aionui-ai-agent/src/manager/acp/agent.rs b/crates/aionui-ai-agent/src/manager/acp/agent.rs index 19d963010..5caf1e21d 100644 --- a/crates/aionui-ai-agent/src/manager/acp/agent.rs +++ b/crates/aionui-ai-agent/src/manager/acp/agent.rs @@ -438,7 +438,14 @@ impl AcpAgentManager { // 70ms in — ELECTRON-1BT), so we explicitly watch the child. If // it dies before init completes, surface a `StartupCrash` carrying // the buffered stderr instead of waiting out the timeout. - let connect_fut = AcpProtocol::connect(stdin, stdout, runtime.event_sender(), permission_tx, notification_tx); + let connect_fut = AcpProtocol::connect( + stdin, + stdout, + runtime.event_sender(), + permission_tx, + notification_tx, + params.dynamic_tool_session.clone(), + ); tokio::pin!(connect_fut); let protocol = tokio::select! { biased; diff --git a/crates/aionui-ai-agent/src/manager/acp/agent_session_flow.rs b/crates/aionui-ai-agent/src/manager/acp/agent_session_flow.rs index 031c2d12b..7912ea28e 100644 --- a/crates/aionui-ai-agent/src/manager/acp/agent_session_flow.rs +++ b/crates/aionui-ai-agent/src/manager/acp/agent_session_flow.rs @@ -39,6 +39,9 @@ impl AcpAgentManager { let session_response = self.protocol.new_session(req).await?; let sid = session_response.session_id.to_string(); + if let Some(dynamic_tools) = &self.params.dynamic_tool_session { + dynamic_tools.bind_thread(&sid); + } { let mut session = self.session.write().await; @@ -122,6 +125,9 @@ impl AcpAgentManager { options.insert("resume".into(), Value::String(session_id.to_owned())); claude_code.insert("options".into(), Value::Object(options)); meta.insert("claudeCode".into(), Value::Object(claude_code)); + if let Some(dynamic_tools) = &self.params.dynamic_tool_session { + meta.extend(dynamic_tools.metadata()); + } let req = self.params.new_session_request().meta(meta); let new_response = match self.protocol.new_session(req).await { @@ -132,6 +138,9 @@ impl AcpAgentManager { Err(e) => return Err(e.into()), }; let new_sid = new_response.session_id.to_string(); + if let Some(dynamic_tools) = &self.params.dynamic_tool_session { + dynamic_tools.bind_thread(&new_sid); + } { let mut session = self.session.write().await; @@ -176,6 +185,9 @@ impl AcpAgentManager { if !self.params.mcp_servers.is_empty() { load_req = load_req.mcp_servers(self.params.mcp_servers.clone()); } + if let Some(dynamic_tools) = &self.params.dynamic_tool_session { + load_req = load_req.meta(dynamic_tools.metadata()); + } let load_response = match self.protocol.load_session(load_req).await { Ok(r) => r, Err(e) if is_acp_session_not_found(&e) => { @@ -183,6 +195,9 @@ impl AcpAgentManager { } Err(e) => return Err(e.into()), }; + if let Some(dynamic_tools) = &self.params.dynamic_tool_session { + dynamic_tools.bind_thread(session_id); + } { let mut session = self.session.write().await; @@ -214,6 +229,9 @@ impl AcpAgentManager { // session/load. Seed the aggregate with the stored id and let the // caller prompt — matches pre-refactor behaviour. { + if let Some(dynamic_tools) = &self.params.dynamic_tool_session { + dynamic_tools.bind_thread(session_id); + } let mut session = self.session.write().await; session.set_session_id(DomainSessionId::new(session_id.to_owned())); self.commit_session_changes(&mut session).await; diff --git a/crates/aionui-ai-agent/src/protocol/acp.rs b/crates/aionui-ai-agent/src/protocol/acp.rs index ca6121d77..cf877f288 100644 --- a/crates/aionui-ai-agent/src/protocol/acp.rs +++ b/crates/aionui-ai-agent/src/protocol/acp.rs @@ -26,14 +26,14 @@ use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex, RwLock}; use agent_client_protocol::schema::{ - AGENT_METHOD_NAMES, AuthenticateResponse, ClientNotification, ClientRequest, CloseSessionResponse, ExtResponse, - ForkSessionResponse, Implementation, InitializeRequest, LoadSessionResponse, PromptResponse, ProtocolVersion, - RequestPermissionOutcome, RequestPermissionRequest, RequestPermissionResponse, ResumeSessionResponse, - SelectedPermissionOutcome, SessionNotification, SetSessionConfigOptionResponse, SetSessionModeResponse, - SetSessionModelResponse, + AGENT_METHOD_NAMES, AgentRequest, AuthenticateResponse, ClientNotification, ClientRequest, CloseSessionResponse, + ExtResponse, ForkSessionResponse, Implementation, InitializeRequest, LoadSessionResponse, PromptResponse, + ProtocolVersion, RequestPermissionOutcome, RequestPermissionRequest, RequestPermissionResponse, + ResumeSessionResponse, SelectedPermissionOutcome, SessionNotification, SetSessionConfigOptionResponse, + SetSessionModeResponse, SetSessionModelResponse, }; use agent_client_protocol::{ - Agent, ByteStreams, Client, ConnectionTo, Responder, on_receive_notification, on_receive_request, + Agent, ByteStreams, Client, ConnectionTo, Handled, Responder, on_receive_notification, on_receive_request, }; use aionui_common::ErrorChain; use tokio::process::{ChildStdin, ChildStdout}; @@ -41,6 +41,7 @@ use tokio::sync::{broadcast, mpsc, oneshot}; use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; use tracing::{debug, info, warn}; +use crate::protocol::dynamic_tools::{CODEX_DYNAMIC_TOOL_CALL_METHOD, DynamicToolSession, dynamic_tool_unavailable}; use crate::protocol::error::AcpError; use crate::protocol::events::{self as stream_event, AgentStreamEvent}; @@ -137,6 +138,7 @@ impl AcpProtocol { event_tx: broadcast::Sender, permission_tx: mpsc::Sender, notification_tx: mpsc::Sender, + dynamic_tool_session: Option, ) -> Result { let alive = Arc::new(AtomicBool::new(true)); let replay_suppression = Arc::new(AtomicBool::new(false)); @@ -160,6 +162,7 @@ impl AcpProtocol { event_tx, permission_tx, notification_tx, + dynamic_tool_session, init_tx, ready_tx, shutdown_rx, @@ -406,6 +409,7 @@ async fn run_sdk_background( event_tx: broadcast::Sender, permission_tx: mpsc::Sender, notification_tx: mpsc::Sender, + dynamic_tool_session: Option, init_tx: oneshot::Sender>, ready_tx: oneshot::Sender>, shutdown_rx: oneshot::Receiver<()>, @@ -460,6 +464,37 @@ async fn run_sdk_background( }, on_receive_request!(), ) + .on_receive_request( + { + async move |request: AgentRequest, responder, _cx| { + let AgentRequest::ExtMethodRequest(extension) = request else { + return Ok(Handled::No { + message: (request, responder), + retry: false, + }); + }; + if extension.method.as_ref() != CODEX_DYNAMIC_TOOL_CALL_METHOD { + return Ok(Handled::No { + message: (AgentRequest::ExtMethodRequest(extension), responder), + retry: false, + }); + } + + let response = match ( + dynamic_tool_session.as_ref(), + serde_json::from_str(extension.params.get()), + ) { + (Some(session), Ok(params)) => session.dispatch(params).await, + _ => dynamic_tool_unavailable(), + }; + let response = serde_json::to_value(response) + .unwrap_or_else(|_| serde_json::json!({"success": false, "contentItems": []})); + responder.respond(response)?; + Ok(Handled::Yes) + } + }, + on_receive_request!(), + ) .connect_with(transport, async move |connection: ConnectionTo| { // Step 1 — initialize handshake. main_fn is the canonical place // to call `block_task` (see SDK `connect_with` doc example). diff --git a/crates/aionui-ai-agent/src/protocol/custom_agent_probe.rs b/crates/aionui-ai-agent/src/protocol/custom_agent_probe.rs index 941814f6c..7001b0cd1 100644 --- a/crates/aionui-ai-agent/src/protocol/custom_agent_probe.rs +++ b/crates/aionui-ai-agent/src/protocol/custom_agent_probe.rs @@ -206,7 +206,7 @@ async fn run_handshake(proc: &CliAgentProcess) -> ProbeOutcome { // immediately with a non-zero status; without this race the // `AcpProtocol::connect` call would block on its internal 30 s // timeout waiting for an `initialize` reply that will never arrive. - let connect = AcpProtocol::connect(stdin, stdout, event_tx, permission_tx, notification_tx); + let connect = AcpProtocol::connect(stdin, stdout, event_tx, permission_tx, notification_tx, None); let protocol = tokio::select! { biased; res = connect => match res { diff --git a/crates/aionui-ai-agent/src/protocol/dynamic_tools.rs b/crates/aionui-ai-agent/src/protocol/dynamic_tools.rs new file mode 100644 index 000000000..eedfb3d21 --- /dev/null +++ b/crates/aionui-ai-agent/src/protocol/dynamic_tools.rs @@ -0,0 +1,697 @@ +use std::collections::{HashMap, HashSet}; +use std::fmt; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use aionui_api_types::{ + DynamicToolCallOutputContentItem, DynamicToolCallParams, DynamicToolCallPayload, DynamicToolCallResponse, + DynamicToolResultPayload, DynamicToolSpec, DynamicToolsRegisterRequest, DynamicToolsRegisteredPayload, + WebSocketMessage, +}; +use aionui_realtime::{ConnectionId, MessageRouter, WebSocketManager}; +use serde_json::{Map, Value, json}; +use tokio::sync::oneshot; +use tracing::{info, warn}; + +pub const DYNAMIC_TOOLS_REGISTER_EVENT: &str = "agent.dynamicToolsRegister"; +pub const DYNAMIC_TOOLS_REGISTERED_EVENT: &str = "agent.dynamicToolsRegistered"; +pub const DYNAMIC_TOOL_CALL_EVENT: &str = "agent.dynamicToolCall"; +pub const DYNAMIC_TOOL_RESULT_EVENT: &str = "agent.dynamicToolResult"; +pub const CODEX_DYNAMIC_TOOLS_META_KEY: &str = "codex/dynamic_tools"; +pub const CODEX_DYNAMIC_TOOL_CALL_METHOD: &str = "codex/dynamic_tool_call"; + +const DEFAULT_DYNAMIC_TOOL_TIMEOUT: Duration = Duration::from_secs(60); +const MAX_DYNAMIC_TOOL_SPECS: usize = 64; + +#[derive(Clone)] +pub struct DynamicToolRegistry { + manager: Arc, + state: Arc>, + next_generation: Arc, + call_timeout: Duration, +} + +#[derive(Default)] +struct RegistryState { + registrations: HashMap>, + conflicted_conversations: HashSet, + pending: HashMap, +} + +#[derive(Clone)] +struct Registration { + id: String, + generation: u64, + tools: Vec, + thread_id: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct PendingKey { + registration_id: String, + thread_id: String, + call_id: String, +} + +struct PendingCall { + conversation_id: String, + connection_id: ConnectionId, + turn_id: String, + namespace: Option, + tool: String, + response_tx: oneshot::Sender, +} + +#[derive(Clone)] +pub struct DynamicToolSession { + registry: DynamicToolRegistry, + conversation_id: String, + connection_id: ConnectionId, + registration_id: String, + generation: u64, + tools: Vec, +} + +impl fmt::Debug for DynamicToolSession { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("DynamicToolSession") + .field("conversation_id", &self.conversation_id) + .field("connection_id", &self.connection_id) + .field("registration_id", &self.registration_id) + .field("generation", &self.generation) + .field("tool_count", &self.tools.len()) + .finish() + } +} + +impl DynamicToolRegistry { + pub fn new(manager: Arc) -> Self { + Self::with_timeout(manager, DEFAULT_DYNAMIC_TOOL_TIMEOUT) + } + + pub fn with_timeout(manager: Arc, call_timeout: Duration) -> Self { + Self { + manager, + state: Arc::new(Mutex::new(RegistryState::default())), + next_generation: Arc::new(AtomicU64::new(1)), + call_timeout, + } + } + + pub fn session_for(&self, conversation_id: &str) -> Option { + let state = self.state.lock().unwrap(); + if state.conflicted_conversations.contains(conversation_id) { + return None; + } + let registrations = state.registrations.get(conversation_id)?; + if registrations.len() != 1 { + return None; + } + let (&connection_id, registration) = registrations.iter().next()?; + Some(DynamicToolSession { + registry: self.clone(), + conversation_id: conversation_id.to_owned(), + connection_id, + registration_id: registration.id.clone(), + generation: registration.generation, + tools: registration.tools.clone(), + }) + } + + fn register( + &self, + connection_id: ConnectionId, + request: DynamicToolsRegisterRequest, + ) -> DynamicToolsRegisteredPayload { + let request_id = request.request_id.trim().to_owned(); + let conversation_id = request.conversation_id.trim().to_owned(); + if request_id.is_empty() || conversation_id.is_empty() { + return registration_failure(request_id, conversation_id, "dynamic_tool_registration_invalid"); + } + if let Err(code) = validate_tool_specs(&request.tools) { + return registration_failure(request_id, conversation_id, code); + } + + let generation = self.next_generation.fetch_add(1, Ordering::Relaxed); + let registration_id = format!("dynamic-tools-{generation}"); + let mut state = self.state.lock().unwrap(); + + let was_sole_reregistration = state + .registrations + .get(&conversation_id) + .is_some_and(|owners| owners.len() == 1 && owners.contains_key(&connection_id)); + if was_sole_reregistration { + state.conflicted_conversations.remove(&conversation_id); + } + + let previous_registration_id = state + .registrations + .get_mut(&conversation_id) + .and_then(|owners| owners.remove(&connection_id)) + .map(|registration| registration.id); + if let Some(previous_registration_id) = previous_registration_id { + drop_pending_for_registration(&mut state, &previous_registration_id); + } + + let owners = state.registrations.entry(conversation_id.clone()).or_default(); + owners.insert( + connection_id, + Registration { + id: registration_id.clone(), + generation, + tools: request.tools, + thread_id: None, + }, + ); + let ambiguous = owners.len() != 1; + if ambiguous { + state.conflicted_conversations.insert(conversation_id.clone()); + } + drop(state); + + info!( + %connection_id, + %conversation_id, + %registration_id, + generation, + accepted = !ambiguous, + "dynamic tool registration observed" + ); + DynamicToolsRegisteredPayload { + request_id, + conversation_id, + accepted: !ambiguous, + registration_id: Some(registration_id), + error_code: ambiguous.then(|| "dynamic_tool_owner_ambiguous".to_owned()), + } + } + + fn handle_result(&self, connection_id: ConnectionId, result: DynamicToolResultPayload) -> bool { + let key = PendingKey { + registration_id: result.registration_id.clone(), + thread_id: result.thread_id.clone(), + call_id: result.call_id.clone(), + }; + let mut state = self.state.lock().unwrap(); + let Some(pending) = state.pending.get(&key) else { + warn!( + %connection_id, + thread_id = %result.thread_id, + turn_id = %result.turn_id, + call_id = %result.call_id, + tool = %result.tool, + "dynamic tool result rejected for unknown call" + ); + return false; + }; + let identity_matches = pending.connection_id == connection_id + && pending.conversation_id == result.conversation_id + && pending.turn_id == result.turn_id + && pending.namespace == result.namespace + && pending.tool == result.tool; + if !identity_matches { + warn!( + %connection_id, + thread_id = %result.thread_id, + turn_id = %result.turn_id, + call_id = %result.call_id, + tool = %result.tool, + "dynamic tool result rejected for identity mismatch" + ); + return false; + } + let pending = state.pending.remove(&key).expect("pending call checked above"); + drop(state); + let _ = pending.response_tx.send(result.into_response()); + true + } + + fn disconnect(&self, connection_id: ConnectionId) { + let mut state = self.state.lock().unwrap(); + let mut emptied = Vec::new(); + let mut removed_registration_ids = Vec::new(); + for (conversation_id, owners) in &mut state.registrations { + if let Some(registration) = owners.remove(&connection_id) { + removed_registration_ids.push(registration.id); + } + if owners.is_empty() { + emptied.push(conversation_id.clone()); + } + } + for conversation_id in emptied { + state.registrations.remove(&conversation_id); + state.conflicted_conversations.remove(&conversation_id); + } + for registration_id in &removed_registration_ids { + drop_pending_for_registration(&mut state, registration_id); + } + drop(state); + if !removed_registration_ids.is_empty() { + info!( + %connection_id, + registrations_removed = removed_registration_ids.len(), + "dynamic tool registrations cleared on websocket disconnect" + ); + } + } + + fn send_registration_ack(&self, connection_id: ConnectionId, payload: DynamicToolsRegisteredPayload) { + self.manager.send_to( + connection_id, + WebSocketMessage::new( + DYNAMIC_TOOLS_REGISTERED_EVENT, + serde_json::to_value(payload).unwrap_or_else(|_| { + json!({ + "accepted": false, + "errorCode": "dynamic_tool_registration_internal" + }) + }), + ), + ); + } +} + +impl MessageRouter for DynamicToolRegistry { + fn route(&self, connection_id: ConnectionId, name: &str, data: Value) -> bool { + match name { + DYNAMIC_TOOLS_REGISTER_EVENT => { + let fallback_request_id = data + .get("requestId") + .and_then(Value::as_str) + .unwrap_or_default() + .to_owned(); + let fallback_conversation_id = data + .get("conversationId") + .and_then(Value::as_str) + .unwrap_or_default() + .to_owned(); + let payload = match serde_json::from_value::(data) { + Ok(request) => self.register(connection_id, request), + Err(_) => registration_failure( + fallback_request_id, + fallback_conversation_id, + "dynamic_tool_registration_invalid", + ), + }; + self.send_registration_ack(connection_id, payload); + true + } + DYNAMIC_TOOL_RESULT_EVENT => { + match serde_json::from_value::(data) { + Ok(result) => { + self.handle_result(connection_id, result); + } + Err(error) => { + warn!(%connection_id, error = %error, "malformed dynamic tool result rejected"); + } + } + true + } + _ => false, + } + } + + fn disconnected(&self, connection_id: ConnectionId) { + self.disconnect(connection_id); + } +} + +impl DynamicToolSession { + pub fn metadata(&self) -> Map { + let mut metadata = Map::new(); + metadata.insert( + CODEX_DYNAMIC_TOOLS_META_KEY.to_owned(), + json!({ + "version": 1, + "tools": self.tools, + }), + ); + metadata + } + + pub fn bind_thread(&self, thread_id: &str) -> bool { + let thread_id = thread_id.trim(); + if thread_id.is_empty() { + return false; + } + let mut state = self.registry.state.lock().unwrap(); + if state.conflicted_conversations.contains(&self.conversation_id) { + return false; + } + let Some(registration) = state + .registrations + .get_mut(&self.conversation_id) + .and_then(|owners| owners.get_mut(&self.connection_id)) + else { + return false; + }; + if registration.id != self.registration_id || registration.generation != self.generation { + return false; + } + registration.thread_id = Some(thread_id.to_owned()); + info!( + conversation_id = %self.conversation_id, + registration_id = %self.registration_id, + thread_id, + "dynamic tool registration bound to ACP session" + ); + true + } + + pub async fn dispatch(&self, params: DynamicToolCallParams) -> DynamicToolCallResponse { + if !valid_call_identity(¶ms) || !tool_is_registered(&self.tools, params.namespace.as_deref(), ¶ms.tool) + { + return dynamic_tool_unavailable(); + } + let key = PendingKey { + registration_id: self.registration_id.clone(), + thread_id: params.thread_id.clone(), + call_id: params.call_id.clone(), + }; + let (response_tx, response_rx) = oneshot::channel(); + { + let mut state = self.registry.state.lock().unwrap(); + if state.conflicted_conversations.contains(&self.conversation_id) { + return dynamic_tool_unavailable(); + } + let Some(registration) = state + .registrations + .get(&self.conversation_id) + .and_then(|owners| owners.get(&self.connection_id)) + else { + return dynamic_tool_unavailable(); + }; + let registration_is_current = registration.id == self.registration_id + && registration.generation == self.generation + && registration.thread_id.as_deref() == Some(params.thread_id.as_str()); + if !registration_is_current || state.pending.contains_key(&key) { + return dynamic_tool_unavailable(); + } + state.pending.insert( + key.clone(), + PendingCall { + conversation_id: self.conversation_id.clone(), + connection_id: self.connection_id, + turn_id: params.turn_id.clone(), + namespace: params.namespace.clone(), + tool: params.tool.clone(), + response_tx, + }, + ); + } + + info!( + conversation_id = %self.conversation_id, + thread_id = %params.thread_id, + turn_id = %params.turn_id, + call_id = %params.call_id, + tool = %params.tool, + "dynamic tool call dispatched to websocket owner" + ); + self.registry.manager.send_to( + self.connection_id, + WebSocketMessage::new( + DYNAMIC_TOOL_CALL_EVENT, + serde_json::to_value(DynamicToolCallPayload { + conversation_id: self.conversation_id.clone(), + registration_id: self.registration_id.clone(), + thread_id: params.thread_id, + turn_id: params.turn_id, + call_id: params.call_id, + namespace: params.namespace, + tool: params.tool, + arguments: params.arguments, + }) + .unwrap_or_else(|_| json!({})), + ), + ); + + match tokio::time::timeout(self.registry.call_timeout, response_rx).await { + Ok(Ok(response)) => response, + Ok(Err(_)) | Err(_) => { + self.registry.state.lock().unwrap().pending.remove(&key); + dynamic_tool_unavailable() + } + } + } +} + +fn registration_failure( + request_id: String, + conversation_id: String, + error_code: &str, +) -> DynamicToolsRegisteredPayload { + DynamicToolsRegisteredPayload { + request_id, + conversation_id, + accepted: false, + registration_id: None, + error_code: Some(error_code.to_owned()), + } +} + +fn validate_tool_specs(tools: &[DynamicToolSpec]) -> Result<(), &'static str> { + if tools.is_empty() || tools.len() > MAX_DYNAMIC_TOOL_SPECS { + return Err("dynamic_tool_registry_invalid"); + } + let mut identities = HashSet::new(); + for spec in tools { + match spec { + DynamicToolSpec::Function { name, .. } => { + if name.trim().is_empty() || !identities.insert((None, name.trim().to_owned())) { + return Err("dynamic_tool_registry_invalid"); + } + } + DynamicToolSpec::Namespace { name, tools, .. } => { + let namespace = name.trim(); + if namespace.is_empty() || tools.is_empty() { + return Err("dynamic_tool_registry_invalid"); + } + for nested in tools { + let DynamicToolSpec::Function { name, .. } = nested else { + return Err("dynamic_tool_registry_invalid"); + }; + if name.trim().is_empty() + || !identities.insert((Some(namespace.to_owned()), name.trim().to_owned())) + { + return Err("dynamic_tool_registry_invalid"); + } + } + } + } + } + Ok(()) +} + +fn tool_is_registered(tools: &[DynamicToolSpec], namespace: Option<&str>, tool: &str) -> bool { + match namespace { + None => tools + .iter() + .any(|spec| matches!(spec, DynamicToolSpec::Function { name, .. } if name == tool)), + Some(namespace) => tools.iter().any(|spec| { + matches!( + spec, + DynamicToolSpec::Namespace { name, tools, .. } + if name == namespace + && tools.iter().any(|nested| matches!(nested, DynamicToolSpec::Function { name, .. } if name == tool)) + ) + }), + } +} + +fn valid_call_identity(params: &DynamicToolCallParams) -> bool { + !params.thread_id.trim().is_empty() + && !params.turn_id.trim().is_empty() + && !params.call_id.trim().is_empty() + && !params.tool.trim().is_empty() + && params + .namespace + .as_ref() + .is_none_or(|namespace| !namespace.trim().is_empty()) +} + +fn drop_pending_for_registration(state: &mut RegistryState, registration_id: &str) { + state.pending.retain(|key, _| key.registration_id != registration_id); +} + +pub fn dynamic_tool_unavailable() -> DynamicToolCallResponse { + DynamicToolCallResponse { + success: false, + content_items: vec![DynamicToolCallOutputContentItem::InputText { + text: "Dynamic tool unavailable.".to_owned(), + }], + } +} + +#[cfg(test)] +mod tests { + use super::*; + use aionui_realtime::{PER_CONNECTION_BUFFER, WsOutbound}; + use serde_json::json; + use tokio::sync::mpsc; + + fn test_tools() -> Vec { + vec![DynamicToolSpec::Function { + name: "list_threads".into(), + description: "List tasks".into(), + input_schema: json!({"type": "object"}), + defer_loading: None, + }] + } + + fn add_connection(manager: &Arc) -> (ConnectionId, mpsc::Receiver) { + let (tx, rx) = mpsc::channel(PER_CONNECTION_BUFFER); + (manager.add_client("token".into(), tx), rx) + } + + fn register( + registry: &DynamicToolRegistry, + connection_id: ConnectionId, + conversation_id: &str, + ) -> DynamicToolsRegisteredPayload { + registry.register( + connection_id, + DynamicToolsRegisterRequest { + request_id: format!("request-{conversation_id}"), + conversation_id: conversation_id.into(), + tools: test_tools(), + }, + ) + } + + fn call_params(thread_id: &str) -> DynamicToolCallParams { + DynamicToolCallParams { + thread_id: thread_id.into(), + turn_id: "turn-1".into(), + call_id: "call-1".into(), + namespace: None, + tool: "list_threads".into(), + arguments: json!({"scope": "current_project"}), + } + } + + fn outbound_json(outbound: WsOutbound) -> Value { + let WsOutbound::Text(text) = outbound else { + panic!("expected websocket text"); + }; + serde_json::from_str(&text).unwrap() + } + + #[test] + fn session_metadata_matches_codex_dynamic_tools_extension() { + let manager = Arc::new(WebSocketManager::new()); + let registry = DynamicToolRegistry::new(manager.clone()); + let (connection_id, _rx) = add_connection(&manager); + assert!(register(®istry, connection_id, "conversation-1").accepted); + + let session = registry.session_for("conversation-1").unwrap(); + let meta = session.metadata(); + assert_eq!(meta[CODEX_DYNAMIC_TOOLS_META_KEY]["version"], 1); + assert_eq!(meta[CODEX_DYNAMIC_TOOLS_META_KEY]["tools"][0]["name"], "list_threads"); + } + + #[test] + fn ambiguous_owner_stays_revoked_after_one_connection_disconnects() { + let manager = Arc::new(WebSocketManager::new()); + let registry = DynamicToolRegistry::new(manager.clone()); + let (first, _first_rx) = add_connection(&manager); + let (second, _second_rx) = add_connection(&manager); + assert!(register(®istry, first, "conversation-1").accepted); + assert!(!register(®istry, second, "conversation-1").accepted); + assert!(registry.session_for("conversation-1").is_none()); + + registry.disconnect(second); + assert!(registry.session_for("conversation-1").is_none()); + assert!(register(®istry, first, "conversation-1").accepted); + assert!(registry.session_for("conversation-1").is_some()); + } + + #[tokio::test] + async fn call_result_round_trip_preserves_full_identity() { + let manager = Arc::new(WebSocketManager::new()); + let registry = DynamicToolRegistry::with_timeout(manager.clone(), Duration::from_secs(1)); + let (connection_id, mut rx) = add_connection(&manager); + let registration = register(®istry, connection_id, "conversation-1"); + let session = registry.session_for("conversation-1").unwrap(); + assert!(session.bind_thread("thread-1")); + + let dispatch = tokio::spawn({ + let session = session.clone(); + async move { session.dispatch(call_params("thread-1")).await } + }); + let call = outbound_json(rx.recv().await.unwrap()); + assert_eq!(call["name"], DYNAMIC_TOOL_CALL_EVENT); + assert_eq!(call["data"]["conversationId"], "conversation-1"); + assert_eq!(call["data"]["arguments"]["scope"], "current_project"); + + assert!(registry.handle_result( + connection_id, + DynamicToolResultPayload { + conversation_id: "conversation-1".into(), + registration_id: registration.registration_id.unwrap(), + thread_id: "thread-1".into(), + turn_id: "turn-1".into(), + call_id: "call-1".into(), + namespace: None, + tool: "list_threads".into(), + content_items: vec![DynamicToolCallOutputContentItem::InputText { text: "ready".into() }], + success: true, + }, + )); + assert_eq!(dispatch.await.unwrap().content_items.len(), 1); + } + + #[tokio::test] + async fn wrong_connection_and_thread_fail_closed() { + let manager = Arc::new(WebSocketManager::new()); + let registry = DynamicToolRegistry::with_timeout(manager.clone(), Duration::from_millis(20)); + let (owner, mut owner_rx) = add_connection(&manager); + let (other, _other_rx) = add_connection(&manager); + let registration = register(®istry, owner, "conversation-1"); + let session = registry.session_for("conversation-1").unwrap(); + assert!(session.bind_thread("thread-1")); + + assert!(!session.dispatch(call_params("thread-other")).await.success); + let dispatch = tokio::spawn({ + let session = session.clone(); + async move { session.dispatch(call_params("thread-1")).await } + }); + let _call = owner_rx.recv().await.unwrap(); + assert!(!registry.handle_result( + other, + DynamicToolResultPayload { + conversation_id: "conversation-1".into(), + registration_id: registration.registration_id.unwrap(), + thread_id: "thread-1".into(), + turn_id: "turn-1".into(), + call_id: "call-1".into(), + namespace: None, + tool: "list_threads".into(), + content_items: vec![], + success: true, + }, + )); + assert!(!dispatch.await.unwrap().success); + } + + #[tokio::test] + async fn disconnect_clears_registration_and_pending_call() { + let manager = Arc::new(WebSocketManager::new()); + let registry = DynamicToolRegistry::with_timeout(manager.clone(), Duration::from_secs(1)); + let (connection_id, mut rx) = add_connection(&manager); + register(®istry, connection_id, "conversation-1"); + let session = registry.session_for("conversation-1").unwrap(); + assert!(session.bind_thread("thread-1")); + + let dispatch = tokio::spawn({ + let session = session.clone(); + async move { session.dispatch(call_params("thread-1")).await } + }); + let _call = rx.recv().await.unwrap(); + registry.disconnect(connection_id); + + assert!(!dispatch.await.unwrap().success); + assert!(registry.session_for("conversation-1").is_none()); + } +} diff --git a/crates/aionui-ai-agent/src/protocol/mod.rs b/crates/aionui-ai-agent/src/protocol/mod.rs index 5b908781f..609750332 100644 --- a/crates/aionui-ai-agent/src/protocol/mod.rs +++ b/crates/aionui-ai-agent/src/protocol/mod.rs @@ -1,6 +1,7 @@ pub(crate) mod acp; pub(crate) mod cli_detect; pub(crate) mod custom_agent_probe; +pub mod dynamic_tools; pub(crate) mod error; pub mod events; pub mod send_error; diff --git a/crates/aionui-ai-agent/tests/factory_provider_integration.rs b/crates/aionui-ai-agent/tests/factory_provider_integration.rs index 9792698dd..57f403794 100644 --- a/crates/aionui-ai-agent/tests/factory_provider_integration.rs +++ b/crates/aionui-ai-agent/tests/factory_provider_integration.rs @@ -3,6 +3,7 @@ use std::sync::Arc; use aionui_ai_agent::AcpSessionSyncService; use aionui_ai_agent::AcpSkillManager; +use aionui_ai_agent::DynamicToolRegistry; use aionui_ai_agent::factory::{AgentFactoryDeps, build_agent_factory}; use aionui_ai_agent::registry::AgentRegistry; use aionui_ai_agent::session_context::{ @@ -15,7 +16,7 @@ use aionui_db::{ CreateProviderParams, IAcpSessionRepository, IProviderRepository, SqliteAcpSessionRepository, SqliteAgentMetadataRepository, SqliteProviderRepository, init_database_memory, }; -use aionui_realtime::BroadcastEventBus; +use aionui_realtime::{BroadcastEventBus, WebSocketManager}; fn test_encryption_key() -> [u8; 32] { [0xABu8; 32] @@ -78,6 +79,7 @@ fn make_factory( broadcaster: Arc::new(BroadcastEventBus::new(16)), backend_binary_path: Arc::new(PathBuf::from("/tmp/aionrs-test/aioncore")), mcp_server_repo: None, + dynamic_tool_registry: DynamicToolRegistry::new(Arc::new(WebSocketManager::new())), }) } diff --git a/crates/aionui-api-types/src/dynamic_tools.rs b/crates/aionui-api-types/src/dynamic_tools.rs new file mode 100644 index 000000000..5a6b03120 --- /dev/null +++ b/crates/aionui-api-types/src/dynamic_tools.rs @@ -0,0 +1,144 @@ +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[serde(tag = "type", rename_all = "camelCase")] +pub enum DynamicToolSpec { + Function { + name: String, + description: String, + #[serde(rename = "inputSchema")] + input_schema: Value, + #[serde(rename = "deferLoading", default, skip_serializing_if = "Option::is_none")] + defer_loading: Option, + }, + Namespace { + name: String, + description: String, + tools: Vec, + }, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[serde(rename_all = "camelCase")] +pub struct DynamicToolsRegisterRequest { + pub request_id: String, + pub conversation_id: String, + pub tools: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct DynamicToolsRegisteredPayload { + pub request_id: String, + pub conversation_id: String, + pub accepted: bool, + pub registration_id: Option, + pub error_code: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[serde(rename_all = "camelCase")] +pub struct DynamicToolCallParams { + pub thread_id: String, + pub turn_id: String, + pub call_id: String, + pub namespace: Option, + pub tool: String, + pub arguments: Value, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[serde(rename_all = "camelCase")] +pub struct DynamicToolCallPayload { + pub conversation_id: String, + pub registration_id: String, + pub thread_id: String, + pub turn_id: String, + pub call_id: String, + pub namespace: Option, + pub tool: String, + pub arguments: Value, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(tag = "type", rename_all = "camelCase")] +pub enum DynamicToolCallOutputContentItem { + InputText { + text: String, + }, + InputImage { + #[serde(rename = "imageUrl")] + image_url: String, + }, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct DynamicToolCallResponse { + pub content_items: Vec, + pub success: bool, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct DynamicToolResultPayload { + pub conversation_id: String, + pub registration_id: String, + pub thread_id: String, + pub turn_id: String, + pub call_id: String, + pub namespace: Option, + pub tool: String, + pub content_items: Vec, + pub success: bool, +} + +impl DynamicToolResultPayload { + pub fn into_response(self) -> DynamicToolCallResponse { + DynamicToolCallResponse { + content_items: self.content_items, + success: self.success, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn dynamic_tool_wire_shapes_use_codex_camel_case() { + let spec = DynamicToolSpec::Function { + name: "list_threads".into(), + description: "List tasks".into(), + input_schema: json!({"type": "object"}), + defer_loading: Some(true), + }; + let value = serde_json::to_value(spec).unwrap(); + + assert_eq!(value["type"], "function"); + assert_eq!(value["inputSchema"]["type"], "object"); + assert_eq!(value["deferLoading"], true); + } + + #[test] + fn dynamic_tool_result_round_trips_without_identity_loss() { + let value = json!({ + "conversationId": "conversation-1", + "registrationId": "dynamic-tools-7", + "threadId": "thread-1", + "turnId": "turn-1", + "callId": "call-1", + "namespace": null, + "tool": "list_threads", + "contentItems": [{"type": "inputText", "text": "ready"}], + "success": true + }); + + let result: DynamicToolResultPayload = serde_json::from_value(value).unwrap(); + assert_eq!(result.thread_id, "thread-1"); + assert_eq!(result.into_response().content_items.len(), 1); + } +} diff --git a/crates/aionui-api-types/src/lib.rs b/crates/aionui-api-types/src/lib.rs index 4acca5e33..9438202b3 100644 --- a/crates/aionui-api-types/src/lib.rs +++ b/crates/aionui-api-types/src/lib.rs @@ -14,6 +14,7 @@ mod connection_test; mod conversation; mod cron; mod custom_agent; +mod dynamic_tools; mod extension; mod file; mod lifecycle; @@ -93,6 +94,10 @@ pub use custom_agent::{ AgentOverridesResponse, CustomAgentAdvancedOverrides, CustomAgentUpsertRequest, DeleteCustomAgentResponse, SetAgentOverridesRequest, SetEnabledRequest, }; +pub use dynamic_tools::{ + DynamicToolCallOutputContentItem, DynamicToolCallParams, DynamicToolCallPayload, DynamicToolCallResponse, + DynamicToolResultPayload, DynamicToolSpec, DynamicToolsRegisterRequest, DynamicToolsRegisteredPayload, +}; pub use extension::{ DisableExtensionRequest, EnableExtensionRequest, ExtensionSummaryResponse, GetI18nRequest, GetPermissionsRequest, GetRiskLevelRequest, HubExtensionListItem, HubExtensionListResponse, HubOperationResponse, HubUpdateInfo, diff --git a/crates/aionui-app/src/router/state.rs b/crates/aionui-app/src/router/state.rs index dd06ed362..49eb09cf7 100644 --- a/crates/aionui-app/src/router/state.rs +++ b/crates/aionui-app/src/router/state.rs @@ -36,7 +36,7 @@ use aionui_mcp::{ use aionui_office::{ ConversionService, OfficeRouterState, OfficecliWatchManager, ProxyService, SnapshotService as OfficeSnapshotService, }; -use aionui_realtime::{NoopMessageRouter, WsHandlerState}; +use aionui_realtime::WsHandlerState; use aionui_shell::ShellRouterState; use aionui_system::{ ClientPrefService, ConnectionTestRouterState, ConnectionTestService, FeedbackDiagnosticsService, ModelFetchService, @@ -822,7 +822,7 @@ pub fn build_ws_state(services: &AppServices) -> WsHandlerState { if services.local { return WsHandlerState { manager: services.ws_manager.clone(), - router: Arc::new(NoopMessageRouter), + router: Arc::new(services.dynamic_tool_registry.clone()), token_validator: Arc::new(|_| true), token_extractor: Arc::new(|_| Some("local".into())), }; @@ -835,7 +835,7 @@ pub fn build_ws_state(services: &AppServices) -> WsHandlerState { WsHandlerState { manager: services.ws_manager.clone(), - router: Arc::new(NoopMessageRouter), + router: Arc::new(services.dynamic_tool_registry.clone()), token_validator, token_extractor, } diff --git a/crates/aionui-app/src/services.rs b/crates/aionui-app/src/services.rs index 39b61ba20..cd603e23a 100644 --- a/crates/aionui-app/src/services.rs +++ b/crates/aionui-app/src/services.rs @@ -5,8 +5,8 @@ use std::sync::Arc; use crate::config::{AppConfig, derive_encryption_key}; use aionui_ai_agent::{ - AcpSessionSyncService, AcpSkillManager, ActiveLeaseRegistry, AgentFactoryDeps, AgentRegistry, IWorkerTaskManager, - WorkerTaskManagerImpl, build_agent_factory, + AcpSessionSyncService, AcpSkillManager, ActiveLeaseRegistry, AgentFactoryDeps, AgentRegistry, DynamicToolRegistry, + IWorkerTaskManager, WorkerTaskManagerImpl, build_agent_factory, }; use aionui_auth::{CookieConfig, JwtService, QrTokenStore, resolve_jwt_secret}; use aionui_common::OnConversationDelete; @@ -27,6 +27,7 @@ pub struct AppServices { pub cookie_config: Arc, pub qr_token_store: Arc, pub ws_manager: Arc, + pub dynamic_tool_registry: DynamicToolRegistry, pub event_bus: Arc, pub worker_task_manager: Arc, pub active_lease_registry: Arc, @@ -122,6 +123,8 @@ impl AppServices { let provider_repo = Arc::new(SqliteProviderRepository::new(database.pool().clone())); let event_bus = Arc::new(BroadcastEventBus::new(256)); + let ws_manager = Arc::new(WebSocketManager::new()); + let dynamic_tool_registry = DynamicToolRegistry::new(ws_manager.clone()); // User-configured MCP servers — injected into ACP `session/new` // so the agent gets the operator's tools (ELECTRON-1JG fix). let mcp_server_repo: Arc = @@ -175,6 +178,7 @@ impl AppServices { broadcaster: event_bus.clone(), backend_binary_path: backend_binary_path.clone(), mcp_server_repo: Some(mcp_server_repo), + dynamic_tool_registry: dynamic_tool_registry.clone(), }); // Agent factory is now wired. Future extension/custom agents @@ -208,7 +212,8 @@ impl AppServices { user_repo, cookie_config: Arc::new(CookieConfig::from_env()), qr_token_store: Arc::new(QrTokenStore::new()), - ws_manager: Arc::new(WebSocketManager::new()), + ws_manager, + dynamic_tool_registry, event_bus, worker_task_manager, active_lease_registry, diff --git a/crates/aionui-app/tests/websocket_e2e.rs b/crates/aionui-app/tests/websocket_e2e.rs index 481190926..1a8b30912 100644 --- a/crates/aionui-app/tests/websocket_e2e.rs +++ b/crates/aionui-app/tests/websocket_e2e.rs @@ -7,7 +7,7 @@ use std::net::SocketAddr; use std::sync::Arc; use std::time::Duration; -use aionui_api_types::WebSocketMessage; +use aionui_api_types::{DynamicToolCallParams, WebSocketMessage}; use aionui_app::{AppConfig, AppServices, create_router}; use aionui_realtime::WebSocketManager; use futures_util::{SinkExt, StreamExt}; @@ -542,3 +542,203 @@ async fn t7_2_blacklisted_token_rejected() { let code = read_close(&mut rx).await; assert_eq!(code, Some(1008)); } + +// =========================================================================== +// T8 — Dynamic tool registration and callbacks +// =========================================================================== + +fn dynamic_tools_register(request_id: &str, conversation_id: &str) -> Value { + json!({ + "name": "agent.dynamicToolsRegister", + "data": { + "requestId": request_id, + "conversationId": conversation_id, + "tools": [{ + "type": "function", + "name": "list_threads", + "description": "List tasks", + "inputSchema": {"type": "object"} + }] + } + }) +} + +fn dynamic_tool_result(call: &Value, registration_id: &str, success: bool) -> Value { + json!({ + "name": "agent.dynamicToolResult", + "data": { + "conversationId": call["data"]["conversationId"], + "registrationId": registration_id, + "threadId": call["data"]["threadId"], + "turnId": call["data"]["turnId"], + "callId": call["data"]["callId"], + "namespace": call["data"]["namespace"], + "tool": call["data"]["tool"], + "contentItems": [{"type": "inputText", "text": "ready"}], + "success": success + } + }) +} + +#[tokio::test] +async fn t8_1_dynamic_tool_round_trip_is_connection_and_identity_bound() { + let app = start_app().await; + let token = sign_token(&app, "user1"); + let (mut owner_tx, mut owner_rx) = connect_bearer(app.addr, &token).await; + let (mut other_tx, mut other_rx) = connect_bearer(app.addr, &token).await; + + let register = dynamic_tools_register("request-1", "conversation-1"); + owner_tx.send(send_json(®ister.to_string())).await.unwrap(); + let ack = read_text(&mut owner_rx).await; + assert_eq!(ack["name"], "agent.dynamicToolsRegistered"); + assert_eq!(ack["data"]["requestId"], "request-1"); + assert_eq!(ack["data"]["accepted"], true); + let registration_id = ack["data"]["registrationId"].as_str().unwrap().to_owned(); + + let session = app + .services + .dynamic_tool_registry + .session_for("conversation-1") + .unwrap(); + assert_eq!(session.metadata()["codex/dynamic_tools"]["version"], 1); + assert!(session.bind_thread("thread-1")); + let dispatch = tokio::spawn(async move { + session + .dispatch(DynamicToolCallParams { + thread_id: "thread-1".into(), + turn_id: "turn-1".into(), + call_id: "call-1".into(), + namespace: None, + tool: "list_threads".into(), + arguments: json!({"scope": "current_project"}), + }) + .await + }); + + let call = read_text(&mut owner_rx).await; + assert_eq!(call["name"], "agent.dynamicToolCall"); + assert_eq!(call["data"]["registrationId"], registration_id); + assert_eq!(call["data"]["arguments"]["scope"], "current_project"); + assert!( + tokio::time::timeout(Duration::from_millis(100), other_rx.next()) + .await + .is_err() + ); + + let result = dynamic_tool_result(&call, ®istration_id, true); + other_tx.send(send_json(&result.to_string())).await.unwrap(); + tokio::time::sleep(Duration::from_millis(50)).await; + assert!( + !dispatch.is_finished(), + "a result from another connection must be ignored" + ); + + owner_tx.send(send_json(&result.to_string())).await.unwrap(); + let response = tokio::time::timeout(Duration::from_secs(1), dispatch) + .await + .expect("dynamic tool dispatch timed out") + .unwrap(); + assert!(response.success); + assert_eq!(response.content_items.len(), 1); +} + +#[tokio::test] +async fn t8_2_one_connection_routes_multiple_conversations_without_identity_loss() { + let app = start_app().await; + let token = sign_token(&app, "user1"); + let (mut tx, mut rx) = connect_bearer(app.addr, &token).await; + + let mut registration_ids = std::collections::HashMap::new(); + for conversation_id in ["conversation-1", "conversation-2"] { + let register = dynamic_tools_register(&format!("request-{conversation_id}"), conversation_id); + tx.send(send_json(®ister.to_string())).await.unwrap(); + let ack = read_text(&mut rx).await; + assert_eq!(ack["data"]["accepted"], true); + registration_ids.insert( + conversation_id.to_owned(), + ack["data"]["registrationId"].as_str().unwrap().to_owned(), + ); + } + + let mut dispatches = Vec::new(); + for (conversation_id, thread_id, call_id) in [ + ("conversation-1", "thread-1", "call-1"), + ("conversation-2", "thread-2", "call-2"), + ] { + let session = app.services.dynamic_tool_registry.session_for(conversation_id).unwrap(); + assert!(session.bind_thread(thread_id)); + dispatches.push(tokio::spawn({ + let thread_id = thread_id.to_owned(); + let call_id = call_id.to_owned(); + async move { + session + .dispatch(DynamicToolCallParams { + thread_id, + turn_id: "turn-1".into(), + call_id, + namespace: None, + tool: "list_threads".into(), + arguments: json!({}), + }) + .await + } + })); + } + + for _ in 0..2 { + let call = read_text(&mut rx).await; + let conversation_id = call["data"]["conversationId"].as_str().unwrap(); + let registration_id = registration_ids.get(conversation_id).unwrap(); + let result = dynamic_tool_result(&call, registration_id, true); + tx.send(send_json(&result.to_string())).await.unwrap(); + } + for dispatch in dispatches { + assert!(dispatch.await.unwrap().success); + } +} + +#[tokio::test] +async fn t8_3_websocket_disconnect_revokes_registration_and_pending_call() { + let app = start_app().await; + let token = sign_token(&app, "user1"); + let (mut tx, mut rx) = connect_bearer(app.addr, &token).await; + + let register = dynamic_tools_register("request-1", "conversation-1"); + tx.send(send_json(®ister.to_string())).await.unwrap(); + let ack = read_text(&mut rx).await; + assert_eq!(ack["data"]["accepted"], true); + + let session = app + .services + .dynamic_tool_registry + .session_for("conversation-1") + .unwrap(); + assert!(session.bind_thread("thread-1")); + let dispatch = tokio::spawn(async move { + session + .dispatch(DynamicToolCallParams { + thread_id: "thread-1".into(), + turn_id: "turn-1".into(), + call_id: "call-1".into(), + namespace: None, + tool: "list_threads".into(), + arguments: json!({}), + }) + .await + }); + let call = read_text(&mut rx).await; + assert_eq!(call["name"], "agent.dynamicToolCall"); + + tx.send(tungstenite::Message::Close(None)).await.unwrap(); + let response = tokio::time::timeout(Duration::from_secs(1), dispatch) + .await + .expect("disconnect did not release pending call") + .unwrap(); + assert!(!response.success); + assert!( + app.services + .dynamic_tool_registry + .session_for("conversation-1") + .is_none() + ); +} diff --git a/crates/aionui-realtime/src/handler.rs b/crates/aionui-realtime/src/handler.rs index 16ccc5e38..f22070eae 100644 --- a/crates/aionui-realtime/src/handler.rs +++ b/crates/aionui-realtime/src/handler.rs @@ -87,6 +87,7 @@ async fn handle_socket(socket: WebSocket, token: Option, state: WsHandle // Recv loop exited — client disconnected or errored. send_handle.abort(); + state.router.disconnected(conn_id); state.manager.remove_client(conn_id); info!(%conn_id, "websocket connection closed"); } diff --git a/crates/aionui-realtime/src/router.rs b/crates/aionui-realtime/src/router.rs index dd5d479d3..8b080eccd 100644 --- a/crates/aionui-realtime/src/router.rs +++ b/crates/aionui-realtime/src/router.rs @@ -11,6 +11,9 @@ pub trait MessageRouter: Send + Sync { /// Called for any message whose `name` is not handled internally /// by the WebSocket layer (i.e. not `pong` or `subscribe-show-open`). fn route(&self, conn_id: ConnectionId, name: &str, data: serde_json::Value) -> bool; + + /// Release any state owned by a WebSocket connection. + fn disconnected(&self, _conn_id: ConnectionId) {} } /// A no-op message router that reports every message as unhandled.