From 956bd2f94ba3d4a960292452b144a7e91ff06346 Mon Sep 17 00:00:00 2001 From: Eva Date: Sat, 2 May 2026 00:58:29 +0700 Subject: [PATCH] fix: honor later finalize retry candidates --- .../harness/lifecycle-hook-helpers.test.ts | 47 +++++++++++ src/agents/harness/lifecycle-hook-helpers.ts | 79 +++++++++++++------ .../hooks.before-agent-finalize.test.ts | 49 ++++++++++++ src/plugins/hooks.ts | 59 +++++++++++--- 4 files changed, 198 insertions(+), 36 deletions(-) diff --git a/src/agents/harness/lifecycle-hook-helpers.test.ts b/src/agents/harness/lifecycle-hook-helpers.test.ts index 6038126b396..6f1ec3ff7e1 100644 --- a/src/agents/harness/lifecycle-hook-helpers.test.ts +++ b/src/agents/harness/lifecycle-hook-helpers.test.ts @@ -214,6 +214,53 @@ describe("agent harness lifecycle hook helpers", () => { }); }); + it("honors a later finalize retry candidate after an earlier candidate is spent", async () => { + const firstRetry = { + instruction: "regenerate artifacts", + idempotencyKey: "artifacts", + maxAttempts: 1, + }; + const secondRetry = { + instruction: "rerun focused tests", + idempotencyKey: "tests", + maxAttempts: 1, + }; + const result = { + action: "revise", + reason: "retry generated artifacts\n\nretry focused tests", + retry: firstRetry, + }; + Object.defineProperty(result, "retryCandidates", { + enumerable: false, + value: [firstRetry, secondRetry], + }); + const hookRunner = { + hasHooks: () => true, + runBeforeAgentFinalize: vi.fn().mockResolvedValue(result), + }; + + await expect( + runAgentHarnessBeforeAgentFinalizeHook({ + event: EVENT, + ctx: { runId: "run-1", sessionKey: "agent:main:session-1" }, + hookRunner: hookRunner as never, + }), + ).resolves.toEqual({ + action: "revise", + reason: "retry generated artifacts\n\nretry focused tests\n\nregenerate artifacts", + }); + await expect( + runAgentHarnessBeforeAgentFinalizeHook({ + event: EVENT, + ctx: { runId: "run-1", sessionKey: "agent:main:session-1" }, + hookRunner: hookRunner as never, + }), + ).resolves.toEqual({ + action: "revise", + reason: "retry generated artifacts\n\nretry focused tests\n\nrerun focused tests", + }); + }); + it("falls back to retry instruction keys when retry idempotency keys are malformed", async () => { const hookRunner = { hasHooks: () => true, diff --git a/src/agents/harness/lifecycle-hook-helpers.ts b/src/agents/harness/lifecycle-hook-helpers.ts index 18b48f12e0f..cbad11d9aa4 100644 --- a/src/agents/harness/lifecycle-hook-helpers.ts +++ b/src/agents/harness/lifecycle-hook-helpers.ts @@ -146,33 +146,40 @@ function normalizeBeforeAgentFinalizeResult( return reason ? { action: "finalize", reason } : { action: "finalize" }; } if (result?.action === "revise") { - const retryInstruction = normalizeTrimmedString(result.retry?.instruction); - if (retryInstruction) { - const maxAttempts = - typeof result.retry?.maxAttempts === "number" && Number.isFinite(result.retry.maxAttempts) - ? Math.max(1, Math.floor(result.retry.maxAttempts)) - : 1; - const retryRunId = event?.runId ?? event?.sessionId ?? "unknown-run"; - const retryKey = - normalizeTrimmedString(result.retry?.idempotencyKey) || - buildFinalizeRetryInstructionKey(retryInstruction); - const budget = getFinalizeRetryBudget(); - const runBudget = budget.get(retryRunId) ?? new Map(); - const nextCount = (runBudget.get(retryKey) ?? 0) + 1; - runBudget.delete(retryKey); - runBudget.set(retryKey, nextCount); - budget.delete(retryRunId); - budget.set(retryRunId, runBudget); - pruneFinalizeRetryBudget(budget); - if (nextCount > maxAttempts) { - return { action: "continue" }; - } + const retryCandidates = readBeforeAgentFinalizeRetryCandidates(result); + if (retryCandidates.length > 0) { const reason = normalizeTrimmedString(result.reason); - const revisedReason = - reason && reason.includes(retryInstruction) - ? reason - : [reason, retryInstruction].filter(Boolean).join("\n\n"); - return { action: "revise", reason: revisedReason }; + for (const retry of retryCandidates) { + const retryInstruction = normalizeTrimmedString(retry.instruction); + if (!retryInstruction) { + continue; + } + const maxAttempts = + typeof retry.maxAttempts === "number" && Number.isFinite(retry.maxAttempts) + ? Math.max(1, Math.floor(retry.maxAttempts)) + : 1; + const retryRunId = event?.runId ?? event?.sessionId ?? "unknown-run"; + const retryKey = + normalizeTrimmedString(retry.idempotencyKey) || + buildFinalizeRetryInstructionKey(retryInstruction); + const budget = getFinalizeRetryBudget(); + const runBudget = budget.get(retryRunId) ?? new Map(); + const nextCount = (runBudget.get(retryKey) ?? 0) + 1; + runBudget.delete(retryKey); + runBudget.set(retryKey, nextCount); + budget.delete(retryRunId); + budget.set(retryRunId, runBudget); + pruneFinalizeRetryBudget(budget); + if (nextCount > maxAttempts) { + continue; + } + const revisedReason = + reason && reason.includes(retryInstruction) + ? reason + : [reason, retryInstruction].filter(Boolean).join("\n\n"); + return { action: "revise", reason: revisedReason }; + } + return { action: "continue" }; } const reason = normalizeTrimmedString(result.reason); return reason ? { action: "revise", reason } : { action: "continue" }; @@ -180,6 +187,26 @@ function normalizeBeforeAgentFinalizeResult( return { action: "continue" }; } +function readBeforeAgentFinalizeRetryCandidates( + result: PluginHookBeforeAgentFinalizeResult, +): NonNullable[] { + const candidateList = ( + result as { + retryCandidates?: unknown; + } + ).retryCandidates; + if (Array.isArray(candidateList) && candidateList.length > 0) { + return candidateList.filter(isBeforeAgentFinalizeRetry); + } + return isBeforeAgentFinalizeRetry(result.retry) ? [result.retry] : []; +} + +function isBeforeAgentFinalizeRetry( + value: unknown, +): value is NonNullable { + return Boolean(value) && typeof value === "object" && !Array.isArray(value); +} + function normalizeTrimmedString(value: unknown): string | undefined { if (typeof value !== "string") { return undefined; diff --git a/src/plugins/hooks.before-agent-finalize.test.ts b/src/plugins/hooks.before-agent-finalize.test.ts index 7985818c5cb..fad46d1884b 100644 --- a/src/plugins/hooks.before-agent-finalize.test.ts +++ b/src/plugins/hooks.before-agent-finalize.test.ts @@ -132,6 +132,55 @@ describe("before_agent_finalize hook runner", () => { }); }); + it("preserves multiple valid retry candidates for budget evaluation", async () => { + const runner = createHookRunner( + createMockPluginRegistry([ + { + hookName: "before_agent_finalize", + handler: vi.fn().mockResolvedValue({ + action: "revise", + reason: "retry generated artifacts", + retry: { + instruction: "regenerate artifacts", + idempotencyKey: "artifacts", + maxAttempts: 1, + }, + }), + }, + { + hookName: "before_agent_finalize", + handler: vi.fn().mockResolvedValue({ + action: "revise", + reason: "retry focused tests", + retry: { + instruction: "rerun focused tests", + idempotencyKey: "tests", + maxAttempts: 1, + }, + }), + }, + ]), + ); + + const result = await runner.runBeforeAgentFinalize(EVENT, TEST_PLUGIN_AGENT_CTX); + + expect(result).toEqual({ + action: "revise", + reason: "retry generated artifacts\n\nretry focused tests", + retry: { + instruction: "regenerate artifacts", + idempotencyKey: "artifacts", + maxAttempts: 1, + }, + }); + expect(Object.getOwnPropertyDescriptor(result, "retryCandidates")?.enumerable).toBe(false); + expect( + (Object.getOwnPropertyDescriptor(result, "retryCandidates")?.value as unknown[])?.map( + (retry) => (retry as { idempotencyKey?: string }).idempotencyKey, + ), + ).toEqual(["artifacts", "tests"]); + }); + it("lets finalize override earlier revise decisions", async () => { const runner = createHookRunner( createMockPluginRegistry([ diff --git a/src/plugins/hooks.ts b/src/plugins/hooks.ts index 5e1a092363e..fe3035e08ab 100644 --- a/src/plugins/hooks.ts +++ b/src/plugins/hooks.ts @@ -154,6 +154,11 @@ export type HookRunnerLogger = { export type HookFailurePolicy = "fail-open" | "fail-closed"; +type BeforeAgentFinalizeRetry = NonNullable; +type BeforeAgentFinalizeResultWithRetryCandidates = PluginHookBeforeAgentFinalizeResult & { + retryCandidates?: BeforeAgentFinalizeRetry[]; +}; + export type HookRunnerOptions = { logger?: HookRunnerLogger; /** If true, errors in hooks will be caught and logged instead of thrown */ @@ -321,7 +326,7 @@ export function createHookRunner( ): PluginHookBeforeAgentFinalizeResult => { const normalizeRetry = ( retry: PluginHookBeforeAgentFinalizeResult["retry"] | undefined, - ): PluginHookBeforeAgentFinalizeResult["retry"] | undefined => { + ): BeforeAgentFinalizeRetry | undefined => { const instruction = typeof retry?.instruction === "string" ? retry.instruction.trim() : ""; if (!instruction) { return undefined; @@ -331,6 +336,36 @@ export function createHookRunner( instruction, }; }; + const readRetryCandidates = ( + result: PluginHookBeforeAgentFinalizeResult | undefined, + ): BeforeAgentFinalizeRetry[] => { + if (!result || result.action !== "revise") { + return []; + } + const candidateList = (result as BeforeAgentFinalizeResultWithRetryCandidates) + .retryCandidates; + if (Array.isArray(candidateList) && candidateList.length > 0) { + return candidateList + .map((retry) => normalizeRetry(retry)) + .filter((retry): retry is BeforeAgentFinalizeRetry => retry !== undefined); + } + const retry = normalizeRetry(result.retry); + return retry ? [retry] : []; + }; + const attachRetryCandidates = ( + result: PluginHookBeforeAgentFinalizeResult, + candidates: BeforeAgentFinalizeRetry[], + ): PluginHookBeforeAgentFinalizeResult => { + if (result.action !== "revise" || candidates.length <= 1) { + return result; + } + Object.defineProperty(result, "retryCandidates", { + configurable: true, + enumerable: false, + value: candidates, + }); + return result; + }; if (acc?.action === "finalize") { return acc; } @@ -338,15 +373,19 @@ export function createHookRunner( return { action: "finalize", reason: next.reason }; } if (acc?.action === "revise" && next.action === "revise") { - const retry = normalizeRetry(acc.retry) ?? normalizeRetry(next.retry); - return { - action: "revise", - reason: concatOptionalTextSegments({ - left: acc.reason, - right: next.reason, - }), - ...(retry ? { retry } : {}), - }; + const retryCandidates = [...readRetryCandidates(acc), ...readRetryCandidates(next)]; + const retry = retryCandidates[0]; + return attachRetryCandidates( + { + action: "revise", + reason: concatOptionalTextSegments({ + left: acc.reason, + right: next.reason, + }), + ...(retry ? { retry } : {}), + }, + retryCandidates, + ); } if (acc?.action === "revise") { return acc;