Fix generators, use text/event-stream protocol

package.json

@@ -113,6 +113,7 @@
 		"eslint": "^8.49.0",
 		"eslint-plugin-security": "^2.1.0",
 		"eslint-plugin-sonarjs": "^0.23.0",
+		"eventsource-parser": "^1.1.2",
 		"expect-type": "^0.16.0",
 		"memoirist": "^0.2.0",
 		"prettier": "^3.3.3",

src/handler.ts

@@ -117,18 +117,24 @@ export const serializeCookie = (cookies: Context['set']['cookie']) => {
 
 // 	return arr
 // }
-
 const handleStream = async (
 	generator: Generator | AsyncGenerator,
 	set?: Context['set'],
 	request?: Request
 ) => {
-	let init = generator.next()
-	if (init instanceof Promise) init = await init
-
-	if (init.done) {
-		if (set) return mapResponse(init.value, set, request)
-		return mapCompactResponse(init.value, request)
+	let init
+	try {
+		init = generator.next()
+		if (init instanceof Promise) init = await init
+
+		if (init?.done) {
+			if (set) return mapResponse(init.value, set, request)
+			return mapCompactResponse(init.value, request)
+		}
+	} catch (error) {
+		// TODO should call app.onError if set
+		if (set) return mapResponse(error, set, request)
+		return mapCompactResponse(error, request)
 	}
 
 	return new Response(
@@ -146,38 +152,29 @@ const handleStream = async (
 					}
 				})
 
-				if (init.value !== undefined && init.value !== null) {
-					if (
-						typeof init.value === "object"
+				if (init?.value !== undefined && init?.value !== null)
+					controller.enqueue(
+						Buffer.from(
+							`event: message\ndata: ${JSON.stringify(init.value)}\n\n`
+						)
 					)
-						try {
-							controller.enqueue(
-								Buffer.from(JSON.stringify(init.value))
-							)
-						} catch {
-							controller.enqueue(Buffer.from(init.value.toString()))
-						}
-					else controller.enqueue(Buffer.from(init.value.toString()))
-				}
 
-				for await (const chunk of generator) {
-					if (end) break
-					if (chunk === undefined || chunk === null) continue
+				try {
+					for await (const chunk of generator) {
+						if (end) break
+						if (chunk === undefined || chunk === null) continue
 
-					if (typeof chunk === 'object')
-						try {
-							controller.enqueue(
-								Buffer.from(JSON.stringify(chunk))
+						controller.enqueue(
+							Buffer.from(
+								`event: message\ndata: ${JSON.stringify(chunk)}\n\n`
 							)
-						} catch {
-							controller.enqueue(Buffer.from(chunk.toString()))
-						}
-					else controller.enqueue(Buffer.from(chunk.toString()))
-
-					// Wait for the next event loop
-					// Otherwise the data will be mixed up
-					await new Promise<void>((resolve) =>
-						setTimeout(() => resolve(), 0)
+						)
+					}
+				} catch (error: any) {
+					controller.enqueue(
+						Buffer.from(
+							`event: error\ndata: ${JSON.stringify(error.message || error.name || 'Error')}\n\n`
+						)
 					)
 				}
 

test/response/stream.test.ts

@@ -1,8 +1,20 @@
 import { describe, it, expect } from 'bun:test'
 import { req } from '../utils'
+import { createParser } from 'eventsource-parser'
 
 import { Elysia } from '../../src'
 
+function textEventStream(items: string[]) {
+	return items
+		.map((item) => `event: message\ndata: ${JSON.stringify(item)}\n\n`)
+		.join('')
+}
+
+function parseTextEventStreamItem(item: string) {
+	const data = item.split('data: ')[1].split('\n')[0]
+	return JSON.parse(data)
+}
+
 describe('Stream', () => {
 	it('handle stream', async () => {
 		const expected = ['a', 'b', 'c']
@@ -31,7 +43,9 @@ describe('Stream', () => {
 				reader.read().then(function pump({ done, value }): unknown {
 					if (done) return resolve(acc)
 
-					expect(value.toString()).toBe(expected.shift()!)
+					expect(parseTextEventStreamItem(value.toString())).toBe(
+						expected.shift()!
+					)
 
 					acc += value.toString()
 					return reader.read().then(pump)
@@ -41,7 +55,46 @@ describe('Stream', () => {
 			})
 
 		expect(expected).toHaveLength(0)
-		expect(response).toBe('abc')
+		expect(response).toBe(textEventStream(['a', 'b', 'c']))
+	})
+	it('handle errors after yield', async () => {
+		const app = new Elysia().get('/', async function* () {
+			yield 'a'
+			await Bun.sleep(10)
+
+			throw new Error('an error')
+		})
+
+		const response = await app.handle(req('/')).then((x) => x.text())
+
+		expect(response).toBe(
+			'event: message\ndata: "a"\n\nevent: error\ndata: "an error"\n\n'
+		)
+	})
+	// TODO
+	it.skip('handle errors before yield', async () => {
+		const app = new Elysia().get('/', async function* () {
+			throw new Error('an error')
+		})
+
+		const response = await app.handle(req('/')).then((x) => x.text())
+
+		expect(response).toInclude('an error')
+	})
+
+	it.todo('handle errors before yield with onError', async () => {
+		const expected = 'error expected'
+		const app = new Elysia()
+			.onError(({}) => {
+				return new Response(expected)
+			})
+			.get('/', async function* () {
+				throw new Error('an error')
+			})
+
+		const response = await app.handle(req('/')).then((x) => x.text())
+
+		expect(response).toBe(expected)
 	})
 
 	it('stop stream on canceled request', async () => {
@@ -79,9 +132,13 @@ describe('Stream', () => {
 				const { promise, resolve } = Promise.withResolvers()
 
 				reader.read().then(function pump({ done, value }): unknown {
-					if (done) return resolve(acc)
+					if (done) {
+						return resolve(acc)
+					}
 
-					expect(value.toString()).toBe(expected.shift()!)
+					expect(parseTextEventStreamItem(value.toString())).toBe(
+						expected.shift()!
+					)
 
 					acc += value.toString()
 					return reader.read().then(pump)
@@ -91,7 +148,7 @@ describe('Stream', () => {
 			})
 
 		expect(expected).toHaveLength(0)
-		expect(response).toBe('ab')
+		expect(response).toBe(textEventStream(['a', 'b']))
 	})
 
 	it('mutate set before yield is called', async () => {
@@ -111,6 +168,42 @@ describe('Stream', () => {
 			'http://saltyaom.com'
 		)
 	})
+	it('handle stream with objects', async () => {
+		const objects = [
+			{ message: 'hello' },
+			{ response: 'world' },
+			{ data: [1, 2, 3] },
+			{ result: [4, 5, 6] }
+		]
+		const app = new Elysia().get('/', async function* ({}) {
+			for (const obj of objects) {
+				yield obj
+			}
+		})
+
+		const body = await app.handle(req('/')).then((x) => x.body)
+
+		let events = [] as any[]
+		const parser = createParser((event) => {
+			events.push(event)
+		})
+		const { promise, resolve } = Promise.withResolvers()
+		const reader = body?.getReader()!
+
+		reader.read().then(function pump({ done, value }): unknown {
+			if (done) {
+				return resolve()
+			}
+			const text = value.toString()
+			parser.feed(text)
+			return reader.read().then(pump)
+		})
+		await promise
+
+		expect(events.map((x) => x.data)).toEqual(
+			objects.map((x) => JSON.stringify(x))
+		)
+	})
 
 	it('mutate set before yield is called', async () => {
 		const expected = ['a', 'b', 'c']
@@ -216,7 +309,9 @@ describe('Stream', () => {
 				reader.read().then(function pump({ done, value }): unknown {
 					if (done) return resolve()
 
-					expect(value.toString()).toBe(JSON.stringify(expected[i++]))
+					expect(parseTextEventStreamItem(value.toString())).toEqual(
+						expected[i++]
+					)
 
 					return reader.read().then(pump)
 				})