fix(coding-agent): team-scope and serialize the stale-auth entitlement preserve

Review follow-ups on the stale-preserve mechanism:

- The keyless refresh restored entitlements even after the user selected a
  different prime-inference team, so wrong-team private models validated
  and then failed at the provider. The preserve now requires the current
  team selection to match the snapshot's team; a switch clears instead.
- Concurrent refreshAvailableModels calls could snapshot the entitlement
  set between another call's clear (inside refresh()) and its restore,
  then restore that empty snapshot - losing the cached authorization
  permanently despite the preserve rule. Refresh sequences that snapshot
  and restore the sets now run serialized through one promise chain
  (getExecutableModels' standalone entitlement refresh included).

Author: snimu

packages/coding-agent/src/core/model-registry.ts

@@ -770,11 +770,27 @@ export class ModelRegistry {
 	}
 
 	async refreshAvailableModels(): Promise<Model<Api>[]> {
-		const previousPrivateModelIds = new Set(this.authorizedPrivatePrimeInferenceModelIds);
-		const previousTeamId = this.authorizedPrivatePrimeInferenceTeamId;
-		this.refresh();
-		await this.refreshPrivatePrimeInferenceAuthorization(previousPrivateModelIds, previousTeamId);
-		return this.getAvailable();
+		// Serialized: a concurrent call snapshotting between this call's clear
+		// (inside refresh()) and its restore would capture an empty entitlement
+		// set and lose the cached private-model authorization for good.
+		return this.runSerializedEntitlementRefresh(async () => {
+			const previousPrivateModelIds = new Set(this.authorizedPrivatePrimeInferenceModelIds);
+			const previousTeamId = this.authorizedPrivatePrimeInferenceTeamId;
+			this.refresh();
+			await this.refreshPrivatePrimeInferenceAuthorization(previousPrivateModelIds, previousTeamId);
+			return this.getAvailable();
+		});
+	}
+
+	private entitlementRefreshChain: Promise<unknown> = Promise.resolve();
+
+	private runSerializedEntitlementRefresh<T>(task: () => Promise<T>): Promise<T> {
+		const run = this.entitlementRefreshChain.then(task, task);
+		this.entitlementRefreshChain = run.then(
+			() => undefined,
+			() => undefined,
+		);
+		return run;
 	}
 
 	private async refreshPrivatePrimeInferenceAuthorization(
@@ -786,8 +802,12 @@ export class ModelRegistry {
 		const teamId = teamHeaders?.["X-Prime-Team-ID"];
 		if (!apiKey || !teamHeaders || !teamId) {
 			// Stale is not logout: keep fetched entitlements for explicit re-selection
-			// (the auth filter still hides the models while stale).
-			if (this.authStorage.getAuthStatus(PRIME_INFERENCE_PROVIDER_ID).source === "stale") {
+			// (the auth filter still hides the models while stale) — but only for
+			// the team they were fetched for; a team switch invalidates them.
+			if (
+				this.authStorage.getAuthStatus(PRIME_INFERENCE_PROVIDER_ID).source === "stale" &&
+				teamId === previousTeamId
+			) {
 				this.authorizedPrivatePrimeInferenceModelIds = previousPrivateModelIds;
 				this.authorizedPrivatePrimeInferenceTeamId = previousTeamId;
 				return;
@@ -976,7 +996,7 @@ export class ModelRegistry {
 	}
 
 	async getExecutableModels(): Promise<Model<Api>[]> {
-		await this.refreshPrivatePrimeInferenceAuthorization();
+		await this.runSerializedEntitlementRefresh(() => this.refreshPrivatePrimeInferenceAuthorization());
 		const availableModels = this.getAvailable();
 		const codexModels = availableModels.filter((model) => model.provider === "openai-codex");
 		if (codexModels.length === 0) {

packages/coding-agent/test/model-registry.test.ts

@@ -4,7 +4,7 @@ import { join } from "node:path";
 import type { AnthropicMessagesCompat, Api, Context, Model, OpenAICompletionsCompat } from "@earendil-works/pi-ai";
 import { getApiProvider } from "@earendil-works/pi-ai";
 import { getOAuthProvider, registerOAuthProvider } from "@earendil-works/pi-ai/oauth";
-import { afterEach, beforeEach, describe, expect, test } from "vitest";
+import { afterEach, beforeEach, describe, expect, test, vi } from "vitest";
 import { AuthStorage } from "../src/core/auth-storage.js";
 import { ModelRegistry, type ProviderConfigInput } from "../src/core/model-registry.js";
 
@@ -1502,6 +1502,41 @@ describe("ModelRegistry", () => {
 				await expect(registry.getApiKeyForProvider("custom-provider")).resolves.toBe("literal_api_key_value");
 			});
 
+			test("preserves stale-auth entitlements only for the same prime-inference team", async () => {
+				authStorage.setRuntimeApiKey("prime-inference", "prime-key");
+				const headerSpy = vi
+					.spyOn(authStorage, "getProviderHeaders")
+					.mockReturnValue({ "X-Prime-Team-ID": "team-a" });
+				const registry = ModelRegistry.create(authStorage, modelsJsonPath);
+				const internals = registry as unknown as {
+					authorizedPrivatePrimeInferenceModelIds: Set<string>;
+					authorizedPrivatePrimeInferenceTeamId: string | undefined;
+				};
+				internals.authorizedPrivatePrimeInferenceModelIds.add("internal/private-model");
+				internals.authorizedPrivatePrimeInferenceTeamId = "team-a";
+				expect(registry.markProviderAuthStale("prime-inference")).toBe(true);
+
+				await registry.refreshAvailableModels();
+				expect(internals.authorizedPrivatePrimeInferenceModelIds.has("internal/private-model")).toBe(true);
+
+				// A team switch invalidates entitlements fetched for the old team.
+				headerSpy.mockReturnValue({ "X-Prime-Team-ID": "team-b" });
+				await registry.refreshAvailableModels();
+				expect(internals.authorizedPrivatePrimeInferenceModelIds.size).toBe(0);
+			});
+
+			test("concurrent stale-auth refreshes do not drop preserved entitlements", async () => {
+				authStorage.setRuntimeApiKey("prime-inference", "prime-key");
+				const registry = ModelRegistry.create(authStorage, modelsJsonPath);
+				const internals = registry as unknown as { authorizedPrivatePrimeInferenceModelIds: Set<string> };
+				internals.authorizedPrivatePrimeInferenceModelIds.add("internal/private-model");
+				expect(registry.markProviderAuthStale("prime-inference")).toBe(true);
+
+				await Promise.all([registry.refreshAvailableModels(), registry.refreshAvailableModels()]);
+
+				expect(internals.authorizedPrivatePrimeInferenceModelIds.has("internal/private-model")).toBe(true);
+			});
+
 			test("provider auth status reports command apiKey values from models.json without executing them", () => {
 				const counterFile = join(tempDir, "status-counter");
 				writeFileSync(counterFile, "0");