From a796a328c9773c8557808e8cfe768468ffc3556e Mon Sep 17 00:00:00 2001 From: mini2s Date: Tue, 12 May 2026 19:32:13 +0800 Subject: [PATCH] feat(compact): add CoStrict model support and prevent infinite compaction loops --- bun.lock | 3 + package.json | 1 + src/components/Settings/Config.tsx | 2 +- src/query.ts | 85 ++++-- .../compact/__tests__/autoCompact.test.ts | 247 ++++++++++++++++++ src/services/compact/autoCompact.ts | 47 +++- src/utils/__tests__/context.test.ts | 145 ++++++++++ src/utils/config.ts | 2 +- src/utils/context.ts | 13 + 9 files changed, 514 insertions(+), 31 deletions(-) create mode 100644 src/services/compact/__tests__/autoCompact.test.ts create mode 100644 src/utils/__tests__/context.test.ts diff --git a/bun.lock b/bun.lock index f1d66d4c0..ba08fc583 100644 --- a/bun.lock +++ b/bun.lock @@ -104,6 +104,7 @@ "he": "^1.2.0", "highlight.js": "^11.11.1", "https-proxy-agent": "^8.0.0", + "husky": "^9.1.7", "ignore": "^7.0.5", "image-processor-napi": "workspace:*", "indent-string": "^5.0.0", @@ -2065,6 +2066,8 @@ "human-signals": ["human-signals@8.0.1", "https://registry.npmmirror.com/human-signals/-/human-signals-8.0.1.tgz", {}, "sha512-eKCa6bwnJhvxj14kZk5NCPc6Hb6BdsU9DZcOnmQKSnO1VKrfV0zCvtttPZUsBvjmNDn8rpcJfpwSYnHBjc95MQ=="], + "husky": ["husky@9.1.7", "", { "bin": { "husky": "bin.js" } }, "sha512-5gs5ytaNjBrh5Ow3zrvdUUY+0VxIuWVL4i9irt6friV+BqdCfmV11CQTWMiBYWHbXhco+J1kHfTOUkePhCDvMA=="], + "iconv-lite": ["iconv-lite@0.7.2", "https://registry.npmmirror.com/iconv-lite/-/iconv-lite-0.7.2.tgz", { "dependencies": { "safer-buffer": ">= 2.1.2 < 3.0.0" } }, "sha512-im9DjEDQ55s9fL4EYzOAv0yMqmMBSZp6G0VvFyTMPKWxiSBHUj9NW/qqLmXUwXrrM7AvqSlTCfvqRb0cM8yYqw=="], "ignore": ["ignore@7.0.5", "https://registry.npmmirror.com/ignore/-/ignore-7.0.5.tgz", {}, "sha512-Hs59xBNfUIunMFgWAbGX5cq6893IbWg4KnrjbYwX3tx0ztorVgTDA6B2sxf8ejHJ4wz8BqGUMYlnzNBer5NvGg=="], diff --git a/package.json b/package.json index 9d5f00fbd..82133d852 100644 --- a/package.json +++ b/package.json @@ -170,6 +170,7 @@ "fflate": "^0.8.2", "figures": "^6.1.0", "fuse.js": "^7.1.0", + "husky": "^9.1.7", "get-east-asian-width": "^1.5.0", "google-auth-library": "^10.6.2", "he": "^1.2.0", diff --git a/src/components/Settings/Config.tsx b/src/components/Settings/Config.tsx index a2031bde2..66dbfedb6 100644 --- a/src/components/Settings/Config.tsx +++ b/src/components/Settings/Config.tsx @@ -302,7 +302,7 @@ export function Config({ { id: 'autoCompactEnabled', label: 'Auto-compact', - value: globalConfig.autoCompactEnabled, + value: globalConfig.autoCompactEnabled ?? true, type: 'boolean' as const, onChange(autoCompactEnabled: boolean) { saveGlobalConfig(current => ({ ...current, autoCompactEnabled })); diff --git a/src/query.ts b/src/query.ts index 9182a5b83..2e23a7065 100644 --- a/src/query.ts +++ b/src/query.ts @@ -114,7 +114,11 @@ import { } from './bootstrap/state.js' import { createBudgetTracker, checkTokenBudget } from './query/tokenBudget.js' import { count } from './utils/array.js' -import { createTrace, endTrace, isLangfuseEnabled } from './services/langfuse/index.js' +import { + createTrace, + endTrace, + isLangfuseEnabled, +} from './services/langfuse/index.js' import { getAPIProvider } from './utils/model/providers.js' import { uploadSessionTurn } from './utils/sessionDataUploader.js' @@ -133,7 +137,11 @@ function* yieldMissingToolResultBlocks( ) { for (const assistantMessage of assistantMessages) { // Extract all tool use blocks from this assistant message - const toolUseBlocks = (Array.isArray(assistantMessage.message?.content) ? assistantMessage.message.content : []).filter( + const toolUseBlocks = ( + Array.isArray(assistantMessage.message?.content) + ? assistantMessage.message.content + : [] + ).filter( (content: { type: string }) => content.type === 'tool_use', ) as ToolUseBlock[] @@ -242,8 +250,9 @@ export async function* query( logForDebugging( `[query] ownsTrace=${ownsTrace} incoming langfuseTrace=${params.toolUseContext.langfuseTrace ? 'present' : 'null/undefined'} isLangfuseEnabled=${isLangfuseEnabled()}`, ) - const langfuseTrace = params.toolUseContext.langfuseTrace - ?? (isLangfuseEnabled() + const langfuseTrace = + params.toolUseContext.langfuseTrace ?? + (isLangfuseEnabled() ? createTrace({ sessionId: getSessionId(), model: params.toolUseContext.options.mainLoopModel, @@ -497,20 +506,21 @@ async function* queryLoop( ) queryCheckpoint('query_autocompact_start') - const { compactionResult, consecutiveFailures } = await deps.autocompact( - messagesForQuery, - toolUseContext, - { - systemPrompt, - userContext, - systemContext, + const { compactionResult, consecutiveFailures, consecutiveCompactions } = + await deps.autocompact( + messagesForQuery, toolUseContext, - forkContextMessages: messagesForQuery, - }, - querySource, - tracking, - snipTokensFreed, - ) + { + systemPrompt, + userContext, + systemContext, + toolUseContext, + forkContextMessages: messagesForQuery, + }, + querySource, + tracking, + snipTokensFreed, + ) queryCheckpoint('query_autocompact_end') if (compactionResult) { @@ -564,11 +574,14 @@ async function* queryLoop( // compact. recompactionInfo (autoCompact.ts:190) already captured the // old values for turnsSincePreviousCompact/previousCompactTurnId before // the call, so this reset doesn't lose those. + // Preserve consecutiveCompactions so the circuit breaker can stop + // infinite compaction loops when context stays above threshold. tracking = { compacted: true, turnId: deps.uuid(), turnCounter: 0, consecutiveFailures: 0, + consecutiveCompactions, } const postCompactMessages = buildPostCompactMessages(compactionResult) @@ -794,7 +807,14 @@ async function* queryLoop( let yieldMessage: typeof message = message if (message.type === 'assistant') { const assistantMsg = message as AssistantMessage - const contentArr = Array.isArray(assistantMsg.message?.content) ? assistantMsg.message.content as unknown as Array<{ type: string; input?: unknown; name?: string; [key: string]: unknown }> : [] + const contentArr = Array.isArray(assistantMsg.message?.content) + ? (assistantMsg.message.content as unknown as Array<{ + type: string + input?: unknown + name?: string + [key: string]: unknown + }>) + : [] let clonedContent: typeof contentArr | undefined for (let i = 0; i < contentArr.length; i++) { const block = contentArr[i]! @@ -830,7 +850,10 @@ async function* queryLoop( if (clonedContent) { yieldMessage = { ...message, - message: { ...(assistantMsg.message ?? {}), content: clonedContent }, + message: { + ...(assistantMsg.message ?? {}), + content: clonedContent, + }, } as typeof message } } @@ -876,7 +899,11 @@ async function* queryLoop( const assistantMessage = message as AssistantMessage assistantMessages.push(assistantMessage) - const msgToolUseBlocks = (Array.isArray(assistantMessage.message?.content) ? assistantMessage.message.content : []).filter( + const msgToolUseBlocks = ( + Array.isArray(assistantMessage.message?.content) + ? assistantMessage.message.content + : [] + ).filter( (content: { type: string }) => content.type === 'tool_use', ) as ToolUseBlock[] if (msgToolUseBlocks.length > 0) { @@ -1016,7 +1043,10 @@ async function* queryLoop( logEvent('tengu_query_error', { assistantMessages: assistantMessages.length, toolUses: assistantMessages.flatMap(_ => - (Array.isArray(_.message?.content) ? _.message.content as Array<{ type: string }> : []).filter(content => content.type === 'tool_use'), + (Array.isArray(_.message?.content) + ? (_.message.content as Array<{ type: string }>) + : [] + ).filter(content => content.type === 'tool_use'), ).length, queryChainId: queryChainIdForAnalytics, @@ -1419,7 +1449,6 @@ async function* queryLoop( queryCheckpoint('query_tool_execution_start') - if (streamingToolExecutor) { logEvent('tengu_streaming_tool_execution_used', { tool_count: toolUseBlocks.length, @@ -1479,9 +1508,14 @@ async function* queryLoop( const lastAssistantMessage = assistantMessages.at(-1) let lastAssistantText: string | undefined if (lastAssistantMessage) { - const textBlocks = (Array.isArray(lastAssistantMessage.message?.content) ? lastAssistantMessage.message.content as Array<{ type: string; text?: string }> : []).filter( - block => block.type === 'text', - ) + const textBlocks = ( + Array.isArray(lastAssistantMessage.message?.content) + ? (lastAssistantMessage.message.content as Array<{ + type: string + text?: string + }>) + : [] + ).filter(block => block.type === 'text') if (textBlocks.length > 0) { const lastTextBlock = textBlocks.at(-1) if (lastTextBlock && 'text' in lastTextBlock) { @@ -1670,7 +1704,6 @@ async function* queryLoop( pendingMemoryPrefetch.consumedOnIteration = turnCount - 1 } - // Inject prefetched skill discovery. collectSkillDiscoveryPrefetch emits // hidden_by_main_turn — true when the prefetch resolved before this point // (should be >98% at AKI@250ms / Haiku@573ms vs turn durations of 2-30s). diff --git a/src/services/compact/__tests__/autoCompact.test.ts b/src/services/compact/__tests__/autoCompact.test.ts new file mode 100644 index 000000000..7800bb0cc --- /dev/null +++ b/src/services/compact/__tests__/autoCompact.test.ts @@ -0,0 +1,247 @@ +import { mock, describe, expect, test, beforeEach, afterEach } from 'bun:test' +import { logMock } from '../../../../tests/mocks/log.js' +import { debugMock } from '../../../../tests/mocks/debug.js' + +// Mock bun:bundle first (feature flags all off) +mock.module('bun:bundle', () => ({ feature: () => false })) + +// Mock side-effect modules to avoid bootstrap/state.ts init chain +mock.module('src/utils/log.ts', logMock) +mock.module('src/utils/debug.ts', debugMock) + +// Mock bootstrap/state.js to avoid realpathSync/randomUUID side effects +mock.module('src/bootstrap/state.js', () => ({ + markPostCompaction: () => {}, + getSdkBetas: () => [] as string[], + getSessionId: () => 'test-session-id', +})) + +// Mock config to control isAutoCompactEnabled behavior +let _mockAutoCompactEnabled: boolean | undefined = true +mock.module('src/utils/config.ts', () => ({ + getGlobalConfig: () => ({ + autoCompactEnabled: _mockAutoCompactEnabled, + showTurnDuration: false, + }), + isConfigEnabled: () => false, +})) + +// Mock context to avoid model resolution chain +mock.module('src/utils/context.ts', () => ({ + getContextWindowForModel: () => 200_000, + MODEL_CONTEXT_WINDOW_DEFAULT: 200_000, +})) + +// Mock tokens to avoid dependency on log.ts chain +mock.module('src/utils/tokens.ts', () => ({ + tokenCountWithEstimation: () => 1000, +})) + +// Mock analytics/growthbook +mock.module('src/services/analytics/growthbook.js', () => ({ + getFeatureValue_CACHED_MAY_BE_STALE: () => false, +})) + +// Mock API modules +mock.module('src/services/api/claude.js', () => ({ + getMaxOutputTokensForModel: () => 4096, +})) + +mock.module('src/services/api/promptCacheBreakDetection.js', () => ({ + notifyCompaction: () => {}, +})) + +mock.module('src/services/SessionMemory/sessionMemoryUtils.js', () => ({ + setLastSummarizedMessageId: () => {}, +})) + +mock.module('../compact.js', () => ({ + compactConversation: async () => ({ + summary: 'test summary', + messages: [], + }), + ERROR_MESSAGE_USER_ABORT: 'User aborted compaction', +})) + +mock.module('../postCompactCleanup.js', () => ({ + runPostCompactCleanup: () => {}, +})) + +mock.module('../sessionMemoryCompact.js', () => ({ + trySessionMemoryCompaction: async () => null, +})) + +// Dynamic import after all mocks are in place +const { + AUTOCOMPACT_BUFFER_TOKENS, + MANUAL_COMPACT_BUFFER_TOKENS, + isAutoCompactEnabled, + autoCompactIfNeeded, +} = await import('../autoCompact.js') + +// --- Helpers --- + +function makeToolUseContext(model = 'claude-sonnet-4-20250514') { + return { options: { mainLoopModel: model } } as any +} + +function makeCacheSafeParams(): any { + return { + systemPrompt: [''], + userContext: {}, + systemContext: {}, + toolUseContext: {}, + forkContextMessages: [], + } +} + +// --- Constants --- + +describe('AUTOCOMPACT_BUFFER_TOKENS', () => { + test('is 25_000 (increased from 13K for system prompt + tool definition headroom)', () => { + expect(AUTOCOMPACT_BUFFER_TOKENS).toBe(25_000) + }) +}) + +describe('MANUAL_COMPACT_BUFFER_TOKENS', () => { + test('is 10_000 (increased from 3K for /compact headroom)', () => { + expect(MANUAL_COMPACT_BUFFER_TOKENS).toBe(10_000) + }) +}) + +// --- isAutoCompactEnabled --- + +describe('isAutoCompactEnabled', () => { + const originalDisableCompact = process.env.DISABLE_COMPACT + const originalDisableAutoCompact = process.env.DISABLE_AUTO_COMPACT + + afterEach(() => { + process.env.DISABLE_COMPACT = originalDisableCompact + process.env.DISABLE_AUTO_COMPACT = originalDisableAutoCompact + _mockAutoCompactEnabled = true + }) + + test('returns true by default when autoCompactEnabled is undefined', () => { + _mockAutoCompactEnabled = undefined + delete process.env.DISABLE_COMPACT + delete process.env.DISABLE_AUTO_COMPACT + expect(isAutoCompactEnabled()).toBe(true) + }) + + test('returns true when autoCompactEnabled is explicitly true', () => { + _mockAutoCompactEnabled = true + delete process.env.DISABLE_COMPACT + delete process.env.DISABLE_AUTO_COMPACT + expect(isAutoCompactEnabled()).toBe(true) + }) + + test('returns false when autoCompactEnabled is false', () => { + _mockAutoCompactEnabled = false + delete process.env.DISABLE_COMPACT + delete process.env.DISABLE_AUTO_COMPACT + expect(isAutoCompactEnabled()).toBe(false) + }) + + test('returns false when DISABLE_COMPACT is set', () => { + _mockAutoCompactEnabled = true + process.env.DISABLE_COMPACT = '1' + delete process.env.DISABLE_AUTO_COMPACT + expect(isAutoCompactEnabled()).toBe(false) + }) + + test('returns false when DISABLE_AUTO_COMPACT is set', () => { + _mockAutoCompactEnabled = true + delete process.env.DISABLE_COMPACT + process.env.DISABLE_AUTO_COMPACT = '1' + expect(isAutoCompactEnabled()).toBe(false) + }) + + test('DISABLE_COMPACT takes precedence over autoCompactEnabled setting', () => { + _mockAutoCompactEnabled = true + process.env.DISABLE_COMPACT = '1' + expect(isAutoCompactEnabled()).toBe(false) + }) +}) + +// --- autoCompactIfNeeded --- + +describe('autoCompactIfNeeded', () => { + const originalDisableCompact = process.env.DISABLE_COMPACT + + beforeEach(() => { + delete process.env.DISABLE_COMPACT + _mockAutoCompactEnabled = true + }) + + afterEach(() => { + process.env.DISABLE_COMPACT = originalDisableCompact + _mockAutoCompactEnabled = true + }) + + test('returns wasCompacted=false when DISABLE_COMPACT is set', async () => { + process.env.DISABLE_COMPACT = '1' + const result = await autoCompactIfNeeded( + [], + makeToolUseContext(), + makeCacheSafeParams(), + ) + expect(result.wasCompacted).toBe(false) + }) + + test('trips consecutiveCompactions circuit breaker at limit (2)', async () => { + const result = await autoCompactIfNeeded( + [], + makeToolUseContext(), + makeCacheSafeParams(), + undefined, + { compacted: true, turnCounter: 5, turnId: 't1', consecutiveCompactions: 2 }, + ) + expect(result.wasCompacted).toBe(false) + }) + + test('allows compaction when consecutiveCompactions is below limit', async () => { + const result = await autoCompactIfNeeded( + [], + makeToolUseContext(), + makeCacheSafeParams(), + undefined, + { compacted: false, turnCounter: 0, turnId: 't1', consecutiveCompactions: 1 }, + ) + // Passes the circuit breaker check but shouldAutoCompact returns false (low token count) + expect(result.wasCompacted).toBe(false) + expect(result.consecutiveCompactions).toBeUndefined() + }) + + test('trips consecutiveFailures circuit breaker at limit (3)', async () => { + const result = await autoCompactIfNeeded( + [], + makeToolUseContext(), + makeCacheSafeParams(), + undefined, + { compacted: false, turnCounter: 0, turnId: 't1', consecutiveFailures: 3 }, + ) + expect(result.wasCompacted).toBe(false) + }) + + test('allows attempt when consecutiveFailures is below limit', async () => { + const result = await autoCompactIfNeeded( + [], + makeToolUseContext(), + makeCacheSafeParams(), + undefined, + { compacted: false, turnCounter: 0, turnId: 't1', consecutiveFailures: 2 }, + ) + // Passes the failure circuit breaker but shouldAutoCompact returns false (low tokens) + expect(result.wasCompacted).toBe(false) + }) + + test('returns wasCompacted=false when tracking is undefined', async () => { + const result = await autoCompactIfNeeded( + [], + makeToolUseContext(), + makeCacheSafeParams(), + ) + // No tracking → no circuit breakers tripped, but shouldAutoCompact returns false + expect(result.wasCompacted).toBe(false) + }) +}) diff --git a/src/services/compact/autoCompact.ts b/src/services/compact/autoCompact.ts index b77fb8b61..8ba19310e 100644 --- a/src/services/compact/autoCompact.ts +++ b/src/services/compact/autoCompact.ts @@ -57,12 +57,25 @@ export type AutoCompactTrackingState = { // Used as a circuit breaker to stop retrying when the context is // irrecoverably over the limit (e.g., prompt_too_long). consecutiveFailures?: number + // Consecutive successful autocompact runs. Reset when the user sends a + // new message (turnCounter resets in query.ts on user-driven turns). + // Prevents infinite autocompact loops when post-compact token count + // remains above the threshold (e.g. large CLAUDE.md re-injection). + consecutiveCompactions?: number } -export const AUTOCOMPACT_BUFFER_TOKENS = 13_000 +// Increased from 13_000 to 25_000 to account for system prompt (~15-25K) and +// tool definitions (~5-20K) that are NOT counted in message-body token estimates +// but ARE part of every API request. The old 13K buffer was routinely exhausted +// before autocompact could fire, causing prompt-too-long errors in long sessions. +export const AUTOCOMPACT_BUFFER_TOKENS = 25_000 export const WARNING_THRESHOLD_BUFFER_TOKENS = 20_000 export const ERROR_THRESHOLD_BUFFER_TOKENS = 20_000 -export const MANUAL_COMPACT_BUFFER_TOKENS = 3_000 +// Increased from 3_000 to 10_000 so manual /compact has enough headroom to +// successfully send the full conversation to the summary model. At 3K, the +// compact API call itself would often hit prompt-too-long, making /compact +// unusable right when it was most needed. +export const MANUAL_COMPACT_BUFFER_TOKENS = 10_000 // Conservative estimate for tool result growth per turn. // Typical tool results (file reads, grep, bash) average ~5-10K tokens; @@ -98,6 +111,13 @@ export function estimateMaxTurnGrowth(model: string): number { // in a single session, wasting ~250K API calls/day globally. const MAX_CONSECUTIVE_AUTOCOMPACT_FAILURES = 3 +// Stop trying autocompact after this many consecutive successful compactions +// on the same turn chain. A successful compaction that doesn't reduce context +// below the threshold is just as wasteful as a failed one — each attempt burns +// a full API call (~15K output tokens) for no net progress. After the limit, +// the system falls through to reactive compact on actual prompt-too-long errors. +const MAX_CONSECUTIVE_AUTOCOMPACT_COMPACTIONS = 2 + export function getAutoCompactThreshold(model: string): number { const effectiveContextWindow = getEffectiveContextWindowSize(model) @@ -183,7 +203,7 @@ export function isAutoCompactEnabled(): boolean { } // Check if user has disabled auto-compact in their settings const userConfig = getGlobalConfig() - return userConfig.autoCompactEnabled + return userConfig.autoCompactEnabled ?? true } export async function shouldAutoCompact( @@ -278,6 +298,7 @@ export async function autoCompactIfNeeded( wasCompacted: boolean compactionResult?: CompactionResult consecutiveFailures?: number + consecutiveCompactions?: number }> { if (isEnvTruthy(process.env.DISABLE_COMPACT)) { return { wasCompacted: false } @@ -293,6 +314,23 @@ export async function autoCompactIfNeeded( return { wasCompacted: false } } + // Circuit breaker: stop after N consecutive successful compactions on the + // same turn chain. A successful compaction that doesn't reduce context below + // the autocompact threshold will immediately re-trigger shouldAutoCompact on + // the next iteration, creating an infinite loop that burns API calls with no + // user-visible progress. Common causes: large CLAUDE.md re-injection, many + // MCP tools, or post-compact attachments exceeding the budget. + if ( + tracking?.consecutiveCompactions !== undefined && + tracking.consecutiveCompactions >= MAX_CONSECUTIVE_AUTOCOMPACT_COMPACTIONS + ) { + logForDebugging( + `autocompact: consecutive-compactions circuit breaker tripped after ${tracking.consecutiveCompactions} successful compactions — context still above threshold, falling through to reactive/recovery path`, + { level: 'warn' }, + ) + return { wasCompacted: false } + } + const model = toolUseContext.options.mainLoopModel const shouldCompact = await shouldAutoCompact( messages, @@ -335,6 +373,8 @@ export async function autoCompactIfNeeded( return { wasCompacted: true, compactionResult: sessionMemoryResult, + consecutiveFailures: 0, + consecutiveCompactions: (tracking?.consecutiveCompactions ?? 0) + 1, } } @@ -359,6 +399,7 @@ export async function autoCompactIfNeeded( compactionResult, // Reset failure count on success consecutiveFailures: 0, + consecutiveCompactions: (tracking?.consecutiveCompactions ?? 0) + 1, } } catch (error) { if (!hasExactErrorMessage(error, ERROR_MESSAGE_USER_ABORT)) { diff --git a/src/utils/__tests__/context.test.ts b/src/utils/__tests__/context.test.ts new file mode 100644 index 000000000..fe2dcc395 --- /dev/null +++ b/src/utils/__tests__/context.test.ts @@ -0,0 +1,145 @@ +import { mock, describe, expect, test, beforeEach, afterEach } from 'bun:test' +import { logMock } from '../../../tests/mocks/log.js' +import { debugMock } from '../../../tests/mocks/debug.js' + +// Mock bun:bundle (feature flags all off) +mock.module('bun:bundle', () => ({ feature: () => false })) + +// Mock side-effect modules +mock.module('src/utils/log.ts', logMock) +mock.module('src/utils/debug.ts', debugMock) + +// Mock bootstrap/state.js to avoid realpathSync/randomUUID side effects +mock.module('src/bootstrap/state.js', () => ({ + markPostCompaction: () => {}, + getSdkBetas: () => [] as string[], + getSessionId: () => 'test-session-id', +})) + +// Mock config to avoid file system reads +mock.module('src/utils/config.ts', () => ({ + getGlobalConfig: () => ({ + showTurnDuration: false, + clientDataCache: {}, + }), +})) + +// Mock model modules to control resolution chain +mock.module('src/utils/model/model.js', () => ({ + getCanonicalName: (m: string) => m, +})) + +mock.module('src/utils/model/antModels.js', () => ({ + resolveAntModel: () => null, +})) + +// Control what getCachedCoStrictModels returns via mutable state +let _mockCoStrictModels: Array<{ id: string; contextWindow?: number }> = [] +mock.module('src/costrict/provider/models.js', () => ({ + getCachedCoStrictModels: () => _mockCoStrictModels, +})) + +// Control what resolveCoStrictModel returns +let _mockResolveCoStrictModelResult = '' +mock.module('src/costrict/provider/modelMapping.js', () => ({ + resolveCoStrictModel: (model: string) => + _mockResolveCoStrictModelResult || model, +})) + +mock.module('src/utils/model/modelCapabilities.js', () => ({ + getModelCapability: () => null, +})) + +// Dynamic import after all mocks +const { getContextWindowForModel, MODEL_CONTEXT_WINDOW_DEFAULT } = + await import('../context.js') + +// --- getContextWindowForModel (CoStrict model support) --- + +describe('getContextWindowForModel', () => { + const originalUserType = process.env.USER_TYPE + + beforeEach(() => { + // Ensure USER_TYPE is not 'ant' so ant-only paths are skipped + delete process.env.USER_TYPE + delete process.env.CLAUDE_CODE_MAX_CONTEXT_TOKENS + delete process.env.CLAUDE_CODE_DISABLE_1M_CONTEXT + _mockCoStrictModels = [] + _mockResolveCoStrictModelResult = '' + }) + + afterEach(() => { + process.env.USER_TYPE = originalUserType + _mockCoStrictModels = [] + _mockResolveCoStrictModelResult = '' + }) + + test('returns default context window when no CoStrict models cached', () => { + _mockCoStrictModels = [] + const result = getContextWindowForModel('some-unknown-model') + expect(result).toBe(MODEL_CONTEXT_WINDOW_DEFAULT) + }) + + test('returns CoStrict model contextWindow when model is found by exact id', () => { + _mockCoStrictModels = [ + { id: 'costrict-model-x', contextWindow: 128_000 }, + { id: 'costrict-model-y', contextWindow: 256_000 }, + ] + const result = getContextWindowForModel('costrict-model-x') + expect(result).toBe(128_000) + }) + + test('returns CoStrict model contextWindow when resolved via resolveCoStrictModel', () => { + _mockCoStrictModels = [ + { id: 'costrict-mapped-model', contextWindow: 512_000 }, + ] + _mockResolveCoStrictModelResult = 'costrict-mapped-model' + const result = getContextWindowForModel('claude-sonnet-4') + expect(result).toBe(512_000) + }) + + test('prefers exact id match over resolved name when both exist', () => { + _mockCoStrictModels = [ + { id: 'claude-sonnet-4', contextWindow: 300_000 }, + { id: 'costrict-mapped', contextWindow: 400_000 }, + ] + _mockResolveCoStrictModelResult = 'costrict-mapped' + // exact id 'claude-sonnet-4' matches first element + const result = getContextWindowForModel('claude-sonnet-4') + expect(result).toBe(300_000) + }) + + test('skips CoStrict model with zero contextWindow', () => { + _mockCoStrictModels = [ + { id: 'bad-model', contextWindow: 0 }, + ] + const result = getContextWindowForModel('bad-model') + expect(result).toBe(MODEL_CONTEXT_WINDOW_DEFAULT) + }) + + test('skips CoStrict model with undefined contextWindow', () => { + _mockCoStrictModels = [ + { id: 'no-window-model' }, + ] + const result = getContextWindowForModel('no-window-model') + expect(result).toBe(MODEL_CONTEXT_WINDOW_DEFAULT) + }) + + test('skips CoStrict model with negative contextWindow', () => { + _mockCoStrictModels = [ + { id: 'negative-model', contextWindow: -1 }, + ] + const result = getContextWindowForModel('negative-model') + expect(result).toBe(MODEL_CONTEXT_WINDOW_DEFAULT) + }) + + test('returns default when CoStrict models exist but none match', () => { + _mockCoStrictModels = [ + { id: 'other-model-a', contextWindow: 100_000 }, + { id: 'other-model-b', contextWindow: 200_000 }, + ] + _mockResolveCoStrictModelResult = 'yet-another-model' + const result = getContextWindowForModel('nothing-matches') + expect(result).toBe(MODEL_CONTEXT_WINDOW_DEFAULT) + }) +}) diff --git a/src/utils/config.ts b/src/utils/config.ts index 1f6a89cd4..7e5425ad5 100644 --- a/src/utils/config.ts +++ b/src/utils/config.ts @@ -241,7 +241,7 @@ export type GlobalConfig = { editorMode?: EditorMode bypassPermissionsModeAccepted?: boolean hasUsedBackslashReturn?: boolean - autoCompactEnabled: boolean // Controls whether auto-compact is enabled + autoCompactEnabled?: boolean // Controls whether auto-compact is enabled showTurnDuration: boolean // Controls whether to show turn duration message (e.g., "Cooked for 1m 6s") /** * @deprecated Use settings.env instead. diff --git a/src/utils/context.ts b/src/utils/context.ts index 1c0680bd5..abe5e2ecb 100644 --- a/src/utils/context.ts +++ b/src/utils/context.ts @@ -5,6 +5,8 @@ import { isEnvTruthy } from './envUtils.js' import { getCanonicalName } from './model/model.js' import { resolveAntModel } from './model/antModels.js' import { getModelCapability } from './model/modelCapabilities.js' +import { getCachedCoStrictModels } from '../costrict/provider/models.js' +import { resolveCoStrictModel } from '../costrict/provider/modelMapping.js' // Model context window size (200k tokens for all models right now) export const MODEL_CONTEXT_WINDOW_DEFAULT = 200_000 @@ -99,6 +101,17 @@ export function getContextWindowForModel( return antModel.contextWindow } } + // CoStrict 模型：从运行时缓存获取 contextWindow + const costrictModels = getCachedCoStrictModels() + if (costrictModels.length > 0) { + const costrictModelId = resolveCoStrictModel(model) + const costrictModel = costrictModels.find( + m => m.id === model || m.id === costrictModelId, + ) + if (costrictModel?.contextWindow && costrictModel.contextWindow > 0) { + return costrictModel.contextWindow + } + } return MODEL_CONTEXT_WINDOW_DEFAULT }