diff --git a/src/arbitration.ts b/src/arbitration.ts new file mode 100644 index 0000000..84fa07e --- /dev/null +++ b/src/arbitration.ts @@ -0,0 +1,72 @@ +/** + * Provider-anchored context arbitration. + * + * chars/4×density estimates of the sent view can run far below the provider's + * real token count (CJK + minified code under-report by 4×+; incident session + * 01a02d90: estimated 57K = 51.8% while the provider was already rejecting a + * 134.5K prompt). pi's `getContextUsage().tokens` anchors on the provider- + * reported usage of the last assistant message (plus an estimated trailing + * tail), so when that anchor is trustworthy it beats the local estimate + * outright for nudge/emergency arbitration. + * + * The anchor is NOT trustworthy in two regimes, both observed in the wild: + * + * 1. Transient post-compression turn: the anchor assistant predates the + * compress, so its usage still reflects the pre-compression (much larger) + * sent view. Consuming it would fire false EMERGENCY nudges right after a + * successful compress (omp issue #18 family). Density's postCompression + * skip (density.ts) uses the same guard. + * + * 2. Provider-never-reports-usage regime: when no assistant message ever + * carried usage, pi falls back to summing the whole session tree + * (originals included, never shrinks) — after compression that number can + * exceed the window many times over while the real sent view is a few + * percent, producing permanent false EMERGENCY nudges (omp issue #18). + * Detect it the same way pi's compaction does (compaction.ts + * getAssistantUsage): an assistant entry only counts as provider-backed + * when it carries a non-zero usage record. + */ + +export interface UsageLike { + tokens: number | null; +} + +export interface EntryLike { + type?: string; + message?: { + role?: string; + usage?: { + input?: number; + totalTokens?: number; + }; + }; +} + +/** + * Resolve the provider-anchored token count for arbitration, or null when the + * anchor must not be consumed this turn. + * + * Returns `realUsage.tokens` when ALL hold: + * a) it is a positive finite number (pi yields null when it cannot anchor, + * e.g. right after a pi-side compaction with no post-compaction usage); + * b) this is not a post-compression transient turn (guard 1 above); + * c) at least one assistant entry in the merged session view carries a + * non-zero provider usage record (guard 2 above) — mirroring how pi + * decides between usage-anchoring and tree-summing. + */ +export function providerAnchoredTokens( + realUsage: UsageLike | undefined, + entries: readonly EntryLike[], + postCompression: boolean, +): number | null { + const tokens = realUsage?.tokens; + if (typeof tokens !== "number" || !Number.isFinite(tokens) || tokens <= 0) return null; + if (postCompression) return null; + for (let i = entries.length - 1; i >= 0; i--) { + const message = entries[i]?.message; + if (message?.role !== "assistant") continue; + const usage = message.usage; + if ((usage?.totalTokens ?? 0) > 0 || (usage?.input ?? 0) > 0) return tokens; + } + return null; +} diff --git a/src/index.ts b/src/index.ts index a041e4e..d6dadd6 100644 --- a/src/index.ts +++ b/src/index.ts @@ -20,6 +20,7 @@ import { delegateStatusWidget } from "./fleet-widget.js"; import { wireToolGuardrails } from "./tool-guardrails.js"; import { debug, logError, logInfo, logWarn, logThrow, closeLogStream } from "./log.js"; import { collectCoveredMessageIds, estimateTokens, lastUserMessageId, calibrateTokens } from "./tokens.js"; +import { providerAnchoredTokens } from "./arbitration.js"; import { checkForUpdate } from "./update.js"; import { THROTTLE_RETRY_ERROR_MESSAGE, @@ -195,6 +196,30 @@ function wireContextTransform(pi: ExtensionAPI, runtime: AcpRuntime): void { // below is fed the RAW sentTokens — its samples must stay on the // raw basis or density would chase its own calibration. let tokenCount = calibrateTokens(sentTokens, runtime.density.densityFor(modelId)); + // A compress happened since the previous context round: blocks are + // created out-of-band by the compress tool, so detect new active + // blocks on the LOADED state vs. the previous round (comparing a + // single processTurn's input/output can never see them). Hoisted ABOVE + // the provider-anchored arbitration below: post-compression turns must + // not consume pi's usage anchor (it predates the shrink — see + // arbitration.ts). + const postCompression = runtime.noteActiveBlocks( + sid, + state.blocks.filter((b) => b.active).map((b) => b.blockId), + ); + // Provider-anchored arbitration (incident 01a02d90): when pi's usage + // anchor is trustworthy it beats the calibrated estimate outright — + // that session arbitrated on 57K (51.8%) while the provider was already + // rejecting a 134.5K prompt (4.5× underestimate; density was clamped at + // its 2.5 ceiling and could never catch up). Take the LARGER of the + // two so a stale-low estimate can never mask a real overflow; the + // anchored number already includes pi's estimated trailing tail. + // Guarded against post-compression transients and the + // provider-never-reports-usage tree-sum regime (omp #18). + const anchoredTokens = providerAnchoredTokens(realUsage, entries, postCompression); + if (anchoredTokens !== null && anchoredTokens > tokenCount) { + tokenCount = anchoredTokens; + } // Self-heal (armed): after an overflow, force this turn's usage to >=95% // so the kernel's emergency nudge + tool-result truncate fire immediately, // even if the density-calibrated estimate under-reports the sent view. @@ -208,14 +233,6 @@ function wireContextTransform(pi: ExtensionAPI, runtime: AcpRuntime): void { logWarn("overflow-selfheal", { sid, event: "armed-emergency", tokenCount, limit: config.modelContextLimit }); } } - // A compress happened since the previous context round: blocks are - // created out-of-band by the compress tool, so detect new active - // blocks on the LOADED state vs. the previous round (comparing a - // single processTurn's input/output can never see them). - const postCompression = runtime.noteActiveBlocks( - sid, - state.blocks.filter((b) => b.active).map((b) => b.blockId), - ); debug.event("context-in", { sid, @@ -226,6 +243,7 @@ function wireContextTransform(pi: ExtensionAPI, runtime: AcpRuntime): void { coreMsgs: coreMessages.length, tokenCount, sessionTokens: realUsage?.tokens ?? null, + anchoredTokens, limit: config.modelContextLimit, blocksBefore: state.blocks.length, activeBefore: state.blocks.filter((b) => b.active).length, diff --git a/tests/arbitration.test.ts b/tests/arbitration.test.ts new file mode 100644 index 0000000..17a32b3 --- /dev/null +++ b/tests/arbitration.test.ts @@ -0,0 +1,60 @@ +import test from "node:test"; +import assert from "node:assert/strict"; + +import { providerAnchoredTokens } from "../src/arbitration.js"; + +function assistantEntry(usage?: { input?: number; totalTokens?: number }) { + return { type: "message", id: "a1", message: { role: "assistant", usage } }; +} +function userEntry() { + return { type: "message", id: "u1", message: { role: "user" } }; +} + +test("providerAnchoredTokens: returns anchored tokens when last assistant carries usage", () => { + const entries = [userEntry(), assistantEntry({ input: 8858, totalTokens: 8875 }), userEntry()]; + assert.equal(providerAnchoredTokens({ tokens: 134569 }, entries, false), 134569); +}); + +test("providerAnchoredTokens: skips non-assistant entries in the back-scan", () => { + const entries = [assistantEntry({ input: 100 }), userEntry(), userEntry(), { type: "message", id: "t1", message: { role: "toolResult" } }]; + assert.equal(providerAnchoredTokens({ tokens: 5000 }, entries, false), 5000); +}); + +test("providerAnchoredTokens: usage-less assistant tail falls through to an earlier provider-backed assistant", () => { + // Tail assistant aborted/error (no usage record) — pi anchors on the earlier + // good one, so the number is still provider-backed. + const entries = [assistantEntry({ input: 100 }), userEntry(), assistantEntry(undefined), userEntry()]; + assert.equal(providerAnchoredTokens({ tokens: 45000 }, entries, false), 45000); +}); + +test("providerAnchoredTokens: omp #18 — no assistant ever reported usage → distrust (tree-sum regime)", () => { + const entries = [assistantEntry(undefined), userEntry(), assistantEntry({ input: 0 }), userEntry()]; + assert.equal(providerAnchoredTokens({ tokens: 999999 }, entries, false), null); +}); + +test("providerAnchoredTokens: empty entries → null", () => { + assert.equal(providerAnchoredTokens({ tokens: 1000 }, [], false), null); +}); + +test("providerAnchoredTokens: post-compression transient turn → null regardless of anchor", () => { + const entries = [assistantEntry({ input: 100 }), userEntry()]; + assert.equal(providerAnchoredTokens({ tokens: 134569 }, entries, true), null); +}); + +test("providerAnchoredTokens: null/zero/non-finite tokens → null (pi cannot anchor)", () => { + const entries = [assistantEntry({ input: 100 })]; + assert.equal(providerAnchoredTokens(undefined, entries, false), null); + assert.equal(providerAnchoredTokens({ tokens: null }, entries, false), null); + assert.equal(providerAnchoredTokens({ tokens: 0 }, entries, false), null); + assert.equal(providerAnchoredTokens({ tokens: Number.NaN }, entries, false), null); +}); + +test("providerAnchoredTokens: totalTokens alone counts as provider-backed (OpenAI-style usage)", () => { + const entries = [assistantEntry({ totalTokens: 8875 })]; + assert.equal(providerAnchoredTokens({ tokens: 8875 }, entries, false), 8875); +}); + +test("providerAnchoredTokens: all-zero usage record does not count as provider-backed", () => { + const entries = [assistantEntry({ input: 0, totalTokens: 0 })]; + assert.equal(providerAnchoredTokens({ tokens: 70000 }, entries, false), null); +});