fix(ai): broker refresh sentinel is not row identity; org-rotation cache regression tests

Addresses the third review round (Codex on 31cd0e4f5 + roboomp):

- Usage-path refresh row matching no longer compares the shared
  REMOTE_REFRESH_SENTINEL as a refresh-token identity — with broker-backed
  rows every credential carries the same sentinel, so the comparison
  matched the FIRST OAuth row and could persist another org's rotated
  token into the wrong row. The sentinel is now non-identifying; matching
  falls through to access token, then account/email/project/org equality.
  Regression: two same-email sentinel rows, expired Max row — the refresh
  persists into the Max row only.
- roboomp's status-line finding (usage context key without orgId) was
  already fixed in 31cd0e4f5 — its review raced the push. Added the
  requested regression: rotating orgs under one email/account invalidates
  the cached usage immediately; same-org calls stay cached.
This commit is contained in:
chan1103
2026-07-11 22:49:12 +09:00
parent a47c3c90ec
commit 3df97d5097
4 changed files with 112 additions and 7 deletions
+1
View File
@@ -9,6 +9,7 @@
- Fixed Anthropic logins silently replacing the stored credential when one account email holds multiple organizations (e.g. a Team seat plus a personal Max plan). Credentials are now identified by email + organization: the login flow captures the organization from the token exchange (with a `claude_cli/bootstrap` fallback), both subscriptions store side by side, and the existing multi-account rotation treats them as separate accounts. Legacy email-keyed rows are claimed in place by the first org-scoped login with the same email, and an org-less credential never clobbers org-scoped rows. Usage reports and the per-credential usage cache also partition by organization so the two subscriptions' limit pools no longer merge into one confused row.
- Fixed broker-served usage routing for org-scoped credentials: `RemoteAuthCredentialStore` now matches aggregate reports and keys header-ingest overlays by organization first, so with a Team seat exhausted and a personal Max healthy under one email, each credential receives its own pool instead of whichever report appeared first. The Anthropic usage-cache key version was bumped so pre-org cache entries (including the 24h last-good fallback) cannot be replayed across organizations.
- Fixed OAuth access results (`getOAuthAccess`, `getOAuthAccesses`, `getOAuthAccessAt`) and `checkCredentials` health results dropping the organization: they now carry `orgId`/`orgName` so consumers that key or label per-account results (e.g. `omp dry-balance --bench`, `omp auth-gateway check`) can tell two same-email subscriptions apart. Active-account matching is org-decisive whenever either side carries an organization — an org-scoped session no longer flags the legacy bare-email row, and a legacy-row session no longer flags org-scoped siblings, via the shared email. `getOAuthAccountIdentity` also preserves org-only identities instead of discarding them.
- Fixed usage-path OAuth refreshes on broker-backed credentials persisting the rotated token into the wrong row: the shared `REMOTE_REFRESH_SENTINEL` refresh value is no longer treated as row identity when locating the row to update, so with two same-email organizations the refreshed token lands on the org that was actually refreshed.
## [16.4.3] - 2026-07-11
+10 -2
View File
@@ -2569,9 +2569,14 @@ export class AuthStorage {
*/
#findStoredCredentialIdForUsageCredential(provider: Provider, previous: UsageCredential): number | undefined {
const entries = this.#getStoredCredentials(provider);
// Broker-backed rows all carry REMOTE_REFRESH_SENTINEL as their refresh
// token — it identifies nothing, and comparing it would match the FIRST
// OAuth row regardless of which account/org is being refreshed.
const previousRefresh =
previous.refreshToken && previous.refreshToken !== REMOTE_REFRESH_SENTINEL ? previous.refreshToken : undefined;
const match = entries.find(entry => {
if (entry.credential.type !== "oauth") return false;
if (previous.refreshToken && entry.credential.refresh === previous.refreshToken) return true;
if (previousRefresh && entry.credential.refresh === previousRefresh) return true;
if (previous.accessToken && entry.credential.access === previous.accessToken) return true;
return (
entry.credential.accountId === previous.accountId &&
@@ -2585,9 +2590,12 @@ export class AuthStorage {
#persistRefreshedUsageCredential(provider: Provider, previous: UsageCredential, next: UsageCredential): void {
const entries = this.#getStoredCredentials(provider);
// Same sentinel rule as #findStoredCredentialIdForUsageCredential above.
const previousRefresh =
previous.refreshToken && previous.refreshToken !== REMOTE_REFRESH_SENTINEL ? previous.refreshToken : undefined;
const index = entries.findIndex(entry => {
if (entry.credential.type !== "oauth") return false;
if (previous.refreshToken && entry.credential.refresh === previous.refreshToken) return true;
if (previousRefresh && entry.credential.refresh === previousRefresh) return true;
if (previous.accessToken && entry.credential.access === previous.accessToken) return true;
return (
entry.credential.accountId === previous.accountId &&
@@ -22,6 +22,7 @@ import {
type AuthCredential,
type AuthCredentialStore,
AuthStorage,
REMOTE_REFRESH_SENTINEL,
SqliteAuthCredentialStore,
type StoredAuthCredential,
} from "@oh-my-pi/pi-ai/auth-storage";
@@ -137,14 +138,19 @@ interface CacheEntry {
expiresAtSec: number;
}
function makeStore(rows: StoredAuthCredential[]): AuthCredentialStore {
function makeStore(
rows: StoredAuthCredential[],
onUpdate?: (id: number, credential: AuthCredential) => void,
): AuthCredentialStore {
const cache = new Map<string, CacheEntry>();
return {
close() {},
listAuthCredentials() {
return rows;
},
updateAuthCredential() {},
updateAuthCredential(id, credential) {
onUpdate?.(id, credential);
},
deleteAuthCredential() {},
tryDisableAuthCredentialIfMatches() {
return false;
@@ -169,15 +175,20 @@ function makeStore(rows: StoredAuthCredential[]): AuthCredentialStore {
};
}
function oauthRow(id: number, orgId?: string, orgName?: string): StoredAuthCredential {
function oauthRow(
id: number,
orgId?: string,
orgName?: string,
overrides?: { refresh?: string; expires?: number },
): StoredAuthCredential {
return {
id,
provider: "anthropic",
credential: {
type: "oauth",
access: `oat-${id}`,
refresh: `refresh-${id}`,
expires: Date.now() + 3_600_000,
refresh: overrides?.refresh ?? `refresh-${id}`,
expires: overrides?.expires ?? Date.now() + 3_600_000,
accountId: "account-shared",
email: EMAIL,
orgId,
@@ -266,3 +277,46 @@ describe("anthropic usage report dedupe partitions by org", () => {
expect(reports).toHaveLength(1);
});
});
describe("broker-backed refresh row addressing", () => {
let storage: AuthStorage | null = null;
afterEach(() => {
storage?.close();
storage = null;
vi.restoreAllMocks();
});
it("persists a refresh into the refreshed org's row, not the first sentinel row", async () => {
// Broker snapshots replace every refresh token with the shared
// REMOTE_REFRESH_SENTINEL. The sentinel must not act as row identity:
// with two same-email orgs, refreshing the expired Max row previously
// matched the Team row first and persisted the new token there.
const teamRow = oauthRow(1, TEAM_ORG, "Team Workspace", { refresh: REMOTE_REFRESH_SENTINEL });
const maxRow = oauthRow(2, MAX_ORG, "Personal Max", {
refresh: REMOTE_REFRESH_SENTINEL,
expires: Date.now() - 1_000,
});
const updates: number[] = [];
storage = new AuthStorage(
makeStore([teamRow, maxRow], id => {
updates.push(id);
}),
{
usageProviderResolver: provider => (provider === "anthropic" ? claudeUsage.claudeUsageProvider : undefined),
refreshOAuthCredential: async () => ({
access: "refreshed-access",
refresh: REMOTE_REFRESH_SENTINEL,
expires: Date.now() + 3_600_000,
}),
},
);
await storage.reload();
vi.spyOn(claudeUsage.claudeUsageProvider, "fetchUsage").mockImplementation(async () => emailOnlyReport());
await storage.fetchUsageReports();
// Only the expired Max row (id 2) was rewritten; the Team row was never touched.
expect(updates).toEqual([2]);
});
});
@@ -158,4 +158,46 @@ describe("StatusLineComponent usage refresh", () => {
expect(plain(component.getTopBorder(80).content)).toContain("5h 42%");
});
it("re-fetches usage immediately when the session rotates to another org under the same email", async () => {
let calls = 0;
let orgId = "org-team";
const base = makeSession(async () => {
calls++;
return usageReport(10);
}) as unknown as Record<string, unknown>;
// Same provider + email + account throughout — only the org rotates.
base.state = {
messages: [],
model: { contextWindow: 200_000, provider: "anthropic" },
};
base.modelRegistry = {
authStorage: {
getOAuthAccountIdentity: () => ({
email: "shared@example.com",
accountId: "account-shared",
orgId,
}),
},
};
const component = new StatusLineComponent(base as unknown as AgentSession);
component.refreshUsageInBackground();
vi.advanceTimersByTime(0);
await flushMicrotasks();
expect(calls).toBe(1);
// Same org within the cache TTL: served from cache, no refetch.
component.refreshUsageInBackground();
vi.advanceTimersByTime(0);
await flushMicrotasks();
expect(calls).toBe(1);
// Org rotation under the same email/account must invalidate the cache.
orgId = "org-max";
component.refreshUsageInBackground();
vi.advanceTimersByTime(0);
await flushMicrotasks();
expect(calls).toBe(2);
});
});