feat(hooks): Hook Session Lifecycle & Compression Integration (#14151)

This commit is contained in:
Edilmo Palencia
2025-12-03 09:04:13 -08:00
committed by GitHub
parent 7a6d3067c6
commit 1c12da1fad
27 changed files with 1026 additions and 302 deletions
+11
View File
@@ -257,6 +257,7 @@ describe('gemini.tsx main function', () => {
getMessageBus: () => ({
subscribe: vi.fn(),
}),
getEnableHooks: () => false,
getToolRegistry: vi.fn(),
getContentGeneratorConfig: vi.fn(),
getModel: () => 'gemini-pro',
@@ -489,6 +490,7 @@ describe('gemini.tsx main function kitty protocol', () => {
getMessageBus: () => ({
subscribe: vi.fn(),
}),
getEnableHooks: () => false,
getToolRegistry: vi.fn(),
getContentGeneratorConfig: vi.fn(),
getModel: () => 'gemini-pro',
@@ -588,6 +590,7 @@ describe('gemini.tsx main function kitty protocol', () => {
getExtensions: () => [{ name: 'ext1' }],
getPolicyEngine: vi.fn(),
getMessageBus: () => ({ subscribe: vi.fn() }),
getEnableHooks: () => false,
initialize: vi.fn(),
getContentGeneratorConfig: vi.fn(),
getMcpServers: () => ({}),
@@ -668,6 +671,7 @@ describe('gemini.tsx main function kitty protocol', () => {
getExtensions: () => [],
getPolicyEngine: vi.fn(),
getMessageBus: () => ({ subscribe: vi.fn() }),
getEnableHooks: () => false,
initialize: vi.fn(),
getContentGeneratorConfig: vi.fn(),
getMcpServers: () => ({}),
@@ -733,6 +737,7 @@ describe('gemini.tsx main function kitty protocol', () => {
getDebugMode: () => false,
getPolicyEngine: vi.fn(),
getMessageBus: () => ({ subscribe: vi.fn() }),
getEnableHooks: () => false,
initialize: vi.fn(),
getContentGeneratorConfig: vi.fn(),
getMcpServers: () => ({}),
@@ -814,6 +819,7 @@ describe('gemini.tsx main function kitty protocol', () => {
getDebugMode: () => false,
getPolicyEngine: vi.fn(),
getMessageBus: () => ({ subscribe: vi.fn() }),
getEnableHooks: () => false,
initialize: vi.fn(),
getContentGeneratorConfig: vi.fn(),
getMcpServers: () => ({}),
@@ -890,6 +896,7 @@ describe('gemini.tsx main function kitty protocol', () => {
getDebugMode: () => false,
getPolicyEngine: vi.fn(),
getMessageBus: () => ({ subscribe: vi.fn() }),
getEnableHooks: () => false,
initialize: vi.fn(),
getContentGeneratorConfig: vi.fn(),
getMcpServers: () => ({}),
@@ -961,6 +968,7 @@ describe('gemini.tsx main function kitty protocol', () => {
getDebugMode: () => false,
getPolicyEngine: vi.fn(),
getMessageBus: () => ({ subscribe: vi.fn() }),
getEnableHooks: () => false,
initialize: vi.fn(),
getContentGeneratorConfig: vi.fn(),
getMcpServers: () => ({}),
@@ -1130,6 +1138,7 @@ describe('gemini.tsx main function exit codes', () => {
getGeminiMdFileCount: () => 0,
getPolicyEngine: vi.fn(),
getMessageBus: () => ({ subscribe: vi.fn() }),
getEnableHooks: () => false,
getToolRegistry: vi.fn(),
getContentGeneratorConfig: vi.fn(),
getModel: () => 'gemini-pro',
@@ -1191,6 +1200,7 @@ describe('gemini.tsx main function exit codes', () => {
getGeminiMdFileCount: () => 0,
getPolicyEngine: vi.fn(),
getMessageBus: () => ({ subscribe: vi.fn() }),
getEnableHooks: () => false,
getToolRegistry: vi.fn(),
getContentGeneratorConfig: vi.fn(),
getModel: () => 'gemini-pro',
@@ -1302,6 +1312,7 @@ describe('startInteractiveUI', () => {
registerCleanup: vi.fn(),
runExitCleanup: vi.fn(),
registerSyncCleanup: vi.fn(),
registerTelemetryConfig: vi.fn(),
}));
afterEach(() => {
+33
View File
@@ -31,6 +31,7 @@ import {
registerCleanup,
registerSyncCleanup,
runExitCleanup,
registerTelemetryConfig,
} from './utils/cleanup.js';
import { getCliVersion } from './utils/version.js';
import {
@@ -58,6 +59,10 @@ import {
shouldEnterAlternateScreen,
startupProfiler,
ExitCodes,
SessionStartSource,
SessionEndReason,
fireSessionStartHook,
fireSessionEndHook,
} from '@google/gemini-cli-core';
import {
initializeApp,
@@ -459,10 +464,22 @@ export async function main() {
const config = await loadCliConfig(settings.merged, sessionId, argv);
loadConfigHandle?.end();
// Register config for telemetry shutdown
// This ensures telemetry (including SessionEnd hooks) is properly flushed on exit
registerTelemetryConfig(config);
const policyEngine = config.getPolicyEngine();
const messageBus = config.getMessageBus();
createPolicyUpdater(policyEngine, messageBus);
// Register SessionEnd hook to fire on graceful exit
// This runs before telemetry shutdown in runExitCleanup()
if (config.getEnableHooks() && messageBus) {
registerCleanup(async () => {
await fireSessionEndHook(messageBus, SessionEndReason.Exit);
});
}
// Cleanup sessions after config initialization
try {
await cleanupExpiredSessions(config, settings.merged);
@@ -586,6 +603,22 @@ export async function main() {
await config.initialize();
startupProfiler.flush(config);
// 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();
const hookMessageBus = config.getMessageBus();
if (hooksEnabled && hookMessageBus) {
const sessionStartSource = resumedSessionData
? SessionStartSource.Resume
: SessionStartSource.Startup;
await fireSessionStartHook(hookMessageBus, sessionStartSource);
// Register SessionEnd hook for graceful exit
registerCleanup(async () => {
await fireSessionEndHook(hookMessageBus, SessionEndReason.Exit);
});
}
// If not a TTY, read from stdin
// This is for cases where the user pipes input directly into the command
if (!process.stdin.isTTY) {
+1
View File
@@ -186,6 +186,7 @@ describe('gemini.tsx main function cleanup', () => {
getDebugMode: vi.fn(() => false),
getPolicyEngine: vi.fn(),
getMessageBus: () => ({ subscribe: vi.fn() }),
getEnableHooks: vi.fn(() => false),
initialize: vi.fn(),
getContentGeneratorConfig: vi.fn(),
getMcpServers: () => ({}),
+2 -5
View File
@@ -16,7 +16,6 @@ import type {
import {
executeToolCall,
ToolErrorType,
shutdownTelemetry,
GeminiEventType,
OutputFormat,
uiTelemetryService,
@@ -61,7 +60,6 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => {
return {
...original,
executeToolCall: vi.fn(),
shutdownTelemetry: vi.fn(),
isTelemetrySdkInitialized: vi.fn().mockReturnValue(true),
ChatRecordingService: MockChatRecordingService,
uiTelemetryService: {
@@ -91,7 +89,6 @@ describe('runNonInteractive', () => {
let mockSettings: LoadedSettings;
let mockToolRegistry: ToolRegistry;
let mockCoreExecuteToolCall: Mock;
let mockShutdownTelemetry: Mock;
let consoleErrorSpy: MockInstance;
let processStdoutSpy: MockInstance;
let processStderrSpy: MockInstance;
@@ -123,7 +120,6 @@ describe('runNonInteractive', () => {
beforeEach(async () => {
mockCoreExecuteToolCall = vi.mocked(executeToolCall);
mockShutdownTelemetry = vi.mocked(shutdownTelemetry);
mockCommandServiceCreate.mockResolvedValue({
getCommands: mockGetCommands,
@@ -247,7 +243,8 @@ describe('runNonInteractive', () => {
'prompt-id-1',
);
expect(getWrittenOutput()).toBe('Hello World\n');
expect(mockShutdownTelemetry).toHaveBeenCalled();
// Note: Telemetry shutdown is now handled in runExitCleanup() in cleanup.ts
// so we no longer expect shutdownTelemetry to be called directly here
});
it('should handle a single tool call and respond', async () => {
-5
View File
@@ -15,8 +15,6 @@ import { isSlashCommand } from './ui/utils/commandUtils.js';
import type { LoadedSettings } from './config/settings.js';
import {
executeToolCall,
shutdownTelemetry,
isTelemetrySdkInitialized,
GeminiEventType,
FatalInputError,
promptIdContext,
@@ -445,9 +443,6 @@ export async function runNonInteractive({
consolePatcher.cleanup();
coreEvents.off(CoreEvent.UserFeedback, handleUserFeedback);
if (isTelemetrySdkInitialized()) {
await shutdownTelemetry(config);
}
}
if (errorToHandle) {
+23 -1
View File
@@ -59,6 +59,10 @@ import {
disableLineWrapping,
shouldEnterAlternateScreen,
startupProfiler,
SessionStartSource,
SessionEndReason,
fireSessionStartHook,
fireSessionEndHook,
} from '@google/gemini-cli-core';
import { validateAuthMethod } from '../config/auth.js';
import process from 'node:process';
@@ -284,14 +288,32 @@ export const AppContainer = (props: AppContainerProps) => {
await config.initialize();
setConfigInitialized(true);
startupProfiler.flush(config);
// 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();
const hookMessageBus = config.getMessageBus();
if (hooksEnabled && hookMessageBus) {
const sessionStartSource = resumedSessionData
? SessionStartSource.Resume
: SessionStartSource.Startup;
await fireSessionStartHook(hookMessageBus, sessionStartSource);
}
})();
registerCleanup(async () => {
// Turn off mouse scroll.
disableMouseEvents();
const ideClient = await IdeClient.getInstance();
await ideClient.disconnect();
// Fire SessionEnd hook on cleanup (only if hooks are enabled)
const hooksEnabled = config.getEnableHooks();
const hookMessageBus = config.getMessageBus();
if (hooksEnabled && hookMessageBus) {
await fireSessionEndHook(hookMessageBus, SessionEndReason.Exit);
}
});
}, [config]);
}, [config, resumedSessionData]);
useEffect(
() => setUpdateHandler(historyManager.addItem, setUpdateInfo),
@@ -44,6 +44,8 @@ describe('clearCommand', () => {
}),
}) as unknown as GeminiClient,
setSessionId: vi.fn(),
getEnableHooks: vi.fn().mockReturnValue(false),
getMessageBus: vi.fn().mockReturnValue(undefined),
},
},
});
+29 -1
View File
@@ -4,7 +4,14 @@
* SPDX-License-Identifier: Apache-2.0
*/
import { uiTelemetryService } from '@google/gemini-cli-core';
import {
uiTelemetryService,
fireSessionEndHook,
fireSessionStartHook,
SessionEndReason,
SessionStartSource,
flushTelemetry,
} from '@google/gemini-cli-core';
import type { SlashCommand } from './types.js';
import { CommandKind } from './types.js';
import { randomUUID } from 'node:crypto';
@@ -21,6 +28,12 @@ export const clearCommand: SlashCommand = {
?.getGeminiClient()
?.getChat()
.getChatRecordingService();
const messageBus = config?.getMessageBus();
// Fire SessionEnd hook before clearing
if (config?.getEnableHooks() && messageBus) {
await fireSessionEndHook(messageBus, SessionEndReason.Clear);
}
if (geminiClient) {
context.ui.setDebugMessage('Clearing terminal and resetting chat.');
@@ -38,6 +51,21 @@ export const clearCommand: SlashCommand = {
chatRecordingService.initialize();
}
// Fire SessionStart hook after clearing
if (config?.getEnableHooks() && messageBus) {
await fireSessionStartHook(messageBus, SessionStartSource.Clear);
}
// Give the event loop a chance to process any pending telemetry operations
// This ensures logger.emit() calls have fully propagated to the BatchLogRecordProcessor
await new Promise((resolve) => setImmediate(resolve));
// Flush telemetry to ensure hooks are written to disk immediately
// This is critical for tests and environments with I/O latency
if (config) {
await flushTelemetry(config);
}
uiTelemetryService.setLastPromptTokenCount(0);
context.ui.clear();
},
+25 -1
View File
@@ -6,10 +6,16 @@
import { promises as fs } from 'node:fs';
import { join } from 'node:path';
import { Storage } from '@google/gemini-cli-core';
import {
Storage,
shutdownTelemetry,
isTelemetrySdkInitialized,
} from '@google/gemini-cli-core';
import type { Config } from '@google/gemini-cli-core';
const cleanupFunctions: Array<(() => void) | (() => Promise<void>)> = [];
const syncCleanupFunctions: Array<() => void> = [];
let configForTelemetry: Config | null = null;
export function registerCleanup(fn: (() => void) | (() => Promise<void>)) {
cleanupFunctions.push(fn);
@@ -30,6 +36,14 @@ export function runSyncCleanup() {
syncCleanupFunctions.length = 0;
}
/**
* Register the config instance for telemetry shutdown.
* This must be called early in the application lifecycle.
*/
export function registerTelemetryConfig(config: Config) {
configForTelemetry = config;
}
export async function runExitCleanup() {
runSyncCleanup();
for (const fn of cleanupFunctions) {
@@ -40,6 +54,16 @@ export async function runExitCleanup() {
}
}
cleanupFunctions.length = 0; // Clear the array
// IMPORTANT: Shutdown telemetry AFTER all other cleanup functions have run
// This ensures SessionEnd hooks and other telemetry are properly flushed
if (configForTelemetry && isTelemetrySdkInitialized()) {
try {
await shutdownTelemetry(configForTelemetry);
} catch (_) {
// Ignore errors during telemetry shutdown
}
}
}
export async function cleanupCheckpoints() {