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

/**
 * What the chat server tells the user, and the ledger, about a retried turn.
 *
 * `inference-retry.test.ts` pins the retry itself against the real harness.
 * This file pins the half of the same production failure that lived in PIG's
 * own code, and it is the half that was doing the visible damage.
 *
 * Measured on 2026-08-14: the harness retries a rate-limited turn of its own
 * accord and often succeeds, but `translateSessionEvent` latched
 * `state.errorMessage` on the errored `turn_end` and never cleared it, so a turn
 * that recovered and streamed a perfectly good answer was still closed as
 * `inference_failed` with the 429 in `agent_runs.error`. The reader was told
 * Piggy could not finish an answer they had just been given.
 *
 * Every session here is a double, for the same reason the stall guard's are: an
 * endpoint cannot be asked to rate limit on demand, and the point of these tests
 * is the server's reading of the events, not the transport underneath them.
 */
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';

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

const MODELS: PiggyModelOption[] = [
  {
    id: 'nvidia/nemotron-3-super-120b-a12b',
    label: 'Nemotron 3 Super',
    costPerMTokIn: 0.3,
    costPerMTokOut: 0.9,
    contextWindow: 131_072,
    reasoning: true,
    isDefault: true,
  },
];

/** The body Prime Inference really sends, verbatim from the production log. */
const RATE_LIMIT_ERROR =
  '429: {"message":"Rate limit reached. Please retry shortly.","type":"rate_limit_exceeded","code":"rate_limited"}';

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 {
  aborted: number;
}

function sessions(script: TurnScript, watched: SessionSpy) {
  return async (options: { tools: readonly ToolDefinition[]; modelId?: string }): Promise<PiggySession> => {
    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 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;
}

/** A model call the endpoint refused. This is what a 429 looks like from here. */
function failedTurn(errorMessage: string): AgentSessionEvent {
  return {
    type: 'turn_end',
    message: { role: 'assistant', usage: { input: 0, output: 0 }, stopReason: 'error', errorMessage },
    toolResults: [],
  } as unknown as AgentSessionEvent;
}

/** The harness announcing that it is about to restart the turn. */
function retryStart(errorMessage: string, attempt = 1): AgentSessionEvent {
  return {
    type: 'auto_retry_start',
    attempt,
    maxAttempts: 1,
    delayMs: 1_500,
    errorMessage,
  } as unknown as AgentSessionEvent;
}

/** 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 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-retry',
    ...overrides,
  });
}

async function turnFrames(base: string, body = chatBody()): Promise<PiggyChatEvent[]> {
  const response = await fetch(`${base}/internal/chat`, { method: 'POST', headers: authorised, body });
  return (await response.text())
    .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;
}

function answerText(frames: PiggyChatEvent[]): string {
  return frames
    .filter((frame): frame is Extract<PiggyChatEvent, { type: 'content_delta' }> => frame.type === 'content_delta')
    .map((frame) => frame.delta)
    .join('');
}

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

// ------------------------------------------------- the turn that recovered anyway

test('a turn the harness retried and finished is reported as finished', async (t) => {
  const runs: RecordedRun[] = [];
  const watched: SessionSpy = { aborted: 0 };
  const base = await startForTest(t, runs, {
    createSession: sessions(async (_tools, emit) => {
      // The 429 arrives before a byte of the answer, which is the ordinary
      // shape of one: the endpoint refuses the request rather than dropping a
      // response half way through.
      emit(failedTurn(RATE_LIMIT_ERROR));
      emit(retryStart(RATE_LIMIT_ERROR));
      emit(textDelta('Idle is $12,000.'));
      emit(turnEnd(1_240, 180));
    }, watched),
  });

  const frames = await turnFrames(base);

  // The whole of the visible bug: this used to end in an error frame with the
  // 429 in the ledger, after the reader had already been given the answer.
  assert.deepEqual(
    frames.map((frame) => frame.type),
    ['meta', 'content_delta', 'done'],
  );
  assert.equal(answerText(frames), 'Idle is $12,000.');
  assert.equal(watched.aborted, 0, 'a turn that was recovering was stopped');

  const closed = runs[0]?.closed;
  assert.equal(closed?.status, 'succeeded');
  assert.equal(closed?.error, null);
  // And an operator can still see that it cost two goes, which is the trend
  // they are watching even when every turn eventually answers.
  assert.equal(inference(closed)?.attempts, 2);
  assert.match(String(inference(closed)?.retryReason), /Rate limit reached/);
});

test('a healthy turn records one attempt rather than none', async (t) => {
  const runs: RecordedRun[] = [];
  const base = await startForTest(t, runs, {
    createSession: sessions(async (_tools, emit) => {
      emit(textDelta('Idle is $12,000.'));
      emit(turnEnd(1_240, 180));
    }, { aborted: 0 }),
  });

  const frames = await turnFrames(base);

  assert.equal(frames.at(-1)?.type, 'done');
  // Written on every turn, not only the failed ones: a day where every turn
  // needed two attempts and succeeded must not look like a day where none did.
  assert.equal(inference(runs[0]?.closed)?.attempts, 1);
  assert.equal(inference(runs[0]?.closed)?.retryReason, undefined);
});

// --------------------------------------------------- when the retries run out

test('an exhausted rate limit is its own code, and says what to do about it', async (t) => {
  const runs: RecordedRun[] = [];
  const base = await startForTest(t, runs, {
    createSession: sessions(async (_tools, emit) => {
      emit(failedTurn(RATE_LIMIT_ERROR));
      emit(retryStart(RATE_LIMIT_ERROR));
      emit(failedTurn(RATE_LIMIT_ERROR));
    }, { aborted: 0 }),
  });

  const frames = await turnFrames(base);

  // Distinct from `inference_failed`, because it wants a different response:
  // waiting ten seconds genuinely fixes it, and it is not worth a pager.
  assert.equal(errorFrame(frames)?.code, 'inference_rate_limited');
  assert.match(String(errorFrame(frames)?.message), /rate limiting us/);
  assert.match(String(errorFrame(frames)?.message), /2 times/);
  assert.match(String(errorFrame(frames)?.message), /ask again/i);
  assert.equal(answerText(frames), '');

  const closed = runs[0]?.closed;
  assert.equal(closed?.status, 'failed');
  // The ledger keeps the upstream body; the browser is never shown it.
  assert.match(String(closed?.error), /rate_limit_exceeded/);
  assert.match(String(closed?.error), /2 attempts/);
  assert.equal(inference(closed)?.attempts, 2);
});

test('a fault that is not a rate limit keeps the generic code', async (t) => {
  const runs: RecordedRun[] = [];
  const base = await startForTest(t, runs, {
    createSession: sessions(async (_tools, emit) => {
      emit(failedTurn('502: {"message":"upstream connect error"}'));
    }, { aborted: 0 }),
  });

  const frames = await turnFrames(base);

  // Somebody should look at this one, so it must not wear the name of the fault
  // that fixes itself.
  assert.equal(errorFrame(frames)?.code, 'inference_failed');
  assert.equal(errorFrame(frames)?.message, 'Piggy could not finish this answer.');
  assert.equal(inference(runs[0]?.closed)?.attempts, 1);
});

// ------------------------------------------------- what a retry may never replay

test('a retry that would repeat a delivered answer is refused', async (t) => {
  const runs: RecordedRun[] = [];
  const watched: SessionSpy = { aborted: 0 };
  const base = await startForTest(t, runs, {
    createSession: sessions(async (tools, emit, signal) => {
      // Measured against a stubbed endpoint: the harness's session-level retry
      // discards the errored assistant message and generates a replacement, so
      // a turn that had streamed "Idle is " came back as
      // "Idle is Idle is $12,000." in the transcript.
      emit(textDelta('Idle is '));
      emit(failedTurn(RATE_LIMIT_ERROR));
      emit(retryStart(RATE_LIMIT_ERROR));
      // And this script does not stop when it is told to, which is the nastier
      // shape of the same fault and the one the stall guard already assumes: a
      // harness that ignores the abort would stream the replacement answer over
      // the top of the half the reader already has. Neither the abort nor the
      // suppression is sufficient on its own.
      await untilAborted(tools, emit, signal);
      emit(textDelta('Idle is $12,000.'));
      emit(turnEnd(1_240, 180));
    }, watched),
  });

  const frames = await turnFrames(base);

  assert.equal(answerText(frames), 'Idle is ', 'the reader was shown the answer twice');
  assert.equal(watched.aborted, 1, 'the replay was allowed to proceed');
  assert.equal(errorFrame(frames)?.code, 'inference_rate_limited');
  assert.match(String(errorFrame(frames)?.message), /incomplete/);
  assert.match(String(errorFrame(frames)?.message), /already been shown/);
  assert.equal(
    frames.some((frame) => frame.type === 'done'),
    false,
    'an incomplete answer must not also report itself finished',
  );

  const closed = runs[0]?.closed;
  assert.equal(closed?.status, 'failed');
  assert.equal(closed?.summary, 'Idle is');
  assert.match(String(closed?.error), /retry refused/);
  assert.equal(inference(closed)?.attempts, 2);
});

test('a retry before anything has been delivered is left alone', async (t) => {
  const runs: RecordedRun[] = [];
  const watched: SessionSpy = { aborted: 0 };
  const base = await startForTest(t, runs, {
    createSession: sessions(async (_tools, emit) => {
      // A tool ran, so the turn is not untouched — but nothing has reached the
      // reader's transcript, so there is nothing to say twice. Stopping here
      // would throw away a recoverable turn for no gain.
      emit({
        type: 'tool_execution_start',
        toolCallId: 'call_1',
        toolName: 'pig_get_idle_capacity',
        args: {},
      } as unknown as AgentSessionEvent);
      emit(failedTurn(RATE_LIMIT_ERROR));
      emit(retryStart(RATE_LIMIT_ERROR));
      emit(textDelta('Idle is $12,000.'));
      emit(turnEnd(1_240, 180));
    }, watched),
  });

  const frames = await turnFrames(base);

  assert.equal(watched.aborted, 0, 'a safe retry was refused');
  assert.equal(frames.at(-1)?.type, 'done');
  assert.equal(answerText(frames), 'Idle is $12,000.');
  assert.equal(runs[0]?.closed?.status, 'succeeded');
});

// -------------------------------------------- the guards that outrank the retry

test('the stall watchdog outranks a pending retry', async (t) => {
  const runs: RecordedRun[] = [];
  const watched: SessionSpy = { aborted: 0 };
  const base = await startForTest(t, runs, {
    stallLimits: stallLimits({ firstProgressMs: 30_000, idleMs: 120 }),
    createSession: sessions(async (tools, emit, signal) => {
      emit(failedTurn(RATE_LIMIT_ERROR));
      emit(retryStart(RATE_LIMIT_ERROR));
      // The retry was announced and then nothing ever happened, which is the
      // shape of a backoff into an endpoint that has stopped answering
      // altogether. A retry loop that could outlive the watchdog would hang the
      // browser exactly the way the missing deadline used to.
      await untilAborted(tools, emit, signal);
    }, watched),
  });

  const frames = await turnFrames(base);

  assert.equal(errorFrame(frames)?.code, 'inference_stalled');
  assert.equal(watched.aborted, 1);
  const closed = runs[0]?.closed;
  assert.match(String(closed?.error), /idle deadline/);
  // The attempt count is still recorded: the turn really did try twice before
  // the silence, and that is what an operator is counting.
  assert.equal(inference(closed)?.attempts, 2);
});

test('the turn ceiling outranks a pending retry', async (t) => {
  const runs: RecordedRun[] = [];
  const watched: SessionSpy = { aborted: 0 };
  const base = await startForTest(t, runs, {
    limits: { maxModelCalls: 2, maxTurnTokens: 40_000, dailyLimitCents: 0 },
    createSession: sessions(async (tools, emit, signal) => {
      emit(turnEnd(1_000, 100, 'toolUse'));
      emit(failedTurn(RATE_LIMIT_ERROR));
      emit(retryStart(RATE_LIMIT_ERROR));
      await untilAborted(tools, emit, signal);
    }, watched),
  });

  const frames = await turnFrames(base);

  // A retry that resurrected a turn already stopped for cost would spend money
  // the ceiling exists to refuse.
  assert.equal(errorFrame(frames)?.code, 'turn_limit_exceeded');
  assert.equal(runs[0]?.closed?.status, 'aborted');
});