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

import assert from 'node:assert/strict';
import type { AddressInfo } from 'node:net';
import test from 'node:test';
import { z } from 'zod';
import type { Database } from '@pig/db';
import type { PiggyChatEvent, PiggyChatRequest } from '../src/chat';
import { startPiggyChatServer, type PiggyChatServerOptions } from '../src/chat-server';

const TOKEN = 'test-internal-token-for-piggy-000000';

/**
 * The chat server writes exactly two statements per turn — one insert, one
 * update — so a fake that records them is enough to assert the whole ledger.
 * The tools are built against this handle too, but tool construction never
 * touches it and the provider here is a fake, so nothing else is reached.
 */
interface RecordedRun {
  values: Record<string, unknown>;
  closed?: Record<string, unknown>;
}

function fakeDatabase(runs: RecordedRun[]): Database {
  return {
    insert: () => ({
      values: (values: Record<string, unknown>) => ({
        returning: async () => {
          runs.push({ values });
          return [{ id: `run-${runs.length}` }];
        },
      }),
    }),
    update: () => ({
      set: (closed: Record<string, unknown>) => ({
        where: async () => {
          const run = runs.at(-1);
          if (run) run.closed = closed;
        },
      }),
    }),
  } as unknown as Database;
}

function providerYielding(events: PiggyChatEvent[], thrown?: Error): PiggyChatServerOptions['provider'] {
  return {
    model: 'nvidia/nemotron-3-nano-30b-a3b',
    run: async function* (_request: PiggyChatRequest) {
      for (const event of events) yield event;
      if (thrown) throw thrown;
    },
  };
}

async function startForTest(
  t: { after: (fn: () => void) => void },
  provider: PiggyChatServerOptions['provider'],
  runs: RecordedRun[],
): Promise<string> {
  const server = startPiggyChatServer(fakeDatabase(runs), {
    port: 0,
    internalToken: TOKEN,
    provider,
    tokenPricing: { inputCentsPerMillionTokens: 5, outputCentsPerMillionTokens: 20 },
  });
  t.after(() => server.close());
  // Port 0 is only resolved once the socket is bound.
  await new Promise((resolve) => server.once('listening', resolve));
  const { port } = server.address() as AddressInfo;
  return `http://127.0.0.1:${port}`;
}

function chatBody(message = 'What is idle costing us?') {
  return JSON.stringify({
    principalUserId: '20000000-0000-4000-8000-000000000001',
    message,
    context: { type: 'page', route: '/capacity' },
  });
}

const authorised = { authorization: `Bearer ${TOKEN}`, 'content-type': 'application/json' };

test('health answers without a token, and nothing else does', async (t) => {
  const runs: RecordedRun[] = [];
  const base = await startForTest(t, providerYielding([]), runs);

  const health = await fetch(`${base}/internal/health`);
  assert.equal(health.status, 200);
  assert.deepEqual(await health.json(), {
    ok: true,
    service: 'piggy-chat',
    model: 'nvidia/nemotron-3-nano-30b-a3b',
  });

  assert.equal((await fetch(`${base}/internal/anything`)).status, 404);
  assert.equal(
    (await fetch(`${base}/internal/chat`, { method: 'POST', body: chatBody() })).status,
    401,
  );
});

test('a chat turn is recorded in agent_runs with its tokens and cost', async (t) => {
  const runs: RecordedRun[] = [];
  const base = await startForTest(
    t,
    providerYielding([
      { type: 'meta', model: 'nvidia/nemotron-3-nano-30b-a3b' },
      { type: 'tool_call', id: 'call_1', name: 'pig_get_idle_capacity', arguments: {} },
      { type: 'tool_result', id: 'call_1', name: 'pig_get_idle_capacity', ok: true, result: {} },
      { type: 'content_delta', delta: 'Idle is $12,000.' },
      { type: 'done', inputTokens: 1_240, outputTokens: 180 },
    ]),
    runs,
  );

  const response = await fetch(`${base}/internal/chat`, {
    method: 'POST',
    headers: authorised,
    body: chatBody(),
  });
  assert.equal(response.status, 200);
  const frames = (await response.text()).trim().split('\n').map((line) => JSON.parse(line));
  assert.equal(frames.length, 5);

  const run = runs[0];
  assert.equal(run?.values.model, 'nvidia/nemotron-3-nano-30b-a3b');
  assert.equal(run?.values.principalUserId, '20000000-0000-4000-8000-000000000001');
  assert.equal(run?.closed?.status, 'succeeded');
  assert.equal(run?.closed?.summary, 'Idle is $12,000.');
  assert.equal(run?.closed?.inputTokens, 1_240);
  assert.equal(run?.closed?.outputTokens, 180);
  // 1240 x 5 + 180 x 20 micro-cents, at $0.05/$0.20 per million tokens.
  assert.equal(run?.closed?.costMicroCents, 9_800);
  assert.ok(run?.closed?.finishedAt instanceof Date);
});

test('a malformed request is the only thing called an invalid request', async (t) => {
  const runs: RecordedRun[] = [];
  const base = await startForTest(t, providerYielding([]), runs);

  const response = await fetch(`${base}/internal/chat`, {
    method: 'POST',
    headers: authorised,
    body: JSON.stringify({ principalUserId: 'not-a-uuid', message: '' }),
  });

  assert.equal(response.status, 400);
  assert.deepEqual(await response.json(), { error: 'Invalid Piggy chat request.' });
  // No inference was attempted, so no run should have been opened for it.
  assert.equal(runs.length, 0);
});

test('a fault raised mid-stream is not blamed on the user, and closes its run', async (t) => {
  const runs: RecordedRun[] = [];
  // A ZodError, because that is the one the old code mistook for bad input:
  // a schema failure inside the turn reported "Invalid Piggy chat request" to
  // someone whose request was perfectly valid.
  const upstreamFault = new z.ZodError([]);
  const base = await startForTest(
    t,
    providerYielding(
      [
        { type: 'meta', model: 'nvidia/nemotron-3-nano-30b-a3b' },
        { type: 'content_delta', delta: 'Idle is ' },
      ],
      upstreamFault,
    ),
    runs,
  );

  const response = await fetch(`${base}/internal/chat`, {
    method: 'POST',
    headers: authorised,
    body: chatBody(),
  });

  // The stream had already begun, so the turn ends as an error frame on a 200.
  assert.equal(response.status, 200);
  const frames = (await response.text()).trim().split('\n').map((line) => JSON.parse(line));
  assert.deepEqual(frames.at(-1), { type: 'error', message: 'Piggy chat failed.' });
  assert.equal(runs[0]?.closed?.status, 'failed');
  assert.equal(runs[0]?.closed?.summary, 'Idle is');
});

test('a reader who leaves mid-answer closes the run as abandoned, not as running', async (t) => {
  const runs: RecordedRun[] = [];
  const base = await startForTest(
    t,
    {
      model: 'nvidia/nemotron-3-nano-30b-a3b',
      // A real provider notices the abort at its next await; this one at its
      // next yield, which is the same thing at this scale.
      run: async function* (request: PiggyChatRequest) {
        for (let index = 0; index < 20; index += 1) {
          if (request.signal?.aborted) throw request.signal.reason;
          await new Promise((resolve) => setTimeout(resolve, 20));
          yield { type: 'content_delta', delta: `chunk ${index} ` } as PiggyChatEvent;
        }
      },
    },
    runs,
  );

  const abort = new AbortController();
  setTimeout(() => abort.abort(), 80);
  await assert.rejects(
    fetch(`${base}/internal/chat`, {
      method: 'POST',
      headers: authorised,
      body: chatBody(),
      signal: abort.signal,
    }).then((response) => response.text()),
  );

  await waitFor(() => runs[0]?.closed !== undefined);
  // Without the finally this row stayed `running` for ever, and no later query
  // could tell it from a turn still in flight.
  assert.equal(runs[0]?.closed?.status, 'aborted');
});

async function waitFor(condition: () => boolean): Promise<void> {
  for (let attempt = 0; attempt < 100; attempt += 1) {
    if (condition()) return;
    await new Promise((resolve) => setTimeout(resolve, 10));
  }
  assert.fail('the run was never closed');
}