diff --git a/src/agents/code-mode-execution.ts b/src/agents/code-mode-execution.ts index f4afafe8840d..a2d525032f05 100644 --- a/src/agents/code-mode-execution.ts +++ b/src/agents/code-mode-execution.ts @@ -36,6 +36,7 @@ import { settledBridgeRequestsInCompletionOrder, snapshotState, storeSnapshotState, + takeUndeliveredCodeModeRunOutput, telemetry, waitForPendingBridgeSettlement, type PendingBridgeState, @@ -215,6 +216,7 @@ async function settleCodeModeResult(params: { runtime: ToolSearchRuntime; namespaceRuntime: CodeModeNamespaceRuntime; deadlineMs: number; + deliveredOutputCount?: number; pending?: PendingBridgeState[]; activeRunId?: string; signal?: AbortSignal; @@ -224,6 +226,7 @@ async function settleCodeModeResult(params: { let pending = params.pending ?? []; const activeRunId = params.activeRunId ?? `cm_${randomUUID()}`; const output = params.output; + const deliveredOutputCount = params.deliveredOutputCount ?? 0; // One exec/wait call shares a single wall-clock deadline across its initial // worker run and this inline settle phase, so auto-draining bridge calls // cannot stack a second full `timeoutMs` budget on top of the run that @@ -235,7 +238,7 @@ async function settleCodeModeResult(params: { status: "failed" as const, error: "code mode execution aborted", code: "aborted" as const, - output, + output: output.slice(deliveredOutputCount), replaySafe: params.replaySafe, telemetry: telemetry(params.runtime), }); @@ -259,7 +262,7 @@ async function settleCodeModeResult(params: { status: "failed" as const, error: "restart-safe code mode cannot call namespace tools.", code: "invalid_input" as const, - output, + output: output.slice(deliveredOutputCount), replaySafe: true, telemetry: telemetry(params.runtime), }; @@ -328,6 +331,7 @@ async function settleCodeModeResult(params: { runtime: params.runtime, namespaceRuntime: params.namespaceRuntime, output, + deliveredOutputCount, }); } // Deliver the settled frontier only. Unresolved sibling promises remain @@ -379,7 +383,7 @@ async function settleCodeModeResult(params: { status: "failed" as const, error: "restart-safe code mode cannot call side-effecting tools.", code: "invalid_input" as const, - output, + output: output.slice(deliveredOutputCount), replaySafe: true, telemetry: telemetry(params.runtime), }; @@ -426,6 +430,7 @@ async function settleCodeModeResult(params: { runtime: params.runtime, namespaceRuntime: params.namespaceRuntime, output, + deliveredOutputCount, }); } catch (error) { cancelPendingBridgeStates(pending); @@ -444,6 +449,7 @@ async function settleCodeModeResult(params: { runtime: params.runtime, namespaceRuntime: params.namespaceRuntime, output, + deliveredOutputCount, replaySafe: params.replaySafe, settlementMode: result.settlementMode, signal: params.signal, @@ -460,7 +466,7 @@ async function settleCodeModeResult(params: { }); return { ...result, - output, + output: output.slice(deliveredOutputCount), replaySafe: params.replaySafe, telemetry: telemetry(params.runtime), }; @@ -512,7 +518,7 @@ export async function runWait(params: { status: "failed" as const, error: "code mode execution aborted", code: "aborted" as const, - output: state.output, + output: takeUndeliveredCodeModeRunOutput(state), replaySafe: state.replaySafe, telemetry: telemetry(state.runtime), }; @@ -527,7 +533,7 @@ export async function runWait(params: { reason: codeModeWaitingReason(pending.length > 0 ? pending : state.pending), pendingToolCalls: pendingToolCalls(pending.length > 0 ? pending : state.pending), replaySafe: state.replaySafe, - output: state.output, + output: takeUndeliveredCodeModeRunOutput(state), telemetry: telemetry(state.runtime), }; } @@ -571,6 +577,7 @@ export async function runWait(params: { config: state.config, runtime: state.runtime, namespaceRuntime: state.namespaceRuntime, + deliveredOutputCount: state.deliveredOutputCount, pending, activeRunId: state.runId, signal: params.signal, @@ -586,7 +593,7 @@ export async function runWait(params: { status: "failed" as const, error: codeModeFailureMessage(error), code: codeModeFailureCode(error), - output: state.output, + output: takeUndeliveredCodeModeRunOutput(state), replaySafe: state.replaySafe, telemetry: telemetry(state.runtime), }; diff --git a/src/agents/code-mode-state.ts b/src/agents/code-mode-state.ts index c6fe3252f7f2..0e1ad5ced373 100644 --- a/src/agents/code-mode-state.ts +++ b/src/agents/code-mode-state.ts @@ -36,6 +36,8 @@ type CodeModeRunState = { // True only when every future bridge call is enforced read-only before execution. replaySafe: boolean; output: unknown[]; + // Retain all output for cumulative limits, but never replay blocks already returned to the model. + deliveredOutputCount: number; createdAt: number; expiresAt: number; agentWaitRetainUntil?: number; @@ -105,6 +107,13 @@ export function disposeCodeModeRun(runId: string): void { scheduleActiveRunExpiry(); } +/** Advance the snapshot frontier before exposing output to a wait observer. */ +export function takeUndeliveredCodeModeRunOutput(state: CodeModeRunState): unknown[] { + const output = state.output.slice(state.deliveredOutputCount); + state.deliveredOutputCount = state.output.length; + return output; +} + /** Abort each bridge call whose result has not already reached its guest. */ export function cancelPendingBridgeStates(pending: readonly PendingBridgeState[]): void { for (const entry of pending) { @@ -195,6 +204,7 @@ export function snapshotState(params: { runtime: ToolSearchRuntime; namespaceRuntime: CodeModeNamespaceRuntime; output: unknown[]; + deliveredOutputCount?: number; replaySafe: boolean; settlementMode: CodeModeSettlementMode; signal?: AbortSignal; @@ -322,6 +332,7 @@ export function storeSnapshotState(params: { runtime: ToolSearchRuntime; namespaceRuntime: CodeModeNamespaceRuntime; output: unknown[]; + deliveredOutputCount?: number; }) { const now = Date.now(); const expiresAt = resolveCodeModeSnapshotExpiresAt(now, params.config.snapshotTtlSeconds); @@ -348,6 +359,7 @@ export function storeSnapshotState(params: { settlementMode: params.settlementMode, replaySafe: params.replaySafe, output: params.output, + deliveredOutputCount: params.output.length, createdAt: now, expiresAt, agentWaitRetainUntil, @@ -361,7 +373,7 @@ export function storeSnapshotState(params: { reason: codeModeWaitingReason(params.pending), pendingToolCalls: pendingToolCalls(params.pending), replaySafe: params.replaySafe, - output: params.output, + output: params.output.slice(params.deliveredOutputCount ?? 0), telemetry: telemetry(params.runtime), }; } diff --git a/src/agents/code-mode.test.ts b/src/agents/code-mode.test.ts index 1b0c25b75e2e..39b3e64c87ec 100644 --- a/src/agents/code-mode.test.ts +++ b/src/agents/code-mode.test.ts @@ -2177,10 +2177,93 @@ describe("Code Mode", () => { expect(resumed.status).toBe("completed"); expect(resumed.value).toBe("done"); - expect(resumed.output).toEqual([ - { type: "text", text: "before" }, - { type: "text", text: "after" }, - ]); + expect(resumed.output).toEqual([{ type: "text", text: "after" }]); + }); + + it("delivers each yielded output block exactly once across repeated waits", async () => { + const { config, catalogRef, tools: codeModeTools } = createCodeModeHarness(); + applyCodeModeCatalog({ + tools: [...codeModeTools, pluginTool("fake_noop", "Noop")], + config, + sessionId: "session-code-mode", + sessionKey: "agent:main:main", + runId: "run-code-mode", + catalogRef, + }); + + const execTool = expectDefined(codeModeTools[0], "Code Mode exec test invariant"); + const waitTool = expectDefined(codeModeTools[1], "Code Mode wait test invariant"); + const first = resultDetails( + await execTool.execute("code-call-incremental-output", { + code: ` + text("phase 1"); + await yield_control("first pause"); + text("phase 2"); + await yield_control("second pause"); + text("phase 3"); + return "done"; + `, + }), + ); + + expect(first.status).toBe("waiting"); + expect(first.output).toEqual([{ type: "text", text: "phase 1" }]); + + const second = resultDetails( + await waitTool.execute("code-wait-incremental-output-1", { runId: first.runId }), + ); + + expect(second.status).toBe("waiting"); + expect(second.output).toEqual([{ type: "text", text: "phase 2" }]); + + const third = resultDetails( + await waitTool.execute("code-wait-incremental-output-2", { runId: second.runId }), + ); + + expect(third.status).toBe("completed"); + expect(third.value).toBe("done"); + expect(third.output).toEqual([{ type: "text", text: "phase 3" }]); + }); + + it("returns only newly emitted output when a resumed guest fails", async () => { + const { config, catalogRef, tools: codeModeTools } = createCodeModeHarness(); + applyCodeModeCatalog({ + tools: [...codeModeTools, pluginTool("fake_noop", "Noop")], + config, + sessionId: "session-code-mode", + sessionKey: "agent:main:main", + runId: "run-code-mode", + catalogRef, + }); + + const first = resultDetails( + await expectDefined(codeModeTools[0], "Code Mode exec test invariant").execute( + "code-call-incremental-failure", + { + code: ` + text("before pause"); + await yield_control("pause"); + text("before failure"); + throw new Error("resumed failure"); + `, + }, + ), + ); + + expect(first.status).toBe("waiting"); + expect(first.output).toEqual([{ type: "text", text: "before pause" }]); + + const second = resultDetails( + await expectDefined(codeModeTools[1], "Code Mode wait test invariant").execute( + "code-wait-incremental-failure", + { runId: first.runId }, + ), + ); + + expect(second.status).toBe("failed"); + expect(second.error).toContain("resumed failure"); + expect(second.output).toEqual([{ type: "text", text: "before failure" }]); + expect(testing.activeRuns.has(first.runId as string)).toBe(false); }); it("preserves the original exec identity for tool calls after yield and wait", async () => { @@ -2701,6 +2784,7 @@ describe("Code Mode", () => { "code-call-timeout", { code: ` + text("before timeout"); const fast = tools.fake_fast({}); const slow = tools.fake_slow({}); await fast; @@ -2711,6 +2795,7 @@ describe("Code Mode", () => { ), ); expect(first.status).toBe("waiting"); + expect(first.output).toEqual([{ type: "text", text: "before timeout" }]); expect(first.pendingToolCalls).toEqual([expect.objectContaining({ method: "callValue" })]); const runId = first.runId; expect(typeof runId).toBe("string"); @@ -2730,7 +2815,19 @@ describe("Code Mode", () => { ); expect(second.status).toBe("waiting"); + expect(second.output).toEqual([]); expect(second.pendingToolCalls).toEqual([expect.objectContaining({ method: "callValue" })]); + + const third = resultDetails( + await expectDefined(codeModeTools[1], "codeModeTools[1] test invariant").execute( + "code-wait-timeout-again", + { runId }, + ), + ); + + expect(third.status).toBe("waiting"); + expect(third.output).toEqual([]); + expect(third.pendingToolCalls).toEqual([expect.objectContaining({ method: "callValue" })]); }); it("does not load TypeScript for plain JavaScript code mode runs", async () => { @@ -3285,6 +3382,63 @@ describe("Code Mode", () => { expect(testing.activeRuns.size).toBe(beforeRunCount); }); + it("enforces the cumulative output limit across yielded waits", async () => { + const catalogRef = createToolSearchCatalogRef(); + const config = { + tools: { + codeMode: { + enabled: true, + maxOutputBytes: 1024, + }, + }, + } as never; + const ctx = { + config, + runtimeConfig: config, + sessionId: "session-code-mode", + sessionKey: "agent:main:main", + runId: "run-code-mode", + catalogRef, + }; + const tools = createCodeModeTools(ctx); + applyCodeModeCatalog({ + tools: [...tools, pluginTool("fake_noop", "Noop")], + config, + sessionId: "session-code-mode", + sessionKey: "agent:main:main", + runId: "run-code-mode", + catalogRef, + }); + + const first = resultDetails( + await expectDefined(tools[0], "Code Mode exec test invariant").execute( + "code-call-cumulative-output", + { + code: ` + text("a".repeat(600)); + await yield_control("pause"); + text("b".repeat(600)); + return "done"; + `, + }, + ), + ); + + expect(first.status).toBe("waiting"); + expect(first.output).toEqual([{ type: "text", text: "a".repeat(600) }]); + + const second = resultDetails( + await expectDefined(tools[1], "Code Mode wait test invariant").execute( + "code-wait-cumulative-output", + { runId: first.runId }, + ), + ); + + expect(second.status).toBe("failed"); + expect(second.code).toBe("output_limit_exceeded"); + expect(testing.activeRuns.has(first.runId as string)).toBe(false); + }); + it("enforces output limits before auto-draining namespace calls", async () => { const catalogRef = createToolSearchCatalogRef(); const config = {