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/5] 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/5] 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/5] 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/5] 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/5] 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` |