diff --git a/apps/server/src/provider/Drivers/ClaudeDriver.ts b/apps/server/src/provider/Drivers/ClaudeDriver.ts index c2d4b63861074..aecb0f9eae866 100644 --- a/apps/server/src/provider/Drivers/ClaudeDriver.ts +++ b/apps/server/src/provider/Drivers/ClaudeDriver.ts @@ -177,38 +177,58 @@ export const ClaudeDriver: ProviderDriver(undefined); + // Per-instance capabilities cache: keyed on binary + resolved HOME + logged-in account so + // probes never share auth metadata across instances or across a re-login. + // A transient miss may reuse the same key's last probe once; its usage windows keep their age. + const capabilitiesCacheKey = yield* makeClaudeCapabilitiesCacheKey( + effectiveConfig, + stateDir, + processEnv, + ); + const readCapabilitiesKey = ClaudeResetCredits.readClaudeAccountKey(accountConfigPath).pipe( + Effect.map((accountKey) => `${capabilitiesCacheKey}\0${accountKey}`), + Effect.provideService(FileSystem.FileSystem, fileSystem), + Effect.provideService(Path.Path, path), + ); + const lastGoodCapabilities = yield* Ref.make< + { readonly key: string; readonly capabilities: ClaudeCapabilitiesProbe } | undefined + >(undefined); const capabilitiesProbeCache = yield* Cache.makeWith( - () => + (key: string) => probeClaudeCapabilities(effectiveConfig, processEnv, stateDir).pipe( Effect.flatMap((capabilities) => - ClaudeResetCredits.readClaudeOrganizationId(accountConfigPath).pipe( - Effect.map((workspaceId) => ({ - _tag: "fresh" as const, - capabilities: - workspaceId === undefined ? capabilities : { ...capabilities, workspaceId }, - })), + Effect.all([ + ClaudeResetCredits.readClaudeOrganizationId(accountConfigPath), + readCapabilitiesKey, + ]).pipe( + Effect.map(([workspaceId, currentKey]) => + // A login that landed mid-probe leaves this probe's account unknown: show it once + // without the new login's organization, never cache it, re-probe next check. + currentKey !== key + ? { _tag: "raced" as const, capabilities } + : { + _tag: "fresh" as const, + capabilities: + workspaceId === undefined ? capabilities : { ...capabilities, workspaceId }, + }, + ), ), ), - Effect.tap(({ capabilities }) => Ref.set(lastGoodCapabilities, capabilities)), + Effect.tap((result) => + result._tag === "fresh" + ? Ref.set(lastGoodCapabilities, { key, capabilities: result.capabilities }) + : Effect.void, + ), Effect.catch((error) => - Ref.getAndSet(lastGoodCapabilities, undefined).pipe( - Effect.flatMap((capabilities) => - capabilities === undefined - ? Effect.fail(error) - : ClaudeResetCredits.readClaudeOrganizationId(accountConfigPath).pipe( - Effect.flatMap((workspaceId) => - workspaceId === capabilities.workspaceId - ? Effect.succeed({ - _tag: "fallback" as const, - capabilities: withoutUsage(capabilities), - }) - : Effect.fail(error), - ), - ), + // Re-read the account: a login can land while the failed probe was running. + Effect.all([Ref.getAndSet(lastGoodCapabilities, undefined), readCapabilitiesKey]).pipe( + Effect.flatMap(([lastGood, currentKey]) => + lastGood?.key === key && currentKey === key + ? Effect.succeed({ + _tag: "fallback" as const, + capabilities: withoutUsage(lastGood.capabilities), + }) + : Effect.fail(error), ), ), ), @@ -226,12 +246,8 @@ export const ClaudeDriver: ProviderDriver Cache.get(capabilitiesProbeCache, key)), Effect.map((result) => result.capabilities), ); diff --git a/apps/server/src/provider/ProviderInstanceRegistry.test.ts b/apps/server/src/provider/ProviderInstanceRegistry.test.ts index e998332d3e92f..95db418db2f38 100644 --- a/apps/server/src/provider/ProviderInstanceRegistry.test.ts +++ b/apps/server/src/provider/ProviderInstanceRegistry.test.ts @@ -2,6 +2,8 @@ import * as ServerSecretStore from "../auth/ServerSecretStore.ts"; import * as CodexInstallation from "./CodexInstallation.ts"; import * as CodexManagedTokenBroker from "./CodexManagedTokenBroker.ts"; import * as ServerEnvironment from "../environment/ServerEnvironment.ts"; +// @effect-diagnostics-next-line nodeBuiltinImport:off - the SDK mock rewrites the login mid-probe, outside Effect. +import * as NodeFS from "node:fs"; import * as ClaudeSdk from "@anthropic-ai/claude-agent-sdk"; import { describe, expect, it } from "@effect/vitest"; import { vi } from "vite-plus/test"; @@ -484,14 +486,23 @@ describe("ProviderInstanceRegistry — multi-instance codex slice", () => { let probeCount = 0; let successfulProbeCount = 0; let failProbe = false; + let loginDuringFailedProbe: string | undefined; const query = vi.spyOn(ClaudeSdk, "query").mockImplementation( () => ({ - initializationResult: () => { + initializationResult: async () => { probeCount++; - if (failProbe) return Promise.reject(new Error("initialization failed")); + if (failProbe) { + if (loginDuringFailedProbe !== undefined) { + NodeFS.writeFileSync( + accountConfigPath, + JSON.stringify({ oauthAccount: { organizationUuid: loginDuringFailedProbe } }), + ); + } + throw new Error("initialization failed"); + } successfulProbeCount++; - return Promise.resolve({ + return { commands: [ { name: successfulProbeCount === 1 ? "last-good" : "recovered", @@ -505,7 +516,7 @@ describe("ProviderInstanceRegistry — multi-instance codex slice", () => { subscriptionType: "pro", tokenSource: "oauth", }, - }); + }; }, usage_EXPERIMENTAL_MAY_CHANGE_DO_NOT_RELY_ON_THIS_API_YET: async () => ({ rate_limits_available: true, @@ -600,6 +611,139 @@ describe("ProviderInstanceRegistry — multi-instance codex slice", () => { expect.arrayContaining([expect.objectContaining({ name: "recovered" })]), ); expect(probeCount).toBeGreaterThan(sameAccountFailureProbeCount); + + failProbe = false; + yield* claude!.snapshot.refresh; + yield* TestClock.adjust("1 second"); + yield* claude!.invalidateCaches ?? Effect.void; + failProbe = true; + loginDuringFailedProbe = "org-c"; + const loginDuringFailure = yield* claude!.snapshot.refresh; + expect(loginDuringFailure).toMatchObject({ + status: "warning", + auth: { status: "unknown" }, + message: "Claude Agent CLI failed while checking account status.", + }); + expect(loginDuringFailure.usageLimits).toBeUndefined(); + }).pipe(Effect.provide(layerTest)), + ); + + it.effect("re-probes Claude after a login switches the account", () => + Effect.gen(function* () { + if (yield* HostProcess.isWindows) return; + + const fixtures = yield* makeTildeProviderFixtures(); + const fs = yield* FileSystem.FileSystem; + const path = yield* Path.Path; + const accountConfigPath = path.join(fixtures.claudeHomePath, ".claude.json"); + const logIn = (account: { email: string; organizationUuid: string; utilization: number }) => + fs + .writeFileString( + accountConfigPath, + JSON.stringify({ + oauthAccount: { + accountUuid: `${account.organizationUuid}-user`, + organizationUuid: account.organizationUuid, + emailAddress: account.email, + }, + }), + ) + .pipe(Effect.tap(() => Effect.sync(() => (loggedIn = account)))); + let loggedIn = { + email: "account-a@example.com", + organizationUuid: "org-a", + utilization: 100, + }; + yield* logIn(loggedIn); + let probeCount = 0; + let loginDuringProbe: string | undefined; + const query = vi.spyOn(ClaudeSdk, "query").mockImplementation( + () => + ({ + initializationResult: () => { + probeCount++; + const email = loggedIn.email; + if (loginDuringProbe !== undefined) { + const organizationUuid = loginDuringProbe; + loginDuringProbe = undefined; + loggedIn = { + email: `${organizationUuid}@example.com`, + organizationUuid, + utilization: 50, + }; + NodeFS.writeFileSync( + accountConfigPath, + JSON.stringify({ + oauthAccount: { accountUuid: `${organizationUuid}-user`, organizationUuid }, + }), + ); + } + return Promise.resolve({ + commands: [] as Array, + account: { email, subscriptionType: "pro", tokenSource: "oauth" }, + }); + }, + usage_EXPERIMENTAL_MAY_CHANGE_DO_NOT_RELY_ON_THIS_API_YET: async () => ({ + rate_limits_available: true, + rate_limits: { + five_hour: { utilization: loggedIn.utilization, resets_at: "2026-07-18T14:39:00Z" }, + }, + }), + }) as ReturnType, + ); + yield* Effect.addFinalizer(() => Effect.sync(() => query.mockRestore())); + + const claudeId = ProviderInstanceId.make("claude_account_switch"); + const { registry } = yield* makeProviderInstanceRegistry({ + drivers: [ClaudeDriver], + configMap: { + [claudeId]: { + driver: ProviderDriverKind.make("claudeAgent"), + enabled: true, + config: makeClaudeConfig({ + enabled: true, + binaryPath: fixtures.claudeBinaryPath, + homePath: fixtures.claudeHomePath, + }), + }, + }, + }); + const claude = yield* registry.getInstance(claudeId); + expect(claude).toBeDefined(); + + const before = yield* claude!.snapshot.refresh; + expect(before.auth).toMatchObject({ + status: "authenticated", + email: "account-a@example.com", + workspaceId: "org-a", + }); + const probesBeforeLogin = probeCount; + + yield* TestClock.adjust("1 second"); + yield* logIn({ email: "account-b@example.com", organizationUuid: "org-b", utilization: 10 }); + const after = yield* claude!.snapshot.refresh; + + expect( + { auth: after.auth, usageWindows: after.usageLimits?.windows }, + "the check after a re-login must report the new account, not the old account's cached probe", + ).toMatchObject({ + auth: { status: "authenticated", email: "account-b@example.com", workspaceId: "org-b" }, + usageWindows: [expect.objectContaining({ usedPercent: 10 })], + }); + expect(probeCount).toBeGreaterThan(probesBeforeLogin); + + yield* TestClock.adjust("1 second"); + yield* logIn({ email: "account-a@example.com", organizationUuid: "org-a", utilization: 100 }); + loginDuringProbe = "org-c"; + const raced = yield* claude!.snapshot.refresh; + expect( + raced.auth, + "a probe that raced a login must not pair its account with the new login's organization", + ).toMatchObject({ email: "account-a@example.com" }); + expect(raced.auth.workspaceId).toBeUndefined(); + + const settled = yield* claude!.snapshot.refresh; + expect(settled.auth).toMatchObject({ email: "org-c@example.com", workspaceId: "org-c" }); }).pipe(Effect.provide(layerTest)), ); diff --git a/apps/server/src/provider/claudeResetCredits.ts b/apps/server/src/provider/claudeResetCredits.ts index 35973378edada..30a83acad46c0 100644 --- a/apps/server/src/provider/claudeResetCredits.ts +++ b/apps/server/src/provider/claudeResetCredits.ts @@ -22,7 +22,10 @@ const Credentials = Schema.Struct({ }); const Config = Schema.Struct({ oauthAccount: Schema.optional( - Schema.Struct({ organizationUuid: Schema.optional(Schema.String) }), + Schema.Struct({ + accountUuid: Schema.optional(Schema.String), + organizationUuid: Schema.optional(Schema.String), + }), ), }); const Grant = Schema.Struct({ @@ -172,12 +175,26 @@ export const readClaudeResetCredits = Effect.fn("readClaudeResetCredits")( Effect.orElseSucceed(() => undefined), ); -export const readClaudeOrganizationId = (accountConfigPath: string) => +const readClaudeOauthAccount = (accountConfigPath: string) => readJson(Config, accountConfigPath).pipe( - Effect.map((config) => config.oauthAccount?.organizationUuid?.trim() || undefined), + Effect.map((config) => config.oauthAccount), Effect.orElseSucceed(() => undefined), ); +export const readClaudeOrganizationId = (accountConfigPath: string) => + readClaudeOauthAccount(accountConfigPath).pipe( + Effect.map((account) => account?.organizationUuid?.trim() || undefined), + ); + +/** Changes when a login switches account; constant when signed out, on an API key, or unreadable. */ +export const readClaudeAccountKey = (accountConfigPath: string) => + readClaudeOauthAccount(accountConfigPath).pipe( + Effect.map( + (account) => + `${account?.accountUuid?.trim() ?? ""}\0${account?.organizationUuid?.trim() ?? ""}`, + ), + ); + export const claudeAccountConfigPath = Effect.fn("claudeAccountConfigPath")(function* ( configDir: string | undefined, ) {