Przeglądaj źródła

fix(app): start MCP servers only for open directories (#28937)

Luke Parker 2 miesięcy temu
rodzic
commit
e16bfd745d

+ 1 - 1
packages/app/src/context/directory-sync.ts

@@ -178,7 +178,7 @@ export const createDirSyncContext = (directory: string, serverSync: ReturnType<t
   type Child = ReturnType<(typeof serverSync)["child"]>
   type Setter = Child[1]
 
-  const current = createMemo(() => serverSync.child(directory))
+  const current = createMemo(() => serverSync.child(directory, { mcp: true }))
   const target = (directory?: string) => {
     if (!directory || directory === directory) return current()
     return serverSync.child(directory)

+ 15 - 2
packages/app/src/context/global-sync/bootstrap.test.ts

@@ -10,6 +10,7 @@ const provider = { all: new Map(), connected: [], default: {} } satisfies Normal
 
 describe("bootstrapDirectory", () => {
   test("marks a loading directory partial during bootstrap and complete after success", async () => {
+    const mcpReads: string[] = []
     const [store, setStore] = createStore<State>({
       status: "loading",
       agent: [],
@@ -44,6 +45,7 @@ describe("bootstrapDirectory", () => {
 
     await bootstrapDirectory({
       directory: "/project",
+      mcp: false,
       global: {
         config: {} satisfies Config,
         path: { state: "", config: "", worktree: "/project", directory: "/project", home: "/home" },
@@ -55,10 +57,20 @@ describe("bootstrapDirectory", () => {
         config: { get: async () => ({ data: {} }) },
         session: { status: async () => ({ data: {} }) },
         vcs: { get: async () => ({ data: undefined }) },
-        command: { list: async () => ({ data: [] }) },
+        command: {
+          list: async () => {
+            mcpReads.push("command")
+            return { data: [] }
+          },
+        },
         permission: { list: async () => ({ data: [] }) },
         question: { list: async () => ({ data: [] }) },
-        mcp: { status: async () => ({ data: {} }) },
+        mcp: {
+          status: async () => {
+            mcpReads.push("status")
+            return { data: {} }
+          },
+        },
         provider: { list: async () => ({ data: { all: [], connected: [], default: {} } }) },
       } as unknown as OpencodeClient,
       store,
@@ -74,5 +86,6 @@ describe("bootstrapDirectory", () => {
     await new Promise((resolve) => setTimeout(resolve, 80))
 
     expect(store.status).toBe("complete")
+    expect(mcpReads).toEqual([])
   })
 })

+ 3 - 2
packages/app/src/context/global-sync/bootstrap.ts

@@ -198,6 +198,7 @@ export const loadPathQuery = (directory: string | null, sdk: OpencodeClient) =>
 
 export async function bootstrapDirectory(input: {
   directory: string
+  mcp: boolean
   sdk: OpencodeClient
   store: Store<State>
   setStore: SetStoreFunction<State>
@@ -250,7 +251,7 @@ export async function bootstrapDirectory(input: {
             if (next) input.vcsCache.setStore("value", next)
           }),
         ),
-      () => retry(() => input.sdk.command.list().then((x) => input.setStore("command", x.data ?? []))),
+      input.mcp && (() => retry(() => input.sdk.command.list().then((x) => input.setStore("command", x.data ?? [])))),
       () =>
         retry(() =>
           input.sdk.permission.list().then((x) => {
@@ -304,7 +305,7 @@ export async function bootstrapDirectory(input: {
           }),
         ),
       () => Promise.resolve(input.loadSessions(input.directory)),
-      () => input.queryClient.fetchQuery(loadMcpQuery(input.directory, input.sdk)),
+      input.mcp && (() => input.queryClient.fetchQuery(loadMcpQuery(input.directory, input.sdk))),
       () =>
         input.queryClient.fetchQuery(loadProvidersQuery(input.directory, input.sdk)).catch((err) => {
           const project = getFilename(input.directory)

+ 53 - 6
packages/app/src/context/global-sync/child-store.test.ts

@@ -6,6 +6,7 @@ import type { State } from "./types"
 import type { QueryOptionsApi } from "../server-sync"
 
 let createChildStoreManager: typeof import("./child-store").createChildStoreManager
+const queryGroups: Array<() => { queries: Array<{ enabled?: boolean }> }> = []
 
 const child = () => createStore({} as State)
 const provider = { all: new Map(), connected: [], default: {} } satisfies NormalizedProviderListResponse
@@ -48,12 +49,15 @@ beforeAll(async () => {
     persisted: (_target: string, store: unknown[]) => [store[0], store[1], null, () => true],
   }))
   mock.module("@tanstack/solid-query", () => ({
-    useQueries: () => [
-      { isLoading: false, data: { state: "", config: "", worktree: "", directory: "", home: "" } },
-      { isLoading: false, data: {} },
-      { isLoading: false, data: [] },
-      { isLoading: false, data: provider },
-    ],
+    useQueries: (options: () => { queries: Array<{ enabled?: boolean }> }) => {
+      queryGroups.push(options)
+      return [
+        { isLoading: false, data: { state: "", config: "", worktree: "", directory: "", home: "" } },
+        { isLoading: false, data: {} },
+        { isLoading: false, data: [] },
+        { isLoading: false, data: provider },
+      ]
+    },
   }))
 
   createChildStoreManager = (await import("./child-store")).createChildStoreManager
@@ -73,6 +77,7 @@ describe("createChildStoreManager", () => {
       isBooting: () => false,
       isLoadingSessions: () => false,
       onBootstrap() {},
+      onMcp() {},
       onDispose() {},
       translate: (key) => key,
       queryOptions: queryOptionsApi,
@@ -103,6 +108,7 @@ describe("createChildStoreManager", () => {
         onBootstrap(directory) {
           bootstraps.push(directory)
         },
+        onMcp() {},
         onDispose() {},
         translate: (key) => key,
         queryOptions: queryOptionsApi,
@@ -121,4 +127,45 @@ describe("createChildStoreManager", () => {
       dispose()
     }
   })
+
+  test("enables MCP only when requested for the directory", () => {
+    let manager: ReturnType<typeof createChildStoreManager> | undefined
+    const offset = queryGroups.length
+    const mcpLoads: string[] = []
+
+    const dispose = createOwner((owner) => {
+      manager = createChildStoreManager({
+        owner,
+        isBooting: () => false,
+        isLoadingSessions: () => false,
+        onBootstrap() {},
+        onMcp(directory) {
+          mcpLoads.push(directory)
+        },
+        onDispose() {},
+        translate: (key) => key,
+        queryOptions: queryOptionsApi,
+        global: { provider },
+      })
+    })
+
+    try {
+      if (!manager) throw new Error("manager required")
+      const [, setStore] = manager.child("/project", { bootstrap: false })
+      const queries = queryGroups[offset]
+      if (!queries) throw new Error("queries required")
+      expect(queries().queries[1]?.enabled).toBe(false)
+
+      setStore("status", "complete")
+      manager.child("/project", { bootstrap: false, mcp: true })
+      expect(queries().queries[1]?.enabled).toBe(true)
+      expect(mcpLoads).toEqual(["/project"])
+
+      manager.disableMcp("/project")
+      expect(queries().queries[1]?.enabled).toBe(false)
+      expect(manager.mcp("/project")).toBe(false)
+    } finally {
+      dispose()
+    }
+  })
 })

+ 26 - 2
packages/app/src/context/global-sync/child-store.ts

@@ -1,4 +1,4 @@
-import { createRoot, getOwner, onCleanup, runWithOwner, type Owner } from "solid-js"
+import { createRoot, createSignal, getOwner, onCleanup, runWithOwner, type Owner } from "solid-js"
 import { createStore, type SetStoreFunction, type Store } from "solid-js/store"
 import { Persist, persisted } from "@/utils/persist"
 import type { VcsInfo } from "@opencode-ai/sdk/v2/client"
@@ -24,6 +24,7 @@ export function createChildStoreManager(input: {
   isBooting: (directory: string) => boolean
   isLoadingSessions: (directory: string) => boolean
   onBootstrap: (directory: string) => void
+  onMcp: (directory: string, setStore: SetStoreFunction<State>) => void
   onDispose: (directory: string) => void
   translate: (key: string, vars?: Record<string, string | number>) => string
   queryOptions: QueryOptionsApi
@@ -39,6 +40,8 @@ export function createChildStoreManager(input: {
   const pins = new Map<string, number>()
   const ownerPins = new WeakMap<object, Set<string>>()
   const disposers = new Map<string, () => void>()
+  const mcpDirectories = new Set<string>()
+  const mcpToggles = new Map<string, (enabled: boolean) => void>()
 
   const markKey = (key: DirectoryKey) => {
     if (!key) return
@@ -110,6 +113,8 @@ export function createChildStoreManager(input: {
     metaCache.delete(key)
     iconCache.delete(key)
     lifecycle.delete(key)
+    mcpDirectories.delete(key)
+    mcpToggles.delete(key)
     const dispose = disposers.get(key)
     if (dispose) {
       dispose()
@@ -173,11 +178,12 @@ export function createChildStoreManager(input: {
         createRoot((dispose) => {
           const initialMeta = meta[0].value
           const initialIcon = icon[0].value
+          const [mcpEnabled, setMcpEnabled] = createSignal(false)
 
           const [pathQuery, mcpQuery, lspQuery, providerQuery] = useQueries(() => ({
             queries: [
               input.queryOptions.path(key),
-              input.queryOptions.mcp(key),
+              { ...input.queryOptions.mcp(key), enabled: mcpEnabled() },
               input.queryOptions.lsp(key),
               input.queryOptions.providers(key),
             ],
@@ -236,6 +242,7 @@ export function createChildStoreManager(input: {
           })
           children[key] = child
           disposers.set(key, dispose)
+          mcpToggles.set(key, setMcpEnabled)
 
           const onPersistedInit = (init: Promise<string> | string | null, run: () => void) => {
             if (!(init instanceof Promise)) return
@@ -274,6 +281,7 @@ export function createChildStoreManager(input: {
     const key = directoryKey(directory)
     const childStore = ensureChild(directory)
     pinForOwner(key)
+    if (options.mcp) enableMcp(directory, key, childStore)
     const shouldBootstrap = options.bootstrap ?? true
     if (shouldBootstrap && childStore[0].status === "loading") {
       input.onBootstrap(directory)
@@ -284,6 +292,7 @@ export function createChildStoreManager(input: {
   function peek(directory: string, options: ChildOptions = {}) {
     const key = directoryKey(directory)
     const childStore = ensureChild(directory)
+    if (options.mcp) enableMcp(directory, key, childStore)
     const shouldBootstrap = options.bootstrap ?? true
     if (shouldBootstrap && childStore[0].status === "loading") {
       input.onBootstrap(directory)
@@ -291,6 +300,19 @@ export function createChildStoreManager(input: {
     return childStore
   }
 
+  function enableMcp(directory: string, key: DirectoryKey, childStore: [Store<State>, SetStoreFunction<State>]) {
+    if (mcpDirectories.has(key)) return
+    mcpDirectories.add(key)
+    mcpToggles.get(key)?.(true)
+    if (childStore[0].status !== "loading") input.onMcp(directory, childStore[1])
+  }
+
+  function disableMcp(directory: string) {
+    const key = directoryKey(directory)
+    if (!mcpDirectories.delete(key)) return
+    mcpToggles.get(key)?.(false)
+  }
+
   function projectMeta(directory: string, patch: ProjectMeta) {
     const key = directoryKey(directory)
     const [store, setStore] = ensureChild(directory)
@@ -330,6 +352,8 @@ export function createChildStoreManager(input: {
     pin,
     unpin,
     pinned,
+    mcp: (directory: string) => mcpDirectories.has(directoryKey(directory)),
+    disableMcp,
     disposeDirectory,
     runEviction,
     vcsCache,

+ 1 - 0
packages/app/src/context/global-sync/types.ts

@@ -98,6 +98,7 @@ export type IconCache = {
 
 export type ChildOptions = {
   bootstrap?: boolean
+  mcp?: boolean
 }
 
 export type DirState = {

+ 17 - 1
packages/app/src/context/server-sync.tsx

@@ -31,6 +31,7 @@ import { PathKey } from "@/utils/path-key"
 import { createDirSyncContext } from "./directory-sync"
 import { createSimpleContext, NormalizedProviderListResponse } from "@opencode-ai/ui/context"
 import { createRefCountMap } from "@/utils/refcount"
+import { retry } from "@opencode-ai/core/util/retry"
 
 type GlobalStore = {
   ready: boolean
@@ -205,6 +206,15 @@ export function createServerSyncContext() {
     onBootstrap: (directory) => {
       void bootstrapInstance(directory)
     },
+    onMcp: (directory, setStore) => {
+      void retry(() => sdkFor(directory).command.list().then((x) => setStore("command", x.data ?? []))).catch((err) => {
+        showToast({
+          variant: "error",
+          title: language.t("toast.project.reloadFailed.title", { project: getFilename(directory) }),
+          description: formatServerError(err, language.t),
+        })
+      })
+    },
     onDispose: (directory) => {
       const key = directoryKey(directory)
       queue.clear(key)
@@ -311,6 +321,7 @@ export function createServerSyncContext() {
       const sdk = sdkFor(directory)
       await bootstrapDirectory({
         directory,
+        mcp: children.mcp(key),
         global: {
           config: globalStore.config,
           path: globalStore.path,
@@ -437,6 +448,7 @@ export function createServerSyncContext() {
     },
     child: children.child,
     peek: children.peek,
+    disableMcp: children.disableMcp,
     queryOptions: queryOptionsApi,
     // bootstrap,
     updateConfig: updateConfigMutation.mutateAsync,
@@ -454,7 +466,11 @@ export const { use: useServerSync, provider: ServerSyncProvider } = createSimple
 
     return {
       ...sync,
-      createDirSyncContext: createRefCountMap((dir) => createDirSyncContext(dir, sync)),
+      createDirSyncContext: createRefCountMap(
+        (dir) => createDirSyncContext(dir, sync),
+        (dir) => sync.disableMcp(dir),
+        directoryKey,
+      ),
     }
   },
 })

+ 49 - 0
packages/app/src/utils/refcount.test.ts

@@ -0,0 +1,49 @@
+import { describe, expect, test } from "bun:test"
+import { createRoot } from "solid-js"
+import { createRefCountMap } from "./refcount"
+import { pathKey } from "./path-key"
+
+describe("createRefCountMap", () => {
+  test("removes an item after its last owner is disposed", () => {
+    const removed: string[] = []
+    const map = createRefCountMap(
+      (key) => key,
+      (key) => removed.push(key),
+    )
+    const first = createRoot((dispose) => {
+      map("/project")
+      return dispose
+    })
+    const second = createRoot((dispose) => {
+      map("/project")
+      return dispose
+    })
+
+    first()
+    expect(removed).toEqual([])
+    second()
+    expect(removed).toEqual(["/project"])
+  })
+
+  test("keeps equivalent path consumers until the last owner is disposed", () => {
+    const removed: string[] = []
+    const map = createRefCountMap(
+      (key) => key,
+      (key) => removed.push(key),
+      pathKey,
+    )
+    const first = createRoot((dispose) => {
+      map("C:\\repo")
+      return dispose
+    })
+    const second = createRoot((dispose) => {
+      map("C:/repo/")
+      return dispose
+    })
+
+    first()
+    expect(removed).toEqual([])
+    second()
+    expect(removed).toEqual(["C:/repo"])
+  })
+})

+ 15 - 9
packages/app/src/utils/refcount.ts

@@ -1,26 +1,32 @@
 import { onCleanup } from "solid-js"
 
-export function createRefCountMap<T>(create: (key: string) => T) {
+export function createRefCountMap<T>(
+  create: (key: string) => T,
+  remove?: (key: string) => void,
+  identity: (key: string) => string = (key) => key,
+) {
   const items = new Map<string, T>()
   const refCounts = new Map<string, number>()
 
   return (key: string) => {
+    const id = identity(key)
     onCleanup(() => {
-      refCounts.set(key, (refCounts.get(key) ?? 0) - 1)
-      if (refCounts.get(key) === 0) {
-        items.delete(key)
-        refCounts.delete(key)
+      refCounts.set(id, (refCounts.get(id) ?? 0) - 1)
+      if (refCounts.get(id) === 0) {
+        remove?.(id)
+        items.delete(id)
+        refCounts.delete(id)
       }
     })
 
-    const cached = items.get(key)
+    const cached = items.get(id)
     if (cached) {
-      refCounts.set(key, (refCounts.get(key) ?? 0) + 1)
+      refCounts.set(id, (refCounts.get(id) ?? 0) + 1)
       return cached
     }
     const item = create(key)
-    items.set(key, item)
-    refCounts.set(key, 1)
+    items.set(id, item)
+    refCounts.set(id, 1)
     return item
   }
 }