From ca8a5f4cac3f5ae6fc38fd545e4e8598a3ba472f Mon Sep 17 00:00:00 2001 From: sungl <81428141+Sun-GLiang@users.noreply.github.com> Date: Fri, 28 Aug 2026 16:01:29 +0800 Subject: [PATCH 1/7] fix(runtime-host): serve revision-consistent usage snapshots --- .../runtime-host-client-usage.test.ts | 253 ++++++++++++ .../runtime-host-usage-ipc-main.test.ts | 378 ++++-------------- apps/desktop/src/main/runtime-host-client.ts | 173 +++++++- .../src/main/runtime-host-usage-ipc-main.ts | 113 +----- .../__tests__/handshake-compatibility.test.ts | 4 +- .../__tests__/usage-pricing-protocol.test.ts | 281 +++++++++++++ .../usage-pricing-two-client-uds.test.ts | 71 ++++ packages/runtime-host/src/protocol/index.ts | 5 +- .../src/protocol/usage-pricing.ts | 240 ++++++++++- .../src/server/usage-pricing-coordinator.ts | 302 ++++++++++---- .../src/server/usage-snapshot-cache.ts | 124 ++++++ .../src/__tests__/usage-stores.test.ts | 38 ++ packages/storage/src/model-call-ledger.ts | 15 + packages/storage/src/usage-stores.ts | 66 ++- 14 files changed, 1599 insertions(+), 464 deletions(-) create mode 100644 apps/desktop/src/main/__tests__/runtime-host-client-usage.test.ts create mode 100644 packages/runtime-host/src/server/usage-snapshot-cache.ts diff --git a/apps/desktop/src/main/__tests__/runtime-host-client-usage.test.ts b/apps/desktop/src/main/__tests__/runtime-host-client-usage.test.ts new file mode 100644 index 0000000000..bdec51dca8 --- /dev/null +++ b/apps/desktop/src/main/__tests__/runtime-host-client-usage.test.ts @@ -0,0 +1,253 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import assert from 'node:assert/strict'; +import test from 'node:test'; +import type { RuntimeHostConnection } from '@maka/runtime-host/client'; +import type { OperationInput, OperationKey } from '@maka/runtime-host/protocol'; +import { + DesktopRuntimeHostClient, + DesktopRuntimeHostClientError, +} from '../runtime-host-client.js'; + +test('loads all Usage snapshot pages behind one start revision', async () => { + const requests: Array<{ operation: OperationKey; input: unknown }> = []; + const client = usageClient(async (operation, input) => { + requests.push({ operation, input }); + assert.equal(operation, 'usage.query'); + if (input.kind === 'snapshot_start') return started('revision-1', 2); + assert.equal(input.revision, 'revision-1'); + if (input.kind === 'snapshot_logs' && input.source === 'llm') { + return input.offset === 0 + ? logPage('revision-1', 'llm', [llmLog('llm-1', 2)], 0, 2, 1, false) + : logPage('revision-1', 'llm', [llmLog('llm-2', 1)], 1, 2, null, false); + } + if (input.kind === 'snapshot_logs') { + return logPage('revision-1', 'tool', [toolLog('tool-1', 3)], 0, 1, null, false); + } + if (input.kind === 'snapshot_pricing') { + return input.offset === 0 + ? pricingPage('revision-1', [pricing('a:model')], 0, 2, 1) + : pricingPage('revision-1', [pricing('b:model')], 1, 2, null); + } + throw new Error('Unexpected Usage request'); + }); + + assert.deepEqual(await client.loadUsageSnapshot({ from: 0, to: 10 }), { + revision: 'revision-1', + summary: validSummary(2), + provenance: validProvenance(), + llmLogs: [llmLog('llm-1', 2), llmLog('llm-2', 1)], + toolLogs: [toolLog('tool-1', 3)], + pricingEntries: [pricing('a:model'), pricing('b:model')], + llmLogsTruncated: false, + toolLogsTruncated: false, + }); + assert.equal( + requests.filter(({ input }) => (input as { kind?: string }).kind === 'snapshot_start').length, + 1, + ); +}); + +test('discards every partial Usage result and restarts after revision_changed', async () => { + let starts = 0; + const client = usageClient(async (_operation, input) => { + if (input.kind === 'snapshot_start') { + starts += 1; + return started(`revision-${starts}`, starts); + } + if (input.revision === 'revision-1' && input.kind === 'snapshot_logs' && input.source === 'llm') { + return { kind: 'revision_changed', expectedRevision: 'revision-1' }; + } + if (input.kind === 'snapshot_logs') { + const row = input.source === 'llm' ? llmLog('fresh-llm', 2) : toolLog('fresh-tool', 1); + return logPage(input.revision, input.source, [row], 0, 1, null, false); + } + if (input.kind === 'snapshot_pricing') { + return pricingPage(input.revision, [pricing('fresh:model')], 0, 1, null); + } + throw new Error('Unexpected Usage request'); + }); + + const snapshot = await client.loadUsageSnapshot('all'); + assert.equal(starts, 2); + assert.equal(snapshot.revision, 'revision-2'); + assert.deepEqual(snapshot.llmLogs.map((row) => row.id), ['fresh-llm']); + assert.deepEqual(snapshot.toolLogs.map((row) => row.id), ['fresh-tool']); +}); + +test('fails with usage_unstable after three complete Usage snapshot attempts', async () => { + let starts = 0; + const client = usageClient(async (_operation, input) => { + if (input.kind === 'snapshot_start') { + starts += 1; + return started(`revision-${starts}`, 0); + } + return { kind: 'revision_changed', expectedRevision: input.revision }; + }); + + await assert.rejects( + () => client.loadUsageSnapshot('all'), + (error: unknown) => + error instanceof DesktopRuntimeHostClientError && error.code === 'usage_unstable', + ); + assert.equal(starts, 3); +}); + +test('rejects non-progressing or identity-changing Usage snapshot pages', async () => { + const client = usageClient(async (_operation, input) => { + if (input.kind === 'snapshot_start') return started('revision-1', 1); + if (input.kind === 'snapshot_logs' && input.source === 'llm') { + return logPage('wrong-revision', 'llm', [llmLog('llm-1', 1)], 0, 2, 0, false); + } + if (input.kind === 'snapshot_logs') { + return logPage('revision-1', 'tool', [], 0, 0, null, false); + } + return pricingPage('revision-1', [], 0, 0, null); + }); + + await assert.rejects( + () => client.loadUsageSnapshot('all'), + (error: unknown) => + error instanceof DesktopRuntimeHostClientError && error.code === 'projection_unstable', + ); +}); + +function usageClient( + respond: (operation: OperationKey, input: any) => Promise, +): DesktopRuntimeHostClient { + const connection = { + hostEpoch: 'host-current', + connectionId: 'connection-current', + rootId: 'root-current', + request: (operation: K, input: OperationInput) => + respond(operation, input), + close: async () => undefined, + } as unknown as RuntimeHostConnection; + return new DesktopRuntimeHostClient(connection); +} + +function started(revision: string, totalRequests: number) { + return { + kind: 'snapshot_started' as const, + revision, + summary: validSummary(totalRequests), + provenance: validProvenance(), + }; +} + +function logPage( + revision: string, + source: 'llm' | 'tool', + rows: readonly unknown[], + offset: number, + total: number, + nextOffset: number | null, + truncated: boolean, +) { + return { kind: 'snapshot_logs' as const, revision, source, rows, offset, total, nextOffset, truncated }; +} + +function pricingPage( + revision: string, + entries: readonly unknown[], + offset: number, + total: number, + nextOffset: number | null, +) { + return { kind: 'snapshot_pricing' as const, revision, entries, offset, total, nextOffset }; +} + +function validSummary(totalRequests: number) { + return { + range: { from: 0, to: 10 }, + totalRequests, + totalCostUsd: 0, + totalTokens: { + input: 0, + output: 0, + cacheMiss: 0, + cacheRead: 0, + cacheWrite: 0, + reasoning: 0, + total: 0, + }, + cacheHitRequests: 0, + cacheCreateRequests: 0, + errorRequests: 0, + }; +} + +function validProvenance() { + return { + coverage: { + attempts: 0, + pricedAttempts: 0, + unpricedAttempts: 0, + usageReportedAttempts: 0, + usagePartialAttempts: 0, + usageMissingAttempts: 0, + }, + legacyRecords: 0, + unreadableRecords: 0, + pendingRepairs: 0, + }; +} + +function llmLog(id: string, ts: number) { + return { + source: 'llm' as const, + id, + ts, + providerId: 'provider', + modelId: 'model', + inputTokens: 1, + outputTokens: 1, + cacheMissTokens: 1, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0, + totalTokens: 2, + costUsd: 0, + latencyMs: 1, + status: 'success' as const, + }; +} + +function toolLog(id: string, ts: number) { + return { + source: 'tool' as const, + id, + ts, + toolName: 'Read', + durationMs: 1, + status: 'success' as const, + bytesIn: 0, + bytesOut: 0, + startedAt: ts, + }; +} + +function pricing(modelKey: string) { + return { + source: 'custom' as const, + resetEffect: 'become_unpriced' as const, + pricing: { modelKey, inputUsdPer1M: 1, outputUsdPer1M: 2 }, + }; +} diff --git a/apps/desktop/src/main/__tests__/runtime-host-usage-ipc-main.test.ts b/apps/desktop/src/main/__tests__/runtime-host-usage-ipc-main.test.ts index 6199baa55b..6a26a749a9 100644 --- a/apps/desktop/src/main/__tests__/runtime-host-usage-ipc-main.test.ts +++ b/apps/desktop/src/main/__tests__/runtime-host-usage-ipc-main.test.ts @@ -20,77 +20,31 @@ import assert from "node:assert/strict"; import { test } from "node:test"; import type { UsageStats } from "@maka/core/settings"; -import type { UsageQueryInput, UsageQueryResult } from "@maka/runtime-host/protocol"; import type { IpcHandler } from "../ipc-reconnect-policy.js"; -import type { DesktopRuntimeHostClient } from "../runtime-host-client.js"; +import { + DesktopRuntimeHostClientError, + type DesktopRuntimeHostClient, +} from "../runtime-host-client.js"; import { registerRuntimeHostUsageIpc } from "../runtime-host-usage-ipc-main.js"; test("settings usage stats use the canonical model-call total and load every activity page", async () => { const handlers = new Map(); - const calls: Array<{ source?: "llm" | "tool"; offset?: number }> = []; - const ranges: UsageQueryInput["query"]["range"][] = []; + const ranges: unknown[] = []; registerRuntimeHostUsageIpc({ ipcMain: { handle: (channel, listener) => handlers.set(channel, listener), handleReconnectableRead: (channel, listener) => handlers.set(channel, listener), }, client: { - queryUsage: async (input: UsageQueryInput) => { - ranges.push(input.query.range); - if (input.kind === "summary") { - return { - kind: "summary", - summary: { - range: { from: 1, to: 2 }, - totalRequests: 151, - totalCostUsd: 12.5, - totalTokens: { - input: 3_000_000, - output: 500_000, - cacheMiss: 100_000, - cacheRead: 400_000, - cacheWrite: 43_090, - reasoning: 90, - total: 4_043_090, - }, - cacheHitRequests: 10, - cacheCreateRequests: 5, - errorRequests: 2, - }, - provenance: provenance(), - } satisfies UsageQueryResult; - } - if (input.kind !== "logs") throw new Error("unexpected usage query"); - calls.push({ source: input.source, offset: input.offset }); - if (input.source === "llm") { - const offset = input.offset ?? 0; - const count = offset === 0 ? 100 : 51; - return { - kind: "logs", - source: "llm", - rows: Array.from({ length: count }, (_, index) => llmRow(offset + index)), - offset, - total: 151, - nextOffset: offset === 0 ? 100 : null, - provenance: provenance(), - } satisfies UsageQueryResult; - } - const offset = input.offset ?? 0; - const count = offset === 0 ? 100 : 71; + loadUsageSnapshot: async (range: unknown) => { + ranges.push(range); return { - kind: "logs", - source: "tool", - rows: Array.from({ length: count }, (_, index) => toolRow(offset + index)), - offset, - total: 171, - nextOffset: offset === 0 ? 100 : null, - } satisfies UsageQueryResult; - }, - loadPricingSnapshot: async () => ({ - hostEpoch: "host-epoch", - connectionId: "connection-id", - revision: 1, - entries: [ + revision: "snapshot-1", + summary: usageSummary(151), + provenance: provenance(), + llmLogs: Array.from({ length: 151 }, (_, index) => llmRow(index)), + toolLogs: Array.from({ length: 171 }, (_, index) => toolRow(index)), + pricingEntries: [ { source: "custom", resetEffect: "become_unpriced", @@ -100,8 +54,11 @@ test("settings usage stats use the canonical model-call total and load every act outputUsdPer1M: 2, }, }, - ], - }), + ], + llmLogsTruncated: false, + toolLogsTruncated: false, + }; + }, } as unknown as DesktopRuntimeHostClient, sendToRenderer: () => undefined, }); @@ -115,15 +72,8 @@ test("settings usage stats use the canonical model-call total and load every act assert.equal(stats.logs.length, 322); assert.equal(stats.logs.filter((row) => row.kind === "model").length, 151); assert.equal(stats.logs.filter((row) => row.kind === "tool").length, 171); - const expectedCalls: Array<{ source?: "llm" | "tool"; offset?: number }> = [ - { source: "llm", offset: 0 }, - { source: "llm", offset: 100 }, - { source: "tool", offset: 0 }, - { source: "tool", offset: 100 }, - ]; - assert.deepEqual(calls.sort(compareCall), expectedCalls.sort(compareCall)); - assert.ok(ranges.every((range) => typeof range === "object")); - assert.ok(ranges.every((range) => JSON.stringify(range) === JSON.stringify(ranges[0]))); + assert.equal(ranges.length, 1); + assert.equal(typeof ranges[0], "object"); assert.equal(stats.logs.find((row) => row.id === "llm-150")?.status, "aborted"); assert.equal(stats.logs.find((row) => row.id === "llm-150")?.sessionId, undefined); assert.equal(stats.logs.find((row) => row.id === "llm-150")?.costUsd, undefined); @@ -148,7 +98,7 @@ test("settings usage stats use the canonical model-call total and load every act assert.equal(stats.logsTruncated, undefined); }); -test("settings usage stats reject a non-advancing activity page", async () => { +test("settings usage stats propagate an invalid snapshot projection", async () => { const handlers = new Map(); registerRuntimeHostUsageIpc({ ipcMain: { @@ -156,66 +106,26 @@ test("settings usage stats reject a non-advancing activity page", async () => { handleReconnectableRead: (channel, listener) => handlers.set(channel, listener), }, client: { - queryUsage: async (input: UsageQueryInput) => { - if (input.kind === "summary") { - return { - kind: "summary", - summary: { - range: { from: 1, to: 2 }, - totalRequests: 0, - totalCostUsd: 0, - totalTokens: { - input: 0, - output: 0, - cacheMiss: 0, - cacheRead: 0, - cacheWrite: 0, - reasoning: 0, - total: 0, - }, - cacheHitRequests: 0, - cacheCreateRequests: 0, - errorRequests: 0, - }, - provenance: provenance(), - } satisfies UsageQueryResult; - } - if (input.kind !== "logs") throw new Error("unexpected usage query"); - return input.source === "llm" - ? ({ - kind: "logs", - source: "llm", - rows: [], - offset: 0, - total: 1, - nextOffset: 0, - provenance: provenance(), - } satisfies UsageQueryResult) - : ({ - kind: "logs", - source: "tool", - rows: [], - offset: 0, - total: 0, - nextOffset: null, - } satisfies UsageQueryResult); + loadUsageSnapshot: async () => { + throw new DesktopRuntimeHostClientError( + "projection_unstable", + "Runtime Host returned an invalid Usage snapshot projection", + ); }, - loadPricingSnapshot: async () => ({ - hostEpoch: "host-epoch", - connectionId: "connection-id", - revision: 0, - entries: [], - }), } as unknown as DesktopRuntimeHostClient, sendToRenderer: () => undefined, }); const handler = handlers.get("settings:usageStats"); assert.ok(handler); - await assert.rejects(() => handler({} as never, "24h"), /invalid Usage projection/); + await assert.rejects( + () => handler({} as never, "24h"), + (error: unknown) => + error instanceof DesktopRuntimeHostClientError && error.code === "projection_unstable", + ); }); -test("settings usage stats degrade instead of erroring when logs disagree with the canonical summary", async () => { +test("settings usage stats fail when a coherent snapshot cannot be retained", async () => { const handlers = new Map(); registerRuntimeHostUsageIpc({ ipcMain: { @@ -223,69 +133,23 @@ test("settings usage stats degrade instead of erroring when logs disagree with t handleReconnectableRead: (channel, listener) => handlers.set(channel, listener), }, client: { - queryUsage: async (input: UsageQueryInput) => { - if (input.kind === "summary") { - return { - kind: "summary", - summary: { - range: { from: 1, to: 2 }, - totalRequests: 2, - totalCostUsd: 0, - totalTokens: { - input: 0, - output: 0, - cacheMiss: 0, - cacheRead: 0, - cacheWrite: 0, - reasoning: 0, - total: 0, - }, - cacheHitRequests: 0, - cacheCreateRequests: 0, - errorRequests: 0, - }, - provenance: provenance(), - } satisfies UsageQueryResult; - } - if (input.kind !== "logs") throw new Error("unexpected usage query"); - return input.source === "llm" - ? ({ - kind: "logs", - source: "llm", - rows: [llmRow(0)], - offset: 0, - total: 1, - nextOffset: null, - provenance: provenance(), - } satisfies UsageQueryResult) - : ({ - kind: "logs", - source: "tool", - rows: [], - offset: 0, - total: 0, - nextOffset: null, - } satisfies UsageQueryResult); + loadUsageSnapshot: async () => { + throw new DesktopRuntimeHostClientError( + "usage_unstable", + "Usage snapshot kept expiring while Desktop read it", + ); }, - loadPricingSnapshot: async () => ({ - hostEpoch: "host-epoch", - connectionId: "connection-id", - revision: 0, - entries: [], - }), } as unknown as DesktopRuntimeHostClient, sendToRenderer: () => undefined, }); const handler = handlers.get("settings:usageStats"); assert.ok(handler); - // A catch-up race (summary read before a repair commits, logs read after) must - // not error the whole page. The canonical summary total stays authoritative, - // the activity list holds what actually loaded, and provenance still rides along. - const stats = await handler({} as never, "all") as UsageStats; - assert.equal(stats.summary.totalRequests, 2); - assert.equal(stats.logs.filter((row) => row.kind === "model").length, 1); - assert.deepEqual(stats.provenance, provenance()); + await assert.rejects( + () => handler({} as never, "all"), + (error: unknown) => + error instanceof DesktopRuntimeHostClientError && error.code === "usage_unstable", + ); }); test("settings usage stats group the provider breakdown by connection", async () => { @@ -296,59 +160,18 @@ test("settings usage stats group the provider breakdown by connection", async () handleReconnectableRead: (channel, listener) => handlers.set(channel, listener), }, client: { - queryUsage: async (input: UsageQueryInput) => { - if (input.kind === "summary") { - return { - kind: "summary", - summary: { - range: { from: 1, to: 2 }, - totalRequests: 2, - totalCostUsd: 0, - totalTokens: { - input: 0, - output: 0, - cacheMiss: 0, - cacheRead: 0, - cacheWrite: 0, - reasoning: 0, - total: 0, - }, - cacheHitRequests: 0, - cacheCreateRequests: 0, - errorRequests: 0, - }, - provenance: provenance(), - } satisfies UsageQueryResult; - } - if (input.kind !== "logs") throw new Error("unexpected usage query"); - // Two connections to the SAME provider type must stay two rows. - return input.source === "llm" - ? ({ - kind: "logs", - source: "llm", - rows: [ - { ...llmRow(0), connectionSlug: "conn-a", providerId: "provider-x" }, - { ...llmRow(1), connectionSlug: "conn-b", providerId: "provider-x" }, - ], - offset: 0, - total: 2, - nextOffset: null, - provenance: provenance(), - } satisfies UsageQueryResult) - : ({ - kind: "logs", - source: "tool", - rows: [], - offset: 0, - total: 0, - nextOffset: null, - } satisfies UsageQueryResult); - }, - loadPricingSnapshot: async () => ({ - hostEpoch: "host-epoch", - connectionId: "connection-id", - revision: 0, - entries: [], + loadUsageSnapshot: async () => ({ + revision: "snapshot-1", + summary: usageSummary(2), + provenance: provenance(), + llmLogs: [ + { ...llmRow(0), connectionSlug: "conn-a", providerId: "provider-x" }, + { ...llmRow(1), connectionSlug: "conn-b", providerId: "provider-x" }, + ], + toolLogs: [], + pricingEntries: [], + llmLogsTruncated: false, + toolLogsTruncated: false, }), } as unknown as DesktopRuntimeHostClient, sendToRenderer: () => undefined, @@ -365,68 +188,22 @@ test("settings usage stats group the provider breakdown by connection", async () test("settings usage stats truncate the activity log at the cap instead of erroring", async () => { const handlers = new Map(); - const PAGE = 100; - // Above MAX_ACTIVITY_RECORDS (50_000) so paging must stop and flag truncation. - const TOTAL = 50_150; + const TOTAL = 50_000; registerRuntimeHostUsageIpc({ ipcMain: { handle: (channel, listener) => handlers.set(channel, listener), handleReconnectableRead: (channel, listener) => handlers.set(channel, listener), }, client: { - queryUsage: async (input: UsageQueryInput) => { - if (input.kind === "summary") { - return { - kind: "summary", - summary: { - range: { from: 1, to: 2 }, - totalRequests: TOTAL, - totalCostUsd: 0, - totalTokens: { - input: 0, - output: 0, - cacheMiss: 0, - cacheRead: 0, - cacheWrite: 0, - reasoning: 0, - total: 0, - }, - cacheHitRequests: 0, - cacheCreateRequests: 0, - errorRequests: 0, - }, - provenance: provenance(), - } satisfies UsageQueryResult; - } - if (input.kind !== "logs") throw new Error("unexpected usage query"); - if (input.source === "llm") { - const offset = input.offset ?? 0; - const count = Math.min(PAGE, TOTAL - offset); - const nextOffset = offset + count < TOTAL ? offset + count : null; - return { - kind: "logs", - source: "llm", - rows: Array.from({ length: count }, (_, index) => llmRow(offset + index)), - offset, - total: TOTAL, - nextOffset, - provenance: provenance(), - } satisfies UsageQueryResult; - } - return { - kind: "logs", - source: "tool", - rows: [], - offset: 0, - total: 0, - nextOffset: null, - } satisfies UsageQueryResult; - }, - loadPricingSnapshot: async () => ({ - hostEpoch: "host-epoch", - connectionId: "connection-id", - revision: 0, - entries: [], + loadUsageSnapshot: async () => ({ + revision: "snapshot-1", + summary: usageSummary(TOTAL + 150), + provenance: provenance(), + llmLogs: Array.from({ length: TOTAL }, (_, index) => llmRow(index)), + toolLogs: [], + pricingEntries: [], + llmLogsTruncated: true, + toolLogsTruncated: false, }), } as unknown as DesktopRuntimeHostClient, sendToRenderer: () => undefined, @@ -476,6 +253,26 @@ function toolRow(index: number) { }; } +function usageSummary(totalRequests: number) { + return { + range: { from: 1, to: 2 }, + totalRequests, + totalCostUsd: 12.5, + totalTokens: { + input: 3_000_000, + output: 500_000, + cacheMiss: 100_000, + cacheRead: 400_000, + cacheWrite: 43_090, + reasoning: 90, + total: 4_043_090, + }, + cacheHitRequests: 10, + cacheCreateRequests: 5, + errorRequests: 2, + }; +} + function provenance() { return { coverage: { @@ -491,10 +288,3 @@ function provenance() { pendingRepairs: 0, }; } - -function compareCall( - left: { source?: "llm" | "tool"; offset?: number }, - right: { source?: "llm" | "tool"; offset?: number }, -): number { - return `${left.source}:${left.offset}`.localeCompare(`${right.source}:${right.offset}`); -} diff --git a/apps/desktop/src/main/runtime-host-client.ts b/apps/desktop/src/main/runtime-host-client.ts index e07d625122..e042b65187 100644 --- a/apps/desktop/src/main/runtime-host-client.ts +++ b/apps/desktop/src/main/runtime-host-client.ts @@ -27,6 +27,7 @@ import { } from "@maka/core/session"; import { markPersisted } from "@maka/core/persisted-value"; import type { Task } from "@maka/core/task-ledger"; +import type { UsageProvenance } from "@maka/core/usage-ledger-merge"; import type { ConnectionCatalogSnapshot, @@ -40,7 +41,7 @@ import { canonicalPricingConfigsEqual, comparePricingModelKeys, } from "@maka/core/usage-stats/pricing"; -import type { PricingConfig } from "@maka/core/usage-stats/types"; +import type { PricingConfig, TimeRange, UsageSummaryV2 } from "@maka/core/usage-stats/types"; import { type ClientCapabilityProvider, type DecodedSessionTranscriptPage, @@ -134,6 +135,10 @@ import { type TurnInterruptResult, type TurnMessageSubmitInput, type TurnMessageSubmitResult, + type LlmUsageLogProjection, + type ToolUsageLogProjection, + PRICING_PAGE_MAX_ITEMS, + USAGE_PAGE_MAX_ITEMS, type WorkspaceProjection, } from "@maka/runtime-host/protocol"; @@ -142,6 +147,8 @@ const decodeStoredMessage = (value: unknown): StoredMessage => const MAX_OPTIMISTIC_ATTEMPTS = 3; const MAX_SESSION_REVISION_ATTEMPTS = 8; const MAX_PRICING_SNAPSHOT_ATTEMPTS = 3; +const MAX_USAGE_SNAPSHOT_ATTEMPTS = 3; +const MAX_USAGE_SNAPSHOT_ACTIVITY_RECORDS = 50_000; export type DesktopSessionConfigurationPatch = Partial; @@ -160,6 +167,7 @@ export type DesktopRuntimeHostClientErrorCode = | "revision_conflict" | "session_not_found" | "skill_catalog_unstable" + | "usage_unstable" | "unsupported_session"; export class DesktopRuntimeHostClientError extends Error { @@ -202,6 +210,17 @@ export interface DesktopPricingSnapshot { readonly entries: readonly EffectivePricingEntry[]; } +export interface DesktopUsageSnapshot { + readonly revision: string; + readonly summary: UsageSummaryV2; + readonly provenance: UsageProvenance; + readonly llmLogs: readonly LlmUsageLogProjection[]; + readonly toolLogs: readonly ToolUsageLogProjection[]; + readonly pricingEntries: readonly EffectivePricingEntry[]; + readonly llmLogsTruncated: boolean; + readonly toolLogsTruncated: boolean; +} + export interface DesktopSkillCatalogSnapshot { readonly revision: SkillCatalogRevision; readonly view: SkillCatalogView; @@ -1264,6 +1283,17 @@ export class DesktopRuntimeHostClient { return this.request("usage.query", input); } + async loadUsageSnapshot(range: TimeRange): Promise { + for (let attempt = 0; attempt < MAX_USAGE_SNAPSHOT_ATTEMPTS; attempt += 1) { + const snapshot = await this.#readUsageSnapshot(range); + if (snapshot) return snapshot; + } + throw new DesktopRuntimeHostClientError( + "usage_unstable", + "Usage snapshot kept expiring while Desktop read it", + ); + } + queryGoal(sessionId: string): Promise> { return this.request("goal.query", { sessionId }); } @@ -1605,6 +1635,147 @@ export class DesktopRuntimeHostClient { }; } + async #readUsageSnapshot(range: TimeRange): Promise { + this.#assertOpen(); + const started = await this.request("usage.query", { kind: "snapshot_start", range }); + if ( + started.kind !== "snapshot_started" || + (typeof range === "object" && + (started.summary.range.from !== range.from || started.summary.range.to !== range.to)) + ) { + throw invalidProjection("Usage snapshot start"); + } + const [llm, tool, pricing] = await Promise.all([ + this.#readUsageSnapshotLogs(started.revision, "llm"), + this.#readUsageSnapshotLogs(started.revision, "tool"), + this.#readUsageSnapshotPricing(started.revision), + ]); + if (!llm || !tool || !pricing) return undefined; + return { + revision: started.revision, + summary: started.summary, + provenance: started.provenance, + llmLogs: llm.rows, + toolLogs: tool.rows, + pricingEntries: pricing, + llmLogsTruncated: llm.truncated, + toolLogsTruncated: tool.truncated, + }; + } + + async #readUsageSnapshotLogs( + revision: string, + source: "llm", + ): Promise<{ readonly rows: readonly LlmUsageLogProjection[]; readonly truncated: boolean } | undefined>; + async #readUsageSnapshotLogs( + revision: string, + source: "tool", + ): Promise<{ readonly rows: readonly ToolUsageLogProjection[]; readonly truncated: boolean } | undefined>; + async #readUsageSnapshotLogs( + revision: string, + source: "llm" | "tool", + ): Promise< + | { + readonly rows: readonly (LlmUsageLogProjection | ToolUsageLogProjection)[]; + readonly truncated: boolean; + } + | undefined + > { + const rows: Array = []; + let offset = 0; + let total: number | undefined; + let truncated: boolean | undefined; + while (true) { + const page = await this.request("usage.query", { + kind: "snapshot_logs", + revision, + source, + offset, + limit: USAGE_PAGE_MAX_ITEMS, + }); + if (page.kind === "revision_changed") { + if (page.expectedRevision !== revision) throw invalidProjection("Usage snapshot revision"); + return undefined; + } + if ( + page.kind !== "snapshot_logs" || + page.revision !== revision || + page.source !== source || + page.offset !== offset || + page.rows.length > USAGE_PAGE_MAX_ITEMS || + page.total > MAX_USAGE_SNAPSHOT_ACTIVITY_RECORDS + ) { + throw invalidProjection("Usage snapshot logs"); + } + total ??= page.total; + truncated ??= page.truncated; + if (page.total !== total || page.truncated !== truncated || rows.length !== offset) { + throw invalidProjection("Usage snapshot logs"); + } + rows.push(...page.rows); + if (rows.length > total) throw invalidProjection("Usage snapshot logs"); + if (page.nextOffset === null) { + if (rows.length !== total) throw invalidProjection("Usage snapshot logs"); + return { rows, truncated }; + } + if ( + page.rows.length === 0 || + page.nextOffset !== offset + page.rows.length || + page.nextOffset >= total + ) { + throw invalidProjection("Usage snapshot logs"); + } + offset = page.nextOffset; + } + } + + async #readUsageSnapshotPricing( + revision: string, + ): Promise { + const entries: EffectivePricingEntry[] = []; + let offset = 0; + let total: number | undefined; + while (true) { + const page = await this.request("usage.query", { + kind: "snapshot_pricing", + revision, + offset, + limit: PRICING_PAGE_MAX_ITEMS, + }); + if (page.kind === "revision_changed") { + if (page.expectedRevision !== revision) throw invalidProjection("Usage snapshot revision"); + return undefined; + } + if ( + page.kind !== "snapshot_pricing" || + page.revision !== revision || + page.offset !== offset || + page.entries.length > PRICING_PAGE_MAX_ITEMS || + entries.length !== offset + ) { + throw invalidProjection("Usage snapshot pricing"); + } + total ??= page.total; + if (page.total !== total) throw invalidProjection("Usage snapshot pricing"); + entries.push(...page.entries); + if (entries.length > total) throw invalidProjection("Usage snapshot pricing"); + if (page.nextOffset === null) { + if (entries.length !== total || !pricingEntriesAreCanonical(entries)) { + throw invalidProjection("Usage snapshot pricing"); + } + return entries; + } + if ( + page.entries.length === 0 || + page.nextOffset !== offset + page.entries.length || + page.nextOffset >= total + ) { + throw invalidProjection("Usage snapshot pricing"); + } + offset = page.nextOffset; + } + } + async #reconcilePricingMutation( target: PricingReconciliationTarget, reason: "revision_conflict" | "outcome_unknown", diff --git a/apps/desktop/src/main/runtime-host-usage-ipc-main.ts b/apps/desktop/src/main/runtime-host-usage-ipc-main.ts index 91a162a0ad..cf46b7e1b2 100644 --- a/apps/desktop/src/main/runtime-host-usage-ipc-main.ts +++ b/apps/desktop/src/main/runtime-host-usage-ipc-main.ts @@ -26,7 +26,6 @@ import { } from "@maka/core/usage-stats/pricing"; import type { PricingConfig, - TimeRange, UsageGroupBy, UsageQuery, } from "@maka/core/usage-stats/types"; @@ -48,8 +47,6 @@ interface RuntimeHostUsageIpcDeps { readonly sendToRenderer: (channel: string, ...args: unknown[]) => void; } -const MAX_ACTIVITY_RECORDS = 50_000; - export function registerRuntimeHostUsageIpc( deps: RuntimeHostUsageIpcDeps, ): void { @@ -160,39 +157,24 @@ async function loadUsageStats( client: DesktopRuntimeHostClient, range: UsageRange, ): Promise { - const query = { range: resolveUsageRange(range, Date.now()) } satisfies UsageQuery; - const [summaryResult, llmResult, toolResult, pricing] = await Promise.all([ - client.queryUsage({ kind: "summary", query }), - loadAllLogs(client, "llm", query), - loadAllLogs(client, "tool", query), - client.loadPricingSnapshot(), - ]); - if (summaryResult.kind !== "summary") throw invalidUsageProjection(); - const llmLogs = llmResult.rows; - const toolLogs = toolResult.rows; - const logsTruncated = llmResult.truncated || toolResult.truncated; - // The canonical summary is the authoritative headline count. We no longer - // throw when it disagrees with the number of activity rows we managed to - // load: a Host restart with pending repairs can make the summary read land - // before a catch-up commits and the logs read land after, and truncation - // (above) deliberately shortens the list. Either way the summary total stays - // correct; `provenance`/`logsTruncated` tell the page the activity list may - // be incomplete instead of erroring the whole page. + const snapshot = await client.loadUsageSnapshot(resolveUsageRange(range, Date.now())); + const llmLogs = snapshot.llmLogs; + const toolLogs = snapshot.toolLogs; + const logsTruncated = snapshot.llmLogsTruncated || snapshot.toolLogsTruncated; return { summary: { - totalRequests: summaryResult.summary.totalRequests, - totalCostUsd: summaryResult.summary.totalCostUsd, - totalTokens: summaryResult.summary.totalTokens.total, - inputTokens: summaryResult.summary.totalTokens.input, - outputTokens: summaryResult.summary.totalTokens.output, + totalRequests: snapshot.summary.totalRequests, + totalCostUsd: snapshot.summary.totalCostUsd, + totalTokens: snapshot.summary.totalTokens.total, + inputTokens: snapshot.summary.totalTokens.input, + outputTokens: snapshot.summary.totalTokens.output, cacheTokens: - summaryResult.summary.totalTokens.cacheRead + - summaryResult.summary.totalTokens.cacheWrite, - cacheMiss: summaryResult.summary.totalTokens.cacheMiss, - cacheRead: summaryResult.summary.totalTokens.cacheRead, - cacheCreation: summaryResult.summary.totalTokens.cacheWrite, - reasoning: summaryResult.summary.totalTokens.reasoning, + snapshot.summary.totalTokens.cacheRead + snapshot.summary.totalTokens.cacheWrite, + cacheMiss: snapshot.summary.totalTokens.cacheMiss, + cacheRead: snapshot.summary.totalTokens.cacheRead, + cacheCreation: snapshot.summary.totalTokens.cacheWrite, + reasoning: snapshot.summary.totalTokens.reasoning, }, logs: [...llmLogs.map(projectLlmLog), ...toolLogs.map(projectToolLog)].sort( (left, right) => right.ts - left.ts, @@ -200,81 +182,18 @@ async function loadUsageStats( byProvider: aggregateModelLogs(llmLogs, "provider"), byModel: aggregateModelLogs(llmLogs, "model"), byTool: aggregateToolLogs(toolLogs), - pricing: pricing.entries + pricing: snapshot.pricingEntries .filter((entry) => entry.source === "custom") .map(({ pricing: entry }) => projectPricing(entry)) .sort( (left, right) => left.provider.localeCompare(right.provider) || left.model.localeCompare(right.model), ), - provenance: summaryResult.provenance, + provenance: snapshot.provenance, ...(logsTruncated ? { logsTruncated: true } : {}), }; } -async function loadAllLogs( - client: DesktopRuntimeHostClient, - source: "llm", - query: UsageQuery & { range: TimeRange }, -): Promise<{ rows: LlmUsageLogProjection[]; truncated: boolean }>; -async function loadAllLogs( - client: DesktopRuntimeHostClient, - source: "tool", - query: UsageQuery & { range: TimeRange }, -): Promise<{ rows: ToolUsageLogProjection[]; truncated: boolean }>; -async function loadAllLogs( - client: DesktopRuntimeHostClient, - source: "llm" | "tool", - query: UsageQuery & { range: TimeRange }, -): Promise<{ - rows: Array; - truncated: boolean; -}> { - const rows: Array = []; - let offset = 0; - let total: number | undefined; - while (true) { - const result = await client.queryUsage( - source === "llm" - ? { - kind: "logs", - source, - query: toLlmQuery(query), - offset, - limit: USAGE_PAGE_MAX_ITEMS, - } - : { - kind: "logs", - source, - query: toToolQuery(query), - offset, - limit: USAGE_PAGE_MAX_ITEMS, - }, - ); - if (result.kind !== "logs" || result.source !== source || result.offset !== offset) { - throw invalidUsageProjection(); - } - total ??= result.total; - if (result.total !== total) throw invalidUsageProjection(); - rows.push(...result.rows); - // Structural integrity: the Host must never return more rows than it claims. - if (rows.length > total) throw invalidUsageProjection(); - // Client-side cap: when a range holds more activity than we render, keep the - // newest MAX_ACTIVITY_RECORDS and stop paging. This is truncation, not a - // protocol error, and the exhaustiveness check below is skipped for it — the - // caller surfaces `logsTruncated` so the page can say the list is partial. - if (total > MAX_ACTIVITY_RECORDS && rows.length >= MAX_ACTIVITY_RECORDS) { - return { rows: rows.slice(0, MAX_ACTIVITY_RECORDS), truncated: true }; - } - if (result.nextOffset === null) { - if (rows.length !== total) throw invalidUsageProjection(); - return { rows, truncated: false }; - } - if (result.nextOffset <= offset) throw invalidUsageProjection(); - offset = result.nextOffset; - } -} - function projectLlmLog(row: LlmUsageLogProjection): UsageStats["logs"][number] { return { id: row.id, diff --git a/packages/runtime-host/src/__tests__/handshake-compatibility.test.ts b/packages/runtime-host/src/__tests__/handshake-compatibility.test.ts index 69e9dddd9d..fbe9f208fa 100644 --- a/packages/runtime-host/src/__tests__/handshake-compatibility.test.ts +++ b/packages/runtime-host/src/__tests__/handshake-compatibility.test.ts @@ -108,7 +108,7 @@ test('receives structured incompatibility guidance from the released v0.1.11 Hos ); }); -test('rejects an epoch-39 Host before any domain command', async () => { +test('rejects an epoch-59 Host before any domain command', async () => { let admittedRequest: RequestFrame | undefined; await withForgedHandshakePeer( async (transport, hostEpoch, rootId) => { @@ -120,7 +120,7 @@ test('rejects an epoch-39 Host before any domain command', async () => { hostEpoch, connectionId: 'forged-epoch-connection', selectedProtocol: RUNTIME_HOST_PROTOCOL_VERSION, - compatibilityEpoch: 39, + compatibilityEpoch: 59, compositionId: 'maka.interactive', compositionRevision: '1', state: 'ready', diff --git a/packages/runtime-host/src/__tests__/usage-pricing-protocol.test.ts b/packages/runtime-host/src/__tests__/usage-pricing-protocol.test.ts index 9a5ed0dba2..66d6020ebe 100644 --- a/packages/runtime-host/src/__tests__/usage-pricing-protocol.test.ts +++ b/packages/runtime-host/src/__tests__/usage-pricing-protocol.test.ts @@ -50,6 +50,7 @@ import { type EffectivePricingEntry, type LlmUsageLogProjection, type ToolUsageLogProjection, + type UsageQueryResult, } from '../protocol/index.js'; import type { ConnectionContext } from '../server/operation-dispatcher.js'; import { HostUsagePricingCoordinator } from '../server/usage-pricing-coordinator.js'; @@ -171,6 +172,130 @@ describe('Usage/Pricing protocol', () => { ]) { assert.throws(() => usageRequest(input), invalidFrame); } + for (const result of [ + { + kind: 'snapshot_started', + revision: 'snapshot-revision-1', + summary: validSummary(), + provenance: validProvenance(), + extra: true, + }, + { + kind: 'snapshot_logs', + revision: 'snapshot-revision-1', + source: 'llm', + rows: [validLog()], + offset: 0, + total: 2, + nextOffset: 0, + truncated: false, + }, + { + kind: 'snapshot_logs', + revision: 'snapshot-revision-1', + source: 'llm', + rows: [validLog()], + offset: 0, + total: 1, + nextOffset: null, + truncated: 'no', + }, + { + kind: 'snapshot_pricing', + revision: 'snapshot-revision-1', + entries: Array.from({ length: PRICING_PAGE_MAX_ITEMS + 1 }, (_, index) => + customPricingEntry(`provider:model-${index}`), + ), + offset: 0, + total: PRICING_PAGE_MAX_ITEMS + 1, + nextOffset: null, + }, + { kind: 'revision_changed', expectedRevision: '' }, + ]) { + assert.throws(() => usageResponse(result), invalidFrame); + } + }); + + test('decodes revision-pinned Usage snapshot start, log, and pricing pages', () => { + assert.doesNotThrow(() => usageRequest({ kind: 'snapshot_start', range: { from: 1, to: 2 } })); + assert.doesNotThrow(() => + usageRequest({ + kind: 'snapshot_logs', + revision: 'snapshot-revision-1', + source: 'llm', + offset: 0, + limit: 3, + }), + ); + assert.doesNotThrow(() => + usageRequest({ + kind: 'snapshot_pricing', + revision: 'snapshot-revision-1', + offset: 0, + limit: 3, + }), + ); + + assert.doesNotThrow(() => + usageResponse({ + kind: 'snapshot_started', + revision: 'snapshot-revision-1', + summary: validSummary(), + provenance: validProvenance(), + }), + ); + assert.doesNotThrow(() => + usageResponse({ + kind: 'snapshot_logs', + revision: 'snapshot-revision-1', + source: 'llm', + rows: [validLog()], + offset: 0, + total: 1, + nextOffset: null, + truncated: false, + }), + ); + assert.doesNotThrow(() => + usageResponse({ + kind: 'snapshot_pricing', + revision: 'snapshot-revision-1', + entries: [customPricingEntry('provider:model')], + offset: 0, + total: 1, + nextOffset: null, + }), + ); + assert.doesNotThrow(() => + usageResponse({ kind: 'revision_changed', expectedRevision: 'snapshot-revision-1' }), + ); + + for (const input of [ + { kind: 'snapshot_start', range: 'all', revision: 'unexpected' }, + { kind: 'snapshot_logs', revision: '', source: 'llm', offset: 0, limit: 1 }, + { + kind: 'snapshot_logs', + revision: 'x'.repeat(129), + source: 'llm', + offset: 0, + limit: 1, + }, + { + kind: 'snapshot_logs', + revision: 'snapshot-revision-1', + source: 'model', + offset: 0, + limit: 1, + }, + { + kind: 'snapshot_pricing', + revision: 'snapshot-revision-1', + offset: 0, + limit: PRICING_PAGE_MAX_ITEMS + 1, + }, + ]) { + assert.throws(() => usageRequest(input), invalidFrame); + } }); test('enforces exact usage results and both page bounds', () => { @@ -441,6 +566,90 @@ describe('Usage/Pricing protocol', () => { } }); + test('pins every Usage authority behind one expiring LRU snapshot revision', async () => { + const base = await mkdtemp(join(tmpdir(), 'maka-usage-snapshot-')); + const capability = await resolveStorageRoot({ + path: join(base, 'interactive-root'), + kind: 'interactive', + }); + const owner = await tryAcquireInteractiveRootOwner(capability); + assert.ok(owner); + const stores = await openInteractiveUsageStoresForWrite(owner.lease); + let now = 1_000; + let nextRevision = 0; + try { + await stores.telemetry.recordLlmCall(longUsageRecord('old-llm', 1)); + await stores.telemetry.recordToolInvocation(longToolRecord('old-tool', 1)); + await stores.pricing.upsert(0, pricing('snapshot:old')); + const coordinator = new HostUsagePricingCoordinator( + stores, + () => {}, + new RuntimePolicyActivationGate(), + () => {}, + { + now: () => now, + createRevision: () => `snapshot-${++nextRevision}`, + ttlMs: 100, + capacity: 2, + activityLimit: 1, + }, + ); + + const first = await expectUsageSnapshotStart(coordinator); + assert.equal(first.revision, 'snapshot-1'); + assert.equal(first.summary.totalRequests, 1); + + await stores.telemetry.recordLlmCall(longUsageRecord('new-llm', 2)); + await stores.telemetry.recordToolInvocation(longToolRecord('new-tool', 2)); + await stores.pricing.upsert(1, pricing('snapshot:new')); + + const oldLlm = await expectUsageSnapshotLogs(coordinator, first.revision, 'llm'); + const oldTool = await expectUsageSnapshotLogs(coordinator, first.revision, 'tool'); + const oldPricing = await expectUsageSnapshotPricing(coordinator, first.revision); + assert.deepEqual( + oldLlm.rows.map((row) => row.id), + ['old-llm'], + ); + assert.deepEqual( + oldTool.rows.map((row) => row.id), + ['old-tool'], + ); + assert.equal(oldLlm.total, 1); + assert.equal(oldLlm.truncated, false); + assert.ok(oldPricing.entries.some((entry) => entry.pricing.modelKey === 'snapshot:old')); + assert.ok(!oldPricing.entries.some((entry) => entry.pricing.modelKey === 'snapshot:new')); + + const second = await expectUsageSnapshotStart(coordinator); + const newLlm = await expectUsageSnapshotLogs(coordinator, second.revision, 'llm'); + assert.deepEqual( + newLlm.rows.map((row) => row.id), + ['new-llm'], + ); + assert.equal(newLlm.total, 1, 'total describes retained rows'); + assert.equal(newLlm.truncated, true, 'truncation describes discarded authority rows'); + + await expectUsageSnapshotLogs(coordinator, first.revision, 'llm'); + await expectUsageSnapshotStart(coordinator); + assert.equal( + (await queryUsageSnapshotLogs(coordinator, second.revision, 'llm')).kind, + 'revision_changed', + 'the least recently used snapshot is evicted', + ); + + now += 101; + const expired = await queryUsageSnapshotLogs(coordinator, first.revision, 'llm'); + assert.deepEqual(expired, { kind: 'revision_changed', expectedRevision: first.revision }); + } finally { + await stores.close().catch(() => undefined); + await owner.close(); + await rm(join(resolveRootControlNamespace(), capability.rootId), { + recursive: true, + force: true, + }); + await rm(base, { recursive: true, force: true }); + } + }); + test('decodes revision-pinned numeric-offset pricing pages and revision-CAS mutation', () => { assert.doesNotThrow(() => pricingRequest('pricing.query', { kind: 'start' })); assert.doesNotThrow(() => @@ -870,6 +1079,78 @@ async function queryUsageBuckets( return frame.result.buckets; } +async function expectUsageSnapshotStart( + coordinator: HostUsagePricingCoordinator, +): Promise> { + const outcome = await coordinator.handlers['usage.query']( + { kind: 'snapshot_start', range: 'all' }, + CONNECTION_CONTEXT, + ); + assert.equal(outcome.ok, true); + if (!outcome.ok || outcome.result.kind !== 'snapshot_started') { + throw new Error('Expected a started Usage snapshot'); + } + return outcome.result; +} + +async function queryUsageSnapshotLogs( + coordinator: HostUsagePricingCoordinator, + revision: string, + source: 'llm' | 'tool', +): Promise> { + const outcome = await coordinator.handlers['usage.query']( + { kind: 'snapshot_logs', revision, source, offset: 0, limit: USAGE_PAGE_MAX_ITEMS }, + CONNECTION_CONTEXT, + ); + assert.equal(outcome.ok, true); + if ( + !outcome.ok || + (outcome.result.kind !== 'snapshot_logs' && outcome.result.kind !== 'revision_changed') + ) { + throw new Error('Expected a Usage snapshot log page'); + } + return outcome.result; +} + +async function expectUsageSnapshotLogs( + coordinator: HostUsagePricingCoordinator, + revision: string, + source: 'llm' | 'tool', +): Promise> { + const result = await queryUsageSnapshotLogs(coordinator, revision, source); + if (result.kind !== 'snapshot_logs') throw new Error('Expected a retained Usage snapshot'); + assert.equal(result.source, source); + return result; +} + +async function expectUsageSnapshotPricing( + coordinator: HostUsagePricingCoordinator, + revision: string, +): Promise<{ readonly entries: readonly EffectivePricingEntry[] }> { + const entries: EffectivePricingEntry[] = []; + let offset = 0; + let total: number | undefined; + do { + const outcome = await coordinator.handlers['usage.query']( + { kind: 'snapshot_pricing', revision, offset, limit: PRICING_PAGE_MAX_ITEMS }, + CONNECTION_CONTEXT, + ); + assert.equal(outcome.ok, true); + if (!outcome.ok || outcome.result.kind !== 'snapshot_pricing') { + throw new Error('Expected a Usage snapshot pricing page'); + } + assert.equal(outcome.result.revision, revision); + assert.equal(outcome.result.offset, offset); + total ??= outcome.result.total; + assert.equal(outcome.result.total, total); + entries.push(...outcome.result.entries); + if (outcome.result.nextOffset === null) break; + offset = outcome.result.nextOffset; + } while (true); + assert.equal(entries.length, total); + return { entries }; +} + function assertDistinctBoundedIdentities(values: readonly (string | undefined)[]): void { assert.equal(values.length, 6); assert.ok(values.every((value): value is string => typeof value === 'string')); diff --git a/packages/runtime-host/src/__tests__/usage-pricing-two-client-uds.test.ts b/packages/runtime-host/src/__tests__/usage-pricing-two-client-uds.test.ts index ede75459dd..e45b9d718d 100644 --- a/packages/runtime-host/src/__tests__/usage-pricing-two-client-uds.test.ts +++ b/packages/runtime-host/src/__tests__/usage-pricing-two-client-uds.test.ts @@ -447,6 +447,7 @@ describe('production Usage/Pricing UDS', () => { | undefined; const clients: RuntimeHostConnection[] = []; let endpoint: string | undefined; + let firstHostSnapshotRevision: string | undefined; try { firstOwner = await tryAcquireInteractiveRootOwner(capability); @@ -504,6 +505,15 @@ describe('production Usage/Pricing UDS', () => { }, ]); + const pinnedUsage = await desktop.request( + 'usage.query', + { kind: 'snapshot_start', range: 'all' }, + REQUEST_TIMEOUT_MS, + ); + assert.equal(pinnedUsage.kind, 'snapshot_started'); + if (pinnedUsage.kind !== 'snapshot_started') throw new Error('Usage snapshot did not start'); + firstHostSnapshotRevision = pinnedUsage.revision; + const initial = await readPricing(desktop); assert.equal(initial.revision, 0); assert.deepEqual(initial.entries, builtinPricingEntries()); @@ -560,6 +570,27 @@ describe('production Usage/Pricing UDS', () => { ); assert.deepEqual(retry, { kind: 'committed', revision: 2 }); + const pinnedPricing = await readUsageSnapshotPricing(desktop, pinnedUsage.revision); + assert.deepEqual(pinnedPricing, builtinPricingEntries()); + const pinnedLogs = await desktop.request( + 'usage.query', + { + kind: 'snapshot_logs', + revision: pinnedUsage.revision, + source: 'llm', + offset: 0, + limit: 100, + }, + REQUEST_TIMEOUT_MS, + ); + assert.equal(pinnedLogs.kind, 'snapshot_logs'); + if (pinnedLogs.kind === 'snapshot_logs') { + assert.deepEqual( + pinnedLogs.rows.map((row) => row.id), + ['usage-b', 'usage-a'], + ); + } + const [desktopPricing, tuiPricing] = await Promise.all([ readPricing(desktop), readPricing(tui), @@ -662,6 +693,21 @@ describe('production Usage/Pricing UDS', () => { connectClient(root), ]); clients.push(desktopAfterRestart, tuiAfterRestart); + assert.ok(firstHostSnapshotRevision); + assert.deepEqual( + await desktopAfterRestart.request( + 'usage.query', + { + kind: 'snapshot_logs', + revision: firstHostSnapshotRevision, + source: 'llm', + offset: 0, + limit: 100, + }, + REQUEST_TIMEOUT_MS, + ), + { kind: 'revision_changed', expectedRevision: firstHostSnapshotRevision }, + ); const [usageAfterRestart, pricingAfterRestart, pricingFromSecondClient] = await Promise.all([ readUsage(desktopAfterRestart), readPricing(desktopAfterRestart), @@ -772,6 +818,31 @@ async function readPricing(client: RuntimeHostConnection): Promise<{ return { revision: first.revision, entries, pageCount }; } +async function readUsageSnapshotPricing( + client: RuntimeHostConnection, + revision: string, +): Promise { + const entries: EffectivePricingEntry[] = []; + let offset = 0; + while (true) { + const result = await client.request( + 'usage.query', + { kind: 'snapshot_pricing', revision, offset, limit: 128 }, + REQUEST_TIMEOUT_MS, + ); + assert.equal(result.kind, 'snapshot_pricing'); + if (result.kind !== 'snapshot_pricing') throw new Error('Usage snapshot pricing disappeared'); + assert.equal(result.revision, revision); + assert.equal(result.offset, offset); + entries.push(...result.entries); + if (result.nextOffset === null) { + assert.equal(entries.length, result.total); + return entries; + } + offset = result.nextOffset; + } +} + async function readCoordinatorPricing( coordinator: HostUsagePricingCoordinator, ): Promise> { diff --git a/packages/runtime-host/src/protocol/index.ts b/packages/runtime-host/src/protocol/index.ts index a02f619545..0127a7740e 100644 --- a/packages/runtime-host/src/protocol/index.ts +++ b/packages/runtime-host/src/protocol/index.ts @@ -93,7 +93,10 @@ export const RUNTIME_HOST_REGISTRATION_SCHEMA_VERSION = 1 as const; export const RUNTIME_HOST_PROTOCOL_VERSION = 0 as const; // Increment when the same protocol version no longer guarantees safe Client-Host // interoperability. Mismatches are rejected before domain commands are admitted. -export const RUNTIME_HOST_COMPATIBILITY_EPOCH = 59 as const; +export const RUNTIME_HOST_COMPATIBILITY_EPOCH = 60 as const; +// 60: `usage.query` adds opaque revision-pinned snapshot start, activity, and +// pricing pages. Epoch-59 peers reject these exact new input/output variants, +// so mixed peers must fail the handshake before Settings Usage is requested. // 59: Scheduled Turn provider-retry frames may carry an optional host-clock // `ts`, letting a mid-wait re-projection recompute the authoritative // remaining duration. Older peers decode the frame with an exact key list diff --git a/packages/runtime-host/src/protocol/usage-pricing.ts b/packages/runtime-host/src/protocol/usage-pricing.ts index 266fa11bab..56c3c89290 100644 --- a/packages/runtime-host/src/protocol/usage-pricing.ts +++ b/packages/runtime-host/src/protocol/usage-pricing.ts @@ -34,7 +34,7 @@ import type { } from '@maka/core/usage-stats/types'; import { MODEL_CALL_KINDS } from '@maka/core/usage-stats/types'; import type { UsageProvenance } from '@maka/core/usage-ledger-merge'; -import { requireCount, requireExactRecord, requireRecord } from './codec.js'; +import { requireCount, requireExactRecord, requireId, requireRecord } from './codec.js'; import { invalidProtocolFrame } from './errors.js'; import { defineOperation } from './operation-spec.js'; @@ -193,6 +193,20 @@ export interface ToolUsageLogProjection { export type UsageLogProjection = LlmUsageLogProjection | ToolUsageLogProjection; export type UsageQueryInput = + | { readonly kind: 'snapshot_start'; readonly range: UsageQuery['range'] } + | { + readonly kind: 'snapshot_logs'; + readonly revision: string; + readonly source: 'llm' | 'tool'; + readonly offset?: number; + readonly limit?: number; + } + | { + readonly kind: 'snapshot_pricing'; + readonly revision: string; + readonly offset?: number; + readonly limit?: number; + } | { readonly kind: 'summary'; readonly query: LlmUsageQuery } | { readonly kind: 'buckets'; @@ -224,6 +238,41 @@ export type UsageQueryInput = }; export type UsageQueryResult = + | { + readonly kind: 'snapshot_started'; + readonly revision: string; + readonly summary: UsageSummaryV2; + readonly provenance: UsageProvenance; + } + | { + readonly kind: 'snapshot_logs'; + readonly revision: string; + readonly source: 'llm'; + readonly rows: readonly LlmUsageLogProjection[]; + readonly offset: number; + readonly total: number; + readonly nextOffset: number | null; + readonly truncated: boolean; + } + | { + readonly kind: 'snapshot_logs'; + readonly revision: string; + readonly source: 'tool'; + readonly rows: readonly ToolUsageLogProjection[]; + readonly offset: number; + readonly total: number; + readonly nextOffset: number | null; + readonly truncated: boolean; + } + | { + readonly kind: 'snapshot_pricing'; + readonly revision: string; + readonly entries: readonly EffectivePricingEntry[]; + readonly offset: number; + readonly total: number; + readonly nextOffset: number | null; + } + | { readonly kind: 'revision_changed'; readonly expectedRevision: string } | { readonly kind: 'summary'; readonly summary: UsageSummaryV2; @@ -346,6 +395,42 @@ export const USAGE_PRICING_OPERATION_SPECS = { export function decodeUsageQueryInput(value: unknown): UsageQueryInput { const input = requireRecord(value, 'usage query input'); + if (input.kind === 'snapshot_start') { + const exact = requireExactRecord(input, 'usage snapshot start input', ['kind', 'range']); + return { kind: 'snapshot_start', range: decodeUsageRange(exact.range) }; + } + if (input.kind === 'snapshot_logs') { + assertOptionalExactKeys( + input, + 'usage snapshot logs input', + ['kind', 'revision', 'source'], + ['offset', 'limit'], + ); + if (input.source !== 'llm' && input.source !== 'tool') { + throw invalidProtocolFrame('Invalid usage snapshot log source'); + } + return { + kind: 'snapshot_logs', + revision: requireId(input.revision, 'usage snapshot revision'), + source: input.source, + offset: decodeOffset(input.offset), + limit: decodeLimit(input.limit), + }; + } + if (input.kind === 'snapshot_pricing') { + assertOptionalExactKeys( + input, + 'usage snapshot pricing input', + ['kind', 'revision'], + ['offset', 'limit'], + ); + return { + kind: 'snapshot_pricing', + revision: requireId(input.revision, 'usage snapshot revision'), + offset: decodeOffset(input.offset), + limit: decodePricingLimit(input.limit), + }; + } if (input.kind === 'summary') { const exact = requireExactRecord(input, 'usage summary input', ['kind', 'query']); return { kind: 'summary', query: decodeLlmUsageQuery(exact.query) }; @@ -394,6 +479,60 @@ export function decodeUsageQueryInput(value: unknown): UsageQueryInput { export function decodeUsageQueryResult(value: unknown): UsageQueryResult { const result = requireRecord(value, 'usage query result'); + if (result.kind === 'snapshot_started') { + const exact = requireExactRecord(result, 'usage snapshot started result', [ + 'kind', + 'revision', + 'summary', + 'provenance', + ]); + return { + kind: 'snapshot_started', + revision: requireId(exact.revision, 'usage snapshot revision'), + summary: decodeUsageSummary(exact.summary), + provenance: decodeUsageProvenance(exact.provenance), + }; + } + if (result.kind === 'snapshot_logs') { + const exact = requireExactRecord(result, 'usage snapshot logs result', [ + 'kind', + 'revision', + 'source', + 'rows', + 'offset', + 'total', + 'nextOffset', + 'truncated', + ]); + if (exact.source === 'llm') { + return decodeUsageSnapshotLogPage('llm', exact, decodeLlmUsageLog); + } + if (exact.source === 'tool') { + return decodeUsageSnapshotLogPage('tool', exact, decodeToolUsageLog); + } + throw invalidProtocolFrame('Invalid usage snapshot log source'); + } + if (result.kind === 'snapshot_pricing') { + const exact = requireExactRecord(result, 'usage snapshot pricing result', [ + 'kind', + 'revision', + 'entries', + 'offset', + 'total', + 'nextOffset', + ]); + return decodeUsageSnapshotPricingPage(exact); + } + if (result.kind === 'revision_changed') { + const exact = requireExactRecord(result, 'usage snapshot revision changed result', [ + 'kind', + 'expectedRevision', + ]); + return { + kind: 'revision_changed', + expectedRevision: requireId(exact.expectedRevision, 'expected usage snapshot revision'), + }; + } if (result.kind === 'summary') { const exact = requireExactRecord(result, 'usage summary result', [ 'kind', @@ -606,6 +745,34 @@ export function decodePricingMutateResult(value: unknown): PricingMutateResult { } function assertUsageQueryOutputForInput(input: UsageQueryInput, output: UsageQueryResult): void { + if (input.kind === 'snapshot_start') { + if (output.kind !== 'snapshot_started') { + throw invalidProtocolFrame('Usage snapshot start response does not match its request'); + } + return; + } + if (input.kind === 'snapshot_logs' || input.kind === 'snapshot_pricing') { + if (output.kind === 'revision_changed') { + if (output.expectedRevision !== input.revision) { + throw invalidProtocolFrame('Usage snapshot revision change does not match its request'); + } + return; + } + if (output.kind !== input.kind) { + throw invalidProtocolFrame('Usage snapshot response kind does not match its request'); + } + if (output.revision !== input.revision || output.offset !== (input.offset ?? 0)) { + throw invalidProtocolFrame('Usage snapshot page does not match its request'); + } + if ( + input.kind === 'snapshot_logs' && + output.kind === 'snapshot_logs' && + output.source !== input.source + ) { + throw invalidProtocolFrame('Usage snapshot log source does not match its request'); + } + return; + } if (output.kind !== input.kind) { throw invalidProtocolFrame('Usage response kind does not match its request'); } @@ -729,6 +896,15 @@ function decodeLimit(value: unknown): number { return limit; } +function decodePricingLimit(value: unknown): number { + if (value === undefined) return PRICING_PAGE_MAX_ITEMS; + const limit = requireCount(value, 'usage snapshot pricing limit'); + if (limit === 0 || limit > PRICING_PAGE_MAX_ITEMS) { + throw invalidProtocolFrame('Invalid usage snapshot pricing limit'); + } + return limit; +} + function decodeUsagePage( kind: 'buckets', result: Record, @@ -787,6 +963,68 @@ function decodeUsageLogPage( return decoded; } +function decodeUsageSnapshotLogPage( + source: 'llm', + result: Record, + decodeItem: (value: unknown) => LlmUsageLogProjection, +): Extract; +function decodeUsageSnapshotLogPage( + source: 'tool', + result: Record, + decodeItem: (value: unknown) => ToolUsageLogProjection, +): Extract; +function decodeUsageSnapshotLogPage( + source: 'llm' | 'tool', + result: Record, + decodeItem: (value: unknown) => UsageLogProjection, +): Extract { + const rawItems = result.rows; + if (!Array.isArray(rawItems) || rawItems.length > USAGE_PAGE_MAX_ITEMS) { + throw invalidProtocolFrame('Usage snapshot page exceeds item limit'); + } + if (typeof result.truncated !== 'boolean') { + throw invalidProtocolFrame('Invalid usage snapshot truncation flag'); + } + const rows = rawItems.map(decodeItem); + const decoded = { + kind: 'snapshot_logs', + revision: requireId(result.revision, 'usage snapshot revision'), + source, + rows, + ...decodeUsagePagePosition(result, rows.length), + truncated: result.truncated, + } as Extract; + assertJsonBytes(decoded, USAGE_PAGE_MAX_BYTES, 'Usage snapshot page'); + return decoded; +} + +function decodeUsageSnapshotPricingPage( + result: Record, +): Extract { + const rawItems = result.entries; + if (!Array.isArray(rawItems) || rawItems.length > PRICING_PAGE_MAX_ITEMS) { + throw invalidProtocolFrame('Usage snapshot pricing page exceeds item limit'); + } + const entries = rawItems.map(decodeEffectivePricingEntry); + if ( + entries.some( + (item, index) => + index > 0 && + comparePricingModelKeys(entries[index - 1]!.pricing.modelKey, item.pricing.modelKey) >= 0, + ) + ) { + throw invalidProtocolFrame('Usage snapshot pricing entries are not canonically ordered'); + } + const decoded = { + kind: 'snapshot_pricing', + revision: requireId(result.revision, 'usage snapshot revision'), + entries, + ...decodeUsagePagePosition(result, entries.length), + } as const; + assertJsonBytes(decoded, PRICING_PAGE_MAX_BYTES, 'Usage snapshot pricing page'); + return decoded; +} + function decodeUsagePagePosition( result: Record, itemCount: number, diff --git a/packages/runtime-host/src/server/usage-pricing-coordinator.ts b/packages/runtime-host/src/server/usage-pricing-coordinator.ts index 01557bd642..c7adaf61ee 100644 --- a/packages/runtime-host/src/server/usage-pricing-coordinator.ts +++ b/packages/runtime-host/src/server/usage-pricing-coordinator.ts @@ -18,6 +18,7 @@ */ import { createHash } from 'node:crypto'; +import { resolveUsageRange } from '@maka/core/model-call-usage-projection'; import type { PricingConfig, ToolInvocationRecord, @@ -62,6 +63,7 @@ import { import type { UsagePricingOperationHandlerMap } from './operation-dispatcher.js'; import { RuntimePolicyActivationGate } from './runtime-policy-activation-gate.js'; import { readCanonicalUsage } from './canonical-usage-reader.js'; +import { UsageSnapshotCache, type UsageSnapshotCacheOptions } from './usage-snapshot-cache.js'; /** Root-scoped projection over the authentic lease-bound usage stores. */ export class HostUsagePricingCoordinator { @@ -75,6 +77,7 @@ export class HostUsagePricingCoordinator { readonly #requestDrain: () => void; readonly #activation: RuntimePolicyActivationGate; readonly #onCommittedPricingMutation: () => void; + readonly #usageSnapshots: UsageSnapshotCache; #poisonDrainRequested = false; constructor( @@ -82,11 +85,13 @@ export class HostUsagePricingCoordinator { requestDrain: () => void, activation: RuntimePolicyActivationGate, onCommittedPricingMutation: () => void = () => {}, + usageSnapshotOptions: UsageSnapshotCacheOptions = {}, ) { this.#stores = authenticateInteractiveUsageStoresWriter(stores); this.#requestDrain = requestDrain; this.#activation = activation; this.#onCommittedPricingMutation = onCommittedPricingMutation; + this.#usageSnapshots = new UsageSnapshotCache(usageSnapshotOptions); } /** @@ -104,6 +109,40 @@ export class HostUsagePricingCoordinator { async #queryUsage(input: UsageQueryInput): Promise> { try { const now = Date.now(); + if (input.kind === 'snapshot_start') { + return { ok: true, result: await this.#startUsageSnapshot(input.range, now) }; + } + if (input.kind === 'snapshot_logs') { + const snapshot = this.#usageSnapshots.get(input.revision); + if (!snapshot) return usageRevisionChanged(input.revision); + const rows = input.source === 'llm' ? snapshot.llmRows : snapshot.toolRows; + if ((input.offset ?? 0) > rows.length) return invalidUsageOffset(); + return { + ok: true, + result: usageSnapshotLogPage( + input.revision, + input.source, + rows, + input.offset ?? 0, + input.limit ?? USAGE_PAGE_MAX_ITEMS, + input.source === 'llm' ? snapshot.llmTruncated : snapshot.toolTruncated, + ), + }; + } + if (input.kind === 'snapshot_pricing') { + const snapshot = this.#usageSnapshots.get(input.revision); + if (!snapshot) return usageRevisionChanged(input.revision); + if ((input.offset ?? 0) > snapshot.pricingEntries.length) return invalidUsageOffset(); + return { + ok: true, + result: usageSnapshotPricingPage( + input.revision, + snapshot.pricingEntries, + input.offset ?? 0, + input.limit ?? PRICING_PAGE_MAX_ITEMS, + ), + }; + } if (input.kind === 'summary') { const merged = mergeUsageSummary( await this.#stores.telemetry.summary(input.query), @@ -192,6 +231,50 @@ export class HostUsagePricingCoordinator { } } + async #startUsageSnapshot( + range: UsageQuery['range'], + now: number, + ): Promise> { + const query: UsageQuery = { range: resolveUsageRange(range, now) }; + const captureLimit = this.#usageSnapshots.activityLimit + 1; + const captured = await this.#stores.captureUsageSnapshot({ + query, + activityLimit: captureLimit, + }); + const canonical: CanonicalUsageSource = { + attempts: captured.canonical.attempts, + unreadableRecords: captured.canonical.unreadableRecords + captured.repair.unreadableEvents, + pendingRepairs: captured.repair.pendingRuns, + }; + const mergedSummary = mergeUsageSummary(captured.legacySummary, canonical, query, now); + const { provenance, ...summary } = mergedSummary; + const mergedLogs = mergeUsageLogs( + captured.legacyLlmLogs, + canonical, + query, + now, + 0, + captureLimit, + ); + const retained = this.#usageSnapshots.retain({ + summary, + provenance, + llmRows: mergedLogs.rows.slice(0, this.#usageSnapshots.activityLimit).map(projectUsageLog), + llmTruncated: mergedLogs.total > this.#usageSnapshots.activityLimit, + toolRows: captured.toolLogs.rows + .slice(0, this.#usageSnapshots.activityLimit) + .map(projectToolUsageLog), + toolTruncated: captured.toolLogs.total > this.#usageSnapshots.activityLimit, + pricingEntries: projectEffectivePricingEntries(captured.pricing.overrides), + }); + return encodeUsageQueryResult({ + kind: 'snapshot_started', + revision: retained.revision, + summary: retained.summary, + provenance: retained.provenance, + }) as Extract; + } + async #queryPricing(input: PricingQueryInput): Promise> { try { const snapshot = await this.#stores.pricing.snapshot(); @@ -361,33 +444,111 @@ function invalidUsageOffset(): OperationOutcome<'usage.query'> { }; } +function usageRevisionChanged(revision: string): OperationOutcome<'usage.query'> { + return { + ok: true, + result: encodeUsageQueryResult({ kind: 'revision_changed', expectedRevision: revision }), + }; +} + +function usageSnapshotLogPage( + revision: string, + source: 'llm' | 'tool', + allRows: readonly UsageLogProjection[], + offset: number, + limit: number, + truncated: boolean, +): Extract { + const rows = fitBoundedPageItems( + allRows.slice(offset, offset + limit), + offset < allRows.length, + USAGE_PAGE_MAX_BYTES, + (candidate) => { + const nextOffset = offset + candidate.length; + return { + kind: 'snapshot_logs', + revision, + source, + rows: candidate, + offset, + total: allRows.length, + nextOffset: nextOffset < allRows.length ? nextOffset : null, + truncated, + } as const; + }, + 'Canonical Usage snapshot item', + ); + const nextOffset = offset + rows.length; + return encodeUsageQueryResult({ + kind: 'snapshot_logs', + revision, + source, + rows, + offset, + total: allRows.length, + nextOffset: nextOffset < allRows.length ? nextOffset : null, + truncated, + } as Extract) as Extract< + UsageQueryResult, + { kind: 'snapshot_logs' } + >; +} + +function usageSnapshotPricingPage( + revision: string, + allEntries: readonly EffectivePricingEntry[], + offset: number, + limit: number, +): Extract { + const entries = fitBoundedPageItems( + allEntries.slice(offset, offset + limit), + offset < allEntries.length, + PRICING_PAGE_MAX_BYTES, + (candidate) => { + const nextOffset = offset + candidate.length; + return { + kind: 'snapshot_pricing', + revision, + entries: candidate, + offset, + total: allEntries.length, + nextOffset: nextOffset < allEntries.length ? nextOffset : null, + } as const; + }, + 'Canonical Usage snapshot pricing entry', + ); + const nextOffset = offset + entries.length; + return encodeUsageQueryResult({ + kind: 'snapshot_pricing', + revision, + entries, + offset, + total: allEntries.length, + nextOffset: nextOffset < allEntries.length ? nextOffset : null, + }) as Extract; +} + function createPricingPage( revision: number, entries: readonly EffectivePricingEntry[], offset: number, ): PricingQueryResult { - const items: EffectivePricingEntry[] = []; - for (let index = offset; index < entries.length; index += 1) { - if (items.length >= PRICING_PAGE_MAX_ITEMS) break; - const item = entries[index]; - if (!item) break; - const candidate = [...items, item]; - const nextOffset = offset + candidate.length; - const page: PricingQueryResult = { - kind: 'page', - revision, - offset, - entries: candidate, - nextOffset: nextOffset < entries.length ? nextOffset : null, - }; - if (jsonBytes(page) > PRICING_PAGE_MAX_BYTES) { - if (items.length === 0) { - throw new Error('Canonical pricing entry exceeds the wire page limit'); - } - break; - } - items.push(item); - } + const items = fitBoundedPageItems( + entries.slice(offset, offset + PRICING_PAGE_MAX_ITEMS), + offset < entries.length, + PRICING_PAGE_MAX_BYTES, + (candidate) => { + const nextOffset = offset + candidate.length; + return { + kind: 'page', + revision, + offset, + entries: candidate, + nextOffset: nextOffset < entries.length ? nextOffset : null, + } satisfies PricingQueryResult; + }, + 'Canonical pricing entry', + ); const nextOffset = offset + items.length; return encodePricingQueryResult({ kind: 'page', @@ -427,28 +588,22 @@ function usagePage( provenance: UsageProvenance, ): Extract { const source = allItems.slice(offset, offset + limit); - const items: UsageBucket[] = []; - for (const item of source) { - const candidate = [...items, item]; - const nextOffset = offset + candidate.length; - if ( - jsonBytes( - bucketPageResult( - candidate, - total, - offset, - nextOffset < total ? nextOffset : null, - provenance, - ), - ) > USAGE_PAGE_MAX_BYTES - ) { - break; - } - items.push(item); - } - if (items.length === 0 && offset < total) { - throw new Error('Canonical usage item exceeds the wire page limit'); - } + const items = fitBoundedPageItems( + source, + offset < total, + USAGE_PAGE_MAX_BYTES, + (candidate) => { + const nextOffset = offset + candidate.length; + return bucketPageResult( + candidate, + total, + offset, + nextOffset < total ? nextOffset : null, + provenance, + ); + }, + 'Canonical usage item', + ); const nextOffset = offset + items.length; return bucketPageResult(items, total, offset, nextOffset < total ? nextOffset : null, provenance); } @@ -486,29 +641,23 @@ function usageLogPage( limit: number, provenance?: UsageProvenance, ): Extract { - const items: UsageLogProjection[] = []; - for (const item of allItems.slice(0, limit)) { - const candidate = [...items, item]; - const nextOffset = offset + candidate.length; - if ( - jsonBytes( - logPageResult( - source, - candidate, - total, - offset, - nextOffset < total ? nextOffset : null, - provenance, - ), - ) > USAGE_PAGE_MAX_BYTES - ) { - break; - } - items.push(item); - } - if (items.length === 0 && offset < total) { - throw new Error('Canonical usage item exceeds the wire page limit'); - } + const items = fitBoundedPageItems( + allItems.slice(0, limit), + offset < total, + USAGE_PAGE_MAX_BYTES, + (candidate) => { + const nextOffset = offset + candidate.length; + return logPageResult( + source, + candidate, + total, + offset, + nextOffset < total ? nextOffset : null, + provenance, + ); + }, + 'Canonical usage item', + ); const nextOffset = offset + items.length; return logPageResult( source, @@ -520,6 +669,25 @@ function usageLogPage( ); } +function fitBoundedPageItems( + candidates: readonly T[], + itemRequired: boolean, + maxBytes: number, + createPage: (items: readonly T[]) => unknown, + itemLabel: string, +): T[] { + const items: T[] = []; + for (const item of candidates) { + const next = [...items, item]; + if (jsonBytes(createPage(next)) > maxBytes) break; + items.push(item); + } + if (items.length === 0 && itemRequired) { + throw new Error(`${itemLabel} exceeds the wire page limit`); + } + return items; +} + function logPageResult( source: 'llm' | 'tool', rows: readonly UsageLogProjection[], diff --git a/packages/runtime-host/src/server/usage-snapshot-cache.ts b/packages/runtime-host/src/server/usage-snapshot-cache.ts new file mode 100644 index 0000000000..07e9ab1ef9 --- /dev/null +++ b/packages/runtime-host/src/server/usage-snapshot-cache.ts @@ -0,0 +1,124 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import { randomUUID } from 'node:crypto'; +import type { UsageSummaryV2 } from '@maka/core/usage-stats/types'; +import type { UsageProvenance } from '@maka/core/usage-ledger-merge'; +import type { + EffectivePricingEntry, + LlmUsageLogProjection, + ToolUsageLogProjection, +} from '../protocol/index.js'; + +export const USAGE_SNAPSHOT_TTL_MS = 5 * 60 * 1_000; +export const USAGE_SNAPSHOT_CAPACITY = 4; +export const USAGE_SNAPSHOT_ACTIVITY_LIMIT = 50_000; + +export interface UsageSnapshotCacheOptions { + readonly now?: () => number; + readonly createRevision?: () => string; + readonly ttlMs?: number; + readonly capacity?: number; + readonly activityLimit?: number; +} + +export interface UsageSnapshotContents { + readonly summary: UsageSummaryV2; + readonly provenance: UsageProvenance; + readonly llmRows: readonly LlmUsageLogProjection[]; + readonly llmTruncated: boolean; + readonly toolRows: readonly ToolUsageLogProjection[]; + readonly toolTruncated: boolean; + readonly pricingEntries: readonly EffectivePricingEntry[]; +} + +export interface RetainedUsageSnapshot extends UsageSnapshotContents { + readonly revision: string; +} + +interface CacheEntry extends RetainedUsageSnapshot { + readonly expiresAt: number; +} + +/** Host-epoch-local, absolute-TTL cache for coherent Settings Usage reads. */ +export class UsageSnapshotCache { + readonly activityLimit: number; + readonly #now: () => number; + readonly #createRevision: () => string; + readonly #ttlMs: number; + readonly #capacity: number; + readonly #entries = new Map(); + + constructor(options: UsageSnapshotCacheOptions = {}) { + this.#now = options.now ?? Date.now; + this.#createRevision = options.createRevision ?? randomUUID; + this.#ttlMs = options.ttlMs ?? USAGE_SNAPSHOT_TTL_MS; + this.#capacity = options.capacity ?? USAGE_SNAPSHOT_CAPACITY; + this.activityLimit = options.activityLimit ?? USAGE_SNAPSHOT_ACTIVITY_LIMIT; + if ( + !Number.isSafeInteger(this.#ttlMs) || + this.#ttlMs <= 0 || + !Number.isSafeInteger(this.#capacity) || + this.#capacity <= 0 || + !Number.isSafeInteger(this.activityLimit) || + this.activityLimit <= 0 + ) { + throw new TypeError('Invalid Usage snapshot cache limits'); + } + } + + retain(contents: UsageSnapshotContents): RetainedUsageSnapshot { + const now = this.#now(); + this.#pruneExpired(now); + while (this.#entries.size >= this.#capacity) { + const oldestRevision = this.#entries.keys().next().value; + if (oldestRevision === undefined) break; + this.#entries.delete(oldestRevision); + } + const revision = this.#createRevision(); + if (revision.length === 0 || revision.length > 128 || this.#entries.has(revision)) { + throw new Error('Usage snapshot revision generator returned an invalid revision'); + } + const entry: CacheEntry = { + revision, + ...contents, + expiresAt: now + this.#ttlMs, + }; + this.#entries.set(revision, entry); + return entry; + } + + get(revision: string): RetainedUsageSnapshot | undefined { + const now = this.#now(); + this.#pruneExpired(now); + const entry = this.#entries.get(revision); + if (!entry) return undefined; + // Map insertion order is the LRU order. Reinsert without changing expiresAt: + // page access affects eviction priority, never the absolute lifetime. + this.#entries.delete(revision); + this.#entries.set(revision, entry); + return entry; + } + + #pruneExpired(now: number): void { + for (const [revision, entry] of this.#entries) { + if (entry.expiresAt <= now) this.#entries.delete(revision); + } + } +} diff --git a/packages/storage/src/__tests__/usage-stores.test.ts b/packages/storage/src/__tests__/usage-stores.test.ts index c00e6f0b5b..6ac04f8fab 100644 --- a/packages/storage/src/__tests__/usage-stores.test.ts +++ b/packages/storage/src/__tests__/usage-stores.test.ts @@ -335,6 +335,44 @@ describe('InteractiveUsageStores', () => { }); }); + test('captures one repaired Usage authority snapshot behind the writer lease', async () => { + await withInteractiveRoot(async ({ root, capability }) => { + const owner = await tryAcquireInteractiveRootOwner(capability); + assert(owner); + const stores = await openInteractiveUsageStoresForWrite(owner.lease); + try { + await stores.telemetry.recordLlmCall( + llmRecord({ id: 'legacy-snapshot', sessionId: 'session-legacy' }), + ); + await stores.telemetry.recordToolInvocation(toolRecord()); + appendModelCallAuthorityEvent(root, modelCallAttempt('session-canonical')); + const pricing = { + modelKey: 'openai:gpt-5', + inputUsdPer1M: 1, + outputUsdPer1M: 2, + }; + await stores.pricing.upsert(0, pricing); + + const snapshot = await stores.captureUsageSnapshot({ + query: { range: 'all' }, + activityLimit: 10, + }); + + assert.equal(snapshot.legacySummary.totalRequests, 1); + assert.equal(snapshot.legacyLlmLogs.total, 1); + assert.equal(snapshot.legacyLlmLogs.rows[0]?.id, 'legacy-snapshot'); + assert.equal(snapshot.toolLogs.total, 1); + assert.equal(snapshot.toolLogs.rows[0]?.id, 'tool_1'); + assert.equal(snapshot.canonical.attempts[0]?.sessionId, 'session-canonical'); + assert.equal(snapshot.repair.pendingRuns, 0); + assert.deepEqual(snapshot.pricing, { revision: 1, overrides: [pricing] }); + } finally { + await stores.close(); + await owner.close(); + } + }); + }); + test('legacy summary clamps each cache reading to its own input', async () => { await withInteractiveRoot(async ({ capability }) => { const owner = await tryAcquireInteractiveRootOwner(capability); diff --git a/packages/storage/src/model-call-ledger.ts b/packages/storage/src/model-call-ledger.ts index 19749e1d40..c9a5153036 100644 --- a/packages/storage/src/model-call-ledger.ts +++ b/packages/storage/src/model-call-ledger.ts @@ -120,6 +120,21 @@ export function createSqliteModelCallLedger(workspaceRoot: string): ModelCallLed return new SqliteModelCallLedger(workspaceRoot); } +/** + * Runs the same bounded repair used by the ledger writer inside a caller-owned + * operational-state write transaction. This lets a cross-repository snapshot + * read the repaired projection before any other SQLite writer can intervene. + */ +export function catchUpModelCallProjectionInTransaction( + database: DatabaseSync, +): CatchUpModelCallProjectionResult { + try { + return catchUpModelCallProjection(database, {}, 16, 512); + } catch (cause) { + throw new ModelCallLedgerPublicationError(false, { cause }); + } +} + class SqliteModelCallLedger implements ModelCallLedger { readonly #lease: OperationalStateDatabaseLease; #state: 'open' | 'draining' | 'closed' = 'open'; diff --git a/packages/storage/src/usage-stores.ts b/packages/storage/src/usage-stores.ts index 19823028e3..54f8a94cf1 100644 --- a/packages/storage/src/usage-stores.ts +++ b/packages/storage/src/usage-stores.ts @@ -27,6 +27,7 @@ import type { } from '@maka/core/usage-stats/types'; import { throwDeduplicatedFailures } from './failure-utils.js'; import { + catchUpModelCallProjectionInTransaction, createSqliteModelCallLedger, type CatchUpModelCallProjectionInput, type CatchUpModelCallProjectionResult, @@ -47,6 +48,7 @@ import { type PricingSnapshot, type PricingStore, } from './pricing-store.js'; +import { acquireOperationalStateDatabase } from './operational-state-store.js'; import { runWithStorageRootLease, StorageRootAuthorityError, @@ -62,6 +64,7 @@ import { type PersistedToolInvocationRecord, type TelemetryRepo, type ToolUsageQuery, + resolveRange, } from './telemetry-repo.js'; import { createSqlitePricingStore, createSqliteTelemetryRepo } from './sqlite-usage-store.js'; @@ -123,6 +126,23 @@ export interface PricingAuthorityWriter extends PricingAuthorityReader { delete(expectedRevision: number, modelKey: string): Promise; } +export interface CaptureUsageSnapshotInput { + readonly query: UsageQuery; + readonly activityLimit: number; +} + +export interface CapturedUsageSnapshot { + readonly legacySummary: UsageSummaryV2; + readonly legacyLlmLogs: { readonly rows: readonly UsageLogRow[]; readonly total: number }; + readonly toolLogs: { + readonly rows: readonly PersistedToolInvocationRecord[]; + readonly total: number; + }; + readonly canonical: ModelCallLedgerPage; + readonly repair: CatchUpModelCallProjectionResult; + readonly pricing: PricingSnapshot; +} + export interface InteractiveUsageStoresReader { readonly kind: 'interactive'; readonly access: 'read'; @@ -140,6 +160,7 @@ export interface InteractiveUsageStoresWriter { readonly telemetry: Readonly; readonly modelCalls: Readonly; readonly pricing: Readonly; + captureUsageSnapshot(input: CaptureUsageSnapshotInput): Promise; subscribeSessionUsageChanges(listener: (sessionId: string) => void): () => void; beginDrain(): Promise; flush(): Promise; @@ -291,7 +312,13 @@ export async function openInteractiveUsageStoresForWrite( if (opening) return opening; const pending = runWithStorageRootLease(lease, 'interactive', 'write', async (root) => { const repos = await openRepos(root, true); - const stores = createWriterFacade(lease, repos.telemetry, repos.modelCalls, repos.pricing); + const stores = createWriterFacade( + root, + lease, + repos.telemetry, + repos.modelCalls, + repos.pricing, + ); writers.add(stores); writerByLease.set(lease, stores); return stores; @@ -328,6 +355,7 @@ async function openRepos( } function createWriterFacade( + root: string, lease: StorageRootLease<'interactive', 'write'>, telemetry: TelemetryRepo, modelCalls: ModelCallLedger, @@ -478,6 +506,42 @@ function createWriterFacade( isExpectedPricingFailure, ), }, + captureUsageSnapshot(input) { + if (!Number.isSafeInteger(input.activityLimit) || input.activityLimit <= 0) { + return Promise.reject(new TypeError('Usage snapshot activity limit must be positive')); + } + return admit(() => + run(() => { + const snapshotLease = acquireOperationalStateDatabase(root); + try { + const snapshot = snapshotLease.transaction('write', () => { + const repair = catchUpModelCallProjectionInTransaction(snapshotLease.database); + return snapshotLease.transaction('read', () => ({ + legacySummary: telemetry.summary(input.query), + legacyLlmLogs: telemetry.logs(input.query, 0, input.activityLimit), + toolLogs: telemetry.toolLogs( + { + range: input.query.range, + ...(input.query.status === undefined ? {} : { status: input.query.status }), + }, + 0, + input.activityLimit, + ), + canonical: modelCalls.read(resolveRange(input.query.range), input.query.sessionId), + repair, + pricing: pricing.snapshot(), + })); + }); + for (const sessionId of snapshot.repair.changedSessionIds) { + publishSessionUsageChange(sessionId); + } + return snapshot; + } finally { + snapshotLease.close(); + } + }), + ); + }, subscribeSessionUsageChanges(listener) { assertOpen(); sessionUsageChangeListeners.add(listener); From acf807777d1f3a90473e0cb8942a8e3593e83ea5 Mon Sep 17 00:00:00 2001 From: sungl <81428141+Sun-GLiang@users.noreply.github.com> Date: Mon, 31 Aug 2026 10:09:39 +0800 Subject: [PATCH 2/7] fix(runtime-host): lease usage snapshots to connections Generated-by: Codex --- .../__tests__/usage-pricing-protocol.test.ts | 142 +++++++++++++++--- .../usage-pricing-two-client-uds.test.ts | 115 +++++++++++++- .../__tests__/usage-snapshot-cache.test.ts | 142 ++++++++++++++++++ .../runtime-host/src/protocol/operations.ts | 1 + .../src/protocol/usage-pricing.ts | 32 +++- .../src/server/execution-composition.ts | 5 +- .../src/server/operation-dispatcher.ts | 5 +- .../src/server/usage-pricing-coordinator.ts | 44 +++++- .../src/server/usage-snapshot-cache.ts | 56 +++++-- 9 files changed, 495 insertions(+), 47 deletions(-) create mode 100644 packages/runtime-host/src/__tests__/usage-snapshot-cache.test.ts 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 66d6020ebe..c21a6e08dc 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, @@ -64,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({ @@ -566,7 +596,7 @@ describe('Usage/Pricing protocol', () => { } }); - test('pins every Usage authority behind one expiring LRU snapshot revision', async () => { + 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'), @@ -590,12 +620,16 @@ describe('Usage/Pricing protocol', () => { now: () => now, createRevision: () => `snapshot-${++nextRevision}`, ttlMs: 100, + hardTtlMs: 250, capacity: 2, activityLimit: 1, }, ); - const first = await expectUsageSnapshotStart(coordinator); + 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); @@ -603,9 +637,25 @@ describe('Usage/Pricing protocol', () => { await stores.telemetry.recordToolInvocation(longToolRecord('new-tool', 2)); await stores.pricing.upsert(1, pricing('snapshot:new')); - const oldLlm = await expectUsageSnapshotLogs(coordinator, first.revision, 'llm'); - const oldTool = await expectUsageSnapshotLogs(coordinator, first.revision, 'tool'); - const oldPricing = await expectUsageSnapshotPricing(coordinator, first.revision); + assert.deepEqual( + 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'], @@ -619,8 +669,13 @@ describe('Usage/Pricing protocol', () => { assert.ok(oldPricing.entries.some((entry) => entry.pricing.modelKey === 'snapshot:old')); assert.ok(!oldPricing.entries.some((entry) => entry.pricing.modelKey === 'snapshot:new')); - const second = await expectUsageSnapshotStart(coordinator); - const newLlm = await expectUsageSnapshotLogs(coordinator, second.revision, 'llm'); + 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'], @@ -628,17 +683,58 @@ describe('Usage/Pricing protocol', () => { assert.equal(newLlm.total, 1, 'total describes retained rows'); assert.equal(newLlm.truncated, true, 'truncation describes discarded authority rows'); - await expectUsageSnapshotLogs(coordinator, first.revision, 'llm'); - await expectUsageSnapshotStart(coordinator); + assert.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')).kind, - 'revision_changed', - 'the least recently used snapshot is evicted', + (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 }, ); - now += 101; - const expired = await queryUsageSnapshotLogs(coordinator, first.revision, 'llm'); - assert.deepEqual(expired, { 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(); @@ -1081,10 +1177,11 @@ async function queryUsageBuckets( async function expectUsageSnapshotStart( coordinator: HostUsagePricingCoordinator, + context: ConnectionContext = CONNECTION_CONTEXT, ): Promise> { const outcome = await coordinator.handlers['usage.query']( { kind: 'snapshot_start', range: 'all' }, - CONNECTION_CONTEXT, + context, ); assert.equal(outcome.ok, true); if (!outcome.ok || outcome.result.kind !== 'snapshot_started') { @@ -1097,10 +1194,11 @@ 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 }, - CONNECTION_CONTEXT, + context, ); assert.equal(outcome.ok, true); if ( @@ -1116,8 +1214,9 @@ async function expectUsageSnapshotLogs( coordinator: HostUsagePricingCoordinator, revision: string, source: 'llm' | 'tool', + context: ConnectionContext = CONNECTION_CONTEXT, ): Promise> { - const result = await queryUsageSnapshotLogs(coordinator, revision, source); + 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; @@ -1126,6 +1225,7 @@ async function expectUsageSnapshotLogs( async function expectUsageSnapshotPricing( coordinator: HostUsagePricingCoordinator, revision: string, + context: ConnectionContext = CONNECTION_CONTEXT, ): Promise<{ readonly entries: readonly EffectivePricingEntry[] }> { const entries: EffectivePricingEntry[] = []; let offset = 0; @@ -1133,7 +1233,7 @@ async function expectUsageSnapshotPricing( do { const outcome = await coordinator.handlers['usage.query']( { kind: 'snapshot_pricing', revision, offset, limit: PRICING_PAGE_MAX_ITEMS }, - CONNECTION_CONTEXT, + context, ); assert.equal(outcome.ok, true); if (!outcome.ok || outcome.result.kind !== 'snapshot_pricing') { @@ -1151,6 +1251,10 @@ async function expectUsageSnapshotPricing( return { entries }; } +function connectionContext(connectionId: string): ConnectionContext { + return { ...CONNECTION_CONTEXT, connectionId }; +} + 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 e45b9d718d..ab6dfa0201 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 @@ -38,7 +38,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, @@ -431,6 +435,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, @@ -514,6 +604,29 @@ describe('production Usage/Pricing UDS', () => { 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()); 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..e48ef083ec --- /dev/null +++ b/packages/runtime-host/src/__tests__/usage-snapshot-cache.test.ts @@ -0,0 +1,142 @@ +/* + * 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); +}); diff --git a/packages/runtime-host/src/protocol/operations.ts b/packages/runtime-host/src/protocol/operations.ts index d9176c2657..d1d0a624eb 100644 --- a/packages/runtime-host/src/protocol/operations.ts +++ b/packages/runtime-host/src/protocol/operations.ts @@ -326,6 +326,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 56c3c89290..1eafed8bb2 100644 --- a/packages/runtime-host/src/protocol/usage-pricing.ts +++ b/packages/runtime-host/src/protocol/usage-pricing.ts @@ -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([ @@ -354,6 +354,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, @@ -367,6 +375,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, @@ -393,6 +412,17 @@ 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') { diff --git a/packages/runtime-host/src/server/execution-composition.ts b/packages/runtime-host/src/server/execution-composition.ts index 2468419635..3efcae5698 100644 --- a/packages/runtime-host/src/server/execution-composition.ts +++ b/packages/runtime-host/src/server/execution-composition.ts @@ -1564,7 +1564,10 @@ export async function createExecutionRuntimeHostComposition( unsubscribeTaskLedger?.(); }, ], - 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/operation-dispatcher.ts b/packages/runtime-host/src/server/operation-dispatcher.ts index 8208c5ad08..c87733b88e 100644 --- a/packages/runtime-host/src/server/operation-dispatcher.ts +++ b/packages/runtime-host/src/server/operation-dispatcher.ts @@ -130,7 +130,10 @@ export type SessionCatalogOperationKey = Exclude< export type TaskLedgerOperationKey = Extract; export type ArtifactOperationKey = Extract; export type SkillCatalogOperationKey = Extract; -export type UsagePricingOperationKey = Extract; +export type UsagePricingOperationKey = Extract< + OperationKey, + 'usage.query' | 'usage.snapshot.release' | `pricing.${string}` +>; export type MemoryOperationKey = Extract; export type OAuthOperationKey = Extract; export type RuntimeResourceOperationKey = Extract; diff --git a/packages/runtime-host/src/server/usage-pricing-coordinator.ts b/packages/runtime-host/src/server/usage-pricing-coordinator.ts index c7adaf61ee..cb4aee042d 100644 --- a/packages/runtime-host/src/server/usage-pricing-coordinator.ts +++ b/packages/runtime-host/src/server/usage-pricing-coordinator.ts @@ -60,15 +60,23 @@ 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, type UsageSnapshotCacheOptions } from './usage-snapshot-cache.js'; +import { + UsageSnapshotCache, + UsageSnapshotCapacityError, + type UsageSnapshotCacheOptions, +} from './usage-snapshot-cache.js'; /** 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), }; @@ -94,6 +102,10 @@ export class HostUsagePricingCoordinator { this.#usageSnapshots = new UsageSnapshotCache(usageSnapshotOptions); } + releaseConnection(connectionId: string): void { + this.#usageSnapshots.releaseConnection(connectionId); + } + /** * Reads the canonical ledger for the window a query addresses (#1679). The * range is resolved once here so both sources answer the same window. @@ -106,14 +118,20 @@ 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(input.range, now) }; + return { + ok: true, + result: await this.#startUsageSnapshot(context.connectionId, input.range, now), + }; } if (input.kind === 'snapshot_logs') { - const snapshot = this.#usageSnapshots.get(input.revision); + 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(); @@ -130,7 +148,7 @@ export class HostUsagePricingCoordinator { }; } if (input.kind === 'snapshot_pricing') { - const snapshot = this.#usageSnapshots.get(input.revision); + 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 { @@ -227,11 +245,21 @@ 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> { @@ -256,7 +284,7 @@ export class HostUsagePricingCoordinator { 0, captureLimit, ); - const retained = this.#usageSnapshots.retain({ + const retained = this.#usageSnapshots.retain(connectionId, { summary, provenance, llmRows: mergedLogs.rows.slice(0, this.#usageSnapshots.activityLimit).map(projectUsageLog), diff --git a/packages/runtime-host/src/server/usage-snapshot-cache.ts b/packages/runtime-host/src/server/usage-snapshot-cache.ts index 07e9ab1ef9..714f4d07ca 100644 --- a/packages/runtime-host/src/server/usage-snapshot-cache.ts +++ b/packages/runtime-host/src/server/usage-snapshot-cache.ts @@ -27,13 +27,22 @@ import type { } 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; } @@ -53,15 +62,18 @@ export interface RetainedUsageSnapshot extends UsageSnapshotContents { } interface CacheEntry extends RetainedUsageSnapshot { - readonly expiresAt: number; + readonly connectionId: string; + idleExpiresAt: number; + readonly hardExpiresAt: number; } -/** Host-epoch-local, absolute-TTL cache for coherent Settings Usage reads. */ +/** 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(); @@ -69,11 +81,14 @@ export class UsageSnapshotCache { 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) || @@ -83,42 +98,51 @@ export class UsageSnapshotCache { } } - retain(contents: UsageSnapshotContents): RetainedUsageSnapshot { + retain(connectionId: string, contents: UsageSnapshotContents): RetainedUsageSnapshot { const now = this.#now(); this.#pruneExpired(now); - while (this.#entries.size >= this.#capacity) { - const oldestRevision = this.#entries.keys().next().value; - if (oldestRevision === undefined) break; - this.#entries.delete(oldestRevision); - } + 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: CacheEntry = { revision, ...contents, - expiresAt: now + this.#ttlMs, + connectionId, + idleExpiresAt: Math.min(now + this.#ttlMs, hardExpiresAt), + hardExpiresAt, }; this.#entries.set(revision, entry); return entry; } - get(revision: string): RetainedUsageSnapshot | undefined { + get(connectionId: string, revision: string): RetainedUsageSnapshot | undefined { const now = this.#now(); this.#pruneExpired(now); const entry = this.#entries.get(revision); - if (!entry) return undefined; - // Map insertion order is the LRU order. Reinsert without changing expiresAt: - // page access affects eviction priority, never the absolute lifetime. - this.#entries.delete(revision); - this.#entries.set(revision, entry); + if (!entry || entry.connectionId !== connectionId) return undefined; + entry.idleExpiresAt = Math.min(now + this.#ttlMs, entry.hardExpiresAt); return entry; } + 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.expiresAt <= now) this.#entries.delete(revision); + if (entry.idleExpiresAt <= now || entry.hardExpiresAt <= now) { + this.#entries.delete(revision); + } } } } From 046828c5cfdef020b0676eec478dd5c642747fda Mon Sep 17 00:00:00 2001 From: sungl <81428141+Sun-GLiang@users.noreply.github.com> Date: Mon, 31 Aug 2026 10:23:25 +0800 Subject: [PATCH 3/7] fix(desktop): release usage snapshot leases Generated-by: Codex --- .../runtime-host-client-usage.test.ts | 129 +++++++++++++++++- apps/desktop/src/main/runtime-host-client.ts | 52 ++++--- 2 files changed, 159 insertions(+), 22 deletions(-) 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 index bdec51dca8..a322f2cad9 100644 --- a/apps/desktop/src/main/__tests__/runtime-host-client-usage.test.ts +++ b/apps/desktop/src/main/__tests__/runtime-host-client-usage.test.ts @@ -30,6 +30,7 @@ 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'); @@ -63,11 +64,67 @@ test('loads all Usage snapshot pages behind one start revision', async () => { 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); @@ -90,6 +147,76 @@ test('discards every partial Usage result and restarts after revision_changed', 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 () => { @@ -130,7 +257,7 @@ test('rejects non-progressing or identity-changing Usage snapshot pages', async }); function usageClient( - respond: (operation: OperationKey, input: any) => Promise, + respond: (operation: OperationKey, input: any) => Promise | any, ): DesktopRuntimeHostClient { const connection = { hostEpoch: 'host-current', diff --git a/apps/desktop/src/main/runtime-host-client.ts b/apps/desktop/src/main/runtime-host-client.ts index 43b973fed2..a7ba1b71fb 100644 --- a/apps/desktop/src/main/runtime-host-client.ts +++ b/apps/desktop/src/main/runtime-host-client.ts @@ -1622,29 +1622,39 @@ export class DesktopRuntimeHostClient { async #readUsageSnapshot(range: TimeRange): Promise { this.#assertOpen(); const started = await this.request("usage.query", { kind: "snapshot_start", range }); - if ( - started.kind !== "snapshot_started" || - (typeof range === "object" && - (started.summary.range.from !== range.from || started.summary.range.to !== range.to)) - ) { + if (started.kind !== "snapshot_started") { 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, - }; + 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( From 6bb4474c574027e1bbaf6c7719122d46abffb926 Mon Sep 17 00:00:00 2001 From: sungl <81428141+Sun-GLiang@users.noreply.github.com> Date: Mon, 31 Aug 2026 11:27:00 +0800 Subject: [PATCH 4/7] fix(runtime-host): reserve usage snapshot capacity Generated-by: Codex --- .../__tests__/usage-pricing-protocol.test.ts | 121 ++++++++++++++++++ .../__tests__/usage-snapshot-cache.test.ts | 73 +++++++++++ .../src/server/usage-pricing-coordinator.ts | 85 ++++++------ .../src/server/usage-snapshot-cache.ts | 69 +++++++++- 4 files changed, 305 insertions(+), 43 deletions(-) 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 caf0ed376b..05571231e0 100644 --- a/packages/runtime-host/src/__tests__/usage-pricing-protocol.test.ts +++ b/packages/runtime-host/src/__tests__/usage-pricing-protocol.test.ts @@ -834,6 +834,105 @@ describe('Usage/Pricing protocol', () => { } }); + 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(() => @@ -1343,6 +1442,28 @@ 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-snapshot-cache.test.ts b/packages/runtime-host/src/__tests__/usage-snapshot-cache.test.ts index e48ef083ec..bec3afa6f7 100644 --- a/packages/runtime-host/src/__tests__/usage-snapshot-cache.test.ts +++ b/packages/runtime-host/src/__tests__/usage-snapshot-cache.test.ts @@ -140,3 +140,76 @@ test('expires an idle lease before its hard deadline', () => { 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/server/usage-pricing-coordinator.ts b/packages/runtime-host/src/server/usage-pricing-coordinator.ts index a0baa9d588..8b09d957c2 100644 --- a/packages/runtime-host/src/server/usage-pricing-coordinator.ts +++ b/packages/runtime-host/src/server/usage-pricing-coordinator.ts @@ -316,45 +316,52 @@ export class HostUsagePricingCoordinator { range: UsageQuery['range'], now: number, ): Promise> { - const query: UsageQuery = { range: resolveUsageRange(range, now) }; - const captureLimit = this.#usageSnapshots.activityLimit + 1; - const captured = await this.#stores.captureUsageSnapshot({ - query, - activityLimit: captureLimit, - }); - const canonical: CanonicalUsageSource = { - attempts: captured.canonical.attempts, - unreadableRecords: captured.canonical.unreadableRecords + captured.repair.unreadableEvents, - pendingRepairs: captured.repair.pendingRuns, - }; - const mergedSummary = mergeUsageSummary(captured.legacySummary, canonical, query, now); - const { provenance, ...summary } = mergedSummary; - const mergedLogs = mergeUsageLogs( - captured.legacyLlmLogs, - canonical, - query, - now, - 0, - captureLimit, - ); - const 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.retain(connectionId, { - 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), - }); - return encodeUsageQueryResult({ - kind: 'snapshot_started', - revision: retained.revision, - summary: retained.summary, - provenance: retained.provenance, - }) as Extract; + 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> { diff --git a/packages/runtime-host/src/server/usage-snapshot-cache.ts b/packages/runtime-host/src/server/usage-snapshot-cache.ts index 714f4d07ca..e2054a6b13 100644 --- a/packages/runtime-host/src/server/usage-snapshot-cache.ts +++ b/packages/runtime-host/src/server/usage-snapshot-cache.ts @@ -61,12 +61,26 @@ export interface RetainedUsageSnapshot extends UsageSnapshotContents { readonly revision: string; } -interface CacheEntry extends RetainedUsageSnapshot { +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; @@ -99,6 +113,18 @@ export class UsageSnapshotCache { } 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(); @@ -107,10 +133,12 @@ export class UsageSnapshotCache { throw new Error('Usage snapshot revision generator returned an invalid revision'); } const hardExpiresAt = now + this.#hardTtlMs; - const entry: CacheEntry = { + const entry: PendingCacheEntry = { revision, - ...contents, 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, }; @@ -118,15 +146,48 @@ export class UsageSnapshotCache { 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.connectionId !== connectionId) return undefined; + 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); From 4b717f354496856734b11212818b0e2a32944ef9 Mon Sep 17 00:00:00 2001 From: sungl <81428141+Sun-GLiang@users.noreply.github.com> Date: Mon, 31 Aug 2026 11:40:40 +0800 Subject: [PATCH 5/7] chore(test): refresh Windows skip inventory Generated-by: Codex --- docs/windows-test-inventory.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/docs/windows-test-inventory.md b/docs/windows-test-inventory.md index b530c6553d..5a0447710c 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 | 11 | | platform-contract | 31 | -Total Windows-excluded declarations: **69** +Total Windows-excluded declarations: **70** ## Inventory @@ -59,6 +59,7 @@ Total Windows-excluded declarations: **69** | 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` | From d54e03e4da0238e41ed03fe40b42c4f08502097c Mon Sep 17 00:00:00 2001 From: sungl <81428141+Sun-GLiang@users.noreply.github.com> Date: Thu, 3 Sep 2026 12:54:11 +0800 Subject: [PATCH 6/7] ci: retry transient Node test runner failure The prior run stopped when Node 24 could not deserialize its test-runner child payload. The affected release-contract file is unchanged and passes 10/10 isolated repetitions locally.\n\nGenerated-by: Codex From 3415280800066bab9dfa8f7f788e9c2bbcc81edf Mon Sep 17 00:00:00 2001 From: sungl <81428141+Sun-GLiang@users.noreply.github.com> Date: Thu, 3 Sep 2026 13:25:02 +0800 Subject: [PATCH 7/7] fix(runtime-host): retry busy usage snapshots Bound snapshot leases per connection while preserving one replacement load, retry capacity conflicts as whole Desktop loads, and enforce the shared activity ceiling at the protocol boundary.\n\nGenerated-by: Codex --- .../runtime-host-client-usage.test.ts | 55 ++++++++++++++++++- apps/desktop/src/main/runtime-host-client.ts | 24 ++++++-- .../__tests__/handshake-compatibility.test.ts | 4 +- .../__tests__/usage-pricing-protocol.test.ts | 19 +++++++ .../__tests__/usage-snapshot-cache.test.ts | 49 ++++++++++++++--- .../src/protocol/usage-pricing.ts | 7 ++- .../src/server/usage-snapshot-cache.ts | 39 ++++++------- packages/storage/src/usage-stores.ts | 4 +- 8 files changed, 162 insertions(+), 39 deletions(-) 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 index a322f2cad9..8931e07605 100644 --- a/apps/desktop/src/main/__tests__/runtime-host-client-usage.test.ts +++ b/apps/desktop/src/main/__tests__/runtime-host-client-usage.test.ts @@ -19,7 +19,10 @@ import assert from 'node:assert/strict'; import test from 'node:test'; -import type { RuntimeHostConnection } from '@maka/runtime-host/client'; +import { + type RuntimeHostConnection, + RuntimeHostOperationError, +} from '@maka/runtime-host/client'; import type { OperationInput, OperationKey } from '@maka/runtime-host/protocol'; import { DesktopRuntimeHostClient, @@ -150,6 +153,56 @@ test('discards every partial Usage result and restarts after revision_changed', assert.deepEqual(released, ['revision-1', 'revision-2']); }); +test('retries Usage snapshot capacity conflicts as whole loads', async () => { + let starts = 0; + const client = usageClient(async (operation, input) => { + if (operation === 'usage.snapshot.release') return { released: true }; + if (input.kind === 'snapshot_start') { + starts += 1; + if (starts < 3) { + throw new RuntimeHostOperationError( + 'usage.query', + 'operation_conflict', + 'Usage snapshot capacity is occupied', + ); + } + return started('revision-3', 0); + } + if (input.kind === 'snapshot_logs') { + return logPage('revision-3', input.source, [], 0, 0, null, false); + } + if (input.kind === 'snapshot_pricing') { + return pricingPage('revision-3', [], 0, 0, null); + } + throw new Error('Unexpected Usage request'); + }); + + const snapshot = await client.loadUsageSnapshot('all'); + assert.equal(snapshot.revision, 'revision-3'); + assert.equal(starts, 3); +}); + +test('bounds repeated Usage snapshot capacity conflicts', async () => { + let starts = 0; + const client = usageClient(async (operation, input) => { + assert.equal(operation, 'usage.query'); + assert.equal(input.kind, 'snapshot_start'); + starts += 1; + throw new RuntimeHostOperationError( + 'usage.query', + 'operation_conflict', + 'Usage snapshot capacity is occupied', + ); + }); + + await assert.rejects( + () => client.loadUsageSnapshot('all'), + (error: unknown) => + error instanceof DesktopRuntimeHostClientError && error.code === 'usage_unstable', + ); + assert.equal(starts, 3); +}); + 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[] = []; diff --git a/apps/desktop/src/main/runtime-host-client.ts b/apps/desktop/src/main/runtime-host-client.ts index 2473f3b64d..2e160279be 100644 --- a/apps/desktop/src/main/runtime-host-client.ts +++ b/apps/desktop/src/main/runtime-host-client.ts @@ -155,6 +155,7 @@ import { type ToolUsageLogProjection, PRICING_PAGE_MAX_ITEMS, USAGE_PAGE_MAX_ITEMS, + USAGE_SNAPSHOT_ACTIVITY_MAX_ITEMS, type WorkspaceProjection, } from "@maka/runtime-host/protocol"; @@ -164,7 +165,7 @@ 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; +const USAGE_SNAPSHOT_RETRY_DELAY_MS = 50; export type DesktopSessionConfigurationPatch = SessionConfigurationPatch; @@ -1385,12 +1386,25 @@ export class DesktopRuntimeHostClient { 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; + try { + const snapshot = await this.#readUsageSnapshot(range); + if (snapshot) return snapshot; + } catch (error) { + if ( + !(error instanceof RuntimeHostOperationError) || + error.operation !== "usage.query" || + error.code !== "operation_conflict" + ) { + throw error; + } + } + if (attempt + 1 < MAX_USAGE_SNAPSHOT_ATTEMPTS) { + await new Promise((resolve) => setTimeout(resolve, USAGE_SNAPSHOT_RETRY_DELAY_MS)); + } } throw new DesktopRuntimeHostClientError( "usage_unstable", - "Usage snapshot kept expiring while Desktop read it", + "Usage snapshot stayed unavailable across bounded retries", ); } @@ -1813,7 +1827,7 @@ export class DesktopRuntimeHostClient { page.source !== source || page.offset !== offset || page.rows.length > USAGE_PAGE_MAX_ITEMS || - page.total > MAX_USAGE_SNAPSHOT_ACTIVITY_RECORDS + page.total > USAGE_SNAPSHOT_ACTIVITY_MAX_ITEMS ) { throw invalidProjection("Usage snapshot logs"); } diff --git a/packages/runtime-host/src/__tests__/handshake-compatibility.test.ts b/packages/runtime-host/src/__tests__/handshake-compatibility.test.ts index 98c0c56280..af70582447 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-65 Host before any domain command', async () => { +test('rejects an epoch-99 Host before any domain command', async () => { let admittedRequest: RequestFrame | undefined; await withForgedHandshakePeer( async (transport, hostEpoch, rootId) => { @@ -121,7 +121,7 @@ test('rejects an epoch-65 Host before any domain command', async () => { hostEpoch, connectionId: 'forged-epoch-connection', selectedProtocol: RUNTIME_HOST_PROTOCOL_VERSION, - compatibilityEpoch: 65, + compatibilityEpoch: 99, 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 05571231e0..b94f839b9b 100644 --- a/packages/runtime-host/src/__tests__/usage-pricing-protocol.test.ts +++ b/packages/runtime-host/src/__tests__/usage-pricing-protocol.test.ts @@ -41,6 +41,7 @@ import { decodeUsageSnapshotReleaseInput, decodeUsageSnapshotReleaseResult, decodeUsageQueryInput, + decodeUsageQueryResult, encodePricingQueryResult, encodeProtocolMessage, PRICING_PAGE_MAX_BYTES, @@ -50,6 +51,7 @@ import { USAGE_PAGE_MAX_BYTES, USAGE_PAGE_MAX_ITEMS, USAGE_PROJECTION_TEXT_MAX_BYTES, + USAGE_SNAPSHOT_ACTIVITY_MAX_ITEMS, type EffectivePricingEntry, type LlmUsageLogProjection, type ToolUsageLogProjection, @@ -246,6 +248,23 @@ describe('Usage/Pricing protocol', () => { } }); + test('rejects Usage snapshot totals above the activity maximum', () => { + assert.throws( + () => + decodeUsageQueryResult({ + kind: 'snapshot_logs', + revision: 'snapshot-revision-1', + source: 'llm', + rows: [validLog()], + offset: 0, + total: USAGE_SNAPSHOT_ACTIVITY_MAX_ITEMS + 1, + nextOffset: 1, + truncated: true, + }), + 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(() => diff --git a/packages/runtime-host/src/__tests__/usage-snapshot-cache.test.ts b/packages/runtime-host/src/__tests__/usage-snapshot-cache.test.ts index bec3afa6f7..3d72c0a73a 100644 --- a/packages/runtime-host/src/__tests__/usage-snapshot-cache.test.ts +++ b/packages/runtime-host/src/__tests__/usage-snapshot-cache.test.ts @@ -24,6 +24,7 @@ import { UsageSnapshotCapacityError, type UsageSnapshotContents, } from '../server/usage-snapshot-cache.js'; +import { USAGE_SNAPSHOT_ACTIVITY_MAX_ITEMS } from '../protocol/index.js'; const CONTENTS: UsageSnapshotContents = { summary: { @@ -63,6 +64,13 @@ const CONTENTS: UsageSnapshotContents = { pricingEntries: [], }; +test('rejects an activity limit above the protocol maximum', () => { + assert.throws( + () => new UsageSnapshotCache({ activityLimit: USAGE_SNAPSHOT_ACTIVITY_MAX_ITEMS + 1 }), + TypeError, + ); +}); + test('preserves four owned leases at capacity instead of evicting an active revision', () => { let revision = 0; const cache = new UsageSnapshotCache({ @@ -70,11 +78,11 @@ test('preserves four owned leases at capacity instead of evicting an active revi createRevision: () => `revision-${++revision}`, }); const retained = Array.from({ length: 4 }, (_, index) => - cache.retain(`connection-${index}`, CONTENTS), + completeReservation(cache, `connection-${index}`), ); assert.throws( - () => cache.retain('connection-5', CONTENTS), + () => completeReservation(cache, 'connection-5'), (error: unknown) => error instanceof UsageSnapshotCapacityError && error.message === 'Usage snapshot capacity is occupied', @@ -84,26 +92,42 @@ test('preserves four owned leases at capacity instead of evicting an active revi } }); +test('limits one connection to two leases without consuming global capacity', () => { + let revision = 0; + const cache = new UsageSnapshotCache({ + capacity: 4, + createRevision: () => `revision-${++revision}`, + }); + + cache.reserve('connection-a'); + cache.reserve('connection-a'); + assert.throws(() => cache.reserve('connection-a'), UsageSnapshotCapacityError); + + assert.doesNotThrow(() => cache.reserve('connection-b')); + assert.doesNotThrow(() => cache.reserve('connection-b')); + assert.throws(() => cache.reserve('connection-c'), UsageSnapshotCapacityError); +}); + 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); + const first = completeReservation(cache, 'connection-a'); + completeReservation(cache, 'connection-b'); 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); + assert.throws(() => completeReservation(cache, 'connection-c'), UsageSnapshotCapacityError); cache.release('connection-a', first.revision); - const third = cache.retain('connection-c', CONTENTS); + const third = completeReservation(cache, 'connection-c'); assert.equal(cache.get('connection-c', third.revision)?.revision, third.revision); cache.releaseConnection('connection-b'); - const fourth = cache.retain('connection-d', CONTENTS); + const fourth = completeReservation(cache, 'connection-d'); assert.equal(cache.get('connection-d', fourth.revision)?.revision, fourth.revision); }); @@ -115,7 +139,7 @@ test('renews idle lifetime on owner access without extending the hard deadline', hardTtlMs: 250, createRevision: () => 'revision-1', }); - const retained = cache.retain('connection-a', CONTENTS); + const retained = completeReservation(cache, 'connection-a'); now = 90; assert.equal(cache.get('connection-a', retained.revision)?.revision, retained.revision); @@ -135,7 +159,7 @@ test('expires an idle lease before its hard deadline', () => { hardTtlMs: 250, createRevision: () => 'revision-1', }); - const retained = cache.retain('connection-a', CONTENTS); + const retained = completeReservation(cache, 'connection-a'); now = 100; assert.equal(cache.get('connection-a', retained.revision), undefined); @@ -213,3 +237,10 @@ test('finalization never revives a wrong, released, or expired reservation', () assert.equal(cache.finalize('connection-a', expiring.revision, CONTENTS), undefined); assert.equal(cache.get('connection-a', expiring.revision), undefined); }); + +function completeReservation(cache: UsageSnapshotCache, connectionId: string) { + const reservation = cache.reserve(connectionId); + const retained = cache.finalize(connectionId, reservation.revision, CONTENTS); + assert.ok(retained); + return retained; +} diff --git a/packages/runtime-host/src/protocol/usage-pricing.ts b/packages/runtime-host/src/protocol/usage-pricing.ts index 7016da2599..d5794c850e 100644 --- a/packages/runtime-host/src/protocol/usage-pricing.ts +++ b/packages/runtime-host/src/protocol/usage-pricing.ts @@ -40,6 +40,7 @@ import { defineOperation } from './operation-spec.js'; export const USAGE_PAGE_MAX_ITEMS = 100; export const USAGE_PAGE_MAX_BYTES = 48 * 1024; +export const USAGE_SNAPSHOT_ACTIVITY_MAX_ITEMS = 50_000; export const USAGE_PROJECTION_TEXT_MAX_BYTES = 1024; export const PRICING_PAGE_MAX_ITEMS = 128; export const PRICING_PAGE_MAX_BYTES = 48 * 1024; @@ -1022,12 +1023,16 @@ function decodeUsageSnapshotLogPage( throw invalidProtocolFrame('Invalid usage snapshot truncation flag'); } const rows = rawItems.map(decodeItem); + const page = decodeUsagePagePosition(result, rows.length); + if (page.total > USAGE_SNAPSHOT_ACTIVITY_MAX_ITEMS) { + throw invalidProtocolFrame('Usage snapshot exceeds activity limit'); + } const decoded = { kind: 'snapshot_logs', revision: requireId(result.revision, 'usage snapshot revision'), source, rows, - ...decodeUsagePagePosition(result, rows.length), + ...page, truncated: result.truncated, } as Extract; assertJsonBytes(decoded, USAGE_PAGE_MAX_BYTES, 'Usage snapshot page'); diff --git a/packages/runtime-host/src/server/usage-snapshot-cache.ts b/packages/runtime-host/src/server/usage-snapshot-cache.ts index e2054a6b13..a591a713fb 100644 --- a/packages/runtime-host/src/server/usage-snapshot-cache.ts +++ b/packages/runtime-host/src/server/usage-snapshot-cache.ts @@ -20,16 +20,19 @@ 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, +import { + USAGE_SNAPSHOT_ACTIVITY_MAX_ITEMS, + type EffectivePricingEntry, + type LlmUsageLogProjection, + type 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; +// Admit one current and one replacement load without letting one connection +// occupy every globally available lease. +const USAGE_SNAPSHOT_CONNECTION_CAPACITY = 2; export class UsageSnapshotCapacityError extends Error { constructor() { @@ -89,6 +92,7 @@ export class UsageSnapshotCache { readonly #ttlMs: number; readonly #hardTtlMs: number; readonly #capacity: number; + readonly #connectionCapacity: number; readonly #entries = new Map(); constructor(options: UsageSnapshotCacheOptions = {}) { @@ -97,7 +101,8 @@ export class UsageSnapshotCache { 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; + this.#connectionCapacity = Math.min(USAGE_SNAPSHOT_CONNECTION_CAPACITY, this.#capacity); + this.activityLimit = options.activityLimit ?? USAGE_SNAPSHOT_ACTIVITY_MAX_ITEMS; if ( !Number.isSafeInteger(this.#ttlMs) || this.#ttlMs <= 0 || @@ -106,27 +111,23 @@ export class UsageSnapshotCache { !Number.isSafeInteger(this.#capacity) || this.#capacity <= 0 || !Number.isSafeInteger(this.activityLimit) || - this.activityLimit <= 0 + this.activityLimit <= 0 || + this.activityLimit > USAGE_SNAPSHOT_ACTIVITY_MAX_ITEMS ) { 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); + let connectionEntries = 0; + for (const entry of this.#entries.values()) { + if (entry.connectionId === connectionId) connectionEntries += 1; + } + if (connectionEntries >= this.#connectionCapacity) { + throw new UsageSnapshotCapacityError(); + } if (this.#entries.size >= this.#capacity) throw new UsageSnapshotCapacityError(); const revision = this.#createRevision(); if (revision.length === 0 || revision.length > 128 || this.#entries.has(revision)) { diff --git a/packages/storage/src/usage-stores.ts b/packages/storage/src/usage-stores.ts index 54f8a94cf1..bf21e99695 100644 --- a/packages/storage/src/usage-stores.ts +++ b/packages/storage/src/usage-stores.ts @@ -516,7 +516,7 @@ function createWriterFacade( try { const snapshot = snapshotLease.transaction('write', () => { const repair = catchUpModelCallProjectionInTransaction(snapshotLease.database); - return snapshotLease.transaction('read', () => ({ + return { legacySummary: telemetry.summary(input.query), legacyLlmLogs: telemetry.logs(input.query, 0, input.activityLimit), toolLogs: telemetry.toolLogs( @@ -530,7 +530,7 @@ function createWriterFacade( canonical: modelCalls.read(resolveRange(input.query.range), input.query.sessionId), repair, pricing: pricing.snapshot(), - })); + }; }); for (const sessionId of snapshot.repair.changedSessionIds) { publishSessionUsageChange(sessionId);