pig/apps/piggy/test/agent-session.test.ts

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 ToolDefinition } from '@earendil-works/pi-coding-agent';
import { Type } from 'typebox';

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

before(() => {
  // The runtime reads its configuration from the environment, so the test has
  // to supply one. The key is deliberately fake: nothing below reaches the
  // endpoint, and a test that needs a live key is a test that fails in CI.
  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 fakePigTool(name: string): ToolDefinition {
  return defineTool({
    name,
    label: name,
    description: `Test double for ${name}.`,
    promptSnippet: `${name}: test double.`,
    parameters: Type.Object({}),
    async execute() {
      return { content: [{ type: 'text' as const, text: '{}' }], details: {} };
    },
  });
}

test('the session exposes exactly the tools it was handed, and nothing else', async () => {
  const { createPiggySession } = await import('../src/agent/session');
  const tools = [fakePigTool('pig_get_workspace_summary'), fakePigTool('pig_log_activity')];
  const piggy = await createPiggySession({ mode: 'confirm', tools });

  try {
    const live = piggy.session.agent.state.tools.map((tool) => tool.name).sort();

    // This is the security property of the whole harness swap, pinned rather
    // than assumed. `noTools: 'all'` plus an explicit allowlist should make it
    // impossible for a built-in to survive; if a future SDK changes the
    // precedence between its tool sources, this is what notices.
    assert.deepEqual(live, ['pig_get_workspace_summary', 'pig_log_activity']);
    for (const forbidden of ['bash', 'ipython', 'python', 'read', 'write', 'edit', 'ls', 'grep', 'find']) {
      assert.equal(live.includes(forbidden), false, `${forbidden} leaked into the tool set`);
    }
  } finally {
    piggy.dispose();
  }
});

test('a tool outside the PIG boundary never reaches the harness', async () => {
  const { createPiggySession } = await import('../src/agent/session');

  await assert.rejects(
    () => createPiggySession({ mode: 'auto', tools: [fakePigTool('bash')] }),
    /outside the PIG tool boundary/,
  );
  await assert.rejects(
    () => createPiggySession({ mode: 'auto', tools: [fakePigTool('pig_run_shell')] }),
    /outside the PIG tool boundary/,
  );
  await assert.rejects(
    () => createPiggySession({ mode: 'auto', tools: [fakePigTool('summarise')] }),
    /outside the PIG tool boundary/,
  );
});

test('a tool that reads like a shell is refused however it is spelt', async () => {
  const { createPiggySession } = await import('../src/agent/session');

  // The prefix is a convention and a convention alone is not a boundary: the
  // interesting attack is not a tool called `bash`, it is a tool called
  // `pig_bash` added by somebody who read the rule as "start it with pig_".
  for (const name of [
    'pig_bash',
    'pig_bash_run',
    'pig_BASH',
    'pig_shell_exec',
    'pig_filesystem_list',
    'pig_file_read',
    'pig_file_write',
    // Not `pig_` at all, which is the ordinary case: an agent tool from
    // somewhere else in the repo wired in by mistake.
    'PIG_get_margin_summary',
    'get_margin_summary',
  ]) {
    await assert.rejects(
      () => createPiggySession({ mode: 'auto', tools: [fakePigTool(name)] }),
      /outside the PIG tool boundary/,
      `${name} was allowed through`,
    );
  }
});

test('two tools of the same name are refused rather than silently shadowed', async () => {
  const { createPiggySession } = await import('../src/agent/session');

  await assert.rejects(
    () =>
      createPiggySession({
        mode: 'confirm',
        tools: [fakePigTool('pig_log_activity'), fakePigTool('pig_log_activity')],
      }),
    /two tools named 'pig_log_activity'/,
  );

  // The realistic version: the same name arriving from the read set and the
  // write set, with different descriptions and different bodies. Registered
  // together, one silently shadows the other inside the harness — which is how
  // a read tool ends up answering for a write tool of the same name — so the
  // check is on the name alone and cannot be talked out of it by a tool that
  // looks different in every other respect.
  const readShaped = fakePigTool('pig_log_activity');
  const writeShaped: ToolDefinition = {
    ...fakePigTool('pig_log_activity'),
    description: 'A different tool that happens to share a name.',
  };
  await assert.rejects(
    () => createPiggySession({ mode: 'confirm', tools: [readShaped, writeShaped] }),
    /two tools named 'pig_log_activity'/,
  );
});

test('a tool added after the session exists never becomes callable', async () => {
  const { createPiggySession } = await import('../src/agent/session');
  // Deliberately mutable, and deliberately the same array the caller keeps.
  const tools: ToolDefinition[] = [fakePigTool('pig_get_workspace_summary')];
  const piggy = await createPiggySession({ mode: 'confirm', tools });

  try {
    // The allowlist is decided once, at construction: `createPiggySession`
    // copies the array into `customTools` and names it in `tools`. A caller who
    // keeps a reference and pushes onto it later — a tool assembled per turn, a
    // list built up as pages are visited — must not be able to widen a session
    // that has already been checked.
    tools.push(fakePigTool('pig_delete_everything'));
    tools.push(fakePigTool('bash'));

    const live = piggy.session.agent.state.tools.map((tool) => tool.name);
    assert.deepEqual(live, ['pig_get_workspace_summary']);
  } finally {
    piggy.dispose();
  }
});

test('the system prompt is Piggy, not the harness coding assistant', async () => {
  const { createPiggySession } = await import('../src/agent/session');
  const piggy = await createPiggySession({
    mode: 'confirm',
    tools: [fakePigTool('pig_get_workspace_summary')],
  });

  try {
    // Without `await loader.reload()` the harness serves its stock preamble —
    // "an expert coding assistant operating inside pi" — with no warning of any
    // kind. The absence of that phrase is the only externally visible sign the
    // reload happened.
    assert.match(piggy.systemPrompt, /^You are Piggy/);
    assert.equal(/coding assistant/i.test(piggy.session.systemPrompt), false);
    assert.match(piggy.session.systemPrompt, /You are Piggy/);
    // The tool has to appear in the live prompt, or a 30B model never calls
    // it. The harness will not do this for us: `buildSystemPrompt` emits its
    // own "Available tools" section only when no customPrompt is supplied, and
    // replacing the coding preamble is not optional here — so the snippet is
    // rendered by prompt.ts or it is dropped in silence.
    assert.match(piggy.session.systemPrompt, /- pig_get_workspace_summary: test double\./);
  } finally {
    piggy.dispose();
  }
});

test('the mode is in the prompt, because the tool list alone does not say it', async () => {
  const { createPiggySession } = await import('../src/agent/session');
  const tools = [fakePigTool('pig_log_activity')];

  const confirm = await createPiggySession({ mode: 'confirm', tools });
  const auto = await createPiggySession({ mode: 'auto', tools });
  const readOnly = await createPiggySession({ mode: 'read_only', tools });

  try {
    assert.match(confirm.systemPrompt, /PROPOSES a change/);
    assert.match(auto.systemPrompt, /take effect immediately/);
    assert.match(readOnly.systemPrompt, /read-only mode/);
    // The measured failure: nemotron rendering breakEvenPriceCents: 112 as
    // "112 cents". Every mode carries the correction.
    for (const prompt of [confirm.systemPrompt, auto.systemPrompt, readOnly.systemPrompt]) {
      assert.match(prompt, /breakEvenPriceCents: 112 is \$1\.12/);
      assert.match(prompt, /Never write a money figure in cents/);
    }
  } finally {
    confirm.dispose();
    auto.dispose();
    readOnly.dispose();
  }
});

test('history is replayed so a second turn knows what the first one said', async () => {
  const { createPiggySession } = await import('../src/agent/session');
  const piggy = await createPiggySession({
    mode: 'read_only',
    tools: [fakePigTool('pig_get_workspace_summary')],
    history: [
      { role: 'user', content: 'What is utilisation on Northwind?' },
      { role: 'assistant', content: 'Northwind is at 38 per cent.' },
    ],
  });

  try {
    const messages = piggy.session.agent.state.messages;
    assert.equal(messages.length, 2);
    assert.equal(messages[0]?.role, 'user');
    assert.equal(messages[1]?.role, 'assistant');
  } finally {
    piggy.dispose();
  }
});

test('a model outside the catalogue is refused before a request is made', async () => {
  const { createPiggySession } = await import('../src/agent/session');

  await assert.rejects(
    () =>
      createPiggySession({
        mode: 'read_only',
        modelId: 'openai/gpt-4o',
        tools: [fakePigTool('pig_get_workspace_summary')],
      }),
    /not in the Piggy catalogue/,
  );
});

test('the default model is the configured one', async () => {
  const { createPiggySession } = await import('../src/agent/session');
  const { piggyDefaultModelId } = await import('../src/agent/models');
  const piggy = await createPiggySession({
    mode: 'read_only',
    tools: [fakePigTool('pig_get_workspace_summary')],
  });

  try {
    assert.equal(piggy.modelId, piggyDefaultModelId());
  } finally {
    piggy.dispose();
  }
});