refactor(core): share patch write path (#45588)

Share the identical file write, error mapping, and result recording path used by additions and non-moving updates. Keep deletion and moving-update behavior explicit.

packages/core/src/mcp/client.ts

@@ -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) })

packages/core/src/tool/plugin/patch.ts

@@ -215,17 +215,6 @@ export const Plugin = {
                 prepared,
                 (change) =>
                   Effect.gen(function* () {
-                    if (change.type === "add") {
-                      yield* environment.files
-                        .write(change.target.absolute, new TextEncoder().encode(change.content))
-                        .pipe(Effect.mapError((error) => fail(`Failed to write ${change.target.resource}`, error)))
-                      applied.push({
-                        type: change.type,
-                        resource: change.target.resource,
-                        target: change.target.absolute,
-                      })
-                      return
-                    }
                     if (change.type === "delete") {
                       yield* environment.files
                         .remove(change.target.absolute)
@@ -237,7 +226,7 @@ export const Plugin = {
                       })
                       return
                     }
-                    if (change.moveTarget) {
+                    if (change.type === "update" && change.moveTarget) {
                       const moveTarget = change.moveTarget
                       yield* environment.files
                         .write(moveTarget.absolute, new TextEncoder().encode(change.content))

packages/core/test/mcp.test.ts

@@ -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(