fix(mcp): persist immediate oauth connections (#22376)

2 files changed, 108 insertions, 14 deletions

diff --git a/packages/opencode/src/mcp/index.ts b/packages/opencode/src/mcp/index.ts
index 39c7ed98e..3de427d7e 100644
--- a/packages/opencode/src/mcp/index.ts
+++ b/packages/opencode/src/mcp/index.ts
@@ -204,6 +204,12 @@ export namespace MCP {
     defs?: MCPToolDef[]
   }
 
+  interface AuthResult {
+    authorizationUrl: string
+    oauthState: string
+    client?: MCPClient
+  }
+
   // --- Effect Service ---
 
   interface State {
@@ -552,6 +558,21 @@ export namespace MCP {
         return Effect.tryPromise(() => client.close()).pipe(Effect.ignore)
       }
 
+      const storeClient = Effect.fnUntraced(function* (
+        s: State,
+        name: string,
+        client: MCPClient,
+        listed: MCPToolDef[],
+        timeout?: number,
+      ) {
+        yield* closeClient(s, name)
+        s.status[name] = { status: "connected" }
+        s.clients[name] = client
+        s.defs[name] = listed
+        watch(s, name, client, timeout)
+        return s.status[name]
+      })
+
       const status = Effect.fn("MCP.status")(function* () {
         const s = yield* InstanceState.get(state)
 
@@ -583,11 +604,7 @@ export namespace MCP {
           return result.status
         }
 
-        yield* closeClient(s, name)
-        s.clients[name] = result.mcpClient
-        s.defs[name] = result.defs!
-        watch(s, name, result.mcpClient, mcp.timeout)
-        return result.status
+        return yield* storeClient(s, name, result.mcpClient, result.defs!, mcp.timeout)
       })
 
       const add = Effect.fn("MCP.add")(function* (name: string, mcp: Config.Mcp) {
@@ -753,14 +770,16 @@ export namespace MCP {
         return yield* Effect.tryPromise({
           try: () => {
             const client = new Client({ name: "opencode", version: Installation.VERSION })
-            return client.connect(transport).then(() => ({ authorizationUrl: "", oauthState }))
+            return client
+              .connect(transport)
+              .then(() => ({ authorizationUrl: "", oauthState, client }) satisfies AuthResult)
           },
           catch: (error) => error,
         }).pipe(
           Effect.catch((error) => {
             if (error instanceof UnauthorizedError && capturedUrl) {
               pendingOAuthTransports.set(mcpName, transport)
-              return Effect.succeed({ authorizationUrl: capturedUrl.toString(), oauthState })
+              return Effect.succeed({ authorizationUrl: capturedUrl.toString(), oauthState } satisfies AuthResult)
             }
             return Effect.die(error)
           }),
@@ -768,14 +787,31 @@ export namespace MCP {
       })
 
       const authenticate = Effect.fn("MCP.authenticate")(function* (mcpName: string) {
-        const { authorizationUrl, oauthState } = yield* startAuth(mcpName)
-        if (!authorizationUrl) return { status: "connected" } as Status
+        const result = yield* startAuth(mcpName)
+        if (!result.authorizationUrl) {
+          const client = "client" in result ? result.client : undefined
+          const mcpConfig = yield* getMcpConfig(mcpName)
+          if (!mcpConfig) {
+            yield* Effect.tryPromise(() => client?.close() ?? Promise.resolve()).pipe(Effect.ignore)
+            return { status: "failed", error: "MCP config not found after auth" } as Status
+          }
+
+          const listed = client ? yield* defs(mcpName, client, mcpConfig.timeout) : undefined
+          if (!client || !listed) {
+            yield* Effect.tryPromise(() => client?.close() ?? Promise.resolve()).pipe(Effect.ignore)
+            return { status: "failed", error: "Failed to get tools" } as Status
+          }
+
+          const s = yield* InstanceState.get(state)
+          yield* auth.clearOAuthState(mcpName)
+          return yield* storeClient(s, mcpName, client, listed, mcpConfig.timeout)
+        }
 
-        log.info("opening browser for oauth", { mcpName, url: authorizationUrl, state: oauthState })
+        log.info("opening browser for oauth", { mcpName, url: result.authorizationUrl, state: result.oauthState })
 
-        const callbackPromise = McpOAuthCallback.waitForCallback(oauthState, mcpName)
+        const callbackPromise = McpOAuthCallback.waitForCallback(result.oauthState, mcpName)
 
-        yield* Effect.tryPromise(() => open(authorizationUrl)).pipe(
+        yield* Effect.tryPromise(() => open(result.authorizationUrl)).pipe(
           Effect.flatMap((subprocess) =>
             Effect.callback<void, Error>((resume) => {
               const timer = setTimeout(() => resume(Effect.void), 500)
@@ -793,14 +829,14 @@ export namespace MCP {
           ),
           Effect.catch(() => {
             log.warn("failed to open browser, user must open URL manually", { mcpName })
-            return bus.publish(BrowserOpenFailed, { mcpName, url: authorizationUrl }).pipe(Effect.ignore)
+            return bus.publish(BrowserOpenFailed, { mcpName, url: result.authorizationUrl }).pipe(Effect.ignore)
           }),
         )
 
         const code = yield* Effect.promise(() => callbackPromise)
 
         const storedState = yield* auth.getOAuthState(mcpName)
-        if (storedState !== oauthState) {
+        if (storedState !== result.oauthState) {
           yield* auth.clearOAuthState(mcpName)
           throw new Error("OAuth state mismatch - potential CSRF attack")
         }
diff --git a/packages/opencode/test/mcp/oauth-auto-connect.test.ts b/packages/opencode/test/mcp/oauth-auto-connect.test.ts
index 7f267e9c3..13ae0bb34 100644
--- a/packages/opencode/test/mcp/oauth-auto-connect.test.ts
+++ b/packages/opencode/test/mcp/oauth-auto-connect.test.ts
@@ -20,6 +20,7 @@ const transportCalls: Array<{
 // Controls whether the mock transport simulates a 401 that triggers the SDK
 // auth flow (which calls provider.state()) or a simple UnauthorizedError.
 let simulateAuthFlow = true
+let connectSucceedsImmediately = false
 
 // Mock the transport constructors to simulate OAuth auto-auth on 401
 mock.module("@modelcontextprotocol/sdk/client/streamableHttp.js", () => ({
@@ -40,6 +41,8 @@ mock.module("@modelcontextprotocol/sdk/client/streamableHttp.js", () => ({
       })
     }
     async start() {
+      if (connectSucceedsImmediately) return
+
       // Simulate what the real SDK transport does on 401:
       // It calls auth() which eventually calls provider.state(), then
       // provider.redirectToAuthorization(), then throws UnauthorizedError.
@@ -85,6 +88,14 @@ mock.module("@modelcontextprotocol/sdk/client/index.js", () => ({
     async connect(transport: { start: () => Promise<void> }) {
       await transport.start()
     }
+
+    setNotificationHandler() {}
+
+    async listTools() {
+      return { tools: [{ name: "test_tool", inputSchema: { type: "object", properties: {} } }] }
+    }
+
+    async close() {}
   },
 }))
 
@@ -96,6 +107,7 @@ mock.module("@modelcontextprotocol/sdk/client/auth.js", () => ({
 beforeEach(() => {
   transportCalls.length = 0
   simulateAuthFlow = true
+  connectSucceedsImmediately = false
 })
 
 // Import modules after mocking
@@ -222,3 +234,49 @@ test("state() returns existing state when one is saved", async () => {
     },
   })
 })
+
+test("authenticate() stores a connected client when auth completes without redirect", async () => {
+  await using tmp = await tmpdir({
+    init: async (dir) => {
+      await Bun.write(
+        `${dir}/opencode.json`,
+        JSON.stringify({
+          $schema: "https://opencode.ai/config.json",
+          mcp: {
+            "test-oauth-connect": {
+              type: "remote",
+              url: "https://example.com/mcp",
+            },
+          },
+        }),
+      )
+    },
+  })
+
+  await Instance.provide({
+    directory: tmp.path,
+    fn: async () => {
+      await Effect.runPromise(
+        MCP.Service.use((mcp) =>
+          Effect.gen(function* () {
+            const added = yield* mcp.add("test-oauth-connect", {
+              type: "remote",
+              url: "https://example.com/mcp",
+            })
+            const before = added.status as Record<string, { status: string; error?: string }>
+            expect(before["test-oauth-connect"]?.status).toBe("needs_auth")
+
+            simulateAuthFlow = false
+            connectSucceedsImmediately = true
+
+            const result = yield* mcp.authenticate("test-oauth-connect")
+            expect(result.status).toBe("connected")
+
+            const after = yield* mcp.status()
+            expect(after["test-oauth-connect"]?.status).toBe("connected")
+          }),
+        ).pipe(Effect.provide(MCP.defaultLayer)),
+      )
+    },
+  })
+})