src/mcp_oauth/registry.test.ts
Rama main · réplica de lectura del repositorio original.
import { describe, expect, it, vi } from "vitest";
import {
createFakeClock,
createSequentialIdSource,
} from "@/state_machines/testing";
import type { IdSource } from "@/state_machines/clock";
import {
createMcpOAuthRegistry,
type McpOAuthListenerRequest,
} from "./registry";
async function flush(): Promise<void> {
for (let index = 0; index < 10; index += 1) await Promise.resolve();
}
function createHarness(ids: IdSource = createSequentialIdSource()) {
const clock = createFakeClock();
const bindings: McpOAuthListenerRequest[] = [];
const closed: string[] = [];
const settled = vi.fn();
const registry = createMcpOAuthRegistry({
clock,
ids,
bindListener(request) {
bindings.push(request);
return {
settled: Promise.resolve({
boundHosts: ["127.0.0.1"],
anyInUse: false,
}),
async close() {
closed.push(request.invocationRef.operationId);
},
};
},
onSettled: settled,
observer: undefined,
});
return { registry, clock, bindings, closed, settled };
}
function request(
serverId: number,
messageId = `message-${serverId}`,
authorize: (code?: string) => Promise<"AUTHORIZED" | "REDIRECT"> = vi.fn(
async () => "REDIRECT" as const,
),
port = 53682,
) {
return {
port,
serverId,
rendererMessageId: messageId,
expectedState: `state-${serverId}`,
authorize,
onAbort: vi.fn(),
};
}
describe("MCP OAuth registry", () => {
it("keeps the existing happy path with typed listener/callback correlation", async () => {
const { registry, bindings, settled } = createHarness();
const authorize = vi.fn(async (code?: string) =>
code === undefined ? ("REDIRECT" as const) : ("AUTHORIZED" as const),
);
const connectRequest = request(7, "message-7", authorize);
const result = registry.connect(connectRequest);
await flush();
expect(bindings).toHaveLength(1);
const listener = bindings[0];
expect(listener.invocationRef).toEqual({
kind: "mcp-oauth",
entityKey: 7,
operationId: "mcp-oauth:1",
});
expect(
listener.onCallback({
invocationRef: listener.invocationRef,
state: "state-7",
code: "authorization-code",
}),
).toMatchObject({ claimed: true });
await flush();
await expect(result).resolves.toEqual({ success: true, error: null });
expect(connectRequest.authorize).toHaveBeenNthCalledWith(1);
expect(connectRequest.authorize).toHaveBeenNthCalledWith(
2,
"authorization-code",
);
expect(settled).toHaveBeenCalledWith(7, listener.invocationRef, {
success: true,
error: null,
});
});
it("deduplicates renderer retries by message ID without conflating the invocation ref", async () => {
const { registry, bindings } = createHarness();
const original = request(4, "stable-renderer-message");
const retry = request(4, "stable-renderer-message");
const first = registry.connect(original);
const duplicate = registry.connect(retry);
expect(duplicate).toBe(first);
await flush();
expect(bindings).toHaveLength(1);
expect(retry.onAbort).not.toHaveBeenCalled();
await registry.dispose();
await expect(first).resolves.toMatchObject({ success: false });
});
it("never evicts unresolved renderer receipts", async () => {
const { registry, bindings } = createHarness();
const requests = Array.from({ length: 128 }, (_, index) =>
request(index + 1, `pending-${index + 1}`, undefined, 40_000 + index),
);
const pending = requests.map((entry) => registry.connect(entry));
await flush();
expect(bindings).toHaveLength(128);
const overflowRequest = request(999, "overflow", undefined, 50_000);
await expect(registry.connect(overflowRequest)).resolves.toMatchObject({
success: false,
error: expect.stringMatching(/too many/i),
});
expect(overflowRequest.onAbort).toHaveBeenCalledOnce();
expect(registry.connect(requests[0])).toBe(pending[0]);
expect(bindings).toHaveLength(128);
await registry.dispose();
await Promise.all(pending);
});
it("keys resources by the complete typed ref when operation IDs collide", async () => {
const { registry, bindings } = createHarness({
next: () => "collision",
});
const first = registry.connect(request(1, "first", undefined, 45_001));
const second = registry.connect(request(2, "second", undefined, 45_002));
await flush();
expect(bindings).toHaveLength(2);
expect(bindings.map((binding) => binding.invocationRef.entityKey)).toEqual([
1, 2,
]);
await registry.dispose();
await expect(first).resolves.toMatchObject({ success: false });
await expect(second).resolves.toMatchObject({ success: false });
});
it("preserves last-request-wins on a shared callback port", async () => {
const { registry, bindings } = createHarness();
const firstRequest = request(1, "first");
const secondRequest = request(2, "second");
const first = registry.connect(firstRequest);
await flush();
const second = registry.connect(secondRequest);
await flush();
await expect(first).resolves.toMatchObject({
success: false,
error: expect.stringMatching(/superseded/i),
});
expect(firstRequest.onAbort).toHaveBeenCalledOnce();
expect(bindings).toHaveLength(2);
expect(bindings[1].invocationRef.entityKey).toBe(2);
await registry.dispose();
await expect(second).resolves.toMatchObject({ success: false });
});
it("cancels and settles a server while waiting for its callback", async () => {
const { registry, closed } = createHarness();
const connectRequest = request(11);
const result = registry.connect(connectRequest);
await flush();
await registry.cancelServer(11, "server deleted");
await expect(result).resolves.toEqual({
success: false,
error: "server deleted",
});
expect(connectRequest.onAbort).toHaveBeenCalledOnce();
expect(closed).toEqual(["mcp-oauth:1"]);
});
it("keeps cancellation authoritative while token exchange is in flight", async () => {
const { registry, bindings, settled } = createHarness();
let finishExchange!: (value: "AUTHORIZED") => void;
const exchange = new Promise<"AUTHORIZED">((resolve) => {
finishExchange = resolve;
});
const authorize = vi.fn((code?: string) =>
code === undefined ? Promise.resolve("REDIRECT" as const) : exchange,
);
const connectRequest = request(12, "exchange", authorize);
const result = registry.connect(connectRequest);
await flush();
const listener = bindings[0];
listener.onCallback({
invocationRef: listener.invocationRef,
state: "state-12",
code: "code",
});
await flush();
expect(registry.getState(53682).status).toBe("exchanging");
await registry.cancelServer(12, "OAuth disabled");
finishExchange("AUTHORIZED");
await flush();
await expect(result).resolves.toEqual({
success: false,
error: "OAuth disabled",
});
expect(settled).toHaveBeenCalledTimes(1);
});
it("rejects a stale typed callback without consuming the active flow", async () => {
const { registry, bindings } = createHarness();
const result = registry.connect(request(3));
await flush();
const active = bindings[0];
const staleRef = { ...active.invocationRef, operationId: "stale" };
expect(
active.onCallback({
invocationRef: staleRef,
state: "state-3",
code: "code",
}),
).toEqual({ claimed: false, reason: "inactive" });
expect(registry.getState(53682).status).toBe("awaitingCallback");
await registry.dispose();
await expect(result).resolves.toMatchObject({ success: false });
});
it("settles waiters and closes listeners during explicit disposal", async () => {
const { registry, clock, closed } = createHarness();
const result = registry.connect(request(9));
await flush();
expect(clock.pendingTimerCount()).toBe(1);
await registry.dispose();
expect(clock.pendingTimerCount()).toBe(0);
expect(closed).toEqual(["mcp-oauth:1"]);
await expect(result).resolves.toEqual({
success: false,
error: "OAuth flow registry disposed.",
});
});
it("awaits an active authorization effect during explicit disposal", async () => {
const { registry } = createHarness();
let finishAuthorization!: (value: "REDIRECT") => void;
const authorization = new Promise<"REDIRECT">((resolve) => {
finishAuthorization = resolve;
});
const result = registry.connect(
request(10, "deferred", () => authorization),
);
await flush();
let disposed = false;
const disposal = registry.dispose().then(() => {
disposed = true;
});
await flush();
expect(disposed).toBe(false);
finishAuthorization("REDIRECT");
await disposal;
await expect(result).resolves.toMatchObject({ success: false });
});
});