deepline 0.1.200 → 0.1.202

package/dist/bundling-sources/shared_libs/play-runtime/runtime-api.ts

@@ -115,7 +115,11 @@ type RuntimeApiContext = {
   runId?: string | null;
   userEmail?: string | null;
   integrationMode?: 'live' | 'eval_stub' | 'fixture' | null;
-  dbSessionStrategy?: 'preloaded' | 'sandbox_public_key' | null;
+  dbSessionStrategy?:
+    | 'preloaded'
+    | 'sandbox_public_key'
+    | 'gateway_only'
+    | null;
   preloadedDbSessions?: PreloadedRuntimeDbSession[] | null;
   vercelProtectionBypassToken?: string | null;
   runtimeTestFaultHeader?: string | null;
@@ -443,9 +447,33 @@ type RuntimeApiActionRequest =
       action: 'repair_runtime_storage_grants';
       playName: string;
       runId?: string | null;
+    }
+  | {
+      action: 'runtime_sheet_start';
+      input: Parameters<typeof startRuntimeSheetDataset>[1];
+    }
+  | {
+      action: 'runtime_sheet_complete_map_rows';
+      input: Parameters<typeof completeRuntimeMapRows>[1];
+    }
+  | {
+      action: 'runtime_sheet_read_rows';
+      input: Parameters<typeof readRuntimeSheetDatasetRows>[1];
+    }
+  | {
+      action: 'runtime_sheet_read_row_keys';
+      input: Parameters<typeof readRuntimeSheetDatasetRowKeys>[1];
+    }
+  | {
+      action: 'runtime_sheet_release_attempt';
+      input: Parameters<typeof releaseRuntimeSheetAttempt>[1];
+    }
+  | {
+      action: 'runtime_receipts_release_attempt';
+      input: Parameters<typeof releaseRuntimeWorkReceipts>[1];
     };
 
-type RuntimeApiRowRecord = MapRowOutcome & {
+export type RuntimeApiRowRecord = MapRowOutcome & {
   inputIndex?: number | null;
 };
 
@@ -614,8 +642,8 @@ function runtimeApiMaxAttempts(
   return runtimeApiRetryDelaysForAction(action).length + 1;
 }
 
-function isTransientRuntimeApiErrorMessage(message: string): boolean {
-  return /timeout exceeded when trying to connect|timed out|fetch failed|ECONNRESET|ECONNREFUSED|UND_ERR_CONNECT_TIMEOUT|requested endpoint could not be found, or you don't have access/i.test(
+export function isTransientRuntimeApiErrorMessage(message: string): boolean {
+  return /timeout exceeded when trying to connect|timed out|fetch failed|ECONNRESET|ECONNREFUSED|UND_ERR_CONNECT_TIMEOUT|requested endpoint could not be found, or you don't have access|RUNTIME_SCHEDULER_SATURATED|Runtime scheduler DB pool is saturated|Runtime scheduler DB circuit breaker is open/i.test(
     message,
   );
 }
@@ -726,6 +754,75 @@ async function isRuntimeSheetSchemaReady(
   }
 }
 
+/**
+ * Report exactly what the scoped runtime role can see for the sheet table, from
+ * the same connection the readiness probe uses. When ensure_sheet reports the
+ * table as provisioned but the runner probe still returns not-ready, the cause
+ * is almost always that the runtime role lacks a privilege on the freshly
+ * created table (so information_schema.columns hides the attempt columns) or a
+ * schema/table-name mismatch. This turns the previously blind "still not ready"
+ * error into an actionable one that names the concrete missing piece.
+ */
+async function readRuntimeSheetSchemaDiagnostic(
+  session: RuntimePostgresSession,
+): Promise<Record<string, unknown>> {
+  const query = async (
+    client: RuntimeQueryClient,
+  ): Promise<Record<string, unknown>> => {
+    const visibleColumns = await client.query<{ column_name: string }>(
+      `
+        SELECT column_name
+        FROM information_schema.columns
+        WHERE table_schema = $1
+          AND table_name = $2
+        ORDER BY ordinal_position
+      `,
+      [session.postgres.schema, session.postgres.sheetTable],
+    );
+    const privileges = await client.query<Record<string, unknown>>(
+      `
+        SELECT
+          current_user AS probe_role,
+          to_regclass($3) IS NOT NULL AS regclass_visible,
+          has_table_privilege(current_user, $3, 'SELECT') AS can_select,
+          has_table_privilege(current_user, $3, 'INSERT') AS can_insert,
+          has_table_privilege(current_user, $3, 'UPDATE') AS can_update
+      `,
+      [
+        session.postgres.schema,
+        session.postgres.sheetTable,
+        `${quoteIdentifier(session.postgres.schema)}.${quoteIdentifier(session.postgres.sheetTable)}`,
+      ],
+    );
+    const visible = new Set(visibleColumns.rows.map((row) => row.column_name));
+    return {
+      schema: session.postgres.schema,
+      table: session.postgres.sheetTable,
+      probeRole: privileges.rows[0]?.probe_role ?? null,
+      regclassVisible: privileges.rows[0]?.regclass_visible ?? null,
+      canSelect: privileges.rows[0]?.can_select ?? null,
+      canInsert: privileges.rows[0]?.can_insert ?? null,
+      canUpdate: privileges.rows[0]?.can_update ?? null,
+      visibleColumnCount: visible.size,
+      missingAttemptColumns: RUNTIME_SHEET_ATTEMPT_COLUMNS.filter(
+        (column) => !visible.has(column),
+      ),
+    };
+  };
+  try {
+    if (isRuntimeOneShotQueryFactoryRegistered()) {
+      return await withRuntimeOneShotPostgres(session, query);
+    }
+    return await withRuntimePostgres(session, query);
+  } catch (error) {
+    return {
+      schema: session.postgres.schema,
+      table: session.postgres.sheetTable,
+      diagnosticError: error instanceof Error ? error.message : String(error),
+    };
+  }
+}
+
 async function waitForRuntimeSheetSchemaReadyAfterEnsure(
   session: RuntimePostgresSession,
   input: {
@@ -839,7 +936,7 @@ async function ensureRuntimeSheetForPreloadedSession(
       phase: 'ensure_sheet_for_preloaded_session',
       ms: Date.now() - ensureStartedAt,
     });
-    const readyAfterEnsure = await waitForRuntimeSheetSchemaReadyAfterEnsure(
+    let readyAfterEnsure = await waitForRuntimeSheetSchemaReadyAfterEnsure(
       input.session,
       {
         sheetContract: input.sheetContract,
@@ -847,8 +944,39 @@ async function ensureRuntimeSheetForPreloadedSession(
       },
     );
     if (!readyAfterEnsure) {
+      // ensure_sheet provisions the sheet as the tenant admin role, but the
+      // preloaded session probes as the scoped runtime role. When the tenant
+      // role contract was only partially applied (e.g. role-membership grants
+      // skipped on a data plane where the owner cannot administer roles), the
+      // runtime role never inherits access to the freshly created table, so its
+      // information_schema view hides the attempt columns and the probe stays
+      // not-ready even though the table exists. Re-apply the storage grants
+      // through the same self-heal the pre-ensure probe uses, then re-check
+      // once. This is a grant repair, not a blind retry: if the runtime role
+      // still cannot see the sheet, we fail loudly below with the concrete
+      // runtime-role view.
+      const repairStartedAt = Date.now();
+      await repairRuntimeStorageGrants(context, {
+        playName: context.playName,
+      });
+      input.timings?.push({
+        phase: 'repair_storage_grants_after_preloaded_session_ensure',
+        ms: Date.now() - repairStartedAt,
+        retried: true,
+      });
+      readyAfterEnsure = await waitForRuntimeSheetSchemaReadyAfterEnsure(
+        input.session,
+        {
+          sheetContract: input.sheetContract,
+          timings: input.timings,
+        },
+      );
+    }
+    if (!readyAfterEnsure) {
+      const diagnostic = await readRuntimeSheetSchemaDiagnostic(input.session);
       throw new Error(
-        `Runtime sheet schema for ctx.dataset("${input.tableNamespace}") is still not ready after ensure_sheet.`,
+        `Runtime sheet schema for ctx.dataset("${input.tableNamespace}") is still not ready after ensure_sheet. ` +
+          `Runtime-role view: ${JSON.stringify(diagnostic)}`,
       );
     }
   })();
@@ -4158,6 +4286,36 @@ export async function markRuntimeWorkReceiptsQueued(
   );
 }
 
+/**
+ * Replay-idempotency fence for a receipt whose active complete/fail UPDATE
+ * matched 0 rows. This fires when the durable write committed but the HTTP
+ * response was lost and the SAME caller retries: the row is already terminal,
+ * so the active-owner predicate (which requires status IN (queued, running))
+ * can never match again. We surface the already-terminal row as success ONLY
+ * when its run identity matches the retrying caller, so the retry does not
+ * spuriously fail a receipt whose provider spend and durable write both
+ * succeeded.
+ *
+ * The predicate mirrors the run-identity half of {@link
+ * workReceiptActiveOwnerPredicateSql}: `COALESCE(lease_owner_run_id, run_id)`
+ * must equal the caller's run id. complete/fail null the lease-owner and
+ * attempt columns as part of the terminal write, so run id is the only durable
+ * ownership marker left on a terminal row — a different run (the stale-owner
+ * rejection case) never matches and still resolves to null exactly as before.
+ * This is replay-idempotency only; it never lets a foreign run read or claim a
+ * receipt it did not complete.
+ */
+function workReceiptTerminalReplayPredicateSql(input: {
+  receiptTable: string;
+  terminalStatusSql: string;
+  ownerRunIdSql: string;
+}): string {
+  const table = input.receiptTable;
+  return `${table}.status = ${input.terminalStatusSql}::smallint
+            /* replay-idempotency: same-run retry after a lost response; run identity is the only durable owner marker on a terminal row */
+            AND COALESCE(${table}.lease_owner_run_id, ${table}.run_id) = ${input.ownerRunIdSql}`;
+}
+
 export async function completeRuntimeWorkReceipt(
   context: RuntimeApiContext,
   input: {
@@ -4240,8 +4398,31 @@ export async function completeRuntimeWorkReceipt(
 	                    lease_owner_attempt,
 	                    lease_expires_at,
 	                    updated_at
+	        ),
+	        replay AS (
+	          SELECT convert_from(k, 'UTF8') AS k,
+	                 status,
+	                 output,
+	                 error,
+	                 failure_kind,
+	                 run_id,
+	                 lease_id,
+	                 lease_owner_run_id,
+	                 lease_owner_attempt,
+	                 lease_expires_at,
+	                 updated_at
+	            FROM ${workReceiptTable(session)}
+	           WHERE k = decode($1, 'hex')
+	             AND NOT EXISTS (SELECT 1 FROM completed)
+	             AND ${workReceiptTerminalReplayPredicateSql({
+                 receiptTable: workReceiptTable(session),
+                 terminalStatusSql: '$2',
+                 ownerRunIdSql: '$4',
+               })}
 	        )
 	        SELECT k, status, output, error, failure_kind, run_id, lease_id, lease_owner_run_id, lease_owner_attempt, lease_expires_at, updated_at FROM completed
+	        UNION ALL
+	        SELECT k, status, output, error, failure_kind, run_id, lease_id, lease_owner_run_id, lease_owner_attempt, lease_expires_at, updated_at FROM replay
       `,
         [
           workReceiptKeyHex(input.key),
@@ -4378,21 +4559,49 @@ export async function completeRuntimeWorkReceipts(
                   target.updated_at,
                   inputs.ord
       ),
+      replay AS (
+        SELECT target.k,
+               target.status,
+               target.output,
+               target.error,
+               target.failure_kind,
+               target.run_id,
+               target.lease_id,
+               target.lease_owner_run_id,
+               target.lease_owner_attempt,
+               target.lease_expires_at,
+               target.updated_at,
+               inputs.ord
+          FROM inputs
+          JOIN ${workReceiptTable(session)} AS target
+            ON target.k = decode(inputs.key_hex, 'hex')
+         WHERE NOT EXISTS (SELECT 1 FROM completed WHERE completed.ord = inputs.ord)
+           AND ${workReceiptTerminalReplayPredicateSql({
+             receiptTable: 'target',
+             terminalStatusSql: '$6',
+             ownerRunIdSql: 'inputs.run_id',
+           })}
+      ),
+      resolved AS (
+        SELECT ord, k, status, output, error, failure_kind, run_id, lease_id, lease_owner_run_id, lease_owner_attempt, lease_expires_at, updated_at FROM completed
+        UNION ALL
+        SELECT ord, k, status, output, error, failure_kind, run_id, lease_id, lease_owner_run_id, lease_owner_attempt, lease_expires_at, updated_at FROM replay
+      ),
       returned AS (
-        SELECT completed.k,
-               completed.status,
-	               completed.output,
-	               completed.error,
-	               completed.failure_kind,
-	               completed.run_id,
-               completed.lease_id,
-               completed.lease_owner_run_id,
-               completed.lease_owner_attempt,
-               completed.lease_expires_at,
-               completed.updated_at,
+        SELECT resolved.k,
+               resolved.status,
+	               resolved.output,
+	               resolved.error,
+	               resolved.failure_kind,
+	               resolved.run_id,
+               resolved.lease_id,
+               resolved.lease_owner_run_id,
+               resolved.lease_owner_attempt,
+               resolved.lease_expires_at,
+               resolved.updated_at,
                inputs.ord
           FROM inputs
-          LEFT JOIN completed ON completed.ord = inputs.ord
+          LEFT JOIN resolved ON resolved.ord = inputs.ord
       )
       SELECT convert_from(returned.k, 'UTF8') AS k,
              returned.status,
@@ -4484,8 +4693,31 @@ export async function failRuntimeWorkReceipt(
 	                    lease_owner_attempt,
 	                    lease_expires_at,
 	                    updated_at
+	        ),
+	        replay AS (
+	          SELECT convert_from(k, 'UTF8') AS k,
+	                 status,
+	                 output,
+	                 error,
+	                 failure_kind,
+	                 run_id,
+	                 lease_id,
+	                 lease_owner_run_id,
+	                 lease_owner_attempt,
+	                 lease_expires_at,
+	                 updated_at
+	            FROM ${workReceiptTable(session)}
+	           WHERE k = decode($1, 'hex')
+	             AND NOT EXISTS (SELECT 1 FROM failed)
+	             AND ${workReceiptTerminalReplayPredicateSql({
+                 receiptTable: workReceiptTable(session),
+                 terminalStatusSql: '$2',
+                 ownerRunIdSql: '$4',
+               })}
 	        )
 	        SELECT k, status, output, error, failure_kind, run_id, lease_id, lease_owner_run_id, lease_owner_attempt, lease_expires_at, updated_at FROM failed
+	        UNION ALL
+	        SELECT k, status, output, error, failure_kind, run_id, lease_id, lease_owner_run_id, lease_owner_attempt, lease_expires_at, updated_at FROM replay
       `,
         [
           workReceiptKeyHex(input.key),
@@ -4649,21 +4881,49 @@ export async function failRuntimeWorkReceipts(
                   target.updated_at,
                   inputs.ord
       ),
+      replay AS (
+        SELECT target.k,
+               target.status,
+               target.output,
+               target.error,
+               target.failure_kind,
+               target.run_id,
+               target.lease_id,
+               target.lease_owner_run_id,
+               target.lease_owner_attempt,
+               target.lease_expires_at,
+               target.updated_at,
+               inputs.ord
+          FROM inputs
+          JOIN ${workReceiptTable(session)} AS target
+            ON target.k = decode(inputs.key_hex, 'hex')
+         WHERE NOT EXISTS (SELECT 1 FROM failed WHERE failed.ord = inputs.ord)
+           AND ${workReceiptTerminalReplayPredicateSql({
+             receiptTable: 'target',
+             terminalStatusSql: '$7',
+             ownerRunIdSql: 'inputs.run_id',
+           })}
+      ),
+      resolved AS (
+        SELECT ord, k, status, output, error, failure_kind, run_id, lease_id, lease_owner_run_id, lease_owner_attempt, lease_expires_at, updated_at FROM failed
+        UNION ALL
+        SELECT ord, k, status, output, error, failure_kind, run_id, lease_id, lease_owner_run_id, lease_owner_attempt, lease_expires_at, updated_at FROM replay
+      ),
       returned AS (
-        SELECT failed.k,
-               failed.status,
-	               failed.output,
-	               failed.error,
-	               failed.failure_kind,
-	               failed.run_id,
-               failed.lease_id,
-               failed.lease_owner_run_id,
-               failed.lease_owner_attempt,
-               failed.lease_expires_at,
-               failed.updated_at,
+        SELECT resolved.k,
+               resolved.status,
+	               resolved.output,
+	               resolved.error,
+	               resolved.failure_kind,
+	               resolved.run_id,
+               resolved.lease_id,
+               resolved.lease_owner_run_id,
+               resolved.lease_owner_attempt,
+               resolved.lease_expires_at,
+               resolved.updated_at,
                inputs.ord
           FROM inputs
-          LEFT JOIN failed ON failed.ord = inputs.ord
+          LEFT JOIN resolved ON resolved.ord = inputs.ord
       )
       SELECT convert_from(returned.k, 'UTF8') AS k,
              returned.status,
@@ -4892,6 +5152,12 @@ export async function startRuntimeSheetDataset(
     force?: boolean;
   },
 ): Promise<PrepareRuntimeSheetResult> {
+  if (context.dbSessionStrategy === 'gateway_only') {
+    return await postRuntimeApi<PrepareRuntimeSheetResult>(context, {
+      action: 'runtime_sheet_start',
+      input,
+    });
+  }
   const totalStartedAt = Date.now();
   const timings: RuntimeSheetTiming[] = [];
   const playName = context.playName?.trim() || input.playName;
@@ -5283,6 +5549,12 @@ export async function releaseRuntimeSheetAttempt(
     attemptSeq?: number | null;
   },
 ): Promise<RuntimeSheetAttemptReleaseResult> {
+  if (context.dbSessionStrategy === 'gateway_only') {
+    return await postRuntimeApi<RuntimeSheetAttemptReleaseResult>(context, {
+      action: 'runtime_sheet_release_attempt',
+      input,
+    });
+  }
   const playName = context.playName?.trim() || input.playName;
   if (!playName) {
     throw new Error('Runtime sheet attempt release requires a playName.');
@@ -5371,6 +5643,12 @@ export async function releaseRuntimeWorkReceipts(
     runAttempt?: number | null;
   },
 ): Promise<RuntimeWorkReceiptReleaseResult> {
+  if (context.dbSessionStrategy === 'gateway_only') {
+    return await postRuntimeApi<RuntimeWorkReceiptReleaseResult>(context, {
+      action: 'runtime_receipts_release_attempt',
+      input,
+    });
+  }
   const runId = input.runId.trim();
   if (!runId) {
     return { released: 0, releasedKeys: [] };
@@ -5423,7 +5701,7 @@ type CompleteRuntimeMapRowChunksInput = {
   outputFields: string[];
 };
 
-type RuntimeMapRowsWriteResult = {
+export type RuntimeMapRowsWriteResult = {
   updated: number;
   fencedKeys: string[];
 };
@@ -6344,6 +6622,12 @@ export async function completeRuntimeMapRows(
     forceFailedRows?: boolean;
   },
 ): Promise<RuntimeMapRowsWriteResult> {
+  if (context.dbSessionStrategy === 'gateway_only') {
+    return await postRuntimeApi<RuntimeMapRowsWriteResult>(context, {
+      action: 'runtime_sheet_complete_map_rows',
+      input,
+    });
+  }
   if (input.rows.length === 0) {
     return { updated: 0, fencedKeys: [] };
   }
@@ -6498,6 +6782,12 @@ export async function readRuntimeSheetDatasetRows(
     offset: number;
   },
 ): Promise<{ rows: Record<string, unknown>[]; limit: number; offset: number }> {
+  if (context.dbSessionStrategy === 'gateway_only') {
+    return await postRuntimeApi(context, {
+      action: 'runtime_sheet_read_rows',
+      input,
+    });
+  }
   const limit = Math.max(
     1,
     Math.min(DIRECT_POSTGRES_BATCH_SIZE, Math.floor(input.limit)),
@@ -6540,6 +6830,12 @@ export async function readRuntimeSheetDatasetRowKeys(
     rowMode?: 'output' | 'all' | 'terminalAnyRun';
   },
 ): Promise<{ keys: string[] }> {
+  if (context.dbSessionStrategy === 'gateway_only') {
+    return await postRuntimeApi(context, {
+      action: 'runtime_sheet_read_row_keys',
+      input,
+    });
+  }
   const keys = [
     ...new Set(
       input.keys.map((key) => key.trim()).filter((key) => key.length > 0),

package/dist/bundling-sources/shared_libs/play-runtime/runtime-scheduler-topology.ts

@@ -0,0 +1,26 @@
+/**
+ * Runtime-scheduler topology predicate shared by the app-side lane helpers
+ * (`src/lib/plays/scheduler-backends/absurd-shared.ts`) and the runner
+ * backends (`runner-backends/backends/daytona-lifecycle.ts`).
+ *
+ * ONE distinction drives several guards: a run whose scheduler schema is the
+ * prod default rides the production topology (base queue names, customer code,
+ * production egress policy), while any other schema is an ISOLATED topology —
+ * a per-PR preview fleet or a per-worktree dev fleet running synthetic
+ * internal test plays. Keeping the predicate here (dependency-free) lets
+ * shared_libs code use it without importing the app-side scheduler stack.
+ */
+export const PROD_DEFAULT_RUNTIME_SCHEDULER_SCHEMA = 'runtime_scheduler';
+
+/**
+ * True when the scheduler schema names an isolated (preview CI / worktree)
+ * topology. Null/empty/prod-default => false (production topology). Callers
+ * that gate SECURITY policy on this must treat `false` as the fail-closed
+ * production posture — an absent schema is production, never "unknown".
+ */
+export function isIsolatedRuntimeSchedulerSchema(
+  schedulerSchema: string | null | undefined,
+): boolean {
+  const schema = schedulerSchema?.trim();
+  return Boolean(schema) && schema !== PROD_DEFAULT_RUNTIME_SCHEDULER_SCHEMA;
+}

package/dist/bundling-sources/shared_libs/play-runtime/scheduler-backend.ts

@@ -18,6 +18,7 @@ import type {
 import type { ExecutionPlan } from './execution-plan';
 import type { PlayRuntimeManifestMap } from '../plays/compiler-manifest';
 import type { PreloadedRuntimeDbSession } from './db-session';
+import type { RuntimeAuthorityDescriptor } from './execution-capabilities';
 import type { PlayRunnerRuntimeTiming } from './protocol';
 import type { RuntimeTestPolicyOverrides } from './test-runtime-seams';
 import type { PlayRunInputPayload } from './play-input';
@@ -123,9 +124,21 @@ export type PlaySchedulerSubmitInput = {
   childPlayManifests?: PlayRuntimeManifestMap | null;
   playCallGovernance?: PlayCallGovernanceSnapshot | null;
   preloadedDbSessions?: PreloadedRuntimeDbSession[] | null;
+  /**
+   * Secret that decrypts `preloadedDbSessions` postgres URLs. Equals the run's
+   * original submit-time executor (authority) token that AES-GCM sealed the
+   * sessions. Held separately from `executorToken` because durable schedulers
+   * re-mint a per-attempt `executorToken` (adding `run_attempt`/capabilities
+   * claims) for API auth; that re-minted string is NOT valid decryption key
+   * material for sessions sealed under the original token. The worker/runner use
+   * this field for DB-session unwrap and the re-minted `executorToken` for auth.
+   */
+  preloadedDbSessionUnwrapKey?: string | null;
   /** Optional immutable Worker module source for local Dynamic Worker loading. */
   dynamicWorkerCode?: string | null;
   executorToken: string;
+  /** Immutable facts used by durable schedulers to mint a fresh leg token. */
+  runtimeAuthority: RuntimeAuthorityDescriptor;
   /**
    * Public app origin used by Workers to call Deepline runtime API routes.
    * When omitted, legacy schedulers fall back to baseUrl.
@@ -151,6 +164,14 @@ export type PlaySchedulerSubmitInput = {
   coordinatorInternalToken?: string | null;
   /** Runtime deploy generation/version that owns this run, when known. */
   runtimeDeployVersion?: string | null;
+  /**
+   * Absurd release lane that owns this run for its whole life (children, parks,
+   * wakes, retries). Stamped at launch by the absurd scheduler backend and
+   * persisted inside launch_json. The worker derives the run's queue from this,
+   * never from process-local config. Absent on pre-release launch rows, in which
+   * case the worker falls back to the dev-collapse id (today's queue).
+   */
+  absurdReleaseId?: string | null;
   /** Request-scoped Vercel Deployment Protection bypass for preview runtime callbacks. */
   vercelProtectionBypassToken?: string | null;
   /** Request-scoped, dev-only runtime fault injection header for black-box durability tests. */

package/dist/bundling-sources/shared_libs/plays/dataset.ts

@@ -4,6 +4,7 @@ const PLAY_DATASET_BRAND = Symbol.for('deepline.play.dataset');
 const NODE_INSPECT_CUSTOM = Symbol.for('nodejs.util.inspect.custom');
 const DEFAULT_MATERIALIZE_LIMIT = 10_000;
 export const PLAY_DATASET_EXECUTION_PAGE_ROWS = 1_000;
+export const PLAY_DATASET_EXECUTION_PAGE_BYTES = 64 * 1024 * 1024;
 
 export type PlayDatasetKind = 'csv' | 'map';
 
@@ -545,6 +546,8 @@ export type MaterializePlayDatasetInputOptions<T> = {
 
 export type IteratePlayDatasetInputPagesOptions<T> = {
   pageSize?: number;
+  maxPageBytes?: number;
+  estimateRowBytes?: (row: T) => number;
   onRow?: (row: T, index: number) => void;
 };
 
@@ -560,21 +563,40 @@ export async function* iteratePlayDatasetInputPages<T>(
   options?: IteratePlayDatasetInputPagesOptions<T>,
 ): AsyncIterable<{ rows: T[]; offset: number }> {
   const pageSize = normalizeDatasetExecutionPageSize(options?.pageSize);
+  const maxPageBytes = Math.max(
+    1,
+    Math.floor(options?.maxPageBytes ?? PLAY_DATASET_EXECUTION_PAGE_BYTES),
+  );
+  const estimateRowBytes =
+    options?.estimateRowBytes ??
+    ((row: T) => {
+      const serialized = JSON.stringify(row);
+      return serialized === undefined ? 0 : serialized.length;
+    });
   let page: T[] = [];
+  let pageBytes = 0;
   let offset = 0;
   let index = 0;
-  const pushRow = (row: T): { rows: T[]; offset: number } | null => {
+  const pushRow = (row: T): Array<{ rows: T[]; offset: number }> => {
     options?.onRow?.(row, index);
     index += 1;
+    const rowBytes = Math.max(0, Math.ceil(estimateRowBytes(row)));
+    const ready: Array<{ rows: T[]; offset: number }> = [];
+    if (page.length > 0 && pageBytes + rowBytes > maxPageBytes) {
+      ready.push({ rows: page, offset });
+      offset += page.length;
+      page = [];
+      pageBytes = 0;
+    }
     page.push(row);
-    if (page.length < pageSize) {
-      return null;
+    pageBytes += rowBytes;
+    if (page.length >= pageSize || pageBytes >= maxPageBytes) {
+      ready.push({ rows: page, offset });
+      offset += page.length;
+      page = [];
+      pageBytes = 0;
     }
-    const fullPage = page;
-    const fullOffset = offset;
-    page = [];
-    offset += fullPage.length;
-    return { rows: fullPage, offset: fullOffset };
+    return ready;
   };
 
   const flush = (): { rows: T[]; offset: number } | null => {
@@ -582,13 +604,13 @@ export async function* iteratePlayDatasetInputPages<T>(
     const finalPage = page;
     const finalOffset = offset;
     page = [];
+    pageBytes = 0;
     offset += finalPage.length;
     return { rows: finalPage, offset: finalOffset };
   };
 
   const emitRow = async function* (row: T) {
-    const ready = pushRow(row);
-    if (ready) yield ready;
+    for (const ready of pushRow(row)) yield ready;
   };
 
   if (isPlayDataset(input)) {