Fix generators not respecting headers

src/handler.ts

@@ -140,38 +140,31 @@ const handleStream = (
 	set?: Context['set'],
 	res?: HttpResponse
 ): ElysiaNodeResponse => {
-	if (!set)
-		set = {
-			status: 200,
-			headers: {
-				'transfer-encoding': 'chunked',
-				'content-type': 'text/event-stream;charset=utf8'
-			}
-		}
-	else {
-		set.headers['transfer-encoding'] = 'chunked'
-		set.headers['content-type'] = 'text/event-stream;charset=utf8'
-	}
-
-	if (res) res.writeHead(set.status as number, set.headers)
-
-	return [handleStreamResponse(generator, set, res), set as any]
-}
-
-export const handleStreamResponse = (
-	generator: Generator | AsyncGenerator,
-	set?: Context['set'],
-	res?: HttpResponse
-) => {
 	const readable = new Readable({
 		read() {}
 	})
-
-	if (res) readable.pipe(res)
+	if (res) {
+		readable.pipe(res)
+	}
 	;(async () => {
 		let init = generator.next()
 		if (init instanceof Promise) init = await init
-
+		if (!set)
+			set = {
+				status: 200,
+				headers: {
+					'transfer-encoding': 'chunked',
+					'content-type': 'text/event-stream;charset=utf8'
+				}
+			}
+		else {
+			set.headers['transfer-encoding'] = 'chunked'
+			if (!set.headers['content-type']) {
+				set.headers['content-type'] = 'text/event-stream;charset=utf8'
+			}
+		}
+
+		if (res) res.writeHead(set.status as number, set.headers)
 		if (init.done) {
 			if (set) return mapResponse(init.value, set, res)
 			return mapCompactResponse(init.value, res)
@@ -216,9 +209,13 @@ export const handleStreamResponse = (
 		readable.push(null)
 	})()
 
-	return readable
+	return [readable, set as any]
+
+
 }
 
+
+
 export async function* streamResponse(response: Response) {
 	const body = response.body
 

src/index.ts

@@ -410,7 +410,7 @@ export const node = () => {
 										// @ts-expect-error private property
 										app._handle
 									).listen(
-										typeof options === 'number'
+										typeof options === 'number' || typeof options === "string"
 											? options
 											: {
 													...options,

test/core/generators.test.ts

@@ -0,0 +1,65 @@
+import Elysia from 'elysia'
+import { inject } from 'light-my-request'
+import { describe, it, expect } from 'vitest'
+
+import node from '../../src'
+
+function delay(n: number) {
+    return new Promise((resolve)=>{
+        setTimeout(()=>{
+            resolve(null)
+        }, n)
+    })
+}
+
+const app =  new Elysia({ adapter: node() })
+    .get("/what", ()=>"feafe")
+    .get('/', function* gen({set}) {
+        set.headers["content-type"] = "application/json"
+        set.headers["content-disposition"] = 'attachment; filename="test.json"'
+        yield JSON.stringify({x: "hi"})
+    })
+    .get('/async-simple', async function* gen({set}) {
+        set.headers["content-type"] = "application/json"
+        set.headers["content-disposition"] = 'attachment; filename="test.json"'
+        await delay(100)
+        yield JSON.stringify({x: "hi"})
+    })
+    .get('/async-hard', async function* gen({set}) {
+        await delay(100)
+        set.headers["content-type"] = "application/json"
+        set.headers["content-disposition"] = 'attachment; filename="test.json"'
+        await delay(100)
+        yield JSON.stringify({x: "hi"})
+    })
+    .compile()
+
+
+// @ts-expect-error
+const handle = app._handle!
+
+describe("Async generators", ()=>{
+
+    it("handle /", async ()=>{
+        await inject(handle, { path: "/"}, (error, res)=>{
+            console.log(res)
+            expect(res?.headers["content-type"]).toBe("application/json")
+            expect(res?.headers["content-disposition"]).toBe('attachment; filename="test.json"')
+        })
+    })
+
+    it("handle /async-simple", async ()=>{
+        await inject(handle, { path: "/async-simple"}, (error, res)=>{
+            expect(true).toBe(false)
+            expect(res?.headers["content-type"]).toBe("application/json")
+            expect(res?.headers["content-disposition"]).toBe('attachment; filename="test.json"')
+        })
+    })
+
+    it("handle /async-hard", async ()=>{
+        await inject(handle, { path: "/async-hard"}, (error, res)=>{
+            expect(res?.headers["content-type"]).toBe("application/json")
+            expect(res?.headers["content-disposition"]).toBe('attachment; filename="test.json"')
+        })
+    })
+})