199 lines
11 KiB
TypeScript
199 lines
11 KiB
TypeScript
import { afterEach, expect, it, spyOn } from "bun:test";
|
|
import { McpQueryClient, type McpHttpFetch } from "../../src/product/mcp-query-client.js";
|
|
import { bindSessionCredential, releaseSessionCredential } from "../../src/product/credential-refresh.js";
|
|
import { PRODUCT_MCP_READ_TOOLS, PRODUCT_MCP_TOOLS } from "../../src/runtime/product-mcp-tools.js";
|
|
|
|
const settings = { issuer: "http://localhost/realms/retry", resource: "http://localhost/mcp", audience: "retry-mcp", clientId: "retry-host", clientSecret: "fixture-only" };
|
|
const sessions: string[] = [];
|
|
afterEach(() => { for (const id of sessions.splice(0)) releaseSessionCredential(id); });
|
|
const token = (overrides: Record<string, unknown> = {}) => `e30.${Buffer.from(JSON.stringify({ sub: "retry-user", iss: settings.issuer, aud: settings.audience, azp: settings.clientId, ...overrides })).toString("base64url")}.fixture`;
|
|
const failure = (code: string) => ({ content: [], isError: true, structuredContent: { error: { code } } });
|
|
const success = { content: [], structuredContent: { instructions: "fixture skill" } };
|
|
|
|
function fixture(outcomes: unknown[], options: { tool?: string; annotations?: "missing" | "destructive" | "misleading"; bindingStatuses?: number[]; tokenStatuses?: number[]; exchangedIdentity?: Record<string, unknown>; metadataIssuer?: string; discoveryOutcomes?: Array<number | Error> } = {}) {
|
|
const id = `retry-${sessions.length}-${crypto.randomUUID()}`;
|
|
sessions.push(id);
|
|
bindSessionCredential(id, { authorization: `Bearer ${token()}`, projectId: "project" });
|
|
const owner = { context: { sessionId: id, actorKey: "actor", clientSessionId: id, projectId: "project", projectKey: "project", traceId: "test", purpose: "task_execution" as const }, runtimeSessionId: id };
|
|
const calls: Array<{ name: string; arguments: Record<string, unknown>; _meta: unknown }> = [];
|
|
const revoked: string[] = [];
|
|
const signals: AbortSignal[] = [];
|
|
const trace = { binds: 0, exchanges: 0, discoveries: 0 };
|
|
const client = new McpQueryClient(settings, (async (url, init) => {
|
|
const path = String(url);
|
|
if (path === settings.resource && !new Headers(init?.headers).has("Authorization")) {
|
|
const outcome = options.discoveryOutcomes?.[trace.discoveries];
|
|
trace.discoveries++;
|
|
if (outcome instanceof Error) throw outcome;
|
|
return new Response(null, { status: outcome ?? 401 });
|
|
}
|
|
if (path.includes("oauth-protected-resource")) return Response.json({ resource: settings.resource, authorization_servers: [settings.issuer] });
|
|
if (path.includes(".well-known")) return Response.json({ issuer: options.metadataIssuer ?? settings.issuer, authorization_endpoint: `${settings.issuer}/authorize`, token_endpoint: `${settings.issuer}/token`, response_types_supported: ["code"], token_endpoint_auth_methods_supported: ["client_secret_post"] });
|
|
if (path.endsWith("/token")) {
|
|
const status = options.tokenStatuses?.[trace.exchanges] ?? 200;
|
|
trace.exchanges++;
|
|
return status === 200 ? Response.json({ access_token: token(options.exchangedIdentity), expires_in: 300, token_type: "Bearer" }) : Response.json({ error: "invalid_client" }, { status });
|
|
}
|
|
if (init?.method === "POST") {
|
|
const status = options.bindingStatuses?.[trace.binds] ?? 200;
|
|
const handle = `context${++trace.binds}`.padEnd(43, "x");
|
|
return Response.json(status === 200 ? { context_handle: handle } : {}, { status });
|
|
}
|
|
revoked.push(path.split("/").at(-1)!);
|
|
return new Response(null, { status: 204 });
|
|
}) as McpHttpFetch);
|
|
const tools = PRODUCT_MCP_TOOLS.map((name) => ({ name, inputSchema: { type: "object", properties: { context_handle: { type: "string" }, section: { type: "string" } } },
|
|
...(options.annotations === "missing" ? {} : { annotations: { readOnlyHint: options.annotations === "misleading" || (PRODUCT_MCP_READ_TOOLS as readonly string[]).includes(name), destructiveHint: options.annotations === "destructive" } }),
|
|
}));
|
|
Object.assign(client, { tools, withClient: async (_token: string, _project: string, signal: AbortSignal, run: (transport: unknown) => Promise<unknown>) => {
|
|
signals.push(signal);
|
|
return run({
|
|
listTools: async () => ({ tools }), callTool: async (params: typeof calls[number]) => {
|
|
const outcome = outcomes[Math.min(calls.length, outcomes.length - 1)];
|
|
calls.push(params);
|
|
if (outcome instanceof Error) throw outcome;
|
|
return outcome;
|
|
},
|
|
});
|
|
} });
|
|
return { client, calls, revoked, trace, tools, signals, query: (signal?: AbortSignal) => client.query(id, owner, options.tool ?? "get_task_context", { section: "skill" }, "same-native-call", signal) };
|
|
}
|
|
|
|
it("recovers a transient skill read with a new context and the same native call ID", async () => {
|
|
const test = fixture([failure("SOURCE_UNAVAILABLE"), success]);
|
|
expect(await test.query()).toEqual(success);
|
|
expect(test.trace.binds).toBe(2);
|
|
expect(test.calls.map((call) => call._meta)).toEqual(Array(2).fill({ "io.tjwater/callId": "same-native-call" }));
|
|
expect(new Set(test.calls.map((call) => call.arguments.context_handle)).size).toBe(2);
|
|
expect(test.revoked).toEqual(test.calls.map((call) => call.arguments.context_handle as string));
|
|
});
|
|
|
|
it("recovers a fetch transport error before exposing it to the model", async () => {
|
|
const test = fixture([new TypeError("fetch failed"), success]);
|
|
expect(await test.query()).toEqual(success);
|
|
expect(test.calls).toHaveLength(2);
|
|
});
|
|
|
|
it("returns the final structured failure after three total attempts", async () => {
|
|
const result = failure("QUERY_TIMEOUT");
|
|
const test = fixture([result]);
|
|
expect(await test.query()).toEqual(result);
|
|
expect(test.calls).toHaveLength(3);
|
|
expect(test.revoked).toHaveLength(3);
|
|
});
|
|
|
|
it.each(["run_burst_detection", "clean_scada_readings", "list_calculation_jobs"])("does not replay %s on a transient failure", async (tool) => {
|
|
const result = failure("SOURCE_UNAVAILABLE");
|
|
const test = fixture([result, success], { tool, annotations: "misleading" });
|
|
expect(await test.query()).toEqual(result);
|
|
expect(test.calls).toHaveLength(1);
|
|
});
|
|
|
|
it.each(["missing", "destructive"] as const)("does not trust %s read annotations", async (annotations) => {
|
|
const result = failure("SOURCE_UNAVAILABLE");
|
|
const test = fixture([result, success], { annotations });
|
|
expect(await test.query()).toEqual(result);
|
|
expect(test.calls).toHaveLength(1);
|
|
});
|
|
|
|
it.each([429, 502, 503, 504])("recovers HTTP %s while binding a read context", async (status) => {
|
|
const test = fixture([success], { bindingStatuses: [status, 200] });
|
|
expect(await test.query()).toEqual(success);
|
|
expect(test.trace.binds).toBe(2);
|
|
expect(test.calls).toHaveLength(1);
|
|
expect(test.trace.exchanges).toBe(1);
|
|
});
|
|
|
|
it.each(["QUERY_TIMEOUT", "RATE_LIMIT", "RATE_LIMITED"])("recovers a structured %s response", async (code) => {
|
|
const test = fixture([failure(code), success]);
|
|
expect(await test.query()).toEqual(success);
|
|
expect(test.calls).toHaveLength(2);
|
|
});
|
|
|
|
it.each([403, 404, 422])("does not refresh or retry HTTP %s authorization and query failures", async (status) => {
|
|
const test = fixture([success], { bindingStatuses: [status, 200] });
|
|
await expect(test.query()).rejects.toMatchObject({ status });
|
|
expect(test.trace.binds).toBe(1);
|
|
expect(test.trace.exchanges).toBe(1);
|
|
});
|
|
|
|
it.each(["QUERY_CONTEXT_DENIED", "INVALID_QUERY", "NOT_FOUND", "MCP_STORAGE_UNAVAILABLE"])("does not retry structured %s failures", async (code) => {
|
|
const result = failure(code);
|
|
const test = fixture([result, success]);
|
|
expect(await test.query()).toEqual(result);
|
|
expect(test.calls).toHaveLength(1);
|
|
expect(test.trace.exchanges).toBe(1);
|
|
});
|
|
|
|
it.each([400, 401, 403])("does not mistake OAuth configuration HTTP %s for transient 503", async (status) => {
|
|
const test = fixture([success], { tokenStatuses: [status, 200] });
|
|
await expect(test.query()).rejects.toMatchObject({ status: 503, retryable: false });
|
|
expect(test.trace.exchanges).toBe(1);
|
|
expect(test.trace.binds).toBe(0);
|
|
});
|
|
|
|
it("does not retry an incompatible issuer or exchanged identity", async () => {
|
|
const issuer = fixture([success], { metadataIssuer: "http://localhost/realms/other" });
|
|
await expect(issuer.query()).rejects.toMatchObject({ status: 503, retryable: false });
|
|
expect(issuer.trace.exchanges).toBe(0);
|
|
const identity = fixture([success], { exchangedIdentity: { sub: "another-user" } });
|
|
await expect(identity.query()).rejects.toMatchObject({ status: 503, retryable: false });
|
|
expect(identity.trace.exchanges).toBe(1);
|
|
expect(identity.trace.binds).toBe(0);
|
|
});
|
|
|
|
it.each([503, new TypeError("fetch failed")])("recovers transient discovery failure before exchanging credentials", async (outage) => {
|
|
const test = fixture([success], { discoveryOutcomes: [outage, 401] });
|
|
expect(await test.query()).toEqual(success);
|
|
expect(test.trace.discoveries).toBe(2);
|
|
expect(test.trace.exchanges).toBe(1);
|
|
});
|
|
|
|
it("preserves an explicit 403 even if a contradictory structured code requests authentication", async () => {
|
|
const denied = { content: [], isError: true, structuredContent: { error: { code: "AUTH_REQUIRED", status: 403 } } };
|
|
const test = fixture([denied, success]);
|
|
expect(await test.query()).toEqual(denied);
|
|
expect(test.calls).toHaveLength(1);
|
|
expect(test.trace.exchanges).toBe(1);
|
|
});
|
|
|
|
it("refreshes structured authentication only once, sharing the three-attempt ceiling", async () => {
|
|
const test = fixture([failure("AUTH_REQUIRED"), failure("SOURCE_UNAVAILABLE"), success]);
|
|
expect(await test.query()).toEqual(success);
|
|
expect(test.calls).toHaveLength(3);
|
|
expect(test.trace.exchanges).toBe(2);
|
|
const denied = failure("AUTH_REQUIRED");
|
|
const rejected = fixture([denied]);
|
|
expect(await rejected.query()).toEqual(denied);
|
|
expect(rejected.calls).toHaveLength(2);
|
|
expect(rejected.trace.exchanges).toBe(2);
|
|
});
|
|
|
|
it("does not restart OAuth when cancelled during read backoff", async () => {
|
|
const test = fixture([failure("SOURCE_UNAVAILABLE"), success]);
|
|
const controller = new AbortController();
|
|
const pending = test.query(controller.signal);
|
|
while (!test.revoked.length) await Bun.sleep(1);
|
|
await Bun.sleep(5);
|
|
controller.abort(new Error("cancelled during backoff"));
|
|
await expect(pending).rejects.toThrow("cancelled during backoff");
|
|
expect(test.calls).toHaveLength(1);
|
|
expect(test.trace.binds).toBe(1);
|
|
expect(test.trace.exchanges).toBe(1);
|
|
});
|
|
|
|
it("keeps one request deadline through all attempts and backoff", async () => {
|
|
const original = AbortSignal.timeout.bind(AbortSignal);
|
|
const timeout = spyOn(AbortSignal, "timeout").mockImplementation((ms) => original(ms === 61_000 ? 250 : ms));
|
|
try {
|
|
const test = fixture([failure("QUERY_TIMEOUT")]);
|
|
for (const tool of test.tools) Object.assign(tool, { _meta: { "io.tjwater/request-timeout": { defaultSeconds: 1, maxSeconds: 1 } } });
|
|
await expect(test.query()).rejects.toMatchObject({ name: "TimeoutError" });
|
|
expect(test.calls.length).toBeGreaterThan(0);
|
|
expect(test.calls.length).toBeLessThan(3);
|
|
expect(new Set(test.signals).size).toBe(1);
|
|
expect(test.signals[0].aborted).toBeTrue();
|
|
expect(timeout.mock.calls.filter(([ms]) => ms === 61_000)).toHaveLength(1);
|
|
} finally { timeout.mockRestore(); }
|
|
});
|