diff --git a/apps/desktop/src/main/__tests__/runtime-host-client-usage.test.ts b/apps/desktop/src/main/__tests__/runtime-host-client-usage.test.ts new file mode 100644 index 0000000000..bdec51dca8 --- /dev/null +++ b/apps/desktop/src/main/__tests__/runtime-host-client-usage.test.ts @@ -0,0 +1,253 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import assert from 'node:assert/strict'; +import test from 'node:test'; +import type { RuntimeHostConnection } from '@maka/runtime-host/client'; +import type { OperationInput, OperationKey } from '@maka/runtime-host/protocol'; +import { + DesktopRuntimeHostClient, + DesktopRuntimeHostClientError, +} from '../runtime-host-client.js'; + +test('loads all Usage snapshot pages behind one start revision', async () => { + const requests: Array<{ operation: OperationKey; input: unknown }> = []; + const client = usageClient(async (operation, input) => { + requests.push({ operation, input }); + assert.equal(operation, 'usage.query'); + if (input.kind === 'snapshot_start') return started('revision-1', 2); + assert.equal(input.revision, 'revision-1'); + if (input.kind === 'snapshot_logs' && input.source === 'llm') { + return input.offset === 0 + ? logPage('revision-1', 'llm', [llmLog('llm-1', 2)], 0, 2, 1, false) + : logPage('revision-1', 'llm', [llmLog('llm-2', 1)], 1, 2, null, false); + } + if (input.kind === 'snapshot_logs') { + return logPage('revision-1', 'tool', [toolLog('tool-1', 3)], 0, 1, null, false); + } + if (input.kind === 'snapshot_pricing') { + return input.offset === 0 + ? pricingPage('revision-1', [pricing('a:model')], 0, 2, 1) + : pricingPage('revision-1', [pricing('b:model')], 1, 2, null); + } + throw new Error('Unexpected Usage request'); + }); + + assert.deepEqual(await client.loadUsageSnapshot({ from: 0, to: 10 }), { + revision: 'revision-1', + summary: validSummary(2), + provenance: validProvenance(), + llmLogs: [llmLog('llm-1', 2), llmLog('llm-2', 1)], + toolLogs: [toolLog('tool-1', 3)], + pricingEntries: [pricing('a:model'), pricing('b:model')], + llmLogsTruncated: false, + toolLogsTruncated: false, + }); + assert.equal( + requests.filter(({ input }) => (input as { kind?: string }).kind === 'snapshot_start').length, + 1, + ); +}); + +test('discards every partial Usage result and restarts after revision_changed', async () => { + let starts = 0; + const client = usageClient(async (_operation, input) => { + if (input.kind === 'snapshot_start') { + starts += 1; + return started(`revision-${starts}`, starts); + } + if (input.revision === 'revision-1' && input.kind === 'snapshot_logs' && input.source === 'llm') { + return { kind: 'revision_changed', expectedRevision: 'revision-1' }; + } + if (input.kind === 'snapshot_logs') { + const row = input.source === 'llm' ? llmLog('fresh-llm', 2) : toolLog('fresh-tool', 1); + return logPage(input.revision, input.source, [row], 0, 1, null, false); + } + if (input.kind === 'snapshot_pricing') { + return pricingPage(input.revision, [pricing('fresh:model')], 0, 1, null); + } + throw new Error('Unexpected Usage request'); + }); + + const snapshot = await client.loadUsageSnapshot('all'); + assert.equal(starts, 2); + assert.equal(snapshot.revision, 'revision-2'); + assert.deepEqual(snapshot.llmLogs.map((row) => row.id), ['fresh-llm']); + assert.deepEqual(snapshot.toolLogs.map((row) => row.id), ['fresh-tool']); +}); + +test('fails with usage_unstable after three complete Usage snapshot attempts', async () => { + let starts = 0; + const client = usageClient(async (_operation, input) => { + if (input.kind === 'snapshot_start') { + starts += 1; + return started(`revision-${starts}`, 0); + } + return { kind: 'revision_changed', expectedRevision: input.revision }; + }); + + await assert.rejects( + () => client.loadUsageSnapshot('all'), + (error: unknown) => + error instanceof DesktopRuntimeHostClientError && error.code === 'usage_unstable', + ); + assert.equal(starts, 3); +}); + +test('rejects non-progressing or identity-changing Usage snapshot pages', async () => { + const client = usageClient(async (_operation, input) => { + if (input.kind === 'snapshot_start') return started('revision-1', 1); + if (input.kind === 'snapshot_logs' && input.source === 'llm') { + return logPage('wrong-revision', 'llm', [llmLog('llm-1', 1)], 0, 2, 0, false); + } + if (input.kind === 'snapshot_logs') { + return logPage('revision-1', 'tool', [], 0, 0, null, false); + } + return pricingPage('revision-1', [], 0, 0, null); + }); + + await assert.rejects( + () => client.loadUsageSnapshot('all'), + (error: unknown) => + error instanceof DesktopRuntimeHostClientError && error.code === 'projection_unstable', + ); +}); + +function usageClient( + respond: (operation: OperationKey, input: any) => Promise, +): DesktopRuntimeHostClient { + const connection = { + hostEpoch: 'host-current', + connectionId: 'connection-current', + rootId: 'root-current', + request: (operation: K, input: OperationInput) => + respond(operation, input), + close: async () => undefined, + } as unknown as RuntimeHostConnection; + return new DesktopRuntimeHostClient(connection); +} + +function started(revision: string, totalRequests: number) { + return { + kind: 'snapshot_started' as const, + revision, + summary: validSummary(totalRequests), + provenance: validProvenance(), + }; +} + +function logPage( + revision: string, + source: 'llm' | 'tool', + rows: readonly unknown[], + offset: number, + total: number, + nextOffset: number | null, + truncated: boolean, +) { + return { kind: 'snapshot_logs' as const, revision, source, rows, offset, total, nextOffset, truncated }; +} + +function pricingPage( + revision: string, + entries: readonly unknown[], + offset: number, + total: number, + nextOffset: number | null, +) { + return { kind: 'snapshot_pricing' as const, revision, entries, offset, total, nextOffset }; +} + +function validSummary(totalRequests: number) { + return { + range: { from: 0, to: 10 }, + totalRequests, + totalCostUsd: 0, + totalTokens: { + input: 0, + output: 0, + cacheMiss: 0, + cacheRead: 0, + cacheWrite: 0, + reasoning: 0, + total: 0, + }, + cacheHitRequests: 0, + cacheCreateRequests: 0, + errorRequests: 0, + }; +} + +function validProvenance() { + return { + coverage: { + attempts: 0, + pricedAttempts: 0, + unpricedAttempts: 0, + usageReportedAttempts: 0, + usagePartialAttempts: 0, + usageMissingAttempts: 0, + }, + legacyRecords: 0, + unreadableRecords: 0, + pendingRepairs: 0, + }; +} + +function llmLog(id: string, ts: number) { + return { + source: 'llm' as const, + id, + ts, + providerId: 'provider', + modelId: 'model', + inputTokens: 1, + outputTokens: 1, + cacheMissTokens: 1, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0, + totalTokens: 2, + costUsd: 0, + latencyMs: 1, + status: 'success' as const, + }; +} + +function toolLog(id: string, ts: number) { + return { + source: 'tool' as const, + id, + ts, + toolName: 'Read', + durationMs: 1, + status: 'success' as const, + bytesIn: 0, + bytesOut: 0, + startedAt: ts, + }; +} + +function pricing(modelKey: string) { + return { + source: 'custom' as const, + resetEffect: 'become_unpriced' as const, + pricing: { modelKey, inputUsdPer1M: 1, outputUsdPer1M: 2 }, + }; +} diff --git a/apps/desktop/src/main/__tests__/runtime-host-usage-ipc-main.test.ts b/apps/desktop/src/main/__tests__/runtime-host-usage-ipc-main.test.ts index 6199baa55b..6a26a749a9 100644 --- a/apps/desktop/src/main/__tests__/runtime-host-usage-ipc-main.test.ts +++ b/apps/desktop/src/main/__tests__/runtime-host-usage-ipc-main.test.ts @@ -20,77 +20,31 @@ import assert from "node:assert/strict"; import { test } from "node:test"; import type { UsageStats } from "@maka/core/settings"; -import type { UsageQueryInput, UsageQueryResult } from "@maka/runtime-host/protocol"; import type { IpcHandler } from "../ipc-reconnect-policy.js"; -import type { DesktopRuntimeHostClient } from "../runtime-host-client.js"; +import { + DesktopRuntimeHostClientError, + type DesktopRuntimeHostClient, +} from "../runtime-host-client.js"; import { registerRuntimeHostUsageIpc } from "../runtime-host-usage-ipc-main.js"; test("settings usage stats use the canonical model-call total and load every activity page", async () => { const handlers = new Map(); - const calls: Array<{ source?: "llm" | "tool"; offset?: number }> = []; - const ranges: UsageQueryInput["query"]["range"][] = []; + const ranges: unknown[] = []; registerRuntimeHostUsageIpc({ ipcMain: { handle: (channel, listener) => handlers.set(channel, listener), handleReconnectableRead: (channel, listener) => handlers.set(channel, listener), }, client: { - queryUsage: async (input: UsageQueryInput) => { - ranges.push(input.query.range); - if (input.kind === "summary") { - return { - kind: "summary", - summary: { - range: { from: 1, to: 2 }, - totalRequests: 151, - totalCostUsd: 12.5, - totalTokens: { - input: 3_000_000, - output: 500_000, - cacheMiss: 100_000, - cacheRead: 400_000, - cacheWrite: 43_090, - reasoning: 90, - total: 4_043_090, - }, - cacheHitRequests: 10, - cacheCreateRequests: 5, - errorRequests: 2, - }, - provenance: provenance(), - } satisfies UsageQueryResult; - } - if (input.kind !== "logs") throw new Error("unexpected usage query"); - calls.push({ source: input.source, offset: input.offset }); - if (input.source === "llm") { - const offset = input.offset ?? 0; - const count = offset === 0 ? 100 : 51; - return { - kind: "logs", - source: "llm", - rows: Array.from({ length: count }, (_, index) => llmRow(offset + index)), - offset, - total: 151, - nextOffset: offset === 0 ? 100 : null, - provenance: provenance(), - } satisfies UsageQueryResult; - } - const offset = input.offset ?? 0; - const count = offset === 0 ? 100 : 71; + loadUsageSnapshot: async (range: unknown) => { + ranges.push(range); return { - kind: "logs", - source: "tool", - rows: Array.from({ length: count }, (_, index) => toolRow(offset + index)), - offset, - total: 171, - nextOffset: offset === 0 ? 100 : null, - } satisfies UsageQueryResult; - }, - loadPricingSnapshot: async () => ({ - hostEpoch: "host-epoch", - connectionId: "connection-id", - revision: 1, - entries: [ + revision: "snapshot-1", + summary: usageSummary(151), + provenance: provenance(), + llmLogs: Array.from({ length: 151 }, (_, index) => llmRow(index)), + toolLogs: Array.from({ length: 171 }, (_, index) => toolRow(index)), + pricingEntries: [ { source: "custom", resetEffect: "become_unpriced", @@ -100,8 +54,11 @@ test("settings usage stats use the canonical model-call total and load every act outputUsdPer1M: 2, }, }, - ], - }), + ], + llmLogsTruncated: false, + toolLogsTruncated: false, + }; + }, } as unknown as DesktopRuntimeHostClient, sendToRenderer: () => undefined, }); @@ -115,15 +72,8 @@ test("settings usage stats use the canonical model-call total and load every act assert.equal(stats.logs.length, 322); assert.equal(stats.logs.filter((row) => row.kind === "model").length, 151); assert.equal(stats.logs.filter((row) => row.kind === "tool").length, 171); - const expectedCalls: Array<{ source?: "llm" | "tool"; offset?: number }> = [ - { source: "llm", offset: 0 }, - { source: "llm", offset: 100 }, - { source: "tool", offset: 0 }, - { source: "tool", offset: 100 }, - ]; - assert.deepEqual(calls.sort(compareCall), expectedCalls.sort(compareCall)); - assert.ok(ranges.every((range) => typeof range === "object")); - assert.ok(ranges.every((range) => JSON.stringify(range) === JSON.stringify(ranges[0]))); + assert.equal(ranges.length, 1); + assert.equal(typeof ranges[0], "object"); assert.equal(stats.logs.find((row) => row.id === "llm-150")?.status, "aborted"); assert.equal(stats.logs.find((row) => row.id === "llm-150")?.sessionId, undefined); assert.equal(stats.logs.find((row) => row.id === "llm-150")?.costUsd, undefined); @@ -148,7 +98,7 @@ test("settings usage stats use the canonical model-call total and load every act assert.equal(stats.logsTruncated, undefined); }); -test("settings usage stats reject a non-advancing activity page", async () => { +test("settings usage stats propagate an invalid snapshot projection", async () => { const handlers = new Map(); registerRuntimeHostUsageIpc({ ipcMain: { @@ -156,66 +106,26 @@ test("settings usage stats reject a non-advancing activity page", async () => { handleReconnectableRead: (channel, listener) => handlers.set(channel, listener), }, client: { - queryUsage: async (input: UsageQueryInput) => { - if (input.kind === "summary") { - return { - kind: "summary", - summary: { - range: { from: 1, to: 2 }, - totalRequests: 0, - totalCostUsd: 0, - totalTokens: { - input: 0, - output: 0, - cacheMiss: 0, - cacheRead: 0, - cacheWrite: 0, - reasoning: 0, - total: 0, - }, - cacheHitRequests: 0, - cacheCreateRequests: 0, - errorRequests: 0, - }, - provenance: provenance(), - } satisfies UsageQueryResult; - } - if (input.kind !== "logs") throw new Error("unexpected usage query"); - return input.source === "llm" - ? ({ - kind: "logs", - source: "llm", - rows: [], - offset: 0, - total: 1, - nextOffset: 0, - provenance: provenance(), - } satisfies UsageQueryResult) - : ({ - kind: "logs", - source: "tool", - rows: [], - offset: 0, - total: 0, - nextOffset: null, - } satisfies UsageQueryResult); + loadUsageSnapshot: async () => { + throw new DesktopRuntimeHostClientError( + "projection_unstable", + "Runtime Host returned an invalid Usage snapshot projection", + ); }, - loadPricingSnapshot: async () => ({ - hostEpoch: "host-epoch", - connectionId: "connection-id", - revision: 0, - entries: [], - }), } as unknown as DesktopRuntimeHostClient, sendToRenderer: () => undefined, }); const handler = handlers.get("settings:usageStats"); assert.ok(handler); - await assert.rejects(() => handler({} as never, "24h"), /invalid Usage projection/); + await assert.rejects( + () => handler({} as never, "24h"), + (error: unknown) => + error instanceof DesktopRuntimeHostClientError && error.code === "projection_unstable", + ); }); -test("settings usage stats degrade instead of erroring when logs disagree with the canonical summary", async () => { +test("settings usage stats fail when a coherent snapshot cannot be retained", async () => { const handlers = new Map(); registerRuntimeHostUsageIpc({ ipcMain: { @@ -223,69 +133,23 @@ test("settings usage stats degrade instead of erroring when logs disagree with t handleReconnectableRead: (channel, listener) => handlers.set(channel, listener), }, client: { - queryUsage: async (input: UsageQueryInput) => { - if (input.kind === "summary") { - return { - kind: "summary", - summary: { - range: { from: 1, to: 2 }, - totalRequests: 2, - totalCostUsd: 0, - totalTokens: { - input: 0, - output: 0, - cacheMiss: 0, - cacheRead: 0, - cacheWrite: 0, - reasoning: 0, - total: 0, - }, - cacheHitRequests: 0, - cacheCreateRequests: 0, - errorRequests: 0, - }, - provenance: provenance(), - } satisfies UsageQueryResult; - } - if (input.kind !== "logs") throw new Error("unexpected usage query"); - return input.source === "llm" - ? ({ - kind: "logs", - source: "llm", - rows: [llmRow(0)], - offset: 0, - total: 1, - nextOffset: null, - provenance: provenance(), - } satisfies UsageQueryResult) - : ({ - kind: "logs", - source: "tool", - rows: [], - offset: 0, - total: 0, - nextOffset: null, - } satisfies UsageQueryResult); + loadUsageSnapshot: async () => { + throw new DesktopRuntimeHostClientError( + "usage_unstable", + "Usage snapshot kept expiring while Desktop read it", + ); }, - loadPricingSnapshot: async () => ({ - hostEpoch: "host-epoch", - connectionId: "connection-id", - revision: 0, - entries: [], - }), } as unknown as DesktopRuntimeHostClient, sendToRenderer: () => undefined, }); const handler = handlers.get("settings:usageStats"); assert.ok(handler); - // A catch-up race (summary read before a repair commits, logs read after) must - // not error the whole page. The canonical summary total stays authoritative, - // the activity list holds what actually loaded, and provenance still rides along. - const stats = await handler({} as never, "all") as UsageStats; - assert.equal(stats.summary.totalRequests, 2); - assert.equal(stats.logs.filter((row) => row.kind === "model").length, 1); - assert.deepEqual(stats.provenance, provenance()); + await assert.rejects( + () => handler({} as never, "all"), + (error: unknown) => + error instanceof DesktopRuntimeHostClientError && error.code === "usage_unstable", + ); }); test("settings usage stats group the provider breakdown by connection", async () => { @@ -296,59 +160,18 @@ test("settings usage stats group the provider breakdown by connection", async () handleReconnectableRead: (channel, listener) => handlers.set(channel, listener), }, client: { - queryUsage: async (input: UsageQueryInput) => { - if (input.kind === "summary") { - return { - kind: "summary", - summary: { - range: { from: 1, to: 2 }, - totalRequests: 2, - totalCostUsd: 0, - totalTokens: { - input: 0, - output: 0, - cacheMiss: 0, - cacheRead: 0, - cacheWrite: 0, - reasoning: 0, - total: 0, - }, - cacheHitRequests: 0, - cacheCreateRequests: 0, - errorRequests: 0, - }, - provenance: provenance(), - } satisfies UsageQueryResult; - } - if (input.kind !== "logs") throw new Error("unexpected usage query"); - // Two connections to the SAME provider type must stay two rows. - return input.source === "llm" - ? ({ - kind: "logs", - source: "llm", - rows: [ - { ...llmRow(0), connectionSlug: "conn-a", providerId: "provider-x" }, - { ...llmRow(1), connectionSlug: "conn-b", providerId: "provider-x" }, - ], - offset: 0, - total: 2, - nextOffset: null, - provenance: provenance(), - } satisfies UsageQueryResult) - : ({ - kind: "logs", - source: "tool", - rows: [], - offset: 0, - total: 0, - nextOffset: null, - } satisfies UsageQueryResult); - }, - loadPricingSnapshot: async () => ({ - hostEpoch: "host-epoch", - connectionId: "connection-id", - revision: 0, - entries: [], + loadUsageSnapshot: async () => ({ + revision: "snapshot-1", + summary: usageSummary(2), + provenance: provenance(), + llmLogs: [ + { ...llmRow(0), connectionSlug: "conn-a", providerId: "provider-x" }, + { ...llmRow(1), connectionSlug: "conn-b", providerId: "provider-x" }, + ], + toolLogs: [], + pricingEntries: [], + llmLogsTruncated: false, + toolLogsTruncated: false, }), } as unknown as DesktopRuntimeHostClient, sendToRenderer: () => undefined, @@ -365,68 +188,22 @@ test("settings usage stats group the provider breakdown by connection", async () test("settings usage stats truncate the activity log at the cap instead of erroring", async () => { const handlers = new Map(); - const PAGE = 100; - // Above MAX_ACTIVITY_RECORDS (50_000) so paging must stop and flag truncation. - const TOTAL = 50_150; + const TOTAL = 50_000; registerRuntimeHostUsageIpc({ ipcMain: { handle: (channel, listener) => handlers.set(channel, listener), handleReconnectableRead: (channel, listener) => handlers.set(channel, listener), }, client: { - queryUsage: async (input: UsageQueryInput) => { - if (input.kind === "summary") { - return { - kind: "summary", - summary: { - range: { from: 1, to: 2 }, - totalRequests: TOTAL, - totalCostUsd: 0, - totalTokens: { - input: 0, - output: 0, - cacheMiss: 0, - cacheRead: 0, - cacheWrite: 0, - reasoning: 0, - total: 0, - }, - cacheHitRequests: 0, - cacheCreateRequests: 0, - errorRequests: 0, - }, - provenance: provenance(), - } satisfies UsageQueryResult; - } - if (input.kind !== "logs") throw new Error("unexpected usage query"); - if (input.source === "llm") { - const offset = input.offset ?? 0; - const count = Math.min(PAGE, TOTAL - offset); - const nextOffset = offset + count < TOTAL ? offset + count : null; - return { - kind: "logs", - source: "llm", - rows: Array.from({ length: count }, (_, index) => llmRow(offset + index)), - offset, - total: TOTAL, - nextOffset, - provenance: provenance(), - } satisfies UsageQueryResult; - } - return { - kind: "logs", - source: "tool", - rows: [], - offset: 0, - total: 0, - nextOffset: null, - } satisfies UsageQueryResult; - }, - loadPricingSnapshot: async () => ({ - hostEpoch: "host-epoch", - connectionId: "connection-id", - revision: 0, - entries: [], + loadUsageSnapshot: async () => ({ + revision: "snapshot-1", + summary: usageSummary(TOTAL + 150), + provenance: provenance(), + llmLogs: Array.from({ length: TOTAL }, (_, index) => llmRow(index)), + toolLogs: [], + pricingEntries: [], + llmLogsTruncated: true, + toolLogsTruncated: false, }), } as unknown as DesktopRuntimeHostClient, sendToRenderer: () => undefined, @@ -476,6 +253,26 @@ function toolRow(index: number) { }; } +function usageSummary(totalRequests: number) { + return { + range: { from: 1, to: 2 }, + totalRequests, + totalCostUsd: 12.5, + totalTokens: { + input: 3_000_000, + output: 500_000, + cacheMiss: 100_000, + cacheRead: 400_000, + cacheWrite: 43_090, + reasoning: 90, + total: 4_043_090, + }, + cacheHitRequests: 10, + cacheCreateRequests: 5, + errorRequests: 2, + }; +} + function provenance() { return { coverage: { @@ -491,10 +288,3 @@ function provenance() { pendingRepairs: 0, }; } - -function compareCall( - left: { source?: "llm" | "tool"; offset?: number }, - right: { source?: "llm" | "tool"; offset?: number }, -): number { - return `${left.source}:${left.offset}`.localeCompare(`${right.source}:${right.offset}`); -} diff --git a/apps/desktop/src/main/runtime-host-client.ts b/apps/desktop/src/main/runtime-host-client.ts index e07d625122..e042b65187 100644 --- a/apps/desktop/src/main/runtime-host-client.ts +++ b/apps/desktop/src/main/runtime-host-client.ts @@ -27,6 +27,7 @@ import { } from "@maka/core/session"; import { markPersisted } from "@maka/core/persisted-value"; import type { Task } from "@maka/core/task-ledger"; +import type { UsageProvenance } from "@maka/core/usage-ledger-merge"; import type { ConnectionCatalogSnapshot, @@ -40,7 +41,7 @@ import { canonicalPricingConfigsEqual, comparePricingModelKeys, } from "@maka/core/usage-stats/pricing"; -import type { PricingConfig } from "@maka/core/usage-stats/types"; +import type { PricingConfig, TimeRange, UsageSummaryV2 } from "@maka/core/usage-stats/types"; import { type ClientCapabilityProvider, type DecodedSessionTranscriptPage, @@ -134,6 +135,10 @@ import { type TurnInterruptResult, type TurnMessageSubmitInput, type TurnMessageSubmitResult, + type LlmUsageLogProjection, + type ToolUsageLogProjection, + PRICING_PAGE_MAX_ITEMS, + USAGE_PAGE_MAX_ITEMS, type WorkspaceProjection, } from "@maka/runtime-host/protocol"; @@ -142,6 +147,8 @@ const decodeStoredMessage = (value: unknown): StoredMessage => const MAX_OPTIMISTIC_ATTEMPTS = 3; const MAX_SESSION_REVISION_ATTEMPTS = 8; const MAX_PRICING_SNAPSHOT_ATTEMPTS = 3; +const MAX_USAGE_SNAPSHOT_ATTEMPTS = 3; +const MAX_USAGE_SNAPSHOT_ACTIVITY_RECORDS = 50_000; export type DesktopSessionConfigurationPatch = Partial; @@ -160,6 +167,7 @@ export type DesktopRuntimeHostClientErrorCode = | "revision_conflict" | "session_not_found" | "skill_catalog_unstable" + | "usage_unstable" | "unsupported_session"; export class DesktopRuntimeHostClientError extends Error { @@ -202,6 +210,17 @@ export interface DesktopPricingSnapshot { readonly entries: readonly EffectivePricingEntry[]; } +export interface DesktopUsageSnapshot { + readonly revision: string; + readonly summary: UsageSummaryV2; + readonly provenance: UsageProvenance; + readonly llmLogs: readonly LlmUsageLogProjection[]; + readonly toolLogs: readonly ToolUsageLogProjection[]; + readonly pricingEntries: readonly EffectivePricingEntry[]; + readonly llmLogsTruncated: boolean; + readonly toolLogsTruncated: boolean; +} + export interface DesktopSkillCatalogSnapshot { readonly revision: SkillCatalogRevision; readonly view: SkillCatalogView; @@ -1264,6 +1283,17 @@ export class DesktopRuntimeHostClient { return this.request("usage.query", input); } + async loadUsageSnapshot(range: TimeRange): Promise { + for (let attempt = 0; attempt < MAX_USAGE_SNAPSHOT_ATTEMPTS; attempt += 1) { + const snapshot = await this.#readUsageSnapshot(range); + if (snapshot) return snapshot; + } + throw new DesktopRuntimeHostClientError( + "usage_unstable", + "Usage snapshot kept expiring while Desktop read it", + ); + } + queryGoal(sessionId: string): Promise> { return this.request("goal.query", { sessionId }); } @@ -1605,6 +1635,147 @@ export class DesktopRuntimeHostClient { }; } + async #readUsageSnapshot(range: TimeRange): Promise { + this.#assertOpen(); + const started = await this.request("usage.query", { kind: "snapshot_start", range }); + if ( + started.kind !== "snapshot_started" || + (typeof range === "object" && + (started.summary.range.from !== range.from || started.summary.range.to !== range.to)) + ) { + throw invalidProjection("Usage snapshot start"); + } + const [llm, tool, pricing] = await Promise.all([ + this.#readUsageSnapshotLogs(started.revision, "llm"), + this.#readUsageSnapshotLogs(started.revision, "tool"), + this.#readUsageSnapshotPricing(started.revision), + ]); + if (!llm || !tool || !pricing) return undefined; + return { + revision: started.revision, + summary: started.summary, + provenance: started.provenance, + llmLogs: llm.rows, + toolLogs: tool.rows, + pricingEntries: pricing, + llmLogsTruncated: llm.truncated, + toolLogsTruncated: tool.truncated, + }; + } + + async #readUsageSnapshotLogs( + revision: string, + source: "llm", + ): Promise<{ readonly rows: readonly LlmUsageLogProjection[]; readonly truncated: boolean } | undefined>; + async #readUsageSnapshotLogs( + revision: string, + source: "tool", + ): Promise<{ readonly rows: readonly ToolUsageLogProjection[]; readonly truncated: boolean } | undefined>; + async #readUsageSnapshotLogs( + revision: string, + source: "llm" | "tool", + ): Promise< + | { + readonly rows: readonly (LlmUsageLogProjection | ToolUsageLogProjection)[]; + readonly truncated: boolean; + } + | undefined + > { + const rows: Array = []; + let offset = 0; + let total: number | undefined; + let truncated: boolean | undefined; + while (true) { + const page = await this.request("usage.query", { + kind: "snapshot_logs", + revision, + source, + offset, + limit: USAGE_PAGE_MAX_ITEMS, + }); + if (page.kind === "revision_changed") { + if (page.expectedRevision !== revision) throw invalidProjection("Usage snapshot revision"); + return undefined; + } + if ( + page.kind !== "snapshot_logs" || + page.revision !== revision || + page.source !== source || + page.offset !== offset || + page.rows.length > USAGE_PAGE_MAX_ITEMS || + page.total > MAX_USAGE_SNAPSHOT_ACTIVITY_RECORDS + ) { + throw invalidProjection("Usage snapshot logs"); + } + total ??= page.total; + truncated ??= page.truncated; + if (page.total !== total || page.truncated !== truncated || rows.length !== offset) { + throw invalidProjection("Usage snapshot logs"); + } + rows.push(...page.rows); + if (rows.length > total) throw invalidProjection("Usage snapshot logs"); + if (page.nextOffset === null) { + if (rows.length !== total) throw invalidProjection("Usage snapshot logs"); + return { rows, truncated }; + } + if ( + page.rows.length === 0 || + page.nextOffset !== offset + page.rows.length || + page.nextOffset >= total + ) { + throw invalidProjection("Usage snapshot logs"); + } + offset = page.nextOffset; + } + } + + async #readUsageSnapshotPricing( + revision: string, + ): Promise { + const entries: EffectivePricingEntry[] = []; + let offset = 0; + let total: number | undefined; + while (true) { + const page = await this.request("usage.query", { + kind: "snapshot_pricing", + revision, + offset, + limit: PRICING_PAGE_MAX_ITEMS, + }); + if (page.kind === "revision_changed") { + if (page.expectedRevision !== revision) throw invalidProjection("Usage snapshot revision"); + return undefined; + } + if ( + page.kind !== "snapshot_pricing" || + page.revision !== revision || + page.offset !== offset || + page.entries.length > PRICING_PAGE_MAX_ITEMS || + entries.length !== offset + ) { + throw invalidProjection("Usage snapshot pricing"); + } + total ??= page.total; + if (page.total !== total) throw invalidProjection("Usage snapshot pricing"); + entries.push(...page.entries); + if (entries.length > total) throw invalidProjection("Usage snapshot pricing"); + if (page.nextOffset === null) { + if (entries.length !== total || !pricingEntriesAreCanonical(entries)) { + throw invalidProjection("Usage snapshot pricing"); + } + return entries; + } + if ( + page.entries.length === 0 || + page.nextOffset !== offset + page.entries.length || + page.nextOffset >= total + ) { + throw invalidProjection("Usage snapshot pricing"); + } + offset = page.nextOffset; + } + } + async #reconcilePricingMutation( target: PricingReconciliationTarget, reason: "revision_conflict" | "outcome_unknown", diff --git a/apps/desktop/src/main/runtime-host-usage-ipc-main.ts b/apps/desktop/src/main/runtime-host-usage-ipc-main.ts index 91a162a0ad..cf46b7e1b2 100644 --- a/apps/desktop/src/main/runtime-host-usage-ipc-main.ts +++ b/apps/desktop/src/main/runtime-host-usage-ipc-main.ts @@ -26,7 +26,6 @@ import { } from "@maka/core/usage-stats/pricing"; import type { PricingConfig, - TimeRange, UsageGroupBy, UsageQuery, } from "@maka/core/usage-stats/types"; @@ -48,8 +47,6 @@ interface RuntimeHostUsageIpcDeps { readonly sendToRenderer: (channel: string, ...args: unknown[]) => void; } -const MAX_ACTIVITY_RECORDS = 50_000; - export function registerRuntimeHostUsageIpc( deps: RuntimeHostUsageIpcDeps, ): void { @@ -160,39 +157,24 @@ async function loadUsageStats( client: DesktopRuntimeHostClient, range: UsageRange, ): Promise { - const query = { range: resolveUsageRange(range, Date.now()) } satisfies UsageQuery; - const [summaryResult, llmResult, toolResult, pricing] = await Promise.all([ - client.queryUsage({ kind: "summary", query }), - loadAllLogs(client, "llm", query), - loadAllLogs(client, "tool", query), - client.loadPricingSnapshot(), - ]); - if (summaryResult.kind !== "summary") throw invalidUsageProjection(); - const llmLogs = llmResult.rows; - const toolLogs = toolResult.rows; - const logsTruncated = llmResult.truncated || toolResult.truncated; - // The canonical summary is the authoritative headline count. We no longer - // throw when it disagrees with the number of activity rows we managed to - // load: a Host restart with pending repairs can make the summary read land - // before a catch-up commits and the logs read land after, and truncation - // (above) deliberately shortens the list. Either way the summary total stays - // correct; `provenance`/`logsTruncated` tell the page the activity list may - // be incomplete instead of erroring the whole page. + const snapshot = await client.loadUsageSnapshot(resolveUsageRange(range, Date.now())); + const llmLogs = snapshot.llmLogs; + const toolLogs = snapshot.toolLogs; + const logsTruncated = snapshot.llmLogsTruncated || snapshot.toolLogsTruncated; return { summary: { - totalRequests: summaryResult.summary.totalRequests, - totalCostUsd: summaryResult.summary.totalCostUsd, - totalTokens: summaryResult.summary.totalTokens.total, - inputTokens: summaryResult.summary.totalTokens.input, - outputTokens: summaryResult.summary.totalTokens.output, + totalRequests: snapshot.summary.totalRequests, + totalCostUsd: snapshot.summary.totalCostUsd, + totalTokens: snapshot.summary.totalTokens.total, + inputTokens: snapshot.summary.totalTokens.input, + outputTokens: snapshot.summary.totalTokens.output, cacheTokens: - summaryResult.summary.totalTokens.cacheRead + - summaryResult.summary.totalTokens.cacheWrite, - cacheMiss: summaryResult.summary.totalTokens.cacheMiss, - cacheRead: summaryResult.summary.totalTokens.cacheRead, - cacheCreation: summaryResult.summary.totalTokens.cacheWrite, - reasoning: summaryResult.summary.totalTokens.reasoning, + snapshot.summary.totalTokens.cacheRead + snapshot.summary.totalTokens.cacheWrite, + cacheMiss: snapshot.summary.totalTokens.cacheMiss, + cacheRead: snapshot.summary.totalTokens.cacheRead, + cacheCreation: snapshot.summary.totalTokens.cacheWrite, + reasoning: snapshot.summary.totalTokens.reasoning, }, logs: [...llmLogs.map(projectLlmLog), ...toolLogs.map(projectToolLog)].sort( (left, right) => right.ts - left.ts, @@ -200,81 +182,18 @@ async function loadUsageStats( byProvider: aggregateModelLogs(llmLogs, "provider"), byModel: aggregateModelLogs(llmLogs, "model"), byTool: aggregateToolLogs(toolLogs), - pricing: pricing.entries + pricing: snapshot.pricingEntries .filter((entry) => entry.source === "custom") .map(({ pricing: entry }) => projectPricing(entry)) .sort( (left, right) => left.provider.localeCompare(right.provider) || left.model.localeCompare(right.model), ), - provenance: summaryResult.provenance, + provenance: snapshot.provenance, ...(logsTruncated ? { logsTruncated: true } : {}), }; } -async function loadAllLogs( - client: DesktopRuntimeHostClient, - source: "llm", - query: UsageQuery & { range: TimeRange }, -): Promise<{ rows: LlmUsageLogProjection[]; truncated: boolean }>; -async function loadAllLogs( - client: DesktopRuntimeHostClient, - source: "tool", - query: UsageQuery & { range: TimeRange }, -): Promise<{ rows: ToolUsageLogProjection[]; truncated: boolean }>; -async function loadAllLogs( - client: DesktopRuntimeHostClient, - source: "llm" | "tool", - query: UsageQuery & { range: TimeRange }, -): Promise<{ - rows: Array; - truncated: boolean; -}> { - const rows: Array = []; - let offset = 0; - let total: number | undefined; - while (true) { - const result = await client.queryUsage( - source === "llm" - ? { - kind: "logs", - source, - query: toLlmQuery(query), - offset, - limit: USAGE_PAGE_MAX_ITEMS, - } - : { - kind: "logs", - source, - query: toToolQuery(query), - offset, - limit: USAGE_PAGE_MAX_ITEMS, - }, - ); - if (result.kind !== "logs" || result.source !== source || result.offset !== offset) { - throw invalidUsageProjection(); - } - total ??= result.total; - if (result.total !== total) throw invalidUsageProjection(); - rows.push(...result.rows); - // Structural integrity: the Host must never return more rows than it claims. - if (rows.length > total) throw invalidUsageProjection(); - // Client-side cap: when a range holds more activity than we render, keep the - // newest MAX_ACTIVITY_RECORDS and stop paging. This is truncation, not a - // protocol error, and the exhaustiveness check below is skipped for it — the - // caller surfaces `logsTruncated` so the page can say the list is partial. - if (total > MAX_ACTIVITY_RECORDS && rows.length >= MAX_ACTIVITY_RECORDS) { - return { rows: rows.slice(0, MAX_ACTIVITY_RECORDS), truncated: true }; - } - if (result.nextOffset === null) { - if (rows.length !== total) throw invalidUsageProjection(); - return { rows, truncated: false }; - } - if (result.nextOffset <= offset) throw invalidUsageProjection(); - offset = result.nextOffset; - } -} - function projectLlmLog(row: LlmUsageLogProjection): UsageStats["logs"][number] { return { id: row.id, diff --git a/packages/runtime-host/src/__tests__/handshake-compatibility.test.ts b/packages/runtime-host/src/__tests__/handshake-compatibility.test.ts index 69e9dddd9d..fbe9f208fa 100644 --- a/packages/runtime-host/src/__tests__/handshake-compatibility.test.ts +++ b/packages/runtime-host/src/__tests__/handshake-compatibility.test.ts @@ -108,7 +108,7 @@ test('receives structured incompatibility guidance from the released v0.1.11 Hos ); }); -test('rejects an epoch-39 Host before any domain command', async () => { +test('rejects an epoch-59 Host before any domain command', async () => { let admittedRequest: RequestFrame | undefined; await withForgedHandshakePeer( async (transport, hostEpoch, rootId) => { @@ -120,7 +120,7 @@ test('rejects an epoch-39 Host before any domain command', async () => { hostEpoch, connectionId: 'forged-epoch-connection', selectedProtocol: RUNTIME_HOST_PROTOCOL_VERSION, - compatibilityEpoch: 39, + compatibilityEpoch: 59, compositionId: 'maka.interactive', compositionRevision: '1', state: 'ready', diff --git a/packages/runtime-host/src/__tests__/usage-pricing-protocol.test.ts b/packages/runtime-host/src/__tests__/usage-pricing-protocol.test.ts index 9a5ed0dba2..66d6020ebe 100644 --- a/packages/runtime-host/src/__tests__/usage-pricing-protocol.test.ts +++ b/packages/runtime-host/src/__tests__/usage-pricing-protocol.test.ts @@ -50,6 +50,7 @@ import { type EffectivePricingEntry, type LlmUsageLogProjection, type ToolUsageLogProjection, + type UsageQueryResult, } from '../protocol/index.js'; import type { ConnectionContext } from '../server/operation-dispatcher.js'; import { HostUsagePricingCoordinator } from '../server/usage-pricing-coordinator.js'; @@ -171,6 +172,130 @@ describe('Usage/Pricing protocol', () => { ]) { assert.throws(() => usageRequest(input), invalidFrame); } + for (const result of [ + { + kind: 'snapshot_started', + revision: 'snapshot-revision-1', + summary: validSummary(), + provenance: validProvenance(), + extra: true, + }, + { + kind: 'snapshot_logs', + revision: 'snapshot-revision-1', + source: 'llm', + rows: [validLog()], + offset: 0, + total: 2, + nextOffset: 0, + truncated: false, + }, + { + kind: 'snapshot_logs', + revision: 'snapshot-revision-1', + source: 'llm', + rows: [validLog()], + offset: 0, + total: 1, + nextOffset: null, + truncated: 'no', + }, + { + kind: 'snapshot_pricing', + revision: 'snapshot-revision-1', + entries: Array.from({ length: PRICING_PAGE_MAX_ITEMS + 1 }, (_, index) => + customPricingEntry(`provider:model-${index}`), + ), + offset: 0, + total: PRICING_PAGE_MAX_ITEMS + 1, + nextOffset: null, + }, + { kind: 'revision_changed', expectedRevision: '' }, + ]) { + assert.throws(() => usageResponse(result), invalidFrame); + } + }); + + test('decodes revision-pinned Usage snapshot start, log, and pricing pages', () => { + assert.doesNotThrow(() => usageRequest({ kind: 'snapshot_start', range: { from: 1, to: 2 } })); + assert.doesNotThrow(() => + usageRequest({ + kind: 'snapshot_logs', + revision: 'snapshot-revision-1', + source: 'llm', + offset: 0, + limit: 3, + }), + ); + assert.doesNotThrow(() => + usageRequest({ + kind: 'snapshot_pricing', + revision: 'snapshot-revision-1', + offset: 0, + limit: 3, + }), + ); + + assert.doesNotThrow(() => + usageResponse({ + kind: 'snapshot_started', + revision: 'snapshot-revision-1', + summary: validSummary(), + provenance: validProvenance(), + }), + ); + assert.doesNotThrow(() => + usageResponse({ + kind: 'snapshot_logs', + revision: 'snapshot-revision-1', + source: 'llm', + rows: [validLog()], + offset: 0, + total: 1, + nextOffset: null, + truncated: false, + }), + ); + assert.doesNotThrow(() => + usageResponse({ + kind: 'snapshot_pricing', + revision: 'snapshot-revision-1', + entries: [customPricingEntry('provider:model')], + offset: 0, + total: 1, + nextOffset: null, + }), + ); + assert.doesNotThrow(() => + usageResponse({ kind: 'revision_changed', expectedRevision: 'snapshot-revision-1' }), + ); + + for (const input of [ + { kind: 'snapshot_start', range: 'all', revision: 'unexpected' }, + { kind: 'snapshot_logs', revision: '', source: 'llm', offset: 0, limit: 1 }, + { + kind: 'snapshot_logs', + revision: 'x'.repeat(129), + source: 'llm', + offset: 0, + limit: 1, + }, + { + kind: 'snapshot_logs', + revision: 'snapshot-revision-1', + source: 'model', + offset: 0, + limit: 1, + }, + { + kind: 'snapshot_pricing', + revision: 'snapshot-revision-1', + offset: 0, + limit: PRICING_PAGE_MAX_ITEMS + 1, + }, + ]) { + assert.throws(() => usageRequest(input), invalidFrame); + } }); test('enforces exact usage results and both page bounds', () => { @@ -441,6 +566,90 @@ describe('Usage/Pricing protocol', () => { } }); + test('pins every Usage authority behind one expiring LRU snapshot revision', async () => { + const base = await mkdtemp(join(tmpdir(), 'maka-usage-snapshot-')); + const capability = await resolveStorageRoot({ + path: join(base, 'interactive-root'), + kind: 'interactive', + }); + const owner = await tryAcquireInteractiveRootOwner(capability); + assert.ok(owner); + const stores = await openInteractiveUsageStoresForWrite(owner.lease); + let now = 1_000; + let nextRevision = 0; + try { + await stores.telemetry.recordLlmCall(longUsageRecord('old-llm', 1)); + await stores.telemetry.recordToolInvocation(longToolRecord('old-tool', 1)); + await stores.pricing.upsert(0, pricing('snapshot:old')); + const coordinator = new HostUsagePricingCoordinator( + stores, + () => {}, + new RuntimePolicyActivationGate(), + () => {}, + { + now: () => now, + createRevision: () => `snapshot-${++nextRevision}`, + ttlMs: 100, + capacity: 2, + activityLimit: 1, + }, + ); + + const first = await expectUsageSnapshotStart(coordinator); + assert.equal(first.revision, 'snapshot-1'); + assert.equal(first.summary.totalRequests, 1); + + await stores.telemetry.recordLlmCall(longUsageRecord('new-llm', 2)); + await stores.telemetry.recordToolInvocation(longToolRecord('new-tool', 2)); + await stores.pricing.upsert(1, pricing('snapshot:new')); + + const oldLlm = await expectUsageSnapshotLogs(coordinator, first.revision, 'llm'); + const oldTool = await expectUsageSnapshotLogs(coordinator, first.revision, 'tool'); + const oldPricing = await expectUsageSnapshotPricing(coordinator, first.revision); + assert.deepEqual( + oldLlm.rows.map((row) => row.id), + ['old-llm'], + ); + assert.deepEqual( + oldTool.rows.map((row) => row.id), + ['old-tool'], + ); + assert.equal(oldLlm.total, 1); + assert.equal(oldLlm.truncated, false); + assert.ok(oldPricing.entries.some((entry) => entry.pricing.modelKey === 'snapshot:old')); + assert.ok(!oldPricing.entries.some((entry) => entry.pricing.modelKey === 'snapshot:new')); + + const second = await expectUsageSnapshotStart(coordinator); + const newLlm = await expectUsageSnapshotLogs(coordinator, second.revision, 'llm'); + assert.deepEqual( + newLlm.rows.map((row) => row.id), + ['new-llm'], + ); + assert.equal(newLlm.total, 1, 'total describes retained rows'); + assert.equal(newLlm.truncated, true, 'truncation describes discarded authority rows'); + + await expectUsageSnapshotLogs(coordinator, first.revision, 'llm'); + await expectUsageSnapshotStart(coordinator); + assert.equal( + (await queryUsageSnapshotLogs(coordinator, second.revision, 'llm')).kind, + 'revision_changed', + 'the least recently used snapshot is evicted', + ); + + now += 101; + const expired = await queryUsageSnapshotLogs(coordinator, first.revision, 'llm'); + assert.deepEqual(expired, { kind: 'revision_changed', expectedRevision: first.revision }); + } finally { + await stores.close().catch(() => undefined); + await owner.close(); + await rm(join(resolveRootControlNamespace(), capability.rootId), { + recursive: true, + force: true, + }); + await rm(base, { recursive: true, force: true }); + } + }); + test('decodes revision-pinned numeric-offset pricing pages and revision-CAS mutation', () => { assert.doesNotThrow(() => pricingRequest('pricing.query', { kind: 'start' })); assert.doesNotThrow(() => @@ -870,6 +1079,78 @@ async function queryUsageBuckets( return frame.result.buckets; } +async function expectUsageSnapshotStart( + coordinator: HostUsagePricingCoordinator, +): Promise> { + const outcome = await coordinator.handlers['usage.query']( + { kind: 'snapshot_start', range: 'all' }, + CONNECTION_CONTEXT, + ); + assert.equal(outcome.ok, true); + if (!outcome.ok || outcome.result.kind !== 'snapshot_started') { + throw new Error('Expected a started Usage snapshot'); + } + return outcome.result; +} + +async function queryUsageSnapshotLogs( + coordinator: HostUsagePricingCoordinator, + revision: string, + source: 'llm' | 'tool', +): Promise> { + const outcome = await coordinator.handlers['usage.query']( + { kind: 'snapshot_logs', revision, source, offset: 0, limit: USAGE_PAGE_MAX_ITEMS }, + CONNECTION_CONTEXT, + ); + assert.equal(outcome.ok, true); + if ( + !outcome.ok || + (outcome.result.kind !== 'snapshot_logs' && outcome.result.kind !== 'revision_changed') + ) { + throw new Error('Expected a Usage snapshot log page'); + } + return outcome.result; +} + +async function expectUsageSnapshotLogs( + coordinator: HostUsagePricingCoordinator, + revision: string, + source: 'llm' | 'tool', +): Promise> { + const result = await queryUsageSnapshotLogs(coordinator, revision, source); + if (result.kind !== 'snapshot_logs') throw new Error('Expected a retained Usage snapshot'); + assert.equal(result.source, source); + return result; +} + +async function expectUsageSnapshotPricing( + coordinator: HostUsagePricingCoordinator, + revision: string, +): Promise<{ readonly entries: readonly EffectivePricingEntry[] }> { + const entries: EffectivePricingEntry[] = []; + let offset = 0; + let total: number | undefined; + do { + const outcome = await coordinator.handlers['usage.query']( + { kind: 'snapshot_pricing', revision, offset, limit: PRICING_PAGE_MAX_ITEMS }, + CONNECTION_CONTEXT, + ); + assert.equal(outcome.ok, true); + if (!outcome.ok || outcome.result.kind !== 'snapshot_pricing') { + throw new Error('Expected a Usage snapshot pricing page'); + } + assert.equal(outcome.result.revision, revision); + assert.equal(outcome.result.offset, offset); + total ??= outcome.result.total; + assert.equal(outcome.result.total, total); + entries.push(...outcome.result.entries); + if (outcome.result.nextOffset === null) break; + offset = outcome.result.nextOffset; + } while (true); + assert.equal(entries.length, total); + return { entries }; +} + function assertDistinctBoundedIdentities(values: readonly (string | undefined)[]): void { assert.equal(values.length, 6); assert.ok(values.every((value): value is string => typeof value === 'string')); diff --git a/packages/runtime-host/src/__tests__/usage-pricing-two-client-uds.test.ts b/packages/runtime-host/src/__tests__/usage-pricing-two-client-uds.test.ts index ede75459dd..e45b9d718d 100644 --- a/packages/runtime-host/src/__tests__/usage-pricing-two-client-uds.test.ts +++ b/packages/runtime-host/src/__tests__/usage-pricing-two-client-uds.test.ts @@ -447,6 +447,7 @@ describe('production Usage/Pricing UDS', () => { | undefined; const clients: RuntimeHostConnection[] = []; let endpoint: string | undefined; + let firstHostSnapshotRevision: string | undefined; try { firstOwner = await tryAcquireInteractiveRootOwner(capability); @@ -504,6 +505,15 @@ describe('production Usage/Pricing UDS', () => { }, ]); + const pinnedUsage = await desktop.request( + 'usage.query', + { kind: 'snapshot_start', range: 'all' }, + REQUEST_TIMEOUT_MS, + ); + assert.equal(pinnedUsage.kind, 'snapshot_started'); + if (pinnedUsage.kind !== 'snapshot_started') throw new Error('Usage snapshot did not start'); + firstHostSnapshotRevision = pinnedUsage.revision; + const initial = await readPricing(desktop); assert.equal(initial.revision, 0); assert.deepEqual(initial.entries, builtinPricingEntries()); @@ -560,6 +570,27 @@ describe('production Usage/Pricing UDS', () => { ); assert.deepEqual(retry, { kind: 'committed', revision: 2 }); + const pinnedPricing = await readUsageSnapshotPricing(desktop, pinnedUsage.revision); + assert.deepEqual(pinnedPricing, builtinPricingEntries()); + const pinnedLogs = await desktop.request( + 'usage.query', + { + kind: 'snapshot_logs', + revision: pinnedUsage.revision, + source: 'llm', + offset: 0, + limit: 100, + }, + REQUEST_TIMEOUT_MS, + ); + assert.equal(pinnedLogs.kind, 'snapshot_logs'); + if (pinnedLogs.kind === 'snapshot_logs') { + assert.deepEqual( + pinnedLogs.rows.map((row) => row.id), + ['usage-b', 'usage-a'], + ); + } + const [desktopPricing, tuiPricing] = await Promise.all([ readPricing(desktop), readPricing(tui), @@ -662,6 +693,21 @@ describe('production Usage/Pricing UDS', () => { connectClient(root), ]); clients.push(desktopAfterRestart, tuiAfterRestart); + assert.ok(firstHostSnapshotRevision); + assert.deepEqual( + await desktopAfterRestart.request( + 'usage.query', + { + kind: 'snapshot_logs', + revision: firstHostSnapshotRevision, + source: 'llm', + offset: 0, + limit: 100, + }, + REQUEST_TIMEOUT_MS, + ), + { kind: 'revision_changed', expectedRevision: firstHostSnapshotRevision }, + ); const [usageAfterRestart, pricingAfterRestart, pricingFromSecondClient] = await Promise.all([ readUsage(desktopAfterRestart), readPricing(desktopAfterRestart), @@ -772,6 +818,31 @@ async function readPricing(client: RuntimeHostConnection): Promise<{ return { revision: first.revision, entries, pageCount }; } +async function readUsageSnapshotPricing( + client: RuntimeHostConnection, + revision: string, +): Promise { + const entries: EffectivePricingEntry[] = []; + let offset = 0; + while (true) { + const result = await client.request( + 'usage.query', + { kind: 'snapshot_pricing', revision, offset, limit: 128 }, + REQUEST_TIMEOUT_MS, + ); + assert.equal(result.kind, 'snapshot_pricing'); + if (result.kind !== 'snapshot_pricing') throw new Error('Usage snapshot pricing disappeared'); + assert.equal(result.revision, revision); + assert.equal(result.offset, offset); + entries.push(...result.entries); + if (result.nextOffset === null) { + assert.equal(entries.length, result.total); + return entries; + } + offset = result.nextOffset; + } +} + async function readCoordinatorPricing( coordinator: HostUsagePricingCoordinator, ): Promise> { diff --git a/packages/runtime-host/src/protocol/index.ts b/packages/runtime-host/src/protocol/index.ts index a02f619545..0127a7740e 100644 --- a/packages/runtime-host/src/protocol/index.ts +++ b/packages/runtime-host/src/protocol/index.ts @@ -93,7 +93,10 @@ export const RUNTIME_HOST_REGISTRATION_SCHEMA_VERSION = 1 as const; export const RUNTIME_HOST_PROTOCOL_VERSION = 0 as const; // Increment when the same protocol version no longer guarantees safe Client-Host // interoperability. Mismatches are rejected before domain commands are admitted. -export const RUNTIME_HOST_COMPATIBILITY_EPOCH = 59 as const; +export const RUNTIME_HOST_COMPATIBILITY_EPOCH = 60 as const; +// 60: `usage.query` adds opaque revision-pinned snapshot start, activity, and +// pricing pages. Epoch-59 peers reject these exact new input/output variants, +// so mixed peers must fail the handshake before Settings Usage is requested. // 59: Scheduled Turn provider-retry frames may carry an optional host-clock // `ts`, letting a mid-wait re-projection recompute the authoritative // remaining duration. Older peers decode the frame with an exact key list diff --git a/packages/runtime-host/src/protocol/usage-pricing.ts b/packages/runtime-host/src/protocol/usage-pricing.ts index 266fa11bab..56c3c89290 100644 --- a/packages/runtime-host/src/protocol/usage-pricing.ts +++ b/packages/runtime-host/src/protocol/usage-pricing.ts @@ -34,7 +34,7 @@ import type { } from '@maka/core/usage-stats/types'; import { MODEL_CALL_KINDS } from '@maka/core/usage-stats/types'; import type { UsageProvenance } from '@maka/core/usage-ledger-merge'; -import { requireCount, requireExactRecord, requireRecord } from './codec.js'; +import { requireCount, requireExactRecord, requireId, requireRecord } from './codec.js'; import { invalidProtocolFrame } from './errors.js'; import { defineOperation } from './operation-spec.js'; @@ -193,6 +193,20 @@ export interface ToolUsageLogProjection { export type UsageLogProjection = LlmUsageLogProjection | ToolUsageLogProjection; export type UsageQueryInput = + | { readonly kind: 'snapshot_start'; readonly range: UsageQuery['range'] } + | { + readonly kind: 'snapshot_logs'; + readonly revision: string; + readonly source: 'llm' | 'tool'; + readonly offset?: number; + readonly limit?: number; + } + | { + readonly kind: 'snapshot_pricing'; + readonly revision: string; + readonly offset?: number; + readonly limit?: number; + } | { readonly kind: 'summary'; readonly query: LlmUsageQuery } | { readonly kind: 'buckets'; @@ -224,6 +238,41 @@ export type UsageQueryInput = }; export type UsageQueryResult = + | { + readonly kind: 'snapshot_started'; + readonly revision: string; + readonly summary: UsageSummaryV2; + readonly provenance: UsageProvenance; + } + | { + readonly kind: 'snapshot_logs'; + readonly revision: string; + readonly source: 'llm'; + readonly rows: readonly LlmUsageLogProjection[]; + readonly offset: number; + readonly total: number; + readonly nextOffset: number | null; + readonly truncated: boolean; + } + | { + readonly kind: 'snapshot_logs'; + readonly revision: string; + readonly source: 'tool'; + readonly rows: readonly ToolUsageLogProjection[]; + readonly offset: number; + readonly total: number; + readonly nextOffset: number | null; + readonly truncated: boolean; + } + | { + readonly kind: 'snapshot_pricing'; + readonly revision: string; + readonly entries: readonly EffectivePricingEntry[]; + readonly offset: number; + readonly total: number; + readonly nextOffset: number | null; + } + | { readonly kind: 'revision_changed'; readonly expectedRevision: string } | { readonly kind: 'summary'; readonly summary: UsageSummaryV2; @@ -346,6 +395,42 @@ export const USAGE_PRICING_OPERATION_SPECS = { export function decodeUsageQueryInput(value: unknown): UsageQueryInput { const input = requireRecord(value, 'usage query input'); + if (input.kind === 'snapshot_start') { + const exact = requireExactRecord(input, 'usage snapshot start input', ['kind', 'range']); + return { kind: 'snapshot_start', range: decodeUsageRange(exact.range) }; + } + if (input.kind === 'snapshot_logs') { + assertOptionalExactKeys( + input, + 'usage snapshot logs input', + ['kind', 'revision', 'source'], + ['offset', 'limit'], + ); + if (input.source !== 'llm' && input.source !== 'tool') { + throw invalidProtocolFrame('Invalid usage snapshot log source'); + } + return { + kind: 'snapshot_logs', + revision: requireId(input.revision, 'usage snapshot revision'), + source: input.source, + offset: decodeOffset(input.offset), + limit: decodeLimit(input.limit), + }; + } + if (input.kind === 'snapshot_pricing') { + assertOptionalExactKeys( + input, + 'usage snapshot pricing input', + ['kind', 'revision'], + ['offset', 'limit'], + ); + return { + kind: 'snapshot_pricing', + revision: requireId(input.revision, 'usage snapshot revision'), + offset: decodeOffset(input.offset), + limit: decodePricingLimit(input.limit), + }; + } if (input.kind === 'summary') { const exact = requireExactRecord(input, 'usage summary input', ['kind', 'query']); return { kind: 'summary', query: decodeLlmUsageQuery(exact.query) }; @@ -394,6 +479,60 @@ export function decodeUsageQueryInput(value: unknown): UsageQueryInput { export function decodeUsageQueryResult(value: unknown): UsageQueryResult { const result = requireRecord(value, 'usage query result'); + if (result.kind === 'snapshot_started') { + const exact = requireExactRecord(result, 'usage snapshot started result', [ + 'kind', + 'revision', + 'summary', + 'provenance', + ]); + return { + kind: 'snapshot_started', + revision: requireId(exact.revision, 'usage snapshot revision'), + summary: decodeUsageSummary(exact.summary), + provenance: decodeUsageProvenance(exact.provenance), + }; + } + if (result.kind === 'snapshot_logs') { + const exact = requireExactRecord(result, 'usage snapshot logs result', [ + 'kind', + 'revision', + 'source', + 'rows', + 'offset', + 'total', + 'nextOffset', + 'truncated', + ]); + if (exact.source === 'llm') { + return decodeUsageSnapshotLogPage('llm', exact, decodeLlmUsageLog); + } + if (exact.source === 'tool') { + return decodeUsageSnapshotLogPage('tool', exact, decodeToolUsageLog); + } + throw invalidProtocolFrame('Invalid usage snapshot log source'); + } + if (result.kind === 'snapshot_pricing') { + const exact = requireExactRecord(result, 'usage snapshot pricing result', [ + 'kind', + 'revision', + 'entries', + 'offset', + 'total', + 'nextOffset', + ]); + return decodeUsageSnapshotPricingPage(exact); + } + if (result.kind === 'revision_changed') { + const exact = requireExactRecord(result, 'usage snapshot revision changed result', [ + 'kind', + 'expectedRevision', + ]); + return { + kind: 'revision_changed', + expectedRevision: requireId(exact.expectedRevision, 'expected usage snapshot revision'), + }; + } if (result.kind === 'summary') { const exact = requireExactRecord(result, 'usage summary result', [ 'kind', @@ -606,6 +745,34 @@ export function decodePricingMutateResult(value: unknown): PricingMutateResult { } function assertUsageQueryOutputForInput(input: UsageQueryInput, output: UsageQueryResult): void { + if (input.kind === 'snapshot_start') { + if (output.kind !== 'snapshot_started') { + throw invalidProtocolFrame('Usage snapshot start response does not match its request'); + } + return; + } + if (input.kind === 'snapshot_logs' || input.kind === 'snapshot_pricing') { + if (output.kind === 'revision_changed') { + if (output.expectedRevision !== input.revision) { + throw invalidProtocolFrame('Usage snapshot revision change does not match its request'); + } + return; + } + if (output.kind !== input.kind) { + throw invalidProtocolFrame('Usage snapshot response kind does not match its request'); + } + if (output.revision !== input.revision || output.offset !== (input.offset ?? 0)) { + throw invalidProtocolFrame('Usage snapshot page does not match its request'); + } + if ( + input.kind === 'snapshot_logs' && + output.kind === 'snapshot_logs' && + output.source !== input.source + ) { + throw invalidProtocolFrame('Usage snapshot log source does not match its request'); + } + return; + } if (output.kind !== input.kind) { throw invalidProtocolFrame('Usage response kind does not match its request'); } @@ -729,6 +896,15 @@ function decodeLimit(value: unknown): number { return limit; } +function decodePricingLimit(value: unknown): number { + if (value === undefined) return PRICING_PAGE_MAX_ITEMS; + const limit = requireCount(value, 'usage snapshot pricing limit'); + if (limit === 0 || limit > PRICING_PAGE_MAX_ITEMS) { + throw invalidProtocolFrame('Invalid usage snapshot pricing limit'); + } + return limit; +} + function decodeUsagePage( kind: 'buckets', result: Record, @@ -787,6 +963,68 @@ function decodeUsageLogPage( return decoded; } +function decodeUsageSnapshotLogPage( + source: 'llm', + result: Record, + decodeItem: (value: unknown) => LlmUsageLogProjection, +): Extract; +function decodeUsageSnapshotLogPage( + source: 'tool', + result: Record, + decodeItem: (value: unknown) => ToolUsageLogProjection, +): Extract; +function decodeUsageSnapshotLogPage( + source: 'llm' | 'tool', + result: Record, + decodeItem: (value: unknown) => UsageLogProjection, +): Extract { + const rawItems = result.rows; + if (!Array.isArray(rawItems) || rawItems.length > USAGE_PAGE_MAX_ITEMS) { + throw invalidProtocolFrame('Usage snapshot page exceeds item limit'); + } + if (typeof result.truncated !== 'boolean') { + throw invalidProtocolFrame('Invalid usage snapshot truncation flag'); + } + const rows = rawItems.map(decodeItem); + const decoded = { + kind: 'snapshot_logs', + revision: requireId(result.revision, 'usage snapshot revision'), + source, + rows, + ...decodeUsagePagePosition(result, rows.length), + truncated: result.truncated, + } as Extract; + assertJsonBytes(decoded, USAGE_PAGE_MAX_BYTES, 'Usage snapshot page'); + return decoded; +} + +function decodeUsageSnapshotPricingPage( + result: Record, +): Extract { + const rawItems = result.entries; + if (!Array.isArray(rawItems) || rawItems.length > PRICING_PAGE_MAX_ITEMS) { + throw invalidProtocolFrame('Usage snapshot pricing page exceeds item limit'); + } + const entries = rawItems.map(decodeEffectivePricingEntry); + if ( + entries.some( + (item, index) => + index > 0 && + comparePricingModelKeys(entries[index - 1]!.pricing.modelKey, item.pricing.modelKey) >= 0, + ) + ) { + throw invalidProtocolFrame('Usage snapshot pricing entries are not canonically ordered'); + } + const decoded = { + kind: 'snapshot_pricing', + revision: requireId(result.revision, 'usage snapshot revision'), + entries, + ...decodeUsagePagePosition(result, entries.length), + } as const; + assertJsonBytes(decoded, PRICING_PAGE_MAX_BYTES, 'Usage snapshot pricing page'); + return decoded; +} + function decodeUsagePagePosition( result: Record, itemCount: number, diff --git a/packages/runtime-host/src/server/usage-pricing-coordinator.ts b/packages/runtime-host/src/server/usage-pricing-coordinator.ts index 01557bd642..c7adaf61ee 100644 --- a/packages/runtime-host/src/server/usage-pricing-coordinator.ts +++ b/packages/runtime-host/src/server/usage-pricing-coordinator.ts @@ -18,6 +18,7 @@ */ import { createHash } from 'node:crypto'; +import { resolveUsageRange } from '@maka/core/model-call-usage-projection'; import type { PricingConfig, ToolInvocationRecord, @@ -62,6 +63,7 @@ import { import type { UsagePricingOperationHandlerMap } from './operation-dispatcher.js'; import { RuntimePolicyActivationGate } from './runtime-policy-activation-gate.js'; import { readCanonicalUsage } from './canonical-usage-reader.js'; +import { UsageSnapshotCache, type UsageSnapshotCacheOptions } from './usage-snapshot-cache.js'; /** Root-scoped projection over the authentic lease-bound usage stores. */ export class HostUsagePricingCoordinator { @@ -75,6 +77,7 @@ export class HostUsagePricingCoordinator { readonly #requestDrain: () => void; readonly #activation: RuntimePolicyActivationGate; readonly #onCommittedPricingMutation: () => void; + readonly #usageSnapshots: UsageSnapshotCache; #poisonDrainRequested = false; constructor( @@ -82,11 +85,13 @@ export class HostUsagePricingCoordinator { requestDrain: () => void, activation: RuntimePolicyActivationGate, onCommittedPricingMutation: () => void = () => {}, + usageSnapshotOptions: UsageSnapshotCacheOptions = {}, ) { this.#stores = authenticateInteractiveUsageStoresWriter(stores); this.#requestDrain = requestDrain; this.#activation = activation; this.#onCommittedPricingMutation = onCommittedPricingMutation; + this.#usageSnapshots = new UsageSnapshotCache(usageSnapshotOptions); } /** @@ -104,6 +109,40 @@ export class HostUsagePricingCoordinator { async #queryUsage(input: UsageQueryInput): Promise> { try { const now = Date.now(); + if (input.kind === 'snapshot_start') { + return { ok: true, result: await this.#startUsageSnapshot(input.range, now) }; + } + if (input.kind === 'snapshot_logs') { + const snapshot = this.#usageSnapshots.get(input.revision); + if (!snapshot) return usageRevisionChanged(input.revision); + const rows = input.source === 'llm' ? snapshot.llmRows : snapshot.toolRows; + if ((input.offset ?? 0) > rows.length) return invalidUsageOffset(); + return { + ok: true, + result: usageSnapshotLogPage( + input.revision, + input.source, + rows, + input.offset ?? 0, + input.limit ?? USAGE_PAGE_MAX_ITEMS, + input.source === 'llm' ? snapshot.llmTruncated : snapshot.toolTruncated, + ), + }; + } + if (input.kind === 'snapshot_pricing') { + const snapshot = this.#usageSnapshots.get(input.revision); + if (!snapshot) return usageRevisionChanged(input.revision); + if ((input.offset ?? 0) > snapshot.pricingEntries.length) return invalidUsageOffset(); + return { + ok: true, + result: usageSnapshotPricingPage( + input.revision, + snapshot.pricingEntries, + input.offset ?? 0, + input.limit ?? PRICING_PAGE_MAX_ITEMS, + ), + }; + } if (input.kind === 'summary') { const merged = mergeUsageSummary( await this.#stores.telemetry.summary(input.query), @@ -192,6 +231,50 @@ export class HostUsagePricingCoordinator { } } + async #startUsageSnapshot( + range: UsageQuery['range'], + now: number, + ): Promise> { + const query: UsageQuery = { range: resolveUsageRange(range, now) }; + const captureLimit = this.#usageSnapshots.activityLimit + 1; + const captured = await this.#stores.captureUsageSnapshot({ + query, + activityLimit: captureLimit, + }); + const canonical: CanonicalUsageSource = { + attempts: captured.canonical.attempts, + unreadableRecords: captured.canonical.unreadableRecords + captured.repair.unreadableEvents, + pendingRepairs: captured.repair.pendingRuns, + }; + const mergedSummary = mergeUsageSummary(captured.legacySummary, canonical, query, now); + const { provenance, ...summary } = mergedSummary; + const mergedLogs = mergeUsageLogs( + captured.legacyLlmLogs, + canonical, + query, + now, + 0, + captureLimit, + ); + const retained = this.#usageSnapshots.retain({ + summary, + provenance, + llmRows: mergedLogs.rows.slice(0, this.#usageSnapshots.activityLimit).map(projectUsageLog), + llmTruncated: mergedLogs.total > this.#usageSnapshots.activityLimit, + toolRows: captured.toolLogs.rows + .slice(0, this.#usageSnapshots.activityLimit) + .map(projectToolUsageLog), + toolTruncated: captured.toolLogs.total > this.#usageSnapshots.activityLimit, + pricingEntries: projectEffectivePricingEntries(captured.pricing.overrides), + }); + return encodeUsageQueryResult({ + kind: 'snapshot_started', + revision: retained.revision, + summary: retained.summary, + provenance: retained.provenance, + }) as Extract; + } + async #queryPricing(input: PricingQueryInput): Promise> { try { const snapshot = await this.#stores.pricing.snapshot(); @@ -361,33 +444,111 @@ function invalidUsageOffset(): OperationOutcome<'usage.query'> { }; } +function usageRevisionChanged(revision: string): OperationOutcome<'usage.query'> { + return { + ok: true, + result: encodeUsageQueryResult({ kind: 'revision_changed', expectedRevision: revision }), + }; +} + +function usageSnapshotLogPage( + revision: string, + source: 'llm' | 'tool', + allRows: readonly UsageLogProjection[], + offset: number, + limit: number, + truncated: boolean, +): Extract { + const rows = fitBoundedPageItems( + allRows.slice(offset, offset + limit), + offset < allRows.length, + USAGE_PAGE_MAX_BYTES, + (candidate) => { + const nextOffset = offset + candidate.length; + return { + kind: 'snapshot_logs', + revision, + source, + rows: candidate, + offset, + total: allRows.length, + nextOffset: nextOffset < allRows.length ? nextOffset : null, + truncated, + } as const; + }, + 'Canonical Usage snapshot item', + ); + const nextOffset = offset + rows.length; + return encodeUsageQueryResult({ + kind: 'snapshot_logs', + revision, + source, + rows, + offset, + total: allRows.length, + nextOffset: nextOffset < allRows.length ? nextOffset : null, + truncated, + } as Extract) as Extract< + UsageQueryResult, + { kind: 'snapshot_logs' } + >; +} + +function usageSnapshotPricingPage( + revision: string, + allEntries: readonly EffectivePricingEntry[], + offset: number, + limit: number, +): Extract { + const entries = fitBoundedPageItems( + allEntries.slice(offset, offset + limit), + offset < allEntries.length, + PRICING_PAGE_MAX_BYTES, + (candidate) => { + const nextOffset = offset + candidate.length; + return { + kind: 'snapshot_pricing', + revision, + entries: candidate, + offset, + total: allEntries.length, + nextOffset: nextOffset < allEntries.length ? nextOffset : null, + } as const; + }, + 'Canonical Usage snapshot pricing entry', + ); + const nextOffset = offset + entries.length; + return encodeUsageQueryResult({ + kind: 'snapshot_pricing', + revision, + entries, + offset, + total: allEntries.length, + nextOffset: nextOffset < allEntries.length ? nextOffset : null, + }) as Extract; +} + function createPricingPage( revision: number, entries: readonly EffectivePricingEntry[], offset: number, ): PricingQueryResult { - const items: EffectivePricingEntry[] = []; - for (let index = offset; index < entries.length; index += 1) { - if (items.length >= PRICING_PAGE_MAX_ITEMS) break; - const item = entries[index]; - if (!item) break; - const candidate = [...items, item]; - const nextOffset = offset + candidate.length; - const page: PricingQueryResult = { - kind: 'page', - revision, - offset, - entries: candidate, - nextOffset: nextOffset < entries.length ? nextOffset : null, - }; - if (jsonBytes(page) > PRICING_PAGE_MAX_BYTES) { - if (items.length === 0) { - throw new Error('Canonical pricing entry exceeds the wire page limit'); - } - break; - } - items.push(item); - } + const items = fitBoundedPageItems( + entries.slice(offset, offset + PRICING_PAGE_MAX_ITEMS), + offset < entries.length, + PRICING_PAGE_MAX_BYTES, + (candidate) => { + const nextOffset = offset + candidate.length; + return { + kind: 'page', + revision, + offset, + entries: candidate, + nextOffset: nextOffset < entries.length ? nextOffset : null, + } satisfies PricingQueryResult; + }, + 'Canonical pricing entry', + ); const nextOffset = offset + items.length; return encodePricingQueryResult({ kind: 'page', @@ -427,28 +588,22 @@ function usagePage( provenance: UsageProvenance, ): Extract { const source = allItems.slice(offset, offset + limit); - const items: UsageBucket[] = []; - for (const item of source) { - const candidate = [...items, item]; - const nextOffset = offset + candidate.length; - if ( - jsonBytes( - bucketPageResult( - candidate, - total, - offset, - nextOffset < total ? nextOffset : null, - provenance, - ), - ) > USAGE_PAGE_MAX_BYTES - ) { - break; - } - items.push(item); - } - if (items.length === 0 && offset < total) { - throw new Error('Canonical usage item exceeds the wire page limit'); - } + const items = fitBoundedPageItems( + source, + offset < total, + USAGE_PAGE_MAX_BYTES, + (candidate) => { + const nextOffset = offset + candidate.length; + return bucketPageResult( + candidate, + total, + offset, + nextOffset < total ? nextOffset : null, + provenance, + ); + }, + 'Canonical usage item', + ); const nextOffset = offset + items.length; return bucketPageResult(items, total, offset, nextOffset < total ? nextOffset : null, provenance); } @@ -486,29 +641,23 @@ function usageLogPage( limit: number, provenance?: UsageProvenance, ): Extract { - const items: UsageLogProjection[] = []; - for (const item of allItems.slice(0, limit)) { - const candidate = [...items, item]; - const nextOffset = offset + candidate.length; - if ( - jsonBytes( - logPageResult( - source, - candidate, - total, - offset, - nextOffset < total ? nextOffset : null, - provenance, - ), - ) > USAGE_PAGE_MAX_BYTES - ) { - break; - } - items.push(item); - } - if (items.length === 0 && offset < total) { - throw new Error('Canonical usage item exceeds the wire page limit'); - } + const items = fitBoundedPageItems( + allItems.slice(0, limit), + offset < total, + USAGE_PAGE_MAX_BYTES, + (candidate) => { + const nextOffset = offset + candidate.length; + return logPageResult( + source, + candidate, + total, + offset, + nextOffset < total ? nextOffset : null, + provenance, + ); + }, + 'Canonical usage item', + ); const nextOffset = offset + items.length; return logPageResult( source, @@ -520,6 +669,25 @@ function usageLogPage( ); } +function fitBoundedPageItems( + candidates: readonly T[], + itemRequired: boolean, + maxBytes: number, + createPage: (items: readonly T[]) => unknown, + itemLabel: string, +): T[] { + const items: T[] = []; + for (const item of candidates) { + const next = [...items, item]; + if (jsonBytes(createPage(next)) > maxBytes) break; + items.push(item); + } + if (items.length === 0 && itemRequired) { + throw new Error(`${itemLabel} exceeds the wire page limit`); + } + return items; +} + function logPageResult( source: 'llm' | 'tool', rows: readonly UsageLogProjection[], diff --git a/packages/runtime-host/src/server/usage-snapshot-cache.ts b/packages/runtime-host/src/server/usage-snapshot-cache.ts new file mode 100644 index 0000000000..07e9ab1ef9 --- /dev/null +++ b/packages/runtime-host/src/server/usage-snapshot-cache.ts @@ -0,0 +1,124 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import { randomUUID } from 'node:crypto'; +import type { UsageSummaryV2 } from '@maka/core/usage-stats/types'; +import type { UsageProvenance } from '@maka/core/usage-ledger-merge'; +import type { + EffectivePricingEntry, + LlmUsageLogProjection, + ToolUsageLogProjection, +} from '../protocol/index.js'; + +export const USAGE_SNAPSHOT_TTL_MS = 5 * 60 * 1_000; +export const USAGE_SNAPSHOT_CAPACITY = 4; +export const USAGE_SNAPSHOT_ACTIVITY_LIMIT = 50_000; + +export interface UsageSnapshotCacheOptions { + readonly now?: () => number; + readonly createRevision?: () => string; + readonly ttlMs?: number; + readonly capacity?: number; + readonly activityLimit?: number; +} + +export interface UsageSnapshotContents { + readonly summary: UsageSummaryV2; + readonly provenance: UsageProvenance; + readonly llmRows: readonly LlmUsageLogProjection[]; + readonly llmTruncated: boolean; + readonly toolRows: readonly ToolUsageLogProjection[]; + readonly toolTruncated: boolean; + readonly pricingEntries: readonly EffectivePricingEntry[]; +} + +export interface RetainedUsageSnapshot extends UsageSnapshotContents { + readonly revision: string; +} + +interface CacheEntry extends RetainedUsageSnapshot { + readonly expiresAt: number; +} + +/** Host-epoch-local, absolute-TTL cache for coherent Settings Usage reads. */ +export class UsageSnapshotCache { + readonly activityLimit: number; + readonly #now: () => number; + readonly #createRevision: () => string; + readonly #ttlMs: number; + readonly #capacity: number; + readonly #entries = new Map(); + + constructor(options: UsageSnapshotCacheOptions = {}) { + this.#now = options.now ?? Date.now; + this.#createRevision = options.createRevision ?? randomUUID; + this.#ttlMs = options.ttlMs ?? USAGE_SNAPSHOT_TTL_MS; + this.#capacity = options.capacity ?? USAGE_SNAPSHOT_CAPACITY; + this.activityLimit = options.activityLimit ?? USAGE_SNAPSHOT_ACTIVITY_LIMIT; + if ( + !Number.isSafeInteger(this.#ttlMs) || + this.#ttlMs <= 0 || + !Number.isSafeInteger(this.#capacity) || + this.#capacity <= 0 || + !Number.isSafeInteger(this.activityLimit) || + this.activityLimit <= 0 + ) { + throw new TypeError('Invalid Usage snapshot cache limits'); + } + } + + retain(contents: UsageSnapshotContents): RetainedUsageSnapshot { + const now = this.#now(); + this.#pruneExpired(now); + while (this.#entries.size >= this.#capacity) { + const oldestRevision = this.#entries.keys().next().value; + if (oldestRevision === undefined) break; + this.#entries.delete(oldestRevision); + } + const revision = this.#createRevision(); + if (revision.length === 0 || revision.length > 128 || this.#entries.has(revision)) { + throw new Error('Usage snapshot revision generator returned an invalid revision'); + } + const entry: CacheEntry = { + revision, + ...contents, + expiresAt: now + this.#ttlMs, + }; + this.#entries.set(revision, entry); + return entry; + } + + get(revision: string): RetainedUsageSnapshot | undefined { + const now = this.#now(); + this.#pruneExpired(now); + const entry = this.#entries.get(revision); + if (!entry) return undefined; + // Map insertion order is the LRU order. Reinsert without changing expiresAt: + // page access affects eviction priority, never the absolute lifetime. + this.#entries.delete(revision); + this.#entries.set(revision, entry); + return entry; + } + + #pruneExpired(now: number): void { + for (const [revision, entry] of this.#entries) { + if (entry.expiresAt <= now) this.#entries.delete(revision); + } + } +} diff --git a/packages/storage/src/__tests__/usage-stores.test.ts b/packages/storage/src/__tests__/usage-stores.test.ts index c00e6f0b5b..6ac04f8fab 100644 --- a/packages/storage/src/__tests__/usage-stores.test.ts +++ b/packages/storage/src/__tests__/usage-stores.test.ts @@ -335,6 +335,44 @@ describe('InteractiveUsageStores', () => { }); }); + test('captures one repaired Usage authority snapshot behind the writer lease', async () => { + await withInteractiveRoot(async ({ root, capability }) => { + const owner = await tryAcquireInteractiveRootOwner(capability); + assert(owner); + const stores = await openInteractiveUsageStoresForWrite(owner.lease); + try { + await stores.telemetry.recordLlmCall( + llmRecord({ id: 'legacy-snapshot', sessionId: 'session-legacy' }), + ); + await stores.telemetry.recordToolInvocation(toolRecord()); + appendModelCallAuthorityEvent(root, modelCallAttempt('session-canonical')); + const pricing = { + modelKey: 'openai:gpt-5', + inputUsdPer1M: 1, + outputUsdPer1M: 2, + }; + await stores.pricing.upsert(0, pricing); + + const snapshot = await stores.captureUsageSnapshot({ + query: { range: 'all' }, + activityLimit: 10, + }); + + assert.equal(snapshot.legacySummary.totalRequests, 1); + assert.equal(snapshot.legacyLlmLogs.total, 1); + assert.equal(snapshot.legacyLlmLogs.rows[0]?.id, 'legacy-snapshot'); + assert.equal(snapshot.toolLogs.total, 1); + assert.equal(snapshot.toolLogs.rows[0]?.id, 'tool_1'); + assert.equal(snapshot.canonical.attempts[0]?.sessionId, 'session-canonical'); + assert.equal(snapshot.repair.pendingRuns, 0); + assert.deepEqual(snapshot.pricing, { revision: 1, overrides: [pricing] }); + } finally { + await stores.close(); + await owner.close(); + } + }); + }); + test('legacy summary clamps each cache reading to its own input', async () => { await withInteractiveRoot(async ({ capability }) => { const owner = await tryAcquireInteractiveRootOwner(capability); diff --git a/packages/storage/src/model-call-ledger.ts b/packages/storage/src/model-call-ledger.ts index 19749e1d40..c9a5153036 100644 --- a/packages/storage/src/model-call-ledger.ts +++ b/packages/storage/src/model-call-ledger.ts @@ -120,6 +120,21 @@ export function createSqliteModelCallLedger(workspaceRoot: string): ModelCallLed return new SqliteModelCallLedger(workspaceRoot); } +/** + * Runs the same bounded repair used by the ledger writer inside a caller-owned + * operational-state write transaction. This lets a cross-repository snapshot + * read the repaired projection before any other SQLite writer can intervene. + */ +export function catchUpModelCallProjectionInTransaction( + database: DatabaseSync, +): CatchUpModelCallProjectionResult { + try { + return catchUpModelCallProjection(database, {}, 16, 512); + } catch (cause) { + throw new ModelCallLedgerPublicationError(false, { cause }); + } +} + class SqliteModelCallLedger implements ModelCallLedger { readonly #lease: OperationalStateDatabaseLease; #state: 'open' | 'draining' | 'closed' = 'open'; diff --git a/packages/storage/src/usage-stores.ts b/packages/storage/src/usage-stores.ts index 19823028e3..54f8a94cf1 100644 --- a/packages/storage/src/usage-stores.ts +++ b/packages/storage/src/usage-stores.ts @@ -27,6 +27,7 @@ import type { } from '@maka/core/usage-stats/types'; import { throwDeduplicatedFailures } from './failure-utils.js'; import { + catchUpModelCallProjectionInTransaction, createSqliteModelCallLedger, type CatchUpModelCallProjectionInput, type CatchUpModelCallProjectionResult, @@ -47,6 +48,7 @@ import { type PricingSnapshot, type PricingStore, } from './pricing-store.js'; +import { acquireOperationalStateDatabase } from './operational-state-store.js'; import { runWithStorageRootLease, StorageRootAuthorityError, @@ -62,6 +64,7 @@ import { type PersistedToolInvocationRecord, type TelemetryRepo, type ToolUsageQuery, + resolveRange, } from './telemetry-repo.js'; import { createSqlitePricingStore, createSqliteTelemetryRepo } from './sqlite-usage-store.js'; @@ -123,6 +126,23 @@ export interface PricingAuthorityWriter extends PricingAuthorityReader { delete(expectedRevision: number, modelKey: string): Promise; } +export interface CaptureUsageSnapshotInput { + readonly query: UsageQuery; + readonly activityLimit: number; +} + +export interface CapturedUsageSnapshot { + readonly legacySummary: UsageSummaryV2; + readonly legacyLlmLogs: { readonly rows: readonly UsageLogRow[]; readonly total: number }; + readonly toolLogs: { + readonly rows: readonly PersistedToolInvocationRecord[]; + readonly total: number; + }; + readonly canonical: ModelCallLedgerPage; + readonly repair: CatchUpModelCallProjectionResult; + readonly pricing: PricingSnapshot; +} + export interface InteractiveUsageStoresReader { readonly kind: 'interactive'; readonly access: 'read'; @@ -140,6 +160,7 @@ export interface InteractiveUsageStoresWriter { readonly telemetry: Readonly; readonly modelCalls: Readonly; readonly pricing: Readonly; + captureUsageSnapshot(input: CaptureUsageSnapshotInput): Promise; subscribeSessionUsageChanges(listener: (sessionId: string) => void): () => void; beginDrain(): Promise; flush(): Promise; @@ -291,7 +312,13 @@ export async function openInteractiveUsageStoresForWrite( if (opening) return opening; const pending = runWithStorageRootLease(lease, 'interactive', 'write', async (root) => { const repos = await openRepos(root, true); - const stores = createWriterFacade(lease, repos.telemetry, repos.modelCalls, repos.pricing); + const stores = createWriterFacade( + root, + lease, + repos.telemetry, + repos.modelCalls, + repos.pricing, + ); writers.add(stores); writerByLease.set(lease, stores); return stores; @@ -328,6 +355,7 @@ async function openRepos( } function createWriterFacade( + root: string, lease: StorageRootLease<'interactive', 'write'>, telemetry: TelemetryRepo, modelCalls: ModelCallLedger, @@ -478,6 +506,42 @@ function createWriterFacade( isExpectedPricingFailure, ), }, + captureUsageSnapshot(input) { + if (!Number.isSafeInteger(input.activityLimit) || input.activityLimit <= 0) { + return Promise.reject(new TypeError('Usage snapshot activity limit must be positive')); + } + return admit(() => + run(() => { + const snapshotLease = acquireOperationalStateDatabase(root); + try { + const snapshot = snapshotLease.transaction('write', () => { + const repair = catchUpModelCallProjectionInTransaction(snapshotLease.database); + return snapshotLease.transaction('read', () => ({ + legacySummary: telemetry.summary(input.query), + legacyLlmLogs: telemetry.logs(input.query, 0, input.activityLimit), + toolLogs: telemetry.toolLogs( + { + range: input.query.range, + ...(input.query.status === undefined ? {} : { status: input.query.status }), + }, + 0, + input.activityLimit, + ), + canonical: modelCalls.read(resolveRange(input.query.range), input.query.sessionId), + repair, + pricing: pricing.snapshot(), + })); + }); + for (const sessionId of snapshot.repair.changedSessionIds) { + publishSessionUsageChange(sessionId); + } + return snapshot; + } finally { + snapshotLease.close(); + } + }), + ); + }, subscribeSessionUsageChanges(listener) { assertOpen(); sessionUsageChangeListeners.add(listener);