fix(pi): resume a sleep that falls due while other work runs

packages/pi/src/runtime.ts

@@ -519,7 +519,8 @@ function reportBlockedTasks(c: PiContext, inspection: HarnessInspection): void {
  * Keeps the actor awake while Pi Durable can run work, and schedules a wake
  * for its earliest sleeping task, such as a retry backoff. Pi calls `onWake`
  * after open and whenever the value changes. With `onWake`, Pi keeps no timer
- * for sleeps, so the scheduled `pi.wake` action resumes them.
+ * for sleeps, so the scheduled `pi.wake` action resumes them. While other work
+ * is active the actor cannot sleep, so an in-process timer resumes them.
  *
  * Pi runs a sleep only once its time has passed, and reports a wake only when
  * the value changes. So a wake that fires early, even by a millisecond,
@@ -541,9 +542,8 @@ function wakeActor(
 	let release: (() => void) | undefined;
 	let scheduledAt: number | undefined;
 	let scheduling = false;
-	let fallback:
-		| { at: number; timer: ReturnType<typeof setTimeout> }
-		| undefined;
+	let timer: { at: number; handle: ReturnType<typeof setTimeout> } | undefined;
+	let scheduleFailedAt: number | undefined;
 	let stopped = false;
 
 	const stayAwake = () => {
@@ -556,14 +556,31 @@ function wakeActor(
 		release?.();
 		release = undefined;
 	};
+	const clearTimer = () => {
+		if (timer) clearTimeout(timer.handle);
+		timer = undefined;
+		scheduleFailedAt = undefined;
+	};
+	const armTimer = (at: number) => {
+		if (timer?.at === at) return;
+		clearTimer();
+		// setTimeout fires at once past 2^31 - 1 ms, so a far sleep re-arms from here.
+		const delay = Math.min(Math.max(0, at - now()), 2 ** 31 - 1);
+		timer = { at, handle: setTimeout(resumeDue, delay) };
+	};
 	const apply = () => {
 		if (stopped) return;
 		const { active, at } = latest;
-		if (active) return stayAwake();
+		if (timer?.at !== at) clearTimer();
+		if (active) {
+			stayAwake();
+			if (at !== undefined) armTimer(at);
+			return;
+		}
 		if (at === undefined || at === scheduledAt) return letGo();
 		// The actor may sleep only once its wake is written.
 		stayAwake();
-		if (scheduling || fallback?.at === at) return;
+		if (scheduling || scheduleFailedAt === at) return;
 		scheduling = true;
 		c.schedule
 			.at(at, WAKE_ACTION)
@@ -576,11 +593,8 @@ function wakeActor(
 						msg: "pi durable could not schedule its wake, so the actor stays awake until it",
 						error,
 					});
-					if (fallback) clearTimeout(fallback.timer);
-					fallback = {
-						at,
-						timer: setTimeout(resumeDue, Math.max(0, at - now())),
-					};
+					armTimer(at);
+					scheduleFailedAt = at;
 				},
 			)
 			.finally(() => {
@@ -591,15 +605,14 @@ function wakeActor(
 	const resumeDue = () => {
 		if (stopped || !harness) return;
 		harness.resume();
-		if (fallback) clearTimeout(fallback.timer);
-		fallback = undefined;
+		clearTimer();
 		scheduledAt = undefined;
 		latest = harness.wake();
 		apply();
 	};
 	const stop = () => {
 		stopped = true;
-		if (fallback) clearTimeout(fallback.timer);
+		clearTimer();
 		letGo();
 	};
 	// Keep-awake only blocks idle sleep. Once shutdown starts it would only delay it.

packages/pi/tests/wake-while-busy.test.ts

@@ -0,0 +1,103 @@
+import { randomUUID } from "node:crypto";
+import { BACKGROUND_CONTEXT } from "@earendil-works/chord/context";
+import { fauxAssistantMessage } from "@earendil-works/pi-ai";
+import {
+	createRegistry,
+	defineExtension,
+	defineTask,
+	ROOT_CONVERSATION_ID,
+} from "@earendil-works/pi-durable";
+import { setup } from "rivetkit";
+import { setupTest } from "rivetkit/test";
+import { expect, test, vi } from "vitest";
+import { pi } from "../src/index.js";
+import { createMockModel } from "./helpers/mock-model.js";
+
+const BACKOFF_MS = 2_000;
+
+/** Runs until its actor stops and never commits, so nothing else starts a scheduler pass. */
+const StuckTask = defineTask<Record<string, never>, { phase: "run" }, null>({
+	name: "app.stuck",
+	version: 1,
+	initial: () => ({ phase: "run" }),
+	phases: {
+		run: (_task, _runtime, context) =>
+			new Promise<never>((_, reject) => {
+				context.abortSignal?.addEventListener(
+					"abort",
+					() => reject(context.abortSignal?.reason),
+					{ once: true },
+				);
+			}),
+	},
+	abort: async (_task, runtime, context) => {
+		await runtime.commit(
+			() => ({ status: "terminal", outcome: { status: "aborted" } }),
+			context,
+		);
+	},
+});
+
+const mock = createMockModel();
+mock.reply(
+	"flaky",
+	fauxAssistantMessage("", {
+		stopReason: "error",
+		errorMessage: "overloaded_error: Overloaded",
+	}),
+	fauxAssistantMessage("recovered"),
+);
+
+const piRegistry = createRegistry();
+piRegistry.install(defineExtension({ name: "stuck", tasks: [StuckTask] }));
+
+const agent = pi({
+	model: "mock/mock-model",
+	providers: { mock: mock.providerConfig },
+	apiKeys: { mock: "mock" },
+	registry: piRegistry,
+	settings: {
+		retry: { enabled: true, maxRetries: 1, baseDelayMs: BACKOFF_MS },
+	},
+	actions: {
+		startStuckTask: async (c) => {
+			await c.pi.root(BACKGROUND_CONTEXT);
+			await c.pi.commit(
+				(tx) =>
+					tx.createTask(
+						StuckTask,
+						{},
+						{
+							ownership: { kind: "conversation" },
+							conversationId: ROOT_CONVERSATION_ID,
+						},
+					),
+				BACKGROUND_CONTEXT,
+			);
+		},
+	},
+});
+const registry = setup({ use: { agent } });
+
+test("a retry whose backoff ends while another task runs happens on time, with no client call", async (c) => {
+	const { client } = await setupTest(c, registry);
+	const handle = client.agent.getOrCreate(["busy", randomUUID()]);
+	const root = await handle.harness.root();
+	await handle.startStuckTask();
+	const callsBefore = mock.requests.length;
+
+	const submission = await handle.conversation.submit(root.id, {
+		type: "input",
+		content: "flaky",
+	});
+
+	await vi.waitFor(() => expect(mock.requests.length).toBe(callsBefore + 1), {
+		timeout: 10_000,
+		interval: 20,
+	});
+	await vi.waitFor(() => expect(mock.requests.length).toBe(callsBefore + 2), {
+		timeout: BACKOFF_MS * 4,
+		interval: 50,
+	});
+	expect((await handle.submission.wait(submission.id)).status).toBe("done");
+});