fix(zoo-gateway): address PR review feedback on auth, seeding, and errors

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

Author: James Mtendamema

src/activate/__tests__/handleUri.spec.ts

@@ -177,4 +177,24 @@ describe("handleUri", () => {
 
 		expect(order).toEqual(["a:start", "a:end", "b:start", "b:end"])
 	})
+
+	it("continues fan-out when one instance fails to persist the callback token", async () => {
+		mockHandleZooCodeAuthCallback.mockResolvedValue(true)
+
+		const failingProvider = {
+			handleZooCodeCallback: vi.fn(async () => {
+				throw new Error("profile store unavailable")
+			}),
+		} as any
+		const healthyProvider = { handleZooCodeCallback: vi.fn() } as any
+		mockGetAllInstances.mockReturnValue([failingProvider, healthyProvider])
+
+		await handleUri({
+			path: "/auth-callback",
+			query: "token=zoo_ext_test_token",
+		} as any)
+
+		expect(failingProvider.handleZooCodeCallback).toHaveBeenCalledWith("zoo_ext_test_token")
+		expect(healthyProvider.handleZooCodeCallback).toHaveBeenCalledWith("zoo_ext_test_token")
+	})
 })

src/activate/handleUri.ts

@@ -63,7 +63,14 @@ export const handleUri = async (uri: vscode.Uri) => {
 					// 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)
+						try {
+							await instance.handleZooCodeCallback(token)
+						} catch (error) {
+							console.error(
+								"Failed to persist Zoo Gateway token for a provider instance:",
+								error instanceof Error ? error.message : error,
+							)
+						}
 					}
 				}
 			}

src/api/providers/__tests__/zoo-gateway.spec.ts

@@ -56,9 +56,11 @@ vitest.mock("../fetchers/modelCache", () => ({
 	getModelsFromCache: vitest.fn().mockReturnValue(undefined),
 }))
 
+const mockGetCachedZooCodeToken = vitest.hoisted(() => vitest.fn<() => string | undefined>(() => undefined))
+
 vitest.mock("../../../services/zoo-code-auth", () => ({
 	getZooCodeBaseUrl: vitest.fn(() => "https://www.zoocode.dev"),
-	getCachedZooCodeToken: vitest.fn(() => undefined),
+	getCachedZooCodeToken: mockGetCachedZooCodeToken,
 	clearZooCodeToken: vitest.fn(async () => undefined),
 }))
 
@@ -91,6 +93,7 @@ describe("ZooGatewayHandler", () => {
 
 	beforeEach(() => {
 		vitest.clearAllMocks()
+		mockGetCachedZooCodeToken.mockReturnValue(undefined)
 		mockCreate.mockClear()
 		showErrorMessage.mockReset()
 		showErrorMessage.mockResolvedValue(undefined)
@@ -126,6 +129,21 @@ describe("ZooGatewayHandler", () => {
 			expect(OpenAI).not.toHaveBeenCalled()
 		})
 
+		it("prefers the secret-storage cache over a persisted profile token", () => {
+			mockGetCachedZooCodeToken.mockReturnValue("zoo_ext_cached_token")
+
+			new ZooGatewayHandler({
+				zooSessionToken: "zoo_ext_stale_profile_token",
+				zooGatewayModelId: mockOptions.zooGatewayModelId,
+			})
+
+			expect(OpenAI).toHaveBeenCalledWith(
+				expect.objectContaining({
+					apiKey: "zoo_ext_cached_token",
+				}),
+			)
+		})
+
 		it("initializes OpenAI with Zoo enrichment headers and session token", () => {
 			const handler = new ZooGatewayHandler({
 				...mockOptions,

src/api/providers/zoo-gateway.ts

@@ -99,10 +99,9 @@ export class ZooGatewayHandler extends RouterProvider implements SingleCompletio
 	constructor(options: ApiHandlerOptions) {
 		const baseURL = options.zooGatewayBaseUrl ?? `${getZooCodeBaseUrl()}/api/gateway/v1`
 
-		// Prefer the profile-persisted token; fall back to the secret-storage cache so
-		// requests work when the user is signed in but the profile hasn't been seeded yet
-		// (e.g. auth callback arrived before any webview instance was open).
-		const sessionToken = options.zooSessionToken || getCachedZooCodeToken()
+		// Prefer the secret-storage cache so a 401 clear takes effect immediately; fall back
+		// to the profile-persisted token when the user is signed in but seeding hasn't run yet.
+		const sessionToken = getCachedZooCodeToken() || options.zooSessionToken
 		if (!sessionToken) {
 			throw new Error("Zoo Gateway requires authentication. Please sign in to Zoo Code first.")
 		}
@@ -210,7 +209,12 @@ export class ZooGatewayHandler extends RouterProvider implements SingleCompletio
 				}
 			}
 		} catch (error) {
-			void surfaceGatewayApiError(error)
+			void surfaceGatewayApiError(error).catch((surfaceError) => {
+				console.error(
+					"Failed to surface Zoo Gateway error:",
+					surfaceError instanceof Error ? surfaceError.message : surfaceError,
+				)
+			})
 			throw error
 		}
 	}
@@ -234,7 +238,12 @@ export class ZooGatewayHandler extends RouterProvider implements SingleCompletio
 			const response = await this.client.chat.completions.create(requestOptions)
 			return response.choices[0]?.message.content || ""
 		} catch (error) {
-			void surfaceGatewayApiError(error)
+			void surfaceGatewayApiError(error).catch((surfaceError) => {
+				console.error(
+					"Failed to surface Zoo Gateway error:",
+					surfaceError instanceof Error ? surfaceError.message : surfaceError,
+				)
+			})
 			if (error instanceof Error) {
 				throw new Error(`Zoo Gateway completion error: ${error.message}`)
 			}

src/core/webview/ClineProvider.ts

@@ -871,9 +871,10 @@ export class ClineProvider
 	 * 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 { getCachedZooCodeToken, getZooCodeBaseUrl } = await import("../../services/zoo-code-auth")
 		const token = getCachedZooCodeToken()
 		if (!token) return
+		const expectedGatewayBaseUrl = `${getZooCodeBaseUrl()}/api/gateway/v1`
 
 		// 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
@@ -889,13 +890,12 @@ export class ClineProvider
 			for (const entry of zooGatewayProfiles) {
 				try {
 					const fullProfile = await this.providerSettingsManager.getProfile({ name: entry.name })
-					if (fullProfile.zooSessionToken !== token) {
+					if (
+						fullProfile.zooSessionToken !== token ||
+						fullProfile.zooGatewayBaseUrl !== expectedGatewayBaseUrl
+					) {
 						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",
-						)
+						this.log("[ensureZooGatewayProfileSeeded] Existing zoo-gateway profile is stale, updating")
 						break
 					}
 				} catch {

src/core/webview/webviewMessageHandler.ts

@@ -955,13 +955,16 @@ export const webviewMessageHandler = async (
 			if (!zooGatewayToken) {
 				try {
 					const allProfiles = await provider.providerSettingsManager.listConfig()
-					const zooGatewayProfile = allProfiles.find((p) => p.apiProvider === "zoo-gateway")
-					if (zooGatewayProfile) {
+					const zooGatewayProfiles = allProfiles.filter((p) => p.apiProvider === "zoo-gateway")
+					for (const profileMeta of zooGatewayProfiles) {
 						const fullProfile = await provider.providerSettingsManager.getProfile({
-							name: zooGatewayProfile.name,
+							name: profileMeta.name,
 						})
-						zooGatewayToken = fullProfile.zooSessionToken
-						zooGatewayBaseUrl = fullProfile.zooGatewayBaseUrl ?? zooGatewayBaseUrl
+						if (fullProfile.zooSessionToken) {
+							zooGatewayToken = fullProfile.zooSessionToken
+							zooGatewayBaseUrl = fullProfile.zooGatewayBaseUrl ?? zooGatewayBaseUrl
+							break
+						}
 					}
 				} catch (error) {
 					console.debug("Failed to look up zoo-gateway profile for model fetch:", error)

webview-ui/src/components/welcome/WelcomeViewProvider.tsx

@@ -20,14 +20,17 @@ const DEFAULT_WELCOME_API_CONFIGURATION: ProviderSettings = {
 	openRouterModelId: openRouterDefaultModelId,
 }
 
-const getWelcomeApiConfiguration = (apiConfiguration?: ProviderSettings): ProviderSettings => {
+const getWelcomeApiConfiguration = (
+	apiConfiguration?: ProviderSettings,
+	zooCodeIsAuthenticated?: boolean,
+): ProviderSettings => {
 	// validateApiConfiguration treats a missing apiProvider as valid (no switch case matches),
 	// so we explicitly fall back here before delegating to it for incomplete-but-set configs.
 	if (!apiConfiguration?.apiProvider) {
 		return DEFAULT_WELCOME_API_CONFIGURATION
 	}
 
-	const validationError = validateApiConfiguration(apiConfiguration)
+	const validationError = validateApiConfiguration(apiConfiguration, undefined, undefined, zooCodeIsAuthenticated)
 	if (validationError) {
 		return DEFAULT_WELCOME_API_CONFIGURATION
 	}
@@ -42,7 +45,8 @@ const WelcomeViewProvider = () => {
 	const [errorMessage, setErrorMessage] = useState<string | undefined>(undefined)
 	const [showProviderSetup, setShowProviderSetup] = useState(false)
 	const [welcomeApiConfiguration, setWelcomeApiConfiguration] = useState<ProviderSettings>()
-	const effectiveApiConfiguration = welcomeApiConfiguration ?? getWelcomeApiConfiguration(apiConfiguration)
+	const effectiveApiConfiguration =
+		welcomeApiConfiguration ?? getWelcomeApiConfiguration(apiConfiguration, zooCodeIsAuthenticated)
 
 	const setApiConfigurationFieldForApiOptions = useCallback(
 		<K extends keyof ProviderSettings>(field: K, value: ProviderSettings[K]) => {
@@ -57,7 +61,7 @@ const WelcomeViewProvider = () => {
 
 	const handleGetStarted = useCallback(() => {
 		if (!showProviderSetup) {
-			const initialApiConfiguration = getWelcomeApiConfiguration(apiConfiguration)
+			const initialApiConfiguration = getWelcomeApiConfiguration(apiConfiguration, zooCodeIsAuthenticated)
 			setWelcomeApiConfiguration(initialApiConfiguration)
 
 			setApiConfiguration(initialApiConfiguration)