From 162c73484426af31f03a208bc05ef4b4c52eb30b Mon Sep 17 00:00:00 2001 From: Matthew Evanusa Date: Thu, 20 Aug 2026 14:43:58 -0700 Subject: [PATCH 1/3] feat(ext): budget routed compaction off the served model's context window Pi budgets auto-compaction off the model the client requested, but the router routinely serves a different (often larger-context) model. Without the served window, a session requesting claude-sonnet-4-6 compacted at ~200K even while a 1M-context model served it. Follow the router's x-router-context-window header (emitted by #968): capture it per successful response and use it as the compaction budget at agent_end, falling back to the requested model's static window when absent. The served window is authoritative in both directions - a larger served model avoids premature compaction; a smaller one still compacts early enough. - config: add ROUTED_CONTEXT_WINDOW_HEADER - compaction: track servedContextWindow per run (reset on start/compact) and apply it in the agent_end budget; ignore non-2xx responses and malformed values - tests: 3 new cases (larger served window no pre-compaction; smaller served window mid-loop compaction; absent header keeps requested budget); the harness now parameterizes the requested model window - e2e.sh: unit-suite pass count 22 -> 25 --- install/pi-router/README.md | 5 +- install/pi-router/src/compaction.ts | 18 +++++++- install/pi-router/src/config.ts | 1 + install/pi-router/test/compaction.test.ts | 56 +++++++++++++++++++++-- install/pi-router/test/e2e.sh | 4 +- 5 files changed, 77 insertions(+), 7 deletions(-) diff --git a/install/pi-router/README.md b/install/pi-router/README.md index 19e47b679..ea22f04c7 100644 --- a/install/pi-router/README.md +++ b/install/pi-router/README.md @@ -36,7 +36,10 @@ from npm on next start and loads this extension via its `pi.extensions` field. an uninterrupted tool loop before its normal post-run compaction check. The extension preserves a usable output budget for the real continuation and compacts once the loop settles, while leaving ordinary threshold compaction - to Pi. + to Pi. When the router reports the served model's real window via + `x-router-context-window`, that budget replaces the requested model's static + 200K assumption, so a 1M-context model serving a 200K session is not + pre-compacted (and a smaller served model still compacts early enough). - **Sticky sessions.** `metadata.user_id = "pi:"` pins the main loop to one model for the session; subagents get their own pins. - **`dispatch` tool — parallel, context-isolated subagents.** pi has none diff --git a/install/pi-router/src/compaction.ts b/install/pi-router/src/compaction.ts index 42d158e9e..3c0bc244f 100644 --- a/install/pi-router/src/compaction.ts +++ b/install/pi-router/src/compaction.ts @@ -21,6 +21,7 @@ import type { TurnEndEvent, } from "@mariozechner/pi-coding-agent"; import type { AssistantMessage } from "@mariozechner/pi-ai"; +import { ROUTED_CONTEXT_WINDOW_HEADER } from "./config.js"; const PROBE_MAX_TOKENS = 4; const CONTINUATION_MAX_TOKENS = 16_384; @@ -86,11 +87,15 @@ export function registerCompaction(pi: ExtensionAPI, schedule: Schedule = (callb let lastTurnTokens = 0; let repairedContinuation = false; let compactionScheduled = false; + // Effective context window reported by the router for the model that + // actually served the response; undefined when the header is absent. + let servedContextWindow: number | undefined; const resetRun = () => { highWaterTokens = 0; lastTurnTokens = 0; repairedContinuation = false; + servedContextWindow = undefined; }; const finishCompaction = (ctx: ExtensionContext) => { @@ -107,6 +112,14 @@ export function registerCompaction(pi: ExtensionAPI, schedule: Schedule = (callb if (repairClampedToolContinuation(event.payload)) repairedContinuation = true; }); + pi.on("after_provider_response", (event) => { + if (event.status < 200 || event.status >= 300) return; + const raw = event.headers?.[ROUTED_CONTEXT_WINDOW_HEADER]; + if (typeof raw !== "string" || !/^\d+$/.test(raw)) return; + const servedWindow = Number(raw); + if (servedWindow > COMPACTION_RESERVE_TOKENS) servedContextWindow = servedWindow; + }); + pi.on("turn_end", (event: TurnEndEvent) => { if (event.message.role !== "assistant") return; lastTurnTokens = contextTokens(event.message as AssistantMessage); @@ -115,7 +128,10 @@ export function registerCompaction(pi: ExtensionAPI, schedule: Schedule = (callb pi.on("agent_end", (_event: AgentEndEvent, ctx: ExtensionContext) => { if (process.env.WEAVE_PI_AUTO_COMPACTION === "0" || compactionScheduled) return; - const contextWindow = ctx.model?.contextWindow ?? ctx.getContextUsage()?.contextWindow ?? 0; + // The routed provider's reported window is authoritative for what is + // actually holding context; the requested model's static window is only + // the fallback when the router did not report one. + const contextWindow = servedContextWindow ?? ctx.model?.contextWindow ?? ctx.getContextUsage()?.contextWindow ?? 0; if (contextWindow <= COMPACTION_RESERVE_TOKENS) return; const threshold = contextWindow - COMPACTION_RESERVE_TOKENS; // Pi's built-in check runs immediately after this event and owns the diff --git a/install/pi-router/src/config.ts b/install/pi-router/src/config.ts index 48c5a50fe..be598f16a 100644 --- a/install/pi-router/src/config.ts +++ b/install/pi-router/src/config.ts @@ -246,6 +246,7 @@ function model(id: string, name: string, maxTokens: number, contextWindow: numbe export const ROUTED_MODEL_HEADER = (process.env.WEAVE_ROUTED_MODEL_HEADER || "x-router-model").toLowerCase(); export const ROUTED_PROVIDER_HEADER = "x-router-provider"; export const ROUTER_DECISION_HEADER = "x-router-decision"; +export const ROUTED_CONTEXT_WINDOW_HEADER = "x-router-context-window"; /** Marker a headless child prints to stderr so the parent dispatch can read its routed model. */ export const ROUTED_MODEL_STDERR_PREFIX = "weave-routed-model:"; diff --git a/install/pi-router/test/compaction.test.ts b/install/pi-router/test/compaction.test.ts index ad54956b1..8c860c0ee 100644 --- a/install/pi-router/test/compaction.test.ts +++ b/install/pi-router/test/compaction.test.ts @@ -37,14 +37,14 @@ function assistant(totalTokens: number) { }; } -function contextHarness() { +function contextHarness(modelWindow = 200_000) { let compactCalls = 0; let status: string | undefined; const branch: any[] = []; const ctx = { hasUI: true, - model: { contextWindow: 200_000 }, - getContextUsage: () => ({ tokens: 0, contextWindow: 200_000, percent: 0 }), + model: { contextWindow: modelWindow }, + getContextUsage: () => ({ tokens: 0, contextWindow: modelWindow, percent: 0 }), sessionManager: { getBranch: () => branch }, ui: { setStatus(_key: string, value: string | undefined) { @@ -123,3 +123,53 @@ test("leaves an over-threshold final turn to Pi's built-in compaction", () => { assert.equal(compactCalls(), 0); }); + +test("budgets compaction off the served x-router-context-window, not the requested model", () => { + const extension = extensionHarness((callback) => callback()); + // ctx.model is the client's requested model (200K); the router's header + // reports the window of the model that actually served the response (1M). + const { ctx, compactCalls } = contextHarness(); + extension.emit("agent_start", { type: "agent_start" }, ctx); + extension.emit( + "after_provider_response", + { type: "after_provider_response", status: 200, headers: { "x-router-context-window": "1000000" } }, + ctx, + ); + extension.emit("turn_end", { type: "turn_end", message: assistant(250_000), toolResults: [] }, ctx); + extension.emit("agent_end", { type: "agent_end", messages: [] }, ctx); + + // 250K is above the 200K requested-model budget but inside the served 1M + // budget, so the extension must not pre-compact. + assert.equal(compactCalls(), 0); +}); + +test("compacts when the served window is smaller than the requested model's", () => { + const extension = extensionHarness((callback) => callback()); + const { ctx, compactCalls } = contextHarness(1_000_000); + extension.emit("agent_start", { type: "agent_start" }, ctx); + extension.emit( + "after_provider_response", + { type: "after_provider_response", status: 200, headers: { "x-router-context-window": "200000" } }, + ctx, + ); + extension.emit("turn_end", { type: "turn_end", message: assistant(250_000), toolResults: [] }, ctx); + extension.emit("turn_end", { type: "turn_end", message: assistant(50_000), toolResults: [] }, ctx); + extension.emit("agent_end", { type: "agent_end", messages: [] }, ctx); + + // highWater 250K exceeds the served 200K budget even though the requested + // model window is 1M, so a mid-loop compaction is required. + assert.equal(compactCalls(), 1); +}); + +test("does not shrink the budget when the served window header is absent", () => { + const extension = extensionHarness((callback) => callback()); + const { ctx, compactCalls } = contextHarness(1_000_000); + extension.emit("agent_start", { type: "agent_start" }, ctx); + extension.emit("turn_end", { type: "turn_end", message: assistant(250_000), toolResults: [] }, ctx); + extension.emit("turn_end", { type: "turn_end", message: assistant(50_000), toolResults: [] }, ctx); + extension.emit("agent_end", { type: "agent_end", messages: [] }, ctx); + + // Without the routed window, the 1M requested budget applies and the run + // stays below it, so Pi's ordinary threshold compaction owns the case. + assert.equal(compactCalls(), 0); +}); diff --git a/install/pi-router/test/e2e.sh b/install/pi-router/test/e2e.sh index d753ef75f..27845ce9c 100755 --- a/install/pi-router/test/e2e.sh +++ b/install/pi-router/test/e2e.sh @@ -133,9 +133,9 @@ phase "Phase 2 — generated pricing + savings contract" if with_timeout 30 env PI_CODING_AGENT_DIR="$PI_DIR" \ pi -e "$UNIT_SUITE" --no-session --offline --model weave/claude-sonnet-4-6 \ -p "Run the unit suite." >"$WORK/unit.out" 2>&1 Date: Thu, 20 Aug 2026 14:49:51 -0700 Subject: [PATCH 2/3] ci: re-trigger routing-review workflow_run handoff From d902e3418eab604996a3629fc7fb4c5eba867c8c Mon Sep 17 00:00:00 2001 From: Matthew Evanusa Date: Thu, 20 Aug 2026 14:53:20 -0700 Subject: [PATCH 3/3] fix(ext): compact final turns above a served window smaller than the requested budget Pi's built-in compaction only fires when the requested model's static window is exceeded. When the router serves a smaller window, a final turn above the served threshold but below the requested budget previously fell through both sides: agent_end returned early for Pi to own it, Pi never fired, and the served model could overflow. Only defer the over-threshold final turn to Pi when Pi's own budget is also exceeded (or the served window is at least the requested window). Add a regression test for the served-window-smaller-than-requested final-turn case. Addresses Cursor Bugbot review on #970. --- install/pi-router/src/compaction.ts | 12 +++++++++--- install/pi-router/test/compaction.test.ts | 17 +++++++++++++++++ install/pi-router/test/e2e.sh | 4 ++-- 3 files changed, 28 insertions(+), 5 deletions(-) diff --git a/install/pi-router/src/compaction.ts b/install/pi-router/src/compaction.ts index 3c0bc244f..79ea00683 100644 --- a/install/pi-router/src/compaction.ts +++ b/install/pi-router/src/compaction.ts @@ -134,10 +134,16 @@ export function registerCompaction(pi: ExtensionAPI, schedule: Schedule = (callb const contextWindow = servedContextWindow ?? ctx.model?.contextWindow ?? ctx.getContextUsage()?.contextWindow ?? 0; if (contextWindow <= COMPACTION_RESERVE_TOKENS) return; const threshold = contextWindow - COMPACTION_RESERVE_TOKENS; + const requestedWindow = ctx.model?.contextWindow ?? ctx.getContextUsage()?.contextWindow ?? 0; + const requestedThreshold = requestedWindow - COMPACTION_RESERVE_TOKENS; // Pi's built-in check runs immediately after this event and owns the - // ordinary final-turn threshold case. Starting another compaction while - // that async summary is in flight would race it. - if (lastTurnTokens > threshold) return; + // ordinary final-turn threshold case; starting another compaction while + // that async summary is in flight would race it. But Pi only fires when + // the requested model's budget is exceeded. When the served window is + // smaller than the requested budget, a final turn above the served + // threshold and below the requested one would otherwise fall through + // both sides and overflow the served model - so compact it here. + if (lastTurnTokens > threshold && lastTurnTokens > requestedThreshold) return; if (!repairedContinuation && highWaterTokens <= threshold) return; compactionScheduled = true; diff --git a/install/pi-router/test/compaction.test.ts b/install/pi-router/test/compaction.test.ts index 8c860c0ee..50bb56307 100644 --- a/install/pi-router/test/compaction.test.ts +++ b/install/pi-router/test/compaction.test.ts @@ -173,3 +173,20 @@ test("does not shrink the budget when the served window header is absent", () => // stays below it, so Pi's ordinary threshold compaction owns the case. assert.equal(compactCalls(), 0); }); + +test("compacts an over-threshold final turn when the served window is smaller than the requested budget", () => { + const extension = extensionHarness((callback) => callback()); + const { ctx, compactCalls } = contextHarness(1_000_000); + extension.emit("agent_start", { type: "agent_start" }, ctx); + extension.emit( + "after_provider_response", + { type: "after_provider_response", status: 200, headers: { "x-router-context-window": "200000" } }, + ctx, + ); + extension.emit("turn_end", { type: "turn_end", message: assistant(190_000), toolResults: [] }, ctx); + extension.emit("agent_end", { type: "agent_end", messages: [] }, ctx); + + // 190K sits above the served 200K budget threshold (183,616) but below the + // requested 1M budget - Pi would not compact, so the extension must. + assert.equal(compactCalls(), 1); +}); diff --git a/install/pi-router/test/e2e.sh b/install/pi-router/test/e2e.sh index 27845ce9c..e51e0f164 100755 --- a/install/pi-router/test/e2e.sh +++ b/install/pi-router/test/e2e.sh @@ -133,9 +133,9 @@ phase "Phase 2 — generated pricing + savings contract" if with_timeout 30 env PI_CODING_AGENT_DIR="$PI_DIR" \ pi -e "$UNIT_SUITE" --no-session --offline --model weave/claude-sonnet-4-6 \ -p "Run the unit suite." >"$WORK/unit.out" 2>&1