Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/components/UserInput.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,7 @@ export const UserInput: React.FC = () => {

return (
<Box>
<BlinkCaret />
<BlinkCaret enabled={!store.isProcessing} />

<TextInput
value={store.input}
Expand Down
6 changes: 4 additions & 2 deletions src/hooks/useAgent.ts
Original file line number Diff line number Diff line change
Expand Up @@ -6,11 +6,12 @@ import { messageTypes, runAgentLoop } from "utils/runAgentLoop"
export function useAgent() {
const messageQueue = AgentStore.useStoreState((state) => state.messageQueue)
const config = AgentStore.useStoreState((state) => state.config)
const mcpServers = AgentStore.useStoreState((state) => state.mcpServers)
const actions = AgentStore.useStoreActions((actions) => actions)
const currentAssistantMessageRef = useRef("")
const sessionIdRef = useRef<string | undefined>(undefined)
const abortControllerRef = useRef<AbortController | undefined>(undefined)
const connectedServersRef = useRef<Set<string>>(new Set())
const inferredServersRef = useRef<Set<string>>(new Set())

const runQuery = useCallback(
async (userMessage: string) => {
Expand All @@ -32,7 +33,8 @@ export function useAgent() {
const agentLoop = runAgentLoop({
abortController,
config,
connectedServers: connectedServersRef.current,
inferredServers: inferredServersRef.current,
mcpServers,
messageQueue,
onToolPermissionRequest: (toolName, input) => {
actions.setPendingToolPermission({ toolName, input })
Expand Down
4 changes: 2 additions & 2 deletions src/mcp/getAgentStatus.ts
Original file line number Diff line number Diff line change
Expand Up @@ -7,12 +7,12 @@ export const getAgentStatus = async (mcpServer?: McpServer) => {
const config = await loadConfig()
const messageQueue = new MessageQueue()
const abortController = new AbortController()
const connectedServers = new Set<string>()
const inferredServers = new Set<string>()

const agentLoop = runAgentLoop({
abortController,
config,
connectedServers,
inferredServers,
messageQueue,
userMessage: "status",
})
Expand Down
14 changes: 10 additions & 4 deletions src/mcp/getMcpServer.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"
import type { McpServerStatus } from "store"
import { registerAskAgentTool } from "mcp/tools/askAgent"
import { registerAskAgentSlackbotTool } from "mcp/tools/askAgentSlackbot"
import { registerGetAgentStatusTool } from "mcp/tools/getAgentStatus"
Expand All @@ -10,8 +11,11 @@ export const getMcpServer = () => {
// Map thread IDs to Claude Agent SDK session IDs for per-thread isolation
const threadSessions = new Map<string, string>()

// Map session IDs to connected MCP servers for persistence across requests
const sessionConnectedServers = new Map<string, Set<string>>()
// Map session IDs to inferred MCP servers for persistence across requests
const sessionInferredServers = new Map<string, Set<string>>()

// Map session IDs to MCP server statuses for persistence across requests
const sessionMcpServers = new Map<string, McpServerStatus[]>()

const mcpServer = new McpServer(
{
Expand All @@ -32,7 +36,8 @@ export const getMcpServer = () => {
get sessionId() {
return sessionId
},
sessionConnectedServers,
sessionInferredServers,
sessionMcpServers,
onSessionIdUpdate: (newSessionId) => {
sessionId = newSessionId
},
Expand All @@ -43,7 +48,8 @@ export const getMcpServer = () => {
mcpServer,
context: {
threadSessions,
sessionConnectedServers,
sessionInferredServers,
sessionMcpServers,
},
})

Expand Down
20 changes: 12 additions & 8 deletions src/mcp/runStandaloneAgentLoop.ts
Original file line number Diff line number Diff line change
@@ -1,12 +1,14 @@
import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"
import type { McpServerStatus } from "store"
import { loadConfig } from "utils/loadConfig"
import { log } from "utils/logger"
import { MessageQueue } from "utils/MessageQueue"
import { contentTypes, messageTypes, runAgentLoop } from "utils/runAgentLoop"

interface RunQueryOptions {
additionalSystemPrompt?: string
existingConnectedServers?: Set<string>
existingInferredServers?: Set<string>
existingMcpServers?: McpServerStatus[]
mcpServer: McpServer
onSessionIdReceived?: (sessionId: string) => void
prompt: string
Expand All @@ -15,7 +17,8 @@ interface RunQueryOptions {

export const runStandaloneAgentLoop = async ({
additionalSystemPrompt,
existingConnectedServers,
existingInferredServers,
existingMcpServers,
mcpServer,
onSessionIdReceived,
prompt,
Expand All @@ -25,17 +28,16 @@ export const runStandaloneAgentLoop = async ({
const messageQueue = new MessageQueue()
const streamEnabled = config.stream ?? false

const connectedServers = existingConnectedServers ?? new Set<string>()
const inferredServers = existingInferredServers ?? new Set<string>()
const abortController = new AbortController()

const agentLoop = runAgentLoop({
abortController,
additionalSystemPrompt,
config,
connectedServers,
inferredServers,
mcpServers: existingMcpServers,
messageQueue,
sessionId,
userMessage: prompt,
onServerConnection: async (status) => {
await mcpServer.sendLoggingMessage({
level: "info",
Expand All @@ -45,6 +47,8 @@ export const runStandaloneAgentLoop = async ({
}),
})
},
sessionId,
userMessage: prompt,
})

let finalResponse = ""
Expand Down Expand Up @@ -155,7 +159,7 @@ export const runStandaloneAgentLoop = async ({

return {
response: finalResponse,
connectedServers,
inferredServers,
}
}
}
Expand All @@ -166,6 +170,6 @@ export const runStandaloneAgentLoop = async ({

return {
response: finalResponse,
connectedServers,
inferredServers,
}
}
21 changes: 14 additions & 7 deletions src/mcp/tools/askAgent.ts
Original file line number Diff line number Diff line change
@@ -1,10 +1,12 @@
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"
import type { McpServerStatus } from "store"
import { runStandaloneAgentLoop } from "mcp/runStandaloneAgentLoop"
import { z } from "zod"

export interface AskAgentContext {
sessionId?: string
sessionConnectedServers: Map<string, Set<string>>
sessionInferredServers: Map<string, Set<string>>
sessionMcpServers: Map<string, McpServerStatus[]>
onSessionIdUpdate: (sessionId: string) => void
}

Expand All @@ -27,23 +29,28 @@ export const registerAskAgentTool = ({
},
},
async ({ query }) => {
const existingConnectedServers = context.sessionId
? context.sessionConnectedServers.get(context.sessionId)
const existingInferredServers = context.sessionId
? context.sessionInferredServers.get(context.sessionId)
: undefined

const { response, connectedServers } = await runStandaloneAgentLoop({
const existingMcpServers = context.sessionId
? context.sessionMcpServers.get(context.sessionId)
: undefined

const { response, inferredServers } = await runStandaloneAgentLoop({
prompt: query,
mcpServer,
sessionId: context.sessionId,
existingConnectedServers,
existingInferredServers,
existingMcpServers,
onSessionIdReceived: (newSessionId) => {
context.onSessionIdUpdate(newSessionId)
},
})

// Update the session's connected servers
// Update the session's inferred servers
if (context.sessionId) {
context.sessionConnectedServers.set(context.sessionId, connectedServers)
context.sessionInferredServers.set(context.sessionId, inferredServers)
}

return {
Expand Down
21 changes: 14 additions & 7 deletions src/mcp/tools/askAgentSlackbot.ts
Original file line number Diff line number Diff line change
@@ -1,10 +1,12 @@
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"
import type { McpServerStatus } from "store"
import { runStandaloneAgentLoop } from "mcp/runStandaloneAgentLoop"
import { z } from "zod"

export interface AskAgentSlackbotContext {
threadSessions: Map<string, string>
sessionConnectedServers: Map<string, Set<string>>
sessionInferredServers: Map<string, Set<string>>
sessionMcpServers: Map<string, McpServerStatus[]>
}

export interface RegisterAskAgentSlackbotToolProps {
Expand Down Expand Up @@ -41,29 +43,34 @@ export const registerAskAgentSlackbotTool = ({
? context.threadSessions.get(threadId)
: undefined

const existingConnectedServers = existingSessionId
? context.sessionConnectedServers.get(existingSessionId)
const existingInferredServers = existingSessionId
? context.sessionInferredServers.get(existingSessionId)
: undefined

const { response, connectedServers } = await runStandaloneAgentLoop({
const existingMcpServers = existingSessionId
? context.sessionMcpServers.get(existingSessionId)
: undefined

const { response, inferredServers } = await runStandaloneAgentLoop({
prompt: query,
mcpServer,
sessionId: existingSessionId,
additionalSystemPrompt: systemPrompt,
existingConnectedServers,
existingInferredServers,
existingMcpServers,
onSessionIdReceived: (newSessionId) => {
if (threadId) {
context.threadSessions.set(threadId, newSessionId)
}
},
})

// Update the session's connected servers
// Update the session's inferred servers
if (existingSessionId || threadId) {
const sessionId =
existingSessionId || context.threadSessions.get(threadId!)
if (sessionId) {
context.sessionConnectedServers.set(sessionId, connectedServers)
context.sessionInferredServers.set(sessionId, inferredServers)
}
}

Expand Down
59 changes: 46 additions & 13 deletions src/utils/__tests__/getPrompt.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -145,47 +145,80 @@ describe("buildSystemPrompt", () => {
expect(prompt).toContain("GitHub instructions")
})

test("should include connected MCP servers in system prompt", async () => {
test("should include inferred MCP servers in system prompt", async () => {
const config: AgentChatConfig = {
mcpServers: {},
}
const connectedServers = new Set(["github", "gitlab"])
const inferredServers = new Set(["github", "gitlab"])

const prompt = await buildSystemPrompt({
config,
connectedServers,
inferredServers,
})

expect(prompt).toContain("Connected MCP Servers")
expect(prompt).toContain("Server Selection Context")
expect(prompt).toContain("github, gitlab")
})

test("should handle empty connected servers set", async () => {
test("should include connected and failed MCP servers sections", async () => {
const config: AgentChatConfig = {
mcpServers: {},
}
const connectedServers = new Set<string>()
const mcpServers = [
{ name: "github", status: "connected" },
{ name: "gitlab", status: "failed" },
]

const prompt = await buildSystemPrompt({
config,
connectedServers,
mcpServers,
})

expect(prompt).toContain("Connected MCP Servers")
expect(prompt).toContain(
"The following MCP servers are currently connected and available:"
)
expect(prompt).toContain("Available MCP Servers")
expect(prompt).toContain("- github")
expect(prompt).toContain("ONLY servers with available tools")
expect(prompt).toContain("Unavailable MCP Servers")
expect(prompt).toContain("- gitlab")
expect(prompt).toContain("These servers have NO tools available")
})

test("should work without connectedServers parameter", async () => {
test("should not include connection status sections when no mcpServers provided", async () => {
const config: AgentChatConfig = {
mcpServers: {},
}

const prompt = await buildSystemPrompt({
config,
})

expect(prompt).not.toContain("Available MCP Servers")
expect(prompt).not.toContain("Unavailable MCP Servers")
})

test("should handle empty inferred servers set", async () => {
const config: AgentChatConfig = {
mcpServers: {},
}
const inferredServers = new Set<string>()

const prompt = await buildSystemPrompt({
config,
inferredServers,
})

expect(prompt).toContain("Current date:")
expect(prompt).not.toContain("Server Selection Context")
})

test("should work without inferredServers parameter", async () => {
const config: AgentChatConfig = {
mcpServers: {},
}

const prompt = await buildSystemPrompt({ config })

expect(prompt).toContain("Connected MCP Servers")
expect(prompt).toContain("Current date:")
expect(prompt).not.toContain("Server Selection Context")
})

test("should skip disabled MCP servers", async () => {
Expand Down
Loading