fix(mcp): Fix memory leaks in OAuth transport and process cleanup

packages/opencode/src/index.ts

@@ -26,6 +26,33 @@ import { EOL } from "os"
 import { WebCommand } from "./cli/cmd/web"
 import { PrCommand } from "./cli/cmd/pr"
 import { SessionCommand } from "./cli/cmd/session"
+import { Instance } from "./project/instance"
+
+// Track whether cleanup has been performed to avoid duplicate cleanup
+let cleanupPerformed = false
+
+// Cleanup handler to dispose all instances (closes MCP clients, LSP servers, etc.)
+async function performCleanup(signal: string) {
+  if (cleanupPerformed) return
+  cleanupPerformed = true
+
+  Log.Default.info("cleanup triggered by signal", { signal })
+  try {
+    await Instance.disposeAll()
+  } catch (error) {
+    Log.Default.error("error during cleanup", { error })
+  }
+}
+
+// Register signal handlers for cleanup
+const signals = ["SIGTERM", "SIGINT"] as const
+for (const signal of signals) {
+  process.on(signal, () => {
+    performCleanup(signal).finally(() => {
+      process.exit(128 + (signal === "SIGINT" ? 2 : 15))
+    })
+  })
+}
 
 process.on("unhandledRejection", (e) => {
   Log.Default.error("rejection", {
@@ -39,121 +66,133 @@ process.on("uncaughtException", (e) => {
   })
 })
 
-const cli = yargs(hideBin(process.argv))
-  .parserConfiguration({ "populate--": true })
-  .scriptName("opencode")
-  .wrap(100)
-  .help("help", "show help")
-  .alias("help", "h")
-  .version("version", "show version number", Installation.VERSION)
-  .alias("version", "v")
-  .option("print-logs", {
-    describe: "print logs to stderr",
-    type: "boolean",
-  })
-  .option("log-level", {
-    describe: "log level",
-    type: "string",
-    choices: ["DEBUG", "INFO", "WARN", "ERROR"],
-  })
-  .middleware(async (opts) => {
-    await Log.init({
-      print: process.argv.includes("--print-logs"),
-      dev: Installation.isLocal(),
-      level: (() => {
-        if (opts.logLevel) return opts.logLevel as Log.Level
-        if (Installation.isLocal()) return "DEBUG"
-        return "INFO"
-      })(),
+// Main function that wraps the CLI with proper cleanup
+async function main() {
+  const cli = yargs(hideBin(process.argv))
+    .parserConfiguration({ "populate--": true })
+    .scriptName("opencode")
+    .wrap(100)
+    .help("help", "show help")
+    .alias("help", "h")
+    .version("version", "show version number", Installation.VERSION)
+    .alias("version", "v")
+    .option("print-logs", {
+      describe: "print logs to stderr",
+      type: "boolean",
+    })
+    .option("log-level", {
+      describe: "log level",
+      type: "string",
+      choices: ["DEBUG", "INFO", "WARN", "ERROR"],
     })
+    .middleware(async (opts) => {
+      await Log.init({
+        print: process.argv.includes("--print-logs"),
+        dev: Installation.isLocal(),
+        level: (() => {
+          if (opts.logLevel) return opts.logLevel as Log.Level
+          if (Installation.isLocal()) return "DEBUG"
+          return "INFO"
+        })(),
+      })
 
-    process.env.AGENT = "1"
-    process.env.OPENCODE = "1"
+      process.env.AGENT = "1"
+      process.env.OPENCODE = "1"
 
-    Log.Default.info("opencode", {
-      version: Installation.VERSION,
-      args: process.argv.slice(2),
+      Log.Default.info("opencode", {
+        version: Installation.VERSION,
+        args: process.argv.slice(2),
+      })
     })
-  })
-  .usage("\n" + UI.logo())
-  .completion("completion", "generate shell completion script")
-  .command(AcpCommand)
-  .command(McpCommand)
-  .command(TuiThreadCommand)
-  .command(AttachCommand)
-  .command(RunCommand)
-  .command(GenerateCommand)
-  .command(DebugCommand)
-  .command(AuthCommand)
-  .command(AgentCommand)
-  .command(UpgradeCommand)
-  .command(UninstallCommand)
-  .command(ServeCommand)
-  .command(WebCommand)
-  .command(ModelsCommand)
-  .command(StatsCommand)
-  .command(ExportCommand)
-  .command(ImportCommand)
-  .command(GithubCommand)
-  .command(PrCommand)
-  .command(SessionCommand)
-  .fail((msg, err) => {
-    if (
-      msg?.startsWith("Unknown argument") ||
-      msg?.startsWith("Not enough non-option arguments") ||
-      msg?.startsWith("Invalid values:")
-    ) {
+    .usage("\n" + UI.logo())
+    .completion("completion", "generate shell completion script")
+    .command(AcpCommand)
+    .command(McpCommand)
+    .command(TuiThreadCommand)
+    .command(AttachCommand)
+    .command(RunCommand)
+    .command(GenerateCommand)
+    .command(DebugCommand)
+    .command(AuthCommand)
+    .command(AgentCommand)
+    .command(UpgradeCommand)
+    .command(UninstallCommand)
+    .command(ServeCommand)
+    .command(WebCommand)
+    .command(ModelsCommand)
+    .command(StatsCommand)
+    .command(ExportCommand)
+    .command(ImportCommand)
+    .command(GithubCommand)
+    .command(PrCommand)
+    .command(SessionCommand)
+    .fail((msg, err) => {
+      if (
+        msg?.startsWith("Unknown argument") ||
+        msg?.startsWith("Not enough non-option arguments") ||
+        msg?.startsWith("Invalid values:")
+      ) {
+        if (err) throw err
+        cli.showHelp("log")
+      }
       if (err) throw err
-      cli.showHelp("log")
-    }
-    if (err) throw err
-    process.exit(1)
-  })
-  .strict()
-
-try {
-  await cli.parse()
-} catch (e) {
-  let data: Record<string, any> = {}
-  if (e instanceof NamedError) {
-    const obj = e.toObject()
-    Object.assign(data, {
-      ...obj.data,
+      process.exit(1)
     })
-  }
+    .strict()
 
-  if (e instanceof Error) {
-    Object.assign(data, {
-      name: e.name,
-      message: e.message,
-      cause: e.cause?.toString(),
-      stack: e.stack,
-    })
-  }
+  try {
+    await cli.parse()
+  } catch (e) {
+    let data: Record<string, any> = {}
+    if (e instanceof NamedError) {
+      const obj = e.toObject()
+      Object.assign(data, {
+        ...obj.data,
+      })
+    }
 
-  if (e instanceof ResolveMessage) {
-    Object.assign(data, {
-      name: e.name,
-      message: e.message,
-      code: e.code,
-      specifier: e.specifier,
-      referrer: e.referrer,
-      position: e.position,
-      importKind: e.importKind,
-    })
-  }
-  Log.Default.error("fatal", data)
-  const formatted = FormatError(e)
-  if (formatted) UI.error(formatted)
-  if (formatted === undefined) {
-    UI.error("Unexpected error, check log file at " + Log.file() + " for more details" + EOL)
-    console.error(e instanceof Error ? e.message : String(e))
+    if (e instanceof Error) {
+      Object.assign(data, {
+        name: e.name,
+        message: e.message,
+        cause: e.cause?.toString(),
+        stack: e.stack,
+      })
+    }
+
+    if (e instanceof ResolveMessage) {
+      Object.assign(data, {
+        name: e.name,
+        message: e.message,
+        code: e.code,
+        specifier: e.specifier,
+        referrer: e.referrer,
+        position: e.position,
+        importKind: e.importKind,
+      })
+    }
+    Log.Default.error("fatal", data)
+    const formatted = FormatError(e)
+    if (formatted) UI.error(formatted)
+    if (formatted === undefined) {
+      UI.error("Unexpected error, check log file at " + Log.file() + " for more details" + EOL)
+      console.error(e instanceof Error ? e.message : String(e))
+    }
+    process.exitCode = 1
   }
-  process.exitCode = 1
-} finally {
-  // Some subprocesses don't react properly to SIGTERM and similar signals.
-  // Most notably, some docker-container-based MCP servers don't handle such signals unless
-  // run using `docker run --init`.
-  // Explicitly exit to avoid any hanging subprocesses.
-  process.exit()
 }
+
+// Run main with cleanup
+main()
+  .then(async () => {
+    await performCleanup("main")
+    // Some subprocesses don't react properly to SIGTERM and similar signals.
+    // Most notably, some docker-container-based MCP servers don't handle such signals unless
+    // run using `docker run --init`.
+    // Explicitly exit to avoid any hanging subprocesses.
+    process.exit(process.exitCode ?? 0)
+  })
+  .catch(async (error) => {
+    await performCleanup("error")
+    process.exit(1)
+  })

packages/opencode/src/mcp/index.ts

@@ -151,6 +151,28 @@ export namespace MCP {
   type TransportWithAuth = StreamableHTTPClientTransport | SSEClientTransport
   const pendingOAuthTransports = new Map<string, TransportWithAuth>()
 
+  // Helper function to close a transport properly
+  async function closeTransport(transport: TransportWithAuth | undefined): Promise<void> {
+    if (!transport) return
+    try {
+      // The transport should have a close method or similar cleanup
+      if (typeof (transport as any).close === "function") {
+        await (transport as any).close()
+      }
+    } catch (error) {
+      log.error("Failed to close transport", { error })
+    }
+  }
+
+  // Helper function to set a transport while properly cleaning up the old one
+  async function setPendingOAuthTransport(key: string, transport: TransportWithAuth): Promise<void> {
+    const existing = pendingOAuthTransports.get(key)
+    if (existing) {
+      await closeTransport(existing)
+    }
+    pendingOAuthTransports.set(key, transport)
+  }
+
   // Prompt cache types
   type PromptInfo = Awaited<ReturnType<MCPClient["listPrompts"]>>["prompts"][number]
 
@@ -378,7 +400,7 @@ export namespace MCP {
               }).catch((e) => log.debug("failed to show toast", { error: e }))
             } else {
               // Store transport for later finishAuth call
-              pendingOAuthTransports.set(key, transport)
+              await setPendingOAuthTransport(key, transport)
               status = { status: "needs_auth" as const }
               // Show toast for needs_auth
               Bus.publish(TuiEvent.ToastShow, {
@@ -766,7 +788,7 @@ export namespace MCP {
     } catch (error) {
       if (error instanceof UnauthorizedError && capturedUrl) {
         // Store transport for finishAuth
-        pendingOAuthTransports.set(mcpName, transport)
+        await setPendingOAuthTransport(mcpName, transport)
         return { authorizationUrl: capturedUrl.toString() }
       }
       throw error
@@ -888,6 +910,9 @@ export namespace MCP {
   export async function removeAuth(mcpName: string): Promise<void> {
     await McpAuth.remove(mcpName)
     McpOAuthCallback.cancelPending(mcpName)
+    // Properly close the transport before removing it
+    const transport = pendingOAuthTransports.get(mcpName)
+    await closeTransport(transport)
     pendingOAuthTransports.delete(mcpName)
     await McpAuth.clearOAuthState(mcpName)
     log.info("removed oauth credentials", { mcpName })