fix(peer, shared, fetch): end a waiting next() with done: true when an event stream is stopped

packages/fetch/src/event-stream.test.ts

@@ -1,5 +1,5 @@
 import { ErrorEvent, getEventMeta, withEventMeta } from '@standard-server/core'
-import { AbortError, isAsyncIteratorObject, sleep } from '@standard-server/shared'
+import { isAsyncIteratorObject, sleep } from '@standard-server/shared'
 import { toAsyncIteratorObject, toEventStream } from './event-stream'
 
 beforeEach(() => {
@@ -160,8 +160,8 @@ describe('toAsyncIteratorObject', () => {
       return true
     })
 
-    // should throw if .return is called while waiting for .next()
-    const nextPromise = expect(generator.next()).rejects.toBeInstanceOf(AbortError)
+    // should end if .return is called while waiting for .next()
+    const nextPromise = expect(generator.next()).resolves.toEqual({ done: true, value: undefined })
     await vi.advanceTimersByTimeAsync(0)
 
     await generator.return(undefined)

packages/fetch/src/event-stream.ts

@@ -1,5 +1,5 @@
 import { encodeEventStreamMessage, ErrorEvent, EventStreamDecoderStream, getEventMeta, unwrapEvent, withEventMeta } from '@standard-server/core'
-import { AbortError, AsyncIteratorClass, isTypescriptObject, parseEmptyableJSON, stringifyJSON } from '@standard-server/shared'
+import { AsyncIteratorClass, isTypescriptObject, parseEmptyableJSON, stringifyJSON } from '@standard-server/shared'
 
 export function toAsyncIteratorObject(
   stream: ReadableStream<Uint8Array<ArrayBuffer>> | null,
@@ -9,7 +9,6 @@ export function toAsyncIteratorObject(
     .pipeThrough(new EventStreamDecoderStream())
 
   const reader = eventStream?.getReader()
-  let isCancelled = false
 
   return new AsyncIteratorClass(async () => {
     while (true) {
@@ -20,22 +19,11 @@ export function toAsyncIteratorObject(
       const { done, value } = await reader.read()
 
       /**
-       * Handle stream completion scenarios:
-       *
-       * 1. If the reader is cancelled while waiting for the next value,
-       *    reader.read() will resolve as { done: true, value: undefined }.
-       *    However, this behavior is unreliable and we should only resolve
-       *    a value when the sender explicitly indicates completion.
-       *
-       * 2. The only implicit behavior we allow is when the sender successfully
-       *    closes the stream without sending a 'close' event - in this case,
-       *    we resolve with { done: true, value: undefined }.
+       * The sender closed the stream without sending a 'close' event.
+       * A read cancelled by `return()` also ends here, and AsyncIteratorClass
+       * resolves it as done either way.
        */
       if (done) {
-        if (isCancelled) {
-          throw new AbortError('Stream was cancelled')
-        }
-
         return { done: true, value: undefined }
       }
 
@@ -69,11 +57,7 @@ export function toAsyncIteratorObject(
         }
       }
     }
-  }, async (state) => {
-    if (state.kind === 'cancelled') {
-      isCancelled = true
-    }
-
+  }, async () => {
     await reader?.cancel()
   })
 }

packages/peer/src/server.test.ts

@@ -246,6 +246,28 @@ describe('serverPeer', () => {
         expect(success).toBeTruthy()
       })
 
+      it('ends a pending next() when the handler stops reading event-stream', async () => {
+        const { handler, box } = deferredHandler()
+
+        const msg = makeRequestMessage({ headers: { 'standard-server': 'event-stream' } })
+        const promise = peer.message(msg, handler)
+
+        await vi.waitFor(() => expect(handler).toHaveBeenCalled())
+        const request = handler.mock.calls[0]![0]
+        const iter = await request.resolveBody() as AsyncIterator<unknown>
+
+        const pending = iter.next()
+        await sleep(0) // let next() start waiting for a message
+        await iter.return?.()
+
+        await expect(pending).resolves.toEqual({ value: undefined, done: true })
+        await expect(iter.next()).resolves.toEqual({ value: undefined, done: true })
+        expect(send).toHaveBeenCalledWith({ id: '1', kind: 'stream/cancel' })
+
+        box.resolve(jsonResponse())
+        await promise
+      })
+
       it('asyncIterator error if receive cancel message', async () => {
         const { handler, box } = deferredHandler()
 
@@ -441,6 +463,27 @@ describe('serverPeer', () => {
         expect(cancelMsgs[0]!.id).toBe('1')
       })
 
+      it('settles a pending read when the handler stops reading octet-stream', async () => {
+        const { handler, box } = deferredHandler()
+
+        const msg = makeRequestMessage({ headers: { 'standard-server': 'octet-stream', 'content-type': 'application/octet-stream' } })
+        const promise = peer.message(msg, handler)
+
+        await vi.waitFor(() => expect(handler).toHaveBeenCalled())
+        const request = handler.mock.calls[0]![0]
+        const body = await request.resolveBody() as ReadableStream<Uint8Array>
+        const reader = body.getReader()
+
+        const pending = reader.read()
+        await reader.cancel()
+
+        await expect(pending).resolves.toEqual({ value: undefined, done: true })
+        expect(send).toHaveBeenCalledWith({ id: '1', kind: 'stream/cancel' })
+
+        box.resolve(jsonResponse())
+        await promise
+      })
+
       it('readableStream error if receive cancel message', async () => {
         const { handler, box } = deferredHandler()
 

packages/peer/src/server.ts

@@ -57,13 +57,19 @@ export class ServerPeer {
     const signal = controller.signal
 
     try {
-      const decoded = toStandardBody(message, async ({ kind }) => {
+      const decoded = toStandardBody(message, async ({ kind, error }) => {
         /**
          * The request body is finished (fully read, errored, or cancelled).
          * Drop the queues so late stream messages are ignored instead of
          * buffered forever.
          */
         const streamActive = state.eventStreamMessageQueue !== undefined || state.octetStreamMessageQueue !== undefined
+
+        if (kind === 'cancelled' && streamActive) {
+          state.eventStreamMessageQueue?.abort(error)
+          state.octetStreamMessageQueue?.abort(error)
+        }
+
         state.eventStreamMessageQueue = undefined
         state.octetStreamMessageQueue = undefined
 

packages/shared/src/iterator.test.ts

@@ -1,4 +1,5 @@
 import { AsyncIteratorClass, isAsyncIteratorObject } from './iterator'
+import { promiseWithResolvers } from './promise'
 
 beforeEach(() => {
   vi.clearAllMocks()
@@ -25,6 +26,12 @@ describe('asyncIteratorClass', () => {
   const cleanup = vi.fn()
   let iterator: AsyncGenerator
 
+  // how the underlying call settles after the consumer has stopped reading
+  const lateSettles: [string, (deferred: ReturnType<typeof promiseWithResolvers<unknown>>) => void][] = [
+    ['resolves', deferred => deferred.resolve({ done: false, value: 42 })],
+    ['rejects', deferred => deferred.reject(new Error('Late'))],
+  ]
+
   beforeEach(() => {
     next.mockReset()
     cleanup.mockReset()
@@ -126,6 +133,20 @@ describe('asyncIteratorClass', () => {
       expect(next).toHaveBeenCalledTimes(0)
     })
 
+    it.each(lateSettles)('should end a waiting next() when its call %s after return()', async (_, settle) => {
+      const deferred = promiseWithResolvers<unknown>()
+      next.mockReturnValueOnce(deferred.promise)
+
+      const pending = iterator.next()
+      await vi.waitFor(() => expect(next).toHaveBeenCalledTimes(1))
+
+      await iterator.return(undefined)
+      settle(deferred)
+
+      await expect(pending).resolves.toEqual({ done: true, value: undefined })
+      expect(cleanup).toHaveBeenCalledWith({ kind: 'cancelled' })
+    })
+
     it('should call cleanup({ kind: cancelled })', async () => {
       await Promise.all([
         iterator.return('done'),
@@ -145,6 +166,22 @@ describe('asyncIteratorClass', () => {
       expect(next).toHaveBeenCalledTimes(0)
     })
 
+    it.each(lateSettles)('should end a waiting next() when its call %s after throw()', async (_, settle) => {
+      const deferred = promiseWithResolvers<unknown>()
+      next.mockReturnValueOnce(deferred.promise)
+
+      const pending = iterator.next()
+      await vi.waitFor(() => expect(next).toHaveBeenCalledTimes(1))
+
+      const error = new Error('Forced error')
+      await expect(iterator.throw(error)).rejects.toBe(error)
+      settle(deferred)
+
+      // like a native async generator, only the caller of throw() sees the error
+      await expect(pending).resolves.toEqual({ done: true, value: undefined })
+      expect(cleanup).toHaveBeenCalledWith({ kind: 'cancelled', error })
+    })
+
     it('should call cleanup({ kind: cancelled, reason })', async () => {
       const error = new Error('Forced error')
 

packages/shared/src/iterator.ts

@@ -37,13 +37,23 @@ export class AsyncIteratorClass<T, TReturn = unknown, TNext = unknown> implement
       try {
         const result = await next()
 
+        // the consumer stopped reading (e.g. `return()`) while this call was waiting
+        if (this.isDone) {
+          return { done: true, value: undefined as any }
+        }
+
         if (result.done) {
           this.isDone = true
         }
 
         return result
       }
       catch (error) {
+        // the consumer stopped reading (e.g. `return()`) while this call was waiting
+        if (this.isDone) {
+          return { done: true, value: undefined as any }
+        }
+
         errorRef = { value: error }
         this.isDone = true