From 6bfd447fc9ee50c0455ea43014c0676cdc5cb48f Mon Sep 17 00:00:00 2001 From: Chandika Jayasundara Date: Wed, 18 Feb 2026 22:16:50 +0000 Subject: [PATCH 1/2] feat: make llm_input/llm_output modifying hooks for middleware patterns Convert llm_input and llm_output from void (observational) hooks to modifying hooks that can transform data flowing to/from LLM providers. This enables an entire class of plugins that need to intercept and modify LLM traffic: - PII redaction (strip personal info before provider, rehydrate after) - Content filtering and guardrails - Translation middleware - Token optimization - Provider-agnostic safety layers Changes: - types.ts: Add PluginHookLlmInputResult and PluginHookLlmOutputResult - types.ts: Update PluginHookHandlerMap to allow returning results - hooks.ts: Switch runLlmInput/runLlmOutput from runVoidHook to runModifyingHook with merge functions - attempt.ts: Await hook results and apply prompt/assistantTexts modifications when returned - wired-hooks-llm.test.ts: Add tests for modifying behavior and backward compatibility (void returns still work) Backward compatible: existing plugins that return void from llm_input/llm_output continue to work unchanged. Closes #20416 --- src/agents/pi-embedded-runner/run/attempt.ts | 37 ++++++--- src/plugins/hooks.ts | 43 ++++++++-- src/plugins/types.ts | 21 ++++- src/plugins/wired-hooks-llm.test.ts | 84 ++++++++++++++++++++ 4 files changed, 162 insertions(+), 23 deletions(-) diff --git a/src/agents/pi-embedded-runner/run/attempt.ts b/src/agents/pi-embedded-runner/run/attempt.ts index 364448eb89d..f30dd01fa26 100644 --- a/src/agents/pi-embedded-runner/run/attempt.ts +++ b/src/agents/pi-embedded-runner/run/attempt.ts @@ -779,7 +779,7 @@ export async function runEmbeddedAttempt( }); const { - assistantTexts, + assistantTexts: rawAssistantTexts, toolMetas, unsubscribe, waitForCompactionRetry, @@ -792,6 +792,7 @@ export async function runEmbeddedAttempt( getUsageTotals, getCompactionCount, } = subscription; + let assistantTexts = rawAssistantTexts; const queueHandle: EmbeddedPiQueueHandle = { queueMessage: async (text: string) => { @@ -1014,9 +1015,10 @@ export async function runEmbeddedAttempt( ); } + // Run llm_input hook — plugins may modify prompt/systemPrompt if (hookRunner?.hasHooks("llm_input")) { - hookRunner - .runLlmInput( + try { + const llmInputResult = await hookRunner.runLlmInput( { runId: params.runId, sessionId: params.sessionId, @@ -1034,10 +1036,14 @@ export async function runEmbeddedAttempt( workspaceDir: params.workspaceDir, messageProvider: params.messageProvider ?? undefined, }, - ) - .catch((err) => { - log.warn(`llm_input hook failed: ${String(err)}`); - }); + ); + // Apply any modifications from the hook + if (llmInputResult?.prompt) { + effectivePrompt = llmInputResult.prompt; + } + } catch (err) { + log.warn(`llm_input hook failed: ${String(err)}`); + } } // Only pass images option if there are actually images to pass @@ -1208,9 +1214,10 @@ export async function runEmbeddedAttempt( ) .map((entry) => ({ toolName: entry.toolName, meta: entry.meta })); + // Run llm_output hook — plugins may modify assistantTexts if (hookRunner?.hasHooks("llm_output")) { - hookRunner - .runLlmOutput( + try { + const llmOutputResult = await hookRunner.runLlmOutput( { runId: params.runId, sessionId: params.sessionId, @@ -1227,10 +1234,14 @@ export async function runEmbeddedAttempt( workspaceDir: params.workspaceDir, messageProvider: params.messageProvider ?? undefined, }, - ) - .catch((err) => { - log.warn(`llm_output hook failed: ${String(err)}`); - }); + ); + // Apply any modifications from the hook + if (llmOutputResult?.assistantTexts) { + assistantTexts = llmOutputResult.assistantTexts; + } + } catch (err) { + log.warn(`llm_output hook failed: ${String(err)}`); + } } return { diff --git a/src/plugins/hooks.ts b/src/plugins/hooks.ts index 19b10404262..636929a2ac2 100644 --- a/src/plugins/hooks.ts +++ b/src/plugins/hooks.ts @@ -19,7 +19,9 @@ import type { PluginHookBeforePromptBuildResult, PluginHookBeforeCompactionEvent, PluginHookLlmInputEvent, + PluginHookLlmInputResult, PluginHookLlmOutputEvent, + PluginHookLlmOutputResult, PluginHookBeforeResetEvent, PluginHookBeforeToolCallEvent, PluginHookBeforeToolCallResult, @@ -54,7 +56,9 @@ export type { PluginHookBeforePromptBuildEvent, PluginHookBeforePromptBuildResult, PluginHookLlmInputEvent, + PluginHookLlmInputResult, PluginHookLlmOutputEvent, + PluginHookLlmOutputResult, PluginHookAgentEndEvent, PluginHookBeforeCompactionEvent, PluginHookBeforeResetEvent, @@ -278,20 +282,43 @@ export function createHookRunner(registry: PluginRegistry, options: HookRunnerOp /** * Run llm_input hook. - * Allows plugins to observe the exact input payload sent to the LLM. - * Runs in parallel (fire-and-forget). + * Allows plugins to observe or modify the input payload sent to the LLM. + * Plugins can return `{ prompt, systemPrompt }` to transform the input, + * or return void/undefined for observation-only (backward compatible). */ - async function runLlmInput(event: PluginHookLlmInputEvent, ctx: PluginHookAgentContext) { - return runVoidHook("llm_input", event, ctx); + async function runLlmInput( + event: PluginHookLlmInputEvent, + ctx: PluginHookAgentContext, + ): Promise { + return runModifyingHook<"llm_input", PluginHookLlmInputResult>( + "llm_input", + event, + ctx, + (acc, next) => ({ + prompt: next.prompt ?? acc?.prompt, + systemPrompt: next.systemPrompt ?? acc?.systemPrompt, + }), + ); } /** * Run llm_output hook. - * Allows plugins to observe the exact output payload returned by the LLM. - * Runs in parallel (fire-and-forget). + * Allows plugins to observe or modify the output payload returned by the LLM. + * Plugins can return `{ assistantTexts }` to transform the output, + * or return void/undefined for observation-only (backward compatible). */ - async function runLlmOutput(event: PluginHookLlmOutputEvent, ctx: PluginHookAgentContext) { - return runVoidHook("llm_output", event, ctx); + async function runLlmOutput( + event: PluginHookLlmOutputEvent, + ctx: PluginHookAgentContext, + ): Promise { + return runModifyingHook<"llm_output", PluginHookLlmOutputResult>( + "llm_output", + event, + ctx, + (acc, next) => ({ + assistantTexts: next.assistantTexts ?? acc?.assistantTexts, + }), + ); } /** diff --git a/src/plugins/types.ts b/src/plugins/types.ts index fc54fdece8a..68686753e7a 100644 --- a/src/plugins/types.ts +++ b/src/plugins/types.ts @@ -373,6 +373,14 @@ export type PluginHookLlmInputEvent = { imagesCount: number; }; +// llm_input hook result (when used as a modifying hook) +export type PluginHookLlmInputResult = { + /** Modified prompt text. If set, replaces the original prompt. */ + prompt?: string; + /** Modified system prompt. If set, replaces the original system prompt. */ + systemPrompt?: string; +}; + // llm_output hook export type PluginHookLlmOutputEvent = { runId: string; @@ -390,6 +398,12 @@ export type PluginHookLlmOutputEvent = { }; }; +// llm_output hook result (when used as a modifying hook) +export type PluginHookLlmOutputResult = { + /** Modified assistant response texts. If set, replaces the originals. */ + assistantTexts?: string[]; +}; + // agent_end hook export type PluginHookAgentEndEvent = { messages: unknown[]; @@ -579,11 +593,14 @@ export type PluginHookHandlerMap = { event: PluginHookBeforeAgentStartEvent, ctx: PluginHookAgentContext, ) => Promise | PluginHookBeforeAgentStartResult | void; - llm_input: (event: PluginHookLlmInputEvent, ctx: PluginHookAgentContext) => Promise | void; + llm_input: ( + event: PluginHookLlmInputEvent, + ctx: PluginHookAgentContext, + ) => Promise | PluginHookLlmInputResult | void; llm_output: ( event: PluginHookLlmOutputEvent, ctx: PluginHookAgentContext, - ) => Promise | void; + ) => Promise | PluginHookLlmOutputResult | void; agent_end: (event: PluginHookAgentEndEvent, ctx: PluginHookAgentContext) => Promise | void; before_compaction: ( event: PluginHookBeforeCompactionEvent, diff --git a/src/plugins/wired-hooks-llm.test.ts b/src/plugins/wired-hooks-llm.test.ts index a20a40aa84c..63b69385e42 100644 --- a/src/plugins/wired-hooks-llm.test.ts +++ b/src/plugins/wired-hooks-llm.test.ts @@ -69,4 +69,88 @@ describe("llm hook runner methods", () => { expect(runner.hasHooks("llm_input")).toBe(true); expect(runner.hasHooks("llm_output")).toBe(false); }); + + it("runLlmInput returns modified prompt from hook", async () => { + const handler = vi.fn().mockResolvedValue({ prompt: "redacted prompt" }); + const registry = createMockPluginRegistry([{ hookName: "llm_input", handler }]); + const runner = createHookRunner(registry); + + const result = await runner.runLlmInput( + { + runId: "run-1", + sessionId: "session-1", + provider: "openai", + model: "gpt-5", + systemPrompt: "be helpful", + prompt: "original prompt", + historyMessages: [], + imagesCount: 0, + }, + { agentId: "main", sessionId: "session-1" }, + ); + + expect(result).toEqual(expect.objectContaining({ prompt: "redacted prompt" })); + }); + + it("runLlmInput returns undefined when hook returns void (backward compat)", async () => { + const handler = vi.fn(); // returns undefined + const registry = createMockPluginRegistry([{ hookName: "llm_input", handler }]); + const runner = createHookRunner(registry); + + const result = await runner.runLlmInput( + { + runId: "run-1", + sessionId: "session-1", + provider: "openai", + model: "gpt-5", + prompt: "hello", + historyMessages: [], + imagesCount: 0, + }, + { agentId: "main", sessionId: "session-1" }, + ); + + expect(result).toBeUndefined(); + }); + + it("runLlmOutput returns modified assistantTexts from hook", async () => { + const handler = vi.fn().mockResolvedValue({ assistantTexts: ["rehydrated response"] }); + const registry = createMockPluginRegistry([{ hookName: "llm_output", handler }]); + const runner = createHookRunner(registry); + + const result = await runner.runLlmOutput( + { + runId: "run-1", + sessionId: "session-1", + provider: "openai", + model: "gpt-5", + assistantTexts: ["raw «PERSON_001» response"], + lastAssistant: { role: "assistant", content: "raw" }, + usage: { input: 10, output: 20, total: 30 }, + }, + { agentId: "main", sessionId: "session-1" }, + ); + + expect(result).toEqual(expect.objectContaining({ assistantTexts: ["rehydrated response"] })); + }); + + it("runLlmOutput returns undefined when hook returns void (backward compat)", async () => { + const handler = vi.fn(); // returns undefined + const registry = createMockPluginRegistry([{ hookName: "llm_output", handler }]); + const runner = createHookRunner(registry); + + const result = await runner.runLlmOutput( + { + runId: "run-1", + sessionId: "session-1", + provider: "openai", + model: "gpt-5", + assistantTexts: ["hi"], + usage: { input: 10, output: 20, total: 30 }, + }, + { agentId: "main", sessionId: "session-1" }, + ); + + expect(result).toBeUndefined(); + }); }); From 32720e044ff80329f23ef54ed7e47d157df9c81e Mon Sep 17 00:00:00 2001 From: Chandika Jayasundara Date: Wed, 18 Feb 2026 23:03:41 +0000 Subject: [PATCH 2/2] fix: remove systemPrompt from llm_input result type (not yet plumbed) Address greptile review: systemPrompt was declared in PluginHookLlmInputResult and merged in the hook runner, but never applied in attempt.ts. Since the system prompt is finalised before the llm_input hook runs, applying it requires restructuring that's out of scope for this PR. Removes the misleading field to avoid a silent no-op for plugin authors. Will be re-added when late-stage system prompt overrides are supported. --- src/agents/pi-embedded-runner/run/attempt.ts | 2 +- src/plugins/hooks.ts | 3 +-- src/plugins/types.ts | 5 +++-- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/agents/pi-embedded-runner/run/attempt.ts b/src/agents/pi-embedded-runner/run/attempt.ts index f30dd01fa26..10d4ab675b3 100644 --- a/src/agents/pi-embedded-runner/run/attempt.ts +++ b/src/agents/pi-embedded-runner/run/attempt.ts @@ -1015,7 +1015,7 @@ export async function runEmbeddedAttempt( ); } - // Run llm_input hook — plugins may modify prompt/systemPrompt + // Run llm_input hook — plugins may modify the user prompt if (hookRunner?.hasHooks("llm_input")) { try { const llmInputResult = await hookRunner.runLlmInput( diff --git a/src/plugins/hooks.ts b/src/plugins/hooks.ts index 636929a2ac2..d27199a1c92 100644 --- a/src/plugins/hooks.ts +++ b/src/plugins/hooks.ts @@ -283,7 +283,7 @@ export function createHookRunner(registry: PluginRegistry, options: HookRunnerOp /** * Run llm_input hook. * Allows plugins to observe or modify the input payload sent to the LLM. - * Plugins can return `{ prompt, systemPrompt }` to transform the input, + * Plugins can return `{ prompt }` to transform the input, * or return void/undefined for observation-only (backward compatible). */ async function runLlmInput( @@ -296,7 +296,6 @@ export function createHookRunner(registry: PluginRegistry, options: HookRunnerOp ctx, (acc, next) => ({ prompt: next.prompt ?? acc?.prompt, - systemPrompt: next.systemPrompt ?? acc?.systemPrompt, }), ); } diff --git a/src/plugins/types.ts b/src/plugins/types.ts index 68686753e7a..061f39dabf2 100644 --- a/src/plugins/types.ts +++ b/src/plugins/types.ts @@ -377,8 +377,9 @@ export type PluginHookLlmInputEvent = { export type PluginHookLlmInputResult = { /** Modified prompt text. If set, replaces the original prompt. */ prompt?: string; - /** Modified system prompt. If set, replaces the original system prompt. */ - systemPrompt?: string; + // Note: systemPrompt modification is not yet supported — the system prompt + // is finalised earlier in the pipeline. Will be added when late-stage + // system prompt overrides are plumbed through. }; // llm_output hook