Refactor: Unified context-management architecture with improved UX (#9795)

Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>

Author: hannesrudolph

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/__tests__/context-management.spec.ts

@@ -9,7 +9,13 @@ import { BaseProvider } from "../../../api/providers/base-provider"
 import { ApiMessage } from "../../task-persistence/apiMessages"
 import * as condenseModule from "../../condense"
 
-import { TOKEN_BUFFER_PERCENTAGE, estimateTokenCount, truncateConversation, manageContext } from "../index"
+import {
+	TOKEN_BUFFER_PERCENTAGE,
+	estimateTokenCount,
+	truncateConversation,
+	manageContext,
+	willManageContext,
+} from "../index"
 
 // Create a mock ApiHandler for testing
 class MockApiHandler extends BaseProvider {
@@ -1280,4 +1286,125 @@ describe("Context Management", () => {
 			expect(result2.truncationId).toBeDefined()
 		})
 	})
+
+	/**
+	 * Tests for the willManageContext helper function
+	 */
+	describe("willManageContext", () => {
+		it("should return true when context percent exceeds threshold", () => {
+			const result = willManageContext({
+				totalTokens: 60000,
+				contextWindow: 100000, // 60% of context window
+				maxTokens: 30000,
+				autoCondenseContext: true,
+				autoCondenseContextPercent: 50, // 50% threshold
+				profileThresholds: {},
+				currentProfileId: "default",
+				lastMessageTokens: 0,
+			})
+			expect(result).toBe(true)
+		})
+
+		it("should return false when context percent is below threshold", () => {
+			const result = willManageContext({
+				totalTokens: 40000,
+				contextWindow: 100000, // 40% of context window
+				maxTokens: 30000,
+				autoCondenseContext: true,
+				autoCondenseContextPercent: 50, // 50% threshold
+				profileThresholds: {},
+				currentProfileId: "default",
+				lastMessageTokens: 0,
+			})
+			expect(result).toBe(false)
+		})
+
+		it("should return true when tokens exceed allowedTokens even if autoCondenseContext is false", () => {
+			// allowedTokens = contextWindow * (1 - 0.1) - reservedTokens = 100000 * 0.9 - 30000 = 60000
+			const result = willManageContext({
+				totalTokens: 60001, // Exceeds allowedTokens
+				contextWindow: 100000,
+				maxTokens: 30000,
+				autoCondenseContext: false, // Even with auto-condense disabled
+				autoCondenseContextPercent: 50,
+				profileThresholds: {},
+				currentProfileId: "default",
+				lastMessageTokens: 0,
+			})
+			expect(result).toBe(true)
+		})
+
+		it("should return false when autoCondenseContext is false and tokens are below allowedTokens", () => {
+			// allowedTokens = contextWindow * (1 - 0.1) - reservedTokens = 100000 * 0.9 - 30000 = 60000
+			const result = willManageContext({
+				totalTokens: 59999, // Below allowedTokens
+				contextWindow: 100000,
+				maxTokens: 30000,
+				autoCondenseContext: false,
+				autoCondenseContextPercent: 50, // This shouldn't matter since autoCondenseContext is false
+				profileThresholds: {},
+				currentProfileId: "default",
+				lastMessageTokens: 0,
+			})
+			expect(result).toBe(false)
+		})
+
+		it("should use profile-specific threshold when available", () => {
+			const result = willManageContext({
+				totalTokens: 55000,
+				contextWindow: 100000, // 55% of context window
+				maxTokens: 30000,
+				autoCondenseContext: true,
+				autoCondenseContextPercent: 80, // Global threshold 80%
+				profileThresholds: { "test-profile": 50 }, // Profile threshold 50%
+				currentProfileId: "test-profile",
+				lastMessageTokens: 0,
+			})
+			// Should trigger because 55% > 50% (profile threshold)
+			expect(result).toBe(true)
+		})
+
+		it("should fall back to global threshold when profile threshold is -1", () => {
+			const result = willManageContext({
+				totalTokens: 55000,
+				contextWindow: 100000, // 55% of context window
+				maxTokens: 30000,
+				autoCondenseContext: true,
+				autoCondenseContextPercent: 80, // Global threshold 80%
+				profileThresholds: { "test-profile": -1 }, // Profile uses global
+				currentProfileId: "test-profile",
+				lastMessageTokens: 0,
+			})
+			// Should NOT trigger because 55% < 80% (global threshold)
+			expect(result).toBe(false)
+		})
+
+		it("should include lastMessageTokens in the calculation", () => {
+			// Without lastMessageTokens: 49000 tokens = 49%
+			// With lastMessageTokens: 49000 + 2000 = 51000 tokens = 51%
+			const resultWithoutLastMessage = willManageContext({
+				totalTokens: 49000,
+				contextWindow: 100000,
+				maxTokens: 30000,
+				autoCondenseContext: true,
+				autoCondenseContextPercent: 50, // 50% threshold
+				profileThresholds: {},
+				currentProfileId: "default",
+				lastMessageTokens: 0,
+			})
+			expect(resultWithoutLastMessage).toBe(false)
+
+			const resultWithLastMessage = willManageContext({
+				totalTokens: 49000,
+				contextWindow: 100000,
+				maxTokens: 30000,
+				autoCondenseContext: true,
+				autoCondenseContextPercent: 50, // 50% threshold
+				profileThresholds: {},
+				currentProfileId: "default",
+				lastMessageTokens: 2000, // Pushes total to 51%
+			})
+			expect(resultWithLastMessage).toBe(true)
+		})
+	})
 })