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

/**
 * The cost ceiling, proved against the real harness rather than argued for.
 *
 * `@earendil-works/pi-agent-core`'s `agent-loop.js` is a `while (true)` with
 * four exits: the model stops asking for tools, it errors, the run is aborted,
 * or `shouldStopAfterTurn` returns true. Nothing in it counts iterations and
 * nothing in it counts tokens, so a model that keeps asking for one more tool
 * call keeps buying model calls until somebody stops it.
 *
 * Every test here drives that real loop — real `createAgentSession`, real tool
 * execution, real event stream — with the provider swapped for a stand-in that
 * always asks for another call. `Agent.streamFunction` is a public, mutable
 * property and is the only seam that lets an offline test spend "money": the
 * alternative is a live endpoint and a real bill, which is not a test.
 */
import assert from 'node:assert/strict';
import { mkdtempSync, rmSync } from 'node:fs';
import { tmpdir } from 'node:os';
import { join } from 'node:path';
import test, { after, before } from 'node:test';
import { defineTool, type AgentSession, type ToolDefinition } from '@earendil-works/pi-coding-agent';
import { Type } from 'typebox';
import { createTurnBudget, observeTurn, type PiggySession } from '../src/agent/session';
import type { PiggyTurnLimits } from '../src/config';

const agentDir = mkdtempSync(join(tmpdir(), 'piggy-budget-test-'));

before(() => {
  process.env.DATABASE_URL = 'postgres://pig:pig@localhost:54330/pig';
  process.env.PIGGY_INTERNAL_TOKEN = 'test-internal-token-for-piggy-000000';
  process.env.PRIME_API_KEY = 'test-key-not-used-offline';
  process.env.PIGGY_AGENT_DIR = agentDir;
});

after(() => {
  rmSync(agentDir, { recursive: true, force: true });
});

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

/** A tool that always succeeds, so the loop is never stopped by a tool failing. */
function alwaysAnswers(): ToolDefinition {
  return defineTool({
    name: 'pig_get_workspace_summary',
    label: 'Workspace summary',
    description: 'Test double: always answers.',
    promptSnippet: 'pig_get_workspace_summary: test double.',
    parameters: Type.Object({}),
    async execute() {
      return { content: [{ type: 'text' as const, text: '{"ok":true}' }], details: { ok: true } };
    },
  });
}

/** The harness's stream function, reached through the object that owns it. */
type StreamFunction = AgentSession['agent']['streamFunction'];
type StreamResult = Awaited<ReturnType<StreamFunction>>;

interface Provocation {
  /** How many times the loop asked the provider for another response. */
  calls: number;
}

/**
 * A provider that always asks for another tool call.
 *
 * This is the runaway in its purest form: every response is a well-formed
 * assistant message whose only content is a tool call, which is precisely the
 * condition `agent-loop.js` uses to decide it has more to do. `relentUntil`
 * exists only so the control test — the one that shows nothing else stops this
 * — terminates: without a cap of our own, the loop's own stopping condition
 * never arrives.
 */
function provokeAnotherCall(
  session: PiggySession,
  usagePerCall: { input: number; output: number },
  relentAfter = Number.POSITIVE_INFINITY,
): Provocation {
  const provocation: Provocation = { calls: 0 };
  const model = session.session.agent.state.model;
  const stream: StreamFunction = () => {
    provocation.calls += 1;
    const relent = provocation.calls >= relentAfter;
    const message = {
      role: 'assistant',
      content: relent
        ? [{ type: 'text', text: 'Done.' }]
        : [
            {
              type: 'toolCall',
              id: `call_${provocation.calls}`,
              name: 'pig_get_workspace_summary',
              arguments: {},
            },
          ],
      api: model.api,
      provider: model.provider,
      model: model.id,
      usage: {
        input: usagePerCall.input,
        output: usagePerCall.output,
        cacheRead: 0,
        cacheWrite: 0,
        totalTokens: usagePerCall.input + usagePerCall.output,
        cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
      },
      stopReason: relent ? 'stop' : 'toolUse',
      timestamp: Date.now(),
    };
    // An empty event sequence with a result is a shape the loop handles: it
    // falls through to `response.result()` and emits the message itself. The
    // cast is the same one the chat-server tests make — building all forty
    // fields of a streamed AssistantMessage would test the double, not the cap.
    return {
      [Symbol.asyncIterator]: () => ({ next: async () => ({ done: true as const, value: undefined }) }),
      result: async () => message,
    } as unknown as StreamResult;
  };
  session.session.agent.streamFunction = stream;
  return provocation;
}

test('nothing in the harness stops a model that keeps asking for another call', async () => {
  const { createPiggySession } = await import('../src/agent/session');
  // Deliberately no budget: this is the finding, reproduced. The loop runs as
  // many model calls as the model asks for, and the only reason this test
  // terminates is that the stand-in provider gives up after twenty.
  const piggy = await createPiggySession({ mode: 'read_only', tools: [alwaysAnswers()] });
  try {
    const provocation = provokeAnotherCall(piggy, { input: 5_000, output: 150 }, 20);
    await piggy.session.prompt('How are we doing?');

    assert.equal(provocation.calls, 20);
  } finally {
    piggy.dispose();
  }
});

test('the model-call ceiling stops the runaway at exactly its ceiling', async () => {
  const { createPiggySession } = await import('../src/agent/session');
  const budget = createTurnBudget(limits({ maxModelCalls: 3 }));
  const piggy = await createPiggySession({
    mode: 'read_only',
    tools: [alwaysAnswers()],
    budget,
  });
  try {
    // Never relents. Without the ceiling this call does not return.
    const provocation = provokeAnotherCall(piggy, { input: 5_000, output: 150 });
    await piggy.session.prompt('How are we doing?');

    assert.equal(provocation.calls, 3, 'the loop bought more calls than the ceiling allows');
    assert.equal(budget.breach?.limit, 'model_calls');
    assert.equal(budget.breach?.ceiling, 3);
    assert.equal(budget.breach?.modelCalls, 3);
    // The stop is graceful: the loop ends of its own accord rather than being
    // aborted, so the turn settles instead of spinning.
    assert.equal(budget.overran, false);
  } finally {
    piggy.dispose();
  }
});

test('the token ceiling stops a turn whose calls are few and enormous', async () => {
  const { createPiggySession } = await import('../src/agent/session');
  // A cap on calls alone is escapable: eight calls of a hundred thousand tokens
  // is a hundred times a normal turn while never reaching the call ceiling.
  const budget = createTurnBudget(limits({ maxModelCalls: 100, maxTurnTokens: 30_000 }));
  const piggy = await createPiggySession({
    mode: 'read_only',
    tools: [alwaysAnswers()],
    budget,
  });
  try {
    const provocation = provokeAnotherCall(piggy, { input: 12_000, output: 500 });
    await piggy.session.prompt('Summarise everything.');

    // 12,500 per call, so the third call is the one that passes 30,000.
    assert.equal(provocation.calls, 3);
    assert.equal(budget.breach?.limit, 'tokens');
    assert.equal(budget.breach?.tokens, 37_500);
    assert.equal(budget.breach?.ceiling, 30_000);
  } finally {
    piggy.dispose();
  }
});

test('input tokens count, because input is what a tool-heavy turn is billed for', async () => {
  const { createPiggySession } = await import('../src/agent/session');
  // Measured on the live stack: a two-tool turn on the default model is 12,099
  // input and 166 output. A ceiling that counted only output would have let
  // that turn run 70 times over before noticing.
  const budget = createTurnBudget(limits({ maxModelCalls: 100, maxTurnTokens: 12_000 }));
  const piggy = await createPiggySession({
    mode: 'read_only',
    tools: [alwaysAnswers()],
    budget,
  });
  try {
    const provocation = provokeAnotherCall(piggy, { input: 6_000, output: 20 });
    await piggy.session.prompt('Summarise everything.');

    assert.equal(provocation.calls, 2);
    assert.equal(budget.breach?.limit, 'tokens');
  } finally {
    piggy.dispose();
  }
});

test('a turn well inside both ceilings is never interfered with', async () => {
  const { createPiggySession } = await import('../src/agent/session');
  const budget = createTurnBudget(limits());
  const piggy = await createPiggySession({
    mode: 'read_only',
    tools: [alwaysAnswers()],
    budget,
  });
  try {
    // The measured shape of a real two-tool turn: three model calls, ~12,265
    // tokens. It must finish on the model's own terms.
    const provocation = provokeAnotherCall(piggy, { input: 4_000, output: 90 }, 3);
    await piggy.session.prompt('Which supplier has the lowest utilisation?');

    assert.equal(provocation.calls, 3);
    assert.equal(budget.breach, undefined);
    assert.equal(budget.modelCalls, 3);
    assert.equal(budget.tokens, 12_270);
  } finally {
    piggy.dispose();
  }
});

test('two counters of the same turn merge rather than halving the ceiling', () => {
  // The in-loop hook and the chat server both report what they have seen, and
  // they are describing the same model calls. Summing them would cut every
  // ceiling in half and stop honest turns; `observeTurn` takes the larger
  // reading instead.
  const budget = createTurnBudget(limits({ maxModelCalls: 4 }));
  observeTurn(budget, 1, 3_000);
  observeTurn(budget, 1, 3_000);
  observeTurn(budget, 2, 6_000);
  observeTurn(budget, 2, 6_000);
  assert.equal(budget.modelCalls, 2);
  assert.equal(budget.tokens, 6_000);
  assert.equal(budget.breach, undefined);
});

test('a model call after the ceiling is recorded as an overrun, not ignored', () => {
  // What it looks like when the in-loop stop does not hold — a harness upgrade
  // that claims `shouldStopAfterTurn` for itself, say. The operator has to be
  // able to see that the graceful brake failed and the hard one was needed.
  const budget = createTurnBudget(limits({ maxModelCalls: 2 }));
  observeTurn(budget, 1, 1_000);
  observeTurn(budget, 2, 2_000);
  assert.equal(budget.breach?.limit, 'model_calls');
  assert.equal(budget.overran, false);
  observeTurn(budget, 3, 3_000);
  assert.equal(budget.overran, true);
  // The breach itself is never rewritten: it records where the line was crossed.
  assert.equal(budget.breach?.modelCalls, 2);
});