Route replies and key sessions by a per-message dispatch context

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_;

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_;
 };

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

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_;
 };

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;

include/mcp/message_dispatch_context.h

@@ -0,0 +1,81 @@
+#pragma once
+
+#include <string>
+
+#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

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<std::mutex> 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

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={}",