Skip to content

Commit d7ae811

Browse files
authored
[Condense] Track telemetry for condense and truncate operations (RooCodeInc#3796)
* [Condense] Track telemetry for condense and truncate operations * update tests * test fix nits
1 parent fc9a42d commit d7ae811

File tree

7 files changed

+63
-17
lines changed

7 files changed

+63
-17
lines changed

src/core/condense/__tests__/index.test.ts

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@ jest.mock("../../../api/transform/image-cleaning", () => ({
99
maybeRemoveImageBlocks: jest.fn((messages: ApiMessage[], _apiHandler: ApiHandler) => [...messages]),
1010
}))
1111

12+
const taskId = "test-task-id"
13+
1214
describe("getMessagesSinceLastSummary", () => {
1315
it("should return all messages when there is no summary", () => {
1416
const messages: ApiMessage[] = [
@@ -106,7 +108,7 @@ describe("summarizeConversation", () => {
106108
{ role: "assistant", content: "Hi there", ts: 2 },
107109
]
108110

109-
const result = await summarizeConversation(messages, mockApiHandler, defaultSystemPrompt)
111+
const result = await summarizeConversation(messages, mockApiHandler, defaultSystemPrompt, taskId)
110112
expect(result.messages).toEqual(messages)
111113
expect(result.cost).toBe(0)
112114
expect(result.summary).toBe("")
@@ -125,7 +127,7 @@ describe("summarizeConversation", () => {
125127
{ role: "user", content: "Tell me more", ts: 7 },
126128
]
127129

128-
const result = await summarizeConversation(messages, mockApiHandler, defaultSystemPrompt)
130+
const result = await summarizeConversation(messages, mockApiHandler, defaultSystemPrompt, taskId)
129131
expect(result.messages).toEqual(messages)
130132
expect(result.cost).toBe(0)
131133
expect(result.summary).toBe("")
@@ -144,7 +146,7 @@ describe("summarizeConversation", () => {
144146
{ role: "user", content: "Tell me more", ts: 7 },
145147
]
146148

147-
const result = await summarizeConversation(messages, mockApiHandler, defaultSystemPrompt)
149+
const result = await summarizeConversation(messages, mockApiHandler, defaultSystemPrompt, taskId)
148150

149151
// Check that the API was called correctly
150152
expect(mockApiHandler.createMessage).toHaveBeenCalled()
@@ -202,7 +204,7 @@ describe("summarizeConversation", () => {
202204
return messages.map(({ role, content }: { role: string; content: any }) => ({ role, content }))
203205
})
204206

205-
const result = await summarizeConversation(messages, mockApiHandler, defaultSystemPrompt)
207+
const result = await summarizeConversation(messages, mockApiHandler, defaultSystemPrompt, taskId)
206208

207209
// Should return original messages when summary is empty
208210
expect(result.messages).toEqual(messages)
@@ -225,7 +227,7 @@ describe("summarizeConversation", () => {
225227
{ role: "user", content: "Tell me more", ts: 7 },
226228
]
227229

228-
await summarizeConversation(messages, mockApiHandler, defaultSystemPrompt)
230+
await summarizeConversation(messages, mockApiHandler, defaultSystemPrompt, taskId)
229231

230232
// Verify the final request message
231233
const expectedFinalMessage = {
@@ -266,7 +268,7 @@ describe("summarizeConversation", () => {
266268
// Override the mock for this test
267269
mockApiHandler.createMessage = jest.fn().mockReturnValue(streamWithUsage) as any
268270

269-
const result = await summarizeConversation(messages, mockApiHandler, systemPrompt)
271+
const result = await summarizeConversation(messages, mockApiHandler, systemPrompt, taskId)
270272

271273
// Verify that countTokens was called with the correct messages including system prompt
272274
expect(mockApiHandler.countTokens).toHaveBeenCalled()

src/core/condense/index.ts

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import Anthropic from "@anthropic-ai/sdk"
22
import { ApiHandler } from "../../api"
33
import { ApiMessage } from "../task-persistence/apiMessages"
44
import { maybeRemoveImageBlocks } from "../../api/transform/image-cleaning"
5+
import { telemetryService } from "../../services/telemetry/TelemetryService"
56

67
export const N_MESSAGES_TO_KEEP = 3
78

@@ -58,13 +59,16 @@ export type SummarizeResponse = {
5859
* @param {ApiMessage[]} messages - The conversation messages
5960
* @param {ApiHandler} apiHandler - The API handler to use for token counting.
6061
* @param {string} systemPrompt - The system prompt for API requests, which should be considered in the context token count
62+
* @param {string} taskId - The task ID for the conversation, used for telemetry
6163
* @returns {SummarizeResponse} - The result of the summarization operation (see above)
6264
*/
6365
export async function summarizeConversation(
6466
messages: ApiMessage[],
6567
apiHandler: ApiHandler,
6668
systemPrompt: string,
69+
taskId: string,
6770
): Promise<SummarizeResponse> {
71+
telemetryService.captureContextCondensed(taskId)
6872
const response: SummarizeResponse = { messages, cost: 0, summary: "" }
6973
const messagesToSummarize = getMessagesSinceLastSummary(messages.slice(0, -N_MESSAGES_TO_KEEP))
7074
if (messagesToSummarize.length <= 1) {

src/core/sliding-window/__tests__/sliding-window.test.ts

Lines changed: 31 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ class MockApiHandler extends BaseProvider {
3737

3838
// Create a singleton instance for tests
3939
const mockApiHandler = new MockApiHandler()
40+
const taskId = "test-task-id"
4041

4142
/**
4243
* Tests for the truncateConversation function
@@ -49,7 +50,7 @@ describe("truncateConversation", () => {
4950
{ role: "user", content: "Third message" },
5051
]
5152

52-
const result = truncateConversation(messages, 0.5)
53+
const result = truncateConversation(messages, 0.5, taskId)
5354

5455
// With 2 messages after the first, 0.5 fraction means remove 1 message
5556
// But 1 is odd, so it rounds down to 0 (to make it even)
@@ -70,7 +71,7 @@ describe("truncateConversation", () => {
7071

7172
// 4 messages excluding first, 0.5 fraction = 2 messages to remove
7273
// 2 is already even, so no rounding needed
73-
const result = truncateConversation(messages, 0.5)
74+
const result = truncateConversation(messages, 0.5, taskId)
7475

7576
expect(result.length).toBe(3)
7677
expect(result[0]).toEqual(messages[0])
@@ -91,7 +92,7 @@ describe("truncateConversation", () => {
9192

9293
// 6 messages excluding first, 0.3 fraction = 1.8 messages to remove
9394
// 1.8 rounds down to 1, then to 0 to make it even
94-
const result = truncateConversation(messages, 0.3)
95+
const result = truncateConversation(messages, 0.3, taskId)
9596

9697
expect(result.length).toBe(7) // No messages removed
9798
expect(result).toEqual(messages)
@@ -104,7 +105,7 @@ describe("truncateConversation", () => {
104105
{ role: "user", content: "Third message" },
105106
]
106107

107-
const result = truncateConversation(messages, 0)
108+
const result = truncateConversation(messages, 0, taskId)
108109

109110
expect(result).toEqual(messages)
110111
})
@@ -119,7 +120,7 @@ describe("truncateConversation", () => {
119120

120121
// 3 messages excluding first, 1.0 fraction = 3 messages to remove
121122
// But 3 is odd, so it rounds down to 2 to make it even
122-
const result = truncateConversation(messages, 1)
123+
const result = truncateConversation(messages, 1, taskId)
123124

124125
expect(result.length).toBe(2)
125126
expect(result[0]).toEqual(messages[0])
@@ -251,6 +252,7 @@ describe("truncateConversationIfNeeded", () => {
251252
autoCondenseContext: false,
252253
autoCondenseContextPercent: 100,
253254
systemPrompt: "System prompt",
255+
taskId,
254256
})
255257

256258
// Check the new return type
@@ -282,6 +284,7 @@ describe("truncateConversationIfNeeded", () => {
282284
autoCondenseContext: false,
283285
autoCondenseContextPercent: 100,
284286
systemPrompt: "System prompt",
287+
taskId,
285288
})
286289

287290
expect(result).toEqual({
@@ -311,6 +314,7 @@ describe("truncateConversationIfNeeded", () => {
311314
autoCondenseContext: false,
312315
autoCondenseContextPercent: 100,
313316
systemPrompt: "System prompt",
317+
taskId,
314318
})
315319

316320
const result2 = await truncateConversationIfNeeded({
@@ -322,6 +326,7 @@ describe("truncateConversationIfNeeded", () => {
322326
autoCondenseContext: false,
323327
autoCondenseContextPercent: 100,
324328
systemPrompt: "System prompt",
329+
taskId,
325330
})
326331

327332
expect(result1.messages).toEqual(result2.messages)
@@ -340,6 +345,7 @@ describe("truncateConversationIfNeeded", () => {
340345
autoCondenseContext: false,
341346
autoCondenseContextPercent: 100,
342347
systemPrompt: "System prompt",
348+
taskId,
343349
})
344350

345351
const result4 = await truncateConversationIfNeeded({
@@ -351,6 +357,7 @@ describe("truncateConversationIfNeeded", () => {
351357
autoCondenseContext: false,
352358
autoCondenseContextPercent: 100,
353359
systemPrompt: "System prompt",
360+
taskId,
354361
})
355362

356363
expect(result3.messages).toEqual(result4.messages)
@@ -384,6 +391,7 @@ describe("truncateConversationIfNeeded", () => {
384391
autoCondenseContext: false,
385392
autoCondenseContextPercent: 100,
386393
systemPrompt: "System prompt",
394+
taskId,
387395
})
388396
expect(resultWithSmall).toEqual({
389397
messages: messagesWithSmallContent,
@@ -416,6 +424,7 @@ describe("truncateConversationIfNeeded", () => {
416424
autoCondenseContext: false,
417425
autoCondenseContextPercent: 100,
418426
systemPrompt: "System prompt",
427+
taskId,
419428
})
420429
expect(resultWithLarge.messages).not.toEqual(messagesWithLargeContent) // Should truncate
421430
expect(resultWithLarge.summary).toBe("")
@@ -441,6 +450,7 @@ describe("truncateConversationIfNeeded", () => {
441450
autoCondenseContext: false,
442451
autoCondenseContextPercent: 100,
443452
systemPrompt: "System prompt",
453+
taskId,
444454
})
445455
expect(resultWithVeryLarge.messages).not.toEqual(messagesWithVeryLargeContent) // Should truncate
446456
expect(resultWithVeryLarge.summary).toBe("")
@@ -469,6 +479,7 @@ describe("truncateConversationIfNeeded", () => {
469479
autoCondenseContext: false,
470480
autoCondenseContextPercent: 100,
471481
systemPrompt: "System prompt",
482+
taskId,
472483
})
473484
expect(result).toEqual({
474485
messages: expectedResult,
@@ -510,10 +521,11 @@ describe("truncateConversationIfNeeded", () => {
510521
autoCondenseContext: true,
511522
autoCondenseContextPercent: 100,
512523
systemPrompt: "System prompt",
524+
taskId,
513525
})
514526

515527
// Verify summarizeConversation was called with the right parameters
516-
expect(summarizeSpy).toHaveBeenCalledWith(messagesWithSmallContent, mockApiHandler, "System prompt")
528+
expect(summarizeSpy).toHaveBeenCalledWith(messagesWithSmallContent, mockApiHandler, "System prompt", taskId)
517529

518530
// Verify the result contains the summary information
519531
expect(result).toMatchObject({
@@ -557,6 +569,7 @@ describe("truncateConversationIfNeeded", () => {
557569
autoCondenseContext: true,
558570
autoCondenseContextPercent: 100,
559571
systemPrompt: "System prompt",
572+
taskId,
560573
})
561574

562575
// Verify summarizeConversation was called
@@ -594,6 +607,7 @@ describe("truncateConversationIfNeeded", () => {
594607
autoCondenseContext: false,
595608
autoCondenseContextPercent: 50, // This shouldn't matter since autoCondenseContext is false
596609
systemPrompt: "System prompt",
610+
taskId,
597611
})
598612

599613
// Verify summarizeConversation was not called
@@ -645,10 +659,11 @@ describe("truncateConversationIfNeeded", () => {
645659
autoCondenseContext: true,
646660
autoCondenseContextPercent: 50, // Set threshold to 50% - our tokens are at 60%
647661
systemPrompt: "System prompt",
662+
taskId,
648663
})
649664

650665
// Verify summarizeConversation was called with the right parameters
651-
expect(summarizeSpy).toHaveBeenCalledWith(messagesWithSmallContent, mockApiHandler, "System prompt")
666+
expect(summarizeSpy).toHaveBeenCalledWith(messagesWithSmallContent, mockApiHandler, "System prompt", taskId)
652667

653668
// Verify the result contains the summary information
654669
expect(result).toMatchObject({
@@ -682,6 +697,7 @@ describe("truncateConversationIfNeeded", () => {
682697
autoCondenseContext: true,
683698
autoCondenseContextPercent: 50, // Set threshold to 50% - our tokens are at 40%
684699
systemPrompt: "System prompt",
700+
taskId,
685701
})
686702

687703
// Verify summarizeConversation was not called
@@ -738,6 +754,7 @@ describe("getMaxTokens", () => {
738754
autoCondenseContext: false,
739755
autoCondenseContextPercent: 100,
740756
systemPrompt: "System prompt",
757+
taskId,
741758
})
742759
expect(result1).toEqual({
743760
messages: messagesWithSmallContent,
@@ -756,6 +773,7 @@ describe("getMaxTokens", () => {
756773
autoCondenseContext: false,
757774
autoCondenseContextPercent: 100,
758775
systemPrompt: "System prompt",
776+
taskId,
759777
})
760778
expect(result2.messages).not.toEqual(messagesWithSmallContent)
761779
expect(result2.messages.length).toBe(3) // Truncated with 0.5 fraction
@@ -782,6 +800,7 @@ describe("getMaxTokens", () => {
782800
autoCondenseContext: false,
783801
autoCondenseContextPercent: 100,
784802
systemPrompt: "System prompt",
803+
taskId,
785804
})
786805
expect(result1).toEqual({
787806
messages: messagesWithSmallContent,
@@ -800,6 +819,7 @@ describe("getMaxTokens", () => {
800819
autoCondenseContext: false,
801820
autoCondenseContextPercent: 100,
802821
systemPrompt: "System prompt",
822+
taskId,
803823
})
804824
expect(result2.messages).not.toEqual(messagesWithSmallContent)
805825
expect(result2.messages.length).toBe(3) // Truncated with 0.5 fraction
@@ -825,6 +845,7 @@ describe("getMaxTokens", () => {
825845
autoCondenseContext: false,
826846
autoCondenseContextPercent: 100,
827847
systemPrompt: "System prompt",
848+
taskId,
828849
})
829850
expect(result1.messages).toEqual(messagesWithSmallContent)
830851

@@ -838,6 +859,7 @@ describe("getMaxTokens", () => {
838859
autoCondenseContext: false,
839860
autoCondenseContextPercent: 100,
840861
systemPrompt: "System prompt",
862+
taskId,
841863
})
842864
expect(result2).not.toEqual(messagesWithSmallContent)
843865
expect(result2.messages.length).toBe(3) // Truncated with 0.5 fraction
@@ -861,6 +883,7 @@ describe("getMaxTokens", () => {
861883
autoCondenseContext: false,
862884
autoCondenseContextPercent: 100,
863885
systemPrompt: "System prompt",
886+
taskId,
864887
})
865888
expect(result1.messages).toEqual(messagesWithSmallContent)
866889

@@ -874,6 +897,7 @@ describe("getMaxTokens", () => {
874897
autoCondenseContext: false,
875898
autoCondenseContextPercent: 100,
876899
systemPrompt: "System prompt",
900+
taskId,
877901
})
878902
expect(result2).not.toEqual(messagesWithSmallContent)
879903
expect(result2.messages.length).toBe(3) // Truncated with 0.5 fraction

src/core/sliding-window/index.ts

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import { Anthropic } from "@anthropic-ai/sdk"
22
import { ApiHandler } from "../../api"
33
import { summarizeConversation, SummarizeResponse } from "../condense"
44
import { ApiMessage } from "../task-persistence/apiMessages"
5+
import { telemetryService } from "../../services/telemetry/TelemetryService"
56

67
/**
78
* Default percentage of the context window to use as a buffer when deciding when to truncate
@@ -31,9 +32,11 @@ export async function estimateTokenCount(
3132
*
3233
* @param {ApiMessage[]} messages - The conversation messages.
3334
* @param {number} fracToRemove - The fraction (between 0 and 1) of messages (excluding the first) to remove.
35+
* @param {string} taskId - The task ID for the conversation, used for telemetry
3436
* @returns {ApiMessage[]} The truncated conversation messages.
3537
*/
36-
export function truncateConversation(messages: ApiMessage[], fracToRemove: number): ApiMessage[] {
38+
export function truncateConversation(messages: ApiMessage[], fracToRemove: number, taskId: string): ApiMessage[] {
39+
telemetryService.captureSlidingWindowTruncation(taskId)
3740
const truncatedMessages = [messages[0]]
3841
const rawMessagesToRemove = Math.floor((messages.length - 1) * fracToRemove)
3942
const messagesToRemove = rawMessagesToRemove - (rawMessagesToRemove % 2)
@@ -66,6 +69,7 @@ type TruncateOptions = {
6669
autoCondenseContext: boolean
6770
autoCondenseContextPercent: number
6871
systemPrompt: string
72+
taskId: string
6973
}
7074

7175
type TruncateResponse = SummarizeResponse & { prevContextTokens: number }
@@ -86,6 +90,7 @@ export async function truncateConversationIfNeeded({
8690
autoCondenseContext,
8791
autoCondenseContextPercent,
8892
systemPrompt,
93+
taskId,
8994
}: TruncateOptions): Promise<TruncateResponse> {
9095
// Calculate the maximum tokens reserved for response
9196
const reservedTokens = maxTokens || contextWindow * 0.2
@@ -108,7 +113,7 @@ export async function truncateConversationIfNeeded({
108113
const contextPercent = (100 * prevContextTokens) / contextWindow
109114
if (contextPercent >= autoCondenseContextPercent || prevContextTokens > allowedTokens) {
110115
// Attempt to intelligently condense the context
111-
const result = await summarizeConversation(messages, apiHandler, systemPrompt)
116+
const result = await summarizeConversation(messages, apiHandler, systemPrompt, taskId)
112117
if (result.summary) {
113118
return { ...result, prevContextTokens }
114119
}
@@ -117,7 +122,7 @@ export async function truncateConversationIfNeeded({
117122

118123
// Fall back to sliding window truncation if needed
119124
if (prevContextTokens > allowedTokens) {
120-
const truncatedMessages = truncateConversation(messages, 0.5)
125+
const truncatedMessages = truncateConversation(messages, 0.5, taskId)
121126
return { messages: truncatedMessages, prevContextTokens, summary: "", cost: 0 }
122127
}
123128
// No truncation or condensation needed

0 commit comments

Comments
 (0)