@linxiraos/pi-ai 1.0.0

package/src/providers/pi-native-client.ts

@@ -0,0 +1,275 @@
+/**
+ * Client half of the pi-native auth-gateway protocol.
+ *
+ * Dispatches a {@link streamSimple}-shaped request to an `omp auth-gateway`
+ * via `POST /v1/pi/stream`, reads the SSE event stream back, and pushes the
+ * parsed events into a local {@link AssistantMessageEventStream} — the same
+ * stream type every other provider client produces. Callers downstream of
+ * `streamSimple` cannot tell whether the events came from a real provider
+ * SDK or from a gateway hop; they consume `AssistantMessageEvent`s either
+ * way.
+ *
+ * Activated when a {@link Model} has `transport: "pi-native"` set; the
+ * dispatch hook lives in `streamSimple()` (see `../stream.ts`). Used by
+ * containerized omp deployments (such as robomp slots) that
+ * route every LLM call through a credential-holding sidecar so the slot
+ * itself stays credential-free.
+ */
+import { readSseJson } from "@linxiraos/pi-utils";
+import * as AIError from "../error";
+import type {
+	Api,
+	AssistantMessage,
+	AssistantMessageEvent,
+	AssistantMessageEventStream as AssistantMessageEventStreamType,
+	Context,
+	Model,
+	SimpleStreamOptions,
+} from "../types";
+import { createAbortSourceTracker } from "../utils/abort";
+import { AssistantMessageEventStream } from "../utils/event-stream";
+import { getStreamFirstEventTimeoutMs, getStreamIdleTimeoutMs, iterateWithIdleTimeout } from "../utils/idle-iterator";
+import { notifyProviderResponse } from "../utils/provider-response";
+
+/**
+ * Fields that must not cross the wire — either non-serializable (functions,
+ * `AbortSignal`, the provider-session `Map`) or server-controlled
+ * (`apiKey`, which the gateway injects from its own credential store; the
+ * client's `apiKey` is the gateway *bearer*, sent in the `Authorization`
+ * header rather than the request body).
+ */
+const NON_WIRE_KEYS = new Set<keyof SimpleStreamOptions>([
+	"signal",
+	"apiKey",
+	"fetch",
+	"onPayload",
+	"onResponse",
+	"onSseEvent",
+	"execHandlers",
+	"cursorExecHandlers",
+	"cursorOnToolResult",
+	"providerSessionState",
+]);
+const PI_NATIVE_STREAM_IDLE_TIMEOUT_ERROR = "pi-native stream stalled while waiting for the next event";
+const PI_NATIVE_STREAM_FIRST_EVENT_TIMEOUT_ERROR = "pi-native stream timed out while waiting for the first event";
+
+function isPiNativeProgressEvent(event: unknown): boolean {
+	if (typeof event !== "object" || event === null || !("type" in event)) return true;
+	return event.type !== "start";
+}
+
+function buildWireOptions(options: SimpleStreamOptions | undefined): Record<string, unknown> {
+	if (!options) return {};
+	const wire: Record<string, unknown> = {};
+	for (const [k, v] of Object.entries(options)) {
+		if (v === undefined) continue;
+		if (NON_WIRE_KEYS.has(k as keyof SimpleStreamOptions)) continue;
+		wire[k] = v;
+	}
+	return wire;
+}
+
+async function decodeGatewayError(response: Response): Promise<AIError.AuthGatewayError> {
+	const status = response.status;
+	let body: unknown;
+	try {
+		body = await response.json();
+	} catch {
+		body = await response.text().catch(() => "");
+	}
+	if (typeof body === "object" && body !== null && "error" in body) {
+		const err = (body as { error: unknown }).error;
+		if (typeof err === "object" && err !== null) {
+			const message = (err as { message?: unknown }).message;
+			const type = (err as { type?: unknown }).type;
+			return new AIError.AuthGatewayError(
+				typeof message === "string" ? message : `auth-gateway ${status}`,
+				status,
+				response.headers,
+				typeof type === "string" ? type : undefined,
+			);
+		}
+	}
+	const text = typeof body === "string" ? body : JSON.stringify(body);
+	return new AIError.AuthGatewayError(
+		`auth-gateway ${status}: ${text || response.statusText}`,
+		status,
+		response.headers,
+	);
+}
+
+/**
+ * Resolve the `/v1/pi/stream` endpoint URL from the model's `baseUrl`.
+ * Trims a trailing slash so concatenation can't double-slash; throws when
+ * the baseUrl is missing (transport=pi-native without a gateway target is
+ * a configuration error, not a runtime recoverable one).
+ */
+function resolveStreamUrl(model: Model<Api>): string {
+	if (!model.baseUrl) {
+		throw new AIError.ConfigurationError(
+			`pi-native transport requires \`baseUrl\` on model ${model.id} (set it on the provider config in models.yml)`,
+		);
+	}
+	return `${model.baseUrl.replace(/\/+$/, "")}/v1/pi/stream`;
+}
+
+function buildHeaders(model: Model<Api>, apiKey: string | undefined): Record<string, string> {
+	const headers: Record<string, string> = {
+		"Content-Type": "application/json",
+		Accept: "text/event-stream",
+		...(model.headers ?? {}),
+	};
+	if (apiKey && !headers.Authorization) {
+		headers.Authorization = `Bearer ${apiKey}`;
+	}
+	return headers;
+}
+
+/**
+ * Stream a turn through an `omp auth-gateway` over the pi-native protocol.
+ *
+ * The returned {@link AssistantMessageEventStream} receives each parsed
+ * `AssistantMessageEvent` verbatim from the gateway; the terminal `done` /
+ * `error` event resolves `.result()` automatically via the base class's
+ * completion check. Non-streaming consumers just call `.result()` and pay
+ * for SSE framing they don't use — that overhead is dominated by provider
+ * latency, so we always stream rather than maintaining a parallel
+ * non-streaming path.
+ */
+export function streamPiNative<TApi extends Api>(
+	model: Model<TApi>,
+	context: Context,
+	options?: SimpleStreamOptions,
+): AssistantMessageEventStreamType {
+	const stream = new AssistantMessageEventStream();
+
+	void (async () => {
+		const callerSignal = options?.signal;
+		const abortTracker = createAbortSourceTracker(callerSignal);
+		// Abort propagation: cancel the response body when the caller's signal
+		// fires. Mirror `streamProxy`'s shape — explicit listener + finally
+		// cleanup — so we don't leak listeners on the long-running case.
+		let response: Response | null = null;
+		const onAbort = (): void => {
+			const body = response?.body;
+			if (body) body.cancel("Request aborted by caller").catch(() => {});
+		};
+		if (callerSignal) {
+			if (callerSignal.aborted) {
+				stream.fail(
+					callerSignal.reason instanceof Error
+						? callerSignal.reason
+						: new Error(String(callerSignal.reason ?? "aborted")),
+				);
+				return;
+			}
+			callerSignal.addEventListener("abort", onAbort, { once: true });
+		}
+
+		try {
+			const url = resolveStreamUrl(model as Model<Api>);
+			const fetchImpl = options?.fetch ?? globalThis.fetch;
+			const headers = buildHeaders(
+				model as Model<Api>,
+				typeof options?.apiKey === "string" ? options.apiKey : undefined,
+			);
+			const body = JSON.stringify({
+				modelId: `${model.provider}/${model.id}`,
+				context,
+				options: buildWireOptions(options),
+				stream: true,
+			});
+
+			response = await fetchImpl(url, { method: "POST", headers, body, signal: abortTracker.requestSignal });
+			if (!response.ok) {
+				stream.fail(await decodeGatewayError(response));
+				return;
+			}
+			// Callers can truthfully inspect the gateway HTTP response, but its
+			// request body is opaque here; callbacks themselves never cross the wire.
+			await notifyProviderResponse(
+				options,
+				response,
+				model,
+				response.headers.get("x-request-id") ?? response.headers.get("request-id"),
+			);
+			if (!response.body) {
+				stream.fail(
+					new AIError.AuthGatewayError("auth-gateway returned empty body", response.status, response.headers),
+				);
+				return;
+			}
+
+			const idleTimeoutMs = options?.streamIdleTimeoutMs ?? getStreamIdleTimeoutMs();
+			const firstEventTimeoutMs = options?.streamFirstEventTimeoutMs ?? getStreamFirstEventTimeoutMs(idleTimeoutMs);
+			const source = readSseJson<AssistantMessageEvent>(
+				response.body as ReadableStream<Uint8Array>,
+				abortTracker.requestSignal,
+			);
+			const watchedSource = iterateWithIdleTimeout(source, {
+				idleTimeoutMs,
+				firstItemTimeoutMs: firstEventTimeoutMs,
+				errorMessage: PI_NATIVE_STREAM_IDLE_TIMEOUT_ERROR,
+				firstItemErrorMessage: PI_NATIVE_STREAM_FIRST_EVENT_TIMEOUT_ERROR,
+				onIdle: () =>
+					abortTracker.abortLocally(new AIError.StreamTimeoutError(PI_NATIVE_STREAM_IDLE_TIMEOUT_ERROR)),
+				onFirstItemTimeout: () =>
+					abortTracker.abortLocally(new AIError.StreamTimeoutError(PI_NATIVE_STREAM_FIRST_EVENT_TIMEOUT_ERROR)),
+				isProgressItem: isPiNativeProgressEvent,
+			});
+			let sawTerminal = false;
+			for await (const event of watchedSource) {
+				if (event.type === "done" || event.type === "error") sawTerminal = true;
+				stream.push(event);
+				// `stream.push` resolves `.result()` on `done`/`error`; subsequent
+				// pushes are silently dropped by the base class. We still iterate
+				// to drain any trailing bytes from the wire so the underlying TCP
+				// stream closes cleanly.
+			}
+
+			if (!sawTerminal) {
+				// SSE closed before a terminal event reached us — synthesize one
+				// so awaiters of `.result()` resolve instead of hanging forever.
+				// Matches the gateway's own defensive fallback in
+				// `pi-native-server.encodeStream`.
+				const aborted = abortTracker.wasCallerAbort();
+				const partial = makeSyntheticAssistant(model as Model<Api>);
+				if (aborted) {
+					partial.stopReason = "aborted";
+					partial.errorMessage = "stream closed without terminal event";
+					stream.push({ type: "error", reason: "aborted", error: partial });
+				} else {
+					partial.stopReason = "stop";
+					stream.push({ type: "done", reason: "stop", message: partial });
+				}
+			}
+			stream.end();
+		} catch (err) {
+			stream.fail(err);
+		} finally {
+			if (callerSignal) callerSignal.removeEventListener("abort", onAbort);
+		}
+	})();
+
+	return stream;
+}
+
+function makeSyntheticAssistant(model: Model<Api>): AssistantMessage {
+	return {
+		role: "assistant",
+		content: [],
+		api: model.api,
+		provider: model.provider,
+		model: model.id,
+		usage: {
+			input: 0,
+			output: 0,
+			cacheRead: 0,
+			cacheWrite: 0,
+			totalTokens: 0,
+			cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+		},
+		stopReason: "stop",
+		timestamp: Date.now(),
+	};
+}

package/src/providers/pi-native-server.ts

@@ -0,0 +1,245 @@
+/**
+ * Pi-native wire format for the auth-gateway.
+ *
+ * Where the OpenAI / Anthropic / Responses route modules translate foreign
+ * wire shapes through pi-ai's canonical {@link Context}, this module accepts
+ * the canonical shape *directly* — for clients that already speak pi-ai
+ * (containerized omp, robomp's sidecar auth-gateway).
+ * Skipping the wire-format → Context → wire-format round-trip cuts
+ * per-request CPU but, more importantly, avoids the quantization that those
+ * translations impose on first-class pi-ai fields (service tier, cache
+ * markers, thinking budgets, tool-choice variants, …).
+ *
+ * The streaming wire is {@link AssistantMessageEvent} serialized verbatim and
+ * SSE-framed. Same type pi-ai already produces internally; the client feeds
+ * each parsed event straight into `AssistantMessageEventStream.push()` with
+ * no translation. Including `partial: AssistantMessage` on every delta is
+ * O(N²) in turn length on the wire — acceptable for the loopback / sidecar
+ * topology this transport is designed for; provider latency dominates the
+ * actual cost.
+ *
+ * Endpoint contract:
+ *   POST /v1/pi/stream
+ *   body:    { modelId, context, options?, stream? }   // `stream` defaults to true
+ *   200 SSE: stream of `AssistantMessageEvent` (terminated by `data: [DONE]`)
+ *   200 JSON (stream=false): { message: AssistantMessage }
+ *   4xx/5xx: { error: { type, message } }
+ */
+
+import type { AuthGatewayStreamControl } from "../auth-gateway/types";
+import * as AIError from "../error";
+import type { AssistantMessageEventStream, Context, SimpleStreamOptions } from "../types";
+
+export interface PiNativeParsedRequest {
+	modelId: string;
+	context: Context;
+	options: SimpleStreamOptions;
+	stream: boolean;
+}
+/**
+ * Subset of {@link SimpleStreamOptions} accepted from the wire. Function-valued
+ * fields (`fetch`, `onPayload`, `onResponse`, `onSseEvent`, exec handlers, the
+ * provider-session map) and gateway-owned controls (`apiKey`, `signal`) are
+ * intentionally absent — those are server-side concerns. Anything outside this
+ * allow-list is dropped silently rather than 400ing, so clients can forward
+ * `SimpleStreamOptions` from older / newer omp builds without per-version
+ * conditionals.
+ */
+const ALLOWED_OPTION_KEYS: ReadonlySet<keyof SimpleStreamOptions> = new Set([
+	"temperature",
+	"topP",
+	"topK",
+	"minP",
+	"presencePenalty",
+	"frequencyPenalty",
+	"repetitionPenalty",
+	"stopSequences",
+	"maxTokens",
+	"cacheRetention",
+	"cachedContent",
+	"headers",
+	"initiatorOverride",
+	"maxRetryDelayMs",
+	"metadata",
+	"sessionId",
+	"promptCacheKey",
+	"promptCache",
+	"statefulResponses",
+	"streamFirstEventTimeoutMs",
+	"streamIdleTimeoutMs",
+	"reasoning",
+	"disableReasoning",
+	"hideThinkingSummary",
+	"thinkingBudgets",
+	"toolChoice",
+	"serviceTier",
+	"kimiApiFormat",
+	"syntheticApiFormat",
+	"preferWebsockets",
+	"openrouterVariant",
+	"loopGuard",
+] as const satisfies readonly (keyof SimpleStreamOptions)[]);
+
+// ---------------------------------------------------------------------------
+// parseRequest
+// ---------------------------------------------------------------------------
+
+/**
+ * Parse a pi-native request body. Validation is intentionally minimal — only
+ * the shape the gateway itself reads is checked (`modelId`, `context.messages`
+ * array, options is an object). Everything downstream is the canonical pi-ai
+ * type surface; mis-shaped values surface as a `502 upstream_error` from
+ * `streamSimple` rather than being re-validated here.
+ *
+ * Accepts both `{ modelId: string }` and `{ model: { id: string } }` so the
+ * existing `streamProxy` client (which sends the full Model object) can target
+ * the gateway with only a URL swap.
+ */
+export function parseRequest(body: unknown, _headers?: Headers): PiNativeParsedRequest {
+	if (typeof body !== "object" || body === null || Array.isArray(body)) {
+		throw new AIError.ValidationError("Request body must be a JSON object");
+	}
+	const obj = body as Record<string, unknown>;
+
+	let modelId: string | undefined;
+	if (typeof obj.modelId === "string" && obj.modelId.length > 0) {
+		modelId = obj.modelId;
+	} else if (typeof obj.model === "string" && obj.model.length > 0) {
+		modelId = obj.model;
+	} else if (typeof obj.model === "object" && obj.model !== null) {
+		const m = obj.model as Record<string, unknown>;
+		if (typeof m.id === "string" && m.id.length > 0) modelId = m.id;
+	}
+	if (!modelId) throw new AIError.ValidationError("Missing `modelId` (or `model.id`) field");
+
+	const context = obj.context;
+	if (typeof context !== "object" || context === null || Array.isArray(context)) {
+		throw new AIError.ValidationError("Missing `context` object");
+	}
+	const ctxObj = context as Record<string, unknown>;
+	if (!Array.isArray(ctxObj.messages)) {
+		throw new AIError.ValidationError("`context.messages` must be an array");
+	}
+	if (ctxObj.systemPrompt !== undefined && !Array.isArray(ctxObj.systemPrompt)) {
+		throw new AIError.ValidationError("`context.systemPrompt` must be an array of strings when present");
+	}
+	if (ctxObj.tools !== undefined && !Array.isArray(ctxObj.tools)) {
+		throw new AIError.ValidationError("`context.tools` must be an array when present");
+	}
+
+	const options: SimpleStreamOptions = {};
+	const rawOpts = obj.options;
+	if (typeof rawOpts === "object" && rawOpts !== null && !Array.isArray(rawOpts)) {
+		const optsBag = options as Record<string, unknown>;
+		for (const [k, v] of Object.entries(rawOpts)) {
+			if (v === undefined || v === null) continue;
+			if (!ALLOWED_OPTION_KEYS.has(k as keyof SimpleStreamOptions)) continue;
+			optsBag[k] = v;
+		}
+	}
+
+	// `stream` defaults to true — pi-native clients overwhelmingly stream, and
+	// matching `streamProxy`'s implicit-stream behavior avoids a one-flag papercut.
+	const stream = typeof obj.stream === "boolean" ? obj.stream : true;
+
+	return {
+		modelId,
+		context: context as Context,
+		options,
+		stream,
+	};
+}
+// ---------------------------------------------------------------------------
+// encodeStream (SSE)
+// ---------------------------------------------------------------------------
+
+const SSE_ENCODER = new TextEncoder();
+const SSE_DONE = SSE_ENCODER.encode("data: [DONE]\n\n");
+
+/**
+ * Ship every {@link AssistantMessageEvent} verbatim, SSE-framed.
+ *
+ * No per-event re-shaping: the pi-native client is pi-ai itself, so the
+ * canonical event type IS the wire type. Including the rolling
+ * `partial: AssistantMessage` on every delta is quadratic in turn length
+ * on the wire, but for the loopback / sidecar topology this transport
+ * targets (containerized omp → host gateway, robomp slot → omp-auth-gateway
+ * sidecar) the bandwidth cost is negligible compared to provider latency —
+ * and the client gets to feed the events straight into its existing
+ * `AssistantMessageEventStream.push()` plumbing with zero translation.
+ */
+export function encodeStream(
+	events: AssistantMessageEventStream,
+	_requestedModelId?: string,
+	_options?: SimpleStreamOptions,
+	control?: AuthGatewayStreamControl,
+): ReadableStream<Uint8Array> {
+	let cancelled = control?.signal?.aborted === true;
+	const markCancelled = () => {
+		cancelled = true;
+	};
+	control?.signal?.addEventListener("abort", markCancelled, { once: true });
+	return new ReadableStream<Uint8Array>({
+		async start(controller) {
+			try {
+				if (cancelled) {
+					controller.close();
+					return;
+				}
+				for await (const event of events) {
+					if (cancelled) return;
+					controller.enqueue(SSE_ENCODER.encode(`data: ${JSON.stringify(event)}\n\n`));
+					if (event.type === "done" || event.type === "error") break;
+				}
+				if (!cancelled) {
+					controller.enqueue(SSE_DONE);
+					controller.close();
+				}
+			} catch (err) {
+				if (!cancelled) {
+					// Best-effort error envelope so the client iterator resolves
+					// instead of hanging on the dropped connection. Shape matches the
+					// canonical `error` event minus the unrecoverable `error:
+					// AssistantMessage` payload (we don't have a usable one here).
+					const message = err instanceof Error ? err.message : String(err);
+					controller.enqueue(
+						SSE_ENCODER.encode(
+							`data: ${JSON.stringify({ type: "error", reason: "error", errorMessage: message })}\n\n`,
+						),
+					);
+					controller.enqueue(SSE_DONE);
+					controller.close();
+				}
+			} finally {
+				control?.signal?.removeEventListener("abort", markCancelled);
+			}
+		},
+		cancel(reason) {
+			cancelled = true;
+			control?.signal?.removeEventListener("abort", markCancelled);
+			control?.onCancel?.(reason);
+		},
+	});
+}
+
+// ---------------------------------------------------------------------------
+// formatError
+// ---------------------------------------------------------------------------
+
+/**
+ * Pi-native error envelope:
+ *   `{ error: { type, message } }`
+ *
+ * Mirrors OpenAI's outer shape (which clients/SDKs already parse) without the
+ * provider-specific status taxonomy — pi-native callers consume `type`
+ * directly.
+ */
+export function formatError(status: number, type: string, message: string): Response {
+	return new Response(JSON.stringify({ error: { type, message } }), {
+		status,
+		headers: {
+			"Content-Type": "application/json; charset=utf-8",
+			"Cache-Control": "no-store",
+		},
+	});
+}