refactor(effect): resolve built tools through the registry (#20787)

This commit is contained in:
Kit Langton
2026-04-03 10:31:00 -04:00
committed by GitHub
parent fbfa148e4e
commit 7994dce0f2
10 changed files with 224 additions and 147 deletions

View File

@@ -1,4 +1,5 @@
import z from "zod"
import { Effect } from "effect"
import type { MessageV2 } from "../session/message-v2"
import type { Agent } from "../agent/agent"
import type { Permission } from "../permission"
@@ -45,48 +46,67 @@ export namespace Tool {
init: (ctx?: InitContext) => Promise<Def<Parameters, M>>
}
export type InferParameters<T extends Info> = T extends Info<infer P> ? z.infer<P> : never
export type InferMetadata<T extends Info> = T extends Info<any, infer M> ? M : never
export type InferParameters<T> =
T extends Info<infer P, any>
? z.infer<P>
: T extends Effect.Effect<Info<infer P, any>, any, any>
? z.infer<P>
: never
export type InferMetadata<T> =
T extends Info<any, infer M> ? M : T extends Effect.Effect<Info<any, infer M>, any, any> ? M : never
function wrap<Parameters extends z.ZodType, Result extends Metadata>(
id: string,
init: ((ctx?: InitContext) => Promise<Def<Parameters, Result>>) | Def<Parameters, Result>,
) {
return async (initCtx?: InitContext) => {
const toolInfo = init instanceof Function ? await init(initCtx) : { ...init }
const execute = toolInfo.execute
toolInfo.execute = async (args, ctx) => {
try {
toolInfo.parameters.parse(args)
} catch (error) {
if (error instanceof z.ZodError && toolInfo.formatValidationError) {
throw new Error(toolInfo.formatValidationError(error), { cause: error })
}
throw new Error(
`The ${id} tool was called with invalid arguments: ${error}.\nPlease rewrite the input so it satisfies the expected schema.`,
{ cause: error },
)
}
const result = await execute(args, ctx)
if (result.metadata.truncated !== undefined) {
return result
}
const truncated = await Truncate.output(result.output, {}, initCtx?.agent)
return {
...result,
output: truncated.content,
metadata: {
...result.metadata,
truncated: truncated.truncated,
...(truncated.truncated && { outputPath: truncated.outputPath }),
},
}
}
return toolInfo
}
}
export function define<Parameters extends z.ZodType, Result extends Metadata>(
id: string,
init: Info<Parameters, Result>["init"] | Def<Parameters, Result>,
init: ((ctx?: InitContext) => Promise<Def<Parameters, Result>>) | Def<Parameters, Result>,
): Info<Parameters, Result> {
return {
id,
init: async (initCtx) => {
const toolInfo = init instanceof Function ? await init(initCtx) : { ...init }
const execute = toolInfo.execute
toolInfo.execute = async (args, ctx) => {
try {
toolInfo.parameters.parse(args)
} catch (error) {
if (error instanceof z.ZodError && toolInfo.formatValidationError) {
throw new Error(toolInfo.formatValidationError(error), { cause: error })
}
throw new Error(
`The ${id} tool was called with invalid arguments: ${error}.\nPlease rewrite the input so it satisfies the expected schema.`,
{ cause: error },
)
}
const result = await execute(args, ctx)
// skip truncation for tools that handle it themselves
if (result.metadata.truncated !== undefined) {
return result
}
const truncated = await Truncate.output(result.output, {}, initCtx?.agent)
return {
...result,
output: truncated.content,
metadata: {
...result.metadata,
truncated: truncated.truncated,
...(truncated.truncated && { outputPath: truncated.outputPath }),
},
}
}
return toolInfo
},
init: wrap(id, init),
}
}
export function defineEffect<Parameters extends z.ZodType, Result extends Metadata, R>(
id: string,
init: Effect.Effect<((ctx?: InitContext) => Promise<Def<Parameters, Result>>) | Def<Parameters, Result>, never, R>,
): Effect.Effect<Info<Parameters, Result>, never, R> {
return Effect.map(init, (next) => ({ id, init: wrap(id, next) }))
}
}