From 0a2aa6dd34447b6ace5e226acf58aa88ef92974e Mon Sep 17 00:00:00 2001 From: ualtinok <94532+ualtinok@users.noreply.github.com> Date: Thu, 12 Nov 2026 16:46:21 +0100 Subject: [PATCH 1/3] mason: route vault API-key attempts through Thalamus Adopt common-auth 0.47 opt-in routing refusals for API-key attempts on all hosts, preserving the plugin's usual OAuth subscription path. Refusals and missing answers stop paid requests; qualifying balance refusals report the selected vault identities. Carry OC1 host session ids through a private stripped header. Avoid consumer.send's internal paid-key 401 resend without routing admission. Add filtered unit and real-host no-send regressions; preserve xAI OAuth billing expectations. Pi real-host assertions include its existing SDK 400 JSON envelope around the fixed text. --- bun.lock | 8 +- docs/setup.md | 17 + packages/core/package.json | 2 +- packages/core/src/fetch.ts | 183 +++++++--- packages/core/src/index.ts | 9 + packages/core/src/placement.ts | 2 + packages/core/src/routing.ts | 107 ++++++ packages/core/src/serving.ts | 10 +- packages/core/test/billing-fixtures.ts | 8 +- packages/core/test/billing.test.ts | 5 + packages/core/test/routing-fixtures.ts | 24 ++ packages/core/test/routing-preload.ts | 10 + packages/core/test/routing.test.ts | 342 ++++++++++++++++++ packages/opencode/package.json | 2 +- packages/opencode/src/v1/fetch.ts | 41 ++- packages/opencode/src/v1/messages.ts | 4 +- packages/opencode/src/v1/plugin.ts | 5 + packages/opencode/src/v2/adapter.ts | 83 ++++- packages/opencode/src/v2/messages.ts | 7 +- packages/opencode/test/v1-plugin.test.ts | 14 + packages/opencode/test/v1-routing.test.ts | 56 +++ .../opencode/test/v2/e2e/billing-preload.ts | 4 + .../test/v2/e2e/opencode2.e2e.test.ts | 25 ++ packages/opencode/test/v2/routing.test.ts | 207 +++++++++++ packages/pi/package.json | 2 +- packages/pi/src/extension.ts | 49 ++- packages/pi/src/messages.ts | 8 +- packages/pi/test-entry.ts | 1 + packages/pi/test/billing.test.ts | 44 ++- packages/pi/test/extension.e2e.test.ts | 24 ++ packages/pi/test/routing-retry.test.ts | 23 ++ 31 files changed, 1217 insertions(+), 109 deletions(-) create mode 100644 packages/core/src/routing.ts create mode 100644 packages/core/test/routing-fixtures.ts create mode 100644 packages/core/test/routing-preload.ts create mode 100644 packages/core/test/routing.test.ts create mode 100644 packages/opencode/test/v1-routing.test.ts create mode 100644 packages/opencode/test/v2/routing.test.ts create mode 100644 packages/pi/test/routing-retry.test.ts diff --git a/bun.lock b/bun.lock index 6ed3e82..28c41b3 100644 --- a/bun.lock +++ b/bun.lock @@ -20,7 +20,7 @@ "version": "0.0.0", "dependencies": { "@cortexkit/claustrum-client": "^0.18.0", - "@cortexkit/common-auth": "^0.44.0", + "@cortexkit/common-auth": "^0.47.0", "@cortexkit/private-files": "^0.2.2", "@cortexkit/subc-client": "^0.24.0", }, @@ -33,7 +33,7 @@ }, "dependencies": { "@cortexkit/claustrum-client": "^0.18.0", - "@cortexkit/common-auth": "^0.44.0", + "@cortexkit/common-auth": "^0.47.0", "@cortexkit/private-files": "^0.2.2", "@cortexkit/subc-client": "^0.24.0", }, @@ -54,7 +54,7 @@ "version": "0.3.0", "dependencies": { "@cortexkit/claustrum-client": "^0.18.0", - "@cortexkit/common-auth": "^0.44.0", + "@cortexkit/common-auth": "^0.47.0", "@cortexkit/private-files": "^0.2.2", "@cortexkit/subc-client": "^0.24.0", }, @@ -154,7 +154,7 @@ "@cortexkit/claustrum-client": ["@cortexkit/claustrum-client@0.18.0", "", { "dependencies": { "@cortexkit/store": "^0.2.3", "@cortexkit/subc-client": "^0.24.0" } }, "sha512-tdRTrtnmOk3En1ZT1T+J+8NDwZl5FNxtN4w4Im2KV12QypPwyMi1ZBUXS2X4ZyDn4vPbHiOK8gjP+hQMsZ/V9g=="], - "@cortexkit/common-auth": ["@cortexkit/common-auth@0.44.0", "", { "dependencies": { "@cortexkit/private-files": "^0.2.2", "jsonc-parser": "^3.3.1" }, "peerDependencies": { "@cortexkit/claustrum-client": ">=0.6.2", "@cortexkit/opencode-sidebar": "~0.1.0", "@cortexkit/subc-client": "^0.16.0 || ^0.22.0 || ^0.24.0", "@opencode/plugin": ">=2.0.21", "@opentui/core": ">=0.4.5", "@opentui/solid": ">=0.4.5", "solid-js": ">=1.9.12 <2" }, "optionalPeers": ["@cortexkit/claustrum-client", "@cortexkit/opencode-sidebar", "@cortexkit/subc-client", "@opencode/plugin", "@opentui/core", "@opentui/solid", "solid-js"] }, "sha512-2MGgT3DpyrwuxhbXw7XXfbxQJZR2X2eznPW8KP7GRtHPEdFngx4JY+nAdPAygy5kZ7WpOUWTR1J2kFwLeOK/lg=="], + "@cortexkit/common-auth": ["@cortexkit/common-auth@0.47.0", "", { "dependencies": { "@cortexkit/private-files": "^0.2.2", "jsonc-parser": "^3.3.1" }, "peerDependencies": { "@cortexkit/claustrum-client": ">=0.6.2", "@cortexkit/opencode-sidebar": "~0.1.0", "@cortexkit/subc-client": "^0.16.0 || ^0.22.0 || ^0.24.0", "@opencode/plugin": ">=2.0.21", "@opentui/core": ">=0.4.5", "@opentui/solid": ">=0.4.5", "solid-js": ">=1.9.12 <2" }, "optionalPeers": ["@cortexkit/claustrum-client", "@cortexkit/opencode-sidebar", "@cortexkit/subc-client", "@opencode/plugin", "@opentui/core", "@opentui/solid", "solid-js"] }, "sha512-fHM1hhVcaIS0m9rmX528zF8euEVm5ac1WtDjbjW/vRdlYZs0kbssg91C9paOLdOvHc8Hrr1ijGXYtSY5LAdspQ=="], "@cortexkit/opencode-claustrum-auth": ["@cortexkit/opencode-claustrum-auth@workspace:packages/opencode"], diff --git a/docs/setup.md b/docs/setup.md index 4aec7eb..e6d8865 100644 --- a/docs/setup.md +++ b/docs/setup.md @@ -98,3 +98,20 @@ vault reachability, provider ownership and the accounts available for routing. Use `/claustrum refusals` for rows the plugin cannot serve. See [provider metadata](providers.md#the-operators-vault-rows) if a row is refused despite a successful grant. + +## Thalamus routing + +When Thalamus supplies its routing accessor and the plugin's next row is an +API key, routing chooses that request's vault API key, including for titles, +compaction and retries, from the enabled +accounts for that provider. The plugin never substitutes its own choice. A +routing refusal stops the request with "Thalamus routing declined this request, +so it was not sent." No answer stops it with "Thalamus routing did not answer, +so this API key request was not sent." Slow routing (over 1.5 seconds) counts +as no answer and stops the API-key request too. Documented empty-balance +refusals are reported back to routing; user-set spending caps are not. +Outside Thalamus, without its accessor, account selection is unchanged. + +Subscription (OAuth) rows are not routed yet. If the plugin's usual next row +is an OAuth login, it serves that login as before without asking routing; only +an API-key attempt asks routing to choose among eligible API keys. diff --git a/packages/core/package.json b/packages/core/package.json index c3bda07..6fee4f3 100644 --- a/packages/core/package.json +++ b/packages/core/package.json @@ -13,7 +13,7 @@ }, "dependencies": { "@cortexkit/claustrum-client": "^0.18.0", - "@cortexkit/common-auth": "^0.44.0", + "@cortexkit/common-auth": "^0.47.0", "@cortexkit/private-files": "^0.2.2", "@cortexkit/subc-client": "^0.24.0" } diff --git a/packages/core/src/fetch.ts b/packages/core/src/fetch.ts index c9c9510..379d547 100644 --- a/packages/core/src/fetch.ts +++ b/packages/core/src/fetch.ts @@ -19,6 +19,7 @@ import { redactError, redactResponse, } from './placement.js' +import { reportRoutedBalance, selectRoutedKey } from './routing.js' import { placeReceipt, type ServingOptions } from './serving.js' export const maxBodyBytes = 64 * 1024 * 1024 @@ -33,6 +34,7 @@ export interface FetchOptions extends ServingOptions { // case rules and config precedence. The check is repeated on every request. effectiveCredential?: () => boolean | Promise billing?: (receipt: ClaustrumScopedAttempt) => BillingAttempt | undefined + routing?: { sessionId: string | undefined; provider: string; modelId: string } } async function bufferBody(request: Request): Promise { @@ -101,6 +103,7 @@ export function createFetch(options: FetchOptions): CoreFetch { cache: request.cache, } const tried = new Set() + let routedAttempt = false const rejected = new Set() const secrets: string[] = [] let redirectBilling: @@ -191,12 +194,10 @@ export function createFetch(options: FetchOptions): CoreFetch { try { for (;;) { signal.throwIfAborted() + if (routedAttempt && pending && !pending.parsed) + return await redactResponse(pending.response, secrets) const rows = await candidates() - if (!pending && !rows.some((row) => row.admitted)) - throw new ClaustrumAuthError('no-vault-account') - // A failed provider response stays readable until a replacement send - // is committed. Reloads must not turn it into a fresh-request refusal. - const initial = pending + const ownDecision = pending ? decide( 'after-limit', rows, @@ -205,6 +206,39 @@ export function createFetch(options: FetchOptions): CoreFetch { pending.parsed, ) : decide('pre-send', rows) + if (ownDecision.action === 'pass' && pending) + return await redactResponse(pending.response, secrets) + const nextRoute = + ownDecision.action === 'reroute' + ? ownDecision.account.accountId + : rows.find((row) => row.admitted && row.placementEligible)?.account + .accountId + let routed = await selectRoutedKey( + options.consumer, + options.routing?.provider ?? options.entry.service, + options.routing?.modelId ?? '', + options.routing?.sessionId, + new Set([...tried, ...rejected]), + nextRoute, + ) + routedAttempt = routed !== undefined + const routedId = routed?.routeId + if (!pending && !rows.some((row) => row.admitted)) + throw new ClaustrumAuthError('no-vault-account') + // Keep a failed provider response readable until a replacement send + // is committed, so a vault roster reload cannot turn a retry into a + // fresh-request refusal instead of returning the provider's response. + const initial = routed + ? { + action: 'reroute' as const, + account: rows.find((row) => row.account.accountId === routedId) + ?.account ?? { + namespace: 'claustrum-auth:routing', + accountId: routed.routeId, + incarnation: 'routing', + }, + } + : ownDecision if (initial.action === 'pass' && pending) return await redactResponse(pending.response, secrets) if (initial.action === 'refuse') { @@ -213,12 +247,12 @@ export function createFetch(options: FetchOptions): CoreFetch { } if (initial.action !== 'reroute') throw new Error('claustrum-auth: unsupported limit action') - const route = initial.account.accountId - if (limits && !(await limits.beforeDispatch(route))) { + let route = initial.account.accountId + if (!routed && limits && !(await limits.beforeDispatch(route))) { rejected.add(route) continue } - if (limits) { + if (!routed && limits) { if ( !(await candidates()).find((row) => row.account.accountId === route) ?.ready @@ -229,7 +263,7 @@ export function createFetch(options: FetchOptions): CoreFetch { pending = undefined tried.add(route) const previousMark = rateLimitMarks.get(route) - const observation = limits?.observe(route) + let observation = limits?.observe(route) let target = url let method = template.method let bytes = body @@ -238,54 +272,76 @@ export function createFetch(options: FetchOptions): CoreFetch { for (;;) { signal.throwIfAborted() checkOrigin(target, options.entry) - response = await options.consumer.send( - route, - async (receipt, dispatchSignal) => { - secrets.push(receipt.accessToken) - const draft = placeReceipt( - target, - headers, - options.entry, - receipt, - ) - const dispatchHeaders = new Headers(draft.headers) - if (!bytes) { - dispatchHeaders.delete('content-length') - if (method !== template.method) - dispatchHeaders.delete('content-type') - } - const outbound = new Request(draft.url, { - method, - headers: dispatchHeaders, - body: bytes ? new Uint8Array(bytes) : undefined, - signal: dispatchSignal ?? signal, - redirect: 'manual', - credentials: template.credentials, - cache: template.cache, - }) - let billing: BillingAttempt | undefined - try { - const reply = await options.fetch(outbound) - if ( - options.billing && - [301, 302, 303, 307, 308].includes(reply.status) && - reply.headers.has('location') - ) { - // A followed redirect is a routing hop, not a billable model - // call. Wait until the redirect checks decide to return it. - redirectBilling = { receipt, status: reply.status } - return reply - } - billing = options.billing?.(receipt) - return observeBillingResponse(billing, reply) - } catch (error) { - const failedBilling = billing ?? options.billing?.(receipt) - failedBilling?.end(false, false) - throw error + if (redirects > 0 && routed) { + routed = await selectRoutedKey( + options.consumer, + options.routing?.provider ?? options.entry.service, + options.routing?.modelId ?? '', + options.routing?.sessionId, + new Set([...tried].filter((id) => id !== route)), + ) + if (!routed) throw new ClaustrumAuthError('routing-no-answer') + route = routed.routeId + tried.add(route) + observation = limits?.observe(route) + } + const dispatch = async ( + receipt: ClaustrumScopedAttempt, + dispatchSignal?: AbortSignal, + ) => { + secrets.push(receipt.accessToken) + const draft = placeReceipt(target, headers, options.entry, receipt) + const dispatchHeaders = new Headers(draft.headers) + if (!bytes) { + dispatchHeaders.delete('content-length') + if (method !== template.method) + dispatchHeaders.delete('content-type') + } + const outbound = new Request(draft.url, { + method, + headers: dispatchHeaders, + body: bytes ? new Uint8Array(bytes) : undefined, + signal: dispatchSignal ?? signal, + redirect: 'manual', + credentials: template.credentials, + cache: template.cache, + }) + let billing: BillingAttempt | undefined + try { + const reply = await options.fetch(outbound) + if ( + options.billing && + [301, 302, 303, 307, 308].includes(reply.status) && + reply.headers.has('location') + ) { + // A followed redirect is a routing hop, not a billable model + // call. Wait until the redirect checks decide to return it. + redirectBilling = { receipt, status: reply.status } + return reply } - }, - { site: 'claustrum-auth', signal }, - ) + billing = options.billing?.(receipt) + return observeBillingResponse(billing, reply) + } catch (error) { + const failedBilling = billing ?? options.billing?.(receipt) + failedBilling?.end(false, false) + throw error + } + } + if (routed) { + // Consumer.send can rotate and resend a 401 on its own. Routing + // owns paid-key admission, so authorize exactly one physical send. + const receipt = await options.consumer.authorize(route, signal) + if (receipt.credentialId !== routed.credentialId) + throw new ClaustrumAuthError('routing-no-answer') + response = await dispatch(receipt, signal) + if (response.status === 401) + await options.consumer.reportFailure(receipt, 401, 'direct') + } else { + response = await options.consumer.send(route, dispatch, { + site: 'claustrum-auth', + signal, + }) + } signal.throwIfAborted() const location = response.headers.get('location') if (![301, 302, 303, 307, 308].includes(response.status) || !location) @@ -322,8 +378,12 @@ export function createFetch(options: FetchOptions): CoreFetch { redirectBilling = undefined response = observeBillingResponse(billing, response) } + const account = observation?.binding.account ?? { + ...initial.account, + accountId: route, + } const parsed = await limitFromReply({ - account: observation?.binding.account ?? initial.account, + account, service: options.entry.service, status: response.status, headers: response.headers, @@ -331,6 +391,13 @@ export function createFetch(options: FetchOptions): CoreFetch { now: now(), }) const balance402 = parsed?.kind === 'payment-balance' + await reportRoutedBalance( + routed, + options.routing?.sessionId, + options.entry.service, + response.status, + () => response.clone().text(), + ) if (response.ok || (response.status === 402 && !balance402)) { if (limits && observation) { if ( @@ -360,7 +427,7 @@ export function createFetch(options: FetchOptions): CoreFetch { } pending = { response, - account: observation?.binding.account ?? initial.account, + account, parsed, } } diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index a592342..71e4024 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -57,6 +57,15 @@ export { reservedProviders, supportedProviders, } from './providers.js' +export type { RoutedKey } from './routing.js' +export { + hasVaultRoutingAccessor, + reportRoutedBalance, + routingDeclinedText, + routingNoAnswerText, + routingRefusalText, + selectRoutedKey, +} from './routing.js' export type { CoreOptions, RosterCore } from './runtime.js' export { createCore, unsavedRosterCause } from './runtime.js' export type { ServingConsumer, ServingOptions } from './serving.js' diff --git a/packages/core/src/placement.ts b/packages/core/src/placement.ts index 4e947d5..c7fd7a3 100644 --- a/packages/core/src/placement.ts +++ b/packages/core/src/placement.ts @@ -17,6 +17,8 @@ export type Refusal = | 'route-gone' | 'vault-mode-off' | 'consent-required' + | 'routing-declined' + | 'routing-no-answer' export class ClaustrumAuthError extends Error { constructor(readonly code: Refusal) { diff --git a/packages/core/src/routing.ts b/packages/core/src/routing.ts new file mode 100644 index 0000000..327d307 --- /dev/null +++ b/packages/core/src/routing.ts @@ -0,0 +1,107 @@ +import { + askRouting, + reportRouting, +} from '@cortexkit/common-auth/routing-accessor' +import { isBalanceFailure } from './balance.js' +import { ClaustrumAuthError } from './placement.js' +import type { ServingConsumer } from './serving.js' + +export const routingDeclinedText = + 'Thalamus routing declined this request, so it was not sent.' +export const routingNoAnswerText = + 'Thalamus routing did not answer, so this API key request was not sent.' + +export function routingRefusalText(code: string): string | undefined { + if (code === 'routing-declined') return routingDeclinedText + if (code === 'routing-no-answer') return routingNoAnswerText + return undefined +} + +// Thalamus may install or release its bridge after the plugin starts. +export function hasVaultRoutingAccessor(): boolean { + try { + return ( + typeof Reflect.get(globalThis, Symbol.for('cortexkit.routing')) === + 'function' + ) + } catch { + return false + } +} + +export interface RoutedKey { + routeId: string + accountId: string + credentialId: string +} + +export async function selectRoutedKey( + consumer: ServingConsumer, + provider: string, + modelId: string, + sessionId: string | undefined, + excluded: ReadonlySet = new Set(), + nextRoute?: string, +): Promise { + if (!hasVaultRoutingAccessor()) return undefined + const routes = consumer.routingRows() + const roster = consumer.snapshot?.()?.rows ?? [] + const next = nextRoute ?? routes.find((row) => !excluded.has(row.id))?.id + // Subscription rows keep the plugin's existing order and retry behavior; + // routing has not adopted those vault logins yet. + if (roster.find((row) => row.routeId === next)?.credentialType === 'oauth') + return undefined + if (!sessionId) throw new ClaustrumAuthError('routing-no-answer') + // routingRows applies the user's on/off switches. Snapshot supplies vault + // identities, never the plugin's synthetic route id or a key version. + const enabled = new Set(routes.map((row) => row.id)) + const candidates = roster + .filter( + (row) => + enabled.has(row.routeId) && + !excluded.has(row.routeId) && + row.credentialType === 'api_key', + ) + .map((row) => ({ + routeId: row.routeId, + credentialId: row.credentialId, + accountId: row.accountIdentity ?? row.credentialId, + })) + const answer = await askRouting({ + sessionId, + provider, + modelId, + candidates, + refusals: true, + }) + if (!answer) throw new ClaustrumAuthError('routing-no-answer') + if ('refused' in answer) throw new ClaustrumAuthError('routing-declined') + return answer.candidate +} + +export async function reportRoutedBalance( + key: RoutedKey | undefined, + sessionId: string | undefined, + service: string, + status: number, + body: () => Promise, +): Promise { + if (!key || !sessionId || status !== 402) return + const text = await body().catch(() => '') + if (!isBalanceFailure(service, text)) return + let providerCode = + service === 'openrouter' ? 'openrouter_credits' : 'http_402' + if (service === 'deepseek') { + try { + const code: unknown = JSON.parse(text)?.error?.code + if (typeof code === 'string') providerCode = code + } catch {} + } + reportRouting({ + sessionId, + accountId: key.accountId, + credentialId: key.credentialId, + outcome: 'credential_refused', + providerCode, + }) +} diff --git a/packages/core/src/serving.ts b/packages/core/src/serving.ts index f9673b3..f728f47 100644 --- a/packages/core/src/serving.ts +++ b/packages/core/src/serving.ts @@ -16,7 +16,8 @@ import type { ProviderEntry } from './providers.js' export type ServingConsumer = Pick< ClaustrumConsumer, 'routingRows' | 'authorize' | 'send' | 'reportFailure' -> +> & + Partial> export interface ServingOptions { entry: ProviderEntry placeholder: string @@ -56,7 +57,10 @@ export function placeReceipt( // Host-independent callbacks for the OpenCode 2 adapter. The returned holder // retains the issuing receipt, so a response cannot be attributed to a row // merely because it shares a provider or session with an owned request. -export function servingCallbacks(options: ServingOptions) { +export function servingCallbacks( + options: ServingOptions, + routingAdmitted = false, +) { return { async request(url: URL, headers: Headers, signal?: AbortSignal) { if ( @@ -67,7 +71,7 @@ export function servingCallbacks(options: ServingOptions) { const route = orderedRoutes(options)[0] if (!route) throw new ClaustrumAuthError('no-vault-account') const limits = await options.limits?.() - if (limits && !(await limits.beforeDispatch(route))) + if (!routingAdmitted && limits && !(await limits.beforeDispatch(route))) throw new ClaustrumAuthError('no-vault-account') const observation = limits?.observe(route) const receipt = await options.consumer.authorize(route, signal) diff --git a/packages/core/test/billing-fixtures.ts b/packages/core/test/billing-fixtures.ts index ee7f976..12d99f3 100644 --- a/packages/core/test/billing-fixtures.ts +++ b/packages/core/test/billing-fixtures.ts @@ -7,7 +7,12 @@ afterEach(() => { else Reflect.deleteProperty(globalThis, symbol) }) -export function captureBilling(keepWarm = false) { +export function captureBilling( + keepWarm = false, + accountFor = async ({ provider }: { provider: string }) => ({ + account_id: provider === 'fixture-a' ? 'account-a' : 'account-1', + }), +) { const reports: Record[] = [] const sessions: string[] = [] Object.defineProperty(globalThis, symbol, { @@ -15,6 +20,7 @@ export function captureBilling(keepWarm = false) { value: (session: string) => { sessions.push(session) return { + accountFor, report: (report: Record) => { reports.push(report) }, diff --git a/packages/core/test/billing.test.ts b/packages/core/test/billing.test.ts index 0d185dd..9066047 100644 --- a/packages/core/test/billing.test.ts +++ b/packages/core/test/billing.test.ts @@ -209,6 +209,11 @@ test.each(['terminal-redirect', 'transport-error', 'redirect-limit'] as const)( entry, placeholder, consumer: h.consumer, + routing: { + sessionId: 'session', + provider: 'fixture-a', + modelId: 'sent-model', + }, billing: (receipt) => new BillingAttempt( entry, diff --git a/packages/core/test/routing-fixtures.ts b/packages/core/test/routing-fixtures.ts new file mode 100644 index 0000000..ef3da59 --- /dev/null +++ b/packages/core/test/routing-fixtures.ts @@ -0,0 +1,24 @@ +import { onTestFinished } from 'bun:test' + +const symbol = Symbol.for('cortexkit.routing') +const original = Object.getOwnPropertyDescriptor(globalThis, symbol) + +export function routing(answer: (call: number) => Promise) { + onTestFinished(() => { + if (original) Object.defineProperty(globalThis, symbol, original) + else Reflect.deleteProperty(globalThis, symbol) + }) + const calls: unknown[] = [] + const reports: unknown[] = [] + Object.defineProperty(globalThis, symbol, { + configurable: true, + value: (session: string) => ({ + accountFor: (request: unknown, options: unknown) => { + calls.push({ session, request, options }) + return answer(calls.length) + }, + report: (payload: unknown) => reports.push(payload), + }), + }) + return { calls, reports } +} diff --git a/packages/core/test/routing-preload.ts b/packages/core/test/routing-preload.ts new file mode 100644 index 0000000..8b5c567 --- /dev/null +++ b/packages/core/test/routing-preload.ts @@ -0,0 +1,10 @@ +// Runs inside the isolated real host, before the plugin is initialized. +const mode = process.env.CLAUSTRUM_AUTH_E2E_ROUTING +if (mode) { + Reflect.set(globalThis, Symbol.for('cortexkit.routing'), () => ({ + accountFor: async () => + mode === 'refuse' + ? { refused: { code: 'account_selection_refused' } } + : new Promise(() => {}), + })) +} diff --git a/packages/core/test/routing.test.ts b/packages/core/test/routing.test.ts new file mode 100644 index 0000000..2205deb --- /dev/null +++ b/packages/core/test/routing.test.ts @@ -0,0 +1,342 @@ +import { expect, test } from 'bun:test' +import { createFetch, fixtureEntries } from '../test-entry.ts' +import { bearer, harness, placeholder, row } from './fakes.ts' +import { routing } from './routing-fixtures.ts' + +const symbol = Symbol.for('cortexkit.routing') + +for (const host of ['opencode1', 'pi'] as const) { + test(`routing ${host}: OAuth-only subscriptions serve without asking routing`, async () => { + const h = await harness( + [ + row('oauth', { + credentialType: 'oauth', + authMethod: 'oauth', + refreshAdapter: 'fixture-a', + }), + ], + host, + ) + const { calls } = routing(async () => { + throw new Error('OAuth must not ask') + }) + const serve = createFetch({ + entry: fixtureEntries[0], + placeholder, + consumer: h.consumer, + host, + fetch: async () => new Response('subscription'), + }) + expect( + await ( + await serve('https://fixture.invalid', { headers: bearer }) + ).text(), + ).toBe('subscription') + expect(h.reads).toEqual(['oauth']) + expect(calls).toEqual([]) + }) + + test(`routing ${host}: mixed pools ask only when the usual next row is an API key`, async () => { + const h = await harness( + [ + row('oauth', { + credentialType: 'oauth', + authMethod: 'oauth', + refreshAdapter: 'fixture-a', + }), + row('apikey'), + ], + host, + ) + const sent: string[] = [] + const { calls } = routing(async () => ({ account_id: 'account-apikey' })) + const order = ['oauth', 'apikey'].map( + (id) => + h.consumer.snapshot()?.rows.find((row) => row.credentialId === id) + ?.routeId as string, + ) + await h.consumer.reorder(order) + const serve = createFetch({ + entry: fixtureEntries[0], + placeholder, + consumer: h.consumer, + host, + routing: { + sessionId: 'session', + provider: 'fixture-a', + modelId: 'model', + }, + fetch: async (input) => { + sent.push(new Request(input).headers.get('authorization') ?? '') + return sent.length === 1 + ? new Response('limited', { status: 429 }) + : new Response('paid') + }, + }) + expect( + await ( + await serve('https://fixture.invalid', { headers: bearer }) + ).text(), + ).toBe('paid') + expect(sent).toEqual(['Bearer secret-oauth-1', 'Bearer secret-apikey-1']) + expect(calls).toHaveLength(1) + }) +} + +test('routing no-answer stop never sends the plugin first key', async () => { + const h = await harness() + routing(async () => undefined) + let sends = 0 + const serve = createFetch({ + entry: fixtureEntries[0], + placeholder, + consumer: h.consumer, + routing: { sessionId: 'session', provider: 'fixture-a', modelId: 'model' }, + fetch: async () => { + sends++ + return new Response('not admitted') + }, + }) + await expect( + serve('https://fixture.invalid', { headers: bearer }), + ).rejects.toMatchObject({ code: 'routing-no-answer' }) + expect(sends).toBe(0) + expect(h.reads).toEqual([]) +}) + +test('routing candidates exclude user-disabled rows OAuth and other providers', async () => { + const h = await harness([ + row('a'), + row('b'), + row('oauth', { + credentialType: 'oauth', + authMethod: 'oauth', + refreshAdapter: 'fixture-a', + }), + row('other', { providerIds: ['fixture-b'] }), + ]) + const a = h.consumer.snapshot()?.rows.find((row) => row.credentialId === 'a') + if (!a) throw new Error('missing row') + await h.consumer.decline(a.routeId) + let sends = 0 + const serve = createFetch({ + entry: fixtureEntries[0], + placeholder, + consumer: h.consumer, + routing: { sessionId: 'session', provider: 'fixture-a', modelId: 'model' }, + fetch: async () => { + sends++ + return new Response('ok') + }, + }) + for (const id of ['a', 'oauth', 'other']) { + routing(async () => ({ account_id: `account-${id}` })) + await expect( + serve('https://fixture.invalid', { headers: bearer }), + ).rejects.toMatchObject({ code: 'routing-no-answer' }) + } + expect(sends).toBe(0) +}) + +test('routing rechecks each redirect and never rotates a 401 behind routing', async () => { + const h = await harness() + let sends = 0 + const { calls } = routing(async (call) => + call === 1 + ? { account_id: 'account-a' } + : { refused: { code: 'account_selection_final' } }, + ) + const serve = createFetch({ + entry: fixtureEntries[0], + placeholder, + consumer: h.consumer, + routing: { sessionId: 'session', provider: 'fixture-a', modelId: 'model' }, + fetch: async () => { + sends++ + return new Response(null, { status: 307, headers: { location: '/next' } }) + }, + }) + await expect( + serve('https://fixture.invalid', { headers: bearer }), + ).rejects.toMatchObject({ code: 'routing-declined' }) + expect(calls).toHaveLength(2) + expect(sends).toBe(1) + routing(async () => ({ account_id: 'account-a' })) + const authFailure = createFetch({ + entry: fixtureEntries[0], + placeholder, + consumer: h.consumer, + routing: { sessionId: 'session', provider: 'fixture-a', modelId: 'model' }, + fetch: async () => { + sends++ + return new Response('bad key', { status: 401 }) + }, + }) + expect( + (await authFailure('https://fixture.invalid', { headers: bearer })).status, + ).toBe(401) + expect(sends).toBe(2) +}) + +for (const host of ['opencode1', 'pi'] as const) { + test(`routing ${host}: selection overrides roster and local marks; no accessor keeps today's choice`, async () => { + const h = await harness( + [row('a', { accountId: undefined }), row('b', { accountId: undefined })], + host, + ) + const sent: string[] = [] + const marks = new Map() + const serve = createFetch({ + entry: fixtureEntries[0], + placeholder, + consumer: h.consumer, + host, + routing: { + sessionId: 'session', + provider: 'fixture-a', + modelId: 'model', + }, + rateLimitMarks: marks, + fetch: async (input) => { + sent.push(new Request(input).headers.get('authorization') ?? '') + return new Response('ok') + }, + }) + Reflect.deleteProperty(globalThis, symbol) + await serve('https://fixture.invalid', { headers: bearer }) + expect(sent).toEqual(['Bearer secret-a-1']) + for (const r of h.consumer.routingRows()) + marks.set(r.id, Date.now() + 60000) + const { calls } = routing(async () => ({ + account_id: 'b', + credential_id: 'b', + })) + await serve('https://fixture.invalid', { headers: bearer }) + expect(sent).toEqual(['Bearer secret-a-1', 'Bearer secret-b-1']) + expect(calls).toEqual([ + { + session: 'session', + request: { provider: 'fixture-a', modelId: 'model', refusals: true }, + options: { deadlineMs: 1500 }, + }, + ]) + }) + + test(`routing ${host}: refusal and no answer send nothing`, async () => { + const h = await harness(undefined, host) + let sends = 0 + const serve = createFetch({ + entry: fixtureEntries[0], + placeholder, + consumer: h.consumer, + host, + routing: { + sessionId: 'session', + provider: 'fixture-a', + modelId: 'model', + }, + fetch: async () => { + sends++ + return new Response('ok') + }, + }) + for (const [answer, code] of [ + [ + { + refused: { + code: 'account_selection_refused', + routing_code: 'private-label', + }, + }, + 'routing-declined', + ], + [undefined, 'routing-no-answer'], + ] as const) { + routing(async () => answer) + await expect( + serve('https://fixture.invalid', { headers: bearer }), + ).rejects.toMatchObject({ code }) + } + expect(sends).toBe(0) + expect(h.reads).toEqual([]) + }) + + test(`routing ${host}: retries ask again and exclude every tried key`, async () => { + const h = await harness([row('a'), row('b'), row('c')], host) + const serve = createFetch({ + entry: fixtureEntries[0], + placeholder, + consumer: h.consumer, + host, + routing: { + sessionId: 'session', + provider: 'fixture-a', + modelId: 'model', + }, + fetch: async () => new Response('limited', { status: 429 }), + }) + const { calls } = routing(async (call) => ({ + account_id: call === 2 ? 'account-b' : 'account-a', + })) + await expect( + serve('https://fixture.invalid', { headers: bearer }), + ).rejects.toMatchObject({ code: 'routing-no-answer' }) + expect(calls).toHaveLength(3) + expect(h.reads).toEqual(['a', 'b']) + }) +} + +test('routing reports only qualifying balance refusals with vault identities and provider codes', async () => { + for (const [service, status, body, code] of [ + ['deepseek', 402, '{}', 'http_402'], + ['deepseek', 402, '{"error":{"code":"balance_empty"}}', 'balance_empty'], + [ + 'openrouter', + 402, + '{"error":{"metadata":{"limit_source":"openrouter_credits"}}}', + 'openrouter_credits', + ], + [ + 'openrouter', + 402, + '{"error":{"metadata":{"limit_source":"key_limit"}}}', + undefined, + ], + [ + 'openrouter', + 402, + '{"error":{"metadata":{"limit_source":"openrouter_credits","reason":"weight_exceeds_budget"}}}', + undefined, + ], + ['deepseek', 401, '{}', undefined], + ['deepseek', 403, '{}', undefined], + ['deepseek', 429, '{}', undefined], + ] as const) { + const h = await harness([row('a', { accountId: undefined }), row('b')]) + const { reports } = routing(async (call) => + call === 1 + ? { account_id: 'a' } + : { refused: { code: 'account_selection_final' } }, + ) + const serve = createFetch({ + entry: { ...fixtureEntries[0], service }, + placeholder, + consumer: h.consumer, + routing: { sessionId: 'session', provider: service, modelId: 'model' }, + fetch: async () => new Response(body, { status }), + }) + await serve('https://fixture.invalid', { headers: bearer }).catch(() => {}) + expect(reports).toEqual( + code + ? [ + { + account_id: 'a', + credential_id: 'a', + outcome: 'credential_refused', + provider_code: code, + }, + ] + : [], + ) + } +}) diff --git a/packages/opencode/package.json b/packages/opencode/package.json index ae2b38c..c317ead 100644 --- a/packages/opencode/package.json +++ b/packages/opencode/package.json @@ -34,7 +34,7 @@ }, "dependencies": { "@cortexkit/claustrum-client": "^0.18.0", - "@cortexkit/common-auth": "^0.44.0", + "@cortexkit/common-auth": "^0.47.0", "@cortexkit/private-files": "^0.2.2", "@cortexkit/subc-client": "^0.24.0" }, diff --git a/packages/opencode/src/v1/fetch.ts b/packages/opencode/src/v1/fetch.ts index 60a46d0..896e41a 100644 --- a/packages/opencode/src/v1/fetch.ts +++ b/packages/opencode/src/v1/fetch.ts @@ -2,6 +2,7 @@ import { type AccountLimits, type CoreFetch, createFetch, + hasVaultRoutingAccessor, type ProviderEntry, type ServingConsumer, } from '@cortexkit/claustrum-auth-core' @@ -59,18 +60,50 @@ export function providerFetch(options: ProviderFetchOptions): CoreFetch { init?: RequestInit, ) => { let shadowed: string | undefined + const headers = new Headers( + init?.headers ?? (input instanceof Request ? input.headers : undefined), + ) + let modelId = '' + const body = init?.body + const routed = + hasVaultRoutingAccessor() && + headers.get('authorization') === + `Bearer ${placeholderFor(options.provider)}` + if (routed && typeof body === 'string') { + try { + const model: unknown = JSON.parse(body)?.model + if (typeof model === 'string') modelId = model + } catch {} + } else if (routed && input instanceof Request && input.body) { + try { + const model: unknown = (await input.clone().json())?.model + if (typeof model === 'string') modelId = model + } catch {} + } const core = createFetch({ entry: options.entry, placeholder: placeholderFor(options.provider), consumer: options.consumer, rateLimitMarks, limits: options.limits, + routing: { + sessionId: headers.get('session-id') ?? undefined, + provider: options.provider, + modelId, + }, // The core passes a request it does not own on with the caller's own // arguments; every owned dispatch is a new Request it built. - fetch: (next, nextInit) => - next === input && nextInit === init - ? previous(next, nextInit) - : transport(next, nextInit), + fetch: (next, nextInit) => { + if (next === input && nextInit === init) { + if (!headers.has('session-id')) return previous(next, nextInit) + const foreign = new Request(next, nextInit) + foreign.headers.delete('session-id') + return previous(foreign) + } + const outbound = new Request(next, nextInit) + outbound.headers.delete('session-id') + return transport(outbound) + }, async effectiveCredential() { shadowed = shadowedVariable( (options.env ?? hostEnvironment)(), diff --git a/packages/opencode/src/v1/messages.ts b/packages/opencode/src/v1/messages.ts index 169c344..b4361f6 100644 --- a/packages/opencode/src/v1/messages.ts +++ b/packages/opencode/src/v1/messages.ts @@ -1,4 +1,4 @@ -import { command } from '@cortexkit/claustrum-auth-core' +import { command, routingRefusalText } from '@cortexkit/claustrum-auth-core' import type { EffectiveCredential } from './effective.js' // Removal commands for the platform the host runs on. The plugin cannot tell @@ -47,6 +47,8 @@ export function userCredentialMessage( // Fixed text for each local refusal. Nothing here interpolates request data, // so a message can never echo a header, URL or served secret. export function refusalMessage(code: string, provider: string): string { + const routing = routingRefusalText(code) + if (routing) return routing switch (code) { case 'placement-refused': return `claustrum-auth: this model's driver sends its key for ${provider} outside the Authorization header; the claustrum-auth plugin only serves keys sent in Authorization.` diff --git a/packages/opencode/src/v1/plugin.ts b/packages/opencode/src/v1/plugin.ts index be6c0ea..2397369 100644 --- a/packages/opencode/src/v1/plugin.ts +++ b/packages/opencode/src/v1/plugin.ts @@ -208,6 +208,11 @@ export async function opencode1Hosts( } return { + async 'chat.headers'(input, output) { + // Carry host identity through the SDK to our fetch, never to a provider. + if (listed.has(input.model.providerID)) + output.headers['session-id'] = input.sessionID + }, async config(config) { const commands = (config.command ?? {}) as Record config.command = { diff --git a/packages/opencode/src/v2/adapter.ts b/packages/opencode/src/v2/adapter.ts index 7d1b2f4..2103e99 100644 --- a/packages/opencode/src/v2/adapter.ts +++ b/packages/opencode/src/v2/adapter.ts @@ -14,9 +14,12 @@ import { localLimitRetry, observeBillingResponse, type ProviderEntry, + type RoutedKey, redactError, redactResponse, + reportRoutedBalance, type ServingConsumer, + selectRoutedKey, servingCallbacks, } from '@cortexkit/claustrum-auth-core' import type { ClaustrumScopedAttempt } from '@cortexkit/common-auth/claustrum' @@ -68,6 +71,8 @@ export interface VaultAttempt { frameStatus?: number billing?: BillingAttempt billingAttempt?: number + routed?: RoutedKey + balanceReported?: boolean } /** Identifies a session's socket: the host keeps one per session and kind. */ @@ -296,6 +301,7 @@ export function createVaultAdapter( >() const placeholder = placeholderSecret(options.provider) let billingAttempts: Map | undefined + const routingChoices = new Map() const refuse = (error: unknown) => { if ( error instanceof ClaustrumAuthError && @@ -419,6 +425,28 @@ export function createVaultAdapter( const selected = available.find((row) => row.ready) ?? (input.phase === 'request' ? available[0] : undefined) + try { + const excluded = new Set(input.excluded ?? []) + if (input.rerouteFrom) excluded.add(input.rerouteFrom.accountId) + const routed = await selectRoutedKey( + consumer, + options.provider, + input.modelID, + input.sessionID, + excluded, + selected?.account.accountId, + ) + if (routed) { + routingChoices.set(socketKey(input.sessionID, input.kind), routed) + if (routingChoices.size > 256) { + const oldest = routingChoices.keys().next().value + if (oldest !== undefined) routingChoices.delete(oldest) + } + return routed.routeId + } + } catch (error) { + throw refuse(error) + } const decision = decideAfterLimit( { phase: input.phase === 'request' ? 'pre-send' : 'after-limit', @@ -477,13 +505,19 @@ export function createVaultAdapter( // so the plugin's handshake hook checked the origin before this call // and left the URL here; only then is the vault read now. const key = socketKey(sessionID, kind) + const choice = routingChoices.get(key) + routingChoices.delete(key) + const routed = choice?.routeId === accountId ? choice : undefined const hostAttempt = billingAttempts?.get(key) billingAttempts?.delete(key) const url = options.handshakes?.pending.get(key) if (url === undefined) return { headers: {}, - attempt: { billingAttempt: hostAttempt } satisfies VaultAttempt, + attempt: { + billingAttempt: hostAttempt, + routed, + } satisfies VaultAttempt, } options.handshakes?.pending.delete(key) // The handshake URL has no account segment the plugin could fill. @@ -492,12 +526,14 @@ export function createVaultAdapter( const consumer = await options.consumer() if (!consumer) throw new Error(noAccountText(options.provider)) const limits = await options.limits?.() - if (limits && !(await limits.beforeDispatch(accountId))) + if (!routed && limits && !(await limits.beforeDispatch(accountId))) throw await noAccount() const observation = limits?.observe(accountId) let receipt: ClaustrumScopedAttempt try { receipt = await consumer.authorize(accountId) + if (routed && receipt.credentialId !== routed.credentialId) + throw new ClaustrumAuthError('routing-no-answer') } catch (error) { throw refuse(error) } @@ -507,6 +543,7 @@ export function createVaultAdapter( headers: { authorization: `Bearer ${receipt.accessToken}` }, attempt: { receipt, + routed, consumer, handshakeURL: url, limits, @@ -553,19 +590,22 @@ export function createVaultAdapter( if (!consumer) throw new Error(noAccountText(options.provider)) // Only the chosen account may serve this attempt, so routing inside the // shared callbacks sees that one row and no marks. - const callbacks = servingCallbacks({ - entry: options.entry, - placeholder, - consumer: { - routingRows: () => - consumer.routingRows().filter((row) => row.id === accountId), - authorize: (route, signal) => consumer.authorize(route, signal), - send: (...args) => consumer.send(...args), - reportFailure: (...args) => consumer.reportFailure(...args), + const callbacks = servingCallbacks( + { + entry: options.entry, + placeholder, + consumer: { + routingRows: () => + consumer.routingRows().filter((row) => row.id === accountId), + authorize: (route, signal) => consumer.authorize(route, signal), + send: (...args) => consumer.send(...args), + reportFailure: (...args) => consumer.reportFailure(...args), + }, + now, + limits: options.limits, }, - now, - limits: options.limits, - }) + holder?.routed !== undefined, + ) // In order: placement and origin (no vault read), authorize, the // class-B account rewrite from the served receipt, then Authorization. let served: Awaited> @@ -579,6 +619,11 @@ export function createVaultAdapter( throw refuse(error) } if (!served) return undefined + if ( + holder?.routed && + served.receipt.credentialId !== holder.routed.credentialId + ) + throw refuse(new ClaustrumAuthError('routing-no-answer')) if (holder) { holder.receipt = served.receipt holder.consumer = consumer @@ -624,6 +669,16 @@ export function createVaultAdapter( }, async limitFromResponse({ status, headers, body, attempt }) { if (attempt.data?.localRefusal) return undefined + if (attempt.data?.routed && !attempt.data.balanceReported) { + attempt.data.balanceReported = true + await reportRoutedBalance( + attempt.data.routed, + attempt.sessionID, + options.entry.service, + status, + body ?? (async () => ''), + ) + } const parsed = await limitFromReply({ account: attempt.data?.observation?.binding.account ?? { namespace: 'claustrum-auth:standalone', diff --git a/packages/opencode/src/v2/messages.ts b/packages/opencode/src/v2/messages.ts index 85be8be..aa83904 100644 --- a/packages/opencode/src/v2/messages.ts +++ b/packages/opencode/src/v2/messages.ts @@ -1,4 +1,7 @@ -import { ClaustrumAuthError } from '@cortexkit/claustrum-auth-core' +import { + ClaustrumAuthError, + routingRefusalText, +} from '@cortexkit/claustrum-auth-core' import { ClaustrumCredentialError } from '@cortexkit/common-auth/claustrum' // Every text the OpenCode 2 entry shows is fixed. None of them interpolates @@ -64,6 +67,8 @@ function refusalText( provider: string, origins: readonly string[], ) { + const routing = routingRefusalText(code) + if (routing) return routing switch (code) { case 'origin-refused': return originText(provider, origins) diff --git a/packages/opencode/test/v1-plugin.test.ts b/packages/opencode/test/v1-plugin.test.ts index 391adf7..3c36561 100644 --- a/packages/opencode/test/v1-plugin.test.ts +++ b/packages/opencode/test/v1-plugin.test.ts @@ -385,6 +385,20 @@ test('the config hook wraps every listed provider, keeps an earlier fetch and re await h.hooks.dispose?.() }) +test('routing OC1 chat headers stamp only served providers with the host session', async () => { + const h = await pluginHarness() + for (const providerID of [P, 'openai', 'unknown']) { + const output = { headers: {} as Record } + await h.hooks['chat.headers']?.( + { sessionID: 'host-session', model: { providerID } } as never, + output, + ) + expect(output.headers).toEqual( + providerID === P ? { 'session-id': 'host-session' } : {}, + ) + } +}) + test('the command hook runs use and release and ignores other commands', async () => { const h = await pluginHarness() await h.hooks.config?.({ diff --git a/packages/opencode/test/v1-routing.test.ts b/packages/opencode/test/v1-routing.test.ts new file mode 100644 index 0000000..355c3fe --- /dev/null +++ b/packages/opencode/test/v1-routing.test.ts @@ -0,0 +1,56 @@ +import { expect, test } from 'bun:test' +import { harness, row } from '../../core/test/fakes.ts' +import { routing } from '../../core/test/routing-fixtures.ts' +import { fixtureEntries } from '../../core/test-entry.ts' +import { providerFetch } from '../src/v1/fetch.ts' +import { placeholderFor } from '../src/v1/placeholder.ts' + +test('routing OC1 uses stamped session and routing key, strips private header, fails closed without session', async () => { + const h = await harness([row('a'), row('b')]) + const sent: Request[] = [] + const serve = providerFetch({ + provider: 'fixture-a', + entry: fixtureEntries[0], + consumer: h.consumer, + envNames: async () => [], + env: () => ({}), + transport: async (input, init) => { + sent.push(new Request(input, init)) + return new Response('ok') + }, + }) + const { calls } = routing(async () => ({ account_id: 'account-b' })) + const init = { + method: 'POST', + headers: { + authorization: `Bearer ${placeholderFor('fixture-a')}`, + 'session-id': 'host-session', + }, + body: '{"model":"requested-model"}', + } + await serve('https://fixture.invalid', init) + expect(sent.map((r) => r.headers.get('authorization'))).toEqual([ + 'Bearer secret-b-1', + ]) + expect(sent[0]?.headers.has('session-id')).toBe(false) + expect(calls).toEqual([ + { + session: 'host-session', + request: { + provider: 'fixture-a', + modelId: 'requested-model', + refusals: true, + }, + options: { deadlineMs: 1500 }, + }, + ]) + await expect( + serve('https://fixture.invalid', { + ...init, + headers: { authorization: init.headers.authorization }, + }), + ).rejects.toThrow( + 'Thalamus routing did not answer, so this API key request was not sent.', + ) + expect(sent).toHaveLength(1) +}) diff --git a/packages/opencode/test/v2/e2e/billing-preload.ts b/packages/opencode/test/v2/e2e/billing-preload.ts index af49c85..f63fc53 100644 --- a/packages/opencode/test/v2/e2e/billing-preload.ts +++ b/packages/opencode/test/v2/e2e/billing-preload.ts @@ -1,9 +1,13 @@ import { appendFileSync } from 'node:fs' +import '../../../../core/test/routing-preload.ts' // Loaded before plugin setup, in the real host process rather than the runner. const destination = process.env.CLAUSTRUM_AUTH_E2E_BILLING if (destination) { Reflect.set(globalThis, Symbol.for('cortexkit.routing'), () => ({ + accountFor: async ({ provider }: { provider: string }) => ({ + account_id: `apikey:${provider}`, + }), report(payload: { kind: string }) { if (payload.kind === 'billing') appendFileSync(destination, `${JSON.stringify(payload)}\n`) diff --git a/packages/opencode/test/v2/e2e/opencode2.e2e.test.ts b/packages/opencode/test/v2/e2e/opencode2.e2e.test.ts index 597c910..1abf62d 100644 --- a/packages/opencode/test/v2/e2e/opencode2.e2e.test.ts +++ b/packages/opencode/test/v2/e2e/opencode2.e2e.test.ts @@ -521,6 +521,31 @@ const onWire = (run: Scenario, needle: string) => ) describe.skipIf(!ENABLED)('claustrum-auth plugin on real OpenCode 2', () => { + for (const [mode, message] of [ + ['refuse', 'Thalamus routing declined this request, so it was not sent.'], + [ + 'never', + 'Thalamus routing did not answer, so this API key request was not sent.', + ], + ] as const) { + test(`routing ${mode}: fixed refusal reaches user without a provider request`, async () => { + await scenario( + [{ id: 'a1', providerIds: ['fixture-a'] }], + async (run) => { + expect((await run.login('fixture-a')).exit).toBe(0) + const result = await run.turn('fixture-a') + expect(result.output).toContain(message) + expect(run.mock.records).toEqual([]) + expect(run.vault.reads).toEqual([]) + }, + { + beforeStart: async (run) => { + run.env.CLAUSTRUM_AUTH_E2E_ROUTING = mode + }, + }, + ) + }) + } test('billing: a real vault-served HTTP turn reports completed provider usage', async () => { await scenario( [{ id: 'apikey:cerebras', providerIds: ['cerebras'] }], diff --git a/packages/opencode/test/v2/routing.test.ts b/packages/opencode/test/v2/routing.test.ts new file mode 100644 index 0000000..6fe74bb --- /dev/null +++ b/packages/opencode/test/v2/routing.test.ts @@ -0,0 +1,207 @@ +import { expect, test } from 'bun:test' +import { placeholderSecret } from '@cortexkit/common-auth/opencode2' +import { harness, row } from '../../../core/test/fakes.ts' +import { routing } from '../../../core/test/routing-fixtures.ts' +import { fixtureEntries } from '../../../core/test-entry.ts' +import { createVaultAdapter } from '../../src/v2/adapter.ts' + +const scope = { + sessionID: 'session', + providerID: 'fixture-a', + modelID: 'model', + kind: 'primary' as const, + agent: 'build', +} + +test('routing OC2 OAuth-only and mixed pools preserve subscription choice until an API-key attempt', async () => { + for (const mixed of [false, true]) { + const h = await harness( + [ + row('oauth', { + credentialType: 'oauth', + authMethod: 'oauth', + refreshAdapter: 'fixture-a', + }), + ...(mixed ? [row('key')] : []), + ], + 'opencode2', + ) + const adapter = createVaultAdapter({ + entry: fixtureEntries[0], + provider: 'fixture-a', + consumer: async () => h.consumer, + rateLimitMarks: new Map(), + }) + if (mixed) + await h.consumer.reorder( + ['oauth', 'key'].map( + (id) => + h.consumer.snapshot()?.rows.find((row) => row.credentialId === id) + ?.routeId as string, + ), + ) + const { calls } = routing(async () => ({ account_id: 'account-key' })) + const first = await adapter.chooseAccount({ ...scope, phase: 'request' }) + expect(first).toBe( + h.consumer.snapshot()?.rows.find((row) => row.credentialId === 'oauth') + ?.routeId, + ) + expect(calls).toEqual([]) + if (mixed && typeof first === 'string') { + const next = await adapter.chooseAccount({ + ...scope, + phase: 'retry', + rerouteFrom: { + accountId: first, + limit: { status: 429, reason: 'rate-limit' }, + }, + excluded: [first], + }) + expect(next).toBe( + h.consumer.snapshot()?.rows.find((row) => row.credentialId === 'key') + ?.routeId, + ) + expect(calls).toHaveLength(1) + } + } +}) +async function world(service = 'fixture-a') { + const h = await harness( + [row('a', { accountId: undefined }), row('b', { accountId: undefined })], + 'opencode2', + ) + const adapter = createVaultAdapter({ + entry: { ...fixtureEntries[0], service }, + provider: 'fixture-a', + consumer: async () => h.consumer, + rateLimitMarks: new Map(), + }) + const send = async (input = { ...scope, phase: 'request' as const }) => { + const choice = await adapter.chooseAccount(input) + if (typeof choice !== 'string') throw new Error('expected account') + const headers = await adapter.accountHeaders({ + ...scope, + accountId: choice, + }) + if (!headers.attempt || typeof headers.attempt !== 'object') + throw new Error('missing attempt') + const attempt = { + ...scope, + attemptId: 'attempt', + accountId: choice, + data: headers.attempt, + transport: 'http' as const, + } + const request = new Request('https://fixture.invalid', { + headers: { authorization: `Bearer ${placeholderSecret('fixture-a')}` }, + }) + const sent = await adapter.rewriteRequest?.({ + ...scope, + request, + accountId: choice, + attempt, + }) + return { sent, attempt } + } + return { h, adapter, send } +} + +test('routing OC2 selection serves routing key for primary titles and compaction; absent accessor preserves roster', async () => { + const { send, adapter, h } = await world() + const first = await send() + expect(first.sent?.headers.get('authorization')).toBe('Bearer secret-a-1') + const { calls } = routing(async () => ({ account_id: 'b' })) + for (const kind of ['primary', 'title', 'compaction'] as const) + expect( + await adapter.chooseAccount({ ...scope, kind, phase: 'request' }), + ).toBe( + h.consumer.snapshot()?.rows.find((r) => r.credentialId === 'b')?.routeId, + ) + expect((await send()).sent?.headers.get('authorization')).toBe( + 'Bearer secret-b-1', + ) + expect(calls).toHaveLength(4) +}) + +test('routing OC2 refusal and no answer authorize and send nothing', async () => { + const { send, h } = await world() + for (const [answer, text] of [ + [ + { refused: { code: 'account_selection_final' } }, + 'Thalamus routing declined this request, so it was not sent.', + ], + [ + undefined, + 'Thalamus routing did not answer, so this API key request was not sent.', + ], + ] as const) { + routing(async () => answer) + await expect(send()).rejects.toThrow(text) + } + expect(h.reads).toEqual([]) +}) + +test('routing OC2 retry excludes earlier keys and asks again', async () => { + const { adapter, h } = await world() + const a = h.consumer.snapshot()?.rows.find((r) => r.credentialId === 'a') + ?.routeId as string + const b = h.consumer.snapshot()?.rows.find((r) => r.credentialId === 'b') + ?.routeId as string + const { calls } = routing(async (call) => ({ + account_id: call === 2 ? 'a' : 'b', + })) + const retry = { + ...scope, + phase: 'retry' as const, + excluded: [a], + rerouteFrom: { accountId: a, limit: { status: 429, reason: 'rate-limit' } }, + } + expect(await adapter.chooseAccount(retry)).toBe(b) + await expect( + adapter.chooseAccount({ ...retry, excluded: [a, b] }), + ).rejects.toThrow('did not answer') + expect(calls).toHaveLength(2) +}) + +test('routing OC2 balance reports once and never for auth rate limits or self-set caps', async () => { + for (const [service, status, body, code] of [ + ['deepseek', 402, '{}', 'http_402'], + [ + 'openrouter', + 402, + '{"error":{"metadata":{"limit_source":"openrouter_credits"}}}', + 'openrouter_credits', + ], + [ + 'openrouter', + 402, + '{"error":{"metadata":{"limit_source":"key_limit"}}}', + undefined, + ], + ['deepseek', 401, '{}', undefined], + ['deepseek', 429, '{}', undefined], + ] as const) { + const { send, adapter } = await world(service) + const { reports } = routing(async () => ({ account_id: 'b' })) + const { attempt } = await send() + for (let i = 0; i < 2; i++) + await adapter.limitFromResponse?.({ + status, + headers: new Headers(), + body: async () => body, + attempt, + }) + expect(reports).toEqual( + code + ? [ + { + account_id: 'b', + credential_id: 'b', + outcome: 'credential_refused', + provider_code: code, + }, + ] + : [], + ) + } +}) diff --git a/packages/pi/package.json b/packages/pi/package.json index a08ff73..2d4bfb2 100644 --- a/packages/pi/package.json +++ b/packages/pi/package.json @@ -32,7 +32,7 @@ }, "dependencies": { "@cortexkit/claustrum-client": "^0.18.0", - "@cortexkit/common-auth": "^0.44.0", + "@cortexkit/common-auth": "^0.47.0", "@cortexkit/private-files": "^0.2.2", "@cortexkit/subc-client": "^0.24.0" }, diff --git a/packages/pi/src/extension.ts b/packages/pi/src/extension.ts index cfbb4ed..b36854b 100644 --- a/packages/pi/src/extension.ts +++ b/packages/pi/src/extension.ts @@ -331,31 +331,46 @@ class PiRuntime { // runtime override goes to pi-ai exactly as Pi resolved it. if (options?.apiKey !== served.placeholder) return streamSimple(model, context, options) + const routing = { + sessionId: options.sessionId ?? this.#ctx.sessionManager.getSessionId(), + provider: model.provider, + modelId: model.id, + } const inner = (options.fetch ?? globalThis.fetch) as CoreFetch if (!hasBillingAccessor()) return streamSimple(model, context, { ...options, - fetch: this.ownedFetch(served, inner) as typeof fetch, + fetch: this.ownedFetch( + served, + inner, + undefined, + routing, + ) as typeof fetch, }) let current: BillingAttempt | undefined let attempt = 0 const stream = streamSimple(model, context, { ...options, - fetch: this.ownedFetch(served, inner, (receipt) => { - current?.end() - if (!hasBillingAccessor()) return undefined - current = new BillingAttempt( - served.entry, - model.provider, - options.sessionId ?? this.#ctx.sessionManager.getSessionId(), - 'primary', - model.id, - receipt, - ++attempt, - true, - ) - return current - }) as typeof fetch, + fetch: this.ownedFetch( + served, + inner, + (receipt) => { + current?.end() + if (!hasBillingAccessor()) return undefined + current = new BillingAttempt( + served.entry, + model.provider, + options.sessionId ?? this.#ctx.sessionManager.getSessionId(), + 'primary', + model.id, + receipt, + ++attempt, + true, + ) + return current + }, + routing, + ) as typeof fetch, }) void stream.result().then( (message) => current?.piEnd(message.responseId, message.stopReason), @@ -368,6 +383,7 @@ class PiRuntime { served: Served, inner: CoreFetch, billing?: Parameters[0]['billing'], + routing?: Parameters[0]['routing'], ): CoreFetch { const serve = createFetch({ entry: served.entry, @@ -378,6 +394,7 @@ class PiRuntime { rateLimitMarks: served.marks, limits: () => this.#core.limits(served.provider), billing, + routing, }) return async (input, init) => { try { diff --git a/packages/pi/src/messages.ts b/packages/pi/src/messages.ts index 0ba3747..6023b19 100644 --- a/packages/pi/src/messages.ts +++ b/packages/pi/src/messages.ts @@ -1,4 +1,8 @@ -import { command, type ProviderEntry } from '@cortexkit/claustrum-auth-core' +import { + command, + type ProviderEntry, + routingRefusalText, +} from '@cortexkit/claustrum-auth-core' // Every text here is fixed: it names the provider and its allowed origins from // the plugin's own table, never anything taken from a request or a vault reply. @@ -10,6 +14,8 @@ export function refusalText( provider: string, entry: ProviderEntry, ): string { + const routing = routingRefusalText(code) + if (routing) return routing switch (code) { case 'placement-refused': return `${refused} this model's request carries the ${provider} vault placeholder outside the Authorization header, and vault keys are only sent as Authorization: Bearer. ${nothingSent}` diff --git a/packages/pi/test-entry.ts b/packages/pi/test-entry.ts index c03a5a2..f32180b 100644 --- a/packages/pi/test-entry.ts +++ b/packages/pi/test-entry.ts @@ -1,6 +1,7 @@ // The Pi test build: the real extension with fixture providers, loopback // origins and a fake vault reached over HTTP. It lives outside src, so neither the published bundle nor npm pack // can contain it. +import '../core/test/routing-preload.ts' import { appendFileSync, mkdirSync, rmSync } from 'node:fs' import type { ClaustrumScopedClient } from '@cortexkit/common-auth/claustrum' import type { ExtensionAPI } from '@earendil-works/pi-coding-agent' diff --git a/packages/pi/test/billing.test.ts b/packages/pi/test/billing.test.ts index afbb39e..f291563 100644 --- a/packages/pi/test/billing.test.ts +++ b/packages/pi/test/billing.test.ts @@ -12,6 +12,7 @@ import { } from '../../core/test/billing-fixtures.ts' import { vaultRowFor } from '../../core/test/fakes.ts' import { installVaultHost } from '../../core/test/host.ts' +import { routing } from '../../core/test/routing-fixtures.ts' import { createTestCore } from '../../core/test-entry.ts' import { createPiExtension, placeholderKey } from '../src/extension.ts' import { agentDir } from './harness.ts' @@ -23,6 +24,7 @@ async function turn( redirect = false, onSemantics?: () => void, onReady?: () => void, + routingExpectation: { second?: boolean; refusal?: string } = {}, ) { const base = supportedProviders.find((row) => row.service === provider) if (!base?.billing) throw new Error('missing provider') @@ -39,7 +41,7 @@ async function turn( : base const rows = [ vaultRowFor(entry, { accountId: 'account-1' }), - ...(retry + ...(retry || routingExpectation.second ? [ vaultRowFor(entry, { id: `apikey:${provider}:second`, @@ -109,6 +111,7 @@ async function turn( ) => AssistantMessageEventStream if (!stream) throw new Error('missing registered stream') let sends = 0 + let requests = 0 let includeUsage: unknown const reply = stream( { @@ -128,12 +131,13 @@ async function turn( sessionId: 'pi-session', apiKey: ownKey ? 'user-key' : placeholderKey(provider), fetch: async (input: RequestInfo | URL, init?: RequestInit) => { + requests++ const request = new Request(input, init) includeUsage = (await request.json()).stream_options expect(request.headers.get('authorization')).toBe( ownKey ? 'Bearer user-key' - : `Bearer secret-apikey:${provider}${retry && sends > 0 ? ':second' : ''}-1`, + : `Bearer secret-apikey:${provider}${routingExpectation.second || (retry && sends > 0) ? ':second' : ''}-1`, ) if (retry && ++sends === 1) return new Response('busy', { @@ -152,6 +156,12 @@ async function turn( for await (const _event of reply) { } const message = await reply.result() + if (routingExpectation.refusal) { + expect(message.stopReason).toBe('error') + expect(message.errorMessage).toContain(routingExpectation.refusal) + expect(requests).toBe(0) + return message + } expect(message.stopReason).toBe('stop') expect(includeUsage).toEqual({ include_usage: true }) return message @@ -206,8 +216,36 @@ test('Pi leaves a user key unbilled even when the routing accessor is installed' expect(reports).toEqual([]) }) +test('Pi routing selection serves the second vault key rather than roster first', async () => { + captureBilling(false, async () => ({ account_id: 'account-2' })) + await turn('deepseek', false, false, false, undefined, undefined, { + second: true, + }) +}) + +test('Pi routing refusal and no answer stop before the transport', async () => { + for (const [answer, refusal] of [ + [ + { refused: { code: 'account_selection_final' } }, + 'Thalamus routing declined this request, so it was not sent.', + ], + [ + undefined, + 'Thalamus routing did not answer, so this API key request was not sent.', + ], + ] as const) { + routing(async () => answer) + await turn('deepseek', false, false, false, undefined, undefined, { + refusal, + }) + } +}) + test('Pi reroute bills both physical sends with different UUIDs and attempts', async () => { - const { reports } = captureBilling() + let choices = 0 + const { reports } = captureBilling(false, async () => ({ + account_id: `account-${++choices}`, + })) await turn('deepseek', false, true) expect(reports).toHaveLength(2) expect( diff --git a/packages/pi/test/extension.e2e.test.ts b/packages/pi/test/extension.e2e.test.ts index dc517b9..65b843e 100644 --- a/packages/pi/test/extension.e2e.test.ts +++ b/packages/pi/test/extension.e2e.test.ts @@ -36,6 +36,30 @@ const consentAsked = ( setDefaultTimeout(90_000) describe.skipIf(!e2e)('Pi extension on the real host', () => { + for (const [mode, message] of [ + ['refuse', 'Thalamus routing declined this request, so it was not sent.'], + [ + 'never', + 'Thalamus routing did not answer, so this API key request was not sent.', + ], + ] as const) { + test(`routing ${mode}: fixed refusal reaches user without a provider request`, async () => { + const { mock, vault, launch } = await world({ + rows: [row('a1', 'fixture-a')], + pi: { env: { CLAUSTRUM_AUTH_E2E_ROUTING: mode } }, + }) + const pi = await launch() + await pi.command({ type: 'set_auto_retry', enabled: true }) + await pi.slash('/claustrum use fixture-a') + const result = await pi.turn('fixture-a', 'ok') + // The OpenAI SDK wraps the fixed local text in its HTTP error envelope. + expect(result.error).toBe( + `400: ${JSON.stringify({ message, type: 'claustrum_auth_refused', code: mode === 'refuse' ? 'routing-declined' : 'routing-no-answer' })}`, + ) + expect(mock.seen).toEqual([]) + expect(vault.reads).toEqual([]) + }) + } test.skipIf(process.env.CLAUSTRUM_AUTH_SETUP_E2E !== '1')( 'setup test build enrolls a clean Pi profile and its printed activation asks consent and serves a fixture request', async () => { diff --git a/packages/pi/test/routing-retry.test.ts b/packages/pi/test/routing-retry.test.ts new file mode 100644 index 0000000..4de91c1 --- /dev/null +++ b/packages/pi/test/routing-retry.test.ts @@ -0,0 +1,23 @@ +import { expect, test } from 'bun:test' +import { isRetryableAssistantError } from '@earendil-works/pi-ai/compat' +import { + routingDeclinedText, + routingNoAnswerText, +} from '../../core/src/routing.ts' + +test('routing fixed texts are not retryable by Pi actual retry classifier', () => { + for (const errorMessage of [routingDeclinedText, routingNoAnswerText]) { + expect( + isRetryableAssistantError({ + stopReason: 'error', + errorMessage, + } as Parameters[0]), + ).toBe(false) + } + expect( + isRetryableAssistantError({ + stopReason: 'error', + errorMessage: 'connection refused', + } as Parameters[0]), + ).toBe(true) +}) From f168d5b1e621d15521bc87afaa7a8a0a236b0fe1 Mon Sep 17 00:00:00 2001 From: ualtinok <94532+ualtinok@users.noreply.github.com> Date: Thu, 12 Nov 2026 16:56:32 +0100 Subject: [PATCH 2/3] mason: preserve routed credential rotation and redirect admission Use consumer.send for routed core attempts, validating credential identity inside dispatch so a newer version of the same credential can retry a 401 without another routing call. Preserve the host-owned OpenCode2 send path and explain its cooldown bypass flag. Same-origin redirects belong to one admitted request, so replace the prior redirect re-selection test with one-key/one-routing-call coverage. Replace the no-rotation claim with regressions for same-credential version rotation on OC1 and Pi. DeepSeek balance reports always use http_402, including bodies with generic request error codes. --- packages/core/src/fetch.ts | 40 +++--------- packages/core/src/routing.ts | 8 +-- packages/core/src/serving.ts | 3 + packages/core/test/routing.test.ts | 78 +++++++++++++++-------- packages/opencode/test/v2/routing.test.ts | 1 + 5 files changed, 67 insertions(+), 63 deletions(-) diff --git a/packages/core/src/fetch.ts b/packages/core/src/fetch.ts index 379d547..48d79e0 100644 --- a/packages/core/src/fetch.ts +++ b/packages/core/src/fetch.ts @@ -213,7 +213,7 @@ export function createFetch(options: FetchOptions): CoreFetch { ? ownDecision.account.accountId : rows.find((row) => row.admitted && row.placementEligible)?.account .accountId - let routed = await selectRoutedKey( + const routed = await selectRoutedKey( options.consumer, options.routing?.provider ?? options.entry.service, options.routing?.modelId ?? '', @@ -247,7 +247,7 @@ export function createFetch(options: FetchOptions): CoreFetch { } if (initial.action !== 'reroute') throw new Error('claustrum-auth: unsupported limit action') - let route = initial.account.accountId + const route = initial.account.accountId if (!routed && limits && !(await limits.beforeDispatch(route))) { rejected.add(route) continue @@ -263,7 +263,7 @@ export function createFetch(options: FetchOptions): CoreFetch { pending = undefined tried.add(route) const previousMark = rateLimitMarks.get(route) - let observation = limits?.observe(route) + const observation = limits?.observe(route) let target = url let method = template.method let bytes = body @@ -272,23 +272,12 @@ export function createFetch(options: FetchOptions): CoreFetch { for (;;) { signal.throwIfAborted() checkOrigin(target, options.entry) - if (redirects > 0 && routed) { - routed = await selectRoutedKey( - options.consumer, - options.routing?.provider ?? options.entry.service, - options.routing?.modelId ?? '', - options.routing?.sessionId, - new Set([...tried].filter((id) => id !== route)), - ) - if (!routed) throw new ClaustrumAuthError('routing-no-answer') - route = routed.routeId - tried.add(route) - observation = limits?.observe(route) - } const dispatch = async ( receipt: ClaustrumScopedAttempt, dispatchSignal?: AbortSignal, ) => { + if (routed && receipt.credentialId !== routed.credentialId) + throw new ClaustrumAuthError('routing-no-answer') secrets.push(receipt.accessToken) const draft = placeReceipt(target, headers, options.entry, receipt) const dispatchHeaders = new Headers(draft.headers) @@ -327,21 +316,10 @@ export function createFetch(options: FetchOptions): CoreFetch { throw error } } - if (routed) { - // Consumer.send can rotate and resend a 401 on its own. Routing - // owns paid-key admission, so authorize exactly one physical send. - const receipt = await options.consumer.authorize(route, signal) - if (receipt.credentialId !== routed.credentialId) - throw new ClaustrumAuthError('routing-no-answer') - response = await dispatch(receipt, signal) - if (response.status === 401) - await options.consumer.reportFailure(receipt, 401, 'direct') - } else { - response = await options.consumer.send(route, dispatch, { - site: 'claustrum-auth', - signal, - }) - } + response = await options.consumer.send(route, dispatch, { + site: 'claustrum-auth', + signal, + }) signal.throwIfAborted() const location = response.headers.get('location') if (![301, 302, 303, 307, 308].includes(response.status) || !location) diff --git a/packages/core/src/routing.ts b/packages/core/src/routing.ts index 327d307..9c326c7 100644 --- a/packages/core/src/routing.ts +++ b/packages/core/src/routing.ts @@ -89,14 +89,8 @@ export async function reportRoutedBalance( if (!key || !sessionId || status !== 402) return const text = await body().catch(() => '') if (!isBalanceFailure(service, text)) return - let providerCode = + const providerCode = service === 'openrouter' ? 'openrouter_credits' : 'http_402' - if (service === 'deepseek') { - try { - const code: unknown = JSON.parse(text)?.error?.code - if (typeof code === 'string') providerCode = code - } catch {} - } reportRouting({ sessionId, accountId: key.accountId, diff --git a/packages/core/src/serving.ts b/packages/core/src/serving.ts index f728f47..f24d70b 100644 --- a/packages/core/src/serving.ts +++ b/packages/core/src/serving.ts @@ -59,6 +59,9 @@ export function placeReceipt( // merely because it shares a provider or session with an owned request. export function servingCallbacks( options: ServingOptions, + // Routing owns cooldown decisions for admitted API keys, so skip the + // plugin's rate-limit and balance cooldowns. + // Vault authorization and the user's on/off switch still apply. routingAdmitted = false, ) { return { diff --git a/packages/core/test/routing.test.ts b/packages/core/test/routing.test.ts index 2205deb..60e4f7f 100644 --- a/packages/core/test/routing.test.ts +++ b/packages/core/test/routing.test.ts @@ -138,9 +138,9 @@ test('routing candidates exclude user-disabled rows OAuth and other providers', expect(sends).toBe(0) }) -test('routing rechecks each redirect and never rotates a 401 behind routing', async () => { +test('routing keeps one admitted key across same-origin redirect hops', async () => { const h = await harness() - let sends = 0 + const sent: Request[] = [] const { calls } = routing(async (call) => call === 1 ? { account_id: 'account-a' } @@ -151,34 +151,62 @@ test('routing rechecks each redirect and never rotates a 401 behind routing', as placeholder, consumer: h.consumer, routing: { sessionId: 'session', provider: 'fixture-a', modelId: 'model' }, - fetch: async () => { - sends++ - return new Response(null, { status: 307, headers: { location: '/next' } }) - }, - }) - await expect( - serve('https://fixture.invalid', { headers: bearer }), - ).rejects.toMatchObject({ code: 'routing-declined' }) - expect(calls).toHaveLength(2) - expect(sends).toBe(1) - routing(async () => ({ account_id: 'account-a' })) - const authFailure = createFetch({ - entry: fixtureEntries[0], - placeholder, - consumer: h.consumer, - routing: { sessionId: 'session', provider: 'fixture-a', modelId: 'model' }, - fetch: async () => { - sends++ - return new Response('bad key', { status: 401 }) + fetch: async (input) => { + sent.push(new Request(input)) + return sent.length === 1 + ? new Response(null, { status: 307, headers: { location: '/next' } }) + : new Response('done') }, }) expect( - (await authFailure('https://fixture.invalid', { headers: bearer })).status, - ).toBe(401) - expect(sends).toBe(2) + await (await serve('https://fixture.invalid', { headers: bearer })).text(), + ).toBe('done') + expect(sent.map((request) => request.url)).toEqual([ + 'https://fixture.invalid/', + 'https://fixture.invalid/next', + ]) + expect(sent.map((request) => request.headers.get('authorization'))).toEqual([ + 'Bearer secret-a-1', + 'Bearer secret-a-1', + ]) + expect(calls).toHaveLength(1) }) for (const host of ['opencode1', 'pi'] as const) { + test(`routing ${host}: 401 rotation retries the same credential without asking again`, async () => { + const h = await harness([row('a'), row('b')], host) + const { calls } = routing(async () => ({ account_id: 'account-b' })) + const sent: string[] = [] + const serve = createFetch({ + entry: fixtureEntries[0], + placeholder, + consumer: h.consumer, + host, + routing: { + sessionId: 'session', + provider: 'fixture-a', + modelId: 'model', + }, + fetch: async (input) => { + sent.push(new Request(input).headers.get('authorization') ?? '') + if (sent.length === 1) { + h.setVersion(2) + return new Response('old key', { status: 401 }) + } + return new Response('done') + }, + }) + expect( + await ( + await serve('https://fixture.invalid', { headers: bearer }) + ).text(), + ).toBe('done') + expect(sent).toEqual(['Bearer secret-b-1', 'Bearer secret-b-2']) + expect(h.reads).toEqual(['b', 'b']) + expect(h.reports).toEqual([]) + expect(calls).toHaveLength(1) + }) + test(`routing ${host}: selection overrides roster and local marks; no accessor keeps today's choice`, async () => { const h = await harness( [row('a', { accountId: undefined }), row('b', { accountId: undefined })], @@ -289,7 +317,7 @@ for (const host of ['opencode1', 'pi'] as const) { test('routing reports only qualifying balance refusals with vault identities and provider codes', async () => { for (const [service, status, body, code] of [ ['deepseek', 402, '{}', 'http_402'], - ['deepseek', 402, '{"error":{"code":"balance_empty"}}', 'balance_empty'], + ['deepseek', 402, '{"error":{"code":"request_failed"}}', 'http_402'], [ 'openrouter', 402, diff --git a/packages/opencode/test/v2/routing.test.ts b/packages/opencode/test/v2/routing.test.ts index 6fe74bb..ff337bb 100644 --- a/packages/opencode/test/v2/routing.test.ts +++ b/packages/opencode/test/v2/routing.test.ts @@ -166,6 +166,7 @@ test('routing OC2 retry excludes earlier keys and asks again', async () => { test('routing OC2 balance reports once and never for auth rate limits or self-set caps', async () => { for (const [service, status, body, code] of [ ['deepseek', 402, '{}', 'http_402'], + ['deepseek', 402, '{"error":{"code":"request_failed"}}', 'http_402'], [ 'openrouter', 402, From c7fe6f6f945289ac65a432f28dd09db5c59697cc Mon Sep 17 00:00:00 2001 From: ualtinok <94532+ualtinok@users.noreply.github.com> Date: Thu, 12 Nov 2026 17:18:13 +0100 Subject: [PATCH 3/3] mason: isolate shared fixtures across Bun test files Replace the billing helper's module-level afterEach with per-test routing-accessor restoration, and restore only the original vault relay/host marker after nested installations. Restore direct accessor removals and Pi profile environment changes per test as well. Handle Pi contexts without sessionManager: API-key routing sees no session and stops without sending, while non-routed serving remains unchanged. Add a two-file Bun subprocess regression for cached helpers and a Pi missing-session regression. Authorized CI diagnosis: root bun test went from 482 pass/79 fail/83 skip to 562 pass/0 fail/83 skip; the full Pi CI command went from 38 pass/22 fail/2 skip to 61 pass/0 fail/2 skip under Bun 1.4.4 and Node 24. --- packages/core/test/billing-fixtures.ts | 30 ++++------ packages/core/test/env.ts | 21 +++++++ packages/core/test/host.ts | 24 +++++--- .../core/test/routing-accessor-fixture.ts | 24 ++++++++ packages/core/test/routing-fixtures.ts | 26 +++------ packages/core/test/routing.test.ts | 5 +- packages/core/test/state-isolation.test.ts | 57 +++++++++++++++++++ packages/opencode/test/v2/billing.test.ts | 3 +- packages/pi/src/extension.ts | 8 ++- packages/pi/test/billing.test.ts | 6 +- packages/pi/test/extension.test.ts | 21 ++++++- 11 files changed, 168 insertions(+), 57 deletions(-) create mode 100644 packages/core/test/env.ts create mode 100644 packages/core/test/routing-accessor-fixture.ts create mode 100644 packages/core/test/state-isolation.test.ts diff --git a/packages/core/test/billing-fixtures.ts b/packages/core/test/billing-fixtures.ts index 12d99f3..8b121ac 100644 --- a/packages/core/test/billing-fixtures.ts +++ b/packages/core/test/billing-fixtures.ts @@ -1,11 +1,4 @@ -import { afterEach } from 'bun:test' - -const symbol = Symbol.for('cortexkit.routing') -const original = Object.getOwnPropertyDescriptor(globalThis, symbol) -afterEach(() => { - if (original) Object.defineProperty(globalThis, symbol, original) - else Reflect.deleteProperty(globalThis, symbol) -}) +import { installRoutingAccessor } from './routing-accessor-fixture.ts' export function captureBilling( keepWarm = false, @@ -15,18 +8,15 @@ export function captureBilling( ) { const reports: Record[] = [] const sessions: string[] = [] - Object.defineProperty(globalThis, symbol, { - configurable: true, - value: (session: string) => { - sessions.push(session) - return { - accountFor, - report: (report: Record) => { - reports.push(report) - }, - cache: () => (keepWarm ? { keep_warm: 'always' } : null), - } - }, + installRoutingAccessor((session: string) => { + sessions.push(session) + return { + accountFor, + report: (report: Record) => { + reports.push(report) + }, + cache: () => (keepWarm ? { keep_warm: 'always' } : null), + } }) return { reports, sessions } } diff --git a/packages/core/test/env.ts b/packages/core/test/env.ts new file mode 100644 index 0000000..7d71e22 --- /dev/null +++ b/packages/core/test/env.ts @@ -0,0 +1,21 @@ +import { onTestFinished } from 'bun:test' + +const previous = new Map() +let restoring = false + +export function setTestEnv(name: string, value: string | undefined) { + if (!previous.has(name)) previous.set(name, process.env[name]) + if (!restoring) { + restoring = true + onTestFinished(() => { + for (const [key, original] of previous) { + if (original === undefined) delete process.env[key] + else process.env[key] = original + } + previous.clear() + restoring = false + }) + } + if (value === undefined) delete process.env[name] + else process.env[name] = value +} diff --git a/packages/core/test/host.ts b/packages/core/test/host.ts index d824a79..3ce4db0 100644 --- a/packages/core/test/host.ts +++ b/packages/core/test/host.ts @@ -1,12 +1,26 @@ import { onTestFinished } from 'bun:test' import type { ClaustrumScopedClient } from '@cortexkit/common-auth/claustrum' +let restoring = false + /** Installs a per-call owned relay client and restores the ambient host after a test. */ export function installVaultHost(client?: ClaustrumScopedClient) { const key = Symbol.for('cortexkit.vault') const globals = globalThis as unknown as Record - const previous = globals[key] - const marker = process.env.CORTEXKIT_VAULT_HOST + if (!restoring) { + const previous = Object.getOwnPropertyDescriptor(globalThis, key) + const marker = process.env.CORTEXKIT_VAULT_HOST + restoring = true + // Nested relay installs share one ambient restore. Per-call restores run + // in registration order and would otherwise reinstall a previous relay. + onTestFinished(() => { + restoring = false + if (marker === undefined) delete process.env.CORTEXKIT_VAULT_HOST + else process.env.CORTEXKIT_VAULT_HOST = marker + if (previous) Object.defineProperty(globalThis, key, previous) + else Reflect.deleteProperty(globalThis, key) + }) + } const calls: string[] = [] let closed = 0 process.env.CORTEXKIT_VAULT_HOST = 'thalamus' @@ -21,11 +35,5 @@ export function installVaultHost(client?: ClaustrumScopedClient) { }, } } - onTestFinished(() => { - if (marker === undefined) delete process.env.CORTEXKIT_VAULT_HOST - else process.env.CORTEXKIT_VAULT_HOST = marker - if (previous === undefined) delete globals[key] - else globals[key] = previous - }) return { calls, closed: () => closed } } diff --git a/packages/core/test/routing-accessor-fixture.ts b/packages/core/test/routing-accessor-fixture.ts new file mode 100644 index 0000000..3c1955b --- /dev/null +++ b/packages/core/test/routing-accessor-fixture.ts @@ -0,0 +1,24 @@ +import { onTestFinished } from 'bun:test' + +const symbol = Symbol.for('cortexkit.routing') +let restoring = false + +export function installRoutingAccessor(accessor?: unknown) { + if (!restoring) { + const previous = Object.getOwnPropertyDescriptor(globalThis, symbol) + restoring = true + // Several helpers can replace the accessor in one test. Restore the + // original once, not an intermediate accessor from an earlier helper. + onTestFinished(() => { + restoring = false + if (previous) Object.defineProperty(globalThis, symbol, previous) + else Reflect.deleteProperty(globalThis, symbol) + }) + } + if (accessor === undefined) Reflect.deleteProperty(globalThis, symbol) + else + Object.defineProperty(globalThis, symbol, { + configurable: true, + value: accessor, + }) +} diff --git a/packages/core/test/routing-fixtures.ts b/packages/core/test/routing-fixtures.ts index ef3da59..6b34777 100644 --- a/packages/core/test/routing-fixtures.ts +++ b/packages/core/test/routing-fixtures.ts @@ -1,24 +1,14 @@ -import { onTestFinished } from 'bun:test' - -const symbol = Symbol.for('cortexkit.routing') -const original = Object.getOwnPropertyDescriptor(globalThis, symbol) +import { installRoutingAccessor } from './routing-accessor-fixture.ts' export function routing(answer: (call: number) => Promise) { - onTestFinished(() => { - if (original) Object.defineProperty(globalThis, symbol, original) - else Reflect.deleteProperty(globalThis, symbol) - }) const calls: unknown[] = [] const reports: unknown[] = [] - Object.defineProperty(globalThis, symbol, { - configurable: true, - value: (session: string) => ({ - accountFor: (request: unknown, options: unknown) => { - calls.push({ session, request, options }) - return answer(calls.length) - }, - report: (payload: unknown) => reports.push(payload), - }), - }) + installRoutingAccessor((session: string) => ({ + accountFor: (request: unknown, options: unknown) => { + calls.push({ session, request, options }) + return answer(calls.length) + }, + report: (payload: unknown) => reports.push(payload), + })) return { calls, reports } } diff --git a/packages/core/test/routing.test.ts b/packages/core/test/routing.test.ts index 60e4f7f..82fa592 100644 --- a/packages/core/test/routing.test.ts +++ b/packages/core/test/routing.test.ts @@ -1,10 +1,9 @@ import { expect, test } from 'bun:test' import { createFetch, fixtureEntries } from '../test-entry.ts' import { bearer, harness, placeholder, row } from './fakes.ts' +import { installRoutingAccessor } from './routing-accessor-fixture.ts' import { routing } from './routing-fixtures.ts' -const symbol = Symbol.for('cortexkit.routing') - for (const host of ['opencode1', 'pi'] as const) { test(`routing ${host}: OAuth-only subscriptions serve without asking routing`, async () => { const h = await harness( @@ -230,7 +229,7 @@ for (const host of ['opencode1', 'pi'] as const) { return new Response('ok') }, }) - Reflect.deleteProperty(globalThis, symbol) + installRoutingAccessor() await serve('https://fixture.invalid', { headers: bearer }) expect(sent).toEqual(['Bearer secret-a-1']) for (const r of h.consumer.routingRows()) diff --git a/packages/core/test/state-isolation.test.ts b/packages/core/test/state-isolation.test.ts new file mode 100644 index 0000000..0882c30 --- /dev/null +++ b/packages/core/test/state-isolation.test.ts @@ -0,0 +1,57 @@ +import { expect, onTestFinished, test } from 'bun:test' +import { mkdir, mkdtemp, rm, writeFile } from 'node:fs/promises' +import path from 'node:path' + +test('shared fixtures restore routing and nested vault relays across test files', async () => { + const scratch = path.join(import.meta.dir, '../../../target/scratch') + await mkdir(scratch, { recursive: true }) + const dir = await mkdtemp(path.join(scratch, 'fixture-isolation-')) + onTestFinished(() => rm(dir, { recursive: true, force: true })) + const billing = JSON.stringify( + new URL('./billing-fixtures.ts', import.meta.url).href, + ) + const host = JSON.stringify(new URL('./host.ts', import.meta.url).href) + const environment = JSON.stringify(new URL('./env.ts', import.meta.url).href) + const source = ` +import { expect, test } from 'bun:test' +import { captureBilling } from ${billing} +import { installVaultHost } from ${host} +import { setTestEnv } from ${environment} +test('install shared helpers more than once', () => { + captureBilling() + captureBilling() + installVaultHost() + installVaultHost() + setTestEnv('PI_CODING_AGENT_DIR', 'first-test-profile') + setTestEnv('PI_CODING_AGENT_DIR', 'second-test-profile') +}) +test('ambient state is restored after the test', () => { + expect({ + routing: Reflect.get(globalThis, Symbol.for('cortexkit.routing')), + vault: Reflect.get(globalThis, Symbol.for('cortexkit.vault')), + host: process.env.CORTEXKIT_VAULT_HOST, + profile: process.env.PI_CODING_AGENT_DIR, + }).toEqual({ routing: undefined, vault: undefined, host: undefined, profile: undefined }) +}) +` + // Both files import the same cached helper. A module-level afterEach would + // clean up only the first file, and per-install relay cleanup would retain + // the earlier relay's hosted-mode marker after nested installations. + const files = ['first.test.ts', 'second.test.ts'].map((name) => + path.join(dir, name), + ) + await Promise.all(files.map((file) => writeFile(file, source))) + const env = { ...process.env } + delete env.CORTEXKIT_VAULT_HOST + delete env.CLAUSTRUM_AUTH_E2E_ROUTING + delete env.PI_CODING_AGENT_DIR + const child = Bun.spawnSync([process.execPath, 'test', ...files], { + cwd: dir, + env, + }) + const output = Buffer.concat([child.stdout, child.stderr]).toString() + expect({ exit: child.exitCode, output }).toMatchObject({ + exit: 0, + output: expect.stringContaining('4 pass'), + }) +}) diff --git a/packages/opencode/test/v2/billing.test.ts b/packages/opencode/test/v2/billing.test.ts index aac00a1..ea68df9 100644 --- a/packages/opencode/test/v2/billing.test.ts +++ b/packages/opencode/test/v2/billing.test.ts @@ -11,6 +11,7 @@ import { expectedUsage, } from '../../../core/test/billing-fixtures.ts' import { vaultRowFor } from '../../../core/test/fakes.ts' +import { installRoutingAccessor } from '../../../core/test/routing-accessor-fixture.ts' import { vaultMethodId } from '../../src/v2/setup.ts' import { scope, startServer, trigger } from './fakes.ts' @@ -111,7 +112,7 @@ test('a reply model alias keeps one report at the sent model without usageByMode }) test('without an accessor a served host request never looks up billing semantics', async () => { - Reflect.deleteProperty(globalThis, Symbol.for('cortexkit.routing')) + installRoutingAccessor() const base = supportedProviders.find((entry) => entry.service === 'deepseek') if (!base?.billing) throw new Error('missing deepseek') let parsed = 0 diff --git a/packages/pi/src/extension.ts b/packages/pi/src/extension.ts index b36854b..397fcb8 100644 --- a/packages/pi/src/extension.ts +++ b/packages/pi/src/extension.ts @@ -331,8 +331,10 @@ class PiRuntime { // runtime override goes to pi-ai exactly as Pi resolved it. if (options?.apiKey !== served.placeholder) return streamSimple(model, context, options) + const sessionId = + options.sessionId ?? this.#ctx.sessionManager?.getSessionId() const routing = { - sessionId: options.sessionId ?? this.#ctx.sessionManager.getSessionId(), + sessionId, provider: model.provider, modelId: model.id, } @@ -356,11 +358,11 @@ class PiRuntime { inner, (receipt) => { current?.end() - if (!hasBillingAccessor()) return undefined + if (!sessionId || !hasBillingAccessor()) return undefined current = new BillingAttempt( served.entry, model.provider, - options.sessionId ?? this.#ctx.sessionManager.getSessionId(), + sessionId, 'primary', model.id, receipt, diff --git a/packages/pi/test/billing.test.ts b/packages/pi/test/billing.test.ts index f291563..f4cba36 100644 --- a/packages/pi/test/billing.test.ts +++ b/packages/pi/test/billing.test.ts @@ -10,8 +10,10 @@ import { chatResponse, expectedUsage, } from '../../core/test/billing-fixtures.ts' +import { setTestEnv } from '../../core/test/env.ts' import { vaultRowFor } from '../../core/test/fakes.ts' import { installVaultHost } from '../../core/test/host.ts' +import { installRoutingAccessor } from '../../core/test/routing-accessor-fixture.ts' import { routing } from '../../core/test/routing-fixtures.ts' import { createTestCore } from '../../core/test-entry.ts' import { createPiExtension, placeholderKey } from '../src/extension.ts' @@ -75,7 +77,7 @@ async function turn( enrolled: false, flags: { [provider]: { consented: true } }, }) - process.env.PI_CODING_AGENT_DIR = agent.agent + setTestEnv('PI_CODING_AGENT_DIR', agent.agent) let registered: Record | undefined const handlers: Record< string, @@ -269,7 +271,7 @@ test('Pi followed same-origin redirect reports only the completed model call', a }) test('Pi without an accessor never looks up usage semantics for a served request', async () => { - Reflect.deleteProperty(globalThis, Symbol.for('cortexkit.routing')) + installRoutingAccessor() let parsed = 0 await turn( 'deepseek', diff --git a/packages/pi/test/extension.test.ts b/packages/pi/test/extension.test.ts index 919dd03..1cd7cc0 100644 --- a/packages/pi/test/extension.test.ts +++ b/packages/pi/test/extension.test.ts @@ -9,8 +9,11 @@ import type { ExtensionAPI, ExtensionContext, } from '@earendil-works/pi-coding-agent' +import { setTestEnv } from '../../core/test/env.ts' import { harness, row } from '../../core/test/fakes.ts' import { installVaultHost } from '../../core/test/host.ts' +import { installRoutingAccessor } from '../../core/test/routing-accessor-fixture.ts' +import { routing } from '../../core/test/routing-fixtures.ts' import { createTestCore } from '../../core/test-entry.ts' import { createPiExtension, placeholderKey } from '../src/extension.ts' import { FlagStore } from '../src/flags.ts' @@ -120,7 +123,7 @@ const servedProvider = 'fixture-a' test('Pi startup refusal names model registration and the remedy without echoing the host error', async () => { const h = await harness([]) const agent = await agentDir({}) - process.env.PI_CODING_AGENT_DIR = agent.agent + setTestEnv('PI_CODING_AGENT_DIR', agent.agent) const handlers: Record< string, (event: unknown, ctx: unknown) => Promise @@ -201,7 +204,7 @@ async function servedTurn(options: { : {}), flags: { [servedProvider]: { consented: true } }, }) - process.env.PI_CODING_AGENT_DIR = agent.agent + setTestEnv('PI_CODING_AGENT_DIR', agent.agent) if (options.hosted === 'missing') await openPiSlot({ authPath: agent.authPath, @@ -342,6 +345,20 @@ test('hosted Pi serves the relay vault key without enrollment or native connect' expect(relay?.calls).toEqual(['claustrum-auth']) }) +test('Pi without a session manager stops an API-key request only under routing', async () => { + installRoutingAccessor() + const outside = await servedTurn({ hosted: 'relay' }) + expect(outside.sends).toBe(1) + expect(outside.authorizations).toEqual(['Bearer secret-a-1']) + const { calls } = routing(async () => ({ account_id: 'account-a' })) + const { sends, message } = await servedTurn({ hosted: 'relay' }) + expect(message).toContain( + 'Thalamus routing did not answer, so this API key request was not sent.', + ) + expect(sends).toBe(0) + expect(calls).toEqual([]) +}) + test('a route withdrawn while a turn is served gets the plugin text, and nothing is sent', async () => { const { message, sends } = await servedTurn({ withdraw(core) {