| @@ -1,4 +1,5 @@ |
| 1 | import { createServer, type IncomingMessage, type ServerResponse } from "node:http"; | 1 | import { createServer, type IncomingMessage, type ServerResponse } from "node:http"; |
| | 2 | import { randomUUID } from "node:crypto"; |
| 2 | | 3 | |
| 3 | import { | 4 | import { |
| 4 | HttpApi, | 5 | HttpApi, |
| @@ -16,9 +17,14 @@ import { makeToolInvokerFromTools } from "@executor-v3/codemode-core"; · 11 unmodified lines |
| 16 | import { | 17 | import { |
| 17 | ControlPlaneStorageError, | 18 | ControlPlaneStorageError, |
| 18 | makeControlPlaneClient, | 19 | makeControlPlaneClient, |
| | 20 | type ControlPlaneClient, |
| 19 | type ResolveExecutionEnvironment, | 21 | type ResolveExecutionEnvironment, |
| 20 | } from "@executor-v3/control-plane"; | 22 | } from "@executor-v3/control-plane"; |
| 21 | import { makeInProcessExecutor } from "@executor-v3/runtime-local-inproc"; | 23 | import { makeInProcessExecutor } from "@executor-v3/runtime-local-inproc"; |
| | 24 | import { createMcpExpressApp } from "@modelcontextprotocol/sdk/server/express.js"; |
| | 25 | import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; |
| | 26 | import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js"; |
| | 27 | import { z } from "zod/v4"; |
| 22 | | 28 | |
| 23 | import { | 29 | import { |
| 24 | seedDemoMcpSourceInWorkspace, | 30 | seedDemoMcpSourceInWorkspace, |
| @@ -122,6 +128,304 @@ const startOpenApiDemoServer = async () => { · 97 unmodified lines |
| 122 | }; | 128 | }; |
| 123 | }; | 129 | }; |
| 124 | | 130 | |
| | 131 | const DEMO_OAUTH_ACCESS_TOKEN = "demo-access-token"; |
| | 132 | const DEMO_OAUTH_REFRESH_TOKEN = "demo-refresh-token"; |
| | 133 | |
| | 134 | const startOAuthProtectedMcpServer = async () => { |
| | 135 | const host = "127.0.0.1"; |
| | 136 | const app = createMcpExpressApp({ host }); |
| | 137 | const transports: Record<string, StreamableHTTPServerTransport> = {}; |
| | 138 | const servers: Record<string, McpServer> = {}; |
| | 139 | |
| | 140 | const createAuthorizedServer = () => { |
| | 141 | const server = new McpServer( |
| | 142 | { |
| | 143 | name: "executor-oauth-mcp-demo", |
| | 144 | version: "1.0.0", |
| | 145 | }, |
| | 146 | { |
| | 147 | capabilities: { |
| | 148 | tools: {}, |
| | 149 | }, |
| | 150 | }, |
| | 151 | ); |
| | 152 | |
| | 153 | server.registerTool( |
| | 154 | "whoami", |
| | 155 | { |
| | 156 | description: "Return the active identity for the OAuth demo source.", |
| | 157 | inputSchema: {}, |
| | 158 | }, |
| | 159 | async () => ({ |
| | 160 | content: [{ type: "text", text: "oauth-demo" }], |
| | 161 | }), |
| | 162 | ); |
| | 163 | |
| | 164 | return server; |
| | 165 | }; |
| | 166 | |
| | 167 | const listener = await new Promise<import("node:http").Server>((resolve, reject) => { |
| | 168 | const server = app.listen(0, host); |
| | 169 | |
| | 170 | const onListening = () => { |
| | 171 | server.off("error", onError); |
| | 172 | resolve(server); |
| | 173 | }; |
| | 174 | |
| | 175 | const onError = (error: Error) => { |
| | 176 | server.off("listening", onListening); |
| | 177 | reject(error); |
| | 178 | }; |
| | 179 | |
| | 180 | server.once("listening", onListening); |
| | 181 | server.once("error", onError); |
| | 182 | }); |
| | 183 | |
| | 184 | const address = listener.address(); |
| | 185 | if (!address || typeof address === "string") { |
| | 186 | throw new Error("Failed to bind OAuth MCP demo server"); |
| | 187 | } |
| | 188 | |
| | 189 | const baseUrl = `http://${host}:${address.port}`; |
| | 190 | const endpoint = `${baseUrl}/mcp`; |
| | 191 | const resourceMetadataUrl = `${baseUrl}/.well-known/oauth-protected-resource/mcp`; |
| | 192 | |
| | 193 | const requireBearerAuth = (req: IncomingMessage, res: ServerResponse) => { |
| | 194 | const authorization = req.headers.authorization; |
| | 195 | if (authorization === `Bearer ${DEMO_OAUTH_ACCESS_TOKEN}`) { |
| | 196 | return true; |
| | 197 | } |
| | 198 | |
| | 199 | res.statusCode = 401; |
| | 200 | res.setHeader( |
| | 201 | "WWW-Authenticate", |
| | 202 | `Bearer resource_metadata="${resourceMetadataUrl}"`, |
| | 203 | ); |
| | 204 | res.setHeader("content-type", "application/json"); |
| | 205 | res.end( |
| | 206 | JSON.stringify({ |
| | 207 | error: "unauthorized", |
| | 208 | }), |
| | 209 | ); |
| | 210 | return false; |
| | 211 | }; |
| | 212 | |
| | 213 | app.get("/.well-known/oauth-protected-resource/mcp", (_req: any, res: any) => { |
| | 214 | res.status(200).json({ |
| | 215 | resource: endpoint, |
| | 216 | authorization_servers: [baseUrl], |
| | 217 | scopes_supported: ["openid", "offline_access"], |
| | 218 | bearer_methods_supported: ["header"], |
| | 219 | }); |
| | 220 | }); |
| | 221 | |
| | 222 | app.get("/.well-known/oauth-authorization-server", (_req: any, res: any) => { |
| | 223 | res.status(200).json({ |
| | 224 | issuer: baseUrl, |
| | 225 | authorization_endpoint: `${baseUrl}/authorize`, |
| | 226 | token_endpoint: `${baseUrl}/token`, |
| | 227 | registration_endpoint: `${baseUrl}/register`, |
| | 228 | response_types_supported: ["code"], |
| | 229 | grant_types_supported: ["authorization_code", "refresh_token"], |
| | 230 | code_challenge_methods_supported: ["S256"], |
| | 231 | token_endpoint_auth_methods_supported: ["none"], |
| | 232 | }); |
| | 233 | }); |
| | 234 | |
| | 235 | app.post("/register", (req: any, res: any) => { |
| | 236 | res.status(201).json({ |
| | 237 | ...(req.body ?? {}), |
| | 238 | client_id: `client_${randomUUID()}`, |
| | 239 | client_id_issued_at: Math.floor(Date.now() / 1000), |
| | 240 | }); |
| | 241 | }); |
| | 242 | |
| | 243 | app.get("/authorize", (req: any, res: any) => { |
| | 244 | const redirectUri = new URL(String(req.query.redirect_uri)); |
| | 245 | redirectUri.searchParams.set("code", "demo-code"); |
| | 246 | if (typeof req.query.state === "string") { |
| | 247 | redirectUri.searchParams.set("state", req.query.state); |
| | 248 | } |
| | 249 | |
| | 250 | res.redirect(302, redirectUri.toString()); |
| | 251 | }); |
| | 252 | |
| | 253 | app.post("/token", (_req: any, res: any) => { |
| | 254 | res.status(200).json({ |
| | 255 | access_token: DEMO_OAUTH_ACCESS_TOKEN, |
| | 256 | token_type: "Bearer", |
| | 257 | expires_in: 3600, |
| | 258 | refresh_token: DEMO_OAUTH_REFRESH_TOKEN, |
| | 259 | scope: "openid offline_access", |
| | 260 | }); |
| | 261 | }); |
| | 262 | |
| | 263 | app.post("/mcp", async (req: any, res: any) => { |
| | 264 | if (!requireBearerAuth(req, res)) { |
| | 265 | return; |
| | 266 | } |
| | 267 | |
| | 268 | const sessionIdHeader = req.headers["mcp-session-id"]; |
| | 269 | const sessionId = |
| | 270 | typeof sessionIdHeader === "string" |
| | 271 | ? sessionIdHeader |
| | 272 | : Array.isArray(sessionIdHeader) |
| | 273 | ? sessionIdHeader[0] |
| | 274 | : undefined; |
| | 275 | |
| | 276 | try { |
| | 277 | let transport: StreamableHTTPServerTransport; |
| | 278 | |
| | 279 | if (sessionId && transports[sessionId]) { |
| | 280 | transport = transports[sessionId]; |
| | 281 | } else { |
| | 282 | transport = new StreamableHTTPServerTransport({ |
| | 283 | sessionIdGenerator: () => randomUUID(), |
| | 284 | onsessioninitialized: (newSessionId) => { |
| | 285 | transports[newSessionId] = transport; |
| | 286 | }, |
| | 287 | }); |
| | 288 | |
| | 289 | transport.onclose = () => { |
| | 290 | const closedSessionId = transport.sessionId; |
| | 291 | if (closedSessionId && transports[closedSessionId]) { |
| | 292 | delete transports[closedSessionId]; |
| | 293 | } |
| | 294 | if (closedSessionId && servers[closedSessionId]) { |
| | 295 | void servers[closedSessionId].close().catch(() => undefined); |
| | 296 | delete servers[closedSessionId]; |
| | 297 | } |
| | 298 | }; |
| | 299 | |
| | 300 | const server = createAuthorizedServer(); |
| | 301 | await server.connect(transport); |
| | 302 | const newSessionId = transport.sessionId; |
| | 303 | if (newSessionId) { |
| | 304 | servers[newSessionId] = server; |
| | 305 | } |
| | 306 | } |
| | 307 | |
| | 308 | await transport.handleRequest(req, res, req.body); |
| | 309 | } catch (error) { |
| | 310 | if (!res.headersSent) { |
| | 311 | res.status(500).json({ |
| | 312 | jsonrpc: "2.0", |
| | 313 | error: { |
| | 314 | code: -32603, |
| | 315 | message: error instanceof Error ? error.message : "Internal server error", |
| | 316 | }, |
| | 317 | id: null, |
| | 318 | }); |
| | 319 | } |
| | 320 | } |
| | 321 | }); |
| | 322 | |
| | 323 | app.get("/mcp", async (req: any, res: any) => { |
| | 324 | if (!requireBearerAuth(req, res)) { |
| | 325 | return; |
| | 326 | } |
| | 327 | |
| | 328 | const sessionIdHeader = req.headers["mcp-session-id"]; |
| | 329 | const sessionId = |
| | 330 | typeof sessionIdHeader === "string" |
| | 331 | ? sessionIdHeader |
| | 332 | : Array.isArray(sessionIdHeader) |
| | 333 | ? sessionIdHeader[0] |
| | 334 | : undefined; |
| | 335 | |
| | 336 | if (!sessionId || !transports[sessionId]) { |
| | 337 | res.status(400).send("Invalid or missing session ID"); |
| | 338 | return; |
| | 339 | } |
| | 340 | |
| | 341 | await transports[sessionId].handleRequest(req, res); |
| | 342 | }); |
| | 343 | |
| | 344 | app.delete("/mcp", async (req: any, res: any) => { |
| | 345 | if (!requireBearerAuth(req, res)) { |
| | 346 | return; |
| | 347 | } |
| | 348 | |
| | 349 | const sessionIdHeader = req.headers["mcp-session-id"]; |
| | 350 | const sessionId = |
| | 351 | typeof sessionIdHeader === "string" |
| | 352 | ? sessionIdHeader |
| | 353 | : Array.isArray(sessionIdHeader) |
| | 354 | ? sessionIdHeader[0] |
| | 355 | : undefined; |
| | 356 | |
| | 357 | if (!sessionId || !transports[sessionId]) { |
| | 358 | res.status(400).send("Invalid or missing session ID"); |
| | 359 | return; |
| | 360 | } |
| | 361 | |
| | 362 | const transport = transports[sessionId]; |
| | 363 | await transport.handleRequest(req, res, req.body); |
| | 364 | await transport.close(); |
| | 365 | delete transports[sessionId]; |
| | 366 | |
| | 367 | if (servers[sessionId]) { |
| | 368 | await servers[sessionId].close().catch(() => undefined); |
| | 369 | delete servers[sessionId]; |
| | 370 | } |
| | 371 | }); |
| | 372 | |
| | 373 | return { |
| | 374 | baseUrl, |
| | 375 | endpoint, |
| | 376 | close: async () => { |
| | 377 | for (const transport of Object.values(transports)) { |
| | 378 | await transport.close().catch(() => undefined); |
| | 379 | } |
| | 380 | |
| | 381 | for (const server of Object.values(servers)) { |
| | 382 | await server.close().catch(() => undefined); |
| | 383 | } |
| | 384 | |
| | 385 | await new Promise<void>((resolve, reject) => { |
| | 386 | listener.close((error) => (error ? reject(error) : resolve())); |
| | 387 | }); |
| | 388 | }, |
| | 389 | }; |
| | 390 | }; |
| | 391 | |
| | 392 | const extractUrlInteractionUrl = (payloadJson: string): string => { |
| | 393 | const parsed = JSON.parse(payloadJson) as { |
| | 394 | elicitation?: { |
| | 395 | url?: string; |
| | 396 | }; |
| | 397 | }; |
| | 398 | |
| | 399 | const url = parsed.elicitation?.url; |
| | 400 | if (typeof url !== "string" || url.trim().length === 0) { |
| | 401 | throw new Error("Missing URL elicitation in payloadJson"); |
| | 402 | } |
| | 403 | |
| | 404 | return url; |
| | 405 | }; |
| | 406 | |
| | 407 | const waitForExecutionCompletion = (input: { |
| | 408 | client: ControlPlaneClient; |
| | 409 | workspaceId: string; |
| | 410 | executionId: string; |
| | 411 | }) => |
| | 412 | Effect.gen(function* () { |
| | 413 | while (true) { |
| | 414 | const next = yield* input.client.executions.get({ |
| | 415 | path: { |
| | 416 | workspaceId: input.workspaceId as never, |
| | 417 | executionId: input.executionId as never, |
| | 418 | }, |
| | 419 | }); |
| | 420 | |
| | 421 | if (next.execution.status !== "waiting_for_interaction") { |
| | 422 | return next; |
| | 423 | } |
| | 424 | |
| | 425 | yield* Effect.sleep("100 millis"); |
| | 426 | } |
| | 427 | }); |
| | 428 | |
| 125 | describe("local-executor-server", () => { | 429 | describe("local-executor-server", () => { |
| 126 | it.scoped("serves the control-plane API and executes code", () => | 430 | it.scoped("serves the control-plane API and executes code", () => |
| 127 | Effect.gen(function* () { | 431 | Effect.gen(function* () { |
| @@ -296,6 +600,7 @@ describe("local-executor-server", () => { · 168 unmodified lines |
| 296 | expect(resumed.execution.resultJson).toContain(`approved:${value}`); | 600 | expect(resumed.execution.resultJson).toContain(`approved:${value}`); |
| 297 | } | 601 | } |
| 298 | }), | 602 | }), |
| | 603 | 15_000, |
| 299 | ); | 604 | ); |
| 300 | | 605 | |
| 301 | it.scoped("updates an existing demo MCP source instead of creating a duplicate", () => | 606 | it.scoped("updates an existing demo MCP source instead of creating a duplicate", () => |
| @@ -413,6 +718,89 @@ describe("local-executor-server", () => { · 111 unmodified lines |
| 413 | }), | 718 | }), |
| 414 | ); | 719 | ); |
| 415 | | 720 | |
| | 721 | it.scoped("adds an OAuth-protected MCP source via executor.sources.add and resumes after callback", () => |
| | 722 | Effect.gen(function* () { |
| | 723 | const oauthServer = yield* Effect.acquireRelease( |
| | 724 | Effect.promise(() => startOAuthProtectedMcpServer()), |
| | 725 | (server) => Effect.promise(() => server.close()).pipe(Effect.orDie), |
| | 726 | ); |
| | 727 | |
| | 728 | const server = yield* makeLocalExecutorServer({ |
| | 729 | port: 0, |
| | 730 | localDataDir: ":memory:", |
| | 731 | }); |
| | 732 | |
| | 733 | const bootstrapClient = yield* makeControlPlaneClient({ |
| | 734 | baseUrl: server.baseUrl, |
| | 735 | }); |
| | 736 | const installation = yield* bootstrapClient.local.installation({}); |
| | 737 | const client = yield* makeControlPlaneClient({ |
| | 738 | baseUrl: server.baseUrl, |
| | 739 | accountId: installation.accountId, |
| | 740 | }); |
| | 741 | |
| | 742 | const added = yield* client.executions.create({ |
| | 743 | path: { |
| | 744 | workspaceId: installation.workspaceId, |
| | 745 | }, |
| | 746 | payload: { |
| | 747 | code: `return await tools.executor.sources.add({ endpoint: ${JSON.stringify(oauthServer.endpoint)}, name: "Axiom", namespace: "axiom" });`, |
| | 748 | }, |
| | 749 | }); |
| | 750 | |
| | 751 | expect(added.execution.status).toBe("waiting_for_interaction"); |
| | 752 | expect(added.pendingInteraction).not.toBeNull(); |
| | 753 | if (added.pendingInteraction === null) { |
| | 754 | throw new Error("Expected pending OAuth interaction"); |
| | 755 | } |
| | 756 | expect(added.pendingInteraction.kind).toBe("url"); |
| | 757 | |
| | 758 | const authorizationUrl = extractUrlInteractionUrl(added.pendingInteraction.payloadJson); |
| | 759 | const callbackResponse = yield* Effect.promise(() => |
| | 760 | fetch(authorizationUrl, { |
| | 761 | redirect: "follow", |
| | 762 | }), |
| | 763 | ); |
| | 764 | expect(callbackResponse.ok).toBe(true); |
| | 765 | expect(yield* Effect.promise(() => callbackResponse.text())).toContain("Source connected:"); |
| | 766 | |
| | 767 | const completed = yield* waitForExecutionCompletion({ |
| | 768 | client, |
| | 769 | workspaceId: installation.workspaceId, |
| | 770 | executionId: added.execution.id, |
| | 771 | }); |
| | 772 | |
| | 773 | expect(completed.execution.status).toBe("completed"); |
| | 774 | expect(completed.pendingInteraction).toBeNull(); |
| | 775 | expect(completed.execution.resultJson).toContain('"name":"Axiom"'); |
| | 776 | expect(completed.execution.resultJson).toContain('"status":"connected"'); |
| | 777 | |
| | 778 | const sources = yield* client.sources.list({ |
| | 779 | path: { |
| | 780 | workspaceId: installation.workspaceId, |
| | 781 | }, |
| | 782 | }); |
| | 783 | |
| | 784 | expect(sources).toHaveLength(1); |
| | 785 | expect(sources[0]?.name).toBe("Axiom"); |
| | 786 | expect(sources[0]?.status).toBe("connected"); |
| | 787 | expect(sources[0]?.auth.kind).toBe("oauth2"); |
| | 788 | |
| | 789 | const toolCall = yield* client.executions.create({ |
| | 790 | path: { |
| | 791 | workspaceId: installation.workspaceId, |
| | 792 | }, |
| | 793 | payload: { |
| | 794 | code: "return await tools.axiom.whoami({});", |
| | 795 | }, |
| | 796 | }); |
| | 797 | |
| | 798 | expect(toolCall.execution.status).toBe("completed"); |
| | 799 | expect(toolCall.pendingInteraction).toBeNull(); |
| | 800 | expect(toolCall.execution.resultJson).toContain("oauth-demo"); |
| | 801 | }), |
| | 802 | ); |
| | 803 | |
| 416 | it.scoped("returns a typed storage error when a configured MCP endpoint is invalid", () => | 804 | it.scoped("returns a typed storage error when a configured MCP endpoint is invalid", () => |
| 417 | Effect.gen(function* () { | 805 | Effect.gen(function* () { |
| 418 | const server = yield* makeLocalExecutorServer({ | 806 | const server = yield* makeLocalExecutorServer({ |