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..a322f2cad9 --- /dev/null +++ b/apps/desktop/src/main/__tests__/runtime-host-client-usage.test.ts @@ -0,0 +1,380 @@ +/* + * 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 }); + if (operation === 'usage.snapshot.release') return { released: true }; + 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, + ); + assert.deepEqual(requests.at(-1), { + operation: 'usage.snapshot.release', + input: { revision: 'revision-1' }, + }); + assert.equal( + requests.filter(({ operation }) => operation === 'usage.snapshot.release').length, + 1, + ); +}); + +test('releases an acquired Usage revision when its start range is invalid', async () => { + const released: string[] = []; + const client = usageClient(async (operation, input) => { + if (operation === 'usage.snapshot.release') { + released.push(input.revision); + return { released: true }; + } + if (input.kind === 'snapshot_start') { + const response = started('revision-1', 1); + return { ...response, summary: { ...response.summary, range: { from: 1, to: 10 } } }; + } + throw new Error('Unexpected Usage request'); + }); + + await assert.rejects( + () => client.loadUsageSnapshot({ from: 0, to: 10 }), + (error: unknown) => + error instanceof DesktopRuntimeHostClientError && error.code === 'projection_unstable', + ); + assert.deepEqual(released, ['revision-1']); +}); + +test('does not release an invalid Usage snapshot start without an acquired revision', async () => { + const released: string[] = []; + const client = usageClient(async (operation, input) => { + if (operation === 'usage.snapshot.release') { + released.push(input.revision); + return { released: true }; + } + if (input.kind === 'snapshot_start') { + return { kind: 'revision_changed', expectedRevision: 'revision-1' }; + } + throw new Error('Unexpected Usage request'); + }); + + await assert.rejects( + () => client.loadUsageSnapshot('all'), + (error: unknown) => + error instanceof DesktopRuntimeHostClientError && error.code === 'projection_unstable', + ); + assert.deepEqual(released, []); +}); + +test('discards every partial Usage result and restarts after revision_changed', async () => { + let starts = 0; + const released: string[] = []; + const client = usageClient(async (_operation, input) => { + if (_operation === 'usage.snapshot.release') { + released.push(input.revision); + return { released: true }; + } + 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']); + assert.deepEqual(released, ['revision-1', 'revision-2']); +}); + +test('releases an acquired Usage revision when a page reader throws without replacing its error', async () => { + const pageError = new Error('Usage page failed'); + const released: string[] = []; + const client = usageClient(async (operation, input) => { + if (operation === 'usage.snapshot.release') { + released.push(input.revision); + return { released: true }; + } + if (input.kind === 'snapshot_start') return started('revision-1', 1); + if (input.kind === 'snapshot_logs' && input.source === 'llm') throw pageError; + if (input.kind === 'snapshot_logs') { + return logPage('revision-1', 'tool', [toolLog('tool-1', 1)], 0, 1, null, false); + } + if (input.kind === 'snapshot_pricing') { + return pricingPage('revision-1', [pricing('model')], 0, 1, null); + } + throw new Error('Unexpected Usage request'); + }); + + await assert.rejects(() => client.loadUsageSnapshot('all'), (error: unknown) => error === pageError); + assert.deepEqual(released, ['revision-1']); +}); + +test('keeps a successful Usage snapshot when its release fails', async () => { + const released: string[] = []; + const client = usageClient(async (operation, input) => { + if (operation === 'usage.snapshot.release') { + released.push(input.revision); + throw new Error('Usage release failed'); + } + if (input.kind === 'snapshot_start') return started('revision-1', 1); + if (input.kind === 'snapshot_logs') { + const row = input.source === 'llm' ? llmLog('llm-1', 1) : toolLog('tool-1', 1); + return logPage('revision-1', input.source, [row], 0, 1, null, false); + } + if (input.kind === 'snapshot_pricing') { + return pricingPage('revision-1', [pricing('model')], 0, 1, null); + } + throw new Error('Unexpected Usage request'); + }); + + const snapshot = await client.loadUsageSnapshot('all'); + assert.equal(snapshot.revision, 'revision-1'); + assert.deepEqual(released, ['revision-1']); +}); + +test('keeps a successful Usage snapshot when release throws synchronously', async () => { + const released: string[] = []; + const client = usageClient((operation, input) => { + if (operation === 'usage.snapshot.release') { + released.push(input.revision); + throw new Error('Usage release failed synchronously'); + } + if (input.kind === 'snapshot_start') return started('revision-1', 1); + if (input.kind === 'snapshot_logs') { + const row = input.source === 'llm' ? llmLog('llm-1', 1) : toolLog('tool-1', 1); + return logPage('revision-1', input.source, [row], 0, 1, null, false); + } + if (input.kind === 'snapshot_pricing') { + return pricingPage('revision-1', [pricing('model')], 0, 1, null); + } + throw new Error('Unexpected Usage request'); + }); + + const snapshot = await client.loadUsageSnapshot('all'); + assert.equal(snapshot.revision, 'revision-1'); + assert.deepEqual(released, ['revision-1']); +}); + +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 | any, +): 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 3b7b3df61f..22b8ed3bc7 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, @@ -447,70 +224,28 @@ test("settings usage stats name each row from the Host-resolved session title", 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"); - // The Host carries `sessionTitle` on the projection (or omits it for - // untitled/unreadable sessions). The desktop layer just surfaces it. - return input.source === "llm" - ? ({ - kind: "logs", - source: "llm", - rows: [ - { - ...llmRow(0), - sessionId: "session-named", - sessionTitle: "重构使用统计页请求日志的任务列", - }, - { ...llmRow(1), sessionId: "session-untitled" }, - ], - offset: 0, - total: 2, - nextOffset: null, - provenance: provenance(), - } satisfies UsageQueryResult) - : ({ - kind: "logs", - source: "tool", - rows: [ - { - ...toolRow(0), - sessionId: "session-named", - sessionTitle: "重构使用统计页请求日志的任务列", - }, - ], - offset: 0, - total: 1, - 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), + sessionId: "session-named", + sessionTitle: "重构使用统计页请求日志的任务列", + }, + { ...llmRow(1), sessionId: "session-untitled" }, + ], + toolLogs: [ + { + ...toolRow(0), + sessionId: "session-named", + sessionTitle: "重构使用统计页请求日志的任务列", + }, + ], + pricingEntries: [], + llmLogsTruncated: false, + toolLogsTruncated: false, }), } as unknown as DesktopRuntimeHostClient, sendToRenderer: () => undefined, @@ -569,6 +304,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: { @@ -584,10 +339,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 54e5932172..2473f3b64d 100644 --- a/apps/desktop/src/main/runtime-host-client.ts +++ b/apps/desktop/src/main/runtime-host-client.ts @@ -30,6 +30,7 @@ import { projectSessionTodoItemsForDisplay, type SessionTodoItem, } from "@maka/core/session-todo"; +import type { UsageProvenance } from "@maka/core/usage-ledger-merge"; import type { ConnectionVersionBasis, @@ -42,7 +43,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, @@ -150,6 +151,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"; @@ -158,6 +163,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 = SessionConfigurationPatch; @@ -187,6 +194,7 @@ export type DesktopRuntimeHostClientErrorCode = | "revision_conflict" | "session_not_found" | "skill_catalog_unstable" + | "usage_unstable" | "unsupported_session"; export class DesktopRuntimeHostClientError extends Error { @@ -229,6 +237,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; @@ -1364,6 +1383,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 }); } @@ -1705,6 +1735,157 @@ 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") { + throw invalidProjection("Usage snapshot start"); + } + try { + if ( + 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, + }; + } finally { + try { + await this.request("usage.snapshot.release", { revision: started.revision }); + } catch { + // Usage snapshot release is best-effort cleanup. + } + } + } + + 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 d3b18b6cf3..6036e3e221 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; - } -} - // The Task column names the session each usage row belongs to. The Host resolves // the human-readable title (from the durable session header) and carries it on // the projection as `sessionTitle`; untitled/unreadable sessions omit it, and the diff --git a/docs/windows-test-inventory.md b/docs/windows-test-inventory.md index d13667bfd7..dda9a0de14 100644 --- a/docs/windows-test-inventory.md +++ b/docs/windows-test-inventory.md @@ -15,11 +15,11 @@ Locations intentionally omit line numbers so unrelated edits do not invalidate t | Classification | Count | |---|---:| -| windows-backend-gap | 27 | +| windows-backend-gap | 28 | | portable-candidate | 18 | | platform-contract | 31 | -Total Windows-excluded declarations: **76** +Total Windows-excluded declarations: **77** ## Inventory @@ -61,6 +61,7 @@ Total Windows-excluded declarations: **76** | windows-backend-gap | `packages/runtime-host/src/__tests__/session-revision-two-client-uds.test.ts` two Clients share exact retryable Session branch and revision authority | `process.platform === 'win32' ? 'Windows SQLite shutdown lifecycle' : false` | | windows-backend-gap | `packages/runtime-host/src/__tests__/usage-pricing-client-correlation.test.ts` fails the connection for a canonical response with mismatched ${mismatch.name} | `process.platform === 'win32'` | | windows-backend-gap | `packages/runtime-host/src/__tests__/usage-pricing-client-correlation.test.ts` rejects local invalid input without poisoning transport and correlates a private canonical copy | `process.platform === 'win32'` | +| windows-backend-gap | `packages/runtime-host/src/__tests__/usage-pricing-two-client-uds.test.ts` reclaims Usage snapshot capacity after a lease-owning client disconnects | `process.platform === 'win32'` | | windows-backend-gap | `packages/runtime-host/src/__tests__/usage-pricing-two-client-uds.test.ts` two clients share usage projection and one revision-CAS pricing authority | `process.platform === 'win32'` | | portable-candidate | `packages/runtime/src/__tests__/filesystem-apply-patch.test.ts` deletes a self-referential symlink entry without following it | `process.platform === 'win32'` | | platform-contract | `packages/runtime/src/__tests__/filesystem-worker-process-runner.test.ts` filesystem worker rejects boundedly when a detached descendant retains stdout | `process.platform === 'win32' ? 'POSIX detached process-group semantics required' : false` | diff --git a/packages/runtime-host/src/__tests__/handshake-compatibility.test.ts b/packages/runtime-host/src/__tests__/handshake-compatibility.test.ts index a2e93d3f62..98c0c56280 100644 --- a/packages/runtime-host/src/__tests__/handshake-compatibility.test.ts +++ b/packages/runtime-host/src/__tests__/handshake-compatibility.test.ts @@ -109,7 +109,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-65 Host before any domain command', async () => { let admittedRequest: RequestFrame | undefined; await withForgedHandshakePeer( async (transport, hostEpoch, rootId) => { @@ -121,7 +121,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: 65, 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 bed834df9a..05571231e0 100644 --- a/packages/runtime-host/src/__tests__/usage-pricing-protocol.test.ts +++ b/packages/runtime-host/src/__tests__/usage-pricing-protocol.test.ts @@ -38,11 +38,14 @@ import { openInteractiveUsageStoresForWrite } from '@maka/storage/usage-stores'; import { decodeClientFrame, decodeHostFrame, + decodeUsageSnapshotReleaseInput, + decodeUsageSnapshotReleaseResult, decodeUsageQueryInput, encodePricingQueryResult, encodeProtocolMessage, PRICING_PAGE_MAX_BYTES, PRICING_PAGE_MAX_ITEMS, + REMOTE_OWNER_OPERATION_GRANTS, RUNTIME_HOST_MAX_MESSAGE_BYTES, USAGE_PAGE_MAX_BYTES, USAGE_PAGE_MAX_ITEMS, @@ -50,6 +53,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'; @@ -63,6 +67,33 @@ const CONNECTION_CONTEXT: ConnectionContext = { }; describe('Usage/Pricing protocol', () => { + test('decodes the exact Usage snapshot release input and output', () => { + assert.equal(REMOTE_OWNER_OPERATION_GRANTS.includes('usage.snapshot.release'), true); + assert.deepEqual(decodeUsageSnapshotReleaseInput({ revision: 'snapshot-revision-1' }), { + revision: 'snapshot-revision-1', + }); + assert.deepEqual(decodeUsageSnapshotReleaseResult({ released: true }), { released: true }); + assert.deepEqual( + decodeClientFrame({ + requestId: 'usage-release-request', + operation: 'usage.snapshot.release', + input: { revision: 'snapshot-revision-1' }, + }), + { + requestId: 'usage-release-request', + operation: 'usage.snapshot.release', + input: { revision: 'snapshot-revision-1' }, + }, + ); + + for (const input of [{}, { revision: '' }, { revision: 'snapshot-revision-1', extra: true }]) { + assert.throws(() => decodeUsageSnapshotReleaseInput(input), invalidFrame); + } + for (const result of [{}, { released: false }, { released: true, extra: true }]) { + assert.throws(() => decodeUsageSnapshotReleaseResult(result), invalidFrame); + } + }); + test('decodes exact bounded usage queries', () => { assert.deepEqual( decodeUsageQueryInput({ @@ -171,6 +202,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', () => { @@ -472,6 +627,312 @@ describe('Usage/Pricing protocol', () => { } }); + test('leases Usage snapshots to connections with renewable idle and bounded hard lifetime', 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(), + () => {}, + async (sessionId) => `Title ${sessionId}`, + { + now: () => now, + createRevision: () => `snapshot-${++nextRevision}`, + ttlMs: 100, + hardTtlMs: 250, + capacity: 2, + activityLimit: 1, + }, + ); + + const connectionA = connectionContext('connection-a'); + const connectionB = connectionContext('connection-b'); + const connectionC = connectionContext('connection-c'); + const first = await expectUsageSnapshotStart(coordinator, connectionA); + 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')); + + assert.deepEqual( + await queryUsageSnapshotLogs(coordinator, first.revision, 'llm', connectionB), + { kind: 'revision_changed', expectedRevision: first.revision }, + ); + assert.deepEqual( + await coordinator.handlers['usage.snapshot.release']( + { revision: first.revision }, + connectionB, + ), + { ok: true, result: { released: true } }, + ); + const oldLlm = await expectUsageSnapshotLogs(coordinator, first.revision, 'llm', connectionA); + const oldTool = await expectUsageSnapshotLogs( + coordinator, + first.revision, + 'tool', + connectionA, + ); + const oldPricing = await expectUsageSnapshotPricing(coordinator, first.revision, connectionA); + assert.deepEqual( + oldLlm.rows.map((row) => row.id), + ['old-llm'], + ); + assert.equal(oldLlm.rows[0]?.sessionTitle, 'Title old-llm'); + assert.deepEqual( + oldTool.rows.map((row) => row.id), + ['old-tool'], + ); + assert.equal(oldTool.rows[0]?.sessionTitle, 'Title 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, connectionB); + const newLlm = await expectUsageSnapshotLogs( + coordinator, + second.revision, + 'llm', + connectionB, + ); + 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'); + + assert.deepEqual( + await coordinator.handlers['usage.query']( + { kind: 'snapshot_start', range: 'all' }, + connectionC, + ), + { + ok: false, + error: { + code: 'operation_conflict', + message: 'Usage snapshot capacity is occupied', + }, + }, + ); + assert.equal( + (await queryUsageSnapshotLogs(coordinator, second.revision, 'llm', connectionB)).kind, + 'snapshot_logs', + 'capacity pressure preserves every active lease', + ); + + now = 1_090; + await expectUsageSnapshotLogs(coordinator, first.revision, 'llm', connectionA); + now = 1_180; + await expectUsageSnapshotLogs(coordinator, first.revision, 'llm', connectionA); + now = 1_249; + await expectUsageSnapshotLogs(coordinator, first.revision, 'llm', connectionA); + now = 1_250; + assert.deepEqual( + await queryUsageSnapshotLogs(coordinator, first.revision, 'llm', connectionA), + { kind: 'revision_changed', expectedRevision: first.revision }, + ); + + assert.deepEqual( + await coordinator.handlers['usage.snapshot.release']( + { revision: first.revision }, + connectionA, + ), + { ok: true, result: { released: true } }, + ); + assert.deepEqual( + await coordinator.handlers['usage.snapshot.release']( + { revision: first.revision }, + connectionA, + ), + { ok: true, result: { released: true } }, + ); + const third = await expectUsageSnapshotStart(coordinator, connectionC); + coordinator.releaseConnection(connectionC.connectionId); + assert.deepEqual( + await queryUsageSnapshotLogs(coordinator, third.revision, 'llm', connectionC), + { kind: 'revision_changed', expectedRevision: third.revision }, + ); + await expectUsageSnapshotStart(coordinator, connectionA); + } 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('bounds concurrent Session title reads while retaining every snapshot title', async () => { + const base = await mkdtemp(join(tmpdir(), 'maka-usage-snapshot-titles-')); + 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); + const sessionIds = Array.from({ length: 20 }, (_, index) => `session-${index}`); + let inFlight = 0; + let maxInFlight = 0; + try { + await Promise.all( + sessionIds.map((sessionId, index) => + stores.telemetry.recordLlmCall(longUsageRecord(sessionId, index + 1)), + ), + ); + const coordinator = new HostUsagePricingCoordinator( + stores, + () => {}, + new RuntimePolicyActivationGate(), + () => {}, + async (sessionId) => { + inFlight += 1; + maxInFlight = Math.max(maxInFlight, inFlight); + await Promise.resolve(); + inFlight -= 1; + return `Title ${sessionId}`; + }, + ); + + const snapshot = await expectUsageSnapshotStart(coordinator); + const logs = await expectUsageSnapshotLogs(coordinator, snapshot.revision, 'llm'); + + assert.equal(logs.rows.length, sessionIds.length); + for (const row of logs.rows) { + assert.equal(row.sessionTitle, `Title ${row.sessionId}`); + } + assert.ok( + maxInFlight <= 16, + `Session title read concurrency ${maxInFlight} exceeded the limit`, + ); + } 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('reserves capacity before overlapping snapshot title hydration', async () => { + const base = await mkdtemp(join(tmpdir(), 'maka-usage-snapshot-reservations-')); + 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); + const releaseTitles = deferred(); + const fourTitlesEntered = deferred(); + const inFlight: Promise[] = []; + let titleReads = 0; + try { + await stores.telemetry.recordLlmCall(longUsageRecord('barrier-session', 1)); + const coordinator = new HostUsagePricingCoordinator( + stores, + () => {}, + new RuntimePolicyActivationGate(), + () => {}, + async (sessionId) => { + assert.equal(sessionId, 'barrier-session'); + titleReads += 1; + if (titleReads === 4) fourTitlesEntered.resolve(); + await releaseTitles.promise; + return 'Barrier title'; + }, + ); + const contexts = Array.from({ length: 5 }, (_, index) => + connectionContext(`overlap-${index}`), + ); + const firstStarts = contexts + .slice(0, 4) + .map((context) => + coordinator.handlers['usage.query']({ kind: 'snapshot_start', range: 'all' }, context), + ); + inFlight.push(...firstStarts); + await within( + fourTitlesEntered.promise, + 1_000, + 'Four admitted snapshot starts did not reach title hydration', + ); + + const fifthStart = coordinator.handlers['usage.query']( + { kind: 'snapshot_start', range: 'all' }, + contexts[4]!, + ); + inFlight.push(fifthStart); + let admissionFailure: unknown; + try { + const fifthBeforeRelease = await within( + fifthStart, + 1_000, + 'Fifth snapshot start reached expensive work before capacity conflict', + ); + assert.deepEqual(fifthBeforeRelease, { + ok: false, + error: { + code: 'operation_conflict', + message: 'Usage snapshot capacity is occupied', + }, + }); + assert.equal(titleReads, 4, 'only admitted starts may hydrate Session titles'); + } catch (error) { + admissionFailure = error; + } finally { + releaseTitles.resolve(); + } + + const firstOutcomes = await Promise.all(firstStarts); + await fifthStart; + if (admissionFailure) throw admissionFailure; + for (const [index, outcome] of firstOutcomes.entries()) { + assert.equal(outcome.ok, true); + if (!outcome.ok || outcome.result.kind !== 'snapshot_started') { + throw new Error('Admitted overlapping Usage snapshot did not start'); + } + const logs = await expectUsageSnapshotLogs( + coordinator, + outcome.result.revision, + 'llm', + contexts[index]!, + ); + assert.equal(logs.rows[0]?.sessionTitle, 'Barrier title'); + } + assert.equal(titleReads, 4); + } finally { + releaseTitles.resolve(); + await Promise.allSettled(inFlight); + 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(() => @@ -901,6 +1362,108 @@ async function queryUsageBuckets( return frame.result.buckets; } +async function expectUsageSnapshotStart( + coordinator: HostUsagePricingCoordinator, + context: ConnectionContext = CONNECTION_CONTEXT, +): Promise> { + const outcome = await coordinator.handlers['usage.query']( + { kind: 'snapshot_start', range: 'all' }, + 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', + context: ConnectionContext = CONNECTION_CONTEXT, +): Promise> { + const outcome = await coordinator.handlers['usage.query']( + { kind: 'snapshot_logs', revision, source, offset: 0, limit: USAGE_PAGE_MAX_ITEMS }, + 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', + context: ConnectionContext = CONNECTION_CONTEXT, +): Promise> { + const result = await queryUsageSnapshotLogs(coordinator, revision, source, context); + 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, + context: ConnectionContext = CONNECTION_CONTEXT, +): 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 }, + 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 connectionContext(connectionId: string): ConnectionContext { + return { ...CONNECTION_CONTEXT, connectionId }; +} + +function deferred(): { readonly promise: Promise; resolve(): void } { + let resolve!: () => void; + const promise = new Promise((settle) => { + resolve = settle; + }); + return { promise, resolve }; +} + +async function within(promise: Promise, timeoutMs: number, message: string): Promise { + let timer: ReturnType | undefined; + try { + return await Promise.race([ + promise, + new Promise((_, reject) => { + timer = setTimeout(() => reject(new Error(message)), timeoutMs); + }), + ]); + } finally { + if (timer) clearTimeout(timer); + } +} + 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 afd2d44c48..4e5281bd73 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 @@ -40,7 +40,11 @@ import { type StorageRootLease, type InteractiveRootOwner, } from '@maka/storage/root-authority'; -import { connectRuntimeHost, type RuntimeHostConnection } from '../client/index.js'; +import { + connectRuntimeHost, + RuntimeHostOperationError, + type RuntimeHostConnection, +} from '../client/index.js'; import { RUNTIME_HOST_PROTOCOL_VERSION, type EffectivePricingEntry, @@ -539,6 +543,92 @@ test('pricing query projects built-in and custom authority with reset effects', }); describe('production Usage/Pricing UDS', () => { + test('reclaims Usage snapshot capacity after a lease-owning client disconnects', { + skip: process.platform === 'win32', + timeout: 60_000, + }, async () => { + const base = await mkdtemp(join(tmpdir(), 'maka-usage-snapshot-disconnect-')); + const root = join(base, 'root'); + const capability = await resolveStorageRoot({ path: root, kind: 'interactive' }); + let owner = await tryAcquireInteractiveRootOwner(capability); + let host: RuntimeHostKernel | undefined; + const clients: RuntimeHostConnection[] = []; + + try { + assert.ok(owner, 'test must acquire the real Interactive write lease'); + host = await RuntimeHostKernel.start({ + owner, + idleGraceMs: 30_000, + composition: defineInteractiveRuntimeHostComposition(createExecutionRuntimeHostComposition), + }); + owner = undefined; + clients.push(...(await Promise.all(Array.from({ length: 5 }, () => connectClient(root))))); + + const revisions: string[] = []; + for (const client of clients.slice(0, 4)) { + const snapshot = await client.request( + 'usage.query', + { kind: 'snapshot_start', range: 'all' }, + REQUEST_TIMEOUT_MS, + ); + assert.equal(snapshot.kind, 'snapshot_started'); + if (snapshot.kind !== 'snapshot_started') throw new Error('Usage snapshot did not start'); + revisions.push(snapshot.revision); + } + assert.equal(new Set(revisions).size, 4); + + const contender = clients[4]!; + await assert.rejects( + contender.request( + 'usage.query', + { kind: 'snapshot_start', range: 'all' }, + REQUEST_TIMEOUT_MS, + ), + (error: unknown) => + error instanceof RuntimeHostOperationError && error.code === 'operation_conflict', + ); + + const disconnectedOwner = clients.shift()!; + await disconnectedOwner.close(); + + const deadline = Date.now() + 1_000; + let replacement; + while (true) { + try { + replacement = await contender.request( + 'usage.query', + { kind: 'snapshot_start', range: 'all' }, + REQUEST_TIMEOUT_MS, + ); + break; + } catch (error) { + if ( + !(error instanceof RuntimeHostOperationError) || + error.code !== 'operation_conflict' || + Date.now() >= deadline + ) { + throw error; + } + await new Promise((resolve) => setTimeout(resolve, 10)); + } + } + assert.equal(replacement.kind, 'snapshot_started'); + if (replacement.kind !== 'snapshot_started') { + throw new Error('Replacement Usage snapshot did not start'); + } + assert.equal(revisions.includes(replacement.revision), false); + } finally { + await Promise.allSettled(clients.map((client) => client.close())); + await host?.close().catch(() => undefined); + await owner?.close().catch(() => undefined); + await rm(join(resolveRootControlNamespace(), capability.rootId), { + recursive: true, + force: true, + }); + await rm(base, { recursive: true, force: true }); + } + }); + test('two clients share usage projection and one revision-CAS pricing authority', { skip: process.platform === 'win32', timeout: 60_000, @@ -555,6 +645,7 @@ describe('production Usage/Pricing UDS', () => { | undefined; const clients: RuntimeHostConnection[] = []; let endpoint: string | undefined; + let firstHostSnapshotRevision: string | undefined; try { firstOwner = await tryAcquireInteractiveRootOwner(capability); @@ -612,6 +703,38 @@ 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; + + assert.deepEqual( + await tui.request( + 'usage.query', + { + kind: 'snapshot_logs', + revision: pinnedUsage.revision, + source: 'llm', + offset: 0, + limit: 100, + }, + REQUEST_TIMEOUT_MS, + ), + { kind: 'revision_changed', expectedRevision: pinnedUsage.revision }, + ); + assert.deepEqual( + await tui.request( + 'usage.snapshot.release', + { revision: pinnedUsage.revision }, + REQUEST_TIMEOUT_MS, + ), + { released: true }, + ); + const initial = await readPricing(desktop); assert.equal(initial.revision, 0); assert.deepEqual(initial.entries, builtinPricingEntries()); @@ -668,6 +791,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), @@ -770,6 +914,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), @@ -880,6 +1039,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/__tests__/usage-snapshot-cache.test.ts b/packages/runtime-host/src/__tests__/usage-snapshot-cache.test.ts new file mode 100644 index 0000000000..bec3afa6f7 --- /dev/null +++ b/packages/runtime-host/src/__tests__/usage-snapshot-cache.test.ts @@ -0,0 +1,215 @@ +/* + * 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 { + UsageSnapshotCache, + UsageSnapshotCapacityError, + type UsageSnapshotContents, +} from '../server/usage-snapshot-cache.js'; + +const CONTENTS: UsageSnapshotContents = { + summary: { + range: { from: 0, to: 1 }, + 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: { + coverage: { + attempts: 0, + pricedAttempts: 0, + unpricedAttempts: 0, + usageReportedAttempts: 0, + usagePartialAttempts: 0, + usageMissingAttempts: 0, + }, + legacyRecords: 0, + unreadableRecords: 0, + pendingRepairs: 0, + }, + llmRows: [], + llmTruncated: false, + toolRows: [], + toolTruncated: false, + pricingEntries: [], +}; + +test('preserves four owned leases at capacity instead of evicting an active revision', () => { + let revision = 0; + const cache = new UsageSnapshotCache({ + capacity: 4, + createRevision: () => `revision-${++revision}`, + }); + const retained = Array.from({ length: 4 }, (_, index) => + cache.retain(`connection-${index}`, CONTENTS), + ); + + assert.throws( + () => cache.retain('connection-5', CONTENTS), + (error: unknown) => + error instanceof UsageSnapshotCapacityError && + error.message === 'Usage snapshot capacity is occupied', + ); + for (const [index, snapshot] of retained.entries()) { + assert.equal(cache.get(`connection-${index}`, snapshot.revision)?.revision, snapshot.revision); + } +}); + +test('enforces ownership and reclaims capacity on release and connection teardown', () => { + let revision = 0; + const cache = new UsageSnapshotCache({ + capacity: 2, + createRevision: () => `revision-${++revision}`, + }); + const first = cache.retain('connection-a', CONTENTS); + cache.retain('connection-b', CONTENTS); + + assert.equal(cache.get('connection-b', first.revision), undefined); + cache.release('connection-b', first.revision); + assert.equal(cache.get('connection-a', first.revision)?.revision, first.revision); + assert.throws(() => cache.retain('connection-c', CONTENTS), UsageSnapshotCapacityError); + + cache.release('connection-a', first.revision); + const third = cache.retain('connection-c', CONTENTS); + assert.equal(cache.get('connection-c', third.revision)?.revision, third.revision); + + cache.releaseConnection('connection-b'); + const fourth = cache.retain('connection-d', CONTENTS); + assert.equal(cache.get('connection-d', fourth.revision)?.revision, fourth.revision); +}); + +test('renews idle lifetime on owner access without extending the hard deadline', () => { + let now = 0; + const cache = new UsageSnapshotCache({ + now: () => now, + ttlMs: 100, + hardTtlMs: 250, + createRevision: () => 'revision-1', + }); + const retained = cache.retain('connection-a', CONTENTS); + + now = 90; + assert.equal(cache.get('connection-a', retained.revision)?.revision, retained.revision); + now = 180; + assert.equal(cache.get('connection-a', retained.revision)?.revision, retained.revision); + now = 249; + assert.equal(cache.get('connection-a', retained.revision)?.revision, retained.revision); + now = 250; + assert.equal(cache.get('connection-a', retained.revision), undefined); +}); + +test('expires an idle lease before its hard deadline', () => { + let now = 0; + const cache = new UsageSnapshotCache({ + now: () => now, + ttlMs: 100, + hardTtlMs: 250, + createRevision: () => 'revision-1', + }); + const retained = cache.retain('connection-a', CONTENTS); + + now = 100; + assert.equal(cache.get('connection-a', retained.revision), undefined); +}); + +test('pending reservations occupy capacity until their exact owner finalizes them', () => { + let revision = 0; + const cache = new UsageSnapshotCache({ + capacity: 4, + createRevision: () => `revision-${++revision}`, + }); + const reservations = Array.from({ length: 4 }, (_, index) => + cache.reserve(`connection-${index}`), + ); + + assert.throws(() => cache.reserve('connection-5'), UsageSnapshotCapacityError); + for (const [index, reservation] of reservations.entries()) { + assert.equal(cache.get(`connection-${index}`, reservation.revision), undefined); + } + + const first = cache.finalize('connection-0', reservations[0]!.revision, CONTENTS); + assert.equal(first?.revision, reservations[0]!.revision); + assert.equal(cache.get('connection-0', reservations[0]!.revision)?.revision, first?.revision); + assert.throws(() => cache.reserve('connection-5'), UsageSnapshotCapacityError); +}); + +test('abort, release, and connection teardown reclaim pending reservations', () => { + let revision = 0; + const cache = new UsageSnapshotCache({ + capacity: 2, + createRevision: () => `revision-${++revision}`, + }); + const first = cache.reserve('connection-a'); + const second = cache.reserve('connection-b'); + + cache.release('connection-b', first.revision); + assert.throws(() => cache.reserve('connection-c'), UsageSnapshotCapacityError); + + cache.abort('connection-a', first.revision); + const third = cache.reserve('connection-c'); + cache.release('connection-b', second.revision); + const fourth = cache.reserve('connection-d'); + + cache.releaseConnection('connection-c'); + assert.equal(cache.finalize('connection-c', third.revision, CONTENTS), undefined); + assert.equal( + cache.finalize('connection-d', fourth.revision, CONTENTS)?.revision, + fourth.revision, + ); + assert.doesNotThrow(() => cache.reserve('connection-e')); +}); + +test('finalization never revives a wrong, released, or expired reservation', () => { + let now = 0; + let revision = 0; + const cache = new UsageSnapshotCache({ + now: () => now, + ttlMs: 100, + hardTtlMs: 250, + capacity: 2, + createRevision: () => `revision-${++revision}`, + }); + const exact = cache.reserve('connection-a'); + + assert.equal(cache.finalize('connection-b', exact.revision, CONTENTS), undefined); + assert.equal(cache.finalize('connection-a', 'missing-revision', CONTENTS), undefined); + assert.equal(cache.get('connection-a', exact.revision), undefined); + assert.equal(cache.finalize('connection-a', exact.revision, CONTENTS)?.revision, exact.revision); + + cache.release('connection-a', exact.revision); + assert.equal(cache.finalize('connection-a', exact.revision, CONTENTS), undefined); + + const expiring = cache.reserve('connection-a'); + now = 100; + assert.equal(cache.finalize('connection-a', expiring.revision, CONTENTS), undefined); + assert.equal(cache.get('connection-a', expiring.revision), undefined); +}); diff --git a/packages/runtime-host/src/protocol/index.ts b/packages/runtime-host/src/protocol/index.ts index 623cd0ff10..6c21920bda 100644 --- a/packages/runtime-host/src/protocol/index.ts +++ b/packages/runtime-host/src/protocol/index.ts @@ -100,7 +100,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 = 96 as const; +export const RUNTIME_HOST_COMPATIBILITY_EPOCH = 97 as const; +// 97: `usage.query` adds connection-owned, revision-pinned snapshot start, +// activity, and pricing pages plus explicit `usage.snapshot.release`. Epoch-96 +// peers reject these closed variants, so mixed peers must fail the handshake. // 96: Read image tool results may carry durable `session_context` refs. // 95: Catalog entries carry `describedByMetadata`, so a client asks the // Host-resolved entry — not its own bundled table — whether a model needs a diff --git a/packages/runtime-host/src/protocol/operations.ts b/packages/runtime-host/src/protocol/operations.ts index 0a03ab137e..49cfb6c796 100644 --- a/packages/runtime-host/src/protocol/operations.ts +++ b/packages/runtime-host/src/protocol/operations.ts @@ -345,6 +345,7 @@ export const REMOTE_OWNER_OPERATION_GRANTS = Object.freeze([ 'turn.start', 'turn.stop', 'usage.query', + 'usage.snapshot.release', 'web-search.execute', 'workhub.coordination.answer', 'workhub.coordination.act', diff --git a/packages/runtime-host/src/protocol/usage-pricing.ts b/packages/runtime-host/src/protocol/usage-pricing.ts index 007bcda93d..7016da2599 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'; @@ -51,7 +51,7 @@ const QUERY_ERRORS = [ 'persistence_failed', 'internal_failure', ] as const; -const USAGE_QUERY_ERRORS = [...QUERY_ERRORS, 'invalid_request'] as const; +const USAGE_QUERY_ERRORS = [...QUERY_ERRORS, 'invalid_request', 'operation_conflict'] as const; const PRICING_QUERY_ERRORS = [...QUERY_ERRORS, 'invalid_request'] as const; const MUTATION_ERRORS = [...QUERY_ERRORS, 'invalid_request', 'commit_outcome_unknown'] as const; const LLM_USAGE_QUERY_FIELDS = new Set([ @@ -199,6 +199,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'; @@ -230,6 +244,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; @@ -311,6 +360,14 @@ export type PricingMutateResult = readonly actualRevision: number; }; +export interface UsageSnapshotReleaseInput { + readonly revision: string; +} + +export interface UsageSnapshotReleaseResult { + readonly released: true; +} + export const USAGE_PRICING_OPERATION_SPECS = { 'usage.query': defineOperation< UsageQueryInput, @@ -324,6 +381,17 @@ export const USAGE_PRICING_OPERATION_SPECS = { decodeOutput: decodeUsageQueryResult, assertOutputForInput: assertUsageQueryOutputForInput, }), + 'usage.snapshot.release': defineOperation< + UsageSnapshotReleaseInput, + UsageSnapshotReleaseResult, + (typeof QUERY_ERRORS)[number] + >({ + mode: 'control', + availability: 'ready', + errors: QUERY_ERRORS, + decodeInput: decodeUsageSnapshotReleaseInput, + decodeOutput: decodeUsageSnapshotReleaseResult, + }), 'pricing.query': defineOperation< PricingQueryInput, PricingQueryResult, @@ -350,8 +418,55 @@ export const USAGE_PRICING_OPERATION_SPECS = { }), } as const; +export function decodeUsageSnapshotReleaseInput(value: unknown): UsageSnapshotReleaseInput { + const input = requireExactRecord(value, 'usage snapshot release input', ['revision']); + return { revision: requireId(input.revision, 'usage snapshot revision') }; +} + +export function decodeUsageSnapshotReleaseResult(value: unknown): UsageSnapshotReleaseResult { + const result = requireExactRecord(value, 'usage snapshot release result', ['released']); + if (result.released !== true) throw invalidProtocolFrame('Usage snapshot was not released'); + return { released: true }; +} + 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) }; @@ -400,6 +515,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', @@ -612,6 +781,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'); } @@ -735,6 +932,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, @@ -793,6 +999,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/execution-composition.ts b/packages/runtime-host/src/server/execution-composition.ts index 2c390fed4f..35276d7707 100644 --- a/packages/runtime-host/src/server/execution-composition.ts +++ b/packages/runtime-host/src/server/execution-composition.ts @@ -1748,7 +1748,10 @@ export async function createExecutionRuntimeHostComposition( unsubscribeUsageChanges?.(); }, ], - releaseConnection: [(connectionId) => artifacts.releaseConnection(connectionId)], + releaseConnection: [ + (connectionId) => artifacts.releaseConnection(connectionId), + (connectionId) => usagePricing.releaseConnection(connectionId), + ], }), createRuntimeHostDomainModule({ id: 'client-capability', diff --git a/packages/runtime-host/src/server/usage-pricing-coordinator.ts b/packages/runtime-host/src/server/usage-pricing-coordinator.ts index 1404664657..8b09d957c2 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, @@ -60,14 +61,25 @@ import { type UsageQueryInput, type UsageQueryResult, } from '../protocol/index.js'; -import type { UsagePricingOperationHandlerMap } from './operation-dispatcher.js'; +import type { ConnectionContext, UsagePricingOperationHandlerMap } from './operation-dispatcher.js'; import { RuntimePolicyActivationGate } from './runtime-policy-activation-gate.js'; import { readCanonicalUsage } from './canonical-usage-reader.js'; +import { + UsageSnapshotCache, + UsageSnapshotCapacityError, + type UsageSnapshotCacheOptions, +} from './usage-snapshot-cache.js'; + +const USAGE_SESSION_TITLE_READ_CONCURRENCY = 16; /** Root-scoped projection over the authentic lease-bound usage stores. */ export class HostUsagePricingCoordinator { readonly handlers: UsagePricingOperationHandlerMap = { - 'usage.query': (input) => this.#queryUsage(input), + 'usage.query': (input, context) => this.#queryUsage(input, context), + 'usage.snapshot.release': async (input, context) => { + this.#usageSnapshots.release(context.connectionId, input.revision); + return { ok: true, result: { released: true } }; + }, 'pricing.query': (input) => this.#queryPricing(input), 'pricing.mutate': (input) => this.#mutatePricing(input), }; @@ -76,6 +88,7 @@ export class HostUsagePricingCoordinator { readonly #requestDrain: () => void; readonly #activation: RuntimePolicyActivationGate; readonly #onCommittedPricingMutation: () => void; + readonly #usageSnapshots: UsageSnapshotCache; // Resolves a session's human-readable title for the Task column. Reads the // durable session header directly (unfiltered, in-process), so it covers // reserved-role, coordination, and legacy sessions the catalog omits. @@ -88,14 +101,20 @@ export class HostUsagePricingCoordinator { activation: RuntimePolicyActivationGate, onCommittedPricingMutation: () => void = () => {}, readSessionTitle?: (sessionId: string) => Promise, + usageSnapshotOptions: UsageSnapshotCacheOptions = {}, ) { this.#stores = authenticateInteractiveUsageStoresWriter(stores); this.#requestDrain = requestDrain; this.#activation = activation; this.#onCommittedPricingMutation = onCommittedPricingMutation; + this.#usageSnapshots = new UsageSnapshotCache(usageSnapshotOptions); this.#readSessionTitle = readSessionTitle; } + releaseConnection(connectionId: string): void { + this.#usageSnapshots.releaseConnection(connectionId); + } + // Resolve titles for exactly the sessions on this page. A session that no // longer exists is simply left untitled — one deleted session never blanks // the rest. Store lifecycle, persistence, and malformed-header failures are @@ -110,8 +129,12 @@ export class HostUsagePricingCoordinator { const ids = [ ...new Set(rows.map((row) => row.sessionId).filter((id): id is string => id !== undefined)), ]; - await Promise.all( - ids.map(async (id) => { + let nextIndex = 0; + const worker = async (): Promise => { + while (nextIndex < ids.length) { + const id = ids[nextIndex]; + nextIndex += 1; + if (id === undefined) return; try { const title = (await read(id))?.trim(); if (title) titles.set(id, title); @@ -120,7 +143,10 @@ export class HostUsagePricingCoordinator { // any other failure is a store problem and must reach #queryUsage. if (!isSessionNotFoundError(error)) throw error; } - }), + } + }; + await Promise.all( + Array.from({ length: Math.min(ids.length, USAGE_SESSION_TITLE_READ_CONCURRENCY) }, worker), ); return titles; } @@ -137,9 +163,49 @@ export class HostUsagePricingCoordinator { return readCanonicalUsage(this.#stores, query, now, repair); } - async #queryUsage(input: UsageQueryInput): Promise> { + async #queryUsage( + input: UsageQueryInput, + context: ConnectionContext, + ): Promise> { try { const now = Date.now(); + if (input.kind === 'snapshot_start') { + return { + ok: true, + result: await this.#startUsageSnapshot(context.connectionId, input.range, now), + }; + } + if (input.kind === 'snapshot_logs') { + const snapshot = this.#usageSnapshots.get(context.connectionId, 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(context.connectionId, 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), @@ -232,10 +298,72 @@ export class HostUsagePricingCoordinator { ), }; } catch (error) { + if (error instanceof UsageSnapshotCapacityError) { + return { + ok: false, + error: { + code: 'operation_conflict', + message: 'Usage snapshot capacity is occupied', + }, + }; + } return this.#mapReadFailure<'usage.query'>(error, 'Usage authority'); } } + async #startUsageSnapshot( + connectionId: string, + range: UsageQuery['range'], + now: number, + ): Promise> { + const reservation = this.#usageSnapshots.reserve(connectionId); + try { + 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 llmRows = mergedLogs.rows.slice(0, this.#usageSnapshots.activityLimit); + const toolRows = captured.toolLogs.rows.slice(0, this.#usageSnapshots.activityLimit); + const titles = await this.#resolveSessionTitles([...llmRows, ...toolRows]); + const retained = this.#usageSnapshots.finalize(connectionId, reservation.revision, { + summary, + provenance, + llmRows: llmRows.map((row) => projectUsageLog(row, titles)), + llmTruncated: mergedLogs.total > this.#usageSnapshots.activityLimit, + toolRows: toolRows.map((row) => projectToolUsageLog(row, titles)), + toolTruncated: captured.toolLogs.total > this.#usageSnapshots.activityLimit, + pricingEntries: projectEffectivePricingEntries(captured.pricing.overrides), + }); + if (!retained) throw new Error('Usage snapshot reservation is no longer active'); + return encodeUsageQueryResult({ + kind: 'snapshot_started', + revision: retained.revision, + summary: retained.summary, + provenance: retained.provenance, + }) as Extract; + } catch (error) { + this.#usageSnapshots.abort(connectionId, reservation.revision); + throw error; + } + } + async #queryPricing(input: PricingQueryInput): Promise> { try { const snapshot = await this.#stores.pricing.snapshot(); @@ -405,33 +533,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', @@ -471,28 +677,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); } @@ -530,29 +730,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, @@ -564,6 +758,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..e2054a6b13 --- /dev/null +++ b/packages/runtime-host/src/server/usage-snapshot-cache.ts @@ -0,0 +1,209 @@ +/* + * 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_HARD_TTL_MS = 30 * 60 * 1_000; +export const USAGE_SNAPSHOT_CAPACITY = 4; +export const USAGE_SNAPSHOT_ACTIVITY_LIMIT = 50_000; + +export class UsageSnapshotCapacityError extends Error { + constructor() { + super('Usage snapshot capacity is occupied'); + this.name = 'UsageSnapshotCapacityError'; + } +} + +export interface UsageSnapshotCacheOptions { + readonly now?: () => number; + readonly createRevision?: () => string; + readonly ttlMs?: number; + readonly hardTtlMs?: 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; +} + +export interface UsageSnapshotReservation { + readonly revision: string; +} + +interface BaseCacheEntry extends UsageSnapshotReservation { + readonly connectionId: string; + idleExpiresAt: number; + readonly hardExpiresAt: number; +} + +interface PendingCacheEntry extends BaseCacheEntry { + readonly state: 'pending'; +} + +interface RetainedCacheEntry extends BaseCacheEntry, RetainedUsageSnapshot { + readonly state: 'retained'; +} + +type CacheEntry = PendingCacheEntry | RetainedCacheEntry; + +/** Host-epoch-local, connection-owned lease cache for coherent Settings Usage reads. */ +export class UsageSnapshotCache { + readonly activityLimit: number; + readonly #now: () => number; + readonly #createRevision: () => string; + readonly #ttlMs: number; + readonly #hardTtlMs: 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.#hardTtlMs = options.hardTtlMs ?? USAGE_SNAPSHOT_HARD_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.#hardTtlMs) || + this.#hardTtlMs <= 0 || + !Number.isSafeInteger(this.#capacity) || + this.#capacity <= 0 || + !Number.isSafeInteger(this.activityLimit) || + this.activityLimit <= 0 + ) { + throw new TypeError('Invalid Usage snapshot cache limits'); + } + } + + retain(connectionId: string, contents: UsageSnapshotContents): RetainedUsageSnapshot { + const reservation = this.reserve(connectionId); + try { + const retained = this.finalize(connectionId, reservation.revision, contents); + if (!retained) throw new Error('Usage snapshot reservation is no longer active'); + return retained; + } catch (error) { + this.abort(connectionId, reservation.revision); + throw error; + } + } + + reserve(connectionId: string): UsageSnapshotReservation { + const now = this.#now(); + this.#pruneExpired(now); + if (this.#entries.size >= this.#capacity) throw new UsageSnapshotCapacityError(); + 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 hardExpiresAt = now + this.#hardTtlMs; + const entry: PendingCacheEntry = { + revision, + connectionId, + state: 'pending', + // Both deadlines begin at reservation so capture and projection time can + // never escape either the renewable idle bound or the hard lifetime. + idleExpiresAt: Math.min(now + this.#ttlMs, hardExpiresAt), + hardExpiresAt, + }; + this.#entries.set(revision, entry); + return entry; + } + + finalize( + connectionId: string, + revision: string, + contents: UsageSnapshotContents, + ): RetainedUsageSnapshot | undefined { + const now = this.#now(); + this.#pruneExpired(now); + const reservation = this.#entries.get(revision); + if ( + !reservation || + reservation.state !== 'pending' || + reservation.connectionId !== connectionId + ) { + return undefined; + } + const retained: RetainedCacheEntry = { + revision, + ...contents, + connectionId, + state: 'retained', + idleExpiresAt: reservation.idleExpiresAt, + hardExpiresAt: reservation.hardExpiresAt, + }; + this.#entries.set(revision, retained); + return retained; + } + + get(connectionId: string, revision: string): RetainedUsageSnapshot | undefined { + const now = this.#now(); + this.#pruneExpired(now); + const entry = this.#entries.get(revision); + if (!entry || entry.state !== 'retained' || entry.connectionId !== connectionId) { + return undefined; + } + entry.idleExpiresAt = Math.min(now + this.#ttlMs, entry.hardExpiresAt); + return entry; + } + + abort(connectionId: string, revision: string): void { + this.release(connectionId, revision); + } + + release(connectionId: string, revision: string): void { + const entry = this.#entries.get(revision); + if (entry?.connectionId === connectionId) this.#entries.delete(revision); + } + + releaseConnection(connectionId: string): void { + for (const [revision, entry] of this.#entries) { + if (entry.connectionId === connectionId) this.#entries.delete(revision); + } + } + + #pruneExpired(now: number): void { + for (const [revision, entry] of this.#entries) { + if (entry.idleExpiresAt <= now || entry.hardExpiresAt <= 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);