AgentCompany/tests/unit/agent-document-tools.test.ts

import { randomUUID } from "node:crypto";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { describe, expect, test } from "vitest";
import { DatabaseConnection } from "../../src/main/database/connection";
import { OrganizationService } from "../../src/main/services/organization-service";
import { ChatService } from "../../src/main/services/chat-service";
import { dispatchAgentTool, type AgentToolContext } from "../../src/main/services/agent-tools";
import { RISK_LEVEL, TASK_PRIORITY, TASK_STATUS } from "../../src/shared/domain";

function harness(reviewDeliverables: boolean) {
  const connection = new DatabaseConnection(join(tmpdir(), `ac-docs-${randomUUID()}.db`));
  const organization = new OrganizationService(connection.db);
  let orgId = "";
  let agentId = "";
  connection.transaction(() => {
    orgId = organization.createOrganization({ organizationName: "D", employeeName: "Y", mission: "m", goal: "g" });
    agentId = organization.createAgent({ orgId, name: "Ana", role: "PM", department: "Product", modelProviderId: null, modelId: null, budgetCents: 1000, managerAgentId: null });
  });
  connection.db.prepare("update organizations set review_deliverables = ? where id = ?").run(reviewDeliverables ? 1 : 0, orgId);
  const ctx: AgentToolContext = {
    orgId, agent: { id: agentId, modelProviderId: null, modelId: null },
    taskProjectId: null, taskGoalId: null, taskId: null, sourceThreadId: null, workspace: null,
  };
  const run = (name: string, args: Record<string, unknown>) =>
    dispatchAgentTool({ db: connection.db, organization, chat: new ChatService(connection.db) }, ctx, { id: `c_${randomUUID()}`, name, arguments: args });
  return { connection, orgId, run };
}

describe("document provenance", () => {
  test("a document records the task that produced it, shown when listing", () => {
    const connection = new DatabaseConnection(join(tmpdir(), `ac-docs-${randomUUID()}.db`));
    const organization = new OrganizationService(connection.db);
    let orgId = "";
    let agentId = "";
    let taskId = "";
    connection.transaction(() => {
      orgId = organization.createOrganization({ organizationName: "D", employeeName: "Y", mission: "m", goal: "g" });
      agentId = organization.createAgent({ orgId, name: "Ana", role: "PM", department: "Product", modelProviderId: null, modelId: null, budgetCents: 1000, managerAgentId: null });
      taskId = organization.createTask({
        orgId, title: "Write the booking PRD", description: "d", status: TASK_STATUS.IN_PROGRESS, assigneeAgentId: agentId,
        risk: RISK_LEVEL.LOW, projectId: null, goalId: null, workspaceId: null, priority: TASK_PRIORITY.NORMAL, requiresFilesystem: false,
      });
    });
    const ctx: AgentToolContext = {
      orgId, agent: { id: agentId, modelProviderId: null, modelId: null },
      taskProjectId: null, taskGoalId: null, taskId, sourceThreadId: null, workspace: null,
    };
    const run = (name: string, args: Record<string, unknown>) =>
      dispatchAgentTool({ db: connection.db, organization, chat: new ChatService(connection.db) }, ctx, { id: `c_${randomUUID()}`, name, arguments: args });
    const docId = /doc_id=(doc_\w+)/.exec(run("create_document", { title: "Booking PRD", doc_type: "prd", content: "v1" }))![1];
    expect(connection.db.prepare("select task_id from documents where id = ?").get(docId)).toEqual({ task_id: taskId });
    expect(run("list_documents", {})).toContain(`from task "Write the booking PRD" task_id=${taskId})`);
    connection.close();
  });
});

describe("agent document tools", () => {
  test("documents an agent wrote can be found, read in pages, and updated", () => {
    const h = harness(false);
    const created = h.run("create_document", { title: "Booking PRD", doc_type: "prd", content: Array.from({ length: 3000 }, (_, i) => `Requirement ${i + 1}: booking flow`).join("\n") });
    const docId = /doc_id=(doc_\w+)/.exec(created)![1];
    h.run("create_document", { title: "Q3 Report", doc_type: "report", content: "Revenue grew." });

    expect(h.run("list_documents", { query: "booking" })).toMatch(new RegExp(`^Booking PRD \\(doc_id=${docId}, prd, draft, v1, \\d+ chars`));
    expect(h.run("list_documents", { doc_type: "report" })).toMatch(/^Q3 Report \(doc_id=doc_\w+, report/);
    expect(h.run("list_documents", { query: "kubernetes" })).toBe('No documents match "kubernetes".');

    const page = h.run("read_document", { doc_id: docId });
    expect(page.startsWith("Booking PRD (prd, draft, v1 of 1)\n\n[lines 1-")).toBe(true);
    expect(page).toContain("[more: pass start_line=");
    expect(h.run("read_document", { doc_id: docId, search: "Requirement 2999:" })).toContain("2999: Requirement 2999: booking flow");

    expect(h.run("update_document", { doc_id: docId, content: "Requirement 1: rebooking", reason: "Scope cut" })).toBe(`Updated document "Booking PRD" to v2 (doc_id=${docId}).`);
    expect(h.run("read_document", { doc_id: docId })).toBe("Booking PRD (prd, draft, v2 of 2)\n\nRequirement 1: rebooking");
    expect(h.run("read_document", { doc_id: docId, version: 1, search: "Requirement 2999:" })).toContain("v1 of 2");
    expect(h.run("read_document", { doc_id: docId, version: 1, search: "Requirement 2999:" })).toContain("2999: Requirement 2999: booking flow");
    expect(h.run("read_document", { doc_id: docId, history: true })).toMatch(/^v1 .* by agent agt_\w+ — created \(\d+ chars\)\nv2 .* by agent agt_\w+ — Scope cut \(24 chars\)$/);
    expect(h.run("read_document", { doc_id: docId, version: 3 })).toBe('ERROR: "version" must be a whole number from 1 to 2.');
    expect(h.run("read_document", { doc_id: docId, history: "yes" })).toBe('ERROR: "history" must be true or false.');
    expect(h.run("read_document", { doc_id: docId, history: true, version: 0 })).toBe("ERROR: history=true lists versions; it takes no version or paging arguments.");
    expect(h.run("read_document", { doc_id: docId, history: true, search: 123 })).toBe('ERROR: "search" is required.');
    expect(h.run("update_document", { doc_id: docId, content: "x" })).toBe('ERROR: "reason" is required.');
    expect(h.run("read_document", { doc_id: "doc_missing" })).toBe('ERROR: No document with doc_id "doc_missing".');
    expect(h.run("read_document", { doc_id: docId, search: 123 })).toBe('ERROR: "search" is required.');
    expect(h.run("list_documents", { doc_type: "novel" })).toBe('ERROR: "doc_type" must be one of: prd, spec, memo, report, runbook, proposal, policy, guide.');

    h.connection.db.prepare("delete from documents where id = ?").run(docId);
    expect(h.run("read_document", { doc_id: docId, history: true }).split("\n")).toHaveLength(2);
    expect(h.run("read_document", { doc_id: docId, version: 1, line_count: 1 })).toMatch(/^Booking PRD \(deleted, draft, v1 of 2\)/);
    h.connection.close();
  });

  test("in review mode an update is held in Approvals and the document keeps its version", () => {
    const h = harness(false);
    const docId = /doc_id=(doc_\w+)/.exec(h.run("create_document", { title: "Policy", doc_type: "policy", content: "v1" }))![1];
    h.connection.db.prepare("update organizations set review_deliverables = 1 where id = ?").run(h.orgId);
    expect(h.run("update_document", { doc_id: docId, content: "v2", reason: "Tighter wording" })).toContain("held for board review");
    const approvals = h.connection.db.prepare("select title from approvals where org_id = ? and kind = 'deliverable_review'").all(h.orgId);
    expect(approvals).toEqual([{ title: "Review policy update: Policy" }]);
    expect(h.connection.db.prepare("select version, content from documents where id = ?").get(docId)).toEqual({ version: 1, content: "v1" });
    h.connection.close();
  });
});