109 lines
7.6 KiB
TypeScript
109 lines
7.6 KiB
TypeScript
import { Client } from "@modelcontextprotocol/sdk/client/index.js";
|
|
import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
|
|
import type { Tool } from "@modelcontextprotocol/sdk/types.js";
|
|
import { HarnessError } from "./errors.js";
|
|
|
|
export interface GatewayConnection {
|
|
list(signal: AbortSignal): Promise<Tool[]>;
|
|
call(name: string, args: Record<string, unknown>, signal: AbortSignal): Promise<unknown>;
|
|
close(): Promise<void>;
|
|
}
|
|
export type GatewayConnector = (signal: AbortSignal) => Promise<GatewayConnection>;
|
|
export interface GatewayConfig { url: URL; token: string | undefined; timeoutMs: number }
|
|
export const MAX_GATEWAY_BYTES = 2 * 1024 * 1024;
|
|
export function gatewayConfig(env: NodeJS.ProcessEnv = process.env): GatewayConfig {
|
|
let url: URL;
|
|
try { url = new URL(env.SWADS_MCP_URL || "https://ads.mincode.cn/mcp"); } catch { throw new HarnessError("SWADS_INVALID_CONFIG", "SWADS_MCP_URL must be a valid HTTP(S) URL"); }
|
|
if (!["https:", "http:"].includes(url.protocol) || url.username || url.password || url.search || url.hash) throw new HarnessError("SWADS_INVALID_CONFIG", "SWADS_MCP_URL must not contain credentials, query or fragment");
|
|
const timeoutMs = Number(env.SWADS_MCP_TIMEOUT_MS ?? 30_000);
|
|
if (!Number.isInteger(timeoutMs) || timeoutMs < 100 || timeoutMs > 120_000) throw new HarnessError("SWADS_INVALID_CONFIG", "SWADS_MCP_TIMEOUT_MS must be between 100 and 120000");
|
|
return { url, token: env.SWADS_MCP_TOKEN || undefined, timeoutMs };
|
|
}
|
|
|
|
// Read the wire body with a byte bound before handing it to the official SDK.
|
|
// A connection belongs to one command, so cancellation never affects another run.
|
|
export function mcpConnector(config: GatewayConfig, fetcher: typeof fetch = fetch): GatewayConnector {
|
|
return async (signal) => {
|
|
const boundedFetch: typeof fetch = async (input, init) => {
|
|
const response = await fetcher(input, { ...init, redirect: "error", signal: AbortSignal.any([signal, ...(init?.signal ? [init.signal] : [])]) });
|
|
if (!response.ok) { await response.body?.cancel(); throw new GatewayHttpError(response.status); }
|
|
if (!response.body) return response;
|
|
if (Number(response.headers.get("content-length")) > MAX_GATEWAY_BYTES) { await response.body.cancel(); throw failure("SWADS_RESPONSE_TOO_LARGE"); }
|
|
const reader = response.body.getReader(); let total = 0;
|
|
const body = new ReadableStream<Uint8Array>({
|
|
async pull(controller) {
|
|
try {
|
|
const { value, done } = await reader.read();
|
|
if (done) { controller.close(); return; }
|
|
total += value.byteLength;
|
|
if (total > MAX_GATEWAY_BYTES) { await reader.cancel(); controller.error(failure("SWADS_RESPONSE_TOO_LARGE")); return; }
|
|
controller.enqueue(value);
|
|
} catch (error) { controller.error(error); }
|
|
},
|
|
cancel: () => reader.cancel(),
|
|
});
|
|
return new Response(body, { status: response.status, headers: response.headers });
|
|
};
|
|
const client = new Client({ name: "agent-studio-swads-cli", version: "1.0.0" });
|
|
const transport = new StreamableHTTPClientTransport(config.url, { fetch: boundedFetch, requestInit: { headers: { Authorization: `Bearer ${config.token}` } }, reconnectionOptions: { maxRetries: 0, initialReconnectionDelay: 1000, maxReconnectionDelay: 1000, reconnectionDelayGrowFactor: 1 } });
|
|
const close = async () => { await client.close().catch(() => undefined); };
|
|
const abort = () => { void close(); };
|
|
signal.addEventListener("abort", abort, { once: true });
|
|
try { await client.connect(transport as Parameters<Client["connect"]>[0], { signal, timeout: config.timeoutMs }); } catch (error) { signal.removeEventListener("abort", abort); await close(); throw error; }
|
|
return {
|
|
async list(requestSignal) {
|
|
const all: Tool[] = []; let cursor: string | undefined;
|
|
do {
|
|
const page = await client.listTools(cursor ? { cursor } : {}, { signal: requestSignal, timeout: config.timeoutMs });
|
|
all.push(...page.tools); cursor = page.nextCursor;
|
|
if (all.length > 500 || Buffer.byteLength(JSON.stringify(all)) > MAX_GATEWAY_BYTES) throw failure("SWADS_RESPONSE_TOO_LARGE");
|
|
} while (cursor);
|
|
return all;
|
|
},
|
|
call: (name, args, requestSignal) => client.callTool({ name, arguments: args }, undefined, { signal: requestSignal, timeout: config.timeoutMs }),
|
|
async close() { signal.removeEventListener("abort", abort); await close(); },
|
|
};
|
|
};
|
|
}
|
|
class GatewayHttpError extends Error { constructor(readonly status: number) { super("Gateway HTTP request failed"); } }
|
|
export function failure(code: string): HarnessError { return new HarnessError(code, code, 422); }
|
|
export function gatewayError(error: unknown, signal?: AbortSignal): HarnessError {
|
|
if (error instanceof HarnessError && error.code.startsWith("SWADS_")) return failure(error.code);
|
|
if (signal?.aborted) return failure(signal.reason?.name === "TimeoutError" ? "SWADS_TIMEOUT" : "SWADS_ABORTED");
|
|
const status = error instanceof GatewayHttpError ? error.status : (error as { code?: number; status?: number } | null)?.status ?? (error as { code?: number } | null)?.code;
|
|
return failure(status === 401 ? "SWADS_AUTH_FAILED" : status === 403 ? "SWADS_FORBIDDEN" : status === -32601 ? "SWADS_TOOL_MISSING" : status === -32602 ? "SWADS_TOOL_DRIFT" : status === -32001 ? "SWADS_TIMEOUT" : "SWADS_UNAVAILABLE");
|
|
}
|
|
export class SwadsGateway {
|
|
private capabilities: { expires: number; tools: Tool[] } | undefined;
|
|
private readonly connect: GatewayConnector;
|
|
constructor(private readonly config = gatewayConfig(), connector?: GatewayConnector) { this.connect = connector ?? mcpConnector(config); }
|
|
sanitize<T>(value: T): T {
|
|
if (!this.config.token) return value;
|
|
const scrub = (item: unknown): unknown => typeof item === "string" ? item.split(this.config.token!).join("[REDACTED]") : Array.isArray(item) ? item.map(scrub) : item && typeof item === "object" ? Object.fromEntries(Object.entries(item).map(([key, val]) => [String(scrub(key)), /token|authorization|secret|password/i.test(key) ? "[REDACTED]" : scrub(val)])) : item;
|
|
return scrub(value) as T;
|
|
}
|
|
async execute<T>(operation: (connection: GatewayConnection, signal: AbortSignal) => Promise<T>, runSignal: AbortSignal): Promise<T> {
|
|
if (!this.config.token) throw failure("SWADS_NOT_CONFIGURED");
|
|
const signal = AbortSignal.any([runSignal, AbortSignal.timeout(this.config.timeoutMs)]);
|
|
let connection: GatewayConnection | undefined;
|
|
let abort: (() => void) | undefined;
|
|
try {
|
|
signal.throwIfAborted();
|
|
const cancelled = new Promise<never>((_, reject) => { abort = () => reject(gatewayError(undefined, signal)); signal.addEventListener("abort", abort, { once: true }); });
|
|
const work = (async () => {
|
|
connection = await this.connect(signal);
|
|
if (signal.aborted) { await connection.close(); throw gatewayError(undefined, signal); }
|
|
return operation(connection, signal);
|
|
})();
|
|
return this.sanitize(await Promise.race([work, cancelled]));
|
|
} catch (error) { throw gatewayError(error, signal); }
|
|
finally { if (abort) signal.removeEventListener("abort", abort); await connection?.close().catch(() => undefined); }
|
|
}
|
|
async list(connection: GatewayConnection, signal: AbortSignal, refresh = false): Promise<Tool[]> {
|
|
if (!refresh && this.capabilities && this.capabilities.expires > Date.now()) return this.capabilities.tools;
|
|
const tools = await connection.list(signal);
|
|
if (Buffer.byteLength(JSON.stringify(tools)) > MAX_GATEWAY_BYTES) throw failure("SWADS_RESPONSE_TOO_LARGE");
|
|
this.capabilities = { tools, expires: Date.now() + 60_000 }; return tools;
|
|
}
|
|
}
|