feat: unlimited turns + tests

This commit is contained in:
Willem van den Ende 2026-04-23 15:56:16 +01:00
parent cab445e603
commit a4181af13e
3 changed files with 591 additions and 8 deletions

View File

@ -0,0 +1,484 @@
import { describe, it, beforeEach, afterEach, mock } from "node:test";
import assert from "node:assert/strict";
import { checkTurnLimit, getMaxTurns } from "./turn-limit.ts";
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import initExtension from "./turn-limit.ts";
// ============================================================================
// P: Pure function tests — checkTurnLimit
// ============================================================================
describe("checkTurnLimit", () => {
it("P1: below limit returns exceeded=false", () => {
const result = checkTurnLimit(20, 25);
assert.equal(result.exceeded, false);
assert.equal(result.turnIndex, 20);
assert.equal(result.maxTurns, 25);
});
it("P2: at limit returns exceeded=false (strict >)", () => {
const result = checkTurnLimit(25, 25);
assert.equal(result.exceeded, false);
});
it("P3: above limit returns exceeded=true", () => {
const result = checkTurnLimit(26, 25);
assert.equal(result.exceeded, true);
});
it("P4: zero max — turnIndex 1 exceeds", () => {
const result = checkTurnLimit(1, 0);
assert.equal(result.exceeded, true);
});
it("P5: zero turn — below any positive limit", () => {
const result = checkTurnLimit(0, 25);
assert.equal(result.exceeded, false);
});
});
// ============================================================================
// C: Config tests — getMaxTurns
// ============================================================================
describe("getMaxTurns", () => {
let originalEnv: string | undefined;
beforeEach(() => {
originalEnv = process.env.PI_MAX_TURNS;
});
afterEach(() => {
if (originalEnv === undefined) {
delete process.env.PI_MAX_TURNS;
} else {
process.env.PI_MAX_TURNS = originalEnv;
}
});
it("C1: returns 25 when PI_MAX_TURNS is unset", () => {
delete process.env.PI_MAX_TURNS;
assert.equal(getMaxTurns(), 25);
});
it("C2: returns env value when PI_MAX_TURNS is valid", () => {
process.env.PI_MAX_TURNS = "10";
assert.equal(getMaxTurns(), 10);
});
it("C3: falls back to 25 for non-numeric PI_MAX_TURNS", () => {
process.env.PI_MAX_TURNS = "abc";
assert.equal(getMaxTurns(), 25);
});
it("C4: falls back to 25 for PI_MAX_TURNS=0", () => {
process.env.PI_MAX_TURNS = "0";
assert.equal(getMaxTurns(), 25);
});
it("C5: falls back to 25 for negative PI_MAX_TURNS", () => {
process.env.PI_MAX_TURNS = "-5";
assert.equal(getMaxTurns(), 25);
});
});
// ============================================================================
// Helpers for integration tests
// ============================================================================
type EventHandler = (event: any, ctx: any) => Promise<void>;
type CommandHandler = { description: string; handler: (args: string, ctx: any) => Promise<void> };
function createMockPi() {
const handlers: Record<string, EventHandler> = {};
const commands: Record<string, CommandHandler> = {};
const pi = {
on: (event: string, handler: EventHandler) => {
handlers[event] = handler;
},
registerCommand: (name: string, cmd: CommandHandler) => {
commands[name] = cmd;
},
} as unknown as ExtensionAPI;
return { pi, handlers, commands };
}
function createMockCtx(options: { hasUI?: boolean; confirmResult?: boolean } = {}) {
const { hasUI = true, confirmResult = true } = options;
const calls: { method: string; args: any[] }[] = [];
return {
ctx: {
hasUI,
abort: mock.fn(() => { calls.push({ method: "abort", args: [] }); }),
ui: {
setWidget: mock.fn((...args: any[]) => { calls.push({ method: "setWidget", args }); }),
notify: mock.fn((...args: any[]) => { calls.push({ method: "notify", args }); }),
confirm: mock.fn(async () => { calls.push({ method: "confirm", args: [] }); return confirmResult; }),
},
},
calls,
};
}
// ============================================================================
// E: Entity & state tests
// ============================================================================
describe("extension handler — entity state", () => {
it("E1: counter starts at 0 — widget shows 0/25 on agent_start", async () => {
const { pi, handlers } = createMockPi();
initExtension(pi);
const { ctx } = createMockCtx();
await handlers["agent_start"]({}, ctx);
assert.equal(ctx.ui.setWidget.mock.callCount(), 1);
assert.deepEqual(ctx.ui.setWidget.mock.calls[0].arguments, ["turn-limit", ["Turns: 0/25"]]);
});
it("E2: counter increments each turn_start", async () => {
const { pi, handlers } = createMockPi();
initExtension(pi);
const { ctx } = createMockCtx();
await handlers["agent_start"]({}, ctx);
await handlers["turn_start"]({}, ctx);
// After 1 turn: widget should show 1/25
const lastCall = ctx.ui.setWidget.mock.calls.at(-1);
assert.deepEqual(lastCall!.arguments, ["turn-limit", ["Turns: 1/25"]]);
await handlers["turn_start"]({}, ctx);
const lastCall2 = ctx.ui.setWidget.mock.calls.at(-1);
assert.deepEqual(lastCall2!.arguments, ["turn-limit", ["Turns: 2/25"]]);
});
it("E3: agent_start resets counter", async () => {
const { pi, handlers } = createMockPi();
initExtension(pi);
const { ctx } = createMockCtx();
await handlers["agent_start"]({}, ctx);
await handlers["turn_start"]({}, ctx);
await handlers["turn_start"]({}, ctx);
// New agent_start resets
await handlers["agent_start"]({}, ctx);
const lastCall = ctx.ui.setWidget.mock.calls.at(-1);
assert.deepEqual(lastCall!.arguments, ["turn-limit", ["Turns: 0/25"]]);
});
});
// ============================================================================
// R: Rule tests — TurnLimitReached
// ============================================================================
describe("extension handler — TurnLimitReached rule", () => {
it("R1: at limit, user confirms → counter resets to 0", async () => {
const { pi, handlers, commands } = createMockPi();
initExtension(pi);
// Set max turns to 3 for faster testing
const { ctx } = createMockCtx({ confirmResult: true });
await commands["turn-limit"].handler("3", ctx);
// Fire 3 turns — third should trigger confirmation
await handlers["turn_start"]({}, ctx);
await handlers["turn_start"]({}, ctx);
await handlers["turn_start"]({}, ctx);
// Confirm was called
assert.equal(ctx.ui.confirm.mock.callCount(), 1);
// Counter reset — widget shows 0/3
const lastWidget = ctx.ui.setWidget.mock.calls.at(-1);
assert.deepEqual(lastWidget!.arguments, ["turn-limit", ["Turns: 0/3"]]);
// abort was NOT called
assert.equal(ctx.abort.mock.callCount(), 0);
});
it("R2: at limit, user declines → session aborts", async () => {
const { pi, handlers, commands } = createMockPi();
initExtension(pi);
const { ctx } = createMockCtx({ confirmResult: false });
await commands["turn-limit"].handler("2", ctx);
await handlers["turn_start"]({}, ctx);
await handlers["turn_start"]({}, ctx);
assert.equal(ctx.ui.confirm.mock.callCount(), 1);
assert.equal(ctx.abort.mock.callCount(), 1);
// Notify about abort
const notifyCalls = ctx.ui.notify.mock.calls.filter(
(c) => c.arguments[0] === "Agent aborted by user."
);
assert.equal(notifyCalls.length, 1);
});
it("R3: at limit, no UI → silent abort", async () => {
const { pi, handlers } = createMockPi();
initExtension(pi);
// Set max to 1 via env before init — but we already initialized.
// Instead, use registerCommand with a UI ctx first, then turn_start with no UI.
const uiCtx = createMockCtx().ctx;
await handlers["agent_start"]({}, uiCtx);
// Use a separate noUI ctx for the turn
const { ctx: noUiCtx } = createMockCtx({ hasUI: false });
// We need max=1 — use the command to set it
// But the command uses ctx.ui.setWidget which won't work without UI...
// The command still sets maxTurns regardless
const cmdCtx = createMockCtx().ctx;
const { commands } = createMockPi();
// Re-init to get fresh state with command access
const pi2Mock = createMockPi();
initExtension(pi2Mock.pi);
const setCmdCtx = createMockCtx().ctx;
await pi2Mock.commands["turn-limit"].handler("1", setCmdCtx);
const noUi = createMockCtx({ hasUI: false }).ctx;
await pi2Mock.handlers["turn_start"]({}, noUi);
assert.equal(noUi.abort.mock.callCount(), 1);
// confirm should NOT have been called (no UI)
assert.equal(noUi.ui.confirm.mock.callCount(), 0);
});
it("R4: below limit → no prompt, no abort", async () => {
const { pi, handlers } = createMockPi();
initExtension(pi);
const { ctx } = createMockCtx();
await handlers["agent_start"]({}, ctx);
await handlers["turn_start"]({}, ctx);
assert.equal(ctx.ui.confirm.mock.callCount(), 0);
assert.equal(ctx.abort.mock.callCount(), 0);
});
it("R5: after reset, counting continues from 0", async () => {
const { pi, handlers, commands } = createMockPi();
initExtension(pi);
const { ctx } = createMockCtx({ confirmResult: true });
await commands["turn-limit"].handler("2", ctx);
// First round: 2 turns → confirm → reset
await handlers["turn_start"]({}, ctx);
await handlers["turn_start"]({}, ctx);
assert.equal(ctx.ui.confirm.mock.callCount(), 1);
// Second round: 2 more turns → confirm again
await handlers["turn_start"]({}, ctx);
await handlers["turn_start"]({}, ctx);
assert.equal(ctx.ui.confirm.mock.callCount(), 2);
// Counter reset again
const lastWidget = ctx.ui.setWidget.mock.calls.at(-1);
assert.deepEqual(lastWidget!.arguments, ["turn-limit", ["Turns: 0/2"]]);
});
});
// ============================================================================
// C6-C8: Command tests
// ============================================================================
describe("extension handler — turn-limit command", () => {
it("C6: valid arg updates maxTurns", async () => {
const { pi, commands } = createMockPi();
initExtension(pi);
const { ctx } = createMockCtx();
await commands["turn-limit"].handler("10", ctx);
const notifyCalls = ctx.ui.notify.mock.calls.filter(
(c) => c.arguments[0] === "Turn limit set to 10."
);
assert.equal(notifyCalls.length, 1);
// Widget updated with new max
const widgetCall = ctx.ui.setWidget.mock.calls.at(-1);
assert.deepEqual(widgetCall!.arguments, ["turn-limit", ["Turns: 0/10"]]);
});
it("C7: invalid arg shows error", async () => {
const { pi, commands } = createMockPi();
initExtension(pi);
const { ctx } = createMockCtx();
await commands["turn-limit"].handler("abc", ctx);
const errorCalls = ctx.ui.notify.mock.calls.filter(
(c) => c.arguments[1] === "error"
);
assert.equal(errorCalls.length, 1);
});
it("C8: empty args shows error", async () => {
const { pi, commands } = createMockPi();
initExtension(pi);
const { ctx } = createMockCtx();
await commands["turn-limit"].handler("", ctx);
const errorCalls = ctx.ui.notify.mock.calls.filter(
(c) => c.arguments[1] === "error"
);
assert.equal(errorCalls.length, 1);
});
});
// ============================================================================
// Unlimited / disable feature tests
// ============================================================================
describe("unlimited mode — config", () => {
let originalEnv: string | undefined;
beforeEach(() => {
originalEnv = process.env.PI_MAX_TURNS;
});
afterEach(() => {
if (originalEnv === undefined) {
delete process.env.PI_MAX_TURNS;
} else {
process.env.PI_MAX_TURNS = originalEnv;
}
});
it("CFG-UNLIM-1: getMaxTurns returns Infinity for PI_MAX_TURNS=unlimited", () => {
process.env.PI_MAX_TURNS = "unlimited";
assert.equal(getMaxTurns(), Infinity);
});
});
describe("unlimited mode — command", () => {
it("CMD-UNLIM-1: 'turn-limit unlimited' is accepted", async () => {
const { pi, commands } = createMockPi();
initExtension(pi);
const { ctx } = createMockCtx();
await commands["turn-limit"].handler("unlimited", ctx);
// Should NOT show error
const errorCalls = ctx.ui.notify.mock.calls.filter(
(c) => c.arguments[1] === "error"
);
assert.equal(errorCalls.length, 0);
});
it("CMD-UNLIM-2: 'turn-limit unlimited' notifies user", async () => {
const { pi, commands } = createMockPi();
initExtension(pi);
const { ctx } = createMockCtx();
await commands["turn-limit"].handler("unlimited", ctx);
const infoCalls = ctx.ui.notify.mock.calls.filter(
(c) => c.arguments[1] === "info"
);
assert.equal(infoCalls.length, 1);
assert.match(infoCalls[0].arguments[0] as string, /unlimited/i);
});
it("CMD-UNLIM-3: after 'turn-limit unlimited', widget shows ∞", async () => {
const { pi, commands } = createMockPi();
initExtension(pi);
const { ctx } = createMockCtx();
await commands["turn-limit"].handler("unlimited", ctx);
const lastWidget = ctx.ui.setWidget.mock.calls.at(-1);
assert.match((lastWidget!.arguments[1] as string[])[0], /∞/);
});
});
describe("unlimited mode — no boundary check fires", () => {
it("RUL-UNLIM-2: unlimited mode — no confirmation or abort after many turns", async () => {
const { pi, handlers, commands } = createMockPi();
initExtension(pi);
const { ctx } = createMockCtx();
await commands["turn-limit"].handler("unlimited", ctx);
await handlers["agent_start"]({}, ctx);
// Fire 50 turns — none should trigger confirmation or abort
for (let i = 0; i < 50; i++) {
await handlers["turn_start"]({}, ctx);
}
assert.equal(ctx.ui.confirm.mock.callCount(), 0);
assert.equal(ctx.abort.mock.callCount(), 0);
});
it("RUL-UNLIM-3: unlimited mode — counter still increments", async () => {
const { pi, handlers, commands } = createMockPi();
initExtension(pi);
const { ctx } = createMockCtx();
await commands["turn-limit"].handler("unlimited", ctx);
await handlers["agent_start"]({}, ctx);
await handlers["turn_start"]({}, ctx);
await handlers["turn_start"]({}, ctx);
await handlers["turn_start"]({}, ctx);
// Widget should show counter incrementing with ∞
const lastWidget = ctx.ui.setWidget.mock.calls.at(-1);
assert.match((lastWidget!.arguments[1] as string[])[0], /3/);
assert.match((lastWidget!.arguments[1] as string[])[0], /∞/);
});
});
describe("unlimited mode — switching back to limited (LimitReEnabled)", () => {
it("CMD-INT-1: switching from unlimited to number resets counter", async () => {
const { pi, handlers, commands } = createMockPi();
initExtension(pi);
const { ctx } = createMockCtx();
// Set unlimited
await commands["turn-limit"].handler("unlimited", ctx);
await handlers["agent_start"]({}, ctx);
// Do several turns
for (let i = 0; i < 10; i++) {
await handlers["turn_start"]({}, ctx);
}
// Switch back to limited
await commands["turn-limit"].handler("5", ctx);
// Counter should be reset to 0, widget shows 0/5
const lastWidget = ctx.ui.setWidget.mock.calls.at(-1);
assert.deepEqual(lastWidget!.arguments, ["turn-limit", ["Turns: 0/5"]]);
});
it("INT-1: unlimited → switch to 3 → boundary fires at turn 3", async () => {
const { pi, handlers, commands } = createMockPi();
initExtension(pi);
const { ctx } = createMockCtx({ confirmResult: true });
// Set unlimited, do turns
await commands["turn-limit"].handler("unlimited", ctx);
await handlers["agent_start"]({}, ctx);
for (let i = 0; i < 10; i++) {
await handlers["turn_start"]({}, ctx);
}
assert.equal(ctx.ui.confirm.mock.callCount(), 0);
// Switch to limit=3 → counter resets
await commands["turn-limit"].handler("3", ctx);
// Now 3 turns should trigger confirmation
await handlers["turn_start"]({}, ctx);
await handlers["turn_start"]({}, ctx);
await handlers["turn_start"]({}, ctx);
assert.equal(ctx.ui.confirm.mock.callCount(), 1);
});
});

View File

@ -6,13 +6,18 @@ import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
const DEFAULT_MAX_TURNS = 25;
function getMaxTurns(): number {
export function getMaxTurns(): number {
const env = process.env.PI_MAX_TURNS;
if (!env) return DEFAULT_MAX_TURNS;
if (env.trim().toLowerCase() === "unlimited") return Infinity;
const parsed = parseInt(env, 10);
return Number.isFinite(parsed) && parsed > 0 ? parsed : DEFAULT_MAX_TURNS;
}
function formatMax(maxTurns: number): string {
return maxTurns === Infinity ? "∞" : String(maxTurns);
}
// ============================================================================
// Pure detection logic (testable)
// ============================================================================
@ -39,7 +44,7 @@ export default function (pi: ExtensionAPI) {
pi.on("session_start", async (event, ctx) => {
// On reload, show the widget immediately
if (event.reason === "reload" && ctx.hasUI) {
ctx.ui.setWidget("turn-limit", [`Turns: ${turnCount}/${maxTurns}`]);
ctx.ui.setWidget("turn-limit", [`Turns: ${turnCount}/${formatMax(maxTurns)}`]);
}
});
@ -48,7 +53,7 @@ export default function (pi: ExtensionAPI) {
turnCount = 0;
// Show initial widget state on fresh session
if (ctx.hasUI) {
ctx.ui.setWidget("turn-limit", [`Turns: ${turnCount}/${maxTurns}`]);
ctx.ui.setWidget("turn-limit", [`Turns: ${turnCount}/${formatMax(maxTurns)}`]);
}
});
@ -57,16 +62,26 @@ export default function (pi: ExtensionAPI) {
handler: async (args: string, ctx) => {
const trimmed = args.trim();
if (!trimmed) {
ctx.ui.notify("Invalid turn limit. Must be a positive integer.", "error");
ctx.ui.notify("Invalid turn limit. Must be a positive integer or 'unlimited'.", "error");
return;
}
if (trimmed.toLowerCase() === "unlimited") {
maxTurns = Infinity;
ctx.ui.setWidget("turn-limit", [`Turns: ${turnCount}/${formatMax(maxTurns)}`]);
ctx.ui.notify("Turn limit set to unlimited.", "info");
return;
}
const parsed = parseInt(trimmed, 10);
if (!Number.isFinite(parsed) || parsed <= 0) {
ctx.ui.notify("Invalid turn limit. Must be a positive integer.", "error");
ctx.ui.notify("Invalid turn limit. Must be a positive integer or 'unlimited'.", "error");
return;
}
const wasUnlimited = maxTurns === Infinity;
maxTurns = parsed;
ctx.ui.setWidget("turn-limit", [`Turns: ${turnCount}/${maxTurns}`]);
if (wasUnlimited) {
turnCount = 0;
}
ctx.ui.setWidget("turn-limit", [`Turns: ${turnCount}/${formatMax(maxTurns)}`]);
ctx.ui.notify(`Turn limit set to ${parsed}.`, "info");
},
});
@ -76,9 +91,12 @@ export default function (pi: ExtensionAPI) {
// Update live widget
if (ctx.hasUI) {
ctx.ui.setWidget("turn-limit", [`Turns: ${turnCount}/${maxTurns}`]);
ctx.ui.setWidget("turn-limit", [`Turns: ${turnCount}/${formatMax(maxTurns)}`]);
}
// No boundary check when unlimited
if (maxTurns === Infinity) return;
// Boundary confirmation: when we hit maxTurns exactly
if (turnCount === maxTurns) {
if (ctx.hasUI) {
@ -90,7 +108,7 @@ export default function (pi: ExtensionAPI) {
// Reset counter and let the turn proceed
turnCount = 0;
if (ctx.hasUI) {
ctx.ui.setWidget("turn-limit", [`Turns: ${turnCount}/${maxTurns}`]);
ctx.ui.setWidget("turn-limit", [`Turns: ${turnCount}/${formatMax(maxTurns)}`]);
}
return;
} else {

View File

@ -0,0 +1,81 @@
-- allium: 3
-- turn-limit.allium
-- Scope: Agent turn limit enforcement per session
-- Includes: Turn counting, limit enforcement, session abort, disable/enable
-- Excludes:
-- - Widget display (UI implementation detail)
-- - Environment variable reading (configuration mechanism)
-- - turn-limit command (configuration mechanism)
------------------------------------------------------------
-- Entities
------------------------------------------------------------
entity Session {
turn_count: Integer
status: active | aborted
user_confirms_continuation: Boolean?
transitions status {
active -> aborted
terminal: aborted
}
}
------------------------------------------------------------
-- Config
------------------------------------------------------------
config {
max_turns: Integer | unlimited = 25
@guidance
-- When max_turns is unlimited, no boundary check fires.
-- The turn counter still increments for observability.
-- Transitioning from unlimited to a positive integer
-- resets turn_count to 0.
}
------------------------------------------------------------
-- Rules
------------------------------------------------------------
rule TurnLimitReached {
when: session: Session.turn_count transitions_to config.max_turns
requires:
config.max_turns != unlimited
session.turn_count = config.max_turns
ensures:
if session.user_confirms_continuation:
session.turn_count = 0
else:
session.status = aborted
@guidance
-- The user_confirms_continuation field is set by the
-- implementation when presenting a confirmation prompt to
-- the user at the turn limit boundary. The implementation
-- may use a widget, dialog, or other mechanism to capture
-- the user's choice.
--
-- When the user confirms, the turn_count resets to zero,
-- allowing the agent to continue for another round of turns.
-- When the user declines (field is null or false), the
-- session is aborted.
--
-- Without a UI, the default behaviour is to abort.
}
rule LimitReEnabled {
when: config.max_turns transitions_to Integer
ensures:
session.turn_count = 0
@guidance
-- When the user switches from unlimited back to a positive
-- integer limit, the turn counter resets to zero so the
-- new limit applies from a clean starting point.
}