mirror of
https://github.com/stablyai/orca.git
synced 2026-11-14 01:47:24 +00:00
Run local usage services in a supervised backend process (#40792)
* Move local usage services into a supervised backend process * Own usage recovery and worker retirement lifetimes * Preserve answered usage data during background refresh * Run independent usage scans concurrently in the backend * Limit backend startup paths to the delegated usage service * Show usage failures through backend recovery * Include usage failure strings in the English catalog * Keep usage update details beside refresh failures * Resolve usage freshness failure from the final producer snapshot * Retain the forced refresh completion requirement * Keep managed backend failures out of application crash reports * Preserve usage recovery boundaries across backend reconnects * Match the backend RPC cancellation signal in the usage seam proof * Avoid serializing backend requests before port delivery * Keep receiving request validation in the boundary schemas * test: align usage transport coverage with renderer project * fix: preserve unknown usage state and renderer retirement * fix(usage): revalidate tracking before scan admission * fix(usage): retry failed backend startup on demand * test(usage): isolate stale backend ownership checks
This commit is contained in:
@@ -0,0 +1,319 @@
|
||||
import { mkdtempSync, readFileSync, rmSync } from 'node:fs'
|
||||
|
||||
import { tmpdir } from 'node:os'
|
||||
import { dirname, join } from 'node:path'
|
||||
import type { UsageStoreHost, UsageTelemetryHost } from '../usage-store-host'
|
||||
import { createTestUsageStoreHost } from '../usage-store-host-test-fixture'
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
|
||||
import type { UsageScanWorktreeRef } from './usage-provider-contract'
|
||||
import { UsageProviderStoreLifecycle } from './usage-provider-store-lifecycle'
|
||||
import {
|
||||
splitUsageCacheFile,
|
||||
type UsageCacheSplitRequest,
|
||||
type UsageCacheSplitResult,
|
||||
type UsageSourceCacheRef
|
||||
} from '../cache/usage-source-cache-file'
|
||||
|
||||
type TestSource = { id: string }
|
||||
type TestSession = { id: string }
|
||||
type TestDailyAggregate = { day: string }
|
||||
type TestScanState = {
|
||||
enabled: boolean
|
||||
lastScanStartedAt: number | null
|
||||
lastScanCompletedAt: number | null
|
||||
lastScanError: string | null
|
||||
}
|
||||
type TestState = {
|
||||
schemaVersion: number
|
||||
worktreeFingerprint: string | null
|
||||
processedSources: TestSource[]
|
||||
sessions: TestSession[]
|
||||
dailyAggregates: TestDailyAggregate[]
|
||||
scanState: TestScanState
|
||||
}
|
||||
type TestScanResult = Pick<TestState, 'sessions' | 'dailyAggregates'>
|
||||
type TestScan = (
|
||||
worktrees: UsageScanWorktreeRef[],
|
||||
sourceCache: UsageSourceCacheRef
|
||||
) => Promise<TestScanResult>
|
||||
type TestSplit = (request: UsageCacheSplitRequest) => Promise<UsageCacheSplitResult>
|
||||
|
||||
const NOW = Date.parse('2026-04-10T16:00:00.000Z')
|
||||
|
||||
function makeState(
|
||||
overrides: Partial<Omit<TestState, 'scanState'>> & { scanState?: Partial<TestScanState> } = {}
|
||||
): TestState {
|
||||
const { scanState, ...stateOverrides } = overrides
|
||||
return {
|
||||
schemaVersion: 1,
|
||||
worktreeFingerprint: null,
|
||||
processedSources: [],
|
||||
sessions: [],
|
||||
dailyAggregates: [],
|
||||
...stateOverrides,
|
||||
scanState: {
|
||||
enabled: false,
|
||||
lastScanStartedAt: null,
|
||||
lastScanCompletedAt: null,
|
||||
lastScanError: null,
|
||||
...scanState
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
function emptyScanResult(): TestScanResult {
|
||||
return { sessions: [], dailyAggregates: [] }
|
||||
}
|
||||
|
||||
class TestUsageStore extends UsageProviderStoreLifecycle<
|
||||
'processedSources',
|
||||
TestState,
|
||||
'hasAnyTestData'
|
||||
> {
|
||||
constructor(
|
||||
cacheFile: string,
|
||||
scan: TestScan,
|
||||
splitCacheFile: TestSplit,
|
||||
capture: ReturnType<typeof vi.fn<UsageTelemetryHost['capture']>>,
|
||||
readCatalog?: UsageStoreHost['readCatalog'],
|
||||
isEnabled?: UsageTelemetryHost['isEnabled'],
|
||||
signal?: AbortSignal
|
||||
) {
|
||||
super(
|
||||
{
|
||||
...createTestUsageStoreHost({
|
||||
userDataPath: dirname(cacheFile),
|
||||
telemetryEnabled: true,
|
||||
capture,
|
||||
signal
|
||||
}),
|
||||
...(readCatalog ? { readCatalog } : {}),
|
||||
...(isEnabled ? { telemetry: { isEnabled, capture } } : {})
|
||||
},
|
||||
{
|
||||
tokenUsage: {
|
||||
provider: 'claude',
|
||||
selectSessions: (state) =>
|
||||
state.sessions.map((session) => ({
|
||||
providerSessionId: session.id,
|
||||
input_tokens: 10,
|
||||
output_tokens: 2,
|
||||
cached_input_tokens: 3,
|
||||
cache_write_input_tokens: 1
|
||||
}))
|
||||
},
|
||||
logTag: '[test-usage]',
|
||||
resolveCacheFile: () => cacheFile,
|
||||
createDefaultState: makeState,
|
||||
normalizeState: (state) => state,
|
||||
sourceKey: 'processedSources',
|
||||
dataPresenceKey: 'hasAnyTestData',
|
||||
scan,
|
||||
splitCacheFile
|
||||
}
|
||||
)
|
||||
}
|
||||
|
||||
replaceState(state: TestState): void {
|
||||
this.state = state
|
||||
}
|
||||
|
||||
getState(): TestState {
|
||||
return this.state
|
||||
}
|
||||
}
|
||||
|
||||
describe('UsageProviderStoreLifecycle', () => {
|
||||
let tempDirectory: string
|
||||
let stores: TestUsageStore[]
|
||||
let scan: ReturnType<typeof vi.fn<TestScan>>
|
||||
let split: ReturnType<typeof vi.fn<TestSplit>>
|
||||
let capture: ReturnType<typeof vi.fn<UsageTelemetryHost['capture']>>
|
||||
|
||||
beforeEach(() => {
|
||||
capture = vi.fn<UsageTelemetryHost['capture']>(async () => true)
|
||||
tempDirectory = mkdtempSync(join(tmpdir(), 'orca-usage-lifecycle-'))
|
||||
stores = []
|
||||
scan = vi.fn<TestScan>().mockResolvedValue(emptyScanResult())
|
||||
// Run the real cache split in-process; the worker only wraps it in a thread.
|
||||
split = vi.fn<TestSplit>(splitUsageCacheFile)
|
||||
vi.spyOn(Date, 'now').mockReturnValue(NOW)
|
||||
})
|
||||
|
||||
afterEach(async () => {
|
||||
await Promise.all(stores.map((store) => store.flush()))
|
||||
rmSync(tempDirectory, { recursive: true, force: true })
|
||||
vi.restoreAllMocks()
|
||||
})
|
||||
|
||||
it.each([100, 1000])('uses one consistent catalog per refresh at %i worktrees', async (count) => {
|
||||
const worktrees = Array.from({ length: count }, (_, i) => ({
|
||||
worktreeId: `repo::/work/${i}`,
|
||||
path: `/work/${i}`,
|
||||
displayName: `Work ${i}`
|
||||
}))
|
||||
const readCatalog = vi.fn<UsageStoreHost['readCatalog']>().mockResolvedValue({
|
||||
repoIds: ['repo'],
|
||||
worktreesByRepo: [['repo', worktrees]]
|
||||
})
|
||||
const store = new TestUsageStore(
|
||||
join(tempDirectory, 'catalog.json'),
|
||||
scan,
|
||||
split,
|
||||
capture,
|
||||
readCatalog
|
||||
)
|
||||
stores.push(store)
|
||||
await store.whenLoaded()
|
||||
await store.refresh(true)
|
||||
expect(readCatalog).not.toHaveBeenCalled()
|
||||
await store.setEnabled(true)
|
||||
const blocked = Promise.withResolvers<TestScanResult>()
|
||||
scan.mockReturnValueOnce(blocked.promise)
|
||||
const first = store.refresh(true)
|
||||
const second = store.refresh(true)
|
||||
await vi.waitFor(() => expect(scan).toHaveBeenCalledTimes(1))
|
||||
expect(readCatalog).toHaveBeenCalledTimes(2)
|
||||
expect(scan.mock.calls[0][0]).toHaveLength(count)
|
||||
expect(scan.mock.calls[0][0][0]).toMatchObject({ path: '/work/0' })
|
||||
blocked.resolve({ sessions: [{ id: 'catalog-result' }], dailyAggregates: [] })
|
||||
await Promise.all([first, second])
|
||||
expect(store.getState().sessions).toEqual([{ id: 'catalog-result' }])
|
||||
readCatalog.mockClear()
|
||||
await store.refresh()
|
||||
expect(readCatalog).toHaveBeenCalledTimes(1)
|
||||
expect(scan).toHaveBeenCalledTimes(1)
|
||||
store.stopScanAdmission()
|
||||
readCatalog.mockClear()
|
||||
await store.refresh(true)
|
||||
expect(readCatalog).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('keeps a successful scan when the first async telemetry consent request rejects', async () => {
|
||||
const cacheFile = join(tempDirectory, 'consent-rejection.json')
|
||||
const consent = vi
|
||||
.fn<UsageTelemetryHost['isEnabled']>()
|
||||
.mockRejectedValue(new Error('No RPC credit'))
|
||||
const store = new TestUsageStore(cacheFile, scan, split, capture, undefined, consent)
|
||||
stores.push(store)
|
||||
await store.setEnabled(true)
|
||||
scan.mockResolvedValueOnce({
|
||||
sessions: [{ id: 'fresh' }],
|
||||
dailyAggregates: [{ day: '2026-04-10' }]
|
||||
})
|
||||
vi.spyOn(console, 'warn').mockImplementation(() => {})
|
||||
|
||||
await expect(store.refresh(true)).resolves.toMatchObject({
|
||||
isScanning: false,
|
||||
lastScanCompletedAt: NOW,
|
||||
lastScanError: null,
|
||||
hasAnyTestData: true
|
||||
})
|
||||
await store.flush()
|
||||
expect(store.getState().sessions).toEqual([{ id: 'fresh' }])
|
||||
expect(JSON.parse(readFileSync(cacheFile, 'utf8'))).toMatchObject({
|
||||
sessions: [{ id: 'fresh' }],
|
||||
dailyAggregates: [{ day: '2026-04-10' }],
|
||||
scanState: { lastScanCompletedAt: NOW, lastScanError: null }
|
||||
})
|
||||
expect(consent).toHaveBeenCalledTimes(1)
|
||||
expect(capture).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('stops admission without persisting a host catalog abort as a scan failure', async () => {
|
||||
const cacheFile = join(tempDirectory, 'aborted-catalog.json')
|
||||
const abort = new AbortController()
|
||||
const readCatalog = vi.fn<UsageStoreHost['readCatalog']>().mockResolvedValue({
|
||||
repoIds: ['repo-1'],
|
||||
worktreesByRepo: [
|
||||
['repo-1', [{ worktreeId: 'repo-1::/workspace/a', path: '/workspace/a', displayName: 'A' }]]
|
||||
]
|
||||
})
|
||||
const store = new TestUsageStore(
|
||||
cacheFile,
|
||||
scan,
|
||||
split,
|
||||
capture,
|
||||
readCatalog,
|
||||
undefined,
|
||||
abort.signal
|
||||
)
|
||||
stores.push(store)
|
||||
await store.setEnabled(true)
|
||||
scan.mockResolvedValueOnce({ sessions: [{ id: 'scan-result' }], dailyAggregates: [] })
|
||||
await store.refresh(true)
|
||||
expect(scan.mock.calls[0]?.[0]).toEqual([
|
||||
{
|
||||
repoId: 'repo-1',
|
||||
worktreeId: 'repo-1::/workspace/a',
|
||||
path: '/workspace/a',
|
||||
displayName: 'A'
|
||||
}
|
||||
])
|
||||
const catalog = Promise.withResolvers<Awaited<ReturnType<UsageStoreHost['readCatalog']>>>()
|
||||
readCatalog.mockReturnValueOnce(catalog.promise)
|
||||
const refresh = store.refresh(true)
|
||||
abort.abort(new Error('Backend stopping'))
|
||||
catalog.resolve({ repoIds: [], worktreesByRepo: [] })
|
||||
await expect(refresh).rejects.toThrow('Backend stopping')
|
||||
store.stopScanAdmission()
|
||||
await store.refresh(true)
|
||||
await store.flush()
|
||||
expect(JSON.parse(readFileSync(cacheFile, 'utf8'))).toMatchObject({
|
||||
sessions: [{ id: 'scan-result' }],
|
||||
scanState: { enabled: true, lastScanError: null }
|
||||
})
|
||||
})
|
||||
|
||||
it.each([
|
||||
{ waitingOn: 'catalog', sessions: [{ id: 'prior' }], completedAt: 12_000 },
|
||||
{ waitingOn: 'scanner', sessions: [{ id: 'fresh' }], completedAt: NOW }
|
||||
])(
|
||||
'preserves disable while waiting on $waitingOn',
|
||||
async ({ waitingOn, sessions, completedAt }) => {
|
||||
const cacheFile = join(tempDirectory, `disable-${waitingOn}.json`)
|
||||
const entered = Promise.withResolvers<void>()
|
||||
const release = Promise.withResolvers<void>()
|
||||
const readCatalog = async () => {
|
||||
if (waitingOn === 'catalog') {
|
||||
entered.resolve()
|
||||
await release.promise
|
||||
}
|
||||
return { repoIds: [], worktreesByRepo: [] }
|
||||
}
|
||||
scan.mockImplementationOnce(async () => {
|
||||
if (waitingOn === 'scanner') {
|
||||
entered.resolve()
|
||||
await release.promise
|
||||
}
|
||||
return { sessions: [{ id: 'fresh' }], dailyAggregates: [] }
|
||||
})
|
||||
const store = new TestUsageStore(cacheFile, scan, split, capture, readCatalog)
|
||||
stores.push(store)
|
||||
await store.whenLoaded()
|
||||
store.replaceState(
|
||||
makeState({
|
||||
sessions: [{ id: 'prior' }],
|
||||
scanState: { enabled: true, lastScanCompletedAt: 12_000 }
|
||||
})
|
||||
)
|
||||
const refresh = store.refresh(true)
|
||||
await entered.promise
|
||||
await store.setEnabled(false)
|
||||
release.resolve()
|
||||
await expect(refresh).resolves.toMatchObject({
|
||||
enabled: false,
|
||||
isScanning: false,
|
||||
lastScanCompletedAt: completedAt,
|
||||
lastScanError: null
|
||||
})
|
||||
expect(store.getState().sessions).toEqual(sessions)
|
||||
await store.flush()
|
||||
expect(JSON.parse(readFileSync(cacheFile, 'utf8'))).toMatchObject({
|
||||
sessions,
|
||||
scanState: { enabled: false, lastScanCompletedAt: completedAt, lastScanError: null }
|
||||
})
|
||||
expect(capture).not.toHaveBeenCalled()
|
||||
}
|
||||
)
|
||||
})
|
||||
@@ -1,19 +1,13 @@
|
||||
import { readFileSync, statSync } from 'node:fs'
|
||||
import { readFile } from 'node:fs/promises'
|
||||
import { AgentTokenUsageReporter } from '../reporting/token-usage/agent-token-usage-reporter'
|
||||
|
||||
import { isTelemetryEnabled } from '../../telemetry/client'
|
||||
import { join, parse } from 'node:path'
|
||||
import {
|
||||
AnalyticsSessionIdStore,
|
||||
type AnalyticsSessionId
|
||||
} from '../reporting/token-usage/analytics-session-id-store'
|
||||
import type { UsageStoreHost } from '../usage-store-host'
|
||||
import { UsageCacheSnapshotWriter } from '../cache/usage-cache-snapshot-writer'
|
||||
import {
|
||||
loadKnownUsageWorktreesByRepo,
|
||||
type UsageWorktreeSource
|
||||
} from '../worktree-attribution/usage-worktree-metadata'
|
||||
|
||||
import {
|
||||
createWorktreeRefs,
|
||||
getUsageWorktreeFingerprint
|
||||
@@ -40,20 +34,21 @@ export abstract class UsageProviderStoreLifecycle<
|
||||
private readonly loaded: Promise<void>
|
||||
private readonly schemaVersion: number
|
||||
private scanPromise: Promise<void> | null = null
|
||||
private stopping = false
|
||||
private tokenReporter: AgentTokenUsageReporter | null = null
|
||||
private analyticsSessionIds: AnalyticsSessionIdStore | null = null
|
||||
private readonly writer: UsageCacheSnapshotWriter
|
||||
private readonly inlineCache = new UsageInlineCachePreservation()
|
||||
|
||||
constructor(
|
||||
private readonly store: UsageWorktreeSource,
|
||||
private readonly host: UsageStoreHost,
|
||||
private readonly config: UsageProviderStoreLifecycleConfig<SourceKey, State, DataPresenceKey>
|
||||
) {
|
||||
this.writer = new UsageCacheSnapshotWriter(config.logTag, config.resolveCacheFile)
|
||||
const defaults = config.createDefaultState()
|
||||
this.schemaVersion = defaults.schemaVersion
|
||||
this.state = defaults
|
||||
this.loaded = Promise.all([this.load(), store.whenLoaded()]).then(() => {})
|
||||
this.loaded = this.load()
|
||||
}
|
||||
|
||||
/** Resolves once persisted state is in memory; synchronous readers must wait for it first. */
|
||||
@@ -82,6 +77,14 @@ export abstract class UsageProviderStoreLifecycle<
|
||||
}
|
||||
|
||||
/** Await queued cache writes so quit does not drop the final snapshot. */
|
||||
stopScanAdmission(): void {
|
||||
this.stopping = true
|
||||
}
|
||||
|
||||
async settleScansForShutdown(): Promise<void> {
|
||||
await this.scanPromise
|
||||
}
|
||||
|
||||
async flush(): Promise<void> {
|
||||
await this.loaded
|
||||
await this.inlineCache.settle()
|
||||
@@ -98,17 +101,22 @@ export abstract class UsageProviderStoreLifecycle<
|
||||
|
||||
async refresh(force = false): Promise<PublicUsageProviderScanState<DataPresenceKey>> {
|
||||
await this.loaded
|
||||
if (!this.state.scanState.enabled) {
|
||||
if (this.stopping || !this.state.scanState.enabled) {
|
||||
return this.getScanState()
|
||||
}
|
||||
const currentWorktreeFingerprint = await this.getCurrentWorktreeFingerprint()
|
||||
const catalog = await this.host.readCatalog()
|
||||
this.host.signal.throwIfAborted()
|
||||
const currentWorktreeFingerprint = getUsageWorktreeFingerprint(new Map(catalog.worktreesByRepo))
|
||||
if (!force && this.state.scanState.lastScanCompletedAt) {
|
||||
const ageMs = Date.now() - this.state.scanState.lastScanCompletedAt
|
||||
if (ageMs < STALE_MS && this.state.worktreeFingerprint === currentWorktreeFingerprint) {
|
||||
return this.getScanState()
|
||||
}
|
||||
}
|
||||
await this.runScan()
|
||||
if (this.stopping || !this.state.scanState.enabled) {
|
||||
return this.getScanState()
|
||||
}
|
||||
await this.runScan(catalog)
|
||||
return this.getScanState()
|
||||
}
|
||||
|
||||
@@ -219,7 +227,9 @@ export abstract class UsageProviderStoreLifecycle<
|
||||
})
|
||||
}
|
||||
|
||||
private async runScan(): Promise<void> {
|
||||
private async runScan(
|
||||
catalog: Awaited<ReturnType<UsageStoreHost['readCatalog']>>
|
||||
): Promise<void> {
|
||||
if (this.scanPromise) {
|
||||
await this.scanPromise
|
||||
return
|
||||
@@ -231,15 +241,22 @@ export abstract class UsageProviderStoreLifecycle<
|
||||
// Assign before yielding so concurrent refreshes share one scan.
|
||||
this.scanPromise = (async () => {
|
||||
try {
|
||||
const repos = this.store.getRepos()
|
||||
const worktreesByRepo = loadKnownUsageWorktreesByRepo(this.store, repos)
|
||||
const { repoIds, worktreesByRepo: entries } = catalog
|
||||
this.host.signal.throwIfAborted()
|
||||
const worktreesByRepo = new Map(entries)
|
||||
const worktreeFingerprint = getUsageWorktreeFingerprint(worktreesByRepo)
|
||||
const result = await this.config.scan(createWorktreeRefs(repos, worktreesByRepo), {
|
||||
path: usageSourceCachePath(this.config.resolveCacheFile()),
|
||||
schemaVersion: this.schemaVersion,
|
||||
worktreeFingerprint,
|
||||
reuse: this.state.worktreeFingerprint === worktreeFingerprint
|
||||
})
|
||||
const result = await this.config.scan(
|
||||
createWorktreeRefs(
|
||||
repoIds.map((id) => ({ id })),
|
||||
worktreesByRepo
|
||||
),
|
||||
{
|
||||
path: usageSourceCachePath(this.config.resolveCacheFile()),
|
||||
schemaVersion: this.schemaVersion,
|
||||
worktreeFingerprint,
|
||||
reuse: this.state.worktreeFingerprint === worktreeFingerprint
|
||||
}
|
||||
)
|
||||
this.state.sessions = result.sessions
|
||||
this.state.dailyAggregates = result.dailyAggregates
|
||||
this.state.worktreeFingerprint = worktreeFingerprint
|
||||
@@ -251,6 +268,9 @@ export abstract class UsageProviderStoreLifecycle<
|
||||
await this.writeToDisk().catch(() => {})
|
||||
await this.reportTokenUsage()
|
||||
} catch (error) {
|
||||
if (this.stopping || this.host.signal.aborted) {
|
||||
return
|
||||
}
|
||||
this.state.scanState.lastScanError = error instanceof Error ? error.message : String(error)
|
||||
await this.writeToDisk().catch(() => {})
|
||||
} finally {
|
||||
@@ -263,16 +283,20 @@ export abstract class UsageProviderStoreLifecycle<
|
||||
|
||||
private async reportTokenUsage(): Promise<void> {
|
||||
const config = this.config.tokenUsage
|
||||
if (!config || !isTelemetryEnabled() || !this.state.scanState.enabled) {
|
||||
if (!config) {
|
||||
return
|
||||
}
|
||||
try {
|
||||
if (!(await this.host.telemetry.isEnabled()) || !this.state.scanState.enabled) {
|
||||
return
|
||||
}
|
||||
if (!this.tokenReporter) {
|
||||
const { dir, name } = parse(this.config.resolveCacheFile())
|
||||
this.tokenReporter = new AgentTokenUsageReporter(
|
||||
join(dir, `${name}-token-usage.json`),
|
||||
config.provider,
|
||||
(ids) => this.getAnalyticsSessionIds(ids)
|
||||
(ids) => this.getAnalyticsSessionIds(ids),
|
||||
this.host.telemetry
|
||||
)
|
||||
}
|
||||
await this.tokenReporter.report(config.selectSessions(this.state))
|
||||
@@ -283,11 +307,6 @@ export abstract class UsageProviderStoreLifecycle<
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
private async getCurrentWorktreeFingerprint(): Promise<string> {
|
||||
const repos = this.store.getRepos()
|
||||
return getUsageWorktreeFingerprint(loadKnownUsageWorktreesByRepo(this.store, repos))
|
||||
}
|
||||
}
|
||||
|
||||
function isMissingFileError(error: unknown): boolean {
|
||||
|
||||
@@ -5,7 +5,7 @@ import {
|
||||
agentTokenUsageSchema,
|
||||
type AgentTokenUsage
|
||||
} from '../../../../shared/telemetry-agent-token-usage-schema'
|
||||
import { isTelemetryEnabled, track } from '../../../telemetry/client'
|
||||
import type { UsageTelemetryHost } from '../../usage-store-host'
|
||||
import { UsageCacheSnapshotWriter } from '../../cache/usage-cache-snapshot-writer'
|
||||
import type { AgentTokenSession } from './agent-token-usage'
|
||||
|
||||
@@ -26,7 +26,8 @@ export class AgentTokenUsageReporter {
|
||||
constructor(
|
||||
private readonly file: string,
|
||||
private readonly provider: AgentTokenUsage['provider'],
|
||||
private readonly getSessionIds: (providerSessionIds: readonly string[]) => Promise<string[]>
|
||||
private readonly getSessionIds: (providerSessionIds: readonly string[]) => Promise<string[]>,
|
||||
private readonly telemetry: UsageTelemetryHost
|
||||
) {
|
||||
this.writer = new UsageCacheSnapshotWriter('[agent-token-usage]', () => file)
|
||||
}
|
||||
@@ -43,7 +44,7 @@ export class AgentTokenUsageReporter {
|
||||
}
|
||||
|
||||
private async publish(sessions: AgentTokenSession[]): Promise<void> {
|
||||
if (!isTelemetryEnabled()) {
|
||||
if (!(await this.telemetry.isEnabled())) {
|
||||
return
|
||||
}
|
||||
const rows = sessions.flatMap(({ providerSessionId, ...counts }) => {
|
||||
@@ -51,7 +52,7 @@ export class AgentTokenUsageReporter {
|
||||
return parsed.success ? [{ providerSessionId, counts: parsed.data }] : []
|
||||
})
|
||||
const ids = await this.getSessionIds(rows.map((row) => row.providerSessionId))
|
||||
if (!isTelemetryEnabled()) {
|
||||
if (!(await this.telemetry.isEnabled())) {
|
||||
return
|
||||
}
|
||||
const previous = this.snapshots ?? (await this.load())
|
||||
@@ -91,7 +92,7 @@ export class AgentTokenUsageReporter {
|
||||
if (this.acceptedRevisions.get(snapshot.analytics_session_id) === snapshot.revision) {
|
||||
continue
|
||||
}
|
||||
if (!track('agent_token_usage', snapshot)) {
|
||||
if (!(await this.telemetry.capture(snapshot))) {
|
||||
break
|
||||
}
|
||||
// Queue acceptance only deduplicates this process; restart must retry undelivered snapshots.
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
import { WorkerThreadRequestQueue } from '../../worker-threads/worker-thread-request-queue'
|
||||
import { UsageScanRequestScheduler } from '../usage-scan-request-scheduler'
|
||||
import type { WorkerThreadFactory } from '../../worker-threads/worker-thread-transport'
|
||||
import type { ClaudeUsageDailyAggregate, ClaudeUsageSession } from '../../claude-usage/types'
|
||||
import type { CodexUsageDailyAggregate, CodexUsageSession } from '../../codex-usage/types'
|
||||
@@ -7,9 +7,7 @@ import type { MuseUsageDailyAggregate, MuseUsageSession } from '../../muse-usage
|
||||
import type { UsageScanWorktreeRef } from '../provider-lifecycle/usage-provider-contract'
|
||||
import type {
|
||||
UsageScanWorkerProviderId,
|
||||
UsageScanWorkerRequest,
|
||||
UsageScanWorkerRequestBody,
|
||||
UsageScanWorkerResponse,
|
||||
UsageScanWorkerScanBody,
|
||||
UsageScanWorkerValue
|
||||
} from './usage-scan-worker-protocol'
|
||||
@@ -18,10 +16,9 @@ import type {
|
||||
UsageCacheSplitResult,
|
||||
UsageSourceCacheRef
|
||||
} from '../cache/usage-source-cache-file'
|
||||
import { isUsageScanWorkerProgress } from './usage-scan-worker-protocol'
|
||||
import { createUsageScanWorkerTransport } from './usage-scan-worker-cache-cleanup'
|
||||
|
||||
// #20940: serialize provider scans with WorkerThreadRequestQueue; LazyWorkerThreadHost owns thread lifetime and the spawn module owns layout/singleton.
|
||||
// Provider-local ordering shares one worker; LazyWorkerThreadHost owns its lifetime.
|
||||
|
||||
// Cold corpora can exceed any total deadline; time only gaps between progress reports. A huge rollout is one gap, so keep the window generous.
|
||||
export const USAGE_SCAN_NO_PROGRESS_TIMEOUT_MS = 10 * 60_000
|
||||
@@ -39,25 +36,26 @@ type ProviderScanResult<TSession, TDaily> = {
|
||||
|
||||
/** Unavailable, timed-out or crashed workers reject; stores retain their previous projection instead of parsing on main. */
|
||||
export class UsageScanWorkerClient {
|
||||
private readonly queue: WorkerThreadRequestQueue<UsageScanWorkerRequest, UsageScanWorkerResponse>
|
||||
private readonly queue: UsageScanRequestScheduler
|
||||
|
||||
constructor(options: { workerFactory: WorkerThreadFactory; log?: (message: string) => void }) {
|
||||
const log = options.log ?? ((message: string) => console.warn(message))
|
||||
this.queue = new WorkerThreadRequestQueue({
|
||||
this.queue = new UsageScanRequestScheduler({
|
||||
factory: () => createUsageScanWorkerTransport(options.workerFactory),
|
||||
idleTeardownMs: IDLE_TEARDOWN_MS,
|
||||
maxConsecutiveDeaths: MAX_CONSECUTIVE_DEATHS,
|
||||
createUnavailableError: (message) => new UsageScanWorkerUnavailableError(message),
|
||||
isProgress: isUsageScanWorkerProgress,
|
||||
describeTimeout: (timeoutMs) => `Usage scan worker reported no progress for ${timeoutMs}ms`,
|
||||
describeExit: (code) => `Usage scan worker exited with code ${code}`,
|
||||
describeCrashLoop: (lastError) => `Usage scan worker crashed repeatedly (${lastError})`,
|
||||
// Missing/resource-starved workers must fail closed; main-thread scanning would reintroduce event-loop blocking.
|
||||
onUnavailable: (err) =>
|
||||
log(`[usage-scan] worker unavailable; usage scans will report an error. ${message(err)}`)
|
||||
})
|
||||
}
|
||||
|
||||
async dispose(): Promise<void> {
|
||||
this.queue.dispose()
|
||||
await this.queue.retired()
|
||||
}
|
||||
|
||||
/** sourceCache identifies provider records that stay on the worker. */
|
||||
scan(body: UsageScanWorkerScanBody): Promise<UsageScanWorkerValue> {
|
||||
return this.dispatch(body)
|
||||
@@ -79,10 +77,7 @@ export class UsageScanWorkerClient {
|
||||
}
|
||||
|
||||
private async dispatch(body: UsageScanWorkerRequestBody): Promise<UsageScanWorkerValue> {
|
||||
const response = await this.queue.dispatch(
|
||||
(id) => ({ ...body, id }),
|
||||
USAGE_SCAN_NO_PROGRESS_TIMEOUT_MS
|
||||
)
|
||||
const response = await this.queue.dispatch(body, USAGE_SCAN_NO_PROGRESS_TIMEOUT_MS)
|
||||
if (!response.ok) {
|
||||
throw new Error(response.error)
|
||||
}
|
||||
|
||||
@@ -1,4 +1,3 @@
|
||||
import { claudeProfileTranscriptDirs } from '../../claude-usage/transcript-file-discovery'
|
||||
import { existsSync } from 'node:fs'
|
||||
import type { Worker } from 'node:worker_threads'
|
||||
import { spawnWorkerThread } from '../../worker-threads/worker-thread-spawn'
|
||||
@@ -49,18 +48,15 @@ function getSharedClient(): UsageScanWorkerClient {
|
||||
|
||||
export function scanClaudeUsageFilesViaWorker(
|
||||
worktrees: UsageScanWorktreeRef[],
|
||||
sourceCache: UsageSourceCacheRef
|
||||
sourceCache: UsageSourceCacheRef,
|
||||
profileDirs: string[]
|
||||
): Promise<{
|
||||
sessions: ClaudeUsageSession[]
|
||||
dailyAggregates: ClaudeUsageDailyAggregate[]
|
||||
}> {
|
||||
return scanClaudeUsageOnWorker(
|
||||
(body) =>
|
||||
getSharedClient().scan(
|
||||
body.providerId === 'claude'
|
||||
? { ...body, profileDirs: claudeProfileTranscriptDirs() }
|
||||
: body
|
||||
),
|
||||
getSharedClient().scan(body.providerId === 'claude' ? { ...body, profileDirs } : body),
|
||||
worktrees,
|
||||
sourceCache
|
||||
)
|
||||
@@ -102,3 +98,9 @@ export function splitUsageCacheFileViaWorker(
|
||||
): Promise<UsageCacheSplitResult> {
|
||||
return getSharedClient().splitCacheFile(request)
|
||||
}
|
||||
|
||||
export async function stopUsageScanWorker(): Promise<void> {
|
||||
const client = sharedClient
|
||||
sharedClient = null
|
||||
await client?.dispose()
|
||||
}
|
||||
|
||||
@@ -0,0 +1,274 @@
|
||||
import { EventEmitter } from 'node:events'
|
||||
import type { Worker } from 'node:worker_threads'
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
|
||||
import {
|
||||
UsageScanWorkerClient,
|
||||
UsageScanWorkerUnavailableError,
|
||||
USAGE_SCAN_NO_PROGRESS_TIMEOUT_MS
|
||||
} from './scan-worker/usage-scan-worker-client'
|
||||
import type {
|
||||
UsageScanWorkerProviderId,
|
||||
UsageScanWorkerRequest,
|
||||
UsageScanWorkerScanBody,
|
||||
UsageScanWorkerValue
|
||||
} from './scan-worker/usage-scan-worker-protocol'
|
||||
|
||||
class ScanWorker extends EventEmitter {
|
||||
posted: UsageScanWorkerRequest[] = []
|
||||
exit: Promise<number> = Promise.resolve(1)
|
||||
terminate = vi.fn(() => this.exit)
|
||||
unref(): void {}
|
||||
|
||||
postMessage(request: UsageScanWorkerRequest): void {
|
||||
this.posted.push(request)
|
||||
}
|
||||
|
||||
complete(index: number): void {
|
||||
const request = this.posted[index]
|
||||
const value: UsageScanWorkerValue =
|
||||
request.operation === 'scan'
|
||||
? { operation: 'scan', providerId: request.providerId, sessions: [], dailyAggregates: [] }
|
||||
: { operation: 'splitCacheFile', reportText: 'saved report', migrated: false }
|
||||
this.emit('message', { id: request.id, ok: true, value })
|
||||
}
|
||||
}
|
||||
|
||||
function body(providerId: UsageScanWorkerProviderId, population = 1): UsageScanWorkerScanBody {
|
||||
return {
|
||||
operation: 'scan',
|
||||
providerId,
|
||||
worktrees: Array.from({ length: population }, (_, i) => ({
|
||||
repoId: 'repo',
|
||||
worktreeId: `tree-${i}`,
|
||||
displayName: `tree-${i}`,
|
||||
path: `/tree-${i}`
|
||||
})),
|
||||
sourceCache: {
|
||||
path: `/${providerId}-sources.json`,
|
||||
schemaVersion: 1,
|
||||
worktreeFingerprint: '[]',
|
||||
reuse: true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const clients: UsageScanWorkerClient[] = []
|
||||
function fixture(): { client: UsageScanWorkerClient; workers: ScanWorker[] } {
|
||||
const workers: ScanWorker[] = []
|
||||
const client = new UsageScanWorkerClient({
|
||||
workerFactory: () => {
|
||||
const worker = new ScanWorker()
|
||||
workers.push(worker)
|
||||
// oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: the fixture implements every Worker member the host and transport use; no native Worker API is reached.
|
||||
return worker as unknown as Worker
|
||||
},
|
||||
log: () => {}
|
||||
})
|
||||
clients.push(client)
|
||||
return { client, workers }
|
||||
}
|
||||
|
||||
beforeEach(() => vi.useFakeTimers())
|
||||
afterEach(async () => {
|
||||
await Promise.all(clients.splice(0).map((client) => client.dispose().catch(() => {})))
|
||||
vi.useRealTimers()
|
||||
})
|
||||
|
||||
describe('usage scan scheduling through the production client', () => {
|
||||
it.each([100, 1000])(
|
||||
'admits only four providers at %i worktrees and routes unordered replies',
|
||||
async (population) => {
|
||||
const { client, workers } = fixture()
|
||||
const calls = [
|
||||
client.scan(body('claude', population)),
|
||||
client.scan(body('codex', population)),
|
||||
client.scan(body('opencode', population)),
|
||||
client.scan(body('muse', population))
|
||||
]
|
||||
const worker = workers[0]
|
||||
expect(worker.posted.map((request) => request.providerId)).toEqual([
|
||||
'claude',
|
||||
'codex',
|
||||
'opencode',
|
||||
'muse'
|
||||
])
|
||||
worker.emit('message', { id: worker.posted[1].id, filesScanned: 1 })
|
||||
for (const index of [2, 0, 3, 1]) {
|
||||
worker.complete(index)
|
||||
}
|
||||
expect(await Promise.all(calls)).toEqual([
|
||||
{ operation: 'scan', providerId: 'claude', sessions: [], dailyAggregates: [] },
|
||||
{ operation: 'scan', providerId: 'codex', sessions: [], dailyAggregates: [] },
|
||||
{ operation: 'scan', providerId: 'opencode', sessions: [], dailyAggregates: [] },
|
||||
{ operation: 'scan', providerId: 'muse', sessions: [], dailyAggregates: [] }
|
||||
])
|
||||
}
|
||||
)
|
||||
|
||||
it('keeps same-provider FIFO and makes a split an exclusive barrier without overtaking', async () => {
|
||||
const { client, workers } = fixture()
|
||||
const first = client.scan(body('claude'))
|
||||
const second = client.scan(body('claude'))
|
||||
const codex = client.scan(body('codex'))
|
||||
const split = client.splitCacheFile({ cacheFile: '/claude.json', sourceKey: 'processedFiles' })
|
||||
const muse = client.scan(body('muse'))
|
||||
const worker = workers[0]
|
||||
expect(worker.posted.map((request) => request.providerId)).toEqual(['claude', 'codex'])
|
||||
worker.complete(0)
|
||||
expect(worker.posted.map((request) => request.providerId)).toEqual([
|
||||
'claude',
|
||||
'codex',
|
||||
'claude'
|
||||
])
|
||||
worker.complete(2)
|
||||
expect(worker.posted).toHaveLength(3)
|
||||
worker.complete(1)
|
||||
expect(worker.posted[3].operation).toBe('splitCacheFile')
|
||||
expect(worker.posted).toHaveLength(4)
|
||||
worker.complete(3)
|
||||
expect(worker.posted[4].providerId).toBe('muse')
|
||||
worker.complete(4)
|
||||
await Promise.all([first, second, codex, muse])
|
||||
expect(await split).toEqual({ reportText: 'saved report', migrated: false })
|
||||
})
|
||||
|
||||
it('settles every posted call on one fault, parks only unposted work, and never replays', async () => {
|
||||
const { client, workers } = fixture()
|
||||
const posted = [
|
||||
client.scan(body('claude')),
|
||||
client.scan(body('codex')),
|
||||
client.scan(body('muse'))
|
||||
].map((promise) => promise.catch((error: unknown) => error))
|
||||
const pending = client.scan(body('claude'))
|
||||
const old = workers[0]
|
||||
let proveExit!: (code: number) => void
|
||||
old.exit = new Promise((resolve) => {
|
||||
proveExit = resolve
|
||||
})
|
||||
const failure = new Error('worker crashed')
|
||||
old.emit('error', failure)
|
||||
expect(await Promise.all(posted)).toEqual([failure, failure, failure])
|
||||
expect(old.terminate).toHaveBeenCalledTimes(1)
|
||||
await vi.advanceTimersByTimeAsync(USAGE_SCAN_NO_PROGRESS_TIMEOUT_MS + 1)
|
||||
expect(workers).toHaveLength(1)
|
||||
proveExit(1)
|
||||
await vi.advanceTimersByTimeAsync(0)
|
||||
expect(workers).toHaveLength(2)
|
||||
expect(workers[1].posted.map((request) => request.providerId)).toEqual(['claude'])
|
||||
old.complete(0)
|
||||
workers[1].complete(0)
|
||||
expect(await pending).toEqual({
|
||||
operation: 'scan',
|
||||
providerId: 'claude',
|
||||
sessions: [],
|
||||
dailyAggregates: []
|
||||
})
|
||||
})
|
||||
|
||||
it('times each dispatched ID independently, ignoring other-provider and unknown progress', async () => {
|
||||
const { client, workers } = fixture()
|
||||
const claude = client.scan(body('claude')).catch((error: unknown) => error)
|
||||
const codex = client.scan(body('codex')).catch((error: unknown) => error)
|
||||
await vi.advanceTimersByTimeAsync(USAGE_SCAN_NO_PROGRESS_TIMEOUT_MS / 2)
|
||||
workers[0].emit('message', { id: workers[0].posted[1].id, filesScanned: 1 })
|
||||
workers[0].emit('message', { id: 999, filesScanned: 1 })
|
||||
await vi.advanceTimersByTimeAsync(USAGE_SCAN_NO_PROGRESS_TIMEOUT_MS / 2 + 1)
|
||||
expect(await claude).toEqual(new Error('Usage scan worker reported no progress for 600000ms'))
|
||||
expect(await codex).toEqual(new Error('Usage scan worker reported no progress for 600000ms'))
|
||||
expect(workers[0].terminate).toHaveBeenCalledTimes(1)
|
||||
})
|
||||
|
||||
it('counts a shared death once and caps consecutive worker deaths in one queued burst', async () => {
|
||||
const { client, workers } = fixture()
|
||||
const calls = [
|
||||
client.scan(body('claude')),
|
||||
client.scan(body('codex')),
|
||||
client.scan(body('opencode')),
|
||||
client.scan(body('muse')),
|
||||
client.scan(body('claude')),
|
||||
client.scan(body('claude')),
|
||||
client.scan(body('claude'))
|
||||
].map((promise) => promise.catch((error: unknown) => error))
|
||||
workers[0].emit('error', new Error('fault one'))
|
||||
await vi.advanceTimersByTimeAsync(0)
|
||||
expect(workers).toHaveLength(2)
|
||||
workers[1].emit('error', new Error('fault two'))
|
||||
await vi.advanceTimersByTimeAsync(0)
|
||||
expect(workers).toHaveLength(3)
|
||||
workers[2].emit('error', new Error('fault three'))
|
||||
const results = await Promise.all(calls)
|
||||
expect(results).toEqual([
|
||||
new Error('fault one'),
|
||||
new Error('fault one'),
|
||||
new Error('fault one'),
|
||||
new Error('fault one'),
|
||||
new Error('fault two'),
|
||||
new Error('fault three'),
|
||||
new Error('Usage scan worker crashed repeatedly (fault three)')
|
||||
])
|
||||
expect(workers).toHaveLength(3)
|
||||
})
|
||||
|
||||
it('keeps a request-level error local while the other provider completes', async () => {
|
||||
const { client, workers } = fixture()
|
||||
const claude = client.scan(body('claude')).catch((error: unknown) => error)
|
||||
const codex = client.scan(body('codex'))
|
||||
workers[0].emit('message', { id: workers[0].posted[0].id, ok: false, error: 'scan failed' })
|
||||
expect(await claude).toEqual(new Error('scan failed'))
|
||||
expect(workers[0].terminate).not.toHaveBeenCalled()
|
||||
workers[0].complete(1)
|
||||
expect(await codex).toEqual({
|
||||
operation: 'scan',
|
||||
providerId: 'codex',
|
||||
sessions: [],
|
||||
dailyAggregates: []
|
||||
})
|
||||
})
|
||||
|
||||
it('cancels a parked wake when disposed and cannot spawn after the exact exit', async () => {
|
||||
const { client, workers } = fixture()
|
||||
const active = client.scan(body('claude')).catch((error: unknown) => error)
|
||||
const pending = client.scan(body('claude')).catch((error: unknown) => error)
|
||||
let proveExit!: (code: number) => void
|
||||
workers[0].exit = new Promise((resolve) => {
|
||||
proveExit = resolve
|
||||
})
|
||||
workers[0].emit('error', new Error('worker crashed'))
|
||||
const disposed = client.dispose()
|
||||
expect(await active).toEqual(new Error('worker crashed'))
|
||||
expect(await pending).toEqual(new Error('Usage scan scheduler disposed'))
|
||||
proveExit(1)
|
||||
await disposed
|
||||
await vi.advanceTimersByTimeAsync(0)
|
||||
expect(workers).toHaveLength(1)
|
||||
})
|
||||
|
||||
it('keeps failed termination fenced until actual exit and rejects strict disposal honestly', async () => {
|
||||
const { client, workers } = fixture()
|
||||
const first = client.scan(body('claude')).catch((error: unknown) => error)
|
||||
const pending = client.scan(body('claude')).catch((error: unknown) => error)
|
||||
const failure = new Error('termination failed')
|
||||
workers[0].exit = Promise.reject(failure)
|
||||
workers[0].emit('error', new Error('worker crashed'))
|
||||
expect(await first).toEqual(new Error('worker crashed'))
|
||||
expect(await pending).toBe(failure)
|
||||
expect(await client.scan(body('codex')).catch((error: unknown) => error)).toBe(failure)
|
||||
expect(workers).toHaveLength(1)
|
||||
await expect(client.dispose()).rejects.toBe(failure)
|
||||
workers[0].emit('exit', 1)
|
||||
await client.dispose()
|
||||
})
|
||||
|
||||
it('preserves unavailable classification without falling back to caller-thread scans', async () => {
|
||||
const client = new UsageScanWorkerClient({
|
||||
workerFactory: () => {
|
||||
throw new Error('cannot spawn')
|
||||
},
|
||||
log: () => {}
|
||||
})
|
||||
clients.push(client)
|
||||
await expect(client.scan(body('claude'))).rejects.toBeInstanceOf(
|
||||
UsageScanWorkerUnavailableError
|
||||
)
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,219 @@
|
||||
import { LazyWorkerThreadHost } from '../worker-threads/lazy-worker-thread-host'
|
||||
import type { WorkerThreadFactory } from '../worker-threads/worker-thread-transport'
|
||||
import {
|
||||
isUsageScanWorkerProgress,
|
||||
type UsageScanWorkerMessage,
|
||||
type UsageScanWorkerRequest,
|
||||
type UsageScanWorkerRequestBody,
|
||||
type UsageScanWorkerResponse
|
||||
} from './scan-worker/usage-scan-worker-protocol'
|
||||
|
||||
type ScanCall = {
|
||||
request: UsageScanWorkerRequest
|
||||
timeoutMs: number
|
||||
timer: NodeJS.Timeout | null
|
||||
resolve: (response: UsageScanWorkerResponse) => void
|
||||
reject: (error: Error) => void
|
||||
}
|
||||
|
||||
/** Independent providers share a worker; each provider and cache split retain exclusive ownership. */
|
||||
export class UsageScanRequestScheduler {
|
||||
private readonly active = new Map<number, ScanCall>()
|
||||
private pending: ScanCall[] = []
|
||||
private nextId = 1
|
||||
private consecutiveDeaths = 0
|
||||
private disposed = false
|
||||
private retirementWait: Promise<void> | null = null
|
||||
private readonly host: LazyWorkerThreadHost<UsageScanWorkerMessage>
|
||||
|
||||
constructor(
|
||||
private readonly options: {
|
||||
factory: WorkerThreadFactory
|
||||
idleTeardownMs: number
|
||||
maxConsecutiveDeaths: number
|
||||
createUnavailableError: (message: string) => Error
|
||||
onUnavailable: (error: unknown) => void
|
||||
}
|
||||
) {
|
||||
this.host = new LazyWorkerThreadHost({
|
||||
factory: options.factory,
|
||||
idleTeardownMs: options.idleTeardownMs,
|
||||
onMessage: (response) => this.onMessage(response),
|
||||
onError: (error) => this.onFault(error),
|
||||
onExit: (code) => this.onExit(code),
|
||||
isIdle: () => this.active.size === 0 && this.pending.length === 0,
|
||||
onUnavailable: options.onUnavailable,
|
||||
awaitRetirement: true,
|
||||
strictRetirement: true
|
||||
})
|
||||
}
|
||||
|
||||
dispatch(body: UsageScanWorkerRequestBody, timeoutMs: number): Promise<UsageScanWorkerResponse> {
|
||||
return new Promise((resolve, reject) => {
|
||||
if (this.disposed) {
|
||||
reject(new Error('Usage scan scheduler disposed'))
|
||||
return
|
||||
}
|
||||
if (this.active.size === 0 && this.pending.length === 0) {
|
||||
this.consecutiveDeaths = 0
|
||||
}
|
||||
this.pending.push({
|
||||
request: { ...body, id: this.nextId++ },
|
||||
timeoutMs,
|
||||
timer: null,
|
||||
resolve,
|
||||
reject
|
||||
})
|
||||
this.pump()
|
||||
})
|
||||
}
|
||||
|
||||
dispose(): void {
|
||||
this.disposed = true
|
||||
this.host.destroy()
|
||||
const calls = [...this.active.values(), ...this.pending]
|
||||
this.pending = []
|
||||
for (const call of calls) {
|
||||
this.settle(call, () => call.reject(new Error('Usage scan scheduler disposed')))
|
||||
}
|
||||
}
|
||||
|
||||
retired(): Promise<void> {
|
||||
return this.host.retired()
|
||||
}
|
||||
|
||||
private pump(): void {
|
||||
if (this.disposed) {
|
||||
return
|
||||
}
|
||||
if (this.pending.length === 0) {
|
||||
if (this.active.size === 0) {
|
||||
this.host.scheduleIdleTeardown()
|
||||
}
|
||||
return
|
||||
}
|
||||
if (this.host.isRetiring) {
|
||||
this.waitForRetirement()
|
||||
return
|
||||
}
|
||||
while (this.pending.length > 0) {
|
||||
const index = this.nextEligible()
|
||||
if (index < 0) {
|
||||
return
|
||||
}
|
||||
const worker = this.host.ensure()
|
||||
if (!worker) {
|
||||
this.rejectPending(this.options.createUnavailableError('worker spawn failed'))
|
||||
return
|
||||
}
|
||||
const call = this.pending.splice(index, 1)[0]
|
||||
this.active.set(call.request.id, call)
|
||||
this.host.clearIdleTimer()
|
||||
this.armDeadline(call)
|
||||
try {
|
||||
worker.postMessage(call.request)
|
||||
} catch (error) {
|
||||
this.onFault(error instanceof Error ? error : new Error(String(error)))
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private nextEligible(): number {
|
||||
const requests = [...this.active.values()].map((call) => call.request)
|
||||
if (requests.some((request) => request.operation === 'splitCacheFile')) {
|
||||
return -1
|
||||
}
|
||||
for (let index = 0; index < this.pending.length; index++) {
|
||||
const request = this.pending[index].request
|
||||
if (request.operation === 'splitCacheFile') {
|
||||
return index === 0 && requests.length === 0 ? index : -1
|
||||
}
|
||||
if (!requests.some((active) => active.providerId === request.providerId)) {
|
||||
return index
|
||||
}
|
||||
}
|
||||
return -1
|
||||
}
|
||||
|
||||
private waitForRetirement(): void {
|
||||
if (this.retirementWait) {
|
||||
return
|
||||
}
|
||||
this.retirementWait = this.host.retired().then(
|
||||
() => {
|
||||
this.retirementWait = null
|
||||
this.pump()
|
||||
},
|
||||
(error: unknown) => {
|
||||
this.retirementWait = null
|
||||
this.rejectPending(error instanceof Error ? error : new Error(String(error)))
|
||||
}
|
||||
)
|
||||
}
|
||||
|
||||
private armDeadline(call: ScanCall): void {
|
||||
if (call.timer) {
|
||||
clearTimeout(call.timer)
|
||||
}
|
||||
call.timer = setTimeout(() => {
|
||||
if (this.active.get(call.request.id) === call) {
|
||||
this.onFault(new Error(`Usage scan worker reported no progress for ${call.timeoutMs}ms`))
|
||||
}
|
||||
}, call.timeoutMs)
|
||||
call.timer.unref?.()
|
||||
}
|
||||
|
||||
private onMessage(response: UsageScanWorkerMessage): void {
|
||||
const call = this.active.get(response.id)
|
||||
if (!call) {
|
||||
return
|
||||
}
|
||||
if (isUsageScanWorkerProgress(response)) {
|
||||
this.armDeadline(call)
|
||||
return
|
||||
}
|
||||
this.consecutiveDeaths = 0
|
||||
this.settle(call, () => call.resolve(response))
|
||||
this.pump()
|
||||
}
|
||||
|
||||
private onExit(code: number): void {
|
||||
if (code === 0 && this.active.size === 0 && this.pending.length === 0) {
|
||||
this.host.destroy()
|
||||
return
|
||||
}
|
||||
this.onFault(new Error(`Usage scan worker exited with code ${code}`))
|
||||
}
|
||||
|
||||
private onFault(error: Error): void {
|
||||
const posted = [...this.active.values()]
|
||||
this.host.destroy()
|
||||
for (const call of posted) {
|
||||
this.settle(call, () => call.reject(error))
|
||||
}
|
||||
if (posted.length > 0 && ++this.consecutiveDeaths >= this.options.maxConsecutiveDeaths) {
|
||||
this.consecutiveDeaths = 0
|
||||
this.rejectPending(new Error(`Usage scan worker crashed repeatedly (${error.message})`))
|
||||
return
|
||||
}
|
||||
this.pump()
|
||||
}
|
||||
|
||||
private rejectPending(error: Error): void {
|
||||
const pending = this.pending
|
||||
this.pending = []
|
||||
for (const call of pending) {
|
||||
this.settle(call, () => call.reject(error))
|
||||
}
|
||||
}
|
||||
|
||||
private settle(call: ScanCall, finish: () => void): void {
|
||||
if (call.timer) {
|
||||
clearTimeout(call.timer)
|
||||
call.timer = null
|
||||
}
|
||||
this.active.delete(call.request.id)
|
||||
finish()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,21 @@
|
||||
import type { UsageStoreHost, UsageWorktreeCatalog } from './usage-store-host'
|
||||
|
||||
export function createTestUsageStoreHost(options: {
|
||||
userDataPath: string
|
||||
catalog?: UsageWorktreeCatalog
|
||||
profileDirs?: string[]
|
||||
telemetryEnabled?: boolean
|
||||
capture?: UsageStoreHost['telemetry']['capture']
|
||||
signal?: AbortSignal
|
||||
}): UsageStoreHost {
|
||||
return {
|
||||
userDataPath: options.userDataPath,
|
||||
signal: options.signal ?? new AbortController().signal,
|
||||
readCatalog: async () => options.catalog ?? { repoIds: [], worktreesByRepo: [] },
|
||||
claudeProfileDirs: async () => options.profileDirs ?? [],
|
||||
telemetry: {
|
||||
isEnabled: async () => options.telemetryEnabled ?? false,
|
||||
capture: options.capture ?? (async () => true)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
import type { AgentTokenUsage } from '../../shared/telemetry-agent-token-usage-schema'
|
||||
import type { UsageWorktreeRef } from './worktree-attribution/usage-worktree-metadata'
|
||||
|
||||
export type UsageWorktreeCatalog = {
|
||||
repoIds: string[]
|
||||
worktreesByRepo: [string, UsageWorktreeRef[]][]
|
||||
}
|
||||
|
||||
export type UsageTelemetryHost = {
|
||||
isEnabled(): Promise<boolean>
|
||||
capture(snapshot: AgentTokenUsage): Promise<boolean>
|
||||
}
|
||||
|
||||
export type UsageStoreHost = {
|
||||
signal: AbortSignal
|
||||
userDataPath: string
|
||||
readCatalog(): Promise<UsageWorktreeCatalog>
|
||||
claudeProfileDirs(): Promise<string[]>
|
||||
telemetry: UsageTelemetryHost
|
||||
}
|
||||
@@ -21,7 +21,7 @@ export function getUsageWorktreeFingerprint(
|
||||
}
|
||||
|
||||
export function createWorktreeRefs(
|
||||
repos: Repo[],
|
||||
repos: readonly Pick<Repo, 'id'>[],
|
||||
worktreesByRepo: Map<string, { path: string; worktreeId: string; displayName: string }[]>
|
||||
): UsageScanWorktreeRef[] {
|
||||
const refs: UsageScanWorktreeRef[] = []
|
||||
|
||||
Reference in New Issue
Block a user