Browse Source

fix: Opencode hangs after exit (#3481)

Co-authored-by: Aiden Cline <[email protected]>
Haris Gušić 3 months ago
parent
commit
c1ada302f9

+ 0 - 1
packages/opencode/src/acp/session.ts

@@ -1,5 +1,4 @@
 import type { McpServer } from "@agentclientprotocol/sdk"
-import { Identifier } from "../id/id"
 import { Session } from "../session"
 import { Provider } from "../provider/provider"
 import type { ACPSessionState } from "./types"

+ 6 - 3
packages/opencode/src/cli/bootstrap.ts

@@ -6,9 +6,12 @@ export async function bootstrap<T>(directory: string, cb: () => Promise<T>) {
     directory,
     init: InstanceBootstrap,
     fn: async () => {
-      const result = await cb()
-      await Instance.dispose()
-      return result
+      try {
+        const result = await cb()
+        return result
+      } finally {
+        await Instance.dispose()
+      }
     },
   })
 }

+ 5 - 1
packages/opencode/src/cli/cmd/acp.ts

@@ -56,7 +56,11 @@ export const AcpCommand = cmd({
       }, stream)
 
       log.info("setup connection")
+      process.stdin.resume()
+      await new Promise((resolve, reject) => {
+        process.stdin.on("end", resolve)
+        process.stdin.on("error", reject)
+      })
     })
-    process.stdin.resume()
   },
 })

+ 1 - 1
packages/opencode/src/cli/cmd/serve.ts

@@ -27,6 +27,6 @@ export const ServeCommand = cmd({
     })
     console.log(`opencode server listening on http://${server.hostname}:${server.port}`)
     await new Promise(() => {})
-    server.stop()
+    await server.stop()
   },
 })

+ 3 - 1
packages/opencode/src/config/config.ts

@@ -77,14 +77,16 @@ export namespace Config {
       log.debug("loading config from OPENCODE_CONFIG_DIR", { path: Flag.OPENCODE_CONFIG_DIR })
     }
 
+    const promises: Promise<void>[] = []
     for (const dir of directories) {
       await assertValid(dir)
-      installDependencies(dir)
+      promises.push(installDependencies(dir))
       result.command = mergeDeep(result.command ?? {}, await loadCommand(dir))
       result.agent = mergeDeep(result.agent, await loadAgent(dir))
       result.agent = mergeDeep(result.agent, await loadMode(dir))
       result.plugin.push(...(await loadPlugin(dir)))
     }
+    await Promise.all(promises)
 
     // Migrate deprecated mode field to agent field
     for (const [name, mode] of Object.entries(result.mode)) {

+ 2 - 1
packages/opencode/src/file/watcher.ts

@@ -63,7 +63,8 @@ export namespace FileWatcher {
       return { sub }
     },
     async (state) => {
-      state.sub?.unsubscribe()
+      if (!state.sub) return
+      await state.sub?.unsubscribe()
     },
   )
 

+ 6 - 4
packages/opencode/src/index.ts

@@ -22,8 +22,6 @@ import { AttachCommand } from "./cli/cmd/attach"
 import { AcpCommand } from "./cli/cmd/acp"
 import { EOL } from "os"
 
-const cancel = new AbortController()
-
 process.on("unhandledRejection", (e) => {
   Log.Default.error("rejection", {
     e: e instanceof Error ? e.message : e,
@@ -135,6 +133,10 @@ try {
     console.error(e)
   }
   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();
 }
-
-cancel.abort()

+ 1 - 3
packages/opencode/src/lsp/index.ts

@@ -101,9 +101,7 @@ export namespace LSP {
       }
     },
     async (state) => {
-      for (const client of state.clients) {
-        await client.shutdown()
-      }
+      await Promise.all(state.clients.map((client) => client.shutdown()))
     },
   )
 

+ 1 - 3
packages/opencode/src/mcp/index.ts

@@ -45,9 +45,7 @@ export namespace MCP {
       }
     },
     async (state) => {
-      for (const client of Object.values(state.clients)) {
-        client.close()
-      }
+      await Promise.all(Object.values(state.clients).map((client) => client.close()))
     },
   )
 

+ 1 - 1
packages/opencode/src/project/bootstrap.ts

@@ -11,7 +11,7 @@ export async function InstanceBootstrap() {
   await Plugin.init()
   Share.init()
   Format.init()
-  LSP.init()
+  await LSP.init()
   FileWatcher.init()
   File.init()
 }

+ 41 - 9
packages/opencode/src/project/state.ts

@@ -1,23 +1,26 @@
+import { Log } from "@/util/log"
+
 export namespace State {
   interface Entry {
     state: any
     dispose?: (state: any) => Promise<void>
   }
 
-  const entries = new Map<string, Map<any, Entry>>()
+  const log = Log.create({ service: "state" })
+  const recordsByKey = new Map<string, Map<any, Entry>>()
 
   export function create<S>(root: () => string, init: () => S, dispose?: (state: Awaited<S>) => Promise<void>) {
     return () => {
       const key = root()
-      let collection = entries.get(key)
-      if (!collection) {
-        collection = new Map<string, Entry>()
-        entries.set(key, collection)
+      let entries = recordsByKey.get(key)
+      if (!entries) {
+        entries = new Map<string, Entry>()
+        recordsByKey.set(key, entries)
       }
-      const exists = collection.get(init)
+      const exists = entries.get(init)
       if (exists) return exists.state as S
       const state = init()
-      collection.set(init, {
+      entries.set(init, {
         state,
         dispose,
       })
@@ -26,9 +29,38 @@ export namespace State {
   }
 
   export async function dispose(key: string) {
-    for (const [_, entry] of entries.get(key)?.entries() ?? []) {
+    const entries = recordsByKey.get(key)
+    if (!entries) return
+
+    log.info("waiting for state disposal to complete", { key })
+
+    let disposalFinished = false
+
+    setTimeout(() => {
+      if (!disposalFinished) {
+        log.warn(
+          "state disposal is taking an unusually long time - if it does not complete in a reasonable time, please report this as a bug",
+          { key },
+        )
+      }
+    }, 10000).unref()
+
+    const tasks: Promise<void>[] = []
+    for (const entry of entries.values()) {
       if (!entry.dispose) continue
-      await entry.dispose(await entry.state)
+
+      const task = Promise.resolve(entry.state)
+        .then((state) => entry.dispose!(state))
+        .catch((error) => {
+          log.error("Error while disposing state:", { error, key })
+        })
+
+      tasks.push(task)
     }
+
+    await Promise.all(tasks)
+
+    disposalFinished = true
+    log.info("state disposal completed", { key })
   }
 }

+ 19 - 8
packages/opencode/src/session/prompt.ts

@@ -76,13 +76,22 @@ export namespace SessionPrompt {
           callback: (input: MessageV2.WithParts) => void
         }[]
       >()
+      const pending = new Set<Promise<void>>()
+
+      const track = (promise: Promise<void>) => {
+        pending.add(promise)
+        promise.finally(() => pending.delete(promise))
+      }
 
       return {
         queued,
+        pending,
+        track,
       }
     },
     async (current) => {
       current.queued.clear()
+      await Promise.allSettled([...current.pending])
     },
   )
 
@@ -227,13 +236,15 @@ export namespace SessionPrompt {
       step++
       await processor.next(msgs.findLast((m) => m.info.role === "user")?.info.id!)
       if (step === 1) {
-        ensureTitle({
-          session,
-          history: msgs,
-          message: userMsg,
-          providerID: model.providerID,
-          modelID: model.info.id,
-        })
+        state().track(
+          ensureTitle({
+            session,
+            history: msgs,
+            message: userMsg,
+            providerID: model.providerID,
+            modelID: model.info.id,
+          }),
+        )
         SessionSummary.summarize({
           sessionID: input.sessionID,
           messageID: userMsg.info.id,
@@ -1794,7 +1805,7 @@ export namespace SessionPrompt {
         thinkingBudget: 0,
       }
     }
-    generateText({
+    await generateText({
       maxOutputTokens: small.info.reasoning ? 1500 : 20,
       providerOptions: ProviderTransform.providerOptions(small.npm, small.providerID, options),
       messages: [