From ac95282758d27cfbd87e2ade45e0546b56d48afd Mon Sep 17 00:00:00 2001 From: Abhi <43648792+abhipatel12@users.noreply.github.com> Date: Mon, 23 Mar 2026 15:21:49 -0400 Subject: [PATCH] test(core): migrate hook tests to scheduler (#23496) --- .../src/scheduler/scheduler_hooks.test.ts | 305 ++++++++++++++++++ 1 file changed, 305 insertions(+) create mode 100644 packages/core/src/scheduler/scheduler_hooks.test.ts diff --git a/packages/core/src/scheduler/scheduler_hooks.test.ts b/packages/core/src/scheduler/scheduler_hooks.test.ts new file mode 100644 index 0000000000..b59ffc4ace --- /dev/null +++ b/packages/core/src/scheduler/scheduler_hooks.test.ts @@ -0,0 +1,305 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi } from 'vitest'; +import { Scheduler } from './scheduler.js'; +import type { ErroredToolCall } from './types.js'; +import { CoreToolCallStatus } from './types.js'; +import type { Config, ToolRegistry, AgentLoopContext } from '../index.js'; +import { + ApprovalMode, + DEFAULT_TRUNCATE_TOOL_OUTPUT_THRESHOLD, +} from '../index.js'; +import { createMockMessageBus } from '../test-utils/mock-message-bus.js'; +import { MockTool } from '../test-utils/mock-tool.js'; +import { DEFAULT_GEMINI_MODEL } from '../config/models.js'; +import type { PolicyEngine } from '../policy/policy-engine.js'; +import { HookSystem } from '../hooks/hookSystem.js'; +import { HookType, HookEventName } from '../hooks/types.js'; + +function createMockConfig(overrides: Partial = {}): Config { + const defaultToolRegistry = { + getTool: () => undefined, + getToolByName: () => undefined, + getFunctionDeclarations: () => [], + tools: new Map(), + discovery: {}, + registerTool: () => {}, + getToolByDisplayName: () => undefined, + getTools: () => [], + discoverTools: async () => {}, + getAllTools: () => [], + getToolsByServer: () => [], + getExperiments: () => {}, + } as unknown as ToolRegistry; + + const baseConfig = { + getSessionId: () => 'test-session-id', + getUsageStatisticsEnabled: () => true, + getDebugMode: () => false, + isInteractive: () => true, + getApprovalMode: () => ApprovalMode.DEFAULT, + setApprovalMode: () => {}, + getAllowedTools: () => [], + getContentGeneratorConfig: () => ({ + model: 'test-model', + authType: 'oauth-personal', + }), + getShellExecutionConfig: () => ({ + terminalWidth: 90, + terminalHeight: 30, + sanitizationConfig: { + enableEnvironmentVariableRedaction: true, + allowedEnvironmentVariables: [], + blockedEnvironmentVariables: [], + }, + }), + storage: { + getProjectTempDir: () => '/tmp', + }, + getTruncateToolOutputThreshold: () => + DEFAULT_TRUNCATE_TOOL_OUTPUT_THRESHOLD, + getTruncateToolOutputLines: () => 1000, + getToolRegistry: () => defaultToolRegistry, + getWorkingDir: () => '/mock/dir', + getActiveModel: () => DEFAULT_GEMINI_MODEL, + getGeminiClient: () => null, + getMessageBus: () => createMockMessageBus(), + getEnableHooks: () => true, + getExperiments: () => {}, + getPolicyEngine: () => + ({ + check: async () => ({ decision: 'allow' }), + }) as unknown as PolicyEngine, + } as unknown as Config; + + const mockConfig = Object.assign({}, baseConfig, overrides) as Config; + + (mockConfig as { config?: Config }).config = mockConfig; + + return mockConfig; +} + +describe('Scheduler Hooks', () => { + it('should stop execution if BeforeTool hook requests stop', async () => { + const executeFn = vi.fn().mockResolvedValue({ + llmContent: 'Tool executed', + returnDisplay: 'Tool executed', + }); + const mockTool = new MockTool({ name: 'mockTool', execute: executeFn }); + + const toolRegistry = { + getTool: () => mockTool, + getAllToolNames: () => ['mockTool'], + } as unknown as ToolRegistry; + + const mockMessageBus = createMockMessageBus(); + + const mockConfig = createMockConfig({ + getToolRegistry: () => toolRegistry, + getMessageBus: () => mockMessageBus, + getApprovalMode: () => ApprovalMode.YOLO, + }); + + const hookSystem = new HookSystem(mockConfig); + + (mockConfig as { getHookSystem?: () => HookSystem }).getHookSystem = () => + hookSystem; + + // Register a programmatic runtime hook + hookSystem.registerHook( + { + type: HookType.Runtime, + name: 'test-stop-hook', + action: async () => ({ + continue: false, + stopReason: 'Hook stopped execution', + }), + }, + HookEventName.BeforeTool, + ); + + const scheduler = new Scheduler({ + context: { + config: mockConfig, + messageBus: mockMessageBus, + toolRegistry, + } as unknown as AgentLoopContext, + getPreferredEditor: () => 'vscode', + schedulerId: 'test-scheduler', + }); + + const request = { + callId: '1', + name: 'mockTool', + args: {}, + isClientInitiated: false, + prompt_id: 'prompt-1', + }; + + const results = await scheduler.schedule( + [request], + new AbortController().signal, + ); + + expect(results.length).toBe(1); + const result = results[0]; + expect(result.status).toBe(CoreToolCallStatus.Error); + const erroredCall = result as ErroredToolCall; + + expect(erroredCall.response.error?.message).toContain( + 'Agent execution stopped by hook: Hook stopped execution', + ); + expect(executeFn).not.toHaveBeenCalled(); + }); + + it('should block tool execution if BeforeTool hook requests block', async () => { + const executeFn = vi.fn(); + const mockTool = new MockTool({ name: 'mockTool', execute: executeFn }); + + const toolRegistry = { + getTool: () => mockTool, + getAllToolNames: () => ['mockTool'], + } as unknown as ToolRegistry; + + const mockMessageBus = createMockMessageBus(); + + const mockConfig = createMockConfig({ + getToolRegistry: () => toolRegistry, + getMessageBus: () => mockMessageBus, + getApprovalMode: () => ApprovalMode.YOLO, + }); + + const hookSystem = new HookSystem(mockConfig); + + (mockConfig as { getHookSystem?: () => HookSystem }).getHookSystem = () => + hookSystem; + + hookSystem.registerHook( + { + type: HookType.Runtime, + name: 'test-block-hook', + action: async () => ({ + decision: 'block', + reason: 'Hook blocked execution', + }), + }, + HookEventName.BeforeTool, + ); + + const scheduler = new Scheduler({ + context: { + config: mockConfig, + messageBus: mockMessageBus, + toolRegistry, + } as unknown as AgentLoopContext, + getPreferredEditor: () => 'vscode', + schedulerId: 'test-scheduler', + }); + + const request = { + callId: '1', + name: 'mockTool', + args: {}, + isClientInitiated: false, + prompt_id: 'prompt-1', + }; + + const results = await scheduler.schedule( + [request], + new AbortController().signal, + ); + + expect(results.length).toBe(1); + const result = results[0]; + expect(result.status).toBe(CoreToolCallStatus.Error); + const erroredCall = result as ErroredToolCall; + + expect(erroredCall.response.error?.message).toContain( + 'Tool execution blocked: Hook blocked execution', + ); + expect(executeFn).not.toHaveBeenCalled(); + }); + + it('should update tool input if BeforeTool hook provides modified input', async () => { + const executeFn = vi.fn().mockResolvedValue({ + llmContent: 'Tool executed', + returnDisplay: 'Tool executed', + }); + const mockTool = new MockTool({ name: 'mockTool', execute: executeFn }); + + const toolRegistry = { + getTool: () => mockTool, + getAllToolNames: () => ['mockTool'], + } as unknown as ToolRegistry; + + const mockMessageBus = createMockMessageBus(); + + const mockConfig = createMockConfig({ + getToolRegistry: () => toolRegistry, + getMessageBus: () => mockMessageBus, + getApprovalMode: () => ApprovalMode.YOLO, + }); + + const hookSystem = new HookSystem(mockConfig); + + (mockConfig as { getHookSystem?: () => HookSystem }).getHookSystem = () => + hookSystem; + + hookSystem.registerHook( + { + type: HookType.Runtime, + name: 'test-modify-input-hook', + action: async () => ({ + continue: true, + hookSpecificOutput: { + hookEventName: 'BeforeTool', + tool_input: { newParam: 'modifiedValue' }, + }, + }), + }, + HookEventName.BeforeTool, + ); + + const scheduler = new Scheduler({ + context: { + config: mockConfig, + messageBus: mockMessageBus, + toolRegistry, + } as unknown as AgentLoopContext, + getPreferredEditor: () => 'vscode', + schedulerId: 'test-scheduler', + }); + + const request = { + callId: '1', + name: 'mockTool', + args: { originalParam: 'originalValue' }, + isClientInitiated: false, + prompt_id: 'prompt-1', + }; + + const results = await scheduler.schedule( + [request], + new AbortController().signal, + ); + + expect(results.length).toBe(1); + const result = results[0]; + expect(result.status).toBe(CoreToolCallStatus.Success); + + expect(executeFn).toHaveBeenCalledWith( + { newParam: 'modifiedValue' }, + expect.anything(), + undefined, + expect.anything(), + ); + + expect(result.request.args).toEqual({ + newParam: 'modifiedValue', + }); + }); +});