fix: recover omp resume and blocked states

refs #879

1 files changed, 102 insertions, 16 deletions

diff --git a/src/integration/assets/omp/herdr-agent-state.ts b/src/integration/assets/omp/herdr-agent-state.ts
index e0cfab5e..21c218b7 100644
--- a/src/integration/assets/omp/herdr-agent-state.ts
+++ b/src/integration/assets/omp/herdr-agent-state.ts
@@ -16,7 +16,9 @@ function enabled() {
   return HERDR_ENV === "1" && !!socketPath && !!paneId;
 }
 
-function sendRequest(request: unknown): Promise<void> {
+let requestQueue = Promise.resolve();
+
+function sendRequestNow(request: unknown): Promise<void> {
   if (!enabled()) {
     return Promise.resolve();
   }
@@ -40,6 +42,14 @@ function sendRequest(request: unknown): Promise<void> {
   });
 }
 
+function sendRequest(request: unknown): Promise<void> {
+  requestQueue = requestQueue.then(
+    () => sendRequestNow(request),
+    () => sendRequestNow(request),
+  );
+  return requestQueue;
+}
+
 type AgentState = "working" | "blocked" | "idle";
 
 type QueuedState = {
@@ -110,7 +120,7 @@ function currentSessionRef(): Record<string, unknown> | undefined {
   return undefined;
 }
 
-function reportSession(): Promise<void> {
+function reportSession(sessionStartSource = "startup"): Promise<void> {
   const sessionRef = currentSessionRef();
   if (!sessionRef) {
     return Promise.resolve();
@@ -124,6 +134,7 @@ function reportSession(): Promise<void> {
       source,
       agent: "omp",
       seq: nextReportSeq(),
+      session_start_source: sessionStartSource,
       ...sessionRef,
     },
   });
@@ -221,6 +232,15 @@ function retryableErrorMessage(event: any): string | undefined {
   return errorMessage || "retryable provider error";
 }
 
+function askBlockedMessage(args: any): string {
+  const questions = Array.isArray(args?.questions) ? args.questions : [];
+  const firstQuestion = questions.find((question: any) => typeof question?.question === "string");
+  if (firstQuestion?.question) {
+    return firstQuestion.question;
+  }
+  return "waiting for user input";
+}
+
 export default function (pi) {
   if (!enabled()) {
     return;
@@ -306,37 +326,68 @@ export default function (pi) {
     retryTimer.unref?.();
   }
 
+  function activateRootSession(ctx: any, sessionStartSource = "startup"): boolean {
+    if (ctx?.hasUI !== true) {
+      return false;
+    }
+    rootSession = true;
+    updateSessionRef(ctx);
+    void reportSession(sessionStartSource);
+    return true;
+  }
+
+  function resetSessionState() {
+    clearPendingTimers();
+    clearFailureState();
+    agentActive = false;
+    blockedCount = 0;
+    blockedMessage = undefined;
+  }
+
+  function activateBlocked(message: string | undefined) {
+    clearPendingTimers();
+    blockedCount += 1;
+    blockedMessage = message;
+    publishState();
+  }
+
+  function deactivateBlocked() {
+    blockedCount = Math.max(0, blockedCount - 1);
+    if (blockedCount === 0) {
+      blockedMessage = undefined;
+    }
+    publishState();
+  }
+
   pi.events.on("herdr:blocked", (data) => {
     if (!rootSession) {
       return;
     }
     if (!data?.active) {
-      blockedCount = Math.max(0, blockedCount - 1);
-      if (blockedCount === 0) {
-        blockedMessage = undefined;
-      }
-      publishState();
+      deactivateBlocked();
       return;
     }
 
-    clearPendingTimers();
-    blockedCount += 1;
-    blockedMessage = data.label;
-    publishState();
+    activateBlocked(data.label);
   });
 
   pi.on("session_start", (_event, ctx) => {
-    if (ctx?.hasUI !== true) {
+    if (!activateRootSession(ctx)) {
       return;
     }
-    rootSession = true;
-    updateSessionRef(ctx);
-    void reportSession();
+    publishState(true);
+  });
+
+  pi.on("session_switch", (event, ctx) => {
+    if (!activateRootSession(ctx, event?.reason || "resume")) {
+      return;
+    }
+    resetSessionState();
     publishState(true);
   });
 
   pi.on("agent_start", (_event, ctx) => {
-    if (!rootSession) {
+    if (!rootSession && !activateRootSession(ctx)) {
       return;
     }
     updateSessionRef(ctx);
@@ -347,6 +398,41 @@ export default function (pi) {
     publishState();
   });
 
+  pi.on("tool_approval_requested", (event, ctx) => {
+    if (!rootSession && !activateRootSession(ctx)) {
+      return;
+    }
+    const label = event?.reason || `${event?.toolName || "Tool"} approval`;
+    activateBlocked(label);
+  });
+
+  pi.on("tool_approval_resolved", (_event, ctx) => {
+    if (!rootSession && !activateRootSession(ctx)) {
+      return;
+    }
+    deactivateBlocked();
+  });
+
+  pi.on("tool_execution_start", (event, ctx) => {
+    if (event?.toolName !== "ask") {
+      return;
+    }
+    if (!rootSession && !activateRootSession(ctx)) {
+      return;
+    }
+    activateBlocked(askBlockedMessage(event.args));
+  });
+
+  pi.on("tool_execution_end", (event, ctx) => {
+    if (event?.toolName !== "ask") {
+      return;
+    }
+    if (!rootSession && !activateRootSession(ctx)) {
+      return;
+    }
+    deactivateBlocked();
+  });
+
   pi.on("agent_end", (event) => {
     if (!rootSession) {
       return;