From a1dd361f51fa72d44830d04ab025d948805a009f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Thu, 23 Apr 2026 19:08:06 -0400 Subject: [PATCH] fix(thread): make wait-state side effects atomic --- src/main/database-base.ts | 23 +++ src/main/thread-integration/index.ts | 189 ++++++++++-------- .../integration/thread-wait-heartbeat.spec.ts | 165 +++++++++++++++ tests/unit/database-transaction.test.ts | 62 ++++++ 4 files changed, 356 insertions(+), 83 deletions(-) create mode 100644 tests/unit/database-transaction.test.ts diff --git a/src/main/database-base.ts b/src/main/database-base.ts index b115abb5..17d0b8d3 100644 --- a/src/main/database-base.ts +++ b/src/main/database-base.ts @@ -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(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; } } diff --git a/src/main/thread-integration/index.ts b/src/main/thread-integration/index.ts index d685996f..e2b66111 100644 --- a/src/main/thread-integration/index.ts +++ b/src/main/thread-integration/index.ts @@ -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, +): 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 }); } }, diff --git a/tests/integration/thread-wait-heartbeat.spec.ts b/tests/integration/thread-wait-heartbeat.spec.ts index ebb4e910..235f991b 100644 --- a/tests/integration/thread-wait-heartbeat.spec.ts +++ b/tests/integration/thread-wait-heartbeat.spec.ts @@ -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(); + }); }); diff --git a/tests/unit/database-transaction.test.ts b/tests/unit/database-transaction.test.ts new file mode 100644 index 00000000..34ec7b15 --- /dev/null +++ b/tests/unit/database-transaction.test.ts @@ -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); + }); +});