@@ -26,6 +26,7 @@ vitest.mock("../../../i18n", () => ({
2626 t : ( key : string ) => key ,
2727} ) )
2828
29+ import type { MockInstance } from "vitest"
2930import OpenAI from "openai"
3031
3132import { zooGatewayDefaultModelId , ZOO_GATEWAY_DEFAULT_TEMPERATURE } from "@roo-code/types"
@@ -500,10 +501,17 @@ describe("ZooGatewayHandler", () => {
500501 } )
501502
502503 describe ( "completePrompt" , ( ) => {
504+ let consoleError : MockInstance < typeof console . error >
505+
503506 beforeEach ( ( ) => {
504507 mockCreate . mockImplementation ( async ( ) => ( {
505508 choices : [ { message : { role : "assistant" , content : "Test completion response" } } ] ,
506509 } ) )
510+ consoleError = vitest . spyOn ( console , "error" ) . mockImplementation ( ( ) => { } )
511+ } )
512+
513+ afterEach ( ( ) => {
514+ consoleError . mockRestore ( )
507515 } )
508516
509517 it ( "returns completion text from the gateway" , async ( ) => {
@@ -520,6 +528,7 @@ describe("ZooGatewayHandler", () => {
520528 temperature : ZOO_GATEWAY_DEFAULT_TEMPERATURE ,
521529 max_completion_tokens : 64000 ,
522530 } ) ,
531+ { } ,
523532 )
524533 } )
525534
@@ -534,14 +543,287 @@ describe("ZooGatewayHandler", () => {
534543 )
535544 } )
536545
537- it ( "returns an empty string when the model returns no content " , async ( ) => {
546+ it ( "forwards a per-request signal and the timeout to the SDK " , async ( ) => {
538547 const handler = new ZooGatewayHandler ( mockOptions )
539- mockCreate . mockImplementation ( async ( ) => ( {
540- choices : [ { message : { role : "assistant" , content : null } } ] ,
541- } ) )
548+ const caller = new AbortController ( )
549+
550+ await handler . completePrompt ( "Test" , { abortSignal : caller . signal , timeoutMs : 5_000 } )
551+
552+ const requestOptions = mockCreate . mock . calls [ 0 ] [ 1 ] as { signal ?: AbortSignal ; timeout ?: number }
553+ expect ( requestOptions . signal ) . toBeInstanceOf ( AbortSignal )
554+ expect ( requestOptions . signal ) . not . toBe ( caller . signal )
555+ expect ( requestOptions . timeout ) . toBe ( 5_000 )
556+ } )
557+
558+ it ( "sends no signal when the caller gives none" , async ( ) => {
559+ const handler = new ZooGatewayHandler ( mockOptions )
560+
561+ await handler . completePrompt ( "Test" )
562+
563+ expect ( mockCreate . mock . calls [ 0 ] [ 1 ] ) . toEqual ( { } )
564+ } )
565+
566+ it . each ( [
567+ [ "with a session" , mockOptions ] ,
568+ [ "without a session" , { zooGatewayModelId : "anthropic/claude-sonnet-4" } ] ,
569+ ] ) (
570+ "rejects a pre-aborted signal with an AbortError, %s, before model lookup" ,
571+ async ( _label , options : ApiHandlerOptions ) => {
572+ const { getModels } = await import ( "../fetchers/modelCache" )
573+ const handler = new ZooGatewayHandler ( options )
574+ const caller = new AbortController ( )
575+ caller . abort ( )
576+
577+ await expect ( handler . completePrompt ( "Test" , { abortSignal : caller . signal } ) ) . rejects . toMatchObject ( {
578+ name : "AbortError" ,
579+ } )
580+ expect ( getModels ) . not . toHaveBeenCalled ( )
581+ expect ( mockCreate ) . not . toHaveBeenCalled ( )
582+ } ,
583+ )
584+
585+ it ( "keeps an ordinary model lookup failure while the caller signal is live" , async ( ) => {
586+ const { getModels, refreshModels } = await import ( "../fetchers/modelCache" )
587+ const lookupFailure = new Error ( "model catalog unavailable" )
588+ vitest . mocked ( getModels ) . mockRejectedValue ( lookupFailure )
589+ vitest . mocked ( refreshModels ) . mockRejectedValue ( lookupFailure )
590+ const handler = new ZooGatewayHandler ( mockOptions )
591+
592+ await expect ( handler . completePrompt ( "Test" , { abortSignal : new AbortController ( ) . signal } ) ) . rejects . toBe (
593+ lookupFailure ,
594+ )
595+ } )
596+
597+ it . each ( [
598+ [ "a raw AbortError" , Object . assign ( new Error ( "The user aborted a request" ) , { name : "AbortError" } ) ] ,
599+ [ "the SDK abort error" , new Error ( "Request was aborted." ) ] ,
600+ ] ) (
601+ "normalizes %s from the model lookup to the provider AbortError while the caller signal is live" ,
602+ async ( _label , lookupFailure : Error ) => {
603+ const { getModels, refreshModels } = await import ( "../fetchers/modelCache" )
604+ vitest . mocked ( getModels ) . mockRejectedValue ( lookupFailure )
605+ vitest . mocked ( refreshModels ) . mockRejectedValue ( lookupFailure )
606+ const handler = new ZooGatewayHandler ( mockOptions )
607+
608+ await expect (
609+ handler . completePrompt ( "Test" , { abortSignal : new AbortController ( ) . signal } ) ,
610+ ) . rejects . toMatchObject ( {
611+ name : "AbortError" ,
612+ message : "The Zoo Gateway request was aborted" ,
613+ } )
614+ expect ( mockCreate ) . not . toHaveBeenCalled ( )
615+ } ,
616+ )
617+
618+ it ( "rejects with an AbortError when the timeout alone expires during the SDK call" , async ( ) => {
619+ const handler = new ZooGatewayHandler ( mockOptions )
620+ let sdkSignal : AbortSignal | undefined
621+ mockCreate . mockImplementation ( ( _params : unknown , requestOptions : { signal ?: AbortSignal } ) => {
622+ sdkSignal = requestOptions . signal
623+ return new Promise ( ( ) => { } )
624+ } )
625+
626+ await expect ( handler . completePrompt ( "Test" , { timeoutMs : 20 } ) ) . rejects . toMatchObject ( {
627+ name : "AbortError" ,
628+ message : "The Zoo Gateway request was aborted" ,
629+ } )
630+ expect ( sdkSignal ?. aborted ) . toBe ( true )
631+ expect ( showErrorMessage ) . not . toHaveBeenCalled ( )
632+ } )
633+
634+ it ( "settles on abort while a gateway error notification waits for the user" , async ( ) => {
635+ const handler = new ZooGatewayHandler ( mockOptions )
636+ const caller = new AbortController ( )
637+ mockCreate . mockRejectedValue ( makeApiError ( 402 ) )
638+ showErrorMessage . mockImplementation ( ( ) => new Promise ( ( ) => { } ) )
639+
640+ const pending = handler . completePrompt ( "Test" , { abortSignal : caller . signal } )
641+ await vitest . waitFor ( ( ) => expect ( showErrorMessage ) . toHaveBeenCalled ( ) )
642+ caller . abort ( )
643+
644+ await expect ( pending ) . rejects . toMatchObject ( {
645+ name : "AbortError" ,
646+ message : "The Zoo Gateway request was aborted" ,
647+ } )
648+ // The abort is not a failure of the notification, so nothing is logged.
649+ expect ( consoleError ) . not . toHaveBeenCalled ( )
650+ } )
651+
652+ it . each ( [
653+ [ "with a live caller signal" , true , new Error ( "notification failed" ) , "notification failed" ] ,
654+ [ "without a caller signal" , false , "notification dismissed" , "notification dismissed" ] ,
655+ ] ) (
656+ "logs a failed gateway error notification and still rejects with the gateway error, %s" ,
657+ async ( _label , withSignal : boolean , notificationFailure : unknown , logged : string ) => {
658+ const handler = new ZooGatewayHandler ( mockOptions )
659+ mockCreate . mockRejectedValue ( makeApiError ( 402 , { message : "out of credits" } ) )
660+ showErrorMessage . mockRejectedValue ( notificationFailure )
661+ const options = withSignal ? { abortSignal : new AbortController ( ) . signal } : undefined
662+
663+ await expect ( handler . completePrompt ( "Test" , options ) ) . rejects . toThrow (
664+ "Zoo Gateway completion error: out of credits" ,
665+ )
666+ expect ( consoleError ) . toHaveBeenCalledWith ( "Failed to surface Zoo Gateway error:" , logged )
667+ } ,
668+ )
669+
670+ it ( "rethrows an SDK failure that is not an Error unchanged" , async ( ) => {
671+ const handler = new ZooGatewayHandler ( mockOptions )
672+ const failure = { reason : "upstream failure" }
673+ mockCreate . mockRejectedValue ( failure )
674+
675+ await expect ( handler . completePrompt ( "Test" , { abortSignal : new AbortController ( ) . signal } ) ) . rejects . toBe (
676+ failure ,
677+ )
678+ } )
679+
680+ it ( "rejects with an AbortError when the timeout alone expires during model lookup" , async ( ) => {
681+ const { getModels } = await import ( "../fetchers/modelCache" )
682+ vitest . mocked ( getModels ) . mockImplementation ( ( ) => new Promise ( ( ) => { } ) )
683+ const handler = new ZooGatewayHandler ( mockOptions )
684+
685+ await expect ( handler . completePrompt ( "Test" , { timeoutMs : 20 } ) ) . rejects . toMatchObject ( {
686+ name : "AbortError" ,
687+ message : "The Zoo Gateway request was aborted" ,
688+ } )
689+ expect ( mockCreate ) . not . toHaveBeenCalled ( )
690+ } )
691+
692+ it ( "rejects with an AbortError when the signal aborts during model lookup" , async ( ) => {
693+ const { getModels } = await import ( "../fetchers/modelCache" )
694+ vitest . mocked ( getModels ) . mockImplementation ( ( ) => new Promise ( ( ) => { } ) )
695+ const handler = new ZooGatewayHandler ( mockOptions )
696+ const caller = new AbortController ( )
697+
698+ const pending = handler . completePrompt ( "Test" , { abortSignal : caller . signal } )
699+ caller . abort ( )
700+
701+ await expect ( pending ) . rejects . toMatchObject ( {
702+ name : "AbortError" ,
703+ message : "The Zoo Gateway request was aborted" ,
704+ } )
705+ expect ( mockCreate ) . not . toHaveBeenCalled ( )
706+ } )
707+
708+ it ( "rejects with an AbortError at once when aborted mid-flight, without a gateway error" , async ( ) => {
709+ const handler = new ZooGatewayHandler ( mockOptions )
710+ const caller = new AbortController ( )
711+ let sdkSignal : AbortSignal | undefined
712+ // The SDK call never settles, as when the SDK sleeps between retries.
713+ mockCreate . mockImplementation ( ( _params : unknown , requestOptions : { signal ?: AbortSignal } ) => {
714+ sdkSignal = requestOptions . signal
715+ return new Promise ( ( ) => { } )
716+ } )
717+
718+ const pending = handler . completePrompt ( "Test" , { abortSignal : caller . signal } )
719+ await vitest . waitFor ( ( ) => expect ( mockCreate ) . toHaveBeenCalled ( ) )
720+ caller . abort ( )
721+
722+ await expect ( pending ) . rejects . toMatchObject ( {
723+ name : "AbortError" ,
724+ message : "The Zoo Gateway request was aborted" ,
725+ } )
726+ expect ( sdkSignal ?. aborted ) . toBe ( true )
727+ expect ( showErrorMessage ) . not . toHaveBeenCalled ( )
728+ } )
729+
730+ it ( "reports an SDK abort error after the caller aborts as an AbortError, not a gateway error" , async ( ) => {
731+ const handler = new ZooGatewayHandler ( mockOptions )
732+ const caller = new AbortController ( )
733+ mockCreate . mockImplementation ( async ( ) => {
734+ caller . abort ( )
735+ throw new Error ( "Request was aborted." )
736+ } )
737+
738+ await expect ( handler . completePrompt ( "Test" , { abortSignal : caller . signal } ) ) . rejects . toMatchObject ( {
739+ name : "AbortError" ,
740+ } )
741+ expect ( showErrorMessage ) . not . toHaveBeenCalled ( )
742+ } )
743+
744+ it ( "shows no gateway error for a 402 that rejects just before the caller aborts" , async ( ) => {
745+ const handler = new ZooGatewayHandler ( mockOptions )
746+ const caller = new AbortController ( )
747+ let rejectCompletion : ( reason : unknown ) => void = ( ) => { }
748+ const completion = new Promise ( ( _resolve , reject ) => {
749+ rejectCompletion = reject
750+ } )
751+ mockCreate . mockImplementation ( ( ) => completion )
752+
753+ const pending = handler . completePrompt ( "Test" , { abortSignal : caller . signal } )
754+ await vitest . waitFor ( ( ) => expect ( mockCreate ) . toHaveBeenCalled ( ) )
755+ // This callback runs after the handler's own race settles with the 402 but before completePrompt
756+ // resumes, so only the abort check before the notification keeps the 402 notification away.
757+ completion . catch ( ( ) => caller . abort ( ) )
758+ rejectCompletion ( makeApiError ( 402 ) )
759+
760+ await expect ( pending ) . rejects . toMatchObject ( { name : "AbortError" } )
761+ expect ( showErrorMessage ) . not . toHaveBeenCalled ( )
762+ } )
763+
764+ it ( "does not return a completion that resolves just before the caller aborts" , async ( ) => {
765+ const handler = new ZooGatewayHandler ( mockOptions )
766+ const caller = new AbortController ( )
767+ let resolveCompletion : ( value : unknown ) => void = ( ) => { }
768+ const completion = new Promise ( ( resolve ) => {
769+ resolveCompletion = resolve
770+ } )
771+ mockCreate . mockImplementation ( ( ) => completion )
772+
773+ const pending = handler . completePrompt ( "Test" , { abortSignal : caller . signal } )
774+ await vitest . waitFor ( ( ) => expect ( mockCreate ) . toHaveBeenCalled ( ) )
775+ // This callback runs after the handler's own race settles but before completePrompt resumes,
776+ // so only the check after the await can drop the late result.
777+ void completion . then ( ( ) => caller . abort ( ) )
778+ resolveCompletion ( { choices : [ { message : { role : "assistant" , content : "late" } } ] } )
779+
780+ await expect ( pending ) . rejects . toMatchObject ( { name : "AbortError" } )
781+ } )
782+
783+ it . each ( [
784+ [ "a completed request" , ( ) => { } ] ,
785+ [ "a failed request" , ( ) => mockCreate . mockRejectedValue ( new Error ( "upstream failure" ) ) ] ,
786+ [ "a request aborted mid-flight" , ( ) => mockCreate . mockImplementation ( ( ) => new Promise ( ( ) => { } ) ) ] ,
787+ ] ) ( "removes the same abort listener from the caller's signal after %s" , async ( _label , setup ) => {
788+ setup ( )
789+ const handler = new ZooGatewayHandler ( mockOptions )
790+ const caller = new AbortController ( )
791+ const addListener = vitest . spyOn ( caller . signal , "addEventListener" )
792+ const removeListener = vitest . spyOn ( caller . signal , "removeEventListener" )
793+
794+ const pending = handler . completePrompt ( "Test" , { abortSignal : caller . signal } ) . catch ( ( ) => undefined )
795+ await vitest . waitFor ( ( ) => expect ( mockCreate ) . toHaveBeenCalled ( ) )
796+ caller . abort ( )
797+ await pending
798+
799+ const abortListeners = addListener . mock . calls . filter ( ( [ type ] ) => type === "abort" )
800+ expect ( abortListeners ) . toHaveLength ( 1 )
801+ expect ( removeListener ) . toHaveBeenCalledWith ( "abort" , abortListeners [ 0 ] [ 1 ] )
802+ } )
803+
804+ it . each ( [
805+ [ "no content" , [ { message : { role : "assistant" , content : null } } ] ] ,
806+ [ "no choices" , [ ] ] ,
807+ ] ) ( "returns an empty string when the model returns %s" , async ( _label , choices : unknown [ ] ) => {
808+ const handler = new ZooGatewayHandler ( mockOptions )
809+ mockCreate . mockImplementation ( async ( ) => ( { choices } ) )
542810
543811 await expect ( handler . completePrompt ( "Test" ) ) . resolves . toBe ( "" )
544812 } )
813+
814+ it . each ( [
815+ [ "with a live caller signal" , { abortSignal : new AbortController ( ) . signal } ] ,
816+ [ "without a caller signal" , undefined ] ,
817+ ] ) ( "requires a session before the model lookup, %s" , async ( _label , options ) => {
818+ const { getModels } = await import ( "../fetchers/modelCache" )
819+ const handler = new ZooGatewayHandler ( { zooGatewayModelId : "anthropic/claude-sonnet-4" } )
820+
821+ await expect ( handler . completePrompt ( "Test" , options ) ) . rejects . toThrow (
822+ "Zoo Gateway requires authentication. Please sign in to Zoo Code first." ,
823+ )
824+ expect ( getModels ) . not . toHaveBeenCalled ( )
825+ expect ( mockCreate ) . not . toHaveBeenCalled ( )
826+ } )
545827 } )
546828
547829 describe ( "classifyGatewayApiError" , ( ) => {
0 commit comments