Переглянути джерело

fix(opencode): serialize mcp auth mutations (#29852)

Shoubhit Dash 3 місяців тому
батько
коміт
fa73ec4fa6

+ 40 - 15
packages/opencode/src/mcp/auth.ts

@@ -3,6 +3,7 @@ import { serviceUse } from "@opencode-ai/core/effect/service-use"
 import { Global } from "@opencode-ai/core/global"
 import { Effect, Layer, Context, Option, Schema } from "effect"
 import { AppFileSystem } from "@opencode-ai/core/filesystem"
+import { EffectFlock } from "@opencode-ai/core/util/effect-flock"
 
 export const Tokens = Schema.Struct({
   accessToken: Schema.mutableKey(Schema.String),
@@ -33,6 +34,7 @@ const decodeAuthData = Schema.decodeUnknownOption(Schema.Record(Schema.String, E
 type AuthData = Record<string, Entry>
 
 const filepath = path.join(Global.Path.data, "mcp-auth.json")
+const lockKey = `mcp-auth:${filepath}`
 
 export interface Interface {
   readonly all: () => Effect.Effect<Record<string, Entry>>
@@ -58,14 +60,27 @@ export const layer = Layer.effect(
   Service,
   Effect.gen(function* () {
     const fs = yield* AppFileSystem.Service
+    const flock = yield* EffectFlock.Service
 
-    const all = Effect.fn("McpAuth.all")(function* () {
+    const read = Effect.fn("McpAuth.read")(function* () {
       return yield* fs.readJson(filepath).pipe(
         Effect.map((data): AuthData => Option.getOrElse(decodeAuthData(data), () => ({}) as AuthData) as AuthData),
         Effect.catch(() => Effect.succeed({} as AuthData)),
       )
     })
 
+    const all = Effect.fn("McpAuth.all")(function* () {
+      return yield* read().pipe(flock.withLock(lockKey), Effect.orDie)
+    })
+
+    const mutate = Effect.fn("McpAuth.mutate")(function* (update: (data: AuthData) => AuthData | undefined) {
+      yield* Effect.gen(function* () {
+        const next = update(yield* read())
+        if (!next) return
+        yield* fs.writeJson(filepath, next, 0o600).pipe(Effect.orDie)
+      }).pipe(flock.withLock(lockKey), Effect.orDie)
+    })
+
     const get = Effect.fn("McpAuth.get")(function* (mcpName: string) {
       const data = yield* all()
       return data[mcpName]
@@ -80,31 +95,38 @@ export const layer = Layer.effect(
     })
 
     const set = Effect.fn("McpAuth.set")(function* (mcpName: string, entry: Entry, serverUrl?: string) {
-      const data = yield* all()
-      if (serverUrl) entry.serverUrl = serverUrl
-      yield* fs.writeJson(filepath, { ...data, [mcpName]: entry }, 0o600).pipe(Effect.orDie)
+      yield* mutate((data) => ({
+        ...data,
+        [mcpName]: serverUrl ? { ...entry, serverUrl } : entry,
+      }))
     })
 
     const remove = Effect.fn("McpAuth.remove")(function* (mcpName: string) {
-      const data = yield* all()
-      delete data[mcpName]
-      yield* fs.writeJson(filepath, data, 0o600).pipe(Effect.orDie)
+      yield* mutate((data) => {
+        const next = { ...data }
+        delete next[mcpName]
+        return next
+      })
     })
 
     const updateField = <K extends keyof Entry>(field: K, spanName: string) =>
       Effect.fn(`McpAuth.${spanName}`)(function* (mcpName: string, value: NonNullable<Entry[K]>, serverUrl?: string) {
-        const entry = (yield* get(mcpName)) ?? {}
-        entry[field] = value
-        yield* set(mcpName, entry, serverUrl)
+        yield* mutate((data) => {
+          const entry = data[mcpName] ?? {}
+          entry[field] = value
+          if (serverUrl) entry.serverUrl = serverUrl
+          return { ...data, [mcpName]: entry }
+        })
       })
 
     const clearField = (field: keyof Entry, spanName: string) =>
       Effect.fn(`McpAuth.${spanName}`)(function* (mcpName: string) {
-        const entry = yield* get(mcpName)
-        if (entry) {
+        yield* mutate((data) => {
+          const entry = data[mcpName]
+          if (!entry) return undefined
           delete entry[field]
-          yield* set(mcpName, entry)
-        }
+          return { ...data, [mcpName]: entry }
+        })
       })
 
     const updateTokens = updateField("tokens", "updateTokens")
@@ -144,6 +166,9 @@ export const layer = Layer.effect(
   }),
 )
 
-export const defaultLayer = layer.pipe(Layer.provide(AppFileSystem.defaultLayer))
+export const defaultLayer = layer.pipe(
+  Layer.provide(EffectFlock.defaultLayer),
+  Layer.provide(AppFileSystem.defaultLayer),
+)
 
 export * as McpAuth from "./auth"

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

@@ -971,7 +971,7 @@ export type AuthStatus = "authenticated" | "expired" | "not_authenticated"
 // --- Per-service runtime ---
 
 export const defaultLayer = layer.pipe(
-  Layer.provide(McpAuth.layer),
+  Layer.provide(McpAuth.defaultLayer),
   Layer.provide(Bus.layer),
   Layer.provide(Config.defaultLayer),
   Layer.provide(CrossSpawnSpawner.defaultLayer),

+ 83 - 0
packages/opencode/test/mcp/auth.test.ts

@@ -0,0 +1,83 @@
+import { expect, test } from "bun:test"
+import { setTimeout as sleep } from "node:timers/promises"
+import { Effect, Layer } from "effect"
+import { AppFileSystem } from "@opencode-ai/core/filesystem"
+import { EffectFlock } from "@opencode-ai/core/util/effect-flock"
+import { McpAuth } from "../../src/mcp/auth"
+
+function authFile() {
+  let raw = ""
+  let activeWrites = 0
+  let sawOverlap = false
+
+  const layer = Layer.effect(
+    AppFileSystem.Service,
+    Effect.gen(function* () {
+      const fs = yield* AppFileSystem.Service
+
+      return AppFileSystem.Service.of({
+        ...fs,
+        readJson: (file) =>
+          file.endsWith("mcp-auth.json")
+            ? Effect.try({
+                try: () => {
+                  if (!raw) throw new Error("mcp-auth.json missing")
+                  return JSON.parse(raw)
+                },
+                catch: (cause) => new AppFileSystem.FileSystemError({ method: "readJson", cause }),
+              })
+            : fs.readJson(file),
+        writeJson: (file, value, mode) =>
+          file.endsWith("mcp-auth.json")
+            ? Effect.promise(async () => {
+                activeWrites++
+                sawOverlap = sawOverlap || activeWrites > 1
+                raw = ""
+                await sleep(10)
+                const next = JSON.stringify(value, null, 2)
+                raw = sawOverlap ? `${next}\n}` : next
+                activeWrites--
+              })
+            : fs.writeJson(file, value, mode),
+      })
+    }),
+  ).pipe(Layer.provide(AppFileSystem.defaultLayer))
+
+  return { layer, raw: () => raw }
+}
+
+function authService(layer: Layer.Layer<AppFileSystem.Service>) {
+  return McpAuth.Service.use((auth) => Effect.succeed(auth)).pipe(
+    Effect.provide(
+      McpAuth.layer.pipe(
+        Layer.provide(EffectFlock.defaultLayer),
+        Layer.provide(layer),
+      ),
+    ),
+  )
+}
+
+test("serializes concurrent auth file updates across service instances", async () => {
+  const file = authFile()
+
+  await Effect.runPromise(
+    Effect.gen(function* () {
+      const first = yield* authService(file.layer)
+      const second = yield* authService(file.layer)
+
+      yield* Effect.all(
+        [
+          first.updateTokens("posthog", { accessToken: "access-token" }, "https://mcp.posthog.com/mcp"),
+          second.updateClientInfo("posthog", { clientId: "client-id" }, "https://mcp.posthog.com/mcp"),
+        ],
+        { concurrency: "unbounded" },
+      )
+
+      const entry = yield* first.get("posthog")
+      expect(entry?.tokens?.accessToken).toBe("access-token")
+      expect(entry?.clientInfo?.clientId).toBe("client-id")
+      expect(entry?.serverUrl).toBe("https://mcp.posthog.com/mcp")
+      expect(() => JSON.parse(file.raw())).not.toThrow()
+    }),
+  )
+})