Skip to content

Commit a404a0a

Browse files
ndycodendycode
authored
Merge pull request #64 from sdip15fa/fix/multi-org-account-resolution
fix: preserve multi-org accounts and prevent cascading auth-removal
2 parents cda37b1 + be0039d commit a404a0a

File tree

8 files changed

+547
-343
lines changed

8 files changed

+547
-343
lines changed

‎index.ts‎

Lines changed: 94 additions & 163 deletions
Original file line numberDiff line numberDiff line change
@@ -601,15 +601,6 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
601601
return accountId && accountId.length > 0 ? accountId : undefined;
602602
};
603603

604-
const hasDistinctNonEmptyAccountIds = (
605-
left: { accountId?: string } | undefined,
606-
right: { accountId?: string } | undefined,
607-
): boolean => {
608-
const leftId = normalizeStoredAccountId(left);
609-
const rightId = normalizeStoredAccountId(right);
610-
return !!leftId && !!rightId && leftId !== rightId;
611-
};
612-
613604
const canCollapseWithCandidateAccountId = (
614605
existing: { accountId?: string } | undefined,
615606
candidateAccountId: string | undefined,
@@ -715,19 +706,41 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
715706
return newestExactAccountId ?? newestNoAccountId;
716707
};
717708

718-
const resolveUniqueOrgScopedMatch = (
719-
indexes: IdentityIndexes,
720-
accountId: string | undefined,
721-
refreshToken: string,
722-
): number | undefined => {
723-
const byAccountId = accountId
724-
? asUniqueIndex(indexes.byAccountIdOrgScoped.get(accountId))
725-
: undefined;
726-
if (byAccountId !== undefined) return byAccountId;
709+
const resolveUniqueOrgScopedMatch = (
710+
indexes: IdentityIndexes,
711+
accountId: string | undefined,
712+
refreshToken: string,
713+
): number | undefined => {
714+
const byAccountId = accountId
715+
? asUniqueIndex(indexes.byAccountIdOrgScoped.get(accountId))
716+
: undefined;
717+
if (byAccountId !== undefined) return byAccountId;
718+
719+
if (accountId) {
720+
const accountMatches = indexes.byAccountIdOrgScoped.get(accountId);
721+
if (accountMatches && accountMatches.length > 1) {
722+
let newestRefreshMatch: number | undefined;
723+
for (const index of accountMatches) {
724+
const existing = accounts[index];
725+
if (!existing) continue;
726+
const existingRefresh = existing.refreshToken?.trim();
727+
if (!existingRefresh || existingRefresh !== refreshToken) {
728+
continue;
729+
}
730+
newestRefreshMatch =
731+
typeof newestRefreshMatch === "number"
732+
? pickNewestAccountIndex(newestRefreshMatch, index)
733+
: index;
734+
}
735+
if (typeof newestRefreshMatch === "number") {
736+
return newestRefreshMatch;
737+
}
738+
}
739+
}
727740

728-
// Refresh-token-only fallback is allowed only when accountId is absent.
729-
// This avoids collapsing distinct workspace variants that share refresh token.
730-
if (accountId) return undefined;
741+
// Refresh-token-only fallback is allowed only when accountId is absent.
742+
// This avoids collapsing distinct workspace variants that share refresh token.
743+
if (accountId) return undefined;
731744

732745
return asUniqueIndex(indexes.byRefreshTokenOrgScoped.get(refreshToken));
733746
};
@@ -910,157 +923,48 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
910923

911924
const pruneRefreshTokenCollisions = (): void => {
912925
const indicesToRemove = new Set<number>();
913-
const refreshMap = new Map<
914-
string,
915-
{
916-
byOrg: Map<string, number[]>;
917-
preferredOrgIndex?: number;
918-
fallbackNoAccountIdIndex?: number;
919-
fallbackByAccountId: Map<string, number>;
920-
}
921-
>();
922-
923-
const pickPreferredOrgIndex = (
924-
existingIndex: number | undefined,
925-
candidateIndex: number,
926-
): number => {
927-
if (existingIndex === undefined) return candidateIndex;
928-
return pickNewestAccountIndex(existingIndex, candidateIndex);
929-
};
930-
931-
const collapseFallbackIntoPreferredOrg = (entry: {
932-
byOrg: Map<string, number[]>;
933-
preferredOrgIndex?: number;
934-
fallbackNoAccountIdIndex?: number;
935-
fallbackByAccountId: Map<string, number>;
936-
}): void => {
937-
if (entry.preferredOrgIndex === undefined) {
938-
return;
939-
}
940-
941-
const preferredOrgIndex = entry.preferredOrgIndex;
942-
const collapseFallbackIndex = (fallbackIndex: number): boolean => {
943-
if (preferredOrgIndex === fallbackIndex) return true;
944-
const target = accounts[preferredOrgIndex];
945-
const source = accounts[fallbackIndex];
946-
if (!target || !source) return true;
947-
const targetAccountId = normalizeStoredAccountId(target);
948-
const sourceAccountId = normalizeStoredAccountId(source);
949-
if (!targetAccountId && sourceAccountId) {
950-
return false;
951-
}
952-
if (hasDistinctNonEmptyAccountIds(target, source)) {
953-
return false;
954-
}
955-
mergeAccountRecords(preferredOrgIndex, fallbackIndex);
956-
indicesToRemove.add(fallbackIndex);
957-
return true;
958-
};
959-
960-
if (typeof entry.fallbackNoAccountIdIndex === "number") {
961-
if (collapseFallbackIndex(entry.fallbackNoAccountIdIndex)) {
962-
entry.fallbackNoAccountIdIndex = undefined;
963-
}
964-
}
965-
966-
const fallbackAccountIdsToDelete: string[] = [];
967-
for (const [accountId, fallbackIndex] of entry.fallbackByAccountId) {
968-
if (collapseFallbackIndex(fallbackIndex)) {
969-
fallbackAccountIdsToDelete.push(accountId);
970-
}
971-
}
972-
for (const accountId of fallbackAccountIdsToDelete) {
973-
entry.fallbackByAccountId.delete(accountId);
926+
const exactIdentityToIndex = new Map<string, number>();
927+
928+
const getExactIdentityKey = (
929+
account: {
930+
organizationId?: string;
931+
accountId?: string;
932+
email?: string;
933+
refreshToken?: string;
934+
} | undefined,
935+
): string => {
936+
const organizationId = account?.organizationId?.trim() ?? "";
937+
const accountId = normalizeStoredAccountId(account) ?? "";
938+
const email = account?.email?.trim().toLowerCase() ?? "";
939+
const refreshToken = account?.refreshToken?.trim() ?? "";
940+
if (organizationId || accountId) {
941+
return `org:${organizationId}|account:${accountId}|refresh:${refreshToken}`;
974942
}
943+
return `email:${email}|refresh:${refreshToken}`;
975944
};
976945

977946
for (let i = 0; i < accounts.length; i += 1) {
978947
const account = accounts[i];
979948
if (!account) continue;
980-
const refreshToken = account.refreshToken?.trim();
981-
if (!refreshToken) continue;
982-
const orgKey = account.organizationId?.trim() ?? "";
983-
let entry = refreshMap.get(refreshToken);
984-
if (!entry) {
985-
entry = {
986-
byOrg: new Map<string, number[]>(),
987-
preferredOrgIndex: undefined,
988-
fallbackNoAccountIdIndex: undefined,
989-
fallbackByAccountId: new Map<string, number>(),
990-
};
991-
refreshMap.set(refreshToken, entry);
992-
}
993-
994-
if (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-
);
1009-
if (existingIndex !== undefined) {
1010-
const newestIndex = pickNewestAccountIndex(existingIndex, i);
1011-
const obsoleteIndex = newestIndex === existingIndex ? i : existingIndex;
1012-
mergeAccountRecords(newestIndex, obsoleteIndex);
1013-
indicesToRemove.add(obsoleteIndex);
1014-
const nextOrgMatches = orgMatches.filter(
1015-
(index) => index !== obsoleteIndex && index !== newestIndex,
1016-
);
1017-
nextOrgMatches.push(newestIndex);
1018-
entry.byOrg.set(orgKey, nextOrgMatches);
1019-
entry.preferredOrgIndex = pickPreferredOrgIndex(entry.preferredOrgIndex, newestIndex);
1020-
collapseFallbackIntoPreferredOrg(entry);
1021-
continue;
1022-
}
1023-
entry.byOrg.set(orgKey, [...orgMatches, i]);
1024-
entry.preferredOrgIndex = pickPreferredOrgIndex(entry.preferredOrgIndex, i);
1025-
collapseFallbackIntoPreferredOrg(entry);
1026-
continue;
1027-
}
1028949

1029-
const fallbackAccountId = normalizeStoredAccountId(account);
1030-
if (fallbackAccountId) {
1031-
const existingFallback = entry.fallbackByAccountId.get(fallbackAccountId);
1032-
if (typeof existingFallback === "number") {
1033-
const newestIndex = pickNewestAccountIndex(existingFallback, i);
1034-
const obsoleteIndex = newestIndex === existingFallback ? i : existingFallback;
1035-
mergeAccountRecords(newestIndex, obsoleteIndex);
1036-
indicesToRemove.add(obsoleteIndex);
1037-
entry.fallbackByAccountId.set(fallbackAccountId, newestIndex);
1038-
collapseFallbackIntoPreferredOrg(entry);
1039-
continue;
1040-
}
1041-
entry.fallbackByAccountId.set(fallbackAccountId, i);
1042-
collapseFallbackIntoPreferredOrg(entry);
950+
const identityKey = getExactIdentityKey(account);
951+
const existingIndex = exactIdentityToIndex.get(identityKey);
952+
if (existingIndex === undefined) {
953+
exactIdentityToIndex.set(identityKey, i);
1043954
continue;
1044955
}
1045956

1046-
const existingFallback = entry.fallbackNoAccountIdIndex;
1047-
if (typeof existingFallback === "number") {
1048-
const newestIndex = pickNewestAccountIndex(existingFallback, i);
1049-
const obsoleteIndex = newestIndex === existingFallback ? i : existingFallback;
1050-
mergeAccountRecords(newestIndex, obsoleteIndex);
1051-
indicesToRemove.add(obsoleteIndex);
1052-
entry.fallbackNoAccountIdIndex = newestIndex;
1053-
collapseFallbackIntoPreferredOrg(entry);
1054-
continue;
1055-
}
1056-
entry.fallbackNoAccountIdIndex = i;
1057-
collapseFallbackIntoPreferredOrg(entry);
957+
const newestIndex = pickNewestAccountIndex(existingIndex, i);
958+
const obsoleteIndex = newestIndex === existingIndex ? i : existingIndex;
959+
mergeAccountRecords(newestIndex, obsoleteIndex);
960+
indicesToRemove.add(obsoleteIndex);
961+
exactIdentityToIndex.set(identityKey, newestIndex);
1058962
}
1059963

1060-
if (indicesToRemove.size > 0) {
1061-
accounts = accounts.filter((_, index) => !indicesToRemove.has(index));
1062-
}
1063-
};
964+
if (indicesToRemove.size > 0) {
965+
accounts = accounts.filter((_, index) => !indicesToRemove.has(index));
966+
}
967+
};
1064968

1065969
const collectIdentityKeys = (
1066970
account: { organizationId?: string; accountId?: string; refreshToken?: string } | undefined,
@@ -2179,7 +2083,7 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
21792083
}
21802084

21812085
while (true) {
2182-
const accountCount = accountManager.getAccountCount();
2086+
let accountCount = accountManager.getAccountCount();
21832087
const attempted = new Set<number>();
21842088
let restartAccountTraversalWithFallback = false;
21852089

@@ -2258,13 +2162,36 @@ while (attempted.size < Math.max(1, accountCount)) {
22582162
const accountLabel = formatAccountLabel(account, account.index);
22592163

22602164
if (failures >= ACCOUNT_LIMITS.MAX_AUTH_FAILURES_BEFORE_REMOVAL) {
2261-
accountManager.removeAccount(account);
2165+
const removedCount = accountManager.removeAccountsWithSameRefreshToken(account);
2166+
if (removedCount <= 0) {
2167+
logWarn(
2168+
`[${PLUGIN_NAME}] Expected grouped account removal after auth failures, but removed ${removedCount}.`,
2169+
);
2170+
const cooledCount = accountManager.markAccountsWithRefreshTokenCoolingDown(
2171+
account.refreshToken,
2172+
ACCOUNT_LIMITS.AUTH_FAILURE_COOLDOWN_MS,
2173+
"auth-failure",
2174+
);
2175+
if (cooledCount <= 0) {
2176+
logWarn(
2177+
`[${PLUGIN_NAME}] Unable to apply auth-failure cooldown; no live account found for refresh token.`,
2178+
);
2179+
}
2180+
accountManager.saveToDiskDebounced();
2181+
continue;
2182+
}
22622183
accountManager.saveToDiskDebounced();
2184+
const removalMessage = removedCount > 1
2185+
? `Removed ${removedCount} accounts (same refresh token) after ${failures} consecutive auth failures. Run 'opencode auth login' to re-add.`
2186+
: `Removed ${accountLabel} after ${failures} consecutive auth failures. Run 'opencode auth login' to re-add.`;
22632187
await showToast(
2264-
`Removed ${accountLabel} after ${failures} consecutive auth failures. Run 'opencode auth login' to re-add.`,
2188+
removalMessage,
22652189
"error",
22662190
{ duration: toastDurationMs * 2 },
22672191
);
2192+
// Restart traversal: clear attempted and refresh accountCount to avoid skipping healthy accounts
2193+
attempted.clear();
2194+
accountCount = accountManager.getAccountCount();
22682195
continue;
22692196
}
22702197

@@ -2322,6 +2249,7 @@ while (attempted.size < Math.max(1, accountCount)) {
23222249
{
23232250
model,
23242251
promptCacheKey,
2252+
organizationId: account.organizationId,
23252253
},
23262254
);
23272255

@@ -2959,6 +2887,7 @@ while (attempted.size < Math.max(1, accountCount)) {
29592887
const fetchCodexQuotaSnapshot = async (params: {
29602888
accountId: string;
29612889
accessToken: string;
2890+
organizationId: string | undefined;
29622891
}): Promise<CodexQuotaSnapshot> => {
29632892
const QUOTA_PROBE_MODELS = ["gpt-5-codex", "gpt-5.3-codex", "gpt-5.2-codex"];
29642893
let lastError: Error | null = null;
@@ -2985,6 +2914,7 @@ while (attempted.size < Math.max(1, accountCount)) {
29852914

29862915
const headers = createCodexHeaders(undefined, params.accountId, params.accessToken, {
29872916
model,
2917+
organizationId: params.organizationId,
29882918
});
29892919
headers.set("content-type", "application/json");
29902920

@@ -3258,6 +3188,7 @@ while (attempted.size < Math.max(1, accountCount)) {
32583188
const snapshot = await fetchCodexQuotaSnapshot({
32593189
accountId: requestAccountId,
32603190
accessToken,
3191+
organizationId: account.organizationId,
32613192
});
32623193
ok += 1;
32633194
console.log(

0 commit comments

Comments
 (0)