go-sdk/mcp/sampling_test.go

// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file.

package mcp

// TODO: move other sampling-related tests to this file.

import (
	"context"
	"strings"
	"testing"
)

func TestSamplingWithTools_Integration(t *testing.T) {
	ctx := context.Background()
	ct, st := NewInMemoryTransports()

	// Track what the client received
	var receivedParams *CreateMessageWithToolsParams

	// Client with tools capability, using CreateMessageWithToolsHandler
	client := NewClient(testImpl, &ClientOptions{
		CreateMessageWithToolsHandler: func(_ context.Context, req *CreateMessageWithToolsRequest) (*CreateMessageWithToolsResult, error) {
			receivedParams = req.Params
			// Return a tool use response
			return &CreateMessageWithToolsResult{
				Model: "test-model",
				Role:  "assistant",
				Content: []Content{&ToolUseContent{
					ID:    "tool_call_1",
					Name:  "calculator",
					Input: map[string]any{"x": 1.0, "y": 2.0},
				}},
				StopReason: "toolUse",
			}, nil
		},
		Capabilities: &ClientCapabilities{
			Sampling: &SamplingCapabilities{Tools: &SamplingToolsCapabilities{}},
		},
	})

	server := NewServer(testImpl, nil)
	ss, err := server.Connect(ctx, st, nil)
	if err != nil {
		t.Fatal(err)
	}
	defer ss.Close()

	cs, err := client.Connect(ctx, ct, nil)
	if err != nil {
		t.Fatal(err)
	}
	defer cs.Close()

	// Server sends CreateMessageWithTools
	result, err := ss.CreateMessageWithTools(ctx, &CreateMessageWithToolsParams{
		MaxTokens: 1000,
		Messages: []*SamplingMessageV2{
			{Role: "user", Content: []Content{&TextContent{Text: "Calculate 1+2"}}},
		},
		Tools: []*Tool{
			{
				Name:        "calculator",
				Description: "A calculator",
				InputSchema: map[string]any{
					"type": "object",
					"properties": map[string]any{
						"x": map[string]any{"type": "number"},
						"y": map[string]any{"type": "number"},
					},
				},
			},
		},
		ToolChoice: &ToolChoice{Mode: "auto"},
	})
	if err != nil {
		t.Fatalf("CreateMessageWithTools() error = %v", err)
	}

	// Verify client received the tools
	if receivedParams == nil {
		t.Fatal("client did not receive params")
	}
	if len(receivedParams.Tools) != 1 {
		t.Errorf("client received %d tools, want 1", len(receivedParams.Tools))
	}
	if receivedParams.Tools[0].Name != "calculator" {
		t.Errorf("tool name = %v, want calculator", receivedParams.Tools[0].Name)
	}
	if receivedParams.ToolChoice == nil || receivedParams.ToolChoice.Mode != "auto" {
		t.Errorf("tool choice mode = %v, want auto", receivedParams.ToolChoice)
	}

	// Verify server received the tool use response
	if result.StopReason != "toolUse" {
		t.Errorf("StopReason = %v, want toolUse", result.StopReason)
	}
	if len(result.Content) != 1 {
		t.Fatalf("len(Content) = %d, want 1", len(result.Content))
	}
	toolUse, ok := result.Content[0].(*ToolUseContent)
	if !ok {
		t.Fatalf("Content[0] type = %T, want *ToolUseContent", result.Content[0])
	}
	if toolUse.ID != "tool_call_1" {
		t.Errorf("ToolUse.ID = %v, want tool_call_1", toolUse.ID)
	}
	if toolUse.Name != "calculator" {
		t.Errorf("ToolUse.Name = %v, want calculator", toolUse.Name)
	}
}

func TestSamplingWithToolResult_Integration(t *testing.T) {
	ctx := context.Background()
	ct, st := NewInMemoryTransports()

	// Track messages received by client
	var receivedMessages []*SamplingMessage

	client := NewClient(testImpl, &ClientOptions{
		CreateMessageHandler: func(_ context.Context, req *CreateMessageRequest) (*CreateMessageResult, error) {
			receivedMessages = req.Params.Messages
			return &CreateMessageResult{
				Model:   "test-model",
				Role:    "assistant",
				Content: &TextContent{Text: "The result is 3"},
			}, nil
		},
		Capabilities: &ClientCapabilities{
			Sampling: &SamplingCapabilities{Tools: &SamplingToolsCapabilities{}},
		},
	})

	server := NewServer(testImpl, nil)
	ss, err := server.Connect(ctx, st, nil)
	if err != nil {
		t.Fatal(err)
	}
	defer ss.Close()

	cs, err := client.Connect(ctx, ct, nil)
	if err != nil {
		t.Fatal(err)
	}
	defer cs.Close()

	// Server sends CreateMessage with tool result in messages
	_, err = ss.CreateMessage(ctx, &CreateMessageParams{
		MaxTokens: 1000,
		Messages: []*SamplingMessage{
			{Role: "user", Content: &TextContent{Text: "Calculate 1+2"}},
			{Role: "assistant", Content: &ToolUseContent{
				ID:    "tool_1",
				Name:  "calculator",
				Input: map[string]any{"x": 1.0, "y": 2.0},
			}},
			{Role: "user", Content: &ToolResultContent{
				ToolUseID: "tool_1",
				Content:   []Content{&TextContent{Text: "3"}},
			}},
		},
	})
	if err != nil {
		t.Fatalf("CreateMessage() error = %v", err)
	}

	// Verify client received all messages including tool content
	if len(receivedMessages) != 3 {
		t.Fatalf("received %d messages, want 3", len(receivedMessages))
	}

	// Check first message is text
	if _, ok := receivedMessages[0].Content.(*TextContent); !ok {
		t.Errorf("message[0] content type = %T, want *TextContent", receivedMessages[0].Content)
	}

	// Check second message is tool use
	toolUse, ok := receivedMessages[1].Content.(*ToolUseContent)
	if !ok {
		t.Fatalf("message[1] content type = %T, want *ToolUseContent", receivedMessages[1].Content)
	}
	if toolUse.ID != "tool_1" {
		t.Errorf("toolUse.ID = %v, want tool_1", toolUse.ID)
	}

	// Check third message is tool result
	toolResult, ok := receivedMessages[2].Content.(*ToolResultContent)
	if !ok {
		t.Fatalf("message[2] content type = %T, want *ToolResultContent", receivedMessages[2].Content)
	}
	if toolResult.ToolUseID != "tool_1" {
		t.Errorf("toolResult.ToolUseID = %v, want tool_1", toolResult.ToolUseID)
	}
	if len(toolResult.Content) != 1 {
		t.Fatalf("toolResult.Content len = %d, want 1", len(toolResult.Content))
	}
	if tc, ok := toolResult.Content[0].(*TextContent); !ok || tc.Text != "3" {
		t.Errorf("toolResult.Content[0] = %v, want TextContent with '3'", toolResult.Content[0])
	}
}

func TestSamplingToolsCapabilities(t *testing.T) {
	ctx := context.Background()

	t.Run("client with explicit tools capability", func(t *testing.T) {
		ct, st := NewInMemoryTransports()

		client := NewClient(testImpl, &ClientOptions{
			CreateMessageHandler: func(_ context.Context, _ *CreateMessageRequest) (*CreateMessageResult, error) {
				return &CreateMessageResult{Model: "m", Content: &TextContent{}}, nil
			},
			Capabilities: &ClientCapabilities{
				Sampling: &SamplingCapabilities{
					Tools:   &SamplingToolsCapabilities{},
					Context: &SamplingContextCapabilities{},
				},
			},
		})

		server := NewServer(testImpl, nil)
		ss, err := server.Connect(ctx, st, nil)
		if err != nil {
			t.Fatal(err)
		}
		defer ss.Close()

		cs, err := client.Connect(ctx, ct, nil)
		if err != nil {
			t.Fatal(err)
		}
		defer cs.Close()

		// Check server sees client capabilities
		caps := ss.InitializeParams().Capabilities
		if caps.Sampling == nil {
			t.Fatal("client should advertise sampling capability")
		}
		if caps.Sampling.Tools == nil {
			t.Error("client should advertise sampling.tools capability")
		}
		if caps.Sampling.Context == nil {
			t.Error("client should advertise sampling.context capability")
		}
	})

	t.Run("client without tools capability", func(t *testing.T) {
		ct, st := NewInMemoryTransports()

		client := NewClient(testImpl, &ClientOptions{
			CreateMessageHandler: func(_ context.Context, _ *CreateMessageRequest) (*CreateMessageResult, error) {
				return &CreateMessageResult{Model: "m", Content: &TextContent{}}, nil
			},
			// No Capabilities.Sampling.Tools set
		})

		server := NewServer(testImpl, nil)
		ss, err := server.Connect(ctx, st, nil)
		if err != nil {
			t.Fatal(err)
		}
		defer ss.Close()

		cs, err := client.Connect(ctx, ct, nil)
		if err != nil {
			t.Fatal(err)
		}
		defer cs.Close()

		// Check server sees client capabilities
		caps := ss.InitializeParams().Capabilities
		if caps.Sampling == nil {
			t.Fatal("client should advertise sampling capability")
		}
		if caps.Sampling.Tools != nil {
			t.Error("client should NOT advertise sampling.tools capability")
		}
		if caps.Sampling.Context != nil {
			t.Error("client should NOT advertise sampling.context capability")
		}
	})

	t.Run("client with inferred tools capability", func(t *testing.T) {
		ct, st := NewInMemoryTransports()

		client := NewClient(testImpl, &ClientOptions{
			CreateMessageWithToolsHandler: func(_ context.Context, _ *CreateMessageWithToolsRequest) (*CreateMessageWithToolsResult, error) {
				return &CreateMessageWithToolsResult{Model: "m", Content: []Content{&TextContent{}}}, nil
			},
			// No explicit Capabilities set — tools should be inferred.
		})

		server := NewServer(testImpl, nil)
		ss, err := server.Connect(ctx, st, nil)
		if err != nil {
			t.Fatal(err)
		}
		defer ss.Close()

		cs, err := client.Connect(ctx, ct, nil)
		if err != nil {
			t.Fatal(err)
		}
		defer cs.Close()

		caps := ss.InitializeParams().Capabilities
		if caps.Sampling == nil {
			t.Fatal("client should advertise sampling capability")
		}
		if caps.Sampling.Tools == nil {
			t.Error("client should infer sampling.tools capability from CreateMessageWithToolsHandler")
		}
	})
}

func TestSamplingToolResultWithError_Integration(t *testing.T) {
	ctx := context.Background()
	ct, st := NewInMemoryTransports()

	var receivedMessages []*SamplingMessage

	client := NewClient(testImpl, &ClientOptions{
		CreateMessageHandler: func(_ context.Context, req *CreateMessageRequest) (*CreateMessageResult, error) {
			receivedMessages = req.Params.Messages
			return &CreateMessageResult{
				Model:   "test-model",
				Role:    "assistant",
				Content: &TextContent{Text: "I see the tool failed"},
			}, nil
		},
		Capabilities: &ClientCapabilities{
			Sampling: &SamplingCapabilities{Tools: &SamplingToolsCapabilities{}},
		},
	})

	server := NewServer(testImpl, nil)
	ss, err := server.Connect(ctx, st, nil)
	if err != nil {
		t.Fatal(err)
	}
	defer ss.Close()

	cs, err := client.Connect(ctx, ct, nil)
	if err != nil {
		t.Fatal(err)
	}
	defer cs.Close()

	// Server sends CreateMessage with error tool result, preceded by
	// the original request and tool use for a more realistic scenario.
	_, err = ss.CreateMessage(ctx, &CreateMessageParams{
		MaxTokens: 1000,
		Messages: []*SamplingMessage{
			{Role: "user", Content: &TextContent{Text: "Divide 1 by 0"}},
			{Role: "assistant", Content: &ToolUseContent{
				ID:    "tool_1",
				Name:  "calculator",
				Input: map[string]any{"op": "div", "x": 1.0, "y": 0.0},
			}},
			{Role: "user", Content: &ToolResultContent{
				ToolUseID: "tool_1",
				Content:   []Content{&TextContent{Text: "division by zero"}},
				IsError:   true,
			}},
		},
	})
	if err != nil {
		t.Fatalf("CreateMessage() error = %v", err)
	}

	if len(receivedMessages) != 3 {
		t.Fatalf("received %d messages, want 3", len(receivedMessages))
	}

	toolResult, ok := receivedMessages[2].Content.(*ToolResultContent)
	if !ok {
		t.Fatalf("content type = %T, want *ToolResultContent", receivedMessages[2].Content)
	}
	if !toolResult.IsError {
		t.Error("IsError should be true")
	}
	if toolResult.ToolUseID != "tool_1" {
		t.Errorf("ToolUseID = %v, want tool_1", toolResult.ToolUseID)
	}
}

func TestParallelToolCalls_Integration(t *testing.T) {
	ctx := context.Background()
	ct, st := NewInMemoryTransports()

	// Client returns parallel tool use results
	client := NewClient(testImpl, &ClientOptions{
		CreateMessageWithToolsHandler: func(_ context.Context, req *CreateMessageWithToolsRequest) (*CreateMessageWithToolsResult, error) {
			return &CreateMessageWithToolsResult{
				Model: "test-model",
				Role:  "assistant",
				Content: []Content{
					&ToolUseContent{ID: "call_1", Name: "weather", Input: map[string]any{"city": "SF"}},
					&ToolUseContent{ID: "call_2", Name: "weather", Input: map[string]any{"city": "NY"}},
				},
				StopReason: "toolUse",
			}, nil
		},
		Capabilities: &ClientCapabilities{
			Sampling: &SamplingCapabilities{Tools: &SamplingToolsCapabilities{}},
		},
	})

	server := NewServer(testImpl, nil)
	ss, err := server.Connect(ctx, st, nil)
	if err != nil {
		t.Fatal(err)
	}
	defer ss.Close()

	cs, err := client.Connect(ctx, ct, nil)
	if err != nil {
		t.Fatal(err)
	}
	defer cs.Close()

	result, err := ss.CreateMessageWithTools(ctx, &CreateMessageWithToolsParams{
		MaxTokens: 1000,
		Messages: []*SamplingMessageV2{
			{Role: "user", Content: []Content{&TextContent{Text: "Weather in SF and NY"}}},
		},
		Tools: []*Tool{
			{Name: "weather", InputSchema: map[string]any{"type": "object"}},
		},
	})
	if err != nil {
		t.Fatalf("CreateMessageWithTools() error = %v", err)
	}

	if len(result.Content) != 2 {
		t.Fatalf("len(Content) = %d, want 2", len(result.Content))
	}
	for i, c := range result.Content {
		tu, ok := c.(*ToolUseContent)
		if !ok {
			t.Fatalf("Content[%d] type = %T, want *ToolUseContent", i, c)
		}
		if tu.Name != "weather" {
			t.Errorf("Content[%d].Name = %v, want weather", i, tu.Name)
		}
	}
	if result.Content[0].(*ToolUseContent).ID != "call_1" {
		t.Errorf("Content[0].ID = %v, want call_1", result.Content[0].(*ToolUseContent).ID)
	}
	if result.Content[1].(*ToolUseContent).ID != "call_2" {
		t.Errorf("Content[1].ID = %v, want call_2", result.Content[1].(*ToolUseContent).ID)
	}
}

func TestNewClient_BothHandlersPanics(t *testing.T) {
	defer func() {
		r := recover()
		if r == nil {
			t.Fatal("expected panic when both handlers set")
		}
		msg, ok := r.(string)
		if !ok || !strings.Contains(msg, "CreateMessageHandler") {
			t.Errorf("unexpected panic: %v", r)
		}
	}()
	NewClient(testImpl, &ClientOptions{
		CreateMessageHandler: func(context.Context, *CreateMessageRequest) (*CreateMessageResult, error) {
			return nil, nil
		},
		CreateMessageWithToolsHandler: func(context.Context, *CreateMessageWithToolsRequest) (*CreateMessageWithToolsResult, error) {
			return nil, nil
		},
	})
}

func TestCreateMessage_MultipleContentError(t *testing.T) {
	ctx := context.Background()
	ct, st := NewInMemoryTransports()

	// Client returns multiple content blocks via CreateMessageWithToolsHandler
	client := NewClient(testImpl, &ClientOptions{
		CreateMessageWithToolsHandler: func(_ context.Context, _ *CreateMessageWithToolsRequest) (*CreateMessageWithToolsResult, error) {
			return &CreateMessageWithToolsResult{
				Model: "test",
				Role:  "assistant",
				Content: []Content{
					&TextContent{Text: "a"},
					&TextContent{Text: "b"},
				},
			}, nil
		},
	})

	server := NewServer(testImpl, nil)
	ss, err := server.Connect(ctx, st, nil)
	if err != nil {
		t.Fatal(err)
	}
	defer ss.Close()

	cs, err := client.Connect(ctx, ct, nil)
	if err != nil {
		t.Fatal(err)
	}
	defer cs.Close()

	// Server calls CreateMessage (singular), should get error
	_, err = ss.CreateMessage(ctx, &CreateMessageParams{
		MaxTokens: 100,
		Messages:  []*SamplingMessage{{Role: "user", Content: &TextContent{Text: "hi"}}},
	})
	if err == nil {
		t.Fatal("expected error for multiple content blocks")
	}
	if !strings.Contains(err.Error(), "CreateMessageWithTools") {
		t.Errorf("error should mention CreateMessageWithTools, got: %v", err)
	}
}

func TestClientCapabilities_CloneSampling(t *testing.T) {
	caps := &ClientCapabilities{
		Sampling: &SamplingCapabilities{
			Tools:   &SamplingToolsCapabilities{},
			Context: &SamplingContextCapabilities{},
		},
	}
	cloned := caps.clone()

	// Verify deep copy — Sampling pointer should differ.
	// (Tools and Context are empty structs, so Go may reuse the same address;
	// we just check they're non-nil and that mutating Sampling doesn't alias.)
	if cloned.Sampling == caps.Sampling {
		t.Error("Sampling pointer should differ after clone")
	}
	if cloned.Sampling.Tools == nil {
		t.Error("cloned Sampling.Tools should not be nil")
	}
	if cloned.Sampling.Context == nil {
		t.Error("cloned Sampling.Context should not be nil")
	}
	// Verify mutation doesn't affect original.
	cloned.Sampling.Tools = nil
	if caps.Sampling.Tools == nil {
		t.Error("modifying cloned Sampling.Tools should not affect original")
	}
}