Fix MCP server: create separate McpServer instance per session

Previously shared one McpServer instance across all sessions/transports,
which caused Protocol._transport to be overwritten when new sessions
connected. This broke message routing for requests on previous transports
after long-running SSE streams.

Now creates a new McpServer instance per session, following the stateful
session pattern from the MCP SDK documentation. Each session gets its own
server+transport pair, preventing transport reference conflicts.

Fixes issue where get_agent_status would fail after wait_for_agent completed.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
Mohamed Boudra
2025-11-29 21:39:46 +00:00
parent 6b9fbb83c0
commit 36facf4bfd
4 changed files with 190 additions and 94 deletions

View File

@@ -76,6 +76,7 @@ export type WaitForAgentOptions = {
export type WaitForAgentResult = {
status: AgentLifecycleStatus;
permission: AgentPermissionRequest | null;
lastMessage: string | null;
};
type ManagedAgent = {
@@ -491,6 +492,32 @@ export class AgentManager {
await this.primeHistory(agent);
}
private getLastAssistantMessage(agentId: string): string | null {
const agent = this.agents.get(agentId);
if (!agent) {
return null;
}
// Collect the last contiguous assistant messages (Claude streams chunks)
const chunks: string[] = [];
for (let i = agent.timeline.length - 1; i >= 0; i--) {
const item = agent.timeline[i];
if (item.type !== "assistant_message") {
if (chunks.length) {
break;
}
continue;
}
chunks.push(item.text);
}
if (!chunks.length) {
return null;
}
return chunks.reverse().join("");
}
async waitForAgentEvent(
agentId: string,
options?: WaitForAgentOptions
@@ -500,41 +527,89 @@ export class AgentManager {
throw new Error(`Agent ${agentId} not found`);
}
const immediatePermission = snapshot.pendingPermissions[0] ?? null;
if (immediatePermission) {
return { status: snapshot.status, permission: immediatePermission };
return {
status: snapshot.status,
permission: immediatePermission,
lastMessage: this.getLastAssistantMessage(agentId)
};
}
if (!isAgentBusy(snapshot.status)) {
return { status: snapshot.status, permission: null };
return {
status: snapshot.status,
permission: null,
lastMessage: this.getLastAssistantMessage(agentId)
};
}
if (options?.signal?.aborted) {
throw createAbortError(options.signal, "wait_for_agent aborted");
}
return await new Promise<WaitForAgentResult>((resolve, reject) => {
// Bug #1 Fix: Check abort signal AGAIN inside Promise constructor
// to avoid race condition between pre-Promise check and abort listener registration
if (options?.signal?.aborted) {
reject(createAbortError(options.signal, "wait_for_agent aborted"));
return;
}
let currentStatus: AgentLifecycleStatus = snapshot.status;
const cleanupFns: Array<() => void> = [];
// Bug #3 Fix: Declare unsubscribe and abortHandler upfront so cleanup can reference them
let unsubscribe: (() => void) | null = null;
let abortHandler: (() => void) | null = null;
const cleanup = () => {
while (cleanupFns.length) {
const fn = cleanupFns.pop();
// Clean up subscription
if (unsubscribe) {
try {
fn?.();
unsubscribe();
} catch {
// ignore cleanup errors
}
unsubscribe = null;
}
// Clean up abort listener
if (abortHandler && options?.signal) {
try {
options.signal.removeEventListener("abort", abortHandler);
} catch {
// ignore cleanup errors
}
abortHandler = null;
}
};
const finish = (permission: AgentPermissionRequest | null) => {
cleanup();
resolve({ status: currentStatus, permission });
resolve({
status: currentStatus,
permission,
lastMessage: this.getLastAssistantMessage(agentId)
});
};
const unsubscribe = this.subscribe(
// Bug #3 Fix: Set up abort handler BEFORE subscription
// to ensure cleanup handlers exist before callback can fire
if (options?.signal) {
abortHandler = () => {
cleanup();
reject(createAbortError(options.signal, "wait_for_agent aborted"));
};
options.signal.addEventListener("abort", abortHandler, { once: true });
}
// Bug #3 Fix: Now subscribe with cleanup handlers already in place
// This prevents race condition if callback fires synchronously with replayState: true
unsubscribe = this.subscribe(
(event) => {
// Bug #2 Fix: Only handle agent_state events, remove redundant agent_stream handling
if (event.type === "agent_state") {
currentStatus = event.agent.status;
const pending = event.agent.pendingPermissions[0] ?? null;
@@ -545,48 +620,10 @@ export class AgentManager {
if (!isAgentBusy(event.agent.status)) {
finish(null);
}
return;
}
if (event.type !== "agent_stream") {
return;
}
switch (event.event.type) {
case "permission_requested": {
currentStatus = "running";
finish(event.event.request);
break;
}
case "turn_completed": {
currentStatus = "idle";
finish(null);
break;
}
case "turn_failed": {
currentStatus = "error";
finish(null);
break;
}
default:
break;
}
},
{ agentId, replayState: true }
);
cleanupFns.push(unsubscribe);
if (options?.signal) {
const abortHandler = () => {
cleanup();
reject(createAbortError(options.signal, "wait_for_agent aborted"));
};
options.signal.addEventListener("abort", abortHandler, { once: true });
cleanupFns.push(() =>
options.signal?.removeEventListener("abort", abortHandler)
);
}
});
}

View File

@@ -96,27 +96,6 @@ async function serializeSnapshotWithMetadata(
return serializeAgentSnapshot(snapshot, { title });
}
function buildActivityPayload(
agentManager: AgentManager,
agentId: string
): {
format: "curated";
updateCount: number;
currentModeId: string | null;
content: string;
} {
const timeline = agentManager.getTimeline(agentId);
const snapshot = agentManager.getAgent(agentId);
const curatedText = curateAgentActivity(timeline);
return {
format: "curated",
updateCount: timeline.length,
currentModeId: snapshot?.currentModeId ?? null,
content: curatedText,
};
}
export async function createAgentMcpServer(
options: AgentMcpServerOptions
): Promise<McpServer> {
@@ -159,6 +138,13 @@ export async function createAgentMcpServer(
.describe(
"Optional git worktree branch name (lowercase alphanumerics + hyphen)."
),
background: z
.boolean()
.optional()
.default(false)
.describe(
"Run agent in background. If false (default), waits for completion or permission request. If true, returns immediately."
),
},
outputSchema: {
agentId: z.string(),
@@ -173,9 +159,11 @@ export async function createAgentMcpServer(
description: z.string().nullable().optional(),
})
),
lastMessage: z.string().nullable().optional(),
permission: AgentPermissionRequestPayloadSchema.nullable().optional(),
},
},
async ({ cwd, agentType, initialPrompt, initialMode, worktreeName }) => {
async ({ cwd, agentType, initialPrompt, initialMode, worktreeName, background = false }) => {
let resolvedCwd = expandPath(cwd);
if (worktreeName) {
@@ -203,17 +191,43 @@ export async function createAgentMcpServer(
error
);
}
try {
startAgentRun(agentManager, snapshot.id, initialPrompt);
// If not running in background, wait for completion
if (!background) {
const result = await agentManager.waitForAgentEvent(snapshot.id);
const responseData = {
agentId: snapshot.id,
type: provider,
status: result.status,
cwd: snapshot.cwd,
currentModeId: snapshot.currentModeId,
availableModes: snapshot.availableModes,
lastMessage: result.lastMessage,
permission: result.permission,
};
const validJson = ensureValidJson(responseData);
const response = {
content: [],
structuredContent: validJson,
};
return response;
}
} catch (error) {
console.error(
`[Agent MCP] Failed to run initial prompt for ${snapshot.id}:`,
error
);
}
} else {
}
return {
// Return immediately if background=true or no initialPrompt
const response = {
content: [],
structuredContent: ensureValidJson({
agentId: snapshot.id,
@@ -222,8 +236,11 @@ export async function createAgentMcpServer(
cwd: snapshot.cwd,
currentModeId: snapshot.currentModeId,
availableModes: snapshot.availableModes,
lastMessage: null,
permission: null,
}),
};
return response;
}
);
@@ -240,12 +257,7 @@ export async function createAgentMcpServer(
agentId: z.string(),
status: AgentStatusEnum,
permission: AgentPermissionRequestPayloadSchema.nullable(),
activity: z.object({
format: z.literal("curated"),
updateCount: z.number(),
currentModeId: z.string().nullable(),
content: z.string(),
}),
lastMessage: z.string().nullable(),
},
},
async ({ agentId }, { signal }) => {
@@ -293,17 +305,19 @@ export async function createAgentMcpServer(
await agentManager.waitForAgentEvent(agentId, {
signal: abortController.signal,
});
const activity = buildActivityPayload(agentManager, agentId);
return {
const validJson = ensureValidJson({
agentId,
status: result.status,
permission: result.permission,
lastMessage: result.lastMessage,
});
const response = {
content: [],
structuredContent: ensureValidJson({
agentId,
status: result.status,
permission: result.permission,
activity,
}),
structuredContent: validJson,
};
return response;
} finally {
cleanup();
}
@@ -323,13 +337,23 @@ export async function createAgentMcpServer(
.string()
.optional()
.describe("Optional mode to set before running the prompt."),
background: z
.boolean()
.optional()
.default(false)
.describe(
"Run agent in background. If false (default), waits for completion or permission request. If true, returns immediately."
),
},
outputSchema: {
success: z.boolean(),
status: AgentStatusEnum,
lastMessage: z.string().nullable().optional(),
permission: AgentPermissionRequestPayloadSchema.nullable().optional(),
},
},
async ({ agentId, prompt, sessionMode }) => {
async ({ agentId, prompt, sessionMode, background = false }) => {
if (sessionMode) {
await agentManager.setAgentMode(agentId, sessionMode);
}
@@ -344,15 +368,43 @@ export async function createAgentMcpServer(
}
startAgentRun(agentManager, agentId, prompt);
// If not running in background, wait for completion
if (!background) {
const result = await agentManager.waitForAgentEvent(agentId);
const responseData = {
success: true,
status: result.status,
lastMessage: result.lastMessage,
permission: result.permission,
};
const validJson = ensureValidJson(responseData);
const response = {
content: [],
structuredContent: validJson,
};
return response;
}
// Return immediately if background=true
const snapshot = agentManager.getAgent(agentId);
return {
content: [],
structuredContent: ensureValidJson({
success: true,
status: snapshot?.status ?? "idle",
}),
const responseData = {
success: true,
status: snapshot?.status ?? "idle",
lastMessage: null,
permission: null,
};
const validJson = ensureValidJson(responseData);
const response = {
content: [],
structuredContent: validJson,
};
return response;
}
);

View File

@@ -491,6 +491,9 @@ class ClaudeAgentSession implements AgentSession {
permissionMode: this.currentMode,
agents: this.defaults?.agents,
canUseTool: this.handlePermissionRequest,
stderr: (data: string) => {
console.error("[ClaudeAgentSDK]", data.trim());
},
...this.config.extra?.claude,
};

View File

@@ -112,13 +112,15 @@ async function main() {
await restorePersistedAgents(agentManager, agentRegistry);
console.log("✓ Global agent manager initialized with persisted agents");
const agentMcpServer = await createAgentMcpServer({
agentManager,
agentRegistry,
});
const agentMcpTransports: AgentMcpTransportMap = new Map();
const createAgentMcpTransport = async () => {
// Create a NEW McpServer instance per session (not shared across sessions)
const agentMcpServer = await createAgentMcpServer({
agentManager,
agentRegistry,
});
const transport = new StreamableHTTPServerTransport({
sessionIdGenerator: () => randomUUID(),
onsessioninitialized: (sessionId) => {
@@ -140,7 +142,9 @@ async function main() {
transport.onerror = (error) => {
console.error("[Agent MCP] Transport error:", error);
};
await agentMcpServer.connect(transport);
return transport;
};