@tanstack/ai-sandbox-cloudflare 0.3.14 → 0.4.0

package/src/coordinator.ts

@@ -31,21 +31,30 @@ import { EventType, isTerminalRunStatus } from '@tanstack/ai'
 // migrated on read; see './run-log').
 import { RunController } from '@tanstack/ai-sandbox'
 import { runLogStore, runLogStream } from './durability'
+import { hasInFlightCallback } from './coordinator-callbacks'
 import { DurableObjectRunEventLog } from './run-log-do'
 import type { ModelMessage, StreamChunk } from '@tanstack/ai'
 import type { RunLogRecord } from './run-log'
 
-/** Re-arm window for the liveness watchdog while a run is in flight (ms). */
+/** Upper bound on the watchdog check interval while a run is in flight (ms). */
 const WATCHDOG_MS = 30_000
 
-/**
- * How long a non-terminal run may go without ANY new event before the watchdog
- * presumes the orchestrator driving it is dead (eviction that lost the
- * `waitUntil` promise, an uncaught fault, a hung container) and fails the run so
- * tailing clients stop waiting forever. Generous so a legitimately slow agent
- * step (a long tool call that emits no chunks) is not killed prematurely.
- */
-const WATCHDOG_STALL_MS = 5 * 60_000
+/** Default permitted period without persisted run activity (ms). */
+const DEFAULT_STALL_TIMEOUT_MS = 5 * 60_000
+
+/** @internal Shared validation for direct subclasses and the eager factory path. */
+export function normalizeStallTimeoutMs(
+  stallTimeoutMs: number | false | undefined,
+): number | false {
+  if (stallTimeoutMs === undefined) return DEFAULT_STALL_TIMEOUT_MS
+  if (stallTimeoutMs === false) return false
+  if (!Number.isSafeInteger(stallTimeoutMs) || stallTimeoutMs <= 0) {
+    throw new TypeError(
+      'stallTimeoutMs must be a positive safe integer or false',
+    )
+  }
+  return stallTimeoutMs
+}
 
 /** What the Worker hands the coordinator to start a run. */
 export interface StartRunInput {
@@ -106,9 +115,15 @@ export abstract class SandboxCoordinator<
    * running — double-delivering events and racing the persisted cursor.
    */
   private readonly pumping = new WeakSet<WebSocket>()
+  private readonly stallTimeoutMs: number | false
 
-  constructor(ctx: DurableObjectState, env: TEnv) {
+  constructor(
+    ctx: DurableObjectState,
+    env: TEnv,
+    stallTimeoutMs?: number | false,
+  ) {
     super(ctx, env)
+    this.stallTimeoutMs = normalizeStallTimeoutMs(stallTimeoutMs)
     this.log = new DurableObjectRunEventLog(ctx.storage)
     this.controller = new RunController({
       runs: runLogStore(this.log),
@@ -154,7 +169,12 @@ export abstract class SandboxCoordinator<
 
   async startRun(input: StartRunInput): Promise<{ runId: string }> {
     const existing = await this.log.get(input.runId)
-    if (existing) return { runId: input.runId } // idempotent re-trigger
+    if (existing) {
+      if (!isTerminalRunStatus(existing.status)) await this.armWatchdog()
+      return { runId: input.runId }
+    }
+
+    await this.armWatchdog()
 
     // Open the run BEFORE building the stream. `pipeToRunLog`'s never-rejects
     // guarantee only covers failures AFTER the stream is handed to it — a throw
@@ -189,7 +209,6 @@ export abstract class SandboxCoordinator<
     // running the settle hook.
     const settle = (): void => this.onRunSettled(input.runId)
     this.ctx.waitUntil(done.then(settle, settle))
-    await this.ctx.storage.setAlarm(Date.now() + WATCHDOG_MS)
     return { runId: input.runId }
   }
 
@@ -320,41 +339,58 @@ export abstract class SandboxCoordinator<
   // ===========================================================================
 
   override async alarm(): Promise<void> {
+    // A previously configured alarm may still be delivered once after watchdogs
+    // are disabled. It must self-extinguish before reading storage or entering a
+    // catch path that could re-arm it. Deliberately do not call deleteAlarm().
+    if (this.stallTimeoutMs === false) return
+
     try {
       // Through the log (not a raw `rec:` list) so legacy records are migrated
       // on the way out — the storage layout is the log's private concern.
       const runs = await this.log.list()
-      const now = Date.now()
+      const cutoff = Date.now() - this.stallTimeoutMs
       let active = false
       for (const record of runs) {
         if (isTerminalRunStatus(record.status)) continue
-        if (now - record.updatedAt > WATCHDOG_STALL_MS) {
-          // No progress for too long — the driver is presumed dead. Fail the run
-          // so tailing clients stop waiting forever (the whole point of the
-          // watchdog; without this a stuck run sits at `running` indefinitely).
-          await this.failStalledRun(record.runId)
-        } else {
+        if (
+          hasInFlightCallback(this, record.runId) ||
+          record.updatedAt >= cutoff
+        ) {
           active = true
+          continue
         }
+
+        // Re-check and terminalize atomically against the same strict cutoff.
+        // A callback touch or normal completion racing this alarm wins cleanly.
+        const finished = await this.failStalledRun(record.runId, cutoff)
+        if (finished) this.onRunSettled(record.runId)
+        else active = true
       }
-      if (active) await this.ctx.storage.setAlarm(Date.now() + WATCHDOG_MS)
+      if (active) await this.armWatchdog()
     } catch (error) {
       // Never let the watchdog die silently: a transient storage error must not
       // permanently disable liveness detection. Re-arm and try again next tick.
       console.error('[sandbox-coordinator] watchdog alarm failed:', error)
-      await this.ctx.storage.setAlarm(Date.now() + WATCHDOG_MS)
+      await this.armWatchdog()
     }
   }
 
-  /** Mark a stalled (orchestrator-presumed-dead) run as a terminal error. */
-  private async failStalledRun(runId: string): Promise<void> {
-    const message = 'run watchdog: no progress; orchestrator presumed dead'
-    try {
-      await this.log.append(runId, { type: EventType.RUN_ERROR, message })
-    } catch {
-      // The run may have just reached terminal concurrently; finish is idempotent.
+  /** Arm without allowing a new run to postpone an earlier pending check. */
+  private async armWatchdog(): Promise<void> {
+    if (this.stallTimeoutMs === false) return
+    const next = Date.now() + Math.min(WATCHDOG_MS, this.stallTimeoutMs)
+    const pending = await this.ctx.storage.getAlarm()
+    if (pending === null || pending > next) {
+      await this.ctx.storage.setAlarm(next)
     }
-    await this.log.finish(runId, 'failed', { message })
-    this.onRunSettled(runId)
+  }
+
+  /** Mark a strictly stale run as a terminal error if it still qualifies. */
+  private failStalledRun(runId: string, cutoff: number): Promise<boolean> {
+    const message = 'run watchdog: no progress; orchestrator presumed dead'
+    return this.log.finishIfStale(runId, cutoff, {
+      type: EventType.RUN_ERROR,
+      message,
+    })
   }
 }

package/src/factory.ts

@@ -45,7 +45,7 @@ import { cloudflareSandbox } from './provider'
 import { ChatSandboxCoordinator } from './chat-coordinator'
 import { ContainerSandboxCoordinator } from './container-coordinator'
 import { createSandboxAgentWorker } from './worker'
-import { resolvePreviewHost } from './coordinator'
+import { normalizeStallTimeoutMs, resolvePreviewHost } from './coordinator'
 import type { ChatCoordinatorEnv, ChatRunConfig } from './chat-coordinator'
 import type {
   ContainerCoordinatorEnv,
@@ -83,6 +83,11 @@ export interface SandboxAgentEnv
 interface BaseAgentConfig<TEnv extends SandboxAgentEnv> {
   /** chat()-provided server tools, resolved per run (DO-drives: bridged over MCP). */
   tools?: (input: StartRunInput, env: TEnv) => Array<AnyTool>
+  /**
+   * Fail a run after this many milliseconds without persisted activity. Omitted
+   * defaults to five minutes; `false` disables the watchdog.
+   */
+  stallTimeoutMs?: number | false
 }
 
 /** DO-drives config: the DO runs `chat()` with the given adapter. */
@@ -188,11 +193,16 @@ function resolveCoordinator<TEnv extends SandboxAgentEnv>(
 export function createCloudflareSandboxAgent<
   TEnv extends SandboxAgentEnv = SandboxAgentEnv,
 >(config: CloudflareSandboxAgentConfig<TEnv>): CloudflareSandboxAgent<TEnv> {
+  const stallTimeoutMs = normalizeStallTimeoutMs(config.stallTimeoutMs)
   const worker = createSandboxAgentWorker<TEnv>(resolveCoordinator)
 
   if (config.mode === 'colocated') {
     const colocated = config
     class ConfiguredContainerCoordinator extends ContainerSandboxCoordinator<TEnv> {
+      constructor(ctx: DurableObjectState, env: TEnv) {
+        super(ctx, env, stallTimeoutMs)
+      }
+
       protected override config(input: StartRunInput): ContainerRunConfig {
         return {
           hostTools: colocated.tools?.(input, this.env) ?? [],
@@ -207,6 +217,10 @@ export function createCloudflareSandboxAgent<
 
   const doDrives = config
   class ConfiguredChatCoordinator extends ChatSandboxCoordinator<TEnv> {
+    constructor(ctx: DurableObjectState, env: TEnv) {
+      super(ctx, env, stallTimeoutMs)
+    }
+
     protected override config(input: StartRunInput): ChatRunConfig {
       const tools = doDrives.tools?.(input, this.env)
       return {

package/src/run-log-do.ts

@@ -149,15 +149,79 @@ export class DurableObjectRunEventLog implements RunEventLog {
     this.wake(runId)
   }
 
+  async touch(runId: string): Promise<void> {
+    await this.storage.transaction(async (txn) => {
+      const stored = await txn.get<StoredRunRecord>(recKey(runId))
+      if (!stored) return
+      const { record, migrated } = migrateStoredRunRecord(stored)
+      if (isTerminalRunStatus(record.status)) {
+        if (migrated) await txn.put(recKey(runId), record)
+        return
+      }
+      await txn.put(recKey(runId), { ...record, updatedAt: Date.now() })
+    })
+    // Activity without a new event or status transition gives a tailing reader
+    // nothing to observe, so intentionally do not wake it.
+  }
+
+  async finishIfStale(
+    runId: string,
+    cutoff: number,
+    chunk: Extract<StreamChunk, { type: 'RUN_ERROR' }>,
+  ): Promise<boolean> {
+    const finished = await this.storage.transaction(async (txn) => {
+      const stored = await txn.get<StoredRunRecord>(recKey(runId))
+      if (!stored) return false
+      const { record, migrated } = migrateStoredRunRecord(stored)
+      if (isTerminalRunStatus(record.status) || record.updatedAt >= cutoff) {
+        if (migrated) await txn.put(recKey(runId), record)
+        return false
+      }
+
+      const now = Date.now()
+      const seq = record.lastSeq + 1
+      const next: RunLogRecord = {
+        ...record,
+        status: 'failed',
+        lastSeq: seq,
+        error: {
+          message: chunk.message,
+          ...(chunk.code !== undefined ? { code: chunk.code } : {}),
+        },
+        finishedAt: now,
+        updatedAt: now,
+      }
+      // Read the current activity clock and commit the terminal event + record
+      // in this one transaction. Keep wake-ups outside: transaction callbacks
+      // may be retried and must contain no external work.
+      await txn.put(evtKey(runId, seq), chunk)
+      await txn.put(recKey(runId), next)
+      return true
+    })
+    if (finished) this.wake(runId)
+    return finished
+  }
+
   async update(runId: string, patch: RunRecordPatch): Promise<void> {
-    const record = await this.getRecord(runId)
-    if (!record) return // unknown runId is a no-op
-    const next: RunLogRecord = { ...record, ...patch, updatedAt: Date.now() }
-    await this.storage.put(recKey(runId), next)
+    const updated = await this.storage.transaction(async (txn) => {
+      const stored = await txn.get<StoredRunRecord>(recKey(runId))
+      if (!stored) return false
+      const { record, migrated } = migrateStoredRunRecord(stored)
+      if (isTerminalRunStatus(record.status)) {
+        if (migrated) await txn.put(recKey(runId), record)
+        return false
+      }
+      await txn.put(recKey(runId), {
+        ...record,
+        ...patch,
+        updatedAt: Date.now(),
+      })
+      return true
+    })
     // A patch may terminalize the shared status field (core's driver writes its
-    // terminal status through `RunStore.update`) — parked readers must see it
-    // now, not a TAIL_POLL_MS later.
-    this.wake(runId)
+    // terminal status through `RunStore.update`) — wake only after that update
+    // commits. A late update racing a terminal writer is an atomic no-op.
+    if (updated) this.wake(runId)
   }
 
   async get(runId: string): Promise<RunLogRecord | null> {

package/src/run-log.ts

@@ -97,7 +97,8 @@ export interface RunEventLogReadOptions {
  * - `append` assigns the next `seq` (0, 1, 2, …) and returns it.
  * - `read` yields the backlog after `fromSeq` in order, then live-tails new
  *   events, and RETURNS once the run is terminal and the cursor has caught up.
- * - All methods reject for an unknown `runId` except `get`, which resolves null.
+ * - Unknown-run behavior is method-specific: `append`/`finish`/`read` reject,
+ *   `get` resolves null, and `update` is a no-op.
  */
 export interface RunEventLog {
   /**
@@ -120,9 +121,9 @@ export interface RunEventLog {
     error?: RunError,
   ) => Promise<void>
   /**
-   * Patch the record's mutable fields ({@link RunRecordPatch}). Unknown `runId`
-   * is a NO-OP (never a throw, never a create) — core's `RunStore.update`
-   * invariant, which `runLogStore` maps onto this method.
+   * Patch the record's mutable fields ({@link RunRecordPatch}). Unknown and
+   * already-terminal runs are NO-OPs (never a throw, never a create), preserving
+   * the first terminal status against late driver updates.
    *
    * MUST wake blocked readers, exactly like `append`/`finish`: the record and
    * the event log share one status field here, so a driver that terminalizes
@@ -299,9 +300,49 @@ export class InMemoryRunEventLog implements RunEventLog {
     return Promise.resolve()
   }
 
+  touch(runId: string): Promise<void> {
+    const state = this.runs.get(runId)
+    if (!state || isTerminalRunStatus(state.record.status)) {
+      return Promise.resolve()
+    }
+    state.record.updatedAt = this.now()
+    return Promise.resolve()
+  }
+
+  finishIfStale(
+    runId: string,
+    cutoff: number,
+    chunk: Extract<StreamChunk, { type: 'RUN_ERROR' }>,
+  ): Promise<boolean> {
+    const state = this.runs.get(runId)
+    if (
+      !state ||
+      isTerminalRunStatus(state.record.status) ||
+      state.record.updatedAt >= cutoff
+    ) {
+      return Promise.resolve(false)
+    }
+
+    const now = this.now()
+    const seq = state.record.lastSeq + 1
+    state.chunks.push(chunk)
+    state.record.lastSeq = seq
+    state.record.status = 'failed'
+    state.record.error = {
+      message: chunk.message,
+      ...(chunk.code !== undefined ? { code: chunk.code } : {}),
+    }
+    state.record.finishedAt = now
+    state.record.updatedAt = now
+    this.wake(state)
+    return Promise.resolve(true)
+  }
+
   update(runId: string, patch: RunRecordPatch): Promise<void> {
     const state = this.runs.get(runId)
-    if (!state) return Promise.resolve() // unknown runId is a no-op
+    if (!state || isTerminalRunStatus(state.record.status)) {
+      return Promise.resolve()
+    }
     state.record = { ...state.record, ...patch, updatedAt: this.now() }
     // A patch may terminalize the shared status field (core's driver writes its
     // terminal status through `RunStore.update`) — parked readers must see it.