diff --git a/packages/junior/src/chat/agent-dispatch/heartbeat.ts b/packages/junior/src/chat/agent-dispatch/heartbeat.ts index 736e37bd55..12f15048bc 100644 --- a/packages/junior/src/chat/agent-dispatch/heartbeat.ts +++ b/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 { + 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 { 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, diff --git a/packages/junior/src/chat/agent-invocations/store.ts b/packages/junior/src/chat/agent-invocations/store.ts index 06c9d64d32..300009edbb 100644 --- a/packages/junior/src/chat/agent-invocations/store.ts +++ b/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 { 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 { - 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 { + 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, diff --git a/packages/junior/src/chat/agent-invocations/work.ts b/packages/junior/src/chat/agent-invocations/work.ts index 8180cfd13f..8612c56b6a 100644 --- a/packages/junior/src/chat/agent-invocations/work.ts +++ b/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 { - 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)); diff --git a/packages/junior/src/chat/task-execution/heartbeat.ts b/packages/junior/src/chat/task-execution/heartbeat.ts index 744fcb01b3..aa17dbc367 100644 --- a/packages/junior/src/chat/task-execution/heartbeat.ts +++ b/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; 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, diff --git a/packages/junior/src/chat/task-execution/state.ts b/packages/junior/src/chat/task-execution/state.ts index 4ed46d4a84..e40ad44343 100644 --- a/packages/junior/src/chat/task-execution/state.ts +++ b/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; 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, diff --git a/packages/junior/src/chat/task-execution/store.ts b/packages/junior/src/chat/task-execution/store.ts index eb0f28f299..4a339a8c51 100644 --- a/packages/junior/src/chat/task-execution/store.ts +++ b/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; conversationId: string; conversationStore?: ConversationStore; nowMs?: number; diff --git a/packages/junior/tests/integration/agent-invocation-worker-crash.test.ts b/packages/junior/tests/integration/agent-invocation-worker-crash.test.ts index ac9aac162c..e3f2d7aef6 100644 --- a/packages/junior/tests/integration/agent-invocation-worker-crash.test.ts +++ b/packages/junior/tests/integration/agent-invocation-worker-crash.test.ts @@ -7,12 +7,14 @@ import { import { buildAgentInvocationInboundMessage, createAndEnqueueAgentInvocation, + enqueueAgentInvocation, } from "@/chat/agent-invocations/work"; import { disconnectStateAdapter } from "@/chat/state/adapter"; import { appendAndEnqueueInboundMessage, CONVERSATION_WORK_MAX_RETRIES, getConversationWorkState, + recordConversationRetry, } from "@/chat/task-execution/store"; import { getTurnRecord } from "@/chat/task-execution/turn-cursor"; import { neverRunAgentRunner } from "../fixtures/agent-runner"; @@ -50,7 +52,7 @@ describe("agent invocation worker crash", () => { const { childConversationId, invocationId } = created; let followUpId = ""; - let unappendedId = ""; + const statuses: (string | undefined)[] = []; for (let tries = 1; tries <= CONVERSATION_WORK_MAX_RETRIES; tries++) { // Injected fault: the worker of this try is killed mid-run. await killWorker(queue.takeMessageFor(childConversationId)); @@ -75,38 +77,16 @@ describe("agent invocation worker crash", () => { state, }); followUpId = invocation.invocationId; - // A task whose append has not run yet starts new work instead. - const { invocation: unappended } = - await createAgentInvocationForMessage({ - actor: invocationInput.actor, - childConversationId, - destination: invocationInput.destination, - idempotencyKey: "crash-unappended", - input: "Then check the docs.", - parentConversationId, - parentToolCallId: "call-message-2", - }); - unappendedId = unappended.invocationId; } await runHeartbeatAfterLeasesExpire(queue); + statuses.push((await getAgentInvocation(invocationId))?.status); } - const child = await getConversationWorkState({ - conversationId: childConversationId, - }); - expect(child?.execution.pendingMessages).toEqual([ - expect.objectContaining({ - input: expect.objectContaining({ - metadata: expect.objectContaining({ - agentInvocationId: unappendedId, - }), - }), - }), + // Each try before the last leaves it running for the next try. + expect(statuses).toEqual([ + ...Array(CONVERSATION_WORK_MAX_RETRIES - 1).fill("running"), + "failed", ]); - await expect(getAgentInvocation(unappendedId)).resolves.toMatchObject({ - mailboxStatus: "appended", - status: "pending", - }); for (const id of [invocationId, followUpId]) { await expect(getAgentInvocation(id)).resolves.toMatchObject({ parentMailboxStatus: "appended", @@ -136,4 +116,85 @@ describe("agent invocation worker crash", () => { await fixture.close(); } }); + + it("tells the parent about dropped work when a child worker fails past its last retry", async () => { + vi.useFakeTimers({ + now: 1_800_000_000_000, + toFake: ["Date", "setInterval", "clearInterval"], + }); + const { conversationStore, fixture } = await prepareParentConversation(); + const { killWorker, queue, state } = await createConversationWorkHarness( + neverRunAgentRunner(), + ); + try { + const { childConversationId, invocationId } = + await createAndEnqueueAgentInvocation( + { ...invocationInput, idempotencyKey: "retry-1" }, + { conversationStore, queue, state }, + ); + // Injected fault: the worker is killed mid-run, so it keeps its lease. + await killWorker(queue.takeMessageFor(childConversationId)); + // The parent gives the running child more work. + const { invocation: followUp } = await createAgentInvocationForMessage({ + actor: invocationInput.actor, + childConversationId, + destination: invocationInput.destination, + idempotencyKey: "retry-follow-up", + input: "Also check the retry path.", + parentConversationId, + parentToolCallId: "call-message", + }); + await enqueueAgentInvocation(followUp, { + conversationStore, + queue, + state, + }); + const leaseToken = ( + await getConversationWorkState({ conversationId: childConversationId }) + )?.execution.lease?.token; + // Injected fault: an error escapes the agent invocation worker on each + // try, so the worker records a retry until the work stops and drops its + // mailbox. No harness fault throws past the agent invocation worker. + let retry = "recorded"; + while (retry !== "stopped") { + retry = await recordConversationRetry({ + conversationId: childConversationId, + leaseToken: leaseToken!, + state, + }); + } + // The parent gives the stopped child new work before the heartbeat. + const { invocation: next } = await createAgentInvocationForMessage({ + actor: invocationInput.actor, + childConversationId, + destination: invocationInput.destination, + idempotencyKey: "retry-next", + input: "Now check the docs.", + parentConversationId, + parentToolCallId: "call-next", + }); + await enqueueAgentInvocation(next, { conversationStore, queue, state }); + await runHeartbeatAfterLeasesExpire(queue); + + // The stop dropped the follow-up, so it can never run. + await expect( + getAgentInvocation(followUp.invocationId), + ).resolves.toMatchObject({ + parentMailboxStatus: "appended", + status: "failed", + }); + // The child has new work, so the next agent invocation ends the first. + for (const id of [invocationId, next.invocationId]) { + await expect(getAgentInvocation(id)).resolves.not.toMatchObject({ + status: "failed", + }); + } + const parent = await getConversationWorkState({ + conversationId: parentConversationId, + }); + expect(parent?.messages).toHaveLength(1); + } finally { + await fixture.close(); + } + }); });