[Condense Context] Track metrics around context condensing and show in UI

This commit is contained in:
Canyon Robins 2025-05-19 17:36:17 -07:00
parent ce3e4e883d
commit 603d2eface
11 changed files with 233 additions and 77 deletions

View file

@ -993,6 +993,7 @@ export const clineSays = [
"checkpoint_saved",
"rooignore_error",
"diff_error",
"condense_context",
] as const
export const clineSaySchema = z.enum(clineSays)
@ -1011,6 +1012,18 @@ export const toolProgressStatusSchema = z.object({
export type ToolProgressStatus = z.infer<typeof toolProgressStatusSchema>
/**
* ContextCondense
*/
export const contextCondenseSchema = z.object({
cost: z.number(),
prevContextTokens: z.number(),
newContextTokens: z.number(),
})
export type ContextCondense = z.infer<typeof contextCondenseSchema>
/**
* ClineMessage
*/
@ -1027,6 +1040,7 @@ export const clineMessageSchema = z.object({
conversationHistoryIndex: z.number().optional(),
checkpoint: z.record(z.string(), z.unknown()).optional(),
progressStatus: toolProgressStatusSchema.optional(),
contextCondense: contextCondenseSchema.optional(),
})
export type ClineMessage = z.infer<typeof clineMessageSchema>

View file

@ -45,22 +45,33 @@ Example summary structure:
Output only the summary of the conversation so far, without any additional commentary or explanation.
`
export type SummarizeResponse = {
messages: ApiMessage[] // The messages after summarization
summary: string // The summary text; empty string for no summary
cost: number // The cost of the summarization operation
newContextTokens?: number // The number of tokens in the context for the next API request
}
/**
* Summarizes the conversation messages using an LLM call
*
* @param {ApiMessage[]} messages - The conversation messages
* @param {ApiHandler} apiHandler - The API handler to use for token counting.
* @returns {ApiMessage[]} - The input messages, potentially including a new summary message before the last message.
* @returns {SummarizeResponse} - The result of the summarization operation (see above)
*/
export async function summarizeConversation(messages: ApiMessage[], apiHandler: ApiHandler): Promise<ApiMessage[]> {
export async function summarizeConversation(
messages: ApiMessage[],
apiHandler: ApiHandler,
): Promise<SummarizeResponse> {
const response: SummarizeResponse = { messages, cost: 0, summary: "" }
const messagesToSummarize = getMessagesSinceLastSummary(messages.slice(0, -N_MESSAGES_TO_KEEP))
if (messagesToSummarize.length <= 1) {
return messages // Not enough messages to warrant a summary
return response // Not enough messages to warrant a summary
}
const keepMessages = messages.slice(-N_MESSAGES_TO_KEEP)
for (const message of keepMessages) {
if (message.isSummary) {
return messages // We recently summarized these messages; it's too soon to summarize again.
return response // We recently summarized these messages; it's too soon to summarize again.
}
}
const finalRequestMessage: Anthropic.MessageParam = {
@ -73,16 +84,21 @@ export async function summarizeConversation(messages: ApiMessage[], apiHandler:
// Note: this doesn't need to be a stream, consider using something like apiHandler.completePrompt
const stream = apiHandler.createMessage(SUMMARY_PROMPT, requestMessages)
let summary = ""
// TODO(canyon): compute usage and cost for this operation and update the global metrics.
let cost = 0
let outputTokens = 0
for await (const chunk of stream) {
if (chunk.type === "text") {
summary += chunk.text
} else if (chunk.type === "usage") {
// Record final usage chunk only
cost = chunk.totalCost ?? 0
outputTokens = chunk.outputTokens ?? 0
}
}
summary = summary.trim()
if (summary.length === 0) {
console.warn("Received empty summary from API")
return messages
return { ...response, cost }
}
const summaryMessage: ApiMessage = {
role: "assistant",
@ -90,8 +106,16 @@ export async function summarizeConversation(messages: ApiMessage[], apiHandler:
ts: keepMessages[0].ts,
isSummary: true,
}
const newMessages = [...messages.slice(0, -N_MESSAGES_TO_KEEP), summaryMessage, ...keepMessages]
return [...messages.slice(0, -N_MESSAGES_TO_KEEP), summaryMessage, ...keepMessages]
// Count the tokens in the context for the next API request
// We only estimate the tokens in summaryMesage if outputTokens is 0, otherwise we use outputTokens
const contextMessages = outputTokens ? [...keepMessages] : [summaryMessage, ...keepMessages]
const contextBlocks = contextMessages.flatMap((message) =>
typeof message.content === "string" ? [{ text: message.content, type: "text" as const }] : message.content,
)
const newContextTokens = outputTokens + (await apiHandler.countTokens(contextBlocks))
return { ...response, messages: newMessages, summary, cost, newContextTokens }
}
/* Returns the list of all messages since the last summary message, including the summary. Returns all messages if there is no summary. */

View file

@ -1,6 +1,6 @@
import { Anthropic } from "@anthropic-ai/sdk"
import { ApiHandler } from "../../api"
import { summarizeConversation } from "../condense"
import { summarizeConversation, SummarizeResponse } from "../condense"
import { ApiMessage } from "../task-persistence/apiMessages"
/**
@ -65,6 +65,8 @@ type TruncateOptions = {
autoCondenseContext?: boolean
}
type TruncateResponse = SummarizeResponse & { prevContextTokens: number }
/**
* Conditionally truncates the conversation messages if the total token count
* exceeds the model's limit, considering the size of incoming content.
@ -79,7 +81,7 @@ export async function truncateConversationIfNeeded({
maxTokens,
apiHandler,
autoCondenseContext,
}: TruncateOptions): Promise<ApiMessage[]> {
}: TruncateOptions): Promise<TruncateResponse> {
// Calculate the maximum tokens reserved for response
const reservedTokens = maxTokens || contextWindow * 0.2
@ -99,12 +101,13 @@ export async function truncateConversationIfNeeded({
// Determine if truncation is needed and apply if necessary
if (effectiveTokens <= allowedTokens) {
return messages
return { messages, summary: "", cost: 0, prevContextTokens: effectiveTokens }
} else if (autoCondenseContext) {
const summarizedMessages = await summarizeConversation(messages, apiHandler)
if (messages !== summarizedMessages) {
return summarizedMessages
const result = await summarizeConversation(messages, apiHandler)
if (messages !== result.messages) {
return { ...result, prevContextTokens: effectiveTokens }
}
}
return truncateConversation(messages, 0.5)
const truncatedmessages = truncateConversation(messages, 0.5)
return { messages: truncatedmessages, prevContextTokens: effectiveTokens, summary: "", cost: 0 }
}

View file

@ -9,7 +9,7 @@ import pWaitFor from "p-wait-for"
import { serializeError } from "serialize-error"
// schemas
import { TokenUsage, ToolUsage, ToolName } from "../../schemas"
import { TokenUsage, ToolUsage, ToolName, ContextCondense } from "../../schemas"
// api
import { ApiHandler, buildApiHandler } from "../../api"
@ -490,6 +490,7 @@ export class Task extends EventEmitter<ClineEvents> {
options: {
isNonInteractive?: boolean
} = {},
contextCondense?: ContextCondense,
): Promise<undefined> {
if (this.abort) {
throw new Error(`[RooCode#say] task ${this.taskId}.${this.instanceId} aborted`)
@ -562,7 +563,15 @@ export class Task extends EventEmitter<ClineEvents> {
this.lastMessageTs = sayTs
}
await this.addToClineMessages({ ts: sayTs, type: "say", say: type, text, images, checkpoint })
await this.addToClineMessages({
ts: sayTs,
type: "say",
say: type,
text,
images,
checkpoint,
contextCondense,
})
}
}
@ -985,10 +994,6 @@ export class Task extends EventEmitter<ClineEvents> {
this.consecutiveMistakeCount = 0
}
// Get previous api req's index to check token usage and determine if we
// need to truncate conversation history.
const previousApiReqIndex = findLastIndex(this.clineMessages, (m) => m.say === "api_req_started")
// In this Cline request loop, we need to check if this task instance
// has been asked to wait for a subtask to finish before continuing.
const provider = this.providerRef.deref()
@ -1147,7 +1152,7 @@ export class Task extends EventEmitter<ClineEvents> {
// Yields only if the first chunk is successful, otherwise will
// allow the user to retry the request (most likely due to rate
// limit error, which gets thrown on the first chunk).
const stream = this.attemptApiRequest(previousApiReqIndex)
const stream = this.attemptApiRequest()
let assistantMessage = ""
let reasoningMessage = ""
this.isStreaming = true
@ -1354,7 +1359,7 @@ export class Task extends EventEmitter<ClineEvents> {
}
}
public async *attemptApiRequest(previousApiReqIndex: number, retryAttempt: number = 0): ApiStream {
public async *attemptApiRequest(retryAttempt: number = 0): ApiStream {
let mcpHub: McpHub | undefined
const { apiConfiguration, mcpEnabled, autoApprovalEnabled, alwaysApproveResubmit, requestDelaySeconds } =
@ -1444,25 +1449,8 @@ export class Task extends EventEmitter<ClineEvents> {
)
})()
// If the previous API request's total token usage is close to the
// context window, truncate the conversation history to free up space
// for the new request.
if (previousApiReqIndex >= 0) {
const previousRequest = this.clineMessages[previousApiReqIndex]?.text
if (!previousRequest) {
return
}
const {
tokensIn = 0,
tokensOut = 0,
cacheWrites = 0,
cacheReads = 0,
}: ClineApiReqInfo = JSON.parse(previousRequest)
const totalTokens = tokensIn + tokensOut + cacheWrites + cacheReads
const { contextTokens } = this.getTokenUsage()
if (contextTokens) {
// Default max tokens value for thinking models when no specific
// value is set.
const DEFAULT_THINKING_MODEL_MAX_TOKENS = 16_384
@ -1476,16 +1464,30 @@ export class Task extends EventEmitter<ClineEvents> {
const contextWindow = modelInfo.contextWindow
const autoCondenseContext = experiments?.autoCondenseContext ?? false
const trimmedMessages = await truncateConversationIfNeeded({
const truncateResult = await truncateConversationIfNeeded({
messages: this.apiConversationHistory,
totalTokens,
totalTokens: contextTokens,
maxTokens,
contextWindow,
apiHandler: this.api,
autoCondenseContext,
})
if (trimmedMessages !== this.apiConversationHistory) {
await this.overwriteApiConversationHistory(trimmedMessages)
if (truncateResult.messages !== this.apiConversationHistory) {
await this.overwriteApiConversationHistory(truncateResult.messages)
}
if (truncateResult.summary) {
const { summary, cost, prevContextTokens, newContextTokens = 0 } = truncateResult
const contextCondense: ContextCondense = { summary, cost, newContextTokens, prevContextTokens }
await this.say(
"condense_context",
undefined /* text */,
undefined /* images */,
false /* partial */,
undefined /* checkpoint */,
undefined /* progressStatus */,
undefined /* options */,
contextCondense,
)
}
}
@ -1556,7 +1558,7 @@ export class Task extends EventEmitter<ClineEvents> {
// Delegate generator output from the recursive call with
// incremented retry count.
yield* this.attemptApiRequest(previousApiReqIndex, retryAttempt + 1)
yield* this.attemptApiRequest(retryAttempt + 1)
return
} else {
@ -1574,7 +1576,7 @@ export class Task extends EventEmitter<ClineEvents> {
await this.say("api_req_retried")
// Delegate generator output from the recursive call.
yield* this.attemptApiRequest(previousApiReqIndex)
yield* this.attemptApiRequest()
return
}
}
@ -1610,7 +1612,7 @@ export class Task extends EventEmitter<ClineEvents> {
return combineApiRequests(combineCommandSequences(messages))
}
public getTokenUsage() {
public getTokenUsage(): TokenUsage {
return getApiMetrics(this.combineMessages(this.clineMessages.slice(1)))
}

View file

@ -405,6 +405,7 @@ type ClineMessage = {
| "checkpoint_saved"
| "rooignore_error"
| "diff_error"
| "condense_context"
)
| undefined
text?: string | undefined
@ -423,6 +424,14 @@ type ClineMessage = {
text?: string | undefined
}
| undefined
contextCondense?:
| {
cost: number
prevContextTokens: number
newContextTokens: number
summary: string
}
| undefined
}
type TokenUsage = {
@ -480,6 +489,7 @@ type RooCodeEvents = {
| "checkpoint_saved"
| "rooignore_error"
| "diff_error"
| "condense_context"
)
| undefined
text?: string | undefined
@ -498,6 +508,14 @@ type RooCodeEvents = {
text?: string | undefined
}
| undefined
contextCondense?:
| {
cost: number
prevContextTokens: number
newContextTokens: number
summary: string
}
| undefined
}
},
]
@ -949,6 +967,7 @@ type IpcMessage =
| "checkpoint_saved"
| "rooignore_error"
| "diff_error"
| "condense_context"
)
| undefined
text?: string | undefined
@ -967,6 +986,14 @@ type IpcMessage =
text?: string | undefined
}
| undefined
contextCondense?:
| {
cost: number
prevContextTokens: number
newContextTokens: number
summary: string
}
| undefined
}
},
]
@ -1408,6 +1435,7 @@ type TaskEvent =
| "checkpoint_saved"
| "rooignore_error"
| "diff_error"
| "condense_context"
)
| undefined
text?: string | undefined
@ -1426,6 +1454,14 @@ type TaskEvent =
text?: string | undefined
}
| undefined
contextCondense?:
| {
cost: number
prevContextTokens: number
newContextTokens: number
summary: string
}
| undefined
}
},
]

View file

@ -413,6 +413,7 @@ type ClineMessage = {
| "checkpoint_saved"
| "rooignore_error"
| "diff_error"
| "condense_context"
)
| undefined
text?: string | undefined
@ -431,6 +432,14 @@ type ClineMessage = {
text?: string | undefined
}
| undefined
contextCondense?:
| {
cost: number
prevContextTokens: number
newContextTokens: number
summary: string
}
| undefined
}
export type { ClineMessage }
@ -492,6 +501,7 @@ type RooCodeEvents = {
| "checkpoint_saved"
| "rooignore_error"
| "diff_error"
| "condense_context"
)
| undefined
text?: string | undefined
@ -510,6 +520,14 @@ type RooCodeEvents = {
text?: string | undefined
}
| undefined
contextCondense?:
| {
cost: number
prevContextTokens: number
newContextTokens: number
summary: string
}
| undefined
}
},
]
@ -963,6 +981,7 @@ type IpcMessage =
| "checkpoint_saved"
| "rooignore_error"
| "diff_error"
| "condense_context"
)
| undefined
text?: string | undefined
@ -981,6 +1000,14 @@ type IpcMessage =
text?: string | undefined
}
| undefined
contextCondense?:
| {
cost: number
prevContextTokens: number
newContextTokens: number
summary: string
}
| undefined
}
},
]
@ -1426,6 +1453,7 @@ type TaskEvent =
| "checkpoint_saved"
| "rooignore_error"
| "diff_error"
| "condense_context"
)
| undefined
text?: string | undefined
@ -1444,6 +1472,14 @@ type TaskEvent =
text?: string | undefined
}
| undefined
contextCondense?:
| {
cost: number
prevContextTokens: number
newContextTokens: number
summary: string
}
| undefined
}
},
]

View file

@ -929,6 +929,7 @@ export const clineSays = [
"checkpoint_saved",
"rooignore_error",
"diff_error",
"condense_context",
] as const
export const clineSaySchema = z.enum(clineSays)
@ -946,6 +947,19 @@ export const toolProgressStatusSchema = z.object({
export type ToolProgressStatus = z.infer<typeof toolProgressStatusSchema>
/**
* ContextCondense
*/
export const contextCondenseSchema = z.object({
cost: z.number(),
prevContextTokens: z.number(),
newContextTokens: z.number(),
summary: z.string(),
})
export type ContextCondense = z.infer<typeof contextCondenseSchema>
/**
* ClineMessage
*/
@ -962,6 +976,7 @@ export const clineMessageSchema = z.object({
conversationHistoryIndex: z.number().optional(),
checkpoint: z.record(z.string(), z.unknown()).optional(),
progressStatus: toolProgressStatusSchema.optional(),
contextCondense: contextCondenseSchema.optional(),
})
export type ClineMessage = z.infer<typeof clineMessageSchema>

View file

@ -2,11 +2,19 @@ import { TokenUsage } from "../schemas"
import { ClineMessage } from "./ExtensionMessage"
export type ParsedApiReqStartedTextType = {
tokensIn: number
tokensOut: number
cacheWrites: number
cacheReads: number
cost?: number // Only present if combineApiRequests has been called
}
/**
* Calculates API metrics from an array of ClineMessages.
*
* This function processes 'api_req_started' messages that have been combined with their
* corresponding 'api_req_finished' messages by the combineApiRequests function.
* This function processes 'condense_context' messages and 'api_req_started' messages that have been
* combined with their corresponding 'api_req_finished' messages by the combineApiRequests function.
* It extracts and sums up the tokensIn, tokensOut, cacheWrites, cacheReads, and cost from these messages.
*
* @param messages - An array of ClineMessage objects to process.
@ -29,30 +37,15 @@ export function getApiMetrics(messages: ClineMessage[]) {
contextTokens: 0,
}
// Helper function to get total tokens from a message
const getTotalTokensFromMessage = (message: ClineMessage): number => {
if (!message.text) return 0
try {
const { tokensIn, tokensOut, cacheWrites, cacheReads } = JSON.parse(message.text)
return (tokensIn || 0) + (tokensOut || 0) + (cacheWrites || 0) + (cacheReads || 0)
} catch {
return 0
}
}
// Find the last api_req_started message that has any tokens
const lastApiReq = [...messages].reverse().find((message) => {
if (message.type === "say" && message.say === "api_req_started") {
return getTotalTokensFromMessage(message) > 0
}
return false
})
// Calculate running totals
messages.forEach((message) => {
if (message.type === "say" && message.say === "api_req_started" && message.text) {
if (!message.text || message.type !== "say") {
return
}
if (message.say === "api_req_started") {
try {
const { tokensIn, tokensOut, cacheWrites, cacheReads, cost } = JSON.parse(message.text)
const parsedText: ParsedApiReqStartedTextType = JSON.parse(message.text)
const { tokensIn, tokensOut, cacheWrites, cacheReads, cost } = parsedText
if (typeof tokensIn === "number") {
result.totalTokensIn += tokensIn
@ -69,16 +62,31 @@ export function getApiMetrics(messages: ClineMessage[]) {
if (typeof cost === "number") {
result.totalCost += cost
}
// If this is the last api request with tokens, use its total for context size
if (message === lastApiReq) {
result.contextTokens = getTotalTokensFromMessage(message)
}
} catch (error) {
console.error("Error parsing JSON:", error)
}
} else if (message.say === "condense_context") {
result.totalCost += message.contextCondense?.cost ?? 0
}
})
// Calculate context tokens, from the last API request started or condense context message
result.contextTokens = 0
for (let i = messages.length - 1; i >= 0; i--) {
const message = messages[i]
if (!message.text || message.type !== "say") {
continue
} else if (message.say === "api_req_started") {
const parsedText: ParsedApiReqStartedTextType = JSON.parse(message.text)
const { tokensIn, tokensOut, cacheWrites, cacheReads } = parsedText
result.contextTokens = (tokensIn || 0) + (tokensOut || 0) + (cacheWrites || 0) + (cacheReads || 0)
} else if (message.say === "condense_context") {
result.contextTokens = message.contextCondense?.newContextTokens ?? 0
}
if (result.contextTokens) {
break
}
}
return result
}

View file

@ -31,6 +31,7 @@ import { ProgressIndicator } from "./ProgressIndicator"
import { Markdown } from "./Markdown"
import { CommandExecution } from "./CommandExecution"
import { CommandExecutionError } from "./CommandExecutionError"
import ContextCondenseRow from "./ContextCondenseRow"
interface ChatRowProps {
message: ClineMessage
@ -926,6 +927,8 @@ export const ChatRowContent = ({
checkpoint={message.checkpoint}
/>
)
case "condense_context":
return <ContextCondenseRow ts={message.ts} contextCondense={message.contextCondense} />
default:
return (
<>

View file

@ -0,0 +1,15 @@
import { ContextCondense } from "@roo/schemas"
interface ContextCondenseRowProps {
ts: number
contextCondense?: ContextCondense
}
const ContextCondenseRow = ({ contextCondense }: ContextCondenseRowProps) => {
if (!contextCondense) {
return null
}
return null
}
export default ContextCondenseRow

View file

@ -400,7 +400,7 @@
"warning": "⚠️",
"AUTO_CONDENSE_CONTEXT": {
"name": "Intelligently condense the context window",
"description": "Uses an LLM call to summarize the past conversation when the task's context window is almost full, rather than dropping old messages. Disclaimer: the cost of summarizing is not currently included in the API costs shown in the UI."
"description": "Uses an LLM call to summarize the past conversation when the task's context window is almost full, rather than dropping old messages."
},
"DIFF_STRATEGY_UNIFIED": {
"name": "Use experimental unified diff strategy",