From 14282370ec580726b136630224bf5d1ac499f1b5 Mon Sep 17 00:00:00 2001
From: Willem van den Ende <willem@qwan.eu>
Date: Tue, 21 Apr 2026 12:28:13 +0100
Subject: [PATCH] Added live widget for turn limit

---
 packages/pi-turn-limit/src/turn-limit.ts | 57 +++++++++++++++++++++++-
 1 file changed, 56 insertions(+), 1 deletion(-)

diff --git a/packages/pi-turn-limit/src/turn-limit.ts b/packages/pi-turn-limit/src/turn-limit.ts
index e32584c..97bfb72 100644
--- a/packages/pi-turn-limit/src/turn-limit.ts
+++ b/packages/pi-turn-limit/src/turn-limit.ts
@@ -34,16 +34,71 @@ export function checkTurnLimit(
 
 export default function (pi: ExtensionAPI) {
   let turnCount = 0;
-  const maxTurns = getMaxTurns();
+  let maxTurns = getMaxTurns();
 
   pi.on("agent_start", async (_event, ctx) => {
     // Reset counter for each new user prompt
     turnCount = 0;
+    // Clear widget on new session
+    if (ctx.hasUI) {
+      ctx.ui.setWidget("turn-limit", undefined);
+    }
+  });
+
+  pi.registerCommand("turn-limit", {
+    description: "Set the maximum number of agent turns for this session",
+    handler: async (args: string, ctx) => {
+      const trimmed = args.trim();
+      if (!trimmed) {
+        ctx.ui.notify("Invalid turn limit. Must be a positive integer.", "error");
+        return;
+      }
+      const parsed = parseInt(trimmed, 10);
+      if (!Number.isFinite(parsed) || parsed <= 0) {
+        ctx.ui.notify("Invalid turn limit. Must be a positive integer.", "error");
+        return;
+      }
+      maxTurns = parsed;
+      ctx.ui.setWidget("turn-limit", [`Turns: ${turnCount}/${maxTurns}`]);
+      ctx.ui.notify(`Turn limit set to ${parsed}.`, "info");
+    },
   });
 
   pi.on("turn_start", async (event, ctx) => {
     turnCount++;
 
+    // Update live widget
+    if (ctx.hasUI) {
+      ctx.ui.setWidget("turn-limit", [`Turns: ${turnCount}/${maxTurns}`]);
+    }
+
+    // Boundary confirmation: when we hit maxTurns exactly
+    if (turnCount === maxTurns) {
+      if (ctx.hasUI) {
+        const confirmed = await ctx.ui.confirm(
+          "Turn limit reached",
+          `You've used ${maxTurns} turns. Continue?`,
+        );
+        if (confirmed) {
+          // Reset counter and let the turn proceed
+          turnCount = 0;
+          if (ctx.hasUI) {
+            ctx.ui.setWidget("turn-limit", [`Turns: ${turnCount}/${maxTurns}`]);
+          }
+          return;
+        } else {
+          // User declined — abort
+          ctx.ui.notify("Agent aborted by user.", "error");
+          ctx.abort();
+          return;
+        }
+      } else {
+        // No UI: abort silently (current behavior)
+        ctx.abort();
+        return;
+      }
+    }
+
     const { exceeded } = checkTurnLimit(event.turnIndex, maxTurns);
     if (exceeded && ctx.hasUI) {
       ctx.ui.notify(