diff --git a/packages/core/src/context/contextManager.golden.test.ts b/packages/core/src/context/contextManager.golden.test.ts index e0772d367c..3982672561 100644 --- a/packages/core/src/context/contextManager.golden.test.ts +++ b/packages/core/src/context/contextManager.golden.test.ts @@ -13,19 +13,10 @@ import { beforeAll, afterAll, } from 'vitest'; -import { ContextManager } from './contextManager.js'; -import { ContextEnvironmentImpl } from './sidecar/environmentImpl.js'; -import { SidecarLoader } from './sidecar/SidecarLoader.js'; -import { ContextTracer } from './tracer.js'; -import { ContextEventBus } from './eventBus.js'; -import { PipelineOrchestrator } from './sidecar/orchestrator.js'; -import { AgentChatHistory } from '../core/agentChatHistory.js'; +import type { ContextManager } from './contextManager.js'; import type { Content } from '@google/genai'; -import type { BaseLlmClient } from '../core/baseLlmClient.js'; import type { Episode } from './ir/types.js'; import type { SidecarConfig } from './sidecar/types.js'; -import { SidecarRegistry } from './sidecar/registry.js'; -import { registerBuiltInProcessors } from './sidecar/builtins.js'; import { createMockContextConfig, setupContextComponentTest } from './testing/contextTestUtils.js'; expect.addSnapshotSerializer({ @@ -47,73 +38,10 @@ describe('ContextManager Golden Tests', () => { vi.restoreAllMocks(); }); - let mockConfig: any; // eslint-disable-line @typescript-eslint/no-explicit-any let contextManager: ContextManager; beforeEach(() => { - mockConfig = { - isContextManagementEnabled: vi.fn().mockReturnValue(true), - getExperimentalContextSidecarConfig: vi.fn().mockReturnValue(undefined), - getTargetDir: vi.fn().mockReturnValue('/tmp'), - getSessionId: vi.fn().mockReturnValue('test-session'), - getToolOutputMaskingConfig: vi.fn().mockResolvedValue({ - enabled: true, - minPrunableThresholdTokens: 50, - protectLatestTurn: false, - protectionThresholdTokens: 100, - }), - storage: { getProjectTempDir: vi.fn().mockReturnValue('/tmp') }, - getUsageStatisticsEnabled: vi.fn().mockReturnValue(false), - getBaseLlmClient: vi.fn().mockReturnValue({ - generateJson: vi.fn().mockResolvedValue({ - 'test_file.txt': { level: 'SUMMARY' }, - }), - generateContent: vi.fn().mockResolvedValue({ - candidates: [ - { content: { parts: [{ text: 'This is a summary.' }] } }, - ], - }), - }), - }; - - const registry = new SidecarRegistry(); - registerBuiltInProcessors(registry); - - const sidecar = SidecarLoader.fromConfig(mockConfig, registry); - const tracer = new ContextTracer({ - targetDir: '/tmp', - sessionId: 'test-session', - }); - const eventBus = new ContextEventBus(); - const env = new ContextEnvironmentImpl( - { - generateContent: async () => ({}), - generateJson: async () => ({}), - } as unknown as BaseLlmClient, - 'test-prompt-id', - 'test', - '/tmp', - '/tmp', - tracer, - 4, - eventBus, - ); - const chatHistory = new AgentChatHistory(); - const orchestrator = new PipelineOrchestrator( - sidecar, - env, - eventBus, - tracer, - registry - ); - - contextManager = new ContextManager( - sidecar, - env, - tracer, - orchestrator, - chatHistory - ); + contextManager = setupContextComponentTest(createMockContextConfig()).contextManager; }); const createLargeHistory = (): Content[] => [ @@ -177,3 +105,4 @@ describe('ContextManager Golden Tests', () => { expect(result.length).toEqual(history.length + 1); }); }); + diff --git a/packages/core/src/context/processors/nodeDistillationProcessor.test.ts b/packages/core/src/context/processors/nodeDistillationProcessor.test.ts index aab43b2d61..485fce5a85 100644 --- a/packages/core/src/context/processors/nodeDistillationProcessor.test.ts +++ b/packages/core/src/context/processors/nodeDistillationProcessor.test.ts @@ -5,23 +5,20 @@ */ import assert from 'node:assert'; -import { describe, it, expect, vi } from 'vitest'; +import { describe, it, expect } from 'vitest'; import { NodeDistillationProcessor } from './nodeDistillationProcessor.js'; import { createMockProcessArgs, createMockEnvironment, createDummyNode, createDummyToolNode, - createMockGenerateContentResponse + createMockLlmClient } from '../testing/contextTestUtils.js'; import type { UserPrompt, AgentThought, ToolExecution } from '../ir/types.js'; -import type { BaseLlmClient } from '../../core/baseLlmClient.js'; describe('NodeDistillationProcessor', () => { it('should trigger summarization via LLM for long text parts', async () => { - const mockLlmClient = { - generateContent: vi.fn().mockResolvedValue(createMockGenerateContentResponse('Mocked Summary!')), // length = 15 - } as unknown as BaseLlmClient; + const mockLlmClient = createMockLlmClient(['Mocked Summary!']); // Use charsPerToken=1 naturally. const env = createMockEnvironment({ @@ -75,9 +72,7 @@ describe('NodeDistillationProcessor', () => { }); it('should ignore nodes that are below the threshold', async () => { - const mockLlmClient = { - generateContent: vi.fn().mockResolvedValue(createMockGenerateContentResponse('S')), // length = 1 - } as unknown as BaseLlmClient; + const mockLlmClient = createMockLlmClient(['S']); // length = 1 const env = createMockEnvironment({ llmClient: mockLlmClient, diff --git a/packages/core/src/context/sidecar/environmentImpl.test.ts b/packages/core/src/context/sidecar/environmentImpl.test.ts index d3b43b1674..e990b106b2 100644 --- a/packages/core/src/context/sidecar/environmentImpl.test.ts +++ b/packages/core/src/context/sidecar/environmentImpl.test.ts @@ -9,13 +9,15 @@ import { ContextTracer } from '../tracer.js'; import { ContextEventBus } from '../eventBus.js'; import { InMemoryFileSystem } from '../system/InMemoryFileSystem.js'; import { DeterministicIdGenerator } from '../system/DeterministicIdGenerator.js'; -import type { BaseLlmClient } from '../../core/baseLlmClient.js'; + + +import { createMockLlmClient } from '../testing/contextTestUtils.js'; describe('ContextEnvironmentImpl', () => { it('should initialize with defaults correctly', () => { const tracer = new ContextTracer({ targetDir: '/tmp', sessionId: 'mock' }); const eventBus = new ContextEventBus(); - const mockLlmClient = {} as BaseLlmClient; + const mockLlmClient = createMockLlmClient(); const env = new ContextEnvironmentImpl( mockLlmClient, @@ -49,7 +51,7 @@ describe('ContextEnvironmentImpl', () => { it('should initialize with provided overrides', () => { const tracer = new ContextTracer({ targetDir: '/tmp', sessionId: 'mock' }); const eventBus = new ContextEventBus(); - const mockLlmClient = {} as BaseLlmClient; + const mockLlmClient = createMockLlmClient(); const fileSystem = new InMemoryFileSystem(); const idGenerator = new DeterministicIdGenerator('test-'); diff --git a/packages/core/src/context/system-tests/lifecycle.golden.test.ts b/packages/core/src/context/system-tests/lifecycle.golden.test.ts index 2c56f8119b..0a3437ae06 100644 --- a/packages/core/src/context/system-tests/lifecycle.golden.test.ts +++ b/packages/core/src/context/system-tests/lifecycle.golden.test.ts @@ -6,8 +6,9 @@ import { describe, it, expect, vi, beforeAll, afterAll } from 'vitest'; import { SimulationHarness } from './SimulationHarness.js'; +import { createMockLlmClient } from '../testing/contextTestUtils.js'; import type { SidecarConfig } from '../sidecar/types.js'; -import type { BaseLlmClient } from '../../core/baseLlmClient.js'; + expect.addSnapshotSerializer({ test: (val) => @@ -58,11 +59,7 @@ describe('System Lifecycle Golden Tests', () => { ], }); - const mockLlmClient = { - generateContent: vi.fn().mockResolvedValue({ - text: '', - }), - } as unknown as BaseLlmClient; + const mockLlmClient = createMockLlmClient(['']); it('Scenario 1: Organic Growth with Huge Tool Output & Images', async () => { const harness = await SimulationHarness.create( diff --git a/packages/core/src/context/testing/contextTestUtils.ts b/packages/core/src/context/testing/contextTestUtils.ts index 84163f8aa0..6715e66d4e 100644 --- a/packages/core/src/context/testing/contextTestUtils.ts +++ b/packages/core/src/context/testing/contextTestUtils.ts @@ -21,7 +21,9 @@ import type { ConcreteNode, ToolExecution } from '../ir/types.js'; import type { ContextEnvironment } from '../sidecar/environment.js'; import type { Config } from '../../config/config.js'; import type { BaseLlmClient } from '../../core/baseLlmClient.js'; -import type { Content , GenerateContentResponse } from '@google/genai'; +import type { Content, GenerateContentResponse } from '@google/genai'; +import { InboxSnapshotImpl } from '../sidecar/inbox.js'; +import type { ContextWorkingBuffer, InboxMessage, ProcessArgs } from '../pipeline.js'; /** @@ -92,14 +94,46 @@ export function createDummyToolNode( } as unknown as ToolExecution; } +import type { Mock } from 'vitest'; +import type { SidecarConfig } from '../sidecar/types.js'; + +export interface MockLlmClient extends BaseLlmClient { + generateContent: Mock; +} + +export function createMockLlmClient(responses?: Array): MockLlmClient { + const generateContentMock = vi.fn(); + + if (responses && responses.length > 0) { + for (const response of responses) { + if (typeof response === 'string') { + generateContentMock.mockResolvedValueOnce(createMockGenerateContentResponse(response)); + } else { + generateContentMock.mockResolvedValueOnce(response); + } + } + // Fallback to the last response for any subsequent calls + const lastResponse = responses[responses.length - 1]; + if (typeof lastResponse === 'string') { + generateContentMock.mockResolvedValue(createMockGenerateContentResponse(lastResponse)); + } else { + generateContentMock.mockResolvedValue(lastResponse); + } + } else { + // Default fallback + generateContentMock.mockResolvedValue(createMockGenerateContentResponse('Mock LLM response')); + } + + // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion + return { + generateContent: generateContentMock, + } as unknown as MockLlmClient; +} + export function createMockEnvironment( overrides?: Partial, ): ContextEnvironment { - const mockClient: Partial = { - generateContent: vi.fn().mockResolvedValue(createMockGenerateContentResponse('Mock LLM summary response')), - }; - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion - const llmClient = mockClient as BaseLlmClient; + const llmClient = createMockLlmClient(['Mock LLM summary response']); const tracer = new ContextTracer({ targetDir: '/tmp', sessionId: 'mock-session' }); const eventBus = new ContextEventBus(); @@ -127,9 +161,6 @@ export function createMockEnvironment( * Creates a block of synthetic conversation history designed to consume a specific number of tokens. * Assumes roughly 4 characters per token for standard English text. */ -import { InboxSnapshotImpl } from '../sidecar/inbox.js'; -import type { ContextWorkingBuffer, InboxMessage, ProcessArgs } from '../pipeline.js'; - export class FakeContextWorkingBuffer implements ContextWorkingBuffer { readonly nodes: readonly ConcreteNode[]; private readonly nodesById = new Map(); @@ -221,8 +252,8 @@ export function createMockContextConfig( export function setupContextComponentTest( config: Config, - sidecarOverride?: import('../sidecar/types.js').SidecarConfig, -) { + sidecarOverride?: SidecarConfig, +): {chatHistory: AgentChatHistory, contextManager: ContextManager} { const chatHistory = new AgentChatHistory(); const registry = new SidecarRegistry(); registerBuiltInProcessors(registry); @@ -260,6 +291,5 @@ export function setupContextComponentTest( ); // The async worker is now internally managed by ContextManager - return { chatHistory, contextManager }; }