Просмотр исходного кода

fix(mcp): make client creation failure-safe (#31595)

Aiden Cline 2 месяцев назад
Родитель
Сommit
91073360c6
2 измененных файлов с 65 добавлено и 25 удалено
  1. 37 25
      packages/opencode/src/mcp/index.ts
  2. 28 0
      packages/opencode/test/mcp/lifecycle.test.ts

+ 37 - 25
packages/opencode/src/mcp/index.ts

@@ -27,7 +27,7 @@ import { EventV2Bridge } from "@/event-v2-bridge"
 import { EventV2 } from "@opencode-ai/core/event"
 import { TuiEvent } from "@/server/tui-event"
 import open from "open"
-import { Effect, Exit, Layer, Option, Context, Schema, Stream } from "effect"
+import { Cause, Effect, Exit, Layer, Option, Context, Schema, Stream } from "effect"
 import { EffectBridge } from "@/effect/bridge"
 import { InstanceState } from "@/effect/instance-state"
 import { ChildProcess, ChildProcessSpawner } from "effect/unstable/process"
@@ -446,31 +446,45 @@ export const layer = Layer.effect(
       )
     })
 
-    const create = Effect.fn("MCP.create")(function* (key: string, mcp: ConfigMCPV1.Info) {
-      if (mcp.enabled === false) {
-        return DISABLED_RESULT
-      }
+    const create = Effect.fn("MCP.create")(
+      function* (key: string, mcp: ConfigMCPV1.Info) {
+        if (mcp.enabled === false) {
+          return DISABLED_RESULT
+        }
 
-      const { client: mcpClient, status } =
-        mcp.type === "remote"
-          ? yield* connectRemote(key, mcp as ConfigMCPV1.Info & { type: "remote" })
-          : yield* connectLocal(key, mcp as ConfigMCPV1.Info & { type: "local" })
+        const { client: mcpClient, status } =
+          mcp.type === "remote"
+            ? yield* connectRemote(key, mcp as ConfigMCPV1.Info & { type: "remote" })
+            : yield* connectLocal(key, mcp as ConfigMCPV1.Info & { type: "local" })
 
-      if (!mcpClient) {
-        if (status.status !== "connected" && status.status !== "disabled") {
-          yield* Effect.logWarning("server unavailable", { key, type: mcp.type, status: status.status })
+        if (!mcpClient) {
+          if (status.status !== "connected" && status.status !== "disabled") {
+            yield* Effect.logWarning("server unavailable", { key, type: mcp.type, status: status.status })
+          }
+          return { status } satisfies CreateResult
         }
-        return { status } satisfies CreateResult
-      }
 
-      const listed = mcpClient.getServerCapabilities()?.tools ? yield* defs(mcpClient, mcp.timeout) : []
-      if (!listed) {
-        yield* Effect.tryPromise(() => mcpClient.close()).pipe(Effect.ignore)
-        return { status: { status: "failed", error: "Failed to get tools" } } satisfies CreateResult
-      }
-
-      return { mcpClient, status, defs: listed } satisfies CreateResult
-    })
+        return yield* Effect.gen(function* () {
+          const listed = mcpClient.getServerCapabilities()?.tools ? yield* defs(mcpClient, mcp.timeout) : []
+          if (!listed) {
+            return yield* Effect.fail(new Error("Failed to get tools"))
+          }
+          return { mcpClient, status, defs: listed } satisfies CreateResult
+        }).pipe(
+          Effect.catchCause((cause) =>
+            Effect.tryPromise(() => mcpClient.close()).pipe(Effect.ignore, Effect.andThen(Effect.failCause(cause))),
+          ),
+        )
+      },
+      Effect.map((result): CreateResult => result),
+      Effect.catchCause((cause) => {
+        if (Cause.hasInterruptsOnly(cause)) return Effect.interrupt
+        const error = Cause.squash(cause)
+        return Effect.succeed<CreateResult>({
+          status: { status: "failed", error: error instanceof Error ? error.message : String(error) },
+        })
+      }),
+    )
     const cfgSvc = yield* Config.Service
 
     const descendants = Effect.fnUntraced(
@@ -537,9 +551,7 @@ export const layer = Layer.effect(
                 return
               }
 
-              const result = yield* create(key, mcp).pipe(Effect.catch(() => Effect.void))
-              if (!result) return
-
+              const result = yield* create(key, mcp)
               s.status[key] = result.status
               if (result.mcpClient) {
                 s.clients[key] = result.mcpClient

+ 28 - 0
packages/opencode/test/mcp/lifecycle.test.ts

@@ -8,6 +8,7 @@ import { testEffect } from "../lib/effect"
 // Per-client state for controlling mock behavior
 interface MockClientState {
   capabilities: { tools?: object; prompts?: object; resources?: object }
+  capabilitiesShouldThrow: boolean
   tools: Array<{ name: string; description?: string; inputSchema: object; outputSchema?: object }>
   listToolsCalls: number
   listPromptsCalls: number
@@ -51,6 +52,7 @@ function getOrCreateClientState(name?: string): MockClientState {
   if (!state) {
     state = {
       capabilities: { tools: {}, prompts: {}, resources: {} },
+      capabilitiesShouldThrow: false,
       tools: [{ name: "test_tool", description: "A test tool", inputSchema: { type: "object", properties: {} } }],
       listToolsCalls: 0,
       listPromptsCalls: 0,
@@ -155,6 +157,7 @@ void mock.module("@modelcontextprotocol/sdk/client/index.js", () => ({
     }
 
     getServerCapabilities() {
+      if (this._state?.capabilitiesShouldThrow) throw new Error("capability discovery failed")
       return this._state?.capabilities
     }
 
@@ -566,6 +569,31 @@ it.instance(
   },
 )
 
+it.instance(
+  "returns failed and closes the client when SDK initialization throws",
+  () =>
+    MCP.Service.use((mcp: MCPNS.Interface) =>
+      Effect.gen(function* () {
+        lastCreatedClientName = "defective-server"
+        const serverState = getOrCreateClientState("defective-server")
+        serverState.capabilitiesShouldThrow = true
+
+        const result = yield* mcp.add("defective-server", {
+          type: "local",
+          command: ["echo", "test"],
+        })
+
+        expect(statusName(result.status, "defective-server")).toBe("failed")
+        expect((yield* mcp.status())["defective-server"]).toEqual({
+          status: "failed",
+          error: "capability discovery failed",
+        })
+        expect(serverState.closed).toBe(true)
+      }),
+    ),
+  { config: { mcp: {} } },
+)
+
 it.instance(
   "falls back when MCP output schema refs fail SDK tool discovery",
   () =>