1import type lambda from "aws-lambda";
2import * as Effect from "effect/Effect";
3import * as Layer from "effect/Layer";
4
5import * as Namespace from "../../Namespace.ts";
6import * as Output from "../../Output.ts";
7import type { Api } from "../ApiGatewayV2/Api.ts";
8import { IntegrationResource } from "../ApiGatewayV2/Integration.ts";
9import { RouteResource } from "../ApiGatewayV2/Route.ts";
10import {
11 WebSocketEventSource as AGW2WebSocketEventSource,
12 type WebSocketEvent,
13 type WebSocketEventSourceService,
14 type WebSocketResult,
15 type WebSocketRouteProps,
16} from "../ApiGatewayV2/WebSocketEventSource.ts";
17import * as Lambda from "./Function.ts";
18import { Permission as LambdaPermission } from "./Permission.ts";
19
20export const isWebSocketEvent = (
21 event: any,
22): event is lambda.APIGatewayProxyWebsocketEventV2 =>
23 typeof event?.requestContext?.connectionId === "string" &&
24 typeof event?.requestContext?.routeKey === "string" &&
25 event?.requestContext?.messageDirection === "IN";
26
27/** `$connect` → `connect` for stable, readable logical ids. */
28const sanitizeRouteKey = (routeKey: string) =>
29 routeKey.replace(/[^A-Za-z0-9]+/g, "") || "route";
30
31/**
32 * Connects a WebSocket API route to the current Lambda function.
33 *
34 * At deploy time this layer materializes the `AWS_PROXY` Integration, the
35 * Route, and the API Gateway invoke Permission for each registered route
36 * key; at runtime it dispatches matching WebSocket proxy events to the
37 * registered handler.
38 * ### Handling WebSocket routes
39 * **Example:** Echo server
40 * ```typescript
41 * const connections = yield* AWS.ApiGatewayV2.ManageConnections(stage);
42 *
43 * yield* AWS.ApiGatewayV2.onWebSocketRoute(api, { routeKey: "$connect" }, () =>
44 * Effect.succeed({ statusCode: 200 }),
45 * );
46 * yield* AWS.ApiGatewayV2.onWebSocketRoute(api, { routeKey: "$default" }, (event) =>
47 * connections
48 * .postToConnection({
49 * ConnectionId: event.requestContext.connectionId,
50 * Data: `echo:${event.body ?? ""}`,
51 * })
52 * .pipe(
53 * Effect.asVoid,
54 * Effect.catchTag("GoneException", () => Effect.void),
55 * Effect.orDie,
56 * ),
57 * );
58 * ```
59 *
60 * @binding
61 */
62export const WebSocketEventSource = Layer.effect(
63 AGW2WebSocketEventSource,
64 Effect.gen(function* () {
65 // this layer can only be used in a Lambda Function
66 const host = yield* Lambda.Function;
67 const Integration = yield* IntegrationResource;
68 const Route = yield* RouteResource;
69 const Permission = yield* LambdaPermission;
70
71 return Effect.fn(function* <Req = never>(
72 api: Api,
73 props: WebSocketRouteProps,
74 handler: (
75 event: WebSocketEvent,
76 ) => Effect.Effect<WebSocketResult, never, Req>,
77 ) {
78 // Deploy-time: materialize the Integration + Route + invoke
79 // Permission for this route key. Skipped once running inside the
80 // deployed Function (the global guard), where the only work is
81 // registering the runtime dispatcher below. Namespaced under the
82 // host for stable logical identity.
83 if (!globalThis.__ALCHEMY_RUNTIME__) {
84 yield* Namespace.push(
85 host.LogicalId,
86 Effect.gen(function* () {
87 const key = sanitizeRouteKey(props.routeKey);
88
89 const integration = yield* Integration(
90 `${api.LogicalId}-${key}-Integration`,
91 {
92 apiId: api.apiId,
93 integrationType: "AWS_PROXY",
94 // WebSocket AWS_PROXY integrations require the full Lambda
95 // invocation URI. Region is derived from the function ARN
96 // so the binding needs no ambient environment.
97 integrationUri: Output.map(
98 host.functionArn,
99 (arn: string) =>
100 `arn:aws:apigateway:${arn.split(":")[3]}:lambda:path/2015-03-31/functions/${arn}/invocations`,
101 ),
102 integrationMethod: "POST",
103 },
104 );
105
106 yield* Route(`${api.LogicalId}-${key}-Route`, {
107 apiId: api.apiId,
108 routeKey: props.routeKey,
109 target: Output.map(
110 integration.integrationId,
111 (id: string) => `integrations/${id}`,
112 ),
113 });
114
115 yield* Permission(`${api.LogicalId}-${key}-Permission`, {
116 action: "lambda:InvokeFunction",
117 functionName: host.functionName,
118 principal: "apigateway.amazonaws.com",
119 sourceArn: Output.map(
120 Output.all(host.functionArn, api.apiId),
121 ([fnArn, apiId]: [string, string]) => {
122 const [, , , region, accountId] = fnArn.split(":");
123 return `arn:aws:execute-api:${region}:${accountId}:${apiId}/*`;
124 },
125 ),
126 });
127 }),
128 );
129 }
130
131 yield* host.listen(
132 Effect.gen(function* () {
133 return (event: any) => {
134 if (
135 isWebSocketEvent(event) &&
136 event.requestContext.routeKey === props.routeKey
137 ) {
138 return handler(event).pipe(
139 Effect.map((result) => result ?? { statusCode: 200 }),
140 Effect.orDie,
141 );
142 }
143 };
144 }),
145 );
146 }) as WebSocketEventSourceService;
147 }),
148);