From 4d696fbe0d2bf53b8f590743a325667ea94f3f41 Mon Sep 17 00:00:00 2001 From: Vincent Koc Date: Sat, 22 Aug 2026 22:15:31 -0700 Subject: [PATCH] refactor(plugins): unify synchronous message hooks (#127797) --- config/assertion-safety-baseline.txt | 2 +- ...ndled-rejections.plugin-sync-hooks.test.ts | 162 +++++++++++++++++- src/plugins/hooks.ts | 143 +++++++--------- 3 files changed, 227 insertions(+), 80 deletions(-) diff --git a/config/assertion-safety-baseline.txt b/config/assertion-safety-baseline.txt index 65399f950c63..c5ef0bb1cb0d 100644 --- a/config/assertion-safety-baseline.txt +++ b/config/assertion-safety-baseline.txt @@ -3577,7 +3577,7 @@ src/plugins/hook-decision-types.ts 1 src/plugins/hook-isolation.ts 3 src/plugins/hook-runner-global-state.ts 1 src/plugins/hook-types.ts 2 -src/plugins/hooks.ts 19 +src/plugins/hooks.ts 16 src/plugins/host-hook-attachments.ts 2 src/plugins/host-hook-json.ts 1 src/plugins/host-hook-runtime.ts 1 diff --git a/src/infra/unhandled-rejections.plugin-sync-hooks.test.ts b/src/infra/unhandled-rejections.plugin-sync-hooks.test.ts index 2958d19aba62..d032c72dd43e 100644 --- a/src/infra/unhandled-rejections.plugin-sync-hooks.test.ts +++ b/src/infra/unhandled-rejections.plugin-sync-hooks.test.ts @@ -28,10 +28,18 @@ function runSyncHook(params: { hookName: SyncHookName; runner: ReturnType; message: AgentMessage; +}) { + return runSyncHookEvent({ ...params, event: { message: params.message } }); +} + +function runSyncHookEvent(params: { + hookName: SyncHookName; + runner: ReturnType; + event: { message: AgentMessage }; }) { return params.hookName === "tool_result_persist" - ? params.runner.runToolResultPersist({ message: params.message }, {}) - : params.runner.runBeforeMessageWrite({ message: params.message }, {}); + ? params.runner.runToolResultPersist(params.event, {}) + : params.runner.runBeforeMessageWrite(params.event, {}); } describe("sync-only plugin hooks", () => { @@ -103,6 +111,156 @@ describe("sync-only plugin hooks", () => { expect(logger.error).not.toHaveBeenCalled(); }); + it.each(syncHookNames)("composes synchronous %s results after fail-open errors", (hookName) => { + const logger = createLogger(); + const originalMessage = createToolResultMessage("original"); + const replacementMessage = createToolResultMessage("replacement"); + const finalMessage = createToolResultMessage("final"); + const observer = vi.fn((event: unknown) => { + expect((event as { message: AgentMessage }).message).toBe(replacementMessage); + return { message: finalMessage }; + }); + const runner = createHookRunner( + createMockPluginRegistry([ + { + hookName, + pluginId: "replacement-hook", + priority: 30, + handler: () => ({ message: replacementMessage }), + }, + { + hookName, + pluginId: "failed-hook", + priority: 20, + handler: () => { + throw new Error("sync-hook-failure"); + }, + }, + { hookName, pluginId: "observer-hook", priority: 10, handler: observer }, + ]), + { logger }, + ); + + expect(runSyncHook({ hookName, runner, message: originalMessage })).toEqual({ + message: finalMessage, + }); + expect(observer).toHaveBeenCalledOnce(); + expect(logger.error).toHaveBeenCalledWith( + `[hooks] ${hookName} handler from failed-hook failed: Error: sync-hook-failure`, + ); + expect(logger.warn).not.toHaveBeenCalled(); + }); + + it.each([ + ["tool_result_persist", "event"], + ["before_message_write", "event"], + ["tool_result_persist", "message"], + ["before_message_write", "message"], + ["before_message_write", "block"], + ] as const)( + "contains %s %s getter failures and continues composition", + (hookName, getterName) => { + const logger = createLogger(); + const originalMessage = createToolResultMessage("original"); + const finalMessage = createToolResultMessage("final"); + const event = { message: originalMessage }; + if (getterName === "event") { + let reads = 0; + Object.defineProperty(event, "hostile", { + enumerable: true, + get: () => { + if (reads++ === 0) { + throw new Error("event-getter-failure"); + } + return true; + }, + }); + } + const failingHandler = vi.fn(() => { + const result: { message?: AgentMessage; block?: boolean } = {}; + Object.defineProperty(result, getterName, { + enumerable: true, + get: () => { + throw new Error(`${getterName}-getter-failure`); + }, + }); + return result; + }); + const runner = createHookRunner( + createMockPluginRegistry([ + { hookName, pluginId: "getter", priority: 20, handler: failingHandler }, + { + hookName, + pluginId: "observer", + priority: 10, + handler: () => ({ message: finalMessage }), + }, + ]), + { logger }, + ); + + expect(runSyncHookEvent({ hookName, runner, event })).toEqual({ message: finalMessage }); + expect(logger.error).toHaveBeenCalledWith( + `[hooks] ${hookName} handler from getter failed: Error: ${getterName}-getter-failure`, + ); + expect(failingHandler).toHaveBeenCalledTimes(getterName === "event" ? 0 : 1); + }, + ); + + it("does not read block from tool_result_persist results", () => { + const replacementMessage = createToolResultMessage("replacement"); + let blockReads = 0; + const result = { message: replacementMessage }; + Object.defineProperty(result, "block", { + get: () => { + blockReads += 1; + throw new Error("tool-block-getter-failure"); + }, + }); + const runner = createHookRunner( + createMockPluginRegistry([ + { hookName: "tool_result_persist", pluginId: "tool", handler: () => result }, + ]), + ); + + expect( + runner.runToolResultPersist({ message: createToolResultMessage("original") }, {}), + ).toEqual({ message: replacementMessage }); + expect(blockReads).toBe(0); + }); + + it.each(syncHookNames)( + "fails closed on synchronous %s invocation errors and skips later handlers", + (hookName) => { + const cause = new Error("sync-hook-failure"); + const laterHandler = vi.fn(); + const runner = createHookRunner( + createMockPluginRegistry([ + { + hookName, + pluginId: "failed", + priority: 20, + handler: () => { + throw cause; + }, + }, + { hookName, pluginId: "later", priority: 10, handler: laterHandler }, + ]), + { failurePolicyByHook: { [hookName]: "fail-closed" } }, + ); + + expect(() => + runSyncHook({ hookName, runner, message: createToolResultMessage("original") }), + ).toThrow( + expect.objectContaining({ + message: `[hooks] ${hookName} handler from failed failed: Error: sync-hook-failure`, + cause, + }), + ); + expect(laterHandler).not.toHaveBeenCalled(); + }, + ); + it("preserves synchronous secret redaction and subsequent handler composition", () => { const logger = createLogger(); const secret = ["fixture", "secret"].join("-"); diff --git a/src/plugins/hooks.ts b/src/plugins/hooks.ts index 5f92c902f4d4..ff0effeb91c5 100644 --- a/src/plugins/hooks.ts +++ b/src/plugins/hooks.ts @@ -259,7 +259,8 @@ type SyncHookName = "tool_result_persist" | "before_message_write"; type SyncHookHandler = NonNullable["handler"]>; type SyncHookEvent = Parameters>[0]; type SyncHookContext = Parameters>[1]; -type SyncHookResult = ReturnType>; +type SyncHookMessage = PluginHookToolResultPersistEvent["message"]; +type SyncMessageHookStepResult = { message?: SyncHookMessage; block?: true }; /** * Get hooks for a specific hook name, sorted by priority (higher first). @@ -721,27 +722,70 @@ export function createHookRunner( } }; - const runSyncHookHandler = ( + const runSyncMessageHookStep = ( hook: PluginHookRegistration, - event: SyncHookEvent, + hookName: K, + event: SyncHookEvent & { message: SyncHookMessage }, + message: SyncHookMessage, ctx: SyncHookContext, - ): SyncHookResult | undefined => { - const handler = hook.handler as SyncHookHandler; - const out = handler(event, ctx) as SyncHookResult | PromiseLike; - if (!isPromiseLike(out)) { - return out; + ): SyncMessageHookStepResult | undefined => { + try { + const handler = hook.handler as ( + event: SyncHookEvent, + ctx: SyncHookContext, + ) => { message?: SyncHookMessage; block?: boolean } | PromiseLike | void; + const result = handler({ ...event, message }, ctx); + if (isPromiseLike(result)) { + // Sync-only hooks ignore async results; observe rejections so the global fatal handler cannot crash. + void Promise.resolve(result).catch(() => undefined); + const msg = + `[hooks] ${hookName} handler from ${hook.pluginId} returned a Promise; ` + + `this hook is synchronous and the result was ignored.`; + if (!shouldCatchHookErrors(hookName)) { + throw new Error(msg); + } + logger?.warn?.(msg); + return undefined; + } + if (!result) { + return undefined; + } + if (hookName === "before_message_write" && result.block) { + return { block: true }; + } + const nextMessage = result.message; + return nextMessage ? { message: nextMessage } : undefined; + } catch (err) { + const msg = `[hooks] ${hookName} handler from ${hook.pluginId} failed: ${String(err)}`; + if (shouldCatchHookErrors(hookName)) { + logger?.error(msg); + return undefined; + } + throw new Error(msg, { cause: err }); } + }; - // Sync-only hooks ignore async results; observe rejections so the global fatal handler cannot crash. - void Promise.resolve(out).catch(() => undefined); - const msg = - `[hooks] ${hook.hookName} handler from ${hook.pluginId} returned a Promise; ` + - `this hook is synchronous and the result was ignored.`; - if (shouldCatchHookErrors(hook.hookName)) { - logger?.warn?.(msg); + const runSyncMessageHooks = ( + hookName: K, + event: SyncHookEvent & { message: SyncHookMessage }, + ctx: SyncHookContext, + ): { message: SyncHookMessage; block?: true } | undefined => { + const hooks = getHooksForName(registry, hookName); + if (hooks.length === 0) { return undefined; } - throw new Error(msg); + + let current = event.message; + for (const hook of hooks) { + const result = runSyncMessageHookStep(hook, hookName, event, current, ctx); + if (result?.block) { + return { message: current, block: true }; + } + if (result?.message) { + current = result.message; + } + } + return { message: current }; }; /** @@ -1500,31 +1544,8 @@ export function createHookRunner( event: PluginHookToolResultPersistEvent, ctx: PluginHookToolResultPersistContext, ): PluginHookToolResultPersistResult | undefined { - const hooks = getHooksForName(registry, "tool_result_persist"); - if (hooks.length === 0) { - return undefined; - } - - let current = event.message; - - for (const hook of hooks) { - try { - const out = runSyncHookHandler(hook, { ...event, message: current }, ctx); - const next = (out as PluginHookToolResultPersistResult | undefined)?.message; - if (next) { - current = next; - } - } catch (err) { - const msg = `[hooks] tool_result_persist handler from ${hook.pluginId} failed: ${String(err)}`; - if (shouldCatchHookErrors("tool_result_persist")) { - logger?.error(msg); - } else { - throw new Error(msg, { cause: err }); - } - } - } - - return { message: current }; + const result = runSyncMessageHooks("tool_result_persist", event, ctx); + return result ? { message: result.message } : undefined; } // ========================================================================= @@ -1547,43 +1568,11 @@ export function createHookRunner( event: PluginHookBeforeMessageWriteEvent, ctx: { agentId?: string; sessionKey?: string }, ): PluginHookBeforeMessageWriteResult | undefined { - const hooks = getHooksForName(registry, "before_message_write"); - if (hooks.length === 0) { - return undefined; + const result = runSyncMessageHooks("before_message_write", event, ctx); + if (result?.block) { + return { block: true }; } - - let current = event.message; - - for (const hook of hooks) { - try { - const out = runSyncHookHandler(hook, { ...event, message: current }, ctx); - const result = out as PluginHookBeforeMessageWriteResult | undefined; - - // If any handler blocks, return immediately. - if (result?.block) { - return { block: true }; - } - - // If handler provided a modified message, use it for subsequent handlers. - if (result?.message) { - current = result.message; - } - } catch (err) { - const msg = `[hooks] before_message_write handler from ${hook.pluginId} failed: ${String(err)}`; - if (shouldCatchHookErrors("before_message_write")) { - logger?.error(msg); - } else { - throw new Error(msg, { cause: err }); - } - } - } - - // If message was modified by any handler, return it. - if (current !== event.message) { - return { message: current }; - } - - return undefined; + return result && result.message !== event.message ? { message: result.message } : undefined; } // =========================================================================