diff --git a/include/mcp/filter/callback_bridges.h b/include/mcp/filter/callback_bridges.h index 4a97c516e..d949ef5b8 100644 --- a/include/mcp/filter/callback_bridges.h +++ b/include/mcp/filter/callback_bridges.h @@ -99,6 +99,12 @@ class JsonRpcToProtocolBridge : public JsonRpcProtocolFilter::MessageHandler { void onNotification(const jsonrpc::Notification& notification) override; void onResponse(const jsonrpc::Response& response) override; void onProtocolError(const Error& error) override; + // Propagate the per-message origin built by the JSON-RPC filter so the + // application layer can reply on the connection the message arrived on. + void onRequestWithContext(const jsonrpc::Request& request, + MessageDispatchContext& context) override; + void onNotificationWithContext(const jsonrpc::Notification& notification, + MessageDispatchContext& context) override; private: McpProtocolCallbacks& callbacks_; diff --git a/include/mcp/filter/json_rpc_filter_factory.h b/include/mcp/filter/json_rpc_filter_factory.h index 49cd68168..05cb501ed 100644 --- a/include/mcp/filter/json_rpc_filter_factory.h +++ b/include/mcp/filter/json_rpc_filter_factory.h @@ -33,6 +33,18 @@ class DirectJsonRpcCallbacks : public JsonRpcProtocolFilter::MessageHandler { mcp_callbacks_.onError(error); } + // Propagate the per-message origin built by the JSON-RPC filter so the + // application layer can reply on the connection the message arrived on. + void onRequestWithContext(const jsonrpc::Request& request, + MessageDispatchContext& context) override { + mcp_callbacks_.onRequestWithContext(request, context); + } + + void onNotificationWithContext(const jsonrpc::Notification& notification, + MessageDispatchContext& context) override { + mcp_callbacks_.onNotificationWithContext(notification, context); + } + private: McpProtocolCallbacks& mcp_callbacks_; }; diff --git a/include/mcp/filter/json_rpc_protocol_filter.h b/include/mcp/filter/json_rpc_protocol_filter.h index 5e6187bbf..32f06650f 100644 --- a/include/mcp/filter/json_rpc_protocol_filter.h +++ b/include/mcp/filter/json_rpc_protocol_filter.h @@ -4,6 +4,7 @@ #include "mcp/event/event_loop.h" #include "mcp/filter/filter_context.h" #include "mcp/json/json_bridge.h" +#include "mcp/message_dispatch_context.h" #include "mcp/network/filter.h" #include "mcp/types.h" @@ -66,6 +67,30 @@ class JsonRpcProtocolFilter : public network::Filter { * @param error The error details */ virtual void onProtocolError(const Error& error) = 0; + + /** + * Context-carrying variants. The filter that parsed the message builds + * a per-message context (origin connection + a reply path through this + * filter's own encoder) and dispatches through these; the defaults + * forward to the context-free hooks so existing handlers keep working. + * Handlers that forward toward McpProtocolCallbacks should override + * these and propagate the context so origin information travels with + * the message all the way to the application layer. Distinct names + * (not overloads) keep existing handlers outside the + * -Woverloaded-virtual hiding trap. + */ + virtual void onRequestWithContext(const jsonrpc::Request& request, + MessageDispatchContext& context) { + (void)context; + onRequest(request); + } + + virtual void onNotificationWithContext( + const jsonrpc::Notification& notification, + MessageDispatchContext& context) { + (void)context; + onNotification(notification); + } }; /** @@ -148,6 +173,12 @@ class JsonRpcProtocolFilter : public network::Filter { class EncoderImpl; friend class EncoderImpl; // Allow encoder to access private members + // Per-message dispatch context handed to the MessageHandler: origin is + // this filter's connection, reply path is this filter's encoder. Nested + // (like EncoderImpl) so it can read write_callbacks_. + class DispatchContextImpl; + friend class DispatchContextImpl; + /** * Parse messages from buffer * @param buffer Data buffer containing JSON messages diff --git a/include/mcp/mcp_application_base.h b/include/mcp/mcp_application_base.h index 4fcf093ed..bc1f2e323 100644 --- a/include/mcp/mcp_application_base.h +++ b/include/mcp/mcp_application_base.h @@ -143,6 +143,18 @@ class McpToJsonRpcAdapter callbacks_.onError(error); } + // Propagate the per-message origin built by the JSON-RPC filter so the + // application layer can reply on the connection the message arrived on. + void onRequestWithContext(const jsonrpc::Request& request, + MessageDispatchContext& context) override { + callbacks_.onRequestWithContext(request, context); + } + + void onNotificationWithContext(const jsonrpc::Notification& notification, + MessageDispatchContext& context) override { + callbacks_.onNotificationWithContext(notification, context); + } + private: McpProtocolCallbacks& callbacks_; }; diff --git a/include/mcp/mcp_connection_manager.h b/include/mcp/mcp_connection_manager.h index c008a68f1..550aba954 100644 --- a/include/mcp/mcp_connection_manager.h +++ b/include/mcp/mcp_connection_manager.h @@ -7,6 +7,7 @@ #include "mcp/core/result.h" #include "mcp/event/event_loop.h" #include "mcp/json/json_bridge.h" +#include "mcp/message_dispatch_context.h" #include "mcp/network/connection_manager.h" #include "mcp/network/filter.h" #include "mcp/transport/http_sse_transport_socket.h" @@ -70,6 +71,34 @@ class McpProtocolCallbacks { */ virtual void onNotification(const jsonrpc::Notification& notification) = 0; + /** + * Context-carrying variants: the message arrives together with a + * per-message dispatch context describing its origin (connection, + * transport session id) and reply path. Producers that know the origin + * call these; the defaults forward to the context-free hooks so existing + * implementations keep working unchanged. Receivers that route replies or + * key sessions should override these instead of the context-free forms — + * the context makes "respond to the wrong connection" and "inherit a + * stale session binding" unrepresentable, where ambient + * current-connection state cannot. + * + * Distinct names (rather than overloads of onRequest/onNotification) keep + * every existing implementation outside the -Woverloaded-virtual hiding + * trap. + */ + virtual void onRequestWithContext(const jsonrpc::Request& request, + MessageDispatchContext& context) { + (void)context; + onRequest(request); + } + + virtual void onNotificationWithContext( + const jsonrpc::Notification& notification, + MessageDispatchContext& context) { + (void)context; + onNotification(notification); + } + /** * Called when a response is received */ @@ -93,26 +122,6 @@ class McpProtocolCallbacks { (void)endpoint; // Default implementation does nothing } - /** - * Called immediately before a parsed message is dispatched, carrying the - * transport-level session id the message belongs to. For HTTP+SSE servers - * this is the SSE stream id extracted from the POST /callback/{id} path — - * the durable client identity that outlives the one-shot POST connection - * the message physically arrived on. Empty when the transport has no - * session concept (stdio, plain HTTP); receivers must then fall back to - * connection identity. - * - * Always invoked in the dispatcher thread right before the matching - * onRequest/onNotification/onResponse for the same message, so an - * implementation may stash it as request-scoped context without locking. - * It is re-announced per message (not per connection) because reads from - * different connections interleave on the dispatcher thread. - */ - virtual void onTransportSessionBound( - const std::string& transport_session_id) { - (void)transport_session_id; // Default implementation does nothing - } - /** * Send a POST request to the message endpoint * Used by HTTP/SSE transport to send messages on a separate connection @@ -203,6 +212,12 @@ class McpConnectionManager : public McpProtocolCallbacks, void onRequest(const jsonrpc::Request& request) override; void onNotification(const jsonrpc::Notification& notification) override; void onResponse(const jsonrpc::Response& response) override; + // Forward the per-message origin alongside the message so the application + // layer can key sessions and route replies without ambient state. + void onRequestWithContext(const jsonrpc::Request& request, + MessageDispatchContext& context) override; + void onNotificationWithContext(const jsonrpc::Notification& notification, + MessageDispatchContext& context) override; void onConnectionEvent(network::ConnectionEvent event) override; void onError(const Error& error) override; void onMessageEndpoint(const std::string& endpoint) override; diff --git a/include/mcp/message_dispatch_context.h b/include/mcp/message_dispatch_context.h new file mode 100644 index 000000000..cd70bafb0 --- /dev/null +++ b/include/mcp/message_dispatch_context.h @@ -0,0 +1,81 @@ +#pragma once + +#include + +#include "mcp/core/result.h" +#include "mcp/types.h" + +namespace mcp { + +namespace network { +class Connection; +} + +/** + * Per-message dispatch context. + * + * Carries where a JSON-RPC message came from and how to send a reply back + * along the same path. The transport producer (the filter or connection + * manager that parsed the message) constructs one immediately before + * dispatching the message and passes it through the callback chain, so the + * origin travels *with* the message instead of living in ambient state on + * the receiver. A stale binding is unrepresentable by construction: the + * context dies when the dispatch call returns. + * + * Lifetime contract: valid only for the duration of the dispatch call, on + * the dispatcher thread. Receivers must not retain a pointer or reference + * past the callback's return; a handler that wants to finish work + * asynchronously must resolve what it needs (session id, transport session + * id) while the context is live. + */ +class MessageDispatchContext { + public: + virtual ~MessageDispatchContext() = default; + + /** + * The connection the message physically arrived on. May be null when the + * producer is not connection-backed (e.g. a legacy dispatch path with no + * origin information). Only guaranteed valid while the context is live. + */ + virtual network::Connection* originConnection() const = 0; + + /** + * Durable transport-level session id the message belongs to (e.g. the SSE + * stream id from a POST /callback/{id} path — the client identity that + * outlives the one-shot POST connection). Empty when the transport has no + * session concept (stdio, plain HTTP); receivers then fall back to + * connection identity. + */ + virtual const std::string& transportSessionId() const = 0; + + /** + * Send a JSON-RPC response back along this message's own return path. + * Returns an error (rather than silently dropping) when the path is gone, + * e.g. the origin connection already closed. + */ + virtual VoidResult sendResponse(const jsonrpc::Response& response) = 0; +}; + +/** + * Context for dispatch paths that carry no origin information. Session + * resolution falls back to "no connection" and any attempted reply fails + * loudly instead of being written to an unrelated connection. + */ +class NullMessageDispatchContext : public MessageDispatchContext { + public: + network::Connection* originConnection() const override { return nullptr; } + + const std::string& transportSessionId() const override { + static const std::string empty; + return empty; + } + + VoidResult sendResponse(const jsonrpc::Response&) override { + Error err; + err.code = jsonrpc::INTERNAL_ERROR; + err.message = "no dispatch context: response has no return path"; + return makeVoidError(err); + } +}; + +} // namespace mcp diff --git a/include/mcp/server/mcp_server.h b/include/mcp/server/mcp_server.h index 9457085ab..8ea0323ef 100644 --- a/include/mcp/server/mcp_server.h +++ b/include/mcp/server/mcp_server.h @@ -960,21 +960,20 @@ class McpServer : public application::ApplicationBase, void onRequest(const jsonrpc::Request& request) override; void onNotification(const jsonrpc::Notification& notification) override; void onResponse(const jsonrpc::Response& response) override; + + // Context-carrying dispatch entry points. The context travels with the + // message from the filter that parsed it, so session keying and the reply + // both bind to the message's own origin — never to whichever connection + // happened to be accepted or announced most recently. The context-free + // overrides above remain only as a degraded fallback for producers that + // do not supply origin information. + void onRequestWithContext(const jsonrpc::Request& request, + MessageDispatchContext& context); + void onNotificationWithContext(const jsonrpc::Notification& notification, + MessageDispatchContext& context); void onConnectionEvent(network::ConnectionEvent event); void onError(const Error& error) override; - // Request-scoped transport session binding. The transport filter announces - // the transport-level session id (e.g. SSE stream id from the POST - // /callback/{id} path) immediately before dispatching each message, on the - // dispatcher thread. It is single-use: getOrCreateCurrentSession() consumes - // and clears it, so it applies only to the one message whose dispatch it - // preceded and a non-announcing producer can never inherit a stale id. An - // empty id means the current message's transport has no session concept and - // session lookup falls back to connection identity. - void onTransportSessionBound(const std::string& transport_session_id) { - current_transport_session_id_ = transport_session_id; - } - // Request tracking helpers bool isRequestCancelled(const RequestId& id) const { std::lock_guard lock(pending_requests_mutex_); @@ -996,12 +995,18 @@ class McpServer : public application::ApplicationBase, // Register built-in handlers void registerBuiltinHandlers(); - // Resolve the session for the message currently being dispatched. - // Prefers the transport session id (durable across the short-lived POST - // connections of HTTP+SSE) and falls back to connection identity for + // Resolve the session for the message described by `context`. Prefers + // the transport session id (durable across the short-lived POST + // connections of HTTP+SSE) and falls back to the origin connection for // transports without one (stdio, plain HTTP). Returns nullptr only when // the session limit is reached. Dispatcher thread only. - SessionManager::SessionPtr getOrCreateCurrentSession(); + SessionManager::SessionPtr getOrCreateSessionFor( + const MessageDispatchContext& context); + + // Reply path for messages that arrived without origin information (the + // context-free legacy hooks): no connection to key a session on, replies + // fall back to the first connected stdio manager. Defined in the .cc. + class LegacyDispatchContext; // Deliver a notification to one session's client, routed by how the // session is keyed: SSE stream (via the registry), owning connection, @@ -1095,6 +1100,19 @@ class McpServer : public application::ApplicationBase, server_.onNotification(notification); } + void onRequestWithContext(const jsonrpc::Request& request, + MessageDispatchContext& context) override { + GOPHER_LOG_DEBUG( + "ServerProtocolCallbacks::onRequestWithContext for method: {}", + request.method); + server_.onRequestWithContext(request, context); + } + + void onNotificationWithContext(const jsonrpc::Notification& notification, + MessageDispatchContext& context) override { + server_.onNotificationWithContext(notification, context); + } + void onResponse(const jsonrpc::Response& response) override { server_.onResponse(response); } @@ -1105,11 +1123,6 @@ class McpServer : public application::ApplicationBase, void onError(const Error& error) override { server_.onError(error); } - void onTransportSessionBound( - const std::string& transport_session_id) override { - server_.onTransportSessionBound(transport_session_id); - } - private: McpServer& server_; }; @@ -1154,15 +1167,6 @@ class McpServer : public application::ApplicationBase, // Following production pattern: listener owns connections, not threads // Connections tracked by count only, ownership managed by listener - // Current connection being processed (for request context) - // This is set temporarily during request processing in dispatcher thread - network::Connection* current_connection_{nullptr}; - - // Transport session id of the message currently being dispatched (see - // onTransportSessionBound). Dispatcher thread only, like - // current_connection_. Empty for transports without session identity. - std::string current_transport_session_id_; - // Active connections owned by server // Following production pattern: all operations in dispatcher thread, no mutex // needed Connections removed when they close via callbacks diff --git a/src/filter/callback_bridges.cc b/src/filter/callback_bridges.cc index ab589416a..8b9654502 100644 --- a/src/filter/callback_bridges.cc +++ b/src/filter/callback_bridges.cc @@ -153,6 +153,24 @@ void JsonRpcToProtocolBridge::onNotification( callbacks_.onNotification(notification); } +void JsonRpcToProtocolBridge::onRequestWithContext( + const jsonrpc::Request& request, MessageDispatchContext& context) { + GOPHER_LOG_DEBUG( + "JsonRpcToProtocolBridge::onRequestWithContext called with method='{}'", + request.method); + callbacks_.onRequestWithContext(request, context); +} + +void JsonRpcToProtocolBridge::onNotificationWithContext( + const jsonrpc::Notification& notification, + MessageDispatchContext& context) { + GOPHER_LOG_DEBUG( + "JsonRpcToProtocolBridge::onNotificationWithContext called with " + "method='{}'", + notification.method); + callbacks_.onNotificationWithContext(notification, context); +} + void JsonRpcToProtocolBridge::onResponse(const jsonrpc::Response& response) { if (response.id.has_value()) { GOPHER_LOG_DEBUG("JsonRpcToProtocolBridge::onResponse called with id={}", diff --git a/src/filter/enhanced_filter_chain_factory.cc b/src/filter/enhanced_filter_chain_factory.cc index 3cf305832..a428d90b0 100644 --- a/src/filter/enhanced_filter_chain_factory.cc +++ b/src/filter/enhanced_filter_chain_factory.cc @@ -333,6 +333,25 @@ class EnhancedProtocolFilter : public network::Filter, mcp_callbacks_.onNotification(notification); } + // Context-carrying variants: track metrics identically, then propagate + // the per-message origin so the application layer can reply on the + // connection the message arrived on. + void onRequestWithContext(const jsonrpc::Request& request, + MessageDispatchContext& context) override { + if (metrics_collector_) { + metrics_collector_->onRequest(request); + } + mcp_callbacks_.onRequestWithContext(request, context); + } + + void onNotificationWithContext(const jsonrpc::Notification& notification, + MessageDispatchContext& context) override { + if (metrics_collector_) { + metrics_collector_->onNotification(notification); + } + mcp_callbacks_.onNotificationWithContext(notification, context); + } + void onProtocolError(const Error& error) override { // Track errors if (circuit_breaker_) { diff --git a/src/filter/http_sse_filter_chain_factory.cc b/src/filter/http_sse_filter_chain_factory.cc index 3e129dc38..dacb5a864 100644 --- a/src/filter/http_sse_filter_chain_factory.cc +++ b/src/filter/http_sse_filter_chain_factory.cc @@ -925,6 +925,55 @@ class HttpSseJsonRpcProtocolFilter // ===== JsonRpcProtocolFilter::MessageHandler ===== + /** + * Per-message dispatch context for messages decoded on this composite + * chain. Origin is the connection the message physically arrived on (for + * HTTP+SSE, the short-lived POST connection); the transport session id is + * the SSE stream id parsed from POST /callback/{id} — the durable client + * identity the server keys its session on. + * + * The reply sink writes the bare JSON to the origin connection, exactly + * the bytes the server used to write to its ambient current-connection + * pointer: this composite's own onWrite then decides the wire form + * (reroute through the SSE registry for a callback proxy, HTTP-frame for + * plain HTTP). Same wire behavior, but the destination is the message's + * own connection by construction. + */ + class DispatchContext : public MessageDispatchContext { + public: + explicit DispatchContext(HttpSseJsonRpcProtocolFilter& parent) + : parent_(parent) {} + + network::Connection* originConnection() const override { + return parent_.write_callbacks_ ? &parent_.write_callbacks_->connection() + : nullptr; + } + + const std::string& transportSessionId() const override { + return parent_.sse_callback_session_id_; + } + + VoidResult sendResponse(const jsonrpc::Response& response) override { + // Fail loudly when the reply path is gone instead of pretending the + // response went out. + if (!parent_.write_callbacks_) { + Error err; + err.code = jsonrpc::INTERNAL_ERROR; + err.message = "response dropped: origin connection is gone"; + return makeVoidError(err); + } + auto json_val = json::to_json(response); + std::string json_str = json_val.toString(); + OwnedBuffer buffer; + buffer.add(json_str); + parent_.write_callbacks_->connection().write(buffer, false); + return makeVoidSuccess(); + } + + private: + HttpSseJsonRpcProtocolFilter& parent_; + }; + /** * Called by JsonRpcProtocolFilter when a complete JSON-RPC request is parsed * Creates a RequestStream to track this request-response pair @@ -932,20 +981,35 @@ class HttpSseJsonRpcProtocolFilter */ void onRequest(const jsonrpc::Request& request) override { GOPHER_LOG_DEBUG("HttpSseFilter::onRequest for method: {}", request.method); - // Announce which transport session this message belongs to before - // dispatching. On a POST /callback/{id} connection this carries the SSE - // stream id — the durable client identity — so the server can key its - // MCP session on it instead of this short-lived POST connection. It is - // announced per message (not per connection) because dispatcher-thread - // reads from different connections interleave; an empty id explicitly - // clears any previous connection's binding. - mcp_callbacks_.onTransportSessionBound(sse_callback_session_id_); - mcp_callbacks_.onRequest(request); + // The context carries the transport session id (the SSE stream id from + // a POST /callback/{id} path — the durable client identity) with the + // message itself. Built fresh per message because dispatcher-thread + // reads from different connections interleave; a previous message's + // binding cannot leak because the previous context is already gone. + DispatchContext context(*this); + mcp_callbacks_.onRequestWithContext(request, context); + } + + /** + * The JSON-RPC sub-filter dispatches through here with its own generic + * context, but that context knows neither the SSE stream id nor this + * composite's write semantics — replace it with the composite's own. + */ + void onRequestWithContext(const jsonrpc::Request& request, + MessageDispatchContext& context) override { + (void)context; + onRequest(request); + } + + void onNotificationWithContext(const jsonrpc::Notification& notification, + MessageDispatchContext& context) override { + (void)context; + onNotification(notification); } void onNotification(const jsonrpc::Notification& notification) override { - mcp_callbacks_.onTransportSessionBound(sse_callback_session_id_); - mcp_callbacks_.onNotification(notification); + DispatchContext context(*this); + mcp_callbacks_.onNotificationWithContext(notification, context); // For HTTP transport, send HTTP 202 Accepted response // JSON-RPC notifications don't have responses, but HTTP requires one @@ -970,7 +1034,6 @@ class HttpSseJsonRpcProtocolFilter } void onResponse(const jsonrpc::Response& response) override { - mcp_callbacks_.onTransportSessionBound(sse_callback_session_id_); mcp_callbacks_.onResponse(response); } diff --git a/src/filter/json_rpc_protocol_filter.cc b/src/filter/json_rpc_protocol_filter.cc index 792cbd9d7..0550bc3b7 100644 --- a/src/filter/json_rpc_protocol_filter.cc +++ b/src/filter/json_rpc_protocol_filter.cc @@ -165,6 +165,51 @@ class JsonRpcProtocolFilter::EncoderImpl JsonRpcProtocolFilter& parent_; }; +// DispatchContextImpl - per-message origin + reply path for the handler. +// +// Why the encoder is the reply path: the response must be framed exactly +// like any other outbound message on the connection the request arrived on +// (newline-delimited or length-prefixed, per this filter's configuration). +// Routing the reply through the parsing filter's own encoder pairs decode +// and encode on the same connection, so a handler can never answer on a +// different connection than the one that asked. +// +// Stack-constructed immediately before each dispatch; dies when the +// dispatch returns, which is what makes a stale origin unrepresentable. +class JsonRpcProtocolFilter::DispatchContextImpl + : public MessageDispatchContext { + public: + explicit DispatchContextImpl(JsonRpcProtocolFilter& parent) + : parent_(parent) {} + + network::Connection* originConnection() const override { + return parent_.write_callbacks_ ? &parent_.write_callbacks_->connection() + : nullptr; + } + + const std::string& transportSessionId() const override { + // A bare JSON-RPC chain has no transport session concept; composite + // transports (HTTP+SSE) supply their own context with the stream id. + static const std::string empty; + return empty; + } + + VoidResult sendResponse(const jsonrpc::Response& response) override { + // Fail loudly when the reply path is gone: a null result here would + // otherwise read as "sent" to the caller while nothing went out. + if (!parent_.write_callbacks_) { + Error err; + err.code = jsonrpc::INTERNAL_ERROR; + err.message = "response dropped: origin connection is gone"; + return makeVoidError(err); + } + return parent_.encoder_->encodeResponse(response); + } + + private: + JsonRpcProtocolFilter& parent_; +}; + // JsonRpcProtocolFilter implementation JsonRpcProtocolFilter::JsonRpcProtocolFilter(MessageHandler& handler, @@ -328,13 +373,15 @@ bool JsonRpcProtocolFilter::parseMessage(const std::string& json_str) { GOPHER_LOG_DEBUG("JsonRpcFilter dispatching request for method: {}", request.method); requests_received_++; - handler_.onRequest(request); + DispatchContextImpl context(*this); + handler_.onRequestWithContext(request, context); } else { // JSON-RPC Notification jsonrpc::Notification notification = json::from_json(json_val); notifications_received_++; - handler_.onNotification(notification); + DispatchContextImpl context(*this); + handler_.onNotificationWithContext(notification, context); } } else if (json_val.contains("result") || json_val.contains("error")) { // JSON-RPC Response diff --git a/src/filter/protocol_detection_filter_chain_factory.cc b/src/filter/protocol_detection_filter_chain_factory.cc index 28776417c..bfab2436b 100644 --- a/src/filter/protocol_detection_filter_chain_factory.cc +++ b/src/filter/protocol_detection_filter_chain_factory.cc @@ -34,6 +34,18 @@ class ProtocolDetectionJsonRpcCallbacks mcp_callbacks_.onError(error); } + // Propagate the per-message origin built by the JSON-RPC filter so the + // application layer can reply on the connection the message arrived on. + void onRequestWithContext(const jsonrpc::Request& request, + MessageDispatchContext& context) override { + mcp_callbacks_.onRequestWithContext(request, context); + } + + void onNotificationWithContext(const jsonrpc::Notification& notification, + MessageDispatchContext& context) override { + mcp_callbacks_.onNotificationWithContext(notification, context); + } + private: McpProtocolCallbacks& mcp_callbacks_; }; diff --git a/src/filter/stdio_filter_chain_factory.cc b/src/filter/stdio_filter_chain_factory.cc index 994e0064c..f8adcd3fe 100644 --- a/src/filter/stdio_filter_chain_factory.cc +++ b/src/filter/stdio_filter_chain_factory.cc @@ -28,6 +28,18 @@ class DirectJsonRpcCallbacks : public JsonRpcProtocolFilter::MessageHandler { mcp_callbacks_.onNotification(notification); } + // Propagate the per-message origin built by the JSON-RPC filter so the + // application layer can reply on the connection the message arrived on. + void onRequestWithContext(const jsonrpc::Request& request, + MessageDispatchContext& context) override { + mcp_callbacks_.onRequestWithContext(request, context); + } + + void onNotificationWithContext(const jsonrpc::Notification& notification, + MessageDispatchContext& context) override { + mcp_callbacks_.onNotificationWithContext(notification, context); + } + void onResponse(const jsonrpc::Response& response) override { mcp_callbacks_.onResponse(response); } diff --git a/src/mcp_connection_manager.cc b/src/mcp_connection_manager.cc index 6b75169c5..15df01e5f 100644 --- a/src/mcp_connection_manager.cc +++ b/src/mcp_connection_manager.cc @@ -767,6 +767,22 @@ void McpConnectionManager::onRequest(const jsonrpc::Request& request) { } } +void McpConnectionManager::onRequestWithContext( + const jsonrpc::Request& request, MessageDispatchContext& context) { + if (protocol_callbacks_) { + protocol_callbacks_->onRequestWithContext(request, context); + } +} + +void McpConnectionManager::onNotificationWithContext( + const jsonrpc::Notification& notification, + MessageDispatchContext& context) { + if (protocol_callbacks_) { + protocol_callbacks_->onNotificationWithContext(notification, context); + } + // HTTP 202 response is sent by HttpSseJsonRpcProtocolFilter::onNotification +} + void McpConnectionManager::onNotification( const jsonrpc::Notification& notification) { if (protocol_callbacks_) { diff --git a/src/server/mcp_server.cc b/src/server/mcp_server.cc index 4373eab32..bc535ea2d 100644 --- a/src/server/mcp_server.cc +++ b/src/server/mcp_server.cc @@ -757,41 +757,87 @@ void McpServer::setupFilterChain(application::FilterChainBuilder& builder) { } } -SessionManager::SessionPtr McpServer::getOrCreateCurrentSession() { - // The transport session binding is strictly single-use: consume it here - // and clear it immediately, so it can only ever apply to the one message - // whose dispatch it preceded. The transport filter announces it (via - // onTransportSessionBound) right before onRequest/onNotification, but a - // producer that does NOT announce — stdio, or any non-SSE filter chain — - // must not inherit the previous message's id. Clearing on consume makes a - // stale binding unrepresentable rather than relying on every producer to - // remember to announce an empty id. Dispatch is synchronous per message - // on the dispatcher thread, so consume-and-clear is race-free. - std::string transport_session_id; - transport_session_id.swap(current_transport_session_id_); - +SessionManager::SessionPtr McpServer::getOrCreateSessionFor( + const MessageDispatchContext& context) { + // The context is constructed per message by the transport that parsed + // it and dies when the dispatch returns, so a stale binding is + // unrepresentable by construction — no consume-and-clear discipline + // needed, unlike the ambient announce-then-dispatch scheme this + // replaces. + // // Transport session id wins: for HTTP+SSE each request arrives on a // one-shot POST connection, so keying the session on the connection // would hand every request a fresh session and silently drop state // such as resource subscriptions. The SSE stream id is the identity that // actually spans the client's requests — and it is also what the push // path needs to find the client's SSE stream. + const std::string& transport_session_id = context.transportSessionId(); if (!transport_session_id.empty()) { return session_manager_->getOrCreateSessionByTransportId( transport_session_id); } // Connection-keyed fallback for transports where the connection is - // long-lived (stdio) or there is no transport session concept. - auto session = session_manager_->getSessionByConnection(current_connection_); + // long-lived (stdio) or there is no transport session concept. The + // origin comes from the message itself, so interleaved reads from + // concurrent connections each land in their own session. + auto session = + session_manager_->getSessionByConnection(context.originConnection()); if (!session) { - session = session_manager_->createSession(current_connection_); + session = session_manager_->createSession(context.originConnection()); } return session; } +// Reply path for messages that arrived through the context-free legacy +// hooks. There is no origin connection to reply on, so fall back to the +// first connected stdio manager — the historical degraded behavior — and +// fail loudly when there is none, instead of writing to an unrelated +// connection. +class McpServer::LegacyDispatchContext : public MessageDispatchContext { + public: + explicit LegacyDispatchContext(McpServer& server) : server_(server) {} + + network::Connection* originConnection() const override { return nullptr; } + + const std::string& transportSessionId() const override { + static const std::string empty; + return empty; + } + + VoidResult sendResponse(const jsonrpc::Response& response) override { + for (auto& conn_manager : server_.connection_managers_) { + if (conn_manager->isConnected()) { + return conn_manager->sendResponse(response); + } + } + Error err; + err.code = jsonrpc::INTERNAL_ERROR; + err.message = "no dispatch context and no connected transport"; + return makeVoidError(err); + } + + private: + McpServer& server_; +}; + // McpProtocolCallbacks overrides void McpServer::onRequest(const jsonrpc::Request& request) { + // Context-free legacy entry: the producer did not say where the message + // came from. Every in-tree transport dispatches through + // onRequestWithContext; reaching this path means an external producer + // has not been migrated, so route replies through the degraded legacy + // fallback rather than guessing at a connection. + GOPHER_LOG_WARN( + "Request '{}' dispatched without origin context; session and reply " + "routing degraded to the legacy transport fallback", + request.method); + LegacyDispatchContext context(*this); + onRequestWithContext(request, context); +} + +void McpServer::onRequestWithContext(const jsonrpc::Request& request, + MessageDispatchContext& context) { GOPHER_LOG_DEBUG("McpServer::onRequest called with method: {}", request.method); @@ -813,38 +859,25 @@ void McpServer::onRequest(const jsonrpc::Request& request) { } // Resolve the session for this request: transport session id first - // (durable across HTTP+SSE POST connections), connection identity as + // (durable across HTTP+SSE POST connections), origin connection as // fallback. - auto session = getOrCreateCurrentSession(); + auto session = getOrCreateSessionFor(context); if (session) { pending_req->session_id = session->getId(); } if (!session) { - // Max sessions reached + // Max sessions reached. Reply on the requester's own return path — + // the context pins it to the connection the request arrived on. server_stats_.requests_failed++; auto response = jsonrpc::Response::make_error( request.id, Error(jsonrpc::INTERNAL_ERROR, "Max sessions reached")); - // Send response through appropriate mechanism - // For HTTP connections, use filter chain; for stdio, use connection manager - GOPHER_LOG_DEBUG("Sending error response for max sessions"); - - // Send response through the current connection (for TCP/HTTP connections) - // Following production pattern: thread-local connection context - if (current_connection_) { - filter::HttpSseFilterChainFactory::sendHttpResponse(response, - *current_connection_); - } - - // Also try connection managers (for stdio connections) - for (auto& conn_manager : connection_managers_) { - if (conn_manager->isConnected()) { - GOPHER_LOG_DEBUG("Found connected manager, sending response"); - conn_manager->sendResponse(response); - break; - } + auto send_result = context.sendResponse(response); + if (holds_alternative(send_result)) { + GOPHER_LOG_ERROR("Failed to send max-sessions error response: {}", + get(send_result).message); } return; } @@ -901,36 +934,24 @@ void McpServer::onRequest(const jsonrpc::Request& request) { } } - // Send response through appropriate channel - // Following proper architecture: use filter chain for HTTP, connection - // manager for stdio + // Send the response along the request's own return path. The context + // pins the origin connection and its transport's framing, so concurrent + // connections cannot cross-wire replies and a single dispatch can never + // fan out to an unrelated transport. A failed send is surfaced instead + // of silently dropped. GOPHER_LOG_DEBUG("Sending response for request id: {}", holds_alternative(request.id) ? get(request.id) : std::to_string(get(request.id))); - // Send response through the current connection (for TCP/HTTP connections) - // Following production pattern: server sends JSON-RPC, filter handles HTTP - if (current_connection_) { - // Convert response to JSON and send through connection - // The filter chain will handle HTTP protocol wrapping - auto json_val = json::to_json(response); - std::string json_str = json_val.toString(); - - OwnedBuffer response_buffer; - response_buffer.add(json_str); - - // Write JSON-RPC response - HTTP filter will wrap it - current_connection_->write(response_buffer, false); - } - - // Also try connection managers (for stdio transport) - // This is the legacy path for non-HTTP transports - for (auto& conn_manager : connection_managers_) { - if (conn_manager->isConnected()) { - conn_manager->sendResponse(response); - break; - } + auto send_result = context.sendResponse(response); + if (holds_alternative(send_result)) { + server_stats_.errors_total++; + GOPHER_LOG_ERROR("Failed to send response for request id {}: {}", + holds_alternative(request.id) + ? get(request.id) + : std::to_string(get(request.id)), + get(send_result).message); } // Remove request from pending list @@ -945,12 +966,24 @@ void McpServer::onRequest(const jsonrpc::Request& request) { } void McpServer::onNotification(const jsonrpc::Notification& notification) { + // Context-free legacy entry; see onRequest for why this is degraded. + GOPHER_LOG_WARN( + "Notification '{}' dispatched without origin context; session " + "resolution degraded to the legacy transport fallback", + notification.method); + LegacyDispatchContext context(*this); + onNotificationWithContext(notification, context); +} + +void McpServer::onNotificationWithContext( + const jsonrpc::Notification& notification, + MessageDispatchContext& context) { // Handle notification in dispatcher context server_stats_.notifications_total++; // Resolve the session the same way requests do, so a notification sent // on a fresh POST connection still lands in the subscriber's session. - auto session = getOrCreateCurrentSession(); + auto session = getOrCreateSessionFor(context); if (!session) { return; // Can't process notification without session } @@ -1070,8 +1103,7 @@ void McpServer::onConnectionLifecycleEvent(network::Connection* connection, http_sse_factory_->sseRegistry().removeConnection(connection); } - // Tear down session state keyed by the actual closing connection, not a - // global current_connection_ which may point at a different session. + // Tear down session state keyed by the actual closing connection. // Transport-keyed sessions (HTTP+SSE) are untouched here by design: they // are created with a null connection, live across many short POST // connections, and are released by the SSE registry's session-closed @@ -1112,10 +1144,6 @@ void McpServer::onConnectionLifecycleEvent(network::Connection* connection, lifecycle_callbacks_.erase(cb_it); } } - - if (current_connection_ == connection) { - current_connection_ = nullptr; - } } void McpServer::onError(const Error& error) { @@ -1567,9 +1595,6 @@ void McpServer::onNewConnection(network::ConnectionPtr&& connection) { connection_sessions_[conn_ptr] = session; } - // Set current connection for request processing context - current_connection_ = conn_ptr; - // Update connection count. // // We bump the public stats here rather than wait for a Connected event diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 0d429398c..609fa2bf5 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -1668,6 +1668,17 @@ target_link_libraries(test_server_notification_delivery ) add_test(NAME ServerNotificationDeliveryTest COMMAND test_server_notification_delivery) +add_executable(test_mcp_server_request_routing integration/test_mcp_server_request_routing.cc) +target_link_libraries(test_mcp_server_request_routing + gopher-mcp + gopher-mcp-event + gtest + gtest_main + gmock + Threads::Threads +) +add_test(NAME McpServerRequestRoutingTest COMMAND test_mcp_server_request_routing) + add_executable(test_mcp_server_connection_lifecycle integration/test_mcp_server_connection_lifecycle.cc) target_link_libraries(test_mcp_server_connection_lifecycle gopher-mcp diff --git a/tests/filter/test_json_rpc_protocol_filter.cc b/tests/filter/test_json_rpc_protocol_filter.cc index 21fbf8197..024016166 100644 --- a/tests/filter/test_json_rpc_protocol_filter.cc +++ b/tests/filter/test_json_rpc_protocol_filter.cc @@ -39,6 +39,49 @@ class MockJsonRpcCallbacks : public JsonRpcProtocolFilter::MessageHandler { MOCK_METHOD(void, onProtocolError, (const Error&), (override)); }; +/** + * Handler that records the dispatch context contract: the per-message + * context must arrive together with the message (context path, not the + * context-free fallback), report the filter's actual origin, and refuse + * to pretend a reply was sent when there is no connection to send it on. + */ +class ContextCapturingHandler : public JsonRpcProtocolFilter::MessageHandler { + public: + void onRequest(const jsonrpc::Request&) override { legacy_requests_++; } + void onNotification(const jsonrpc::Notification&) override { + legacy_notifications_++; + } + void onResponse(const jsonrpc::Response&) override {} + void onProtocolError(const Error&) override {} + + void onRequestWithContext(const jsonrpc::Request& request, + MessageDispatchContext& context) override { + context_requests_++; + last_origin_ = context.originConnection(); + last_transport_session_id_ = context.transportSessionId(); + // Try to reply while the context is live; record whether the filter + // surfaced the missing reply path as an error. + auto result = context.sendResponse(jsonrpc::Response::make_error( + request.id, Error(jsonrpc::INTERNAL_ERROR, "test reply"))); + last_send_failed_ = holds_alternative(result); + } + + void onNotificationWithContext(const jsonrpc::Notification&, + MessageDispatchContext& context) override { + context_notifications_++; + last_origin_ = context.originConnection(); + last_transport_session_id_ = context.transportSessionId(); + } + + int legacy_requests_{0}; + int legacy_notifications_{0}; + int context_requests_{0}; + int context_notifications_{0}; + network::Connection* last_origin_{reinterpret_cast(1)}; + std::string last_transport_session_id_{""}; + bool last_send_failed_{false}; +}; + /** * Test fixture for JsonRpcProtocolFilter using real I/O */ @@ -306,6 +349,65 @@ TEST_F(JsonRpcProtocolFilterTest, WriteFilterAddsFraming) { }); } +/** + * A parsed request must be dispatched through the context-carrying hook + * with a context that travels with the message. This filter is not wired + * to a connection, so the contract under test is the honest degraded + * form: no origin, no transport session id, and a reply attempt that + * FAILS instead of silently succeeding (the pre-context encoder returned + * success while writing nothing). + */ +TEST_F(JsonRpcProtocolFilterTest, RequestDispatchCarriesContext) { + auto handler = std::make_unique(); + std::unique_ptr filter; + executeInDispatcher([&]() { + filter = std::make_unique(*handler, *dispatcher_, + /*is_server=*/true); + OwnedBuffer buffer; + buffer.add( + std::string(R"({"jsonrpc":"2.0","id":7,"method":"test.method"})") + + "\n"); + filter->onData(buffer, false); + }); + + EXPECT_EQ(handler->context_requests_, 1) + << "request must dispatch through the context-carrying hook"; + EXPECT_EQ(handler->legacy_requests_, 0) + << "context-free fallback must not run when a context was built"; + EXPECT_EQ(handler->last_origin_, nullptr) + << "unwired filter must report no origin, not a stale pointer"; + EXPECT_EQ(handler->last_transport_session_id_, "") + << "bare JSON-RPC chain has no transport session concept"; + EXPECT_TRUE(handler->last_send_failed_) + << "sendResponse with no connection must surface an error, not " + "silently drop the reply"; + + executeInDispatcher([&]() { filter.reset(); }); +} + +/** + * Same contract for notifications: context path, not the fallback. + */ +TEST_F(JsonRpcProtocolFilterTest, NotificationDispatchCarriesContext) { + auto handler = std::make_unique(); + std::unique_ptr filter; + executeInDispatcher([&]() { + filter = std::make_unique(*handler, *dispatcher_, + /*is_server=*/true); + OwnedBuffer buffer; + buffer.add(std::string(R"({"jsonrpc":"2.0","method":"notify.method"})") + + "\n"); + filter->onData(buffer, false); + }); + + EXPECT_EQ(handler->context_notifications_, 1); + EXPECT_EQ(handler->legacy_notifications_, 0); + EXPECT_EQ(handler->last_origin_, nullptr); + EXPECT_EQ(handler->last_transport_session_id_, ""); + + executeInDispatcher([&]() { filter.reset(); }); +} + /** * Integration test with real connection * NOTE: Disabled due to issues with real I/O test infrastructure diff --git a/tests/integration/test_http_sse_filter_server_mode.cc b/tests/integration/test_http_sse_filter_server_mode.cc index 49ea32742..a474b0f3c 100644 --- a/tests/integration/test_http_sse_filter_server_mode.cc +++ b/tests/integration/test_http_sse_filter_server_mode.cc @@ -42,25 +42,30 @@ using namespace std::chrono_literals; class ServerModeCallbacks : public McpProtocolCallbacks { public: void onRequest(const jsonrpc::Request& req) override { + // Context-free fallback: nothing traveled with the message. If the + // filter ever regresses to this path, the seeded sentinel below + // surfaces in binding_at_request_ and fails the assertions. requests_.push_back(req); - // Capture which transport session id was in effect when this request - // was dispatched — the filter's contract is to announce it (possibly - // empty) immediately beforehand, per message. binding_at_request_.push_back(current_binding_); } + void onRequestWithContext(const jsonrpc::Request& req, + MessageDispatchContext& context) override { + requests_.push_back(req); + // Capture the transport session id that traveled WITH this message — + // the filter's contract is to build a fresh context per dispatched + // message (possibly with an empty id). + binding_at_request_.push_back(context.transportSessionId()); + } void onNotification(const jsonrpc::Notification&) override {} void onResponse(const jsonrpc::Response&) override {} void onConnectionEvent(network::ConnectionEvent) override {} void onError(const Error&) override { error_count_++; } void onMessageEndpoint(const std::string&) override {} bool sendHttpPost(const std::string&) override { return true; } - void onTransportSessionBound(const std::string& id) override { - current_binding_ = id; - } std::vector requests_; std::vector binding_at_request_; - std::string current_binding_{""}; + std::string current_binding_{""}; int error_count_{0}; }; @@ -304,7 +309,9 @@ TEST_F(ServerModeFilterTest, PlainPost_AnnouncesEmptyBinding) { network::IoHandlePtr peer; std::shared_ptr factory; - // Pretend a callback POST on another connection just dispatched. + // Seed the context-free sentinel: if dispatch ever bypasses the + // per-message context, this value leaks into binding_at_request_ and + // the empty-id assertion below fails. callbacks.current_binding_ = "client_from_previous_connection"; executeInDispatcher([&]() { diff --git a/tests/integration/test_mcp_server_request_routing.cc b/tests/integration/test_mcp_server_request_routing.cc new file mode 100644 index 000000000..f0e756675 --- /dev/null +++ b/tests/integration/test_mcp_server_request_routing.cc @@ -0,0 +1,321 @@ +/** + * Integration test: responses are routed by the request's own dispatch + * context, not by ambient most-recently-accepted-connection state. + * + * Historical failure modes pinned here (all real with the ambient + * current-connection scheme, all fixed by the per-message dispatch + * context): + * + * 1. Cross-wired responses. The ambient pointer was stamped at accept + * time, so with two live connections (accept A, accept B) a request + * arriving on A was answered on B: client B received client A's + * response and A hung. The dispatch context pins the reply to the + * connection the request physically arrived on. + * + * 2. Silent drop after an unrelated close. Closing B nulled the ambient + * pointer even though A was alive and mid-flight, so A's next + * response was skipped without a log line. With the context, A's + * reply path is A's own connection regardless of what B does. + * + * The tests drive a real McpServer over real TCP sockets with plain HTTP + * POSTs (no SSE handshake), because the plain-HTTP path is exactly the one + * that used to depend on the ambient pointer for both session keying and + * the response write. + */ + +#include +#include +#include +#include +#include +#include + +#include + +#include "mcp/buffer.h" +#include "mcp/network/address.h" +#include "mcp/network/io_handle.h" +#include "mcp/network/socket_interface.h" +#include "mcp/server/mcp_server.h" +#include "mcp/types.h" + +namespace mcp { +namespace { + +using namespace std::chrono_literals; + +// Bind ephemeral 0 to get a port the kernel thinks is free, then let go +// of it and hand the number to the server. Same mild TOCTOU as the other +// integration tests -- accepted on a loopback test bed. +uint16_t pickEphemeralPort() { + auto& iface = network::socketInterface(); + + auto fd_result = + iface.socket(network::SocketType::Stream, network::Address::Type::Ip, + network::Address::IpVersion::v4); + if (!fd_result.ok()) { + throw std::runtime_error("pickEphemeralPort: socket() failed"); + } + + auto handle = iface.ioHandleForFd(*fd_result, /*socket_v6only=*/false); + handle->setBlocking(false); + + auto bind_addr = network::Address::parseInternetAddress("127.0.0.1", 0); + auto bind_result = handle->bind(bind_addr); + if (!bind_result.ok()) { + throw std::runtime_error("pickEphemeralPort: bind() failed"); + } + + auto local_addr_result = handle->localAddress(); + if (!local_addr_result.ok()) { + throw std::runtime_error("pickEphemeralPort: localAddress() failed"); + } + + const auto* ip = + dynamic_cast(local_addr_result->get()); + if (ip == nullptr) { + throw std::runtime_error("pickEphemeralPort: not an IP address"); + } + uint16_t port = ip->port(); + handle->close(); + return port; +} + +class McpServerRequestRoutingTest : public ::testing::Test { + protected: + void SetUp() override { + port_ = pickEphemeralPort(); + + server::McpServerConfig config; + config.server_name = "routing-test-server"; + config.server_version = "0.0.1"; + config.supported_transports = {TransportType::HttpSse}; + config.num_workers = 1; + + server_ = server::createMcpServer(config); + ASSERT_NE(server_, nullptr); + + const std::string listen_address = + "http://127.0.0.1:" + std::to_string(port_); + auto listen_result = server_->listen(listen_address); + ASSERT_TRUE(holds_alternative(listen_result)) + << "McpServer::listen failed"; + + server_thread_ = std::thread([this]() { server_->run(); }); + + ASSERT_TRUE(waitForListenerReady(port_, 5s)) + << "Server did not begin accepting on port " << port_; + } + + void TearDown() override { + if (server_) { + server_->shutdown(); + } + if (server_thread_.joinable()) { + server_thread_.join(); + } + server_.reset(); + } + + static bool waitForListenerReady(uint16_t port, + std::chrono::milliseconds budget) { + auto& iface = network::socketInterface(); + auto addr = network::Address::parseInternetAddress("127.0.0.1", port); + const auto deadline = std::chrono::steady_clock::now() + budget; + while (std::chrono::steady_clock::now() < deadline) { + auto fd_result = + iface.socket(network::SocketType::Stream, network::Address::Type::Ip, + network::Address::IpVersion::v4); + if (fd_result.ok()) { + auto handle = iface.ioHandleForFd(*fd_result, false); + handle->setBlocking(true); + auto connect_result = handle->connect(addr); + handle->close(); + if (connect_result.ok()) { + return true; + } + } + std::this_thread::sleep_for(25ms); + } + return false; + } + + network::IoHandlePtr openClient() { + auto& iface = network::socketInterface(); + auto fd_result = + iface.socket(network::SocketType::Stream, network::Address::Type::Ip, + network::Address::IpVersion::v4); + if (!fd_result.ok()) { + return nullptr; + } + auto handle = iface.ioHandleForFd(*fd_result, /*socket_v6only=*/false); + handle->setBlocking(true); + auto addr = network::Address::parseInternetAddress("127.0.0.1", port_); + auto connect_result = handle->connect(addr); + if (!connect_result.ok()) { + handle->close(); + return nullptr; + } + return handle; + } + + // POST a JSON-RPC body to the server's plain-HTTP RPC path on an + // already-open client socket. + static bool sendRpcPost(network::IoHandle& handle, const std::string& body) { + std::string request = + "POST /rpc HTTP/1.1\r\n" + "Host: localhost\r\n" + "Content-Type: application/json\r\n" + "Content-Length: " + + std::to_string(body.size()) + + "\r\n" + "\r\n" + + body; + OwnedBuffer out; + out.add(request); + return handle.write(out).ok(); + } + + // Accumulate whatever arrives on the socket until it contains `needle` + // or the budget elapses. Returns everything read either way. + static std::string readUntilContains(network::IoHandle& handle, + const std::string& needle, + std::chrono::milliseconds budget) { + handle.setBlocking(false); + std::string received; + const auto deadline = std::chrono::steady_clock::now() + budget; + while (std::chrono::steady_clock::now() < deadline) { + OwnedBuffer buf; + auto r = handle.read(buf, /*max_length=*/8192); + if (r.ok() && *r > 0) { + received += buf.toString(); + if (received.find(needle) != std::string::npos) { + return received; + } + } else if (r.ok() && *r == 0) { + return received; // EOF + } else if (!r.wouldBlock()) { + return received; // hard error; caller's assertion reports it + } + std::this_thread::sleep_for(10ms); + } + return received; + } + + uint16_t port_{0}; + std::unique_ptr server_; + std::thread server_thread_; +}; + +// Two live connections; the request goes in on the FIRST-accepted one. +// Under the ambient scheme the reply went out on the most recently +// accepted connection (B); under the dispatch context it must come back +// on A, and B must stay silent. +TEST_F(McpServerRequestRoutingTest, ResponseReturnsOnOriginConnection) { + auto client_a = openClient(); + ASSERT_NE(client_a, nullptr); + auto client_b = openClient(); + ASSERT_NE(client_b, nullptr); + + // Give the dispatcher a beat to accept both, in order, so B is the + // most recently accepted connection when A's request dispatches — + // exactly the interleaving the ambient pointer got wrong. + std::this_thread::sleep_for(100ms); + + ASSERT_TRUE( + sendRpcPost(*client_a, R"({"jsonrpc":"2.0","id":41,"method":"ping"})")); + + std::string on_a = readUntilContains(*client_a, "\"id\":41", 3s); + EXPECT_NE(on_a.find("\"id\":41"), std::string::npos) + << "Ping response did not return on the connection that sent the " + "request; got instead: " + << on_a; + + // B must not have received A's response. A short read window is + // enough: the response already proved the round trip completed. + std::string on_b = readUntilContains(*client_b, "\"id\":41", 200ms); + EXPECT_EQ(on_b.find("\"id\":41"), std::string::npos) + << "Another connection received this request's response: " << on_b; + + // And the reverse direction: B's own request comes back on B. + ASSERT_TRUE( + sendRpcPost(*client_b, R"({"jsonrpc":"2.0","id":42,"method":"ping"})")); + std::string on_b2 = readUntilContains(*client_b, "\"id\":42", 3s); + EXPECT_NE(on_b2.find("\"id\":42"), std::string::npos) + << "Second connection's response did not return to it; got: " << on_b2; + + client_a->close(); + client_b->close(); +} + +// Closing an unrelated connection must not affect another connection's +// request/response cycle. Under the ambient scheme, B's close nulled the +// shared pointer and A's response was silently skipped. +TEST_F(McpServerRequestRoutingTest, UnrelatedCloseDoesNotDropResponse) { + auto client_a = openClient(); + ASSERT_NE(client_a, nullptr); + auto client_b = openClient(); + ASSERT_NE(client_b, nullptr); + + // B is the most recently accepted connection; its close is the one the + // ambient pointer used to track. + std::this_thread::sleep_for(100ms); + client_b->close(); + client_b.reset(); + std::this_thread::sleep_for(100ms); + + ASSERT_TRUE( + sendRpcPost(*client_a, R"({"jsonrpc":"2.0","id":43,"method":"ping"})")); + + std::string on_a = readUntilContains(*client_a, "\"id\":43", 3s); + EXPECT_NE(on_a.find("\"id\":43"), std::string::npos) + << "Response was dropped after an unrelated connection closed; " + "received: " + << on_a; + + client_a->close(); +} + +// Several requests interleaved across two live connections, each keeping +// its own request/response pairing throughout — the steady-state version +// of the two tests above. +TEST_F(McpServerRequestRoutingTest, InterleavedRequestsKeepTheirConnections) { + auto client_a = openClient(); + ASSERT_NE(client_a, nullptr); + auto client_b = openClient(); + ASSERT_NE(client_b, nullptr); + std::this_thread::sleep_for(100ms); + + for (int i = 0; i < 3; ++i) { + const int id_a = 100 + i; + const int id_b = 200 + i; + + ASSERT_TRUE(sendRpcPost(*client_a, R"({"jsonrpc":"2.0","id":)" + + std::to_string(id_a) + + R"(,"method":"ping"})")); + ASSERT_TRUE(sendRpcPost(*client_b, R"({"jsonrpc":"2.0","id":)" + + std::to_string(id_b) + + R"(,"method":"ping"})")); + + std::string needle_a = "\"id\":" + std::to_string(id_a); + std::string needle_b = "\"id\":" + std::to_string(id_b); + + std::string on_a = readUntilContains(*client_a, needle_a, 3s); + EXPECT_NE(on_a.find(needle_a), std::string::npos) + << "round " << i << ": A's response missing on A; got: " << on_a; + EXPECT_EQ(on_a.find(needle_b), std::string::npos) + << "round " << i << ": B's response leaked onto A: " << on_a; + + std::string on_b = readUntilContains(*client_b, needle_b, 3s); + EXPECT_NE(on_b.find(needle_b), std::string::npos) + << "round " << i << ": B's response missing on B; got: " << on_b; + EXPECT_EQ(on_b.find(needle_a), std::string::npos) + << "round " << i << ": A's response leaked onto B: " << on_b; + } + + client_a->close(); + client_b->close(); +} + +} // namespace +} // namespace mcp