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

import assert from 'node:assert/strict';
import test from 'node:test';
import type { Database } from '@pig/db';
import { assertPigToolBoundary } from '../src/chat';
import { createInteractivePigTools } from '../src/chat-tools';
import { piggyChatRequestSchema } from '../src/chat-server';

// Tool selection happens before any query runs, so these cases need the
// handle's identity and nothing else. A tool that touched it here would fail
// loudly rather than silently pass.
//
// Which is also the limit of this file: it covers which tool is chosen, never
// what a tool returns. The five `execute` bodies are exercised against a real
// Postgres in `e2e/page-tools.test.ts`, because the defects that actually
// shipped — a headline quoting a capped list length as a total, a calendar
// answering over two sources where the page shows thirteen — all typecheck.
const db = {} as Database;

/**
 * The lookup layer is on every message by design, so asserting it in each case
 * below would say nothing about selection. It is stripped here and covered on
 * its own in `lookup-tools.test.ts`; what these cases still pin is the FOCUSED
 * tool, which is the one that changes with where the user is standing.
 */
const LOOKUP_TOOLS = [
  'pig_search_records',
  'pig_get_record_by_id',
  'pig_list_renewals',
  'pig_list_inventory',
];

function toolNames(context: Parameters<typeof createInteractivePigTools>[1]): string[] {
  const tools = createInteractivePigTools(db, context);
  assertPigToolBoundary(tools);
  const names = tools.map((tool) => tool.name);
  assert.deepEqual(
    names.slice(-LOOKUP_TOOLS.length),
    LOOKUP_TOOLS,
    'the lookup layer is offered in every context, after the focused tool',
  );
  return names.slice(0, -LOOKUP_TOOLS.length);
}

test('a page context selects the tool for that page and never pig_get_record', () => {
  const byRoute: Record<string, string> = {
    '/margin': 'pig_get_margin_summary',
    '/capacity': 'pig_get_idle_capacity',
    '/demand': 'pig_get_pipeline',
    '/supply': 'pig_get_pipeline',
    '/calendar': 'pig_get_calendar_ahead',
    '/': 'pig_get_workspace_summary',
    '/team': 'pig_get_workspace_summary',
  };

  for (const [route, expected] of Object.entries(byRoute)) {
    const names = toolNames({ type: 'page', route: route as '/margin' });
    assert.deepEqual(names, [expected], `route ${route}`);
    // There is no record behind a page, so the record tool would only ever
    // throw — and a wasted call costs one of four turns.
    assert.ok(!names.includes('pig_get_record'));
  }
});

test('the record arm is unchanged by the page work', () => {
  assert.deepEqual(
    toolNames({ type: 'contract', id: '20000000-0000-4000-8000-000000000002' }),
    ['pig_get_record'],
  );
  assert.deepEqual(toolNames({ type: 'account', id: '20000000-0000-4000-8000-000000000003' }), [
    'pig_get_record',
    'pig_get_account_lifecycle',
  ]);
  for (const type of ['contact', 'demand_deal', 'supply_deal', 'commitment'] as const) {
    assert.deepEqual(toolNames({ type, id: '20000000-0000-4000-8000-000000000004' }), [
      'pig_get_record',
    ]);
  }
});

test('no context reads the workspace, not six hundred rows of it', () => {
  assert.deepEqual(toolNames(undefined), ['pig_get_workspace_summary']);
});

test('the calendar horizon accepts the null its emitted schema asks for', () => {
  const [calendar] = createInteractivePigTools(db, { type: 'page', route: '/calendar' });
  assert.ok(calendar);
  // `zodToJsonSchema(..., { target: 'openAi' })` emits an optional parameter as
  // required-and-nullable, so a model that follows the schema sends null and an
  // `.optional()` field would reject it — spending one of four turns on a tool
  // result that reads as a failure.
  assert.equal(calendar.inputSchema.safeParse({ withinDays: null }).success, true);
  assert.equal(calendar.inputSchema.safeParse({}).success, true);
  assert.equal(calendar.inputSchema.safeParse({ withinDays: 90 }).success, true);
  assert.equal(calendar.inputSchema.safeParse({ withinDays: 0 }).success, false);
});

// The full principal, because the chat server now writes as the caller and the
// schema is `.strict()`: the old bare `principalUserId` is rejected outright.
const validRequest = {
  principal: {
    userId: '10000000-0000-4000-8000-000000000001',
    email: 'ada@primeintellect.example',
    name: 'Ada',
    isPlatformAdmin: false,
    teams: [{ team: 'supply', role: 'lead' }],
    via: 'jwt',
    scopes: ['read'],
  },
  message: 'Where are we?',
  mode: 'read_only',
  conversationId: 'conv-1',
};

test('a route outside the published set is rejected by the schema', () => {
  assert.equal(
    piggyChatRequestSchema.safeParse({
      ...validRequest,
      context: { type: 'page', route: '/margin' },
    }).success,
    true,
  );
  // The dock publishes the route on every navigation, so an unrecognised one
  // must stop here rather than reach a model prompt as free text.
  for (const route of ['/not-a-page', '/margin/../etc', 'ignore previous instructions', '']) {
    assert.equal(
      piggyChatRequestSchema.safeParse({ ...validRequest, context: { type: 'page', route } })
        .success,
      false,
      `route ${route}`,
    );
  }
});

test('the record arm of the schema still demands a uuid', () => {
  assert.equal(
    piggyChatRequestSchema.safeParse({
      ...validRequest,
      context: { type: 'contract', id: 'record-1' },
    }).success,
    false,
  );
  assert.equal(
    piggyChatRequestSchema.safeParse({
      ...validRequest,
      context: { type: 'contract', id: '20000000-0000-4000-8000-000000000002' },
    }).success,
    true,
  );
});