// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file.
package mcp
import (
"bytes"
"context"
"encoding/json"
"log"
"log/slog"
"slices"
"strings"
"testing"
"time"
"github.com/google/go-cmp/cmp"
"github.com/google/jsonschema-go/jsonschema"
"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
)
type testItem struct {
Name string
Value string
}
type testListParams struct {
Cursor string
}
func (p *testListParams) cursorPtr() *string {
return &p.Cursor
}
type testListResult struct {
Items []*testItem
NextCursor string
}
func (r *testListResult) nextCursorPtr() *string {
return &r.NextCursor
}
var allTestItems = []*testItem{
{"alpha", "val-A"},
{"bravo", "val-B"},
{"charlie", "val-C"},
{"delta", "val-D"},
{"echo", "val-E"},
{"foxtrot", "val-F"},
{"golf", "val-G"},
{"hotel", "val-H"},
{"india", "val-I"},
{"juliet", "val-J"},
{"kilo", "val-K"},
}
// getCursor encodes a string input into a URL-safe base64 cursor,
// fatally logging any encoding errors.
func getCursor(input string) string {
cursor, err := encodeCursor(input)
if err != nil {
log.Fatalf("encodeCursor(%s) error = %v", input, err)
}
return cursor
}
func TestServerPaginateBasic(t *testing.T) {
testCases := []struct {
name string
initialItems []*testItem
inputCursor string
inputPageSize int
wantFeatures []*testItem
wantNextCursor string
wantErr bool
}{
{
name: "FirstPage_DefaultSize_Full",
initialItems: allTestItems,
inputCursor: "",
inputPageSize: 5,
wantFeatures: allTestItems[0:5],
wantNextCursor: getCursor("echo"), // Based on last item of first page
wantErr: false,
},
{
name: "SecondPage_DefaultSize_Full",
initialItems: allTestItems,
inputCursor: getCursor("echo"),
inputPageSize: 5,
wantFeatures: allTestItems[5:10],
wantNextCursor: getCursor("juliet"), // Based on last item of second page
wantErr: false,
},
{
name: "SecondPage_DefaultSize_Full_OutOfOrder",
initialItems: append(allTestItems[5:], allTestItems[0:5]...),
inputCursor: getCursor("echo"),
inputPageSize: 5,
wantFeatures: allTestItems[5:10],
wantNextCursor: getCursor("juliet"), // Based on last item of second page
wantErr: false,
},
{
name: "SecondPage_DefaultSize_Full_Duplicates",
initialItems: append(allTestItems, allTestItems[0:5]...),
inputCursor: getCursor("echo"),
inputPageSize: 5,
wantFeatures: allTestItems[5:10],
wantNextCursor: getCursor("juliet"), // Based on last item of second page
wantErr: false,
},
{
name: "LastPage_Remaining",
initialItems: allTestItems,
inputCursor: getCursor("juliet"),
inputPageSize: 5,
wantFeatures: allTestItems[10:11], // Only 1 item left
wantNextCursor: "", // No more pages
wantErr: false,
},
{
name: "PageSize_1",
initialItems: allTestItems,
inputCursor: "",
inputPageSize: 1,
wantFeatures: allTestItems[0:1],
wantNextCursor: getCursor("alpha"),
wantErr: false,
},
{
name: "PageSize_All",
initialItems: allTestItems,
inputCursor: "",
inputPageSize: len(allTestItems), // Page size equals total
wantFeatures: allTestItems,
wantNextCursor: "", // No more pages
wantErr: false,
},
{
name: "PageSize_LargerThanAll",
initialItems: allTestItems,
inputCursor: "",
inputPageSize: len(allTestItems) + 5, // Page size larger than total
wantFeatures: allTestItems,
wantNextCursor: "",
wantErr: false,
},
{
name: "EmptySet",
initialItems: nil,
inputCursor: "",
inputPageSize: 5,
wantFeatures: nil,
wantNextCursor: "",
wantErr: false,
},
{
name: "InvalidCursor",
initialItems: allTestItems,
inputCursor: "not-a-valid-gob-base64-cursor",
inputPageSize: 5,
wantFeatures: nil, // Should be nil for error cases
wantNextCursor: "",
wantErr: true,
},
{
name: "AboveNonExistentID",
initialItems: allTestItems,
inputCursor: getCursor("dne"), // A UID that doesn't exist
inputPageSize: 5,
wantFeatures: allTestItems[4:9], // Should return elements above UID.
wantNextCursor: getCursor("india"),
wantErr: false,
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
fs := newFeatureSet(func(t *testItem) string { return t.Name })
fs.add(tc.initialItems...)
params := &testListParams{Cursor: tc.inputCursor}
gotResult, err := paginateList(fs, tc.inputPageSize, params, &testListResult{}, func(res *testListResult, items []*testItem) {
res.Items = items
})
if (err != nil) != tc.wantErr {
t.Errorf("paginateList(%s) error, got %v, wantErr %v", tc.name, err, tc.wantErr)
}
if tc.wantErr {
return
}
if diff := cmp.Diff(tc.wantFeatures, gotResult.Items); diff != "" {
t.Errorf("paginateList(%s) mismatch (-want +got):\n%s", tc.name, diff)
}
if tc.wantNextCursor != gotResult.NextCursor {
t.Errorf("paginateList(%s) nextCursor, got %v, want %v", tc.name, gotResult.NextCursor, tc.wantNextCursor)
}
})
}
}
func TestServerPaginateVariousPageSizes(t *testing.T) {
fs := newFeatureSet(func(t *testItem) string { return t.Name })
fs.add(allTestItems...)
// Try all possible page sizes, ensuring we get the correct list of items.
for pageSize := 1; pageSize < len(allTestItems)+1; pageSize++ {
var gotItems []*testItem
var nextCursor string
wantChunks := slices.Collect(slices.Chunk(allTestItems, pageSize))
index := 0
// Iterate through all pages, comparing sub-slices to the paginated list.
for {
params := &testListParams{Cursor: nextCursor}
gotResult, err := paginateList(fs, pageSize, params, &testListResult{}, func(res *testListResult, items []*testItem) {
res.Items = items
})
if err != nil {
t.Fatalf("paginateList() unexpected error for pageSize %d, cursor %q: %v", pageSize, nextCursor, err)
}
if diff := cmp.Diff(wantChunks[index], gotResult.Items); diff != "" {
t.Errorf("paginateList mismatch (-want +got):\n%s", diff)
}
gotItems = append(gotItems, gotResult.Items...)
nextCursor = gotResult.NextCursor
if nextCursor == "" {
break
}
index++
}
if len(gotItems) != len(allTestItems) {
t.Fatalf("paginateList() returned %d items, want %d", len(allTestItems), len(gotItems))
}
}
}
func TestServerCapabilities(t *testing.T) {
tool := &Tool{Name: "t", InputSchema: &jsonschema.Schema{Type: "object"}}
testCases := []struct {
name string
configureServer func(s *Server)
serverOpts ServerOptions
wantCapabilities *ServerCapabilities
}{
{
name: "no capabilities",
configureServer: func(s *Server) {},
wantCapabilities: &ServerCapabilities{
Logging: &LoggingCapabilities{},
},
},
{
name: "with prompts",
configureServer: func(s *Server) {
s.AddPrompt(&Prompt{Name: "p"}, nil)
},
wantCapabilities: &ServerCapabilities{
Logging: &LoggingCapabilities{},
Prompts: &PromptCapabilities{ListChanged: true},
},
},
{
name: "with resources",
configureServer: func(s *Server) {
s.AddResource(&Resource{URI: "file:///r"}, nil)
},
wantCapabilities: &ServerCapabilities{
Logging: &LoggingCapabilities{},
Resources: &ResourceCapabilities{ListChanged: true},
},
},
{
name: "with resource templates",
configureServer: func(s *Server) {
s.AddResourceTemplate(&ResourceTemplate{URITemplate: "file:///rt"}, nil)
},
wantCapabilities: &ServerCapabilities{
Logging: &LoggingCapabilities{},
Resources: &ResourceCapabilities{ListChanged: true},
},
},
{
name: "with resource subscriptions",
configureServer: func(s *Server) {
s.AddResourceTemplate(&ResourceTemplate{URITemplate: "file:///rt"}, nil)
},
serverOpts: ServerOptions{
SubscribeHandler: func(context.Context, *SubscribeRequest) error {
return nil
},
UnsubscribeHandler: func(context.Context, *UnsubscribeRequest) error {
return nil
},
},
wantCapabilities: &ServerCapabilities{
Logging: &LoggingCapabilities{},
Resources: &ResourceCapabilities{ListChanged: true, Subscribe: true},
},
},
{
name: "with tools",
configureServer: func(s *Server) {
s.AddTool(tool, nil)
},
wantCapabilities: &ServerCapabilities{
Logging: &LoggingCapabilities{},
Tools: &ToolCapabilities{ListChanged: true},
},
},
{
name: "with completions",
configureServer: func(s *Server) {},
serverOpts: ServerOptions{
CompletionHandler: func(context.Context, *CompleteRequest) (*CompleteResult, error) {
return nil, nil
},
},
wantCapabilities: &ServerCapabilities{
Logging: &LoggingCapabilities{},
Completions: &CompletionCapabilities{},
},
},
{
name: "all capabilities",
configureServer: func(s *Server) {
s.AddPrompt(&Prompt{Name: "p"}, nil)
s.AddResource(&Resource{URI: "file:///r"}, nil)
s.AddResourceTemplate(&ResourceTemplate{URITemplate: "file:///rt"}, nil)
s.AddTool(tool, nil)
},
serverOpts: ServerOptions{
SubscribeHandler: func(context.Context, *SubscribeRequest) error {
return nil
},
UnsubscribeHandler: func(context.Context, *UnsubscribeRequest) error {
return nil
},
CompletionHandler: func(context.Context, *CompleteRequest) (*CompleteResult, error) {
return nil, nil
},
},
wantCapabilities: &ServerCapabilities{
Completions: &CompletionCapabilities{},
Logging: &LoggingCapabilities{},
Prompts: &PromptCapabilities{ListChanged: true},
Resources: &ResourceCapabilities{ListChanged: true, Subscribe: true},
Tools: &ToolCapabilities{ListChanged: true},
},
},
{
name: "has features",
configureServer: func(s *Server) {},
serverOpts: ServerOptions{
HasPrompts: true,
HasResources: true,
HasTools: true,
},
wantCapabilities: &ServerCapabilities{
Logging: &LoggingCapabilities{},
Prompts: &PromptCapabilities{ListChanged: true},
Resources: &ResourceCapabilities{ListChanged: true},
Tools: &ToolCapabilities{ListChanged: true},
},
},
{
name: "empty capabilities",
configureServer: func(s *Server) {},
serverOpts: ServerOptions{
Capabilities: &ServerCapabilities{},
},
wantCapabilities: &ServerCapabilities{},
},
{
name: "no logging",
configureServer: func(s *Server) {},
serverOpts: ServerOptions{
Capabilities: &ServerCapabilities{
Tools: &ToolCapabilities{ListChanged: true},
},
},
wantCapabilities: &ServerCapabilities{
Tools: &ToolCapabilities{ListChanged: true},
},
},
{
name: "no list",
configureServer: func(s *Server) {},
serverOpts: ServerOptions{
Capabilities: &ServerCapabilities{
Tools: &ToolCapabilities{ListChanged: false},
Prompts: &PromptCapabilities{ListChanged: false},
},
},
wantCapabilities: &ServerCapabilities{
Tools: &ToolCapabilities{ListChanged: false},
Prompts: &PromptCapabilities{ListChanged: false},
},
},
{
name: "adding tools-list",
configureServer: func(s *Server) {
s.AddTool(tool, nil)
},
serverOpts: ServerOptions{
Capabilities: &ServerCapabilities{
Logging: &LoggingCapabilities{},
},
},
wantCapabilities: &ServerCapabilities{
Logging: &LoggingCapabilities{},
Tools: &ToolCapabilities{ListChanged: true},
},
},
{
name: "adding tools-no list",
configureServer: func(s *Server) {
s.AddTool(tool, nil)
},
serverOpts: ServerOptions{
Capabilities: &ServerCapabilities{
Tools: &ToolCapabilities{ListChanged: false},
},
},
wantCapabilities: &ServerCapabilities{
Tools: &ToolCapabilities{ListChanged: false},
},
},
{
name: "experimental preserved",
configureServer: func(s *Server) {},
serverOpts: ServerOptions{
Capabilities: &ServerCapabilities{
Experimental: map[string]any{"custom": "value"},
Logging: &LoggingCapabilities{},
},
},
wantCapabilities: &ServerCapabilities{
Experimental: map[string]any{"custom": "value"},
Logging: &LoggingCapabilities{},
},
},
{
name: "extensions preserved",
configureServer: func(s *Server) {},
serverOpts: func() ServerOptions {
caps := &ServerCapabilities{
Logging: &LoggingCapabilities{},
}
caps.AddExtension("io.example/ext1", map[string]any{"key": "value"})
caps.AddExtension("io.example/ext2", nil)
return ServerOptions{Capabilities: caps}
}(),
wantCapabilities: &ServerCapabilities{
Extensions: map[string]any{
"io.example/ext1": map[string]any{"key": "value"},
"io.example/ext2": map[string]any{},
},
Logging: &LoggingCapabilities{},
},
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
server := NewServer(testImpl, &tc.serverOpts)
tc.configureServer(server)
gotCapabilities := server.capabilities()
if diff := cmp.Diff(tc.wantCapabilities, gotCapabilities); diff != "" {
t.Errorf("capabilities() mismatch (-want +got):\n%s", diff)
}
})
}
}
func TestServerAddResourceTemplate(t *testing.T) {
tests := []struct {
name string
template string
expectPanic bool
}{
{"ValidFileTemplate", "file:///{a}/{b}", false},
{"ValidCustomScheme", "myproto:///{a}", false},
{"EmptyVariable", "file:///{}/{b}", true},
{"UnclosedVariable", "file:///{a", true},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
rt := ResourceTemplate{URITemplate: tt.template}
defer func() {
if r := recover(); r != nil {
if !tt.expectPanic {
t.Errorf("%s: unexpected panic: %v", tt.name, r)
}
} else {
if tt.expectPanic {
t.Errorf("%s: expected panic but did not panic", tt.name)
}
}
}()
s := NewServer(testImpl, nil)
s.AddResourceTemplate(&rt, nil)
})
}
}
// TestServerSessionkeepaliveCancelOverwritten is to verify that `ServerSession.keepaliveCancel` is assigned exactly once,
// ensuring that only a single goroutine is responsible for the session's keepalive ping mechanism.
func TestServerSessionkeepaliveCancelOverwritten(t *testing.T) {
// Set KeepAlive to a long duration to ensure the keepalive
// goroutine stays alive for the duration of the test without actually sending
// ping requests, since we don't have a real client connection established.
server := NewServer(testImpl, &ServerOptions{KeepAlive: 5 * time.Second})
ss := &ServerSession{server: server}
// 1. Initialize the session.
_, err := ss.initialize(context.Background(), &InitializeParams{})
if err != nil {
t.Fatalf("ServerSession initialize failed: %v", err)
}
// 2. Call 'initialized' for the first time. This should start the keepalive mechanism.
_, err = ss.initialized(context.Background(), &InitializedParams{})
if err != nil {
t.Fatalf("First initialized call failed: %v", err)
}
if ss.keepaliveCancel == nil {
t.Fatalf("expected ServerSession.keepaliveCancel to be set after the first call of initialized")
}
// Save the cancel function and use defer to ensure resources are cleaned up.
firstCancel := ss.keepaliveCancel
defer firstCancel()
// 3. Manually set the field to nil.
// Do this to facilitate the test's core assertion. The goal is to verify that
// 'ss.keepaliveCancel' is not assigned a second time. By setting it to nil,
// we can easily check after the next call if a new keepalive goroutine was started.
ss.keepaliveCancel = nil
// 4. Call 'initialized' for the second time. This should return an error.
_, err = ss.initialized(context.Background(), &InitializedParams{})
if err == nil {
t.Fatalf("Expected 'duplicate initialized received' error on second call, got nil")
}
// 5. Re-check the field to ensure it remains nil.
// Since 'initialized' correctly returned an error and did not call
// 'startKeepalive', the field should remain unchanged.
if ss.keepaliveCancel != nil {
t.Fatal("expected ServerSession.keepaliveCancel to be nil after we manually niled it and re-initialized")
}
}
// panicks reports whether f() panics.
func panics(f func()) (b bool) {
defer func() {
b = recover() != nil
}()
f()
return false
}
func TestAddTool(t *testing.T) {
// AddTool should panic if In or Out are not JSON objects.
s := NewServer(testImpl, nil)
if !panics(func() {
AddTool(s, &Tool{Name: "T1"}, func(context.Context, *CallToolRequest, string) (*CallToolResult, any, error) { return nil, nil, nil })
}) {
t.Error("bad In: expected panic")
}
if panics(func() {
AddTool(s, &Tool{Name: "T2"}, func(context.Context, *CallToolRequest, map[string]any) (*CallToolResult, any, error) {
return nil, nil, nil
})
}) {
t.Error("good In: expected no panic")
}
if !panics(func() {
AddTool(s, &Tool{Name: "T2"}, func(context.Context, *CallToolRequest, map[string]any) (*CallToolResult, int, error) {
return nil, 0, nil
})
}) {
t.Error("bad Out: expected panic")
}
}
func TestAddToolNameValidation(t *testing.T) {
tests := []struct {
label string
name string
wantLogContaining string
}{
{
label: "empty name",
name: "",
wantLogContaining: `tool name cannot be empty`,
},
{
label: "long name",
name: strings.Repeat("a", 129),
wantLogContaining: "exceeds maximum length of 128 characters",
},
{
label: "name with spaces",
name: "get user profile",
wantLogContaining: `tool name contains invalid characters: \" \"`,
},
{
label: "name with multiple invalid chars",
name: "user name@domain,com",
wantLogContaining: `tool name contains invalid characters: \" \", \"@\", \",\"`,
},
{
label: "name with unicode",
name: "tool-ñame",
wantLogContaining: `tool name contains invalid characters: \"ñ\"`,
},
{
label: "valid name",
name: "valid-tool_name.123",
wantLogContaining: "", // No log expected
},
}
for _, test := range tests {
t.Run(test.label, func(t *testing.T) {
var buf bytes.Buffer
s := NewServer(testImpl, &ServerOptions{
Logger: slog.New(slog.NewTextHandler(&buf, nil)),
})
// Use the generic AddTool as it also calls validateToolName.
AddTool(s, &Tool{Name: test.name}, func(context.Context, *CallToolRequest, any) (*CallToolResult, any, error) {
return nil, nil, nil
})
logOutput := buf.String()
if test.wantLogContaining != "" {
if !strings.Contains(logOutput, test.wantLogContaining) {
t.Errorf("log output =\n%s\nwant containing %q", logOutput, test.wantLogContaining)
}
} else {
if logOutput != "" {
t.Errorf("expected empty log output, got %q", logOutput)
}
}
})
}
}
type schema = jsonschema.Schema
func testToolForSchema[In, Out any](t *testing.T, tool *Tool, in string, out Out, wantIn, wantOut any, wantErrContaining string) {
t.Helper()
th := func(context.Context, *CallToolRequest, In) (*CallToolResult, Out, error) {
return nil, out, nil
}
gott, goth, err := toolForErr(tool, th, nil)
if err != nil {
t.Fatal(err)
}
if diff := cmp.Diff(wantIn, gott.InputSchema); diff != "" {
t.Errorf("input: mismatch (-want, +got):\n%s", diff)
}
if diff := cmp.Diff(wantOut, gott.OutputSchema); diff != "" {
t.Errorf("output: mismatch (-want, +got):\n%s", diff)
}
ctr := &CallToolRequest{
Params: &CallToolParamsRaw{
Arguments: json.RawMessage(in),
},
}
result, err := goth(context.Background(), ctr)
if wantErrContaining != "" {
if err == nil {
t.Errorf("got nil error, want error containing %q", wantErrContaining)
} else {
if !strings.Contains(err.Error(), wantErrContaining) {
t.Errorf("got error %q, want containing %q", err, wantErrContaining)
}
}
} else if err != nil {
t.Errorf("got error %v, want no error", err)
}
if gott.OutputSchema != nil && err == nil && !result.IsError {
// Check that structured content matches exactly.
unstructured := result.Content[0].(*TextContent).Text
structured := string(result.StructuredContent.(json.RawMessage))
if diff := cmp.Diff(unstructured, structured); diff != "" {
t.Errorf("Unstructured content does not match structured content exactly (-unstructured +structured):\n%s", diff)
}
}
}
// TestClientRootCapabilities verifies that the server correctly observes
// RootsV2 for various client capability configurations. This tests the fix
// for #607.
func TestClientRootCapabilities(t *testing.T) {
testCases := []struct {
name string
capabilities *string // JSON for the capabilities field; nil means omit the field
wantRootsV2 *RootCapabilities
}{
{
name: "capabilities field omitted",
capabilities: nil,
wantRootsV2: nil,
},
{
name: "empty capabilities",
capabilities: ptr(`{}`),
wantRootsV2: nil,
},
{
name: "capabilities with no roots",
capabilities: ptr(`{"sampling": {}}`),
wantRootsV2: nil,
},
{
name: "capabilities with empty roots",
capabilities: ptr(`{"roots": {}}`),
wantRootsV2: &RootCapabilities{ListChanged: false},
},
{
name: "capabilities with roots without listChanged",
capabilities: ptr(`{"roots": {"listChanged": false}}`),
wantRootsV2: &RootCapabilities{ListChanged: false},
},
{
name: "capabilities with roots with listChanged",
capabilities: ptr(`{"roots": {"listChanged": true}}`),
wantRootsV2: &RootCapabilities{ListChanged: true},
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
ctx := context.Background()
// Create a minimal server.
impl := &Implementation{Name: "testServer", Version: "v1.0.0"}
s := NewServer(impl, nil)
// Connect the server.
cTransport, sTransport := NewInMemoryTransports()
ss, err := s.Connect(ctx, sTransport, nil)
if err != nil {
t.Fatal(err)
}
// Connect the client JSON-RPC connection (raw, no client).
cConn, err := cTransport.Connect(ctx)
if err != nil {
t.Fatal(err)
}
// Build initialize params, optionally including capabilities.
var initParams json.RawMessage
if tc.capabilities != nil {
initParams = json.RawMessage(`{
"protocolVersion": "2025-06-18",
"capabilities": ` + *tc.capabilities + `,
"clientInfo": {"name": "TestClient", "version": "1.0.0"}
}`)
} else {
initParams = json.RawMessage(`{
"protocolVersion": "2025-06-18",
"clientInfo": {"name": "TestClient", "version": "1.0.0"}
}`)
}
initReq, err := jsonrpc2.NewCall(jsonrpc2.Int64ID(1), "initialize", initParams)
if err != nil {
t.Fatal(err)
}
if err := cConn.Write(ctx, initReq); err != nil {
t.Fatalf("Write failed: %v", err)
}
// Read the initialize response.
msg, err := cConn.Read(ctx)
if err != nil {
t.Fatalf("Read failed: %v", err)
}
resp, ok := msg.(*jsonrpc2.Response)
if !ok {
t.Fatalf("expected Response, got %T", msg)
}
if resp.Error != nil {
t.Fatalf("initialize failed: %v", resp.Error)
}
// Verify that the server session has the correct RootsV2 value.
params := ss.InitializeParams()
if params == nil {
t.Fatal("InitializeParams is nil")
}
var gotRootsV2 *RootCapabilities
if params.Capabilities != nil {
gotRootsV2 = params.Capabilities.RootsV2
}
if diff := cmp.Diff(tc.wantRootsV2, gotRootsV2); diff != "" {
t.Errorf("RootsV2 mismatch (-want +got):\n%s", diff)
}
// Close the client connection.
if err := cConn.Close(); err != nil {
t.Fatalf("Stream.Close failed: %v", err)
}
ss.Wait()
})
}
}
// TODO: move this to tool_test.go
func TestToolForSchemas(t *testing.T) {
// Validate that toolForErr handles schemas properly.
type in struct {
P int `json:"p,omitempty"`
}
type out struct {
B bool `json:"b,omitempty"`
}
var (
falseSchema = &schema{Not: &schema{}}
inSchema = &schema{
Type: "object",
AdditionalProperties: falseSchema,
Properties: map[string]*schema{"p": {Type: "integer"}},
PropertyOrder: []string{"p"},
}
inSchema2 = &schema{
Type: "object",
AdditionalProperties: falseSchema,
Properties: map[string]*schema{"p": {Type: "string"}},
}
inSchema3 = &schema{
Type: "object",
AdditionalProperties: falseSchema,
Properties: map[string]*schema{}, // empty map is preserved
}
outSchema = &schema{
Type: "object",
AdditionalProperties: falseSchema,
Properties: map[string]*schema{"b": {Type: "boolean"}},
PropertyOrder: []string{"b"},
}
outSchema2 = &schema{
Type: "object",
AdditionalProperties: falseSchema,
Properties: map[string]*schema{"b": {Type: "integer"}},
PropertyOrder: []string{"b"},
}
)
// Infer both schemas.
testToolForSchema[in](t, &Tool{}, `{"p":3}`, out{true}, inSchema, outSchema, "")
// Validate the input schema: expect an error if it's wrong.
// We can't test that the output schema is validated, because it's typed.
testToolForSchema[in](t, &Tool{}, `{"p":"x"}`, out{true}, inSchema, outSchema, `want "integer"`)
// Ignore type any for output.
testToolForSchema[in, any](t, &Tool{}, `{"p":3}`, 0, inSchema, nil, "")
// Input is still validated.
testToolForSchema[in, any](t, &Tool{}, `{"p":"x"}`, 0, inSchema, nil, `want "integer"`)
// Tool sets input schema: that is what's used.
testToolForSchema[in, any](t, &Tool{InputSchema: inSchema2}, `{"p":3}`, 0, inSchema2, nil, `want "string"`)
// Tool sets input schema, empty properties map.
testToolForSchema[in, any](t, &Tool{InputSchema: inSchema3}, `{}`, 0, inSchema3, nil, "")
// Tool sets output schema: that is what's used, and validation happens.
testToolForSchema[in, any](t, &Tool{OutputSchema: outSchema2}, `{"p":3}`, out{true},
inSchema, outSchema2, `want "integer"`)
// Check a slightly more complicated case.
type weatherOutput struct {
Summary string
AsOf time.Time
Source string
}
testToolForSchema[any](t, &Tool{}, `{}`, weatherOutput{},
&schema{Type: "object"},
&schema{
Type: "object",
Required: []string{"Summary", "AsOf", "Source"},
AdditionalProperties: falseSchema,
Properties: map[string]*schema{
"Summary": {Type: "string"},
"AsOf": {Type: "string"},
"Source": {Type: "string"},
},
PropertyOrder: []string{"Summary", "AsOf", "Source"},
},
"")
}
// TestServerCapabilitiesOverWire verifies that server capabilities are
// correctly sent over the wire during initialization.
func TestServerCapabilitiesOverWire(t *testing.T) {
tool := &Tool{Name: "test-tool", InputSchema: &jsonschema.Schema{Type: "object"}}
testCases := []struct {
name string
serverOpts *ServerOptions
configureServer func(s *Server)
wantCapabilities *ServerCapabilities
}{
{
name: "Default capabilities",
serverOpts: nil,
configureServer: func(s *Server) {},
wantCapabilities: &ServerCapabilities{
Logging: &LoggingCapabilities{},
},
},
{
name: "Custom Capabilities with tools",
serverOpts: &ServerOptions{
Capabilities: &ServerCapabilities{
Tools: &ToolCapabilities{ListChanged: false},
},
},
configureServer: func(s *Server) {},
wantCapabilities: &ServerCapabilities{
Tools: &ToolCapabilities{ListChanged: false},
},
},
{
name: "Dynamic tool capability",
serverOpts: &ServerOptions{
Capabilities: &ServerCapabilities{
Logging: &LoggingCapabilities{},
},
},
configureServer: func(s *Server) {
s.AddTool(tool, nil)
},
wantCapabilities: &ServerCapabilities{
Logging: &LoggingCapabilities{},
Tools: &ToolCapabilities{ListChanged: true},
},
},
{
name: "Extensions over wire",
serverOpts: func() *ServerOptions {
caps := &ServerCapabilities{
Logging: &LoggingCapabilities{},
}
caps.AddExtension("io.example/ext", map[string]any{"key": "value"})
return &ServerOptions{Capabilities: caps}
}(),
configureServer: func(s *Server) {},
wantCapabilities: &ServerCapabilities{
Extensions: map[string]any{
"io.example/ext": map[string]any{"key": "value"},
},
Logging: &LoggingCapabilities{},
},
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
ctx := context.Background()
// Create server.
impl := &Implementation{Name: "testServer", Version: "v1.0.0"}
server := NewServer(impl, tc.serverOpts)
tc.configureServer(server)
// Connect client and server.
cTransport, sTransport := NewInMemoryTransports()
ss, err := server.Connect(ctx, sTransport, nil)
if err != nil {
t.Fatal(err)
}
defer ss.Close()
client := NewClient(&Implementation{Name: "testClient", Version: "v1.0.0"}, nil)
cs, err := client.Connect(ctx, cTransport, nil)
if err != nil {
t.Fatal(err)
}
defer cs.Close()
// Check that the client received the expected capabilities.
initResult := cs.InitializeResult()
if initResult == nil {
t.Fatal("InitializeResult is nil")
}
if diff := cmp.Diff(tc.wantCapabilities, initResult.Capabilities); diff != "" {
t.Errorf("Capabilities mismatch (-want +got):\n%s", diff)
}
})
}
}