@vellumai/assistant 0.8.5 → 0.8.6

package/src/background-wake/platform-client.test.ts

@@ -0,0 +1,308 @@
+import { beforeEach, describe, expect, mock, test } from "bun:test";
+
+import type { BackgroundWakeIntent } from "./next-wake.js";
+
+let mockPlatformAssistantId: string;
+let mockClientFetch: ReturnType<typeof mock>;
+let mockCreateClient: ReturnType<typeof mock>;
+
+mock.module("../platform/client.js", () => ({
+  VellumPlatformClient: {
+    create: () => mockCreateClient(),
+  },
+}));
+
+const {
+  clearBackgroundWakeIntent,
+  completeBackgroundWakeLease,
+  publishBackgroundWakeIntent,
+  renewBackgroundWakeLease,
+} = await import("./platform-client.js");
+
+const NOW = 1_800_000_000_000;
+
+describe("background wake platform client", () => {
+  beforeEach(() => {
+    mockPlatformAssistantId = "asst-123";
+    mockClientFetch = mock(async () => new Response('{"status":"stored"}'));
+    mockCreateClient = mock(async () => ({
+      platformAssistantId: mockPlatformAssistantId,
+      fetch: mockClientFetch,
+    }));
+  });
+
+  test("publishes derived wake intent fields to Django", async () => {
+    const result = await publishBackgroundWakeIntent(intentFixture());
+
+    expect(result).toEqual({ status: "published", httpStatus: 200 });
+    expect(mockClientFetch).toHaveBeenCalledTimes(1);
+    const [path, init] = mockClientFetch.mock.calls[0] as [string, RequestInit];
+    expect(path).toBe("/v1/assistants/asst-123/background-wake-intent/");
+    expect(init.method).toBe("PUT");
+    expect(new Headers(init.headers).get("Content-Type")).toBe(
+      "application/json",
+    );
+    expect(JSON.parse(String(init.body))).toEqual({
+      reason: "schedule",
+      source_generation: "bw1:opaque-hash",
+      computed_at: new Date(NOW - 1_000).toISOString(),
+      next_wake_at: new Date(NOW + 60_000).toISOString(),
+      actual_next_due_at: new Date(NOW + 60_000).toISOString(),
+      source_payload: {
+        heartbeat: null,
+        schedules: [
+          {
+            id: "schedule-1",
+            nextRunAt: NOW + 60_000,
+            mode: "wake",
+            createdBy: "defer",
+            status: "active",
+            updatedAt: NOW - 5_000,
+          },
+        ],
+      },
+    });
+  });
+
+  test("preserves stable string sourceGeneration across recomputes", async () => {
+    await publishBackgroundWakeIntent(
+      intentFixture({
+        computedAt: NOW - 1_000,
+        sourceGeneration: "bw1:stable-source-hash",
+      }),
+    );
+    await publishBackgroundWakeIntent(
+      intentFixture({
+        computedAt: NOW,
+        sourceGeneration: "bw1:stable-source-hash",
+      }),
+    );
+
+    const [, firstInit] = mockClientFetch.mock.calls[0] as [
+      string,
+      RequestInit,
+    ];
+    const [, secondInit] = mockClientFetch.mock.calls[1] as [
+      string,
+      RequestInit,
+    ];
+    expect(JSON.parse(String(firstInit.body))).toMatchObject({
+      source_generation: "bw1:stable-source-hash",
+      computed_at: new Date(NOW - 1_000).toISOString(),
+    });
+    expect(JSON.parse(String(secondInit.body))).toMatchObject({
+      source_generation: "bw1:stable-source-hash",
+      computed_at: new Date(NOW).toISOString(),
+    });
+  });
+
+  test("preserves numeric-looking sourceGeneration as a string", async () => {
+    await publishBackgroundWakeIntent(
+      intentFixture({ sourceGeneration: "42" }),
+    );
+
+    const [, init] = mockClientFetch.mock.calls[0] as [string, RequestInit];
+    expect(JSON.parse(String(init.body)).source_generation).toBe("42");
+  });
+
+  test("clears wake intent with the last computed snapshot when available", async () => {
+    mockClientFetch = mock(async () => new Response('{"status":"cleared"}'));
+    mockCreateClient = mock(async () => ({
+      platformAssistantId: mockPlatformAssistantId,
+      fetch: mockClientFetch,
+    }));
+
+    const result = await clearBackgroundWakeIntent(
+      intentFixture({ sourceGeneration: "bw1:opaque-hash" }),
+    );
+
+    expect(result).toEqual({ status: "cleared", httpStatus: 200 });
+    expect(mockClientFetch).toHaveBeenCalledTimes(1);
+    const [path, init] = mockClientFetch.mock.calls[0] as [string, RequestInit];
+    expect(path).toBe("/v1/assistants/asst-123/background-wake-intent/");
+    expect(init.method).toBe("DELETE");
+    expect(JSON.parse(String(init.body))).toEqual({
+      source_generation: "bw1:opaque-hash",
+      computed_at: new Date(NOW - 1_000).toISOString(),
+    });
+  });
+
+  test("clears wake intent without generation data when no snapshot is available", async () => {
+    await clearBackgroundWakeIntent(null);
+
+    const [, init] = mockClientFetch.mock.calls[0] as [string, RequestInit];
+    expect(init.method).toBe("DELETE");
+    expect(JSON.parse(String(init.body))).toEqual({});
+  });
+
+  test("renews background wake lease through Django", async () => {
+    const result = await renewBackgroundWakeLease("lease/with/slash");
+
+    expect(result).toEqual({ status: "renewed", httpStatus: 200 });
+    expect(mockClientFetch).toHaveBeenCalledTimes(1);
+    const [path, init] = mockClientFetch.mock.calls[0] as [string, RequestInit];
+    expect(path).toBe(
+      "/v1/assistants/asst-123/background-wake-leases/lease%2Fwith%2Fslash/renew/",
+    );
+    expect(init.method).toBe("POST");
+    expect(JSON.parse(String(init.body))).toEqual({});
+  });
+
+  test("completes background wake lease with the next intent", async () => {
+    const result = await completeBackgroundWakeLease({
+      leaseId: "lease-123",
+      status: "completed",
+      nextIntent: intentFixture({ sourceGeneration: "bw1:next-hash" }),
+    });
+
+    expect(result).toEqual({ status: "completed", httpStatus: 200 });
+    expect(mockClientFetch).toHaveBeenCalledTimes(1);
+    const [path, init] = mockClientFetch.mock.calls[0] as [string, RequestInit];
+    expect(path).toBe(
+      "/v1/assistants/asst-123/background-wake-leases/lease-123/complete/",
+    );
+    expect(init.method).toBe("POST");
+    expect(JSON.parse(String(init.body))).toEqual({
+      status: "completed",
+      next_intent: {
+        reason: "schedule",
+        source_generation: "bw1:next-hash",
+        computed_at: new Date(NOW - 1_000).toISOString(),
+        next_wake_at: new Date(NOW + 60_000).toISOString(),
+        actual_next_due_at: new Date(NOW + 60_000).toISOString(),
+        source_payload: {
+          heartbeat: null,
+          schedules: [
+            {
+              id: "schedule-1",
+              nextRunAt: NOW + 60_000,
+              mode: "wake",
+              createdBy: "defer",
+              status: "active",
+              updatedAt: NOW - 5_000,
+            },
+          ],
+        },
+      },
+    });
+  });
+
+  test("completes failed background wake lease with an error", async () => {
+    await completeBackgroundWakeLease({
+      leaseId: "lease-123",
+      status: "failed",
+      error: "scheduler exploded",
+    });
+
+    const [, init] = mockClientFetch.mock.calls[0] as [string, RequestInit];
+    expect(JSON.parse(String(init.body))).toEqual({
+      status: "failed",
+      error: "scheduler exploded",
+    });
+  });
+
+  test("completes background wake lease with an explicit null next intent", async () => {
+    await completeBackgroundWakeLease({
+      leaseId: "lease-123",
+      status: "completed",
+      nextIntent: null,
+    });
+
+    const [, init] = mockClientFetch.mock.calls[0] as [string, RequestInit];
+    expect(JSON.parse(String(init.body))).toEqual({
+      status: "completed",
+      next_intent: null,
+    });
+  });
+
+  test("skips when platform prerequisites are missing", async () => {
+    mockCreateClient = mock(async () => null);
+
+    await expect(publishBackgroundWakeIntent(intentFixture())).resolves.toEqual(
+      {
+        status: "skipped",
+        reason: "missing_platform_client",
+      },
+    );
+    expect(mockClientFetch).not.toHaveBeenCalled();
+
+    mockCreateClient = mock(async () => ({
+      platformAssistantId: "",
+      fetch: mockClientFetch,
+    }));
+
+    await expect(publishBackgroundWakeIntent(intentFixture())).resolves.toEqual(
+      {
+        status: "skipped",
+        reason: "missing_platform_assistant_id",
+      },
+    );
+    expect(mockClientFetch).not.toHaveBeenCalled();
+  });
+
+  test("treats Django disabled no-op responses as success", async () => {
+    mockClientFetch = mock(async () => new Response('{"status":"disabled"}'));
+    mockCreateClient = mock(async () => ({
+      platformAssistantId: mockPlatformAssistantId,
+      fetch: mockClientFetch,
+    }));
+
+    await expect(publishBackgroundWakeIntent(intentFixture())).resolves.toEqual(
+      {
+        status: "published",
+        httpStatus: 200,
+      },
+    );
+  });
+
+  test("throws on non-OK platform responses", async () => {
+    mockClientFetch = mock(
+      async () => new Response("bad gateway", { status: 502 }),
+    );
+    mockCreateClient = mock(async () => ({
+      platformAssistantId: mockPlatformAssistantId,
+      fetch: mockClientFetch,
+    }));
+
+    await expect(publishBackgroundWakeIntent(intentFixture())).rejects.toThrow(
+      "Failed to publish background wake intent: HTTP 502: bad gateway",
+    );
+    await expect(clearBackgroundWakeIntent(intentFixture())).rejects.toThrow(
+      "Failed to clear background wake intent: HTTP 502: bad gateway",
+    );
+    await expect(renewBackgroundWakeLease("lease-123")).rejects.toThrow(
+      "Failed to renew background wake lease: HTTP 502: bad gateway",
+    );
+    await expect(
+      completeBackgroundWakeLease({ leaseId: "lease-123", status: "failed" }),
+    ).rejects.toThrow(
+      "Failed to complete background wake lease: HTTP 502: bad gateway",
+    );
+  });
+});
+
+function intentFixture(
+  overrides: Partial<BackgroundWakeIntent> = {},
+): BackgroundWakeIntent {
+  return {
+    nextWakeAt: NOW + 60_000,
+    actualNextDueAt: NOW + 60_000,
+    reason: "schedule",
+    sourceGeneration: "bw1:opaque-hash",
+    computedAt: NOW - 1_000,
+    sourcePayload: {
+      heartbeat: null,
+      schedules: [
+        {
+          id: "schedule-1",
+          nextRunAt: NOW + 60_000,
+          mode: "wake",
+          createdBy: "defer",
+          status: "active",
+          updatedAt: NOW - 5_000,
+        },
+      ],
+    },
+    ...overrides,
+  };
+}

package/src/background-wake/platform-client.ts

@@ -0,0 +1,167 @@
+import { VellumPlatformClient } from "../platform/client.js";
+import type { BackgroundWakeIntent } from "./next-wake.js";
+
+export type BackgroundWakeIntentClientResult = {
+  status: "published" | "cleared" | "skipped";
+  httpStatus?: number;
+  reason?: "missing_platform_client" | "missing_platform_assistant_id";
+};
+
+export type BackgroundWakeLeaseClientResult = {
+  status: "renewed" | "completed" | "skipped";
+  httpStatus?: number;
+  reason?: "missing_platform_client" | "missing_platform_assistant_id";
+};
+
+type BackgroundWakeIntentSnapshot = Pick<
+  BackgroundWakeIntent,
+  "sourceGeneration" | "computedAt"
+> | null;
+
+export async function publishBackgroundWakeIntent(
+  intent: BackgroundWakeIntent,
+): Promise<BackgroundWakeIntentClientResult> {
+  const client = await VellumPlatformClient.create();
+  if (!client) return { status: "skipped", reason: "missing_platform_client" };
+  if (!client.platformAssistantId) {
+    return { status: "skipped", reason: "missing_platform_assistant_id" };
+  }
+
+  const response = await client.fetch(intentPath(client.platformAssistantId), {
+    method: "PUT",
+    headers: { "Content-Type": "application/json" },
+    body: JSON.stringify({
+      reason: intent.reason,
+      source_generation: intent.sourceGeneration,
+      computed_at: toIsoString(intent.computedAt),
+      next_wake_at: toIsoString(intent.nextWakeAt),
+      actual_next_due_at: toIsoString(intent.actualNextDueAt),
+      source_payload: intent.sourcePayload,
+    }),
+  });
+
+  await throwIfNotOk(response, "publish background wake intent");
+  return { status: "published", httpStatus: response.status };
+}
+
+export async function clearBackgroundWakeIntent(
+  intentSnapshot: BackgroundWakeIntentSnapshot = null,
+): Promise<BackgroundWakeIntentClientResult> {
+  const client = await VellumPlatformClient.create();
+  if (!client) return { status: "skipped", reason: "missing_platform_client" };
+  if (!client.platformAssistantId) {
+    return { status: "skipped", reason: "missing_platform_assistant_id" };
+  }
+
+  const body: Record<string, unknown> = {};
+  if (intentSnapshot) {
+    body.source_generation = intentSnapshot.sourceGeneration;
+    body.computed_at = toIsoString(intentSnapshot.computedAt);
+  }
+
+  const response = await client.fetch(intentPath(client.platformAssistantId), {
+    method: "DELETE",
+    headers: { "Content-Type": "application/json" },
+    body: JSON.stringify(body),
+  });
+
+  await throwIfNotOk(response, "clear background wake intent");
+  return { status: "cleared", httpStatus: response.status };
+}
+
+export async function renewBackgroundWakeLease(
+  leaseId: string,
+): Promise<BackgroundWakeLeaseClientResult> {
+  const client = await VellumPlatformClient.create();
+  if (!client) return { status: "skipped", reason: "missing_platform_client" };
+  if (!client.platformAssistantId) {
+    return { status: "skipped", reason: "missing_platform_assistant_id" };
+  }
+
+  const response = await client.fetch(
+    leasePath(client.platformAssistantId, leaseId, "renew"),
+    {
+      method: "POST",
+      headers: { "Content-Type": "application/json" },
+      body: JSON.stringify({}),
+    },
+  );
+
+  await throwIfNotOk(response, "renew background wake lease");
+  return { status: "renewed", httpStatus: response.status };
+}
+
+export async function completeBackgroundWakeLease(args: {
+  leaseId: string;
+  status: "completed" | "failed" | "expired";
+  error?: string;
+  nextIntent?: BackgroundWakeIntent | null;
+}): Promise<BackgroundWakeLeaseClientResult> {
+  const client = await VellumPlatformClient.create();
+  if (!client) return { status: "skipped", reason: "missing_platform_client" };
+  if (!client.platformAssistantId) {
+    return { status: "skipped", reason: "missing_platform_assistant_id" };
+  }
+
+  const body: Record<string, unknown> = {
+    status: args.status,
+  };
+  if (args.error) body.error = args.error;
+  if ("nextIntent" in args) {
+    body.next_intent = args.nextIntent
+      ? {
+          reason: args.nextIntent.reason,
+          source_generation: args.nextIntent.sourceGeneration,
+          computed_at: toIsoString(args.nextIntent.computedAt),
+          next_wake_at: toIsoString(args.nextIntent.nextWakeAt),
+          actual_next_due_at: toIsoString(args.nextIntent.actualNextDueAt),
+          source_payload: args.nextIntent.sourcePayload,
+        }
+      : null;
+  }
+
+  const response = await client.fetch(
+    leasePath(client.platformAssistantId, args.leaseId, "complete"),
+    {
+      method: "POST",
+      headers: { "Content-Type": "application/json" },
+      body: JSON.stringify(body),
+    },
+  );
+
+  await throwIfNotOk(response, "complete background wake lease");
+  return { status: "completed", httpStatus: response.status };
+}
+
+function intentPath(assistantId: string): string {
+  const encodedAssistantId = encodeURIComponent(assistantId);
+  return `/v1/assistants/${encodedAssistantId}/background-wake-intent/`;
+}
+
+function leasePath(
+  assistantId: string,
+  leaseId: string,
+  action: "renew" | "complete",
+): string {
+  const encodedAssistantId = encodeURIComponent(assistantId);
+  const encodedLeaseId = encodeURIComponent(leaseId);
+  return `/v1/assistants/${encodedAssistantId}/background-wake-leases/${encodedLeaseId}/${action}/`;
+}
+
+function toIsoString(timestampMs: number): string {
+  return new Date(timestampMs).toISOString();
+}
+
+async function throwIfNotOk(response: Response, action: string): Promise<void> {
+  if (response.ok) return;
+
+  let body = "";
+  try {
+    body = await response.text();
+  } catch {
+    body = "";
+  }
+
+  const detail = body ? `: ${body}` : "";
+  throw new Error(`Failed to ${action}: HTTP ${response.status}${detail}`);
+}

package/src/background-wake/publisher.ts

@@ -0,0 +1,91 @@
+import { getLogger } from "../util/logger.js";
+import {
+  type BackgroundWakeIntent,
+  computeNextBackgroundWakeIntent,
+} from "./next-wake.js";
+import {
+  clearBackgroundWakeIntent,
+  publishBackgroundWakeIntent,
+} from "./platform-client.js";
+
+const log = getLogger("background-wake-publisher");
+const REFRESH_DEBOUNCE_MS = 250;
+
+let refreshTimer: ReturnType<typeof setTimeout> | null = null;
+let pendingReason: string | null = null;
+let inFlightRefresh: Promise<void> | null = null;
+let lastIntentSnapshot: BackgroundWakeIntent | null = null;
+
+export function refreshBackgroundWakeIntent(reason: string): void {
+  pendingReason = reason;
+  schedulePendingRefresh();
+}
+
+function schedulePendingRefresh(): void {
+  if (refreshTimer || inFlightRefresh) return;
+  refreshTimer = setTimeout(runPendingRefresh, REFRESH_DEBOUNCE_MS);
+  refreshTimer.unref?.();
+}
+
+function runPendingRefresh(): void {
+  if (refreshTimer) {
+    clearTimeout(refreshTimer);
+    refreshTimer = null;
+  }
+
+  const reason = pendingReason ?? "unspecified";
+  pendingReason = null;
+  const refresh = runRefresh(reason).finally(() => {
+    if (inFlightRefresh === refresh) {
+      inFlightRefresh = null;
+    }
+    if (pendingReason) {
+      schedulePendingRefresh();
+    }
+  });
+  inFlightRefresh = refresh;
+}
+
+async function runRefresh(reason: string): Promise<void> {
+  try {
+    const intent = computeNextBackgroundWakeIntent();
+    if (intent) {
+      const result = await publishBackgroundWakeIntent(intent);
+      lastIntentSnapshot = intent;
+      log.debug({ reason, result }, "Background wake intent refreshed");
+      return;
+    }
+
+    const result = await clearBackgroundWakeIntent(lastIntentSnapshot);
+    if (result.status === "cleared") {
+      lastIntentSnapshot = null;
+    }
+    log.debug({ reason, result }, "Background wake intent cleared");
+  } catch (err) {
+    log.warn({ err, reason }, "Failed to refresh background wake intent");
+  }
+}
+
+/** @internal Test helper. */
+export async function flushBackgroundWakeIntentRefreshForTest(): Promise<void> {
+  if (refreshTimer) {
+    runPendingRefresh();
+  }
+  if (inFlightRefresh) {
+    await inFlightRefresh;
+  }
+  if (refreshTimer) {
+    await flushBackgroundWakeIntentRefreshForTest();
+  }
+}
+
+/** @internal Test helper. */
+export function resetBackgroundWakeIntentPublisherForTest(): void {
+  if (refreshTimer) {
+    clearTimeout(refreshTimer);
+    refreshTimer = null;
+  }
+  pendingReason = null;
+  inFlightRefresh = null;
+  lastIntentSnapshot = null;
+}

package/src/background-wake/runtime-registry.ts

@@ -2,8 +2,8 @@ import type { HeartbeatService } from "../heartbeat/heartbeat-service.js";
 import type { SchedulerHandle } from "../schedule/scheduler.js";
 
 export interface BackgroundWakeRuntime {
-  scheduler: Pick<SchedulerHandle, "runOnce">;
-  heartbeat: Pick<HeartbeatService, "nextRunAt" | "runOnce">;
+  scheduler: Pick<SchedulerHandle, "runOnce" | "runDueWorkOnce">;
+  heartbeat: Pick<HeartbeatService, "nextRunAt" | "runManagedWakeIfDue">;
 }
 
 let runtime: BackgroundWakeRuntime | null = null;