From e84ef14196786d1cab2aa69d3efe081ad61e9eb2 Mon Sep 17 00:00:00 2001
From: KATT <alexander@n1s.se>
Date: Fri, 13 Oct 2023 22:52:41 +0200
Subject: [PATCH] add SSE support

---
 src/async/asyncErrors.ts                |  7 +++++
 src/async/deserializeAsync.ts           | 37 +++++++++++--------------
 src/async/handlers/tsonAsyncIterable.ts |  6 ++++
 src/async/sse.test.ts                   | 12 ++++++--
 4 files changed, 38 insertions(+), 24 deletions(-)

diff --git a/src/async/asyncErrors.ts b/src/async/asyncErrors.ts
index 1aaec2b8..08a66c4e 100644
--- a/src/async/asyncErrors.ts
+++ b/src/async/asyncErrors.ts
@@ -42,3 +42,10 @@ export class TsonStreamInterruptedError extends TsonError {
 		this.name = "TsonStreamInterruptedError";
 	}
 }
+
+export class TsonAbortError extends TsonError {
+	constructor(cause: unknown) {
+		super("Aborted", { cause });
+		this.name = "TsonAbortError";
+	}
+}
diff --git a/src/async/deserializeAsync.ts b/src/async/deserializeAsync.ts
index e450c8dc..0b79aafb 100644
--- a/src/async/deserializeAsync.ts
+++ b/src/async/deserializeAsync.ts
@@ -9,7 +9,7 @@ import {
 	TsonSerialized,
 	TsonTransformerSerializeDeserialize,
 } from "../sync/syncTypes.js";
-import { TsonStreamInterruptedError } from "./asyncErrors.js";
+import { TsonAbortError, TsonStreamInterruptedError } from "./asyncErrors.js";
 import {
 	TsonAsyncIndex,
 	TsonAsyncOptions,
@@ -18,7 +18,6 @@ import {
 } from "./asyncTypes.js";
 import {
 	createReadableStream,
-	mapIterable,
 	readableStreamToAsyncIterable,
 } from "./iterableUtils.js";
 import { TsonAsyncValueTuple } from "./serializeAsync.js";
@@ -42,9 +41,8 @@ type TsonParseAsync = <TValue>(
 	opts?: TsonParseAsyncOptions,
 ) => Promise<TValue>;
 
-type TsonDeserializeIterable = AsyncIterable<
-	TsonAsyncValueTuple | TsonSerialized
->;
+type TsonDeserializeIterableValue = TsonAsyncValueTuple | TsonSerialized;
+type TsonDeserializeIterable = AsyncIterable<TsonDeserializeIterableValue>;
 function createTsonDeserializer(opts: TsonAsyncOptions) {
 	const typeByKey: Record<string, AnyTsonTransformerSerializeDeserialize> = {};
 
@@ -116,7 +114,9 @@ function createTsonDeserializer(opts: TsonAsyncOptions) {
 					break;
 				}
 
-				const [index, result] = nextValue.value as TsonAsyncValueTuple;
+				const { value } = nextValue;
+
+				const [index, result] = value as TsonAsyncValueTuple;
 
 				const controller = cache.get(index);
 
@@ -269,32 +269,27 @@ export function createTsonParseEventSource(opts: TsonAsyncOptions) {
 			signal?: AbortSignal;
 		} = {},
 	) => {
-		const [stream, controller] = createReadableStream<string>();
+		const [stream, controller] =
+			createReadableStream<TsonDeserializeIterableValue>();
 		const eventSource = new EventSource(url);
 
+		const { signal } = parseOpts;
 		const onAbort = () => {
+			assert(signal);
 			eventSource.close();
-			controller.close();
-			parseOpts.signal?.removeEventListener("abort", onAbort);
+			controller.error(new TsonAbortError("Stream aborted by user"));
+
+			signal.removeEventListener("abort", onAbort);
 		};
 
-		parseOpts.signal?.addEventListener("abort", onAbort);
+		signal?.addEventListener("abort", onAbort);
 
 		eventSource.onmessage = (msg) => {
 			// eslint-disable-next-line @typescript-eslint/no-unsafe-argument
-			controller.enqueue(msg.data);
+			controller.enqueue(JSON.parse(msg.data));
 		};
 
-		const iterable = mapIterable(
-			readableStreamToAsyncIterable(stream),
-			(msg) => {
-				const parsed = JSON.parse(msg) as TsonAsyncValueTuple | TsonSerialized;
-
-				console.log({ parsed });
-				return parsed;
-			},
-		);
-
+		const iterable = readableStreamToAsyncIterable(stream);
 		return (await instance(iterable, parseOpts)) as TValue;
 	};
 }
diff --git a/src/async/handlers/tsonAsyncIterable.ts b/src/async/handlers/tsonAsyncIterable.ts
index 8a99cfc1..d6e56577 100644
--- a/src/async/handlers/tsonAsyncIterable.ts
+++ b/src/async/handlers/tsonAsyncIterable.ts
@@ -1,4 +1,5 @@
 import {
+	TsonAbortError,
 	TsonPromiseRejectionError,
 	TsonStreamInterruptedError,
 } from "../asyncErrors.js";
@@ -31,6 +32,11 @@ export const tsonAsyncIterable: TsonAsyncType<
 			while (((next = await opts.reader.read()), !next.done)) {
 				const { value } = next;
 				if (value instanceof TsonStreamInterruptedError) {
+					if (value.cause instanceof TsonAbortError) {
+						opts.close();
+						return;
+					}
+
 					throw value;
 				}
 
diff --git a/src/async/sse.test.ts b/src/async/sse.test.ts
index 51416f0b..cd7622ed 100644
--- a/src/async/sse.test.ts
+++ b/src/async/sse.test.ts
@@ -99,8 +99,14 @@ test("SSE response test", async () => {
 			}
 		}
 
-		expect(messages).toMatchInlineSnapshot();
+		expect(messages).toMatchInlineSnapshot(`
+			[
+			  0,
+			  1,
+			  2,
+			  3,
+			  4,
+			]
+		`);
 	}
 });
-
-test.todo("parse SSE response");