diff --git a/md/migration-v5.md b/md/migration-v5.md index a6702e37..97020206 100644 --- a/md/migration-v5.md +++ b/md/migration-v5.md @@ -41,7 +41,7 @@ versions (`AgentProtocolRouter`, `ProxyProtocolRouter`, | `SessionSetup` | a session and the response of the request that set it up | a session and the response of the request that set it up | | `ActiveSession` | the session handle with its update stream | the command-only session handle | | `ForkSessionBuilder` | - | `fork_session*` (`unstable_session_fork`) | -| `LoadSessionBuilder`, `RestoreSessionBuilder`, `SessionMessage`, `Blocking`, `NonBlocking`, `SessionBlockState` | v1 session helpers | - | +| `LoadSessionBuilder`, `SessionMessage`, `Blocking`, `NonBlocking`, `SessionBlockState` | v1 session helpers | - | | `InitializeOutcome`, `InitializationVersionMismatch`, `SelectedInitializeResponse` | - | the result of `send_initialize` | | `schema` | `schema::v1` | `schema::v2` | @@ -110,8 +110,21 @@ Every setup method returns one, in both versions: v1 `start_session` for and `session/resume`; v2 `session/new`, `session/resume` and `session/fork`. The v1 `run_until` and `on_session_start` callbacks receive one too, as do the v2 `on_proxy_session_start` callbacks. The v1 -`on_proxy_session_start` callback still receives the `SessionId`, because the -session is proxied rather than handed to the caller. +`on_proxy_session_start` callback receives a clone of the complete +`NewSessionResponse` forwarded to the client, and `start_session_proxy` +returns it. Both used to give only the `SessionId`. There is no +`SessionSetup` here, because the session is proxied rather than handed to the +caller. + +```rust,ignore +// Before +.on_proxy_session_start(responder, async |session_id| { ... }) +let session_id = builder.block_task().start_session_proxy(responder).await?; + +// After +.on_proxy_session_start(responder, async |response| { let session_id = response.session_id; ... }) +let session_id = builder.block_task().start_session_proxy(responder).await?.session_id; +``` The accessors are removed in favor of the fields: `session()`, `session_mut()`, `response()`, `into_parts()` and `into_session()`. The v1 @@ -135,6 +148,49 @@ let SessionSetup { session, response, .. } = restored; let session = opened.session; ``` +### The v1 session stream reports prompt responses and errors + +`v1::ActiveSession::send_prompt` used to close the whole connection, and every +session on it, when the agent answered `session/prompt` with an error. The +caller never saw the error. The response, or the error, now arrives on the +session's update stream, and `v1::SessionMessage` changed to carry it: + +- `SessionMessage::StopReason(StopReason)` is now + `SessionMessage::PromptResponse(Result)`. It keeps + the full `PromptResponse`, including `usage` and `_meta`, which used to be + dropped. +- `SessionMessage::SessionMessage(Dispatch)` is now + `SessionMessage::Message(v1::Dispatch)`, whose request handles are typed + for v1. Call `.into_raw()` before handing it to `MatchDispatch`. +- `read_to_string` returns the prompt's error instead of the connection + closing. +- `send_prompt_blocks(Vec)` sends prompts that are not plain + text. + +`v1::SessionBlockState` is sealed. `v1::Blocking` and `v1::NonBlocking` are +its only implementations. + +```rust,ignore +// Before +match session.read_update().await? { + SessionMessage::SessionMessage(dispatch) => MatchDispatch::new(dispatch) + /* .if_notification(..) */, + SessionMessage::StopReason(stop_reason) => { /* .. */ } + _ => {} +} + +// After +match session.read_update().await? { + SessionMessage::Message(dispatch) => MatchDispatch::new(dispatch.into_raw()) + /* .if_notification(..) */, + SessionMessage::PromptResponse(response) => { + let stop_reason = response?.stop_reason; + /* .. */ + } + _ => {} +} +``` + ### Renamed and moved items Every path is relative to `agent_client_protocol`. Items not listed keep @@ -162,8 +218,10 @@ their path. | `SessionBuilder` | `v1::SessionBuilder` | | `ActiveSession` | `v1::ActiveSession` | | `RestoredSession` | `v1::SessionSetup` | -| `RestoreSessionBuilder` | `v1::RestoreSessionBuilder` (also `v1::LoadSessionBuilder` / `v1::ResumeSessionBuilder` for each request) | +| `RestoreSessionBuilder` | `v1::LoadSessionBuilder` / `v1::ResumeSessionBuilder` (the generic builder is no longer exported) | | `SessionMessage` | `v1::SessionMessage` | +| `SessionMessage::SessionMessage(Dispatch)` | `v1::SessionMessage::Message(v1::Dispatch)` | +| `SessionMessage::StopReason(StopReason)` | `v1::SessionMessage::PromptResponse(Result)` | | `Blocking`, `NonBlocking`, `SessionBlockState` | `v1::Blocking`, `v1::NonBlocking`, `v1::SessionBlockState` | | `schema::InitializeProxyRequest` | `schema::v1::InitializeProxyRequest` | | `OpenedSession` (`v1` and `v2`) | `v1::SessionSetup`, `v2::SessionSetup` (public `session` and `response` fields) | diff --git a/src/agent-client-protocol-conductor/tests/request_cancellation.rs b/src/agent-client-protocol-conductor/tests/request_cancellation.rs index cfe0062f..f1086ec3 100644 --- a/src/agent-client-protocol-conductor/tests/request_cancellation.rs +++ b/src/agent-client-protocol-conductor/tests/request_cancellation.rs @@ -849,7 +849,7 @@ async fn proxy_session_helper_cancellation_propagates_to_agent() -> Result<(), E responder: Responder, cx: ConnectionTo| { cx.build_session_from(request) - .on_proxy_session_start(responder, async |_session_id| Ok::<(), Error>(())) + .on_proxy_session_start(responder, async |_response| Ok::<(), Error>(())) }, agent_client_protocol::on_receive_request!(), ); @@ -1056,7 +1056,7 @@ async fn proxy_session_helper_cleans_up_mcp_handlers_after_cancelled_session() - ); cx.build_session_from(request) .with_mcp_server(mcp_server)? - .on_proxy_session_start(responder, async |_session_id| Ok::<(), Error>(())) + .on_proxy_session_start(responder, async |_response| Ok::<(), Error>(())) }, agent_client_protocol::on_receive_request!(), ); diff --git a/src/agent-client-protocol-conductor/tests/v1_chain_equivalence.rs b/src/agent-client-protocol-conductor/tests/v1_chain_equivalence.rs index 8d94ff49..e60e18b7 100644 --- a/src/agent-client-protocol-conductor/tests/v1_chain_equivalence.rs +++ b/src/agent-client-protocol-conductor/tests/v1_chain_equivalence.rs @@ -11,12 +11,15 @@ use std::{ }; use agent_client_protocol::V1; +use agent_client_protocol::mcp_server::McpServer; use agent_client_protocol::v1::ConnectionTo; use agent_client_protocol::{ - Agent, Client, ConnectTo, DynConnectTo, Error, Lines, Proxy, Responder, UntypedMessage, + Agent, Client, Conductor, ConnectTo, DynConnectTo, Error, Lines, Proxy, Responder, + UntypedMessage, schema::{ProtocolVersion, v1}, }; use agent_client_protocol_conductor::{ConductorImpl, SelectedInitializeRequest}; +use agent_client_protocol_rmcp::rmcp_3::McpServerExt as _; use futures::{SinkExt as _, StreamExt as _, channel::mpsc}; use serde_json::{Value, json}; @@ -78,6 +81,21 @@ struct ChainRun { /// Initialize a v1 chain with `initialize`, then create a session. async fn run_v1_chain(initialize: &Value) -> ChainRun { + run_v1_chain_with( + initialize, + json!({ "cwd": "/v1-equivalence", "mcpServers": [], "_unknownRequestField": true }), + || DynConnectTo::new(Proxy.v1()), + ) + .await +} + +/// Initialize a chain of `proxy` and a v1 agent with `initialize`, then send +/// `session/new` with `session_params`. +async fn run_v1_chain_with( + initialize: &Value, + session_params: Value, + proxy: fn() -> DynConnectTo, +) -> ChainRun { let received = Received::default(); let instantiated_version = Arc::new(Mutex::new(None)); let conductor = ConductorImpl::new_agent("v1-equivalence", { @@ -86,7 +104,7 @@ async fn run_v1_chain(initialize: &Value) -> ChainRun { move |request: SelectedInitializeRequest| async move { *instantiated_version.lock().expect("version lock") = Some(request.protocol_version()); let request = request.into_v1()?; - let proxies = vec![DynConnectTo::new(Proxy.v1())]; + let proxies = vec![proxy()]; let agent = DynConnectTo::new(recording_v1_agent(received)); Ok((request.into(), proxies, agent)) } @@ -122,7 +140,7 @@ async fn run_v1_chain(initialize: &Value) -> ChainRun { "jsonrpc": "2.0", "id": 2, "method": "session/new", - "params": { "cwd": "/v1-equivalence", "mcpServers": [], "_unknownRequestField": true }, + "params": session_params, })); let session_response = recv().await; @@ -239,3 +257,45 @@ async fn v1_chain_negotiates_newer_requests_to_v1() { assert!(received.get("_unknownRequestField").is_none(), "{received}"); assert_session_forwarded_unchanged(&run); } + +/// A v1 proxy built with `with_mcp_server` appends its declaration to the raw +/// `session/new` params. Fields the v1 schema types do not model reach the +/// agent unchanged. +#[tokio::test(flavor = "current_thread")] +async fn v1_mcp_proxy_keeps_unmodeled_session_fields() { + let run = run_v1_chain_with( + &initialize_params(1), + json!({ + "cwd": "/v1-equivalence", + "mcpServers": [], + "futureField": { "nested": ["keep me"] }, + "_unknownRequestField": true, + }), + || { + DynConnectTo::new( + Proxy + .v1() + .with_mcp_server(McpServer::builder("raw-fields").build()), + ) + }, + ) + .await; + + assert_eq!(run.responses[1]["id"], 2, "{}", run.responses[1]); + assert_eq!( + run.responses[1]["result"]["sessionId"], "v1-session", + "{}", + run.responses[1] + ); + let (method, params) = &run.received[1]; + assert_eq!(method, "session/new"); + assert_eq!( + params["futureField"], + json!({ "nested": ["keep me"] }), + "{params}" + ); + assert_eq!(params["_unknownRequestField"], true, "{params}"); + let servers = params["mcpServers"].as_array().expect("mcpServers array"); + assert_eq!(servers.len(), 1, "{params}"); + assert_eq!(servers[0]["name"], "raw-fields", "{params}"); +} diff --git a/src/agent-client-protocol-cookbook/src/lib.rs b/src/agent-client-protocol-cookbook/src/lib.rs index ced77392..a11c442b 100644 --- a/src/agent-client-protocol-cookbook/src/lib.rs +++ b/src/agent-client-protocol-cookbook/src/lib.rs @@ -346,7 +346,9 @@ pub mod connecting_as_client { //! //! - [`send_prompt`] - Send a text prompt to the agent //! - [`read_update`] - Read the next update (text chunk, tool call, etc.) - //! - [`read_to_string`] - Read all text until the turn ends + //! or the prompt's response + //! - [`read_to_string`] - Read all text until the turn ends. If the agent + //! answers the prompt with an error, this returns that error. //! //! With the core SDK's `unstable_mcp_over_acp` feature, the session builder //! also supports adding MCP servers with [`with_mcp_server`]. @@ -826,14 +828,14 @@ pub mod per_session_mcp_server { //! // Build the session and run code after it starts //! connection.build_session_from(request) //! .with_mcp_server(mcp_server)? - //! .on_proxy_session_start(responder, async move |session_id| { + //! .on_proxy_session_start(responder, async move |response| { //! // Session proxying is installed before this callback is spawned. //! // //! // Use this for follow-up work that may wait for later connection //! // traffic. Register ID-independent state in the request handler. //! // For ID-keyed state, preinstall a gate that later handlers await, //! // then populate it here. - //! tracing::info!(%session_id, "Session started"); + //! tracing::info!(session_id = %response.session_id, "Session started"); //! Ok(()) //! }) //! }, agent_client_protocol::on_receive_request!()) @@ -849,7 +851,7 @@ pub mod per_session_mcp_server { //! 1. Sends `NewSessionRequest` to the agent //! 2. When the response arrives, responds to the client automatically //! 3. Sets up message proxying for the session - //! 4. Runs your callback with the `SessionId` + //! 4. Runs your callback with a clone of the `NewSessionResponse` //! //! The callback runs after the session is established but doesn't block //! the message handler. It is suitable for follow-up work and eventual @@ -866,8 +868,8 @@ pub mod per_session_mcp_server { //! `v2::ResumeSessionBuilder` handles `session/resume`. With //! `unstable_session_fork`, `v2::ForkSessionBuilder` handles `session/fork`. //! Their `on_proxy_session_start` callbacks receive a `v2::SessionSetup`, - //! not just a session ID, so they retain both the command-only handle and - //! the complete operation-specific response: + //! which holds the command-only handle as well as the complete + //! operation-specific response: //! //! ```rust,ignore //! use agent_client_protocol::schema::v2; @@ -953,13 +955,13 @@ pub mod per_session_mcp_server { //! //! let task_connection = connection.clone(); //! connection.spawn(async move { - //! let session_id = task_connection.build_session_from(request) + //! let response = task_connection.build_session_from(request) //! .with_mcp_server(mcp_server)? //! .block_task() //! .start_session_proxy(responder) //! .await?; //! - //! tracing::info!(%session_id, "Session started"); + //! tracing::info!(session_id = %response.session_id, "Session started"); //! Ok(()) //! })?; //! Ok(()) diff --git a/src/agent-client-protocol/src/concepts/proxies.rs b/src/agent-client-protocol/src/concepts/proxies.rs index c5962658..ee1953c1 100644 --- a/src/agent-client-protocol/src/concepts/proxies.rs +++ b/src/agent-client-protocol/src/concepts/proxies.rs @@ -119,8 +119,8 @@ //! //! The v1 builder injects the declaration into new, load, resume, and //! feature-gated fork requests. The v2 builder injects it into new, resume, -//! and feature-gated fork requests while preserving unrelated setup fields. -//! Both reuse one connection-scoped server ID. +//! and feature-gated fork requests. Both preserve unrelated setup fields and +//! reuse one connection-scoped server ID. //! //! ## Per-Session MCP Server //! @@ -135,8 +135,8 @@ //! let my_mcp_server = McpServer::::builder("tools").build(); //! cx.build_session_from(req) //! .with_mcp_server(my_mcp_server)? -//! .on_proxy_session_start(responder, async |session_id| { -//! // Session started with MCP server attached +//! .on_proxy_session_start(responder, async |response| { +//! // Session `response.session_id` started with MCP server attached //! Ok(()) //! }) //! }, agent_client_protocol::on_receive_request!()) diff --git a/src/agent-client-protocol/src/concepts/sessions.rs b/src/agent-client-protocol/src/concepts/sessions.rs index c0e0e44e..b8544adb 100644 --- a/src/agent-client-protocol/src/concepts/sessions.rs +++ b/src/agent-client-protocol/src/concepts/sessions.rs @@ -129,6 +129,13 @@ //! # } //! ``` //! +//! The `session/prompt` response arrives on the session's update stream as +//! [`SessionMessage::PromptResponse`](crate::v1::SessionMessage::PromptResponse), +//! after the updates the agent sent for that turn. An error response arrives +//! there too, and `read_to_string` returns it. It does not close the +//! connection. Use `read_update` to see each +//! [`SessionMessage`](crate::v1::SessionMessage), including the stop reason. +//! //! # Adding MCP Servers //! //! You can attach MCP (Model Context Protocol) servers to a session to provide diff --git a/src/agent-client-protocol/src/jsonrpc.rs b/src/agent-client-protocol/src/jsonrpc.rs index 755cf005..9a5a487d 100644 --- a/src/agent-client-protocol/src/jsonrpc.rs +++ b/src/agent-client-protocol/src/jsonrpc.rs @@ -2050,6 +2050,7 @@ impl< /// /// The same native MCP server declaration is added to new, load, and resume /// requests, plus fork requests when `unstable_session_fork` is enabled. + /// Unrelated request fields are preserved exactly. /// /// Only applicable to proxies. Use the same method on `v2::Builder` to /// attach the server to protocol v2 setup requests. diff --git a/src/agent-client-protocol/src/mcp_server/server.rs b/src/agent-client-protocol/src/mcp_server/server.rs index 2fee3be4..7719e9cb 100644 --- a/src/agent-client-protocol/src/mcp_server/server.rs +++ b/src/agent-client-protocol/src/mcp_server/server.rs @@ -28,7 +28,7 @@ use crate::{ util::MatchDispatchFrom, }; -#[cfg(all(feature = "unstable_mcp_over_acp", feature = "unstable_protocol_v2"))] +#[cfg(feature = "unstable_mcp_over_acp")] use crate::{JsonRpcMessage, UntypedMessage}; #[cfg(feature = "unstable_mcp_over_acp")] @@ -37,6 +37,33 @@ use crate::role::HasPeer; #[cfg(all(feature = "unstable_mcp_over_acp", feature = "unstable_session_fork"))] use crate::schema::v1::ForkSessionRequest; +/// Append an MCP server declaration to the `mcpServers` array of a raw +/// session setup request, leaving every other field as received. +#[cfg(feature = "unstable_mcp_over_acp")] +fn append_raw_declaration( + request: &mut UntypedMessage, + declaration: &impl serde::Serialize, +) -> Result<(), crate::Error> { + let serde_json::Value::Object(params) = &mut request.params else { + return Err( + crate::Error::invalid_params().data("session setup parameters must be an object") + ); + }; + let declaration = + serde_json::to_value(declaration).map_err(crate::Error::into_internal_error)?; + match params.get_mut("mcpServers") { + Some(serde_json::Value::Array(servers)) => servers.push(declaration), + Some(value) => *value = serde_json::Value::Array(vec![declaration]), + None => { + params.insert( + "mcpServers".to_string(), + serde_json::Value::Array(vec![declaration]), + ); + } + } + Ok(()) +} + /// A runtime-agnostic MCP server. /// /// `McpServer` wraps an [`McpServerConnect`](`super::McpServerConnect`) @@ -239,12 +266,41 @@ where } } - /// Append this MCP server's native ACP declaration to a session setup request. - fn append_declaration(&self, mcp_servers: &mut Vec) { - mcp_servers.push(SchemaMcpServer::Acp(McpServerAcp::new( + fn declaration(&self) -> SchemaMcpServer { + SchemaMcpServer::Acp(McpServerAcp::new( self.connect.name(), self.server_id.clone(), - ))); + )) + } + + /// Append this MCP server's native ACP declaration to a session setup request. + fn append_declaration(&self, mcp_servers: &mut Vec) { + mcp_servers.push(self.declaration()); + } + + /// Whether `request` is a valid v1 session setup request this server + /// attaches to. A setup request with malformed parameters is an error. + fn validate_session_setup(request: &UntypedMessage) -> Result { + match request.method() { + "session/new" => { + NewSessionRequest::parse_message(request.method(), request.params())?; + Ok(true) + } + "session/load" => { + LoadSessionRequest::parse_message(request.method(), request.params())?; + Ok(true) + } + "session/resume" => { + ResumeSessionRequest::parse_message(request.method(), request.params())?; + Ok(true) + } + #[cfg(feature = "unstable_session_fork")] + "session/fork" => { + ForkSessionRequest::parse_message(request.method(), request.params())?; + Ok(true) + } + _ => Ok(false), + } } } @@ -346,19 +402,7 @@ where ))); } } - let declaration = - serde_json::to_value(self.declaration()).map_err(crate::Error::into_internal_error)?; - match params.get_mut("mcpServers") { - Some(serde_json::Value::Array(servers)) => servers.push(declaration), - Some(value) => *value = serde_json::Value::Array(vec![declaration]), - None => { - params.insert( - "mcpServers".to_string(), - serde_json::Value::Array(vec![declaration]), - ); - } - } - Ok(()) + append_raw_declaration(request, &self.declaration()) } /// Attach this server to a draft protocol v2 session setup request. @@ -486,53 +530,21 @@ where message: Dispatch, cx: ConnectionTo, ) -> Result, crate::Error> { - let matcher = MatchDispatchFrom::new(message, &cx) - .if_request_from(Client, async |mut request: NewSessionRequest, responder| { - self.append_declaration(&mut request.mcp_servers); + // Edit the raw request so fields the v1 schema types do not model + // reach the agent unchanged. + MatchDispatchFrom::new(message, &cx) + .if_request_from(Client, async |mut request: UntypedMessage, responder| { + if Self::validate_session_setup(&request)? { + append_raw_declaration(&mut request, &self.declaration())?; + } Ok(Handled::No { message: (request, responder), retry: false, }) }) .await - .if_request_from( - Client, - async |mut request: LoadSessionRequest, responder| { - self.append_declaration(&mut request.mcp_servers); - Ok(Handled::No { - message: (request, responder), - retry: false, - }) - }, - ) + .otherwise_delegate(&mut self.active_session) .await - .if_request_from( - Client, - async |mut request: ResumeSessionRequest, responder| { - self.append_declaration(&mut request.mcp_servers); - Ok(Handled::No { - message: (request, responder), - retry: false, - }) - }, - ) - .await; - - #[cfg(feature = "unstable_session_fork")] - let matcher = matcher - .if_request_from( - Client, - async |mut request: ForkSessionRequest, responder| { - self.append_declaration(&mut request.mcp_servers); - Ok(Handled::No { - message: (request, responder), - retry: false, - }) - }, - ) - .await; - - matcher.otherwise_delegate(&mut self.active_session).await } fn describe_chain(&self) -> impl std::fmt::Debug { @@ -844,3 +856,115 @@ mod tests { Ok(()) } } + +#[cfg(all(test, feature = "unstable_mcp_over_acp"))] +mod v1_tests { + use std::{path::PathBuf, sync::Arc}; + + use serde::Serialize; + use serde_json::{Value, json}; + + use super::{McpSessionHandler, append_raw_declaration}; + use crate::{ + Conductor, DynConnectTo, Error, UntypedMessage, + mcp_server::{McpConnectionTo, McpServerConnect}, + role, + schema::v1::{ + LoadSessionRequest, McpServer, McpServerAcpId, McpServerHttp, NewSessionRequest, + ResumeSessionRequest, SessionId, + }, + }; + + struct UnusedMcpConnect; + + impl McpServerConnect for UnusedMcpConnect { + fn name(&self) -> String { + "global-v1-server".to_owned() + } + + fn connect(&self, _context: McpConnectionTo) -> DynConnectTo { + panic!("declaration tests must not connect to the MCP server") + } + } + + fn handler() -> McpSessionHandler { + McpSessionHandler::new_with_service( + McpServerAcpId::new("global-v1-server-id"), + Arc::new(UnusedMcpConnect), + None, + ) + } + + fn assert_raw_append_preserves_params( + handler: &McpSessionHandler, + method: &str, + params: impl Serialize, + ) -> Result<(), Error> { + let mut params = serde_json::to_value(params)?; + let Value::Object(params_object) = &mut params else { + panic!("session setup params should serialize as an object"); + }; + // Fields the v1 schema types do not model must reach the agent. + params_object.insert("futureField".to_owned(), json!("kept")); + params_object.insert( + "_futureSessionField".to_owned(), + json!({ "must": ["remain", "untouched"] }), + ); + + let mut expected = params.clone(); + expected + .get_mut("mcpServers") + .and_then(Value::as_array_mut) + .expect("test request should contain mcpServers") + .push(serde_json::to_value(handler.declaration())?); + + let mut request = UntypedMessage::new(method, params)?; + assert!(McpSessionHandler::::validate_session_setup( + &request + )?); + append_raw_declaration(&mut request, &handler.declaration())?; + + assert_eq!( + request.params, expected, + "global attachment must only append its declaration" + ); + Ok(()) + } + + #[test] + fn v1_global_mcp_declaration_preserves_all_session_setup_params() -> Result<(), Error> { + let handler = handler(); + let cwd = PathBuf::from("/tmp/global-v1-mcp"); + let session_id = SessionId::new("session-to-restore"); + let existing = vec![McpServer::Http(McpServerHttp::new( + "client-server", + "https://example.com/mcp", + ))]; + + assert_raw_append_preserves_params( + &handler, + "session/new", + NewSessionRequest::new(cwd.clone()).mcp_servers(existing.clone()), + )?; + assert_raw_append_preserves_params( + &handler, + "session/load", + LoadSessionRequest::new(session_id.clone(), cwd.clone()).mcp_servers(existing.clone()), + )?; + assert_raw_append_preserves_params( + &handler, + "session/resume", + ResumeSessionRequest::new(session_id, cwd).mcp_servers(existing), + )?; + Ok(()) + } + + #[test] + fn v1_global_mcp_declaration_ignores_other_requests() -> Result<(), Error> { + let request = UntypedMessage::new("session/prompt", json!({ "sessionId": "s" }))?; + assert!(!McpSessionHandler::::validate_session_setup( + &request + )?); + Ok(()) + } +} diff --git a/src/agent-client-protocol/src/session.rs b/src/agent-client-protocol/src/session.rs index dca26861..0e539075 100644 --- a/src/agent-client-protocol/src/session.rs +++ b/src/agent-client-protocol/src/session.rs @@ -14,7 +14,7 @@ use crate::{ schema::v1::{ ContentBlock, ContentChunk, LoadSessionRequest, LoadSessionResponse, NewSessionRequest, NewSessionResponse, PromptRequest, PromptResponse, ResumeSessionRequest, - ResumeSessionResponse, SessionId, SessionNotification, SessionUpdate, StopReason, + ResumeSessionResponse, SessionId, SessionNotification, SessionUpdate, }, util::{MatchDispatch, MatchDispatchFrom}, v1::ConnectionTo, @@ -131,16 +131,25 @@ async fn run_session_scope( /// Marker type indicating the session builder will block the current task. #[derive(Debug)] pub struct Blocking; +impl block_state::Sealed for Blocking {} impl SessionBlockState for Blocking {} /// Marker type indicating the session builder will not block the current task. #[derive(Debug)] pub struct NonBlocking; +impl block_state::Sealed for NonBlocking {} impl SessionBlockState for NonBlocking {} +mod block_state { + pub trait Sealed {} +} + /// Trait for marker types that indicate blocking vs non-blocking API. /// See [`SessionBuilder::block_task`]. -pub trait SessionBlockState: Send + 'static + Sync + std::fmt::Debug {} +/// +/// This trait is sealed: [`Blocking`] and [`NonBlocking`] are its only +/// implementations. +pub trait SessionBlockState: block_state::Sealed + Send + 'static + Sync + std::fmt::Debug {} impl ConnectionTo where @@ -327,6 +336,9 @@ impl RestoreRequest for ResumeSessionRequest { /// Stable protocol v1 builder for `session/load` or `session/resume`. /// +/// This type is not exported. Callers name it through the +/// [`LoadSessionBuilder`] and [`ResumeSessionBuilder`] aliases. +/// /// Use [`ConnectionTo::load_session`] or [`ConnectionTo::resume_session`] to /// construct this builder. Use the matching `_from` method to send an existing /// typed request without rebuilding it. @@ -354,12 +366,34 @@ where } /// Stable protocol v1 builder for a `session/load` request, returned by -/// [`ConnectionTo::load_session`]. +/// [`ConnectionTo::load_session`] and [`ConnectionTo::load_session_from`]. +/// +/// The `BlockState` parameter mirrors [`SessionBuilder`]: +/// - [`NonBlocking`] exposes `on_session_start`. +/// - [`Blocking`], selected with `block_task`, exposes `start_session`. +/// +/// Session routing is installed before the request can reach the peer, so +/// replay notifications sent before the response are kept on the restored +/// [`ActiveSession`]. +/// +/// Unlike [`SessionBuilder`], this builder has no `with_mcp_server` and no +/// proxy helpers (`on_proxy_session_start`, `start_session_proxy`). To add an +/// MCP server to load requests a proxy forwards, use `with_mcp_server` on the +/// proxy's [`v1::Builder`](crate::v1::Builder). pub type LoadSessionBuilder = RestoreSessionBuilder; /// Stable protocol v1 builder for a `session/resume` request, returned by -/// [`ConnectionTo::resume_session`]. +/// [`ConnectionTo::resume_session`] and [`ConnectionTo::resume_session_from`]. +/// +/// The `BlockState` parameter mirrors [`SessionBuilder`]: +/// - [`NonBlocking`] exposes `on_session_start`. +/// - [`Blocking`], selected with `block_task`, exposes `start_session`. +/// +/// Unlike [`SessionBuilder`], this builder has no `with_mcp_server` and no +/// proxy helpers (`on_proxy_session_start`, `start_session_proxy`). To add an +/// MCP server to resume requests a proxy forwards, use `with_mcp_server` on +/// the proxy's [`v1::Builder`](crate::v1::Builder). pub type ResumeSessionBuilder = RestoreSessionBuilder; @@ -521,7 +555,7 @@ where /// Publish `session/load`, wait on the current task, and return an /// [`ActiveSession`] together with the exact [`LoadSessionResponse`]. /// - /// Requires [`block_task`](RestoreSessionBuilder::block_task). Dropping + /// Requires [`block_task`](LoadSessionBuilder::block_task). Dropping /// this future while it is pending cancels the request and removes the /// provisional session route. pub async fn start_session( @@ -538,7 +572,7 @@ where /// Publish `session/resume`, wait on the current task, and return an /// [`ActiveSession`] together with the exact [`ResumeSessionResponse`]. /// - /// Requires [`block_task`](RestoreSessionBuilder::block_task). Dropping + /// Requires [`block_task`](ResumeSessionBuilder::block_task). Dropping /// this future while it is pending cancels the request and removes the /// provisional session route. pub async fn start_session( @@ -552,8 +586,8 @@ where /// up. /// /// Every setup method returns one: [`SessionBuilder::start_session`] -/// (`session/new`) and [`RestoreSessionBuilder::start_session`] -/// (`session/load` and `session/resume`), and the `on_session_start` and +/// (`session/new`), and `start_session` on [`LoadSessionBuilder`] +/// (`session/load`) and [`ResumeSessionBuilder`] (`session/resume`). The `on_session_start` and /// [`run_until`](SessionBuilder::run_until) callbacks receive one too. The /// response is the complete, operation-specific setup response; it is kept /// next to the session rather than inside [`ActiveSession`], so the setup @@ -743,7 +777,7 @@ where }) } - /// Spawn a proxy session and run a closure with the session ID. + /// Spawn a proxy session and run a closure with the `session/new` response. /// /// A **proxy session** starts the session with the agent and then automatically /// proxies all session updates (prompts, tool calls, etc.) from the agent back @@ -756,7 +790,9 @@ where /// immediately without blocking the current task. The session handshake, client /// response, and proxy setup all happen in a spawned background task. /// - /// The closure receives the `SessionId` once the session is established. Use it for logging + /// The closure receives a clone of the complete [`NewSessionResponse`], which was also + /// forwarded to the client, once the session is established. Read the session ID from its + /// `session_id` field. Use it for logging /// or eventual tracking; it runs concurrently with later connection traffic. Register /// ID-independent state that later handlers must observe before calling this helper. For /// ID-keyed bookkeeping, install a gate or placeholder first, make later handlers await it, @@ -775,7 +811,7 @@ where /// let mcp = McpServer::::builder("tools").build(); /// cx.build_session_from(request) /// .with_mcp_server(mcp)? - /// .on_proxy_session_start(responder, async |session_id| { + /// .on_proxy_session_start(responder, async |response| { /// // Session started /// Ok(()) /// }) @@ -801,7 +837,7 @@ where op: F, ) -> Result<(), crate::Error> where - F: FnOnce(SessionId) -> Fut + Send + 'static, + F: FnOnce(NewSessionResponse) -> Fut + Send + 'static, Fut: Future> + Send, Counterpart: HasPeer, R: 'static, @@ -837,11 +873,11 @@ where // Extract the session-id from the response and forward // the response back to the client let session_id = response.session_id.clone(); - responder.respond(response)?; + responder.respond(response.clone())?; // Install a dynamic handler to proxy messages from this session connection - .add_dynamic_handler(ProxySessionMessages::new(session_id.clone()))? + .add_dynamic_handler(ProxySessionMessages::new(session_id))? .detach(); // Keep dynamic handlers live for the connection. @@ -849,7 +885,7 @@ where .into_iter() .for_each(DynamicHandlerGuard::detach); - connection.spawn(async move { op(session_id).await }) + connection.spawn(async move { op(response).await }) } }) } @@ -997,11 +1033,14 @@ where /// [`start_session`](Self::start_session) instead and call /// [`proxy_remaining_messages`](ActiveSession::proxy_remaining_messages) manually. /// + /// Returns a clone of the complete [`NewSessionResponse`] forwarded to the + /// client. + /// /// Requires calling [`block_task`](Self::block_task) first. pub async fn start_session_proxy( self, responder: Responder, - ) -> Result + ) -> Result where Counterpart: HasPeer, R: 'static, @@ -1009,15 +1048,26 @@ where let SessionSetup { session, response, .. } = self.start_session().await?; - let session_id = session.session_id().clone(); - responder.respond(response)?; + responder.respond(response.clone())?; session.proxy_remaining_messages()?; - Ok(session_id) + Ok(response) } } /// Stable protocol v1 active session that lets you send prompts and receive updates. /// +/// The session owns a stream of [`SessionMessage`]s: the requests and +/// notifications the agent sends for this session, and the responses to +/// prompts sent with [`send_prompt`](Self::send_prompt) or +/// [`send_prompt_blocks`](Self::send_prompt_blocks). +/// +/// This handle covers prompting and reading the stream only. It has no +/// helpers for `session/cancel`, `session/set_mode`, +/// `session/set_config_option`, `session/close` or `session/fork`. Send those +/// messages through [`connection`](Self::connection) with +/// [`session_id`](Self::session_id), for example +/// `session.connection().send_notification_to(Agent, CancelNotification::new(session.session_id().clone()))`. +/// /// The `'runner` lifetime represents the span during which session support runners /// (such as MCP servers) are active. When created via [`SessionBuilder::start_session`], /// this is `'static` because the runners are spawned into background tasks. @@ -1048,20 +1098,31 @@ where _runner: PhantomData<&'runner ()>, } -/// Incoming stable protocol v1 message from the agent. +/// An item on a stable protocol v1 session's update stream, returned by +/// [`ActiveSession::read_update`]. #[non_exhaustive] #[derive(Debug)] #[allow( clippy::large_enum_variant, - reason = "Dispatch messages vastly outnumber StopReason; boxing would add a heap allocation" + reason = "messages vastly outnumber prompt responses; boxing would add a heap allocation" )] pub enum SessionMessage { - /// Periodic updates with new content, tool requests, etc. - /// Use [`MatchDispatch`] to match on the message type. - SessionMessage(Dispatch), - - /// When a prompt completes, the stop reason. - StopReason(StopReason), + /// A request or notification the agent sent for this session, such as a + /// `session/update` notification or a `session/request_permission` + /// request. + /// + /// Request handles are typed for protocol v1. To match on the message + /// type with [`MatchDispatch`], convert it with + /// [`Dispatch::into_raw`](crate::Dispatch::into_raw) first. + Message(crate::v1::Dispatch), + + /// The response to a `session/prompt` request sent with + /// [`ActiveSession::send_prompt`] or + /// [`ActiveSession::send_prompt_blocks`], which ends that prompt turn. + /// + /// An error response from the agent arrives here as `Err`. It does not + /// close the connection or affect other sessions. + PromptResponse(Result), } impl ActiveSession<'_, Link> @@ -1078,27 +1139,41 @@ where &self.connection } - /// Send a prompt to the agent. You can then read messages sent in response. + /// Send a text prompt to the agent. + /// + /// Read the agent's output with [`read_update`](Self::read_update) or + /// [`read_to_string`](Self::read_to_string). The `session/prompt` + /// response, or the error the agent answered with, arrives on the same + /// stream as [`SessionMessage::PromptResponse`]. pub fn send_prompt(&mut self, prompt: impl ToString) -> Result<(), crate::Error> { + self.send_prompt_blocks(vec![prompt.to_string().into()]) + } + + /// Send a prompt made of arbitrary content blocks to the agent. + /// + /// Content other than text and resource links needs the matching + /// `promptCapabilities` entry from the agent's initialize response. Like + /// [`send_prompt`](Self::send_prompt), the + /// response arrives on the update stream as + /// [`SessionMessage::PromptResponse`]. + pub fn send_prompt_blocks(&mut self, prompt: Vec) -> Result<(), crate::Error> { let update_tx = self.update_tx.clone(); self.connection .raw_connection() .send_ordered_request_to::( Agent, - PromptRequest::new(self.session_id.clone(), vec![prompt.to_string().into()]), + PromptRequest::new(self.session_id.clone(), prompt), ) .on_receiving_result(async move |result| { - let PromptResponse { stop_reason, .. } = result?; - - update_tx - .unbounded_send(SessionMessage::StopReason(stop_reason)) - .map_err(crate::util::internal_error)?; - + // A closed channel means the session was dropped or handed to + // `proxy_remaining_messages`, so nobody reads the response. + drop(update_tx.unbounded_send(SessionMessage::PromptResponse(result))); Ok(()) }) } - /// Read an update from the agent in response to the prompt. + /// Read the next item on the session's update stream: a message the agent + /// sent for this session, or the response to a prompt. pub async fn read_update(&mut self) -> Result { use futures::StreamExt; let message = @@ -1109,15 +1184,21 @@ where Ok(message) } - /// Read all updates until the end of the turn and create a string. - /// Ignores non-text updates. + /// Read updates until the next prompt response and return the text of the + /// agent's message chunks. + /// + /// Ignores non-text updates. Requests from the agent, such as + /// `session/request_permission`, are dropped, so the SDK answers them + /// with an Internal Error (`-32603`). Use + /// [`read_update`](Self::read_update) to answer them. If the agent answers + /// the prompt with an error, this returns that error. pub async fn read_to_string(&mut self) -> Result { let mut output = String::new(); loop { let update = self.read_update().await?; tracing::trace!("read_to_string update"); match update { - SessionMessage::SessionMessage(dispatch) => MatchDispatch::new(dispatch) + SessionMessage::Message(dispatch) => MatchDispatch::new(dispatch.into_raw()) .if_notification(async |notif: SessionNotification| match notif.update { SessionUpdate::AgentMessageChunk(ContentChunk { content: ContentBlock::Text(text), @@ -1130,7 +1211,10 @@ where }) .await .otherwise_ignore()?, - SessionMessage::StopReason(_stop_reason) => break, + SessionMessage::PromptResponse(response) => { + response?; + break; + } } } Ok(output) @@ -1199,12 +1283,12 @@ where // consumed yet. We must forward them to maintain message ordering. while let Ok(message) = update_rx.try_recv() { match message { - SessionMessage::SessionMessage(dispatch) => { + SessionMessage::Message(dispatch) => { // Forward the message to the client connection.send_proxied_message_to(Client, dispatch)?; } - SessionMessage::StopReason(_) => { - // StopReason is internal bookkeeping, not forwarded + SessionMessage::PromptResponse(_) => { + // Responses to the proxy's own prompts are not forwarded } } } @@ -1264,7 +1348,7 @@ where ); if session_id == self.session_id { self.update_tx - .unbounded_send(SessionMessage::SessionMessage(message)) + .unbounded_send(SessionMessage::Message(message.into_protocol()?)) .map_err(crate::util::internal_error)?; return Ok(Handled::Yes); } diff --git a/src/agent-client-protocol/src/v1.rs b/src/agent-client-protocol/src/v1.rs index 1a170c17..98cbdf96 100644 --- a/src/agent-client-protocol/src/v1.rs +++ b/src/agent-client-protocol/src/v1.rs @@ -45,8 +45,8 @@ pub use crate::schema::v1 as schema; pub use crate::jsonrpc::v1_connection::ConnectionTo; pub use crate::session::{ - ActiveSession, Blocking, LoadSessionBuilder, NonBlocking, RestoreSessionBuilder, - ResumeSessionBuilder, SessionBlockState, SessionBuilder, SessionMessage, SessionSetup, + ActiveSession, Blocking, LoadSessionBuilder, NonBlocking, ResumeSessionBuilder, + SessionBlockState, SessionBuilder, SessionMessage, SessionSetup, }; /// A protocol v1 connection builder, returned by diff --git a/src/agent-client-protocol/tests/session_ordering.rs b/src/agent-client-protocol/tests/session_ordering.rs index 0c5714cd..31099e7b 100644 --- a/src/agent-client-protocol/tests/session_ordering.rs +++ b/src/agent-client-protocol/tests/session_ordering.rs @@ -87,8 +87,8 @@ mod callback_future_lifetimes { } fn proxy_session_callback<'a>() - -> impl FnOnce(SessionId) -> LifetimeTaggedFuture<'a> + Send + 'static { - |_session_id| LifetimeTaggedFuture(PhantomData) + -> impl FnOnce(NewSessionResponse) -> LifetimeTaggedFuture<'a> + Send + 'static { + |_response| LifetimeTaggedFuture(PhantomData) } #[cfg(feature = "unstable_protocol_v2")] @@ -280,6 +280,108 @@ async fn on_session_start_callback_can_consume_later_session_messages() { .expect("session connection failed"); } +#[tokio::test(flavor = "current_thread")] +async fn prompt_error_reaches_the_session_stream_and_keeps_the_connection() { + let session_id = SessionId::new("prompt-error-session"); + let new_session_id = session_id.clone(); + let other_session_id = SessionId::new("other-session"); + let other_new_session_id = other_session_id.clone(); + let meta = agent_client_protocol::schema::v1::Meta::from_iter([( + "turn".to_owned(), + serde_json::json!({ "kept": true }), + )]); + let expected_response = PromptResponse::new(StopReason::EndTurn).meta(meta.clone()); + + let agent = Agent + .v1() + .on_receive_request( + { + let created = std::sync::Mutex::new(0_u8); + async move |_request: NewSessionRequest, + responder: Responder, + _connection: ConnectionTo| { + let id = { + let mut created = created.lock().expect("lock poisoned"); + *created += 1; + if *created == 1 { + new_session_id.clone() + } else { + other_new_session_id.clone() + } + }; + responder.respond(NewSessionResponse::new(id)) + } + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async move |request: PromptRequest, + responder: Responder, + connection: ConnectionTo| { + let ContentBlock::Text(text) = &request.prompt[0] else { + panic!("expected a text prompt"); + }; + if text.text == "fail" { + return responder.respond_with_error( + agent_client_protocol::Error::invalid_params().data("rejected prompt"), + ); + } + connection.send_notification(SessionNotification::new( + request.session_id, + SessionUpdate::AgentMessageChunk(ContentChunk::new(ContentBlock::Text( + TextContent::new("ok"), + ))), + ))?; + responder.respond(PromptResponse::new(StopReason::EndTurn).meta(meta.clone())) + }, + agent_client_protocol::on_receive_request!(), + ); + + let client = Client.v1().connect_with(agent, async move |connection| { + let mut session = connection + .build_session_cwd()? + .block_task() + .start_session() + .await? + .session; + let mut other = connection + .build_session_cwd()? + .block_task() + .start_session() + .await? + .session; + assert_eq!(session.session_id(), &session_id); + assert_eq!(other.session_id(), &other_session_id); + + session.send_prompt("fail")?; + let error = session + .read_to_string() + .await + .expect_err("the prompt error should be returned"); + assert_eq!(error.code, agent_client_protocol::ErrorCode::INVALID_PARAMS); + + // The error did not close the connection or the other session. + other.send_prompt("hello")?; + assert_eq!(other.read_to_string().await?, "ok"); + + // The same session keeps working, and the complete response arrives. + session.send_prompt_blocks(vec![ContentBlock::Text(TextContent::new("hello"))])?; + let SessionMessage::Message(_) = session.read_update().await? else { + panic!("expected the agent message chunk first"); + }; + let SessionMessage::PromptResponse(response) = session.read_update().await? else { + panic!("expected the prompt response"); + }; + assert_eq!(response?, expected_response); + Ok(()) + }); + + tokio::time::timeout(TIMEOUT, client) + .await + .expect("prompt error test timed out") + .expect("session connection failed"); +} + #[tokio::test(flavor = "current_thread")] async fn on_session_start_installs_routing_before_later_message() { let session_id = SessionId::new("queued-session"); @@ -294,7 +396,7 @@ async fn on_session_start_installs_routing_before_later_message() { connection.build_session_cwd()?.on_session_start( async move |SessionSetup { mut session, .. }| { let update = session.read_update().await?; - assert!(matches!(update, SessionMessage::SessionMessage(_))); + assert!(matches!(update, SessionMessage::Message(_))); result_tx .send(()) .map_err(|()| agent_client_protocol::Error::internal_error()) @@ -366,7 +468,7 @@ async fn start_session_installs_routing_before_later_message() { .await? .session; let update = session.read_update().await?; - assert!(matches!(update, SessionMessage::SessionMessage(_))); + assert!(matches!(update, SessionMessage::Message(_))); Ok(()) }); diff --git a/src/agent-client-protocol/tests/session_restore.rs b/src/agent-client-protocol/tests/session_restore.rs index 5f46d3b4..6300b998 100644 --- a/src/agent-client-protocol/tests/session_restore.rs +++ b/src/agent-client-protocol/tests/session_restore.rs @@ -79,7 +79,7 @@ fn dispatch_barrier_notification() -> RawJsonRpcMessage { } fn session_message_text(message: SessionMessage) -> String { - let SessionMessage::SessionMessage(dispatch) = message else { + let SessionMessage::Message(dispatch) = message else { panic!("expected a session notification") }; let untyped = dispatch diff --git a/src/yopo/src/lib.rs b/src/yopo/src/lib.rs index f2a98530..51c2a388 100644 --- a/src/yopo/src/lib.rs +++ b/src/yopo/src/lib.rs @@ -122,8 +122,8 @@ pub async fn prompt_with_callback( loop { let update = session.read_update().await?; match update { - agent_client_protocol::v1::SessionMessage::SessionMessage(message) => { - MatchDispatch::new(message) + agent_client_protocol::v1::SessionMessage::Message(message) => { + MatchDispatch::new(message.into_raw()) .if_notification(async |notification: SessionNotification| { tracing::debug!( ?notification, @@ -167,8 +167,8 @@ pub async fn prompt_with_callback( .otherwise(async |_msg| Ok(())) .await?; } - agent_client_protocol::v1::SessionMessage::StopReason(stop_reason) => { - match stop_reason { + agent_client_protocol::v1::SessionMessage::PromptResponse(response) => { + match response?.stop_reason { StopReason::EndTurn => break, StopReason::MaxTokens => { tracing::debug!("Agent hit max tokens limit");