From 9954c407b1171ee3b8157ac5f862149303aeab11 Mon Sep 17 00:00:00 2001 From: Annika <14750837+annikaschmid@users.noreply.github.com> Date: Wed, 13 May 2026 18:11:07 -0400 Subject: [PATCH 1/2] fix(agent): complete automation runs after end_turn --- packages/agent/src/server/agent-server.ts | 33 ++++++ .../agent/src/server/question-relay.test.ts | 111 ++++++++++++++++++ 2 files changed, 144 insertions(+) diff --git a/packages/agent/src/server/agent-server.ts b/packages/agent/src/server/agent-server.ts index 12ea84c474..90e35a8eca 100644 --- a/packages/agent/src/server/agent-server.ts +++ b/packages/agent/src/server/agent-server.ts @@ -1153,6 +1153,7 @@ export class AgentServer { if (result.stopReason === "end_turn") { await this.relayAgentResponse(payload); + await this.maybeCompleteAutomationRun(payload, taskRun); } } catch (error) { this.logger.error("Failed to send initial task message", error); @@ -1276,6 +1277,7 @@ export class AgentServer { if (result.stopReason === "end_turn") { await this.relayAgentResponse(payload); + await this.maybeCompleteAutomationRun(payload, taskRun); } } catch (error) { this.logger.error("Failed to send resume message", error); @@ -1744,6 +1746,37 @@ ${attributionInstructions} } } + private shouldAutoCompleteAutomationRun(taskRun: TaskRun | null): boolean { + const state = taskRun?.state; + return Boolean( + state && + typeof state === "object" && + typeof state.automation_id === "string" && + state.automation_id.trim().length > 0, + ); + } + + private async maybeCompleteAutomationRun( + payload: JwtPayload, + taskRun: TaskRun | null, + ): Promise { + if (!this.shouldAutoCompleteAutomationRun(taskRun)) { + return; + } + + try { + await this.posthogAPI.updateTaskRun(payload.task_id, payload.run_id, { + status: "completed", + }); + this.logger.debug("Automation run marked completed after end_turn", { + taskId: payload.task_id, + runId: payload.run_id, + }); + } catch (error) { + this.logger.error("Failed to mark automation run completed", error); + } + } + private async signalTaskComplete( payload: JwtPayload, stopReason: string, diff --git a/packages/agent/src/server/question-relay.test.ts b/packages/agent/src/server/question-relay.test.ts index daf597ec19..b1b193b12d 100644 --- a/packages/agent/src/server/question-relay.test.ts +++ b/packages/agent/src/server/question-relay.test.ts @@ -551,6 +551,117 @@ describe("Question relay", () => { }); }); + it("marks automation-triggered runs completed after a successful first turn", async () => { + vi.spyOn(server.posthogAPI, "getTask").mockResolvedValue({ + id: "test-task-id", + title: "t", + description: "original task description", + } as unknown as Task); + vi.spyOn(server.posthogAPI, "getTaskRun").mockResolvedValue({ + id: "test-run-id", + task: "test-task-id", + state: { automation_id: "automation-1" }, + } as unknown as TaskRun); + vi.spyOn( + server as unknown as { + syncCloudBranchMetadata: ( + payload: Record, + ) => Promise; + }, + "syncCloudBranchMetadata", + ).mockResolvedValue(undefined); + + const promptSpy = vi.fn().mockResolvedValue({ stopReason: "end_turn" }); + const updateTaskRunSpy = vi + .spyOn(server.posthogAPI, "updateTaskRun") + .mockResolvedValue({} as TaskRun); + const relaySpy = vi + .spyOn(server.posthogAPI, "relayMessage") + .mockResolvedValue(undefined); + server.session = { + payload: TEST_PAYLOAD, + acpSessionId: "acp-session", + clientConnection: { prompt: promptSpy }, + logWriter: { + appendRawLine: vi.fn(), + flushAll: vi.fn().mockResolvedValue(undefined), + getFullAgentResponse: vi + .fn() + .mockReturnValue( + "At 2026-05-13 17:45, the waitlist count is 2827.", + ), + resetTurnMessages: vi.fn(), + flush: vi.fn().mockResolvedValue(undefined), + isRegistered: vi.fn().mockReturnValue(true), + }, + }; + + await server.sendInitialTaskMessage(TEST_PAYLOAD); + + expect(relaySpy).toHaveBeenCalledWith( + "test-task-id", + "test-run-id", + "At 2026-05-13 17:45, the waitlist count is 2827.", + ); + expect(updateTaskRunSpy).toHaveBeenCalledWith( + "test-task-id", + "test-run-id", + { + status: "completed", + }, + ); + }); + + it("keeps non-automation runs open after a successful first turn", async () => { + vi.spyOn(server.posthogAPI, "getTask").mockResolvedValue({ + id: "test-task-id", + title: "t", + description: "original task description", + } as unknown as Task); + vi.spyOn(server.posthogAPI, "getTaskRun").mockResolvedValue({ + id: "test-run-id", + task: "test-task-id", + state: {}, + } as unknown as TaskRun); + vi.spyOn( + server as unknown as { + syncCloudBranchMetadata: ( + payload: Record, + ) => Promise; + }, + "syncCloudBranchMetadata", + ).mockResolvedValue(undefined); + + const promptSpy = vi.fn().mockResolvedValue({ stopReason: "end_turn" }); + const updateTaskRunSpy = vi + .spyOn(server.posthogAPI, "updateTaskRun") + .mockResolvedValue({} as TaskRun); + vi.spyOn(server.posthogAPI, "relayMessage").mockResolvedValue(undefined); + server.session = { + payload: TEST_PAYLOAD, + acpSessionId: "acp-session", + clientConnection: { prompt: promptSpy }, + logWriter: { + appendRawLine: vi.fn(), + flushAll: vi.fn().mockResolvedValue(undefined), + getFullAgentResponse: vi.fn().mockReturnValue("done"), + resetTurnMessages: vi.fn(), + flush: vi.fn().mockResolvedValue(undefined), + isRegistered: vi.fn().mockReturnValue(true), + }, + }; + + await server.sendInitialTaskMessage(TEST_PAYLOAD); + + expect(updateTaskRunSpy).not.toHaveBeenCalledWith( + "test-task-id", + "test-run-id", + { + status: "completed", + }, + ); + }); + it("does not replay a transient upstream termination before any session activity", async () => { vi.spyOn(server.posthogAPI, "getTask").mockResolvedValue({ id: "test-task-id", From 72d3660358bfe2b509051b07e5dda639a9e1a119 Mon Sep 17 00:00:00 2001 From: Annika <14750837+annikaschmid@users.noreply.github.com> Date: Wed, 13 May 2026 18:34:41 -0400 Subject: [PATCH 2/2] test(agent): address automation run review feedback --- .../agent/src/server/question-relay.test.ts | 118 ++++++++++++------ 1 file changed, 77 insertions(+), 41 deletions(-) diff --git a/packages/agent/src/server/question-relay.test.ts b/packages/agent/src/server/question-relay.test.ts index b1b193b12d..d37d536e41 100644 --- a/packages/agent/src/server/question-relay.test.ts +++ b/packages/agent/src/server/question-relay.test.ts @@ -2,6 +2,7 @@ import { type SetupServerApi, setupServer } from "msw/node"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { classifyAgentError } from "../adapters/claude/conversion/sdk-to-acp"; import type { PostHogAPIClient } from "../posthog-api"; +import type { ResumeState } from "../resume"; import { createTestRepo, type TestRepo } from "../test/fixtures/api"; import { createPostHogHandlers } from "../test/mocks/msw-handlers"; import type { Task, TaskRun } from "../types"; @@ -22,9 +23,11 @@ interface TestableAgentServer { }>; }; questionRelayedToSlack: boolean; + resumeState: ResumeState | null; session: unknown; relayAgentResponse: (payload: Record) => Promise; sendInitialTaskMessage: (payload: Record) => Promise; + syncCloudBranchMetadata: (payload: Record) => Promise; } const TEST_PAYLOAD = { @@ -457,6 +460,7 @@ describe("Question relay", () => { acpSessionId: "acp-session", clientConnection: { prompt: promptSpy }, logWriter: { + appendRawLine: vi.fn(), flushAll: vi.fn().mockResolvedValue(undefined), getFullAgentResponse: vi.fn().mockReturnValue(null), resetTurnMessages: vi.fn(), @@ -501,6 +505,7 @@ describe("Question relay", () => { acpSessionId: "acp-session", clientConnection: { prompt: promptSpy }, logWriter: { + appendRawLine: vi.fn(), flushAll: vi.fn().mockResolvedValue(undefined), getFullAgentResponse: vi.fn().mockReturnValue(null), resetTurnMessages: vi.fn(), @@ -535,6 +540,7 @@ describe("Question relay", () => { acpSessionId: "acp-session", clientConnection: { prompt: promptSpy }, logWriter: { + appendRawLine: vi.fn(), flushAll: vi.fn().mockResolvedValue(undefined), getFullAgentResponse: vi.fn().mockReturnValue(null), resetTurnMessages: vi.fn(), @@ -551,7 +557,20 @@ describe("Question relay", () => { }); }); - it("marks automation-triggered runs completed after a successful first turn", async () => { + it.each([ + { + name: "marks automation-triggered runs completed after a successful first turn", + state: { automation_id: "automation-1" }, + agentResponse: "At 2026-05-13 17:45, the waitlist count is 2827.", + expectCompleted: true, + }, + { + name: "keeps non-automation runs open after a successful first turn", + state: {}, + agentResponse: "done", + expectCompleted: false, + }, + ])("$name", async ({ state, agentResponse, expectCompleted }) => { vi.spyOn(server.posthogAPI, "getTask").mockResolvedValue({ id: "test-task-id", title: "t", @@ -560,16 +579,9 @@ describe("Question relay", () => { vi.spyOn(server.posthogAPI, "getTaskRun").mockResolvedValue({ id: "test-run-id", task: "test-task-id", - state: { automation_id: "automation-1" }, + state, } as unknown as TaskRun); - vi.spyOn( - server as unknown as { - syncCloudBranchMetadata: ( - payload: Record, - ) => Promise; - }, - "syncCloudBranchMetadata", - ).mockResolvedValue(undefined); + vi.spyOn(server, "syncCloudBranchMetadata").mockResolvedValue(undefined); const promptSpy = vi.fn().mockResolvedValue({ stopReason: "end_turn" }); const updateTaskRunSpy = vi @@ -585,11 +597,7 @@ describe("Question relay", () => { logWriter: { appendRawLine: vi.fn(), flushAll: vi.fn().mockResolvedValue(undefined), - getFullAgentResponse: vi - .fn() - .mockReturnValue( - "At 2026-05-13 17:45, the waitlist count is 2827.", - ), + getFullAgentResponse: vi.fn().mockReturnValue(agentResponse), resetTurnMessages: vi.fn(), flush: vi.fn().mockResolvedValue(undefined), isRegistered: vi.fn().mockReturnValue(true), @@ -601,42 +609,54 @@ describe("Question relay", () => { expect(relaySpy).toHaveBeenCalledWith( "test-task-id", "test-run-id", - "At 2026-05-13 17:45, the waitlist count is 2827.", - ); - expect(updateTaskRunSpy).toHaveBeenCalledWith( - "test-task-id", - "test-run-id", - { - status: "completed", - }, + agentResponse, ); + + if (expectCompleted) { + expect(updateTaskRunSpy).toHaveBeenCalledWith( + "test-task-id", + "test-run-id", + { + status: "completed", + }, + ); + } else { + expect(updateTaskRunSpy).not.toHaveBeenCalledWith( + "test-task-id", + "test-run-id", + { + status: "completed", + }, + ); + } }); - it("keeps non-automation runs open after a successful first turn", async () => { - vi.spyOn(server.posthogAPI, "getTask").mockResolvedValue({ - id: "test-task-id", - title: "t", - description: "original task description", - } as unknown as Task); + it("marks automation-triggered runs completed after a successful resume turn", async () => { vi.spyOn(server.posthogAPI, "getTaskRun").mockResolvedValue({ id: "test-run-id", task: "test-task-id", - state: {}, + state: { automation_id: "automation-1" }, } as unknown as TaskRun); - vi.spyOn( - server as unknown as { - syncCloudBranchMetadata: ( - payload: Record, - ) => Promise; - }, - "syncCloudBranchMetadata", - ).mockResolvedValue(undefined); + vi.spyOn(server, "syncCloudBranchMetadata").mockResolvedValue(undefined); const promptSpy = vi.fn().mockResolvedValue({ stopReason: "end_turn" }); const updateTaskRunSpy = vi .spyOn(server.posthogAPI, "updateTaskRun") .mockResolvedValue({} as TaskRun); - vi.spyOn(server.posthogAPI, "relayMessage").mockResolvedValue(undefined); + const relaySpy = vi + .spyOn(server.posthogAPI, "relayMessage") + .mockResolvedValue(undefined); + server.resumeState = { + conversation: [ + { + role: "user", + content: [{ type: "text", text: "Please continue from here." }], + }, + ], + latestGitCheckpoint: null, + interrupted: false, + logEntryCount: 1, + }; server.session = { payload: TEST_PAYLOAD, acpSessionId: "acp-session", @@ -644,7 +664,7 @@ describe("Question relay", () => { logWriter: { appendRawLine: vi.fn(), flushAll: vi.fn().mockResolvedValue(undefined), - getFullAgentResponse: vi.fn().mockReturnValue("done"), + getFullAgentResponse: vi.fn().mockReturnValue("resume reply"), resetTurnMessages: vi.fn(), flush: vi.fn().mockResolvedValue(undefined), isRegistered: vi.fn().mockReturnValue(true), @@ -653,7 +673,23 @@ describe("Question relay", () => { await server.sendInitialTaskMessage(TEST_PAYLOAD); - expect(updateTaskRunSpy).not.toHaveBeenCalledWith( + expect(promptSpy).toHaveBeenCalledWith({ + sessionId: "acp-session", + prompt: [ + expect.objectContaining({ + type: "text", + text: expect.stringContaining( + "You are resuming a previous conversation.", + ), + }), + ], + }); + expect(relaySpy).toHaveBeenCalledWith( + "test-task-id", + "test-run-id", + "resume reply", + ); + expect(updateTaskRunSpy).toHaveBeenCalledWith( "test-task-id", "test-run-id", {