diff --git a/packages/typescript/ai-client/src/chat-client.ts b/packages/typescript/ai-client/src/chat-client.ts index 3cb0f4f7..b80acfd0 100644 --- a/packages/typescript/ai-client/src/chat-client.ts +++ b/packages/typescript/ai-client/src/chat-client.ts @@ -430,6 +430,10 @@ export class ChatClient { this.setStatus('submitted') this.setError(undefined) this.abortController = new AbortController() + // Capture the signal immediately so that a concurrent stop() or + // sendMessage() that reassigns this.abortController cannot cause + // connect() to receive a stale or null signal. + const signal = this.abortController.signal // Reset pending tool executions for the new stream this.pendingToolExecutions.clear() let streamCompletedSuccessfully = false @@ -456,7 +460,7 @@ export class ChatClient { const stream = this.connection.connect( messages, mergedBody, - this.abortController.signal, + signal, ) await this.processStream(stream) diff --git a/packages/typescript/ai-client/tests/chat-client-abort.test.ts b/packages/typescript/ai-client/tests/chat-client-abort.test.ts index 2adffb1c..48b694b4 100644 --- a/packages/typescript/ai-client/tests/chat-client-abort.test.ts +++ b/packages/typescript/ai-client/tests/chat-client-abort.test.ts @@ -306,4 +306,101 @@ describe('ChatClient - Abort Signal Handling', () => { // Each should be a different signal instance expect(abortSignals[0]).not.toBe(abortSignals[1]) }) + + it('should pass the original signal to connect() even if stop() is called during onResponse', async () => { + let signalPassedToConnect: AbortSignal | undefined + + const adapter: ConnectionAdapter = { + // eslint-disable-next-line @typescript-eslint/require-await + async *connect(_messages, _data, abortSignal) { + signalPassedToConnect = abortSignal + yield { + type: 'RUN_FINISHED', + runId: 'run-1', + model: 'test', + timestamp: Date.now(), + finishReason: 'stop', + } + }, + } + + const client = new ChatClient({ + connection: adapter, + onResponse: () => { + // Simulate a concurrent stop() during the onResponse callback, + // which sets this.abortController to null. Without the fix, + // the code would dereference this.abortController.signal after + // this point and crash with a null reference. + client.stop() + }, + }) + + await client.append({ + id: 'user-1', + role: 'user', + parts: [{ type: 'text', content: 'Hello' }], + createdAt: new Date(), + }) + + // The signal should still be a valid AbortSignal instance + // (captured before the await), not undefined/null + expect(signalPassedToConnect).toBeInstanceOf(AbortSignal) + }) + + it('should pass the original signal to connect() even if sendMessage() reassigns abortController during onResponse', async () => { + const signalsPassedToConnect: Array = [] + let secondAppendPromise: Promise | undefined + + const adapter: ConnectionAdapter = { + // eslint-disable-next-line @typescript-eslint/require-await + async *connect(_messages, _data, abortSignal) { + if (abortSignal) { + signalsPassedToConnect.push(abortSignal) + } + yield { + type: 'RUN_FINISHED', + runId: 'run-1', + model: 'test', + timestamp: Date.now(), + finishReason: 'stop', + } + }, + } + + let firstCall = true + const client = new ChatClient({ + connection: adapter, + onResponse: () => { + if (firstCall) { + firstCall = false + // Trigger a second message during onResponse callback. + // This queues a new streamResponse that would create a new + // AbortController, potentially overwriting this.abortController + // before the first connect() call reads the signal. + secondAppendPromise = client.append({ + id: 'user-2', + role: 'user', + parts: [{ type: 'text', content: 'Second message' }], + createdAt: new Date(), + }) + } + }, + }) + + await client.append({ + id: 'user-1', + role: 'user', + parts: [{ type: 'text', content: 'Hello' }], + createdAt: new Date(), + }) + + // Deterministically wait for the queued second stream + await secondAppendPromise + + // Both calls should have received valid, distinct AbortSignal instances + expect(signalsPassedToConnect.length).toBe(2) + expect(signalsPassedToConnect[0]).toBeInstanceOf(AbortSignal) + expect(signalsPassedToConnect[1]).toBeInstanceOf(AbortSignal) + expect(signalsPassedToConnect[0]).not.toBe(signalsPassedToConnect[1]) + }) })