Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
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
5 changes: 5 additions & 0 deletions .changeset/abort-handlers-on-close.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
'@modelcontextprotocol/core': patch
---

Abort in-flight request handlers when the connection closes. Previously, request handlers would continue running after the transport disconnected, wasting resources and preventing proper cleanup. Also fixes `InMemoryTransport.close()` firing `onclose` twice on the initiating side.
7 changes: 7 additions & 0 deletions packages/core/src/shared/protocol.ts
Original file line number Diff line number Diff line change
Expand Up @@ -730,6 +730,9 @@ export abstract class Protocol<ContextT extends BaseContext> {
this._taskProgressTokens.clear();
this._pendingDebouncedNotifications.clear();

Comment thread
claude[bot] marked this conversation as resolved.
const requestHandlerAbortControllers = this._requestHandlerAbortControllers;
this._requestHandlerAbortControllers = new Map();
Comment thread
felixweinberger marked this conversation as resolved.

const error = new SdkError(SdkErrorCode.ConnectionClosed, 'Connection closed');

this._transport = undefined;
Comment thread
claude[bot] marked this conversation as resolved.
Expand All @@ -738,6 +741,10 @@ export abstract class Protocol<ContextT extends BaseContext> {
for (const handler of responseHandlers.values()) {
handler(error);
}

for (const controller of requestHandlerAbortControllers.values()) {
controller.abort(error);
}
Comment thread
claude[bot] marked this conversation as resolved.
Outdated
}

private _onerror(error: Error): void {
Expand Down
4 changes: 4 additions & 0 deletions packages/core/src/util/inMemory.ts
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ interface QueuedMessage {
export class InMemoryTransport implements Transport {
private _otherTransport?: InMemoryTransport;
private _messageQueue: QueuedMessage[] = [];
private _closed = false;

onclose?: () => void;
onerror?: (error: Error) => void;
Expand All @@ -39,6 +40,9 @@ export class InMemoryTransport implements Transport {
}

async close(): Promise<void> {
if (this._closed) return;
this._closed = true;

const other = this._otherTransport;
this._otherTransport = undefined;
await other?.close();
Comment thread
claude[bot] marked this conversation as resolved.
Outdated
Expand Down
36 changes: 36 additions & 0 deletions packages/core/test/inMemory.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,42 @@ describe('InMemoryTransport', () => {
await expect(clientTransport.send({ jsonrpc: '2.0', method: 'test', id: 1 })).rejects.toThrow('Not connected');
});

test('should fire onclose exactly once per transport', async () => {
let clientCloseCount = 0;
let serverCloseCount = 0;

clientTransport.onclose = () => clientCloseCount++;
serverTransport.onclose = () => serverCloseCount++;

await clientTransport.close();

expect(clientCloseCount).toBe(1);
expect(serverCloseCount).toBe(1);
});

test('should handle double close idempotently', async () => {
let clientCloseCount = 0;
clientTransport.onclose = () => clientCloseCount++;

await clientTransport.close();
await clientTransport.close();

expect(clientCloseCount).toBe(1);
});

test('should handle concurrent close from both sides', async () => {
let clientCloseCount = 0;
let serverCloseCount = 0;

clientTransport.onclose = () => clientCloseCount++;
serverTransport.onclose = () => serverCloseCount++;

await Promise.all([clientTransport.close(), serverTransport.close()]);

expect(clientCloseCount).toBe(1);
expect(serverCloseCount).toBe(1);
});

test('should queue messages sent before start', async () => {
const message: JSONRPCMessage = {
jsonrpc: '2.0',
Expand Down
30 changes: 30 additions & 0 deletions packages/core/test/shared/protocol.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -209,6 +209,36 @@ describe('protocol tests', () => {
expect(oncloseMock).toHaveBeenCalled();
});

test('should abort in-flight request handlers when the connection is closed', async () => {
await protocol.connect(transport);

let abortReason: unknown;
let handlerStarted = false;
const handlerDone = new Promise<void>(resolve => {
protocol.setRequestHandler('ping', async (_request, ctx) => {
handlerStarted = true;
await new Promise<void>(resolveInner => {
ctx.mcpReq.signal.addEventListener('abort', () => {
abortReason = ctx.mcpReq.signal.reason;
resolveInner();
});
});
resolve();
return {};
});
});

transport.onmessage?.({ jsonrpc: '2.0', id: 1, method: 'ping', params: {} });

await vi.waitFor(() => expect(handlerStarted).toBe(true));

await transport.close();
await handlerDone;

expect(abortReason).toBeInstanceOf(SdkError);
expect((abortReason as SdkError).code).toBe(SdkErrorCode.ConnectionClosed);
});

test('should not overwrite existing hooks when connecting transports', async () => {
const oncloseMock = vi.fn();
const onerrorMock = vi.fn();
Expand Down
Loading