From f0cc28f80fe5bf7462005067fc279b084e5ec5e9 Mon Sep 17 00:00:00 2001 From: DennisYu07 <617072224@qq.com> Date: Tue, 3 Mar 2026 23:39:57 -0800 Subject: [PATCH] implementation SessionStart and SessionEnd hook --- .../hook-integration/hooks.test.ts | 1395 +++++++++++++++++ packages/cli/src/ui/AppContainer.tsx | 37 + .../cli/src/ui/commands/clearCommand.test.ts | 62 + packages/cli/src/ui/commands/clearCommand.ts | 27 +- .../core/src/hooks/hookEventHandler.test.ts | 237 ++- packages/core/src/hooks/hookEventHandler.ts | 42 + packages/core/src/hooks/hookSystem.test.ts | 139 ++ packages/core/src/hooks/hookSystem.ts | 32 + packages/core/src/hooks/types.ts | 10 +- .../services/chatCompressionService.test.ts | 114 ++ .../src/services/chatCompressionService.ts | 11 + 11 files changed, 2099 insertions(+), 7 deletions(-) diff --git a/integration-tests/hook-integration/hooks.test.ts b/integration-tests/hook-integration/hooks.test.ts index f134dc1ab..1b1ba8468 100644 --- a/integration-tests/hook-integration/hooks.test.ts +++ b/integration-tests/hook-integration/hooks.test.ts @@ -7,6 +7,8 @@ import { TestRig, validateModelOutput } from '../test-helper.js'; * Tests for complete hook system flow including: * - UserPromptSubmit hooks: Triggered before prompt is sent to LLM * - Stop hooks: Triggered when agent is about to stop + * - SessionStart hooks: Triggered when a new session starts (Startup, Resume, Clear, Compact) + * - SessionEnd hooks: Triggered when a session ends (Clear, Logout, PromptInputExit) * * Test categories: * - Single hook scenarios (allow, block, modify, context, etc.) @@ -1835,6 +1837,1059 @@ console.log(JSON.stringify({ }); }); + // ========================================================================== + // SessionStart Hooks + // Triggered when a new session starts (Startup, Resume, Clear, Compact) + // ========================================================================== + describe('SessionStart Hooks', () => { + describe('Allow Decision', () => { + it('should allow session start when hook returns allow decision (Startup source)', async () => { + const allowScript = `console.log(JSON.stringify({decision: 'allow', reason: 'Session startup approved'}));`; + + await rig.setup('session-start-allow-startup', { + settings: { + hooks: { + enabled: true, + SessionStart: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${allowScript}"`, + name: 'session-start-allow-hook', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say hello'); + expect(result).toBeDefined(); + expect(result.length).toBeGreaterThan(0); + }); + + it('should allow session start with additional context', async () => { + const contextScript = `console.log(JSON.stringify({decision: 'allow', hookSpecificOutput: {additionalContext: 'Session context from hook'}}));`; + + await rig.setup('session-start-add-context', { + settings: { + hooks: { + enabled: true, + SessionStart: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${contextScript}"`, + name: 'session-start-context-hook', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say context test'); + expect(result).toBeDefined(); + }); + }); + + describe('Block Decision', () => { + it('should block session start when hook returns block decision', async () => { + const blockScript = `console.log(JSON.stringify({decision: 'block', reason: 'Session blocked by security policy'}));`; + + await rig.setup('session-start-block-decision', { + settings: { + hooks: { + enabled: true, + SessionStart: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${blockScript}"`, + name: 'session-start-block-hook', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say hello'); + expect(result).toBeDefined(); + expect(result.toLowerCase()).toContain('block'); + }); + + it('should block session start with custom reason', async () => { + const blockReasonScript = `console.log(JSON.stringify({decision: 'block', reason: 'Custom block reason: unauthorized user'}));`; + + await rig.setup('session-start-block-custom-reason', { + settings: { + hooks: { + enabled: true, + SessionStart: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${blockReasonScript}"`, + name: 'session-start-block-reason-hook', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say test'); + expect(result).toBeDefined(); + expect(result.toLowerCase()).toContain('block'); + }); + }); + + describe('System Message', () => { + it('should include system message when hook provides it', async () => { + const systemMsgScript = `console.log(JSON.stringify({decision: 'allow', systemMessage: 'System message from SessionStart hook'}));`; + + await rig.setup('session-start-system-message', { + settings: { + hooks: { + enabled: true, + SessionStart: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${systemMsgScript}"`, + name: 'session-start-system-msg-hook', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say system message test'); + expect(result).toBeDefined(); + }); + }); + + describe('Input Format Validation', () => { + it('should receive properly formatted input with session start source', async () => { + const inputValidationScript = ` +const input = JSON.parse(process.argv[2] || '{}'); +const hasRequired = input.session_id && input.cwd && input.hook_event_name && input.source && input.model; +console.log(JSON.stringify({ + decision: 'allow', + hookSpecificOutput: { + additionalContext: hasRequired ? 'Valid SessionStart input: ' + input.source : 'Invalid input format' + } +})); +`; + + await rig.setup('session-start-correct-input', { + settings: { + hooks: { + enabled: true, + SessionStart: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${inputValidationScript.replace(/\n/g, ' ')}"`, + name: 'session-start-input-hook', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say input test'); + expect(result).toBeDefined(); + }); + }); + + describe('Timeout Handling', () => { + it('should continue session start when hook times out', async () => { + await rig.setup('session-start-timeout', { + settings: { + hooks: { + enabled: true, + SessionStart: [ + { + hooks: [ + { + type: 'command', + command: 'sleep 60', + name: 'session-start-timeout-hook', + timeout: 1000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say timeout test'); + expect(result).toBeDefined(); + }); + }); + + describe('Error Handling', () => { + it('should continue session start when hook exits with non-blocking error', async () => { + await rig.setup('session-start-nonblocking-error', { + settings: { + hooks: { + enabled: true, + SessionStart: [ + { + hooks: [ + { + type: 'command', + command: 'echo warning && exit 1', + name: 'session-start-error-hook', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say error test'); + expect(result).toBeDefined(); + }); + + it('should continue session start when hook command does not exist', async () => { + await rig.setup('session-start-missing-command', { + settings: { + hooks: { + enabled: true, + SessionStart: [ + { + hooks: [ + { + type: 'command', + command: '/nonexistent/session/start/command', + name: 'session-start-missing-hook', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say missing test'); + expect(result).toBeDefined(); + }); + }); + + describe('Multiple SessionStart Hooks', () => { + it('should block when one of multiple parallel hooks returns block', async () => { + const allowScript = `console.log(JSON.stringify({decision: 'allow', reason: 'Allowed'}));`; + const blockScript = `console.log(JSON.stringify({decision: 'block', reason: 'Blocked by security policy'}));`; + + await rig.setup('session-start-multi-one-blocks', { + settings: { + hooks: { + enabled: true, + SessionStart: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${allowScript}"`, + name: 'session-start-allow-hook', + timeout: 5000, + }, + { + type: 'command', + command: `node -e "${blockScript}"`, + name: 'session-start-block-hook', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say hello'); + expect(result).toBeDefined(); + expect(result.toLowerCase()).toContain('block'); + }); + + it('should block when first sequential hook returns block', async () => { + const blockScript = `console.log(JSON.stringify({decision: 'block', reason: 'First hook blocks'}));`; + const allowScript = `console.log(JSON.stringify({decision: 'allow'}));`; + + await rig.setup('session-start-seq-first-blocks', { + settings: { + hooks: { + enabled: true, + SessionStart: [ + { + sequential: true, + hooks: [ + { + type: 'command', + command: `node -e "${blockScript}"`, + name: 'session-start-seq-block-hook', + timeout: 5000, + }, + { + type: 'command', + command: `node -e "${allowScript}"`, + name: 'session-start-seq-allow-hook', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say test'); + expect(result).toBeDefined(); + expect(result.toLowerCase()).toContain('block'); + }); + + it('should handle multiple hooks all returning allow', async () => { + const allow1Script = `console.log(JSON.stringify({decision: 'allow', reason: 'First allows'}));`; + const allow2Script = `console.log(JSON.stringify({decision: 'allow', reason: 'Second allows'}));`; + + await rig.setup('session-start-multi-all-allow', { + settings: { + hooks: { + enabled: true, + SessionStart: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${allow1Script}"`, + name: 'session-start-allow-1', + timeout: 5000, + }, + { + type: 'command', + command: `node -e "${allow2Script}"`, + name: 'session-start-allow-2', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say hello'); + expect(result).toBeDefined(); + expect(result.length).toBeGreaterThan(0); + }); + + it('should concatenate additional context from multiple hooks', async () => { + const context1Script = `console.log(JSON.stringify({decision: 'allow', hookSpecificOutput: {additionalContext: 'context from session start hook 1'}}));`; + const context2Script = `console.log(JSON.stringify({decision: 'allow', hookSpecificOutput: {additionalContext: 'context from session start hook 2'}}));`; + + await rig.setup('session-start-multi-context', { + settings: { + hooks: { + enabled: true, + SessionStart: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${context1Script}"`, + name: 'session-start-context-1', + timeout: 5000, + }, + { + type: 'command', + command: `node -e "${context2Script}"`, + name: 'session-start-context-2', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say hello'); + expect(result).toBeDefined(); + }); + + it('should handle hook with error alongside blocking hook', async () => { + const blockScript = `console.log(JSON.stringify({decision: 'block', reason: 'Blocked'}));`; + + await rig.setup('session-start-error-with-block', { + settings: { + hooks: { + enabled: true, + SessionStart: [ + { + hooks: [ + { + type: 'command', + command: '/nonexistent/command', + name: 'session-start-error-hook', + timeout: 5000, + }, + { + type: 'command', + command: `node -e "${blockScript}"`, + name: 'session-start-block-hook', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say test'); + expect(result).toBeDefined(); + expect(result.toLowerCase()).toContain('block'); + }); + + it('should handle hook timeout alongside blocking hook', async () => { + const blockScript = `console.log(JSON.stringify({decision: 'block', reason: 'Blocked while other times out'}));`; + + await rig.setup('session-start-timeout-with-block', { + settings: { + hooks: { + enabled: true, + SessionStart: [ + { + hooks: [ + { + type: 'command', + command: 'sleep 60', + name: 'session-start-timeout-hook', + timeout: 1000, + }, + { + type: 'command', + command: `node -e "${blockScript}"`, + name: 'session-start-block-hook', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say test'); + expect(result).toBeDefined(); + expect(result.toLowerCase()).toContain('block'); + }); + + it('should handle system messages from multiple hooks', async () => { + const msg1Script = `console.log(JSON.stringify({decision: 'allow', systemMessage: 'System message 1 from SessionStart'}));`; + const msg2Script = `console.log(JSON.stringify({decision: 'allow', systemMessage: 'System message 2 from SessionStart'}));`; + + await rig.setup('session-start-multi-system-msg', { + settings: { + hooks: { + enabled: true, + SessionStart: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${msg1Script}"`, + name: 'session-start-msg-1', + timeout: 5000, + }, + { + type: 'command', + command: `node -e "${msg2Script}"`, + name: 'session-start-msg-2', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say hello'); + expect(result).toBeDefined(); + }); + }); + }); + + // ========================================================================== + // SessionEnd Hooks + // Triggered when a session ends (Clear, Logout, PromptInputExit) + // ========================================================================== + describe('SessionEnd Hooks', () => { + describe('Allow Decision', () => { + it('should allow session end when hook returns allow decision', async () => { + const allowScript = `console.log(JSON.stringify({decision: 'allow', reason: 'Session end approved'}));`; + + await rig.setup('session-end-allow', { + settings: { + hooks: { + enabled: true, + SessionEnd: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${allowScript}"`, + name: 'session-end-allow-hook', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say hello'); + expect(result).toBeDefined(); + }); + + it('should allow session end with additional context', async () => { + const contextScript = `console.log(JSON.stringify({decision: 'allow', hookSpecificOutput: {additionalContext: 'Session end context from hook'}}));`; + + await rig.setup('session-end-add-context', { + settings: { + hooks: { + enabled: true, + SessionEnd: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${contextScript}"`, + name: 'session-end-context-hook', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say context test'); + expect(result).toBeDefined(); + }); + }); + + describe('Block Decision', () => { + it('should block session end when hook returns block decision', async () => { + const blockScript = `console.log(JSON.stringify({decision: 'block', reason: 'Session end blocked by security policy'}));`; + + await rig.setup('session-end-block-decision', { + settings: { + hooks: { + enabled: true, + SessionEnd: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${blockScript}"`, + name: 'session-end-block-hook', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say hello'); + expect(result).toBeDefined(); + expect(result.toLowerCase()).toContain('block'); + }); + + it('should block session end with custom reason', async () => { + const blockReasonScript = `console.log(JSON.stringify({decision: 'block', reason: 'Custom block reason: session audit required'}));`; + + await rig.setup('session-end-block-custom-reason', { + settings: { + hooks: { + enabled: true, + SessionEnd: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${blockReasonScript}"`, + name: 'session-end-block-reason-hook', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say test'); + expect(result).toBeDefined(); + expect(result.toLowerCase()).toContain('block'); + }); + }); + + describe('System Message', () => { + it('should include system message when hook provides it', async () => { + const systemMsgScript = `console.log(JSON.stringify({decision: 'allow', systemMessage: 'System message from SessionEnd hook'}));`; + + await rig.setup('session-end-system-message', { + settings: { + hooks: { + enabled: true, + SessionEnd: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${systemMsgScript}"`, + name: 'session-end-system-msg-hook', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say system message test'); + expect(result).toBeDefined(); + }); + }); + + describe('Input Format Validation', () => { + it('should receive properly formatted input with session end reason', async () => { + const inputValidationScript = ` +const input = JSON.parse(process.argv[2] || '{}'); +const hasRequired = input.session_id && input.cwd && input.hook_event_name && input.reason; +console.log(JSON.stringify({ + decision: 'allow', + hookSpecificOutput: { + additionalContext: hasRequired ? 'Valid SessionEnd input: ' + input.reason : 'Invalid input format' + } +})); +`; + + await rig.setup('session-end-correct-input', { + settings: { + hooks: { + enabled: true, + SessionEnd: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${inputValidationScript.replace(/\n/g, ' ')}"`, + name: 'session-end-input-hook', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say input test'); + expect(result).toBeDefined(); + }); + }); + + describe('Timeout Handling', () => { + it('should continue session end when hook times out', async () => { + await rig.setup('session-end-timeout', { + settings: { + hooks: { + enabled: true, + SessionEnd: [ + { + hooks: [ + { + type: 'command', + command: 'sleep 60', + name: 'session-end-timeout-hook', + timeout: 1000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say timeout test'); + expect(result).toBeDefined(); + }); + }); + + describe('Error Handling', () => { + it('should continue session end when hook exits with non-blocking error', async () => { + await rig.setup('session-end-nonblocking-error', { + settings: { + hooks: { + enabled: true, + SessionEnd: [ + { + hooks: [ + { + type: 'command', + command: 'echo warning && exit 1', + name: 'session-end-error-hook', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say error test'); + expect(result).toBeDefined(); + }); + + it('should continue session end when hook command does not exist', async () => { + await rig.setup('session-end-missing-command', { + settings: { + hooks: { + enabled: true, + SessionEnd: [ + { + hooks: [ + { + type: 'command', + command: '/nonexistent/session/end/command', + name: 'session-end-missing-hook', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say missing test'); + expect(result).toBeDefined(); + }); + }); + + describe('Multiple SessionEnd Hooks', () => { + it('should block when one of multiple parallel hooks returns block', async () => { + const allowScript = `console.log(JSON.stringify({decision: 'allow', reason: 'Allowed'}));`; + const blockScript = `console.log(JSON.stringify({decision: 'block', reason: 'Blocked by security policy'}));`; + + await rig.setup('session-end-multi-one-blocks', { + settings: { + hooks: { + enabled: true, + SessionEnd: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${allowScript}"`, + name: 'session-end-allow-hook', + timeout: 5000, + }, + { + type: 'command', + command: `node -e "${blockScript}"`, + name: 'session-end-block-hook', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say hello'); + expect(result).toBeDefined(); + expect(result.toLowerCase()).toContain('block'); + }); + + it('should block when first sequential hook returns block', async () => { + const blockScript = `console.log(JSON.stringify({decision: 'block', reason: 'First hook blocks'}));`; + const allowScript = `console.log(JSON.stringify({decision: 'allow'}));`; + + await rig.setup('session-end-seq-first-blocks', { + settings: { + hooks: { + enabled: true, + SessionEnd: [ + { + sequential: true, + hooks: [ + { + type: 'command', + command: `node -e "${blockScript}"`, + name: 'session-end-seq-block-hook', + timeout: 5000, + }, + { + type: 'command', + command: `node -e "${allowScript}"`, + name: 'session-end-seq-allow-hook', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say test'); + expect(result).toBeDefined(); + expect(result.toLowerCase()).toContain('block'); + }); + + it('should handle multiple hooks all returning allow', async () => { + const allow1Script = `console.log(JSON.stringify({decision: 'allow', reason: 'First allows'}));`; + const allow2Script = `console.log(JSON.stringify({decision: 'allow', reason: 'Second allows'}));`; + + await rig.setup('session-end-multi-all-allow', { + settings: { + hooks: { + enabled: true, + SessionEnd: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${allow1Script}"`, + name: 'session-end-allow-1', + timeout: 5000, + }, + { + type: 'command', + command: `node -e "${allow2Script}"`, + name: 'session-end-allow-2', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say hello'); + expect(result).toBeDefined(); + expect(result.length).toBeGreaterThan(0); + }); + + it('should concatenate additional context from multiple hooks', async () => { + const context1Script = `console.log(JSON.stringify({decision: 'allow', hookSpecificOutput: {additionalContext: 'context from session end hook 1'}}));`; + const context2Script = `console.log(JSON.stringify({decision: 'allow', hookSpecificOutput: {additionalContext: 'context from session end hook 2'}}));`; + + await rig.setup('session-end-multi-context', { + settings: { + hooks: { + enabled: true, + SessionEnd: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${context1Script}"`, + name: 'session-end-context-1', + timeout: 5000, + }, + { + type: 'command', + command: `node -e "${context2Script}"`, + name: 'session-end-context-2', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say hello'); + expect(result).toBeDefined(); + }); + + it('should handle hook with error alongside blocking hook', async () => { + const blockScript = `console.log(JSON.stringify({decision: 'block', reason: 'Blocked'}));`; + + await rig.setup('session-end-error-with-block', { + settings: { + hooks: { + enabled: true, + SessionEnd: [ + { + hooks: [ + { + type: 'command', + command: '/nonexistent/command', + name: 'session-end-error-hook', + timeout: 5000, + }, + { + type: 'command', + command: `node -e "${blockScript}"`, + name: 'session-end-block-hook', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say test'); + expect(result).toBeDefined(); + expect(result.toLowerCase()).toContain('block'); + }); + + it('should handle hook timeout alongside blocking hook', async () => { + const blockScript = `console.log(JSON.stringify({decision: 'block', reason: 'Blocked while other times out'}));`; + + await rig.setup('session-end-timeout-with-block', { + settings: { + hooks: { + enabled: true, + SessionEnd: [ + { + hooks: [ + { + type: 'command', + command: 'sleep 60', + name: 'session-end-timeout-hook', + timeout: 1000, + }, + { + type: 'command', + command: `node -e "${blockScript}"`, + name: 'session-end-block-hook', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say test'); + expect(result).toBeDefined(); + expect(result.toLowerCase()).toContain('block'); + }); + + it('should handle system messages from multiple hooks', async () => { + const msg1Script = `console.log(JSON.stringify({decision: 'allow', systemMessage: 'System message 1 from SessionEnd'}));`; + const msg2Script = `console.log(JSON.stringify({decision: 'allow', systemMessage: 'System message 2 from SessionEnd'}));`; + + await rig.setup('session-end-multi-system-msg', { + settings: { + hooks: { + enabled: true, + SessionEnd: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${msg1Script}"`, + name: 'session-end-msg-1', + timeout: 5000, + }, + { + type: 'command', + command: `node -e "${msg2Script}"`, + name: 'session-end-msg-2', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say hello'); + expect(result).toBeDefined(); + }); + }); + }); + // ========================================================================== // Combined Hooks // Tests for using multiple hook types (UserPromptSubmit + Stop) together @@ -1880,6 +2935,346 @@ console.log(JSON.stringify({ const result = await rig.run('Say both hooks'); expect(result).toBeDefined(); }); + + it('should execute SessionStart, SessionEnd, UserPromptSubmit, and Stop hooks in same session', async () => { + const sessionStartScript = `console.log(JSON.stringify({decision: 'allow', hookSpecificOutput: {additionalContext: 'Session start hook executed'}}));`; + const sessionEndScript = `console.log(JSON.stringify({decision: 'allow', hookSpecificOutput: {additionalContext: 'Session end hook executed'}}));`; + const upsScript = `console.log(JSON.stringify({decision: 'allow', hookSpecificOutput: {additionalContext: 'UPS hook executed'}}));`; + const stopScript = `console.log(JSON.stringify({decision: 'allow', hookSpecificOutput: {additionalContext: 'Stop hook executed'}}));`; + + await rig.setup('combined-all-hooks', { + settings: { + hooks: { + enabled: true, + SessionStart: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${sessionStartScript}"`, + name: 'session-start-hook', + timeout: 5000, + }, + ], + }, + ], + SessionEnd: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${sessionEndScript}"`, + name: 'session-end-hook', + timeout: 5000, + }, + ], + }, + ], + UserPromptSubmit: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${upsScript}"`, + name: 'ups-hook', + timeout: 5000, + }, + ], + }, + ], + Stop: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${stopScript}"`, + name: 'stop-hook', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say all hooks test'); + expect(result).toBeDefined(); + }); + + it('should block session when SessionStart hook returns block', async () => { + const sessionStartBlockScript = `console.log(JSON.stringify({decision: 'block', reason: 'Session start blocked'}));`; + const upsScript = `console.log(JSON.stringify({decision: 'allow'}));`; + + await rig.setup('combined-session-start-blocks', { + settings: { + hooks: { + enabled: true, + SessionStart: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${sessionStartBlockScript}"`, + name: 'session-start-block-hook', + timeout: 5000, + }, + ], + }, + ], + UserPromptSubmit: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${upsScript}"`, + name: 'ups-hook', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say test'); + expect(result).toBeDefined(); + expect(result.toLowerCase()).toContain('block'); + }); + + it('should block session when SessionEnd hook returns block', async () => { + const sessionEndBlockScript = `console.log(JSON.stringify({decision: 'block', reason: 'Session end blocked'}));`; + const upsScript = `console.log(JSON.stringify({decision: 'allow'}));`; + + await rig.setup('combined-session-end-blocks', { + settings: { + hooks: { + enabled: true, + SessionEnd: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${sessionEndBlockScript}"`, + name: 'session-end-block-hook', + timeout: 5000, + }, + ], + }, + ], + UserPromptSubmit: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${upsScript}"`, + name: 'ups-hook', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say test'); + expect(result).toBeDefined(); + expect(result.toLowerCase()).toContain('block'); + }); + + it('should handle multiple hooks of different types all returning allow', async () => { + const sessionStartScript = `console.log(JSON.stringify({decision: 'allow', hookSpecificOutput: {additionalContext: 'Session start allows'}}));`; + const sessionEndScript = `console.log(JSON.stringify({decision: 'allow', hookSpecificOutput: {additionalContext: 'Session end allows'}}));`; + const upsScript = `console.log(JSON.stringify({decision: 'allow', hookSpecificOutput: {additionalContext: 'UPS allows'}}));`; + const stopScript = `console.log(JSON.stringify({decision: 'allow', hookSpecificOutput: {additionalContext: 'Stop allows'}}));`; + + await rig.setup('combined-multi-all-allow', { + settings: { + hooks: { + enabled: true, + SessionStart: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${sessionStartScript}"`, + name: 'session-start-allow', + timeout: 5000, + }, + ], + }, + ], + SessionEnd: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${sessionEndScript}"`, + name: 'session-end-allow', + timeout: 5000, + }, + ], + }, + ], + UserPromptSubmit: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${upsScript}"`, + name: 'ups-allow', + timeout: 5000, + }, + ], + }, + ], + Stop: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${stopScript}"`, + name: 'stop-allow', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say all allow test'); + expect(result).toBeDefined(); + expect(result.length).toBeGreaterThan(0); + }); + + it('should handle error in one hook type while others succeed', async () => { + const sessionStartErrorScript = `node -e "console.log(JSON.stringify({decision: 'allow'})); process.exit(1)"`; + const upsScript = `console.log(JSON.stringify({decision: 'allow'}));`; + const stopScript = `console.log(JSON.stringify({decision: 'allow'}));`; + + await rig.setup('combined-error-one-type', { + settings: { + hooks: { + enabled: true, + SessionStart: [ + { + hooks: [ + { + type: 'command', + command: sessionStartErrorScript, + name: 'session-start-error-hook', + timeout: 5000, + }, + ], + }, + ], + UserPromptSubmit: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${upsScript}"`, + name: 'ups-hook', + timeout: 5000, + }, + ], + }, + ], + Stop: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${stopScript}"`, + name: 'stop-hook', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say error test'); + expect(result).toBeDefined(); + }); + + it('should concatenate additional context from all hook types', async () => { + const sessionStartScript = `console.log(JSON.stringify({decision: 'allow', hookSpecificOutput: {additionalContext: 'Context from SessionStart'}}));`; + const sessionEndScript = `console.log(JSON.stringify({decision: 'allow', hookSpecificOutput: {additionalContext: 'Context from SessionEnd'}}));`; + const upsScript = `console.log(JSON.stringify({decision: 'allow', hookSpecificOutput: {additionalContext: 'Context from UPS'}}));`; + const stopScript = `console.log(JSON.stringify({decision: 'allow', hookSpecificOutput: {additionalContext: 'Context from Stop'}}));`; + + await rig.setup('combined-all-context', { + settings: { + hooks: { + enabled: true, + SessionStart: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${sessionStartScript}"`, + name: 'session-start-context', + timeout: 5000, + }, + ], + }, + ], + SessionEnd: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${sessionEndScript}"`, + name: 'session-end-context', + timeout: 5000, + }, + ], + }, + ], + UserPromptSubmit: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${upsScript}"`, + name: 'ups-context', + timeout: 5000, + }, + ], + }, + ], + Stop: [ + { + hooks: [ + { + type: 'command', + command: `node -e "${stopScript}"`, + name: 'stop-context', + timeout: 5000, + }, + ], + }, + ], + }, + trusted: true, + }, + }); + + const result = await rig.run('Say context test'); + expect(result).toBeDefined(); + }); }); // ========================================================================== diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index 781aab375..3738f55df 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/packages/cli/src/ui/AppContainer.tsx @@ -39,6 +39,8 @@ import { getAllGeminiMdFilenames, ShellExecutionService, Storage, + SessionEndReason, + SessionStartSource, } from '@qwen-code/qwen-code-core'; import { buildResumedHistoryItems } from './utils/resumeHistoryUtils.js'; import { validateAuthMethod } from '../config/auth.js'; @@ -287,7 +289,42 @@ export const AppContainer = (props: AppContainerProps) => { ); historyManager.loadHistory(historyItems); } + + // Fire SessionStart event after config is initialized + const sessionStartSource = resumedSessionData + ? SessionStartSource.Resume + : SessionStartSource.Startup; + + const hookSystem = config.getHookSystem(); + + if (hookSystem) { + hookSystem + .fireSessionStartEvent(sessionStartSource, config.getModel() ?? '') + .then(() => { + debugLogger.debug('SessionStart event completed successfully'); + }) + .catch((err) => { + debugLogger.warn(`SessionStart hook failed: ${err}`); + }); + } else { + debugLogger.debug( + 'SessionStart: HookSystem not available, skipping event', + ); + } })(); + + // Register SessionEnd cleanup for process exit + registerCleanup(async () => { + try { + await config + .getHookSystem() + ?.fireSessionEndEvent(SessionEndReason.PromptInputExit); + debugLogger.debug('SessionEnd event completed successfully!!!'); + } catch (err) { + debugLogger.error(`SessionEnd hook failed: ${err}`); + } + }); + registerCleanup(async () => { const ideClient = await IdeClient.getInstance(); await ideClient.disconnect(); diff --git a/packages/cli/src/ui/commands/clearCommand.test.ts b/packages/cli/src/ui/commands/clearCommand.test.ts index e94c974fb..1cbd7c012 100644 --- a/packages/cli/src/ui/commands/clearCommand.test.ts +++ b/packages/cli/src/ui/commands/clearCommand.test.ts @@ -8,6 +8,10 @@ import { vi, describe, it, expect, beforeEach } from 'vitest'; import { clearCommand } from './clearCommand.js'; import { type CommandContext } from './types.js'; import { createMockCommandContext } from '../../test-utils/mockCommandContext.js'; +import { + SessionEndReason, + SessionStartSource, +} from '@qwen-code/qwen-code-core'; // Mock the telemetry service vi.mock('@qwen-code/qwen-code-core', async () => { @@ -26,10 +30,19 @@ describe('clearCommand', () => { let mockContext: CommandContext; let mockResetChat: ReturnType; let mockStartNewSession: ReturnType; + let mockFireSessionEndEvent: ReturnType; + let mockFireSessionStartEvent: ReturnType; + let mockGetHookSystem: ReturnType; beforeEach(() => { mockResetChat = vi.fn().mockResolvedValue(undefined); mockStartNewSession = vi.fn().mockReturnValue('new-session-id'); + mockFireSessionEndEvent = vi.fn().mockResolvedValue(undefined); + mockFireSessionStartEvent = vi.fn().mockResolvedValue(undefined); + mockGetHookSystem = vi.fn().mockReturnValue({ + fireSessionEndEvent: mockFireSessionEndEvent, + fireSessionStartEvent: mockFireSessionStartEvent, + }); vi.clearAllMocks(); mockContext = createMockCommandContext({ @@ -40,6 +53,11 @@ describe('clearCommand', () => { resetChat: mockResetChat, }) as unknown as GeminiClient, startNewSession: mockStartNewSession, + getHookSystem: mockGetHookSystem, + getDebugLogger: () => ({ + warn: vi.fn(), + }), + getModel: () => 'test-model', }, }, session: { @@ -75,6 +93,50 @@ describe('clearCommand', () => { expect(mockContext.ui.clear).toHaveBeenCalled(); }); + it('should fire SessionEnd event before clearing and SessionStart event after clearing', async () => { + if (!clearCommand.action) { + throw new Error('clearCommand must have an action.'); + } + + await clearCommand.action(mockContext, ''); + + expect(mockGetHookSystem).toHaveBeenCalled(); + expect(mockFireSessionEndEvent).toHaveBeenCalledWith( + SessionEndReason.Clear, + ); + expect(mockFireSessionStartEvent).toHaveBeenCalledWith( + SessionStartSource.Clear, + 'test-model', + ); + + // SessionEnd should be called before SessionStart + const sessionEndCallOrder = + mockFireSessionEndEvent.mock.invocationCallOrder[0]; + const sessionStartCallOrder = + mockFireSessionStartEvent.mock.invocationCallOrder[0]; + expect(sessionEndCallOrder).toBeLessThan(sessionStartCallOrder); + }); + + it('should handle hook errors gracefully and continue execution', async () => { + if (!clearCommand.action) { + throw new Error('clearCommand must have an action.'); + } + + mockFireSessionEndEvent.mockRejectedValue( + new Error('SessionEnd hook failed'), + ); + mockFireSessionStartEvent.mockRejectedValue( + new Error('SessionStart hook failed'), + ); + + await clearCommand.action(mockContext, ''); + + // Should still complete the clear operation despite hook errors + expect(mockStartNewSession).toHaveBeenCalledTimes(1); + expect(mockResetChat).toHaveBeenCalledTimes(1); + expect(mockContext.ui.clear).toHaveBeenCalledTimes(1); + }); + it('should not attempt to reset chat if config service is not available', async () => { if (!clearCommand.action) { throw new Error('clearCommand must have an action.'); diff --git a/packages/cli/src/ui/commands/clearCommand.ts b/packages/cli/src/ui/commands/clearCommand.ts index dd774934b..e1a529ceb 100644 --- a/packages/cli/src/ui/commands/clearCommand.ts +++ b/packages/cli/src/ui/commands/clearCommand.ts @@ -7,7 +7,11 @@ import type { SlashCommand } from './types.js'; import { CommandKind } from './types.js'; import { t } from '../../i18n/index.js'; -import { uiTelemetryService } from '@qwen-code/qwen-code-core'; +import { + uiTelemetryService, + SessionEndReason, + SessionStartSource, +} from '@qwen-code/qwen-code-core'; export const clearCommand: SlashCommand = { name: 'clear', @@ -20,6 +24,15 @@ export const clearCommand: SlashCommand = { const { config } = context.services; if (config) { + // Fire SessionEnd event before clearing (current session ends) + try { + await config + .getHookSystem() + ?.fireSessionEndEvent(SessionEndReason.Clear); + } catch (err) { + config.getDebugLogger().warn(`SessionEnd hook failed: ${err}`); + } + const newSessionId = config.startNewSession(); // Reset UI telemetry metrics for the new session @@ -40,6 +53,18 @@ export const clearCommand: SlashCommand = { } else { context.ui.setDebugMessage(t('Starting a new session and clearing.')); } + + // Fire SessionStart event after clearing (new session starts) + try { + await config + .getHookSystem() + ?.fireSessionStartEvent( + SessionStartSource.Clear, + config.getModel() ?? '', + ); + } catch (err) { + config.getDebugLogger().warn(`SessionStart hook failed: ${err}`); + } } else { context.ui.setDebugMessage(t('Starting a new session and clearing.')); } diff --git a/packages/core/src/hooks/hookEventHandler.test.ts b/packages/core/src/hooks/hookEventHandler.test.ts index f556a8c30..82a4d2fe3 100644 --- a/packages/core/src/hooks/hookEventHandler.test.ts +++ b/packages/core/src/hooks/hookEventHandler.test.ts @@ -6,7 +6,15 @@ import { describe, it, expect, vi, beforeEach, type Mock } from 'vitest'; import { HookEventHandler } from './hookEventHandler.js'; -import { HookEventName, HookType, HooksConfigSource } from './types.js'; +import { + HookEventName, + HookType, + HooksConfigSource, + SessionStartSource, + SessionEndReason, + PermissionMode, + AgentType, +} from './types.js'; import type { Config } from '../config/config.js'; import type { HookPlanner, @@ -192,6 +200,204 @@ describe('HookEventHandler', () => { }); }); + describe('fireSessionStartEvent', () => { + it('should execute hooks for SessionStart event', async () => { + const mockPlan = createMockExecutionPlan([]); + const mockAggregated = createMockAggregatedResult(true); + + vi.mocked(mockHookPlanner.createExecutionPlan).mockReturnValue(mockPlan); + vi.mocked(mockHookRunner.executeHooksParallel).mockResolvedValue([]); + vi.mocked(mockHookAggregator.aggregateResults).mockReturnValue( + mockAggregated, + ); + + const result = await hookEventHandler.fireSessionStartEvent( + SessionStartSource.Startup, + 'test-model', + ); + + expect(mockHookPlanner.createExecutionPlan).toHaveBeenCalledWith( + HookEventName.SessionStart, + undefined, + ); + expect(result.success).toBe(true); + }); + + it('should include all session start parameters in the hook input', async () => { + const mockPlan = createMockExecutionPlan([ + { + type: HookType.Command, + command: 'echo test', + source: HooksConfigSource.Project, + }, + ]); + vi.mocked(mockHookPlanner.createExecutionPlan).mockReturnValue(mockPlan); + vi.mocked(mockHookRunner.executeHooksParallel).mockResolvedValue([]); + vi.mocked(mockHookAggregator.aggregateResults).mockReturnValue( + createMockAggregatedResult(true), + ); + + await hookEventHandler.fireSessionStartEvent( + SessionStartSource.Resume, + 'test-model', + PermissionMode.Plan, + AgentType.Bash, + ); + + const mockCalls = (mockHookRunner.executeHooksParallel as Mock).mock + .calls; + const input = mockCalls[0][2] as { + permission_mode: PermissionMode; + source: SessionStartSource; + model: string; + agent_type?: AgentType; + }; + expect(input.permission_mode).toBe(PermissionMode.Plan); + expect(input.source).toBe(SessionStartSource.Resume); + expect(input.model).toBe('test-model'); + expect(input.agent_type).toBe(AgentType.Bash); + }); + + it('should use default permission mode when not provided', async () => { + const mockPlan = createMockExecutionPlan([ + { + type: HookType.Command, + command: 'echo test', + source: HooksConfigSource.Project, + }, + ]); + vi.mocked(mockHookPlanner.createExecutionPlan).mockReturnValue(mockPlan); + vi.mocked(mockHookRunner.executeHooksParallel).mockResolvedValue([]); + vi.mocked(mockHookAggregator.aggregateResults).mockReturnValue( + createMockAggregatedResult(true), + ); + + await hookEventHandler.fireSessionStartEvent( + SessionStartSource.Clear, + 'test-model', + ); + + const mockCalls = (mockHookRunner.executeHooksParallel as Mock).mock + .calls; + const input = mockCalls[0][2] as { + permission_mode: PermissionMode; + }; + expect(input.permission_mode).toBe(PermissionMode.Default); + }); + + it('should handle session start event with undefined agent type', async () => { + const mockPlan = createMockExecutionPlan([ + { + type: HookType.Command, + command: 'echo test', + source: HooksConfigSource.Project, + }, + ]); + vi.mocked(mockHookPlanner.createExecutionPlan).mockReturnValue(mockPlan); + vi.mocked(mockHookRunner.executeHooksParallel).mockResolvedValue([]); + vi.mocked(mockHookAggregator.aggregateResults).mockReturnValue( + createMockAggregatedResult(true), + ); + + await hookEventHandler.fireSessionStartEvent( + SessionStartSource.Compact, + 'test-model', + ); + + const mockCalls = (mockHookRunner.executeHooksParallel as Mock).mock + .calls; + const input = mockCalls[0][2] as { + permission_mode: PermissionMode; + source: SessionStartSource; + model: string; + agent_type?: AgentType; + }; + expect(input.source).toBe(SessionStartSource.Compact); + expect(input.model).toBe('test-model'); + expect(input.agent_type).toBeUndefined(); + }); + }); + + describe('fireSessionEndEvent', () => { + it('should execute hooks for SessionEnd event', async () => { + const mockPlan = createMockExecutionPlan([]); + const mockAggregated = createMockAggregatedResult(true); + + vi.mocked(mockHookPlanner.createExecutionPlan).mockReturnValue(mockPlan); + vi.mocked(mockHookRunner.executeHooksParallel).mockResolvedValue([]); + vi.mocked(mockHookAggregator.aggregateResults).mockReturnValue( + mockAggregated, + ); + + const result = await hookEventHandler.fireSessionEndEvent( + SessionEndReason.Clear, + ); + + expect(mockHookPlanner.createExecutionPlan).toHaveBeenCalledWith( + HookEventName.SessionEnd, + undefined, + ); + expect(result.success).toBe(true); + }); + + it('should include reason in the hook input', async () => { + const mockPlan = createMockExecutionPlan([ + { + type: HookType.Command, + command: 'echo test', + source: HooksConfigSource.Project, + }, + ]); + vi.mocked(mockHookPlanner.createExecutionPlan).mockReturnValue(mockPlan); + vi.mocked(mockHookRunner.executeHooksParallel).mockResolvedValue([]); + vi.mocked(mockHookAggregator.aggregateResults).mockReturnValue( + createMockAggregatedResult(true), + ); + + await hookEventHandler.fireSessionEndEvent(SessionEndReason.Logout); + + const mockCalls = (mockHookRunner.executeHooksParallel as Mock).mock + .calls; + const input = mockCalls[0][2] as { reason: SessionEndReason }; + expect(input.reason).toBe(SessionEndReason.Logout); + }); + + it('should handle different session end reasons', async () => { + const mockPlan = createMockExecutionPlan([ + { + type: HookType.Command, + command: 'echo test', + source: HooksConfigSource.Project, + }, + ]); + vi.mocked(mockHookPlanner.createExecutionPlan).mockReturnValue(mockPlan); + vi.mocked(mockHookRunner.executeHooksParallel).mockResolvedValue([]); + vi.mocked(mockHookAggregator.aggregateResults).mockReturnValue( + createMockAggregatedResult(true), + ); + + // Test all possible session end reasons + const testReasons = [ + SessionEndReason.Clear, + SessionEndReason.Logout, + SessionEndReason.PromptInputExit, + SessionEndReason.Bypass_permissions_disabled, + SessionEndReason.Other, + ]; + + for (const reason of testReasons) { + await hookEventHandler.fireSessionEndEvent(reason); + + const mockCalls = (mockHookRunner.executeHooksParallel as Mock).mock + .calls; + const input = mockCalls[mockCalls.length - 1][2] as { + reason: SessionEndReason; + }; + expect(input.reason).toBe(reason); + } + }); + }); + describe('sequential vs parallel execution', () => { it('should execute hooks sequentially when plan.sequential is true', async () => { const mockPlan = createMockExecutionPlan( @@ -274,5 +480,34 @@ describe('HookEventHandler', () => { expect(result.errors).toHaveLength(1); expect(result.errors[0].message).toBe('Runner error'); }); + + it('should handle errors for SessionStart event', async () => { + vi.mocked(mockHookPlanner.createExecutionPlan).mockImplementation(() => { + throw new Error('SessionStart planner error'); + }); + + const result = await hookEventHandler.fireSessionStartEvent( + SessionStartSource.Startup, + 'test-model', + ); + + expect(result.success).toBe(false); + expect(result.errors).toHaveLength(1); + expect(result.errors[0].message).toBe('SessionStart planner error'); + }); + + it('should handle errors for SessionEnd event', async () => { + vi.mocked(mockHookPlanner.createExecutionPlan).mockImplementation(() => { + throw new Error('SessionEnd planner error'); + }); + + const result = await hookEventHandler.fireSessionEndEvent( + SessionEndReason.Clear, + ); + + expect(result.success).toBe(false); + expect(result.errors).toHaveLength(1); + expect(result.errors[0].message).toBe('SessionEnd planner error'); + }); }); }); diff --git a/packages/core/src/hooks/hookEventHandler.ts b/packages/core/src/hooks/hookEventHandler.ts index 2fd5f2892..95b285d37 100644 --- a/packages/core/src/hooks/hookEventHandler.ts +++ b/packages/core/src/hooks/hookEventHandler.ts @@ -15,7 +15,13 @@ import type { HookExecutionResult, UserPromptSubmitInput, StopInput, + SessionStartInput, + SessionEndInput, + SessionStartSource, + SessionEndReason, + AgentType, } from './types.js'; +import { PermissionMode } from './types.js'; import { createDebugLogger } from '../utils/debugLogger.js'; const debugLogger = createDebugLogger('TRUSTED_HOOKS'); @@ -73,6 +79,42 @@ export class HookEventHandler { return this.executeHooks(HookEventName.Stop, input); } + /** + * Fire a SessionStart event + * Called when a new session starts or resumes + */ + async fireSessionStartEvent( + source: SessionStartSource, + model: string, + permissionMode?: PermissionMode, + agentType?: AgentType, + ): Promise { + const input: SessionStartInput = { + ...this.createBaseInput(HookEventName.SessionStart), + permission_mode: permissionMode ?? PermissionMode.Default, + source, + model, + agent_type: agentType, + }; + + return this.executeHooks(HookEventName.SessionStart, input); + } + + /** + * Fire a SessionEnd event + * Called when a session ends + */ + async fireSessionEndEvent( + reason: SessionEndReason, + ): Promise { + const input: SessionEndInput = { + ...this.createBaseInput(HookEventName.SessionEnd), + reason, + }; + + return this.executeHooks(HookEventName.SessionEnd, input); + } + /** * Execute hooks for a specific event (direct execution without MessageBus) * Used as fallback when MessageBus is not available diff --git a/packages/core/src/hooks/hookSystem.test.ts b/packages/core/src/hooks/hookSystem.test.ts index 51f2d3050..0a77a81ca 100644 --- a/packages/core/src/hooks/hookSystem.test.ts +++ b/packages/core/src/hooks/hookSystem.test.ts @@ -63,6 +63,8 @@ describe('HookSystem', () => { mockHookEventHandler = { fireUserPromptSubmitEvent: vi.fn(), fireStopEvent: vi.fn(), + fireSessionStartEvent: vi.fn(), + fireSessionEndEvent: vi.fn(), } as unknown as HookEventHandler; vi.mocked(HookRegistry).mockImplementation(() => mockHookRegistry); @@ -325,4 +327,141 @@ describe('HookSystem', () => { expect(result?.getAdditionalContext()).toBe('Some additional context'); }); }); + + describe('fireSessionStartEvent', () => { + it('should fire session start event and return output', async () => { + const mockResult = { + success: true, + allOutputs: [], + errors: [], + totalDuration: 50, + finalOutput: { + continue: true, + decision: 'allow' as HookDecision, + }, + }; + vi.mocked(mockHookEventHandler.fireSessionStartEvent).mockResolvedValue( + mockResult, + ); + + const result = await hookSystem.fireSessionStartEvent('manual', 'gpt-4'); + + expect(mockHookEventHandler.fireSessionStartEvent).toHaveBeenCalledWith( + 'manual', + 'gpt-4', + undefined, + undefined, + ); + expect(result).toBeDefined(); + }); + + it('should pass all parameters to event handler', async () => { + const mockResult = { + success: true, + allOutputs: [], + errors: [], + totalDuration: 0, + finalOutput: { + decision: 'allow' as HookDecision, + }, + }; + vi.mocked(mockHookEventHandler.fireSessionStartEvent).mockResolvedValue( + mockResult, + ); + + await hookSystem.fireSessionStartEvent( + 'api_call', + 'claude-3', + 'auto_edit', // Using actual enum value from PermissionMode + 'chat', + ); + + expect(mockHookEventHandler.fireSessionStartEvent).toHaveBeenCalledWith( + 'api_call', + 'claude-3', + 'auto_edit', + 'chat', + ); + }); + + it('should return undefined when no final output', async () => { + const mockResult = { + success: true, + allOutputs: [], + errors: [], + totalDuration: 0, + finalOutput: undefined, + }; + vi.mocked(mockHookEventHandler.fireSessionStartEvent).mockResolvedValue( + mockResult, + ); + + const result = await hookSystem.fireSessionStartEvent('manual', 'gpt-4'); + + expect(result).toBeUndefined(); + }); + }); + + describe('fireSessionEndEvent', () => { + it('should fire session end event and return output', async () => { + const mockResult = { + success: true, + allOutputs: [], + errors: [], + totalDuration: 50, + finalOutput: { + continue: true, + decision: 'allow' as HookDecision, + }, + }; + vi.mocked(mockHookEventHandler.fireSessionEndEvent).mockResolvedValue( + mockResult, + ); + + const result = await hookSystem.fireSessionEndEvent('user_quit'); + + expect(mockHookEventHandler.fireSessionEndEvent).toHaveBeenCalledWith( + 'user_quit', + ); + expect(result).toBeDefined(); + }); + + it('should pass reason to event handler', async () => { + const mockResult = { + success: true, + allOutputs: [], + errors: [], + totalDuration: 0, + finalOutput: { + decision: 'allow' as HookDecision, + }, + }; + vi.mocked(mockHookEventHandler.fireSessionEndEvent).mockResolvedValue( + mockResult, + ); + + await hookSystem.fireSessionEndEvent('timeout'); + + expect(mockHookEventHandler.fireSessionEndEvent).toHaveBeenCalledWith( + 'timeout', + ); + }); + + it('should return undefined when no final output', async () => { + const mockResult = { + success: true, + allOutputs: [], + errors: [], + totalDuration: 0, + finalOutput: undefined, + }; + vi.mocked(mockHookEventHandler.fireSessionEndEvent).mockResolvedValue( + mockResult, + ); + + const result = await hookSystem.fireSessionEndEvent('normal_exit'); + + expect(result).toBeUndefined(); + }); + }); }); diff --git a/packages/core/src/hooks/hookSystem.ts b/packages/core/src/hooks/hookSystem.ts index 8a40cbd9e..3922cf1ec 100644 --- a/packages/core/src/hooks/hookSystem.ts +++ b/packages/core/src/hooks/hookSystem.ts @@ -14,6 +14,12 @@ import type { HookRegistryEntry } from './hookRegistry.js'; import { createDebugLogger } from '../utils/debugLogger.js'; import type { DefaultHookOutput } from './types.js'; import { createHookOutput } from './types.js'; +import type { + SessionStartSource, + SessionEndReason, + PermissionMode, + AgentType, +} from './types.js'; const debugLogger = createDebugLogger('TRUSTED_HOOKS'); @@ -100,4 +106,30 @@ export class HookSystem { ? createHookOutput('Stop', result.finalOutput) : undefined; } + + async fireSessionStartEvent( + source: SessionStartSource, + model: string, + permissionMode?: PermissionMode, + agentType?: AgentType, + ): Promise { + const result = await this.hookEventHandler.fireSessionStartEvent( + source, + model, + permissionMode, + agentType, + ); + return result.finalOutput + ? createHookOutput('SessionStart', result.finalOutput) + : undefined; + } + + async fireSessionEndEvent( + reason: SessionEndReason, + ): Promise { + const result = await this.hookEventHandler.fireSessionEndEvent(reason); + return result.finalOutput + ? createHookOutput('SessionEnd', result.finalOutput) + : undefined; + } } diff --git a/packages/core/src/hooks/types.ts b/packages/core/src/hooks/types.ts index 49ac7a5ef..bd9883767 100644 --- a/packages/core/src/hooks/types.ts +++ b/packages/core/src/hooks/types.ts @@ -524,18 +524,18 @@ export enum SessionStartSource { export enum PermissionMode { Default = 'default', Plan = 'plan', - AcceptEdit = 'accept_edit', - DontAsk = 'dont_ask', - BypassPermissions = 'bypass_permissions', + AutoEdit = 'auto_edit', + Yolo = 'yolo', } /** * SessionStart hook input */ export interface SessionStartInput extends HookInput { - permission_mode?: PermissionMode; + permission_mode: PermissionMode; source: SessionStartSource; - model?: string; + model: string; + agent_type?: AgentType; } /** diff --git a/packages/core/src/services/chatCompressionService.test.ts b/packages/core/src/services/chatCompressionService.test.ts index 8f19fe9cf..777619172 100644 --- a/packages/core/src/services/chatCompressionService.test.ts +++ b/packages/core/src/services/chatCompressionService.test.ts @@ -16,6 +16,7 @@ import { tokenLimit } from '../core/tokenLimits.js'; import type { GeminiChat } from '../core/geminiChat.js'; import type { Config } from '../config/config.js'; import type { ContentGenerator } from '../core/contentGenerator.js'; +import { SessionStartSource } from '../hooks/types.js'; vi.mock('../telemetry/uiTelemetry.js'); vi.mock('../core/tokenLimits.js'); @@ -107,16 +108,27 @@ describe('ChatCompressionService', () => { let mockConfig: Config; const mockModel = 'gemini-pro'; const mockPromptId = 'test-prompt-id'; + let mockFireSessionStartEvent: ReturnType; + let mockGetHookSystem: ReturnType; beforeEach(() => { service = new ChatCompressionService(); mockChat = { getHistory: vi.fn(), } as unknown as GeminiChat; + mockFireSessionStartEvent = vi.fn().mockResolvedValue(undefined); + mockGetHookSystem = vi.fn().mockReturnValue({ + fireSessionStartEvent: mockFireSessionStartEvent, + }); mockConfig = { getChatCompression: vi.fn(), getContentGenerator: vi.fn(), getContentGeneratorConfig: vi.fn().mockReturnValue({}), + getHookSystem: mockGetHookSystem, + getModel: () => 'test-model', + getDebugLogger: () => ({ + warn: vi.fn(), + }), } as unknown as Config; vi.mocked(tokenLimit).mockReturnValue(1000); @@ -274,6 +286,11 @@ describe('ChatCompressionService', () => { expect(result.newHistory).not.toBeNull(); expect(result.newHistory![0].parts![0].text).toBe('Summary'); expect(mockGenerateContent).toHaveBeenCalled(); + expect(mockGetHookSystem).toHaveBeenCalled(); + expect(mockFireSessionStartEvent).toHaveBeenCalledWith( + SessionStartSource.Compact, + 'test-model', + ); }); it('should force compress even if under threshold', async () => { @@ -317,6 +334,10 @@ describe('ChatCompressionService', () => { expect(result.info.compressionStatus).toBe(CompressionStatus.COMPRESSED); expect(result.newHistory).not.toBeNull(); + expect(mockFireSessionStartEvent).toHaveBeenCalledWith( + SessionStartSource.Compact, + 'test-model', + ); }); it('should return FAILED if new token count is inflated', async () => { @@ -481,4 +502,97 @@ describe('ChatCompressionService', () => { ); expect(result.newHistory).toBeNull(); }); + + it('should not fire SessionStart event when compression fails', async () => { + const history: Content[] = [ + { role: 'user', parts: [{ text: 'msg1' }] }, + { role: 'model', parts: [{ text: 'msg2' }] }, + ]; + vi.mocked(mockChat.getHistory).mockReturnValue(history); + vi.mocked(uiTelemetryService.getLastPromptTokenCount).mockReturnValue(10); + vi.mocked(tokenLimit).mockReturnValue(1000); + + const mockGenerateContent = vi.fn().mockResolvedValue({ + candidates: [ + { + content: { + parts: [{ text: 'Summary' }], + }, + }, + ], + usageMetadata: { + promptTokenCount: 1, + candidatesTokenCount: 20, + totalTokenCount: 21, + }, + } as unknown as GenerateContentResponse); + vi.mocked(mockConfig.getContentGenerator).mockReturnValue({ + generateContent: mockGenerateContent, + } as unknown as ContentGenerator); + + const result = await service.compress( + mockChat, + mockPromptId, + true, + mockModel, + mockConfig, + false, + ); + + expect(result.info.compressionStatus).toBe( + CompressionStatus.COMPRESSION_FAILED_INFLATED_TOKEN_COUNT, + ); + expect(result.newHistory).toBeNull(); + expect(mockFireSessionStartEvent).not.toHaveBeenCalled(); + }); + + it('should handle SessionStart hook errors gracefully', async () => { + const history: Content[] = [ + { role: 'user', parts: [{ text: 'msg1' }] }, + { role: 'model', parts: [{ text: 'msg2' }] }, + { role: 'user', parts: [{ text: 'msg3' }] }, + { role: 'model', parts: [{ text: 'msg4' }] }, + ]; + vi.mocked(mockChat.getHistory).mockReturnValue(history); + vi.mocked(uiTelemetryService.getLastPromptTokenCount).mockReturnValue(800); + vi.mocked(mockConfig.getContentGeneratorConfig).mockReturnValue({ + model: 'gemini-pro', + contextWindowSize: 1000, + } as unknown as ReturnType); + + mockFireSessionStartEvent.mockRejectedValue( + new Error('SessionStart hook failed'), + ); + + const mockGenerateContent = vi.fn().mockResolvedValue({ + candidates: [ + { + content: { + parts: [{ text: 'Summary' }], + }, + }, + ], + usageMetadata: { + promptTokenCount: 1600, + candidatesTokenCount: 50, + totalTokenCount: 1650, + }, + } as unknown as GenerateContentResponse); + vi.mocked(mockConfig.getContentGenerator).mockReturnValue({ + generateContent: mockGenerateContent, + } as unknown as ContentGenerator); + + const result = await service.compress( + mockChat, + mockPromptId, + false, + mockModel, + mockConfig, + false, + ); + + // Should still complete compression despite hook error + expect(result.info.compressionStatus).toBe(CompressionStatus.COMPRESSED); + expect(result.newHistory).not.toBeNull(); + }); }); diff --git a/packages/core/src/services/chatCompressionService.ts b/packages/core/src/services/chatCompressionService.ts index 3a89ee103..53b8d4d10 100644 --- a/packages/core/src/services/chatCompressionService.ts +++ b/packages/core/src/services/chatCompressionService.ts @@ -14,6 +14,7 @@ import { getCompressionPrompt } from '../core/prompts.js'; import { getResponseText } from '../utils/partUtils.js'; import { logChatCompression } from '../telemetry/loggers.js'; import { makeChatCompressionEvent } from '../telemetry/types.js'; +import { SessionStartSource } from '../hooks/types.js'; /** * Threshold for compression token count as a fraction of the model's token limit. @@ -261,6 +262,16 @@ export class ChatCompressionService { }; } else { uiTelemetryService.setLastPromptTokenCount(newTokenCount); + + // Fire SessionStart event after successful compression + try { + await config + .getHookSystem() + ?.fireSessionStartEvent(SessionStartSource.Compact, model ?? ''); + } catch (err) { + config.getDebugLogger().warn(`SessionStart hook failed: ${err}`); + } + return { newHistory: extraHistory, info: {