diff --git a/packages/server/src/server/agent/providers/pi/extensions/pi-subagents/index.test.ts b/packages/server/src/server/agent/providers/pi/extensions/pi-subagents/index.test.ts index 12eb85f66b..f21e16ae96 100644 --- a/packages/server/src/server/agent/providers/pi/extensions/pi-subagents/index.test.ts +++ b/packages/server/src/server/agent/providers/pi/extensions/pi-subagents/index.test.ts @@ -1,6 +1,6 @@ import { describe, expect, test, vi } from "vitest"; import { fileURLToPath } from "node:url"; -import { mkdtemp, writeFile, rm } from "node:fs/promises"; +import { mkdtemp, writeFile, appendFile, rm } from "node:fs/promises"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { createPiExtensionHost } from "../index.js"; @@ -792,7 +792,7 @@ describe("pi-subagents adapter", () => { expect( subagentUpserts(received).filter((event) => event.id === "call-wf:alpha" && event.status), - ).toEqual([]); + ).toEqual([expect.objectContaining({ id: "call-wf:alpha", status: "completed" })]); // The child that never reported a verdict of its own is still settled by that tick. expect( subagentUpserts(received).findLast((event) => event.id === "call-wf:beta")?.status, @@ -982,7 +982,13 @@ describe("pi-subagents adapter", () => { const alphaUpserts = () => subagentUpserts(received).filter((event) => event.id === "call-wf:alpha"); expect(alphaUpserts()).toEqual([ - { type: "upsert", id: "call-wf:alpha", title: "scout", liveness: "poll" }, + { + type: "upsert", + id: "call-wf:alpha", + title: "scout", + status: "completed", + liveness: "poll", + }, ]); // Reported once, not on every tick. @@ -1031,4 +1037,115 @@ describe("pi-subagents adapter", () => { expect.objectContaining({ id: "fanout:beta", title: "reviewer", status: "completed" }), ]); }); + test("replays keyed async workflow child transcripts without a live poll", async () => { + const dir = await mkdtemp(join(tmpdir(), "paseo-pi-workflow-replay-")); + const host = createPiExtensionHost(); + try { + const sessionFile = join(dir, "child.jsonl"); + await writeFile( + sessionFile, + JSON.stringify({ + type: "message", + message: { + role: "assistant", + content: [{ type: "text", text: "Saved reviewer answer" }], + }, + }) + "\n", + ); + await writeFile( + join(dir, "status.json"), + JSON.stringify({ + state: "complete", + steps: [{ agent: "reviewer", workflowKey: "review", status: "complete", sessionFile }], + }), + ); + const output = host.mapToolCall({ + callId: "replayed-workflow", + toolName: "subagent", + args: { workflowScript: "return await runs.all(items);", async: true }, + status: "completed", + result: { + details: { + mode: "workflow", + runId: "saved-run", + asyncId: "saved-run", + asyncDir: dir, + results: [], + }, + }, + }); + const events = await output!.hydration; + expect( + events.some( + (event) => + event.type === "provider_subagent" && + event.event.type === "timeline" && + event.event.id === "replayed-workflow:review", + ), + ).toBe(true); + } finally { + host.close(); + await rm(dir, { recursive: true, force: true }); + } + }); + + test("stops following a child whose completion arrived before its transcript was attached", async () => { + const dir = await mkdtemp(join(tmpdir(), "paseo-pi-finished-reader-")); + const host = createPiExtensionHost(); + const received: AgentStreamEvent[] = []; + try { + const sessionFile = join(dir, "child.jsonl"); + const message = (text: string) => + JSON.stringify({ + type: "message", + message: { + role: "assistant", + content: [{ type: "text", text }], + }, + }) + "\n"; + await writeFile(sessionFile, message("Finished child")); + host.follow((event) => received.push(event)); + const spawn = host.mapToolCall({ + callId: "fast-workflow", + toolName: "subagent", + args: { workflow: true, async: true }, + status: "completed", + result: { + details: { + mode: "workflow", + runId: "fast-run", + asyncId: "fast-run", + asyncDir: dir, + results: [], + }, + }, + }); + await spawn!.hydration; + const notice = host.mapCustomMessage({ + role: "custom", + customType: "subagent-incremental-child-notify", + content: "Workflow child completed: **review**\nWorkflow run: fast-run", + }); + await notice!.hydration; + await writeFile( + join(dir, "status.json"), + JSON.stringify({ + state: "running", + steps: [{ agent: "reviewer", workflowKey: "review", status: "running", sessionFile }], + }), + ); + await expect.poll(() => hasChildTimeline(received)).toBe(true); + const timelineCount = () => + received.filter( + (event) => event.type === "provider_subagent" && event.event.type === "timeline", + ).length; + const count = timelineCount(); + await appendFile(sessionFile, message("Later content must not be followed")); + await new Promise((resolve) => setTimeout(resolve, 600)); + expect(timelineCount()).toBe(count); + } finally { + host.close(); + await rm(dir, { recursive: true, force: true }); + } + }); }); diff --git a/packages/server/src/server/agent/providers/pi/extensions/pi-subagents/index.ts b/packages/server/src/server/agent/providers/pi/extensions/pi-subagents/index.ts index 2c84ff35d1..59cd96f8a9 100644 --- a/packages/server/src/server/agent/providers/pi/extensions/pi-subagents/index.ts +++ b/packages/server/src/server/agent/providers/pi/extensions/pi-subagents/index.ts @@ -177,11 +177,11 @@ export const piSubagents: PiExtension = { const callsByRun = new Map(); const readSessions = new Set(); /** Rows already reported terminal; a late give-up must not rewrite one of them. */ - const settledRows = new Set(); + const settledRows = new Map>(); /** Rows whose poll already reported the step's agent. A settled row is still named once. */ const namedRows = new Set(); - const recordTerminal = (id: string, status: string) => { - if (status !== "running") settledRows.add(id); + const recordTerminal = (id: string, status: NonNullable) => { + if (status !== "running") settledRows.set(id, status); }; /** Keyed by async directory, which both status.json and the completion notice carry. */ const asyncRuns = new Map(); @@ -281,6 +281,66 @@ export const piSubagents: PiExtension = { } return subagents.length ? { subagents, childSessions } : undefined; }; + /** Project a saved status both when replaying a launch and when its live poll advances. */ + const projectAsyncStatus = (run: AsyncRun, status: z.infer) => { + const subagents: NonNullable = []; + const childSessions: NonNullable = []; + const runTerminal = TERMINAL_RUN_STATES.has(status.state); + for (const [index, step] of status.steps.entries()) { + const id = rowId(run.owner, status.steps.length, step.workflowKey, index); + run.rows.add(id); + // A child's notice supplies its verdict even when its saved status is stale. Reuse that + // verdict when attaching a transcript or filling its name after the notice arrived. + const settledStatus = settledRows.get(id); + const childSession = + step.sessionFile && !readSessions.has(step.sessionFile) + ? { id, file: step.sessionFile } + : undefined; + if (!settledStatus) { + const childStatus = asyncChildStatus(status.state, step.status); + recordTerminal(id, childStatus); + namedRows.add(id); + subagents.push({ + type: "upsert", + id, + title: step.agent, + status: childStatus, + liveness: "poll" as const, + }); + } else if (!namedRows.has(id) || childSession) { + // Attach a late transcript with its known verdict, so the host reads it once instead of + // starting a follower that no later terminal update will release. + namedRows.add(id); + subagents.push({ + type: "upsert", + id, + title: step.agent, + status: settledStatus, + liveness: "poll" as const, + }); + } + if (childSession) { + readSessions.add(childSession.file); + childSessions.push(childSession); + } + } + if (runTerminal) run.polling = false; + // The launch row stands for the whole run, and a single unnamed child already took its + // id. Re-reporting it every tick is also what proves to the manager that this run is alive. + // A row that already settled keeps that verdict: a result that landed before the run's + // steps named its children must not be reverted to running by the next tick. + if (!run.rows.has(run.owner) && !settledRows.has(run.owner)) { + const ownerStatus = run.polling ? ("running" as const) : asyncRunStatus(status.state); + recordTerminal(run.owner, ownerStatus); + subagents.push({ + type: "upsert", + id: run.owner, + status: ownerStatus, + liveness: "poll", + }); + } + return { subagents, childSessions }; + }; /** A settled spawn: Pi either reports the rows it ran, or an async directory the poll follows. */ const mapSettledSpawn = ( call: PiExtensionToolCall, @@ -310,6 +370,10 @@ export const piSubagents: PiExtension = { } // The placeholder row is now backed by a per-tick poll, which is what proves it alive. rememberAsyncRun(details.data, call); + const dir = details.data.asyncDir; + const run = dir ? asyncRuns.get(dir) : undefined; + const status = dir ? readAsyncStatus(dir) : undefined; + if (run && status) return projectAsyncStatus(run, status); return { subagents: [{ ...base, status: "running" as const, liveness: "poll" as const }] }; } const collected = collectRows(call.callId, details.data.results, description, failed); @@ -437,51 +501,9 @@ export const piSubagents: PiExtension = { continue; } run.misses = 0; - const runTerminal = TERMINAL_RUN_STATES.has(status.state); - for (const [index, step] of status.steps.entries()) { - const id = rowId(run.owner, status.steps.length, step.workflowKey, index); - run.rows.add(id); - // A child notice is this child's verdict, so a settled row gets no further status write - // from the poll. While the run is going, a stale step entry must not revert it to - // running; on the terminal tick, the run's own outcome must not overwrite it. - if (!settledRows.has(id)) { - const childStatus = asyncChildStatus(status.state, step.status); - recordTerminal(id, childStatus); - namedRows.add(id); - subagents.push({ - type: "upsert", - id, - title: step.agent, - status: childStatus, - liveness: "poll" as const, - }); - } else if (!namedRows.has(id)) { - // A child notice can settle a row before any poll saw it, and that notice carries no - // agent name. Omit `status`: the store keeps the settled verdict and only takes the - // title the run's status file names. - namedRows.add(id); - subagents.push({ type: "upsert", id, title: step.agent, liveness: "poll" as const }); - } - if (step.sessionFile && !readSessions.has(step.sessionFile)) { - readSessions.add(step.sessionFile); - childSessions.push({ id, file: step.sessionFile }); - } - } - if (runTerminal) run.polling = false; - // The launch row stands for the whole run, and a single unnamed child already took its - // id. Re-reporting it every tick is also what proves to the manager that this run is alive. - // A row that already settled keeps that verdict: a result that landed before the run's - // steps named its children must not be reverted to running by the next tick. - if (!run.rows.has(run.owner) && !settledRows.has(run.owner)) { - const ownerStatus = run.polling ? ("running" as const) : asyncRunStatus(status.state); - recordTerminal(run.owner, ownerStatus); - subagents.push({ - type: "upsert", - id: run.owner, - status: ownerStatus, - liveness: "poll", - }); - } + const output = projectAsyncStatus(run, status); + subagents.push(...output.subagents); + childSessions.push(...output.childSessions); } return subagents.length || childSessions.length ? { subagents, childSessions } : undefined; },