Repository navigation
mcp: bound server/discover probe on stdio transports #1334
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -7,12 +7,15 @@ | |
| package mcp | ||
|
|
||
| import ( | ||
| "bufio" | ||
| "context" | ||
| "encoding/json" | ||
| "fmt" | ||
| "io" | ||
| "log/slog" | ||
| "sync/atomic" | ||
| "testing" | ||
| "time" | ||
|
|
||
| "github.com/google/go-cmp/cmp" | ||
| "github.com/google/go-cmp/cmp/cmpopts" | ||
|
|
@@ -1094,6 +1097,159 @@ func TestClientConnectDiscover_UnsupportedVersionNegotiation(t *testing.T) { | |
| } | ||
| } | ||
|
|
||
| // serveSilentStdio emulates a handshake-era stdio server that answers only | ||
| // initialize and tools/list, silently dropping every other message. | ||
| func serveSilentStdio(r io.Reader, w io.Writer) { | ||
| sc := bufio.NewScanner(r) | ||
| for sc.Scan() { | ||
| var msg struct { | ||
| ID json.RawMessage `json:"id"` | ||
| Method string `json:"method"` | ||
| } | ||
| if json.Unmarshal(sc.Bytes(), &msg) != nil { | ||
| continue | ||
| } | ||
| var result any | ||
| switch msg.Method { | ||
| case methodInitialize: | ||
| result = map[string]any{ | ||
| "protocolVersion": protocolVersion20251125, | ||
| "capabilities": map[string]any{"tools": map[string]any{}}, | ||
| "serverInfo": map[string]any{"name": "silent", "version": "v1"}, | ||
| } | ||
| case methodListTools: | ||
| result = map[string]any{"tools": []any{}} | ||
| default: | ||
| continue | ||
| } | ||
| data, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": msg.ID, "result": result}) | ||
| if _, err := w.Write(append(data, '\n')); err != nil { | ||
| return | ||
| } | ||
| } | ||
| } | ||
|
|
||
| // delegatingTransport is a user-defined Transport wrapping another one, as an | ||
| // instrumentation or filtering decorator would. | ||
| type delegatingTransport struct { | ||
| Transport | ||
| } | ||
|
|
||
| // connWrappingTransport is a user-defined Transport that also wraps the | ||
| // Connection, which hides the underlying stream connection from Connect. | ||
| type connWrappingTransport struct { | ||
| Transport | ||
| } | ||
|
|
||
| func (t connWrappingTransport) Connect(ctx context.Context) (Connection, error) { | ||
| conn, err := t.Transport.Connect(ctx) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| return struct{ Connection }{conn}, nil | ||
| } | ||
|
|
||
| func silentStdioTransport() *IOTransport { | ||
| cr, sw := io.Pipe() | ||
| sr, cw := io.Pipe() | ||
| go serveSilentStdio(sr, sw) | ||
| return &IOTransport{Reader: cr, Writer: cw} | ||
| } | ||
|
|
||
| func connectWithFallback(t *testing.T, ctx context.Context, tr Transport, opts *ClientSessionOptions) { | ||
| t.Helper() | ||
| client := NewClient(&Implementation{Name: "client", Version: "v1"}, nil) | ||
| cs, err := client.Connect(ctx, tr, opts) | ||
| if err != nil { | ||
| t.Fatalf("client.Connect: %v", err) | ||
| } | ||
| defer cs.Close() | ||
|
|
||
| if got, want := cs.InitializeResult().ProtocolVersion, protocolVersion20251125; got != want { | ||
| t.Errorf("InitializeResult.ProtocolVersion = %q, want %q", got, want) | ||
| } | ||
| if _, err := cs.ListTools(ctx, nil); err != nil { | ||
| t.Errorf("ListTools after fallback initialize: %v", err) | ||
| } | ||
| } | ||
|
|
||
| // TestClientConnectDiscover_SilentStdioServer verifies that Connect falls back | ||
| // to initialize when a stdio server never answers server/discover (#1332). | ||
| func TestClientConnectDiscover_SilentStdioServer(t *testing.T) { | ||
| ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) | ||
| defer cancel() | ||
| connectWithFallback(t, ctx, silentStdioTransport(), &ClientSessionOptions{DiscoverTimeout: 50 * time.Millisecond}) | ||
| } | ||
|
|
||
| // TestClientConnectDiscover_SilentStdioServerShortDeadline verifies that the | ||
| // default discover probe leaves part of a short caller deadline for | ||
| // initialize, including when the stdio transport is wrapped. | ||
| func TestClientConnectDiscover_SilentStdioServerShortDeadline(t *testing.T) { | ||
| transports := map[string]func() Transport{ | ||
| "plain": func() Transport { return silentStdioTransport() }, | ||
| "logging": func() Transport { return &LoggingTransport{Transport: silentStdioTransport(), Writer: io.Discard} }, | ||
| "custom": func() Transport { return delegatingTransport{silentStdioTransport()} }, | ||
| } | ||
| for name, newTransport := range transports { | ||
| t.Run(name, func(t *testing.T) { | ||
| ctx, cancel := context.WithTimeout(context.Background(), time.Second) | ||
| defer cancel() | ||
| connectWithFallback(t, ctx, newTransport(), nil) | ||
| }) | ||
| } | ||
| } | ||
|
|
||
| // TestClientConnectDiscover_ExplicitTimeout verifies that an explicit | ||
| // DiscoverTimeout also applies to connections the default doesn't recognize. | ||
| func TestClientConnectDiscover_ExplicitTimeout(t *testing.T) { | ||
| ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) | ||
| defer cancel() | ||
| tr := connWrappingTransport{silentStdioTransport()} | ||
| connectWithFallback(t, ctx, tr, &ClientSessionOptions{DiscoverTimeout: 50 * time.Millisecond}) | ||
| } | ||
|
|
||
| // TestClientConnectDiscover_NegativeTimeout verifies that a negative | ||
| // DiscoverTimeout disables the bound, so a silent server blocks Connect until | ||
| // the caller's deadline. | ||
| func TestClientConnectDiscover_NegativeTimeout(t *testing.T) { | ||
| ctx, cancel := context.WithTimeout(context.Background(), 300*time.Millisecond) | ||
| defer cancel() | ||
| client := NewClient(&Implementation{Name: "client", Version: "v1"}, nil) | ||
| cs, err := client.Connect(ctx, silentStdioTransport(), &ClientSessionOptions{DiscoverTimeout: -1}) | ||
| if err == nil { | ||
| cs.Close() | ||
| t.Fatal("client.Connect succeeded, want an error from the caller's deadline") | ||
| } | ||
| } | ||
|
|
||
| // TestClientConnectDiscover_SlowStdioServer verifies that a modern stdio | ||
| // server answering server/discover after the probe timeout still connects | ||
| // through the initialize fallback, and its late reply is ignored. | ||
| func TestClientConnectDiscover_SlowStdioServer(t *testing.T) { | ||
| ctx := context.Background() | ||
|
|
||
| server := NewServer(&Implementation{Name: "slow-server", Version: "v1"}, nil) | ||
| server.AddReceivingMiddleware(func(next MethodHandler) MethodHandler { | ||
| return func(ctx context.Context, method string, req Request) (Result, error) { | ||
| if method == methodDiscover { | ||
| time.Sleep(200 * time.Millisecond) | ||
| } | ||
| return next(ctx, method, req) | ||
| } | ||
| }) | ||
|
|
||
| cr, sw := io.Pipe() | ||
| sr, cw := io.Pipe() | ||
| ss, err := server.Connect(ctx, &IOTransport{Reader: sr, Writer: sw}, nil) | ||
| if err != nil { | ||
| t.Fatalf("server.Connect: %v", err) | ||
| } | ||
| defer ss.Close() | ||
|
|
||
| tr := &IOTransport{Reader: cr, Writer: cw} | ||
| connectWithFallback(t, ctx, tr, &ClientSessionOptions{DiscoverTimeout: 50 * time.Millisecond}) | ||
| } | ||
|
|
||
| // TestClientRejectsMissingRequiredParams checks that server-to-client messages | ||
| // whose params are required by the schema are rejected when params is omitted | ||
| // or null, instead of reaching handlers that dereference them. | ||
|
|
||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
this default value seems too short, maybe 10 sec would be a better default