1import * as AWS from "@/AWS";
2import * as Test from "@/Test/Alchemy";
3import { expect } from "alchemy-test";
4import * as Data from "effect/Data";
5import * as Duration from "effect/Duration";
6import * as Effect from "effect/Effect";
7import * as Schedule from "effect/Schedule";
8import * as HttpClient from "effect/http/HttpClient";
9import WebSocket from "ws";
10import WebSocketTestFunctionLive, { WebSocketTestFunction } from "./ws-handler";
11
12const { test } = Test.make({ providers: AWS.providers() });
13
14class WebSocketFailure extends Data.TaggedError("WebSocketFailure")<{
15 readonly reason: string;
16}> {}
17
18class FixtureNotConsistent extends Data.TaggedError("FixtureNotConsistent") {}
19
20/**
21 * Connect a wss client, send one frame, and resolve with the first frame
22 * received back. Encapsulates the socket in `Effect.callback` so fiber
23 * interruption (the outer `Effect.timeout`) closes the socket.
24 *
25 * An `{"message": "Internal server error"}` frame is API Gateway reporting
26 * that the route's Lambda invocation failed — right after a fresh deploy
27 * that is IAM/binding propagation, so it fails as a retryable
28 * `WebSocketFailure` rather than resolving.
29 */
30const wsSendAndReceive = (url: string, message: string) =>
31 Effect.callback<string, WebSocketFailure>((resume, signal) => {
32 const socket = new WebSocket(url);
33 let settled = false;
34 const settle = (effect: Effect.Effect<string, WebSocketFailure>) => {
35 if (settled) return;
36 settled = true;
37 try {
38 socket.close();
39 } catch {
40 // already closed
41 }
42 resume(effect);
43 };
44 signal.addEventListener("abort", () =>
45 settle(Effect.fail(new WebSocketFailure({ reason: "interrupted" }))),
46 );
47 socket.on("open", () => socket.send(message));
48 socket.on("message", (data) => {
49 const frame = String(data);
50 if (frame.includes('"Internal server error"')) {
51 settle(
52 Effect.fail(
53 new WebSocketFailure({
54 reason: `route invocation failed: ${frame}`,
55 }),
56 ),
57 );
58 } else {
59 settle(Effect.succeed(frame));
60 }
61 });
62 socket.on("error", (error) =>
63 settle(Effect.fail(new WebSocketFailure({ reason: String(error) }))),
64 );
65 socket.on("close", (code) =>
66 settle(
67 Effect.fail(
68 new WebSocketFailure({ reason: `closed before echo (${code})` }),
69 ),
70 ),
71 );
72 });
73
74test.provider(
75 "WebSocket API echoes via ManageConnections",
76 (stack) =>
77 Effect.gen(function* () {
78 yield* stack.destroy();
79
80 const fn = yield* stack.deploy(
81 Effect.gen(function* () {
82 return yield* WebSocketTestFunction.pipe(
83 Effect.provide(WebSocketTestFunctionLive),
84 );
85 }),
86 );
87 expect(fn.functionUrl).toBeTruthy();
88
89 // Discover the wss:// URL through the fixture's function URL.
90 // Two eventual-consistency windows to ride out: function-URL cold
91 // start (non-200s) and Lambda env-config propagation — an instance
92 // initialized before the binding env update serves `{}` until it is
93 // recycled, so poll until `wsUrl` is actually present.
94 const info = yield* HttpClient.get(fn.functionUrl!).pipe(
95 Effect.flatMap((response) =>
96 response.status === 200
97 ? response.json
98 : Effect.fail(new Error(`fixture returned ${response.status}`)),
99 ),
100 Effect.flatMap((body) =>
101 typeof (body as { wsUrl?: string }).wsUrl === "string"
102 ? Effect.succeed(body as { wsUrl: string; callbackUrl: string })
103 : Effect.fail(new FixtureNotConsistent()),
104 ),
105 Effect.retry({
106 schedule: Schedule.max([
107 Schedule.exponential("500 millis").pipe(
108 Schedule.modifyDelay(({ duration }) =>
109 Effect.succeed(
110 Duration.isGreaterThan(duration, Duration.seconds(3))
111 ? Duration.seconds(3)
112 : duration,
113 ),
114 ),
115 ),
116 Schedule.recurs(30),
117 ]),
118 }),
119 );
120 const { wsUrl, callbackUrl } = info;
121 expect(wsUrl).toMatch(/^wss:\/\//);
122 expect(callbackUrl).toMatch(/^https:\/\//);
123
124 // Connect, send, and expect the Lambda to push the echo back over
125 // the socket. The first connections retry through route/permission/
126 // IAM propagation (failed invocations surface as retryable
127 // "Internal server error" frames).
128 const echoed = yield* wsSendAndReceive(wsUrl, "hello-websocket").pipe(
129 Effect.timeout(Duration.seconds(15)),
130 Effect.retry({
131 schedule: Schedule.max([
132 Schedule.exponential("1 second").pipe(
133 Schedule.modifyDelay(({ duration }) =>
134 Effect.succeed(
135 Duration.isGreaterThan(duration, Duration.seconds(5))
136 ? Duration.seconds(5)
137 : duration,
138 ),
139 ),
140 ),
141 Schedule.recurs(10),
142 ]),
143 }),
144 );
145 expect(echoed).toBe("echo:hello-websocket");
146
147 // A second frame on a fresh connection round-trips too (no
148 // propagation retry needed once the first succeeded).
149 const second = yield* wsSendAndReceive(wsUrl, "again").pipe(
150 Effect.timeout(Duration.seconds(15)),
151 Effect.retry({
152 schedule: Schedule.max([
153 Schedule.exponential("1 second"),
154 Schedule.recurs(3),
155 ]),
156 }),
157 );
158 expect(second).toBe("echo:again");
159
160 yield* stack.destroy();
161 }),
162 {
163 tags: [
164 "provider:aws",
165 "provider:aws:apigatewayv2",
166 "provider:aws:lambda",
167 "live",
168 ],
169 timeout: 600_000,
170 },
171);