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

/**
 * What Piggy does when Prime Inference says "please retry shortly".
 *
 * The failure this file pins was measured on production on 2026-08-14, roughly
 * every other turn:
 *
 *     [piggy] chat turn ended in an inference error: 429:
 *     {"message":"Rate limit reached. Please retry shortly.",
 *      "type":"rate_limit_exceeded","code":"rate_limited"}
 *
 * A `curl` a second later succeeded, so these were transient bursts and the
 * endpoint was telling us what to do about them. Nothing did.
 *
 * The endpoint cannot be asked to rate limit on demand, and a test that waited
 * for it to happen would be untrustworthy in exactly the conditions it exists
 * for, so every upstream here is a stub installed over `globalThis.fetch`. That
 * is a real seam and not a convenience: the OpenAI client the harness builds
 * resolves its fetch through `getDefaultFetch()` at construction, and it
 * constructs one per model call (openai@6.26.0 internal/shims.js:9-14), so a
 * stub installed before `prompt()` is the transport the harness genuinely uses.
 * Everything below therefore runs the real `createPiggySession`, the real
 * harness and the real OpenAI SDK against a fake endpoint — the retry is the
 * only thing under test, and none of it is mocked.
 */
import assert from 'node:assert/strict';
import { mkdtempSync, rmSync, writeFileSync } from 'node:fs';
import { tmpdir } from 'node:os';
import { join } from 'node:path';
import test, { after, before } from 'node:test';
import {
  createAgentSession,
  defineTool,
  ModelRuntime,
  SessionManager,
  SettingsManager,
  type ToolDefinition,
} from '@earendil-works/pi-coding-agent';
import { Type } from 'typebox';
import { piggyDefaultModelId, piggyModelsJsonText, PIGGY_PROVIDER_ID } from '../src/agent/models';
import {
  piggyAgentSettings,
  PIGGY_INFERENCE_RETRY,
  type PiggyInferenceRetryPolicy,
} from '../src/agent/session';

const agentDir = mkdtempSync(join(tmpdir(), 'piggy-retry-test-'));
const realFetch = globalThis.fetch;

before(() => {
  process.env.DATABASE_URL = 'postgres://pig:pig@localhost:54330/pig';
  process.env.PIGGY_INTERNAL_TOKEN = 'test-internal-token-for-piggy-000000';
  // Deliberately fake. Nothing below leaves the process, and a test that needs
  // a live key is a test that fails in CI.
  process.env.PRIME_API_KEY = 'test-key-not-used-offline';
  process.env.PIGGY_AGENT_DIR = agentDir;
});

after(() => {
  globalThis.fetch = realFetch;
  rmSync(agentDir, { recursive: true, force: true });
});

// ------------------------------------------------------------- the fake endpoint

const MODEL = piggyDefaultModelId();

function chunk(delta: unknown, finish: string | null, usage?: unknown): string {
  return JSON.stringify({
    id: 'chatcmpl-test',
    object: 'chat.completion.chunk',
    created: 1,
    model: MODEL,
    choices: [{ index: 0, delta, finish_reason: finish }],
    ...(usage ? { usage } : {}),
  });
}

function eventStream(chunks: string[], terminated = true): Response {
  const body = chunks.map((line) => `data: ${line}\n\n`).join('') + (terminated ? 'data: [DONE]\n\n' : '');
  return new Response(body, { status: 200, headers: { 'content-type': 'text/event-stream' } });
}

/** A complete, ordinary answer. */
function answers(text = 'Idle is $12,000.'): Response {
  return eventStream([
    chunk({ role: 'assistant', content: text }, null),
    chunk({}, 'stop', { prompt_tokens: 100, completion_tokens: 8, total_tokens: 108 }),
  ]);
}

/** One tool call and nothing else, which is how a tool-using turn starts. */
function callsTool(name: string): Response {
  return eventStream([
    chunk(
      {
        role: 'assistant',
        tool_calls: [
          { index: 0, id: 'call_1', type: 'function', function: { name, arguments: '{}' } },
        ],
      },
      null,
    ),
    chunk({}, 'tool_calls', { prompt_tokens: 100, completion_tokens: 8, total_tokens: 108 }),
  ]);
}

/** The body Prime Inference really sends, verbatim from the production log. */
function rateLimited(retryAfterSeconds?: number): Response {
  return new Response(
    JSON.stringify({
      message: 'Rate limit reached. Please retry shortly.',
      type: 'rate_limit_exceeded',
      code: 'rate_limited',
    }),
    {
      status: 429,
      headers: {
        'content-type': 'application/json',
        ...(retryAfterSeconds === undefined ? {} : { 'retry-after': String(retryAfterSeconds) }),
      },
    },
  );
}

function failsWith(status: number, message: string): Response {
  return new Response(JSON.stringify({ message }), {
    status,
    headers: { 'content-type': 'application/json' },
  });
}

interface Upstream {
  /** When each request arrived, in milliseconds since the stub was installed. */
  readonly at: number[];
  readonly count: number;
}

/** Installs a stub over the global fetch and records every request it sees. */
function upstream(reply: (attempt: number) => Response | Promise<Response>): Upstream {
  const at: number[] = [];
  const started = Date.now();
  globalThis.fetch = (async (_input: unknown, init?: RequestInit) => {
    at.push(Date.now() - started);
    const response = await reply(at.length);
    // The caller's signal is honoured so that a stub which never answers can
    // still be cancelled by a deadline, which is the whole point of one.
    if (init?.signal?.aborted) throw init.signal.reason;
    return response;
  }) as typeof fetch;
  return {
    at,
    get count() {
      return at.length;
    },
  };
}

/** A stub that never answers, and unblocks only when the request is abandoned. */
function silence(): Upstream {
  const at: number[] = [];
  const started = Date.now();
  globalThis.fetch = ((_input: unknown, init?: RequestInit) => {
    at.push(Date.now() - started);
    return new Promise<Response>((_resolve, reject) => {
      const signal = init?.signal;
      if (!signal) return;
      if (signal.aborted) {
        reject(signal.reason);
        return;
      }
      signal.addEventListener('abort', () => reject(signal.reason), { once: true });
    });
  }) as typeof fetch;
  return {
    at,
    get count() {
      return at.length;
    },
  };
}

// ------------------------------------------------------------------ the fixtures

function countingTool(name: string, runs: { count: number }): ToolDefinition {
  return defineTool({
    name,
    label: name,
    description: `Test double for ${name}.`,
    promptSnippet: `${name}: test double.`,
    parameters: Type.Object({}),
    async execute() {
      runs.count += 1;
      return { content: [{ type: 'text' as const, text: '{"ok":true}' }], details: { tool: name } };
    },
  });
}

interface TurnResult {
  /** Everything the reader would have been shown, concatenated. */
  text: string;
  /** How the last model call ended, as the harness reports it. */
  errorMessage?: string;
  stopReason?: string;
  /** Retries the harness announced, which are the ones that replay work. */
  announcedRetries: number;
  elapsedMs: number;
}

/** One real Piggy turn, driven through the real `createPiggySession`. */
async function drive(tools: ToolDefinition[], message = 'What is idle costing us?'): Promise<TurnResult> {
  const { createPiggySession } = await import('../src/agent/session');
  const piggy = await createPiggySession({ mode: 'read_only', tools });
  const result: TurnResult = { text: '', announcedRetries: 0, elapsedMs: 0 };
  const started = Date.now();
  const unsubscribe = piggy.session.subscribe((event) => {
    if (event.type === 'message_update' && event.assistantMessageEvent.type === 'text_delta') {
      result.text += event.assistantMessageEvent.delta;
    }
    if (event.type === 'auto_retry_start') result.announcedRetries += 1;
    if (event.type === 'turn_end') {
      const assistant = event.message as { stopReason?: string; errorMessage?: string };
      result.stopReason = assistant.stopReason;
      result.errorMessage = assistant.errorMessage;
    }
  });
  try {
    await piggy.session.prompt(message);
  } finally {
    unsubscribe();
    result.elapsedMs = Date.now() - started;
    piggy.dispose();
  }
  return result;
}

// ------------------------------------------------------- the measured production bug

test('a 429 that clears on the next attempt is answered rather than reported', async () => {
  // The bug, in one test. Before the policy existed the harness made exactly
  // one attempt per model call — `retryProviderRequest` defaults `maxRetries`
  // to 0 and the settings supplied none — so this turn ended as
  // `inference_failed` with no answer at all.
  const runs = { count: 0 };
  const endpoint = upstream((attempt) => (attempt === 1 ? rateLimited() : answers()));

  const turn = await drive([countingTool('pig_get_workspace_summary', runs)]);

  assert.equal(endpoint.count, 2, 'the refusal was not retried');
  assert.equal(turn.errorMessage, undefined);
  assert.equal(turn.stopReason, 'stop');
  assert.equal(turn.text, 'Idle is $12,000.');
});

test('a retried turn shows the reader one answer, not two', async () => {
  // The constraint that makes the seam matter. The retry happens where the
  // response has not begun, so there is nothing to replay — no delta is emitted
  // twice, and the harness never has to announce a retry at all.
  const runs = { count: 0 };
  upstream((attempt) => (attempt <= 2 ? rateLimited() : answers('Idle is $12,000.')));

  const turn = await drive([countingTool('pig_get_workspace_summary', runs)]);

  assert.equal(turn.text, 'Idle is $12,000.');
  assert.equal(
    turn.text.indexOf('Idle is'),
    turn.text.lastIndexOf('Idle is'),
    'the answer was streamed to the reader twice',
  );
  assert.equal(turn.announcedRetries, 0, 'the turn was restarted when it did not need to be');
});

test('a retry never re-runs a tool that has already run', async () => {
  // The expensive property. `pig_log_activity` writes a row; a retry that
  // re-executed it would write it twice and no diff card would be shown for the
  // second one. The tool is called on the first model call, the SECOND model
  // call is the one that is rate limited, and the tool must not move.
  const runs = { count: 0 };
  const endpoint = upstream((attempt) => {
    if (attempt === 1) return callsTool('pig_log_activity');
    if (attempt === 2) return rateLimited();
    return answers('Logged.');
  });

  const turn = await drive([countingTool('pig_log_activity', runs)], 'Log a call on Northwind.');

  assert.equal(endpoint.count, 3);
  assert.equal(runs.count, 1, 'the tool ran again on the retry');
  assert.equal(turn.text, 'Logged.');
  assert.equal(turn.errorMessage, undefined);
});

test('Retry-After is honoured when the endpoint sends one', async () => {
  const runs = { count: 0 };
  const endpoint = upstream((attempt) => (attempt === 1 ? rateLimited(1) : answers()));

  const turn = await drive([countingTool('pig_get_workspace_summary', runs)]);

  assert.equal(endpoint.count, 2);
  assert.equal(turn.errorMessage, undefined);
  // A second is far longer than the jittered backoff this attempt would have
  // chosen for itself (500ms, minus up to a quarter), so waiting it out is only
  // possible if the header was read.
  const waited = endpoint.at[1]! - endpoint.at[0]!;
  assert.ok(waited >= 900, `waited ${waited}ms, so Retry-After was ignored`);
  assert.ok(waited < 3_000, `waited ${waited}ms, which is longer than was asked for`);
});

test('a refusal with no Retry-After still backs off, and briefly', async () => {
  // Jitter matters more than the curve: without it every open chat that hit the
  // same limit retries in lockstep and reproduces the limit that caused it.
  const runs = { count: 0 };
  const endpoint = upstream((attempt) => (attempt === 1 ? rateLimited() : answers()));

  await drive([countingTool('pig_get_workspace_summary', runs)]);

  const waited = endpoint.at[1]! - endpoint.at[0]!;
  assert.ok(waited > 0, 'the retry was fired immediately, which reproduces the limit');
  assert.ok(waited < 2_000, `waited ${waited}ms without being asked to`);
});

test('a rate limit that never clears is reported, and inside a bearable wait', async () => {
  const runs = { count: 0 };
  const endpoint = upstream(() => rateLimited());

  const turn = await drive([countingTool('pig_get_workspace_summary', runs)]);

  assert.match(String(turn.errorMessage), /429/);
  assert.equal(turn.stopReason, 'error');
  assert.equal(turn.text, '');
  // Every attempt the policy buys was spent: the request-level budget, twice
  // over, because the turn-level budget allows one restart of a turn that got
  // nothing from the endpoint.
  assert.equal(endpoint.count, PIGGY_INFERENCE_RETRY.attempts * PIGGY_INFERENCE_RETRY.streamAttempts);
  // Nobody may be left staring at a docked panel for a minute to be told no.
  assert.ok(turn.elapsedMs < 30_000, `the failure took ${turn.elapsedMs}ms to arrive`);
});

test('a 500 is retried and a 400 is not', async () => {
  const runs = { count: 0 };
  const serverError = upstream((attempt) =>
    attempt === 1 ? failsWith(500, 'internal error') : answers(),
  );
  const recovered = await drive([countingTool('pig_get_workspace_summary', runs)]);
  assert.equal(serverError.count, 2, 'a 5xx is transient and should have been retried');
  assert.equal(recovered.errorMessage, undefined);

  // A 4xx that is not 429 will fail identically however often it is retried,
  // and each attempt costs a round trip and a place in the queue.
  const badRequest = upstream(() => failsWith(400, 'unknown parameter'));
  const refused = await drive([countingTool('pig_get_workspace_summary', runs)]);
  assert.equal(badRequest.count, 1, 'a 400 was retried, which can only ever fail again');
  assert.equal(refused.stopReason, 'error');
  assert.match(String(refused.errorMessage), /400/);
});

test('a caller who hangs up wins over the retry', async () => {
  // A retry loop that resurrects an abandoned turn is worse than the bug: it
  // spends credit generating an answer nobody will read, and it does it while
  // the reader has already gone.
  const { createPiggySession } = await import('../src/agent/session');
  const endpoint = upstream(() => rateLimited());
  const runs = { count: 0 };
  const piggy = await createPiggySession({
    mode: 'read_only',
    tools: [countingTool('pig_get_workspace_summary', runs)],
  });
  try {
    const prompt = piggy.session.prompt('What is idle costing us?');
    // Long enough for the first attempt to have been refused and the second to
    // be sleeping on its backoff, which is where an abort has to be honoured.
    await new Promise((resolve) => setTimeout(resolve, 250));
    const seenBeforeAbort = endpoint.count;
    await piggy.session.abort();
    await prompt;
    await new Promise((resolve) => setTimeout(resolve, 400));

    assert.ok(seenBeforeAbort >= 1, 'the turn had not started, so nothing was proved');
    assert.equal(
      endpoint.count,
      seenBeforeAbort,
      'the retry carried on asking after the caller had gone',
    );
  } finally {
    piggy.dispose();
  }
});

// ------------------------------------------ the deadline the model entry cannot carry

/**
 * A bare harness session, wired the way `createPiggySession` wires one but with
 * a policy of the test's choosing.
 *
 * Built by hand rather than through `createPiggySession` because the shipped
 * deadline is twenty seconds and a test may not take twenty seconds to prove
 * one. What it proves is a fact about the INSTALLED package rather than about
 * PIG's wiring — that `retry.provider.timeoutMs` and `retry.provider.maxRetries`
 * are read and acted on — and the wiring itself is proved by every test above,
 * all of which go through the real `createPiggySession`.
 */
async function bareSession(policy: PiggyInferenceRetryPolicy, tools: ToolDefinition[]) {
  const modelsPath = join(agentDir, 'models-for-timeout-test.json');
  writeFileSync(modelsPath, piggyModelsJsonText(), { mode: 0o600 });
  const modelRuntime = await ModelRuntime.create({ modelsPath, allowModelNetwork: false });
  await modelRuntime.setRuntimeApiKey(PIGGY_PROVIDER_ID, 'test-key-not-used-offline');
  const model = modelRuntime.getModel(PIGGY_PROVIDER_ID, MODEL);
  assert.ok(model, 'the default model should be registered');
  const { session } = await createAgentSession({
    agentDir,
    cwd: agentDir,
    modelRuntime,
    model,
    settingsManager: SettingsManager.inMemory(piggyAgentSettings(policy)),
    thinkingLevel: 'off',
    noTools: 'all',
    tools: tools.map((tool) => tool.name),
    customTools: tools,
    sessionManager: SessionManager.inMemory(),
  });
  return session;
}

test('the request deadline is read from the settings the runtime is built with', async () => {
  // The stall watchdog is the outer guard and it stays; this is the deadline
  // underneath it, on one HTTP request rather than on the turn. Without it a
  // hung fetch has only the harness's own five-minute idle default.
  const endpoint = silence();
  const runs = { count: 0 };
  const session = await bareSession(
    { ...PIGGY_INFERENCE_RETRY, headersTimeoutMs: 150, streamAttempts: 1 },
    [countingTool('pig_get_workspace_summary', runs)],
  );
  let errorMessage: string | undefined;
  session.subscribe((event) => {
    if (event.type === 'turn_end') {
      errorMessage = (event.message as { errorMessage?: string }).errorMessage;
    }
  });

  const started = Date.now();
  await session.prompt('What is idle costing us?');
  const elapsed = Date.now() - started;

  // Every attempt was abandoned at its own deadline and the next one started,
  // which is only possible if BOTH fields reached the transport.
  assert.equal(endpoint.count, PIGGY_INFERENCE_RETRY.attempts);
  assert.ok(elapsed >= 150, `gave up after ${elapsed}ms, before the deadline it was given`);
  assert.ok(elapsed < 10_000, `took ${elapsed}ms, so the deadline was not honoured`);
  assert.ok(errorMessage, 'a hung request ended as a success');
  await session.abort();
  session.dispose();
});

test('the settings the harness reads are exactly the policy PIG declares', () => {
  // Read back through the installed `SettingsManager` rather than compared to
  // the object we wrote, because the field names and their nesting are the
  // whole risk: a policy under a key the harness has never heard of parses,
  // loads and does nothing, and there is no error anywhere to say so.
  const manager = SettingsManager.inMemory(piggyAgentSettings());
  const provider = manager.getProviderRetrySettings();
  const turn = manager.getRetrySettings();

  assert.equal(provider.timeoutMs, PIGGY_INFERENCE_RETRY.headersTimeoutMs);
  assert.equal(provider.maxRetries, PIGGY_INFERENCE_RETRY.attempts - 1);
  assert.equal(provider.maxRetryDelayMs, PIGGY_INFERENCE_RETRY.maxRetryDelayMs);
  assert.equal(turn.enabled, true);
  assert.equal(turn.maxRetries, PIGGY_INFERENCE_RETRY.streamAttempts - 1);
  assert.equal(turn.baseDelayMs, PIGGY_INFERENCE_RETRY.streamBackoffMs);

  // The default this replaces, and the reason the bug existed: the harness
  // ships no provider retry budget at all, and `retryProviderRequest` reads a
  // missing budget as zero.
  assert.equal(SettingsManager.inMemory().getProviderRetrySettings().maxRetries, undefined);
});

test('models.json carries no request timeout, because the harness would ignore one', () => {
  // The obvious place to put a request deadline is beside `contextWindow`, and
  // it does nothing there. `ModelDefinitionSchema` in the installed harness has
  // no `timeoutMs`; neither does `Model` in `@earendil-works/pi-ai`; and the
  // only reader is `options.timeoutMs`, which the agent loop never populates.
  // A `timeoutMs` written into a model entry validates, loads, freezes and is
  // dropped in silence, so this asserts its absence rather than its presence.
  const document = JSON.parse(piggyModelsJsonText()) as {
    providers: Record<string, { models: Record<string, unknown>[] }>;
  };
  for (const model of document.providers[PIGGY_PROVIDER_ID]?.models ?? []) {
    assert.equal(
      'timeoutMs' in model,
      false,
      `${String(model.id)} declares a timeoutMs that nothing reads; the deadline belongs in piggyAgentSettings()`,
    );
  }
});