AgentCompany/tests/unit/connector-grants.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 { ConnectorGrantsService, grantLimitRefusal } from "../../src/main/services/composio/connector-grants";
import { dispatchAgentTool, type AgentToolContext } from "../../src/main/services/agent-tools";

function call(name: string, args: Record<string, unknown>) {
  return { id: `call_${randomUUID()}`, name, arguments: args };
}

interface Harness {
  connection: DatabaseConnection;
  orgId: string;
  ceoId: string;
  danaId: string;
  employeeId: string;
  ceoCtx: AgentToolContext;
  danaCtx: AgentToolContext;
  deps: { db: DatabaseConnection["db"]; organization: OrganizationService; chat: ChatService };
}

function harness(): Harness {
  const connection = new DatabaseConnection(join(tmpdir(), `ac-connector-${randomUUID()}.db`));
  const organization = new OrganizationService(connection.db);
  let orgId = "";
  let ceoId = "";
  let danaId = "";
  connection.transaction(() => {
    orgId = organization.createOrganization({ organizationName: "OrbitOps", employeeName: "Yang", mission: "m", goal: "g" });
    const providerId = (connection.db.prepare("select id from providers where org_id = ? and provider_key = 'openai'").get(orgId) as { id: string }).id;
    ceoId = organization.createAgent({
      orgId, name: "CEO", role: "CEO", department: "Executive",
      modelProviderId: providerId, modelId: "gpt-5.4", budgetCents: 25000, managerAgentId: null, roleLevel: "c_suite",
    });
    danaId = organization.createAgent({
      orgId, name: "Dana", role: "Engineer", department: "Engineering",
      modelProviderId: providerId, modelId: "gpt-5.4", budgetCents: 1000, managerAgentId: ceoId,
    });
  });
  const context = (agentId: string): AgentToolContext => ({
    orgId,
    agent: { id: agentId, modelProviderId: null, modelId: null },
    taskProjectId: null,
    taskGoalId: null, taskId: null, sourceThreadId: null,
    workspace: null,
  });
  return {
    connection,
    orgId,
    ceoId,
    danaId,
    employeeId: (connection.db.prepare("select id from employees where org_id = ?").get(orgId) as { id: string }).id,
    ceoCtx: context(ceoId),
    danaCtx: context(danaId),
    deps: { db: connection.db, organization, chat: new ChatService(connection.db) },
  };
}

function connectToolkit(h: Harness, slug: string, name: string): void {
  const now = new Date().toISOString();
  h.connection.db.prepare(`
    insert into connector_accounts(id, org_id, toolkit_slug, toolkit_name, label, auth_config_id, status, created_at, updated_at)
    values (?, ?, ?, ?, 'main', ?, 'ACTIVE', ?, ?)
  `).run(`ca_${slug}`, h.orgId, slug, name, `ac_${slug}`, now, now);
}

describe("connector grants", () => {
  test("granting requires an actively connected toolkit", () => {
    const h = harness();
    const grants = new ConnectorGrantsService(h.connection.db);
    expect(() => grants.grant({
      orgId: h.orgId, agentId: h.danaId, toolkitSlug: "github", accountLabel: null, limits: { actions: null, fixedArguments: null }, grantedByType: "employee", grantedById: h.employeeId,
    })).toThrow(/No active connection/);
  });

  test("the CEO grants and revokes toolkits for teammates", () => {
    const h = harness();
    connectToolkit(h, "github", "GitHub");
    const granted = dispatchAgentTool(h.deps, h.ceoCtx, call("grant_connector", { agent_name: "Dana", toolkit: "github" }));
    expect(granted).toContain("Dana now has access to the github connector");

    const grants = new ConnectorGrantsService(h.connection.db);
    expect(grants.toolkitsForAgent(h.orgId, h.danaId)).toEqual(["github"]);
    // The CEO holds every active connection by role, without a grant row.
    expect(grants.toolkitsForAgent(h.orgId, h.ceoId)).toEqual(["github"]);

    const listing = dispatchAgentTool(h.deps, h.danaCtx, call("list_connectors", {}));
    expect(listing).toContain("GitHub (github), account \"main\": granted to Dana");

    const revoked = dispatchAgentTool(h.deps, h.ceoCtx, call("revoke_connector", { agent_name: "Dana", toolkit: "github" }));
    expect(revoked).toContain("no longer has access");
    expect(grants.toolkitsForAgent(h.orgId, h.danaId)).toEqual([]);
  });

  test("a toolkit with several accounts: grants and actions name the account", () => {
    const h = harness();
    const now = new Date().toISOString();
    for (const label of ["main", "sales"]) {
      h.connection.db.prepare(`
        insert into connector_accounts(id, org_id, toolkit_slug, toolkit_name, label, auth_config_id, status, created_at, updated_at)
        values (?, ?, 'github', 'GitHub', ?, 'ac_github', 'ACTIVE', ?, ?)
      `).run(`ca_gh_${label}`, h.orgId, label, now, now);
    }
    const grants = new ConnectorGrantsService(h.connection.db);
    expect(dispatchAgentTool(h.deps, h.ceoCtx, call("grant_connector", { agent_name: "Dana", toolkit: "github" })))
      .toBe('ERROR: "github" has several accounts (main, sales); name one with account.');
    expect(dispatchAgentTool(h.deps, h.ceoCtx, call("grant_connector", { agent_name: "Dana", toolkit: "github", account: "Sales" })))
      .toContain('through its "sales" account');
    expect(dispatchAgentTool(h.deps, h.ceoCtx, call("grant_connector", { agent_name: "Dana", toolkit: "github", account: null })))
      .toBe('ERROR: "account" is required.');
    expect(grants.accountForAgent(h.orgId, h.danaId, "github", null).id).toBe("ca_gh_sales");
    expect(() => grants.accountForAgent(h.orgId, h.danaId, "github", "main")).toThrow('Your grant for "github" is the "sales" account, not "main".');
    expect(() => grants.accountForAgent(h.orgId, h.ceoId, "github", null)).toThrow('"github" has several accounts (main, sales)');
    expect(grants.accountForAgent(h.orgId, h.ceoId, "github", "main").id).toBe("ca_gh_main");
    const listing = dispatchAgentTool(h.deps, h.danaCtx, call("list_connectors", {}));
    expect(listing).toContain('GitHub (github), account "main": no grants yet');
    expect(listing).toContain('GitHub (github), account "sales": granted to Dana');
    h.connection.db.prepare("update connector_accounts set status = 'EXPIRED' where id = 'ca_gh_sales'").run();
    expect(grants.toolkitsForAgent(h.orgId, h.danaId)).toEqual([]);
    expect(grants.toolkitsForAgent(h.orgId, h.ceoId)).toEqual(["github"]);
  });

  test("the CEO grants a toolkit narrowed to some actions and argument values", () => {
    const h = harness();
    connectToolkit(h, "gmail", "Gmail");
    expect(dispatchAgentTool(h.deps, h.ceoCtx, call("grant_connector", { agent_name: "Dana", toolkit: "gmail", actions: ["gmail_fetch_emails"], fixed_arguments: { label: "support" } })))
      .toContain('through its "main" account (only GMAIL_FETCH_EMAILS; limited to {"label":"support"})');
    expect(dispatchAgentTool(h.deps, h.danaCtx, call("list_connectors", {})))
      .toContain('Gmail (gmail), account "main": granted to Dana (only GMAIL_FETCH_EMAILS; limited to {"label":"support"})');
    expect(new ConnectorGrantsService(h.connection.db).limitsForAgent(h.orgId, h.danaId, "gmail")).toEqual({ actions: ["GMAIL_FETCH_EMAILS"], fixedArguments: { label: "support" } });
    expect(dispatchAgentTool(h.deps, h.ceoCtx, call("grant_connector", { agent_name: "Dana", toolkit: "gmail", actions: [] })))
      .toBe("ERROR: actions must name at least one action, none blank.");
    expect(dispatchAgentTool(h.deps, h.ceoCtx, call("grant_connector", { agent_name: "Dana", toolkit: "gmail", fixed_arguments: ["label"] })))
      .toBe('ERROR: "fixed_arguments" must be an object of argument values.');
    const nullRepo = { actions: null, fixedArguments: { repo: null } };
    const action = { name: "GITHUB_LIST_ISSUES", parameters: { properties: { repo: {} } } };
    expect(grantLimitRefusal(nullRepo, "github", action, {})).toBe("Your grant for github is limited to repo=null; pass repo with that value.");
    expect(grantLimitRefusal(nullRepo, "github", action, { repo: null })).toBeNull();
    const proto = { actions: null, fixedArguments: JSON.parse('{"__proto__": {}}') as Record<string, unknown> };
    expect(grantLimitRefusal(proto, "github", action, {})).toBe("Your grant for github is limited to __proto__={}, and GITHUB_LIST_ISSUES takes no __proto__.");
    // A grant without limits replaces a limited one.
    dispatchAgentTool(h.deps, h.ceoCtx, call("grant_connector", { agent_name: "Dana", toolkit: "gmail" }));
    expect(new ConnectorGrantsService(h.connection.db).limitsForAgent(h.orgId, h.danaId, "gmail")).toEqual({ actions: null, fixedArguments: null });
  });

  test("the board grants and revokes as a named employee of the company", () => {
    const h = harness();
    connectToolkit(h, "github", "GitHub");
    const grants = new ConnectorGrantsService(h.connection.db);
    const grant = (grantedById: string | null) => grants.grant({
      orgId: h.orgId, agentId: h.danaId, toolkitSlug: "github", accountLabel: "main", limits: { actions: null, fixedArguments: null }, grantedByType: "employee", grantedById,
    });
    expect(() => grant(null)).toThrow("A grant by an employee must name an employee of this organization.");
    expect(() => grant("emp_elsewhere")).toThrow("A grant by an employee must name an employee of this organization.");
    expect(grants.toolkitsForAgent(h.orgId, h.danaId)).toEqual([]);
    grant(h.employeeId);
    expect(grants.toolkitsForAgent(h.orgId, h.danaId)).toEqual(["github"]);
    grants.revoke({ orgId: h.orgId, agentId: h.danaId, toolkitSlug: "github" });
    expect(grants.toolkitsForAgent(h.orgId, h.danaId)).toEqual([]);
    // The page never replaces a grant the CEO shaped, even from a stale selection.
    grants.grant({ orgId: h.orgId, agentId: h.danaId, toolkitSlug: "github", accountLabel: null, limits: { actions: ["GITHUB_LIST_ISSUES"], fixedArguments: null }, grantedByType: "agent", grantedById: h.ceoId });
    expect(() => grants.grantFromBoard({ orgId: h.orgId, agentId: h.danaId, toolkitSlug: "github", accountLabel: "main", employeeId: h.employeeId }))
      .toThrow('This agent already holds github through its "main" account; revoke that grant first.');
    expect(grants.limitsForAgent(h.orgId, h.danaId, "github")).toEqual({ actions: ["GITHUB_LIST_ISSUES"], fixedArguments: null });
  });

  test("staff cannot hand out connector access", () => {
    const h = harness();
    connectToolkit(h, "github", "GitHub");
    const denied = dispatchAgentTool(h.deps, h.danaCtx, call("grant_connector", { agent_name: "Dana", toolkit: "github" }));
    expect(denied).toContain("ERROR");
    expect(new ConnectorGrantsService(h.connection.db).toolkitsForAgent(h.orgId, h.danaId)).toEqual([]);
  });

  test("a disconnected toolkit stops counting for its grants", () => {
    const h = harness();
    connectToolkit(h, "slack", "Slack");
    const grants = new ConnectorGrantsService(h.connection.db);
    grants.grant({ orgId: h.orgId, agentId: h.danaId, toolkitSlug: "slack", accountLabel: null, limits: { actions: null, fixedArguments: null }, grantedByType: "employee", grantedById: h.employeeId });
    expect(grants.toolkitsForAgent(h.orgId, h.danaId)).toEqual(["slack"]);
    h.connection.db.prepare("update connector_accounts set status = 'EXPIRED' where org_id = ? and toolkit_slug = 'slack'").run(h.orgId);
    expect(grants.toolkitsForAgent(h.orgId, h.danaId)).toEqual([]);
  });
});