Sybil-2/server/src/agent.ts

import type { IncomingMessage } from 'node:http';
import type { Duplex } from 'node:stream';
import type { FastifyInstance } from 'fastify';
import { WebSocket, WebSocketServer } from 'ws';
import { z } from 'zod';
import { requireAdmin } from './auth.js';
import { prisma } from './db.js';
import { env } from './env.js';

const Cursor = z.string().regex(/^\d{1,16}$/).refine(value => Number.isSafeInteger(Number(value)), 'Invalid cursor')
  .transform(value => String(Number(value)));
const PageQuery = z.object({ after: Cursor.optional(), before: Cursor.optional(), limit: z.coerce.number().int().min(1).max(200).optional() })
  .refine(value => !(value.after && value.before), 'Use after or before, not both');
const MessageBody = z.object({ content: z.string().min(1).max(40_000).refine(value => value.trim().length > 0, 'content is required'), requestId: z.string().min(1).max(256) });
const ReadBody = z.object({ cursor: Cursor });
const AuthFrame = z.object({ type: z.literal('auth'), token: z.string().nullable().optional(), after: Cursor.optional() });
const SubscribeFrame = z.object({ type: z.literal('subscribe'), after: Cursor.optional() });
const SubmissionReceipt = z.object({ id: z.number().int().positive().max(Number.MAX_SAFE_INTEGER),
  status: z.enum(['queued', 'placed', 'done', 'unanswered']), messageId: z.string() });

type AgentOptions = { baseUrl?: string; apiKey?: string };
const parse = <T>(schema: z.ZodType<T>, value: unknown): T => {
  const result = schema.safeParse(value);
  if (!result.success) throw Object.assign(new Error(result.error.issues.map(issue => issue.message).join('; ')), { statusCode: 400 });
  return result.data;
};

async function readState() {
  const value = await prisma.agentReadState.findUnique({ where: { id: 'root' } });
  return { cursor: value?.cursor || '0', updatedAt: value?.updatedAt.toISOString() || null };
}

/** Pi owns the conversation; Sybil stores only its shared personal read cursor. */
export async function registerAgentRoutes(app: FastifyInstance, options: AgentOptions = {}) {
  const baseUrl = options.baseUrl || env.PI_AGENT_BASE_URL;
  const apiKey = options.apiKey ?? env.PI_AGENT_API_KEY;
  const proxyJson = async (path: string, method = 'GET', body?: unknown) => {
    if (!apiKey) throw app.httpErrors.serviceUnavailable('Pi agent is not configured');
    let response: Response;
    try {
      response = await fetch(new URL(path, baseUrl), { method,
        headers: { Authorization: `Bearer ${apiKey}`, ...(body === undefined ? {} : { 'Content-Type': 'application/json' }) },
        ...(body === undefined ? {} : { body: JSON.stringify(body) }), signal: AbortSignal.timeout(15_000),
      });
    } catch { throw app.httpErrors.badGateway('Pi agent is unavailable'); }
    let data: unknown;
    try { data = await response.json(); }
    catch { throw app.httpErrors.badGateway('Pi agent returned invalid JSON'); }
    return { status: response.status, data };
  };

  app.get('/v1/agent', async (req, reply) => {
    reply.header('Cache-Control', 'no-store');
    requireAdmin(req);
    const result = await proxyJson('/api/conversations/root');
    reply.code(result.status);
    if (result.status !== 200) return result.data;
    return { ...result.data as object, websocket: '/v1/agent/ws' };
  });
  app.get('/v1/agent/messages', async (req, reply) => {
    reply.header('Cache-Control', 'no-store');
    requireAdmin(req);
    const query = parse(PageQuery, req.query);
    const params = new URLSearchParams();
    for (const [name, value] of Object.entries(query)) if (value !== undefined) params.set(name, String(value));
    const result = await proxyJson(`/api/conversations/root/messages?${params}`);
    reply.code(result.status);
    return result.data;
  });
  app.post('/v1/agent/messages', { bodyLimit: 256 * 1024 }, async (req, reply) => {
    reply.header('Cache-Control', 'no-store');
    requireAdmin(req);
    const body = parse(MessageBody, req.body);
    const result = await proxyJson('/api/conversations/root/messages', 'POST', { ...body, channel: 'sybil', whenBusy: 'followUp', wait: false });
    if (result.status >= 200 && result.status < 300) {
      const receipt = SubmissionReceipt.safeParse(result.data);
      if (!receipt.success || receipt.data.messageId !== `inbound:sybil:${body.requestId}`) {
        throw app.httpErrors.badGateway('Pi agent returned an invalid message receipt');
      }
    }
    reply.code(result.status);
    return result.data;
  });
  app.get('/v1/agent/read-state', async (req, reply) => {
    reply.header('Cache-Control', 'no-store');
    requireAdmin(req);
    return readState();
  });
  app.put('/v1/agent/read-state', async (req, reply) => {
    reply.header('Cache-Control', 'no-store');
    requireAdmin(req);
    const { cursor } = parse(ReadBody, req.body);
    // One SQLite statement makes concurrent tabs monotonic. Never write an older
    // cursor over one saved by another connection between a read and an update.
    await prisma.$executeRaw`INSERT INTO "AgentReadState" ("id", "cursor", "updatedAt") VALUES ('root', ${cursor}, CURRENT_TIMESTAMP)
      ON CONFLICT("id") DO UPDATE SET
        "cursor" = CASE WHEN CAST(excluded."cursor" AS INTEGER) > CAST("AgentReadState"."cursor" AS INTEGER) THEN excluded."cursor" ELSE "AgentReadState"."cursor" END,
        "updatedAt" = CASE WHEN CAST(excluded."cursor" AS INTEGER) > CAST("AgentReadState"."cursor" AS INTEGER) THEN CURRENT_TIMESTAMP ELSE "AgentReadState"."updatedAt" END`;
    return readState();
  });

  attachAgentSocket(app, baseUrl, apiKey);
}

function attachAgentSocket(app: FastifyInstance, baseUrl: string, apiKey?: string) {
  const sockets = new WebSocketServer({ noServer: true, maxPayload: 16 * 1024 });
  const cleanup = new Set<() => void>();
  const upgrade = (request: IncomingMessage, socket: Duplex, head: Buffer) => {
    let url: URL;
    try { url = new URL(request.url || '/', 'http://localhost'); }
    catch { socket.end('HTTP/1.1 400 Bad Request\r\nConnection: close\r\n\r\n'); return; }
    if (url.pathname !== '/v1/agent/ws') { socket.end('HTTP/1.1 404 Not Found\r\nConnection: close\r\n\r\n'); return; }
    // Credentials must never enter URLs and access logs.
    if (url.search) { socket.end('HTTP/1.1 400 Bad Request\r\nConnection: close\r\n\r\n'); return; }
    sockets.handleUpgrade(request, socket, head, browser => {
      let upstream: WebSocket | undefined;
      let authenticated = false, disposed = false, alive = true;
      const send = (value: unknown) => {
        if (browser.readyState !== WebSocket.OPEN) return;
        if (browser.bufferedAmount > 2 * 1024 * 1024) { browser.close(1013, 'Slow consumer; reconnect with cursor'); return; }
        browser.send(typeof value === 'string' ? value : JSON.stringify(value));
      };
      const authTimeout = setTimeout(() => browser.close(4401, 'Authentication required'), 5000);
      const heartbeat = setInterval(() => { if (!alive) browser.terminate(); else { alive = false; browser.ping(); } }, 25_000);
      browser.on('pong', () => { alive = true; });
      const dispose = () => {
        if (disposed) return;
        disposed = true; clearTimeout(authTimeout); clearInterval(heartbeat); cleanup.delete(dispose);
        upstream?.terminate(); browser.terminate();
      };
      cleanup.add(dispose);
      browser.on('close', dispose); browser.on('error', dispose);
      browser.on('message', (raw, isBinary) => {
        try {
          if (isBinary) throw new Error('JSON text frames are required');
          const frame = JSON.parse(raw.toString());
          if (!authenticated) {
            const auth = parse(AuthFrame, frame);
            if (env.ADMIN_TOKEN && auth.token !== env.ADMIN_TOKEN) { browser.close(4401, 'Invalid token'); return; }
            if (!apiKey) { browser.close(4503, 'Pi agent is not configured'); return; }
            authenticated = true; clearTimeout(authTimeout);
            const target = new URL('/api/ws', baseUrl); target.protocol = target.protocol === 'https:' ? 'wss:' : 'ws:';
            upstream = new WebSocket(target, { headers: { Authorization: `Bearer ${apiKey}` }, handshakeTimeout: 10_000, maxPayload: 16 * 1024 * 1024 });
            upstream.on('open', () => upstream!.send(JSON.stringify({ type: 'subscribe', ...(auth.after === undefined ? {} : { after: auth.after }) })));
            upstream.on('message', (data, binary) => {
              if (binary) { browser.close(1011, 'Invalid Pi frame'); return; }
              send(data.toString());
            });
            upstream.on('close', () => { if (!disposed) browser.close(1012, 'Pi agent disconnected; reconnect with cursor'); });
            upstream.on('error', () => { send({ type: 'error', error: 'Pi agent is unavailable' }); browser.close(1011, 'Pi agent is unavailable'); });
            return;
          }
          if (frame.type === 'ping') { send({ type: 'pong' }); return; }
          const subscribe = parse(SubscribeFrame, frame);
          if (upstream?.readyState !== WebSocket.OPEN) { send({ type: 'error', error: 'Pi agent is still connecting' }); return; }
          upstream.send(JSON.stringify(subscribe));
        } catch (error) {
          if (!authenticated) browser.close(4400, 'Invalid authentication frame');
          else send({ type: 'error', error: error instanceof Error ? error.message : 'Invalid frame' });
        }
      });
    });
  };
  app.server.on('upgrade', upgrade);
  app.addHook('preClose', async () => {
    app.server.off('upgrade', upgrade);
    for (const stop of cleanup) stop();
    for (const client of sockets.clients) client.terminate();
    sockets.close();
  });
}