Merge branch 'main' into mb/beforemodel-context

docs/hooks/reference.md

@@ -197,6 +197,9 @@ request format.
     outgoing request (e.g., changing models or temperature).
   - `hookSpecificOutput.llm_response`: A **Synthetic Response** object. If
     provided, the CLI skips the LLM call entirely and uses this as the response.
+  - `hookSpecificOutput.additionalContext`: (`string`) Text that is **appended**
+    to the end of the user message. If `llm_request` also modifies contents, it
+    will be appended to the modified contents.
   - `decision`: Set to `"deny"` to block the request and abort the turn.
 - **Exit Code 2 (Block Turn)**: Aborts the turn and skips the LLM call. Uses
   `stderr` as the error message.

packages/cli/src/gemini.tsx

@@ -665,8 +665,9 @@ export async function main() {
         const additionalContext = result.getAdditionalContext();
         if (additionalContext) {
           // Prepend context to input (System Context -> Stdin -> Question)
-          const wrappedContext = `<hook_context>${additionalContext}</hook_context>`;
-          input = input ? `${wrappedContext}\n\n${input}` : wrappedContext;
+          input = input
+            ? `${additionalContext}\n\n${input}`
+            : additionalContext;
         }
       }
     }

packages/cli/src/ui/AppContainer.tsx

@@ -498,13 +498,11 @@ export const AppContainer = (props: AppContainerProps) => {
         }
 
         const additionalContext = result.getAdditionalContext();
-        const geminiClient = config.getGeminiClient();
+
         if (additionalContext && geminiClient) {
           await geminiClient.addHistory({
             role: 'user',
-            parts: [
-              { text: `<hook_context>${additionalContext}</hook_context>` },
-            ],
+            parts: [{ text: `\n\n${additionalContext}` }],
           });
         }
       }

packages/core/src/core/client.ts

@@ -922,10 +922,7 @@ export class GeminiClient {
           const additionalContext = hookResult.additionalContext;
           if (additionalContext) {
             const requestArray = Array.isArray(request) ? request : [request];
-            request = [
-              ...requestArray,
-              { text: `<hook_context>${additionalContext}</hook_context>` },
-            ];
+            request = [...requestArray, { text: `\n\n${additionalContext}` }];
           }
         }
       }

packages/core/src/core/coreToolHookTriggers.ts

@@ -220,7 +220,7 @@ export async function executeToolWithHooks(
     // Add additional context from hooks to the tool result
     const additionalContext = afterOutput?.getAdditionalContext();
     if (additionalContext) {
-      const wrappedContext = `\n\n<hook_context>${additionalContext}</hook_context>`;
+      const wrappedContext = `\n\n${additionalContext}`;
       if (typeof toolResult.llmContent === 'string') {
         toolResult.llmContent += wrappedContext;
       } else if (Array.isArray(toolResult.llmContent)) {

packages/core/src/hooks/hookAggregator.test.ts

@@ -423,6 +423,54 @@ describe('HookAggregator', () => {
       const llmRequest = output.hookSpecificOutput?.llm_request;
       expect(llmRequest?.['model']).toBe('model2'); // Later value wins
     });
+
+    it('should aggregate additionalContext for BeforeModel hooks', () => {
+      const results: HookExecutionResult[] = [
+        {
+          hookConfig: {
+            type: HookType.Command,
+            command: 'h1',
+            timeout: 30000,
+          },
+          eventName: HookEventName.BeforeModel,
+          success: true,
+          output: {
+            hookSpecificOutput: {
+              hookEventName: 'BeforeModel',
+              additionalContext: 'Context 1',
+            },
+          },
+          duration: 10,
+        },
+        {
+          hookConfig: {
+            type: HookType.Command,
+            command: 'h2',
+            timeout: 30000,
+          },
+          eventName: HookEventName.BeforeModel,
+          success: true,
+          output: {
+            hookSpecificOutput: {
+              hookEventName: 'BeforeModel',
+              additionalContext: 'Context 2',
+            },
+          },
+          duration: 10,
+        },
+      ];
+
+      const aggregated = aggregator.aggregateResults(
+        results,
+        HookEventName.BeforeModel,
+      );
+
+      expect(
+        aggregated.finalOutput?.hookSpecificOutput?.['additionalContext'],
+      ).toBe(
+        '<hook_context hook="h1">\nContext 1\n</hook_context>\n<hook_context hook="h2">\nContext 2\n</hook_context>',
+      );
+    });
   });
 
   describe('extractAdditionalContext', () => {
@@ -470,7 +518,40 @@ describe('HookAggregator', () => {
       expect(aggregated.success).toBe(true);
       expect(
         aggregated.finalOutput?.hookSpecificOutput?.['additionalContext'],
-      ).toBe('Context from hook 1\nContext from hook 2');
+      ).toBe(
+        '<hook_context hook="test-command">\nContext from hook 1\n</hook_context>\n<hook_context hook="test-command">\nContext from hook 2\n</hook_context>',
+      );
+    });
+
+    it('should sanitize additional context by escaping < and > tags', () => {
+      const results: HookExecutionResult[] = [
+        {
+          hookConfig: {
+            type: HookType.Command,
+            command: 'test-hook',
+          },
+          eventName: HookEventName.AfterTool,
+          success: true,
+          output: {
+            hookSpecificOutput: {
+              hookEventName: 'AfterTool',
+              additionalContext: 'context with <b>bold</b> and <script> tags',
+            },
+          },
+          duration: 10,
+        },
+      ];
+
+      const aggregated = aggregator.aggregateResults(
+        results,
+        HookEventName.AfterTool,
+      );
+
+      expect(
+        aggregated.finalOutput?.hookSpecificOutput?.['additionalContext'],
+      ).toBe(
+        '<hook_context hook="test-hook">\ncontext with &lt;b&gt;bold&lt;/b&gt; and &lt;script&gt; tags\n</hook_context>',
+      );
     });
   });
 });

packages/core/src/hooks/hookAggregator.ts

@@ -58,7 +58,7 @@ export class HookAggregator {
     }
 
     // Merge outputs using event-specific strategy
-    const mergedOutput = this.mergeOutputs(allOutputs, eventName);
+    const mergedOutput = this.mergeOutputs(results, eventName);
     const finalOutput = mergedOutput
       ? this.createSpecificHookOutput(mergedOutput, eventName)
       : undefined;
@@ -79,10 +79,11 @@ export class HookAggregator {
    * consistent default behaviors (e.g., default decision='allow' for OR logic)
    */
   private mergeOutputs(
-    outputs: HookOutput[],
+    results: HookExecutionResult[],
     eventName: HookEventName,
   ): HookOutput | undefined {
-    if (outputs.length === 0) {
+    const resultsWithOutput = results.filter((r) => r.output);
+    if (resultsWithOutput.length === 0) {
       return undefined;
     }
 
@@ -92,28 +93,25 @@ export class HookAggregator {
       case HookEventName.BeforeAgent:
       case HookEventName.AfterAgent:
       case HookEventName.SessionStart:
-        return this.mergeWithOrDecision(outputs);
+        return this.mergeWithOrDecision(resultsWithOutput);
 
       case HookEventName.BeforeModel:
       case HookEventName.AfterModel:
-        return this.mergeWithFieldReplacement(outputs);
+        return this.mergeWithFieldReplacement(resultsWithOutput);
 
       case HookEventName.BeforeToolSelection:
-        return this.mergeToolSelectionOutputs(
-          // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
-          outputs as BeforeToolSelectionOutput[],
-        );
+        return this.mergeToolSelectionOutputs(resultsWithOutput);
 
       default:
         // For other events, use simple merge
-        return this.mergeSimple(outputs);
+        return this.mergeSimple(resultsWithOutput);
     }
   }
 
   /**
    * Merge outputs with OR decision logic and message concatenation
    */
-  private mergeWithOrDecision(outputs: HookOutput[]): HookOutput {
+  private mergeWithOrDecision(results: HookExecutionResult[]): HookOutput {
     const merged: HookOutput = {
       continue: true,
       suppressOutput: false,
@@ -128,7 +126,9 @@ export class HookAggregator {
     let hasAskDecision = false;
     let hasContinueFalse = false;
 
-    for (const output of outputs) {
+    for (const result of results) {
+      const output = result.output!;
+
       // Handle continue flag
       if (output.continue === false) {
         hasContinueFalse = true;
@@ -184,7 +184,7 @@ export class HookAggregator {
       }
 
       // Collect additional context from hook-specific outputs
-      this.extractAdditionalContext(output, additionalContexts);
+      this.extractAdditionalContext(result, additionalContexts);
     }
 
     // Set final decision if no blocking or ask decision was found
@@ -219,10 +219,18 @@ export class HookAggregator {
   /**
    * Merge outputs with later fields replacing earlier fields
    */
-  private mergeWithFieldReplacement(outputs: HookOutput[]): HookOutput {
+  private mergeWithFieldReplacement(
+    results: HookExecutionResult[],
+  ): HookOutput {
     let merged: HookOutput = {};
+    const additionalContexts: string[] = [];
+
+    for (const result of results) {
+      const output = result.output!;
+
+      // Collect additional context
+      this.extractAdditionalContext(result, additionalContexts);
 
-    for (const output of outputs) {
       // Later outputs override earlier ones
       merged = {
         ...merged,
@@ -234,6 +242,14 @@ export class HookAggregator {
       };
     }
 
+    // Add merged additional context
+    if (additionalContexts.length > 0) {
+      merged.hookSpecificOutput = {
+        ...(merged.hookSpecificOutput || {}),
+        additionalContext: additionalContexts.join('\n'),
+      };
+    }
+
     return merged;
   }
 
@@ -251,7 +267,7 @@ export class HookAggregator {
    * If one hook restricts and another re-enables, the union takes the re-enabled tool.
    */
   private mergeToolSelectionOutputs(
-    outputs: BeforeToolSelectionOutput[],
+    results: HookExecutionResult[],
   ): BeforeToolSelectionOutput {
     const merged: BeforeToolSelectionOutput = {};
 
@@ -259,7 +275,9 @@ export class HookAggregator {
     let hasNoneMode = false;
     let hasAnyMode = false;
 
-    for (const output of outputs) {
+    for (const result of results) {
+      // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
+      const output = result.output as BeforeToolSelectionOutput;
       const toolConfig = output.hookSpecificOutput?.toolConfig;
       if (!toolConfig) {
         continue;
@@ -314,11 +332,11 @@ export class HookAggregator {
   /**
    * Simple merge for events without special logic
    */
-  private mergeSimple(outputs: HookOutput[]): HookOutput {
+  private mergeSimple(results: HookExecutionResult[]): HookOutput {
     let merged: HookOutput = {};
 
-    for (const output of outputs) {
-      merged = { ...merged, ...output };
+    for (const result of results) {
+      merged = { ...merged, ...result.output! };
     }
 
     return merged;
@@ -351,10 +369,10 @@ export class HookAggregator {
    * Extract additional context from hook-specific outputs
    */
   private extractAdditionalContext(
-    output: HookOutput,
+    result: HookExecutionResult,
     contexts: string[],
   ): void {
-    const specific = output.hookSpecificOutput;
+    const specific = result.output?.hookSpecificOutput;
     if (!specific) {
       return;
     }
@@ -365,7 +383,15 @@ export class HookAggregator {
       // eslint-disable-next-line no-restricted-syntax
       typeof specific['additionalContext'] === 'string'
     ) {
-      contexts.push(specific['additionalContext']);
+      const hookName =
+        result.hookConfig.name || result.hookConfig.command || 'unknown-hook';
+      // Sanitize the context text before wrapping to prevent tag injection
+      const contextText = specific['additionalContext']
+        .replace(/</g, '&lt;')
+        .replace(/>/g, '&gt;');
+      contexts.push(
+        `<hook_context hook="${hookName}">\n${contextText}\n</hook_context>`,
+      );
     }
   }
 }

packages/core/src/hooks/hookRunner.test.ts

@@ -659,10 +659,9 @@ describe('HookRunner', () => {
 
       // Verify that the second hook received modified input
       const secondHookInput = JSON.parse(
-        vi.mocked(mockSpawn.stdin.write).mock.calls[1][0],
+        vi.mocked(mockSpawn.stdin.write).mock.calls[1][0] as string,
       );
-      expect(secondHookInput.prompt).toContain('Original prompt');
-      expect(secondHookInput.prompt).toContain('Context from hook 1');
+      expect(secondHookInput.prompt).toBe('Original prompt');
     });
 
     it('should pass modified LLM request from one hook to the next for BeforeModel', async () => {

packages/core/src/hooks/hookRunner.ts

@@ -15,7 +15,6 @@ import {
   type HookInput,
   type HookOutput,
   type HookExecutionResult,
-  type BeforeAgentInput,
   type BeforeModelInput,
   type BeforeModelOutput,
   type BeforeToolInput,
@@ -180,22 +179,6 @@ export class HookRunner {
     // Apply modifications based on hook output and event type
     if (hookOutput.hookSpecificOutput) {
       switch (eventName) {
-        case HookEventName.BeforeAgent:
-          if ('additionalContext' in hookOutput.hookSpecificOutput) {
-            // For BeforeAgent, we could modify the prompt with additional context
-            const additionalContext =
-              hookOutput.hookSpecificOutput['additionalContext'];
-            if (
-              typeof additionalContext === 'string' &&
-              'prompt' in modifiedInput
-            ) {
-              // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
-              (modifiedInput as BeforeAgentInput).prompt +=
-                '\n\n' + additionalContext;
-            }
-          }
-          break;
-
         case HookEventName.BeforeModel:
           if ('llm_request' in hookOutput.hookSpecificOutput) {
             // For BeforeModel, we update the LLM request

packages/core/src/hooks/types.test.ts

@@ -24,6 +24,7 @@ import {
 import type {
   GenerateContentParameters,
   GenerateContentResponse,
+  Content,
   ToolConfig,
 } from '@google/genai';
 
@@ -187,14 +188,14 @@ describe('Hook Output Classes', () => {
       expect(output.getAdditionalContext()).toBe('some context');
     });
 
-    it('getAdditionalContext should sanitize context by escaping <', () => {
+    it('getAdditionalContext should return raw context without sanitization (handled by aggregator)', () => {
       const output = new DefaultHookOutput({
         hookSpecificOutput: {
           additionalContext: 'context with <tag> and </hook_context>',
         },
       });
       expect(output.getAdditionalContext()).toBe(
-        'context with &lt;tag&gt; and &lt;/hook_context&gt;',
+        'context with <tag> and </hook_context>',
       );
     });
 
@@ -271,6 +272,73 @@ describe('Hook Output Classes', () => {
       const output = new BeforeModelHookOutput({});
       expect(output.applyLLMRequestModifications(target)).toBe(target);
     });
+
+    it('applyLLMRequestModifications should append additionalContext to contents', () => {
+      const target: GenerateContentParameters = {
+        model: 'gemini-pro',
+        contents: [{ role: 'user', parts: [{ text: 'Hello' }] }],
+      };
+      const output = new BeforeModelHookOutput({
+        hookSpecificOutput: {
+          additionalContext:
+            '<hook_context hook="test">\nNew Context\n</hook_context>',
+        },
+      });
+      const result = output.applyLLMRequestModifications(target);
+      expect(result.contents).toHaveLength(1);
+      const contents = result.contents as Content[];
+      expect(contents[0].parts!).toHaveLength(2);
+      expect(contents[0].parts![1]).toEqual({
+        text: '\n\n<hook_context hook="test">\nNew Context\n</hook_context>',
+      });
+    });
+
+    it('applyLLMRequestModifications should create new user message if none exists', () => {
+      const target: GenerateContentParameters = {
+        model: 'gemini-pro',
+        contents: [{ role: 'model', parts: [{ text: 'Hi' }] }],
+      };
+      const output = new BeforeModelHookOutput({
+        hookSpecificOutput: {
+          additionalContext:
+            '<hook_context hook="test">\nNew Context\n</hook_context>',
+        },
+      });
+      const result = output.applyLLMRequestModifications(target);
+      expect(result.contents).toHaveLength(2);
+      const contents = result.contents as Content[];
+      expect(contents[1].role).toBe('user');
+      expect(contents[1].parts![0]).toEqual({
+        text: '\n\n<hook_context hook="test">\nNew Context\n</hook_context>',
+      });
+    });
+
+    it('applyLLMRequestModifications should handle both llm_request and additionalContext', () => {
+      const target: GenerateContentParameters = {
+        model: 'gemini-pro',
+        contents: [{ role: 'user', parts: [{ text: 'original' }] }],
+      };
+      const mockRequest = {
+        // Our mock fromHookLLMRequest just spreads the request, so we can use contents here for convenience in tests
+        // even though LLMRequest uses messages.
+        contents: [{ role: 'user', parts: [{ text: 'modified' }] }],
+      } as unknown as Partial<LLMRequest>;
+      const output = new BeforeModelHookOutput({
+        hookSpecificOutput: {
+          llm_request: mockRequest as LLMRequest,
+          additionalContext:
+            '<hook_context hook="test">\nNew Context\n</hook_context>',
+        },
+      });
+      const result = output.applyLLMRequestModifications(target);
+      expect(result.contents).toHaveLength(1);
+      const contents = result.contents as Content[];
+      expect(contents[0].parts!).toHaveLength(2);
+      expect(contents[0].parts![0]).toEqual({ text: 'modified' });
+      expect(contents[0].parts![1]).toEqual({
+        text: '\n\n<hook_context hook="test">\nNew Context\n</hook_context>',
+      });
+    });
   });
 
   describe('BeforeToolSelectionHookOutput', () => {

packages/core/src/hooks/types.ts

@@ -7,6 +7,8 @@
 import type {
   GenerateContentResponse,
   GenerateContentParameters,
+  Content,
+  Part,
   ToolConfig as GenAIToolConfig,
   ToolListUnion,
 } from '@google/genai';
@@ -270,8 +272,7 @@ export class DefaultHookOutput implements HookOutput {
         return undefined;
       }
 
-      // Sanitize by escaping < and > to prevent tag injection
-      return context.replace(/</g, '&lt;').replace(/>/g, '&gt;');
+      return context;
     }
     return undefined;
   }
@@ -374,6 +375,8 @@ export class BeforeModelHookOutput extends DefaultHookOutput {
   override applyLLMRequestModifications(
     target: GenerateContentParameters,
   ): GenerateContentParameters {
+    let resultTarget = target;
+
     if (this.hookSpecificOutput && 'llm_request' in this.hookSpecificOutput) {
       // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
       const hookRequest = this.hookSpecificOutput[
@@ -386,13 +389,114 @@ export class BeforeModelHookOutput extends DefaultHookOutput {
           hookRequest as LLMRequest,
           target,
         );
-        return {
+        resultTarget = {
           ...target,
           ...sdkRequest,
         };
       }
     }
-    return target;
+
+    const additionalContext = this.getAdditionalContext();
+    if (additionalContext) {
+      const originalContents = resultTarget.contents;
+      let contents: Content[];
+
+      const isPart = (obj: unknown): obj is Part =>
+        typeof obj === 'object' &&
+        obj !== null &&
+        !('role' in obj) &&
+        !('parts' in obj);
+
+      if (Array.isArray(originalContents)) {
+        const first = originalContents[0];
+        if (
+          first &&
+          typeof first !== 'string' &&
+          ('role' in first || 'parts' in first)
+        ) {
+          // It's already Content[]
+          contents = originalContents.map((c) => {
+            const contentObj =
+              typeof c === 'string'
+                ? { role: 'user', parts: [{ text: c }] }
+                : c;
+            if (
+              typeof contentObj === 'object' &&
+              contentObj !== null &&
+              ('role' in contentObj || 'parts' in contentObj)
+            ) {
+              return {
+                role:
+                  'role' in contentObj && typeof contentObj.role === 'string'
+                    ? contentObj.role
+                    : 'user',
+                parts:
+                  'parts' in contentObj && Array.isArray(contentObj.parts)
+                    ? [...contentObj.parts]
+                    : [],
+              };
+            }
+            return { role: 'user', parts: [] };
+          });
+        } else {
+          // It's Part[] or (Part | string)[]
+          const parts: Part[] = originalContents.reduce((acc, c) => {
+            if (typeof c === 'string') {
+              acc.push({ text: c });
+            } else if (isPart(c)) {
+              acc.push(c); // Safe because it lacks Content fields
+            }
+            return acc;
+          }, [] as Part[]);
+          contents = [{ role: 'user', parts }];
+        }
+      } else if (typeof originalContents === 'string') {
+        contents = [{ role: 'user', parts: [{ text: originalContents }] }];
+      } else if (
+        originalContents &&
+        typeof originalContents === 'object' &&
+        ('role' in originalContents || 'parts' in originalContents)
+      ) {
+        const c = originalContents;
+        contents = [
+          {
+            role: 'role' in c && typeof c.role === 'string' ? c.role : 'user',
+            parts: 'parts' in c && Array.isArray(c.parts) ? [...c.parts] : [],
+          },
+        ];
+      } else if (isPart(originalContents)) {
+        // It's a single Part
+        contents = [{ role: 'user', parts: [originalContents] }]; // Safe because it lacks Content fields
+      } else {
+        contents = [];
+      }
+
+      const wrappedContext = `\n\n${additionalContext}`;
+
+      let lastUserMessageIndex = -1;
+      for (let i = contents.length - 1; i >= 0; i--) {
+        if (contents[i].role === 'user') {
+          lastUserMessageIndex = i;
+          break;
+        }
+      }
+
+      if (lastUserMessageIndex !== -1) {
+        if (!contents[lastUserMessageIndex].parts) {
+          contents[lastUserMessageIndex].parts = [];
+        }
+        contents[lastUserMessageIndex].parts!.push({ text: wrappedContext });
+      } else {
+        contents.push({ role: 'user', parts: [{ text: wrappedContext }] });
+      }
+
+      resultTarget = {
+        ...resultTarget,
+        contents,
+      };
+    }
+
+    return resultTarget;
   }
 }
 
@@ -683,6 +787,7 @@ export interface BeforeModelOutput extends HookOutput {
     hookEventName: 'BeforeModel';
     llm_request?: Partial<LLMRequest>;
     llm_response?: LLMResponse;
+    additionalContext?: string;
   };
 }