AgentCompany/tests/integration/database.spec.ts

import { mkdtemp, rm } from "node:fs/promises";
import { join } from "node:path";
import { tmpdir } from "node:os";
import { afterEach, beforeAll, 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";

const tempDirs: string[] = [];

async function createTempDir(prefix: string) {
  const dir = await mkdtemp(join(tmpdir(), prefix));
  tempDirs.push(dir);
  return dir;
}

afterEach(async () => {
  await Promise.all(tempDirs.splice(0).map((dir) => rm(dir, { recursive: true, force: true })));
});

describe("AppDatabase", () => {
  it("persists core entities, secrets, run logs, and backup restore state", async () => {
    const profileDir = await createTempDir("agentcompany-profile-");
    const backupDir = await createTempDir("agentcompany-backup-");
    const db = new AppDatabase(profileDir);
    db.init();

    const companyId = db.company.saveCompany({
      name: "Orbit Foundry",
      description: "Test company",
      status: "active",
    });

    const workspaceId = db.workspacesMod.saveWorkspace({
      companyId,
      name: "Primary",
      localPath: profileDir,
      repoUrl: "https://example.com/repo.git",
      repoRef: "main",
      isPrimary: true,
    });

    const agentId = db.agents.saveAgent({
      companyId,
      name: "Helios",
      role: "CEO",
      title: "Lead agent",
      status: "active",
      reportsTo: null,
      connectorId: "openai-codex",
      workspaceId,
      model: null,
      capabilities: "Operate the company",
      budgetMonthlyUsd: 200,
    });

    const goalId = db.entityCrud.saveGoal({
      companyId,
      title: "Ship the desktop app",
      description: "Reach packaged desktop parity.",
      ownerAgentId: agentId,
      status: "active",
    });

    const projectId = db.entityCrud.saveProject({
      companyId,
      goalId,
      name: "Desktop rebuild",
      description: "Rebuild as Electron desktop app.",
      leadAgentId: agentId,
      status: "in_progress",
    });

    const taskId = db.tasks.saveTask({
      companyId,
      projectId,
      goalId,
      title: "Validate connector state",
      description: "Run a health check and capture output.",
      assigneeAgentId: agentId,
      workspaceId,
      priority: "high",
      status: "todo",
    });

    db.entityCrud.saveSecret({
      companyId,
      name: "vault-openai",
      description: "Test key",
      value: "sk-test-123",
    });

    db.connectorsCrud.saveConnector({
      id: "openai-codex",
      model: null,
      envBindingText: "OPENAI_API_KEY=vault-openai",
      notes: "Connector note",
    });

    expect(db.entityCrud.resolveSecretEnv("OPENAI_API_KEY=vault-openai", companyId)).toEqual({
      OPENAI_API_KEY: "sk-test-123",
    });

    const run = db.runs.createRun({
      companyId,
      taskId,
      agentId,
      workspaceId,
      connectorId: "openai-codex",
    });

    await db.runs.appendRunLog(run.id, "stdout", "hello from run\n");
    expect(db.runs.readRunLog(run.id)).toContain("hello from run");

    db.runs.markRunStatus(run.id, "running", "Connector launched");
    db.runs.finishRun({
      runId: run.id,
      status: "succeeded",
      summary: "Completed successfully",
      errorMessage: null,
      exitCode: 0,
      signal: null,
      model: "gpt-5-codex",
      unattributedCost: false,
    });

    const snapshotBeforeBackup = db.snap.listSnapshot();
    expect(snapshotBeforeBackup.companies).toHaveLength(1);
    expect(snapshotBeforeBackup.runs[0]?.status).toBe("succeeded");
    expect(snapshotBeforeBackup.tasks[0]?.activeRunId).toBeNull();

    await db.backupMod.exportBackup(backupDir);

    db.company.saveCompany({
      name: "Temporary company",
      description: "Should disappear after restore",
      status: "active",
    });
    expect(db.snap.listSnapshot().companies).toHaveLength(2);

    await db.backupMod.restoreFromBackup(backupDir);
    const restoredSnapshot = db.snap.listSnapshot();
    expect(restoredSnapshot.companies).toHaveLength(1);
    expect(restoredSnapshot.companies[0]?.name).toBe("Orbit Foundry");
    expect(restoredSnapshot.runs[0]?.summary).toBe("Completed successfully");

    db.close();
  });

  it("returns a renderer snapshot scoped to the current company while preserving global selectors", async () => {
    const profileDir = await createTempDir("agentcompany-renderer-snapshot-");
    const db = new AppDatabase(profileDir);
    db.init();

    const alphaCompanyId = db.company.saveCompany({
      name: "Alpha",
      description: "Alpha company",
      status: "active",
    });

    const betaCompanyId = db.company.saveCompany({
      name: "Beta",
      description: "Beta company",
      status: "active",
    });

    const alphaWorkspaceId = db.workspacesMod.saveWorkspace({
      companyId: alphaCompanyId,
      name: "Alpha Workspace",
      localPath: profileDir,
      repoUrl: "",
      repoRef: "",
      isPrimary: true,
    });

    const betaWorkspaceId = db.workspacesMod.saveWorkspace({
      companyId: betaCompanyId,
      name: "Beta Workspace",
      localPath: profileDir,
      repoUrl: "",
      repoRef: "",
      isPrimary: true,
    });

    const alphaAgentId = db.agents.saveAgent({
      companyId: alphaCompanyId,
      name: "Alpha CEO",
      role: "CEO",
      title: "",
      status: "active",
      reportsTo: null,
      connectorId: "openai-codex",
      workspaceId: alphaWorkspaceId,
      model: null,
      capabilities: "Lead alpha",
      budgetMonthlyUsd: 100,
    });

    db.agents.saveAgent({
      companyId: betaCompanyId,
      name: "Beta CEO",
      role: "CEO",
      title: "",
      status: "active",
      reportsTo: null,
      connectorId: "openai-codex",
      workspaceId: betaWorkspaceId,
      model: null,
      capabilities: "Lead beta",
      budgetMonthlyUsd: 100,
    });

    db.tasks.saveTask({
      companyId: alphaCompanyId,
      title: "Alpha task",
      description: "Alpha work",
      assigneeAgentId: alphaAgentId,
      workspaceId: alphaWorkspaceId,
      priority: "medium",
      status: "todo",
    });

    db.tasks.saveTask({
      companyId: betaCompanyId,
      title: "Beta task",
      description: "Beta work",
      assigneeAgentId: null,
      workspaceId: betaWorkspaceId,
      priority: "medium",
      status: "todo",
    });

    db.settings.setCurrentCompany(alphaCompanyId);

    const rendererSnapshot = db.snap.listRendererSnapshot();
    expect(rendererSnapshot.currentCompanyId).toBe(alphaCompanyId);
    expect(rendererSnapshot.companies).toHaveLength(2);
    expect(rendererSnapshot.connectors.length).toBeGreaterThan(0);
    expect(rendererSnapshot.workspaces.map((workspace) => workspace.companyId)).toEqual([alphaCompanyId]);
    expect(rendererSnapshot.agents.every((agent) => agent.companyId === alphaCompanyId)).toBe(true);
    expect(rendererSnapshot.tasks.every((task) => task.companyId === alphaCompanyId)).toBe(true);

    db.close();
  });
});