pig/apps/piggy/test/turn-limits.test.ts

/**
 * What the chat server does about a turn that costs too much.
 *
 * `turn-budget.test.ts` proves the in-loop brake against the real harness. This
 * proves the other half: that the server has a brake of its own for a harness
 * that ignores it, that the user is told what happened rather than handed a
 * truncated answer dressed as a finished one, that the run row says the turn
 * was stopped rather than that it failed — and that none of it fires on a turn
 * that is merely slow because a human is thinking about an approval.
 *
 * The sessions here are deliberately hook-free doubles: they never call
 * `shouldStopAfterTurn`, which is exactly the condition the server's counter
 * exists for.
 */
import assert from 'node:assert/strict';
import type { AddressInfo } from 'node:net';
import test from 'node:test';
import type { AgentSession, AgentSessionEvent, ToolDefinition } from '@earendil-works/pi-coding-agent';
import type { PiggyChatEvent, PiggyModelOption } from '@pig/core';
import type { Database } from '@pig/db';
import type { PiggySession } from '../src/agent/session';
import { startPiggyChatServer, type PiggyChatServerOptions } from '../src/chat-server';
import type { PiggyTurnLimits } from '../src/config';
import type { PigWriteToolDeps } from '../src/write-tools';

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

const MODELS: PiggyModelOption[] = [
  {
    id: 'nvidia/nemotron-3-nano-30b-a3b',
    label: 'Nemotron 3 Nano',
    costPerMTokIn: 0.05,
    costPerMTokOut: 0.2,
    contextWindow: 131_072,
    reasoning: true,
    isDefault: true,
  },
];

interface RecordedRun {
  values: Record<string, unknown>;
  closed?: Record<string, unknown>;
}

/**
 * The two statements the chat server writes, plus the one it reads: the daily
 * spend. `spentMicroCents` is what the sum comes back as — a string, because
 * that is how the driver hands over a numeric so a bigint cannot be rounded.
 */
function fakeDatabase(runs: RecordedRun[], spentMicroCents = '0'): 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;
        },
      }),
    }),
    select: () => ({
      from: () => ({
        where: async () => [{ spent: spentMicroCents }],
      }),
    }),
  } as unknown as Database;
}

type TurnScript = (
  tools: readonly ToolDefinition[],
  emit: (event: AgentSessionEvent) => void,
  signal: AbortSignal,
) => Promise<void>;

interface SessionSpy {
  created: number;
  aborted: number;
}

/**
 * A session double with no `shouldStopAfterTurn` at all.
 *
 * `abort()` is the only thing that can stop its script, which is the point: it
 * stands in for a harness whose in-loop hooks we do not control, and it is how
 * the server's own brake gets tested rather than the harness's.
 */
function hookFreeSessions(script: TurnScript, watched: SessionSpy) {
  return async (options: { tools: readonly ToolDefinition[]; modelId?: string }): Promise<PiggySession> => {
    watched.created += 1;
    const listeners = new Set<(event: AgentSessionEvent) => void>();
    const aborted = new AbortController();
    const session = {
      subscribe(listener: (event: AgentSessionEvent) => void) {
        listeners.add(listener);
        return () => listeners.delete(listener);
      },
      async prompt() {
        await script(
          options.tools,
          (event) => {
            for (const listener of [...listeners]) listener(event);
          },
          aborted.signal,
        );
      },
      async abort() {
        watched.aborted += 1;
        aborted.abort();
      },
      dispose() {},
    } as unknown as AgentSession;

    return {
      session,
      modelId: options.modelId ?? MODELS[0]!.id,
      systemPrompt: 'You are Piggy.',
      dispose: () => aborted.abort(),
    } satisfies PiggySession;
  };
}

function turnEnd(input: number, output: number, stopReason = 'toolUse'): AgentSessionEvent {
  return {
    type: 'turn_end',
    message: { role: 'assistant', usage: { input, output }, stopReason },
    toolResults: [],
  } as unknown as AgentSessionEvent;
}

function toolStart(id: string, name: string): AgentSessionEvent {
  return { type: 'tool_execution_start', toolCallId: id, toolName: name, args: {} } as unknown as AgentSessionEvent;
}

function limits(overrides: Partial<PiggyTurnLimits> = {}): PiggyTurnLimits {
  return { maxModelCalls: 8, maxTurnTokens: 40_000, dailyLimitCents: 0, ...overrides };
}

async function startForTest(
  t: { after: (fn: () => void) => void },
  db: Database,
  options: Partial<PiggyChatServerOptions>,
): Promise<string> {
  const server = startPiggyChatServer(db, {
    port: 0,
    internalToken: TOKEN,
    models: MODELS,
    createReadTools: () => [],
    createWriteTools: () => [],
    limits: limits(),
    ...options,
  });
  t.after(() => server.close());
  await new Promise((resolve) => server.once('listening', resolve));
  const { port } = server.address() as AddressInfo;
  return `http://127.0.0.1:${port}`;
}

const PRINCIPAL = {
  userId: '20000000-0000-4000-8000-000000000001',
  email: 'ada@primeintellect.example',
  name: 'Ada',
  isPlatformAdmin: false,
  teams: [{ team: 'supply', role: 'lead' }],
  via: 'jwt',
  scopes: ['read', 'write'],
};

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

function chatBody(overrides: Record<string, unknown> = {}): string {
  return JSON.stringify({
    principal: PRINCIPAL,
    message: 'What is idle costing us?',
    mode: 'read_only',
    conversationId: 'conv-limit',
    ...overrides,
  });
}

function parseFrames(body: string): PiggyChatEvent[] {
  return body
    .trim()
    .split('\n')
    .filter((line) => line.length > 0)
    .map((line) => JSON.parse(line) as PiggyChatEvent);
}

/** The runaway: a turn that asks for another tool call for ever. */
function relentless(counted: { calls: number }, usage = { input: 4_000, output: 100 }): TurnScript {
  return async (_tools, emit, signal) => {
    while (!signal.aborted) {
      counted.calls += 1;
      emit(toolStart(`call_${counted.calls}`, 'pig_get_workspace_summary'));
      emit(turnEnd(usage.input, usage.output));
      // Yield, so an abort raised inside the event handling above is observed
      // rather than starved by a tight synchronous loop.
      await new Promise((resolve) => setImmediate(resolve));
    }
  };
}

test('a harness that ignores the in-loop stop is aborted by the server', async (t) => {
  const runs: RecordedRun[] = [];
  const counted = { calls: 0 };
  const watched: SessionSpy = { created: 0, aborted: 0 };
  const base = await startForTest(t, fakeDatabase(runs), {
    limits: limits({ maxModelCalls: 4 }),
    createSession: hookFreeSessions(relentless(counted), watched),
  });

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

  // The double would have run for ever. Something stopped it, and it was not
  // the double.
  assert.equal(watched.aborted, 1);
  assert.ok(counted.calls >= 4, 'the ceiling was not reached at all');
  assert.ok(counted.calls <= 6, `the abort did not take hold: ${counted.calls} model calls`);

  // The user is told, in their own terms, and the transcript settles on an
  // error rather than on a `done` that would present a truncated answer as
  // the whole of it.
  const last = frames.at(-1);
  assert.equal(last?.type, 'error');
  assert.equal(last?.type === 'error' ? last.code : null, 'turn_limit_exceeded');
  assert.match(last?.type === 'error' ? last.message : '', /incomplete/);
  assert.equal(
    frames.some((frame) => frame.type === 'done'),
    false,
    'a cut-off turn must not also report itself finished',
  );

  // And the operator can tell "stopped for cost" from "failed".
  const closed = runs[0]?.closed;
  assert.equal(closed?.status, 'aborted');
  assert.match(String(closed?.error), /model_calls ceiling/);
  const result = closed?.result as { limit?: Record<string, unknown>; modelCalls?: number };
  assert.equal(result?.limit?.reason, 'model_calls');
  assert.equal(result?.limit?.ceiling, 4);
  assert.equal(typeof result?.modelCalls, 'number');
});

test('the token ceiling stops a turn whose model calls are few and enormous', async (t) => {
  const runs: RecordedRun[] = [];
  const counted = { calls: 0 };
  const watched: SessionSpy = { created: 0, aborted: 0 };
  const base = await startForTest(t, fakeDatabase(runs), {
    // Far more calls than the tokens allow, so only the token ceiling can bite.
    limits: limits({ maxModelCalls: 500, maxTurnTokens: 25_000 }),
    createSession: hookFreeSessions(
      relentless(counted, { input: 12_000, output: 500 }),
      watched,
    ),
  });

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

  assert.equal(watched.aborted, 1);
  assert.ok(counted.calls <= 4, `${counted.calls} model calls before the tokens ran out`);
  const last = frames.at(-1);
  assert.equal(last?.type === 'error' ? last.code : null, 'turn_limit_exceeded');
  assert.match(last?.type === 'error' ? last.message : '', /size limit/);

  const closed = runs[0]?.closed;
  assert.equal(closed?.status, 'aborted');
  assert.match(String(closed?.error), /tokens ceiling/);
  const result = closed?.result as { limit?: Record<string, unknown> };
  assert.equal(result?.limit?.reason, 'tokens');
  assert.equal(result?.limit?.ceiling, 25_000);
  // The tokens generated before the stop are still billed to the ledger: they
  // were spent whether or not the answer arrived.
  assert.ok(Number(closed?.inputTokens) > 0);
  assert.ok(Number(closed?.costMicroCents) > 0);
});

test('a turn that finishes on the very call that reaches the ceiling still reports done', async (t) => {
  const runs: RecordedRun[] = [];
  const base = await startForTest(t, fakeDatabase(runs), {
    limits: limits({ maxModelCalls: 2 }),
    createSession: hookFreeSessions(async (_tools, emit) => {
      emit(toolStart('call_1', 'pig_get_workspace_summary'));
      emit(turnEnd(4_000, 100));
      // The second call is the ceiling AND the answer. Nothing was taken away
      // from the reader, so telling them their answer is incomplete would be a
      // lie in the other direction.
      emit(turnEnd(4_200, 140, 'stop'));
    }, { created: 0, aborted: 0 }),
  });

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

  assert.equal(frames.at(-1)?.type, 'done');
  const closed = runs[0]?.closed;
  assert.equal(closed?.status, 'succeeded');
  // The reading is still kept, because it is what an operator tuning the
  // ceiling needs to see.
  const result = closed?.result as { limit?: Record<string, unknown>; modelCalls?: number };
  assert.equal(result?.modelCalls, 2);
  assert.equal(result?.limit?.reason, 'model_calls');
});

/** A write tool that parks on a human, the way `confirm` mode really does. */
function proposingWriteTools(): (deps: PigWriteToolDeps) => ToolDefinition[] {
  return ({ propose }) => [
    {
      name: 'pig_log_activity',
      async execute() {
        const decision = await propose({
          tool: 'pig_log_activity',
          kind: 'activity',
          summary: 'Log a call on Northwind Robotics',
          fields: [{ label: 'Subject', value: 'Capacity review' }],
        });
        return {
          content: [{ type: 'text', text: `The change was ${decision}.` }],
          details: { tool: 'pig_log_activity', status: decision },
        };
      },
    } as unknown as ToolDefinition,
  ];
}

test('a write waiting on a human is not model work, and is not cut off for cost', async (t) => {
  const runs: RecordedRun[] = [];
  const started = Date.now();
  // Two model calls allowed and two made, with a human sitting in the middle of
  // them. A ceiling that measured wall-clock, or that counted the parked tool
  // as work, would kill precisely the turn that matters most — the one about to
  // change the CRM.
  const base = await startForTest(t, fakeDatabase(runs), {
    limits: limits({ maxModelCalls: 2, maxTurnTokens: 12_000 }),
    createWriteTools: proposingWriteTools(),
    createSession: hookFreeSessions(async (tools, emit, signal) => {
      const tool = tools.find((candidate) => candidate.name === 'pig_log_activity');
      assert.ok(tool, 'the write tool should have been handed over');
      emit(turnEnd(4_000, 120));
      emit(toolStart('call_1', 'pig_log_activity'));
      await tool.execute('call_1', {}, signal, undefined, undefined as never);
      emit(turnEnd(4_500, 160, 'stop'));
    }, { created: 0, aborted: 0 }),
  });

  const response = await fetch(`${base}/internal/chat`, {
    method: 'POST',
    headers: authorised,
    body: chatBody({ mode: 'confirm', message: 'Log a call on Northwind.' }),
  });

  // Read up to the approval card, answer it after a deliberate pause, then read
  // the rest.
  const body = response.body;
  assert.ok(body);
  const reader = body.getReader();
  const decoder = new TextDecoder();
  let buffered = '';
  const frames: PiggyChatEvent[] = [];
  const drain = (chunk: Uint8Array | undefined): void => {
    buffered += decoder.decode(chunk, { stream: true });
    const lines = buffered.split('\n');
    buffered = lines.pop() ?? '';
    for (const line of lines) if (line) frames.push(JSON.parse(line) as PiggyChatEvent);
  };
  while (!frames.some((frame) => frame.type === 'approval_required')) {
    const { done, value } = await reader.read();
    if (done) break;
    drain(value);
  }
  const asked = frames.find((frame) => frame.type === 'approval_required');
  assert.ok(asked && asked.type === 'approval_required');

  await new Promise((resolve) => setTimeout(resolve, 150));
  const decision = await fetch(`${base}/internal/approve`, {
    method: 'POST',
    headers: authorised,
    body: JSON.stringify({
      conversationId: 'conv-limit',
      changeId: asked.change.id,
      decision: 'apply',
    }),
  });
  assert.equal(decision.status, 202);

  while (true) {
    const { done, value } = await reader.read();
    if (done) break;
    drain(value);
  }

  assert.ok(Date.now() - started >= 150, 'the turn did not actually wait on the human');
  assert.equal(frames.at(-1)?.type, 'done');
  assert.equal(
    frames.some((frame) => frame.type === 'error'),
    false,
    'the pending approval was charged against a ceiling',
  );
  assert.equal(runs[0]?.closed?.status, 'succeeded');
});

test("a user who has spent the day's ceiling is refused before anything is opened", async (t) => {
  const runs: RecordedRun[] = [];
  const watched: SessionSpy = { created: 0, aborted: 0 };
  // 250 cents spent against a 200 cent ceiling.
  const base = await startForTest(t, fakeDatabase(runs, '250000000'), {
    limits: limits({ dailyLimitCents: 200 }),
    createSession: hookFreeSessions(async () => {
      assert.fail('a refused turn must not open a session');
    }, watched),
  });

  const response = await fetch(`${base}/internal/chat`, {
    method: 'POST',
    headers: authorised,
    body: chatBody(),
  });
  assert.equal(response.status, 200, 'the relay turns a non-200 into an unreadable 502');
  const frames = parseFrames(await response.text());

  assert.equal(frames[0]?.type, 'meta');
  const last = frames.at(-1);
  assert.equal(last?.type === 'error' ? last.code : null, 'daily_spend_exceeded');
  assert.match(last?.type === 'error' ? last.message : '', /\$2\.50/);
  assert.equal(watched.created, 0);
  // Nothing was spent, so nothing is written to the ledger.
  assert.equal(runs.length, 0);
});

test('a user inside the daily ceiling is answered as usual', async (t) => {
  const runs: RecordedRun[] = [];
  const base = await startForTest(t, fakeDatabase(runs, '150000000'), {
    limits: limits({ dailyLimitCents: 200 }),
    createSession: hookFreeSessions(async (_tools, emit) => {
      emit(turnEnd(4_000, 120, 'stop'));
    }, { created: 0, aborted: 0 }),
  });

  const response = await fetch(`${base}/internal/chat`, {
    method: 'POST',
    headers: authorised,
    body: chatBody(),
  });
  const frames = parseFrames(await response.text());
  assert.equal(frames.at(-1)?.type, 'done');
  assert.equal(runs[0]?.closed?.status, 'succeeded');
});

test('a daily ceiling that cannot be read allows the turn rather than denying everyone', async (t) => {
  const runs: RecordedRun[] = [];
  const broken = {
    ...fakeDatabase(runs),
    select: () => {
      throw new Error('relation "agent_runs" does not exist');
    },
  } as unknown as Database;
  const base = await startForTest(t, broken, {
    limits: limits({ dailyLimitCents: 200 }),
    createSession: hookFreeSessions(async (_tools, emit) => {
      emit(turnEnd(4_000, 120, 'stop'));
    }, { created: 0, aborted: 0 }),
  });

  const response = await fetch(`${base}/internal/chat`, {
    method: 'POST',
    headers: authorised,
    body: chatBody(),
  });
  const frames = parseFrames(await response.text());
  // A bookkeeping sum that will not come back is not a reason to stop talking
  // to anybody: the per-turn ceilings still hold, and if the database is really
  // gone the turn fails on its own merits a moment later.
  assert.equal(frames.at(-1)?.type, 'done');
});