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

import { timingSafeEqual } from 'node:crypto';
import { createServer, type IncomingMessage, type Server, type ServerResponse } from 'node:http';
import { eq } from 'drizzle-orm';
import { z } from 'zod';
import { PIGGY_PAGE_ROUTES, PIGGY_RECORD_TYPES } from '@pig/core';
import { agentRuns, 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>;
}

/**
 * What a token costs, in cents per million, so the cost arithmetic stays
 * integral: micro-cents = tokens x cents-per-million. Omitted, the tokens are
 * still recorded and the cost is left null — an unpriced run is honest, an
 * invented price is not.
 */
export interface ChatTokenPricing {
  inputCentsPerMillionTokens: number;
  outputCentsPerMillionTokens: number;
}

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

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) => {
    // Unauthenticated on purpose: a container healthcheck and a load balancer
    // have no token, and this says nothing an attacker on loopback could not
    // learn by watching the port.
    if (request.method === 'GET' && request.url === '/internal/health') {
      respondJson(response, 200, {
        ok: true,
        service: 'piggy-chat',
        model: options.provider.model,
      });
      return;
    }
    if (request.method !== 'POST' || request.url !== '/internal/chat') {
      response.writeHead(404).end();
      return;
    }
    if (!tokenMatches(request.headers.authorization, options.internalToken)) {
      respondJson(response, 401, { error: 'Unauthorised internal request.' });
      return;
    }

    let body: z.infer<typeof piggyChatRequestSchema>;
    try {
      body = piggyChatRequestSchema.parse(JSON.parse(await readBoundedBody(request, 32_768)));
    } catch {
      // Every failure reachable here — an oversized body, malformed JSON, a
      // context arm this schema has not been told about — genuinely is the
      // caller's. Nothing below may borrow this message: a ZodError raised
      // mid-stream is an upstream fault, and reporting it as invalid input
      // told the user their question was malformed when it was not.
      respondJson(response, 400, { error: 'Invalid Piggy chat request.' });
      return;
    }

    const abort = new AbortController();
    response.on('close', () => abort.abort());
    const run = await startChatRun(db, {
      principalUserId: body.principalUserId,
      model: options.provider.model,
      message: body.message,
      context: body.context,
      historyTurns: body.history?.length ?? 0,
    });
    const spend: ChatRunOutcome = { toolCalls: 0 };

    response.writeHead(200, {
      'content-type': 'application/x-ndjson; charset=utf-8',
      'cache-control': 'no-cache, no-transform',
      'x-content-type-options': 'nosniff',
    });

    try {
      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,
      })) {
        recordEvent(spend, event);
        response.write(`${JSON.stringify(event)}\n`);
      }
      spend.completed = true;
      response.end();
    } catch (error) {
      // Read before the error frame is written: ending the response fires
      // 'close' as well, so a reading of the abort state taken afterwards
      // cannot tell a reader who walked away from one who got the answer.
      spend.aborted = abort.signal.aborted;
      spend.error = error instanceof Error ? error.message : String(error);
      // Server-side, with the real reason. The client gets none of it: the
      // upstream body is echoed into these messages and is not ours to relay.
      console.error('[piggy] chat turn failed:', spend.error);
      response.end(`${JSON.stringify({ type: 'error', message: 'Piggy chat failed.' })}\n`);
    } finally {
      // In a finally so that every exit closes the row, including the exit
      // that is not a fault at all: a reader who navigates away aborts the
      // turn mid-answer. A row left `running` cannot be told from a turn still
      // in flight by any later query — which is exactly the query a per-user
      // daily cap would have to make.
      await finishChatRun(db, run, spend, options.tokenPricing);
    }
  });
  server.listen(options.port, host);
  return server;
}

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

/**
 * The chat's cost ledger.
 *
 * `agent_runs` existed and only the queued worker ever wrote to it, so every
 * token the docked panel spent was invisible: nothing in the API or the web app
 * could answer "what has Piggy cost today", let alone cap it per user. A chat
 * turn is one run, with `agent_task_id` left null — the column is nullable for
 * precisely this case, a run with no queued task behind it.
 *
 * A failure to write the ledger never fails the answer. Losing the accounting
 * for one turn is a smaller harm than refusing to talk to the user because a
 * bookkeeping insert did not land.
 */
interface ChatRunOutcome {
  toolCalls: number;
  answer?: string;
  inputTokens?: number | null;
  outputTokens?: number | null;
  /** The stream ran to its end. */
  completed?: boolean;
  /** The reader hung up before it did. */
  aborted?: boolean;
  error?: string;
}

function recordEvent(outcome: ChatRunOutcome, event: PiggyChatEvent): void {
  if (event.type === 'content_delta') outcome.answer = (outcome.answer ?? '') + event.delta;
  if (event.type === 'tool_call') outcome.toolCalls += 1;
  if (event.type === 'done') {
    outcome.inputTokens = event.inputTokens;
    outcome.outputTokens = event.outputTokens;
  }
  if (event.type === 'error') outcome.error = event.message;
}

async function startChatRun(
  db: Database,
  input: {
    principalUserId: string;
    model: string;
    message: string;
    context?: z.infer<typeof contextSchema>;
    historyTurns: number;
  },
): Promise<string | null> {
  try {
    const [run] = await db
      .insert(agentRuns)
      .values({
        principalUserId: input.principalUserId,
        model: input.model,
        input: {
          surface: 'chat',
          message: input.message,
          context: input.context ?? null,
          historyTurns: input.historyTurns,
        },
      })
      .returning({ id: agentRuns.id });
    return run?.id ?? null;
  } catch (error) {
    console.error('[piggy] could not open an agent run for this chat turn:', error);
    return null;
  }
}

async function finishChatRun(
  db: Database,
  runId: string | null,
  outcome: ChatRunOutcome,
  pricing?: ChatTokenPricing,
): Promise<void> {
  if (!runId) return;
  const summary = outcome.answer?.trim();
  try {
    await db
      .update(agentRuns)
      .set({
        // An abandoned turn is not a failed one — the answer was fine, the
        // reader left — and counting it as failed would make the failure rate
        // read as an outage every time somebody closed a tab.
        status: outcome.completed ? 'succeeded' : outcome.aborted ? 'aborted' : 'failed',
        summary: summary || null,
        result: { toolCalls: outcome.toolCalls },
        inputTokens: outcome.inputTokens ?? null,
        outputTokens: outcome.outputTokens ?? null,
        costMicroCents: costMicroCents(outcome, pricing),
        error: outcome.error ?? null,
        finishedAt: new Date(),
      })
      .where(eq(agentRuns.id, runId));
  } catch (error) {
    console.error(`[piggy] could not close agent run ${runId}:`, error);
  }
}

/**
 * Tokens are billed per million, so cents-per-million multiplied by tokens is
 * already micro-cents. Doing it that way keeps the whole calculation in
 * integers rather than rounding a fraction of a cent per turn and drifting.
 */
function costMicroCents(outcome: ChatRunOutcome, pricing?: ChatTokenPricing): number | null {
  if (!pricing) return null;
  const input = outcome.inputTokens ?? null;
  const output = outcome.outputTokens ?? null;
  if (input === null && output === null) return null;
  return Math.round(
    (input ?? 0) * pricing.inputCentsPerMillionTokens +
      (output ?? 0) * pricing.outputCentsPerMillionTokens,
  );
}

function respondJson(response: ServerResponse, status: number, body: unknown): void {
  response.writeHead(status, { 'content-type': 'application/json' });
  response.end(JSON.stringify(body));
}

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';
}