codex-chief-of-staff/test/tools.test.ts

import assert from "node:assert/strict";
import test from "node:test";

import { Client } from "@modelcontextprotocol/sdk/client/index.js";
import { InMemoryTransport } from "@modelcontextprotocol/sdk/inMemory.js";

import { AppServerClient } from "../src/app-server-client.js";
import { normalizeThread, type ContextBounds, type ThreadInspection } from "../src/domain.js";
import { createMcpServer, SERVER_INSTRUCTIONS } from "../src/mcp-server.js";
import { ThreadMonitor } from "../src/monitor.js";
import { MonitorStore, type WaitResult } from "../src/store.js";
import { ChiefOfStaffTools, type MonitorFacade } from "../src/tools.js";
import { FakeAppServer, rawThread, successfulHandshake } from "./fake-app-server.js";

class StoreFacade implements MonitorFacade {
  readonly store: MonitorStore;

  constructor(store: MonitorStore) {
    this.store = store;
  }

  snapshot() { return this.store.snapshot(); }
  ensureFresh() { return Promise.resolve(); }
  inspectThread(_threadId: string, _bounds: ContextBounds): Promise<ThreadInspection> {
    return Promise.reject(new Error("inspection not configured"));
  }
  waitForChange(
    afterRevision: number,
    timeoutMs: number,
    options: { signal?: AbortSignal; maxChanges?: number } = {},
  ): Promise<WaitResult> {
    return this.store.waitForChange(afterRevision, timeoutMs, options);
  }
}

function storeWith(rawThreads: Record<string, unknown>[]): MonitorStore {
  const store = new MonitorStore();
  store.replaceSnapshot(rawThreads.map(normalizeThread), new Set());
  return store;
}

async function connectMcp(tools: ChiefOfStaffTools) {
  const server = createMcpServer(tools);
  const client = new Client({ name: "test-client", version: "1.0.0" });
  const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
  await Promise.all([server.connect(serverTransport), client.connect(clientTransport)]);
  return { server, client };
}

test("overview is recency ordered and excludes exec, subagent, guardian, and message bodies", async () => {
  const store = storeWith([
    rawThread("interactive-old", { updatedAt: 10, recencyAt: 10, source: "cli" }),
    rawThread("interactive-new", { updatedAt: 50, recencyAt: 50, source: "vscode" }),
    rawThread("exec", { updatedAt: 60, recencyAt: 60, source: "exec" }),
    rawThread("sub", { updatedAt: 70, recencyAt: 70, source: { subAgent: { type: "thread_spawn" } }, parentThreadId: "parent" }),
    rawThread("guardian", { updatedAt: 80, recencyAt: 80, source: "cli", agentRole: "guardian" }),
  ]);
  const result = await new ChiefOfStaffTools(new StoreFacade(store)).overview({});
  const threads = result.threads as Array<Record<string, unknown>>;
  assert.deepEqual(threads.map((thread) => thread.id), ["interactive-new", "interactive-old"]);
  assert.equal(JSON.stringify(result).includes("turns"), false);
  assert.equal((result.coverage as Record<string, unknown>).freshness, "current");
});

test("list_threads covers cwd roots, descendants, non-matches, filters, pagination, and truncation", async () => {
  const nowSeconds = 2_000;
  const store = storeWith([
    rawThread("root", { cwd: "/Users/dnabracz/Work", recencyAt: 1_990, updatedAt: 1_990, status: { type: "active", activeFlags: [] } }),
    rawThread("descendant", { cwd: "/Users/dnabracz/Work/scm/repo", recencyAt: 1_980, updatedAt: 1_980, status: { type: "idle" } }),
    rawThread("similar", { cwd: "/Users/dnabracz/Workspace", recencyAt: 1_970, updatedAt: 1_970 }),
    rawThread("old", { cwd: "/Users/dnabracz/Work/old", recencyAt: 100, updatedAt: 100 }),
  ]);
  const tools = new ChiefOfStaffTools(new StoreFacade(store), () => new Date(nowSeconds * 1_000));
  const first = await tools.listThreads({ cwdPrefix: "/Users/dnabracz/Work", recentSeconds: 100, limit: 1 });
  assert.deepEqual((first.threads as Array<Record<string, unknown>>).map((thread) => thread.id), ["root"]);
  assert.equal(first.truncated, true);
  assert.equal(first.nextCursor, "1");
  const second = await tools.listThreads({ cwdPrefix: "/Users/dnabracz/Work", recentSeconds: 100, limit: 1, cursor: "1" });
  assert.deepEqual((second.threads as Array<Record<string, unknown>>).map((thread) => thread.id), ["descendant"]);
  assert.equal(second.nextCursor, null);
  const active = await tools.listThreads({ statuses: ["active"] });
  assert.deepEqual((active.threads as Array<Record<string, unknown>>).map((thread) => thread.id), ["root"]);
});

test("MCP client lists exactly four observational tools", async () => {
  const { client, server } = await connectMcp(new ChiefOfStaffTools(new StoreFacade(storeWith([]))));
  assert.equal(client.getInstructions(), SERVER_INSTRUCTIONS);
  assert.match(client.getInstructions() ?? "", /call overview first/i);
  assert.match(client.getInstructions() ?? "", /Do not browse the web/i);
  const listed = await client.listTools();
  assert.deepEqual(listed.tools.map((tool) => tool.name), [
    "overview",
    "list_threads",
    "inspect_thread",
    "wait_for_change",
  ]);
  assert.match(listed.tools[0]?.description ?? "", /Use first for Chief-of-Staff requests/i);
  await Promise.all([client.close(), server.close()]);
});

test("wait_for_change MCP integration covers status, connection, and timeout", async () => {
  const store = storeWith([rawThread("a")]);
  const { client, server } = await connectMcp(new ChiefOfStaffTools(new StoreFacade(store)));

  let revision = store.revision;
  queueMicrotask(() => store.updateStatus("a", { type: "active", activeFlags: [] }));
  const status = await client.callTool({ name: "wait_for_change", arguments: { afterRevision: revision, timeoutMs: 100 } });
  assert.equal((status.structuredContent as Record<string, unknown>).timedOut, false);

  revision = store.revision;
  queueMicrotask(() => store.setConnection("disconnected", "test disconnect"));
  const connection = await client.callTool({ name: "wait_for_change", arguments: { afterRevision: revision, timeoutMs: 100 } });
  const changes = (connection.structuredContent as Record<string, unknown>).changes as Array<Record<string, unknown>>;
  assert.equal(changes[0]?.kind, "connection");

  const timedOut = await client.callTool({ name: "wait_for_change", arguments: { afterRevision: store.revision, timeoutMs: 5 } });
  assert.equal((timedOut.structuredContent as Record<string, unknown>).timedOut, true);
  await Promise.all([client.close(), server.close()]);
});

test("every published tool keeps upstream traffic inside the read-only allowlist", async () => {
  const fake = new FakeAppServer(successfulHandshake((request, server) => {
    if (request.method === "thread/list") {
      server.respond(request, { data: [rawThread("a")], nextCursor: null, backwardsCursor: null });
    } else if (request.method === "thread/loaded/list") {
      server.respond(request, { data: [], nextCursor: null });
    } else if (request.method === "thread/read") {
      server.respond(request, { thread: rawThread("a") });
    }
  }));
  const monitor = new ThreadMonitor({
    clientFactory: () => new AppServerClient({
      url: "ws://127.0.0.1:4500",
      webSocketFactory: fake.factory,
      requestTimeoutMs: 50,
      connectTimeoutMs: 50,
    }),
    waitRefreshIntervalMs: 10,
  });
  monitor.start();
  await monitor.waitUntilSynchronized();
  const { client, server } = await connectMcp(new ChiefOfStaffTools(monitor));
  await client.callTool({ name: "overview", arguments: {} });
  await client.callTool({ name: "list_threads", arguments: { includeSubagents: true } });
  const inspection = await client.callTool({ name: "inspect_thread", arguments: { threadId: "a" } });
  assert.equal(inspection.isError, undefined);
  await client.callTool({ name: "wait_for_change", arguments: { afterRevision: monitor.snapshot().revision, timeoutMs: 1 } });
  const methods = fake.frames.map((frame) => frame.method as string);
  const permitted = new Set(["initialize", "initialized", "thread/list", "thread/loaded/list", "thread/read"]);
  assert.deepEqual(methods.filter((method) => !permitted.has(method)), []);
  assert.equal(methods.includes("thread/read"), true);
  await Promise.all([client.close(), server.close()]);
  await monitor.stop();
});

test("inspect_thread returns structured upstream errors without resuming or subscribing", async () => {
  const fake = new FakeAppServer(successfulHandshake((request, server) => {
    if (request.method === "thread/list") server.respond(request, { data: [], nextCursor: null, backwardsCursor: null });
    else if (request.method === "thread/loaded/list") server.respond(request, { data: [], nextCursor: null });
    else if (request.method === "thread/read") server.fail(request, -32602, "thread not found");
  }));
  const monitor = new ThreadMonitor({ clientFactory: () => new AppServerClient({
    url: "ws://127.0.0.1:4500",
    webSocketFactory: fake.factory,
    requestTimeoutMs: 50,
  }) });
  monitor.start();
  await monitor.waitUntilSynchronized();
  const result = await new ChiefOfStaffTools(monitor).call("inspect_thread", { threadId: "missing" });
  assert.equal(result.isError, true);
  assert.equal(((result.structuredContent.error as Record<string, unknown>).code), "not_found");
  const methods = fake.frames.map((frame) => frame.method);
  assert.equal(methods.includes("thread/resume"), false);
  assert.equal(methods.includes("thread/unsubscribe"), false);
  await monitor.stop();
});