fix(identity): verify org-scoped preservation and no-org dedupe compatibility

Author: sdip15fa

index.ts

@@ -782,7 +782,43 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
 
 			const pruneRefreshTokenCollisions = (): void => {
 				const indicesToRemove = new Set<number>();
-				const refreshMap = new Map<string, { byOrg: Map<string, number>; fallbackIndex?: number }>();
+				const refreshMap = new Map<
+					string,
+					{ byOrg: Map<string, number>; preferredOrgIndex?: number; fallbackIndex?: number }
+				>();
+
+				const pickPreferredOrgIndex = (
+					existingIndex: number | undefined,
+					candidateIndex: number,
+				): number => {
+					if (existingIndex === undefined) return candidateIndex;
+					return pickNewestAccountIndex(existingIndex, candidateIndex);
+				};
+
+				const collapseFallbackIntoPreferredOrg = (entry: {
+					byOrg: Map<string, number>;
+					preferredOrgIndex?: number;
+					fallbackIndex?: number;
+				}): void => {
+					if (entry.preferredOrgIndex === undefined || entry.fallbackIndex === undefined) {
+						return;
+					}
+
+					const preferredOrgIndex = entry.preferredOrgIndex;
+					const fallbackIndex = entry.fallbackIndex;
+					if (preferredOrgIndex === fallbackIndex) {
+						entry.fallbackIndex = undefined;
+						return;
+					}
+
+					const target = accounts[preferredOrgIndex];
+					const source = accounts[fallbackIndex];
+					if (target && source) {
+						mergeAccountRecords(preferredOrgIndex, fallbackIndex);
+						indicesToRemove.add(fallbackIndex);
+						entry.fallbackIndex = undefined;
+					}
+				};
 
 				for (let i = 0; i < accounts.length; i += 1) {
 					const account = accounts[i];
@@ -792,7 +828,11 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
 					const orgKey = account.organizationId?.trim() ?? "";
 					let entry = refreshMap.get(refreshToken);
 					if (!entry) {
-						entry = { byOrg: new Map<string, number>(), fallbackIndex: undefined };
+						entry = {
+							byOrg: new Map<string, number>(),
+							preferredOrgIndex: undefined,
+							fallbackIndex: undefined,
+						};
 						refreshMap.set(refreshToken, entry);
 					}
 
@@ -804,9 +844,13 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
 							mergeAccountRecords(newestIndex, obsoleteIndex);
 							indicesToRemove.add(obsoleteIndex);
 							entry.byOrg.set(orgKey, newestIndex);
+							entry.preferredOrgIndex = pickPreferredOrgIndex(entry.preferredOrgIndex, newestIndex);
+							collapseFallbackIntoPreferredOrg(entry);
 							continue;
 						}
 						entry.byOrg.set(orgKey, i);
+						entry.preferredOrgIndex = pickPreferredOrgIndex(entry.preferredOrgIndex, i);
+						collapseFallbackIntoPreferredOrg(entry);
 						continue;
 					}
 
@@ -817,9 +861,11 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
 						mergeAccountRecords(newestIndex, obsoleteIndex);
 						indicesToRemove.add(obsoleteIndex);
 						entry.fallbackIndex = newestIndex;
+						collapseFallbackIntoPreferredOrg(entry);
 						continue;
 					}
 					entry.fallbackIndex = i;
+					collapseFallbackIntoPreferredOrg(entry);
 				}
 
 			if (indicesToRemove.size > 0) {
@@ -893,7 +939,13 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
 
 			const clampedActiveIndex = Math.max(0, Math.min(Math.floor(activeIndex), accounts.length - 1));
 			const activeIndexByFamily: Partial<Record<ModelFamily, number>> = {};
-			for (const family of MODEL_FAMILIES) {
+			const familiesToPersist = replaceAll
+				? []
+				: MODEL_FAMILIES.filter((family) => {
+					const storedFamilyIndex = stored?.activeIndexByFamily?.[family];
+					return typeof storedFamilyIndex === "number" && Number.isFinite(storedFamilyIndex);
+				});
+			for (const family of familiesToPersist) {
 				const storedFamilyIndex = stored?.activeIndexByFamily?.[family];
 				const remappedFamilyIndex = replaceAll
 					? undefined

lib/storage.ts

@@ -412,7 +412,46 @@ function mergeAccountRecords<T extends AccountLike>(target: T, source: T): T {
 function deduplicateAccountsByRefreshToken<T extends AccountLike>(accounts: T[]): T[] {
   const working = [...accounts];
   const indicesToRemove = new Set<number>();
-  const refreshMap = new Map<string, { byOrg: Map<string, number>; fallbackIndex?: number }>();
+  const refreshMap = new Map<string, {
+    byOrg: Map<string, number>;
+    preferredOrgIndex?: number;
+    fallbackIndex?: number;
+  }>();
+
+  const pickPreferredOrgIndex = (
+    existingIndex: number | undefined,
+    candidateIndex: number,
+  ): number => {
+    if (existingIndex === undefined) return candidateIndex;
+    return pickNewestAccountIndex(working, existingIndex, candidateIndex);
+  };
+
+  const collapseFallbackIntoPreferredOrg = (
+    entry: {
+      byOrg: Map<string, number>;
+      preferredOrgIndex?: number;
+      fallbackIndex?: number;
+    },
+  ): void => {
+    if (entry.preferredOrgIndex === undefined || entry.fallbackIndex === undefined) {
+      return;
+    }
+
+    const preferredOrgIndex = entry.preferredOrgIndex;
+    const fallbackIndex = entry.fallbackIndex;
+    if (preferredOrgIndex === fallbackIndex) {
+      entry.fallbackIndex = undefined;
+      return;
+    }
+
+    const target = working[preferredOrgIndex];
+    const source = working[fallbackIndex];
+    if (target && source) {
+      working[preferredOrgIndex] = mergeAccountRecords(target, source);
+      indicesToRemove.add(fallbackIndex);
+      entry.fallbackIndex = undefined;
+    }
+  };
 
   for (let i = 0; i < working.length; i += 1) {
     const account = working[i];
@@ -424,7 +463,11 @@ function deduplicateAccountsByRefreshToken<T extends AccountLike>(accounts: T[])
 
     let entry = refreshMap.get(refreshToken);
     if (!entry) {
-      entry = { byOrg: new Map<string, number>(), fallbackIndex: undefined };
+      entry = {
+        byOrg: new Map<string, number>(),
+        preferredOrgIndex: undefined,
+        fallbackIndex: undefined,
+      };
       refreshMap.set(refreshToken, entry);
     }
 
@@ -440,9 +483,13 @@ function deduplicateAccountsByRefreshToken<T extends AccountLike>(accounts: T[])
         }
         indicesToRemove.add(obsoleteIndex);
         entry.byOrg.set(orgKey, newestIndex);
+        entry.preferredOrgIndex = pickPreferredOrgIndex(entry.preferredOrgIndex, newestIndex);
+        collapseFallbackIntoPreferredOrg(entry);
         continue;
       }
       entry.byOrg.set(orgKey, i);
+      entry.preferredOrgIndex = pickPreferredOrgIndex(entry.preferredOrgIndex, i);
+      collapseFallbackIntoPreferredOrg(entry);
       continue;
     }
 
@@ -457,9 +504,11 @@ function deduplicateAccountsByRefreshToken<T extends AccountLike>(accounts: T[])
       }
       indicesToRemove.add(obsoleteIndex);
       entry.fallbackIndex = newestIndex;
+      collapseFallbackIntoPreferredOrg(entry);
       continue;
     }
     entry.fallbackIndex = i;
+    collapseFallbackIntoPreferredOrg(entry);
   }
 
   const result: T[] = [];