diff --git a/mcp/streamable.go b/mcp/streamable.go index e7777eb..7a40753 100644 --- a/mcp/streamable.go +++ b/mcp/streamable.go @@ -1234,7 +1234,14 @@ func (c *streamableClientConn) Write(ctx context.Context, msg jsonrpc.Message) e default: resp.Body.Close() - return fmt.Errorf("unsupported content type %q", ct) + switch msg := msg.(type) { + case *jsonrpc.Request: + return fmt.Errorf("unsupported content type %q when sending %q (status: %d)", ct, msg.Method, resp.StatusCode) + case *jsonrpc.Response: + return fmt.Errorf("unsupported content type %q when sending jsonrpc response #%d (status: %d)", ct, msg.ID, resp.StatusCode) + default: + panic("unreachable") + } } return nil } diff --git a/mcp/streamable_client_test.go b/mcp/streamable_client_test.go new file mode 100644 index 0000000..ee89df5 --- /dev/null +++ b/mcp/streamable_client_test.go @@ -0,0 +1,185 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +package mcp + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "sync" + "testing" + + "github.com/google/go-cmp/cmp" + "github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2" + "github.com/modelcontextprotocol/go-sdk/jsonrpc" +) + +type streamableRequestKey struct { + httpMethod string // http method + sessionID string // session ID header + jsonrpcMethod string // jsonrpc method, or "" for non-requests +} + +type header map[string]string + +type streamableResponse struct { + header header + status int // or http.StatusOK + body string // or "" + optional bool // if set, request need not be sent + wantProtocolVersion string // if "", unchecked + callback func() // if set, called after the request is handled +} + +type fakeResponses map[streamableRequestKey]*streamableResponse + +type fakeStreamableServer struct { + t *testing.T + responses fakeResponses + + callMu sync.Mutex + calls map[streamableRequestKey]int +} + +func (s *fakeStreamableServer) missingRequests() []streamableRequestKey { + s.callMu.Lock() + defer s.callMu.Unlock() + + var unused []streamableRequestKey + for k, resp := range s.responses { + if s.calls[k] == 0 && !resp.optional { + unused = append(unused, k) + } + } + return unused +} + +func (s *fakeStreamableServer) ServeHTTP(w http.ResponseWriter, req *http.Request) { + key := streamableRequestKey{ + httpMethod: req.Method, + sessionID: req.Header.Get(sessionIDHeader), + } + if req.Method == http.MethodPost { + body, err := io.ReadAll(req.Body) + if err != nil { + s.t.Errorf("failed to read body: %v", err) + http.Error(w, "failed to read body", http.StatusInternalServerError) + return + } + msg, err := jsonrpc.DecodeMessage(body) + if err != nil { + s.t.Errorf("invalid body: %v", err) + http.Error(w, "invalid body", http.StatusInternalServerError) + return + } + if r, ok := msg.(*jsonrpc.Request); ok { + key.jsonrpcMethod = r.Method + } + } + + s.callMu.Lock() + if s.calls == nil { + s.calls = make(map[streamableRequestKey]int) + } + s.calls[key]++ + s.callMu.Unlock() + + resp, ok := s.responses[key] + if !ok { + s.t.Errorf("missing response for %v", key) + http.Error(w, "no response", http.StatusInternalServerError) + return + } + if resp.callback != nil { + defer resp.callback() + } + for k, v := range resp.header { + w.Header().Set(k, v) + } + status := resp.status + if status == 0 { + status = http.StatusOK + } + w.WriteHeader(status) + + if v := req.Header.Get(protocolVersionHeader); v != resp.wantProtocolVersion && resp.wantProtocolVersion != "" { + s.t.Errorf("%v: bad protocol version header: got %q, want %q", key, v, resp.wantProtocolVersion) + } + w.Write([]byte(resp.body)) +} + +var ( + initResult = &InitializeResult{ + Capabilities: &ServerCapabilities{ + Completions: &CompletionCapabilities{}, + Logging: &LoggingCapabilities{}, + Tools: &ToolCapabilities{ListChanged: true}, + }, + ProtocolVersion: latestProtocolVersion, + ServerInfo: &Implementation{Name: "testServer", Version: "v1.0.0"}, + } + initResp = resp(1, initResult, nil) +) + +func jsonBody(t *testing.T, msg jsonrpc2.Message) string { + data, err := jsonrpc2.EncodeMessage(msg) + if err != nil { + t.Fatalf("encoding failed: %v", err) + } + return string(data) +} + +func TestStreamableClientTransportLifecycle(t *testing.T) { + ctx := context.Background() + + // The lifecycle test verifies various behavior of the streamable client + // initialization: + // - check that it can handle application/json responses + // - check that it sends the negotiated protocol version + fake := &fakeStreamableServer{ + t: t, + responses: fakeResponses{ + {"POST", "", methodInitialize}: { + header: header{ + "Content-Type": "application/json", + sessionIDHeader: "123", + }, + body: jsonBody(t, initResp), + }, + {"POST", "123", notificationInitialized}: { + status: http.StatusAccepted, + wantProtocolVersion: latestProtocolVersion, + }, + {"GET", "123", ""}: { + header: header{ + "Content-Type": "text/event-stream", + }, + optional: true, + wantProtocolVersion: latestProtocolVersion, + }, + {"DELETE", "123", ""}: {}, + }, + } + + httpServer := httptest.NewServer(fake) + defer httpServer.Close() + + transport := &StreamableClientTransport{Endpoint: httpServer.URL} + client := NewClient(testImpl, nil) + session, err := client.Connect(ctx, transport, nil) + if err != nil { + t.Fatalf("client.Connect() failed: %v", err) + } + if err := session.Close(); err != nil { + t.Errorf("closing session: %v", err) + } + if missing := fake.missingRequests(); len(missing) > 0 { + t.Errorf("did not receive expected requests: %v", missing) + } + if diff := cmp.Diff(initResult, session.state.InitializeResult); diff != "" { + t.Errorf("mismatch (-want, +got):\n%s", diff) + } +} diff --git a/mcp/streamable_test.go b/mcp/streamable_test.go index 0d171d8..2963a04 100644 --- a/mcp/streamable_test.go +++ b/mcp/streamable_test.go @@ -1035,77 +1035,6 @@ func mustMarshal(v any) json.RawMessage { return data } -func TestStreamableClientTransport(t *testing.T) { - // This test verifies various behavior of the streamable client transport: - // - check that it can handle application/json responses - // - check that it sends the negotiated protocol version - // - // TODO(rfindley): make this test more comprehensive, similar to - // [TestStreamableServerTransport]. - ctx := context.Background() - resp := func(id int64, result any, err error) *jsonrpc.Response { - return &jsonrpc.Response{ - ID: jsonrpc2.Int64ID(id), - Result: mustMarshal(result), - Error: err, - } - } - initResult := &InitializeResult{ - Capabilities: &ServerCapabilities{ - Completions: &CompletionCapabilities{}, - Logging: &LoggingCapabilities{}, - Tools: &ToolCapabilities{ListChanged: true}, - }, - ProtocolVersion: latestProtocolVersion, - ServerInfo: &Implementation{Name: "testServer", Version: "v1.0.0"}, - } - initResp := resp(1, initResult, nil) - - var reqN atomic.Int32 // request count - serverHandler := func(w http.ResponseWriter, r *http.Request) { - rN := reqN.Add(1) - - // TODO(rfindley): if the status code is NoContent or Accepted, we should - // probably be tolerant of when the content type is not application/json. - w.Header().Set("Content-Type", "application/json") - if rN == 1 { - data, err := jsonrpc2.EncodeMessage(initResp) - if err != nil { - t.Errorf("encoding failed: %v", err) - } - w.Header().Set("Mcp-Session-Id", "123") - w.Write(data) - } else { - if v := r.Header.Get(protocolVersionHeader); v != latestProtocolVersion { - t.Errorf("bad protocol version header: got %q, want %q", v, latestProtocolVersion) - } - } - } - - httpServer := httptest.NewServer(http.HandlerFunc(serverHandler)) - defer httpServer.Close() - - transport := &StreamableClientTransport{Endpoint: httpServer.URL} - client := NewClient(testImpl, nil) - session, err := client.Connect(ctx, transport, nil) - if err != nil { - t.Fatalf("client.Connect() failed: %v", err) - } - if err := session.Close(); err != nil { - t.Errorf("closing session: %v", err) - } - - if got, want := reqN.Load(), int32(3); got < want { - // Expect at least 3 requests: initialize, initialized, and DELETE. - // We may or may not observe the GET, depending on timing. - t.Errorf("unexpected number of requests: got %d, want at least %d", got, want) - } - - if diff := cmp.Diff(initResult, session.state.InitializeResult); diff != "" { - t.Errorf("mismatch (-want, +got):\n%s", diff) - } -} - func TestEventID(t *testing.T) { tests := []struct { sid StreamID