Skip to content

fix: recover stale context limits after model switches #341

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
coleleavitt wants to merge 8 commits into
base: master
Choose a base branch
from
Open

fix: recover stale context limits after model switches #341

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
fdc66ff
fix(plugin): trust successful context usage samples
coleleavitt Aug 19, 2026
7eff875
fix(plugin): preserve tokenless response timing
coleleavitt Aug 19, 2026
ce09be2
fix(plugin): keep prompt bounds out of cache geometry
coleleavitt Aug 19, 2026
2d610d5
fix(plugin): preserve overflow state across model switches
coleleavitt Aug 19, 2026
f4f6b55
fix(plugin): preserve legacy usage timestamps
coleleavitt Aug 19, 2026
6a5e30d
fix(plugin): guard legacy usage backfill
coleleavitt Aug 19, 2026
af22bd4
test(plugin): support stacked schema fences
coleleavitt Aug 19, 2026
8549d11
fix(plugin): restore usage before tokenless refresh
coleleavitt Aug 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
231 changes: 230 additions & 1 deletion packages/plugin/src/hooks/magic-context/event-handler.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ import {
getHistorianFailureState,
getMaxCompressionDepth,
getOrCreateSessionMeta,
getOverflowState,
getStrippedPlaceholderIds,
getTagsBySession,
incrementCompressionDepth,
Expand Down Expand Up @@ -42,6 +43,7 @@ import type { ContextUsage } from "../../features/magic-context/types";
import { getWindowReportsPath } from "../../features/magic-context/window-report-ledger";
import { clearModelsDevCache, refreshModelLimitsFromApi } from "../../shared/models-dev-cache";
import { createEventHandler } from "./event-handler";
import { createEventHook } from "./hook-handlers";

type ContextUsageCacheEntry = {
usage: ContextUsage;
Expand Down Expand Up @@ -159,6 +161,29 @@ function providersClient(limit: number, prompt?: ReturnType<typeof mock>) {
};
}

function createHostEventHook(
deps: ReturnType<typeof createDeps>,
contextUsageMap: Map<string, ContextUsageCacheEntry>,
) {
return createEventHook({
eventHandler: createEventHandler(deps),
contextUsageMap,
db: deps.db,
liveModelBySession: new Map(),
variantBySession: new Map(),
agentBySession: new Map(),
sessionDirectoryBySession: new Map(),
historyRefreshSessions: new Set(),
deferredHistoryRefreshSessions: new Set(),
systemPromptRefreshSessions: new Set(),
pendingMaterializationSessions: new Set(),
deferredMaterializationSessions: new Set(),
lastHeuristicsTurnId: new Map(),
client: deps.client as never,
protectedTags: 5,
});
}

describe("createEventHandler", () => {
it("normalizes transform decision reasons across harnesses", () => {
expect(normalizeMaterializeReason("opencode", "system_hash", true)).toBe("system_hash");
Expand Down Expand Up @@ -546,6 +571,209 @@ describe("createEventHandler", () => {
expect(contextUsageMap.get("ses-regression-recovered")?.usage.percentage).toBe(90);
});

it("uses an output-length completion after the real model-switch reset to correct a stale low limit", async () => {
// Given: a prior model established a small safe-token baseline while the
// catalog advertises only 128k for the next model.
useTempDataHome("context-event-model-switch-first-sample-");
await refreshModelLimitsFromApi(providersClient(128_000));
const contextUsageMap = new Map<string, ContextUsageCacheEntry>();
const deps = createDeps(contextUsageMap);
deps.client = providersClient(128_000);
const hook = createHostEventHook(deps, contextUsageMap);
await hook({
event: {
type: "message.updated",
properties: {
info: {
role: "assistant",
finish: "stop",
sessionID: "ses-model-switch-first-sample",
providerID: "test-provider",
modelID: "old-model",
tokens: { input: 1_000, cache: { read: 0, write: 0 } },
},
},
},
});

// When: an error-free output-length response on the new model proves
// a 10M-token prompt, including cache reads and writes, succeeded.
await hook({
event: {
type: "message.updated",
properties: {
info: {
role: "assistant",
finish: "length",
sessionID: "ses-model-switch-first-sample",
providerID: "test-provider",
modelID: "test-model",
tokens: { input: 6_000_000, cache: { read: 3_000_000, write: 1_000_000 } },
},
},
},
});

// Then: shared metadata and live pressure both use the proven lower bound.
const meta = getOrCreateSessionMeta(deps.db, "ses-model-switch-first-sample");
expect(contextUsageMap.get("ses-model-switch-first-sample")?.usage).toEqual({
inputTokens: 10_000_000,
percentage: 100,
});
expect(meta.lastUsageContextLimit).toBe(10_000_000);
expect(getOverflowState(deps.db, "ses-model-switch-first-sample")).toMatchObject({
detectedContextLimit: 10_000_000,
detectedContextLimitModelKey: "test-provider/test-model",
detectedContextLimitProvenance: "prompt_only",
});
});

for (const completion of [
{
name: "an error-bearing completed event",
finish: undefined,
completed: 1,
error: "boom",
},
{ name: "finish:error", finish: "error", completed: undefined, error: undefined },
{
name: "a nonterminal numeric usage event",
finish: undefined,
completed: undefined,
error: undefined,
},
]) {
it(`does not update ordinary usage from ${completion.name}`, async () => {
// Given: ordinary usage, a pending transform decision, and historian
// failure state established by the prior successful model.
useTempDataHome("context-event-model-switch-failed-sample-");
await refreshModelLimitsFromApi(providersClient(128_000));
const contextUsageMap = new Map<string, ContextUsageCacheEntry>();
const deps = createDeps(contextUsageMap);
deps.client = providersClient(128_000);
const hook = createHostEventHook(deps, contextUsageMap);
await hook({
event: {
type: "message.updated",
properties: {
info: {
role: "assistant",
finish: "stop",
sessionID: "ses-model-switch-failed-sample",
providerID: "test-provider",
modelID: "old-model",
tokens: { input: 100_000, cache: { read: 0, write: 0 } },
},
},
},
});
incrementHistorianFailure(deps.db, "ses-model-switch-failed-sample", "failed");
recordPendingTransformDecision("ses-model-switch-failed-sample", {
tsMs: 1,
decision: "execute",
materialized: true,
materializeReason: "pressure_refold",
emergency: false,
droppedTokens: 0,
droppedCount: 1,
inputTokens: 100_000,
bustedThisPass: true,
});
const baselineMeta = getOrCreateSessionMeta(deps.db, "ses-model-switch-failed-sample");
const baselineUsage = contextUsageMap.get("ses-model-switch-failed-sample");

// When: the new model reports usage on a failed completion shape.
await hook({
event: {
type: "message.updated",
properties: {
info: {
id: "msg-model-switch-failed-sample",
role: "assistant",
finish: completion.finish,
time:
completion.completed === undefined
? undefined
: { completed: completion.completed },
error: completion.error,
sessionID: "ses-model-switch-failed-sample",
providerID: "test-provider",
modelID: "test-model",
tokens: { input: 258_901, cache: { read: 0, write: 0 } },
},
},
},
});
await waitForTimers();

// Then: live pressure, persisted pressure, historian recovery, and
// the pending transform decision remain exactly as they were.
expect(contextUsageMap.get("ses-model-switch-failed-sample")).toEqual(baselineUsage);
const meta = getOrCreateSessionMeta(deps.db, "ses-model-switch-failed-sample");
expect(meta.lastContextPercentage).toBe(baselineMeta.lastContextPercentage);
expect(meta.lastInputTokens).toBe(baselineMeta.lastInputTokens);
expect(meta.lastUsageContextLimit).toBe(baselineMeta.lastUsageContextLimit);
expect(meta.lastObservedModelKey).toBe(baselineMeta.lastObservedModelKey);
expect(meta.observedSafeInputTokens).toBe(baselineMeta.observedSafeInputTokens);
expect(
getHistorianFailureState(deps.db, "ses-model-switch-failed-sample").failureCount,
).toBe(1);
const row = deps.db
.prepare("SELECT COUNT(*) AS count FROM transform_decisions WHERE session_id = ?")
.get("ses-model-switch-failed-sample") as { count: number };
expect(row.count).toBe(0);
expect(
transformDecisionLogTest.getPending("ses-model-switch-failed-sample"),
).toBeDefined();
});
}

for (const malformed of [
{ name: "fractional input", tokens: { input: 1.5, cache: { read: 0, write: 0 } } },
{ name: "negative cache usage", tokens: { input: 1, cache: { read: -1, write: 0 } } },
{
name: "cache-inclusive integer overflow",
tokens: { input: Number.MAX_SAFE_INTEGER, cache: { read: 1, write: 0 } },
},
]) {
it(`rejects ${malformed.name}`, async () => {
// Given: stale 128k catalog metadata and no prior usage.
useTempDataHome("context-event-model-switch-malformed-sample-");
await refreshModelLimitsFromApi(providersClient(128_000));
const contextUsageMap = new Map<string, ContextUsageCacheEntry>();
const deps = createDeps(contextUsageMap);
deps.client = providersClient(128_000);
const hook = createHostEventHook(deps, contextUsageMap);

// When: the host reports malformed cache-inclusive usage.
await hook({
event: {
type: "message.updated",
properties: {
info: {
role: "assistant",
finish: "stop",
sessionID: "ses-model-switch-malformed-sample",
providerID: "test-provider",
modelID: "test-model",
tokens: malformed.tokens,
},
},
},
});

// Then: malformed usage cannot overwrite metadata or live pressure.
expect(
getOverflowState(deps.db, "ses-model-switch-malformed-sample").detectedContextLimit,
).toBe(0);
expect(contextUsageMap.get("ses-model-switch-malformed-sample")).toBeUndefined();
expect(
getOrCreateSessionMeta(deps.db, "ses-model-switch-malformed-sample")
.observedSafeInputTokens,
).toBe(0);
});
}

it("alerts once when a cache-regressed context limit stays wrong after refresh", async () => {
useTempDataHome("context-event-cache-regression-alert-");
const contextUsageMap = new Map<string, ContextUsageCacheEntry>();
Expand Down Expand Up @@ -596,7 +824,7 @@ describe("createEventHandler", () => {

const meta = getOrCreateSessionMeta(openDatabase(), "ses-regression-alert");
expect(meta.cacheAlertSent).toBe(true);
expect(meta.lastContextPercentage).toBe(400);
expect(meta.lastContextPercentage).toBe(100);
expect(prompt).toHaveBeenCalledTimes(1);
const call = prompt.mock.calls[0]?.[0] as { body?: { parts?: Array<{ text?: string }> } };
expect(call.body?.parts?.[0]?.text).toContain("context limit of 30,000 tokens");
Expand Down Expand Up @@ -781,6 +1009,7 @@ describe("createEventHandler", () => {
properties: {
info: {
role: "assistant",
finish: "stop",
sessionID: "ses-bg",
tokens: { input: 120_000, cache: { read: 12_000, write: 0 } },
},
Expand Down
57 changes: 36 additions & 21 deletions packages/plugin/src/hooks/magic-context/event-handler.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@ import {
getSessionCreatedInfo,
getSessionErrorInfo,
getSessionProperties,
isSuccessfulHostEvent,
} from "./event-payloads";
import {
resolveCacheTtl,
Expand Down Expand Up @@ -424,8 +425,6 @@ export function createEventHandler(deps: EventHandlerDeps) {
});
}

let messageHadOverflowError = false;

// Secondary overflow-detection path: OpenCode attaches overflow
// errors to the assistant message itself in addition to emitting
// session.error. Checking both ensures we catch the error no
Expand All @@ -435,7 +434,6 @@ export function createEventHandler(deps: EventHandlerDeps) {
if (info.error !== undefined && info.error !== null) {
const detection = detectOverflow(info.error);
if (detection.isOverflow) {
messageHadOverflowError = true;
try {
captureWindowReport({
db: deps.db,
Expand Down Expand Up @@ -518,15 +516,24 @@ export function createEventHandler(deps: EventHandlerDeps) {
}
}

if (!isSuccessfulHostEvent(info)) return;
Comment thread
greptile-apps[bot] marked this conversation as resolved.
Outdated
Comment thread
cubic-dev-ai[bot] marked this conversation as resolved.
Outdated

const now = Date.now();
const usageTokens = [
info.tokens?.input,
info.tokens?.cache?.read,
info.tokens?.cache?.write,
];
const hasUsageTokens = usageTokens.some(
(value) => typeof value === "number" && value > 0,
const totalInputTokens = usageTokens.reduce<number>(
(total, value) => total + (value ?? 0),
0,
);
const hasUsageTokens =
usageTokens.every(
(value) => value === undefined || (Number.isSafeInteger(value) && value >= 0),
) &&
usageTokens.some((value) => typeof value === "number" && value > 0) &&
Number.isSafeInteger(totalInputTokens);
const terminalAssistantUpdate =
info.messageID !== undefined &&
hasUsageTokens &&
Expand All @@ -536,10 +543,7 @@ export function createEventHandler(deps: EventHandlerDeps) {
db: deps.db,
sessionId: info.sessionID,
messageId: info.messageID,
inputTokens:
(info.tokens?.input ?? 0) +
(info.tokens?.cache?.read ?? 0) +
(info.tokens?.cache?.write ?? 0),
inputTokens: totalInputTokens,
});
}

Expand Down Expand Up @@ -570,10 +574,6 @@ export function createEventHandler(deps: EventHandlerDeps) {
}

if (hasUsageTokens) {
const totalInputTokens =
(info.tokens?.input ?? 0) +
(info.tokens?.cache?.read ?? 0) +
(info.tokens?.cache?.write ?? 0);
// Auth is provably live now (a request returned usage), so
// re-warm the model-limit cache once per process to overwrite
// any stale pre-auth limit (e.g. gpt-5.5 cached at the raw
Expand All @@ -599,8 +599,9 @@ export function createEventHandler(deps: EventHandlerDeps) {
const observedSafeInputTokens = sessionMeta.observedSafeInputTokens ?? 0;
if (
percentage > 100 &&
observedSafeInputTokens > 0 &&
totalInputTokens <= observedSafeInputTokens * 2
modelKey !== undefined &&
(observedSafeInputTokens === 0 ||
totalInputTokens <= observedSafeInputTokens * 2)
) {
const oldLimit = contextLimit;
if (deps.client) {
Expand Down Expand Up @@ -635,6 +636,22 @@ export function createEventHandler(deps: EventHandlerDeps) {
updates.cacheAlertSent = true;
}
}

if (contextLimit < totalInputTokens) {
Comment thread
cubic-dev-ai[bot] marked this conversation as resolved.
recordDetectedContextLimit(
deps.db,
info.sessionID,
totalInputTokens,
modelKey,
"prompt_only",
);
contextLimit = resolveContextLimit(info.providerID, info.modelID, {
db: deps.db,
sessionID: info.sessionID,
});
percentage = (totalInputTokens / contextLimit) * 100;
if (sessionMeta.cacheAlertSent) updates.cacheAlertSent = true;
}
}

deps.contextUsageMap.set(info.sessionID, {
Expand All @@ -651,12 +668,10 @@ export function createEventHandler(deps: EventHandlerDeps) {
updates.lastInputTokens = totalInputTokens;
updates.lastUsageContextLimit = contextLimit;
updates.lastObservedModelKey = modelKey ?? null;
if (!messageHadOverflowError) {
updates.observedSafeInputTokens = Math.max(
observedSafeInputTokens,
totalInputTokens,
);
}
updates.observedSafeInputTokens = Math.max(
observedSafeInputTokens,
totalInputTokens,
);

const historianFailureState = getHistorianFailureState(deps.db, info.sessionID);
if (historianFailureState.failureCount > 0 && percentage < 90) {
Expand Down
Loading
Loading