@prometheus-ai/ai 0.5.4 → 0.5.8

package/src/providers/openai-responses-server-schema.ts

@@ -7,6 +7,8 @@
  * Real clients (codex, openai-python, llm-git) routinely send these and a 400
  * is a worse outcome than dropping them on the floor.
  */
+
+import { z } from "zod/v4";
 import type {
 	EasyInputMessage,
 	ResponseCreateParams,
@@ -16,8 +18,7 @@ import type {
 	ResponseOutputMessage,
 	ResponseReasoningItem,
 	Tool as ResponsesTool,
-} from "openai/resources/responses/responses";
-import * as z from "zod/v4";
+} from "./openai-responses-wire";
 
 // ─── Input content blocks ───────────────────────────────────────────────────
 
@@ -97,12 +98,17 @@ const assistantMessageItemSchema = z.object({
 	content: z.union([z.string(), z.array(outputContentBlockSchema)]).optional(),
 });
 
-const reasoningItemSchema = z.object({
-	type: z.literal("reasoning"),
-	id: z.string().optional(),
-	summary: z.array(summaryTextSchema).optional(),
-	content: z.array(reasoningTextSchema).optional(),
-});
+const reasoningItemSchema = z
+	.object({
+		type: z.literal("reasoning"),
+		id: z.string().optional(),
+		summary: z.array(summaryTextSchema).optional(),
+		content: z.array(reasoningTextSchema).optional(),
+	})
+	// Loose: unknown keys like `encrypted_content` must survive the parse —
+	// the outbound encoder replays them verbatim (buildReasoningItem spreads
+	// the persisted item to preserve encrypted reasoning round-trips).
+	.loose();
 
 const functionCallItemSchema = z.object({
 	type: z.literal("function_call"),
@@ -234,7 +240,7 @@ export const toolChoiceSchema = z.union([
 export const reasoningConfigSchema = z.object({
 	effort: z.string().optional(),
 	// `none` maps to hideThinkingSummary; auto/concise/detailed mean "show
-	// summary". @prometheus-ai/ai has no per-level plumbing for the latter — walker logs
+	// summary". Prometheus AI has no per-level plumbing for the latter — walker logs
 	// once and treats them as default.
 	summary: z.enum(["auto", "concise", "detailed", "none"]).optional(),
 });

package/src/providers/openai-responses-server.ts

@@ -11,7 +11,7 @@
 
 import { logger } from "@prometheus-ai/utils";
 import { resolvePromptCacheKey } from "../auth-gateway/http";
-import type { AuthGatewayParsedRequest as ParsedRequest } from "../auth-gateway/types";
+import type { AuthGatewayStreamControl, AuthGatewayParsedRequest as ParsedRequest } from "../auth-gateway/types";
 import type {
 	AssistantMessage,
 	AssistantMessageEventStream,
@@ -102,9 +102,9 @@ type InputBlockUnion =
 	| { type: "input_file"; file_id?: string; filename?: string; file_data?: string };
 
 /**
- * Walk an input message's content array and produce @prometheus-ai/ai's `TextContent[]`.
+ * Walk an input message's content array and produce Prometheus AI's `TextContent[]`.
  * `input_image`/`input_file` blocks become bracketed text placeholders since
- * @prometheus-ai/ai's `ImageContent` only carries inline base64 data and we have no
+ * Prometheus AI's `ImageContent` only carries inline base64 data and we have no
  * resolver for OpenAI `image_url` / `file_id` references. Logs once per kind.
  */
 function inputContentParts(blocks: OpenAIResponsesInputContent[] | string | undefined): string | TextContent[] {
@@ -118,17 +118,20 @@ function inputContentParts(blocks: OpenAIResponsesInputContent[] | string | unde
 		} else if (block.type === "input_image") {
 			if (!warnedImageNotSupported) {
 				warnedImageNotSupported = true;
-				logger.warn("openai-responses-server: input_image dropped (no canonical bridge for image_url/file_id)", {
-					hasUrl: typeof block.image_url === "string",
-					hasFileId: typeof block.file_id === "string",
-				});
+				logger.warn(
+					"openai-responses-server: input_image dropped (no Prometheus AI bridge for image_url/file_id)",
+					{
+						hasUrl: typeof block.image_url === "string",
+						hasFileId: typeof block.file_id === "string",
+					},
+				);
 			}
 			const ref = block.image_url ?? block.file_id ?? "?";
 			parts.push({ type: "text", text: `[image: ${ref}]` });
 		} else if (block.type === "input_file") {
 			if (!warnedFileNotSupported) {
 				warnedFileNotSupported = true;
-				logger.warn("openai-responses-server: input_file dropped (no canonical bridge for file_id/file_data)", {
+				logger.warn("openai-responses-server: input_file dropped (no Prometheus AI bridge for file_id/file_data)", {
 					hasFileId: typeof block.file_id === "string",
 					hasFileData: typeof block.file_data === "string",
 				});
@@ -185,10 +188,10 @@ function mapToolChoice(value: ParsedToolChoice | undefined): ParsedRequest["opti
 	if (value === "auto" || value === "none" || value === "required") return value;
 	if ("type" in value) {
 		// `custom` (codex apply_patch) and `function` both resolve to the same
-		// @prometheus-ai/ai shape: @prometheus-ai/ai's dispatcher matches `Tool.name` AND `customWireName`,
+		// Prometheus AI shape: Prometheus AI's dispatcher matches `Tool.name` AND `customWireName`,
 		// so passing the wire name works for either.
 		if (value.type === "function" || value.type === "custom") return { name: value.name };
-		// Hosted tools + allowed_tools — we don't surface these to @prometheus-ai/ai; fall
+		// Hosted tools + allowed_tools — we don't surface these to Prometheus AI; fall
 		// back to letting the model pick a tool freely.
 		return "auto";
 	}
@@ -360,7 +363,7 @@ export function parseRequest(body: unknown, headers?: Headers): ParsedRequest {
 			if (effectiveType === "custom_tool_call") {
 				const call = item as { id?: string; call_id: string; name: string; input: string };
 				// Custom tools carry a raw input string. We stash it in `arguments.input`
-				// matching @prometheus-ai/ai's openai-responses-shared convention, and tag the call
+				// matching Prometheus AI's openai-responses-shared convention, and tag the call
 				// with `customWireName` so encoders re-emit it as `custom_tool_call`.
 				const toolCall: ToolCall = {
 					type: "toolCall",
@@ -428,7 +431,7 @@ export function parseRequest(body: unknown, headers?: Headers): ParsedRequest {
 		options.reasoning = data.reasoning.effort;
 	}
 	// OpenAI summary: `none` → suppress; `auto`/`concise`/`detailed` → request
-	// visible summary. @prometheus-ai/ai has no per-level plumbing — log once and let the
+	// visible summary. Prometheus AI has no per-level plumbing — log once and let the
 	// provider default kick in.
 	if (data.reasoning?.summary === "none") {
 		options.hideThinkingSummary = true;
@@ -573,6 +576,17 @@ function reasoningItemId(part: ThinkingContent): string {
 	return makeReasoningId();
 }
 
+/**
+ * Prometheus AI responses providers mint composite `"{call_id}|{item_id}"` tool-call
+ * ids ({@link encodeResponsesToolCallId}). Only the call_id half belongs on
+ * the wire: third-party clients validate the `call_id` charset
+ * (`^[a-zA-Z0-9_-]+$`) or echo it to other backends, and `|` fails both.
+ */
+function wireCallId(id: string): string {
+	const sep = id.indexOf("|");
+	return sep >= 0 ? id.slice(0, sep) : id;
+}
+
 /**
  * Walk the assistant content array and group consecutive TextContent into a
  * single message item; each ThinkingContent / ToolCall is its own item.
@@ -609,7 +623,7 @@ function buildOutputItems(message: AssistantMessage): OutputItem[] {
 				out.push({
 					type: "custom_tool_call",
 					id: part.thoughtSignature ?? makeCustomCallId(),
-					call_id: part.id,
+					call_id: wireCallId(part.id),
 					name: part.customWireName,
 					input: rawInput,
 					status: "completed",
@@ -618,7 +632,7 @@ function buildOutputItems(message: AssistantMessage): OutputItem[] {
 				out.push({
 					type: "function_call",
 					id: part.thoughtSignature ?? makeFuncCallId(),
-					call_id: part.id,
+					call_id: wireCallId(part.id),
 					name: part.name,
 					arguments: JSON.stringify(part.arguments ?? {}),
 					status: "completed",
@@ -698,6 +712,7 @@ interface OpenFunctionCall {
 	kind: "function_call";
 	itemId: string;
 	outputIndex: number;
+	contentIndex: number;
 	callId: string;
 	name: string;
 	argsText: string;
@@ -713,23 +728,35 @@ function sseEvent(name: string, data: unknown): string {
 export function encodeStream(
 	events: AssistantMessageEventStream,
 	requestedModelId: string,
+	_options?: ParsedRequest["options"],
+	control?: AuthGatewayStreamControl,
 ): ReadableStream<Uint8Array> {
 	const encoder = new TextEncoder();
 	const responseId = makeRespId();
 	let sequenceNumber = 0;
+	let cancelled = control?.signal?.aborted === true;
+	const markCancelled = () => {
+		cancelled = true;
+	};
+	control?.signal?.addEventListener("abort", markCancelled, { once: true });
 	const seq = () => sequenceNumber++;
 
 	return new ReadableStream<Uint8Array>({
 		async start(controller) {
 			const emit = (name: string, data: Record<string, unknown>) => {
-				controller.enqueue(encoder.encode(sseEvent(name, { type: name, sequence_number: seq(), ...data })));
+				if (!cancelled)
+					controller.enqueue(encoder.encode(sseEvent(name, { type: name, sequence_number: seq(), ...data })));
+			};
+			const emitDone = () => {
+				if (!cancelled) controller.enqueue(encoder.encode("data: [DONE]\n\n"));
 			};
-			const emitDone = () => controller.enqueue(encoder.encode("data: [DONE]\n\n"));
 
 			let createdAt = Math.floor(Date.now() / 1000);
 			let outputIndex = 0;
 			const state: { open: OpenItem | null } = { open: null };
+			const openFunctionCalls = new Map<number, OpenFunctionCall>();
 			const finishedItems: OutputItem[] = [];
+			const allocateOutputIndex = (): number => outputIndex++;
 
 			const responseSnapshot = (status: ResponseStatus, output: OutputItem[] | []) => ({
 				id: responseId,
@@ -742,6 +769,7 @@ export function encodeStream(
 			});
 
 			const openMessage = (): OpenMessage => {
+				const itemOutputIndex = allocateOutputIndex();
 				const itemId = makeMsgId();
 				const item = {
 					type: "message" as const,
@@ -750,11 +778,11 @@ export function encodeStream(
 					role: "assistant" as const,
 					content: [] as Array<{ type: "output_text"; text: string; annotations: never[] }>,
 				};
-				emit("response.output_item.added", { output_index: outputIndex, item });
+				emit("response.output_item.added", { output_index: itemOutputIndex, item });
 				const next: OpenMessage = {
 					kind: "message",
 					itemId,
-					outputIndex,
+					outputIndex: itemOutputIndex,
 					contentIndex: 0,
 					currentPartText: "",
 					content: [],
@@ -764,6 +792,7 @@ export function encodeStream(
 			};
 
 			const openReasoning = (partial: AssistantMessage, contentIndex: number): OpenReasoning => {
+				const itemOutputIndex = allocateOutputIndex();
 				const part = partial.content[contentIndex];
 				const itemId = part && part.type === "thinking" ? reasoningItemId(part) : makeReasoningId();
 				const item = {
@@ -771,22 +800,23 @@ export function encodeStream(
 					id: itemId,
 					summary: [] as Array<{ type: "summary_text"; text: string }>,
 				};
-				emit("response.output_item.added", { output_index: outputIndex, item });
+				emit("response.output_item.added", { output_index: itemOutputIndex, item });
 				// Open the summary part. Real OpenAI streams summary text in the
-				// canonical `reasoning_summary_*` lifecycle; @prometheus-ai/ai's own decoder
+				// canonical `reasoning_summary_*` lifecycle; Prometheus AI's own decoder
 				// reads `summary[].text` from the eventual `output_item.done`.
 				emit("response.reasoning_summary_part.added", {
 					item_id: itemId,
-					output_index: outputIndex,
+					output_index: itemOutputIndex,
 					summary_index: 0,
 					part: { type: "summary_text", text: "" },
 				});
-				const next: OpenReasoning = { kind: "reasoning", itemId, outputIndex, reasoningText: "" };
+				const next: OpenReasoning = { kind: "reasoning", itemId, outputIndex: itemOutputIndex, reasoningText: "" };
 				state.open = next;
 				return next;
 			};
 
 			const openToolCall = (partial: AssistantMessage, contentIndex: number): OpenFunctionCall => {
+				const itemOutputIndex = allocateOutputIndex();
 				const part = partial.content[contentIndex];
 				const tc = part && part.type === "toolCall" ? part : undefined;
 				const customWireName: string | undefined =
@@ -795,7 +825,7 @@ export function encodeStream(
 						: undefined;
 				const isCustom = customWireName !== undefined;
 				const itemId = tc?.thoughtSignature ?? (isCustom ? makeCustomCallId() : makeFuncCallId());
-				const callId = tc?.id ?? "";
+				const callId = wireCallId(tc?.id ?? "");
 				const name = customWireName ?? tc?.name ?? "";
 				const item = isCustom
 					? {
@@ -814,20 +844,65 @@ export function encodeStream(
 							arguments: "",
 							status: "in_progress",
 						};
-				emit("response.output_item.added", { output_index: outputIndex, item });
+				emit("response.output_item.added", { output_index: itemOutputIndex, item });
 				const next: OpenFunctionCall = {
 					kind: "function_call",
 					itemId,
-					outputIndex,
+					outputIndex: itemOutputIndex,
+					contentIndex,
 					callId,
 					name,
 					argsText: "",
 					...(isCustom ? { customWireName } : {}),
 				};
+				openFunctionCalls.set(contentIndex, next);
 				state.open = next;
 				return next;
 			};
 
+			const closeFunctionCall = (call: OpenFunctionCall): void => {
+				const text = call.argsText ?? "";
+				if (call.customWireName) {
+					const item = {
+						type: "custom_tool_call",
+						id: call.itemId,
+						call_id: call.callId ?? "",
+						name: call.customWireName,
+						input: text,
+						status: "completed",
+					};
+					emit("response.output_item.done", { output_index: call.outputIndex, item });
+					finishedItems.push({
+						type: "custom_tool_call",
+						id: call.itemId,
+						call_id: call.callId ?? "",
+						name: call.customWireName,
+						input: text,
+						status: "completed",
+					});
+				} else {
+					const item = {
+						type: "function_call",
+						id: call.itemId,
+						call_id: call.callId ?? "",
+						name: call.name ?? "",
+						arguments: text,
+						status: "completed",
+					};
+					emit("response.output_item.done", { output_index: call.outputIndex, item });
+					finishedItems.push({
+						type: "function_call",
+						id: call.itemId,
+						call_id: call.callId ?? "",
+						name: call.name ?? "",
+						arguments: text,
+						status: "completed",
+					});
+				}
+				openFunctionCalls.delete(call.contentIndex);
+				if (state.open === call) state.open = null;
+			};
+
 			const closeOpen = () => {
 				if (!state.open) return;
 				if (state.open.kind === "message") {
@@ -846,6 +921,7 @@ export function encodeStream(
 						status: "completed",
 						content: state.open.content,
 					});
+					state.open = null;
 				} else if (state.open.kind === "reasoning") {
 					const summary = [{ type: "summary_text" as const, text: state.open.reasoningText ?? "" }];
 					const item = {
@@ -859,79 +935,40 @@ export function encodeStream(
 						id: state.open.itemId,
 						summary,
 					});
+					state.open = null;
 				} else {
-					const text = state.open.argsText ?? "";
-					if (state.open.customWireName) {
-						const item = {
-							type: "custom_tool_call",
-							id: state.open.itemId,
-							call_id: state.open.callId ?? "",
-							name: state.open.customWireName,
-							input: text,
-							status: "completed",
-						};
-						emit("response.output_item.done", { output_index: state.open.outputIndex, item });
-						finishedItems.push({
-							type: "custom_tool_call",
-							id: state.open.itemId,
-							call_id: state.open.callId ?? "",
-							name: state.open.customWireName,
-							input: text,
-							status: "completed",
-						});
-					} else {
-						const item = {
-							type: "function_call",
-							id: state.open.itemId,
-							call_id: state.open.callId ?? "",
-							name: state.open.name ?? "",
-							arguments: text,
-							status: "completed",
-						};
-						emit("response.output_item.done", { output_index: state.open.outputIndex, item });
-						finishedItems.push({
-							type: "function_call",
-							id: state.open.itemId,
-							call_id: state.open.callId ?? "",
-							name: state.open.name ?? "",
-							arguments: text,
-							status: "completed",
-						});
-					}
+					closeFunctionCall(state.open);
 				}
-				outputIndex++;
-				state.open = null;
 			};
 
-			try {
-				let finalMessage: AssistantMessage | null = null;
-				let failureMessage: AssistantMessage | null = null;
+			const closeOpenFunctionCalls = (): void => {
+				for (const call of [...openFunctionCalls.values()]) {
+					closeFunctionCall(call);
+				}
+			};
 
+			const functionCallForEvent = (contentIndex: number): OpenFunctionCall | undefined => {
+				const byIndex = openFunctionCalls.get(contentIndex);
+				if (byIndex) return byIndex;
+				return state.open?.kind === "function_call" ? state.open : undefined;
+			};
+			let finalMessage: AssistantMessage | undefined;
+			let failureMessage: AssistantMessage | undefined;
+			try {
+				if (cancelled) {
+					controller.close();
+					return;
+				}
 				for await (const ev of events) {
+					if (cancelled) return;
 					switch (ev.type) {
 						case "start": {
 							createdAt = Math.floor((ev.partial.timestamp || Date.now()) / 1000);
 							// response.created — initial envelope.
-							controller.enqueue(
-								encoder.encode(
-									sseEvent("response.created", {
-										type: "response.created",
-										sequence_number: seq(),
-										response: responseSnapshot("in_progress", []),
-									}),
-								),
-							);
+							emit("response.created", { response: responseSnapshot("in_progress", []) });
 							// response.in_progress — mirrors real OpenAI; some clients gate
 							// on it before reading items.
-							controller.enqueue(
-								encoder.encode(
-									sseEvent("response.in_progress", {
-										type: "response.in_progress",
-										sequence_number: seq(),
-										response: responseSnapshot("in_progress", []),
-									}),
-								),
-							);
+							emit("response.in_progress", { response: responseSnapshot("in_progress", []) });
 							break;
 						}
 						case "text_start": {
@@ -941,7 +978,7 @@ export function encodeStream(
 								cur = state.open;
 								cur.currentPartText = "";
 							} else {
-								if (state.open) closeOpen();
+								if (state.open && state.open.kind !== "function_call") closeOpen();
 								cur = openMessage();
 							}
 							const part = { type: "output_text", text: "", annotations: [] as never[] };
@@ -964,7 +1001,7 @@ export function encodeStream(
 								delta: ev.delta,
 								logprobs: [],
 							});
-							// TODO: when @prometheus-ai/ai surfaces output_text annotations
+							// TODO: when Prometheus AI surfaces output_text annotations
 							// (web_search citations, …), emit
 							// `response.output_text.annotation.added` here.
 							break;
@@ -992,7 +1029,7 @@ export function encodeStream(
 							break;
 						}
 						case "thinking_start": {
-							if (state.open) closeOpen();
+							if (state.open && state.open.kind !== "function_call") closeOpen();
 							openReasoning(ev.partial, ev.contentIndex);
 							break;
 						}
@@ -1029,13 +1066,13 @@ export function encodeStream(
 							break;
 						}
 						case "toolcall_start": {
-							if (state.open) closeOpen();
+							if (state.open && state.open.kind !== "function_call") closeOpen();
 							openToolCall(ev.partial, ev.contentIndex);
 							break;
 						}
 						case "toolcall_delta": {
-							if (state.open?.kind !== "function_call") break;
-							const cur: OpenFunctionCall = state.open;
+							const cur = functionCallForEvent(ev.contentIndex);
+							if (!cur) break;
 							cur.argsText += ev.delta;
 							if (cur.customWireName) {
 								emit("response.custom_tool_call_input.delta", {
@@ -1053,8 +1090,8 @@ export function encodeStream(
 							break;
 						}
 						case "toolcall_end": {
-							if (state.open?.kind !== "function_call") break;
-							const cur: OpenFunctionCall = state.open;
+							const cur = functionCallForEvent(ev.contentIndex);
+							if (!cur) break;
 							// Promote possibly-late info from the canonical ToolCall.
 							const tc = ev.toolCall;
 							if (tc.customWireName && !cur.customWireName) cur.customWireName = tc.customWireName;
@@ -1087,7 +1124,7 @@ export function encodeStream(
 									name: cur.name,
 								});
 							}
-							closeOpen();
+							closeFunctionCall(cur);
 							break;
 						}
 						case "done": {
@@ -1102,6 +1139,7 @@ export function encodeStream(
 				}
 
 				if (failureMessage) {
+					closeOpenFunctionCalls();
 					if (state.open) closeOpen();
 					controller.enqueue(
 						encoder.encode(
@@ -1120,6 +1158,7 @@ export function encodeStream(
 					return;
 				}
 
+				closeOpenFunctionCalls();
 				if (state.open) closeOpen();
 				const message = finalMessage ?? ((await events.result().catch(() => null)) as AssistantMessage | null);
 
@@ -1158,26 +1197,35 @@ export function encodeStream(
 				emitDone();
 				controller.close();
 			} catch (err) {
-				controller.enqueue(
-					encoder.encode(
-						sseEvent("response.failed", {
-							type: "response.failed",
-							sequence_number: seq(),
-							response: {
-								id: responseId,
-								object: "response",
-								created_at: Math.floor(Date.now() / 1000),
-								status: "failed",
-								model: requestedModelId,
-								output: [],
-								error: { message: err instanceof Error ? err.message : String(err) },
-							},
-						}),
-					),
-				);
-				emitDone();
-				controller.close();
+				if (!cancelled) {
+					controller.enqueue(
+						encoder.encode(
+							sseEvent("response.failed", {
+								type: "response.failed",
+								sequence_number: seq(),
+								response: {
+									id: responseId,
+									object: "response",
+									created_at: Math.floor(Date.now() / 1000),
+									status: "failed",
+									model: requestedModelId,
+									output: [],
+									error: { message: err instanceof Error ? err.message : String(err) },
+								},
+							}),
+						),
+					);
+					emitDone();
+					controller.close();
+				}
+			} finally {
+				control?.signal?.removeEventListener("abort", markCancelled);
 			}
 		},
+		cancel(reason) {
+			cancelled = true;
+			control?.signal?.removeEventListener("abort", markCancelled);
+			control?.onCancel?.(reason);
+		},
 	});
 }