Add new eventStream API for async generators

src/server/event-stream.ts

@@ -1,13 +1,13 @@
-interface SendFunctionArgs {
+interface SendFunctionArgs<T = string> {
 	/**
 	 * @default "message"
 	 */
 	event?: string;
-	data: string;
+	data: T;
 }
 
-interface SendFunction {
-	(args: SendFunctionArgs): void;
+interface SendFunction<T = string> {
+	(args: SendFunctionArgs<T>): void;
 }
 
 interface CleanupFunction {
@@ -18,8 +18,17 @@ interface AbortFunction {
 	(): void;
 }
 
-interface InitFunction {
-	(send: SendFunction, abort: AbortFunction): CleanupFunction;
+interface InitFunction<T = string> {
+	(send: SendFunction<T>, abort: AbortFunction): CleanupFunction;
+}
+
+interface HandleFunction<T = string> {
+	(send: SendFunction<T>, abort: AbortFunction): Promise<void> | void;
+}
+
+interface Options<T = string> {
+	cleanup: CleanupFunction;
+	handle: HandleFunction<T>;
 }
 
 /**
@@ -28,21 +37,33 @@ interface InitFunction {
  * @param init The function that will be called to initialize the stream, here you can subscribe to your events
  * @returns A Response object that can be returned from a loader
  */
-export function eventStream(
+export function eventStream<T = string>(
 	signal: AbortSignal,
-	init: InitFunction,
-	options: ResponseInit = {},
+	init: InitFunction<T>,
+	responseInit?: ResponseInit,
+): Response;
+export function eventStream<T = string>(
+	signal: AbortSignal,
+	options: Options<T>,
+	responseInit?: ResponseInit,
+): Response;
+export function eventStream<T = string>(
+	signal: AbortSignal,
+	initOrOptions: InitFunction<T> | Options<T>,
+	responseInit: ResponseInit = {},
 ) {
 	let stream = new ReadableStream({
 		start(controller) {
 			let encoder = new TextEncoder();
 
-			function send({ event = "message", data }: SendFunctionArgs) {
-				controller.enqueue(encoder.encode(`event: ${event}\n`));
-				controller.enqueue(encoder.encode(`data: ${data}\n\n`));
-			}
+			let cleanup =
+				typeof initOrOptions === "function"
+					? initOrOptions(send, close)
+					: initOrOptions.cleanup;
 
-			let cleanup = init(send, close);
+			if (typeof initOrOptions !== "function") {
+				initOrOptions.handle(send, close);
+			}
 
 			let closed = false;
 
@@ -57,10 +78,15 @@ export function eventStream(
 			signal.addEventListener("abort", close);
 
 			if (signal.aborted) return close();
+
+			async function send({ event = "message", data }: SendFunctionArgs<T>) {
+				controller.enqueue(encoder.encode(`event: ${event}\n`));
+				controller.enqueue(encoder.encode(`data: ${data}\n\n`));
+			}
 		},
 	});
 
-	let headers = new Headers(options.headers);
+	let headers = new Headers(responseInit.headers);
 
 	if (headers.has("Content-Type")) {
 		console.warn("Overriding Content-Type header to `text/event-stream`");
@@ -78,5 +104,5 @@ export function eventStream(
 	headers.set("Cache-Control", "no-cache");
 	headers.set("Connection", "keep-alive");
 
-	return new Response(stream, { headers });
+	return new Response(stream, { ...responseInit, headers });
 }