1/// <reference types="@cloudflare/workers-types" />
2
3import * as Effect from "effect/Effect";
4import * as Layer from "effect/Layer";
5import * as Stream from "effect/Stream";
6import {
7 AiError,
8 LanguageModel as AiLanguageModel,
9 IdGenerator,
10 Prompt,
11 Response,
12 Tool,
13} from "effect/ai";
14import * as Sse from "effect/encoding/Sse";
15import type { RuntimeContext } from "../../RuntimeContext.ts";
16
17/**
18 * The slice of an AI client the LanguageModel adapter needs: the raw Workers
19 * AI handle plus, when routed through an AI Gateway, the gateway id.
20 *
21 * Both `Cloudflare.AI.QueryGateway(gateway)` (gateway-routed) and
22 * `Cloudflare.Workers.AI()` (plain Workers AI binding) clients satisfy it.
23 */
24export interface LanguageModelClient {
25 /** Effect resolving to the raw Workers AI runtime binding. */
26 readonly raw: Effect.Effect<Ai, never, RuntimeContext>;
27 /**
28 * Effect resolving to the AI Gateway id to route `ai.run` calls through.
29 * Omit for a plain Workers AI binding — requests then go directly to
30 * Workers AI without a gateway.
31 */
32 readonly id?: Effect.Effect<string, never, RuntimeContext>;
33}
34
35/**
36 * Options for constructing a Workers AI-backed LanguageModel.
37 */
38export interface LanguageModelOptions {
39 /**
40 * Already-bound AI client — from `Cloudflare.AI.QueryGateway(gateway)`
41 * (routed through the gateway) or `Cloudflare.Workers.AI()` (direct).
42 */
43 readonly client: LanguageModelClient;
44 /** Workers AI model id, e.g. `@cf/meta/llama-3.3-70b-instruct-fp8-fast`. */
45 readonly model: string;
46 /** Optional per-call defaults; overridable per request via `providerOptions`. */
47 readonly parameters?: {
48 readonly temperature?: number;
49 readonly maxTokens?: number;
50 readonly topP?: number;
51 readonly topK?: number;
52 readonly seed?: number;
53 readonly frequencyPenalty?: number;
54 readonly presencePenalty?: number;
55 };
56}
57
58/**
59 * Provide a {@link AiLanguageModel.LanguageModel} layer backed by the supplied
60 * AI Gateway client and Workers AI model.
61 */
62export const makeLanguageModelLayer = (
63 options: LanguageModelOptions,
64): Layer.Layer<AiLanguageModel.LanguageModel, never, RuntimeContext> =>
65 Layer.effect(AiLanguageModel.LanguageModel, makeLanguageModel(options));
66
67/**
68 * Build a {@link AiLanguageModel.LanguageModel} that proxies generateText/streamText
69 * through the supplied AI Gateway client to a Workers AI model.
70 */
71export const makeLanguageModel = ({
72 client,
73 model,
74 parameters,
75}: LanguageModelOptions): Effect.Effect<
76 AiLanguageModel.LanguageModel,
77 never,
78 RuntimeContext
79> =>
80 Effect.gen(function* () {
81 const ai = yield* client.raw;
82 const gatewayId = client.id === undefined ? undefined : yield* client.id;
83
84 const callRaw = (
85 body: WorkersAiInputs,
86 method: "generateText" | "streamText",
87 ): Effect.Effect<Response, AiError.AiError> =>
88 Effect.tryPromise({
89 try: () =>
90 ai.run(
91 model as keyof AiModels,
92 body as unknown as AiModels[keyof AiModels]["inputs"],
93 {
94 ...(gatewayId === undefined
95 ? {}
96 : { gateway: { id: gatewayId } }),
97 returnRawResponse: true,
98 },
99 ),
100 catch: (cause) => toAiError(cause, method),
101 });
102
103 return yield* AiLanguageModel.make({
104 generateText: (options) =>
105 Effect.gen(function* () {
106 const body = toRequestBody({ options, parameters, stream: false });
107 const resp = yield* callRaw(body, "generateText");
108 const json = yield* Effect.tryPromise({
109 try: () => resp.json() as Promise<Record<string, unknown>>,
110 catch: (cause) => toAiError(cause, "generateText"),
111 });
112 return yield* parseGenerateText(json);
113 }),
114 streamText: (options) =>
115 Stream.unwrap(
116 Effect.gen(function* () {
117 const idGen = yield* IdGenerator.IdGenerator;
118 const body = toRequestBody({ options, parameters, stream: true });
119 const resp = yield* callRaw(body, "streamText");
120 const hasTools =
121 options.tools.length > 0 && options.toolChoice !== "none";
122 return parseStreamText(resp, idGen, hasTools);
123 }),
124 ),
125 });
126 });
127
128// ---------------------------------------------------------------------------
129// Wire format types (Workers AI request)
130//
131// Workers AI returns two response shapes depending on the model:
132// - Native: { response: "...", tool_calls: [...] , usage }
133// - OpenAI: { choices: [{ message: { content, tool_calls, reasoning_content } }], usage }
134//
135// We accept both defensively — schemas would over-constrain.
136// ---------------------------------------------------------------------------
137
138interface WorkersAiMessage {
139 readonly role: "system" | "user" | "assistant" | "tool";
140 readonly content?: unknown;
141 readonly name?: string;
142 readonly tool_call_id?: string;
143 readonly tool_calls?: ReadonlyArray<{
144 readonly id: string;
145 readonly type: "function";
146 readonly function: { readonly name: string; readonly arguments: string };
147 }>;
148 readonly reasoning?: string;
149}
150
151interface WorkersAiToolDef {
152 readonly type: "function";
153 readonly function: {
154 readonly name: string;
155 readonly description?: string;
156 readonly parameters: unknown;
157 };
158}
159
160type WorkersAiToolChoice =
161 | "auto"
162 | "required"
163 | "none"
164 | { readonly type: "function"; readonly function: { readonly name: string } };
165
166interface WorkersAiInputs {
167 readonly messages: ReadonlyArray<WorkersAiMessage>;
168 readonly tools?: ReadonlyArray<WorkersAiToolDef>;
169 readonly tool_choice?: WorkersAiToolChoice;
170 readonly stream?: boolean;
171 readonly stream_options?: { readonly include_usage: boolean };
172 readonly max_tokens?: number;
173 readonly temperature?: number;
174 readonly top_p?: number;
175 readonly top_k?: number;
176 readonly random_seed?: number;
177 readonly frequency_penalty?: number;
178 readonly presence_penalty?: number;
179}
180
181// ---------------------------------------------------------------------------
182// Prompt → Workers AI messages (pure, no .push mutation)
183// ---------------------------------------------------------------------------
184
185const uint8ArrayToBase64 = (bytes: Uint8Array): string => {
186 let binary = "";
187 const chunkSize = 8192;
188 for (let i = 0; i < bytes.length; i += chunkSize) {
189 const chunk = bytes.subarray(i, Math.min(i + chunkSize, bytes.length));
190 binary += String.fromCharCode(...chunk);
191 }
192 return btoa(binary);
193};
194
195const fileToImageUrl = (
196 data: string | Uint8Array | URL,
197 mediaType: string,
198): string => {
199 if (data instanceof URL) return data.toString();
200 if (data instanceof Uint8Array) {
201 return `data:${mediaType};base64,${uint8ArrayToBase64(data)}`;
202 }
203 if (data.startsWith("data:") || data.startsWith("http")) return data;
204 return `data:${mediaType};base64,${data}`;
205};
206
207const convertPromptToMessages = (
208 prompt: Prompt.Prompt,
209): ReadonlyArray<WorkersAiMessage> =>
210 prompt.content.flatMap((m): ReadonlyArray<WorkersAiMessage> => {
211 switch (m.role) {
212 case "system":
213 return [{ role: "system", content: m.content }];
214 case "user":
215 return [toUserMessage(m.content)];
216 case "assistant":
217 return [toAssistantMessage(m.content)];
218 case "tool":
219 return m.content.flatMap(toToolMessage);
220 }
221 });
222
223const toUserMessage = (
224 parts: Prompt.UserMessage["content"],
225): WorkersAiMessage => {
226 const text = parts
227 .flatMap((p) => (p.type === "text" ? [p.text] : []))
228 .join("\n");
229 const images = parts.flatMap((p) =>
230 p.type === "file"
231 ? [
232 {
233 type: "image_url" as const,
234 image_url: { url: fileToImageUrl(p.data, p.mediaType) },
235 },
236 ]
237 : [],
238 );
239 if (images.length === 0) return { role: "user", content: text };
240 return {
241 role: "user",
242 content: [
243 ...(text.length > 0 ? [{ type: "text" as const, text }] : []),
244 ...images,
245 ],
246 };
247};
248
249const toAssistantMessage = (
250 parts: Prompt.AssistantMessage["content"],
251): WorkersAiMessage => {
252 const text = parts
253 .flatMap((p) => (p.type === "text" ? [p.text] : []))
254 .join("");
255 const reasoning = parts
256 .flatMap((p) => (p.type === "reasoning" ? [p.text] : []))
257 .join("");
258 const toolCalls = parts.flatMap((p) =>
259 p.type === "tool-call"
260 ? [
261 {
262 id: p.id,
263 type: "function" as const,
264 function: {
265 name: p.name,
266 arguments:
267 typeof p.params === "string"
268 ? p.params
269 : JSON.stringify(p.params ?? {}),
270 },
271 },
272 ]
273 : [],
274 );
275 return {
276 role: "assistant",
277 content: text,
278 ...(reasoning ? { reasoning } : {}),
279 ...(toolCalls.length > 0 ? { tool_calls: toolCalls } : {}),
280 };
281};
282
283const toToolMessage = (
284 part: Prompt.ToolMessage["content"][number],
285): ReadonlyArray<WorkersAiMessage> =>
286 part.type === "tool-result"
287 ? [
288 {
289 role: "tool",
290 name: part.name,
291 tool_call_id: part.id,
292 content:
293 typeof part.result === "string"
294 ? part.result
295 : JSON.stringify(part.result),
296 },
297 ]
298 : [];
299
300// ---------------------------------------------------------------------------
301// Tools / tool_choice
302// ---------------------------------------------------------------------------
303
304const prepareTools = (
305 tools: ReadonlyArray<Tool.Any>,
306 toolChoice: AiLanguageModel.ProviderOptions["toolChoice"],
307): {
308 tools?: ReadonlyArray<WorkersAiToolDef>;
309 tool_choice?: WorkersAiToolChoice;
310} => {
311 if (tools.length === 0) return {};
312 const mapped: ReadonlyArray<WorkersAiToolDef> = tools.map((tool) => ({
313 type: "function",
314 function: {
315 name: tool.name,
316 description: Tool.getDescription(tool),
317 parameters: Tool.getJsonSchema(tool),
318 },
319 }));
320
321 if (toolChoice === "auto" || toolChoice == null) {
322 return { tools: mapped, tool_choice: "auto" };
323 }
324 if (toolChoice === "none") return { tools: mapped, tool_choice: "none" };
325 if (toolChoice === "required") {
326 return { tools: mapped, tool_choice: "required" };
327 }
328 if (typeof toolChoice === "object" && "tool" in toolChoice) {
329 return {
330 tools: mapped.filter((t) => t.function.name === toolChoice.tool),
331 tool_choice: "required",
332 };
333 }
334 if (typeof toolChoice === "object" && "oneOf" in toolChoice) {
335 const allowed = new Set(toolChoice.oneOf);
336 return {
337 tools: mapped.filter((t) => allowed.has(t.function.name)),
338 tool_choice: toolChoice.mode === "required" ? "required" : "auto",
339 };
340 }
341 return { tools: mapped, tool_choice: "auto" };
342};
343
344// ---------------------------------------------------------------------------
345// Request body
346// ---------------------------------------------------------------------------
347
348const toRequestBody = ({
349 options,
350 parameters,
351 stream,
352}: {
353 readonly options: AiLanguageModel.ProviderOptions;
354 readonly parameters: LanguageModelOptions["parameters"];
355 readonly stream: boolean;
356}): WorkersAiInputs => {
357 const messages = convertPromptToMessages(options.prompt);
358 const { tools, tool_choice } = prepareTools(
359 options.tools,
360 options.toolChoice,
361 );
362 return {
363 messages,
364 ...(tools !== undefined ? { tools } : {}),
365 ...(tool_choice !== undefined ? { tool_choice } : {}),
366 // `stream_options.include_usage` is the OpenAI-compatible opt-in for
367 // usage tokens to appear in the final streamed chunk. Without it most
368 // Workers AI models omit `usage` from the stream entirely, leaving the
369 // `finish` part with zeroed counts.
370 ...(stream
371 ? { stream: true, stream_options: { include_usage: true } }
372 : {}),
373 ...(parameters?.maxTokens !== undefined
374 ? { max_tokens: parameters.maxTokens }
375 : {}),
376 ...(parameters?.temperature !== undefined
377 ? { temperature: parameters.temperature }
378 : {}),
379 ...(parameters?.topP !== undefined ? { top_p: parameters.topP } : {}),
380 ...(parameters?.topK !== undefined ? { top_k: parameters.topK } : {}),
381 ...(parameters?.seed !== undefined ? { random_seed: parameters.seed } : {}),
382 ...(parameters?.frequencyPenalty !== undefined
383 ? { frequency_penalty: parameters.frequencyPenalty }
384 : {}),
385 ...(parameters?.presencePenalty !== undefined
386 ? { presence_penalty: parameters.presencePenalty }
387 : {}),
388 };
389};
390
391// ---------------------------------------------------------------------------
392// Finish reason / usage mapping
393// ---------------------------------------------------------------------------
394
395const mapFinishReason = (raw: unknown): Response.FinishReason => {
396 switch (raw) {
397 case "stop":
398 return "stop";
399 case "length":
400 case "model_length":
401 return "length";
402 case "tool_calls":
403 return "tool-calls";
404 case "content_filter":
405 case "content-filter":
406 return "content-filter";
407 case "error":
408 return "error";
409 case undefined:
410 case null:
411 return "unknown";
412 default:
413 return "other";
414 }
415};
416
417const mapUsage = (raw: Record<string, unknown> | undefined): Response.Usage => {
418 const usage = (raw?.usage as Record<string, unknown> | undefined) ?? {};
419 const promptTokens = (usage.prompt_tokens as number | undefined) ?? 0;
420 const completionTokens = (usage.completion_tokens as number | undefined) ?? 0;
421 const cached = (
422 usage.prompt_tokens_details as { cached_tokens?: number } | undefined
423 )?.cached_tokens;
424 // Construct an actual `Response.Usage` instance — `Schema.Class<Usage>`
425 // encodes by going through the class constructor / `isInstance` check, so a
426 // plain struct that "matches" the encoded shape isn't enough.
427 return new Response.Usage({
428 inputTokens: {
429 uncached:
430 cached !== undefined
431 ? Math.max(0, promptTokens - cached)
432 : promptTokens,
433 total: promptTokens,
434 cacheRead: cached ?? 0,
435 cacheWrite: 0,
436 },
437 outputTokens: {
438 total: completionTokens,
439 text: 0,
440 reasoning: 0,
441 },
442 });
443};
444
445// ---------------------------------------------------------------------------
446// generateText: JSON → Response.PartEncoded[]
447//
448// Normalize the dual-shape (native + OpenAI) response into a single
449// `DecodedResponse` once, then build the part list with pure spreads.
450// ---------------------------------------------------------------------------
451
452interface DecodedToolCall {
453 readonly rawId: string;
454 readonly name: string;
455 readonly arguments: unknown;
456}
457
458interface DecodedResponse {
459 readonly text: string | undefined;
460 readonly reasoning: string | undefined;
461 readonly toolCalls: ReadonlyArray<DecodedToolCall>;
462 readonly finishReason: string | undefined;
463}
464
465const decodeResponse = (raw: Record<string, unknown>): DecodedResponse => {
466 const choice = (
467 raw.choices as
468 | Array<{
469 message?: {
470 content?: string | null;
471 reasoning_content?: string;
472 reasoning?: string;
473 tool_calls?: ReadonlyArray<Record<string, unknown>>;
474 };
475 finish_reason?: string;
476 }>
477 | undefined
478 )?.[0];
479 const message = choice?.message;
480
481 const openAiText = message?.content;
482 const text =
483 typeof openAiText === "string" && openAiText.length > 0
484 ? openAiText
485 : nativeTextOf(raw.response);
486 const reasoning = message?.reasoning_content ?? message?.reasoning;
487 const rawToolCalls =
488 message?.tool_calls ??
489 (Array.isArray(raw.tool_calls)
490 ? (raw.tool_calls as ReadonlyArray<Record<string, unknown>>)
491 : []);
492
493 return {
494 text,
495 reasoning: reasoning && reasoning.length > 0 ? reasoning : undefined,
496 toolCalls: rawToolCalls.flatMap(decodeToolCall),
497 finishReason:
498 choice?.finish_reason ?? (raw.finish_reason as string | undefined),
499 };
500};
501
502const nativeTextOf = (raw: unknown): string | undefined => {
503 if (raw == null) return undefined;
504 if (typeof raw === "object") return JSON.stringify(raw);
505 const text = String(raw);
506 return text.length > 0 ? text : undefined;
507};
508
509const decodeToolCall = (
510 tc: Record<string, unknown>,
511): ReadonlyArray<DecodedToolCall> => {
512 const fn = tc.function as { name?: string; arguments?: unknown } | undefined;
513 const rawId = (tc.id as string | undefined) ?? "";
514 if (fn?.name) {
515 return [{ rawId, name: fn.name, arguments: fn.arguments ?? "" }];
516 }
517 const flatName = tc.name as string | undefined;
518 if (flatName) {
519 return [{ rawId, name: flatName, arguments: tc.arguments ?? "" }];
520 }
521 return [];
522};
523
524const tryParseJsonArgs = (raw: unknown): unknown => {
525 if (typeof raw !== "string") return raw;
526 try {
527 return JSON.parse(raw);
528 } catch {
529 // Leave as raw string; the framework's tool-result decoder will fail loudly.
530 return raw;
531 }
532};
533
534const parseGenerateText = Effect.fn(function* (raw: Record<string, unknown>) {
535 const idGen = yield* IdGenerator.IdGenerator;
536 const decoded = decodeResponse(raw);
537
538 const toolCallParts = yield* Effect.forEach(decoded.toolCalls, (tc) =>
539 Effect.gen(function* () {
540 const id = tc.rawId || (yield* idGen.generateId());
541 return {
542 type: "tool-call" as const,
543 id,
544 name: tc.name,
545 params: tryParseJsonArgs(tc.arguments),
546 };
547 }),
548 );
549
550 const finish = mapFinishReason(
551 decoded.finishReason ??
552 (decoded.toolCalls.length > 0 ? "tool_calls" : "stop"),
553 );
554
555 return [
556 ...(decoded.reasoning !== undefined
557 ? [{ type: "reasoning" as const, text: decoded.reasoning }]
558 : []),
559 ...(decoded.text !== undefined && decoded.text.length > 0
560 ? [{ type: "text" as const, text: decoded.text }]
561 : []),
562 ...toolCallParts,
563 {
564 type: "finish" as const,
565 reason: finish,
566 usage: mapUsage(raw),
567 response: undefined,
568 },
569 ] satisfies ReadonlyArray<Response.PartEncoded>;
570});
571
572// ---------------------------------------------------------------------------
573// streamText: SSE byte stream → Stream<Response.StreamPartEncoded>
574//
575// Immutable `StreamState` is threaded through `Stream.mapAccumEffect`. The
576// per-chunk output buffer (`parts: Array<StreamPartEncoded>`) is mutable for
577// performance — it's scoped to one chunk, never escapes the handler, and lets
578// us avoid the O(n²) array-spread that pure threading would force in the hot
579// path. This matches the pattern Effect's own `@effect/ai-*` adapters use.
580// ---------------------------------------------------------------------------
581
582interface StreamState {
583 readonly textId: string | undefined;
584 readonly reasoningId: string | undefined;
585 readonly toolCalls: ReadonlyMap<
586 number,
587 { readonly id: string; readonly name: string; readonly arguments: string }
588 >;
589 readonly lastToolIndex: number | undefined;
590 readonly closedToolIndices: ReadonlySet<number>;
591 readonly usage: Record<string, unknown> | undefined;
592 readonly finishReason: string | undefined;
593 readonly receivedAnyData: boolean;
594 readonly receivedDone: boolean;
595 // Workers AI's *native* streaming for function-calling models (e.g.
596 // `@cf/meta/llama-3.3-70b`) does not emit structured `tool_calls` deltas
597 // — it streams the call as a JSON document inside the `response` text
598 // field (`{"name":"get_weather","parameters":{...}}`). When tools were
599 // requested we buffer that text here instead of emitting it as
600 // text-delta, then parse it into tool-params parts on finalize.
601 readonly nativeToolBuffer: string;
602 readonly nativeToolId: string | undefined;
603}
604
605const initialStreamState = (): StreamState => ({
606 textId: undefined,
607 reasoningId: undefined,
608 toolCalls: new Map(),
609 lastToolIndex: undefined,
610 closedToolIndices: new Set(),
611 usage: undefined,
612 finishReason: undefined,
613 receivedAnyData: false,
614 receivedDone: false,
615 nativeToolBuffer: "",
616 nativeToolId: undefined,
617});
618
619type StreamParts = Array<Response.StreamPartEncoded>;
620
621const tryParseJson = (data: string): Record<string, unknown> | undefined => {
622 try {
623 const v = JSON.parse(data);
624 return v && typeof v === "object"
625 ? (v as Record<string, unknown>)
626 : undefined;
627 } catch {
628 return undefined;
629 }
630};
631
632const isNullFinalizationToolCall = (tc: Record<string, unknown>): boolean => {
633 const fn = tc.function as Record<string, unknown> | undefined;
634 const name = fn?.name ?? tc.name ?? null;
635 const args = fn?.arguments ?? tc.arguments ?? null;
636 const id = tc.id ?? null;
637 return !id && !name && (!args || args === "");
638};
639
640const closeReasoning = (
641 state: StreamState,
642 parts: StreamParts,
643): StreamState => {
644 if (state.reasoningId === undefined) return state;
645 parts.push({ type: "reasoning-end", id: state.reasoningId });
646 return { ...state, reasoningId: undefined };
647};
648
649const closeToolCall = (
650 state: StreamState,
651 index: number,
652 parts: StreamParts,
653): StreamState => {
654 if (state.closedToolIndices.has(index)) return state;
655 const tc = state.toolCalls.get(index);
656 if (!tc) return state;
657 parts.push({ type: "tool-params-end", id: tc.id });
658 const closed = new Set(state.closedToolIndices);
659 closed.add(index);
660 return { ...state, closedToolIndices: closed };
661};
662
663const emitTextDelta = (
664 state: StreamState,
665 delta: string,
666 parts: StreamParts,
667 idGen: IdGenerator.Service,
668): Effect.Effect<StreamState> =>
669 Effect.gen(function* () {
670 let s = closeReasoning(state, parts);
671 if (s.textId === undefined) {
672 const id = yield* idGen.generateId();
673 parts.push({ type: "text-start", id });
674 s = { ...s, textId: id };
675 }
676 parts.push({ type: "text-delta", id: s.textId!, delta });
677 return s;
678 });
679
680const emitReasoningDelta = (
681 state: StreamState,
682 delta: string,
683 parts: StreamParts,
684 idGen: IdGenerator.Service,
685): Effect.Effect<StreamState> =>
686 Effect.gen(function* () {
687 let s = state;
688 if (s.reasoningId === undefined) {
689 const id = yield* idGen.generateId();
690 parts.push({ type: "reasoning-start", id });
691 s = { ...s, reasoningId: id };
692 }
693 parts.push({ type: "reasoning-delta", id: s.reasoningId!, delta });
694 return s;
695 });
696
697const handleToolDeltas = (
698 state: StreamState,
699 deltas: ReadonlyArray<Record<string, unknown>>,
700 parts: StreamParts,
701 idGen: IdGenerator.Service,
702): Effect.Effect<StreamState> =>
703 Effect.gen(function* () {
704 let s = state;
705 for (const d of deltas) {
706 if (isNullFinalizationToolCall(d)) {
707 if (s.lastToolIndex !== undefined) {
708 s = closeToolCall(s, s.lastToolIndex, parts);
709 }
710 continue;
711 }
712 const idx = (d.index as number | undefined) ?? 0;
713 const fn = d.function as
714 | { name?: string; arguments?: string }
715 | undefined;
716 const name = fn?.name ?? (d.name as string | undefined) ?? "";
717 const args = fn?.arguments ?? (d.arguments as string | undefined) ?? "";
718 const rawId = (d.id as string | undefined) ?? "";
719
720 const existing = s.toolCalls.get(idx);
721 if (existing === undefined) {
722 if (s.lastToolIndex !== undefined && s.lastToolIndex !== idx) {
723 s = closeToolCall(s, s.lastToolIndex, parts);
724 }
725 const id = rawId || (yield* idGen.generateId());
726 const entry = { id, name, arguments: args };
727 const next = new Map(s.toolCalls);
728 next.set(idx, entry);
729 s = { ...s, toolCalls: next, lastToolIndex: idx };
730 parts.push({ type: "tool-params-start", id, name });
731 if (args.length > 0) {
732 parts.push({ type: "tool-params-delta", id, delta: args });
733 }
734 } else {
735 // OpenAI-compatible providers stream argument fragments, while
736 // Workers AI can resend the complete accumulated JSON on every
737 // chunk. Normalize both shapes to incremental deltas.
738 const delta = args.startsWith(existing.arguments)
739 ? args.slice(existing.arguments.length)
740 : args;
741 const accumulated = args.startsWith(existing.arguments)
742 ? args
743 : existing.arguments + args;
744 const next = new Map(s.toolCalls);
745 next.set(idx, { ...existing, arguments: accumulated });
746 s = { ...s, toolCalls: next, lastToolIndex: idx };
747 if (delta.length > 0) {
748 parts.push({
749 type: "tool-params-delta",
750 id: existing.id,
751 delta,
752 });
753 }
754 }
755 }
756 return s;
757 });
758
759const hasNonZeroUsage = (raw: unknown): boolean => {
760 if (raw == null || typeof raw !== "object") return false;
761 const u = raw as Record<string, unknown>;
762 const prompt = (u.prompt_tokens as number | undefined) ?? 0;
763 const completion = (u.completion_tokens as number | undefined) ?? 0;
764 const total = (u.total_tokens as number | undefined) ?? 0;
765 return prompt > 0 || completion > 0 || total > 0;
766};
767
768const updateChunkMeta = (
769 state: StreamState,
770 chunk: Record<string, unknown>,
771): StreamState => {
772 let s = state;
773 // Workers AI's native stream emits the real usage chunk, then a
774 // "zero-valued terminator" chunk where every count is 0 (it also re-emits
775 // `usage` with all zeros). Treat the zero chunk as a no-op so we keep the
776 // meaningful counts.
777 if (chunk.usage !== undefined && hasNonZeroUsage(chunk.usage)) {
778 s = { ...s, usage: chunk };
779 }
780 const choices = chunk.choices as
781 | Array<{ finish_reason?: string }>
782 | undefined;
783 const finish =
784 choices?.[0]?.finish_reason ?? (chunk.finish_reason as string | undefined);
785 if (finish != null) s = { ...s, finishReason: finish };
786 return s;
787};
788
789const handleNativeText = (
790 state: StreamState,
791 chunk: Record<string, unknown>,
792 parts: StreamParts,
793 idGen: IdGenerator.Service,
794 hasTools: boolean,
795): Effect.Effect<StreamState> => {
796 const native = chunk.response;
797 if (native == null || native === "") return Effect.succeed(state);
798 const text =
799 typeof native === "object" ? JSON.stringify(native) : String(native);
800 if (text.length === 0) return Effect.succeed(state);
801 // When tools were requested, the native `response` stream is the
802 // tool-call JSON, not prose — buffer it and decide on finalize. We
803 // pre-allocate the tool id here (we're in an Effect, finalize is sync).
804 if (hasTools) {
805 return Effect.gen(function* () {
806 const id = state.nativeToolId ?? (yield* idGen.generateId());
807 return {
808 ...state,
809 nativeToolId: id,
810 nativeToolBuffer: state.nativeToolBuffer + text,
811 };
812 });
813 }
814 return emitTextDelta(state, text, parts, idGen);
815};
816
817const handleNativeToolCalls = (
818 state: StreamState,
819 chunk: Record<string, unknown>,
820 parts: StreamParts,
821 idGen: IdGenerator.Service,
822): Effect.Effect<StreamState> => {
823 const openAiToolCalls = (
824 chunk.choices as
825 | Array<{ delta?: { tool_calls?: ReadonlyArray<unknown> } }>
826 | undefined
827 )?.[0]?.delta?.tool_calls;
828 // Workers AI can mirror the same fragments in both its top-level native
829 // field and the OpenAI-compatible delta. Prefer the latter so every
830 // fragment is emitted exactly once.
831 if (Array.isArray(openAiToolCalls)) return Effect.succeed(state);
832 if (!Array.isArray(chunk.tool_calls)) return Effect.succeed(state);
833 return Effect.gen(function* () {
834 const s = closeReasoning(state, parts);
835 return yield* handleToolDeltas(
836 s,
837 chunk.tool_calls as ReadonlyArray<Record<string, unknown>>,
838 parts,
839 idGen,
840 );
841 });
842};
843
844const handleOpenAiDelta = (
845 state: StreamState,
846 chunk: Record<string, unknown>,
847 parts: StreamParts,
848 idGen: IdGenerator.Service,
849): Effect.Effect<StreamState> => {
850 const delta = (
851 chunk.choices as Array<{ delta?: Record<string, unknown> }> | undefined
852 )?.[0]?.delta;
853 if (!delta) return Effect.succeed(state);
854 return Effect.gen(function* () {
855 let s = state;
856 const reasoning = (delta.reasoning_content ?? delta.reasoning) as
857 | string
858 | undefined;
859 if (reasoning && reasoning.length > 0) {
860 s = yield* emitReasoningDelta(s, reasoning, parts, idGen);
861 }
862 const text = delta.content as string | undefined;
863 if (text && text.length > 0) {
864 s = yield* emitTextDelta(s, text, parts, idGen);
865 }
866 const toolDeltas = delta.tool_calls as
867 | ReadonlyArray<Record<string, unknown>>
868 | undefined;
869 if (Array.isArray(toolDeltas)) {
870 s = closeReasoning(s, parts);
871 s = yield* handleToolDeltas(s, toolDeltas, parts, idGen);
872 }
873 return s;
874 });
875};
876
877const handleStreamChunk = (
878 state: StreamState,
879 data: string,
880 idGen: IdGenerator.Service,
881 hasTools: boolean,
882): Effect.Effect<
883 readonly [StreamState, ReadonlyArray<Response.StreamPartEncoded>]
884> =>
885 Effect.gen(function* () {
886 if (data === "") return [state, []] as const;
887 if (data === "[DONE]") {
888 return [{ ...state, receivedDone: true }, []] as const;
889 }
890 const chunk = tryParseJson(data);
891 if (chunk === undefined) return [state, []] as const;
892
893 const parts: StreamParts = [];
894 let s: StreamState = { ...state, receivedAnyData: true };
895 s = updateChunkMeta(s, chunk);
896 s = yield* handleNativeText(s, chunk, parts, idGen, hasTools);
897 s = yield* handleNativeToolCalls(s, chunk, parts, idGen);
898 s = yield* handleOpenAiDelta(s, chunk, parts, idGen);
899 return [s, parts] as const;
900 });
901
902/**
903 * Decode a buffered native-streaming tool call. Workers AI's native shape
904 * for a function call is a JSON document with `name` plus `parameters` (or
905 * `arguments`), optionally wrapped in an array. Returns one entry per call,
906 * or `undefined` if the buffer isn't a tool-call document (so the caller
907 * can fall back to treating it as plain text).
908 */
909const decodeNativeToolBuffer = (
910 buffer: string,
911): ReadonlyArray<{ name: string; args: string }> | undefined => {
912 const trimmed = buffer.trim();
913 if (trimmed.length === 0 || !(trimmed[0] === "{" || trimmed[0] === "[")) {
914 return undefined;
915 }
916 let parsed: unknown;
917 try {
918 parsed = JSON.parse(trimmed);
919 } catch {
920 return undefined;
921 }
922 const items = Array.isArray(parsed) ? parsed : [parsed];
923 const calls = items.flatMap(
924 (
925 item,
926 ): ReadonlyArray<{
927 name: string;
928 args: string;
929 }> => {
930 if (item == null || typeof item !== "object") return [];
931 const rec = item as Record<string, unknown>;
932 const fn = rec.function as Record<string, unknown> | undefined;
933 const name = (fn?.name ?? rec.name) as string | undefined;
934 if (typeof name !== "string" || name.length === 0) return [];
935 const rawArgs = fn?.arguments ?? rec.arguments ?? rec.parameters ?? {};
936 const args =
937 typeof rawArgs === "string" ? rawArgs : JSON.stringify(rawArgs);
938 return [{ name, args }];
939 },
940 );
941 return calls.length > 0 ? calls : undefined;
942};
943
944const flushNativeToolBuffer = (
945 state: StreamState,
946 parts: StreamParts,
947): StreamState => {
948 if (state.nativeToolBuffer.length === 0) return state;
949 const calls = decodeNativeToolBuffer(state.nativeToolBuffer);
950 // Not a tool-call document — the model returned prose despite tools being
951 // available; surface it as text so nothing is dropped.
952 if (calls === undefined) {
953 const id = state.nativeToolId ?? "text-0";
954 parts.push({ type: "text-start", id });
955 parts.push({ type: "text-delta", id, delta: state.nativeToolBuffer });
956 parts.push({ type: "text-end", id });
957 return { ...state, nativeToolBuffer: "", nativeToolId: undefined };
958 }
959 const baseId = state.nativeToolId ?? "tool-0";
960 calls.forEach((call, i) => {
961 const id = i === 0 ? baseId : `${baseId}-${i}`;
962 parts.push({ type: "tool-params-start", id, name: call.name });
963 if (call.args.length > 0) {
964 parts.push({ type: "tool-params-delta", id, delta: call.args });
965 }
966 parts.push({ type: "tool-params-end", id });
967 });
968 return { ...state, nativeToolBuffer: "", nativeToolId: undefined };
969};
970
971const finalizeStream = (
972 state: StreamState,
973): ReadonlyArray<Response.StreamPartEncoded> => {
974 const parts: StreamParts = [];
975 let s = state;
976 for (const [idx] of s.toolCalls) s = closeToolCall(s, idx, parts);
977 s = closeReasoning(s, parts);
978 if (s.textId !== undefined) parts.push({ type: "text-end", id: s.textId });
979 // Emit any buffered native-streaming tool call (llama-style) before the
980 // terminal finish part.
981 s = flushNativeToolBuffer(s, parts);
982
983 // Three cases for the final reason:
984 // 1. The model emitted an explicit `finish_reason` → map it.
985 // 2. The stream ended cleanly (`[DONE]` seen) but no reason → "stop".
986 // Workers AI's native shape never includes `finish_reason`,
987 // so without this rule every native-mode stream would report
988 // `unknown` despite completing successfully.
989 // 3. The stream ended abnormally (no `[DONE]`, no reason) → "error".
990 const reason: Response.FinishReason =
991 s.finishReason !== undefined
992 ? mapFinishReason(s.finishReason)
993 : s.receivedDone
994 ? "stop"
995 : s.receivedAnyData
996 ? "error"
997 : "unknown";
998
999 parts.push({
1000 type: "finish",
1001 reason,
1002 usage: mapUsage(s.usage),
1003 response: undefined,
1004 });
1005 return parts;
1006};
1007
1008const parseStreamText = (
1009 resp: Response,
1010 idGen: IdGenerator.Service,
1011 hasTools: boolean,
1012): Stream.Stream<Response.StreamPartEncoded, AiError.AiError> => {
1013 const body = resp.body;
1014 if (body === null) {
1015 return Stream.fromIterable<Response.StreamPartEncoded>(
1016 finalizeStream(initialStreamState()),
1017 );
1018 }
1019 return Stream.fromReadableStream<Uint8Array, AiError.AiError>({
1020 evaluate: () => body,
1021 onError: (cause) => toAiError(cause, "streamText"),
1022 }).pipe(
1023 Stream.decodeText(),
1024 Stream.pipeThroughChannel(Sse.decode<AiError.AiError, unknown>()),
1025 Stream.catchTag("Retry", (retry) => Stream.die(retry)),
1026 Stream.catchTag("SseError", (error) =>
1027 Stream.fail(toAiError(error, "streamText")),
1028 ),
1029 Stream.mapAccumEffect(
1030 initialStreamState,
1031 (state, event) => handleStreamChunk(state, event.data, idGen, hasTools),
1032 { onHalt: (state) => finalizeStream(state) },
1033 ),
1034 );
1035};
1036
1037// ---------------------------------------------------------------------------
1038// Error mapping
1039// ---------------------------------------------------------------------------
1040
1041const toAiError = (
1042 cause: unknown,
1043 method: "generateText" | "streamText",
1044): AiError.AiError =>
1045 AiError.AiError.make({
1046 module: "Cloudflare.AI.LanguageModel",
1047 method,
1048 reason: new AiError.UnknownError({
1049 description:
1050 cause instanceof Error ? cause.message : "AI Gateway request failed",
1051 }),
1052 });