1import type { QueryAIGatewayClient, LanguageModelOptions } from "alchemy/Neon";
2import { FunctionRequest } from "alchemy/Neon";
3import { RuntimeContext } from "alchemy/RuntimeContext";
4import * as NodeServices from "@effect/platform-node/NodeServices";
5import { expect, test as bunTest } from "bun:test";
6import * as Deferred from "effect/Deferred";
7import * as Effect from "effect/Effect";
8import * as Fiber from "effect/Fiber";
9import * as FileSystem from "effect/FileSystem";
10import * as Layer from "effect/Layer";
11import * as Path from "effect/Path";
12import * as Schema from "effect/Schema";
13import * as Stream from "effect/Stream";
14import { AiError, LanguageModel, Response } from "effect/ai";
15import * as Sse from "effect/encoding/Sse";
16import * as HttpServerRequest from "effect/http/HttpServerRequest";
17import type * as HttpServerResponse from "effect/http/HttpServerResponse";
18import { chat } from "../src/ai/EffectApi.ts";
19
20const test = {
21 effect: (
22 name: string,
23 body: () => Effect.Effect<unknown, unknown>,
24 options?: { timeout: number },
25 ) => bunTest(name, () => Effect.runPromise(body()), options?.timeout),
26};
27
28const environment = {
29 NEON_EXAMPLE_API_KEY: "example-secret",
30 NEON_AI_MODEL: "mock-model",
31 NEON_AI_ALLOW_PAID: "false",
32};
33const usage = new Response.Usage({
34 inputTokens: {
35 total: 1,
36 uncached: 1,
37 cacheRead: undefined,
38 cacheWrite: undefined,
39 },
40 outputTokens: { total: 1, text: 1, reasoning: undefined },
41});
42const textParts: Array<Response.StreamPartEncoded> = [
43 { type: "text-start", id: "text" },
44 { type: "text-delta", id: "text", delta: "Hello 🌍" },
45 { type: "text-end", id: "text" },
46 { type: "finish", reason: "stop", usage },
47];
48const mockModel = (
49 stream: Stream.Stream<
50 Response.StreamPartEncoded,
51 AiError.AiError
52 > = Stream.fromIterable(textParts),
53) => {
54 const calls: Array<LanguageModelOptions> = [];
55 const prompts: Array<LanguageModel.ProviderOptions["prompt"]> = [];
56 const runtimeIds: Array<string> = [];
57 const ai: Pick<QueryAIGatewayClient, "model"> = {
58 model: (options) => {
59 calls.push(options);
60 return Layer.effect(
61 LanguageModel.LanguageModel,
62 Effect.gen(function* () {
63 const runtime = yield* RuntimeContext;
64 runtimeIds.push(runtime.id);
65 return yield* LanguageModel.make({
66 generateText: () =>
67 Effect.die("Example should use LanguageModel.streamText"),
68 streamText: (options) => {
69 prompts.push(options.prompt);
70 return stream;
71 },
72 });
73 }),
74 );
75 },
76 };
77 return { ai, calls, prompts, runtimeIds };
78};
79const invoke = (
80 ai: Pick<QueryAIGatewayClient, "model">,
81 env: Readonly<Record<string, string | undefined>> = environment,
82 body = '{"prompt":"Hello"}',
83 authorization: string | null | undefined = "Bearer example-secret",
84 signal?: AbortSignal,
85) =>
86 Effect.gen(function* () {
87 const request = yield* Effect.sync(
88 () =>
89 new Request("https://example.invalid/chat", {
90 method: "POST",
91 body,
92 signal,
93 headers: {
94 "content-type": "application/json",
95 ...(authorization == null ? {} : { authorization }),
96 },
97 }),
98 );
99 return yield* chat(ai, env).pipe(
100 Effect.provideService(
101 HttpServerRequest.HttpServerRequest,
102 HttpServerRequest.fromWeb(request),
103 ),
104 Effect.provideService(FunctionRequest, request),
105 Effect.provideService(RuntimeContext, {
106 Type: "test",
107 id: "captured-request",
108 env: {},
109 get: () => Effect.succeed(undefined),
110 set: (id) => Effect.succeed(id),
111 }),
112 );
113 });
114const bodyStream = (response: HttpServerResponse.HttpServerResponse) => {
115 if (response.body._tag !== "Stream")
116 return Stream.die("Expected SSE response body");
117 return response.body.stream;
118};
119const text = (response: HttpServerResponse.HttpServerResponse) =>
120 bodyStream(response).pipe(
121 Stream.decodeText(),
122 Stream.runCollect,
123 Effect.map((chunks) => chunks.join("")),
124 );
125const UiEvent = Schema.Struct({
126 type: Schema.String,
127 id: Schema.optional(Schema.String),
128 delta: Schema.optional(Schema.String),
129 errorText: Schema.optional(Schema.String),
130 finishReason: Schema.optional(Schema.String),
131});
132const events = (wire: string) =>
133 Stream.succeed(wire).pipe(
134 Stream.pipeThroughChannel(Sse.decode<never, unknown>()),
135 Stream.map((event) => event.data),
136 Stream.filter((data) => data !== "[DONE]"),
137 Stream.mapEffect((data) =>
138 Schema.decodeUnknownEffect(Schema.fromJsonString(UiEvent))(data),
139 ),
140 Stream.runCollect,
141 );
142
143for (const [name, authorization] of [
144 ["missing", undefined],
145 ["wrong", "Bearer wrong"],
146] as const) {
147 test.effect(
148 `Effect example returns 401 for ${name} authorization before validation`,
149 () =>
150 Effect.gen(function* () {
151 const mock = mockModel();
152 const response = yield* invoke(
153 mock.ai,
154 environment,
155 "not JSON",
156 authorization ?? null,
157 );
158 expect(response.status).toBe(401);
159 expect(mock.calls).toHaveLength(0);
160 }).pipe(Effect.scoped),
161 );
162}
163
164for (const [name, body] of [
165 ["malformed JSON", "{"],
166 ["missing prompt", "{}"],
167 ["null", "null"],
168 ["non-string prompt", '{"prompt":4}'],
169 ["blank prompt", '{"prompt":" "}'],
170 ["oversized prompt", JSON.stringify({ prompt: "x".repeat(4001) })],
171] as const) {
172 test.effect(
173 `Effect example returns 400 for ${name} before the paid gate`,
174 () =>
175 Effect.gen(function* () {
176 const mock = mockModel();
177 const response = yield* invoke(mock.ai, environment, body);
178 expect(response.status).toBe(400);
179 expect(mock.calls).toHaveLength(0);
180 }).pipe(Effect.scoped),
181 );
182}
183
184for (const [name, env] of [
185 ["disabled inference", environment],
186 [
187 "missing model",
188 { ...environment, NEON_AI_ALLOW_PAID: "true", NEON_AI_MODEL: undefined },
189 ],
190 ["non-explicit opt-in", { ...environment, NEON_AI_ALLOW_PAID: "1" }],
191] as const) {
192 test.effect(
193 `Effect example returns 503 for ${name} without constructing a model`,
194 () =>
195 Effect.gen(function* () {
196 const mock = mockModel();
197 const response = yield* invoke(mock.ai, env);
198 expect(response.status).toBe(503);
199 expect(mock.calls).toHaveLength(0);
200 }).pipe(Effect.scoped),
201 );
202}
203
204test.effect(
205 "Effect example calls the bound LanguageModel and preserves context in its deferred SSE body",
206 () =>
207 Effect.gen(function* () {
208 const mock = mockModel();
209 const response = yield* invoke(mock.ai, {
210 ...environment,
211 NEON_AI_ALLOW_PAID: "true",
212 });
213 expect(response.status).toBe(200);
214 expect(response.headers["content-type"]).toBe("text/event-stream");
215 expect(response.headers["x-vercel-ai-ui-message-stream"]).toBe("v1");
216 expect(mock.calls).toEqual([
217 { model: "mock-model", parameters: { maxTokens: 128 } },
218 ]);
219 expect(mock.prompts).toHaveLength(0);
220 // Consume after invoke's request-service provision has ended.
221 const wire = yield* text(response);
222 expect(mock.runtimeIds).toEqual(["captured-request"]);
223 expect(mock.prompts).toHaveLength(1);
224 expect(mock.prompts[0]?.content).toMatchObject([
225 { role: "user", content: [{ type: "text", text: "Hello" }] },
226 ]);
227 const parsed = yield* events(wire);
228 expect(parsed.map((event) => event.type)).toEqual([
229 "start",
230 "start-step",
231 "text-start",
232 "text-delta",
233 "text-end",
234 "finish-step",
235 "finish",
236 ]);
237 expect(parsed.find((event) => event.type === "text-delta")?.delta).toBe(
238 "Hello 🌍",
239 );
240 expect(wire).toContain("data: [DONE]");
241 }).pipe(Effect.scoped),
242);
243
244test.effect(
245 "Effect example sanitizes AiError SSE failures without retrying or leaking credentials",
246 () =>
247 Effect.gen(function* () {
248 const mock = mockModel(
249 Stream.concat(
250 Stream.fromIterable(textParts.slice(0, 2)),
251 Stream.fail(
252 AiError.make({
253 module: "test",
254 method: "streamText",
255 reason: new AiError.InvalidRequestError({
256 description: "UPSTREAM_TOKEN_SENTINEL",
257 }),
258 }),
259 ),
260 ),
261 );
262 const response = yield* invoke(mock.ai, {
263 ...environment,
264 NEON_AI_ALLOW_PAID: "true",
265 });
266 const wire = yield* text(response);
267 const parsed = yield* events(wire);
268 expect(parsed.filter((event) => event.type === "error")).toHaveLength(1);
269 expect(
270 parsed.find((event) => event.type === "error")?.errorText,
271 ).toContain("Gateway rejected the request");
272 expect(parsed.some((event) => event.type === "finish")).toBe(false);
273 expect(wire).not.toContain("UPSTREAM_TOKEN_SENTINEL");
274 expect(wire).not.toContain("example-secret");
275 expect(wire).toContain("data: [DONE]");
276 expect(mock.prompts).toHaveLength(1);
277 }).pipe(Effect.scoped),
278);
279
280test.effect(
281 "Effect example request abort interrupts upstream model consumption",
282 () =>
283 Effect.gen(function* () {
284 const controller = yield* Effect.sync(() => new AbortController());
285 const started = yield* Deferred.make<void>();
286 let finalized = false;
287 const mock = mockModel(
288 Stream.fromEffect(Deferred.succeed(started, undefined)).pipe(
289 Stream.drain,
290 Stream.concat(Stream.never),
291 Stream.ensuring(
292 Effect.sync(() => {
293 finalized = true;
294 }),
295 ),
296 ),
297 );
298 const response = yield* invoke(
299 mock.ai,
300 { ...environment, NEON_AI_ALLOW_PAID: "true" },
301 undefined,
302 undefined,
303 controller.signal,
304 );
305 const consumer = yield* bodyStream(response).pipe(
306 Stream.runDrain,
307 Effect.forkChild,
308 );
309 yield* Deferred.await(started);
310 yield* Effect.sync(() => controller.abort());
311 yield* Fiber.join(consumer);
312 expect(finalized).toBe(true);
313 }).pipe(Effect.scoped),
314 { timeout: 5000 },
315);
316
317test.effect(
318 "Effect example does not start model inference for an already-aborted request",
319 () =>
320 Effect.gen(function* () {
321 const controller = yield* Effect.sync(() => new AbortController());
322 yield* Effect.sync(() => controller.abort());
323 const mock = mockModel();
324 const response = yield* invoke(
325 mock.ai,
326 { ...environment, NEON_AI_ALLOW_PAID: "true" },
327 undefined,
328 undefined,
329 controller.signal,
330 );
331 yield* bodyStream(response).pipe(Stream.runDrain);
332 expect(mock.prompts).toHaveLength(0);
333 }).pipe(Effect.scoped),
334 { timeout: 5000 },
335);
336
337test.effect(
338 "Effect example downstream cancellation closes the model stream",
339 () =>
340 Effect.gen(function* () {
341 let finalized = false;
342 const mock = mockModel(
343 Stream.fromIterable(textParts.slice(0, 2)).pipe(
344 Stream.concat(Stream.never),
345 Stream.ensuring(
346 Effect.sync(() => {
347 finalized = true;
348 }),
349 ),
350 ),
351 );
352 const response = yield* invoke(mock.ai, {
353 ...environment,
354 NEON_AI_ALLOW_PAID: "true",
355 });
356 yield* bodyStream(response).pipe(
357 Stream.decodeText(),
358 Stream.takeUntil((chunk) => chunk.includes("text-delta")),
359 Stream.runDrain,
360 );
361 expect(finalized).toBe(true);
362 }).pipe(Effect.scoped),
363 { timeout: 5000 },
364);
365
366test.effect(
367 "example source uses the canonical Effect model and gateway reference documentation is discoverable",
368 () =>
369 Effect.gen(function* () {
370 const fs = yield* FileSystem.FileSystem;
371 const path = yield* Path.Path;
372 const examplePath = yield* path.fromFileUrl(
373 new URL("../src/ai/EffectApi.ts", import.meta.url),
374 );
375 const gatewayPath = yield* path.fromFileUrl(
376 new URL(
377 "../../../packages/alchemy/src/Neon/AIGateway.ts",
378 import.meta.url,
379 ),
380 );
381 const source = yield* fs.readFileString(examplePath);
382 expect(source).toContain("const model = ai.model(");
383 expect(source).toContain("LanguageModel.streamText(");
384 expect(source).toContain("Stream.provideContext(context)");
385 expect(source).not.toContain("@neon/ai-sdk-provider");
386 expect(source).not.toContain('from "ai"');
387 for (const name of [
388 "NEON_EXAMPLE_API_KEY",
389 "NEON_AI_MODEL",
390 "NEON_AI_ALLOW_PAID",
391 ]) {
392 expect(source).toContain(`environment.${name}`);
393 expect(source).toContain(`("${name}")`);
394 }
395 expect(yield* fs.readFileString(gatewayPath)).toContain(
396 " * @resource\n * @category AI Gateway",
397 );
398 }).pipe(Effect.provide(NodeServices.layer)),
399);