Skip to content

Commit 51aeabc

Browse files
hannesrudolphhannesrudolph
committed
refactor: centralize context management threshold logic with willManageContext helper
- Added willManageContext() helper function to src/core/context-management/index.ts - Simplified Task.ts by using the helper instead of duplicated threshold calculation - Added 7 comprehensive tests for the new helper function - All 53 context management tests passing Addresses PR review feedback about duplicated logic in Task.ts lines 3519-3543
1 parent a6f5626 commit 51aeabc

File tree

3 files changed

+204
-22
lines changed

3 files changed

+204
-22
lines changed

‎src/core/context-management/__tests__/context-management.spec.ts‎

Lines changed: 128 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,13 @@ import { BaseProvider } from "../../../api/providers/base-provider"
99
import { ApiMessage } from "../../task-persistence/apiMessages"
1010
import * as condenseModule from "../../condense"
1111

12-
import { TOKEN_BUFFER_PERCENTAGE, estimateTokenCount, truncateConversation, manageContext } from "../index"
12+
import {
13+
TOKEN_BUFFER_PERCENTAGE,
14+
estimateTokenCount,
15+
truncateConversation,
16+
manageContext,
17+
willManageContext,
18+
} from "../index"
1319

1420
// Create a mock ApiHandler for testing
1521
class MockApiHandler extends BaseProvider {
@@ -1280,4 +1286,125 @@ describe("Context Management", () => {
12801286
expect(result2.truncationId).toBeDefined()
12811287
})
12821288
})
1289+
1290+
/**
1291+
* Tests for the willManageContext helper function
1292+
*/
1293+
describe("willManageContext", () => {
1294+
it("should return true when context percent exceeds threshold", () => {
1295+
const result = willManageContext({
1296+
totalTokens: 60000,
1297+
contextWindow: 100000, // 60% of context window
1298+
maxTokens: 30000,
1299+
autoCondenseContext: true,
1300+
autoCondenseContextPercent: 50, // 50% threshold
1301+
profileThresholds: {},
1302+
currentProfileId: "default",
1303+
lastMessageTokens: 0,
1304+
})
1305+
expect(result).toBe(true)
1306+
})
1307+
1308+
it("should return false when context percent is below threshold", () => {
1309+
const result = willManageContext({
1310+
totalTokens: 40000,
1311+
contextWindow: 100000, // 40% of context window
1312+
maxTokens: 30000,
1313+
autoCondenseContext: true,
1314+
autoCondenseContextPercent: 50, // 50% threshold
1315+
profileThresholds: {},
1316+
currentProfileId: "default",
1317+
lastMessageTokens: 0,
1318+
})
1319+
expect(result).toBe(false)
1320+
})
1321+
1322+
it("should return true when tokens exceed allowedTokens even if autoCondenseContext is false", () => {
1323+
// allowedTokens = contextWindow * (1 - 0.1) - reservedTokens = 100000 * 0.9 - 30000 = 60000
1324+
const result = willManageContext({
1325+
totalTokens: 60001, // Exceeds allowedTokens
1326+
contextWindow: 100000,
1327+
maxTokens: 30000,
1328+
autoCondenseContext: false, // Even with auto-condense disabled
1329+
autoCondenseContextPercent: 50,
1330+
profileThresholds: {},
1331+
currentProfileId: "default",
1332+
lastMessageTokens: 0,
1333+
})
1334+
expect(result).toBe(true)
1335+
})
1336+
1337+
it("should return false when autoCondenseContext is false and tokens are below allowedTokens", () => {
1338+
// allowedTokens = contextWindow * (1 - 0.1) - reservedTokens = 100000 * 0.9 - 30000 = 60000
1339+
const result = willManageContext({
1340+
totalTokens: 59999, // Below allowedTokens
1341+
contextWindow: 100000,
1342+
maxTokens: 30000,
1343+
autoCondenseContext: false,
1344+
autoCondenseContextPercent: 50, // This shouldn't matter since autoCondenseContext is false
1345+
profileThresholds: {},
1346+
currentProfileId: "default",
1347+
lastMessageTokens: 0,
1348+
})
1349+
expect(result).toBe(false)
1350+
})
1351+
1352+
it("should use profile-specific threshold when available", () => {
1353+
const result = willManageContext({
1354+
totalTokens: 55000,
1355+
contextWindow: 100000, // 55% of context window
1356+
maxTokens: 30000,
1357+
autoCondenseContext: true,
1358+
autoCondenseContextPercent: 80, // Global threshold 80%
1359+
profileThresholds: { "test-profile": 50 }, // Profile threshold 50%
1360+
currentProfileId: "test-profile",
1361+
lastMessageTokens: 0,
1362+
})
1363+
// Should trigger because 55% > 50% (profile threshold)
1364+
expect(result).toBe(true)
1365+
})
1366+
1367+
it("should fall back to global threshold when profile threshold is -1", () => {
1368+
const result = willManageContext({
1369+
totalTokens: 55000,
1370+
contextWindow: 100000, // 55% of context window
1371+
maxTokens: 30000,
1372+
autoCondenseContext: true,
1373+
autoCondenseContextPercent: 80, // Global threshold 80%
1374+
profileThresholds: { "test-profile": -1 }, // Profile uses global
1375+
currentProfileId: "test-profile",
1376+
lastMessageTokens: 0,
1377+
})
1378+
// Should NOT trigger because 55% < 80% (global threshold)
1379+
expect(result).toBe(false)
1380+
})
1381+
1382+
it("should include lastMessageTokens in the calculation", () => {
1383+
// Without lastMessageTokens: 49000 tokens = 49%
1384+
// With lastMessageTokens: 49000 + 2000 = 51000 tokens = 51%
1385+
const resultWithoutLastMessage = willManageContext({
1386+
totalTokens: 49000,
1387+
contextWindow: 100000,
1388+
maxTokens: 30000,
1389+
autoCondenseContext: true,
1390+
autoCondenseContextPercent: 50, // 50% threshold
1391+
profileThresholds: {},
1392+
currentProfileId: "default",
1393+
lastMessageTokens: 0,
1394+
})
1395+
expect(resultWithoutLastMessage).toBe(false)
1396+
1397+
const resultWithLastMessage = willManageContext({
1398+
totalTokens: 49000,
1399+
contextWindow: 100000,
1400+
maxTokens: 30000,
1401+
autoCondenseContext: true,
1402+
autoCondenseContextPercent: 50, // 50% threshold
1403+
profileThresholds: {},
1404+
currentProfileId: "default",
1405+
lastMessageTokens: 2000, // Pushes total to 51%
1406+
})
1407+
expect(resultWithLastMessage).toBe(true)
1408+
})
1409+
})
12831410
})

‎src/core/context-management/index.ts‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -133,6 +133,68 @@ export function truncateConversation(messages: ApiMessage[], fracToRemove: numbe
133133
}
134134
}
135135

136+
/**
137+
* Options for checking if context management will likely run.
138+
* A subset of ContextManagementOptions with only the fields needed for threshold calculation.
139+
*/
140+
export type WillManageContextOptions = {
141+
totalTokens: number
142+
contextWindow: number
143+
maxTokens?: number | null
144+
autoCondenseContext: boolean
145+
autoCondenseContextPercent: number
146+
profileThresholds: Record<string, number>
147+
currentProfileId: string
148+
lastMessageTokens: number
149+
}
150+
151+
/**
152+
* Checks whether context management (condensation or truncation) will likely run based on current token usage.
153+
*
154+
* This is useful for showing UI indicators before `manageContext` is actually called,
155+
* without duplicating the threshold calculation logic.
156+
*
157+
* @param {WillManageContextOptions} options - The options for threshold calculation
158+
* @returns {boolean} True if context management will likely run, false otherwise
159+
*/
160+
export function willManageContext({
161+
totalTokens,
162+
contextWindow,
163+
maxTokens,
164+
autoCondenseContext,
165+
autoCondenseContextPercent,
166+
profileThresholds,
167+
currentProfileId,
168+
lastMessageTokens,
169+
}: WillManageContextOptions): boolean {
170+
if (!autoCondenseContext) {
171+
// When auto-condense is disabled, only truncation can occur
172+
const reservedTokens = maxTokens || ANTHROPIC_DEFAULT_MAX_TOKENS
173+
const prevContextTokens = totalTokens + lastMessageTokens
174+
const allowedTokens = contextWindow * (1 - TOKEN_BUFFER_PERCENTAGE) - reservedTokens
175+
return prevContextTokens > allowedTokens
176+
}
177+
178+
const reservedTokens = maxTokens || ANTHROPIC_DEFAULT_MAX_TOKENS
179+
const prevContextTokens = totalTokens + lastMessageTokens
180+
const allowedTokens = contextWindow * (1 - TOKEN_BUFFER_PERCENTAGE) - reservedTokens
181+
182+
// Determine the effective threshold to use
183+
let effectiveThreshold = autoCondenseContextPercent
184+
const profileThreshold = profileThresholds[currentProfileId]
185+
if (profileThreshold !== undefined) {
186+
if (profileThreshold === -1) {
187+
effectiveThreshold = autoCondenseContextPercent
188+
} else if (profileThreshold >= MIN_CONDENSE_THRESHOLD && profileThreshold <= MAX_CONDENSE_THRESHOLD) {
189+
effectiveThreshold = profileThreshold
190+
}
191+
// Invalid values fall back to global setting (effectiveThreshold already set)
192+
}
193+
194+
const contextPercent = (100 * prevContextTokens) / contextWindow
195+
return contextPercent >= effectiveThreshold || prevContextTokens > allowedTokens
196+
}
197+
136198
/**
137199
* Context Management: Conditionally manages the conversation context when approaching limits.
138200
*

‎src/core/task/Task.ts‎

Lines changed: 14 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -99,7 +99,7 @@ import { RooProtectedController } from "../protect/RooProtectedController"
9999
import { type AssistantMessageContent, presentAssistantMessage } from "../assistant-message"
100100
import { AssistantMessageParser } from "../assistant-message/AssistantMessageParser"
101101
import { NativeToolCallParser } from "../assistant-message/NativeToolCallParser"
102-
import { manageContext } from "../context-management"
102+
import { manageContext, willManageContext } from "../context-management"
103103
import { ClineProvider } from "../webview/ClineProvider"
104104
import { MultiSearchReplaceDiffStrategy } from "../diff/strategies/multi-search-replace"
105105
import { MultiFileSearchReplaceDiffStrategy } from "../diff/strategies/multi-file-search-replace"
@@ -3514,8 +3514,7 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
35143514

35153515
// Check if context management will likely run (threshold check)
35163516
// This allows us to show an in-progress indicator to the user
3517-
// Important: Match the exact calculation from manageContext to avoid threshold mismatch
3518-
// manageContext uses: prevContextTokens = totalTokens + lastMessageTokens
3517+
// We use the centralized willManageContext helper to avoid duplicating threshold logic
35193518
const lastMessage = this.apiConversationHistory[this.apiConversationHistory.length - 1]
35203519
const lastMessageContent = lastMessage?.content
35213520
let lastMessageTokens = 0
@@ -3524,28 +3523,22 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
35243523
? await this.api.countTokens(lastMessageContent)
35253524
: await this.api.countTokens([{ type: "text", text: lastMessageContent as string }])
35263525
}
3527-
const prevContextTokens = contextTokens + lastMessageTokens
3528-
const estimatedUsagePercent = (100 * prevContextTokens) / contextWindow
35293526

3530-
// Match manageContext's threshold logic
3531-
const profileThresholdSettings = profileThresholds[currentProfileId] as
3532-
| { autoCondenseContextPercent?: number }
3533-
| undefined
3534-
const effectiveThreshold =
3535-
profileThresholdSettings?.autoCondenseContextPercent ?? autoCondenseContextPercent
3536-
3537-
// Calculate allowedTokens the same way manageContext does
3538-
const TOKEN_BUFFER_PERCENTAGE = 0.1
3539-
const reservedTokens = maxTokens ?? 8192 // ANTHROPIC_DEFAULT_MAX_TOKENS
3540-
const allowedTokens = contextWindow * (1 - TOKEN_BUFFER_PERCENTAGE) - reservedTokens
3541-
3542-
// Match manageContext's condition: contextPercent >= effectiveThreshold || prevContextTokens > allowedTokens
3543-
const willManageContext = estimatedUsagePercent >= effectiveThreshold || prevContextTokens > allowedTokens
3527+
const contextManagementWillRun = willManageContext({
3528+
totalTokens: contextTokens,
3529+
contextWindow,
3530+
maxTokens,
3531+
autoCondenseContext,
3532+
autoCondenseContextPercent,
3533+
profileThresholds,
3534+
currentProfileId,
3535+
lastMessageTokens,
3536+
})
35443537

35453538
// Send condenseTaskContextStarted BEFORE manageContext to show in-progress indicator
35463539
// This notification must be sent here (not earlier) because the early check uses stale token count
35473540
// (before user message is added to history), which could incorrectly skip showing the indicator
3548-
if (willManageContext && autoCondenseContext) {
3541+
if (contextManagementWillRun && autoCondenseContext) {
35493542
await this.providerRef
35503543
.deref()
35513544
?.postMessageToWebview({ type: "condenseTaskContextStarted", text: this.taskId })
@@ -3614,7 +3607,7 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
36143607

36153608
// Notify webview that context management is complete (sets isCondensing = false)
36163609
// This removes the in-progress spinner and allows the completed result to show
3617-
if (willManageContext && autoCondenseContext) {
3610+
if (contextManagementWillRun && autoCondenseContext) {
36183611
await this.providerRef
36193612
.deref()
36203613
?.postMessageToWebview({ type: "condenseTaskContextResponse", text: this.taskId })

0 commit comments

Comments
 (0)