set_test.go

Overview

The set_test.go file is a Go test suite focused on validating the integration and functionality of the MCP Toolset Integration within an agent framework. It demonstrates how an MCP (Model Context Protocol) server can be hosted in-memory, how MCP tools are exposed to an agent, and how agents invoke these tools to respond to natural language queries.

This file exercises key flows such as:

The tests use mocks and recorded HTTP round-trippers to simulate LLM model interactions without real network calls.

Key Types and Functions

Input and Output Structs

type Input struct {
    City string `json:"city" jsonschema:"city name"`
}

type Output struct {
    WeatherSummary string `json:"weather_summary" jsonschema:"weather summary in the given city"`
}

weatherFunc

func weatherFunc(ctx context.Context, req *mcp.CallToolRequest, input Input) (*mcp.CallToolResult, Output, error)

TestMCPToolSet

func TestMCPToolSet(t *testing.T)

newGeminiTestClientConfig

func newGeminiTestClientConfig(t *testing.T, rrfile string) (http.RoundTripper, bool)

newGeminiModel

func newGeminiModel(t *testing.T, modelName string) model.LLM

testAgentRunner

A helper struct to encapsulate agent execution within tests.

type testAgentRunner struct {
    agent          agent.Agent
    sessionService session.Service
    lastSession    session.Session
    appName        string
    runner         *runner.Runner
}

TestToolFilter

func TestToolFilter(t *testing.T)

Important Implementation Details

Interaction with Other Parts of the System

Usage Example (from TestMCPToolSet)

// Create in-memory MCP transports
clientTransport, serverTransport := mcp.NewInMemoryTransports()

// Setup MCP server and add tools
server := mcp.NewServer(&mcp.Implementation{Name: "weather_server", Version: "v1.0.0"}, nil)
mcp.AddTool(server, &mcp.Tool{Name: "get_weather", Description: "returns weather in the given city"}, weatherFunc)
_, err := server.Connect(t.Context(), serverTransport, nil)

// Create MCP ToolSet client
ts, err := mcptoolset.New(mcptoolset.Config{
    Transport: clientTransport,
})

// Create LLM agent with MCP toolset
agent, err := llmagent.New(llmagent.Config{
    Name:        "weather_time_agent",
    Model:       newGeminiModel(t, modelName),
    Description: "Agent to answer questions about the time and weather in a city.",
    Instruction: "I can answer your questions about the time and weather in a city.",
    Toolsets:    []tool.Toolset{ ts },
})

// Run agent session with prompt
runner := newTestAgentRunner(t, agent)
events := runner.Run(t, "session1", "what is the weather in london?")

This example demonstrates the setup of an MCP server and client, creation of an LLM agent with MCP tools, and execution of a session that triggers the get_weather tool.

Mermaid Diagram: Structure and Workflow of set_test.go

flowchart TD
A[MCP Server Setup]
B[MCP Tools Registration]
C[MCP In-Memory Transport]
D[MCP Toolset Client]
E[LLM Agent Creation]
F[Agent Runner]
G[Run Agent Session]
H[Collect Session Events]
I[Verify Events]
A --> B
B --> C
C --> D
D --> E
E --> F
F --> G
G --> H
H --> I

References to Related Topics

This file exemplifies the integration of MCP tools into an agent-driven workflow, demonstrating in-memory testing of tool invocation via LLM function calls and validating agent event sequences.