430 lines · 15.1 KB · TypeScript
RawDownload
import { describe, expect, it } from "@effect/vitest";import { Effect } from "effect"; import type { ExecutionEngine } from "@executor-js/execution"; import {  makeInMemoryMcpSessionStore,  McpEngineBuildError,  type McpBuildServer,  type McpBuildServerOptions,} from "./in-memory-session-store";import { defaultMcpResource, type Principal } from "./seams";import { createExecutorMcpServer } from "./tool-server"; const TEST_PRINCIPAL: Principal = {  accountId: "acct_test",  organizationId: "org_test",  organizationName: "Test Org",  email: "test@example.com",  name: "Test",  avatarUrl: null,  roles: ["user"],}; it("preserves native elicitation mode when creating an in-memory MCP session", async () => {  let buildOptions: McpBuildServerOptions | undefined;  const sessions = makeInMemoryMcpSessionStore((_principal, options) => {    buildOptions = options;    return Effect.fail(new McpEngineBuildError({ cause: "stop after capturing options" }));  });   const result = await Effect.runPromise(    sessions.store.dispatch({      request: new Request("https://executor.test/mcp?elicitation_mode=native", {        method: "POST",        headers: { "content-type": "application/json" },        body: JSON.stringify({          jsonrpc: "2.0",          id: 1,          method: "initialize",          params: {            protocolVersion: "2025-06-18",            capabilities: { elicitation: { form: {} } },            clientInfo: { name: "test-client", version: "1.0.0" },          },        }),      }),      principal: TEST_PRINCIPAL,      resource: defaultMcpResource,      sessionId: null,      method: "POST",    }),  );   expect(result).toBeInstanceOf(Response);  expect((result as Response).status).toBe(500);  expect(buildOptions?.elicitationMode).toEqual({ mode: "native" });}); /** A do-nothing engine: the eviction test drives session lifetime, not tools. */const makeIdleTestEngine = (): ExecutionEngine => ({  execute: () => Effect.succeed({ result: "unused" }),  executeWithPause: () => Effect.succeed({ status: "completed", result: { result: "unused" } }),  resume: () => Effect.succeed(null),  getPausedExecution: () => Effect.succeed(null),  pausedExecutionCount: () => Effect.succeed(0),  hasPausedExecutions: () => Effect.succeed(false),  getDescription: Effect.succeed("idle-eviction test executor"),  shutdown: Effect.void,}); /** * An engine whose `execute` parks until the test releases it, so a request can * be held inside `transport.handleRequest` while the sweep runs. `shutdowns` * counts `engine.shutdown` runs — the disposal step that ends the detached * sandbox fibers, and which dropping the engine reference does not do. */const makeLatchedTestEngine = (): {  readonly engine: ExecutionEngine;  readonly started: Promise<void>;  readonly release: () => void;  readonly shutdowns: () => number;} => {  let signalStarted: () => void = () => {};  const started = new Promise<void>((resolve) => {    signalStarted = resolve;  });  let openGate: () => void = () => {};  const gate = new Promise<void>((resolve) => {    openGate = resolve;  });  let shutdowns = 0;  const park = <A>(value: A): Effect.Effect<A> =>    Effect.promise(async () => {      signalStarted();      await gate;      return value;    });  const engine: ExecutionEngine = {    ...makeIdleTestEngine(),    execute: () => park({ result: "released" }),    executeWithPause: () => park({ status: "completed", result: { result: "released" } }),    shutdown: Effect.sync(() => {      shutdowns += 1;    }),  };  return { engine, started, release: () => openGate(), shutdowns: () => shutdowns };}; // A long TTL keeps the sweep's own timer out of the way; the assertions drive// `sweepIdleSessions` directly with an explicit instant instead of sleeping// through a real window, so the test is deterministic rather than timing-raced.const IDLE_TTL_MS = 60_000; type TestSessionStore = ReturnType<typeof makeInMemoryMcpSessionStore>; /** Open a session on `sessions` and return its minted id. */const openSession = async (sessions: TestSessionStore): Promise<string> => {  const response = (await Effect.runPromise(    sessions.store.dispatch({      request: new Request("https://executor.test/mcp", {        method: "POST",        headers: {          "content-type": "application/json",          accept: "application/json, text/event-stream",        },        body: JSON.stringify({          jsonrpc: "2.0",          id: 1,          method: "initialize",          params: {            protocolVersion: "2025-06-18",            capabilities: {},            clientInfo: { name: "idle-test", version: "1.0.0" },          },        }),      }),      principal: TEST_PRINCIPAL,      resource: defaultMcpResource,      sessionId: null,      method: "POST",    }),  )) as Response;  expect(response.status).toBe(200);  const sessionId = response.headers.get("mcp-session-id") ?? "";  expect(sessionId).not.toBe("");  return sessionId;}; it("evicts a session that goes idle past the TTL and keeps a busy one", async () => {  const engine = makeIdleTestEngine();  const sessions = makeInMemoryMcpSessionStore(    () =>      createExecutorMcpServer({ engine }).pipe(Effect.map((mcpServer) => ({ mcpServer, engine }))),    { sessionIdleTtlMs: IDLE_TTL_MS },  );   const open = (): Promise<string> => openSession(sessions);   const call = (sessionId: string, id: number) =>    Effect.runPromise(      sessions.store.dispatch({        request: new Request("https://executor.test/mcp", {          method: "POST",          headers: {            "content-type": "application/json",            accept: "application/json, text/event-stream",            "mcp-session-id": sessionId,          },          body: JSON.stringify({ jsonrpc: "2.0", id, method: "tools/list" }),        }),        principal: TEST_PRINCIPAL,        resource: defaultMcpResource,        sessionId,        method: "POST",      }),    );   // oxlint-disable-next-line executor/no-try-catch-or-throw -- test boundary: always close the store  try {    const idle = await open();    const busy = await open();    expect(sessions.sessionCount()).toBe(2);     // Neither is stale yet, so a sweep at the current instant takes nothing.    expect(await sessions.sweepIdleSessions()).toBe(0);    expect(sessions.sessionCount()).toBe(2);     // Let the wall clock advance so the two sessions' stamps are separable,    // then keep working on one of them: `forward` restamps that one and only    // that one.    await new Promise((resolve) => setTimeout(resolve, 25));    const restampedAt = Date.now();    await call(busy, 2);     // Sweep one TTL after the restamp, less a millisecond: `busy` was stamped    // at or after `restampedAt` so it cannot have aged a full TTL, while `idle`    // was stamped at least 25ms earlier and must have. Exactly one goes.    expect(await sessions.sweepIdleSessions(restampedAt + IDLE_TTL_MS - 1)).toBe(1);    expect(sessions.sessionCount()).toBe(1);     // The evicted id is gone; the store reports it the way the envelope 404s.    expect(await call(idle, 3)).toBe("not-found");    // The busy one still serves.    expect(await call(busy, 4)).toBeInstanceOf(Response);  } finally {    await sessions.close();  }}); it("never evicts a session while one of its requests is still in flight", async () => {  const latched = makeLatchedTestEngine();  const sessions = makeInMemoryMcpSessionStore(    () =>      createExecutorMcpServer({ engine: latched.engine }).pipe(        Effect.map((mcpServer) => ({ mcpServer, engine: latched.engine })),      ),    { sessionIdleTtlMs: IDLE_TTL_MS },  );   const callExecute = (sessionId: string): Promise<unknown> =>    Effect.runPromise(      sessions.store.dispatch({        request: new Request("https://executor.test/mcp", {          method: "POST",          headers: {            "content-type": "application/json",            accept: "application/json, text/event-stream",            "mcp-session-id": sessionId,          },          body: JSON.stringify({            jsonrpc: "2.0",            id: 2,            method: "tools/call",            params: { name: "execute", arguments: { code: "return 1" } },          }),        }),        principal: TEST_PRINCIPAL,        resource: defaultMcpResource,        sessionId,        method: "POST",      }),    );   // oxlint-disable-next-line executor/no-try-catch-or-throw -- test boundary: always release the latch and close the store  try {    const sessionId = await openSession(sessions);     // Start a call and park it inside the engine. `forward` stamps last-seen    // BEFORE it awaits the transport, so from here on the stamp only ages — a    // request slower than the TTL is indistinguishable from an abandoned    // session unless the store also counts what is in flight.    const startedAt = Date.now();    const inFlight = callExecute(sessionId);    await latched.started;     // Sweep a full TTL past the moment the call began. Without the in-flight    // counter this evicts the session and closes the transport, the server, and    // the engine underneath the request that is still using them.    expect(await sessions.sweepIdleSessions(startedAt + IDLE_TTL_MS)).toBe(0);    expect(sessions.sessionCount()).toBe(1);    expect(latched.shutdowns()).toBe(0);     // The parked request still completes, on the transport it started on.    latched.release();    const response = await inFlight;    expect(response).toBeInstanceOf(Response);    expect((response as Response).status).toBe(200);     // And the reprieve is only for the duration of the call: the session is    // restamped as it ends, so the next idle window still reclaims it — engine    // shutdown included, which is what ends the detached sandbox fibers.    expect(await sessions.sweepIdleSessions(Date.now() + IDLE_TTL_MS)).toBe(1);    expect(sessions.sessionCount()).toBe(0);    expect(latched.shutdowns()).toBe(1);  } finally {    latched.release();    await sessions.close();  }}); // ---------------------------------------------------------------------------// The pre-initialize guard, through the real store path.//// `store.dispatch` with no session id runs the guard and, when the guard// declines, builds a real MCP server and drives a real streamable-HTTP// transport. So these assert BOTH halves of the contract: the one answer the// guard replaces, and the transport answers it must not shadow.// --------------------------------------------------------------------------- /** The headers a streamable-HTTP client must send on a POST; less is a 406/415. */const MCP_POST_HEADERS = {  "content-type": "application/json",  accept: "application/json, text/event-stream",} as const; /** No code ever runs here: these requests are answered before any tool call. */const stubEngine: ExecutionEngine<never> = {  execute: () => Effect.succeed({ result: "unused" }),  executeWithPause: () => Effect.succeed({ status: "completed", result: { result: "unused" } }),  resume: () => Effect.succeed(null),  getPausedExecution: () => Effect.succeed(null),  pausedExecutionCount: () => Effect.succeed(0),  hasPausedExecutions: () => Effect.succeed(false),  getDescription: Effect.succeed("test executor"),  shutdown: Effect.void,}; /** A store whose sessions are real: a real MCP server on a real transport. */const makeServingStore = () => {  let builds = 0;  const buildServer: McpBuildServer = () =>    Effect.sync(() => {      builds += 1;    }).pipe(      Effect.flatMap(() => createExecutorMcpServer({ engine: stubEngine })),      Effect.map((mcpServer) => ({ mcpServer, engine: stubEngine })),    );  return { sessions: makeInMemoryMcpSessionStore(buildServer), buildCount: (): number => builds };}; const dispatchPost = (  sessions: ReturnType<typeof makeServingStore>["sessions"],  body: unknown,  headers: Record<string, string> = MCP_POST_HEADERS,): Promise<Response> =>  Effect.runPromise(    sessions.store      .dispatch({        request: new Request("https://executor.test/mcp", {          method: "POST",          headers,          body: JSON.stringify(body),        }),        principal: TEST_PRINCIPAL,        resource: defaultMcpResource,        sessionId: null,        method: "POST",      })      .pipe(        Effect.map((result) => {          expect(result).toBeInstanceOf(Response);          return result as Response;        }),      ),  ); interface JsonRpcErrorBody {  readonly error: { readonly code: number; readonly message: string };} describe("pre-initialize dispatch through the in-memory session store", () => {  it("answers a valid unknown pre-session method with -32601 on a 200", async () => {    const { sessions, buildCount } = makeServingStore();    const response = await dispatchPost(sessions, {      jsonrpc: "2.0",      id: 7,      method: "server/discover",      params: {},    });     // 200, not the transport's 400: a per-request error the client survives.    expect(response.status).toBe(200);    expect(await response.json()).toEqual({      jsonrpc: "2.0",      id: 7,      error: { code: -32601, message: "Method not found" },    });    // The guard short-circuits before any engine is built.    expect(buildCount()).toBe(0);    await sessions.close();  });   it("passes a pre-session notification to the transport", async () => {    const { sessions, buildCount } = makeServingStore();    const response = await dispatchPost(sessions, {      jsonrpc: "2.0",      method: "notifications/initialized",    });     // A notification carries no id, so the guard may not answer it at all;    // whatever comes back is the transport's own answer.    const body = (await response.json()) as JsonRpcErrorBody;    expect(body.error.code).not.toBe(-32601);    expect(body.error.code).toBe(-32000);    expect(buildCount()).toBe(1);    await sessions.close();  });   it("leaves a structurally invalid request to the transport's parse error", async () => {    const { sessions } = makeServingStore();    // A fractional id is not a JSON-RPC id, so this is not a request the guard    // may report an unknown method for.    const response = await dispatchPost(sessions, {      jsonrpc: "2.0",      id: 1.5,      method: "server/discover",    });     expect(response.status).toBe(400);    const body = (await response.json()) as JsonRpcErrorBody;    expect(body.error.code).toBe(-32700);    expect(body.error.code).not.toBe(-32601);    await sessions.close();  });   it("leaves a wrong Content-Type to the transport's 415", async () => {    const { sessions } = makeServingStore();    const response = await dispatchPost(      sessions,      { jsonrpc: "2.0", id: 1, method: "server/discover" },      { "content-type": "text/plain", accept: MCP_POST_HEADERS.accept },    );     expect(response.status).toBe(415);    await sessions.close();  });   it("leaves an incomplete Accept to the transport's 406", async () => {    const { sessions } = makeServingStore();    const response = await dispatchPost(      sessions,      { jsonrpc: "2.0", id: 1, method: "server/discover" },      { "content-type": "application/json", accept: "application/json" },    );     expect(response.status).toBe(406);    await sessions.close();  });});