@vellumai/assistant 0.8.5 → 0.8.6

package/src/background-wake/background-wake-routes.test.ts

@@ -1,5 +1,6 @@
 import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test";
 
+import type { SchedulerDueWorkResult } from "../schedule/scheduler.js";
 import type { BackgroundWakeIntent } from "./next-wake.js";
 import {
   clearBackgroundWakeRuntime,
@@ -7,12 +8,32 @@ import {
 } from "./runtime-registry.js";
 
 type MockIntent = BackgroundWakeIntent;
+type MockCompletionPayload = {
+  leaseId: string;
+  status: "completed" | "failed" | "expired";
+  error?: string;
+  nextIntent?: MockIntent | null;
+};
 
 let computedIntent: MockIntent | null;
 let computeCalls: number;
-
-const { ROUTES, setBackgroundWakeIntentComputerForTest } =
-  await import("../runtime/routes/background-wake-routes.js");
+const mockRenewBackgroundWakeLease = mock(async (_leaseId: string) => ({
+  status: "renewed" as const,
+  httpStatus: 200,
+}));
+const mockCompleteBackgroundWakeLease = mock(
+  async (_args: MockCompletionPayload) => ({
+    status: "completed" as const,
+    httpStatus: 200,
+  }),
+);
+
+const {
+  ROUTES,
+  flushBackgroundWakeDrainsForTest,
+  setBackgroundWakeIntentComputerForTest,
+  setBackgroundWakeLeaseClientForTest,
+} = await import("../runtime/routes/background-wake-routes.js");
 
 function findHandler(operationId: string) {
   const route = ROUTES.find(
@@ -24,9 +45,74 @@ function findHandler(operationId: string) {
 
 const NOW = 1_800_000_000_000;
 
+function schedulerResultFixture(
+  overrides: Partial<SchedulerDueWorkResult> = {},
+): SchedulerDueWorkResult {
+  return {
+    claimed: 0,
+    completed: 0,
+    failed: 0,
+    skipped: 0,
+    stillPending: 0,
+    ...overrides,
+  };
+}
+
+function firstCallArg(fn: {
+  mock: { calls: unknown[] };
+}): Record<string, unknown> {
+  const call = fn.mock.calls[0] as unknown[] | undefined;
+  return (call?.[0] ?? {}) as Record<string, unknown>;
+}
+
+function acceptedResponse(
+  overrides: Partial<{
+    leaseId: string;
+    reason: string;
+    sourceGeneration: string;
+    startedAt: number;
+    deadlineAt: number;
+  }> = {},
+) {
+  return {
+    accepted: true,
+    leaseId: "lease-123",
+    reason: "schedule",
+    sourceGeneration: "source-generation",
+    startedAt: NOW,
+    deadlineAt: NOW + 30_000,
+    ...overrides,
+  };
+}
+
+function lastCompletionPayload(): MockCompletionPayload {
+  const calls = mockCompleteBackgroundWakeLease.mock.calls;
+  const call = calls[calls.length - 1] as unknown[] | undefined;
+  if (!call) throw new Error("completeBackgroundWakeLease was not called");
+  return call[0] as MockCompletionPayload;
+}
+
 describe("background wake runtime routes", () => {
   beforeEach(() => {
     clearBackgroundWakeRuntime();
+    mockRenewBackgroundWakeLease.mockClear();
+    mockCompleteBackgroundWakeLease.mockClear();
+    mockRenewBackgroundWakeLease.mockImplementation(
+      async (_leaseId: string) => ({
+        status: "renewed" as const,
+        httpStatus: 200,
+      }),
+    );
+    mockCompleteBackgroundWakeLease.mockImplementation(
+      async (_args: MockCompletionPayload) => ({
+        status: "completed" as const,
+        httpStatus: 200,
+      }),
+    );
+    setBackgroundWakeLeaseClientForTest({
+      renew: mockRenewBackgroundWakeLease,
+      complete: mockCompleteBackgroundWakeLease,
+    });
     computeCalls = 0;
     computedIntent = intentFixture({ nextWakeAt: Date.now() + 10 * 60_000 });
     setBackgroundWakeIntentComputerForTest(() => {
@@ -35,8 +121,11 @@ describe("background wake runtime routes", () => {
     });
   });
 
-  afterEach(() => {
+  afterEach(async () => {
+    await flushBackgroundWakeDrainsForTest();
     setBackgroundWakeIntentComputerForTest(null);
+    setBackgroundWakeLeaseClientForTest(null);
+    clearBackgroundWakeRuntime();
   });
 
   test("intent route returns the current computed wake intent", async () => {
@@ -82,67 +171,84 @@ describe("background wake runtime routes", () => {
   });
 
   test("drain-due invokes due heartbeat and scheduler work", async () => {
-    const heartbeatRunOnce = mock(async () => true);
-    const schedulerRunOnce = mock(async () => 2);
+    const dueAt = Date.now() - 1;
+    const recomputedIntent = intentFixture({
+      nextWakeAt: Date.now() + 60_000,
+      sourceGeneration: "next-generation",
+    });
+    computedIntent = intentFixture({
+      nextWakeAt: dueAt,
+      actualNextDueAt: dueAt,
+      reason: "mixed",
+    });
+    const heartbeatRunManaged = mock(async () => ({
+      due: true,
+      completed: 1,
+      skipped: 0,
+    }));
+    const schedulerRunDue = mock(async () => {
+      computedIntent = recomputedIntent;
+      return schedulerResultFixture({ claimed: 2, completed: 2 });
+    });
     registerBackgroundWakeRuntime({
       heartbeat: {
-        nextRunAt: Date.now() - 1,
-        runOnce: heartbeatRunOnce,
+        nextRunAt: dueAt,
+        runManagedWakeIfDue: heartbeatRunManaged,
       },
       scheduler: {
-        runOnce: schedulerRunOnce,
+        runOnce: mock(async () => 2),
+        runDueWorkOnce: schedulerRunDue,
       },
     });
-    computedIntent = intentFixture({
-      nextWakeAt: NOW + 60_000,
-      sourceGeneration: "next-generation",
-    });
     const handler = findHandler("drainDueBackgroundWake");
 
-    const response = await handler({ body: drainBodyFixture() });
+    const response = await handler({
+      body: drainBodyFixture({ reason: "mixed" }),
+    });
+    await flushBackgroundWakeDrainsForTest();
 
-    expect(heartbeatRunOnce).toHaveBeenCalledTimes(1);
-    expect(heartbeatRunOnce).toHaveBeenCalledWith();
-    expect(schedulerRunOnce).toHaveBeenCalledTimes(1);
-    expect(response).toEqual({
+    expect(heartbeatRunManaged).toHaveBeenCalledTimes(1);
+    expect(heartbeatRunManaged).toHaveBeenCalledWith(
+      expect.objectContaining({ scheduledFor: NOW }),
+    );
+    expect(firstCallArg(heartbeatRunManaged)).not.toHaveProperty("assumeDue");
+    expect(schedulerRunDue).toHaveBeenCalledTimes(1);
+    expect(response).toEqual(acceptedResponse({ reason: "mixed" }));
+    expect(lastCompletionPayload()).toEqual({
       leaseId: "lease-123",
-      reason: "schedule",
-      sourceGeneration: "source-generation",
-      startedAt: NOW,
-      deadlineAt: NOW + 30_000,
-      counts: {
-        heartbeat: 1,
-        scheduler: 2,
-        total: 3,
-      },
-      nextIntent: computedIntent,
+      status: "completed",
+      nextIntent: recomputedIntent,
     });
   });
 
   test("drain-due skips heartbeat when it is not due and reports no scheduler work", async () => {
-    const heartbeatRunOnce = mock(async () => true);
-    const schedulerRunOnce = mock(async () => 0);
+    const heartbeatRunManaged = mock(async () => ({
+      due: true,
+      completed: 1,
+      skipped: 0,
+    }));
+    const schedulerRunDue = mock(async () => schedulerResultFixture());
     registerBackgroundWakeRuntime({
       heartbeat: {
         nextRunAt: Date.now() + 5 * 60_000,
-        runOnce: heartbeatRunOnce,
+        runManagedWakeIfDue: heartbeatRunManaged,
       },
       scheduler: {
-        runOnce: schedulerRunOnce,
+        runOnce: mock(async () => 0),
+        runDueWorkOnce: schedulerRunDue,
       },
     });
     const handler = findHandler("drainDueBackgroundWake");
 
-    const response = (await handler({ body: drainBodyFixture() })) as {
-      counts: { heartbeat: number; scheduler: number; total: number };
-    };
+    const response = await handler({ body: drainBodyFixture() });
+    await flushBackgroundWakeDrainsForTest();
 
-    expect(heartbeatRunOnce).not.toHaveBeenCalled();
-    expect(schedulerRunOnce).toHaveBeenCalledTimes(1);
-    expect(response.counts).toEqual({
-      heartbeat: 0,
-      scheduler: 0,
-      total: 0,
+    expect(heartbeatRunManaged).not.toHaveBeenCalled();
+    expect(schedulerRunDue).toHaveBeenCalledTimes(1);
+    expect(response).toEqual(acceptedResponse());
+    expect(lastCompletionPayload()).toMatchObject({
+      leaseId: "lease-123",
+      status: "completed",
     });
   });
 
@@ -153,35 +259,549 @@ describe("background wake runtime routes", () => {
     });
     const schedulerRunOnce = mock(async () => {
       computedIntent = recomputedIntent;
-      return 1;
+      return schedulerResultFixture({ claimed: 1, completed: 1 });
     });
     registerBackgroundWakeRuntime({
       heartbeat: {
         nextRunAt: null,
-        runOnce: mock(async () => true),
+        runManagedWakeIfDue: mock(async () => ({
+          due: false,
+          completed: 0,
+          skipped: 0,
+        })),
       },
       scheduler: {
-        runOnce: schedulerRunOnce,
+        runOnce: mock(async () => 1),
+        runDueWorkOnce: schedulerRunOnce,
       },
     });
     const handler = findHandler("drainDueBackgroundWake");
 
-    const response = (await handler({ body: drainBodyFixture() })) as {
-      nextIntent: MockIntent | null;
-    };
+    const response = await handler({ body: drainBodyFixture() });
+    await flushBackgroundWakeDrainsForTest();
 
-    expect(response.nextIntent).toEqual(recomputedIntent);
+    expect(response).toEqual(acceptedResponse());
+    expect(lastCompletionPayload().nextIntent).toEqual(recomputedIntent);
+  });
+
+  test("drain-due runs scheduler backlog on heartbeat-only drains", async () => {
+    const dueAt = Date.now() - 1;
+    const recomputedIntent = intentFixture({
+      nextWakeAt: Date.now() + 60_000,
+      sourceGeneration: "after-heartbeat-drain",
+    });
+    computedIntent = intentFixture({
+      nextWakeAt: dueAt,
+      actualNextDueAt: dueAt,
+      reason: "heartbeat",
+    });
+    const heartbeatRunManaged = mock(async () => ({
+      due: true,
+      completed: 1,
+      skipped: 0,
+    }));
+    const schedulerRunDue = mock(async () => {
+      computedIntent = recomputedIntent;
+      return schedulerResultFixture({ claimed: 1, completed: 1 });
+    });
+    registerBackgroundWakeRuntime({
+      heartbeat: {
+        nextRunAt: dueAt,
+        runManagedWakeIfDue: heartbeatRunManaged,
+      },
+      scheduler: {
+        runOnce: mock(async () => 1),
+        runDueWorkOnce: schedulerRunDue,
+      },
+    });
+    const handler = findHandler("drainDueBackgroundWake");
+
+    const response = await handler({
+      body: drainBodyFixture({ reason: "heartbeat" }),
+    });
+    await flushBackgroundWakeDrainsForTest();
+
+    expect(heartbeatRunManaged).toHaveBeenCalledWith(
+      expect.objectContaining({
+        scheduledFor: NOW,
+      }),
+    );
+    expect(firstCallArg(heartbeatRunManaged)).not.toHaveProperty("assumeDue");
+    expect(schedulerRunDue).toHaveBeenCalledTimes(1);
+    expect(response).toEqual(acceptedResponse({ reason: "heartbeat" }));
+    expect(lastCompletionPayload()).toMatchObject({
+      leaseId: "lease-123",
+      status: "completed",
+      nextIntent: recomputedIntent,
+    });
+  });
+
+  test("drain-due does not rerun heartbeat for stale heartbeat drain requests", async () => {
+    const heartbeatRunManaged = mock(async () => ({
+      due: true,
+      completed: 1,
+      skipped: 0,
+    }));
+    const schedulerRunDue = mock(async () => schedulerResultFixture());
+    registerBackgroundWakeRuntime({
+      heartbeat: {
+        nextRunAt: Date.now() + 30 * 60_000,
+        runManagedWakeIfDue: heartbeatRunManaged,
+      },
+      scheduler: {
+        runOnce: mock(async () => 0),
+        runDueWorkOnce: schedulerRunDue,
+      },
+    });
+    const handler = findHandler("drainDueBackgroundWake");
+
+    const response = await handler({
+      body: drainBodyFixture({ reason: "heartbeat" }),
+    });
+    await flushBackgroundWakeDrainsForTest();
+
+    expect(heartbeatRunManaged).not.toHaveBeenCalled();
+    expect(schedulerRunDue).not.toHaveBeenCalled();
+    expect(response).toEqual(acceptedResponse({ reason: "heartbeat" }));
+    expect(lastCompletionPayload()).toMatchObject({
+      leaseId: "lease-123",
+      status: "completed",
+    });
+  });
+
+  test("drain-due runs heartbeat when the local intent is due at wake", async () => {
+    const dueAt = Date.now() - 1;
+    computedIntent = intentFixture({
+      nextWakeAt: dueAt,
+      actualNextDueAt: dueAt,
+      reason: "heartbeat",
+    });
+    const heartbeatRunManaged = mock(async () => ({
+      due: true,
+      completed: 1,
+      skipped: 0,
+    }));
+    registerBackgroundWakeRuntime({
+      heartbeat: {
+        nextRunAt: dueAt,
+        runManagedWakeIfDue: heartbeatRunManaged,
+      },
+      scheduler: {
+        runOnce: mock(async () => 0),
+        runDueWorkOnce: mock(async () => schedulerResultFixture()),
+      },
+    });
+    const handler = findHandler("drainDueBackgroundWake");
+
+    const response = await handler({
+      body: drainBodyFixture({ reason: "refresh" }),
+    });
+    await flushBackgroundWakeDrainsForTest();
+
+    expect(heartbeatRunManaged).toHaveBeenCalledWith(
+      expect.objectContaining({ scheduledFor: NOW }),
+    );
+    expect(firstCallArg(heartbeatRunManaged)).not.toHaveProperty("assumeDue");
+    expect(response).toEqual(acceptedResponse({ reason: "refresh" }));
+  });
+
+  test("drain-due runs scheduler when schedule work is due at wake", async () => {
+    const dueAt = Date.now() - 1;
+    computedIntent = intentFixture({
+      nextWakeAt: dueAt,
+      actualNextDueAt: dueAt,
+      reason: "schedule",
+    });
+    const schedulerRunDue = mock(async () =>
+      schedulerResultFixture({ claimed: 1, completed: 1 }),
+    );
+    registerBackgroundWakeRuntime({
+      heartbeat: {
+        nextRunAt: Date.now() + 60_000,
+        runManagedWakeIfDue: mock(async () => ({
+          due: false,
+          completed: 0,
+          skipped: 0,
+        })),
+      },
+      scheduler: {
+        runOnce: mock(async () => 1),
+        runDueWorkOnce: schedulerRunDue,
+      },
+    });
+    const handler = findHandler("drainDueBackgroundWake");
+
+    const response = await handler({
+      body: drainBodyFixture({ reason: "refresh" }),
+    });
+    await flushBackgroundWakeDrainsForTest();
+
+    expect(schedulerRunDue).toHaveBeenCalledTimes(1);
+    expect(response).toEqual(acceptedResponse({ reason: "refresh" }));
+    expect(lastCompletionPayload()).toMatchObject({
+      leaseId: "lease-123",
+      status: "completed",
+    });
+  });
+
+  test("drain-due runs due heartbeat even when schedule is the earliest intent source", async () => {
+    const now = Date.now();
+    const scheduleDueAt = now - 60_000;
+    const heartbeatDueAt = now - 1_000;
+    const recomputedIntent = intentFixture({
+      nextWakeAt: now + 60_000,
+      sourceGeneration: "after-mixed-backlog",
+    });
+    computedIntent = intentFixture({
+      nextWakeAt: scheduleDueAt,
+      actualNextDueAt: scheduleDueAt,
+      reason: "schedule",
+    });
+    const heartbeatRunManaged = mock(async () => ({
+      due: true,
+      completed: 1,
+      skipped: 0,
+    }));
+    const schedulerRunDue = mock(async () => {
+      computedIntent = recomputedIntent;
+      return schedulerResultFixture({ claimed: 1, completed: 1 });
+    });
+    registerBackgroundWakeRuntime({
+      heartbeat: {
+        nextRunAt: heartbeatDueAt,
+        runManagedWakeIfDue: heartbeatRunManaged,
+      },
+      scheduler: {
+        runOnce: mock(async () => 1),
+        runDueWorkOnce: schedulerRunDue,
+      },
+    });
+    const handler = findHandler("drainDueBackgroundWake");
+
+    const response = await handler({
+      body: drainBodyFixture({ reason: "schedule" }),
+    });
+    await flushBackgroundWakeDrainsForTest();
+
+    expect(heartbeatRunManaged).toHaveBeenCalledWith(
+      expect.objectContaining({ scheduledFor: NOW }),
+    );
+    expect(schedulerRunDue).toHaveBeenCalledTimes(1);
+    expect(response).toEqual(acceptedResponse({ reason: "schedule" }));
+    expect(lastCompletionPayload()).toMatchObject({
+      leaseId: "lease-123",
+      status: "completed",
+      nextIntent: recomputedIntent,
+    });
+  });
+
+  test("drain-due does no work for refresh-only wakes with no due intent", async () => {
+    const heartbeatRunManaged = mock(async () => ({
+      due: true,
+      completed: 1,
+      skipped: 0,
+    }));
+    const schedulerRunDue = mock(async () =>
+      schedulerResultFixture({ completed: 1 }),
+    );
+    registerBackgroundWakeRuntime({
+      heartbeat: {
+        nextRunAt: Date.now() + 60_000,
+        runManagedWakeIfDue: heartbeatRunManaged,
+      },
+      scheduler: {
+        runOnce: mock(async () => 1),
+        runDueWorkOnce: schedulerRunDue,
+      },
+    });
+    const handler = findHandler("drainDueBackgroundWake");
+
+    const response = await handler({
+      body: drainBodyFixture({ reason: "refresh" }),
+    });
+    await flushBackgroundWakeDrainsForTest();
+
+    expect(heartbeatRunManaged).not.toHaveBeenCalled();
+    expect(schedulerRunDue).not.toHaveBeenCalled();
+    expect(response).toEqual(acceptedResponse({ reason: "refresh" }));
+    expect(lastCompletionPayload()).toMatchObject({
+      leaseId: "lease-123",
+      status: "completed",
+    });
+  });
+
+  test("drain-due honors deadline exhaustion before starting heartbeat work", async () => {
+    const dueAt = Date.now() - 1;
+    computedIntent = intentFixture({
+      nextWakeAt: dueAt,
+      actualNextDueAt: dueAt,
+      reason: "mixed",
+    });
+    const heartbeatRunManaged = mock(async () => ({
+      due: true,
+      completed: 1,
+      skipped: 0,
+    }));
+    const schedulerRunDue = mock(async () =>
+      schedulerResultFixture({ skipped: 2, stillPending: 2 }),
+    );
+    registerBackgroundWakeRuntime({
+      heartbeat: {
+        nextRunAt: dueAt,
+        runManagedWakeIfDue: heartbeatRunManaged,
+      },
+      scheduler: {
+        runOnce: mock(async () => 0),
+        runDueWorkOnce: schedulerRunDue,
+      },
+    });
+    const handler = findHandler("drainDueBackgroundWake");
+
+    const response = await handler({
+      body: drainBodyFixture({
+        reason: "mixed",
+        deadlineAt: Date.now() + 1_000,
+      }),
+    });
+    await flushBackgroundWakeDrainsForTest();
+
+    expect(heartbeatRunManaged).not.toHaveBeenCalled();
+    expect(schedulerRunDue).toHaveBeenCalledWith(
+      expect.objectContaining({ minStartBudgetMs: 5_000 }),
+    );
+    expect(response).toMatchObject({
+      accepted: true,
+      reason: "mixed",
+    });
+    expect(lastCompletionPayload()).toMatchObject({
+      leaseId: "lease-123",
+      status: "completed",
+    });
+  });
+
+  test("drain-due completes expired without starting when the lease deadline already elapsed", async () => {
+    const dueAt = Date.now() - 1;
+    computedIntent = intentFixture({
+      nextWakeAt: dueAt,
+      actualNextDueAt: dueAt,
+      reason: "mixed",
+    });
+    const heartbeatRunManaged = mock(async () => ({
+      due: true,
+      completed: 1,
+      skipped: 0,
+    }));
+    const schedulerRunDue = mock(async () =>
+      schedulerResultFixture({ claimed: 1, completed: 1 }),
+    );
+    registerBackgroundWakeRuntime({
+      heartbeat: {
+        nextRunAt: dueAt,
+        runManagedWakeIfDue: heartbeatRunManaged,
+      },
+      scheduler: {
+        runOnce: mock(async () => 1),
+        runDueWorkOnce: schedulerRunDue,
+      },
+    });
+    const handler = findHandler("drainDueBackgroundWake");
+
+    const response = await handler({
+      body: drainBodyFixture({
+        reason: "mixed",
+        deadlineAt: Date.now() - 1,
+      }),
+    });
+    await flushBackgroundWakeDrainsForTest();
+
+    expect(heartbeatRunManaged).not.toHaveBeenCalled();
+    expect(schedulerRunDue).not.toHaveBeenCalled();
+    expect(response).toMatchObject({
+      accepted: true,
+      reason: "mixed",
+    });
+    expect(lastCompletionPayload()).toMatchObject({
+      leaseId: "lease-123",
+      status: "expired",
+      error: "background wake lease deadline elapsed before starting work",
+    });
+  });
+
+  test("drain-due keeps a lease active until slow work settles after the original deadline", async () => {
+    const schedulerResolvers: Array<() => void> = [];
+    const schedulerRunDue = mock(async () => {
+      await new Promise<void>((resolve) => schedulerResolvers.push(resolve));
+      return schedulerResultFixture({ claimed: 1, completed: 1 });
+    });
+    registerBackgroundWakeRuntime({
+      heartbeat: {
+        nextRunAt: Date.now() + 60_000,
+        runManagedWakeIfDue: mock(async () => ({
+          due: false,
+          completed: 0,
+          skipped: 0,
+        })),
+      },
+      scheduler: {
+        runOnce: mock(async () => 1),
+        runDueWorkOnce: schedulerRunDue,
+      },
+    });
+    const handler = findHandler("drainDueBackgroundWake");
+    const deadlineAt = Date.now() + 10;
+
+    const firstResponse = await handler({
+      body: drainBodyFixture({
+        leaseId: "lease-long-running",
+        deadlineAt,
+      }),
+    });
+
+    try {
+      await new Promise((resolve) => setTimeout(resolve, 25));
+      expect(mockCompleteBackgroundWakeLease).not.toHaveBeenCalled();
+
+      const duplicateResponse = await handler({
+        body: drainBodyFixture({
+          leaseId: "lease-long-running",
+          deadlineAt,
+        }),
+      });
+
+      expect(firstResponse).toMatchObject({
+        accepted: true,
+        leaseId: "lease-long-running",
+      });
+      expect(duplicateResponse).toMatchObject({
+        accepted: true,
+        leaseId: "lease-long-running",
+      });
+      expect(schedulerRunDue).toHaveBeenCalledTimes(1);
+    } finally {
+      for (const resolve of schedulerResolvers) resolve();
+    }
+
+    await flushBackgroundWakeDrainsForTest();
+    expect(schedulerRunDue).toHaveBeenCalledTimes(1);
+    expect(mockCompleteBackgroundWakeLease).toHaveBeenCalledTimes(1);
+    expect(lastCompletionPayload()).toMatchObject({
+      leaseId: "lease-long-running",
+      status: "completed",
+    });
+  });
+
+  test("drain-due runs scheduler work even when pending work remains", async () => {
+    const schedulerRunDue = mock(async () =>
+      schedulerResultFixture({ claimed: 1, completed: 1, stillPending: 1 }),
+    );
+    registerBackgroundWakeRuntime({
+      heartbeat: {
+        nextRunAt: Date.now() + 60_000,
+        runManagedWakeIfDue: mock(async () => ({
+          due: false,
+          completed: 0,
+          skipped: 0,
+        })),
+      },
+      scheduler: {
+        runOnce: mock(async () => 1),
+        runDueWorkOnce: schedulerRunDue,
+      },
+    });
+    const handler = findHandler("drainDueBackgroundWake");
+
+    const response = await handler({ body: drainBodyFixture() });
+    await flushBackgroundWakeDrainsForTest();
+
+    expect(schedulerRunDue).toHaveBeenCalledTimes(1);
+    expect(response).toEqual(acceptedResponse());
+    expect(lastCompletionPayload()).toMatchObject({
+      leaseId: "lease-123",
+      status: "completed",
+    });
+  });
+
+  test("drain-due completes failed when due work throws", async () => {
+    const schedulerRunDue = mock(async () => {
+      throw new Error("scheduler exploded");
+    });
+    registerBackgroundWakeRuntime({
+      heartbeat: {
+        nextRunAt: null,
+        runManagedWakeIfDue: mock(async () => ({
+          due: false,
+          completed: 0,
+          skipped: 0,
+        })),
+      },
+      scheduler: {
+        runOnce: mock(async () => 0),
+        runDueWorkOnce: schedulerRunDue,
+      },
+    });
+    const handler = findHandler("drainDueBackgroundWake");
+
+    const response = await handler({ body: drainBodyFixture() });
+    await flushBackgroundWakeDrainsForTest();
+
+    expect(response).toEqual(acceptedResponse());
+    expect(lastCompletionPayload()).toMatchObject({
+      leaseId: "lease-123",
+      status: "failed",
+      error: "scheduler exploded",
+    });
+  });
+
+  test("drain-due does not reclassify successful work when completion reporting fails", async () => {
+    const completionError = new Error("platform unavailable");
+    mockCompleteBackgroundWakeLease.mockImplementation(async (args) => {
+      if (args.status === "completed" || args.status === "expired") {
+        throw completionError;
+      }
+      return { status: "completed" as const, httpStatus: 200 };
+    });
+    const schedulerRunDue = mock(async () => schedulerResultFixture());
+    registerBackgroundWakeRuntime({
+      heartbeat: {
+        nextRunAt: null,
+        runManagedWakeIfDue: mock(async () => ({
+          due: false,
+          completed: 0,
+          skipped: 0,
+        })),
+      },
+      scheduler: {
+        runOnce: mock(async () => 0),
+        runDueWorkOnce: schedulerRunDue,
+      },
+    });
+    const handler = findHandler("drainDueBackgroundWake");
+
+    const response = await handler({ body: drainBodyFixture() });
+    await flushBackgroundWakeDrainsForTest();
+
+    expect(response).toEqual(acceptedResponse());
+    expect(mockCompleteBackgroundWakeLease).toHaveBeenCalledTimes(1);
+    expect(lastCompletionPayload()).toMatchObject({
+      leaseId: "lease-123",
+      status: "completed",
+    });
   });
 
   test("drain-due accepts vembda snake_case payload aliases", async () => {
-    const schedulerRunOnce = mock(async () => 0);
+    const schedulerRunDue = mock(async () => schedulerResultFixture());
     registerBackgroundWakeRuntime({
       heartbeat: {
         nextRunAt: null,
-        runOnce: mock(async () => false),
+        runManagedWakeIfDue: mock(async () => ({
+          due: false,
+          completed: 0,
+          skipped: 0,
+        })),
       },
       scheduler: {
-        runOnce: schedulerRunOnce,
+        runOnce: mock(async () => 0),
+        runDueWorkOnce: schedulerRunDue,
       },
     });
     const handler = findHandler("drainDueBackgroundWake");
@@ -195,24 +815,39 @@ describe("background wake runtime routes", () => {
         deadline_at: String(NOW + 30_000),
       },
     });
+    await flushBackgroundWakeDrainsForTest();
 
     expect(response).toMatchObject({
+      accepted: true,
       leaseId: "lease-snake",
       reason: "refresh",
       sourceGeneration: "generation-snake",
       startedAt: NOW,
       deadlineAt: NOW + 30_000,
     });
+    expect(lastCompletionPayload()).toMatchObject({
+      leaseId: "lease-snake",
+      status: "completed",
+    });
   });
 });
 
-function drainBodyFixture() {
+function drainBodyFixture(
+  overrides: Partial<{
+    leaseId: string;
+    reason: string;
+    sourceGeneration: string;
+    startedAt: number;
+    deadlineAt: number;
+  }> = {},
+) {
   return {
     leaseId: "lease-123",
     reason: "schedule",
     sourceGeneration: "source-generation",
     startedAt: NOW,
     deadlineAt: NOW + 30_000,
+    ...overrides,
   };
 }