pig/apps/piggy/test/stall-guard.test.ts

/**
 * What the chat server does about a turn the endpoint stops answering.
 *
 * The failure this file pins was observed in production: `POST
 * /chat/completions` began hanging while `GET /models` still answered in 0.2s,
 * so the stream emitted its `meta` frame and then nothing at all, for ever, and
 * the transcript span until the browser gave up. A direct `fetch` from Node ran
 * past 180 seconds without settling. The harness owns the HTTP call now and sets
 * no deadline on it, so the guard has to live where PIG can see the turn: the
 * session's event stream.
 *
 * Every session here is a double, and deliberately so — the endpoint that
 * caused this cannot be asked to stall on demand, and a test that depended on it
 * would be untrustworthy in exactly the conditions it exists for. A double that
 * never settles is the same silence, and it is deterministic besides.
 */
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 { PiggyStallLimits } 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>;
}

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;
        },
      }),
    }),
    select: () => ({ from: () => ({ where: async () => [{ spent: '0' }] }) }),
  } as unknown as Database;
}

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

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

function spy(): SessionSpy {
  return { created: 0, aborted: 0, disposed: 0 };
}

/**
 * A session whose `prompt()` does whatever the script does, including nothing.
 *
 * `abort()` fires the script's signal, which is how the real harness tells a
 * turn to stop; a script that ignores it stands in for a harness that cannot
 * unwind because the socket underneath it has no deadline either.
 */
function sessions(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: () => {
        watched.disposed += 1;
        aborted.abort();
      },
    } satisfies PiggySession;
  };
}

function textDelta(delta: string): AgentSessionEvent {
  return {
    type: 'message_update',
    message: { role: 'assistant' },
    assistantMessageEvent: { type: 'text_delta', contentIndex: 0, delta },
  } as unknown as AgentSessionEvent;
}

function turnEnd(input: number, output: number, stopReason = 'stop'): 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;
}

/** The harness's own bookkeeping, which is not the model doing any work. */
function turnStart(): AgentSessionEvent {
  return { type: 'turn_start' } as unknown as AgentSessionEvent;
}

function stallLimits(overrides: Partial<PiggyStallLimits> = {}): PiggyStallLimits {
  return { firstProgressMs: 5_000, idleMs: 5_000, ...overrides };
}

async function startForTest(
  t: { after: (fn: () => void) => void },
  runs: RecordedRun[],
  options: Partial<PiggyChatServerOptions>,
): Promise<string> {
  const server = startPiggyChatServer(fakeDatabase(runs), {
    port: 0,
    internalToken: TOKEN,
    models: MODELS,
    createReadTools: () => [],
    createWriteTools: () => [],
    limits: { maxModelCalls: 8, maxTurnTokens: 40_000, dailyLimitCents: 0 },
    stallLimits: stallLimits(),
    ...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-stall',
    ...overrides,
  });
}

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

function errorFrame(frames: PiggyChatEvent[]): { message: string; code?: string } | null {
  const frame = frames.at(-1);
  return frame?.type === 'error' ? { message: frame.message, ...(frame.code ? { code: frame.code } : {}) } : null;
}

/** Silence, until somebody tells the turn to stop. A harness that unwinds. */
const untilAborted: TurnScript = (_tools, _emit, signal) =>
  new Promise<void>((resolve) => {
    if (signal.aborted) {
      resolve();
      return;
    }
    signal.addEventListener('abort', () => resolve(), { once: true });
  });

function readStall(closed: Record<string, unknown> | undefined): Record<string, unknown> | undefined {
  return (closed?.result as { stall?: Record<string, unknown> } | undefined)?.stall;
}

// ------------------------------------------------------- the endpoint goes quiet

test('a turn the endpoint never answers is ended by the first-progress deadline', async (t) => {
  const runs: RecordedRun[] = [];
  const watched = spy();
  const base = await startForTest(t, runs, {
    stallLimits: stallLimits({ firstProgressMs: 120 }),
    createSession: sessions(untilAborted, watched),
  });

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

  // The whole bug, in one assertion: this used to hang until the browser gave
  // up, and now it settles inside the deadline it was given.
  assert.ok(elapsed < 2_000, `the turn took ${elapsed}ms to give up`);
  assert.equal(frames[0]?.type, 'meta');
  assert.equal(errorFrame(frames)?.code, 'inference_stalled');
  assert.match(String(errorFrame(frames)?.message), /never answered/);
  assert.equal(
    frames.some((frame) => frame.type === 'done'),
    false,
    'a stalled turn must not also report itself finished',
  );

  // The session is told to stop rather than left generating into nothing.
  assert.equal(watched.aborted, 1);
  assert.ok(watched.disposed >= 1);

  // And an operator can tell a silent endpoint from a fault without a log: the
  // reason names the deadline, and `result.stall` names which of the two it was.
  const closed = runs[0]?.closed;
  assert.equal(closed?.status, 'failed');
  assert.match(String(closed?.error), /first_progress deadline/);
  assert.equal(readStall(closed)?.phase, 'first_progress');
  assert.equal(readStall(closed)?.ceilingMs, 120);
  assert.ok(Number(readStall(closed)?.waitedMs) >= 120);
});

test("the harness's own bookkeeping does not count as the model working", async (t) => {
  const runs: RecordedRun[] = [];
  const watched = spy();
  const base = await startForTest(t, runs, {
    stallLimits: stallLimits({ firstProgressMs: 150, idleMs: 30_000 }),
    createSession: sessions(async (tools, emit, signal) => {
      // `turn_start` is announced the instant a prompt is submitted, before a
      // byte has left the process. If it counted as progress the turn would
      // fall into the far more generous idle window and the hang would be back.
      emit(turnStart());
      await untilAborted(tools, emit, signal);
    }, watched),
  });

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

  assert.equal(errorFrame(frames)?.code, 'inference_stalled');
  assert.equal(readStall(runs[0]?.closed)?.phase, 'first_progress');
});

test('a turn that goes quiet part way through is ended by the idle deadline', async (t) => {
  const runs: RecordedRun[] = [];
  const watched = spy();
  const base = await startForTest(t, runs, {
    stallLimits: stallLimits({ firstProgressMs: 30_000, idleMs: 120 }),
    createSession: sessions(async (tools, emit, signal) => {
      emit(toolStart('call_1', 'pig_get_idle_capacity'));
      emit(textDelta('Idle is '));
      // The socket dies here, mid-sentence, and never says another word.
      await untilAborted(tools, emit, signal);
    }, watched),
  });

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

  // What did arrive is still shown; the reader is told it is not the whole of
  // the answer rather than being left with a truncated one that looks finished.
  assert.ok(frames.some((frame) => frame.type === 'content_delta'));
  assert.equal(errorFrame(frames)?.code, 'inference_stalled');
  assert.match(String(errorFrame(frames)?.message), /went quiet/);
  assert.equal(watched.aborted, 1);

  const closed = runs[0]?.closed;
  assert.equal(closed?.status, 'failed');
  assert.equal(closed?.summary, 'Idle is');
  assert.match(String(closed?.error), /idle deadline/);
  assert.equal(readStall(closed)?.phase, 'idle');
  assert.equal(readStall(closed)?.ceilingMs, 120);
});

test('a stall is not reported as a fault, and a fault is not reported as a stall', async (t) => {
  // Three things can end a turn early and they want three different responses
  // from whoever reads the code: wait, investigate, and do nothing. They must
  // not share a name.
  const runs: RecordedRun[] = [];
  const base = await startForTest(t, runs, {
    stallLimits: stallLimits({ firstProgressMs: 30_000, idleMs: 30_000 }),
    createSession: sessions(async (_tools, emit) => {
      emit(textDelta('Idle is '));
      emit({
        type: 'turn_end',
        message: {
          role: 'assistant',
          usage: { input: 120, output: 4 },
          stopReason: 'error',
          errorMessage: 'upstream returned 502',
        },
        toolResults: [],
      } as unknown as AgentSessionEvent);
    }, spy()),
  });

  const frames = parseFrames(
    await fetch(`${base}/internal/chat`, { method: 'POST', headers: authorised, body: chatBody() })
      .then((response) => response.text()),
  );
  assert.equal(errorFrame(frames)?.code, 'inference_failed');
  assert.equal(readStall(runs[0]?.closed), undefined);
});

// ----------------------------------------------------- what must NOT be killed

test('a slow but progressing answer is never cut off, however long it takes', async (t) => {
  const runs: RecordedRun[] = [];
  const watched = spy();
  // Twelve chunks, 40ms apart: 480ms in total, which is four times the idle
  // deadline and twice the first-progress one. A flat deadline over the turn —
  // the obvious implementation, and the wrong one — would kill this, and it is
  // precisely the long answer the product exists to give.
  const base = await startForTest(t, runs, {
    stallLimits: stallLimits({ firstProgressMs: 250, idleMs: 120 }),
    createSession: sessions(async (_tools, emit) => {
      for (let index = 0; index < 12; index += 1) {
        await new Promise((resolve) => setTimeout(resolve, 40));
        emit(textDelta(`part ${index} `));
      }
      emit(turnEnd(4_000, 400));
    }, watched),
  });

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

  assert.ok(Date.now() - started >= 400, 'the turn did not actually run long');
  assert.equal(frames.at(-1)?.type, 'done');
  assert.equal(
    frames.some((frame) => frame.type === 'error'),
    false,
    'a turn that kept arriving was killed for taking a while',
  );
  assert.equal(watched.aborted, 0);
  assert.equal(runs[0]?.closed?.status, 'succeeded');
  assert.equal(readStall(runs[0]?.closed), undefined);
});

/** A write tool that parks on a human, the way `confirm` mode really does. */
function proposingWriteTools(applied: string[]): (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' }],
        });
        if (decision === 'apply') applied.push('applied');
        return {
          content: [{ type: 'text', text: `The change was ${decision}.` }],
          details: { tool: 'pig_log_activity', status: decision === 'apply' ? 'applied' : 'declined' },
        };
      },
    } as unknown as ToolDefinition,
  ];
}

test('a write parked on a human outlives the idle deadline and still applies', async (t) => {
  const runs: RecordedRun[] = [];
  const applied: string[] = [];
  const watched = spy();
  // The card is left on screen for five times the idle deadline. A turn parked
  // on `propose()` emits nothing at all by design, so a watchdog that could not
  // see the rendezvous would kill every write Piggy ever proposed — and it
  // would do it to the one flow where being killed loses real work.
  const base = await startForTest(t, runs, {
    stallLimits: stallLimits({ firstProgressMs: 500, idleMs: 100 }),
    approvalTimeoutMs: 30_000,
    createWriteTools: proposingWriteTools(applied),
    createSession: sessions(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(toolStart('call_1', 'pig_log_activity'));
      const result = await tool.execute('call_1', {}, signal, undefined, undefined as never);
      emit({
        type: 'tool_execution_end',
        toolCallId: 'call_1',
        toolName: 'pig_log_activity',
        result,
        isError: false,
      } as unknown as AgentSessionEvent);
      emit(textDelta('Logged.'));
      emit(turnEnd(200, 20));
    }, watched),
  });

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

  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');

  const thinking = Date.now();
  await new Promise((resolve) => setTimeout(resolve, 500));
  const decision = await fetch(`${base}/internal/approve`, {
    method: 'POST',
    headers: authorised,
    body: JSON.stringify({
      conversationId: 'conv-stall',
      changeId: asked.change.id,
      decision: 'apply',
    }),
  });
  assert.equal(decision.status, 202);
  assert.ok(Date.now() - thinking >= 500, 'the human did not actually take their time');

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

  assert.equal(frames.at(-1)?.type, 'done');
  assert.equal(
    frames.some((frame) => frame.type === 'error'),
    false,
    'a turn waiting on a person was reported as a silent endpoint',
  );
  // And it did not merely survive: the change the human approved was applied.
  assert.deepEqual(applied, ['applied']);
  const result = frames.find((frame) => frame.type === 'tool_result');
  assert.deepEqual(result?.type === 'tool_result' ? result.result : null, {
    tool: 'pig_log_activity',
    status: 'applied',
  });
  assert.equal(watched.aborted, 0);
  assert.equal(runs[0]?.closed?.status, 'succeeded');
});

test('the happy path is untouched', async (t) => {
  const runs: RecordedRun[] = [];
  const watched = spy();
  const base = await startForTest(t, runs, {
    createSession: sessions(async (_tools, emit) => {
      emit(toolStart('call_1', 'pig_get_idle_capacity'));
      emit(textDelta('Idle is $12,000.'));
      emit(turnEnd(1_240, 180));
    }, watched),
  });

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

  assert.deepEqual(
    frames.map((frame) => frame.type),
    ['meta', 'tool_call', 'content_delta', 'done'],
  );
  assert.equal(watched.aborted, 0);
  const closed = runs[0]?.closed;
  assert.equal(closed?.status, 'succeeded');
  assert.equal(closed?.error, null);
  assert.equal(readStall(closed), undefined);
});

// ---------------------------------------------------- a harness that will not stop

test('a harness that ignores the abort still gives the browser its answer', async (t) => {
  const runs: RecordedRun[] = [];
  const watched = spy();
  // The nastier shape of the same fault: the session is told to stop and the
  // request underneath it has no deadline either, so `prompt()` never settles.
  // Trusting that promise would rebuild the hang one level up, so the turn is
  // raced against the stall and ends anyway.
  const base = await startForTest(t, runs, {
    stallLimits: stallLimits({ firstProgressMs: 100 }),
    createSession: sessions(() => new Promise<void>(() => {}), watched),
  });

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

  assert.equal(errorFrame(frames)?.code, 'inference_stalled');
  assert.equal(watched.aborted, 1, 'the session was told to stop, even though it did not');
  // Long enough to have waited for a clean unwind, short enough to be nothing
  // like the three minutes the endpoint spent not answering.
  assert.ok(elapsed >= 100, `the turn ended in ${elapsed}ms, before its own deadline`);
  assert.ok(elapsed < 10_000, `the turn took ${elapsed}ms to give up`);
  assert.equal(runs[0]?.closed?.status, 'failed');
  assert.equal(readStall(runs[0]?.closed)?.phase, 'first_progress');
});