fix(zoo-gateway): sign-out clears stale profile tokens, simplify model fetch

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

Author: 2 people

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

@@ -57,11 +57,21 @@ vitest.mock("../fetchers/modelCache", () => ({
 }))
 
 const mockGetCachedZooCodeToken = vitest.hoisted(() => vitest.fn<() => string | undefined>(() => undefined))
+const mockSessionCleared = vitest.hoisted(() => ({ value: false }))
 
 vitest.mock("../../../services/zoo-code-auth", () => ({
 	getZooCodeBaseUrl: vitest.fn(() => "https://www.zoocode.dev"),
-	getCachedZooCodeToken: mockGetCachedZooCodeToken,
-	clearZooCodeToken: vitest.fn(async () => undefined),
+	getCachedZooCodeToken: () => mockGetCachedZooCodeToken() ?? "",
+	resolveZooGatewaySessionToken: (profileToken?: string) => {
+		const cached = mockGetCachedZooCodeToken()
+		if (cached) return cached
+		if (mockSessionCleared.value) return undefined
+		return profileToken
+	},
+	clearZooCodeToken: vitest.fn(async () => {
+		mockSessionCleared.value = true
+		mockGetCachedZooCodeToken.mockReturnValue(undefined)
+	}),
 }))
 
 vitest.mock("../../transform/caching/vercel-ai-gateway", () => ({
@@ -93,6 +103,7 @@ describe("ZooGatewayHandler", () => {
 
 	beforeEach(() => {
 		vitest.clearAllMocks()
+		mockSessionCleared.value = false
 		mockGetCachedZooCodeToken.mockReturnValue(undefined)
 		mockCreate.mockClear()
 		showErrorMessage.mockReset()
@@ -195,6 +206,13 @@ describe("ZooGatewayHandler", () => {
 	})
 
 	describe("createMessage", () => {
+		it("requires authentication at request time when no session token is available", async () => {
+			const handler = new ZooGatewayHandler({})
+			await expect(drainCreateMessage(handler)).rejects.toThrow(
+				"Zoo Gateway requires authentication. Please sign in to Zoo Code first.",
+			)
+		})
+
 		beforeEach(() => {
 			mockCreate.mockImplementation(async () => ({
 				[Symbol.asyncIterator]: async function* () {

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

@@ -8,6 +8,7 @@ vitest.mock("axios")
 vitest.mock("../../../../services/zoo-code-auth", () => ({
 	getCachedZooCodeToken: vitest.fn(() => ""),
 	getZooCodeBaseUrl: vitest.fn(() => "https://example.test"),
+	resolveZooGatewaySessionToken: vitest.fn((profileToken?: string) => profileToken || undefined),
 }))
 const mockedAxios = axios as any
 

src/api/providers/fetchers/zoo-gateway.ts

@@ -3,7 +3,7 @@ import axios from "axios"
 import type { ModelInfo } from "@roo-code/types"
 
 import type { ApiHandlerOptions } from "../../../shared/api"
-import { getCachedZooCodeToken, getZooCodeBaseUrl } from "../../../services/zoo-code-auth"
+import { getZooCodeBaseUrl, resolveZooGatewaySessionToken } from "../../../services/zoo-code-auth"
 
 import {
 	type VercelAiGatewayModel,
@@ -24,9 +24,7 @@ export async function getZooGatewayModels(options?: ApiHandlerOptions): Promise<
 	const models: Record<string, ModelInfo> = {}
 	const baseURL = options?.zooGatewayBaseUrl ?? `${getZooCodeBaseUrl()}/api/gateway/v1`
 
-	// Build headers - Zoo Gateway requires authentication via the zoo_ext_ session token.
-	// Fall back to the secret-storage cache when the profile hasn't been seeded yet.
-	const sessionToken = options?.zooSessionToken || getCachedZooCodeToken()
+	const sessionToken = resolveZooGatewaySessionToken(options?.zooSessionToken)
 	const headers: Record<string, string> = {}
 	if (sessionToken) {
 		headers["Authorization"] = `Bearer ${sessionToken}`

src/api/providers/zoo-gateway.ts

@@ -10,7 +10,7 @@ import {
 } from "@roo-code/types"
 
 import { ApiHandlerOptions } from "../../shared/api"
-import { clearZooCodeToken, getCachedZooCodeToken, getZooCodeBaseUrl } from "../../services/zoo-code-auth"
+import { clearZooCodeToken, getZooCodeBaseUrl, resolveZooGatewaySessionToken } from "../../services/zoo-code-auth"
 import { Package } from "../../shared/package"
 import { t } from "../../i18n"
 
@@ -101,9 +101,7 @@ export class ZooGatewayHandler extends RouterProvider implements SingleCompletio
 	constructor(options: ApiHandlerOptions) {
 		const baseURL = options.zooGatewayBaseUrl ?? `${getZooCodeBaseUrl()}/api/gateway/v1`
 
-		// 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
+		const sessionToken = resolveZooGatewaySessionToken(options.zooSessionToken)
 
 		// Merge Zoo-specific enrichment headers into openAiHeaders so they flow through
 		// the parent's single OpenAI client. We avoid reassigning `this.client` (which
@@ -128,8 +126,7 @@ export class ZooGatewayHandler extends RouterProvider implements SingleCompletio
 	}
 
 	private ensureAuthenticated(): void {
-		const sessionToken = getCachedZooCodeToken() || this.options.zooSessionToken
-		if (!sessionToken) {
+		if (!resolveZooGatewaySessionToken(this.options.zooSessionToken)) {
 			throw new Error(ZOO_GATEWAY_AUTH_ERROR)
 		}
 	}
@@ -217,12 +214,14 @@ export class ZooGatewayHandler extends RouterProvider implements SingleCompletio
 				}
 			}
 		} catch (error) {
-			void surfaceGatewayApiError(error).catch((surfaceError) => {
+			try {
+				await surfaceGatewayApiError(error)
+			} catch (surfaceError) {
 				console.error(
 					"Failed to surface Zoo Gateway error:",
 					surfaceError instanceof Error ? surfaceError.message : surfaceError,
 				)
-			})
+			}
 			throw error
 		}
 	}
@@ -248,12 +247,14 @@ 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).catch((surfaceError) => {
+			try {
+				await 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/webviewMessageHandler.ts

@@ -946,31 +946,6 @@ export const webviewMessageHandler = async (
 				}
 			}
 
-			// For zoo-gateway, the token may be stored in a separate zoo-gateway profile
-			// (not the currently active profile). Look it up so the model list populates
-			// even when zoo-gateway isn't the active provider.
-			let zooGatewayToken = apiConfiguration.zooSessionToken
-			let zooGatewayBaseUrl = apiConfiguration.zooGatewayBaseUrl
-
-			if (!zooGatewayToken) {
-				try {
-					const allProfiles = await provider.providerSettingsManager.listConfig()
-					const zooGatewayProfiles = allProfiles.filter((p) => p.apiProvider === "zoo-gateway")
-					for (const profileMeta of zooGatewayProfiles) {
-						const fullProfile = await provider.providerSettingsManager.getProfile({
-							name: profileMeta.name,
-						})
-						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)
-				}
-			}
-
 			// Base candidates (only those handled by this aggregate fetcher)
 			const candidates: { key: RouterName; options: GetModelsOptions }[] = [
 				{ key: "openrouter", options: { provider: "openrouter" } },
@@ -994,8 +969,7 @@ export const webviewMessageHandler = async (
 					key: "zoo-gateway",
 					options: {
 						provider: "zoo-gateway",
-						apiKey: zooGatewayToken,
-						baseUrl: zooGatewayBaseUrl,
+						baseUrl: apiConfiguration.zooGatewayBaseUrl,
 					},
 				},
 			]

src/services/__tests__/zoo-code-auth.test.ts

@@ -12,6 +12,7 @@ import {
 	getZooCodeBaseUrl,
 	handleAuthCallback,
 	initZooCodeAuth,
+	resolveZooGatewaySessionToken,
 	setZooCodeToken,
 	setZooCodeUserInfo,
 	verifyZooCodeToken,
@@ -432,6 +433,29 @@ describe("zoo-code-auth", () => {
 		})
 	})
 
+	describe("resolveZooGatewaySessionToken", () => {
+		it("prefers the cached token over a profile token", async () => {
+			await initZooCodeAuth(mockContext)
+			await setZooCodeToken("zoo_ext_cached")
+
+			expect(resolveZooGatewaySessionToken("zoo_ext_profile")).toBe("zoo_ext_cached")
+		})
+
+		it("ignores profile tokens after an explicit sign-out clear", async () => {
+			await initZooCodeAuth(mockContext)
+			await setZooCodeToken("zoo_ext_cached")
+			await clearZooCodeToken()
+
+			expect(resolveZooGatewaySessionToken("zoo_ext_stale_profile")).toBeUndefined()
+		})
+
+		it("falls back to the profile token when the cache is empty and not cleared", async () => {
+			await initZooCodeAuth(mockContext)
+
+			expect(resolveZooGatewaySessionToken("zoo_ext_profile")).toBe("zoo_ext_profile")
+		})
+	})
+
 	describe("disconnectZooCode", () => {
 		it("revokes the current token and clears cached auth state", async () => {
 			await initZooCodeAuth(mockContext)

src/services/zoo-code-auth.ts

@@ -11,6 +11,7 @@ let secretStorage: vscode.SecretStorage | undefined
 
 // In-memory cache for synchronous access in ZooCodeHandler hot path
 let _cachedToken: string | undefined = undefined
+let _sessionCleared = false
 let _cachedUserName: string | undefined = undefined
 let _cachedUserEmail: string | undefined = undefined
 let _cachedUserImage: string | undefined = undefined
@@ -28,6 +29,7 @@ export async function initZooCodeAuth(context: vscode.ExtensionContext): Promise
 
 	// Pre-load the token and user info into memory on init so ZooCodeHandler can access them synchronously
 	_cachedToken = await secretStorage.get(ZOO_CODE_TOKEN_KEY)
+	_sessionCleared = false
 	_cachedUserName = await secretStorage.get(ZOO_CODE_USER_NAME_KEY)
 	_cachedUserEmail = await secretStorage.get(ZOO_CODE_USER_EMAIL_KEY)
 	_cachedUserImage = await secretStorage.get(ZOO_CODE_USER_IMAGE_KEY)
@@ -85,6 +87,21 @@ export function getCachedZooCodeToken(): string {
 	return _cachedToken ?? ""
 }
 
+/**
+ * Resolves the Zoo Gateway session token for API calls.
+ * Secret-storage cache wins over profile-persisted tokens; after an explicit sign-out
+ * or 401 clear, profile tokens are ignored so stale credentials cannot be reused.
+ */
+export function resolveZooGatewaySessionToken(profileToken?: string): string | undefined {
+	if (_cachedToken) {
+		return _cachedToken
+	}
+	if (_sessionCleared) {
+		return undefined
+	}
+	return profileToken || undefined
+}
+
 export function getCachedZooCodeUserInfo(): { name?: string; email?: string; image?: string } {
 	return {
 		name: _cachedUserName,
@@ -153,6 +170,7 @@ export async function setZooCodeToken(token: string): Promise<void> {
 	if (!secretStorage) return
 	await secretStorage.store(ZOO_CODE_TOKEN_KEY, token)
 	_cachedToken = token
+	_sessionCleared = false
 	// Reset subscription status when token is set
 	_cachedSubscriptionStatus = "unknown"
 	_lastSubscriptionCheck = 0
@@ -204,6 +222,7 @@ export async function clearZooCodeToken(): Promise<void> {
 	if (!secretStorage) return
 	await secretStorage.delete(ZOO_CODE_TOKEN_KEY)
 	_cachedToken = undefined
+	_sessionCleared = true
 	_cachedSubscriptionStatus = "unknown"
 	_lastSubscriptionCheck = 0
 }