diff --git a/CHANGELOG.md b/CHANGELOG.md index da4881dd..36fc07b1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,10 @@ This repo is a CortexKit-maintained Anthropic auth monorepo for OpenCode and Pi. The OpenCode package is a fork of the original `@ex-machina/opencode-anthropic-auth` plugin, so older entries below the initial CortexKit release are inherited from upstream package history. +## Unreleased + +- Requests no longer wait on relay or session-pin bookkeeping locks; relay settings and tokens are checked together before sending. + ## 2.1.1 ### Packaging diff --git a/mutations.toml b/mutations.toml index 40476c5a..034fbd2f 100644 --- a/mutations.toml +++ b/mutations.toml @@ -1354,3 +1354,55 @@ expect_message = "error: expect(received).toEqual(expected)" broad_command = ["bun", "scripts/run-mutation-bun.ts", "--broad", "packages/core", "tmp/mutations/core.xml"] broad_report = "tmp/mutations/core.xml" broad_id = "{name}" + +[[control]] +id = "native-relay-read-never-waits-for-writer" +guards = "Reading published relay configuration never acquires the exclusive relay writer lock." +file = "packages/core/src/native-account-runtime.ts" +old = ''' async getRelayConfig() { + const { mode, paths, store } = await selectedStorage()''' +new = ''' async getRelayConfig() { + const { mode, paths, store } = await selectedStorage() + // NON-VACUITY BREAK: restore relay writer contention on the request path. + await withLock(paths.config, { ...POOL_LOCK_DEFAULTS, name: 'native-relay' }, async () => {})''' +test_file = "packages/core/src/tests/native-account-runtime.test.ts" +runner = "command" +command = ["bun", "scripts/run-mutation-bun.ts", "packages/core/src/tests/native-account-runtime.test.ts", "{test}"] +test_count_pattern = "CKDEV executed {count} tests" +expect_red = ["relay config reads published state while another writer holds native-relay"] +expect_message = "error: expect(received).toBe(expected)" +broad_command = ["bun", "scripts/run-mutation-bun.ts", "--broad", "packages/core", "tmp/mutations/core.xml"] +broad_report = "tmp/mutations/core.xml" +broad_id = "{name}" + +[[control]] +id = "native-relay-generation-agreement" +guards = "A relay reader never combines the old URL with a token published by another generation." +file = "packages/core/src/native-account-runtime.ts" +old = "if (generation !== state?.relayGeneration) continue" +new = "if (false && generation !== state?.relayGeneration) continue // NON-VACUITY BREAK: ignore crossed relay generations." +test_file = "packages/core/src/tests/native-account-runtime.test.ts" +runner = "command" +command = ["bun", "scripts/run-mutation-bun.ts", "packages/core/src/tests/native-account-runtime.test.ts", "{test}"] +test_count_pattern = "CKDEV executed {count} tests" +expect_red = ["relay reader racing a writer never combines old URL and new token (stamped)"] +expect_message = "error: expect(received).toMatchObject(expected)" +broad_command = ["bun", "scripts/run-mutation-bun.ts", "--broad", "packages/core", "tmp/mutations/core.xml"] +broad_report = "tmp/mutations/core.xml" +broad_id = "{name}" + +[[control]] +id = "sticky-routing-decision-never-joins-writer" +guards = "A request reads persisted session pins and decides its route without joining a writer lease." +file = "packages/core/src/sticky-routing.ts" +old = " if (!input.sessionId) return null\n this.recoverLocalPins()" +new = " if (!input.sessionId) return null\n // NON-VACUITY BREAK: decide a request only after joining the writer lock.\n const decisionLease = await this.acquireLock()\n await decisionLease.release()\n this.recoverLocalPins()" +test_file = "packages/core/src/tests/sticky-routing-contention.test.ts" +runner = "command" +command = ["bun", "scripts/run-mutation-bun.ts", "packages/core/src/tests/sticky-routing-contention.test.ts", "{test}"] +test_count_pattern = "CKDEV executed {count} tests" +expect_red = ["sticky routing decisions read published pins without waiting and publish after the writer releases"] +expect_message = "error: expect(received).toBe(expected)" +broad_command = ["bun", "scripts/run-mutation-bun.ts", "--broad", "packages/core", "tmp/mutations/core.xml"] +broad_report = "tmp/mutations/core.xml" +broad_id = "{name}" diff --git a/packages/core/scripts/check-native-type-closure.test.ts b/packages/core/scripts/check-native-type-closure.test.ts index 6a5c042d..a400f72a 100644 --- a/packages/core/scripts/check-native-type-closure.test.ts +++ b/packages/core/scripts/check-native-type-closure.test.ts @@ -422,7 +422,7 @@ test('packed Core types and Node runtime work in a fresh consumer without common hold: (point, id) => { const name: 'refresh-before-provider' | 'pull-before-request' = point; void [name, id]; }, logger: { warn: (message) => { const text: string = message; void text; } }, onLockEvent: (event) => { - const phase: 'acquired' | 'released' | 'contended' = event.type; + const phase: 'acquiring' | 'acquired' | 'released' | 'contended' = event.type; const diagnostic: NativeLockEvent = event; const fields: string[] = [event.name, event.path]; void [phase, diagnostic, fields]; @@ -436,7 +436,7 @@ test('packed Core types and Node runtime work in a fresh consumer without common compatible.onLockEvent?.(oldProducerRecord); const contention: Parameters>[0] = { type: 'contended', name: 'synthetic', path: paths.state }; options.onLockEvent?.(contention); - type ClosedEvent = NativeLockEvent['type'] extends 'acquired' | 'released' | 'contended' ? true : false; + type ClosedEvent = NativeLockEvent['type'] extends 'acquiring' | 'acquired' | 'released' | 'contended' ? true : false; const closed: ClosedEvent = true; void closed; const roster: string = paths.roster; diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index 8682e351..66839158 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -116,6 +116,7 @@ export type { } from './native-local-credential-service.ts' export { createNativeLocalCredentialService } from './native-local-credential-service.ts' export type { NativeLocalExternalPolicy } from './native-local-runtime-readers.ts' +export type { NativeLockEvent } from './native-lock-observer.ts' export type { NativeMenuActionValues, NativeMenuCapabilityOutcome, @@ -191,6 +192,7 @@ export type { NativeLocalFailurePolicy, NativeLocalFailureSkipReason, } from './native-runtime.ts' +export { NativeRuntimeError } from './native-runtime.ts' export { nativeRuntimeOptionIdentity } from './native-runtime-registry.ts' export * from './native-settings.ts' export { diff --git a/packages/core/src/native-account-runtime.ts b/packages/core/src/native-account-runtime.ts index 2de93b08..c347f154 100644 --- a/packages/core/src/native-account-runtime.ts +++ b/packages/core/src/native-account-runtime.ts @@ -1,12 +1,11 @@ import { randomUUID } from 'node:crypto' -import { withLock } from '@cortexkit/common-auth/fs' +import { withLock as acquireLock } from '@cortexkit/common-auth/fs' import { fingerprintOf, POOL_LOCK_DEFAULTS, type PoolCredential, type PoolRow, } from '@cortexkit/common-auth/store' - import { buildRefreshOperationError, fetchOAuthQuotaSnapshot, @@ -52,6 +51,11 @@ import { type NativeLocalCredentialServiceOptions, } from './native-local-credential-service.ts' import type { NativeLocalExternalPolicy } from './native-local-runtime-readers.ts' +import { + type NativeLockEvent, + nativeLockObserver, + observeNativeLocks, +} from './native-lock-observer.ts' import { fromNativeQuotaMap, nativeQuotaCodec, @@ -126,7 +130,11 @@ export interface NativeAccountRuntimeOptions { local?: Partial< Pick< NativeLocalCredentialServiceOptions, - 'refreshToken' | 'resolveIdentity' | 'failurePolicy' | 'now' + | 'refreshToken' + | 'resolveIdentity' + | 'failurePolicy' + | 'now' + | 'onLockEvent' > > vault?: Omit @@ -160,9 +168,11 @@ export type NativeRelayUpdate = Partial> & { } export interface NativeAccountRuntime { + readonly onLockEvent?: (event: NativeLockEvent) => void readonly vault: NativeVaultRuntime readonly quotaObservations: NativeQuotaObservations drainQuotaPublications(): Promise + drainRoutingPublications(): Promise /** Fresh private relay authorization for transport only; never include it in a status snapshot. */ getRelayConfig(): Promise updateRelay(patch: NativeRelayUpdate): Promise @@ -267,7 +277,7 @@ export interface NativeAccountRuntime { fetchImpl?: typeof fetch, signal?: AbortSignal, ): Promise - close(): void + close(): void | Promise getSettingsPath(): Promise } @@ -289,6 +299,7 @@ export function acquireNativeAccountRuntime( options.local?.resolveIdentity, options.local?.failurePolicy, options.local?.now, + options.local?.onLockEvent, ].map(nativeRuntimeOptionIdentity), ]) const { use, close } = sharedAccounts.acquire(key, () => @@ -296,10 +307,14 @@ export function acquireNativeAccountRuntime( ) const vault = leaseNativeVaultRuntime(() => use().vault, close) const quotaObservations = use().quotaObservations + const onLockEvent = use().onLockEvent + const drainRoutingPublications = use().drainRoutingPublications return { + onLockEvent, vault, quotaObservations, drainQuotaPublications: () => quotaObservations.drain(), + drainRoutingPublications, getRelayConfig: () => use().getRelayConfig(), getSettingsPath: () => use().getSettingsPath(), updateRelay: (patch) => use().updateRelay(patch), @@ -329,7 +344,10 @@ export function acquireNativeAccountRuntime( use().fetchProfile(id, transport, signal), fetchProfileForDisplay: (id, transport, signal) => use().fetchProfileForDisplay(id, transport, signal), - close: () => vault.close(), + close: () => { + vault.close() + return drainRoutingPublications() + }, } } @@ -360,6 +378,10 @@ export function createNativeAccountRuntime( options: NativeAccountRuntimeOptions, ): NativeAccountRuntime { const paths = Object.freeze({ ...options.paths }) + const onLockEvent = nativeLockObserver( + options.local?.onLockEvent ?? options.vault?.onLockEvent, + ) + const withLock = observeNativeLocks(() => acquireLock, onLockEvent) const now = options.local?.now ?? options.vault?.now ?? Date.now const env = { ...(options.vault?.env ?? process.env) } const vaultPaths = nativeVaultPaths(paths) @@ -370,6 +392,7 @@ export function createNativeAccountRuntime( const vaultRouter = new StickySessionRouter({ path: vaultPaths.routing, vaultSeed: routingSeed, + onLockEvent, now, }) let settingsBootstrap: Promise | undefined @@ -377,6 +400,7 @@ export function createNativeAccountRuntime( paths: vaultPaths, quota: nativeQuotaCodec, now, + onLockEvent: options.local?.onLockEvent, onStep: async (point) => { if (point === 'before-config-write' || point === 'before-state-write') await requireMode('claustrum') @@ -387,6 +411,7 @@ export function createNativeAccountRuntime( paths, quota: nativeQuotaCodec, now, + onLockEvent: options.local?.onLockEvent, onStep: async (point) => { if (point === 'before-config-write' || point === 'before-state-write') { if (options.beforePoolWrite) { @@ -401,6 +426,7 @@ export function createNativeAccountRuntime( ...options.vault, paths, host: options.host, + onLockEvent, }) function check(signal?: AbortSignal) { if (shutdown.signal.aborted) throw new NativeCustodyError('closed') @@ -504,11 +530,19 @@ export function createNativeAccountRuntime( try { return await withLock( config.path, - { ...POOL_LOCK_DEFAULTS, name: config.name }, + { + ...POOL_LOCK_DEFAULTS, + name: config.name, + onLockEvent: options.local?.onLockEvent, + }, async (configLease) => withLock( state.path, - { ...POOL_LOCK_DEFAULTS, name: state.name }, + { + ...POOL_LOCK_DEFAULTS, + name: state.name, + onLockEvent: options.local?.onLockEvent, + }, async (stateLease) => { if (!(await matches())) return false await updateNativeRuntime( @@ -529,6 +563,7 @@ export function createNativeAccountRuntime( return current }, { + onLockEvent: options.local?.onLockEvent, beforeRename: async () => { if (!(await matches())) throw new NativeRuntimeError('publication-refused') @@ -763,28 +798,34 @@ export function createNativeAccountRuntime( ) } const runtime: NativeAccountRuntime = { + onLockEvent, vault, quotaObservations: vault.quotaObservations, drainQuotaPublications: () => vault.quotaObservations.drain(), + drainRoutingPublications: () => vaultRouter.drain(), async getRelayConfig() { const { mode, paths, store } = await selectedStorage() - return withLock( - paths.config, - { ...POOL_LOCK_DEFAULTS, name: 'native-relay' }, - async (lease) => { - const settings = await store.readSettings() - if (settings.status !== 'ready') - throw new NativeRuntimeError('publication-refused') - const flags = relayFlags(settings.settings.relay) - const read = await readNativeRuntime(paths.runtime, paths.storageId) - const token = - read.status === 'ready' ? read.state.relay?.token : undefined - await requireMode(mode) - await lease.assertOwned() - if (!flags.enabled || !flags.url || !token) return null - return { ...flags, token } - }, - ) + // Both files are atomically published, but a reader can cross a relay + // change between reading its URL and private token. Never send that torn + // pair: all hosts retry a mismatched generation briefly, then refuse + // locally without dispatch or direct fallback. Two unstamped files are + // legacy data; a mixed stamped/unstamped pair must also wait for agreement. + for (let attempt = 0; attempt < 4; attempt++) { + if (attempt) await new Promise((resolve) => setTimeout(resolve, 10)) + const settings = await store.readSettings() + if (settings.status !== 'ready') + throw new NativeRuntimeError('publication-refused') + const flags = relayFlags(settings.settings.relay) + const generation = relayGeneration(settings.settings.relay) + const read = await readNativeRuntime(paths.runtime, paths.storageId) + await requireMode(mode) + const state = read.status === 'ready' ? read.state : undefined + if (generation !== state?.relayGeneration) continue + const token = state?.relay?.token + if (!flags.enabled || !flags.url || !token) return null + return { ...flags, token } + } + throw new NativeRuntimeError('relay-inconsistent') }, async updateRelay(input) { const patch = captureNativePlainData(input) as NativeRelayUpdate @@ -831,6 +872,8 @@ export function createNativeAccountRuntime( if (read.status !== 'ready') throw new NativeRuntimeError('publication-refused') const previous = relayFlags(read.settings.relay) + const previousGeneration = relayGeneration(read.settings.relay) + const generation = randomUUID() const metadata = await readNativeRuntime( paths.runtime, paths.storageId, @@ -851,36 +894,42 @@ export function createNativeAccountRuntime( }) if (next.enabled && (!next.url || !token)) throw new NativeRuntimeError('publication-refused') - async function writeFlags(flags: Omit) { + async function writeFlags( + flags: Omit, + stamp: string | undefined, + ) { await store.updateSettings(async (settings) => { await requireMode(mode) await lease.assertOwned() - settings.relay = flags + settings.relay = { + ...flags, + ...(stamp === undefined ? {} : { generation: stamp }), + } }) } - // Relay flags and its token live in separate files. Disable first so a crash - // cannot send the new token to the old endpoint. Credential-store locks - // protect the flags, and the relay lock ties them to the private token. - // Before publishing that token, recheck both relay ownership and that - // the selected local or vault credential source has not changed. - await writeFlags({ ...previous, enabled: false }) - if (patch.token !== undefined) - await updateNativeRuntime( - paths.runtime, - paths.storageId, - (current) => { - if (token) current.relay = { token } - else delete current.relay - return current - }, - { - beforeRename: async () => { - await requireMode(mode) - await lease.assertOwned() - }, + // Disable with the old stamp first. If publication stops before the + // runtime stamp changes, readers can still go direct; after it changes, + // a missing final settings write makes readers refuse locally instead. + // Writers stay serialized, including when a token is kept or removed. + // Recheck custody and lease ownership before each save. + await writeFlags({ ...previous, enabled: false }, previousGeneration) + await updateNativeRuntime( + paths.runtime, + paths.storageId, + (current) => { + if (token) current.relay = { token } + else delete current.relay + current.relayGeneration = generation + return current + }, + { + beforeRename: async () => { + await requireMode(mode) + await lease.assertOwned() }, - ) - await writeFlags(next) + }, + ) + await writeFlags(next, generation) }, ) }, @@ -912,6 +961,7 @@ export function createNativeAccountRuntime( ? [shutdown.signal, request.signal] : [shutdown.signal], () => check(request.signal), + onLockEvent, ) if (joined === 'deadline') { logger.warn('native-auth', 'account write did not finish in time', { @@ -1628,6 +1678,7 @@ export function createNativeAccountRuntime( shutdown.abort(new NativeCustodyError('closed')) vault.close() } + return vaultRouter.drain() }, } function readLocalProfile( @@ -1789,3 +1840,13 @@ function relayFlags(input: unknown): Omit { fallbackToDirect: value.fallbackToDirect !== false, } } + +function relayGeneration(input: unknown): string | undefined { + if (!input || typeof input !== 'object' || Array.isArray(input)) + return undefined + const generation = (input as Record).generation + if (generation === undefined) return undefined + if (typeof generation !== 'string' || !generation.trim()) + throw new NativeRuntimeError('relay-inconsistent') + return generation +} diff --git a/packages/core/src/native-local-credential-service.ts b/packages/core/src/native-local-credential-service.ts index cbd567ba..701c6f4f 100644 --- a/packages/core/src/native-local-credential-service.ts +++ b/packages/core/src/native-local-credential-service.ts @@ -1,9 +1,9 @@ import { types as utilTypes } from 'node:util' - import { createNativeLocalRuntimeReaders, type NativeLocalExternalPolicy, } from './native-local-runtime-readers.ts' +import type { NativeLockEvent } from './native-lock-observer.ts' import { nativeQuotaCodec } from './native-quota-codec.ts' import { createNativeRefreshCoordinator, @@ -82,6 +82,7 @@ export interface NativeLocalCredentialServiceOptions { readonly refreshToken?: NativeRefreshCoordinatorOptions['refreshToken'] /** Account lookup for new access; defaults to the Claude Code identity lookup. */ readonly resolveIdentity?: NativeRefreshCoordinatorOptions['resolveIdentity'] + readonly onLockEvent?: (event: NativeLockEvent) => void } /** @@ -101,6 +102,7 @@ const optionKeys = [ 'now', 'refreshToken', 'resolveIdentity', + 'onLockEvent', ] as const const pathKeys = [ 'legacyConfig', @@ -242,6 +244,9 @@ export function createNativeLocalCredentialService( const failurePolicy = requiredFunction< NativeLocalCredentialServiceOptions['failurePolicy'] >(descriptors, 'failurePolicy') + const onLockEvent = optionalFunction< + NonNullable + >(descriptors, 'onLockEvent') const now = optionalFunction<() => number>(descriptors, 'now') const refreshToken = optionalFunction< NativeLocalCredentialServiceOptions['refreshToken'] @@ -253,11 +258,12 @@ export function createNativeLocalCredentialService( async function publish( event: NativeRefreshObservation, operation: ( - hooks: Pick, + hooks: Pick, ) => Promise, ) { if ((await publicationGuard(event)) !== true) return const hooks = { + onLockEvent, beforePublish: async () => { if ((await publicationGuard(event)) !== true) throw new NativeRuntimeError('publication-refused') @@ -318,6 +324,7 @@ export function createNativeLocalCredentialService( const coordinator = createNativeRefreshCoordinator({ paths, quota: nativeQuotaCodec, + onLockEvent, ...(now ? { now } : {}), ...(refreshToken ? { refreshToken } : {}), ...(resolveIdentity ? { resolveIdentity } : {}), diff --git a/packages/core/src/native-lock-observer.ts b/packages/core/src/native-lock-observer.ts index 9000f461..3356f43b 100644 --- a/packages/core/src/native-lock-observer.ts +++ b/packages/core/src/native-lock-observer.ts @@ -1,6 +1,18 @@ +import { AsyncLocalStorage } from 'node:async_hooks' +import { withLock } from '@cortexkit/common-auth/fs' + +const detachedLocks = new AsyncLocalStorage() + +/** Mark work as background-origin without changing how it acquires or releases locks. */ +export function withDetachedNativeLocks(work: () => Promise): Promise { + return detachedLocks.run(true, work) +} + /** Diagnostic evidence, not an ownership assertion or a control hook. */ export type NativeLockEvent = { - type: 'acquired' | 'released' | 'contended' + type: 'acquiring' | 'acquired' | 'released' | 'contended' + /** Marks background-origin work; callers may still await the same shared refresh. */ + detached?: boolean name: string path: string } @@ -16,7 +28,9 @@ export function nativeLockObserver( if (!observer) return undefined return (event) => { try { - const result: unknown = observer(event) + const result: unknown = observer( + detachedLocks.getStore() ? { ...event, detached: true } : event, + ) if ( result !== null && (typeof result === 'object' || typeof result === 'function') && @@ -31,3 +45,45 @@ export function nativeLockObserver( } } } + +/** Observe the real lease without changing acquisition, ownership or release. */ +export async function withNativeLock( + path: string, + options: Parameters[1] & { + onLockEvent?: (event: NativeLockEvent) => void + }, + run: (lease: { assertOwned(): Promise }) => Promise, + acquire: typeof withLock = withLock, +): Promise { + const observer = nativeLockObserver(options.onLockEvent) + if (!observer) return acquire(path, options, run) + observer({ type: 'acquiring', name: options.name, path }) + let acquired = false + try { + return await acquire(path, options, async (lease) => { + acquired = true + observer({ type: 'acquired', name: options.name, path }) + return run(lease) + }) + } finally { + if (acquired) observer({ type: 'released', name: options.name, path }) + } +} + +/** Resolve the supplied acquisition function for each operation. */ +export function observeNativeLocks( + acquire: () => typeof withLock, + observer?: (event: NativeLockEvent) => void, +) { + return ( + path: string, + options: Parameters[1], + run: (lease: { assertOwned(): Promise }) => Promise, + ) => + withNativeLock( + path, + { ...options, onLockEvent: options.onLockEvent ?? observer }, + run, + acquire(), + ) +} diff --git a/packages/core/src/native-runtime.ts b/packages/core/src/native-runtime.ts index a9f53715..54bc9bb8 100644 --- a/packages/core/src/native-runtime.ts +++ b/packages/core/src/native-runtime.ts @@ -3,9 +3,8 @@ import { constants } from 'node:fs' import { lstat, mkdir, open, rename, rm } from 'node:fs/promises' import { basename, dirname, isAbsolute, join } from 'node:path' import { types as utilTypes } from 'node:util' -import { withLock } from '@cortexkit/common-auth/fs' +import { withLock as acquireLock } from '@cortexkit/common-auth/fs' import { fingerprintOf, POOL_LOCK_DEFAULTS } from '@cortexkit/common-auth/store' - import type { AccountOperationError, OAuthAccountProfile, @@ -18,6 +17,10 @@ import { type NativeLocalCredentialValidation, nativeLocalCredentialValidationMatches, } from './native-credential-validation.ts' +import { + type NativeLockEvent, + observeNativeLocks, +} from './native-lock-observer.ts' import { fromNativeQuotaMap, type NativeQuotaMap, @@ -48,6 +51,8 @@ import { } from './publication-rename.ts' import { tokenFingerprint } from './token-fingerprint.ts' +const withLock = observeNativeLocks(() => acquireLock) + /** Identifies the vault credential and Claude account that supplied an observation. This binding contains no access token. */ export interface NativeCustodyRuntimeBinding { kind: 'custody' @@ -122,6 +127,8 @@ export interface NativeRuntimeState { storageId: string accounts: Record relay?: { token: string } + /** Binds the private relay token (including its removal) to relay settings. */ + relayGeneration?: string } export type NativeRuntimeRead = @@ -129,13 +136,17 @@ export type NativeRuntimeRead = | { status: 'ready'; state: NativeRuntimeState } export class NativeRuntimeError extends Error { + readonly status?: 400 + readonly retryable?: false + constructor( public readonly code: | 'invalid-runtime' | 'unsafe-runtime' | 'runtime-io' | 'runtime-conflict' - | 'publication-refused', + | 'publication-refused' + | 'relay-inconsistent', ) { super( { @@ -146,9 +157,15 @@ export class NativeRuntimeError extends Error { 'runtime-conflict': 'Anthropic runtime observation is older than persisted state', 'publication-refused': 'Anthropic runtime publication was refused', + 'relay-inconsistent': + 'request refused: relay configuration is inconsistent; no request was sent', }[code], ) this.name = 'NativeRuntimeError' + if (code === 'relay-inconsistent') { + this.status = 400 + this.retryable = false + } } } @@ -468,11 +485,16 @@ export function decodeNativeRuntime( ): NativeRuntimeState { if ( !record(value) || - !keys(value, ['version', 'storageId', 'accounts'], ['relay']) || + !keys( + value, + ['version', 'storageId', 'accounts'], + ['relay', 'relayGeneration'], + ) || value.version !== 1 || !digest(storageId) || value.storageId !== storageId || !record(value.accounts) || + !optional(value, 'relayGeneration', text) || !optional( value, 'relay', @@ -493,6 +515,9 @@ export function decodeNativeRuntime( version: 1, storageId, accounts, + ...(typeof value.relayGeneration === 'string' + ? { relayGeneration: value.relayGeneration } + : {}), ...(record(value.relay) && typeof value.relay.token === 'string' ? { relay: { token: value.relay.token } } : {}), @@ -583,6 +608,7 @@ export async function readNativeRuntime( } export interface NativeRuntimeWriteHooks { + onLockEvent?: (event: NativeLockEvent) => void publication?: PublicationOptions /** Verify that the observed account still matches the current pool record or vault roster and that the caller still owns the locks protecting those records. */ beforeRename?: () => Promise @@ -741,6 +767,7 @@ export type NativeRuntimeOptionalPublication = | { status: 'skipped' } export interface NativeRuntimeOptionalWriteHooks { + onLockEvent?: (event: NativeLockEvent) => void publication?: PublicationOptions /** * Same checks as NativeRuntimeWriteHooks.beforeRename, run at the same @@ -767,6 +794,7 @@ export async function updateNativeRuntimeOptional( return writeNativeRuntimeOutcome(path, storageId, change, { beforeRename: hooks.beforeRename, publication: hooks.publication, + onLockEvent: hooks.onLockEvent, }) } @@ -796,6 +824,7 @@ async function writeNativeRuntime( beforePublish: hooks.beforePublish, afterRename: hooks.afterRename, publication: hooks.publication, + onLockEvent: hooks.onLockEvent, }, firstIdentity, ) @@ -822,7 +851,13 @@ async function writeNativeRuntimeOutcome( await mkdir(dirname(path), { recursive: true, mode: 0o700 }) return await withLock( path, - { name: 'native-runtime', ttlMs: 10_000, timeoutMs: 15_000, renew: true }, + { + name: 'native-runtime', + ttlMs: 10_000, + timeoutMs: 15_000, + renew: true, + onLockEvent: hooks.onLockEvent, + }, async (lock) => { const read = await readNativeRuntime(path, storageId) const current: NativeRuntimeState = @@ -1090,7 +1125,7 @@ function firstIdentityEntry( export async function publishNativeLocalFirstIdentity( paths: NativePoolPaths, observation: NativeRefreshObservation, - hooks: Pick = {}, + hooks: Pick = {}, ): Promise { try { paths = { ...paths } @@ -1149,7 +1184,7 @@ function knownValidationEntry( export async function publishNativeLocalCredentialValidation( paths: NativePoolPaths, observation: NativeRefreshObservation, - hooks: Pick = {}, + hooks: Pick = {}, ): Promise { try { paths = { ...paths } @@ -1491,7 +1526,7 @@ export async function publishNativeLocalRefreshFailure( paths: NativePoolPaths, observation: NativeRefreshObservation, policy: NativeLocalFailurePolicy, - hooks: Pick = {}, + hooks: Pick = {}, ): Promise { let reason: NativeLocalFailureSupersededReason | undefined const superseded = (value: NativeLocalFailureSupersededReason): never => { @@ -1655,10 +1690,14 @@ async function publishLocalEntry( replacement: (old: NativeRuntimeEntry | undefined) => NativeRuntimeEntry, firstIdentity?: FirstIdentityPublication, superseded?: (reason: 'row-binding' | 'row-credential') => never, - hooks: Pick = {}, + hooks: Pick = {}, ): Promise { try { - const store = createNativePoolStore({ paths, quota: nativeQuotaCodec }) + const store = createNativePoolStore({ + paths, + quota: nativeQuotaCodec, + onLockEvent: hooks.onLockEvent, + }) const [config, state] = nativePoolStoreLocks(paths) const checkPool = async () => { const read = await store.read() @@ -1679,11 +1718,19 @@ async function publishLocalEntry( } return await withLock( config.path, - { ...POOL_LOCK_DEFAULTS, name: config.name }, + { + ...POOL_LOCK_DEFAULTS, + name: config.name, + onLockEvent: hooks.onLockEvent, + }, async (configLease) => withLock( state.path, - { ...POOL_LOCK_DEFAULTS, name: state.name }, + { + ...POOL_LOCK_DEFAULTS, + name: state.name, + onLockEvent: hooks.onLockEvent, + }, async (stateLease) => { await checkPool() return writeNativeRuntime( @@ -1703,6 +1750,7 @@ async function publishLocalEntry( await stateLease.assertOwned() }, beforePublish: hooks.beforePublish, + onLockEvent: hooks.onLockEvent, }, firstIdentity, ) diff --git a/packages/core/src/native-transient-account.ts b/packages/core/src/native-transient-account.ts index 300551d4..e4e91b4c 100644 --- a/packages/core/src/native-transient-account.ts +++ b/packages/core/src/native-transient-account.ts @@ -3,6 +3,10 @@ import { type RefreshFileLock, } from '@cortexkit/common-auth/fs' import { POOL_LOCK_DEFAULTS } from '@cortexkit/common-auth/store' +import { + type NativeLockEvent, + nativeLockObserver, +} from './native-lock-observer.ts' import { type NativePoolStorePaths, @@ -90,7 +94,9 @@ export async function joinNativePoolWrite( paths: NativePoolStorePaths, signals: readonly AbortSignal[], check: () => void, + onLockEvent?: (event: NativeLockEvent) => void, ): Promise<'joined' | 'deadline'> { + const observer = nativeLockObserver(onLockEvent) // Wall-clock adjustments must not change the total wait budget. const started = performance.now() const remaining = () => @@ -108,13 +114,17 @@ export async function joinNativePoolWrite( outcome = 'deadline' break } + observer?.({ type: 'acquiring', path: lock.path, name: lock.name }) const lease = await acquireRefreshFileLock({ path: lock.path, name: lock.name, ttlMs: POOL_LOCK_DEFAULTS.ttlMs, renew: false, }) - if (lease) held.push(lease) + if (lease) { + held.push(lease) + observer?.({ type: 'acquired', path: lock.path, name: lock.name }) + } // An acquisition that finished after the budget ran out does not // count: a writer that released that late is reported as too slow. if (remaining() <= 0) { diff --git a/packages/core/src/native-vault-runtime.ts b/packages/core/src/native-vault-runtime.ts index 612772da..0b8193b6 100644 --- a/packages/core/src/native-vault-runtime.ts +++ b/packages/core/src/native-vault-runtime.ts @@ -14,7 +14,6 @@ import { SAFETY_POLL_MS, } from '@cortexkit/common-auth/claustrum' import { SubcClient } from '@cortexkit/subc-client' - import { fetchOAuthQuotaSnapshot, type OAuthAccountProfile, @@ -52,6 +51,11 @@ import { type NativeRosterWatch, } from './native-custody.ts' import { readNativeEnrollmentMode } from './native-enrollment-authority.ts' +import { + type NativeLockEvent, + nativeLockObserver, + withDetachedNativeLocks, +} from './native-lock-observer.ts' import { fromNativeQuotaMap, nativeQuotaCodec } from './native-quota-codec.ts' import { NativeQuotaObservations, @@ -105,6 +109,7 @@ export interface NativeVaultRuntimeOptions { readToken?: () => Promise<{ token: string; token_generation: number }> env?: NodeJS.ProcessEnv now?: () => number + onLockEvent?: (event: NativeLockEvent) => void /** * Timer functions for the vault custody's outage probe and roster watch * retries. Defaults to the global ones; tests pass their own. @@ -418,7 +423,18 @@ export function createNativeVaultRuntime( const connectionFile = nativeRuntimeConnectionFile(options.host, env) const identity = options.identity ? { ...options.identity } : undefined const shutdown = new AbortController() - const store = createNativePoolStore({ paths, quota: nativeQuotaCodec }) + const store = createNativePoolStore({ + paths, + quota: nativeQuotaCodec, + onLockEvent: options.onLockEvent, + }) + const observeLock = nativeLockObserver(options.onLockEvent) + const observeRosterLock = () => + observeLock?.({ + type: 'acquiring', + name: 'claustrum-roster-write', + path: paths.roster, + }) // The vault watch for the current custody, opened while a host watches the // roster. While it reports `watching`, the list timer below runs only as a // slow safety net; in any other mode it keeps the host's own pace. @@ -631,7 +647,7 @@ export function createNativeVaultRuntime( } function poll(): Promise { if (polling) return polling - polling = (async () => { + polling = withDetachedNativeLocks(async () => { const allowed = await Promise.all( [...watchers].map( async (watch) => @@ -658,7 +674,7 @@ export function createNativeVaultRuntime( }) .catch(() => {}) } - })().finally(() => { + }).finally(() => { polling = undefined }) return polling @@ -692,6 +708,7 @@ export function createNativeVaultRuntime( // Existing settings need no initialization. Avoid the shared roster write // lock so concurrent plugin starts cannot delay reads of that complete file. if (absentOnly && !(await nativeVaultFileAbsent(paths.config))) return + observeRosterLock() await mutateVaultRoster(paths.roster, async (roster, lease) => { if (absentOnly && !(await nativeVaultFileAbsent(paths.config))) return { result: undefined } @@ -723,6 +740,7 @@ export function createNativeVaultRuntime( const writer = createNativePoolStore({ paths, quota: nativeQuotaCodec, + onLockEvent: options.onLockEvent, onStep: async (point) => { if ( point === 'before-config-write' || @@ -831,6 +849,12 @@ export function createNativeVaultRuntime( check(signal) // Abort this caller's wait, not the discovery operation: another caller // may still need the result of that shared operation. + if (!refreshing) + observeLock?.({ + type: 'acquiring', + name: 'claustrum-roster', + path: paths.roster, + }) refreshing ??= refreshVaultRoster({ path: paths.roster, custody: options.discover ? { discover: options.discover } : custody, @@ -1271,6 +1295,7 @@ export function createNativeVaultRuntime( receipt.assertedAccountIdentity, ]) if ((versions.get(key) ?? 0) > receipt.recordVersion) return false + observeRosterLock() return mutateVaultRoster(paths.roster, async (roster, lease) => { const row = await rowFor(roster, route.logicalId) if ( @@ -1315,6 +1340,7 @@ export function createNativeVaultRuntime( return current }, { + onLockEvent: options.onLockEvent, beforeRename: async () => { await requireActive() const pinned = await rowFor(roster, route.logicalId) @@ -1533,6 +1559,7 @@ export function nativeVaultRuntimeKey( options.discover, options.readToken, options.now, + options.onLockEvent, options.setTimeout, options.clearTimeout, options.onRosterNoticeHandled, diff --git a/packages/core/src/sticky-routing.ts b/packages/core/src/sticky-routing.ts index b01f7014..54abe473 100644 --- a/packages/core/src/sticky-routing.ts +++ b/packages/core/src/sticky-routing.ts @@ -18,6 +18,10 @@ import { } from './accounts.ts' import { logger } from './logger.ts' import { isClaudeFableOrMythos5Model } from './models.ts' +import { + type NativeLockEvent, + nativeLockObserver, +} from './native-lock-observer.ts' import type { NativeVaultRosterDocument } from './native-roster-seed.ts' import { type NativeVaultRoutingSeed, @@ -544,7 +548,16 @@ export class StickySessionRouter { private loaded = false private state: StickyRouteState private stateMtimeNs: bigint | undefined - private touchChain: Promise = Promise.resolve() + private publicationChain: Promise = Promise.resolve() + private seedPublication: Promise | undefined + private readonly changes = new Map< + string, + { + assignment?: StickyRouteAssignment + operation: 'touch' | 'assign' | 'delete' + queued: boolean + } + >() private seedFinalized = false constructor( @@ -554,6 +567,7 @@ export class StickySessionRouter { assignmentTtlMs?: number vaultSeed?: NativeVaultRoutingSeed localRecoveryPath?: string + onLockEvent?: (event: NativeLockEvent) => void }, ) { this.state = emptyState(this.now()) @@ -601,18 +615,23 @@ export class StickySessionRouter { const after = await this.readStateMtimeNs() if (before !== after) continue this.state = state + this.overlayChanges() this.stateMtimeNs = after this.loaded = true return } this.state = await this.readState() + this.overlayChanges() this.stateMtimeNs = await this.readStateMtimeNs() this.loaded = true } - private async acquireLock(wait = true) { + private async acquireLock(wait = true, detached = false) { + const observer = nativeLockObserver(this.options.onLockEvent) + const event = { name: 'write', path: this.options.path, detached } + observer?.({ ...event, type: 'acquiring' }) await mkdir(dirname(this.options.path), { recursive: true, mode: 0o700 }) - const deadline = this.now() + STICKY_ROUTING_LOCK_WAIT_MS + const deadline = performance.now() + STICKY_ROUTING_LOCK_WAIT_MS while (true) { const lock = await acquireRefreshFileLock({ name: 'write', @@ -620,8 +639,21 @@ export class StickySessionRouter { ttlMs: STICKY_ROUTING_LOCK_TTL_MS, renew: true, }) - if (lock) return lock - if (!wait || this.now() >= deadline) { + if (lock) { + observer?.({ ...event, type: 'acquired' }) + return { + assertOwned: () => lock.assertOwned(), + release: async () => { + try { + await lock.release() + } finally { + observer?.({ ...event, type: 'released' }) + } + }, + } + } + observer?.({ ...event, type: 'contended' }) + if (!wait || performance.now() >= deadline) { throw new Error('Timed out acquiring sticky routing state lock') } await new Promise((resolve) => @@ -630,27 +662,24 @@ export class StickySessionRouter { } } - private async recoverLocalPins() { + private recoverLocalPins() { const source = this.options.localRecoveryPath if (!source || source === this.options.path) return const key = JSON.stringify([this.options.path, source]) let recovery = localRecoveries.get(key) if (!recovery) { - // Recovery is bookkeeping, not a prerequisite for using destination pins. - // Remember failures as completed attempts so later requests do no recovery I/O. - recovery = this.fillLocalPins(source).catch(() => { - logger.warn('sticky-routing', 'local pin recovery unavailable', { - reason: 'recovery-failed', - }) - }) + recovery = this.enqueuePublication('recovery', () => + this.fillLocalPins(source), + ) localRecoveries.set(key, recovery) + } else { + const shared = recovery + this.publicationChain = this.publicationChain.then(() => shared) } - await recovery } private async fillLocalPins(source: string) { - // Do not make a request wait for another writer just to recover old pins. - const lock = await this.acquireLock(false) + const lock = await this.acquireLock(true, true) try { let legacy: StickyRouteState | undefined try { @@ -719,9 +748,9 @@ export class StickySessionRouter { } finally { await rm(tempPath, { force: true }).catch(() => {}) } - this.state = state - this.stateMtimeNs = await this.readStateMtimeNs() - this.loaded = true + // A background merge must not replace newer in-memory session changes. + // Re-read the atomically published file before the next route decision. + this.loaded = false } private async writeResolutionState( @@ -731,89 +760,84 @@ export class StickySessionRouter { ) { let ownershipLost = false try { - await this.writeState(state, async () => { - try { - await assertOwned() - } catch (error) { - ownershipLost = true - throw error - } - }) - } catch (error) { - if (!ownershipLost) throw error - // Saving an assignment is bookkeeping: a displaced writer must not fail - // the already-computed route or overwrite the successor's state. + await this.writeState( + state, + async () => { + try { + await assertOwned() + } catch (error) { + ownershipLost = true + throw error + } + }, + false, + ) + return true + } catch { + // A failed publication keeps its session change pending for a later retry. logger.warn('sticky-routing', 'assignment write refused', { operation, - reason: 'ownership-lost', + reason: ownershipLost ? 'ownership-lost' : 'storage-unavailable', }) + return false } } - /** Import legacy session pins only after vault account discovery produces a complete verified roster. No empty routing file is published while that roster is incomplete. */ + /** Queue legacy-pin import in the background; take its writer lease only after verifying a complete roster. */ async prepareLegacySeed(knownRoster?: NativeVaultRosterDocument | null) { const seed = this.options.vaultSeed - if (!seed || this.seedFinalized) return - try { - await this.refreshStateIfChanged() - } catch { - // An unreadable existing file cannot prove it opted into a deferred import. - // Leave it alone; ordinary sticky routing still reports its read error. - this.seedFinalized = true - return - } - // An unmarked existing file predates this importer and belongs to the user. - if (this.state.vaultLegacySeed !== 'pending') { - this.seedFinalized = true - return - } - const roster = - knownRoster === undefined - ? await readVaultRoster(seed.roster) - : knownRoster - if (!roster?.complete) return - const lock = await this.acquireLock() - try { - const state = await this.readState() - if (state.vaultLegacySeed !== 'pending') { - this.seedFinalized = true - return - } - const source = await readNativeVaultLegacyJson(seed.legacyConfig) - const routes = nativeVaultLegacyRoutes(source, roster) - const legacy = await readNativeVaultLegacyJson( - getStickyRoutingStatePath(seed.legacyConfig), - ) - if ( - isRecord(legacy) && - legacy.version === 1 && - isRecord(legacy.assignments) - ) { - for (const [key, raw] of Object.entries(legacy.assignments)) { - // Both namespaces hash the raw session ID with SHA-256, without a salt. - if ( - !/^[0-9a-f]{64}$/.test(key) || - Object.hasOwn(state.assignments, key) - ) - continue - const assignment = normalizeAssignment(raw) - const target = assignment && routes.get(assignment.accountId) - if (assignment && target) - state.assignments[key] = { - ...assignment, - accountId: target.routeId, - accountIdentity: target.accountIdentity, - } + if (!seed || this.seedFinalized || this.seedPublication) return + if (knownRoster !== undefined && !knownRoster?.complete) return + this.seedPublication = this.enqueuePublication('seed', async () => { + const roster = + knownRoster === undefined + ? await readVaultRoster(seed.roster) + : knownRoster + if (!roster?.complete) return + const lock = await this.acquireLock(true, true) + try { + const state = await this.readState() + if (state.vaultLegacySeed !== 'pending') { + this.seedFinalized = true + return } + const source = await readNativeVaultLegacyJson(seed.legacyConfig) + const routes = nativeVaultLegacyRoutes(source, roster) + const legacy = await readNativeVaultLegacyJson( + getStickyRoutingStatePath(seed.legacyConfig), + ) + if ( + isRecord(legacy) && + legacy.version === 1 && + isRecord(legacy.assignments) + ) { + for (const [key, raw] of Object.entries(legacy.assignments)) { + if ( + !/^[0-9a-f]{64}$/.test(key) || + Object.hasOwn(state.assignments, key) + ) + continue + const assignment = normalizeAssignment(raw) + const target = assignment && routes.get(assignment.accountId) + if (assignment && target) + state.assignments[key] = { + ...assignment, + accountId: target.routeId, + accountIdentity: target.accountIdentity, + } + } + } + state.vaultLegacySeed = 'complete' + // Imported pins fill only missing assignments, leaving existing ones + // unchanged. Publish the marker and additions in one owned rename. + await this.writeState(state, () => lock.assertOwned(), false) + this.seedFinalized = true + } finally { + await lock.release() } - state.vaultLegacySeed = 'complete' - // Seed and marker share one rename. Pruning here would remove live pins - // written during deferral, so leave their lifetime to ordinary routing. - await this.writeState(state, () => lock.assertOwned(), false) - this.seedFinalized = true - } finally { - await lock.release() - } + }).finally(() => { + this.seedPublication = undefined + }) } private logicalAccountId(assignment: StickyRouteAssignment) { @@ -827,21 +851,102 @@ export class StickySessionRouter { } private scheduleTouch(key: string) { - this.touchChain = this.touchChain - .catch(() => {}) - .then(async () => { - const lock = await this.acquireLock() - try { - const state = await this.readState() - const assignment = state.assignments[key] - if (!assignment) return - assignment.lastSeenAt = this.now() - await this.writeState(state, () => lock.assertOwned()) - } finally { - await lock.release() - } + const assignment = this.state.assignments[key] + if (assignment) this.scheduleChange(key, assignment, 'touch') + } + + private overlayChanges() { + for (const [key, change] of this.changes) { + if (change.assignment) + this.state.assignments[key] = { ...change.assignment } + else delete this.state.assignments[key] + } + } + + private enqueuePublication(operation: string, work: () => Promise) { + const pending = this.publicationChain.then(work).catch(() => { + if (operation === 'recovery') { + logger.warn('sticky-routing', 'local pin recovery unavailable', { + reason: 'recovery-failed', + }) + return + } + logger.warn('sticky-routing', 'assignment write refused', { + operation, + reason: 'storage-unavailable', }) - .catch(() => {}) + }) + this.publicationChain = pending + return pending + } + + private scheduleChange( + key: string, + assignment: StickyRouteAssignment | undefined, + operation: 'touch' | 'assign' | 'delete', + ) { + const change = { + assignment: assignment ? { ...assignment } : undefined, + operation, + queued: false, + } + this.changes.set(key, change) + this.overlayChanges() + this.retryChange(key) + } + + private retryChange(key: string) { + const change = this.changes.get(key) + if (!change || change.queued) return + change.queued = true + void this.enqueuePublication(change.operation, async () => { + if (this.changes.get(key) !== change) return + const lock = await this.acquireLock(true, true) + try { + const state = await this.readState() + // A reset or reassignment may supersede a change while publication waits. + if (this.changes.get(key) !== change) return + // Re-read under the publication lease and merge only this session, so + // another project's pins cannot be overwritten by a request snapshot. + if (change.assignment) { + if (change.operation === 'touch') { + const current = state.assignments[key] + if ( + current?.accountId !== change.assignment.accountId || + current.accountIdentity !== change.assignment.accountIdentity || + current.affinityModelId !== change.assignment.affinityModelId || + current.family !== change.assignment.family + ) { + if (this.changes.get(key) === change) this.changes.delete(key) + return + } + current.lastSeenAt = Math.max( + current.lastSeenAt, + change.assignment.lastSeenAt, + ) + } else state.assignments[key] = { ...change.assignment } + } else delete state.assignments[key] + const saved = await this.writeResolutionState( + state, + () => lock.assertOwned(), + change.operation, + ) + if (saved && this.changes.get(key) === change) this.changes.delete(key) + } finally { + await lock.release().catch(() => {}) + } + }).finally(() => { + change.queued = false + }) + } + + /** Wait for owned background publications during host disposal or fixture cleanup. */ + async drain() { + let pending: Promise + do { + pending = this.publicationChain + await pending + } while (pending !== this.publicationChain) } private selectCandidate(input: { @@ -927,9 +1032,9 @@ export class StickySessionRouter { excludeAccountIds?: ReadonlySet }): Promise { if (!input.sessionId) return null - await this.recoverLocalPins() + this.recoverLocalPins() await this.prepareLegacySeed() - await this.refreshStateIfChanged() + await this.refreshStateIfChanged().catch(() => {}) const key = sessionKey(input.sessionId) const affinityModelId = input.affinityModelId ?? input.modelId const matchesAffinity = (assignment: StickyRouteAssignment) => @@ -962,6 +1067,7 @@ export class StickySessionRouter { cached.lastSeenAt = this.now() this.scheduleTouch(key) } + this.retryChange(key) return { accountId: this.logicalAccountId(cached), assignment: { ...cached }, @@ -970,113 +1076,93 @@ export class StickySessionRouter { } } - const lock = await this.acquireLock() - try { - const state = await this.readState() - this.prune(state) - const current = state.assignments[key] - if ( - current && - matchesAffinity(current) && - matchesAccountIdentity(current) && - input.retainAccountIds.has(this.logicalAccountId(current)) && - !input.excludeAccountIds?.has(this.logicalAccountId(current)) - ) { - current.lastSeenAt = this.now() - await this.writeResolutionState( - state, - () => lock.assertOwned(), - 'touch', - ) - return { - accountId: this.logicalAccountId(current), - assignment: { ...current }, - created: false, - migrated: false, - } - } - const affinityChanged = Boolean(current && !matchesAffinity(current)) - const identityChanged = Boolean( - current && !matchesAccountIdentity(current), - ) - if (affinityChanged || identityChanged) { - delete state.assignments[key] + // Route decisions use the atomically published state, never a writer lease. + const state = structuredClone(this.state) + const save = (operation: 'touch' | 'delete' | 'assign') => { + this.scheduleChange(key, state.assignments[key], operation) + } + this.prune(state) + const current = state.assignments[key] + if ( + current && + matchesAffinity(current) && + matchesAccountIdentity(current) && + input.retainAccountIds.has(this.logicalAccountId(current)) && + !input.excludeAccountIds?.has(this.logicalAccountId(current)) + ) { + current.lastSeenAt = this.now() + save('touch') + return { + accountId: this.logicalAccountId(current), + assignment: { ...current }, + created: false, + migrated: false, } - const candidates = input.candidates.filter( - (candidate) => - !input.excludeAccountIds?.has(candidate.accountId) && - (!input.accountIdentities || - identityForAccount(candidate.accountId) !== undefined), - ) - const preferred = - !affinityChanged && !identityChanged && input.preferredAccountId - ? candidates.find( - (candidate) => - candidate.accountId === input.preferredAccountId && - stickyRouteCandidateWeight({ - candidate, - family: input.family, - modelId: input.modelId, - storage: input.storage, - now: this.now(), - }) > 0, - ) - : undefined - const selected = - preferred ?? this.selectCandidate({ ...input, candidates, state }) - if (!selected) { - if (current) { - delete state.assignments[key] - await this.writeResolutionState( - state, - () => lock.assertOwned(), - 'delete', + } + const affinityChanged = Boolean(current && !matchesAffinity(current)) + const identityChanged = Boolean(current && !matchesAccountIdentity(current)) + if (affinityChanged || identityChanged) { + delete state.assignments[key] + } + const candidates = input.candidates.filter( + (candidate) => + !input.excludeAccountIds?.has(candidate.accountId) && + (!input.accountIdentities || + identityForAccount(candidate.accountId) !== undefined), + ) + const preferred = + !affinityChanged && !identityChanged && input.preferredAccountId + ? candidates.find( + (candidate) => + candidate.accountId === input.preferredAccountId && + stickyRouteCandidateWeight({ + candidate, + family: input.family, + modelId: input.modelId, + storage: input.storage, + now: this.now(), + }) > 0, ) - } - return null - } - const now = this.now() - const selectedQuota = knownQuota(selected.quota) - const assignment: StickyRouteAssignment = { - accountId: selected.accountId, - ...(input.accountIdentities - ? { accountIdentity: identityForAccount(selected.accountId) } - : {}), - family: input.family, - affinityModelId, - assignedAt: now, - lastSeenAt: now, - initialInputBytes: Math.max(1, input.inputBytes), - quotaCheckedAt: selectedQuota ? snapshotCheckedAt(selectedQuota) : 0, - } - state.assignments[key] = assignment - await this.writeResolutionState(state, () => lock.assertOwned(), 'assign') - return { - accountId: selected.accountId, - assignment: { ...assignment }, - created: !current, - migrated: Boolean( - current && - (identityChanged || current.accountId !== selected.accountId), - ), + : undefined + const selected = + preferred ?? this.selectCandidate({ ...input, candidates, state }) + if (!selected) { + if (current) { + delete state.assignments[key] + save('delete') } - } finally { - await lock.release() + return null + } + const now = this.now() + const selectedQuota = knownQuota(selected.quota) + const assignment: StickyRouteAssignment = { + accountId: selected.accountId, + ...(input.accountIdentities + ? { accountIdentity: identityForAccount(selected.accountId) } + : {}), + family: input.family, + affinityModelId, + assignedAt: now, + lastSeenAt: now, + initialInputBytes: Math.max(1, input.inputBytes), + quotaCheckedAt: selectedQuota ? snapshotCheckedAt(selectedQuota) : 0, + } + state.assignments[key] = assignment + save('assign') + return { + accountId: selected.accountId, + assignment: { ...assignment }, + created: !current, + migrated: Boolean( + current && + (identityChanged || current.accountId !== selected.accountId), + ), } } async clear(sessionId: string) { if (!sessionId) return - await this.recoverLocalPins() - const key = sessionKey(sessionId) - const lock = await this.acquireLock() - try { - const state = await this.readState() - if (!(key in state.assignments)) return - delete state.assignments[key] - await this.writeState(state, () => lock.assertOwned()) - } finally { - await lock.release() - } + this.recoverLocalPins() + this.scheduleChange(sessionKey(sessionId), undefined, 'delete') } } diff --git a/packages/core/src/tests/locked-publication.test.ts b/packages/core/src/tests/locked-publication.test.ts index c088a166..7560f16b 100644 --- a/packages/core/src/tests/locked-publication.test.ts +++ b/packages/core/src/tests/locked-publication.test.ts @@ -67,23 +67,23 @@ test('sticky queued touch refuses a replaced lease and preserves successor assig inputBytes: 1, } await router.resolve(input) + await router.drain() + lifetime.deferCleanup(() => router.drain()) const entered = lifetime.gate() const resume = lifetime.gate() - // Pause the read after it captures routing state while the lock is held. + // Pause staging while the background publication owns the lease. const seam = router as unknown as { - readState(): Promise - touchChain: Promise + writeState(...args: unknown[]): Promise } - const readState = seam.readState.bind(router) - seam.readState = async () => { - const state = await readState() + const writeState = seam.writeState.bind(router) + seam.writeState = async (...args) => { entered.open() await resume.wait - return state + return writeState(...args) } now += 60 * 60_000 await router.resolve(input) - lifetime.trackDetached(seam.touchChain) + lifetime.trackDetached(router.drain()) await entered.wait const successor = await replaceLease(path, 'write') const state = JSON.parse(await readFile(path, 'utf8')) @@ -99,7 +99,7 @@ test('sticky queued touch refuses a replaced lease and preserves successor assig await successor.assertOwned() await writeFile(path, published, { mode: 0o600 }) resume.open() - await seam.touchChain + await router.drain() expect(await readFile(path, 'utf8')).toBe(published) await successor.assertOwned() }) @@ -120,6 +120,7 @@ test('sticky writes publish normally while their lease remains owned', async () }), ).toMatchObject({ accountId: 'account', created: true }) await router.clear('session') + await router.drain() expect(JSON.parse(await readFile(path, 'utf8')).assignments).toEqual({}) }) @@ -153,14 +154,14 @@ for (const operation of ['touch', 'delete', 'assign'] as const) { accountId: 'sensitive-account', family: operation === 'delete' ? 'fable' : 'general', assignedAt: now, - lastSeenAt: now, + lastSeenAt: operation === 'touch' ? now - 61 * 60_000 : now, initialInputBytes: 1, quotaCheckedAt: 0, }, }, }) async function routerAt(path: string) { - if (operation === 'delete') await writeFile(path, seed, { mode: 0o600 }) + if (operation !== 'assign') await writeFile(path, seed, { mode: 0o600 }) const router = new StickySessionRouter({ path, now: () => now }) const seam = router as unknown as { refreshStateIfChanged(): Promise @@ -170,19 +171,12 @@ for (const operation of ['touch', 'delete', 'assign'] as const) { prune?: boolean, ): Promise } - if (operation === 'touch') { - const refreshState = seam.refreshStateIfChanged.bind(router) - seam.refreshStateIfChanged = async () => { - await refreshState() - // An assignment can arrive after the initial cache read but before - // the locked read. That makes resolve use its locked touch branch. - await writeFile(path, seed, { mode: 0o600 }) - } - } + return { router, seam } } const healthy = await routerAt(join(directory, 'healthy.json')) const expected = await healthy.router.resolve(input) + await healthy.router.drain() if (operation === 'delete') expect(expected).toBeNull() else expect(expected).toMatchObject({ @@ -225,6 +219,7 @@ for (const operation of ['touch', 'delete', 'assign'] as const) { await writeFile(path, published, { mode: 0o600 }) resume.open() expect(await running).toEqual(expected) + await router.drain() expect(await readFile(path, 'utf8')).toBe(published) expect(logs).toEqual([ { diff --git a/packages/core/src/tests/native-account-runtime.test.ts b/packages/core/src/tests/native-account-runtime.test.ts index fa615445..43d046b1 100644 --- a/packages/core/src/tests/native-account-runtime.test.ts +++ b/packages/core/src/tests/native-account-runtime.test.ts @@ -3,7 +3,6 @@ import * as filesystem from 'node:fs/promises' import { mkdir, mkdtemp, readFile, rm, writeFile } from 'node:fs/promises' import { join } from 'node:path' import { disconnectClaustrumEnrollment } from '@cortexkit/common-auth/claustrum' - import { ClaudeOAuthRefreshError } from '../auth.ts' import { type ProviderAccountUuid, @@ -38,6 +37,7 @@ import { resolveNativePoolPaths } from '../pool-paths.ts' import { createNativePoolStore } from '../pool-store.ts' import { effectiveRelayStatus, effectiveRelayTransport } from '../relay.ts' import { initializeNativeTestAuthority } from './native-authority-fixture.ts' +import { holdNativeRelayLock } from './native-relay-lock-fixture.ts' import { createTestLifetimeSuite } from './test-lifetime.ts' const { test, deferCleanup, gate, trackDetached } = createTestLifetimeSuite() @@ -502,6 +502,184 @@ test('local quota retry admits only after deferred body cleanup so disable preve expect(f.counts.refresh).toBe(0) }) +test('relay config reads published state while another writer holds native-relay', async () => { + const f = await fixture() + await f.runtime.updateRelay({ + enabled: true, + url: 'https://held.example.test', + token: 'synthetic-held-token', + }) + const held = await holdNativeRelayLock(f.paths.config, { deferCleanup }) + try { + const read = await f.runtime.getRelayConfig().catch(() => null) + expect(read?.url).toBe('https://held.example.test') + expect(read?.token).toBe('synthetic-held-token') + await held.assertOwned() + } finally { + await held.pastTimeout + await held.release() + } +}, 25_000) + +test('relay writer serializes behind a held native-relay writer lease', async () => { + const f = await fixture() + const held = await holdNativeRelayLock(f.paths.config, { deferCleanup }) + const entered = gate() + const original = runtimeWriter.readNativeRuntime + const hook = spyOn(runtimeWriter, 'readNativeRuntime').mockImplementation( + async (...args) => { + entered.open() + return original(...args) + }, + ) + deferCleanup(() => hook.mockRestore()) + let settled = false + const pending = f.runtime + .updateRelay({ + enabled: true, + url: 'https://writer.example.test', + token: 'synthetic-writer-token', + }) + .then(() => { + settled = true + }) + trackDetached(pending) + await new Promise((resolve) => setTimeout(resolve, 50)) + expect(settled).toBe(false) + expect(await readFile(f.paths.config, 'utf8')).not.toContain( + 'writer.example.test', + ) + await held.assertOwned() + await held.release() + await entered.wait + await pending + expect(await f.runtime.getRelayConfig()).toMatchObject({ + url: 'https://writer.example.test', + token: 'synthetic-writer-token', + }) +}) + +for (const legacy of [false, true]) { + test(`relay reader racing a writer never combines old URL and new token (${legacy ? 'legacy' : 'stamped'})`, async () => { + const f = await fixture() + await f.runtime.updateRelay({ + enabled: true, + url: 'https://old.example.test', + token: 'synthetic-old-token', + }) + if (legacy) { + await f.store.updateSettings((settings) => { + delete (settings.relay as Record).generation + }) + await runtimeWriter.updateNativeRuntime( + f.paths.runtime, + f.paths.storageId, + (state) => { + delete state.relayGeneration + return state + }, + ) + } + const original = runtimeWriter.readNativeRuntime + let raced = false + const hook = spyOn(runtimeWriter, 'readNativeRuntime').mockImplementation( + async (...args) => { + if (!raced) { + raced = true + await f.runtime.updateRelay({ + url: 'https://new.example.test', + token: 'synthetic-new-token', + }) + } + return original(...args) + }, + ) + deferCleanup(() => hook.mockRestore()) + expect(await f.runtime.getRelayConfig()).toMatchObject({ + url: 'https://new.example.test', + token: 'synthetic-new-token', + }) + expect(raced).toBe(true) + }) +} + +for (const mismatch of [ + 'different', + 'settings-only', + 'runtime-only', +] as const) { + test(`relay generations that never agree refuse locally (${mismatch})`, async () => { + const f = await fixture() + await f.runtime.updateRelay({ + enabled: true, + url: 'https://private.example.test', + token: 'synthetic-never-send-token', + }) + await f.store.updateSettings((settings) => { + const relay = settings.relay as Record + if (mismatch === 'runtime-only') delete relay.generation + else relay.generation = 'different-generation' + }) + if (mismatch === 'settings-only') + await runtimeWriter.updateNativeRuntime( + f.paths.runtime, + f.paths.storageId, + (state) => { + delete state.relayGeneration + return state + }, + ) + await expect(f.runtime.getRelayConfig()).rejects.toMatchObject({ + code: 'relay-inconsistent', + message: + 'request refused: relay configuration is inconsistent; no request was sent', + }) + }) +} + +test('legacy relay pairs remain readable and every writer update stamps both files', async () => { + const f = await fixture() + await f.runtime.updateRelay({ + enabled: true, + url: 'https://legacy.example.test', + token: 'synthetic-legacy-relay-token', + }) + await f.store.updateSettings((settings) => { + delete (settings.relay as Record).generation + }) + await runtimeWriter.updateNativeRuntime( + f.paths.runtime, + f.paths.storageId, + (state) => { + delete state.relayGeneration + return state + }, + ) + expect(await f.runtime.getRelayConfig()).toMatchObject({ + url: 'https://legacy.example.test', + token: 'synthetic-legacy-relay-token', + }) + let previous: string | undefined + for (const patch of [ + { fallbackToDirect: false }, + { url: 'https://new.example.test' }, + { token: null }, + ]) { + await f.runtime.updateRelay(patch) + const settings = await f.store.readSettings() + const runtime = await readNativeRuntime(f.paths.runtime, f.paths.storageId) + if (settings.status !== 'ready' || runtime.status !== 'ready') + throw new Error('Fixture not readable') + const generation = (settings.settings.relay as Record) + .generation + expect(typeof generation).toBe('string') + expect(generation).not.toBe(previous) + expect(generation).toBe(runtime.state.relayGeneration) + previous = runtime.state.relayGeneration + } + expect(await f.runtime.getRelayConfig()).toBeNull() +}) + test('relay transport reads fresh private token while native public snapshots and config stay secret-free', async () => { const f = await fixture() expect(await f.runtime.getRelayConfig()).toBeNull() diff --git a/packages/core/src/tests/native-relay-lock-fixture.ts b/packages/core/src/tests/native-relay-lock-fixture.ts new file mode 100644 index 00000000..9051c175 --- /dev/null +++ b/packages/core/src/tests/native-relay-lock-fixture.ts @@ -0,0 +1,31 @@ +import { acquireRefreshFileLock } from '@cortexkit/common-auth/fs' +import { POOL_LOCK_DEFAULTS } from '@cortexkit/common-auth/store' + +/** Hold a native-relay writer lease beyond the pool-lock wait budget to exercise requests while it remains held. */ +export async function holdNativeRelayLock( + path: string, + owner: { deferCleanup(cleanup: () => void | Promise): void }, +) { + const lease = await acquireRefreshFileLock({ + path, + name: 'native-relay', + ttlMs: 60_000, + renew: true, + }) + if (!lease) throw new Error('Relay fixture could not acquire writer lease') + let finish!: () => void + const pastTimeout = new Promise((resolve) => { + finish = resolve + }) + const timer = setTimeout(finish, POOL_LOCK_DEFAULTS.timeoutMs + 100) + owner.deferCleanup(async () => { + clearTimeout(timer) + finish() + await lease.release() + }) + return { + pastTimeout, + assertOwned: () => lease.assertOwned(), + release: () => lease.release(), + } +} diff --git a/packages/core/src/tests/native-request-lock-fixture.ts b/packages/core/src/tests/native-request-lock-fixture.ts new file mode 100644 index 00000000..960dc7bc --- /dev/null +++ b/packages/core/src/tests/native-request-lock-fixture.ts @@ -0,0 +1,35 @@ +import type { NativeLockEvent } from '../native-lock-observer.ts' + +/** Collect acquisition attempts and acquired events, grouped as foreground, background-origin or post-dispatch. */ +export function requestLockProbe() { + let active = false + let dispatched = false + const foreground: NativeLockEvent[] = [] + const detached: NativeLockEvent[] = [] + const afterDispatch: NativeLockEvent[] = [] + return { + foreground, + detached, + afterDispatch, + onLockEvent(event: NativeLockEvent) { + if (!active) return + if ( + event.type !== 'acquired' && + event.type !== 'acquiring' && + event.type !== 'contended' + ) + return + if (event.detached) detached.push(event) + else if (dispatched) afterDispatch.push(event) + else foreground.push(event) + }, + begin() { + active = true + dispatched = false + foreground.length = detached.length = afterDispatch.length = 0 + }, + dispatch() { + dispatched = true + }, + } +} diff --git a/packages/core/src/tests/native-routing-migration.test.ts b/packages/core/src/tests/native-routing-migration.test.ts index e6bf45af..d0b09d29 100644 --- a/packages/core/src/tests/native-routing-migration.test.ts +++ b/packages/core/src/tests/native-routing-migration.test.ts @@ -225,7 +225,7 @@ test('preserved quota timestamps causally retain only matching deficits in the r const resolve = async (checkedAt: number) => { await writeFile(path, migrateStickyRoutingState(input, proven)) const router = new StickySessionRouter({ path, now: () => 100 }) - return router.resolve({ + const result = await router.resolve({ sessionId: 'fresh', family: 'general', inputBytes: 1, @@ -236,6 +236,8 @@ test('preserved quota timestamps causally retain only matching deficits in the r quota: { ...candidate.quota, checkedAt }, })), }) + await router.drain() + return result } expect((await resolve(37))?.accountId).toBe('fallback-row') // Moving the snapshot clock beyond the preserved clock removes the charge; diff --git a/packages/core/src/tests/native-vault-affinity.test.ts b/packages/core/src/tests/native-vault-affinity.test.ts index 6c1c240d..cf9faa9f 100644 --- a/packages/core/src/tests/native-vault-affinity.test.ts +++ b/packages/core/src/tests/native-vault-affinity.test.ts @@ -19,7 +19,10 @@ const model = 'claude-sonnet-5-5' async function fixture() { const root = await mkdtemp(join(tmpdir(), 'native-vault-affinity-')) - deferCleanup(() => rm(root, { recursive: true, force: true })) + deferCleanup(async () => { + await router.drain() + await rm(root, { recursive: true, force: true }) + }) const path = join(root, 'routing.json') let now = 2_000_000_000_000 const router = new StickySessionRouter({ path, now: () => now }) @@ -44,13 +47,13 @@ async function fixture() { quota: quota(identity), order: id === 'main' ? 0 : 1, }) - const resolve = ( + const resolveAndDrain = async ( sessionId: string, candidates: StickyRouteCandidate[], identities?: ReadonlyMap, preferredAccountId?: string, - ) => - router.resolve({ + ) => { + const result = await router.resolve({ sessionId, family: 'general', modelId: model, @@ -62,10 +65,13 @@ async function fixture() { inputBytes: 77, preferredAccountId, }) + await router.drain() + return result + } return { path, router, - resolve, + resolveAndDrain, candidate, advance: () => { now += 1 @@ -90,12 +96,17 @@ test('replacing the vault default UUID re-places only its session while an unrel ['fallback', fallbackUuid], ]) expect( - (await f.resolve('main-session', [f.candidate('main', originalUuid)], old)) - ?.accountId, + ( + await f.resolveAndDrain( + 'main-session', + [f.candidate('main', originalUuid)], + old, + ) + )?.accountId, ).toBe('main') expect( ( - await f.resolve( + await f.resolveAndDrain( 'fallback-session', [f.candidate('fallback', fallbackUuid)], old, @@ -108,7 +119,7 @@ test('replacing the vault default UUID re-places only its session while an unrel ['main', replacedUuid], ['fallback', fallbackUuid], ]) - const result = await f.resolve( + const result = await f.resolveAndDrain( 'main-session', [ { accountId: 'main', quota: { kind: 'unknown' }, order: 0 }, @@ -127,14 +138,14 @@ test('replacing the vault default UUID re-places only its session while an unrel test('a replaced vault default using the same logical main id gets a new UUID-bound pin and no old preference or timestamps', async () => { const f = await fixture() - const first = await f.resolve( + const first = await f.resolveAndDrain( 'session', [f.candidate('main', originalUuid)], new Map([['main', originalUuid]]), ) expect(first?.assignment.accountIdentity).toBe(originalUuid) f.advance() - const next = await f.resolve( + const next = await f.resolveAndDrain( 'session', [f.candidate('main', replacedUuid)], new Map([['main', replacedUuid]]), @@ -148,8 +159,13 @@ test('a vault route without a verified UUID is never pinned, including retained const f = await fixture() const candidates = [f.candidate('main', originalUuid)] expect( - (await f.resolve('session', candidates, new Map([['main', originalUuid]]))) - ?.accountId, + ( + await f.resolveAndDrain( + 'session', + candidates, + new Map([['main', originalUuid]]), + ) + )?.accountId, ).toBe('main') for (const identities of [ new Map(), @@ -157,7 +173,7 @@ test('a vault route without a verified UUID is never pinned, including retained new Map([['main', ' ']]), ]) { expect( - await f.resolve('session', candidates, identities, 'main'), + await f.resolveAndDrain('session', candidates, identities, 'main'), ).toBeNull() } }) @@ -166,7 +182,7 @@ test('vault UUID changes refuse old retention even while the quota candidate lis const f = await fixture() expect( ( - await f.resolve( + await f.resolveAndDrain( 'session', [f.candidate('main', originalUuid)], new Map([['main', originalUuid]]), @@ -174,18 +190,20 @@ test('vault UUID changes refuse old retention even while the quota candidate lis )?.accountId, ).toBe('main') expect( - await f.resolve('session', [], new Map([['main', replacedUuid]])), + await f.resolveAndDrain('session', [], new Map([['main', replacedUuid]])), ).toBeNull() }) test('local retention ignores account UUID changes and leaves existing routing bytes untouched', async () => { const f = await fixture() expect( - (await f.resolve('session', [f.candidate('main', originalUuid)])) + (await f.resolveAndDrain('session', [f.candidate('main', originalUuid)])) ?.accountId, ).toBe('main') const before = await readFile(f.path) - const result = await f.resolve('session', [f.candidate('main', replacedUuid)]) + const result = await f.resolveAndDrain('session', [ + f.candidate('main', replacedUuid), + ]) expect(result?.accountId).toBe('main') expect(result?.assignment).not.toHaveProperty('accountIdentity') expect(await readFile(f.path)).toEqual(before) @@ -193,14 +211,14 @@ test('local retention ignores account UUID changes and leaves existing routing b test('migration preserves optional vault UUID metadata without adding any identity to legacy local assignments', async () => { const f = await fixture() - await f.resolve('local', [f.candidate('main', originalUuid)]) + await f.resolveAndDrain('local', [f.candidate('main', originalUuid)]) const localSource = JSON.parse(await readFile(f.path, 'utf8')) const local = migrateStickyRoutingState( localSource, new Map([['main', 'imported-main']]), ) expect(local).not.toContain('accountIdentity') - await f.resolve( + await f.resolveAndDrain( 'vault', [f.candidate('fallback', fallbackUuid)], new Map([['fallback', fallbackUuid]]), @@ -219,8 +237,8 @@ test('migration preserves optional vault UUID metadata without adding any identi test('a legacy vault pin without UUID is re-placed instead of inheriting its logical main id', async () => { const f = await fixture() - await f.resolve('session', [f.candidate('main', originalUuid)]) - const next = await f.resolve( + await f.resolveAndDrain('session', [f.candidate('main', originalUuid)]) + const next = await f.resolveAndDrain( 'session', [ { accountId: 'main', quota: { kind: 'unknown' }, order: 0 }, @@ -238,7 +256,7 @@ test('a legacy vault pin without UUID is re-placed instead of inheriting its log test('a replaced vault UUID cannot inherit a warm-cache preference for logical main', async () => { const f = await fixture() - await f.resolve( + await f.resolveAndDrain( 'session', [f.candidate('main', originalUuid)], new Map([['main', originalUuid]]), @@ -257,7 +275,7 @@ test('a replaced vault UUID cannot inherit a warm-cache preference for logical m remainingPercent: 1, resetsAt: new Date(f.now() + 86_400_000).toISOString(), } - const next = await f.resolve( + const next = await f.resolveAndDrain( 'session', [main, f.candidate('fallback', fallbackUuid)], new Map([ diff --git a/packages/core/src/tests/native-vault-legacy-seed.test.ts b/packages/core/src/tests/native-vault-legacy-seed.test.ts index de8a553c..c9ad63a4 100644 --- a/packages/core/src/tests/native-vault-legacy-seed.test.ts +++ b/packages/core/src/tests/native-vault-legacy-seed.test.ts @@ -66,7 +66,13 @@ const preferences = { async function fixture(verified = true) { const root = await mkdtemp(join(tmpdir(), 'vault-legacy-seed-')) - deferCleanup(() => rm(root, { recursive: true, force: true })) + const routers: StickySessionRouter[] = [] + let accountRuntime: ReturnType | undefined + deferCleanup(async () => { + await accountRuntime?.drainRoutingPublications() + await Promise.all(routers.map((router) => router.drain())) + await rm(root, { recursive: true, force: true }) + }) const paths = await resolveNativePoolPaths(join(root, 'anthropic-auth.json')) const vault = nativeVaultPaths(paths) const now = Date.now() @@ -184,9 +190,15 @@ async function fixture(verified = true) { if (!mainRoute || !workRoute) throw new Error('Fixture roster has no routes') const options = { paths, host: 'opencode' as const, vault: { env } } const runtime = ownRuntime(createNativeAccountRuntime(options)) + accountRuntime = runtime + const readAndDrain = async () => { + const snapshot = await runtime.read() + await runtime.drainRoutingPublications() + return snapshot + } const store = createNativePoolStore({ paths: vault, quota: nativeQuotaCodec }) - const router = () => - new StickySessionRouter({ + const router = () => { + const value = new StickySessionRouter({ path: vault.routing, vaultSeed: { legacyConfig: paths.legacyConfig, @@ -194,15 +206,18 @@ async function fixture(verified = true) { mainRouteId: mainRoute, }, }) - const resolve = ( + routers.push(value) + return value + } + const resolveAndDrain = async ( session: string, routerInstance = router(), identities = new Map([ ['main', mainUuid], [workRoute, workUuid], ]), - ) => - routerInstance.resolve({ + ) => { + const resolution = await routerInstance.resolve({ sessionId: session, family: 'general', modelId: model, @@ -214,6 +229,9 @@ async function fixture(verified = true) { storage: null, inputBytes: 1, }) + await routerInstance.drain() + return resolution + } const legacyBefore = await Promise.all( [paths.legacyConfig, paths.legacyState, legacyRouting].map((path) => readFile(path), @@ -238,7 +256,8 @@ async function fixture(verified = true) { mainRoute, workRoute, router, - resolve, + resolveAndDrain, + readAndDrain, assertLegacy, config, legacyRouting, @@ -248,7 +267,7 @@ async function fixture(verified = true) { test('first vault snapshot seeds validated preferences and UUID-translated pins without touching legacy bytes', async () => { const f = await fixture() - const snapshot = await f.runtime.read() + const snapshot = await f.readAndDrain() for (const [section, value] of Object.entries(preferences)) expect(snapshot.settings[section]).toEqual(value) const settings = JSON.parse(await readFile(f.vault.config, 'utf8')) @@ -273,8 +292,8 @@ test('first vault snapshot seeds validated preferences and UUID-translated pins accountIdentity: workUuid, assignedAt: f.now - 10, }) - expect((await f.resolve('work-session'))?.created).toBe(false) - expect((await f.resolve('main-session'))?.accountId).toBe('main') + expect((await f.resolveAndDrain('work-session'))?.created).toBe(false) + expect((await f.resolveAndDrain('main-session'))?.accountId).toBe('main') for (const path of [f.vault.config, f.vault.routing]) expect((await stat(path)).mode & 0o777).toBe(0o600) await f.assertLegacy() @@ -297,7 +316,7 @@ test('existing pin-only Mac settings and unmarked live routing remain byte-ident const before = await Promise.all( [f.vault.config, f.vault.routing].map((path) => readFile(path)), ) - await f.runtime.read() + await f.readAndDrain() expect( await Promise.all( [f.vault.config, f.vault.routing].map((path) => readFile(path)), @@ -321,7 +340,7 @@ test('existing pin-only Mac settings and unmarked live routing remain byte-ident routing: { mode: 'fallback-first' }, logging: { level: 'error' }, }) - await f.resolve('live-session') + await f.resolveAndDrain('live-session') const live = await Promise.all( [f.vault.config, f.vault.routing].map((path) => readFile(path)), ) @@ -338,13 +357,17 @@ test('existing pin-only Mac settings and unmarked live routing remain byte-ident test('request during routing deferral routes normally without creating an empty file or marker', async () => { const f = await fixture(false) - expect((await f.runtime.read()).settings.routing).toEqual(preferences.routing) + expect((await f.readAndDrain()).settings.routing).toEqual(preferences.routing) await expect(stat(f.vault.routing)).rejects.toMatchObject({ code: 'ENOENT' }) - const result = await f.resolve('no-verified-route', f.router(), new Map()) + const result = await f.resolveAndDrain( + 'no-verified-route', + f.router(), + new Map(), + ) expect(result).toBeNull() await expect(stat(f.vault.routing)).rejects.toMatchObject({ code: 'ENOENT' }) await f.publishRoster() - await f.runtime.read() + await f.readAndDrain() expect( JSON.parse(await readFile(f.vault.routing, 'utf8')).vaultLegacySeed, ).toBe('complete') @@ -356,14 +379,14 @@ test('a live pin during deferral survives the one-time seed and later starts nev next: { ...f.roster, complete: false }, result: undefined, })) - await f.runtime.read() - expect((await f.resolve('main-session'))?.accountId).toBe(f.workRoute) + await f.readAndDrain() + expect((await f.resolveAndDrain('main-session'))?.accountId).toBe(f.workRoute) const pending = JSON.parse(await readFile(f.vault.routing, 'utf8')) expect(pending.vaultLegacySeed).toBe('pending') const livePin = pending.assignments[key('main-session')] expect(livePin.accountId).toBe(f.workRoute) await f.publishRoster() - await f.runtime.read() + await f.readAndDrain() const seeded = JSON.parse(await readFile(f.vault.routing, 'utf8')) expect(seeded.vaultLegacySeed).toBe('complete') expect(seeded.assignments[key('main-session')]).toEqual(livePin) @@ -372,7 +395,9 @@ test('a live pin during deferral survives the one-time seed and later starts nev ...s, routing: { mode: 'fallback-first' }, })) - await f.router().clear('work-session') + const clearing = f.router() + await clearing.clear('work-session') + await clearing.drain() const before = await Promise.all( [f.vault.config, f.vault.routing].map((path) => readFile(path)), ) @@ -451,7 +476,7 @@ test('invalid settings fields are dropped while valid siblings survive and no cr }), { mode: 0o600 }, ) - const settings = (await f.runtime.read()).settings + const settings = (await f.readAndDrain()).settings expect(settings).not.toHaveProperty('routing') expect(settings.claudeCache).toEqual({ enabled: true }) expect(settings.quota).toEqual({ @@ -474,7 +499,7 @@ test('missing unreadable or malformed legacy input never refuses vault serving', if (bad === 'directory') await fs.mkdir(f.paths.legacyConfig) if (bad === 'malformed') await writeFile(f.paths.legacyConfig, '{invalid', { mode: 0o600 }) - expect((await f.runtime.read()).mode).toBe('claustrum') + expect((await f.readAndDrain()).mode).toBe('claustrum') await f.store.updateSettings((s) => ({ ...s, logging: { level: 'error' } })) expect((await f.store.readSettings()).status).toBe('ready') } @@ -501,7 +526,7 @@ test('settings seed interruption before rename leaves no partial file and a retr await f.assertLegacy() }) -test('routing seed rename failure leaves no partial file or marker and retry seeds completely', async () => { +test('routing seed rename failure is nonfatal, leaves no partial file or marker, and retry seeds completely', async () => { const f = await fixture() const originalRename = fs.rename const rename = spyOn(fs, 'rename').mockImplementation(async (from, to) => { @@ -509,10 +534,10 @@ test('routing seed rename failure leaves no partial file or marker and retry see return originalRename(from, to) }) deferCleanup(() => rename.mockRestore()) - await expect(f.runtime.read()).rejects.toThrow('synthetic crash at rename') + expect((await f.readAndDrain()).mode).toBe('claustrum') await expect(stat(f.vault.routing)).rejects.toMatchObject({ code: 'ENOENT' }) rename.mockRestore() - await f.runtime.read() + await f.readAndDrain() expect( JSON.parse(await readFile(f.vault.routing, 'utf8')).vaultLegacySeed, ).toBe('complete') @@ -521,15 +546,15 @@ test('routing seed rename failure leaves no partial file or marker and retry see test('completed routing seed adds no legacy or routing file reads to later Core snapshots', async () => { const f = await fixture() - await f.runtime.read() + await f.readAndDrain() const read = spyOn(fs, 'readFile') const metadata = spyOn(fs, 'stat') deferCleanup(() => { read.mockRestore() metadata.mockRestore() }) - await f.runtime.read() - await f.runtime.read() + await f.readAndDrain() + await f.readAndDrain() expect( metadata.mock.calls.filter(([path]) => path === f.vault.routing), ).toEqual([]) @@ -545,13 +570,13 @@ test('completed routing seed adds no legacy or routing file reads to later Core expect(reads.filter((path) => path === f.vault.roster)).toHaveLength(2) }) -test('existing malformed routing is left untouched by the importer while ordinary sticky routing still refuses it', async () => { +test('malformed routing stays untouched while requests can still select an admitted account', async () => { const f = await fixture() const bytes = 'pre-existing malformed vault routing' await writeFile(f.vault.routing, bytes, { mode: 0o600 }) - expect((await f.runtime.read()).mode).toBe('claustrum') + expect((await f.readAndDrain()).mode).toBe('claustrum') expect(await readFile(f.vault.routing, 'utf8')).toBe(bytes) - await expect(f.resolve('session')).rejects.toThrow() + expect((await f.resolveAndDrain('session'))?.accountId).toBe(f.workRoute) expect(await readFile(f.vault.routing, 'utf8')).toBe(bytes) await f.assertLegacy() }) @@ -579,7 +604,7 @@ test('first snapshot with existing settings completes while another writer still trackDetached(holder) await entered.wait try { - const snapshot = await f.runtime.read() + const snapshot = await f.readAndDrain() expect(holding).toBe(true) expect(snapshot.settings.routing).toEqual({ mode: 'fallback-first' }) expect(await readFile(f.vault.config, 'utf8')).toBe(existing) diff --git a/packages/core/src/tests/sticky-routing-contention.test.ts b/packages/core/src/tests/sticky-routing-contention.test.ts new file mode 100644 index 00000000..db3e9da7 --- /dev/null +++ b/packages/core/src/tests/sticky-routing-contention.test.ts @@ -0,0 +1,260 @@ +import { expect } from 'bun:test' +import { createHash } from 'node:crypto' +import { mkdtemp, readFile, rm, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { acquireRefreshFileLock } from '@cortexkit/common-auth/fs' +import { + __setLogTestSink, + getLogLevel, + type LogTestRecord, + setLogLevel, +} from '../logger.ts' +import type { NativeLockEvent } from '../native-lock-observer.ts' +import { StickySessionRouter } from '../sticky-routing.ts' +import { createTestLifetimeSuite } from './test-lifetime.ts' + +const lifetimes = createTestLifetimeSuite() +const key = (session: string) => + createHash('sha256').update(session).digest('hex') +const input = (sessionId: string, accountId = 'main') => ({ + sessionId, + family: 'general' as const, + modelId: 'claude-sonnet-5', + candidates: [{ accountId, order: 0, quota: { kind: 'unknown' as const } }], + retainAccountIds: new Set(['main', 'other']), + storage: null, + inputBytes: 1, +}) +async function fixture() { + const directory = await mkdtemp(join(tmpdir(), 'sticky-held-lock-')) + const path = join(directory, 'routing.json') + const routers: StickySessionRouter[] = [] + const leases: Array< + NonNullable>> + > = [] + lifetimes.deferCleanup(async () => { + await Promise.all(leases.map((lease) => lease.release())) + await Promise.all(routers.map((router) => router.drain())) + await rm(directory, { recursive: true, force: true }) + }) + const router = (onLockEvent?: (event: NativeLockEvent) => void) => { + const value = new StickySessionRouter({ path, onLockEvent }) + routers.push(value) + return value + } + const hold = async () => { + const lease = await acquireRefreshFileLock({ + path, + name: 'write', + ttlMs: 60_000, + renew: true, + }) + if (!lease) throw new Error('Fixture lock not acquired') + leases.push(lease) + return lease + } + return { directory, path, router, hold } +} + +lifetimes.test( + 'sticky routing decisions read published pins without waiting and publish after the writer releases', + async () => { + const f = await fixture() + const a = f.router() + const attempted = lifetimes.gate() + let measuring = false + const b = f.router((event) => { + if (measuring && !event.detached && event.type === 'acquiring') + attempted.open() + }) + await b.resolve(input('unrelated-b', 'other')) + await b.drain() + await a.resolve(input('pinned-by-a')) + await a.drain() + const held = await f.hold() + try { + measuring = true + const pending = b.resolve(input('pinned-by-a', 'other')) + lifetimes.trackDetached(pending) + const result = await Promise.race([ + pending.then((route) => ({ settled: true as const, route })), + attempted.wait.then(() => ({ settled: false as const })), + ]) + measuring = false + expect(result.settled).toBe(true) + if (!result.settled) return + expect(result.route?.accountId).toBe('main') + expect((await b.resolve(input('new-b', 'other')))?.accountId).toBe( + 'other', + ) + await held.assertOwned() + const before = JSON.parse(await readFile(f.path, 'utf8')) + expect(before.assignments[key('new-b')]).toBeUndefined() + before.assignments[key('late-a')] = { + ...before.assignments[key('pinned-by-a')], + } + await writeFile(f.path, JSON.stringify(before), { mode: 0o600 }) + let drained = false + const closing = b.drain().then(() => { + drained = true + }) + lifetimes.trackDetached(closing) + await Promise.resolve() + expect(drained).toBe(false) + await held.release() + await closing + const after = JSON.parse(await readFile(f.path, 'utf8')) + expect(after.assignments[key('new-b')].accountId).toBe('other') + expect(after.assignments[key('pinned-by-a')].accountId).toBe('main') + expect(after.assignments[key('late-a')].accountId).toBe('main') + expect(after.assignments[key('unrelated-b')].accountId).toBe('other') + } finally { + await held.release() + } + }, + 10_000, +) + +for (const operation of ['assign', 'clear', 'seed'] as const) { + lifetimes.test( + `sticky routing ${operation} publishes after a held bookkeeping lock is released`, + async () => { + const f = await fixture() + const router = f.router() + const assignment = { + accountId: 'main', + family: 'general', + affinityModelId: 'claude-sonnet-5', + assignedAt: Date.now(), + lastSeenAt: Date.now(), + initialInputBytes: 1, + quotaCheckedAt: 0, + } + const state = { + version: 1, + updatedAt: Date.now(), + assignments: + operation === 'clear' ? { [key('session')]: assignment } : {}, + ...(operation === 'seed' ? { vaultLegacySeed: 'pending' } : {}), + } + const bytes = JSON.stringify(state) + await writeFile(f.path, bytes, { mode: 0o600 }) + const held = await f.hold() + if (operation === 'assign') + expect((await router.resolve(input('session')))?.accountId).toBe('main') + else if (operation === 'clear') { + await router.clear('session') + expect( + await router.resolve({ ...input('session'), candidates: [] }), + ).toBeNull() + } else { + router.setVaultSeed({ + legacyConfig: join(f.directory, 'legacy.json'), + roster: join(f.directory, 'roster.json'), + }) + await router.prepareLegacySeed({ + version: 1, + complete: true, + rows: [], + declined: [], + }) + } + await held.assertOwned() + expect(await readFile(f.path, 'utf8')).toBe(bytes) + await held.release() + await router.drain() + const published = JSON.parse(await readFile(f.path, 'utf8')) + if (operation === 'assign') + expect(published.assignments[key('session')].accountId).toBe('main') + else if (operation === 'clear') expect(published.assignments).toEqual({}) + else expect(published.vaultLegacySeed).toBe('complete') + }, + ) +} + +lifetimes.test( + 'failed sticky publication logs once, preserves its choice, and retries on the next resolve', + async () => { + const f = await fixture() + const router = f.router() + const seam = router as unknown as { + writeState(...args: unknown[]): Promise + } + const write = seam.writeState.bind(router) + let refuse = true + seam.writeState = async (...args) => { + if (refuse) { + refuse = false + throw new Error('synthetic private write failure') + } + return write(...args) + } + const records: LogTestRecord[] = [] + const previous = getLogLevel() + setLogLevel('warn') + __setLogTestSink((record) => records.push(record)) + lifetimes.deferCleanup(() => { + __setLogTestSink(null) + setLogLevel(previous) + }) + expect((await router.resolve(input('retry')))?.accountId).toBe('main') + await router.drain() + expect(await readFile(f.path, 'utf8').catch(() => null)).toBeNull() + expect(records).toHaveLength(1) + expect(records[0]).toMatchObject({ + channel: 'sticky-routing', + message: 'assignment write refused', + payload: { operation: 'assign', reason: 'storage-unavailable' }, + }) + expect((await router.resolve(input('retry', 'other')))?.accountId).toBe( + 'main', + ) + await router.drain() + expect( + JSON.parse(await readFile(f.path, 'utf8')).assignments[key('retry')] + .accountId, + ).toBe('main') + expect(records).toHaveLength(1) + }, +) + +lifetimes.test( + 'a clear supersedes an assignment already waiting for the writer lease', + async () => { + const f = await fixture() + const entered = lifetimes.gate() + const router = f.router((event) => { + if (event.detached && event.type === 'acquiring') entered.open() + }) + const seam = router as unknown as { + writeState( + state: { assignments: Record }, + beforeRename: () => Promise, + prune?: boolean, + ): Promise + } + const write = seam.writeState.bind(router) + const published: Array = [] + seam.writeState = async (...args) => { + published.push(args[0].assignments[key('superseded')]?.accountId) + return write(...args) + } + const held = await f.hold() + try { + expect((await router.resolve(input('superseded')))?.accountId).toBe( + 'main', + ) + await entered.wait + await router.clear('superseded') + await held.assertOwned() + } finally { + await held.release() + } + await router.drain() + expect(published).toEqual([undefined]) + expect( + JSON.parse(await readFile(f.path, 'utf8')).assignments[key('superseded')], + ).toBeUndefined() + }, +) diff --git a/packages/opencode/src/index.ts b/packages/opencode/src/index.ts index c220e57f..9f3ad2a1 100644 --- a/packages/opencode/src/index.ts +++ b/packages/opencode/src/index.ts @@ -137,6 +137,7 @@ import { type NativeLocalCredentialValidation as NativeKnownCredentialSubject, type NativeMenuDispatch, type NativeQuotaObservation, + NativeRuntimeError, type NativeUiOptions, NativeVaultNoAccountError, nativeAccountCanServe, @@ -356,7 +357,7 @@ async function endIfCancelled( signal.throwIfAborted() } -// Selecting another account cannot repair locally invalid conversation history. +// Selecting another account cannot repair a locally refused request. // Provider-generated 400 responses still follow the configured fallbackOn list. function shouldFallbackResponse( response: Response, @@ -1494,6 +1495,10 @@ const anthropicAuthPlugin = async ( vaultWatch?.close() await usagePublications.drain() await nativeAccounts.drainQuotaPublications() + await Promise.all( + [...stickySessionRouters.values()].map((router) => router.drain()), + ) + await nativeAccounts.drainRoutingPublications() nativeAccounts.close() } @@ -1776,6 +1781,7 @@ const anthropicAuthPlugin = async ( router = new StickySessionRouter({ path, vaultSeed: snapshot.routingSeed, + onLockEvent: nativeAccounts.onLockEvent, localRecoveryPath: snapshot.mode === 'local' ? snapshot.routingRecoveryPath : undefined, }) @@ -7405,7 +7411,19 @@ const anthropicAuthPlugin = async ( } const requestStorageForIdentity = await getRequestStorage() - const relayConfig = await nativeAccounts.getRelayConfig() + let relayConfig: Awaited< + ReturnType + > + try { + relayConfig = await nativeAccounts.getRelayConfig() + } catch (error) { + if ( + error instanceof NativeRuntimeError && + error.code === 'relay-inconsistent' + ) + return localInvalidRequestResponse(error.message) + throw error + } const persistedFallbackAccount = oauthAccountId === 'main' ? undefined diff --git a/packages/opencode/src/tests/local-routing-migration.test.ts b/packages/opencode/src/tests/local-routing-migration.test.ts index d6ab7e8f..17e4fb11 100644 --- a/packages/opencode/src/tests/local-routing-migration.test.ts +++ b/packages/opencode/src/tests/local-routing-migration.test.ts @@ -273,7 +273,7 @@ async function oc1(f: HostFixture) { provider, ) const handler = options.fetch as typeof fetch - return async (session: string) => { + const send = async (session: string) => { const request = handler('https://api.anthropic.com/v1/messages', { method: 'POST', headers: { @@ -292,6 +292,11 @@ async function oc1(f: HostFixture) { await response.text() return f.sent.at(-1) } + return Object.assign(send, { + dispose: async () => { + await plugin.dispose?.() + }, + }) } async function oc2(f: HostFixture) { @@ -307,7 +312,7 @@ async function oc2(f: HostFixture) { fetch: f.offlineFetch, }) f.cleanups.push(() => source.close()) - return async (sessionID: string) => { + const send = async (sessionID: string) => { const scope = { providerID: 'anthropic', modelID: 'claude-sonnet-4-6', @@ -326,6 +331,7 @@ async function oc2(f: HostFixture) { tests.trackDetached(work) return work } + return Object.assign(send, { dispose: () => source.close() }) } async function emptyFixture(override: boolean): Promise { @@ -472,6 +478,7 @@ for (const [host, open] of [ const send = await open(f) expect(await send(session1)).toContain('sk-ant-oat01-secondary') await send(session2) + await send.dispose() const destination = resolveLocalRoutingDestination( 'opencode', f.paths, @@ -516,19 +523,23 @@ tests.test( await writeFile(legacy, legacyBytes) const send = await open(f) expect(await send(session1)).toContain('sk-ant-oat01-primary') + await send.dispose() const state = JSON.parse(await readFile(destination, 'utf8')) expect(state.assignments[hash(session1)].accountId).toBe('main') expect(state.assignments[hash(session2)].accountId).toBe('secondary') expect(state.assignments[hash('missing-account-session')].accountId).toBe( 'removed-account', ) - expect(await send(session2)).toContain('sk-ant-oat01-secondary') + const resumed = await open(f) + expect(await resumed(session2)).toContain('sk-ant-oat01-secondary') + await resumed.dispose() expect(await readFile(legacy, 'utf8')).toBe(legacyBytes) // A second router in the same process must not import newly added source pins. const lateBytes = routing({ 'late-session': 'secondary' }) await writeFile(legacy, lateBytes) const second = await open(f) await second(session1) + await second.dispose() expect( JSON.parse(await readFile(destination, 'utf8')).assignments[ hash('late-session') @@ -566,7 +577,9 @@ tests.test( for (const failure of ['malformed', 'unreadable', 'locked'] as const) { tests.test( - `local routing serves destination pins after ${failure} legacy recovery fails`, + failure === 'locked' + ? 'local routing serves destination pins while legacy recovery waits in the background' + : `local routing serves destination pins after ${failure} legacy recovery fails`, async () => { for (const open of [oc1, oc2]) { const f = await fixture() @@ -607,6 +620,7 @@ for (const failure of ['malformed', 'unreadable', 'locked'] as const) { restoreRead = () => read.mockRestore() f.cleanups.push(restoreRead) } + let held: Awaited> = null if (failure === 'locked') { const lock = await acquireRefreshFileLock({ path: f.journal.routingPaths.destination, @@ -616,23 +630,31 @@ for (const failure of ['malformed', 'unreadable', 'locked'] as const) { }) if (!lock) throw new Error('Fixture could not hold the routing lock') + held = lock f.cleanups.push(() => lock.release()) } expect(await send(session1)).toContain('sk-ant-oat01-secondary') expect(await send(session1)).toContain('sk-ant-oat01-secondary') + await held?.assertOwned() + await held?.release() + await send.dispose() expect( records.filter( (record) => record.channel === 'sticky-routing' && record.level === 'warn', ), - ).toEqual([ - { - level: 'warn', - channel: 'sticky-routing', - message: 'local pin recovery unavailable', - payload: { reason: 'recovery-failed' }, - }, - ]) + ).toEqual( + failure === 'locked' + ? [] + : [ + { + level: 'warn', + channel: 'sticky-routing', + message: 'local pin recovery unavailable', + payload: { reason: 'recovery-failed' }, + }, + ], + ) if (failure === 'unreadable') expect(deniedReads).toBe(1) restoreRead() expect(await readFile(legacy, 'utf8')).toBe(bytes) diff --git a/packages/opencode/src/tests/native-serving.test.ts b/packages/opencode/src/tests/native-serving.test.ts index a4321e5d..19567ab6 100644 --- a/packages/opencode/src/tests/native-serving.test.ts +++ b/packages/opencode/src/tests/native-serving.test.ts @@ -41,6 +41,8 @@ import { toNativeQuotaMap, } from '../../../core/src/native-quota-codec.ts' import { nativeVaultPaths } from '../../../core/src/native-vault-paths.ts' +import { holdNativeRelayLock } from '../../../core/src/tests/native-relay-lock-fixture.ts' +import { requestLockProbe } from '../../../core/src/tests/native-request-lock-fixture.ts' import { createTestLifetimeSuite } from '../../../core/src/tests/test-lifetime.ts' import { MockRelayServer } from '../../../e2e-tests/src/mock-relay.ts' import { AnthropicAuthPlugin } from '../index.ts' @@ -1026,6 +1028,160 @@ test('supervised content after real migration refuses foreground and background expect(fixture.records).toHaveLength(0) }) +for (const scenario of ['local-healthy', 'vault-healthy', 'refresh'] as const) { + test(`OC1 ${scenario} request observes only required foreground locks`, async () => { + const fixture = await migrateServingFixture( + scenario === 'vault-healthy' ? 'claustrum' : 'local', + ) + const probe = requestLockProbe() + const runtime = createNativeAccountRuntime({ + paths: fixture.paths, + host: 'opencode', + vault: { connect: async () => fixture.scopedClient }, + }) + lifetimes.deferCleanup(async () => { + await runtime.drainRoutingPublications() + runtime.close() + }) + if (scenario === 'local-healthy') + expect((await runtime.authorizeLocal('main')).status).toBe('usable') + if (scenario !== 'refresh') { + await runtime.fetchQuota( + 'main', + Object.assign(fixture.fetchProvider, { + preconnect: originalFetch.preconnect, + }), + ) + await runtime.drainQuotaPublications() + } + await runtime.updateSettings((settings) => { + settings.routing = { mode: 'sticky-balanced' } + }) + await runtime.read() + await runtime.drainRoutingPublications() + network.mockImplementation((input, init) => { + const url = input instanceof Request ? input.url : input.toString() + if (url.includes('/v1/messages')) probe.dispatch() + return fixture.fetchProvider(input, init) + }) + const plugin = await createPlugin({ + nativeLocal: { onLockEvent: probe.onLockEvent }, + ...(scenario === 'vault-healthy' + ? { claustrumScopedConnect: async () => fixture.scopedClient } + : {}), + }) + probe.begin() + const response = await sendNative(plugin) + expect(response.status).toBe(200) + await response.text() + if (scenario === 'refresh') + expect( + probe.foreground.some( + (event) => + event.name.startsWith('native-account-provider-') && + event.path === fixture.paths.state, + ), + ).toBe(true) + else expect(probe.foreground).toEqual([]) + await plugin.dispose?.() + expect(probe.detached.some((event) => event.name === 'write')).toBe(true) + expect(probe.afterDispatch.length).toBeGreaterThan(0) + expect( + probe.afterDispatch.some((event) => event.name === 'native-runtime'), + ).toBe(true) + if (scenario === 'vault-healthy') + expect( + probe.afterDispatch.some( + (event) => event.name === 'claustrum-roster-write', + ), + ).toBe(true) + }) +} + +test('OC1 inconsistent relay pair refuses locally without any dispatch', async () => { + const fixture = await migrateServingFixture('local') + const runtime = createNativeAccountRuntime({ + paths: fixture.paths, + host: 'opencode', + }) + lifetimes.deferCleanup(() => runtime.close()) + await runtime.updateRelay({ + enabled: true, + url: 'https://never-send.invalid', + token: 'synthetic-never-send-token', + }) + await createNativePoolStore({ + paths: fixture.paths, + quota: nativeQuotaCodec, + }).updateSettings((settings) => { + ;(settings.relay as Record).generation = 'mismatched' + }) + const plugin = await createPlugin() + const response = await sendNative(plugin) + expect(response.status).toBe(400) + expect(await response.json()).toMatchObject({ + error: { + type: 'invalid_request_error', + message: + 'request refused: relay configuration is inconsistent; no request was sent', + }, + }) + expect( + fixture.records.filter((record) => record.url.includes('/v1/messages')), + ).toHaveLength(0) +}) + +for (const configured of [false, true]) { + test(`OC1 request ignores held native-relay lock (${configured ? 'relay' : 'direct'})`, async () => { + const fixture = await migrateServingFixture('local') + const runtime = createNativeAccountRuntime({ + paths: fixture.paths, + host: 'opencode', + }) + lifetimes.deferCleanup(() => runtime.close()) + const relay = new MockRelayServer() + relays.add(relay) + const address = await relay.start({ token: 'synthetic-held-relay-token' }) + if (configured) + await runtime.updateRelay({ + enabled: true, + url: address.url, + token: 'synthetic-held-relay-token', + transport: 'http', + fallbackToDirect: false, + }) + network.mockImplementation((input, init) => { + const url = input instanceof Request ? input.url : input.toString() + return url.startsWith(address.url) + ? originalFetch(input, init) + : fixture.fetchProvider(input, init) + }) + const plugin = await createPlugin() + const held = await holdNativeRelayLock(fixture.paths.config, lifetimes) + try { + const response = await sendNative(plugin) + expect(response.status).toBe(200) + expect(await response.text()).toContain('native-serving-ok') + expect(relay.acceptedRequests()).toBe(configured ? 1 : 0) + expect(await runtime.getRelayConfig()).toEqual( + configured + ? { + enabled: true, + url: address.url, + token: 'synthetic-held-relay-token', + transport: 'http', + fallbackToDirect: false, + } + : null, + ) + await held.assertOwned() + } finally { + await held.pastTimeout + await held.release() + } + }, 25_000) +} + for (const transport of ['http', 'websocket'] as const) { test(`migrated vault ${transport} relay uses native private relay config and fresh vault receipt`, async () => { const fixture = await migrateServingFixture('claustrum') @@ -1447,6 +1603,7 @@ test('a replaced default vault account loses its session pins and quota while an // Earlier sessions pinned through the vault routing file: one to the // default account, one to the other vault account. const earlier = new StickySessionRouter({ path: routingPath }) + lifetimes.deferCleanup(() => earlier.drain()) const identities = new Map([ ['main', mainIdentity], [fallbackRoute, fallbackIdentity], @@ -1487,6 +1644,7 @@ test('a replaced default vault account loses its session pins and quota while an }) expect(pinned?.accountId).toBe(accountId) } + await earlier.drain() const defaultPin = await stickyPin(routingPath, 'session-on-default') const fallbackPin = await stickyPin(routingPath, 'session-on-fallback') expect(defaultPin?.accountIdentity).toBe(mainIdentity) @@ -1519,6 +1677,7 @@ test('a replaced default vault account loses its session pins and quota while an expect((await sendNative(plugin, model, 'session-on-fallback')).status).toBe( 200, ) + await plugin.dispose?.() // The session pinned to the replaced account is placed again and its new // pin names the account that now serves it. const replaced = await stickyPin(routingPath, 'session-on-default') @@ -2621,6 +2780,7 @@ test('OpenCode routes normally during roster deferral without publishing an empt try { await runtime.vault.refresh() await runtime.read() + await runtime.drainRoutingPublications() expect( JSON.parse(await readFile(vault.routing, 'utf8')).vaultLegacySeed, ).toBe('complete') diff --git a/packages/opencode/src/tests/opencode2-native-source.test.ts b/packages/opencode/src/tests/opencode2-native-source.test.ts index 6e85fe87..5ff90f6d 100644 --- a/packages/opencode/src/tests/opencode2-native-source.test.ts +++ b/packages/opencode/src/tests/opencode2-native-source.test.ts @@ -20,6 +20,7 @@ import { type NativeAccountRuntime, type NativeCustodyClient, type NativePoolPaths, + type ProviderAccountUuid, readNativeAccountsForRequest, readNativeEnrollmentMode, resolveNativePoolPaths, @@ -31,6 +32,8 @@ import { fingerprintOf } from '@cortexkit/common-auth/store' import type { IntegrationOAuthMethodRegistration } from '@opencode/plugin/promise/integration' import { nativeQuotaCodec } from '../../../core/src/native-quota-codec' import { nativeVaultPaths } from '../../../core/src/native-vault-paths' +import { holdNativeRelayLock } from '../../../core/src/tests/native-relay-lock-fixture.ts' +import { requestLockProbe } from '../../../core/src/tests/native-request-lock-fixture.ts' import { createTestLifetimeSuite } from '../../../core/src/tests/test-lifetime.ts' import { tokenFingerprint } from '../../../core/src/token-fingerprint' import { @@ -134,6 +137,166 @@ usageTests.test( }, ) +usageTests.test( + 'OC2 inconsistent relay pair refuses locally before dispatch authorization', + async () => { + const owned: Array<() => Promise | void> = [] + usageTests.deferCleanup(async () => { + for (const cleanup of owned.reverse()) await cleanup() + }) + const fixture = await local(undefined, owned) + await fixture.runtime.updateRelay({ + enabled: true, + url: 'https://never-send.invalid', + token: 'synthetic-never-send-token', + }) + await fixture.store.updateSettings((settings) => { + ;(settings.relay as Record).generation = 'mismatched' + }) + const native = source(fixture.runtime, fixture.paths, owned) + await expect(native.authorizeAccount(scope, 'main')).rejects.toMatchObject({ + code: 'relay-inconsistent', + status: 400, + retryable: false, + message: + 'request refused: relay configuration is inconsistent; no request was sent', + }) + }, +) + +for (const scenario of ['local-healthy', 'vault-healthy', 'refresh'] as const) { + usageTests.test( + `OC2 ${scenario} request observes only required foreground locks`, + async () => { + const owned: Array<() => Promise | void> = [] + usageTests.deferCleanup(async () => { + for (const cleanup of owned.reverse()) await cleanup() + }) + const probe = requestLockProbe() + const now = () => Date.now() + 2 * 3600_000 + const fixture = + scenario === 'vault-healthy' + ? await vault( + [7], + () => identity, + { mode: 'sticky-balanced' }, + probe.onLockEvent, + ) + : await local( + { + onLockEvent: probe.onLockEvent, + ...(scenario === 'refresh' + ? { + now, + refreshToken: async () => ({ + access: 'sk-ant-oat01-refreshed-control', + refresh: 'synthetic-control-refresh', + expires: now() + 3600_000, + expiresIn: 3600, + }), + resolveIdentity: async () => ({ + accountUuid: identity as ProviderAccountUuid, + deviceId: 'fixture', + sessionId: 'fixture', + }), + } + : {}), + }, + owned, + ) + await fixture.runtime.updateSettings((settings) => { + settings.routing = { mode: 'sticky-balanced' } + }) + const native = source(fixture.runtime, fixture.paths, owned) + probe.begin() + expect(await native.chooseAccount(scope)).toBe('main') + const auth = await native.authorizeAccount(scope, 'main') + expect(auth).toBeDefined() + if (!auth) return + probe.dispatch() + if (scenario === 'refresh') + expect( + probe.foreground.some( + (event) => + event.name.startsWith('native-account-provider-') && + event.path === fixture.paths.state, + ), + ).toBe(true) + else expect(probe.foreground).toEqual([]) + await native.onAttemptEnd?.(context(auth.provenance), { + status: 200, + outputStarted: true, + }) + await native.close() + expect(probe.detached.some((event) => event.name === 'write')).toBe(true) + expect(probe.afterDispatch.length).toBeGreaterThan(0) + expect( + probe.afterDispatch.some((event) => event.name === 'native-runtime'), + ).toBe(true) + if (scenario === 'vault-healthy') + expect( + probe.afterDispatch.some( + (event) => event.name === 'claustrum-roster-write', + ), + ).toBe(true) + }, + ) +} + +for (const configured of [false, true]) { + usageTests.test( + `OC2 request ignores held native-relay lock (${configured ? 'relay' : 'direct'})`, + async () => { + const owned: Array<() => Promise | void> = [] + usageTests.deferCleanup(async () => { + for (const cleanup of owned.reverse()) await cleanup() + }) + const fixture = await local(undefined, owned) + if (configured) + await fixture.runtime.updateRelay({ + enabled: true, + url: 'https://held-relay.invalid', + token: 'synthetic-held-relay-token', + transport: 'http', + fallbackToDirect: false, + }) + const native = source(fixture.runtime, fixture.paths, owned) + const held = await holdNativeRelayLock(fixture.paths.config, usageTests) + try { + const account = await native.chooseAccount(scope) + expect(account).toBe('main') + const auth = await native.authorizeAccount(scope, 'main') + if (!auth) throw new Error('Native authorization refused') + expect(auth.dispatch).toEqual( + configured + ? { + kind: 'http-relay', + url: 'https://held-relay.invalid', + token: 'synthetic-held-relay-token', + } + : { kind: 'direct' }, + ) + expect(await fixture.runtime.getRelayConfig()).toEqual( + configured + ? { + enabled: true, + url: 'https://held-relay.invalid', + token: 'synthetic-held-relay-token', + transport: 'http', + fallbackToDirect: false, + } + : null, + ) + await held.assertOwned() + } finally { + await held.pastTimeout + await held.release() + } + }, + 25_000, + ) +} + async function storage(owned = cleanups) { const parent = resolve( import.meta.dir, @@ -314,6 +477,7 @@ async function vault( versions: number[] = [7, 8], servedIdentity: (get: number) => string = () => identity, routing: { mode: string } = { mode: 'main-first' }, + onLockEvent?: NativeAccountRuntime['onLockEvent'], ) { const { paths, env, store: localStore } = await storage() await writeFile( @@ -378,6 +542,7 @@ async function vault( const runtime = createNativeAccountRuntime({ paths, host: 'opencode', + local: { onLockEvent }, vault: { env, connect: async () => client, @@ -1305,9 +1470,11 @@ describe('OC2 source over actual native services', () => { } ).assignments, ).map((assignment) => assignment.accountId) + await native.drainRoutingPublications() expect(await assignments(snapshot.routingPath)).toEqual(['main']) expect(await exists(localRouting)).toBe(false) await native.resetRouting(scope.sessionID) + await native.drainRoutingPublications() expect(await assignments(snapshot.routingPath)).toEqual([]) expect(await exists(localRouting)).toBe(false) }) @@ -1334,6 +1501,7 @@ describe('OC2 source over actual native services', () => { ).toBe(fallbackRoute) const routing = before.routingPath if (!routing) throw new Error('Vault mode named no routing file') + await native.drainRoutingPublications() const pinned = await readFile(routing, 'utf8') expect(JSON.parse(rawPin(pinned, mainSession.sessionID))).toMatchObject({ accountId: 'main', @@ -1353,6 +1521,7 @@ describe('OC2 source over actual native services', () => { ) if (!replaced) throw new Error('Replaced default account is missing') const choice = await native.chooseAccount(mainSession) + await native.drainRoutingPublications() const repinned = JSON.parse( rawPin(await readFile(routing, 'utf8'), mainSession.sessionID), ) @@ -1378,6 +1547,7 @@ describe('OC2 source over actual native services', () => { expect(snapshot.routingPath).toBe( join(dirname(fixture.paths.config), 'native-routing.json'), ) + await native.drainRoutingPublications() const pin = JSON.parse( rawPin(await readFile(snapshot.routingPath!, 'utf8'), scope.sessionID), ) diff --git a/packages/opencode/src/tests/opencode2-vault-legacy-seed.test.ts b/packages/opencode/src/tests/opencode2-vault-legacy-seed.test.ts index 55669af9..4495f0dc 100644 --- a/packages/opencode/src/tests/opencode2-vault-legacy-seed.test.ts +++ b/packages/opencode/src/tests/opencode2-vault-legacy-seed.test.ts @@ -17,9 +17,10 @@ import { nativeVaultPaths } from '../../../core/src/native-vault-paths.ts' import { createTestLifetimeSuite } from '../../../core/src/tests/test-lifetime.ts' import { createNativeAnthropicSource } from '../v2/source.ts' -const owners = new Set<{ close(): void }>() -afterEach(() => { - for (const owner of owners) owner.close() +type Owner = { close(): void | Promise } +const owners = new Set() +afterEach(async () => { + await Promise.all([...owners].map((owner) => owner.close())) owners.clear() }) const { test, deferCleanup } = createTestLifetimeSuite() @@ -30,7 +31,16 @@ const hash = (session: string) => async function fixture() { const root = await mkdtemp(join(tmpdir(), 'oc2-vault-legacy-seed-')) - deferCleanup(() => rm(root, { recursive: true, force: true })) + const owned: Owner[] = [] + const own = (owner: T): T => { + owned.push(owner) + owners.add(owner) + return owner + } + deferCleanup(async () => { + await Promise.all(owned.map((owner) => owner.close())) + await rm(root, { recursive: true, force: true }) + }) const paths = await resolveNativePoolPaths( join(root, 'anthropic-auth.json'), join(root, 'anthropic-auth-state.json'), @@ -141,16 +151,14 @@ async function fixture() { host: 'opencode' as const, vault: { env, connect: async () => client }, } - const runtime = createNativeAccountRuntime(options) - owners.add(runtime) - deferCleanup(() => runtime.close()) - const native = createNativeAnthropicSource({ - runtime, - paths, - requestOptions: () => ({}), - }) - owners.add(native) - deferCleanup(() => native.close()) + const runtime = own(createNativeAccountRuntime(options)) + const native = own( + createNativeAnthropicSource({ + runtime, + paths, + requestOptions: () => ({}), + }), + ) const scope = { providerID: 'anthropic', modelID: 'claude-sonnet-5-5', @@ -160,6 +168,7 @@ async function fixture() { } return { paths, + own, vault, runtime, native, @@ -179,6 +188,7 @@ test('OpenCode 2 live pin during roster deferral carries pending provenance and expect((await f.native.authorizeAccount(f.scope, 'main'))?.route).toBe( 'oauth', ) + await f.native.drainRoutingPublications() const pending = JSON.parse(await readFile(f.vault.routing, 'utf8')) expect(pending.vaultLegacySeed).toBe('pending') expect(Object.keys(pending.assignments)).toEqual([hash('live-session')]) @@ -188,12 +198,16 @@ test('OpenCode 2 live pin during roster deferral carries pending provenance and next: f.roster, result: undefined, })) + // A ready roster schedules the one-time import; persistence checks join it. + await f.runtime.read() + await f.native.drainRoutingPublications() expect( await f.native.chooseAccount({ ...f.scope, sessionID: 'legacy-only-session', }), ).toBe(f.workRoute) + await f.native.drainRoutingPublications() const complete = JSON.parse(await readFile(f.vault.routing, 'utf8')) expect(complete.vaultLegacySeed).toBe('complete') expect(complete.assignments[hash('live-session')]).toEqual(live) @@ -203,11 +217,14 @@ test('OpenCode 2 live pin during roster deferral carries pending provenance and initialInputBytes: 22, }) await f.native.resetRouting('legacy-only-session') + await f.native.drainRoutingPublications() const before = await readFile(f.vault.routing) - const restarted = createNativeAccountRuntime(f.options) - owners.add(restarted) - deferCleanup(() => restarted.close()) + const reset = JSON.parse(before.toString('utf8')) + expect(reset.assignments[hash('live-session')]).toEqual(live) + expect(reset.assignments[hash('legacy-only-session')]).toBeUndefined() + const restarted = f.own(createNativeAccountRuntime(f.options)) await restarted.read() + await restarted.drainRoutingPublications() expect(await readFile(f.vault.routing)).toEqual(before) expect(await readFile(f.paths.legacyConfig, 'utf8')).toBe(f.legacyConfig) expect(await readFile(f.routingPath, 'utf8')).toBe(f.legacyRouting) diff --git a/packages/opencode/src/tests/sticky-routing.test.ts b/packages/opencode/src/tests/sticky-routing.test.ts index 0f3da537..786f02df 100644 --- a/packages/opencode/src/tests/sticky-routing.test.ts +++ b/packages/opencode/src/tests/sticky-routing.test.ts @@ -15,6 +15,14 @@ import { } from '@cortexkit/anthropic-auth-core' const NOW = Date.parse('2026-07-18T08:00:00Z') +const routers: StickySessionRouter[] = [] +function ownRouter( + options: ConstructorParameters[0], +) { + const router = new StickySessionRouter(options) + routers.push(router) + return router +} const directories: string[] = [] type KnownStickyRouteCandidate = Omit & { quota: OAuthQuotaSnapshot @@ -81,6 +89,7 @@ const storage = { } afterEach(async () => { + await Promise.all(routers.splice(0).map((router) => router.drain())) await Promise.all( directories .splice(0) @@ -246,7 +255,7 @@ describe('sticky-balanced session routing', () => { test('distributes cold Fable sessions proportionally instead of draining account order', async () => { const path = await statePath() - const router = new StickySessionRouter({ path, now: () => NOW }) + const router = ownRouter({ path, now: () => NOW }) const scarce = candidate({ accountId: 'yiyi', order: 0, @@ -283,7 +292,7 @@ describe('sticky-balanced session routing', () => { test('seeds an existing cachekeep session from its current OAuth route', async () => { const path = await statePath() - const router = new StickySessionRouter({ path, now: () => NOW }) + const router = ownRouter({ path, now: () => NOW }) const scarce = candidate({ accountId: 'existing-route', order: 0, @@ -315,7 +324,7 @@ describe('sticky-balanced session routing', () => { test('does not seed a cold session onto an ineligible preferred account', async () => { const path = await statePath() - const router = new StickySessionRouter({ path, now: () => NOW }) + const router = ownRouter({ path, now: () => NOW }) const exhausted = candidate({ accountId: 'old-cache-route', order: 0, @@ -345,7 +354,7 @@ describe('sticky-balanced session routing', () => { test('persists an assignment and keeps it sticky when weights change', async () => { const path = await statePath() - const router = new StickySessionRouter({ path, now: () => NOW }) + const router = ownRouter({ path, now: () => NOW }) const scarce = candidate({ accountId: 'scarce', order: 0, @@ -371,11 +380,12 @@ describe('sticky-balanced session routing', () => { }) expect(first?.accountId).toBe('abundant') expect(first?.created).toBe(true) + await router.drain() const persisted = await readFile(path, 'utf8') expect(persisted).not.toContain('session-1') expect(persisted).toContain('abundant') - const reloaded = new StickySessionRouter({ path, now: () => NOW + 1_000 }) + const reloaded = ownRouter({ path, now: () => NOW + 1_000 }) const second = await reloaded.resolve({ sessionId: 'session-1', family: 'fable', @@ -407,7 +417,7 @@ describe('sticky-balanced session routing', () => { test('reallocates a legacy assignment whose selected model is unknown', async () => { const path = await statePath() - const router = new StickySessionRouter({ path, now: () => NOW }) + const router = ownRouter({ path, now: () => NOW }) const original = candidate({ accountId: 'original', order: 0, @@ -433,13 +443,14 @@ describe('sticky-balanced session routing', () => { } expect((await router.resolve(base))?.accountId).toBe('original') + await router.drain() const legacyState = JSON.parse(await readFile(path, 'utf8')) as { assignments: Record } delete Object.values(legacyState.assignments)[0]!.affinityModelId await writeFile(path, JSON.stringify(legacyState)) - const reloaded = new StickySessionRouter({ path, now: () => NOW + 1_000 }) + const reloaded = ownRouter({ path, now: () => NOW + 1_000 }) const changed = await reloaded.resolve({ ...base, candidates: [ @@ -452,7 +463,7 @@ describe('sticky-balanced session routing', () => { test('reselects the account when the user changes model families', async () => { const path = await statePath() - const router = new StickySessionRouter({ path, now: () => NOW }) + const router = ownRouter({ path, now: () => NOW }) const fableDepleted = candidate({ accountId: 'fable-depleted', order: 0, @@ -497,7 +508,7 @@ describe('sticky-balanced session routing', () => { test('reselects the account for a model change within one routing family', async () => { const path = await statePath() - const router = new StickySessionRouter({ path, now: () => NOW }) + const router = ownRouter({ path, now: () => NOW }) const first = candidate({ accountId: 'first', order: 0, @@ -560,7 +571,7 @@ describe('sticky-balanced session routing', () => { test('retains affinity when only an internal recovery model changes', async () => { const path = await statePath() - const router = new StickySessionRouter({ path, now: () => NOW }) + const router = ownRouter({ path, now: () => NOW }) const original = candidate({ accountId: 'original', order: 0, @@ -604,7 +615,7 @@ describe('sticky-balanced session routing', () => { test('does not reserve a direct Opus session from stale Fable exhaustion', async () => { const path = await statePath() - const router = new StickySessionRouter({ path, now: () => NOW }) + const router = ownRouter({ path, now: () => NOW }) const staleDepleted = candidate({ accountId: 'stale-depleted', order: 0, @@ -634,7 +645,7 @@ describe('sticky-balanced session routing', () => { test('uses known Fable-depleted accounts first for direct Opus sessions', async () => { const path = await statePath() - const router = new StickySessionRouter({ path, now: () => NOW }) + const router = ownRouter({ path, now: () => NOW }) const depleted = candidate({ accountId: 'main', order: 0, @@ -664,7 +675,7 @@ describe('sticky-balanced session routing', () => { test('expires an inactive assignment even while the router process stays alive', async () => { const path = await statePath() let now = NOW - const router = new StickySessionRouter({ + const router = ownRouter({ path, now: () => now, assignmentTtlMs: 1_000, @@ -696,8 +707,8 @@ describe('sticky-balanced session routing', () => { test('serializes concurrent first assignment across router instances', async () => { const path = await statePath() - const left = new StickySessionRouter({ path }) - const right = new StickySessionRouter({ path }) + const left = ownRouter({ path }) + const right = ownRouter({ path }) const candidates = [ candidate({ accountId: 'a', @@ -730,8 +741,8 @@ describe('sticky-balanced session routing', () => { test('observes assignment resets written by another router process', async () => { const path = await statePath() - const left = new StickySessionRouter({ path, now: () => NOW }) - const right = new StickySessionRouter({ path, now: () => NOW + 1 }) + const left = ownRouter({ path, now: () => NOW }) + const right = ownRouter({ path, now: () => NOW + 1 }) const available = candidate({ accountId: 'main', order: 0, @@ -749,14 +760,16 @@ describe('sticky-balanced session routing', () => { inputBytes: 1_000, } await left.resolve(input) + await left.drain() expect((await right.resolve(input))?.created).toBe(false) await left.clear(input.sessionId) + await left.drain() expect((await right.resolve(input))?.created).toBe(true) }) test('migrates atomically when the assigned account is explicitly excluded', async () => { const path = await statePath() - const router = new StickySessionRouter({ path, now: () => NOW }) + const router = ownRouter({ path, now: () => NOW }) const candidates = [ candidate({ accountId: 'a', @@ -884,7 +897,7 @@ describe('sticky-balanced session routing', () => { test('excludes an unknown-quota candidate from cold weighted selection when a known candidate exists', async () => { const path = await statePath() - const router = new StickySessionRouter({ path, now: () => NOW }) + const router = ownRouter({ path, now: () => NOW }) const known = candidate({ accountId: 'known', order: 1, @@ -910,7 +923,7 @@ describe('sticky-balanced session routing', () => { test('router-level escape degrades to ordered selection when every candidate quota is unknown', async () => { const path = await statePath() - const router = new StickySessionRouter({ path, now: () => NOW }) + const router = ownRouter({ path, now: () => NOW }) const resolution = await router.resolve({ sessionId: 'all-unknown-quota-selection', family: 'general', @@ -928,7 +941,7 @@ describe('sticky-balanced session routing', () => { test('does not escape to an exhausted candidate when every known quota is exhausted', async () => { const path = await statePath() - const router = new StickySessionRouter({ path, now: () => NOW }) + const router = ownRouter({ path, now: () => NOW }) const resolution = await router.resolve({ sessionId: 'all-exhausted-quota-selection', family: 'general', @@ -958,7 +971,7 @@ describe('sticky-balanced session routing', () => { test('router-level pending bytes reset when an unknown quota becomes known', async () => { const path = await statePath() - const router = new StickySessionRouter({ path, now: () => NOW }) + const router = ownRouter({ path, now: () => NOW }) const unknown = { accountId: 'a', quota: { kind: 'unknown' as const }, @@ -1001,7 +1014,7 @@ describe('sticky-balanced session routing', () => { test('keeps an excluded route out of a new assignment', async () => { const path = await statePath() - const router = new StickySessionRouter({ path, now: () => NOW }) + const router = ownRouter({ path, now: () => NOW }) const input = { sessionId: 'restricted-unknown-identity', family: 'general', diff --git a/packages/opencode/src/v2/source.ts b/packages/opencode/src/v2/source.ts index 75009cc0..87b21686 100644 --- a/packages/opencode/src/v2/source.ts +++ b/packages/opencode/src/v2/source.ts @@ -109,6 +109,7 @@ export interface NativeSourceOptions { export interface NativeSource extends NativeAnthropicSource { readonly close: () => Promise + readonly drainRoutingPublications: () => Promise readonly resetRouting: (sessionID: string) => Promise readonly recordPrime: ( provenance: NativeSourceProvenance, @@ -134,6 +135,7 @@ export function createNativeAnthropicSource( router = new StickySessionRouter({ path, vaultSeed: snapshot.routingSeed, + onLockEvent: runtime.onLockEvent, localRecoveryPath: snapshot.mode === 'local' ? snapshot.routingRecoveryPath : undefined, }) @@ -823,13 +825,21 @@ export function createNativeAnthropicSource( ) return false }, + drainRoutingPublications: () => + Promise.all([ + runtime.drainRoutingPublications(), + ...[...routers.values()].map((router) => router.drain()), + ]).then(() => {}), close() { closed = true for (const active of sessions.values()) active.abort() sessions.clear() choices.clear() retries.clear() - return usagePublications.drain() + return Promise.all([ + usagePublications.drain(), + source.drainRoutingPublications(), + ]).then(() => {}) }, } const publications = new WeakMap>() diff --git a/packages/pi/src/native.ts b/packages/pi/src/native.ts index f03c6a73..73d77864 100644 --- a/packages/pi/src/native.ts +++ b/packages/pi/src/native.ts @@ -8,6 +8,7 @@ import { killswitchPassesPolicy, logger, type NativeAccountRuntime, + type NativeAccountRuntimeOptions, type NativeAccountSnapshot, type NativeCustodyClient, type NativeCustodyReceipt, @@ -145,11 +146,27 @@ export class PiNativeRuntime { private closed = false private closing: Promise | undefined readonly usagePublications = createNativeUsagePublisher('native-auth') + private readonly routingPublications = new Set<{ drain(): Promise }>() + + ownRoutingPublications(router: { drain(): Promise }) { + this.routingPublications.add(router) + } + + async drainRoutingPublications() { + await Promise.all([ + ...[...this.routingPublications].map((router) => router.drain()), + this.initialization.then( + (runtime) => runtime.drainRoutingPublications(), + () => {}, + ), + ]) + } constructor( storagePath: string, options: { connect?: () => Promise + local?: NativeAccountRuntimeOptions['local'] vaultPoll?: PiVaultPollOptions } = {}, ) { @@ -162,6 +179,7 @@ export class PiNativeRuntime { const runtime = acquireNativeAccountRuntime({ paths, host: 'pi', + local: options.local, vault: { // Core checks Pi enrollment and reads a fresh scoped receipt on every send. // Keep the injected connector's identity rather than wrapping it per instance. @@ -399,9 +417,12 @@ export class PiNativeRuntime { this.vaultWatch?.close() this.closing ??= (async () => { await this.usagePublications.drain() + await this.drainRoutingPublications() + this.routingPublications.clear() await this.initialization.then( async (runtime) => { await runtime.drainQuotaPublications() + await runtime.drainRoutingPublications() runtime.close() }, () => {}, @@ -415,6 +436,7 @@ export function getPiNativeRuntime( storagePath = getPiAccountStoragePath(), options: { connect?: () => Promise + local?: NativeAccountRuntimeOptions['local'] vaultPoll?: PiVaultPollOptions } = {}, ): PiNativeRuntime { @@ -422,10 +444,14 @@ export function getPiNativeRuntime( const key = JSON.stringify([ contextKey, nativeRuntimeOptionIdentity(options.connect), + nativeRuntimeOptionIdentity(options.local?.onLockEvent), + nativeRuntimeOptionIdentity(options.local?.refreshToken), + nativeRuntimeOptionIdentity(options.local?.resolveIdentity), + nativeRuntimeOptionIdentity(options.local?.now), ]) const existing = runtimes.get(key) if (existing) return existing.runtime - if (!options.connect) { + if (!options.connect && !options.local) { // A supplied connector installs the context's test transport. Ordinary host // lookups reuse it, but an explicitly different connector owns another key. const configured = [...runtimes.values()].find( diff --git a/packages/pi/src/stream.ts b/packages/pi/src/stream.ts index 7967c3b6..21eaa37d 100644 --- a/packages/pi/src/stream.ts +++ b/packages/pi/src/stream.ts @@ -237,10 +237,16 @@ async function getPiStickyRouter(storagePath: string) { await getPiNativeRuntime(storagePath).routingStorage() let router = stickyRouters.get(path) if (!router) { - router = new StickySessionRouter({ path, vaultSeed }) + const runtime = await getPiNativeRuntime(storagePath).service() + router = new StickySessionRouter({ + path, + vaultSeed, + onLockEvent: runtime.onLockEvent, + }) setBoundedService(stickyRouters, path, router) } router.setVaultSeed(vaultSeed) + getPiNativeRuntime(storagePath).ownRoutingPublications(router) return router } diff --git a/packages/pi/src/tests/commands.test.ts b/packages/pi/src/tests/commands.test.ts index 6f4dac8b..b1e7067c 100644 --- a/packages/pi/src/tests/commands.test.ts +++ b/packages/pi/src/tests/commands.test.ts @@ -331,8 +331,10 @@ describe('native CacheKeep and routing commands', () => { retainAccountIds: new Set(['main']), storage: await readNativePiFixture(storagePath), }) + await router.drain() const before = await readNativePiFixture(storagePath) await controller().execute({ action: 'routing-reset' }) + await getPiNativeRuntime(storagePath).drainRoutingPublications() const cleared = await router.resolve({ sessionId: 'command-session', family: 'opus', diff --git a/packages/pi/src/tests/native-relay-lock.test.ts b/packages/pi/src/tests/native-relay-lock.test.ts new file mode 100644 index 00000000..9f8401ba --- /dev/null +++ b/packages/pi/src/tests/native-relay-lock.test.ts @@ -0,0 +1,149 @@ +import { expect } from 'bun:test' +import { mkdtemp, rm } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { + createNativeAccountRuntime, + createNativePoolStore, + resolveNativePoolPaths, +} from '@cortexkit/anthropic-auth-core' +import type { Api, Context, Model } from '@earendil-works/pi-ai' +import { nativeQuotaCodec } from '../../../core/src/native-quota-codec.ts' +import { holdNativeRelayLock } from '../../../core/src/tests/native-relay-lock-fixture.ts' +import { createTestLifetimeSuite } from '../../../core/src/tests/test-lifetime.ts' +import { MockRelayServer } from '../../../e2e-tests/src/mock-relay.ts' +import { closePiNativeRuntime } from '../native.ts' +import { streamCortexKitAnthropic } from '../stream.ts' +import { saveNativePiFixture } from './native-fixture.ts' +import { trackPiTestBody } from './setup.ts' + +const lifetimes = createTestLifetimeSuite() +const model: Model = { + id: 'claude-fable-5-1', + name: 'Fable', + provider: 'anthropic', + api: 'cortexkit-anthropic-messages', + baseUrl: 'https://api.anthropic.com', + reasoning: true, + input: ['text'], + cost: { input: 1, output: 1, cacheRead: 1, cacheWrite: 1 }, + contextWindow: 1_000_000, + maxTokens: 128_000, +} +const context: Context = { + messages: [{ role: 'user', content: 'hello', timestamp: 0 }], + tools: [], +} +for (const scenario of ['direct', 'relay', 'inconsistent'] as const) { + const configured = scenario !== 'direct' + lifetimes.test( + scenario === 'inconsistent' + ? 'Pi inconsistent relay pair refuses locally without any dispatch' + : `Pi request ignores held native-relay lock (${scenario})`, + () => + trackPiTestBody( + (async () => { + const directory = await mkdtemp(join(tmpdir(), 'pi-held-relay-')) + const storagePath = join(directory, 'anthropic-auth.json') + const originalFetch = globalThis.fetch + const saved = { + PI_CODING_AGENT_DIR: process.env.PI_CODING_AGENT_DIR, + PI_ANTHROPIC_AUTH_FILE: process.env.PI_ANTHROPIC_AUTH_FILE, + OPENCODE_ANTHROPIC_AUTH_STATE_FILE: + process.env.OPENCODE_ANTHROPIC_AUTH_STATE_FILE, + } + const relay = new MockRelayServer() + lifetimes.deferCleanup(async () => { + await closePiNativeRuntime(storagePath) + relay.stop() + globalThis.fetch = originalFetch + for (const [key, value] of Object.entries(saved)) { + if (value === undefined) delete process.env[key] + else process.env[key] = value + } + await rm(directory, { recursive: true, force: true }) + }) + process.env.PI_CODING_AGENT_DIR = directory + process.env.PI_ANTHROPIC_AUTH_FILE = storagePath + await saveNativePiFixture( + { version: 1, accounts: [], quota: { enabled: false } }, + storagePath, + { + mainAccess: 'sk-ant-oat01-synthetic-held-access', + }, + ) + const paths = await resolveNativePoolPaths(storagePath) + const runtime = createNativeAccountRuntime({ paths, host: 'pi' }) + lifetimes.deferCleanup(() => runtime.close()) + const address = await relay.start({ + token: 'synthetic-held-relay-token', + }) + if (configured) + await runtime.updateRelay({ + enabled: true, + url: address.url, + token: 'synthetic-held-relay-token', + transport: 'http', + fallbackToDirect: false, + }) + let upstream = 0 + globalThis.fetch = Object.assign( + async (input: Parameters[0], init?: RequestInit) => { + const url = + input instanceof Request ? input.url : input.toString() + if (url.startsWith(address.url)) return originalFetch(input, init) + if (!url.includes('/v1/messages')) return Response.json({}) + upstream++ + return new Response( + 'event: message_start\ndata: {"type":"message_start","message":{"usage":{"input_tokens":1,"output_tokens":0}}}\n\nevent: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":1}}\n\nevent: message_stop\ndata: {"type":"message_stop"}\n\n', + ) + }, + { preconnect: originalFetch.preconnect }, + ) + if (scenario === 'inconsistent') + await createNativePoolStore({ + paths, + quota: nativeQuotaCodec, + }).updateSettings((settings) => { + ;(settings.relay as Record).generation = + 'mismatched' + }) + const held = await holdNativeRelayLock(paths.config, lifetimes) + try { + const result = await streamCortexKitAnthropic(model, context, { + sessionId: `held-relay-${scenario}-${directory}`, + }).result() + if (scenario === 'inconsistent') { + expect(result.stopReason).toBe('error') + expect(result.errorMessage).toBe( + 'request refused: relay configuration is inconsistent; no request was sent', + ) + expect(relay.acceptedRequests()).toBe(0) + expect(upstream).toBe(0) + return + } + expect(result.stopReason).toBe('stop') + expect(relay.acceptedRequests()).toBe(configured ? 1 : 0) + // The loopback relay forwards through the same mocked fetch. + expect(upstream).toBe(1) + expect(await runtime.getRelayConfig()).toEqual( + configured + ? { + enabled: true, + url: address.url, + token: 'synthetic-held-relay-token', + transport: 'http', + fallbackToDirect: false, + } + : null, + ) + await held.assertOwned() + } finally { + await held.pastTimeout + await held.release() + } + })(), + ), + 25_000, + ) +} diff --git a/packages/pi/src/tests/native-request-locks.test.ts b/packages/pi/src/tests/native-request-locks.test.ts new file mode 100644 index 00000000..f483446b --- /dev/null +++ b/packages/pi/src/tests/native-request-locks.test.ts @@ -0,0 +1,214 @@ +import { expect } from 'bun:test' +import { mkdtemp, rm, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { + getClaudeCodeIdentityForVerifiedAccount, + type NativeCustodyClient, +} from '@cortexkit/anthropic-auth-core' +import type { Api, Context, Model } from '@earendil-works/pi-ai' +import { requestLockProbe } from '../../../core/src/tests/native-request-lock-fixture.ts' +import { createTestLifetimeSuite } from '../../../core/src/tests/test-lifetime.ts' +import { closePiNativeRuntime, getPiNativeRuntime } from '../native.ts' +import { streamCortexKitAnthropic } from '../stream.ts' +import { + fixtureAccountIdentity, + saveNativePiFixture, +} from './native-fixture.ts' +import { trackPiTestBody } from './setup.ts' + +const lifetimes = createTestLifetimeSuite() +const model: Model = { + id: 'claude-fable-5-1', + name: 'Fable', + provider: 'anthropic', + api: 'cortexkit-anthropic-messages', + baseUrl: 'https://api.anthropic.com', + reasoning: true, + input: ['text'], + cost: { input: 1, output: 1, cacheRead: 1, cacheWrite: 1 }, + contextWindow: 1_000_000, + maxTokens: 128_000, +} +const context: Context = { + messages: [{ role: 'user', content: 'hello', timestamp: 0 }], + tools: [], +} +const quota = () => ({ + checkedAt: Date.now(), + five_hour: { + checkedAt: Date.now(), + remainingPercent: 100, + usedPercent: 0, + resetsAt: new Date(Date.now() + 3600_000).toISOString(), + }, + seven_day: { + checkedAt: Date.now(), + remainingPercent: 100, + usedPercent: 0, + resetsAt: new Date(Date.now() + 86_400_000).toISOString(), + }, +}) +for (const scenario of ['local-healthy', 'vault-healthy', 'refresh'] as const) { + lifetimes.test( + `Pi ${scenario} request observes only required foreground locks`, + () => + trackPiTestBody( + (async () => { + const directory = await mkdtemp(join(tmpdir(), 'pi-request-locks-')) + const storagePath = join(directory, 'anthropic-auth.json') + const originalFetch = globalThis.fetch + const keys = [ + 'PI_CODING_AGENT_DIR', + 'PI_ANTHROPIC_AUTH_FILE', + 'PI_ANTHROPIC_AUTH_CLAUSTRUM_ENROLLMENT_FILE', + 'OPENCODE_ANTHROPIC_AUTH_STATE_FILE', + ] as const + const saved = new Map(keys.map((key) => [key, process.env[key]])) + lifetimes.deferCleanup(async () => { + await closePiNativeRuntime(storagePath) + globalThis.fetch = originalFetch + for (const [key, value] of saved) { + if (value === undefined) delete process.env[key] + else process.env[key] = value + } + await rm(directory, { recursive: true, force: true }) + }) + process.env.PI_CODING_AGENT_DIR = directory + process.env.PI_ANTHROPIC_AUTH_FILE = storagePath + await saveNativePiFixture( + { + version: 1, + accounts: [], + routing: { mode: 'sticky-balanced' }, + quota: { enabled: false, mainQuota: quota() }, + }, + storagePath, + ) + const probe = requestLockProbe() + const identity = fixtureAccountIdentity('main') + const client: NativeCustodyClient = { + async listScoped() { + return { + view: 'ready-lock-proof', + rows: [ + { + id: 'oauth:anthropic', + accountId: identity, + credentialType: 'oauth', + categories: ['anthropic-native'], + serves: ['anthropic'], + providerIds: [], + refreshAdapter: 'anthropic', + operations: ['read'], + state: 'active', + recordVersion: 1, + createdAtMs: 1, + }, + ], + grants: undefined, + grantTuples: undefined, + } + }, + async getScoped() { + return { + credentialId: 'oauth:anthropic', + accountId: identity, + material: 'sk-ant-oat01-ready-vault', + recordVersion: 1, + expiresAtMs: Date.now() + 3600_000, + } + }, + async reportAuthFailureScoped() {}, + close() {}, + } + if (scenario === 'vault-healthy') { + const enrollment = join(directory, 'enrollment.json') + process.env.PI_ANTHROPIC_AUTH_CLAUSTRUM_ENROLLMENT_FILE = enrollment + await writeFile( + enrollment, + JSON.stringify({ token: '01'.repeat(32), token_generation: 1 }), + { mode: 0o600 }, + ) + } + const now = () => Date.now() + 2 * 3600_000 + const native = getPiNativeRuntime(storagePath, { + vaultPoll: { intervalMs: 0 }, + ...(scenario === 'vault-healthy' + ? { connect: async () => client } + : {}), + local: { + onLockEvent: probe.onLockEvent, + ...(scenario === 'refresh' + ? { + now, + refreshToken: async () => ({ + access: 'sk-ant-oat01-refreshed-control', + refresh: 'synthetic-control-refresh', + expires: now() + 3600_000, + expiresIn: 3600, + }), + resolveIdentity: async () => + getClaudeCodeIdentityForVerifiedAccount('main', identity), + } + : {}), + }, + }) + const runtime = await native.service() + if (scenario === 'vault-healthy') { + await runtime.vault.refresh() + await runtime.updateSettings((settings) => { + settings.routing = { mode: 'sticky-balanced' } + settings.quota = { enabled: false } + }) + const receipt = await runtime.authorizeVault('main') + await runtime.vault.publish(receipt, { + quota: { ...quota(), accountIdentity: identity }, + }) + } + await runtime.read() + await runtime.drainRoutingPublications() + globalThis.fetch = Object.assign( + async (input: Parameters[0]) => { + const url = + input instanceof Request ? input.url : input.toString() + if (!url.includes('/v1/messages')) return Response.json({}) + probe.dispatch() + return new Response( + 'event: message_start\ndata: {"type":"message_start","message":{"usage":{"input_tokens":1,"output_tokens":0}}}\n\nevent: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":1}}\n\nevent: message_stop\ndata: {"type":"message_stop"}\n\n', + ) + }, + { preconnect: originalFetch.preconnect }, + ) + probe.begin() + const result = await streamCortexKitAnthropic(model, context, { + sessionId: `lock-proof-${directory}`, + }).result() + expect(result.stopReason).toBe('stop') + if (scenario === 'refresh') + expect( + probe.foreground.some((event) => + event.name.startsWith('native-account-provider-'), + ), + ).toBe(true) + else expect(probe.foreground).toEqual([]) + await closePiNativeRuntime(storagePath) + expect(probe.detached.some((event) => event.name === 'write')).toBe( + true, + ) + expect(probe.afterDispatch.length).toBeGreaterThan(0) + expect( + probe.afterDispatch.some( + (event) => event.name === 'native-runtime', + ), + ).toBe(true) + if (scenario === 'vault-healthy') + expect( + probe.afterDispatch.some( + (event) => event.name === 'claustrum-roster-write', + ), + ).toBe(true) + })(), + ), + ) +}