- 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.
Unified · Split
| @@ -1,6 +1,5 @@ | |||
| 1 | import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; | 1 | import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; |
| 2 | import { WebStandardStreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/webStandardStreamableHttp.js"; | 2 | import { WebStandardStreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/webStandardStreamableHttp.js"; |
| 3 | import type { HandleRequestOptions } from "@modelcontextprotocol/sdk/server/webStandardStreamableHttp.js"; | ||
| 4 | import { z } from "zod"; | 3 | import { z } from "zod"; |
| 5 | import { generateToolDeclarations, generateToolInventory, typecheckCode } from "./typechecker"; | 4 | import { generateToolDeclarations, generateToolInventory, typecheckCode } from "./typechecker"; |
| 6 | import type { LiveTaskEvent } from "./events"; | 5 | import type { LiveTaskEvent } from "./events"; |
| @@ -146,7 +145,6 @@ function waitForTerminalTask( · 139 unmodified lines | |||
| 146 | if (pending.length === 0) { | 145 | if (pending.length === 0) { |
| 147 | return; | 146 | return; |
| 148 | } | 147 | } |
| 149 | |||
| 150 | for (const approval of pending) { | 148 | for (const approval of pending) { |
| 151 | let decision: ApprovalPromptDecision | null; | 149 | let decision: ApprovalPromptDecision | null; |
| 152 | try { | 150 | try { |
| @@ -426,65 +424,13 @@ async function createMcpServer( · 273 unmodified lines | |||
| 426 | }); | 424 | }); |
| 427 | } | 425 | } |
| 428 | 426 | ||
| 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 | |||
| 481 | mcp.registerTool( | 427 | mcp.registerTool( |
| 482 | "run_code", | 428 | "run_code", |
| 483 | { | 429 | { |
| 484 | description: buildRunCodeDescription(tools), | 430 | description: buildRunCodeDescription(tools), |
| 485 | inputSchema: context ? BOUND_INPUT : FULL_INPUT, | 431 | inputSchema: context ? BOUND_INPUT : FULL_INPUT, |
| 486 | }, | 432 | }, |
| 487 | createRunCodeTool(service, context, approvalPrompt), | 433 | createRunCodeTool(service, context), |
| 488 | ); | 434 | ); |
| 489 | 435 | ||
| 490 | return mcp; | 436 | return mcp; |
| @@ -494,178 +440,22 @@ async function createMcpServer( · 3 unmodified lines | |||
| 494 | // Request handler | 440 | // Request handler |
| 495 | // --------------------------------------------------------------------------- | 441 | // --------------------------------------------------------------------------- |
| 496 | 442 | ||
| 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 | |||
| 640 | export async function handleMcpRequest( | 443 | export async function handleMcpRequest( |
| 641 | service: McpExecutorService, | 444 | service: McpExecutorService, |
| 642 | request: Request, | 445 | request: Request, |
| 643 | context?: McpWorkspaceContext, | 446 | context?: McpWorkspaceContext, |
| 644 | ): Promise<Response> { | 447 | ): Promise<Response> { |
| 645 | const requestOptions = await parseRequestBody(request); | 448 | const transport = new WebStandardStreamableHTTPServerTransport({ |
| 646 | const sessionId = request.headers.get("mcp-session-id")?.trim() ?? ""; | 449 | sessionIdGenerator: undefined, |
| 647 | 450 | enableJsonResponse: true, | |
| 648 | if (sessionId) { | 451 | }); |
| 649 | const existing = mcpSessions.get(sessionId); | 452 | const mcp = await createMcpServer(service, context); |
| 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 | } | ||
| 659 | 453 | ||
| 660 | const session = await createSessionState(service, context); | ||
| 661 | try { | 454 | try { |
| 662 | const response = await handleWithSession(session, service, request, requestOptions); | 455 | await mcp.connect(transport); |
| 663 | if (!session.transport.sessionId) { | 456 | return await transport.handleRequest(request); |
| 664 | await closeSessionState(session); | 457 | } finally { |
| 665 | } | 458 | await transport.close().catch(() => {}); |
| 666 | return response; | 459 | await mcp.close().catch(() => {}); |
| 667 | } catch (error) { | ||
| 668 | await closeSessionState(session); | ||
| 669 | throw error; | ||
| 670 | } | 460 | } |
| 671 | } | 461 | } |