pig/apps/piggy/src/chat-server.ts

import { timingSafeEqual } from 'node:crypto';
import { createServer, type IncomingMessage, type Server, type ServerResponse } from 'node:http';
import { z } from 'zod';
import { PIGGY_PAGE_ROUTES, PIGGY_RECORD_TYPES } from '@pig/core';
import type { Database } from '@pig/db';
import {
  PrimeOpenAIChatProvider,
  type PiggyChatEvent,
  type PiggyChatRequest,
} from './chat';
import { createInteractivePigTools } from './chat-tools';

/**
 * Derived from the @pig/core tuples rather than retyped, because this schema
 * is `.strict()` and so is the relay's: a context shape one of them has not
 * been told about is a 400, not a degraded answer. `route` is a closed set
 * because a docked panel publishes it on every navigation, and free text there
 * would put arbitrary client strings into a model prompt on every page change.
 */
const contextSchema = z.discriminatedUnion('type', [
  z
    .object({
      type: z.enum(PIGGY_RECORD_TYPES),
      id: z.string().uuid(),
      label: z.string().max(240).optional(),
    })
    .strict(),
  z
    .object({
      type: z.literal('page'),
      route: z.enum(PIGGY_PAGE_ROUTES),
      label: z.string().max(240).optional(),
    })
    .strict(),
]);

export const piggyChatRequestSchema = z
  .object({
    principalUserId: z.string().uuid(),
    message: z.string().trim().min(1).max(4_000),
    history: z
      .array(
        z.object({
          role: z.enum(['user', 'assistant']),
          content: z.string().min(1).max(8_000),
        }),
      )
      .max(20)
      .optional(),
    context: contextSchema.optional(),
  })
  .strict();

interface ChatRunner {
  readonly model: string;
  run(request: PiggyChatRequest): AsyncIterable<PiggyChatEvent>;
}

export interface PiggyChatServerOptions {
  host?: string;
  port: number;
  internalToken: string;
  provider: ChatRunner;
  allowNonLoopback?: boolean;
}

export function startPiggyChatServer(
  db: Database,
  options: PiggyChatServerOptions,
): Server {
  const host = options.host ?? '127.0.0.1';
  if (!isLoopback(host) && !options.allowNonLoopback) {
    throw new Error('Piggy chat must bind to loopback; expose it only through the authenticated CRM API.');
  }
  if (options.internalToken.length < 32) {
    throw new Error('PIGGY_INTERNAL_TOKEN must contain at least 32 characters.');
  }

  const server = createServer(async (request, response) => {
    if (request.method !== 'POST' || request.url !== '/internal/chat') {
      response.writeHead(404).end();
      return;
    }
    if (!tokenMatches(request.headers.authorization, options.internalToken)) {
      response.writeHead(401, { 'content-type': 'application/json' });
      response.end(JSON.stringify({ error: 'Unauthorised internal request.' }));
      return;
    }

    try {
      const body = piggyChatRequestSchema.parse(JSON.parse(await readBoundedBody(request, 32_768)));
      const abort = new AbortController();
      response.on('close', () => abort.abort());
      response.writeHead(200, {
        'content-type': 'application/x-ndjson; charset=utf-8',
        'cache-control': 'no-cache, no-transform',
        'x-content-type-options': 'nosniff',
      });

      for await (const event of options.provider.run({
        message: body.message,
        history: body.history,
        context: body.context,
        tools: createInteractivePigTools(db, body.context),
        signal: abort.signal,
      })) {
        response.write(`${JSON.stringify(event)}\n`);
      }
      response.end();
    } catch (error) {
      const invalidRequest = error instanceof z.ZodError;
      const message = invalidRequest ? 'Invalid Piggy chat request.' : 'Piggy chat failed.';
      if (!response.headersSent) {
        response.writeHead(invalidRequest ? 400 : 500, {
          'content-type': 'application/json',
        });
        response.end(JSON.stringify({ error: message }));
        return;
      }
      response.end(`${JSON.stringify({ type: 'error', message })}\n`);
    }
  });
  server.listen(options.port, host);
  return server;
}

export function createPrimeChatProvider(options: ConstructorParameters<typeof PrimeOpenAIChatProvider>[0]) {
  return new PrimeOpenAIChatProvider(options);
}

function tokenMatches(header: string | undefined, expected: string): boolean {
  const supplied = header?.startsWith('Bearer ') ? header.slice(7) : '';
  const suppliedBytes = Buffer.from(supplied);
  const expectedBytes = Buffer.from(expected);
  return (
    suppliedBytes.length === expectedBytes.length &&
    timingSafeEqual(suppliedBytes, expectedBytes)
  );
}

async function readBoundedBody(request: IncomingMessage, maximumBytes: number): Promise<string> {
  const chunks: Buffer[] = [];
  let size = 0;
  for await (const chunk of request) {
    const bytes = Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk);
    size += bytes.length;
    if (size > maximumBytes) throw new Error('Piggy chat request is too large.');
    chunks.push(bytes);
  }
  return Buffer.concat(chunks).toString('utf8');
}

function isLoopback(host: string): boolean {
  return host === '127.0.0.1' || host === '::1' || host === 'localhost';
}