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

fix(core): sanitize registered tool names (#34512)

Aiden Cline 1 месяц назад
Родитель
Сommit
f928b5be07

+ 1 - 2
packages/core/src/tool/application-tools.ts

@@ -41,9 +41,8 @@ export const layer = Layer.effect(
 
     return Service.of({
       register: Effect.fn("ApplicationTools.register")(function* (tools) {
-        const entries = Object.entries(tools)
+        const entries = Tool.registrationEntries(tools)
         if (entries.length === 0) return
-        yield* Effect.forEach(entries, ([name]) => Tool.validateName(name), { discard: true })
         const registrations = entries.map(([name, tool]) => [name, { identity: {}, tool }] as const)
         yield* state.transform((draft) => {
           for (const [name, entry] of registrations) draft.set(name, entry)

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

@@ -9,7 +9,7 @@ import { SessionSchema } from "../session/schema"
 import { ToolOutputStore } from "../tool-output-store"
 import { Wildcard } from "../util/wildcard"
 import { ApplicationTools } from "./application-tools"
-import { definition, permission, settle, validateName, type AnyTool, type RegistrationError } from "./tool"
+import { definition, permission, registrationEntries, settle, type AnyTool, type RegistrationError } from "./tool"
 import { Tools } from "./tools"
 import { makeLocationNode } from "../effect/app-node"
 
@@ -83,9 +83,8 @@ const registryLayer = Layer.effect(
 
     return Service.of({
       register: Effect.fn("ToolRegistry.register")(function* (tools) {
-        const entries = Object.entries(tools)
+        const entries = registrationEntries(tools)
         if (entries.length === 0) return
-        yield* Effect.forEach(entries, ([name]) => validateName(name), { discard: true })
         yield* Effect.uninterruptible(
           Effect.gen(function* () {
             const token = {}

+ 3 - 0
packages/core/src/tool/tool.ts

@@ -136,6 +136,9 @@ export const validateName = (name: string) =>
     ? Effect.void
     : Effect.fail(new RegistrationError({ name, message: `Invalid tool name: ${name}` }))
 
+export const registrationEntries = (tools: Readonly<Record<string, AnyTool>>) =>
+  Object.entries(tools).map(([name, tool]) => [name.replace(/[^a-zA-Z0-9_-]/g, "_"), tool] as const)
+
 export const withPermission = <Input extends SchemaType<any>, Output extends SchemaType<any>>(
   tool: Definition<Input, Output>,
   permission: string,

+ 3 - 6
packages/core/test/application-tools.test.ts

@@ -70,17 +70,14 @@ describe("ApplicationTools", () => {
     }),
   )
 
-  it.effect("exposes narrow scoped Location registration and validates names", () =>
+  it.effect("exposes narrow scoped Location registration and sanitizes names", () =>
     Effect.gen(function* () {
       const tools: Tools.Interface = yield* Tools.Service
       const registry = yield* ToolRegistry.Service
       const scope = yield* Scope.make()
 
-      yield* tools.register({ location_tool: contextual([]) }).pipe(Scope.provide(scope))
-      expect((yield* toolDefinitions(registry)).map((tool) => tool.name)).toEqual(["location_tool"])
-      expect(yield* Effect.flip(tools.register({ "invalid name": contextual([]) }))).toBeInstanceOf(
-        Tool.RegistrationError,
-      )
+      yield* tools.register({ "location.tool/search": contextual([]) }).pipe(Scope.provide(scope))
+      expect((yield* toolDefinitions(registry)).map((tool) => tool.name)).toEqual(["location_tool_search"])
 
       yield* Scope.close(scope, Exit.void)
       expect(yield* toolDefinitions(registry)).toEqual([])