diff --git a/.changeset/calm-streams-stop.md b/.changeset/calm-streams-stop.md new file mode 100644 index 0000000000..76cd7d53dd --- /dev/null +++ b/.changeset/calm-streams-stop.md @@ -0,0 +1,5 @@ +--- +'@workflow/core': patch +--- + +Stop framed stream reconnects after the consumer cancels, including while completion checks or reconnect acquisition are still pending. diff --git a/packages/core/src/reconnecting-framed-stream.test.ts b/packages/core/src/reconnecting-framed-stream.test.ts index fb64466c36..52f499ac75 100644 --- a/packages/core/src/reconnecting-framed-stream.test.ts +++ b/packages/core/src/reconnecting-framed-stream.test.ts @@ -290,6 +290,110 @@ describe('createReconnectingFramedStream', () => { expect(cancelSpy).toHaveBeenCalled(); }); + it('does not reconnect when canceled during completion verification', async () => { + const infoStarted = Promise.withResolvers(); + const infoGate = Promise.withResolvers(); + const { world, calls } = makeWorldWithScriptedStreams( + { + 0: () => + scriptedStream([ + { kind: 'value', value: payloadFrame(1) }, + { kind: 'close' }, + ]), + }, + async () => { + infoStarted.resolve(); + await infoGate.promise; + return { tailIndex: 0, done: false }; + } + ); + setWorld(world); + + const reader = createReconnectingFramedStream(RUN_ID, 's', 0).getReader(); + expect((await reader.read()).value).toEqual(payloadFrame(1)); + const pendingRead = reader.read(); + await infoStarted.promise; + + await reader.cancel('client abort'); + infoGate.resolve(); + await pendingRead; + + expect(calls).toEqual([0]); + }); + + it('cancels a reconnect source acquired after the consumer cancels', async () => { + const reconnectStarted = Promise.withResolvers(); + const reconnectGate = Promise.withResolvers(); + const cancelSpy = vi.fn(); + let connections = 0; + const readFromStream = vi.fn(async () => { + connections++; + if (connections === 1) { + return scriptedStream([ + { kind: 'value', value: payloadFrame(1) }, + { kind: 'error', err: new Error('connection dropped') }, + ]); + } + reconnectStarted.resolve(); + await reconnectGate.promise; + return new ReadableStream({ + async pull() { + await new Promise(() => {}); + }, + cancel(reason) { + cancelSpy(reason); + }, + }); + }); + const world = { readFromStream } as unknown as World; + setWorld(world); + + const reader = createReconnectingFramedStream(RUN_ID, 's', 0).getReader(); + expect((await reader.read()).value).toEqual(payloadFrame(1)); + const pendingRead = reader.read(); + await reconnectStarted.promise; + + await reader.cancel('client abort'); + reconnectGate.resolve(); + await pendingRead; + + await vi.waitFor(() => { + expect(cancelSpy).toHaveBeenCalledWith('client abort'); + }); + expect(readFromStream).toHaveBeenCalledTimes(2); + }); + + it('stops retrying when a pending reconnect rejects after cancellation', async () => { + const reconnectStarted = Promise.withResolvers(); + const reconnectGate = Promise.withResolvers(); + let connections = 0; + const readFromStream = vi.fn(async () => { + connections++; + if (connections === 1) { + return scriptedStream([ + { kind: 'value', value: payloadFrame(1) }, + { kind: 'error', err: new Error('connection dropped') }, + ]); + } + reconnectStarted.resolve(); + await reconnectGate.promise; + throw new Error('reconnect failed'); + }); + const world = { readFromStream } as unknown as World; + setWorld(world); + + const reader = createReconnectingFramedStream(RUN_ID, 's', 0).getReader(); + expect((await reader.read()).value).toEqual(payloadFrame(1)); + const pendingRead = reader.read(); + await reconnectStarted.promise; + + await reader.cancel('client abort'); + reconnectGate.resolve(); + await pendingRead; + + expect(readFromStream).toHaveBeenCalledTimes(2); + }); + it('emits every complete frame packed into a single read', async () => { // One transport read carrying three back-to-back frames must surface as // three separate downstream chunks — exercises the inner drain loop. diff --git a/packages/core/src/serialization.ts b/packages/core/src/serialization.ts index 10a562e7a0..6492d74779 100644 --- a/packages/core/src/serialization.ts +++ b/packages/core/src/serialization.ts @@ -710,15 +710,23 @@ export function createReconnectingFramedStream( let reconnectCount = 0; let totalReconnectCount = 0; let reader: ReadableStreamDefaultReader | undefined; + let canceled = false; + let cancelReason: unknown; let buffer = new Uint8Array(0); - async function connect(): Promise { + async function connect(): Promise { + if (canceled) return false; const world = getWorld(); const effectiveStartIndex = reconnectSupported ? currentStartIndex + consumedFrames : startIndex; const stream = await world.readFromStream(name, effectiveStartIndex); + if (canceled) { + await stream.cancel(cancelReason).catch(() => {}); + return false; + } reader = stream.getReader(); + return true; } /** @@ -739,11 +747,13 @@ export function createReconnectingFramedStream( } } - async function reconnect(): Promise { + async function reconnect(): Promise { + if (canceled) return false; if (reader) { await reader.cancel().catch(() => {}); reader = undefined; } + if (canceled) return false; // Advance the resume position past the frames already delivered, then // drop any partial-frame bytes — the reopened connection re-sends from a // frame boundary at the new index. @@ -771,9 +781,10 @@ export function createReconnectingFramedStream( ); } try { - await connect(); - return; + if (!(await connect())) return false; + return true; } catch { + if (canceled) return false; // Reopen failed transiently; loop to retry, counting against the // budget so a server that never recovers still terminates the stream. } @@ -782,14 +793,16 @@ export function createReconnectingFramedStream( return new ReadableStream({ pull: async (controller) => { + if (canceled) return; // Loop until we emit something, hit EOF, or fatally error. Reads that // only extend the in-flight-frame buffer don't enqueue anything — we // keep reading rather than returning empty-handed. for (;;) { if (!reader) { try { - await connect(); + if (!(await connect())) return; } catch (err) { + if (canceled) return; controller.error(err); return; } @@ -800,12 +813,13 @@ export function createReconnectingFramedStream( // biome-ignore lint/style/noNonNullAssertion: connect() guarantees reader result = await reader!.read(); } catch (err) { + if (canceled) return; if (!reconnectSupported) { controller.error(err); return; } try { - await reconnect(); + if (!(await reconnect())) return; } catch (reconnectErr) { controller.error(reconnectErr); return; @@ -813,6 +827,7 @@ export function createReconnectingFramedStream( continue; } + if (canceled) return; if (result.done || !result.value) { reader = undefined; // A clean EOF is only trustworthy if the stream is actually @@ -822,9 +837,12 @@ export function createReconnectingFramedStream( // errored body, but on some paths it reaches the client as a clean // EOF), and a completed stream can still be cut mid-body — both // would otherwise be silently read as a shorter, complete stream. - if (reconnectSupported && !(await isVerifiedComplete())) { + const verifiedComplete = + !reconnectSupported || (await isVerifiedComplete()); + if (canceled) return; + if (!verifiedComplete) { try { - await reconnect(); + if (!(await reconnect())) return; } catch (reconnectErr) { controller.error(reconnectErr); return; @@ -873,12 +891,15 @@ export function createReconnectingFramedStream( // Only partial bytes — read more. } }, - cancel: async () => { - if (reader) { - await reader.cancel().catch((err) => { + cancel: async (reason) => { + canceled = true; + cancelReason = reason; + const currentReader = reader; + reader = undefined; + if (currentReader) { + await currentReader.cancel(reason).catch((err) => { console.warn('Error closing ReadableStream reader:', err); }); - reader = undefined; } }, });