From 85a9aabc56b8cbb34d8d1e799aadfde674aca79d Mon Sep 17 00:00:00 2001 From: Eva Date: Fri, 1 May 2026 21:29:29 +0700 Subject: [PATCH] fix: harden cleanup and finalize retry keys --- .../harness/lifecycle-hook-helpers.test.ts | 48 +++++++++++++++++ src/agents/harness/lifecycle-hook-helpers.ts | 8 ++- src/plugins/host-hook-cleanup.config.test.ts | 52 +++++++++++++++++++ src/plugins/host-hook-cleanup.ts | 26 ++++++---- 4 files changed, 124 insertions(+), 10 deletions(-) create mode 100644 src/plugins/host-hook-cleanup.config.test.ts diff --git a/src/agents/harness/lifecycle-hook-helpers.test.ts b/src/agents/harness/lifecycle-hook-helpers.test.ts index 210130747cd..6038126b396 100644 --- a/src/agents/harness/lifecycle-hook-helpers.test.ts +++ b/src/agents/harness/lifecycle-hook-helpers.test.ts @@ -245,4 +245,52 @@ describe("agent harness lifecycle hook helpers", () => { }), ).resolves.toEqual({ action: "continue" }); }); + + it("does not collide fallback retry keys for long instructions with shared prefixes", async () => { + const sharedPrefix = "x".repeat(180); + const firstInstruction = `${sharedPrefix} first`; + const secondInstruction = `${sharedPrefix} second`; + const hookRunner = { + hasHooks: () => true, + runBeforeAgentFinalize: vi + .fn() + .mockResolvedValueOnce({ + action: "revise", + retry: { + instruction: firstInstruction, + idempotencyKey: { invalid: true }, + maxAttempts: 1, + }, + }) + .mockResolvedValueOnce({ + action: "revise", + retry: { + instruction: secondInstruction, + idempotencyKey: { invalid: true }, + maxAttempts: 1, + }, + }), + }; + + await expect( + runAgentHarnessBeforeAgentFinalizeHook({ + event: EVENT, + ctx: { runId: "run-1", sessionKey: "agent:main:session-1" }, + hookRunner: hookRunner as never, + }), + ).resolves.toEqual({ + action: "revise", + reason: firstInstruction, + }); + await expect( + runAgentHarnessBeforeAgentFinalizeHook({ + event: EVENT, + ctx: { runId: "run-1", sessionKey: "agent:main:session-1" }, + hookRunner: hookRunner as never, + }), + ).resolves.toEqual({ + action: "revise", + reason: secondInstruction, + }); + }); }); diff --git a/src/agents/harness/lifecycle-hook-helpers.ts b/src/agents/harness/lifecycle-hook-helpers.ts index 5f39e70bf26..18b48f12e0f 100644 --- a/src/agents/harness/lifecycle-hook-helpers.ts +++ b/src/agents/harness/lifecycle-hook-helpers.ts @@ -1,3 +1,4 @@ +import { createHash } from "node:crypto"; import { createSubsystemLogger } from "../../logging/subsystem.js"; import { getGlobalHookRunner } from "../../plugins/hook-runner-global.js"; import type { @@ -46,6 +47,10 @@ function pruneFinalizeRetryBudget(budget: FinalizeRetryBudget): void { } } +function buildFinalizeRetryInstructionKey(instruction: string): string { + return `instruction:${createHash("sha256").update(instruction).digest("hex")}`; +} + export function clearAgentHarnessFinalizeRetryBudget(params?: { runId?: string }): void { const budget = getFinalizeRetryBudget(); if (!params?.runId) { @@ -149,7 +154,8 @@ function normalizeBeforeAgentFinalizeResult( : 1; const retryRunId = event?.runId ?? event?.sessionId ?? "unknown-run"; const retryKey = - normalizeTrimmedString(result.retry?.idempotencyKey) || retryInstruction.slice(0, 160); + normalizeTrimmedString(result.retry?.idempotencyKey) || + buildFinalizeRetryInstructionKey(retryInstruction); const budget = getFinalizeRetryBudget(); const runBudget = budget.get(retryRunId) ?? new Map(); const nextCount = (runBudget.get(retryKey) ?? 0) + 1; diff --git a/src/plugins/host-hook-cleanup.config.test.ts b/src/plugins/host-hook-cleanup.config.test.ts new file mode 100644 index 00000000000..024bd277f5d --- /dev/null +++ b/src/plugins/host-hook-cleanup.config.test.ts @@ -0,0 +1,52 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { runPluginHostCleanup } from "./host-hook-cleanup.js"; +import { createEmptyPluginRegistry } from "./registry-empty.js"; + +const mocks = vi.hoisted(() => ({ + getRuntimeConfig: vi.fn(), +})); + +vi.mock("../config/config.js", () => ({ + getRuntimeConfig: mocks.getRuntimeConfig, +})); + +describe("plugin host cleanup config fallback", () => { + afterEach(() => { + mocks.getRuntimeConfig.mockReset(); + }); + + it("records session store config failures while continuing runtime cleanup", async () => { + const registry = createEmptyPluginRegistry(); + const cleanup = vi.fn(); + registry.runtimeLifecycles.push({ + pluginId: "cleanup-plugin", + pluginName: "Cleanup Plugin", + lifecycle: { + id: "runtime-cleanup", + cleanup, + }, + }); + mocks.getRuntimeConfig.mockImplementation(() => { + throw new Error("invalid config"); + }); + + const result = await runPluginHostCleanup({ + registry, + pluginId: "cleanup-plugin", + reason: "disable", + }); + + expect(cleanup).toHaveBeenCalledWith( + expect.objectContaining({ + reason: "disable", + }), + ); + expect(result.cleanupCount).toBe(1); + expect(result.failures).toEqual([ + expect.objectContaining({ + pluginId: "cleanup-plugin", + hookId: "session-store", + }), + ]); + }); +}); diff --git a/src/plugins/host-hook-cleanup.ts b/src/plugins/host-hook-cleanup.ts index 913fa5e27d5..b558c61eece 100644 --- a/src/plugins/host-hook-cleanup.ts +++ b/src/plugins/host-hook-cleanup.ts @@ -115,16 +115,24 @@ export async function runPluginHostCleanup(params: { runId?: string; preserveSchedulerJobIds?: ReadonlySet; }): Promise { - const persistentCleanupCount = - params.reason === "restart" - ? 0 - : await clearPluginOwnedSessionStores({ - cfg: params.cfg ?? getRuntimeConfig(), - pluginId: params.pluginId, - sessionKey: params.sessionKey, - }); - const registry = params.registry; const failures: PluginHostCleanupFailure[] = []; + let persistentCleanupCount = 0; + if (params.reason !== "restart") { + try { + persistentCleanupCount = await clearPluginOwnedSessionStores({ + cfg: params.cfg ?? getRuntimeConfig(), + pluginId: params.pluginId, + sessionKey: params.sessionKey, + }); + } catch (error) { + failures.push({ + pluginId: params.pluginId ?? "plugin-host", + hookId: "session-store", + error, + }); + } + } + const registry = params.registry; let cleanupCount = persistentCleanupCount; if (registry) { for (const registration of registry.sessionExtensions ?? []) {