Skip to content

Commit ea99443

Browse files
committed
fix(mcp): make provider clients Worker-safe
1 parent 22b5946 commit ea99443

8 files changed

Lines changed: 230 additions & 25 deletions

File tree

src/providers/excalidraw_mcp/runtime.ts

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import { UnauthorizedError } from "@modelcontextprotocol/sdk/client/auth.js";
55
import { Client } from "@modelcontextprotocol/sdk/client/index.js";
66
import { StreamableHTTPClientTransport, StreamableHTTPError } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
77
import { ErrorCode, McpError } from "@modelcontextprotocol/sdk/types.js";
8+
import { CfWorkerJsonSchemaValidator } from "@modelcontextprotocol/sdk/validation/cfworker";
89
import { createHash } from "node:crypto";
910
import { optionalRecord, optionalString, requiredString } from "../../core/cast.ts";
1011
import { assertPublicHttpUrl, isPrivateNetworkAccessAllowed } from "../../core/request.ts";
@@ -30,6 +31,7 @@ type ExcalidrawMcpToolResult = {
3031
};
3132
const defaultEndpoint = "https://mcp.excalidraw.com";
3233
const requestTimeoutMs = 30_000;
34+
const excalidrawMcpJsonSchemaValidator = new CfWorkerJsonSchemaValidator();
3335

3436
export const excalidrawMcpActionHandlers: Record<string, ExcalidrawMcpActionHandler> = {
3537
read_me(_input, context) {
@@ -153,7 +155,10 @@ async function withExcalidrawMcpClient<T>(
153155
},
154156
},
155157
});
156-
const client = new Client({ name: "oomol-connect-excalidraw-mcp", version: "1.0.0" });
158+
const client = new Client(
159+
{ name: "oomol-connect-excalidraw-mcp", version: "1.0.0" },
160+
{ jsonSchemaValidator: excalidrawMcpJsonSchemaValidator },
161+
);
157162

158163
try {
159164
await client.connect(transport, { timeout: requestTimeoutMs, signal: context.signal });

src/providers/flomo/executors.ts

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import { UnauthorizedError } from "@modelcontextprotocol/sdk/client/auth.js";
1111
import { Client } from "@modelcontextprotocol/sdk/client/index.js";
1212
import { StreamableHTTPClientTransport, StreamableHTTPError } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
1313
import { McpError } from "@modelcontextprotocol/sdk/types.js";
14+
import { CfWorkerJsonSchemaValidator } from "@modelcontextprotocol/sdk/validation/cfworker";
1415
import { createHash } from "node:crypto";
1516
import { optionalString, requiredString } from "../../core/cast.ts";
1617
import {
@@ -30,6 +31,7 @@ const flomoWebhookPathPrefix = "/iwh/";
3031
const flomoMcpEndpoint = "https://flomoapp.com/mcp";
3132
const flomoMcpTokenField = "token";
3233
const flomoRequestTimeoutMs = 30_000;
34+
const flomoMcpJsonSchemaValidator = new CfWorkerJsonSchemaValidator();
3335

3436
type FlomoActionHandler = (input: Record<string, unknown>, context: FlomoActionContext) => Promise<unknown>;
3537
type FlomoMcpToolResult = Awaited<ReturnType<Client["callTool"]>>;
@@ -441,10 +443,13 @@ async function withFlomoMcpClient<T>(
441443
signal: input.signal,
442444
},
443445
});
444-
const client = new Client({
445-
name: "oomol-connect-flomo",
446-
version: "1.0.0",
447-
});
446+
const client = new Client(
447+
{
448+
name: "oomol-connect-flomo",
449+
version: "1.0.0",
450+
},
451+
{ jsonSchemaValidator: flomoMcpJsonSchemaValidator },
452+
);
448453

449454
try {
450455
await client.connect(transport, {

src/providers/jin10/executors.ts

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import { UnauthorizedError } from "@modelcontextprotocol/sdk/client/auth.js";
66
import { Client } from "@modelcontextprotocol/sdk/client/index.js";
77
import { StreamableHTTPClientTransport, StreamableHTTPError } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
88
import { McpError } from "@modelcontextprotocol/sdk/types.js";
9+
import { CfWorkerJsonSchemaValidator } from "@modelcontextprotocol/sdk/validation/cfworker";
910
import { createHash } from "node:crypto";
1011
import {
1112
defineApiKeyProviderExecutors,
@@ -19,6 +20,7 @@ const jin10McpOrigin = "https://mcp.jin10.com";
1920
const jin10McpEndpoint = "https://mcp.jin10.com/mcp";
2021
const jin10QuoteCodesResourceUri = "quote://codes";
2122
const jin10RequestTimeoutMs = 30_000;
23+
const jin10McpJsonSchemaValidator = new CfWorkerJsonSchemaValidator();
2224

2325
type Jin10ActionContext = Pick<ApiKeyProviderContext, "apiKey" | "fetcher" | "signal">;
2426
type Jin10ActionHandler = (input: Record<string, unknown>, context: Jin10ActionContext) => Promise<unknown>;
@@ -176,10 +178,13 @@ async function withJin10McpClient<T>(
176178
signal: input.signal,
177179
},
178180
});
179-
const client = new Client({
180-
name: "oomol-connect-jin10",
181-
version: "1.0.0",
182-
});
181+
const client = new Client(
182+
{
183+
name: "oomol-connect-jin10",
184+
version: "1.0.0",
185+
},
186+
{ jsonSchemaValidator: jin10McpJsonSchemaValidator },
187+
);
183188

184189
try {
185190
await client.connect(transport, {

src/providers/jumpserver/definition.ts

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ import { jumpServerActions } from "./actions.ts";
44

55
const service = "jumpserver";
66

7-
/** JumpServer provider backed by the official jumpserver/mcp SSE server. */
7+
/** JumpServer provider backed by the official jumpserver/mcp server. */
88
export const provider: ProviderDefinition = {
99
service,
1010
displayName: "JumpServer",
@@ -18,13 +18,13 @@ export const provider: ProviderDefinition = {
1818
fields: [
1919
{
2020
key: "mcpEndpoint",
21-
label: "MCP SSE Endpoint",
21+
label: "MCP Endpoint",
2222
inputType: "text",
2323
required: true,
2424
secret: false,
2525
placeholder: "https://jumpserver-mcp.example.com/sse",
2626
description:
27-
"The SSE endpoint of the official jumpserver/mcp server. Public HTTPS endpoints are supported by default. Private-network, Tailscale, and NetBird endpoints require OOMOL_CONNECT_ALLOW_PRIVATE_NETWORK. Loopback endpoints remain blocked. See https://github.com/jumpserver/mcp.",
27+
"The Streamable HTTP endpoint, or the legacy SSE endpoint exposed by official jumpserver/mcp deployments. Public HTTPS endpoints are supported by default. Private-network, Tailscale, and NetBird endpoints require OOMOL_CONNECT_ALLOW_PRIVATE_NETWORK. Loopback endpoints remain blocked. See https://github.com/jumpserver/mcp.",
2828
},
2929
{
3030
key: "token",
Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
1+
import { afterEach, describe, expect, it, vi } from "vitest";
2+
import { validateJumpServerCredential } from "./runtime.ts";
3+
4+
afterEach(() => {
5+
vi.unstubAllGlobals();
6+
});
7+
8+
describe("JumpServer MCP runtime", () => {
9+
it("uses Streamable HTTP with Worker-safe tool schema validation", async () => {
10+
const credential = {
11+
mcpEndpoint: "https://jumpserver.example.com/mcp",
12+
token: "jumpserver-token",
13+
};
14+
15+
// Warm the SDK's Node-only Zod fast path so this test isolates tool output validation.
16+
await validateJumpServerCredential(credential, createStreamableMcpFetch());
17+
vi.stubGlobal("Function", function disabledFunctionConstructor() {
18+
throw new EvalError("Code generation from strings disallowed for this context");
19+
});
20+
21+
const result = await validateJumpServerCredential(credential, createStreamableMcpFetch());
22+
23+
expect(result.metadata).toMatchObject({
24+
mcpEndpoint: "https://jumpserver.example.com/mcp",
25+
availableActions: ["assets_assets_list"],
26+
});
27+
});
28+
29+
it("falls back to legacy SSE when Streamable HTTP is unavailable", async () => {
30+
const requests: Array<{ method: string; pathname: string }> = [];
31+
32+
const result = await validateJumpServerCredential(
33+
{
34+
mcpEndpoint: "https://jumpserver.example.com/sse",
35+
token: "jumpserver-token",
36+
},
37+
createLegacySseMcpFetch(requests),
38+
);
39+
40+
expect(result.metadata).toMatchObject({
41+
mcpEndpoint: "https://jumpserver.example.com/sse",
42+
availableActions: ["assets_assets_list"],
43+
});
44+
expect(requests).toContainEqual({ method: "POST", pathname: "/sse" });
45+
expect(requests).toContainEqual({ method: "GET", pathname: "/sse" });
46+
expect(requests).toContainEqual({ method: "POST", pathname: "/messages" });
47+
});
48+
});
49+
50+
function createStreamableMcpFetch(): typeof fetch {
51+
return vi.fn(async (_input: RequestInfo | URL, init?: RequestInit): Promise<Response> => {
52+
const method = init?.method ?? "GET";
53+
if (method === "GET") {
54+
return new Response(null, { status: 405 });
55+
}
56+
57+
const request = readRequest(init);
58+
if (!("id" in request)) {
59+
return new Response(null, { status: 202 });
60+
}
61+
62+
const result =
63+
request.method === "initialize"
64+
? {
65+
protocolVersion: "2025-03-26",
66+
capabilities: {},
67+
serverInfo: { name: "jumpserver", version: "1.0.0" },
68+
}
69+
: {
70+
tools: [
71+
{
72+
name: "assets_assets_list",
73+
inputSchema: { type: "object" },
74+
outputSchema: { type: "object" },
75+
},
76+
],
77+
};
78+
79+
return new Response(JSON.stringify({ jsonrpc: "2.0", id: request.id, result }), {
80+
headers: {
81+
"content-type": "application/json",
82+
"mcp-session-id": "test-session",
83+
},
84+
});
85+
}) as typeof fetch;
86+
}
87+
88+
function createLegacySseMcpFetch(requests: Array<{ method: string; pathname: string }>): typeof fetch {
89+
const encoder = new TextEncoder();
90+
let streamController: ReadableStreamDefaultController<Uint8Array> | undefined;
91+
92+
return vi.fn(async (input: RequestInfo | URL, init?: RequestInit): Promise<Response> => {
93+
const url = new URL(input instanceof Request ? input.url : input.toString());
94+
const method = init?.method ?? "GET";
95+
requests.push({ method, pathname: url.pathname });
96+
97+
if (method === "POST" && url.pathname === "/sse") {
98+
return new Response(null, { status: 405 });
99+
}
100+
if (method === "GET") {
101+
return new Response(
102+
new ReadableStream<Uint8Array>({
103+
start(controller) {
104+
streamController = controller;
105+
controller.enqueue(encoder.encode("event: endpoint\ndata: /messages?session_id=test-session\n\n"));
106+
},
107+
}),
108+
{ headers: { "content-type": "text/event-stream" } },
109+
);
110+
}
111+
112+
const request = readRequest(init);
113+
if ("id" in request) {
114+
const result =
115+
request.method === "initialize"
116+
? {
117+
protocolVersion: "2024-11-05",
118+
capabilities: {},
119+
serverInfo: { name: "jumpserver", version: "1.0.0" },
120+
}
121+
: {
122+
tools: [
123+
{
124+
name: "assets_assets_list",
125+
inputSchema: { type: "object" },
126+
},
127+
],
128+
};
129+
streamController!.enqueue(
130+
encoder.encode(`event: message\ndata: ${JSON.stringify({ jsonrpc: "2.0", id: request.id, result })}\n\n`),
131+
);
132+
}
133+
return new Response(null, { status: 202 });
134+
}) as typeof fetch;
135+
}
136+
137+
function readRequest(init?: RequestInit): Record<string, unknown> {
138+
return typeof init?.body === "string" ? (JSON.parse(init.body) as Record<string, unknown>) : {};
139+
}

src/providers/jumpserver/runtime.ts

Lines changed: 49 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,9 @@ import type { CredentialValidationResult } from "../../core/types.ts";
33
import { UnauthorizedError } from "@modelcontextprotocol/sdk/client/auth.js";
44
import { Client } from "@modelcontextprotocol/sdk/client/index.js";
55
import { SSEClientTransport, SseError } from "@modelcontextprotocol/sdk/client/sse.js";
6+
import { StreamableHTTPClientTransport, StreamableHTTPError } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
67
import { McpError } from "@modelcontextprotocol/sdk/types.js";
8+
import { CfWorkerJsonSchemaValidator } from "@modelcontextprotocol/sdk/validation/cfworker";
79
import { createHash } from "node:crypto";
810
import { requiredString } from "../../core/cast.ts";
911
import { assertPublicHttpUrl, isPrivateNetworkAccessAllowed } from "../../core/request.ts";
@@ -21,6 +23,7 @@ export interface JumpServerMcpContext {
2123
}
2224

2325
const requestTimeoutMs = 60_000;
26+
const jumpServerMcpJsonSchemaValidator = new CfWorkerJsonSchemaValidator();
2427

2528
export const jumpServerActionHandlers: Record<string, JumpServerActionHandler> = {};
2629
for (const toolName of jumpServerMcpToolNames) {
@@ -123,20 +126,58 @@ async function withJumpServerMcpClient<T>(
123126
Authorization: `Bearer ${context.token}`,
124127
"user-agent": providerUserAgent,
125128
});
126-
const transport = new SSEClientTransport(context.endpoint, {
127-
fetch: context.fetcher,
128-
requestInit: { headers, signal: context.signal },
129-
});
130-
const client = new Client({ name: "oomol-connect-jumpserver", version: "1.0.0" });
129+
let client: Client | undefined;
131130

132131
try {
133-
await client.connect(transport, { timeout: requestTimeoutMs });
132+
client = await connectJumpServerMcpClient(context, headers);
134133
return await run(client);
135134
} catch (error) {
136135
throw mapJumpServerMcpError(error);
137136
} finally {
138-
await client.close().catch(() => undefined);
137+
await client?.close().catch(() => undefined);
138+
}
139+
}
140+
141+
async function connectJumpServerMcpClient(context: JumpServerMcpContext, headers: Headers): Promise<Client> {
142+
const streamableClient = createJumpServerMcpClient();
143+
const streamableTransport = new StreamableHTTPClientTransport(context.endpoint, {
144+
fetch: context.fetcher,
145+
requestInit: { headers, signal: context.signal },
146+
});
147+
148+
try {
149+
await streamableClient.connect(streamableTransport, { timeout: requestTimeoutMs });
150+
return streamableClient;
151+
} catch (error) {
152+
await streamableClient.close().catch(() => undefined);
153+
if (!isUnsupportedStreamableHttp(error)) {
154+
throw error;
155+
}
139156
}
157+
158+
const legacyClient = createJumpServerMcpClient();
159+
const legacyTransport = new SSEClientTransport(context.endpoint, {
160+
fetch: context.fetcher,
161+
requestInit: { headers, signal: context.signal },
162+
});
163+
try {
164+
await legacyClient.connect(legacyTransport, { timeout: requestTimeoutMs });
165+
return legacyClient;
166+
} catch (error) {
167+
await legacyClient.close().catch(() => undefined);
168+
throw error;
169+
}
170+
}
171+
172+
function createJumpServerMcpClient(): Client {
173+
return new Client(
174+
{ name: "oomol-connect-jumpserver", version: "1.0.0" },
175+
{ jsonSchemaValidator: jumpServerMcpJsonSchemaValidator },
176+
);
177+
}
178+
179+
function isUnsupportedStreamableHttp(error: unknown): boolean {
180+
return error instanceof StreamableHTTPError && (error.code === 404 || error.code === 405);
140181
}
141182

142183
function normalizeJumpServerMcpToolResult(toolName: string, result: JumpServerMcpToolResult): unknown {
@@ -182,7 +223,7 @@ function mapJumpServerMcpError(error: unknown): ProviderRequestError {
182223
if (error instanceof UnauthorizedError) {
183224
return new ProviderRequestError(401, "JumpServer MCP token is invalid or expired", error);
184225
}
185-
if (error instanceof SseError) {
226+
if (error instanceof SseError || error instanceof StreamableHTTPError) {
186227
const status = error.code;
187228
return new ProviderRequestError(
188229
status === 401 || status === 403 ? 401 : status && status >= 400 && status < 500 ? 400 : 502,

src/providers/lingxing/runtime.ts

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import { UnauthorizedError } from "@modelcontextprotocol/sdk/client/auth.js";
55
import { Client } from "@modelcontextprotocol/sdk/client/index.js";
66
import { StreamableHTTPClientTransport, StreamableHTTPError } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
77
import { ErrorCode, McpError } from "@modelcontextprotocol/sdk/types.js";
8+
import { CfWorkerJsonSchemaValidator } from "@modelcontextprotocol/sdk/validation/cfworker";
89
import { createHash } from "node:crypto";
910
import { optionalRecord, requiredString } from "../../core/cast.ts";
1011
import { assertPublicHttpUrl } from "../../core/request.ts";
@@ -14,6 +15,7 @@ const lingxingMcpHost = "openmcp.lingxing.com";
1415
const lingxingRequestTimeoutMs = 30_000;
1516
const lingxingToolIntervalMs = 1_000;
1617
const maximumTrackedRateLimitKeys = 1_024;
18+
const lingxingMcpJsonSchemaValidator = new CfWorkerJsonSchemaValidator();
1719

1820
interface LingxingCredential {
1921
endpoint: URL;
@@ -254,10 +256,13 @@ async function withLingxingMcpClient<T>(context: LingxingContext, run: (client:
254256
signal: context.signal,
255257
},
256258
});
257-
const client = new Client({
258-
name: "oomol-connect-lingxing",
259-
version: "1.0.0",
260-
});
259+
const client = new Client(
260+
{
261+
name: "oomol-connect-lingxing",
262+
version: "1.0.0",
263+
},
264+
{ jsonSchemaValidator: lingxingMcpJsonSchemaValidator },
265+
);
261266

262267
try {
263268
await client.connect(transport, {

0 commit comments

Comments
 (0)