feat(ext): budget routed compaction off the served model's context window

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:<sessionId>"` pins the main loop
   to one model for the session; subagents get their own pins.
 - **`dispatch` tool — parallel, context-isolated subagents.** pi has none

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

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:";
 

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);
+});

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 </dev/null; then
-  [ "$(grep -Ec '^(✔ |ok [0-9]+ - )' "$WORK/unit.out" || true)" = "22" ] \
+  [ "$(grep -Ec '^(✔ |ok [0-9]+ - )' "$WORK/unit.out" || true)" = "25" ] \
     && ok "pricing, force-model, UI, and compaction unit suite passed" \
-    || bad "unit suite did not report all 22 passes (see $WORK/unit.out)"
+    || bad "unit suite did not report all 25 passes (see $WORK/unit.out)"
 else
   bad "unit suite failed to load through pi (see $WORK/unit.out)"
 fi