feat(zoo-gateway): auth callback, profile token sync, and sign-out

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: James Mtendamema

src/activate/__tests__/handleUri.spec.ts

@@ -6,25 +6,32 @@ vi.mock("vscode", () => ({
 
 import * as vscode from "vscode"
 
-const { mockGetVisibleInstance, mockHandleZooCodeAuthCallback, mockSetZooCodeUserInfo, mockVisibleProvider } =
-	vi.hoisted(() => {
-		const mockVisibleProvider = {
-			handleOpenRouterCallback: vi.fn(),
-			handleRequestyCallback: vi.fn(),
-			handleZooCodeCallback: vi.fn(),
-		} as any
-
-		return {
-			mockGetVisibleInstance: vi.fn(() => mockVisibleProvider),
-			mockHandleZooCodeAuthCallback: vi.fn(),
-			mockSetZooCodeUserInfo: vi.fn(),
-			mockVisibleProvider,
-		}
-	})
+const {
+	mockGetVisibleInstance,
+	mockGetAllInstances,
+	mockHandleZooCodeAuthCallback,
+	mockSetZooCodeUserInfo,
+	mockVisibleProvider,
+} = vi.hoisted(() => {
+	const mockVisibleProvider = {
+		handleOpenRouterCallback: vi.fn(),
+		handleRequestyCallback: vi.fn(),
+		handleZooCodeCallback: vi.fn(),
+	} as any
+
+	return {
+		mockGetVisibleInstance: vi.fn(() => mockVisibleProvider),
+		mockGetAllInstances: vi.fn(() => [mockVisibleProvider]),
+		mockHandleZooCodeAuthCallback: vi.fn(),
+		mockSetZooCodeUserInfo: vi.fn(),
+		mockVisibleProvider,
+	}
+})
 
 vi.mock("../../core/webview/ClineProvider", () => ({
 	ClineProvider: {
 		getVisibleInstance: mockGetVisibleInstance,
+		getAllInstances: mockGetAllInstances,
 	},
 }))
 
@@ -39,6 +46,7 @@ describe("handleUri", () => {
 	beforeEach(() => {
 		vi.clearAllMocks()
 		mockGetVisibleInstance.mockReturnValue(mockVisibleProvider)
+		mockGetAllInstances.mockReturnValue([mockVisibleProvider])
 	})
 
 	it("ignores legacy cloud auth callback", async () => {
@@ -54,8 +62,9 @@ describe("handleUri", () => {
 		)
 	})
 
-	it("stores callback user info even when no webview is visible", async () => {
+	it("stores callback user info even when no provider instances exist", async () => {
 		mockGetVisibleInstance.mockReturnValue(null)
+		mockGetAllInstances.mockReturnValue([])
 		mockHandleZooCodeAuthCallback.mockResolvedValue(true)
 
 		await handleUri({
@@ -69,6 +78,7 @@ describe("handleUri", () => {
 			email: "jane@example.com",
 			image: "https://example.com/avatar.png",
 		})
+		// No provider instances exist, so handleZooCodeCallback should not be called
 		expect(mockVisibleProvider.handleZooCodeCallback).not.toHaveBeenCalled()
 	})
 
@@ -116,4 +126,55 @@ describe("handleUri", () => {
 		expect(mockSetZooCodeUserInfo).not.toHaveBeenCalled()
 		expect(mockVisibleProvider.handleZooCodeCallback).not.toHaveBeenCalled()
 	})
+
+	it("propagates the callback token to every ClineProvider instance, not just the visible one", async () => {
+		// Regression: prior to multi-instance fan-out, hidden providers (sidebar collapsed,
+		// secondary panels) never received the zooSessionToken, so their profile settings
+		// stayed unauthenticated until reload.
+		mockHandleZooCodeAuthCallback.mockResolvedValue(true)
+
+		const hiddenProvider = { handleZooCodeCallback: vi.fn() } as any
+		const secondHidden = { handleZooCodeCallback: vi.fn() } as any
+		mockGetAllInstances.mockReturnValue([mockVisibleProvider, hiddenProvider, secondHidden])
+
+		await handleUri({
+			path: "/auth-callback",
+			query: "token=zoo_ext_test_token",
+		} as any)
+
+		expect(mockHandleZooCodeAuthCallback).toHaveBeenCalledWith("zoo_ext_test_token")
+		expect(mockSetZooCodeUserInfo).toHaveBeenCalled()
+		expect(mockVisibleProvider.handleZooCodeCallback).toHaveBeenCalledWith("zoo_ext_test_token")
+		expect(hiddenProvider.handleZooCodeCallback).toHaveBeenCalledWith("zoo_ext_test_token")
+		expect(secondHidden.handleZooCodeCallback).toHaveBeenCalledWith("zoo_ext_test_token")
+	})
+
+	it("serializes callbacks across instances to avoid concurrent profile-store writes", async () => {
+		// Regression: a previous implementation used Promise.all which fanned out concurrent
+		// read-modify-write operations on the same provider settings store. Verify the
+		// callbacks are invoked sequentially.
+		mockHandleZooCodeAuthCallback.mockResolvedValue(true)
+
+		const order: string[] = []
+		const makeProvider = (name: string) =>
+			({
+				handleZooCodeCallback: vi.fn(async () => {
+					order.push(`${name}:start`)
+					// Yield to the event loop so a concurrent call would interleave.
+					await new Promise((resolve) => setTimeout(resolve, 0))
+					order.push(`${name}:end`)
+				}),
+			}) as any
+
+		const a = makeProvider("a")
+		const b = makeProvider("b")
+		mockGetAllInstances.mockReturnValue([a, b])
+
+		await handleUri({
+			path: "/auth-callback",
+			query: "token=zoo_ext_test_token",
+		} as any)
+
+		expect(order).toEqual(["a:start", "a:end", "b:start", "b:end"])
+	})
 })

src/activate/handleUri.ts

@@ -50,9 +50,20 @@ export const handleUri = async (uri: vscode.Uri) => {
 						email,
 						image,
 					})
-					// Refresh webview state if a panel is currently open
-					if (visibleProvider) {
-						await visibleProvider.handleZooCodeCallback(token)
+					// Write the token to all active provider instances regardless of visibility.
+					// The profile settings write (handleZooCodeCallback) must run on any active
+					// instance — not just the visible one — so the zoo-gateway zooSessionToken
+					// is persisted even when the sidebar/panel is hidden at callback time.
+					//
+					// Run sequentially (NOT Promise.all): each ClineProvider's
+					// handleZooCodeCallback does a read-modify-write on the same backing
+					// provider settings store (listConfig → getProfile → saveConfig /
+					// upsertProviderProfile). Fanning out concurrently across N instances
+					// can interleave reads/writes and clobber updates. Serialization here
+					// is cheap (at most a handful of instances) and avoids the race.
+					const allInstances = ClineProvider.getAllInstances()
+					for (const instance of allInstances) {
+						await instance.handleZooCodeCallback(token)
 					}
 				}
 			}

src/core/webview/ClineProvider.ts

@@ -620,6 +620,10 @@ export class ClineProvider
 		return findLast(Array.from(this.activeInstances), (instance) => instance.view?.visible === true)
 	}
 
+	public static getAllInstances(): ClineProvider[] {
+		return Array.from(this.activeInstances)
+	}
+
 	public static async getInstance(): Promise<ClineProvider | undefined> {
 		let visibleProvider = ClineProvider.getVisibleInstance()
 
@@ -851,6 +855,64 @@ export class ClineProvider
 		if (!currentTask || currentTask.abandoned || currentTask.abort) {
 			await this.removeClineFromStack()
 		}
+
+		// Ensure zoo-gateway profile is seeded for users who signed in before this feature existed.
+		// Without this, users with a valid cached token but no zoo-gateway profile would need to
+		// re-authenticate to use Zoo Gateway. Fire-and-forget to avoid blocking webview init.
+		void this.ensureZooGatewayProfileSeeded().catch((err) => {
+			this.log(`[ensureZooGatewayProfileSeeded] Error: ${err instanceof Error ? err.message : String(err)}`)
+		})
+	}
+
+	/**
+	 * Seeds the zoo-gateway provider profile for users who have a cached auth token
+	 * but no profile (e.g., users who signed in before Zoo Gateway was added), or
+	 * who have an empty/imported profile without a token.
+	 * Called once per webview init; handleZooCodeCallback is idempotent so repeated calls are safe.
+	 */
+	private async ensureZooGatewayProfileSeeded(): Promise<void> {
+		const { getCachedZooCodeToken } = await import("../../services/zoo-code-auth")
+		const token = getCachedZooCodeToken()
+		if (!token) return
+
+		// Check ALL zoo-gateway profiles — only skip seeding if every profile has the current token.
+		// Using .find() would miss stale tokens in duplicate/renamed profiles since handleZooCodeCallback
+		// uses .filter() and updates all of them — the early-return guard must match.
+		const allProfiles = await this.providerSettingsManager.listConfig()
+		const zooGatewayProfiles = allProfiles.filter((p) => p.apiProvider === "zoo-gateway")
+
+		if (zooGatewayProfiles.length === 0) {
+			this.log("[ensureZooGatewayProfileSeeded] No zoo-gateway profile found, creating one")
+		} else {
+			let allUpToDate = true
+
+			for (const entry of zooGatewayProfiles) {
+				try {
+					const fullProfile = await this.providerSettingsManager.getProfile({ name: entry.name })
+					if (fullProfile.zooSessionToken !== token) {
+						allUpToDate = false
+						this.log(
+							fullProfile.zooSessionToken
+								? "[ensureZooGatewayProfileSeeded] Token mismatch (stale session?), updating with current token"
+								: "[ensureZooGatewayProfileSeeded] Existing zoo-gateway profile has no token, updating with cached token",
+						)
+						break
+					}
+				} catch {
+					allUpToDate = false
+					this.log("[ensureZooGatewayProfileSeeded] Failed to read existing profile, will re-seed")
+					break
+				}
+			}
+
+			if (allUpToDate) {
+				// All profiles have the current token — nothing to do
+				return
+			}
+		}
+
+		// User has token but either no profile, some profiles without token, or stale tokens — seed all
+		await this.handleZooCodeCallback(token)
 	}
 
 	public async createTaskWithHistoryItem(
@@ -1641,12 +1703,80 @@ export class ClineProvider
 		await this.upsertProviderProfile(currentApiConfigName, newConfiguration)
 	}
 
-	// Zoo Code Auth (for observability telemetry)
+	// Zoo Code Auth
 
-	async handleZooCodeCallback(_token: string) {
+	async handleZooCodeCallback(token: string) {
 		// Auth mutation (token storage, subscription check, success toast) was already
 		// performed by handleAuthCallback() in handleUri.ts before this method was called.
-		// This method only needs to refresh the webview state to reflect the new auth status.
+		// Save the zoo-gateway provider profile with the session token so that
+		// ZooGatewayHandler can authenticate without any manual user input.
+		//
+		// activate: true ONLY if Zoo Gateway is already the active profile — this pushes
+		// the new token to the in-memory handler so the current task picks it up immediately.
+		// Otherwise activate: false — do NOT switch providers mid-conversation. The user
+		// must explicitly select Zoo Gateway in settings if they want to use it.
+		try {
+			const { apiConfiguration } = await this.getState()
+			const currentSettings = this.contextProxy.getProviderSettings()
+			const currentApiConfigName = this.contextProxy.getValues().currentApiConfigName
+
+			// Derive the gateway base URL from ZOO_CODE_BASE_URL so that non-prod environments
+			// (staging, local dev) route completions to the correct backend instead of always
+			// hard-coding production. An already-set value in the profile is NOT preserved here —
+			// it must always align with the auth server the user just authenticated against.
+			const { getZooCodeBaseUrl } = await import("../../services/zoo-code-auth")
+			const derivedGatewayBaseUrl = `${getZooCodeBaseUrl()}/api/gateway/v1`
+
+			// Check if Zoo Gateway is the currently active profile by apiProvider identity,
+			// not by profile name (profile names are user-renameable).
+			const isZooGatewayActive = currentSettings.apiProvider === "zoo-gateway"
+
+			// Always scan ALL profiles and update every zoo-gateway profile with the new token.
+			// This ensures renamed profiles, duplicate profiles, and inactive profiles all stay
+			// in sync. The model lookup in requestRouterModels uses .find() which returns the
+			// first zoo-gateway profile it finds — if that profile has a stale token, requests fail.
+			const allProfiles = await this.providerSettingsManager.listConfig()
+			const zooProfiles = allProfiles.filter((p) => p.apiProvider === "zoo-gateway")
+
+			if (zooProfiles.length === 0) {
+				// No existing zoo-gateway profile — create the canonical default.
+				const newConfiguration: ProviderSettings = {
+					apiProvider: "zoo-gateway",
+					zooSessionToken: token,
+					zooGatewayModelId: apiConfiguration.zooGatewayModelId,
+					zooGatewayBaseUrl: derivedGatewayBaseUrl,
+				}
+				// Activate only if zoo-gateway was the active provider (shouldn't happen if
+				// no profiles exist, but defensive).
+				await this.upsertProviderProfile("Zoo Gateway", newConfiguration, isZooGatewayActive)
+			} else {
+				// Update every existing zoo-gateway profile with the new token and the
+				// derived base URL so that environment-specific routing stays consistent.
+				for (const entry of zooProfiles) {
+					const isActiveProfile = isZooGatewayActive && entry.name === currentApiConfigName
+					const existing = await this.providerSettingsManager.getProfile({ name: entry.name })
+					const updated: ProviderSettings = {
+						...existing,
+						zooSessionToken: token,
+						zooGatewayBaseUrl: derivedGatewayBaseUrl,
+					}
+					if (isActiveProfile) {
+						// Use upsertProviderProfile with activate: true so the in-memory handler
+						// picks up the new token immediately for the current task.
+						await this.upsertProviderProfile(entry.name, updated, true)
+					} else {
+						// Non-active profiles just need the token saved to disk.
+						await this.providerSettingsManager.saveConfig(entry.name, updated)
+					}
+				}
+			}
+		} catch (error) {
+			this.log(
+				`[handleZooCodeCallback] Failed to save zoo-gateway profile: ${
+					error instanceof Error ? error.message : String(error)
+				}`,
+			)
+		}
 		await this.postStateToWebview()
 	}
 

src/core/webview/__tests__/ClineProvider.spec.ts

@@ -2473,10 +2473,12 @@ describe("ClineProvider - Router Models", () => {
 				openrouter: mockModels,
 				requesty: mockModels,
 				unbound: mockModels,
+				"vercel-ai-gateway": mockModels,
+				"zoo-gateway": mockModels,
+				roo: {},
 				litellm: mockModels,
 				ollama: {},
 				lmstudio: {},
-				"vercel-ai-gateway": mockModels,
 				poe: {},
 				deepseek: {},
 			},
@@ -2508,6 +2510,7 @@ describe("ClineProvider - Router Models", () => {
 			.mockRejectedValueOnce(new Error("Requesty API error")) // requesty fail
 			.mockResolvedValueOnce(mockModels) // unbound success
 			.mockResolvedValueOnce(mockModels) // vercel-ai-gateway success
+			.mockResolvedValueOnce(mockModels) // zoo-gateway success
 			.mockRejectedValueOnce(new Error("LiteLLM connection failed")) // litellm fail
 
 		await messageHandler({ type: "requestRouterModels" })
@@ -2519,10 +2522,12 @@ describe("ClineProvider - Router Models", () => {
 				openrouter: mockModels,
 				requesty: {},
 				unbound: mockModels,
+				"vercel-ai-gateway": mockModels,
+				"zoo-gateway": mockModels,
+				roo: {},
 				ollama: {},
 				lmstudio: {},
 				litellm: {},
-				"vercel-ai-gateway": mockModels,
 				poe: {},
 				deepseek: {},
 			},
@@ -2614,10 +2619,12 @@ describe("ClineProvider - Router Models", () => {
 				openrouter: mockModels,
 				requesty: mockModels,
 				unbound: mockModels,
+				"vercel-ai-gateway": mockModels,
+				"zoo-gateway": mockModels,
+				roo: {},
 				litellm: {},
 				ollama: {},
 				lmstudio: {},
-				"vercel-ai-gateway": mockModels,
 				poe: {},
 				deepseek: {},
 			},