mcp

package
v0.0.0-...-a256278 Latest Latest
Warning

This package is not in the latest version of its module.

Go to latest
Published: Aug 27, 2026 License: Apache-2.0 Imports: 33 Imported by: 0

Documentation

Overview

Package mcp provides MCP (Model Context Protocol) client infrastructure for connecting to and executing tools on MCP servers.

Index

Constants

View Source
const (
	// MaxRetries is the number of retry attempts after the initial failure.
	MaxRetries = 1

	// ReinitTimeout is the deadline for recreating an MCP session during recovery.
	ReinitTimeout = 10 * time.Second

	// OperationTimeout is the per-call deadline for ListTools and reinit.
	// Tool calls are separately governed by ToolCallTimeout (default 1m)
	// applied in the controller layer; this remains as defense-in-depth.
	OperationTimeout = 90 * time.Second

	// RetryBackoffMin is the minimum jittered backoff between retries.
	RetryBackoffMin = 250 * time.Millisecond

	// RetryBackoffMax is the maximum jittered backoff between retries.
	RetryBackoffMax = 750 * time.Millisecond

	// MCPInitTimeout is the per-server initialization timeout (transport + handshake).
	MCPInitTimeout = 30 * time.Second

	// MCPHealthPingTimeout is the health check ping timeout.
	MCPHealthPingTimeout = 5 * time.Second

	// MCPHealthInterval is the health check loop interval.
	MCPHealthInterval = 15 * time.Second
)

Recovery configuration constants.

View Source
const DefaultStorageMaxTokens = 8000

DefaultStorageMaxTokens is the maximum token count for storage-truncated tool output. Protects the dashboard from rendering massive text blobs.

View Source
const DefaultSummarizationMaxTokens = 100000

DefaultSummarizationMaxTokens is the maximum token count for summarization LLM input. Safety net — summarization prompt + truncated output must fit in the model's context window.

Variables

This section is empty.

Functions

func CleanupSessions

func CleanupSessions(
	ctx context.Context,
	registry *config.MCPServerRegistry,
	mcpSessionID string,
	serverIDs []string,
	logger *slog.Logger,
)

CleanupSessions best-effort deletes per-execution sandbox sessions on the given MCP servers that declare session_cleanup_url (e.g. cli-mcp-server). mcpSessionID is the agent execution ID used as X-Session-ID. Only serverIDs requested by the closing client are considered — other registry servers are left alone. Failures are logged and never returned — idle TTL remains the safety net.

func EstimateTokens

func EstimateTokens(text string) int

EstimateTokens returns an approximate token count for the given text. Uses the common heuristic of ~4 characters per token for English text. This is intentionally approximate — exact counts would require a tokenizer library and add a dependency for minimal benefit (the threshold is a configurable soft limit, not a hard boundary).

Note: len(text) counts bytes, not Unicode characters. For multi-byte UTF-8 content (CJK, emoji), this overestimates the character count and therefore the token count. This is a safe direction to err — summarization triggers slightly earlier than necessary, which is preferable to missing it.

func NormalizeBuiltinPlainToolName

func NormalizeBuiltinPlainToolName(name string) string

NormalizeBuiltinPlainToolName strips a mistaken provider: prefix when the suffix is a known built-in plain tool (e.g. "google:load_skill" → "load_skill").

func NormalizeToolName

func NormalizeToolName(name string) string

NormalizeToolName converts tool names between controller formats. FunctionCalling uses "server__tool" (API name restriction for Gemini/LangChain). Normalizes to "server.tool" for routing.

func ParseActionInput

func ParseActionInput(input string) (map[string]any, error)

ParseActionInput parses a raw ActionInput string into structured parameters.

Parsing cascade (first successful parse wins):

  1. JSON object → map[string]any
  2. JSON non-object (string, number, array) → {"input": value}
  3. YAML with complex structures (arrays, nested maps) → map[string]any
  4. Key-value pairs (key: value or key=value, comma/newline separated)
  5. Single raw string → {"input": string}

Empty input returns empty map (for no-parameter tools).

func SplitToolName

func SplitToolName(name string) (serverID, toolName string, err error)

SplitToolName splits "server.tool" into (serverID, toolName, error). Validates format with strict regex: server and tool parts must be word characters and hyphens, non-empty.

func TruncateForStorage

func TruncateForStorage(content string) string

TruncateForStorage truncates tool output for llm_tool_call completion content and MCPInteraction records. Protects the UI from rendering massive text blobs. Applied to ALL raw results, regardless of whether summarization is triggered.

func TruncateForSummarization

func TruncateForSummarization(content string) string

TruncateForSummarization truncates tool output before sending to the summarization LLM. Safety net — summarization prompt + truncated output must fit in the model's context window. Uses a larger limit than storage truncation to give the summarizer maximum data.

Types

type Client

type Client struct {
	// contains filtered or unexported fields
}

Client manages MCP SDK sessions for multiple servers. Each Client instance is scoped to a single agent execution (or health check). Thread-safe: sessions may be accessed from multiple goroutines during parallel stages.

func (*Client) CallTool

func (c *Client) CallTool(ctx context.Context, serverID, toolName string, args map[string]any) (*mcpsdk.CallToolResult, error)

CallTool executes a tool call on the specified server. Handles recovery (retry with session recreation) on transport failures. At most one retry is attempted after a jittered backoff; if the retry also fails the error is returned to the caller.

func (*Client) Close

func (c *Client) Close() error

Close shuts down all sessions and transports gracefully, then best-effort deletes per-execution MCP sandboxes (e.g. cli-mcp-server) when mcpSessionID is set. Idempotent: a second Close skips sandbox DELETE.

func (*Client) FailedServers

func (c *Client) FailedServers() map[string]string

FailedServers returns the map of servers that failed to initialize.

func (*Client) HasSession

func (c *Client) HasSession(serverID string) bool

HasSession checks if a server has an active session.

func (*Client) Initialize

func (c *Client) Initialize(ctx context.Context, serverIDs []string)

Initialize connects to all configured MCP servers. Servers that fail to connect are recorded in failedServers (retrievable via FailedServers()). The caller decides how to handle partial failures:

  • Startup: check FailedServers() and log warnings (non-fatal, TARSy starts degraded)
  • Per agent execution: partial initialization is acceptable

func (*Client) InitializeServer

func (c *Client) InitializeServer(ctx context.Context, serverID string) error

InitializeServer connects to a single MCP server. Returns nil if already connected. Used for lazy initialization and recovery. Uses per-server mutex to prevent concurrent initialization of the same server.

func (*Client) InjectSession

func (c *Client) InjectSession(serverID string, sdkClient *mcpsdk.Client, session *mcpsdk.ClientSession)

InjectSession injects a pre-connected MCP SDK session into the Client. This is intended for test infrastructure that needs to wire in-memory MCP servers without going through the real Initialize() transport creation path.

func (*Client) InvalidateToolCache

func (c *Client) InvalidateToolCache(serverID string)

InvalidateToolCache removes the cached tool list for a server, forcing the next ListTools call to re-probe the server. Lock ordering: never acquire c.mu while holding toolCacheMu.

func (*Client) ListAllTools

func (c *Client) ListAllTools(ctx context.Context) (map[string][]*mcpsdk.Tool, error)

ListAllTools returns tools from all connected servers. Returns partial results if some servers fail (logs errors, does not abort). Returns an error only when every server fails (no tools available at all).

func (*Client) ListTools

func (c *Client) ListTools(ctx context.Context, serverID string) ([]*mcpsdk.Tool, error)

ListTools returns tools from a specific server. Uses cache if available.

type ClientFactory

type ClientFactory struct {
	// contains filtered or unexported fields
}

ClientFactory creates Client instances for agent executions.

func NewClientFactory

func NewClientFactory(registry *config.MCPServerRegistry, maskingService *masking.Service) *ClientFactory

NewClientFactory creates a new factory. maskingService may be nil (masking disabled).

func NewTestClientFactory

func NewTestClientFactory(registry *config.MCPServerRegistry, injectFn func(c *Client, mcpSessionID string)) *ClientFactory

NewTestClientFactory creates a ClientFactory that uses injectFn to wire sessions into each new Client instead of calling Initialize(). Each call to CreateClient/CreateToolExecutor invokes injectFn on the freshly-created Client with the mcpSessionID used for that client.

func (*ClientFactory) CreateClient

func (f *ClientFactory) CreateClient(ctx context.Context, serverIDs []string, mcpSessionID string) (*Client, error)

CreateClient creates a new Client connected to the specified servers. mcpSessionID is the agent execution ID used to resolve per-execution custom_headers (e.g. X-Session-ID). Pass an empty string for health checks and startup validation (no sandbox header; cli-mcp only requires X-Session-ID for bash tool calls). The caller is responsible for calling Close() when done (also triggers sandbox cleanup for requested servers that declare session_cleanup_url).

func (*ClientFactory) CreateToolExecutor

func (f *ClientFactory) CreateToolExecutor(
	ctx context.Context,
	serverIDs []string,
	toolFilter map[string][]string,
	mcpSessionID string,
) (*ToolExecutor, *Client, error)

CreateToolExecutor creates a fully-wired ToolExecutor for an agent execution. This is the primary entry point used by the session executor. mcpSessionID is the agent execution ID forwarded for custom_headers resolution.

type HealthMonitor

type HealthMonitor struct {
	// contains filtered or unexported fields
}

HealthMonitor periodically checks MCP server health. Runs a background goroutine that probes each server with ListTools.

func NewHealthMonitor

func NewHealthMonitor(
	factory *ClientFactory,
	registry *config.MCPServerRegistry,
	warningService *services.SystemWarningsService,
) *HealthMonitor

NewHealthMonitor creates a new health monitor.

func (*HealthMonitor) GetCachedTools

func (m *HealthMonitor) GetCachedTools() map[string][]*mcpsdk.Tool

GetCachedTools returns the cached tools from the last successful health check. The returned map is a shallow copy: slices share the same underlying Tool pointers with the monitor's cache. Callers must not mutate the slices.

func (*HealthMonitor) GetStatuses

func (m *HealthMonitor) GetStatuses() map[string]*HealthStatus

GetStatuses returns the current health status of all monitored servers.

func (*HealthMonitor) IsHealthy

func (m *HealthMonitor) IsHealthy() bool

IsHealthy returns true if all monitored servers are healthy. Returns true when no statuses exist (no servers configured or before first check completes) so that an empty registry doesn't cause spurious failures.

func (*HealthMonitor) Start

func (m *HealthMonitor) Start(ctx context.Context)

Start launches the background health check loop. Calling Start on an already-running monitor is a no-op.

func (*HealthMonitor) Stop

func (m *HealthMonitor) Stop()

Stop gracefully shuts down the health monitor. After Stop returns, Start may be called again.

type HealthStatus

type HealthStatus struct {
	ServerID  string    `json:"server_id"`
	Healthy   bool      `json:"healthy"`
	LastCheck time.Time `json:"last_check"`
	Error     string    `json:"error,omitempty"`
	ToolCount int       `json:"tool_count"`
}

HealthStatus captures the health check result for a single MCP server.

type RecoveryAction

type RecoveryAction int

RecoveryAction determines how to handle an MCP operation failure.

const (
	// NoRetry — the error is not recoverable (bad request, auth failure, timeout).
	NoRetry RecoveryAction = iota
	// RetrySameSession — transient error, retry with existing session (rate limit).
	// Reserved for future use: ClassifyError does not currently return this value.
	// Intended for rate-limit / throttle errors once server-side rate limiting is detected.
	RetrySameSession
	// RetryNewSession — transport failure, recreate session and retry.
	RetryNewSession
)

func ClassifyError

func ClassifyError(err error) RecoveryAction

ClassifyError determines the recovery action for an MCP operation error.

type ToolExecutor

type ToolExecutor struct {
	// contains filtered or unexported fields
}

ToolExecutor implements agent.ToolExecutor backed by real MCP servers. Created per agent execution by ClientFactory.

func NewToolExecutor

func NewToolExecutor(
	client *Client,
	registry *config.MCPServerRegistry,
	serverIDs []string,
	toolFilter map[string][]string,
	maskingService *masking.Service,
) *ToolExecutor

NewToolExecutor creates a new executor for the given servers. maskingService may be nil (masking disabled).

func (*ToolExecutor) Close

func (e *ToolExecutor) Close() error

Close releases MCP transports/subprocesses and best-effort deletes per-execution sandboxes for servers with session_cleanup_url.

func (*ToolExecutor) Execute

func (e *ToolExecutor) Execute(ctx context.Context, call agent.ToolCall) (*agent.ToolResult, error)

Execute runs a tool call via MCP.

Flow:

  1. Normalize tool name (server__tool → server.tool for GoogleNative)
  2. Split and validate server.tool name
  3. Check server is in allowed serverIDs
  4. Check tool is in allowed tools (if filter set)
  5. Parse Arguments string into map[string]any
  6. Call Client.CallTool(ctx, serverID, toolName, params)
  7. Convert MCP result to ToolResult
  8. Apply data masking (if masking service configured)
  9. Return ToolResult (summarization is handled at the controller level)

func (*ToolExecutor) ListTools

func (e *ToolExecutor) ListTools(ctx context.Context) ([]agent.ToolDefinition, error)

ListTools returns all available tools from configured MCP servers. Tools are returned with server-prefixed names (e.g., "kubernetes-server.get_pods").

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL