@aexol/spectral 0.0.1

package/dist/server/session-stream.js

@@ -0,0 +1,579 @@
+/**
+ * Session-scoped streaming layer for `spectral serve`.
+ *
+ * Background: prior to this module each WebSocket owned its own `PiBridge`
+ * instance and the routes layer enforced single-writer-wins (4001 eviction)
+ * to keep that bridge unique per session. That model lost data on browser
+ * refresh — the WS close torn down the pi process mid-stream, and a re-open
+ * couldn't recover what hadn't yet hit `agent_end` (and thus SQLite).
+ *
+ * New model:
+ *   - Pi lifecycle is per **Spectral session**, not per WS.
+ *   - 0..N WebSockets may attach to the same session simultaneously. Each
+ *     gets the same broadcast stream of events.
+ *   - When a WS detaches (close, error, refresh), the pi process keeps
+ *     running. Closing every tab does NOT cancel the in-flight turn —
+ *     it runs to completion and persists on `agent_end` as before.
+ *   - On `attach`, the manager hands back a replay payload: full DB history
+ *     plus a snapshot of the currently in-flight turn (if any). The client
+ *     replays the snapshot through the same reducer it uses for live events
+ *     and continues streaming naturally.
+ *   - Persistence shape is unchanged: only the final assistant message is
+ *     written to SQLite on `agent_end`. In-flight events live in memory only.
+ *     A server crash mid-turn discards the in-flight state — acceptable
+ *     for MVP.
+ *
+ * Failure modes:
+ *   - Pi throws synchronously in `prompt()` → bridge surfaces as `error`
+ *     event; manager broadcasts and clears `currentTurn`.
+ *   - One subscriber's `ws.send` throws → caught, logged, removed from the
+ *     subscriber set; broadcast continues to the rest.
+ *   - `agent_end` arrives without a current turn (defensive) → broadcast
+ *     anyway so any late attachers don't get stuck.
+ *
+ * TODO (future): idle GC. A `SessionStream` with `subscribers.size === 0`
+ *   and no current turn could be disposed after some grace window (e.g. 5
+ *   minutes) to release pi resources for chronically-idle sessions. Skipped
+ *   for now — streams accumulate for the lifetime of the server process.
+ */
+import { randomUUID } from "node:crypto";
+import { PiBridge } from "./pi-bridge.js";
+import { generateSessionTitle, isDefaultTitle, } from "./title-generator.js";
+const DEFAULT_BRIDGE_FACTORY = (args) => new PiBridge(args);
+/** Safety limit for autonomous refactor loop iterations per session. */
+const MAX_AEXOL_ITERATIONS = 100;
+export class SessionStreamManager {
+    store;
+    cwd;
+    backendUrl;
+    machineJwt;
+    bridgeFactory;
+    agentDir;
+    titleLlmCall;
+    disableAutoTitle;
+    publishMetaEvent;
+    streams = new Map();
+    /**
+     * Sessions for which we've already attempted (or queued) auto-title
+     * generation in this server process. Per-process is intentional: a server
+     * restart resets the set, but `isDefaultTitle()` still gates the work so a
+     * since-renamed session is never overwritten.
+     */
+    titleGenerationAttempted = new Set();
+    disposed = false;
+    constructor(opts) {
+        this.store = opts.store;
+        this.cwd = opts.cwd;
+        this.backendUrl = opts.backendUrl;
+        this.machineJwt = opts.machineJwt;
+        this.bridgeFactory = opts.bridgeFactory ?? DEFAULT_BRIDGE_FACTORY;
+        this.agentDir = opts.agentDir;
+        this.titleLlmCall = opts.titleLlmCall;
+        this.disableAutoTitle = opts.disableAutoTitle === true;
+        this.publishMetaEvent = opts.publishMetaEvent;
+    }
+    /**
+     * Attach a subscriber to a session. Lazily creates the underlying pi
+     * session on first attach. The caller is responsible for sending the
+     * initial `session_ready` frame using the returned replay payload (this
+     * keeps wire-protocol concerns in the routes layer).
+     *
+     * Throws if the session id is unknown in SQLite (caller should turn this
+     * into a wire-level error frame + close).
+     */
+    attach(sessionId, subscriber) {
+        if (this.disposed)
+            throw new Error("SessionStreamManager disposed");
+        const detail = this.store.getSession(sessionId);
+        if (!detail)
+            throw new Error(`Unknown sessionId: ${sessionId}`);
+        let stream = this.streams.get(sessionId);
+        if (!stream) {
+            stream = this.createStream(sessionId);
+            this.streams.set(sessionId, stream);
+        }
+        stream.subscribers.add(subscriber);
+        return {
+            history: detail.messages,
+            currentTurn: stream.currentTurn ? snapshotTurn(stream.currentTurn) : null,
+            ready: stream.ready,
+        };
+    }
+    /**
+     * Detach a subscriber. Idempotent. Does NOT dispose the underlying pi
+     * session — even when subscribers reach zero, the in-flight turn must
+     * complete and persist.
+     */
+    detach(sessionId, subscriber) {
+        const stream = this.streams.get(sessionId);
+        if (!stream)
+            return;
+        stream.subscribers.delete(subscriber);
+        // Intentional: do NOT dispose the bridge here. See file-level docs.
+    }
+    /** True if the session has an in-flight turn (manager-side; not WS-side). */
+    hasActiveTurn(sessionId) {
+        return this.streams.get(sessionId)?.currentTurn != null;
+    }
+    /**
+     * Persist a user message and forward it to pi. Resolves after the user
+     * message is persisted + pi is invoked (NOT after the turn completes —
+     * the turn lifetime is observed via the broadcast stream).
+     *
+     * Broadcast ordering:
+     *   1. user message persisted to SQLite
+     *   2. `user_message_appended` broadcast to all subscribers (including
+     *      the originating tab)
+     *   3. new `currentTurn` opened
+     *   4. `bridge.prompt()` invoked (events arrive asynchronously and are
+     *      buffered + broadcast as they come)
+     *
+     * Sticky model selection (Phase 3 — Available Models whitelist):
+     *   - When `modelId` is provided, we apply it via `bridge.setModel()` and
+     *     persist to SQLite for cross-restart recovery, BEFORE invoking
+     *     `bridge.prompt()`. If `setModel` fails (unknown model, registry
+     *     unavailable, pi-side error) the bridge has already emitted an
+     *     `error` wire event and we drop the prompt to avoid running it
+     *     against the wrong model.
+     *   - When `modelId` is omitted, we look up SQLite. If a previous turn
+     *     persisted a value, we reapply it on this turn (this is the
+     *     cross-restart recovery path: a fresh server process has lost pi's
+     *     in-memory model state, so we re-pin from durable storage).
+     *   - When neither envelope nor SQLite have a value, we leave model
+     *     selection to pi's own settings file (pre-Phase-3 behaviour).
+     */
+    async prompt(sessionId, content, modelId) {
+        if (this.disposed)
+            throw new Error("SessionStreamManager disposed");
+        const stream = this.streams.get(sessionId);
+        if (!stream)
+            throw new Error(`No active stream for session: ${sessionId}`);
+        // Wait for pi to be ready before we persist + invoke. If start failed,
+        // surface to all subscribers instead of throwing into the route handler.
+        try {
+            await stream.ready;
+        }
+        catch (err) {
+            const e = err instanceof Error ? err : new Error(String(err));
+            this.broadcast(stream, { type: "error", message: `Agent not ready: ${e.message}` });
+            return;
+        }
+        // Sticky-model resolution & application. Phase 3 (Available Models
+        // whitelist). Order:
+        //   a) If envelope carried a `modelId`, use it.
+        //   b) Else, look up the per-session persisted modelId in SQLite
+        //      (cross-restart recovery — server restart wipes pi's in-memory
+        //      session model state, but our durable store has the last value).
+        //   c) Else, leave model selection to pi (pre-Phase-3 behaviour).
+        //
+        // We apply BEFORE persisting the user message: if the bridge can't
+        // resolve the model (unknown id, registry unavailable), it has already
+        // emitted an `error` wire event and we drop the prompt rather than
+        // recording a user message we know we won't be able to respond to.
+        //
+        // Persistence: only the envelope-supplied value is written back. A
+        // recovery-only application (case b) doesn't update the row — the
+        // value is already there.
+        const effectiveModelId = modelId ?? this.store.getSessionModel(sessionId) ?? undefined;
+        if (effectiveModelId && stream.bridge.setModel) {
+            const ok = await stream.bridge.setModel(effectiveModelId);
+            if (!ok) {
+                // Bridge already emitted an error; nothing else to do. We
+                // intentionally do not persist the user message — the turn never
+                // ran.
+                return;
+            }
+        }
+        if (modelId) {
+            try {
+                this.store.setSessionModel(sessionId, modelId);
+            }
+            catch (err) {
+                // Persisting the sticky model is best-effort: the live turn will
+                // still run with the model already applied above. A failure here
+                // only affects cross-restart recovery, which is non-critical.
+                console.warn(`[spectral] warn: failed to persist sticky model for ${sessionId}: ${err instanceof Error ? err.message : String(err)}`);
+            }
+        }
+        // 1. Persist user message first (survives mid-prompt failures).
+        let stored;
+        try {
+            stored = this.store.appendMessage(sessionId, { role: "user", content });
+        }
+        catch (err) {
+            const e = err instanceof Error ? err : new Error(String(err));
+            this.broadcast(stream, {
+                type: "error",
+                message: `Failed to persist user message: ${e.message}`,
+            });
+            return;
+        }
+        // 2. Broadcast the persisted message so every tab — including the one
+        //    that sent the prompt — appends an authoritative user turn.
+        this.broadcast(stream, { type: "user_message_appended", message: stored });
+        // 3. Open a new in-flight turn. Allocated even before pi emits anything
+        //    so a re-attach immediately after `prompt` sees the turn.
+        stream.currentTurn = {
+            turnId: randomUUID(),
+            startedAt: Date.now(),
+            events: [],
+            assistantText: "",
+        };
+        // 4. Fire pi. `prompt` resolves on agent_end; errors are handled inside
+        //    PiBridge (it emits `error` for us). We don't await — broadcast is
+        //    driven by the bridge's emit callback.
+        void stream.bridge.prompt(content);
+    }
+    /**
+     * Tear down everything. Best-effort: disposes every bridge, drops all
+     * subscribers. After this the manager is unusable.
+     */
+    dispose() {
+        if (this.disposed)
+            return;
+        this.disposed = true;
+        for (const stream of this.streams.values()) {
+            stream.aexolActive = false;
+            try {
+                stream.bridge.dispose();
+            }
+            catch {
+                // ignore
+            }
+            stream.subscribers.clear();
+        }
+        this.streams.clear();
+    }
+    /** Test/inspection helper: how many streams are currently tracked. */
+    streamCount() {
+        return this.streams.size;
+    }
+    /**
+     * Count of sessions with an in-flight turn (i.e. a `currentTurn` set).
+     * Used by `gracefulShutdown` to decide whether to keep waiting before
+     * tearing down — a non-zero count means at least one assistant response
+     * is mid-stream and we'd rather let it finish (within the grace window)
+     * than orphan a half-streamed message in the UI.
+     *
+     * Cheap O(streams) scan; we only call it ~50× during a 5 s graceful
+     * shutdown so the linear walk is fine.
+     */
+    activeTurnCount() {
+        let n = 0;
+        for (const s of this.streams.values()) {
+            if (s.currentTurn != null)
+                n++;
+        }
+        return n;
+    }
+    /**
+     * Tear down a single session's stream — disposes the pi bridge and clears
+     * subscribers. Idempotent. Called by the routes layer right before
+     * `DELETE /api/sessions/:id` so the SQL cascade doesn't leave a zombie
+     * pi process driving events at a session that no longer exists.
+     *
+     * Does NOT remove the session from the store — that's the caller's job.
+     */
+    disposeSessionStream(sessionId) {
+        const stream = this.streams.get(sessionId);
+        if (!stream)
+            return;
+        stream.aexolActive = false;
+        try {
+            stream.bridge.dispose();
+        }
+        catch {
+            // ignore
+        }
+        // Best-effort: notify any still-open subscribers so they close cleanly
+        // rather than hanging on a dead pi process. We don't broadcast through
+        // `broadcast()` because that would re-enter the dead-subscriber pruning
+        // loop on a stream we're about to drop anyway.
+        for (const sub of stream.subscribers) {
+            if (!sub.isOpen())
+                continue;
+            try {
+                sub.send({ type: "error", message: "Session deleted" });
+            }
+            catch {
+                // ignore
+            }
+        }
+        stream.subscribers.clear();
+        this.streams.delete(sessionId);
+    }
+    /**
+     * Tear down every stream whose session belongs to the given list of ids.
+     * Used by the project-delete path: the route layer reads the project's
+     * session ids from `deleteProject()` and passes them here BEFORE the SQL
+     * cascade fires, so no pi process ever observes the FK cascade.
+     */
+    disposeProjectStreams(sessionIds) {
+        for (const sid of sessionIds) {
+            this.disposeSessionStream(sid);
+        }
+    }
+    /**
+     * Set the autonomous refactor loop state for a session. When `active` is
+     * true, the manager will auto-send "continue" after each `agent_end` event
+     * until deactivated or the safety limit is reached.
+     */
+    setAexolActive(sessionId, active) {
+        const stream = this.streams.get(sessionId);
+        if (stream) {
+            stream.aexolActive = active;
+            if (!active)
+                stream.aexolIterationCount = 0;
+        }
+    }
+    // --- internals ----------------------------------------------------------
+    createStream(sessionId) {
+        // Resolve cwd from the owning project. Sessions without a project
+        // shouldn't exist (FK enforces it), but we fall back to the manager's
+        // default cwd if the lookup somehow fails — better than crashing the
+        // attach.
+        const projectId = this.store.getSessionProjectId(sessionId);
+        let cwd = this.cwd;
+        if (projectId) {
+            const project = this.store.getProject(projectId);
+            if (project)
+                cwd = project.path;
+        }
+        // Forward declaration so the bridge factory's emit callback can refer to
+        // the stream object that's still being assembled.
+        const stream = {
+            sessionId,
+            cwd,
+            bridge: undefined,
+            ready: Promise.resolve(),
+            startError: null,
+            subscribers: new Set(),
+            currentTurn: null,
+            aexolActive: false,
+            aexolIterationCount: 0,
+        };
+        const bridgeOpts = {
+            cwd,
+            agentDir: this.agentDir,
+            backendUrl: this.backendUrl,
+            machineJwt: this.machineJwt,
+            emit: (event) => this.handleBridgeEvent(stream, event),
+            onAssistantMessageComplete: ({ messageId, content, eventsJsonl }) => {
+                try {
+                    this.store.appendMessage(sessionId, {
+                        id: messageId,
+                        role: "assistant",
+                        content,
+                        eventsJsonl,
+                    });
+                }
+                catch (err) {
+                    console.error(`[spectral] error: failed to persist assistant message: ${err instanceof Error ? err.message : String(err)}`);
+                }
+            },
+            onError: (err) => {
+                console.error(`[spectral] error: pi bridge error: ${err.message}`);
+            },
+        };
+        stream.bridge = this.bridgeFactory(bridgeOpts);
+        stream.ready = stream.bridge
+            .start()
+            .catch((err) => {
+            const e = err instanceof Error ? err : new Error(String(err));
+            stream.startError = e;
+            // Notify any already-attached subscribers so they don't sit on a
+            // dead connection. Late attachers see startError via attach->ready.
+            this.broadcast(stream, {
+                type: "error",
+                message: `Failed to start agent: ${e.message}`,
+            });
+            throw e;
+        });
+        return stream;
+    }
+    handleBridgeEvent(stream, event) {
+        // Buffer replayable events into the in-flight turn. We intentionally
+        // accept events even if currentTurn is null (rare race: pi emits before
+        // prompt() opened the turn), in which case we open one defensively so
+        // late attachers see the events. The first event in such a case is
+        // typically `message_start`.
+        if (isReplayable(event)) {
+            if (!stream.currentTurn) {
+                stream.currentTurn = {
+                    turnId: randomUUID(),
+                    startedAt: Date.now(),
+                    events: [],
+                    assistantText: "",
+                };
+            }
+            stream.currentTurn.events.push(event);
+            if (event.type === "text_delta") {
+                stream.currentTurn.assistantText += event.delta;
+            }
+        }
+        // Broadcast first, then maybe close out the turn. agent_end clears the
+        // buffer because by that point the assistant message is already in
+        // SQLite (PiBridge calls onAssistantMessageComplete on message_end,
+        // which fires before agent_end).
+        this.broadcast(stream, event);
+        if (event.type === "agent_end") {
+            const finishedTurn = stream.currentTurn;
+            stream.currentTurn = null;
+            // Fire-and-forget auto-title generation. Runs only once per session
+            // per server lifetime, only when the session is still wearing its
+            // default title, and never blocks the user's stream (the user's
+            // turn is already complete by the time this runs).
+            this.maybeGenerateTitle(stream, finishedTurn);
+            // Autonomous refactor loop: when aexolActive is set, auto-send
+            // "continue" after each agent_end to keep the loop going.
+            if (stream.aexolActive && stream.aexolIterationCount < MAX_AEXOL_ITERATIONS) {
+                stream.aexolIterationCount++;
+                console.log(`[aexol-loop] iteration ${stream.aexolIterationCount}/${MAX_AEXOL_ITERATIONS}`);
+                void this.prompt(stream.sessionId, "continue", undefined).catch((err) => {
+                    console.error(`[aexol-loop] iteration failed: ${err instanceof Error ? err.message : String(err)}`);
+                    stream.aexolActive = false;
+                });
+            }
+            else if (stream.aexolActive) {
+                console.log("[aexol-loop] max iterations reached, stopping");
+                stream.aexolActive = false;
+            }
+        }
+        else if (event.type === "error") {
+            // An error event arriving outside a turn (or bubbling out of one) —
+            // discard partial buffer to avoid replaying half a turn that the
+            // client has already shown an error for. The error event itself is
+            // still broadcast above.
+            stream.currentTurn = null;
+        }
+    }
+    /**
+     * Auto-title the session if it's still wearing the default title and we
+     * haven't already attempted generation in this process. Fire-and-forget
+     * (errors are caught, logged, and swallowed) — the user's stream finished
+     * before this runs, so blocking would only delay the broadcast.
+     */
+    maybeGenerateTitle(stream, finishedTurn) {
+        if (this.disableAutoTitle)
+            return;
+        if (this.titleGenerationAttempted.has(stream.sessionId))
+            return;
+        // Check the persisted title now (manual rename takes precedence).
+        const detail = this.store.getSession(stream.sessionId);
+        if (!detail || !isDefaultTitle(detail.title))
+            return;
+        // Find the first user message + the assistant text from the just-finished
+        // turn. We deliberately read user content from SQLite (authoritative)
+        // and assistant content from the in-memory turn buffer (cheap, and
+        // matches what the user just saw).
+        const firstUser = detail.messages.find((m) => m.role === "user");
+        if (!firstUser || !firstUser.content.trim())
+            return;
+        let assistantText = finishedTurn?.assistantText ?? "";
+        if (!assistantText) {
+            // Fallback: pull the most recent assistant message from SQLite. This
+            // can happen if `agent_end` fires for a turn whose buffer was cleared
+            // by an intervening error event, or when this code path is reached
+            // via a synthetic test event.
+            const lastAssistant = [...detail.messages]
+                .reverse()
+                .find((m) => m.role === "assistant");
+            assistantText = lastAssistant?.content ?? "";
+        }
+        // Mark BEFORE awaiting so a second `agent_end` arriving while we're
+        // generating doesn't double-fire. Even if generation throws, we leave
+        // the entry in place — one-shot semantics, no retries.
+        this.titleGenerationAttempted.add(stream.sessionId);
+        void this.runTitleGeneration(stream, firstUser.content, assistantText);
+    }
+    async runTitleGeneration(stream, firstUserMessage, firstAssistantMessage) {
+        try {
+            const title = await generateSessionTitle(firstUserMessage, firstAssistantMessage, {
+                cwd: stream.cwd,
+                agentDir: this.agentDir,
+                llmCall: this.titleLlmCall,
+            });
+            if (this.disposed)
+                return;
+            if (!title)
+                return;
+            // Re-check the title hasn't been changed underneath us while we were
+            // waiting on the LLM (e.g. user manually renamed mid-generation).
+            const current = this.store.getSession(stream.sessionId);
+            if (!current || !isDefaultTitle(current.title))
+                return;
+            const updated = this.store.renameSession(stream.sessionId, title);
+            if (!updated)
+                return;
+            // Broadcast to every subscriber of this session so all open tabs
+            // update their sidebar in real time. The wire event is independent
+            // of the in-flight turn lifecycle, so it's safe to fire post-agent_end.
+            this.broadcast(stream, {
+                type: "session_renamed",
+                sessionId: stream.sessionId,
+                title: updated.title,
+            });
+            // Cross-tab fan-out hint: tabs that don't have THIS session open
+            // (and so don't have a per-session ws subscription) still want to
+            // refresh their sidebar to show the new title. Best-effort; a
+            // failed publish never undoes the rename.
+            if (this.publishMetaEvent) {
+                try {
+                    this.publishMetaEvent({
+                        type: "session_renamed",
+                        projectId: updated.projectId,
+                        sessionId: stream.sessionId,
+                    });
+                }
+                catch (err) {
+                    const msg = err instanceof Error ? err.message : String(err);
+                    console.warn(`[spectral] warn: meta publish for auto-title failed: ${msg}`);
+                }
+            }
+        }
+        catch (err) {
+            // Defensive: generateSessionTitle already swallows LLM errors. This
+            // catches any unexpected throw from rename/broadcast so the manager
+            // is never destabilized by a background title task.
+            const msg = err instanceof Error ? err.message : String(err);
+            console.warn(`[spectral] warn: auto-title pipeline failed: ${msg}`);
+        }
+    }
+    broadcast(stream, event) {
+        const dead = [];
+        for (const sub of stream.subscribers) {
+            if (!sub.isOpen()) {
+                dead.push(sub);
+                continue;
+            }
+            try {
+                sub.send(event);
+            }
+            catch (err) {
+                console.error(`[spectral] error: failed to send WS frame to subscriber: ${err instanceof Error ? err.message : String(err)}`);
+                dead.push(sub);
+            }
+        }
+        for (const sub of dead)
+            stream.subscribers.delete(sub);
+    }
+}
+function isReplayable(event) {
+    return (event.type === "message_start" ||
+        event.type === "text_delta" ||
+        event.type === "thinking_delta" ||
+        event.type === "tool_call" ||
+        event.type === "tool_result" ||
+        event.type === "message_end" ||
+        event.type === "error");
+}
+function snapshotTurn(turn) {
+    // Defensive copy of the events array so the snapshot can't be mutated by
+    // subsequent buffer pushes. Individual event objects are safe to share —
+    // they're never mutated after creation.
+    return {
+        turnId: turn.turnId,
+        startedAt: turn.startedAt,
+        events: turn.events.slice(),
+    };
+}