diff --git a/.env.example b/.env.example index 0aa0b11..20636d4 100644 --- a/.env.example +++ b/.env.example @@ -1,5 +1,8 @@ -# Anthropic API Key +# Required ANTHROPIC_API_KEY= -# GitHub Access Token +ENABLE_LOGGING=false + +# Generate at https://github.com/settings/personal-access-tokens GITHUB_ACCESS_TOKEN= + diff --git a/.prettierignore b/.prettierignore new file mode 100644 index 0000000..2defb5c --- /dev/null +++ b/.prettierignore @@ -0,0 +1,5 @@ +node_modules/ +dist/ +*.env +*.env.example +.prettierignore diff --git a/CLAUDE.md b/CLAUDE.md index 1589474..3560bf0 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -1,5 +1,3 @@ ---- - Default to using Bun instead of Node.js. - Use `bun ` instead of `node ` or `ts-node ` @@ -17,7 +15,12 @@ Default to using Bun instead of Node.js. - Prefer named exports vs default exports - Prefer try/catch over vanilla promises - If a function has two or more arguments, use an object. And use an interface, vs adding types inline. -- Never use implicit returns in react components. They should ONLY be used in one-liners.- +- Never use implicit returns in react components. They should ONLY be used in one-liners. + +## Common Developer Commands + +- `bun test` +- `bun type-check` ## APIs diff --git a/README.md b/README.md index 72d1072..9a129f3 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,10 @@ -# Agent Chat CLI +# Agent Chat Cli -A bare-bones, terminal-based chat CLI built to explore the new [Claude Agent SDK](https://docs.claude.com/en/api/agent-sdk/overview). Terminal rendering is built on top of [React Ink](https://github.com/vadimdemedes/ink). +A minimalist, terminal-based chat CLI built to explore the new [Claude Agent SDK](https://docs.claude.com/en/api/agent-sdk/overview) and based on [damassi/agent-chat-cli](https://github.com/damassi/agent-chat-cli). Terminal rendering is built on top of [React Ink](https://github.com/vadimdemedes/ink). + +Additionally, via inference, Agent Chat CLI supports lazy, turn-based MCP connections to keep token costs down. The agent will only use those MCP servers you ask about, limiting the context that is sent up to the LLM. (After an MCP server is connected it remains connected, however.) + +## Overview The app has three modes: @@ -12,7 +16,7 @@ The agent, including MCP server setup, is configured in [agent-chat-cli.config.t The MCP _client_ is configured in [mcp-client.config.ts](mcp-client.config.ts). -https://github.com/user-attachments/assets/c2026c47-c798-4a1d-a68a-54e4abe73c63 +https://github.com/user-attachments/assets/f9a82631-ee26-4a7b-9d89-a732d2605513 ### Why? @@ -49,6 +53,18 @@ OAuth support works out of the box via `mcp-remote`: See the config above for an example. +### Example MCP Servers + +For demonstration purposes, Agent is configured with the following MCP servers: + +- **Chrome DevTools MCP**: https://developer.chrome.com/blog/chrome-devtools-mcp +- **Github MCP**: https://github.com/github/github-mcp-server + - [Generate a Github PAT token](https://github.com/settings/personal-access-tokens) +- **Notion MCP**: https://developers.notion.com/docs/mcp + - Authenticate via OAuth, which will launch a browser when attempting to connect + +**Note**: OAuth-based MCP servers (Notion, JIRA, etc) require browser-based authentication and cannot be deployed remotely. These servers are only accessible in the CLI version of the agent. + ### Usage #### Interactive Agent Mode @@ -80,11 +96,11 @@ Configure the MCP server connection in `mcp-client.config.ts`. HTTP is also supp Run as a stand-alone MCP server, using one of two modes: ```bash -bun server -bun server:http +bun server:http # streaming HTTP (use this for deployments) +bun server # stdio ``` -The server exposes an `ask_agent` tool that other MCP clients can use to interact with the agent. The agent has access to all configured MCP servers and can use their tools. +The MCP server exposes an `ask_agent` and `ask_agent_slackbot` tools that other MCP clients can use to interact with the agent. The agent has access to all configured MCP servers and can use their tools. ### Configuration @@ -96,10 +112,10 @@ To add specific instructions for each MCP server, create a markdown file in `src ```ts const config = { - systemPrompt: "You are a helpful agent." + systemPrompt: getPrompt("system.md"), mcpServers: { someMcpServer: { - command: "npx", + command: "bunx", args: ["..."], prompt: getPrompt("someMcpServer.md"), }, @@ -107,20 +123,130 @@ const config = { } ``` +#### Remote Prompts + +Prompts can be loaded from remote sources (e.g., APIs) using `getRemotePrompt`. This enables dynamic prompt management where prompts are stored in a database or CMS rather than in files. + +Both `getPrompt` (for local files) and `getRemotePrompt` (for API calls) return lazy functions that are only evaluated when the agent needs them, ensuring prompts are fetched on-demand during each LLM turn, enabling iteration in real time. + +```ts +import { getRemotePrompt } from "./src/utils/getRemotePrompt" + +const config = { + systemPrompt: getRemotePrompt(), + mcpServers: { + someMcpServer: { + command: "bunx", + args: ["..."], + prompt: getRemotePrompt({ + fetchPrompt: async () => { + const response = await fetch("https://some-prompt/name") + + if (!response.ok) { + throw new Error( + `[agent] [getRemotePrompt] [ERROR HTTP] status: ${response.status}` + ) + } + + const text = await response.text() + return text + }, + }), + }, + }, +} +``` + +You can also provide a fallback to a local file if the remote fetch fails: + +```ts +const config = { + mcpServers: { + github: { + prompt: getRemotePrompt({ + fallback: "github.md" + fetchPrompt: ... + }), + }, + }, +} +``` + #### Denying Tools -You can prevent specific MCP tools from being used by adding a `denyTools` array to your server configuration: +You can limit what tools the claude-agent-sdk has access to by adding a `disallowedTools` config: + +```ts +const config = { + disallowedTools: ["Bash"], +} +``` + +You can also prevent specific MCP tools from being used by adding a `disallowedTools` array to your server configuration: ```ts const config = { mcpServers: { github: { - command: "npx", + command: "bunx", args: ["..."], - denyTools: ["delete_repository", "update_secrets"], + disallowedTools: ["delete_repository", "update_secrets"], }, }, } ``` Denied tools are filtered at the SDK level and won't be available to the agent. + +In CLI mode, if `permissionMode` is set to "ask" then a prompt will appear to confirm when tools need to be invoked. + +### Specialized Subagents + +You can define specialized subagents in `agent-chat-cli.config.ts` to handle domain-specific tasks, leveraging the powerful [Claude Subagent SDK](https://docs.claude.com/en/docs/claude-code/sub-agents). Subagents are automatically invoked when user queries match their domain, and they have access to specific MCP servers. + +#### Example + +```ts +import { createAgent } from "./src/utils/createAgent" +import { getPrompt } from "./src/utils/getPrompt" + +const config = { + agents: { + "sales-partner-sentiment-agent": createAgent({ + description: + "An expert SalesForce partner sentiment agent, designed to produce insights for renewal and churn conversations", + prompt: getPrompt("agents/sales-partner-sentiment-agent.md"), + mcpServers: ["salesforce"], + }), + }, + mcpServers: { + salesforce: { + description: "Salesforce CRM: leads, opportunities, accounts...", + command: "bunx", + args: ["-y", "@tsmztech/mcp-server-salesforce@0.0.3"], + enabled: true, + }, + }, +} +``` + +When a user asks something like "Analyze partner churn", the routing agent will: + +1. Match the query to the `sales-partner-sentiment-agent` based on its description +2. Automatically connect to the required `salesforce` MCP server +3. Invoke the subagent with its specialized prompt and tools + +The `description` field is **critical**; it's used by the routing agent to determine when to invoke the subagent. + +**Note:** Subagents also support remote prompts via `getRemotePrompt`, allowing you to manage agent prompts dynamically from an API or database. + +### Note on Lazy MCP Server Initialization + +In order to keep LLM costs low and response times quick, a specialized sub-agent sits in front of user queries to infer which MCP servers are needed; the result is then forwarded on to the main agent, lazily initializing required MCP servers. Without this, we would need to initialize _all_ MCP servers defined in the config upfront, and for every query that we send to Anthropic, we'd _also_ be sending along a huge system prompt, and this is very expensive! + +#### The Flow + +- User sends a message, something like "In Salesforce, tell me about some recent leads" +- Sub-agent forwards message onto Anthropic's light-weight Haiku model and asks which MCP servers seem to be necessary +- Returns result as JSON, and based on the result, mcpServers are passed to the main agent query +- Agent now boots quickly and responds in a timely way, vs having to wait for every MCP server to initialize before being able to chat diff --git a/agent-chat-cli.config.ts b/agent-chat-cli.config.ts index 14b7422..5952452 100644 --- a/agent-chat-cli.config.ts +++ b/agent-chat-cli.config.ts @@ -1,31 +1,80 @@ import type { AgentChatConfig } from "./src/store" +import { createAgent } from "./src/utils/createAgent" import { getPrompt } from "./src/utils/getPrompt" const config: AgentChatConfig = { systemPrompt: getPrompt("system.md"), + model: "sonnet", + stream: true, + + agents: { + "demo-agent": createAgent({ + description: "A claude subagent designed to show off functionality", + prompt: getPrompt("agents/demo-agent.md"), + mcpServers: [], + }), + }, + mcpServers: { + chrome: { + description: + "The Chrome DevTools MCP server adds web browser automation and debugging capabilities to your AI agent", + command: "bunx", + args: ["chrome-devtools-mcp@latest"], + }, github: { + description: + "GitHub MCP tools to search code, PRs, issues; discover documentation in repo docs/; find deployment guides and code examples.", prompt: getPrompt("github.md"), - command: "npx", + command: "bunx", args: [ "mcp-remote@0.1.29", "https://api.githubcopilot.com/mcp/readonly", "--header", `Authorization: Bearer ${process.env.GITHUB_ACCESS_TOKEN}`, ], - env: { - GITHUB_ACCESS_TOKEN: process.env.GITHUB_ACCESS_TOKEN!, - }, - denyTools: [], + disallowedTools: [], + enabled: true, }, + notion: { + description: + "Notion workspace for documentation, wikis, OKRs, department pages, onboarding guides. Navigate hierarchies, search pages, retrieve structured content.", prompt: getPrompt("notion.md"), - command: "npx", + command: "bunx", args: ["mcp-remote@0.1.29", "https://mcp.notion.com/mcp"], + enabled: true, + }, + + /** + // Example of how to use getRemotePrompt + + someOtherServer: { + description: "Some description", + command: "bunx", + args: ["mcp-remote@0.1.29", "https://mcp.some-server.com/mcp"], + prompt: getRemotePrompt({ + fetchPrompt: async () => { + const response = await fetch("https://some-prompt/name") + + if (!response.ok) { + throw new Error( + `[agent] [getRemotePrompt] [ERROR HTTP] status: ${response.status}` + ) + } + + const text = await response.text() + return text + }, + }), }, + */ }, + + disallowedTools: ["Bash"], + + // Ungate MCP tools, which have their own disallowedTools array. permissionMode: "bypassPermissions", - stream: false, } export default config diff --git a/bun.lock b/bun.lock index 275d1bb..aa73a8a 100644 --- a/bun.lock +++ b/bun.lock @@ -2,9 +2,9 @@ "lockfileVersion": 1, "workspaces": { "": { - "name": "artsy-agent-claude", + "name": "agent-claude", "dependencies": { - "@anthropic-ai/claude-agent-sdk": "^0.1.1", + "@anthropic-ai/claude-agent-sdk": "^0.1.30", "@modelcontextprotocol/sdk": "^1.18.2", "cors": "^2.8.5", "cosmiconfig": "^9.0.0", @@ -35,7 +35,7 @@ "packages": { "@alcalzone/ansi-tokenize": ["@alcalzone/ansi-tokenize@0.2.0", "", { "dependencies": { "ansi-styles": "^6.2.1", "is-fullwidth-code-point": "^5.0.0" } }, "sha512-qI/5TaaaCZE4yeSZ83lu0+xi1r88JSxUjnH4OP/iZF7+KKZ75u3ee5isd0LxX+6N8U0npL61YrpbthILHB6BnA=="], - "@anthropic-ai/claude-agent-sdk": ["@anthropic-ai/claude-agent-sdk@0.1.1", "", { "optionalDependencies": { "@img/sharp-darwin-arm64": "^0.33.5", "@img/sharp-darwin-x64": "^0.33.5", "@img/sharp-linux-arm": "^0.33.5", "@img/sharp-linux-arm64": "^0.33.5", "@img/sharp-linux-x64": "^0.33.5", "@img/sharp-win32-x64": "^0.33.5" } }, "sha512-+12GQktMFc5Uqz6oVjJbj7Q+GD5QDorKEKtInALKD7VleJwLlFbMYIlm4586owIV5veFvb6bAVofKn9CnYWtvw=="], + "@anthropic-ai/claude-agent-sdk": ["@anthropic-ai/claude-agent-sdk@0.1.30", "", { "optionalDependencies": { "@img/sharp-darwin-arm64": "^0.33.5", "@img/sharp-darwin-x64": "^0.33.5", "@img/sharp-linux-arm": "^0.33.5", "@img/sharp-linux-arm64": "^0.33.5", "@img/sharp-linux-x64": "^0.33.5", "@img/sharp-win32-x64": "^0.33.5" }, "peerDependencies": { "zod": "^3.24.1" } }, "sha512-lo1tqxCr2vygagFp6kUMHKSN6AAWlULCskwGKtLB/JcIXy/8H8GsLSKX54anTsvc9mBbCR8wWASdFmiiL9NSKA=="], "@babel/code-frame": ["@babel/code-frame@7.27.1", "", { "dependencies": { "@babel/helper-validator-identifier": "^7.27.1", "js-tokens": "^4.0.0", "picocolors": "^1.1.1" } }, "sha512-cjQ7ZlQ0Mv3b47hABuTevyTuYN4i+loJKGeV9flcCgIK37cCXRh+L1bd3iBHlynerhQ7BhCkn2BPbQUL+rGqFg=="], diff --git a/docs/architecture.md b/docs/architecture.md index f79a268..88b1a6e 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -39,21 +39,29 @@ #### Agent Integration -- [src/utils/runAgentLoop.ts](../src/utils/runAgentLoop.ts) - Shared agent logic: - - `runAgentLoop()` - Creates agent query with configuration - - `startConversation()` - Async generator for message queue +- [src/utils/runAgentLoop.ts](../src/utils/runAgentLoop.ts) - Resume-based agent loop: + - `runAgentLoop()` - Async function that returns a conversation generator + - Implements turn-by-turn query loop with dynamic MCP server selection + - Maintains `connectedServers` Set across turns + - Uses SDK's `resume` to preserve conversation history + - Yields messages from each turn's query - `messageTypes` - Constants for message type checking - [src/hooks/useAgent.ts](../src/hooks/useAgent.ts) - React hook for interactive agent mode: - - Uses shared agent logic - - Manages the agent SDK query loop + - Calls `runAgentLoop()` and processes the conversation generator - Handles streaming responses - Processes tool uses - Tracks session state + - Manages server connection notifications via `onServerConnection` callback - Updates UI store - [src/hooks/useMcpClient.ts](../src/hooks/useMcpClient.ts) - React hook for MCP client mode: - Connects to an MCP server via stdio/HTTP/SSE - - Calls `ask_agent` tool on the server - - Displays tool uses via logging notifications + - Calls `get_agent_status` on initialization to fetch available servers + - Calls `ask_agent` tool with 10-minute timeout for user queries + - Handles logging notifications: + - `system_message`: Server connection status + - `text_message`: Streaming responses + - `tool_use`: Tool invocations + - Accumulates response text and properly clears between queries - Updates UI store with responses - Configured via `mcp-client.config.ts` @@ -92,22 +100,123 @@ Configurable streaming responses via `config.stream`. The CLI can connect to external MCP servers as a client: - Multiple MCP servers can be configured -- Per-server custom system prompts +- Per-server custom system prompts and descriptions - Server status tracking - Tool use from any connected server - Configured via `mcpServers` in config file +#### Dynamic MCP Server Selection + +The agent uses an intelligent routing system to load MCP servers on-demand instead of connecting to all servers upfront: + +**Key Components:** + +- [src/utils/mcpServerSelectionAgent.ts](../src/utils/mcpServerSelectionAgent.ts) - MCP server routing agent + - Uses Claude Agent SDK's custom tool system (`createSdkMcpServer`) + - Defines `select_mcp_servers` tool with Zod schema for structured output + - Analyzes user message + server descriptions to determine needed servers + - Case-insensitive server name matching + - Tracks already-connected servers to avoid redundant connections + - Returns both accumulated servers and newly selected servers + +- [src/utils/getEnabledMcpServers.ts](../src/utils/getEnabledMcpServers.ts) - Server filtering utility + - Filters MCP servers where `enabled !== false` + - Shared by `runAgentLoop` and store's computed properties + +- [src/utils/logger.ts](../src/utils/logger.ts) - Logging configuration + - Exports `enableLogging` constant based on `ENABLE_LOGGING` env var + - Centralized logging control for all routing logs + +**Configuration:** + +Each MCP server in `agent-chat-cli.config.ts` includes: + +```typescript +{ + description: string // Used by routing agent for intelligent matching + enabled: boolean // Filter servers before routing + prompt?: string // Optional system prompt + // ... other config +} +``` + +**Resume-Based Conversation Loop:** + +Instead of running one continuous `query()`, the agent runs a loop where each user message: + +1. Waits for user input from message queue +2. Calls `selectMcpServers()` with: + - Current user message + - Enabled MCP servers list + - Set of already-connected servers +3. Routing agent uses custom tool to return needed servers +4. New servers are accumulated into `connectedServers` Set +5. `query()` runs with: + - `resume: sessionId` (preserves conversation history) + - `mcpServers`: all accumulated servers +6. Processes messages until RESULT, then loops back + +**Server Accumulation:** + +Servers persist across conversation turns: + +- Turn 1: User mentions "github" → Connects to `[github]` +- Turn 2: User mentions "sentry" → Connects to `[github, sentry]` +- Turn 3: User mentions "notion" → Connects to `[github, sentry, notion]` + +The SDK's `resume` functionality maintains full conversation context while servers are dynamically added. + +**User Notifications:** + +System messages appear in chat when new servers connect: + +``` +[system] Connecting to github... +[system] Connecting to notion, salesforce... +``` + +**Routing Intelligence:** + +The routing agent matches servers based on: + +- Explicit mentions: "using github" → github +- Capability inference: "show me docs on OKRs" → notion +- Multiple servers: "query sales data" → redshift, salesforce +- No match: "what's the weather?" → no servers + #### MCP Client Flow -1. MCP client calls `query_agent` tool with prompt +**Initialization:** + +1. MCP client connects to server via stdio/HTTP/SSE transport +2. Client calls `get_agent_status` tool on initialization +3. Server's `getAgentStatus()` runs `runAgentLoop()` to get available MCP servers +4. Available servers list sent to client and displayed in header + +**User Query Flow:** + +1. User sends message → client calls `ask_agent` tool with query 2. Zod validates input schema -3. `runQuery()` creates a message queue -4. `runAgentLoop()` initializes agent with shared logic -5. Prompt is resolved into the message queue -6. Agent processes via async generator -7. Response messages are collected -8. Session state persists for follow-up queries -9. Final text response returned to MCP client +3. `runStandaloneAgentLoop()` creates message queue and calls `runAgentLoop()` +4. **Dynamic Server Selection** (same as interactive mode): + - User message triggers `selectMcpServers()` routing agent + - Routing agent analyzes message against server descriptions + - New servers accumulated into Set across conversation turns + - System message notification sent via `onServerConnection` callback +5. Server emits logging notifications for: + - `system_message`: Server connection status (`[system] Connecting to...`) + - `text_message`: Streaming assistant responses + - `tool_use`: Tool invocations with parameters +6. Client receives notifications and updates UI in real-time +7. Session state persists for follow-up queries +8. Final text response returned to MCP client + +**Key Implementation Details:** + +- `runStandaloneAgentLoop()` passes `onServerConnection` callback to `runAgentLoop()` +- Server connection notifications sent as MCP logging messages +- Client accumulates response text from `text_message` notifications +- Response properly added to chat history and display cleared between queries #### MCP Server Mode @@ -115,15 +224,21 @@ The CLI can also run as an MCP server itself, exposing the agent as a tool to ot **Shared MCP Infrastructure:** -- [src/mcp/utils/getServer.ts](../src/mcp/utils/getServer.ts) - MCP server factory +- [src/mcp/utils/getMcpServer.ts](../src/mcp/utils/getMcpServer.ts) - MCP server factory - Creates `McpServer` instance - - Registers `query_agent` tool with Zod validation + - Registers `ask_agent` tool with Zod validation + - Registers `get_agent_status` tool for initialization - Shared by both stdio and HTTP modes -- [src/mcp/utils/runQuery.ts](../src/mcp/utils/runQuery.ts) - Query execution +- [src/mcp/utils/runStandaloneAgentLoop.ts](../src/mcp/utils/runStandaloneAgentLoop.ts) - Query execution - Uses shared `runAgentLoop()` from agent integration - Manages message queue and session state - - Processes streaming responses + - Processes streaming responses with logging notifications + - Implements `onServerConnection` callback for dynamic server selection - Returns final text response +- [src/mcp/utils/getAgentStatus.ts](../src/mcp/utils/getAgentStatus.ts) - Status utility + - Initializes agent to get available MCP servers + - Called by `get_agent_status` tool on client initialization + - Returns session ID and MCP servers list **Transport Implementations:** @@ -176,22 +291,44 @@ The CLI can also run as an MCP server itself, exposing the agent as a tool to ot 1. User submits input via TextInput 2. Input added to chat history and message queue -3. `useAgent` hook processes queue via async generator -4. Agent SDK processes message with MCP servers -5. Streaming events update current assistant message -6. Tool uses are tracked and displayed separately -7. Final result includes stats (cost, duration, turns) -8. UI updates reactively via easy-peasy store +3. `runAgentLoop` waits for message from queue +4. `selectMcpServers()` routing agent determines needed servers: + - Analyzes user message against server descriptions + - Returns servers to connect (case-insensitive matching) + - Tracks already-connected servers +5. New servers added to accumulated `connectedServers` Set +6. System message shown: "[system] Connecting to server1, server2..." +7. Agent SDK `query()` runs with: + - Resume session ID (preserves history) + - All accumulated MCP servers +8. Streaming events update current assistant message +9. Tool uses are tracked and displayed separately +10. Final RESULT message includes stats (cost, duration, turns) +11. Loop returns to step 3 for next user message +12. UI updates reactively via easy-peasy store #### MCP Client Mode -1. User submits input via TextInput -2. Input added to chat history and message queue -3. `useMcpClient` hook sends to MCP server's `ask_agent` tool -4. Server emits logging notifications for tool uses -5. Client receives and displays tool uses in real-time -6. Final response text added to chat history -7. UI updates reactively via easy-peasy store +1. **Initialization:** Client calls `get_agent_status` to fetch available MCP servers +2. User submits input via TextInput +3. Input added to chat history and message queue +4. `useMcpClient` hook sends to MCP server's `ask_agent` tool with 10-minute timeout +5. **Server-side Dynamic Selection:** Server runs same routing logic as interactive mode: + - `selectMcpServers()` analyzes message against server descriptions + - New servers accumulated into Set across turns + - Session preserved via `resume` for conversation continuity +6. Server emits logging notifications in real-time: + - `system_message`: Server connections (`[system] Connecting to...`) + - `text_message`: Streaming assistant responses + - `tool_use`: Tool invocations with parameters +7. Client receives notifications and updates UI: + - System messages added to chat history + - Text messages displayed via `currentAssistantMessage` + - Tool uses shown separately +8. After `ask_agent` completes: + - Accumulated response text added to chat history + - Display cleared for next query (`clearCurrentAssistantMessage()`) +9. UI updates reactively via easy-peasy store #### MCP Server Mode @@ -215,7 +352,7 @@ Uses `cosmiconfig` for flexible configuration loading: args: string[] env?: Record prompt?: string // Optional system prompt - denyTools?: string[] // Optional list of tools to deny + disallowedTools?: string[] // Optional list of tools to deny } } } @@ -223,7 +360,7 @@ Uses `cosmiconfig` for flexible configuration loading: #### Tool Filtering -The `denyTools` configuration allows blocking specific MCP tools: +The `disallowedTools` configuration allows blocking specific MCP tools: - Tool names are exact matches (wildcards not supported) - Denied tools are passed to the SDK as `disallowedTools` option diff --git a/package.json b/package.json index d5e84c5..61f175c 100644 --- a/package.json +++ b/package.json @@ -28,7 +28,7 @@ "typescript": "^5" }, "dependencies": { - "@anthropic-ai/claude-agent-sdk": "^0.1.1", + "@anthropic-ai/claude-agent-sdk": "^0.1.30", "@modelcontextprotocol/sdk": "^1.18.2", "cors": "^2.8.5", "cosmiconfig": "^9.0.0", diff --git a/src/__tests__/store.test.tsx b/src/__tests__/store.test.tsx index 68c4cce..a465eb4 100644 --- a/src/__tests__/store.test.tsx +++ b/src/__tests__/store.test.tsx @@ -48,7 +48,7 @@ describe("Store", () => { expect(getState().sessionId).toBeUndefined() expect(getState().stats).toBeUndefined() expect(getState().pendingToolPermission).toBeUndefined() - expect(getState().abortController).toBeUndefined() + expect(getState().abortController).toBeInstanceOf(AbortController) }) test("should have MessageQueue instance", () => { @@ -73,13 +73,14 @@ describe("Store", () => { "setPendingToolPermission", "setAbortController", "setConfig", - "setcurrentAssistantMessage", + "setCurrentAssistantMessage", "setCurrentToolUses", "setInput", "setIsProcessing", "setMcpServers", "setSessionId", "setStats", + "handleMcpServerStatus", ]) ) }) @@ -143,10 +144,10 @@ describe("Store", () => { }) describe("assistant message actions", () => { - test("setcurrentAssistantMessage should set message", () => { + test("setCurrentAssistantMessage should set message", () => { const { getState, actions } = setup() - actions.setcurrentAssistantMessage("Hello") + actions.setCurrentAssistantMessage("Hello") expect(getState().currentAssistantMessage).toBe("Hello") }) @@ -154,7 +155,7 @@ describe("Store", () => { test("appendCurrentAssistantMessage should append to current message", () => { const { getState, actions } = setup() - actions.setcurrentAssistantMessage("Hello") + actions.setCurrentAssistantMessage("Hello") actions.appendCurrentAssistantMessage(" world") expect(getState().currentAssistantMessage).toBe("Hello world") @@ -163,7 +164,7 @@ describe("Store", () => { test("clearCurrentAssistantMessage should clear message", () => { const { getState, actions } = setup() - actions.setcurrentAssistantMessage("Hello") + actions.setCurrentAssistantMessage("Hello") actions.clearCurrentAssistantMessage() expect(getState().currentAssistantMessage).toBe("") @@ -397,7 +398,7 @@ describe("Store", () => { role: "user", content: "Hello", }) - actions.setcurrentAssistantMessage("Assistant message") + actions.setCurrentAssistantMessage("Assistant message") actions.addToolUse({ type: "tool_use", name: "tool", @@ -414,4 +415,89 @@ describe("Store", () => { expect(getState().stats).toBeNull() }) }) + + describe("handleMcpServerStatus thunk", () => { + test("should set mcp servers", () => { + const { getState, actions } = setup() + + const servers: McpServerStatus[] = [ + { name: "github", status: "connected" }, + { name: "notion", status: "connected" }, + ] + + actions.handleMcpServerStatus(servers) + + expect(getState().mcpServers).toEqual(servers) + }) + + test("should add error message for failed servers", () => { + const { getState, actions } = setup() + + const servers: McpServerStatus[] = [ + { name: "github", status: "connected" }, + { name: "postgres", status: "failed" }, + { name: "redis", status: "failed" }, + ] + + actions.handleMcpServerStatus(servers) + + const chatHistory = getState().chatHistory + const errorMessage = chatHistory.find( + (entry) => + entry.type === "message" && + entry.role === "system" && + "content" in entry && + entry.content.includes("Failed to connect") + ) + + expect(errorMessage).toBeDefined() + expect(errorMessage?.type).toBe("message") + if (errorMessage?.type === "message") { + expect(errorMessage.content).toBe( + "[Error] Failed to connect to postgres, redis" + ) + } + }) + + test("should not add any messages when all servers connect successfully", () => { + const { getState, actions } = setup() + + const servers: McpServerStatus[] = [ + { name: "github", status: "connected" }, + { name: "notion", status: "connected" }, + ] + + actions.handleMcpServerStatus(servers) + + const chatHistory = getState().chatHistory + expect(chatHistory.length).toBe(0) + }) + + test("should not add message when no servers provided", () => { + const { getState, actions } = setup() + + actions.handleMcpServerStatus([]) + + const chatHistory = getState().chatHistory + expect(chatHistory.length).toBe(0) + }) + + test("should handle only failed servers", () => { + const { getState, actions } = setup() + + const servers: McpServerStatus[] = [ + { name: "postgres", status: "failed" }, + ] + + actions.handleMcpServerStatus(servers) + + const chatHistory = getState().chatHistory + expect(chatHistory.length).toBe(1) + expect(chatHistory[0]).toEqual({ + type: "message", + role: "system", + content: "[Error] Failed to connect to postgres", + }) + }) + }) }) diff --git a/src/components/BlinkCaret.tsx b/src/components/BlinkCaret.tsx index d6faacf..8679f77 100644 --- a/src/components/BlinkCaret.tsx +++ b/src/components/BlinkCaret.tsx @@ -11,7 +11,7 @@ interface BlinkCaretProps { export const BlinkCaret: React.FC = ({ color = "purple", - enabled = false, + enabled = true, interval = BLINK_INTERVAL, }) => { const [visible, setVisible] = useState(false) diff --git a/src/components/ChatHeader.tsx b/src/components/ChatHeader.tsx index e0515a8..137010b 100644 --- a/src/components/ChatHeader.tsx +++ b/src/components/ChatHeader.tsx @@ -1,9 +1,14 @@ import { Box, Text } from "ink" -import Spinner from "ink-spinner" import { AgentStore } from "store" export const ChatHeader: React.FC = () => { const store = AgentStore.useStoreState((state) => state) + const availableServers = AgentStore.useStoreState( + (state) => state.availableMcpServers + ) + const availableAgents = AgentStore.useStoreState( + (state) => state.availableAgents + ) return ( @@ -25,12 +30,19 @@ export const ChatHeader: React.FC = () => { ))} ) : ( - - - - - {" Connecting to MCP servers..."} - + <> + + Available MCP Servers: + {availableServers.join(", ")} + + + )} + + {availableAgents.length > 0 && ( + + Agents: + {availableAgents.join(", ")} + )} diff --git a/src/components/Markdown.tsx b/src/components/Markdown.tsx index 5af8942..abb3639 100644 --- a/src/components/Markdown.tsx +++ b/src/components/Markdown.tsx @@ -48,6 +48,50 @@ const renderInline = (tokens: any[]): React.ReactNode[] => { }) } +const renderListItem = ( + item: any, + itemIdx: number, + depth: number +): React.ReactNode => { + const indent = " ".repeat(depth) + const bullet = "• " + + return ( + + {item.tokens.map((t: any, tIdx: number) => { + if (t.type === "text" && t.tokens) { + return ( + + {indent} + {bullet} + {renderInline(t.tokens)} + + ) + } else if (t.type === "list") { + // Handle nested lists + return ( + + {t.items.map((nestedItem: any, nestedIdx: number) => + renderListItem(nestedItem, nestedIdx, depth + 1) + )} + + ) + } else if (t.type === "paragraph") { + // Handle paragraphs in list items + return ( + + {indent} + {bullet} + {renderInline(t.tokens)} + + ) + } + return null + })} + + ) +} + const renderToken = (token: any, idx: number): React.ReactNode => { switch (token.type) { case "paragraph": @@ -62,21 +106,7 @@ const renderToken = (token: any, idx: number): React.ReactNode => { return ( {token.items.map((item: any, itemIdx: number) => { - // List items contain nested tokens (usually paragraphs with text) - const content = item.tokens - .map((t: any) => { - if (t.type === "text" && t.tokens) { - return renderInline(t.tokens) - } - return null - }) - .filter(Boolean) - - return ( - - • {content} - - ) + return renderListItem(item, itemIdx, 0) })} ) diff --git a/src/components/UserInput.tsx b/src/components/UserInput.tsx index 254ebc2..fbf81e4 100644 --- a/src/components/UserInput.tsx +++ b/src/components/UserInput.tsx @@ -1,7 +1,7 @@ -import { Box } from "ink" -import TextInput from "ink-text-input" import { BlinkCaret } from "components/BlinkCaret" import { useCycleMessages } from "hooks/useCycleMessages" +import { Box } from "ink" +import TextInput from "ink-text-input" import { AgentStore } from "store" const commands = { @@ -40,7 +40,7 @@ export const UserInput: React.FC = () => { return ( - 0} /> + { @@ -26,14 +25,14 @@ describe("ChatHeader", () => { expect(lastFrame()).toContain("Type 'exit' to quit") }) - test("should show connecting message when no servers", () => { + test("should show available servers when no servers connected", () => { const { lastFrame } = render( ) - expect(lastFrame()).toContain("Connecting to MCP servers...") + expect(lastFrame()).toContain("Available MCP Servers:") }) test("should display connected servers", () => { diff --git a/src/components/__tests__/Markdown.test.tsx b/src/components/__tests__/Markdown.test.tsx index 191d4e3..fab0bdd 100644 --- a/src/components/__tests__/Markdown.test.tsx +++ b/src/components/__tests__/Markdown.test.tsx @@ -90,4 +90,23 @@ describe("Markdown", () => { expect(lastFrame()).toContain("italic") expect(lastFrame()).toContain("code") }) + + test("should render nested lists", () => { + const markdown = `1. **Configure AWS credentials** using one of these methods: + - Set the AWS_PROFILE environment variable + - Run aws sso login --profile + - Ensure credentials are in ~/.aws/credentials + +2. **Verify IAM permissions** - Your credentials need: + - redshift:DescribeClusters + - redshift-data:ExecuteStatement` + + const { lastFrame } = render({markdown}) + + expect(lastFrame()).toContain("Configure AWS credentials") + expect(lastFrame()).toContain("Set the AWS_PROFILE") + expect(lastFrame()).toContain("Run aws sso login") + expect(lastFrame()).toContain("Verify IAM permissions") + expect(lastFrame()).toContain("redshift:DescribeClusters") + }) }) diff --git a/src/components/__tests__/ToolUses.test.tsx b/src/components/__tests__/ToolUses.test.tsx index 973a09c..b378858 100644 --- a/src/components/__tests__/ToolUses.test.tsx +++ b/src/components/__tests__/ToolUses.test.tsx @@ -1,9 +1,8 @@ -import React from "react" -import { render } from "ink-testing-library" -import { test, expect, describe } from "bun:test" +import { describe, expect, test } from "bun:test" import { ToolUses } from "components/ToolUses" -import { AgentStore } from "store" +import { render } from "ink-testing-library" import type { ToolUse } from "store" +import { AgentStore } from "store" describe("ToolUses", () => { test("should display MCP tool with server name", () => { @@ -13,15 +12,10 @@ describe("ToolUses", () => { input: { query: "test" }, } - const { lastFrame, rerender } = render( - - - - ) - - rerender( + const config = { mcpServers: {} } + const { lastFrame } = render( - + ) @@ -92,7 +86,7 @@ describe("ToolUses", () => { github: { command: "node", args: [], - denyTools: ["search"], + disallowedTools: ["search"], }, }, }} @@ -109,7 +103,7 @@ describe("ToolUses", () => { github: { command: "node", args: [], - denyTools: ["search"], + disallowedTools: ["search"], }, }, }} @@ -136,7 +130,7 @@ describe("ToolUses", () => { github: { command: "node", args: [], - denyTools: ["search"], + disallowedTools: ["search"], }, }, }} @@ -153,7 +147,7 @@ describe("ToolUses", () => { github: { command: "node", args: [], - denyTools: ["search"], + disallowedTools: ["search"], }, }, }} @@ -219,6 +213,11 @@ const TestWrapper = ({ config: any }) => { const actions = AgentStore.useStoreActions((actions) => actions) - actions.setConfig(config) + const currentConfig = AgentStore.useStoreState((state) => state.config) + + if (!currentConfig) { + actions.setConfig(config) + } + return } diff --git a/src/hooks/useAgent.ts b/src/hooks/useAgent.ts index e83e697..83ab4c6 100644 --- a/src/hooks/useAgent.ts +++ b/src/hooks/useAgent.ts @@ -1,43 +1,50 @@ import { useEffect, useRef } from "react" import { AgentStore } from "store" -import { useMcpServers } from "hooks/useMcpServers" -import { runAgentLoop, messageTypes } from "utils/runAgentLoop" +import { messageTypes, runAgentLoop } from "utils/runAgentLoop" export function useAgent() { - const { initMcpServers } = useMcpServers() - const messageQueue = AgentStore.useStoreState((state) => state.messageQueue) const sessionId = AgentStore.useStoreState((state) => state.sessionId) const config = AgentStore.useStoreState((state) => state.config) + const abortController = AgentStore.useStoreState( + (state) => state.abortController + ) const actions = AgentStore.useStoreActions((actions) => actions) const currentAssistantMessageRef = useRef("") + const abortControllerRef = useRef(abortController) + + // Update ref when abort controller changes + abortControllerRef.current = abortController useEffect(() => { const streamEnabled = config.stream ?? false - const abortController = new AbortController() - actions.setAbortController(abortController) const runAgent = async () => { - const { response } = runAgentLoop({ + const { agentLoop } = await runAgentLoop({ messageQueue, sessionId, config, - abortController, + abortControllerRef, onToolPermissionRequest: (toolName, input) => { actions.setPendingToolPermission({ toolName, input }) }, + onServerConnection: (status) => { + actions.addChatHistoryEntry({ + type: "message", + role: "system", + content: status, + }) + }, setIsProcessing: actions.setIsProcessing, }) - await initMcpServers(response) - try { - for await (const message of response) { + for await (const message of agentLoop) { switch (true) { case message.type === messageTypes.SYSTEM && message.subtype === messageTypes.INIT: { actions.setSessionId(message.session_id) - actions.setMcpServers(message.mcp_servers) + actions.handleMcpServerStatus(message.mcp_servers) continue } @@ -118,7 +125,7 @@ export function useAgent() { }` ) } else { - actions.setStats(`[agent-chat-cli] Error: ${message.subtype}`) + actions.setStats(`[agent-cli] Error: ${message.subtype}`) } actions.setIsProcessing(false) break @@ -130,7 +137,7 @@ export function useAgent() { error instanceof Error && !error.message.includes("process aborted by user") ) { - actions.setStats(`[agent-chat-cli] ${error}`) + actions.setStats(`[agent-cli] ${error}`) } actions.setIsProcessing(false) @@ -138,9 +145,5 @@ export function useAgent() { } runAgent() - - return () => { - abortController.abort() - } }, []) } diff --git a/src/hooks/useConfig.ts b/src/hooks/useConfig.ts index 8e95e76..fc7cb86 100644 --- a/src/hooks/useConfig.ts +++ b/src/hooks/useConfig.ts @@ -14,7 +14,7 @@ export const useConfig = () => { const loadedConfig = await loadConfig() actions.setConfig(loadedConfig) } catch (error) { - console.error("[agent-chat-cli] Failed to load config:", error) + console.error("[agent-cli] Failed to load config:", error) process.exit(1) } } diff --git a/src/hooks/useMcpClient.ts b/src/hooks/useMcpClient.ts index 6332f90..ce794ac 100644 --- a/src/hooks/useMcpClient.ts +++ b/src/hooks/useMcpClient.ts @@ -1,4 +1,3 @@ -import { useEffect } from "react" import { Client } from "@modelcontextprotocol/sdk/client/index.js" import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js" import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js" @@ -6,6 +5,7 @@ import { LoggingMessageNotificationSchema, type LoggingMessageNotification, } from "@modelcontextprotocol/sdk/types.js" +import { useEffect } from "react" import { AgentStore } from "store" import config from "../../mcp-client.config" @@ -19,10 +19,11 @@ export const useMcpClient = () => { const client = new Client( { name: "agent-chat-cli-client", - version: "1.0.0", + version: "0.1.0", }, { capabilities: { + // Allow bot to ask for input elicitation: {}, }, } @@ -53,7 +54,19 @@ export const useMcpClient = () => { input: data.input as Record, }) } else if (data.type === "mcp_servers") { - actions.setMcpServers(data.servers) + actions.handleMcpServerStatus(data.servers) + } else if (data.type === "system_message") { + actions.addChatHistoryEntry({ + type: "message", + role: "system", + content: data.content, + }) + } else if (data.type === "text_message") { + actions.addChatHistoryEntry({ + type: "message", + role: "assistant", + content: data.content, + }) } } catch { // noop @@ -78,7 +91,7 @@ export const useMcpClient = () => { default: { throw new Error( - `[agent-chat-cli] Unsupported transport: ${config.transport}` + `[agent-cli] Unsupported transport: ${config.transport}` ) } } @@ -88,6 +101,7 @@ export const useMcpClient = () => { await client.connect(transport) await client.setLoggingLevel("debug") + await client.callTool({ name: "get_agent_status", arguments: {}, @@ -107,7 +121,7 @@ export const useMcpClient = () => { try { const startTime = Date.now() - const result = await client.callTool({ + await client.callTool({ name: "ask_agent", arguments: { query: userMessage, @@ -116,36 +130,22 @@ export const useMcpClient = () => { const duration = Date.now() - startTime - let responseText = "" - - if (result.content && Array.isArray(result.content)) { - for (const item of result.content) { - if (item.type === "text") { - responseText += item.text - } - } - } - - if (responseText) { - actions.addChatHistoryEntry({ - type: "message", - role: "assistant", - content: responseText, - }) - } + // Messages are already added via logging notifications during execution + // No need to add a final message here since text_message notifications + // already added each response as it came in actions.setStats(`Completed in ${(duration / 1000).toFixed(2)}s`) actions.setIsProcessing(false) } catch (error) { actions.setStats( - `[agent-chat-cli] Error: ${error instanceof Error ? error.message : String(error)}` + `[agent-cli] Error: ${error instanceof Error ? error.message : String(error)}` ) actions.setIsProcessing(false) } } } catch (error) { actions.setStats( - `[agent-chat-cli] Client error: ${error instanceof Error ? error.message : String(error)}` + `[agent-cli] Client error: ${error instanceof Error ? error.message : String(error)}` ) actions.setIsProcessing(false) } diff --git a/src/hooks/useMcpServers.ts b/src/hooks/useMcpServers.ts deleted file mode 100644 index 98baa81..0000000 --- a/src/hooks/useMcpServers.ts +++ /dev/null @@ -1,71 +0,0 @@ -import { AgentStore } from "store" - -const CONNECTION_TIMEOUT = 10000 -const MAX_RETRIES = 3 -const RETRY_DELAY = 2000 - -export const useMcpServers = () => { - const config = AgentStore.useStoreState((state) => state.config) - const actions = AgentStore.useStoreActions((actions) => actions) - - const initMcpServers = async ( - response: any, - retryCount = 0 - ): Promise => { - const connectionTimeout = config.connectionTimeout ?? CONNECTION_TIMEOUT - const maxRetries = config.maxRetries ?? MAX_RETRIES - const retryDelay = config.retryDelay ?? RETRY_DELAY - const mcpServers = config.mcpServers - - try { - const timeoutPromise = new Promise((_, reject) => - setTimeout( - () => reject(new Error("MCP server connection timeout")), - connectionTimeout - ) - ) - - const servers = await Promise.race([ - response.mcpServerStatus(), - timeoutPromise, - ]) - - actions.setMcpServers(servers) - } catch (error) { - const errorMessage = - error instanceof Error ? error.message : String(error) - - if (retryCount < maxRetries) { - const newRetryCount = retryCount + 1 - - actions.addChatHistoryEntry({ - type: "message", - role: "system", - content: `Failed to get MCP server status: ${errorMessage}. Retrying connection (${newRetryCount}/${maxRetries})...`, - }) - - await new Promise((resolve) => setTimeout(resolve, retryDelay)) - - // Retry - await initMcpServers(response, newRetryCount) - } else { - actions.addChatHistoryEntry({ - type: "message", - role: "system", - content: `Failed to connect to MCP servers after ${maxRetries} attempts: ${errorMessage}. Servers marked as disconnected.`, - }) - - const disconnectedServers = Object.keys(mcpServers).map((name) => ({ - name, - status: "disconnected", - })) - - actions.setMcpServers(disconnectedServers) - } - } - } - - return { - initMcpServers, - } -} diff --git a/src/index.tsx b/src/index.tsx index 19893b3..4f7658d 100644 --- a/src/index.tsx +++ b/src/index.tsx @@ -1,7 +1,7 @@ import { render } from "ink" -import { App } from "./App" -import { validateEnv } from "utils/validateEnv" import { AgentStore } from "store" +import { validateEnv } from "utils/validateEnv" +import { App } from "./App" const main = () => { validateEnv() @@ -18,5 +18,5 @@ const main = () => { try { main() } catch (error) { - console.error(error) + console.error("[agent-cli] [ERROR]:", error) } diff --git a/src/mcp/utils/getAgentStatus.ts b/src/mcp/getAgentStatus.ts similarity index 84% rename from src/mcp/utils/getAgentStatus.ts rename to src/mcp/getAgentStatus.ts index e553e40..eb04b6d 100644 --- a/src/mcp/utils/getAgentStatus.ts +++ b/src/mcp/getAgentStatus.ts @@ -1,13 +1,13 @@ import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js" import { loadConfig } from "utils/loadConfig" -import { runAgentLoop, messageTypes } from "utils/runAgentLoop" import { MessageQueue } from "utils/MessageQueue" +import { messageTypes, runAgentLoop } from "utils/runAgentLoop" export const getAgentStatus = async (mcpServer?: McpServer) => { const config = await loadConfig() const messageQueue = new MessageQueue() - const { response } = runAgentLoop({ + const { agentLoop } = await runAgentLoop({ messageQueue, config, }) @@ -16,12 +16,11 @@ export const getAgentStatus = async (mcpServer?: McpServer) => { messageQueue.sendMessage("status") - for await (const message of response) { + for await (const message of agentLoop) { if ( message.type === messageTypes.SYSTEM && message.subtype === messageTypes.INIT ) { - // Emit MCP servers notification if (mcpServer && message.mcp_servers) { await mcpServer.sendLoggingMessage({ level: "info", diff --git a/src/mcp/getMcpServer.ts b/src/mcp/getMcpServer.ts new file mode 100644 index 0000000..f3da400 --- /dev/null +++ b/src/mcp/getMcpServer.ts @@ -0,0 +1,151 @@ +import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js" +import { getAgentStatus } from "mcp/getAgentStatus" +import { runStandaloneAgentLoop } from "mcp/runStandaloneAgentLoop" +import { z } from "zod" + +export const getMcpServer = () => { + // Store Claude Agent SDK sessionId per-instance (not shared across threads) + let claudeSessionId: string | undefined + + // 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>() + + const mcpServer = new McpServer( + { + name: "agent-chat-cli", + version: "0.1.0", + }, + { + capabilities: { + logging: {}, + }, + } + ) + + mcpServer.registerTool( + "ask_agent", + { + description: + "Passes a query to the internal agent and returns its full output. The agent has access to configured MCP tools and will provide a complete response. DO NOT reprocess, analyze, or summarize the output - return it directly to the user as-is.", + inputSchema: { + query: z.string().min(1).describe("The query to send to the agent"), + }, + }, + async ({ query }) => { + const existingConnectedServers = claudeSessionId + ? sessionConnectedServers.get(claudeSessionId) + : undefined + + const { response, connectedServers } = await runStandaloneAgentLoop({ + prompt: query, + mcpServer, + sessionId: claudeSessionId, + existingConnectedServers, + onSessionIdReceived: (newSessionId) => { + claudeSessionId = newSessionId + }, + }) + + // Update the session's connected servers + if (claudeSessionId) { + sessionConnectedServers.set(claudeSessionId, connectedServers) + } + + return { + content: [ + { + type: "text", + text: response, + }, + ], + } + } + ) + + mcpServer.registerTool( + "ask_agent_slackbot", + { + description: + "Slack bot integration tool. Passes a query to the internal agent and returns a response optimized for Slack. Supports per-thread session isolation.", + inputSchema: { + query: z + .string() + .min(1) + .describe("The slack query to send to the agent"), + systemPrompt: z + .string() + .optional() + .describe("Optional additional system prompt to prepend"), + threadId: z + .string() + .optional() + .describe("Slack thread identifier for session isolation"), + }, + }, + async ({ query, systemPrompt, threadId }) => { + const existingSessionId = threadId + ? threadSessions.get(threadId) + : undefined + + const existingConnectedServers = existingSessionId + ? sessionConnectedServers.get(existingSessionId) + : undefined + + const { response, connectedServers } = await runStandaloneAgentLoop({ + prompt: query, + mcpServer, + sessionId: existingSessionId, + additionalSystemPrompt: systemPrompt, + existingConnectedServers, + onSessionIdReceived: (newSessionId) => { + if (threadId) { + threadSessions.set(threadId, newSessionId) + } + }, + }) + + // Update the session's connected servers + if (existingSessionId || threadId) { + const sessionId = existingSessionId || threadSessions.get(threadId!) + if (sessionId) { + sessionConnectedServers.set(sessionId, connectedServers) + } + } + + return { + content: [ + { + type: "text", + text: response, + }, + ], + } + } + ) + + mcpServer.registerTool( + "get_agent_status", + { + description: + "Get the status of the agent including which MCP servers it has access to. Call this on initialization to see available servers.", + inputSchema: {}, + }, + async () => { + const status = await getAgentStatus(mcpServer) + + return { + content: [ + { + type: "text", + text: JSON.stringify(status, null, 2), + }, + ], + } + } + ) + + return mcpServer +} diff --git a/src/mcp/http.ts b/src/mcp/http.ts index db67394..a3c9185 100644 --- a/src/mcp/http.ts +++ b/src/mcp/http.ts @@ -1,12 +1,13 @@ import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js" import { isInitializeRequest } from "@modelcontextprotocol/sdk/types.js" +import cors from "cors" +import express from "express" +import { getMcpServer } from "mcp/getMcpServer" import { randomUUID } from "node:crypto" import { loadConfig } from "utils/loadConfig" -import { getMcpServer } from "mcp/utils/getMcpServer" -import express from "express" -import cors from "cors" +import { log } from "utils/logger" -const PORT = 3000 +const PORT = 8080 export const main = async () => { await loadConfig() @@ -24,15 +25,17 @@ export const main = async () => { }) ) + app.get("/health", (_req, res) => { + res.json({ status: "ok" }) + }) + app.post("/mcp", async (req, res) => { const sessionId = req.headers["mcp-session-id"] as string | undefined if (sessionId) { - console.log( - `[agent-chat-cli] Received MCP request for session: ${sessionId}` - ) + log(`Received MCP request for session: ${sessionId}`) } else { - console.log("[agent-chat-cli] New MCP request") + log("New MCP request") } try { @@ -40,13 +43,24 @@ export const main = async () => { if (sessionId && transports[sessionId]) { transport = transports[sessionId] + } else if (sessionId && !transports[sessionId]) { + log(`[/mcp POST] Session not found: ${sessionId}`) + + // Per MCP spec + res.status(404).json({ + jsonrpc: "2.0", + error: { + code: -32000, + message: "Session not found", + }, + id: null, + }) + return } else if (!sessionId && isInitializeRequest(req.body)) { transport = new StreamableHTTPServerTransport({ sessionIdGenerator: () => randomUUID(), onsessioninitialized: (newSessionId) => { - console.log( - `[agent-chat-cli] Session initialized with ID: ${newSessionId}` - ) + log(`[/mcp POST] Session initialized with ID: ${newSessionId}`) transports[newSessionId] = transport }, @@ -56,13 +70,14 @@ export const main = async () => { const sid = transport.sessionId if (sid && transports[sid]) { - console.log(`[agent-chat-cli] Transport closed for session ${sid}`) + log(`[onclose] Transport closed for session ${sid}`) delete transports[sid] } } const server = getMcpServer() + // Connect await server.connect(transport) await transport.handleRequest(req, res, req.body) @@ -79,9 +94,10 @@ export const main = async () => { return } + // Req / Res handlers await transport.handleRequest(req, res, req.body) } catch (error) { - console.error("[agent-chat-cli] Error handling MCP request:", error) + console.error("[agent-cli] Error handling MCP request:", error) if (!res.headersSent) { res.status(500).json({ @@ -100,20 +116,26 @@ export const main = async () => { const sessionId = req.headers["mcp-session-id"] as string | undefined if (!sessionId || !transports[sessionId]) { - res.status(400).send("Invalid or missing session ID") + log(`[/mcp GET] Session not found: ${sessionId}`) + + res.status(404).json({ + jsonrpc: "2.0", + error: { + code: -32000, + message: "Session not found", + }, + id: null, + }) + return } const lastEventId = req.headers["last-event-id"] if (lastEventId) { - console.log( - `[agent-chat-cli] Client reconnecting with Last-Event-ID: ${lastEventId}` - ) + log(`[/mcp GET] Client reconnecting with Last-Event-ID: ${lastEventId}`) } else { - console.log( - `[agent-chat-cli] Establishing new HTTP stream for session ${sessionId}` - ) + log(`[/mcp GET] Establishing new HTTP stream for session ${sessionId}`) } const transport = transports[sessionId] @@ -124,20 +146,28 @@ export const main = async () => { const sessionId = req.headers["mcp-session-id"] as string | undefined if (!sessionId || !transports[sessionId]) { - res.status(400).send("Invalid or missing session ID") + log(`[/mcp DELETE] Session not found: ${sessionId}`) + + res.status(404).json({ + jsonrpc: "2.0", + error: { + code: -32000, + message: "Session not found", + }, + id: null, + }) + return } - console.log( - `[agent-chat-cli] Received session termination request for session ${sessionId}` - ) + log(`[/mcp DELETE] Deleting session: ${sessionId}`) try { const transport = transports[sessionId] await transport.handleRequest(req, res) } catch (error) { - console.error( - "[agent-chat-cli] Error handling session termination:", + console.log( + "[agent-cli] [/mcp DELETE] [ERROR] Error deleting session:", error ) @@ -149,30 +179,28 @@ export const main = async () => { app.listen(PORT, () => { console.log( - `\n[agent-chat-cli] MCP HTTP Server running on port http://localhost:${PORT}\n` + `\n[agent-cli] MCP HTTP Server running on port http://localhost:${PORT}\n` ) }) process.on("SIGINT", async () => { - console.log("[agent-chat-cli] Shutting down server...") + console.log("[agent-cli] Shutting down server...") for (const sessionId in transports) { try { - console.log( - `[agent-chat-cli] Closing transport for session ${sessionId}` - ) + console.log(`[agent-cli] Closing transport for session ${sessionId}`) await transports[sessionId]?.close() delete transports[sessionId] } catch (error) { console.error( - `[agent-chat-cli] Error closing transport for session ${sessionId}:`, + `[agent-cli] Error closing transport for session ${sessionId}:`, error ) } } - console.log("[agent-chat-cli] Server shutdown complete") + console.log("[agent-cli] Server shutdown complete.") process.exit(0) }) } @@ -180,6 +208,6 @@ export const main = async () => { try { main() } catch (error) { - console.error("[agent-chat-cli] Fatal error starting HTTP server:", error) + console.error("[agent-cli] Fatal error starting HTTP server:", error) process.exit(1) } diff --git a/src/mcp/runStandaloneAgentLoop.ts b/src/mcp/runStandaloneAgentLoop.ts new file mode 100644 index 0000000..2836f31 --- /dev/null +++ b/src/mcp/runStandaloneAgentLoop.ts @@ -0,0 +1,170 @@ +import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js" +import { loadConfig } from "utils/loadConfig" +import { log } from "utils/logger" +import { MessageQueue } from "utils/MessageQueue" +import { contentTypes, messageTypes, runAgentLoop } from "utils/runAgentLoop" + +interface RunQueryOptions { + additionalSystemPrompt?: string + existingConnectedServers?: Set + mcpServer: McpServer + onSessionIdReceived?: (sessionId: string) => void + prompt: string + sessionId?: string +} + +export const runStandaloneAgentLoop = async ({ + additionalSystemPrompt, + existingConnectedServers, + mcpServer, + onSessionIdReceived, + prompt, + sessionId, +}: RunQueryOptions) => { + const config = await loadConfig() + const messageQueue = new MessageQueue() + const streamEnabled = config.stream ?? false + + const { agentLoop, connectedServers } = await runAgentLoop({ + additionalSystemPrompt, + config, + existingConnectedServers, + messageQueue, + sessionId, + onServerConnection: async (status) => { + await mcpServer.sendLoggingMessage({ + level: "info", + data: JSON.stringify({ + type: "system_message", + content: status, + }), + }) + }, + }) + + await new Promise((resolve) => setTimeout(resolve, 0)) + + messageQueue.sendMessage(prompt) + + let finalResponse = "" + let assistantMessage = "" + + try { + for await (const message of agentLoop) { + switch (true) { + case message.type === messageTypes.SYSTEM && + message.subtype === messageTypes.INIT: { + log( + "[runStandaloneAgentLoop] [messageTypes.INIT]:\n", + JSON.stringify(message, null, 2) + ) + + if (onSessionIdReceived) { + onSessionIdReceived(message.session_id) + } + continue + } + + case message.type === messageTypes.STREAM_EVENT: { + if (streamEnabled) { + const event = message.event + + if (event.type === "content_block_delta") { + if (event.delta.type === "text_delta") { + assistantMessage += event.delta.text + finalResponse += event.delta.text + } + } else if ( + event.type === "content_block_stop" && + assistantMessage + ) { + // Flush text when content block ends (before potential tool use) + await mcpServer.sendLoggingMessage({ + level: "info", + data: JSON.stringify({ + type: "text_message", + content: assistantMessage, + }), + }) + assistantMessage = "" + } + } + continue + } + + case message.type === messageTypes.ASSISTANT: { + for (const content of message.message.content) { + log( + "[runStandaloneAgentLoop] [messageTypes.ASSISTANT]:\n", + JSON.stringify(content, null, 2) + ) + + switch (true) { + case content.type === contentTypes.TEXT: { + if (!streamEnabled) { + assistantMessage += content.text + finalResponse += content.text + } + break + } + + case content.type === contentTypes.TOOL_USE: { + if (assistantMessage) { + await mcpServer.sendLoggingMessage({ + level: "info", + data: JSON.stringify({ + type: "text_message", + content: assistantMessage, + }), + }) + + // Flush after tool use + assistantMessage = "" + } + + await mcpServer.sendLoggingMessage({ + level: "info", + data: JSON.stringify({ + type: "tool_use", + name: content.name, + input: content.input, + }), + }) + + break + } + } + } + + break + } + + case message.type === messageTypes.RESULT: { + if (!streamEnabled) { + if (assistantMessage) { + await mcpServer.sendLoggingMessage({ + level: "info", + data: JSON.stringify({ + type: "text_message", + content: assistantMessage, + }), + }) + } + } + + return { + response: finalResponse, + connectedServers, + } + } + } + } + } catch (error) { + console.error(`[agent-cli] [runStandaloneAgentLoop] Error: ${error}`) + } + + return { + response: finalResponse, + connectedServers, + } +} diff --git a/src/mcp/stdio.ts b/src/mcp/stdio.ts index f2df275..fbcdf99 100644 --- a/src/mcp/stdio.ts +++ b/src/mcp/stdio.ts @@ -1,6 +1,6 @@ import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js" import { loadConfig } from "utils/loadConfig" -import { getMcpServer } from "mcp/utils/getMcpServer" +import { getMcpServer } from "mcp/getMcpServer" export const main = async () => { try { @@ -10,9 +10,9 @@ export const main = async () => { const transport = new StdioServerTransport() await mcpServer.connect(transport) - console.log("\n[agent-chat-cli] MCP Server running on stdio\n") + console.log("\n[agent-cli] MCP Server running on stdio\n") } catch (error) { - console.error("[agent-chat-cli] Fatal error running server:", error) + console.error("[agent-cli] Fatal error running server:", error) process.exit(1) } } diff --git a/src/mcp/utils/getMcpServer.ts b/src/mcp/utils/getMcpServer.ts deleted file mode 100644 index 8a8cd3c..0000000 --- a/src/mcp/utils/getMcpServer.ts +++ /dev/null @@ -1,67 +0,0 @@ -import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js" -import { getAgentStatus } from "mcp/utils/getAgentStatus" -import { runStandaloneAgentLoop } from "mcp/utils/runStandaloneAgentLoop" -import { z } from "zod" - -export const getMcpServer = () => { - const mcpServer = new McpServer( - { - name: "agent-chat-cli", - version: "0.1.0", - }, - { - capabilities: { - logging: {}, - }, - } - ) - - mcpServer.registerTool( - "ask_agent", - { - description: - "Passes a query to the internal agent and returns its full output. The agent has access to configured MCP tools and will provide a complete response. DO NOT reprocess, analyze, or summarize the output - return it directly to the user as-is.", - inputSchema: { - query: z.string().min(1).describe("The query to send to the agent"), - }, - }, - async ({ query }) => { - const result = await runStandaloneAgentLoop({ - prompt: query, - mcpServer, - }) - - return { - content: [ - { - type: "text", - text: result, - }, - ], - } - } - ) - - mcpServer.registerTool( - "get_agent_status", - { - description: - "Get the status of the agent including which MCP servers it has access to. Call this on initialization to see available servers.", - inputSchema: {}, - }, - async () => { - const status = await getAgentStatus(mcpServer) - - return { - content: [ - { - type: "text", - text: JSON.stringify(status, null, 2), - }, - ], - } - } - ) - - return mcpServer -} diff --git a/src/mcp/utils/runStandaloneAgentLoop.ts b/src/mcp/utils/runStandaloneAgentLoop.ts deleted file mode 100644 index a9860b8..0000000 --- a/src/mcp/utils/runStandaloneAgentLoop.ts +++ /dev/null @@ -1,92 +0,0 @@ -import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js" -import { loadConfig } from "utils/loadConfig" -import { runAgentLoop, messageTypes } from "utils/runAgentLoop" -import { MessageQueue } from "utils/MessageQueue" - -let sessionId: string | undefined - -interface RunQueryOptions { - prompt: string - mcpServer?: McpServer -} - -export const runStandaloneAgentLoop = async ({ - prompt, - mcpServer, -}: RunQueryOptions) => { - const config = await loadConfig() - const messageQueue = new MessageQueue() - const streamEnabled = config.stream ?? false - - const { response } = runAgentLoop({ - messageQueue, - sessionId, - config, - }) - - await new Promise((resolve) => setTimeout(resolve, 0)) - - messageQueue.sendMessage(prompt) - - let fullResponse = "" - let assistantMessage = "" - - try { - for await (const message of response) { - switch (true) { - case message.type === messageTypes.SYSTEM && - message.subtype === messageTypes.INIT: { - sessionId = message.session_id - continue - } - - case message.type === messageTypes.STREAM_EVENT: { - if (streamEnabled) { - const event = message.event - - if (event.type === "content_block_delta") { - if (event.delta.type === "text_delta") { - assistantMessage += event.delta.text - } - } - } - continue - } - - case message.type === messageTypes.ASSISTANT: { - for (const content of message.message.content) { - if (content.type === "text") { - if (!streamEnabled) { - assistantMessage += content.text - } - } else if (content.type === "tool_use") { - // Emit tool use notification if mcpServer is available - if (mcpServer) { - await mcpServer.sendLoggingMessage({ - level: "info", - data: JSON.stringify({ - type: "tool_use", - name: content.name, - input: content.input, - }), - }) - } - } - } - break - } - - case message.type === messageTypes.RESULT: { - if (assistantMessage) { - fullResponse = assistantMessage - } - return fullResponse - } - } - } - } catch (error) { - throw new Error(`[agent-chat-cli] Error: ${error}`) - } - - return fullResponse -} diff --git a/src/prompts/agents/demo-agent.md b/src/prompts/agents/demo-agent.md new file mode 100644 index 0000000..88c53df --- /dev/null +++ b/src/prompts/agents/demo-agent.md @@ -0,0 +1,5 @@ +# Demo Agent + +## Instructions + +- Your sole purpose is to respond to all queries in the style of William Faulkner. diff --git a/src/prompts/artsy-mcp.md b/src/prompts/artsy-mcp.md deleted file mode 100644 index 170d1bb..0000000 --- a/src/prompts/artsy-mcp.md +++ /dev/null @@ -1,15 +0,0 @@ -# System Prompt for Artsy MCP Server Agent - -You are an expert Artsy GraphQL assistant with access to Artsy's complete production GraphQL API through the artsy-mcp server. You can query, analyze and execute mutations. - -**CRITICAL**: If a user mentions "artsy mcp", invoke "artsymcp"-related tools. Never, ever, under any circumstance, return generalized knowledge when "artsy mcp" is included in a users prompt. - -**NEVER** generate a "best guess" based on known GraphQL patterns instead of using artsymcp-introspect or artsymcp-search to see whether the necessary fields and structure exist on the live server. - -## Instructions - -- Always use `internalID` (not `id` or `_id`) when referencing records between operations. Understand that `id` is a -- Use the introspect tool to understand schema structure before complex queries -- Under no circumstances will you execute a mutation unless the user has explicitly asked to perform an action. Eg "update andy warhol's birthday to x/y/z". -- If a graphql query fails with no results, use the introspect tool to search the schema for a similar or more appropriate query to execute -- Frequently slugs can be used as ids and vice versa. If one fails, try the other. Search the schema for ways to find the correct slug or ID. Do not give up right away. diff --git a/src/prompts/github.md b/src/prompts/github.md index a6d7cf2..af0f193 100644 --- a/src/prompts/github.md +++ b/src/prompts/github.md @@ -1,14 +1,8 @@ # System Prompt for Github MCP Server Agent -You are a GitHub MCP server agent, optimized for performing **READ ONLY** operations on artsy github repos. - -**ALL** GitHub search queries performed by Artsy Agent must include the org:artsy qualifier in the query string (e.g., org:artsy is:pr ), unless the user explicitly asks for a global or broader search. - -Example logic in natural language: - -> Whenever a user asks for a search across PRs (or issues, or code) in GitHub, always prepend/add org:artsy to the query string so that search is scoped strictly to the Artsy organization. For example, failed login becomes org:artsy is:pr failed login for PRs. +You are a GitHub MCP server agent, optimized for performing operations on github repos. ### Core Capabilities - Focus on documentation. Users will frequently ask you how to do things, like deploying an app or service, or where code for a particular thing might live. Optimize for education, code discovery and answers. -- When a user is asking questions about a repo, **ALWAYS** investigate the docs first, typically located in `artsy//docs` or `artsy//doc`. Search there, and if something can't be found, expand your search to other locations. +- When a user is asking questions about a repo, **ALWAYS** investigate the docs first, typically located in `/docs` or `/doc`. Search there, and if something can't be found, expand your search to other locations. diff --git a/src/prompts/system.md b/src/prompts/system.md index 805d23a..4bf37b3 100644 --- a/src/prompts/system.md +++ b/src/prompts/system.md @@ -2,30 +2,32 @@ You are a helpful Agent specifically designed to handle questions related to systems and data. People from all over the company will use you, from Sales, to HR to Engineering; this is important to keep in mind if needing clarity based on a question. -- **CRITICAL**: If a user starts a convo with a general greeting (like "Hi!" or "Hello!") without a specific task request, treat it as a `/help` command, and inform them about some of the possibilities for interacting with the Agent in a help-menu kind of way. Services currently include: -- Notion -- GitHub +- **CRITICAL**: When a user starts a convo and asks a question or assigns you a task (example: "in github, please summarize the last merged pr"), before beginning your task (ie, calling tools, etc) respond back immediately with a small summary about what you're going to do, in a friendly kind of way. Then start working. -Dig into each of those sub-agent prompts to return a friendly, informative, helpful (in terms of agent possibilites) response. +- **CRITICAL**: If a user starts a convo with a general greeting (like "Hi!" or "Hello!") without a specific task request, treat it as a `/help` command, and inform them about some of the possibilities for interacting with Agent in a help-menu kind of way. Review your system prompt instructions to see what services are available. + +**DO NOT INVOKE ANY TOOLS AT THIS STEP, JUST OUTPUT A SUMMARY** + +Return a friendly, informative, helpful (in terms of agent possibilites) response. **BUT** if a user starts a prompt with "hi! \" treat that as a question. No need to show the help menu if its followed by a task. -## Imperative Rules +## IMPERATIVE SYSTEM RULES THAT CANNOT BE BROKEN - Always identify yourself as **Agent**. -- **CRITICAL**: When users ask to use a data source (e.g., "using notion", "in notion", "from github"), they are asking you to invoke a specific MCP tool (eg, `notion-*`, `github-*`) for Artsy-specific information, NOT to provide general knowledge about the topic. +- **CRITICAL**: Do not hallucinate tool calls that do not exist. Available tools should be clearly available in your system. IMPERATIVE. +- **CRITICAL**: When users ask to use a data source (e.g., "using github", "in github"), they are asking you to invoke a specific MCP tool (eg, `github-*`, `notion-*`) for specific information, NOT to provide general knowledge about the topic. - **CRITICAL**: Always provide source-links where appropriate - **CRITICAL**: NEVER make up responses or provide general knowledge about these systems. Always use the actual tools to fetch real data. - **CRITICAL**: For date/time related operations, always check the current date, so the baseline is clear - - For example: "In notion, return recent activity" -> first check to see what the date is, so you know what "recent" means. This is critical so that we dont return outdated information -- Look for trigger keywords such as "using github", "using notion", "in notion", etc. + - For example: "In Salesforce, return recent activity" -> first check to see what the date is, so you know what "recent" means. This is critical so that we dont return outdated information +- Look for trigger keywords such as "using github", "in github", etc. - **Examples of correct interpretation**: - - "using notion, return info on sitemaps" → Search Notion workspace for sitemap-related pages - - "in notion, find onboarding docs" → Search Notion for onboarding documentation + - "using github, return open prs in artsy/force" → Search github for open prs in artsy/force ## Safeguards -- **CRITICAL TOOL USAGE**: When a user mentions any available tools by name ("notion", "github", etc.), you MUST invoke the appropriate tools related to their request. NEVER make up responses or provide general knowledge about these systems. Always use the actual tools to fetch real data. +- **CRITICAL TOOL USAGE**: When a user mentions any available tools by name, you MUST invoke the appropriate tools related to their request. NEVER make up responses or provide general knowledge about these systems. Always use the actual tools to fetch real data. - Do not fabricate answers. If unsure, say you don't know. - Prefer canonical documents (handbooks, wikis, root dashboards) over stale or duplicate pages. - If multiple plausible results exist, group and present them clearly for disambiguation. diff --git a/src/store.ts b/src/store.ts index aa356b8..74a828c 100644 --- a/src/store.ts +++ b/src/store.ts @@ -6,9 +6,13 @@ import { action, computed, createContextStore, + thunk, type Action, type Computed, + type Thunk, } from "easy-peasy" +import type { AgentConfig } from "utils/createAgent" +import { getEnabledMcpServers } from "utils/getEnabledMcpServers" import { MessageQueue } from "utils/MessageQueue" export interface Message { @@ -37,18 +41,23 @@ export interface ToolDenied { export type ChatHistoryEntry = Message | ToolUse | ToolDenied type McpServerConfigWithPrompt = McpServerConfig & { - prompt?: string - denyTools?: string[] + description: string + prompt?: () => Promise + disallowedTools?: string[] + enabled?: boolean } export interface AgentChatConfig { + agents?: Record + disallowedTools?: string[] connectionTimeout?: number maxRetries?: number mcpServers: Record + model?: "sonnet" | "haiku" permissionMode?: PermissionMode retryDelay?: number stream?: boolean - systemPrompt?: string + systemPrompt?: () => Promise } export interface PendingToolPermission { @@ -72,6 +81,8 @@ export interface StoreModel { // Computed isBooted: Computed + availableMcpServers: Computed + availableAgents: Computed // Actions abortRequest: Action @@ -88,17 +99,18 @@ export interface StoreModel { > setAbortController: Action setConfig: Action - setcurrentAssistantMessage: Action + setCurrentAssistantMessage: Action setCurrentToolUses: Action setInput: Action setIsProcessing: Action setMcpServers: Action setSessionId: Action setStats: Action + handleMcpServerStatus: Thunk } export const AgentStore = createContextStore({ - abortController: undefined, + abortController: new AbortController(), chatHistory: [], messageQueue: new MessageQueue(), sessionId: undefined, @@ -116,9 +128,19 @@ export const AgentStore = createContextStore({ return !!state.config }), + availableMcpServers: computed((state) => { + const enabled = getEnabledMcpServers(state.config?.mcpServers) + return enabled ? Object.keys(enabled) : [] + }), + + availableAgents: computed((state) => { + return state.config?.agents ? Object.keys(state.config.agents) : [] + }), + // Actions abortRequest: action((state) => { state.abortController?.abort() + state.abortController = new AbortController() state.isProcessing = false }), @@ -149,7 +171,7 @@ export const AgentStore = createContextStore({ state.isProcessing = payload }), - setcurrentAssistantMessage: action((state, payload) => { + setCurrentAssistantMessage: action((state, payload) => { state.currentAssistantMessage = payload }), @@ -196,4 +218,24 @@ export const AgentStore = createContextStore({ setAbortController: action((state, payload) => { state.abortController = payload }), + + handleMcpServerStatus: thunk((actions, mcpServers) => { + actions.setMcpServers(mcpServers) + + if (mcpServers.length === 0) { + return + } + + const failedServers = mcpServers + .filter((s) => s.status === "failed") + .map((s) => s.name) + + if (failedServers.length > 0) { + actions.addChatHistoryEntry({ + type: "message", + role: "system", + content: `[Error] Failed to connect to ${failedServers.join(", ")}`, + }) + } + }), }) diff --git a/src/utils/__tests__/createAgent.test.ts b/src/utils/__tests__/createAgent.test.ts new file mode 100644 index 0000000..f37e8ab --- /dev/null +++ b/src/utils/__tests__/createAgent.test.ts @@ -0,0 +1,232 @@ +import { test, expect, describe } from "bun:test" +import { createAgent, createSDKAgents } from "utils/createAgent" + +describe("createAgent", () => { + test("should return agent config as-is", () => { + const config = { + description: "Test agent", + prompt: async () => "Test prompt", + mcpServers: ["server1", "server2"], + } + + const result = createAgent(config) + + expect(result).toEqual(config) + }) + + test("should work without mcpServers", () => { + const config = { + description: "Test agent", + prompt: async () => "Test prompt", + } + + const result = createAgent(config) + + expect(result).toEqual(config) + expect(result.mcpServers).toBeUndefined() + }) + + test("should preserve function reference", () => { + const promptFn = async () => "Test prompt" + const config = { + description: "Test agent", + prompt: promptFn, + } + + const result = createAgent(config) + + expect(result.prompt).toBe(promptFn) + }) +}) + +describe("createSDKAgents", () => { + test("should convert agents to SDK format", async () => { + const agents = { + agent1: createAgent({ + description: "First agent", + prompt: async () => "Prompt 1", + mcpServers: ["server1"], + }), + agent2: createAgent({ + description: "Second agent", + prompt: async () => "Prompt 2", + mcpServers: ["server2"], + }), + } + + const result = await createSDKAgents(agents) + + expect(result).toEqual({ + agent1: { + description: "First agent", + prompt: "Prompt 1", + }, + agent2: { + description: "Second agent", + prompt: "Prompt 2", + }, + }) + }) + + test("should return undefined when no agents provided", async () => { + const result = await createSDKAgents(undefined) + + expect(result).toBeUndefined() + }) + + test("should handle empty agents object", async () => { + const result = await createSDKAgents({}) + + expect(result).toEqual({}) + }) + + test("should await lazy prompt functions", async () => { + let callCount = 0 + const agents = { + agent1: createAgent({ + description: "Test agent", + prompt: async () => { + callCount++ + return "Lazy prompt" + }, + }), + } + + expect(callCount).toBe(0) + + const result = await createSDKAgents(agents) + + expect(callCount).toBe(1) + expect(result?.agent1?.prompt).toBe("Lazy prompt") + }) + + test("should fetch all prompts in parallel", async () => { + const timestamps: number[] = [] + + const agents = { + agent1: createAgent({ + description: "Agent 1", + prompt: async () => { + timestamps.push(Date.now()) + await new Promise((resolve) => setTimeout(resolve, 10)) + return "Prompt 1" + }, + }), + agent2: createAgent({ + description: "Agent 2", + prompt: async () => { + timestamps.push(Date.now()) + await new Promise((resolve) => setTimeout(resolve, 10)) + return "Prompt 2" + }, + }), + agent3: createAgent({ + description: "Agent 3", + prompt: async () => { + timestamps.push(Date.now()) + await new Promise((resolve) => setTimeout(resolve, 10)) + return "Prompt 3" + }, + }), + } + + await createSDKAgents(agents) + + // All prompts should start roughly at the same time (within 5ms) + const maxDiff = Math.max(...timestamps) - Math.min(...timestamps) + expect(maxDiff).toBeLessThan(5) + }) + + test("should handle agents without mcpServers", async () => { + const agents = { + agent1: createAgent({ + description: "Test agent", + prompt: async () => "Test prompt", + }), + } + + const result = await createSDKAgents(agents) + + expect(result).toEqual({ + agent1: { + description: "Test agent", + prompt: "Test prompt", + }, + }) + }) + + test("should preserve agent order", async () => { + const agents = { + zebra: createAgent({ + description: "Z agent", + prompt: async () => "Z prompt", + }), + alpha: createAgent({ + description: "A agent", + prompt: async () => "A prompt", + }), + beta: createAgent({ + description: "B agent", + prompt: async () => "B prompt", + }), + } + + const result = await createSDKAgents(agents) + + expect(Object.keys(result!)).toEqual(["zebra", "alpha", "beta"]) + }) + + test("should handle errors in prompt functions", async () => { + const agents = { + failingAgent: createAgent({ + description: "Failing agent", + prompt: async () => { + throw new Error("Prompt fetch failed") + }, + }), + } + + await expect(createSDKAgents(agents)).rejects.toThrow("Prompt fetch failed") + }) + + test("should work with single agent", async () => { + const agents = { + solo: createAgent({ + description: "Solo agent", + prompt: async () => "Solo prompt", + }), + } + + const result = await createSDKAgents(agents) + + expect(result).toEqual({ + solo: { + description: "Solo agent", + prompt: "Solo prompt", + }, + }) + }) + + test("should handle complex prompt strings", async () => { + const complexPrompt = ` + # System Instructions + + You are a specialized agent. + + ## Guidelines + - Be helpful + - Be concise + ` + + const agents = { + agent1: createAgent({ + description: "Complex agent", + prompt: async () => complexPrompt, + }), + } + + const result = await createSDKAgents(agents) + + expect(result?.agent1?.prompt).toBe(complexPrompt) + }) +}) diff --git a/src/utils/__tests__/getPrompt.test.ts b/src/utils/__tests__/getPrompt.test.ts index 370c50d..b0308b4 100644 --- a/src/utils/__tests__/getPrompt.test.ts +++ b/src/utils/__tests__/getPrompt.test.ts @@ -3,87 +3,97 @@ import { getPrompt, buildSystemPrompt } from "utils/getPrompt" import type { AgentChatConfig } from "store" describe("getPrompt", () => { - test("should load system prompt from file", () => { - const prompt = getPrompt("system.md") + test("should return a lazy function", () => { + const lazyPrompt = getPrompt("system.md") + expect(typeof lazyPrompt).toBe("function") + }) + + test("should load system prompt from file when invoked", async () => { + const lazyPrompt = getPrompt("system.md") + const prompt = await lazyPrompt() expect(typeof prompt).toBe("string") expect(prompt.length).toBeGreaterThan(0) }) - test("should trim whitespace from prompt", () => { - const prompt = getPrompt("system.md") + test("should trim whitespace from prompt", async () => { + const lazyPrompt = getPrompt("system.md") + const prompt = await lazyPrompt() expect(prompt).toBe(prompt.trim()) }) }) describe("buildSystemPrompt", () => { - test("should include current date", () => { + test("should include current date", async () => { const config: AgentChatConfig = { mcpServers: {}, } - const prompt = buildSystemPrompt(config) + const prompt = await buildSystemPrompt({ config }) expect(prompt).toContain("Current date:") }) - test("should use default prompt when systemPrompt is not provided", () => { + test("should use default prompt when systemPrompt is not provided", async () => { const config: AgentChatConfig = { mcpServers: {}, } - const prompt = buildSystemPrompt(config) + const prompt = await buildSystemPrompt({ config }) expect(prompt).toContain("You are a helpful agent.") }) - test("should use custom systemPrompt when provided", () => { + test("should use custom systemPrompt when provided", async () => { const config: AgentChatConfig = { mcpServers: {}, - systemPrompt: "You are a custom agent.", + systemPrompt: async () => "You are a custom agent.", } - const prompt = buildSystemPrompt(config) + const prompt = await buildSystemPrompt({ config }) expect(prompt).toContain("You are a custom agent.") expect(prompt).not.toContain("You are a helpful agent.") }) - test("should include MCP server prompts", () => { + test("should include MCP server prompts", async () => { const config: AgentChatConfig = { mcpServers: { github: { command: "node", args: [], - prompt: "GitHub server instructions", + description: "GitHub server", + prompt: async () => "GitHub server instructions", }, }, } - const prompt = buildSystemPrompt(config) + const prompt = await buildSystemPrompt({ config }) expect(prompt).toContain("# github MCP Server") expect(prompt).toContain("GitHub server instructions") }) - test("should combine multiple MCP server prompts", () => { + test("should combine multiple MCP server prompts", async () => { const config: AgentChatConfig = { mcpServers: { github: { command: "node", args: [], - prompt: "GitHub instructions", + description: "GitHub server", + prompt: async () => "GitHub instructions", }, gitlab: { command: "node", args: [], - prompt: "GitLab instructions", + description: "GitLab server", + prompt: async () => "GitLab instructions", }, }, } - const prompt = buildSystemPrompt(config) + const prompt = await buildSystemPrompt({ config }) expect(prompt).toContain("# github MCP Server") expect(prompt).toContain("GitHub instructions") @@ -91,44 +101,118 @@ describe("buildSystemPrompt", () => { expect(prompt).toContain("GitLab instructions") }) - test("should skip MCP servers without prompts", () => { + test("should skip MCP servers without prompts", async () => { const config: AgentChatConfig = { mcpServers: { github: { command: "node", args: [], - prompt: "GitHub instructions", + description: "GitHub server", + prompt: async () => "GitHub instructions", }, gitlab: { command: "node", args: [], + description: "GitLab server", }, }, } - const prompt = buildSystemPrompt(config) + const prompt = await buildSystemPrompt({ config }) expect(prompt).toContain("# github MCP Server") expect(prompt).not.toContain("# gitlab MCP Server") }) - test("should build prompt with custom system prompt and MCP prompts", () => { + test("should build prompt with custom system prompt and MCP prompts", async () => { const config: AgentChatConfig = { mcpServers: { github: { command: "node", args: [], - prompt: "GitHub instructions", + description: "GitHub server", + prompt: async () => "GitHub instructions", }, }, - systemPrompt: "Custom base prompt", + systemPrompt: async () => "Custom base prompt", } - const prompt = buildSystemPrompt(config) + const prompt = await buildSystemPrompt({ config }) expect(prompt).toContain("Current date:") expect(prompt).toContain("Custom base prompt") expect(prompt).toContain("# github MCP Server") expect(prompt).toContain("GitHub instructions") }) + + test("should include connected MCP servers in system prompt", async () => { + const config: AgentChatConfig = { + mcpServers: {}, + } + const connectedServers = new Set(["github", "gitlab"]) + + const prompt = await buildSystemPrompt({ + config, + connectedServers, + }) + + expect(prompt).toContain("Connected MCP Servers") + expect(prompt).toContain("github, gitlab") + }) + + test("should handle empty connected servers set", async () => { + const config: AgentChatConfig = { + mcpServers: {}, + } + const connectedServers = new Set() + + const prompt = await buildSystemPrompt({ + config, + connectedServers, + }) + + expect(prompt).toContain("Connected MCP Servers") + expect(prompt).toContain( + "The following MCP servers are currently connected and available:" + ) + }) + + test("should work without connectedServers parameter", async () => { + const config: AgentChatConfig = { + mcpServers: {}, + } + + const prompt = await buildSystemPrompt({ config }) + + expect(prompt).toContain("Connected MCP Servers") + expect(prompt).toContain("Current date:") + }) + + test("should skip disabled MCP servers", async () => { + const config: AgentChatConfig = { + mcpServers: { + github: { + command: "node", + args: [], + description: "GitHub server", + prompt: async () => "GitHub instructions", + enabled: true, + }, + gitlab: { + command: "node", + args: [], + description: "GitLab server", + prompt: async () => "GitLab instructions", + enabled: false, + }, + }, + } + + const prompt = await buildSystemPrompt({ config }) + + expect(prompt).toContain("# github MCP Server") + expect(prompt).toContain("GitHub instructions") + expect(prompt).not.toContain("# gitlab MCP Server") + expect(prompt).not.toContain("GitLab instructions") + }) }) diff --git a/src/utils/__tests__/getToolInfo.test.ts b/src/utils/__tests__/getToolInfo.test.ts index 0170632..7103d94 100644 --- a/src/utils/__tests__/getToolInfo.test.ts +++ b/src/utils/__tests__/getToolInfo.test.ts @@ -1,10 +1,10 @@ -import { test, expect, describe } from "bun:test" +import { describe, expect, test } from "bun:test" +import type { AgentChatConfig } from "store" import { - getToolInfo, getDisallowedTools, + getToolInfo, isToolDisallowed, } from "utils/getToolInfo" -import type { AgentChatConfig } from "store" describe("getToolInfo", () => { test("should extract server name and tool name from MCP format", () => { @@ -50,7 +50,8 @@ describe("getDisallowedTools", () => { github: { command: "node", args: [], - denyTools: ["search_repositories", "create_issue"], + description: "GitHub server", + disallowedTools: ["search_repositories", "create_issue"], }, }, } @@ -60,6 +61,7 @@ describe("getDisallowedTools", () => { expect(result).toEqual([ "mcp__github__search_repositories", "mcp__github__create_issue", + "Bash", ]) }) @@ -69,13 +71,14 @@ describe("getDisallowedTools", () => { github: { command: "node", args: [], + description: "GitHub server", }, }, } const result = getDisallowedTools(config) - expect(result).toEqual([]) + expect(result).toEqual(["Bash"]) }) test("should handle multiple servers", () => { @@ -84,39 +87,47 @@ describe("getDisallowedTools", () => { github: { command: "node", args: [], - denyTools: ["search"], + description: "GitHub server", + disallowedTools: ["search"], }, gitlab: { command: "node", args: [], - denyTools: ["merge"], + description: "GitLab server", + disallowedTools: ["merge"], }, }, } const result = getDisallowedTools(config) - expect(result).toEqual(["mcp__github__search", "mcp__gitlab__merge"]) + expect(result).toEqual([ + "mcp__github__search", + "mcp__gitlab__merge", + "Bash", + ]) }) - test("should handle servers with no denyTools", () => { + test("should handle servers with no disallowedTools", () => { const config: AgentChatConfig = { mcpServers: { github: { command: "node", args: [], - denyTools: ["search"], + description: "GitHub server", + disallowedTools: ["search"], }, gitlab: { command: "node", args: [], + description: "GitLab server", }, }, } const result = getDisallowedTools(config) - expect(result).toEqual(["mcp__github__search"]) + expect(result).toEqual(["mcp__github__search", "Bash"]) }) }) @@ -127,7 +138,8 @@ describe("isToolDisallowed", () => { github: { command: "node", args: [], - denyTools: ["search_repositories"], + description: "GitHub server", + disallowedTools: ["search_repositories"], }, }, } @@ -146,7 +158,8 @@ describe("isToolDisallowed", () => { github: { command: "node", args: [], - denyTools: ["search_repositories"], + description: "GitHub server", + disallowedTools: ["search_repositories"], }, }, } @@ -165,6 +178,7 @@ describe("isToolDisallowed", () => { github: { command: "node", args: [], + description: "GitHub server", }, }, } @@ -183,7 +197,8 @@ describe("isToolDisallowed", () => { github: { command: "node", args: [], - denyTools: ["search"], + description: "GitHub server", + disallowedTools: ["search"], }, }, } diff --git a/src/utils/__tests__/validateEnv.test.ts b/src/utils/__tests__/validateEnv.test.ts index 4f87f66..039bbeb 100644 --- a/src/utils/__tests__/validateEnv.test.ts +++ b/src/utils/__tests__/validateEnv.test.ts @@ -1,4 +1,4 @@ -import { test, expect, describe, beforeEach, afterEach, spyOn } from "bun:test" +import { afterEach, beforeEach, describe, expect, spyOn, test } from "bun:test" import { validateEnv } from "../validateEnv" describe("validateEnv", () => { @@ -21,7 +21,6 @@ describe("validateEnv", () => { test("should pass when all required environment variables are present", () => { process.env.ANTHROPIC_API_KEY = "test-key" - process.env.GITHUB_ACCESS_TOKEN = "test-github" validateEnv() @@ -31,27 +30,6 @@ describe("validateEnv", () => { test("should exit when ANTHROPIC_API_KEY is missing", () => { delete process.env.ANTHROPIC_API_KEY - process.env.GITHUB_ACCESS_TOKEN = "test-github" - - validateEnv() - - expect(consoleErrorSpy).toHaveBeenCalled() - expect(processExitSpy).toHaveBeenCalledWith(1) - }) - - test("should exit when GITHUB_ACCESS_TOKEN is missing", () => { - process.env.ANTHROPIC_API_KEY = "test-key" - delete process.env.GITHUB_ACCESS_TOKEN - - validateEnv() - - expect(consoleErrorSpy).toHaveBeenCalled() - expect(processExitSpy).toHaveBeenCalledWith(1) - }) - - test("should exit when all environment variables are missing", () => { - delete process.env.ANTHROPIC_API_KEY - delete process.env.GITHUB_ACCESS_TOKEN validateEnv() @@ -61,12 +39,10 @@ describe("validateEnv", () => { test("should display helpful error message with missing variables", () => { delete process.env.ANTHROPIC_API_KEY - delete process.env.GITHUB_ACCESS_TOKEN validateEnv() const errorMessage = consoleErrorSpy.mock.calls[0][0] expect(errorMessage).toContain("ANTHROPIC_API_KEY") - expect(errorMessage).toContain("GITHUB_ACCESS_TOKEN") }) }) diff --git a/src/utils/createAgent.ts b/src/utils/createAgent.ts new file mode 100644 index 0000000..d08c584 --- /dev/null +++ b/src/utils/createAgent.ts @@ -0,0 +1,41 @@ +import type { AgentDefinition } from "@anthropic-ai/claude-agent-sdk" + +export interface AgentConfig { + description: string + prompt: () => Promise + mcpServers?: string[] +} + +export const createAgent = (options: AgentConfig): AgentConfig => { + return options +} + +/** + * We need to convert our async config format into the format that + * claude-agent-sdk accepts. + */ +export const createSDKAgents = async ( + agents?: Record +): Promise | undefined> => { + if (!agents) { + return undefined + } + + const entries = await Promise.all( + Object.entries(agents).map(async ([name, agent]) => { + const prompt = await agent.prompt() + + return [ + name, + { + description: agent.description, + prompt, + }, + ] as const + }) + ) + + const sdkAgents = Object.fromEntries(entries) + + return sdkAgents +} diff --git a/src/utils/getEnabledMcpServers.ts b/src/utils/getEnabledMcpServers.ts new file mode 100644 index 0000000..ba0c782 --- /dev/null +++ b/src/utils/getEnabledMcpServers.ts @@ -0,0 +1,13 @@ +import type { AgentChatConfig } from "store" + +export const getEnabledMcpServers = ( + mcpServers: AgentChatConfig["mcpServers"] | undefined +) => { + if (!mcpServers) { + return undefined + } + + return Object.fromEntries( + Object.entries(mcpServers).filter(([_, server]) => server.enabled !== false) + ) +} diff --git a/src/utils/getPrompt.ts b/src/utils/getPrompt.ts index edb7ae3..bc91584 100644 --- a/src/utils/getPrompt.ts +++ b/src/utils/getPrompt.ts @@ -1,16 +1,29 @@ -import type { AgentChatConfig } from "store" +import type { Options } from "@anthropic-ai/claude-agent-sdk" import { readFileSync } from "node:fs" -import { resolve, dirname } from "node:path" +import { dirname, resolve } from "node:path" import { fileURLToPath } from "node:url" +import type { AgentChatConfig } from "store" const __dirname = dirname(fileURLToPath(import.meta.url)) -export const getPrompt = (filename: string): string => { - const path = resolve(__dirname, "../prompts", filename) - return readFileSync(path, "utf-8").trim() +export const getPrompt = (filename: string) => { + return async (): Promise => { + const path = resolve(__dirname, "../prompts", filename) + return readFileSync(path, "utf-8").trim() + } } -export const buildSystemPrompt = (config: AgentChatConfig): string => { +interface BuildSystemPromptProps { + config: AgentChatConfig + additionalSystemPrompt?: string + connectedServers?: Set +} + +export const buildSystemPrompt = async ({ + config, + additionalSystemPrompt = "", + connectedServers = new Set(), +}: BuildSystemPromptProps): Promise => { const currentDate = new Date().toLocaleDateString("en-US", { weekday: "long", year: "numeric", @@ -19,18 +32,47 @@ export const buildSystemPrompt = (config: AgentChatConfig): string => { }) const dateHeader = `Current date: ${currentDate}` - const basePrompt = config.systemPrompt ?? "You are a helpful agent." + const basePrompt = config.systemPrompt + ? await config.systemPrompt() + : "You are a helpful agent." + + const mcpPromptEntries = await Promise.all( + Object.entries(config.mcpServers) + .filter( + ([_, serverConfig]) => + serverConfig.enabled !== false && serverConfig.prompt + ) + .map(async ([name, serverConfig]) => { + const prompt = serverConfig.prompt ? await serverConfig.prompt() : "" + return `# ${name} MCP Server\n\n${prompt}` + }) + ) + + const mcpPrompts = mcpPromptEntries.join("\n\n") + + const parts = [dateHeader] + + if (additionalSystemPrompt) { + parts.push(additionalSystemPrompt) + } + + const connectedMCPServers = Array.from(connectedServers).join(", ") + + parts.push( + `# Connected MCP Servers + +**The following MCP servers are currently connected and available: ${connectedMCPServers}** + +- **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. +` + ) - const mcpPrompts = Object.entries(config.mcpServers) - .filter(([_, serverConfig]) => serverConfig.prompt) - .map( - ([name, serverConfig]) => `# ${name} MCP Server\n\n${serverConfig.prompt}` - ) - .join("\n\n") + parts.push(basePrompt) - if (!mcpPrompts) { - return `${dateHeader}\n\n${basePrompt}` + if (mcpPrompts) { + parts.push(mcpPrompts) } - return `${dateHeader}\n\n${basePrompt}\n\n${mcpPrompts}` + return parts.join("\n\n") } diff --git a/src/utils/getRemotePrompt.ts b/src/utils/getRemotePrompt.ts new file mode 100644 index 0000000..8cd6eaf --- /dev/null +++ b/src/utils/getRemotePrompt.ts @@ -0,0 +1,25 @@ +import { getPrompt } from "./getPrompt" + +interface GetRemotePromptOptions { + fallback?: string + fetchPrompt: () => Promise +} + +export const getRemotePrompt = ({ + fetchPrompt, + fallback, +}: GetRemotePromptOptions) => { + return async () => { + try { + return await fetchPrompt() + } catch (error) { + if (fallback) { + return await getPrompt(fallback)() + } + + throw new Error( + `[agent] [getRemotePrompt] Failed to fetch remote prompt: ${error instanceof Error ? error.message : String(error)}` + ) + } + } +} diff --git a/src/utils/getToolInfo.ts b/src/utils/getToolInfo.ts index 9576a9f..0cc120e 100644 --- a/src/utils/getToolInfo.ts +++ b/src/utils/getToolInfo.ts @@ -12,21 +12,23 @@ export const getToolInfo = (tool: string) => { } /** - * Get disallowedTools list from MCP server denyTools config. + * Get disallowedTools list from MCP server disallowedTools config. * Converts tool names like "search_repositories" to "mcp__github__search_repositories" */ export const getDisallowedTools = (config: AgentChatConfig): string[] => { - const disallowed = Object.entries(config.mcpServers).flatMap( - ([serverName, serverConfig]) => { - if (!serverConfig.denyTools) { + const disallowedSystemTools = config.disallowedTools ?? ["Bash"] + + const disallowed = Object.entries(config.mcpServers) + .flatMap(([serverName, serverConfig]) => { + if (!serverConfig.disallowedTools) { return [] } - return serverConfig.denyTools.map( + return serverConfig.disallowedTools.map( (toolName) => `mcp__${serverName}__${toolName}` ) - } - ) + }) + .concat(disallowedSystemTools) return disallowed } diff --git a/src/utils/loadConfig.ts b/src/utils/loadConfig.ts index c8bc87b..05dd948 100644 --- a/src/utils/loadConfig.ts +++ b/src/utils/loadConfig.ts @@ -5,7 +5,7 @@ export const loadConfig = async (): Promise => { const result = await cosmiconfig("agent-chat-cli").search() if (!result || result.isEmpty) { - throw new Error("[agent-chat-cli] No configuration file found") + throw new Error("[agent-cli] No configuration file found") } return result.config diff --git a/src/utils/logger.ts b/src/utils/logger.ts new file mode 100644 index 0000000..3c24bb2 --- /dev/null +++ b/src/utils/logger.ts @@ -0,0 +1,7 @@ +export const enableLogging = process.env.ENABLE_LOGGING === "true" + +export const log = (...messages) => { + if (enableLogging) { + console.log("[agent]", ...messages) + } +} diff --git a/src/utils/mcpServerSelectionAgent.ts b/src/utils/mcpServerSelectionAgent.ts new file mode 100644 index 0000000..2a8073e --- /dev/null +++ b/src/utils/mcpServerSelectionAgent.ts @@ -0,0 +1,222 @@ +import { + createSdkMcpServer, + query, + tool, + type SDKUserMessage, +} from "@anthropic-ai/claude-agent-sdk" +import type { AgentConfig } from "utils/createAgent" +import { log } from "utils/logger" +import { z } from "zod" +import { messageTypes } from "./runAgentLoop" + +interface SelectMcpServersOptions { + abortController?: AbortController + agents?: Record + alreadyConnectedServers?: Set + enabledMcpServers: Record | undefined + onServerConnection?: (status: string) => void + sessionId?: string + userMessage: string +} + +export const selectMcpServers = async ({ + abortController, + agents, + alreadyConnectedServers = new Set(), + enabledMcpServers, + onServerConnection, + sessionId, + userMessage, +}: SelectMcpServersOptions) => { + if (!enabledMcpServers) { + return { mcpServers: undefined, newServers: [] } + } + + const mcpServerNames = Object.keys(enabledMcpServers).join(", ") + + log("[mcpServerSelectionAgent] Available servers:", mcpServerNames) + + log( + "[mcpServerSelectionAgent] Already connected:", + Array.from(alreadyConnectedServers).join(", ") || "none" + ) + + const serverCapabilities = Object.entries(enabledMcpServers) + .map(([name, server]) => { + const description = server.description || "No description" + return `- ${name}: ${description}` + }) + .join("\n") + + // Build agent capabilities section + const agentCapabilities = agents + ? Object.entries(agents) + .map(([name, agent]) => { + const description = agent.description || "No description" + const requiredServers = agent.mcpServers || [] + return `- ${name}: ${description}${requiredServers.length > 0 ? ` (requires: ${requiredServers.join(", ")})` : ""}` + }) + .join("\n") + : "" + + let selectedServers: string[] = [] + + /** + * Create a custom MCP server tool. + */ + const selectionServer = createSdkMcpServer({ + name: "mcp-router", + version: "0.0.1", + tools: [ + tool( + "select_mcp_servers", + "Select which MCP servers are needed for the user's request", + { + servers: z + .array(z.string()) + .describe( + "Array of MCP server names needed. Use exact names from the available list. Return empty array if no servers are needed." + ), + }, + async (args) => { + selectedServers = args.servers.map((s) => s.trim()) + + return { + content: [ + { + type: "text", + text: `Selected servers: ${selectedServers.join(", ") || "none"}`, + }, + ], + } + } + ), + ], + }) + + const routingSystemPrompt = `You are an MCP server router. Your job is to determine which MCP servers are needed for a user's request. + +Available MCP servers: ${mcpServerNames} + +SERVER CAPABILITIES: +${serverCapabilities} +${ + agentCapabilities + ? ` +AVAILABLE SUBAGENTS: +${agentCapabilities} + +When a user request matches a subagent's domain, include that subagent's required MCP servers in your selection. +` + : "" +} +INSTRUCTIONS: +- You do not need to respond with a friendly greeting. Your sole purpose is to return results in the form requested below. +- Intelligently infer which servers are needed based on the request context and server capabilities +- Consider if the request might invoke a subagent and include its required servers +- Match case-insensitively when user explicitly mentions server names (e.g., "github" matches "github") +- Use the select_mcp_servers tool to return your selection +- If no relevant servers are available, return an empty array + +Examples: +- "Show me GitHub issues" → ["github"] +- "Show me some docs on OKRs" → ["notion"] +- "What's the weather?" → [] +` + + const routingResponse = query({ + prompt: (async function* () { + yield { + type: "user" as const, + session_id: sessionId || "", + message: { + role: "user" as const, + content: userMessage, + }, + } as SDKUserMessage + })(), + options: { + model: "haiku", + systemPrompt: routingSystemPrompt, + abortController, + mcpServers: { + "mcp-router": selectionServer, + }, + allowedTools: ["mcp__mcp-router__select_mcp_servers"], + maxTurns: 1, + }, + }) + + for await (const message of routingResponse) { + if (message.type === messageTypes.ASSISTANT) { + log( + "[mcpServerSelectionAgent] [messageTypes.ASSISTANT]:", + JSON.stringify(message.message.content, null, 2) + ) + } + } + + log("[mcpServerSelectionAgent] Selected MCP servers:", selectedServers) + + const newServers = selectedServers.filter( + (server) => !alreadyConnectedServers.has(server.toLowerCase()) + ) + + if (newServers.length > 0) { + log( + "[mcpServerSelectionAgent] New MCP servers to connect:", + newServers.join(", ") + ) + } else { + log("[mcpServerSelectionAgent] No new MCP servers needed") + } + + const allServers = new Set([ + ...Array.from(alreadyConnectedServers), + ...selectedServers, + ]) + + const mcpServers = + allServers.size > 0 + ? Object.fromEntries( + Object.entries(enabledMcpServers).filter(([name]) => + Array.from(allServers).some( + (s) => s.toLowerCase() === name.toLowerCase() + ) + ) + ) + : undefined + + log( + "[mcpServerSelectionAgent] Final MCP servers:", + mcpServers ? Object.keys(mcpServers).join(", ") : "none" + ) + + // Log servers selected for this turn + log( + "[mcpServerSelectionAgent] Servers selected for this turn:", + newServers.length > 0 ? newServers : "none (reusing existing)" + ) + + // Log total accumulated servers + log( + "[mcpServerSelectionAgent] Total accumulated servers:", + Array.from(allServers).join(", ") || "none" + ) + + // Notify about new server connections + if (newServers.length > 0) { + const serverList = newServers.join(", ") + onServerConnection?.(`Connecting to ${serverList}...`) + } + + // Update the connected servers set with new servers + newServers.forEach((server) => { + alreadyConnectedServers.add(server.toLowerCase()) + }) + + return { + mcpServers, + newServers, + } +} diff --git a/src/utils/runAgentLoop.ts b/src/utils/runAgentLoop.ts index 221e8fd..9d2e0b7 100644 --- a/src/utils/runAgentLoop.ts +++ b/src/utils/runAgentLoop.ts @@ -1,8 +1,12 @@ import { query, type SDKUserMessage } from "@anthropic-ai/claude-agent-sdk" import type { AgentChatConfig } from "store" -import { buildSystemPrompt } from "utils/getPrompt" 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 type { MessageQueue } from "utils/MessageQueue" export const messageTypes = { @@ -13,25 +17,34 @@ export const messageTypes = { SYSTEM: "system", } as const +export const contentTypes = { + TEXT: "text", + TOOL_USE: "tool_use", +} as const + export interface RunAgentLoopOptions { - abortController?: AbortController + abortControllerRef?: { current: AbortController | undefined } + additionalSystemPrompt?: string config: AgentChatConfig + existingConnectedServers?: Set messageQueue: MessageQueue - sessionId?: string + onServerConnection?: (status: string) => void onToolPermissionRequest?: (toolName: string, input: any) => void + sessionId?: string setIsProcessing?: (value: boolean) => void } -export const runAgentLoop = ({ - abortController, +export const runAgentLoop = async ({ + abortControllerRef, + additionalSystemPrompt, config, + existingConnectedServers, messageQueue, + onServerConnection, onToolPermissionRequest, - sessionId, + sessionId: initialSessionId, setIsProcessing, }: RunAgentLoopOptions) => { - const systemPrompt = buildSystemPrompt(config) - const canUseTool = createCanUseTool({ messageQueue, onToolPermissionRequest, @@ -39,48 +52,99 @@ export const runAgentLoop = ({ }) const disallowedTools = getDisallowedTools(config) + const enabledMcpServers = getEnabledMcpServers(config.mcpServers) - const response = query({ - prompt: startConversation(messageQueue, sessionId), - options: { - model: "haiku", - permissionMode: config.permissionMode ?? "default", - includePartialMessages: config.stream ?? false, - mcpServers: config.mcpServers, - abortController, - canUseTool, - systemPrompt, - disallowedTools, - }, - }) + let currentSessionId = initialSessionId + const connectedServers = existingConnectedServers ?? new Set() - return { - response, - } -} + async function* agentLoop() { + while (true) { + const userMessage = await messageQueue.waitForMessage() -const startConversation = async function* ( - messageQueue: MessageQueue, - sessionId?: string -) { - while (true) { - const userMessage = await messageQueue.waitForMessage() + log("\n[runAgentLoop] USER:", userMessage, "\n") - if (userMessage.toLowerCase() === "exit") { - break - } + if (userMessage.toLowerCase() === "exit") { + break + } + + if (!userMessage.trim()) { + continue + } + + const systemPrompt = await buildSystemPrompt({ + config, + additionalSystemPrompt, + connectedServers, + }) + + const { mcpServers } = await selectMcpServers({ + abortController: abortControllerRef?.current, + agents: config.agents, + alreadyConnectedServers: connectedServers, + enabledMcpServers, + onServerConnection, + sessionId: currentSessionId, + userMessage, + }) + + const agents = await createSDKAgents(config.agents) + + try { + const turnResponse = query({ + prompt: (async function* () { + yield { + type: "user" as const, + session_id: currentSessionId || "", + message: { + role: "user" as const, + content: userMessage, + }, + } as SDKUserMessage + })(), + options: { + model: config.model ?? "haiku", + permissionMode: config.permissionMode ?? "default", + includePartialMessages: config.stream ?? false, + mcpServers, + agents, + abortController: abortControllerRef?.current, + canUseTool, + systemPrompt, + disallowedTools, + resume: currentSessionId, + }, + }) - if (!userMessage.trim()) { - continue + for await (const message of turnResponse) { + if ( + message.type === messageTypes.SYSTEM && + message.subtype === messageTypes.INIT + ) { + log( + "[runAgentLoop] [messageTypes.INIT]:", + JSON.stringify(message, null, 2) + ) + + currentSessionId = message.session_id + } + + yield message + + // If we hit a RESULT, this turn is complete + if (message.type === messageTypes.RESULT) { + break + } + } + } catch (error) { + log("[ERROR] [runAgentLoop] Query aborted or failed:", error) + + // Continue to next message + } } + } - yield { - type: "user" as const, - session_id: sessionId || "", - message: { - role: "user" as const, - content: userMessage, - }, - } as SDKUserMessage + return { + agentLoop: agentLoop(), + connectedServers, } } diff --git a/src/utils/validateEnv.ts b/src/utils/validateEnv.ts index 0a536d0..40de498 100644 --- a/src/utils/validateEnv.ts +++ b/src/utils/validateEnv.ts @@ -1,7 +1,6 @@ export function validateEnv() { const requiredEnvVars = { ANTHROPIC_API_KEY: process.env.ANTHROPIC_API_KEY, - GITHUB_ACCESS_TOKEN: process.env.GITHUB_ACCESS_TOKEN, } const missingVars = Object.entries(requiredEnvVars) @@ -14,7 +13,10 @@ export function validateEnv() { .map((v) => ` - ${v}`) .join("\n")}\n` ) - console.error("Please copy .env.example to .env and fill in the values.\n") + console.error( + "[agent-cli] Please copy .env.example to .env and fill in the values.\n" + ) + process.exit(1) } }