Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 22 additions & 0 deletions src/subagent/index.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import {
evaluateSubAgentStop,
forcedStopReport,
formatSubAgentReport,
formatTurnTokenNotice,
parseSubAgentReport,
appendSubAgentParentHints,
EMPTY_THRASH_STATE,
Expand Down Expand Up @@ -1297,6 +1298,27 @@ describe("SubAgentDirector stall management", () => {
});
});

describe("submit_result turn token notice", () => {
test("the dispatch brief embeds the shared token notice verbatim", () => {
const token = "01a09856-4dd3-7209-a3df-d7e543dc4ffe";
const brief = buildDispatchBrief({
description: "token probe",
prompt: "do the thing",
turnToken: token,
});
// Byte-identity: the brief and followup steers render the same contract
// through one shared function, so a worker can never see two wordings.
expect(brief).toContain(formatTurnTokenNotice(token));
expect(formatTurnTokenNotice(token)).toContain(
"A mismatched token means this turn was superseded",
);
// Non-leaf dispatches state no token.
expect(
buildDispatchBrief({ description: "plain", prompt: "do the thing" }),
).not.toContain("## Turn token");
});
});

describe("buildDispatchBrief typed spawn contract", () => {
test("renders Intent, Success criteria, Do not, and report_focus only when set", () => {
const full = buildDispatchBrief({
Expand Down
1 change: 1 addition & 0 deletions src/subagent/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@ export {
buildDispatchBrief,
demoteNestedReportHeadings,
formatSubAgentReport,
formatTurnTokenNotice,
hasPlanFindings,
hasReportEnvelope,
parseSubAgentReport,
Expand Down
19 changes: 13 additions & 6 deletions src/subagent/report.ts
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,18 @@ export interface DispatchBrief {
turnToken?: string;
}

/**
* States the active submit_result token. Shared by the initial dispatch brief
* and followup steers so both turns state the same contract.
*/
export function formatTurnTokenNotice(turnToken: string): string {
return [
"## Turn token",
turnToken,
`If you call submit_result, pass turn_token="${turnToken}" exactly. A mismatched token means this turn was superseded — do not resubmit under it.`,
].join("\n");
}

export function buildDispatchBrief(brief: DispatchBrief): string {
const parts: string[] = [
`# Dispatch brief: ${brief.description}`,
Expand Down Expand Up @@ -103,12 +115,7 @@ export function buildDispatchBrief(brief: DispatchBrief): string {
}
parts.push("", "## Report shape", ...reportLines);
if (brief.turnToken !== undefined && brief.turnToken.length > 0) {
parts.push(
"",
"## Turn token",
brief.turnToken,
`If you call submit_result, pass turn_token="${brief.turnToken}" exactly. A mismatched token means this turn was superseded — do not resubmit under it.`,
);
parts.push("", formatTurnTokenNotice(brief.turnToken));
}
return parts.join("\n");
}
Expand Down
151 changes: 151 additions & 0 deletions src/subagent/run-submit-result-rotation.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,151 @@
/**
* Steering rotation: runSubAgent mints the submit_result turn token at leaf
* dispatch and rotates it on every followup steer, so a worker holding the
* dispatched token cannot submit after its turn was superseded. The pure
* evaluator half (old token rejected, budget reset) is covered in
* submit-result.test.ts; this test drives the real runSubAgent wiring end to
* end — the one seam the pure tests cannot see is whether followup actually
* mints, swaps, and re-states the token. Same stub-agent pattern as
* followup-live-agent.test.ts.
*/
import { describe, expect, test } from "bun:test";
import { mkdtemp } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";

import { withMockedModuleDuring } from "../../tests/helpers/mock-module.js";
import { defined } from "../../tests/helpers/defined.js";
import { createPermissionGate } from "../permission/gate.js";
import type { RunSubAgentParams } from "./types.js";

const testPermissionGate = createPermissionGate({
approvals: [],
interactive: false,
skipPermissions: true,
reactorGated: false,
});

/** First send hangs (the run stays alive for steering); later sends resolve. */
function createRotatingStubAgent(sendLog: string[]) {
return {
async send(content: string, optsSend?: { signal?: AbortSignal }) {
sendLog.push(content);
if (sendLog.length === 1) {
return await new Promise((_, reject) => {
if (optsSend?.signal?.aborted === true) {
reject(
optsSend.signal.reason instanceof Error
? optsSend.signal.reason
: new Error("aborted"),
);
return;
}
optsSend?.signal?.addEventListener(
"abort",
() => {
const reason = defined(optsSend.signal).reason;
reject(reason instanceof Error ? reason : new Error("aborted"));
},
{ once: true },
);
});
}
return {
type: "reply" as const,
reply: `reply #${sendLog.length}`,
turn: { role: "assistant", content: [] },
};
},
stream: () =>
(async function* () {
yield* [];
})(),
deliver: () => undefined,
close: async () => undefined,
setSource: () => undefined,
setSources: () => undefined,
history: async () => [],
checkpoints: async () => [],
readAt: async () => [],
blobReader: {},
};
}

function tokenOfSend(send: string): string | undefined {
return send.match(/^## Turn token\n(.+)$/m)?.[1];
}

describe("submit_result token rotation on steering", () => {
test("followup rotates the leaf turn token and states the replacement in the steer", async () => {
const cwd = await mkdtemp(join(tmpdir(), "cl6946-token-rotation-"));
const sendLog: string[] = [];

const outcome = await withMockedModuleDuring(
import.meta.resolve("../agent/live-tool-dispatch.js"),
(real: typeof import("../agent/live-tool-dispatch.js")) => ({
...real,
createAgentWithLiveToolDispatch: async () =>
createRotatingStubAgent(sendLog) as unknown as Awaited<
ReturnType<typeof real.createAgentWithLiveToolDispatch>
>,
}),
async () => {
const { runSubAgent } = await import("./run.js");

let handles:
| {
close: (ms?: number) => Promise<void>;
interrupt: () => void;
followup: (message: string) => Promise<string>;
}
| undefined;

const params: RunSubAgentParams = {
cwd,
workdirBase: join(cwd, ".ctx"),
permissionGate: testPermissionGate,
provider: {
providerName: "test",
baseURL: "http://localhost",
model: "test-model",
},
description: "token rotation probe",
prompt: "hold for steering",
persist: true,
tier: "leaf",
onAgentReady: (h) => {
handles = h;
},
};

const runPromise = runSubAgent(params);
for (let i = 0; i < 500 && sendLog.length < 1; i++) {
await new Promise((resolve) => setTimeout(resolve, 1));
}
if (handles === undefined) throw new Error("onAgentReady never fired");

const reply = await handles.followup("new orders: pivot to X");
// followup replaced the per-turn interrupt controller, so interrupt()
// can no longer reach the still-hung first send — close() aborts the
// run controller instead and settles the run for cleanup. Both can
// reject with the abort reason; settlement is all we need.
await handles.close().catch(() => undefined);
await runPromise.catch(() => undefined);
return { reply };
},
);

expect(outcome.reply).toBe("reply #2");
expect(sendLog.length).toBe(2);

// The dispatched brief states the first token; the steer carries the
// followup message plus a DIFFERENT token — the old one dies here.
const dispatched = tokenOfSend(defined(sendLog[0]));
const steered = tokenOfSend(defined(sendLog[1]));
expect(dispatched).toBeDefined();
expect(steered).toBeDefined();
expect(defined(steered)).not.toBe(defined(dispatched));
expect(defined(sendLog[1]).startsWith("new orders: pivot to X")).toBe(true);
expect(defined(sendLog[1])).toContain(`turn_token="${defined(steered)}"`);
});
});
27 changes: 23 additions & 4 deletions src/subagent/run.ts
Original file line number Diff line number Diff line change
Expand Up @@ -116,6 +116,7 @@ import {
appendActivitySummary,
buildDispatchBrief,
formatSubAgentReport,
formatTurnTokenNotice,
parseSubAgentReport,
subAgentToolName,
} from "./report.js";
Expand All @@ -140,6 +141,7 @@ import { createReadAgentTraceTool } from "./trace-tool.js";
import {
createSubmitResultState,
evaluateSubmitResult,
resetSubmitResultTurn,
SUBMIT_RESULT_MAX_CORRECTIONS,
} from "./submit-result.js";
import {
Expand Down Expand Up @@ -578,8 +580,9 @@ async function runSubAgentInner(
const permissionGate = workerPermissionGate(params.permissionGate);
// Identifies this dispatch to submit_result so a submission survives
// only for the turn it was spawned under — a stale call from a redirected
// orchestrator (echoing an old token) is rejected.
const turnToken = params.tier === "leaf" ? generateSessionId() : undefined;
// orchestrator (echoing an old token) is rejected. Steering (followup)
// rotates it: the old token dies with the superseded turn.
let turnToken = params.tier === "leaf" ? generateSessionId() : undefined;
const submitResultState = createSubmitResultState();
const askDirectorState = createAskDirectorState();
const spawnRegistry = createSubAgentSpawnRegistryPlugin();
Expand Down Expand Up @@ -747,8 +750,15 @@ async function runSubAgentInner(
handler: async (
rawArgs: Record<string, unknown>,
): Promise<string> => {
// Read the live binding (not the dispatch-time value): steering
// rotates turnToken, so a submission under the old token lands here
// and must be rejected as stale.
const currentToken = turnToken;
if (currentToken === undefined) {
return "Error: this run has no turn token, so submit_result cannot verify freshness.";
}
const outcome = evaluateSubmitResult({
turnToken,
turnToken: currentToken,
submittedToken: rawArgs.turn_token,
result: rawArgs.result,
...(params.reportType !== undefined
Expand Down Expand Up @@ -1353,7 +1363,16 @@ async function runSubAgentInner(
const followup = async (message: string): Promise<string> => {
resetAskDirectorTurn(askDirectorState);
interruptController = new AbortController();
const result = await sendWithProviderFailure(message, {
// A steer supersedes the dispatched turn: mint a fresh submit_result
// token (the handler closure reads this binding, so the old token is
// rejected from here on) and hand the worker the replacement.
let steered = message;
if (turnToken !== undefined) {
turnToken = generateSessionId();
resetSubmitResultTurn(submitResultState);
steered = `${message}\n\n${formatTurnTokenNotice(turnToken)}`;
}
const result = await sendWithProviderFailure(steered, {
signal: sendAbortSignal(),
});
if (terminalProviderError !== undefined) {
Expand Down
44 changes: 44 additions & 0 deletions src/subagent/submit-result.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import { type } from "arktype";
import {
createSubmitResultState,
evaluateSubmitResult,
resetSubmitResultTurn,
} from "./submit-result.js";

const TOKEN = "turn-abc123";
Expand Down Expand Up @@ -86,3 +87,46 @@ describe("evaluateSubmitResult", () => {
expect(capped.message).toContain("correction cap");
});
});

describe("submit_result turn rotation", () => {
test("a steered turn rejects the old token and resets the correction budget", () => {
const state = createSubmitResultState();
const outputType = type({ verdict: "'pass'|'fail'" });

// The dispatched turn burns one correction on a malformed result.
const bad = evaluateSubmitResult({
turnToken: TOKEN,
submittedToken: TOKEN,
result: { verdict: "maybe" },
outputType,
state,
});
expect(bad.ok).toBe(false);
expect(state.corrections).toBe(1);

// Steering rotates the token and resets the budget (run.ts followup).
const rotated = "turn-def456";
resetSubmitResultTurn(state);

const stale = evaluateSubmitResult({
turnToken: rotated,
submittedToken: TOKEN,
result: { verdict: "pass" },
outputType,
state,
});
expect(stale.ok).toBe(false);
expect(stale.message).toContain("turn_token does not match");
expect(state.corrections).toBe(0);

const fresh = evaluateSubmitResult({
turnToken: rotated,
submittedToken: rotated,
result: { verdict: "pass" },
outputType,
state,
});
expect(fresh.ok).toBe(true);
expect(fresh.message).toBe("Result accepted.");
});
});
8 changes: 8 additions & 0 deletions src/subagent/submit-result.ts
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,14 @@ export function createSubmitResultState(): SubmitResultState {
return { corrections: 0 };
}

/**
* Steering starts a new turn: drop the finished turn's correction count so the
* replacement token gets a full budget. The caller mints the new token.
*/
export function resetSubmitResultTurn(state: SubmitResultState): void {
state.corrections = 0;
}

export interface SubmitResultInput {
/** This turn's token, generated by runSubAgent at dispatch time. */
turnToken: string;
Expand Down
Loading