registry.ts 6.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139
  1. export * as ToolRegistry from "./registry"
  2. import { ToolOutput, type ToolCall, type ToolDefinition, type ToolResultValue } from "@opencode-ai/llm"
  3. import { Context, Effect, Layer, Scope } from "effect"
  4. import { AgentV2 } from "../agent"
  5. import { PermissionV2 } from "../permission"
  6. import { SessionMessage } from "../session/message"
  7. import { SessionSchema } from "../session/schema"
  8. import { ToolOutputStore } from "../tool-output-store"
  9. import { Wildcard } from "../util/wildcard"
  10. import { ApplicationTools } from "./application-tools"
  11. import { definition, permission, settle, validateName, type AnyTool, type RegistrationError } from "./tool"
  12. import { Tools } from "./tools"
  13. export type ExecuteInput = {
  14. readonly sessionID: SessionSchema.ID
  15. readonly agent: AgentV2.ID
  16. readonly assistantMessageID: SessionMessage.ID
  17. readonly call: ToolCall
  18. }
  19. export interface Interface {
  20. readonly materialize: (permissions?: PermissionV2.Ruleset) => Effect.Effect<Materialization>
  21. /** Internal registration capability exposed publicly only through Tools.Service. */
  22. readonly register: (tools: Readonly<Record<string, AnyTool>>) => Effect.Effect<void, RegistrationError, Scope.Scope>
  23. }
  24. export interface Materialization {
  25. readonly definitions: ReadonlyArray<ToolDefinition>
  26. readonly settle: (input: ExecuteInput) => Effect.Effect<Settlement, ToolOutputStore.Error>
  27. }
  28. export interface Settlement {
  29. readonly result: ToolResultValue
  30. readonly output?: ToolOutput
  31. readonly outputPaths?: ReadonlyArray<string>
  32. }
  33. export class Service extends Context.Service<Service, Interface>()("@opencode/v2/ToolRegistry") {}
  34. const registryLayer = Layer.effect(
  35. Service,
  36. Effect.gen(function* () {
  37. const applications = yield* ApplicationTools.Service
  38. const resources = yield* ToolOutputStore.Service
  39. type Registration = { readonly identity: object; readonly tool: AnyTool }
  40. const local = new Map<string, Array<{ readonly token: object; readonly registration: Registration }>>()
  41. const settleWith = Effect.fn("ToolRegistry.settle")(function* (input: ExecuteInput, advertised?: object) {
  42. const registration =
  43. local.get(input.call.name)?.at(-1)?.registration ?? applications.entries().get(input.call.name)
  44. if (!registration)
  45. return {
  46. result: {
  47. type: "error" as const,
  48. value: advertised ? `Stale tool call: ${input.call.name}` : `Unknown tool: ${input.call.name}`,
  49. },
  50. }
  51. if (advertised && registration.identity !== advertised)
  52. return { result: { type: "error" as const, value: `Stale tool call: ${input.call.name}` } }
  53. const pending = yield* settle(registration.tool, input.call, {
  54. sessionID: input.sessionID,
  55. agent: input.agent,
  56. assistantMessageID: input.assistantMessageID,
  57. toolCallID: input.call.id,
  58. }).pipe(
  59. Effect.map((output) => ({ output })),
  60. Effect.catchTag("LLM.ToolFailure", (failure) =>
  61. Effect.succeed({ result: { type: "error" as const, value: failure.message } }),
  62. ),
  63. )
  64. if ("result" in pending) return pending
  65. const output = pending.output
  66. const bounded = yield* resources.bound({ sessionID: input.sessionID, toolCallID: input.call.id, output })
  67. const result = ToolOutput.toResultValue(bounded.output)
  68. if (result.type === "error")
  69. return bounded.outputPaths.length > 0 ? { result, outputPaths: bounded.outputPaths } : { result }
  70. return bounded.outputPaths.length > 0
  71. ? { result, output: bounded.output, outputPaths: bounded.outputPaths }
  72. : { result, output: bounded.output }
  73. })
  74. return Service.of({
  75. register: Effect.fn("ToolRegistry.register")(function* (tools) {
  76. const entries = Object.entries(tools)
  77. if (entries.length === 0) return
  78. yield* Effect.forEach(entries, ([name]) => validateName(name), { discard: true })
  79. yield* Effect.uninterruptible(
  80. Effect.gen(function* () {
  81. const token = {}
  82. for (const [name, tool] of entries)
  83. local.set(name, [...(local.get(name) ?? []), { token, registration: { identity: {}, tool } }])
  84. yield* Effect.addFinalizer(() =>
  85. Effect.sync(() => {
  86. for (const [name] of entries) {
  87. const registrations = local.get(name)?.filter((registration) => registration.token !== token) ?? []
  88. if (registrations.length > 0) local.set(name, registrations)
  89. else local.delete(name)
  90. }
  91. }),
  92. )
  93. }),
  94. )
  95. }),
  96. materialize: Effect.fn("ToolRegistry.materialize")(function* (permissions = []) {
  97. const registrations = new Map(applications.entries())
  98. for (const [name, entries] of local) {
  99. const registration = entries.at(-1)?.registration
  100. if (registration) registrations.set(name, registration)
  101. }
  102. for (const [name, registration] of registrations)
  103. if (whollyDisabled(permission(registration.tool, name), permissions)) registrations.delete(name)
  104. return {
  105. definitions: Array.from(registrations, ([name, registration]) => definition(name, registration.tool)),
  106. settle: (input) => {
  107. const registration = registrations.get(input.call.name)
  108. if (registration) return settleWith(input, registration.identity)
  109. return Effect.succeed({ result: { type: "error", value: `Unknown tool: ${input.call.name}` } })
  110. },
  111. }
  112. }),
  113. })
  114. }),
  115. )
  116. export const layer = Layer.effect(
  117. Tools.Service,
  118. Service.use((registry) => Effect.succeed(Tools.Service.of({ register: registry.register }))),
  119. ).pipe(Layer.provideMerge(registryLayer))
  120. function whollyDisabled(action: string, rules: PermissionV2.Ruleset) {
  121. const rule = rules.findLast((rule) => Wildcard.match(action, rule.action))
  122. return rule?.resource === "*" && rule.effect === "deny"
  123. }
  124. export const defaultLayer = layer.pipe(
  125. Layer.provide(ApplicationTools.layer),
  126. Layer.provide(ToolOutputStore.defaultLayer),
  127. )