307 lines
9.8 KiB
TypeScript
307 lines
9.8 KiB
TypeScript
/**
|
|
* Per-agent JSON-RPC client. Talks to the deployed agent's POST /mcp endpoint.
|
|
*
|
|
* We hand-roll the wire calls (instead of using @modelcontextprotocol/sdk's
|
|
* Client) because (a) the surface we need is tiny — tools/list, tools/call —
|
|
* and (b) we want to send the user's bearer token as a plain Authorization
|
|
* header without coupling to a transport class.
|
|
*/
|
|
const JSON_RPC = "2.0";
|
|
const DEFAULT_RPC_TIMEOUT_MS = 30_000;
|
|
const DEFAULT_STREAM_IDLE_TIMEOUT_MS = 120_000;
|
|
|
|
export interface UpstreamConfig {
|
|
name: string;
|
|
url: string;
|
|
token: string | null;
|
|
requestTimeoutMs?: number;
|
|
streamIdleTimeoutMs?: number;
|
|
/** Forwarded to the upstream as ``params._meta.cp_jwt`` on tools/call.
|
|
* Lets the upstream agent act on the caller's behalf (file/agent CRUD
|
|
* on /v1/me/*) — same path the platform orchestrator uses. */
|
|
cpJwt?: string | null;
|
|
/** Paired with ``cpJwt``. The upstream agent uses this as the base URL
|
|
* when calling back into /v1/me/*. */
|
|
cpUrl?: string | null;
|
|
/** User's MinIO bucket (``user-<id>-files``). Forwarded as
|
|
* ``params._meta.bucket`` so agents that touch the workspace
|
|
* (agent-builder, file tools) get a context whose
|
|
* ``ctx.workspace.bucket`` resolves to the right bucket. */
|
|
bucket?: string | null;
|
|
}
|
|
|
|
export interface UpstreamTool {
|
|
name: string;
|
|
description?: string;
|
|
inputSchema: unknown;
|
|
outputSchema?: unknown;
|
|
}
|
|
|
|
export interface UpstreamCallResult {
|
|
content: Array<{ type: string; text?: string; [k: string]: unknown }>;
|
|
structuredContent?: unknown;
|
|
isError?: boolean;
|
|
}
|
|
|
|
export class UpstreamError extends Error {
|
|
constructor(
|
|
public agent: string,
|
|
public status: number,
|
|
message: string,
|
|
) {
|
|
super(`upstream ${agent}: ${message}`);
|
|
this.name = "UpstreamError";
|
|
}
|
|
}
|
|
|
|
let _nextId = 1;
|
|
|
|
export class UpstreamAgent {
|
|
constructor(private cfg: UpstreamConfig) {
|
|
this.cfg.url = cfg.url.replace(/\/+$/, "");
|
|
}
|
|
|
|
get name(): string {
|
|
return this.cfg.name;
|
|
}
|
|
|
|
private async rpc<T>(method: string, params: unknown): Promise<T> {
|
|
const id = _nextId++;
|
|
const timeoutMs = this.cfg.requestTimeoutMs ?? DEFAULT_RPC_TIMEOUT_MS;
|
|
const controller = new AbortController();
|
|
let timedOut = false;
|
|
const timeout = setTimeout(() => {
|
|
timedOut = true;
|
|
controller.abort();
|
|
}, timeoutMs);
|
|
const headers: Record<string, string> = {
|
|
"content-type": "application/json",
|
|
accept: "application/json",
|
|
};
|
|
if (this.cfg.token) headers["authorization"] = `Bearer ${this.cfg.token}`;
|
|
|
|
try {
|
|
const resp = await fetch(`${this.cfg.url}/mcp`, {
|
|
method: "POST",
|
|
headers,
|
|
body: JSON.stringify({ jsonrpc: JSON_RPC, id, method, params }),
|
|
signal: controller.signal,
|
|
});
|
|
|
|
if (!resp.ok) {
|
|
const text = await resp.text().catch(() => "");
|
|
throw new UpstreamError(this.cfg.name, resp.status, text || resp.statusText);
|
|
}
|
|
|
|
const env: any = await resp.json();
|
|
if (env?.error) {
|
|
throw new UpstreamError(
|
|
this.cfg.name,
|
|
500,
|
|
`${env.error.code}: ${env.error.message}`,
|
|
);
|
|
}
|
|
return env?.result as T;
|
|
} catch (err) {
|
|
if (timedOut) {
|
|
throw new UpstreamError(
|
|
this.cfg.name,
|
|
504,
|
|
`request timed out after ${timeoutMs}ms`,
|
|
);
|
|
}
|
|
throw err;
|
|
} finally {
|
|
clearTimeout(timeout);
|
|
}
|
|
}
|
|
|
|
async listTools(): Promise<UpstreamTool[]> {
|
|
const out = await this.rpc<{ tools: UpstreamTool[] }>("tools/list", {});
|
|
return out?.tools ?? [];
|
|
}
|
|
|
|
async callTool(
|
|
name: string,
|
|
arguments_: Record<string, unknown>,
|
|
opts: {
|
|
onNotification?: (notif: {
|
|
method: string;
|
|
params?: Record<string, unknown>;
|
|
}) => void | Promise<void>;
|
|
progressToken?: string | number;
|
|
} = {},
|
|
): Promise<UpstreamCallResult> {
|
|
const meta: Record<string, unknown> = {};
|
|
if (this.cfg.cpJwt) meta.cp_jwt = this.cfg.cpJwt;
|
|
if (this.cfg.cpUrl) meta.cp_url = this.cfg.cpUrl;
|
|
if (this.cfg.bucket) meta.bucket = this.cfg.bucket;
|
|
if (opts.progressToken !== undefined) meta.progressToken = opts.progressToken;
|
|
const params: Record<string, unknown> = { name, arguments: arguments_ };
|
|
if (Object.keys(meta).length > 0) params._meta = meta;
|
|
return this.streamingCallTool(params, opts.onNotification);
|
|
}
|
|
|
|
/**
|
|
* tools/call always negotiates SSE on the wire — upstream agents can
|
|
* take minutes (agent-builder.build runs an LLM loop + sandbox tests +
|
|
* a deploy + a wait-for-live-card poll), and a plain JSON POST gets
|
|
* killed by ingress idle timeouts (~60s default on traefik). The
|
|
* upstream's SSE response interleaves any elicitation/create requests
|
|
* and a final tools/call result; here we drain the stream and return
|
|
* the final result. (Elicit relaying to the MCP client is not wired
|
|
* yet — any elicit that fires will time out on the upstream side and
|
|
* the skill will surface a soft error.)
|
|
*/
|
|
private async streamingCallTool(
|
|
params: Record<string, unknown>,
|
|
onNotification?: (notif: {
|
|
method: string;
|
|
params?: Record<string, unknown>;
|
|
}) => void | Promise<void>,
|
|
): Promise<UpstreamCallResult> {
|
|
const id = _nextId++;
|
|
const idleTimeoutMs =
|
|
this.cfg.streamIdleTimeoutMs ?? DEFAULT_STREAM_IDLE_TIMEOUT_MS;
|
|
const controller = new AbortController();
|
|
let timedOut = false;
|
|
let idleTimer: ReturnType<typeof setTimeout> | undefined;
|
|
const resetIdleTimer = () => {
|
|
if (idleTimer) clearTimeout(idleTimer);
|
|
idleTimer = setTimeout(() => {
|
|
timedOut = true;
|
|
controller.abort();
|
|
}, idleTimeoutMs);
|
|
};
|
|
resetIdleTimer();
|
|
const headers: Record<string, string> = {
|
|
"content-type": "application/json",
|
|
accept: "text/event-stream, application/json",
|
|
};
|
|
if (this.cfg.token) headers["authorization"] = `Bearer ${this.cfg.token}`;
|
|
|
|
let resp: Response;
|
|
try {
|
|
resp = await fetch(`${this.cfg.url}/mcp`, {
|
|
method: "POST",
|
|
headers,
|
|
body: JSON.stringify({
|
|
jsonrpc: JSON_RPC, id, method: "tools/call", params,
|
|
}),
|
|
signal: controller.signal,
|
|
});
|
|
} catch (err) {
|
|
clearTimeout(idleTimer);
|
|
if (timedOut) {
|
|
throw new UpstreamError(
|
|
this.cfg.name,
|
|
504,
|
|
`stream idle timeout after ${idleTimeoutMs}ms`,
|
|
);
|
|
}
|
|
throw err;
|
|
}
|
|
resetIdleTimer();
|
|
|
|
if (!resp.ok) {
|
|
const text = await resp.text().catch(() => "");
|
|
clearTimeout(idleTimer);
|
|
throw new UpstreamError(this.cfg.name, resp.status, text || resp.statusText);
|
|
}
|
|
|
|
const ctype = resp.headers.get("content-type") || "";
|
|
if (!ctype.includes("text/event-stream")) {
|
|
// Server fell back to JSON (older a2a-pack or non-SSE upstream).
|
|
const env: any = await resp.json();
|
|
if (env?.error) {
|
|
clearTimeout(idleTimer);
|
|
throw new UpstreamError(
|
|
this.cfg.name, 500,
|
|
`${env.error.code}: ${env.error.message}`,
|
|
);
|
|
}
|
|
clearTimeout(idleTimer);
|
|
return env?.result as UpstreamCallResult;
|
|
}
|
|
|
|
if (!resp.body) {
|
|
clearTimeout(idleTimer);
|
|
throw new UpstreamError(this.cfg.name, 502, "empty stream body");
|
|
}
|
|
|
|
const reader = resp.body.getReader();
|
|
const decoder = new TextDecoder();
|
|
let buf = "";
|
|
try {
|
|
while (true) {
|
|
let done: boolean;
|
|
let value: Uint8Array | undefined;
|
|
try {
|
|
({ done, value } = await reader.read());
|
|
} catch (err) {
|
|
if (timedOut) {
|
|
throw new UpstreamError(
|
|
this.cfg.name,
|
|
504,
|
|
`stream idle timeout after ${idleTimeoutMs}ms`,
|
|
);
|
|
}
|
|
throw err;
|
|
}
|
|
if (value) {
|
|
resetIdleTimer();
|
|
buf += decoder.decode(value, { stream: true });
|
|
}
|
|
let sep = buf.indexOf("\n\n");
|
|
while (sep !== -1) {
|
|
const frame = buf.slice(0, sep);
|
|
buf = buf.slice(sep + 2);
|
|
const dataLines: string[] = [];
|
|
for (const ln of frame.split("\n")) {
|
|
if (ln.startsWith("data:")) dataLines.push(ln.slice(5).trimStart());
|
|
}
|
|
if (dataLines.length > 0) {
|
|
const raw = dataLines.join("\n");
|
|
if (raw === "[DONE]") {
|
|
sep = buf.indexOf("\n\n");
|
|
continue;
|
|
}
|
|
let msg: any;
|
|
try { msg = JSON.parse(raw); } catch { msg = null; }
|
|
if (msg && msg.id === id) {
|
|
if (msg.error) {
|
|
throw new UpstreamError(
|
|
this.cfg.name, 500,
|
|
`${msg.error.code}: ${msg.error.message}`,
|
|
);
|
|
}
|
|
return msg.result as UpstreamCallResult;
|
|
}
|
|
// Otherwise it's a server→client message. Notifications
|
|
// (no id) → forward to the MCP client so long-running
|
|
// skills keep the client's tool-call timer alive. Requests
|
|
// (has id, has method) — e.g. elicitation/create — are
|
|
// not relayed yet, so they will time out upstream.
|
|
if (msg && typeof msg.method === "string" && !("id" in msg) && onNotification) {
|
|
try {
|
|
await onNotification({
|
|
method: msg.method,
|
|
params: msg.params as Record<string, unknown> | undefined,
|
|
});
|
|
} catch {
|
|
// Best-effort; never break the stream on a notify error.
|
|
}
|
|
}
|
|
}
|
|
sep = buf.indexOf("\n\n");
|
|
}
|
|
if (done) break;
|
|
}
|
|
} finally {
|
|
clearTimeout(idleTimer);
|
|
try { reader.releaseLock(); } catch { /* ignore */ }
|
|
}
|
|
throw new UpstreamError(this.cfg.name, 502, "stream ended without result");
|
|
}
|
|
}
|