next-tjwater-agent/tests/product/mcpQueryClient.test.ts

import { afterEach, describe, expect, it } from "bun:test";
import { McpQueryClient, validateMcpSettings, type McpHttpFetch, assertExchangedIdentity } from "../../src/product/mcp-query-client.js";
import { bindSessionCredential, credentialRefreshCoordinator, refreshSessionCredential, releaseSessionCredential } from "../../src/product/credential-refresh.js";
import { PRODUCT_MCP_TOOLS } from "../../src/runtime/product-mcp-tools.js";

const sessions: string[] = [];
afterEach(() => { for (const id of sessions.splice(0)) releaseSessionCredential(id); });
const settings = { issuer: "http://localhost/realms/test", resource: "http://localhost/mcp", audience: "tjwater-mcp", clientId: "tjwater-agent-mcp", clientSecret: "fixture-secret" };
const owner = (id: string) => ({ context: { sessionId: id, actorKey: "actor", clientSessionId: id, projectId: "project", projectKey: "project", traceId: "test", purpose: "ordinary_chat" as const }, runtimeSessionId: id });
const catalog = PRODUCT_MCP_TOOLS.map((name) => ({ name, inputSchema: { type: "object" as const, additionalProperties: false, properties: { context_handle: { type: "string" }, element_type: { type: "string" }, ...(name === "run_simulation" ? { timeout_seconds: { type: "integer", default: 120, maximum: 3600 } } : {}) }, required: ["context_handle", "element_type"] } }));
// Test through a transport-independent fake only where the concern is token ownership/renewal.
function fixture(http: McpHttpFetch, authMethod = "client_secret_post") {
  const client = new McpQueryClient(settings, (async (url, init) => {
    if (String(url) === settings.resource && !new Headers(init?.headers).has("Authorization")) return new Response(null, {
      status: 401, headers: { "WWW-Authenticate": 'Bearer resource_metadata="http://localhost/.well-known/oauth-protected-resource/mcp"' },
    });
    if (String(url).includes("oauth-protected-resource")) return Response.json({ resource: settings.resource, authorization_servers: [settings.issuer] });
    if (String(url).includes(".well-known")) return Response.json({ issuer: settings.issuer, authorization_endpoint: `${settings.issuer}/authorize`, response_types_supported: ["code"], token_endpoint: `${settings.issuer}/protocol/openid-connect/token`, token_endpoint_auth_methods_supported: [authMethod] });
    return http(url, init);
  }) as McpHttpFetch);
  Object.assign(client, { tools: catalog, withClient: async (_token: string, _project: string, _signal: AbortSignal, run: (client: unknown) => Promise<unknown>) => run({ listTools: async () => ({ tools: catalog }), callTool: async (params: unknown) => ({ content: [], structuredContent: params }) }) });
  return client;
}
const jwt = (id: string, overrides: Record<string, unknown> = {}) => `e30.${Buffer.from(JSON.stringify({ sub: "test-subject", iss: settings.issuer, aud: settings.audience, azp: settings.clientId, jti: id, ...overrides })).toString("base64url")}.fixture`;
const resourceToken = jwt("resource");
const token = () => Response.json({ access_token: resourceToken, token_type: "Bearer", expires_in: 300 });
function credential(id: string) { sessions.push(id); bindSessionCredential(id, { authorization: `Bearer ${jwt(id)}`, projectId: "project" }); }

it("forwards the selected simulation budget through HTTP and SDK while rejecting excess before OAuth", async () => {
  const id = "mcp-calculation-budget"; credential(id);
  let requests = 0, transportTimeout = 0, sdkTimeout = 0;
  const client = fixture((async (url, init) => {
    requests++;
    if (String(url).includes("/token")) return token();
    if (init?.method === "POST") return Response.json({ context_handle: "x".repeat(43) });
    return new Response(null, { status: 204 });
  }) as McpHttpFetch);
  Object.assign(client, { withClient: async (_token: string, _project: string, _signal: AbortSignal, run: (client: unknown) => Promise<unknown>, timeout: number) => {
    transportTimeout = timeout;
    return run({ listTools: async () => ({ tools: catalog }), callTool: async (params: unknown, options: { timeout: number }) => {
      sdkTimeout = options.timeout;
      return { content: [], structuredContent: params };
    } });
  } });
  await expect(client.query(id, owner(id), "run_simulation", { timeout_seconds: 3601 }, "excess")).rejects.toMatchObject({ status: 400 });
  expect(requests).toBe(0);
  const result = await client.query(id, owner(id), "run_simulation", { timeout_seconds: 900 }, "long");
  expect(transportTimeout).toBe(905_000);
  expect(sdkTimeout).toBe(905_000);
  expect(result.structuredContent).toMatchObject({ arguments: { timeout_seconds: 900, context_handle: "x".repeat(43) } });
});

it("requires HTTPS except loopback, and fixed resource paths", () => {
  expect(() => validateMcpSettings({ ...settings, issuer: "http://remote/realm" })).toThrow();
  expect(() => validateMcpSettings({ ...settings, resource: "http://localhost/api" })).toThrow();
  expect(() => validateMcpSettings({ ...settings, resource: "https://user:password@example/mcp" })).toThrow();
});

it("allows only an explicit private development issuer while keeping MCP resources protected", () => {
  const issuer = "http://192.168.1.114:8088/realms/tjwater";
  const local = { ...settings, issuer, devHttpIssuer: issuer, runtimeEnvironment: "development" };
  expect(() => validateMcpSettings(local)).not.toThrow();
  expect(() => validateMcpSettings({ ...local, runtimeEnvironment: "production" })).toThrow();
  expect(() => validateMcpSettings({ ...local, runtimeEnvironment: undefined })).toThrow();
  expect(() => validateMcpSettings({ ...local, devHttpIssuer: undefined })).toThrow();
  expect(() => validateMcpSettings({ ...local, issuer: `${issuer}/other` })).toThrow();
  expect(() => validateMcpSettings({ ...local, resource: "http://192.168.1.114:8088/mcp" })).toThrow();
  expect(() => validateMcpSettings({ ...local, issuer: "http://example.com/realm", devHttpIssuer: "http://example.com/realm" })).toThrow();
});

it("rejects personal queries and model supplied scope before any network call", async () => {
  let count = 0;
  const client = fixture((async () => { count++; return token(); }) as McpHttpFetch);
  await expect(client.query("p", { ...owner("p"), context: { ...owner("p").context, projectKey: "personal-actor" } }, "list_elements", {}, "call")).rejects.toMatchObject({ status: 403 });
  await expect(client.query("p", owner("p"), "list_elements", { context_handle: "MODEL", project_id: "other" }, "call")).rejects.toMatchObject({ status: 400 });
  expect(count).toBe(0);
  expect(client.catalog()[0].inputSchema.required).toEqual(["element_type"]);
  expect(client.catalog()[0].inputSchema.properties).not.toHaveProperty("context_handle");
});

it("exchanges source credentials only with issuer and injects a host context into MCP", async () => {
  const id = "mcp-owner"; credential(id); const requests: Array<{ url: string; init?: RequestInit }> = [];
  const client = fixture((async (url, init) => {
    requests.push({ url: String(url), init });
    if (String(url).includes("/token")) return token();
    if (init?.method === "POST") return Response.json({ context_handle: "x".repeat(43) });
    return new Response(null, { status: 204 });
  }) as McpHttpFetch);
  const result = await client.query(id, owner(id), "list_elements", { element_type: "pipe" }, "native-call");
  const exchange = new URLSearchParams(String(requests[0].init?.body));
  expect(exchange.get("subject_token")).toBe(jwt(id));
  expect(exchange.get("audience")).toBe("tjwater-mcp");
  expect(exchange.get("resource")).toBe(settings.resource);
  const bind = requests.find((request) => request.init?.method === "POST" && !request.url.includes("/token"))!;
  expect(new Headers(bind.init?.headers).get("Authorization")).toBe(`Bearer ${resourceToken}`);
  expect(JSON.parse(String(bind.init?.body))).toMatchObject({ purpose: "ordinary_chat", native_session_id: id });
  expect(result.structuredContent).toMatchObject({ arguments: { context_handle: "x".repeat(43) }, _meta: { "io.tjwater/callId": "native-call" } });
  expect(requests.at(-1)?.init?.method).toBe("DELETE");
});

it("coalesces concurrent exchange, isolates owners and discards cached tokens on release", async () => {
  const a = "mcp-a", b = "mcp-b"; credential(a); credential(b); let exchanges = 0;
  const client = fixture((async (url, init) => {
    if (String(url).includes("/token")) { exchanges++; await Bun.sleep(5); return token(); }
    if (init?.method === "POST") return Response.json({ context_handle: "x".repeat(43) });
    return new Response(null, { status: 204 });
  }) as McpHttpFetch);
  await Promise.all([client.query(a, owner(a), "list_elements", {}, "one"), client.query(a, owner(a), "list_elements", {}, "two"), client.query(b, owner(b), "list_elements", {}, "one")]);
  expect(exchanges).toBe(2);
  client.dropSession(a);
  await client.query(a, owner(a), "list_elements", {}, "three");
  expect(exchanges).toBe(3);
});

it("renews an expired source token through existing credential events", async () => {
  const id = "mcp-renew"; credential(id); let exchanges = 0, renewals = 0;
  const unsubscribe = credentialRefreshCoordinator.subscribe(id, (event) => {
    if (event.type === "credential_refresh_required") {
      renewals++; refreshSessionCredential(id, { authorization: `Bearer ${jwt("renewed")}`, projectId: "project" }, event.requestId);
    }
  });
  const client = fixture((async (url, init) => {
    if (String(url).includes("/token")) {
      exchanges++;
      const subject = new URLSearchParams(String(init?.body)).get("subject_token");
      return subject === jwt("renewed") ? token() : Response.json({ error: "invalid_grant" }, { status: 400 });
    }
    if (init?.method === "POST") return Response.json({ context_handle: "x".repeat(43) });
    return new Response(null, { status: 204 });
  }) as McpHttpFetch);
  try { await client.query(id, owner(id), "list_elements", {}, "call"); expect(exchanges).toBe(2); expect(renewals).toBe(1); }
  finally { unsubscribe(); }
});

it("does not refresh for scope rejection and renews resource authentication only once", async () => {
  const id = "mcp-rejected"; credential(id); let exchanges = 0, binds = 0;
  const client = fixture((async (url, init) => {
    if (String(url).includes("/token")) { exchanges++; return token(); }
    if (init?.method === "POST") { binds++; return Response.json({}, { status: 401 }); }
    return new Response(null, { status: 204 });
  }) as McpHttpFetch);
  await expect(client.query(id, owner(id), "list_elements", {}, "call")).rejects.toMatchObject({ status: 401 });
  expect(exchanges).toBe(2); expect(binds).toBe(2);
});

it("cancels a caller waiting for shared credential renewal", async () => {
  const id = "mcp-cancel"; credential(id);
  let resolve!: (response: Response) => void;
  let notifyStarted!: () => void;
  const started = new Promise<void>((done) => { notifyStarted = done; });
  const client = fixture((async () => new Promise<Response>((done) => { resolve = done; notifyStarted(); })) as McpHttpFetch);
  const controller = new AbortController();
  const pending = client.query(id, owner(id), "list_elements", {}, "call", controller.signal);
  await started;
  controller.abort(new Error("cancelled"));
  await expect(pending).rejects.toThrow("cancelled");
  client.dropSession(id); resolve(token());
});

it("sends resource on both grants and supports advertised Basic client authentication", async () => {
  const id = "mcp-basic"; credential(id);
  const grants: Array<{ headers: Headers; body: URLSearchParams }> = [];
  const client = fixture((async (url, init) => {
    if (String(url).endsWith("/token")) {
      grants.push({ headers: new Headers(init?.headers), body: new URLSearchParams(String(init?.body)) });
      return token();
    }
    if (init?.method === "POST") return Response.json({ context_handle: "x".repeat(43) });
    return new Response(null, { status: 204 });
  }) as McpHttpFetch, "client_secret_basic");
  await client.prepare();
  await client.query(id, owner(id), "list_elements", {}, "call");
  expect(grants.map((g) => g.body.get("grant_type"))).toEqual(["client_credentials", "urn:ietf:params:oauth:grant-type:token-exchange"]);
  for (const grant of grants) {
    expect(grant.body.get("resource")).toBe(settings.resource);
    expect(grant.body.has("client_secret")).toBeFalse();
    expect(grant.headers.get("Authorization")).toBe(`Basic ${Buffer.from(`${settings.clientId}:${settings.clientSecret}`).toString("base64")}`);
  }
});

it("reports an unsupported resource indicator as configuration failure without a weaker retry", async () => {
  const id = "mcp-invalid-resource"; credential(id);
  let attempts = 0;
  const client = fixture((async (_url, init) => {
    attempts++;
    expect(new URLSearchParams(String(init?.body)).get("resource")).toBe(settings.resource);
    return Response.json({ error: "invalid_target" }, { status: 400 });
  }) as McpHttpFetch);
  await expect(client.query(id, owner(id), "list_elements", {}, "call")).rejects.toThrow("RFC 8707");
  expect(attempts).toBe(1);
});

it("rejects an exchanged identity, issuer, client or audience that differs from the verified source", () => {
  assertExchangedIdentity(jwt("source"), resourceToken, settings);
  for (const overrides of [{ sub: "other" }, { aud: "api" }, { iss: "https://other/realm" }, { azp: "browser" }]) {
    expect(() => assertExchangedIdentity(jwt("source"), jwt("resource", overrides), settings)).toThrow();
  }
});