fix(usage): measure server-side fallback replies by the last run (#877)

packages/plugin/src/shared/provider-response-completion.test.ts

@@ -0,0 +1,45 @@
+import { expect, test } from "bun:test";
+import * as completion from "./provider-response-completion";
+
+for (const location of ["native", "state"]) {
+    test(`detects usage and boundary markers in ${location}.signature`, () => {
+        for (const signature of [
+            'cortexkit-server-fallback-usage-v1:{"iterations":[]}',
+            "cortexkit-server-fallback-v1:primary|fallback",
+        ]) {
+            expect(
+                completion.hasServerFallbackUsage([
+                    { type: "reasoning", [location]: { signature } },
+                ]),
+            ).toBe(true);
+        }
+    });
+}
+
+test("detects server-fallback usage only in the marked reasoning part", () => {
+    const content = [
+        { type: "reasoning", text: "thinking", native: { signature: "real-signature" } },
+        {
+            type: "reasoning",
+            text: " ",
+            native: {
+                signature: "cortexkit-server-fallback-v1:claude-opus-5-5|claude-opus-4-8",
+            },
+        },
+        { type: "text", text: "answer" },
+        { type: "tool", name: "read" },
+    ];
+    expect(completion.hasServerFallbackUsage(content)).toBe(true);
+    expect(completion.hasServerFallbackUsage([content[0], content[2]])).toBe(false);
+    for (const invalid of [
+        undefined,
+        null,
+        {},
+        [null, {}, { type: "reasoning", native: { signature: 3 } }],
+        [{ type: "text", native: content[1]!.native }],
+        [{ type: "reasoning", native: { signature: "cortexkit-server-fallback-v10:models" } }],
+        [{ type: "reasoning", native: { signature: "other-cortexkit-server-fallback-v1:models" } }],
+    ]) {
+        expect(completion.hasServerFallbackUsage(invalid)).toBe(false);
+    }
+});

packages/plugin/src/shared/provider-response-completion.ts

@@ -1,3 +1,67 @@
+import { isRecord } from "./record-type-guard";
+
+const SERVER_FALLBACK_USAGE_SIGNATURE_PREFIX = "cortexkit-server-fallback-usage-v1:";
+const SERVER_FALLBACK_SIGNATURE_PREFIX = "cortexkit-server-fallback-v1:";
+
+export function hasServerFallbackUsage(content: unknown): boolean {
+    return readServerFallbackUsage(content) !== undefined;
+}
+
+type ServerFallbackUsage =
+    | { kind: "usage"; inputTokens: number }
+    | { kind: "boundary"; invalidUsage: boolean };
+
+function parseServerFallbackInputTokens(json: string): number | undefined {
+    try {
+        const payload: unknown = JSON.parse(json);
+        if (!isRecord(payload) || !Array.isArray(payload.iterations)) return undefined;
+        let last: unknown;
+        for (const entry of payload.iterations) {
+            if (isRecord(entry) && entry.type === "fallback_message") last = entry;
+        }
+        if (!isRecord(last)) return undefined;
+        const counts = [
+            last.input_tokens,
+            last.cache_read_input_tokens ?? 0,
+            last.cache_creation_input_tokens ?? 0,
+        ];
+        if (
+            !counts.every(
+                (count) => typeof count === "number" && Number.isFinite(count) && count >= 0,
+            )
+        )
+            return undefined;
+        const inputTokens = counts.reduce<number>((sum, count) => sum + Number(count), 0);
+        return Number.isFinite(inputTokens) ? inputTokens : undefined;
+    } catch {
+        // Invalid usage must not replace the previous pressure measurement.
+        return undefined;
+    }
+}
+
+export function readServerFallbackUsage(content: unknown): ServerFallbackUsage | undefined {
+    if (!Array.isArray(content)) return undefined;
+    let boundary = false;
+    let invalidUsage = false;
+    for (const part of content) {
+        if (!isRecord(part) || part.type !== "reasoning") continue;
+        for (const location of [part.native, part.state]) {
+            if (!isRecord(location) || typeof location.signature !== "string") continue;
+            const signature = location.signature;
+            if (signature.startsWith(SERVER_FALLBACK_USAGE_SIGNATURE_PREFIX)) {
+                const inputTokens = parseServerFallbackInputTokens(
+                    signature.slice(SERVER_FALLBACK_USAGE_SIGNATURE_PREFIX.length),
+                );
+                if (inputTokens !== undefined) return { kind: "usage", inputTokens };
+                invalidUsage = true;
+            } else if (signature.startsWith(SERVER_FALLBACK_SIGNATURE_PREFIX)) {
+                boundary = true;
+            }
+        }
+    }
+    return boundary || invalidUsage ? { kind: "boundary", invalidUsage } : undefined;
+}
+
 /** Failure markers take precedence over a completion timestamp or token usage. */
 export function providerResponseFailed(input: { finish?: unknown; error?: unknown }): boolean {
     if (input.error !== undefined && input.error !== null && input.error !== "") return true;

packages/plugin/src/v2/hooks/context.ts

@@ -1517,6 +1517,7 @@ export async function registerContext(context: V2Context) {
                     rowModel: latest?.data.model,
                     draftModel: { providerID: draft.model.providerID, id: draft.model.id },
                     tokens: latest?.data.tokens,
+                    content: latest?.data.content,
                     // Completion is authoritative. Legacy finish-only rows use a
                     // stable stored message time, never the time of this reread.
                     completed:

packages/plugin/src/v2/hooks/server-fallback-usage.test.ts

@@ -0,0 +1,115 @@
+import { expect, test } from "bun:test";
+import { getOrCreateSessionMeta } from "../../features/magic-context/storage";
+import { oc2ContextFixture } from "./__tests__/oc2-context.fixture";
+
+const markers = [
+    "boundary",
+    "usage-native",
+    "usage-state",
+    "invalid-json",
+    "missing-fallback",
+    "none",
+];
+for (const marker of markers) {
+    const marked = marker !== "none";
+    const usage = marker.startsWith("usage");
+    test(`OC2 newest ${marker} reply ${marked ? "admits at per-run pressure without execute" : "enters the force band"}`, async () => {
+        const f = await oc2ContextFixture(801 + markers.indexOf(marker), {
+            area: "server-fallback-usage",
+            provider: "anthropic",
+            model: "claude-opus-5-5",
+            cacheTtl: "1h",
+        });
+        try {
+            await f.pass();
+            f.respond(1, 507_000);
+            await f.pass();
+            const before = getOrCreateSessionMeta(f.db, f.session);
+            expect(before.lastInputTokens).toBe(507_000);
+            f.respond(1, 963_300);
+            const latest = f.rows.at(-1)!;
+            const row = f.raw
+                .prepare("SELECT data FROM session_message WHERE id=?")
+                .get(latest.id) as { data: string };
+            const data = JSON.parse(row.data);
+            data.content.unshift({
+                type: "reasoning",
+                text: " ",
+                [marker === "usage-state" ? "state" : "native"]: {
+                    signature: usage
+                        ? `cortexkit-server-fallback-usage-v1:${JSON.stringify({
+                              iterations: [
+                                  {
+                                      type: "message",
+                                      model: "primary",
+                                      input_tokens: 453_300,
+                                      output_tokens: 50,
+                                  },
+                                  {
+                                      type: "fallback_message",
+                                      model: "fallback",
+                                      input_tokens: 100_000,
+                                      output_tokens: 50,
+                                      cache_read_input_tokens: 400_000,
+                                      cache_creation_input_tokens: 10_000,
+                                  },
+                              ],
+                          })}`
+                        : marker === "invalid-json"
+                          ? "cortexkit-server-fallback-usage-v1:{"
+                          : marker === "missing-fallback"
+                            ? 'cortexkit-server-fallback-usage-v1:{"iterations":[]}'
+                            : marked
+                              ? "cortexkit-server-fallback-v1:claude-opus-5-5|claude-opus-4-8"
+                              : "real-signature",
+                },
+            });
+            if (usage)
+                data.content.unshift({
+                    type: "reasoning",
+                    native: {
+                        signature: "cortexkit-server-fallback-v1:primary|fallback",
+                    },
+                });
+            f.raw
+                .prepare("UPDATE session_message SET data=? WHERE id=?")
+                .run(JSON.stringify(data), latest.id);
+            const result = await f.attempt();
+            const logs = result.logs.join("\n");
+            const meta = getOrCreateSessionMeta(f.db, f.session);
+            if (marked) {
+                expect(result.error).toBeUndefined();
+                expect(meta.lastInputTokens).toBe(usage ? 510_000 : 507_000);
+                expect(meta.lastContextPercentage).toBeCloseTo(
+                    usage
+                        ? (510_000 / before.lastUsageContextLimit) * 100
+                        : before.lastContextPercentage,
+                    5,
+                );
+                expect(meta.lastUsageContextLimit).toBe(before.lastUsageContextLimit);
+                expect(meta.lastResponseTime).toBe(latest.created + 1000);
+                expect(logs).toMatch(/transform scheduler: .*decision=defer/);
+                expect(logs).not.toMatch(/transform scheduler: .*decision=execute/);
+                if (!usage)
+                    expect(logs).toContain("ignored server-fallback reply inputTokens=963300");
+                if (marker === "invalid-json" || marker === "missing-fallback") {
+                    expect(
+                        result.logs.filter((line) =>
+                            line.includes("invalid server-fallback usage marker"),
+                        ),
+                    ).toHaveLength(1);
+                    const repeat = await f.attempt();
+                    expect(repeat.logs.join("\n")).not.toContain(
+                        "invalid server-fallback usage marker",
+                    );
+                }
+            } else {
+                expect(meta.lastInputTokens).toBe(963_300);
+                expect(meta.lastContextPercentage).toBeGreaterThan(95);
+                expect(logs).toMatch(/transform scheduler: .*decision=execute/);
+            }
+        } finally {
+            await f.cleanup();
+        }
+    });
+}

packages/plugin/src/v2/hooks/usage-persist.test.ts

@@ -65,6 +65,104 @@ async function configure272kWindow(): Promise<void> {
 }
 
 describe("persistV2UsageReading", () => {
+    it("records last-iteration pressure instead of billed usage and accepts the next normal reply", () => {
+        process.env.XDG_DATA_HOME = makeTempDir("v2-per-run-usage-");
+        const db = openDatabase();
+        const sessionID = "v2-per-run-usage";
+        const contextUsageMap: TransformDeps["contextUsageMap"] = new Map();
+        const draftModel = { providerID: "anthropic", id: "claude-opus-5-5" };
+        const tokens = { input: 963_300 };
+        const persist = (completed: number, content: unknown) => {
+            const reading = resolveUsageReading({
+                draftModel,
+                tokens,
+                completed,
+                finish: "stop",
+                content,
+                limitFor: () => 872_000,
+            })!;
+            persistV2UsageReading({ db, sessionID, draftModel, reading, contextUsageMap });
+        };
+        persist(1000, [
+            {
+                type: "reasoning",
+                native: {
+                    signature: `cortexkit-server-fallback-usage-v1:${JSON.stringify({
+                        iterations: [
+                            {
+                                type: "message",
+                                model: "primary",
+                                input_tokens: 453_300,
+                                output_tokens: 50,
+                            },
+                            {
+                                type: "fallback_message",
+                                model: "fallback",
+                                input_tokens: 100_000,
+                                output_tokens: 50,
+                                cache_read_input_tokens: 400_000,
+                                cache_creation_input_tokens: 10_000,
+                            },
+                        ],
+                    })}`,
+                },
+            },
+        ]);
+        const meta = getOrCreateSessionMeta(db, sessionID);
+        expect(meta.lastInputTokens).toBe(510_000);
+        expect(meta.lastContextPercentage).toBeCloseTo(58.4862, 3);
+        expect(meta.lastResponseTime).toBe(1000);
+        expect(contextUsageMap.get(sessionID)?.usage.inputTokens).toBe(510_000);
+        expect(contextUsageMap.get(sessionID)?.usage.percentage).toBeCloseTo(58.4862, 3);
+        expect(tokens.input).toBe(963_300);
+        tokens.input = 520_000;
+        persist(2000, []);
+        expect(getOrCreateSessionMeta(db, sessionID).lastInputTokens).toBe(520_000);
+        expect(contextUsageMap.get(sessionID)?.usage.inputTokens).toBe(520_000);
+    });
+
+    it("keeps pressure and the map entry for a server-fallback reply but moves completion time", () => {
+        process.env.XDG_DATA_HOME = makeTempDir("v2-fallback-usage-");
+        const db = openDatabase();
+        const sessionID = "v2-fallback-usage";
+        const contextUsageMap: TransformDeps["contextUsageMap"] = new Map();
+        const draftModel = { providerID: "anthropic", id: "claude-opus-5-5" };
+        const persist = (input: number, completed: number, marked: boolean) => {
+            const reading = resolveUsageReading({
+                draftModel,
+                tokens: { input },
+                completed,
+                finish: "stop",
+                content: marked
+                    ? [
+                          {
+                              type: "reasoning",
+                              text: " ",
+                              native: {
+                                  signature: "cortexkit-server-fallback-v1:primary|fallback",
+                              },
+                          },
+                      ]
+                    : [{ type: "reasoning", native: { signature: "real-signature" } }],
+                limitFor: () => 872_000,
+            })!;
+            persistV2UsageReading({ db, sessionID, draftModel, reading, contextUsageMap });
+        };
+        persist(507_000, 1000, false);
+        const previous = getOrCreateSessionMeta(db, sessionID);
+        const entry = contextUsageMap.get(sessionID);
+        persist(963_300, 2000, true);
+        expect(getOrCreateSessionMeta(db, sessionID)).toEqual({
+            ...previous,
+            lastResponseTime: 2000,
+        });
+        expect(contextUsageMap.get(sessionID)).toBe(entry);
+        persist(510_000, 3000, false);
+        expect(getOrCreateSessionMeta(db, sessionID).lastInputTokens).toBe(510_000);
+        expect(contextUsageMap.get(sessionID)?.usage.inputTokens).toBe(510_000);
+        expect(getOrCreateSessionMeta(db, sessionID).lastResponseTime).toBe(3000);
+    });
+
     it("dates provider usage by completion when the store is read after a fold", () => {
         process.env.XDG_DATA_HOME = makeTempDir("v2-usage-age-");
         const db = openDatabase();

packages/plugin/src/v2/hooks/usage-persist.ts

@@ -22,7 +22,7 @@ export interface PersistV2UsageReadingArgs {
 /**
  * Record the usage OpenCode stored for the latest completed reply.
  *
- * The reading is the prompt size of a request the provider accepted, so it is
+ * An unmarked `reading` is the prompt size of a request the provider accepted, so it is
  * real pressure at any size. It is never compared with the configured window:
  * that window can be smaller than what the model actually serves, and a reading
  * past it is real overflow of the user's limit for the scheduler to handle.
@@ -68,6 +68,15 @@ export function persistV2UsageReading(args: PersistV2UsageReadingArgs): void {
         !providerResponseFailed(completion) &&
         (reading.inputTokens > 0 || isSuccessfulProviderCompletion(completion));
     const responseTime = served && reading.completed;
+    // Without usage for each run, the server sum does not measure one context size.
+    if (reading.serverFallback) {
+        if (responseTime) updateSessionMeta(db, sessionID, { lastResponseTime: responseTime });
+        sessionLog(
+            sessionID,
+            `v2 usage: ignored server-fallback reply inputTokens=${reading.inputTokens} (${reading.serverFallbackInvalidUsage ? "invalid server-fallback usage marker" : "sum of two server runs"}); keeping previous reading`,
+        );
+        return;
+    }
     if (reading.inputTokens <= 0 || !Number.isFinite(reading.limit) || reading.limit <= 0) {
         if (responseTime) updateSessionMeta(db, sessionID, { lastResponseTime: responseTime });
         return;