1 files changed, 38 insertions, 9 deletions

diff --git a/crates/dinggerd/assets/pi/pi-agent-state.ts b/crates/dinggerd/assets/pi/pi-agent-state.ts
index fc70638..62f29db 100644
--- a/crates/dinggerd/assets/pi/pi-agent-state.ts
+++ b/crates/dinggerd/assets/pi/pi-agent-state.ts
@@ -167,7 +167,7 @@ function requestId(): string {
   return `${source}:${Date.now()}:${Math.random().toString(36).slice(2)}`;
 }
 
-async function report(state: AgentState | undefined, message?: string): Promise<void> {
+async function report(state: AgentState | undefined, message?: string, aborted = false): Promise<void> {
   await sendRequest({
     id: requestId(),
     method: "pane.report_agent",
@@ -177,6 +177,8 @@ async function report(state: AgentState | undefined, message?: string): Promise<
       agent: "pi",
       ...(state ? { state } : {}),
       message,
+      // The run this idle report ends was cancelled: not done.
+      ...(aborted ? { aborted: true } : {}),
       importance,
       extension: extensionVersion,
       ...(controlListening ? { control: controlPath } : {}),
@@ -189,10 +191,10 @@ async function report(state: AgentState | undefined, message?: string): Promise<
 }
 
 let sendInFlight = false;
-let queuedState: { state: AgentState; message?: string } | undefined;
+let queuedState: { state: AgentState; message?: string; aborted: boolean } | undefined;
 
-function queueState(state: AgentState, message?: string): void {
-  queuedState = { state, message };
+function queueState(state: AgentState, message?: string, aborted = false): void {
+  queuedState = { state, message, aborted };
   if (!sendInFlight) {
     void drainStateQueue();
   }
@@ -208,7 +210,7 @@ async function drainStateQueue(): Promise<void> {
     while (queuedState) {
       const next = queuedState;
       queuedState = undefined;
-      await report(next.state, next.message);
+      await report(next.state, next.message, next.aborted);
     }
   } finally {
     sendInFlight = false;
@@ -239,16 +241,42 @@ export default function (pi) {
     return { state: "idle" as const, message: undefined };
   }
 
-  function publishState(force = false) {
+  function publishState(force = false, aborted = false) {
     const next = desiredState();
     if (!force && next.state === lastState && next.message === lastMessage) {
       return;
     }
     lastState = next.state;
     lastMessage = next.message;
-    queueState(next.state, next.message);
+    queueState(next.state, next.message, aborted && next.state === "idle");
   }
 
+  // A dialog that opens during a run (an extension asking before it goes
+  // on): the agent waits on you, so it is blocked until the dialog closes.
+  // One opened while idle comes from what you just did, and changes nothing.
+  let dialogsBlocking = 0;
+  pi.on("ui_prompt_start", (event, ctx) => {
+    latestCtx = ctx;
+    if (!agentActive) {
+      return;
+    }
+    dialogsBlocking += 1;
+    blockedCount += 1;
+    blockedMessage = event?.title;
+    publishState();
+  });
+  pi.on("ui_prompt_end", () => {
+    if (dialogsBlocking === 0) {
+      return;
+    }
+    dialogsBlocking -= 1;
+    blockedCount = Math.max(0, blockedCount - 1);
+    if (blockedCount === 0) {
+      blockedMessage = undefined;
+    }
+    publishState();
+  });
+
   pi.events.on("blocked", (data) => {
     if (!data?.active) {
       blockedCount = Math.max(0, blockedCount - 1);
@@ -516,7 +544,7 @@ export default function (pi) {
     });
   }
 
-  pi.on("agent_settled", (_event, ctx) => {
+  pi.on("agent_settled", (event, ctx) => {
     latestCtx = ctx;
     if (ctx?.isIdle?.() !== true) {
       return;
@@ -525,7 +553,8 @@ export default function (pi) {
     agentActive = false;
     // The last reply's usage: message_end comes before pi counts it.
     checkContext(ctx);
-    publishState();
+    // A cancelled run (Escape) is idle, not done.
+    publishState(false, event?.aborted === true);
     // Escape puts queued messages back into the input box.
     checkDraftSoon();
   });