History

Add executor MCP source OAuth flow

Rhys Sullivan committed 8 months ago

Commit d48e6bc1 parent ea7c60b25 files · +4295 −23Browse files

Showing the changes to apps (4 of 25 files). Show all

Unified · Split

4 files
  1. apps/executor/package.json +3 −1
  2. apps/executor/src/cli/main.ts +47 −2
  3. apps/executor/src/server/index.ts +6 −1
  4. apps/executor/src/server/server.test.ts +388 −0
apps/executor/package.json+3−1
@@ -25,8 +25,10 @@ · 24 unmodified lines
2525 "devDependencies": {
2626 "@executor-v3/mcp-elicitation-demo": "workspace:*",
2727 "@effect/vitest": "catalog:",
28+ "@modelcontextprotocol/sdk": "^1.26.0",
2829 "@types/node": "catalog:",
2930 "bun-types": "catalog:",
30− "vitest": "catalog:"
31+ "vitest": "catalog:",
32+ "zod": "catalog:"
3133 }
3234 }
apps/executor/src/cli/main.ts+47−2
@@ -205,8 +205,7 @@ const promptInteraction = (interaction: ExecutionInteraction) => · 204 unmodified lines
205205 yield* Effect.sync(() => {
206206 process.stdout.write(`${parsed.message}\n${parsed.url ?? ""}\n`);
207207 });
208− yield* promptLine("Press enter when done. ");
209− return JSON.stringify({ action: "accept" });
208+ return null;
210209 }
211210
212211 const line = yield* promptLine(`${parsed.message} [y/N] `);
@@ -227,6 +226,33 @@ const promptInteraction = (interaction: ExecutionInteraction) => · 14 unmodified lines
227226 });
228227 });
229228
229+const waitForExecutionProgress = (input: {
230+ client: ControlPlaneClient;
231+ workspaceId: ExecutionEnvelope["execution"]["workspaceId"];
232+ executionId: ExecutionEnvelope["execution"]["id"];
233+ pendingInteractionId: ExecutionInteraction["id"];
234+}) =>
235+ Effect.gen(function* () {
236+ while (true) {
237+ yield* sleep(SERVER_POLL_INTERVAL_MS);
238+
239+ const next = yield* input.client.executions.get({
240+ path: {
241+ workspaceId: input.workspaceId,
242+ executionId: input.executionId,
243+ },
244+ });
245+
246+ if (
247+ next.execution.status !== "waiting_for_interaction"
248+ || next.pendingInteraction === null
249+ || next.pendingInteraction.id !== input.pendingInteractionId
250+ ) {
251+ return next;
252+ }
253+ }
254+ });
255+
230256 const printExecution = (envelope: ExecutionEnvelope) =>
231257 Effect.sync(() => {
232258 const execution = envelope.execution;
@@ -314,6 +340,25 @@ const driveExecution = (input: { · 81 unmodified lines
314340 return current;
315341 }
316342
343+ const parsed = parseInteractionPayload(pending);
344+ if (parsed?.mode === "url") {
345+ yield* Effect.sync(() => {
346+ process.stdout.write(`${parsed.message}\n${parsed.url ?? ""}\n`);
347+ });
348+
349+ if (!process.stdin.isTTY || !process.stdout.isTTY) {
350+ return current;
351+ }
352+
353+ current = yield* waitForExecutionProgress({
354+ client: input.client,
355+ workspaceId: input.workspaceId,
356+ executionId: current.execution.id,
357+ pendingInteractionId: pending.id,
358+ });
359+ continue;
360+ }
361+
317362 const responseJson = yield* promptInteraction(pending);
318363 if (responseJson === null) {
319364 yield* Effect.sync(() => {
apps/executor/src/server/index.ts+6−1
@@ -80,11 +80,14 @@ export const makeLocalExecutorServer = ( · 79 unmodified lines
8080 });
8181 }
8282
83+ let baseUrlRef: string | undefined;
84+
8385 const runtime = yield* Effect.acquireRelease(
8486 makeSqlControlPlaneRuntime({
8587 localDataDir,
8688 executionResolver: options.executionResolver,
8789 resolveSecretMaterial: options.resolveSecretMaterial,
90+ getLocalServerBaseUrl: () => baseUrlRef,
8891 }).pipe(
8992 Effect.mapError((cause) =>
9093 cause instanceof Error ? cause : new Error(String(cause)),
@@ -110,11 +113,13 @@ export const makeLocalExecutorServer = ( · 19 unmodified lines
110113 return yield* Effect.fail(new Error("Local executor server did not bind to a TCP address"));
111114 }
112115
116+ baseUrlRef = `http://${server.address.hostname}:${server.address.port}`;
117+
113118 return {
114119 runtime,
115120 host: server.address.hostname,
116121 port: server.address.port,
117− baseUrl: `http://${server.address.hostname}:${server.address.port}`,
122+ baseUrl: baseUrlRef,
118123 } satisfies LocalExecutorServer;
119124 });
120125
apps/executor/src/server/server.test.ts+388−0
@@ -1,4 +1,5 @@
11 import { createServer, type IncomingMessage, type ServerResponse } from "node:http";
2+import { randomUUID } from "node:crypto";
23
34 import {
45 HttpApi,
@@ -16,9 +17,14 @@ import { makeToolInvokerFromTools } from "@executor-v3/codemode-core"; · 11 unmodified lines
1617 import {
1718 ControlPlaneStorageError,
1819 makeControlPlaneClient,
20+ type ControlPlaneClient,
1921 type ResolveExecutionEnvironment,
2022 } from "@executor-v3/control-plane";
2123 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";
2228
2329 import {
2430 seedDemoMcpSourceInWorkspace,
@@ -122,6 +128,304 @@ const startOpenApiDemoServer = async () => { · 97 unmodified lines
122128 };
123129 };
124130
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+
125429 describe("local-executor-server", () => {
126430 it.scoped("serves the control-plane API and executes code", () =>
127431 Effect.gen(function* () {
@@ -296,6 +600,7 @@ describe("local-executor-server", () => { · 168 unmodified lines
296600 expect(resumed.execution.resultJson).toContain(`approved:${value}`);
297601 }
298602 }),
603+ 15_000,
299604 );
300605
301606 it.scoped("updates an existing demo MCP source instead of creating a duplicate", () =>
@@ -413,6 +718,89 @@ describe("local-executor-server", () => { · 111 unmodified lines
413718 }),
414719 );
415720
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+
416804 it.scoped("returns a typed storage error when a configured MCP endpoint is invalid", () =>
417805 Effect.gen(function* () {
418806 const server = yield* makeLocalExecutorServer({