1 /* 2 * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 * SPDX-License-Identifier: MIT 4 */ 5 import type { AgentMessage, AgentToolResult } from "@earendil-works/pi-agent-core"; 6 import { 7 buildSessionContext, 8 estimateTokens, 9 findCutPoint, 10 sessionEntryToContextMessages, 11 type ExtensionContext, 12 type ExtensionFactory, 13 type SessionEntry, 14 } from "@earendil-works/pi-coding-agent"; 15 import { formatSavingsCount, showSolPiSavings } from "../../tui.ts"; 16 import { 17 DEFAULT_COMPACTION_ECONOMICS, 18 decideCompaction, 19 type CompactionDecision, 20 } from "./economics.ts"; 21 import { analyzePlanTransition, formatPlanSnapshot, parsePlanSteps } from "./plan.ts"; 22 import { 23 appendOnlineState, 24 initialOnlineState, 25 recordBoundary, 26 recordCompaction, 27 recordCorrection, 28 recordProviderRequest, 29 restoreOnlineState, 30 type OnlineState, 31 type ProgressSummary, 32 } from "./state.ts"; 33 import { registerOnlineTools, type PlanUpdateInput } from "./tools.ts"; 34 35 export const DEFAULT_KEEP_RECENT_TOKENS = 20_000; 36 export const DEFAULT_NATIVE_SUMMARY_TOKEN_ESTIMATE = 1_000; 37 export const BOUNDARY_COMPACTION_INSTRUCTIONS = 38 "Preserve completed work, verification results, important decisions, and remaining work."; 39 export const POST_COMPACTION_PLAN_REMINDER = 40 "Online context compaction finished. The parent task is still active. " + 41 "Before continuing work, call update_plan with a fresh plan for the remaining work."; 42 43 export type OnlineContextCompactOptions = { 44 readonly cacheWriteReadRatio?: number | null; 45 readonly keepRecentTokens?: number; 46 }; 47 48 type PendingBoundary = { readonly toolCallId: string }; 49 type SelectedCompaction = { readonly decision: CompactionDecision }; 50 type CacheDebt = { readonly debtTokens: number; readonly repaymentTokens: number }; 51 type PendingContinuation = { readonly promise: Promise; readonly resolve: () => void }; 52 53 export function resolveKeepRecentTokens(value: number | undefined): number { 54 const resolved = value ?? DEFAULT_KEEP_RECENT_TOKENS; 55 if (!Number.isSafeInteger(resolved) || resolved < 1) { 56 throw new Error("Online Context Compact keepRecentTokens must be a positive safe integer"); 57 } 58 return resolved; 59 } 60 61 function resolveCacheWriteReadRatio(value: number | null | undefined): number | null { 62 if (value === undefined || value === null) return null; 63 if (!Number.isFinite(value) || value < 0) { 64 throw new Error("Online Context Compact cacheWriteReadRatio must be finite and non-negative"); 65 } 66 return value; 67 } 68 69 function tokenEstimate(text: string): number { 70 return Math.ceil(Buffer.byteLength(text) / 4); 71 } 72 73 function result(text: string, details: Readonly>): AgentToolResult>> { 74 return { content: [{ type: "text", text }], details }; 75 } 76 77 function progressSummary(input: PlanUpdateInput, completedStepId: string): ProgressSummary | undefined { 78 const step = input.steps.find((item) => item.id === completedStepId); 79 if (!step || !input.progress) return; 80 return { 81 stepId: step.id, 82 goal: step.goal, 83 filesChanged: [...input.progress.files_changed], 84 verification: [...input.progress.verification], 85 decisions: [...input.progress.decisions], 86 nextWork: input.steps.filter((item) => item.status !== "completed").map((item) => item.goal), 87 }; 88 } 89 90 function compactionMessageCount(entries: readonly SessionEntry[], startIndex: number, endIndex: number): number { 91 let count = 0; 92 for (let index = startIndex; index < endIndex; index++) { 93 const entry = entries[index]; 94 if (entry && entry.type !== "compaction" && sessionEntryToContextMessages(entry).length > 0) count++; 95 } 96 return count; 97 } 98 99 function branchAfterAbort(entries: readonly SessionEntry[]): SessionEntry[] { 100 const last = entries.at(-1); 101 const markerProvider = ["sol", "pi"].join("-"); 102 return [ 103 ...entries, 104 { 105 type: "message", 106 id: "sol-pi-online-context-compact-abort-marker", 107 parentId: last?.id ?? null, 108 timestamp: new Date(0).toISOString(), 109 message: { 110 role: "assistant", 111 content: [], 112 api: markerProvider, 113 provider: markerProvider, 114 model: "aborted", 115 usage: { 116 input: 0, 117 output: 0, 118 cacheRead: 0, 119 cacheWrite: 0, 120 totalTokens: 0, 121 cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, 122 }, 123 stopReason: "aborted", 124 timestamp: 0, 125 }, 126 } as SessionEntry, 127 ]; 128 } 129 130 function nativeCompactionFeasible(entries: readonly SessionEntry[], keepRecentTokens: number): boolean { 131 const path = branchAfterAbort(entries); 132 let startIndex = 0; 133 for (let index = path.length - 1; index >= 0; index--) { 134 const entry = path[index]; 135 if (entry?.type !== "compaction") continue; 136 const keptIndex = path.findIndex((item) => item.id === entry.firstKeptEntryId); 137 startIndex = keptIndex >= 0 ? keptIndex : index + 1; 138 break; 139 } 140 141 const cut = findCutPoint(path, startIndex, path.length, keepRecentTokens); 142 const historyEnd = cut.isSplitTurn ? cut.turnStartIndex : cut.firstKeptEntryIndex; 143 const historyMessages = historyEnd > startIndex ? compactionMessageCount(path, startIndex, historyEnd) : 0; 144 const prefixMessages = 145 cut.isSplitTurn && cut.turnStartIndex >= 0 146 ? compactionMessageCount(path, cut.turnStartIndex, cut.firstKeptEntryIndex) 147 : 0; 148 return historyMessages > 0 || prefixMessages > 0; 149 } 150 151 function validPositiveInteger(value: unknown): value is number { 152 return typeof value === "number" && Number.isSafeInteger(value) && value > 0; 153 } 154 155 export function createOnlineContextCompactExtension(options: OnlineContextCompactOptions = {}): ExtensionFactory { 156 const keepRecentTokens = resolveKeepRecentTokens(options.keepRecentTokens); 157 const cacheWriteReadRatio = resolveCacheWriteReadRatio(options.cacheWriteReadRatio); 158 159 return (pi) => { 160 let state: OnlineState = initialOnlineState(); 161 let restored = false; 162 let observedMessages: readonly AgentMessage[] = []; 163 let pendingBoundary: PendingBoundary | undefined; 164 let selected: SelectedCompaction | undefined; 165 let activeDebt: CacheDebt | undefined; 166 let nextContinuation: PendingContinuation | undefined; 167 let compactionInFlight = false; 168 169 const releaseContinuation = (): void => { 170 const continuation = nextContinuation; 171 nextContinuation = undefined; 172 continuation?.resolve(); 173 }; 174 const releaseParentContinuation = (continuation: PendingContinuation | undefined): void => { 175 if (continuation) setTimeout(continuation.resolve, 0); 176 }; 177 178 const restore = (context: ExtensionContext): void => { 179 releaseContinuation(); 180 state = restoreOnlineState(context.sessionManager.getBranch()); 181 restored = true; 182 observedMessages = buildSessionContext( 183 context.sessionManager.getEntries(), 184 context.sessionManager.getLeafId(), 185 ).messages; 186 pendingBoundary = undefined; 187 selected = undefined; 188 activeDebt = undefined; 189 compactionInFlight = false; 190 }; 191 const ensureRestored = (context: ExtensionContext): void => { 192 if (!restored) restore(context); 193 }; 194 const save = (): void => appendOnlineState(pi, state); 195 const contextTokens = (context: ExtensionContext): number => { 196 const visible = observedMessages.reduce((total, message) => total + estimateTokens(message), 0); 197 const estimated = visible + tokenEstimate(context.getSystemPrompt()); 198 const reported = context.getContextUsage()?.tokens; 199 return validPositiveInteger(reported) ? Math.max(reported, estimated) : estimated; 200 }; 201 202 registerOnlineTools(pi, { 203 updatePlan: async (input) => { 204 ensureRestored(input.context); 205 if (input.signal?.aborted) throw new Error("Plan update was aborted"); 206 const steps = parsePlanSteps(input.steps); 207 if (!steps || steps.length === 0) throw new Error("Plan must contain at least one valid step"); 208 209 const transition = analyzePlanTransition(state.plan, steps); 210 const completedIds = transition.completedSteps.map((step) => step.id); 211 if (completedIds.length > 0) { 212 state = recordBoundary(state, steps, progressSummary(input, completedIds[0] ?? "")); 213 if (!pendingBoundary) pendingBoundary = { toolCallId: input.toolCallId }; 214 } else if (JSON.stringify(state.plan) !== JSON.stringify(steps)) { 215 state = { ...state, plan: [...steps] }; 216 } 217 save(); 218 219 return result( 220 [formatPlanSnapshot(steps), ...transition.advice].join("\n"), 221 { 222 boundary: completedIds.length > 0, 223 completed_step_ids: completedIds, 224 progress_recorded: completedIds.length > 0 && input.progress !== undefined, 225 task_status: "active", 226 plan: steps, 227 }, 228 ); 229 }, 230 }); 231 232 pi.on("session_start", (_event, context) => restore(context)); 233 pi.on("session_before_tree", () => (compactionInFlight ? { cancel: true } : undefined)); 234 pi.on("session_tree", (_event, context) => restore(context)); 235 236 pi.on("context", (event, context) => { 237 ensureRestored(context); 238 observedMessages = [...event.messages]; 239 }); 240 241 pi.on("before_provider_request", (_event, context) => { 242 ensureRestored(context); 243 state = recordProviderRequest(state, contextTokens(context)); 244 save(); 245 }); 246 247 pi.on("input", (event, context) => { 248 if (event.streamingBehavior !== "steer" && !event.text.startsWith("CORRECTION:")) { 249 return { action: "continue" as const }; 250 } 251 ensureRestored(context); 252 pendingBoundary = undefined; 253 selected = undefined; 254 activeDebt = undefined; 255 state = recordCorrection(state); 256 save(); 257 return { action: "continue" as const }; 258 }); 259 260 pi.on("turn_end", (event, context) => { 261 const boundary = pendingBoundary; 262 pendingBoundary = undefined; 263 if (!boundary || selected) return; 264 const toolResult = event.toolResults.find((item) => item.toolCallId === boundary.toolCallId); 265 if ( 266 event.message.role !== "assistant" || 267 event.message.stopReason === "error" || 268 event.message.stopReason === "aborted" || 269 context.signal?.aborted || 270 !toolResult || 271 toolResult.isError 272 ) { 273 return; 274 } 275 276 const usage = context.getContextUsage(); 277 const writeTokens = contextTokens(context); 278 const fixedTokens = tokenEstimate(context.getSystemPrompt()); 279 const archiveTokens = Math.max(0, writeTokens - fixedTokens - keepRecentTokens); 280 const contextWindowTokens = validPositiveInteger(usage?.contextWindow) 281 ? usage.contextWindow 282 : validPositiveInteger(context.model?.contextWindow) 283 ? context.model.contextWindow 284 : null; 285 const averageContextTokenIncrement = 286 state.positiveContextDeltaCount === 0 287 ? null 288 : state.positiveContextDeltaTotal / state.positiveContextDeltaCount; 289 const priced = decideCompaction({ 290 writeTokens, 291 archiveTokens, 292 memoTokens: DEFAULT_NATIVE_SUMMARY_TOKEN_ESTIMATE, 293 contextTokens: writeTokens, 294 completedBoundaryRequestCounts: state.completedBoundaryRequestCounts, 295 remainingBoundaries: state.plan.filter((step) => step.status !== "completed").length, 296 averageContextTokenIncrement, 297 contextWindowTokens, 298 priorCompactionCount: state.nativeCompactionCount, 299 carriedDebtTokens: state.cacheDebtTokens, 300 cacheDebtRepaymentTokens: state.cacheDebtRepaymentTokens, 301 cacheWriteReadRatio, 302 economics: DEFAULT_COMPACTION_ECONOMICS, 303 }); 304 const decision: CompactionDecision = 305 priced.compact && !nativeCompactionFeasible(context.sessionManager.getBranch(), keepRecentTokens) 306 ? { ...priced, compact: false, reason: "native_not_compactable" } 307 : priced; 308 if (!decision.compact) return; 309 310 selected = { decision }; 311 context.abort(); 312 }); 313 314 pi.on("agent_settled", async (_event, context) => { 315 // sendMessage() starts a turn without returning its promise. Capture the 316 // child settlement so print/JSON mode cannot dispose while it is running. 317 const parentContinuation = nextContinuation; 318 nextContinuation = undefined; 319 const pending = selected; 320 selected = undefined; 321 if (!context.isIdle()) { 322 selected = pending; 323 nextContinuation = parentContinuation; 324 return; 325 } 326 if (!pending) { 327 releaseParentContinuation(parentContinuation); 328 return; 329 } 330 331 activeDebt = { 332 debtTokens: pending.decision.writeTokens * (pending.decision.incrementalCacheCostRatio ?? 0), 333 repaymentTokens: Math.max(0, pending.decision.archiveTokens - pending.decision.memoTokens), 334 }; 335 let compacted = false; 336 let compactionError: Error | undefined; 337 try { 338 compactionInFlight = true; 339 await new Promise((resolve) => { 340 let finished = false; 341 const finish = (): void => { 342 if (finished) return; 343 finished = true; 344 resolve(); 345 }; 346 context.compact({ 347 customInstructions: BOUNDARY_COMPACTION_INSTRUCTIONS, 348 onComplete: (compaction) => { 349 try { 350 compacted = true; 351 const removed = Math.max( 352 0, 353 pending.decision.archiveTokens - tokenEstimate(compaction.summary), 354 ); 355 if (removed > 0) { 356 showSolPiSavings( 357 context, 358 "Online Context Compact", 359 formatSavingsCount(removed, "context tokens removed"), 360 ); 361 } 362 } finally { 363 finish(); 364 } 365 }, 366 onError: (error) => { 367 compactionError = error; 368 finish(); 369 }, 370 }); 371 }); 372 compactionInFlight = false; 373 if ( 374 compactionError && 375 compactionError.name !== "AbortError" && 376 compactionError.message !== "Compaction cancelled" 377 ) { 378 throw compactionError; 379 } 380 381 if (compacted) { 382 let resolveContinuation!: () => void; 383 const continuation: PendingContinuation = { 384 promise: new Promise((resolve) => { 385 resolveContinuation = resolve; 386 }), 387 resolve: () => resolveContinuation(), 388 }; 389 nextContinuation = continuation; 390 try { 391 pi.sendMessage( 392 { 393 customType: "sol-pi-online-context-compact", 394 content: POST_COMPACTION_PLAN_REMINDER, 395 display: false, 396 }, 397 { triggerTurn: true }, 398 ); 399 } catch (error) { 400 if (nextContinuation === continuation) nextContinuation = undefined; 401 continuation.resolve(); 402 throw error; 403 } 404 if (context.isIdle() && nextContinuation === continuation) { 405 nextContinuation = undefined; 406 continuation.resolve(); 407 throw new Error("Online context compact continuation did not start"); 408 } 409 await continuation.promise; 410 } 411 } finally { 412 compactionInFlight = false; 413 activeDebt = undefined; 414 releaseParentContinuation(parentContinuation); 415 } 416 }); 417 418 pi.on("session_compact", (event, context) => { 419 ensureRestored(context); 420 state = recordCompaction( 421 state, 422 event.fromExtension || !activeDebt ? { debtTokens: 0, repaymentTokens: 0 } : activeDebt, 423 ); 424 save(); 425 pendingBoundary = undefined; 426 selected = undefined; 427 activeDebt = undefined; 428 observedMessages = buildSessionContext( 429 context.sessionManager.getEntries(), 430 context.sessionManager.getLeafId(), 431 ).messages; 432 }); 433 434 pi.on("session_shutdown", () => { 435 releaseContinuation(); 436 pendingBoundary = undefined; 437 selected = undefined; 438 activeDebt = undefined; 439 compactionInFlight = false; 440 }); 441 }; 442 } 443