refactor(plugins): unify synchronous message hooks (#127797)

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

src/infra/unhandled-rejections.plugin-sync-hooks.test.ts

@@ -28,10 +28,18 @@ function runSyncHook(params: {
   hookName: SyncHookName;
   runner: ReturnType<typeof createHookRunner>;
   message: AgentMessage;
+}) {
+  return runSyncHookEvent({ ...params, event: { message: params.message } });
+}
+
+function runSyncHookEvent(params: {
+  hookName: SyncHookName;
+  runner: ReturnType<typeof createHookRunner>;
+  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("-");

src/plugins/hooks.ts

@@ -259,7 +259,8 @@ type SyncHookName = "tool_result_persist" | "before_message_write";
 type SyncHookHandler<K extends SyncHookName> = NonNullable<PluginHookRegistration<K>["handler"]>;
 type SyncHookEvent<K extends SyncHookName> = Parameters<SyncHookHandler<K>>[0];
 type SyncHookContext<K extends SyncHookName> = Parameters<SyncHookHandler<K>>[1];
-type SyncHookResult<K extends SyncHookName> = ReturnType<SyncHookHandler<K>>;
+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 = <K extends SyncHookName>(
+  const runSyncMessageHookStep = <K extends SyncHookName>(
     hook: PluginHookRegistration<K>,
-    event: SyncHookEvent<K>,
+    hookName: K,
+    event: SyncHookEvent<K> & { message: SyncHookMessage },
+    message: SyncHookMessage,
     ctx: SyncHookContext<K>,
-  ): SyncHookResult<K> | undefined => {
-    const handler = hook.handler as SyncHookHandler<K>;
-    const out = handler(event, ctx) as SyncHookResult<K> | PromiseLike<unknown>;
-    if (!isPromiseLike(out)) {
-      return out;
+  ): SyncMessageHookStepResult | undefined => {
+    try {
+      const handler = hook.handler as (
+        event: SyncHookEvent<K>,
+        ctx: SyncHookContext<K>,
+      ) => { message?: SyncHookMessage; block?: boolean } | PromiseLike<unknown> | 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 = <K extends SyncHookName>(
+    hookName: K,
+    event: SyncHookEvent<K> & { message: SyncHookMessage },
+    ctx: SyncHookContext<K>,
+  ): { 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;
   }
 
   // =========================================================================