Refactor: Unified context-management architecture with improved UX

packages/types/src/__tests__/context-management.test.ts

@@ -0,0 +1,28 @@
+import { describe, it, expect } from "vitest"
+import { CONTEXT_MANAGEMENT_EVENTS, isContextManagementEvent } from "../context-management.js"
+
+describe("context-management", () => {
+	describe("CONTEXT_MANAGEMENT_EVENTS", () => {
+		it("should contain all expected event types", () => {
+			expect(CONTEXT_MANAGEMENT_EVENTS).toContain("condense_context")
+			expect(CONTEXT_MANAGEMENT_EVENTS).toContain("condense_context_error")
+			expect(CONTEXT_MANAGEMENT_EVENTS).toContain("sliding_window_truncation")
+			expect(CONTEXT_MANAGEMENT_EVENTS).toHaveLength(3)
+		})
+	})
+
+	describe("isContextManagementEvent", () => {
+		it("should return true for valid context management events", () => {
+			expect(isContextManagementEvent("condense_context")).toBe(true)
+			expect(isContextManagementEvent("condense_context_error")).toBe(true)
+			expect(isContextManagementEvent("sliding_window_truncation")).toBe(true)
+		})
+
+		it("should return false for non-context-management events", () => {
+			expect(isContextManagementEvent("text")).toBe(false)
+			expect(isContextManagementEvent("error")).toBe(false)
+			expect(isContextManagementEvent(null)).toBe(false)
+			expect(isContextManagementEvent(undefined)).toBe(false)
+		})
+	})
+})

packages/types/src/context-management.ts

@@ -0,0 +1,34 @@
+/**
+ * Context Management Types
+ *
+ * This module provides type definitions for context management events.
+ * These events are used to handle different strategies for managing conversation context
+ * when approaching token limits.
+ *
+ * Event Types:
+ * - `condense_context`: Context was condensed using AI summarization
+ * - `condense_context_error`: An error occurred during context condensation
+ * - `sliding_window_truncation`: Context was truncated using sliding window strategy
+ */
+
+/**
+ * Array of all context management event types.
+ * Used for runtime type checking.
+ */
+export const CONTEXT_MANAGEMENT_EVENTS = [
+	"condense_context",
+	"condense_context_error",
+	"sliding_window_truncation",
+] as const
+
+/**
+ * Union type representing all possible context management event types.
+ */
+export type ContextManagementEvent = (typeof CONTEXT_MANAGEMENT_EVENTS)[number]
+
+/**
+ * Type guard function to check if a value is a valid context management event.
+ */
+export function isContextManagementEvent(value: unknown): value is ContextManagementEvent {
+	return typeof value === "string" && (CONTEXT_MANAGEMENT_EVENTS as readonly string[]).includes(value)
+}

packages/types/src/index.ts

@@ -1,6 +1,7 @@
 export * from "./api.js"
 export * from "./cloud.js"
 export * from "./codebase-index.js"
+export * from "./context-management.js"
 export * from "./cookie-consent.js"
 export * from "./events.js"
 export * from "./experiment.js"

packages/types/src/message.ts

@@ -197,8 +197,17 @@ export type ToolProgressStatus = z.infer<typeof toolProgressStatusSchema>
 
 /**
  * ContextCondense
+ *
+ * Data associated with a successful context condensation event.
+ * This is attached to messages with `say: "condense_context"` when
+ * the condensation operation completes successfully.
+ *
+ * @property cost - The API cost incurred for the condensation operation
+ * @property prevContextTokens - Token count before condensation
+ * @property newContextTokens - Token count after condensation
+ * @property summary - The condensed summary that replaced the original context
+ * @property condenseId - Optional unique identifier for this condensation operation
  */
-
 export const contextCondenseSchema = z.object({
 	cost: z.number(),
 	prevContextTokens: z.number(),
@@ -212,21 +221,39 @@ export type ContextCondense = z.infer<typeof contextCondenseSchema>
 /**
  * ContextTruncation
  *
- * Used to track sliding window truncation events for the UI.
+ * Data associated with a sliding window truncation event.
+ * This is attached to messages with `say: "sliding_window_truncation"` when
+ * messages are removed from the conversation history to stay within token limits.
+ *
+ * Unlike condensation, truncation simply removes older messages without
+ * summarizing them. This is a faster but less context-preserving approach.
+ *
+ * @property truncationId - Unique identifier for this truncation operation
+ * @property messagesRemoved - Number of conversation messages that were removed
+ * @property prevContextTokens - Token count before truncation occurred
+ * @property newContextTokens - Token count after truncation occurred
  */
-
 export const contextTruncationSchema = z.object({
 	truncationId: z.string(),
 	messagesRemoved: z.number(),
 	prevContextTokens: z.number(),
+	newContextTokens: z.number(),
 })
 
 export type ContextTruncation = z.infer<typeof contextTruncationSchema>
 
 /**
  * ClineMessage
+ *
+ * The main message type used for communication between the extension and webview.
+ * Messages can either be "ask" (requiring user response) or "say" (informational).
+ *
+ * Context Management Fields:
+ * - `contextCondense`: Present when `say: "condense_context"` and condensation succeeded
+ * - `contextTruncation`: Present when `say: "sliding_window_truncation"` and truncation occurred
+ *
+ * Note: These fields are mutually exclusive - a message will have at most one of them.
  */
-
 export const clineMessageSchema = z.object({
 	ts: z.number(),
 	type: z.union([z.literal("ask"), z.literal("say")]),
@@ -239,7 +266,15 @@ export const clineMessageSchema = z.object({
 	conversationHistoryIndex: z.number().optional(),
 	checkpoint: z.record(z.string(), z.unknown()).optional(),
 	progressStatus: toolProgressStatusSchema.optional(),
+	/**
+	 * Data for successful context condensation.
+	 * Present when `say: "condense_context"` and `partial: false`.
+	 */
 	contextCondense: contextCondenseSchema.optional(),
+	/**
+	 * Data for sliding window truncation.
+	 * Present when `say: "sliding_window_truncation"`.
+	 */
 	contextTruncation: contextTruncationSchema.optional(),
 	isProtected: z.boolean().optional(),
 	apiProtocol: z.union([z.literal("openai"), z.literal("anthropic")]).optional(),

src/core/context-management/index.ts

@@ -164,6 +164,7 @@ export type ContextManagementResult = SummarizeResponse & {
 	prevContextTokens: number
 	truncationId?: string
 	messagesRemoved?: number
+	newContextTokensAfterTruncation?: number
 }
 
 /**
@@ -254,6 +255,25 @@ export async function manageContext({
 	// Fall back to sliding window truncation if needed
 	if (prevContextTokens > allowedTokens) {
 		const truncationResult = truncateConversation(messages, 0.5, taskId)
+
+		// Calculate new context tokens after truncation by counting non-truncated messages
+		// Messages with truncationParent are hidden, so we count only those without it
+		const effectiveMessages = truncationResult.messages.filter(
+			(msg) => !msg.truncationParent && !msg.isTruncationMarker,
+		)
+		let newContextTokensAfterTruncation = 0
+		for (const msg of effectiveMessages) {
+			const content = msg.content
+			if (Array.isArray(content)) {
+				newContextTokensAfterTruncation += await estimateTokenCount(content, apiHandler)
+			} else if (typeof content === "string") {
+				newContextTokensAfterTruncation += await estimateTokenCount(
+					[{ type: "text", text: content }],
+					apiHandler,
+				)
+			}
+		}
+
 		return {
 			messages: truncationResult.messages,
 			prevContextTokens,
@@ -262,6 +282,7 @@ export async function manageContext({
 			error,
 			truncationId: truncationResult.truncationId,
 			messagesRemoved: truncationResult.messagesRemoved,
+			newContextTokensAfterTruncation,
 		}
 	}
 	// No truncation or condensation needed

src/core/task/Task.ts

@@ -3366,6 +3366,9 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 		const protocol = resolveToolProtocol(this.apiConfiguration, modelInfo)
 		const useNativeTools = isNativeProtocol(protocol)
 
+		// Send condenseTaskContextStarted to show in-progress indicator
+		await this.providerRef.deref()?.postMessageToWebview({ type: "condenseTaskContextStarted", text: this.taskId })
+
 		// Force aggressive truncation by keeping only 75% of the conversation history
 		const truncateResult = await manageContext({
 			messages: this.apiConversationHistory,
@@ -3405,6 +3408,7 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 				truncationId: truncateResult.truncationId,
 				messagesRemoved: truncateResult.messagesRemoved ?? 0,
 				prevContextTokens: truncateResult.prevContextTokens,
+				newContextTokens: truncateResult.newContextTokensAfterTruncation ?? 0,
 			}
 			await this.say(
 				"sliding_window_truncation",
@@ -3418,6 +3422,9 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 				contextTruncation,
 			)
 		}
+
+		// Notify webview that context management is complete (removes in-progress spinner)
+		await this.providerRef.deref()?.postMessageToWebview({ type: "condenseTaskContextResponse", text: this.taskId })
 	}
 
 	public async *attemptApiRequest(retryAttempt: number = 0): ApiStream {
@@ -3505,6 +3512,45 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 			const protocol = resolveToolProtocol(this.apiConfiguration, modelInfoForProtocol)
 			const useNativeTools = isNativeProtocol(protocol)
 
+			// Check if context management will likely run (threshold check)
+			// This allows us to show an in-progress indicator to the user
+			// Important: Match the exact calculation from manageContext to avoid threshold mismatch
+			// manageContext uses: prevContextTokens = totalTokens + lastMessageTokens
+			const lastMessage = this.apiConversationHistory[this.apiConversationHistory.length - 1]
+			const lastMessageContent = lastMessage?.content
+			let lastMessageTokens = 0
+			if (lastMessageContent) {
+				lastMessageTokens = Array.isArray(lastMessageContent)
+					? await this.api.countTokens(lastMessageContent)
+					: await this.api.countTokens([{ type: "text", text: lastMessageContent as string }])
+			}
+			const prevContextTokens = contextTokens + lastMessageTokens
+			const estimatedUsagePercent = (100 * prevContextTokens) / contextWindow
+
+			// Match manageContext's threshold logic
+			const profileThresholdSettings = profileThresholds[currentProfileId] as
+				| { autoCondenseContextPercent?: number }
+				| undefined
+			const effectiveThreshold =
+				profileThresholdSettings?.autoCondenseContextPercent ?? autoCondenseContextPercent
+
+			// Calculate allowedTokens the same way manageContext does
+			const TOKEN_BUFFER_PERCENTAGE = 0.1
+			const reservedTokens = maxTokens ?? 8192 // ANTHROPIC_DEFAULT_MAX_TOKENS
+			const allowedTokens = contextWindow * (1 - TOKEN_BUFFER_PERCENTAGE) - reservedTokens
+
+			// Match manageContext's condition: contextPercent >= effectiveThreshold || prevContextTokens > allowedTokens
+			const willManageContext = estimatedUsagePercent >= effectiveThreshold || prevContextTokens > allowedTokens
+
+			// Send condenseTaskContextStarted BEFORE manageContext to show in-progress indicator
+			// This notification must be sent here (not earlier) because the early check uses stale token count
+			// (before user message is added to history), which could incorrectly skip showing the indicator
+			if (willManageContext && autoCondenseContext) {
+				await this.providerRef
+					.deref()
+					?.postMessageToWebview({ type: "condenseTaskContextStarted", text: this.taskId })
+			}
+
 			const truncateResult = await manageContext({
 				messages: this.apiConversationHistory,
 				totalTokens: contextTokens,
@@ -3551,6 +3597,7 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 					truncationId: truncateResult.truncationId,
 					messagesRemoved: truncateResult.messagesRemoved ?? 0,
 					prevContextTokens: truncateResult.prevContextTokens,
+					newContextTokens: truncateResult.newContextTokensAfterTruncation ?? 0,
 				}
 				await this.say(
 					"sliding_window_truncation",
@@ -3564,6 +3611,14 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 					contextTruncation,
 				)
 			}
+
+			// Notify webview that context management is complete (sets isCondensing = false)
+			// This removes the in-progress spinner and allows the completed result to show
+			if (willManageContext && autoCondenseContext) {
+				await this.providerRef
+					.deref()
+					?.postMessageToWebview({ type: "condenseTaskContextResponse", text: this.taskId })
+			}
 		}
 
 		// Get the effective API history by filtering out condensed messages

src/shared/ExtensionMessage.ts

@@ -110,6 +110,7 @@ export interface ExtensionMessage {
 		| "mcpExecutionStatus"
 		| "vsCodeSetting"
 		| "authenticatedUser"
+		| "condenseTaskContextStarted"
 		| "condenseTaskContextResponse"
 		| "singleRouterModelFetchResponse"
 		| "rooCreditBalance"

webview-ui/src/components/chat/ChatRow.tsx

@@ -38,7 +38,7 @@ import { Markdown } from "./Markdown"
 import { CommandExecution } from "./CommandExecution"
 import { CommandExecutionError } from "./CommandExecutionError"
 import { AutoApprovedRequestLimitWarning } from "./AutoApprovedRequestLimitWarning"
-import { CondensingContextRow, ContextCondenseRow } from "./ContextCondenseRow"
+import { InProgressRow, CondensationResultRow, CondensationErrorRow, TruncationResultRow } from "./context-management"
 import CodebaseSearchResultsDisplay from "./CodebaseSearchResultsDisplay"
 import { appendImages } from "@src/utils/imageUtils"
 import { McpExecution } from "./McpExecution"
@@ -1280,18 +1280,27 @@ export const ChatRowContent = ({
 						/>
 					)
 				case "condense_context":
+					// In-progress state
 					if (message.partial) {
-						return <CondensingContextRow />
+						return <InProgressRow eventType="condense_context" />
 					}
-					return message.contextCondense ? <ContextCondenseRow {...message.contextCondense} /> : null
+					// Completed state
+					if (message.contextCondense) {
+						return <CondensationResultRow data={message.contextCondense} />
+					}
+					return null
 				case "condense_context_error":
-					return (
-						<ErrorRow
-							type="error"
-							title={t("chat:contextCondense.errorHeader")}
-							message={message.text || ""}
-						/>
-					)
+					return <CondensationErrorRow errorText={message.text} />
+				case "sliding_window_truncation":
+					// In-progress state
+					if (message.partial) {
+						return <InProgressRow eventType="sliding_window_truncation" />
+					}
+					// Completed state
+					if (message.contextTruncation) {
+						return <TruncationResultRow data={message.contextTruncation} />
+					}
+					return null
 				case "codebase_search_result":
 					let parsed: {
 						content: {

webview-ui/src/components/chat/ChatView.tsx

@@ -826,8 +826,21 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 							break
 					}
 					break
+				case "condenseTaskContextStarted":
+					// Handle both manual and automatic condensation start
+					// We don't check the task ID because:
+					// 1. There can only be one active task at a time
+					// 2. Task switching resets isCondensing to false (see useEffect with task?.ts dependency)
+					// 3. For new tasks, currentTaskItem may not be populated yet due to async state updates
+					if (message.text) {
+						setIsCondensing(true)
+						// Note: sendingDisabled is only set for manual condensation via handleCondenseContext
+						// Automatic condensation doesn't disable sending since the task is already running
+					}
+					break
 				case "condenseTaskContextResponse":
-					if (message.text && message.text === currentTaskItem?.id) {
+					// Same reasoning as above - we trust this is for the current task
+					if (message.text) {
 						if (isCondensing && sendingDisabled) {
 							setSendingDisabled(false)
 						}
@@ -850,7 +863,6 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 			isHidden,
 			sendingDisabled,
 			enableButtons,
-			currentTaskItem,
 			handleChatReset,
 			handleSendMessage,
 			handleSetChatBoxMessage,