spb/chat-spboucher Public

1// Author: Simon-Pierre Boucher2// Contact: contact@spboucher.ai3// Project: chat.spboucher.ai45// Generation runner: explicit state machine per generation6// (queued → starting → streaming → completed | cancelled | failed),7// SSE emission with keep-alives, incremental persistence, cancellation registry.89import crypto from "node:crypto";10import { getDb } from "@/lib/db/database";11import { streamGeneration, estimateCost } from "@/lib/openrouter";12import type { ChatStreamEvent } from "@/lib/openrouter";13import { compileContext } from "@/lib/context";14import {15  getConversation,16  insertMessage,17  threadToLeaf,18  touchConversation,19  updateConversation,20  type MessageRow,21} from "@/lib/conversations";22import { getModel, touchModelUsage, type CatalogModel } from "@/lib/catalog";2324const KEEPALIVE_MS = 15_000;25const FLUSH_MS = 700;2627// Registry of live generations for cancellation; on globalThis to survive dev HMR.28const registry: Map<string, AbortController> = ((globalThis as Record<string, unknown>).__genRegistry ??=29  new Map()) as Map<string, AbortController>;3031export function cancelGeneration(generationId: string): boolean {32  const ctrl = registry.get(generationId);33  if (!ctrl) return false;34  ctrl.abort();35  return true;36}3738export function generationState(generationId: string) {39  return getDb()40    .prepare("SELECT id, message_id, state, error_code, error_message, finished_at FROM generations WHERE id = ?")41    .get(generationId) as42    | { id: string; message_id: string; state: string; error_code: string | null; error_message: string | null; finished_at: number | null }43    | undefined;44}4546function setGenState(id: string, state: string, errorCode?: string, errorMessage?: string): void {47  const finished = ["completed", "cancelled", "failed"].includes(state) ? Date.now() : null;48  getDb()49    .prepare(50      "UPDATE generations SET state = ?, error_code = COALESCE(?, error_code), error_message = COALESCE(?, error_message), finished_at = COALESCE(?, finished_at) WHERE id = ?"51    )52    .run(state, errorCode ?? null, errorMessage ?? null, finished, id);53}5455export interface StartGenerationArgs {56  conversationId: string;57  parentMessage: MessageRow | null; // the user message the assistant answers58  model: CatalogModel;59}6061export interface StartedGeneration {62  generationId: string;63  assistantMessage: MessageRow;64  stream: ReadableStream<Uint8Array>;65}6667function sseFrame(event: object): string {68  return `data: ${JSON.stringify(event)}\n\n`;69}7071/**72 * Create the assistant placeholder + generation row, then return an SSE stream73 * that runs the generation, persists deltas incrementally, and finalizes state.74 */75export function startGeneration(args: StartGenerationArgs): StartedGeneration {76  const db = getDb();77  const generationId = crypto.randomUUID();78  const { conversationId, parentMessage, model } = args;7980  const assistantMessage = insertMessage({81    conversationId,82    parentId: parentMessage?.id ?? null,83    role: "assistant",84    modelId: model.id,85    modelName: model.name,86    provider: model.provider,87    generationId,88    status: "pending",89  });9091  db.prepare(92    "INSERT INTO generations (id, message_id, conversation_id, model_id, state, created_at) VALUES (?, ?, ?, ?, 'queued', ?)"93  ).run(generationId, assistantMessage.id, conversationId, model.id, Date.now());9495  updateConversation(conversationId, { currentLeafId: assistantMessage.id });96  touchConversation(conversationId);97  touchModelUsage(model.id);9899  const ctrl = new AbortController();100  registry.set(generationId, ctrl);101102  const thread = threadToLeaf(conversationId, parentMessage?.id ?? null);103  const compiled = compileContext(thread, model.contextLength);104105  const encoder = new TextEncoder();106107  const stream = new ReadableStream<Uint8Array>({108    async start(controller) {109      const send = (e: object) => {110        try {111          controller.enqueue(encoder.encode(sseFrame(e)));112        } catch {113          /* client went away; generation continues until aborted or done */114        }115      };116117      // keep-alive comments so ngrok and mobile radios don't kill idle streams118      const keepAlive = setInterval(() => {119        try {120          controller.enqueue(encoder.encode(`: ping\n\n`));121        } catch {122          /* ignore */123        }124      }, KEEPALIVE_MS);125126      let content = "";127      let reasoning = "";128      let dirty = false;129      const flush = () => {130        if (!dirty) return;131        db.prepare("UPDATE messages SET content = ?, reasoning = ?, status = 'streaming' WHERE id = ?").run(132          content,133          reasoning || null,134          assistantMessage.id135        );136        dirty = false;137      };138      const flusher = setInterval(flush, FLUSH_MS);139140      send({141        type: "meta",142        conversationId,143        assistantMessageId: assistantMessage.id,144        generationId,145        model: model.id,146        estimatedPromptTokens: compiled.estimatedPromptTokens,147      });148149      setGenState(generationId, "starting");150151      let finalState: "completed" | "cancelled" | "failed" = "completed";152      let errorMessage: string | null = null;153      let sawUsage = false;154155      try {156        const events = streamGeneration(157          { model: model.id, messages: compiled.messages, signal: ctrl.signal },158          generationId159        );160        for await (const event of events) {161          switch (event.type) {162            case "generation.start":163              setGenState(generationId, "streaming");164              break;165            case "content.delta":166              content += event.text;167              dirty = true;168              break;169            case "reasoning.delta":170              reasoning += event.text;171              dirty = true;172              break;173            case "usage":174              sawUsage = true;175              recordUsage(generationId, model, event);176              break;177            case "generation.error":178              finalState = "failed";179              errorMessage = event.message;180              break;181            case "generation.end":182            case "tool.start":183            case "tool.delta":184              break;185          }186          send(event);187        }188        if (ctrl.signal.aborted) finalState = "cancelled";189      } catch {190        finalState = ctrl.signal.aborted ? "cancelled" : "failed";191        if (finalState === "failed") errorMessage = "The stream failed unexpectedly.";192        send({ type: "generation.error", message: errorMessage ?? "Cancelled", retryable: finalState === "failed" });193      } finally {194        clearInterval(keepAlive);195        clearInterval(flusher);196        registry.delete(generationId);197198        flush();199        db.prepare(200          "UPDATE messages SET content = ?, reasoning = ?, status = ?, error_message = ? WHERE id = ?"201        ).run(content, reasoning || null, finalState, errorMessage, assistantMessage.id);202        setGenState(generationId, finalState, undefined, errorMessage ?? undefined);203        touchConversation(conversationId);204205        // If the provider never reported usage, store an estimate for cost tracking.206        if (!sawUsage && finalState === "completed") {207          recordUsage(generationId, model, {208            type: "usage",209            promptTokens: compiled.estimatedPromptTokens,210            completionTokens: Math.ceil(content.length / 3.6),211            totalTokens: compiled.estimatedPromptTokens + Math.ceil(content.length / 3.6),212          });213        }214215        send({ type: "state.final", state: finalState, messageId: assistantMessage.id });216        try {217          controller.close();218        } catch {219          /* already closed */220        }221      }222    },223    cancel() {224      // Client disconnected (tab closed, network drop). Abort upstream — never keep paying.225      ctrl.abort();226    },227  });228229  return { generationId, assistantMessage, stream };230}231232function recordUsage(233  generationId: string,234  model: CatalogModel,235  usage: Extract<ChatStreamEvent, { type: "usage" }>236): void {237  const estimated = estimateCost(model.pricing, usage.promptTokens, usage.completionTokens);238  getDb()239    .prepare(240      `INSERT INTO generation_usage241        (generation_id, model_id, prompt_tokens, completion_tokens, reasoning_tokens, cached_tokens, total_tokens, estimated_cost_usd, reported_cost_usd, created_at)242       VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`243    )244    .run(245      generationId,246      model.id,247      usage.promptTokens ?? null,248      usage.completionTokens ?? null,249      usage.reasoningTokens ?? null,250      usage.cachedTokens ?? null,251      usage.totalTokens ?? null,252      estimated ?? null,253      usage.cost ?? null,254      Date.now()255    );256}257258/** Guard used by the chat route: conversation must exist. */259export function assertConversation(id: string) {260  const conv = getConversation(id);261  if (!conv) throw new Error("Conversation not found");262  return conv;263}264