From 4c49a14d2d75fffed0f13909fae9b80d95e05a40 Mon Sep 17 00:00:00 2001
From: Artem Zakharchenko <kettanaito@gmail.com>
Date: Tue, 1 Oct 2024 18:30:10 +0200
Subject: [PATCH] feat: support `EventMap` type argument

---
 src/browser/sse.ts                            | 49 +++++++++++++------
 .../{sse.test.ts => sse.client.send.test.ts}  |  0
 test/typings/sse.test-d.ts                    | 39 +++++++++++++++
 test/typings/vitest.config.mts                |  4 +-
 4 files changed, 75 insertions(+), 17 deletions(-)
 rename test/browser/sse-api/{sse.test.ts => sse.client.send.test.ts} (100%)
 create mode 100644 test/typings/sse.test-d.ts

diff --git a/src/browser/sse.ts b/src/browser/sse.ts
index 95400413c..7a1cbf3e4 100644
--- a/src/browser/sse.ts
+++ b/src/browser/sse.ts
@@ -7,21 +7,26 @@ import {
 } from '~/core/handlers/HttpHandler'
 import type { Path, PathParams } from '~/core/utils/matching/matchRequestUrl'
 
-export type ServerSentEventResolverExtras<Params extends PathParams> =
-  HttpRequestResolverExtras<Params> & {
-    client: ServerSentEventClient
-    server: ServerSentEventServer
-  }
+export type ServerSentEventResolverExtras<
+  EventMap extends Record<string, unknown>,
+  Params extends PathParams,
+> = HttpRequestResolverExtras<Params> & {
+  client: ServerSentEventClient<EventMap>
+  server: ServerSentEventServer
+}
 
-export type ServerSentEventResolver<Params extends PathParams> =
-  ResponseResolver<ServerSentEventResolverExtras<Params>, any, any>
+export type ServerSentEventResolver<
+  EventMap extends Record<string, unknown>,
+  Params extends PathParams,
+> = ResponseResolver<ServerSentEventResolverExtras<EventMap, Params>, any, any>
 
 export type ServerSentEventRequestHandler = <
+  EventMap extends Record<string, unknown>,
   Params extends PathParams<keyof Params> = PathParams,
   RequestPath extends Path = Path,
 >(
   path: RequestPath,
-  resolver: ServerSentEventResolver<Params>,
+  resolver: ServerSentEventResolver<EventMap, Params>,
 ) => HttpHandler
 
 /**
@@ -38,8 +43,10 @@ export const sse: ServerSentEventRequestHandler = (path, resolver) => {
   return new ServerSentEventHandler(path, resolver)
 }
 
-class ServerSentEventHandler extends HttpHandler {
-  constructor(path: Path, resolver: ServerSentEventResolver<any>) {
+class ServerSentEventHandler<
+  EventMap extends Record<string, unknown>,
+> extends HttpHandler {
+  constructor(path: Path, resolver: ServerSentEventResolver<EventMap, any>) {
     invariant(
       typeof EventSource !== 'undefined',
       'Failed to construct a Server-Sent Event handler for path "%s": your environment does not support the EventSource API',
@@ -49,7 +56,7 @@ class ServerSentEventHandler extends HttpHandler {
     super('GET', path, (info) => {
       const stream = new ReadableStream({
         start(controller) {
-          const client = new ServerSentEventClient({
+          const client = new ServerSentEventClient<EventMap>({
             controller,
           })
           const server = new ServerSentEventServer({
@@ -92,7 +99,7 @@ class ServerSentEventHandler extends HttpHandler {
 
 const kSend = Symbol('kSend')
 
-class ServerSentEventClient {
+class ServerSentEventClient<EventMap extends Record<string, unknown>> {
   private encoder: TextEncoder
   protected controller: ReadableStreamDefaultController
 
@@ -104,7 +111,19 @@ class ServerSentEventClient {
   /**
    * Sends the given payload to the underlying `EventSource`.
    */
-  public send(payload: { id?: string; event?: string; data: unknown }): void {
+  public send<EventType extends keyof EventMap & string>(
+    payload:
+      | {
+          id?: string
+          event: EventType
+          data: EventMap[EventType]
+        }
+      | {
+          id?: string
+          event?: never
+          data?: EventMap['message']
+        },
+  ): void {
     this[kSend]({
       id: payload.id,
       event: payload.event,
@@ -142,9 +161,9 @@ class ServerSentEventClient {
     this.controller.error()
   }
 
-  private [kSend](payload: {
+  private [kSend]<EventType extends keyof EventMap & string>(payload: {
     id?: string
-    event?: string
+    event?: EventType
     data: string
   }): void {
     const frames: Array<string> = []
diff --git a/test/browser/sse-api/sse.test.ts b/test/browser/sse-api/sse.client.send.test.ts
similarity index 100%
rename from test/browser/sse-api/sse.test.ts
rename to test/browser/sse-api/sse.client.send.test.ts
diff --git a/test/typings/sse.test-d.ts b/test/typings/sse.test-d.ts
new file mode 100644
index 000000000..59ce96238
--- /dev/null
+++ b/test/typings/sse.test-d.ts
@@ -0,0 +1,39 @@
+import { it } from 'vitest'
+import { sse } from 'msw/browser'
+
+it('supports custom event map type argument', () => {
+  sse<{ myevent: string }>('/stream', ({ client }) => {
+    client.send({
+      event: 'myevent',
+      data: 'hello',
+    })
+
+    client.send({
+      // @ts-expect-error Unknown event type "unknown".
+      event: 'unknown',
+      data: 'hello',
+    })
+  })
+})
+
+it('supports event map type argument for unnamed events', () => {
+  sse<{ message: number; custom: 'goodbye' }>('/stream', ({ client }) => {
+    client.send({
+      data: 123,
+    })
+    client.send({
+      // @ts-expect-error Unexpected data type for "message" event.
+      data: 'goodbye',
+    })
+
+    client.send({
+      event: 'custom',
+      data: 'goodbye',
+    })
+    client.send({
+      event: 'custom',
+      // @ts-expect-error Unexpected data type for "custom" event.
+      data: 'invalid',
+    })
+  })
+})
diff --git a/test/typings/vitest.config.mts b/test/typings/vitest.config.mts
index 4f879ab8e..b5db2bcc8 100644
--- a/test/typings/vitest.config.mts
+++ b/test/typings/vitest.config.mts
@@ -1,8 +1,8 @@
 import * as path from 'node:path'
+import * as fs from 'fs'
 import { defineConfig } from 'vitest/config'
-import tsPackageJson from 'typescript/package.json' assert { type: 'json' }
 import { invariant } from 'outvariant'
-import * as fs from 'fs'
+import tsPackageJson from 'typescript/package.json' assert { type: 'json' }
 
 const LIB_DIR = path.resolve(__dirname, '../../lib')