fix(agent-invocations): End stranded agent invocations from the heartbeat

packages/junior/src/chat/agent-dispatch/heartbeat.ts

@@ -7,13 +7,15 @@ import type { ConversationWorkQueue } from "@/chat/task-execution/queue";
 import { getVercelConversationWorkQueue } from "@/chat/task-execution/vercel-queue";
 import { createHeartbeatContext } from "./context";
 import {
+  listAppendedOpenAgentInvocations,
   listPendingAgentInvocationMailboxAppends,
   listPendingAgentInvocationResults,
 } from "@/chat/agent-invocations/store";
 import { enqueueAgentInvocationResult } from "@/chat/agent-invocations/result";
 import { recordSubagentEnded } from "@/chat/agent-invocations/markers";
+import type { AgentInvocation } from "@/chat/agent-invocations/types";
 import {
-  endAgentInvocationsOfStoppedWork,
+  endStrandedAgentInvocation,
   enqueueAgentInvocation,
 } from "@/chat/agent-invocations/work";
 import {
@@ -160,6 +162,34 @@ export async function recoverPendingAgentInvocationMailboxAppends(args: {
   }
 }
 
+/**
+ * End open agent invocations whose child work can no longer run them. It reads
+ * every open agent invocation, because healthy ones can stay open for a long
+ * time and would otherwise hide the stranded ones.
+ */
+export async function recoverStrandedAgentInvocations(): Promise<void> {
+  let after: AgentInvocation | undefined;
+  for (;;) {
+    const page = await listAppendedOpenAgentInvocations({
+      after,
+      limit: AGENT_INVOCATION_MAILBOX_APPEND_LIMIT,
+    });
+    for (const invocation of page) {
+      try {
+        await endStrandedAgentInvocation(invocation);
+      } catch (error) {
+        logException(error, "agent.invocation.stranded_recovery.failed", {
+          "app.agent.invocation_id": invocation.invocationId,
+        });
+      }
+    }
+    if (page.length < AGENT_INVOCATION_MAILBOX_APPEND_LIMIT) {
+      return;
+    }
+    after = page.at(-1);
+  }
+}
+
 /** Repair terminal agent invocation results that never reached their parent. */
 export async function recoverPendingAgentInvocationResults(args: {
   conversationWorkQueue: ConversationWorkQueue;
@@ -189,10 +219,7 @@ export async function runHeartbeat(args: {
   nowMs: number;
 }): Promise<void> {
   const queue = args.conversationWorkQueue ?? getVercelConversationWorkQueue();
-  // Runs before result recovery, so an invocation that this ends is sent to
-  // its parent in the same heartbeat.
   await recoverConversationWork({
-    beforeWorkStops: endAgentInvocationsOfStoppedWork,
     nowMs: args.nowMs,
     queue,
   });
@@ -204,6 +231,9 @@ export async function runHeartbeat(args: {
     conversationWorkQueue: queue,
     nowMs: args.nowMs,
   });
+  // Runs after the phases that stop work or append input, and before result
+  // recovery, so an agent invocation that this ends reports in this heartbeat.
+  await recoverStrandedAgentInvocations();
   await recoverPendingAgentInvocationResults({
     conversationWorkQueue: queue,
     nowMs: args.nowMs,

packages/junior/src/chat/agent-invocations/store.ts

@@ -6,6 +6,7 @@ import {
   countDistinct,
   desc,
   eq,
+  gt,
   inArray,
   lt,
   ne,
@@ -51,11 +52,6 @@ const TERMINAL_AGENT_INVOCATION_STATUSES = [
   "failed",
   "stopped",
 ] as const satisfies readonly AgentInvocationStatus[];
-/** Statuses of an agent invocation that owns its child's Turn. */
-export const ACTIVE_AGENT_INVOCATION_STATUSES = [
-  "running",
-  "paused",
-] as const satisfies readonly AgentInvocationStatus[];
 /** Statuses of an agent invocation that has not finished. */
 export const NON_TERMINAL_AGENT_INVOCATION_STATUSES = [
   "pending",
@@ -233,10 +229,9 @@ export async function getAgentInvocation(
   return rows[0] ? invocationFromRow(rows[0]) : undefined;
 }
 
-/** List the agent invocations of one child Conversation in the given statuses. */
-export async function listAgentInvocationsForConversation(
+/** List the running and paused agent invocations of one child Conversation. */
+export async function listActiveAgentInvocationsForConversation(
   childConversationId: string,
-  statuses: readonly AgentInvocationStatus[],
 ): Promise<AgentInvocation[]> {
   const rows = await getSqlExecutor()
     .db()
@@ -245,7 +240,7 @@ export async function listAgentInvocationsForConversation(
     .where(
       and(
         eq(juniorAgentInvocations.childConversationId, childConversationId),
-        inArray(juniorAgentInvocations.status, [...statuses]),
+        inArray(juniorAgentInvocations.status, ["running", "paused"]),
       ),
     )
     .orderBy(asc(juniorAgentInvocations.createdAt));
@@ -256,10 +251,8 @@ export async function listAgentInvocationsForConversation(
 export async function getActiveAgentInvocationForConversation(
   childConversationId: string,
 ): Promise<AgentInvocation | undefined> {
-  const active = await listAgentInvocationsForConversation(
-    childConversationId,
-    ACTIVE_AGENT_INVOCATION_STATUSES,
-  );
+  const active =
+    await listActiveAgentInvocationsForConversation(childConversationId);
   if (active.length > 1) {
     throw new Error(
       `Child conversation ${childConversationId} has multiple active agent invocations`,
@@ -757,6 +750,48 @@ export async function listPendingAgentInvocationMailboxAppends(
   return rows.map(invocationFromRow);
 }
 
+/**
+ * List open agent invocations whose input reached the child mailbox, oldest
+ * first. Pass the last agent invocation of a page to read the next page.
+ */
+export async function listAppendedOpenAgentInvocations(args: {
+  after?: AgentInvocation;
+  limit: number;
+}): Promise<AgentInvocation[]> {
+  const { after } = args;
+  const rows = await getSqlExecutor()
+    .db()
+    .select()
+    .from(juniorAgentInvocations)
+    .where(
+      and(
+        eq(juniorAgentInvocations.mailboxStatus, "appended"),
+        inArray(
+          juniorAgentInvocations.status,
+          NON_TERMINAL_AGENT_INVOCATION_STATUSES,
+        ),
+        after
+          ? or(
+              gt(juniorAgentInvocations.createdAt, new Date(after.createdAtMs)),
+              and(
+                eq(
+                  juniorAgentInvocations.createdAt,
+                  new Date(after.createdAtMs),
+                ),
+                gt(juniorAgentInvocations.invocationId, after.invocationId),
+              ),
+            )
+          : undefined,
+      ),
+    )
+    .orderBy(
+      asc(juniorAgentInvocations.createdAt),
+      asc(juniorAgentInvocations.invocationId),
+    )
+    .limit(args.limit);
+  return rows.map(invocationFromRow);
+}
+
 /** List terminal agent invocation results that still need a parent mailbox append. */
 export async function listPendingAgentInvocationResults(
   limit = 100,

packages/junior/src/chat/agent-invocations/work.ts

@@ -60,6 +60,7 @@ import { loadForkHistory } from "./fork-history";
 import type { AgentRunResult } from "@/chat/services/turn-result";
 import {
   appendAndEnqueueInboundMessage,
+  getConversationWorkState,
   type InboundMessage,
 } from "@/chat/task-execution/store";
 import type { ConversationWorkQueue } from "@/chat/task-execution/queue";
@@ -76,9 +77,7 @@ import {
   getAgentInvocationTurnId,
   getLatestAgentInvocationForConversation,
   isTerminalAgentInvocation,
-  ACTIVE_AGENT_INVOCATION_STATUSES,
-  listAgentInvocationsForConversation,
-  NON_TERMINAL_AGENT_INVOCATION_STATUSES,
+  listActiveAgentInvocationsForConversation,
   markAgentInvocationPaused,
   markAgentInvocationMailboxAppended,
   markAgentInvocationRunning,
@@ -500,38 +499,55 @@ async function endOtherActiveAgentInvocation(
 }
 
 /**
- * End every agent invocation of a child whose last attempt lost its worker and
- * whose input reached the child's mailbox. The worker ran no failure path, and
- * the stop drops the mailbox, so none of them can run. The heartbeat sends
- * each result.
+ * End an open agent invocation that can no longer run. Its input left the
+ * child mailbox, and no worker holds the child's lease. A worker that dies, or
+ * that fails past its last retry, leaves this state and runs no failure path.
+ * Each consumed id stays in `inboundMessageIds`, so the input cannot be
+ * appended again. The heartbeat sends the result.
  */
-export async function endAgentInvocationsOfStoppedWork(
-  conversationId: string,
-  inboundMessageIds: readonly string[],
+export async function endStrandedAgentInvocation(
+  listed: AgentInvocation,
 ): Promise<void> {
-  for (const invocation of await listAgentInvocationsForConversation(
-    conversationId,
-    NON_TERMINAL_AGENT_INVOCATION_STATUSES,
-  )) {
-    // Its append runs after the stop and starts new work for the child.
-    if (
-      !inboundMessageIds.includes(
-        getAgentInvocationMessageId(invocation.invocationId),
-      )
-    ) {
-      continue;
-    }
-    const terminal =
-      (await projectTerminalTurn(invocation, { reportNow: true })) ??
-      (await completeAgentInvocation({
-        invocationId: invocation.invocationId,
-        errorMessage: "This task stopped because its worker failed repeatedly.",
-        reportNow: true,
-        status: "failed",
-      }));
-    if (terminal && isTerminalAgentInvocation(terminal)) {
-      await persistTerminalLifecycle(terminal);
-    }
+  const work = await getConversationWorkState({
+    conversationId: listed.childConversationId,
+  });
+  const messageId = getAgentInvocationMessageId(listed.invocationId);
+  if (
+    !work ||
+    work.execution.lease ||
+    !work.execution.inboundMessageIds.includes(messageId) ||
+    work.execution.pendingMessages.some(
+      (message) => message.inboundMessageId === messageId,
+    )
+  ) {
+    return;
+  }
+  // Read after the work state: a worker can start it after the listing.
+  const invocation = await getAgentInvocation(listed.invocationId);
+  if (!invocation || isTerminalAgentInvocation(invocation)) {
+    return;
+  }
+  // Only a worker that holds the lease starts a pending agent invocation, so
+  // it can never start. A running or paused one keeps its Turn while the
+  // child can still resume it. The next agent invocation of the child ends
+  // it, or this ends it when the child has no work left.
+  if (
+    invocation.status !== "pending" &&
+    work.execution.status !== "failed" &&
+    work.execution.status !== "idle"
+  ) {
+    return;
+  }
+  const terminal =
+    (await projectTerminalTurn(invocation, { reportNow: true })) ??
+    (await completeAgentInvocation({
+      invocationId: invocation.invocationId,
+      errorMessage: "This task stopped because its worker failed repeatedly.",
+      reportNow: true,
+      status: "failed",
+    }));
+  if (terminal && isTerminalAgentInvocation(terminal)) {
+    await persistTerminalLifecycle(terminal);
   }
 }
 
@@ -754,9 +770,8 @@ export function createAgentInvocationWorker(
       }
       // A child has at most one running or paused agent invocation. End any
       // other active one first. Each one reports its own result.
-      for (const active of await listAgentInvocationsForConversation(
+      for (const active of await listActiveAgentInvocationsForConversation(
         invocation.childConversationId,
-        ACTIVE_AGENT_INVOCATION_STATUSES,
       )) {
         if (active.invocationId !== invocation.invocationId) {
           await finish(await endOtherActiveAgentInvocation(active));

packages/junior/src/chat/task-execution/heartbeat.ts

@@ -31,11 +31,6 @@ function heartbeatIdempotencyKey(
 
 /** Requeue expired leases and stranded mailbox work without running the agent. */
 export async function recoverConversationWork(args: {
-  /** Ends what a conversation owns when its last attempt loses its worker. */
-  beforeWorkStops?: (
-    conversationId: string,
-    inboundMessageIds: readonly string[],
-  ) => Promise<void>;
   limit?: number;
   nowMs: number;
   queue: ConversationWorkQueue;
@@ -99,7 +94,6 @@ export async function recoverConversationWork(args: {
           work.execution.lease.expiresAtMs <= args.nowMs
         ) {
           const recovery = await clearExpiredConversationLease({
-            beforeWorkStops: args.beforeWorkStops,
             conversationId,
             nowMs: args.nowMs,
             state: args.state,

packages/junior/src/chat/task-execution/state.ts

@@ -2317,10 +2317,6 @@ export async function deadLetterAttempt(args: {
 
 /** Recover an expired lease, consuming the retry budget under the mutation lock. */
 export async function clearExpiredConversationLease(args: {
-  beforeWorkStops?: (
-    conversationId: string,
-    inboundMessageIds: readonly string[],
-  ) => Promise<void>;
   conversationId: string;
   nowMs?: number;
   state?: StateAdapter;
@@ -2336,16 +2332,6 @@ export async function clearExpiredConversationLease(args: {
     }
     const retryCount = (current.execution.retryCount ?? 0) + 1;
     const stopped = retryCount >= CONVERSATION_WORK_MAX_RETRIES;
-    if (stopped) {
-      // A lost worker runs no failure path. This runs under the lock, so the
-      // worker cannot renew the lease and no message can reach the mailbox
-      // after it. When it throws, the lease stays expired and the next
-      // heartbeat runs it again.
-      await args.beforeWorkStops?.(
-        args.conversationId,
-        current.execution.inboundMessageIds,
-      );
-    }
     await writeConversation(
       state,
       lock,

packages/junior/src/chat/task-execution/store.ts

@@ -620,10 +620,6 @@ export async function deadLetterAttempt(args: {
 
 /** Recover an expired lease and persist the resulting retry state. */
 export async function clearExpiredConversationLease(args: {
-  beforeWorkStops?: (
-    conversationId: string,
-    inboundMessageIds: readonly string[],
-  ) => Promise<void>;
   conversationId: string;
   conversationStore?: ConversationStore;
   nowMs?: number;