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

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

Author: James Mtendamema

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()
@@ -122,11 +133,13 @@ describe("ZooGatewayHandler", () => {
 	}
 
 	describe("constructor", () => {
-		it("requires authentication before constructing the client", () => {
-			expect(() => new ZooGatewayHandler({})).toThrow(
-				"Zoo Gateway requires authentication. Please sign in to Zoo Code first.",
+		it("allows construction without a session token (auth is enforced at request time)", () => {
+			expect(() => new ZooGatewayHandler({})).not.toThrow()
+			expect(OpenAI).toHaveBeenCalledWith(
+				expect.objectContaining({
+					apiKey: "not-provided",
+				}),
 			)
-			expect(OpenAI).not.toHaveBeenCalled()
 		})
 
 		it("prefers the secret-storage cache over a persisted profile token", () => {
@@ -193,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/vercel-ai-gateway.ts

@@ -42,7 +42,7 @@ export type VercelAiGatewayModel = z.infer<typeof vercelAiGatewayModelSchema>
  * VercelAiGatewayModelsResponse
  */
 
-const vercelAiGatewayModelsResponseSchema = z.object({
+export const vercelAiGatewayModelsResponseSchema = z.object({
 	object: z.string(),
 	data: z.array(vercelAiGatewayModelSchema),
 })

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

@@ -3,52 +3,13 @@ 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"
 
-// Reuse the same schemas and parsing logic from vercel-ai-gateway since the API format is identical
-import { type VercelAiGatewayModel, parseVercelAiGatewayModel } from "./vercel-ai-gateway"
-
-import { z } from "zod"
-
-/**
- * ZooGatewayPricing (same format as Vercel AI Gateway)
- */
-
-const zooGatewayPricingSchema = z.object({
-	input: z.string().optional(),
-	output: z.string().optional(),
-	input_cache_write: z.string().optional(),
-	input_cache_read: z.string().optional(),
-	image: z.string().optional(),
-})
-
-/**
- * ZooGatewayModel (same format as Vercel AI Gateway)
- */
-
-const zooGatewayModelSchema = z.object({
-	id: z.string(),
-	object: z.string(),
-	created: z.number(),
-	owned_by: z.string(),
-	name: z.string(),
-	description: z.string(),
-	context_window: z.number(),
-	max_tokens: z.number(),
-	type: z.string(),
-	pricing: zooGatewayPricingSchema,
-})
-
-/**
- * ZooGatewayModelsResponse
- */
-
-const zooGatewayModelsResponseSchema = z.object({
-	object: z.string(),
-	data: z.array(zooGatewayModelSchema),
-})
-
-type ZooGatewayModelsResponse = z.infer<typeof zooGatewayModelsResponseSchema>
+import {
+	type VercelAiGatewayModel,
+	parseVercelAiGatewayModel,
+	vercelAiGatewayModelsResponseSchema,
+} from "./vercel-ai-gateway"
 
 // Bound model discovery so a network stall can't hang provider initialization paths.
 const MODEL_DISCOVERY_TIMEOUT_MS = 15_000
@@ -63,30 +24,25 @@ 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}`
 	}
 
 	try {
-		const response = await axios.get<ZooGatewayModelsResponse>(`${baseURL}/models`, {
+		const response = await axios.get(`${baseURL}/models`, {
 			headers,
 			timeout: MODEL_DISCOVERY_TIMEOUT_MS,
 		})
-		const result = zooGatewayModelsResponseSchema.safeParse(response.data)
-
-		// Fall back to the raw response only when it looks structurally sound; otherwise return
-		// an empty list rather than crashing on `response.data.data` being undefined.
-		const data = result.success ? result.data.data : Array.isArray(response.data?.data) ? response.data.data : []
+		const result = vercelAiGatewayModelsResponseSchema.safeParse(response.data)
 
 		if (!result.success) {
 			console.error(`Zoo Gateway models response is invalid ${JSON.stringify(result.error.format())}`)
+			return models
 		}
 
-		for (const model of data) {
+		for (const model of result.data.data) {
 			const { id } = model
 
 			// Only include language models for chat inference.
@@ -95,8 +51,7 @@ export async function getZooGatewayModels(options?: ApiHandlerOptions): Promise<
 				continue
 			}
 
-			// Parse model using the same logic as Vercel AI Gateway since formats are identical
-			models[id] = parseZooGatewayModel({ id, model: model as VercelAiGatewayModel })
+			models[id] = parseZooGatewayModel({ id, model })
 		}
 	} catch (error) {
 		// Log only safe fields; never serialize the full error object because it
@@ -123,6 +78,5 @@ export async function getZooGatewayModels(options?: ApiHandlerOptions): Promise<
  */
 
 export const parseZooGatewayModel = ({ id, model }: { id: string; model: VercelAiGatewayModel }): ModelInfo => {
-	// Reuse the parsing logic from vercel-ai-gateway
 	return parseVercelAiGatewayModel({ id, model })
 }

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"
 
@@ -95,16 +95,13 @@ interface ZooGatewayUsage extends OpenAI.CompletionUsage {
 	cost?: number
 }
 
+const ZOO_GATEWAY_AUTH_ERROR = "Zoo Gateway requires authentication. Please sign in to Zoo Code first."
+
 export class ZooGatewayHandler extends RouterProvider implements SingleCompletionHandler {
 	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
-		if (!sessionToken) {
-			throw new Error("Zoo Gateway requires authentication. Please sign in to Zoo Code first.")
-		}
+		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
@@ -121,18 +118,26 @@ export class ZooGatewayHandler extends RouterProvider implements SingleCompletio
 			},
 			name: "zoo-gateway",
 			baseURL,
-			apiKey: sessionToken,
+			apiKey: sessionToken || "not-provided",
 			modelId: options.zooGatewayModelId,
 			defaultModelId: zooGatewayDefaultModelId,
 			defaultModelInfo: zooGatewayDefaultModelInfo,
 		})
 	}
 
+	private ensureAuthenticated(): void {
+		if (!resolveZooGatewaySessionToken(this.options.zooSessionToken)) {
+			throw new Error(ZOO_GATEWAY_AUTH_ERROR)
+		}
+	}
+
 	override async *createMessage(
 		systemPrompt: string,
 		messages: Anthropic.Messages.MessageParam[],
 		metadata?: ApiHandlerCreateMessageMetadata,
 	): ApiStream {
+		this.ensureAuthenticated()
+
 		const { id: modelId, info } = await this.fetchModel()
 
 		const openAiMessages: OpenAI.Chat.ChatCompletionMessageParam[] = [
@@ -209,17 +214,21 @@ 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
 		}
 	}
 
 	async completePrompt(prompt: string): Promise<string> {
+		this.ensureAuthenticated()
+
 		const { id: modelId, info } = await this.fetchModel()
 
 		try {
@@ -238,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)