update

src/main/worker-service.ts

@@ -0,0 +1,337 @@
+import { utilityProcess } from "electron";
+import { z } from "zod";
+import { getConnectorDefinition } from "./connectors";
+import type { AppDatabase } from "./database";
+import type { RunRecord } from "@shared/types";
+
+const runStatusSchema = z.enum(["queued", "running", "succeeded", "failed", "cancelled", "timed_out", "interrupted"]);
+export const workerMessageSchema = z.union([
+  z.object({
+    type: z.literal("run-log"),
+    runId: z.string().uuid(),
+    stream: z.enum(["stdout", "stderr"]),
+    chunk: z.string(),
+  }),
+  z.object({
+    type: z.literal("run-status"),
+    runId: z.string().uuid(),
+    status: runStatusSchema,
+    message: z.string(),
+  }),
+  z.object({
+    type: z.literal("run-finished"),
+    runId: z.string().uuid(),
+    status: runStatusSchema,
+    summary: z.string().nullable(),
+    errorMessage: z.string().nullable(),
+    exitCode: z.number().int().nullable(),
+    signal: z.string().nullable(),
+    model: z.string().nullable(),
+    sessionParams: z.record(z.string(), z.unknown()).nullable(),
+    sessionDisplayId: z.string().nullable(),
+    costUsd: z.number().nullable(),
+    unattributedCost: z.boolean(),
+  }),
+]);
+type WorkerUiEvent = Extract<z.infer<typeof workerMessageSchema>, { type: "run-log" | "run-status" }>;
+
+interface WorkerServiceDependencies {
+  db: AppDatabase;
+  logger: {
+    info: (...args: unknown[]) => void;
+    warn: (...args: unknown[]) => void;
+    error: (...args: unknown[]) => void;
+  };
+  workerPath: string;
+  isQuitting: () => boolean;
+  publishDomainChanged: () => void;
+  emitEvent: (event: WorkerUiEvent) => void;
+  wakeAgentIfPossible: (agentId: string | null | undefined, companyId: string, trigger: string) => string | null;
+  dispatchRunToWorker: (params: {
+    runId: string;
+    companyId: string;
+    agent: ReturnType<AppDatabase["getAgent"]>;
+    task: ReturnType<AppDatabase["getTask"]>;
+    connector: ReturnType<AppDatabase["getConnector"]>;
+    workspace: ReturnType<AppDatabase["getWorkspace"]>;
+    workspaceMode: "exclusive" | "shared" | "none";
+    companyName: string;
+    companyDescription: string;
+    snapshot: ReturnType<AppDatabase["listSnapshot"]>;
+    trigger: string;
+  }) => void;
+  processRunCompletion: (run: RunRecord, event: z.infer<typeof workerMessageSchema> & { type: "run-finished" }) => void;
+  notify: (title: string, body: string) => void;
+}
+
+export function createWorkerService({
+  db,
+  logger,
+  workerPath,
+  isQuitting,
+  publishDomainChanged,
+  emitEvent,
+  wakeAgentIfPossible,
+  dispatchRunToWorker,
+  processRunCompletion,
+  notify,
+}: WorkerServiceDependencies) {
+  let worker: ReturnType<typeof utilityProcess.fork> | null = null;
+
+  function getWorker() {
+    return worker;
+  }
+
+  function createWorker() {
+    worker = utilityProcess.fork(workerPath, [], {
+      stdio: "pipe",
+      serviceName: "AgentCompany Orchestrator",
+    });
+
+    worker.stdout?.on("data", (chunk) => logger.info(`[worker] ${String(chunk).trim()}`));
+    worker.stderr?.on("data", (chunk) => logger.error(`[worker] ${String(chunk).trim()}`));
+
+    worker.on("exit", (code) => {
+      if (isQuitting()) return;
+
+      logger.warn(`[worker] Worker exited with code ${code}, recovering...`);
+
+      const snapshot = db.listSnapshot();
+      const affectedAgentIds = new Set<string>();
+      const requeueCandidates: Array<{ id: string; taskId: string; agentId: string; companyId: string; workspaceId: string | null; connectorId: string }> = [];
+
+      for (const run of snapshot.runs) {
+        if (run.status === "running") {
+          db.finishRun({
+            runId: run.id,
+            status: "interrupted",
+            summary: "Worker process crashed.",
+            errorMessage: "Worker process exited unexpectedly.",
+            exitCode: code,
+            signal: null,
+            model: null,
+            sessionDisplayId: null,
+            costUsd: null,
+            unattributedCost: true,
+          });
+          if (run.workspaceId) {
+            db.releaseWorkspaceLock(run.workspaceId, run.id);
+          }
+          affectedAgentIds.add(run.agentId);
+
+          const chain = db.getChainOfCommand(run.agentId);
+          if (chain.length > 0) {
+            wakeAgentIfPossible(chain[0].id, run.companyId, "report_failed");
+          }
+        } else if (run.status === "queued") {
+          requeueCandidates.push({
+            id: run.id,
+            taskId: run.taskId,
+            agentId: run.agentId,
+            companyId: run.companyId,
+            workspaceId: run.workspaceId,
+            connectorId: run.connectorId,
+          });
+        }
+      }
+
+      for (const agentId of affectedAgentIds) {
+        const pendingWakes = db.consumePendingWakes(agentId);
+        for (const wake of pendingWakes) {
+          wakeAgentIfPossible(agentId, wake.companyId, wake.trigger);
+        }
+      }
+
+      publishDomainChanged();
+
+      createWorker();
+
+      if (requeueCandidates.length > 0 && worker) {
+        logger.info(`[worker] Re-enqueuing ${requeueCandidates.length} queued run(s) after worker restart.`);
+        const recoverySnapshot = db.listSnapshot();
+        for (const candidate of requeueCandidates) {
+          try {
+            const task = db.getTask(candidate.taskId);
+            const agent = db.getAgent(candidate.agentId);
+            const connector = db.getConnector(agent.connectorId);
+            const workspaceResolution = {
+              taskWorkspace: db.getWorkspace(task.workspaceId),
+              agentWorkspace: db.getWorkspace(agent.workspaceId),
+              projectWorkspace: db.getProjectPrimaryWorkspace(task.projectId ?? null, candidate.companyId),
+            };
+            const workspace = workspaceResolution.taskWorkspace
+              ?? workspaceResolution.agentWorkspace
+              ?? workspaceResolution.projectWorkspace
+              ?? db.getAnyCompanyWorkspace(candidate.companyId);
+            const workspaceMode = workspaceResolution.taskWorkspace || workspaceResolution.agentWorkspace || workspaceResolution.projectWorkspace
+              ? "exclusive"
+              : workspace
+                ? "shared"
+                : "none";
+            const recoveryCompany = recoverySnapshot.companies.find((company) => company.id === candidate.companyId);
+
+            dispatchRunToWorker({
+              runId: candidate.id,
+              companyId: candidate.companyId,
+              agent,
+              task,
+              connector,
+              workspace,
+              workspaceMode,
+              companyName: recoveryCompany?.name ?? "Company",
+              companyDescription: recoveryCompany?.description ?? "",
+              snapshot: recoverySnapshot,
+              trigger: "recovery",
+            });
+          } catch (error) {
+            logger.warn(`[worker] Failed to re-enqueue run ${candidate.id}: ${error instanceof Error ? error.message : String(error)}`);
+            db.finishRun({
+              runId: candidate.id,
+              status: "interrupted",
+              summary: "Could not re-enqueue after worker restart.",
+              errorMessage: error instanceof Error ? error.message : String(error),
+              exitCode: null,
+              signal: null,
+              model: null,
+              sessionDisplayId: null,
+              costUsd: null,
+              unattributedCost: true,
+            });
+            if (candidate.workspaceId) {
+              db.releaseWorkspaceLock(candidate.workspaceId, candidate.id);
+            }
+          }
+        }
+      }
+    });
+
+    worker.on("message", async (message) => {
+      let event: z.infer<typeof workerMessageSchema>;
+      try {
+        event = workerMessageSchema.parse(message);
+      } catch (error) {
+        logger.error("[worker] invalid message payload", error);
+        return;
+      }
+
+      if (event.type === "run-log") {
+        await db.appendRunLog(event.runId, event.stream, event.chunk);
+        emitEvent(event);
+        return;
+      }
+
+      if (event.type === "run-status") {
+        db.markRunStatus(event.runId, event.status, event.message);
+        if (event.status === "running") {
+          const run = db.listSnapshot().runs.find((entry) => entry.id === event.runId);
+          if (run?.workspaceId && !db.claimWorkspaceLock(run.workspaceId, run.id)) {
+            logger.error(`[workspace] Failed to claim lock for workspace ${run.workspaceId} when run ${run.id} entered running state. Cancelling run.`);
+            worker?.postMessage({ type: "cancel-run", runId: run.id });
+          }
+        }
+        emitEvent(event);
+        publishDomainChanged();
+        return;
+      }
+
+      const snapshot = db.listSnapshot();
+      const run = snapshot.runs.find((entry) => entry.id === event.runId);
+      if (!run) {
+        return;
+      }
+
+      const connector = getConnectorDefinition(run.connectorId);
+
+      let task: ReturnType<AppDatabase["getTask"]> | null = null;
+      try {
+        task = db.getTask(run.taskId);
+      } catch {
+        task = null;
+      }
+
+      if (task) {
+        if (event.sessionParams && connector.module.sessionCodec) {
+          const serialized = connector.module.sessionCodec.serialize(event.sessionParams);
+          db.updateTaskSession(task.id, serialized, connector.module.sessionCodec.getDisplayId?.(serialized ?? null) ?? event.sessionDisplayId ?? null);
+        }
+        if (!event.sessionParams) {
+          db.updateTaskSession(task.id, null, null);
+        }
+      }
+
+      const runOutcome = db.finishRun(event);
+      if (run.workspaceId) {
+        db.releaseWorkspaceLock(run.workspaceId, run.id);
+      }
+      db.addActivity({
+        companyId: run.companyId,
+        actor: run.agentId,
+        action: `run.${event.status}`,
+        entityType: "run",
+        entityId: run.id,
+        detail: event.summary ?? event.errorMessage ?? "",
+      });
+      publishDomainChanged();
+      emitEvent({
+        type: "run-status",
+        runId: run.id,
+        status: event.status,
+        message: event.summary ?? event.errorMessage ?? event.status,
+      });
+      if ((event.status === "failed" || event.status === "timed_out" || event.status === "interrupted") && task) {
+        notify("Run attention required", task.title);
+      }
+
+      let updatedAgent: ReturnType<AppDatabase["getAgent"]> | null = null;
+      try {
+        updatedAgent = db.getAgent(run.agentId);
+      } catch {
+        updatedAgent = null;
+      }
+
+      if (updatedAgent && runOutcome.budgetHardStopped) {
+        db.addActivity({
+          companyId: updatedAgent.companyId,
+          actor: "system",
+          action: "agent.auto_paused",
+          entityType: "agent",
+          entityId: updatedAgent.id,
+          detail: `Budget exhausted at ${updatedAgent.spentMonthlyUsd.toFixed(2)} / ${updatedAgent.budgetMonthlyUsd.toFixed(2)} USD.`,
+        });
+        notify("Agent paused for budget limit", updatedAgent.name);
+        publishDomainChanged();
+      }
+
+      if (event.status === "failed" || event.status === "timed_out") {
+        try {
+          const chain = db.getChainOfCommand(run.agentId);
+          if (chain.length > 0) {
+            wakeAgentIfPossible(chain[0].id, run.companyId, "report_failed");
+          }
+        } catch {
+          // Agent may have been deleted while resolving the escalation chain.
+        }
+      }
+
+      const pendingWakes = db.consumePendingWakes(run.agentId);
+      if (pendingWakes.length > 0) {
+        for (const wake of pendingWakes) {
+          wakeAgentIfPossible(run.agentId, wake.companyId, wake.trigger);
+        }
+      }
+
+      processRunCompletion(run, event as z.infer<typeof workerMessageSchema> & { type: "run-finished" });
+    });
+  }
+
+  function destroy() {
+    worker?.kill();
+    worker = null;
+  }
+
+  return {
+    createWorker,
+    getWorker,
+    destroy,
+  };
+}