feat(simulation): control arbitrary tool lifecycles (#37816)

packages/core/src/tool/registry.ts

@@ -1,7 +1,7 @@
 export * as ToolRegistry from "./registry"
 
 import { ToolOutput, type ToolCall, type ToolDefinition, type ToolResultValue } from "@opencode-ai/ai"
-import { Context, Effect, Layer, Scope } from "effect"
+import { Context, Effect, Layer, Scope, Semaphore } from "effect"
 import type { AgentV2 } from "../agent"
 import { Image } from "../image"
 import { PermissionV2 } from "../permission"
@@ -45,6 +45,13 @@ export interface Interface {
     tools: Readonly<Record<string, AnyTool>>,
     options?: Tools.RegisterOptions,
   ) => Effect.Effect<void, RegistrationError, Scope.Scope>
+  /** Internal atomic registration capability used by plugin transforms. */
+  readonly registerBatch: (
+    registrations: ReadonlyArray<{
+      readonly tools: Readonly<Record<string, AnyTool>>
+      readonly options?: Tools.RegisterOptions
+    }>,
+  ) => Effect.Effect<void, RegistrationError, Scope.Scope>
 }
 
 export interface Materialization {
@@ -107,6 +114,7 @@ const registryLayer = Layer.effect(
       readonly codemode: boolean
     }
     const local = new Map<string, Array<{ readonly token: object; readonly registration: Registration }>>()
+    const registrationLock = Semaphore.makeUnsafe(1)
 
     const settleTool = Effect.fn("ToolRegistry.settleTool")(function* (input: ExecuteInput, tool: AnyTool) {
       // Hooks fire only for hosted/local tools; provider-executed calls never reach settleTool.
@@ -192,82 +200,111 @@ const registryLayer = Layer.effect(
       }
     })
 
-    return Service.of({
-      register: Effect.fn("ToolRegistry.register")(function* (tools, options) {
-        if (options?.namespace !== undefined) yield* validateNamespace(options.namespace)
-        const entries = registrationEntries(tools, options?.namespace)
-        if (entries.length === 0) return
-        const codemode = options?.codemode ?? true
-        const reserved = codemode ? undefined : entries.find((entry) => entry.key === "execute")
-        if (reserved)
-          return yield* Effect.fail(
-            new RegistrationError({ name: reserved.key, message: 'Tool name "execute" is reserved for CodeMode' }),
-          )
-        yield* Effect.uninterruptible(
+    const registerBatch: Interface["registerBatch"] = Effect.fn("ToolRegistry.registerBatch")(
+      function* (registrations) {
+        const planned = yield* Effect.forEach(registrations, ({ tools, options }) =>
           Effect.gen(function* () {
-            const token = {}
-            for (const entry of entries)
-              local.set(entry.key, [
-                ...(local.get(entry.key) ?? []),
-                {
-                  token,
-                  registration: {
-                    tool: entry.tool,
-                    name: entry.name,
-                    namespace: entry.namespace,
-                    codemode,
-                  },
-                },
-              ])
-            yield* Effect.addFinalizer(() =>
-              Effect.sync(() => {
-                for (const entry of entries) {
-                  const registrations =
-                    local.get(entry.key)?.filter((registration) => registration.token !== token) ?? []
-                  if (registrations.length > 0) local.set(entry.key, registrations)
-                  else local.delete(entry.key)
-                }
-              }),
-            )
+            if (options?.namespace !== undefined) yield* validateNamespace(options.namespace)
+            const entries = registrationEntries(tools, options?.namespace)
+            const codemode = options?.codemode ?? true
+            const reserved = codemode ? undefined : entries.find((entry) => entry.key === "execute")
+            if (reserved)
+              return yield* Effect.fail(
+                new RegistrationError({ name: reserved.key, message: 'Tool name "execute" is reserved for CodeMode' }),
+              )
+            return { entries, codemode }
           }),
         )
-      }),
-      materialize: Effect.fn("ToolRegistry.materialize")(function* (permissions) {
-        const direct = new Map<string, Registration>()
-        const codemode = new Map<string, Registration>()
-        const rules = permissions ?? []
-        for (const [name, entries] of local) {
-          const registration = entries.at(-1)?.registration
-          if (!registration) continue
-          if (whollyDisabled(permission(registration.tool, name), rules)) continue
-          if (registration.codemode) codemode.set(name, registration)
-          else direct.set(name, registration)
-        }
-        const execute =
-          codemode.size > 0 && !whollyDisabled("execute", rules) ? ExecuteTool.create(codemode) : undefined
-        return {
-          definitions: [
-            ...Array.from(direct, ([name, registration]) => definition(name, registration.tool)),
-            ...(execute ? [definition("execute", execute)] : []),
-          ],
-          settle: (input) => {
-            if (input.call.name === "execute" && execute) return settleTool(input, execute)
-            const registration = direct.get(input.call.name)
-            if (registration) return settleTool(input, registration.tool)
-            return Effect.succeed({
-              result: { type: "error", value: `Unknown tool: ${input.call.name}` },
-              error: { type: "tool.unknown", message: `Unknown tool: ${input.call.name}` },
-            })
+        if (planned.every((plan) => plan.entries.length === 0)) return
+        yield* Effect.uninterruptible(
+          registrationLock.withPermit(
+            Effect.gen(function* () {
+              const token = {}
+              for (const { entries, codemode } of planned)
+                for (const entry of entries)
+                  local.set(entry.key, [
+                    ...(local.get(entry.key) ?? []),
+                    {
+                      token,
+                      registration: {
+                        tool: entry.tool,
+                        name: entry.name,
+                        namespace: entry.namespace,
+                        codemode,
+                      },
+                    },
+                  ])
+              yield* Effect.addFinalizer(() =>
+                registrationLock.withPermit(
+                  Effect.sync(() => {
+                    for (const { entries } of planned)
+                      for (const entry of entries) {
+                        const registrations =
+                          local.get(entry.key)?.filter((registration) => registration.token !== token) ?? []
+                        if (registrations.length > 0) local.set(entry.key, registrations)
+                        else local.delete(entry.key)
+                      }
+                  }),
+                ),
+              )
+            }),
+          ),
+        )
+      },
+    )
+
+    return Service.of({
+      register: Effect.fn("ToolRegistry.register")((tools, options) =>
+        registerBatch([
+          {
+            tools,
+            ...(options === undefined ? {} : { options }),
           },
-        }
-      }),
+        ]),
+      ),
+      registerBatch,
+      materialize: Effect.fn("ToolRegistry.materialize")((permissions) =>
+        registrationLock.withPermit(
+          Effect.sync(() => {
+            const direct = new Map<string, Registration>()
+            const codemode = new Map<string, Registration>()
+            const rules = permissions ?? []
+            for (const [name, entries] of local) {
+              const registration = entries.at(-1)?.registration
+              if (!registration) continue
+              if (whollyDisabled(permission(registration.tool, name), rules)) continue
+              if (registration.codemode) codemode.set(name, registration)
+              else direct.set(name, registration)
+            }
+            const execute =
+              codemode.size > 0 && !whollyDisabled("execute", rules) ? ExecuteTool.create(codemode) : undefined
+            return {
+              definitions: [
+                ...Array.from(direct, ([name, registration]) => definition(name, registration.tool)),
+                ...(execute ? [definition("execute", execute)] : []),
+              ],
+              settle: (input: ExecuteInput) => {
+                if (input.call.name === "execute" && execute) return settleTool(input, execute)
+                const registration = direct.get(input.call.name)
+                if (registration) return settleTool(input, registration.tool)
+                return Effect.succeed({
+                  result: { type: "error", value: `Unknown tool: ${input.call.name}` },
+                  error: { type: "tool.unknown", message: `Unknown tool: ${input.call.name}` },
+                })
+              },
+            }
+          }),
+        ),
+      ),
     })
   }),
 )
 
 const layer = Layer.effect(
   Tools.Service,
-  Service.use((registry) => Effect.succeed(Tools.Service.of({ register: registry.register }))),
+  Service.use((registry) =>
+    Effect.succeed(Tools.Service.of({ register: registry.register, registerBatch: registry.registerBatch })),
+  ),
 ).pipe(Layer.provideMerge(registryLayer))
 
 function whollyDisabled(action: string, rules: PermissionV2.Ruleset) {

packages/core/src/tool/tools.ts

@@ -10,6 +10,13 @@ export interface Interface {
     tools: Readonly<Record<string, Tool.AnyTool>>,
     options?: Tool.RegisterOptions,
   ) => Effect.Effect<void, Tool.RegistrationError, Scope.Scope>
+  /** Internal atomic registration capability used by plugin transforms. */
+  readonly registerBatch: (
+    registrations: ReadonlyArray<{
+      readonly tools: Readonly<Record<string, Tool.AnyTool>>
+      readonly options?: Tool.RegisterOptions
+    }>,
+  ) => Effect.Effect<void, Tool.RegistrationError, Scope.Scope>
 }
 
 /** Narrow registration-only Location capability. */