1import type * as lambda from "aws-lambda";
2import * as Context from "effect/Context";
3import type * as Effect from "effect/Effect";
4import type { Api } from "./Api.ts";
5
6/**
7 * The event shape a WebSocket route handler receives — the API Gateway v2
8 * WebSocket proxy event (`requestContext.connectionId`, `routeKey`,
9 * `eventType`, plus `body` for `MESSAGE` events).
10 */
11export type WebSocketEvent = lambda.APIGatewayProxyWebsocketEventV2;
12
13/**
14 * The result a WebSocket route handler may return. For `$connect` routes a
15 * non-2xx `statusCode` rejects the connection; for other routes the result
16 * is only surfaced to the client when the route has a route response
17 * configured. Returning `void` means `{ statusCode: 200 }`.
18 */
19export type WebSocketResult =
20 | { statusCode: number; body?: string }
21 | undefined
22 | void;
23
24export interface WebSocketRouteProps {
25 /**
26 * The WebSocket route key this handler serves: `$connect`,
27 * `$disconnect`, `$default`, or a custom action name matched by the
28 * API's route selection expression.
29 */
30 routeKey: string;
31}
32
33export type WebSocketEventSourceService = <Req = never>(
34 api: Api,
35 props: WebSocketRouteProps,
36 handler: (
37 event: WebSocketEvent,
38 ) => Effect.Effect<WebSocketResult, never, Req>,
39) => Effect.Effect<void, never, never>;
40
41/**
42 * Event source connecting a WebSocket {@link Api} route to the hosting
43 * Lambda function.
44 *
45 * At deploy time the Lambda implementation (`Lambda.WebSocketEventSource`)
46 * materializes the `Integration`, `Route`, and invoke `Permission` for the
47 * route; at runtime it dispatches matching WebSocket proxy events to the
48 * handler. Subscribe routes with {@link onWebSocketRoute} and provide
49 * `Lambda.WebSocketEventSource` on the hosting function.
50 *
51 * ### Handling WebSocket Routes
52 * **Example:** Echo server on a WEBSOCKET Api
53 * ```typescript
54 * export default MyFunction.make(
55 * { main: import.meta.url },
56 * Effect.gen(function* () {
57 * const api = yield* ApiGatewayV2.Api("WsApi", {
58 * protocolType: "WEBSOCKET",
59 * routeSelectionExpression: "$request.body.action",
60 * });
61 * const stage = yield* ApiGatewayV2.Stage("WsStage", {
62 * api,
63 * stageName: "prod",
64 * autoDeploy: true,
65 * });
66 * const connections = yield* ApiGatewayV2.ManageConnections(stage);
67 *
68 * yield* ApiGatewayV2.onWebSocketRoute(api, { routeKey: "$connect" }, () =>
69 * Effect.succeed({ statusCode: 200 }),
70 * );
71 *
72 * yield* ApiGatewayV2.onWebSocketRoute(api, { routeKey: "$default" }, (event) =>
73 * connections
74 * .postToConnection({
75 * ConnectionId: event.requestContext.connectionId,
76 * Data: `echo:${event.body ?? ""}`,
77 * })
78 * .pipe(
79 * Effect.asVoid,
80 * Effect.catchTag("GoneException", () => Effect.void),
81 * Effect.orDie,
82 * ),
83 * );
84 *
85 * return {};
86 * }).pipe(
87 * Effect.provide(
88 * Layer.mergeAll(
89 * Lambda.WebSocketEventSource,
90 * ApiGatewayV2.ManageConnectionsHttp,
91 * ),
92 * ),
93 * ),
94 * );
95 * ```
96 *
97 * @binding
98 */
99export class WebSocketEventSource extends Context.Service<
100 WebSocketEventSource,
101 WebSocketEventSourceService
102>()("AWS.ApiGatewayV2.WebSocketEventSource") {}
103
104/**
105 * Subscribe an Effect handler to a WebSocket route of an API Gateway v2
106 * WEBSOCKET {@link Api}.
107 *
108 * Provide `Lambda.WebSocketEventSource` on the hosting function to
109 * satisfy the requirement.
110 *
111 * @param api The WebSocket API.
112 * @param props The route key to serve.
113 * @param handler Invoked once per WebSocket event for the route.
114 *
115 * @example
116 * ```typescript
117 * yield* ApiGatewayV2.onWebSocketRoute(api, { routeKey: "$connect" }, () =>
118 * Effect.succeed({ statusCode: 200 }),
119 * );
120 * ```
121 */
122export function onWebSocketRoute<Req = never>(
123 api: Api,
124 props: WebSocketRouteProps,
125 handler: (
126 event: WebSocketEvent,
127 ) => Effect.Effect<WebSocketResult, never, Req>,
128): Effect.Effect<void, never, WebSocketEventSource> {
129 return WebSocketEventSource.use((source) => source(api, props, handler));
130}