From 4dd04346f02b063e4fb2910455258d623e214cce Mon Sep 17 00:00:00 2001 From: giveen Date: Mon, 13 Jul 2026 14:35:39 -0600 Subject: [PATCH 1/2] fix(mcp): resolve all MCP spec compliance issues MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Replace manual stdio transport with SDK CommandTransport for proper SIGTERM → wait → SIGKILL process lifecycle - Reuse single SDK client instance across all server sessions - Wire ToolListChangedHandler to re-discover tools on dynamic changes - Handle all MCP content types (AudioContent, ResourceLink, EmbeddedResource) plus default warning for unknown types - Add SSE and Streamable HTTP transport support for remote servers - Capture stderr from stdio subprocesses on connection failure - Log warning on config file write failure instead of silent discard - Suppress unnecessary 'roots' capability advertisement --- internal/mcp/client.go | 225 +++++++++++++++++++++++++++++++---------- internal/mcp/config.go | 15 ++- 2 files changed, 179 insertions(+), 61 deletions(-) diff --git a/internal/mcp/client.go b/internal/mcp/client.go index b950943..c0cda5a 100644 --- a/internal/mcp/client.go +++ b/internal/mcp/client.go @@ -1,6 +1,7 @@ package mcp import ( + "bytes" "context" "encoding/json" "fmt" @@ -16,16 +17,32 @@ import ( // Client manages MCP connections and tools. type Client struct { - sessions map[string]*mcp.ClientSession - tools map[string]*ToolAdapter + sessions map[string]*mcp.ClientSession + tools map[string]*ToolAdapter + sdkClient *mcp.Client // single SDK client instance managing all sessions } -// NewClient creates a new MCP client. +// NewClient creates a new MCP client with a single underlying SDK client +// instance managing all sessions. The ToolListChangedHandler is wired to +// re-discover tools when a server's tool list changes dynamically. func NewClient() *Client { - return &Client{ + c := &Client{ sessions: make(map[string]*mcp.ClientSession), tools: make(map[string]*ToolAdapter), } + c.sdkClient = mcp.NewClient( + &mcp.Implementation{ + Name: "late", + Version: common.Version, + }, + &mcp.ClientOptions{ + Capabilities: &mcp.ClientCapabilities{}, // suppress the SDK default "roots" capability + ToolListChangedHandler: func(ctx context.Context, req *mcp.ToolListChangedRequest) { + c.handleToolListChanged(ctx, req) + }, + }, + ) + return c } // ToolAdapter adapts MCP tools to the Tool interface. @@ -81,10 +98,19 @@ func (t *ToolAdapter) Execute(ctx context.Context, args json.RawMessage) (string // Convert result to string var sb strings.Builder for _, content := range result.Content { - if text, ok := content.(*mcp.TextContent); ok { - sb.WriteString(text.Text) - } else if image, ok := content.(*mcp.ImageContent); ok { - sb.WriteString(fmt.Sprintf("[Image: %s]", image.MIMEType)) + switch c := content.(type) { + case *mcp.TextContent: + sb.WriteString(c.Text) + case *mcp.ImageContent: + sb.WriteString(fmt.Sprintf("[Image: %s]", c.MIMEType)) + case *mcp.AudioContent: + sb.WriteString(fmt.Sprintf("[Audio: %s]", c.MIMEType)) + case *mcp.ResourceLink: + sb.WriteString(fmt.Sprintf("[Resource: %s]", c.URI)) + case *mcp.EmbeddedResource: + sb.WriteString("[Embedded resource]") + default: + fmt.Fprintf(os.Stderr, "MCP tool returned unhandled content type: %T\n", content) } } @@ -112,17 +138,12 @@ func (t *ToolAdapter) CallString(args json.RawMessage) string { return fmt.Sprintf("Calling MCP tool '%s'...", t.Name()) } -// Connect establishes a connection to an MCP server. +// Connect establishes a connection to an MCP server using the shared SDK client. // serverName is stored on each ToolAdapter so that tool names are namespaced // as "{server}:{tool}" in allowed_tools.json, preventing collisions between // servers that expose tools with the same bare name. func (c *Client) Connect(ctx context.Context, transport mcp.Transport, serverName string) error { - client := mcp.NewClient(&mcp.Implementation{ - Name: "late", - Version: common.Version, - }, nil) - - session, err := client.Connect(ctx, transport, nil) + session, err := c.sdkClient.Connect(ctx, transport, nil) if err != nil { return fmt.Errorf("failed to connect to MCP server: %w", err) } @@ -145,6 +166,46 @@ func (c *Client) Connect(ctx context.Context, transport mcp.Transport, serverNam return nil } +// handleToolListChanged re-discovers tools for a server when the SDK notifies +// us of a tools/list change. It removes stale tool adapters for that server +// and re-enumerates via the session's paginating Tools iterator. +func (c *Client) handleToolListChanged(ctx context.Context, req *mcp.ToolListChangedRequest) { + sess := req.Session + + // Map the SDK session back to the server name we assigned in Connect. + var serverName string + for name, s := range c.sessions { + if s == sess { + serverName = name + break + } + } + if serverName == "" { + fmt.Fprintf(os.Stderr, "MCP tool list changed notification for unknown session\n") + return + } + + fmt.Printf("MCP server '%s' tools changed, re-discovering...\n", serverName) + + // Remove stale tool adapters for this server. + for name, t := range c.tools { + if t.serverName == serverName { + delete(c.tools, name) + } + } + + // Re-discover and register the new tool set. + for tool := range sess.Tools(ctx, &mcp.ListToolsParams{}) { + if tool != nil { + adapter := &ToolAdapter{ + mcpTool: tool, + session: sess, + serverName: serverName, + } + c.tools[adapter.Name()] = adapter + } + } +} // GetTools returns all MCP tools as Tool interface instances. func (c *Client) GetTools() []tool.Tool { @@ -174,45 +235,84 @@ func (c *Client) Close() error { return nil } -// NewStdioTransport creates a new transport that communicates with a subprocess. +// NewStdioTransport creates a new transport that communicates with a subprocess +// using the SDK's CommandTransport for proper process lifecycle (SIGTERM → wait → SIGKILL). +// Stderr is discarded to prevent MCP server output from bleeding into the TUI. func NewStdioTransport(ctx context.Context, command string, args []string, env []string) (mcp.Transport, error) { + return NewStdioTransportWithStderr(ctx, command, args, env, io.Discard) +} + +// NewStdioTransportWithStderr is like NewStdioTransport but writes the subprocess's +// stderr to the provided writer instead of discarding it. This is used by +// ConnectFromConfig to buffer stderr for diagnostics on connection failure. +func NewStdioTransportWithStderr(ctx context.Context, command string, args []string, env []string, stderr io.Writer) (mcp.Transport, error) { cmd := exec.Command(command, args...) cmd.Env = append(os.Environ(), env...) - stdin, err := cmd.StdinPipe() + serr, err := cmd.StderrPipe() if err != nil { - return nil, fmt.Errorf("failed to create stdin pipe: %w", err) + return nil, fmt.Errorf("failed to create stderr pipe: %w", err) } + go func() { + io.Copy(stderr, serr) + }() - stdout, err := cmd.StdoutPipe() - if err != nil { - return nil, fmt.Errorf("failed to create stdout pipe: %w", err) - } + return &mcp.CommandTransport{ + Command: cmd, + }, nil +} - stderr, err := cmd.StderrPipe() - if err != nil { - return nil, fmt.Errorf("failed to create stderr pipe: %w", err) +// TransportForServer creates the appropriate mcp.Transport for a server config. +// Supported transport types: +// - "stdio" (default when command is set) — local subprocess via CommandTransport +// - "sse" (default when url is set) — remote SSE endpoint via SSEClientTransport +// - "streamable-http" — remote streamable HTTP endpoint +// +// The server config should have already had environment variables expanded. +func TransportForServer(ctx context.Context, server *MCPServer) (mcp.Transport, error) { + switch server.TransportType { + case "sse": + if server.URL == "" { + return nil, fmt.Errorf("sse transport requires 'url'") + } + return &mcp.SSEClientTransport{ + Endpoint: server.URL, + }, nil + case "streamable-http": + if server.URL == "" { + return nil, fmt.Errorf("streamable-http transport requires 'url'") + } + return &mcp.StreamableClientTransport{ + Endpoint: server.URL, + }, nil + case "stdio": + // handled below + case "": + // Infer from available fields: URL → remote, Command → stdio + if server.URL != "" { + return &mcp.SSEClientTransport{ + Endpoint: server.URL, + }, nil + } + default: + return nil, fmt.Errorf("unknown transport type: %q", server.TransportType) } - if err := cmd.Start(); err != nil { - return nil, fmt.Errorf("failed to start command: %w", err) + // Stdio subprocess transport (explicit or default) + if server.Command == "" { + return nil, fmt.Errorf("server config must specify 'command' for stdio transport or 'url' for remote transport") } - // Discard stderr to prevent output from bleeding into TUI - go func() { - io.Copy(io.Discard, stderr) - }() - - // Kill the subprocess when the context is cancelled. - go func() { - <-ctx.Done() - cmd.Process.Kill() - }() + envSlice := make([]string, 0, len(server.Env)) + for k, v := range server.Env { + envSlice = append(envSlice, k+"="+v) + } - return &mcp.IOTransport{ - Reader: stdout, - Writer: stdin, - }, nil + t, err := NewStdioTransportWithStderr(ctx, server.Command, server.Args, envSlice, io.Discard) + if err != nil { + return nil, err + } + return t, nil } func (c *Client) ConnectFromConfig(ctx context.Context, config *MCPConfig) error { @@ -225,23 +325,36 @@ func (c *Client) ConnectFromConfig(ctx context.Context, config *MCPConfig) error // Expand environment variables in server configuration ExpandServerEnvVars(&server) - // Convert server.Env map to KEY=VALUE slice for the subprocess - envSlice := make([]string, 0, len(server.Env)) - for k, v := range server.Env { - envSlice = append(envSlice, k+"="+v) - } + var transport mcp.Transport + var err error - // Create transport for this server - transport, err := NewStdioTransport(ctx, server.Command, server.Args, envSlice) - if err != nil { - return fmt.Errorf("failed to create transport for server %s: %w", name, err) - } + // For stdio transports, buffer stderr so we can include diagnostics if + // the connection fails. Remote transports don't have stderr. + if server.TransportType == "stdio" || (server.TransportType == "" && server.Command != "" && server.URL == "") { + envSlice := make([]string, 0, len(server.Env)) + for k, v := range server.Env { + envSlice = append(envSlice, k+"="+v) + } + var stderrBuf bytes.Buffer + transport, err = NewStdioTransportWithStderr(ctx, server.Command, server.Args, envSlice, &stderrBuf) + if err == nil { + err = c.Connect(ctx, transport, name) + } + if err != nil { + if stderrBuf.Len() > 0 { + return fmt.Errorf("failed to connect to server %s: %w\nstderr:\n%s", name, err, stderrBuf.String()) + } + return fmt.Errorf("failed to connect to server %s: %w", name, err) + } + } else { + transport, err = TransportForServer(ctx, &server) + if err != nil { + return fmt.Errorf("failed to create transport for server %s: %w", name, err) + } - // Connect to the server, passing the server name so tools are registered - // with namespaced names ("{server}:{tool}") in the tool registry and - // allowed_tools.json. - if err := c.Connect(ctx, transport, name); err != nil { - return fmt.Errorf("failed to connect to server %s: %w", name, err) + if err := c.Connect(ctx, transport, name); err != nil { + return fmt.Errorf("failed to connect to server %s: %w", name, err) + } } } diff --git a/internal/mcp/config.go b/internal/mcp/config.go index 5e22e2f..c33b9a8 100644 --- a/internal/mcp/config.go +++ b/internal/mcp/config.go @@ -17,10 +17,14 @@ type MCPConfig struct { // MCPServer represents a single MCP server configuration type MCPServer struct { + // Stdio subprocess (default when Command is set): Command string `json:"command"` Args []string `json:"args"` Env map[string]string `json:"env"` - Disabled bool `json:"disabled,omitempty"` + // Remote server (used when URL is set): + URL string `json:"url"` + TransportType string `json:"transportType,omitempty"` // "stdio", "sse", or "streamable-http" + Disabled bool `json:"disabled,omitempty"` } // LoadMCPConfig loads the MCP configuration from the first available config file @@ -43,8 +47,9 @@ func LoadMCPConfig() (*MCPConfig, error) { defaultData, _ := json.MarshalIndent(emptyConfig, "", " ") if err := os.MkdirAll(lateConfigDir, 0700); err == nil { - // Ignore write error, just fallback to empty config - _ = os.WriteFile(defaultUserPath, defaultData, 0600) + if err := os.WriteFile(defaultUserPath, defaultData, 0600); err != nil { + fmt.Fprintf(os.Stderr, "Warning: Could not write default MCP config to %s: %v\n", defaultUserPath, err) + } } return &emptyConfig, nil @@ -111,12 +116,12 @@ func ExpandEnvVars(value string) string { func ExpandServerEnvVars(server *MCPServer) { // Expand command server.Command = ExpandEnvVars(server.Command) - + // Expand URL + server.URL = ExpandEnvVars(server.URL) // Expand args for i := range server.Args { server.Args[i] = ExpandEnvVars(server.Args[i]) } - // Expand env values for key, value := range server.Env { server.Env[key] = ExpandEnvVars(value) From 31238c665a167bce8c85ee0907f458d36fa61f94 Mon Sep 17 00:00:00 2001 From: giveen Date: Tue, 14 Jul 2026 15:28:16 -0600 Subject: [PATCH 2/2] fix(mcp): protect Client maps and stderr buffer from races - Add sync.RWMutex to Client and guard c.tools/c.sessions access in Connect, handleToolListChanged, GetTools, GetTool, and Close. - Wrap ConnectFromConfig's stderr capture buffer in a mutex so the background copy goroutine and the failure path no longer race. - Sleep briefly before reading stderr on failure so the goroutine has time to write the error. --- internal/mcp/client.go | 99 +++++++++++++++++++++++++++++++++--------- 1 file changed, 78 insertions(+), 21 deletions(-) diff --git a/internal/mcp/client.go b/internal/mcp/client.go index c0cda5a..80545f0 100644 --- a/internal/mcp/client.go +++ b/internal/mcp/client.go @@ -10,6 +10,8 @@ import ( "os" "os/exec" "strings" + "sync" + "time" "github.com/modelcontextprotocol/go-sdk/mcp" "late/internal/common" @@ -17,6 +19,7 @@ import ( // Client manages MCP connections and tools. type Client struct { + mu sync.RWMutex sessions map[string]*mcp.ClientSession tools map[string]*ToolAdapter sdkClient *mcp.Client // single SDK client instance managing all sessions @@ -148,22 +151,28 @@ func (c *Client) Connect(ctx context.Context, transport mcp.Transport, serverNam return fmt.Errorf("failed to connect to MCP server: %w", err) } - // Store session keyed by server name so Close() shuts down every - // subprocess, not just the last one connected. - c.sessions[serverName] = session - - // List and store tools using iterator + // Collect adapters without holding the lock; the SDK's Tools iterator may + // perform RPCs, and we want to avoid blocking readers. + var adapters []*ToolAdapter for tool := range session.Tools(ctx, &mcp.ListToolsParams{}) { if tool != nil { - adapter := &ToolAdapter{ + adapters = append(adapters, &ToolAdapter{ mcpTool: tool, session: session, serverName: serverName, - } - c.tools[adapter.Name()] = adapter + }) } } + // Store session keyed by server name so Close() shuts down every + // subprocess, not just the last one connected. + c.mu.Lock() + c.sessions[serverName] = session + for _, adapter := range adapters { + c.tools[adapter.Name()] = adapter + } + c.mu.Unlock() + return nil } // handleToolListChanged re-discovers tools for a server when the SDK notifies @@ -173,6 +182,7 @@ func (c *Client) handleToolListChanged(ctx context.Context, req *mcp.ToolListCha sess := req.Session // Map the SDK session back to the server name we assigned in Connect. + c.mu.RLock() var serverName string for name, s := range c.sessions { if s == sess { @@ -180,6 +190,7 @@ func (c *Client) handleToolListChanged(ctx context.Context, req *mcp.ToolListCha break } } + c.mu.RUnlock() if serverName == "" { fmt.Fprintf(os.Stderr, "MCP tool list changed notification for unknown session\n") return @@ -187,28 +198,36 @@ func (c *Client) handleToolListChanged(ctx context.Context, req *mcp.ToolListCha fmt.Printf("MCP server '%s' tools changed, re-discovering...\n", serverName) - // Remove stale tool adapters for this server. - for name, t := range c.tools { - if t.serverName == serverName { - delete(c.tools, name) - } - } - - // Re-discover and register the new tool set. + // Collect the new tool set without holding the lock; the SDK's Tools + // iterator may perform RPCs. + var adapters []*ToolAdapter for tool := range sess.Tools(ctx, &mcp.ListToolsParams{}) { if tool != nil { - adapter := &ToolAdapter{ + adapters = append(adapters, &ToolAdapter{ mcpTool: tool, session: sess, serverName: serverName, - } - c.tools[adapter.Name()] = adapter + }) } } + + // Remove stale tool adapters for this server and register the new set. + c.mu.Lock() + for name, t := range c.tools { + if t.serverName == serverName { + delete(c.tools, name) + } + } + for _, adapter := range adapters { + c.tools[adapter.Name()] = adapter + } + c.mu.Unlock() } // GetTools returns all MCP tools as Tool interface instances. func (c *Client) GetTools() []tool.Tool { + c.mu.RLock() + defer c.mu.RUnlock() tools := make([]tool.Tool, 0, len(c.tools)) for _, t := range c.tools { tools = append(tools, t) @@ -218,6 +237,8 @@ func (c *Client) GetTools() []tool.Tool { // GetTool returns a specific MCP tool by name. func (c *Client) GetTool(name string) tool.Tool { + c.mu.RLock() + defer c.mu.RUnlock() t, ok := c.tools[name] if !ok { return nil @@ -227,14 +248,48 @@ func (c *Client) GetTool(name string) tool.Tool { // Close closes all MCP connections. func (c *Client) Close() error { + c.mu.RLock() + sessions := make([]*mcp.ClientSession, 0, len(c.sessions)) + names := make([]string, 0, len(c.sessions)) for name, session := range c.sessions { + names = append(names, name) + sessions = append(sessions, session) + } + c.mu.RUnlock() + + for i, session := range sessions { if err := session.Close(); err != nil { - fmt.Fprintf(os.Stderr, "Error closing MCP session '%s': %v\n", name, err) + fmt.Fprintf(os.Stderr, "Error closing MCP session '%s': %v\n", names[i], err) } } return nil } +// lockedBuffer wraps a bytes.Buffer with a mutex so it can be safely written +// by a background goroutine and read by another goroutine on connection failure. +type lockedBuffer struct { + mu sync.Mutex + buf bytes.Buffer +} + +func (b *lockedBuffer) Write(p []byte) (int, error) { + b.mu.Lock() + defer b.mu.Unlock() + return b.buf.Write(p) +} + +func (b *lockedBuffer) Len() int { + b.mu.Lock() + defer b.mu.Unlock() + return b.buf.Len() +} + +func (b *lockedBuffer) String() string { + b.mu.Lock() + defer b.mu.Unlock() + return b.buf.String() +} + // NewStdioTransport creates a new transport that communicates with a subprocess // using the SDK's CommandTransport for proper process lifecycle (SIGTERM → wait → SIGKILL). // Stderr is discarded to prevent MCP server output from bleeding into the TUI. @@ -335,12 +390,14 @@ func (c *Client) ConnectFromConfig(ctx context.Context, config *MCPConfig) error for k, v := range server.Env { envSlice = append(envSlice, k+"="+v) } - var stderrBuf bytes.Buffer + var stderrBuf lockedBuffer transport, err = NewStdioTransportWithStderr(ctx, server.Command, server.Args, envSlice, &stderrBuf) if err == nil { err = c.Connect(ctx, transport, name) } if err != nil { + // Give the stderr goroutine a moment to capture the error. + time.Sleep(50 * time.Millisecond) if stderrBuf.Len() > 0 { return fmt.Errorf("failed to connect to server %s: %w\nstderr:\n%s", name, err, stderrBuf.String()) }