fix(mcp): retry initial 404 without injected codemode (#45563)

Co-authored-by: rekram1-node <rekram1-node@users.noreply.github.com>
2026-08-27 22:10:11 -04:00 · 2026-08-27 11:51:34 -05:00
parent 497a24c17d
commit 95c3c3f962
2 changed files with 139 additions and 33 deletions
@@ -3,7 +3,7 @@ export * as MCPClient from "./client.js"
 import path from "node:path"
 import { pathToFileURL } from "node:url"
 import { Client } from "@modelcontextprotocol/sdk/client/index.js"
-import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js"
+import { StreamableHTTPClientTransport, StreamableHTTPError } from "@modelcontextprotocol/sdk/client/streamableHttp.js"
 import type { Transport } from "@modelcontextprotocol/sdk/shared/transport.js"
 import { UnauthorizedError, type OAuthClientProvider } from "@modelcontextprotocol/sdk/client/auth.js"
 import {
@@ -191,10 +191,38 @@ export const connect = Effect.fnUntraced(function* (
  elicitation?: ElicitationHandler,
  clientInfo: Implementation = { name: "opencode", version: "unknown" },
 ) {
-  const transport: Transport = yield* Effect.gen(function* () {
+  const initialize = Effect.fnUntraced(function* (transport: Transport) {
+    const client = new Client(clientInfo, {
+      capabilities: {
+        ...(elicitation ? { elicitation: { form: { applyDefaults: true }, url: {} } } : {}),
+        // https://github.com/anomalyco/opencode/issues/2308
+        roots: {},
+      },
+    })
+    client.setRequestHandler(ListRootsRequestSchema, () =>
+      Promise.resolve({ roots: [{ uri: pathToFileURL(directory).href }] }),
+    )
+    if (elicitation) {
+      client.setRequestHandler(ElicitRequestSchema, (request, extra) =>
+        Effect.runPromise(elicitation.create({ server, params: request.params, signal: extra.signal })),
+      )
+      client.setNotificationHandler(ElicitationCompleteNotificationSchema, (notification) =>
+        Effect.runPromise(elicitation.complete({ server, elicitationID: notification.params.elicitationId })),
+      )
+    }
+
+    yield* Effect.tryPromise({
+      try: (signal) =>
+        client.connect(transport, { timeout: config.timeout?.startup ?? DEFAULT_STARTUP_TIMEOUT, signal }),
+      catch: (error) => error,
+    }).pipe(Effect.onError(() => Effect.promise(() => transport.close()).pipe(Effect.ignore)))
+    return client
+  })
+
+  const exit = yield* Effect.gen(function* () {
    if (config.type === "local") {
      const [command, ...args] = config.command
-      return yield* MCPStdio.make({
+      const transport = yield* MCPStdio.make({
        server,
        command,
        args,
@@ -204,41 +232,32 @@ export const connect = Effect.fnUntraced(function* (
          ...config.environment,
        },
      })
+      return yield* initialize(transport)
    }
    if (!URL.canParse(config.url))
      return yield* new ConnectError({ server, message: `Invalid MCP URL for "${server}"` })
    // Prefer raw tools for our Code Mode without changing the configured URL used for OAuth identity.
    const url = new URL(config.url)
-    if (config.codemode !== false && !url.searchParams.has("codemode")) url.searchParams.set("codemode", "false")
-    return new StreamableHTTPClientTransport(url, {
-      requestInit: config.headers ? { headers: config.headers } : undefined,
-      authProvider,
-    })
-  })
-  const client = new Client(clientInfo, {
-    capabilities: {
-      ...(elicitation ? { elicitation: { form: { applyDefaults: true }, url: {} } } : {}),
-      // https://github.com/anomalyco/opencode/issues/2308
-      roots: {},
-    },
-  })
-  client.setRequestHandler(ListRootsRequestSchema, () =>
-    Promise.resolve({ roots: [{ uri: pathToFileURL(directory).href }] }),
-  )
-  if (elicitation) {
-    client.setRequestHandler(ElicitRequestSchema, (request, extra) =>
-      Effect.runPromise(elicitation.create({ server, params: request.params, signal: extra.signal })),
-    )
-    client.setNotificationHandler(ElicitationCompleteNotificationSchema, (notification) =>
-      Effect.runPromise(elicitation.complete({ server, elicitationID: notification.params.elicitationId })),
-    )
-  }
+    const addedCodemode = config.codemode !== false && !url.searchParams.has("codemode")
+    if (addedCodemode) url.searchParams.set("codemode", "false")
+    const open = (url: URL) =>
+      initialize(
+        new StreamableHTTPClientTransport(url, {
+          requestInit: config.headers ? { headers: config.headers } : undefined,
+          authProvider,
+        }),
+      )

-  const exit = yield* Effect.tryPromise({
-    try: (signal) => client.connect(transport, { timeout: config.timeout?.startup ?? DEFAULT_STARTUP_TIMEOUT, signal }),
-    catch: (error) => error,
+    return yield* open(url).pipe(
+      Effect.catch((error) => {
+        if (!addedCodemode || !(error instanceof StreamableHTTPError) || error.code !== 404) return Effect.fail(error)
+        // Some servers reject unknown query params. Retry once with the user's original URL.
+        return open(new URL(config.url))
+      }),
+    )
  }).pipe(Effect.exit)
  if (Exit.isSuccess(exit)) {
+    const client = exit.value
    // Closing the client closes the transport, which ends stdin and then kills through the spawner
    // handle if the server does not exit cleanly. The process scope remains a final backstop.
    yield* Effect.addFinalizer(() => Effect.promise(() => client.close()).pipe(Effect.ignore))
@@ -440,7 +459,6 @@ export const connect = Effect.fnUntraced(function* (
    } satisfies Connection
  }

-  yield* Effect.promise(() => transport.close()).pipe(Effect.ignore)
  const error = Cause.squash(exit.cause)
  if (error instanceof UnauthorizedError) return yield* new NeedsAuthError({ server })
  return yield* new ConnectError({ server, message: error instanceof Error ? error.message : String(error) })
@@ -61,7 +61,13 @@ type ResourceTemplatePage = {
 }

 function resourceServer(
-  input: { resources?: boolean; listChanged?: boolean; emptyElicitation?: boolean; urlElicitation?: boolean } = {},
+  input: {
+    resources?: boolean
+    listChanged?: boolean
+    emptyElicitation?: boolean
+    urlElicitation?: boolean
+    respond?: (request: Request) => Response | undefined
+  } = {},
 ) {
  return Effect.acquireRelease(
    Effect.promise(async () => {
@@ -152,7 +158,7 @@ function resourceServer(
          if (typeof body === "object" && body !== null && "method" in body && body.method === "initialize") {
            state.initializations += 1
          }
-          return transport.handleRequest(request)
+          return input.respond?.(request) ?? transport.handleRequest(request)
        },
      })
      return {
@@ -755,6 +761,88 @@ for (const entry of [
  )
 }

+for (const query of ["", "?source=hello%20world&tag=a&tag=b"]) {
+  testEffect(Layer.empty).live(`retries an MCP initialization 404 with the original URL: ${query || "no query"}`, () =>
+    Effect.gen(function* () {
+      const headers: Array<string | null> = []
+      const server = yield* resourceServer({
+        respond: (request) => {
+          headers.push(request.headers.get("x-mcp-test"))
+          return new URL(request.url).searchParams.has("codemode") ? new Response(null, { status: 404 }) : undefined
+        },
+      })
+      const config = new ConfigMCP.Remote({
+        type: "remote",
+        url: server.url + query,
+        headers: { "x-mcp-test": "preserved" },
+        oauth: false,
+      })
+      const connection = yield* connect("resources", config, import.meta.dir)
+      yield* connection.tools()
+      yield* connection.resources()
+
+      expect(server.state.initializations).toBe(2)
+      expect(new URL(server.state.urls[0]).searchParams.get("codemode")).toBe("false")
+      expect(new Set(server.state.urls.slice(1))).toEqual(new Set([config.url]))
+      expect(new Set(headers)).toEqual(new Set(["preserved"]))
+      expect(server.state.toolLists).toBe(1)
+      expect(server.state.resourceLists).toBe(1)
+      expect(config.url).toBe(server.url + query)
+    }),
+  )
+}
+
+for (const entry of [
+  { name: "second 404", status: 404, query: "", codemode: undefined, attempts: 2 },
+  { name: "400", status: 400, query: "", codemode: undefined, attempts: 1 },
+  { name: "401", status: 401, query: "", codemode: undefined, attempts: 1 },
+  { name: "403", status: 403, query: "", codemode: undefined, attempts: 1 },
+  { name: "500", status: 500, query: "", codemode: undefined, attempts: 1 },
+  { name: "user codemode=true", status: 404, query: "?codemode=true", codemode: undefined, attempts: 1 },
+  { name: "user codemode=false", status: 404, query: "?codemode=false", codemode: undefined, attempts: 1 },
+  { name: "empty user codemode", status: 404, query: "?codemode=", codemode: undefined, attempts: 1 },
+  { name: "direct tools", status: 404, query: "", codemode: false, attempts: 1 },
+]) {
+  testEffect(Layer.empty).live(`does not retry MCP beyond the query fallback: ${entry.name}`, () =>
+    Effect.gen(function* () {
+      const server = yield* resourceServer({
+        respond: () => new Response(null, { status: entry.status }),
+      })
+      const config = new ConfigMCP.Remote({
+        type: "remote",
+        url: server.url + entry.query,
+        codemode: entry.codemode,
+        oauth: false,
+      })
+      const error = yield* connect("resources", config, import.meta.dir).pipe(Effect.flip)
+
+      expect(error).toBeInstanceOf(MCPClient.ConnectError)
+      expect(server.state.initializations).toBe(entry.attempts)
+      expect(server.state.urls).toHaveLength(entry.attempts)
+      if (entry.query || entry.codemode === false) expect(server.state.urls).toEqual([config.url])
+      if (entry.attempts === 2) expect(server.state.urls[1]).toBe(config.url)
+    }),
+  )
+}
+
+testEffect(Layer.empty).live("does not strip codemode for an MCP 404 after initialization", () =>
+  Effect.gen(function* () {
+    let expired = false
+    const server = yield* resourceServer({
+      respond: (request) => (expired && request.method === "POST" ? new Response(null, { status: 404 }) : undefined),
+    })
+    const config = new ConfigMCP.Remote({ type: "remote", url: server.url, oauth: false })
+    const connection = yield* connect("resources", config, import.meta.dir)
+    expired = true
+    expect(yield* connection.tools().pipe(Effect.flip)).toBeInstanceOf(Error)
+
+    // The SDK tries to recover the expired session, but must keep the same URL.
+    expect(server.state.initializations).toBe(2)
+    expect(new Set(server.state.urls)).toEqual(new Set([server.url + "?codemode=false"]))
+    expect(server.state.toolLists).toBe(0)
+  }),
+)
+
 test("lists, reads, and reports MCP resource changes", async () => {
  await Effect.runPromise(
    Effect.scoped(