@aexol/spectral 0.0.1

package/dist/server/pi-bridge.js

@@ -0,0 +1,572 @@
+/**
+ * Per-connection pi SDK lifecycle.
+ *
+ * One `PiBridge` instance per active WebSocket connection. Wraps:
+ *  - `createAgentSession` (in-memory session manager — we own persistence in
+ *    SQLite; pi doesn't need to write its own JSONL files).
+ *  - `subscribe` listener that translates pi `AgentSessionEvent`s into our
+ *    own `ServerEvent` wire format and pushes them through a caller-supplied
+ *    sink.
+ *  - `prompt(text)` to send user input.
+ *  - `dispose()` for clean teardown on WS close.
+ *
+ * Event mapping (pi → wire):
+ *  - `message_start` (assistant)        → emit our own `message_start` with a
+ *    fresh UUID `messageId`. Pi's AssistantMessage has no stable id field, so
+ *    we mint one per turn and use it for all subsequent deltas/tool events
+ *    until `message_end`.
+ *  - `message_update` w/ inner text_delta / thinking_delta → wire `text_delta`
+ *    or `thinking_delta` carrying the `delta` and our `messageId`.
+ *  - `tool_execution_start`             → wire `tool_call`.
+ *  - `tool_execution_end`               → wire `tool_result`.
+ *  - `message_end` (assistant)          → wire `message_end` and persist the
+ *    final assembled text + JSONL of every WIRE event we emitted for this
+ *    message into SQLite. Content is taken from the final
+ *    `AssistantMessage.content` (concatenating `text` blocks); we fall back
+ *    to the `text_delta` accumulator when the provider didn't populate
+ *    final blocks (e.g. deepseek via openai-compatible).
+ *  - `agent_end`                        → wire `agent_end`.
+ *
+ * Persistence shape:
+ *  `events_jsonl` is the newline-delimited JSON of the wire-format
+ *  `ServerEvent`s we emitted for this message — NOT raw pi
+ *  `AgentSessionEvent`s. This guarantees the client's `parseWireEvents`
+ *  reducer can rehydrate the turn after a refresh using the exact same
+ *  reducer it uses for the live broadcast.
+ *
+ * Errors thrown synchronously from `session.prompt()` (e.g. no model
+ * configured) are caught by the caller in `routes.ts` and re-emitted as
+ * `{type:"error"}`.
+ *
+ * History rehydration limitation:
+ *  Pi's SDK exposes `messages` and `sendUserMessage` but reconstructing a
+ *  fresh AgentSession from a transcript of `WireMessage`s is non-trivial —
+ *  the SDK's internal state (tool registry, system prompt, model context)
+ *  expects to own message creation. For the MVP we accept that pi sees a
+ *  fresh context on reconnect; the user still sees the full transcript in
+ *  the UI because we send `session_ready.history` from SQLite. Multi-turn
+ *  conversations within a single WS connection work normally.
+ */
+import { AuthStorage, createAgentSession, DefaultResourceLoader, ModelRegistry, SessionManager, } from "@mariozechner/pi-coding-agent";
+import { randomUUID } from "node:crypto";
+import aexolMcpExtension from "../extensions/aexol-mcp.js";
+import { fetchAllowedModels as defaultFetchAllowedModels, } from "../relay/models-fetch.js";
+/**
+ * Synthetic provider names registered with pi's `ModelRegistry`. They route
+ * 100% of inference traffic through the backend (`${backendUrl}/v1`), which
+ * authenticates with the machine JWT and forwards to the upstream provider
+ * with centrally-managed API keys.
+ *
+ * Two providers (rather than one) because pi's `ProviderConfigInput.api`
+ * picks the request shape per registration. Backend supports both, but a
+ * single bag would force all models onto one shape; instead we group by
+ * upstream provider type.
+ */
+const SPECTRAL_PROXY_ANTHROPIC = "spectral-proxy-anthropic";
+const SPECTRAL_PROXY_OPENAI = "spectral-proxy-openai";
+const SPECTRAL_PROXY_USER_MODEL = "spectral-proxy-user-model";
+const SPECTRAL_PROXY_AEXOL_REFACTOR = "spectral-proxy-aexol-refactor";
+/**
+ * Concatenate text from an `AssistantMessage.content` array. Returns the
+ * empty string when no text blocks are present (tool-only turns) or when
+ * the input is missing/non-array (defensive — pi's `message_end` always
+ * carries an array, but we don't want to crash on a future SDK change).
+ */
+function extractTextFromContent(content) {
+    if (!Array.isArray(content))
+        return "";
+    let out = "";
+    for (const block of content) {
+        if (block &&
+            typeof block === "object" &&
+            block.type === "text" &&
+            typeof block.text === "string") {
+            out += block.text;
+        }
+    }
+    return out;
+}
+export class PiBridge {
+    session;
+    unsubscribe;
+    pending;
+    disposed = false;
+    opts;
+    /**
+     * Pi's model registry. Built lazily in `start()` so we can resolve a
+     * `modelId` (envelope-supplied or SQLite-persisted) to a concrete `Model`
+     * via `registry.getAll().find(m => m.id === modelId)` before invoking
+     * `session.setModel()`. Phase 3 (Available Models whitelist).
+     */
+    modelRegistry;
+    /**
+     * Last `modelId` we successfully applied via `session.setModel()`, or
+     * `undefined` if we never applied one (pi falls back to its own settings
+     * file in that case, matching pre-Phase-3 behaviour). Tracked so repeated
+     * envelopes carrying the same modelId don't churn pi's internal state.
+     */
+    lastAppliedModelId;
+    constructor(opts) {
+        this.opts = opts;
+    }
+    /**
+     * Create the pi session, wire up subscription, and return.
+     * Throws on creation failure (caller should surface to client).
+     */
+    async start() {
+        if (this.disposed)
+            throw new Error("PiBridge already disposed");
+        // ResourceLoader with the Aexol MCP extension wired in via factory.
+        // The extension's signature `(pi: ExtensionAPI) => Promise<void>` matches
+        // the ExtensionFactory type exactly, so we can pass it directly.
+        const resourceLoader = new DefaultResourceLoader({
+            cwd: this.opts.cwd,
+            agentDir: this.opts.agentDir ?? `${process.env.HOME ?? ""}/.pi/agent`,
+            extensionFactories: [aexolMcpExtension],
+            // Skip on-disk extension/skill discovery so the server is self-contained.
+            noExtensions: false,
+            noSkills: true,
+            noPromptTemplates: true,
+            noThemes: true,
+        });
+        await resourceLoader.reload();
+        // In-memory session: SQLite is our source of truth.
+        const sessionManager = SessionManager.inMemory(this.opts.cwd);
+        // Build a model registry that does NOT touch ~/.pi/agent/auth.json or
+        // ~/.pi/agent/models.json — the backend is now the only source of
+        // provider credentials and the only allowed inference target. We then
+        // register synthetic providers (`spectral-proxy-anthropic` /
+        // `spectral-proxy-openai`) whose `baseUrl` points at the backend's
+        // `/v1` proxy and whose `apiKey` is the machine JWT. Pi will then
+        // POST `${baseUrl}/messages` (or `/chat/completions`) with
+        // `Authorization: Bearer <machineJwt>` for every turn — the backend
+        // verifies the JWT, looks up the requested `model` in the BaseModel
+        // whitelist, and forwards to the upstream provider with its own
+        // centrally-managed API keys. There is intentionally NO fallback to
+        // disk-based auth: this is the single path for `spectral serve`.
+        const authStorage = AuthStorage.inMemory();
+        this.modelRegistry = ModelRegistry.inMemory(authStorage);
+        const fetchModels = this.opts.fetchAllowedModels ?? defaultFetchAllowedModels;
+        let allowedModels;
+        try {
+            allowedModels = await fetchModels({
+                backendUrl: this.opts.backendUrl,
+                machineJwt: this.opts.machineJwt,
+            });
+        }
+        catch (err) {
+            const e = err instanceof Error ? err : new Error(String(err));
+            throw new Error(`Failed to fetch allowed models from backend; check SPECTRAL_BACKEND_URL ` +
+                `and machine JWT. Underlying error: ${e.message}`);
+        }
+        this.registerSyntheticProviders(allowedModels);
+        console.info(`✓ Inference routed via backend proxy (${allowedModels.length} model(s) available)`);
+        const result = await createAgentSession({
+            cwd: this.opts.cwd,
+            resourceLoader,
+            sessionManager,
+            authStorage,
+            modelRegistry: this.modelRegistry,
+        });
+        this.session = result.session;
+        // Subscribe BEFORE any prompt fires.
+        this.unsubscribe = this.session.subscribe((ev) => this.handleEvent(ev));
+    }
+    /**
+     * Register one synthetic provider per upstream API shape. Anthropic models
+     * go to `${backendUrl}/v1/messages` (Messages API); everything else (OpenAI,
+     * Google, Cerebras, etc.) goes to `${backendUrl}/v1/chat/completions`
+     * (OpenAI-compatible API). The backend supports both endpoints natively
+     * (verified in F1).
+     *
+     * Pi will send `Authorization: Bearer ${apiKey}` (because `authHeader: true`)
+     * which carries the machine JWT — the only credential the backend trusts.
+     *
+     * The `id` we register is the raw `modelId` (e.g. `claude-3-5-haiku-latest`),
+     * which is exactly what the backend expects in `body.model`.
+     */
+    registerSyntheticProviders(allowedModels) {
+        if (!this.modelRegistry)
+            return;
+        const baseUrl = `${this.opts.backendUrl.replace(/\/$/, "")}/v1`;
+        const anthropicModels = allowedModels.filter((m) => m.provider === "anthropic");
+        const openaiCompatModels = allowedModels.filter((m) => m.provider !== "anthropic" && m.provider !== "built-in");
+        if (anthropicModels.length > 0) {
+            this.modelRegistry.registerProvider(SPECTRAL_PROXY_ANTHROPIC, {
+                baseUrl,
+                apiKey: this.opts.machineJwt,
+                authHeader: true,
+                api: "anthropic-messages",
+                models: anthropicModels.map((m) => ({
+                    id: m.modelId,
+                    name: m.displayName,
+                    api: "anthropic-messages",
+                    // Pin provider/baseUrl explicitly so pi's ModelRegistry doesn't
+                    // auto-derive `provider` from a slash-prefixed id (e.g. treating
+                    // `deepseek/deepseek-v4-pro` as provider `"deepseek"`), which would
+                    // make `hasConfiguredAuth(model)` look up the wrong provider key
+                    // and surface "No API key for deepseek/...". Both must point back
+                    // at our synthetic proxy provider so auth resolves to the machine JWT.
+                    provider: SPECTRAL_PROXY_ANTHROPIC,
+                    baseUrl,
+                    reasoning: false,
+                    input: ["text", "image"],
+                    // The cost block is required by pi's typing but unused for routing;
+                    // the backend enforces real billing/limits server-side, not pi.
+                    cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+                    contextWindow: 0,
+                    maxTokens: 0,
+                })),
+            });
+        }
+        if (openaiCompatModels.length > 0) {
+            this.modelRegistry.registerProvider(SPECTRAL_PROXY_OPENAI, {
+                baseUrl,
+                apiKey: this.opts.machineJwt,
+                authHeader: true,
+                api: "openai-completions",
+                models: openaiCompatModels.map((m) => ({
+                    id: m.modelId,
+                    name: m.displayName,
+                    api: "openai-completions",
+                    // See anthropic batch above for rationale — without these, pi
+                    // auto-derives `provider` from slash-prefixed ids like
+                    // `deepseek/deepseek-v4-pro` or `meta-llama/llama-3.3-70b-instruct`,
+                    // breaking auth lookup against our synthetic proxy provider.
+                    provider: SPECTRAL_PROXY_OPENAI,
+                    baseUrl,
+                    reasoning: false,
+                    input: ["text", "image"],
+                    cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+                    contextWindow: 0,
+                    maxTokens: 0,
+                })),
+            });
+        }
+        // Built-in UserModel entries — custom models registered by the team.
+        // These route through `${backendUrl}/v1` (OpenAI-compatible) and carry
+        // the `name` as the model id, which the backend resolves to the actual
+        // provider/configuration via the UserModel record.
+        const userModelEntries = allowedModels.filter((m) => m.provider === "built-in");
+        if (userModelEntries.length > 0) {
+            this.modelRegistry.registerProvider(SPECTRAL_PROXY_USER_MODEL, {
+                baseUrl,
+                apiKey: this.opts.machineJwt,
+                authHeader: true,
+                api: "openai-completions",
+                models: userModelEntries.map((m) => ({
+                    id: m.modelId,
+                    name: m.displayName,
+                    api: "openai-completions",
+                    provider: SPECTRAL_PROXY_USER_MODEL,
+                    baseUrl,
+                    reasoning: false,
+                    input: ["text", "image"],
+                    cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+                    contextWindow: 0,
+                    maxTokens: 0,
+                })),
+            });
+        }
+        // Refactor-loop model — dedicated endpoint for Aexol agent chat toggle.
+        // Routes to the team user model at /models/team/aexol/refactor-loop/v1
+        // using machine JWT auth, same as other synthetic providers.
+        {
+            const refactorBaseUrl = `${this.opts.backendUrl.replace(/\/$/, "")}/models/built-in/refactor-loop/v1`;
+            this.modelRegistry.registerProvider(SPECTRAL_PROXY_AEXOL_REFACTOR, {
+                baseUrl: refactorBaseUrl,
+                apiKey: this.opts.machineJwt,
+                authHeader: true,
+                api: "openai-completions",
+                models: [
+                    {
+                        id: "__aexol_refactor_loop__",
+                        name: "Aexol Refactor Loop",
+                        api: "openai-completions",
+                        baseUrl: refactorBaseUrl,
+                        reasoning: false,
+                        input: ["text", "image"],
+                        cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+                        contextWindow: 0,
+                        maxTokens: 0,
+                    },
+                ],
+            });
+        }
+    }
+    /**
+     * Apply a sticky model selection to the underlying pi session, if it
+     * differs from what was last applied. No-ops when:
+     *  - `modelId` is null/undefined (caller passed nothing to apply)
+     *  - the same modelId was already applied to this session
+     *  - the registry can't resolve the modelId (emits an `error` wire event
+     *    so the client surfaces the failure, but does NOT throw — the caller
+     *    is expected to drop the prompt)
+     *
+     * Returns true when the requested model is now in effect (either because
+     * we just applied it or because it was already applied). Returns false
+     * on resolution failure so the caller can skip `prompt()` and avoid
+     * driving pi against the wrong model.
+     *
+     * Phase 3 (Available Models whitelist).
+     */
+    async setModel(modelId) {
+        if (!modelId)
+            return true; // nothing to apply — pi keeps its current model
+        if (!this.session)
+            throw new Error("PiBridge.start() not called");
+        if (this.lastAppliedModelId === modelId)
+            return true; // idempotent: same model already in effect
+        if (!this.modelRegistry) {
+            // Defensive: start() always populates this; if it didn't we can't
+            // resolve and must surface to the client rather than silently using
+            // pi's default.
+            this.opts.emit({
+                type: "error",
+                message: `Cannot apply modelId "${modelId}": model registry unavailable`,
+            });
+            return false;
+        }
+        const model = this.modelRegistry
+            .getAvailable()
+            .find((m) => m.id === modelId);
+        if (!model) {
+            this.opts.emit({
+                type: "error",
+                message: `Unknown modelId "${modelId}" — not found in pi model registry`,
+            });
+            return false;
+        }
+        try {
+            await this.session.setModel(model);
+            this.lastAppliedModelId = modelId;
+            return true;
+        }
+        catch (err) {
+            const e = err instanceof Error ? err : new Error(String(err));
+            this.opts.onError?.(e);
+            this.opts.emit({
+                type: "error",
+                message: `Failed to switch to modelId "${modelId}": ${e.message}`,
+            });
+            return false;
+        }
+    }
+    /**
+     * Forward a user message to pi. Resolves when the full turn ends.
+     * The caller is responsible for persisting the user message to SQLite
+     * BEFORE invoking this — we don't do it here because pi's `prompt` may
+     * fail and we still want the user message recorded.
+     */
+    async prompt(text) {
+        if (!this.session)
+            throw new Error("PiBridge.start() not called");
+        try {
+            await this.session.prompt(text);
+        }
+        catch (err) {
+            const e = err instanceof Error ? err : new Error(String(err));
+            this.opts.onError?.(e);
+            this.opts.emit({ type: "error", message: e.message });
+        }
+    }
+    dispose() {
+        if (this.disposed)
+            return;
+        this.disposed = true;
+        try {
+            this.unsubscribe?.();
+        }
+        catch {
+            // ignore
+        }
+        try {
+            this.session?.dispose();
+        }
+        catch {
+            // ignore
+        }
+        this.session = undefined;
+        this.unsubscribe = undefined;
+        this.pending = undefined;
+    }
+    // --- internals -----------------------------------------------------------
+    /**
+     * Emit a wire event AND record it in the pending message's audit log so
+     * the persisted JSONL exactly matches the live broadcast. Use this in
+     * place of `this.opts.emit` for any event that should appear in
+     * `events_jsonl` for the current assistant message.
+     */
+    emitAndBuffer(event) {
+        if (this.pending)
+            this.pending.wireEvents.push(event);
+        this.opts.emit(event);
+    }
+    /**
+     * Finalize the current `pending` assistant message: persist its assembled
+     * content + wire events (including any tool_call / tool_result events that
+     * arrived after `message_end`) to SQLite via `onAssistantMessageComplete`.
+     *
+     * Idempotent: marks the pending as `finalized` on first call; subsequent
+     * calls are no-ops (guards against double-persisting when both
+     * `message_start` and `agent_end` would otherwise finalize the same pending).
+     *
+     * Skips empty framing-only messages (messages with no text, thinking, or
+     * tool events that contribute nothing the client can render).
+     */
+    finalizePendingMessage() {
+        if (!this.pending || this.pending.finalized)
+            return;
+        this.pending.finalized = true;
+        const { messageId, text, wireEvents, content } = this.pending;
+        // Skip persistence for pure-framing intermediate "messages". Pi
+        // emits a `message_end` for each internal step (e.g. between two
+        // tool-calling rounds) and many of those carry no visible content
+        // at all — only `message_start` + `message_end` framing. Persisting
+        // them creates orphan rows that are dropped on hydration anyway,
+        // and they pollute the sidebar message counts. The decision is:
+        //   - persist if `content` is non-empty (text the user saw), OR
+        //   - persist if any "meaningful" wire event was buffered
+        //     (text_delta / thinking_delta / tool_call / tool_result).
+        // Otherwise the message contributes nothing the client can render
+        // — drop it. Live broadcast is unchanged; subscribers already saw
+        // the framing on the wire.
+        const finalContent = content ?? text;
+        const hasMeaningfulEvent = wireEvents.some((e) => e.type === "text_delta" ||
+            e.type === "thinking_delta" ||
+            e.type === "tool_call" ||
+            e.type === "tool_result");
+        if (!finalContent && !hasMeaningfulEvent) {
+            console.debug("[pi-bridge] skipping empty intermediate message");
+            return;
+        }
+        const eventsJsonl = wireEvents.map((e) => JSON.stringify(e)).join("\n");
+        try {
+            this.opts.onAssistantMessageComplete({
+                messageId,
+                content: finalContent,
+                eventsJsonl,
+            });
+        }
+        catch (err) {
+            const e = err instanceof Error ? err : new Error(String(err));
+            this.opts.onError?.(e);
+        }
+    }
+    /**
+     * Subscriber callback. Public so tests can drive event flow without
+     * spinning up a real pi session — production code never calls this
+     * directly; pi's `subscribe()` does, via the closure registered in
+     * `start()`.
+     */
+    handleEvent(ev) {
+        if (this.disposed)
+            return;
+        switch (ev.type) {
+            case "message_start": {
+                // Only assistant messages get our message_start frame on the wire.
+                // User messages are persisted separately by the routes layer.
+                if (ev.message.role !== "assistant")
+                    return;
+                // Finalize the previous pending message (if any) before starting a
+                // new one. This captures tool events that arrived after the previous
+                // `message_end` but before this `message_start`. Deferred persistence
+                // is necessary because pi fires tool_execution_* events BETWEEN
+                // messages — after the previous `message_end` nulled the pending in
+                // the old code, those tool events were lost from history.
+                this.finalizePendingMessage();
+                const messageId = randomUUID();
+                const wireEvent = {
+                    type: "message_start",
+                    messageId,
+                    role: "assistant",
+                };
+                this.pending = { messageId, text: "", wireEvents: [wireEvent] };
+                this.opts.emit(wireEvent);
+                return;
+            }
+            case "message_update": {
+                if (!this.pending)
+                    return;
+                const inner = ev.assistantMessageEvent;
+                if (inner.type === "text_delta") {
+                    this.pending.text += inner.delta;
+                    this.emitAndBuffer({
+                        type: "text_delta",
+                        messageId: this.pending.messageId,
+                        delta: inner.delta,
+                    });
+                }
+                else if (inner.type === "thinking_delta") {
+                    this.emitAndBuffer({
+                        type: "thinking_delta",
+                        messageId: this.pending.messageId,
+                        delta: inner.delta,
+                    });
+                }
+                // Other inner types (text_start/end, toolcall_*, done, error) we
+                // don't surface as their own wire frames — tool calls come through
+                // the top-level tool_execution_* events instead.
+                return;
+            }
+            case "tool_execution_start": {
+                this.emitAndBuffer({
+                    type: "tool_call",
+                    messageId: this.pending?.messageId ?? "",
+                    id: ev.toolCallId,
+                    name: ev.toolName,
+                    args: ev.args,
+                });
+                return;
+            }
+            case "tool_execution_end": {
+                this.emitAndBuffer({
+                    type: "tool_result",
+                    messageId: this.pending?.messageId ?? "",
+                    id: ev.toolCallId,
+                    result: ev.result,
+                    isError: ev.isError,
+                });
+                return;
+            }
+            case "message_end": {
+                if (ev.message.role !== "assistant" || !this.pending)
+                    return;
+                const { messageId, text } = this.pending;
+                // Prefer the assembled text from pi's final AssistantMessage
+                // (authoritative — providers like deepseek only populate this and
+                // skip per-token `text_delta`s entirely). Fall back to the
+                // accumulator for providers that stream deltas without re-asserting
+                // the final content, and finally to the empty string for tool-only
+                // turns that legitimately have no text.
+                const assembled = extractTextFromContent(ev.message.content);
+                const content = assembled || text;
+                const endEvent = { type: "message_end", messageId };
+                this.pending.wireEvents.push(endEvent);
+                this.opts.emit(endEvent);
+                // Defer persistence: keep `this.pending` alive so tool events that
+                // arrive after `message_end` (pi fires tool_execution_* events
+                // BETWEEN messages) are buffered into `pending.wireEvents`. We store
+                // the final content and persist later — when the next `message_start`
+                // signals a new step, or when `agent_end` closes the turn.
+                this.pending.content = content;
+                return;
+            }
+            case "agent_end": {
+                // Finalize the last pending message before closing the turn. After
+                // `agent_end`, no more tool events will arrive, so this is the final
+                // persistence opportunity for the current message.
+                this.finalizePendingMessage();
+                this.pending = undefined;
+                this.opts.emit({ type: "agent_end" });
+                return;
+            }
+            default:
+                // Other pi-internal events (turn_start, queue_update, compaction_*,
+                // auto_retry_*, tool_execution_update) are intentionally not on the
+                // wire surface for MVP and are NOT persisted — the wire format is
+                // the source of truth for replay.
+                return;
+        }
+    }
+}