Cyrene_v2/internal/capability/mcp/client.go

package mcp

import (
	"context"
	"fmt"
	"log/slog"

	"git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk"
)

// initializeParams is the MCP initialize request body.
type initializeParams struct {
	ProtocolVersion string         `json:"protocolVersion"`
	Capabilities    map[string]any `json:"capabilities"`
	ClientInfo      clientInfo     `json:"clientInfo"`
}

// initializeResult is the MCP initialize response body.
type initializeResult struct {
	ProtocolVersion string         `json:"protocolVersion"`
	Capabilities    map[string]any `json:"capabilities"`
	ServerInfo      serverInfo     `json:"serverInfo"`
}

// clientInfo identifies the client to the server.
type clientInfo struct {
	Name    string `json:"name"`
	Version string `json:"version"`
}

// serverInfo identifies the server to the client.
type serverInfo struct {
	Name    string `json:"name"`
	Version string `json:"version"`
}

// Connect selects the transport (unless one was injected), starts it, completes
// the MCP initialize handshake, and marks the client connected. It is
// idempotent; a second call while connected returns nil.
func (c *Client) Connect(ctx context.Context) error {
	c.mu.Lock()
	switch c.state {
	case stateClosed:
		c.mu.Unlock()
		return fmt.Errorf("mcp connect %q: %w", c.cfg.Name, ErrClosed)
	case stateConnected:
		c.mu.Unlock()
		return nil
	}
	c.mu.Unlock()

	tr := c.tr
	if tr == nil {
		var err error
		if tr, err = newTransport(c.cfg, c.log); err != nil {
			return fmt.Errorf("mcp connect %q: %w", c.cfg.Name, err)
		}
	}

	if err := tr.Start(ctx); err != nil {
		tr.Close()
		return fmt.Errorf("mcp connect %q: %w", c.cfg.Name, err)
	}
	info, err := c.initialize(ctx, tr)
	if err != nil {
		tr.Close()
		return fmt.Errorf("mcp connect %q: %w", c.cfg.Name, err)
	}
	c.sendInitialized(ctx, tr)

	c.mu.Lock()
	c.tr = tr
	c.info = info
	c.state = stateConnected
	c.mu.Unlock()
	c.log.Debug("mcp connected", "name", c.cfg.Name, "server", info.Name)
	return nil
}

// Initialize performs the MCP initialize handshake on an already-connected
// client. It returns the server identity.
func (c *Client) Initialize(ctx context.Context) (ServerInfo, error) {
	tr := c.transport()
	if tr == nil {
		return ServerInfo{}, fmt.Errorf("mcp initialize %q: %w", c.cfg.Name, ErrNotConnected)
	}
	return c.initialize(ctx, tr)
}

// initialize sends the initialize request over tr and records the server
// identity.
func (c *Client) initialize(ctx context.Context, tr Transport) (ServerInfo, error) {
	params := initializeParams{
		ProtocolVersion: protocolVersion,
		Capabilities:    map[string]any{},
		ClientInfo:      clientInfo{Name: "cyrene", Version: "0.1.0"},
	}
	var res initializeResult
	if err := tr.Call(ctx, "initialize", params, &res); err != nil {
		return ServerInfo{}, err
	}
	return ServerInfo{Name: res.ServerInfo.Name, Version: res.ServerInfo.Version}, nil
}

// sendInitialized best-effort sends the notifications/initialized advisory ack.
func (c *Client) sendInitialized(ctx context.Context, tr Transport) {
	n, ok := tr.(notifier)
	if !ok {
		return
	}
	if err := n.notify(ctx, "notifications/initialized", map[string]any{}); err != nil {
		c.log.Debug("mcp initialized notification failed", "name", c.cfg.Name, "err", err)
	}
}

// ListTools returns the tools exposed by the connected MCP server.
func (c *Client) ListTools(ctx context.Context) ([]ToolInfo, error) {
	tr := c.transport()
	if tr == nil {
		return nil, fmt.Errorf("mcp list tools %q: %w", c.cfg.Name, ErrNotConnected)
	}
	var res listToolsResult
	if err := tr.Call(ctx, "tools/list", map[string]any{}, &res); err != nil {
		return nil, fmt.Errorf("mcp list tools %q: %w", c.cfg.Name, err)
	}
	out := make([]ToolInfo, 0, len(res.Tools))
	for _, t := range res.Tools {
		out = append(out, ToolInfo{Name: t.Name, Description: t.Description, InputSchema: t.InputSchema})
	}
	return out, nil
}

// CallTool invokes a tool on the connected MCP server.
func (c *Client) CallTool(ctx context.Context, name string, args any) (ToolResult, error) {
	tr := c.transport()
	if tr == nil {
		return ToolResult{}, fmt.Errorf("mcp call tool %q: %w", name, ErrNotConnected)
	}
	arguments, err := normalizeArgs(args)
	if err != nil {
		return ToolResult{}, fmt.Errorf("mcp call tool %q: %w", name, err)
	}
	params := callToolParams{Name: name, Arguments: arguments}
	var res callToolResult
	if err := tr.Call(ctx, "tools/call", params, &res); err != nil {
		return ToolResult{}, fmt.Errorf("mcp call tool %q: %w", name, err)
	}
	return ToolResult{Content: toContentBlocks(res.Content), IsError: res.IsError}, nil
}

// DiscoverTools lists the server tools and maps each to an sdk.Tool so the MCP
// server appears as a normal tool provider.
func (c *Client) DiscoverTools(ctx context.Context) ([]sdk.Tool, error) {
	tools, err := c.ListTools(ctx)
	if err != nil {
		return nil, err
	}
	out := make([]sdk.Tool, 0, len(tools))
	for _, t := range tools {
		out = append(out, &mcpTool{client: c, info: t})
	}
	return out, nil
}

// Close closes the transport and marks the client closed. It is idempotent.
func (c *Client) Close() error {
	c.mu.Lock()
	if c.state == stateClosed {
		c.mu.Unlock()
		return nil
	}
	c.state = stateClosed
	tr := c.tr
	c.mu.Unlock()
	if tr != nil {
		return tr.Close()
	}
	return nil
}

// newTransport selects the transport matching cfg: stdio when Command is set, SSE
// when URL is set.
func newTransport(cfg MCPConfig, log *slog.Logger) (Transport, error) {
	switch {
	case cfg.Command != "":
		return newStdioTransport(cfg, log), nil
	case cfg.URL != "":
		return newSSETransport(cfg, log), nil
	default:
		return nil, fmt.Errorf("mcp: %w", ErrNoTransport)
	}
}

// normalizeArgs converts an agent-supplied argument value into the MCP arguments
// object, requiring a JSON object.
func normalizeArgs(args any) (map[string]any, error) {
	if args == nil {
		return map[string]any{}, nil
	}
	if m, ok := args.(map[string]any); ok {
		return m, nil
	}
	return nil, fmt.Errorf("arguments must be a JSON object, got %T", args)
}

// listToolsResult is the MCP tools/list response body.
type listToolsResult struct {
	Tools []toolSpec `json:"tools"`
}

// toolSpec is a single tool in a tools/list response.
type toolSpec struct {
	Name        string         `json:"name"`
	Description string         `json:"description"`
	InputSchema map[string]any `json:"inputSchema"`
}

// callToolParams is the MCP tools/call request body.
type callToolParams struct {
	Name      string `json:"name"`
	Arguments any    `json:"arguments,omitempty"`
}

// callToolResult is the MCP tools/call response body.
type callToolResult struct {
	Content []contentBlock `json:"content"`
	IsError bool           `json:"isError"`
}

// contentBlock is a wire content block in a tools/call response.
type contentBlock struct {
	Type string `json:"type"`
	Text string `json:"text"`
}