From 6d1e27633a32ce778127f1cf8fd88845284ae507 Mon Sep 17 00:00:00 2001 From: Christian Gunderman Date: Mon, 5 Jan 2026 13:27:53 -0800 Subject: [PATCH] Support context injection via SessionStart hook. (#15746) --- integration-tests/hooks-system.test.ts | 164 ++++++++++++++++++ packages/cli/src/gemini.tsx | 36 +++- packages/cli/src/ui/AppContainer.tsx | 32 +++- packages/cli/src/ui/commands/clearCommand.ts | 15 +- .../core/src/core/contentGenerator.test.ts | 4 +- packages/core/src/core/contentGenerator.ts | 5 +- packages/core/src/core/sessionHookTriggers.ts | 24 ++- 7 files changed, 261 insertions(+), 19 deletions(-) diff --git a/integration-tests/hooks-system.test.ts b/integration-tests/hooks-system.test.ts index bb7d2cd565..c5ab2b024d 100644 --- a/integration-tests/hooks-system.test.ts +++ b/integration-tests/hooks-system.test.ts @@ -950,6 +950,170 @@ try { } }); + it('should fire SessionStart hook and inject context', async () => { + // Create hook script that outputs JSON with additionalContext + const hookScript = `const fs = require('fs'); +console.log(JSON.stringify({ + decision: 'allow', + systemMessage: 'Context injected via SessionStart hook', + hookSpecificOutput: { + hookEventName: 'SessionStart', + additionalContext: 'If asked about your identity, you must respond with the following: I am a protocol droid that speaks the bacchi.' + } +}));`; + + await rig.setup('should fire SessionStart hook and inject context', { + fakeResponsesPath: join( + import.meta.dirname, + 'hooks-system.session-startup.responses', + ), + }); + + const scriptPath = join(rig.testDir!, 'session_start_context_hook.cjs'); + writeFileSync(scriptPath, hookScript); + + await rig.setup('should fire SessionStart hook and inject context', { + settings: { + tools: { + enableHooks: true, + }, + hooks: { + SessionStart: [ + { + matcher: 'startup', + hooks: [ + { + type: 'command', + command: `node "${scriptPath}"`, + timeout: 5000, + }, + ], + }, + ], + }, + }, + }); + + // Run a query - the SessionStart hook will fire during app initialization + const result = await rig.run({ args: 'Who are you?' }); + + // Check if systemMessage was displayed (in stderr, which rig.run captures) + expect(result).toContain('Context injected via SessionStart hook'); + + // Check if additionalContext influenced the model response + // Note: We use fake responses, but the rig records interactions. + // If we are using fake responses, the model won't actually respond unless we provide a fake response for the injected context. + // But the test rig setup uses 'hooks-system.session-startup.responses'. + // If I'm adding a new test, I might need to generate new fake responses or expect the context to be sent to the model (verify API logs). + + // Verify hook executed + const hookLogs = rig.readHookLogs(); + const sessionStartLog = hookLogs.find( + (log) => log.hookCall.hook_event_name === 'SessionStart', + ); + + expect(sessionStartLog).toBeDefined(); + + // Verify the API request contained the injected context + // rig.readAllApiRequest() gives us telemetry on API requests. + const apiRequests = rig.readAllApiRequest(); + // We expect at least one API request + expect(apiRequests.length).toBeGreaterThan(0); + + // The injected context should be in the request text + // For non-interactive mode, I prepended it to input: "context\n\ninput" + // The telemetry `request_text` should contain it. + const requestText = apiRequests[0].attributes?.request_text || ''; + expect(requestText).toContain('protocol droid'); + }); + + it('should fire SessionStart hook and display systemMessage in interactive mode', async () => { + // Create hook script that outputs JSON with systemMessage and additionalContext + const hookScript = `const fs = require('fs'); +console.log(JSON.stringify({ + decision: 'allow', + systemMessage: 'Interactive Session Start Message', + hookSpecificOutput: { + hookEventName: 'SessionStart', + additionalContext: 'The user is a Jedi Master.' + } +}));`; + + await rig.setup( + 'should fire SessionStart hook and display systemMessage in interactive mode', + { + fakeResponsesPath: join( + import.meta.dirname, + 'hooks-system.session-startup.responses', + ), + }, + ); + + const scriptPath = join( + rig.testDir!, + 'session_start_interactive_hook.cjs', + ); + writeFileSync(scriptPath, hookScript); + + await rig.setup( + 'should fire SessionStart hook and display systemMessage in interactive mode', + { + settings: { + tools: { + enableHooks: true, + }, + hooks: { + SessionStart: [ + { + matcher: 'startup', + hooks: [ + { + type: 'command', + command: `node "${scriptPath}"`, + timeout: 5000, + }, + ], + }, + ], + }, + }, + }, + ); + + const run = await rig.runInteractive(); + + // Verify systemMessage is displayed + await run.expectText('Interactive Session Start Message', 10000); + + // Send a prompt to establish a session and trigger an API call + await run.sendKeys('Hello'); + await run.sendKeys('\r'); + + // Wait for response to ensure API call happened + await run.expectText('Hello', 15000); + + // Wait for telemetry to be written to disk + await rig.waitForTelemetryReady(); + + // Verify the API request contained the injected context + // We may need to poll for API requests as they are written asynchronously + const pollResult = await poll( + () => { + const apiRequests = rig.readAllApiRequest(); + return apiRequests.length > 0; + }, + 15000, + 500, + ); + + expect(pollResult).toBe(true); + + const apiRequests = rig.readAllApiRequest(); + // The injected context should be in the request_text of the API request + const requestText = apiRequests[0].attributes?.request_text || ''; + expect(requestText).toContain('Jedi Master'); + }); + it('should fire SessionEnd and SessionStart hooks on /clear command', async () => { // Create inline hook commands for both SessionEnd and SessionStart const sessionEndCommand = diff --git a/packages/cli/src/gemini.tsx b/packages/cli/src/gemini.tsx index eacef49cb3..4a0096150a 100644 --- a/packages/cli/src/gemini.tsx +++ b/packages/cli/src/gemini.tsx @@ -634,6 +634,16 @@ export async function main() { await config.initialize(); startupProfiler.flush(config); + // If not a TTY, read from stdin + // This is for cases where the user pipes input directly into the command + let stdinData: string | undefined = undefined; + if (!process.stdin.isTTY) { + stdinData = await readStdin(); + if (stdinData) { + input = input ? `${stdinData}\n\n${input}` : stdinData; + } + } + // Fire SessionStart hook through MessageBus (only if hooks are enabled) // Must be called AFTER config.initialize() to ensure HookRegistry is loaded const hooksEnabled = config.getEnableHooks(); @@ -642,7 +652,23 @@ export async function main() { const sessionStartSource = resumedSessionData ? SessionStartSource.Resume : SessionStartSource.Startup; - await fireSessionStartHook(hookMessageBus, sessionStartSource); + const result = await fireSessionStartHook( + hookMessageBus, + sessionStartSource, + ); + + if (result) { + if (result.systemMessage) { + writeToStderr(result.systemMessage + '\n'); + } + const additionalContext = result.getAdditionalContext(); + if (additionalContext) { + // Prepend context to input (System Context -> Stdin -> Question) + input = input + ? `${additionalContext}\n\n${input}` + : additionalContext; + } + } // Register SessionEnd hook for graceful exit registerCleanup(async () => { @@ -650,14 +676,6 @@ export async function main() { }); } - // If not a TTY, read from stdin - // This is for cases where the user pipes input directly into the command - if (!process.stdin.isTTY) { - const stdinData = await readStdin(); - if (stdinData) { - input = `${stdinData}\n\n${input}`; - } - } if (!input) { debugLogger.error( `No input provided via stdin. Input can be provided by piping data into gemini or using the --prompt option.`, diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index bba0f1fd4e..a7ac42ce66 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/packages/cli/src/ui/AppContainer.tsx @@ -300,7 +300,31 @@ export const AppContainer = (props: AppContainerProps) => { const sessionStartSource = resumedSessionData ? SessionStartSource.Resume : SessionStartSource.Startup; - await fireSessionStartHook(hookMessageBus, sessionStartSource); + const result = await fireSessionStartHook( + hookMessageBus, + sessionStartSource, + ); + + if (result) { + if (result.systemMessage) { + historyManager.addItem( + { + type: MessageType.INFO, + text: result.systemMessage, + }, + Date.now(), + ); + } + + const additionalContext = result.getAdditionalContext(); + const geminiClient = config.getGeminiClient(); + if (additionalContext && geminiClient) { + await geminiClient.addHistory({ + role: 'user', + parts: [{ text: additionalContext }], + }); + } + } } // Fire-and-forget: generate summary for previous session in background @@ -321,6 +345,12 @@ export const AppContainer = (props: AppContainerProps) => { await fireSessionEndHook(hookMessageBus, SessionEndReason.Exit); } }); + // Disable the dependencies check here. historyManager gets flagged + // but we don't want to react to changes to it because each new history + // item, including the ones from the start session hook will cause a + // re-render and an error when we try to reload config. + // + // eslint-disable-next-line react-hooks/exhaustive-deps }, [config, resumedSessionData]); useEffect( diff --git a/packages/cli/src/ui/commands/clearCommand.ts b/packages/cli/src/ui/commands/clearCommand.ts index d2edbebbf2..ec8b7a52ef 100644 --- a/packages/cli/src/ui/commands/clearCommand.ts +++ b/packages/cli/src/ui/commands/clearCommand.ts @@ -4,6 +4,7 @@ * SPDX-License-Identifier: Apache-2.0 */ +import type { DefaultHookOutput } from '@google/gemini-cli-core'; import { uiTelemetryService, fireSessionEndHook, @@ -14,6 +15,7 @@ import { } from '@google/gemini-cli-core'; import type { SlashCommand } from './types.js'; import { CommandKind } from './types.js'; +import { MessageType } from '../types.js'; import { randomUUID } from 'node:crypto'; export const clearCommand: SlashCommand = { @@ -52,8 +54,9 @@ export const clearCommand: SlashCommand = { } // Fire SessionStart hook after clearing + let result: DefaultHookOutput | undefined; if (config?.getEnableHooks() && messageBus) { - await fireSessionStartHook(messageBus, SessionStartSource.Clear); + result = await fireSessionStartHook(messageBus, SessionStartSource.Clear); } // Give the event loop a chance to process any pending telemetry operations @@ -68,5 +71,15 @@ export const clearCommand: SlashCommand = { uiTelemetryService.setLastPromptTokenCount(0); context.ui.clear(); + + if (result?.systemMessage) { + context.ui.addItem( + { + type: MessageType.INFO, + text: result.systemMessage, + }, + Date.now(), + ); + } }, }; diff --git a/packages/core/src/core/contentGenerator.test.ts b/packages/core/src/core/contentGenerator.test.ts index 6acae4f57e..9e10558a18 100644 --- a/packages/core/src/core/contentGenerator.test.ts +++ b/packages/core/src/core/contentGenerator.test.ts @@ -61,7 +61,9 @@ describe('createContentGenerator', () => { expect(FakeContentGenerator.fromFile).toHaveBeenCalledWith( fakeResponsesFile, ); - expect(generator).toEqual(mockGenerator); + expect(generator).toEqual( + new LoggingContentGenerator(mockGenerator, mockConfigWithFake), + ); }); it('should create a RecordingContentGenerator', async () => { diff --git a/packages/core/src/core/contentGenerator.ts b/packages/core/src/core/contentGenerator.ts index e4b568b871..12e07790cc 100644 --- a/packages/core/src/core/contentGenerator.ts +++ b/packages/core/src/core/contentGenerator.ts @@ -114,7 +114,10 @@ export async function createContentGenerator( ): Promise { const generator = await (async () => { if (gcConfig.fakeResponses) { - return FakeContentGenerator.fromFile(gcConfig.fakeResponses); + return new LoggingContentGenerator( + await FakeContentGenerator.fromFile(gcConfig.fakeResponses), + gcConfig, + ); } const version = await getVersion(); const model = resolveModel( diff --git a/packages/core/src/core/sessionHookTriggers.ts b/packages/core/src/core/sessionHookTriggers.ts index 524fe9beb9..149a84edbd 100644 --- a/packages/core/src/core/sessionHookTriggers.ts +++ b/packages/core/src/core/sessionHookTriggers.ts @@ -10,10 +10,12 @@ import { type HookExecutionRequest, type HookExecutionResponse, } from '../confirmation-bus/types.js'; -import type { - SessionStartSource, - SessionEndReason, - PreCompressTrigger, +import { + type SessionStartSource, + type SessionEndReason, + type PreCompressTrigger, + createHookOutput, + type DefaultHookOutput, } from '../hooks/types.js'; import { debugLogger } from '../utils/debugLogger.js'; @@ -22,13 +24,17 @@ import { debugLogger } from '../utils/debugLogger.js'; * * @param messageBus The message bus to use for hook communication * @param source The source/trigger of the session start + * @returns The output from the SessionStart hook, or undefined if failed/no output */ export async function fireSessionStartHook( messageBus: MessageBus, source: SessionStartSource, -): Promise { +): Promise { try { - await messageBus.request( + const response = await messageBus.request< + HookExecutionRequest, + HookExecutionResponse + >( { type: MessageBusType.HOOK_EXECUTION_REQUEST, eventName: 'SessionStart', @@ -38,8 +44,14 @@ export async function fireSessionStartHook( }, MessageBusType.HOOK_EXECUTION_RESPONSE, ); + + if (response.output) { + return createHookOutput('SessionStart', response.output); + } + return undefined; } catch (error) { debugLogger.debug(`SessionStart hook failed:`, error); + return undefined; } }