diff --git a/packages/pi-plugin/src/dreamer/index.test.ts b/packages/pi-plugin/src/dreamer/index.test.ts index b41772baba..b2662b11ad 100644 --- a/packages/pi-plugin/src/dreamer/index.test.ts +++ b/packages/pi-plugin/src/dreamer/index.test.ts @@ -430,6 +430,90 @@ describe("Pi dreamer wiring", () => { await runPiDreamForProject(identity, "curate", owner); expect(caps).toEqual([4096, 8192]); }); + async function runnerTimeoutsFor( + identity: string, + curate: Record, + ): Promise<{ + timeouts: Array; + fallbackModels: Array; + }> { + db = createDb(); + const owner = {}; + const timeouts: Array = []; + const fallbackModels: Array = []; + __test.setStartDreamScheduleTimerFactory(async () => mock(() => {})); + __test.setPiSubagentRunnerFactory( + () => + ({ + run: mock( + async (args: { + timeoutMs?: number; + fallbackModels?: readonly unknown[]; + }) => { + timeouts.push(args.timeoutMs); + fallbackModels.push(args.fallbackModels); + return { ok: true, assistantText: "curation complete" }; + }, + ), + }) as never, + ); + const config = DreamerConfigSchema.parse({ + pi: { model: "test/model", tasks: { curate } }, + tasks: { curate: { schedule: "0 4 * * *" } }, + }); + insertMemory(db, { + projectPath: identity, + category: "PROJECT_RULES", + content: "Keep configured task timeout.", + }); + registerPiDreamerProject( + dreamerOptions({ + database: db, + projectIdentity: identity, + projectDir: process.cwd(), + registrationOwner: owner, + config, + }), + ); + const result = await runPiDreamForProject(identity, "curate", owner); + expect(result.failed).toEqual([]); + return { timeouts, fallbackModels }; + } + + test("Pi dreamer gives a task configured above 30 minutes a runner ceiling above its timeout", async () => { + const { timeouts, fallbackModels } = await runnerTimeoutsFor( + "git:pi-task-timeout", + { timeout_minutes: 120 }, + ); + // 120 minutes of task timeout plus 5 minutes of headroom, so the + // executor's own timeout ends the run, not the subprocess ceiling. + expect(timeouts).toEqual([(120 + 5) * 60 * 1000]); + expect(fallbackModels).toEqual([undefined]); + }); + + test("Pi dreamer maps every child title to its configured task", () => { + // The task name decides which timeout and model config a child gets, + // so a title that maps to no task falls back to the unknown-task ceiling. + expect(__test.accountingTaskFromTitle("magic-context-dream-curate")).toBe( + "curate", + ); + expect(__test.accountingTaskFromTitle("magic-context-dream-classify")).toBe( + "classify-memories", + ); + expect( + __test.accountingTaskFromTitle("magic-context-dream-user-memories"), + ).toBe("review-user-memories"); + expect( + __test.accountingTaskFromTitle("magic-context-smart-note-confirm-1"), + ).toBe("evaluate-smart-notes"); + expect(__test.accountingTaskFromTitle("something-else")).toBeNull(); + }); + + test("Pi dreamer keeps the runner ceiling above the default task timeout", async () => { + const { timeouts } = await runnerTimeoutsFor("git:pi-task-default", {}); + // An unconfigured task runs with the shared 20-minute default. + expect(timeouts).toEqual([(20 + 5) * 60 * 1000]); + }); test("drops an unknown fallback with a warning while the valid primary runs", async () => { db = createDb(); const identity = "git:pi-model-validation"; diff --git a/packages/pi-plugin/src/dreamer/index.ts b/packages/pi-plugin/src/dreamer/index.ts index e9a05d4ff7..bd014c0ece 100644 --- a/packages/pi-plugin/src/dreamer/index.ts +++ b/packages/pi-plugin/src/dreamer/index.ts @@ -537,6 +537,17 @@ export async function awaitInFlightDreamers( await Promise.allSettled(runs); } +/** + * Minutes the Pi subprocess may run past its task's own timeout. The executor's + * abort signal ends a task at its configured timeout; the subprocess ceiling + * sits this far above it, so the signal always fires first and the run is + * reported as the task's timeout rather than a killed runner. + */ +const RUNNER_TIMEOUT_HEADROOM_MINUTES = 5; + +/** Task timeout assumed when the dream session's task can't be identified. */ +const UNKNOWN_TASK_TIMEOUT_MINUTES = 30; + function createPiDreamerClient( opts: PiDreamerOptions, onAdjunctsRefreshNeeded = opts.onAdjunctsRefreshNeeded, @@ -628,10 +639,14 @@ function createPiDreamerClient( userMessage, model: perTaskModel, fallbackModels: undefined, - // The executor enforces the per-task timeout via its abort signal; - // give the subprocess a generous ceiling so the signal is the - // authority (not a second, conflicting wall-clock here). - timeoutMs: 30 * 60 * 1000, + // The executor enforces the per-task timeout via its abort signal. + // The subprocess ceiling follows that timeout plus headroom, so it + // never ends a task the signal would still let run. + timeoutMs: + ((taskConfig?.timeoutMinutes ?? UNKNOWN_TASK_TIMEOUT_MINUTES) + + RUNNER_TIMEOUT_HEADROOM_MINUTES) * + 60 * + 1000, cwd: dreamSession.directory, signal: args.signal ? AbortSignal.any([args.signal, controller.signal]) @@ -823,7 +838,10 @@ function accountingTaskFromTitle(title: string | undefined): string | null { const task = title.slice("magic-context-dream-".length); if (!title.startsWith("magic-context-dream-") || task.length === 0) return null; - return task === "classify" ? "classify-memories" : task; + // Two child titles use a shorter name than their configured task. + if (task === "classify") return "classify-memories"; + if (task === "user-memories") return "review-user-memories"; + return task; } function readSessionTitle(args: { body?: unknown }): string | undefined { @@ -955,6 +973,7 @@ function makeMessage( } export const __test = { + accountingTaskFromTitle, registeredProjectCount: () => registeredProjects.size, clearRegistrationGeneration: (projectIdentity: string) => { const registration = registeredProjects.get(projectIdentity);