Merge pull request #14809 from kr1shna-exe/fix/14339-mcp-interrupted-startup

fix(cli): clean up MCP servers when their startup is interrupted

.changeset/mcp-interrupted-startup.md

@@ -0,0 +1,5 @@
+---
+"@kilocode/cli": patch
+---
+
+Stop MCP servers from piling up when a request is cancelled while they are still starting. Servers that had already started are now shut down, and the next request starts MCP again instead of failing until the backend restarts.

packages/opencode/src/effect/instance-state.ts

@@ -1,5 +1,6 @@
 import { Effect, ScopedCache, Scope } from "effect"
 import { capture, type InstanceContext } from "@/kilocode/instance" // kilocode_change
+import * as KiloInstanceState from "@/kilocode/effect/instance-state" // kilocode_change
 import { InstanceRef, WorkspaceRef } from "./instance-ref"
 import { registerDisposer } from "./instance-registry"
 import { WorkspaceContext } from "@/control-plane/workspace-context"
@@ -27,8 +28,11 @@ export const make = <A, E = never, R = never>(
   init: (ctx: InstanceContext) => Effect.Effect<A, E, R | Scope.Scope>,
 ): Effect.Effect<InstanceState<A, E, Exclude<R, Scope.Scope>>, never, R | Scope.Scope> =>
   Effect.gen(function* () {
-    const cache = yield* ScopedCache.make<string, A, E, R>({
+    // kilocode_change start - do not cache a lookup interrupted by a disconnected caller
+    const cache = yield* ScopedCache.makeWith<string, A, E, R>({
       capacity: Number.POSITIVE_INFINITY,
+      timeToLive: KiloInstanceState.ttl,
+      // kilocode_change end
       lookup: () =>
         Effect.gen(function* () {
           return yield* init(yield* context)

packages/opencode/src/kilocode/effect/instance-state.ts

@@ -0,0 +1,9 @@
+import { Duration, Exit } from "effect"
+
+/**
+ * Cache lifetime for an InstanceState lookup. The lookup runs on the fiber of the first
+ * caller, so a client that disconnects mid-lookup interrupts it. That interruption is not
+ * a result: expire it at once so the next caller closes the abandoned scope (running the
+ * finalizers it registered) and builds the state again. Real values and failures stay cached.
+ */
+export const ttl = <A, E>(exit: Exit.Exit<A, E>) => (Exit.hasInterrupts(exit) ? Duration.zero : Duration.infinity)

packages/opencode/src/kilocode/mcp/cleanup.ts

@@ -0,0 +1,45 @@
+import { Effect } from "effect"
+
+type Client = {
+  readonly transport?: unknown
+  close(): Promise<void>
+}
+
+/**
+ * Closes the clients the MCP state has connected so far. Registered before any server
+ * starts, so it also covers a state lookup that is interrupted part way through, when the
+ * upstream finalizer has not been registered yet and connected servers would otherwise
+ * keep running with no owner. When the lookup completes, the upstream finalizer runs
+ * first and empties `clients`, which makes this a no-op.
+ */
+export const release = (
+  state: { clients: Record<string, Client> },
+  descendants: (pid: number) => Effect.Effect<number[]>,
+) =>
+  Effect.suspend(() => {
+    const clients = Object.values(state.clients)
+    state.clients = {}
+    return Effect.forEach(
+      clients,
+      (client) =>
+        Effect.gen(function* () {
+          // Only stdio transports have a server process.
+          const transport = client.transport
+          if (
+            typeof transport === "object" &&
+            transport !== null &&
+            "pid" in transport &&
+            typeof transport.pid === "number"
+          ) {
+            for (const child of yield* descendants(transport.pid)) {
+              // The child may already have exited.
+              yield* Effect.try({ try: () => process.kill(child, "SIGTERM"), catch: () => undefined }).pipe(
+                Effect.ignore,
+              )
+            }
+          }
+          yield* Effect.tryPromise(() => client.close()).pipe(Effect.ignore)
+        }),
+      { concurrency: "unbounded", discard: true },
+    )
+  })

packages/opencode/src/mcp/index.ts

@@ -34,6 +34,7 @@ import { McpOAuthPendingProvider, McpOAuthProvider, OAUTH_CALLBACK_PATH } from "
 import { McpOAuthCallback } from "./oauth-callback"
 import { McpAuth } from "./auth"
 import { probe } from "@/kilocode/mcp/sse-probe" // kilocode_change - normalize the optional GET stream probe
+import * as KiloMcpCleanup from "@/kilocode/mcp/cleanup" // kilocode_change
 import { EventV2Bridge } from "@/event-v2-bridge"
 import { TuiEvent } from "@/server/tui-event"
 import { Cause, Effect, Exit, Layer, Context, Schema, Stream } from "effect"
@@ -589,6 +590,7 @@ const layer = Layer.effect(
           defs: {},
           instructions: {},
         }
+        yield* Effect.addFinalizer(() => KiloMcpCleanup.release(s, descendants)) // kilocode_change - close servers if this lookup is interrupted
 
         yield* Effect.forEach(
           Object.entries(config),

packages/opencode/test/kilocode/fixture/mcp-pid-stdio.ts

@@ -0,0 +1,36 @@
+import { Server } from "@modelcontextprotocol/sdk/server/index.js"
+import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js"
+import { ListToolsRequestSchema } from "@modelcontextprotocol/sdk/types.js"
+
+// Stdio MCP server that reports its pid. With --hang it never answers initialize.
+const file = process.env.MCP_PID_FILE
+if (!file) throw new Error("MCP_PID_FILE is required")
+
+if (process.argv.includes("--hang")) {
+  await Bun.write(file, String(process.pid))
+  await new Promise(() => {})
+}
+
+const server = new Server(
+  { name: "mcp-pid-stdio", version: "1.0.0" },
+  { capabilities: { tools: { listChanged: true } } },
+)
+let timer: ReturnType<typeof setInterval> | undefined
+let listed = 0
+
+server.setRequestHandler(ListToolsRequestSchema, async () => {
+  listed++
+  if (listed === 1) {
+    await Bun.write(file, String(process.pid))
+    // The client only re-lists tools on a change notification once the MCP state holds it
+    // as a connected client, so keep announcing a change until that second request arrives.
+    timer = setInterval(() => void server.sendToolListChanged(), 20)
+  }
+  if (listed === 2) {
+    clearInterval(timer)
+    await Bun.write(`${file}.stored`, "")
+  }
+  return { tools: [{ name: "ping", description: "ping", inputSchema: { type: "object", properties: {} } }] }
+})
+
+await server.connect(new StdioServerTransport())

packages/opencode/test/kilocode/instance-state-interrupt.test.ts

@@ -0,0 +1,56 @@
+import { expect } from "bun:test"
+import { CrossSpawnSpawner } from "@opencode-ai/core/cross-spawn-spawner"
+import { LayerNode } from "@opencode-ai/core/effect/layer-node"
+import { Deferred, Effect, Fiber, Layer } from "effect"
+import { InstanceState } from "@/effect/instance-state"
+import { provideInstanceEffect, testInstanceStoreLayer, tmpdirScoped } from "../fixture/fixture"
+import { testEffect } from "../lib/effect"
+
+const it = testEffect(Layer.mergeAll(LayerNode.compile(CrossSpawnSpawner.node), testInstanceStoreLayer))
+
+const access = <A, E>(state: InstanceState.InstanceState<A, E>, dir: string) =>
+  InstanceState.get(state).pipe(provideInstanceEffect(dir))
+
+it.live("InstanceState retries a lookup whose caller was interrupted", () =>
+  Effect.gen(function* () {
+    const dir = yield* tmpdirScoped()
+    const started = yield* Deferred.make<void>()
+    const released: number[] = []
+    let n = 0
+    const state = yield* InstanceState.make(() =>
+      Effect.gen(function* () {
+        const id = ++n
+        yield* Effect.addFinalizer(() => Effect.sync(() => released.push(id)))
+        if (id === 1) {
+          yield* Deferred.succeed(started, undefined)
+          return yield* Effect.never
+        }
+        return id
+      }),
+    )
+
+    // A client that disconnects mid-lookup interrupts the request fiber running it.
+    const fiber = yield* access(state, dir).pipe(Effect.forkScoped)
+    yield* Deferred.await(started)
+    yield* Fiber.interrupt(fiber)
+
+    // The next caller gets a fresh value instead of the cached interruption, and the
+    // abandoned lookup's resources are released.
+    expect(yield* access(state, dir)).toBe(2)
+    expect(released).toEqual([1])
+  }),
+)
+
+it.live("InstanceState still caches a lookup that failed", () =>
+  Effect.gen(function* () {
+    const dir = yield* tmpdirScoped()
+    let n = 0
+    const state = yield* InstanceState.make(() => Effect.suspend(() => Effect.fail(++n)))
+
+    const a = yield* access(state, dir).pipe(Effect.flip)
+    const b = yield* access(state, dir).pipe(Effect.flip)
+
+    expect(a).toBe(1)
+    expect(b).toBe(1)
+  }),
+)

packages/opencode/test/kilocode/mcp-interrupted-state.test.ts

@@ -0,0 +1,87 @@
+import path from "node:path"
+import os from "node:os"
+import fs from "node:fs/promises"
+import { expect } from "bun:test"
+import { LayerNode } from "@opencode-ai/core/effect/layer-node"
+import { Effect, Fiber } from "effect"
+import { MCP } from "../../src/mcp/index"
+import { pollWithTimeout, testEffect } from "../lib/effect"
+
+const it = testEffect(LayerNode.compile(MCP.node))
+const fixture = path.join(import.meta.dir, "fixture/mcp-pid-stdio.ts")
+const dir = path.join(os.tmpdir(), `kilo-mcp-interrupt-${process.pid}`)
+const ready = path.join(dir, "ready.pid")
+const hanging = path.join(dir, "hanging.pid")
+
+const pid = (file: string) =>
+  pollWithTimeout(
+    Effect.promise(async () => {
+      const handle = Bun.file(file)
+      return (await handle.exists()) ? Number(await handle.text()) : undefined
+    }),
+    `${path.basename(file)} was not written`,
+  )
+
+const alive = (id: number) => {
+  try {
+    process.kill(id, 0)
+    return true
+  } catch {
+    return false
+  }
+}
+
+it.instance(
+  "an interrupted MCP startup closes servers that already connected and retries",
+  () =>
+    Effect.gen(function* () {
+      const mcp = yield* MCP.Service
+
+      // A client that disconnects while MCP servers are still starting interrupts the lookup.
+      const fiber = yield* mcp.status().pipe(Effect.forkScoped)
+      const first = yield* pid(ready)
+      yield* pid(hanging)
+      // The fixture writes this once the MCP state holds `ready` as a connected client.
+      yield* pollWithTimeout(
+        Effect.promise(() => Bun.file(`${ready}.stored`).exists()).pipe(
+          Effect.map((done) => (done ? true : undefined)),
+        ),
+        "ready client was never stored",
+      )
+      yield* Fiber.interrupt(fiber)
+
+      const status = yield* mcp.status()
+      expect(status["ready"]?.status).toBe("connected")
+      expect(status["hanging"]?.status).toBe("failed")
+
+      // The server that connected before the interruption is not left running.
+      yield* pollWithTimeout(
+        Effect.sync(() => (alive(first) ? undefined : true)),
+        "server from the interrupted startup is still running",
+      )
+    }),
+  {
+    init: () =>
+      Effect.promise(async () => {
+        await fs.rm(dir, { recursive: true, force: true })
+        await fs.mkdir(dir, { recursive: true })
+      }),
+    config: {
+      mcp: {
+        ready: {
+          type: "local",
+          command: [process.execPath, fixture],
+          environment: { MCP_PID_FILE: ready },
+        },
+        hanging: {
+          type: "local",
+          command: [process.execPath, fixture, "--hang"],
+          environment: { MCP_PID_FILE: hanging },
+          // Keeps the lookup in progress until the test interrupts it; also bounds the retry.
+          timeout: 5000,
+        },
+      },
+    },
+  },
+  30_000,
+)