History

Refactor MCP server by removing approval prompt logic and simplifying tool registration

- Eliminated the approval prompt function to streamline tool invocation. - Updated the `createMcpServer` function to register tools without the approval prompt. - Removed unused interfaces and session management code to enhance clarity and maintainability.

Rhys Sullivan committed 9 months ago

Commit c7a02461 parent cd860c91 file · +11 −221Browse files

Unified · Split

executor/convex/lib/mcp_server.ts+11−221
@@ -1,6 +1,5 @@
11 import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
22 import { WebStandardStreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/webStandardStreamableHttp.js";
3−import type { HandleRequestOptions } from "@modelcontextprotocol/sdk/server/webStandardStreamableHttp.js";
43 import { z } from "zod";
54 import { generateToolDeclarations, generateToolInventory, typecheckCode } from "./typechecker";
65 import type { LiveTaskEvent } from "./events";
@@ -146,7 +145,6 @@ function waitForTerminalTask( · 139 unmodified lines
146145 if (pending.length === 0) {
147146 return;
148147 }
149−
150148 for (const approval of pending) {
151149 let decision: ApprovalPromptDecision | null;
152150 try {
@@ -426,65 +424,13 @@ async function createMcpServer( · 273 unmodified lines
426424 });
427425 }
428426
429− const approvalPrompt: ApprovalPrompt = async (approval) => {
430− const inputPreview = (() => {
431− try {
432− const text = JSON.stringify(approval.input ?? {});
433− return text.length > 500 ? `${text.slice(0, 500)}...` : text;
434− } catch {
435− return "{}";
436− }
437− })();
438−
439− const result = await mcp.elicitInput({
440− mode: "form",
441− message: `Approval required for tool call '${approval.toolPath}'. Input: ${inputPreview}`,
442− requestedSchema: {
443− type: "object",
444− properties: {
445− decision: {
446− type: "string",
447− oneOf: [
448− { const: "approve", title: "Approve" },
449− { const: "deny", title: "Deny" },
450− ],
451− description: "Choose whether to allow this tool invocation",
452− default: "deny",
453− },
454− reason: {
455− type: "string",
456− description: "Optional reason",
457− },
458− },
459− required: ["decision"],
460− },
461− });
462−
463− if (result.action !== "accept") {
464− return {
465− decision: "denied",
466− reason: result.action === "cancel" ? "Cancelled via MCP elicitation" : "Declined via MCP elicitation",
467− } satisfies ApprovalPromptDecision;
468− }
469−
470− const content = result.content && typeof result.content === "object"
471− ? (result.content as Record<string, unknown>)
472− : {};
473− const decision = content.decision === "approve" ? "approved" : "denied";
474− const reason = typeof content.reason === "string" && content.reason.trim().length > 0
475− ? content.reason.trim()
476− : undefined;
477−
478− return { decision, reason } satisfies ApprovalPromptDecision;
479− };
480−
481427 mcp.registerTool(
482428 "run_code",
483429 {
484430 description: buildRunCodeDescription(tools),
485431 inputSchema: context ? BOUND_INPUT : FULL_INPUT,
486432 },
487− createRunCodeTool(service, context, approvalPrompt),
433+ createRunCodeTool(service, context),
488434 );
489435
490436 return mcp;
@@ -494,178 +440,22 @@ async function createMcpServer( · 3 unmodified lines
494440 // Request handler
495441 // ---------------------------------------------------------------------------
496442
497−interface SessionServiceRef {
498− current: McpExecutorService;
499−}
500−
501−interface McpSessionState {
502− mcp: McpServer;
503− transport: WebStandardStreamableHTTPServerTransport;
504− serviceRef: SessionServiceRef;
505− queue: Promise<void>;
506− lastSeenAt: number;
507−}
508−
509−const mcpSessions = new Map<string, McpSessionState>();
510−
511−function jsonRpcErrorResponse(status: number, code: number, message: string): Response {
512− return Response.json(
513− {
514− jsonrpc: "2.0",
515− error: { code, message },
516− id: null,
517− },
518− { status },
519− );
520−}
521−
522−function createServiceProxy(serviceRef: SessionServiceRef): McpExecutorService {
523− const listPendingApprovals = serviceRef.current.listPendingApprovals
524− ? async (workspaceId: string) => {
525− const fn = serviceRef.current.listPendingApprovals;
526− if (!fn) return [];
527− return await fn(workspaceId);
528− }
529− : undefined;
530−
531− const resolveApproval = serviceRef.current.resolveApproval
532− ? async (input: {
533− workspaceId: string;
534− approvalId: string;
535− decision: "approved" | "denied";
536− reviewerId?: string;
537− reason?: string;
538− }) => {
539− const fn = serviceRef.current.resolveApproval;
540− if (!fn) {
541− throw new Error("resolveApproval is not available");
542− }
543− return await fn(input);
544− }
545− : undefined;
546−
547− return {
548− createTask: async (input) => await serviceRef.current.createTask(input),
549− getTask: async (taskId, workspaceId) => await serviceRef.current.getTask(taskId, workspaceId),
550− subscribe: (taskId, listener) => serviceRef.current.subscribe(taskId, listener),
551− bootstrapAnonymousContext: async (sessionId) => await serviceRef.current.bootstrapAnonymousContext(sessionId),
552− listTools: async (context) => await serviceRef.current.listTools(context),
553− listPendingApprovals,
554− resolveApproval,
555− };
556−}
557−
558−async function createSessionState(
559− service: McpExecutorService,
560− context?: McpWorkspaceContext,
561−): Promise<McpSessionState> {
562− const serviceRef: SessionServiceRef = { current: service };
563− const proxyService = createServiceProxy(serviceRef);
564− let state: McpSessionState;
565−
566− const transport = new WebStandardStreamableHTTPServerTransport({
567− sessionIdGenerator: () => crypto.randomUUID(),
568− onsessioninitialized: (sessionId) => {
569− mcpSessions.set(sessionId, state);
570− },
571− onsessionclosed: async (sessionId) => {
572− const existing = mcpSessions.get(sessionId);
573− if (existing === state) {
574− mcpSessions.delete(sessionId);
575− }
576− await state.mcp.close().catch(() => {});
577− },
578− });
579−
580− const mcp = await createMcpServer(proxyService, context);
581− state = {
582− mcp,
583− transport,
584− serviceRef,
585− queue: Promise.resolve(),
586− lastSeenAt: Date.now(),
587− };
588−
589− await mcp.connect(transport);
590− return state;
591−}
592−
593−async function closeSessionState(state: McpSessionState): Promise<void> {
594− await state.transport.close().catch(() => {});
595− await state.mcp.close().catch(() => {});
596−}
597−
598−async function handleWithSession(
599− state: McpSessionState,
600− service: McpExecutorService,
601− request: Request,
602− requestOptions?: HandleRequestOptions,
603−): Promise<Response> {
604− const run = async () => {
605− state.serviceRef.current = service;
606− state.lastSeenAt = Date.now();
607− return await state.transport.handleRequest(request, requestOptions);
608− };
609−
610− const next = state.queue.then(run, run);
611− state.queue = next.then(() => undefined, () => undefined);
612− return await next;
613−}
614−
615−async function parseRequestBody(request: Request): Promise<HandleRequestOptions | undefined> {
616− if (request.method !== "POST") return undefined;
617− try {
618− const parsedBody = await request.clone().json();
619− const messages = Array.isArray(parsedBody) ? parsedBody : [parsedBody];
620− for (const message of messages) {
621− if (!message || typeof message !== "object") continue;
622− const method = (message as { method?: unknown }).method;
623− if (method !== "initialize") continue;
624− const params = (message as { params?: unknown }).params;
625− const capabilities = params && typeof params === "object"
626− ? (params as { capabilities?: unknown }).capabilities
627− : undefined;
628− console.warn(
629− `[executor] MCP initialize client capabilities: ${JSON.stringify(capabilities ?? {})}`,
630− );
631− break;
632− }
633−
634− return { parsedBody };
635− } catch {
636− return undefined;
637− }
638−}
639−
640443 export async function handleMcpRequest(
641444 service: McpExecutorService,
642445 request: Request,
643446 context?: McpWorkspaceContext,
644447 ): Promise<Response> {
645− const requestOptions = await parseRequestBody(request);
646− const sessionId = request.headers.get("mcp-session-id")?.trim() ?? "";
647−
648− if (sessionId) {
649− const existing = mcpSessions.get(sessionId);
650− if (!existing) {
651− return jsonRpcErrorResponse(404, -32001, "Session not found");
652− }
653− return await handleWithSession(existing, service, request, requestOptions);
654− }
655−
656− if (request.method !== "POST") {
657− return jsonRpcErrorResponse(400, -32000, "Bad Request: Mcp-Session-Id header is required");
658− }
448+ const transport = new WebStandardStreamableHTTPServerTransport({
449+ sessionIdGenerator: undefined,
450+ enableJsonResponse: true,
451+ });
452+ const mcp = await createMcpServer(service, context);
659453
660− const session = await createSessionState(service, context);
661454 try {
662− const response = await handleWithSession(session, service, request, requestOptions);
663− if (!session.transport.sessionId) {
664− await closeSessionState(session);
665− }
666− return response;
667− } catch (error) {
668− await closeSessionState(session);
669− throw error;
455+ await mcp.connect(transport);
456+ return await transport.handleRequest(request);
457+ } finally {
458+ await transport.close().catch(() => {});
459+ await mcp.close().catch(() => {});
670460 }
671461 }