From 2e80fad7a4db113ffcf945147219296eed406f00 Mon Sep 17 00:00:00 2001 From: Your Name Date: Mon, 6 Apr 2026 19:18:17 +0000 Subject: [PATCH] refactor environment --- .../src/context/contextManager.golden.test.ts | 16 +++++- packages/core/src/context/contextManager.ts | 7 +-- packages/core/src/context/ir/projector.ts | 2 +- .../processors/blobDegradationProcessor.ts | 6 +-- .../semanticCompressionProcessor.test.ts | 5 +- .../semanticCompressionProcessor.ts | 2 +- .../processors/stateSnapshotProcessor.ts | 6 +-- .../processors/toolMaskingProcessor.ts | 4 +- .../core/src/context/sidecar/environment.ts | 17 +++--- .../src/context/sidecar/environmentImpl.ts | 54 +++---------------- .../src/context/testing/contextTestUtils.ts | 21 ++++---- 11 files changed, 57 insertions(+), 83 deletions(-) diff --git a/packages/core/src/context/contextManager.golden.test.ts b/packages/core/src/context/contextManager.golden.test.ts index 4b4d45b058..9eb05ff620 100644 --- a/packages/core/src/context/contextManager.golden.test.ts +++ b/packages/core/src/context/contextManager.golden.test.ts @@ -17,6 +17,7 @@ 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 type { Content } from '@google/genai'; @@ -69,6 +70,7 @@ describe('ContextManager Golden Tests', () => { const sidecar = SidecarLoader.fromLegacyConfig(mockConfig as any); const tracer = new ContextTracer('/tmp', 'test-session'); + const eventBus = new ContextEventBus(); const env = new ContextEnvironmentImpl( {} as any, 'test-prompt-id', @@ -77,6 +79,7 @@ describe('ContextManager Golden Tests', () => { '/tmp', tracer, 4, + eventBus ); contextManager = new ContextManager(sidecar, env, tracer); }); @@ -128,12 +131,23 @@ describe('ContextManager Golden Tests', () => { // In Golden Tests, we just want to ensure the logic doesn't throw or alter unprotected history in weird ways. // Since we're skipping processors due to being under budget, it should equal history. const tracer2 = new ContextTracer('/tmp', 'test2'); + const eventBus2 = new ContextEventBus(); + const env2 = new ContextEnvironmentImpl( + {} as any, + 'test-prompt-id', + 'test', + '/tmp', + '/tmp', + tracer2, + 4, + eventBus2 + ); contextManager = new ContextManager( { budget: { retainedTokens: 100000, maxTokens: 150000 }, pipelines: [], } as any, - {} as any, + env2, tracer2, ); diff --git a/packages/core/src/context/contextManager.ts b/packages/core/src/context/contextManager.ts index 01652d0ac0..f5f23abccb 100644 --- a/packages/core/src/context/contextManager.ts +++ b/packages/core/src/context/contextManager.ts @@ -51,12 +51,7 @@ export class ContextManager { constructor(private sidecar: SidecarConfig, private env: ContextEnvironment, private readonly tracer: ContextTracer) { - - - this.eventBus = new ContextEventBus(); - if ('setEventBus' in this.env) { - (this.env as any).setEventBus(this.eventBus); - } + this.eventBus = env.eventBus; // Register built-ins BEFORE creating Orchestrator ProcessorRegistry.register({ id: 'ToolMaskingProcessor', create: (env, opts) => new ToolMaskingProcessor(env, opts as any) }); diff --git a/packages/core/src/context/ir/projector.ts b/packages/core/src/context/ir/projector.ts index 22fc1a3221..75aae8d85b 100644 --- a/packages/core/src/context/ir/projector.ts +++ b/packages/core/src/context/ir/projector.ts @@ -68,7 +68,7 @@ export class IrProjector { try { const fs = await import('node:fs/promises'); const path = await import('node:path'); - const dumpPath = path.join(env.getTraceDir(), '.gemini', 'projected_context.json'); + const dumpPath = path.join(env.traceDir, '.gemini', 'projected_context.json'); await fs.mkdir(path.dirname(dumpPath), { recursive: true }); await fs.writeFile(dumpPath, JSON.stringify(contents, null, 2), 'utf-8'); debugLogger.log(`[Observability] Context successfully dumped to ${dumpPath}`); diff --git a/packages/core/src/context/processors/blobDegradationProcessor.ts b/packages/core/src/context/processors/blobDegradationProcessor.ts index 7dc54848cf..2857e5764d 100644 --- a/packages/core/src/context/processors/blobDegradationProcessor.ts +++ b/packages/core/src/context/processors/blobDegradationProcessor.ts @@ -33,10 +33,10 @@ export class BlobDegradationProcessor implements ContextProcessor { let directoryCreated = false; let blobOutputsDir = path.join( - this.env.getProjectTempDir(), + this.env.projectTempDir, 'degraded-blobs', ); - const sessionId = this.env.getSessionId(); + const sessionId = this.env.sessionId; if (sessionId) { blobOutputsDir = path.join( blobOutputsDir, @@ -102,7 +102,7 @@ export class BlobDegradationProcessor implements ContextProcessor { if (newText && tokensSaved > 0) { const newTokens = estimateTokenCountSync([{ text: newText }], 0, { - charsPerToken: this.env.getCharsPerToken(), + charsPerToken: this.env.charsPerToken, }); part.presentation = { text: newText, tokens: newTokens }; diff --git a/packages/core/src/context/processors/semanticCompressionProcessor.test.ts b/packages/core/src/context/processors/semanticCompressionProcessor.test.ts index 5af28443ce..9692189292 100644 --- a/packages/core/src/context/processors/semanticCompressionProcessor.test.ts +++ b/packages/core/src/context/processors/semanticCompressionProcessor.test.ts @@ -15,6 +15,7 @@ import type { } from '../ir/types.js'; import type { ContextAccountingState } from '../pipeline.js'; import { randomUUID } from 'node:crypto'; +import type { BaseLlmClient } from 'src/core/baseLlmClient.js'; describe('SemanticCompressionProcessor', () => { let processor: SemanticCompressionProcessor; @@ -26,9 +27,7 @@ describe('SemanticCompressionProcessor', () => { }); const env = createMockEnvironment(); - env.getLlmClient = vi - .fn() - .mockReturnValue({ generateContent: generateContentMock }) as any; + vi.spyOn(env, 'llmClient', 'get').mockReturnValue({ generateContent: generateContentMock } as unknown as BaseLlmClient); processor = new SemanticCompressionProcessor(env, { nodeThresholdTokens: 2000, }); diff --git a/packages/core/src/context/processors/semanticCompressionProcessor.ts b/packages/core/src/context/processors/semanticCompressionProcessor.ts index 74509d2791..7c4a56546a 100644 --- a/packages/core/src/context/processors/semanticCompressionProcessor.ts +++ b/packages/core/src/context/processors/semanticCompressionProcessor.ts @@ -181,7 +181,7 @@ export class SemanticCompressionProcessor implements ContextProcessor { ): Promise { const promptMessage = `You are compressing an old episodic context buffer for an AI assistant.\nSummarize this ${contentType} block in 2-3 highly technical sentences. Keep all critical facts, file names, dependencies, and architectural decisions. Discard conversational filler and boilerplate.\n\nContent:\n${content.slice(0, 30000)}`; - const client = this.env.getLlmClient(); + const client = this.env.llmClient; try { const response = await client.generateContent({ modelConfigKey: { model: this.modelToUse }, diff --git a/packages/core/src/context/processors/stateSnapshotProcessor.ts b/packages/core/src/context/processors/stateSnapshotProcessor.ts index 929138bf51..b158b1c80d 100644 --- a/packages/core/src/context/processors/stateSnapshotProcessor.ts +++ b/packages/core/src/context/processors/stateSnapshotProcessor.ts @@ -19,7 +19,7 @@ export interface StateSnapshotProcessorOptions { export class StateSnapshotProcessor implements ContextProcessor { static create(env: ContextEnvironment, options: StateSnapshotProcessorOptions): StateSnapshotProcessor { - return new StateSnapshotProcessor(env, options, (env as any).getEventBus()); + return new StateSnapshotProcessor(env, options, env.eventBus); } readonly id = 'StateSnapshotProcessor'; readonly name = 'StateSnapshotProcessor'; @@ -76,7 +76,7 @@ export class StateSnapshotProcessor implements ContextProcessor { } private async synthesizeSnapshot(episodes: Episode[]): Promise { - const client = this.env.getLlmClient(); + const client = this.env.llmClient; const systemPrompt = this.options.systemInstruction ?? `You are an expert Context Memory Manager. You will be provided with a raw transcript of older conversation turns between a user and an AI assistant. @@ -106,7 +106,7 @@ Output ONLY the raw factual snapshot, formatted compactly. Do not include markdo modelConfigKey: { model: 'state-snapshot-processor' }, contents: [{ role: 'user', parts: [{ text: userPromptText }] }], systemInstruction: { role: 'system', parts: [{ text: systemPrompt }] }, - promptId: this.env.getPromptId(), + promptId: this.env.promptId, role: LlmRole.UTILITY_STATE_SNAPSHOT_PROCESSOR, abortSignal: new AbortController().signal, }, diff --git a/packages/core/src/context/processors/toolMaskingProcessor.ts b/packages/core/src/context/processors/toolMaskingProcessor.ts index f68095fced..e254600d8c 100644 --- a/packages/core/src/context/processors/toolMaskingProcessor.ts +++ b/packages/core/src/context/processors/toolMaskingProcessor.ts @@ -53,10 +53,10 @@ export class ToolMaskingProcessor implements ContextProcessor { const limitChars = maskingConfig.stringLengthThresholdTokens * 4; let toolOutputsDir = path.join( - this.env.getProjectTempDir(), + this.env.projectTempDir, 'tool-outputs', ); - const sessionId = this.env.getSessionId(); + const sessionId = this.env.sessionId; if (sessionId) { toolOutputsDir = path.join( toolOutputsDir, diff --git a/packages/core/src/context/sidecar/environment.ts b/packages/core/src/context/sidecar/environment.ts index 5700f0232e..f1e0f99b23 100644 --- a/packages/core/src/context/sidecar/environment.ts +++ b/packages/core/src/context/sidecar/environment.ts @@ -9,12 +9,13 @@ export type { ContextTracer, ContextEventBus }; export interface ContextEnvironment { - getLlmClient(): BaseLlmClient; - getPromptId(): string; - getSessionId(): string; - getTraceDir(): string; - getProjectTempDir(): string; - getEventBus(): ContextEventBus; - getTracer(): ContextTracer; - getCharsPerToken(): number; + readonly llmClient: BaseLlmClient; + readonly promptId: string; + readonly sessionId: string; + readonly traceDir: string; + readonly projectTempDir: string; + readonly tracer: ContextTracer; + readonly charsPerToken: number; + + readonly eventBus: ContextEventBus; } diff --git a/packages/core/src/context/sidecar/environmentImpl.ts b/packages/core/src/context/sidecar/environmentImpl.ts index 0cbfc12bba..fabd089cfc 100644 --- a/packages/core/src/context/sidecar/environmentImpl.ts +++ b/packages/core/src/context/sidecar/environmentImpl.ts @@ -11,52 +11,14 @@ import type { ContextEnvironment } from './environment.js'; import type { ContextEventBus } from '../eventBus.js'; export class ContextEnvironmentImpl implements ContextEnvironment { - private eventBus?: ContextEventBus; - constructor( - private llmClient: BaseLlmClient, - private sessionId: string, - private promptId: string, - private traceDir: string, - private tempDir: string, - private tracer: ContextTracer, - private charsPerToken: number, + public readonly llmClient: BaseLlmClient, + public readonly sessionId: string, + public readonly promptId: string, + public readonly traceDir: string, + public readonly projectTempDir: string, + public readonly tracer: ContextTracer, + public readonly charsPerToken: number, + public readonly eventBus: ContextEventBus, ) {} - - setEventBus(bus: ContextEventBus) { - this.eventBus = bus; - } - - getEventBus(): ContextEventBus { - if (!this.eventBus) throw new Error('EventBus not bound'); - return this.eventBus; - } - - getLlmClient(): BaseLlmClient { - return this.llmClient; - } - - getSessionId(): string { - return this.sessionId; - } - - getTraceDir(): string { - return this.traceDir; - } - - getProjectTempDir(): string { - return this.tempDir; - } - - getTracer(): ContextTracer { - return this.tracer; - } - - getCharsPerToken(): number { - return this.charsPerToken; - } - - getPromptId(): string { - return this.promptId; - } } diff --git a/packages/core/src/context/testing/contextTestUtils.ts b/packages/core/src/context/testing/contextTestUtils.ts index 918c1917c8..9001bc3f4b 100644 --- a/packages/core/src/context/testing/contextTestUtils.ts +++ b/packages/core/src/context/testing/contextTestUtils.ts @@ -14,18 +14,18 @@ import { ContextManager } from '../contextManager.js'; export function createMockEnvironment(): ContextEnvironment { return { // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion - getLlmClient: vi.fn().mockReturnValue({ + llmClient: vi.fn().mockReturnValue({ generateContent: vi.fn().mockResolvedValue({ text: 'Mock LLM summary response', }), - } as unknown as BaseLlmClient), - getPromptId: vi.fn().mockReturnValue('mock-prompt-id'), - getSessionId: vi.fn().mockReturnValue('mock-session'), - getTraceDir: vi.fn().mockReturnValue('/tmp/.gemini/trace'), - getProjectTempDir: vi.fn().mockReturnValue('/tmp/.gemini/tool-outputs'), - getEventBus: vi.fn(), - getTracer: vi.fn(), - getCharsPerToken: vi.fn().mockReturnValue(1), + })() as unknown as BaseLlmClient, + promptId: 'mock-prompt-id', + sessionId: 'mock-session', + traceDir: '/tmp/.gemini/trace', + projectTempDir: '/tmp/.gemini/tool-outputs', + eventBus: new ContextEventBus(), + tracer: new ContextTracer('/tmp', 'mock-session'), + charsPerToken: 1, }; } @@ -88,12 +88,14 @@ export function createMockContextConfig( import { ContextTracer } from '../tracer.js'; import { ContextEnvironmentImpl } from '../sidecar/environmentImpl.js'; import { SidecarLoader } from '../sidecar/SidecarLoader.js'; +import { ContextEventBus } from '../eventBus.js'; import type { BaseLlmClient } from 'src/core/baseLlmClient.js'; export function setupContextComponentTest(config: Config) { const chatHistory = new AgentChatHistory(); const sidecar = SidecarLoader.fromLegacyConfig(config); const tracer = new ContextTracer('/tmp', 'test-session'); + const eventBus = new ContextEventBus(); const env = new ContextEnvironmentImpl( config.getBaseLlmClient(), 'test prompt-id', @@ -102,6 +104,7 @@ export function setupContextComponentTest(config: Config) { '/tmp/gemini-test', tracer, 1, + eventBus ); const contextManager = new ContextManager(sidecar, env, tracer);