pi-openai-toolkit 0.23.1 → 0.23.2

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (42) hide show
  1. package/LICENSE +21 -21
  2. package/NOTICE +83 -83
  3. package/README.md +6 -2
  4. package/README.zh.md +6 -2
  5. package/config.schema.json +278 -278
  6. package/docs/configuration.md +306 -298
  7. package/docs/internals.md +184 -181
  8. package/package.json +136 -135
  9. package/src/compact-client.ts +412 -412
  10. package/src/compaction-projection.ts +35 -35
  11. package/src/config/migration.ts +103 -103
  12. package/src/config.ts +310 -310
  13. package/src/context-management/codex-provider.ts +270 -270
  14. package/src/context-management/history-notes.ts +357 -357
  15. package/src/context-management/manual-compact.ts +247 -599
  16. package/src/context-management/scoped-compaction.ts +89 -0
  17. package/src/context-management/tools.ts +371 -371
  18. package/src/context-management/types.ts +168 -168
  19. package/src/context-management/window-manager.ts +912 -905
  20. package/src/debug.ts +238 -238
  21. package/src/deferred-tool-carryover.ts +282 -282
  22. package/src/details-store.ts +151 -151
  23. package/src/extension-runtime.ts +1919 -1986
  24. package/src/image-generation/extension.ts +189 -189
  25. package/src/image-generation/protocol.ts +579 -579
  26. package/src/image-generation/types.ts +300 -300
  27. package/src/native-fallback.ts +32 -29
  28. package/src/payload-rewrite.ts +515 -515
  29. package/src/pi-context-hook.ts +117 -117
  30. package/src/remote-v2-client.ts +811 -811
  31. package/src/request-context-cache.ts +155 -155
  32. package/src/responses-headers.ts +180 -180
  33. package/src/runtime.ts +865 -865
  34. package/src/serializer.ts +882 -882
  35. package/src/supported-environment.ts +27 -27
  36. package/src/types.ts +565 -565
  37. package/src/web-search/alpha-client.ts +876 -876
  38. package/src/web-search/extension.ts +530 -530
  39. package/src/web-search/payload.ts +299 -299
  40. package/src/web-search/prompt.ts +60 -60
  41. package/src/web-search/types.ts +166 -166
  42. package/tsconfig.check.json +50 -50
@@ -1,905 +1,912 @@
1
- import { AsyncLocalStorage } from "node:async_hooks";
2
- import { randomUUID } from "node:crypto";
3
- import { isDeepStrictEqual } from "node:util";
4
- import type { AgentMessage } from "@earendil-works/pi-agent-core";
5
- import type { ProviderHeaders } from "@earendil-works/pi-ai";
6
- import {
7
- buildSessionProjection,
8
- calculateContextTokens,
9
- estimateTokens,
10
- type CompactionResult,
11
- type ExtensionAPI,
12
- type ExtensionContext,
13
- type SessionBeforeCompactEvent,
14
- type SessionEntry,
15
- } from "@earendil-works/pi-coding-agent";
16
- import { ContextWindowBudget, type ContextRemaining } from "./window-budget";
17
- import {
18
- rewriteEncryptedToolOutputs,
19
- rewriteWindowHeaders,
20
- rewriteWindowPayload,
21
- } from "./window-request";
22
- import {
23
- CONTEXT_WINDOW_COMPACTION_STRATEGY,
24
- CONTEXT_WINDOW_COMPACTION_SUMMARY,
25
- isContextWindowBoundary,
26
- preserveSystemHead,
27
- renderContextWindowMessage,
28
- sendContextWindowMessage,
29
- } from "./messages";
30
- import {
31
- CODEX_CONTEXT_WINDOW_MESSAGE_TYPE,
32
- CONTEXT_MANAGEMENT_PROTOCOL,
33
- type CodexContextManagementMessageDetails,
34
- type ContextWindowCompactionDetails,
35
- type ContextWindowIdentity,
36
- isCodexContextManagementMessageDetails,
37
- isNonEmptyString,
38
- isRecord,
39
- isContextWindowCompactionDetails,
40
- type NotesCheckpointReceipt,
41
- } from "./types";
42
- import {
43
- encodeEncryptedOutputForContext,
44
- isSuccessfulHistoryNotesToolResult,
45
- loadHistoryNotesThreadHint,
46
- } from "./history-notes";
47
- import { rewriteContextNamespaceTools } from "./namespace-tools";
48
- import { hasInvalidatedCompactionContext } from "../payload-rewrite";
49
-
50
- interface StartContextWindowOptions {
51
- triggerTurn: boolean;
52
- signal?: AbortSignal;
53
- trimPreviousWindow: boolean;
54
- /** Captured with the tool activation policy, before any awaited backend work. */
55
- gatewayModels?: readonly string[];
56
- /** Persist an owned handoff's exact target before the marker can be queued. */
57
- beforeSchedule?: (identity: ContextWindowIdentity) => void;
58
- }
59
-
60
- type ThreadHintLoader = (
61
- ctx: ExtensionContext,
62
- signal?: AbortSignal,
63
- gatewayModels?: readonly string[],
64
- ) => Promise<string | undefined>;
65
-
66
- type WindowBoundaryEntry = Extract<SessionEntry, { type: "custom_message" }> & {
67
- details: CodexContextManagementMessageDetails;
68
- };
69
-
70
- /** Canonical input and real source anchors for one read-only manual compression. */
71
- export type RemoteWindowCompactionInput = {
72
- messages: AgentMessage[];
73
- firstKeptEntryId: string;
74
- /** Only a newer compaction can own this window; pre-window checkpoints are retired. */
75
- owningCompaction?: Extract<SessionEntry, { type: "compaction" }>;
76
- /** Entries whose visible contributions Pi will retain alongside the new checkpoint. */
77
- retainedMessages: AgentMessage[];
78
- windowBoundaryIndex: number;
79
- };
80
-
81
- /**
82
- * In-process duplicate guard for one scheduled rollover. It is anchored to the
83
- * session that scheduled it and to the exact target window id, so it can never
84
- * be satisfied by an older or foreign marker. It is not checkpoint evidence:
85
- * only a persisted notes result unlocks the gate.
86
- */
87
- type PendingRollover = {
88
- sessionId?: string;
89
- targetWindowId: string;
90
- /** The current request already contains this marker; never persistence proof. */
91
- projected?: boolean;
92
- };
93
-
94
- export class CodexContextWindowManager {
95
- private identity: ContextWindowIdentity | undefined;
96
- private sessionId: string | undefined;
97
- private restoredMarkerId: string | undefined;
98
- private restoredCompactionId: string | undefined;
99
- private boundary: WindowBoundaryResolution = { kind: "none" };
100
- private readonly budget = new ContextWindowBudget();
101
- private pendingRollover: PendingRollover | undefined;
102
- private readonly projectionDiagnostics = new Set<string>();
103
- private lastKnownSystemHead: AgentMessage | undefined;
104
- private readonly compactionProjection = new AsyncLocalStorage<boolean>();
105
-
106
- get isPreparingCompaction(): boolean { return this.compactionProjection.getStore() === true; }
107
-
108
- /** Keep read-only preparation scoped across awaits without affecting other live requests. */
109
- withCompactionProjection<T>(prepare: () => Promise<T>): Promise<T> {
110
- return this.compactionProjection.run(true, prepare);
111
- }
112
- /**
113
- * Pi's getContextUsage() measures the durable transcript. Remote windows are
114
- * projected only by this extension, so remember the last provider-visible
115
- * estimate for budget decisions and get_context_remaining.
116
- */
117
- private projectedContextTokens: number | undefined;
118
- private readonly loadThreadHint: ThreadHintLoader;
119
-
120
- constructor(loadThreadHint?: ThreadHintLoader) {
121
- this.loadThreadHint = loadThreadHint ?? ((ctx, signal, gatewayModels) => loadHistoryNotesThreadHint(ctx, signal, gatewayModels));
122
- }
123
-
124
- reset(): void {
125
- this.resetWindowState();
126
- this.pendingRollover = undefined;
127
- this.projectionDiagnostics.clear();
128
- this.lastKnownSystemHead = undefined;
129
- }
130
-
131
- currentIdentity(): ContextWindowIdentity | undefined {
132
- return this.identity ? { ...this.identity } : undefined;
133
- }
134
- /** Derived eligibility, never a replayed marker command. */
135
- hasPendingTrim(): boolean {
136
- return this.boundary.kind === "effective" && !this.boundary.cleaned &&
137
- this.boundary.entry.details.contextManagement.trimPreviousWindow === true;
138
- }
139
-
140
- /** Historical identity does not own a newer native/opaque checkpoint. */
141
- hasEffectiveWindow(): boolean {
142
- if (!this.identity || this.boundary.kind === "ambiguous") return false;
143
- return this.boundary.kind !== "superseded" ||
144
- (this.pendingRollover?.projected === true && this.pendingRollover.targetWindowId === this.identity.currentWindowId);
145
- }
146
-
147
- hasSupersededWindow(): boolean {
148
- return this.boundary.kind === "superseded";
149
- }
150
-
151
- private resetWindowState(): void {
152
- this.identity = undefined;
153
- this.sessionId = undefined;
154
- this.restoredMarkerId = undefined;
155
- this.restoredCompactionId = undefined;
156
- this.boundary = { kind: "none" };
157
- this.budget.reset();
158
- this.projectedContextTokens = undefined;
159
- }
160
-
161
- restore(entries: readonly SessionEntry[], sessionId?: string): void {
162
- // A rebuild of derived state must not drop an outstanding rollover guard;
163
- // only a changed session or an observed target marker retires it below.
164
- this.resetWindowState();
165
- this.sessionId = sessionId;
166
- for (const entry of entries) {
167
- if (entry.type === "compaction") {
168
- this.restoredCompactionId = entry.id;
169
- continue;
170
- }
171
- if (entry.type !== "custom_message" || entry.customType !== CODEX_CONTEXT_WINDOW_MESSAGE_TYPE) continue;
172
- if (!couldBelongToSession(entry.details, sessionId)) continue;
173
- if (!isCodexContextManagementMessageDetails(entry.details)) {
174
- throw new Error("Malformed persisted Codex context-window message");
175
- }
176
- if (!matchesSession(entry.details.sessionId, sessionId)) continue;
177
- const details = entry.details.contextManagement;
178
- this.restoredMarkerId = entry.id;
179
- if (details.kind === "window") {
180
- this.identity = identityFromDetails(entry.details);
181
- }
182
- this.budget.restore(details.kind, details.currentWindowId);
183
- }
184
- this.boundary = resolveWindowBoundary(entries, sessionId);
185
- this.retireSatisfiedOrStalePending(entries, sessionId);
186
- }
187
-
188
- /** Rebuild state when Pi navigates to a different session/branch. */
189
- synchronize(ctx: Pick<ExtensionContext, "sessionManager">): void {
190
- const entries = ctx.sessionManager.getBranch();
191
- const sessionId = ctx.sessionManager.getSessionId();
192
- this.retireSatisfiedOrStalePending(entries, sessionId);
193
- const latestMarkerId = findLatestContextMarkerId(entries, sessionId);
194
- let latestCompactionId: string | undefined;
195
- for (let index = entries.length - 1; index >= 0; index--) {
196
- if (entries[index]!.type === "compaction") {
197
- latestCompactionId = entries[index]!.id;
198
- break;
199
- }
200
- }
201
- // Hook-provided compactions may have no success callback. The persisted
202
- // branch also detects navigation before/after a commit with the same marker.
203
- if (this.sessionId !== sessionId || this.restoredMarkerId !== latestMarkerId || this.restoredCompactionId !== latestCompactionId) {
204
- this.restore(entries, sessionId);
205
- }
206
- // Edits and sibling navigation may leave marker/compaction IDs unchanged.
207
- this.boundary = resolveWindowBoundary(entries, sessionId);
208
- }
209
-
210
- /**
211
- * Drop a pending rollover only for durable reasons: the target window marker
212
- * is now persisted in this session, or the session changed. A queued marker
213
- * visible in an in-memory message list is not durable evidence.
214
- */
215
- private retireSatisfiedOrStalePending(
216
- entries: readonly SessionEntry[],
217
- sessionId: string | undefined,
218
- ): void {
219
- const pending = this.pendingRollover;
220
- if (!pending) return;
221
- if (pending.sessionId !== undefined && pending.sessionId !== sessionId) {
222
- this.pendingRollover = undefined;
223
- return;
224
- }
225
- if (hasPersistedWindowBoundary(entries, sessionId, pending.targetWindowId)) {
226
- this.pendingRollover = undefined;
227
- }
228
- }
229
-
230
- ensureInitialized(pi: ExtensionAPI, ctx: ExtensionContext, active: boolean): void {
231
- if (!active) return;
232
- this.synchronize(ctx);
233
- if (this.identity) return;
234
- const windowId = randomUUID();
235
- this.sendWindowMessage(
236
- pi,
237
- ctx,
238
- { firstWindowId: windowId, currentWindowId: windowId, windowNumber: 0 },
239
- { triggerTurn: false, trimPreviousWindow: false },
240
- );
241
- }
242
-
243
- project(
244
- messages: readonly AgentMessage[],
245
- mode: "off" | "local" | "remote",
246
- /** Delay encoding until verified retained copies have been removed. */
247
- encodeEncryptedOutputs = true,
248
- ): AgentMessage[] {
249
- if (this.isPreparingCompaction) return this.projectCompactionContext(messages);
250
- // Recompute the request-only preview on every projection. Observing a
251
- // queued marker must not authorize maintenance or retire the duplicate guard.
252
- if (this.pendingRollover) this.pendingRollover.projected = false;
253
- if (mode === "off") {
254
- this.projectedContextTokens = undefined;
255
- return messages.filter(
256
- (message) => message.role !== "custom" || message.customType !== CODEX_CONTEXT_WINDOW_MESSAGE_TYPE,
257
- );
258
- }
259
- if (this.boundary.kind === "ambiguous") {
260
- throw new Error("Persisted Codex context-window boundary missing or edited in canonical projection");
261
- }
262
- if (this.boundary.kind === "superseded") this.identity = identityFromDetails(this.boundary.entry.details);
263
- let boundaryIndex = -1;
264
- for (let index = 0; index < messages.length; index += 1) {
265
- const message = messages[index]!;
266
- if (
267
- message.role === "custom" &&
268
- message.customType === CODEX_CONTEXT_WINDOW_MESSAGE_TYPE
269
- ) {
270
- if (!couldBelongToSession(message.details, this.sessionId)) continue;
271
- if (!isCodexContextManagementMessageDetails(message.details)) {
272
- throw new Error("Malformed persisted Codex context-window message");
273
- }
274
- if (!matchesSession(message.details.sessionId, this.sessionId)) continue;
275
- }
276
- if (!isContextWindowBoundary(message)) continue;
277
- const queuedRollover = mode === "remote" && this.pendingRollover?.targetWindowId === message.details.contextManagement.currentWindowId;
278
- if (this.boundary.kind === "superseded" && !queuedRollover) continue;
279
- // Local consumers follow the synchronized persisted identity, not a
280
- // queued marker from the request view. They never own the Remote session.
281
- if (mode === "local" && message.details.contextManagement.currentWindowId !== this.identity?.currentWindowId) continue;
282
- boundaryIndex = index;
283
- if (mode === "remote") {
284
- this.identity = identityFromDetails(message.details);
285
- if (this.pendingRollover) this.pendingRollover.projected = queuedRollover;
286
- }
287
- }
288
- if (boundaryIndex < 0 && this.boundary.kind === "effective") {
289
- throw new Error("Persisted Codex context-window boundary missing from effective projection");
290
- }
291
- // Projection observes the request view, never durable session state. A
292
- // queued rollover marker must not clear the duplicate guard here.
293
- const trimmed = boundaryIndex < 0 ? [...messages] : this.trimToWindow(messages, boundaryIndex, mode === "local");
294
- const projected = mode === "remote" ? (encodeEncryptedOutputs ? projectEncryptedToolResults(trimmed) : trimmed) : trimmed.filter(
295
- (message) => message.role !== "custom" || message.customType !== CODEX_CONTEXT_WINDOW_MESSAGE_TYPE,
296
- );
297
- this.projectedContextTokens = mode === "remote" && this.identity
298
- ? estimateProjectedWindowTokens(projected)
299
- : undefined;
300
- return projected;
301
- }
302
-
303
- /**
304
- * Strict, read-only projection. Unlike live Remote projection, an unsafe head
305
- * never falls back to the full transcript, updates identity, or records budget.
306
- */
307
- projectCompactionContext(messages: readonly AgentMessage[], encodeEncryptedOutputs = false): AgentMessage[] {
308
- if (this.pendingRollover || (this.boundary.kind !== "effective" && this.boundary.kind !== "superseded")) {
309
- throw new Error("unsafe-compaction-window-boundary");
310
- }
311
- let boundaryIndex = 0;
312
- if (this.boundary.kind === "effective") {
313
- const boundary = this.boundary.entry;
314
- boundaryIndex = messages.findIndex((message) => isContextWindowBoundary(message)
315
- && isDeepStrictEqual(message.details, boundary.details)
316
- && isDeepStrictEqual(message.content, boundary.content));
317
- if (boundaryIndex < 0) throw new Error("compaction-window-boundary-missing");
318
- }
319
- const repair = preserveSystemHead({ messages, boundaryIndex });
320
- if (!repair.head || repair.shrinkRejected) throw new Error("compaction-system-tool-head-lost");
321
- return encodeEncryptedOutputs ? projectEncryptedToolResults(repair.messages) : [...repair.messages];
322
- }
323
-
324
- prepareRemoteCompaction(event: SessionBeforeCompactEvent, ctx: ExtensionContext): RemoteWindowCompactionInput {
325
- this.synchronize(ctx);
326
- if (this.hasPendingRollover(ctx) || (this.boundary.kind !== "effective" && this.boundary.kind !== "superseded")) {
327
- throw new Error("unsafe-compaction-window-boundary");
328
- }
329
- const branch = ctx.sessionManager.getBranch();
330
- if (!isDeepStrictEqual(event.branchEntries, branch)) throw new Error("compaction-branch-changed");
331
- const projection = ctx.sessionManager.buildSessionProjection();
332
- // Never manufacture source IDs or substitute a raw-history reconstruction.
333
- if (!projection || !Array.isArray(projection.entries) || !Array.isArray(projection.messages)
334
- || !isDeepStrictEqual(projection.messages, projection.entries.flatMap((entry) => entry.messages))
335
- || new Set(projection.entries.map((entry) => entry.sourceEntry.id)).size !== projection.entries.length
336
- || projection.entries.some((entry) => !branch.some((raw) => raw.id === entry.sourceEntry.id && isDeepStrictEqual(raw, entry.sourceEntry)))) {
337
- throw new Error("compaction-source-projection-unavailable");
338
- }
339
- const windowBoundaryIndex = branch.findIndex((entry) => entry.id === this.boundaryEntryId());
340
- if (windowBoundaryIndex < 0) throw new Error("compaction-window-boundary-missing");
341
- const latest = [...branch].reverse().find((entry) => entry.type === "compaction");
342
- if (this.boundary.kind === "effective" && this.boundary.cleaned) {
343
- // Details alone do not turn an arbitrary summary into Toolkit cleanup.
344
- if (!latest || latest.type !== "compaction" || latest.fromHook !== true
345
- || latest.summary !== CONTEXT_WINDOW_COMPACTION_SUMMARY
346
- || !isContextWindowCompactionDetails(latest.details)
347
- || latest.details.windowId !== this.boundary.entry.details.contextManagement.currentWindowId
348
- || latest.firstKeptEntryId !== this.boundary.entry.id) throw new Error("unverified-window-cleanup");
349
- }
350
- const owningCompaction = this.boundary.kind === "superseded" && latest?.type === "compaction" ? latest : undefined;
351
- if (this.boundary.kind === "superseded" && !owningCompaction) throw new Error("compaction-owner-missing");
352
- const lowerId = owningCompaction?.id ?? this.boundary.entry.id;
353
- const lowerIndex = branch.findIndex((entry) => entry.id === lowerId);
354
- const requestedIndex = branch.findIndex((entry) => entry.id === event.preparation.firstKeptEntryId);
355
- if (requestedIndex < 0 || lowerIndex < windowBoundaryIndex) throw new Error("unsafe-compaction-retained-start");
356
- const firstKeptEntryId = requestedIndex < lowerIndex ? lowerId : event.preparation.firstKeptEntryId;
357
- const keptIndex = Math.max(requestedIndex, lowerIndex);
358
- const safeSources = new Set(projection.entries.filter((entry) => entry.messages.length > 0).map((entry) => entry.sourceEntry.id));
359
- if (!safeSources.has(firstKeptEntryId)) throw new Error("unsafe-compaction-retained-start");
360
- const retainedIds = new Set(branch.slice(keptIndex).map((entry) => entry.id));
361
- return {
362
- messages: this.projectCompactionContext(structuredClone(projection.messages)),
363
- firstKeptEntryId,
364
- owningCompaction,
365
- retainedMessages: structuredClone(projection.entries.filter((entry) => retainedIds.has(entry.sourceEntry.id)
366
- && entry.sourceEntry.type !== "compaction").flatMap((entry) => entry.messages)),
367
- windowBoundaryIndex,
368
- };
369
- }
370
-
371
- private boundaryEntryId(): string | undefined {
372
- return this.boundary.kind === "effective" || this.boundary.kind === "superseded" ? this.boundary.entry.id : undefined;
373
- }
374
-
375
- /**
376
- * Trim everything before the window boundary while keeping Pi's prompt/tool head.
377
- *
378
- * When the head cannot be preserved the trim is refused and the full message list is
379
- * sent, because a window without tools silently degrades into a chat the model can
380
- * only narrate. Each refusal is reported once per window and reason.
381
- */
382
- private trimToWindow(messages: readonly AgentMessage[], boundaryIndex: number, local = false): AgentMessage[] {
383
- const repair = preserveSystemHead({ messages, boundaryIndex, fallbackHead: this.lastKnownSystemHead });
384
- if (repair.head) this.lastKnownSystemHead = repair.head;
385
- if (repair.shrinkRejected) {
386
- // Remote requests retain the existing tool-preservation fallback. A
387
- // local request must never recover tools by leaking retired windows.
388
- if (local) throw new Error("Cannot preserve the system/tool head in local context-window projection");
389
- this.recordProjectionDiagnostic(repair.lostToolNames);
390
- }
391
- return [...repair.messages];
392
- }
393
-
394
- private recordProjectionDiagnostic(lostToolNames: readonly string[]): void {
395
- const key = `${this.identity?.currentWindowId ?? "unknown"}|${[...lostToolNames].sort().join(",")}`;
396
- this.projectionDiagnostics.add(key);
397
- }
398
-
399
- /**
400
- * Drain the pending `tool-loadout-shrink` reports. Returns one entry per window and
401
- * lost-tool set, so a caller can surface the incident without flooding the UI.
402
- */
403
- takeProjectionDiagnostics(): string[] {
404
- if (this.projectionDiagnostics.size === 0) return [];
405
- const pending = [...this.projectionDiagnostics];
406
- this.projectionDiagnostics.clear();
407
- return pending;
408
- }
409
-
410
- async startNewWindow(
411
- pi: ExtensionAPI,
412
- ctx: ExtensionContext,
413
- options: StartContextWindowOptions,
414
- ): Promise<boolean> {
415
- this.synchronize(ctx);
416
- // A same-session rollover whose target marker is not persisted yet is
417
- // still in flight: report it instead of scheduling a second window.
418
- if (this.hasPendingRollover(ctx)) return false;
419
- if (options.signal?.aborted) throw new Error("Remote context rollover was aborted");
420
- const current = this.identity;
421
- const liveSessionId = ctx.sessionManager.getSessionId();
422
- const currentWindowId = randomUUID();
423
- const next: ContextWindowIdentity = current
424
- ? {
425
- firstWindowId: current.firstWindowId,
426
- currentWindowId,
427
- previousWindowId: current.currentWindowId,
428
- windowNumber: current.windowNumber + 1,
429
- }
430
- : { firstWindowId: currentWindowId, currentWindowId, windowNumber: 0 };
431
- this.pendingRollover = { sessionId: liveSessionId, targetWindowId: currentWindowId };
432
- try {
433
- const checkpoint = current
434
- ? findLatestNotesCheckpointSinceBoundary(ctx.sessionManager.getBranch(), liveSessionId)
435
- : undefined;
436
- let threadHint: string | undefined;
437
- if (current) {
438
- try {
439
- threadHint = await this.loadThreadHint(ctx, options.signal, options.gatewayModels);
440
- } catch (error) {
441
- if (options.signal?.aborted) throw error;
442
- }
443
- }
444
- if (options.signal?.aborted) throw new Error("Remote context rollover was aborted");
445
- options.beforeSchedule?.(next);
446
- this.sendWindowMessage(pi, ctx, next, options, threadHint, checkpoint);
447
- return true;
448
- } catch (error) {
449
- this.pendingRollover = undefined;
450
- throw error;
451
- }
452
- }
453
-
454
- recordBudget(
455
- pi: ExtensionAPI,
456
- ctx: ExtensionContext,
457
- active: boolean,
458
- contextReminderThresholdPercent: number,
459
- contextTokens?: number,
460
- ): void {
461
- if (!active || !this.identity || contextReminderThresholdPercent <= 0) return;
462
- // Right after a rollover, Pi's usage anchor still reports the previous
463
- // window's last request until the new window's first request completes.
464
- // Budget decisions taken in that gap act on stale numbers: they burn the
465
- // once-per-window reminder on a false alarm seconds after a successful
466
- // rollover, leaving the window silent for the rest of its life.
467
- if (!hasAssistantUsageSinceWindowBoundary(ctx.sessionManager.getBranch(), this.sessionId)) return;
468
- const reminder = this.budget.record(
469
- ctx,
470
- this.identity,
471
- contextTokens ?? this.projectedContextTokens,
472
- contextReminderThresholdPercent,
473
- );
474
- if (!reminder) return;
475
- sendContextWindowMessage(
476
- pi,
477
- reminder.content,
478
- reminder.kind,
479
- this.identity,
480
- { triggerTurn: reminder.kind === "fallback", sessionId: ctx.sessionManager.getSessionId() },
481
- );
482
- }
483
-
484
- remaining(ctx: ExtensionContext, contextTokens?: number): ContextRemaining {
485
- return this.budget.remaining(ctx, this.identity, contextTokens ?? this.projectedContextTokens);
486
- }
487
-
488
- /**
489
- * True when a notes checkpoint succeeded in the current window (after the
490
- * latest boundary). The scan always uses the live session id: a cached id can
491
- * belong to a session Pi already navigated away from, which would filter the
492
- * current window's marker and result out as foreign.
493
- */
494
- hasNotesCheckpointSinceBoundary(
495
- ctx: Pick<ExtensionContext, "sessionManager">,
496
- ): boolean {
497
- return findNotesCheckpointSinceBoundary(
498
- ctx.sessionManager.getBranch(),
499
- ctx.sessionManager.getSessionId(),
500
- );
501
- }
502
-
503
- /**
504
- * True while a rollover scheduled for the current session is awaiting its
505
- * persisted target marker. Call `synchronize()` first so a marker that has
506
- * already been persisted or a session switch retires the guard.
507
- */
508
- hasPendingRollover(ctx: Pick<ExtensionContext, "sessionManager">): boolean {
509
- const pending = this.pendingRollover;
510
- if (!pending) return false;
511
- return pending.sessionId === undefined || pending.sessionId === ctx.sessionManager.getSessionId();
512
- }
513
-
514
- /**
515
- * Whether a rollover may be scheduled from the current window.
516
- *
517
- * A window that was entered by a context switch has to do real work before it switches
518
- * again: without this guard a model can checkpoint and roll over on its very first turn,
519
- * which is how long sessions start looping between notes and new_context.
520
- */
521
- canRolloverFromCurrentWindow(ctx: ExtensionContext): boolean {
522
- if (!this.identity || this.identity.windowNumber === 0) return true;
523
- // An exhausted window must always be able to escape: the budget fallback tells the
524
- // model to checkpoint and switch, and refusing that would strand the session. An
525
- // unmeasurable window keeps the guard, since nothing there is pushing for a switch.
526
- const { remainingTokens } = this.remaining(ctx);
527
- if (remainingTokens !== undefined && remainingTokens <= 0) return true;
528
- return hasSubstantiveToolResultSinceBoundary(ctx.sessionManager.getBranch(), this.sessionId);
529
- }
530
-
531
- prepareCompaction(
532
- event: SessionBeforeCompactEvent,
533
- ): { cancel: true } | { compaction: CompactionResult<ContextWindowCompactionDetails> } {
534
- const boundary = resolveWindowBoundary(event.branchEntries, this.sessionId);
535
- if (boundary.kind !== "effective" || boundary.cleaned || !boundary.entry.details.contextManagement.trimPreviousWindow) {
536
- return { cancel: true };
537
- }
538
- // An unpersisted/cancelled proposal changes no eligibility.
539
- return { compaction: this.createCompaction(event, boundary.entry) };
540
- }
541
-
542
- /**
543
- * Detach a local consumer at the durable Remote boundary, without summarizing
544
- * any Remote-managed history. Unlike a scheduled rollover trim, this needs
545
- * no pending trim and remains retryable until Pi persists the result.
546
- *
547
- * A consumed boundary is not an error. Once Pi has persisted the no-summary
548
- * boundary, `buildSessionProjection()` already starts at that marker, so
549
- * retired Remote history is outside every producer's input and the ordinary
550
- * compaction chain may summarize the remaining local window. Returning
551
- * `undefined` hands the event back to that chain; returning `{ cancel: true }`
552
- * would strand the session with no working compaction path at all.
553
- */
554
- prepareDetachedCompaction(
555
- event: SessionBeforeCompactEvent,
556
- ): { compaction: CompactionResult<ContextWindowCompactionDetails> } | undefined {
557
- const boundary = resolveWindowBoundary(event.branchEntries, this.sessionId);
558
- if (boundary.kind === "ambiguous") throw new Error("Persisted Codex context-window boundary missing or edited");
559
- if (boundary.kind !== "effective") return undefined;
560
- // A matching no-summary boundary is already the effective compaction anchor.
561
- // Nothing is left to detach, and writing a second boundary for the same window
562
- // would move the retained start backwards.
563
- if (boundary.cleaned) return undefined;
564
- const compaction = this.createCompaction(event, boundary.entry);
565
- // Pi's preparation/tokensBefore describes the durable branch, which includes
566
- // Remote-managed history. Detached local accounting must use the same current
567
- // window projection that the non-window provider sees instead.
568
- return {
569
- compaction: {
570
- ...compaction,
571
- tokensBefore: this.estimateLocalContextTokens(event.branchEntries),
572
- },
573
- };
574
- }
575
-
576
- private estimateLocalContextTokens(entries: readonly SessionEntry[]): number {
577
- const messages = buildSessionProjection([...entries]).messages;
578
- return this.project(messages, "local").reduce((total, message) => total + estimateTokens(message as never), 0);
579
- }
580
-
581
- private createCompaction(event: SessionBeforeCompactEvent, boundary: WindowBoundaryEntry): CompactionResult<ContextWindowCompactionDetails> {
582
- return {
583
- summary: CONTEXT_WINDOW_COMPACTION_SUMMARY,
584
- firstKeptEntryId: boundary.id,
585
- tokensBefore: event.preparation.tokensBefore,
586
- details: {
587
- protocol: CONTEXT_MANAGEMENT_PROTOCOL,
588
- strategy: CONTEXT_WINDOW_COMPACTION_STRATEGY,
589
- windowId: boundary.details.contextManagement.currentWindowId,
590
- },
591
- };
592
- }
593
-
594
- rewritePayload(payload: unknown, ctx: ExtensionContext): unknown {
595
- const withMetadata = rewriteWindowPayload(payload, ctx, this.identity);
596
- return rewriteContextNamespaceTools(withMetadata, { encrypted: true });
597
- }
598
-
599
- rewriteHeaders(headers: ProviderHeaders, ctx: ExtensionContext): void {
600
- rewriteWindowHeaders(headers, ctx, this.identity);
601
- }
602
-
603
- private sendWindowMessage(
604
- pi: ExtensionAPI,
605
- ctx: ExtensionContext,
606
- identity: ContextWindowIdentity,
607
- options: StartContextWindowOptions,
608
- threadHint?: string,
609
- checkpoint?: NotesCheckpointReceipt,
610
- ): void {
611
- sendContextWindowMessage(
612
- pi,
613
- renderContextWindowMessage(identity, threadHint, checkpoint),
614
- "window",
615
- identity,
616
- { triggerTurn: options.triggerTurn, sessionId: ctx.sessionManager.getSessionId() },
617
- options.trimPreviousWindow,
618
- );
619
- this.identity = identity;
620
- this.sessionId = ctx.sessionManager.getSessionId();
621
- this.restoredMarkerId = undefined;
622
- // sendMessage() acceptance is not persistence. The pending rollover stays
623
- // armed until synchronize() observes this target window in the persisted
624
- // branch (or the session changes).
625
- }
626
- }
627
-
628
- type WindowBoundaryResolution =
629
- | { kind: "none" | "ambiguous" }
630
- | { kind: "superseded"; entry: WindowBoundaryEntry }
631
- | { kind: "effective"; entry: WindowBoundaryEntry; cleaned: boolean };
632
-
633
- /** Pi's projection owns context; raw chronology only establishes window ownership. */
634
- function resolveWindowBoundary(entries: readonly SessionEntry[], sessionId?: string): WindowBoundaryResolution {
635
- const boundary = findLatestWindowBoundaryEntry(entries, sessionId);
636
- if (!boundary) return { kind: "none" };
637
- const projection = buildSessionProjection([...entries]);
638
- const latest = projection.entries.find((entry) => entry.sourceEntry.type === "compaction")?.sourceEntry;
639
- const boundaryIndex = entries.findIndex((entry) => entry.id === boundary.id);
640
- const laterCompaction = latest?.type === "compaction" && entries.findIndex((entry) => entry.id === latest.id) > boundaryIndex
641
- ? latest : undefined;
642
- if (laterCompaction && hasInvalidatedCompactionContext([...entries], laterCompaction.id)) return { kind: "ambiguous" };
643
- const marker = projection.entries.find((entry) => entry.sourceEntry.id === boundary.id);
644
- // A surviving edit can replace content without changing custom details.
645
- const edited = entries.some((entry) => entry.type === "context_edit" && entry.targetId === boundary.id &&
646
- projection.entries.some((projected) => projected.sourceEntry.id === entry.id));
647
- if (edited) return { kind: "ambiguous" };
648
- const cleaned = laterCompaction !== undefined &&
649
- isContextWindowCompactionDetails(laterCompaction.details) &&
650
- laterCompaction.details.windowId === boundary.details.contextManagement.currentWindowId &&
651
- laterCompaction.firstKeptEntryId === boundary.id;
652
- if (laterCompaction && !cleaned) return { kind: "superseded", entry: boundary };
653
- if (!marker?.messages.some(isContextWindowBoundary)) return { kind: "ambiguous" };
654
- return { kind: "effective", entry: boundary, cleaned };
655
- }
656
-
657
- function identityFromDetails(details: CodexContextManagementMessageDetails): ContextWindowIdentity {
658
- const context = details.contextManagement;
659
- return {
660
- firstWindowId: context.firstWindowId,
661
- currentWindowId: context.currentWindowId,
662
- ...(context.previousWindowId ? { previousWindowId: context.previousWindowId } : {}),
663
- windowNumber: context.windowNumber,
664
- };
665
- }
666
-
667
- const NOTES_CHECKPOINT_ACTIONS: ReadonlySet<string> = new Set(["append_to_file", "write_file"]);
668
-
669
- /** Tools that only move context around; they never count as work done in a window. */
670
- const CONTEXT_MANAGEMENT_TOOL_NAMES: ReadonlySet<string> = new Set([
671
- "new_context",
672
- "get_context_remaining",
673
- "history",
674
- "notes",
675
- ]);
676
-
677
- /**
678
- * Whether the current window produced a tool result from anything other than the
679
- * context-handoff tools. Derived from persisted entries so it survives restarts.
680
- */
681
- export function hasSubstantiveToolResultSinceBoundary(
682
- entries: readonly SessionEntry[],
683
- sessionId?: string,
684
- ): boolean {
685
- const boundaryIndex = findLatestWindowBoundaryIndex(entries, sessionId);
686
- if (boundaryIndex < 0) return true;
687
- for (let index = boundaryIndex + 1; index < entries.length; index += 1) {
688
- const entry = entries[index]!;
689
- if (entry.type !== "message") continue;
690
- const message = entry.message as unknown as { role?: string; toolName?: string };
691
- if (message.role !== "toolResult" || typeof message.toolName !== "string") continue;
692
- if (!CONTEXT_MANAGEMENT_TOOL_NAMES.has(message.toolName)) return true;
693
- }
694
- return false;
695
- }
696
-
697
- /**
698
- * Whether the branch contains a successful assistant usage after the latest
699
- * window boundary. Before that first usage lands, any context measurement is
700
- * still anchored to the previous window's last request.
701
- */
702
- export function hasAssistantUsageSinceWindowBoundary(
703
- entries: readonly SessionEntry[],
704
- sessionId?: string,
705
- ): boolean {
706
- const boundaryIndex = findLatestWindowBoundaryIndex(entries, sessionId);
707
- if (boundaryIndex < 0) return false;
708
- for (let index = boundaryIndex + 1; index < entries.length; index += 1) {
709
- const entry = entries[index]!;
710
- if (entry.type !== "message") continue;
711
- const message = entry.message;
712
- if (message.role !== "assistant") continue;
713
- const candidate = message as unknown as {
714
- stopReason?: string;
715
- usage?: { totalTokens?: number; input?: number; output?: number; cacheRead?: number; cacheWrite?: number };
716
- };
717
- if (candidate.stopReason === "aborted" || candidate.stopReason === "error") continue;
718
- const usage = candidate.usage;
719
- if (!usage) continue;
720
- const tokens = usage.totalTokens ?? (usage.input ?? 0) + (usage.output ?? 0) + (usage.cacheRead ?? 0) + (usage.cacheWrite ?? 0);
721
- if (tokens > 0) return true;
722
- }
723
- return false;
724
- }
725
-
726
- /**
727
- * Return the most recent successful notes write after the latest window
728
- * boundary. The receipt is derived from persisted session entries so it
729
- * survives restarts and can be used for an explicit rollover handoff.
730
- */
731
- export function findLatestNotesCheckpointSinceBoundary(
732
- entries: readonly SessionEntry[],
733
- sessionId?: string,
734
- afterEntryId?: string,
735
- ): NotesCheckpointReceipt | undefined {
736
- const windowBoundaryIndex = findLatestWindowBoundaryIndex(entries, sessionId);
737
- const afterIndex = afterEntryId === undefined ? -1 : entries.findIndex((entry) => entry.id === afterEntryId);
738
- if (windowBoundaryIndex < 0 || (afterEntryId !== undefined && afterIndex < 0)) return undefined;
739
- const boundaryIndex = Math.max(windowBoundaryIndex, afterIndex);
740
- const checkpointCalls = new Map<string, NotesCheckpointReceipt>();
741
- let latest: NotesCheckpointReceipt | undefined;
742
- for (let index = boundaryIndex + 1; index < entries.length; index += 1) {
743
- const entry = entries[index]!;
744
- if (entry.type !== "message") continue;
745
- const message = entry.message;
746
- if (message.role === "assistant") {
747
- const parts = Array.isArray(message.content) ? message.content : [];
748
- for (const part of parts) {
749
- if (!isRecord(part) || part.type !== "toolCall" || part.name !== "notes") continue;
750
- const args = isRecord(part.arguments) ? part.arguments : undefined;
751
- const action = typeof args?.action === "string" ? args.action : "";
752
- const path = isNonEmptyString(args?.path) ? args.path : undefined;
753
- if (isNonEmptyString(part.id) && path !== undefined && NOTES_CHECKPOINT_ACTIONS.has(action)) {
754
- checkpointCalls.set(part.id, { path, toolCallId: part.id });
755
- }
756
- }
757
- continue;
758
- }
759
- if (message.role !== "toolResult" || message.toolName !== "notes") continue;
760
- const receipt = checkpointCalls.get(message.toolCallId);
761
- if (!receipt) continue;
762
- checkpointCalls.delete(message.toolCallId);
763
- if (message.isError !== false || !isSuccessfulHistoryNotesToolResult(message.details)) continue;
764
- latest = receipt;
765
- }
766
- return latest;
767
- }
768
-
769
- /**
770
- * Whether the branch contains a successful notes append/write result after the
771
- * latest window boundary. Reads and failed writes never count as checkpoints.
772
- */
773
- export function findNotesCheckpointSinceBoundary(
774
- entries: readonly SessionEntry[],
775
- sessionId?: string,
776
- ): boolean {
777
- return findLatestNotesCheckpointSinceBoundary(entries, sessionId) !== undefined;
778
- }
779
-
780
- function findLatestWindowBoundaryIndex(
781
- entries: readonly SessionEntry[],
782
- sessionId?: string,
783
- ): number {
784
- for (let index = entries.length - 1; index >= 0; index -= 1) {
785
- const entry = entries[index]!;
786
- if (entry.type !== "custom_message" || entry.customType !== CODEX_CONTEXT_WINDOW_MESSAGE_TYPE) continue;
787
- if (!couldBelongToSession(entry.details, sessionId)) continue;
788
- if (!isCodexContextManagementMessageDetails(entry.details)) {
789
- throw new Error("Malformed persisted Codex context-window message");
790
- }
791
- if (!matchesSession(entry.details.sessionId, sessionId)) continue;
792
- if (entry.details.contextManagement.kind === "window") return index;
793
- }
794
- return -1;
795
- }
796
-
797
- export function findLatestWindowBoundaryEntry(
798
- entries: readonly SessionEntry[],
799
- sessionId?: string,
800
- ): WindowBoundaryEntry | undefined {
801
- for (let index = entries.length - 1; index >= 0; index -= 1) {
802
- const entry = entries[index]!;
803
- if (entry.type !== "custom_message" || entry.customType !== CODEX_CONTEXT_WINDOW_MESSAGE_TYPE) continue;
804
- if (!couldBelongToSession(entry.details, sessionId)) continue;
805
- if (!isCodexContextManagementMessageDetails(entry.details)) {
806
- throw new Error("Malformed persisted Codex context-window message");
807
- }
808
- if (!matchesSession(entry.details.sessionId, sessionId)) continue;
809
- if (entry.details.contextManagement.kind === "window") return entry as WindowBoundaryEntry;
810
- }
811
- return undefined;
812
- }
813
-
814
- /**
815
- * Whether this exact window boundary is already persisted for the session.
816
- * Only durable entries count; a queued/visible marker does not.
817
- */
818
- function hasPersistedWindowBoundary(
819
- entries: readonly SessionEntry[],
820
- sessionId: string | undefined,
821
- windowId: string,
822
- ): boolean {
823
- for (const entry of entries) {
824
- if (entry.type !== "custom_message" || entry.customType !== CODEX_CONTEXT_WINDOW_MESSAGE_TYPE) continue;
825
- if (!couldBelongToSession(entry.details, sessionId)) continue;
826
- if (!isCodexContextManagementMessageDetails(entry.details)) continue;
827
- if (!matchesSession(entry.details.sessionId, sessionId)) continue;
828
- const details = entry.details.contextManagement;
829
- if (details.kind === "window" && details.currentWindowId === windowId) return true;
830
- }
831
- return false;
832
- }
833
-
834
- function findLatestContextMarkerId(
835
- entries: readonly SessionEntry[],
836
- sessionId?: string,
837
- ): string | undefined {
838
- for (let index = entries.length - 1; index >= 0; index -= 1) {
839
- const entry = entries[index]!;
840
- if (entry.type !== "custom_message" || entry.customType !== CODEX_CONTEXT_WINDOW_MESSAGE_TYPE) continue;
841
- if (!couldBelongToSession(entry.details, sessionId)) continue;
842
- if (!isCodexContextManagementMessageDetails(entry.details)) {
843
- throw new Error("Malformed persisted Codex context-window message");
844
- }
845
- if (matchesSession(entry.details.sessionId, sessionId)) return entry.id;
846
- }
847
- return undefined;
848
- }
849
-
850
- function couldBelongToSession(details: unknown, sessionId: string | undefined): boolean {
851
- if (sessionId === undefined || !isRecord(details)) return true;
852
- const markerSessionId = details.sessionId;
853
- return !isNonEmptyString(markerSessionId) || markerSessionId === sessionId;
854
- }
855
-
856
- function matchesSession(markerSessionId: string | undefined, sessionId: string | undefined): boolean {
857
- return sessionId === undefined || markerSessionId === sessionId;
858
- }
859
-
860
- function estimateProjectedWindowTokens(messages: readonly AgentMessage[]): number {
861
- let lastUsageIndex = -1;
862
- let usageTokens = 0;
863
- for (let index = messages.length - 1; index >= 0; index -= 1) {
864
- const message = messages[index]!;
865
- if (message.role !== "assistant" || message.stopReason === "aborted" || message.stopReason === "error" || !message.usage) continue;
866
- const tokens = calculateContextTokens(message.usage);
867
- if (tokens <= 0) continue;
868
- lastUsageIndex = index;
869
- usageTokens = tokens;
870
- break;
871
- }
872
- if (lastUsageIndex < 0) return messages.reduce((total, message) => total + estimateProjectedMessageTokens(message), 0);
873
- return usageTokens + messages
874
- .slice(lastUsageIndex + 1)
875
- .reduce((total, message) => total + estimateProjectedMessageTokens(message), 0);
876
- }
877
-
878
- function estimateProjectedMessageTokens(message: AgentMessage): number {
879
- // A synthetic custom marker may be represented without content in a test or
880
- // by an older host. It contributes no measurable text; it must not make the
881
- // context hook fail closed merely because the diagnostic estimate is missing.
882
- if (
883
- !("content" in message) ||
884
- (typeof message.content !== "string" && !Array.isArray(message.content))
885
- ) return 0;
886
- return estimateTokens(message);
887
- }
888
-
889
- function projectEncryptedToolResults(messages: readonly AgentMessage[]): AgentMessage[] {
890
- let changed = false;
891
- const projected = messages.map((message) => {
892
- if (message.role !== "toolResult" || !isRecord(message.details)) return message;
893
- const historyNotes = message.details.codexHistoryNotes;
894
- if (!isRecord(historyNotes) || typeof historyNotes.encrypted_output !== "string") return message;
895
- changed = true;
896
- return {
897
- ...message,
898
- content: [
899
- { type: "text" as const, text: encodeEncryptedOutputForContext(historyNotes.encrypted_output) },
900
- ...message.content.filter((item) => item.type === "image"),
901
- ],
902
- };
903
- });
904
- return changed ? projected : [...messages];
905
- }
1
+ import { AsyncLocalStorage } from "node:async_hooks";
2
+ import { randomUUID } from "node:crypto";
3
+ import { isDeepStrictEqual } from "node:util";
4
+ import type { AgentMessage } from "@earendil-works/pi-agent-core";
5
+ import type { ProviderHeaders } from "@earendil-works/pi-ai";
6
+ import {
7
+ buildSessionProjection,
8
+ calculateContextTokens,
9
+ estimateTokens,
10
+ type CompactionResult,
11
+ type ExtensionAPI,
12
+ type ExtensionContext,
13
+ type SessionBeforeCompactEvent,
14
+ type SessionEntry,
15
+ type ProjectedSessionEntry,
16
+ } from "@earendil-works/pi-coding-agent";
17
+ import { ContextWindowBudget, type ContextRemaining } from "./window-budget";
18
+ import {
19
+ rewriteEncryptedToolOutputs,
20
+ rewriteWindowHeaders,
21
+ rewriteWindowPayload,
22
+ } from "./window-request";
23
+ import {
24
+ CONTEXT_WINDOW_COMPACTION_STRATEGY,
25
+ CONTEXT_WINDOW_COMPACTION_SUMMARY,
26
+ isContextWindowBoundary,
27
+ preserveSystemHead,
28
+ renderContextWindowMessage,
29
+ sendContextWindowMessage,
30
+ } from "./messages";
31
+ import {
32
+ CODEX_CONTEXT_WINDOW_MESSAGE_TYPE,
33
+ CONTEXT_MANAGEMENT_PROTOCOL,
34
+ type CodexContextManagementMessageDetails,
35
+ type ContextWindowCompactionDetails,
36
+ type ContextWindowIdentity,
37
+ isCodexContextManagementMessageDetails,
38
+ isNonEmptyString,
39
+ isRecord,
40
+ isContextWindowCompactionDetails,
41
+ type NotesCheckpointReceipt,
42
+ } from "./types";
43
+ import {
44
+ encodeEncryptedOutputForContext,
45
+ isSuccessfulHistoryNotesToolResult,
46
+ loadHistoryNotesThreadHint,
47
+ } from "./history-notes";
48
+ import { rewriteContextNamespaceTools } from "./namespace-tools";
49
+ import { hasInvalidatedCompactionContext } from "../payload-rewrite";
50
+
51
+ interface StartContextWindowOptions {
52
+ triggerTurn: boolean;
53
+ signal?: AbortSignal;
54
+ trimPreviousWindow: boolean;
55
+ /** Captured with the tool activation policy, before any awaited backend work. */
56
+ gatewayModels?: readonly string[];
57
+ /** Persist an owned handoff's exact target before the marker can be queued. */
58
+ beforeSchedule?: (identity: ContextWindowIdentity) => void;
59
+ }
60
+
61
+ type ThreadHintLoader = (
62
+ ctx: ExtensionContext,
63
+ signal?: AbortSignal,
64
+ gatewayModels?: readonly string[],
65
+ ) => Promise<string | undefined>;
66
+
67
+ type WindowBoundaryEntry = Extract<SessionEntry, { type: "custom_message" }> & {
68
+ details: CodexContextManagementMessageDetails;
69
+ };
70
+
71
+ /** Canonical input and real source anchors for one read-only manual compression. */
72
+ export type RemoteWindowCompactionInput = {
73
+ messages: AgentMessage[];
74
+ /** Verified canonical source contributions within this window, including edits. */
75
+ entries: ProjectedSessionEntry[];
76
+ firstKeptEntryId: string;
77
+ /** Only a newer compaction can own this window; pre-window checkpoints are retired. */
78
+ owningCompaction?: Extract<SessionEntry, { type: "compaction" }>;
79
+ /** Entries whose visible contributions Pi will retain alongside the new checkpoint. */
80
+ retainedMessages: AgentMessage[];
81
+ windowBoundaryIndex: number;
82
+ };
83
+
84
+ /**
85
+ * In-process duplicate guard for one scheduled rollover. It is anchored to the
86
+ * session that scheduled it and to the exact target window id, so it can never
87
+ * be satisfied by an older or foreign marker. It is not checkpoint evidence:
88
+ * only a persisted notes result unlocks the gate.
89
+ */
90
+ type PendingRollover = {
91
+ sessionId?: string;
92
+ targetWindowId: string;
93
+ /** The current request already contains this marker; never persistence proof. */
94
+ projected?: boolean;
95
+ };
96
+
97
+ export class CodexContextWindowManager {
98
+ private identity: ContextWindowIdentity | undefined;
99
+ private sessionId: string | undefined;
100
+ private restoredMarkerId: string | undefined;
101
+ private restoredCompactionId: string | undefined;
102
+ private boundary: WindowBoundaryResolution = { kind: "none" };
103
+ private readonly budget = new ContextWindowBudget();
104
+ private pendingRollover: PendingRollover | undefined;
105
+ private readonly projectionDiagnostics = new Set<string>();
106
+ private lastKnownSystemHead: AgentMessage | undefined;
107
+ private readonly compactionProjection = new AsyncLocalStorage<boolean>();
108
+
109
+ get isPreparingCompaction(): boolean { return this.compactionProjection.getStore() === true; }
110
+
111
+ /** Keep read-only preparation scoped across awaits without affecting other live requests. */
112
+ withCompactionProjection<T>(prepare: () => Promise<T>): Promise<T> {
113
+ return this.compactionProjection.run(true, prepare);
114
+ }
115
+ /**
116
+ * Pi's getContextUsage() measures the durable transcript. Remote windows are
117
+ * projected only by this extension, so remember the last provider-visible
118
+ * estimate for budget decisions and get_context_remaining.
119
+ */
120
+ private projectedContextTokens: number | undefined;
121
+ private readonly loadThreadHint: ThreadHintLoader;
122
+
123
+ constructor(loadThreadHint?: ThreadHintLoader) {
124
+ this.loadThreadHint = loadThreadHint ?? ((ctx, signal, gatewayModels) => loadHistoryNotesThreadHint(ctx, signal, gatewayModels));
125
+ }
126
+
127
+ reset(): void {
128
+ this.resetWindowState();
129
+ this.pendingRollover = undefined;
130
+ this.projectionDiagnostics.clear();
131
+ this.lastKnownSystemHead = undefined;
132
+ }
133
+
134
+ currentIdentity(): ContextWindowIdentity | undefined {
135
+ return this.identity ? { ...this.identity } : undefined;
136
+ }
137
+ /** Derived eligibility, never a replayed marker command. */
138
+ hasPendingTrim(): boolean {
139
+ return this.boundary.kind === "effective" && !this.boundary.cleaned &&
140
+ this.boundary.entry.details.contextManagement.trimPreviousWindow === true;
141
+ }
142
+
143
+ /** Historical identity does not own a newer native/opaque checkpoint. */
144
+ hasEffectiveWindow(): boolean {
145
+ if (!this.identity || this.boundary.kind === "ambiguous") return false;
146
+ return this.boundary.kind !== "superseded" ||
147
+ (this.pendingRollover?.projected === true && this.pendingRollover.targetWindowId === this.identity.currentWindowId);
148
+ }
149
+
150
+ hasSupersededWindow(): boolean {
151
+ return this.boundary.kind === "superseded";
152
+ }
153
+
154
+ private resetWindowState(): void {
155
+ this.identity = undefined;
156
+ this.sessionId = undefined;
157
+ this.restoredMarkerId = undefined;
158
+ this.restoredCompactionId = undefined;
159
+ this.boundary = { kind: "none" };
160
+ this.budget.reset();
161
+ this.projectedContextTokens = undefined;
162
+ }
163
+
164
+ restore(entries: readonly SessionEntry[], sessionId?: string): void {
165
+ // A rebuild of derived state must not drop an outstanding rollover guard;
166
+ // only a changed session or an observed target marker retires it below.
167
+ this.resetWindowState();
168
+ this.sessionId = sessionId;
169
+ for (const entry of entries) {
170
+ if (entry.type === "compaction") {
171
+ this.restoredCompactionId = entry.id;
172
+ continue;
173
+ }
174
+ if (entry.type !== "custom_message" || entry.customType !== CODEX_CONTEXT_WINDOW_MESSAGE_TYPE) continue;
175
+ if (!couldBelongToSession(entry.details, sessionId)) continue;
176
+ if (!isCodexContextManagementMessageDetails(entry.details)) {
177
+ throw new Error("Malformed persisted Codex context-window message");
178
+ }
179
+ if (!matchesSession(entry.details.sessionId, sessionId)) continue;
180
+ const details = entry.details.contextManagement;
181
+ this.restoredMarkerId = entry.id;
182
+ if (details.kind === "window") {
183
+ this.identity = identityFromDetails(entry.details);
184
+ }
185
+ this.budget.restore(details.kind, details.currentWindowId);
186
+ }
187
+ this.boundary = resolveWindowBoundary(entries, sessionId);
188
+ this.retireSatisfiedOrStalePending(entries, sessionId);
189
+ }
190
+
191
+ /** Rebuild state when Pi navigates to a different session/branch. */
192
+ synchronize(ctx: Pick<ExtensionContext, "sessionManager">): void {
193
+ const entries = ctx.sessionManager.getBranch();
194
+ const sessionId = ctx.sessionManager.getSessionId();
195
+ this.retireSatisfiedOrStalePending(entries, sessionId);
196
+ const latestMarkerId = findLatestContextMarkerId(entries, sessionId);
197
+ let latestCompactionId: string | undefined;
198
+ for (let index = entries.length - 1; index >= 0; index--) {
199
+ if (entries[index]!.type === "compaction") {
200
+ latestCompactionId = entries[index]!.id;
201
+ break;
202
+ }
203
+ }
204
+ // Hook-provided compactions may have no success callback. The persisted
205
+ // branch also detects navigation before/after a commit with the same marker.
206
+ if (this.sessionId !== sessionId || this.restoredMarkerId !== latestMarkerId || this.restoredCompactionId !== latestCompactionId) {
207
+ this.restore(entries, sessionId);
208
+ }
209
+ // Edits and sibling navigation may leave marker/compaction IDs unchanged.
210
+ this.boundary = resolveWindowBoundary(entries, sessionId);
211
+ }
212
+
213
+ /**
214
+ * Drop a pending rollover only for durable reasons: the target window marker
215
+ * is now persisted in this session, or the session changed. A queued marker
216
+ * visible in an in-memory message list is not durable evidence.
217
+ */
218
+ private retireSatisfiedOrStalePending(
219
+ entries: readonly SessionEntry[],
220
+ sessionId: string | undefined,
221
+ ): void {
222
+ const pending = this.pendingRollover;
223
+ if (!pending) return;
224
+ if (pending.sessionId !== undefined && pending.sessionId !== sessionId) {
225
+ this.pendingRollover = undefined;
226
+ return;
227
+ }
228
+ if (hasPersistedWindowBoundary(entries, sessionId, pending.targetWindowId)) {
229
+ this.pendingRollover = undefined;
230
+ }
231
+ }
232
+
233
+ ensureInitialized(pi: ExtensionAPI, ctx: ExtensionContext, active: boolean): void {
234
+ if (!active) return;
235
+ this.synchronize(ctx);
236
+ if (this.identity) return;
237
+ const windowId = randomUUID();
238
+ this.sendWindowMessage(
239
+ pi,
240
+ ctx,
241
+ { firstWindowId: windowId, currentWindowId: windowId, windowNumber: 0 },
242
+ { triggerTurn: false, trimPreviousWindow: false },
243
+ );
244
+ }
245
+
246
+ project(
247
+ messages: readonly AgentMessage[],
248
+ mode: "off" | "local" | "remote",
249
+ /** Delay encoding until verified retained copies have been removed. */
250
+ encodeEncryptedOutputs = true,
251
+ ): AgentMessage[] {
252
+ if (this.isPreparingCompaction) return this.projectCompactionContext(messages);
253
+ // Recompute the request-only preview on every projection. Observing a
254
+ // queued marker must not authorize maintenance or retire the duplicate guard.
255
+ if (this.pendingRollover) this.pendingRollover.projected = false;
256
+ if (mode === "off") {
257
+ this.projectedContextTokens = undefined;
258
+ return messages.filter(
259
+ (message) => message.role !== "custom" || message.customType !== CODEX_CONTEXT_WINDOW_MESSAGE_TYPE,
260
+ );
261
+ }
262
+ if (this.boundary.kind === "ambiguous") {
263
+ throw new Error("Persisted Codex context-window boundary missing or edited in canonical projection");
264
+ }
265
+ if (this.boundary.kind === "superseded") this.identity = identityFromDetails(this.boundary.entry.details);
266
+ let boundaryIndex = -1;
267
+ for (let index = 0; index < messages.length; index += 1) {
268
+ const message = messages[index]!;
269
+ if (
270
+ message.role === "custom" &&
271
+ message.customType === CODEX_CONTEXT_WINDOW_MESSAGE_TYPE
272
+ ) {
273
+ if (!couldBelongToSession(message.details, this.sessionId)) continue;
274
+ if (!isCodexContextManagementMessageDetails(message.details)) {
275
+ throw new Error("Malformed persisted Codex context-window message");
276
+ }
277
+ if (!matchesSession(message.details.sessionId, this.sessionId)) continue;
278
+ }
279
+ if (!isContextWindowBoundary(message)) continue;
280
+ const queuedRollover = mode === "remote" && this.pendingRollover?.targetWindowId === message.details.contextManagement.currentWindowId;
281
+ if (this.boundary.kind === "superseded" && !queuedRollover) continue;
282
+ // Local consumers follow the synchronized persisted identity, not a
283
+ // queued marker from the request view. They never own the Remote session.
284
+ if (mode === "local" && message.details.contextManagement.currentWindowId !== this.identity?.currentWindowId) continue;
285
+ boundaryIndex = index;
286
+ if (mode === "remote") {
287
+ this.identity = identityFromDetails(message.details);
288
+ if (this.pendingRollover) this.pendingRollover.projected = queuedRollover;
289
+ }
290
+ }
291
+ if (boundaryIndex < 0 && this.boundary.kind === "effective") {
292
+ throw new Error("Persisted Codex context-window boundary missing from effective projection");
293
+ }
294
+ // Projection observes the request view, never durable session state. A
295
+ // queued rollover marker must not clear the duplicate guard here.
296
+ const trimmed = boundaryIndex < 0 ? [...messages] : this.trimToWindow(messages, boundaryIndex, mode === "local");
297
+ const projected = mode === "remote" ? (encodeEncryptedOutputs ? projectEncryptedToolResults(trimmed) : trimmed) : trimmed.filter(
298
+ (message) => message.role !== "custom" || message.customType !== CODEX_CONTEXT_WINDOW_MESSAGE_TYPE,
299
+ );
300
+ this.projectedContextTokens = mode === "remote" && this.identity
301
+ ? estimateProjectedWindowTokens(projected)
302
+ : undefined;
303
+ return projected;
304
+ }
305
+
306
+ /**
307
+ * Strict, read-only projection. Unlike live Remote projection, an unsafe head
308
+ * never falls back to the full transcript, updates identity, or records budget.
309
+ */
310
+ projectCompactionContext(messages: readonly AgentMessage[], encodeEncryptedOutputs = false): AgentMessage[] {
311
+ if (this.pendingRollover || (this.boundary.kind !== "effective" && this.boundary.kind !== "superseded")) {
312
+ throw new Error("unsafe-compaction-window-boundary");
313
+ }
314
+ let boundaryIndex = 0;
315
+ if (this.boundary.kind === "effective") {
316
+ const boundary = this.boundary.entry;
317
+ boundaryIndex = messages.findIndex((message) => isContextWindowBoundary(message)
318
+ && isDeepStrictEqual(message.details, boundary.details)
319
+ && isDeepStrictEqual(message.content, boundary.content));
320
+ if (boundaryIndex < 0) throw new Error("compaction-window-boundary-missing");
321
+ }
322
+ const repair = preserveSystemHead({ messages, boundaryIndex });
323
+ if (!repair.head || repair.shrinkRejected) throw new Error("compaction-system-tool-head-lost");
324
+ return encodeEncryptedOutputs ? projectEncryptedToolResults(repair.messages) : [...repair.messages];
325
+ }
326
+
327
+ prepareRemoteCompaction(event: SessionBeforeCompactEvent, ctx: ExtensionContext): RemoteWindowCompactionInput {
328
+ this.synchronize(ctx);
329
+ if (this.hasPendingRollover(ctx) || (this.boundary.kind !== "effective" && this.boundary.kind !== "superseded")) {
330
+ throw new Error("unsafe-compaction-window-boundary");
331
+ }
332
+ const branch = ctx.sessionManager.getBranch();
333
+ if (!isDeepStrictEqual(event.branchEntries, branch)) throw new Error("compaction-branch-changed");
334
+ const projection = ctx.sessionManager.buildSessionProjection();
335
+ // Never manufacture source IDs or substitute a raw-history reconstruction.
336
+ if (!projection || !Array.isArray(projection.entries) || !Array.isArray(projection.messages)
337
+ || !isDeepStrictEqual(projection.messages, projection.entries.flatMap((entry) => entry.messages))
338
+ || new Set(projection.entries.map((entry) => entry.sourceEntry.id)).size !== projection.entries.length
339
+ || projection.entries.some((entry) => !branch.some((raw) => raw.id === entry.sourceEntry.id && isDeepStrictEqual(raw, entry.sourceEntry)))) {
340
+ throw new Error("compaction-source-projection-unavailable");
341
+ }
342
+ const windowBoundaryIndex = branch.findIndex((entry) => entry.id === this.boundaryEntryId());
343
+ if (windowBoundaryIndex < 0) throw new Error("compaction-window-boundary-missing");
344
+ const latest = [...branch].reverse().find((entry) => entry.type === "compaction");
345
+ if (latest?.type === "compaction" && branch.findIndex((entry) => entry.id === latest.id) > windowBoundaryIndex
346
+ && (isContextWindowCompactionDetails(latest.details) || latest.summary === CONTEXT_WINDOW_COMPACTION_SUMMARY)) {
347
+ // A malformed cleanup cannot masquerade as a new native/text owner.
348
+ // Details alone do not turn an arbitrary summary into Toolkit cleanup.
349
+ if (!latest || latest.type !== "compaction" || latest.fromHook !== true
350
+ || latest.summary !== CONTEXT_WINDOW_COMPACTION_SUMMARY
351
+ || !isContextWindowCompactionDetails(latest.details)
352
+ || latest.details.windowId !== this.boundary.entry.details.contextManagement.currentWindowId
353
+ || latest.firstKeptEntryId !== this.boundary.entry.id) throw new Error("unverified-window-cleanup");
354
+ }
355
+ const owningCompaction = this.boundary.kind === "superseded" && latest?.type === "compaction" ? latest : undefined;
356
+ if (this.boundary.kind === "superseded" && !owningCompaction) throw new Error("compaction-owner-missing");
357
+ const lowerId = owningCompaction?.id ?? this.boundary.entry.id;
358
+ const lowerIndex = branch.findIndex((entry) => entry.id === lowerId);
359
+ const requestedIndex = branch.findIndex((entry) => entry.id === event.preparation.firstKeptEntryId);
360
+ if (requestedIndex < 0 || lowerIndex < windowBoundaryIndex) throw new Error("unsafe-compaction-retained-start");
361
+ const firstKeptEntryId = requestedIndex < lowerIndex ? lowerId : event.preparation.firstKeptEntryId;
362
+ const keptIndex = Math.max(requestedIndex, lowerIndex);
363
+ const safeSources = new Set(projection.entries.filter((entry) => entry.messages.length > 0).map((entry) => entry.sourceEntry.id));
364
+ if (!safeSources.has(firstKeptEntryId)) throw new Error("unsafe-compaction-retained-start");
365
+ const retainedIds = new Set(branch.slice(keptIndex).map((entry) => entry.id));
366
+ return {
367
+ messages: this.projectCompactionContext(structuredClone(projection.messages)),
368
+ entries: structuredClone(projection.entries.filter((entry) =>
369
+ branch.findIndex((raw) => raw.id === entry.sourceEntry.id) >= windowBoundaryIndex)),
370
+ firstKeptEntryId,
371
+ owningCompaction,
372
+ retainedMessages: structuredClone(projection.entries.filter((entry) => retainedIds.has(entry.sourceEntry.id)
373
+ && entry.sourceEntry.type !== "compaction").flatMap((entry) => entry.messages)),
374
+ windowBoundaryIndex,
375
+ };
376
+ }
377
+
378
+ private boundaryEntryId(): string | undefined {
379
+ return this.boundary.kind === "effective" || this.boundary.kind === "superseded" ? this.boundary.entry.id : undefined;
380
+ }
381
+
382
+ /**
383
+ * Trim everything before the window boundary while keeping Pi's prompt/tool head.
384
+ *
385
+ * When the head cannot be preserved the trim is refused and the full message list is
386
+ * sent, because a window without tools silently degrades into a chat the model can
387
+ * only narrate. Each refusal is reported once per window and reason.
388
+ */
389
+ private trimToWindow(messages: readonly AgentMessage[], boundaryIndex: number, local = false): AgentMessage[] {
390
+ const repair = preserveSystemHead({ messages, boundaryIndex, fallbackHead: this.lastKnownSystemHead });
391
+ if (repair.head) this.lastKnownSystemHead = repair.head;
392
+ if (repair.shrinkRejected) {
393
+ // Remote requests retain the existing tool-preservation fallback. A
394
+ // local request must never recover tools by leaking retired windows.
395
+ if (local) throw new Error("Cannot preserve the system/tool head in local context-window projection");
396
+ this.recordProjectionDiagnostic(repair.lostToolNames);
397
+ }
398
+ return [...repair.messages];
399
+ }
400
+
401
+ private recordProjectionDiagnostic(lostToolNames: readonly string[]): void {
402
+ const key = `${this.identity?.currentWindowId ?? "unknown"}|${[...lostToolNames].sort().join(",")}`;
403
+ this.projectionDiagnostics.add(key);
404
+ }
405
+
406
+ /**
407
+ * Drain the pending `tool-loadout-shrink` reports. Returns one entry per window and
408
+ * lost-tool set, so a caller can surface the incident without flooding the UI.
409
+ */
410
+ takeProjectionDiagnostics(): string[] {
411
+ if (this.projectionDiagnostics.size === 0) return [];
412
+ const pending = [...this.projectionDiagnostics];
413
+ this.projectionDiagnostics.clear();
414
+ return pending;
415
+ }
416
+
417
+ async startNewWindow(
418
+ pi: ExtensionAPI,
419
+ ctx: ExtensionContext,
420
+ options: StartContextWindowOptions,
421
+ ): Promise<boolean> {
422
+ this.synchronize(ctx);
423
+ // A same-session rollover whose target marker is not persisted yet is
424
+ // still in flight: report it instead of scheduling a second window.
425
+ if (this.hasPendingRollover(ctx)) return false;
426
+ if (options.signal?.aborted) throw new Error("Remote context rollover was aborted");
427
+ const current = this.identity;
428
+ const liveSessionId = ctx.sessionManager.getSessionId();
429
+ const currentWindowId = randomUUID();
430
+ const next: ContextWindowIdentity = current
431
+ ? {
432
+ firstWindowId: current.firstWindowId,
433
+ currentWindowId,
434
+ previousWindowId: current.currentWindowId,
435
+ windowNumber: current.windowNumber + 1,
436
+ }
437
+ : { firstWindowId: currentWindowId, currentWindowId, windowNumber: 0 };
438
+ this.pendingRollover = { sessionId: liveSessionId, targetWindowId: currentWindowId };
439
+ try {
440
+ const checkpoint = current
441
+ ? findLatestNotesCheckpointSinceBoundary(ctx.sessionManager.getBranch(), liveSessionId)
442
+ : undefined;
443
+ let threadHint: string | undefined;
444
+ if (current) {
445
+ try {
446
+ threadHint = await this.loadThreadHint(ctx, options.signal, options.gatewayModels);
447
+ } catch (error) {
448
+ if (options.signal?.aborted) throw error;
449
+ }
450
+ }
451
+ if (options.signal?.aborted) throw new Error("Remote context rollover was aborted");
452
+ options.beforeSchedule?.(next);
453
+ this.sendWindowMessage(pi, ctx, next, options, threadHint, checkpoint);
454
+ return true;
455
+ } catch (error) {
456
+ this.pendingRollover = undefined;
457
+ throw error;
458
+ }
459
+ }
460
+
461
+ recordBudget(
462
+ pi: ExtensionAPI,
463
+ ctx: ExtensionContext,
464
+ active: boolean,
465
+ contextReminderThresholdPercent: number,
466
+ contextTokens?: number,
467
+ ): void {
468
+ if (!active || !this.identity || contextReminderThresholdPercent <= 0) return;
469
+ // Right after a rollover, Pi's usage anchor still reports the previous
470
+ // window's last request until the new window's first request completes.
471
+ // Budget decisions taken in that gap act on stale numbers: they burn the
472
+ // once-per-window reminder on a false alarm seconds after a successful
473
+ // rollover, leaving the window silent for the rest of its life.
474
+ if (!hasAssistantUsageSinceWindowBoundary(ctx.sessionManager.getBranch(), this.sessionId)) return;
475
+ const reminder = this.budget.record(
476
+ ctx,
477
+ this.identity,
478
+ contextTokens ?? this.projectedContextTokens,
479
+ contextReminderThresholdPercent,
480
+ );
481
+ if (!reminder) return;
482
+ sendContextWindowMessage(
483
+ pi,
484
+ reminder.content,
485
+ reminder.kind,
486
+ this.identity,
487
+ { triggerTurn: reminder.kind === "fallback", sessionId: ctx.sessionManager.getSessionId() },
488
+ );
489
+ }
490
+
491
+ remaining(ctx: ExtensionContext, contextTokens?: number): ContextRemaining {
492
+ return this.budget.remaining(ctx, this.identity, contextTokens ?? this.projectedContextTokens);
493
+ }
494
+
495
+ /**
496
+ * True when a notes checkpoint succeeded in the current window (after the
497
+ * latest boundary). The scan always uses the live session id: a cached id can
498
+ * belong to a session Pi already navigated away from, which would filter the
499
+ * current window's marker and result out as foreign.
500
+ */
501
+ hasNotesCheckpointSinceBoundary(
502
+ ctx: Pick<ExtensionContext, "sessionManager">,
503
+ ): boolean {
504
+ return findNotesCheckpointSinceBoundary(
505
+ ctx.sessionManager.getBranch(),
506
+ ctx.sessionManager.getSessionId(),
507
+ );
508
+ }
509
+
510
+ /**
511
+ * True while a rollover scheduled for the current session is awaiting its
512
+ * persisted target marker. Call `synchronize()` first so a marker that has
513
+ * already been persisted or a session switch retires the guard.
514
+ */
515
+ hasPendingRollover(ctx: Pick<ExtensionContext, "sessionManager">): boolean {
516
+ const pending = this.pendingRollover;
517
+ if (!pending) return false;
518
+ return pending.sessionId === undefined || pending.sessionId === ctx.sessionManager.getSessionId();
519
+ }
520
+
521
+ /**
522
+ * Whether a rollover may be scheduled from the current window.
523
+ *
524
+ * A window that was entered by a context switch has to do real work before it switches
525
+ * again: without this guard a model can checkpoint and roll over on its very first turn,
526
+ * which is how long sessions start looping between notes and new_context.
527
+ */
528
+ canRolloverFromCurrentWindow(ctx: ExtensionContext): boolean {
529
+ if (!this.identity || this.identity.windowNumber === 0) return true;
530
+ // An exhausted window must always be able to escape: the budget fallback tells the
531
+ // model to checkpoint and switch, and refusing that would strand the session. An
532
+ // unmeasurable window keeps the guard, since nothing there is pushing for a switch.
533
+ const { remainingTokens } = this.remaining(ctx);
534
+ if (remainingTokens !== undefined && remainingTokens <= 0) return true;
535
+ return hasSubstantiveToolResultSinceBoundary(ctx.sessionManager.getBranch(), this.sessionId);
536
+ }
537
+
538
+ prepareCompaction(
539
+ event: SessionBeforeCompactEvent,
540
+ ): { cancel: true } | { compaction: CompactionResult<ContextWindowCompactionDetails> } {
541
+ const boundary = resolveWindowBoundary(event.branchEntries, this.sessionId);
542
+ if (boundary.kind !== "effective" || boundary.cleaned || !boundary.entry.details.contextManagement.trimPreviousWindow) {
543
+ return { cancel: true };
544
+ }
545
+ // An unpersisted/cancelled proposal changes no eligibility.
546
+ return { compaction: this.createCompaction(event, boundary.entry) };
547
+ }
548
+
549
+ /**
550
+ * Detach a local consumer at the durable Remote boundary, without summarizing
551
+ * any Remote-managed history. Unlike a scheduled rollover trim, this needs
552
+ * no pending trim and remains retryable until Pi persists the result.
553
+ *
554
+ * A consumed boundary is not an error. Once Pi has persisted the no-summary
555
+ * boundary, `buildSessionProjection()` already starts at that marker, so
556
+ * retired Remote history is outside every producer's input and the ordinary
557
+ * compaction chain may summarize the remaining local window. Returning
558
+ * `undefined` hands the event back to that chain; returning `{ cancel: true }`
559
+ * would strand the session with no working compaction path at all.
560
+ */
561
+ prepareDetachedCompaction(
562
+ event: SessionBeforeCompactEvent,
563
+ ): { compaction: CompactionResult<ContextWindowCompactionDetails> } | undefined {
564
+ const boundary = resolveWindowBoundary(event.branchEntries, this.sessionId);
565
+ if (boundary.kind === "ambiguous") throw new Error("Persisted Codex context-window boundary missing or edited");
566
+ if (boundary.kind !== "effective") return undefined;
567
+ // A matching no-summary boundary is already the effective compaction anchor.
568
+ // Nothing is left to detach, and writing a second boundary for the same window
569
+ // would move the retained start backwards.
570
+ if (boundary.cleaned) return undefined;
571
+ const compaction = this.createCompaction(event, boundary.entry);
572
+ // Pi's preparation/tokensBefore describes the durable branch, which includes
573
+ // Remote-managed history. Detached local accounting must use the same current
574
+ // window projection that the non-window provider sees instead.
575
+ return {
576
+ compaction: {
577
+ ...compaction,
578
+ tokensBefore: this.estimateLocalContextTokens(event.branchEntries),
579
+ },
580
+ };
581
+ }
582
+
583
+ private estimateLocalContextTokens(entries: readonly SessionEntry[]): number {
584
+ const messages = buildSessionProjection([...entries]).messages;
585
+ return this.project(messages, "local").reduce((total, message) => total + estimateTokens(message as never), 0);
586
+ }
587
+
588
+ private createCompaction(event: SessionBeforeCompactEvent, boundary: WindowBoundaryEntry): CompactionResult<ContextWindowCompactionDetails> {
589
+ return {
590
+ summary: CONTEXT_WINDOW_COMPACTION_SUMMARY,
591
+ firstKeptEntryId: boundary.id,
592
+ tokensBefore: event.preparation.tokensBefore,
593
+ details: {
594
+ protocol: CONTEXT_MANAGEMENT_PROTOCOL,
595
+ strategy: CONTEXT_WINDOW_COMPACTION_STRATEGY,
596
+ windowId: boundary.details.contextManagement.currentWindowId,
597
+ },
598
+ };
599
+ }
600
+
601
+ rewritePayload(payload: unknown, ctx: ExtensionContext): unknown {
602
+ const withMetadata = rewriteWindowPayload(payload, ctx, this.identity);
603
+ return rewriteContextNamespaceTools(withMetadata, { encrypted: true });
604
+ }
605
+
606
+ rewriteHeaders(headers: ProviderHeaders, ctx: ExtensionContext): void {
607
+ rewriteWindowHeaders(headers, ctx, this.identity);
608
+ }
609
+
610
+ private sendWindowMessage(
611
+ pi: ExtensionAPI,
612
+ ctx: ExtensionContext,
613
+ identity: ContextWindowIdentity,
614
+ options: StartContextWindowOptions,
615
+ threadHint?: string,
616
+ checkpoint?: NotesCheckpointReceipt,
617
+ ): void {
618
+ sendContextWindowMessage(
619
+ pi,
620
+ renderContextWindowMessage(identity, threadHint, checkpoint),
621
+ "window",
622
+ identity,
623
+ { triggerTurn: options.triggerTurn, sessionId: ctx.sessionManager.getSessionId() },
624
+ options.trimPreviousWindow,
625
+ );
626
+ this.identity = identity;
627
+ this.sessionId = ctx.sessionManager.getSessionId();
628
+ this.restoredMarkerId = undefined;
629
+ // sendMessage() acceptance is not persistence. The pending rollover stays
630
+ // armed until synchronize() observes this target window in the persisted
631
+ // branch (or the session changes).
632
+ }
633
+ }
634
+
635
+ type WindowBoundaryResolution =
636
+ | { kind: "none" | "ambiguous" }
637
+ | { kind: "superseded"; entry: WindowBoundaryEntry }
638
+ | { kind: "effective"; entry: WindowBoundaryEntry; cleaned: boolean };
639
+
640
+ /** Pi's projection owns context; raw chronology only establishes window ownership. */
641
+ function resolveWindowBoundary(entries: readonly SessionEntry[], sessionId?: string): WindowBoundaryResolution {
642
+ const boundary = findLatestWindowBoundaryEntry(entries, sessionId);
643
+ if (!boundary) return { kind: "none" };
644
+ const projection = buildSessionProjection([...entries]);
645
+ const latest = projection.entries.find((entry) => entry.sourceEntry.type === "compaction")?.sourceEntry;
646
+ const boundaryIndex = entries.findIndex((entry) => entry.id === boundary.id);
647
+ const laterCompaction = latest?.type === "compaction" && entries.findIndex((entry) => entry.id === latest.id) > boundaryIndex
648
+ ? latest : undefined;
649
+ if (laterCompaction && hasInvalidatedCompactionContext([...entries], laterCompaction.id)) return { kind: "ambiguous" };
650
+ const marker = projection.entries.find((entry) => entry.sourceEntry.id === boundary.id);
651
+ // A surviving edit can replace content without changing custom details.
652
+ const edited = entries.some((entry) => entry.type === "context_edit" && entry.targetId === boundary.id &&
653
+ projection.entries.some((projected) => projected.sourceEntry.id === entry.id));
654
+ if (edited) return { kind: "ambiguous" };
655
+ const cleaned = laterCompaction !== undefined &&
656
+ isContextWindowCompactionDetails(laterCompaction.details) &&
657
+ laterCompaction.details.windowId === boundary.details.contextManagement.currentWindowId &&
658
+ laterCompaction.firstKeptEntryId === boundary.id;
659
+ if (laterCompaction && !cleaned) return { kind: "superseded", entry: boundary };
660
+ if (!marker?.messages.some(isContextWindowBoundary)) return { kind: "ambiguous" };
661
+ return { kind: "effective", entry: boundary, cleaned };
662
+ }
663
+
664
+ function identityFromDetails(details: CodexContextManagementMessageDetails): ContextWindowIdentity {
665
+ const context = details.contextManagement;
666
+ return {
667
+ firstWindowId: context.firstWindowId,
668
+ currentWindowId: context.currentWindowId,
669
+ ...(context.previousWindowId ? { previousWindowId: context.previousWindowId } : {}),
670
+ windowNumber: context.windowNumber,
671
+ };
672
+ }
673
+
674
+ const NOTES_CHECKPOINT_ACTIONS: ReadonlySet<string> = new Set(["append_to_file", "write_file"]);
675
+
676
+ /** Tools that only move context around; they never count as work done in a window. */
677
+ const CONTEXT_MANAGEMENT_TOOL_NAMES: ReadonlySet<string> = new Set([
678
+ "new_context",
679
+ "get_context_remaining",
680
+ "history",
681
+ "notes",
682
+ ]);
683
+
684
+ /**
685
+ * Whether the current window produced a tool result from anything other than the
686
+ * context-handoff tools. Derived from persisted entries so it survives restarts.
687
+ */
688
+ export function hasSubstantiveToolResultSinceBoundary(
689
+ entries: readonly SessionEntry[],
690
+ sessionId?: string,
691
+ ): boolean {
692
+ const boundaryIndex = findLatestWindowBoundaryIndex(entries, sessionId);
693
+ if (boundaryIndex < 0) return true;
694
+ for (let index = boundaryIndex + 1; index < entries.length; index += 1) {
695
+ const entry = entries[index]!;
696
+ if (entry.type !== "message") continue;
697
+ const message = entry.message as unknown as { role?: string; toolName?: string };
698
+ if (message.role !== "toolResult" || typeof message.toolName !== "string") continue;
699
+ if (!CONTEXT_MANAGEMENT_TOOL_NAMES.has(message.toolName)) return true;
700
+ }
701
+ return false;
702
+ }
703
+
704
+ /**
705
+ * Whether the branch contains a successful assistant usage after the latest
706
+ * window boundary. Before that first usage lands, any context measurement is
707
+ * still anchored to the previous window's last request.
708
+ */
709
+ export function hasAssistantUsageSinceWindowBoundary(
710
+ entries: readonly SessionEntry[],
711
+ sessionId?: string,
712
+ ): boolean {
713
+ const boundaryIndex = findLatestWindowBoundaryIndex(entries, sessionId);
714
+ if (boundaryIndex < 0) return false;
715
+ for (let index = boundaryIndex + 1; index < entries.length; index += 1) {
716
+ const entry = entries[index]!;
717
+ if (entry.type !== "message") continue;
718
+ const message = entry.message;
719
+ if (message.role !== "assistant") continue;
720
+ const candidate = message as unknown as {
721
+ stopReason?: string;
722
+ usage?: { totalTokens?: number; input?: number; output?: number; cacheRead?: number; cacheWrite?: number };
723
+ };
724
+ if (candidate.stopReason === "aborted" || candidate.stopReason === "error") continue;
725
+ const usage = candidate.usage;
726
+ if (!usage) continue;
727
+ const tokens = usage.totalTokens ?? (usage.input ?? 0) + (usage.output ?? 0) + (usage.cacheRead ?? 0) + (usage.cacheWrite ?? 0);
728
+ if (tokens > 0) return true;
729
+ }
730
+ return false;
731
+ }
732
+
733
+ /**
734
+ * Return the most recent successful notes write after the latest window
735
+ * boundary. The receipt is derived from persisted session entries so it
736
+ * survives restarts and can be used for an explicit rollover handoff.
737
+ */
738
+ export function findLatestNotesCheckpointSinceBoundary(
739
+ entries: readonly SessionEntry[],
740
+ sessionId?: string,
741
+ afterEntryId?: string,
742
+ ): NotesCheckpointReceipt | undefined {
743
+ const windowBoundaryIndex = findLatestWindowBoundaryIndex(entries, sessionId);
744
+ const afterIndex = afterEntryId === undefined ? -1 : entries.findIndex((entry) => entry.id === afterEntryId);
745
+ if (windowBoundaryIndex < 0 || (afterEntryId !== undefined && afterIndex < 0)) return undefined;
746
+ const boundaryIndex = Math.max(windowBoundaryIndex, afterIndex);
747
+ const checkpointCalls = new Map<string, NotesCheckpointReceipt>();
748
+ let latest: NotesCheckpointReceipt | undefined;
749
+ for (let index = boundaryIndex + 1; index < entries.length; index += 1) {
750
+ const entry = entries[index]!;
751
+ if (entry.type !== "message") continue;
752
+ const message = entry.message;
753
+ if (message.role === "assistant") {
754
+ const parts = Array.isArray(message.content) ? message.content : [];
755
+ for (const part of parts) {
756
+ if (!isRecord(part) || part.type !== "toolCall" || part.name !== "notes") continue;
757
+ const args = isRecord(part.arguments) ? part.arguments : undefined;
758
+ const action = typeof args?.action === "string" ? args.action : "";
759
+ const path = isNonEmptyString(args?.path) ? args.path : undefined;
760
+ if (isNonEmptyString(part.id) && path !== undefined && NOTES_CHECKPOINT_ACTIONS.has(action)) {
761
+ checkpointCalls.set(part.id, { path, toolCallId: part.id });
762
+ }
763
+ }
764
+ continue;
765
+ }
766
+ if (message.role !== "toolResult" || message.toolName !== "notes") continue;
767
+ const receipt = checkpointCalls.get(message.toolCallId);
768
+ if (!receipt) continue;
769
+ checkpointCalls.delete(message.toolCallId);
770
+ if (message.isError !== false || !isSuccessfulHistoryNotesToolResult(message.details)) continue;
771
+ latest = receipt;
772
+ }
773
+ return latest;
774
+ }
775
+
776
+ /**
777
+ * Whether the branch contains a successful notes append/write result after the
778
+ * latest window boundary. Reads and failed writes never count as checkpoints.
779
+ */
780
+ export function findNotesCheckpointSinceBoundary(
781
+ entries: readonly SessionEntry[],
782
+ sessionId?: string,
783
+ ): boolean {
784
+ return findLatestNotesCheckpointSinceBoundary(entries, sessionId) !== undefined;
785
+ }
786
+
787
+ function findLatestWindowBoundaryIndex(
788
+ entries: readonly SessionEntry[],
789
+ sessionId?: string,
790
+ ): number {
791
+ for (let index = entries.length - 1; index >= 0; index -= 1) {
792
+ const entry = entries[index]!;
793
+ if (entry.type !== "custom_message" || entry.customType !== CODEX_CONTEXT_WINDOW_MESSAGE_TYPE) continue;
794
+ if (!couldBelongToSession(entry.details, sessionId)) continue;
795
+ if (!isCodexContextManagementMessageDetails(entry.details)) {
796
+ throw new Error("Malformed persisted Codex context-window message");
797
+ }
798
+ if (!matchesSession(entry.details.sessionId, sessionId)) continue;
799
+ if (entry.details.contextManagement.kind === "window") return index;
800
+ }
801
+ return -1;
802
+ }
803
+
804
+ export function findLatestWindowBoundaryEntry(
805
+ entries: readonly SessionEntry[],
806
+ sessionId?: string,
807
+ ): WindowBoundaryEntry | undefined {
808
+ for (let index = entries.length - 1; index >= 0; index -= 1) {
809
+ const entry = entries[index]!;
810
+ if (entry.type !== "custom_message" || entry.customType !== CODEX_CONTEXT_WINDOW_MESSAGE_TYPE) continue;
811
+ if (!couldBelongToSession(entry.details, sessionId)) continue;
812
+ if (!isCodexContextManagementMessageDetails(entry.details)) {
813
+ throw new Error("Malformed persisted Codex context-window message");
814
+ }
815
+ if (!matchesSession(entry.details.sessionId, sessionId)) continue;
816
+ if (entry.details.contextManagement.kind === "window") return entry as WindowBoundaryEntry;
817
+ }
818
+ return undefined;
819
+ }
820
+
821
+ /**
822
+ * Whether this exact window boundary is already persisted for the session.
823
+ * Only durable entries count; a queued/visible marker does not.
824
+ */
825
+ function hasPersistedWindowBoundary(
826
+ entries: readonly SessionEntry[],
827
+ sessionId: string | undefined,
828
+ windowId: string,
829
+ ): boolean {
830
+ for (const entry of entries) {
831
+ if (entry.type !== "custom_message" || entry.customType !== CODEX_CONTEXT_WINDOW_MESSAGE_TYPE) continue;
832
+ if (!couldBelongToSession(entry.details, sessionId)) continue;
833
+ if (!isCodexContextManagementMessageDetails(entry.details)) continue;
834
+ if (!matchesSession(entry.details.sessionId, sessionId)) continue;
835
+ const details = entry.details.contextManagement;
836
+ if (details.kind === "window" && details.currentWindowId === windowId) return true;
837
+ }
838
+ return false;
839
+ }
840
+
841
+ function findLatestContextMarkerId(
842
+ entries: readonly SessionEntry[],
843
+ sessionId?: string,
844
+ ): string | undefined {
845
+ for (let index = entries.length - 1; index >= 0; index -= 1) {
846
+ const entry = entries[index]!;
847
+ if (entry.type !== "custom_message" || entry.customType !== CODEX_CONTEXT_WINDOW_MESSAGE_TYPE) continue;
848
+ if (!couldBelongToSession(entry.details, sessionId)) continue;
849
+ if (!isCodexContextManagementMessageDetails(entry.details)) {
850
+ throw new Error("Malformed persisted Codex context-window message");
851
+ }
852
+ if (matchesSession(entry.details.sessionId, sessionId)) return entry.id;
853
+ }
854
+ return undefined;
855
+ }
856
+
857
+ function couldBelongToSession(details: unknown, sessionId: string | undefined): boolean {
858
+ if (sessionId === undefined || !isRecord(details)) return true;
859
+ const markerSessionId = details.sessionId;
860
+ return !isNonEmptyString(markerSessionId) || markerSessionId === sessionId;
861
+ }
862
+
863
+ function matchesSession(markerSessionId: string | undefined, sessionId: string | undefined): boolean {
864
+ return sessionId === undefined || markerSessionId === sessionId;
865
+ }
866
+
867
+ function estimateProjectedWindowTokens(messages: readonly AgentMessage[]): number {
868
+ let lastUsageIndex = -1;
869
+ let usageTokens = 0;
870
+ for (let index = messages.length - 1; index >= 0; index -= 1) {
871
+ const message = messages[index]!;
872
+ if (message.role !== "assistant" || message.stopReason === "aborted" || message.stopReason === "error" || !message.usage) continue;
873
+ const tokens = calculateContextTokens(message.usage);
874
+ if (tokens <= 0) continue;
875
+ lastUsageIndex = index;
876
+ usageTokens = tokens;
877
+ break;
878
+ }
879
+ if (lastUsageIndex < 0) return messages.reduce((total, message) => total + estimateProjectedMessageTokens(message), 0);
880
+ return usageTokens + messages
881
+ .slice(lastUsageIndex + 1)
882
+ .reduce((total, message) => total + estimateProjectedMessageTokens(message), 0);
883
+ }
884
+
885
+ function estimateProjectedMessageTokens(message: AgentMessage): number {
886
+ // A synthetic custom marker may be represented without content in a test or
887
+ // by an older host. It contributes no measurable text; it must not make the
888
+ // context hook fail closed merely because the diagnostic estimate is missing.
889
+ if (
890
+ !("content" in message) ||
891
+ (typeof message.content !== "string" && !Array.isArray(message.content))
892
+ ) return 0;
893
+ return estimateTokens(message);
894
+ }
895
+
896
+ function projectEncryptedToolResults(messages: readonly AgentMessage[]): AgentMessage[] {
897
+ let changed = false;
898
+ const projected = messages.map((message) => {
899
+ if (message.role !== "toolResult" || !isRecord(message.details)) return message;
900
+ const historyNotes = message.details.codexHistoryNotes;
901
+ if (!isRecord(historyNotes) || typeof historyNotes.encrypted_output !== "string") return message;
902
+ changed = true;
903
+ return {
904
+ ...message,
905
+ content: [
906
+ { type: "text" as const, text: encodeEncryptedOutputForContext(historyNotes.encrypted_output) },
907
+ ...message.content.filter((item) => item.type === "image"),
908
+ ],
909
+ };
910
+ });
911
+ return changed ? projected : [...messages];
912
+ }