11// npx vitest run src/api/providers/__tests__/zoo-gateway.spec.ts
22
3- vitest . mock ( "vscode" , ( ) => ( { } ) )
3+ const { showErrorMessage, openExternal } = vitest . hoisted ( ( ) => ( {
4+ showErrorMessage : vitest . fn ( async ( ) => undefined as string | undefined ) ,
5+ openExternal : vitest . fn ( async ( ) => true ) ,
6+ } ) )
7+
8+ vitest . mock ( "vscode" , ( ) => ( {
9+ window : { showErrorMessage } ,
10+ env : { openExternal, uriScheme : "vscode" , appName : "VS Code" } ,
11+ Uri : { parse : ( value : string ) => ( { toString : ( ) => value } ) } ,
12+ } ) )
13+
14+ vitest . mock ( "../../../i18n" , ( ) => ( {
15+ t : ( key : string ) => key ,
16+ } ) )
417
518import OpenAI from "openai"
619
@@ -9,6 +22,7 @@ import { zooGatewayDefaultModelId, ZOO_GATEWAY_DEFAULT_TEMPERATURE } from "@roo-
922import { ZooGatewayHandler } from "../zoo-gateway"
1023import { ApiHandlerOptions } from "../../../shared/api"
1124import { Package } from "../../../shared/package"
25+ import { clearZooCodeToken } from "../../../services/zoo-code-auth"
1226
1327vitest . mock ( "openai" )
1428vitest . mock ( "delay" , ( ) => ( { default : vitest . fn ( ( ) => Promise . resolve ( ) ) } ) )
@@ -78,9 +92,32 @@ describe("ZooGatewayHandler", () => {
7892 beforeEach ( ( ) => {
7993 vitest . clearAllMocks ( )
8094 mockCreate . mockClear ( )
95+ showErrorMessage . mockReset ( )
96+ showErrorMessage . mockResolvedValue ( undefined )
97+ openExternal . mockReset ( )
98+ openExternal . mockResolvedValue ( true )
8199 mockOpenAIClient ( )
82100 } )
83101
102+ function makeApiError ( status : number , options : { code ?: string ; message ?: string } = { } ) {
103+ const err = new Error ( options . message ?? `HTTP ${ status } ` ) as Error & {
104+ status : number
105+ code ?: string
106+ }
107+ err . status = status
108+ if ( options . code ) err . code = options . code
109+ return err
110+ }
111+
112+ async function drainCreateMessage ( handler : ZooGatewayHandler ) {
113+ const stream = handler . createMessage ( "system" , [ { role : "user" , content : "hi" } ] )
114+ const out : unknown [ ] = [ ]
115+ for await ( const chunk of stream ) {
116+ out . push ( chunk )
117+ }
118+ return out
119+ }
120+
84121 describe ( "constructor" , ( ) => {
85122 it ( "requires authentication before constructing the client" , ( ) => {
86123 expect ( ( ) => new ZooGatewayHandler ( { } ) ) . toThrow (
@@ -323,4 +360,112 @@ describe("ZooGatewayHandler", () => {
323360 await expect ( handler . completePrompt ( "Test" ) ) . resolves . toBe ( "" )
324361 } )
325362 } )
363+
364+ describe ( "surfaceGatewayApiError" , ( ) => {
365+ it ( "clears the cached token and offers re-sign-in on 401" , async ( ) => {
366+ const handler = new ZooGatewayHandler ( mockOptions )
367+ mockCreate . mockImplementation ( ( ) => {
368+ throw makeApiError ( 401 )
369+ } )
370+ showErrorMessage . mockResolvedValueOnce ( "common:zooAuth.buttons.sign_in" )
371+
372+ await expect ( drainCreateMessage ( handler ) ) . rejects . toThrow ( )
373+ expect ( clearZooCodeToken ) . toHaveBeenCalledTimes ( 1 )
374+ expect ( showErrorMessage ) . toHaveBeenCalledWith (
375+ "common:zooAuth.errors.session_expired" ,
376+ "common:zooAuth.buttons.sign_in" ,
377+ )
378+ expect ( openExternal ) . toHaveBeenCalledTimes ( 1 )
379+ } )
380+
381+ it ( "does not open a URL on 401 when the user dismisses the prompt" , async ( ) => {
382+ const handler = new ZooGatewayHandler ( mockOptions )
383+ mockCreate . mockImplementation ( ( ) => {
384+ throw makeApiError ( 401 )
385+ } )
386+ showErrorMessage . mockResolvedValueOnce ( undefined )
387+
388+ await expect ( drainCreateMessage ( handler ) ) . rejects . toThrow ( )
389+ expect ( clearZooCodeToken ) . toHaveBeenCalledTimes ( 1 )
390+ expect ( openExternal ) . not . toHaveBeenCalled ( )
391+ } )
392+
393+ it ( "prompts to add credits on 402" , async ( ) => {
394+ const handler = new ZooGatewayHandler ( mockOptions )
395+ mockCreate . mockImplementation ( ( ) => {
396+ throw makeApiError ( 402 )
397+ } )
398+ showErrorMessage . mockResolvedValueOnce ( "common:zooAuth.buttons.add_credits" )
399+
400+ await expect ( drainCreateMessage ( handler ) ) . rejects . toThrow ( )
401+ expect ( clearZooCodeToken ) . not . toHaveBeenCalled ( )
402+ expect ( showErrorMessage ) . toHaveBeenCalledWith (
403+ "common:zooAuth.errors.out_of_credits" ,
404+ "common:zooAuth.buttons.add_credits" ,
405+ )
406+ expect ( openExternal ) . toHaveBeenCalledTimes ( 1 )
407+ } )
408+
409+ it ( "shows the budget message on 429 with a budget code" , async ( ) => {
410+ const handler = new ZooGatewayHandler ( mockOptions )
411+ mockCreate . mockImplementation ( ( ) => {
412+ throw makeApiError ( 429 , { code : "monthly_budget_exceeded" } )
413+ } )
414+
415+ await expect ( drainCreateMessage ( handler ) ) . rejects . toThrow ( )
416+ expect ( showErrorMessage ) . toHaveBeenCalledWith (
417+ "common:zooAuth.errors.budget_exceeded" ,
418+ "common:zooAuth.buttons.add_credits" ,
419+ )
420+ } )
421+
422+ it ( "does not surface a notification on 429 without a budget code" , async ( ) => {
423+ const handler = new ZooGatewayHandler ( mockOptions )
424+ mockCreate . mockImplementation ( ( ) => {
425+ throw makeApiError ( 429 , { code : "rate_limited" } )
426+ } )
427+
428+ await expect ( drainCreateMessage ( handler ) ) . rejects . toThrow ( )
429+ expect ( showErrorMessage ) . not . toHaveBeenCalled ( )
430+ } )
431+
432+ it ( "offers contact support on 403" , async ( ) => {
433+ const handler = new ZooGatewayHandler ( mockOptions )
434+ mockCreate . mockImplementation ( ( ) => {
435+ throw makeApiError ( 403 )
436+ } )
437+ showErrorMessage . mockResolvedValueOnce ( "common:zooAuth.buttons.contact_support" )
438+
439+ await expect ( drainCreateMessage ( handler ) ) . rejects . toThrow ( )
440+ expect ( showErrorMessage ) . toHaveBeenCalledWith (
441+ "common:zooAuth.errors.account_unavailable" ,
442+ "common:zooAuth.buttons.contact_support" ,
443+ )
444+ expect ( openExternal ) . toHaveBeenCalledTimes ( 1 )
445+ } )
446+
447+ it ( "ignores errors without an HTTP status" , async ( ) => {
448+ const handler = new ZooGatewayHandler ( mockOptions )
449+ mockCreate . mockImplementation ( ( ) => {
450+ throw new Error ( "network down" )
451+ } )
452+
453+ await expect ( drainCreateMessage ( handler ) ) . rejects . toThrow ( "network down" )
454+ expect ( showErrorMessage ) . not . toHaveBeenCalled ( )
455+ expect ( clearZooCodeToken ) . not . toHaveBeenCalled ( )
456+ } )
457+
458+ it ( "surfaces the gateway error then wraps the message in completePrompt" , async ( ) => {
459+ const handler = new ZooGatewayHandler ( mockOptions )
460+ mockCreate . mockImplementation ( ( ) => {
461+ throw makeApiError ( 402 , { message : "out of credits" } )
462+ } )
463+
464+ await expect ( handler . completePrompt ( "ping" ) ) . rejects . toThrow ( "Zoo Gateway completion error: out of credits" )
465+ expect ( showErrorMessage ) . toHaveBeenCalledWith (
466+ "common:zooAuth.errors.out_of_credits" ,
467+ "common:zooAuth.buttons.add_credits" ,
468+ )
469+ } )
470+ } )
326471} )
0 commit comments