3785c500b5CallToolResult currently decodes structuredContent through an any field, which converts JSON numbers to float64 and silently rounds integers outside the IEEE-754 safe range. This change decodes structuredContent separately with UseNumber, preserving wire numbers as json.Number while retaining the existing decoded object, array, and primitive shapes. The field documentation now describes that representation. Users who depend on the previous float64 behavior can temporarily restore it with MCPGODEBUG=structuredcontentfloat64=1. An end-to-end in-memory client/server regression covers 9007199254740993 from a low-level Server.AddTool result. Additional coverage verifies the compatibility option restores the legacy rounded float64 value. Verification: - gofmt -l . - go test ./... - go vet ./... Fixes #1236 --------- Co-authored-by: Guglielmo Colombo <guglielmoc@google.com>
4455 lines
145 KiB
Go
4455 lines
145 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.
|
|
|
|
//lint:file-ignore SA1019 tests exercise deprecated SEP-2577 APIs (roots, sampling, logging).
|
|
|
|
package mcp
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"maps"
|
|
"net"
|
|
"net/http"
|
|
"net/http/cookiejar"
|
|
"net/http/httptest"
|
|
"net/http/httputil"
|
|
"net/url"
|
|
"os"
|
|
"runtime"
|
|
"slices"
|
|
"sort"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/google/go-cmp/cmp"
|
|
"github.com/google/go-cmp/cmp/cmpopts"
|
|
"github.com/google/jsonschema-go/jsonschema"
|
|
"github.com/modelcontextprotocol/go-sdk/auth"
|
|
"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
|
|
"github.com/modelcontextprotocol/go-sdk/jsonrpc"
|
|
)
|
|
|
|
func TestStreamableTransports(t *testing.T) {
|
|
// This test checks that the streamable server and client transports can
|
|
// communicate.
|
|
|
|
ctx := context.Background()
|
|
|
|
tests := []struct {
|
|
useJSON bool
|
|
replay bool
|
|
}{
|
|
{false, false},
|
|
{false, true},
|
|
{true, false},
|
|
{true, true},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
t.Run(fmt.Sprintf("JSONResponse=%v;replay=%v", test.useJSON, test.replay), func(t *testing.T) {
|
|
// Create a server with some simple tools.
|
|
server := NewServer(testImpl, nil)
|
|
AddTool(server, &Tool{Name: "greet", Description: "say hi"}, sayHi)
|
|
// The "hang" tool checks that context cancellation is propagated.
|
|
// It hangs until the context is cancelled.
|
|
var (
|
|
start = make(chan struct{})
|
|
cancelled = make(chan struct{}, 1) // don't block the request
|
|
)
|
|
hang := func(ctx context.Context, req *CallToolRequest, args any) (*CallToolResult, any, error) {
|
|
start <- struct{}{}
|
|
select {
|
|
case <-ctx.Done():
|
|
cancelled <- struct{}{}
|
|
case <-time.After(5 * time.Second):
|
|
return nil, nil, nil
|
|
}
|
|
return nil, nil, nil
|
|
}
|
|
AddTool(server, &Tool{Name: "hang"}, hang)
|
|
// We use sampling to test server->client requests, both before and after
|
|
// the related client->server request completes.
|
|
sampleDone := make(chan struct{})
|
|
var sampleWG sync.WaitGroup
|
|
AddTool(server, &Tool{Name: "sample"}, func(ctx context.Context, req *CallToolRequest, args any) (*CallToolResult, any, error) {
|
|
type testCase struct {
|
|
label string
|
|
ctx context.Context
|
|
wantSuccess bool
|
|
}
|
|
testSample := func(tc testCase) {
|
|
res, err := req.Session.CreateMessage(tc.ctx, &CreateMessageParams{})
|
|
if gotSuccess := err == nil; gotSuccess != tc.wantSuccess {
|
|
t.Errorf("%s: CreateMessage success=%v, want %v", tc.label, gotSuccess, tc.wantSuccess)
|
|
}
|
|
if err != nil {
|
|
return
|
|
}
|
|
if g, w := res.Model, "aModel"; g != w {
|
|
t.Errorf("%s: got model %q, want %q", tc.label, g, w)
|
|
}
|
|
}
|
|
// Test that we can make sampling requests during tool handling.
|
|
//
|
|
// Try this on both the request context and a background context, so
|
|
// that messages may be delivered on either the POST or GET connection.
|
|
for _, test := range []testCase{
|
|
{"request context", ctx, true},
|
|
{"background context", context.Background(), true},
|
|
} {
|
|
testSample(test)
|
|
}
|
|
// Now, spin off a goroutine that runs after the sampling request, to
|
|
// check behavior when the client request has completed.
|
|
sampleWG.Go(func() {
|
|
<-sampleDone
|
|
// Test that sampling requests in the tool context fail outside of
|
|
// tool handling, but succeed on the background context.
|
|
for _, test := range []testCase{
|
|
{"request context", ctx, false},
|
|
{"background context", context.Background(), true},
|
|
} {
|
|
testSample(test)
|
|
}
|
|
})
|
|
return &CallToolResult{}, nil, nil
|
|
})
|
|
|
|
// Start an httptest.Server with the StreamableHTTPHandler, wrapped in a
|
|
// cookie-checking middleware.
|
|
opts := &StreamableHTTPOptions{
|
|
JSONResponse: test.useJSON,
|
|
}
|
|
if test.replay {
|
|
opts.EventStore = NewMemoryEventStore(nil)
|
|
}
|
|
handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, opts)
|
|
|
|
var (
|
|
headerMu sync.Mutex
|
|
lastHeader http.Header
|
|
)
|
|
httpServer := httptest.NewServer(mustNotPanic(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
headerMu.Lock()
|
|
lastHeader = r.Header
|
|
headerMu.Unlock()
|
|
cookie, err := r.Cookie("test-cookie")
|
|
if err != nil {
|
|
t.Errorf("missing cookie: %v", err)
|
|
} else if cookie.Value != "test-value" {
|
|
t.Errorf("got cookie %q, want %q", cookie.Value, "test-value")
|
|
}
|
|
handler.ServeHTTP(w, r)
|
|
})))
|
|
defer httpServer.Close()
|
|
|
|
// Create a client and connect it to the server using our StreamableClientTransport.
|
|
// Check that all requests honor a custom client.
|
|
jar, err := cookiejar.New(nil)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
u, err := url.Parse(httpServer.URL)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
jar.SetCookies(u, []*http.Cookie{{Name: "test-cookie", Value: "test-value"}})
|
|
httpClient := &http.Client{Jar: jar}
|
|
transport := &StreamableClientTransport{
|
|
Endpoint: httpServer.URL,
|
|
HTTPClient: httpClient,
|
|
}
|
|
client := NewClient(testImpl, &ClientOptions{
|
|
CreateMessageHandler: func(context.Context, *CreateMessageRequest) (*CreateMessageResult, error) {
|
|
return &CreateMessageResult{Model: "aModel", Content: &TextContent{}}, nil
|
|
},
|
|
})
|
|
session, err := client.Connect(ctx, transport, nil)
|
|
if err != nil {
|
|
t.Fatalf("client.Connect() failed: %v", err)
|
|
}
|
|
defer session.Close()
|
|
sid := session.ID()
|
|
if sid == "" {
|
|
t.Fatalf("empty session ID")
|
|
}
|
|
// Stateful HTTP servers reject 2026-07-28 (which requires
|
|
// StreamableHTTPOptions.Stateless = true), so the client falls
|
|
// back to the legacy initialize handshake which explicitly
|
|
// requests 2025-11-25.
|
|
if g, w := session.mcpConn.(*streamableClientConn).initializedResult.ProtocolVersion, protocolVersion20251125; g != w {
|
|
t.Fatalf("got protocol version %q, want %q", g, w)
|
|
}
|
|
|
|
// Verify the behavior of various tools.
|
|
|
|
// The "greet" tool should just work.
|
|
params := &CallToolParams{
|
|
Name: "greet",
|
|
Arguments: map[string]any{"Name": "foo"},
|
|
}
|
|
got, err := session.CallTool(ctx, params)
|
|
if err != nil {
|
|
t.Fatalf("CallTool() failed: %v", err)
|
|
}
|
|
if g := session.ID(); g != sid {
|
|
t.Errorf("session ID: got %q, want %q", g, sid)
|
|
}
|
|
// Stateful HTTP requires the legacy protocol (see the earlier
|
|
// initializeResult assertion); the client sends the matching
|
|
// MCP-Protocol-Version header.
|
|
if g, w := lastHeader.Get(protocolVersionHeader), protocolVersion20251125; g != w {
|
|
t.Errorf("got protocol version header %q, want %q", g, w)
|
|
}
|
|
want := &CallToolResult{
|
|
Content: []Content{&TextContent{Text: "hi foo"}},
|
|
}
|
|
if diff := cmp.Diff(want, got, ctrCmpOpts...); diff != "" {
|
|
t.Errorf("CallTool() returned unexpected content (-want +got):\n%s", diff)
|
|
}
|
|
|
|
// The "hang" tool should be cancellable.
|
|
ctx2, cancel := context.WithCancel(context.Background())
|
|
go session.CallTool(ctx2, &CallToolParams{Name: "hang"})
|
|
<-start
|
|
cancel()
|
|
select {
|
|
case <-cancelled:
|
|
case <-time.After(5 * time.Second):
|
|
t.Fatal("timeout waiting for cancellation")
|
|
}
|
|
|
|
// The "sampling" tool checks the validity of server->client requests
|
|
// both within and without the tool context.
|
|
result, err := session.CallTool(ctx, &CallToolParams{
|
|
Name: "sample",
|
|
Arguments: map[string]any{},
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
// Run the out-of-band checks.
|
|
close(sampleDone)
|
|
sampleWG.Wait()
|
|
|
|
if result.IsError {
|
|
t.Fatalf("tool failed: %s", result.Content[0].(*TextContent).Text)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestStreamableConcurrentHandling(t *testing.T) {
|
|
// This test checks that the streamable server and client transports can
|
|
// communicate.
|
|
type count struct {
|
|
Count int
|
|
}
|
|
|
|
var mu sync.Mutex
|
|
counts := make(map[string]int)
|
|
|
|
server := NewServer(testImpl, nil)
|
|
AddTool(server, &Tool{Name: "inc"}, func(ctx context.Context, req *CallToolRequest, _ any) (*CallToolResult, count, error) {
|
|
id := req.Session.ID()
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
c := counts[id]
|
|
counts[id] = c + 1
|
|
return nil, count{c}, nil
|
|
})
|
|
handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, nil)
|
|
httpServer := httptest.NewServer(mustNotPanic(t, handler))
|
|
defer httpServer.Close()
|
|
|
|
ctx := context.Background()
|
|
client := NewClient(testImpl, nil)
|
|
var wg sync.WaitGroup
|
|
for range 100 {
|
|
wg.Go(func() {
|
|
clientSession, err := client.Connect(ctx, &StreamableClientTransport{Endpoint: httpServer.URL}, nil)
|
|
if err != nil {
|
|
t.Errorf("Connect failed: %v", err)
|
|
return
|
|
}
|
|
defer clientSession.Close()
|
|
for i := range 10 {
|
|
res, err := clientSession.CallTool(ctx, &CallToolParams{Name: "inc"})
|
|
if err != nil {
|
|
t.Errorf("CallTool failed: %v", err)
|
|
return
|
|
}
|
|
number, ok := res.StructuredContent.(map[string]any)["Count"].(json.Number)
|
|
if !ok {
|
|
t.Errorf("count type = %T, want json.Number", res.StructuredContent.(map[string]any)["Count"])
|
|
continue
|
|
}
|
|
got, err := number.Int64()
|
|
if err != nil || got != int64(i) {
|
|
t.Errorf("got count %d, %v; want %d", got, err, i)
|
|
}
|
|
}
|
|
})
|
|
}
|
|
wg.Wait()
|
|
}
|
|
|
|
func TestStreamableServerShutdown(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
// This test checks that closing the streamable HTTP server actually results
|
|
// in client session termination, provided one of following holds:
|
|
// 1. The server is stateful, and therefore the hanging GET fails the connection.
|
|
// 2. The server is stateless, and the client uses a KeepAlive.
|
|
tests := []struct {
|
|
name string
|
|
stateless, keepalive bool
|
|
}{
|
|
{"stateful", false, false},
|
|
{"stateless with keepalive", true, true},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
server := NewServer(testImpl, nil)
|
|
// Add a tool, just so we can check things are working.
|
|
AddTool(server, &Tool{Name: "greet"}, sayHi)
|
|
|
|
handler := NewStreamableHTTPHandler(
|
|
func(req *http.Request) *Server { return server },
|
|
&StreamableHTTPOptions{Stateless: test.stateless})
|
|
|
|
// When we shut down the server, we need to explicitly close ongoing
|
|
// connections. Otherwise, the hanging GET may never terminate.
|
|
httpServer := httptest.NewUnstartedServer(handler)
|
|
httpServer.Config.RegisterOnShutdown(func() {
|
|
for session := range server.Sessions() {
|
|
session.Close()
|
|
}
|
|
})
|
|
httpServer.Start()
|
|
defer httpServer.Close()
|
|
|
|
// Connect and run a tool.
|
|
var opts ClientOptions
|
|
if test.keepalive {
|
|
opts.KeepAlive = 50 * time.Millisecond
|
|
}
|
|
client := NewClient(testImpl, &opts)
|
|
// Pin to 2025-11-25: this test relies on the legacy hanging GET
|
|
// (stateful case) or Ping-based keepalive (stateless case) for
|
|
// the client to notice the server shutdown. Both mechanisms are
|
|
// removed in 2026-07-28 (SEP-2575), so the new protocol would
|
|
// leave the client waiting forever.
|
|
clientSession, err := client.Connect(ctx, &StreamableClientTransport{
|
|
Endpoint: httpServer.URL,
|
|
MaxRetries: -1, // avoid slow tests during exponential retries
|
|
}, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer clientSession.Close()
|
|
|
|
params := &CallToolParams{
|
|
Name: "greet",
|
|
Arguments: map[string]any{"Name": "foo"},
|
|
}
|
|
// Verify that we can call a tool.
|
|
if _, err := clientSession.CallTool(ctx, params); err != nil {
|
|
t.Fatalf("CallTool() failed: %v", err)
|
|
}
|
|
|
|
// Shut down the server. Sessions should terminate.
|
|
go func() {
|
|
if err := httpServer.Config.Shutdown(ctx); err != nil && !errors.Is(err, http.ErrServerClosed) {
|
|
t.Errorf("closing http server: %v", err)
|
|
}
|
|
}()
|
|
|
|
// Wait may return an error (after all, the connection failed), but it
|
|
// should not hang.
|
|
t.Log("Client waiting")
|
|
_ = clientSession.Wait()
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestStreamableStatelessKeepaliveRace verifies that there is no data race between
|
|
// ServerSession.startKeepalive and ServerSession.Close in stateless servers.
|
|
func TestStreamableStatelessKeepaliveRace(t *testing.T) {
|
|
ctx := context.Background()
|
|
server := NewServer(testImpl, &ServerOptions{KeepAlive: time.Hour})
|
|
AddTool(server, &Tool{Name: "greet"}, sayHi)
|
|
handler := NewStreamableHTTPHandler(
|
|
func(*http.Request) *Server { return server },
|
|
&StreamableHTTPOptions{Stateless: true},
|
|
)
|
|
httpServer := httptest.NewServer(mustNotPanic(t, handler))
|
|
defer httpServer.Close()
|
|
|
|
for range 50 {
|
|
cs, err := NewClient(testImpl, nil).Connect(ctx, &StreamableClientTransport{
|
|
Endpoint: httpServer.URL,
|
|
}, nil)
|
|
if err != nil {
|
|
t.Fatalf("NewClient() failed: %v", err)
|
|
}
|
|
_, _ = cs.CallTool(ctx, &CallToolParams{
|
|
Name: "greet",
|
|
Arguments: map[string]any{"Name": "world"},
|
|
})
|
|
_ = cs.Close()
|
|
}
|
|
}
|
|
|
|
// TestClientReplay verifies that the client can recover from a mid-stream
|
|
// network failure and receive replayed messages (if replay is configured). It
|
|
// uses a proxy that is killed and restarted to simulate a recoverable network
|
|
// outage.
|
|
//
|
|
// TODO: Until we have a way to clean up abandoned sessions, this test will leak goroutines (see #499)
|
|
func TestClientReplay(t *testing.T) {
|
|
for _, test := range []clientReplayTest{
|
|
{"default", 0, true},
|
|
{"no retries", -1, false},
|
|
} {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
testClientReplay(t, test)
|
|
})
|
|
}
|
|
}
|
|
|
|
type clientReplayTest struct {
|
|
name string
|
|
maxRetries int
|
|
wantRecovered bool
|
|
}
|
|
|
|
func testClientReplay(t *testing.T, test clientReplayTest) {
|
|
notifications := make(chan string)
|
|
// Configure the real MCP server.
|
|
server := NewServer(testImpl, nil)
|
|
|
|
// Use a channel to synchronize the server's message sending with the test's
|
|
// proxy-killing action.
|
|
serverReadyToKillProxy := make(chan struct{})
|
|
serverClosed := make(chan struct{})
|
|
AddTool(server, &Tool{Name: "multiMessageTool", InputSchema: &jsonschema.Schema{Type: "object"}},
|
|
func(ctx context.Context, req *CallToolRequest, args map[string]any) (*CallToolResult, any, error) {
|
|
// Send one message to the request context, and another to a background
|
|
// context (which will end up on the hanging GET).
|
|
|
|
bgCtx := context.Background()
|
|
req.Session.NotifyProgress(ctx, &ProgressNotificationParams{Message: "msg1"})
|
|
req.Session.NotifyProgress(bgCtx, &ProgressNotificationParams{Message: "msg2"})
|
|
|
|
// Signal the test that it can now kill the proxy.
|
|
close(serverReadyToKillProxy)
|
|
<-serverClosed
|
|
|
|
// These messages should be queued for replay by the server after
|
|
// the client's connection drops.
|
|
req.Session.NotifyProgress(ctx, &ProgressNotificationParams{Message: "msg3"})
|
|
req.Session.NotifyProgress(bgCtx, &ProgressNotificationParams{Message: "msg4"})
|
|
return new(CallToolResult), nil, nil
|
|
})
|
|
|
|
realServer := httptest.NewServer(mustNotPanic(t, NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, &StreamableHTTPOptions{
|
|
EventStore: NewMemoryEventStore(nil), // necessary for replay
|
|
})))
|
|
t.Cleanup(func() {
|
|
t.Log("Closing real HTTP server")
|
|
realServer.Close()
|
|
})
|
|
realServerURL, err := url.Parse(realServer.URL)
|
|
if err != nil {
|
|
t.Fatalf("Failed to parse real server URL: %v", err)
|
|
}
|
|
|
|
// Configure a proxy that sits between the client and the real server.
|
|
proxyHandler := httputil.NewSingleHostReverseProxy(realServerURL)
|
|
// note: don't use mustNotPanic here as the proxy WILL panic when killed.
|
|
proxy := httptest.NewServer(proxyHandler)
|
|
proxyAddr := proxy.Listener.Addr().String() // Get the address to restart it later.
|
|
|
|
// Configure the client to connect to the proxy with default options.
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
client := NewClient(testImpl, &ClientOptions{
|
|
ProgressNotificationHandler: func(ctx context.Context, req *ProgressNotificationClientRequest) {
|
|
notifications <- req.Params.Message
|
|
},
|
|
})
|
|
clientSession, err := client.Connect(ctx, &StreamableClientTransport{
|
|
Endpoint: proxy.URL,
|
|
MaxRetries: test.maxRetries,
|
|
}, nil)
|
|
if err != nil {
|
|
t.Fatalf("client.Connect() failed: %v", err)
|
|
}
|
|
t.Cleanup(func() {
|
|
t.Log("Closing clientSession")
|
|
clientSession.Close()
|
|
})
|
|
|
|
toolCallResult := make(chan error, 1)
|
|
go func() {
|
|
_, callErr := clientSession.CallTool(ctx, &CallToolParams{Name: "multiMessageTool"})
|
|
toolCallResult <- callErr
|
|
}()
|
|
|
|
select {
|
|
case <-serverReadyToKillProxy:
|
|
t.Log("Server has sent the first two messages and is paused.")
|
|
case <-ctx.Done():
|
|
t.Fatalf("Context timed out before server was ready to kill proxy")
|
|
}
|
|
|
|
// We should always get the first two notifications.
|
|
msgs := readNotifications(t, ctx, notifications, 2)
|
|
sort.Strings(msgs) // notifications may arrive in either order
|
|
want := []string{"msg1", "msg2"}
|
|
if diff := cmp.Diff(want, msgs); diff != "" {
|
|
t.Errorf("Recovered notifications mismatch (-want +got):\n%s", diff)
|
|
}
|
|
|
|
// Simulate a total network failure by closing the proxy.
|
|
t.Log("--- Killing proxy to simulate network failure ---")
|
|
proxy.CloseClientConnections()
|
|
proxy.Close()
|
|
close(serverClosed)
|
|
|
|
// Simulate network recovery by restarting the proxy on the same address.
|
|
t.Logf("--- Restarting proxy on %s ---", proxyAddr)
|
|
listener, err := net.Listen("tcp", proxyAddr)
|
|
if err != nil {
|
|
t.Fatalf("Failed to listen on proxy address: %v", err)
|
|
}
|
|
|
|
restartedProxy := &http.Server{Handler: proxyHandler}
|
|
go restartedProxy.Serve(listener)
|
|
t.Cleanup(func() { restartedProxy.Close() })
|
|
|
|
callErr := <-toolCallResult
|
|
|
|
if test.wantRecovered {
|
|
// If we've recovered, we should get all 4 notifications and the tool call
|
|
// should have succeeded.
|
|
msgs := readNotifications(t, ctx, notifications, 2)
|
|
sort.Strings(msgs)
|
|
want := []string{"msg3", "msg4"}
|
|
if diff := cmp.Diff(want, msgs); diff != "" {
|
|
t.Errorf("Recovered notifications mismatch (-want +got):\n%s", diff)
|
|
}
|
|
if callErr != nil {
|
|
t.Errorf("CallTool failed unexpectedly: %v", err)
|
|
}
|
|
} else {
|
|
// Otherwise, the call should fail.
|
|
if callErr == nil {
|
|
t.Errorf("CallTool succeeded unexpectedly")
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestStreamableServerDisconnect(t *testing.T) {
|
|
server := NewServer(testImpl, nil)
|
|
|
|
// Test that client replayability allows the server to terminate incoming
|
|
// requests immediately, and have the client replay them.
|
|
|
|
// testStream exercises stream resumption by interleaving stream termination
|
|
// with progress notifications.
|
|
testStream := func(ctx context.Context, session *ServerSession, extra *RequestExtra) {
|
|
// Close the stream before the first message. We should have sent an
|
|
// initial priming message already, so the client will be able to replay
|
|
extra.CloseSSEStream(CloseSSEStreamArgs{RetryAfter: 10 * time.Millisecond})
|
|
session.NotifyProgress(ctx, &ProgressNotificationParams{Message: "msg1"})
|
|
time.Sleep(20 * time.Millisecond)
|
|
extra.CloseSSEStream(CloseSSEStreamArgs{RetryAfter: 10 * time.Millisecond}) // Closing twice should still be supported.
|
|
session.NotifyProgress(ctx, &ProgressNotificationParams{Message: "msg2"})
|
|
}
|
|
|
|
AddTool(server, &Tool{Name: "disconnect"},
|
|
func(ctx context.Context, req *CallToolRequest, args map[string]any) (*CallToolResult, map[string]any, error) {
|
|
testStream(ctx, req.Session, req.Extra)
|
|
return new(CallToolResult), nil, nil
|
|
})
|
|
|
|
server.AddPrompt(&Prompt{Name: "disconnect"}, func(ctx context.Context, req *GetPromptRequest) (*GetPromptResult, error) {
|
|
testStream(ctx, req.Session, req.Extra)
|
|
return nil, nil
|
|
})
|
|
|
|
tests := []struct {
|
|
name string
|
|
doCall func(context.Context, *ClientSession) error
|
|
}{
|
|
{
|
|
"tool",
|
|
func(ctx context.Context, cs *ClientSession) error {
|
|
_, err := cs.CallTool(ctx, &CallToolParams{Name: "disconnect"})
|
|
return err
|
|
},
|
|
},
|
|
{
|
|
"prompt",
|
|
func(ctx context.Context, cs *ClientSession) error {
|
|
_, err := cs.GetPrompt(ctx, &GetPromptParams{Name: "disconnect"})
|
|
return err
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
notifications := make(chan string, 2)
|
|
handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, &StreamableHTTPOptions{
|
|
EventStore: NewMemoryEventStore(nil), // support replayability
|
|
})
|
|
httpServer := httptest.NewServer(mustNotPanic(t, handler))
|
|
defer httpServer.Close()
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
client := NewClient(testImpl, &ClientOptions{
|
|
ProgressNotificationHandler: func(ctx context.Context, req *ProgressNotificationClientRequest) {
|
|
notifications <- req.Params.Message
|
|
},
|
|
})
|
|
clientSession, err := client.Connect(ctx, &StreamableClientTransport{
|
|
Endpoint: httpServer.URL,
|
|
}, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125})
|
|
if err != nil {
|
|
t.Fatalf("client.Connect() failed: %v", err)
|
|
}
|
|
defer clientSession.Close()
|
|
|
|
if err = test.doCall(ctx, clientSession); err != nil {
|
|
t.Fatalf("CallTool failed: %v", err)
|
|
}
|
|
|
|
got := readNotifications(t, ctx, notifications, 2)
|
|
want := []string{"msg1", "msg2"}
|
|
if diff := cmp.Diff(want, got); diff != "" {
|
|
t.Errorf("got unexpected notifications (-want +got):\n%s", diff)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestServerTransportCleanup(t *testing.T) {
|
|
nClient := 3
|
|
|
|
var mu sync.Mutex
|
|
var id int = -1 // session id starting from "0", "1", "2"...
|
|
chans := make(map[string]chan struct{}, nClient)
|
|
|
|
server := NewServer(testImpl, &ServerOptions{
|
|
KeepAlive: 10 * time.Millisecond,
|
|
GetSessionID: func() string {
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
id++
|
|
if id == nClient {
|
|
t.Errorf("creating more than %v session", nClient)
|
|
}
|
|
chans[fmt.Sprint(id)] = make(chan struct{}, 1)
|
|
return fmt.Sprint(id)
|
|
},
|
|
})
|
|
|
|
handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, nil)
|
|
handler.onTransportDeletion = func(sessionID string) {
|
|
mu.Lock()
|
|
ch := chans[sessionID]
|
|
mu.Unlock()
|
|
ch <- struct{}{}
|
|
}
|
|
|
|
httpServer := httptest.NewServer(mustNotPanic(t, handler))
|
|
defer httpServer.Close()
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
|
|
// Spin up clients connect to the same server but refuse to ping request.
|
|
for range nClient {
|
|
client := NewClient(testImpl, nil)
|
|
pingMiddleware := func(next MethodHandler) MethodHandler {
|
|
return func(
|
|
ctx context.Context,
|
|
method string,
|
|
req Request,
|
|
) (Result, error) {
|
|
if method == "ping" {
|
|
return &emptyResult{}, errors.New("ping error")
|
|
}
|
|
return next(ctx, method, req)
|
|
}
|
|
}
|
|
client.AddReceivingMiddleware(pingMiddleware)
|
|
// Pin to 2025-11-25: this test counts server-side sessions, but
|
|
// the 2026-07-28 client probes server/discover first against a
|
|
// stateful HTTP server (which rejects the new protocol), opening
|
|
// an extra HTTP connection and therefore an extra session. Pinning
|
|
// to 2025-11-25 skips the discover probe.
|
|
clientSession, err := client.Connect(ctx, &StreamableClientTransport{Endpoint: httpServer.URL}, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125})
|
|
if err != nil {
|
|
t.Fatalf("client.Connect() failed: %v", err)
|
|
}
|
|
t.Cleanup(func() { _ = clientSession.Close() })
|
|
}
|
|
|
|
mu.Lock()
|
|
channels := slices.Collect(maps.Values(chans))
|
|
mu.Unlock()
|
|
|
|
for _, ch := range channels {
|
|
select {
|
|
case <-ctx.Done():
|
|
t.Errorf("did not capture transport deletion event from all session in 10 seconds")
|
|
case <-ch:
|
|
t.Log("Received session transport deletion signal")
|
|
}
|
|
}
|
|
|
|
handler.mu.Lock()
|
|
if len(handler.sessions) != 0 {
|
|
t.Errorf("want empty transports map, find %v entries from handler's transports map", len(handler.sessions))
|
|
}
|
|
handler.mu.Unlock()
|
|
}
|
|
|
|
// TestServerInitiatedSSE verifies that the persistent SSE connection remains
|
|
// open and can receive server-initiated events.
|
|
func TestServerInitiatedSSE(t *testing.T) {
|
|
notifications := make(chan string)
|
|
server := NewServer(testImpl, nil)
|
|
|
|
httpServer := httptest.NewServer(mustNotPanic(t, NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, nil)))
|
|
defer httpServer.Close()
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
client := NewClient(testImpl, &ClientOptions{
|
|
ToolListChangedHandler: func(context.Context, *ToolListChangedRequest) {
|
|
notifications <- "toolListChanged"
|
|
},
|
|
})
|
|
clientSession, err := client.Connect(ctx, &StreamableClientTransport{Endpoint: httpServer.URL}, nil)
|
|
if err != nil {
|
|
t.Fatalf("client.Connect() failed: %v", err)
|
|
}
|
|
defer clientSession.Close()
|
|
AddTool(server, &Tool{Name: "testTool", InputSchema: &jsonschema.Schema{Type: "object"}},
|
|
func(context.Context, *CallToolRequest, map[string]any) (*CallToolResult, any, error) {
|
|
return &CallToolResult{}, nil, nil
|
|
})
|
|
receivedNotifications := readNotifications(t, ctx, notifications, 1)
|
|
wantReceived := []string{"toolListChanged"}
|
|
if diff := cmp.Diff(wantReceived, receivedNotifications); diff != "" {
|
|
t.Errorf("Received notifications mismatch (-want +got):\n%s", diff)
|
|
}
|
|
}
|
|
|
|
// Helper to read a specific number of notifications.
|
|
func readNotifications(t *testing.T, ctx context.Context, notifications chan string, count int) []string {
|
|
t.Helper()
|
|
var collectedNotifications []string
|
|
for {
|
|
select {
|
|
case n := <-notifications:
|
|
collectedNotifications = append(collectedNotifications, n)
|
|
if len(collectedNotifications) == count {
|
|
return collectedNotifications
|
|
}
|
|
case <-ctx.Done():
|
|
if len(collectedNotifications) != count {
|
|
t.Fatalf("readProgressNotifications(): did not receive expected notifications, got %d, want %d", len(collectedNotifications), count)
|
|
}
|
|
return collectedNotifications
|
|
}
|
|
}
|
|
}
|
|
|
|
// JSON-RPC message constructors.
|
|
func req(id int64, method string, params any) *jsonrpc.Request {
|
|
r := &jsonrpc.Request{
|
|
Method: method,
|
|
Params: mustMarshal(params),
|
|
}
|
|
if id > 0 {
|
|
r.ID = jsonrpc2.Int64ID(id)
|
|
}
|
|
return r
|
|
}
|
|
|
|
func completeCallToolResult() *CallToolResult {
|
|
r := &CallToolResult{Content: []Content{}}
|
|
r.resultType = resultTypeComplete
|
|
return r
|
|
}
|
|
|
|
// completeCallToolResultWithServerInfo returns a completeCallToolResult
|
|
// annotated with the auto-populated [MetaKeyServerInfo] entry that servers
|
|
// add to every new-protocol response.
|
|
func completeCallToolResultWithServerInfo(impl *Implementation) *CallToolResult {
|
|
r := completeCallToolResult()
|
|
r.Meta = Meta{MetaKeyServerInfo: impl}
|
|
return r
|
|
}
|
|
|
|
func resp(id int64, result any, err error) *jsonrpc.Response {
|
|
return &jsonrpc.Response{
|
|
ID: jsonrpc2.Int64ID(id),
|
|
Result: mustMarshal(result),
|
|
Error: err,
|
|
}
|
|
}
|
|
|
|
func TestStreamableServerTransport(t *testing.T) {
|
|
// This test checks detailed behavior of the streamable server transport, by
|
|
// faking the behavior of a streamable client using a sequence of HTTP
|
|
// requests.
|
|
|
|
// Predefined steps, to avoid repetition below.
|
|
initReq := req(1, methodInitialize, &InitializeParams{ProtocolVersion: protocolVersion20250618})
|
|
initResp := resp(1, &InitializeResult{
|
|
Capabilities: &ServerCapabilities{
|
|
Logging: &LoggingCapabilities{},
|
|
Tools: &ToolCapabilities{ListChanged: true},
|
|
},
|
|
ProtocolVersion: protocolVersion20250618,
|
|
ServerInfo: &Implementation{Name: "testServer", Version: "v1.0.0"},
|
|
}, nil)
|
|
initializedMsg := req(0, notificationInitialized, &InitializedParams{})
|
|
initialize := streamableRequest{
|
|
method: "POST",
|
|
messages: []jsonrpc.Message{initReq},
|
|
wantStatusCode: http.StatusOK,
|
|
wantMessages: []jsonrpc.Message{initResp},
|
|
wantSessionID: true,
|
|
}
|
|
initialized := streamableRequest{
|
|
method: "POST",
|
|
messages: []jsonrpc.Message{initializedMsg},
|
|
wantStatusCode: http.StatusAccepted,
|
|
}
|
|
|
|
// Protocol version 2025-11-25 variants, for testing prime/close events (SEP-1699).
|
|
initReq20251125 := req(1, methodInitialize, &InitializeParams{ProtocolVersion: protocolVersion20251125})
|
|
initResp20251125 := resp(1, &InitializeResult{
|
|
Capabilities: &ServerCapabilities{
|
|
Logging: &LoggingCapabilities{},
|
|
Tools: &ToolCapabilities{ListChanged: true},
|
|
},
|
|
ProtocolVersion: protocolVersion20251125,
|
|
ServerInfo: &Implementation{Name: "testServer", Version: "v1.0.0"},
|
|
}, nil)
|
|
initialize20251125 := streamableRequest{
|
|
method: "POST",
|
|
messages: []jsonrpc.Message{initReq20251125},
|
|
wantStatusCode: http.StatusOK,
|
|
wantMessages: []jsonrpc.Message{initResp20251125},
|
|
wantSessionID: true,
|
|
}
|
|
initialized20251125 := streamableRequest{
|
|
method: "POST",
|
|
headers: http.Header{protocolVersionHeader: {protocolVersion20251125}},
|
|
messages: []jsonrpc.Message{initializedMsg},
|
|
wantStatusCode: http.StatusAccepted,
|
|
}
|
|
|
|
tests := []struct {
|
|
name string
|
|
replay bool // if set, use a MemoryEventStore to enable replay
|
|
tool func(*testing.T, context.Context, *CallToolRequest) // if set, called during execution
|
|
requests []streamableRequest
|
|
wantSessions int // number of sessions expected after the test
|
|
}{
|
|
{
|
|
name: "basic",
|
|
requests: []streamableRequest{
|
|
initialize,
|
|
initialized,
|
|
{
|
|
method: "POST",
|
|
messages: []jsonrpc.Message{req(2, "tools/call", &CallToolParams{Name: "tool"})},
|
|
wantStatusCode: http.StatusOK,
|
|
wantMessages: []jsonrpc.Message{resp(2, &CallToolResult{Content: []Content{}}, nil)},
|
|
},
|
|
},
|
|
wantSessions: 1,
|
|
},
|
|
{
|
|
name: "uninitialized",
|
|
requests: []streamableRequest{
|
|
{
|
|
method: "POST",
|
|
messages: []jsonrpc.Message{req(2, "tools/call", &CallToolParams{Name: "tool"})},
|
|
wantStatusCode: http.StatusOK,
|
|
wantBodyContaining: "invalid during session initialization",
|
|
},
|
|
},
|
|
wantSessions: 0,
|
|
},
|
|
{
|
|
name: "content type headers",
|
|
requests: []streamableRequest{
|
|
initialize,
|
|
initialized,
|
|
{
|
|
// Request with incorrect Content-Type should be rejected.
|
|
method: "POST",
|
|
headers: http.Header{"Content-Type": {"text/plain"}},
|
|
messages: []jsonrpc.Message{req(3, "tools/call", &CallToolParams{Name: "tool"})},
|
|
wantStatusCode: http.StatusUnsupportedMediaType,
|
|
},
|
|
{
|
|
// Request with empty Content-Type should be rejected.
|
|
method: "POST",
|
|
headers: http.Header{"Content-Type": {""}},
|
|
messages: []jsonrpc.Message{req(4, "tools/call", &CallToolParams{Name: "tool"})},
|
|
wantStatusCode: http.StatusUnsupportedMediaType,
|
|
},
|
|
{
|
|
// Correct Content-Type should pass.
|
|
method: "POST",
|
|
headers: http.Header{"Content-Type": {"application/json"}},
|
|
messages: []jsonrpc.Message{req(5, "tools/call", &CallToolParams{Name: "tool"})},
|
|
wantStatusCode: http.StatusOK,
|
|
wantMessages: []jsonrpc.Message{resp(5, &CallToolResult{Content: []Content{}}, nil)},
|
|
},
|
|
{
|
|
// Correct Content-Type with parameters should pass.
|
|
method: "POST",
|
|
headers: http.Header{"Content-Type": {"application/json; charset=utf-8"}},
|
|
messages: []jsonrpc.Message{req(5, "tools/call", &CallToolParams{Name: "tool"})},
|
|
wantStatusCode: http.StatusOK,
|
|
wantMessages: []jsonrpc.Message{resp(5, &CallToolResult{Content: []Content{}}, nil)},
|
|
},
|
|
},
|
|
wantSessions: 1,
|
|
},
|
|
{
|
|
name: "accept headers",
|
|
requests: []streamableRequest{
|
|
initialize,
|
|
initialized,
|
|
// Test various accept headers.
|
|
{
|
|
method: "POST",
|
|
headers: http.Header{"Accept": {"text/plain", "application/*"}},
|
|
messages: []jsonrpc.Message{req(3, "tools/call", &CallToolParams{Name: "tool"})},
|
|
wantStatusCode: http.StatusBadRequest, // missing text/event-stream
|
|
},
|
|
{
|
|
method: "POST",
|
|
headers: http.Header{"Accept": {"text/event-stream"}},
|
|
messages: []jsonrpc.Message{req(3, "tools/call", &CallToolParams{Name: "tool"})},
|
|
wantStatusCode: http.StatusBadRequest, // missing application/json
|
|
},
|
|
{
|
|
method: "POST",
|
|
headers: http.Header{"Accept": {"text/plain", "*/*"}},
|
|
messages: []jsonrpc.Message{req(4, "tools/call", &CallToolParams{Name: "tool"})},
|
|
wantStatusCode: http.StatusOK,
|
|
wantMessages: []jsonrpc.Message{resp(4, &CallToolResult{Content: []Content{}}, nil)},
|
|
},
|
|
{
|
|
method: "POST",
|
|
headers: http.Header{"Accept": {"application/json;charset=utf-8, text/event-stream"}},
|
|
messages: []jsonrpc.Message{req(5, "tools/call", &CallToolParams{Name: "tool"})},
|
|
wantStatusCode: http.StatusOK,
|
|
wantMessages: []jsonrpc.Message{resp(5, &CallToolResult{Content: []Content{}}, nil)},
|
|
},
|
|
{
|
|
method: "POST",
|
|
headers: http.Header{"Accept": {"application/json;charset=utf-8", "text/event-stream"}},
|
|
messages: []jsonrpc.Message{req(6, "tools/call", &CallToolParams{Name: "tool"})},
|
|
wantStatusCode: http.StatusOK,
|
|
wantMessages: []jsonrpc.Message{resp(6, &CallToolResult{Content: []Content{}}, nil)},
|
|
},
|
|
{
|
|
method: "POST",
|
|
headers: http.Header{"Accept": {"text/*, application/*"}},
|
|
messages: []jsonrpc.Message{req(7, "tools/call", &CallToolParams{Name: "tool"})},
|
|
wantStatusCode: http.StatusOK,
|
|
wantMessages: []jsonrpc.Message{resp(7, &CallToolResult{Content: []Content{}}, nil)},
|
|
},
|
|
},
|
|
wantSessions: 1,
|
|
},
|
|
{
|
|
name: "protocol version headers",
|
|
requests: []streamableRequest{
|
|
initialize,
|
|
initialized,
|
|
{
|
|
method: "POST",
|
|
headers: http.Header{"mcp-protocol-version": {"2025-01-01"}}, // an invalid protocol version
|
|
messages: []jsonrpc.Message{req(2, "tools/call", &CallToolParams{Name: "tool"})},
|
|
wantStatusCode: http.StatusBadRequest,
|
|
wantBodyContaining: "2025-03-26", // a supported version
|
|
wantSessionID: false, // could be true, but shouldn't matter
|
|
},
|
|
},
|
|
wantSessions: 1,
|
|
},
|
|
{
|
|
name: "batch rejected on 2025-06-18",
|
|
requests: []streamableRequest{
|
|
initialize,
|
|
initialized,
|
|
{
|
|
method: "POST",
|
|
// Explicitly set the protocol version header
|
|
headers: http.Header{"MCP-Protocol-Version": {"2025-06-18"}},
|
|
// Two messages => batch. Expect reject.
|
|
messages: []jsonrpc.Message{
|
|
req(101, "tools/call", &CallToolParams{Name: "tool"}),
|
|
req(102, "tools/call", &CallToolParams{Name: "tool"}),
|
|
},
|
|
wantStatusCode: http.StatusBadRequest,
|
|
wantBodyContaining: "batch",
|
|
},
|
|
},
|
|
wantSessions: 1,
|
|
},
|
|
{
|
|
name: "batch accepted on 2025-03-26",
|
|
requests: []streamableRequest{
|
|
initialize,
|
|
initialized,
|
|
{
|
|
method: "POST",
|
|
headers: http.Header{"MCP-Protocol-Version": {"2025-03-26"}},
|
|
// Two messages => batch. Expect OK with two responses in order.
|
|
messages: []jsonrpc.Message{
|
|
// Note: only include one request here, because responses are not
|
|
// necessarily sorted.
|
|
req(201, "tools/call", &CallToolParams{Name: "tool"}),
|
|
req(0, "notifications/roots/list_changed", &RootsListChangedParams{}),
|
|
},
|
|
wantStatusCode: http.StatusOK,
|
|
wantMessages: []jsonrpc.Message{
|
|
resp(201, &CallToolResult{Content: []Content{}}, nil),
|
|
},
|
|
},
|
|
},
|
|
wantSessions: 1,
|
|
},
|
|
{
|
|
name: "tool notification",
|
|
tool: func(t *testing.T, ctx context.Context, req *CallToolRequest) {
|
|
// Send an arbitrary notification.
|
|
if err := req.Session.NotifyProgress(ctx, &ProgressNotificationParams{}); err != nil {
|
|
t.Errorf("Notify failed: %v", err)
|
|
}
|
|
},
|
|
requests: []streamableRequest{
|
|
initialize,
|
|
initialized,
|
|
{
|
|
method: "POST",
|
|
messages: []jsonrpc.Message{
|
|
req(2, "tools/call", &CallToolParams{Name: "tool"}),
|
|
},
|
|
wantStatusCode: http.StatusOK,
|
|
wantMessages: []jsonrpc.Message{
|
|
req(0, "notifications/progress", &ProgressNotificationParams{}),
|
|
resp(2, &CallToolResult{Content: []Content{}}, nil),
|
|
},
|
|
},
|
|
},
|
|
wantSessions: 1,
|
|
},
|
|
{
|
|
name: "tool upcall",
|
|
tool: func(t *testing.T, ctx context.Context, req *CallToolRequest) {
|
|
// Make an arbitrary call.
|
|
if _, err := req.Session.ListRoots(ctx, &ListRootsParams{}); err != nil {
|
|
t.Errorf("Call failed: %v", err)
|
|
}
|
|
},
|
|
requests: []streamableRequest{
|
|
initialize,
|
|
initialized,
|
|
{
|
|
method: "POST",
|
|
onRequest: 1,
|
|
messages: []jsonrpc.Message{
|
|
resp(1, &ListRootsResult{}, nil),
|
|
},
|
|
wantStatusCode: http.StatusAccepted,
|
|
},
|
|
{
|
|
method: "POST",
|
|
messages: []jsonrpc.Message{
|
|
req(2, "tools/call", &CallToolParams{Name: "tool"}),
|
|
},
|
|
wantStatusCode: http.StatusOK,
|
|
wantMessages: []jsonrpc.Message{
|
|
req(1, "roots/list", &ListRootsParams{}),
|
|
resp(2, &CallToolResult{Content: []Content{}}, nil),
|
|
},
|
|
},
|
|
},
|
|
wantSessions: 1,
|
|
},
|
|
{
|
|
name: "background",
|
|
// Enabling replay is necessary here because the standalone "GET" request
|
|
// is fully asynchronous. Replay is needed to guarantee message delivery.
|
|
//
|
|
// TODO(rfindley): this should no longer be necessary.
|
|
replay: true,
|
|
tool: func(t *testing.T, _ context.Context, req *CallToolRequest) {
|
|
// Perform operations on a background context, and ensure the client
|
|
// receives it.
|
|
ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
|
|
defer cancel()
|
|
|
|
if err := req.Session.NotifyProgress(ctx, &ProgressNotificationParams{}); err != nil {
|
|
t.Errorf("Notify failed: %v", err)
|
|
}
|
|
// TODO(rfindley): finish implementing logging.
|
|
// if err := ss.LoggingMessage(ctx, &LoggingMessageParams{}); err != nil {
|
|
// t.Errorf("Logging failed: %v", err)
|
|
// }
|
|
if _, err := req.Session.ListRoots(ctx, &ListRootsParams{}); err != nil {
|
|
t.Errorf("ListRoots failed: %v", err)
|
|
}
|
|
},
|
|
requests: []streamableRequest{
|
|
initialize,
|
|
initialized,
|
|
{
|
|
method: "POST",
|
|
onRequest: 1,
|
|
messages: []jsonrpc.Message{
|
|
resp(1, &ListRootsResult{}, nil),
|
|
},
|
|
wantStatusCode: http.StatusAccepted,
|
|
},
|
|
{
|
|
method: "GET",
|
|
async: true,
|
|
wantStatusCode: http.StatusOK,
|
|
closeAfter: 2,
|
|
wantMessages: []jsonrpc.Message{
|
|
req(0, "notifications/progress", &ProgressNotificationParams{}),
|
|
req(1, "roots/list", &ListRootsParams{}),
|
|
},
|
|
},
|
|
{
|
|
method: "POST",
|
|
messages: []jsonrpc.Message{
|
|
req(2, "tools/call", &CallToolParams{Name: "tool"}),
|
|
},
|
|
wantStatusCode: http.StatusOK,
|
|
wantMessages: []jsonrpc.Message{
|
|
resp(2, &CallToolResult{Content: []Content{}}, nil),
|
|
},
|
|
},
|
|
{
|
|
method: "DELETE",
|
|
wantStatusCode: http.StatusNoContent,
|
|
// Delete request expects 204 No Content with empty body. So override
|
|
// the default "accept: application/json, text/event-stream" header.
|
|
headers: map[string][]string{"Accept": nil},
|
|
},
|
|
},
|
|
wantSessions: 0, // session deleted
|
|
},
|
|
{
|
|
name: "no priming message on old protocol",
|
|
replay: true,
|
|
requests: []streamableRequest{
|
|
initialize,
|
|
initialized,
|
|
{
|
|
method: "POST",
|
|
messages: []jsonrpc.Message{req(2, "tools/call", &CallToolParams{Name: "tool"})},
|
|
wantStatusCode: http.StatusOK,
|
|
wantMessages: []jsonrpc.Message{resp(2, &CallToolResult{Content: []Content{}}, nil)},
|
|
wantBodyNotContaining: "prime",
|
|
},
|
|
},
|
|
wantSessions: 1,
|
|
},
|
|
{
|
|
name: "no close message on old protocol",
|
|
replay: true,
|
|
tool: func(t *testing.T, _ context.Context, req *CallToolRequest) {
|
|
req.Extra.CloseSSEStream(CloseSSEStreamArgs{RetryAfter: time.Millisecond})
|
|
},
|
|
requests: []streamableRequest{
|
|
initialize,
|
|
initialized,
|
|
{
|
|
method: "POST",
|
|
messages: []jsonrpc.Message{req(2, "tools/call", &CallToolParams{Name: "tool"})},
|
|
wantStatusCode: http.StatusOK,
|
|
wantMessages: []jsonrpc.Message{resp(2, &CallToolResult{Content: []Content{}}, nil)},
|
|
wantBodyNotContaining: "close",
|
|
},
|
|
},
|
|
wantSessions: 1,
|
|
},
|
|
{
|
|
name: "priming message on 2025-11-25",
|
|
replay: true,
|
|
requests: []streamableRequest{
|
|
initialize20251125,
|
|
initialized20251125,
|
|
{
|
|
method: "POST",
|
|
headers: http.Header{protocolVersionHeader: {protocolVersion20251125}},
|
|
messages: []jsonrpc.Message{req(2, "tools/call", &CallToolParams{Name: "tool"})},
|
|
wantStatusCode: http.StatusOK,
|
|
wantMessages: []jsonrpc.Message{resp(2, &CallToolResult{Content: []Content{}}, nil)},
|
|
wantBodyContaining: "prime",
|
|
},
|
|
},
|
|
wantSessions: 1,
|
|
},
|
|
{
|
|
name: "close message on 2025-11-25",
|
|
replay: true,
|
|
tool: func(t *testing.T, _ context.Context, req *CallToolRequest) {
|
|
req.Extra.CloseSSEStream(CloseSSEStreamArgs{RetryAfter: time.Millisecond})
|
|
},
|
|
requests: []streamableRequest{
|
|
initialize20251125,
|
|
initialized20251125,
|
|
{
|
|
method: "POST",
|
|
headers: http.Header{protocolVersionHeader: {protocolVersion20251125}},
|
|
messages: []jsonrpc.Message{req(2, "tools/call", &CallToolParams{Name: "tool"})},
|
|
wantStatusCode: http.StatusOK,
|
|
wantMessages: []jsonrpc.Message{resp(2, &CallToolResult{Content: []Content{}}, nil)},
|
|
wantBodyContaining: "close",
|
|
wantBodyNotContaining: "result",
|
|
},
|
|
},
|
|
wantSessions: 1,
|
|
},
|
|
{
|
|
name: "no close message",
|
|
replay: true,
|
|
tool: func(t *testing.T, _ context.Context, req *CallToolRequest) {
|
|
req.Extra.CloseSSEStream(CloseSSEStreamArgs{})
|
|
},
|
|
requests: []streamableRequest{
|
|
initialize,
|
|
initialized,
|
|
{
|
|
method: "POST",
|
|
messages: []jsonrpc.Message{req(2, "tools/call", &CallToolParams{Name: "tool"})},
|
|
wantStatusCode: http.StatusOK,
|
|
wantMessages: []jsonrpc.Message{resp(2, &CallToolResult{Content: []Content{}}, nil)},
|
|
wantBodyNotContaining: "close",
|
|
},
|
|
},
|
|
wantSessions: 1,
|
|
},
|
|
{
|
|
name: "errors",
|
|
requests: []streamableRequest{
|
|
{
|
|
method: "PUT",
|
|
wantStatusCode: http.StatusMethodNotAllowed,
|
|
},
|
|
{
|
|
method: "DELETE",
|
|
wantStatusCode: http.StatusBadRequest,
|
|
},
|
|
{
|
|
method: "POST",
|
|
messages: []jsonrpc.Message{req(1, "notamethod", nil)},
|
|
wantStatusCode: http.StatusBadRequest, // notamethod is an invalid method
|
|
},
|
|
{
|
|
method: "POST",
|
|
messages: []jsonrpc.Message{req(0, "tools/call", &CallToolParams{Name: "tool"})},
|
|
wantStatusCode: http.StatusBadRequest, // tools/call must have an ID
|
|
},
|
|
{
|
|
method: "POST",
|
|
messages: []jsonrpc.Message{req(2, "tools/call", &CallToolParams{Name: "tool"})},
|
|
wantStatusCode: http.StatusOK,
|
|
wantMessages: []jsonrpc.Message{resp(2, nil, &jsonrpc.Error{
|
|
Message: `method "tools/call" is invalid during session initialization`,
|
|
})},
|
|
},
|
|
},
|
|
wantSessions: 0,
|
|
},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
// Create a server containing a single tool, which runs the test tool
|
|
// behavior, if any.
|
|
server := NewServer(&Implementation{Name: "testServer", Version: "v1.0.0"}, nil)
|
|
server.AddTool(
|
|
&Tool{Name: "tool", InputSchema: &jsonschema.Schema{Type: "object"}},
|
|
func(ctx context.Context, req *CallToolRequest) (*CallToolResult, error) {
|
|
if test.tool != nil {
|
|
test.tool(t, ctx, req)
|
|
}
|
|
return &CallToolResult{}, nil
|
|
})
|
|
|
|
opts := &StreamableHTTPOptions{}
|
|
if test.replay {
|
|
opts.EventStore = NewMemoryEventStore(nil)
|
|
}
|
|
// Start the streamable handler.
|
|
handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, opts)
|
|
defer handler.closeAll()
|
|
|
|
testStreamableHandler(t, handler, test.requests)
|
|
if got := len(slices.Collect(server.Sessions())); got != test.wantSessions {
|
|
t.Errorf("after test, got %d sessions, want %d", got, test.wantSessions)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func testStreamableHandler(t *testing.T, handler http.Handler, requests []streamableRequest) {
|
|
httpServer := httptest.NewServer(mustNotPanic(t, handler))
|
|
defer httpServer.Close()
|
|
|
|
// blocks records request blocks by jsonrpc. ID.
|
|
//
|
|
// When an OnRequest step is encountered, it waits on the corresponding
|
|
// block. When a request with that ID is received, the block is closed.
|
|
var mu sync.Mutex
|
|
blocks := make(map[int64]chan struct{})
|
|
for _, req := range requests {
|
|
if req.onRequest > 0 {
|
|
blocks[req.onRequest] = make(chan struct{})
|
|
}
|
|
}
|
|
|
|
// signal when all synchronous requests have executed, so we can fail
|
|
// async requests that are blocked.
|
|
syncRequestsDone := make(chan struct{})
|
|
|
|
// To avoid complicated accounting for session ID, just set the first
|
|
// non-empty session ID from a response.
|
|
var sessionID atomic.Value
|
|
sessionID.Store("")
|
|
|
|
// doStep executes a single step.
|
|
doStep := func(t *testing.T, i int, request streamableRequest) {
|
|
if request.onRequest > 0 {
|
|
// Block the step until we've received the server->client request.
|
|
mu.Lock()
|
|
block := blocks[request.onRequest]
|
|
mu.Unlock()
|
|
select {
|
|
case <-block:
|
|
case <-syncRequestsDone:
|
|
t.Errorf("after all sync requests are complete, request still blocked on %d", request.onRequest)
|
|
return
|
|
}
|
|
}
|
|
|
|
// Collect messages received during this request, unblock other steps
|
|
// when requests are received.
|
|
var got []jsonrpc.Message
|
|
out := make(chan jsonrpc.Message)
|
|
// Cancel the step if we encounter a request that isn't going to be
|
|
// handled.
|
|
//
|
|
// Also, add a timeout (hopefully generous).
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
|
|
var wg sync.WaitGroup
|
|
wg.Go(func() {
|
|
|
|
for m := range out {
|
|
if req, ok := m.(*jsonrpc.Request); ok && req.IsCall() {
|
|
// Encountered a server->client request. We should have a
|
|
// response queued. Otherwise, we may deadlock.
|
|
mu.Lock()
|
|
if block, ok := blocks[req.ID.Raw().(int64)]; ok {
|
|
close(block)
|
|
} else {
|
|
t.Errorf("no queued response for %v", req.ID)
|
|
cancel()
|
|
}
|
|
mu.Unlock()
|
|
}
|
|
got = append(got, m)
|
|
if request.closeAfter > 0 && len(got) == request.closeAfter {
|
|
cancel()
|
|
}
|
|
}
|
|
})
|
|
|
|
gotSessionID, gotStatusCode, gotBody, err := request.do(ctx, httpServer.URL, sessionID.Load().(string), out)
|
|
|
|
// Don't fail on cancelled requests: error (if any) is handled
|
|
// elsewhere.
|
|
if err != nil && ctx.Err() == nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
if gotStatusCode != request.wantStatusCode {
|
|
t.Errorf("request #%d: got status %d, want %d", i, gotStatusCode, request.wantStatusCode)
|
|
}
|
|
if got := gotSessionID != ""; got != request.wantSessionID {
|
|
t.Errorf("request #%d: got session id: %t, want %t", i, got, request.wantSessionID)
|
|
}
|
|
wg.Wait()
|
|
|
|
if request.wantBodyContaining != "" || request.wantBodyNotContaining != "" {
|
|
body := string(gotBody)
|
|
if request.wantBodyContaining != "" && !strings.Contains(body, request.wantBodyContaining) {
|
|
t.Errorf("body does not contain %q:\n%s", request.wantBodyContaining, body)
|
|
}
|
|
if request.wantBodyNotContaining != "" && strings.Contains(body, request.wantBodyNotContaining) {
|
|
t.Errorf("body contains %q:\n%s", request.wantBodyNotContaining, body)
|
|
}
|
|
} else {
|
|
transform := cmpopts.AcyclicTransformer("jsonrpcid", func(id jsonrpc.ID) any { return id.Raw() })
|
|
if diff := cmp.Diff(request.wantMessages, got, transform); diff != "" {
|
|
t.Errorf("request #%d: received unexpected messages (-want +got):\n%s", i, diff)
|
|
}
|
|
}
|
|
sessionID.CompareAndSwap("", gotSessionID)
|
|
}
|
|
|
|
var wg sync.WaitGroup
|
|
for i, request := range requests {
|
|
if request.async || request.onRequest > 0 {
|
|
wg.Go(func() {
|
|
doStep(t, i, request)
|
|
})
|
|
} else {
|
|
doStep(t, i, request)
|
|
}
|
|
}
|
|
|
|
// Fail any blocked responses if they weren't needed by a synchronous
|
|
// request.
|
|
close(syncRequestsDone)
|
|
|
|
wg.Wait()
|
|
}
|
|
|
|
// A streamableRequest describes a single streamable HTTP request, consisting
|
|
// of a request payload and expected response.
|
|
type streamableRequest struct {
|
|
// If onRequest is > 0, this step only executes after a request with the
|
|
// given ID is received.
|
|
//
|
|
// All onRequest steps must occur before the step that creates the request.
|
|
//
|
|
// To avoid tests hanging when there's a bug, it's expected that this
|
|
// request is received in the course of a *synchronous* request to the
|
|
// server (otherwise, we wouldn't be able to terminate the test without
|
|
// analyzing a dependency graph).
|
|
onRequest int64
|
|
// If set, async causes the step to run asynchronously to other steps.
|
|
// Redundant with OnRequest: all OnRequest steps are asynchronous.
|
|
async bool
|
|
|
|
// Request attributes
|
|
method string // HTTP request method (required)
|
|
headers http.Header // additional headers to set, overlaid on top of the default headers
|
|
messages []jsonrpc.Message // messages to send
|
|
|
|
closeAfter int // if nonzero, close after receiving this many messages
|
|
wantStatusCode int // expected status code
|
|
wantBodyContaining string // if set, expect the response body to contain this text; overrides wantMessages
|
|
wantBodyNotContaining string // if set, a negative assertion on the body; overrides wantMessages
|
|
wantMessages []jsonrpc.Message // expected messages to receive; ignored if wantBodyContaining is set
|
|
wantSessionID bool // whether or not a session ID is expected in the response
|
|
}
|
|
|
|
// streamingRequest makes a request to the given streamable server with the
|
|
// given url, sessionID, and method.
|
|
//
|
|
// If provided, the in messages are encoded in the request body. A single
|
|
// message is encoded as a JSON object. Multiple messages are batched as a JSON
|
|
// array.
|
|
//
|
|
// Any received messages are sent to the out channel, which is closed when the
|
|
// request completes.
|
|
//
|
|
// Returns the sessionID and http status code from the response. If an error is
|
|
// returned, sessionID and status code may still be set if the error occurs
|
|
// after the response headers have been received.
|
|
func (s streamableRequest) do(ctx context.Context, serverURL, sessionID string, out chan<- jsonrpc.Message) (string, int, []byte, error) {
|
|
defer close(out)
|
|
|
|
var body []byte
|
|
if len(s.messages) == 1 {
|
|
data, err := jsonrpc2.EncodeMessage(s.messages[0])
|
|
if err != nil {
|
|
return "", 0, nil, fmt.Errorf("encoding message: %w", err)
|
|
}
|
|
body = data
|
|
} else {
|
|
var rawMsgs []json.RawMessage
|
|
for _, msg := range s.messages {
|
|
data, err := jsonrpc2.EncodeMessage(msg)
|
|
if err != nil {
|
|
return "", 0, nil, fmt.Errorf("encoding message: %w", err)
|
|
}
|
|
rawMsgs = append(rawMsgs, data)
|
|
}
|
|
data, err := json.Marshal(rawMsgs)
|
|
if err != nil {
|
|
return "", 0, nil, fmt.Errorf("marshaling batch: %w", err)
|
|
}
|
|
body = data
|
|
}
|
|
|
|
req, err := http.NewRequestWithContext(ctx, s.method, serverURL, bytes.NewReader(body))
|
|
if err != nil {
|
|
return "", 0, nil, fmt.Errorf("creating request: %w", err)
|
|
}
|
|
if sessionID != "" {
|
|
req.Header.Set(sessionIDHeader, sessionID)
|
|
}
|
|
if s.method == http.MethodPost {
|
|
req.Header.Set("Content-Type", "application/json")
|
|
}
|
|
req.Header.Set("Accept", "application/json, text/event-stream")
|
|
maps.Copy(req.Header, s.headers)
|
|
|
|
if req.Header.Get("Content-Type") == "" {
|
|
req.Header.Del("Content-Type")
|
|
}
|
|
|
|
resp, err := http.DefaultClient.Do(req)
|
|
if err != nil {
|
|
return "", 0, nil, fmt.Errorf("request failed: %v", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
newSessionID := resp.Header.Get(sessionIDHeader)
|
|
|
|
contentType := baseMediaType(resp.Header.Get("Content-Type"))
|
|
var respBody []byte
|
|
if contentType == "text/event-stream" {
|
|
r := readerInto{resp.Body, new(bytes.Buffer)}
|
|
for evt, err := range scanEventsLimited(r, DefaultMaxEventSize) {
|
|
if err != nil {
|
|
return newSessionID, resp.StatusCode, nil, fmt.Errorf("reading events: %v", err)
|
|
}
|
|
if evt.Name == "" || evt.Name == "message" { // ordinary message
|
|
// TODO(rfindley): do we need to check evt.name?
|
|
// Does the MCP spec say anything about this?
|
|
msg, err := jsonrpc2.DecodeMessage(evt.Data)
|
|
if err != nil {
|
|
return newSessionID, resp.StatusCode, nil, fmt.Errorf("decoding message: %w", err)
|
|
}
|
|
out <- msg
|
|
}
|
|
}
|
|
respBody = r.w.Bytes()
|
|
} else if contentType == "application/json" {
|
|
data, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return newSessionID, resp.StatusCode, nil, fmt.Errorf("reading json body: %w", err)
|
|
}
|
|
respBody = data
|
|
msg, err := jsonrpc2.DecodeMessage(data)
|
|
if err != nil {
|
|
return newSessionID, resp.StatusCode, nil, fmt.Errorf("decoding message: %w", err)
|
|
}
|
|
out <- msg
|
|
} else {
|
|
respBody, err = io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return newSessionID, resp.StatusCode, nil, fmt.Errorf("reading response: %v", err)
|
|
}
|
|
}
|
|
|
|
return newSessionID, resp.StatusCode, respBody, nil
|
|
}
|
|
|
|
// readerInto is an io.Reader that writes any bytes read from r into w.
|
|
type readerInto struct {
|
|
r io.Reader
|
|
w *bytes.Buffer
|
|
}
|
|
|
|
// Read implements io.Reader.
|
|
func (r readerInto) Read(p []byte) (n int, err error) {
|
|
n, err = r.r.Read(p)
|
|
if err == nil || err == io.EOF {
|
|
n2, err2 := r.w.Write(p[:n])
|
|
if err2 != nil {
|
|
return n, fmt.Errorf("failed to write: %v", err)
|
|
}
|
|
if n2 != n {
|
|
return n, fmt.Errorf("short write: %d != %d", n2, n)
|
|
}
|
|
}
|
|
return n, err
|
|
}
|
|
|
|
func mustMarshal(v any) json.RawMessage {
|
|
if v == nil {
|
|
return nil
|
|
}
|
|
data, err := json.Marshal(v)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
return data
|
|
}
|
|
|
|
func TestEventID(t *testing.T) {
|
|
tests := []struct {
|
|
sid string
|
|
idx int
|
|
}{
|
|
{"0", 0},
|
|
{"0", 1},
|
|
{"1", 0},
|
|
{"1", 1},
|
|
{"", 1},
|
|
{"1234", 5678},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
t.Run(fmt.Sprintf("%s_%d", test.sid, test.idx), func(t *testing.T) {
|
|
eventID := formatEventID(test.sid, test.idx)
|
|
gotSID, gotIdx, ok := parseEventID(eventID)
|
|
if !ok {
|
|
t.Fatalf("parseEventID(%q) failed, want ok", eventID)
|
|
}
|
|
if gotSID != test.sid || gotIdx != test.idx {
|
|
t.Errorf("parseEventID(%q) = %s, %d, want %s, %d", eventID, gotSID, gotIdx, test.sid, test.idx)
|
|
}
|
|
})
|
|
}
|
|
|
|
invalid := []string{
|
|
"",
|
|
"_",
|
|
"1_",
|
|
"1_a",
|
|
"1_-1",
|
|
}
|
|
|
|
for _, eventID := range invalid {
|
|
t.Run(fmt.Sprintf("invalid_%q", eventID), func(t *testing.T) {
|
|
if _, _, ok := parseEventID(eventID); ok {
|
|
t.Errorf("parseEventID(%q) succeeded, want failure", eventID)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestStreamableStateless(t *testing.T) {
|
|
initReq := req(1, methodInitialize, &InitializeParams{})
|
|
initResp := resp(1, &InitializeResult{
|
|
Capabilities: &ServerCapabilities{
|
|
Logging: &LoggingCapabilities{},
|
|
Tools: &ToolCapabilities{ListChanged: true},
|
|
},
|
|
ProtocolVersion: protocolVersion20251125,
|
|
ServerInfo: &Implementation{Name: "test", Version: "v1.0.0"},
|
|
}, nil)
|
|
// This version of sayHi expects
|
|
// that request from our client).
|
|
sayHi := func(ctx context.Context, req *CallToolRequest, args hiParams) (*CallToolResult, any, 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 " + args.Name}}}, nil, nil
|
|
}
|
|
|
|
requests := []streamableRequest{
|
|
{
|
|
method: "POST",
|
|
messages: []jsonrpc.Message{initReq},
|
|
wantStatusCode: http.StatusOK,
|
|
wantMessages: []jsonrpc.Message{initResp},
|
|
wantSessionID: false, // sessionless
|
|
},
|
|
{
|
|
method: "POST",
|
|
wantStatusCode: http.StatusOK,
|
|
messages: []jsonrpc.Message{req(1, "tools/list", struct{}{})},
|
|
wantBodyContaining: "greet",
|
|
},
|
|
{
|
|
method: "GET",
|
|
wantStatusCode: http.StatusMethodNotAllowed,
|
|
},
|
|
{
|
|
method: "POST",
|
|
wantStatusCode: http.StatusOK,
|
|
messages: []jsonrpc.Message{
|
|
req(2, "tools/call", &CallToolParams{Name: "greet", Arguments: hiParams{Name: "World"}}),
|
|
},
|
|
wantMessages: []jsonrpc.Message{
|
|
resp(2, &CallToolResult{
|
|
Content: []Content{&TextContent{Text: "hi World"}},
|
|
}, nil),
|
|
},
|
|
},
|
|
{
|
|
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),
|
|
},
|
|
},
|
|
}
|
|
|
|
testClientCompatibility := func(t *testing.T, handler http.Handler) {
|
|
ctx := context.Background()
|
|
httpServer := httptest.NewServer(mustNotPanic(t, handler))
|
|
defer httpServer.Close()
|
|
cs, err := NewClient(testImpl, nil).Connect(ctx, &StreamableClientTransport{Endpoint: httpServer.URL}, nil)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
t.Cleanup(func() { cs.Close() })
|
|
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)
|
|
}
|
|
}
|
|
|
|
statelessHandler := NewStreamableHTTPHandler(func(*http.Request) *Server {
|
|
server := NewServer(testImpl, nil)
|
|
AddTool(server, &Tool{Name: "greet", Description: "say hi"}, sayHi)
|
|
return server
|
|
}, &StreamableHTTPOptions{
|
|
Stateless: true,
|
|
})
|
|
|
|
testStreamableHandler(t, statelessHandler, requests)
|
|
testClientCompatibility(t, statelessHandler)
|
|
}
|
|
|
|
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
|
|
}
|
|
|
|
func TestSessionHijackingPrevention(t *testing.T) {
|
|
// This test verifies that sessions bound to a user ID cannot be accessed
|
|
// by a different user (session hijacking prevention).
|
|
ctx := context.Background()
|
|
|
|
server := NewServer(testImpl, nil)
|
|
streamHandler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, nil)
|
|
|
|
// Use the bearer token directly as the user ID. This simulates how a real
|
|
// verifier might extract a user ID from a JWT "sub" claim or introspection.
|
|
verifier := func(_ context.Context, token string, _ *http.Request) (*auth.TokenInfo, error) {
|
|
return &auth.TokenInfo{
|
|
Scopes: []string{"scope"},
|
|
UserID: token,
|
|
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()
|
|
|
|
// Helper to send a JSON-RPC request as a given user.
|
|
doRequest := func(msg jsonrpc.Message, sessionID, userID string) *http.Response {
|
|
t.Helper()
|
|
data, _ := jsonrpc2.EncodeMessage(msg)
|
|
req, _ := http.NewRequestWithContext(ctx, http.MethodPost, httpServer.URL, bytes.NewReader(data))
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("Accept", "application/json, text/event-stream")
|
|
req.Header.Set("Authorization", "Bearer "+userID)
|
|
if sessionID != "" {
|
|
req.Header.Set(sessionIDHeader, sessionID)
|
|
}
|
|
resp, err := http.DefaultClient.Do(req)
|
|
if err != nil {
|
|
t.Fatalf("request failed: %v", err)
|
|
}
|
|
return resp
|
|
}
|
|
|
|
// Create a session as user1.
|
|
initReq := &jsonrpc.Request{Method: "initialize", ID: jsonrpc2.Int64ID(1)}
|
|
initReq.Params, _ = json.Marshal(&InitializeParams{
|
|
ProtocolVersion: protocolVersion20250618,
|
|
ClientInfo: &Implementation{Name: "test", Version: "1.0"},
|
|
})
|
|
resp := doRequest(initReq, "", "user1")
|
|
defer resp.Body.Close()
|
|
if resp.StatusCode != http.StatusOK {
|
|
body, _ := io.ReadAll(resp.Body)
|
|
t.Fatalf("initialize failed with status %d: %s", resp.StatusCode, body)
|
|
}
|
|
sessionID := resp.Header.Get(sessionIDHeader)
|
|
if sessionID == "" {
|
|
t.Fatal("no session ID in response")
|
|
}
|
|
|
|
pingReq := &jsonrpc.Request{Method: "ping", ID: jsonrpc2.Int64ID(2)}
|
|
pingReq.Params, _ = json.Marshal(&PingParams{})
|
|
|
|
// Try to access the session as user2 - should fail.
|
|
resp2 := doRequest(pingReq, sessionID, "user2")
|
|
defer resp2.Body.Close()
|
|
if resp2.StatusCode != http.StatusForbidden {
|
|
body, _ := io.ReadAll(resp2.Body)
|
|
t.Errorf("expected status %d for user mismatch, got %d: %s", http.StatusForbidden, resp2.StatusCode, body)
|
|
}
|
|
|
|
// Access as original user1 should succeed.
|
|
resp3 := doRequest(pingReq, sessionID, "user1")
|
|
defer resp3.Body.Close()
|
|
if resp3.StatusCode != http.StatusOK {
|
|
body, _ := io.ReadAll(resp3.Body)
|
|
t.Errorf("expected status %d for matching user, got %d: %s", http.StatusOK, resp3.StatusCode, body)
|
|
}
|
|
}
|
|
|
|
func TestStreamableGET(t *testing.T) {
|
|
// This test checks the fix for problematic behavior described in #410:
|
|
// Hanging GET headers should be written immediately, even if there are no
|
|
// messages.
|
|
server := NewServer(testImpl, nil)
|
|
|
|
handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, nil)
|
|
httpServer := httptest.NewServer(mustNotPanic(t, handler))
|
|
defer httpServer.Close()
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
|
|
newReq := func(method string, msg jsonrpc.Message) *http.Request {
|
|
var body io.Reader
|
|
if msg != nil {
|
|
data, err := jsonrpc2.EncodeMessage(msg)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
body = bytes.NewReader(data)
|
|
}
|
|
req, err := http.NewRequestWithContext(ctx, method, httpServer.URL, body)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
req.Header.Set("Accept", "application/json, text/event-stream")
|
|
if msg != nil {
|
|
req.Header.Set("Content-Type", "application/json")
|
|
}
|
|
return req
|
|
}
|
|
|
|
get1 := newReq(http.MethodGet, nil)
|
|
resp, err := http.DefaultClient.Do(get1)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
// GET without session should return 400 Bad Request (not 405) because
|
|
// GET is a valid method - it just requires a session ID.
|
|
if got, want := resp.StatusCode, http.StatusBadRequest; got != want {
|
|
t.Errorf("initial GET: got status %d, want %d", got, want)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
post1 := newReq(http.MethodPost, req(1, methodInitialize, &InitializeParams{}))
|
|
resp, err = http.DefaultClient.Do(post1)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer resp.Body.Close()
|
|
if got, want := resp.StatusCode, http.StatusOK; got != want {
|
|
body, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
t.Errorf("initialize POST: got status %d, want %d; body:\n%s", got, want, string(body))
|
|
}
|
|
|
|
sessionID := resp.Header.Get(sessionIDHeader)
|
|
if sessionID == "" {
|
|
t.Fatalf("initialized missing session ID")
|
|
}
|
|
|
|
get2 := newReq("GET", nil)
|
|
get2.Header.Set(sessionIDHeader, sessionID)
|
|
resp, err = http.DefaultClient.Do(get2)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer resp.Body.Close()
|
|
if got, want := resp.StatusCode, http.StatusOK; got != want {
|
|
t.Errorf("GET with session ID: got status %d, want %d", got, want)
|
|
}
|
|
|
|
t.Log("Sending final DELETE request to close session and release resources")
|
|
del := newReq("DELETE", nil)
|
|
del.Header.Set(sessionIDHeader, sessionID)
|
|
resp, err = http.DefaultClient.Do(del)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer resp.Body.Close()
|
|
if got, want := resp.StatusCode, http.StatusNoContent; got != want {
|
|
t.Errorf("DELETE with session ID: got status %d, want %d", got, want)
|
|
}
|
|
}
|
|
|
|
// TestStreamableMcpHeaderValidation tests the server-side Mcp-Method,
|
|
// Mcp-Name, and Mcp-Param header validation through the full HTTP handler,
|
|
// as specified in SEP-2243.
|
|
func TestStreamableMcpHeaderValidation(t *testing.T) {
|
|
// Temporarily register the future version so the handler accepts it.
|
|
|
|
server := NewServer(&Implementation{Name: "testServer", Version: "v1.0.0"}, nil)
|
|
server.AddTool(
|
|
&Tool{Name: "my-tool", InputSchema: &jsonschema.Schema{Type: "object"}},
|
|
func(ctx context.Context, req *CallToolRequest) (*CallToolResult, error) {
|
|
return &CallToolResult{}, nil
|
|
})
|
|
server.AddTool(
|
|
&Tool{
|
|
Name: "execute_sql",
|
|
InputSchema: map[string]any{
|
|
"type": "object",
|
|
"properties": map[string]any{
|
|
"region": map[string]any{
|
|
"type": "string",
|
|
"x-mcp-header": "Region",
|
|
},
|
|
"query": map[string]any{
|
|
"type": "string",
|
|
},
|
|
},
|
|
},
|
|
},
|
|
func(ctx context.Context, req *CallToolRequest) (*CallToolResult, error) {
|
|
return &CallToolResult{}, nil
|
|
})
|
|
|
|
handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, &StreamableHTTPOptions{
|
|
Stateless: true,
|
|
})
|
|
defer handler.closeAll()
|
|
|
|
testMeta := Meta{
|
|
MetaKeyProtocolVersion: minVersionForStandardHeaders,
|
|
MetaKeyClientInfo: map[string]any{"name": "testClient", "version": "v1.0.0"},
|
|
MetaKeyClientCapabilities: map[string]any{},
|
|
}
|
|
serverImpl := &Implementation{Name: "testServer", Version: "v1.0.0"}
|
|
|
|
testStreamableHandler(t, handler, []streamableRequest{
|
|
{
|
|
method: "POST",
|
|
headers: http.Header{
|
|
protocolVersionHeader: {minVersionForStandardHeaders},
|
|
methodHeader: {"tools/call"},
|
|
nameHeader: {"my-tool"},
|
|
},
|
|
messages: []jsonrpc.Message{req(2, "tools/call", &CallToolParams{Meta: testMeta, Name: "my-tool"})},
|
|
wantStatusCode: http.StatusOK,
|
|
wantMessages: []jsonrpc.Message{resp(2, completeCallToolResultWithServerInfo(serverImpl), nil)},
|
|
},
|
|
{
|
|
method: "POST",
|
|
headers: http.Header{
|
|
protocolVersionHeader: {minVersionForStandardHeaders},
|
|
methodHeader: {"prompts/get"},
|
|
nameHeader: {"my-tool"},
|
|
},
|
|
messages: []jsonrpc.Message{req(3, "tools/call", &CallToolParams{Meta: testMeta, Name: "my-tool"})},
|
|
wantStatusCode: http.StatusBadRequest,
|
|
wantBodyContaining: "Mcp-Method header value",
|
|
},
|
|
{
|
|
method: "POST",
|
|
headers: http.Header{
|
|
protocolVersionHeader: {minVersionForStandardHeaders},
|
|
methodHeader: {"tools/call"},
|
|
nameHeader: {"wrong-tool"},
|
|
},
|
|
messages: []jsonrpc.Message{req(4, "tools/call", &CallToolParams{Meta: testMeta, Name: "my-tool"})},
|
|
wantStatusCode: http.StatusBadRequest,
|
|
wantBodyContaining: "Mcp-Name header value",
|
|
},
|
|
{
|
|
method: "POST",
|
|
headers: http.Header{
|
|
protocolVersionHeader: {minVersionForStandardHeaders},
|
|
methodHeader: {"TOOLS/CALL"},
|
|
nameHeader: {"my-tool"},
|
|
},
|
|
messages: []jsonrpc.Message{req(5, "tools/call", &CallToolParams{Meta: testMeta, Name: "my-tool"})},
|
|
wantStatusCode: http.StatusBadRequest,
|
|
wantBodyContaining: "Mcp-Method header value",
|
|
},
|
|
{
|
|
method: "POST",
|
|
headers: http.Header{
|
|
protocolVersionHeader: {minVersionForStandardHeaders},
|
|
methodHeader: {"tools/call"},
|
|
nameHeader: {"my-tool"},
|
|
},
|
|
messages: []jsonrpc.Message{req(6, "tools/call", &CallToolParams{Meta: testMeta, Name: "my-tool"})},
|
|
wantStatusCode: http.StatusOK,
|
|
wantMessages: []jsonrpc.Message{resp(6, completeCallToolResultWithServerInfo(serverImpl), nil)},
|
|
},
|
|
{
|
|
method: "POST",
|
|
headers: http.Header{
|
|
protocolVersionHeader: {minVersionForStandardHeaders},
|
|
methodHeader: {"tools/call"},
|
|
nameHeader: {"execute_sql"},
|
|
paramHeaderPrefix + "Region": {"us-west1"},
|
|
},
|
|
messages: []jsonrpc.Message{req(7, "tools/call", &CallToolParams{
|
|
Meta: testMeta,
|
|
Name: "execute_sql",
|
|
Arguments: map[string]any{"region": "us-west1", "query": "SELECT 1"},
|
|
})},
|
|
wantStatusCode: http.StatusOK,
|
|
wantMessages: []jsonrpc.Message{resp(7, completeCallToolResultWithServerInfo(serverImpl), nil)},
|
|
},
|
|
{
|
|
method: "POST",
|
|
headers: http.Header{
|
|
protocolVersionHeader: {minVersionForStandardHeaders},
|
|
methodHeader: {"tools/call"},
|
|
nameHeader: {"execute_sql"},
|
|
paramHeaderPrefix + "Region": {"eu-central1"},
|
|
},
|
|
messages: []jsonrpc.Message{req(8, "tools/call", &CallToolParams{
|
|
Meta: testMeta,
|
|
Name: "execute_sql",
|
|
Arguments: map[string]any{"region": "us-west1"},
|
|
})},
|
|
wantStatusCode: http.StatusBadRequest,
|
|
wantBodyContaining: "header mismatch",
|
|
},
|
|
{
|
|
method: "POST",
|
|
headers: http.Header{
|
|
protocolVersionHeader: {minVersionForStandardHeaders},
|
|
methodHeader: {"tools/call"},
|
|
nameHeader: {"execute_sql"},
|
|
},
|
|
messages: []jsonrpc.Message{req(9, "tools/call", &CallToolParams{
|
|
Meta: testMeta,
|
|
Name: "execute_sql",
|
|
Arguments: map[string]any{"region": "us-west1"},
|
|
})},
|
|
wantStatusCode: http.StatusBadRequest,
|
|
wantBodyContaining: "missing",
|
|
},
|
|
})
|
|
}
|
|
|
|
// TestStreamableMcpHeaderValidationErrorFormat verifies that header
|
|
// validation errors return a JSON-RPC error with code -32020 and
|
|
// Content-Type application/json, per SEP-2243.
|
|
func TestStreamableMcpHeaderValidationErrorFormat(t *testing.T) {
|
|
|
|
server := NewServer(&Implementation{Name: "testServer", Version: "v1.0.0"}, nil)
|
|
server.AddTool(
|
|
&Tool{Name: "my-tool", InputSchema: &jsonschema.Schema{Type: "object"}},
|
|
func(ctx context.Context, req *CallToolRequest) (*CallToolResult, error) {
|
|
return &CallToolResult{}, nil
|
|
})
|
|
|
|
handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, &StreamableHTTPOptions{
|
|
Stateless: true,
|
|
})
|
|
defer handler.closeAll()
|
|
|
|
httpServer := httptest.NewServer(mustNotPanic(t, handler))
|
|
defer httpServer.Close()
|
|
|
|
// Use the MCP client with a custom RoundTripper to inject a bad header.
|
|
var toolCallResp *http.Response
|
|
var toolCallRespBody []byte
|
|
|
|
customClient := &http.Client{
|
|
Transport: roundTripperFunc(func(req *http.Request) (*http.Response, error) {
|
|
var originalMethodHeader string
|
|
if req.Header.Get(methodHeader) == "tools/call" {
|
|
originalMethodHeader = req.Header.Get(methodHeader)
|
|
req.Header.Set(methodHeader, "wrong-method")
|
|
}
|
|
resp, err := http.DefaultTransport.RoundTrip(req)
|
|
if err == nil && originalMethodHeader == "tools/call" {
|
|
toolCallResp = resp
|
|
toolCallRespBody, _ = io.ReadAll(resp.Body)
|
|
}
|
|
return resp, err
|
|
}),
|
|
}
|
|
|
|
clientTransport := &StreamableClientTransport{
|
|
Endpoint: httpServer.URL,
|
|
HTTPClient: customClient,
|
|
}
|
|
|
|
client := NewClient(&Implementation{Name: "testClient", Version: "v1.0.0"}, nil)
|
|
ctx := context.Background()
|
|
session, err := client.Connect(ctx, clientTransport, &ClientSessionOptions{ProtocolVersion: minVersionForStandardHeaders})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer session.Close()
|
|
|
|
_, err = session.CallTool(ctx, &CallToolParams{Name: "my-tool"})
|
|
// We expect an error because the server should reject it.
|
|
if err == nil {
|
|
t.Error("CallTool succeeded unexpectedly")
|
|
}
|
|
|
|
if toolCallResp == nil {
|
|
t.Fatal("no response captured")
|
|
}
|
|
|
|
// Verify HTTP status code.
|
|
if toolCallResp.StatusCode != http.StatusBadRequest {
|
|
t.Errorf("status code = %d, want %d", toolCallResp.StatusCode, http.StatusBadRequest)
|
|
}
|
|
|
|
// Verify Content-Type.
|
|
if baseMediaType(toolCallResp.Header.Get("Content-Type")) != "application/json" {
|
|
t.Errorf("Content-Type = %q, want %q", baseMediaType(toolCallResp.Header.Get("Content-Type")), "application/json")
|
|
}
|
|
|
|
// Verify JSON-RPC error body contains error code -32020.
|
|
msg, err := jsonrpc2.DecodeMessage(toolCallRespBody)
|
|
if err != nil {
|
|
t.Fatalf("failed to decode message: %v", err)
|
|
}
|
|
resp, ok := msg.(*jsonrpc2.Response)
|
|
if !ok {
|
|
t.Fatalf("expected *jsonrpc2.Response, got %T", msg)
|
|
}
|
|
var wireErr *jsonrpc2.WireError
|
|
if !errors.As(resp.Error, &wireErr) {
|
|
t.Fatalf("expected *jsonrpc2.WireError, got %T", resp.Error)
|
|
}
|
|
if wireErr.Code != CodeHeaderMismatch {
|
|
t.Errorf("wireErr.Code = %d, want %d", wireErr.Code, CodeHeaderMismatch)
|
|
}
|
|
if !strings.Contains(wireErr.Message, "Mcp-Method header value") {
|
|
t.Errorf("wireErr.Message = %q, want it to contain %q", wireErr.Message, "Mcp-Method header value")
|
|
}
|
|
}
|
|
|
|
// TestStreamableMcpHeaderVersionGating verifies that header validation
|
|
// is skipped for protocol versions older than minVersionForStandardHeaders.
|
|
func TestStreamableMcpHeaderVersionGating(t *testing.T) {
|
|
server := NewServer(&Implementation{Name: "testServer", Version: "v1.0.0"}, nil)
|
|
server.AddTool(
|
|
&Tool{Name: "my-tool", InputSchema: &jsonschema.Schema{Type: "object"}},
|
|
func(ctx context.Context, req *CallToolRequest) (*CallToolResult, error) {
|
|
return &CallToolResult{}, nil
|
|
})
|
|
|
|
handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, nil)
|
|
defer handler.closeAll()
|
|
|
|
initReq := req(1, methodInitialize, &InitializeParams{ProtocolVersion: protocolVersion20251125})
|
|
initResp := resp(1, &InitializeResult{
|
|
Capabilities: &ServerCapabilities{
|
|
Logging: &LoggingCapabilities{},
|
|
Tools: &ToolCapabilities{ListChanged: true},
|
|
},
|
|
ProtocolVersion: protocolVersion20251125,
|
|
ServerInfo: &Implementation{Name: "testServer", Version: "v1.0.0"},
|
|
}, nil)
|
|
|
|
testStreamableHandler(t, handler, []streamableRequest{
|
|
{
|
|
method: "POST",
|
|
messages: []jsonrpc.Message{initReq},
|
|
wantStatusCode: http.StatusOK,
|
|
wantMessages: []jsonrpc.Message{initResp},
|
|
wantSessionID: true,
|
|
},
|
|
{
|
|
method: "POST",
|
|
headers: http.Header{protocolVersionHeader: {protocolVersion20251125}},
|
|
messages: []jsonrpc.Message{req(0, notificationInitialized, &InitializedParams{})},
|
|
wantStatusCode: http.StatusAccepted,
|
|
},
|
|
// Requests with deliberately wrong MCP headers should still succeed
|
|
// because the protocol version is too old for validation.
|
|
{
|
|
method: "POST",
|
|
headers: http.Header{
|
|
protocolVersionHeader: {protocolVersion20251125},
|
|
methodHeader: {"wrong-method"},
|
|
nameHeader: {"wrong-name"},
|
|
},
|
|
messages: []jsonrpc.Message{req(2, "tools/call", &CallToolParams{Name: "my-tool"})},
|
|
wantStatusCode: http.StatusOK,
|
|
wantMessages: []jsonrpc.Message{resp(2, &CallToolResult{Content: []Content{}}, nil)},
|
|
},
|
|
})
|
|
}
|
|
|
|
// TestStreamableParamHeadersClientSetsHeaders verifies that the client sets
|
|
// Mcp-Param-* headers on tool calls when the tool has x-mcp-header annotations.
|
|
func TestStreamableParamHeadersClientSetsHeaders(t *testing.T) {
|
|
|
|
server := NewServer(&Implementation{Name: "testServer", Version: "v1.0.0"}, nil)
|
|
server.AddTool(
|
|
&Tool{
|
|
Name: "execute_sql",
|
|
InputSchema: map[string]any{
|
|
"type": "object",
|
|
"properties": map[string]any{
|
|
"region": map[string]any{
|
|
"type": "string",
|
|
"x-mcp-header": "Region",
|
|
},
|
|
"query": map[string]any{
|
|
"type": "string",
|
|
},
|
|
},
|
|
},
|
|
},
|
|
func(ctx context.Context, req *CallToolRequest) (*CallToolResult, error) {
|
|
return &CallToolResult{Content: []Content{&TextContent{Text: "ok"}}}, nil
|
|
})
|
|
|
|
// Install a receiving middleware on the server that opts the tools/list
|
|
// response caching by setting a positive ttlMs.
|
|
server.AddReceivingMiddleware(func(next MethodHandler) MethodHandler {
|
|
return func(ctx context.Context, method string, req Request) (Result, error) {
|
|
res, err := next(ctx, method, req)
|
|
if err == nil && method == methodListTools {
|
|
if r, ok := res.(*ListToolsResult); ok {
|
|
r.TTLMs = 60_000
|
|
}
|
|
}
|
|
return res, err
|
|
}
|
|
})
|
|
|
|
handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, &StreamableHTTPOptions{
|
|
Stateless: true,
|
|
})
|
|
defer handler.closeAll()
|
|
httpServer := httptest.NewServer(mustNotPanic(t, handler))
|
|
defer httpServer.Close()
|
|
|
|
var capturedHeaders http.Header
|
|
customClient := &http.Client{
|
|
Transport: roundTripperFunc(func(req *http.Request) (*http.Response, error) {
|
|
if req.Header.Get(methodHeader) == "tools/call" {
|
|
capturedHeaders = req.Header.Clone()
|
|
}
|
|
return http.DefaultTransport.RoundTrip(req)
|
|
}),
|
|
}
|
|
|
|
clientTransport := &StreamableClientTransport{
|
|
Endpoint: httpServer.URL,
|
|
HTTPClient: customClient,
|
|
}
|
|
|
|
client := NewClient(&Implementation{Name: "testClient", Version: "v1.0.0"}, nil)
|
|
ctx := context.Background()
|
|
session, err := client.Connect(ctx, clientTransport, &ClientSessionOptions{ProtocolVersion: minVersionForStandardHeaders})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer session.Close()
|
|
|
|
// ListTools to populate the tool cache (needed for param headers).
|
|
// Pass a non-nil params so the SEP-2575 per-request _meta triple is
|
|
// injected; injectMeta is a no-op when params is nil.
|
|
if _, err := session.ListTools(ctx, &ListToolsParams{}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
_, err = session.CallTool(ctx, &CallToolParams{
|
|
Name: "execute_sql",
|
|
Arguments: map[string]any{"region": "us-west1", "query": "SELECT 1"},
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
if capturedHeaders == nil {
|
|
t.Fatal("no tool call headers captured")
|
|
}
|
|
if got := capturedHeaders.Get(methodHeader); got != "tools/call" {
|
|
t.Errorf("Mcp-Method = %q, want %q", got, "tools/call")
|
|
}
|
|
if got := capturedHeaders.Get(nameHeader); got != "execute_sql" {
|
|
t.Errorf("Mcp-Name = %q, want %q", got, "execute_sql")
|
|
}
|
|
if got := capturedHeaders.Get(paramHeaderPrefix + "Region"); got != "us-west1" {
|
|
t.Errorf("Mcp-Param-Region = %q, want %q", got, "us-west1")
|
|
}
|
|
}
|
|
|
|
// TestStreamableFilterValidToolsIntegration verifies that AddTool rejects
|
|
// tools with invalid x-mcp-header annotations at registration time, and
|
|
// that valid tools are returned by ListTools.
|
|
func TestStreamableFilterValidToolsIntegration(t *testing.T) {
|
|
|
|
server := NewServer(&Implementation{Name: "testServer", Version: "v1.0.0"}, nil)
|
|
noop := func(ctx context.Context, req *CallToolRequest) (*CallToolResult, error) {
|
|
return &CallToolResult{}, nil
|
|
}
|
|
|
|
// Valid tool with correct x-mcp-header annotation.
|
|
server.AddTool(&Tool{
|
|
Name: "valid-tool",
|
|
InputSchema: map[string]any{
|
|
"type": "object",
|
|
"properties": map[string]any{
|
|
"region": map[string]any{
|
|
"type": "string",
|
|
"x-mcp-header": "Region",
|
|
},
|
|
},
|
|
},
|
|
}, noop)
|
|
|
|
// Invalid tool: x-mcp-header on an array type should panic.
|
|
func() {
|
|
defer func() {
|
|
r := recover()
|
|
if r == nil {
|
|
t.Fatal("AddTool with invalid x-mcp-header annotation did not panic")
|
|
}
|
|
}()
|
|
server.AddTool(&Tool{
|
|
Name: "invalid-tool",
|
|
InputSchema: map[string]any{
|
|
"type": "object",
|
|
"properties": map[string]any{
|
|
"items": map[string]any{
|
|
"type": "array",
|
|
"x-mcp-header": "Items",
|
|
},
|
|
},
|
|
},
|
|
}, noop)
|
|
}()
|
|
|
|
// Tool with no x-mcp-header annotations (always valid).
|
|
server.AddTool(&Tool{
|
|
Name: "plain-tool",
|
|
InputSchema: &jsonschema.Schema{Type: "object"},
|
|
}, noop)
|
|
|
|
handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, &StreamableHTTPOptions{
|
|
Stateless: true,
|
|
})
|
|
defer handler.closeAll()
|
|
httpServer := httptest.NewServer(mustNotPanic(t, handler))
|
|
defer httpServer.Close()
|
|
|
|
client := NewClient(&Implementation{Name: "testClient", Version: "v1.0.0"}, nil)
|
|
ctx := context.Background()
|
|
session, err := client.Connect(ctx, &StreamableClientTransport{
|
|
Endpoint: httpServer.URL,
|
|
}, &ClientSessionOptions{ProtocolVersion: minVersionForStandardHeaders})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer session.Close()
|
|
|
|
// Pass non-nil params so the SEP-2575 per-request _meta triple is
|
|
// injected; injectMeta is a no-op when params is nil.
|
|
result, err := session.ListTools(ctx, &ListToolsParams{})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
toolNames := make([]string, len(result.Tools))
|
|
for i, tool := range result.Tools {
|
|
toolNames[i] = tool.Name
|
|
}
|
|
sort.Strings(toolNames)
|
|
|
|
wantNames := []string{"plain-tool", "valid-tool"}
|
|
if !slices.Equal(toolNames, wantNames) {
|
|
t.Errorf("ListTools returned %v, want %v", toolNames, wantNames)
|
|
}
|
|
}
|
|
|
|
// TestStreamable405AllowHeader verifies RFC 9110 §15.5.6 compliance:
|
|
// 405 Method Not Allowed responses MUST include an Allow header.
|
|
func TestStreamable405AllowHeader(t *testing.T) {
|
|
server := NewServer(testImpl, nil)
|
|
|
|
tests := []struct {
|
|
name string
|
|
stateless bool
|
|
method string
|
|
wantStatus int
|
|
wantAllow string
|
|
}{
|
|
{
|
|
name: "unsupported method stateful",
|
|
stateless: false,
|
|
method: "PUT",
|
|
wantStatus: http.StatusMethodNotAllowed,
|
|
wantAllow: "GET, POST, DELETE",
|
|
},
|
|
{
|
|
name: "GET in stateless mode",
|
|
stateless: true,
|
|
method: "GET",
|
|
wantStatus: http.StatusMethodNotAllowed,
|
|
wantAllow: "POST",
|
|
},
|
|
{
|
|
// In stateless mode, only POST is supported.
|
|
// DELETE returns 405, consistent with Allow: POST.
|
|
name: "DELETE without session stateless",
|
|
stateless: true,
|
|
method: "DELETE",
|
|
wantStatus: http.StatusMethodNotAllowed,
|
|
wantAllow: "POST",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
opts := &StreamableHTTPOptions{Stateless: tt.stateless}
|
|
handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, opts)
|
|
httpServer := httptest.NewServer(mustNotPanic(t, handler))
|
|
defer httpServer.Close()
|
|
|
|
req, err := http.NewRequest(tt.method, httpServer.URL, nil)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
req.Header.Set("Accept", "application/json, text/event-stream")
|
|
|
|
resp, err := http.DefaultClient.Do(req)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if got := resp.StatusCode; got != tt.wantStatus {
|
|
t.Errorf("status code: got %d, want %d", got, tt.wantStatus)
|
|
}
|
|
|
|
allow := resp.Header.Get("Allow")
|
|
if allow != tt.wantAllow {
|
|
t.Errorf("Allow header: got %q, want %q", allow, tt.wantAllow)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestStreamableGETWithoutSession verifies that GET without session ID in stateful mode
|
|
// returns 400 Bad Request (not 405), since GET is a supported method that requires a session.
|
|
func TestStreamableGETWithoutSession(t *testing.T) {
|
|
server := NewServer(testImpl, nil)
|
|
handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, nil)
|
|
httpServer := httptest.NewServer(mustNotPanic(t, handler))
|
|
defer httpServer.Close()
|
|
|
|
req, err := http.NewRequest("GET", httpServer.URL, nil)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
req.Header.Set("Accept", "text/event-stream")
|
|
|
|
resp, err := http.DefaultClient.Do(req)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer resp.Body.Close()
|
|
body, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
// GET without session should return 400 Bad Request, not 405 Method Not Allowed,
|
|
// because GET is a valid method - it just requires a session ID.
|
|
if got, want := resp.StatusCode, http.StatusBadRequest; got != want {
|
|
t.Errorf("status code: got %d, want %d", got, want)
|
|
}
|
|
if got, want := strings.TrimSpace(string(body)), "Bad Request: GET requires an Mcp-Session-Id header"; got != want {
|
|
t.Errorf("body: got %q, want %q", got, want)
|
|
}
|
|
}
|
|
|
|
func TestStreamableGETWithoutEventStreamAccept(t *testing.T) {
|
|
server := NewServer(testImpl, nil)
|
|
handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, nil)
|
|
httpServer := httptest.NewServer(mustNotPanic(t, handler))
|
|
defer httpServer.Close()
|
|
|
|
req, err := http.NewRequest("GET", httpServer.URL, nil)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
req.Header.Set("Accept", "application/json")
|
|
|
|
resp, err := http.DefaultClient.Do(req)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer resp.Body.Close()
|
|
body, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
if got, want := resp.StatusCode, http.StatusBadRequest; got != want {
|
|
t.Errorf("status code: got %d, want %d", got, want)
|
|
}
|
|
if got, want := strings.TrimSpace(string(body)), "Accept must contain 'text/event-stream' for GET requests"; got != want {
|
|
t.Errorf("body: got %q, want %q", got, want)
|
|
}
|
|
}
|
|
|
|
func TestBaseMediaType(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
value string
|
|
want string
|
|
}{
|
|
{name: "empty", want: ""},
|
|
{name: "json", value: "application/json", want: "application/json"},
|
|
{name: "json with params", value: "Application/JSON; charset=utf-8", want: "application/json"},
|
|
{name: "event stream with params", value: "Text/Event-Stream; charset=utf-8", want: "text/event-stream"},
|
|
{name: "invalid", value: "application/json; charset", want: ""},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
if got := baseMediaType(test.value); got != test.want {
|
|
t.Errorf("baseMediaType(%q) = %q, want %q", test.value, got, test.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestStreamableClientContextPropagation(t *testing.T) {
|
|
type contextKey string
|
|
const testKey = contextKey("test-key")
|
|
const testValue = "test-value"
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
ctx2 := context.WithValue(ctx, testKey, testValue)
|
|
|
|
server := httptest.NewServer(mustNotPanic(t, http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
|
|
switch req.Method {
|
|
case "POST":
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.Header().Set(sessionIDHeader, "test-session")
|
|
w.WriteHeader(http.StatusOK)
|
|
w.Write([]byte(`{"jsonrpc":"2.0","id":1,"result":{"protocolVersion":"2025-03-26","capabilities":{},"serverInfo":{"name":"test","version":"1.0"}}}`))
|
|
case "GET":
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
w.WriteHeader(http.StatusOK)
|
|
case "DELETE":
|
|
w.WriteHeader(http.StatusNoContent)
|
|
}
|
|
})))
|
|
defer server.Close()
|
|
|
|
transport := &StreamableClientTransport{Endpoint: server.URL}
|
|
conn, err := transport.Connect(ctx2)
|
|
if err != nil {
|
|
t.Fatalf("Connect failed: %v", err)
|
|
}
|
|
defer conn.Close()
|
|
|
|
streamableConn, ok := conn.(*streamableClientConn)
|
|
if !ok {
|
|
t.Fatalf("Expected *streamableClientConn, got %T", conn)
|
|
}
|
|
|
|
if got := streamableConn.ctx.Value(testKey); got != testValue {
|
|
t.Errorf("Context value not propagated: got %v, want %v", got, testValue)
|
|
}
|
|
|
|
if streamableConn.ctx.Done() == nil {
|
|
t.Error("Connection context is not cancellable")
|
|
}
|
|
|
|
cancel()
|
|
select {
|
|
case <-streamableConn.ctx.Done():
|
|
t.Errorf("cancelling the connection context after successful connection broke the connection")
|
|
default:
|
|
}
|
|
}
|
|
|
|
func TestStreamableSessionTimeout(t *testing.T) {
|
|
// TODO: this test relies on timing and may be flaky.
|
|
// Fixing with testing/synctest is challenging because it uses real I/O (via
|
|
// httptest.NewServer).
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
|
|
server := NewServer(testImpl, nil)
|
|
|
|
deleted := make(chan string, 1)
|
|
handler := NewStreamableHTTPHandler(
|
|
func(req *http.Request) *Server { return server },
|
|
&StreamableHTTPOptions{
|
|
SessionTimeout: 50 * time.Millisecond,
|
|
},
|
|
)
|
|
handler.onTransportDeletion = func(sessionID string) {
|
|
deleted <- sessionID
|
|
}
|
|
|
|
httpServer := httptest.NewServer(mustNotPanic(t, handler))
|
|
defer httpServer.Close()
|
|
|
|
// Connect a client to create a session.
|
|
// Pin to 2025-11-25: under 2026-07-28 the client probes server/discover
|
|
// first against a stateful HTTP server (which rejects the new
|
|
// protocol), opening an extra HTTP connection and therefore an extra
|
|
// session. Pinning to 2025-11-25 skips the discover probe.
|
|
client := NewClient(testImpl, nil)
|
|
session, err := client.Connect(ctx, &StreamableClientTransport{Endpoint: httpServer.URL}, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125})
|
|
if err != nil {
|
|
t.Fatalf("client.Connect() failed: %v", err)
|
|
}
|
|
defer session.Close()
|
|
|
|
sessionID := session.ID()
|
|
if sessionID == "" {
|
|
t.Fatal("client session has empty ID")
|
|
}
|
|
|
|
// Verify the session exists on the server.
|
|
serverSessions := slices.Collect(server.Sessions())
|
|
if len(serverSessions) != 1 {
|
|
t.Fatalf("got %d sessions, want 1", len(serverSessions))
|
|
}
|
|
if got := serverSessions[0].ID(); got != sessionID {
|
|
t.Fatalf("server session is %q, want %q", got, sessionID)
|
|
}
|
|
|
|
// Test that (possibly concurrent) requests keep the session alive.
|
|
//
|
|
// Spin up two goroutines, each making a request every 10ms. These requests
|
|
// should keep the server from timing out.
|
|
var wg sync.WaitGroup
|
|
for range 2 {
|
|
wg.Go(func() {
|
|
for range 20 {
|
|
if _, err := session.ListTools(ctx, nil); err != nil {
|
|
t.Errorf("ListTools failed: %v", err)
|
|
}
|
|
time.Sleep(10 * time.Millisecond)
|
|
}
|
|
})
|
|
}
|
|
|
|
wg.Wait()
|
|
|
|
// Wait for the session to be cleaned up.
|
|
select {
|
|
case deletedID := <-deleted:
|
|
if deletedID != sessionID {
|
|
t.Errorf("deleted session ID = %q, want %q", deletedID, sessionID)
|
|
}
|
|
case <-ctx.Done():
|
|
t.Fatal("timed out waiting for session cleanup")
|
|
}
|
|
|
|
// Verify the session is gone from both handler and server.
|
|
handler.mu.Lock()
|
|
if len(handler.sessions) != 0 {
|
|
t.Errorf("handler.sessions is not empty; length %d", len(handler.sessions))
|
|
}
|
|
if ss := slices.Collect(server.Sessions()); len(ss) != 0 {
|
|
t.Errorf("server.Sessions() is not empty; length %d", len(ss))
|
|
}
|
|
handler.mu.Unlock()
|
|
}
|
|
|
|
// mustNotPanic is a helper to enforce that test handlers do not panic (see
|
|
// issue #556).
|
|
func mustNotPanic(t *testing.T, h http.Handler) http.Handler {
|
|
return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
|
|
defer func() {
|
|
if r := recover(); r != nil {
|
|
buf := make([]byte, 1<<20)
|
|
n := runtime.Stack(buf, false)
|
|
fmt.Fprintf(os.Stderr, "handler panic: %v\n\n%s", r, buf[:n])
|
|
t.Errorf("handler panicked: %v", r)
|
|
}
|
|
}()
|
|
h.ServeHTTP(w, req)
|
|
})
|
|
}
|
|
|
|
// TestPingEventFiltering verifies that the streamable client correctly filters
|
|
// out SSE "ping" events, which are used for keep-alive but should not be
|
|
// treated as JSON-RPC messages.
|
|
//
|
|
// This test addresses issue #636: the client should skip non-"message" events
|
|
// according to the SSE specification. It tests the fix in processStream where
|
|
// events with evt.Name != "" && evt.Name != "message" are skipped.
|
|
func TestPingEventFiltering(t *testing.T) {
|
|
// This test verifies the low-level processStream filtering.
|
|
// We create a mock response with ping and message events.
|
|
|
|
sseData := `event: ping
|
|
data: ping
|
|
|
|
event: message
|
|
id: 1
|
|
data: {"jsonrpc":"2.0","id":1,"result":{}}
|
|
|
|
event: ping
|
|
data: keepalive
|
|
|
|
`
|
|
|
|
resp := &http.Response{
|
|
StatusCode: http.StatusOK,
|
|
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
|
Body: io.NopCloser(bytes.NewReader([]byte(sseData))),
|
|
}
|
|
|
|
// Create a minimal streamableClientConn for testing
|
|
ctx := t.Context()
|
|
|
|
incoming := make(chan jsonrpc.Message, 10)
|
|
done := make(chan struct{})
|
|
|
|
conn := &streamableClientConn{
|
|
ctx: ctx,
|
|
done: done,
|
|
incoming: incoming,
|
|
}
|
|
|
|
// Create a test request
|
|
testReq := &jsonrpc.Request{
|
|
ID: jsonrpc2.Int64ID(1),
|
|
Method: "test",
|
|
}
|
|
|
|
// Process the stream
|
|
go conn.processStream(ctx, "test", resp, testReq)
|
|
|
|
// Collect messages with timeout
|
|
var messages []jsonrpc.Message
|
|
timeout := time.After(1 * time.Second)
|
|
|
|
collectLoop:
|
|
for {
|
|
select {
|
|
case msg := <-incoming:
|
|
messages = append(messages, msg)
|
|
// We expect only 1 message (the response), not the ping events
|
|
if len(messages) >= 1 {
|
|
break collectLoop
|
|
}
|
|
case <-timeout:
|
|
break collectLoop
|
|
}
|
|
}
|
|
|
|
// Verify we only received the actual message, not the ping events
|
|
if len(messages) != 1 {
|
|
t.Errorf("got %d messages, want 1 (ping events should be filtered)", len(messages))
|
|
for i, msg := range messages {
|
|
t.Logf("message %d: %T", i, msg)
|
|
}
|
|
}
|
|
|
|
// Verify the message is the response
|
|
if len(messages) > 0 {
|
|
resp, ok := messages[0].(*jsonrpc.Response)
|
|
if !ok {
|
|
t.Errorf("first message is %T, want *jsonrpc.Response", messages[0])
|
|
} else if resp.ID.Raw() != int64(1) {
|
|
t.Errorf("response ID is %v, want 1", resp.ID.Raw())
|
|
}
|
|
}
|
|
}
|
|
|
|
// TestProcessStreamPrimingEvent verifies that the streamable client correctly ignores
|
|
// SSE events with empty data buffers, which are used as priming events (e.g. SEP-1699).
|
|
func TestProcessStreamPrimingEvent(t *testing.T) {
|
|
// We create a mock response with a priming event (empty data, with an ID),
|
|
// followed by a normal event.
|
|
sseData := `id: 123
|
|
|
|
id: 124
|
|
data: {"jsonrpc":"2.0","id":1,"result":{}}
|
|
|
|
`
|
|
|
|
ctx := t.Context()
|
|
resp := &http.Response{
|
|
StatusCode: http.StatusOK,
|
|
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
|
Body: io.NopCloser(strings.NewReader(sseData)),
|
|
}
|
|
|
|
incoming := make(chan jsonrpc.Message, 10)
|
|
done := make(chan struct{})
|
|
|
|
conn := &streamableClientConn{
|
|
ctx: ctx,
|
|
done: done,
|
|
incoming: incoming,
|
|
failed: make(chan struct{}),
|
|
logger: ensureLogger(nil),
|
|
}
|
|
|
|
lastID, _, clientClosed := conn.processStream(ctx, "test", resp, nil)
|
|
|
|
if clientClosed {
|
|
t.Fatalf("processStream was unexpectedly closed by client")
|
|
}
|
|
|
|
if lastID != "124" {
|
|
t.Errorf("lastEventID = %q, want %q", lastID, "124")
|
|
}
|
|
|
|
select {
|
|
case msg := <-incoming:
|
|
if res, ok := msg.(*jsonrpc.Response); !(ok && res.ID == jsonrpc2.Int64ID(1)) {
|
|
t.Errorf("got unexpected message: %v", msg)
|
|
}
|
|
default:
|
|
t.Errorf("expected a JSON-RPC message to be produced")
|
|
}
|
|
}
|
|
|
|
// TestProcessStreamMaxEventSize verifies that a streamableClientConn honors its
|
|
// configured maxEventSize.
|
|
func TestProcessStreamMaxEventSize(t *testing.T) {
|
|
jsonrpcEventOfSize := func(padSize int) string {
|
|
msg := fmt.Sprintf(`{"jsonrpc":"2.0","id":1,"result":{"pad":%q}}`, strings.Repeat("A", padSize))
|
|
return "data: " + msg + "\n\n"
|
|
}
|
|
|
|
tests := []struct {
|
|
name string
|
|
maxEventSize int
|
|
event string
|
|
wantFail bool
|
|
}{
|
|
{
|
|
name: "event over cap is rejected",
|
|
event: jsonrpcEventOfSize(4096),
|
|
maxEventSize: 1024,
|
|
wantFail: true,
|
|
},
|
|
{
|
|
name: "event under cap is accepted",
|
|
event: jsonrpcEventOfSize(1024),
|
|
maxEventSize: 4096,
|
|
wantFail: false,
|
|
},
|
|
{
|
|
name: "negative cap disables the limit",
|
|
event: jsonrpcEventOfSize(4096),
|
|
maxEventSize: -1,
|
|
wantFail: false,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
ctx := t.Context()
|
|
resp := &http.Response{
|
|
StatusCode: http.StatusOK,
|
|
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
|
Body: io.NopCloser(strings.NewReader(tt.event)),
|
|
}
|
|
conn := &streamableClientConn{
|
|
ctx: ctx,
|
|
done: make(chan struct{}),
|
|
incoming: make(chan jsonrpc.Message, 10),
|
|
failed: make(chan struct{}),
|
|
logger: ensureLogger(nil),
|
|
maxEventSize: tt.maxEventSize,
|
|
}
|
|
conn.processStream(ctx, "test", resp, nil)
|
|
|
|
err := conn.failure()
|
|
if tt.wantFail && err == nil {
|
|
t.Fatal("failure() = nil, want a non-nil error for an oversized event")
|
|
}
|
|
if !tt.wantFail && err != nil {
|
|
t.Fatalf("failure() = %v, want nil", err)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestScanEventsPingFiltering is a unit test for the low-level event scanning
|
|
// with ping events to verify scanEvents properly parses all event types.
|
|
func TestScanEventsPingFiltering(t *testing.T) {
|
|
// Create SSE stream with mixed events
|
|
sseData := `event: ping
|
|
data: ping
|
|
|
|
event: message
|
|
data: {"jsonrpc":"2.0","method":"test","params":{}}
|
|
|
|
event: ping
|
|
data: keepalive
|
|
|
|
event: message
|
|
data: {"jsonrpc":"2.0","method":"test2","params":{}}
|
|
|
|
`
|
|
|
|
reader := strings.NewReader(sseData)
|
|
var events []Event
|
|
|
|
// Scan all events
|
|
for evt, err := range scanEventsLimited(reader, DefaultMaxEventSize) {
|
|
if err != nil {
|
|
if err != io.EOF {
|
|
t.Fatalf("scanEvents error: %v", err)
|
|
}
|
|
break
|
|
}
|
|
events = append(events, evt)
|
|
}
|
|
|
|
// Verify we got all 4 events
|
|
if len(events) != 4 {
|
|
t.Fatalf("got %d events, want 4", len(events))
|
|
}
|
|
|
|
// Verify event types
|
|
expectedNames := []string{"ping", "message", "ping", "message"}
|
|
for i, evt := range events {
|
|
if evt.Name != expectedNames[i] {
|
|
t.Errorf("event %d: got name %q, want %q", i, evt.Name, expectedNames[i])
|
|
}
|
|
}
|
|
|
|
// Verify that we can decode the message events but would fail on ping events
|
|
for i, evt := range events {
|
|
switch evt.Name {
|
|
case "message":
|
|
_, err := jsonrpc.DecodeMessage(evt.Data)
|
|
if err != nil {
|
|
t.Errorf("event %d: failed to decode message event: %v", i, err)
|
|
}
|
|
case "ping":
|
|
// Ping events have non-JSON data and should fail decoding
|
|
_, err := jsonrpc.DecodeMessage(evt.Data)
|
|
if err == nil {
|
|
t.Errorf("event %d: ping event unexpectedly decoded as valid JSON-RPC", i)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func Test_ExportErrSessionMissing(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
// 1. Setup server
|
|
impl := &Implementation{Name: "test", Version: "1.0.0"}
|
|
server := NewServer(impl, nil)
|
|
handler := NewStreamableHTTPHandler(func(r *http.Request) *Server { return server }, nil)
|
|
ts := httptest.NewServer(handler)
|
|
defer ts.Close()
|
|
|
|
// 2. Setup client
|
|
clientTransport := &StreamableClientTransport{
|
|
Endpoint: ts.URL,
|
|
}
|
|
client := NewClient(impl, nil)
|
|
// Pin to 2025-11-25: under 2026-07-28 the client probes server/discover
|
|
// first against a stateful HTTP server (which rejects the new
|
|
// protocol), opening an extra HTTP connection and therefore an extra
|
|
// session. Pinning to 2025-11-25 skips the discover probe.
|
|
session, err := client.Connect(ctx, clientTransport, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125})
|
|
if err != nil {
|
|
t.Fatalf("Connect failed: %v", err)
|
|
}
|
|
defer session.Close()
|
|
|
|
// 3. Manually invalidate session on server
|
|
handler.mu.Lock()
|
|
if len(handler.sessions) != 1 {
|
|
handler.mu.Unlock()
|
|
t.Fatalf("expected 1 session, got %d", len(handler.sessions))
|
|
}
|
|
for id := range handler.sessions {
|
|
delete(handler.sessions, id)
|
|
}
|
|
handler.mu.Unlock()
|
|
|
|
// 4. Try to call a tool (or any request)
|
|
_, err = session.ListTools(ctx, nil)
|
|
if err == nil {
|
|
t.Fatal("expected error, got nil")
|
|
}
|
|
|
|
// 5. Verify it's ErrSessionMissing
|
|
if !errors.Is(err, ErrSessionMissing) {
|
|
t.Errorf("expected error to wrap ErrSessionMissing, got: %v", err)
|
|
}
|
|
}
|
|
|
|
// TestStreamableLocalhostProtection verifies that DNS rebinding protection
|
|
// is automatically enabled for localhost servers.
|
|
func TestStreamableLocalhostProtection(t *testing.T) {
|
|
server := NewServer(testImpl, nil)
|
|
|
|
tests := []struct {
|
|
name string
|
|
listenAddr string // Address to listen on
|
|
hostHeader string // Host header in request
|
|
disableProtection bool // DisableLocalhostProtection setting
|
|
wantStatus int
|
|
}{
|
|
// Auto-enabled for localhost listeners (127.0.0.1).
|
|
{
|
|
name: "127.0.0.1 accepts 127.0.0.1",
|
|
listenAddr: "127.0.0.1:0",
|
|
hostHeader: "127.0.0.1:1234",
|
|
disableProtection: false,
|
|
wantStatus: http.StatusOK,
|
|
},
|
|
{
|
|
name: "127.0.0.1 accepts localhost",
|
|
listenAddr: "127.0.0.1:0",
|
|
hostHeader: "localhost:1234",
|
|
disableProtection: false,
|
|
wantStatus: http.StatusOK,
|
|
},
|
|
{
|
|
name: "127.0.0.1 rejects evil.com",
|
|
listenAddr: "127.0.0.1:0",
|
|
hostHeader: "evil.com",
|
|
disableProtection: false,
|
|
wantStatus: http.StatusForbidden,
|
|
},
|
|
{
|
|
name: "127.0.0.1 rejects evil.com:80",
|
|
listenAddr: "127.0.0.1:0",
|
|
hostHeader: "evil.com:80",
|
|
disableProtection: false,
|
|
wantStatus: http.StatusForbidden,
|
|
},
|
|
{
|
|
name: "127.0.0.1 rejects localhost.evil.com",
|
|
listenAddr: "127.0.0.1:0",
|
|
hostHeader: "localhost.evil.com",
|
|
disableProtection: false,
|
|
wantStatus: http.StatusForbidden,
|
|
},
|
|
|
|
// When listening on 0.0.0.0, requests arriving via localhost are still protected
|
|
// because LocalAddrContextKey returns the actual connection's local address.
|
|
// This is actually more secure - DNS rebinding attacks target localhost regardless
|
|
// of the listener configuration.
|
|
{
|
|
name: "0.0.0.0 via localhost rejects evil.com",
|
|
listenAddr: "0.0.0.0:0",
|
|
hostHeader: "evil.com",
|
|
disableProtection: false,
|
|
wantStatus: http.StatusForbidden,
|
|
},
|
|
|
|
// Explicit disable
|
|
{
|
|
name: "disabled accepts evil.com",
|
|
listenAddr: "127.0.0.1:0",
|
|
hostHeader: "evil.com",
|
|
disableProtection: true,
|
|
wantStatus: http.StatusOK,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
opts := &StreamableHTTPOptions{
|
|
Stateless: true, // Simpler for testing
|
|
DisableLocalhostProtection: tt.disableProtection,
|
|
}
|
|
handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, opts)
|
|
|
|
listener, err := net.Listen("tcp", tt.listenAddr)
|
|
if err != nil {
|
|
t.Fatalf("Failed to listen on %s: %v", tt.listenAddr, err)
|
|
}
|
|
defer listener.Close()
|
|
|
|
srv := &http.Server{Handler: handler}
|
|
go srv.Serve(listener)
|
|
defer srv.Close()
|
|
|
|
reqReader := strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2024-11-05","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}`)
|
|
req, err := http.NewRequest("POST", fmt.Sprintf("http://%s", listener.Addr().String()), reqReader)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
req.Host = tt.hostHeader
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("Accept", "application/json, text/event-stream")
|
|
|
|
resp, err := http.DefaultClient.Do(req)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if got := resp.StatusCode; got != tt.wantStatus {
|
|
t.Errorf("Status code: got %d, want %d", got, tt.wantStatus)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestStreamableOriginProtection(t *testing.T) {
|
|
server := NewServer(testImpl, nil)
|
|
|
|
tests := []struct {
|
|
name string
|
|
protection *http.CrossOriginProtection
|
|
requestOrigin string
|
|
wantStatusCode int
|
|
}{
|
|
{
|
|
name: "no protection with Origin header",
|
|
protection: nil,
|
|
requestOrigin: "https://example.com",
|
|
wantStatusCode: http.StatusOK,
|
|
},
|
|
{
|
|
name: "custom protection with trusted origin and same Origin",
|
|
protection: func() *http.CrossOriginProtection {
|
|
p := http.NewCrossOriginProtection()
|
|
if err := p.AddTrustedOrigin("https://example.com"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
return p
|
|
}(),
|
|
requestOrigin: "https://example.com",
|
|
wantStatusCode: http.StatusOK,
|
|
},
|
|
{
|
|
name: "custom protection with trusted origin and different Origin",
|
|
protection: func() *http.CrossOriginProtection {
|
|
p := http.NewCrossOriginProtection()
|
|
if err := p.AddTrustedOrigin("https://example.com"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
return p
|
|
}(),
|
|
requestOrigin: "https://malicious.com",
|
|
wantStatusCode: http.StatusForbidden,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
opts := &StreamableHTTPOptions{
|
|
Stateless: true, // avoid session ID requirement
|
|
CrossOriginProtection: tt.protection,
|
|
}
|
|
handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, opts)
|
|
httpServer := httptest.NewServer(handler)
|
|
defer httpServer.Close()
|
|
|
|
reqReader := strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}`)
|
|
req, err := http.NewRequest(http.MethodPost, httpServer.URL, reqReader)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("Origin", tt.requestOrigin)
|
|
req.Header.Set("Accept", "application/json, text/event-stream")
|
|
|
|
resp, err := http.DefaultClient.Do(req)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if got := resp.StatusCode; got != tt.wantStatusCode {
|
|
t.Errorf("Status code: got %d, want %d", got, tt.wantStatusCode)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestStandaloneSSEEmitsCommentForHTTP2Flush is a regression test for the
|
|
// HTTP/2 reverse-proxy header-buffering bug. The standalone SSE GET stream
|
|
// must emit at least one body byte (an SSE comment) immediately after the
|
|
// response headers. Without a DATA frame, HTTP/2 reverse proxies (Envoy,
|
|
// Caddy, net/http/httputil, etc.) buffer the HEADERS frame indefinitely
|
|
// because they have nothing to coalesce it with — calling Flush() alone is
|
|
// not sufficient.
|
|
//
|
|
// SSE explicitly defines lines starting with ":" as comments that clients
|
|
// ignore, so this is behavior-preserving for the SSE protocol while
|
|
// producing the DATA frame that HTTP/2 proxies need.
|
|
func TestStandaloneSSEEmitsCommentForHTTP2Flush(t *testing.T) {
|
|
server := NewServer(testImpl, nil)
|
|
handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, nil)
|
|
httpServer := httptest.NewServer(mustNotPanic(t, handler))
|
|
defer httpServer.Close()
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
|
|
// Initialize a session so we have a valid Mcp-Session-Id to use on the
|
|
// standalone SSE GET.
|
|
initBody := strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}`)
|
|
initReq, err := http.NewRequestWithContext(ctx, http.MethodPost, httpServer.URL, initBody)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
initReq.Header.Set("Content-Type", "application/json")
|
|
initReq.Header.Set("Accept", "application/json, text/event-stream")
|
|
initResp, err := http.DefaultClient.Do(initReq)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
io.Copy(io.Discard, initResp.Body)
|
|
initResp.Body.Close()
|
|
sessionID := initResp.Header.Get(sessionIDHeader)
|
|
if sessionID == "" {
|
|
t.Fatalf("initialize response missing %s header", sessionIDHeader)
|
|
}
|
|
|
|
// Open the standalone SSE stream.
|
|
getReq, err := http.NewRequestWithContext(ctx, http.MethodGet, httpServer.URL, nil)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
getReq.Header.Set("Accept", "text/event-stream")
|
|
getReq.Header.Set(sessionIDHeader, sessionID)
|
|
|
|
resp, err := http.DefaultClient.Do(getReq)
|
|
if err != nil {
|
|
t.Fatalf("GET standalone SSE: %v", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
t.Fatalf("status = %d; want 200", resp.StatusCode)
|
|
}
|
|
if got := resp.Header.Get("Content-Type"); !strings.HasPrefix(got, "text/event-stream") {
|
|
t.Fatalf("Content-Type = %q; want text/event-stream", got)
|
|
}
|
|
|
|
// Read the first chunk. With the fix, an SSE comment (": ok\n\n") is
|
|
// written immediately. Without the fix, the body has no bytes to read
|
|
// until the server pushes an event, which won't happen in this test.
|
|
type readResult struct {
|
|
n int
|
|
err error
|
|
buf []byte
|
|
}
|
|
ch := make(chan readResult, 1)
|
|
go func() {
|
|
buf := make([]byte, 64)
|
|
n, err := resp.Body.Read(buf)
|
|
ch <- readResult{n: n, err: err, buf: buf}
|
|
}()
|
|
|
|
select {
|
|
case r := <-ch:
|
|
if r.err != nil && r.err != io.EOF {
|
|
t.Fatalf("reading first SSE chunk: %v", r.err)
|
|
}
|
|
if r.n == 0 {
|
|
t.Fatal("first SSE chunk was empty; expected an SSE comment to flush HTTP/2 proxy HEADERS frame")
|
|
}
|
|
got := string(r.buf[:r.n])
|
|
// SSE spec: lines starting with ":" are comments, ignored by clients.
|
|
// We don't pin the exact comment text; just require the first byte to
|
|
// be the SSE comment marker so HTTP/2 proxies have a DATA frame.
|
|
if !strings.HasPrefix(got, ":") {
|
|
t.Errorf("first SSE chunk = %q; want it to start with %q (SSE comment marker, so HTTP/2 proxies receive a DATA frame and forward HEADERS)", got, ":")
|
|
}
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("timed out waiting for first SSE bytes; the standalone SSE stream must emit a DATA frame immediately so HTTP/2 reverse proxies don't buffer the HEADERS frame")
|
|
}
|
|
}
|
|
|
|
// newProtocolBody builds a raw JSON body for a tools/call request that
|
|
// carries the >= 2026-07-28 per-request _meta fields.
|
|
func newProtocolBody(t *testing.T, toolName string, args any) []byte {
|
|
t.Helper()
|
|
rawArgs, err := json.Marshal(args)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
body, err := json.Marshal(map[string]any{
|
|
"jsonrpc": "2.0",
|
|
"id": 1,
|
|
"method": "tools/call",
|
|
"params": map[string]any{
|
|
"_meta": map[string]any{
|
|
MetaKeyProtocolVersion: protocolVersion20260728,
|
|
MetaKeyClientInfo: map[string]any{"name": "new-proto-client", "version": "9.9"},
|
|
MetaKeyClientCapabilities: map[string]any{"sampling": map[string]any{}},
|
|
},
|
|
"name": toolName,
|
|
"arguments": json.RawMessage(rawArgs),
|
|
},
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
return body
|
|
}
|
|
|
|
func TestEphemeralConnectOpts(t *testing.T) {
|
|
mkReq := func(body []byte) *http.Request {
|
|
r := httptest.NewRequest(http.MethodPost, "/", bytes.NewReader(body))
|
|
r.Header.Set("Content-Type", "application/json")
|
|
return r
|
|
}
|
|
|
|
h := &StreamableHTTPHandler{opts: StreamableHTTPOptions{}}
|
|
|
|
oldProtocolBody, err := json.Marshal(map[string]any{
|
|
"jsonrpc": "2.0",
|
|
"id": 1,
|
|
"method": "tools/call",
|
|
"params": map[string]any{"name": "x", "arguments": map[string]any{}},
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
initializeBody, err := json.Marshal(map[string]any{
|
|
"jsonrpc": "2.0",
|
|
"id": 1,
|
|
"method": methodInitialize,
|
|
"params": map[string]any{"protocolVersion": protocolVersion20250618},
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
tests := []struct {
|
|
name string
|
|
body []byte
|
|
wantUsesNew bool
|
|
wantInitializeParams bool
|
|
wantInitializedParams bool
|
|
}{
|
|
{
|
|
name: "new-protocol request: no synthetic state",
|
|
body: newProtocolBody(t, "x", struct{}{}),
|
|
wantUsesNew: true,
|
|
wantInitializeParams: false,
|
|
wantInitializedParams: false,
|
|
},
|
|
{
|
|
name: "old-protocol request: synthetic state populated",
|
|
body: oldProtocolBody,
|
|
wantUsesNew: false,
|
|
wantInitializeParams: true,
|
|
wantInitializedParams: true,
|
|
},
|
|
{
|
|
name: "initialize request: no synthetic InitializeParams",
|
|
body: initializeBody,
|
|
wantUsesNew: false,
|
|
wantInitializeParams: false,
|
|
wantInitializedParams: true,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
req := mkReq(tt.body)
|
|
var pver string
|
|
if tt.wantUsesNew {
|
|
pver = protocolVersion20260728
|
|
} else {
|
|
pver = protocolVersion20250326
|
|
}
|
|
req.Header.Set(protocolVersionHeader, pver)
|
|
req = req.WithContext(context.WithValue(req.Context(), protocolVersionContextKey{}, pver))
|
|
info, err := h.ephemeralConnectOpts(req)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if info.usesNewProtocol != tt.wantUsesNew {
|
|
t.Errorf("usesNewProtocol = %v, want %v", info.usesNewProtocol, tt.wantUsesNew)
|
|
}
|
|
if got := info.opts.State.InitializeParams != nil; got != tt.wantInitializeParams {
|
|
t.Errorf("InitializeParams non-nil = %v, want %v (value = %+v)",
|
|
got, tt.wantInitializeParams, info.opts.State.InitializeParams)
|
|
}
|
|
if got := info.opts.State.InitializedParams != nil; got != tt.wantInitializedParams {
|
|
t.Errorf("InitializedParams non-nil = %v, want %v (value = %+v)",
|
|
got, tt.wantInitializedParams, info.opts.State.InitializedParams)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// statelessHandlerCapture builds a stateless server with a single tool whose
|
|
// handler captures everything we want to assert about the per-request view of
|
|
// the session and the new-protocol accessors.
|
|
type statelessHandlerCapture struct {
|
|
mu sync.Mutex
|
|
sessionInitParams *InitializeParams
|
|
reqProtocolVersion string
|
|
reqClientInfo *Implementation
|
|
reqClientCapabilities *ClientCapabilities
|
|
}
|
|
|
|
func TestStreamableStateless_NewProtocolSession_NoFakeInit(t *testing.T) {
|
|
// SEP-2575: the MCP-Protocol-Version header is mandatory for new-protocol
|
|
// requests and must be a supported version. The 2026-07-28 version is
|
|
// not yet in the global list, so register it for the duration of the test.
|
|
|
|
capture := &statelessHandlerCapture{}
|
|
mcpServer := NewServer(testImpl, nil)
|
|
AddTool(mcpServer, &Tool{Name: "capture", Description: "captures request info"},
|
|
func(ctx context.Context, req *CallToolRequest, args struct{}) (*CallToolResult, any, error) {
|
|
capture.mu.Lock()
|
|
defer capture.mu.Unlock()
|
|
capture.sessionInitParams = req.Session.InitializeParams()
|
|
capture.reqProtocolVersion = req.ProtocolVersion()
|
|
capture.reqClientInfo = req.ClientInfo()
|
|
capture.reqClientCapabilities = req.ClientCapabilities()
|
|
return &CallToolResult{Content: []Content{&TextContent{Text: "ok"}}}, nil, nil
|
|
})
|
|
|
|
handler := NewStreamableHTTPHandler(
|
|
func(*http.Request) *Server { return mcpServer },
|
|
&StreamableHTTPOptions{Stateless: true},
|
|
)
|
|
httpServer := httptest.NewServer(handler)
|
|
defer httpServer.Close()
|
|
|
|
body := newProtocolBody(t, "capture", struct{}{})
|
|
httpReq, err := http.NewRequest(http.MethodPost, httpServer.URL, bytes.NewReader(body))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
httpReq.Header.Set("Content-Type", "application/json")
|
|
httpReq.Header.Set("Accept", "application/json, text/event-stream")
|
|
httpReq.Header.Set(protocolVersionHeader, protocolVersion20260728)
|
|
// >= 2026-07-28 also requires the Mcp-Method and Mcp-Name standard
|
|
// headers (see streamable_headers.go).
|
|
httpReq.Header.Set(methodHeader, "tools/call")
|
|
httpReq.Header.Set(nameHeader, "capture")
|
|
|
|
resp, err := http.DefaultClient.Do(httpReq)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer resp.Body.Close()
|
|
if resp.StatusCode != http.StatusOK {
|
|
respBody, _ := io.ReadAll(resp.Body)
|
|
t.Fatalf("status = %d, want 200; body = %s", resp.StatusCode, respBody)
|
|
}
|
|
|
|
capture.mu.Lock()
|
|
defer capture.mu.Unlock()
|
|
if capture.sessionInitParams == nil {
|
|
t.Errorf("Session.InitializeParams() is nil, want populated initializeParams for new-protocol session")
|
|
} else {
|
|
if got, want := capture.sessionInitParams.ProtocolVersion, protocolVersion20260728; got != want {
|
|
t.Errorf("Session.InitializeParams().ProtocolVersion = %q, want %q", got, want)
|
|
}
|
|
if got, want := capture.sessionInitParams.ClientInfo.Name, "new-proto-client"; got != want {
|
|
t.Errorf("Session.InitializeParams().ClientInfo.Name = %q, want %q", got, want)
|
|
}
|
|
}
|
|
if got, want := capture.reqProtocolVersion, protocolVersion20260728; got != want {
|
|
t.Errorf("req.ProtocolVersion() = %q, want %q", got, want)
|
|
}
|
|
if capture.reqClientInfo == nil || capture.reqClientInfo.Name != "new-proto-client" {
|
|
t.Errorf("req.ClientInfo() = %+v, want Name=new-proto-client", capture.reqClientInfo)
|
|
}
|
|
if capture.reqClientCapabilities == nil || capture.reqClientCapabilities.Sampling == nil {
|
|
t.Errorf("req.ClientCapabilities() = %+v, want non-nil Sampling", capture.reqClientCapabilities)
|
|
}
|
|
}
|
|
|
|
// TestStreamableStateful_RejectsNewProtocol verifies that a stateful HTTP
|
|
// server rejects requests carrying _meta.protocolVersion (i.e. >= 2026-07-28
|
|
// requests) with HTTP 400. The new protocol is
|
|
// supported on HTTP only when StreamableHTTPOptions.Stateless=true.
|
|
func TestStreamableStateful_RejectsNewProtocol(t *testing.T) {
|
|
// Make 2026-07-28 a "known" version so that the request reaches servePOST
|
|
// (otherwise the early header validation at ServeHTTP rejects it).
|
|
|
|
server := NewServer(testImpl, nil)
|
|
AddTool(server, &Tool{Name: "noop"},
|
|
func(ctx context.Context, req *CallToolRequest, args struct{}) (*CallToolResult, any, error) {
|
|
return &CallToolResult{Content: []Content{&TextContent{Text: "ok"}}}, nil, nil
|
|
})
|
|
handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, nil)
|
|
httpServer := httptest.NewServer(handler)
|
|
defer httpServer.Close()
|
|
|
|
// Initialize a legacy session first.
|
|
initBody := strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}`)
|
|
initReq, err := http.NewRequest(http.MethodPost, httpServer.URL, initBody)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
initReq.Header.Set("Content-Type", "application/json")
|
|
initReq.Header.Set("Accept", "application/json, text/event-stream")
|
|
initResp, err := http.DefaultClient.Do(initReq)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
io.Copy(io.Discard, initResp.Body)
|
|
initResp.Body.Close()
|
|
sessionID := initResp.Header.Get(sessionIDHeader)
|
|
if sessionID == "" {
|
|
t.Fatalf("initialize response missing %s header", sessionIDHeader)
|
|
}
|
|
|
|
// Drive the existing session with a new-protocol request whose header and
|
|
// body agree. The cross-check passes; the stateful-rejection check fires.
|
|
body := newProtocolBody(t, "noop", struct{}{})
|
|
req, err := http.NewRequest(http.MethodPost, httpServer.URL, bytes.NewReader(body))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("Accept", "application/json, text/event-stream")
|
|
req.Header.Set(sessionIDHeader, sessionID)
|
|
req.Header.Set(protocolVersionHeader, protocolVersion20260728)
|
|
req.Header.Set(methodHeader, "tools/call")
|
|
req.Header.Set(nameHeader, "noop")
|
|
|
|
resp, err := http.DefaultClient.Do(req)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer resp.Body.Close()
|
|
respBody, _ := io.ReadAll(resp.Body)
|
|
if resp.StatusCode != http.StatusBadRequest {
|
|
t.Fatalf("status = %d, want 400; body = %s", resp.StatusCode, respBody)
|
|
}
|
|
msg, err := jsonrpc.DecodeMessage(respBody)
|
|
if err != nil {
|
|
t.Fatalf("decoding response body as JSON-RPC: %v; body = %s", err, respBody)
|
|
}
|
|
jresp, ok := msg.(*jsonrpc.Response)
|
|
if !ok || jresp.Error == nil {
|
|
t.Fatalf("response is not a JSON-RPC error: %s", respBody)
|
|
}
|
|
var jerr *jsonrpc.Error
|
|
if !errors.As(jresp.Error, &jerr) {
|
|
t.Fatalf("response error is not *jsonrpc.Error: %v", jresp.Error)
|
|
}
|
|
if jerr.Code != CodeUnsupportedProtocolVersion {
|
|
t.Errorf("error code = %d, want %d (CodeUnsupportedProtocolVersion)", jerr.Code, CodeUnsupportedProtocolVersion)
|
|
}
|
|
var data UnsupportedProtocolVersionData
|
|
if err := json.Unmarshal(jerr.Data, &data); err != nil {
|
|
t.Fatalf("decoding error data: %v; data = %s", err, jerr.Data)
|
|
}
|
|
if data.Requested != protocolVersion20260728 {
|
|
t.Errorf("data.Requested = %q, want %q", data.Requested, protocolVersion20260728)
|
|
}
|
|
if slices.Contains(data.Supported, protocolVersion20260728) {
|
|
t.Errorf("data.Supported = %v, must not contain the rejected version %q",
|
|
data.Supported, protocolVersion20260728)
|
|
}
|
|
if len(data.Supported) == 0 {
|
|
t.Errorf("data.Supported is empty; expected at least one legacy version")
|
|
}
|
|
}
|
|
|
|
// TestStreamableStateful_RejectsNewProtocol_LegacyPlainText verifies that
|
|
// MCPGODEBUG=plaintextstatefulrejection=1 restores the pre-v1.8.0 plain-text
|
|
// 400 body for the stateful/new-protocol rejection.
|
|
func TestStreamableStateful_RejectsNewProtocol_LegacyPlainText(t *testing.T) {
|
|
prev := plaintextstatefulrejection
|
|
plaintextstatefulrejection = "1"
|
|
t.Cleanup(func() { plaintextstatefulrejection = prev })
|
|
|
|
server := NewServer(testImpl, nil)
|
|
AddTool(server, &Tool{Name: "noop"},
|
|
func(ctx context.Context, req *CallToolRequest, args struct{}) (*CallToolResult, any, error) {
|
|
return &CallToolResult{Content: []Content{&TextContent{Text: "ok"}}}, nil, nil
|
|
})
|
|
handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, nil)
|
|
httpServer := httptest.NewServer(handler)
|
|
defer httpServer.Close()
|
|
|
|
initBody := strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}`)
|
|
initReq, err := http.NewRequest(http.MethodPost, httpServer.URL, initBody)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
initReq.Header.Set("Content-Type", "application/json")
|
|
initReq.Header.Set("Accept", "application/json, text/event-stream")
|
|
initResp, err := http.DefaultClient.Do(initReq)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
io.Copy(io.Discard, initResp.Body)
|
|
initResp.Body.Close()
|
|
sessionID := initResp.Header.Get(sessionIDHeader)
|
|
if sessionID == "" {
|
|
t.Fatalf("initialize response missing %s header", sessionIDHeader)
|
|
}
|
|
|
|
body := newProtocolBody(t, "noop", struct{}{})
|
|
req, err := http.NewRequest(http.MethodPost, httpServer.URL, bytes.NewReader(body))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("Accept", "application/json, text/event-stream")
|
|
req.Header.Set(sessionIDHeader, sessionID)
|
|
req.Header.Set(protocolVersionHeader, protocolVersion20260728)
|
|
req.Header.Set(methodHeader, "tools/call")
|
|
req.Header.Set(nameHeader, "noop")
|
|
|
|
resp, err := http.DefaultClient.Do(req)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer resp.Body.Close()
|
|
respBody, _ := io.ReadAll(resp.Body)
|
|
if resp.StatusCode != http.StatusBadRequest {
|
|
t.Fatalf("status = %d, want 400; body = %s", resp.StatusCode, respBody)
|
|
}
|
|
if ct := resp.Header.Get("Content-Type"); !strings.HasPrefix(ct, "text/plain") {
|
|
t.Errorf("Content-Type = %q, want text/plain; body = %s", ct, respBody)
|
|
}
|
|
if !strings.Contains(string(respBody), "stateless") {
|
|
t.Errorf("body = %q, want a message mentioning 'stateless'", respBody)
|
|
}
|
|
}
|
|
|
|
// TestCallCancellation_FastReturn verifies that a cancelled tool call
|
|
// returns as soon as its context is cancelled, even when the peer is slow
|
|
// to accept the follow-up notifications/cancelled POST.
|
|
func TestCallCancellation_FastReturn(t *testing.T) {
|
|
// stallDuration is the time the middleware holds a
|
|
// notifications/cancelled POST. It must be larger than
|
|
// notifyCancellationTimeout so the blocking branch clearly overshoots
|
|
// even after cushioning for scheduler jitter.
|
|
const stallDuration = notifyCancellationTimeout + 2*time.Second
|
|
|
|
// callerDeadline is the deadline the client passes into CallTool. It
|
|
// must be short enough that any waiting on the cancel notification is
|
|
// obvious in the measured elapsed time.
|
|
const callerDeadline = 100 * time.Millisecond
|
|
|
|
// slack accounts for scheduling jitter around the deadline.
|
|
const slack = 500 * time.Millisecond
|
|
|
|
stallCancelNotify := func(h http.Handler) http.Handler {
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.Method == http.MethodPost && r.Body != nil {
|
|
body, _ := io.ReadAll(r.Body)
|
|
r.Body = io.NopCloser(bytes.NewReader(body))
|
|
if bytes.Contains(body, []byte(notificationCancelled)) {
|
|
time.Sleep(stallDuration)
|
|
}
|
|
}
|
|
h.ServeHTTP(w, r)
|
|
})
|
|
}
|
|
|
|
server := NewServer(testImpl, nil)
|
|
AddTool(server, &Tool{Name: "slow"},
|
|
func(ctx context.Context, req *CallToolRequest, args struct{}) (*CallToolResult, any, error) {
|
|
// Bound the handler above the caller deadline so the test
|
|
// always exercises the cancellation path, but not so long
|
|
// that a stuck cancellation makes the test slow.
|
|
select {
|
|
case <-ctx.Done():
|
|
case <-time.After(2 * stallDuration):
|
|
}
|
|
return &CallToolResult{Content: []Content{&TextContent{Text: "ok"}}}, nil, nil
|
|
})
|
|
|
|
handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, nil)
|
|
httpServer := httptest.NewServer(stallCancelNotify(handler))
|
|
t.Cleanup(httpServer.Close)
|
|
|
|
client := NewClient(testImpl, nil)
|
|
cs, err := client.Connect(context.Background(),
|
|
&StreamableClientTransport{Endpoint: httpServer.URL}, nil)
|
|
if err != nil {
|
|
t.Fatalf("client connect: %v", err)
|
|
}
|
|
t.Cleanup(func() { cs.Close() })
|
|
|
|
callCtx, cancel := context.WithTimeout(context.Background(), callerDeadline)
|
|
defer cancel()
|
|
|
|
start := time.Now()
|
|
_, err = cs.CallTool(callCtx, &CallToolParams{Name: "slow"})
|
|
elapsed := time.Since(start)
|
|
|
|
if !errors.Is(err, context.DeadlineExceeded) {
|
|
t.Fatalf("CallTool error = %v, want context.DeadlineExceeded", err)
|
|
}
|
|
// The caller should return within slack of its own deadline.
|
|
if elapsed > callerDeadline+slack {
|
|
t.Errorf("CallTool returned after %v; want <= %v (caller deadline %v + slack %v)",
|
|
elapsed, callerDeadline+slack, callerDeadline, slack)
|
|
}
|
|
}
|
|
|
|
// TestStreamableStateless_AcceptsNewProtocol is the positive control:
|
|
// confirms that a stateless server still accepts new-protocol requests
|
|
// (the rejection in TestStreamableStateful_RejectsNewProtocol must not
|
|
// fire on Stateless: true).
|
|
func TestStreamableStateless_AcceptsNewProtocol(t *testing.T) {
|
|
|
|
server := NewServer(testImpl, nil)
|
|
AddTool(server, &Tool{Name: "noop"},
|
|
func(ctx context.Context, req *CallToolRequest, args struct{}) (*CallToolResult, any, error) {
|
|
return &CallToolResult{Content: []Content{&TextContent{Text: "ok"}}}, nil, nil
|
|
})
|
|
handler := NewStreamableHTTPHandler(
|
|
func(*http.Request) *Server { return server },
|
|
&StreamableHTTPOptions{Stateless: true},
|
|
)
|
|
httpServer := httptest.NewServer(handler)
|
|
defer httpServer.Close()
|
|
|
|
body := newProtocolBody(t, "noop", struct{}{})
|
|
req, err := http.NewRequest(http.MethodPost, httpServer.URL, bytes.NewReader(body))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("Accept", "application/json, text/event-stream")
|
|
req.Header.Set(protocolVersionHeader, protocolVersion20260728)
|
|
req.Header.Set(methodHeader, "tools/call")
|
|
req.Header.Set(nameHeader, "noop")
|
|
|
|
resp, err := http.DefaultClient.Do(req)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer resp.Body.Close()
|
|
if resp.StatusCode != http.StatusOK {
|
|
respBody, _ := io.ReadAll(resp.Body)
|
|
t.Fatalf("status = %d, want 200; body = %s", resp.StatusCode, respBody)
|
|
}
|
|
}
|
|
|
|
// TestStreamableStateless_NotificationMetaValidation checks that the SEP-2575
|
|
// per-request `_meta` triple is required of calls only. NotificationParams
|
|
// declares `_meta` optional with no protocolVersion, so a notification that
|
|
// omits it is well-formed and must be accepted with 202.
|
|
func TestStreamableStateless_NotificationMetaValidation(t *testing.T) {
|
|
newProtocolMeta := map[string]any{
|
|
MetaKeyProtocolVersion: protocolVersion20260728,
|
|
MetaKeyClientInfo: map[string]any{"name": "new-proto-client", "version": "9.9"},
|
|
MetaKeyClientCapabilities: map[string]any{},
|
|
}
|
|
|
|
tests := []struct {
|
|
name string
|
|
message map[string]any
|
|
wantStatus int
|
|
}{
|
|
{
|
|
name: "cancelled without meta",
|
|
message: map[string]any{
|
|
"jsonrpc": "2.0",
|
|
"method": notificationCancelled,
|
|
"params": map[string]any{"requestID": 1, "reason": "context canceled"},
|
|
},
|
|
wantStatus: http.StatusAccepted,
|
|
},
|
|
{
|
|
name: "progress without meta",
|
|
message: map[string]any{
|
|
"jsonrpc": "2.0",
|
|
"method": notificationProgress,
|
|
"params": map[string]any{"progressToken": "t", "progress": 1},
|
|
},
|
|
wantStatus: http.StatusAccepted,
|
|
},
|
|
{
|
|
name: "notification with matching meta",
|
|
message: map[string]any{
|
|
"jsonrpc": "2.0",
|
|
"method": notificationCancelled,
|
|
"params": map[string]any{
|
|
"_meta": newProtocolMeta,
|
|
"requestID": 1,
|
|
},
|
|
},
|
|
wantStatus: http.StatusAccepted,
|
|
},
|
|
{
|
|
// A notification that volunteers a version is still held to the
|
|
// header-match rule, so relaxing the requirement does not open a
|
|
// hole for inconsistent messages.
|
|
name: "notification with mismatched meta",
|
|
message: map[string]any{
|
|
"jsonrpc": "2.0",
|
|
"method": notificationCancelled,
|
|
"params": map[string]any{
|
|
"_meta": map[string]any{
|
|
MetaKeyProtocolVersion: "2025-06-18",
|
|
MetaKeyClientCapabilities: map[string]any{},
|
|
},
|
|
"requestID": 1,
|
|
},
|
|
},
|
|
wantStatus: http.StatusBadRequest,
|
|
},
|
|
{
|
|
// Calls remain subject to the requirement.
|
|
name: "call without meta",
|
|
message: map[string]any{
|
|
"jsonrpc": "2.0",
|
|
"id": 1,
|
|
"method": "tools/list",
|
|
"params": map[string]any{},
|
|
},
|
|
wantStatus: http.StatusBadRequest,
|
|
},
|
|
}
|
|
|
|
server := NewServer(testImpl, nil)
|
|
AddTool(server, &Tool{Name: "noop"},
|
|
func(ctx context.Context, req *CallToolRequest, args struct{}) (*CallToolResult, any, error) {
|
|
return &CallToolResult{Content: []Content{&TextContent{Text: "ok"}}}, nil, nil
|
|
})
|
|
handler := NewStreamableHTTPHandler(
|
|
func(*http.Request) *Server { return server },
|
|
&StreamableHTTPOptions{Stateless: true},
|
|
)
|
|
httpServer := httptest.NewServer(handler)
|
|
defer httpServer.Close()
|
|
|
|
for _, test := range tests {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
body, err := json.Marshal(test.message)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
req, err := http.NewRequest(http.MethodPost, httpServer.URL, bytes.NewReader(body))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("Accept", "application/json, text/event-stream")
|
|
req.Header.Set(protocolVersionHeader, protocolVersion20260728)
|
|
req.Header.Set(methodHeader, test.message["method"].(string))
|
|
|
|
resp, err := http.DefaultClient.Do(req)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer resp.Body.Close()
|
|
respBody, _ := io.ReadAll(resp.Body)
|
|
if resp.StatusCode != test.wantStatus {
|
|
t.Fatalf("status = %d, want %d; body = %s", resp.StatusCode, test.wantStatus, respBody)
|
|
}
|
|
if test.wantStatus == http.StatusAccepted && len(respBody) > 0 {
|
|
t.Errorf("accepted notification returned body %q, want empty", respBody)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestStreamableClientUnsupportedVersionFallback exercises the full
|
|
// SEP-2575 fallback. The client requests protocolVersion20260728, which the
|
|
// server is configured to NOT advertise in supportedProtocolVersions for the
|
|
// duration of this test. The server therefore rejects the server/discover
|
|
// POST at the transport-level header validation with a plain HTTP 400 ("Bad
|
|
// Request: Unsupported protocol version ..."). The streamable client must
|
|
// recognize this body, keep the connection alive, and successfully complete
|
|
// the legacy initialize handshake.
|
|
func TestStreamableClientUnsupportedVersionFallback(t *testing.T) {
|
|
// Remove protocolVersion20260728 from the server-side supported list so
|
|
// the fallback path is exercised.
|
|
orig := supportedProtocolVersions
|
|
supportedProtocolVersions = slices.DeleteFunc(slices.Clone(orig), func(v string) bool {
|
|
return v == protocolVersion20260728
|
|
})
|
|
t.Cleanup(func() { supportedProtocolVersions = orig })
|
|
|
|
ctx := context.Background()
|
|
|
|
server := NewServer(testImpl, nil)
|
|
handler := NewStreamableHTTPHandler(
|
|
func(*http.Request) *Server { return server },
|
|
nil,
|
|
)
|
|
httpServer := httptest.NewServer(handler)
|
|
defer httpServer.Close()
|
|
|
|
client := NewClient(testImpl, nil)
|
|
transport := &StreamableClientTransport{Endpoint: httpServer.URL}
|
|
|
|
session, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728})
|
|
if err != nil {
|
|
t.Fatalf("Connect: %v", err)
|
|
}
|
|
defer session.Close()
|
|
|
|
ir := session.InitializeResult()
|
|
if ir == nil {
|
|
t.Fatal("InitializeResult is nil; expected legacy initialize to populate it")
|
|
}
|
|
// The fallback initialize explicitly requests protocolVersion20251125
|
|
// (see client.go), so the server negotiates that version.
|
|
if ir.ProtocolVersion != protocolVersion20251125 {
|
|
t.Errorf("InitializeResult.ProtocolVersion = %q, want %q (legacy fallback)",
|
|
ir.ProtocolVersion, protocolVersion20251125)
|
|
}
|
|
|
|
// Verify the session is fully usable after the fallback by issuing a
|
|
// real call against the server.
|
|
if err := session.Ping(ctx, nil); err != nil {
|
|
t.Errorf("Ping after fallback initialize: %v", err)
|
|
}
|
|
}
|
|
|
|
// TestStreamableStateful_AcceptsDiscover verifies that a stateful HTTP server
|
|
// accepts a server/discover probe carrying MCP-Protocol-Version: 2026-07-28
|
|
// (and the matching _meta.protocolVersion), instead of rejecting it with the
|
|
// "stateless required" 400. The SEP-2575 client flow has the client probing
|
|
// the server with the new protocol version to learn which versions are
|
|
// supported.
|
|
func TestStreamableStateful_AcceptsDiscover(t *testing.T) {
|
|
|
|
server := NewServer(testImpl, nil)
|
|
handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, nil)
|
|
httpServer := httptest.NewServer(mustNotPanic(t, handler))
|
|
defer httpServer.Close()
|
|
|
|
body, err := json.Marshal(map[string]any{
|
|
"jsonrpc": "2.0",
|
|
"id": 1,
|
|
"method": methodDiscover,
|
|
"params": map[string]any{
|
|
"_meta": map[string]any{
|
|
MetaKeyProtocolVersion: protocolVersion20260728,
|
|
MetaKeyClientInfo: map[string]any{"name": "new-proto-client", "version": "9.9"},
|
|
MetaKeyClientCapabilities: map[string]any{"sampling": map[string]any{}},
|
|
},
|
|
},
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
req, err := http.NewRequest(http.MethodPost, httpServer.URL, bytes.NewReader(body))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("Accept", "application/json, text/event-stream")
|
|
req.Header.Set(protocolVersionHeader, protocolVersion20260728)
|
|
req.Header.Set(methodHeader, methodDiscover)
|
|
|
|
resp, err := http.DefaultClient.Do(req)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer resp.Body.Close()
|
|
respBody, _ := io.ReadAll(resp.Body)
|
|
if resp.StatusCode != http.StatusOK {
|
|
t.Fatalf("status = %d, want 200; body = %s", resp.StatusCode, respBody)
|
|
}
|
|
|
|
// Parse the JSON-RPC response. The body may arrive as a bare JSON object
|
|
// or as a single SSE event depending on the Accept negotiation; both
|
|
// shapes are valid here.
|
|
jsonPayload := respBody
|
|
if i := bytes.Index(respBody, []byte("data: ")); i >= 0 {
|
|
jsonPayload = respBody[i+len("data: "):]
|
|
if j := bytes.IndexByte(jsonPayload, '\n'); j >= 0 {
|
|
jsonPayload = jsonPayload[:j]
|
|
}
|
|
}
|
|
var rpcResp struct {
|
|
Result *DiscoverResult `json:"result"`
|
|
Error *jsonrpc.Error `json:"error"`
|
|
}
|
|
if err := json.Unmarshal(jsonPayload, &rpcResp); err != nil {
|
|
t.Fatalf("unmarshal response %q: %v", respBody, err)
|
|
}
|
|
if rpcResp.Error != nil {
|
|
t.Fatalf("discover returned error: %+v (body = %s)", rpcResp.Error, respBody)
|
|
}
|
|
if rpcResp.Result == nil {
|
|
t.Fatalf("discover returned no result; body = %s", respBody)
|
|
}
|
|
if slices.Contains(rpcResp.Result.SupportedVersions, protocolVersion20260728) {
|
|
t.Errorf("DiscoverResult.SupportedVersions = %v, must not include %q on a stateful transport",
|
|
rpcResp.Result.SupportedVersions, protocolVersion20260728)
|
|
}
|
|
if len(rpcResp.Result.SupportedVersions) == 0 {
|
|
t.Errorf("DiscoverResult.SupportedVersions is empty; want at least one legacy version")
|
|
}
|
|
}
|
|
|
|
// On a stateful StreamableHTTPHandler, a v1.7.0-style client always
|
|
// probes server/discover before falling back to initialize. That probe used
|
|
// to create a ServerSession that was registered in the handler's session
|
|
// map but never surfaced via Mcp-Session-Id, so the client could not DELETE
|
|
// it and it was leaked for the lifetime of the process.
|
|
//
|
|
// The fix in Server.discover skips populating InitializeParams when the
|
|
// transport cannot serve the new protocol, which lets the safety-net
|
|
// cleanup in serveStatefulPOST close the otherwise-unaddressable session.
|
|
func TestStreamableStateful_DiscoverDoesNotLeakSession(t *testing.T) {
|
|
ctx := context.Background()
|
|
server := NewServer(testImpl, nil)
|
|
AddTool(server, &Tool{Name: "noop", Description: "noop"},
|
|
func(context.Context, *CallToolRequest, struct{}) (*CallToolResult, any, error) {
|
|
return &CallToolResult{}, nil, nil
|
|
})
|
|
|
|
// Stateful (default). SessionTimeout unset so nothing masks the leak
|
|
// via idle eviction.
|
|
handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, nil)
|
|
httpServer := httptest.NewServer(handler)
|
|
defer httpServer.Close()
|
|
|
|
countSessions := func() int {
|
|
n := 0
|
|
for range server.Sessions() {
|
|
n++
|
|
}
|
|
return n
|
|
}
|
|
|
|
const clients = 3
|
|
for i := 1; i <= clients; i++ {
|
|
c := NewClient(&Implementation{Name: "c", Version: "1.0.0"}, nil)
|
|
sess, err := c.Connect(ctx, &StreamableClientTransport{
|
|
Endpoint: httpServer.URL,
|
|
DisableStandaloneSSE: true,
|
|
MaxRetries: -1,
|
|
}, nil)
|
|
if err != nil {
|
|
t.Fatalf("client %d Connect: %v", i, err)
|
|
}
|
|
if _, err := sess.ListTools(ctx, nil); err != nil {
|
|
t.Fatalf("client %d ListTools: %v", i, err)
|
|
}
|
|
// Exactly one addressable session should exist while this client
|
|
// is connected. Anything more means a discover-only session
|
|
// slipped through.
|
|
if got := countSessions(); got != 1 {
|
|
t.Errorf("client %d connected: sessions = %d, want 1", i, got)
|
|
}
|
|
if err := sess.Close(); err != nil {
|
|
t.Fatalf("client %d Close: %v", i, err)
|
|
}
|
|
if got := countSessions(); got != 0 {
|
|
t.Errorf("client %d closed: sessions = %d, want 0", i, got)
|
|
}
|
|
}
|
|
|
|
if got := countSessions(); got != 0 {
|
|
t.Errorf("final leaked sessions = %d, want 0", got)
|
|
}
|
|
}
|
|
|
|
// TestStreamableHTTP_E2E_DiscoverSuccess is a full end-to-end smoke test for
|
|
// SEP-2575 over the streamable HTTP transport.
|
|
func TestStreamableHTTP_E2E_DiscoverSuccess(t *testing.T) {
|
|
ctx := context.Background()
|
|
server := NewServer(&Implementation{Name: "e2e-server", Version: "v1"}, nil)
|
|
// Register a simple tool so we can prove the session is usable end-to-end.
|
|
AddTool(server, &Tool{Name: "echo", Description: "echoes its input"},
|
|
func(_ context.Context, _ *CallToolRequest, args struct {
|
|
Msg string `json:"msg"`
|
|
}) (*CallToolResult, struct{}, error) {
|
|
return &CallToolResult{
|
|
Content: []Content{&TextContent{Text: args.Msg}},
|
|
}, struct{}{}, nil
|
|
},
|
|
)
|
|
|
|
handler := NewStreamableHTTPHandler(
|
|
func(*http.Request) *Server { return server },
|
|
&StreamableHTTPOptions{Stateless: true},
|
|
)
|
|
httpServer := httptest.NewServer(handler)
|
|
defer httpServer.Close()
|
|
|
|
client := NewClient(&Implementation{Name: "e2e-client", Version: "v1"}, nil)
|
|
transport := &StreamableClientTransport{Endpoint: httpServer.URL}
|
|
cs, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728})
|
|
if err != nil {
|
|
t.Fatalf("Connect: %v", err)
|
|
}
|
|
defer cs.Close()
|
|
|
|
ir := cs.InitializeResult()
|
|
if ir == nil {
|
|
t.Fatal("InitializeResult is nil after Connect; discover should have populated it")
|
|
}
|
|
if ir.ProtocolVersion != protocolVersion20260728 {
|
|
t.Errorf("InitializeResult.ProtocolVersion = %q, want %q (negotiated via discover)",
|
|
ir.ProtocolVersion, protocolVersion20260728)
|
|
}
|
|
if ir.ServerInfo == nil || ir.ServerInfo.Name != "e2e-server" {
|
|
t.Errorf("InitializeResult.ServerInfo = %+v, want name=e2e-server", ir.ServerInfo)
|
|
}
|
|
|
|
// Prove the session is fully usable: list tools and call one.
|
|
tools, err := cs.ListTools(ctx, nil)
|
|
if err != nil {
|
|
t.Fatalf("ListTools: %v", err)
|
|
}
|
|
if len(tools.Tools) != 1 || tools.Tools[0].Name != "echo" {
|
|
t.Errorf("ListTools = %+v, want one tool named 'echo'", tools.Tools)
|
|
}
|
|
res, err := cs.CallTool(ctx, &CallToolParams{
|
|
Name: "echo",
|
|
Arguments: map[string]any{"msg": "hello"},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("CallTool: %v", err)
|
|
}
|
|
if len(res.Content) != 1 {
|
|
t.Fatalf("CallTool result content = %+v, want 1 entry", res.Content)
|
|
}
|
|
tc, ok := res.Content[0].(*TextContent)
|
|
if !ok || tc.Text != "hello" {
|
|
t.Errorf("CallTool result[0] = %+v, want TextContent{Text:\"hello\"}", res.Content[0])
|
|
}
|
|
}
|
|
|
|
// A POST whose call ID is already in flight on the same session must be rejected
|
|
// with a JSON-RPC error, must not overwrite the existing request-to-stream
|
|
// mapping, and must not publish the duplicate message to the session's incoming
|
|
// channel.
|
|
func TestStreamableServerRejectsDuplicateInFlightRequestID(t *testing.T) {
|
|
id := jsonrpc2.Int64ID(1)
|
|
conn := &streamableServerConn{
|
|
logger: ensureLogger(nil),
|
|
incoming: make(chan jsonrpc.Message, 1),
|
|
done: make(chan struct{}),
|
|
streams: map[string]*stream{
|
|
"existing": {
|
|
id: "existing",
|
|
logger: ensureLogger(nil),
|
|
requests: map[jsonrpc.ID]struct{}{id: {}},
|
|
},
|
|
},
|
|
requestStreams: map[jsonrpc.ID]string{id: "existing"},
|
|
}
|
|
|
|
data, err := jsonrpc2.EncodeMessage(req(1, methodPing, &PingParams{}))
|
|
if err != nil {
|
|
t.Fatalf("EncodeMessage() error = %v", err)
|
|
}
|
|
ctx, cancel := context.WithTimeout(t.Context(), 20*time.Millisecond)
|
|
defer cancel()
|
|
httpReq := httptest.NewRequestWithContext(ctx, http.MethodPost, "/", bytes.NewReader(data))
|
|
httpReq.Header.Set("Content-Type", "application/json")
|
|
httpReq.Header.Set("Accept", "application/json, text/event-stream")
|
|
|
|
rec := httptest.NewRecorder()
|
|
conn.servePOST(rec, httpReq)
|
|
|
|
if rec.Code != http.StatusBadRequest {
|
|
t.Fatalf("status = %d, want %d; body = %s", rec.Code, http.StatusBadRequest, rec.Body.String())
|
|
}
|
|
msg, err := jsonrpc2.DecodeMessage(rec.Body.Bytes())
|
|
if err != nil {
|
|
t.Fatalf("DecodeMessage() error = %v; body = %s", err, rec.Body.String())
|
|
}
|
|
resp, ok := msg.(*jsonrpc.Response)
|
|
if !ok {
|
|
t.Fatalf("response type = %T, want *jsonrpc.Response", msg)
|
|
}
|
|
if got := resp.ID.Raw(); got != int64(1) {
|
|
t.Fatalf("response ID = %v, want 1", got)
|
|
}
|
|
var jerr *jsonrpc.Error
|
|
if !errors.As(resp.Error, &jerr) {
|
|
t.Fatalf("response error = %v, want *jsonrpc.Error", resp.Error)
|
|
}
|
|
if jerr.Code != jsonrpc.CodeInvalidRequest {
|
|
t.Fatalf("error code = %d, want %d", jerr.Code, jsonrpc.CodeInvalidRequest)
|
|
}
|
|
if got := conn.requestStreams[id]; got != "existing" {
|
|
t.Fatalf("requestStreams[%v] = %q, want %q", id, got, "existing")
|
|
}
|
|
select {
|
|
case msg := <-conn.incoming:
|
|
t.Fatalf("duplicate request was published to incoming: %v", msg)
|
|
default:
|
|
}
|
|
}
|
|
|
|
// TestStreamableServerPreservesSingleRequestBatchResponse verifies that a
|
|
// single-request JSON-RPC batch is returned as a one-element JSON array.
|
|
func TestStreamableServerPreservesSingleRequestBatchResponse(t *testing.T) {
|
|
server := NewServer(&Implementation{Name: "testServer", Version: "v1.0.0"}, nil)
|
|
handler := NewStreamableHTTPHandler(
|
|
func(*http.Request) *Server { return server },
|
|
&StreamableHTTPOptions{JSONResponse: true},
|
|
)
|
|
defer handler.closeAll()
|
|
|
|
httpServer := httptest.NewServer(handler)
|
|
defer httpServer.Close()
|
|
|
|
body := []byte(`[{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{},"clientInfo":{"name":"test-client","version":"1.0"}}}]`)
|
|
req, err := http.NewRequest(http.MethodPost, httpServer.URL, bytes.NewReader(body))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("Accept", "application/json, text/event-stream")
|
|
|
|
resp, err := http.DefaultClient.Do(req)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer resp.Body.Close()
|
|
data, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if resp.StatusCode != http.StatusOK {
|
|
t.Fatalf("status = %d, want %d; body = %s", resp.StatusCode, http.StatusOK, data)
|
|
}
|
|
|
|
var messages []json.RawMessage
|
|
if err := json.Unmarshal(data, &messages); err != nil {
|
|
t.Fatalf("response is not a JSON array: %v; body = %s", err, data)
|
|
}
|
|
if len(messages) != 1 {
|
|
t.Fatalf("response contains %d messages, want 1; body = %s", len(messages), data)
|
|
}
|
|
if _, err := jsonrpc2.DecodeMessage(messages[0]); err != nil {
|
|
t.Fatalf("response item is not a JSON-RPC message: %v", err)
|
|
}
|
|
}
|
|
|
|
// TestStreamableMaxRequestBodyBytes verifies that the streamable HTTP handler
|
|
// enforces StreamableHTTPOptions.MaxRequestBodyBytes on incoming request
|
|
// bodies. The limit must apply uniformly regardless of transfer encoding:
|
|
// Content-Length, Transfer-Encoding: chunked, or HTTP/2 (which has no
|
|
// Content-Length).
|
|
func TestStreamableMaxRequestBodyBytes(t *testing.T) {
|
|
// A minimally-valid JSON-RPC initialize request. The body-size check runs
|
|
// before parsing, so exact contents don't matter for the reject cases,
|
|
// but we use a real one for the "under limit" success case.
|
|
validInit := []byte(`{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"t","version":"v1"}}}`)
|
|
|
|
// Build a payload of the requested size by padding a JSON object. The
|
|
// server doesn't need to parse it; the MaxBytesReader trips first.
|
|
oversizedBody := func(n int) []byte {
|
|
b := make([]byte, n)
|
|
b[0] = '{'
|
|
for i := 1; i < n-1; i++ {
|
|
b[i] = ' '
|
|
}
|
|
b[n-1] = '}'
|
|
return b
|
|
}
|
|
|
|
tests := []struct {
|
|
name string
|
|
limit int64 // value assigned to MaxRequestBodyBytes
|
|
body []byte
|
|
wantStatus int
|
|
wantStatusDesc string
|
|
}{
|
|
{
|
|
name: "default limit rejects oversized",
|
|
limit: 0, // triggers DefaultMaxRequestBodyBytes (4 MiB)
|
|
body: oversizedBody(DefaultMaxRequestBodyBytes + 1),
|
|
wantStatus: http.StatusRequestEntityTooLarge,
|
|
wantStatusDesc: "413 for body > 4MiB default",
|
|
},
|
|
{
|
|
name: "custom limit rejects oversized",
|
|
limit: 1024,
|
|
body: oversizedBody(2048),
|
|
wantStatus: http.StatusRequestEntityTooLarge,
|
|
wantStatusDesc: "413 for body > custom 1KiB limit",
|
|
},
|
|
{
|
|
name: "chunked encoding does not bypass limit",
|
|
limit: 1024,
|
|
body: oversizedBody(4096),
|
|
wantStatus: http.StatusRequestEntityTooLarge,
|
|
wantStatusDesc: "413 for chunked body > limit (regression: Content-Length-only checks are bypassable)",
|
|
},
|
|
{
|
|
name: "negative limit disables the check",
|
|
limit: -1,
|
|
body: oversizedBody(2 * DefaultMaxRequestBodyBytes),
|
|
wantStatus: http.StatusBadRequest, // body is not valid JSON-RPC; but the size check must not fire
|
|
wantStatusDesc: "not 413 when limit is disabled",
|
|
},
|
|
{
|
|
name: "under limit succeeds",
|
|
limit: 64 * 1024,
|
|
body: validInit,
|
|
wantStatus: http.StatusOK,
|
|
wantStatusDesc: "200 for a valid initialize under the limit",
|
|
},
|
|
}
|
|
|
|
for _, tc := range tests {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
server := NewServer(testImpl, nil)
|
|
handler := NewStreamableHTTPHandler(
|
|
func(*http.Request) *Server { return server },
|
|
&StreamableHTTPOptions{
|
|
Stateless: true,
|
|
MaxRequestBodyBytes: tc.limit,
|
|
},
|
|
)
|
|
ts := httptest.NewServer(handler)
|
|
defer ts.Close()
|
|
|
|
req, err := http.NewRequest(http.MethodPost, ts.URL, bytes.NewReader(tc.body))
|
|
if err != nil {
|
|
t.Fatalf("NewRequest: %v", err)
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("Accept", "application/json, text/event-stream")
|
|
|
|
resp, err := ts.Client().Do(req)
|
|
if err != nil {
|
|
t.Fatalf("Do: %v", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
body, _ := io.ReadAll(resp.Body)
|
|
|
|
if resp.StatusCode != tc.wantStatus {
|
|
t.Fatalf("%s: status = %d, want %d (body=%q)",
|
|
tc.wantStatusDesc, resp.StatusCode, tc.wantStatus, string(body))
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestStreamableSupportedProtocolVersions_Header verifies that the
|
|
// MCP-Protocol-Version header is validated against
|
|
// [ServerOptions.SupportedProtocolVersions], and not merely against the
|
|
// versions the SDK knows about.
|
|
func TestStreamableSupportedProtocolVersions_Header(t *testing.T) {
|
|
const initBody = `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-11-25","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}`
|
|
|
|
newServer := func(t *testing.T, stateless bool) *httptest.Server {
|
|
server := NewServer(testImpl, &ServerOptions{
|
|
SupportedProtocolVersions: []string{protocolVersion20251125},
|
|
})
|
|
handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server },
|
|
&StreamableHTTPOptions{Stateless: stateless})
|
|
httpServer := httptest.NewServer(handler)
|
|
t.Cleanup(httpServer.Close)
|
|
return httpServer
|
|
}
|
|
post := func(t *testing.T, url, version, sessionID string) *http.Response {
|
|
t.Helper()
|
|
req, err := http.NewRequest(http.MethodPost, url, strings.NewReader(initBody))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("Accept", "application/json, text/event-stream")
|
|
req.Header.Set(protocolVersionHeader, version)
|
|
if sessionID != "" {
|
|
req.Header.Set(sessionIDHeader, sessionID)
|
|
}
|
|
resp, err := http.DefaultClient.Do(req)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
return resp
|
|
}
|
|
|
|
for _, stateless := range []bool{false, true} {
|
|
for _, test := range []struct {
|
|
version string
|
|
wantStatus int
|
|
}{
|
|
{protocolVersion20251125, http.StatusOK},
|
|
// Supported by the SDK, but excluded by the server.
|
|
{protocolVersion20250618, http.StatusBadRequest},
|
|
// Unknown to the SDK: rejected before any server is consulted.
|
|
{"1999-01-01", http.StatusBadRequest},
|
|
} {
|
|
t.Run(fmt.Sprintf("stateless=%v/%s", stateless, test.version), func(t *testing.T) {
|
|
resp := post(t, newServer(t, stateless).URL, test.version, "")
|
|
defer resp.Body.Close()
|
|
body, _ := io.ReadAll(resp.Body)
|
|
if resp.StatusCode != test.wantStatus {
|
|
t.Errorf("status = %d, want %d; body = %s", resp.StatusCode, test.wantStatus, body)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
t.Run("existing session", func(t *testing.T) {
|
|
httpServer := newServer(t, false)
|
|
resp := post(t, httpServer.URL, protocolVersion20251125, "")
|
|
io.Copy(io.Discard, resp.Body)
|
|
resp.Body.Close()
|
|
sessionID := resp.Header.Get(sessionIDHeader)
|
|
if sessionID == "" {
|
|
t.Fatalf("initialize response missing %s header", sessionIDHeader)
|
|
}
|
|
|
|
resp = post(t, httpServer.URL, protocolVersion20250618, sessionID)
|
|
defer resp.Body.Close()
|
|
body, _ := io.ReadAll(resp.Body)
|
|
if resp.StatusCode != http.StatusBadRequest {
|
|
t.Errorf("status = %d, want %d; body = %s", resp.StatusCode, http.StatusBadRequest, body)
|
|
}
|
|
})
|
|
}
|