From fa92f3630c04dea2c71a0c755ae9f4d29443fc63 Mon Sep 17 00:00:00 2001 From: Wolf Mermelstein Date: Tue, 8 Sep 2026 21:39:57 +0000 Subject: [PATCH 1/5] feat(agent): consolidate inmem product tools behind MCP elicitation --- Cargo.lock | 26 ++ .../component/parts/ElicitationPart.tsx | 5 + .../state/elicitation-review-sink.test.ts | 25 ++ .../state/elicitation-review-sink.ts | 13 +- crates/agent/src/agent_loop.rs | 21 +- crates/agent/src/hook.rs | 89 +--- crates/agent/src/lib.rs | 1 - .../src/test/agent_loop/test_user_tools.rs | 73 +-- crates/agent/src/test/test_hook.rs | 209 +-------- crates/agent_egress/src/domain/service.rs | 8 +- .../agent_egress/src/domain/service/test.rs | 47 +- crates/agent_inmem/Cargo.toml | 2 +- crates/agent_inmem/src/domain/agent.rs | 231 ++++------ crates/agent_inmem/src/domain/agent/test.rs | 315 ++++++------- crates/agent_inmem/src/domain/engine.rs | 10 +- crates/agent_inmem/src/domain/mcp.rs | 57 +-- crates/agent_inmem/src/domain/mcp/test.rs | 4 +- crates/agent_inmem/src/domain/user_input.rs | 11 + crates/agent_inmem/src/outbound/acp_mcp.rs | 147 +++++- .../agent_inmem/src/outbound/acp_mcp/test.rs | 74 +++ crates/agent_inmem/src/outbound/manager.rs | 8 +- .../src/outbound/rig_engine/test.rs | 62 ++- crates/agent_inmem/src/rig_engine.rs | 85 ++-- crates/ai_tools/src/lib.rs | 49 +- crates/ai_tools/src/test.rs | 34 +- crates/ai_tools/src/user_tool_review.rs | 424 ------------------ crates/ai_tools/src/user_tool_review/test.rs | 338 -------------- crates/mcp_toolset/Cargo.toml | 2 + crates/mcp_toolset/src/lib.rs | 2 +- crates/mcp_toolset/src/toolset.rs | 65 ++- crates/mcp_toolset/src/toolset/test.rs | 84 ++++ crates/pipedream_mcp/src/outbound/api.rs | 2 +- docs/ACP_ELICITATION.md | 131 +++--- docs/AGENT_GUIDE/ai-chat.md | 9 + docs/MCP_TOOL_CONSOLIDATION.md | 67 +++ infra/stacks/mcp-server/mcp-server.ts | 14 + .../mcp_auth_proxy/src/inbound/axum_router.rs | 4 + .../src/inbound/axum_router/test.rs | 31 ++ services/mcp_service/Cargo.toml | 11 +- services/mcp_service/src/main.rs | 81 +++- services/mcp_service/src/session_routing.rs | 7 + .../src/session_routing/directory.rs | 53 +++ .../src/session_routing/directory/test.rs | 36 ++ .../mcp_service/src/session_routing/http.rs | 253 +++++++++++ .../mcp_service/src/session_routing/redis.rs | 154 +++++++ .../src/session_routing/redis/test.rs | 31 ++ services/mcp_service/src/tool_service.rs | 138 +++++- .../mcp_service/src/tool_service/review.rs | 205 +++++++++ .../src/tool_service/review/test.rs | 71 +++ services/mcp_service/src/tool_service/test.rs | 2 + .../src/tool_service/test/transport.rs | 352 +++++++++++++++ 51 files changed, 2422 insertions(+), 1751 deletions(-) create mode 100644 apps/web/src/features/block-agent/state/elicitation-review-sink.test.ts delete mode 100644 crates/ai_tools/src/user_tool_review.rs delete mode 100644 crates/ai_tools/src/user_tool_review/test.rs create mode 100644 crates/mcp_toolset/src/toolset/test.rs create mode 100644 docs/MCP_TOOL_CONSOLIDATION.md create mode 100644 services/mcp_auth_proxy/src/inbound/axum_router/test.rs create mode 100644 services/mcp_service/src/session_routing.rs create mode 100644 services/mcp_service/src/session_routing/directory.rs create mode 100644 services/mcp_service/src/session_routing/directory/test.rs create mode 100644 services/mcp_service/src/session_routing/http.rs create mode 100644 services/mcp_service/src/session_routing/redis.rs create mode 100644 services/mcp_service/src/session_routing/redis/test.rs create mode 100644 services/mcp_service/src/tool_service/review.rs create mode 100644 services/mcp_service/src/tool_service/review/test.rs create mode 100644 services/mcp_service/src/tool_service/test/transport.rs diff --git a/Cargo.lock b/Cargo.lock index 873940d601c..c05c5f40e55 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10246,6 +10246,7 @@ dependencies = [ "aws-sdk-secretsmanager", "aws-sdk-sqs", "axum", + "base64 0.22.1", "call", "channels", "chat", @@ -10261,6 +10262,7 @@ dependencies = [ "foreign_entity", "frecency", "fusionauth", + "futures", "http 1.4.0", "lexical_client", "macro_auth", @@ -10273,16 +10275,20 @@ dependencies = [ "macro_queues", "macro_service_urls", "macro_user_id", + "macro_uuid", "mcp_auth_proxy", "notification", "prompt", + "pulldown-cmark", "readonly_pool", "redis", "reminders", "reqwest 0.13.4", "rmcp", + "schemars 1.2.1", "search_service_client", "secretsmanager_client", + "serde", "serde_json", "soup", "sqlx", @@ -10290,6 +10296,7 @@ dependencies = [ "sync_service_client", "tokio", "tokio-util", + "tower 0.5.3", "tracing", "url", "workspace-hack", @@ -10307,6 +10314,7 @@ dependencies = [ "schemars 1.2.1", "serde_json", "thiserror 2.0.18", + "tokio", "tracing", "workspace-hack", ] @@ -12975,6 +12983,24 @@ dependencies = [ "prost 0.12.6", ] +[[package]] +name = "pulldown-cmark" +version = "0.13.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e9f068eba8e7071c5f9511831b44f32c740d5adf574e990f946ddb53db2f314e" +dependencies = [ + "bitflags 2.11.1", + "memchr", + "pulldown-cmark-escape", + "unicase", +] + +[[package]] +name = "pulldown-cmark-escape" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "007d8adb5ddab6f8e3f491ac63566a7d5002cc7ed73901f72057943fa71ae1ae" + [[package]] name = "pxfm" version = "0.1.29" diff --git a/apps/web/src/features/block-agent/component/parts/ElicitationPart.tsx b/apps/web/src/features/block-agent/component/parts/ElicitationPart.tsx index 5394b4ce1f5..dcf126a587d 100644 --- a/apps/web/src/features/block-agent/component/parts/ElicitationPart.tsx +++ b/apps/web/src/features/block-agent/component/parts/ElicitationPart.tsx @@ -233,6 +233,11 @@ function LiveUserTool(props: { }); const sink = () => createElicitationReviewSink({ + encodedEmailBody: + props.request.tool === 'SendEmail' && + props.request.schema.properties.some( + (property) => property.name === 'bodyFormat' + ), canAnswer: elicitation.canAnswer, ownerName: elicitation.ownerName, answering: elicitation.answering, diff --git a/apps/web/src/features/block-agent/state/elicitation-review-sink.test.ts b/apps/web/src/features/block-agent/state/elicitation-review-sink.test.ts new file mode 100644 index 00000000000..c27abcf15c8 --- /dev/null +++ b/apps/web/src/features/block-agent/state/elicitation-review-sink.test.ts @@ -0,0 +1,25 @@ +import { describe, expect, it, vi } from 'vitest'; +import { createElicitationReviewSink } from './elicitation-review-sink'; + +describe('MCP composer answers', () => { + it('declares encoded email bodies only when the form supports that field', async () => { + for (const encodedEmailBody of [false, true]) { + const respond = vi.fn().mockResolvedValue(true); + const sink = createElicitationReviewSink({ + encodedEmailBody, + canAnswer: () => true, + ownerName: () => 'Alice', + answering: () => false, + respond, + }); + await sink.onExecute({ body: 'PHA-SGVsbG88L3A-' }); + expect(respond).toHaveBeenCalledWith({ + action: 'accept', + content: { + draft: JSON.stringify({ body: 'PHA-SGVsbG88L3A-' }), + ...(encodedEmailBody ? { bodyFormat: 'base64url_html' } : {}), + }, + }); + } + }); +}); diff --git a/apps/web/src/features/block-agent/state/elicitation-review-sink.ts b/apps/web/src/features/block-agent/state/elicitation-review-sink.ts index e03ea948055..2a7678560ba 100644 --- a/apps/web/src/features/block-agent/state/elicitation-review-sink.ts +++ b/apps/web/src/features/block-agent/state/elicitation-review-sink.ts @@ -2,8 +2,8 @@ * The agent session's half of a user tool's composer: answering the review * elicitation the agent is blocked on. * - * Accept sends the whole edited draft under the `draft` field (Macro's - * `_macro/json` extension); the agent's finisher runs the tool with it. + * Accept sends the whole edited draft as a standard string `draft` field. + * The MCP server validates and executes the reviewed tool arguments. * Reject declines. There is nothing to persist between edits - the draft * lives in the form until the user decides - so `onEdit` is left out. */ @@ -16,6 +16,8 @@ import type { Accessor } from 'solid-js'; export const DRAFT_FIELD = 'draft'; export function createElicitationReviewSink(options: { + /** New MCP email forms declare the composer body encoding explicitly. */ + encodedEmailBody?: boolean; canAnswer: Accessor; ownerName: Accessor; answering: Accessor; @@ -32,7 +34,12 @@ export function createElicitationReviewSink(options: { canAct() ? options.respond({ action: 'accept', - content: { [DRAFT_FIELD]: JSON.stringify(args) }, + content: { + [DRAFT_FIELD]: JSON.stringify(args), + ...(options.encodedEmailBody + ? { bodyFormat: 'base64url_html' } + : {}), + }, }) : Promise.resolve(false), onReject: () => diff --git a/crates/agent/src/agent_loop.rs b/crates/agent/src/agent_loop.rs index 406c9175b47..48f86d5f773 100644 --- a/crates/agent/src/agent_loop.rs +++ b/crates/agent/src/agent_loop.rs @@ -1,6 +1,6 @@ /// The main entry point: [`AgentLoop`] and [`Session`]. use crate::error::AgentError; -use crate::hook::{BridgeInputs, RegisterFn, ToolRouter, UserToolFinisher}; +use crate::hook::{BridgeInputs, RegisterFn, ToolRouter}; use crate::model::PredefinedModel; use crate::model::router::{ModelRouter, ProviderAgent}; use crate::stream::ChatCompletionStream; @@ -30,7 +30,6 @@ pub struct AgentLoop { max_turns: usize, max_tokens: u64, recorder: Arc, - user_tool_finisher: Option, } impl AgentLoop { @@ -47,24 +46,9 @@ impl AgentLoop { max_turns: DEFAULT_MAX_TURNS, max_tokens: DEFAULT_MAX_TOKENS, recorder, - user_tool_finisher: None, } } - /// Finish user tools inside the turn. - /// - /// A user tool (`ai_toolset::UserTool`) answers `"PendingUserExecution"` - /// and leaves the call for the host to finish. Without a finisher that - /// answer reaches the model as-is and the host finishes the call later, - /// as chat does over HTTP. With one, the bridge hands each pending call - /// to `finisher` before the model reads it, and the model sees what the - /// user decided instead - the shape a host that can reach its user - /// mid-turn wants. - pub fn with_user_tool_finisher(mut self, finisher: UserToolFinisher) -> Self { - self.user_tool_finisher = Some(finisher); - self - } - /// Override the model. /// /// Accepts any stringifiable id — an [`AgentModel`] (backend) or a raw @@ -248,7 +232,6 @@ impl AgentLoop { routing, loaded_buffer, register_loaded, - user_tool_finisher: self.user_tool_finisher.clone(), }, recorder: self.recorder.clone(), usage_ctx, @@ -293,7 +276,7 @@ pub struct Session { history: Vec, max_turns: usize, /// What every turn's stream bridge is built from: tool routing, the - /// on-demand tool loading pair, and the user-tool finisher if any. + /// on-demand tool loading pair. bridge_inputs: BridgeInputs, recorder: Arc, usage_ctx: UsageContext, diff --git a/crates/agent/src/hook.rs b/crates/agent/src/hook.rs index 68cb08d9a43..f3cb3efd24e 100644 --- a/crates/agent/src/hook.rs +++ b/crates/agent/src/hook.rs @@ -30,44 +30,6 @@ pub type ToolRouter = Arc Option + Send + Sync>; pub type RegisterFn = Arc) -> Pin + Send>> + Send + Sync>; -/// A user tool the model called, as the host's [`UserToolFinisher`] sees it. -/// -/// A user tool (`ai_toolset::UserTool`) answers `"PendingUserExecution"` and -/// does nothing: the host is meant to finish it - let the user review the -/// call, then execute or reject it. Chat does that after the turn, over HTTP; -/// a host that can reach its user mid-turn does it here, before the model -/// reads the result. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct PendingUserTool { - /// The tool's name as the toolset knows it. - pub tool_name: String, - /// The call's id as the stream reported it ([`ToolCall::id`]): the - /// provider's, or rig's correlation id when the provider gave none. The - /// id the host's transcript shows the call under. - pub tool_call_id: String, - /// The arguments the model called the tool with. - pub args: serde_json::Value, -} - -/// What a finished user tool comes back as, in place of the pending answer. -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum FinishedUserTool { - /// The tool's own result (`UserToolResponse` as JSON: the user's - /// action, or their rejection). - Result(serde_json::Value), - /// Finishing failed; the description is what the model reads. - Error(String), -} - -/// Finishes a user tool inside the turn. Returns `None` to leave the pending -/// answer as it is - the host will finish the call some other way, or not at -/// all. Context-erased for the same reason as [`RegisterFn`]. -pub type UserToolFinisher = Arc< - dyn Fn(PendingUserTool) -> Pin> + Send>> - + Send - + Sync, ->; - /// Everything the session hands the stream bridge besides the request context. #[derive(Clone)] pub struct BridgeInputs { @@ -77,13 +39,8 @@ pub struct BridgeInputs { pub loaded_buffer: Arc>>, /// Registers loaded tools with the live tool server. pub register_loaded: RegisterFn, - /// Finishes user tools mid-turn, when the host can. - pub user_tool_finisher: Option, } -/// The answer a user tool gives when it has not been finished. -const PENDING_USER_EXECUTION: &str = "PendingUserExecution"; - static CANCELLED_REASON: &str = "user cancelled"; /// Retry budget for invalid tool calls recovered via @@ -110,9 +67,6 @@ pub struct StreamBridge { /// [`Self::on_invalid_tool_call`] to recover calls to tools the model /// discovered but never loaded. searchable_catalog: Arc>, - /// Finishes a user tool's pending answer before the model reads it, on - /// hosts that can reach the user mid-turn (see [`UserToolFinisher`]). - user_tool_finisher: Option, /// the user has requested the stream stop cancel: CancellationToken, } @@ -137,7 +91,6 @@ impl StreamBridge { routing, loaded_buffer, register_loaded, - user_tool_finisher, } = inputs; let (tx, rx) = mpsc::unbounded_channel(); ( @@ -147,7 +100,6 @@ impl StreamBridge { loaded_buffer, register_loaded, searchable_catalog, - user_tool_finisher, cancel, }, rx, @@ -177,11 +129,6 @@ fn presentation_json(presentation: &ToolOutput) -> Option { .and_then(|text| serde_json::from_str(text).ok()) } -/// Whether a tool's JSON result is a user tool's unfinished answer. -fn is_pending_user_execution(json: &serde_json::Value) -> bool { - json.as_str() == Some(PENDING_USER_EXECUTION) -} - /// The hook bodies, as inherent methods so tests can exercise them directly: /// rig's [`HookContext`] has no public constructor, so the [`AgentHook`] impl /// below is a thin delegation layer over these. @@ -274,7 +221,7 @@ impl StreamBridge { tool_name: &str, tool_call_id: Option<&str>, internal_call_id: &str, - args: &str, + _args: &str, presentation: &ToolOutput, is_success: bool, ) -> ToolResultAction { @@ -297,40 +244,6 @@ impl StreamBridge { None }; - // A user tool's pending answer is finished here when the host can: - // the user reviews the call while the turn waits, and the model reads - // what they decided instead of a "pending" it would take for success. - if let Some(finisher) = &self.user_tool_finisher - && json.as_ref().is_some_and(is_pending_user_execution) - { - let call = PendingUserTool { - tool_name: tool_name.to_owned(), - tool_call_id: id.clone(), - args: serde_json::from_str(args).unwrap_or(serde_json::Value::Null), - }; - match finisher(call).await { - Some(FinishedUserTool::Result(result)) => { - let _ = self - .tx - .send(Ok(StreamPart::ToolResponse(ToolResponse::Json { - id, - json: result.clone(), - name: tool_name.to_owned(), - }))); - return ToolResultAction::Rewrite(ToolOutput::json(result)); - } - Some(FinishedUserTool::Error(description)) => { - let _ = self.tx.send(Ok(StreamPart::ToolResponse(ToolResponse::Err { - id, - name: tool_name.to_owned(), - description: description.clone(), - }))); - return ToolResultAction::rewrite(description); - } - None => {} - } - } - let response = if let Some(json) = json { ToolResponse::Json { id, diff --git a/crates/agent/src/lib.rs b/crates/agent/src/lib.rs index 99154451b29..c0732f6d90d 100644 --- a/crates/agent/src/lib.rs +++ b/crates/agent/src/lib.rs @@ -23,7 +23,6 @@ pub use agent_loop::{AgentLoop, Session}; pub use completion::{complete, complete_with_history}; pub use convert::{merge_consecutive_parts, to_rig_messages}; pub use error::AgentError; -pub use hook::{FinishedUserTool, PendingUserTool, UserToolFinisher}; pub use model::PredefinedModel; pub use stream::{ChatCompletionStream, McpInfo, StreamPart, ToolCall, ToolResponse, Usage}; pub use tool_adapter::{DynToolSetAdapter, ToolsetToolAdapter, normalize_request_schema}; diff --git a/crates/agent/src/test/agent_loop/test_user_tools.rs b/crates/agent/src/test/agent_loop/test_user_tools.rs index c692d8286c0..fbb75013e47 100644 --- a/crates/agent/src/test/agent_loop/test_user_tools.rs +++ b/crates/agent/src/test/agent_loop/test_user_tools.rs @@ -1,9 +1,6 @@ -//! User tools through the whole loop: the tool answers pending, the session's -//! finisher (when it has one) settles the call before the model reads it, and -//! the model's next request carries what the user decided. +//! Chat user tools still return pending for the existing composer flow. use super::util; -use crate::hook::{FinishedUserTool, PendingUserTool, UserToolFinisher}; use crate::stream::ToolResponse; use ai_toolset::{ AsyncTool, AsyncToolCollection, RequestContext, ServiceContext, ToolAnnotated, ToolAnnotations, @@ -15,8 +12,7 @@ use rig_core::message::{Message, UserContent}; use rig_core::test_utils::{MockCompletionModel, MockStreamEvent}; use schemars::JsonSchema; use serde::{Deserialize, Serialize}; -use std::pin::Pin; -use std::sync::{Arc, Mutex}; +use std::sync::Arc; #[derive(Deserialize, JsonSchema)] #[schemars( @@ -91,7 +87,7 @@ fn tool_result_shown(requests: &[CompletionRequest]) -> String { } #[tokio::test] -async fn without_a_finisher_the_model_reads_the_pending_answer() { +async fn chat_reads_the_deferred_user_tool_answer() { let model = call_then_done(); let toolset = util::tool_set(AsyncToolCollection::<()>::new().add_user_tool::()); let mut session = util::session(toolset, Arc::new(()), model.clone()).await; @@ -109,66 +105,3 @@ async fn without_a_finisher_the_model_reads_the_pending_answer() { assert!(tool_result_shown(&model.requests()).contains("PendingUserExecution")); assert_eq!(result.content(), "done"); } - -#[tokio::test] -async fn with_a_finisher_the_model_reads_what_the_user_decided() { - let model = call_then_done(); - let seen: Arc>> = Arc::new(Mutex::new(Vec::new())); - let finisher: UserToolFinisher = { - let seen = Arc::clone(&seen); - Arc::new(move |call: PendingUserTool| { - let seen = Arc::clone(&seen); - Box::pin(async move { - seen.lock().unwrap().push(call); - Some(FinishedUserTool::Result( - serde_json::json!({"UserAction": {"delivered": "hello, edited"}}), - )) - }) as Pin> + Send>> - }) - }; - let toolset = util::tool_set(AsyncToolCollection::<()>::new().add_user_tool::()); - let mut session = util::test_loop() - .with_user_tool_finisher(finisher) - .test_session( - toolset, - Arc::new(()), - "test preamble", - util::usage_ctx(), - model.clone(), - ) - .await; - - let result = util::drive(&mut session, "send hello").await; - - let call = result.tool_calls()[0].clone(); - assert_eq!( - &*seen.lock().unwrap(), - &[PendingUserTool { - tool_name: "SendNote".to_owned(), - tool_call_id: call.id.clone(), - args: serde_json::json!({ "text": "hello" }), - }], - "the finisher saw the call as the model made it, under the id the stream gave it" - ); - let response = result - .tool_response(&call.id) - .expect("the call was answered"); - assert!( - matches!( - response, - ToolResponse::Json { json, .. } - if json == &serde_json::json!({"UserAction": {"delivered": "hello, edited"}}) - ), - "the stream records the finished result, got {response:?}" - ); - let shown = tool_result_shown(&model.requests()); - assert!( - shown.contains("hello, edited"), - "the model read the finished result: {shown}" - ); - assert!( - !shown.contains("PendingUserExecution"), - "the pending answer never reached the model: {shown}" - ); - assert_eq!(result.content(), "done"); -} diff --git a/crates/agent/src/test/test_hook.rs b/crates/agent/src/test/test_hook.rs index 3852d366ecf..a839ce777a0 100644 --- a/crates/agent/src/test/test_hook.rs +++ b/crates/agent/src/test/test_hook.rs @@ -1,7 +1,7 @@ use crate::hook::*; -use crate::stream::{StreamPart, ToolResponse}; +use crate::stream::StreamPart; use ai_toolset::SearchableTool; -use rig_agent::agent::{InvalidToolCallAction, ToolCallAction, ToolResultAction}; +use rig_agent::agent::{InvalidToolCallAction, ToolCallAction}; use rig_agent::tool::ToolOutput; use schemars::Schema; use std::pin::Pin; @@ -33,37 +33,15 @@ fn recording_register() -> (RegisterFn, Arc>>) { } /// Bridge inputs with no routing and `register` as the registrar; the caller -/// supplies the loaded-tool buffer and, optionally, a user-tool finisher. -fn inputs( - loaded_buffer: Arc>>, - register: RegisterFn, - user_tool_finisher: Option, -) -> BridgeInputs { +/// supplies the loaded-tool buffer. +fn inputs(loaded_buffer: Arc>>, register: RegisterFn) -> BridgeInputs { BridgeInputs { routing: Arc::new(|_| None), loaded_buffer, register_loaded: register, - user_tool_finisher, } } -/// A finisher that records what it was handed and answers with `answer`. -fn recording_finisher( - answer: Option, -) -> (UserToolFinisher, Arc>>) { - let recorded = Arc::new(Mutex::new(Vec::new())); - let sink = recorded.clone(); - let finisher: UserToolFinisher = Arc::new(move |call: PendingUserTool| { - let sink = sink.clone(); - let answer = answer.clone(); - Box::pin(async move { - sink.lock().unwrap().push(call); - answer - }) as Pin> + Send>> - }); - (finisher, recorded) -} - #[tokio::test] async fn on_tool_result_drains_buffer_and_registers_loaded_tools() { let buffer = Arc::new(Mutex::new(vec![ @@ -72,11 +50,8 @@ async fn on_tool_result_drains_buffer_and_registers_loaded_tools() { ])); let (register, registered) = recording_register(); let token = CancellationToken::new(); - let (bridge, _rx) = StreamBridge::channel( - inputs(buffer.clone(), register, None), - Arc::new(vec![]), - token, - ); + let (bridge, _rx) = + StreamBridge::channel(inputs(buffer.clone(), register), Arc::new(vec![]), token); bridge .handle_tool_result( @@ -107,8 +82,7 @@ async fn on_tool_result_registers_nothing_when_buffer_empty() { let buffer = Arc::new(Mutex::new(Vec::new())); let (register, registered) = recording_register(); let token = CancellationToken::new(); - let (bridge, _rx) = - StreamBridge::channel(inputs(buffer, register, None), Arc::new(vec![]), token); + let (bridge, _rx) = StreamBridge::channel(inputs(buffer, register), Arc::new(vec![]), token); bridge .handle_tool_result( @@ -133,27 +107,13 @@ fn bare_bridge() -> ( ) { let (register, _registered) = recording_register(); StreamBridge::channel( - inputs(Arc::new(Mutex::new(Vec::new())), register, None), + inputs(Arc::new(Mutex::new(Vec::new())), register), Arc::new(vec![]), CancellationToken::new(), ) } /// A bare bridge whose user tools `finisher` finishes. -fn finishing_bridge( - finisher: UserToolFinisher, -) -> ( - StreamBridge, - tokio::sync::mpsc::UnboundedReceiver>, -) { - let (register, _registered) = recording_register(); - StreamBridge::channel( - inputs(Arc::new(Mutex::new(Vec::new())), register, Some(finisher)), - Arc::new(vec![]), - CancellationToken::new(), - ) -} - #[tokio::test] async fn on_tool_call_parses_object_args() { let (bridge, mut rx) = bare_bridge(); @@ -196,7 +156,7 @@ async fn invalid_call_to_searchable_tool_loads_it_and_retries() { let (register, registered) = recording_register(); let catalog = Arc::new(vec![searchable("mcp__linear__create_issue")]); let (bridge, _rx) = StreamBridge::channel( - inputs(Arc::new(Mutex::new(Vec::new())), register, None), + inputs(Arc::new(Mutex::new(Vec::new())), register), catalog, CancellationToken::new(), ); @@ -218,7 +178,7 @@ async fn invalid_call_to_unknown_tool_retries_with_feedback_without_loading() { let (register, registered) = recording_register(); let catalog = Arc::new(vec![searchable("mcp__linear__create_issue")]); let (bridge, _rx) = StreamBridge::channel( - inputs(Arc::new(Mutex::new(Vec::new())), register, None), + inputs(Arc::new(Mutex::new(Vec::new())), register), catalog, CancellationToken::new(), ); @@ -237,152 +197,3 @@ async fn invalid_call_to_unknown_tool_retries_with_feedback_without_loading() { } // --- user tools --- - -const PENDING: &str = "\"PendingUserExecution\""; - -/// The finisher gets the call as the model made it, and what it returns is -/// what the model reads and what the stream records - the pending answer -/// never leaves the bridge. -#[tokio::test] -async fn a_pending_user_tool_is_finished_before_the_model_reads_it() { - let created = serde_json::json!({"UserAction": {"eventId": "evt-1", "title": "Sync"}}); - let (finisher, seen) = recording_finisher(Some(FinishedUserTool::Result(created.clone()))); - let (bridge, mut rx) = finishing_bridge(finisher); - - let action = bridge - .handle_tool_result( - "CreateCalendarEvent", - Some("toolu_1"), - "internal-1", - "{\"title\":\"Sync\"}", - &ToolOutput::json(serde_json::Value::String("PendingUserExecution".into())), - true, - ) - .await; - - assert_eq!( - &*seen.lock().unwrap(), - &[PendingUserTool { - tool_name: "CreateCalendarEvent".to_owned(), - tool_call_id: "toolu_1".to_owned(), - args: serde_json::json!({"title": "Sync"}), - }] - ); - let ToolResultAction::Rewrite(shown) = action else { - panic!("the model's view is rewritten, got {action:?}"); - }; - assert_eq!(shown.as_json(), Some(&created)); - let Ok(StreamPart::ToolResponse(ToolResponse::Json { id, json, name })) = - rx.try_recv().unwrap() - else { - panic!("the stream records the finished result"); - }; - assert_eq!( - (id.as_str(), name.as_str()), - ("toolu_1", "CreateCalendarEvent") - ); - assert_eq!(json, created); -} - -#[tokio::test] -async fn a_user_tool_the_finisher_fails_reads_as_a_tool_error() { - let (finisher, _seen) = recording_finisher(Some(FinishedUserTool::Error( - "the user is already being asked something".to_owned(), - ))); - let (bridge, mut rx) = finishing_bridge(finisher); - - let action = bridge - .handle_tool_result( - "SendEmail", - Some("toolu_2"), - "internal-2", - "{}", - &ToolOutput::text(PENDING), - true, - ) - .await; - - let ToolResultAction::Rewrite(shown) = action else { - panic!("the model's view is rewritten, got {action:?}"); - }; - assert_eq!( - shown.as_text(), - Some("the user is already being asked something") - ); - let Ok(StreamPart::ToolResponse(ToolResponse::Err { description, .. })) = - rx.try_recv().unwrap() - else { - panic!("the stream records the failure"); - }; - assert_eq!(description, "the user is already being asked something"); -} - -/// A finisher that declines to act, and a bridge with no finisher at all, -/// both leave the pending answer for the host to finish later - chat's flow. -#[tokio::test] -async fn an_unfinished_user_tool_keeps_its_pending_answer() { - let (finisher, seen) = recording_finisher(None); - for (bridge, mut rx) in [finishing_bridge(finisher), bare_bridge()] { - let action = bridge - .handle_tool_result( - "SendEmail", - None, - "internal-3", - "{}", - &ToolOutput::text(PENDING), - true, - ) - .await; - assert!(matches!(action, ToolResultAction::Keep), "got {action:?}"); - let Ok(StreamPart::ToolResponse(ToolResponse::Json { json, .. })) = rx.try_recv().unwrap() - else { - panic!("the stream records the pending answer"); - }; - assert_eq!(json, serde_json::json!("PendingUserExecution")); - } - assert_eq!( - seen.lock().unwrap().len(), - 1, - "the declining finisher was asked once" - ); -} - -/// Only the pending answer is a user tool's: every other result, and every -/// failure, passes the finisher by. -#[tokio::test] -async fn results_other_than_pending_never_reach_the_finisher() { - let (finisher, seen) = - recording_finisher(Some(FinishedUserTool::Result(serde_json::json!("never")))); - let (bridge, mut rx) = finishing_bridge(finisher); - - bridge - .handle_tool_result( - "ListCalendars", - None, - "internal-4", - "{}", - &ToolOutput::json(serde_json::json!({"calendars": []})), - true, - ) - .await; - bridge - .handle_tool_result( - "CreateCalendarEvent", - None, - "internal-5", - "{}", - &ToolOutput::text(PENDING), - false, - ) - .await; - - assert!(seen.lock().unwrap().is_empty()); - assert!(matches!( - rx.try_recv().unwrap(), - Ok(StreamPart::ToolResponse(ToolResponse::Json { .. })) - )); - assert!(matches!( - rx.try_recv().unwrap(), - Ok(StreamPart::ToolResponse(ToolResponse::Err { .. })) - )); -} diff --git a/crates/agent_egress/src/domain/service.rs b/crates/agent_egress/src/domain/service.rs index e7b0fe24d97..ba9783e9e5c 100644 --- a/crates/agent_egress/src/domain/service.rs +++ b/crates/agent_egress/src/domain/service.rs @@ -92,11 +92,13 @@ where span.record("session", tracing::field::display(&grant.session)); span.record("owner", tracing::field::display(&grant.owner)); - // Staff-only for now, checked here so every target - git, connected - // MCP servers, Macro's own - passes one gate. The refusal names + // Macro tools belong to every authenticated session owner. Git and + // third-party MCP retain their staff-only rollout gate. The refusal names // itself ("not Macro staff") so the sandbox can report an actionable // reason; the reason is our own static wording, never the request's. - if !is_macro_staff(&grant.owner) { + if !is_macro_staff(&grant.owner) + && !matches!(&target, EgressTarget::McpServer(McpDestination::Macro)) + { tracing::warn!(owner = %grant.owner, "refusing egress for a session owned outside macro.com"); return Err(EgressError::Unauthenticated( "the session owner is not Macro staff", diff --git a/crates/agent_egress/src/domain/service/test.rs b/crates/agent_egress/src/domain/service/test.rs index 05a63d52ca8..dcaf3c41494 100644 --- a/crates/agent_egress/src/domain/service/test.rs +++ b/crates/agent_egress/src/domain/service/test.rs @@ -146,8 +146,18 @@ impl McpCredentials for SpyCredentials { owner: &MacroUserIdStr<'static>, destination: &McpDestination, ) -> Result { + if matches!(destination, McpDestination::Macro) { + self.asked + .lock() + .unwrap() + .push((owner.to_string(), "macro".into())); + return Ok(McpResolution::Connected(UpstreamCall::bearer( + Url::parse(&self.url).expect("url"), + BearerToken::new("upstream-token"), + )?)); + } let McpDestination::Connected(slug) = destination else { - unreachable!("these tests only dial connected servers"); + unreachable!() }; self.asked .lock() @@ -668,11 +678,9 @@ async fn strips_hop_by_hop_headers_from_the_response() { assert_eq!(names(response.headers()), ["mcp-session-id"]); } -/// The proxy is staff-only for now: a session owned outside macro.com gets -/// nothing, whatever its token says - told only, in our words, that staff -/// membership is what it lacks. +/// Third-party integrations remain restricted to staff. #[tokio::test] -async fn a_session_owned_outside_macro_gets_nothing() { +async fn a_session_owned_outside_macro_cannot_use_third_party_integrations() { let service = EgressServiceImpl::new( StubSessions(Ok(SessionGrant { session: AgentSessionId::new(), @@ -977,3 +985,32 @@ async fn a_cleartext_git_base_is_refused_too() { assert!(matches!(error, EgressError::InsecureUpstream(_))); assert!(!service.forward.was_called()); } + +#[tokio::test] +async fn a_nonstaff_session_can_reach_macro_with_its_verified_owner() { + let owner = MacroUserIdStr::try_from_email("visitor@example.com").unwrap(); + let service = EgressServiceImpl::new( + StubSessions(Ok(SessionGrant { + session: AgentSessionId::new(), + owner: owner.clone(), + repo: session_repo(), + mcp_servers: Vec::new(), + })), + SpyCredentials::knowing(), + SpyGithubTokens::default(), + SpyForwarder::answering(&[]), + ); + service + .proxy( + &SessionToken::new("token"), + EgressTarget::McpServer(McpDestination::Macro), + request(Method::POST, &[]), + ) + .await + .unwrap(); + assert_eq!( + *service.credentials.asked.lock().unwrap(), + vec![(owner.to_string(), "macro".to_owned())] + ); + assert!(service.forward.was_called()); +} diff --git a/crates/agent_inmem/Cargo.toml b/crates/agent_inmem/Cargo.toml index ce7bb90f9b6..efe7083ede1 100644 --- a/crates/agent_inmem/Cargo.toml +++ b/crates/agent_inmem/Cargo.toml @@ -29,7 +29,7 @@ mcp_select = { path = "../mcp_select" } mcp_toolset = { path = "../mcp_toolset" } memory = { path = "../memory" } prompt = { path = "../prompt" } -rmcp = { workspace = true, features = ["client", "transport-streamable-http-client"] } +rmcp = { workspace = true, features = ["client", "transport-streamable-http-client", "elicitation"] } schemars = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } diff --git a/crates/agent_inmem/src/domain/agent.rs b/crates/agent_inmem/src/domain/agent.rs index 0b1a695c3c9..ad076d84940 100644 --- a/crates/agent_inmem/src/domain/agent.rs +++ b/crates/agent_inmem/src/domain/agent.rs @@ -16,31 +16,26 @@ use std::time::Duration; use agent::types::{AssistantMessagePart, ChatMessage}; use agent::{StreamAccumulator, StreamPart, ToolResponse}; use agent_client_protocol::schema::v1::{ - AgentCapabilities, BooleanPropertySchema, CancelNotification, ContentBlock, ContentChunk, - CreateElicitationRequest, ElicitationAction, ElicitationFormMode, ElicitationPropertySchema, - ElicitationSchema, ElicitationSessionScope, EnumOption, Implementation, InitializeRequest, - InitializeResponse, IntegerPropertySchema, Meta, NewSessionRequest, NewSessionResponse, - NumberPropertySchema, OtherElicitationPropertySchema, PromptRequest, PromptResponse, - ResumeSessionRequest, ResumeSessionResponse, SessionCapabilities, SessionId, - SessionNotification, SessionResumeCapabilities, SessionUpdate, SetSessionConfigOptionRequest, - SetSessionConfigOptionResponse, StopReason, StringFormat, StringPropertySchema, - ToolCall as AcpToolCall, ToolCallId, ToolCallStatus, ToolCallUpdate, ToolCallUpdateFields, - ToolKind, + AgentCapabilities, CancelNotification, ContentBlock, ContentChunk, CreateElicitationRequest, + ElicitationAction, ElicitationFormMode, ElicitationSchema, ElicitationSessionScope, EnumOption, + Implementation, InitializeRequest, InitializeResponse, Meta, NewSessionRequest, + NewSessionResponse, PromptRequest, PromptResponse, ResumeSessionRequest, ResumeSessionResponse, + SessionCapabilities, SessionId, SessionNotification, SessionResumeCapabilities, SessionUpdate, + SetSessionConfigOptionRequest, SetSessionConfigOptionResponse, StopReason, + StringPropertySchema, ToolCall as AcpToolCall, ToolCallStatus, ToolCallUpdate, + ToolCallUpdateFields, ToolKind, }; use agent_client_protocol::{ Agent, Channel as AcpChannel, Client, ConnectionTo, Error as AcpError, }; use agent_runtime_protocol::domain::action::{COMPACT_COMMAND, MODEL_CONFIG_ID}; use agent_session::domain::model::AgentSessionId; -use ai_tools::user_tool_review::{ - ReviewError, ReviewFieldKind, ReviewForm, ReviewOutcome, ReviewRequest, UserToolReviewer, -}; use async_trait::async_trait; use macro_user_id::user_id::MacroUserIdStr; use tokio_util::sync::CancellationToken; use crate::domain::engine::{TurnEngine, TurnRequest}; -use crate::domain::mcp::{DynMcpToolConnector, dialable_servers}; +use crate::domain::mcp::{McpToolConnector, dialable_servers}; use crate::domain::session::{HistoryEntry, SessionStore, messages_for_turn}; use crate::domain::user_input::{ SharedUserInputRequester, UserInputError, UserInputOutcome, UserInputRequest, @@ -112,10 +107,14 @@ pub struct AgentState { /// a time. pub turn_lock: tokio::sync::Mutex<()>, /// Dials the MCP servers `session/new` and `session/resume` hand over. - pub mcp: Arc, + pub mcp: Arc, /// The tools of those servers, once dialed; `None` until then or when /// there were none. pub mcp_tools: Mutex>, + /// Shared across direct questions and requests forwarded from MCP. + pub(crate) awaiting_user: Arc, + /// The ACP client can display one form at a time. + pub(crate) input_gate: Arc>, /// Whether the client advertised `elicitation.form` on `initialize`. The /// protocol forbids asking a mode the client did not advertise. pub client_renders_forms: AtomicBool, @@ -129,12 +128,24 @@ impl AgentState { /// every turn that follows. Done at `session/new`/`session/resume`, the /// same moment a sandboxed harness connects its servers, so the first /// turn already has them. - async fn connect_mcp(&self, servers: Vec) { - let tools = self.mcp.connect_dyn(dialable_servers(servers)).await; + async fn connect_mcp( + &self, + servers: Vec, + connection: &ConnectionTo, + session_id: SessionId, + ) -> Result<(), AcpError> { + let input = user_input_requester(self, connection, session_id, &self.awaiting_user) + .map(|input| input as SharedUserInputRequester); + let tools = self + .mcp + .connect(dialable_servers(servers), input) + .await + .map_err(|error| AcpError::internal_error().data(error))?; *self .mcp_tools .lock() .expect("mcp tools lock should not be poisoned") = tools; + Ok(()) } fn current_mcp_tools(&self) -> Option { @@ -232,7 +243,7 @@ impl AgentState { /// the turn's requester, which counts each question it is waiting on, and the /// turn loop, which reads it to tell "waiting on the user" from "hung". #[derive(Default)] -struct AwaitingUser(AtomicUsize); +pub(crate) struct AwaitingUser(AtomicUsize); impl AwaitingUser { fn is_waiting(&self) -> bool { @@ -257,145 +268,49 @@ impl Drop for AwaitingGuard<'_> { /// ACP-backed user-input port for one connected session: the one place this /// agent sends `elicitation/create`, whether a tool is asking a question -/// (`AskUser`, [`UserInputRequester`]) or a user tool wants its call reviewed -/// ([`UserToolReviewer`]). +/// (`AskUser`, [`UserInputRequester`]) or an MCP server requests a form. struct AcpUserInputRequester { connection: ConnectionTo, session_id: SessionId, awaiting: Arc, + gate: Arc>, } -/// The key under `_meta.macro` naming the user tool an elicitation reviews, -/// so a Macro client can render the tool's own composer instead of the form. -const USER_TOOL_META_KEY: &str = "userTool"; - -/// The custom property type carrying the whole edited draft as a JSON string -/// (`_`-prefixed, as ACP reserves for implementation-specific extensions). -const JSON_PROPERTY_TYPE: &str = "_macro/json"; - #[async_trait] -impl UserToolReviewer for AcpUserInputRequester { - async fn review(&self, request: ReviewRequest) -> Result { +impl UserInputRequester for AcpUserInputRequester { + async fn form( + &self, + message: String, + schema: serde_json::Value, + meta: Option>, + ) -> Result { let _waiting = self.awaiting.begin(); - let scope = ElicitationSessionScope::new(self.session_id.clone()) - .tool_call_id(ToolCallId::new(request.tool_call_id.as_str())); - // The name lets a Macro client pick the tool's composer; the draft - // rides along so the fold has it even when the call the review is - // scoped to is not one it opened. - let mut ours = serde_json::Map::new(); - ours.insert( - USER_TOOL_META_KEY.to_owned(), - serde_json::json!({ "name": request.tool_name, "draft": request.draft }), + let _slot = self.gate.lock().await; + let schema = serde_json::from_value(schema).map_err(|error| { + UserInputError::InvalidAnswer(format!("unsupported form schema: {error}")) + })?; + let mut request = CreateElicitationRequest::new( + ElicitationFormMode::new( + ElicitationSessionScope::new(self.session_id.clone()), + schema, + ), + message, ); - let mut meta = Meta::new(); - meta.insert(META_NAMESPACE.to_owned(), serde_json::Value::Object(ours)); - let elicitation = CreateElicitationRequest::new( - ElicitationFormMode::new(scope, review_form_schema(&request.form)), - request.message, - ) - .meta(meta); - + if let Some(meta) = meta { + request = request.meta(meta); + } let response = self .connection - .send_request(elicitation) + .send_request(request) .block_task() .await - .map_err(|error| ReviewError::Unavailable(error.to_string()))?; - Ok(match response.action { - ElicitationAction::Accept(accept) => ReviewOutcome::Accepted( - accept - .content - .unwrap_or_default() - .into_iter() - .filter_map(|(name, value)| { - serde_json::to_value(value).ok().map(|value| (name, value)) - }) - .collect(), - ), - ElicitationAction::Decline => ReviewOutcome::Declined, - ElicitationAction::Cancel => ReviewOutcome::Cancelled, - _ => { - return Err(ReviewError::Failed( - "the client returned an unknown elicitation action".to_owned(), - )); - } - }) - } -} - -/// A review form as ACP's restricted schema. Fields are the draft's flat -/// arguments with their current values as defaults; the draft field is the -/// `_macro/json` extension a Macro client fills from its own composer. -fn review_form_schema(form: &ReviewForm) -> ElicitationSchema { - let mut schema = ElicitationSchema::new().title(form.title.clone()); - for field in &form.fields { - let required = form.required.contains(&field.name); - let property: ElicitationPropertySchema = match &field.kind { - ReviewFieldKind::Text { default, format } => StringPropertySchema::new() - .title(field.name.clone()) - .description(field.description.clone()) - .default_value(default.clone()) - .format(format.as_deref().and_then(string_format)) - .into(), - ReviewFieldKind::Boolean { default } => BooleanPropertySchema::new() - .title(field.name.clone()) - .description(field.description.clone()) - .default_value(*default) - .into(), - ReviewFieldKind::Number { default } => NumberPropertySchema::new() - .title(field.name.clone()) - .description(field.description.clone()) - .default_value(*default) - .into(), - ReviewFieldKind::Integer { default } => IntegerPropertySchema::new() - .title(field.name.clone()) - .description(field.description.clone()) - .default_value(*default) - .into(), - ReviewFieldKind::Choice { options, default } => StringPropertySchema::new() - .title(field.name.clone()) - .description(field.description.clone()) - .enum_values(options.clone()) - .default_value(default.clone()) - .into(), - ReviewFieldKind::Json => { - let mut fields = std::collections::BTreeMap::new(); - fields.insert( - "title".to_owned(), - serde_json::Value::String(field.name.clone()), - ); - if let Some(description) = &field.description { - fields.insert( - "description".to_owned(), - serde_json::Value::String(description.clone()), - ); - } - ElicitationPropertySchema::Other(OtherElicitationPropertySchema::new( - JSON_PROPERTY_TYPE, - fields, - )) - } - }; - schema = schema.property(field.name.clone(), property, required); - } - schema -} - -/// ACP's string format for a JSON Schema `format`, for the ones it names. -fn string_format(format: &str) -> Option { - match format { - "email" => Some(StringFormat::Email), - "uri" => Some(StringFormat::Uri), - "date" => Some(StringFormat::Date), - "date-time" => Some(StringFormat::DateTime), - _ => None, + .map_err(|error| UserInputError::RequestFailed(error.to_string()))?; + serde_json::to_value(response) + .map_err(|error| UserInputError::RequestFailed(error.to_string())) } -} - -#[async_trait] -impl UserInputRequester for AcpUserInputRequester { async fn ask(&self, request: UserInputRequest) -> Result { let _waiting = self.awaiting.begin(); + let _slot = self.gate.lock().await; let options = request.options; let mut field = StringPropertySchema::new().title("Answer"); if !options.is_empty() { @@ -464,6 +379,7 @@ fn user_input_requester( connection: connection.clone(), session_id, awaiting: Arc::clone(awaiting), + gate: Arc::clone(&state.input_gate), }) }) } @@ -499,11 +415,17 @@ pub async fn serve(state: Arc, acp: AcpChannel) -> Result<(), AcpErr .on_receive_request( { let state = Arc::clone(&state); - async move |request: NewSessionRequest, responder, _connection| { + async move |request: NewSessionRequest, responder, connection| { let state = Arc::clone(&state); + let _turn = state.turn_lock.lock().await; let acp_id = SessionId::new(macro_uuid::generate_uuid_v7().to_string()); + if let Err(error) = state + .connect_mcp(request.mcp_servers, &connection, acp_id.clone()) + .await + { + return responder.respond_with_error(error); + } state.bind_acp_session(acp_id.clone(), false); - state.connect_mcp(request.mcp_servers).await; responder.respond(NewSessionResponse::new(acp_id)) } }, @@ -512,14 +434,20 @@ pub async fn serve(state: Arc, acp: AcpChannel) -> Result<(), AcpErr .on_receive_request( { let state = Arc::clone(&state); - async move |request: ResumeSessionRequest, responder, _connection| { + async move |request: ResumeSessionRequest, responder, connection| { let state = Arc::clone(&state); + let _turn = state.turn_lock.lock().await; // Kept when the state already belongs to this ACP id - // either this process served the session, or a cold // attach replayed the frame log back into it (see // `domain::replay`). + if let Err(error) = state + .connect_mcp(request.mcp_servers, &connection, request.session_id.clone()) + .await + { + return responder.respond_with_error(error); + } state.bind_acp_session(request.session_id, true); - state.connect_mcp(request.mcp_servers).await; responder.respond(ResumeSessionResponse::new()) } }, @@ -645,7 +573,7 @@ async fn run_turn( model, instructions, } = state.turn_input(&prompt); - let awaiting = Arc::new(AwaitingUser::default()); + let awaiting = Arc::clone(&state.awaiting_user); let requester = user_input_requester(state, connection, acp_session_id.clone(), &awaiting); let mut parts = state.engine.run_turn(TurnRequest { owner: state.owner.clone(), @@ -657,7 +585,6 @@ async fn run_turn( user_input: requester .clone() .map(|requester| requester as SharedUserInputRequester), - reviewer: requester.map(|requester| requester as Arc), }); let mut accumulator = StreamAccumulator::new(); @@ -766,7 +693,8 @@ async fn run_ask( session_id: acp_session_id.clone(), // `/ask` has no idle timeout to hold off: it waits on the answer // directly rather than through the turn loop. - awaiting: Arc::new(AwaitingUser::default()), + awaiting: Arc::clone(&state.awaiting_user), + gate: Arc::clone(&state.input_gate), }; let text = match requester.ask(UserInputRequest { question, options }).await { Ok(UserInputOutcome::Answered(value)) => format!("You answered: {value}"), @@ -855,7 +783,12 @@ fn tool_call_meta(call: &agent::ToolCall) -> Meta { }; let mut ours = serde_json::Map::new(); ours.insert("toolName".to_owned(), serde_json::Value::String(tool_name)); - if call.mcp.is_none() && call.name == SUBAGENT_TOOL { + if (call.mcp.is_none() && call.name == SUBAGENT_TOOL) + || call + .mcp + .as_ref() + .is_some_and(|mcp| mcp.service == "macro" && mcp.tool_name == SUBAGENT_TOOL) + { ours.insert("subagent".to_owned(), serde_json::Value::Bool(true)); } let mut meta = Meta::new(); diff --git a/crates/agent_inmem/src/domain/agent/test.rs b/crates/agent_inmem/src/domain/agent/test.rs index f0bd38a021b..08b25d9bfe5 100644 --- a/crates/agent_inmem/src/domain/agent/test.rs +++ b/crates/agent_inmem/src/domain/agent/test.rs @@ -48,6 +48,14 @@ async fn with_agent( where Engine: TurnEngine, { + with_agent_connector(engine, Arc::new(crate::domain::mcp::NoMcpServers), scenario).await +} + +async fn with_agent_connector( + engine: Arc, + connector: Arc, + scenario: impl AsyncFnOnce(ConnectionTo, SessionId) -> Out, +) -> (Vec, Out) { let store = Arc::new(SessionStore::new()); let session_id = AgentSessionId::new(); store.insert( @@ -61,8 +69,10 @@ where store, active_cancel: std::sync::Mutex::new(Vec::new()), turn_lock: tokio::sync::Mutex::new(()), - mcp: Arc::new(crate::domain::mcp::NoMcpServers), + mcp: connector, mcp_tools: std::sync::Mutex::new(None), + awaiting_user: Default::default(), + input_gate: Default::default(), client_renders_forms: AtomicBool::new(false), enable_dev_commands: true, }); @@ -363,23 +373,24 @@ struct SpyConnector { asked: std::sync::Mutex>>, } +#[async_trait::async_trait] impl crate::domain::mcp::McpToolConnector for Arc { async fn connect( &self, servers: Vec, - ) -> Option { + _input: Option, + ) -> Result, String> { self.asked .lock() .expect("asked lock") .push(servers.into_iter().map(|server| server.name).collect()); - None + Ok(None) } } -/// The servers `session/new` carries are dialed then and there, minus Macro's -/// own, whose tools this runtime already has natively. +/// Session creation dials every advertised HTTP server, including Macro. #[tokio::test] -async fn session_new_dials_the_advertised_servers_except_macros_own() { +async fn session_new_dials_macro_and_other_advertised_servers() { use agent_client_protocol::schema::v1::{HttpHeader, McpServer as AcpMcpServer, McpServerHttp}; let spy = Arc::new(SpyConnector { @@ -401,6 +412,8 @@ async fn session_new_dials_the_advertised_servers_except_macros_own() { enable_dev_commands: false, mcp: Arc::new(Arc::clone(&spy)), mcp_tools: std::sync::Mutex::new(None), + awaiting_user: Default::default(), + input_gate: Default::default(), client_renders_forms: AtomicBool::new(false), }); let (client_channel, agent_channel) = AcpChannel::duplex(); @@ -438,7 +451,11 @@ async fn session_new_dials_the_advertised_servers_except_macros_own() { assert_eq!( *spy.asked.lock().expect("asked lock"), - vec![vec!["linear".to_owned(), "notion".to_owned()]] + vec![vec![ + "macro".to_owned(), + "linear".to_owned(), + "notion".to_owned() + ]] ); } @@ -489,6 +506,8 @@ where store, active_cancel: std::sync::Mutex::new(Vec::new()), turn_lock: tokio::sync::Mutex::new(()), + awaiting_user: Default::default(), + input_gate: Default::default(), client_renders_forms: AtomicBool::new(false), mcp: Arc::new(crate::domain::mcp::NoMcpServers), mcp_tools: std::sync::Mutex::new(None), @@ -789,167 +808,6 @@ async fn a_turn_waiting_on_the_user_outlasts_the_idle_timeout() { assert_eq!(spoken(¬ifications), "You said teal."); } -/// An engine whose turn puts one user tool call to the reviewer and says what -/// came back - the agent loop's finisher reduced to the ACP surface. -struct ReviewingEngine; - -impl TurnEngine for ReviewingEngine { - fn run_turn( - &self, - request: TurnRequest, - ) -> tokio::sync::mpsc::Receiver> { - let (parts, receiver) = tokio::sync::mpsc::channel(4); - tokio::spawn(async move { - let reviewer = request - .reviewer - .expect("the client advertised forms, so the turn can ask for a review"); - let outcome = reviewer - .review(ReviewRequest { - tool_name: "CreateCalendarEvent".to_owned(), - tool_call_id: "toolu_7".to_owned(), - message: "Create calendar event?".to_owned(), - draft: serde_json::json!({"title": "Q3 sync", "addGoogleMeet": false}), - form: ReviewForm { - title: Some("Create calendar event".to_owned()), - fields: vec![ - ai_tools::user_tool_review::ReviewField { - name: "title".to_owned(), - description: Some("The event title.".to_owned()), - kind: ReviewFieldKind::Text { - default: Some("Q3 sync".to_owned()), - format: None, - }, - }, - ai_tools::user_tool_review::ReviewField { - name: "addGoogleMeet".to_owned(), - description: None, - kind: ReviewFieldKind::Boolean { - default: Some(false), - }, - }, - ai_tools::user_tool_review::ReviewField { - name: "draft".to_owned(), - description: None, - kind: ReviewFieldKind::Json, - }, - ], - required: vec!["title".to_owned()], - }, - }) - .await; - let text = match outcome { - Ok(ReviewOutcome::Accepted(content)) => format!( - "accepted {}", - serde_json::to_string(&content).expect("content serializes") - ), - Ok(other) => format!("{other:?}"), - Err(error) => error.to_string(), - }; - let _ = parts.send(Ok(StreamPart::Content(text))).await; - }); - receiver - } -} - -/// A user tool's review goes out as a tool-call-scoped form elicitation the -/// fold and a Macro client can recognize: the draft's flat fields with their -/// values as defaults, the `_macro/json` draft field, and `_meta.macro.userTool` -/// naming the tool. The accepted content comes back as submitted. -#[tokio::test] -async fn a_user_tool_review_is_a_tool_scoped_form_elicitation_naming_the_tool() { - use agent_client_protocol::schema::v1::{ - ElicitationAcceptAction, ElicitationAction, ElicitationContentValue, ElicitationMode, - ElicitationPropertySchema, ElicitationScope, - }; - use std::collections::BTreeMap; - - let answer = - ElicitationAction::Accept(ElicitationAcceptAction::new().content(BTreeMap::from([ - ( - "title".to_owned(), - ElicitationContentValue::String("Q3 planning".to_owned()), - ), - ( - "addGoogleMeet".to_owned(), - ElicitationContentValue::Boolean(true), - ), - ]))); - let (notifications, asked, response) = with_asking_engine( - Arc::new(ReviewingEngine), - false, - answer, - Duration::ZERO, - async |connection, session| { - connection - .send_request(text_prompt(&session, "create the event")) - .block_task() - .await - .expect("the turn should complete") - }, - ) - .await; - - assert_eq!(response.stop_reason, StopReason::EndTurn); - assert_eq!(asked.len(), 1, "one review was asked"); - let request = &asked[0]; - assert_eq!(request.message, "Create calendar event?"); - let ElicitationMode::Form(form) = &request.mode else { - panic!("a form was asked"); - }; - let ElicitationScope::Session(scope) = &form.scope else { - panic!("session scoped"); - }; - assert_eq!( - scope.tool_call_id.as_ref().map(|id| id.0.as_ref()), - Some("toolu_7"), - "the elicitation names the call it reviews" - ); - let user_tool = request - .meta - .as_ref() - .and_then(|meta| meta.get(META_NAMESPACE)) - .and_then(|ours| ours.get(USER_TOOL_META_KEY)) - .expect("_meta.macro.userTool names the tool under review"); - assert_eq!( - user_tool.get("name").and_then(|name| name.as_str()), - Some("CreateCalendarEvent"), - "a Macro client learns which tool's composer to show" - ); - assert_eq!( - user_tool.get("draft"), - Some(&serde_json::json!({"title": "Q3 sync", "addGoogleMeet": false})), - "and has the draft whether or not it opened the call" - ); - let schema = &form.requested_schema; - assert_eq!(schema.title.as_deref(), Some("Create calendar event")); - let ElicitationPropertySchema::String(title) = &schema.properties["title"] else { - panic!("title is a string field"); - }; - assert_eq!(title.default.as_deref(), Some("Q3 sync")); - assert_eq!(title.description.as_deref(), Some("The event title.")); - let ElicitationPropertySchema::Boolean(meet) = &schema.properties["addGoogleMeet"] else { - panic!("addGoogleMeet is a boolean field"); - }; - assert_eq!(meet.default, Some(false)); - let ElicitationPropertySchema::Other(draft) = &schema.properties["draft"] else { - panic!( - "the draft is a custom field, got {:?}", - schema.properties["draft"] - ); - }; - assert_eq!(draft.type_, JSON_PROPERTY_TYPE); - assert_eq!( - schema.required.as_deref(), - Some(&["title".to_owned()][..]), - "the draft field is never required" - ); - assert_eq!( - spoken(¬ifications), - r#"accepted {"addGoogleMeet":true,"title":"Q3 planning"}"# - ); -} - -/// The timeout still guards a turn that is silent with nothing asked. #[tokio::test(start_paused = true)] async fn a_silent_turn_with_no_question_out_is_stopped_by_the_idle_timeout() { let (notifications, response) = @@ -989,3 +847,124 @@ async fn ask_without_form_support_explains_instead_of_asking() { spoken(¬ifications) ); } + +struct FormEngine; +impl TurnEngine for FormEngine { + fn run_turn( + &self, + request: TurnRequest, + ) -> tokio::sync::mpsc::Receiver> { + let (parts, receiver) = tokio::sync::mpsc::channel(1); + tokio::spawn(async move { + let input = request.user_input.unwrap(); + let forms = (0..2).map(|index| { + let input = input.clone(); + async move { + input.form(format!("macro: Review {index}"), serde_json::json!({ + "type":"object", "properties":{"draft":{"type":"string"}} + }), Some(serde_json::json!({"macro":{"userTool":{"name":"SendEmail","draft":{"subject":"hello"}}}}).as_object().unwrap().clone())).await.unwrap() + } + }); + let responses = tokio::select! { + _ = request.cancel.cancelled() => return, + responses = futures::future::join_all(forms) => responses, + }; + let _ = parts + .send(Ok(StreamPart::Content( + serde_json::to_string(&responses).unwrap(), + ))) + .await; + }); + receiver + } +} + +#[tokio::test(start_paused = true)] +async fn concurrent_mcp_forms_queue_and_wait_beyond_the_turn_idle_deadline() { + let (notifications, asked, response) = with_asking_engine( + Arc::new(FormEngine), + false, + ElicitationAction::Decline, + TURN_IDLE_TIMEOUT * 2, + async |connection, session| { + connection + .send_request(text_prompt(&session, "review drafts")) + .block_task() + .await + .unwrap() + }, + ) + .await; + assert_eq!(asked.len(), 2); + assert_eq!(asked[0].message, "macro: Review 0"); + assert_eq!(asked[1].message, "macro: Review 1"); + assert_eq!( + serde_json::to_value(&asked[0]).unwrap()["_meta"]["macro"]["userTool"]["name"], + "SendEmail" + ); + assert_eq!(response.stop_reason, StopReason::EndTurn); + assert_eq!( + spoken(¬ifications), + r#"[{"action":"decline"},{"action":"decline"}]"# + ); +} + +struct FailingReconnect(std::sync::atomic::AtomicUsize); +#[async_trait::async_trait] +impl crate::domain::mcp::McpToolConnector for FailingReconnect { + async fn connect( + &self, + _: Vec, + _: Option, + ) -> Result, String> { + if self.0.fetch_add(1, Ordering::SeqCst) == 0 { + Ok(None) + } else { + Err("MCP unavailable".into()) + } + } +} + +#[tokio::test] +async fn failed_session_setup_preserves_the_previous_binding_and_history() { + let engine = Arc::new(ScriptedEngine::new(vec![StreamPart::Content( + "remembered".into(), + )])); + with_agent_connector( + engine.clone(), + Arc::new(FailingReconnect(Default::default())), + async |connection, session| { + connection + .send_request(text_prompt(&session, "first")) + .block_task() + .await + .unwrap(); + assert!( + connection + .send_request(NewSessionRequest::new("/")) + .block_task() + .await + .is_err() + ); + assert!( + connection + .send_request(ResumeSessionRequest::new( + SessionId::new("wrong-session"), + "/" + )) + .block_task() + .await + .is_err() + ); + connection + .send_request(text_prompt(&session, "second")) + .block_task() + .await + .unwrap(); + }, + ) + .await; + let requests = engine.requests(); + assert_eq!(requests.len(), 2); + assert_eq!(requests[1].messages, vec!["first", "remembered", "second"]); +} diff --git a/crates/agent_inmem/src/domain/engine.rs b/crates/agent_inmem/src/domain/engine.rs index 3f1997917eb..5105cceac92 100644 --- a/crates/agent_inmem/src/domain/engine.rs +++ b/crates/agent_inmem/src/domain/engine.rs @@ -1,10 +1,7 @@ //! The seam between the ACP surface and the agentic loop that serves it. -use std::sync::Arc; - use agent::types::ChatMessage; use agent::{AgentError, StreamPart}; -use ai_tools::user_tool_review::UserToolReviewer; use macro_user_id::user_id::MacroUserIdStr; use mcp_toolset::RemoteMcpToolSet; use tokio::sync::mpsc; @@ -27,7 +24,7 @@ pub struct TurnRequest { /// answered. pub messages: Vec, /// Tools of the MCP servers the session was handed, composed next to the - /// native Macro tools. `None` when the session has none. + /// local harness utilities. `None` when the session has none. pub mcp_tools: Option, /// Cancelling this token stops the turn; the stream ends after the /// engine has drained cooperatively. @@ -35,11 +32,6 @@ pub struct TurnRequest { /// User-input capability for model-callable tools. Absent when the ACP /// client did not advertise form elicitation. pub user_input: Option, - /// Puts a user tool's call (`SendEmail`, `CreateCalendarEvent`) to the - /// user for review mid-turn, so the tool is finished - run as edited, or - /// rejected - before the model reads its result. Absent for the same - /// reason as `user_input`; a pending call then stays pending. - pub reviewer: Option>, } /// Runs one conversational turn and streams its parts back. diff --git a/crates/agent_inmem/src/domain/mcp.rs b/crates/agent_inmem/src/domain/mcp.rs index 8569349b4ef..64df526f1ac 100644 --- a/crates/agent_inmem/src/domain/mcp.rs +++ b/crates/agent_inmem/src/domain/mcp.rs @@ -14,23 +14,14 @@ use mcp_toolset::RemoteMcpToolSet; #[cfg(test)] mod test; -/// The name the harness gives Macro's own MCP server in the ACP list. -/// -/// Restated from `agent_harness::MACRO_MCP_NAME` rather than imported - this -/// crate is a runtime the harness drives, not a dependant of it - and pinned -/// equal by a test in the composition root, which sees both. Macro's tools are -/// native here (see [`ai_tools::all_tools`]), so the entry is skipped: dialing -/// it would give the model every Macro tool twice. +/// Reserved name of the Macro MCP server supplied by the harness. pub const MACRO_MCP_NAME: &str = "macro"; -/// The servers this runtime should dial out of an ACP server list: HTTP -/// entries other than Macro's own. SSE and stdio entries are not something an -/// in-process runtime can serve and are dropped with a warning. +/// HTTP servers this runtime can dial, including Macro itself. pub fn dialable_servers(servers: Vec) -> Vec { servers .into_iter() .filter_map(|server| match server { - AcpMcpServer::Http(http) if http.name == MACRO_MCP_NAME => None, AcpMcpServer::Http(http) => Some(http), other => { tracing::warn!( @@ -43,43 +34,27 @@ pub fn dialable_servers(servers: Vec) -> Vec { .collect() } -/// Opens sessions on the HTTP MCP servers a session was handed and builds the -/// toolset over them. -/// -/// A port because it is transport work: the production adapter speaks -/// streamable HTTP through the egress proxy, tests hand back nothing. +/// Connects the supplied servers. Production requires a working Macro catalog. +#[async_trait::async_trait] pub trait McpToolConnector: Send + Sync + 'static { - /// Connect to every server, skipping any that fail, and return the tools - /// found. `None` when no server yielded any tool. - fn connect( + /// Optional integrations can fail independently of the required Macro server. + async fn connect( &self, servers: Vec, - ) -> impl Future> + Send; + input: Option, + ) -> Result, String>; } -/// Erased form of [`McpToolConnector`] for storage on the agent state. -pub trait DynMcpToolConnector: Send + Sync + 'static { - /// See [`McpToolConnector::connect`]. - fn connect_dyn( - &self, - servers: Vec, - ) -> std::pin::Pin> + Send + '_>>; -} - -impl DynMcpToolConnector for C { - fn connect_dyn( - &self, - servers: Vec, - ) -> std::pin::Pin> + Send + '_>> { - Box::pin(self.connect(servers)) - } -} - -/// A connector for deployments and tests that hand the runtime no servers. +/// Connector for scripted engines that do not execute product tools. pub struct NoMcpServers; +#[async_trait::async_trait] impl McpToolConnector for NoMcpServers { - async fn connect(&self, _servers: Vec) -> Option { - None + async fn connect( + &self, + _servers: Vec, + _input: Option, + ) -> Result, String> { + Ok(None) } } diff --git a/crates/agent_inmem/src/domain/mcp/test.rs b/crates/agent_inmem/src/domain/mcp/test.rs index 20d201044e3..9b39fd4773f 100644 --- a/crates/agent_inmem/src/domain/mcp/test.rs +++ b/crates/agent_inmem/src/domain/mcp/test.rs @@ -10,10 +10,10 @@ fn http(name: &str) -> AcpMcpServer { } #[test] -fn macros_own_server_is_never_dialed() { +fn macros_own_server_is_dialed() { let dialable = dialable_servers(vec![http("macro"), http("linear"), http("notion")]); let names: Vec<&str> = dialable.iter().map(|server| server.name.as_str()).collect(); - assert_eq!(names, ["linear", "notion"]); + assert_eq!(names, ["macro", "linear", "notion"]); } #[test] diff --git a/crates/agent_inmem/src/domain/user_input.rs b/crates/agent_inmem/src/domain/user_input.rs index a0b64ece7d7..9ae53382b2b 100644 --- a/crates/agent_inmem/src/domain/user_input.rs +++ b/crates/agent_inmem/src/domain/user_input.rs @@ -55,6 +55,17 @@ impl std::error::Error for UserInputError {} /// Capability used by a turn to ask its connected user a question. #[async_trait] pub trait UserInputRequester: Send + Sync { + /// Present a flat schema supplied by a connected tool server. + /// The response contains the user's action and submitted content. + async fn form( + &self, + _message: String, + _schema: serde_json::Value, + _meta: Option>, + ) -> Result { + Err(UserInputError::Unsupported) + } + /// Ask one question and wait for the user's decision. async fn ask(&self, request: UserInputRequest) -> Result; } diff --git a/crates/agent_inmem/src/outbound/acp_mcp.rs b/crates/agent_inmem/src/outbound/acp_mcp.rs index 71cab724980..3e32e686a73 100644 --- a/crates/agent_inmem/src/outbound/acp_mcp.rs +++ b/crates/agent_inmem/src/outbound/acp_mcp.rs @@ -11,7 +11,67 @@ use rmcp::transport::streamable_http_client::{ StreamableHttpClient, StreamableHttpClientTransportConfig, }; -use crate::domain::mcp::McpToolConnector; +use crate::domain::mcp::{MACRO_MCP_NAME, McpToolConnector}; +use crate::domain::user_input::SharedUserInputRequester; + +struct ElicitationClient { + server: String, + input: Option, +} + +impl rmcp::ClientHandler for ElicitationClient { + fn get_info(&self) -> rmcp::model::ClientInfo { + let mut info = client_info(); + if self.input.is_some() { + info.capabilities.elicitation = Some(rmcp::model::ElicitationCapability { + form: Some(Default::default()), + url: None, + }); + } + info + } + + async fn create_elicitation( + &self, + params: rmcp::model::CreateElicitationRequestParams, + context: rmcp::service::RequestContext, + ) -> Result { + let input = self.input.as_ref().ok_or_else(|| { + rmcp::ErrorData::invalid_params("form elicitation is unsupported", None) + })?; + let rmcp::model::CreateElicitationRequestParams::FormElicitationParams { + message, + requested_schema, + meta, + } = params + else { + return Err(rmcp::ErrorData::invalid_params( + "only form elicitation is supported", + None, + )); + }; + let schema = serde_json::to_value(requested_schema) + .map_err(|error| rmcp::ErrorData::internal_error(error.to_string(), None))?; + // rmcp moves request metadata into RequestContext before dispatch. + let mut metadata = context.meta.0.clone(); + if let Some(meta) = meta { + metadata.extend(meta.0); + } + let mut meta = (!metadata.is_empty()).then_some(metadata); + if self.server != MACRO_MCP_NAME + && let Some(meta) = &mut meta + { + meta.remove("macro"); + } + let response = tokio::select! { + biased; + _ = context.ct.cancelled() => return Ok(rmcp::model::CreateElicitationResult::new(rmcp::model::ElicitationAction::Cancel)), + response = input.form(format!("{}: {message}", self.server), schema, meta) => response, + }.map_err(|error| rmcp::ErrorData::internal_error(error.to_string(), None))?; + serde_json::from_value(response) + .map_err(|error| rmcp::ErrorData::invalid_params(error.to_string(), None)) + } +} #[cfg(test)] mod test; @@ -65,7 +125,11 @@ where Self { client } } - async fn connect_one(&self, server: McpServerHttp) -> Option { + async fn connect_one( + &self, + server: McpServerHttp, + input: Option, + ) -> Option { let mut config = StreamableHttpClientTransportConfig::with_uri(server.url.clone()); let mut custom = HashMap::new(); for header in &server.headers { @@ -83,44 +147,81 @@ where } config.custom_headers = custom; + // Reconnect on a later request, never replay a potentially mutating call. + config.reinit_on_expired_session = false; let transport = StreamableHttpClientTransport::with_client(self.client.clone(), config); - match client_info().serve(transport).await { - Ok(client) => Some(ConnectedServer { + let handler = ElicitationClient { + server: server.name.clone(), + input, + }; + match tokio::time::timeout( + std::time::Duration::from_secs(20), + handler.into_dyn().serve(transport), + ) + .await + { + Ok(Ok(client)) => Some(ConnectedServer { name: server.name, client, }), Err(error) => { - // One server that will not answer must not cost the session - // the others, nor the session itself. - tracing::warn!(server = %server.name, error = ?error, "failed to connect to an MCP server; skipping it"); + tracing::warn!(server = %server.name, error = ?error, "timed out initializing MCP server"); + None + } + Ok(Err(error)) => { + tracing::warn!(server = %server.name, error = ?error, "failed to initialize MCP server"); None } } } } +#[async_trait::async_trait] impl McpToolConnector for AcpMcpConnector where - Client: StreamableHttpClient + Send + Sync, + Client: StreamableHttpClient + Send + Sync + 'static, { - #[tracing::instrument(skip_all, fields(servers = servers.len()))] - async fn connect(&self, servers: Vec) -> Option { - if servers.is_empty() { - return None; + async fn connect( + &self, + servers: Vec, + input: Option, + ) -> Result, String> { + if servers + .iter() + .filter(|server| server.name == MACRO_MCP_NAME) + .count() + != 1 + { + return Err("Exactly one Macro MCP server is required".to_owned()); } - let connected: Vec = - futures::future::join_all(servers.into_iter().map(|server| self.connect_one(server))) - .await + let connected: Vec<_> = futures::future::join_all( + servers .into_iter() - .flatten() - .collect(); - if connected.is_empty() { - return None; + .map(|server| self.connect_one(server, input.clone())), + ) + .await + .into_iter() + .flatten() + .collect(); + if !connected.iter().any(|server| server.name == MACRO_MCP_NAME) { + return Err("Could not connect to Macro MCP; retry the session".to_owned()); } - let tools = RemoteMcpToolSet::from_connected(connected, None).await; - if tools.is_empty() { - return None; + let tools = tokio::time::timeout( + std::time::Duration::from_secs(25), + RemoteMcpToolSet::from_connected(connected, None), + ) + .await + .map_err(|_| "MCP tool discovery timed out".to_owned())?; + let catalog = tools.searchable_catalog(); + if !["mcp__macro__SendEmail", "mcp__macro__CreateCalendarEvent"] + .iter() + .all(|name| catalog.iter().any(|tool| tool.name == *name)) + { + return Err( + "Macro MCP must expose its full reviewed tool catalog before inmem can run" + .to_owned(), + ); } - Some(tools) + Ok(Some(tools)) } } diff --git a/crates/agent_inmem/src/outbound/acp_mcp/test.rs b/crates/agent_inmem/src/outbound/acp_mcp/test.rs index c152dc3a389..a219539796f 100644 --- a/crates/agent_inmem/src/outbound/acp_mcp/test.rs +++ b/crates/agent_inmem/src/outbound/acp_mcp/test.rs @@ -38,3 +38,77 @@ fn invalid_headers_are_dropped() { assert_eq!(place_header("not a header", "x"), None); assert_eq!(place_header("X-Custom", "line\nbreak"), None); } + +#[derive(Default)] +struct FormRecorder(std::sync::Mutex>); +#[async_trait::async_trait] +impl crate::domain::user_input::UserInputRequester for FormRecorder { + async fn ask( + &self, + _: crate::domain::user_input::UserInputRequest, + ) -> Result< + crate::domain::user_input::UserInputOutcome, + crate::domain::user_input::UserInputError, + > { + unreachable!() + } + async fn form( + &self, + message: String, + schema: serde_json::Value, + meta: Option>, + ) -> Result { + self.0 + .lock() + .unwrap() + .push(serde_json::json!({"message":message,"schema":schema,"meta":meta})); + Ok(serde_json::json!({"action":"accept","content":{"answer":"edited"}})) + } +} +struct ProbeServer; +impl rmcp::ServerHandler for ProbeServer { + async fn call_tool( + &self, + _: rmcp::model::CallToolRequestParams, + context: rmcp::service::RequestContext, + ) -> Result { + let response = context.peer.create_elicitation(serde_json::from_value(serde_json::json!({ + "mode":"form", "message":"Question", "requestedSchema":{"type":"object","properties":{"answer":{"type":"string"}}}, + "_meta":{"macro":{"userTool":{"name":"SendEmail","draft":{}}},"other":"kept"} + })).unwrap()).await.unwrap(); + Ok(rmcp::model::CallToolResult::structured( + serde_json::to_value(response).unwrap(), + )) + } +} + +#[tokio::test] +async fn mcp_form_bridge_preserves_answers_and_only_trusts_macro_composer_metadata() { + use std::sync::Arc; + for server in ["macro", "thirdparty"] { + let input = Arc::new(FormRecorder::default()); + let client = ElicitationClient { + server: server.into(), + input: Some(input.clone()), + }; + let (server_io, client_io) = tokio::io::duplex(8192); + let server_task = tokio::spawn(async move { ProbeServer.serve(server_io).await.unwrap() }); + let client = client.serve(client_io).await.unwrap(); + let service = server_task.await.unwrap(); + let result = client + .call_tool(rmcp::model::CallToolRequestParams::new("probe")) + .await + .unwrap(); + assert_eq!( + result.structured_content.unwrap(), + serde_json::json!({"action":"accept","content":{"answer":"edited"}}) + ); + let forms = input.0.lock().unwrap(); + assert_eq!(forms[0]["message"], format!("{server}: Question")); + assert_eq!(forms[0]["meta"]["other"], "kept"); + assert_eq!(forms[0]["meta"].get("macro").is_some(), server == "macro"); + drop(forms); + client.cancel().await.unwrap(); + service.cancel().await.unwrap(); + } +} diff --git a/crates/agent_inmem/src/outbound/manager.rs b/crates/agent_inmem/src/outbound/manager.rs index 31cc07a8575..ffbfeb6b3bb 100644 --- a/crates/agent_inmem/src/outbound/manager.rs +++ b/crates/agent_inmem/src/outbound/manager.rs @@ -18,7 +18,7 @@ use macro_user_id::user_id::MacroUserIdStr; use crate::domain::agent::{AgentState, serve}; use crate::domain::engine::TurnEngine; -use crate::domain::mcp::DynMcpToolConnector; +use crate::domain::mcp::McpToolConnector; use crate::domain::replay::{FrameSource, replay_history}; use crate::domain::session::{SessionState, SessionStore}; @@ -64,7 +64,7 @@ impl Drop for LiveAgent { pub struct InMemAgentManager { engine: Arc, frames: Arc, - mcp: Arc, + mcp: Arc, enable_dev_commands: bool, store: Arc, live: DashMap, @@ -84,7 +84,7 @@ impl InMemAgentManager { pub fn new( engine: Arc, frames: Arc, - mcp: Arc, + mcp: Arc, ) -> Self { Self { engine, @@ -155,6 +155,8 @@ impl InMemAgentManager { turn_lock: tokio::sync::Mutex::new(()), mcp: Arc::clone(&self.mcp), mcp_tools: std::sync::Mutex::new(None), + awaiting_user: Default::default(), + input_gate: Default::default(), client_renders_forms: AtomicBool::new(false), enable_dev_commands: self.enable_dev_commands, }); diff --git a/crates/agent_inmem/src/outbound/rig_engine/test.rs b/crates/agent_inmem/src/outbound/rig_engine/test.rs index 867d8aef49e..cd283b5f183 100644 --- a/crates/agent_inmem/src/outbound/rig_engine/test.rs +++ b/crates/agent_inmem/src/outbound/rig_engine/test.rs @@ -4,9 +4,7 @@ use super::*; const TOOLS: &str = "TOOLS"; fn has_ask_user(supports_user_input: bool) -> bool { - let tools = tools_for(AiHost::AgentSession); - let base_tools = Arc::into_inner(tools.toolset) - .expect("tools_for should return a fresh, uniquely owned collection"); + let base_tools = ai_tools::harness_tools(); tools_for_turn(base_tools, supports_user_input) .request_schemas() .expect("tool schemas should be valid") @@ -71,3 +69,61 @@ fn empty_instructions_add_no_section() { assert!(!prompt.contains("session_instructions")); } + +#[test] +fn main_agent_registers_only_harness_utilities() { + let tools = tools_for_turn(ai_tools::harness_tools(), true); + let mut names: Vec<_> = tools + .request_schemas() + .unwrap() + .into_iter() + .map(|s| s.name) + .collect(); + names.sort(); + assert_eq!( + names, + ["AskUser", "DisplayResults", "LoadTools", "SearchTools"] + ); +} + +#[test] +fn historical_native_calls_and_results_use_current_mcp_names_only_in_model_context() { + use agent::types::{AssistantMessagePart as Part, ChatMessage, ChatMessageContent, Role}; + let original = vec![ + Part::ToolCall { + name: "SendEmail".into(), + id: "call".into(), + json: serde_json::json!({}), + }, + Part::ToolCallResponseJson { + name: "SendEmail".into(), + id: "call".into(), + json: serde_json::json!("Rejected"), + }, + Part::ToolCall { + name: "AskUser".into(), + id: "ask".into(), + json: serde_json::json!({}), + }, + ]; + let mut messages = vec![ChatMessage { + content: ChatMessageContent::AssistantMessageParts(original.clone()), + role: Role::Assistant, + attachments: None, + }]; + normalize_tool_history( + &mut messages, + &std::collections::HashSet::from(["mcp__macro__SendEmail".into()]), + ); + let ChatMessageContent::AssistantMessageParts(parts) = &messages[0].content else { + panic!("parts"); + }; + assert!( + matches!(&parts[0], Part::ToolCall {name, id, ..} if name == "mcp__macro__SendEmail" && id == "call") + ); + assert!( + matches!(&parts[1], Part::ToolCallResponseJson {name, ..} if name == "mcp__macro__SendEmail") + ); + assert_eq!(parts[2], original[2]); + assert!(matches!(&original[0], Part::ToolCall {name, ..} if name == "SendEmail")); +} diff --git a/crates/agent_inmem/src/rig_engine.rs b/crates/agent_inmem/src/rig_engine.rs index 917b81b53bf..67f11d178c0 100644 --- a/crates/agent_inmem/src/rig_engine.rs +++ b/crates/agent_inmem/src/rig_engine.rs @@ -7,13 +7,7 @@ //! prompt (immediately before any session instructions), and the owner's //! memory, with usage recorded per turn against the session owner. //! -//! User tools (`SendEmail`, `CreateCalendarEvent`) are the chat host's -//! deferring ones, finished inside the turn: the turn's [`TurnRequest`] -//! carries a reviewer over the ACP connection, and the agent loop's -//! user-tool finisher puts each pending call to it - the session renders the -//! elicitation - then runs or rejects the tool before the model reads the -//! result. Without a reviewer (a client with no form support) a pending call -//! stays pending, as in chat. +//! Product tools execute through MCP; this engine only registers harness utilities. //! //! This is also where a session's own instructions become a system prompt. //! Nothing has to be transported for it - the loop runs in this process - which @@ -23,7 +17,6 @@ use std::sync::Arc; use agent::{AgentError, AgentLoop, StreamPart}; -use ai_tools::user_tool_review::user_tool_finisher; use ai_tools::{AiHost, ToolServiceContext, ToolSetWithPrompt, tools_for}; use ai_toolset::{AsyncToolCollection, ToolSet as AiToolSet}; use axum::extract::FromRef; @@ -116,30 +109,28 @@ async fn drive_turn( owner, model, instructions, - messages, + mut messages, mcp_tools, cancel, user_input, - reviewer, } = request; - // Chat's tools with the session's prompt: the user tools (`SendEmail`, - // `CreateCalendarEvent`) defer to the user, and this runtime finishes - // them in the turn through `reviewer`. - let tools = tools_for(AiHost::AgentSession); + let mcp_tools = mcp_tools.ok_or_else(|| { + AgentError::Other(anyhow::anyhow!( + "Macro MCP is unavailable; reconnect the session" + )) + })?; let user_memory = fetch_user_memory(&db, &base_context, &owner).await; let system_prompt = system_prompt( - &tools.prompt, + &prompt::SESSION_TOOL_USE_PROMPT, instructions.as_deref(), user_memory.as_deref(), ); - // `tools_for` returns a fresh Arc. Take its collection back so the - // in-memory runtime can widen it onto the session-specific context and - // add the one tool that needs the active ACP connection. - let base_tools = Arc::into_inner(tools.toolset) - .expect("tools_for should return a fresh, uniquely owned collection"); - let toolset = Arc::new(tools_for_turn(base_tools, user_input.is_some())); + let toolset = Arc::new(tools_for_turn( + ai_tools::harness_tools(), + user_input.is_some(), + )); let usage_ctx = ai_usage::UsageContext::new(ai_usage::AiFeature::AgentSession, owner.clone()); // Carry the feature on the context so tool-spawned subagents attribute to it. let mut tool_context = base_context.clone(); @@ -151,22 +142,17 @@ async fn drive_turn( }, }; - let mut agent_loop = AgentLoop::new(base_context.recorder.clone()).with_model(&model); - if let Some(reviewer) = reviewer { - agent_loop = agent_loop.with_user_tool_finisher(user_tool_finisher( - Arc::clone(&toolset), - tool_context.clone(), - owner, - reviewer, - cancel.clone(), - )); - } - // Keep remote MCP tools alongside the native and AskUser tools. The - // finisher above reviews only Macro's native user tools. - let toolset: Arc + Send + Sync> = match mcp_tools { - Some(mcp) => Arc::new(mcp_select::CombinedToolSet::new(toolset, mcp)), - None => toolset, - }; + normalize_tool_history( + &mut messages, + &mcp_tools + .searchable_catalog() + .into_iter() + .map(|tool| tool.name) + .collect(), + ); + let agent_loop = AgentLoop::new(base_context.recorder.clone()).with_model(&model); + let toolset: Arc + Send + Sync> = + Arc::new(mcp_select::CombinedToolSet::new(toolset, mcp_tools)); let session = agent_loop .session(toolset, Arc::new(tool_context), &system_prompt, usage_ctx) .await; @@ -253,3 +239,28 @@ async fn fetch_user_memory( } } } + +/// Adapt model context from before the MCP cutover without rewriting stored logs. +fn normalize_tool_history( + messages: &mut [agent::types::ChatMessage], + catalog: &std::collections::HashSet, +) { + use agent::types::{AssistantMessagePart, ChatMessageContent}; + for message in messages { + let ChatMessageContent::AssistantMessageParts(parts) = &mut message.content else { + continue; + }; + for part in parts { + let name = match part { + AssistantMessagePart::ToolCall { name, .. } + | AssistantMessagePart::ToolCallResponseJson { name, .. } + | AssistantMessagePart::ToolCallErr { name, .. } => name, + _ => continue, + }; + let qualified = format!("mcp__macro__{name}"); + if catalog.contains(&qualified) { + *name = qualified; + } + } + } +} diff --git a/crates/ai_tools/src/lib.rs b/crates/ai_tools/src/lib.rs index f82cb211600..d4aaa624543 100644 --- a/crates/ai_tools/src/lib.rs +++ b/crates/ai_tools/src/lib.rs @@ -16,7 +16,6 @@ mod self_knowledge; pub mod serde_utils; mod subagent; mod tool_context; -pub mod user_tool_review; pub use anthropic::toolset::AnthropicToolContext; use anthropic::toolset::anthropic_toolset; @@ -109,36 +108,25 @@ pub(crate) fn subagent_toolset() -> AiToolSet { .add_subtoolset::(anthropic_toolset()) } -/// The host a toolset is assembled for. -/// -/// Hosts differ on two axes. First, whether something finishes a deferred -/// user tool for them: [`AiHost::Chat`] has the composer card after the -/// turn and [`AiHost::AgentSession`] the review elicitation in it, so those -/// two register `SendEmail` and the deferring `CreateCalendarEvent` — on any -/// other host those registrations would return `PendingUserExecution` -/// forever while reading to the model as success. Second, whether the host -/// runs the chat frontend's tool discovery and display tools. +/// The host assembling Macro's canonical product tools. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum AiHost { - /// The AI chat, and any host whose conversation is stored as a chat the - /// frontend can render (scheduled agents, memory generation): composer - /// cards finish deferred user tools there, after the turn. + /// Chat defers reviewed tools to its existing composer. Chat, - /// Macro's in-process agent serving an agent session: the same toolset - /// as [`AiHost::Chat`], but its user tools are reviewed in the turn - the - /// agent loop's finisher puts the call to the user over ACP and the tool - /// returns the outcome - so the prompt describes a review card the agent - /// waits on, never a pending composer. - AgentSession, - /// The channel-mention bot: no composer, so `CreateCalendarEvent` - /// executes directly in the agent loop and `SendEmail` is omitted. + /// The legacy channel bot keeps direct calendar creation and no email sending. ChannelBot, - /// The MCP server: like [`AiHost::ChannelBot`] for user tools — MCP - /// clients apply their own confirmation policy from tool annotations — - /// and without the chat frontend's discovery/display tools. + /// MCP reviews user tools through elicitation before execution. Mcp, } +/// Harness-local discovery and presentation tools, without product actions. +pub fn harness_tools() -> AiToolSet { + AsyncToolCollection::new() + .add_tool::() + .add_tool::() + .add_tool::() +} + /// Assemble the toolset and tool-use prompt for a host. These are actually /// sent to the AI provider. pub fn tools_for(host: AiHost) -> ToolSetWithPrompt { @@ -146,10 +134,10 @@ pub fn tools_for(host: AiHost) -> ToolSetWithPrompt { .add_subtoolset::(notification_toolset()) .add_subtoolset::(reminders_toolset()); let toolset = match host { - AiHost::Chat | AiHost::AgentSession => toolset + AiHost::Chat | AiHost::Mcp => toolset .add_subtoolset::(email_toolset()) .add_subtoolset::(calendar_toolset()), - AiHost::ChannelBot | AiHost::Mcp => toolset + AiHost::ChannelBot => toolset .add_subtoolset::(email_mcp_toolset()) .add_subtoolset::(calendar_mcp_toolset()), }; @@ -157,16 +145,13 @@ pub fn tools_for(host: AiHost) -> ToolSetWithPrompt { .add_subtoolset::(import_toolset()) .add_tool::(); let toolset = match host { - AiHost::Chat | AiHost::AgentSession | AiHost::ChannelBot => toolset - .add_tool::() - .add_tool::() - .add_tool::(), + AiHost::Chat | AiHost::ChannelBot => toolset.add_toolset(harness_tools()), AiHost::Mcp => toolset, }; let prompt: Box = match host { AiHost::Chat => Box::new(&prompt::TOOL_USE_PROMPT), - AiHost::AgentSession => Box::new(&prompt::SESSION_TOOL_USE_PROMPT), - AiHost::ChannelBot | AiHost::Mcp => Box::new(&prompt::DIRECT_TOOL_USE_PROMPT), + AiHost::Mcp => Box::new(&prompt::SESSION_TOOL_USE_PROMPT), + AiHost::ChannelBot => Box::new(&prompt::DIRECT_TOOL_USE_PROMPT), }; ToolSetWithPrompt { toolset: Arc::new(toolset), diff --git a/crates/ai_tools/src/test.rs b/crates/ai_tools/src/test.rs index 5b7e062c0ab..825844e60b4 100644 --- a/crates/ai_tools/src/test.rs +++ b/crates/ai_tools/src/test.rs @@ -22,22 +22,15 @@ fn subagent_toolset_passes_schema_validation() { #[test] fn every_host_toolset_passes_schema_validation() { - for host in [ - AiHost::Chat, - AiHost::AgentSession, - AiHost::ChannelBot, - AiHost::Mcp, - ] { + for host in [AiHost::Chat, AiHost::ChannelBot, AiHost::Mcp] { let _ = tools_for(host); } } -/// An agent session finishes user tools in the turn, so it keeps chat's -/// deferring registrations - and gets the prompt that says a review card, -/// not a pending composer, is what follows the call. +/// MCP and chat share the product catalog; MCP settles reviews before returning. #[test] -fn the_agent_session_host_keeps_chats_user_tools_with_the_review_prompt() { - let session = tools_for(AiHost::AgentSession); +fn mcp_keeps_chats_reviewed_product_tools() { + let session = tools_for(AiHost::Mcp); assert!( session .toolset @@ -48,17 +41,12 @@ fn the_agent_session_host_keeps_chats_user_tools_with_the_review_prompt() { let prompt = session.prompt.to_string(); assert!(prompt.contains("review card")); assert!(!prompt.contains("PendingUserExecution")); - assert_eq!( - session - .toolset - .request_schemas() - .map(|schemas| schemas.len()), - tools_for(AiHost::Chat) - .toolset - .request_schemas() - .map(|schemas| schemas.len()), - "the same tools as chat" - ); + let chat = tools_for(AiHost::Chat); + for name in chat.toolset.tools.keys() { + if !["SearchTools", "LoadTools", "DisplayResults"].contains(&name.as_str()) { + assert!(session.toolset.tools.contains_key(name), "missing {name}"); + } + } } /// Hosts without a composer cannot finish a deferred user tool, so their @@ -67,7 +55,7 @@ fn the_agent_session_host_keeps_chats_user_tools_with_the_review_prompt() { /// nothing can ever execute. #[test] fn composerless_hosts_execute_calendar_create_directly_and_omit_send_email() { - for host in [AiHost::ChannelBot, AiHost::Mcp] { + for host in [AiHost::ChannelBot] { let json = frontend_schemas_builder() .merge(&tools_for(host)) .build() diff --git a/crates/ai_tools/src/user_tool_review.rs b/crates/ai_tools/src/user_tool_review.rs deleted file mode 100644 index 58d2d0c2488..00000000000 --- a/crates/ai_tools/src/user_tool_review.rs +++ /dev/null @@ -1,424 +0,0 @@ -//! Finishing a user tool inside the turn: the user reviews the call the model -//! made, then the host executes or rejects it. -//! -//! A user tool (`ai_toolset::UserTool`, registered with `add_user_tool`) -//! answers `"PendingUserExecution"` and does nothing. The host finishes it -//! with the pieces `ai_toolset` exposes for exactly that: `is_valid_tool` to -//! check edited arguments and `try_user_tool_call` to run the wrapped tool. -//! Chat does so after the turn, over HTTP, from its composer. A host that can -//! reach its user *during* the turn - Macro's in-process agent, whose ACP -//! client renders elicitations - does it here, through the -//! [`UserToolReviewer`] port, before the model reads the result. -//! -//! The port speaks in the restricted form vocabulary elicitation shares -//! across ACP and MCP (flat primitives with defaults), stated in this crate's -//! own types so neither protocol leaks in. A [`ReviewForm`] is projected from -//! the tool's input schema and the call's arguments, so every user tool is -//! reviewable without per-tool code; the whole edited draft can come back in -//! one `Json` field for a client that has a composer of its own. - -use std::collections::BTreeMap; -use std::future::Future; -use std::pin::Pin; -use std::sync::Arc; - -use agent::{FinishedUserTool, PendingUserTool, UserToolFinisher}; -use ai_toolset::tool_object::UserToolResponse; -use ai_toolset::{AsyncToolCollection, RequestContext}; -use async_trait::async_trait; -use macro_user_id::user_id::MacroUserIdStr; -use serde_json::{Map, Value}; -use tokio_util::sync::CancellationToken; - -#[cfg(test)] -mod test; - -/// The name of the form field that carries the whole edited draft as JSON. -/// A client with its own composer for the tool fills this in; a client that -/// renders the form generically leaves it out and edits the flat fields. -pub const DRAFT_FIELD: &str = "draft"; - -/// A user tool's call, put to the user for review. -#[derive(Debug, Clone, PartialEq)] -pub struct ReviewRequest { - /// The tool's name as the toolset knows it. - pub tool_name: String, - /// The call's id as the transcript shows it. - pub tool_call_id: String, - /// What the user is asked, in one line: "Create calendar event?". - pub message: String, - /// The arguments the model called the tool with, whole. - pub draft: Value, - /// The flat form a client renders to edit the draft. - pub form: ReviewForm, -} - -/// A flat form: the restricted schema elicitation allows, in neutral terms. -#[derive(Debug, Clone, PartialEq, Default)] -pub struct ReviewForm { - /// The tool's human title, for the form's heading. - pub title: Option, - /// The fields, in the schema's order. - pub fields: Vec, - /// Names of the fields an answer must fill. - pub required: Vec, -} - -/// One form field. -#[derive(Debug, Clone, PartialEq)] -pub struct ReviewField { - /// The key the answer is sent back under; the argument it edits. - pub name: String, - /// The argument's description from the tool's schema, when it has one. - pub description: Option, - /// The field's type and pre-filled value. - pub kind: ReviewFieldKind, -} - -/// A field's type, with the value the draft holds as its default. -#[derive(Debug, Clone, PartialEq)] -pub enum ReviewFieldKind { - /// Free text; `format` is the schema's (`date-time`, `email`, ...) when given. - Text { - /// The draft's value. - default: Option, - /// The schema's format hint. - format: Option, - }, - /// A yes/no. - Boolean { - /// The draft's value. - default: Option, - }, - /// A number. - Number { - /// The draft's value. - default: Option, - }, - /// A whole number. - Integer { - /// The draft's value. - default: Option, - }, - /// One of a fixed set of strings. - Choice { - /// The allowed values. - options: Vec, - /// The draft's value. - default: Option, - }, - /// The whole edited draft as a JSON string - see [`DRAFT_FIELD`]. - Json, -} - -/// What the user decided. -#[derive(Debug, Clone, PartialEq)] -pub enum ReviewOutcome { - /// The user confirmed, with the form's values as submitted. Absent - /// fields keep the draft's values. - Accepted(BTreeMap), - /// The user said no. - Declined, - /// The question was dismissed, or the turn stopped. - Cancelled, -} - -/// Why a review could not happen. -#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] -pub enum ReviewError { - /// The user cannot be asked right now: the client has no way to show the - /// form, or is already holding another question. - #[error("the user cannot be asked to review this call right now: {0}")] - Unavailable(String), - /// Asking failed on the way. - #[error("asking the user to review this call failed: {0}")] - Failed(String), -} - -/// A host's way of putting a call to its user and waiting for the answer. -#[async_trait] -pub trait UserToolReviewer: Send + Sync { - /// Ask the user to review `request` and wait for their decision. - async fn review(&self, request: ReviewRequest) -> Result; -} - -/// A [`UserToolFinisher`] for `tools`: puts each pending user tool to -/// `reviewer`, and on acceptance runs the tool with the reviewed arguments -/// as `user`, against `context`. -/// -/// The tool's answer is a `UserToolResponse` as JSON, the same shape chat -/// writes when its composer finishes the call: the user's action with the -/// tool's result, or `"Rejected"`. A tool this toolset does not know as a -/// user tool is left alone (`None`) for the host to finish some other way. -pub fn user_tool_finisher( - tools: Arc>, - context: Context, - user: MacroUserIdStr<'static>, - reviewer: Arc, - cancel: CancellationToken, -) -> UserToolFinisher -where - Context: Clone + Send + Sync + 'static, -{ - Arc::new(move |call: PendingUserTool| { - let tools = Arc::clone(&tools); - let context = context.clone(); - let user = user.clone(); - let reviewer = Arc::clone(&reviewer); - let cancel = cancel.clone(); - Box::pin(async move { finish(&tools, context, user, &*reviewer, cancel, call).await }) - as Pin> + Send>> - }) -} - -async fn finish( - tools: &AsyncToolCollection, - context: Context, - user: MacroUserIdStr<'static>, - reviewer: &dyn UserToolReviewer, - cancel: CancellationToken, - call: PendingUserTool, -) -> Option -where - Context: Clone + Send + Sync + 'static, -{ - let tool = tools.user_tools.get(&call.tool_name)?; - let request = ReviewRequest { - tool_name: call.tool_name.clone(), - tool_call_id: call.tool_call_id.clone(), - message: format!("{}?", tool.annotations.title), - draft: call.args.clone(), - form: project_form( - Some(tool.annotations.title.to_owned()), - &tool.input_schema, - &call.args, - ), - }; - - let outcome = match reviewer.review(request).await { - Ok(outcome) => outcome, - Err(error) => return Some(FinishedUserTool::Error(error.to_string())), - }; - let content = match outcome { - ReviewOutcome::Accepted(content) => content, - ReviewOutcome::Declined => { - return Some(FinishedUserTool::Result(rejected())); - } - ReviewOutcome::Cancelled => { - return Some(FinishedUserTool::Error( - "the user cancelled the review; the call was not made".to_owned(), - )); - } - }; - - let args = apply_review(&call.args, &content); - if !tools.is_valid_tool(&call.tool_name, &args) { - return Some(FinishedUserTool::Error(format!( - "the reviewed arguments are not valid for {}; nothing was done", - call.tool_name - ))); - } - let mut request_context = RequestContext::new(user); - request_context.cancel = cancel; - match tools - .try_user_tool_call(context, request_context, &call.tool_name, &args) - .await - { - Ok(Ok(response)) => Some(match serde_json::to_value(response) { - Ok(json) => FinishedUserTool::Result(json), - Err(error) => { - FinishedUserTool::Error(format!("the result could not be encoded: {error}")) - } - }), - Ok(Err(error)) => Some(FinishedUserTool::Error(error.description)), - Err(error) => Some(FinishedUserTool::Error(error.to_string())), - } -} - -/// `UserToolResponse::Rejected` as the JSON chat writes for it. -fn rejected() -> Value { - serde_json::to_value(UserToolResponse::::Rejected).unwrap_or(Value::Null) -} - -/// The arguments to run with: the draft, overwritten by what the user -/// submitted. A whole draft under [`DRAFT_FIELD`] (a JSON string, or the -/// object itself) replaces everything; otherwise each flat field replaces -/// the argument of the same name. -pub fn apply_review(draft: &Value, content: &BTreeMap) -> Value { - if let Some(whole) = content.get(DRAFT_FIELD) { - let parsed = match whole { - Value::String(text) => serde_json::from_str::(text).ok(), - Value::Object(_) => Some(whole.clone()), - _ => None, - }; - if let Some(object @ Value::Object(_)) = parsed { - return object; - } - } - let mut args = match draft { - Value::Object(map) => map.clone(), - _ => Map::new(), - }; - for (name, value) in content { - if name != DRAFT_FIELD { - args.insert(name.clone(), value.clone()); - } - } - Value::Object(args) -} - -/// The flat form for a tool: one field per top-level primitive argument in -/// `schema`, pre-filled from `args`, plus the [`DRAFT_FIELD`]. Arguments the -/// restricted form cannot show - objects, arrays, unions - are left to the -/// draft field. -pub fn project_form( - title: Option, - schema: &Map, - args: &Value, -) -> ReviewForm { - let properties = schema - .get("properties") - .and_then(Value::as_object) - .cloned() - .unwrap_or_default(); - let schema_required: Vec<&str> = schema - .get("required") - .and_then(Value::as_array) - .map(|names| names.iter().filter_map(Value::as_str).collect()) - .unwrap_or_default(); - - let mut fields = Vec::new(); - let mut required = Vec::new(); - for (name, property) in &properties { - let Some(kind) = field_kind(property, args.get(name)) else { - continue; - }; - if schema_required.contains(&name.as_str()) { - required.push(name.clone()); - } - fields.push(ReviewField { - name: name.clone(), - description: property - .get("description") - .and_then(Value::as_str) - .map(str::to_owned), - kind, - }); - } - fields.push(ReviewField { - name: DRAFT_FIELD.to_owned(), - description: Some( - "The complete edited arguments as JSON, for a client with its own editor. \ - When present it replaces every other field." - .to_owned(), - ), - kind: ReviewFieldKind::Json, - }); - ReviewForm { - title, - fields, - required, - } -} - -/// The field a property renders as, or `None` when the restricted form has -/// no shape for it. -fn field_kind(property: &Value, current: Option<&Value>) -> Option { - // `Option` fields come out as `{"type": ["string", "null"]}` or - // `anyOf: [T, null]`; read through to the one non-null type. - let property = non_null(property); - if let Some(options) = choice_options(property) { - return Some(ReviewFieldKind::Choice { - options, - default: current.and_then(Value::as_str).map(str::to_owned), - }); - } - match json_type(property)? { - "string" => Some(ReviewFieldKind::Text { - default: current.and_then(Value::as_str).map(str::to_owned), - format: property - .get("format") - .and_then(Value::as_str) - .map(str::to_owned), - }), - "boolean" => Some(ReviewFieldKind::Boolean { - default: current.and_then(Value::as_bool), - }), - "integer" => Some(ReviewFieldKind::Integer { - default: current.and_then(Value::as_i64), - }), - "number" => Some(ReviewFieldKind::Number { - default: current.and_then(Value::as_f64), - }), - _ => None, - } -} - -/// The string values a property allows, when it is a choice among fixed -/// strings: a plain `enum`, or - how a documented Rust enum comes out - an -/// `anyOf`/`oneOf` whose every real variant is one `const` string (or a -/// one-value `enum`). -fn choice_options(property: &Value) -> Option> { - if let Some(options) = property.get("enum").and_then(Value::as_array) { - let options: Vec = options - .iter() - .filter_map(Value::as_str) - .map(str::to_owned) - .collect(); - return (!options.is_empty()).then_some(options); - } - let variants = ["anyOf", "oneOf"] - .into_iter() - .find_map(|key| property.get(key).and_then(Value::as_array))?; - let mut options = Vec::new(); - for variant in variants { - if json_type(variant) == Some("null") { - continue; - } - let value = variant.get("const").and_then(Value::as_str).or_else(|| { - match variant.get("enum").and_then(Value::as_array)?.as_slice() { - [only] => only.as_str(), - _ => None, - } - })?; - options.push(value.to_owned()); - } - (!options.is_empty()).then_some(options) -} - -/// The property with a nullable wrapper removed: `anyOf`/`oneOf` of one real -/// schema and `null` reads as that schema. -fn non_null(property: &Value) -> &Value { - for key in ["anyOf", "oneOf"] { - if let Some(variants) = property.get(key).and_then(Value::as_array) { - let real: Vec<&Value> = variants - .iter() - .filter(|variant| json_type(variant) != Some("null")) - .collect(); - if let [only] = real[..] { - return only; - } - } - } - property -} - -/// The one JSON type a property declares, reading `["string", "null"]` as -/// `string`. `None` for a property with several real types or none. -fn json_type(property: &Value) -> Option<&str> { - match property.get("type")? { - Value::String(name) => Some(name.as_str()), - Value::Array(names) => { - let real: Vec<&str> = names - .iter() - .filter_map(Value::as_str) - .filter(|name| *name != "null") - .collect(); - match real[..] { - [only] => Some(only), - _ => None, - } - } - _ => None, - } -} diff --git a/crates/ai_tools/src/user_tool_review/test.rs b/crates/ai_tools/src/user_tool_review/test.rs deleted file mode 100644 index 01a40fa4f5f..00000000000 --- a/crates/ai_tools/src/user_tool_review/test.rs +++ /dev/null @@ -1,338 +0,0 @@ -//! The review form is projected from a tool's own schema, the user's answer -//! is applied over the draft, and the finisher runs the wrapped tool with -//! what came back - or reports what stopped it. - -use std::sync::Mutex; - -use ai_toolset::{ - AsyncTool, RequestContext, ServiceContext, ToolAnnotated, ToolAnnotations, ToolCallError, - ToolResult, -}; -use schemars::JsonSchema; -use serde::{Deserialize, Serialize}; -use serde_json::json; - -use super::*; -use crate::{AiHost, tools_for}; - -fn owner() -> MacroUserIdStr<'static> { - MacroUserIdStr::try_from_email("owner@macro.com").expect("a valid user id") -} - -fn field<'form>(form: &'form ReviewForm, name: &str) -> &'form ReviewField { - form.fields - .iter() - .find(|field| field.name == name) - .unwrap_or_else(|| panic!("the form has a {name} field: {form:#?}")) -} - -/// The real `CreateCalendarEvent` schema, from the chat toolset that -/// registers it as a user tool. -fn create_calendar_event_schema() -> Map { - let tools = tools_for(AiHost::Chat); - tools - .toolset - .user_tools - .get("CreateCalendarEvent") - .expect("chat registers CreateCalendarEvent as a user tool") - .input_schema - .clone() -} - -#[test] -fn the_form_shows_a_tools_flat_arguments_prefilled_and_leaves_the_rest_to_the_draft() { - let schema = create_calendar_event_schema(); - let draft = json!({ - "title": "Q3 sync", - "time": {"kind": "timed", "startsAt": "2026-08-20T17:00:00Z", "endsAt": "2026-08-20T17:30:00Z"}, - "location": "Room 4", - "attendees": [{"email": "alice@example.com"}], - "addGoogleMeet": true, - }); - let form = project_form(Some("Create calendar event".to_owned()), &schema, &draft); - - assert_eq!(form.title.as_deref(), Some("Create calendar event")); - assert_eq!( - field(&form, "title").kind, - ReviewFieldKind::Text { - default: Some("Q3 sync".to_owned()), - format: None, - } - ); - assert_eq!( - field(&form, "location").kind, - ReviewFieldKind::Text { - default: Some("Room 4".to_owned()), - format: None, - }, - "an optional string reads through its nullable wrapper" - ); - assert_eq!( - field(&form, "addGoogleMeet").kind, - ReviewFieldKind::Boolean { - default: Some(true) - } - ); - assert!( - field(&form, "description").description.is_some(), - "field help comes from the schema" - ); - assert_eq!( - field(&form, "eventType").kind, - ReviewFieldKind::Choice { - options: vec!["default".to_owned(), "out_of_office".to_owned()], - default: None, - }, - "an enum is a choice, unset when the draft omits it" - ); - // Objects and arrays have no flat shape: they ride in the draft field. - for nested in [ - "time", - "attendees", - "reminders", - "outOfOffice", - "recurrenceLines", - ] { - assert!( - form.fields.iter().all(|field| field.name != nested), - "{nested} is not a flat field" - ); - } - assert_eq!(field(&form, DRAFT_FIELD).kind, ReviewFieldKind::Json); - assert_eq!( - form.required, - vec!["title".to_owned()], - "the schema's requirements, flat ones only" - ); -} - -#[test] -fn the_answer_is_applied_over_the_draft_and_a_whole_draft_wins() { - let draft = json!({"title": "Q3 sync", "location": "Room 4", "time": {"kind": "allDay"}}); - - // Flat fields replace the arguments of the same name; the rest stays. - let flat = BTreeMap::from([ - ("title".to_owned(), json!("Q3 planning")), - ("addGoogleMeet".to_owned(), json!(true)), - ]); - assert_eq!( - apply_review(&draft, &flat), - json!({"title": "Q3 planning", "location": "Room 4", "time": {"kind": "allDay"}, "addGoogleMeet": true}) - ); - - // A whole draft, as the JSON string a composer sends, replaces everything. - let whole = BTreeMap::from([ - ("title".to_owned(), json!("ignored")), - ( - DRAFT_FIELD.to_owned(), - json!(r#"{"title":"From the composer","time":{"kind":"allDay"}}"#), - ), - ]); - assert_eq!( - apply_review(&draft, &whole), - json!({"title": "From the composer", "time": {"kind": "allDay"}}) - ); - - // A draft field that is not an object is ignored, not applied. - let broken = BTreeMap::from([ - ("title".to_owned(), json!("Kept")), - (DRAFT_FIELD.to_owned(), json!("not json")), - ]); - assert_eq!(apply_review(&draft, &broken)["title"], "Kept"); - assert!(apply_review(&draft, &broken).get(DRAFT_FIELD).is_none()); - - // Nothing submitted: the draft as it was. - assert_eq!(apply_review(&draft, &BTreeMap::new()), draft); -} - -// --- the finisher, over a toolset with one user tool --- - -#[derive(Debug, Clone, PartialEq, Deserialize, JsonSchema)] -#[schemars(title = "Greet", description = "Greets someone.")] -struct Greet { - /// Who to greet. - name: String, - /// Whether to shout. - #[serde(default)] - loud: bool, -} - -#[derive(Debug, Serialize, JsonSchema)] -struct Greeting { - text: String, -} - -impl ToolAnnotated for Greet { - const ANNOTATIONS: ToolAnnotations = ToolAnnotations::read_only("Greet someone"); -} - -/// Records every call it runs. -#[derive(Clone, Default)] -struct Ran(Arc>>); - -#[async_trait] -impl AsyncTool for Greet { - type Output = Greeting; - - async fn call( - &self, - context: ServiceContext, - _request: RequestContext, - ) -> ToolResult { - if self.name == "nobody" { - return Err(ToolCallError { - description: "nobody is not a person".to_owned(), - internal_error: anyhow::anyhow!("nobody"), - }); - } - context.0.0.lock().unwrap().push(self.clone()); - let text = format!("Hello, {}{}", self.name, if self.loud { "!" } else { "." }); - Ok(Greeting { text }) - } -} - -/// A reviewer that records what it was asked and answers as scripted. -struct Scripted { - asked: Mutex>, - answer: Result, -} - -#[async_trait] -impl UserToolReviewer for Scripted { - async fn review(&self, request: ReviewRequest) -> Result { - self.asked.lock().unwrap().push(request); - self.answer.clone() - } -} - -fn finisher_over_greet( - answer: Result, -) -> (UserToolFinisher, Arc, Ran) { - let tools = Arc::new(AsyncToolCollection::::new().add_user_tool::()); - let ran = Ran::default(); - let reviewer = Arc::new(Scripted { - asked: Mutex::new(Vec::new()), - answer, - }); - let finisher = user_tool_finisher( - tools, - ran.clone(), - owner(), - Arc::clone(&reviewer) as Arc, - CancellationToken::new(), - ); - (finisher, reviewer, ran) -} - -fn pending(args: Value) -> PendingUserTool { - PendingUserTool { - tool_name: "Greet".to_owned(), - tool_call_id: "toolu_1".to_owned(), - args, - } -} - -#[tokio::test] -async fn an_accepted_review_runs_the_tool_with_the_edited_arguments() { - let (finisher, reviewer, ran) = finisher_over_greet(Ok(ReviewOutcome::Accepted( - BTreeMap::from([("loud".to_owned(), json!(true))]), - ))); - - let finished = finisher(pending(json!({"name": "Alice", "loud": false}))).await; - - let asked = reviewer.asked.lock().unwrap(); - assert_eq!(asked.len(), 1); - assert_eq!(asked[0].tool_name, "Greet"); - assert_eq!(asked[0].tool_call_id, "toolu_1"); - assert_eq!(asked[0].message, "Greet someone?"); - assert_eq!(asked[0].draft, json!({"name": "Alice", "loud": false})); - assert_eq!( - field(&asked[0].form, "name").kind, - ReviewFieldKind::Text { - default: Some("Alice".to_owned()), - format: None, - } - ); - assert_eq!(asked[0].form.required, vec!["name".to_owned()]); - - assert_eq!( - &*ran.0.lock().unwrap(), - &[Greet { - name: "Alice".to_owned(), - loud: true, - }] - ); - assert_eq!( - finished, - Some(FinishedUserTool::Result( - json!({"UserAction": {"text": "Hello, Alice!"}}) - )), - "the answer is the user tool response chat writes" - ); -} - -#[tokio::test] -async fn a_declined_review_rejects_without_running() { - let (finisher, _reviewer, ran) = finisher_over_greet(Ok(ReviewOutcome::Declined)); - let finished = finisher(pending(json!({"name": "Alice"}))).await; - assert!(ran.0.lock().unwrap().is_empty()); - assert_eq!(finished, Some(FinishedUserTool::Result(json!("Rejected")))); -} - -#[tokio::test] -async fn a_cancelled_review_and_an_unavailable_reviewer_fail_the_call_closed() { - let (finisher, _reviewer, ran) = finisher_over_greet(Ok(ReviewOutcome::Cancelled)); - let Some(FinishedUserTool::Error(message)) = finisher(pending(json!({"name": "Alice"}))).await - else { - panic!("a cancelled review is an error the model reads"); - }; - assert!(message.contains("cancelled"), "{message}"); - assert!(ran.0.lock().unwrap().is_empty()); - - let (finisher, _reviewer, ran) = finisher_over_greet(Err(ReviewError::Unavailable( - "another question is pending".to_owned(), - ))); - let Some(FinishedUserTool::Error(message)) = finisher(pending(json!({"name": "Alice"}))).await - else { - panic!("an unavailable reviewer is an error the model reads"); - }; - assert!(message.contains("another question is pending"), "{message}"); - assert!(ran.0.lock().unwrap().is_empty()); -} - -#[tokio::test] -async fn edited_arguments_the_tool_rejects_never_run() { - let (finisher, _reviewer, ran) = finisher_over_greet(Ok(ReviewOutcome::Accepted( - BTreeMap::from([("name".to_owned(), json!(42))]), - ))); - let Some(FinishedUserTool::Error(message)) = finisher(pending(json!({"name": "Alice"}))).await - else { - panic!("invalid arguments are an error the model reads"); - }; - assert!(message.contains("not valid"), "{message}"); - assert!(ran.0.lock().unwrap().is_empty()); -} - -#[tokio::test] -async fn a_tool_failure_after_acceptance_is_the_tools_own_error() { - let (finisher, _reviewer, _ran) = - finisher_over_greet(Ok(ReviewOutcome::Accepted(BTreeMap::new()))); - assert_eq!( - finisher(pending(json!({"name": "nobody"}))).await, - Some(FinishedUserTool::Error("nobody is not a person".to_owned())) - ); -} - -#[tokio::test] -async fn a_tool_the_toolset_does_not_know_as_a_user_tool_is_left_alone() { - let (finisher, reviewer, _ran) = - finisher_over_greet(Ok(ReviewOutcome::Accepted(BTreeMap::new()))); - let finished = finisher(PendingUserTool { - tool_name: "SomethingElse".to_owned(), - tool_call_id: "toolu_2".to_owned(), - args: json!({}), - }) - .await; - assert_eq!(finished, None); - assert!(reviewer.asked.lock().unwrap().is_empty()); -} diff --git a/crates/mcp_toolset/Cargo.toml b/crates/mcp_toolset/Cargo.toml index 0b659d0a104..ab606186947 100644 --- a/crates/mcp_toolset/Cargo.toml +++ b/crates/mcp_toolset/Cargo.toml @@ -12,8 +12,10 @@ rmcp = { workspace = true, features = ["client"] } schemars.workspace = true serde_json.workspace = true thiserror.workspace = true +tokio.workspace = true tracing.workspace = true workspace-hack = { version = "0.1", path = "../workspace-hack" } [dev-dependencies] regex.workspace = true +rmcp = { workspace = true, features = ["server"] } diff --git a/crates/mcp_toolset/src/lib.rs b/crates/mcp_toolset/src/lib.rs index 2b41e282b1d..979ef258213 100644 --- a/crates/mcp_toolset/src/lib.rs +++ b/crates/mcp_toolset/src/lib.rs @@ -25,7 +25,7 @@ use rmcp::service::RunningService; pub const MCP_CLIENT_NAME: &str = "Macro"; /// A connected MCP server session. -pub type McpServer = RunningService; +pub type McpServer = RunningService>>; /// Build the client info sent to MCP servers during initialization. pub fn client_info() -> ClientInfo { diff --git a/crates/mcp_toolset/src/toolset.rs b/crates/mcp_toolset/src/toolset.rs index a9bff2128e2..ab2579a0855 100644 --- a/crates/mcp_toolset/src/toolset.rs +++ b/crates/mcp_toolset/src/toolset.rs @@ -67,9 +67,14 @@ impl RemoteMcpToolSet { #[tracing::instrument(skip_all, fields(servers = servers.len(), subject = ?subject))] pub async fn from_connected(servers: Vec, subject: Option) -> Self { let listings = futures::future::join_all(servers.into_iter().map(|server| async move { - match server.client.list_all_tools().await { - Ok(tools) => Some((server.name, server.client, tools)), - Err(error) => { + match tokio::time::timeout( + std::time::Duration::from_secs(20), + server.client.list_all_tools(), + ) + .await + { + Ok(Ok(tools)) => Some((server.name, server.client, tools)), + error => { tracing::warn!( server = %server.name, error = ?error, @@ -175,6 +180,7 @@ impl RemoteMcpToolSet { &self, name: &str, arguments: serde_json::Map, + request_context: RequestContext, ) -> Result { let key = MangledName(name.to_owned()); let entry = self @@ -183,13 +189,36 @@ impl RemoteMcpToolSet { .get(&key) .ok_or_else(|| Error::UnknownTool(name.to_owned()))?; + if request_context.cancel.is_cancelled() { + return Err(Error::ToolCall("the agent turn was cancelled".to_owned())); + } let params = CallToolRequestParams::new(entry.tool.name.clone()).with_arguments(arguments); - entry + let mut pending = entry .peer - .call_tool(params) + .send_cancellable_request( + rmcp::model::CallToolRequest::new(params).into(), + Default::default(), + ) .await - .map_err(|e| Error::ToolCall(e.to_string())) + .map_err(|error| Error::ToolCall(error.to_string()))?; + let mut cancel_on_drop = CancelOnDrop(Some((entry.peer.clone(), pending.id.clone()))); + let response = tokio::select! { + biased; + _ = request_context.cancel.cancelled() => { + cancel_on_drop.0.take(); + let _ = pending.cancel(Some("the agent turn was cancelled".to_owned())).await; + return Err(Error::ToolCall("the agent turn was cancelled".to_owned())); + } + response = &mut pending.rx => response, + } + .map_err(|error| Error::ToolCall(error.to_string()))? + .map_err(|error| Error::ToolCall(error.to_string()))?; + cancel_on_drop.0.take(); + match response { + rmcp::model::ServerResult::CallToolResult(result) => Ok(result), + _ => Err(Error::ToolCall("unexpected MCP tool response".to_owned())), + } } } @@ -197,7 +226,7 @@ impl ToolSet for RemoteMcpToolSet { fn try_tool_call<'a>( &'a self, _context: Context, - _request_context: RequestContext, + request_context: RequestContext, tool_name: &'a str, json: &'a serde_json::Value, ) -> Pin< @@ -209,7 +238,7 @@ impl ToolSet for RemoteMcpToolSet { _ => serde_json::Map::new(), }; - let result = match self.call_tool(tool_name, arguments).await { + let result = match self.call_tool(tool_name, arguments, request_context).await { Ok(result) => result, Err(Error::UnknownTool(name)) => { return Err(ToolSetError::NotFound(name)); @@ -276,3 +305,23 @@ impl ToolSet for RemoteMcpToolSet { }) } } + +/// Dropping a Rig tool future must also release the server's pending review. +struct CancelOnDrop(Option<(Peer, rmcp::model::RequestId)>); +impl Drop for CancelOnDrop { + fn drop(&mut self) { + if let Some((peer, request_id)) = self.0.take() { + tokio::spawn(async move { + let _ = peer + .notify_cancelled(rmcp::model::CancelledNotificationParam { + request_id, + reason: Some("the tool call was abandoned".to_owned()), + }) + .await; + }); + } + } +} + +#[cfg(test)] +mod test; diff --git a/crates/mcp_toolset/src/toolset/test.rs b/crates/mcp_toolset/src/toolset/test.rs new file mode 100644 index 00000000000..8fe5537fe71 --- /dev/null +++ b/crates/mcp_toolset/src/toolset/test.rs @@ -0,0 +1,84 @@ +use super::*; +use rmcp::ServiceExt; +use std::sync::atomic::{AtomicUsize, Ordering}; + +struct Probe(Arc, Arc); +impl rmcp::ServerHandler for Probe { + async fn list_tools( + &self, + _: Option, + _: rmcp::service::RequestContext, + ) -> Result { + Ok(rmcp::model::ListToolsResult { + tools: vec![Tool::new( + "probe", + "A probe", + Arc::new( + serde_json::json!({"type":"object","properties":{}}) + .as_object() + .unwrap() + .clone(), + ), + )], + ..Default::default() + }) + } + async fn call_tool( + &self, + _: CallToolRequestParams, + context: rmcp::service::RequestContext, + ) -> Result { + self.0.fetch_add(1, Ordering::SeqCst); + context.ct.cancelled().await; + self.1.notify_one(); + Ok(CallToolResult::success(vec![])) + } +} + +#[tokio::test] +async fn cancelled_and_abandoned_calls_do_not_leave_server_work_running() { + let calls = Arc::new(AtomicUsize::new(0)); + let cancelled = Arc::new(tokio::sync::Notify::new()); + let (a, b) = tokio::io::duplex(8192); + let probe = Probe(calls.clone(), cancelled.clone()); + let server = tokio::spawn(async move { probe.serve(a).await.unwrap() }); + let client = crate::client_info().into_dyn().serve(b).await.unwrap(); + let server = server.await.unwrap(); + let tools = RemoteMcpToolSet::from_connected( + vec![ConnectedServer { + name: "macro".into(), + client, + }], + None, + ) + .await; + let context = RequestContext::new("macro|alice@example.com".try_into().unwrap()); + context.cancel.cancel(); + assert!( + tools + .call_tool("mcp__macro__probe", Default::default(), context) + .await + .is_err() + ); + assert_eq!(calls.load(Ordering::SeqCst), 0); + + let context = RequestContext::new("macro|alice@example.com".try_into().unwrap()); + let active_tools = tools.clone(); + let task = tokio::spawn(async move { + active_tools + .call_tool("mcp__macro__probe", Default::default(), context) + .await + }); + tokio::time::timeout(std::time::Duration::from_secs(5), async { + while calls.load(Ordering::SeqCst) == 0 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + task.abort(); + tokio::time::timeout(std::time::Duration::from_secs(5), cancelled.notified()) + .await + .unwrap(); + server.cancel().await.unwrap(); +} diff --git a/crates/pipedream_mcp/src/outbound/api.rs b/crates/pipedream_mcp/src/outbound/api.rs index 2c3a78797fb..05e5f796ec0 100644 --- a/crates/pipedream_mcp/src/outbound/api.rs +++ b/crates/pipedream_mcp/src/outbound/api.rs @@ -372,7 +372,7 @@ impl McpConnection for PipedreamClient { let config = StreamableHttpClientTransportConfig::with_uri(upstream.url.as_str()); let transport = StreamableHttpClientTransport::with_client(client, config); - Ok(client_info().serve(transport).await?) + Ok(client_info().into_dyn().serve(transport).await?) } } diff --git a/docs/ACP_ELICITATION.md b/docs/ACP_ELICITATION.md index e3f56bbee74..53474649d73 100644 --- a/docs/ACP_ELICITATION.md +++ b/docs/ACP_ELICITATION.md @@ -318,29 +318,10 @@ branch: - `scripts/convert_stdio_recording.py` turns a stdio recorder's log into a fold fixture. -Second pass, same branch — Macro's own user tools reviewed through -elicitation (see [User tools reviewed in the turn](#user-tools-reviewed-in-the-turn)): - -- `agent`: `AgentLoop::with_user_tool_finisher`. The stream bridge's - `on_tool_result` hands a `"PendingUserExecution"` answer to the finisher - and rewrites what the model reads (and the stream records) to what the - user decided. -- `ai_tools::user_tool_review`: the `UserToolReviewer` port in a neutral - form vocabulary, the schema→form projection, and `user_tool_finisher` - over `is_valid_tool` / `try_user_tool_call` — the same pair chat's - `/tool/call` uses. -- `agent_inmem`: the ACP requester implements the port; a review is a - form elicitation scoped to the call with `_meta.macro.userTool`. The idle - timeout re-arms while any question is out. `AiHost::AgentSession` pairs - chat's toolset with a prompt that describes the review card. -- `agent_fold`: `ElicitationRequest::UserTool { tool, draft, schema }`, - recognized from the absorbed user-tool call or `_meta.macro.userTool`; - `MessagePart::Elicitation.tool_outcome` from the absorbed call's later - updates. -- Web: the calendar and email composers split from their chat bindings - behind `UserToolReviewSink`; `ElicitationPart` mounts them for a review - and settles into the finished user tool; the owner gate lives in the - controller; the MagicChip gets an `asking` state with a compact card. +Macro product tools now run through MCP for the in-memory agent and external +harnesses. The MCP server reviews user tools before execution, and inmem forwards +forms to ACP. The chat composer retains its deferred user-tool path. See +[User tools reviewed in the turn](#user-tools-reviewed-in-the-turn). Third pass, same branch - the answer surface typed end to end (see [Answers are shapes, not values](#answers-are-shapes-not-values)): @@ -359,68 +340,60 @@ are shapes, not values](#answers-are-shapes-not-values)): The radio needs no synthetic `__custom` value. Decisions use exhaustive matches, and `pattern` / `format` / `looksSuspicious` are gone. -Still to do: request scope, more than one outstanding question per session -(Claude Code's parallel subagents), the MCP server reviewing user tools for -sandboxed harnesses. - ## User tools reviewed in the turn -Macro's user tools (`SendEmail`, `CreateCalendarEvent`) are registered with -`add_user_tool`: calling one returns `"PendingUserExecution"` and does -nothing, and the *host* finishes the call with the pieces `ai_toolset` -exposes for that — `is_valid_tool` for edited arguments, -`try_user_tool_call` to run the wrapped tool, `UserToolResponse` as the -answer. Chat is one host: it finishes after the turn, over HTTP, from its -composer. An agent session is another: it finishes *inside* the turn, -through elicitation, before the model reads the result. +Macro MCP exposes the canonical product catalog, including `SendEmail` and +`CreateCalendarEvent`. The main inmem agent loads product tools through that MCP +server; only `AskUser`, `SearchTools`, `LoadTools`, and `DisplayResults` are local. +Memory lookup and generation remain native. The older channel bot retains its +existing direct calendar creation and does not gain `SendEmail`. ```text -model calls CreateCalendarEvent(draft) - └─ UserTool::call → "PendingUserExecution" - └─ StreamBridge::on_tool_result (agent) - └─ UserToolFinisher (ai_tools) - ├─ form = project(tool input schema, draft) + `draft` (_macro/json) - ├─ reviewer.review(..) → elicitation/create (agent_inmem) - │ sessionId, toolCallId, mode: form, requestedSchema, - │ _meta.macro.userTool = { name, draft } - │ ← accept {content} | decline | cancel - ├─ accept: args = apply(draft, content); is_valid_tool; try_user_tool_call - │ → Rewrite(UserToolResponse::UserAction(result)) - ├─ decline → Rewrite("Rejected") - └─ cancel / unreachable client → tool error, nothing runs +model → mcp__macro__SendEmail → Macro MCP + ├─ elicitation/create (MCP 2025-11-25 form) + │ → inmem MCP client → ACP elicitation/create + │ ← accept / decline / cancel + ├─ accept: validate edited arguments, execute tool + └─ return UserToolResponse or tool error ``` -Why this shape rather than a second, "reviewing" tool wrapper: - -- One contract. `PendingUserExecution` already means "the host finishes - this"; the session just finishes sooner. Toolset, schemas, descriptions - and the `UserToolResponse` output type are identical across chat and - sessions, so the fold's `user_tool_outcome` reader and the generated - frontend types need nothing new. -- `toolCallId` for free. The hook has the call's id, so the elicitation is - properly tool-call-scoped and the fold's absorption replaces the tool row - with the question. -- Generic forms. The flat elicitation form is projected from the tool's - input schema (top-level string / boolean / number / enum arguments, - pre-filled from the call); anything nested rides in one `draft` field of - Macro's `_macro/json` type. Any user tool is reviewable with no per-tool - code; a client without a composer edits the flat fields, Macro's client - renders the tool's composer and sends the whole edited draft. -- Fail closed. No form capability, a slot already taken, or a cancelled - turn all read to the model as an error; nothing is created silently. - -The fold types the review — `ElicitationRequest::UserTool` — so the web -`match`es on it exhaustively: the session block mounts the calendar or -email composer over the draft (through `UserToolReviewSink`, the same -components chat uses) and the MagicChip shows a compact summary with -Create/Cancel and "Edit in session". Once the tool reports, the part's -`tool_outcome` carries its result and the question renders as the finished -user tool. - -Hosts (`ai_tools::AiHost`): `Chat` keeps the deferring registrations and the -composer prompt; `AgentSession` keeps the same tools with the review prompt; -`ChannelBot` and `Mcp` are unchanged — direct `CreateCalendarEvent`, no -`SendEmail` — until the MCP server reviews through `rmcp`'s elicitation. +External MCP clients answer the same form directly. Clients must advertise form +support; missing support, cancellation, invalid content, and a failed review +cannot execute a reviewed tool. A review expires after one hour. The server +returns the existing `UserToolResponse` envelope; chat still receives +`PendingUserExecution` and finishes through its composer endpoint. The agent-loop +finisher and neutral `UserToolReviewer` abstraction have been removed. + +Forms use standard MCP primitive fields, projected from the draft with defaults. +The optional string `draft` accepts the complete edited arguments as JSON, allowing +nested recipients and other complex arguments without a private schema type. +Malformed JSON is rejected rather than falling back to the original draft. +`_meta.macro.userTool = { name, draft }` lets Macro clients open the existing email +or calendar composer. Email forms default to Markdown; the email composer +explicitly submits `bodyFormat: base64url_html`, which preserves its rendered HTML. +The server renders generic Markdown submissions and passes encoded HTML to the +existing email service. Old review forms without this field retain their original +composer response shape. + +The inmem MCP client advertises form support only when the ACP client supports it. +It prefixes the question with the server name, preserves standard schemas and +answers, and trusts Macro composer metadata only from the reserved Macro server. +A per-session gate serializes concurrent forms and native questions; time queued +or awaiting an answer does not consume the turn's five-minute idle allowance. +Stopping the turn cancels the MCP request and pending review. + +These forms are session-scoped. The current Rig dispatch API does not expose an +exact provider tool-call ID to the MCP client, so no guessed ID is attached. The +composer review and final MCP tool-result row remain separate. Existing logs with +tool-scoped reviews still fold as before. Old native tool names are normalized +only in the model's copied context; persisted transcripts are not rewritten. + +The server uses stateful Streamable HTTP with SSE and authenticated session +ownership. Redis locates the process holding each session; requests landing on +another replica are streamed to that owner. Session IDs do not grant access. +Each hop verifies the bearer token, and session routing checks its user. A process +restart expires its pending sessions; it never restores or replays an uncertain +mutation. See [deployment and verification](MCP_TOOL_CONSOLIDATION.md). ## Answers are shapes, not values diff --git a/docs/AGENT_GUIDE/ai-chat.md b/docs/AGENT_GUIDE/ai-chat.md index cffccd664d4..aa546c923c3 100644 --- a/docs/AGENT_GUIDE/ai-chat.md +++ b/docs/AGENT_GUIDE/ai-chat.md @@ -81,6 +81,15 @@ to `Question · ` with `Answered` / `Declined` / `Cancelled` on the right continues. Messages typed while a question is open queue behind it; the composer's `Stop` square cancels the question and the turn. +The new Macro agent uses Macro MCP for workspace tools. Sending email and creating +calendar events opens a review composer; edit the draft and explicitly Send/Create, +or decline. The accepted review and the final tool result may appear as separate +rows. Questions from other MCP servers show the requesting server name and use +ordinary forms. Concurrent questions appear one at a time. If a server session +expires during a deploy, reconnect the agent session; do not blindly repeat an +action whose result is uncertain. The classic in-channel `@Macro` behavior is +unchanged. + ## In channels Mention `@Macro` in any channel message for the classic in-channel reply. Mention diff --git a/docs/MCP_TOOL_CONSOLIDATION.md b/docs/MCP_TOOL_CONSOLIDATION.md new file mode 100644 index 00000000000..e8ce25a2f19 --- /dev/null +++ b/docs/MCP_TOOL_CONSOLIDATION.md @@ -0,0 +1,67 @@ +# Macro MCP consolidation + +The main inmem agent consumes Macro product tools through `/mcp-macro`, the same +authenticated egress path used by sandbox harnesses. Macro MCP is required at +session creation and resume. Optional integrations may fail independently. The +catalog is discovered on connect and loaded on demand through `SearchTools` and +`LoadTools`; tools have only their `mcp__macro__` names. There is no native +product-tool fallback. Memory and harness-local utilities stay in process. + +MCP uses the stable 2025-11-25 server-to-client form elicitation flow. No +`mcp_2026_07_28` feature or MRTR protocol is required. URL elicitation is outside +this migration. The older channel bot keeps its current tool access. + +## Stateful sessions across replicas + +- Redis stores session → verified user, process incarnation, private IP/port. + Entries expire after two hours without HTTP activity. Process heartbeats run + every 15 seconds with a 60-second lease. +- The request is authenticated before routing. Different users receive 403. + A request reaching another replica is streamed over the private service port + to the recorded owner, with bearer authentication repeated there. Forwarding + cannot loop or redirect and never retries a request. +- ECS supplies the container's private IP via `ECS_CONTAINER_METADATA_URI_V4`. + Standalone replicas can set `MCP_REPLICA_ADDRESS=IP:port`; otherwise a local + server uses `127.0.0.1:`. These optional runtime variables are declared + with `macro_env_var`; ECS supplies its own metadata variable. No new Doppler + secret is needed. Redis continues using the existing `REDIS_URL`. +- The service security group permits TCP from its own security group on the + application port. Autoscaling remains enabled; no ALB affinity is required. +- Reviews expire after one hour. The MCP transport closes idle workers after + 61 minutes and removes their handles; routing rejects expired sessions before + dispatch. Sessionless POSTs must be initialization requests before a worker is + allocated. SSE heartbeats keep intermediaries alive while users answer. +- On shutdown, the process stops its heartbeat and removes its liveness entry + before stopping the listener, then cancels outstanding MCP streams. Pending + reviews are cancelled and sessions expire. Abrupt process loss instead expires + after the 60-second lease (transport failures can return 503 sooner). No pending + mutation is restored on another process. Reconnect/resume initializes fresh MCP + sessions; inspect an uncertain action's outcome before requesting it again. + +## Deployment order + +1. Deploy the MCP server and replica network rule. Verify a form review works + through two replicas, including a response landing on the other replica. +2. Deploy inmem and web. Inmem requires Macro MCP's reviewed email/calendar + catalog before allowing a production session to run. +3. Roll back in reverse order: revert the inmem consumer before reverting MCP. + Existing external clients without form support can still use other tools; + reviewed actions return an explicit unsupported error without executing. + +## Verification + +Run crate tests separately for `mcp_service`, `agent_inmem`, `agent`, `ai_tools`, +`mcp_toolset`, `agent_egress`, and `agent_fold`, with `SQLX_OFFLINE` unset and a +local database migrated to this branch. The HTTP regression test drives real +MCP requests across two listening replicas with a deterministic directory, +covering edited acceptance, decline, cancel, invalid content, missing capability, +wrong-user responses, and expired ownership. Domain routing tests cover a new +process at the previous process's address. Inmem tests cover the actual rmcp +form bridge, metadata trust, concurrent forms, waiting beyond five minutes, +and historical tool names. + +For a manual harness probe, configure a legacy Streamable HTTP MCP endpoint with +form capability enabled in the client and call a reviewed tool. Confirm the form +appears in normal and unrestricted execution modes, that editing changes what is +sent, and that decline/cancel performs no action. Harness approval settings are +separate from the server's requirement for an accepted elicitation response. diff --git a/infra/stacks/mcp-server/mcp-server.ts b/infra/stacks/mcp-server/mcp-server.ts index f4a810247b5..d97225fa6e6 100644 --- a/infra/stacks/mcp-server/mcp-server.ts +++ b/infra/stacks/mcp-server/mcp-server.ts @@ -371,6 +371,20 @@ export class McpServer extends pulumi.ComponentResource { { parent: this } ); + new aws.vpc.SecurityGroupIngressRule( + `${BASE_NAME}-replica-in`, + { + securityGroupId: serviceSg.id, + description: 'Route authenticated MCP sessions to their owning replica', + referencedSecurityGroupId: serviceSg.id, + fromPort: serviceContainerPort, + toPort: serviceContainerPort, + ipProtocol: 'tcp', + tags: this.tags, + }, + { parent: this } + ); + new aws.vpc.SecurityGroupEgressRule( `${BASE_NAME}-all-out`, { diff --git a/services/mcp_auth_proxy/src/inbound/axum_router.rs b/services/mcp_auth_proxy/src/inbound/axum_router.rs index 88ddbd1b764..99be8c29d7b 100644 --- a/services/mcp_auth_proxy/src/inbound/axum_router.rs +++ b/services/mcp_auth_proxy/src/inbound/axum_router.rs @@ -266,6 +266,7 @@ fn mcp_cors_layer() -> CorsLayer { header::AUTHORIZATION, HeaderName::from_static("mcp-protocol-version"), HeaderName::from_static("mcp-session-id"), + HeaderName::from_static("last-event-id"), ]) .expose_headers([ HeaderName::from_static("mcp-session-id"), @@ -273,3 +274,6 @@ fn mcp_cors_layer() -> CorsLayer { ]) .max_age(Duration::from_secs(3600)) } + +#[cfg(test)] +mod test; diff --git a/services/mcp_auth_proxy/src/inbound/axum_router/test.rs b/services/mcp_auth_proxy/src/inbound/axum_router/test.rs new file mode 100644 index 00000000000..08e8db0c4e6 --- /dev/null +++ b/services/mcp_auth_proxy/src/inbound/axum_router/test.rs @@ -0,0 +1,31 @@ +use super::*; +use axum::{body::Body, http::Request}; +use tower::ServiceExt; + +#[tokio::test] +async fn browser_sse_resume_preflight_allows_last_event_id() { + let app = Router::new() + .route("/mcp", routing::get(|| async { "ok" })) + .layer(mcp_cors_layer()); + let response = app + .oneshot( + Request::builder() + .method(Method::OPTIONS) + .uri("/mcp") + .header("origin", "https://client.example") + .header("access-control-request-method", "GET") + .header( + "access-control-request-headers", + "authorization,mcp-session-id,last-event-id", + ) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert!(response.status().is_success()); + let allowed = response.headers()["access-control-allow-headers"] + .to_str() + .unwrap(); + assert!(allowed.contains("last-event-id")); +} diff --git a/services/mcp_service/Cargo.toml b/services/mcp_service/Cargo.toml index c3070502be4..9dca971e8b4 100644 --- a/services/mcp_service/Cargo.toml +++ b/services/mcp_service/Cargo.toml @@ -56,14 +56,19 @@ macro_queues = { path = "../../crates/macro_queues" } macro_user_id = { path = "../../crates/macro_user_id" } mcp_auth_proxy = { path = "../mcp_auth_proxy" } notification = { path = "../../crates/notification", features = ["ai_tool"] } +base64.workspace = true +pulldown-cmark = { version = "0.13", default-features = false, features = ["html"] } prompt = { path = "../../crates/prompt" } readonly_pool = { path = "../../crates/readonly_pool" } redis = { workspace = true, features = ["aio", "tokio-comp"] } reminders = { path = "../../crates/reminders", default-features = false, features = ["ports"] } reqwest = { workspace = true } -rmcp = { workspace = true, features = ["transport-streamable-http-server"] } +rmcp = { workspace = true, features = ["transport-streamable-http-server", "elicitation"] } search_service_client = { path = "../../crates/search_service_client" } secretsmanager_client = { path = "../../crates/secretsmanager_client" } +serde = { workspace = true, features = ["derive"] } +tower = { workspace = true, features = ["util"] } +macro_uuid = { path = "../../crates/macro_uuid" } serde_json = { workspace = true } soup = { path = "../../crates/soup", features = ["outbound"] } sqlx = { workspace = true } @@ -82,3 +87,7 @@ url = { workspace = true } macro_service_urls = { path = "../../crates/macro_service_urls" } workspace-hack = { version = "0.1", path = "../../crates/workspace-hack" } + +[dev-dependencies] +schemars.workspace = true +futures.workspace = true diff --git a/services/mcp_service/src/main.rs b/services/mcp_service/src/main.rs index 0571f75dcfc..4af62b112df 100644 --- a/services/mcp_service/src/main.rs +++ b/services/mcp_service/src/main.rs @@ -7,6 +7,7 @@ mod config; mod context; mod markdown_images; +mod session_routing; mod tool_service; use anyhow::Context; use config::Config; @@ -14,10 +15,7 @@ use context::build_context; use macro_entrypoint::MacroEntrypoint; use mcp_auth_proxy::domain::service::McpAuthProxyService; use mcp_auth_proxy::inbound::axum_router::mcp_router; -use rmcp::transport::streamable_http_server::{ - StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager, -}; -use std::sync::Arc; +use rmcp::transport::streamable_http_server::{StreamableHttpServerConfig, StreamableHttpService}; use tokio::time::Duration; use tokio_util::task::TaskTracker; use tool_service::AuthenticatedToolService; @@ -39,6 +37,34 @@ async fn main() -> anyhow::Result<()> { let event_broker_tracker = TaskTracker::new(); let context = build_context(&config, event_broker_tracker.clone()).await?; + let process = macro_uuid::generate_uuid_v7().to_string(); + let address = session_routing::replica_address(config.port) + .await + .map_err(anyhow::Error::msg)?; + let directory = + session_routing::RedisDirectory::new(config.redis_url.as_ref(), &context.mcp_public_host)?; + directory + .heartbeat(&process) + .await + .map_err(anyhow::Error::msg)?; + let heartbeat_directory = directory.clone(); + let heartbeat_process = process.clone(); + let heartbeat = tokio::spawn(async move { + let mut interval = tokio::time::interval(Duration::from_secs(15)); + loop { + interval.tick().await; + if let Err(error) = heartbeat_directory.heartbeat(&heartbeat_process).await { + tracing::error!(error=?error, "MCP replica heartbeat failed"); + } + } + }); + let public_host = context.mcp_public_host.clone(); + let mut sessions = + rmcp::transport::streamable_http_server::session::local::LocalSessionManager::default(); + sessions.session_config.keep_alive = Some(Duration::from_secs(3660)); + let sessions = std::sync::Arc::new(sessions); + let shutdown = tokio_util::sync::CancellationToken::new(); + // Create the MCP service with authenticated tool handler let mcp_service = StreamableHttpService::new( move || { @@ -49,15 +75,16 @@ async fn main() -> anyhow::Result<()> { item_base_url.clone(), )) }, - Arc::new(LocalSessionManager::default()), + sessions.clone(), { let mut config = StreamableHttpServerConfig::default().with_allowed_hosts([ context.mcp_public_host.clone(), "localhost".into(), "127.0.0.1".into(), ]); - config.stateful_mode = false; - config.json_response = true; + config.cancellation_token = shutdown.clone(); + config.stateful_mode = true; + config.json_response = false; config }, ); @@ -74,7 +101,15 @@ async fn main() -> anyhow::Result<()> { } }); - let app = mcp_router(context.auth_proxy, context.jwt_args, mcp_service); + let routed = session_routing::route_sessions( + mcp_service, + directory.clone(), + process.clone(), + address, + public_host, + sessions, + ); + let app = mcp_router(context.auth_proxy, context.jwt_args, routed); let port = config.port; let addr = format!("0.0.0.0:{port}"); @@ -84,10 +119,31 @@ async fn main() -> anyhow::Result<()> { tracing::info!("MCP server listening on http://{addr}/mcp"); - let server_result = axum::serve(listener, app) - .with_graceful_shutdown(macro_entrypoint::shutdown_signal()) - .await - .context("MCP server error"); + let heartbeat_abort = heartbeat.abort_handle(); + let shutdown_observed = shutdown.clone(); + let server = axum::serve(listener, app).with_graceful_shutdown(async move { + macro_entrypoint::shutdown_signal().await; + // Stop advertising before the listener stops accepting answer POSTs. + heartbeat_abort.abort(); + if let Err(error) = directory.retire(&process).await { + tracing::error!(error=?error, "failed to retire MCP replica; lease will expire"); + } + shutdown.cancel(); + }); + let server = std::future::IntoFuture::into_future(server); + tokio::pin!(server); + let server_result = tokio::select! { + result = &mut server => result.context("MCP server error"), + _ = shutdown_observed.cancelled() => { + match tokio::time::timeout(Duration::from_secs(10), &mut server).await { + Ok(result) => result.context("MCP server error"), + Err(_) => { + tracing::warn!("MCP stream drain deadline reached"); + Ok(()) + } + } + } + }; tracing::info!("waiting for event broker publishes to drain"); event_broker_tracker.close(); @@ -102,5 +158,6 @@ async fn main() -> anyhow::Result<()> { } } + heartbeat.abort(); server_result } diff --git a/services/mcp_service/src/session_routing.rs b/services/mcp_service/src/session_routing.rs new file mode 100644 index 00000000000..172e027ffd9 --- /dev/null +++ b/services/mcp_service/src/session_routing.rs @@ -0,0 +1,7 @@ +//! Stateful MCP routing. Only the owning process can resume pending JSON-RPC calls. +pub(crate) mod directory; +mod http; +mod redis; + +pub(crate) use http::route_sessions; +pub(crate) use redis::{RedisDirectory, replica_address}; diff --git a/services/mcp_service/src/session_routing/directory.rs b/services/mcp_service/src/session_routing/directory.rs new file mode 100644 index 00000000000..2bf2d0a5804 --- /dev/null +++ b/services/mcp_service/src/session_routing/directory.rs @@ -0,0 +1,53 @@ +//! Session ownership and routing policy, independent of the HTTP/Redis adapters. +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub(crate) struct Owner { + pub user: String, + pub process: String, + pub address: std::net::SocketAddr, +} + +#[async_trait::async_trait] +pub(crate) trait Directory: Clone + Send + Sync + 'static { + async fn lookup(&self, session: &str) -> Result, String>; + async fn register(&self, session: &str, owner: &Owner) -> Result<(), String>; + async fn remove(&self, session: &str) -> Result<(), String>; +} + +#[derive(Debug, PartialEq, Eq)] +pub(super) enum Route { + Local, + Forward(std::net::SocketAddr), + Forbidden, + Expired, +} + +/// A session ID is a locator, never an authentication credential. A forwarded +/// request must arrive at the exact process incarnation recorded in the directory. +pub(super) fn route( + owner: Option<&Owner>, + user: &str, + process: &str, + forwarded: Option<&str>, +) -> Route { + let Some(owner) = owner else { + return Route::Expired; + }; + if owner.user != user { + return Route::Forbidden; + } + if let Some(expected_process) = forwarded + && (expected_process != process || owner.process != process) + { + return Route::Expired; + } + if owner.process == process { + Route::Local + } else { + Route::Forward(owner.address) + } +} + +#[cfg(test)] +mod test; diff --git a/services/mcp_service/src/session_routing/directory/test.rs b/services/mcp_service/src/session_routing/directory/test.rs new file mode 100644 index 00000000000..87a4dd60993 --- /dev/null +++ b/services/mcp_service/src/session_routing/directory/test.rs @@ -0,0 +1,36 @@ +use super::*; + +#[test] +fn responses_route_only_to_the_owning_user_and_process() { + let owner = Owner { + user: "alice".into(), + process: "process-a".into(), + address: "127.0.0.1:8001".parse().unwrap(), + }; + assert_eq!( + route(Some(&owner), "alice", "process-a", None), + Route::Local + ); + assert_eq!( + route(Some(&owner), "alice", "process-b", None), + Route::Forward(owner.address) + ); + assert_eq!( + route(Some(&owner), "bob", "process-a", None), + Route::Forbidden + ); + assert_eq!( + route(Some(&owner), "bob", "process-b", None), + Route::Forbidden + ); + assert_eq!( + route(Some(&owner), "alice", "process-a", Some("process-a")), + Route::Local + ); + // A replacement process at the same address must not restore pending calls. + assert_eq!( + route(Some(&owner), "alice", "replacement", Some("process-a")), + Route::Expired + ); + assert_eq!(route(None, "alice", "process-b", None), Route::Expired); +} diff --git a/services/mcp_service/src/session_routing/http.rs b/services/mcp_service/src/session_routing/http.rs new file mode 100644 index 00000000000..9c4b0e8c742 --- /dev/null +++ b/services/mcp_service/src/session_routing/http.rs @@ -0,0 +1,253 @@ +//! Authenticated streaming transport routing; JWT middleware runs on every hop. +use super::directory::{Directory, Owner, Route, route}; +use axum::{ + body::Body, + http::{Request, StatusCode, header}, + response::{IntoResponse, Response}, +}; +use macro_user_id::user_id::MacroUserIdStr; +use rmcp::transport::streamable_http_server::session::{ + SessionManager, local::LocalSessionManager, +}; +use std::sync::Arc; +use std::{convert::Infallible, net::SocketAddr}; +use tower::{Service, ServiceExt}; + +const SESSION: &str = "mcp-session-id"; +const FORWARDED: &str = "x-macro-mcp-process"; + +#[derive(Clone)] +struct Replica { + directory: D, + process: String, + address: SocketAddr, + public_host: String, + client: reqwest::Client, + sessions: Arc, +} + +pub(crate) fn route_sessions( + local: S, + directory: D, + process: String, + address: SocketAddr, + public_host: String, + sessions: Arc, +) -> impl Service, Response = Response, Error = Infallible, Future: Send> ++ Clone ++ Send ++ Sync +where + S: Service, Error = Infallible> + Clone + Send + Sync + 'static, + S::Response: IntoResponse, + S::Future: Send, + D: Directory, +{ + // No redirects, retries, or whole-response timeout: an accepted mutation must + // never be replayed, and SSE may be waiting for a human for an hour. + let client = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .retry(reqwest::retry::never()) + .connect_timeout(std::time::Duration::from_secs(5)) + .build() + .expect("valid replica HTTP client configuration"); + let replica = Replica { + directory, + process, + address, + public_host, + client, + sessions, + }; + tower::service_fn(move |request: Request| { + let local = local.clone(); + let replica = replica.clone(); + async move { + let result = handle(request, local, replica).await; + Ok::<_, Infallible>(result.unwrap_or_else(|error| { + tracing::error!(error=?error, "MCP session routing failed"); + ( + StatusCode::SERVICE_UNAVAILABLE, + "MCP session routing unavailable; no automatic retry", + ) + .into_response() + })) + } + }) +} + +async fn handle( + request: Request, + local: S, + Replica { + directory, + process, + address, + public_host, + client, + sessions, + }: Replica, +) -> Result +where + S: Service, Error = Infallible>, + S::Response: IntoResponse, + D: Directory, +{ + let Some(user) = request.extensions().get::>() else { + return Ok(StatusCode::UNAUTHORIZED.into_response()); + }; + let user = user.to_string(); + let session = request + .headers() + .get(SESSION) + .map(|value| value.to_str().map(str::to_owned)) + .transpose() + .map_err(|_| "invalid session header")?; + let forwarded = request + .headers() + .get(FORWARDED) + .map(|value| value.to_str()) + .transpose() + .map_err(|_| "invalid forwarding header")?; + let delete = request.method() == axum::http::Method::DELETE; + if let Some(session) = &session { + let owner = directory.lookup(session).await?; + match route(owner.as_ref(), &user, &process, forwarded) { + Route::Forbidden => return Ok(StatusCode::FORBIDDEN.into_response()), + Route::Expired => { + return Ok(( + StatusCode::NOT_FOUND, + "MCP session expired; initialize a new session", + ) + .into_response()); + } + Route::Forward(address) => { + let owner = owner.expect("forward requires owner"); + let (parts, body) = request.into_parts(); + // Forward protocol headers only; never trust client-supplied routing or Host. + let mut headers = axum::http::HeaderMap::new(); + for name in [ + header::AUTHORIZATION.as_str(), + header::CONTENT_TYPE.as_str(), + header::ACCEPT.as_str(), + SESSION, + "mcp-protocol-version", + "last-event-id", + ] { + if let Some(value) = parts.headers.get(name) { + headers.insert( + header::HeaderName::from_bytes(name.as_bytes()) + .map_err(|e| e.to_string())?, + value.clone(), + ); + } + } + headers.insert( + header::HOST, + public_host.parse().map_err(|_| "invalid public host")?, + ); + headers.insert( + FORWARDED, + owner.process.parse().map_err(|_| "invalid process id")?, + ); + let upstream = client + .request(parts.method, format!("http://{address}/mcp")) + .headers(headers) + .body(reqwest::Body::wrap_stream(body.into_data_stream())) + .send() + .await + .map_err(|e| e.to_string())?; + let status = upstream.status(); + let headers = upstream.headers().clone(); + let mut response = Response::new(Body::from_stream(upstream.bytes_stream())); + *response.status_mut() = status; + for name in [ + header::CONTENT_TYPE.as_str(), + header::CACHE_CONTROL.as_str(), + SESSION, + "mcp-protocol-version", + "www-authenticate", + ] { + if let Some(value) = headers.get(name) { + response.headers_mut().insert( + header::HeaderName::from_bytes(name.as_bytes()) + .map_err(|e| e.to_string())?, + value.clone(), + ); + } + } + return Ok(response); + } + Route::Local => { + if !sessions + .has_session(&session.clone().into()) + .await + .map_err(|e| e.to_string())? + { + directory.remove(session).await?; + return Ok(StatusCode::NOT_FOUND.into_response()); + } + } + } + } else if forwarded.is_some() { + return Ok(StatusCode::BAD_REQUEST.into_response()); + } + // rmcp allocates before rejecting a sessionless non-initialize message. + // Validate that envelope first so rejected calls cannot leak allocations. + let request = if session.is_none() && request.method() == axum::http::Method::POST { + let (parts, body) = request.into_parts(); + let bytes = match tokio::time::timeout( + std::time::Duration::from_secs(10), + axum::body::to_bytes(body, 1024 * 1024), + ) + .await + { + Ok(Ok(bytes)) => bytes, + _ => return Ok(StatusCode::BAD_REQUEST.into_response()), + }; + let valid = matches!( + serde_json::from_slice::(&bytes), + Ok(rmcp::model::JsonRpcMessage::Request( + rmcp::model::JsonRpcRequest { + request: rmcp::model::ClientRequest::InitializeRequest(_), + .. + } + )) + ); + if !valid { + return Ok(StatusCode::BAD_REQUEST.into_response()); + } + Request::from_parts(parts, Body::from(bytes)) + } else { + request + }; + let response = local + .oneshot(request) + .await + .unwrap_or_else(|never| match never {}) + .into_response(); + if session.is_none() { + if let Some(id) = response + .headers() + .get(SESSION) + .and_then(|id| id.to_str().ok()) + && let Err(error) = directory + .register( + id, + &Owner { + user, + process: process.to_owned(), + address, + }, + ) + .await + { + let _ = sessions.close_session(&id.to_owned().into()).await; + return Err(error); + } + } else if delete && response.status().is_success() { + let id = session.as_deref().expect("session exists"); + directory.remove(id).await?; + } + Ok(response) +} diff --git a/services/mcp_service/src/session_routing/redis.rs b/services/mcp_service/src/session_routing/redis.rs new file mode 100644 index 00000000000..58191a358a3 --- /dev/null +++ b/services/mcp_service/src/session_routing/redis.rs @@ -0,0 +1,154 @@ +//! Redis directory and process liveness. No pending tool calls are persisted or replayed. +use super::directory::{Directory, Owner}; +use redis::AsyncCommands; +use std::{net::SocketAddr, time::Duration}; + +macro_env_var::maybe_env_vars! { + /// Explicit private listen address for standalone replicas (IP:port). + pub struct McpReplicaAddress; + /// ECS-provided metadata endpoint for this container. + pub struct EcsContainerMetadataUriV4; +} + +const SESSION_TTL: u64 = 7200; +const PROCESS_TTL: u64 = 60; + +#[derive(Clone)] +pub(crate) struct RedisDirectory { + client: redis::Client, + prefix: String, +} + +impl RedisDirectory { + pub(crate) fn new(url: &str, namespace: &str) -> Result { + Ok(Self { + client: redis::Client::open(url)?, + prefix: format!("mcp:sessions:{namespace}"), + }) + } + + async fn connection(&self) -> Result { + self.client + .get_multiplexed_async_connection_with_config( + &redis::AsyncConnectionConfig::new() + .set_connection_timeout(Some(Duration::from_secs(5))) + .set_response_timeout(Some(Duration::from_secs(5))), + ) + .await + .map_err(|e| e.to_string()) + } + + pub(crate) async fn retire(&self, process: &str) -> Result<(), String> { + let mut conn = self.connection().await?; + conn.set_ex::<_, _, ()>( + format!("{}:process:{process}", self.prefix), + "retired", + SESSION_TTL, + ) + .await + .map_err(|e| e.to_string()) + } + + pub(crate) async fn heartbeat(&self, process: &str) -> Result<(), String> { + let mut conn = self.connection().await?; + // A delayed heartbeat must not resurrect a process after retirement. + let refreshed: bool = redis::Script::new( + "if redis.call('GET', KEYS[1]) == 'retired' then return 0 end \ + redis.call('SET', KEYS[1], 'alive', 'EX', ARGV[1]); return 1", + ) + .key(format!("{}:process:{process}", self.prefix)) + .arg(PROCESS_TTL) + .invoke_async(&mut conn) + .await + .map_err(|e| e.to_string())?; + if refreshed { + Ok(()) + } else { + Err("MCP replica was retired".into()) + } + } +} + +#[async_trait::async_trait] +impl Directory for RedisDirectory { + async fn lookup(&self, session: &str) -> Result, String> { + let mut conn = self.connection().await?; + let key = format!("{}:session:{session}", self.prefix); + let value: Option = conn.get(&key).await.map_err(|e| e.to_string())?; + let Some(value) = value else { + return Ok(None); + }; + let owner: Owner = serde_json::from_str(&value).map_err(|e| e.to_string())?; + let alive: Option = conn + .get(format!("{}:process:{}", self.prefix, owner.process)) + .await + .map_err(|e| e.to_string())?; + if alive.as_deref() != Some("alive") { + return Ok(None); + } + let _: bool = conn + .expire(key, SESSION_TTL as i64) + .await + .map_err(|e| e.to_string())?; + Ok(Some(owner)) + } + + async fn register(&self, session: &str, owner: &Owner) -> Result<(), String> { + let mut conn = self.connection().await?; + let value = serde_json::to_string(owner).map_err(|e| e.to_string())?; + conn.set_ex::<_, _, ()>( + format!("{}:session:{session}", self.prefix), + value, + SESSION_TTL, + ) + .await + .map_err(|e| e.to_string()) + } + + async fn remove(&self, session: &str) -> Result<(), String> { + let mut conn = self.connection().await?; + conn.del::<_, ()>(format!("{}:session:{session}", self.prefix)) + .await + .map_err(|e| e.to_string()) + } +} + +/// ECS supplies the private task IP. Local probes may override it without DNS +/// or a public load balancer becoming a forwarding destination. +pub(crate) async fn replica_address(port: usize) -> Result { + if let Some(address) = McpReplicaAddress::new().as_ref() { + return address.parse::().map_err(|e| e.to_string()); + } + if let Some(metadata) = EcsContainerMetadataUriV4::new().as_ref() { + let value: serde_json::Value = reqwest::Client::builder() + .timeout(Duration::from_secs(5)) + .build() + .map_err(|e| e.to_string())? + .get(metadata.as_ref()) + .send() + .await + .map_err(|e| e.to_string())? + .error_for_status() + .map_err(|e| e.to_string())? + .json() + .await + .map_err(|e| e.to_string())?; + let ip = value["Networks"] + .as_array() + .and_then(|networks| { + networks + .iter() + .find_map(|network| network["IPv4Addresses"][0].as_str()) + }) + .ok_or("ECS metadata has no private IPv4 address")?; + return format!("{ip}:{port}") + .parse::() + .map_err(|e| e.to_string()); + } + format!("127.0.0.1:{port}") + .parse::() + .map_err(|e| e.to_string()) +} + +#[cfg(test)] +mod test; diff --git a/services/mcp_service/src/session_routing/redis/test.rs b/services/mcp_service/src/session_routing/redis/test.rs new file mode 100644 index 00000000000..a9805747f41 --- /dev/null +++ b/services/mcp_service/src/session_routing/redis/test.rs @@ -0,0 +1,31 @@ +use super::*; + +#[tokio::test] +async fn redis_directory_shares_owners_and_retirement_cannot_be_undone_by_a_late_heartbeat() { + let url = crate::config::RedisUrl::new().expect("REDIS_URL for local tests"); + let namespace = format!("test-{}", macro_uuid::generate_uuid_v7()); + let a = RedisDirectory::new(url.as_ref(), &namespace).unwrap(); + let b = RedisDirectory::new(url.as_ref(), &namespace).unwrap(); + let owner = Owner { + user: "alice".into(), + process: "old".into(), + address: "127.0.0.1:1234".parse().unwrap(), + }; + a.heartbeat("old").await.unwrap(); + a.register("session", &owner).await.unwrap(); + assert_eq!(b.lookup("session").await.unwrap().unwrap().process, "old"); + a.retire("old").await.unwrap(); + assert!(a.heartbeat("old").await.is_err()); + assert!(b.lookup("session").await.unwrap().is_none()); + // A different incarnation at the same address has no claim to this session. + b.heartbeat("replacement").await.unwrap(); + assert!(b.lookup("session").await.unwrap().is_none()); + b.remove("session").await.unwrap(); + let mut conn = a.connection().await.unwrap(); + let _: () = redis::cmd("DEL") + .arg(format!("{}:process:old", a.prefix)) + .arg(format!("{}:process:replacement", a.prefix)) + .query_async(&mut conn) + .await + .unwrap(); +} diff --git a/services/mcp_service/src/tool_service.rs b/services/mcp_service/src/tool_service.rs index 740a2d99b30..b5832723b03 100644 --- a/services/mcp_service/src/tool_service.rs +++ b/services/mcp_service/src/tool_service.rs @@ -1,3 +1,5 @@ +mod review; + use crate::markdown_images::{MarkdownImageResolver, tool_result_with_images}; use ai_toolset::{AsyncToolCollection, RequestContext, ToolSet}; use macro_user_id::user_id::MacroUserIdStr; @@ -10,6 +12,24 @@ use rmcp::{ }; use std::sync::Arc; +/// Per-call context preparation keeps usage attribution tied to the verified user. +pub(crate) trait McpToolContext: Clone + Send + Sync + MarkdownImageResolver { + fn for_user(&self, user: MacroUserIdStr<'static>) -> Self; +} + +impl McpToolContext for ai_tools::ToolServiceContext { + fn for_user(&self, user: MacroUserIdStr<'static>) -> Self { + let mut context = self.clone(); + context.usage_context = ai_usage::UsageContext::new(ai_usage::AiFeature::Chat, user); + context + } +} + +#[cfg(test)] +impl McpToolContext for () { + fn for_user(&self, _user: MacroUserIdStr<'static>) -> Self {} +} + /// Maps our protocol-agnostic annotations onto the MCP wire representation. /// /// [`ToolKind`](ai_toolset::ToolKind) collapses `readOnlyHint`/`destructiveHint` @@ -85,7 +105,7 @@ mod test; impl ServerHandler for AuthenticatedToolService where - Context: Clone + Send + Sync + MarkdownImageResolver + 'static, + Context: McpToolContext + 'static, { fn get_info(&self) -> ServerInfo { let mut info = ServerInfo::new(ServerCapabilities::builder().enable_tools().build()); @@ -137,21 +157,119 @@ where ) -> Result { let user_id = Self::authenticated_user_id(&context.extensions)?; - let request_context = RequestContext::new(user_id.clone()); + let tool_context = self.context.for_user(user_id.clone()); + let mut request_context = RequestContext::new(user_id.clone()); + request_context.cancel = context.ct.clone(); let arguments = request .arguments .map(serde_json::Value::Object) .ok_or(rmcp::ErrorData::invalid_params("No params provided", None))?; - let result = self - .toolset - .try_tool_call( - self.context.clone(), - request_context, - &request.name, + if let Some(tool) = self.toolset.user_tools.get(request.name.as_ref()) { + let supports_form = context + .peer + .peer_info() + .and_then(|info| info.capabilities.elicitation.as_ref()) + .is_some_and(|cap| cap.form.is_some() || cap.url.is_none()); + if !supports_form { + return Ok(tool_error( + "This tool requires form elicitation support; nothing was executed.", + )); + } + let schema = review::project_form( + &review::tool_schema(&request.name, &tool.input_schema), &arguments, ) + .map_err(|error| rmcp::ErrorData::internal_error(error, None))?; + let params = rmcp::model::CreateElicitationRequestParams::FormElicitationParams { + meta: Some(rmcp::model::Meta( + serde_json::from_value(serde_json::json!({ + "macro": {"userTool": {"name": request.name, "draft": arguments}} + })) + .map_err(|error| rmcp::ErrorData::internal_error(error.to_string(), None))?, + )), + message: format!("{}?", tool.annotations.title), + requested_schema: schema, + }; + let mut pending = context + .peer + .send_cancellable_request( + rmcp::model::CreateElicitationRequest::new(params).into(), + Default::default(), + ) + .await + .map_err(|error| rmcp::ErrorData::internal_error(error.to_string(), None))?; + let response = tokio::select! { + biased; + _ = context.ct.cancelled() => None, + _ = tokio::time::sleep(std::time::Duration::from_secs(3600)) => None, + response = &mut pending.rx => Some(response), + }; + let response = match response { + None => { + let _ = pending + .cancel(Some("the review was cancelled or expired".to_owned())) + .await; + return Ok(tool_error( + "The review was cancelled or expired; nothing was executed.", + )); + } + Some(Ok(Ok(rmcp::model::ClientResult::CreateElicitationResult(response)))) => { + response + } + _ => return Ok(tool_error("The review failed; nothing was executed.")), + }; + match response.action { + rmcp::model::ElicitationAction::Decline => { + // MCP structuredContent must be an object. Text preserves + // the existing "Rejected" value when the inmem client reads it. + return Ok(rmcp::model::CallToolResult::success(vec![Content::text( + "Rejected", + )])); + } + rmcp::model::ElicitationAction::Cancel => { + return Ok(tool_error( + "The review was cancelled; nothing was executed.", + )); + } + rmcp::model::ElicitationAction::Accept => {} + } + let reviewed = match response + .content + .as_ref() + .ok_or_else(|| "The accepted form had no content".to_owned()) + .and_then(|content| review::reviewed_arguments(&request.name, &arguments, content)) + { + Ok(args) if self.toolset.is_valid_tool(&request.name, &args) => args, + _ => { + return Ok(tool_error( + "The reviewed arguments are invalid; nothing was executed.", + )); + } + }; + if context.ct.is_cancelled() { + return Ok(tool_error("The call was cancelled before execution.")); + } + return match self + .toolset + .try_user_tool_call(tool_context, request_context, &request.name, &reviewed) + .await + { + Ok(Ok(result)) => { + let value = serde_json::to_value(result).map_err(|error| { + rmcp::ErrorData::internal_error(error.to_string(), None) + })?; + Ok(tool_result_with_images(&self.context, &user_id, value).await) + } + Ok(Err(error)) => Ok(tool_error(error.description)), + Err(error) => Ok(tool_error(error.to_string())), + }; + } + + let result = self + .toolset + .try_tool_call(tool_context, request_context, &request.name, &arguments) .await .map_err(|error| match error { ai_toolset::ToolSetError::Deserialization(error) => { @@ -170,3 +288,7 @@ where } } } + +fn tool_error(message: impl Into) -> rmcp::model::CallToolResult { + rmcp::model::CallToolResult::error(vec![Content::text(message.into())]) +} diff --git a/services/mcp_service/src/tool_service/review.rs b/services/mcp_service/src/tool_service/review.rs new file mode 100644 index 00000000000..96ff1883887 --- /dev/null +++ b/services/mcp_service/src/tool_service/review.rs @@ -0,0 +1,205 @@ +//! Standard MCP forms for reviewing Macro product tool arguments. +use rmcp::model::ElicitationSchema; +use serde_json::{Map, Value, json}; + +pub(super) const DRAFT_FIELD: &str = "draft"; + +/// Project primitive arguments; the complete draft remains editable as JSON. +pub(super) fn project_form( + schema: &Map, + draft: &Value, +) -> Result { + let mut properties = Map::new(); + if let Some(fields) = schema.get("properties").and_then(Value::as_object) { + for (name, field) in fields { + let field = non_null(field); + let mut projected = if let Some(options) = choice_options(field) { + json!({"type": "string", "enum": options}) + } else { + let Some(kind @ ("string" | "number" | "integer" | "boolean")) = json_type(field) + else { + continue; + }; + json!({"type": kind}) + }; + for key in [ + "title", + "description", + "minimum", + "maximum", + "minLength", + "maxLength", + ] { + if let Some(value) = field.get(key) { + projected[key] = value.clone(); + } + } + if let Some(value) = draft.get(name).filter(|value| !value.is_null()) { + projected["default"] = value.clone(); + } + properties.insert(name.clone(), projected); + } + } + properties.insert(DRAFT_FIELD.to_owned(), json!({ + "type": "string", "title": "Complete draft (JSON)", + "description": format!("Optional replacement for the entire draft. Otherwise edit the fields above. Current draft: {}", draft) + })); + // Defaults are applied to the original draft. Requiring projected fields would + // prevent a composer from submitting just the complete edited JSON draft. + serde_json::from_value(json!({"type":"object", "properties": properties})) + .map_err(|error| error.to_string()) +} + +/// Merge edited fields, rejecting a malformed replacement instead of executing the old draft. +pub(super) fn apply_review(draft: &Value, content: &Value) -> Result { + let content = content + .as_object() + .ok_or("the accepted form must contain an object")?; + if let Some(whole) = content.get(DRAFT_FIELD) { + let text = whole + .as_str() + .ok_or("the complete draft must be a JSON string")?; + if !text.trim().is_empty() { + let parsed: Value = + serde_json::from_str(text).map_err(|_| "the complete draft is not valid JSON")?; + if !parsed.is_object() { + return Err("the complete draft must be an object".to_owned()); + } + return Ok(parsed); + } + } + let mut args = draft + .as_object() + .cloned() + .ok_or("the tool draft must be an object")?; + for (name, value) in content { + if name != DRAFT_FIELD { + args.insert(name.clone(), value.clone()); + } + } + Ok(Value::Object(args)) +} + +#[cfg(test)] +mod test; + +/// The string values a property allows, when it is a choice among fixed +/// strings: a plain `enum`, or - how a documented Rust enum comes out - an +/// `anyOf`/`oneOf` whose every real variant is one `const` string (or a +/// one-value `enum`). +fn choice_options(property: &Value) -> Option> { + if let Some(options) = property.get("enum").and_then(Value::as_array) { + let options: Vec = options + .iter() + .filter_map(Value::as_str) + .map(str::to_owned) + .collect(); + return (!options.is_empty()).then_some(options); + } + let variants = ["anyOf", "oneOf"] + .into_iter() + .find_map(|key| property.get(key).and_then(Value::as_array))?; + let mut options = Vec::new(); + for variant in variants { + if json_type(variant) == Some("null") { + continue; + } + let value = variant.get("const").and_then(Value::as_str).or_else(|| { + match variant.get("enum").and_then(Value::as_array)?.as_slice() { + [only] => only.as_str(), + _ => None, + } + })?; + options.push(value.to_owned()); + } + (!options.is_empty()).then_some(options) +} + +/// The property with a nullable wrapper removed: `anyOf`/`oneOf` of one real +/// schema and `null` reads as that schema. +fn non_null(property: &Value) -> &Value { + for key in ["anyOf", "oneOf"] { + if let Some(variants) = property.get(key).and_then(Value::as_array) { + let real: Vec<&Value> = variants + .iter() + .filter(|variant| json_type(variant) != Some("null")) + .collect(); + if let [only] = real[..] { + return only; + } + } + } + property +} + +/// The one JSON type a property declares, reading `["string", "null"]` as +/// `string`. `None` for a property with several real types or none. +fn json_type(property: &Value) -> Option<&str> { + match property.get("type")? { + Value::String(name) => Some(name.as_str()), + Value::Array(names) => { + let real: Vec<&str> = names + .iter() + .filter_map(Value::as_str) + .filter(|name| *name != "null") + .collect(); + match real[..] { + [only] => Some(only), + _ => None, + } + } + _ => None, + } +} + +/// Email composers return encoded HTML; standard forms edit Markdown. Keep that +/// distinction explicit rather than guessing an encoding from the user's text. +pub(super) fn tool_schema(name: &str, schema: &Map) -> Map { + let mut schema = schema.clone(); + if name == "SendEmail" + && let Some(properties) = schema.get_mut("properties").and_then(Value::as_object_mut) + { + properties.insert("bodyFormat".into(), json!({ + "type":"string", "title":"Body format", + "enum":["markdown", "base64url_html"], + "description":"Use markdown for text edits. Macro's email composer submits base64url_html." + })); + } + schema +} + +pub(super) fn reviewed_arguments( + name: &str, + draft: &Value, + content: &Value, +) -> Result { + if name != "SendEmail" { + return apply_review(draft, content); + } + use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; + let mut content = content + .as_object() + .cloned() + .ok_or("the accepted form must contain an object")?; + let format = content.remove("bodyFormat").unwrap_or(json!("markdown")); + let mut reviewed = apply_review(draft, &Value::Object(content))?; + let body = reviewed + .get("body") + .and_then(Value::as_str) + .ok_or("the email body must be text")?; + match format.as_str() { + Some("markdown") => { + let mut html = String::new(); + pulldown_cmark::html::push_html(&mut html, pulldown_cmark::Parser::new(body)); + reviewed["body"] = Value::String(URL_SAFE_NO_PAD.encode(html)); + } + Some("base64url_html") => { + let bytes = URL_SAFE_NO_PAD + .decode(body) + .map_err(|_| "the composer body is not base64url HTML")?; + std::str::from_utf8(&bytes).map_err(|_| "the composer body is not UTF-8")?; + } + _ => return Err("unsupported email body format".into()), + } + Ok(reviewed) +} diff --git a/services/mcp_service/src/tool_service/review/test.rs b/services/mcp_service/src/tool_service/review/test.rs new file mode 100644 index 00000000000..91092d4ae08 --- /dev/null +++ b/services/mcp_service/src/tool_service/review/test.rs @@ -0,0 +1,71 @@ +use super::*; + +#[test] +fn standard_form_preserves_nested_draft_and_applies_flat_edits() { + let draft = json!({"subject":"Original", "recipients":[{"email":"a@example.com"}]}); + let schema = json!({"type":"object", "properties":{ + "subject":{"type":"string"}, "recipients":{"type":"array"} + }}); + let form = + serde_json::to_value(project_form(schema.as_object().unwrap(), &draft).unwrap()).unwrap(); + assert_eq!(form["properties"]["subject"]["default"], "Original"); + assert_eq!(form["properties"]["draft"]["type"], "string"); + assert!(form["properties"].get("recipients").is_none()); + assert!(form["properties"]["draft"].get("default").is_none()); + let edited = apply_review(&draft, &json!({"subject":"Edited", "draft":""})).unwrap(); + assert_eq!(edited["subject"], "Edited"); + assert_eq!(edited["recipients"], draft["recipients"]); +} + +#[test] +fn composer_replacement_wins_over_prepopulated_flat_fields() { + let edited = json!({"subject":"Composer", "recipients":[]}); + assert_eq!( + apply_review( + &json!({"subject":"Original"}), + &json!({ + "subject":"Original", "draft":edited.to_string() + }) + ) + .unwrap(), + edited + ); +} + +#[test] +fn malformed_accepted_content_never_falls_back_to_original() { + for content in [ + json!(null), + json!({"draft":{}}), + json!({"draft":"broken"}), + json!({"draft":"[]"}), + ] { + assert!(apply_review(&json!({"subject":"Original"}), &content).is_err()); + } +} + +#[test] +fn email_forms_render_markdown_and_preserve_explicit_composer_html() { + use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; + let draft = json!({"body":"**Hello**"}); + let generic = reviewed_arguments("SendEmail", &draft, &json!({})).unwrap(); + let html = String::from_utf8( + URL_SAFE_NO_PAD + .decode(generic["body"].as_str().unwrap()) + .unwrap(), + ) + .unwrap(); + assert_eq!(html, "

Hello

\n"); + let encoded = URL_SAFE_NO_PAD.encode("

Edited in composer

"); + let composer = reviewed_arguments( + "SendEmail", + &draft, + &json!({"draft":json!({"body":encoded}).to_string(),"bodyFormat":"base64url_html"}), + ) + .unwrap(); + assert_eq!(composer["body"], encoded); + assert!( + reviewed_arguments("SendEmail", &draft, &json!({"bodyFormat":"base64url_html"})).is_err() + ); + assert!(reviewed_arguments("SendEmail", &draft, &json!({"bodyFormat":"unknown"})).is_err()); +} diff --git a/services/mcp_service/src/tool_service/test.rs b/services/mcp_service/src/tool_service/test.rs index 260c46ae3b2..53154c6fa60 100644 --- a/services/mcp_service/src/tool_service/test.rs +++ b/services/mcp_service/src/tool_service/test.rs @@ -211,3 +211,5 @@ fn authenticated_user_id_requires_user_extension_inside_request_parts() { assert_eq!(error.code, ErrorCode::INTERNAL_ERROR); assert_eq!(error.message, "missing user identity — is auth configured?"); } + +mod transport; diff --git a/services/mcp_service/src/tool_service/test/transport.rs b/services/mcp_service/src/tool_service/test/transport.rs new file mode 100644 index 00000000000..8905e2c29f8 --- /dev/null +++ b/services/mcp_service/src/tool_service/test/transport.rs @@ -0,0 +1,352 @@ +//! Real HTTP MCP round trips, including a form response arriving on a second replica. +use super::*; +use crate::session_routing::{ + directory::{Directory, Owner}, + route_sessions, +}; +use ai_toolset::{AsyncTool, ServiceContext, ToolAnnotated, ToolAnnotations, ToolResult}; +use axum::{ + body::Body, + http::{Request, StatusCode}, + response::IntoResponse, +}; +use futures::StreamExt; +use rmcp::transport::streamable_http_server::{StreamableHttpServerConfig, StreamableHttpService}; +use serde_json::{Value, json}; +use std::{collections::HashMap, sync::Mutex}; + +#[derive(Clone, Default)] +struct TestDirectory(Arc>>); +#[async_trait::async_trait] +impl Directory for TestDirectory { + async fn lookup(&self, id: &str) -> Result, String> { + Ok(self.0.lock().unwrap().get(id).cloned()) + } + async fn register(&self, id: &str, owner: &Owner) -> Result<(), String> { + self.0.lock().unwrap().insert(id.to_owned(), owner.clone()); + Ok(()) + } + async fn remove(&self, id: &str) -> Result<(), String> { + self.0.lock().unwrap().remove(id); + Ok(()) + } +} + +#[derive(Clone, Default)] +struct TestContext(Arc>>); +#[async_trait::async_trait] +impl MarkdownImageResolver for TestContext { + async fn resolve_static(&self, _: &str) -> Option { + None + } + async fn resolve_dss( + &self, + _: &MacroUserIdStr<'_>, + _: &str, + ) -> Option { + None + } +} +impl McpToolContext for TestContext { + fn for_user(&self, _: MacroUserIdStr<'static>) -> Self { + self.clone() + } +} +#[derive(serde::Deserialize, schemars::JsonSchema)] +#[schemars(title = "ReviewedNote", description = "Send a reviewed test note.")] +struct ReviewedNote { + text: String, +} +impl ToolAnnotated for ReviewedNote { + const ANNOTATIONS: ToolAnnotations = ToolAnnotations::read_only("Review note"); +} +#[async_trait::async_trait] +impl AsyncTool for ReviewedNote { + type Output = String; + async fn call( + &self, + context: ServiceContext, + _: RequestContext, + ) -> ToolResult { + context.0.0.lock().unwrap().push(self.text.clone()); + Ok(self.text.clone()) + } +} + +async fn replica( + directory: TestDirectory, + context: TestContext, + process: &str, +) -> (String, tokio::task::JoinHandle<()>) { + let (url, task, _) = replica_with_lifetime( + directory, + context, + process, + std::time::Duration::from_secs(3660), + Default::default(), + ) + .await; + (url, task) +} + +async fn replica_with_lifetime( + directory: TestDirectory, + context: TestContext, + process: &str, + idle: std::time::Duration, + shutdown: tokio_util::sync::CancellationToken, +) -> ( + String, + tokio::task::JoinHandle<()>, + Arc, +) { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let mut sessions = + rmcp::transport::streamable_http_server::session::local::LocalSessionManager::default(); + sessions.session_config.keep_alive = Some(idle); + let sessions = Arc::new(sessions); + let local = StreamableHttpService::new( + move || { + Ok(AuthenticatedToolService::new( + Arc::new(AsyncToolCollection::new().add_user_tool::()), + context.clone(), + "https://macro.com".into(), + )) + }, + sessions.clone(), + StreamableHttpServerConfig::default().with_cancellation_token(shutdown.clone()), + ); + let routed = route_sessions( + local, + directory, + process.into(), + address, + "localhost".into(), + sessions.clone(), + ); + // Replace JWT verification with two already-verified principals at the boundary. + let app = axum::Router::new() + .nest_service("/mcp", routed) + .layer(axum::middleware::from_fn( + |mut request: Request, next: axum::middleware::Next| async move { + let email = match request + .headers() + .get("authorization") + .and_then(|h| h.to_str().ok()) + { + Some("Bearer alice") => "alice@example.com", + Some("Bearer bob") => "bob@example.com", + _ => return StatusCode::UNAUTHORIZED.into_response(), + }; + request + .extensions_mut() + .insert(MacroUserIdStr::try_from_email(email).unwrap()); + next.run(request).await + }, + )); + let task = tokio::spawn(async move { + axum::serve(listener, app) + .with_graceful_shutdown(shutdown.cancelled_owned()) + .await + .unwrap(); + }); + (format!("http://{address}/mcp"), task, sessions) +} + +fn post( + client: &reqwest::Client, + url: &str, + session: Option<&str>, + message: Value, +) -> reqwest::RequestBuilder { + let request = client + .post(url) + .bearer_auth("alice") + .header("accept", "application/json, text/event-stream") + .header("mcp-protocol-version", "2025-11-25") + .json(&message); + match session { + Some(id) => request.header("mcp-session-id", id), + None => request, + } +} + +async fn initialize(client: &reqwest::Client, url: &str, form: bool) -> String { + let capabilities = if form { + json!({"elicitation":{"form":{}}}) + } else { + json!({}) + }; + let response = post(client, url, None, json!({"jsonrpc":"2.0","id":1,"method":"initialize","params":{ + "protocolVersion":"2025-11-25", "capabilities":capabilities, "clientInfo":{"name":"probe","version":"1"} + }})).send().await.unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let session = response.headers()["mcp-session-id"] + .to_str() + .unwrap() + .to_owned(); + post( + client, + url, + Some(&session), + json!({"jsonrpc":"2.0","method":"notifications/initialized"}), + ) + .send() + .await + .unwrap() + .error_for_status() + .unwrap(); + session +} + +#[tokio::test] +async fn reviewed_tools_round_trip_across_replicas_and_never_execute_on_refusal() { + tokio::time::timeout(std::time::Duration::from_secs(30), async { + let directory = TestDirectory::default(); + let context = TestContext::default(); + let (a, a_task) = replica(directory.clone(), context.clone(), "a").await; + let (b, b_task) = replica(directory.clone(), context.clone(), "b").await; + let client = reqwest::Client::new(); + let session = initialize(&client, &a, true).await; + let events = client.get(&b).bearer_auth("alice").header("mcp-session-id", &session) + .header("accept", "text/event-stream").header("mcp-protocol-version", "2025-11-25") + .send().await.unwrap().error_for_status().unwrap(); + let mut questions = events.bytes_stream(); + // Tool request, elicitation answer, and subsequent calls can land on different replicas. + for (index, answer) in [ + json!({"action":"accept","content":{"text":"edited"}}), + json!({"action":"decline"}), + json!({"action":"cancel"}), + json!({"action":"accept","content":{"draft":"invalid json"}}), + json!({"action":"accept","content":{"text":42}}), + ].into_iter().enumerate() { + let response = post(&client, &b, Some(&session), json!({"jsonrpc":"2.0","id":index+10,"method":"tools/call","params":{"name":"ReviewedNote","arguments":{"text":"original"}}})) + .send().await.unwrap().error_for_status().unwrap(); + let mut stream = response.bytes_stream(); + let mut buffer = String::new(); + let elicitation = loop { + let chunk = questions.next().await.expect("SSE remains open").unwrap(); + buffer.push_str(std::str::from_utf8(&chunk).unwrap()); + if let Some(message) = buffer.lines().filter_map(|line| line.strip_prefix("data: ")) + .filter_map(|data| serde_json::from_str::(data).ok()) + .find(|v| v["method"] == "elicitation/create") { break message; } + }; + assert_eq!(elicitation["params"]["requestedSchema"]["properties"]["draft"]["type"], "string"); + assert_eq!(elicitation["params"]["_meta"]["macro"]["userTool"]["name"], "ReviewedNote"); + assert_eq!(client.post(&b).bearer_auth("bob").header("mcp-session-id", &session).json(&json!({"jsonrpc":"2.0","id":elicitation["id"],"result":answer})).send().await.unwrap().status(), StatusCode::FORBIDDEN); + post(&client, &b, Some(&session), json!({"jsonrpc":"2.0","id":elicitation["id"],"result":answer})).send().await.unwrap().error_for_status().unwrap(); + while let Some(chunk) = stream.next().await { buffer.push_str(std::str::from_utf8(&chunk.unwrap()).unwrap()); } + let final_result = buffer.lines().filter_map(|line| line.strip_prefix("data: ")) + .filter_map(|data| serde_json::from_str::(data).ok()) + .find(|value| value["id"] == index + 10).expect("the tool result"); + let structured = &final_result["result"]["structuredContent"]; + assert!(structured.is_null() || structured.is_object()); + if index == 1 { assert_eq!(final_result["result"]["content"][0]["text"], "Rejected"); } + assert_eq!(*context.0.lock().unwrap(), vec!["edited"]); + } + let unsupported = initialize(&client, &a, false).await; + let response = post(&client, &b, Some(&unsupported), json!({"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"ReviewedNote","arguments":{"text":"must not run"}}})).send().await.unwrap().text().await.unwrap(); + assert!(response.contains("requires form elicitation"), "{response}"); + assert_eq!(*context.0.lock().unwrap(), vec!["edited"]); + // Dead owners expire the session; never run an uncertain mutation elsewhere. + directory.remove(&session).await.unwrap(); + assert_eq!(post(&client, &b, Some(&session), json!({"jsonrpc":"2.0","id":100,"method":"tools/call","params":{"name":"ReviewedNote","arguments":{"text":"no replay"}}})).send().await.unwrap().status(), StatusCode::NOT_FOUND); + a_task.abort(); b_task.abort(); + }).await.expect("elicitation round trip must complete"); +} + +#[tokio::test] +async fn expired_local_worker_is_removed_and_returns_not_found() { + let directory = TestDirectory::default(); + let shutdown = tokio_util::sync::CancellationToken::new(); + let (url, task, sessions) = replica_with_lifetime( + directory.clone(), + Default::default(), + "short-lived", + std::time::Duration::from_millis(100), + shutdown.clone(), + ) + .await; + let client = reqwest::Client::new(); + let id = initialize(&client, &url, true).await; + tokio::time::sleep(std::time::Duration::from_millis(150)).await; + let response = post( + &client, + &url, + Some(&id), + json!({"jsonrpc":"2.0","id":3,"method":"ping"}), + ) + .send() + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::NOT_FOUND); + assert!(sessions.sessions.read().await.is_empty()); + assert!(directory.lookup(&id).await.unwrap().is_none()); + shutdown.cancel(); + tokio::time::timeout(std::time::Duration::from_secs(5), task) + .await + .unwrap() + .unwrap(); +} + +#[tokio::test] +async fn shutdown_releases_a_pending_review_without_executing_it() { + tokio::time::timeout(std::time::Duration::from_secs(10), async { + let directory = TestDirectory::default(); + let context = TestContext::default(); + let shutdown = tokio_util::sync::CancellationToken::new(); + let (url, task, _) = replica_with_lifetime(directory.clone(), context.clone(), "old-process", std::time::Duration::from_secs(3660), shutdown.clone()).await; + let client = reqwest::Client::new(); + let id = initialize(&client, &url, true).await; + let mut questions = client.get(&url).bearer_auth("alice").header("mcp-session-id", &id).header("accept", "text/event-stream").send().await.unwrap().bytes_stream(); + let response = post(&client, &url, Some(&id), json!({"jsonrpc":"2.0","id":5,"method":"tools/call","params":{"name":"ReviewedNote","arguments":{"text":"never run"}}})).send().await.unwrap(); + let mut seen = String::new(); + while !seen.contains("elicitation/create") { + let chunk = questions.next().await.unwrap().unwrap(); + seen.push_str(std::str::from_utf8(&chunk).unwrap()); + } + // Retirement makes the old session unroutable before streams close. + directory.remove(&id).await.unwrap(); + shutdown.cancel(); + let _ = response.bytes().await; + drop(questions); + task.await.unwrap(); + assert!(context.0.lock().unwrap().is_empty()); + let (new_url, new_task) = replica(directory, context.clone(), "replacement").await; + assert_eq!(post(&client, &new_url, Some(&id), json!({"jsonrpc":"2.0","id":6,"method":"tools/call","params":{"name":"ReviewedNote","arguments":{"text":"no replay"}}})).send().await.unwrap().status(), StatusCode::NOT_FOUND); + assert!(context.0.lock().unwrap().is_empty()); + new_task.abort(); + }).await.unwrap(); +} + +#[tokio::test] +async fn sessionless_tool_calls_are_rejected_without_allocating_workers() { + let shutdown = tokio_util::sync::CancellationToken::new(); + let (url, task, sessions) = replica_with_lifetime( + Default::default(), + Default::default(), + "a", + std::time::Duration::from_secs(3660), + shutdown.clone(), + ) + .await; + let client = reqwest::Client::new(); + for id in 0..3 { + let response = post(&client, &url, None, json!({"jsonrpc":"2.0","id":id,"method":"tools/call","params":{"name":"ReviewedNote","arguments":{"text":"no session"}}})).send().await.unwrap(); + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + } + let response = post( + &client, + &url, + None, + json!({"jsonrpc":"2.0","id":1,"method":"initialize","result":{}}), + ) + .send() + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + assert!(sessions.sessions.read().await.is_empty()); + shutdown.cancel(); + task.await.unwrap(); +} From 9cc3103054f691c5e577424a2eef29bc470b4f0e Mon Sep 17 00:00:00 2001 From: Wolf Mermelstein Date: Wed, 9 Sep 2026 15:56:33 +0000 Subject: [PATCH 2/5] Simplify email elicitation forms for external MCP clients --- crates/agent_inmem/src/outbound/acp_mcp.rs | 6 + .../agent_inmem/src/outbound/acp_mcp/test.rs | 21 ++++ docs/ACP_ELICITATION.md | 6 + docs/AGENT_GUIDE/ai-chat.md | 3 +- services/mcp_service/src/tool_service.rs | 17 ++- .../mcp_service/src/tool_service/review.rs | 117 ++++++++++++++++++ .../src/tool_service/review/test.rs | 58 +++++++++ 7 files changed, 223 insertions(+), 5 deletions(-) diff --git a/crates/agent_inmem/src/outbound/acp_mcp.rs b/crates/agent_inmem/src/outbound/acp_mcp.rs index 3e32e686a73..9de18c695ec 100644 --- a/crates/agent_inmem/src/outbound/acp_mcp.rs +++ b/crates/agent_inmem/src/outbound/acp_mcp.rs @@ -23,6 +23,12 @@ impl rmcp::ClientHandler for ElicitationClient { fn get_info(&self) -> rmcp::model::ClientInfo { let mut info = client_info(); if self.input.is_some() { + if self.server == MACRO_MCP_NAME { + info.capabilities + .experimental + .get_or_insert_with(Default::default) + .insert("macro/composer".into(), Default::default()); + } info.capabilities.elicitation = Some(rmcp::model::ElicitationCapability { form: Some(Default::default()), url: None, diff --git a/crates/agent_inmem/src/outbound/acp_mcp/test.rs b/crates/agent_inmem/src/outbound/acp_mcp/test.rs index a219539796f..c5146c85761 100644 --- a/crates/agent_inmem/src/outbound/acp_mcp/test.rs +++ b/crates/agent_inmem/src/outbound/acp_mcp/test.rs @@ -112,3 +112,24 @@ async fn mcp_form_bridge_preserves_answers_and_only_trusts_macro_composer_metada service.cancel().await.unwrap(); } } + +#[test] +fn composer_capability_is_only_advertised_to_macro_with_user_input() { + use rmcp::ClientHandler; + for (server, has_input) in [("macro", true), ("thirdparty", true), ("macro", false)] { + let client = ElicitationClient { + server: server.into(), + input: has_input.then(|| { + std::sync::Arc::new(FormRecorder::default()) + as crate::domain::user_input::SharedUserInputRequester + }), + }; + let info = client.get_info(); + let composer = info + .capabilities + .experimental + .as_ref() + .is_some_and(|caps| caps.contains_key("macro/composer")); + assert_eq!(composer, server == "macro" && has_input); + } +} diff --git a/docs/ACP_ELICITATION.md b/docs/ACP_ELICITATION.md index 53474649d73..df727910ca4 100644 --- a/docs/ACP_ELICITATION.md +++ b/docs/ACP_ELICITATION.md @@ -364,6 +364,12 @@ returns the existing `UserToolResponse` envelope; chat still receives `PendingUserExecution` and finishes through its composer endpoint. The agent-loop finisher and neutral `UserToolReviewer` abstraction have been removed. +External email reviews show To, Subject, and Body, plus populated Cc/Bcc and explicit +reply/signature settings. Addresses are edited as comma-separated text. Inmem +advertises `experimental["macro/composer"]` to Macro MCP when it supports forms; +that opts into the full-draft fields used by the web composer. External email +forms omit the JSON draft and body encoding controls. + Forms use standard MCP primitive fields, projected from the draft with defaults. The optional string `draft` accepts the complete edited arguments as JSON, allowing nested recipients and other complex arguments without a private schema type. diff --git a/docs/AGENT_GUIDE/ai-chat.md b/docs/AGENT_GUIDE/ai-chat.md index aa546c923c3..fdaeb93d8e4 100644 --- a/docs/AGENT_GUIDE/ai-chat.md +++ b/docs/AGENT_GUIDE/ai-chat.md @@ -85,7 +85,8 @@ The new Macro agent uses Macro MCP for workspace tools. Sending email and creati calendar events opens a review composer; edit the draft and explicitly Send/Create, or decline. The accepted review and the final tool result may appear as separate rows. Questions from other MCP servers show the requesting server name and use -ordinary forms. Concurrent questions appear one at a time. If a server session +ordinary forms. External MCP clients review emails using address, subject, and +body fields; the Macro web UI keeps its email composer. Concurrent questions appear one at a time. If a server session expires during a deploy, reconnect the agent session; do not blindly repeat an action whose result is uncertain. The classic in-channel `@Macro` behavior is unchanged. diff --git a/services/mcp_service/src/tool_service.rs b/services/mcp_service/src/tool_service.rs index b5832723b03..ec3c3297bcf 100644 --- a/services/mcp_service/src/tool_service.rs +++ b/services/mcp_service/src/tool_service.rs @@ -177,10 +177,19 @@ where "This tool requires form elicitation support; nothing was executed.", )); } - let schema = review::project_form( - &review::tool_schema(&request.name, &tool.input_schema), - &arguments, - ) + let composer = context + .peer + .peer_info() + .and_then(|info| info.capabilities.experimental.as_ref()) + .is_some_and(|caps| caps.contains_key("macro/composer")); + let schema = if request.name == "SendEmail" && !composer { + review::email_form(&arguments) + } else { + review::project_form( + &review::tool_schema(&request.name, &tool.input_schema), + &arguments, + ) + } .map_err(|error| rmcp::ErrorData::internal_error(error, None))?; let params = rmcp::model::CreateElicitationRequestParams::FormElicitationParams { meta: Some(rmcp::model::Meta( diff --git a/services/mcp_service/src/tool_service/review.rs b/services/mcp_service/src/tool_service/review.rs index 96ff1883887..807bec75a40 100644 --- a/services/mcp_service/src/tool_service/review.rs +++ b/services/mcp_service/src/tool_service/review.rs @@ -50,6 +50,115 @@ pub(super) fn project_form( .map_err(|error| error.to_string()) } +/// External email clients get editable address lists and the message, without +/// the private composer's transport fields. Explicit reply/signature choices +/// remain visible so the user can review them. +pub(super) fn email_form(draft: &Value) -> Result { + let mut properties = Map::new(); + for (name, title, description) in [ + ("to", "To", "Email addresses, separated by commas."), + ("cc", "Cc", "Email addresses, separated by commas."), + ("bcc", "Bcc", "Email addresses, separated by commas."), + ] { + let recipients = draft.get(name).and_then(Value::as_array); + if name != "to" && recipients.is_none_or(Vec::is_empty) { + continue; + } + let addresses = recipients + .into_iter() + .flatten() + .filter_map(|recipient| recipient.get("email").and_then(Value::as_str)) + .collect::>() + .join(", "); + properties.insert( + name.into(), + json!({"type":"string", "title":title, + "description":description, "default":addresses}), + ); + } + for (name, title, description, kind) in [ + ("subject", "Subject", "Email subject.", "string"), + ( + "body", + "Body", + "Email message (Markdown supported).", + "string", + ), + ( + "replyingToId", + "Reply to message", + "Message ID this email replies to.", + "string", + ), + ( + "includeSignature", + "Include signature", + "Include your email signature.", + "boolean", + ), + ] { + let value = draft.get(name).filter(|value| !value.is_null()); + if value.is_none() && !["subject", "body"].contains(&name) { + continue; + } + let mut field = json!({"type":kind,"title":title,"description":description}); + if let Some(value) = value { + field["default"] = value.clone(); + } + properties.insert(name.into(), field); + } + serde_json::from_value(json!({"type":"object","properties":properties})) + .map_err(|error| error.to_string()) +} + +/// Turn plain address edits back into tool recipients. Keep display names for +/// unchanged addresses; never interpret a malformed address as the old value. +fn email_recipient_edits(draft: &Value, content: &mut Map) -> Result<(), String> { + for name in ["to", "cc", "bcc"] { + let Some(value) = content.get(name) else { + continue; + }; + let text = value + .as_str() + .ok_or("recipient fields must be comma-separated email addresses")?; + let mut recipients = Vec::new(); + for address in text + .split(',') + .map(str::trim) + .filter(|address| !address.is_empty()) + { + let valid = address.split_once('@').is_some_and(|(local, domain)| { + !local.is_empty() && !domain.is_empty() && !domain.contains('@') + }) && !address + .chars() + .any(|c| c.is_whitespace() || matches!(c, '<' | '>' | ';')); + if !valid { + return Err( + "Enter email addresses separated by commas (without display names).".into(), + ); + } + let original = draft + .get(name) + .and_then(Value::as_array) + .and_then(|values| { + values + .iter() + .find(|value| value["email"].as_str() == Some(address)) + }); + recipients.push( + original + .cloned() + .unwrap_or_else(|| json!({"email":address})), + ); + } + if name == "to" && recipients.is_empty() { + return Err("At least one To recipient is required.".into()); + } + content.insert(name.into(), Value::Array(recipients)); + } + Ok(()) +} + /// Merge edited fields, rejecting a malformed replacement instead of executing the old draft. pub(super) fn apply_review(draft: &Value, content: &Value) -> Result { let content = content @@ -182,6 +291,14 @@ pub(super) fn reviewed_arguments( .cloned() .ok_or("the accepted form must contain an object")?; let format = content.remove("bodyFormat").unwrap_or(json!("markdown")); + // A complete composer draft takes precedence over any prepopulated fields. + if content + .get(DRAFT_FIELD) + .and_then(Value::as_str) + .is_none_or(|text| text.trim().is_empty()) + { + email_recipient_edits(draft, &mut content)?; + } let mut reviewed = apply_review(draft, &Value::Object(content))?; let body = reviewed .get("body") diff --git a/services/mcp_service/src/tool_service/review/test.rs b/services/mcp_service/src/tool_service/review/test.rs index 91092d4ae08..1a0ad4b1126 100644 --- a/services/mcp_service/src/tool_service/review/test.rs +++ b/services/mcp_service/src/tool_service/review/test.rs @@ -69,3 +69,61 @@ fn email_forms_render_markdown_and_preserve_explicit_composer_html() { ); assert!(reviewed_arguments("SendEmail", &draft, &json!({"bodyFormat":"unknown"})).is_err()); } + +#[test] +fn email_form_is_readable_and_only_shows_relevant_fields() { + let draft = json!({"to":[{"email":"wolf@example.com","name":"Wolf"}], + "subject":"Frogs","body":"Hello\n\nFrogs!","cc":[],"includeSignature":null}); + let form = serde_json::to_value(email_form(&draft).unwrap()).unwrap(); + let fields = form["properties"].as_object().unwrap(); + assert_eq!( + fields.keys().map(String::as_str).collect::>(), + ["body", "subject", "to"] + ); + assert_eq!(fields["to"]["default"], "wolf@example.com"); + assert_eq!(fields["body"]["title"], "Body"); + assert_eq!(fields["body"]["default"], "Hello\n\nFrogs!"); + let form = serde_json::to_value( + email_form(&json!({ + "to":[],"bcc":[{"email":"private@example.com"}], + "includeSignature":false,"replyingToId":"message-id" + })) + .unwrap(), + ) + .unwrap(); + assert_eq!(form["properties"]["bcc"]["default"], "private@example.com"); + assert_eq!(form["properties"]["includeSignature"]["default"], false); + assert_eq!(form["properties"]["replyingToId"]["default"], "message-id"); +} + +#[test] +fn email_address_edits_preserve_names_clear_cc_and_reject_invalid_input() { + let draft = json!({"body":"Hello", "to":[{"email":"wolf@example.com","name":"Wolf"}], + "cc":[{"email":"old@example.com"}]}); + let edited = reviewed_arguments( + "SendEmail", + &draft, + &json!({ + "to":"wolf@example.com, new@example.com", "cc":"" + }), + ) + .unwrap(); + assert_eq!( + edited["to"], + json!([{"email":"wolf@example.com","name":"Wolf"},{"email":"new@example.com"}]) + ); + assert_eq!(edited["cc"], json!([])); + for invalid in [ + json!(""), + json!("not-an-address"), + json!("a@example.com\nBcc:x@example.com"), + json!(42), + ] { + assert!(reviewed_arguments("SendEmail", &draft, &json!({"to":invalid})).is_err()); + } + // Composer JSON replaces the whole draft, even when other fields were prefilled. + let edited = reviewed_arguments("SendEmail", &draft, &json!({ + "to":"not-an-address", "draft":json!({"body":"Edited","to":[{"email":"composer@example.com"}]}).to_string() + })).unwrap(); + assert_eq!(edited["to"], json!([{"email":"composer@example.com"}])); +} From b80e4a3981f43fffa1199905a267a56963fbc909 Mon Sep 17 00:00:00 2001 From: Wolf Mermelstein Date: Wed, 9 Sep 2026 16:09:42 +0000 Subject: [PATCH 3/5] Show multiline email previews in MCP elicitation messages --- docs/ACP_ELICITATION.md | 5 +- docs/AGENT_GUIDE/ai-chat.md | 4 +- services/mcp_service/src/tool_service.rs | 6 ++- .../mcp_service/src/tool_service/review.rs | 46 ++++++++++++++++--- .../src/tool_service/review/test.rs | 35 ++++++++++++-- 5 files changed, 81 insertions(+), 15 deletions(-) diff --git a/docs/ACP_ELICITATION.md b/docs/ACP_ELICITATION.md index df727910ca4..56f9ffb12ac 100644 --- a/docs/ACP_ELICITATION.md +++ b/docs/ACP_ELICITATION.md @@ -364,8 +364,9 @@ returns the existing `UserToolResponse` envelope; chat still receives `PendingUserExecution` and finishes through its composer endpoint. The agent-loop finisher and neutral `UserToolReviewer` abstraction have been removed. -External email reviews show To, Subject, and Body, plus populated Cc/Bcc and explicit -reply/signature settings. Addresses are edited as comma-separated text. Inmem +External email reviews show the full multiline message above editable To/Subject +fields and an optional Replacement body. Leaving the replacement blank keeps the +previewed body. Populated Cc/Bcc and explicit reply/signature settings remain visible. Addresses are edited as comma-separated text. Inmem advertises `experimental["macro/composer"]` to Macro MCP when it supports forms; that opts into the full-draft fields used by the web composer. External email forms omit the JSON draft and body encoding controls. diff --git a/docs/AGENT_GUIDE/ai-chat.md b/docs/AGENT_GUIDE/ai-chat.md index fdaeb93d8e4..7736192f06e 100644 --- a/docs/AGENT_GUIDE/ai-chat.md +++ b/docs/AGENT_GUIDE/ai-chat.md @@ -85,8 +85,8 @@ The new Macro agent uses Macro MCP for workspace tools. Sending email and creati calendar events opens a review composer; edit the draft and explicitly Send/Create, or decline. The accepted review and the final tool result may appear as separate rows. Questions from other MCP servers show the requesting server name and use -ordinary forms. External MCP clients review emails using address, subject, and -body fields; the Macro web UI keeps its email composer. Concurrent questions appear one at a time. If a server session +ordinary forms. External MCP clients see a multiline email preview with editable addresses, subject, +and an optional replacement body; the Macro web UI keeps its email composer. Concurrent questions appear one at a time. If a server session expires during a deploy, reconnect the agent session; do not blindly repeat an action whose result is uncertain. The classic in-channel `@Macro` behavior is unchanged. diff --git a/services/mcp_service/src/tool_service.rs b/services/mcp_service/src/tool_service.rs index ec3c3297bcf..34a44f9370e 100644 --- a/services/mcp_service/src/tool_service.rs +++ b/services/mcp_service/src/tool_service.rs @@ -198,7 +198,11 @@ where })) .map_err(|error| rmcp::ErrorData::internal_error(error.to_string(), None))?, )), - message: format!("{}?", tool.annotations.title), + message: if request.name == "SendEmail" && !composer { + review::email_message(&arguments) + } else { + format!("{}?", tool.annotations.title) + }, requested_schema: schema, }; let mut pending = context diff --git a/services/mcp_service/src/tool_service/review.rs b/services/mcp_service/src/tool_service/review.rs index 807bec75a40..0e4f2a26465 100644 --- a/services/mcp_service/src/tool_service/review.rs +++ b/services/mcp_service/src/tool_service/review.rs @@ -78,12 +78,6 @@ pub(super) fn email_form(draft: &Value) -> Result { } for (name, title, description, kind) in [ ("subject", "Subject", "Email subject.", "string"), - ( - "body", - "Body", - "Email message (Markdown supported).", - "string", - ), ( "replyingToId", "Reply to message", @@ -98,7 +92,7 @@ pub(super) fn email_form(draft: &Value) -> Result { ), ] { let value = draft.get(name).filter(|value| !value.is_null()); - if value.is_none() && !["subject", "body"].contains(&name) { + if value.is_none() && name != "subject" { continue; } let mut field = json!({"type":kind,"title":title,"description":description}); @@ -107,10 +101,40 @@ pub(super) fn email_form(draft: &Value) -> Result { } properties.insert(name.into(), field); } + properties.insert("replacementBody".into(), json!({ + "type":"string", "title":"Replacement body", + "description":"Optional. Leave blank to keep the message above; enter text to replace it (Markdown supported)." + })); serde_json::from_value(json!({"type":"object","properties":properties})) .map_err(|error| error.to_string()) } +/// Multiline previews belong in the message: terminal clients often render +/// prepopulated string inputs on a single line. +pub(super) fn email_message(draft: &Value) -> String { + let mut lines = vec!["Send this email?".to_owned(), String::new()]; + for (name, title) in [("to", "To"), ("cc", "Cc"), ("bcc", "Bcc")] { + let addresses = draft + .get(name) + .and_then(Value::as_array) + .into_iter() + .flatten() + .filter_map(|recipient| recipient.get("email").and_then(Value::as_str)) + .collect::>() + .join(", "); + if !addresses.is_empty() { + lines.push(format!("{title}: {addresses}")); + } + } + lines.push(format!( + "Subject: {}", + draft["subject"].as_str().unwrap_or_default() + )); + lines.push(String::new()); + lines.push(draft["body"].as_str().unwrap_or_default().to_owned()); + lines.join("\n") +} + /// Turn plain address edits back into tool recipients. Keep display names for /// unchanged addresses; never interpret a malformed address as the old value. fn email_recipient_edits(draft: &Value, content: &mut Map) -> Result<(), String> { @@ -291,6 +315,14 @@ pub(super) fn reviewed_arguments( .cloned() .ok_or("the accepted form must contain an object")?; let format = content.remove("bodyFormat").unwrap_or(json!("markdown")); + if let Some(replacement) = content.remove("replacementBody") { + let text = replacement + .as_str() + .ok_or("the replacement body must be text")?; + if !text.trim().is_empty() { + content.insert("body".into(), replacement); + } + } // A complete composer draft takes precedence over any prepopulated fields. if content .get(DRAFT_FIELD) diff --git a/services/mcp_service/src/tool_service/review/test.rs b/services/mcp_service/src/tool_service/review/test.rs index 1a0ad4b1126..5ab906f9b08 100644 --- a/services/mcp_service/src/tool_service/review/test.rs +++ b/services/mcp_service/src/tool_service/review/test.rs @@ -78,11 +78,15 @@ fn email_form_is_readable_and_only_shows_relevant_fields() { let fields = form["properties"].as_object().unwrap(); assert_eq!( fields.keys().map(String::as_str).collect::>(), - ["body", "subject", "to"] + ["replacementBody", "subject", "to"] ); assert_eq!(fields["to"]["default"], "wolf@example.com"); - assert_eq!(fields["body"]["title"], "Body"); - assert_eq!(fields["body"]["default"], "Hello\n\nFrogs!"); + assert!(fields["replacementBody"].get("default").is_none()); + assert!(!fields.contains_key("body")); + assert_eq!( + email_message(&draft), + "Send this email?\n\nTo: wolf@example.com\nSubject: Frogs\n\nHello\n\nFrogs!" + ); let form = serde_json::to_value( email_form(&json!({ "to":[],"bcc":[{"email":"private@example.com"}], @@ -127,3 +131,28 @@ fn email_address_edits_preserve_names_clear_cc_and_reject_invalid_input() { })).unwrap(); assert_eq!(edited["to"], json!([{"email":"composer@example.com"}])); } + +#[test] +fn optional_replacement_body_keeps_or_replaces_the_previewed_message() { + use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; + let draft = json!({"body":"Original\n\nmessage"}); + for (content, expected) in [ + (json!({}), "

Original

\n

message

\n"), + ( + json!({"replacementBody":" "}), + "

Original

\n

message

\n", + ), + ( + json!({"replacementBody":"**Edited**"}), + "

Edited

\n", + ), + ] { + let result = reviewed_arguments("SendEmail", &draft, &content).unwrap(); + let body = URL_SAFE_NO_PAD + .decode(result["body"].as_str().unwrap()) + .unwrap(); + assert_eq!(String::from_utf8(body).unwrap(), expected); + assert!(result.get("replacementBody").is_none()); + } + assert!(reviewed_arguments("SendEmail", &draft, &json!({"replacementBody":42})).is_err()); +} From abaffec495e319b4bc596da3164a5f8e33742999 Mon Sep 17 00:00:00 2001 From: Wolf Mermelstein Date: Wed, 9 Sep 2026 16:27:24 +0000 Subject: [PATCH 4/5] Request inline MCP email previews with an internal composer exemption --- crates/prompt/src/user_tools.rs | 10 ++++++++++ docs/ACP_ELICITATION.md | 5 +++++ services/mcp_service/src/tool_service.rs | 9 ++++++++- 3 files changed, 23 insertions(+), 1 deletion(-) diff --git a/crates/prompt/src/user_tools.rs b/crates/prompt/src/user_tools.rs index 9ef626da121..7ce449d4b8f 100644 --- a/crates/prompt/src/user_tools.rs +++ b/crates/prompt/src/user_tools.rs @@ -16,6 +16,11 @@ static INSTRUCTIONS: &str = r##"- User tools are tools that must be executed by A user tool will return "PendingUserExecution" until a user chooses to accept / reject the tool. +- Some MCP tools instruct you to display a draft before calling them. When using Macro's + built-in email or calendar composer, you are exempt from that display requirement: call + the tool directly and let the composer present the draft. User review and approval are + still required before execution. + - IMPORTANT: When the user asks you to draft, write, compose, or send an email (or reply to one), you MUST use the `SendEmail` tool to produce it. NEVER write the email body as plain text in the chat. The `SendEmail` tool opens a real draft in the email composer that the user can review, @@ -33,6 +38,11 @@ pub static PROMPT: StaticPrompt<'static> = StaticPrompt::borrowed(TITLE, INSTRUC static SESSION_INSTRUCTIONS: &str = r##"- `SendEmail` and `CreateCalendarEvent` are reviewed by the user before they run. Calling one opens a review card in the session, the turn waits while the user edits, confirms or declines, and the tool then returns what happened: the sent email or created event, or "Rejected". Nothing is pending afterwards and there is no chat composer; do not tell the user to confirm anything, and do not ask for confirmation in prose before calling the tool - the review card is the confirmation. +- Some MCP tools instruct you to display a draft before calling them. When using Macro's + built-in email or calendar composer, you are exempt from that display requirement: call + the tool directly and let the composer present the draft. User review and approval are + still required before execution. + - IMPORTANT: When the user asks you to draft, write, compose, or send an email (or reply to one), you MUST use the `SendEmail` tool to produce it. NEVER write the email body as plain text in your reply. The review card is a real email composer the user can edit before it sends; inline text diff --git a/docs/ACP_ELICITATION.md b/docs/ACP_ELICITATION.md index 56f9ffb12ac..6f821b9fdb6 100644 --- a/docs/ACP_ELICITATION.md +++ b/docs/ACP_ELICITATION.md @@ -364,6 +364,11 @@ returns the existing `UserToolResponse` envelope; chat still receives `PendingUserExecution` and finishes through its composer endpoint. The agent-loop finisher and neutral `UserToolReviewer` abstraction have been removed. +MCP's `SendEmail` description asks external harnesses to display recipients, +subject, and the full body before calling the tool. Macro's internal chat and +session prompts explicitly exempt the built-in composers from that display +requirement; user approval is still required before execution. + External email reviews show the full multiline message above editable To/Subject fields and an optional Replacement body. Leaving the replacement blank keeps the previewed body. Populated Cc/Bcc and explicit reply/signature settings remain visible. Addresses are edited as comma-separated text. Inmem diff --git a/services/mcp_service/src/tool_service.rs b/services/mcp_service/src/tool_service.rs index 34a44f9370e..f34ec3a84e5 100644 --- a/services/mcp_service/src/tool_service.rs +++ b/services/mcp_service/src/tool_service.rs @@ -79,7 +79,14 @@ impl AuthenticatedToolService { .map(|(key, value)| { Tool::new( key.to_owned(), - value.description.to_owned(), + if key == "SendEmail" { + format!( + "{} Before calling this tool, display the proposed recipients, subject, and full email body in your response so the user can read it. Then call this tool to request review. Displaying the draft does not authorize sending; explicit user approval through the tool's review is still required.", + value.description.replace(" — never write the email as plain text in the chat", "") + ) + } else { + value.description.to_owned() + }, Arc::new(value.input_schema.clone()), ) .with_title(value.annotations.title) From 0eed8b10667f72b1db1bc71aae79dd93529871ef Mon Sep 17 00:00:00 2001 From: Wolf Mermelstein Date: Wed, 9 Sep 2026 16:48:28 +0000 Subject: [PATCH 5/5] Apply MCP previews to all reviewed tools including calendar creation --- .../email/src/inbound/toolset/send_email.rs | 2 +- docs/ACP_ELICITATION.md | 10 ++++--- services/mcp_service/src/tool_service.rs | 6 ++-- .../src/tool_service/test/transport.rs | 29 +++++++++++++++++++ 4 files changed, 39 insertions(+), 8 deletions(-) diff --git a/crates/email/src/inbound/toolset/send_email.rs b/crates/email/src/inbound/toolset/send_email.rs index 5c4ab541055..b5462aff3de 100644 --- a/crates/email/src/inbound/toolset/send_email.rs +++ b/crates/email/src/inbound/toolset/send_email.rs @@ -40,7 +40,7 @@ impl From for ContactInfo { #[derive(Debug, Deserialize, JsonSchema, Clone)] #[schemars( title = "SendEmail", - description = "Draft, compose, and send an email. ALWAYS use this tool whenever the user asks you to draft, write, compose, or send an email (or reply to one) — never write the email as plain text in the chat. This tool opens the email draft in the composer for the user to review, edit, and confirm before it is sent, so it is the correct tool even when the user only wants a draft. To reply to an existing message, provide the replying_to_id. Write the body in Markdown — use **bold**, *italics*, lists, links, and other standard Markdown formatting. The draft composer renders the Markdown for the user to review and edit; the composer produces HTML that is sent as the actual email body." + description = "Draft, compose, and send an email. ALWAYS use this tool whenever the user asks you to draft, write, compose, or send an email (or reply to one). This tool opens the email draft in the composer for the user to review, edit, and confirm before it is sent, so it is the correct tool even when the user only wants a draft. To reply to an existing message, provide the replying_to_id. Write the body in Markdown — use **bold**, *italics*, lists, links, and other standard Markdown formatting. The draft composer renders the Markdown for the user to review and edit; the composer produces HTML that is sent as the actual email body." )] #[serde(rename_all = "camelCase")] pub struct SendEmail { diff --git a/docs/ACP_ELICITATION.md b/docs/ACP_ELICITATION.md index 6f821b9fdb6..67684d91b5b 100644 --- a/docs/ACP_ELICITATION.md +++ b/docs/ACP_ELICITATION.md @@ -364,10 +364,12 @@ returns the existing `UserToolResponse` envelope; chat still receives `PendingUserExecution` and finishes through its composer endpoint. The agent-loop finisher and neutral `UserToolReviewer` abstraction have been removed. -MCP's `SendEmail` description asks external harnesses to display recipients, -subject, and the full body before calling the tool. Macro's internal chat and -session prompts explicitly exempt the built-in composers from that display -requirement; user approval is still required before execution. +MCP appends a preview instruction to every tool registered for user review, +currently `SendEmail` and `CreateCalendarEvent`. External harnesses display the +proposed content and key details before calling the tool. Macro's internal chat +and session prompts exempt the built-in composers from that display requirement; +user approval is still required before execution. Shared tool descriptions do +not impose the host's inline-preview behavior. External email reviews show the full multiline message above editable To/Subject fields and an optional Replacement body. Leaving the replacement blank keeps the diff --git a/services/mcp_service/src/tool_service.rs b/services/mcp_service/src/tool_service.rs index f34ec3a84e5..bca56f47d59 100644 --- a/services/mcp_service/src/tool_service.rs +++ b/services/mcp_service/src/tool_service.rs @@ -79,10 +79,10 @@ impl AuthenticatedToolService { .map(|(key, value)| { Tool::new( key.to_owned(), - if key == "SendEmail" { + if self.toolset.user_tools.contains_key(key) { format!( - "{} Before calling this tool, display the proposed recipients, subject, and full email body in your response so the user can read it. Then call this tool to request review. Displaying the draft does not authorize sending; explicit user approval through the tool's review is still required.", - value.description.replace(" — never write the email as plain text in the chat", "") + "{} Before calling this tool, display the proposed content in full and the key details in your response so the user can review them. For an email, include recipients, subject, and the full body. For a calendar event, include title, date and time with timezone, attendees, location, and description. Then call the tool to request approval. Displaying a preview does not authorize execution; explicit user approval through the tool's review is still required.", + value.description ) } else { value.description.to_owned() diff --git a/services/mcp_service/src/tool_service/test/transport.rs b/services/mcp_service/src/tool_service/test/transport.rs index 8905e2c29f8..e10e0812860 100644 --- a/services/mcp_service/src/tool_service/test/transport.rs +++ b/services/mcp_service/src/tool_service/test/transport.rs @@ -350,3 +350,32 @@ async fn sessionless_tool_calls_are_rejected_without_allocating_workers() { shutdown.cancel(); task.await.unwrap(); } + +#[test] +fn preview_instructions_follow_review_registration_not_tool_names() { + for reviewed in [false, true] { + let toolset = if reviewed { + AsyncToolCollection::new().add_user_tool::() + } else { + AsyncToolCollection::new().add_tool::() + }; + let original = toolset.tools["ReviewedNote"].description.clone(); + let service = AuthenticatedToolService::new( + Arc::new(toolset), + TestContext::default(), + "https://macro.com".into(), + ); + let tools = service.tool_definitions(); + let description = tools[0].description.as_deref().unwrap(); + assert_eq!( + description.contains("Before calling this tool, display"), + reviewed + ); + if reviewed { + assert!(description.contains("explicit user approval")); + } else { + assert_eq!(description, original); + } + assert_eq!(service.toolset.tools["ReviewedNote"].description, original); + } +}