1import * as AWS from "@/AWS";
2import * as Test from "@/Test/Alchemy";
3import * as Core from "@/Test/Core";
4import { describe, expect } from "alchemy-test";
5import * as Data from "effect/Data";
6import * as Effect from "effect/Effect";
7import * as Schedule from "effect/Schedule";
8import * as HttpClient from "effect/http/HttpClient";
9import * as HttpClientRequest from "effect/http/HttpClientRequest";
10import BedrockLanguageModelFunctionLive, {
11 BedrockLanguageModelFunction,
12} from "./language-model-handler";
13
14const testOptions = { providers: AWS.providers() };
15const { test, beforeAll, afterAll } = Test.make(testOptions);
16const sharedStack = Core.scratchStack(testOptions, "BedrockLanguageModel");
17
18// Lambda function URL cold-start (DNS, IAM propagation, init) can take well
19// over 60s on a fresh deploy under parallel-suite load. Budget ~150s of
20// readiness polling.
21const readinessPolicy = Schedule.max([
22 Schedule.fixed("2 seconds"),
23 Schedule.recurs(75),
24]);
25
26let baseUrl: string;
27
28class TransientUpstream extends Data.TaggedError("TransientUpstream")<{
29 readonly status: number;
30 readonly body: string;
31}> {}
32
33// The shared Lambda fixture occasionally answers a transient 5xx under
34// parallel load (cold re-init, IAM propagation, Bedrock throttling). Retry
35// 5xx only; a genuine 4xx/assertion failure surfaces immediately.
36const send = (request: HttpClientRequest.HttpClientRequest) =>
37 HttpClient.execute(request).pipe(
38 Effect.flatMap((response) =>
39 response.status >= 500
40 ? response.text.pipe(
41 Effect.flatMap((body) =>
42 Effect.fail(
43 new TransientUpstream({ status: response.status, body }),
44 ),
45 ),
46 )
47 : Effect.succeed(response),
48 ),
49 Effect.retry({
50 while: (e) => e._tag === "TransientUpstream",
51 schedule: Schedule.max([
52 Schedule.exponential("1 second"),
53 Schedule.recurs(5),
54 ]),
55 }),
56 );
57
58interface StreamPart {
59 type: string;
60 id?: string;
61 name?: string;
62 delta?: string;
63 reason?: string;
64 usage?: {
65 inputTokens?: { total?: number };
66 outputTokens?: { total?: number };
67 };
68}
69
70const parseSse = (sse: string): ReadonlyArray<StreamPart> =>
71 sse
72 .split("\n\n")
73 .map((frame) => frame.replace(/^data:\s*/, "").trim())
74 .filter((line) => line.length > 0)
75 .map((line) => JSON.parse(line) as StreamPart);
76
77describe(
78 "Bedrock LanguageModel",
79 {
80 tags: [
81 "provider:aws",
82 "provider:aws:bedrock",
83 "provider:aws:lambda",
84 "live",
85 ],
86 },
87 () => {
88 beforeAll(
89 Effect.gen(function* () {
90 yield* Effect.logInfo(
91 "Bedrock LanguageModel setup: destroying previous resources",
92 );
93 yield* sharedStack.destroy();
94
95 yield* Effect.logInfo("Bedrock LanguageModel setup: deploying fixture");
96 const { functionUrl } = yield* sharedStack.deploy(
97 Effect.gen(function* () {
98 return yield* BedrockLanguageModelFunction;
99 }).pipe(Effect.provide(BedrockLanguageModelFunctionLive)),
100 );
101
102 expect(functionUrl).toBeTruthy();
103 baseUrl = functionUrl!.replace(/\/+$/, "");
104 const readinessUrl = `${baseUrl}/ping`;
105
106 yield* Effect.logInfo(
107 `Bedrock LanguageModel setup: probing readiness at ${readinessUrl}`,
108 );
109 yield* HttpClient.get(readinessUrl).pipe(
110 Effect.flatMap((response) =>
111 response.status === 200
112 ? Effect.succeed(response)
113 : Effect.fail(
114 new Error(`Function not ready: ${response.status}`),
115 ),
116 ),
117 Effect.tapError((error) =>
118 Effect.logWarning(
119 `Bedrock LanguageModel setup: fixture not ready yet (${String(error)})`,
120 ),
121 ),
122 Effect.retry({ schedule: readinessPolicy }),
123 );
124 }),
125 { timeout: 240_000 },
126 );
127
128 afterAll.skipIf(!!process.env.NO_DESTROY)(sharedStack.destroy(), {
129 timeout: 120_000,
130 });
131
132 test.provider(
133 "generateText answers with text, usage, and a terminal finish reason",
134 (_stack) =>
135 Effect.gen(function* () {
136 const response = (yield* send(
137 HttpClientRequest.get(
138 `${baseUrl}/generate?prompt=${encodeURIComponent("Say pong.")}`,
139 ),
140 ).pipe(Effect.flatMap((r) => r.json))) as {
141 text: string;
142 finishReason: string;
143 usage: { inputTokens: number; outputTokens: number };
144 };
145
146 expect(typeof response.text).toBe("string");
147 expect(response.text.length).toBeGreaterThan(0);
148 expect(response.usage.inputTokens).toBeGreaterThan(0);
149 expect(response.usage.outputTokens).toBeGreaterThan(0);
150 // A normal Converse completion (`end_turn`) maps to `stop`.
151 expect(["stop", "length"]).toContain(response.finishReason);
152 }),
153 { timeout: 120_000 },
154 );
155
156 test.provider(
157 "withModelParameters overrides maxTokens at runtime",
158 (_stack) =>
159 Effect.gen(function* () {
160 // The binding's default is maxTokens: 1024; the route clamps to 8
161 // via withModelParameters. Truncation proves the override reached
162 // Bedrock: the model runs out of budget mid-answer.
163 const response = (yield* send(
164 HttpClientRequest.get(
165 `${baseUrl}/generate-short?prompt=${encodeURIComponent(
166 "Write a detailed multi-paragraph essay about the history of infrastructure as code.",
167 )}`,
168 ),
169 ).pipe(Effect.flatMap((r) => r.json))) as {
170 finishReason: string;
171 outputTokens: number;
172 };
173
174 expect(response.finishReason).toBe("length");
175 expect(response.outputTokens).toBeLessThanOrEqual(8);
176 }),
177 { timeout: 120_000 },
178 );
179
180 test.provider(
181 "withModelParameters routes a call to another bound model",
182 (_stack) =>
183 Effect.gen(function* () {
184 // The layer binds [nova-micro, nova-lite]; this route overrides
185 // modelId to nova-lite. A 200 with text proves the per-call model
186 // selection and the multi-model IAM grant both work.
187 const response = (yield* send(
188 HttpClientRequest.get(
189 `${baseUrl}/generate-lite?prompt=${encodeURIComponent("Say pong.")}`,
190 ),
191 ).pipe(Effect.flatMap((r) => r.json))) as {
192 text: string;
193 finishReason: string;
194 };
195
196 expect(response.text.length).toBeGreaterThan(0);
197 expect(["stop", "length"]).toContain(response.finishReason);
198 }),
199 { timeout: 120_000 },
200 );
201
202 test.provider(
203 "streamText emits ordered parts: text-start → text-delta+ → text-end → finish",
204 (_stack) =>
205 Effect.gen(function* () {
206 const sse = yield* send(
207 HttpClientRequest.get(
208 `${baseUrl}/stream?prompt=${encodeURIComponent(
209 "Write a short paragraph (around 80 words) about why TypeScript developers might enjoy Effect TS.",
210 )}`,
211 ),
212 ).pipe(Effect.flatMap((r) => r.text));
213 const parts = parseSse(sse);
214
215 const startIdx = parts.findIndex((p) => p.type === "text-start");
216 const firstDeltaIdx = parts.findIndex((p) => p.type === "text-delta");
217 const deltas = parts.filter((p) => p.type === "text-delta");
218 const lastDeltaIdx = parts.lastIndexOf(deltas[deltas.length - 1]!);
219 const endIdx = parts.findIndex((p) => p.type === "text-end");
220 const finishIdx = parts.findIndex((p) => p.type === "finish");
221
222 // Adapter invariant: a text block opens before its first delta and
223 // closes before the terminal finish part.
224 expect(startIdx).toBeGreaterThanOrEqual(0);
225 expect(firstDeltaIdx).toBeGreaterThan(startIdx);
226 expect(endIdx).toBeGreaterThan(lastDeltaIdx);
227 expect(finishIdx).toBe(parts.length - 1);
228
229 // Bedrock streams one text-delta per contentBlockDelta event — a
230 // long response must produce many (a single fused delta would mean
231 // the adapter buffered the stream).
232 expect(deltas.length).toBeGreaterThan(3);
233 const text = deltas.map((p) => p.delta ?? "").join("");
234 expect(text.length).toBeGreaterThan(20);
235
236 // The metadata event's usage must survive into the finish part.
237 const finish = parts[finishIdx]!;
238 expect(finish.usage?.inputTokens?.total).toBeGreaterThan(0);
239 expect(finish.usage?.outputTokens?.total).toBeGreaterThan(0);
240 expect(["stop", "length"]).toContain(finish.reason);
241 }),
242 { timeout: 120_000 },
243 );
244
245 test.provider(
246 "generateText invokes a tool and returns its result",
247 (_stack) =>
248 Effect.gen(function* () {
249 const response = (yield* send(
250 HttpClientRequest.get(
251 `${baseUrl}/tool?prompt=${encodeURIComponent(
252 "What's the weather in San Francisco?",
253 )}`,
254 ),
255 ).pipe(Effect.flatMap((r) => r.json))) as {
256 finishReason: string;
257 toolCalls: Array<{
258 id: string;
259 name: string;
260 params: { city: string };
261 }>;
262 toolResults: Array<{
263 id: string;
264 name: string;
265 result: { city: string; temperatureF: number; condition: string };
266 isFailure: boolean;
267 }>;
268 };
269
270 expect(response.toolCalls.length).toBeGreaterThan(0);
271 const call = response.toolCalls[0]!;
272 expect(call.name).toBe("get_weather");
273 expect(typeof call.params.city).toBe("string");
274 expect(call.params.city.toLowerCase()).toContain("san francisco");
275
276 expect(response.toolResults.length).toBeGreaterThan(0);
277 const result = response.toolResults[0]!;
278 expect(result.name).toBe("get_weather");
279 expect(result.isFailure).toBe(false);
280 expect(result.result.temperatureF).toBe(72);
281 expect(result.result.condition).toBe("sunny");
282 }),
283 { timeout: 120_000 },
284 );
285
286 test.provider(
287 "streamText emits tool-params parts whose deltas concatenate to the arguments",
288 (_stack) =>
289 Effect.gen(function* () {
290 const sse = yield* send(
291 HttpClientRequest.get(
292 `${baseUrl}/tool-stream?prompt=${encodeURIComponent(
293 "What's the weather in Portland?",
294 )}`,
295 ),
296 ).pipe(Effect.flatMap((r) => r.text));
297 const parts = parseSse(sse);
298
299 const starts = parts.filter((p) => p.type === "tool-params-start");
300 const ends = parts.filter((p) => p.type === "tool-params-end");
301 expect(starts.length).toBeGreaterThan(0);
302 expect(starts[0]?.name).toBe("get_weather");
303 // Adapter invariant: every contentBlockStart-opened tool block is
304 // closed by contentBlockStop (or finalize), matched by id.
305 expect(new Set(ends.map((p) => p.id))).toEqual(
306 new Set(starts.map((p) => p.id)),
307 );
308
309 const firstId = starts[0]!.id!;
310 const joined = parts
311 .filter((p) => p.type === "tool-params-delta" && p.id === firstId)
312 .map((p) => p.delta ?? "")
313 .join("");
314 const args = yield* Effect.try({
315 try: () => JSON.parse(joined) as { city?: string },
316 catch: (cause) =>
317 new Error(
318 `Invalid concatenated tool arguments ${JSON.stringify(joined)}: ${cause}`,
319 ),
320 });
321 expect(typeof args.city).toBe("string");
322 expect(args.city!.toLowerCase()).toContain("portland");
323
324 expect(parts[parts.length - 1]?.type).toBe("finish");
325 }),
326 { timeout: 120_000 },
327 );
328 },
329);