From 4667e3cb26833c7f3ed711c2e2f1dd967fb7c9c2 Mon Sep 17 00:00:00 2001
From: bogeychan <bogeychan@outlook.com>
Date: Wed, 17 Jul 2024 14:21:27 +0200
Subject: [PATCH] fix: generator set additional hooks header

---
 src/handler.ts               | 18 +++++++++---------
 test/response/stream.test.ts | 18 ++++++++++++++++++
 2 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/src/handler.ts b/src/handler.ts
index eec37c6a..2354092c 100644
--- a/src/handler.ts
+++ b/src/handler.ts
@@ -125,6 +125,12 @@ const handleStream = (
 ) => {
 	let end = false
 
+	set ??= { headers: {} }
+
+	// Manually set transfer-encoding for direct response, eg. app.handle, eden
+	set.headers['transfer-encoding'] = 'chunked'
+	set.headers['content-type'] = 'text/event-stream; charset=utf-8'
+
 	return new Response(
 		new ReadableStream({
 			async start(controller) {
@@ -158,14 +164,7 @@ const handleStream = (
 				}
 			}
 		}),
-		{
-			...(set as ResponseInit),
-			headers: {
-				// Manually set transfer-encoding for direct response, eg. app.handle, eden
-				'transfer-encoding': 'chunked',
-				'content-type': 'text/event-stream; charset=utf-8'
-			}
-		}
+		set as ResponseInit
 	)
 }
 
@@ -1055,7 +1054,8 @@ export const errorToResponse = (error: Error, set?: Context['set']) =>
 			cause: error?.cause
 		}),
 		{
-			status: set?.status !== 200 ? (set?.status as number) ?? 500 : 500,
+			status:
+				set?.status !== 200 ? ((set?.status as number) ?? 500) : 500,
 			headers: set?.headers
 		}
 	)
diff --git a/test/response/stream.test.ts b/test/response/stream.test.ts
index 08002327..09a03e6b 100644
--- a/test/response/stream.test.ts
+++ b/test/response/stream.test.ts
@@ -44,6 +44,24 @@ describe('Stream', () => {
 		expect(response).toBe('abc')
 	})
 
+	it('handle stream with hook header', async () => {
+		const app = new Elysia()
+			.onRequest(({ set }) => {
+				set.headers['my'] = 'header'
+			})
+			.get('/', async function* () {
+				yield 'a'
+			})
+
+		const response = await app.handle(req('/'))
+
+		expect(response.headers.get('my')).toBe('header')
+		expect(response.headers.get('transfer-encoding')).toBe('chunked')
+		expect(response.headers.get('content-type')).toBe(
+			'text/event-stream; charset=utf-8'
+		)
+	})
+
 	it('stop stream on canceled request', async () => {
 		const expected = ['a', 'b']