Skip to content

Commit 9ba8e5d

Browse files
ndycodendycode
committed
fix(identity): preserve org-scoped variants with distinct account ids
1 parent 2f3e2f4 commit 9ba8e5d

File tree

2 files changed

+151
-23
lines changed

2 files changed

+151
-23
lines changed

‎index.ts‎

Lines changed: 81 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -624,7 +624,7 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
624624

625625

626626
type IdentityIndexes = {
627-
byOrganizationId: Map<string, number>;
627+
byOrganizationId: Map<string, number[]>;
628628
byAccountIdNoOrg: Map<string, number>;
629629
byRefreshTokenNoOrg: Map<string, number[]>;
630630
byEmailNoOrg: Map<string, number>;
@@ -633,6 +633,56 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
633633
byRefreshTokenGlobal: Map<string, number[]>;
634634
};
635635

636+
const resolveOrganizationMatch = (
637+
indexes: IdentityIndexes,
638+
organizationId: string,
639+
candidateAccountId: string | undefined,
640+
): number | undefined => {
641+
const matches = indexes.byOrganizationId.get(organizationId);
642+
if (!matches || matches.length === 0) return undefined;
643+
644+
const candidateId = candidateAccountId?.trim() || undefined;
645+
let newestNoAccountId: number | undefined;
646+
let newestExactAccountId: number | undefined;
647+
let newestAnyNonEmptyAccountId: number | undefined;
648+
let nonEmptyAccountIdCount = 0;
649+
650+
for (const index of matches) {
651+
const existing = accounts[index];
652+
if (!existing) continue;
653+
const existingAccountId = normalizeStoredAccountId(existing);
654+
if (!existingAccountId) {
655+
newestNoAccountId =
656+
typeof newestNoAccountId === "number"
657+
? pickNewestAccountIndex(newestNoAccountId, index)
658+
: index;
659+
continue;
660+
}
661+
nonEmptyAccountIdCount += 1;
662+
newestAnyNonEmptyAccountId =
663+
typeof newestAnyNonEmptyAccountId === "number"
664+
? pickNewestAccountIndex(newestAnyNonEmptyAccountId, index)
665+
: index;
666+
if (candidateId && existingAccountId === candidateId) {
667+
newestExactAccountId =
668+
typeof newestExactAccountId === "number"
669+
? pickNewestAccountIndex(newestExactAccountId, index)
670+
: index;
671+
}
672+
}
673+
674+
if (candidateId) {
675+
return newestExactAccountId ?? newestNoAccountId;
676+
}
677+
if (typeof newestNoAccountId === "number") {
678+
return newestNoAccountId;
679+
}
680+
if (nonEmptyAccountIdCount === 1) {
681+
return newestAnyNonEmptyAccountId;
682+
}
683+
return undefined;
684+
};
685+
636686
const resolveNoOrgRefreshMatch = (
637687
indexes: IdentityIndexes,
638688
refreshToken: string,
@@ -683,7 +733,7 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
683733
};
684734

685735
const buildIdentityIndexes = (): IdentityIndexes => {
686-
const byOrganizationId = new Map<string, number>();
736+
const byOrganizationId = new Map<string, number[]>();
687737
const byAccountIdNoOrg = new Map<string, number>();
688738
const byRefreshTokenNoOrg = new Map<string, number[]>();
689739
const byEmailNoOrg = new Map<string, number>();
@@ -707,7 +757,7 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
707757
}
708758

709759
if (organizationId) {
710-
byOrganizationId.set(organizationId, i);
760+
pushIndex(byOrganizationId, organizationId, i);
711761
if (accountId) {
712762
pushIndex(byAccountIdOrgScoped, accountId, i);
713763
}
@@ -755,7 +805,11 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
755805

756806
const existingIndex = (() => {
757807
if (organizationId) {
758-
return identityIndexes.byOrganizationId.get(organizationId);
808+
return resolveOrganizationMatch(
809+
identityIndexes,
810+
organizationId,
811+
normalizedAccountId,
812+
);
759813
}
760814
if (normalizedAccountId) {
761815
const byAccountId = identityIndexes.byAccountIdNoOrg.get(normalizedAccountId);
@@ -859,7 +913,7 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
859913
const refreshMap = new Map<
860914
string,
861915
{
862-
byOrg: Map<string, number>;
916+
byOrg: Map<string, number[]>;
863917
preferredOrgIndex?: number;
864918
fallbackNoAccountIdIndex?: number;
865919
fallbackByAccountId: Map<string, number>;
@@ -875,7 +929,7 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
875929
};
876930

877931
const collapseFallbackIntoPreferredOrg = (entry: {
878-
byOrg: Map<string, number>;
932+
byOrg: Map<string, number[]>;
879933
preferredOrgIndex?: number;
880934
fallbackNoAccountIdIndex?: number;
881935
fallbackByAccountId: Map<string, number>;
@@ -929,7 +983,7 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
929983
let entry = refreshMap.get(refreshToken);
930984
if (!entry) {
931985
entry = {
932-
byOrg: new Map<string, number>(),
986+
byOrg: new Map<string, number[]>(),
933987
preferredOrgIndex: undefined,
934988
fallbackNoAccountIdIndex: undefined,
935989
fallbackByAccountId: new Map<string, number>(),
@@ -938,18 +992,35 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
938992
}
939993

940994
if (orgKey) {
941-
const existingIndex = entry.byOrg.get(orgKey);
995+
const orgMatches = entry.byOrg.get(orgKey) ?? [];
996+
const existingIndex = resolveOrganizationMatch(
997+
{
998+
byOrganizationId: new Map([[orgKey, orgMatches]]),
999+
byAccountIdNoOrg: new Map(),
1000+
byRefreshTokenNoOrg: new Map(),
1001+
byEmailNoOrg: new Map(),
1002+
byAccountIdOrgScoped: new Map(),
1003+
byRefreshTokenOrgScoped: new Map(),
1004+
byRefreshTokenGlobal: new Map(),
1005+
},
1006+
orgKey,
1007+
normalizeStoredAccountId(account),
1008+
);
9421009
if (existingIndex !== undefined) {
9431010
const newestIndex = pickNewestAccountIndex(existingIndex, i);
9441011
const obsoleteIndex = newestIndex === existingIndex ? i : existingIndex;
9451012
mergeAccountRecords(newestIndex, obsoleteIndex);
9461013
indicesToRemove.add(obsoleteIndex);
947-
entry.byOrg.set(orgKey, newestIndex);
1014+
const nextOrgMatches = orgMatches.filter(
1015+
(index) => index !== obsoleteIndex && index !== newestIndex,
1016+
);
1017+
nextOrgMatches.push(newestIndex);
1018+
entry.byOrg.set(orgKey, nextOrgMatches);
9481019
entry.preferredOrgIndex = pickPreferredOrgIndex(entry.preferredOrgIndex, newestIndex);
9491020
collapseFallbackIntoPreferredOrg(entry);
9501021
continue;
9511022
}
952-
entry.byOrg.set(orgKey, i);
1023+
entry.byOrg.set(orgKey, [...orgMatches, i]);
9531024
entry.preferredOrgIndex = pickPreferredOrgIndex(entry.preferredOrgIndex, i);
9541025
collapseFallbackIntoPreferredOrg(entry);
9551026
continue;

‎test/index.test.ts‎

Lines changed: 70 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1916,7 +1916,7 @@ describe("OpenAIOAuthPlugin persistAccountPool", () => {
19161916
]);
19171917
});
19181918

1919-
it("collapses duplicate organization candidates during persistence", async () => {
1919+
it("preserves duplicate organization candidates when accountId differs", async () => {
19201920
const accountsModule = await import("../lib/accounts.js");
19211921
const authModule = await import("../lib/auth/auth.js");
19221922

@@ -1958,13 +1958,14 @@ describe("OpenAIOAuthPlugin persistAccountPool", () => {
19581958
const organizationEntries = mockStorage.accounts.filter(
19591959
(account) => account.organizationId === "organization-shared",
19601960
);
1961-
expect(organizationEntries).toHaveLength(1);
1962-
expect(organizationEntries[0]?.accountId).toBe("org-variant-b");
1963-
expect(mockStorage.accounts).toHaveLength(2);
1961+
expect(organizationEntries).toHaveLength(2);
1962+
const organizationAccountIds = organizationEntries.map((account) => account.accountId).sort();
1963+
expect(organizationAccountIds).toEqual(["org-variant-a", "org-variant-b"]);
1964+
expect(mockStorage.accounts).toHaveLength(3);
19641965
expect(mockStorage.accounts.some((account) => account.accountId === "token-personal")).toBe(true);
19651966
});
19661967

1967-
it("preserves entries with different accountId values even when they share the same refresh token (org-scoped vs no-org)", async () => {
1968+
it("preserves org/no-org shared-refresh entries with different accountId values in pre-populated storage", async () => {
19681969
const accountsModule = await import("../lib/accounts.js");
19691970
const authModule = await import("../lib/auth/auth.js");
19701971

@@ -2216,7 +2217,7 @@ describe("OpenAIOAuthPlugin persistAccountPool", () => {
22162217
},
22172218
},
22182219
{
2219-
accountId: "token-shared",
2220+
accountId: "org-shared",
22202221
organizationId: "org-keep",
22212222
email: "token@example.com",
22222223
refreshToken: "shared-refresh",
@@ -2252,11 +2253,12 @@ describe("OpenAIOAuthPlugin persistAccountPool", () => {
22522253

22532254
await autoMethod.authorize({ loginMode: "add", accountCount: "1" });
22542255

2255-
const mergedOrg = mockStorage.accounts.find(
2256+
const mergedOrgEntries = mockStorage.accounts.filter(
22562257
(account) => account.organizationId === "org-keep",
22572258
);
2258-
expect(mergedOrg).toBeDefined();
2259-
expect(mergedOrg?.accountId).toBe("token-shared");
2259+
expect(mergedOrgEntries).toHaveLength(1);
2260+
const mergedOrg = mergedOrgEntries[0];
2261+
expect(mergedOrg?.accountId).toBe("org-shared");
22602262
expect(mergedOrg?.rateLimitResetTimes?.codex).toBe(9_000);
22612263
expect(mergedOrg?.rateLimitResetTimes?.["codex-max"]).toBe(5_000);
22622264
expect(mergedOrg?.rateLimitResetTimes?.["gpt-5.1"]).toBe(8_000);
@@ -2277,7 +2279,7 @@ describe("OpenAIOAuthPlugin persistAccountPool", () => {
22772279
lastUsed: 10,
22782280
},
22792281
{
2280-
accountId: "token-shared",
2282+
accountId: "org-shared",
22812283
organizationId: "org-keep",
22822284
email: "token@example.com",
22832285
refreshToken: "shared-refresh",
@@ -2312,16 +2314,71 @@ describe("OpenAIOAuthPlugin persistAccountPool", () => {
23122314

23132315
await autoMethod.authorize({ loginMode: "add", accountCount: "1" });
23142316

2315-
const mergedOrg = mockStorage.accounts.find(
2317+
const mergedOrgEntries = mockStorage.accounts.filter(
23162318
(account) => account.organizationId === "org-keep",
23172319
);
2318-
expect(mergedOrg).toBeDefined();
2319-
expect(mergedOrg?.accountId).toBe("token-shared");
2320+
expect(mergedOrgEntries).toHaveLength(1);
2321+
const mergedOrg = mergedOrgEntries[0];
2322+
expect(mergedOrg?.accountId).toBe("org-shared");
23202323
expect(mergedOrg?.enabled).toBe(false);
23212324
expect(mergedOrg?.coolingDownUntil).toBe(12_000);
23222325
expect(mergedOrg?.cooldownReason).toBe("auth-failure");
23232326
});
23242327

2328+
it("preserves same-organization entries when accountId differs", async () => {
2329+
const accountsModule = await import("../lib/accounts.js");
2330+
const authModule = await import("../lib/auth/auth.js");
2331+
2332+
mockStorage.accounts = [
2333+
{
2334+
accountId: "org-shared-a",
2335+
organizationId: "org-keep",
2336+
email: "org-a@example.com",
2337+
refreshToken: "shared-refresh",
2338+
addedAt: 10,
2339+
lastUsed: 10,
2340+
},
2341+
{
2342+
accountId: "org-shared-b",
2343+
organizationId: "org-keep",
2344+
email: "org-b@example.com",
2345+
refreshToken: "shared-refresh",
2346+
addedAt: 20,
2347+
lastUsed: 20,
2348+
},
2349+
];
2350+
2351+
vi.mocked(authModule.exchangeAuthorizationCode).mockResolvedValueOnce({
2352+
type: "success",
2353+
access: "access-unrelated-preserve",
2354+
refresh: "refresh-unrelated-preserve",
2355+
expires: Date.now() + 300_000,
2356+
idToken: "id-unrelated-preserve",
2357+
});
2358+
vi.mocked(accountsModule.getAccountIdCandidates).mockReturnValueOnce([]);
2359+
vi.mocked(accountsModule.extractAccountId).mockImplementation((accessToken) =>
2360+
accessToken === "access-unrelated-preserve" ? "unrelated-preserve" : "account-1",
2361+
);
2362+
2363+
const mockClient = createMockClient();
2364+
const { OpenAIOAuthPlugin } = await import("../index.js");
2365+
const plugin = (await OpenAIOAuthPlugin({
2366+
client: mockClient,
2367+
} as never)) as unknown as PluginType;
2368+
const autoMethod = plugin.auth.methods[0] as unknown as {
2369+
authorize: (inputs?: Record<string, string>) => Promise<{ instructions: string }>;
2370+
};
2371+
2372+
await autoMethod.authorize({ loginMode: "add", accountCount: "1" });
2373+
2374+
const orgEntries = mockStorage.accounts.filter(
2375+
(account) => account.organizationId === "org-keep",
2376+
);
2377+
expect(orgEntries).toHaveLength(2);
2378+
const accountIds = orgEntries.map((account) => account.accountId).sort();
2379+
expect(accountIds).toEqual(["org-shared-a", "org-shared-b"]);
2380+
});
2381+
23252382
it("persists non-team login and updates same record via accountId/refresh fallback", async () => {
23262383
const accountsModule = await import("../lib/accounts.js");
23272384
const authModule = await import("../lib/auth/auth.js");

0 commit comments

Comments
 (0)