refactor(core): Centralize context management logic into src/context (#24380)

packages/core/src/agents/local-executor.test.ts

@@ -108,7 +108,7 @@ import {
 } from '../scheduler/types.js';
 
 import { CompressionStatus } from '../core/turn.js';
-import { ChatCompressionService } from '../services/chatCompressionService.js';
+import { ChatCompressionService } from '../context/chatCompressionService.js';
 import type {
   ModelConfigKey,
   ResolvedModelConfig,
@@ -121,7 +121,7 @@ const mockSetHistory = vi.fn((newHistory: Content[]) => {
   mockChatHistory = newHistory;
 });
 
-vi.mock('../services/chatCompressionService.js', () => ({
+vi.mock('../context/chatCompressionService.js', () => ({
   ChatCompressionService: vi.fn().mockImplementation(() => ({
     compress: mockCompress,
   })),

packages/core/src/agents/local-executor.ts

@@ -30,7 +30,7 @@ import {
 } from '../tools/mcp-tool.js';
 import { CompressionStatus } from '../core/turn.js';
 import { type ToolCallRequestInfo } from '../scheduler/types.js';
-import { ChatCompressionService } from '../services/chatCompressionService.js';
+import { ChatCompressionService } from '../context/chatCompressionService.js';
 import { getDirectoryContextString } from '../utils/environmentContext.js';
 import { renderUserMemory } from '../prompts/snippets.js';
 import { promptIdContext } from '../utils/promptIdContext.js';

packages/core/src/config/config.test.ts

@@ -221,7 +221,7 @@ vi.mock('../utils/fetch.js', () => ({
   setGlobalProxy: mockSetGlobalProxy,
 }));
 
-vi.mock('../services/contextManager.js', () => ({
+vi.mock('../context/contextManager.js', () => ({
   ContextManager: vi.fn().mockImplementation(() => ({
     refresh: vi.fn(),
     getGlobalMemory: vi.fn().mockReturnValue(''),
@@ -237,7 +237,7 @@ import { tokenLimit } from '../core/tokenLimits.js';
 import { getCodeAssistServer } from '../code_assist/codeAssist.js';
 import { getExperiments } from '../code_assist/experiments/experiments.js';
 import type { CodeAssistServer } from '../code_assist/server.js';
-import { ContextManager } from '../services/contextManager.js';
+import { ContextManager } from '../context/contextManager.js';
 import { UserTierId } from '../code_assist/types.js';
 import type {
   ModelConfigService,

packages/core/src/config/config.ts

@@ -116,7 +116,7 @@ import {
   type ModelConfigServiceConfig,
 } from '../services/modelConfigService.js';
 import { DEFAULT_MODEL_CONFIGS } from './defaultModelConfigs.js';
-import { ContextManager } from '../services/contextManager.js';
+import { ContextManager } from '../context/contextManager.js';
 import { TrackerService } from '../services/trackerService.js';
 import type { GenerateContentParameters } from '@google/genai';
 
@@ -451,7 +451,7 @@ import {
   DEFAULT_TOOL_PROTECTION_THRESHOLD,
   DEFAULT_MIN_PRUNABLE_TOKENS_THRESHOLD,
   DEFAULT_PROTECT_LATEST_TURN,
-} from '../services/toolOutputMaskingService.js';
+} from '../context/toolOutputMaskingService.js';
 
 import {
   type ExtensionLoader,

packages/core/src/context/__snapshots__/toolOutputMaskingService.test.ts.snap

@@ -0,0 +1,31 @@
+// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
+
+exports[`ToolOutputMaskingService > should match the expected snapshot for a masked tool output 1`] = `
+"<tool_output_masked>
+Line
+Line
+Line
+Line
+Line
+Line
+Line
+Line
+Line
+Line
+
+... [6 lines omitted] ...
+
+Line
+Line
+Line
+Line
+Line
+Line
+Line
+Line
+Line
+
+
+Output too large. Full output available at: /mock/temp/tool-outputs/session-mock-session/run_shell_command_deterministic.txt
+</tool_output_masked>"
+`;

packages/core/src/context/agentHistoryProvider.test.ts

@@ -17,12 +17,12 @@ vi.mock('../utils/tokenCalculation.js', () => ({
 import type { Content, GenerateContentResponse, Part } from '@google/genai';
 import type { Config, ContextManagementConfig } from '../config/config.js';
 import type { BaseLlmClient } from '../core/baseLlmClient.js';
-import type { AgentHistoryProviderConfig } from './types.js';
+import type { AgentHistoryProviderConfig } from '../services/types.js';
 import {
   TEXT_TRUNCATION_PREFIX,
   TOOL_TRUNCATION_PREFIX,
   truncateProportionally,
-} from 'src/utils/truncation.js';
+} from './truncation.js';
 
 describe('AgentHistoryProvider', () => {
   let config: Config;

packages/core/src/context/agentHistoryProvider.ts

@@ -9,7 +9,7 @@ import { getResponseText } from '../utils/partUtils.js';
 import { estimateTokenCountSync } from '../utils/tokenCalculation.js';
 import { LlmRole } from '../telemetry/llmRole.js';
 import { debugLogger } from '../utils/debugLogger.js';
-import type { AgentHistoryProviderConfig } from './types.js';
+import type { AgentHistoryProviderConfig } from '../services/types.js';
 import type { Config } from '../config/config.js';
 import {
   MIN_TARGET_TOKENS,
@@ -18,7 +18,7 @@ import {
   estimateCharsFromTokens,
   truncateProportionally,
   normalizeFunctionResponse,
-} from '../utils/truncation.js';
+} from './truncation.js';
 
 export class AgentHistoryProvider {
   // TODO(joshualitt): just pass the BaseLlmClient instead of the whole Config.

packages/core/src/context/toolDistillationService.ts

@@ -25,7 +25,7 @@ import {
   MIN_TARGET_TOKENS,
   estimateCharsFromTokens,
   normalizeFunctionResponse,
-} from '../utils/truncation.js';
+} from './truncation.js';
 
 // Skip structural map generation for outputs larger than this threshold (in characters)
 // as it consumes excessive tokens and may not be representative of the full content.

packages/core/src/context/truncation.ts

@@ -9,7 +9,7 @@ import {
   estimateTokenCountSync,
   ASCII_TOKENS_PER_CHAR,
   NON_ASCII_TOKENS_PER_CHAR,
-} from './tokenCalculation.js';
+} from '../utils/tokenCalculation.js';
 
 export const MIN_TARGET_TOKENS = 10;
 export const MIN_CHARS_FOR_TRUNCATION = 100;

packages/core/src/core/client.test.ts

@@ -39,7 +39,7 @@ import { tokenLimit } from './tokenLimits.js';
 import { ideContextStore } from '../ide/ideContext.js';
 import type { ModelRouterService } from '../routing/modelRouterService.js';
 import { uiTelemetryService } from '../telemetry/uiTelemetry.js';
-import { ChatCompressionService } from '../services/chatCompressionService.js';
+import { ChatCompressionService } from '../context/chatCompressionService.js';
 import type { ChatRecordingService } from '../services/chatRecordingService.js';
 import { createAvailabilityServiceMock } from '../availability/testUtils.js';
 import type { ModelAvailabilityService } from '../availability/modelAvailabilityService.js';

packages/core/src/core/client.ts

@@ -43,8 +43,8 @@ import type {
 } from '../services/chatRecordingService.js';
 import type { ContentGenerator } from './contentGenerator.js';
 import { LoopDetectionService } from '../services/loopDetectionService.js';
-import { ChatCompressionService } from '../services/chatCompressionService.js';
-import { AgentHistoryProvider } from '../services/agentHistoryProvider.js';
+import { ChatCompressionService } from '../context/chatCompressionService.js';
+import { AgentHistoryProvider } from '../context/agentHistoryProvider.js';
 import { ideContextStore } from '../ide/ideContext.js';
 import {
   logContentRetryFailure,
@@ -65,7 +65,7 @@ import { handleFallback } from '../fallback/handler.js';
 import type { RoutingContext } from '../routing/routingStrategy.js';
 import { debugLogger } from '../utils/debugLogger.js';
 import type { ModelConfigKey } from '../services/modelConfigService.js';
-import { ToolOutputMaskingService } from '../services/toolOutputMaskingService.js';
+import { ToolOutputMaskingService } from '../context/toolOutputMaskingService.js';
 import { calculateRequestTokenCount } from '../utils/tokenCalculation.js';
 import {
   applyModelSelection,

packages/core/src/index.ts

@@ -134,7 +134,7 @@ export * from './services/fileSystemService.js';
 export * from './services/sandboxedFileSystemService.js';
 export * from './sandbox/windows/WindowsSandboxManager.js';
 export * from './services/sessionSummaryUtils.js';
-export * from './services/contextManager.js';
+export * from './context/contextManager.js';
 export * from './services/trackerService.js';
 export * from './services/trackerTypes.js';
 export * from './services/keychainService.js';

packages/core/src/scheduler/tool-executor.ts

@@ -19,7 +19,7 @@ import {
 import { isAbortError } from '../utils/errors.js';
 import { SHELL_TOOL_NAME } from '../tools/tool-names.js';
 import { DiscoveredMCPTool } from '../tools/mcp-tool.js';
-import { ToolOutputDistillationService } from '../services/toolDistillationService.js';
+import { ToolOutputDistillationService } from '../context/toolDistillationService.js';
 import { executeToolWithHooks } from '../core/coreToolHookTriggers.js';
 import {
   saveTruncatedToolOutput,

packages/core/src/tools/jit-context.test.ts

@@ -7,7 +7,7 @@
 import { describe, it, expect, vi, beforeEach } from 'vitest';
 import { discoverJitContext, appendJitContext } from './jit-context.js';
 import type { Config } from '../config/config.js';
-import type { ContextManager } from '../services/contextManager.js';
+import type { ContextManager } from '../context/contextManager.js';
 
 describe('jit-context', () => {
   describe('discoverJitContext', () => {