mock llm

packages/core/src/context/contextManager.golden.test.ts

@@ -13,19 +13,10 @@ import {
   beforeAll,
   afterAll,
 } from 'vitest';
-import { ContextManager } from './contextManager.js';
+import type { 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 { 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 = {
+    contextManager = setupContextComponentTest(createMockContextConfig()).contextManager;
-      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
-    );
   });
 
   const createLargeHistory = (): Content[] => [
@@ -177,3 +105,4 @@ describe('ContextManager Golden Tests', () => {
     expect(result.length).toEqual(history.length + 1);
   });
 });
+

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 = {
+    const mockLlmClient = createMockLlmClient(['Mocked Summary!']);
-      generateContent: vi.fn().mockResolvedValue(createMockGenerateContentResponse('Mocked Summary!')), // length = 15
-    } as unknown as BaseLlmClient;
 
     // 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 = {
+    const mockLlmClient = createMockLlmClient(['S']); // length = 1
-      generateContent: vi.fn().mockResolvedValue(createMockGenerateContentResponse('S')), // length = 1
-    } as unknown as BaseLlmClient;
 
     const env = createMockEnvironment({
        llmClient: mockLlmClient,

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-');
 

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 = {
+  const mockLlmClient = createMockLlmClient(['<MOCKED_STATE_SNAPSHOT_SUMMARY>']);
-    generateContent: vi.fn().mockResolvedValue({
-      text: '<MOCKED_STATE_SNAPSHOT_SUMMARY>',
-    }),
-  } as unknown as BaseLlmClient;
 
   it('Scenario 1: Organic Growth with Huge Tool Output & Images', async () => {
     const harness = await SimulationHarness.create(

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<string | GenerateContentResponse>): 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>,
 ): ContextEnvironment {
-  const mockClient: Partial<BaseLlmClient> = {
+  const llmClient = createMockLlmClient(['Mock LLM summary response']);
-    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 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<string, ConcreteNode>();
@@ -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 };
 }