@gajae-code/ai 0.1.1

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

@@ -0,0 +1,228 @@
+/**
+ * Client half of the pi-native auth-gateway protocol.
+ *
+ * Dispatches a {@link streamSimple}-shaped request to an `gjc 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 gjc deployments (robogjc slots, the swarm extension) that
+ * route every LLM call through a credential-holding sidecar so the slot
+ * itself stays credential-free.
+ */
+import { readSseJson } from "@gajae-code/utils";
+import type {
+	Api,
+	AssistantMessage,
+	AssistantMessageEvent,
+	AssistantMessageEventStream as AssistantMessageEventStreamType,
+	Context,
+	Model,
+	SimpleStreamOptions,
+} from "../types";
+import { AssistantMessageEventStream } from "../utils/event-stream";
+
+/**
+ * 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",
+]);
+
+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<Error> {
+	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;
+			const out = new Error(typeof message === "string" ? message : `auth-gateway ${status}`);
+			(out as { status?: number; type?: string }).status = status;
+			if (typeof type === "string") (out as { type?: string }).type = type;
+			return out;
+		}
+	}
+	const text = typeof body === "string" ? body : JSON.stringify(body);
+	const err = new Error(`auth-gateway ${status}: ${text || response.statusText}`);
+	(err as { status?: number }).status = status;
+	return err;
+}
+
+/**
+ * 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 Error(
+			`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 `gjc 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 signal = options?.signal;
+		// 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 (signal) {
+			if (signal.aborted) {
+				stream.fail(signal.reason instanceof Error ? signal.reason : new Error(String(signal.reason ?? "aborted")));
+				return;
+			}
+			signal.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>, options?.apiKey);
+			const body = JSON.stringify({
+				modelId: model.id,
+				context,
+				options: buildWireOptions(options),
+				stream: true,
+			});
+
+			response = await fetchImpl(url, { method: "POST", headers, body, signal });
+			if (!response.ok) {
+				stream.fail(await decodeGatewayError(response));
+				return;
+			}
+			if (!response.body) {
+				stream.fail(new Error("auth-gateway returned empty body"));
+				return;
+			}
+
+			let sawTerminal = false;
+			for await (const event of readSseJson<AssistantMessageEvent>(
+				response.body as ReadableStream<Uint8Array>,
+				signal,
+			)) {
+				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 = signal?.aborted === true;
+				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 (signal) signal.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,210 @@
+/**
+ * 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 gjc, the swarm extension, robogjc'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 { 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 gjc builds without per-version
+ * conditionals.
+ */
+const ALLOWED_OPTION_KEYS: ReadonlySet<keyof SimpleStreamOptions> = new Set([
+	"temperature",
+	"topP",
+	"topK",
+	"minP",
+	"presencePenalty",
+	"frequencyPenalty",
+	"repetitionPenalty",
+	"stopSequences",
+	"maxTokens",
+	"cacheRetention",
+	"headers",
+	"initiatorOverride",
+	"maxRetryDelayMs",
+	"metadata",
+	"sessionId",
+	"streamFirstEventTimeoutMs",
+	"streamIdleTimeoutMs",
+	"reasoning",
+	"disableReasoning",
+	"hideThinkingSummary",
+	"thinkingBudgets",
+	"toolChoice",
+	"serviceTier",
+	"kimiApiFormat",
+	"syntheticApiFormat",
+	"preferWebsockets",
+] 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 Error("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 Error("Missing `modelId` (or `model.id`) field");
+
+	const context = obj.context;
+	if (typeof context !== "object" || context === null || Array.isArray(context)) {
+		throw new Error("Missing `context` object");
+	}
+	const ctxObj = context as Record<string, unknown>;
+	if (!Array.isArray(ctxObj.messages)) {
+		throw new Error("`context.messages` must be an array");
+	}
+	if (ctxObj.systemPrompt !== undefined && !Array.isArray(ctxObj.systemPrompt)) {
+		throw new Error("`context.systemPrompt` must be an array of strings when present");
+	}
+	if (ctxObj.tools !== undefined && !Array.isArray(ctxObj.tools)) {
+		throw new Error("`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 gjc → host gateway, robogjc slot → gjc-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): ReadableStream<Uint8Array> {
+	return new ReadableStream<Uint8Array>({
+		async start(controller) {
+			try {
+				for await (const event of events) {
+					controller.enqueue(SSE_ENCODER.encode(`data: ${JSON.stringify(event)}\n\n`));
+					if (event.type === "done" || event.type === "error") break;
+				}
+				controller.enqueue(SSE_DONE);
+				controller.close();
+			} catch (err) {
+				// 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();
+			}
+		},
+	});
+}
+
+// ---------------------------------------------------------------------------
+// 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",
+		},
+	});
+}