From b728453eb6fddaf495b10b807dcd69b9df312ac1 Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Thu, 7 May 2026 21:52:39 +0100 Subject: [PATCH] feat(sdk): chat.agent oomMachine for automatic OOM-retry on a larger machine Setting oomMachine opts a chat.agent into one-shot OOM recovery: the failed turn re-runs on the larger machine, derives a session.in cutoff from the latest trigger:turn-complete chunk on session.out, and skips the turns that already completed on the prior attempt. Adds the setMinTimestamp filter on the session-stream manager and force-kills the dev worker between attempts so local behavior matches prod. --- packages/cli-v3/src/dev/taskRunProcessPool.ts | 19 ++- .../src/entryPoints/dev-run-controller.ts | 37 ++++- packages/core/src/v3/sessionStreams/index.ts | 8 + .../src/v3/sessionStreams/manager.test.ts | 151 ++++++++++++++++++ .../core/src/v3/sessionStreams/manager.ts | 62 +++++-- .../core/src/v3/sessionStreams/noopManager.ts | 6 + packages/core/src/v3/sessionStreams/types.ts | 14 ++ .../v3/test/test-session-stream-manager.ts | 9 ++ packages/trigger-sdk/src/v3/ai.ts | 139 +++++++++++++++- 9 files changed, 425 insertions(+), 20 deletions(-) create mode 100644 packages/core/src/v3/sessionStreams/manager.test.ts diff --git a/packages/cli-v3/src/dev/taskRunProcessPool.ts b/packages/cli-v3/src/dev/taskRunProcessPool.ts index 810be7acb..414626e64 100644 --- a/packages/cli-v3/src/dev/taskRunProcessPool.ts +++ b/packages/cli-v3/src/dev/taskRunProcessPool.ts @@ -127,7 +127,11 @@ export class TaskRunProcessPool { return { taskRunProcess: newProcess, isReused: false }; } - async returnProcess(process: TaskRunProcess, version: string): Promise { + async returnProcess( + process: TaskRunProcess, + version: string, + options?: { forceKill?: boolean } + ): Promise { // Remove from busy processes for this version const busyProcesses = this.busyProcessesByVersion.get(version); if (busyProcesses) { @@ -141,6 +145,19 @@ export class TaskRunProcessPool { ); } + // `forceKill` skips the reuse heuristic and tears the process down. Used + // on outcomes that leave the process in a state we can't safely reuse + // (OOM in particular — production would get a fresh container, so local + // dev should match that). + if (options?.forceKill) { + logger.debug("[TaskRunProcessPool] Force-killing process", { + version, + pid: process.pid, + }); + await this.killProcess(process); + return; + } + if (this.shouldReuseProcess(process, version)) { const availableCount = this.availableProcessesByVersion.get(version)?.length || 0; const busyCount = this.busyProcessesByVersion.get(version)?.size || 0; diff --git a/packages/cli-v3/src/entryPoints/dev-run-controller.ts b/packages/cli-v3/src/entryPoints/dev-run-controller.ts index fe4866774..ce73da66f 100644 --- a/packages/cli-v3/src/entryPoints/dev-run-controller.ts +++ b/packages/cli-v3/src/entryPoints/dev-run-controller.ts @@ -2,6 +2,8 @@ import { CompleteRunAttemptResult, DequeuedMessage, IntervalService, + isManualOutOfMemoryError, + isOOMRunError, LogLevel, RunExecutionData, SuspendedProcessError, @@ -52,6 +54,12 @@ export class DevRunController { private readonly cwd?: string; private isCompletingRun = false; private isShuttingDown = false; + // Set when the current attempt's outcome means the worker process can't + // safely be reused (OOM in particular). Production gives every retry a + // fresh container; local dev's process pool needs the same on these + // outcomes or in-process state (e.g. session.in cursors) leaks across + // attempts and the OOM retry skips the message that triggered it. + private discardProcessOnReturn = false; private state: | { @@ -539,6 +547,13 @@ export class DevRunController { error: TaskRunProcess.parseExecuteError(error), } satisfies TaskRunFailedExecutionResult; + // Same OOM check as the success path: if the thrown error parses to + // an OOM, force-kill the process when it's eventually returned (via + // runFinished / stop) instead of recycling it. + if (isOOMRunError(completion.error) || isManualOutOfMemoryError(completion.error)) { + this.discardProcessOnReturn = true; + } + const completionResult = await this.httpClient.dev.completeRunAttempt( run.friendlyId, this.snapshotFriendlyId ?? snapshot.friendlyId, @@ -664,10 +679,22 @@ export class DevRunController { this.isCompletingRun = true; + // Detect OOM in the failure result so we can force-kill the worker + // instead of returning it to the pool. Mirrors the production behavior + // where OOM retry happens on a brand-new container. + if ( + !completion.ok && + (isOOMRunError(completion.error) || isManualOutOfMemoryError(completion.error)) + ) { + this.discardProcessOnReturn = true; + } + // Return process to pool instead of killing it try { const version = this.opts.worker.serverWorker?.version || "unknown"; - await this.opts.taskRunProcessPool.returnProcess(this.taskRunProcess, version); + await this.opts.taskRunProcessPool.returnProcess(this.taskRunProcess, version, { + forceKill: this.discardProcessOnReturn, + }); this.taskRunProcess = undefined; } catch (error) { logger.debug("Failed to return task run process to pool, submitting completion anyway", { @@ -820,7 +847,9 @@ export class DevRunController { if (this.taskRunProcess) { try { const version = this.opts.worker.serverWorker?.version || "unknown"; - await this.opts.taskRunProcessPool.returnProcess(this.taskRunProcess, version); + await this.opts.taskRunProcessPool.returnProcess(this.taskRunProcess, version, { + forceKill: this.discardProcessOnReturn, + }); this.taskRunProcess = undefined; } catch (error) { logger.debug("Failed to return task run process to pool during runFinished", { error }); @@ -854,7 +883,9 @@ export class DevRunController { if (this.taskRunProcess && !this.taskRunProcess.isBeingKilled) { try { const version = this.opts.worker.serverWorker?.version || "unknown"; - await this.opts.taskRunProcessPool.returnProcess(this.taskRunProcess, version); + await this.opts.taskRunProcessPool.returnProcess(this.taskRunProcess, version, { + forceKill: this.discardProcessOnReturn, + }); this.taskRunProcess = undefined; } catch (error) { logger.debug("Failed to return task run process to pool during stop", { error }); diff --git a/packages/core/src/v3/sessionStreams/index.ts b/packages/core/src/v3/sessionStreams/index.ts index 91cdfe7eb..75b372c83 100644 --- a/packages/core/src/v3/sessionStreams/index.ts +++ b/packages/core/src/v3/sessionStreams/index.ts @@ -59,6 +59,14 @@ export class SessionStreamsAPI implements SessionStreamManager { this.#getManager().setLastSeqNum(sessionId, io, seqNum); } + public setMinTimestamp( + sessionId: string, + io: SessionChannelIO, + minTimestamp: number | undefined + ): void { + this.#getManager().setMinTimestamp(sessionId, io, minTimestamp); + } + public shiftBuffer(sessionId: string, io: SessionChannelIO): boolean { return this.#getManager().shiftBuffer(sessionId, io); } diff --git a/packages/core/src/v3/sessionStreams/manager.test.ts b/packages/core/src/v3/sessionStreams/manager.test.ts new file mode 100644 index 000000000..6089d7057 --- /dev/null +++ b/packages/core/src/v3/sessionStreams/manager.test.ts @@ -0,0 +1,151 @@ +import { describe, expect, it } from "vitest"; +import { StandardSessionStreamManager } from "./manager.js"; +import type { ApiClient } from "../apiClient/index.js"; +import type { SSEStreamPart } from "../apiClient/runStream.js"; + +// Single-shot mock that mimics S2's long-poll: delivers `records` once via +// `onPart` on the first subscribe call, then keeps the returned async +// iterable OPEN until the abort signal fires. Real S2 keeps the SSE +// connection alive on a long-poll; the manager's `runTail` finally / +// reconnect path only fires when the connection actually closes. Returning +// an empty stream synchronously triggers a tight reconnect loop, so the +// mock parks indefinitely instead. +function singleShotApiClient( + records: Array<{ id: string; chunk: unknown; timestamp: number }> +): ApiClient { + let delivered = false; + return { + async subscribeToSessionStream( + _sessionIdOrExternalId: string, + _io: "out" | "in", + options?: { onPart?: (part: SSEStreamPart) => void; signal?: AbortSignal } + ) { + if (!delivered) { + delivered = true; + for (const record of records) { + options?.onPart?.(record as SSEStreamPart); + } + } + const signal = options?.signal; + return (async function* () { + if (signal?.aborted) return; + await new Promise((resolve) => { + if (!signal) { + // No signal — block the stream forever; tests must + // explicitly call `disconnectStream` / `disconnect` to + // unblock. + return; + } + signal.addEventListener("abort", () => resolve(), { once: true }); + }); + })() as unknown as Awaited>; + }, + } as unknown as ApiClient; +} + +describe("StandardSessionStreamManager — minTimestamp filter", () => { + const sessionId = "session-1"; + const io = "in" as const; + + it("dispatches records when no filter is set", async () => { + const records = [ + { id: "0", chunk: { kind: "message", payload: { id: "u1" } }, timestamp: 1000 }, + { id: "1", chunk: { kind: "message", payload: { id: "u2" } }, timestamp: 2000 }, + ]; + const manager = new StandardSessionStreamManager(singleShotApiClient(records), "http://localhost"); + + const first = await manager.once(sessionId, io); + expect(first).toEqual({ ok: true, output: { kind: "message", payload: { id: "u1" } } }); + + const second = await manager.once(sessionId, io); + expect(second).toEqual({ ok: true, output: { kind: "message", payload: { id: "u2" } } }); + + manager.disconnectStream(sessionId, io); // stop reconnect loop + manager.disconnect(); + }); + + it("drops records whose timestamp is <= minTimestamp", async () => { + const records = [ + { id: "0", chunk: { kind: "message", payload: { id: "u1" } }, timestamp: 1000 }, + { id: "1", chunk: { kind: "message", payload: { id: "u2" } }, timestamp: 2000 }, + { id: "2", chunk: { kind: "message", payload: { id: "u3" } }, timestamp: 3000 }, + ]; + const manager = new StandardSessionStreamManager(singleShotApiClient(records), "http://localhost"); + + // Cutoff at 2000 (inclusive: `<=` is dropped). Only u3 should pass. + manager.setMinTimestamp(sessionId, io, 2000); + + const passed = await manager.once(sessionId, io, { timeoutMs: 200 }); + expect(passed).toEqual({ ok: true, output: { kind: "message", payload: { id: "u3" } } }); + + manager.disconnectStream(sessionId, io); + manager.disconnect(); + }); + + it("clears the filter when set to undefined", async () => { + const records = [ + { id: "0", chunk: { kind: "message", payload: { id: "u1" } }, timestamp: 1000 }, + ]; + const manager = new StandardSessionStreamManager(singleShotApiClient(records), "http://localhost"); + + manager.setMinTimestamp(sessionId, io, 5000); + manager.setMinTimestamp(sessionId, io, undefined); + + const passed = await manager.once(sessionId, io, { timeoutMs: 200 }); + expect(passed).toEqual({ ok: true, output: { kind: "message", payload: { id: "u1" } } }); + + manager.disconnectStream(sessionId, io); + manager.disconnect(); + }); + + it("filter is per-(sessionId, io) and doesn't bleed across streams", async () => { + const inApi = singleShotApiClient([ + { id: "0", chunk: { kind: "in-record" }, timestamp: 1000 }, + ]); + const manager = new StandardSessionStreamManager(inApi, "http://localhost"); + + manager.setMinTimestamp(sessionId, "in", 5000); + + // The "out" stream uses the same singleShotApiClient instance — its + // single-shot delivers the same fixture, but the filter doesn't apply + // to "out" so the record passes. + const outResult = await manager.once(sessionId, "out", { timeoutMs: 200 }); + expect(outResult).toEqual({ ok: true, output: { kind: "in-record" } }); + + // The "in" stream is filtered (minTimestamp=5000, record ts=1000): the + // once() call should idle-timeout instead of resolving with the record. + // But the singleShot instance has already delivered to the "out" tail, + // so the "in" tail will get nothing on first connect anyway. Use a + // separate manager+api to keep the assertion crisp. + const inApi2 = singleShotApiClient([ + { id: "0", chunk: { kind: "in-record-2" }, timestamp: 1000 }, + ]); + const manager2 = new StandardSessionStreamManager(inApi2, "http://localhost"); + manager2.setMinTimestamp(sessionId, "in", 5000); + + const inResult = await manager2.once(sessionId, "in", { timeoutMs: 100 }); + expect(inResult.ok).toBe(false); // filter-dropped → idle timeout + + manager.disconnectStream(sessionId, "in"); + manager.disconnectStream(sessionId, "out"); + manager.disconnect(); + manager2.disconnectStream(sessionId, "in"); + manager2.disconnect(); + }); + + it("reset() clears all per-stream timestamp filters", async () => { + const records = [ + { id: "0", chunk: { kind: "message", payload: { id: "u1" } }, timestamp: 1000 }, + ]; + const manager = new StandardSessionStreamManager(singleShotApiClient(records), "http://localhost"); + + manager.setMinTimestamp(sessionId, io, 5000); + manager.reset(); + + const passed = await manager.once(sessionId, io, { timeoutMs: 200 }); + expect(passed).toEqual({ ok: true, output: { kind: "message", payload: { id: "u1" } } }); + + manager.disconnectStream(sessionId, io); + manager.disconnect(); + }); +}); diff --git a/packages/core/src/v3/sessionStreams/manager.ts b/packages/core/src/v3/sessionStreams/manager.ts index 92d8459e6..d40af9e7b 100644 --- a/packages/core/src/v3/sessionStreams/manager.ts +++ b/packages/core/src/v3/sessionStreams/manager.ts @@ -36,6 +36,13 @@ export class StandardSessionStreamManager implements SessionStreamManager { private onceWaiters = new Map(); private buffer = new Map(); private tails = new Map(); + // Per-stream lower-bound timestamp filter. When set, records whose + // SSE timestamp is <= the bound are dropped before dispatch — used by + // chat.agent on OOM-retry boot to skip session.in records belonging + // to turns that already completed on the prior attempt. The filter + // is consulted in `runTail`'s `onPart` so the buffer never sees the + // dropped records. + private minTimestamps = new Map(); // Keys that were explicitly torn down by `disconnectStream`. The tail's // `.finally` reconnect path checks this so a long-lived persistent handler // (e.g. `chat.agent`'s run-level `stopInput.on(...)`) doesn't silently @@ -164,6 +171,19 @@ export class StandardSessionStreamManager implements SessionStreamManager { } } + setMinTimestamp( + sessionId: string, + io: SessionChannelIO, + minTimestamp: number | undefined + ): void { + const key = keyFor(sessionId, io); + if (minTimestamp === undefined) { + this.minTimestamps.delete(key); + } else { + this.minTimestamps.set(key, minTimestamp); + } + } + shiftBuffer(sessionId: string, io: SessionChannelIO): boolean { const key = keyFor(sessionId, io); const buffered = this.buffer.get(key); @@ -213,6 +233,7 @@ export class StandardSessionStreamManager implements SessionStreamManager { reset(): void { this.disconnect(); this.seqNums.clear(); + this.minTimestamps.clear(); this.handlers.clear(); for (const [, waiters] of this.onceWaiters) { @@ -271,16 +292,40 @@ export class StandardSessionStreamManager implements SessionStreamManager { const key = keyFor(sessionId, io); try { const lastSeq = this.seqNums.get(key); + // Dispatch is driven from `onPart` (not the for-await loop) so each + // record reaches dispatch with its full SSE metadata in scope — + // specifically the timestamp, which we need for the per-stream + // min-timestamp filter. The for-await loop below just drains the + // pipeThrough output to keep the source flowing. const stream = await this.apiClient.subscribeToSessionStream(sessionId, io, { signal, baseUrl: this.baseUrl, timeoutInSeconds: 600, lastEventId: lastSeq !== undefined ? String(lastSeq) : undefined, onPart: (part) => { + if (signal.aborted) return; const seqNum = parseInt(part.id, 10); if (Number.isFinite(seqNum)) { this.seqNums.set(key, seqNum); } + + // Min-timestamp filter: drop records older than (or at) the + // bound. Used to skip already-processed records on OOM-retry + // boot. + const minTs = this.minTimestamps.get(key); + if (minTs !== undefined && part.timestamp <= minTs) { + return; + } + + let data: unknown = part.chunk; + if (typeof data === "string") { + try { + data = JSON.parse(data); + } catch { + // keep as string + } + } + this.#dispatch(key, data); }, onComplete: () => { if (this.debug) { @@ -294,21 +339,10 @@ export class StandardSessionStreamManager implements SessionStreamManager { }, }); - for await (const record of stream) { + // Drain to keep the pipeThrough flowing. Records were already + // dispatched in `onPart`, so the body here is a no-op. + for await (const _record of stream) { if (signal.aborted) break; - - let data: unknown; - if (typeof record === "string") { - try { - data = JSON.parse(record); - } catch { - data = record; - } - } else { - data = record; - } - - this.#dispatch(key, data); } } catch (error) { if (error instanceof Error && error.name === "AbortError") return; diff --git a/packages/core/src/v3/sessionStreams/noopManager.ts b/packages/core/src/v3/sessionStreams/noopManager.ts index 388598494..c1c3c38dc 100644 --- a/packages/core/src/v3/sessionStreams/noopManager.ts +++ b/packages/core/src/v3/sessionStreams/noopManager.ts @@ -31,6 +31,12 @@ export class NoopSessionStreamManager implements SessionStreamManager { setLastSeqNum(_sessionId: string, _io: SessionChannelIO, _seqNum: number): void {} + setMinTimestamp( + _sessionId: string, + _io: SessionChannelIO, + _minTimestamp: number | undefined + ): void {} + shiftBuffer(_sessionId: string, _io: SessionChannelIO): boolean { return false; } diff --git a/packages/core/src/v3/sessionStreams/types.ts b/packages/core/src/v3/sessionStreams/types.ts index 81820dae5..2310fabae 100644 --- a/packages/core/src/v3/sessionStreams/types.ts +++ b/packages/core/src/v3/sessionStreams/types.ts @@ -45,6 +45,20 @@ export interface SessionStreamManager { /** Advance the last-seen sequence number (prevents SSE replay after `.wait` resume). */ setLastSeqNum(sessionId: string, io: SessionChannelIO, seqNum: number): void; + /** + * Set a per-stream lower-bound SSE timestamp. Records whose timestamp + * is `<= minTimestamp` are dropped before dispatch. Used by chat.agent + * on OOM-retry boot to skip session.in records belonging to turns + * that already completed on the prior attempt. + * + * Pass `undefined` to clear the filter. + */ + setMinTimestamp( + sessionId: string, + io: SessionChannelIO, + minTimestamp: number | undefined + ): void; + /** Remove and discard the first buffered record. Returns true if one was removed. */ shiftBuffer(sessionId: string, io: SessionChannelIO): boolean; diff --git a/packages/core/src/v3/test/test-session-stream-manager.ts b/packages/core/src/v3/test/test-session-stream-manager.ts index b8db1cb8e..fa022573f 100644 --- a/packages/core/src/v3/test/test-session-stream-manager.ts +++ b/packages/core/src/v3/test/test-session-stream-manager.ts @@ -138,6 +138,15 @@ export class TestSessionStreamManager implements SessionStreamManager { this.seqNums.set(keyFor(sessionId, io), seqNum); } + setMinTimestamp( + _sessionId: string, + _io: SessionChannelIO, + _minTimestamp: number | undefined + ): void { + // No filter applied in tests; the test harness drives records directly + // and the chat.agent retry path is exercised separately. + } + shiftBuffer(sessionId: string, io: SessionChannelIO): boolean { const key = keyFor(sessionId, io); const buffered = this.buffer.get(key); diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 6116e44eb..a306cf453 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -9,7 +9,10 @@ import { type InputStreamWaitWithIdleTimeoutOptions, isSchemaZodEsque, logger, + type MachinePresetName, ManualWaitpointPromise, + OutOfMemoryError, + sessionStreams, type PipeStreamResult, type RealtimeDefinedInputStream, type RealtimeDefinedStream, @@ -134,6 +137,53 @@ const chatTurnContextKey = locals.create("chat.turnContext"); */ const chatSessionHandleKey = locals.create("chat.sessionHandle"); +/** + * Scan `session.out` for the latest `trigger:turn-complete` chunk and + * return its SSE timestamp. Used at OOM-retry boot to derive a + * lower-bound timestamp for the `session.in` filter — records older + * than `T_last_complete` belong to turns that already completed on the + * prior attempt and are dropped before they reach the turn loop. + * + * Implementation is a streaming scan: subscribes via the existing SSE + * endpoint with a short `timeoutInSeconds`, processes each part inline, + * and discards the chunk body so memory stays O(1) regardless of how + * many records are on `session.out`. Bandwidth scales linearly with + * stream length but the scan only fires on retry — a rare event. + * + * Returns `undefined` if no `trigger:turn-complete` chunk has been + * written yet (first-turn OOM, no completed turns to dedup against). + * @internal + */ +async function findLatestTurnCompleteTimestamp( + chatId: string +): Promise { + const apiClient = apiClientManager.clientOrThrow(); + let latestTs: number | undefined; + const stream = await apiClient.subscribeToSessionStream(chatId, "out", { + timeoutInSeconds: 1, + onPart: (part) => { + let chunk: unknown = part.chunk; + if (typeof chunk === "string") { + try { + chunk = JSON.parse(chunk); + } catch { + return; + } + } + if (chunk && typeof chunk === "object" && (chunk as { type?: unknown }).type === "trigger:turn-complete") { + latestTs = part.timestamp; + } + }, + }); + // Drain the stream to drive `onPart`. We don't accumulate the chunks — + // each iteration discards the data immediately, so a long session.out + // doesn't blow memory on the retry-boot worker. + for await (const _ of stream) { + // intentionally empty + } + return latestTs; +} + /** * Resolve the Session handle for the current chat.agent run. Throws if * called outside of a chat.agent `run()` — every internal consumer is @@ -3378,8 +3428,42 @@ export type ChatAgentOptions< ChatTaskWirePayload>, unknown >, - "run" + "run" | "retry" > & { + /** + * Fallback machine preset to use when an attempt fails with an + * out-of-memory (OOM) error. Setting this enables a single OOM retry: + * the next attempt boots on the larger machine, and the chat picks + * up via the standard continuation path (same `chatId` / Session, + * accumulator rebuilds via `hydrateMessages` or post-`onTurnStart` + * persisted state). + * + * Set `machine` (top-level `TaskOptions`) to control the *default* + * machine the agent runs on. `oomMachine` is the *retry-only* swap. + * + * Note: an OOM retry restarts the entire turn from the top — the + * model call and any in-flight tool executes re-run on the larger + * machine. Make tool executes idempotent or persist results before + * returning if you can't tolerate re-execution. + * + * Generic `retry` options are not exposed on `chat.agent` because + * arbitrary retries against an LLM-driven loop tend to be expensive + * and side-effecting. If you need richer retry semantics, drop down + * to `chat.task` (the raw primitive). + * + * @example + * ```ts + * chat.agent({ + * id: "my-chat", + * machine: "small-1x", + * oomMachine: "medium-2x", + * run: async ({ messages, signal }) => + * streamText({ model, messages, abortSignal: signal }), + * }); + * ``` + */ + oomMachine?: MachinePresetName; + /** * Schema for validating `clientData` from the frontend. * Accepts Zod, ArkType, Valibot, or any supported schema library. @@ -4025,19 +4109,28 @@ function chatAgent< onChatSuspend, onChatResume, exitAfterPreloadIdle = false, + oomMachine, ...restOptions } = options; const parseClientData = clientDataSchema ? getSchemaParseFn(clientDataSchema) : undefined; const parseAction = actionSchema ? getSchemaParseFn(actionSchema) : undefined; + // chat.agent does not expose generic retry options (see docstring on + // `oomMachine`). The only opt-in is an OOM-triggered machine swap. If + // `oomMachine` is set we allow one retry on a larger machine; otherwise + // we keep the historical no-retry default. + const retry = oomMachine + ? { maxAttempts: 2, outOfMemory: { machine: oomMachine } } + : { maxAttempts: 1 }; + const task = createTask< TIdentifier, ChatTaskWirePayload>, unknown >({ - retry: { maxAttempts: 1 }, ...restOptions, + retry, triggerSource: "agent", agentConfig: { type: "ai-sdk-chat" }, run: async ( @@ -4104,6 +4197,37 @@ function chatAgent< // but in frontend-friendly UIMessage format (with parts, id, etc.). let accumulatedUIMessages: TUIMessage[] = []; + // OOM-retry boot: a fresh worker subscribes to `session.in` from + // seq 0 and would re-deliver every record ever appended to that + // session — including the messages from turns that already + // completed on the prior attempt. Without dedup, the loop would + // re-process them as fresh turns. + // + // We derive the cutoff from `session.out`: the latest + // `trigger:turn-complete` chunk's timestamp is the high-water + // mark of completed work. Any session.in record with a timestamp + // at or below that mark belongs to a completed turn and is + // dropped by the SessionStreamManager filter before dispatch. + // + // No customer setup required. With `hydrateMessages` configured, + // the OOM'd turn re-runs against the full prior conversation (the + // common pattern). Without it, the OOM'd turn re-runs against + // whatever the chat.agent's default accumulator can rebuild from + // `payload.messages` — degraded continuity but no duplicate work. + if (ctx.attempt.number > 1) { + try { + const cutoff = await findLatestTurnCompleteTimestamp(payload.chatId); + if (cutoff !== undefined) { + sessionStreams.setMinTimestamp(payload.chatId, "in", cutoff); + } + } catch (error) { + logger.warn( + "chat.agent OOM-retry session.out scan failed; session.in dedup not applied", + { error: error instanceof Error ? error.message : String(error) } + ); + } + } + // Token usage tracking across turns let previousTurnUsage: LanguageModelUsage | undefined; let cumulativeUsage: LanguageModelUsage = emptyUsage(); @@ -5711,6 +5835,17 @@ function chatAgent< throw turnError; } + // OOM errors must escape the turn loop so the task runtime can + // honor `retry.outOfMemory.machine` (set on chat.agent via + // `oomMachine`). Catching them here would keep the dead worker + // alive and defeat the machine swap. Re-throw and let the + // runtime dispatch the retry on a larger machine; recovery on + // attempt 2 picks up via the standard continuation path + // (same chatId / Session, accumulator rehydrates). + if (turnError instanceof OutOfMemoryError) { + throw turnError; + } + try { await withChatWriter(async (writer) => { const errorText =