c68ad9a4e6When connecting to remote stateless servers (such as GitHub's MCP server), optional streams like subscriptions/listen may be rejected by middleware with a plain-text HTTP 404 before reaching a JSON-RPC handler. Previously, checkResponse misclassified these 404s as lost stateful sessions (ErrSessionMissing) and tore down the transport connection because the error didn't wrap jsonrpc2.ErrRejected. This change aligns with the MCP 2026-07-28 Streamable HTTP specification: 1. Conditionally check for an active session ID before returning ErrSessionMissing, since a 404 on a stateless connection cannot mean a session was lost. 2. Wrap subscriptions/listen failures with ErrRejected (matching server/discover) so the jsonrpc2 layer doesn't permanently break the transport when encountering non-compliant plain-text 404s. --- I have manually tested this change against GitHub's hosted MCP server and confirmed that it fixes https://github.com/github/github-mcp-server/issues/3129
2031 lines
68 KiB
Go
2031 lines
68 KiB
Go
// 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"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/google/go-cmp/cmp"
|
|
"github.com/modelcontextprotocol/go-sdk/auth"
|
|
"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
|
|
"github.com/modelcontextprotocol/go-sdk/jsonrpc"
|
|
"golang.org/x/oauth2"
|
|
)
|
|
|
|
type streamableRequestKey struct {
|
|
httpMethod string // http method
|
|
sessionID string // session ID header
|
|
jsonrpcMethod string // jsonrpc method, or "" for non-requests
|
|
lastEventID string // Last-Event-ID header
|
|
}
|
|
|
|
type header map[string]string
|
|
|
|
// TODO: replace body and status fields with responseFunc; add helpers to reduce duplication.
|
|
type streamableResponse struct {
|
|
header header // response headers
|
|
status int // or http.StatusOK; ignored if responseFunc is set
|
|
body string // or ""; ignored if responseFunc is set
|
|
responseFunc func(r *jsonrpc.Request) (string, int) // if set, overrides body and status
|
|
optional bool // if set, request need not be sent
|
|
wantProtocolVersion string // if "", unchecked
|
|
done chan struct{} // if set, receive from this channel before terminating the request
|
|
}
|
|
|
|
type fakeResponses map[streamableRequestKey]*streamableResponse
|
|
|
|
type fakeStreamableServer struct {
|
|
t *testing.T
|
|
responses fakeResponses
|
|
|
|
calledMu sync.Mutex
|
|
called map[streamableRequestKey]bool
|
|
}
|
|
|
|
func (s *fakeStreamableServer) missingRequests() []streamableRequestKey {
|
|
s.calledMu.Lock()
|
|
defer s.calledMu.Unlock()
|
|
|
|
var unused []streamableRequestKey
|
|
for k, resp := range s.responses {
|
|
if !s.called[k] && !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),
|
|
lastEventID: req.Header.Get(lastEventIDHeader),
|
|
}
|
|
var jsonrpcReq *jsonrpc.Request
|
|
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
|
|
jsonrpcReq = r
|
|
}
|
|
}
|
|
|
|
s.calledMu.Lock()
|
|
if s.called == nil {
|
|
s.called = make(map[streamableRequestKey]bool)
|
|
}
|
|
s.called[key] = true
|
|
s.calledMu.Unlock()
|
|
|
|
resp, ok := s.responses[key]
|
|
if !ok {
|
|
if key.jsonrpcMethod == "server/discover" {
|
|
// Return MethodNotFound to trigger fallback to legacy initialize.
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(http.StatusOK)
|
|
respErr := &jsonrpc.Error{
|
|
Code: jsonrpc.CodeMethodNotFound,
|
|
Message: `method not found: "server/discover"`,
|
|
}
|
|
var id jsonrpc.ID
|
|
if jsonrpcReq != nil {
|
|
id = jsonrpcReq.ID
|
|
}
|
|
respMsg, _ := jsonrpc2.NewResponse(id, nil, respErr)
|
|
data, _ := jsonrpc2.EncodeMessage(respMsg)
|
|
w.Write(data)
|
|
return
|
|
}
|
|
s.t.Errorf("missing response for %v", key)
|
|
http.Error(w, "no response", http.StatusInternalServerError)
|
|
return
|
|
}
|
|
|
|
// Determine body and status, potentially using responseFunc for dynamic responses.
|
|
body := resp.body
|
|
status := resp.status
|
|
if resp.responseFunc != nil {
|
|
body, status = resp.responseFunc(jsonrpcReq)
|
|
}
|
|
if status == 0 {
|
|
status = http.StatusOK
|
|
}
|
|
|
|
for k, v := range resp.header {
|
|
w.Header().Set(k, v)
|
|
}
|
|
rc := http.NewResponseController(w)
|
|
w.WriteHeader(status)
|
|
rc.Flush() // flush response headers
|
|
|
|
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(body))
|
|
rc.Flush() // flush response
|
|
|
|
if resp.done != nil {
|
|
<-resp.done
|
|
}
|
|
}
|
|
|
|
var (
|
|
initResult = &InitializeResult{
|
|
Capabilities: &ServerCapabilities{
|
|
Completions: &CompletionCapabilities{},
|
|
Logging: &LoggingCapabilities{},
|
|
Tools: &ToolCapabilities{ListChanged: true},
|
|
},
|
|
// Pin negotiated version to 2025-11-25
|
|
ProtocolVersion: protocolVersion20251125,
|
|
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: protocolVersion20251125,
|
|
},
|
|
{"GET", "123", "", ""}: {
|
|
header: header{
|
|
"Content-Type": "text/event-stream",
|
|
},
|
|
wantProtocolVersion: protocolVersion20251125,
|
|
},
|
|
{"DELETE", "123", "", ""}: {},
|
|
},
|
|
}
|
|
|
|
httpServer := httptest.NewServer(fake)
|
|
defer httpServer.Close()
|
|
|
|
transport := &StreamableClientTransport{Endpoint: httpServer.URL}
|
|
client := NewClient(testImpl, nil)
|
|
// Pin to 2025-11-25: the fixture's canned initialize response uses
|
|
// hardcoded id=1, which only matches when initialize is the first
|
|
// request. Under 2026-07-28 the client probes server/discover first.
|
|
session, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125})
|
|
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)
|
|
}
|
|
}
|
|
|
|
func TestStreamableClientRedundantDelete(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: protocolVersion20251125,
|
|
},
|
|
{"GET", "123", "", ""}: {
|
|
status: http.StatusMethodNotAllowed,
|
|
},
|
|
{"POST", "123", methodListTools, ""}: {
|
|
status: http.StatusNotFound,
|
|
},
|
|
},
|
|
}
|
|
|
|
httpServer := httptest.NewServer(fake)
|
|
defer httpServer.Close()
|
|
|
|
transport := &StreamableClientTransport{Endpoint: httpServer.URL}
|
|
client := NewClient(testImpl, nil)
|
|
// Pin to 2025-11-25: the fixture's canned initialize response uses
|
|
// hardcoded id=1, which only matches when initialize is the first
|
|
// request. Under 2026-07-28 the client probes server/discover first.
|
|
session, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125})
|
|
if err != nil {
|
|
t.Fatalf("client.Connect() failed: %v", err)
|
|
}
|
|
_, err = session.ListTools(ctx, nil)
|
|
if err == nil {
|
|
t.Errorf("Listing tools: got nil error, want non-nil")
|
|
}
|
|
_ = session.Wait() // must not hang
|
|
if missing := fake.missingRequests(); len(missing) > 0 {
|
|
t.Errorf("did not receive expected requests: %v", missing)
|
|
}
|
|
}
|
|
|
|
// TestStreamableClientCloseUnresponsiveDelete checks that Close stays bounded
|
|
// when the peer accepts the session-termination DELETE but never answers it.
|
|
func TestStreamableClientCloseUnresponsiveDelete(t *testing.T) {
|
|
base := NewStreamableHTTPHandler(func(*http.Request) *Server {
|
|
return NewServer(testImpl, nil)
|
|
}, nil)
|
|
httpServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.Method == http.MethodDelete {
|
|
<-r.Context().Done()
|
|
return
|
|
}
|
|
base.ServeHTTP(w, r)
|
|
}))
|
|
t.Cleanup(func() {
|
|
// The DELETE handler only returns once its request is cancelled, so
|
|
// drop the connections before waiting for outstanding requests.
|
|
httpServer.CloseClientConnections()
|
|
httpServer.Close()
|
|
})
|
|
|
|
client := NewClient(testImpl, nil)
|
|
session, err := client.Connect(context.Background(),
|
|
&StreamableClientTransport{Endpoint: httpServer.URL}, nil)
|
|
if err != nil {
|
|
t.Fatalf("client.Connect() failed: %v", err)
|
|
}
|
|
if session.ID() == "" {
|
|
t.Fatal("empty session ID: Close would not send DELETE")
|
|
}
|
|
|
|
done := make(chan struct{})
|
|
go func() {
|
|
session.Close()
|
|
close(done)
|
|
}()
|
|
select {
|
|
case <-done:
|
|
case <-time.After(closeDeleteTimeout + 5*time.Second):
|
|
t.Fatal("Close() blocked on an unanswered DELETE")
|
|
}
|
|
}
|
|
|
|
func TestStreamableClientGETHandling(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
tests := []struct {
|
|
status int
|
|
wantErrorContaining string
|
|
contentType string
|
|
}{
|
|
{http.StatusOK, "", "text/event-stream"},
|
|
{http.StatusOK, "", "text/event-stream; charset=utf-8"},
|
|
{http.StatusMethodNotAllowed, "", "text/event-stream"},
|
|
//// The client error status code is not treated as an error in non-strict
|
|
//// mode.
|
|
{http.StatusNotFound, "", "text/event-stream"},
|
|
{http.StatusBadRequest, "", "text/event-stream"},
|
|
{http.StatusInternalServerError, "standalone SSE", "text/event-stream"},
|
|
{http.StatusOK, "", "text/html; charset=utf-8"},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
t.Run(fmt.Sprintf("status=%d content_type=%q", test.status, test.contentType), func(t *testing.T) {
|
|
fake := &fakeStreamableServer{
|
|
t: t,
|
|
responses: fakeResponses{
|
|
{"POST", "", methodInitialize, ""}: {
|
|
header: header{
|
|
"Content-Type": "application/json; charset=utf-8", // should ignore the charset
|
|
sessionIDHeader: "123",
|
|
},
|
|
body: jsonBody(t, initResp),
|
|
},
|
|
{"POST", "123", notificationInitialized, ""}: {
|
|
status: http.StatusAccepted,
|
|
wantProtocolVersion: protocolVersion20251125,
|
|
},
|
|
{"GET", "123", "", ""}: {
|
|
header: header{
|
|
"Content-Type": test.contentType,
|
|
},
|
|
status: test.status,
|
|
wantProtocolVersion: protocolVersion20251125,
|
|
},
|
|
{"DELETE", "123", "", ""}: {optional: true},
|
|
},
|
|
}
|
|
httpServer := httptest.NewServer(fake)
|
|
defer httpServer.Close()
|
|
|
|
transport := &StreamableClientTransport{Endpoint: httpServer.URL}
|
|
client := NewClient(testImpl, nil)
|
|
session, err := client.Connect(ctx, transport, &ClientSessionOptions{
|
|
ProtocolVersion: protocolVersion20251125,
|
|
})
|
|
if err == nil {
|
|
defer session.Close()
|
|
}
|
|
if test.wantErrorContaining != "" {
|
|
if err == nil {
|
|
t.Fatalf("Connect succeeded unexpectedly, want error containing %q", test.wantErrorContaining)
|
|
}
|
|
if got := err.Error(); !strings.Contains(got, test.wantErrorContaining) {
|
|
t.Errorf("Connect error = %q, want containing %q", got, test.wantErrorContaining)
|
|
}
|
|
} else if err != nil {
|
|
t.Fatalf("Connect failed: %v", err)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestStreamableClientStrictness(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
tests := []struct {
|
|
label string
|
|
strict bool
|
|
initializedStatus int
|
|
getStatus int
|
|
wantConnectError bool
|
|
}{
|
|
{"conformant server", true, http.StatusAccepted, http.StatusMethodNotAllowed, false},
|
|
{"strict initialized", true, http.StatusOK, http.StatusMethodNotAllowed, true},
|
|
{"unstrict initialized", false, http.StatusOK, http.StatusMethodNotAllowed, false},
|
|
{"strict GET", true, http.StatusAccepted, http.StatusNotFound, true},
|
|
// The client error status code is not treated as an error in non-strict
|
|
// mode.
|
|
{"unstrict GET on StatusNotFound", false, http.StatusOK, http.StatusNotFound, false},
|
|
{"unstrict GET on StatusBadRequest", false, http.StatusOK, http.StatusBadRequest, false},
|
|
{"GET on InternlServerError", false, http.StatusOK, http.StatusInternalServerError, true},
|
|
}
|
|
for _, test := range tests {
|
|
t.Run(test.label, func(t *testing.T) {
|
|
fake := &fakeStreamableServer{
|
|
t: t,
|
|
responses: fakeResponses{
|
|
{"POST", "", methodInitialize, ""}: {
|
|
header: header{
|
|
"Content-Type": "application/json",
|
|
sessionIDHeader: "123",
|
|
},
|
|
body: jsonBody(t, initResp),
|
|
},
|
|
{"POST", "123", notificationInitialized, ""}: {
|
|
status: test.initializedStatus,
|
|
wantProtocolVersion: protocolVersion20251125,
|
|
},
|
|
{"GET", "123", "", ""}: {
|
|
header: header{
|
|
"Content-Type": "text/event-stream",
|
|
},
|
|
status: test.getStatus,
|
|
wantProtocolVersion: protocolVersion20251125,
|
|
},
|
|
{"POST", "123", methodListTools, ""}: {
|
|
header: header{
|
|
"Content-Type": "application/json",
|
|
sessionIDHeader: "123",
|
|
},
|
|
body: jsonBody(t, resp(2, &ListToolsResult{Tools: []*Tool{}}, nil)),
|
|
optional: true,
|
|
},
|
|
{"DELETE", "123", "", ""}: {optional: true},
|
|
},
|
|
}
|
|
httpServer := httptest.NewServer(fake)
|
|
defer httpServer.Close()
|
|
|
|
transport := &StreamableClientTransport{Endpoint: httpServer.URL, strict: test.strict}
|
|
client := NewClient(testImpl, nil)
|
|
// Pin to 2025-11-25: the fixture's canned initialize response
|
|
// uses hardcoded id=1, which only matches when initialize is
|
|
// the first request. Under 2026-07-28 the client probes
|
|
// server/discover first.
|
|
session, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125})
|
|
if (err != nil) != test.wantConnectError {
|
|
t.Errorf("client.Connect() returned error %v; want error: %t", err, test.wantConnectError)
|
|
}
|
|
if err != nil {
|
|
return
|
|
}
|
|
_, err = session.ListTools(ctx, nil)
|
|
if err != nil {
|
|
t.Errorf("ListTools failed: %v", err)
|
|
}
|
|
if err := session.Close(); err != nil {
|
|
t.Errorf("closing session: %v", err)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestStreamableClientUnresumableRequest(t *testing.T) {
|
|
// This test verifies that the client fails fast when making a request that
|
|
// is unresumable, because it does not contain any events.
|
|
ctx := context.Background()
|
|
fake := &fakeStreamableServer{
|
|
t: t,
|
|
responses: fakeResponses{
|
|
{"POST", "", methodInitialize, ""}: {
|
|
header: header{
|
|
"Content-Type": "text/event-stream",
|
|
sessionIDHeader: "123",
|
|
},
|
|
body: "",
|
|
},
|
|
{"DELETE", "123", "", ""}: {optional: true},
|
|
},
|
|
}
|
|
httpServer := httptest.NewServer(fake)
|
|
defer httpServer.Close()
|
|
|
|
transport := &StreamableClientTransport{Endpoint: httpServer.URL}
|
|
client := NewClient(testImpl, nil)
|
|
cs, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125})
|
|
if err == nil {
|
|
cs.Close()
|
|
t.Fatalf("Connect succeeded unexpectedly")
|
|
}
|
|
// This may be a bit of a change detector, but for now check that we're
|
|
// actually exercising the early failure codepath.
|
|
msg := "terminated without response"
|
|
if !strings.Contains(err.Error(), msg) {
|
|
t.Errorf("Connect: got error %v, want containing %q", err, msg)
|
|
}
|
|
}
|
|
|
|
func TestStreamableClientResumption_Cancelled(t *testing.T) {
|
|
// This test verifies that the resumed requests are closed when their context
|
|
// is cancelled (issue #662).
|
|
|
|
// This test (unfortunately) relies on timing, so may have false positives.
|
|
//
|
|
// Set the reconnect initial delay to some small(ish) value so that the test
|
|
// doesn't take too long. But this value must be large enough that we mostly
|
|
// avoid races in the tests below, where one test cases is intended to be in
|
|
// between the initial attempt and first reconnection.
|
|
//
|
|
// For easier tuning (and debugging), factor out the tick size.
|
|
//
|
|
// TODO(#680): experiment with instead using synctest.
|
|
const tick = 10 * time.Millisecond
|
|
defer func(delay int64) {
|
|
reconnectInitialDelay.Store(delay)
|
|
}(reconnectInitialDelay.Load())
|
|
reconnectInitialDelay.Store(int64(2 * tick))
|
|
|
|
// The setup: terminate a request stream and make the resumed request hang
|
|
// indefinitely. CallTool should still exit when its context is canceled.
|
|
//
|
|
// This should work whether we're handling the initial request, waiting to
|
|
// retry, or handling the retry.
|
|
//
|
|
// Furthermore, closing the client connection should not hang, because there
|
|
// should be no ongoing requests.
|
|
|
|
tests := []struct {
|
|
label string
|
|
cancelAfter time.Duration
|
|
}{
|
|
{"in process", 1 * tick}, // cancel while the request is being handled
|
|
// initial request terminates at 2 ticks (see below)
|
|
{"awaiting retry", 3 * tick}, // cancel in-between first and second attempt
|
|
// retry starts at 4 ticks (=2+2)
|
|
{"in retry", 5 * tick}, // cancel while second attempt is hanging
|
|
}
|
|
|
|
for _, test := range tests {
|
|
t.Run(test.label, func(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
// done will be closed when the test exits: used to simulate requests that
|
|
// hang indefinitely.
|
|
initialRequestDone := make(chan struct{}) // closed below
|
|
allDone := make(chan struct{})
|
|
|
|
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: protocolVersion20251125,
|
|
},
|
|
{"GET", "123", "", ""}: {
|
|
header: header{
|
|
"Content-Type": "text/event-stream",
|
|
},
|
|
status: http.StatusMethodNotAllowed, // don't allow the standalone stream
|
|
},
|
|
{"POST", "123", methodCallTool, ""}: {
|
|
header: header{
|
|
"Content-Type": "text/event-stream",
|
|
},
|
|
status: http.StatusOK,
|
|
body: `id: 1
|
|
data: { "jsonrpc": "2.0", "method": "notifications/message", "params": { "level": "error", "data": "bad" } }
|
|
|
|
`,
|
|
done: initialRequestDone,
|
|
},
|
|
{"POST", "123", methodListTools, ""}: {
|
|
header: header{
|
|
"Content-Type": "application/json",
|
|
sessionIDHeader: "123",
|
|
},
|
|
body: jsonBody(t, resp(3, &ListToolsResult{Tools: []*Tool{}}, nil)),
|
|
},
|
|
{"GET", "123", "", "1"}: {
|
|
header: header{
|
|
"Content-Type": "text/event-stream",
|
|
},
|
|
status: http.StatusOK,
|
|
done: allDone, // hang indefinitely
|
|
},
|
|
{"POST", "123", notificationCancelled, ""}: {status: http.StatusAccepted},
|
|
{"DELETE", "123", "", ""}: {optional: true},
|
|
},
|
|
}
|
|
httpServer := httptest.NewServer(fake)
|
|
defer httpServer.Close()
|
|
defer close(allDone) // must be deferred *after* httpServer.Close, to avoid deadlock
|
|
|
|
transport := &StreamableClientTransport{Endpoint: httpServer.URL}
|
|
client := NewClient(testImpl, nil)
|
|
// Pin to 2025-11-25: the fixture's canned initialize response
|
|
// uses hardcoded id=1, which only matches when initialize is
|
|
// the first request. Under 2026-07-28 the client probes
|
|
// server/discover first.
|
|
cs, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer cs.Close() // ensure the session is closed, though we're also closing below
|
|
|
|
// start the timer on the initial request
|
|
go func() {
|
|
<-time.After(2 * tick)
|
|
close(initialRequestDone)
|
|
}()
|
|
|
|
// start the timer on the call cancellation
|
|
timeoutCtx, cancel := context.WithTimeout(ctx, test.cancelAfter)
|
|
defer cancel()
|
|
|
|
go func() {
|
|
<-timeoutCtx.Done()
|
|
}()
|
|
|
|
if _, err := cs.CallTool(timeoutCtx, &CallToolParams{
|
|
Name: "tool",
|
|
}); err == nil {
|
|
t.Errorf("CallTool succeeded unexpectedly")
|
|
}
|
|
|
|
// ...but cancellation should not break the session.
|
|
// Check that an arbitrary request succeeds.
|
|
if _, err := cs.ListTools(ctx, nil); err != nil {
|
|
t.Errorf("ListTools failed after cancellation")
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestStreamableClientTransientErrors verifies that transient errors (timeouts,
|
|
// 5xx HTTP status codes) do not permanently break the client connection.
|
|
// This tests the fix for issue #683.
|
|
func TestStreamableClientTransientErrors(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
tests := []struct {
|
|
transientStatus int // HTTP status to return for the transient call
|
|
wantCallError bool // whether the transient call should error
|
|
wantSessionBroken bool // whether the session should be broken after
|
|
wantErrorContains string // substring expected in error message
|
|
}{
|
|
{
|
|
transientStatus: http.StatusServiceUnavailable,
|
|
wantCallError: true,
|
|
wantSessionBroken: false,
|
|
wantErrorContains: "Service Unavailable",
|
|
},
|
|
{
|
|
transientStatus: http.StatusBadGateway,
|
|
wantCallError: true,
|
|
wantSessionBroken: false,
|
|
wantErrorContains: "Bad Gateway",
|
|
},
|
|
{
|
|
transientStatus: http.StatusGatewayTimeout,
|
|
wantCallError: true,
|
|
wantSessionBroken: false,
|
|
wantErrorContains: "Gateway Timeout",
|
|
},
|
|
{
|
|
transientStatus: http.StatusTooManyRequests,
|
|
wantCallError: true,
|
|
wantSessionBroken: false,
|
|
wantErrorContains: "Too Many Requests",
|
|
},
|
|
{
|
|
transientStatus: http.StatusUnauthorized,
|
|
wantCallError: true,
|
|
wantSessionBroken: true,
|
|
wantErrorContains: "Unauthorized",
|
|
},
|
|
{
|
|
transientStatus: http.StatusNotFound,
|
|
wantCallError: true,
|
|
wantSessionBroken: true,
|
|
wantErrorContains: "not found", // NotFound has special handling
|
|
},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
t.Run(http.StatusText(test.transientStatus), func(t *testing.T) {
|
|
var returnedError atomic.Bool
|
|
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: protocolVersion20251125,
|
|
},
|
|
{"GET", "123", "", ""}: {
|
|
status: http.StatusMethodNotAllowed,
|
|
},
|
|
{"POST", "123", methodListTools, ""}: {
|
|
header: header{
|
|
"Content-Type": "application/json",
|
|
sessionIDHeader: "123",
|
|
},
|
|
responseFunc: func(r *jsonrpc.Request) (string, int) {
|
|
// First call returns transient error, subsequent calls succeed.
|
|
if !returnedError.Swap(true) && test.transientStatus != 0 {
|
|
return "", test.transientStatus
|
|
}
|
|
return jsonBody(t, resp(r.ID.Raw().(int64), &ListToolsResult{Tools: []*Tool{}}, nil)), 0
|
|
},
|
|
optional: true,
|
|
},
|
|
{"DELETE", "123", "", ""}: {optional: true},
|
|
},
|
|
}
|
|
|
|
httpServer := httptest.NewServer(fake)
|
|
defer httpServer.Close()
|
|
|
|
transport := &StreamableClientTransport{Endpoint: httpServer.URL}
|
|
client := NewClient(testImpl, nil)
|
|
// Pin to 2025-11-25: the fixture's canned initialize response
|
|
// uses hardcoded id=1, which only matches when initialize is
|
|
// the first request. Under 2026-07-28 the client probes
|
|
// server/discover first.
|
|
session, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125})
|
|
if err != nil {
|
|
t.Fatalf("Connect failed: %v", err)
|
|
}
|
|
defer session.Close()
|
|
|
|
// First call: should trigger transient error.
|
|
_, err = session.ListTools(ctx, nil)
|
|
if test.wantCallError {
|
|
if err == nil {
|
|
t.Error("ListTools succeeded unexpectedly, want error")
|
|
} else if test.wantErrorContains != "" && !strings.Contains(err.Error(), test.wantErrorContains) {
|
|
t.Errorf("ListTools error = %q, want containing %q", err.Error(), test.wantErrorContains)
|
|
}
|
|
} else if err != nil {
|
|
t.Errorf("ListTools failed unexpectedly: %v", err)
|
|
}
|
|
|
|
// Second call: verifies whether the session is still usable.
|
|
_, err = session.ListTools(ctx, nil)
|
|
if test.wantSessionBroken {
|
|
if err == nil {
|
|
t.Error("second ListTools succeeded unexpectedly, want session broken")
|
|
}
|
|
} else {
|
|
if err != nil {
|
|
t.Errorf("second ListTools failed unexpectedly: %v (session should survive transient errors)", err)
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestStreamableClientRetryWithoutProgress verifies that the client fails after
|
|
// exceeding the retry limit when no progress is made (Last-Event-ID does not advance).
|
|
// This tests the fix for issue #679.
|
|
func TestStreamableClientRetryWithoutProgress(t *testing.T) {
|
|
// Speed up reconnection delays for testing.
|
|
const tick = 10 * time.Millisecond
|
|
defer func(delay int64) {
|
|
reconnectInitialDelay.Store(delay)
|
|
}(reconnectInitialDelay.Load())
|
|
reconnectInitialDelay.Store(int64(tick))
|
|
|
|
// Use the fakeStreamableServer pattern like other tests to avoid race conditions.
|
|
ctx := context.Background()
|
|
const maxRetries = 2
|
|
var retryCount atomic.Int32
|
|
|
|
fake := &fakeStreamableServer{
|
|
t: t,
|
|
responses: fakeResponses{
|
|
{"POST", "", methodInitialize, ""}: {
|
|
header: header{
|
|
"Content-Type": "application/json",
|
|
sessionIDHeader: "test-session",
|
|
},
|
|
body: jsonBody(t, initResp),
|
|
},
|
|
{"POST", "test-session", notificationInitialized, ""}: {
|
|
status: http.StatusAccepted,
|
|
wantProtocolVersion: protocolVersion20251125,
|
|
},
|
|
{"GET", "test-session", "", ""}: {
|
|
// Disable standalone SSE stream to simplify the test.
|
|
status: http.StatusMethodNotAllowed,
|
|
},
|
|
{"POST", "test-session", methodCallTool, ""}: {
|
|
header: header{
|
|
"Content-Type": "text/event-stream",
|
|
},
|
|
// Return SSE stream with fixed event ID.
|
|
body: `id: fixed_1
|
|
data: {"jsonrpc":"2.0","method":"notifications/message","params":{"level":"info","data":"test"}}
|
|
|
|
`,
|
|
},
|
|
// Resumption attempts with the same event ID (no progress).
|
|
{"GET", "test-session", "", "fixed_1"}: {
|
|
header: header{
|
|
"Content-Type": "text/event-stream",
|
|
},
|
|
responseFunc: func(r *jsonrpc.Request) (string, int) {
|
|
retryCount.Add(1)
|
|
// Return the same event ID - no progress.
|
|
return `id: fixed_1
|
|
data: {"jsonrpc":"2.0","method":"notifications/message","params":{"level":"info","data":"retry"}}
|
|
|
|
`, http.StatusOK
|
|
},
|
|
},
|
|
{"DELETE", "test-session", "", ""}: {optional: true},
|
|
},
|
|
}
|
|
|
|
httpServer := httptest.NewServer(fake)
|
|
defer httpServer.Close()
|
|
|
|
transport := &StreamableClientTransport{
|
|
Endpoint: httpServer.URL,
|
|
MaxRetries: maxRetries,
|
|
}
|
|
client := NewClient(testImpl, nil)
|
|
// Pin to 2025-11-25: the fixture's canned initialize response uses
|
|
// hardcoded id=1, which only matches when initialize is the first
|
|
// request. Under 2026-07-28 the client probes server/discover first.
|
|
session, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125})
|
|
if err != nil {
|
|
t.Fatalf("Connect failed: %v", err)
|
|
}
|
|
defer session.Close()
|
|
|
|
// Make a call that will trigger reconnections without progress.
|
|
_, err = session.CallTool(ctx, &CallToolParams{Name: "test"})
|
|
if err == nil {
|
|
t.Fatal("CallTool succeeded unexpectedly, want error due to exceeded retries")
|
|
}
|
|
|
|
// Check that the error mentions exceeding retries without progress.
|
|
wantErr := "exceeded"
|
|
if !strings.Contains(err.Error(), wantErr) {
|
|
t.Errorf("CallTool error = %q, want containing %q", err.Error(), wantErr)
|
|
}
|
|
|
|
// Verify that we actually retried the expected number of times.
|
|
// We expect maxRetries+1 attempts because we increment before checking the limit.
|
|
if got := retryCount.Load(); got != int32(maxRetries+1) {
|
|
t.Errorf("retry count = %d, want exactly %d", got, maxRetries+1)
|
|
}
|
|
}
|
|
|
|
func TestStreamableClientDisableStandaloneSSE(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
tests := []struct {
|
|
name string
|
|
disableStandaloneSSE bool
|
|
expectGETRequest bool
|
|
}{
|
|
{
|
|
name: "default behavior (standalone SSE enabled)",
|
|
disableStandaloneSSE: false,
|
|
expectGETRequest: true,
|
|
},
|
|
{
|
|
name: "standalone SSE disabled",
|
|
disableStandaloneSSE: true,
|
|
expectGETRequest: false,
|
|
},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
getRequestKey := streamableRequestKey{"GET", "123", "", ""}
|
|
|
|
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: protocolVersion20251125,
|
|
},
|
|
getRequestKey: {
|
|
header: header{
|
|
"Content-Type": "text/event-stream",
|
|
},
|
|
wantProtocolVersion: protocolVersion20251125,
|
|
optional: !test.expectGETRequest,
|
|
},
|
|
{"DELETE", "123", "", ""}: {
|
|
optional: true,
|
|
},
|
|
},
|
|
}
|
|
|
|
httpServer := httptest.NewServer(fake)
|
|
defer httpServer.Close()
|
|
|
|
transport := &StreamableClientTransport{
|
|
Endpoint: httpServer.URL,
|
|
DisableStandaloneSSE: test.disableStandaloneSSE,
|
|
}
|
|
client := NewClient(testImpl, nil)
|
|
// Pin to 2025-11-25: the fixture's canned initialize response
|
|
// uses hardcoded id=1, which only matches when initialize is
|
|
// the first request. Under 2026-07-28 the client probes
|
|
// server/discover first.
|
|
session, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125})
|
|
if err != nil {
|
|
t.Fatalf("client.Connect() failed: %v", err)
|
|
}
|
|
|
|
// Give some time for the standalone SSE connection to be established (if enabled)
|
|
time.Sleep(100 * time.Millisecond)
|
|
|
|
// Verify the connection state
|
|
streamableConn, ok := session.mcpConn.(*streamableClientConn)
|
|
if !ok {
|
|
t.Fatalf("Expected *streamableClientConn, got %T", session.mcpConn)
|
|
}
|
|
|
|
if got, want := streamableConn.disableStandaloneSSE, test.disableStandaloneSSE; got != want {
|
|
t.Errorf("disableStandaloneSSE field: got %v, want %v", got, want)
|
|
}
|
|
|
|
// Clean up
|
|
if err := session.Close(); err != nil {
|
|
t.Errorf("closing session: %v", err)
|
|
}
|
|
|
|
// Check if GET request was received
|
|
fake.calledMu.Lock()
|
|
getRequestReceived := false
|
|
if fake.called != nil {
|
|
getRequestReceived = fake.called[getRequestKey]
|
|
}
|
|
fake.calledMu.Unlock()
|
|
|
|
if got, want := getRequestReceived, test.expectGETRequest; got != want {
|
|
t.Errorf("GET request received: got %v, want %v", got, want)
|
|
}
|
|
|
|
// If we expected a GET request, verify it was actually received
|
|
if test.expectGETRequest {
|
|
if missing := fake.missingRequests(); len(missing) > 0 {
|
|
// Filter out optional requests
|
|
var requiredMissing []streamableRequestKey
|
|
for _, key := range missing {
|
|
if resp, ok := fake.responses[key]; ok && !resp.optional {
|
|
requiredMissing = append(requiredMissing, key)
|
|
}
|
|
}
|
|
if len(requiredMissing) > 0 {
|
|
t.Errorf("did not receive expected requests: %v", requiredMissing)
|
|
}
|
|
}
|
|
} else {
|
|
// If we didn't expect a GET request, verify it wasn't sent
|
|
if getRequestReceived {
|
|
t.Error("GET request was sent unexpectedly when DisableStandaloneSSE is true")
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
type mockOAuthHandler struct {
|
|
token *oauth2.Token
|
|
authorizeErr error
|
|
authorizeCalled bool
|
|
}
|
|
|
|
func (h *mockOAuthHandler) TokenSource(ctx context.Context) (oauth2.TokenSource, error) {
|
|
if h.token == nil {
|
|
return nil, nil
|
|
}
|
|
return oauth2.StaticTokenSource(h.token), nil
|
|
}
|
|
|
|
func (h *mockOAuthHandler) Authorize(ctx context.Context, req *http.Request, resp *http.Response) error {
|
|
h.authorizeCalled = true
|
|
return h.authorizeErr
|
|
}
|
|
|
|
func TestStreamableClientOAuth_AuthorizationHeader(t *testing.T) {
|
|
ctx := context.Background()
|
|
token := &oauth2.Token{AccessToken: "test-token"}
|
|
oauthHandler := &mockOAuthHandler{token: token}
|
|
|
|
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: protocolVersion20251125,
|
|
},
|
|
{"GET", "123", "", ""}: {
|
|
header: header{
|
|
"Content-Type": "text/event-stream",
|
|
},
|
|
},
|
|
{"DELETE", "123", "", ""}: {},
|
|
},
|
|
}
|
|
verifier := func(ctx context.Context, token string, req *http.Request) (*auth.TokenInfo, error) {
|
|
if token != "test-token" {
|
|
return nil, auth.ErrInvalidToken
|
|
}
|
|
return &auth.TokenInfo{Expiration: time.Now().Add(time.Hour)}, nil
|
|
}
|
|
httpServer := httptest.NewServer(auth.RequireBearerToken(verifier, nil)(fake))
|
|
t.Cleanup(httpServer.Close)
|
|
|
|
transport := &StreamableClientTransport{
|
|
Endpoint: httpServer.URL,
|
|
OAuthHandler: oauthHandler,
|
|
}
|
|
client := NewClient(testImpl, nil)
|
|
// Pin to 2025-11-25: the fixture's canned initialize response uses
|
|
// hardcoded id=1, which only matches when initialize is the first
|
|
// request. Under 2026-07-28 the client probes server/discover first.
|
|
session, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125})
|
|
if err != nil {
|
|
t.Fatalf("client.Connect() failed: %v", err)
|
|
}
|
|
session.Close()
|
|
}
|
|
|
|
func TestStreamableClientOAuth_401(t *testing.T) {
|
|
ctx := context.Background()
|
|
oauthHandler := &mockOAuthHandler{token: nil}
|
|
|
|
fake := &fakeStreamableServer{
|
|
t: t,
|
|
responses: fakeResponses{
|
|
{"POST", "", methodInitialize, ""}: {
|
|
header: header{
|
|
"Content-Type": "application/json",
|
|
sessionIDHeader: "123",
|
|
},
|
|
body: jsonBody(t, initResp),
|
|
},
|
|
},
|
|
}
|
|
verifier := func(ctx context.Context, token string, req *http.Request) (*auth.TokenInfo, error) {
|
|
// Accept any token.
|
|
return &auth.TokenInfo{Expiration: time.Now().Add(time.Hour)}, nil
|
|
}
|
|
httpServer := httptest.NewServer(auth.RequireBearerToken(verifier, nil)(fake))
|
|
t.Cleanup(httpServer.Close)
|
|
|
|
transport := &StreamableClientTransport{
|
|
Endpoint: httpServer.URL,
|
|
OAuthHandler: oauthHandler,
|
|
}
|
|
client := NewClient(testImpl, nil)
|
|
_, err := client.Connect(ctx, transport, nil)
|
|
if err == nil || !strings.Contains(err.Error(), "Unauthorized") {
|
|
t.Fatalf("client.Connect() error does not contain 'Unauthorized': %v", err)
|
|
}
|
|
|
|
if !oauthHandler.authorizeCalled {
|
|
t.Errorf("expected Authorize to be called")
|
|
}
|
|
}
|
|
|
|
// blockingCountingOAuthHandler is an OAuthHandler that blocks inside
|
|
// Authorize until the caller's context is cancelled, then returns a custom
|
|
// error that does NOT wrap context.Canceled. This mirrors real-world OAuth
|
|
// handlers that catch the cancellation internally and surface their own
|
|
// error type. The fix for #882 checks ctx.Err() directly rather than
|
|
// relying on the error from Authorize, so this must still trigger c.fail().
|
|
// It records how many times Authorize is invoked.
|
|
type blockingCountingOAuthHandler struct {
|
|
mu sync.Mutex
|
|
callCount int
|
|
}
|
|
|
|
func (h *blockingCountingOAuthHandler) TokenSource(ctx context.Context) (oauth2.TokenSource, error) {
|
|
return nil, nil
|
|
}
|
|
|
|
func (h *blockingCountingOAuthHandler) Authorize(ctx context.Context, req *http.Request, resp *http.Response) error {
|
|
h.mu.Lock()
|
|
h.callCount++
|
|
h.mu.Unlock()
|
|
// Block until the caller's context is cancelled, mirroring an
|
|
// interactive OAuth flow that the user has abandoned.
|
|
<-ctx.Done()
|
|
// Return a custom error that does not wrap context.Canceled, as a
|
|
// real-world handler might. The code under test must check ctx.Err()
|
|
// to detect the cancellation, not this error.
|
|
return fmt.Errorf("oauth flow interrupted")
|
|
}
|
|
|
|
func (h *blockingCountingOAuthHandler) Calls() int {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
return h.callCount
|
|
}
|
|
|
|
// TestStreamableClientOAuth_CancelledAuthorize_NoReprompt is a regression
|
|
// test for #882. When OAuthHandler.Authorize returns a context-cancelled
|
|
// error, the connection must enter a failed state so that the cancellation
|
|
// notification the call layer sends in response to ctx cancellation does
|
|
// not flow back through the same broken auth path and re-invoke Authorize.
|
|
func TestStreamableClientOAuth_CancelledAuthorize_NoReprompt(t *testing.T) {
|
|
handler := &blockingCountingOAuthHandler{}
|
|
|
|
fake := &fakeStreamableServer{
|
|
t: t,
|
|
responses: fakeResponses{
|
|
{"POST", "", methodInitialize, ""}: {
|
|
header: header{
|
|
"Content-Type": "application/json",
|
|
sessionIDHeader: "123",
|
|
},
|
|
body: jsonBody(t, initResp),
|
|
},
|
|
},
|
|
}
|
|
verifier := func(ctx context.Context, token string, req *http.Request) (*auth.TokenInfo, error) {
|
|
return &auth.TokenInfo{Expiration: time.Now().Add(time.Hour)}, nil
|
|
}
|
|
httpServer := httptest.NewServer(auth.RequireBearerToken(verifier, nil)(fake))
|
|
t.Cleanup(httpServer.Close)
|
|
|
|
transport := &StreamableClientTransport{
|
|
Endpoint: httpServer.URL,
|
|
OAuthHandler: handler,
|
|
}
|
|
client := NewClient(testImpl, nil)
|
|
|
|
// Use a context with a tight deadline so the cancellation path runs
|
|
// while the auth flow is in progress.
|
|
ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
|
|
defer cancel()
|
|
|
|
_, err := client.Connect(ctx, transport, nil)
|
|
if err == nil {
|
|
t.Fatal("expected client.Connect to fail")
|
|
}
|
|
|
|
// Give the cancellation Notify path a moment to (try to) run.
|
|
time.Sleep(50 * time.Millisecond)
|
|
|
|
// Authorize should be invoked exactly once. The bug in #882 caused
|
|
// it to be invoked a second time when the call layer sent the
|
|
// cancellation notification through the same auth-broken connection.
|
|
if got := handler.Calls(); got != 1 {
|
|
t.Errorf("expected Authorize to be called exactly 1 time, got %d", got)
|
|
}
|
|
}
|
|
|
|
func TestTokenInfo(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
// Create a server with a tool that returns TokenInfo.
|
|
tokenInfo := func(ctx context.Context, req *CallToolRequest, _ struct{}) (*CallToolResult, any, error) {
|
|
return &CallToolResult{Content: []Content{&TextContent{Text: fmt.Sprintf("%v", req.Extra.TokenInfo)}}}, nil, nil
|
|
}
|
|
server := NewServer(testImpl, nil)
|
|
AddTool(server, &Tool{Name: "tokenInfo", Description: "return token info"}, tokenInfo)
|
|
|
|
streamHandler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, nil)
|
|
verifier := func(ctx context.Context, token string, req *http.Request) (*auth.TokenInfo, error) {
|
|
if token != "test-token" {
|
|
return nil, auth.ErrInvalidToken
|
|
}
|
|
return &auth.TokenInfo{
|
|
Scopes: []string{"scope"},
|
|
// Expiration is far, far in the future.
|
|
Expiration: time.Date(5000, 1, 2, 3, 4, 5, 0, time.UTC),
|
|
}, nil
|
|
}
|
|
handler := auth.RequireBearerToken(verifier, nil)(streamHandler)
|
|
httpServer := httptest.NewServer(mustNotPanic(t, handler))
|
|
defer httpServer.Close()
|
|
|
|
transport := &StreamableClientTransport{
|
|
Endpoint: httpServer.URL,
|
|
OAuthHandler: &mockOAuthHandler{token: &oauth2.Token{AccessToken: "test-token"}},
|
|
}
|
|
client := NewClient(testImpl, nil)
|
|
session, err := client.Connect(ctx, transport, nil)
|
|
if err != nil {
|
|
t.Fatalf("client.Connect() failed: %v", err)
|
|
}
|
|
defer session.Close()
|
|
|
|
res, err := session.CallTool(ctx, &CallToolParams{Name: "tokenInfo"})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(res.Content) == 0 {
|
|
t.Fatal("missing content")
|
|
}
|
|
tc, ok := res.Content[0].(*TextContent)
|
|
if !ok {
|
|
t.Fatal("not TextContent")
|
|
}
|
|
if g, w := tc.Text, "&{[scope] 5000-01-02 03:04:05 +0000 UTC map[]}"; g != w {
|
|
t.Errorf("got %q, want %q", g, w)
|
|
}
|
|
}
|
|
|
|
// errTestAuthorizeFailed is a sentinel error returned by
|
|
// retrieveErrorOAuthHandler.Authorize().
|
|
var errTestAuthorizeFailed = errors.New("authorize intentionally failed for test")
|
|
|
|
// retrieveErrorOAuthHandler is a mock OAuthHandler that always returns
|
|
// an oauth2.RetrieveError from its TokenSource's Token() method.
|
|
type retrieveErrorOAuthHandler struct{}
|
|
|
|
func (h *retrieveErrorOAuthHandler) TokenSource(ctx context.Context) (oauth2.TokenSource, error) {
|
|
return h, nil
|
|
}
|
|
|
|
func (h *retrieveErrorOAuthHandler) Token() (*oauth2.Token, error) {
|
|
return nil, &oauth2.RetrieveError{
|
|
Response: &http.Response{StatusCode: http.StatusBadRequest},
|
|
Body: []byte("test retrieve error"),
|
|
ErrorCode: "invalid_grant",
|
|
}
|
|
}
|
|
|
|
func (h *retrieveErrorOAuthHandler) Authorize(ctx context.Context, req *http.Request, resp *http.Response) error {
|
|
return errTestAuthorizeFailed
|
|
}
|
|
|
|
// TestStreamableClientOAuth_RetrieveError verifies that an invalid_grant RetrieveError
|
|
// from the OAuth token source correctly skips sending Authorization header and relies on
|
|
// the server's 401 response to trigger the Authorize fallback flow.
|
|
func TestStreamableClientOAuth_RetrieveError(t *testing.T) {
|
|
ctx := context.Background()
|
|
oauthHandler := &retrieveErrorOAuthHandler{}
|
|
|
|
// Mock MCP server returns 401 Unauthorized to simulate a server rejecting
|
|
// the request that omitted the Authorization header.
|
|
httpServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.WriteHeader(http.StatusUnauthorized)
|
|
}))
|
|
t.Cleanup(httpServer.Close)
|
|
|
|
transport := &StreamableClientTransport{
|
|
Endpoint: httpServer.URL,
|
|
OAuthHandler: oauthHandler,
|
|
}
|
|
client := NewClient(testImpl, nil)
|
|
|
|
// Attempt to connect. The Connect call will trigger the initialization request,
|
|
// which will fail to retrieve the token and proceed without auth header, receive 401,
|
|
// and invoke Authorize().
|
|
_, err := client.Connect(ctx, transport, nil)
|
|
|
|
// Expect the connection to fail with the sentinel error, not the RetrieveError.
|
|
if !errors.Is(err, errTestAuthorizeFailed) {
|
|
t.Fatalf("client.Connect() error = %v, want %v", err, errTestAuthorizeFailed)
|
|
}
|
|
}
|
|
|
|
// discoverResult is the canned successful DiscoverResult returned by
|
|
// fakeStreamableServer setups in the tests below.
|
|
var discoverResult = &DiscoverResult{
|
|
Meta: Meta{
|
|
MetaKeyServerInfo: &Implementation{Name: "discoverServer", Version: "v1.0.0"},
|
|
},
|
|
SupportedVersions: []string{protocolVersion20260728},
|
|
Capabilities: &ServerCapabilities{
|
|
Tools: &ToolCapabilities{ListChanged: true},
|
|
},
|
|
Instructions: "test discover",
|
|
}
|
|
|
|
// TestStreamableClientConnect_DiscoverSuccess verifies that Client.Connect on
|
|
// the streamable transport:
|
|
// - sends a POST server/discover with Mcp-Protocol-Version: 2026-07-28 and
|
|
// the SEP-2575 per-request _meta triple in the body,
|
|
// - on a successful DiscoverResult, skips the legacy initialize handshake
|
|
// entirely, and
|
|
// - seeds ClientSession.InitializeResult() from the discover response,
|
|
// picking a mutually-supported protocol version.
|
|
func TestStreamableClientConnect_DiscoverSuccess(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
var (
|
|
gotDiscoverMu sync.Mutex
|
|
gotDiscover *jsonrpc.Request
|
|
)
|
|
|
|
fake := &fakeStreamableServer{
|
|
t: t,
|
|
responses: fakeResponses{
|
|
{"POST", "", methodDiscover, ""}: {
|
|
header: header{
|
|
"Content-Type": "application/json",
|
|
sessionIDHeader: "sess-1",
|
|
},
|
|
wantProtocolVersion: protocolVersion20260728,
|
|
responseFunc: func(r *jsonrpc.Request) (string, int) {
|
|
gotDiscoverMu.Lock()
|
|
gotDiscover = r
|
|
gotDiscoverMu.Unlock()
|
|
return jsonBody(t, resp(1, discoverResult, nil)), http.StatusOK
|
|
},
|
|
},
|
|
{"DELETE", "sess-1", "", ""}: {optional: true},
|
|
},
|
|
}
|
|
|
|
httpServer := httptest.NewServer(fake)
|
|
defer httpServer.Close()
|
|
|
|
transport := &StreamableClientTransport{Endpoint: httpServer.URL}
|
|
client := NewClient(testImpl, nil)
|
|
session, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728})
|
|
if err != nil {
|
|
t.Fatalf("Connect: %v", err)
|
|
}
|
|
defer session.Close()
|
|
|
|
if missing := fake.missingRequests(); len(missing) > 0 {
|
|
t.Errorf("missing expected requests: %v", missing)
|
|
}
|
|
|
|
gotDiscoverMu.Lock()
|
|
defer gotDiscoverMu.Unlock()
|
|
if gotDiscover == nil {
|
|
t.Fatal("server did not receive server/discover")
|
|
}
|
|
|
|
// Inspect the discover request body for the SEP-2575 _meta triple.
|
|
var body struct {
|
|
Meta map[string]any `json:"_meta"`
|
|
}
|
|
if err := json.Unmarshal(gotDiscover.Params, &body); err != nil {
|
|
t.Fatalf("decoding discover params: %v", err)
|
|
}
|
|
if v, _ := body.Meta[MetaKeyProtocolVersion].(string); v != protocolVersion20260728 {
|
|
t.Errorf("_meta[%s] = %q, want %q", MetaKeyProtocolVersion, v, protocolVersion20260728)
|
|
}
|
|
if _, ok := body.Meta[MetaKeyClientInfo]; !ok {
|
|
t.Errorf("_meta[%s] missing", MetaKeyClientInfo)
|
|
}
|
|
if _, ok := body.Meta[MetaKeyClientCapabilities]; !ok {
|
|
t.Errorf("_meta[%s] missing", MetaKeyClientCapabilities)
|
|
}
|
|
|
|
ir := session.InitializeResult()
|
|
if ir == nil {
|
|
t.Fatal("InitializeResult is nil after Connect")
|
|
}
|
|
if got, want := ir.ProtocolVersion, protocolVersion20260728; got != want {
|
|
t.Errorf("InitializeResult.ProtocolVersion = %q, want %q", got, want)
|
|
}
|
|
if ir.ServerInfo == nil || ir.ServerInfo.Name != "discoverServer" {
|
|
t.Errorf("InitializeResult.ServerInfo = %+v, want Name=discoverServer", ir.ServerInfo)
|
|
}
|
|
if ir.Instructions != "test discover" {
|
|
t.Errorf("InitializeResult.Instructions = %q, want %q", ir.Instructions, "test discover")
|
|
}
|
|
}
|
|
|
|
// TestStreamableClientConnSetMCPHeaders_ProtocolVersion covers
|
|
// streamableClientConn.setMCPHeaders' selection of the Mcp-Protocol-Version
|
|
// header value.
|
|
//
|
|
// Ordinarily initializedResult is populated by sessionUpdated, called
|
|
// through a type assertion to the unexported clientConnection interface
|
|
// (see Client.Connect). That assertion silently fails, leaving
|
|
// initializedResult nil for the life of the session, whenever the
|
|
// Connection returned by a Transport is wrapped by another type exposing
|
|
// only the exported Connection interface (a real pattern for transports
|
|
// that intercept traffic, e.g. to filter notifications): Go does not
|
|
// promote unexported interface methods across an embedded interface
|
|
// boundary. Every SEP-2575 (>= 2026-07-28) request already carries its own
|
|
// `_meta.protocolVersion` field, so setMCPHeaders falls back to reading it
|
|
// from the outgoing message when initializedResult is unset.
|
|
func TestStreamableClientConnSetMCPHeaders_ProtocolVersion(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
initializedResult *InitializeResult
|
|
msg jsonrpc.Message
|
|
ctxVersion string
|
|
want string
|
|
}{
|
|
{
|
|
name: "nil checked",
|
|
initializedResult: nil,
|
|
msg: (*jsonrpc.Request)(nil),
|
|
want: "",
|
|
},
|
|
{
|
|
name: "message meta wins when initializedResult unset",
|
|
initializedResult: nil,
|
|
msg: req(1, methodListTools, &ListToolsParams{Meta: Meta{MetaKeyProtocolVersion: protocolVersion20260728}}),
|
|
want: protocolVersion20260728,
|
|
},
|
|
{
|
|
name: "initializedResult used when message has no meta",
|
|
initializedResult: &InitializeResult{ProtocolVersion: protocolVersion20251125},
|
|
msg: req(1, methodListTools, &ListToolsParams{}),
|
|
want: protocolVersion20251125,
|
|
},
|
|
{
|
|
name: "initializedResult used for nil message (GET/DELETE)",
|
|
initializedResult: &InitializeResult{ProtocolVersion: protocolVersion20251125},
|
|
msg: nil,
|
|
want: protocolVersion20251125,
|
|
},
|
|
{
|
|
name: "message meta preferred over stale initializedResult",
|
|
initializedResult: &InitializeResult{ProtocolVersion: protocolVersion20251125},
|
|
msg: req(1, methodListTools, &ListToolsParams{Meta: Meta{MetaKeyProtocolVersion: protocolVersion20260728}}),
|
|
want: protocolVersion20260728,
|
|
},
|
|
{
|
|
name: "no header when neither source is set",
|
|
initializedResult: nil,
|
|
msg: req(1, methodListTools, &ListToolsParams{}),
|
|
want: "",
|
|
},
|
|
{
|
|
// A process that is both a server and a client (a proxy) carries the
|
|
// inbound request's version on the context of the outgoing call. The
|
|
// version this connection negotiated must win over it.
|
|
name: "initializedResult preferred over request context",
|
|
initializedResult: &InitializeResult{ProtocolVersion: protocolVersion20251125},
|
|
msg: req(1, methodListTools, &ListToolsParams{}),
|
|
ctxVersion: protocolVersion20260728,
|
|
want: protocolVersion20251125,
|
|
},
|
|
{
|
|
name: "request context used when initializedResult unset",
|
|
initializedResult: nil,
|
|
msg: req(1, methodListTools, &ListToolsParams{}),
|
|
ctxVersion: protocolVersion20260728,
|
|
want: protocolVersion20260728,
|
|
},
|
|
{
|
|
name: "message meta preferred over request context",
|
|
initializedResult: nil,
|
|
msg: req(1, methodListTools, &ListToolsParams{Meta: Meta{MetaKeyProtocolVersion: protocolVersion20251125}}),
|
|
ctxVersion: protocolVersion20260728,
|
|
want: protocolVersion20251125,
|
|
},
|
|
}
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
conn := &streamableClientConn{initializedResult: tt.initializedResult}
|
|
httpReq, err := http.NewRequest(http.MethodPost, "http://test.invalid", nil)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if tt.ctxVersion != "" {
|
|
httpReq = httpReq.WithContext(context.WithValue(httpReq.Context(), protocolVersionContextKey{}, tt.ctxVersion))
|
|
}
|
|
if err := conn.setMCPHeaders(httpReq, tt.msg); err != nil {
|
|
t.Fatalf("setMCPHeaders: %v", err)
|
|
}
|
|
if got := httpReq.Header.Get(protocolVersionHeader); got != tt.want {
|
|
t.Errorf("Mcp-Protocol-Version header = %q, want %q", got, tt.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestStreamableClientConnect_DiscoverMethodNotFound verifies that Client.Connect
|
|
// falls back to the legacy initialize handshake when the server responds to
|
|
// server/discover with a JSON-RPC "Method not found" error.
|
|
func TestStreamableClientConnect_DiscoverMethodNotFound(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
// Each request gets a fresh jsonrpc2 ID from the same client connection.
|
|
// Use responseFunc to echo the request's ID back so the client matches
|
|
// the response to the in-flight call regardless of ordering.
|
|
echoErr := func(err error) func(*jsonrpc.Request) (string, int) {
|
|
return func(r *jsonrpc.Request) (string, int) {
|
|
return jsonBody(t, &jsonrpc.Response{ID: r.ID, Error: err.(*jsonrpc.Error)}), http.StatusOK
|
|
}
|
|
}
|
|
echoResult := func(result any) func(*jsonrpc.Request) (string, int) {
|
|
return func(r *jsonrpc.Request) (string, int) {
|
|
return jsonBody(t, &jsonrpc.Response{ID: r.ID, Result: mustMarshal(result)}), http.StatusOK
|
|
}
|
|
}
|
|
|
|
fake := &fakeStreamableServer{
|
|
t: t,
|
|
responses: fakeResponses{
|
|
{"POST", "", methodDiscover, ""}: {
|
|
header: header{"Content-Type": "application/json"},
|
|
wantProtocolVersion: protocolVersion20260728,
|
|
responseFunc: echoErr(&jsonrpc.Error{
|
|
Code: jsonrpc.CodeMethodNotFound,
|
|
Message: "method not found",
|
|
}),
|
|
},
|
|
{"POST", "", methodInitialize, ""}: {
|
|
header: header{
|
|
"Content-Type": "application/json",
|
|
sessionIDHeader: "fallback",
|
|
},
|
|
responseFunc: echoResult(initResult),
|
|
},
|
|
{"POST", "fallback", notificationInitialized, ""}: {
|
|
status: http.StatusAccepted,
|
|
wantProtocolVersion: protocolVersion20251125,
|
|
},
|
|
{"GET", "fallback", "", ""}: {
|
|
header: header{"Content-Type": "text/event-stream"},
|
|
wantProtocolVersion: protocolVersion20251125,
|
|
optional: true,
|
|
},
|
|
{"DELETE", "fallback", "", ""}: {optional: true},
|
|
},
|
|
}
|
|
|
|
httpServer := httptest.NewServer(fake)
|
|
defer httpServer.Close()
|
|
|
|
transport := &StreamableClientTransport{Endpoint: httpServer.URL}
|
|
client := NewClient(testImpl, nil)
|
|
session, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125})
|
|
if err != nil {
|
|
t.Fatalf("Connect: %v", err)
|
|
}
|
|
defer session.Close()
|
|
|
|
// The fallback initialize explicitly requests protocolVersion20251125
|
|
// (see client.go), so the server negotiates that version.
|
|
if got := session.InitializeResult().ProtocolVersion; got != protocolVersion20251125 {
|
|
t.Errorf("InitializeResult.ProtocolVersion = %q, want %q (initialize fallback)", got, protocolVersion20251125)
|
|
}
|
|
}
|
|
|
|
// TestStreamableClientConnect_DiscoverUnsupportedVersion verifies that
|
|
// Client.Connect falls back to the legacy initialize handshake when the
|
|
// server responds to server/discover with the SEP-2575
|
|
// UnsupportedProtocolVersionError JSON-RPC code.
|
|
func TestStreamableClientConnect_DiscoverUnsupportedVersion(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
echoErr := func(err error) func(*jsonrpc.Request) (string, int) {
|
|
return func(r *jsonrpc.Request) (string, int) {
|
|
return jsonBody(t, &jsonrpc.Response{ID: r.ID, Error: err.(*jsonrpc.Error)}), http.StatusOK
|
|
}
|
|
}
|
|
echoResult := func(result any) func(*jsonrpc.Request) (string, int) {
|
|
return func(r *jsonrpc.Request) (string, int) {
|
|
return jsonBody(t, &jsonrpc.Response{ID: r.ID, Result: mustMarshal(result)}), http.StatusOK
|
|
}
|
|
}
|
|
|
|
fake := &fakeStreamableServer{
|
|
t: t,
|
|
responses: fakeResponses{
|
|
{"POST", "", methodDiscover, ""}: {
|
|
header: header{"Content-Type": "application/json"},
|
|
wantProtocolVersion: protocolVersion20260728,
|
|
responseFunc: echoErr(&jsonrpc.Error{
|
|
Code: CodeUnsupportedProtocolVersion,
|
|
Message: "unsupported protocol version",
|
|
}),
|
|
},
|
|
{"POST", "", methodInitialize, ""}: {
|
|
header: header{
|
|
"Content-Type": "application/json",
|
|
sessionIDHeader: "fallback",
|
|
},
|
|
responseFunc: echoResult(initResult),
|
|
},
|
|
{"POST", "fallback", notificationInitialized, ""}: {
|
|
status: http.StatusAccepted,
|
|
wantProtocolVersion: protocolVersion20251125,
|
|
},
|
|
{"GET", "fallback", "", ""}: {
|
|
header: header{"Content-Type": "text/event-stream"},
|
|
wantProtocolVersion: protocolVersion20251125,
|
|
optional: true,
|
|
},
|
|
{"DELETE", "fallback", "", ""}: {optional: true},
|
|
},
|
|
}
|
|
|
|
httpServer := httptest.NewServer(fake)
|
|
defer httpServer.Close()
|
|
|
|
transport := &StreamableClientTransport{Endpoint: httpServer.URL}
|
|
client := NewClient(testImpl, nil)
|
|
session, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125})
|
|
if err != nil {
|
|
t.Fatalf("Connect: %v", err)
|
|
}
|
|
defer session.Close()
|
|
|
|
// The fallback initialize explicitly requests protocolVersion20251125
|
|
// (see client.go), so the server negotiates that version.
|
|
if got := session.InitializeResult().ProtocolVersion; got != protocolVersion20251125 {
|
|
t.Errorf("InitializeResult.ProtocolVersion = %q, want %q (initialize fallback)", got, protocolVersion20251125)
|
|
}
|
|
}
|
|
|
|
// TestStreamableClientConnect_DiscoverMethodNotFoundVPre verifies that
|
|
// Client.Connect falls back to the legacy initialize handshake when a
|
|
// pre-SEP-2575 (vPre) server rejects server/discover.
|
|
func TestStreamableClientConnect_DiscoverMethodNotFoundVPre(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
echoResult := func(result any) func(*jsonrpc.Request) (string, int) {
|
|
return func(r *jsonrpc.Request) (string, int) {
|
|
return jsonBody(t, &jsonrpc.Response{ID: r.ID, Result: mustMarshal(result)}), http.StatusOK
|
|
}
|
|
}
|
|
|
|
fake := &fakeStreamableServer{
|
|
t: t,
|
|
responses: fakeResponses{
|
|
{"POST", "", methodDiscover, ""}: {
|
|
wantProtocolVersion: protocolVersion20260728,
|
|
// Reproduce the exact body a vPre server produces via
|
|
// http.Error(w, err.Error(), 400) where err comes from
|
|
// checkRequest. http.Error appends a trailing newline.
|
|
body: "JSON RPC not handled: \"server/discover\" unsupported\n",
|
|
status: http.StatusBadRequest,
|
|
header: header{"Content-Type": "text/plain; charset=utf-8"},
|
|
},
|
|
{"POST", "", methodInitialize, ""}: {
|
|
header: header{
|
|
"Content-Type": "application/json",
|
|
sessionIDHeader: "fallback",
|
|
},
|
|
responseFunc: echoResult(initResult),
|
|
},
|
|
{"POST", "fallback", notificationInitialized, ""}: {
|
|
status: http.StatusAccepted,
|
|
wantProtocolVersion: protocolVersion20251125,
|
|
},
|
|
{"DELETE", "fallback", "", ""}: {optional: true},
|
|
},
|
|
}
|
|
|
|
httpServer := httptest.NewServer(fake)
|
|
defer httpServer.Close()
|
|
|
|
transport := &StreamableClientTransport{
|
|
Endpoint: httpServer.URL,
|
|
DisableStandaloneSSE: true,
|
|
}
|
|
client := NewClient(testImpl, nil)
|
|
session, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728})
|
|
if err != nil {
|
|
t.Fatalf("Connect: %v", err)
|
|
}
|
|
defer session.Close()
|
|
|
|
// The fallback initialize explicitly requests protocolVersion20251125
|
|
// (see client.go), so the server negotiates that version.
|
|
if got := session.InitializeResult().ProtocolVersion; got != protocolVersion20251125 {
|
|
t.Errorf("InitializeResult.ProtocolVersion = %q, want %q (initialize fallback)", got, protocolVersion20251125)
|
|
}
|
|
}
|
|
|
|
// TestStreamableClientConnect_DiscoverUnsupportedProtocolVersion verifies that
|
|
// Client.Connect falls back to the legacy initialize handshake when a
|
|
// server rejects server/discover with a plain HTTP 400 containing "Unsupported protocol version".
|
|
func TestStreamableClientConnect_DiscoverUnsupportedVersionVPre(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
echoResult := func(result any) func(*jsonrpc.Request) (string, int) {
|
|
return func(r *jsonrpc.Request) (string, int) {
|
|
return jsonBody(t, &jsonrpc.Response{ID: r.ID, Result: mustMarshal(result)}), http.StatusOK
|
|
}
|
|
}
|
|
|
|
fake := &fakeStreamableServer{
|
|
t: t,
|
|
responses: fakeResponses{
|
|
{"POST", "", methodDiscover, ""}: {
|
|
wantProtocolVersion: protocolVersion20260728,
|
|
body: "Bad Request: Unsupported protocol version\n",
|
|
status: http.StatusBadRequest,
|
|
header: header{"Content-Type": "text/plain; charset=utf-8"},
|
|
},
|
|
{"POST", "", methodInitialize, ""}: {
|
|
header: header{
|
|
"Content-Type": "application/json",
|
|
sessionIDHeader: "fallback",
|
|
},
|
|
responseFunc: echoResult(initResult),
|
|
},
|
|
{"POST", "fallback", notificationInitialized, ""}: {
|
|
status: http.StatusAccepted,
|
|
wantProtocolVersion: protocolVersion20251125,
|
|
},
|
|
{"DELETE", "fallback", "", ""}: {optional: true},
|
|
},
|
|
}
|
|
|
|
httpServer := httptest.NewServer(fake)
|
|
defer httpServer.Close()
|
|
|
|
transport := &StreamableClientTransport{
|
|
Endpoint: httpServer.URL,
|
|
DisableStandaloneSSE: true,
|
|
}
|
|
client := NewClient(testImpl, nil)
|
|
session, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728})
|
|
if err != nil {
|
|
t.Fatalf("Connect: %v", err)
|
|
}
|
|
defer session.Close()
|
|
|
|
// The fallback initialize explicitly requests protocolVersion20251125
|
|
// (see client.go), so the server negotiates that version.
|
|
if got := session.InitializeResult().ProtocolVersion; got != protocolVersion20251125 {
|
|
t.Errorf("InitializeResult.ProtocolVersion = %q, want %q (initialize fallback)", got, protocolVersion20251125)
|
|
}
|
|
}
|
|
|
|
// TestStreamableClientConnect_DiscoverUnsupportedVersionNegotiation verifies that
|
|
// when Client.Connect over a streamable transport receives an
|
|
// UnsupportedProtocolVersion error containing Data.Supported, it negotiates a
|
|
// mutually supported version and retries server/discover.
|
|
func TestStreamableClientConnect_DiscoverUnsupportedVersionNegotiation(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
const unsupportedClientVersion = "2099-12-31"
|
|
|
|
var discoverCalls atomic.Int32
|
|
|
|
fake := &fakeStreamableServer{
|
|
t: t,
|
|
responses: fakeResponses{
|
|
{"POST", "", methodDiscover, ""}: {
|
|
header: header{
|
|
"Content-Type": "application/json",
|
|
},
|
|
responseFunc: func(r *jsonrpc.Request) (string, int) {
|
|
n := discoverCalls.Add(1)
|
|
if n == 1 {
|
|
data, _ := json.Marshal(UnsupportedProtocolVersionData{
|
|
Supported: []string{protocolVersion20260728},
|
|
})
|
|
respMsg := &jsonrpc.Response{
|
|
ID: r.ID,
|
|
Error: &jsonrpc.Error{
|
|
Code: CodeUnsupportedProtocolVersion,
|
|
Message: "unsupported protocol version",
|
|
Data: data,
|
|
},
|
|
}
|
|
return jsonBody(t, respMsg), http.StatusOK
|
|
}
|
|
return jsonBody(t, &jsonrpc.Response{ID: r.ID, Result: mustMarshal(discoverResult)}), http.StatusOK
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
httpServer := httptest.NewServer(fake)
|
|
defer httpServer.Close()
|
|
|
|
transport := &StreamableClientTransport{Endpoint: httpServer.URL}
|
|
client := NewClient(testImpl, nil)
|
|
session, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: unsupportedClientVersion})
|
|
if err != nil {
|
|
t.Fatalf("Connect: %v", err)
|
|
}
|
|
defer session.Close()
|
|
|
|
if got, want := discoverCalls.Load(), int32(2); got != want {
|
|
t.Errorf("discover call count = %d, want %d", got, want)
|
|
}
|
|
if got := session.InitializeResult().ProtocolVersion; got != protocolVersion20260728 {
|
|
t.Errorf("InitializeResult.ProtocolVersion = %q, want %q", got, protocolVersion20260728)
|
|
}
|
|
}
|
|
|
|
// TestStreamableClientHandlerErrorPropagation verifies that per-call
|
|
// handler-level HTTP errors carrying a JSON-RPC error body do not tear down
|
|
// the session, and that setting MCPGODEBUG=noprotocolerrorbody=1 restores the
|
|
// pre-fix behavior in which non-transient errors permanently failed the
|
|
// connection.
|
|
func TestStreamableClientHandlerErrorPropagation(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
// Build a JSON-RPC error response body for a given tools/call request.
|
|
jsonRPCErrorBody := func(id int64, code int64, msg string) string {
|
|
return jsonBody(t, resp(id, nil, &jsonrpc.Error{Code: code, Message: msg}))
|
|
}
|
|
|
|
tests := []struct {
|
|
name string
|
|
callStatus int // HTTP status returned for the tools/call
|
|
callBody string // response body for the tools/call ("" for empty; may be a JSON-RPC error)
|
|
disableBodyDecode bool // if true, sets noprotocolerrorbody="1" for the duration of the test
|
|
wantErrRejected bool // whether the returned error should match errors.Is(err, jsonrpc2.ErrRejected)
|
|
wantSessionAlive bool // whether a subsequent call on the same session should succeed
|
|
}{
|
|
{
|
|
// SEP-2575 §"Missing Required Capabilities": server returns
|
|
// -32021 at HTTP 400. Session must survive.
|
|
name: "400 with JSON-RPC error body (default)",
|
|
callStatus: http.StatusBadRequest,
|
|
callBody: jsonRPCErrorBody(1, CodeMissingRequiredClientCapabilities, "missing capability"),
|
|
wantErrRejected: true,
|
|
wantSessionAlive: true,
|
|
},
|
|
{
|
|
// Same server response, but the user opted out via
|
|
// MCPGODEBUG=noprotocolerrorbody=1. Pre-fix behavior: the
|
|
// session is torn down.
|
|
name: "400 with JSON-RPC error body (noprotocolerrorbody=1)",
|
|
callStatus: http.StatusBadRequest,
|
|
callBody: jsonRPCErrorBody(1, CodeMissingRequiredClientCapabilities, "missing capability"),
|
|
disableBodyDecode: true,
|
|
wantErrRejected: false,
|
|
wantSessionAlive: false,
|
|
},
|
|
{
|
|
// Legacy server returns plain-text 400 with no JSON-RPC
|
|
// body: cannot be classified as a per-call rejection, so the
|
|
// session is still torn down.
|
|
name: "400 with plain-text body",
|
|
callStatus: http.StatusBadRequest,
|
|
callBody: "Bad Request",
|
|
wantErrRejected: false,
|
|
wantSessionAlive: false,
|
|
},
|
|
{
|
|
// SEP-2575: server returns -32601 for an unimplemented
|
|
// method at HTTP 404. Must be treated as a per-call
|
|
// rejection, not a terminated session.
|
|
name: "404 with JSON-RPC MethodNotFound body",
|
|
callStatus: http.StatusNotFound,
|
|
callBody: jsonRPCErrorBody(1, jsonrpc.CodeMethodNotFound, `method not found: "tools/call"`),
|
|
wantErrRejected: true,
|
|
wantSessionAlive: true,
|
|
},
|
|
{
|
|
// A bare 404 with no JSON-RPC body means the session has
|
|
// been terminated. Session must not survive.
|
|
name: "404 with empty body (terminated session)",
|
|
callStatus: http.StatusNotFound,
|
|
callBody: "",
|
|
wantErrRejected: false,
|
|
wantSessionAlive: false,
|
|
},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
if test.disableBodyDecode {
|
|
prev := noprotocolerrorbody
|
|
noprotocolerrorbody = "1"
|
|
t.Cleanup(func() { noprotocolerrorbody = prev })
|
|
}
|
|
|
|
// Track how many tools/call requests we've served: first
|
|
// returns the error under test, subsequent calls succeed so
|
|
// we can observe whether the session is still alive.
|
|
var callsServed atomic.Int32
|
|
fake := &fakeStreamableServer{
|
|
t: t,
|
|
responses: fakeResponses{
|
|
{"POST", "", methodDiscover, ""}: {
|
|
header: header{
|
|
"Content-Type": "application/json",
|
|
},
|
|
wantProtocolVersion: protocolVersion20260728,
|
|
responseFunc: func(r *jsonrpc.Request) (string, int) {
|
|
return jsonBody(t, resp(r.ID.Raw().(int64), discoverResult, nil)), http.StatusOK
|
|
},
|
|
},
|
|
{"POST", "", "tools/call", ""}: {
|
|
header: header{"Content-Type": "application/json"},
|
|
responseFunc: func(r *jsonrpc.Request) (string, int) {
|
|
if callsServed.Add(1) == 1 {
|
|
return test.callBody, test.callStatus
|
|
}
|
|
return jsonBody(t, resp(r.ID.Raw().(int64), &CallToolResult{}, nil)), 0
|
|
},
|
|
optional: true,
|
|
},
|
|
},
|
|
}
|
|
|
|
httpServer := httptest.NewServer(fake)
|
|
defer httpServer.Close()
|
|
|
|
transport := &StreamableClientTransport{Endpoint: httpServer.URL}
|
|
client := NewClient(testImpl, nil)
|
|
session, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728})
|
|
if err != nil {
|
|
t.Fatalf("Connect: %v", err)
|
|
}
|
|
defer session.Close()
|
|
|
|
// First call: triggers the error under test.
|
|
_, err = session.CallTool(ctx, &CallToolParams{Name: "nonexistent"})
|
|
if err == nil {
|
|
t.Fatal("first CallTool succeeded unexpectedly, want error")
|
|
}
|
|
if got := errors.Is(err, jsonrpc2.ErrRejected); got != test.wantErrRejected {
|
|
t.Errorf("errors.Is(err, ErrRejected) = %v, want %v (err = %v)", got, test.wantErrRejected, err)
|
|
}
|
|
|
|
// Second call: verifies whether the session survived.
|
|
_, err = session.CallTool(ctx, &CallToolParams{Name: "nonexistent"})
|
|
if test.wantSessionAlive {
|
|
if err != nil {
|
|
t.Errorf("second CallTool failed: %v (session should survive per-call rejection)", err)
|
|
}
|
|
} else {
|
|
if err == nil {
|
|
t.Error("second CallTool succeeded unexpectedly, want session torn down")
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestStreamableClient_StatelessSubscriptionsListen404 verifies that in stateless mode
|
|
// (protocol 2026-07-28), when a server rejects subscriptions/listen with 404 (e.g. Copilot MCP server),
|
|
// the client connection survives and subsequent RPCs (such as tools/list) succeed.
|
|
func TestStreamableClient_StatelessSubscriptionsListen404(t *testing.T) {
|
|
ctx := t.Context()
|
|
|
|
var listenServed atomic.Bool
|
|
fake := &fakeStreamableServer{
|
|
t: t,
|
|
responses: fakeResponses{
|
|
{"POST", "", methodDiscover, ""}: {
|
|
header: header{
|
|
"Content-Type": "application/json",
|
|
},
|
|
wantProtocolVersion: protocolVersion20260728,
|
|
responseFunc: func(r *jsonrpc.Request) (string, int) {
|
|
return jsonBody(t, resp(r.ID.Raw().(int64), discoverResult, nil)), http.StatusOK
|
|
},
|
|
},
|
|
{"POST", "", methodSubscriptionsListen, ""}: {
|
|
header: header{"Content-Type": "text/plain"},
|
|
responseFunc: func(r *jsonrpc.Request) (string, int) {
|
|
listenServed.Store(true)
|
|
return "404 Not Found", http.StatusNotFound
|
|
},
|
|
optional: true,
|
|
},
|
|
{"POST", "", methodListTools, ""}: {
|
|
header: header{"Content-Type": "application/json"},
|
|
responseFunc: func(r *jsonrpc.Request) (string, int) {
|
|
return jsonBody(t, resp(r.ID.Raw().(int64), &ListToolsResult{Tools: []*Tool{}}, nil)), http.StatusOK
|
|
},
|
|
optional: true,
|
|
},
|
|
},
|
|
}
|
|
|
|
httpServer := httptest.NewServer(fake)
|
|
defer httpServer.Close()
|
|
|
|
transport := &StreamableClientTransport{Endpoint: httpServer.URL}
|
|
client := NewClient(testImpl, &ClientOptions{
|
|
ToolListChangedHandler: func(ctx context.Context, req *ToolListChangedRequest) {},
|
|
})
|
|
|
|
session, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728})
|
|
if err != nil {
|
|
t.Fatalf("Connect failed: %v", err)
|
|
}
|
|
defer session.Close()
|
|
|
|
// tools/list must succeed even though subscriptions/listen returned 404.
|
|
res, err := session.ListTools(ctx, nil)
|
|
if err != nil {
|
|
t.Fatalf("ListTools failed: %v", err)
|
|
}
|
|
if res == nil {
|
|
t.Fatal("ListTools result is nil")
|
|
}
|
|
if !listenServed.Load() {
|
|
t.Fatal("subscriptions/listen was not called")
|
|
}
|
|
}
|