Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 39 additions & 0 deletions src/api/providers/__tests__/vercel-ai-gateway.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -190,6 +190,45 @@ describe("VercelAiGatewayHandler", () => {
})
})

it("throws the upstream reason when an in-stream error chunk is received", async () => {
mockCreate.mockImplementation(async () => ({
[Symbol.asyncIterator]: async function* () {
yield {
error: {
message: "Too many requests, please wait before trying again",
code: 429,
},
}
},
}))

const handler = new VercelAiGatewayHandler(mockOptions)
const stream = handler.createMessage("You are a helpful assistant.", [{ role: "user", content: "Hello" }])

await expect(async () => {
for await (const _chunk of stream) {
// drain
}
}).rejects.toThrow("Too many requests, please wait before trying again")
})

it("throws a default message when an in-stream error chunk has no message", async () => {
mockCreate.mockImplementation(async () => ({
[Symbol.asyncIterator]: async function* () {
yield { error: {} }
},
}))

const handler = new VercelAiGatewayHandler(mockOptions)
const stream = handler.createMessage("You are a helpful assistant.", [{ role: "user", content: "Hello" }])

await expect(async () => {
for await (const _chunk of stream) {
// drain
}
}).rejects.toThrow("Vercel AI Gateway stream error")
})

it("uses correct temperature from options", async () => {
const customTemp = 0.5
const handler = new VercelAiGatewayHandler({
Expand Down
70 changes: 69 additions & 1 deletion src/api/providers/__tests__/zoo-gateway.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ import OpenAI from "openai"

import { zooGatewayDefaultModelId, ZOO_GATEWAY_DEFAULT_TEMPERATURE } from "@roo-code/types"

import { ZooGatewayHandler, classifyGatewayApiError } from "../zoo-gateway"
import { ZooGatewayHandler, classifyGatewayApiError, toGatewayStreamError } from "../zoo-gateway"
import { ApiHandlerOptions } from "../../../shared/api"
import { Package } from "../../../shared/package"
import { clearZooCodeToken } from "../../../services/zoo-code-auth"
Expand Down Expand Up @@ -363,6 +363,48 @@ describe("ZooGatewayHandler", () => {
},
])
})

it("throws the upstream reason when the gateway sends an in-stream error chunk", async () => {
mockCreate.mockImplementation(async () => ({
[Symbol.asyncIterator]: async function* () {
yield {
error: {
message: "Too many requests, please wait before trying again",
status: 429,
code: "rate_limited",
},
}
},
}))

const handler = new ZooGatewayHandler(mockOptions)

await expect(drainCreateMessage(handler)).rejects.toThrow(
"Too many requests, please wait before trying again",
)
})

it("surfaces the add-credits prompt when an in-stream error carries a budget code", async () => {
mockCreate.mockImplementation(async () => ({
[Symbol.asyncIterator]: async function* () {
yield {
error: {
message: "Monthly budget exceeded",
status: 429,
code: "monthly_budget_exceeded",
},
}
},
}))

const handler = new ZooGatewayHandler(mockOptions)

await expect(drainCreateMessage(handler)).rejects.toThrow()
expect(showErrorMessage).toHaveBeenCalledWith(
"common:zooAuth.errors.budget_exceeded",
"common:zooAuth.buttons.add_credits",
)
})
})

describe("completePrompt", () => {
Expand Down Expand Up @@ -443,6 +485,32 @@ describe("ZooGatewayHandler", () => {
})
})

describe("toGatewayStreamError", () => {
it("preserves the message, status, and code from the chunk", () => {
const error = toGatewayStreamError({
message: "rate limited",
status: 429,
code: "rate_limited",
}) as Error & {
status?: number
code?: string
}

expect(error).toBeInstanceOf(Error)
expect(error.message).toBe("rate limited")
expect(error.status).toBe(429)
expect(error.code).toBe("rate_limited")
})

it("falls back to a default message and leaves status/code undefined", () => {
const error = toGatewayStreamError({}) as Error & { status?: number; code?: string }

expect(error.message).toBe("Zoo Gateway stream error")
expect(error.status).toBeUndefined()
expect(error.code).toBeUndefined()
})
})

describe("surfaceGatewayApiError", () => {
it("clears the cached token and offers re-sign-in on 401", async () => {
const handler = new ZooGatewayHandler(mockOptions)
Expand Down
11 changes: 11 additions & 0 deletions src/api/providers/vercel-ai-gateway.ts
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,17 @@ export class VercelAiGatewayHandler extends RouterProvider implements SingleComp
const completion = await this.client.chat.completions.create(body)

for await (const chunk of completion) {
// Vercel AI Gateway reports mid-stream failures as an in-band error chunk
// rather than throwing, so surface it instead of returning an empty response.
if ("error" in chunk && chunk.error) {
const raw = chunk.error as { message?: unknown }
const message =
typeof raw.message === "string" && raw.message.length > 0
? raw.message
: "Vercel AI Gateway stream error"
throw new Error(message)
}

const delta = chunk.choices[0]?.delta
if (delta?.content) {
yield {
Expand Down
21 changes: 21 additions & 0 deletions src/api/providers/zoo-gateway.ts
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,20 @@ function getApiErrorCode(error: unknown): string | undefined {
return undefined
}

// The gateway sends in-band stream errors as `{ message, status?, code? }`. Rebuild
// them into an Error carrying status/code so the same classify/surface logic that
// handles thrown HTTP errors applies to mid-stream failures too.
// Exported for unit tests.
export function toGatewayStreamError(raw: unknown): Error {
const err = raw as { message?: unknown; status?: unknown; code?: unknown } | null
const message =
typeof err?.message === "string" && err.message.length > 0 ? err.message : "Zoo Gateway stream error"
return Object.assign(new Error(message), {
status: typeof err?.status === "number" ? err.status : undefined,
code: typeof err?.code === "string" ? err.code : undefined,
})
}

function buildZooCodeSignInUrl(): string {
const callbackUri = encodeURIComponent(
`${vscode.env.uriScheme}://${Package.publisher}.${Package.name}/auth-callback`,
Expand Down Expand Up @@ -209,6 +223,13 @@ export class ZooGatewayHandler extends RouterProvider implements SingleCompletio
})

for await (const chunk of completion) {
// Once the gateway starts streaming the HTTP status is already 200, so it
// reports upstream failures (e.g. provider rate limits) as an in-band error
// chunk. Surface it so the user sees the real reason instead of an empty reply.
if ("error" in chunk && chunk.error) {
throw toGatewayStreamError(chunk.error)
}

const delta = chunk.choices[0]?.delta
if (delta?.content) {
yield {
Expand Down
Loading