322 lines
8.5 KiB
Go
322 lines
8.5 KiB
Go
package mcptests
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"os"
|
|
"os/exec"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/mark3labs/mcp-go/mcp"
|
|
"github.com/mark3labs/mcp-go/server"
|
|
"github.com/maximhq/bifrost/core/schemas"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
// =============================================================================
|
|
// MOCK STDIO MCP SERVER
|
|
// =============================================================================
|
|
|
|
// STDIOServerManager manages a STDIO MCP server for testing
|
|
type STDIOServerManager struct {
|
|
server *server.MCPServer
|
|
cmd *exec.Cmd
|
|
stdinPipe *os.File
|
|
stdoutPipe *os.File
|
|
isRunning bool
|
|
mu sync.RWMutex
|
|
serverPath string // Path to the compiled server executable
|
|
t *testing.T
|
|
}
|
|
|
|
// NewSTDIOServerManager creates a new STDIO server manager for testing
|
|
func NewSTDIOServerManager(t *testing.T) *STDIOServerManager {
|
|
t.Helper()
|
|
|
|
return &STDIOServerManager{
|
|
t: t,
|
|
}
|
|
}
|
|
|
|
// Start starts the STDIO server
|
|
func (m *STDIOServerManager) Start() error {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
|
|
if m.isRunning {
|
|
return fmt.Errorf("server already running")
|
|
}
|
|
|
|
// Create a new MCP server
|
|
m.server = server.NewMCPServer(
|
|
"test-stdio-server",
|
|
"1.0.0",
|
|
server.WithToolCapabilities(true),
|
|
)
|
|
|
|
// Register test tools
|
|
if err := m.registerTestTools(); err != nil {
|
|
return fmt.Errorf("failed to register tools: %w", err)
|
|
}
|
|
|
|
m.isRunning = true
|
|
return nil
|
|
}
|
|
|
|
// Stop stops the STDIO server
|
|
func (m *STDIOServerManager) Stop() error {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
|
|
if !m.isRunning {
|
|
return nil
|
|
}
|
|
|
|
if m.cmd != nil && m.cmd.Process != nil {
|
|
if err := m.cmd.Process.Kill(); err != nil {
|
|
return fmt.Errorf("failed to kill process: %w", err)
|
|
}
|
|
}
|
|
|
|
m.isRunning = false
|
|
return nil
|
|
}
|
|
|
|
// IsRunning returns whether the server is running
|
|
func (m *STDIOServerManager) IsRunning() bool {
|
|
m.mu.RLock()
|
|
defer m.mu.RUnlock()
|
|
return m.isRunning
|
|
}
|
|
|
|
// registerTestTools registers test tools on the server
|
|
func (m *STDIOServerManager) registerTestTools() error {
|
|
// Calculator tool
|
|
calculatorTool := mcp.NewTool("calculator",
|
|
mcp.WithDescription("Performs basic arithmetic operations"),
|
|
mcp.WithString("operation",
|
|
mcp.Required(),
|
|
mcp.Description("The operation to perform: add, subtract, multiply, divide"),
|
|
mcp.Enum("add", "subtract", "multiply", "divide"),
|
|
),
|
|
mcp.WithNumber("x",
|
|
mcp.Required(),
|
|
mcp.Description("First number"),
|
|
),
|
|
mcp.WithNumber("y",
|
|
mcp.Required(),
|
|
mcp.Description("Second number"),
|
|
),
|
|
)
|
|
|
|
m.server.AddTool(calculatorTool, m.handleCalculator)
|
|
|
|
// Echo tool
|
|
echoTool := mcp.NewTool("echo",
|
|
mcp.WithDescription("Echoes back the input message"),
|
|
mcp.WithString("message",
|
|
mcp.Required(),
|
|
mcp.Description("The message to echo"),
|
|
),
|
|
)
|
|
|
|
m.server.AddTool(echoTool, m.handleEcho)
|
|
|
|
// Weather tool (for testing external data)
|
|
weatherTool := mcp.NewTool("get_weather",
|
|
mcp.WithDescription("Gets the weather for a location"),
|
|
mcp.WithString("location",
|
|
mcp.Required(),
|
|
mcp.Description("The location to get weather for"),
|
|
),
|
|
mcp.WithString("units",
|
|
mcp.Description("Temperature units: celsius or fahrenheit"),
|
|
mcp.Enum("celsius", "fahrenheit"),
|
|
),
|
|
)
|
|
|
|
m.server.AddTool(weatherTool, m.handleWeather)
|
|
|
|
// Delay tool (for timeout testing)
|
|
delayTool := mcp.NewTool("delay",
|
|
mcp.WithDescription("Delays for a specified duration"),
|
|
mcp.WithNumber("seconds",
|
|
mcp.Required(),
|
|
mcp.Description("Number of seconds to delay"),
|
|
),
|
|
)
|
|
|
|
m.server.AddTool(delayTool, m.handleDelay)
|
|
|
|
// Error tool (for error testing)
|
|
errorTool := mcp.NewTool("throw_error",
|
|
mcp.WithDescription("Throws an error for testing"),
|
|
mcp.WithString("error_message",
|
|
mcp.Required(),
|
|
mcp.Description("The error message to throw"),
|
|
),
|
|
)
|
|
|
|
m.server.AddTool(errorTool, m.handleError)
|
|
|
|
return nil
|
|
}
|
|
|
|
// Tool handlers
|
|
|
|
func (m *STDIOServerManager) handleCalculator(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
|
var args struct {
|
|
Operation string `json:"operation"`
|
|
X float64 `json:"x"`
|
|
Y float64 `json:"y"`
|
|
}
|
|
|
|
argsBytes, ok := request.Params.Arguments.(string)
|
|
if !ok {
|
|
return mcp.NewToolResultError("invalid arguments type"), nil
|
|
}
|
|
if err := json.Unmarshal([]byte(argsBytes), &args); err != nil {
|
|
return mcp.NewToolResultError(fmt.Sprintf("invalid arguments: %v", err)), nil
|
|
}
|
|
|
|
var result float64
|
|
switch args.Operation {
|
|
case "add":
|
|
result = args.X + args.Y
|
|
case "subtract":
|
|
result = args.X - args.Y
|
|
case "multiply":
|
|
result = args.X * args.Y
|
|
case "divide":
|
|
if args.Y == 0 {
|
|
return mcp.NewToolResultError("division by zero"), nil
|
|
}
|
|
result = args.X / args.Y
|
|
default:
|
|
return mcp.NewToolResultError(fmt.Sprintf("unknown operation: %s", args.Operation)), nil
|
|
}
|
|
|
|
return mcp.NewToolResultText(fmt.Sprintf("%.2f", result)), nil
|
|
}
|
|
|
|
func (m *STDIOServerManager) handleEcho(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
|
var args struct {
|
|
Message string `json:"message"`
|
|
}
|
|
|
|
argsBytes, ok := request.Params.Arguments.(string)
|
|
if !ok {
|
|
return mcp.NewToolResultError("invalid arguments type"), nil
|
|
}
|
|
if err := json.Unmarshal([]byte(argsBytes), &args); err != nil {
|
|
return mcp.NewToolResultError(fmt.Sprintf("invalid arguments: %v", err)), nil
|
|
}
|
|
|
|
return mcp.NewToolResultText(args.Message), nil
|
|
}
|
|
|
|
func (m *STDIOServerManager) handleWeather(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
|
var args struct {
|
|
Location string `json:"location"`
|
|
Units string `json:"units"`
|
|
}
|
|
|
|
argsBytes, ok := request.Params.Arguments.(string)
|
|
if !ok {
|
|
return mcp.NewToolResultError("invalid arguments type"), nil
|
|
}
|
|
if err := json.Unmarshal([]byte(argsBytes), &args); err != nil {
|
|
return mcp.NewToolResultError(fmt.Sprintf("invalid arguments: %v", err)), nil
|
|
}
|
|
|
|
if args.Units == "" {
|
|
args.Units = "celsius"
|
|
}
|
|
|
|
// Mock weather response
|
|
temp := "22"
|
|
if args.Units == "fahrenheit" {
|
|
temp = "72"
|
|
}
|
|
|
|
response := fmt.Sprintf("The weather in %s is sunny with a temperature of %s°%s",
|
|
args.Location, temp, args.Units)
|
|
|
|
return mcp.NewToolResultText(response), nil
|
|
}
|
|
|
|
func (m *STDIOServerManager) handleDelay(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
|
var args struct {
|
|
Seconds float64 `json:"seconds"`
|
|
}
|
|
|
|
argsBytes, ok := request.Params.Arguments.(string)
|
|
if !ok {
|
|
return mcp.NewToolResultError("invalid arguments type"), nil
|
|
}
|
|
if err := json.Unmarshal([]byte(argsBytes), &args); err != nil {
|
|
return mcp.NewToolResultError(fmt.Sprintf("invalid arguments: %v", err)), nil
|
|
}
|
|
|
|
duration := time.Duration(args.Seconds * float64(time.Second))
|
|
select {
|
|
case <-time.After(duration):
|
|
return mcp.NewToolResultText(fmt.Sprintf("Delayed for %.2f seconds", args.Seconds)), nil
|
|
case <-ctx.Done():
|
|
return mcp.NewToolResultError("delay cancelled"), nil
|
|
}
|
|
}
|
|
|
|
func (m *STDIOServerManager) handleError(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
|
var args struct {
|
|
ErrorMessage string `json:"error_message"`
|
|
}
|
|
|
|
argsBytes, ok := request.Params.Arguments.(string)
|
|
if !ok {
|
|
return mcp.NewToolResultError("invalid arguments type"), nil
|
|
}
|
|
if err := json.Unmarshal([]byte(argsBytes), &args); err != nil {
|
|
return mcp.NewToolResultError(fmt.Sprintf("invalid arguments: %v", err)), nil
|
|
}
|
|
|
|
return mcp.NewToolResultError(args.ErrorMessage), nil
|
|
}
|
|
|
|
// GetServerExecutablePath returns the path where the STDIO server will be compiled
|
|
func (m *STDIOServerManager) GetServerExecutablePath() string {
|
|
m.mu.RLock()
|
|
defer m.mu.RUnlock()
|
|
return m.serverPath
|
|
}
|
|
|
|
// =============================================================================
|
|
// HELPER FUNCTIONS
|
|
// =============================================================================
|
|
|
|
// createTestContext creates a BifrostContext for testing
|
|
func createTestContext() *schemas.BifrostContext {
|
|
return schemas.NewBifrostContext(context.Background(), schemas.NoDeadline)
|
|
}
|
|
|
|
// createTestContextWithTimeout creates a BifrostContext with timeout
|
|
func createTestContextWithTimeout(timeout time.Duration) (*schemas.BifrostContext, context.CancelFunc) {
|
|
ctx, cancel := context.WithTimeout(context.Background(), timeout)
|
|
return schemas.NewBifrostContext(ctx, schemas.NoDeadline), cancel
|
|
}
|
|
|
|
// assertNoError asserts that error is nil
|
|
func assertNoError(t *testing.T, err error, msgAndArgs ...interface{}) {
|
|
t.Helper()
|
|
require.NoError(t, err, msgAndArgs...)
|
|
}
|
|
|
|
// assertError asserts that error is not nil
|
|
func assertError(t *testing.T, err error, msgAndArgs ...interface{}) {
|
|
t.Helper()
|
|
require.Error(t, err, msgAndArgs...)
|
|
}
|