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; } /** * 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; 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[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; historyTurns: number; }, ): Promise { 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 { 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 { 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'; }