fix: support async generators as response resolvers

src/core/handlers/RequestHandler.ts

@@ -1,4 +1,3 @@
-import { invariant } from 'outvariant'
 import { getCallFrame } from '../utils/internal/getCallFrame'
 import { isIterable } from '../utils/internal/isIterable'
 import type { ResponseResolutionContext } from '../utils/executeHandlers'
@@ -57,6 +56,11 @@ export type AsyncResponseResolverReturnType<
       MaybeAsyncResponseResolverReturnType<ResponseBodyType>,
       MaybeAsyncResponseResolverReturnType<ResponseBodyType>
     >
+  | AsyncGenerator<
+      MaybeAsyncResponseResolverReturnType<ResponseBodyType>,
+      MaybeAsyncResponseResolverReturnType<ResponseBodyType>,
+      MaybeAsyncResponseResolverReturnType<ResponseBodyType>
+    >
 >
 
 export type ResponseResolverInfo<
@@ -117,11 +121,17 @@ export abstract class RequestHandler<
   public isUsed: boolean
 
   protected resolver: ResponseResolver<ResolverExtras, any, any>
-  private resolverGenerator?: Generator<
-    MaybeAsyncResponseResolverReturnType<any>,
-    MaybeAsyncResponseResolverReturnType<any>,
-    MaybeAsyncResponseResolverReturnType<any>
-  >
+  private resolverGenerator?:
+    | Generator<
+        MaybeAsyncResponseResolverReturnType<any>,
+        MaybeAsyncResponseResolverReturnType<any>,
+        MaybeAsyncResponseResolverReturnType<any>
+      >
+    | AsyncGenerator<
+        MaybeAsyncResponseResolverReturnType<any>,
+        MaybeAsyncResponseResolverReturnType<any>,
+        MaybeAsyncResponseResolverReturnType<any>
+      >
   private resolverGeneratorResult?: Response | StrictResponse<any>
   private options?: HandlerOptions
 
@@ -256,6 +266,9 @@ export abstract class RequestHandler<
       return null
     }
 
+    // Preemptively mark the handler as used.
+    // Generators will undo this because only when the resolver reaches the
+    // "done" state of the generator that it considers the handler used.
     this.isUsed = true
 
     // Create a response extraction wrapper around the resolver
@@ -304,39 +317,28 @@ export abstract class RequestHandler<
       const result = this.resolverGenerator || (await resolver(info))
 
       if (isIterable<AsyncResponseResolverReturnType<any>>(result)) {
-        // Immediately mark this handler as unused.
-        // Only when the generator is done, the handler will be
-        // considered used.
+        // Opt-out from marking this handler as used.
         this.isUsed = false
 
-        const { value, done } = result[Symbol.iterator]().next()
-        const nextResponse = await value
-
-        if (done) {
-          this.isUsed = true
-        }
-
-        // If the generator is done and there is no next value,
-        // return the previous generator's value.
-        if (!nextResponse && done) {
-          invariant(
-            this.resolverGeneratorResult,
-            'Failed to returned a previously stored generator response: the value is not a valid Response.',
-          )
-
-          // Clone the previously stored response from the generator
-          // so that it could be read again.
-          return this.resolverGeneratorResult.clone() as StrictResponse<any>
-        }
-
         if (!this.resolverGenerator) {
           this.resolverGenerator = result
         }
 
+        const { done, value } = await this.resolverGenerator.next()
+        const nextResponse = await value
+
         if (nextResponse) {
-          // Also clone the response before storing it
-          // so it could be read again.
-          this.resolverGeneratorResult = nextResponse?.clone()
+          this.resolverGeneratorResult = nextResponse.clone()
+        }
+
+        if (done) {
+          // A one-time generator resolver stops affecting the network
+          // only after it's been completely exhausted.
+          this.isUsed = true
+
+          // Clone the previously stored response so it can be read
+          // when receiving it repeatedly from the "done" generator.
+          return this.resolverGeneratorResult?.clone()
         }
 
         return nextResponse

src/core/utils/internal/isIterable.ts

@@ -3,10 +3,14 @@
  */
 export function isIterable<IteratorType>(
   fn: any,
-): fn is Generator<IteratorType, IteratorType, IteratorType> {
+): fn is
+  | Generator<IteratorType, IteratorType, IteratorType>
+  | AsyncGenerator<IteratorType, IteratorType, IteratorType> {
   if (!fn) {
     return false
   }
 
-  return typeof (fn as Generator<unknown>)[Symbol.iterator] == 'function'
+  return (
+    Reflect.has(fn, Symbol.iterator) || Reflect.has(fn, Symbol.asyncIterator)
+  )
 }

test/node/rest-api/response/generator.test.ts

@@ -0,0 +1,106 @@
+/**
+ * @vitest-environment node
+ */
+import { http, HttpResponse, delay } from 'msw'
+import { setupServer } from 'msw/node'
+
+const server = setupServer()
+const toJson = (response: Response) => response.json()
+
+beforeAll(() => {
+  server.listen()
+})
+
+afterEach(() => {
+  server.resetHandlers()
+})
+
+afterAll(() => {
+  server.close()
+})
+
+it('supports generator function as response resolver', async () => {
+  server.use(
+    http.get('https://example.com/weather', function* () {
+      let degree = 10
+
+      while (degree < 13) {
+        degree++
+        yield HttpResponse.json(degree)
+      }
+
+      degree++
+      return HttpResponse.json(degree)
+    }),
+  )
+
+  // Must respond with yielded responses.
+  expect(await fetch('https://example.com/weather').then(toJson)).toEqual(11)
+  expect(await fetch('https://example.com/weather').then(toJson)).toEqual(12)
+  expect(await fetch('https://example.com/weather').then(toJson)).toEqual(13)
+  // Must respond with the final "done" response.
+  expect(await fetch('https://example.com/weather').then(toJson)).toEqual(14)
+  // Must keep responding with the final "done" response.
+  expect(await fetch('https://example.com/weather').then(toJson)).toEqual(14)
+})
+
+it('supports async generator function as response resolver', async () => {
+  server.use(
+    http.get('https://example.com/weather', async function* () {
+      await delay(20)
+
+      let degree = 10
+
+      while (degree < 13) {
+        degree++
+        yield HttpResponse.json(degree)
+      }
+
+      degree++
+      return HttpResponse.json(degree)
+    }),
+  )
+
+  expect(await fetch('https://example.com/weather').then(toJson)).toEqual(11)
+  expect(await fetch('https://example.com/weather').then(toJson)).toEqual(12)
+  expect(await fetch('https://example.com/weather').then(toJson)).toEqual(13)
+  expect(await fetch('https://example.com/weather').then(toJson)).toEqual(14)
+  expect(await fetch('https://example.com/weather').then(toJson)).toEqual(14)
+})
+
+it('supports generator function as one-time response resolver', async () => {
+  server.use(
+    http.get(
+      'https://example.com/weather',
+      function* () {
+        let degree = 10
+
+        while (degree < 13) {
+          degree++
+          yield HttpResponse.json(degree)
+        }
+
+        degree++
+        return HttpResponse.json(degree)
+      },
+      { once: true },
+    ),
+    http.get('*', () => {
+      return HttpResponse.json('fallback')
+    }),
+  )
+
+  // Must respond with the yielded incrementing responses.
+  expect(await fetch('https://example.com/weather').then(toJson)).toEqual(11)
+  expect(await fetch('https://example.com/weather').then(toJson)).toEqual(12)
+  expect(await fetch('https://example.com/weather').then(toJson)).toEqual(13)
+  // Must respond with the "done" final response from the iterator.
+  expect(await fetch('https://example.com/weather').then(toJson)).toEqual(14)
+  // Must respond with the other handler since the generator one is used.
+  expect(await fetch('https://example.com/weather').then(toJson)).toEqual(
+    'fallback',
+  )
+  expect(await fetch('https://example.com/weather').then(toJson)).toEqual(
+    'fallback',
+  )
+})

test/typings/resolver-generator.test-d.ts

@@ -0,0 +1,36 @@
+import { it } from 'vitest'
+import { http, HttpResponse } from 'msw'
+
+it('supports generator function as response resolver', () => {
+  http.get<never, never, { value: number }>('/', function* () {
+    yield HttpResponse.json({ value: 1 })
+    yield HttpResponse.json({ value: 2 })
+    return HttpResponse.json({ value: 3 })
+  })
+
+  http.get<never, never, { value: string }>('/', function* () {
+    yield HttpResponse.json({ value: 'one' })
+    yield HttpResponse.json({
+      // @ts-expect-error Expected string, got number.
+      value: 2,
+    })
+    return HttpResponse.json({ value: 'three' })
+  })
+})
+
+it('supports async generator function as response resolver', () => {
+  http.get<never, never, { value: number }>('/', async function* () {
+    yield HttpResponse.json({ value: 1 })
+    yield HttpResponse.json({ value: 2 })
+    return HttpResponse.json({ value: 3 })
+  })
+
+  http.get<never, never, { value: string }>('/', async function* () {
+    yield HttpResponse.json({ value: 'one' })
+    yield HttpResponse.json({
+      // @ts-expect-error Expected string, got number.
+      value: 2,
+    })
+    return HttpResponse.json({ value: 'three' })
+  })
+})