next-tjwater-agent/tests/mcp/queryServer.test.ts

import { afterEach, expect, it, mock } from "bun:test";
import { Client, StreamableHTTPClientTransport, type Tool } from "@modelcontextprotocol/client";
import { startQueryMcpServer } from "../../src/mcp/query-server.js";
import { createLocalQueryExecutor } from "../../src/mcp/local-query-tools.js";
import { setRuntimeSessionContext, removeRuntimeSessionContext } from "../../src/runtime/sessionContext.js";
import { LOCAL_QUERY_TOOLS, PRODUCT_MCP_ANALYSIS_TOOLS, queryMcpToolId } from "../../src/runtime/product-mcp-tools.js";
import type { QueryOwner } from "../../src/product/mcp-query-client.js";
import type { SessionTranscriptStore } from "../../src/sessions/transcriptStore.js";

const owner = { sessionId: "query-owner", clientSessionId: "chat-owner", actorKey: "actor-a", projectKey: "project-a", projectId: "project-a", traceId: "test", purpose: "ordinary_chat" as const };
const servers: ReturnType<typeof startQueryMcpServer>[] = [];
const clients: Client[] = [];
afterEach(async () => {
  removeRuntimeSessionContext(owner.sessionId);
  await Promise.all(clients.splice(0).map((client) => client.close()));
  await Promise.all(servers.splice(0).map((server) => server.close()));
});
const productTool: Tool = { name: "get_task_context", description: "Bound task", inputSchema: { type: "object", properties: { report_revision: { type: "integer", minimum: 1 } }, additionalProperties: false } };

it("advertises simulation budgets and derives the native connection ceiling from the catalog", async () => {
  const simulation: Tool = { name: "run_simulation", inputSchema: { type: "object", properties: {
    timeout_seconds: { type: "integer", minimum: 1, maximum: 3600, default: 120 },
  }, additionalProperties: false } };
  const { server, client, call, query } = await setup({ productTools: [simulation] });
  setRuntimeSessionContext(owner);
  expect(server.executionTimeoutMs).toBe(3_665_000);
  expect((await client.listTools()).tools[0].inputSchema.properties?.timeout_seconds).toMatchObject({ maximum: 3600 });
  expect((await call("run_simulation", { timeout_seconds: 900 })).isError).not.toBeTrue();
  expect((query.mock.calls[0] as unknown as unknown[])[3]).toEqual({ timeout_seconds: 900 });
  try { await call("run_simulation", { timeout_seconds: 3601 }); } catch { /* SDK schema rejection */ }
  expect(query).toHaveBeenCalledTimes(1);
});
it("discovers calculations natively but denies verified readers before upstream dispatch", async () => {
  const context = { ...owner, purpose: "task_execution" as const, taskId: "task-a" };
  const productTools = PRODUCT_MCP_ANALYSIS_TOOLS.map((name) => ({ name, inputSchema: { type: "object" as const, properties: {}, additionalProperties: false } }));
  const { client, call, query } = await setup({ productTools, resolveReader: async () => ({ context, runtimeSessionId: owner.sessionId, taskId: "task-a" }) });
  setRuntimeSessionContext(context);
  expect((await client.listTools()).tools.map((tool) => tool.name)).toEqual(expect.arrayContaining([...PRODUCT_MCP_ANALYSIS_TOOLS]));
  for (const name of PRODUCT_MCP_ANALYSIS_TOOLS) {
    expect((await call(name)).isError).not.toBeTrue();
    expect((await call(name, {}, "verified-reader")).isError).toBeTrue();
  }
  expect(query).toHaveBeenCalledTimes(PRODUCT_MCP_ANALYSIS_TOOLS.length);
});
async function setup(extra: Partial<Parameters<typeof startQueryMcpServer>[0]> = {}) {
  const executeLocal = mock(async () => ({ safe: true }));
  const query = mock(async () => ({ content: [{ type: "text" as const, text: "evidence" }], structuredContent: { evidence_id: "E1", data: { value: 1 } } }));
  const server = startQueryMcpServer({ internalToken: "MCP_TEST_ONLY", productClient: { query }, productTools: [productTool], resolveReader: async () => null, executeLocal, ...extra });
  servers.push(server);
  const client = new Client({ name: "query-test", version: "1" }, { versionNegotiation: { mode: { pin: "2026-07-28" } } });
  clients.push(client);
  await client.connect(new StreamableHTTPClientTransport(new URL(server.url), { requestInit: { headers: { "x-agent-internal-token": "MCP_TEST_ONLY" } } }));
  const call = (name: string, args: Record<string, unknown> = {}, id: string | undefined = owner.sessionId) => client.callTool({ name, arguments: args, ...(id ? { _meta: { "ai.opencode/sessionID": id } } : {}) });
  return { server, client, call, executeLocal, query };
}

it("discovers native MCP schemas without duplicate plugins, hidden identity or mixed managers", async () => {
  const { client, server } = await setup();
  const catalog = await client.listTools();
  expect(catalog.tools.map((tool) => tool.name).sort()).toEqual([...LOCAL_QUERY_TOOLS, "get_task_context"].sort());
  for (const tool of catalog.tools) {
    expect(tool.inputSchema.properties).not.toHaveProperty("session_id");
    expect(tool.inputSchema.properties).not.toHaveProperty("authorization");
  }
  expect(JSON.stringify(catalog)).not.toContain("memory_manager");
  expect(JSON.stringify(catalog)).not.toContain("skill_manager");
  expect(catalog.tools.some((tool) => tool.name === "web_search")).toBeFalse();
  for (const headers of [{}, { "x-agent-internal-token": "wrong" }, { "x-agent-internal-token": "MCP_TEST_ONLY", Origin: "http://localhost" }]) {
    expect((await fetch(server.url, { method: "POST", headers: headers as Record<string, string>, body: "{}" })).status).toBe(403);
  }
});

it("requires host-authenticated session metadata and blocks personal, review and unbound project queries", async () => {
  const { call, executeLocal, query } = await setup();
  setRuntimeSessionContext(owner);
  expect((await call("session_search", { query: "pressure" })).isError).not.toBeTrue();
  expect((await call("session_search", { query: "pressure" }, "foreign")).isError).toBeTrue();
  expect((await call("session_search", { query: "pressure" }, "")).isError).toBeTrue();
  setRuntimeSessionContext({ ...owner, projectId: undefined, projectKey: "personal-actor-a" });
  expect((await call("get_task_context")).isError).toBeTrue();
  expect((await call("geocode", { keyword: "北京" })).isError).toBeTrue();
  expect((await call("session_search", { query: "pressure" })).isError).not.toBeTrue();
  setRuntimeSessionContext({ ...owner, learningMode: "review" });
  expect((await call("session_search", { query: "pressure" })).isError).toBeTrue();
  expect(executeLocal).toHaveBeenCalledTimes(2);
  expect(query).not.toHaveBeenCalled();
});

it("delegates only product reads to a verified child and retains parent revision and evidence", async () => {
  const taskOwner = { ...owner, purpose: "task_execution" as const, taskId: "task-a", reportRevision: 3 };
  const { call, query, executeLocal } = await setup({ resolveReader: async (id) => id === "verified-child" ? { context: taskOwner, runtimeSessionId: owner.sessionId, taskId: "task-a" } : null });
  const result = await call("get_task_context", {}, "verified-child");
  expect(result.structuredContent).toEqual({ evidence_id: "E1", data: { value: 1 } });
  const args = query.mock.calls[0] as unknown as [string, QueryOwner, string, unknown];
  expect(args[0]).toBe("verified-child");
  expect(args[1].runtimeSessionId).toBe(owner.sessionId);
  expect(args[3]).toEqual({ report_revision: 3 });
  expect((await call("session_search", { query: "private" }, "verified-child")).isError).toBeTrue();
  expect((await call("read_agent_messages", {}, "verified-child")).isError).toBeTrue();
  expect(executeLocal).not.toHaveBeenCalled();
});

it("rejects identity overrides, arbitrary URLs and oversized queries before executing", async () => {
  setRuntimeSessionContext(owner);
  const { call, executeLocal } = await setup();
  for (const [name, args] of [
    ["session_search", { query: "x", session_id: "foreign" }],
    ["geocode", { keyword: "北京", url: "https://foreign.test" }],
    ["geocode", { keyword: "x".repeat(501) }],
    ["session_search", { query: "x", max_results: 51 }],
    ["list_investigation_tasks", { offset: 10001 }],
    ["get_investigation_task", { task_id: "made-up" }],
    ["read_agent_messages", { limit: 0 }],
  ] as const) {
    let rejected = false;
    try { rejected = (await call(name, args)).isError === true; } catch { rejected = true; }
    expect(rejected).toBeTrue();
  }
  expect(executeLocal).not.toHaveBeenCalled();
});

it("searches only the owning transcript scope and pages the owning native checklist", async () => {
  const search = mock(async () => []);
  const readTodos = mock(async () => Array.from({ length: 25 }, (_, id) => ({ id })));
  const executeLocal = createLocalQueryExecutor({ transcripts: { search } as unknown as SessionTranscriptStore, readTodos, messageQuery: async () => ({}) });
  const { call } = await setup({ executeLocal });
  setRuntimeSessionContext(owner);
  await call("session_search", { query: "pressure" });
  expect(search.mock.calls[0] as unknown).toEqual([{ actorKey: "actor-a", projectKey: "project-a" }, "pressure", 20]);
  const first = await call("todoread");
  const body = JSON.parse((first.content[0] as { text: string }).text);
  expect(body.items).toHaveLength(20);
  expect(body.page).toEqual({ limit: 20, offset: 0, total: 25, has_more: true });
  expect((readTodos.mock.calls[0] as unknown as unknown[])[0]).toBe(owner.sessionId);
});

it("returns safe MCP errors instead of leaking upstream credentials", async () => {
  setRuntimeSessionContext(owner);
  const { call } = await setup({ executeLocal: async () => { throw Object.assign(new Error("Bearer PRIVATE_TOKEN https://private.test"), { status: 401 }); } });
  const result = await call("session_search", { query: "x" });
  expect(result.isError).toBeTrue();
  expect(JSON.stringify(result)).toContain("401");
  expect(JSON.stringify(result)).not.toContain("PRIVATE_TOKEN");
  expect(queryMcpToolId("session_search")).toBe("mcp_tjwater_session_search");
});

it("replaces unavailable evidence resources with the authorized paging tool without changing provenance", async () => {
  setRuntimeSessionContext(owner);
  const uri = "tjwater://tasks/task-a/evidence/E1";
  const upstream = { content: [{ type: "text" as const, text: `Read ${uri}` }, { type: "resource_link" as const, name: "query-evidence", uri }],
    structuredContent: { source: "query_scada_readings", source_ref: "saved:1", evidence_id: "E1", origin: "synthetic", truncated: true,
      data: { resource_uri: uri, preview: "The complete query snapshot is saved in task evidence" },
      data_gaps: ["Inline result abbreviated; read the saved evidence resource"] } };
  const { call } = await setup({ productClient: { query: async () => upstream } });
  const result = await call("get_task_context");
  expect(JSON.stringify(result)).not.toContain("tjwater://");
  expect(result.content.some((item) => item.type === "resource_link")).toBeFalse();
  expect(result.structuredContent).toMatchObject({ source_ref: "saved:1", evidence_id: "E1", origin: "synthetic", truncated: true,
    data: { evidence_read: { tool: "mcp_tjwater_get_task_evidence", arguments: { evidence_id: "E1", limit: 20, offset: 0 } } } });
  expect(upstream.structuredContent.data.resource_uri).toBe(uri);
});