fix: honor later finalize retry candidates

2026-05-06 12:00:44 +00:00 · 2026-05-02 00:58:29 +07:00
parent 5412cf4efe
commit 956bd2f94b
4 changed files with 198 additions and 36 deletions
--- 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,
--- 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<string, number>();
-      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<string, number>();
+        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<PluginHookBeforeAgentFinalizeResult["retry"]>[] {
+  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<PluginHookBeforeAgentFinalizeResult["retry"]> {
+  return Boolean(value) && typeof value === "object" && !Array.isArray(value);
+}
+
 function normalizeTrimmedString(value: unknown): string | undefined {
  if (typeof value !== "string") {
    return undefined;
--- 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([
--- 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<PluginHookBeforeAgentFinalizeResult["retry"]>;
+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;