diff --git a/internal/jsonrpc2/conn.go b/internal/jsonrpc2/conn.go index 6bacfa7..537be47 100644 --- a/internal/jsonrpc2/conn.go +++ b/internal/jsonrpc2/conn.go @@ -739,7 +739,23 @@ func (c *Connection) processResult(from any, req *incomingRequest, result any, e // write is used by all things that write outgoing messages, including replies. // it makes sure that writes are atomic func (c *Connection) write(ctx context.Context, msg Message) error { - err := c.writer.Write(ctx, msg) + var err error + // Fail writes immediately if the connection is shutting down. + // + // TODO(rfindley): should we allow cancellation notifations through? It could + // be the case that writes can still succeed. + c.updateInFlight(func(s *inFlightState) { + err = s.shuttingDown(ErrServerClosing) + }) + if err == nil { + err = c.writer.Write(ctx, msg) + } + + // For rejected requests, we don't set the writeErr (which would break the + // connection). They can just be returned to the caller. + if errors.Is(err, ErrRejected) { + return err + } if err != nil && ctx.Err() == nil { // The call to Write failed, and since ctx.Err() is nil we can't attribute diff --git a/internal/jsonrpc2/wire.go b/internal/jsonrpc2/wire.go index b143dcd..8be2872 100644 --- a/internal/jsonrpc2/wire.go +++ b/internal/jsonrpc2/wire.go @@ -37,6 +37,17 @@ var ( ErrServerClosing = NewError(-32004, "server is closing") // ErrClientClosing is a dummy error returned for calls initiated while the client is closing. ErrClientClosing = NewError(-32003, "client is closing") + + // The following errors have special semantics for MCP transports + + // ErrRejected may be wrapped to return errors from calls to Writer.Write + // that signal that the request was rejected by the transport layer as + // invalid. + // + // Such failures do not indicate that the connection is broken, but rather + // should be returned to the caller to indicate that the specific request is + // invalid in the current context. + ErrRejected = NewError(-32004, "rejected by transport") ) const wireVersion = "2.0" diff --git a/mcp/streamable.go b/mcp/streamable.go index 526ee51..789c7a0 100644 --- a/mcp/streamable.go +++ b/mcp/streamable.go @@ -38,18 +38,29 @@ type StreamableHTTPHandler struct { getServer func(*http.Request) *Server opts StreamableHTTPOptions - mu sync.Mutex + mu sync.Mutex + // TODO: we should store the ServerSession along with the transport, because + // we need to cancel keepalive requests when closing the transport. transports map[string]*StreamableServerTransport // keyed by IDs (from Mcp-Session-Id header) } // StreamableHTTPOptions configures the StreamableHTTPHandler. type StreamableHTTPOptions struct { // GetSessionID provides the next session ID to use for an incoming request. + // If nil, a default randomly generated ID will be used. // - // If GetSessionID returns an empty string, the session is 'stateless', - // meaning it is not persisted and no session validation is performed. + // As a special case, if GetSessionID returns the empty string, the + // Mcp-Session-Id header will not be set. GetSessionID func() string + // Stateless controls whether the session is 'stateless'. + // + // A stateless server does not validate the Mcp-Session-Id header, and uses a + // temporary session with default initialization parameters. Any + // server->client request is rejected immediately as there's no way for the + // client to respond. + Stateless bool + // TODO: support session retention (?) // jsonResponse is forwarded to StreamableServerTransport.jsonResponse. @@ -118,36 +129,39 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque return } + sessionID := req.Header.Get(sessionIDHeader) var transport *StreamableServerTransport - if id := req.Header.Get(sessionIDHeader); id != "" { + if sessionID != "" { h.mu.Lock() - transport = h.transports[id] + transport, _ = h.transports[sessionID] h.mu.Unlock() - if transport == nil { + if transport == nil && !h.opts.Stateless { + // In stateless mode we allow a missing transport. + // + // A synthetic transport will be created below for the transient session. http.Error(w, "session not found", http.StatusNotFound) return } } - // TODO(rfindley): simplify the locking so that each request has only one - // critical section. if req.Method == http.MethodDelete { - if transport == nil { - // => Mcp-Session-Id was not set; else we'd have returned NotFound above. + if sessionID == "" { http.Error(w, "DELETE requires an Mcp-Session-Id header", http.StatusBadRequest) return } - h.mu.Lock() - delete(h.transports, transport.SessionID) - h.mu.Unlock() - transport.connection.Close() + if transport != nil { // transport may be nil in stateless mode + h.mu.Lock() + delete(h.transports, transport.SessionID) + h.mu.Unlock() + transport.connection.Close() + } w.WriteHeader(http.StatusNoContent) return } switch req.Method { case http.MethodPost, http.MethodGet: - if req.Method == http.MethodGet && transport == nil { + if req.Method == http.MethodGet && sessionID == "" { http.Error(w, "GET requires an active session", http.StatusMethodNotAllowed) return } @@ -164,37 +178,83 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque http.Error(w, "no server available", http.StatusBadRequest) return } - sessionID := h.opts.GetSessionID() - s := &StreamableServerTransport{SessionID: sessionID, jsonResponse: h.opts.jsonResponse} + if sessionID == "" { + // In stateless mode, sessionID may be nonempty even if there's no + // existing transport. + sessionID = h.opts.GetSessionID() + } + transport = &StreamableServerTransport{ + SessionID: sessionID, + Stateless: h.opts.Stateless, + jsonResponse: h.opts.jsonResponse, + } // To support stateless mode, we initialize the session with a default // state, so that it doesn't reject subsequent requests. var connectOpts *ServerSessionOptions - if sessionID == "" { + if h.opts.Stateless { + // Peek at the body to see if it is initialize or initialized. + // We want those to be handled as usual. + var hasInitialize, hasInitialized bool + { + // TODO: verify that this allows protocol version negotiation for + // stateless servers. + body, err := io.ReadAll(req.Body) + if err != nil { + http.Error(w, "failed to read body", http.StatusBadRequest) + return + } + req.Body.Close() + + // Reset the body so that it can be read later. + req.Body = io.NopCloser(bytes.NewBuffer(body)) + + msgs, _, err := readBatch(body) + if err == nil { + for _, msg := range msgs { + if req, ok := msg.(*jsonrpc.Request); ok { + switch req.Method { + case methodInitialize: + hasInitialize = true + case notificationInitialized: + hasInitialized = true + } + } + } + } + } + + // If we don't have InitializeParams or InitializedParams in the request, + // set the initial state to a default value. + state := new(ServerSessionState) + if !hasInitialize { + state.InitializeParams = new(InitializeParams) + } + if !hasInitialized { + state.InitializedParams = new(InitializedParams) + } connectOpts = &ServerSessionOptions{ - State: &ServerSessionState{ - InitializeParams: new(InitializeParams), - InitializedParams: new(InitializedParams), - }, + State: state, } } + // Pass req.Context() here, to allow middleware to add context values. // The context is detached in the jsonrpc2 library when handling the // long-running stream. - ss, err := server.Connect(req.Context(), s, connectOpts) + ss, err := server.Connect(req.Context(), transport, connectOpts) if err != nil { http.Error(w, "failed connection", http.StatusInternalServerError) return } - if sessionID == "" { + if h.opts.Stateless { // Stateless mode: close the session when the request exits. defer ss.Close() // close the fake session after handling the request } else { + // Otherwise, save the transport so that it can be reused h.mu.Lock() - h.transports[s.SessionID] = s + h.transports[transport.SessionID] = transport h.mu.Unlock() } - transport = s } transport.ServeHTTP(w, req) @@ -225,6 +285,13 @@ type StreamableServerTransport struct { // generator to produce one, as with [crypto/rand.Text].) SessionID string + // Stateless controls whether the eventstore is 'Stateless'. Servers sessions + // connected to a stateless transport are disallowed from making outgoing + // requests. + // + // See also [StreamableHTTPOptions.Stateless]. + Stateless bool + // Storage for events, to enable stream resumption. // If nil, a [MemoryEventStore] with the default maximum size will be used. EventStore EventStore @@ -265,6 +332,7 @@ func (t *StreamableServerTransport) Connect(context.Context) (Connection, error) } t.connection = &streamableServerConn{ sessionID: t.SessionID, + stateless: t.Stateless, eventStore: t.EventStore, jsonResponse: t.jsonResponse, incoming: make(chan jsonrpc.Message, 10), @@ -285,6 +353,7 @@ func (t *StreamableServerTransport) Connect(context.Context) (Connection, error) type streamableServerConn struct { sessionID string + stateless bool jsonResponse bool eventStore EventStore @@ -755,6 +824,10 @@ func (c *streamableServerConn) Read(ctx context.Context) (jsonrpc.Message, error // Write implements the [Connection] interface. func (c *streamableServerConn) Write(ctx context.Context, msg jsonrpc.Message) error { + if req, ok := msg.(*jsonrpc.Request); ok && req.ID.IsValid() && (c.stateless || c.sessionID == "") { + // Requests aren't possible with stateless servers, or when there's no session ID. + return fmt.Errorf("%w: stateless servers cannot make requests", jsonrpc2.ErrRejected) + } // Find the incoming request that this write relates to, if any. var forRequest jsonrpc.ID isResponse := false diff --git a/mcp/streamable_test.go b/mcp/streamable_test.go index 7c77938..8334bc0 100644 --- a/mcp/streamable_test.go +++ b/mcp/streamable_test.go @@ -748,7 +748,7 @@ func testStreamableHandler(t *testing.T, handler http.Handler, requests []stream if !request.ignoreResponse { transform := cmpopts.AcyclicTransformer("jsonrpcid", func(id jsonrpc.ID) any { return id.Raw() }) if diff := cmp.Diff(request.wantMessages, got, transform); diff != "" { - t.Errorf("received unexpected messages (-want +got):\n%s", diff) + t.Errorf("request #%d: received unexpected messages (-want +got):\n%s", i, diff) } } sessionID.CompareAndSwap("", gotSessionID) @@ -996,19 +996,18 @@ func TestEventID(t *testing.T) { } func TestStreamableStateless(t *testing.T) { - // This version of sayHi doesn't make a ping request (we can't respond to + // This version of sayHi expects // that request from our client). sayHi := func(ctx context.Context, req *ServerRequest[*CallToolParamsFor[hiParams]]) (*CallToolResult, error) { + if err := req.Session.Ping(ctx, nil); err == nil { + // ping should fail, but not break the connection + t.Errorf("ping succeeded unexpectedly") + } return &CallToolResult{Content: []Content{&TextContent{Text: "hi " + req.Params.Arguments.Name}}}, nil } server := NewServer(testImpl, nil) AddTool(server, &Tool{Name: "greet", Description: "say hi"}, sayHi) - // Test stateless mode. - handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, &StreamableHTTPOptions{ - GetSessionID: func() string { return "" }, - }) - requests := []streamableRequest{ { method: "POST", @@ -1028,7 +1027,74 @@ func TestStreamableStateless(t *testing.T) { }, wantSessionID: false, }, + { + method: "POST", + wantStatusCode: http.StatusOK, + messages: []jsonrpc.Message{ + req(2, "tools/call", &CallToolParams{Name: "greet", Arguments: hiParams{Name: "foo"}}), + }, + wantMessages: []jsonrpc.Message{ + resp(2, &CallToolResult{Content: []Content{&TextContent{Text: "hi foo"}}}, nil), + }, + wantSessionID: false, + }, } - testStreamableHandler(t, handler, requests) + testClientCompatibility := func(t *testing.T, handler http.Handler) { + ctx := context.Background() + httpServer := httptest.NewServer(handler) + defer httpServer.Close() + cs, err := NewClient(testImpl, nil).Connect(ctx, &StreamableClientTransport{Endpoint: httpServer.URL}, nil) + if err != nil { + t.Fatal(err) + } + res, err := cs.CallTool(ctx, &CallToolParams{Name: "greet", Arguments: hiParams{Name: "bar"}}) + if err != nil { + t.Fatal(err) + } + if got, want := textContent(t, res), "hi bar"; got != want { + t.Errorf("Result = %q, want %q", got, want) + } + } + + handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, &StreamableHTTPOptions{ + GetSessionID: func() string { return "" }, + Stateless: true, + }) + + // Test the default stateless mode. + t.Run("stateless", func(t *testing.T) { + testStreamableHandler(t, handler, requests) + testClientCompatibility(t, handler) + }) + + // Test a "distributed" variant of stateless mode, where it has non-empty + // session IDs, but is otherwise stateless. + // + // This can be used by tools to look up application state preserved across + // subsequent requests. + for i, req := range requests { + // Now, we want a session for all requests. + req.wantSessionID = true + requests[i] = req + } + distributableHandler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, &StreamableHTTPOptions{ + Stateless: true, + }) + t.Run("distributed", func(t *testing.T) { + testStreamableHandler(t, distributableHandler, requests) + testClientCompatibility(t, handler) + }) +} + +func textContent(t *testing.T, res *CallToolResult) string { + t.Helper() + if len(res.Content) != 1 { + t.Fatalf("len(Content) = %d, want 1", len(res.Content)) + } + text, ok := res.Content[0].(*TextContent) + if !ok { + t.Fatalf("Content[0] is %T, want *TextContent", res.Content[0]) + } + return text.Text } diff --git a/mcp/transport.go b/mcp/transport.go index 6d25de3..8018910 100644 --- a/mcp/transport.go +++ b/mcp/transport.go @@ -40,8 +40,8 @@ type Transport interface { type Connection interface { // Read reads the next message to process off the connection. // - // Read need not be safe for concurrent use: Read is called in a - // concurrency-safe manner by the JSON-RPC library. + // Connections must allow Read to be called concurrently with Close. In + // particular, calling Close should unblock a Read waiting for input. Read(context.Context) (jsonrpc.Message, error) // Write writes a new message to the connection.