basant307/AI_Governance_Project
048
1/**2 * @license3 * Copyright 2025 Qwen Team4 * SPDX-License-Identifier: Apache-2.05 */6 7import {8 describe,9 it,10 expect,11 vi,12 beforeEach,13 afterEach,14 type Mock,15} from 'vitest';16 17// Force UTC timezone so toLocaleDateString('en-US', ...) produces consistent18// output regardless of the developer's local timezone.19process.env.TZ = 'UTC';20 21import { mkdtemp, writeFile, rm } from 'node:fs/promises';22import { tmpdir } from 'node:os';23import { join } from 'node:path';24import type { Content, GenerateContentResponse, Part } from '@google/genai';25import { GeminiClient, SendMessageType } from './client.js';26import { getRecentGitStatus } from '../utils/gitUtils.js';27import {28 AuthType,29 createContentGenerator,30 type ContentGenerator,31 type ContentGeneratorConfig,32} from './contentGenerator.js';33import { BaseLlmClient } from './baseLlmClient.js';34import { buildAgentContentGeneratorConfig } from '../models/content-generator-config.js';35import { GeminiChat } from './geminiChat.js';36import type { Config } from '../config/config.js';37import { ApprovalMode } from '../config/config.js';38import {39 createHookOutput,40 PermissionMode,41 SessionStartSource,42} from '../hooks/types.js';43import type { ModelsConfig } from '../models/modelsConfig.js';44import { UnauthorizedError } from '../utils/errors.js';45import { retryWithBackoff } from '../utils/retry.js';46import {47 CompressionStatus,48 GeminiEventType,49 Turn,50 type ServerGeminiStreamEvent,51} from './turn.js';52import { LoopType } from '../telemetry/types.js';53 54type MockSessionStartProfiler = {55 time: Mock;56 timeSync: Mock;57 finish: Mock;58};59 60const sessionStartProfilerMocks = vi.hoisted(() => ({61 createSessionStartProfiler: vi.fn(),62 profilers: [] as MockSessionStartProfiler[],63}));64 65vi.mock('./session-start-profiler.js', () => ({66 createSessionStartProfiler:67 sessionStartProfilerMocks.createSessionStartProfiler,68}));69 70vi.mock('../utils/retry.js', () => ({71 retryWithBackoff: vi.fn(async (fn) => await fn()),72 isUnattendedMode: vi.fn(() => false),73}));74import {75 getCoreSystemPrompt,76 getCustomSystemPrompt,77 getPlanModeSystemReminder,78} from './prompts.js';79import { DEFAULT_QWEN_FLASH_MODEL } from '../config/models.js';80import { FileDiscoveryService } from '../services/fileDiscoveryService.js';81import { promptIdContext } from '../utils/promptIdContext.js';82import { setSimulate429 } from '../utils/testUtils.js';83import { ideContextStore } from '../ide/ideContext.js';84import { uiTelemetryService } from '../telemetry/uiTelemetry.js';85import {86 buildChangedAgentsReminder,87 buildChangedMcpToolsReminder,88 buildChangedSkillsReminder,89 getInitialChatHistory,90} from '../utils/environmentContext.js';91import { collectAvailableSkillEntries } from '../tools/skill-utils.js';92import type { AvailableSkillEntry } from '../tools/skill-utils.js';93import { ToolNames } from '../tools/tool-names.js';94import {95 __resetActiveGoalStoreForTests,96 clearActiveGoal,97 setActiveGoal,98} from '../goals/activeGoalStore.js';99import type { FileHistorySnapshot } from '../services/fileHistoryService.js';100import { runWithAgentContext } from '../agents/runtime/agent-context.js';101 102// Mock fs module to prevent actual file system operations during tests103const mockFileSystem = new Map<string, string>();104 105vi.mock('node:fs', () => {106 const fsModule = {107 mkdirSync: vi.fn(),108 writeFileSync: vi.fn((path: string, data: string) => {109 mockFileSystem.set(path, data);110 }),111 readFileSync: vi.fn((path: string) => {112 if (mockFileSystem.has(path)) {113 return mockFileSystem.get(path);114 }115 throw Object.assign(new Error('ENOENT: no such file or directory'), {116 code: 'ENOENT',117 });118 }),119 existsSync: vi.fn((path: string) => mockFileSystem.has(path)),120 appendFileSync: vi.fn(),121 };122 123 return {124 default: fsModule,125 ...fsModule,126 };127});128 129// --- Mocks ---130const mockTurnRunFn = vi.fn();131 132vi.mock('./turn', async (importOriginal) => {133 const actual = await importOriginal<typeof import('./turn.js')>();134 // Define a mock class that has the same shape as the real Turn135 class MockTurn {136 pendingToolCalls = [];137 // The run method is a property that holds our mock function138 run = mockTurnRunFn;139 140 constructor() {141 // The constructor can be empty or do some mock setup142 }143 }144 // Export the mock class as 'Turn'145 return {146 ...actual,147 Turn: MockTurn,148 };149});150 151vi.mock('../config/config.js');152vi.mock('./prompts');153vi.mock('../models/content-generator-config.js', async (importOriginal) => {154 const actual =155 await importOriginal<156 typeof import('../models/content-generator-config.js')157 >();158 return {159 ...actual,160 buildAgentContentGeneratorConfig: vi161 .fn()162 .mockImplementation(actual.buildAgentContentGeneratorConfig),163 };164});165vi.mock('./contentGenerator.js', async (importOriginal) => {166 const actual = await importOriginal<typeof import('./contentGenerator.js')>();167 return {168 ...actual,169 createContentGenerator: vi.fn(),170 };171});172vi.mock('../utils/getFolderStructure', () => ({173 getFolderStructure: vi.fn().mockResolvedValue('Mock Folder Structure'),174}));175vi.mock('../utils/errorReporting', () => ({ reportError: vi.fn() }));176vi.mock('../utils/gitUtils.js', async (importOriginal) => {177 const actual = await importOriginal<typeof import('../utils/gitUtils.js')>();178 return {179 ...actual,180 getRecentGitStatus: vi.fn().mockReturnValue(null),181 };182});183vi.mock('../utils/nextSpeakerChecker', () => ({184 checkNextSpeaker: vi.fn().mockResolvedValue(null),185}));186vi.mock('../tools/skill-utils.js', async (importOriginal) => {187 const actual =188 await importOriginal<typeof import('../tools/skill-utils.js')>();189 return {190 ...actual,191 collectAvailableSkillEntries: vi.fn(),192 };193});194vi.mock('../utils/environmentContext', async (importOriginal) => {195 const actual =196 await importOriginal<typeof import('../utils/environmentContext.js')>();197 return {198 ...actual,199 getEnvironmentContext: vi200 .fn()201 .mockResolvedValue([{ text: 'Mocked env context' }]),202 getDirectoryContextString: vi203 .fn()204 .mockResolvedValue('Mocked directory context'),205 getInitialChatHistory: vi.fn(async (_config, extraHistory) => [206 [207 {208 role: 'user',209 parts: [210 {211 text: '<system-reminder>\nMocked env context\n</system-reminder>',212 },213 ],214 },215 ...(extraHistory ?? []),216 ],217 [],218 ]),219 buildChangedMcpToolsReminder: vi.fn(220 (221 tools: Array<{ name: string }>,222 removedToolNames: string[],223 ): string | null =>224 tools.length === 0 && removedToolNames.length === 0225 ? null226 : `<system-reminder>\nchanged mcp: added=${tools.map((tool) => tool.name).join(', ')} removed=${removedToolNames.join(', ')}\n</system-reminder>`,227 ),228 buildChangedSkillsReminder: vi.fn(229 (230 entries: Array<{ name: string }>,231 removedNames: string[],232 ): string | null =>233 entries.length === 0 && removedNames.length === 0234 ? null235 : `<system-reminder>\nchanged skills: added=${entries.map((entry) => entry.name).join(', ')} removed=${removedNames.join(', ')}\n</system-reminder>`,236 ),237 buildChangedAgentsReminder: vi.fn(238 (239 addedAgents: Array<{ name: string }>,240 removedAgentNames: string[],241 ): string | null =>242 addedAgents.length === 0 && removedAgentNames.length === 0243 ? null244 : `<system-reminder>\nchanged agents: added=${addedAgents.map((agent) => agent.name).join(', ')} removed=${removedAgentNames.join(', ')}\n</system-reminder>`,245 ),246 getStartupContextLength: vi.fn((history) => {247 const first = history?.[0];248 if (first?.role !== 'user') return 0;249 const text = first.parts?.[0]?.text;250 if (typeof text === 'string' && text.startsWith('<system-reminder>')) {251 return 1;252 }253 if (254 history?.[1]?.role === 'model' &&255 history?.[1]?.parts?.[0]?.text === 'Got it. Thanks for the context!'256 ) {257 return 2;258 }259 return 0;260 }),261 isSystemReminderContent: vi.fn((content) => {262 const parts = content?.parts;263 if (!parts || parts.length === 0) return false;264 return parts.every(265 (part: { text?: string }) =>266 typeof part.text === 'string' &&267 part.text.startsWith('<system-reminder>') &&268 part.text.includes('</system-reminder>'),269 );270 }),271 };272});273vi.mock('../utils/generateContentResponseUtilities', () => ({274 getResponseText: (result: GenerateContentResponse) =>275 result.candidates?.[0]?.content?.parts?.map((part) => part.text).join('') ||276 undefined,277 getFunctionCalls: (result: GenerateContentResponse) => {278 // Extract function calls from the response279 const parts = result.candidates?.[0]?.content?.parts;280 if (!parts) {281 return undefined;282 }283 const functionCallParts = parts284 .filter((part) => !!part.functionCall)285 .map((part) => part.functionCall);286 return functionCallParts.length > 0 ? functionCallParts : undefined;287 },288}));289// Create shared mock for uiTelemetryService that's used by both telemetry mocks290const mockUiTelemetryService = vi.hoisted(() => ({291 setLastPromptTokenCount: vi.fn(),292 getLastPromptTokenCount: vi.fn(),293 setLastCachedContentTokenCount: vi.fn(),294 reset: vi.fn(),295 resetSession: vi.fn(),296 addEvent: vi.fn(),297}));298vi.mock('../telemetry/tracer.js', () => ({299 API_CALL_ABORTED_SPAN_STATUS_MESSAGE: 'API call aborted',300 API_CALL_FAILED_SPAN_STATUS_MESSAGE: 'API call failed',301}));302 303vi.mock('../telemetry/index.js', async (importOriginal) => {304 const actual = await importOriginal<typeof import('../telemetry/index.js')>();305 return {306 ...actual,307 uiTelemetryService: mockUiTelemetryService,308 // We keep the real implementations of logChatCompression, etc.309 // but we can spy on QwenLogger if needed310 };311});312vi.mock('../ide/ideContext.js');313vi.mock('../telemetry/uiTelemetry.js', () => ({314 uiTelemetryService: mockUiTelemetryService,315}));316vi.mock('../telemetry/loggers.js', () => ({317 logChatCompression: vi.fn(),318 logNextSpeakerCheck: vi.fn(),319 logApiRequest: vi.fn(),320 logLoopDetected: vi.fn(),321 logLoopDetectionDisabled: vi.fn(),322}));323 324const { mockClientDebugLogger } = vi.hoisted(() => ({325 mockClientDebugLogger: {326 isEnabled: vi.fn().mockReturnValue(false),327 debug: vi.fn(),328 info: vi.fn(),329 warn: vi.fn(),330 error: vi.fn(),331 },332}));333vi.mock('../utils/debugLogger.js', async (importOriginal) => {334 const actual =335 await importOriginal<typeof import('../utils/debugLogger.js')>();336 return {337 ...actual,338 createDebugLogger: (namespace: string) =>339 namespace === 'CLIENT'340 ? mockClientDebugLogger341 : actual.createDebugLogger(namespace),342 };343});344 345vi.mock(346 '../services/microcompaction/microcompact.js',347 async (importOriginal) => {348 const actual =349 await importOriginal<350 typeof import('../services/microcompaction/microcompact.js')351 >();352 return {353 ...actual,354 microcompactHistory: vi.fn(actual.microcompactHistory),355 };356 },357);358import { microcompactHistory } from '../services/microcompaction/microcompact.js';359 360// Mock RequestTokenizer to use simple character-based estimation361vi.mock('../utils/request-tokenizer/requestTokenizer.js', () => ({362 RequestTokenizer: class {363 async calculateTokens(request: { contents: unknown }) {364 // Simple estimation: count characters in JSON and divide by 4365 const totalChars = JSON.stringify(request.contents).length;366 return {367 totalTokens: Math.floor(totalChars / 4),368 breakdown: {369 textTokens: Math.floor(totalChars / 4),370 imageTokens: 0,371 audioTokens: 0,372 otherTokens: 0,373 },374 processingTime: 0,375 };376 }377 },378}));379 380/**381 * Array.fromAsync ponyfill, which will be available in es 2024.382 *383 * Buffers an async generator into an array and returns the result.384 */385async function fromAsync<T>(promise: AsyncGenerator<T>): Promise<readonly T[]> {386 const results: T[] = [];387 for await (const result of promise) {388 results.push(result);389 }390 return results;391}392 393function getLastTurnRequestText(): string {394 const request = mockTurnRunFn.mock.calls.at(-1)?.[1];395 if (typeof request === 'string') {396 return request;397 }398 if (Array.isArray(request)) {399 return request400 .map((part) => {401 if (typeof part === 'string') {402 return part;403 }404 if (part && typeof part === 'object' && 'text' in part) {405 return part.text ?? '';406 }407 return JSON.stringify(part);408 })409 .join('');410 }411 return JSON.stringify(request ?? '');412}413 414describe('Gemini Client (client.ts)', () => {415 let mockContentGenerator: ContentGenerator;416 let mockConfig: Config;417 let client: GeminiClient;418 let mockGenerateContentFn: Mock;419 let mockFileHistoryService: {420 makeSnapshot: ReturnType<typeof vi.fn>;421 getSnapshots: ReturnType<typeof vi.fn>;422 restoreFromSnapshots: ReturnType<typeof vi.fn>;423 rewind: ReturnType<typeof vi.fn>;424 };425 let mockMemoryManager: {426 scheduleExtract: ReturnType<typeof vi.fn>;427 scheduleDream: ReturnType<typeof vi.fn>;428 recall: ReturnType<typeof vi.fn>;429 scheduleSkillReview: ReturnType<typeof vi.fn>;430 };431 beforeEach(async () => {432 vi.resetAllMocks();433 sessionStartProfilerMocks.profilers.length = 0;434 sessionStartProfilerMocks.createSessionStartProfiler.mockImplementation(435 () => {436 const profiler: MockSessionStartProfiler = {437 time: vi.fn(async (_stage: string, fn: () => Promise<unknown>) =>438 fn(),439 ),440 timeSync: vi.fn((_stage: string, fn: () => unknown) => fn()),441 finish: vi.fn(),442 };443 sessionStartProfilerMocks.profilers.push(profiler);444 return profiler;445 },446 );447 vi.mocked(uiTelemetryService.setLastPromptTokenCount).mockClear();448 449 // Default: createContentGenerator rejects (simulates test env without auth).450 // Individual tests can override with mockResolvedValue for success path.451 vi.mocked(createContentGenerator).mockRejectedValue(452 new Error('no auth in test env'),453 );454 455 mockMemoryManager = {456 scheduleExtract: vi.fn().mockResolvedValue({457 touchedTopics: [],458 cursor: { updatedAt: new Date(0).toISOString() },459 }),460 scheduleDream: vi.fn().mockResolvedValue({461 status: 'skipped',462 skippedReason: 'min_sessions',463 }),464 recall: vi.fn().mockResolvedValue({465 prompt: '',466 selectedDocs: [],467 strategy: 'none',468 }),469 scheduleSkillReview: vi.fn().mockReturnValue({470 status: 'skipped',471 skippedReason: 'below_threshold',472 }),473 };474 475 mockGenerateContentFn = vi.fn().mockResolvedValue({476 candidates: [{ content: { parts: [{ text: '{"key": "value"}' }] } }],477 });478 mockFileHistoryService = {479 makeSnapshot: vi.fn().mockResolvedValue(undefined),480 getSnapshots: vi.fn().mockReturnValue([]),481 restoreFromSnapshots: vi.fn(),482 rewind: vi.fn(),483 };484 485 // Disable 429 simulation for tests486 setSimulate429(false);487 488 mockContentGenerator = {489 generateContent: mockGenerateContentFn,490 generateContentStream: vi.fn(),491 batchEmbedContents: vi.fn(),492 countTokens: vi.fn().mockResolvedValue({ totalTokens: 100 }),493 } as unknown as ContentGenerator;494 495 // Because the GeminiClient constructor kicks off an async process (startChat)496 // that depends on a fully-formed Config object, we need to mock the497 // entire implementation of Config for these tests.498 const mockToolRegistry = {499 warmAll: vi.fn().mockResolvedValue(undefined),500 ensureTool: vi.fn().mockResolvedValue(null),501 getFunctionDeclarations: vi.fn().mockReturnValue([]),502 getDeferredToolSummary: vi.fn().mockReturnValue([]),503 clearRevealedDeferredTools: vi.fn(),504 revealDeferredTool: vi.fn(),505 isDeferredToolRevealed: vi.fn().mockReturnValue(false),506 getTool: vi.fn().mockReturnValue(null),507 getMcpServerInstructions: vi.fn().mockReturnValue(new Map()),508 };509 const fileService = new FileDiscoveryService('/test/dir');510 const contentGeneratorConfig: ContentGeneratorConfig = {511 model: 'test-model',512 apiKey: 'test-key',513 vertexai: false,514 authType: AuthType.USE_GEMINI,515 };516 mockConfig = {517 getContentGeneratorConfig: vi518 .fn()519 .mockReturnValue(contentGeneratorConfig),520 getToolRegistry: vi.fn().mockReturnValue(mockToolRegistry),521 getModel: vi.fn().mockReturnValue('test-model'),522 getEmbeddingModel: vi.fn().mockReturnValue('test-embedding-model'),523 getApiKey: vi.fn().mockReturnValue('test-key'),524 getVertexAI: vi.fn().mockReturnValue(false),525 getUserAgent: vi.fn().mockReturnValue('test-agent'),526 getUserMemory: vi.fn().mockReturnValue(''),527 getSystemPrompt: vi.fn().mockReturnValue(undefined),528 getAppendSystemPrompt: vi.fn().mockReturnValue(undefined),529 getFullContext: vi.fn().mockReturnValue(false),530 getSessionId: vi.fn().mockReturnValue('test-session-id'),531 getProxy: vi.fn().mockReturnValue(undefined),532 getWorkingDir: vi.fn().mockReturnValue('/test/dir'),533 getFileService: vi.fn().mockReturnValue(fileService),534 getMaxSessionTurns: vi.fn().mockReturnValue(0),535 getClearContextOnIdle: vi.fn().mockReturnValue({536 toolResultsThresholdMinutes: 60,537 toolResultsNumToKeep: 5,538 }),539 getSessionTokenLimit: vi.fn().mockReturnValue(32000),540 getNoBrowser: vi.fn().mockReturnValue(false),541 getUsageStatisticsEnabled: vi.fn().mockReturnValue(true),542 getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),543 getSdkMode: vi.fn().mockReturnValue(false),544 getIdeModeFeature: vi.fn().mockReturnValue(false),545 getIdeMode: vi.fn().mockReturnValue(true),546 getDebugMode: vi.fn().mockReturnValue(false),547 getWorkspaceContext: vi.fn().mockReturnValue({548 getDirectories: vi.fn().mockReturnValue(['/test/dir']),549 }),550 getGeminiClient: vi.fn(),551 getModelRouterService: vi.fn().mockReturnValue({552 route: vi.fn().mockResolvedValue({ model: 'default-routed-model' }),553 }),554 getCliVersion: vi.fn().mockReturnValue('1.0.0'),555 getChatCompression: vi.fn().mockReturnValue(undefined),556 getSkipNextSpeakerCheck: vi.fn().mockReturnValue(false),557 getUseModelRouter: vi.fn().mockReturnValue(false),558 getProjectRoot: vi.fn().mockReturnValue('/test/project/root'),559 getCwd: vi.fn().mockReturnValue('/test/project/root'),560 storage: {561 getProjectTempDir: vi.fn().mockReturnValue('/test/temp'),562 getProjectDir: vi563 .fn()564 .mockReturnValue('/test/project/root/.gemini/projects/test-project'),565 },566 getContentGenerator: vi.fn().mockReturnValue(mockContentGenerator),567 getBaseLlmClient: vi.fn(),568 getSkipLoopDetection: vi.fn().mockReturnValue(false),569 // Mimics the resolved Config getter: always a number (Infinity keeps570 // the cap out of the way of unrelated streaming tests).571 getMaxToolCallsPerTurn: vi.fn().mockReturnValue(Number.POSITIVE_INFINITY),572 getChatRecordingService: vi.fn().mockReturnValue(undefined),573 getFileHistoryService: vi.fn().mockReturnValue(mockFileHistoryService),574 getResumedSessionData: vi.fn().mockReturnValue(undefined),575 getArenaAgentClient: vi.fn().mockReturnValue(null),576 getManagedAutoMemoryEnabled: vi.fn().mockReturnValue(true),577 isManagedMemoryAvailable: vi.fn().mockReturnValue(true),578 getMemoryManager: vi.fn().mockReturnValue(mockMemoryManager),579 getAutoSkillEnabled: vi.fn().mockReturnValue(false),580 getAutoSkillConfirmEnabled: vi.fn().mockReturnValue(true),581 getModelsConfig: vi.fn().mockReturnValue({582 getResolvedModel: vi.fn().mockReturnValue(undefined),583 }),584 getAllConfiguredModels: vi.fn().mockReturnValue([]),585 getDisableAllHooks: vi.fn().mockReturnValue(true),586 getStopHookBlockingCap: vi.fn().mockReturnValue(8),587 getArenaManager: vi.fn().mockReturnValue(null),588 getMessageBus: vi.fn().mockReturnValue(undefined),589 hasHooksForEvent: vi.fn().mockReturnValue(false),590 getHookSystem: vi.fn().mockReturnValue(undefined),591 getSkillManager: vi.fn().mockReturnValue(undefined),592 getSubagentManager: vi.fn().mockReturnValue({593 listSubagents: vi.fn().mockResolvedValue([]),594 }),595 consumeInlineAnnouncedSkillKeys: vi596 .fn()597 .mockReturnValue(new Set<string>()),598 getDebugLogger: vi.fn().mockReturnValue({599 isEnabled: vi.fn().mockReturnValue(true),600 debug: vi.fn(),601 info: vi.fn(),602 warn: vi.fn(),603 error: vi.fn(),604 }),605 getFileReadCache: vi.fn().mockReturnValue({606 clear: vi.fn(),607 }),608 } as unknown as Config;609 610 // Real BaseLlmClient routes generateText through mockContentGenerator;611 // generateJson is stubbed only for the next-speaker classifier so the612 // next-speaker schema isn't reproduced in every test.613 const realBaseLlmClient = new BaseLlmClient(614 mockContentGenerator,615 mockConfig,616 );617 realBaseLlmClient.generateJson = vi.fn().mockResolvedValue({618 next_speaker: 'user',619 reasoning: 'test',620 });621 vi.mocked(mockConfig.getBaseLlmClient).mockReturnValue(realBaseLlmClient);622 623 client = new GeminiClient(mockConfig);624 await client.initialize();625 vi.mocked(mockConfig.getGeminiClient).mockReturnValue(client);626 627 // GeminiClient.sendMessageStream calls this.tryCompressChat (which now628 // delegates to chat.tryCompress) before each turn. Most tests use a629 // hand-rolled chat mock that doesn't implement tryCompress; default the630 // wrapper to a NOOP so those tests don't crash. Tests that exercise631 // compression directly (the delegation tests below, the632 // emits-compression-event test) override this spy.633 vi.spyOn(client, 'tryCompressChat').mockResolvedValue({634 originalTokenCount: 0,635 newTokenCount: 0,636 compressionStatus: CompressionStatus.NOOP,637 });638 });639 640 afterEach(() => {641 vi.useRealTimers();642 vi.restoreAllMocks();643 __resetActiveGoalStoreForTests();644 });645 646 describe('initialize', () => {647 it('seeds resumed chat with replayed prompt token count', async () => {648 vi.mocked(mockConfig.getResumedSessionData).mockReturnValue({649 conversation: {650 sessionId: 'resumed-session-id',651 projectHash: 'project-hash',652 startTime: new Date(0).toISOString(),653 lastUpdated: new Date(0).toISOString(),654 messages: [],655 },656 filePath: '/test/session.jsonl',657 lastCompletedUuid: null,658 });659 vi.mocked(uiTelemetryService.getLastPromptTokenCount).mockReturnValue(660 123_456,661 );662 663 const resumedClient = new GeminiClient(mockConfig);664 await resumedClient.initialize();665 666 expect(resumedClient.getChat().getLastPromptTokenCount()).toBe(123_456);667 });668 669 it('seeds resumed chat with previous response output token count', async () => {670 const seedResumeTokenCountsSpy = vi.spyOn(671 GeminiChat.prototype,672 'seedResumeTokenCounts',673 );674 vi.mocked(mockConfig.getResumedSessionData).mockReturnValue({675 conversation: {676 sessionId: 'resumed-session-id',677 projectHash: 'project-hash',678 startTime: new Date(0).toISOString(),679 lastUpdated: new Date(0).toISOString(),680 messages: [681 {682 uuid: 'assistant-1',683 parentUuid: null,684 sessionId: 'resumed-session-id',685 timestamp: new Date(0).toISOString(),686 type: 'assistant',687 cwd: '/test/project',688 version: '1.0.0',689 message: { role: 'model', parts: [{ text: 'done' }] },690 usageMetadata: {691 promptTokenCount: 200,692 candidatesTokenCount: 60,693 thoughtsTokenCount: 20,694 totalTokenCount: 280,695 },696 },697 ],698 },699 filePath: '/test/session.jsonl',700 lastCompletedUuid: null,701 });702 703 const resumedClient = new GeminiClient(mockConfig);704 await resumedClient.initialize();705 706 expect(resumedClient.getChat().getLastPromptTokenCount()).toBe(200);707 expect(seedResumeTokenCountsSpy).toHaveBeenCalledWith(200, 80);708 });709 710 it('seeds recently completed tools from resumed history', async () => {711 vi.mocked(mockConfig.getResumedSessionData).mockReturnValue({712 conversation: {713 sessionId: 'resumed-session-id',714 projectHash: 'project-hash',715 startTime: new Date(0).toISOString(),716 lastUpdated: new Date(0).toISOString(),717 messages: [718 {719 message: {720 role: 'model',721 parts: [722 {723 functionCall: {724 id: 'call_read',725 name: 'read_file',726 args: {},727 },728 },729 ],730 },731 },732 {733 message: {734 role: 'user',735 parts: [736 {737 functionResponse: {738 id: 'call_read',739 name: 'read_file',740 response: { ok: true },741 },742 },743 ],744 },745 },746 {747 message: {748 role: 'model',749 parts: [750 {751 functionCall: {752 id: 'call_pending',753 name: 'write_file',754 args: {},755 },756 },757 ],758 },759 },760 ],761 },762 filePath: '/test/session.jsonl',763 lastCompletedUuid: null,764 } as unknown as ReturnType<Config['getResumedSessionData']>);765 766 const resumedClient = new GeminiClient(mockConfig);767 await resumedClient.initialize();768 769 expect(resumedClient['recentCompletedToolNames']).toEqual(['read_file']);770 });771 772 it('uses Startup SessionStart source for non-resumed initialize without explicit source', async () => {773 const hookSystem = {774 fireSessionStartEvent: vi.fn().mockResolvedValue(775 createHookOutput('SessionStart', {776 hookSpecificOutput: {777 additionalContext: 'Startup hook context',778 },779 }),780 ),781 };782 vi.mocked(mockConfig.getDisableAllHooks).mockReturnValue(false);783 vi.mocked(mockConfig.hasHooksForEvent).mockReturnValue(true);784 vi.mocked(mockConfig.getHookSystem).mockReturnValue(785 hookSystem as unknown as ReturnType<Config['getHookSystem']>,786 );787 788 const freshClient = new GeminiClient(mockConfig);789 await freshClient.initialize();790 791 expect(hookSystem.fireSessionStartEvent).toHaveBeenCalledWith(792 SessionStartSource.Startup,793 'test-model',794 PermissionMode.Default,795 );796 });797 798 it('is idempotent when initialize is called twice on the same session', async () => {799 const hookSystem = {800 fireSessionStartEvent: vi.fn().mockResolvedValue(801 createHookOutput('SessionStart', {802 hookSpecificOutput: {803 additionalContext: 'Startup hook context',804 },805 }),806 ),807 };808 vi.mocked(mockConfig.getDisableAllHooks).mockReturnValue(false);809 vi.mocked(mockConfig.hasHooksForEvent).mockReturnValue(true);810 vi.mocked(mockConfig.getHookSystem).mockReturnValue(811 hookSystem as unknown as ReturnType<Config['getHookSystem']>,812 );813 814 const freshClient = new GeminiClient(mockConfig);815 await freshClient.initialize();816 const firstChat = freshClient.getChat();817 await freshClient.initialize(SessionStartSource.Resume);818 819 expect(freshClient.getChat()).toBe(firstChat);820 expect(hookSystem.fireSessionStartEvent).toHaveBeenCalledTimes(1);821 expect(hookSystem.fireSessionStartEvent).toHaveBeenCalledWith(822 SessionStartSource.Startup,823 'test-model',824 PermissionMode.Default,825 );826 });827 828 it('rebuilds chat when initialize is called after the session id changes', async () => {829 const hookSystem = {830 fireSessionStartEvent: vi.fn().mockResolvedValue(undefined),831 };832 vi.mocked(mockConfig.getDisableAllHooks).mockReturnValue(false);833 vi.mocked(mockConfig.hasHooksForEvent).mockReturnValue(true);834 vi.mocked(mockConfig.getHookSystem).mockReturnValue(835 hookSystem as unknown as ReturnType<Config['getHookSystem']>,836 );837 vi.mocked(mockConfig.getSessionId)838 .mockReturnValueOnce('session-a')839 .mockReturnValueOnce('session-b');840 841 const freshClient = new GeminiClient(mockConfig);842 await freshClient.initialize();843 const firstChat = freshClient.getChat();844 await freshClient.initialize(SessionStartSource.Resume);845 846 expect(freshClient.getChat()).not.toBe(firstChat);847 expect(hookSystem.fireSessionStartEvent).toHaveBeenCalledTimes(2);848 expect(hookSystem.fireSessionStartEvent).toHaveBeenNthCalledWith(849 1,850 SessionStartSource.Startup,851 'test-model',852 PermissionMode.Default,853 );854 expect(hookSystem.fireSessionStartEvent).toHaveBeenNthCalledWith(855 2,856 SessionStartSource.Resume,857 'test-model',858 PermissionMode.Default,859 );860 });861 });862 863 describe('fireSessionStartHook', () => {864 it('returns trimmed additionalContext from the SessionStart hook', async () => {865 const hookSystem = {866 fireSessionStartEvent: vi.fn().mockResolvedValue(867 createHookOutput('SessionStart', {868 hookSpecificOutput: {869 additionalContext: ' hook context ',870 },871 }),872 ),873 };874 vi.mocked(mockConfig.getDisableAllHooks).mockReturnValue(false);875 vi.mocked(mockConfig.hasHooksForEvent).mockReturnValue(true);876 vi.mocked(mockConfig.getHookSystem).mockReturnValue(877 hookSystem as unknown as ReturnType<Config['getHookSystem']>,878 );879 880 await expect(881 client['fireSessionStartHook'](SessionStartSource.Startup),882 ).resolves.toBe('hook context');883 expect(hookSystem.fireSessionStartEvent).toHaveBeenCalledWith(884 SessionStartSource.Startup,885 'test-model',886 PermissionMode.Default,887 );888 });889 890 it('returns undefined without firing when SessionStart hooks are disabled', async () => {891 const hookSystem = {892 fireSessionStartEvent: vi.fn(),893 };894 vi.mocked(mockConfig.getDisableAllHooks).mockReturnValue(true);895 vi.mocked(mockConfig.hasHooksForEvent).mockReturnValue(true);896 vi.mocked(mockConfig.getHookSystem).mockReturnValue(897 hookSystem as unknown as ReturnType<Config['getHookSystem']>,898 );899 900 await expect(901 client['fireSessionStartHook'](SessionStartSource.Startup),902 ).resolves.toBeUndefined();903 expect(hookSystem.fireSessionStartEvent).not.toHaveBeenCalled();904 });905 906 it('logs and returns undefined when the SessionStart hook throws', async () => {907 const fireSessionStartEvent = vi908 .fn()909 .mockRejectedValue(new Error('hook failed'));910 const debugLogger = {911 isEnabled: vi.fn().mockReturnValue(true),912 debug: vi.fn(),913 info: vi.fn(),914 warn: vi.fn(),915 error: vi.fn(),916 };917 vi.mocked(mockConfig.getDisableAllHooks).mockReturnValue(false);918 vi.mocked(mockConfig.hasHooksForEvent).mockReturnValue(true);919 vi.mocked(mockConfig.getHookSystem).mockReturnValue({920 fireSessionStartEvent,921 } as unknown as ReturnType<Config['getHookSystem']>);922 vi.mocked(mockConfig.getDebugLogger).mockReturnValue(debugLogger);923 924 await expect(925 client['fireSessionStartHook'](SessionStartSource.Compact),926 ).resolves.toBeUndefined();927 expect(debugLogger.warn).toHaveBeenCalledWith(928 'SessionStart hook failed: Error: hook failed',929 );930 });931 });932 933 describe('startChat — session start profiling', () => {934 beforeEach(() => {935 sessionStartProfilerMocks.createSessionStartProfiler.mockClear();936 sessionStartProfilerMocks.profilers.length = 0;937 });938 939 it('passes startup, resume, and clear sources to the profiler', async () => {940 await client.startChat();941 await client.startChat([{ role: 'user', parts: [{ text: 'hi' }] }]);942 await client.startChat(undefined, SessionStartSource.Clear);943 944 expect(945 sessionStartProfilerMocks.createSessionStartProfiler.mock.calls.map(946 ([source]) => source,947 ),948 ).toEqual([949 SessionStartSource.Startup,950 SessionStartSource.Resume,951 SessionStartSource.Clear,952 ]);953 expect(954 sessionStartProfilerMocks.profilers[1].finish,955 ).toHaveBeenCalledWith(956 expect.objectContaining({ extraHistoryLength: 1 }),957 );958 for (const profiler of sessionStartProfilerMocks.profilers) {959 expect(profiler.finish).toHaveBeenCalledTimes(1);960 }961 });962 963 it('finalizes successful startChat profiles with bounded counts', async () => {964 const hookSystem = {965 fireSessionStartEvent: vi.fn().mockResolvedValue(966 createHookOutput('SessionStart', {967 hookSpecificOutput: {968 additionalContext: 'hook output',969 },970 }),971 ),972 };973 vi.mocked(mockConfig.getDisableAllHooks).mockReturnValue(false);974 vi.mocked(mockConfig.hasHooksForEvent).mockReturnValue(true);975 vi.mocked(mockConfig.getHookSystem).mockReturnValue(976 hookSystem as unknown as ReturnType<Config['getHookSystem']>,977 );978 979 await client.startChat(undefined, SessionStartSource.Clear);980 981 const profiler = sessionStartProfilerMocks.profilers.at(-1)!;982 expect(profiler.finish).toHaveBeenCalledWith(983 expect.objectContaining({984 ok: true,985 extraHistoryLength: 0,986 historyLength: 1,987 snapshotEntryCount: 0,988 deferredReminderCount: 0,989 }),990 );991 expect(profiler.time.mock.calls.map(([stage]) => stage)).toEqual([992 'tool_registry_warm',993 'initial_chat_history',994 'agent_reminder_seed',995 'session_start_hook',996 'set_tools',997 ]);998 expect(profiler.timeSync.mock.calls.map(([stage]) => stage)).toEqual([999 'resume_deferred_tool_reveal',1000 'deferred_reminder_setup',1001 'skill_reminder_seed',1002 'system_instruction',1003 'gemini_chat_construct',1004 'orphan_tool_use_repair',1005 'session_start_context_apply',1006 ]);1007 });1008 1009 it('records non-zero snapshot and deferred reminder counts', async () => {1010 const toolRegistry = vi.mocked(1011 mockConfig.getToolRegistry,1012 )() as unknown as {1013 getDeferredToolSummary: ReturnType<typeof vi.fn>;1014 getTool: ReturnType<typeof vi.fn>;1015 };1016 toolRegistry.getDeferredToolSummary.mockReturnValue([1017 { name: 'cron_create', description: 'schedule' },1018 ]);1019 toolRegistry.getTool.mockImplementation((name: string) =>1020 name === ToolNames.TOOL_SEARCH ? ({} as never) : null,1021 );1022 vi.mocked(getInitialChatHistory).mockResolvedValueOnce([1023 [1024 {1025 role: 'user',1026 parts: [{ text: '<system-reminder>context</system-reminder>' }],1027 },1028 ],1029 [1030 { name: 'skill-one', description: 'first skill' },1031 { name: 'skill-two', description: 'second skill' },1032 ],1033 ]);1034 1035 await client.startChat();1036 1037 const profiler = sessionStartProfilerMocks.profilers.at(-1)!;1038 expect(profiler.finish).toHaveBeenCalledWith(1039 expect.objectContaining({1040 ok: true,1041 snapshotEntryCount: 2,1042 deferredReminderCount: 1,1043 }),1044 );1045 });1046 1047 it('does not record context apply stage without SessionStart context', async () => {1048 await client.startChat();1049 1050 const profiler = sessionStartProfilerMocks.profilers.at(-1)!;1051 expect(1052 profiler.timeSync.mock.calls.map(([stage]) => stage),1053 ).not.toContain('session_start_context_apply');1054 });1055 1056 it('finalizes failed startChat profiles without changing the thrown error', async () => {1057 vi.mocked(getInitialChatHistory).mockRejectedValueOnce(1058 new Error('history failed'),1059 );1060 1061 await expect(client.startChat()).rejects.toThrow(1062 'Failed to initialize chat: history failed',1063 );1064 1065 const profiler = sessionStartProfilerMocks.profilers.at(-1)!;1066 expect(profiler.finish).toHaveBeenCalledWith(1067 expect.objectContaining({1068 ok: false,1069 extraHistoryLength: 0,1070 historyLength: 0,1071 snapshotEntryCount: 0,1072 deferredReminderCount: 0,1073 }),1074 );1075 });1076 1077 it('finalizes failed startChat profiles for first-stage warm errors', async () => {1078 const toolRegistry = vi.mocked(1079 mockConfig.getToolRegistry,1080 )() as unknown as {1081 warmAll: ReturnType<typeof vi.fn>;1082 };1083 toolRegistry.warmAll.mockRejectedValueOnce(new Error('warm failed'));1084 1085 await expect(client.startChat()).rejects.toThrow(1086 'Failed to initialize chat: warm failed',1087 );1088 1089 const profiler = sessionStartProfilerMocks.profilers.at(-1)!;1090 expect(profiler.time.mock.calls.map(([stage]) => stage)).toContain(1091 'tool_registry_warm',1092 );1093 expect(profiler.finish).toHaveBeenCalledWith(1094 expect.objectContaining({1095 ok: false,1096 extraHistoryLength: 0,1097 historyLength: 0,1098 snapshotEntryCount: 0,1099 deferredReminderCount: 0,1100 }),1101 );1102 });1103 1104 it('finalizes failed startChat profiles for sync stage errors', async () => {1105 vi.spyOn(1106 client as unknown as { getMainSessionSystemInstruction: () => string },1107 'getMainSessionSystemInstruction',1108 ).mockImplementationOnce(() => {1109 throw new Error('system instruction failed');1110 });1111 1112 await expect(client.startChat()).rejects.toThrow(1113 'Failed to initialize chat: system instruction failed',1114 );1115 1116 const profiler = sessionStartProfilerMocks.profilers.at(-1)!;1117 expect(profiler.timeSync.mock.calls.map(([stage]) => stage)).toContain(1118 'system_instruction',1119 );1120 expect(profiler.finish).toHaveBeenCalledWith(1121 expect.objectContaining({1122 ok: false,1123 extraHistoryLength: 0,1124 historyLength: 1,1125 snapshotEntryCount: 0,1126 deferredReminderCount: 0,1127 }),1128 );1129 });1130 1131 it('finalizes failed startChat profiles with partial counts', async () => {1132 const toolRegistry = vi.mocked(1133 mockConfig.getToolRegistry,1134 )() as unknown as {1135 getDeferredToolSummary: ReturnType<typeof vi.fn>;1136 getTool: ReturnType<typeof vi.fn>;1137 };1138 toolRegistry.getDeferredToolSummary.mockReturnValue([1139 { name: 'cron_create', description: 'schedule' },1140 ]);1141 toolRegistry.getTool.mockImplementation((name: string) =>1142 name === ToolNames.TOOL_SEARCH ? ({} as never) : null,1143 );1144 vi.mocked(getInitialChatHistory).mockResolvedValueOnce([1145 [1146 {1147 role: 'user',1148 parts: [{ text: '<system-reminder>context</system-reminder>' }],1149 },1150 ],1151 [{ name: 'skill-one', description: 'first skill' }],1152 ]);1153 vi.spyOn(client, 'setTools').mockRejectedValueOnce(1154 new Error('set tools failed'),1155 );1156 1157 await expect(client.startChat()).rejects.toThrow(1158 'Failed to initialize chat: set tools failed',1159 );1160 1161 const profiler = sessionStartProfilerMocks.profilers.at(-1)!;1162 expect(profiler.finish).toHaveBeenCalledWith(1163 expect.objectContaining({1164 ok: false,1165 extraHistoryLength: 0,1166 historyLength: 1,1167 snapshotEntryCount: 1,1168 deferredReminderCount: 1,1169 }),1170 );1171 });1172 });1173 1174 describe('startChat — deferred tools', () => {1175 // Pulls the registry mock used by the surrounding suite so each test1176 // can stub the deferred-summary + ToolSearch availability per case.1177 function getRegistryMock() {1178 return vi.mocked(mockConfig.getToolRegistry)() as unknown as {1179 getDeferredToolSummary: ReturnType<typeof vi.fn>;1180 getTool: ReturnType<typeof vi.fn>;1181 isDeferredToolRevealed: ReturnType<typeof vi.fn>;1182 revealDeferredTool: ReturnType<typeof vi.fn>;1183 };1184 }1185 1186 it('re-reveals deferred tools that appear in resumed history', async () => {1187 // Resume contract: a transcript referencing `cron_create` (a1188 // deferred tool) must re-reveal it on startChat so the API1189 // declaration list includes its schema — otherwise a follow-up1190 // call to that tool would be rejected as unknown.1191 const reg = getRegistryMock();1192 reg.getDeferredToolSummary.mockReturnValue([1193 { name: 'cron_create', description: 'schedule' },1194 { name: 'cron_list', description: 'list' },1195 ]);1196 // ToolSearch is available so we DON'T enter the eager-reveal branch.1197 reg.getTool.mockImplementation((n: string) =>1198 n === 'tool_search' ? ({} as never) : null,1199 );1200 reg.revealDeferredTool.mockClear();