bopodev-api 0.1.27 → 0.1.29

package/src/server.ts

@@ -1,11 +1,5 @@
 import { createServer } from "node:http";
-import { dirname, resolve } from "node:path";
-import { fileURLToPath } from "node:url";
-import { sql } from "drizzle-orm";
-import { config as loadDotenv } from "dotenv";
-import { bootstrapDatabase, listCompanies, resolveDefaultDbPath } from "bopodev-db";
-import { checkRuntimeCommandHealth } from "bopodev-agent-sdk";
-import type { RuntimeCommandHealth } from "bopodev-agent-sdk";
+import { listCompanies } from "bopodev-db";
 import { createApp } from "./app";
 import { loadGovernanceRealtimeSnapshot } from "./realtime/governance";
 import { loadOfficeSpaceRealtimeSnapshot } from "./realtime/office-space";
@@ -13,7 +7,6 @@ import { loadHeartbeatRunsRealtimeSnapshot } from "./realtime/heartbeat-runs";
 import { loadAttentionRealtimeSnapshot } from "./realtime/attention";
 import { attachRealtimeHub } from "./realtime/hub";
 import {
-  isAuthenticatedMode,
   resolveAllowedHostnames,
   resolveAllowedOrigins,
   resolveDeploymentMode,
@@ -22,6 +15,17 @@ import {
 import { ensureBuiltinPluginsRegistered } from "./services/plugin-runtime";
 import { ensureBuiltinTemplatesRegistered } from "./services/template-catalog";
 import { createHeartbeatScheduler } from "./worker/scheduler";
+import { bootstrapDatabaseWithStartupLogging } from "./startup/database";
+import { validateDeploymentConfiguration } from "./startup/deployment-validation";
+import { loadApiEnv, normalizeOptionalDbPath } from "./startup/env";
+import {
+  buildGetRuntimeHealth,
+  hasCodexAgentsConfigured,
+  hasOpenCodeAgentsConfigured,
+  runStartupRuntimePreflights
+} from "./startup/runtime-health";
+import { resolveSchedulerCompanyId, shouldStartScheduler } from "./startup/scheduler-config";
+import { attachGracefulShutdownHandlers } from "./shutdown/graceful-shutdown";
 
 loadApiEnv();
 
@@ -33,28 +37,7 @@ async function main() {
   validateDeploymentConfiguration(deploymentMode, allowedOrigins, allowedHostnames, publicBaseUrl);
   const dbPath = normalizeOptionalDbPath(process.env.BOPO_DB_PATH);
   const port = Number(process.env.PORT ?? 4020);
-  const effectiveDbPath = dbPath ?? resolveDefaultDbPath();
-  let db: Awaited<ReturnType<typeof bootstrapDatabase>>["db"];
-  let dbClient: Awaited<ReturnType<typeof bootstrapDatabase>>["client"];
-  try {
-    const boot = await bootstrapDatabase(dbPath);
-    db = boot.db;
-    dbClient = boot.client;
-  } catch (error) {
-    if (isProbablyPgliteWasmAbort(error)) {
-      // eslint-disable-next-line no-console
-      console.error(
-        "[startup] PGlite (embedded Postgres) failed during database bootstrap. This is unrelated to Codex or heartbeat prompt settings."
-      );
-      // eslint-disable-next-line no-console
-      console.error(`[startup] Data path in use: ${effectiveDbPath}`);
-      // eslint-disable-next-line no-console
-      console.error(
-        "[startup] Recovery: stop all API/node processes using this DB, back up the path above, delete the file/dir, then restart (schema will be recreated). Or set BOPO_DB_PATH to a fresh path. See docs/operations/troubleshooting.md (PGlite)."
-      );
-    }
-    throw error;
-  }
+  const { db, client: dbClient } = await bootstrapDatabaseWithStartupLogging(dbPath);
   const existingCompanies = await listCompanies(db);
   await ensureBuiltinPluginsRegistered(
     db,
@@ -74,54 +57,20 @@ async function main() {
   const openCodeHealthRequired =
     !skipOpenCodePreflight &&
     (process.env.BOPO_REQUIRE_OPENCODE_HEALTH === "1" || (await hasOpenCodeAgentsConfigured(db)));
-  const getRuntimeHealth = async () => {
-    const codex = codexHealthRequired
-      ? await checkRuntimeCommandHealth(codexCommand, {
-          timeoutMs: 5_000
-        })
-      : {
-          command: codexCommand,
-          available: skipCodexPreflight ? false : true,
-          exitCode: null,
-          elapsedMs: 0,
-          error: skipCodexPreflight
-            ? "Skipped by configuration: BOPO_SKIP_CODEX_PREFLIGHT=1."
-            : "Skipped: no Codex agents configured."
-        };
-    const opencode = openCodeHealthRequired
-      ? await checkRuntimeCommandHealth(openCodeCommand, {
-          timeoutMs: 5_000
-        })
-      : {
-          command: openCodeCommand,
-          available: skipOpenCodePreflight ? false : true,
-          exitCode: null,
-          elapsedMs: 0,
-          error: skipOpenCodePreflight
-            ? "Skipped by configuration: BOPO_SKIP_OPENCODE_PREFLIGHT=1."
-            : "Skipped: no OpenCode agents configured."
-        };
-    return {
-      codex,
-      opencode
-    };
-  };
-  if (codexHealthRequired) {
-    const startupCodexHealth = await checkRuntimeCommandHealth(codexCommand, {
-      timeoutMs: 5_000
-    });
-    if (!startupCodexHealth.available) {
-      emitCodexPreflightWarning(startupCodexHealth);
-    }
-  }
-  if (openCodeHealthRequired) {
-    const startupOpenCodeHealth = await checkRuntimeCommandHealth(openCodeCommand, {
-      timeoutMs: 5_000
-    });
-    if (!startupOpenCodeHealth.available) {
-      emitOpenCodePreflightWarning(startupOpenCodeHealth);
-    }
-  }
+  const getRuntimeHealth = buildGetRuntimeHealth({
+    codexCommand,
+    openCodeCommand,
+    skipCodexPreflight,
+    skipOpenCodePreflight,
+    codexHealthRequired,
+    openCodeHealthRequired
+  });
+  await runStartupRuntimePreflights({
+    codexHealthRequired,
+    openCodeHealthRequired,
+    codexCommand,
+    openCodeCommand
+  });
 
   const server = createServer();
   const realtimeHub = attachRealtimeHub(server, {
@@ -141,210 +90,20 @@ async function main() {
 
   const defaultCompanyId = process.env.BOPO_DEFAULT_COMPANY_ID;
   const schedulerCompanyId = await resolveSchedulerCompanyId(db, defaultCompanyId ?? null);
-  let stopScheduler: (() => void) | undefined;
+  let scheduler: ReturnType<typeof createHeartbeatScheduler> | undefined;
   if (schedulerCompanyId && shouldStartScheduler()) {
-    stopScheduler = createHeartbeatScheduler(db, schedulerCompanyId, realtimeHub);
+    scheduler = createHeartbeatScheduler(db, schedulerCompanyId, realtimeHub);
   } else if (schedulerCompanyId) {
     // eslint-disable-next-line no-console
     console.log("[startup] Scheduler disabled for this instance (BOPO_SCHEDULER_ROLE is follower/off).");
   }
 
-  let shutdownInFlight: Promise<void> | null = null;
-  function shutdown(signal: string) {
-    shutdownInFlight ??= (async () => {
-      // eslint-disable-next-line no-console
-      console.log(`[shutdown] ${signal} — closing realtime, HTTP server, and embedded DB…`);
-      stopScheduler?.();
-      try {
-        await realtimeHub.close();
-      } catch (closeError) {
-        // eslint-disable-next-line no-console
-        console.error("[shutdown] realtime hub close error", closeError);
-      }
-      await new Promise<void>((resolve, reject) => {
-        server.close((err) => {
-          if (err) {
-            reject(err);
-            return;
-          }
-          resolve();
-        });
-      });
-      try {
-        await closePgliteClient(dbClient);
-      } catch (closeDbError) {
-        // eslint-disable-next-line no-console
-        console.error("[shutdown] PGlite close error", closeDbError);
-      }
-      // eslint-disable-next-line no-console
-      console.log("[shutdown] clean exit");
-      process.exit(0);
-    })().catch((error) => {
-      // eslint-disable-next-line no-console
-      console.error("[shutdown] failed", error);
-      process.exit(1);
-    });
-    return shutdownInFlight;
-  }
-
-  process.once("SIGINT", () => void shutdown("SIGINT"));
-  process.once("SIGTERM", () => void shutdown("SIGTERM"));
+  attachGracefulShutdownHandlers({
+    server,
+    realtimeHub,
+    dbClient,
+    scheduler
+  });
 }
 
 void main();
-
-async function closePgliteClient(client: unknown) {
-  if (!client || typeof client !== "object") {
-    return;
-  }
-  const closeFn = (client as { close?: unknown }).close;
-  if (typeof closeFn !== "function") {
-    return;
-  }
-  await (closeFn as () => Promise<void>)();
-}
-
-async function hasCodexAgentsConfigured(db: Awaited<ReturnType<typeof bootstrapDatabase>>["db"]) {
-  const result = await db.execute(sql`
-    SELECT id
-    FROM agents
-    WHERE provider_type = 'codex'
-    LIMIT 1
-  `);
-  return (result.rows ?? []).length > 0;
-}
-
-async function hasOpenCodeAgentsConfigured(db: Awaited<ReturnType<typeof bootstrapDatabase>>["db"]) {
-  const result = await db.execute(sql`
-    SELECT id
-    FROM agents
-    WHERE provider_type = 'opencode'
-    LIMIT 1
-  `);
-  return (result.rows ?? []).length > 0;
-}
-
-async function resolveSchedulerCompanyId(
-  db: Awaited<ReturnType<typeof bootstrapDatabase>>["db"],
-  configuredCompanyId: string | null
-) {
-  if (configuredCompanyId) {
-    const configured = await db.execute(sql`
-      SELECT id
-      FROM companies
-      WHERE id = ${configuredCompanyId}
-      LIMIT 1
-    `);
-    if ((configured.rows ?? []).length > 0) {
-      return configuredCompanyId;
-    }
-    // eslint-disable-next-line no-console
-    console.warn(`[startup] BOPO_DEFAULT_COMPANY_ID='${configuredCompanyId}' was not found; using first available company.`);
-  }
-
-  const fallback = await db.execute(sql`
-    SELECT id
-    FROM companies
-    ORDER BY created_at ASC
-    LIMIT 1
-  `);
-  const id = fallback.rows?.[0]?.id;
-  return typeof id === "string" && id.length > 0 ? id : null;
-}
-
-function emitCodexPreflightWarning(health: RuntimeCommandHealth) {
-  const red = process.stderr.isTTY ? "\x1b[31m" : "";
-  const yellow = process.stderr.isTTY ? "\x1b[33m" : "";
-  const reset = process.stderr.isTTY ? "\x1b[0m" : "";
-  const symbol = `${red}✖${reset}`;
-  process.stderr.write(
-    `${symbol} ${yellow}Codex preflight failed${reset}: command '${health.command}' is unavailable.\n`
-  );
-  process.stderr.write(`  Install Codex CLI or set BOPO_SKIP_CODEX_PREFLIGHT=1 for local dev.\n`);
-  if (process.env.BOPO_VERBOSE_STARTUP_WARNINGS === "1") {
-    process.stderr.write(`  Details: ${JSON.stringify(health)}\n`);
-  }
-}
-
-function emitOpenCodePreflightWarning(health: RuntimeCommandHealth) {
-  const red = process.stderr.isTTY ? "\x1b[31m" : "";
-  const yellow = process.stderr.isTTY ? "\x1b[33m" : "";
-  const reset = process.stderr.isTTY ? "\x1b[0m" : "";
-  const symbol = `${red}✖${reset}`;
-  process.stderr.write(
-    `${symbol} ${yellow}OpenCode preflight failed${reset}: command '${health.command}' is unavailable.\n`
-  );
-  process.stderr.write(`  Install OpenCode CLI or set BOPO_SKIP_OPENCODE_PREFLIGHT=1 for local dev.\n`);
-  if (process.env.BOPO_VERBOSE_STARTUP_WARNINGS === "1") {
-    process.stderr.write(`  Details: ${JSON.stringify(health)}\n`);
-  }
-}
-
-function validateDeploymentConfiguration(
-  deploymentMode: ReturnType<typeof resolveDeploymentMode>,
-  allowedOrigins: string[],
-  allowedHostnames: string[],
-  publicBaseUrl: URL | null
-) {
-  if (deploymentMode === "authenticated_public" && !publicBaseUrl) {
-    throw new Error("BOPO_PUBLIC_BASE_URL is required in authenticated_public mode.");
-  }
-  if (isAuthenticatedMode(deploymentMode) && process.env.BOPO_AUTH_TOKEN_SECRET?.trim() === "") {
-    throw new Error("BOPO_AUTH_TOKEN_SECRET must not be empty when set.");
-  }
-  if (isAuthenticatedMode(deploymentMode) && !process.env.BOPO_AUTH_TOKEN_SECRET?.trim()) {
-    // eslint-disable-next-line no-console
-    console.warn(
-      "[startup] BOPO_AUTH_TOKEN_SECRET is not set. Authenticated modes will require BOPO_TRUST_ACTOR_HEADERS=1 behind a trusted proxy."
-    );
-  }
-  if (isAuthenticatedMode(deploymentMode) && process.env.BOPO_TRUST_ACTOR_HEADERS !== "1" && !process.env.BOPO_AUTH_TOKEN_SECRET?.trim()) {
-    throw new Error(
-      "Authenticated mode requires either BOPO_AUTH_TOKEN_SECRET (token identity) or BOPO_TRUST_ACTOR_HEADERS=1 (trusted proxy headers)."
-    );
-  }
-  if (isAuthenticatedMode(deploymentMode) && process.env.BOPO_ALLOW_LOCAL_BOARD_FALLBACK === "1") {
-    throw new Error("BOPO_ALLOW_LOCAL_BOARD_FALLBACK cannot be enabled in authenticated modes.");
-  }
-  // eslint-disable-next-line no-console
-  console.log(
-    `[startup] Deployment config: mode=${deploymentMode} origins=${allowedOrigins.join(",")} hosts=${allowedHostnames.join(",")}`
-  );
-}
-
-function shouldStartScheduler() {
-  const rawRole = (process.env.BOPO_SCHEDULER_ROLE ?? "auto").trim().toLowerCase();
-  if (rawRole === "off" || rawRole === "follower") {
-    return false;
-  }
-  if (rawRole === "leader" || rawRole === "auto") {
-    return true;
-  }
-  throw new Error(`Invalid BOPO_SCHEDULER_ROLE '${rawRole}'. Expected one of: auto, leader, follower, off.`);
-}
-
-function loadApiEnv() {
-  const sourceDir = dirname(fileURLToPath(import.meta.url));
-  const repoRoot = resolve(sourceDir, "../../../");
-  const candidates = [resolve(repoRoot, ".env.local"), resolve(repoRoot, ".env")];
-  for (const path of candidates) {
-    loadDotenv({ path, override: false, quiet: true });
-  }
-}
-
-function normalizeOptionalDbPath(value: string | undefined) {
-  const normalized = value?.trim();
-  return normalized && normalized.length > 0 ? normalized : undefined;
-}
-
-function isProbablyPgliteWasmAbort(error: unknown): boolean {
-  const message = error instanceof Error ? error.message : String(error);
-  const cause = error instanceof Error ? error.cause : undefined;
-  const causeMessage = cause instanceof Error ? cause.message : String(cause ?? "");
-  return (
-    message.includes("Aborted") ||
-    causeMessage.includes("Aborted") ||
-    message.includes("pglite") ||
-    causeMessage.includes("RuntimeError")
-  );
-}

package/src/services/attention-service.ts

@@ -1,7 +1,10 @@
-import { and, desc, eq, like } from "drizzle-orm";
 import {
+  and,
   issueComments,
   issues,
+  desc,
+  eq,
+  like,
   listApprovalRequests,
   listAttentionInboxStates,
   listHeartbeatRuns,

package/src/services/budget-service.ts

@@ -1,6 +1,5 @@
-import { and, eq, sql } from "drizzle-orm";
 import type { BopoDb } from "bopodev-db";
-import { agents, projects } from "bopodev-db";
+import { agents, and, eq, projects, sql } from "bopodev-db";
 
 export interface BudgetCheckResult {
   allowed: boolean;

package/src/services/comment-recipient-dispatch-service.ts

@@ -1,21 +1,38 @@
-import { and, desc, eq, inArray, like } from "drizzle-orm";
 import {
+  and,
   agents,
+  desc,
+  eq,
+  inArray,
   issueComments,
+  like,
   updateIssueCommentRecipients,
   type BopoDb
 } from "bopodev-db";
+import { createDrainableWorkTracker } from "../lib/drainable-work";
 import { parseIssueCommentRecipients, type PersistedCommentRecipient } from "../lib/comment-recipients";
 import type { RealtimeHub } from "../realtime/hub";
 import { enqueueHeartbeatQueueJob, triggerHeartbeatQueueWorker } from "./heartbeat-queue-service";
 
 const COMMENT_DISPATCH_SWEEP_LIMIT = 100;
 const activeCompanyDispatchRuns = new Set<string>();
+const commentDispatchTracker = createDrainableWorkTracker();
 
 export async function runIssueCommentDispatchSweep(
   db: BopoDb,
   companyId: string,
   options?: { requestId?: string; realtimeHub?: RealtimeHub; limit?: number }
+) {
+  if (commentDispatchTracker.isShuttingDown()) {
+    return;
+  }
+  return commentDispatchTracker.track(runIssueCommentDispatchSweepInternal(db, companyId, options));
+}
+
+async function runIssueCommentDispatchSweepInternal(
+  db: BopoDb,
+  companyId: string,
+  options?: { requestId?: string; realtimeHub?: RealtimeHub; limit?: number }
 ) {
   const rows = await db
     .select({
@@ -34,6 +51,9 @@ export async function runIssueCommentDispatchSweep(
     .limit(options?.limit ?? COMMENT_DISPATCH_SWEEP_LIMIT);
 
   for (const row of rows) {
+    if (commentDispatchTracker.isShuttingDown()) {
+      return;
+    }
     const recipients = parseIssueCommentRecipients(row.recipientsJson);
     if (!recipients.some((recipient) => recipient.deliveryStatus === "pending")) {
       continue;
@@ -60,11 +80,15 @@ export function triggerIssueCommentDispatchWorker(
   companyId: string,
   options?: { requestId?: string; realtimeHub?: RealtimeHub; limit?: number }
 ) {
-  if (activeCompanyDispatchRuns.has(companyId)) {
+  if (commentDispatchTracker.isShuttingDown() || activeCompanyDispatchRuns.has(companyId)) {
     return;
   }
   activeCompanyDispatchRuns.add(companyId);
   queueMicrotask(() => {
+    if (commentDispatchTracker.isShuttingDown()) {
+      activeCompanyDispatchRuns.delete(companyId);
+      return;
+    }
     void runIssueCommentDispatchSweep(db, companyId, options)
       .catch((error) => {
         // eslint-disable-next-line no-console
@@ -156,3 +180,16 @@ async function dispatchCommentRecipients(
   return dispatchedRecipients;
 }
 
+export function beginIssueCommentDispatchShutdown() {
+  commentDispatchTracker.beginShutdown();
+}
+
+export async function waitForIssueCommentDispatchDrain() {
+  await commentDispatchTracker.drain();
+}
+
+export function resetIssueCommentDispatchShutdownForTests() {
+  activeCompanyDispatchRuns.clear();
+  commentDispatchTracker.resetForTests();
+}
+

package/src/services/company-export-service.ts

@@ -0,0 +1,63 @@
+import type { BopoDb } from "bopodev-db";
+import { getCompany, listAgents, listGoals, listIssues, listProjects } from "bopodev-db";
+
+function serializeValue(value: unknown): unknown {
+  if (value instanceof Date) {
+    return value.toISOString();
+  }
+  if (value !== null && typeof value === "object" && !Array.isArray(value)) {
+    return serializeRow(value as Record<string, unknown>);
+  }
+  if (Array.isArray(value)) {
+    return value.map((entry) => serializeValue(entry));
+  }
+  return value;
+}
+
+function serializeRow<T extends Record<string, unknown>>(row: T): Record<string, unknown> {
+  const out: Record<string, unknown> = {};
+  for (const [key, val] of Object.entries(row)) {
+    out[key] = serializeValue(val);
+  }
+  return out;
+}
+
+/** Strip agent fields that commonly hold secrets or session state for shareable exports. */
+function scrubAgentRow(agent: Record<string, unknown>) {
+  const base = serializeRow(agent);
+  return {
+    ...base,
+    runtimeEnvJson: "{}",
+    bootstrapPrompt: null,
+    stateBlob: "{}",
+    runtimeCommand: null,
+    runtimeArgsJson: "[]"
+  };
+}
+
+/**
+ * Portable company snapshot for backup, templates, or handoff (secrets and agent session state redacted).
+ */
+export async function buildCompanyPortabilityExport(db: BopoDb, companyId: string) {
+  const company = await getCompany(db, companyId);
+  if (!company) {
+    return null;
+  }
+
+  const [projects, goals, issues, agents] = await Promise.all([
+    listProjects(db, companyId),
+    listGoals(db, companyId),
+    listIssues(db, companyId),
+    listAgents(db, companyId)
+  ]);
+
+  return {
+    exportedAt: new Date().toISOString(),
+    formatVersion: 1,
+    company: serializeRow(company as unknown as Record<string, unknown>),
+    projects: projects.map((p) => serializeRow(p as unknown as Record<string, unknown>)),
+    goals: goals.map((g) => serializeRow(g as unknown as Record<string, unknown>)),
+    issues: issues.map((issue) => serializeRow(issue as unknown as Record<string, unknown>)),
+    agents: agents.map((agent) => scrubAgentRow(agent as unknown as Record<string, unknown>))
+  };
+}

package/src/services/governance-service.ts

@@ -1,4 +1,3 @@
-import { and, eq } from "drizzle-orm";
 import { mkdir } from "node:fs/promises";
 import { z } from "zod";
 import {
@@ -10,6 +9,7 @@ import {
 } from "bopodev-contracts";
 import type { BopoDb } from "bopodev-db";
 import {
+  and,
   approvalRequests,
   agents,
   appendAuditEvent,
@@ -26,6 +26,7 @@ import {
   listProjectWorkspaces,
   listProjects,
   projects,
+  eq,
   updateProjectWorkspace,
   updatePluginConfig
 } from "bopodev-db";
@@ -93,6 +94,7 @@ const hireAgentPayloadSchema = AgentCreateRequestSchema.extend({
 const activateGoalPayloadSchema = z.object({
   projectId: z.string().optional(),
   parentGoalId: z.string().optional(),
+  ownerAgentId: z.string().optional(),
   level: z.enum(["company", "project", "agent"]),
   title: z.string().min(1),
   description: z.string().optional()
@@ -371,6 +373,7 @@ async function applyApprovalAction(db: BopoDb, companyId: string, action: string
       companyId,
       projectId: parsed.data.projectId,
       parentGoalId: parsed.data.parentGoalId,
+      ownerAgentId: parsed.data.ownerAgentId,
       level: parsed.data.level,
       title: parsed.data.title,
       description: parsed.data.description
@@ -783,7 +786,8 @@ function buildAgentStartupTaskBody(companyId: string, agentId: string) {
     `   - \`${agentOperatingFolder}/HEARTBEAT.md\``,
     `   - \`${agentOperatingFolder}/SOUL.md\``,
     `   - \`${agentOperatingFolder}/TOOLS.md\``,
-    `3. Update your own agent runtime config via \`PUT /agents/:agentId\` and set \`runtimeConfig.bootstrapPrompt\` to reference \`${agentOperatingFolder}/AGENTS.md\` as your primary guide.`,
+    "3. Each heartbeat already includes your operating directory via `$BOPODEV_AGENT_OPERATING_DIR` and directs you to AGENTS.md and related files there when they exist.",
+    "   You do not need to save file paths into `bootstrapPrompt` for operating docs—use `PUT /agents/:agentId` `bootstrapPrompt` only for optional extra standing instructions.",
     "4. Post an issue comment summarizing completed setup artifacts.",
     "",
     "Safety checks:",

package/src/services/heartbeat-queue-service.ts

@@ -1,6 +1,7 @@
-import { and, eq } from "drizzle-orm";
 import {
+  and,
   cancelHeartbeatJob,
+  eq,
   getHeartbeatRun,
   claimNextHeartbeatJob,
   enqueueHeartbeatJob,
@@ -11,6 +12,7 @@ import {
   updateIssueCommentRecipients,
   type BopoDb
 } from "bopodev-db";
+import { createDrainableWorkTracker } from "../lib/drainable-work";
 import { parseIssueCommentRecipients } from "../lib/comment-recipients";
 import type { RealtimeHub } from "../realtime/hub";
 import { runHeartbeatForAgent } from "./heartbeat-service";
@@ -30,6 +32,7 @@ type QueueJobPayload = {
 };
 
 const activeCompanyQueueWorkers = new Set<string>();
+const queueWorkTracker = createDrainableWorkTracker();
 
 export async function enqueueHeartbeatQueueJob(
   db: BopoDb,
@@ -61,11 +64,15 @@ export function triggerHeartbeatQueueWorker(
   companyId: string,
   options?: { requestId?: string; realtimeHub?: RealtimeHub; maxJobsPerSweep?: number }
 ) {
-  if (activeCompanyQueueWorkers.has(companyId)) {
+  if (queueWorkTracker.isShuttingDown() || activeCompanyQueueWorkers.has(companyId)) {
     return;
   }
   activeCompanyQueueWorkers.add(companyId);
   queueMicrotask(() => {
+    if (queueWorkTracker.isShuttingDown()) {
+      activeCompanyQueueWorkers.delete(companyId);
+      return;
+    }
     void runHeartbeatQueueSweep(db, companyId, options)
       .catch((error) => {
         // eslint-disable-next-line no-console
@@ -81,10 +88,21 @@ export async function runHeartbeatQueueSweep(
   db: BopoDb,
   companyId: string,
   options?: { requestId?: string; realtimeHub?: RealtimeHub; maxJobsPerSweep?: number }
+) {
+  if (queueWorkTracker.isShuttingDown()) {
+    return { processed: 0 };
+  }
+  return queueWorkTracker.track(runHeartbeatQueueSweepInternal(db, companyId, options));
+}
+
+async function runHeartbeatQueueSweepInternal(
+  db: BopoDb,
+  companyId: string,
+  options?: { requestId?: string; realtimeHub?: RealtimeHub; maxJobsPerSweep?: number }
 ) {
   const maxJobs = Math.max(1, Math.min(options?.maxJobsPerSweep ?? 50, 500));
   let processed = 0;
-  while (processed < maxJobs) {
+  while (processed < maxJobs && !queueWorkTracker.isShuttingDown()) {
     const job = await claimNextHeartbeatJob(db, companyId);
     if (!job) {
       break;
@@ -117,6 +135,19 @@ export async function runHeartbeatQueueSweep(
   return { processed };
 }
 
+export function beginHeartbeatQueueShutdown() {
+  queueWorkTracker.beginShutdown();
+}
+
+export async function waitForHeartbeatQueueDrain() {
+  await queueWorkTracker.drain();
+}
+
+export function resetHeartbeatQueueShutdownForTests() {
+  activeCompanyQueueWorkers.clear();
+  queueWorkTracker.resetForTests();
+}
+
 async function processHeartbeatQueueJob(
   db: BopoDb,
   input: {

package/src/services/heartbeat-service/active-runs.ts

@@ -0,0 +1,15 @@
+import type { ActiveHeartbeatRun } from "./types";
+
+const activeHeartbeatRuns = new Map<string, ActiveHeartbeatRun>();
+
+export function registerActiveHeartbeatRun(runId: string, run: ActiveHeartbeatRun) {
+  activeHeartbeatRuns.set(runId, run);
+}
+
+export function unregisterActiveHeartbeatRun(runId: string) {
+  activeHeartbeatRuns.delete(runId);
+}
+
+export function getActiveHeartbeatRun(runId: string) {
+  return activeHeartbeatRuns.get(runId);
+}