diff --git a/packages/cli/src/ui/hooks/atCommandProcessor.test.ts b/packages/cli/src/ui/hooks/atCommandProcessor.test.ts index 9b1f9635a..4d9c949e3 100644 --- a/packages/cli/src/ui/hooks/atCommandProcessor.test.ts +++ b/packages/cli/src/ui/hooks/atCommandProcessor.test.ts @@ -247,6 +247,7 @@ describe('handleAtCommand', () => { expect(registry.readMcpResource).toHaveBeenCalledWith( 'github', 'github://repos/owner/repo/issues', + expect.objectContaining({ signal: abortController.signal }), ); expect(mockAddItem).toHaveBeenCalledWith( expect.objectContaining({ @@ -298,6 +299,7 @@ describe('handleAtCommand', () => { expect(registry.readMcpResource).toHaveBeenCalledWith( 'github', 'github://repos/owner/repo/issues', + expect.objectContaining({ signal: abortController.signal }), ); expect(mockAddItem).toHaveBeenCalledWith( expect.objectContaining({ @@ -349,6 +351,7 @@ describe('handleAtCommand', () => { expect(registry.readMcpResource).toHaveBeenCalledWith( 'github', 'github://repos/owner/repo/issues', + expect.objectContaining({ signal: abortController.signal }), ); expect(mockAddItem).toHaveBeenCalledWith( expect.objectContaining({ @@ -417,6 +420,7 @@ describe('handleAtCommand', () => { expect(readMcpResourceSpy).toHaveBeenCalledWith( 'github', 'github://repos/owner/repo/issues', + expect.objectContaining({ signal: abortController.signal }), ); expect(mockAddItem).toHaveBeenCalledWith( expect.objectContaining({ @@ -473,6 +477,7 @@ describe('handleAtCommand', () => { expect(registry.readMcpResource).toHaveBeenCalledWith( 'github', 'github://repos/owner/repo/issues', + expect.objectContaining({ signal: abortController.signal }), ); expect(mockAddItem).toHaveBeenCalledWith( expect.objectContaining({ diff --git a/packages/cli/src/ui/hooks/atCommandProcessor.ts b/packages/cli/src/ui/hooks/atCommandProcessor.ts index a5ea3a46d..a7fb2fe36 100644 --- a/packages/cli/src/ui/hooks/atCommandProcessor.ts +++ b/packages/cli/src/ui/hooks/atCommandProcessor.ts @@ -693,37 +693,17 @@ export async function handleAtCommand({ const ref = mcpResourceRefs[i]; let resourceResult: unknown; try { - resourceResult = await new Promise((resolve, reject) => { - if (signal.aborted) { - const error = new Error('MCP resource read aborted'); - error.name = 'AbortError'; - reject(error); - return; - } + if (signal.aborted) { + const error = new Error('MCP resource read aborted'); + error.name = 'AbortError'; + throw error; + } - const onAbort = () => { - cleanup(); - const error = new Error('MCP resource read aborted'); - error.name = 'AbortError'; - reject(error); - }; - const cleanup = () => { - signal.removeEventListener('abort', onAbort); - }; - - signal.addEventListener('abort', onAbort, { once: true }); - - toolRegistry - .readMcpResource(ref.serverName, ref.uri) - .then((res) => { - cleanup(); - resolve(res); - }) - .catch((err) => { - cleanup(); - reject(err); - }); - }); + resourceResult = await toolRegistry.readMcpResource( + ref.serverName, + ref.uri, + { signal }, + ); toolDisplays.push({ callId: `client-mcp-resource-${userMessageTimestamp}-${i}`, diff --git a/packages/core/src/tools/mcp-client.ts b/packages/core/src/tools/mcp-client.ts index 8277cd63e..cfae34506 100644 --- a/packages/core/src/tools/mcp-client.ts +++ b/packages/core/src/tools/mcp-client.ts @@ -196,7 +196,10 @@ export class McpClient { return this.status; } - async readResource(uri: string): Promise { + async readResource( + uri: string, + options?: { signal?: AbortSignal }, + ): Promise { if (this.status !== MCPServerStatus.CONNECTED) { throw new Error('Client is not connected.'); } @@ -209,6 +212,7 @@ export class McpClient { return this.client.request( { method: 'resources/read', params: { uri } }, ReadResourceResultSchema, + options, ); } diff --git a/packages/core/src/tools/tool-registry.ts b/packages/core/src/tools/tool-registry.ts index c8abf5ee5..a2096a2a5 100644 --- a/packages/core/src/tools/tool-registry.ts +++ b/packages/core/src/tools/tool-registry.ts @@ -474,12 +474,13 @@ export class ToolRegistry { async readMcpResource( serverName: string, uri: string, + options?: { signal?: AbortSignal }, ): Promise { if (!this.config.isTrustedFolder()) { throw new Error('MCP resources are unavailable in untrusted folders.'); } - return this.mcpClientManager.readResource(serverName, uri); + return this.mcpClientManager.readResource(serverName, uri, options); } /**