From 6d0472cfc4008e832fd093d49ad2ccfa73607f6a Mon Sep 17 00:00:00 2001 From: David Cramer Date: Mon, 24 Aug 2026 19:08:55 -0700 Subject: [PATCH] ref(chat): add native Turn execution --- .../src/chat/agent-invocations/store.ts | 7 +- .../junior/src/chat/agent-invocations/work.ts | 205 ++++++++++-------- .../junior/src/chat/app/conversation-work.ts | 4 +- packages/junior/src/chat/runtime/README.md | 12 +- .../junior/src/chat/runtime/turn-execution.ts | 80 +++++++ packages/junior/src/cli/chat.ts | 4 +- .../component/agent-invocation-worker.test.ts | 10 +- .../tests/component/turn-execution.test.ts | 45 ++++ .../agent-invocation-concurrency.test.ts | 4 +- .../integration/agent-invocation-work.test.ts | 107 ++++++++- 10 files changed, 362 insertions(+), 116 deletions(-) create mode 100644 packages/junior/src/chat/runtime/turn-execution.ts create mode 100644 packages/junior/tests/component/turn-execution.test.ts diff --git a/packages/junior/src/chat/agent-invocations/store.ts b/packages/junior/src/chat/agent-invocations/store.ts index 3f893958bd..369db85298 100644 --- a/packages/junior/src/chat/agent-invocations/store.ts +++ b/packages/junior/src/chat/agent-invocations/store.ts @@ -429,10 +429,13 @@ export async function completeAgentInvocation( return await getAgentInvocation(args.invocationId); } -/** Return whether an invocation already owns its immutable terminal result. */ +/** Return whether an agent invocation has finished. */ export function isTerminalAgentInvocation( invocation: AgentInvocation, -): boolean { +): invocation is Extract< + AgentInvocation, + { status: "blocked" | "completed" | "failed" } +> { return TERMINAL_AGENT_INVOCATION_STATUSES.includes( invocation.status as (typeof TERMINAL_AGENT_INVOCATION_STATUSES)[number], ); diff --git a/packages/junior/src/chat/agent-invocations/work.ts b/packages/junior/src/chat/agent-invocations/work.ts index 3b0342f679..cbb88cdcd9 100644 --- a/packages/junior/src/chat/agent-invocations/work.ts +++ b/packages/junior/src/chat/agent-invocations/work.ts @@ -5,6 +5,7 @@ import { openConversationProjection } from "@/chat/conversations/projection"; import { ConversationTurnLifecycleService } from "@/chat/conversations/turn-lifecycle"; import { getConversationEventStore } from "@/chat/db"; import type { AgentRunner } from "@/chat/runtime/agent-runner"; +import { AgentRunError, executeTurn } from "@/chat/runtime/turn-execution"; import { getPersistedSandboxState, getPersistedThreadState, @@ -26,6 +27,7 @@ import { getAssistantReplyText } from "@/chat/services/assistant-reply"; import { getTerminalAssistantMessages } from "@/chat/pi/transcript"; import type { PiMessage } from "@/chat/pi/messages"; import type { SandboxRef } from "@/chat/sandbox/ref"; +import type { AgentRunResult } from "@/chat/services/turn-result"; import { appendAndEnqueueInboundMessage, type InboundMessage, @@ -258,10 +260,72 @@ function isInvocationInputCommitLost(error: unknown): boolean { return isTurnInputCommitLostError(cause); } -/** Build the invocation consumer that advances work through the shared runner. */ -export function createAgentInvocationWorker(options: { - agentRunner: AgentRunner; +/** Save one completed agent result on its child agent invocation. */ +async function saveAgentInvocationResult(args: { + invocation: AgentInvocation; + result: AgentRunResult; + sandboxRef?: SandboxRef; + turnId: string; }) { + const failed = args.result.diagnostics.outcome !== "success"; + await persistThreadStateById(args.invocation.childConversationId, { + sandboxRef: args.result.sandboxRef ?? args.sandboxRef, + }); + if (args.result.piMessages?.length) { + await saveTurnCheckpoint({ + mode: "completed", + conversationId: args.invocation.childConversationId, + turnId: args.turnId, + durationMs: args.result.diagnostics.durationMs, + usage: args.result.diagnostics.usage, + destination: args.invocation.destination, + destinationVisibility: args.invocation.destinationVisibility, + ...(failed + ? { + errorMessage: + args.result.diagnostics.errorMessage ?? "Agent invocation failed", + } + : undefined), + messages: args.result.piMessages, + actor: args.invocation.actor, + source: args.invocation.source, + surface: "internal", + }); + } + const terminal = await completeAgentInvocation({ + invocationId: args.invocation.invocationId, + ...(failed + ? { + errorMessage: + args.result.diagnostics.errorMessage ?? "Agent invocation failed", + status: "failed" as const, + } + : { + result: args.result.text, + status: "completed" as const, + }), + }); + if (!terminal || !isTerminalAgentInvocation(terminal)) { + throw new Error( + `Agent invocation did not finish for ${args.invocation.invocationId}`, + ); + } + return terminal.status === "completed" + ? { + finishedAtMs: terminal.terminalAtMs, + outcome: terminal.result.trim() + ? ("success" as const) + : ("no_reply" as const), + } + : { + finishedAtMs: terminal.terminalAtMs, + failureCode: "model_execution_failed" as const, + outcome: "failed" as const, + }; +} + +/** Build the invocation consumer that advances work through the shared runner. */ +export function createAgentInvocationWorker(agentRunner: AgentRunner) { return async ( context: ConversationWorkerContext, invocationId: string, @@ -354,46 +418,60 @@ export function createAgentInvocationWorker(options: { let outcome; try { - outcome = await options.agentRunner.run({ - conversationId: invocation.childConversationId, - turnId, - runId: invocation.invocationId, - instruction: { - text: invocation.input, - }, - history, - actor: invocation.actor, - credentialContext: invocation.credentialContext, - destination: invocation.destination, - destinationVisibility: invocation.destinationVisibility, - publishExternally: context.publishExternally, - source: invocation.source, - surface: "internal", - // TODO(#881, #883): Child runs may still need a path to force - // interactive auth when a delegated tool requires credentials the - // parent already has authority to request. Today background children - // hard-fail instead of pausing for an OAuth link. - disabledFeatures: ["handoff", "interactive-auth", "subagents"], - reasoning: invocation.reasoningLevel, - state: { - sandboxRef, - }, - durability: { - onInputCommitted: acknowledge, - shouldYield: context.shouldYield, - onSandboxRefChanged: async (nextSandboxRef) => { - sandboxRef = nextSandboxRef; - await persistThreadStateById(invocation.childConversationId, { - sandboxRef, - }); + outcome = await executeTurn( + agentRunner, + { + conversationId: invocation.childConversationId, + turnId, + runId: invocation.invocationId, + instruction: { + text: invocation.input, + }, + history, + actor: invocation.actor, + credentialContext: invocation.credentialContext, + destination: invocation.destination, + destinationVisibility: invocation.destinationVisibility, + publishExternally: context.publishExternally, + source: invocation.source, + surface: "internal", + // TODO(#881, #883): Child runs may still need a path to force + // interactive auth when a delegated tool requires credentials the + // parent already has authority to request. Today background children + // hard-fail instead of pausing for an OAuth link. + disabledFeatures: ["handoff", "interactive-auth", "subagents"], + reasoning: invocation.reasoningLevel, + state: { + sandboxRef, + }, + durability: { + onInputCommitted: acknowledge, + shouldYield: context.shouldYield, + onSandboxRefChanged: async (nextSandboxRef) => { + sandboxRef = nextSandboxRef; + await persistThreadStateById(invocation.childConversationId, { + sandboxRef, + }); + }, }, }, - }); + async (result) => + await saveAgentInvocationResult({ + invocation, + result, + sandboxRef, + turnId, + }), + ); } catch (error) { - if (isInvocationInputCommitLost(error)) { + if (!(error instanceof AgentRunError)) { + throw error; + } + const runError = error.cause; + if (isInvocationInputCommitLost(runError)) { return { status: "lost_lease" }; } - const blocking = blockingInvocationError(error); + const blocking = blockingInvocationError(runError); if (blocking) { const terminal = await completeAgentInvocation({ invocationId: invocation.invocationId, @@ -407,12 +485,14 @@ export function createAgentInvocationWorker(options: { return { status: "completed" }; } if (!context.attempt.isFinalAttempt) { - throw error; + throw runError; } const terminal = await completeAgentInvocation({ invocationId: invocation.invocationId, errorMessage: - error instanceof Error ? error.message : "Agent invocation failed", + runError instanceof Error + ? runError.message + : "Agent invocation failed", status: "failed", }); if (terminal) { @@ -439,51 +519,6 @@ export function createAgentInvocationWorker(options: { return { status: "completed" }; } - const result = outcome.result; - const failed = result.diagnostics.outcome !== "success"; - await persistThreadStateById(invocation.childConversationId, { - sandboxRef: result.sandboxRef ?? sandboxRef, - }); - if (result.piMessages?.length) { - await saveTurnCheckpoint({ - mode: "completed", - conversationId: invocation.childConversationId, - turnId, - durationMs: result.diagnostics.durationMs, - usage: result.diagnostics.usage, - destination: invocation.destination, - destinationVisibility: invocation.destinationVisibility, - ...(failed - ? { - errorMessage: - result.diagnostics.errorMessage ?? "Agent invocation failed", - } - : undefined), - messages: result.piMessages, - actor: invocation.actor, - source: invocation.source, - surface: "internal", - }); - } - const terminal = await completeAgentInvocation({ - invocationId: invocation.invocationId, - ...(failed - ? { - errorMessage: - result.diagnostics.errorMessage ?? "Agent invocation failed", - status: "failed" as const, - } - : { - result: result.text, - status: "completed" as const, - }), - }); - if (!terminal) { - throw new Error( - `Agent invocation disappeared during completion for ${invocation.invocationId}`, - ); - } - await persistTerminalLifecycle(terminal); await acknowledge(); return { status: "completed" }; }; diff --git a/packages/junior/src/chat/app/conversation-work.ts b/packages/junior/src/chat/app/conversation-work.ts index 6d5f776b9b..83943e7ceb 100644 --- a/packages/junior/src/chat/app/conversation-work.ts +++ b/packages/junior/src/chat/app/conversation-work.ts @@ -130,9 +130,7 @@ export function createConversationWork( turnLifecycle: services.replyExecutor?.turnLifecycle, }), fallbackWorker: routeAgentInvocationWork({ - invocationWorker: createAgentInvocationWorker({ - agentRunner: options.agentRunner, - }), + invocationWorker: createAgentInvocationWorker(options.agentRunner), fallbackWorker: providerWorker, }), }), diff --git a/packages/junior/src/chat/runtime/README.md b/packages/junior/src/chat/runtime/README.md index 747d9e0ba0..863d1f2297 100644 --- a/packages/junior/src/chat/runtime/README.md +++ b/packages/junior/src/chat/runtime/README.md @@ -1,8 +1,14 @@ # Chat Runtime -This folder coordinates a chat turn. It loads conversation state, calls the -agent, delivers replies, saves the result, and schedules more work when a turn -must continue later. `../agent/` owns the model and tool loop. +This folder owns native Turn execution and recovery. `../agent/` owns the model +and tool loop. Callers save input, deliver replies, and save their results. + +## Turn Execution + +`turn-execution.ts` advances a started Turn by one Run. The caller saves the +result it owns. Native execution finishes the Turn only after that save works. +A paused Run leaves the Turn open. If the save fails, the worker can retry or +recover the Turn. ## Replies diff --git a/packages/junior/src/chat/runtime/turn-execution.ts b/packages/junior/src/chat/runtime/turn-execution.ts new file mode 100644 index 0000000000..160016d005 --- /dev/null +++ b/packages/junior/src/chat/runtime/turn-execution.ts @@ -0,0 +1,80 @@ +import type { AgentRun } from "@/chat/agent/types"; +import type { + CompleteConversationTurnInput, + FailConversationTurnInput, +} from "@/chat/conversations/turn-lifecycle"; +import { ConversationTurnLifecycleService } from "@/chat/conversations/turn-lifecycle"; +import { getConversationEventStore } from "@/chat/db"; +import type { AgentRunOutcome } from "@/chat/runtime/agent-run-outcome"; +import type { AgentRunner } from "@/chat/runtime/agent-runner"; +import type { AgentRunResult } from "@/chat/services/turn-result"; + +type SavedTurnResult = + | { + finishedAtMs?: number; + outcome: CompleteConversationTurnInput["outcome"]; + } + | { + finishedAtMs?: number; + failureCode: FailConversationTurnInput["failureCode"]; + outcome: "failed"; + }; + +type TurnExecutionOutcome = + | Exclude + | { status: "completed" }; + +/** An error thrown while the agent advances a Run. */ +export class AgentRunError extends Error { + constructor(cause: unknown) { + super(cause instanceof Error ? cause.message : "Agent Run failed", { + cause, + }); + this.name = "AgentRunError"; + } +} + +/** + * Run the agent and finish the Turn after the caller saves its result. + * + * A paused Run, or a Run waiting for authorization, leaves the Turn open. A + * save error also leaves the Turn open so the worker can retry or recover it. + */ +export async function executeTurn( + agentRunner: AgentRunner, + run: AgentRun, + saveResult: (result: AgentRunResult) => Promise, +): Promise { + let outcome: AgentRunOutcome; + try { + outcome = await agentRunner.run(run); + } catch (error) { + throw new AgentRunError(error); + } + if (outcome.status !== "completed") { + return outcome; + } + + const saved = await saveResult(outcome.result); + const lifecycle = new ConversationTurnLifecycleService( + getConversationEventStore(), + ); + const common = { + conversationId: run.conversationId, + createdAtMs: saved.finishedAtMs ?? Date.now(), + turnId: run.turnId, + }; + if (saved.outcome === "failed") { + await lifecycle.fail({ + ...common, + failureCode: saved.failureCode, + }); + } else { + await lifecycle.complete({ + ...common, + outcome: saved.outcome, + }); + } + + return { status: "completed" }; +} diff --git a/packages/junior/src/cli/chat.ts b/packages/junior/src/cli/chat.ts index 7640b33010..1067922689 100644 --- a/packages/junior/src/cli/chat.ts +++ b/packages/junior/src/cli/chat.ts @@ -262,9 +262,7 @@ async function prepareLocalChatRun( fallbackWorker: async () => { throw new Error("Local child queue received non-invocation work"); }, - invocationWorker: createAgentInvocationWorker({ - agentRunner, - }), + invocationWorker: createAgentInvocationWorker(agentRunner), }); await processConversationWork(message, { queue: localConversationWork.queue, diff --git a/packages/junior/tests/component/agent-invocation-worker.test.ts b/packages/junior/tests/component/agent-invocation-worker.test.ts index e3d80674e0..d1bfecf41a 100644 --- a/packages/junior/tests/component/agent-invocation-worker.test.ts +++ b/packages/junior/tests/component/agent-invocation-worker.test.ts @@ -107,9 +107,7 @@ describe("agent invocation worker", () => { state: "running", surface: "internal", }); - const worker = createAgentInvocationWorker({ - agentRunner: neverRunAgentRunner(), - }); + const worker = createAgentInvocationWorker(neverRunAgentRunner()); const context = { attempt: { ack: vi.fn(), @@ -166,7 +164,7 @@ describe("agent invocation worker", () => { const run = vi.fn(async () => { throw new Error("agent runner unavailable"); }); - const worker = createAgentInvocationWorker({ agentRunner: { run } }); + const worker = createAgentInvocationWorker({ run }); const message = buildAgentInvocationInboundMessage(created); const context = (isFinalAttempt: boolean, ack: () => Promise) => ({ @@ -219,9 +217,7 @@ describe("agent invocation worker", () => { idempotencyKey: "invalid-child-1", }); const ack = vi.fn(); - const worker = createAgentInvocationWorker({ - agentRunner: neverRunAgentRunner(), - }); + const worker = createAgentInvocationWorker(neverRunAgentRunner()); const context = { attempt: { ack, diff --git a/packages/junior/tests/component/turn-execution.test.ts b/packages/junior/tests/component/turn-execution.test.ts new file mode 100644 index 0000000000..ca999fea5a --- /dev/null +++ b/packages/junior/tests/component/turn-execution.test.ts @@ -0,0 +1,45 @@ +import { describe, expect, it, vi } from "vitest"; +import { createLocalSource } from "@sentry/junior-plugin-api"; +import type { AgentRun } from "@/chat/agent/types"; +import { executeTurn } from "@/chat/runtime/turn-execution"; + +const conversationId = "local:test:turn-execution"; +const run = { + conversationId, + destination: { conversationId, platform: "local" }, + instruction: { text: "Run the task." }, + source: createLocalSource(conversationId), + turnId: "turn-1", +} satisfies AgentRun; + +describe("Turn execution", () => { + it("leaves save errors for the worker to retry", async () => { + const saveError = new Error("result save failed"); + + await expect( + executeTurn( + { + run: vi.fn(async () => ({ + status: "completed" as const, + result: { + diagnostics: { + assistantMessageCount: 1, + modelId: "test-model", + outcome: "success" as const, + toolCalls: [], + toolErrorCount: 0, + toolResultCount: 0, + usedPrimaryText: true, + }, + text: "Done", + }, + })), + }, + run, + async () => { + throw saveError; + }, + ), + ).rejects.toBe(saveError); + }); +}); diff --git a/packages/junior/tests/integration/agent-invocation-concurrency.test.ts b/packages/junior/tests/integration/agent-invocation-concurrency.test.ts index ec4b91ed78..ab96fb19b3 100644 --- a/packages/junior/tests/integration/agent-invocation-concurrency.test.ts +++ b/packages/junior/tests/integration/agent-invocation-concurrency.test.ts @@ -62,9 +62,7 @@ async function createHarness(streamForRun: (request: AgentRun) => StreamFn) { const run = vi.spyOn(agentRunner, "run"); const route = routeAgentInvocationWork({ fallbackWorker: vi.fn(async () => ({ status: "completed" as const })), - invocationWorker: createAgentInvocationWorker({ - agentRunner, - }), + invocationWorker: createAgentInvocationWorker(agentRunner), }); return { diff --git a/packages/junior/tests/integration/agent-invocation-work.test.ts b/packages/junior/tests/integration/agent-invocation-work.test.ts index 3430b5204c..c16fb6b4d1 100644 --- a/packages/junior/tests/integration/agent-invocation-work.test.ts +++ b/packages/junior/tests/integration/agent-invocation-work.test.ts @@ -4,6 +4,7 @@ import { completeAgentInvocation, createAgentInvocation, getAgentInvocation, + getAgentInvocationMessageId, getAgentInvocationTurnId, } from "@/chat/agent-invocations/store"; import { @@ -17,6 +18,7 @@ import type { AgentRun } from "@/chat/agent/types"; import { migrateSchema } from "@/chat/conversations/sql/migrations"; import { createSqlStore } from "@/chat/conversations/sql/store"; import { loadProjection } from "@/chat/conversations/projection"; +import { getConversationEventStore } from "@/chat/db"; import { disconnectStateAdapter, getStateAdapter } from "@/chat/state/adapter"; import { processConversationQueueMessage } from "@/chat/task-execution/vercel-callback"; import { recoverPendingAgentInvocationMailboxAppends } from "@/chat/agent-dispatch/heartbeat"; @@ -62,6 +64,15 @@ async function prepareParentConversation() { return { conversationStore, fixture }; } +async function loadTurnEvents(conversationId: string) { + return (await getConversationEventStore().loadHistory(conversationId)).filter( + (event) => + event.data.type === "turn_started" || + event.data.type === "turn_completed" || + event.data.type === "turn_failed", + ); +} + describe("agent invocation conversation work", () => { afterEach(async () => { await disconnectStateAdapter(); @@ -242,7 +253,7 @@ describe("agent invocation conversation work", () => { } }); - it("runs destinationless child work once and persists its terminal result", async () => { + it("runs child work once without external delivery and saves its result", async () => { const { conversationStore, fixture } = await prepareParentConversation(); const queue = createConversationWorkQueueTestAdapter(); const state = getStateAdapter(); @@ -269,9 +280,7 @@ describe("agent invocation conversation work", () => { })); const route = routeAgentInvocationWork({ fallbackWorker, - invocationWorker: createAgentInvocationWorker({ - agentRunner, - }), + invocationWorker: createAgentInvocationWorker(agentRunner), }); const queueMessage = queue.takeMessage(); @@ -325,6 +334,27 @@ describe("agent invocation conversation work", () => { }), ]), ); + await expect( + loadTurnEvents(created.childConversationId), + ).resolves.toEqual([ + expect.objectContaining({ + data: expect.objectContaining({ + inputMessageIds: [ + getAgentInvocationMessageId(created.invocationId), + ], + surface: "internal", + turnId: getAgentInvocationTurnId(created.invocationId), + type: "turn_started", + }), + }), + expect.objectContaining({ + data: expect.objectContaining({ + outcome: "success", + turnId: getAgentInvocationTurnId(created.invocationId), + type: "turn_completed", + }), + }), + ]); expect(run).toHaveBeenCalledOnce(); expect(run.mock.calls[0]?.[0]).toMatchObject({ conversationId: created.childConversationId, @@ -344,6 +374,67 @@ describe("agent invocation conversation work", () => { } }); + it("finishes a failed child Turn after saving its result", async () => { + const { conversationStore, fixture } = await prepareParentConversation(); + const queue = createConversationWorkQueueTestAdapter(); + const state = getStateAdapter(); + await state.connect(); + try { + const created = await createAndEnqueueAgentInvocation( + { + ...invocationInput, + idempotencyKey: "failed-result-1", + }, + { conversationStore, queue, state }, + ); + const route = routeAgentInvocationWork({ + fallbackWorker: vi.fn(async () => ({ status: "completed" as const })), + invocationWorker: createAgentInvocationWorker( + createModelAgentRunner( + createModelStream([ + { type: "error", errorMessage: "model unavailable" }, + ]), + ), + ), + }); + + await expect( + processConversationQueueMessage(queue.takeMessage(), { + conversationStore, + queue, + run: route, + state, + }), + ).resolves.toMatchObject({ status: "completed" }); + + await expect( + getAgentInvocation(created.invocationId), + ).resolves.toMatchObject({ + errorMessage: "model unavailable", + status: "failed", + }); + await expect( + loadTurnEvents(created.childConversationId), + ).resolves.toEqual([ + expect.objectContaining({ + data: expect.objectContaining({ + type: "turn_started", + turnId: getAgentInvocationTurnId(created.invocationId), + }), + }), + expect.objectContaining({ + data: expect.objectContaining({ + failureCode: "model_execution_failed", + type: "turn_failed", + turnId: getAgentInvocationTurnId(created.invocationId), + }), + }), + ]); + } finally { + await fixture.close(); + } + }); + it("resumes a yielded invocation from durable child state", async () => { const { conversationStore, fixture } = await prepareParentConversation(); const queue = createConversationWorkQueueTestAdapter(); @@ -369,9 +460,7 @@ describe("agent invocation conversation work", () => { ]), ); const run = vi.spyOn(agentRunner, "run"); - const invocationWorker = createAgentInvocationWorker({ - agentRunner, - }); + const invocationWorker = createAgentInvocationWorker(agentRunner); const route = routeAgentInvocationWork({ fallbackWorker: vi.fn(async () => ({ status: "completed" as const })), invocationWorker, @@ -501,9 +590,7 @@ describe("agent invocation conversation work", () => { }); const route = routeAgentInvocationWork({ fallbackWorker: vi.fn(async () => ({ status: "completed" as const })), - invocationWorker: createAgentInvocationWorker({ - agentRunner: neverRunAgentRunner(), - }), + invocationWorker: createAgentInvocationWorker(neverRunAgentRunner()), }); await expect(