fix(thread): make wait-state side effects atomic

2026-04-23 19:08:06 -04:00
1 parent f8a90f4402
commit a1dd361f51
@@ -82,6 +82,8 @@ export class AppDatabaseBase {
  readonly appVersion: string;
  readonly vault: SecretVault;
  protected db: DatabaseSync;
+  private transactionDepth = 0;
+  private savepointCounter = 0;

  public readonly core: AppDatabaseCoreModules;
  public readonly automation: AppDatabaseAutomationModules;
@@ -362,7 +364,26 @@ export class AppDatabaseBase {
  }

  protected transaction<T>(operation: () => T): T {
+    if (this.transactionDepth > 0) {
+      const savepoint = `agentcompany_tx_${this.savepointCounter}`;
+      this.savepointCounter += 1;
+      this.db.exec(`SAVEPOINT ${savepoint}`);
+      this.transactionDepth += 1;
+      try {
+        const result = operation();
+        this.db.exec(`RELEASE SAVEPOINT ${savepoint}`);
+        return result;
+      } catch (error) {
+        this.db.exec(`ROLLBACK TO SAVEPOINT ${savepoint}`);
+        this.db.exec(`RELEASE SAVEPOINT ${savepoint}`);
+        throw error;
+      } finally {
+        this.transactionDepth -= 1;
+      }
+    }
+
    this.db.exec("BEGIN IMMEDIATE");
+    this.transactionDepth += 1;
    try {
      const result = operation();
      this.db.exec("COMMIT");
@@ -370,6 +391,8 @@ export class AppDatabaseBase {
    } catch (error) {
      this.db.exec("ROLLBACK");
      throw error;
+    } finally {
+      this.transactionDepth -= 1;
    }
  }

@@ -1,6 +1,7 @@
 /** Thread Integration — wires thread runtime into orchestration callbacks. */

 import type { ExecutionMode } from "@shared/types";
+import { classifyUnknownError } from "../errors";
 import {
  type ThreadIntegrationDeps,
  findTaskThread,
@@ -12,6 +13,22 @@ import {

 export type { ThreadIntegrationDeps } from "./thread-helpers";

+function throwThreadIntegrationError(
+  logger: ThreadIntegrationDeps["logger"],
+  operation: string,
+  error: unknown,
+  context: Record<string, unknown>,
+): never {
+  const classified = classifyUnknownError(error, {
+    kind: "upstream",
+    code: `thread.integration.${operation}.failed`,
+    message: `Thread integration ${operation} failed`,
+    context,
+  });
+  logger.error(`[${classified.code}] ${classified.message}: ${error instanceof Error ? error.message : String(error)}`);
+  throw classified;
+}
+
 export function createThreadIntegration({ db, logger, emitEvent, publishDomainChanged, notify }: ThreadIntegrationDeps) {
  return {
    /**
@@ -104,44 +121,46 @@ export function createThreadIntegration({ db, logger, emitEvent, publishDomainCh
     */
    onApprovalCreated(approvalId: string, companyId: string, taskId: string | null, approvalType: string, summary: string): void {
      try {
-        const approval = db.core.approvals.getApprovalRecord(approvalId);
-        const task = taskId ? db.core.entity.getTask(taskId) : null;
-        const thread = task
-          ? ensureTaskThread(db, task.id, companyId, task.title)
-          : ensureApprovalThread(db, approvalId, companyId, summary || approval.payloadSummary);
+        db.withTransaction(() => {
+          const approval = db.core.approvals.getApprovalRecord(approvalId);
+          const task = taskId ? db.core.entity.getTask(taskId) : null;
+          const thread = task
+            ? ensureTaskThread(db, task.id, companyId, task.title)
+            : ensureApprovalThread(db, approvalId, companyId, summary || approval.payloadSummary);

-        db.core.threads.addThreadEvent({
-          threadId: thread.id,
-          actorId: "system",
-          actorType: "system",
-          eventType: "approval_request",
-          body: summary ? `@Board ${summary}` : `@Board approval required: ${approvalType}`,
-          metadataJson: JSON.stringify({ approvalId, approvalType }),
-        });
-
-        if (approval.requestedByAgentId) {
-          db.core.threads.addParticipant({ threadId: thread.id, participantId: approval.requestedByAgentId, participantType: "agent", role: "owner" });
-          db.core.threads.subscribe({ threadId: thread.id, subscriberId: approval.requestedByAgentId, subscriberType: "agent" });
-        }
-        db.core.threads.addParticipant({ threadId: thread.id, participantId: "user", participantType: "user", role: "observer" });
-        db.core.threads.subscribe({ threadId: thread.id, subscriberId: "user", subscriberType: "user" });
-
-        // In co-founder mode, set wait state for user approval
-        const mode = db.core.company.getCompanyExecutionMode(companyId) as ExecutionMode;
-        if (mode === "ai_cofounder") {
-          const approvalReason = `@Board approval needed: ${approvalType}`;
-          db.core.threads.setThreadWaitState({
+          db.core.threads.addThreadEvent({
            threadId: thread.id,
-            companyId,
-            waitingFor: "user_approval",
-            relatedApprovalId: approvalId,
-            reason: approvalReason,
+            actorId: "system",
+            actorType: "system",
+            eventType: "approval_request",
+            body: summary ? `@Board ${summary}` : `@Board approval required: ${approvalType}`,
+            metadataJson: JSON.stringify({ approvalId, approvalType }),
          });
-          notifyThreadSubscribers(db, thread.id, thread.title, approvalReason, notify);
-        }
-        announceApprovalToBoard(db, approval, { emitEvent, publishDomainChanged, logger });
+
+          if (approval.requestedByAgentId) {
+            db.core.threads.addParticipant({ threadId: thread.id, participantId: approval.requestedByAgentId, participantType: "agent", role: "owner" });
+            db.core.threads.subscribe({ threadId: thread.id, subscriberId: approval.requestedByAgentId, subscriberType: "agent" });
+          }
+          db.core.threads.addParticipant({ threadId: thread.id, participantId: "user", participantType: "user", role: "observer" });
+          db.core.threads.subscribe({ threadId: thread.id, subscriberId: "user", subscriberType: "user" });
+
+          // In co-founder mode, set wait state for user approval
+          const mode = db.core.company.getCompanyExecutionMode(companyId) as ExecutionMode;
+          if (mode === "ai_cofounder") {
+            const approvalReason = `@Board approval needed: ${approvalType}`;
+            db.core.threads.setThreadWaitState({
+              threadId: thread.id,
+              companyId,
+              waitingFor: "user_approval",
+              relatedApprovalId: approvalId,
+              reason: approvalReason,
+            });
+            notifyThreadSubscribers(db, thread.id, thread.title, approvalReason, notify);
+          }
+          announceApprovalToBoard(db, approval, { emitEvent, publishDomainChanged, logger });
+        });
      } catch (err) {
-        logger.error(`[thread-integration] onApprovalCreated failed: ${err instanceof Error ? err.message : String(err)}`);
+        throwThreadIntegrationError(logger, "onApprovalCreated", err, { approvalId, companyId, taskId, approvalType });
      }
    },

@@ -151,25 +170,27 @@ export function createThreadIntegration({ db, logger, emitEvent, publishDomainCh
     */
    onApprovalResolved(approvalId: string, companyId: string, decision: string): void {
      try {
-        // Resolve any wait states linked to this approval
-        const waitStates = db.core.threads.listThreadWaitStates(companyId);
-        for (const ws of waitStates) {
-          if (ws.relatedApprovalId === approvalId) {
-            db.core.threads.resolveThreadWaitState(ws.id, companyId);
+        db.withTransaction(() => {
+          // Resolve any wait states linked to this approval
+          const waitStates = db.core.threads.listThreadWaitStates(companyId);
+          for (const ws of waitStates) {
+            if (ws.relatedApprovalId === approvalId) {
+              db.core.threads.resolveThreadWaitState(ws.id, companyId);

-            // Add decision event to the thread
-            db.core.threads.addThreadEvent({
-              threadId: ws.threadId,
-              actorId: "user",
-              actorType: "user",
-              eventType: "approval_decision",
-              body: `Approval ${decision}`,
-              metadataJson: JSON.stringify({ approvalId, decision }),
-            });
+              // Add decision event to the thread
+              db.core.threads.addThreadEvent({
+                threadId: ws.threadId,
+                actorId: "user",
+                actorType: "user",
+                eventType: "approval_decision",
+                body: `Approval ${decision}`,
+                metadataJson: JSON.stringify({ approvalId, decision }),
+              });
+            }
          }
-        }
+        });
      } catch (err) {
-        logger.error(`[thread-integration] onApprovalResolved failed: ${err instanceof Error ? err.message : String(err)}`);
+        throwThreadIntegrationError(logger, "onApprovalResolved", err, { approvalId, companyId, decision });
      }
    },

@@ -278,44 +299,46 @@ export function createThreadIntegration({ db, logger, emitEvent, publishDomainCh
     */
    onHumanMessage(companyId: string, toAgentId?: string | null, threadId?: string | null, waitStateId?: string | null): void {
      try {
-        // Most precise: resolve a single identified wait state
-        if (waitStateId) {
-          db.core.threads.resolveThreadWaitState(waitStateId, companyId);
-          return;
-        }
-
-        // Scoped resolution: threadId or toAgentId resolve all user_* types
-        const SCOPED_WAIT_TYPES = new Set(["user_reply", "user_review"]);
-        const isScoped = Boolean(threadId || toAgentId);
-
-        const waitStates = db.core.threads.listThreadWaitStates(companyId);
-        for (const ws of waitStates) {
-          if (isScoped) {
-            if (!SCOPED_WAIT_TYPES.has(ws.waitingFor)) continue;
-          } else {
-            // Company-wide fallback: only resolve user_reply to prevent
-            // accidental resolution of reviews/approvals from casual messages
-            if (ws.waitingFor !== "user_reply") continue;
+        db.withTransaction(() => {
+          // Most precise: resolve a single identified wait state
+          if (waitStateId) {
+            db.core.threads.resolveThreadWaitState(waitStateId, companyId);
+            return;
          }

-          if (threadId) {
-            if (ws.threadId !== threadId) continue;
-          } else if (toAgentId) {
-            const thread = db.core.threads.getThread(ws.threadId, companyId);
-            if (!thread || thread.relatedEntityType !== "task" || !thread.relatedEntityId) continue;
-            try {
-              const task = db.core.entity.getTask(thread.relatedEntityId);
-              if (task.assigneeAgentId !== toAgentId) continue;
-            } catch (err) {
-              logger.warn(`[thread-integration] getTask failed for thread ${ws.threadId}:`, err instanceof Error ? err.message : String(err));
-              continue;
+          // Scoped resolution: threadId or toAgentId resolve all user_* types
+          const SCOPED_WAIT_TYPES = new Set(["user_reply", "user_review"]);
+          const isScoped = Boolean(threadId || toAgentId);
+
+          const waitStates = db.core.threads.listThreadWaitStates(companyId);
+          for (const ws of waitStates) {
+            if (isScoped) {
+              if (!SCOPED_WAIT_TYPES.has(ws.waitingFor)) continue;
+            } else {
+              // Company-wide fallback: only resolve user_reply to prevent
+              // accidental resolution of reviews/approvals from casual messages
+              if (ws.waitingFor !== "user_reply") continue;
            }
-          }

-          db.core.threads.resolveThreadWaitState(ws.id, companyId);
-        }
+            if (threadId) {
+              if (ws.threadId !== threadId) continue;
+            } else if (toAgentId) {
+              const thread = db.core.threads.getThread(ws.threadId, companyId);
+              if (!thread || thread.relatedEntityType !== "task" || !thread.relatedEntityId) continue;
+              try {
+                const task = db.core.entity.getTask(thread.relatedEntityId);
+                if (task.assigneeAgentId !== toAgentId) continue;
+              } catch (err) {
+                logger.warn(`[thread-integration] getTask failed for thread ${ws.threadId}:`, err instanceof Error ? err.message : String(err));
+                continue;
+              }
+            }
+
+            db.core.threads.resolveThreadWaitState(ws.id, companyId);
+          }
+        });
      } catch (err) {
-        logger.error(`[thread-integration] onHumanMessage failed: ${err instanceof Error ? err.message : String(err)}`);
+        throwThreadIntegrationError(logger, "onHumanMessage", err, { companyId, toAgentId, threadId, waitStateId });
      }
    },

@@ -6,6 +6,7 @@ import { AppDatabase } from "../../src/main/database";
 import { createThreadIntegration } from "../../src/main/thread-integration";
 import { createCheckpointEngine } from "../../src/main/checkpoint-engine";
 import { dispatchAgentMessage, HUMAN_SENDER_ID } from "../../src/main/messaging/dispatch";
+import { isAppError } from "../../src/main/errors";

 const tempDirs: string[] = [];

@@ -89,6 +90,14 @@ function createScaffold(db: AppDatabase) {
  return { companyId, workspaceId, ceoId, engineerId, taskId };
 }

+function expectThreadIntegrationFailure(error: unknown, code: string) {
+  expect(isAppError(error)).toBe(true);
+  if (isAppError(error)) {
+    expect(error.kind).toBe("upstream");
+    expect(error.code).toBe(code);
+  }
+}
+
 describe("thread wait + heartbeat integration", () => {
  it("pauses the assignee heartbeat on successful run completion until a human reply resolves the wait state", async () => {
    const profileDir = await createTempDir("agentcompany-thread-wait-");
@@ -276,4 +285,160 @@ describe("thread wait + heartbeat integration", () => {

    db.close();
  });
+
+  it("rolls back approval-created thread writes when wait-state creation fails", async () => {
+    const profileDir = await createTempDir("agentcompany-approval-create-fail-");
+    const db = new AppDatabase(profileDir);
+    db.init();
+
+    const { companyId, engineerId, taskId } = createScaffold(db);
+    const threadIntegration = createThreadIntegration({
+      db,
+      logger: createNoopLogger(),
+    });
+    const approvalId = db.core.approvals.requestApproval({
+      companyId,
+      relatedTaskId: taskId,
+      requestedByAgentId: engineerId,
+      type: "draft_review",
+      payloadSummary: "Review the draft",
+      impactSummary: "Board approval is required before this work can continue.",
+      payloadJson: "{}",
+    });
+
+    vi.spyOn(db.core.threads, "setThreadWaitState").mockImplementation(() => {
+      throw new Error("wait-state write failed");
+    });
+
+    let thrown: unknown;
+    try {
+      threadIntegration.onApprovalCreated(approvalId, companyId, taskId, "draft_review", "Draft review needed");
+    } catch (error) {
+      thrown = error;
+    }
+
+    expectThreadIntegrationFailure(thrown, "thread.integration.onApprovalCreated.failed");
+    const taskThreads = db.core.threads.listThreads({
+      companyId,
+      type: "task",
+      relatedEntityType: "task",
+      relatedEntityId: taskId,
+      limit: 1,
+    });
+    expect(taskThreads).toHaveLength(0);
+
+    db.close();
+  });
+
+  it("rolls back approval resolution when decision event recording fails", async () => {
+    const profileDir = await createTempDir("agentcompany-approval-resolve-fail-");
+    const db = new AppDatabase(profileDir);
+    db.init();
+
+    const { companyId, engineerId, taskId } = createScaffold(db);
+    const threadIntegration = createThreadIntegration({
+      db,
+      logger: createNoopLogger(),
+    });
+    threadIntegration.onTaskCreated(taskId, companyId);
+    const approvalId = db.core.approvals.requestApproval({
+      companyId,
+      relatedTaskId: taskId,
+      requestedByAgentId: engineerId,
+      type: "draft_review",
+      payloadSummary: "Review the draft",
+      impactSummary: "Board approval is required before this work can continue.",
+      payloadJson: "{}",
+    });
+    threadIntegration.onApprovalCreated(approvalId, companyId, taskId, "draft_review", "Draft review needed");
+    const waitStateId = db.core.threads.listThreadWaitStates(companyId)
+      .find((entry) => entry.relatedApprovalId === approvalId)?.id;
+    expect(waitStateId).toBeTruthy();
+
+    const addThreadEvent = db.core.threads.addThreadEvent.bind(db.core.threads);
+    vi.spyOn(db.core.threads, "addThreadEvent").mockImplementation((input) => {
+      if (input.eventType === "approval_decision") {
+        throw new Error("decision event write failed");
+      }
+      return addThreadEvent(input);
+    });
+
+    let thrown: unknown;
+    try {
+      threadIntegration.onApprovalResolved(approvalId, companyId, "approved");
+    } catch (error) {
+      thrown = error;
+    }
+
+    expectThreadIntegrationFailure(thrown, "thread.integration.onApprovalResolved.failed");
+    const unresolved = db.core.threads.listThreadWaitStates(companyId)
+      .filter((entry) => entry.relatedApprovalId === approvalId);
+    expect(unresolved.map((entry) => entry.id)).toEqual([waitStateId]);
+
+    db.close();
+  });
+
+  it("rolls back human-message wait-state resolution when one resolution fails", async () => {
+    const profileDir = await createTempDir("agentcompany-human-message-fail-");
+    const db = new AppDatabase(profileDir);
+    db.init();
+
+    const { companyId, taskId } = createScaffold(db);
+    const threadIntegration = createThreadIntegration({
+      db,
+      logger: createNoopLogger(),
+    });
+    threadIntegration.onTaskCreated(taskId, companyId);
+    const thread = db.core.threads.listThreads({
+      companyId,
+      type: "task",
+      relatedEntityType: "task",
+      relatedEntityId: taskId,
+      limit: 1,
+    })[0];
+    if (!thread) {
+      throw new Error("Expected task thread for human-message failure testing.");
+    }
+
+    const first = db.core.threads.setThreadWaitState({
+      threadId: thread.id,
+      companyId,
+      waitingFor: "user_reply",
+      relatedApprovalId: null,
+      reason: "First reply needed",
+    });
+    const second = db.core.threads.setThreadWaitState({
+      threadId: thread.id,
+      companyId,
+      waitingFor: "user_review",
+      relatedApprovalId: null,
+      reason: "Second review needed",
+    });
+
+    const resolveThreadWaitState = db.core.threads.resolveThreadWaitState.bind(db.core.threads);
+    let resolutionCount = 0;
+    vi.spyOn(db.core.threads, "resolveThreadWaitState").mockImplementation((waitStateId, currentCompanyId) => {
+      resolutionCount += 1;
+      if (resolutionCount === 2) {
+        throw new Error("wait-state resolve failed");
+      }
+      resolveThreadWaitState(waitStateId, currentCompanyId);
+    });
+
+    let thrown: unknown;
+    try {
+      threadIntegration.onHumanMessage(companyId, null, thread.id, null);
+    } catch (error) {
+      thrown = error;
+    }
+
+    expectThreadIntegrationFailure(thrown, "thread.integration.onHumanMessage.failed");
+    const unresolvedIds = db.core.threads.listThreadWaitStates(companyId)
+      .filter((entry) => entry.threadId === thread.id)
+      .map((entry) => entry.id)
+      .sort();
+    expect(unresolvedIds).toEqual([first.id, second.id].sort());
+
+    db.close();
+  });
 });
@@ -0,0 +1,62 @@
+import { join } from "node:path";
+import { tmpdir } from "node:os";
+import { mkdtemp, rm } from "node:fs/promises";
+import { afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
+
+beforeAll(() => {
+  process.env.AGENTCOMPANY_ALLOW_INSECURE_VAULT = "true";
+});
+
+vi.mock("electron", () => ({
+  safeStorage: {
+    isEncryptionAvailable: () => false,
+    getSelectedStorageBackend: () => "basic_text",
+    encryptString: (value: string) => Buffer.from(value, "utf8"),
+    decryptString: (value: Buffer) => value.toString("utf8"),
+  },
+}));
+
+import { AppDatabase } from "../../src/main/database";
+
+describe("AppDatabase transactions", () => {
+  let tempDir: string;
+  let db: AppDatabase;
+
+  beforeEach(async () => {
+    tempDir = await mkdtemp(join(tmpdir(), "agentcompany-db-tx-"));
+    db = new AppDatabase(tempDir);
+    db.init();
+  });
+
+  afterEach(async () => {
+    db.close();
+    await rm(tempDir, { recursive: true, force: true });
+  });
+
+  it("rolls back nested transactions to their savepoint", () => {
+    const outerCompanyId = db.withTransaction(() => {
+      const companyId = db.core.company.saveCompany({
+        name: "Nested Transaction Co",
+        description: "Outer transaction should survive inner rollback.",
+        status: "active",
+      });
+
+      expect(() => {
+        db.withTransaction(() => {
+          db.core.company.saveCompany({
+            name: "Nested Transaction Rolled Back",
+            description: "Inner transaction should roll back.",
+            status: "active",
+          });
+          throw new Error("inner write failed");
+        });
+      }).toThrow("inner write failed");
+
+      return companyId;
+    });
+
+    const companies = db.queryAll<{ id: string; name: string }>("select id, name from companies order by created_at asc");
+    expect(companies.map((company) => company.id)).toContain(outerCompanyId);
+    expect(companies.some((company) => company.name === "Nested Transaction Rolled Back")).toBe(false);
+  });
+});