Cyrene_v2/internal/capability/mcp/client.go

package mcp

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"time"

	"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)
	}
	// Wire server→client notification routing (list_changed) before the server
	// can start sending them after notifications/initialized.
	attachNotificationSink(tr, c.handleServerNotification)

	info, err := c.initialize(ctx, tr)
	if err != nil {
		tr.Close()
		return fmt.Errorf("mcp connect %q: %w", c.cfg.Name, err)
	}
	// Let the transport echo the negotiated protocol version (Streamable HTTP).
	attachProtocolProvider(tr, c.protocolVersion)
	c.sendInitialized(ctx, tr)

	c.mu.Lock()
	c.tr = tr
	c.info = info
	c.state = stateConnected
	c.mu.Unlock()

	c.startKeepalive()
	c.log.Debug("mcp connected", "name", c.cfg.Name, "server", info.Name, "protocol", c.protocolVersion())
	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 performs the MCP initialize handshake with protocol version
// negotiation (DevDocs/design/09 §1.1). It does not hard-code a single version:
// it first requests latestProtocolVersion, and if the server rejects that version
// (a JSON-RPC error with code codeUnsupportedProtocolVersion carrying the
// supported set) or replies with a version this client cannot speak, it downgrades
// once to the highest mutually supported version and retries. It records the
// agreed version so the HTTP transport can echo it back as Mcp-Protocol-Version.
func (c *Client) initialize(ctx context.Context, tr Transport) (ServerInfo, error) {
	sent := c.preferredProtocolVersion
	if sent == "" {
		sent = latestProtocolVersion
	}
	var lastErr error
	for attempt := 0; attempt < maxInitializeAttempts; attempt++ {
		params := initializeParams{
			ProtocolVersion: sent,
			Capabilities:    c.capabilities(sent),
			ClientInfo:      clientInfo{Name: "cyrene", Version: "0.1.0"},
		}
		var res initializeResult
		err := tr.Call(ctx, "initialize", params, &res)
		if err != nil {
			// Server rejected the requested version and may advertise the versions
			// it supports; downgrade to the highest mutually supported one.
			if supported, ok := unsupportedVersions(err); ok {
				if v := negotiateMutuallySupportedVersion(supported); v != "" && v != sent {
					c.log.Debug("mcp downgrading protocol version",
						"name", c.cfg.Name, "requested", sent, "version", v)
					sent = v
					lastErr = err
					continue
				}
				return ServerInfo{}, fmt.Errorf("mcp initialize %q: %w: %w", c.cfg.Name, ErrUnsupportedProtocolVersion, err)
			}
			return ServerInfo{}, err
		}
		negotiated := res.ProtocolVersion
		if negotiated == "" {
			negotiated = sent
		}
		if !isSupportedVersion(negotiated) {
			// The server replied with a version this client does not speak;
			// downgrade once to the highest version we both accept.
			if v := highestCommonVersion(negotiated); v != "" && v != sent {
				c.log.Debug("mcp renegotiating protocol version",
					"name", c.cfg.Name, "server", negotiated, "version", v)
				sent = v
				lastErr = err
				continue
			}
			return ServerInfo{}, fmt.Errorf("mcp initialize %q: %w: server advertised %q", c.cfg.Name, ErrUnsupportedProtocolVersion, negotiated)
		}
		c.setProtocolVersion(negotiated)
		return ServerInfo{Name: res.ServerInfo.Name, Version: res.ServerInfo.Version}, nil
	}
	return ServerInfo{}, fmt.Errorf("mcp initialize %q: %w: %w", c.cfg.Name, ErrUnsupportedProtocolVersion, lastErr)
}

// isMethodNotFound reports whether err wraps a JSON-RPC "method not found" error
// (code -32601), which a peer returns when it does not implement a method such as
// the keepalive ping.
func isMethodNotFound(err error) bool {
	var re *rpcError
	return errors.As(err, &re) && re.Code == codeMethodNotFound
}

// unsupportedVersions reports the protocol versions the server supports, extracted
// from the data field of a codeUnsupportedProtocolVersion JSON-RPC error. It
// reports ok=false when err does not carry such an error or the data is unusable.
func unsupportedVersions(err error) ([]string, bool) {
	var re *rpcError
	if !errors.As(err, &re) || re.Code != codeUnsupportedProtocolVersion {
		return nil, false
	}
	if re.Data == nil {
		return nil, false
	}
	data, err := json.Marshal(re.Data)
	if err != nil {
		return nil, false
	}
	var payload struct {
		Supported []string `json:"supported"`
	}
	if err := json.Unmarshal(data, &payload); err != nil || len(payload.Supported) == 0 {
		return nil, false
	}
	return payload.Supported, true
}

// 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)
	}
}

// handleServerNotification is the transport-level sink for server→client
// notifications. It offloads the actual work to a fresh goroutine because it runs
// on the transport read path and must not block or re-enter the same transport
// synchronously; refreshPending coalesces bursts so at most one refresh runs at a
// time.
func (c *Client) handleServerNotification(method string) {
	if method != notificationToolsListChanged || c.onToolListChanged == nil {
		return
	}
	if !c.refreshPending.CompareAndSwap(false, true) {
		return
	}
	go func() {
		defer c.refreshPending.Store(false)
		c.onToolListChanged()
	}()
}

// ping sends a JSON-RPC ping request and reports whether the server responded.
// Pings are used by the keepalive loop to detect a dead connection.
func (c *Client) ping(ctx context.Context) error {
	tr := c.transport()
	if tr == nil {
		return fmt.Errorf("mcp ping: %w", ErrNotConnected)
	}
	if err := tr.Call(ctx, "ping", map[string]any{}, nil); err != nil {
		return fmt.Errorf("mcp ping: %w", err)
	}
	return nil
}

// startKeepalive launches the ping keepalive loop when configured. It closes the
// client once keepAliveFailureThreshold consecutive pings fail, so the host can
// rebuild the session. The loop exits when the client is closed.
func (c *Client) startKeepalive() {
	interval := c.keepAlive
	if interval <= 0 {
		return
	}
	ctx, cancel := context.WithCancel(context.Background())
	c.mu.Lock()
	c.keepaliveCancel = cancel
	c.mu.Unlock()
	threshold := c.keepAliveFailures
	if threshold < 1 {
		threshold = keepAliveFailureThreshold
	}
	go func() {
		ticker := time.NewTicker(interval)
		defer ticker.Stop()
		failures := 0
		for {
			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
				pingCtx, pingCancel := context.WithTimeout(context.Background(), interval/2)
				err := c.ping(pingCtx)
				pingCancel()
				if err != nil {
					// A peer that does not implement ping is still healthy; stop
					// keepalive rather than treating every ping as a failure.
					if isMethodNotFound(err) {
						c.log.Debug("mcp keepalive stopped; server does not support ping", "name", c.cfg.Name)
						return
					}
					failures++
					c.log.Warn("mcp keepalive ping failed", "name", c.cfg.Name, "err", err, "failures", failures)
					if failures < threshold {
						continue
					}
					c.log.Error("mcp keepalive ping failed; closing session", "name", c.cfg.Name, "failures", failures)
					_ = c.Close()
					return
				}
				failures = 0
			}
		}
	}()
}

// 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. It
// also cancels the keepalive loop so no ping goroutine outlives the client.
func (c *Client) Close() error {
	c.mu.Lock()
	if c.state == stateClosed {
		c.mu.Unlock()
		return nil
	}
	c.state = stateClosed
	tr := c.tr
	if c.keepaliveCancel != nil {
		c.keepaliveCancel()
		c.keepaliveCancel = nil
	}
	c.mu.Unlock()
	if tr != nil {
		return tr.Close()
	}
	return nil
}

// newTransport selects the transport matching cfg. An explicit TransportType
// ("stdio", "sse" or "http") wins; otherwise the transport is inferred from the
// fields, preserving the legacy behaviour of stdio-when-Command and SSE-when-URL.
func newTransport(cfg MCPConfig, log *slog.Logger) (Transport, error) {
	switch cfg.TransportType {
	case "http", "streamable":
		return newHTTPTransport(cfg, log), nil
	case "sse":
		return newSSETransport(cfg, log), nil
	case "stdio":
		return newStdioTransport(cfg, log), nil
	}
	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"`
}