diff --git a/packages/plugin/src/shared/provider-response-completion.test.ts b/packages/plugin/src/shared/provider-response-completion.test.ts new file mode 100644 index 000000000..e052beef0 --- /dev/null +++ b/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); + } +}); diff --git a/packages/plugin/src/shared/provider-response-completion.ts b/packages/plugin/src/shared/provider-response-completion.ts index 22e0c15c3..1e98dea87 100644 --- a/packages/plugin/src/shared/provider-response-completion.ts +++ b/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((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; diff --git a/packages/plugin/src/v2/hooks/context.ts b/packages/plugin/src/v2/hooks/context.ts index ce8f77afd..f988f6c02 100644 --- a/packages/plugin/src/v2/hooks/context.ts +++ b/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: diff --git a/packages/plugin/src/v2/hooks/server-fallback-usage.test.ts b/packages/plugin/src/v2/hooks/server-fallback-usage.test.ts new file mode 100644 index 000000000..314692a22 --- /dev/null +++ b/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(); + } + }); +} diff --git a/packages/plugin/src/v2/hooks/usage-persist.test.ts b/packages/plugin/src/v2/hooks/usage-persist.test.ts index 49f06e3bb..4cc63052a 100644 --- a/packages/plugin/src/v2/hooks/usage-persist.test.ts +++ b/packages/plugin/src/v2/hooks/usage-persist.test.ts @@ -65,6 +65,104 @@ async function configure272kWindow(): Promise { } 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(); diff --git a/packages/plugin/src/v2/hooks/usage-persist.ts b/packages/plugin/src/v2/hooks/usage-persist.ts index e1ea6447a..2d736b447 100644 --- a/packages/plugin/src/v2/hooks/usage-persist.ts +++ b/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; diff --git a/packages/plugin/src/v2/hooks/usage-reading.test.ts b/packages/plugin/src/v2/hooks/usage-reading.test.ts index c64140624..4ad0815ae 100644 --- a/packages/plugin/src/v2/hooks/usage-reading.test.ts +++ b/packages/plugin/src/v2/hooks/usage-reading.test.ts @@ -4,6 +4,88 @@ import { resolveUsageReading, usageReadingMatchesDraft } from "./usage-reading"; const windows: Record = { old: 200_000, new: 1_000_000 }; const limitFor = (_providerID: string, modelID: string) => windows[modelID] ?? 0; +for (const location of ["native", "state"]) { + for (const cache of [true, false]) { + test(`reads the last fallback iteration from ${location} with cache=${cache}`, () => { + const tokens = { input: 963_300, output: 100, cache: { read: 1, write: 2 } }; + const reading = resolveUsageReading({ + draftModel: { providerID: "p", id: "new" }, + tokens, + completed: 123, + content: [ + { + type: "reasoning", + [location]: { + signature: `cortexkit-server-fallback-usage-v1:${JSON.stringify({ + iterations: [ + { + type: "message", + model: "primary", + input_tokens: 400_000, + output_tokens: 10, + }, + { + type: "fallback_message", + model: "fallback", + input_tokens: 500_000, + output_tokens: 20, + }, + { + type: "fallback_message", + model: "fallback", + input_tokens: 100_000, + output_tokens: 30, + ...(cache + ? { + cache_read_input_tokens: 400_000, + cache_creation_input_tokens: 10_000, + } + : {}), + }, + { + type: "message", + model: "primary", + input_tokens: 900_000, + output_tokens: 40, + }, + ], + })}`, + }, + }, + ], + limitFor, + }); + expect(reading?.inputTokens).toBe(cache ? 510_000 : 100_000); + expect(reading?.serverFallback).not.toBe(true); + expect(tokens).toEqual({ input: 963_300, output: 100, cache: { read: 1, write: 2 } }); + }); + } +} + +for (const payload of [ + "{", + "null", + '{"iterations":[]}', + '{"iterations":[{"type":"message"}]}', + '{"iterations":[{"type":"fallback_message","input_tokens":"bad"}]}', +]) { + test(`invalid usage marker keeps the boundary-only path: ${payload}`, () => { + const reading = resolveUsageReading({ + draftModel: { providerID: "p", id: "new" }, + tokens: { input: 963_300 }, + completed: 123, + content: [ + { + type: "reasoning", + state: { signature: `cortexkit-server-fallback-usage-v1:${payload}` }, + }, + ], + limitFor, + }); + expect(reading?.serverFallback).toBe(true); + }); +} + test("a same-model reading uses one window for attribution and admission", () => { const reading = resolveUsageReading({ rowModel: { providerID: "p", id: "old" }, diff --git a/packages/plugin/src/v2/hooks/usage-reading.ts b/packages/plugin/src/v2/hooks/usage-reading.ts index bdefa8b66..da517a254 100644 --- a/packages/plugin/src/v2/hooks/usage-reading.ts +++ b/packages/plugin/src/v2/hooks/usage-reading.ts @@ -1,4 +1,7 @@ -import { isSuccessfulProviderCompletion } from "../../shared/provider-response-completion"; +import { + isSuccessfulProviderCompletion, + readServerFallbackUsage, +} from "../../shared/provider-response-completion"; export interface UsageReadingInput { rowModel?: { providerID?: unknown; id?: unknown }; @@ -11,6 +14,7 @@ export interface UsageReadingInput { completed?: number; finish?: string; error?: unknown; + content?: unknown; /** Resolve the output-reserved usable window for a model. */ limitFor: (providerID: string, modelID: string) => number; } @@ -23,6 +27,8 @@ export interface UsageReading { admissionLimit: number; /** Absent when a legacy row carries no model metadata. */ modelKey?: string; + serverFallback?: boolean; + serverFallbackInvalidUsage?: boolean; completed?: number; finish?: string; error?: unknown; @@ -58,8 +64,11 @@ export function resolveUsageReading(input: UsageReadingInput): UsageReading | un const rowModelID = typeof input.rowModel?.id === "string" ? input.rowModel.id : undefined; const measuredProviderID = rowProviderID ?? input.draftModel.providerID; const measuredModelID = rowModelID ?? input.draftModel.id; + const fallback = readServerFallbackUsage(input.content); const inputTokens = - numeric(tokens?.input) + numeric(tokens?.cache?.read) + numeric(tokens?.cache?.write); + fallback?.kind === "usage" + ? fallback.inputTokens + : numeric(tokens?.input) + numeric(tokens?.cache?.read) + numeric(tokens?.cache?.write); const limit = input.limitFor(measuredProviderID, measuredModelID); if ((!Number.isFinite(limit) || limit <= 0) && !successfulCompletion) return undefined; const sameModel = @@ -80,5 +89,11 @@ export function resolveUsageReading(input: UsageReadingInput): UsageReading | un : {}), ...(input.finish !== undefined ? { finish: input.finish } : {}), ...(input.error !== undefined ? { error: input.error } : {}), + ...(fallback?.kind === "boundary" + ? { + serverFallback: true, + ...(fallback.invalidUsage ? { serverFallbackInvalidUsage: true } : {}), + } + : {}), }; }