first commit
This commit is contained in:
@@ -0,0 +1,110 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os/exec"
|
||||
|
||||
sdkmcp "github.com/modelcontextprotocol/go-sdk/mcp"
|
||||
)
|
||||
|
||||
type MCPClient struct {
|
||||
name string
|
||||
client *sdkmcp.Client
|
||||
session *sdkmcp.ClientSession
|
||||
cmd *exec.Cmd
|
||||
}
|
||||
|
||||
func Connect(ctx context.Context, name, command string, args []string, env []string, transport, url string) (*MCPClient, error) {
|
||||
client := sdkmcp.NewClient(
|
||||
&sdkmcp.Implementation{Name: "ai-agent", Version: "0.2.0"},
|
||||
nil,
|
||||
)
|
||||
var t sdkmcp.Transport
|
||||
switch transport {
|
||||
case "sse":
|
||||
if url == "" {
|
||||
return nil, fmt.Errorf("sse transport requires url for %s", name)
|
||||
}
|
||||
t = &sdkmcp.SSEClientTransport{Endpoint: url}
|
||||
case "streamable-http":
|
||||
if url == "" {
|
||||
return nil, fmt.Errorf("streamable-http transport requires url for %s", name)
|
||||
}
|
||||
t = &sdkmcp.StreamableClientTransport{Endpoint: url}
|
||||
default:
|
||||
if command == "" {
|
||||
return nil, fmt.Errorf("stdio transport requires command for %s", name)
|
||||
}
|
||||
cmd := exec.Command(command, args...)
|
||||
if len(env) > 0 {
|
||||
cmd.Env = append(cmd.Environ(), env...)
|
||||
}
|
||||
t = &sdkmcp.CommandTransport{Command: cmd}
|
||||
}
|
||||
session, err := client.Connect(ctx, t, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("connect to %s: %w", name, err)
|
||||
}
|
||||
return &MCPClient{
|
||||
name: name,
|
||||
client: client,
|
||||
session: session,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (c *MCPClient) Name() string { return c.name }
|
||||
|
||||
func (c *MCPClient) ListTools(ctx context.Context) ([]*sdkmcp.Tool, error) {
|
||||
caps := c.session.InitializeResult()
|
||||
if caps == nil || caps.Capabilities.Tools == nil {
|
||||
return nil, nil
|
||||
}
|
||||
var tools []*sdkmcp.Tool
|
||||
for tool, err := range c.session.Tools(ctx, nil) {
|
||||
if err != nil {
|
||||
return tools, fmt.Errorf("list tools from %s: %w", c.name, err)
|
||||
}
|
||||
tools = append(tools, tool)
|
||||
}
|
||||
return tools, nil
|
||||
}
|
||||
|
||||
func (c *MCPClient) CallTool(ctx context.Context, name string, args map[string]any) (*ToolResult, error) {
|
||||
result, err := c.session.CallTool(ctx, &sdkmcp.CallToolParams{
|
||||
Name: name,
|
||||
Arguments: args,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("call tool %s on %s: %w", name, c.name, err)
|
||||
}
|
||||
var text string
|
||||
for _, ct := range result.Content {
|
||||
if tc, ok := ct.(*sdkmcp.TextContent); ok {
|
||||
if text != "" {
|
||||
text += "\n"
|
||||
}
|
||||
text += tc.Text
|
||||
}
|
||||
}
|
||||
return &ToolResult{Content: text, IsError: result.IsError}, nil
|
||||
}
|
||||
|
||||
func (c *MCPClient) Close() error {
|
||||
if c.session != nil {
|
||||
return c.session.Close()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *MCPClient) IsConnected() bool {
|
||||
return c.session != nil
|
||||
}
|
||||
|
||||
func (c *MCPClient) Ping(ctx context.Context) error {
|
||||
if c.session == nil {
|
||||
return fmt.Errorf("no session")
|
||||
}
|
||||
_, err := c.ListTools(ctx)
|
||||
return err
|
||||
}
|
||||
@@ -0,0 +1,241 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"ai-agent/internal/config"
|
||||
"ai-agent/internal/llm"
|
||||
)
|
||||
|
||||
type FailedServer struct {
|
||||
Name string
|
||||
Reason string
|
||||
}
|
||||
|
||||
type ServerStatus struct {
|
||||
Name string
|
||||
Connected bool
|
||||
LastError string
|
||||
LastPing time.Time
|
||||
}
|
||||
|
||||
type Registry struct {
|
||||
mu sync.RWMutex
|
||||
clients []*MCPClient
|
||||
toolMap map[string]*MCPClient
|
||||
toolDefs []llm.ToolDef
|
||||
failedServers []FailedServer
|
||||
serverConfigs map[string]config.ServerConfig
|
||||
}
|
||||
|
||||
func NewRegistry() *Registry {
|
||||
return &Registry{toolMap: make(map[string]*MCPClient), serverConfigs: make(map[string]config.ServerConfig)}
|
||||
}
|
||||
|
||||
const connectTimeout = 5 * time.Second
|
||||
|
||||
func (r *Registry) ConnectServer(ctx context.Context, srv config.ServerConfig) (int, error) {
|
||||
connCtx, cancel := context.WithTimeout(ctx, connectTimeout)
|
||||
defer cancel()
|
||||
client, err := Connect(connCtx, srv.Name, srv.Command, srv.Args, srv.Env, srv.Transport, srv.URL)
|
||||
if err != nil {
|
||||
r.mu.Lock()
|
||||
r.failedServers = append(r.failedServers, FailedServer{Name: srv.Name, Reason: err.Error()})
|
||||
r.mu.Unlock()
|
||||
return 0, fmt.Errorf("connect to %s: %w", srv.Name, err)
|
||||
}
|
||||
tools, err := client.ListTools(connCtx)
|
||||
if err != nil {
|
||||
client.Close()
|
||||
r.mu.Lock()
|
||||
r.failedServers = append(r.failedServers, FailedServer{Name: srv.Name, Reason: err.Error()})
|
||||
r.mu.Unlock()
|
||||
return 0, fmt.Errorf("%s tools: %w", srv.Name, err)
|
||||
}
|
||||
r.mu.Lock()
|
||||
r.clients = append(r.clients, client)
|
||||
for _, tool := range tools {
|
||||
r.toolMap[tool.Name] = client
|
||||
r.toolDefs = append(r.toolDefs, ToLLMToolDef(tool.Name, tool.Description, tool.InputSchema))
|
||||
}
|
||||
r.serverConfigs[srv.Name] = srv
|
||||
r.mu.Unlock()
|
||||
return len(tools), nil
|
||||
}
|
||||
|
||||
func (r *Registry) ConnectAll(ctx context.Context, servers []config.ServerConfig, logFn func(string)) {
|
||||
for _, srv := range servers {
|
||||
toolCount, err := r.ConnectServer(ctx, srv)
|
||||
if err != nil {
|
||||
logFn(fmt.Sprintf("skip %s: %v", srv.Name, err))
|
||||
continue
|
||||
}
|
||||
logFn(fmt.Sprintf("connected %s (%d tools)", srv.Name, toolCount))
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Registry) Tools() []llm.ToolDef {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
return r.toolDefs
|
||||
}
|
||||
|
||||
func (r *Registry) ToolCount() int {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
return len(r.toolDefs)
|
||||
}
|
||||
|
||||
func (r *Registry) ServerCount() int {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
return len(r.clients)
|
||||
}
|
||||
|
||||
func (r *Registry) ServerNames() []string {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
names := make([]string, len(r.clients))
|
||||
for i, c := range r.clients {
|
||||
names[i] = c.Name()
|
||||
}
|
||||
return names
|
||||
}
|
||||
|
||||
func (r *Registry) FailedServers() []FailedServer {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
return r.failedServers
|
||||
}
|
||||
|
||||
func (r *Registry) CallTool(ctx context.Context, name string, args map[string]any) (*ToolResult, error) {
|
||||
r.mu.RLock()
|
||||
client, ok := r.toolMap[name]
|
||||
r.mu.RUnlock()
|
||||
if !ok {
|
||||
return &ToolResult{
|
||||
Content: fmt.Sprintf("unknown tool: %s", name),
|
||||
IsError: true,
|
||||
}, nil
|
||||
}
|
||||
return client.CallTool(ctx, name, args)
|
||||
}
|
||||
|
||||
func (r *Registry) Close() {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
for _, c := range r.clients {
|
||||
c.Close()
|
||||
}
|
||||
r.clients = nil
|
||||
r.toolMap = make(map[string]*MCPClient)
|
||||
r.toolDefs = nil
|
||||
}
|
||||
|
||||
func (r *Registry) HealthCheck(ctx context.Context) []ServerStatus {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
var results []ServerStatus
|
||||
for _, client := range r.clients {
|
||||
status := ServerStatus{Name: client.Name()}
|
||||
if client.IsConnected() {
|
||||
pingCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
err := client.Ping(pingCtx)
|
||||
cancel()
|
||||
status.Connected = err == nil
|
||||
if err != nil {
|
||||
status.LastError = err.Error()
|
||||
}
|
||||
status.LastPing = time.Now()
|
||||
}
|
||||
results = append(results, status)
|
||||
}
|
||||
for _, failed := range r.failedServers {
|
||||
results = append(results, ServerStatus{
|
||||
Name: failed.Name,
|
||||
Connected: false,
|
||||
LastError: failed.Reason,
|
||||
})
|
||||
}
|
||||
return results
|
||||
}
|
||||
|
||||
func (r *Registry) ReconnectServer(ctx context.Context, name string) (int, error) {
|
||||
r.mu.RLock()
|
||||
srv, ok := r.serverConfigs[name]
|
||||
r.mu.RUnlock()
|
||||
if !ok {
|
||||
return 0, fmt.Errorf("no config found for server: %s", name)
|
||||
}
|
||||
r.mu.Lock()
|
||||
var remainingFailed []FailedServer
|
||||
for _, f := range r.failedServers {
|
||||
if f.Name != name {
|
||||
remainingFailed = append(remainingFailed, f)
|
||||
}
|
||||
}
|
||||
r.failedServers = remainingFailed
|
||||
r.mu.Unlock()
|
||||
return r.ConnectServer(ctx, srv)
|
||||
}
|
||||
|
||||
type MonitorConfig struct {
|
||||
Interval time.Duration
|
||||
MaxRetries int
|
||||
BackoffBase time.Duration
|
||||
}
|
||||
|
||||
var defaultMonitorConfig = MonitorConfig{
|
||||
Interval: 30 * time.Second,
|
||||
MaxRetries: 3,
|
||||
BackoffBase: 5 * time.Second,
|
||||
}
|
||||
|
||||
func (r *Registry) StartHealthMonitor(ctx context.Context, cfg MonitorConfig, logFn func(string)) context.CancelFunc {
|
||||
if cfg.Interval == 0 {
|
||||
cfg = defaultMonitorConfig
|
||||
}
|
||||
monitorCtx, cancel := context.WithCancel(ctx)
|
||||
go func() {
|
||||
ticker := time.NewTicker(cfg.Interval)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-monitorCtx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
r.healthCheckRound(monitorCtx, cfg, logFn)
|
||||
}
|
||||
}
|
||||
}()
|
||||
return cancel
|
||||
}
|
||||
|
||||
func (r *Registry) healthCheckRound(ctx context.Context, cfg MonitorConfig, logFn func(string)) {
|
||||
statuses := r.HealthCheck(ctx)
|
||||
for _, status := range statuses {
|
||||
if status.Connected {
|
||||
continue
|
||||
}
|
||||
logFn(fmt.Sprintf("server %s unhealthy, attempting reconnect...", status.Name))
|
||||
for attempt := 1; attempt <= cfg.MaxRetries; attempt++ {
|
||||
backoff := cfg.BackoffBase * time.Duration(attempt)
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-time.After(backoff):
|
||||
}
|
||||
_, err := r.ReconnectServer(ctx, status.Name)
|
||||
if err == nil {
|
||||
logFn(fmt.Sprintf("server %s reconnected", status.Name))
|
||||
break
|
||||
}
|
||||
if attempt == cfg.MaxRetries {
|
||||
logFn(fmt.Sprintf("server %s reconnection failed after %d attempts: %v", status.Name, cfg.MaxRetries, err))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestNewRegistry(t *testing.T) {
|
||||
r := NewRegistry()
|
||||
|
||||
if r.ToolCount() != 0 {
|
||||
t.Errorf("ToolCount() = %d, want 0", r.ToolCount())
|
||||
}
|
||||
if r.ServerCount() != 0 {
|
||||
t.Errorf("ServerCount() = %d, want 0", r.ServerCount())
|
||||
}
|
||||
if tools := r.Tools(); len(tools) != 0 {
|
||||
t.Errorf("Tools() = %v, want empty", tools)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistry_CallTool_Unknown(t *testing.T) {
|
||||
r := NewRegistry()
|
||||
|
||||
result, err := r.CallTool(context.Background(), "nonexistent_tool", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("CallTool() unexpected error: %v", err)
|
||||
}
|
||||
if !result.IsError {
|
||||
t.Error("CallTool() IsError = false, want true for unknown tool")
|
||||
}
|
||||
if !strings.Contains(result.Content, "unknown tool") {
|
||||
t.Errorf("CallTool() Content = %q, want to contain 'unknown tool'", result.Content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistry_HealthCheck_Empty(t *testing.T) {
|
||||
r := NewRegistry()
|
||||
|
||||
statuses := r.HealthCheck(context.Background())
|
||||
if len(statuses) != 0 {
|
||||
t.Errorf("HealthCheck() returned %d statuses, want 0", len(statuses))
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistry_HealthCheck_TracksFailedServers(t *testing.T) {
|
||||
r := NewRegistry()
|
||||
|
||||
// Simulate a failed server by directly adding to failedServers
|
||||
r.mu.Lock()
|
||||
r.failedServers = append(r.failedServers, FailedServer{
|
||||
Name: "failed-server",
|
||||
Reason: "connection refused",
|
||||
})
|
||||
r.mu.Unlock()
|
||||
|
||||
statuses := r.HealthCheck(context.Background())
|
||||
if len(statuses) != 1 {
|
||||
t.Fatalf("HealthCheck() returned %d statuses, want 1", len(statuses))
|
||||
}
|
||||
|
||||
status := statuses[0]
|
||||
if status.Name != "failed-server" {
|
||||
t.Errorf("status.Name = %q, want 'failed-server'", status.Name)
|
||||
}
|
||||
if status.Connected {
|
||||
t.Error("status.Connected = true, want false")
|
||||
}
|
||||
if status.LastError != "connection refused" {
|
||||
t.Errorf("status.LastError = %q, want 'connection refused'", status.LastError)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"ai-agent/internal/llm"
|
||||
)
|
||||
|
||||
type ServerInfo struct {
|
||||
Name string
|
||||
ToolCount int
|
||||
}
|
||||
|
||||
type ToolResult struct {
|
||||
Content string
|
||||
IsError bool
|
||||
}
|
||||
|
||||
func ToLLMToolDef(name, description string, inputSchema any) llm.ToolDef {
|
||||
params, _ := inputSchema.(map[string]any)
|
||||
if params == nil {
|
||||
params = map[string]any{"type": "object", "properties": map[string]any{}}
|
||||
}
|
||||
return llm.ToolDef{
|
||||
Name: name,
|
||||
Description: description,
|
||||
Parameters: params,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,71 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestToLLMToolDef(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
toolName string
|
||||
description string
|
||||
inputSchema any
|
||||
wantName string
|
||||
wantDesc string
|
||||
wantParams bool // true = should have non-nil params
|
||||
}{
|
||||
{
|
||||
name: "normal with valid schema",
|
||||
toolName: "read_file",
|
||||
description: "Read a file",
|
||||
inputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"path": map[string]any{"type": "string"},
|
||||
},
|
||||
},
|
||||
wantName: "read_file",
|
||||
wantDesc: "Read a file",
|
||||
wantParams: true,
|
||||
},
|
||||
{
|
||||
name: "nil schema uses default",
|
||||
toolName: "noop",
|
||||
description: "No-op tool",
|
||||
inputSchema: nil,
|
||||
wantName: "noop",
|
||||
wantDesc: "No-op tool",
|
||||
wantParams: true,
|
||||
},
|
||||
{
|
||||
name: "non-map schema uses default",
|
||||
toolName: "bad_schema",
|
||||
description: "Bad schema tool",
|
||||
inputSchema: "not a map",
|
||||
wantName: "bad_schema",
|
||||
wantDesc: "Bad schema tool",
|
||||
wantParams: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := ToLLMToolDef(tt.toolName, tt.description, tt.inputSchema)
|
||||
if result.Name != tt.wantName {
|
||||
t.Errorf("Name = %q, want %q", result.Name, tt.wantName)
|
||||
}
|
||||
if result.Description != tt.wantDesc {
|
||||
t.Errorf("Description = %q, want %q", result.Description, tt.wantDesc)
|
||||
}
|
||||
if tt.wantParams && result.Parameters == nil {
|
||||
t.Error("Parameters should not be nil")
|
||||
}
|
||||
// Nil and non-map schemas should get the default object schema.
|
||||
if tt.inputSchema == nil || func() bool { _, ok := tt.inputSchema.(map[string]any); return !ok }() {
|
||||
if result.Parameters["type"] != "object" {
|
||||
t.Errorf("default schema type = %v, want 'object'", result.Parameters["type"])
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user