fix(agents): prevent repeated Code Mode output across waits (#115213)

* fix(agents): prevent duplicate code mode output across waits

* fix(agents): advance code mode output frontier on wait
This commit is contained in:
Peter Steinberger
2026-07-28 12:10:16 -04:00
committed by GitHub
parent 13f20055f0
commit bb295df6ea
3 changed files with 185 additions and 12 deletions

View File

@@ -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),
};

View File

@@ -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),
};
}

View File

@@ -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 = {