1// Alchemy modifications are licensed under Apache-2.0.
2// This file includes third-party code; see /THIRD_PARTY_LICENSES.md.
3import {
4 InvalidStepReadableStreamError,
5 OversizedStreamChunkError,
6 StreamOutputStorageLimitError,
7 UnsupportedStreamChunkError,
8 WorkflowTimeoutError,
9} from "./errors.ts";
10
11// ── Constants ───────────────────────────────────────────────────────────────
12
13export const DEFAULT_STREAM_OUTPUT_CHUNK_SIZE = 256 * 1024;
14export const STREAM_OUTPUT_META_SUFFIX = "-value-stream-meta";
15export const MAX_STREAM_OUTPUT_INPUT_CHUNK_BYTES = 16 * 1024 * 1024;
16export const STREAMING_STEP_CHUNKS_TABLE = "streaming_step_chunks";
17export const MAX_OUTPUT_SHOWN_IN_LOGS = 1024;
18
19const DO_STORAGE_LIMIT = 1024 * 1024 * 1024 + 100 * 1024 * 1024;
20const STREAM_OUTPUT_STORAGE_WRITE_HEADROOM_BYTES = 16 * 1024;
21
22// ── Types ───────────────────────────────────────────────────────────────────
23
24export enum StreamOutputState {
25 Pending = "pending",
26 Committing = "committing",
27 Complete = "complete",
28}
29
30export type StreamOutputMeta = {
31 version: 1;
32 state: StreamOutputState;
33 attempt: number;
34 startedAt: number;
35 chunkCount: number;
36 totalBytes: number;
37 committedAt: number | null;
38};
39
40export type StoredStreamOutputPreview = { type: "text"; output: string } | { type: "binary" };
41
42export class InvalidStoredStreamOutputError extends Error {
43 name = "InvalidStoredStreamOutputError";
44}
45
46// ── Helpers ─────────────────────────────────────────────────────────────────
47
48export function getStreamOutputMetaKey(cacheKey: string): string {
49 return `${cacheKey}${STREAM_OUTPUT_META_SUFFIX}`;
50}
51
52export function isReadableStreamLike(value: unknown): value is ReadableStream<unknown> {
53 return value instanceof ReadableStream;
54}
55
56function createInvalidStepReadableStreamError(): InvalidStepReadableStreamError {
57 return new InvalidStepReadableStreamError(
58 "Step returned a ReadableStream that is already locked or otherwise unreadable. Return a fresh, unlocked ReadableStream from step.do().",
59 );
60}
61
62function createOversizedStreamChunkError(): OversizedStreamChunkError {
63 return new OversizedStreamChunkError(
64 `Step returned a ReadableStream chunk larger than the maximum allowed size of ${MAX_STREAM_OUTPUT_INPUT_CHUNK_BYTES} bytes. ` +
65 "Return smaller chunks from step.do().",
66 );
67}
68
69/**
70 * Normalize any incoming chunk to Uint8Array.
71 * Accepts ArrayBuffer, TypedArrays (except DataView), and Uint8Array directly.
72 * Rejects strings, objects, and other non-binary types.
73 */
74function normalizeChunkToUint8Array(value: unknown): Uint8Array {
75 if (value instanceof Uint8Array) {
76 return value;
77 }
78
79 if (value instanceof ArrayBuffer) {
80 return new Uint8Array(value);
81 }
82
83 if (ArrayBuffer.isView(value) && !(value instanceof DataView)) {
84 return new Uint8Array(value.buffer, value.byteOffset, value.byteLength);
85 }
86
87 throw new UnsupportedStreamChunkError(
88 "Step returned a ReadableStream with unsupported chunk type. " +
89 "Only ArrayBuffer and TypedArray chunks are supported.",
90 );
91}
92
93// ── Buffer helpers ──────────────────────────────────────────────────────────
94
95function takeBufferedBytes(bufferedChunks: Array<Uint8Array>, byteLength: number): Uint8Array {
96 const output = new Uint8Array(byteLength);
97 let offset = 0;
98
99 while (offset < byteLength) {
100 const chunk = bufferedChunks[0];
101 const remaining = byteLength - offset;
102
103 if (chunk.byteLength <= remaining) {
104 output.set(chunk, offset);
105 offset += chunk.byteLength;
106 bufferedChunks.shift();
107 continue;
108 }
109
110 output.set(chunk.subarray(0, remaining), offset);
111 bufferedChunks[0] = chunk.subarray(remaining);
112 offset += remaining;
113 }
114
115 return output;
116}
117
118// ── Stream iteration ────────────────────────────────────────────────────────
119
120async function* iterateStreamChunks(
121 stream: ReadableStream<unknown>,
122 signal?: AbortSignal,
123): AsyncGenerator<Uint8Array> {
124 if (stream.locked) {
125 throw createInvalidStepReadableStreamError();
126 }
127
128 if (signal?.aborted) {
129 throw signal.reason ?? new DOMException("The operation was aborted.", "AbortError");
130 }
131
132 let reader: ReadableStreamDefaultReader<unknown>;
133 try {
134 reader = stream.getReader();
135 } catch (error) {
136 if (error instanceof TypeError) {
137 throw createInvalidStepReadableStreamError();
138 }
139 throw error;
140 }
141
142 const onAbort = () => {
143 void reader
144 .cancel(signal?.reason ?? new DOMException("The operation was aborted.", "AbortError"))
145 .catch(() => {});
146 };
147
148 signal?.addEventListener("abort", onAbort, { once: true });
149 let fullyRead = false;
150
151 try {
152 while (true) {
153 let readResult: ReadableStreamReadResult<unknown>;
154 try {
155 readResult = await reader.read();
156 } catch (readError) {
157 // When the abort signal has already fired, the read error is an
158 // expected cancellation (e.g. step timeout or engine shutdown).
159 // Re-throw the original abort reason so callers can distinguish
160 // timeouts from genuine stream failures.
161 if (signal?.aborted) {
162 throw signal.reason ?? new DOMException("The operation was aborted.", "AbortError");
163 }
164
165 // Any other read failure (broken pipe, upstream connection
166 // drop, encoding mismatch, etc.) means the ReadableStream the
167 // step returned is unusable.
168 throw new InvalidStepReadableStreamError(
169 "Failed to read from step ReadableStream output. " +
170 (readError instanceof Error ? readError.message : String(readError)),
171 );
172 }
173
174 if (signal?.aborted) {
175 throw signal.reason ?? new DOMException("The operation was aborted.", "AbortError");
176 }
177 if (readResult.done) {
178 fullyRead = true;
179 return;
180 }
181
182 yield normalizeChunkToUint8Array(readResult.value);
183 }
184 } finally {
185 signal?.removeEventListener("abort", onAbort);
186 if (!fullyRead) {
187 await reader
188 .cancel(new Error("stream output consumption stopped before completion"))
189 .catch(() => {});
190 }
191 try {
192 reader.releaseLock();
193 } catch {
194 // Reader may still be processing cancel()
195 }
196 }
197}
198
199// ── SQL helpers ─────────────────────────────────────────────────────────────
200
201function deleteAttemptChunks(
202 storage: DurableObjectStorage,
203 cacheKey: string,
204 attempt: number,
205): void {
206 // eslint-disable-next-line workers-sdk/no-unsafe-command-execution -- DO SQL exec, not child_process
207 storage.sql.exec(
208 `DELETE FROM ${STREAMING_STEP_CHUNKS_TABLE} WHERE cache_key = ? AND attempt = ?`,
209 cacheKey,
210 attempt,
211 );
212}
213
214async function deleteMetaForAttempt(
215 storage: DurableObjectStorage,
216 cacheKey: string,
217 attempt: number,
218): Promise<void> {
219 const metaKey = getStreamOutputMetaKey(cacheKey);
220 const maybeMeta = await storage.get<StreamOutputMeta>(metaKey);
221 if (maybeMeta === undefined) {
222 return;
223 }
224 if (maybeMeta.attempt !== attempt) {
225 return;
226 }
227 await storage.delete(metaKey);
228}
229
230// ── Integrity validation ────────────────────────────────────────────────────
231
232type StreamOutputChunkSummary = {
233 chunkCount: number;
234 minChunkIndex: number | null;
235 maxChunkIndex: number | null;
236 totalBytes: number;
237};
238
239function getStreamOutputChunkSummary(
240 storage: DurableObjectStorage,
241 cacheKey: string,
242 attempt: number,
243): StreamOutputChunkSummary {
244 const row = storage.sql
245 .exec<StreamOutputChunkSummary>(
246 [
247 `SELECT`,
248 ` COUNT(*) AS chunkCount,`,
249 ` MIN(chunk_index) AS minChunkIndex,`,
250 ` MAX(chunk_index) AS maxChunkIndex,`,
251 ` CAST(COALESCE(SUM(LENGTH(chunk)), 0) AS INTEGER) AS totalBytes`,
252 `FROM ${STREAMING_STEP_CHUNKS_TABLE}`,
253 `WHERE cache_key = ? AND attempt = ?`,
254 ].join("\n"),
255 cacheKey,
256 attempt,
257 )
258 .one();
259
260 if (row === null) {
261 throw new Error("Expected stream chunk summary query to return a row");
262 }
263 return row;
264}
265
266export function getInvalidStoredStreamOutputError(
267 storage: DurableObjectStorage,
268 cacheKey: string,
269 meta: StreamOutputMeta,
270): InvalidStoredStreamOutputError | undefined {
271 const summary = getStreamOutputChunkSummary(storage, cacheKey, meta.attempt);
272
273 if (meta.chunkCount === 0) {
274 if (
275 summary.chunkCount === 0 &&
276 summary.totalBytes === 0 &&
277 summary.minChunkIndex === null &&
278 summary.maxChunkIndex === null
279 ) {
280 return undefined;
281 }
282 } else if (
283 summary.chunkCount === meta.chunkCount &&
284 summary.minChunkIndex === 0 &&
285 summary.maxChunkIndex === meta.chunkCount - 1 &&
286 summary.totalBytes === meta.totalBytes
287 ) {
288 return undefined;
289 }
290
291 return new InvalidStoredStreamOutputError(
292 `Stored streamed step output is corrupt or incomplete for cache key ${cacheKey}. ` +
293 `Expected ${meta.chunkCount} chunks / ${meta.totalBytes} bytes, found ` +
294 `${summary.chunkCount} chunks / ${summary.totalBytes} bytes with chunk index range ` +
295 `${summary.minChunkIndex ?? "null"}..${summary.maxChunkIndex ?? "null"}.`,
296 );
297}
298
299// ── Preview ─────────────────────────────────────────────────────────────────
300
301function readStreamOutputPreviewBytes(options: {
302 storage: DurableObjectStorage;
303 cacheKey: string;
304 attempt: number;
305 maxBytes: number;
306}): Uint8Array {
307 const { storage, cacheKey, attempt, maxBytes } = options;
308 // eslint-disable-next-line workers-sdk/no-unsafe-command-execution -- DO SQL exec, not child_process
309 const cursor = storage.sql.exec<{
310 chunk_index: number;
311 chunk: ArrayBuffer;
312 }>(
313 `SELECT chunk_index, chunk FROM ${STREAMING_STEP_CHUNKS_TABLE} WHERE cache_key = ? AND attempt = ? ORDER BY chunk_index`,
314 cacheKey,
315 attempt,
316 );
317 const previewChunks: Array<Uint8Array> = [];
318 let expectedChunkIndex = 0;
319 let totalBytes = 0;
320
321 while (totalBytes < maxBytes) {
322 const row = cursor.next();
323 if (row.done) {
324 break;
325 }
326
327 if (row.value.chunk_index !== expectedChunkIndex) {
328 throw new InvalidStoredStreamOutputError(
329 `Missing chunk ${expectedChunkIndex} for streamed step output`,
330 );
331 }
332
333 if (!(row.value.chunk instanceof ArrayBuffer)) {
334 throw new InvalidStoredStreamOutputError(
335 "Invalid chunk type returned from streaming_step_chunks table",
336 );
337 }
338
339 const chunkBytes = new Uint8Array(row.value.chunk);
340 const remainingBytes = maxBytes - totalBytes;
341 const previewChunk =
342 chunkBytes.byteLength > remainingBytes ? chunkBytes.subarray(0, remainingBytes) : chunkBytes;
343
344 previewChunks.push(previewChunk);
345 totalBytes += previewChunk.byteLength;
346 expectedChunkIndex++;
347 }
348
349 return takeBufferedBytes(previewChunks, totalBytes);
350}
351
352export function getStoredStreamOutputPreview(options: {
353 storage: DurableObjectStorage;
354 cacheKey: string;
355 meta: StreamOutputMeta;
356 maxChars: number;
357}): StoredStreamOutputPreview {
358 const { storage, cacheKey, meta, maxChars } = options;
359 if (meta.state !== StreamOutputState.Complete) {
360 throw new Error("Cannot preview streamed step output before it is complete");
361 }
362
363 // UTF-8 uses at most 4 bytes per code point
364 const maxPreviewBytes = maxChars * 4;
365 const previewBytes = readStreamOutputPreviewBytes({
366 storage,
367 cacheKey,
368 attempt: meta.attempt,
369 maxBytes: maxPreviewBytes,
370 });
371 const previewTruncatedByBytes = meta.totalBytes > previewBytes.byteLength;
372
373 try {
374 const decoded = new TextDecoder("utf-8", {
375 fatal: true,
376 ignoreBOM: false,
377 }).decode(previewBytes, { stream: previewTruncatedByBytes });
378
379 const previewOutput = decoded.substring(0, maxChars);
380 if (decoded.length > maxChars || previewTruncatedByBytes) {
381 return { type: "text", output: previewOutput + "[truncated output]" };
382 }
383 return { type: "text", output: previewOutput };
384 } catch {
385 return { type: "binary" };
386 }
387}
388
389// ── Cleanup ─────────────────────────────────────────────────────────────────
390
391export async function cleanupPendingStreamOutput(
392 storage: DurableObjectStorage,
393 cacheKey: string,
394): Promise<void> {
395 const metaKey = getStreamOutputMetaKey(cacheKey);
396 const maybeMeta = await storage.get<StreamOutputMeta>(metaKey);
397
398 if (maybeMeta === undefined) {
399 return;
400 }
401
402 if (maybeMeta.state === StreamOutputState.Complete) {
403 return;
404 }
405
406 await rollbackStreamOutput(storage, cacheKey, maybeMeta.attempt);
407}
408
409export async function rollbackStreamOutput(
410 storage: DurableObjectStorage,
411 cacheKey: string,
412 attempt: number,
413): Promise<void> {
414 deleteAttemptChunks(storage, cacheKey, attempt);
415 await deleteMetaForAttempt(storage, cacheKey, attempt);
416}
417
418// ── Write ───────────────────────────────────────────────────────────────────
419
420async function doWriteStreamOutput(options: {
421 storage: DurableObjectStorage;
422 cacheKey: string;
423 attempt: number;
424 stream: ReadableStream<unknown>;
425 chunkSizeBytes?: number;
426 signal?: AbortSignal;
427 skipMetaWrite?: boolean;
428}): Promise<StreamOutputMeta> {
429 const { storage, cacheKey, attempt, stream, signal, skipMetaWrite } = options;
430 const chunkSizeBytes = options.chunkSizeBytes ?? DEFAULT_STREAM_OUTPUT_CHUNK_SIZE;
431 const metaKey = getStreamOutputMetaKey(cacheKey);
432
433 const maybeInvalidState = (additionalBytes = 0): unknown => {
434 if (signal?.aborted) {
435 return signal.reason ?? new DOMException("The operation was aborted.", "AbortError");
436 }
437
438 const currentStorageBytes = storage.sql.databaseSize;
439 if (
440 currentStorageBytes + additionalBytes + STREAM_OUTPUT_STORAGE_WRITE_HEADROOM_BYTES >
441 DO_STORAGE_LIMIT
442 ) {
443 return new StreamOutputStorageLimitError("The instance has exceeded the 1GiB storage limit");
444 }
445 };
446
447 const initialInvalidState = maybeInvalidState();
448 if (initialInvalidState !== undefined) {
449 throw initialInvalidState;
450 }
451
452 const startedAt = Date.now();
453 if (!skipMetaWrite) {
454 await storage.put(metaKey, {
455 version: 1,
456 state: StreamOutputState.Pending,
457 attempt,
458 startedAt,
459 chunkCount: 0,
460 totalBytes: 0,
461 committedAt: null,
462 } satisfies StreamOutputMeta);
463 }
464
465 let chunkCount = 0;
466 let totalBytes = 0;
467 const bufferedChunks: Array<Uint8Array> = [];
468 let bufferedBytes = 0;
469 let outputCommitted = false;
470
471 const flushChunk = async (bytes: Uint8Array) => {
472 const invalidState = maybeInvalidState(bytes.byteLength);
473 if (invalidState !== undefined) {
474 throw invalidState;
475 }
476 // eslint-disable-next-line workers-sdk/no-unsafe-command-execution -- DO SQL exec, not child_process
477 storage.sql.exec(
478 `INSERT INTO ${STREAMING_STEP_CHUNKS_TABLE} (cache_key, attempt, chunk_index, chunk) VALUES (?, ?, ?, ?)`,
479 cacheKey,
480 attempt,
481 chunkCount,
482 bytes,
483 );
484 totalBytes += bytes.byteLength;
485 chunkCount++;
486 };
487
488 try {
489 for await (const bytes of iterateStreamChunks(stream, signal)) {
490 if (bytes.byteLength === 0) {
491 continue;
492 }
493
494 if (bytes.byteLength > MAX_STREAM_OUTPUT_INPUT_CHUNK_BYTES) {
495 throw createOversizedStreamChunkError();
496 }
497
498 bufferedChunks.push(bytes);
499 bufferedBytes += bytes.byteLength;
500
501 // NOTE: we want chunks with fixed length,
502 // that's why we buffer them in-memory here
503 while (bufferedBytes >= chunkSizeBytes) {
504 const chunk = takeBufferedBytes(bufferedChunks, chunkSizeBytes);
505 bufferedBytes -= chunk.byteLength;
506 await flushChunk(chunk);
507 }
508 }
509
510 // Last chunk (remainder)
511 if (bufferedBytes > 0) {
512 await flushChunk(takeBufferedBytes(bufferedChunks, bufferedBytes));
513 bufferedBytes = 0;
514 }
515
516 const meta = {
517 version: 1,
518 state: StreamOutputState.Complete,
519 attempt,
520 startedAt,
521 chunkCount,
522 totalBytes,
523 committedAt: Date.now(),
524 } satisfies StreamOutputMeta;
525
526 const invalidState = maybeInvalidState();
527 if (invalidState !== undefined) {
528 throw invalidState;
529 }
530
531 if (!skipMetaWrite) {
532 // Transition to Committing to signal it's safe to await.
533 // Mock stream setup still needs the SQL chunks and returned meta, but
534 // must not publish the normal stream cache key because Context.do()
535 // treats it as a completed prior run and skips step execution logging.
536 await storage.put(metaKey, {
537 version: 1,
538 state: StreamOutputState.Committing,
539 attempt,
540 startedAt,
541 chunkCount,
542 totalBytes,
543 committedAt: null,
544 } satisfies StreamOutputMeta);
545
546 await storage.put(metaKey, meta);
547 }
548 outputCommitted = true;
549 return meta;
550 } catch (error) {
551 if (!outputCommitted) {
552 await rollbackStreamOutput(storage, cacheKey, attempt);
553 }
554 throw error;
555 }
556}
557
558export async function writeStreamOutput(options: {
559 storage: DurableObjectStorage;
560 cacheKey: string;
561 attempt: number;
562 stream: ReadableStream<unknown>;
563 chunkSizeBytes?: number;
564 signal?: AbortSignal;
565 timeoutTask?: Promise<never>;
566 skipMetaWrite?: boolean;
567}): Promise<StreamOutputMeta> {
568 const { storage, cacheKey, attempt, timeoutTask, ...writeOptions } = options;
569 const writeTask = doWriteStreamOutput({
570 storage,
571 cacheKey,
572 attempt,
573 ...writeOptions,
574 });
575
576 if (timeoutTask === undefined) {
577 return writeTask;
578 }
579
580 try {
581 return await Promise.race([writeTask, timeoutTask]);
582 } catch (error) {
583 if (error instanceof WorkflowTimeoutError) {
584 if (options.skipMetaWrite) {
585 void writeTask.catch(() => {});
586 throw error;
587 }
588
589 const maybeMeta = await storage.get<StreamOutputMeta>(getStreamOutputMetaKey(cacheKey));
590 if (
591 maybeMeta?.attempt === attempt &&
592 (maybeMeta.state === StreamOutputState.Committing ||
593 maybeMeta.state === StreamOutputState.Complete)
594 ) {
595 // Safe to await -- not in the middle of writing chunks
596 return await writeTask;
597 }
598
599 void writeTask.catch(() => {});
600 throw error;
601 }
602
603 throw error;
604 }
605}
606
607// ── Replay ──────────────────────────────────────────────────────────────────
608
609export function createReplayReadableStream(options: {
610 storage: DurableObjectStorage;
611 cacheKey: string;
612 meta: StreamOutputMeta;
613}): ReadableStream<Uint8Array> {
614 const { storage, cacheKey, meta } = options;
615 if (meta.state !== StreamOutputState.Complete) {
616 throw new Error("Cannot replay streamed step output before it is complete");
617 }
618
619 // eslint-disable-next-line workers-sdk/no-unsafe-command-execution -- DO SQL exec, not child_process
620 const chunkCursor = storage.sql.exec<{
621 chunk_index: number;
622 chunk: ArrayBuffer;
623 }>(
624 `SELECT chunk_index, chunk FROM ${STREAMING_STEP_CHUNKS_TABLE} WHERE cache_key = ? AND attempt = ? ORDER BY chunk_index`,
625 cacheKey,
626 meta.attempt,
627 );
628 let index = 0;
629
630 return new ReadableStream<Uint8Array>({
631 pull(controller) {
632 if (index >= meta.chunkCount) {
633 controller.close();
634 return;
635 }
636
637 const row = chunkCursor.next();
638 if (row.done) {
639 controller.error(new Error(`Missing chunk ${index} for streamed step output`));
640 return;
641 }
642
643 if (row.value.chunk_index !== index) {
644 controller.error(new Error(`Missing chunk ${index} for streamed step output`));
645 return;
646 }
647
648 if (!(row.value.chunk instanceof ArrayBuffer)) {
649 controller.error(new Error("Invalid chunk type returned from streaming_step_chunks table"));
650 return;
651 }
652
653 controller.enqueue(new Uint8Array(row.value.chunk));
654 index++;
655 },
656 });
657}