fix(zoo-gateway): defer auth check, fail-closed models, reuse Vercel schemas

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

Author: James Mtendamema

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

@@ -82,11 +82,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("initializes OpenAI with Zoo enrichment headers and session token", () => {
@@ -160,6 +162,17 @@ describe("ZooGatewayHandler", () => {
 			}))
 		})
 
+		it("requires authentication at request time when no session token is available", async () => {
+			const handler = new ZooGatewayHandler({})
+			const stream = handler.createMessage("You are helpful.", [{ role: "user", content: "Hello" }])
+
+			await expect(async () => {
+				for await (const _chunk of stream) {
+					// drain
+				}
+			}).rejects.toThrow("Zoo Gateway requires authentication. Please sign in to Zoo Code first.")
+		})
+
 		it("streams text and usage chunks", async () => {
 			const handler = new ZooGatewayHandler(mockOptions)
 			const stream = handler.createMessage("You are helpful.", [{ role: "user", content: "Hello" }])

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 { getZooCodeBaseUrl } from "../../../services/zoo-code-auth"
+import { getCachedZooCodeToken, getZooCodeBaseUrl } 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,28 +24,27 @@ 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
+	// 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 headers: Record<string, string> = {}
-	if (options?.zooSessionToken) {
-		headers["Authorization"] = `Bearer ${options.zooSessionToken}`
+	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.
@@ -93,8 +53,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
@@ -121,6 +80,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

@@ -9,7 +9,7 @@ import {
 } from "@roo-code/types"
 
 import { ApiHandlerOptions } from "../../shared/api"
-import { getZooCodeBaseUrl } from "../../services/zoo-code-auth"
+import { getCachedZooCodeToken, getZooCodeBaseUrl } from "../../services/zoo-code-auth"
 import { Package } from "../../shared/package"
 
 import { ApiStream } from "../transform/stream"
@@ -25,16 +25,15 @@ 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`
 
-		// Fail fast with a clear message instead of waiting for a 401.
-		// The token is set automatically by handleZooCodeCallback() after the user
-		// authenticates via the "Sign in with Zoo Code" flow in the extension.
-		if (!options.zooSessionToken) {
-			throw new Error("Zoo Gateway requires authentication. Please sign in to Zoo Code first.")
-		}
+		// 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
 
 		// Merge Zoo-specific enrichment headers into openAiHeaders so they flow through
 		// the parent's single OpenAI client. We avoid reassigning `this.client` (which
@@ -51,18 +50,27 @@ export class ZooGatewayHandler extends RouterProvider implements SingleCompletio
 			},
 			name: "zoo-gateway",
 			baseURL,
-			apiKey: options.zooSessionToken,
+			apiKey: sessionToken || "not-provided",
 			modelId: options.zooGatewayModelId,
 			defaultModelId: zooGatewayDefaultModelId,
 			defaultModelInfo: zooGatewayDefaultModelInfo,
 		})
 	}
 
+	private ensureAuthenticated(): void {
+		const sessionToken = getCachedZooCodeToken() || this.options.zooSessionToken
+		if (!sessionToken) {
+			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[] = [
@@ -140,6 +148,8 @@ export class ZooGatewayHandler extends RouterProvider implements SingleCompletio
 	}
 
 	async completePrompt(prompt: string): Promise<string> {
+		this.ensureAuthenticated()
+
 		const { id: modelId, info } = await this.fetchModel()
 
 		try {