1import {
2 makeLanguageModel,
3 type LanguageModelOptions,
4} from "@/Neon/LanguageModel.ts";
5import { RuntimeContext } from "@/RuntimeContext.ts";
6import { expect, test } from "alchemy-test";
7import * as Effect from "effect/Effect";
8import * as Layer from "effect/Layer";
9import * as Redacted from "effect/Redacted";
10import * as Result from "effect/Result";
11import * as Schema from "effect/Schema";
12import * as Stream from "effect/Stream";
13import { LanguageModel, Tool, Toolkit } from "effect/ai";
14import * as HttpClient from "effect/http/HttpClient";
15import * as HttpClientResponse from "effect/http/HttpClientResponse";
16
17const secret = "nt_live_SECRET_SENTINEL";
18const Sum = Tool.make("sum", {
19 parameters: Schema.Struct({ a: Schema.Number, b: Schema.Number }),
20 success: Schema.Number,
21});
22const Tools = Toolkit.make(Sum);
23const tools = Tools.toLayer({ sum: ({ a, b }) => Effect.succeed(a + b) });
24const toolCall = {
25 id: "call_1",
26 type: "function",
27 function: { name: "sum", arguments: '{"a":2,"b":3}' },
28};
29const completion = (message: unknown, finish = "stop", usage?: unknown) => ({
30 choices: [{ message, finish_reason: finish }],
31 usage,
32});
33const chunk = (
34 delta: unknown,
35 finish: string | null = null,
36 usage?: unknown,
37) => ({
38 choices: [{ index: 0, delta, finish_reason: finish }],
39 usage,
40});
41const sse = (...events: Array<unknown>) =>
42 events
43 .map(
44 (value) =>
45 `data: ${typeof value === "string" ? value : JSON.stringify(value)}\n\n`,
46 )
47 .join("");
48
49const harness = (
50 respond: () => globalThis.Response,
51 parameters?: LanguageModelOptions["parameters"],
52) => {
53 const requests: Array<{
54 url: string;
55 authorization: string | undefined;
56 body: unknown;
57 }> = [];
58 const http = HttpClient.make((request) =>
59 Effect.sync(() => {
60 if (request.body._tag !== "Uint8Array")
61 throw new Error("Expected a JSON request");
62 requests.push({
63 url: request.url,
64 authorization: request.headers.authorization,
65 body: JSON.parse(new TextDecoder().decode(request.body.body)),
66 });
67 return HttpClientResponse.fromWeb(request, respond());
68 }),
69 );
70 const layer = Layer.effect(
71 LanguageModel.LanguageModel,
72 makeLanguageModel({
73 model: "test-model",
74 parameters,
75 client: {
76 chatBaseUrl: Effect.succeed("https://branch.invalid/v1/"),
77 token: Effect.succeed(Redacted.make(secret)),
78 },
79 }),
80 ).pipe(
81 Layer.provide(Layer.succeed(HttpClient.HttpClient, http)),
82 Layer.provide(RuntimeContext.phantom),
83 );
84 return { layer, requests };
85};
86
87test.effect(
88 "Chat route, scoped bearer, parameters and detailed usage",
89 () =>
90 Effect.gen(function* () {
91 const { layer, requests } = harness(
92 () =>
93 Response.json(
94 completion({ content: "Hello" }, "stop", {
95 prompt_tokens: 20,
96 completion_tokens: 12,
97 prompt_tokens_details: { cached_tokens: 5 },
98 completion_tokens_details: { reasoning_tokens: 7 },
99 }),
100 ),
101 { temperature: 0.2, maxTokens: 30, seed: 4 },
102 );
103 expect(requests).toHaveLength(0);
104 const response = yield* LanguageModel.generateText({
105 prompt: "Hello",
106 }).pipe(Effect.provide(layer));
107 expect(response.text).toBe("Hello");
108 expect(response.usage.inputTokens).toMatchObject({
109 total: 20,
110 uncached: 15,
111 cacheRead: 5,
112 });
113 expect(response.usage.outputTokens).toMatchObject({
114 total: 12,
115 text: 5,
116 reasoning: 7,
117 });
118 expect(requests[0]).toMatchObject({
119 url: "https://branch.invalid/v1/chat/completions",
120 authorization: `Bearer ${secret}`,
121 body: {
122 model: "test-model",
123 stream: false,
124 max_tokens: 30,
125 temperature: 0.2,
126 seed: 4,
127 },
128 });
129 }),
130 { tags: ["unit", "provider:neon", "provider:neon:languagemodel", "local"] },
131);
132
133test.effect(
134 "content blocks and absent usage are not fabricated as zero",
135 () =>
136 Effect.gen(function* () {
137 const { layer } = harness(() =>
138 Response.json(
139 completion({
140 content: [
141 { type: "text", text: "Hello" },
142 { type: "text", text: " world" },
143 ],
144 }),
145 ),
146 );
147 const response = yield* LanguageModel.generateText({ prompt: "Hi" }).pipe(
148 Effect.provide(layer),
149 );
150 expect(response.text).toBe("Hello world");
151 expect(response.usage.inputTokens.total).toBeUndefined();
152 }),
153 { tags: ["unit", "provider:neon", "provider:neon:languagemodel", "local"] },
154);
155
156test.effect(
157 "schema output sends strict JSON Schema and decodes the object",
158 () =>
159 Effect.gen(function* () {
160 const { layer, requests } = harness(() =>
161 Response.json(completion({ content: '{"greeting":"Hello"}' })),
162 );
163 const response = yield* LanguageModel.generateObject({
164 prompt: "Greet",
165 objectName: "Greeting",
166 schema: Schema.Struct({ greeting: Schema.String }),
167 }).pipe(Effect.provide(layer));
168 expect(response.value).toEqual({ greeting: "Hello" });
169 expect(requests[0]?.body).toMatchObject({
170 response_format: {
171 type: "json_schema",
172 json_schema: {
173 name: "Greeting",
174 strict: true,
175 schema: {
176 type: "object",
177 additionalProperties: false,
178 required: ["greeting"],
179 },
180 },
181 },
182 });
183 }),
184 { tags: ["unit", "provider:neon", "provider:neon:languagemodel", "local"] },
185);
186
187test.effect(
188 "function tools execute and forced choice uses the OpenAI wire shape",
189 () =>
190 Effect.gen(function* () {
191 const { layer, requests } = harness(() =>
192 Response.json(
193 completion({ content: null, tool_calls: [toolCall] }, "tool_calls"),
194 ),
195 );
196 const response = yield* LanguageModel.generateText({
197 prompt: "Add two and three",
198 toolkit: Tools,
199 toolChoice: { tool: "sum" },
200 }).pipe(Effect.provide(layer), Effect.provide(tools));
201 expect(response.toolCalls[0]?.params).toEqual({ a: 2, b: 3 });
202 expect(response.toolResults[0]?.result).toBe(5);
203 expect(response.finishReason).toBe("tool-calls");
204 expect(requests[0]?.body).toMatchObject({
205 tool_choice: { type: "function", function: { name: "sum" } },
206 tools: [{ type: "function", function: { name: "sum" } }],
207 });
208 }),
209 { tags: ["unit", "provider:neon", "provider:neon:languagemodel", "local"] },
210);
211
212test.effect(
213 "prompt round-trips assistant calls and tool results",
214 () =>
215 Effect.gen(function* () {
216 const { layer, requests } = harness(() =>
217 Response.json(completion({ content: "Five" })),
218 );
219 yield* LanguageModel.generateText({
220 prompt: [
221 { role: "system", content: "Calculate" },
222 { role: "user", content: [{ type: "text", text: "Add" }] },
223 {
224 role: "assistant",
225 content: [
226 {
227 type: "tool-call",
228 id: "call_1",
229 name: "sum",
230 params: { a: 2, b: 3 },
231 },
232 ],
233 },
234 {
235 role: "tool",
236 content: [
237 {
238 type: "tool-result",
239 id: "call_1",
240 name: "sum",
241 result: 5,
242 isFailure: false,
243 },
244 ],
245 },
246 ],
247 }).pipe(Effect.provide(layer));
248 expect(requests[0]?.body).toMatchObject({
249 messages: [
250 { role: "system", content: "Calculate" },
251 { role: "user", content: [{ type: "text", text: "Add" }] },
252 { role: "assistant", content: null, tool_calls: [toolCall] },
253 { role: "tool", tool_call_id: "call_1", content: "5" },
254 ],
255 });
256 }),
257 { tags: ["unit", "provider:neon", "provider:neon:languagemodel", "local"] },
258);
259
260test.effect(
261 "SSE emits text ordering, late usage and one finish",
262 () =>
263 Effect.gen(function* () {
264 const { layer, requests } = harness(
265 () =>
266 new Response(
267 sse(
268 chunk({ role: "assistant", content: "Hel" }),
269 chunk({ content: "lo" }),
270 chunk({}, "stop"),
271 {
272 choices: [],
273 usage: { prompt_tokens: 4, completion_tokens: 2 },
274 },
275 "[DONE]",
276 ),
277 { headers: { "content-type": "text/event-stream" } },
278 ),
279 );
280 const parts = yield* LanguageModel.streamText({ prompt: "Hi" }).pipe(
281 Stream.provide(layer),
282 Stream.runCollect,
283 );
284 expect(parts.map((part) => part.type)).toEqual([
285 "text-start",
286 "text-delta",
287 "text-delta",
288 "text-end",
289 "finish",
290 ]);
291 expect(
292 parts.find((part) => part.type === "finish")?.usage.inputTokens.total,
293 ).toBe(4);
294 expect(requests[0]?.body).toMatchObject({
295 stream: true,
296 stream_options: { include_usage: true },
297 });
298 }),
299 { tags: ["unit", "provider:neon", "provider:neon:languagemodel", "local"] },
300);
301
302test.effect(
303 "interleaved streamed tools append fragments and execute once per call",
304 () =>
305 Effect.gen(function* () {
306 const { layer } = harness(
307 () =>
308 new Response(
309 sse(
310 chunk({
311 tool_calls: [
312 {
313 index: 0,
314 id: "one",
315 type: "function",
316 function: { name: "sum", arguments: '{"a":2,' },
317 },
318 {
319 index: 1,
320 id: "two",
321 type: "function",
322 function: { name: "sum", arguments: '{"a":4,' },
323 },
324 ],
325 }),
326 chunk({
327 tool_calls: [
328 { index: 1, function: { arguments: '"b":5}' } },
329 { index: 0, function: { arguments: '"b":3}' } },
330 ],
331 }),
332 chunk({}, "tool_calls"),
333 "[DONE]",
334 ),
335 ),
336 );
337 const parts = yield* LanguageModel.streamText({
338 prompt: "Add",
339 toolkit: Tools,
340 toolChoice: "required",
341 }).pipe(Stream.provide(layer), Stream.provide(tools), Stream.runCollect);
342 expect(
343 parts
344 .filter((part) => part.type === "tool-call")
345 .map((part) => part.params),
346 ).toEqual([
347 { a: 2, b: 3 },
348 { a: 4, b: 5 },
349 ]);
350 expect(
351 parts
352 .filter((part) => part.type === "tool-result")
353 .map((part) => part.result)
354 .sort(),
355 ).toEqual([5, 9]);
356 expect(
357 parts.filter((part) => part.type === "tool-params-start"),
358 ).toHaveLength(2);
359 expect(parts.filter((part) => part.type === "finish")).toHaveLength(1);
360 }),
361 { tags: ["unit", "provider:neon", "provider:neon:languagemodel", "local"] },
362);
363
364for (const [name, body] of [
365 ["malformed JSON", "data: not-json\n\n"],
366 ["premature EOF", sse(chunk({ content: "partial" }))],
367 ["missing DONE", sse(chunk({ content: "partial" }, "stop"))],
368 ["empty stream", ""],
369 [
370 "invalid schema",
371 sse({ choices: [{ delta: { content: 42 }, index: 0 }] }, "[DONE]"),
372 ],
373 [
374 "error event",
375 'event: error\ndata: {"error":{"message":"nt_live_SECRET_SENTINEL"}}\n\n',
376 ],
377] as const) {
378 test.effect(
379 `stream rejects ${name} without disclosing upstream payload`,
380 () =>
381 Effect.gen(function* () {
382 const { layer } = harness(() => new Response(body));
383 const result = yield* LanguageModel.streamText({ prompt: "Hi" }).pipe(
384 Stream.provide(layer),
385 Stream.runCollect,
386 Effect.result,
387 );
388 expect(Result.isFailure(result)).toBe(true);
389 expect(JSON.stringify(result)).not.toContain(secret);
390 }),
391 { tags: ["unit", "provider:neon", "provider:neon:languagemodel", "local"] },
392 );
393}
394
395for (const [status, code, tag] of [
396 [401, undefined, "AuthenticationError"],
397 [403, undefined, "AuthenticationError"],
398 [402, undefined, "QuotaExhaustedError"],
399 [429, "REQUEST_LIMIT_EXCEEDED", "QuotaExhaustedError"],
400 [429, undefined, "RateLimitError"],
401 [502, undefined, "InternalProviderError"],
402 [400, undefined, "InvalidRequestError"],
403] as const) {
404 test.effect(
405 `HTTP ${status} ${code ?? ""} maps to ${tag} and scrubs secrets`,
406 () =>
407 Effect.gen(function* () {
408 const { layer } = harness(() =>
409 Response.json(
410 { error: { code, message: `Bearer ${secret}` } },
411 { status, headers: { "retry-after": "2" } },
412 ),
413 );
414 const result = yield* LanguageModel.generateText({ prompt: "Hi" }).pipe(
415 Effect.provide(layer),
416 Effect.result,
417 );
418 expect(Result.isFailure(result)).toBe(true);
419 if (Result.isFailure(result))
420 expect(result.failure.reason._tag).toBe(tag);
421 expect(JSON.stringify(result)).not.toContain(secret);
422 }),
423 { tags: ["unit", "provider:neon", "provider:neon:languagemodel", "local"] },
424 );
425}
426
427test.effect(
428 "malformed tool JSON fails rather than executing a partial tool",
429 () =>
430 Effect.gen(function* () {
431 const { layer } = harness(() =>
432 Response.json(
433 completion(
434 {
435 tool_calls: [
436 {
437 ...toolCall,
438 function: { ...toolCall.function, arguments: "{" },
439 },
440 ],
441 },
442 "tool_calls",
443 ),
444 ),
445 );
446 const result = yield* LanguageModel.generateText({
447 prompt: "Hi",
448 toolkit: Tools,
449 }).pipe(Effect.provide(layer), Effect.provide(tools), Effect.result);
450 expect(Result.isFailure(result)).toBe(true);
451 }),
452 { tags: ["unit", "provider:neon", "provider:neon:languagemodel", "local"] },
453);
454
455test.effect(
456 "invalid JSON schema responses fail Effect's structured-output validation",
457 () =>
458 Effect.gen(function* () {
459 const { layer } = harness(() =>
460 Response.json(completion({ content: '{"greeting":42}' })),
461 );
462 const result = yield* LanguageModel.generateObject({
463 prompt: "Greet",
464 schema: Schema.Struct({ greeting: Schema.String }),
465 }).pipe(Effect.provide(layer), Effect.result);
466 expect(Result.isFailure(result)).toBe(true);
467 }),
468 { tags: ["unit", "provider:neon", "provider:neon:languagemodel", "local"] },
469);
470
471test.effect(
472 "disabled-account rejection is non-retryable quota exhaustion, not a bad token",
473 () =>
474 Effect.gen(function* () {
475 const { layer } = harness(() =>
476 Response.json(
477 { error: { message: "ai gateway not enabled for account" } },
478 { status: 403 },
479 ),
480 );
481 const result = yield* LanguageModel.generateText({ prompt: "Hi" }).pipe(
482 Effect.provide(layer),
483 Effect.result,
484 );
485 expect(Result.isFailure(result)).toBe(true);
486 if (Result.isFailure(result)) {
487 expect(result.failure.reason._tag).toBe("QuotaExhaustedError");
488 expect(result.failure.reason.isRetryable).toBe(false);
489 }
490 }),
491 { tags: ["unit", "provider:neon", "provider:neon:languagemodel", "local"] },
492);
493
494test.effect(
495 "fragmented UTF-8 SSE chunks preserve multibyte text",
496 () =>
497 Effect.gen(function* () {
498 const { layer } = harness(
499 () =>
500 new Response(
501 new ReadableStream<Uint8Array>({
502 start(controller) {
503 const bytes = new TextEncoder().encode(
504 sse(chunk({ content: "héllo 🌍" }, "stop"), "[DONE]"),
505 );
506 for (const byte of bytes)
507 controller.enqueue(new Uint8Array([byte]));
508 controller.close();
509 },
510 }),
511 ),
512 );
513 const parts = yield* LanguageModel.streamText({ prompt: "Hi" }).pipe(
514 Stream.provide(layer),
515 Stream.runCollect,
516 );
517 expect(
518 parts
519 .filter((part) => part.type === "text-delta")
520 .map((part) => part.delta)
521 .join(""),
522 ).toBe("héllo 🌍");
523 }),
524 { tags: ["unit", "provider:neon", "provider:neon:languagemodel", "local"] },
525);
526
527test.effect(
528 "truncated streamed tools never run a handler",
529 () =>
530 Effect.gen(function* () {
531 const { layer } = harness(
532 () =>
533 new Response(
534 sse(
535 chunk({
536 tool_calls: [
537 {
538 index: 0,
539 id: "one",
540 type: "function",
541 function: { name: "sum", arguments: '{"a":2,' },
542 },
543 ],
544 }),
545 chunk({}, "length"),
546 "[DONE]",
547 ),
548 ),
549 );
550 const parts = yield* LanguageModel.streamText({
551 prompt: "Add",
552 toolkit: Tools,
553 }).pipe(Stream.provide(layer), Stream.provide(tools), Stream.runCollect);
554 expect(parts.filter((part) => part.type === "tool-result")).toHaveLength(
555 0,
556 );
557 expect(parts.find((part) => part.type === "finish")?.reason).toBe(
558 "length",
559 );
560 }),
561 { tags: ["unit", "provider:neon", "provider:neon:languagemodel", "local"] },
562);
563
564test.effect(
565 "image inputs preserve order while audio fails before HTTP",
566 () =>
567 Effect.gen(function* () {
568 const { layer, requests } = harness(() =>
569 Response.json(completion({ content: "Image" })),
570 );
571 yield* LanguageModel.generateText({
572 prompt: [
573 {
574 role: "user",
575 content: [
576 { type: "text", text: "Before" },
577 {
578 type: "file",
579 mediaType: "image/png",
580 data: new Uint8Array([1, 2]),
581 },
582 { type: "text", text: "After" },
583 ],
584 },
585 ],
586 }).pipe(Effect.provide(layer));
587 expect(requests[0]?.body).toMatchObject({
588 messages: [
589 {
590 role: "user",
591 content: [
592 { type: "text", text: "Before" },
593 {
594 type: "image_url",
595 image_url: { url: "data:image/png;base64,AQI=" },
596 },
597 { type: "text", text: "After" },
598 ],
599 },
600 ],
601 });
602 const result = yield* LanguageModel.generateText({
603 prompt: [
604 {
605 role: "user",
606 content: [{ type: "file", mediaType: "audio/wav", data: "AA==" }],
607 },
608 ],
609 }).pipe(Effect.provide(layer), Effect.result);
610 expect(Result.isFailure(result)).toBe(true);
611 expect(requests).toHaveLength(1);
612 }),
613 { tags: ["unit", "provider:neon", "provider:neon:languagemodel", "local"] },
614);
615
616test.effect(
617 "downstream cancellation closes the HTTP stream without waiting for DONE",
618 () =>
619 Effect.gen(function* () {
620 let cancelled = false;
621 const { layer } = harness(
622 () =>
623 new Response(
624 new ReadableStream<Uint8Array>({
625 start(controller) {
626 controller.enqueue(
627 new TextEncoder().encode(sse(chunk({ content: "Hello" }))),
628 );
629 },
630 cancel() {
631 cancelled = true;
632 },
633 }),
634 ),
635 );
636 yield* LanguageModel.streamText({ prompt: "Hi" }).pipe(
637 Stream.provide(layer),
638 Stream.take(1),
639 Stream.runDrain,
640 );
641 expect(cancelled).toBe(true);
642 }),
643 { tags: ["unit", "provider:neon", "provider:neon:languagemodel", "local"] },
644);