reverse-engineering-anything/tests/boundary/mcp/sessionAvailabilityRecovery.test.ts

import { Client, InMemoryTransport } from "@modelcontextprotocol/client";
import { expect, it } from "vitest";

import { toolContract } from "../../../src/contracts/toolContracts.js";
import { expandToolAvailability } from "../../../src/contracts/toolOutputSchemaPrimitives.js";
import { ProviderAdapterError } from "../../../src/domain/providerAdapterError.js";
import { err, ok } from "../../../src/domain/result.js";
import { silentLogger } from "../../../src/logger.js";
import { createServer } from "../../../src/server/createServer.js";
import { observed } from "../../fixtures/analysisExecution.js";
import {
  createBinarySessionTargets,
  createCacheProvider,
  createDeferred,
  createTestBinarySession,
} from "../../fixtures/binarySession.js";

it("restores MCP tool availability after cancelling a target switch", async () => {
  const [first, second] = await createBinarySessionTargets();
  const closing = createDeferred<void>();
  const release = createDeferred<void>();
  const cancellationObserved = createDeferred<void>();
  const provider = createCacheProvider([]);
  const resolveProfile = provider.resolveAnalysisProfile;
  if (resolveProfile === undefined)
    throw new Error("Missing test analysis profile");
  const clients: { calls: number; closed: number }[] = [];
  provider.createClient = () => {
    const state = { calls: 0, closed: 0 };
    const firstClient = clients.length === 0;
    clients.push(state);
    return {
      execute: (operation) => {
        state.calls += 1;
        return Promise.resolve(
          firstClient
            ? err(new ProviderAdapterError("fixture", operation))
            : observed("restored-name"),
        );
      },
      close: async () => {
        state.closed += 1;
        if (firstClient) {
          closing.resolve();
          await release.promise;
        }
        return ok(null);
      },
    };
  };
  const session = createTestBinarySession(provider, {
    resolveAnalysisProfile: (target, options) => {
      if (target.path === second) {
        options?.signal?.addEventListener(
          "abort",
          () => cancellationObserved.resolve(),
          { once: true },
        );
      }
      return resolveProfile(target, options);
    },
  });
  const server = createServer(
    { kind: "session", session },
    { environment: {}, logger: silentLogger },
  );
  const client = new Client({ name: "session-recovery", version: "1" });
  const [clientTransport, serverTransport] =
    InMemoryTransport.createLinkedPair();
  const controller = new AbortController();
  const readStatus = async () => {
    const reply = await client.callTool({
      name: "binary_session",
      arguments: {},
    });
    return toolContract("binary_session").outputSchema.parse(
      reply.structuredContent,
    ).result;
  };
  try {
    await server.connect(serverTransport);
    await client.connect(clientTransport);
    const opened = await client.callTool({
      name: "open_binary",
      arguments: { path: first },
    });
    expect(opened.isError).not.toBe(true);
    const query = { address: "0x1000", document: "fixture" };
    expect(
      (await client.callTool({ name: "address_name", arguments: query }))
        .isError,
    ).toBe(true);
    const before = await readStatus();
    expect(
      expandToolAvailability(before.tool_availability).find(
        (tool) => tool.name === "address_name",
      ),
    ).toMatchObject({
      available: false,
      reason: "provider_unavailable",
    });

    const switching = client.callTool(
      { name: "open_binary", arguments: { path: second } },
      { signal: controller.signal },
    );
    const cancelled = expect(switching).rejects.toThrow();
    await closing.promise;
    controller.abort();
    await cancellationObserved.promise;
    release.resolve();
    await cancelled;

    // Cancellation ends the client's wait before the server restores its session.
    const settled = await session.withAdmittedAnalysis(
      "binary_session",
      undefined,
      async () => null,
    );
    if (!settled.ok) throw settled.error;
    const restored = await readStatus();
    expect(restored).toMatchObject({ open: true, path: first });
    expect(restored.analysis_run?.run_id).not.toBe(before.analysis_run?.run_id);
    expect(clients).toEqual([{ calls: 1, closed: 1 }]);
    expect(
      expandToolAvailability(restored.tool_availability).find(
        (tool) => tool.name === "address_name",
      ),
    ).toMatchObject({
      available: true,
      reason: "available",
      remediation: null,
    });

    const retried = await client.callTool({
      name: "address_name",
      arguments: query,
    });
    expect(retried.isError).not.toBe(true);
    expect(retried.structuredContent).toMatchObject({
      normalized_result: "restored-name",
    });
    expect(clients).toEqual([
      { calls: 1, closed: 1 },
      { calls: 1, closed: 0 },
    ]);
  } finally {
    controller.abort();
    release.resolve();
    await client.close();
    await server.close();
    await session.close();
  }
});