Просмотр исходного кода

feat(core): add simulated tool execution seam

Kit Langton 1 месяц назад
Родитель
Сommit
0918b7d805

+ 3 - 2
packages/core/src/tool/execute.ts

@@ -3,7 +3,8 @@ export * as ExecuteTool from "./execute"
 import { CodeMode, Tool, toolError } from "@opencode-ai/codemode"
 import { CodeMode, Tool, toolError } from "@opencode-ai/codemode"
 import { ToolOutput } from "@opencode-ai/llm"
 import { ToolOutput } from "@opencode-ai/llm"
 import { Effect, Ref, Schema } from "effect"
 import { Effect, Ref, Schema } from "effect"
-import { definition, make, settle, type AnyTool } from "./tool"
+import { definition, make, type AnyTool } from "./tool"
+import { ToolExecutionSimulation } from "./execution-simulation"
 
 
 const ExecuteFile = Schema.Struct({
 const ExecuteFile = Schema.Struct({
   data: Schema.String,
   data: Schema.String,
@@ -111,7 +112,7 @@ export const create = (registrations: ReadonlyMap<string, Registration>) => {
           (name, registration, input) =>
           (name, registration, input) =>
             Effect.gen(function* () {
             Effect.gen(function* () {
               const index = yield* Ref.getAndUpdate(callIndex, (index) => index + 1)
               const index = yield* Ref.getAndUpdate(callIndex, (index) => index + 1)
-              const output = yield* settle(
+              const output = yield* ToolExecutionSimulation.settle(
                 registration.tool,
                 registration.tool,
                 { type: "tool-call", id: context.callID, name, input },
                 { type: "tool-call", id: context.callID, name, input },
                 {
                 {

+ 43 - 0
packages/core/src/tool/execution-simulation.ts

@@ -0,0 +1,43 @@
+export * as ToolExecutionSimulation from "./execution-simulation"
+
+import type { ToolCall } from "@opencode-ai/llm"
+import { Effect } from "effect"
+import { Tool } from "./tool"
+
+export interface Invocation {
+  readonly tool: string
+  readonly input: unknown
+  readonly context: Tool.Context
+}
+
+export type Handler = (
+  invocation: Invocation,
+  passthrough: Effect.Effect<Tool.SettledOutput, Tool.Failure>,
+) => Effect.Effect<Tool.SettledOutput, Tool.Failure>
+
+const handlers = new Map<string, ReadonlyArray<{ readonly token: object; readonly handler: Handler }>>()
+
+export const intercept = (tool: string, handler: Handler) =>
+  Effect.acquireRelease(
+    Effect.sync(() => {
+      const token = {}
+      handlers.set(tool, [...(handlers.get(tool) ?? []), { token, handler }])
+      return token
+    }),
+    (token) =>
+      Effect.sync(() => {
+        const remaining = handlers.get(tool)?.filter((entry) => entry.token !== token) ?? []
+        if (remaining.length > 0) handlers.set(tool, remaining)
+        else handlers.delete(tool)
+      }),
+  ).pipe(Effect.asVoid)
+
+export const settle = (tool: Tool.AnyTool, call: ToolCall, context: Tool.Context) => {
+  const handler = handlers.get(call.name)?.at(-1)?.handler
+  return Tool.settle(
+    tool,
+    call,
+    context,
+    handler ? (input, passthrough) => handler({ tool: call.name, input, context }, passthrough) : undefined,
+  )
+}

+ 25 - 22
packages/core/src/tool/registry.ts

@@ -9,9 +9,11 @@ import { SessionSchema } from "../session/schema"
 import { ToolOutputStore } from "../tool-output-store"
 import { ToolOutputStore } from "../tool-output-store"
 import { Wildcard } from "../util/wildcard"
 import { Wildcard } from "../util/wildcard"
 import { ExecuteTool } from "./execute"
 import { ExecuteTool } from "./execute"
-import { definition, permission, registrationEntries, RegistrationError, settle, type AnyTool } from "./tool"
+import { Tool } from "./tool"
+import { definition, permission, registrationEntries, RegistrationError, type AnyTool } from "./tool"
 import { Tools } from "./tools"
 import { Tools } from "./tools"
 import { ToolHooks } from "./hooks"
 import { ToolHooks } from "./hooks"
+import { ToolExecutionSimulation } from "./execution-simulation"
 import { makeLocationNode } from "../effect/app-node"
 import { makeLocationNode } from "../effect/app-node"
 import { SessionError } from "@opencode-ai/schema/session-error"
 import { SessionError } from "@opencode-ai/schema/session-error"
 import { toSessionError } from "../session/to-session-error"
 import { toSessionError } from "../session/to-session-error"
@@ -76,29 +78,30 @@ const registryLayer = Layer.effect(
         input: input.call.input,
         input: input.call.input,
       }
       }
       yield* toolHooks.runBefore(beforeEvent)
       yield* toolHooks.runBefore(beforeEvent)
-      const pending = yield* settle(
+      const context: Tool.Context = {
+        sessionID: input.sessionID,
+        agent: input.agent,
+        messageID: input.messageID,
+        callID: input.call.id,
+        progress: (update) =>
+          input.progress?.({
+            structured: update.structured,
+            content: (update.content ?? []).map((part) =>
+              part.type === "text"
+                ? { type: "text" as const, text: part.text }
+                : {
+                    type: "file" as const,
+                    uri: `data:${part.mime};base64,${part.data}`,
+                    mime: part.mime,
+                    name: part.name,
+                  },
+            ),
+          }) ?? Effect.void,
+      }
+      const pending = yield* ToolExecutionSimulation.settle(
         tool,
         tool,
         { ...input.call, input: beforeEvent.input },
         { ...input.call, input: beforeEvent.input },
-        {
-          sessionID: input.sessionID,
-          agent: input.agent,
-          messageID: input.messageID,
-          callID: input.call.id,
-          progress: (update) =>
-            input.progress?.({
-              structured: update.structured,
-              content: (update.content ?? []).map((part) =>
-                part.type === "text"
-                  ? { type: "text" as const, text: part.text }
-                  : {
-                      type: "file" as const,
-                      uri: `data:${part.mime};base64,${part.data}`,
-                      mime: part.mime,
-                      name: part.name,
-                    },
-              ),
-            }) ?? Effect.void,
-        },
+        context,
       ).pipe(
       ).pipe(
         Effect.map((output) => ({ output })),
         Effect.map((output) => ({ output })),
         Effect.catchTag("LLM.ToolFailure", (failure) =>
         Effect.catchTag("LLM.ToolFailure", (failure) =>

+ 105 - 0
packages/core/test/session-runner-tool-registry.test.ts

@@ -7,6 +7,7 @@ import { SessionV2 } from "@opencode-ai/core/session"
 import { SessionMessage } from "@opencode-ai/core/session/message"
 import { SessionMessage } from "@opencode-ai/core/session/message"
 import { ToolOutputStore } from "@opencode-ai/core/tool-output-store"
 import { ToolOutputStore } from "@opencode-ai/core/tool-output-store"
 import { ToolRegistry } from "@opencode-ai/core/tool/registry"
 import { ToolRegistry } from "@opencode-ai/core/tool/registry"
+import { ToolExecutionSimulation } from "@opencode-ai/core/tool/execution-simulation"
 import { executeTool, settleTool, toolDefinitions } from "./lib/tool"
 import { executeTool, settleTool, toolDefinitions } from "./lib/tool"
 import { Cause, Deferred, Effect, Exit, Fiber, Layer, Option, Schema, SchemaGetter, SchemaIssue, Scope } from "effect"
 import { Cause, Deferred, Effect, Exit, Fiber, Layer, Option, Schema, SchemaGetter, SchemaIssue, Scope } from "effect"
 import { testEffect } from "./lib/effect"
 import { testEffect } from "./lib/effect"
@@ -62,6 +63,91 @@ const constant = (text: string) =>
   })
   })
 
 
 describe("ToolRegistry", () => {
 describe("ToolRegistry", () => {
+  it.effect("intercepts validated input and passes unhandled tools through", () =>
+    Effect.gen(function* () {
+      const service = yield* ToolRegistry.Service
+      const inputs: unknown[] = []
+      let executions = 0
+      const probe = Tool.make({
+        description: "Probe execution",
+        input: Schema.Struct({
+          text: Schema.String,
+          format: Schema.Literals(["text", "markdown"]).pipe(
+            Schema.withDecodingDefault(Effect.succeed("markdown" as const)),
+          ),
+        }),
+        output: Schema.Struct({ text: Schema.String }),
+        execute: ({ text }) =>
+          Effect.sync(() => {
+            executions++
+            return { text }
+          }),
+        toModelOutput: ({ output }) => [{ type: "text", text: output.text }],
+      })
+      yield* service.register({ intercepted: probe, passthrough: probe }, { codemode: false })
+      const simulationScope = yield* Scope.fork(yield* Scope.Scope)
+      yield* ToolExecutionSimulation.intercept("intercepted", ({ input }) =>
+        Effect.sync(() => {
+          inputs.push(input)
+          return {
+            structured: { text: "simulated" },
+            content: [{ type: "text", text: "simulated" }],
+          }
+        }),
+      ).pipe(Scope.provide(simulationScope))
+      const materialized = yield* service.materialize()
+
+      const intercepted = yield* materialized.settle({
+        sessionID,
+        ...identity,
+        call: {
+          type: "tool-call",
+          id: "call-intercepted",
+          name: "intercepted",
+          input: { text: "ignored" },
+        },
+      })
+      const passed = yield* materialized.settle({
+        sessionID,
+        ...identity,
+        call: {
+          type: "tool-call",
+          id: "call-passthrough",
+          name: "passthrough",
+          input: { text: "real", format: "text" },
+        },
+      })
+      const invalid = yield* materialized.settle({
+        sessionID,
+        ...identity,
+        call: {
+          type: "tool-call",
+          id: "call-invalid",
+          name: "intercepted",
+          input: {},
+        },
+      })
+      yield* Scope.close(simulationScope, Exit.void)
+      const restored = yield* materialized.settle({
+        sessionID,
+        ...identity,
+        call: {
+          type: "tool-call",
+          id: "call-restored",
+          name: "intercepted",
+          input: { text: "restored", format: "text" },
+        },
+      })
+
+      expect(inputs).toEqual([{ text: "ignored", format: "markdown" }])
+      expect(intercepted.result).toEqual({ type: "text", value: "simulated" })
+      expect(passed.result).toEqual({ type: "text", value: "real" })
+      expect(invalid.result).toMatchObject({ type: "error", value: expect.stringContaining("Invalid tool input") })
+      expect(restored.result).toEqual({ type: "text", value: "restored" })
+      expect(executions).toBe(2)
+    }),
+  )
+
   it.effect("filters disabled tools with edit aliases and ordered wildcard precedence", () =>
   it.effect("filters disabled tools with edit aliases and ordered wildcard precedence", () =>
     Effect.gen(function* () {
     Effect.gen(function* () {
       const service = yield* ToolRegistry.Service
       const service = yield* ToolRegistry.Service
@@ -384,6 +470,13 @@ describe("ToolRegistry", () => {
           execute: ({ text }) => Effect.sync(() => executed.push(`new:${text}`)).pipe(Effect.as({ text })),
           execute: ({ text }) => Effect.sync(() => executed.push(`new:${text}`)).pipe(Effect.as({ text })),
         }),
         }),
       })
       })
+      const simulationScope = yield* Scope.fork(yield* Scope.Scope)
+      yield* ToolExecutionSimulation.intercept("echo", () =>
+        Effect.succeed({
+          structured: { text: "simulated" },
+          content: [{ type: "text", text: "simulated" }],
+        }),
+      ).pipe(Scope.provide(simulationScope))
 
 
       const settlement = yield* materialized.settle({
       const settlement = yield* materialized.settle({
         ...call("execute"),
         ...call("execute"),
@@ -396,6 +489,18 @@ describe("ToolRegistry", () => {
       })
       })
 
 
       expect(settlement.result).toMatchObject({ type: "text" })
       expect(settlement.result).toMatchObject({ type: "text" })
+      expect(executed).toEqual([])
+
+      yield* Scope.close(simulationScope, Exit.void)
+      yield* materialized.settle({
+        ...call("execute"),
+        call: {
+          type: "tool-call",
+          id: "call-execute-restored",
+          name: "execute",
+          input: { code: 'return await tools.echo({ text: "request" })' },
+        },
+      })
       expect(executed).toEqual(["old:request"])
       expect(executed).toEqual(["old:request"])
     }),
     }),
   )
   )

+ 26 - 14
packages/plugin/src/v2/effect/tool.ts

@@ -4,7 +4,7 @@ import { Agent } from "@opencode-ai/schema/agent"
 import type { LLM } from "@opencode-ai/schema/llm"
 import type { LLM } from "@opencode-ai/schema/llm"
 import { Session } from "@opencode-ai/schema/session"
 import { Session } from "@opencode-ai/schema/session"
 import { SessionMessage } from "@opencode-ai/schema/session-message"
 import { SessionMessage } from "@opencode-ai/schema/session-message"
-import { Effect, JsonSchema, Schema, type Scope } from "effect"
+import { Effect, JsonSchema, Schema } from "effect"
 import type { Hooks, Transform } from "./registration.js"
 import type { Hooks, Transform } from "./registration.js"
 
 
 export interface Context {
 export interface Context {
@@ -40,11 +40,16 @@ type ToolResultValue =
   | { readonly type: "error"; readonly value: unknown }
   | { readonly type: "error"; readonly value: unknown }
   | { readonly type: "content"; readonly value: ReadonlyArray<LLM.ToolContent> }
   | { readonly type: "content"; readonly value: ReadonlyArray<LLM.ToolContent> }
 
 
-type ToolOutput = {
+export type SettledOutput = {
   readonly structured: unknown
   readonly structured: unknown
   readonly content: ReadonlyArray<LLM.ToolContent>
   readonly content: ReadonlyArray<LLM.ToolContent>
 }
 }
 
 
+export type SettlementInterceptor = (
+  input: unknown,
+  passthrough: Effect.Effect<SettledOutput, Failure>,
+) => Effect.Effect<SettledOutput, Failure>
+
 declare const TypeId: unique symbol
 declare const TypeId: unique symbol
 
 
 export interface Definition<Input extends SchemaType<any>, Output extends SchemaType<any>> {
 export interface Definition<Input extends SchemaType<any>, Output extends SchemaType<any>> {
@@ -113,7 +118,11 @@ type DynamicConfig = {
 type Runtime = {
 type Runtime = {
   readonly permission?: string
   readonly permission?: string
   readonly definition: (name: string) => ToolDefinition
   readonly definition: (name: string) => ToolDefinition
-  readonly settle: (call: ToolCall, context: Context) => Effect.Effect<ToolOutput, Failure>
+  readonly settle: (
+    call: ToolCall,
+    context: Context,
+    interceptor?: SettlementInterceptor,
+  ) => Effect.Effect<SettledOutput, Failure>
 }
 }
 
 
 const runtimes = new WeakMap<AnyTool, Runtime>()
 const runtimes = new WeakMap<AnyTool, Runtime>()
@@ -149,11 +158,11 @@ function makeTyped<
       definitions.set(name, definition)
       definitions.set(name, definition)
       return definition
       return definition
     },
     },
-    settle: (call, context) =>
+    settle: (call, context, interceptor) =>
       Schema.decodeUnknownEffect(config.input)(call.input).pipe(
       Schema.decodeUnknownEffect(config.input)(call.input).pipe(
         Effect.mapError((error) => new Failure({ message: `Invalid tool input: ${error.message}` })),
         Effect.mapError((error) => new Failure({ message: `Invalid tool input: ${error.message}` })),
-        Effect.flatMap((input) =>
-          config.execute(input, context).pipe(
+        Effect.flatMap((input) => {
+          const passthrough = Effect.suspend(() => config.execute(input, context)).pipe(
             Effect.flatMap((output) =>
             Effect.flatMap((output) =>
               Schema.encodeEffect(config.output)(output).pipe(
               Schema.encodeEffect(config.output)(output).pipe(
                 Effect.flatMap((output) => {
                 Effect.flatMap((output) => {
@@ -177,8 +186,9 @@ function makeTyped<
                 config.toModelOutput?.({ input, output }).map(toModelContent) ??
                 config.toModelOutput?.({ input, output }).map(toModelContent) ??
                 (typeof output === "string" ? [{ type: "text" as const, text: output }] : []),
                 (typeof output === "string" ? [{ type: "text" as const, text: output }] : []),
             })),
             })),
-          ),
-        ),
+          )
+          return interceptor?.(input, passthrough) ?? passthrough
+        }),
       ),
       ),
   })
   })
   return tool
   return tool
@@ -200,10 +210,11 @@ function makeDynamic(config: DynamicConfig): AnyTool {
       definitions.set(name, definition)
       definitions.set(name, definition)
       return definition
       return definition
     },
     },
-    settle: (call, context) =>
-      config
-        .execute(call.input, context)
-        .pipe(Effect.map((output) => ({ structured: output.structured, content: output.content.map(toModelContent) }))),
+    settle: (call, context, interceptor) => {
+      const passthrough = Effect.suspend(() => config.execute(call.input, context))
+        .pipe(Effect.map((output) => ({ structured: output.structured, content: output.content.map(toModelContent) })))
+      return interceptor?.(call.input, passthrough) ?? passthrough
+    },
   })
   })
   return tool
   return tool
 }
 }
@@ -241,7 +252,8 @@ export const withPermission = <Input extends SchemaType<any>, Output extends Sch
 
 
 export const permission = (tool: AnyTool, name: string) => runtimeOf(tool).permission ?? name
 export const permission = (tool: AnyTool, name: string) => runtimeOf(tool).permission ?? name
 export const definition = (name: string, tool: AnyTool) => runtimeOf(tool).definition(name)
 export const definition = (name: string, tool: AnyTool) => runtimeOf(tool).definition(name)
-export const settle = (tool: AnyTool, call: ToolCall, context: Context) => runtimeOf(tool).settle(call, context)
+export const settle = (tool: AnyTool, call: ToolCall, context: Context, interceptor?: SettlementInterceptor) =>
+  runtimeOf(tool).settle(call, context, interceptor)
 
 
 function runtimeOf(tool: AnyTool) {
 function runtimeOf(tool: AnyTool) {
   const runtime = runtimes.get(tool)
   const runtime = runtimes.get(tool)
@@ -272,7 +284,7 @@ export interface ToolExecuteAfterEvent {
   readonly callID: string
   readonly callID: string
   readonly input: unknown
   readonly input: unknown
   result: ToolResultValue
   result: ToolResultValue
-  output?: ToolOutput
+  output?: SettledOutput
   outputPaths?: ReadonlyArray<string>
   outputPaths?: ReadonlyArray<string>
 }
 }