diff --git a/.changeset/ws-stream-writer-session.md b/.changeset/ws-stream-writer-session.md new file mode 100644 index 0000000000..f8a9431086 --- /dev/null +++ b/.changeset/ws-stream-writer-session.md @@ -0,0 +1,6 @@ +--- +'@workflow/core': patch +'@workflow/world': minor +--- + +Add an optional stateful stream writer-session seam with stable writer identity and sequence tracking. diff --git a/packages/core/src/serialization.ts b/packages/core/src/serialization.ts index aa9778e7cc..0fc3b2b1a0 100644 --- a/packages/core/src/serialization.ts +++ b/packages/core/src/serialization.ts @@ -6,6 +6,7 @@ import { WorkflowRuntimeError, } from '@workflow/errors'; import { once } from '@workflow/utils'; +import type { StreamWriteSession } from '@workflow/world'; import { envNumber } from '@workflow/world/env-config'; import { parse, stringify, unflatten } from 'devalue'; import { monotonicFactory } from 'ulid'; @@ -1310,6 +1311,20 @@ export class WorkflowServerWritableStream extends WritableStream { } }; + // One stable identity and sequence space per in-memory sink. Start session + // construction as soon as the run-ready barrier permits, so transports may + // negotiate eagerly without allowing a write to overtake run creation. + // This eager chain cannot strand a new rejection: ensureRunReady absorbs its + // ordering-only failure, and every path that can observe worldPromise also + // awaits this session promise before it writes, closes, or disposes. + const writerId = `wrtr_${defaultUlid()}` as const; + const writeSessionPromise: Promise = + ensureRunReady().then(async () => { + const world = await worldPromise; + return world.streams.createWriteSession?.(runId, name, { writerId }); + }); + let nextChunkSeq = 0; + // ------------------------------------------------------------------ // Group-commit buffering. // @@ -1427,8 +1442,14 @@ export class WorkflowServerWritableStream extends WritableStream { ): Promise => { await ensureRunReady(); const world = await worldPromise; + const session = await writeSessionPromise; const dispatchAt = Date.now(); - if (typeof world.streams.writeMulti === 'function' && group.length > 1) { + if (session) { + await session.write(nextChunkSeq, group); + } else if ( + typeof world.streams.writeMulti === 'function' && + group.length > 1 + ) { await world.streams.writeMulti(runId, name, group); } else { // Fall back to sequential writes @@ -1436,6 +1457,9 @@ export class WorkflowServerWritableStream extends WritableStream { await world.streams.write(runId, name, chunk); } } + // `inFlight` admits only one dispatch loop, so no second group can read + // this sequence space until the current group has advanced it. + nextChunkSeq += group.length; if (groupT0 !== undefined) { recordStreamWriteFlush( groupT0, @@ -1635,8 +1659,13 @@ export class WorkflowServerWritableStream extends WritableStream { await ensureRunReady(); const world = await worldPromise; + const session = await writeSessionPromise; const closeStart = Date.now(); - await world.streams.close(runId, name); + if (session) { + await session.close(); + } else { + await world.streams.close(runId, name); + } recordStreamClose(closeStart, runId, name); }, async abort(reason) { @@ -1671,6 +1700,11 @@ export class WorkflowServerWritableStream extends WritableStream { // Reject blocked writers and drain waiters so nothing leaks or // hangs on a promise whose timer was just cleared. rejectWaiters(sinkError); + // Abort is transport cleanup, not semantic stream completion. A + // stateful World releases its socket without sending close; stateless + // Worlds keep the existing no-op behavior. + const session = await writeSessionPromise.catch(() => undefined); + await session?.dispose?.(); }, }); diff --git a/packages/core/src/writable-stream.test.ts b/packages/core/src/writable-stream.test.ts index 8dcbe253bd..539ce5dc8b 100644 --- a/packages/core/src/writable-stream.test.ts +++ b/packages/core/src/writable-stream.test.ts @@ -31,6 +31,7 @@ async function waitFor( describe('WorkflowServerWritableStream', () => { let mockStreams: { + createWriteSession?: ReturnType; write: ReturnType; writeMulti: ReturnType; close: ReturnType; @@ -516,6 +517,82 @@ describe('WorkflowServerWritableStream', () => { }); }); + describe('stateful writer sessions', () => { + it('uses one stable writer id and writer-local sequence across groups', async () => { + const session = { + write: vi.fn().mockResolvedValue(undefined), + close: vi.fn().mockResolvedValue(undefined), + }; + mockStreams.createWriteSession = vi.fn(() => session); + + const stream = new WorkflowServerWritableStream('run-123', 'test-stream'); + const writer = stream.getWriter(); + await writer.write(new Uint8Array([1])); + await waitFor(() => expect(session.write).toHaveBeenCalledTimes(1)); + await writer.write(new Uint8Array([2])); + await waitFor(() => expect(session.write).toHaveBeenCalledTimes(2)); + await writer.close(); + + expect(mockStreams.createWriteSession).toHaveBeenCalledTimes(1); + expect(mockStreams.createWriteSession).toHaveBeenCalledWith( + 'run-123', + 'test-stream', + { + writerId: expect.stringMatching( + /^wrtr_[0123456789ABCDEFGHJKMNPQRSTVWXYZ]{26}$/ + ), + } + ); + expect(session.write).toHaveBeenNthCalledWith(1, 0, [ + new Uint8Array([1]), + ]); + expect(session.write).toHaveBeenNthCalledWith(2, 1, [ + new Uint8Array([2]), + ]); + expect(session.close).toHaveBeenCalledTimes(1); + expect(mockStreams.write).not.toHaveBeenCalled(); + expect(mockStreams.writeMulti).not.toHaveBeenCalled(); + expect(mockStreams.close).not.toHaveBeenCalled(); + }); + + it('creates a new writer id for each in-memory sink lifetime', async () => { + const writerIds: string[] = []; + mockStreams.createWriteSession = vi.fn((_runId, _name, options) => { + writerIds.push(options.writerId); + return { + write: vi.fn().mockResolvedValue(undefined), + close: vi.fn().mockResolvedValue(undefined), + }; + }); + + new WorkflowServerWritableStream('run-123', 'test-stream'); + new WorkflowServerWritableStream('run-123', 'test-stream'); + await waitFor(() => expect(writerIds).toHaveLength(2)); + expect(new Set(writerIds).size).toBe(2); + }); + }); + + describe('abort cleanup', () => { + it('drains accepted chunks then disposes without semantic close', async () => { + const session = { + write: vi.fn().mockResolvedValue(undefined), + close: vi.fn().mockResolvedValue(undefined), + dispose: vi.fn().mockResolvedValue(undefined), + }; + mockStreams.createWriteSession = vi.fn(() => session); + + const stream = new WorkflowServerWritableStream('run-123', 'test-stream'); + const writer = stream.getWriter(); + await writer.write(new Uint8Array([1])); + await writer.abort(new Error('producer failed')); + + expect(session.write).toHaveBeenCalledWith(0, [new Uint8Array([1])]); + expect(session.dispose).toHaveBeenCalledTimes(1); + expect(session.close).not.toHaveBeenCalled(); + expect(mockStreams.close).not.toHaveBeenCalled(); + }); + }); + describe('error handling', () => { it('surfaces a dispatch failure at the durability barrier (close) and poisons later writes', async () => { mockStreams.write.mockRejectedValueOnce(new Error('write error')); @@ -549,6 +626,28 @@ describe('WorkflowServerWritableStream', () => { expect(mockStreams.write).toHaveBeenCalledTimes(1); }); + it('keeps a session failure sticky without advancing its sequence', async () => { + const session = { + write: vi.fn().mockRejectedValueOnce(new Error('session failed')), + close: vi.fn().mockResolvedValue(undefined), + }; + mockStreams.createWriteSession = vi.fn(() => session); + + const writer = new WorkflowServerWritableStream( + 'run-123', + 'test-stream' + ).getWriter(); + await writer.write(new Uint8Array([1])); + await waitFor(() => expect(session.write).toHaveBeenCalledTimes(1)); + await expect(writer.write(new Uint8Array([2]))).rejects.toThrow( + 'session failed' + ); + await expect(writer.close()).rejects.toThrow(); + expect(session.write).toHaveBeenCalledTimes(1); + expect(session.write).toHaveBeenCalledWith(0, [new Uint8Array([1])]); + expect(session.close).not.toHaveBeenCalled(); + }); + it('should propagate close errors', async () => { mockStreams.close.mockRejectedValueOnce(new Error('close error')); diff --git a/packages/world/src/interfaces.ts b/packages/world/src/interfaces.ts index e688353996..47a3b4fc2a 100644 --- a/packages/world/src/interfaces.ts +++ b/packages/world/src/interfaces.ts @@ -40,6 +40,25 @@ import type { StepWithoutData, } from './steps.js'; +export interface StreamWriteSession { + /** + * Write one ordered group from this in-memory writer lifetime. + * `chunkSeq` is writer-local and identifies the first chunk in `chunks`. + */ + write(chunkSeq: number, chunks: (string | Uint8Array)[]): Promise; + + /** Close this writer lifetime after all prior writes are durable. */ + close(): Promise; + + /** Release transport resources without semantically closing the stream. */ + dispose?(): Promise | void; +} + +export interface CreateStreamWriteSessionOptions { + /** Stable observational id for this in-memory writer lifetime. */ + writerId: `wrtr_${string}`; +} + export interface Streamer { /** * Number of milliseconds a stream waits for additional chunks to arrive @@ -57,6 +76,17 @@ export interface Streamer { streamFlushIntervalMs?: number; streams: { + /** + * Optionally create a stateful writer session. Core creates at most one + * session per in-memory WritableStream and otherwise uses the stateless + * write/writeMulti/close methods below unchanged. + */ + createWriteSession?( + runId: string, + name: string, + options: CreateStreamWriteSessionOptions + ): StreamWriteSession; + write( runId: string, name: string,