boocode/apps/server/src/routes/chats.ts

import type { FastifyInstance } from 'fastify';
import { z } from 'zod';
import crypto from 'node:crypto';
import type { Sql } from '../db.js';
import type { Config } from '../config.js';
import type { Broker } from '../services/broker.js';
import type { Chat, Message } from '../types/api.js';
import { getModelContext } from '../services/model-context.js';
import { notifyCoderClose } from '../services/coder-notify.js';
import { MESSAGE_COLUMNS } from '../services/message-columns.js';
import { formatJson, formatMarkdown } from '../services/export-formatter.js';
export interface CompareHandlers {
  enqueueCompare: (
    sessionId: string,
    chatId: string,
    assistantMessageId: string,
    modelOverride: string,
    compareGroupId: string,
  ) => void;
  cancelInference: (sessionId: string, chatId: string) => Promise<boolean>;
  hasActiveInference: (chatId: string) => boolean;
}

const CreateBody = z.object({
  name: z.string().min(1).max(200).optional(),
});

const PatchBody = z.object({
  name: z.string().min(1).max(200).optional(),
  model: z.string().min(1).optional(),
});

const ForkBody = z.object({
  message_id: z.string().uuid(),
  name: z.string().min(1).max(200).optional(),
});

const DiscardStaleBody = z.object({
  message_id: z.string().uuid(),
});

const STALE_MIN_AGE_SECONDS = 60;

const CompareBody = z.object({
  message: z.string().min(1).max(64_000),
  models: z.array(z.string().min(1)).min(2).max(3),
});

export function registerChatRoutes(
  app: FastifyInstance,
  sql: Sql,
  broker: Broker,
  config?: Config,
  compareHandlers?: CompareHandlers,
): void {
  app.get<{ Params: { id: string }; Querystring: { status?: string } }>(
    '/api/sessions/:id/chats',
    async (req, reply) => {
      const session = await sql`SELECT id FROM sessions WHERE id = ${req.params.id}`;
      if (session.length === 0) {
        reply.code(404);
        return { error: 'session not found' };
      }
      const status = req.query.status === 'archived' ? 'archived' : 'open';
      // Enriched list: computed per-chat fields via LATERAL joins.
      const rows = await sql<Chat[]>`
        SELECT
          c.id, c.session_id, c.name, c.status, c.created_at, c.updated_at,
          COALESCE(mc.cnt, 0)::int AS message_count,
          lp.preview AS last_message_preview,
          ec.tokens AS effective_context_tokens
        FROM chats c
        LEFT JOIN LATERAL (
          SELECT COUNT(*) AS cnt FROM messages WHERE chat_id = c.id
        ) mc ON TRUE
        LEFT JOIN LATERAL (
          SELECT LEFT(BTRIM(REGEXP_REPLACE(content, E'[\\n\\r]+', ' ', 'g')), 80) AS preview
          FROM messages
          WHERE chat_id = c.id AND kind = 'message' AND content <> ''
          ORDER BY created_at DESC
          LIMIT 1
        ) lp ON TRUE
        LEFT JOIN LATERAL (
          SELECT ctx_used AS tokens
          FROM messages
          WHERE chat_id = c.id AND kind = 'message' AND role = 'assistant'
            AND status = 'complete' AND ctx_used IS NOT NULL
          ORDER BY created_at DESC
          LIMIT 1
        ) ec ON TRUE
        WHERE c.session_id = ${req.params.id} AND c.status = ${status}
        ORDER BY c.updated_at DESC
      `;
      // v1.11.5: enrich each chat with its model's context window so the
      // ContextBar can render a zero-state (and the auto-compaction threshold
      // tooltip) before the first assistant message lands. All chats in a
      // session share the session's model, so we do ONE getModelContext
      // lookup and apply the result to the whole list. Failed lookups
      // (model unknown, llama-swap down) yield null and the frontend falls
      // through to the "model context unknown" placeholder.
      const sessRow = await sql<{ model: string | null }[]>`
        SELECT model FROM sessions WHERE id = ${req.params.id}
      `;
      const sessionModel = sessRow[0]?.model ?? null;
      const mctx = sessionModel ? await getModelContext(sessionModel) : null;
      const modelContextLimit = mctx?.n_ctx ?? null;
      return rows.map((r) => ({ ...r, model_context_limit: modelContextLimit }));
    }
  );

  app.post<{ Params: { id: string } }>(
    '/api/sessions/:id/chats',
    async (req, reply) => {
      const parsed = CreateBody.safeParse(req.body ?? {});
      if (!parsed.success) {
        reply.code(400);
        return { error: 'invalid body', details: parsed.error.flatten() };
      }
      const session = await sql`SELECT id FROM sessions WHERE id = ${req.params.id}`;
      if (session.length === 0) {
        reply.code(404);
        return { error: 'session not found' };
      }
      const [chat] = await sql<Chat[]>`
        INSERT INTO chats (session_id, name, status)
        VALUES (${req.params.id}, ${parsed.data.name ?? null}, 'open')
        RETURNING id, session_id, name, status, created_at, updated_at
      `;
      broker.publishUserFrame('default', {
        type: 'chat_created',
        chat: chat!,
        session_id: req.params.id,
      });
      reply.code(201);
      return chat;
    }
  );

  app.patch<{ Params: { id: string } }>(
    '/api/chats/:id',
    async (req, reply) => {
      const parsed = PatchBody.safeParse(req.body ?? {});
      if (!parsed.success) {
        reply.code(400);
        return { error: 'invalid body', details: parsed.error.flatten() };
      }
      const { name, model } = parsed.data;
      const sets: Array<ReturnType<typeof sql>> = [sql`updated_at = clock_timestamp()`];
      if (name !== undefined) sets.push(sql`name = ${name}`);
      if (model !== undefined) sets.push(sql`model = ${model}`);
      const rows = await sql<Chat[]>`
        UPDATE chats
        SET ${(sql as any).join(sets, sql`, `)}
        WHERE id = ${req.params.id}
        RETURNING id, session_id, name, model, status, created_at, updated_at
      `;
      if (rows.length === 0) {
        reply.code(404);
        return { error: 'chat not found' };
      }
      const chat = rows[0]!;
      broker.publishUserFrame('default', {
        type: 'chat_updated',
        chat_id: chat.id,
        session_id: chat.session_id,
        name: chat.name,
        updated_at: chat.updated_at,
      });
      return chat;
    }
  );

  // v1.9: bulk-archive every open chat in a session. Mirrors the single
  // /chats/:id/archive shape — N chat_archived frames published, useSidebar
  // reducer handles each via the existing case.
  app.post<{ Params: { id: string } }>(
    '/api/sessions/:id/chats/archive-all',
    async (req, reply) => {
      const session = await sql`SELECT id FROM sessions WHERE id = ${req.params.id}`;
      if (session.length === 0) {
        reply.code(404);
        return { error: 'session not found' };
      }
      const rows = await sql<{ id: string }[]>`
        UPDATE chats
           SET status = 'archived', updated_at = clock_timestamp()
         WHERE session_id = ${req.params.id} AND status = 'open'
         RETURNING id
      `;
      const ids = rows.map((r) => r.id);
      for (const id of ids) {
        broker.publishUserFrame('default', {
          type: 'chat_archived',
          chat_id: id,
          session_id: req.params.id,
        });
        // Fire-and-forget per archived chat: tear down its warm agent backends
        // on the coder. Best-effort — never blocks/fails the bulk archive.
        void notifyCoderClose('chat', id, req.log);
      }
      return { archived: ids.length, ids };
    }
  );

  // v1.9: count helper for the confirm dialog.
  app.get<{ Params: { id: string } }>(
    '/api/sessions/:id/chats/open-count',
    async (req, reply) => {
      const session = await sql`SELECT id FROM sessions WHERE id = ${req.params.id}`;
      if (session.length === 0) {
        reply.code(404);
        return { error: 'session not found' };
      }
      const rows = await sql<{ count: number }[]>`
        SELECT COUNT(*)::int AS count
          FROM chats
         WHERE session_id = ${req.params.id} AND status = 'open'
      `;
      return { count: rows[0]?.count ?? 0 };
    }
  );

  app.post<{ Params: { id: string } }>(
    '/api/chats/:id/archive',
    async (req, reply) => {
      const rows = await sql<{ id: string; session_id: string }[]>`
        UPDATE chats SET status = 'archived', updated_at = clock_timestamp()
        WHERE id = ${req.params.id} AND status = 'open'
        RETURNING id, session_id
      `;
      if (rows.length === 0) {
        reply.code(404);
        return { error: 'chat not found or already archived' };
      }
      const row = rows[0]!;
      broker.publishUserFrame('default', {
        type: 'chat_archived',
        chat_id: row.id,
        session_id: row.session_id,
      });
      // Fire-and-forget: tear down this chat's warm agent backends + (last-chat)
      // worktree on the coder. Best-effort — never blocks/fails the archive.
      void notifyCoderClose('chat', row.id, req.log);
      reply.code(204);
      return null;
    }
  );

  app.post<{ Params: { id: string } }>(
    '/api/chats/:id/unarchive',
    async (req, reply) => {
      const rows = await sql<Chat[]>`
        UPDATE chats SET status = 'open', updated_at = clock_timestamp()
        WHERE id = ${req.params.id} AND status = 'archived'
        RETURNING id, session_id, name, status, created_at, updated_at
      `;
      if (rows.length === 0) {
        reply.code(404);
        return { error: 'chat not found or not archived' };
      }
      const chat = rows[0]!;
      broker.publishUserFrame('default', { type: 'chat_unarchived', chat });
      return chat;
    }
  );

  app.delete<{ Params: { id: string } }>(
    '/api/chats/:id',
    async (req, reply) => {
      const result = await sql<{ id: string; session_id: string }[]>`
        DELETE FROM chats WHERE id = ${req.params.id}
        RETURNING id, session_id
      `;
      if (result.length === 0) {
        reply.code(404);
        return { error: 'chat not found' };
      }
      const row = result[0]!;
      broker.publishUserFrame('default', {
        type: 'chat_deleted',
        chat_id: row.id,
        session_id: row.session_id,
      });
      // Fire-and-forget: tear down this chat's warm agent backends + (last-chat)
      // worktree on the coder. Best-effort — never blocks/fails the delete.
      void notifyCoderClose('chat', row.id, req.log);
      reply.code(204);
      return null;
    }
  );

  app.post<{ Params: { id: string } }>(
    '/api/chats/:id/fork',
    async (req, reply) => {
      const parsed = ForkBody.safeParse(req.body ?? {});
      if (!parsed.success) {
        reply.code(400);
        return { error: 'invalid body', details: parsed.error.flatten() };
      }

      const sourceRows = await sql<Chat[]>`
        SELECT id, session_id, name, status, created_at, updated_at
        FROM chats WHERE id = ${req.params.id}
      `;
      if (sourceRows.length === 0) {
        reply.code(404);
        return { error: 'chat not found' };
      }
      const source = sourceRows[0]!;

      const targetRows = await sql<{ created_at: string; status: string }[]>`
        SELECT created_at, status FROM messages
        WHERE chat_id = ${source.id} AND id = ${parsed.data.message_id}
      `;
      if (targetRows.length === 0) {
        reply.code(404);
        return { error: 'message not found in chat' };
      }
      const target = targetRows[0]!;
      if (target.status !== 'complete') {
        reply.code(400);
        return { error: 'can only fork from completed messages' };
      }

      const newName = parsed.data.name ?? `${source.name ?? 'Chat'} (fork)`;

      const newChat = await sql.begin(async (tx) => {
        const [chat] = await tx<Chat[]>`
          INSERT INTO chats (session_id, name, status)
          VALUES (${source.session_id}, ${newName}, 'open')
          RETURNING id, session_id, name, status, created_at, updated_at
        `;
        await tx`
          INSERT INTO messages (
            session_id, chat_id, role, content, kind,
            status, tokens_used, ctx_used, ctx_max, started_at, finished_at,
            created_at, metadata
          )
          SELECT
            ${source.session_id}, ${chat!.id}, role, content, kind,
            status,
            tokens_used, ctx_used, ctx_max, started_at, finished_at,
            clock_timestamp() + (
              ROW_NUMBER() OVER (ORDER BY created_at ASC, id ASC) * INTERVAL '1 microsecond'
            ),
            metadata
          FROM messages
          WHERE chat_id = ${source.id}
            AND created_at <= ${target.created_at}::timestamptz
            AND status = 'complete'
        `;
        // v1.13.0: clone message_parts for the forked messages. Source and
        // destination preserve ordering (the INSERT above orders by created_at,
        // id) so a ROW_NUMBER pairing maps source.id → dest.id deterministically.
        await tx`
          WITH src AS (
            SELECT id, ROW_NUMBER() OVER (ORDER BY created_at ASC, id ASC) AS rn
            FROM messages
            WHERE chat_id = ${source.id}
              AND created_at <= ${target.created_at}::timestamptz
              AND status = 'complete'
          ),
          dst AS (
            SELECT id, ROW_NUMBER() OVER (ORDER BY created_at ASC, id ASC) AS rn
            FROM messages
            WHERE chat_id = ${chat!.id}
          )
          INSERT INTO message_parts (message_id, sequence, kind, payload)
          SELECT dst.id, p.sequence, p.kind, p.payload
          FROM message_parts p
          JOIN src ON p.message_id = src.id
          JOIN dst ON dst.rn = src.rn
        `;
        return chat!;
      });

      broker.publishUserFrame('default', {
        type: 'chat_created',
        chat: newChat,
        session_id: source.session_id,
      });
      reply.code(201);
      return newChat;
    }
  );

  // v1.12.3: explicit recovery from a stuck-streaming assistant row. The
  // frontend gates this behind a 60s no-token-activity timer; the server
  // re-checks the age and current status for safety. Non-streaming rows
  // return 409 (frontend race; idempotent retry is fine).
  app.post<{ Params: { id: string } }>(
    '/api/chats/:id/discard_stale',
    async (req, reply) => {
      const parsed = DiscardStaleBody.safeParse(req.body ?? {});
      if (!parsed.success) {
        reply.code(400);
        return { error: 'invalid body', details: parsed.error.flatten() };
      }
      const rows = await sql<{
        id: string;
        session_id: string;
        chat_id: string;
        status: string;
        age_seconds: number;
      }[]>`
        SELECT id, session_id, chat_id, status,
               EXTRACT(EPOCH FROM (clock_timestamp() - created_at))::int AS age_seconds
        FROM messages
        WHERE id = ${parsed.data.message_id} AND chat_id = ${req.params.id}
      `;
      if (rows.length === 0) {
        reply.code(404);
        return { error: 'message not found in chat' };
      }
      const msg = rows[0]!;
      if (msg.status !== 'streaming') {
        reply.code(409);
        return { error: 'message is no longer streaming', current_status: msg.status };
      }
      if (msg.age_seconds < STALE_MIN_AGE_SECONDS) {
        reply.code(409);
        return { error: 'message is not stale yet', age_seconds: msg.age_seconds };
      }
      const updated = await sql<{ id: string }[]>`
        UPDATE messages
        SET status = 'failed',
            content = COALESCE(content, ''),
            finished_at = clock_timestamp()
        WHERE id = ${msg.id} AND status = 'streaming'
        RETURNING id
      `;
      if (updated.length === 0) {
        // Race: the row flipped out of 'streaming' between our SELECT and UPDATE.
        reply.code(409);
        return { error: 'message status changed mid-request' };
      }
      // v1.13.20: re-fetch via messages_with_parts so the returned shape
      // carries parts-synthesized tool_calls / tool_results. The dropped
      // legacy columns can no longer be selected directly.
      const refreshed = await sql<Message[]>`
        SELECT * FROM messages_with_parts WHERE id = ${msg.id}
      `;
      broker.publishUserFrame('default', {
        type: 'chat_status',
        chat_id: msg.chat_id,
        status: 'idle',
        at: new Date().toISOString(),
      });
      broker.publishFrame(msg.session_id, {
        type: 'message_complete',
        message_id: msg.id,
        chat_id: msg.chat_id,
      });
      return refreshed[0];
    }
  );

  app.get<{ Params: { id: string } }>(
    '/api/chats/:id/messages',
    async (req, reply) => {
      const chat = await sql`SELECT id FROM chats WHERE id = ${req.params.id}`;
      if (chat.length === 0) {
        reply.code(404);
        return { error: 'chat not found' };
      }
      // v1.13.1-B: reads tool_calls/tool_results via the parts-merged view.
      const rows = await sql<Message[]>`
        SELECT ${sql.unsafe(MESSAGE_COLUMNS)}
        FROM messages_with_parts
        WHERE chat_id = ${req.params.id}
        ORDER BY created_at ASC, id ASC
      `;
      return rows;
    }
  );

  app.get<{ Params: { id: string }; Querystring: { format?: string } }>(
    '/api/chats/:id/export',
    async (req, reply) => {
      const format = req.query.format ?? 'json';
      if (format !== 'json' && format !== 'markdown') {
        reply.code(400);
        return { error: 'format must be json or markdown' };
      }

      const chat = await sql<Chat[]>`SELECT * FROM chats WHERE id = ${req.params.id}`;
      if (chat.length === 0) {
        reply.code(404);
        return { error: 'chat not found' };
      }

      const messages = await sql<Message[]>`
        SELECT ${sql.unsafe(MESSAGE_COLUMNS)}
        FROM messages_with_parts
        WHERE chat_id = ${req.params.id}
        ORDER BY created_at ASC, id ASC
      `;

      if (format === 'markdown') {
        reply.header('Content-Type', 'text/markdown');
        return formatMarkdown(chat[0]!, messages, chat[0]!.model);
      }

      reply.header('Content-Type', 'application/json');
      return formatJson(chat[0]!, messages, chat[0]!.model);
    }
  );

  // v2.8-compare: send the same message to N models and stream back parallel
  // responses. Creates N assistant messages (one per model) and launches N
  // parallel inference runs with model overrides. Each publishes frames
  // scoped to the shared compare_group_id so the frontend can group them.
  if (config && compareHandlers) {
    app.post<{ Params: { id: string } }>(
      '/api/chats/:id/compare',
      async (req, reply) => {
        const parsed = CompareBody.safeParse(req.body);
        if (!parsed.success) {
          reply.code(400);
          return { error: 'invalid body', details: parsed.error.flatten() };
        }

        const { message, models } = parsed.data;

        // Check for active inference first.
        if (compareHandlers.hasActiveInference(req.params.id)) {
          reply.code(409);
          return { error: 'chat is currently streaming; stop it first' };
        }

        const chatRows = await sql<Chat[]>`
          SELECT id, session_id FROM chats WHERE id = ${req.params.id} AND status = 'open'
        `;
        if (chatRows.length === 0) {
          reply.code(404);
          return { error: 'chat not found' };
        }
        const chat = chatRows[0]!;
        const sessionId = chat.session_id;
        const compareGroupId = crypto.randomUUID();

        // Insert user message + N assistant messages in a single transaction.
        const result = await sql.begin(async (tx) => {
          const [userMsg] = await tx<{ id: string }[]>`
            INSERT INTO messages (session_id, chat_id, role, content, status, created_at, metadata)
            VALUES (${sessionId}, ${chat.id}, 'user', ${message}, 'complete', clock_timestamp(), NULL)
            RETURNING id
          `;

          const responses: Array<{ model: string; assistant_message_id: string }> = [];
          for (const model of models) {
            const [asst] = await tx<{ id: string }[]>`
              INSERT INTO messages (session_id, chat_id, role, content, status, created_at, metadata)
              VALUES (
                ${sessionId}, ${chat.id}, 'assistant', '', 'streaming', clock_timestamp(),
                ${tx.json({ compare_group_id: compareGroupId, model } as never)}
              )
              RETURNING id
            `;
            responses.push({ model, assistant_message_id: asst!.id });
          }

          await tx`UPDATE sessions SET updated_at = clock_timestamp() WHERE id = ${sessionId}`;
          await tx`UPDATE chats SET updated_at = clock_timestamp() WHERE id = ${chat.id}`;

          return { user_message_id: userMsg!.id, responses };
        });

        // Publish user message frames.
        broker.publishFrame(sessionId, {
          type: 'message_started',
          message_id: result.user_message_id,
          chat_id: chat.id,
          role: 'user',
        });
        broker.publishFrame(sessionId, {
          type: 'delta',
          message_id: result.user_message_id,
          chat_id: chat.id,
          content: message,
        });
        broker.publishFrame(sessionId, {
          type: 'message_complete',
          message_id: result.user_message_id,
          chat_id: chat.id,
        });

        // Enqueue N parallel inference runs with model overrides.
        for (const resp of result.responses) {
          compareHandlers.enqueueCompare(
            sessionId, chat.id, resp.assistant_message_id, resp.model, compareGroupId,
          );
        }

        reply.code(202);
        return { compare_group_id: compareGroupId, ...result };
      },
    );
  }
}