refactor(core): unify v2 tool architecture (#31168)
This commit is contained in:
Kit Langton
@@ -1,91 +1,35 @@
|
||||
export * as ToolRegistry from "./registry"
|
||||
|
||||
import {
|
||||
Tool,
|
||||
ToolFailure,
|
||||
ToolOutput,
|
||||
ToolResultValue as ToolResult,
|
||||
type Tool as TypedTool,
|
||||
type ToolCall,
|
||||
type ToolResultValue,
|
||||
type ToolSchema,
|
||||
type ToolSettlement,
|
||||
} from "@opencode-ai/llm"
|
||||
import { Context, Effect, Layer, Schema, Scope } from "effect"
|
||||
import { castDraft, enableMapSet } from "immer"
|
||||
import { PermissionV2 } from "../permission"
|
||||
import { State } from "../state"
|
||||
import { SessionSchema } from "../session/schema"
|
||||
import type { SessionV2 } from "../session"
|
||||
import { ApplicationTools } from "./application-tools"
|
||||
import { ToolOutputStore } from "../tool-output-store"
|
||||
import { Tool as LlmTool, ToolOutput, type ToolCall, type ToolSettlement } from "@opencode-ai/llm"
|
||||
import { Context, Effect, Layer, Scope } from "effect"
|
||||
import { AgentV2 } from "../agent"
|
||||
import { PermissionV2 } from "../permission"
|
||||
import { SessionMessage } from "../session/message"
|
||||
import { SessionSchema } from "../session/schema"
|
||||
import { ToolOutputStore } from "../tool-output-store"
|
||||
import { Wildcard } from "../util/wildcard"
|
||||
import { ApplicationTools } from "./application-tools"
|
||||
import { Tool } from "./tool"
|
||||
import { Tools } from "./tools"
|
||||
|
||||
export type ExecuteInput = {
|
||||
readonly sessionID: SessionSchema.ID
|
||||
readonly agent?: AgentV2.ID
|
||||
readonly agent: AgentV2.ID
|
||||
readonly assistantMessageID: SessionMessage.ID
|
||||
readonly call: ToolCall
|
||||
}
|
||||
|
||||
/**
|
||||
* Narrow cross-cutting context for one registry invocation. Leaf tools retain
|
||||
* ownership of sequence-sensitive policy decisions; the registry only binds
|
||||
* identity and shared helper behavior consistently.
|
||||
*
|
||||
* TODO: Add `source` when the runner can pass the durable owning assistant
|
||||
* message ID alongside the call ID. Do not infer it from the tool call alone.
|
||||
* TODO: Add cancellation and progress only when the runner exposes a real
|
||||
* signal and durable/live progress sink.
|
||||
*/
|
||||
export type Invocation = ExecuteInput & {
|
||||
readonly source?: PermissionV2.Source
|
||||
readonly assertPermission: (
|
||||
input: Omit<PermissionV2.AssertInput, "sessionID" | "agent" | "source">,
|
||||
) => Effect.Effect<void, PermissionV2.Error | SessionV2.NotFoundError>
|
||||
}
|
||||
|
||||
/** Kept as the leaf entry input name for backwards-compatible execute usage. */
|
||||
export type AuthorizeInput<Parameters = unknown> = Invocation & {
|
||||
readonly parameters: Parameters
|
||||
}
|
||||
|
||||
export type Entry<
|
||||
Parameters extends ToolSchema<any> = ToolSchema<any>,
|
||||
Success extends ToolSchema<any> = ToolSchema<any>,
|
||||
> = {
|
||||
readonly tool: TypedTool<Parameters, Success>
|
||||
/** Catalog visibility only. Execution authorization remains leaf-owned. */
|
||||
readonly permission?: { readonly action: string; readonly resource: "*" }
|
||||
readonly authorize?: (input: AuthorizeInput<Schema.Schema.Type<Parameters>>) => Effect.Effect<void, ToolFailure>
|
||||
readonly execute?: (
|
||||
input: AuthorizeInput<Schema.Schema.Type<Parameters>>,
|
||||
) => Effect.Effect<Schema.Schema.Type<Success>, ToolFailure>
|
||||
readonly outputPaths?: (output: Schema.Schema.Type<Success>) => ReadonlyArray<string>
|
||||
}
|
||||
|
||||
type Data = {
|
||||
readonly entries: Map<string, Entry>
|
||||
}
|
||||
|
||||
export type Editor = {
|
||||
readonly list: () => ReadonlyArray<readonly [string, Entry]>
|
||||
readonly get: (name: string) => Entry | undefined
|
||||
readonly set: <Parameters extends ToolSchema<any>, Success extends ToolSchema<any>>(
|
||||
name: string,
|
||||
entry: Entry<Parameters, Success>,
|
||||
) => void
|
||||
readonly remove: (name: string) => void
|
||||
}
|
||||
|
||||
export interface Interface {
|
||||
readonly transform: State.Interface<Data, Editor>["transform"]
|
||||
readonly contribute: (update: State.Transform<Editor>) => Effect.Effect<void, never, Scope.Scope>
|
||||
readonly definitions: (
|
||||
permissions?: PermissionV2.Ruleset,
|
||||
) => Effect.Effect<ReadonlyArray<ReturnType<typeof Tool.toDefinitions>[number]>>
|
||||
readonly execute: (input: ExecuteInput) => Effect.Effect<ToolResultValue>
|
||||
readonly settle: (input: ExecuteInput) => Effect.Effect<Settlement>
|
||||
readonly materialize: (permissions?: PermissionV2.Ruleset) => Effect.Effect<Materialization>
|
||||
/** Internal registration capability exposed publicly only through Tools.Service. */
|
||||
readonly register: (
|
||||
tools: Readonly<Record<string, Tool.AnyTool>>,
|
||||
) => Effect.Effect<void, Tool.RegistrationError, Scope.Scope>
|
||||
}
|
||||
|
||||
export interface Materialization {
|
||||
readonly definitions: ReadonlyArray<ReturnType<typeof LlmTool.toDefinitions>[number]>
|
||||
readonly settle: (input: ExecuteInput) => Effect.Effect<Settlement, ToolOutputStore.Error>
|
||||
}
|
||||
|
||||
export interface Settlement extends ToolSettlement {
|
||||
@@ -94,153 +38,94 @@ export interface Settlement extends ToolSettlement {
|
||||
|
||||
export class Service extends Context.Service<Service, Interface>()("@opencode/v2/ToolRegistry") {}
|
||||
|
||||
enableMapSet()
|
||||
|
||||
export const layer = Layer.effect(
|
||||
const registryLayer = Layer.effect(
|
||||
Service,
|
||||
Effect.gen(function* () {
|
||||
const permission = yield* PermissionV2.Service
|
||||
const applications = yield* ApplicationTools.Service
|
||||
const resources = yield* ToolOutputStore.Service
|
||||
const state = State.create<Data, Editor>({
|
||||
initial: () => ({ entries: new Map() }),
|
||||
editor: (draft) => ({
|
||||
list: () => Array.from(draft.entries.entries()) as Array<[string, Entry]>,
|
||||
get: (name) => draft.entries.get(name) as Entry | undefined,
|
||||
set: (name, entry) => {
|
||||
draft.entries.set(
|
||||
name,
|
||||
castDraft(entry) as typeof draft.entries extends Map<string, infer Value> ? Value : never,
|
||||
)
|
||||
},
|
||||
remove: (name) => {
|
||||
draft.entries.delete(name)
|
||||
},
|
||||
}),
|
||||
})
|
||||
type Registration = { readonly identity: object; readonly tool: Tool.AnyTool }
|
||||
const local = new Map<string, Array<{ readonly token: object; readonly registration: Registration }>>()
|
||||
|
||||
const definitions = Effect.fn("ToolRegistry.definitions")(function* (permissions: PermissionV2.Ruleset = []) {
|
||||
const tools = new Map(state.get().entries)
|
||||
// Location tools own their names. Application tools fill otherwise-unclaimed names.
|
||||
for (const [name, tool] of applications.entries()) {
|
||||
if (!tools.has(name)) tools.set(name, { tool: tool.definition })
|
||||
}
|
||||
return Tool.toDefinitions(
|
||||
Object.fromEntries(
|
||||
Array.from(tools)
|
||||
.filter(([name, entry]) => !whollyDisabled(entry.permission ?? defaultPermission(name), permissions))
|
||||
.map(([name, entry]) => [name, entry.tool]),
|
||||
),
|
||||
)
|
||||
})
|
||||
|
||||
const entry = (name: string): Entry | undefined => {
|
||||
const local = state.get().entries.get(name)
|
||||
if (local !== undefined) return local
|
||||
const tool = applications.entries().get(name)
|
||||
if (tool === undefined) return
|
||||
return {
|
||||
tool: tool.definition,
|
||||
execute: ({ parameters, sessionID, call }) =>
|
||||
tool.execute(parameters, { sessionID, id: call.id, name: call.name }),
|
||||
}
|
||||
}
|
||||
|
||||
const invocation = (input: ExecuteInput): Invocation => ({
|
||||
...input,
|
||||
// Source needs the durable owning assistant message ID, which the registry does not receive yet.
|
||||
assertPermission: (request) =>
|
||||
permission.assert({ ...request, sessionID: input.sessionID, ...(input.agent ? { agent: input.agent } : {}) }),
|
||||
})
|
||||
|
||||
const settleEntry = Effect.fn("ToolRegistry.settleEntry")(function* (
|
||||
entry: Entry | undefined,
|
||||
input: ExecuteInput,
|
||||
) {
|
||||
if (!entry) return { result: { type: "error" as const, value: `Unknown tool: ${input.call.name}` } }
|
||||
if (!entry.execute && !entry.tool.execute)
|
||||
return { result: { type: "error" as const, value: `Tool has no execute handler: ${input.call.name}` } }
|
||||
|
||||
return yield* entry.tool._decode(input.call.input).pipe(
|
||||
Effect.mapError((error) => new ToolFailure({ message: `Invalid tool input: ${error.message}` })),
|
||||
Effect.flatMap((parameters) => {
|
||||
const context = { ...invocation(input), parameters }
|
||||
const execute =
|
||||
entry.execute?.(context) ?? entry.tool.execute!(parameters, { id: input.call.id, name: input.call.name })
|
||||
return (
|
||||
entry.authorize === undefined ? execute : entry.authorize(context).pipe(Effect.andThen(execute))
|
||||
).pipe(
|
||||
Effect.flatMap((value) =>
|
||||
entry.tool._encode(value).pipe(
|
||||
Effect.mapError(
|
||||
(error) =>
|
||||
new ToolFailure({
|
||||
message: `Tool returned an invalid value for its success schema: ${error.message}`,
|
||||
}),
|
||||
),
|
||||
),
|
||||
),
|
||||
Effect.map((value): Settlement => {
|
||||
const settled = (() => {
|
||||
if (entry.tool._legacyResult && ToolResult.is(value))
|
||||
return { result: value, output: ToolOutput.fromResultValue(value) }
|
||||
const output = entry.tool._project(parameters, input.call.id, value)
|
||||
const result = ToolOutput.toResultValue(output)
|
||||
return result.type === "error" ? { result } : { result, output }
|
||||
})()
|
||||
const retained = entry.outputPaths?.(value) ?? []
|
||||
return retained.length > 0 ? { ...settled, outputPaths: retained } : settled
|
||||
}),
|
||||
)
|
||||
}),
|
||||
const settleWith = Effect.fn("ToolRegistry.settle")(function* (input: ExecuteInput, advertised?: object) {
|
||||
const registration =
|
||||
local.get(input.call.name)?.at(-1)?.registration ?? applications.entries().get(input.call.name)
|
||||
if (!registration)
|
||||
return {
|
||||
result: {
|
||||
type: "error" as const,
|
||||
value: advertised ? `Stale tool call: ${input.call.name}` : `Unknown tool: ${input.call.name}`,
|
||||
},
|
||||
}
|
||||
if (advertised && registration.identity !== advertised)
|
||||
return { result: { type: "error" as const, value: `Stale tool call: ${input.call.name}` } }
|
||||
const pending = yield* Tool.settle(registration.tool, input.call, {
|
||||
sessionID: input.sessionID,
|
||||
agent: input.agent,
|
||||
assistantMessageID: input.assistantMessageID,
|
||||
toolCallID: input.call.id,
|
||||
}).pipe(
|
||||
Effect.map((output) => ({ output })),
|
||||
Effect.catchTag("LLM.ToolFailure", (failure) =>
|
||||
Effect.succeed({ result: { type: "error" as const, value: failure.message } }),
|
||||
),
|
||||
)
|
||||
})
|
||||
|
||||
const settle = Effect.fn("ToolRegistry.settle")((input: ExecuteInput) =>
|
||||
Effect.uninterruptibleMask((restore) =>
|
||||
Effect.gen(function* () {
|
||||
const settled = yield* restore(settleEntry(entry(input.call.name), input))
|
||||
if (!settled.output) return settled
|
||||
const bounded = yield* resources.bound({
|
||||
sessionID: input.sessionID,
|
||||
toolCallID: input.call.id,
|
||||
output: settled.output,
|
||||
})
|
||||
if (bounded.output === settled.output && bounded.outputPaths.length === 0) return settled
|
||||
const retained = [...(settled.outputPaths ?? []), ...bounded.outputPaths]
|
||||
const result = ToolOutput.toResultValue(bounded.output)
|
||||
return result.type === "error"
|
||||
? { result, outputPaths: retained }
|
||||
: { result, output: bounded.output, outputPaths: retained }
|
||||
}),
|
||||
),
|
||||
)
|
||||
const execute = Effect.fn("ToolRegistry.execute")(function* (input: ExecuteInput) {
|
||||
return (yield* settle(input)).result
|
||||
if ("result" in pending) return pending
|
||||
const output = pending.output
|
||||
const bounded = yield* resources.bound({ sessionID: input.sessionID, toolCallID: input.call.id, output })
|
||||
const result = ToolOutput.toResultValue(bounded.output)
|
||||
if (result.type === "error")
|
||||
return bounded.outputPaths.length > 0 ? { result, outputPaths: bounded.outputPaths } : { result }
|
||||
return bounded.outputPaths.length > 0
|
||||
? { result, output: bounded.output, outputPaths: bounded.outputPaths }
|
||||
: { result, output: bounded.output }
|
||||
})
|
||||
|
||||
return Service.of({
|
||||
transform: state.transform,
|
||||
contribute: Effect.fn("ToolRegistry.contribute")(function* (update) {
|
||||
const transform = yield* state.transform()
|
||||
yield* transform(update)
|
||||
register: Effect.fn("ToolRegistry.register")(function* (tools) {
|
||||
const entries = Object.entries(tools)
|
||||
if (entries.length === 0) return
|
||||
yield* Effect.forEach(entries, ([name]) => Tool.validateName(name), { discard: true })
|
||||
const token = {}
|
||||
for (const [name, tool] of entries)
|
||||
local.set(name, [...(local.get(name) ?? []), { token, registration: { identity: {}, tool } }])
|
||||
yield* Effect.addFinalizer(() =>
|
||||
Effect.sync(() => {
|
||||
for (const [name] of entries) {
|
||||
const registrations = local.get(name)?.filter((registration) => registration.token !== token) ?? []
|
||||
if (registrations.length > 0) local.set(name, registrations)
|
||||
else local.delete(name)
|
||||
}
|
||||
}),
|
||||
)
|
||||
}),
|
||||
materialize: Effect.fn("ToolRegistry.materialize")(function* (permissions = []) {
|
||||
const registrations = new Map(applications.entries())
|
||||
for (const [name, entries] of local) {
|
||||
const registration = entries.at(-1)?.registration
|
||||
if (registration) registrations.set(name, registration)
|
||||
}
|
||||
for (const [name, registration] of registrations)
|
||||
if (whollyDisabled(Tool.permission(registration.tool, name), permissions)) registrations.delete(name)
|
||||
return {
|
||||
definitions: Array.from(registrations, ([name, registration]) => Tool.definition(name, registration.tool)),
|
||||
settle: (input) => {
|
||||
const registration = registrations.get(input.call.name)
|
||||
if (registration) return settleWith(input, registration.identity)
|
||||
return Effect.succeed({ result: { type: "error", value: `Unknown tool: ${input.call.name}` } })
|
||||
},
|
||||
}
|
||||
}),
|
||||
definitions,
|
||||
execute,
|
||||
settle,
|
||||
})
|
||||
}),
|
||||
)
|
||||
|
||||
function defaultPermission(name: string) {
|
||||
return { action: ["edit", "write", "apply_patch"].includes(name) ? "edit" : name, resource: "*" as const }
|
||||
}
|
||||
export const layer = Layer.effect(
|
||||
Tools.Service,
|
||||
Service.use((registry) => Effect.succeed(Tools.Service.of({ register: registry.register }))),
|
||||
).pipe(Layer.provideMerge(registryLayer))
|
||||
|
||||
function whollyDisabled(permission: { readonly action: string; readonly resource: "*" }, rules: PermissionV2.Ruleset) {
|
||||
const rule = rules.findLast((rule) => Wildcard.match(permission.action, rule.action))
|
||||
function whollyDisabled(action: string, rules: PermissionV2.Ruleset) {
|
||||
const rule = rules.findLast((rule) => Wildcard.match(action, rule.action))
|
||||
return rule?.resource === "*" && rule.effect === "deny"
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user