Merge pull request #1595 from abhishekg999/ahh/fix-trace-afterhandle-stream-reference

Fix: stream reference should point to teed value

Author: SaltyAom

src/compose.ts

@@ -1797,6 +1797,9 @@ export const composeHandler = ({
 					(maybeAsync ? '' : `(async()=>{`) +
 					`const stream=await tee(r,3)\n` +
 					`r=stream[0]\n` +
+					(hooks.afterHandle?.length
+						? `c.response=c.responseValue=r\n`
+						: '') +
 					`const listener=stream[1]\n` +
 					(hasTrace || hooks.afterResponse?.length
 						? `afterHandlerStreamListener=stream[2]\n`

test/response/stream.test.ts

@@ -1,4 +1,4 @@
-import { describe, it, expect } from 'bun:test'
+import { describe, expect, it } from 'bun:test'
 import { req } from '../utils'
 
 import { Elysia, sse } from '../../src'
@@ -560,4 +560,39 @@ describe('Stream', () => {
 		expect(result).toEqual(['Elysia', 'Eden'].map((x) => `data: ${x}\n\n`))
 		expect(response.headers.get('content-type')).toBe('text/event-stream')
 	})
+
+	it('handle sse with plugin global hooks and trace', async () => {
+		const PluginA = () =>
+			new Elysia({ name: 'PluginA' })
+				.onBeforeHandle(() => {})
+				.onAfterHandle(() => {})
+				.onParse(() => {})
+				.onTransform(() => {})
+				.onError(() => {})
+				.onAfterResponse(() => {})
+				.onStart(() => {})
+				.onStop(() => {})
+				.onRequest(() => {})
+				.trace(() => {})
+				.as('global')
+
+		const app = new Elysia().use(PluginA()).get('/sse', async function* () {
+			yield sse({ event: 'message', data: { meow: '1' } })
+			yield sse({ event: 'message', data: { meow: '2' } })
+			yield sse({ event: 'message', data: { meow: '3' } })
+		})
+
+		const response = await app.handle(req('/sse'))
+		expect(response.headers.get('content-type')).toBe('text/event-stream')
+
+		const result = []
+
+		for await (const chunk of streamResponse(response)) result.push(chunk)
+		expect(result).toHaveLength(3)
+		expect(result).toEqual([
+			'event: message\ndata: {"meow":"1"}\n\n',
+			'event: message\ndata: {"meow":"2"}\n\n',
+			'event: message\ndata: {"meow":"3"}\n\n'
+		])
+	})
 })