From c3ae7775acd6038c143302fa9f8fd1917e4b0bac Mon Sep 17 00:00:00 2001 From: Christopher Pappas Date: Sat, 22 Nov 2025 21:26:13 +0100 Subject: [PATCH] fix: fail faster when mcp servers fail to connect --- src/components/UserInput.tsx | 2 +- src/hooks/useAgent.ts | 6 ++- src/mcp/getAgentStatus.ts | 4 +- src/mcp/getMcpServer.ts | 14 +++++-- src/mcp/runStandaloneAgentLoop.ts | 20 +++++---- src/mcp/tools/askAgent.ts | 21 ++++++---- src/mcp/tools/askAgentSlackbot.ts | 21 ++++++---- src/utils/__tests__/getPrompt.test.ts | 59 +++++++++++++++++++++------ src/utils/getPrompt.ts | 58 +++++++++++++++++++++----- src/utils/mcpServerSelectionAgent.ts | 18 ++++---- src/utils/runAgentLoop.ts | 19 +++++---- 11 files changed, 171 insertions(+), 71 deletions(-) diff --git a/src/components/UserInput.tsx b/src/components/UserInput.tsx index b680026..b7923e3 100644 --- a/src/components/UserInput.tsx +++ b/src/components/UserInput.tsx @@ -46,7 +46,7 @@ export const UserInput: React.FC = () => { return ( - + state.messageQueue) const config = AgentStore.useStoreState((state) => state.config) + const mcpServers = AgentStore.useStoreState((state) => state.mcpServers) const actions = AgentStore.useStoreActions((actions) => actions) const currentAssistantMessageRef = useRef("") const sessionIdRef = useRef(undefined) const abortControllerRef = useRef(undefined) - const connectedServersRef = useRef>(new Set()) + const inferredServersRef = useRef>(new Set()) const runQuery = useCallback( async (userMessage: string) => { @@ -32,7 +33,8 @@ export function useAgent() { const agentLoop = runAgentLoop({ abortController, config, - connectedServers: connectedServersRef.current, + inferredServers: inferredServersRef.current, + mcpServers, messageQueue, onToolPermissionRequest: (toolName, input) => { actions.setPendingToolPermission({ toolName, input }) diff --git a/src/mcp/getAgentStatus.ts b/src/mcp/getAgentStatus.ts index 451192a..00d6e6b 100644 --- a/src/mcp/getAgentStatus.ts +++ b/src/mcp/getAgentStatus.ts @@ -7,12 +7,12 @@ export const getAgentStatus = async (mcpServer?: McpServer) => { const config = await loadConfig() const messageQueue = new MessageQueue() const abortController = new AbortController() - const connectedServers = new Set() + const inferredServers = new Set() const agentLoop = runAgentLoop({ abortController, config, - connectedServers, + inferredServers, messageQueue, userMessage: "status", }) diff --git a/src/mcp/getMcpServer.ts b/src/mcp/getMcpServer.ts index 3cbf545..f34ef5a 100644 --- a/src/mcp/getMcpServer.ts +++ b/src/mcp/getMcpServer.ts @@ -1,4 +1,5 @@ import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js" +import type { McpServerStatus } from "store" import { registerAskAgentTool } from "mcp/tools/askAgent" import { registerAskAgentSlackbotTool } from "mcp/tools/askAgentSlackbot" import { registerGetAgentStatusTool } from "mcp/tools/getAgentStatus" @@ -10,8 +11,11 @@ export const getMcpServer = () => { // Map thread IDs to Claude Agent SDK session IDs for per-thread isolation const threadSessions = new Map() - // Map session IDs to connected MCP servers for persistence across requests - const sessionConnectedServers = new Map>() + // Map session IDs to inferred MCP servers for persistence across requests + const sessionInferredServers = new Map>() + + // Map session IDs to MCP server statuses for persistence across requests + const sessionMcpServers = new Map() const mcpServer = new McpServer( { @@ -32,7 +36,8 @@ export const getMcpServer = () => { get sessionId() { return sessionId }, - sessionConnectedServers, + sessionInferredServers, + sessionMcpServers, onSessionIdUpdate: (newSessionId) => { sessionId = newSessionId }, @@ -43,7 +48,8 @@ export const getMcpServer = () => { mcpServer, context: { threadSessions, - sessionConnectedServers, + sessionInferredServers, + sessionMcpServers, }, }) diff --git a/src/mcp/runStandaloneAgentLoop.ts b/src/mcp/runStandaloneAgentLoop.ts index 0d5978e..dcf1375 100644 --- a/src/mcp/runStandaloneAgentLoop.ts +++ b/src/mcp/runStandaloneAgentLoop.ts @@ -1,4 +1,5 @@ import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js" +import type { McpServerStatus } from "store" import { loadConfig } from "utils/loadConfig" import { log } from "utils/logger" import { MessageQueue } from "utils/MessageQueue" @@ -6,7 +7,8 @@ import { contentTypes, messageTypes, runAgentLoop } from "utils/runAgentLoop" interface RunQueryOptions { additionalSystemPrompt?: string - existingConnectedServers?: Set + existingInferredServers?: Set + existingMcpServers?: McpServerStatus[] mcpServer: McpServer onSessionIdReceived?: (sessionId: string) => void prompt: string @@ -15,7 +17,8 @@ interface RunQueryOptions { export const runStandaloneAgentLoop = async ({ additionalSystemPrompt, - existingConnectedServers, + existingInferredServers, + existingMcpServers, mcpServer, onSessionIdReceived, prompt, @@ -25,17 +28,16 @@ export const runStandaloneAgentLoop = async ({ const messageQueue = new MessageQueue() const streamEnabled = config.stream ?? false - const connectedServers = existingConnectedServers ?? new Set() + const inferredServers = existingInferredServers ?? new Set() const abortController = new AbortController() const agentLoop = runAgentLoop({ abortController, additionalSystemPrompt, config, - connectedServers, + inferredServers, + mcpServers: existingMcpServers, messageQueue, - sessionId, - userMessage: prompt, onServerConnection: async (status) => { await mcpServer.sendLoggingMessage({ level: "info", @@ -45,6 +47,8 @@ export const runStandaloneAgentLoop = async ({ }), }) }, + sessionId, + userMessage: prompt, }) let finalResponse = "" @@ -155,7 +159,7 @@ export const runStandaloneAgentLoop = async ({ return { response: finalResponse, - connectedServers, + inferredServers, } } } @@ -166,6 +170,6 @@ export const runStandaloneAgentLoop = async ({ return { response: finalResponse, - connectedServers, + inferredServers, } } diff --git a/src/mcp/tools/askAgent.ts b/src/mcp/tools/askAgent.ts index 0ff051e..d7d3f30 100644 --- a/src/mcp/tools/askAgent.ts +++ b/src/mcp/tools/askAgent.ts @@ -1,10 +1,12 @@ import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js" +import type { McpServerStatus } from "store" import { runStandaloneAgentLoop } from "mcp/runStandaloneAgentLoop" import { z } from "zod" export interface AskAgentContext { sessionId?: string - sessionConnectedServers: Map> + sessionInferredServers: Map> + sessionMcpServers: Map onSessionIdUpdate: (sessionId: string) => void } @@ -27,23 +29,28 @@ export const registerAskAgentTool = ({ }, }, async ({ query }) => { - const existingConnectedServers = context.sessionId - ? context.sessionConnectedServers.get(context.sessionId) + const existingInferredServers = context.sessionId + ? context.sessionInferredServers.get(context.sessionId) : undefined - const { response, connectedServers } = await runStandaloneAgentLoop({ + const existingMcpServers = context.sessionId + ? context.sessionMcpServers.get(context.sessionId) + : undefined + + const { response, inferredServers } = await runStandaloneAgentLoop({ prompt: query, mcpServer, sessionId: context.sessionId, - existingConnectedServers, + existingInferredServers, + existingMcpServers, onSessionIdReceived: (newSessionId) => { context.onSessionIdUpdate(newSessionId) }, }) - // Update the session's connected servers + // Update the session's inferred servers if (context.sessionId) { - context.sessionConnectedServers.set(context.sessionId, connectedServers) + context.sessionInferredServers.set(context.sessionId, inferredServers) } return { diff --git a/src/mcp/tools/askAgentSlackbot.ts b/src/mcp/tools/askAgentSlackbot.ts index 42594ca..615b749 100644 --- a/src/mcp/tools/askAgentSlackbot.ts +++ b/src/mcp/tools/askAgentSlackbot.ts @@ -1,10 +1,12 @@ import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js" +import type { McpServerStatus } from "store" import { runStandaloneAgentLoop } from "mcp/runStandaloneAgentLoop" import { z } from "zod" export interface AskAgentSlackbotContext { threadSessions: Map - sessionConnectedServers: Map> + sessionInferredServers: Map> + sessionMcpServers: Map } export interface RegisterAskAgentSlackbotToolProps { @@ -41,16 +43,21 @@ export const registerAskAgentSlackbotTool = ({ ? context.threadSessions.get(threadId) : undefined - const existingConnectedServers = existingSessionId - ? context.sessionConnectedServers.get(existingSessionId) + const existingInferredServers = existingSessionId + ? context.sessionInferredServers.get(existingSessionId) : undefined - const { response, connectedServers } = await runStandaloneAgentLoop({ + const existingMcpServers = existingSessionId + ? context.sessionMcpServers.get(existingSessionId) + : undefined + + const { response, inferredServers } = await runStandaloneAgentLoop({ prompt: query, mcpServer, sessionId: existingSessionId, additionalSystemPrompt: systemPrompt, - existingConnectedServers, + existingInferredServers, + existingMcpServers, onSessionIdReceived: (newSessionId) => { if (threadId) { context.threadSessions.set(threadId, newSessionId) @@ -58,12 +65,12 @@ export const registerAskAgentSlackbotTool = ({ }, }) - // Update the session's connected servers + // Update the session's inferred servers if (existingSessionId || threadId) { const sessionId = existingSessionId || context.threadSessions.get(threadId!) if (sessionId) { - context.sessionConnectedServers.set(sessionId, connectedServers) + context.sessionInferredServers.set(sessionId, inferredServers) } } diff --git a/src/utils/__tests__/getPrompt.test.ts b/src/utils/__tests__/getPrompt.test.ts index b0308b4..01c1f22 100644 --- a/src/utils/__tests__/getPrompt.test.ts +++ b/src/utils/__tests__/getPrompt.test.ts @@ -145,47 +145,80 @@ describe("buildSystemPrompt", () => { expect(prompt).toContain("GitHub instructions") }) - test("should include connected MCP servers in system prompt", async () => { + test("should include inferred MCP servers in system prompt", async () => { const config: AgentChatConfig = { mcpServers: {}, } - const connectedServers = new Set(["github", "gitlab"]) + const inferredServers = new Set(["github", "gitlab"]) const prompt = await buildSystemPrompt({ config, - connectedServers, + inferredServers, }) - expect(prompt).toContain("Connected MCP Servers") + expect(prompt).toContain("Server Selection Context") expect(prompt).toContain("github, gitlab") }) - test("should handle empty connected servers set", async () => { + test("should include connected and failed MCP servers sections", async () => { const config: AgentChatConfig = { mcpServers: {}, } - const connectedServers = new Set() + const mcpServers = [ + { name: "github", status: "connected" }, + { name: "gitlab", status: "failed" }, + ] const prompt = await buildSystemPrompt({ config, - connectedServers, + mcpServers, }) - expect(prompt).toContain("Connected MCP Servers") - expect(prompt).toContain( - "The following MCP servers are currently connected and available:" - ) + expect(prompt).toContain("Available MCP Servers") + expect(prompt).toContain("- github") + expect(prompt).toContain("ONLY servers with available tools") + expect(prompt).toContain("Unavailable MCP Servers") + expect(prompt).toContain("- gitlab") + expect(prompt).toContain("These servers have NO tools available") }) - test("should work without connectedServers parameter", async () => { + test("should not include connection status sections when no mcpServers provided", async () => { + const config: AgentChatConfig = { + mcpServers: {}, + } + + const prompt = await buildSystemPrompt({ + config, + }) + + expect(prompt).not.toContain("Available MCP Servers") + expect(prompt).not.toContain("Unavailable MCP Servers") + }) + + test("should handle empty inferred servers set", async () => { + const config: AgentChatConfig = { + mcpServers: {}, + } + const inferredServers = new Set() + + const prompt = await buildSystemPrompt({ + config, + inferredServers, + }) + + expect(prompt).toContain("Current date:") + expect(prompt).not.toContain("Server Selection Context") + }) + + test("should work without inferredServers parameter", async () => { const config: AgentChatConfig = { mcpServers: {}, } const prompt = await buildSystemPrompt({ config }) - expect(prompt).toContain("Connected MCP Servers") expect(prompt).toContain("Current date:") + expect(prompt).not.toContain("Server Selection Context") }) test("should skip disabled MCP servers", async () => { diff --git a/src/utils/getPrompt.ts b/src/utils/getPrompt.ts index bc91584..498e950 100644 --- a/src/utils/getPrompt.ts +++ b/src/utils/getPrompt.ts @@ -2,7 +2,7 @@ import type { Options } from "@anthropic-ai/claude-agent-sdk" import { readFileSync } from "node:fs" import { dirname, resolve } from "node:path" import { fileURLToPath } from "node:url" -import type { AgentChatConfig } from "store" +import type { AgentChatConfig, McpServerStatus } from "store" const __dirname = dirname(fileURLToPath(import.meta.url)) @@ -16,13 +16,15 @@ export const getPrompt = (filename: string) => { interface BuildSystemPromptProps { config: AgentChatConfig additionalSystemPrompt?: string - connectedServers?: Set + inferredServers?: Set + mcpServers?: McpServerStatus[] } export const buildSystemPrompt = async ({ config, additionalSystemPrompt = "", - connectedServers = new Set(), + inferredServers = new Set(), + mcpServers = [], }: BuildSystemPromptProps): Promise => { const currentDate = new Date().toLocaleDateString("en-US", { weekday: "long", @@ -56,17 +58,53 @@ export const buildSystemPrompt = async ({ parts.push(additionalSystemPrompt) } - const connectedMCPServers = Array.from(connectedServers).join(", ") + if (mcpServers.length > 0) { + // Add connection status sections first as these are the source of truth. + // Inference is secondary. + const connectedServers = mcpServers + .filter((s) => s.status === "connected") + .map((s) => s.name) + const failedServers = mcpServers + .filter((s) => s.status === "failed") + .map((s) => s.name) - parts.push( - `# Connected MCP Servers + if (connectedServers.length > 0) { + parts.push( + `# Available MCP Servers -**The following MCP servers are currently connected and available: ${connectedMCPServers}** +The following MCP servers are **currently connected**: +${connectedServers.map((s) => `- ${s}`).join("\n")} -- **IMPORTANT**: Only use tools from these connected servers. -- If a user asks about a tool or server that is not in this list, immediately inform them that the server is not connected and cannot be used. +These are the ONLY servers with available tools. Use mcp_[servername]_[toolname] to call their tools. ` - ) + ) + } + + if (failedServers.length > 0) { + parts.push( + `# Unavailable MCP Servers + +The following MCP servers **failed to connect**: +${failedServers.map((s) => `- ${s}`).join("\n")} + +These servers have NO tools available. If a user requests functionality from these servers, inform them that the connection failed and the feature is temporarily unavailable. +` + ) + } + } + + const inferredMCPServers = Array.from(inferredServers).join(", ") + + if (inferredMCPServers) { + parts.push( + `# Server Selection Context + +The following servers were inferred as potentially needed: ${inferredMCPServers} + +Note: This is context about what was _should_ connect based on inference from the users question, not about what has actually connected. **Refer to the "Available MCP Servers" and "Unavailable MCP Servers" sections above for which servers actually have tools.** +` + ) + } parts.push(basePrompt) diff --git a/src/utils/mcpServerSelectionAgent.ts b/src/utils/mcpServerSelectionAgent.ts index d663152..4056a51 100644 --- a/src/utils/mcpServerSelectionAgent.ts +++ b/src/utils/mcpServerSelectionAgent.ts @@ -12,17 +12,17 @@ import { messageTypes } from "./runAgentLoop" interface SelectMcpServersOptions { abortController?: AbortController agents?: Record - connectedServers?: Set + inferredServers?: Set enabledMcpServers: Record | undefined onServerConnection?: (status: string) => void sessionId?: string userMessage: string } -export const selectMcpServers = async ({ +export const inferMcpServers = async ({ abortController, agents, - connectedServers = new Set(), + inferredServers = new Set(), enabledMcpServers, onServerConnection, sessionId, @@ -37,8 +37,8 @@ export const selectMcpServers = async ({ log("[mcpServerSelectionAgent] Available servers:", mcpServerNames) log( - "[mcpServerSelectionAgent] Already connected:", - Array.from(connectedServers).join(", ") || "none" + "[mcpServerSelectionAgent] Already inferred:", + Array.from(inferredServers).join(", ") || "none" ) const serverCapabilities = Object.entries(enabledMcpServers) @@ -159,7 +159,7 @@ Examples: log("[mcpServerSelectionAgent] Selected MCP servers:", selectedServers) const newServers = selectedServers.filter( - (server) => !connectedServers.has(server.toLowerCase()) + (server) => !inferredServers.has(server.toLowerCase()) ) if (newServers.length > 0) { @@ -172,7 +172,7 @@ Examples: } const allServers = new Set([ - ...Array.from(connectedServers), + ...Array.from(inferredServers), ...selectedServers, ]) @@ -210,9 +210,9 @@ Examples: onServerConnection?.(`Connecting to ${serverList}...`) } - // Update the connected servers set with new servers + // Update the inferred servers set with new servers newServers.forEach((server) => { - connectedServers.add(server.toLowerCase()) + inferredServers.add(server.toLowerCase()) }) return { diff --git a/src/utils/runAgentLoop.ts b/src/utils/runAgentLoop.ts index 34d655d..9fd4c5f 100644 --- a/src/utils/runAgentLoop.ts +++ b/src/utils/runAgentLoop.ts @@ -1,12 +1,12 @@ import { query } from "@anthropic-ai/claude-agent-sdk" -import type { AgentChatConfig } from "store" +import type { AgentChatConfig, McpServerStatus } from "store" import { createCanUseTool } from "utils/canUseTool" import { createSDKAgents } from "utils/createAgent" import { getEnabledMcpServers } from "utils/getEnabledMcpServers" import { buildSystemPrompt } from "utils/getPrompt" import { getDisallowedTools } from "utils/getToolInfo" import { log } from "utils/logger" -import { selectMcpServers } from "utils/mcpServerSelectionAgent" +import { inferMcpServers } from "utils/mcpServerSelectionAgent" import type { MessageQueue } from "utils/MessageQueue" export const messageTypes = { @@ -26,7 +26,8 @@ export interface RunAgentLoopOptions { abortController: AbortController additionalSystemPrompt?: string config: AgentChatConfig - connectedServers: Set + inferredServers: Set + mcpServers?: McpServerStatus[] messageQueue: MessageQueue onServerConnection?: (status: string) => void onToolPermissionRequest?: (toolName: string, input: any) => void @@ -39,7 +40,8 @@ export async function* runAgentLoop({ abortController, additionalSystemPrompt, config, - connectedServers, + inferredServers, + mcpServers, messageQueue, onServerConnection, onToolPermissionRequest, @@ -58,10 +60,10 @@ export async function* runAgentLoop({ const disallowedTools = getDisallowedTools(config) const enabledMcpServers = getEnabledMcpServers(config.mcpServers) - const { mcpServers } = await selectMcpServers({ + const inferenceResult = await inferMcpServers({ abortController, agents: config.agents, - connectedServers, + inferredServers, enabledMcpServers, onServerConnection, sessionId, @@ -71,7 +73,8 @@ export async function* runAgentLoop({ const systemPrompt = await buildSystemPrompt({ additionalSystemPrompt, config, - connectedServers, + inferredServers, + mcpServers, }) const agents = await createSDKAgents(config.agents) @@ -84,7 +87,7 @@ export async function* runAgentLoop({ canUseTool, disallowedTools, includePartialMessages: config.stream ?? false, - mcpServers, + mcpServers: inferenceResult.mcpServers, model: config.model ?? "haiku", permissionMode: config.permissionMode ?? "default", resume: sessionId,