refactor(tool): clean up native-friendly parameter schemas

This commit is contained in:
Kit Langton
2026-04-18 12:54:03 -04:00
parent a49b5adfbd
commit 29eba01658
7 changed files with 110 additions and 26 deletions

View File

@@ -9,9 +9,10 @@ export const Parameters = Schema.Struct({
description:
"Search query to find relevant context for APIs, Libraries, and SDKs. For example, 'React useState hook examples', 'Python pandas dataframe filtering', 'Express.js middleware', 'Next js partial prerendering configuration'",
}),
tokensNum: Schema.Number.check(Schema.isGreaterThanOrEqualTo(1000))
.check(Schema.isLessThanOrEqualTo(50000))
.pipe(Schema.optional, Schema.withDecodingDefault(Effect.succeed(5000)))
tokensNum: Schema.Finite.pipe(
Schema.check(Schema.isGreaterThanOrEqualTo(1000), Schema.isLessThanOrEqualTo(50000)),
Schema.withDecodingDefaultTypeKey(Effect.succeed(5000)),
)
.annotate({
description:
"Number of tokens to return (1000-50000). Default is 5000 tokens. Adjust this value based on how much context you need - use lower values for focused queries and higher values for comprehensive documentation.",
@@ -26,7 +27,7 @@ export const CodeSearchTool = Tool.define(
return {
description: DESCRIPTION,
parameters: Parameters,
execute: (params: { query: string; tokensNum: number }, ctx: Tool.Context) =>
execute: (params: Schema.Schema.Type<typeof Parameters>, ctx: Tool.Context) =>
Effect.gen(function* () {
yield* ctx.ask({
permission: "codesearch",

View File

@@ -11,11 +11,11 @@ const MAX_TIMEOUT = 120 * 1000 // 2 minutes
export const Parameters = Schema.Struct({
url: Schema.String.annotate({ description: "The URL to fetch content from" }),
format: Schema.Literals(["text", "markdown", "html"])
.pipe(Schema.optional, Schema.withDecodingDefault(Effect.succeed("markdown" as const)))
.pipe(Schema.withDecodingDefaultTypeKey(Effect.succeed("markdown" as const)))
.annotate({
description: "The format to return the content in (text, markdown, or html). Defaults to markdown.",
}),
timeout: Schema.optional(Schema.Number).annotate({ description: "Optional timeout in seconds (max 120)" }),
timeout: Schema.optional(Schema.Finite).annotate({ description: "Optional timeout in seconds (max 120)" }),
})
export const WebFetchTool = Tool.define(

View File

@@ -6,17 +6,21 @@ import DESCRIPTION from "./websearch.txt"
export const Parameters = Schema.Struct({
query: Schema.String.annotate({ description: "Websearch query" }),
numResults: Schema.optional(Schema.Number).annotate({
numResults: Schema.Finite.pipe(Schema.withDecodingDefaultTypeKey(Effect.succeed(8))).annotate({
description: "Number of search results to return (default: 8)",
}),
livecrawl: Schema.optional(Schema.Literals(["fallback", "preferred"])).annotate({
livecrawl: Schema.Literals(["fallback", "preferred"]).pipe(
Schema.withDecodingDefaultTypeKey(Effect.succeed("fallback" as const)),
).annotate({
description:
"Live crawl mode - 'fallback': use live crawling as backup if cached content unavailable, 'preferred': prioritize live crawling (default: 'fallback')",
}),
type: Schema.optional(Schema.Literals(["auto", "fast", "deep"])).annotate({
type: Schema.Literals(["auto", "fast", "deep"]).pipe(
Schema.withDecodingDefaultTypeKey(Effect.succeed("auto" as const)),
).annotate({
description: "Search type - 'auto': balanced search (default), 'fast': quick results, 'deep': comprehensive search",
}),
contextMaxCharacters: Schema.optional(Schema.Number).annotate({
contextMaxCharacters: Schema.optional(Schema.Finite).annotate({
description: "Maximum characters for context string optimized for LLMs (default: 10000)",
}),
})
@@ -52,9 +56,9 @@ export const WebSearchTool = Tool.define(
McpExa.SearchArgs,
{
query: params.query,
type: params.type || "auto",
numResults: params.numResults || 8,
livecrawl: params.livecrawl || "fallback",
type: params.type,
numResults: params.numResults,
livecrawl: params.livecrawl,
contextMaxCharacters: params.contextMaxCharacters,
},
"25 seconds",

View File

@@ -88,15 +88,18 @@ function walkUncached(ast: SchemaAST.AST): z.ZodTypeAny {
// Schema.decodeTo / Schema.transform attach encoding to non-Declaration
// nodes, where we do apply the transform.
//
// Schema.withDecodingDefault also attaches encoding, but we want `.default(v)`
// on the inner Zod rather than a transform wrapper — so optional ASTs whose
// encoding resolves a default from Option.none() route through body()/opt().
// Schema.withDecodingDefault and Schema.withDecodingDefaultTypeKey both
// attach encodings. For JSON Schema we want those as plain `.default(v)`
// annotations rather than transform wrappers, so ASTs whose encoding
// resolves a default from Option.none() route through body()/opt().
const hasEncoding = ast.encoding?.length && ast._tag !== "Declaration"
const hasTransform = hasEncoding && !(SchemaAST.isOptional(ast) && extractDefault(ast) !== undefined)
const base = hasTransform ? encoded(ast) : body(ast)
const fallback = hasEncoding ? extractDefault(ast) : undefined
const hasTransform = hasEncoding && fallback === undefined
const base = hasTransform ? encoded(ast) : body(ast, fallback)
const checked = ast.checks?.length ? applyChecks(base, ast.checks, ast) : base
const defaulted = fallback !== undefined && !SchemaAST.isOptional(ast) ? checked.default(fallback.value) : checked
const preprocess = (ast.annotations as { [ZodPreprocess]?: (val: unknown) => unknown } | undefined)?.[ZodPreprocess]
const out = preprocess ? z.preprocess(preprocess, checked) : checked
const out = preprocess ? z.preprocess(preprocess, defaulted) : defaulted
const desc = SchemaAST.resolveDescription(ast)
const ref = SchemaAST.resolveIdentifier(ast)
const described = desc ? out.describe(desc) : out
@@ -234,8 +237,8 @@ function issueMessage(issue: any): string | undefined {
return undefined
}
function body(ast: SchemaAST.AST): z.ZodTypeAny {
if (SchemaAST.isOptional(ast)) return opt(ast)
function body(ast: SchemaAST.AST, fallback?: { value: unknown }): z.ZodTypeAny {
if (SchemaAST.isOptional(ast)) return opt(ast, fallback)
switch (ast._tag) {
case "String":
@@ -268,7 +271,7 @@ function body(ast: SchemaAST.AST): z.ZodTypeAny {
}
}
function opt(ast: SchemaAST.AST): z.ZodTypeAny {
function opt(ast: SchemaAST.AST, fallback = extractDefault(ast)): z.ZodTypeAny {
if (ast._tag !== "Union") return fail(ast)
const items = ast.types.filter((item) => item._tag !== "Undefined")
const inner =
@@ -280,7 +283,6 @@ function opt(ast: SchemaAST.AST): z.ZodTypeAny {
// Schema.withDecodingDefault attaches an encoding `Link` whose transformation
// decode Getter resolves `Option.none()` to `Option.some(default)`. Invoke
// it to extract the default and emit `.default(...)` instead of `.optional()`.
const fallback = extractDefault(ast)
if (fallback !== undefined) return inner.default(fallback.value)
return inner.optional()
}

View File

@@ -487,6 +487,7 @@ exports[`tool parameters JSON Schema (wire shape) websearch 1`] = `
"type": "number",
},
"livecrawl": {
"default": "fallback",
"description": "Live crawl mode - 'fallback': use live crawling as backup if cached content unavailable, 'preferred': prioritize live crawling (default: 'fallback')",
"enum": [
"fallback",
@@ -495,6 +496,7 @@ exports[`tool parameters JSON Schema (wire shape) websearch 1`] = `
"type": "string",
},
"numResults": {
"default": 8,
"description": "Number of search results to return (default: 8)",
"type": "number",
},
@@ -503,6 +505,7 @@ exports[`tool parameters JSON Schema (wire shape) websearch 1`] = `
"type": "string",
},
"type": {
"default": "auto",
"description": "Search type - 'auto': balanced search (default), 'fast': quick results, 'deep': comprehensive search",
"enum": [
"auto",

View File

@@ -34,6 +34,9 @@ const parse = <S extends Schema.Decoder<unknown>>(schema: S, input: unknown): S[
const accepts = (schema: Schema.Decoder<unknown>, input: unknown): boolean =>
Result.isSuccess(Schema.decodeUnknownResult(schema)(input))
const toNativeJsonSchema = <S extends Schema.Decoder<unknown>>(schema: S) =>
Schema.toStandardJSONSchemaV1(schema)["~standard"].jsonSchema.input({ target: "draft-2020-12" })
describe("tool parameters", () => {
describe("JSON Schema (wire shape)", () => {
test("apply_patch", () => expect(toJsonSchema(ApplyPatch)).toMatchSnapshot())
@@ -56,6 +59,39 @@ describe("tool parameters", () => {
test("write", () => expect(toJsonSchema(Write)).toMatchSnapshot())
})
describe("native JSON Schema (experimental tool route)", () => {
test("codesearch uses a plain finite number", () => {
const native = toNativeJsonSchema(CodeSearch) as any
expect(native.properties.tokensNum).toEqual({
type: "number",
allOf: [{ minimum: 1000 }, { maximum: 50000 }],
})
})
test("webfetch format stays a string enum", () => {
const native = toNativeJsonSchema(WebFetch) as any
expect(native.properties.format).toEqual({
type: "string",
enum: ["text", "markdown", "html"],
})
})
test("websearch defaulted fields stay non-nullable", () => {
const native = toNativeJsonSchema(WebSearch) as any
expect(native.properties.numResults).toEqual({
type: "number",
})
expect(native.properties.livecrawl).toEqual({
type: "string",
enum: ["fallback", "preferred"],
})
expect(native.properties.type).toEqual({
type: "string",
enum: ["auto", "fast", "deep"],
})
})
})
describe("apply_patch", () => {
test("accepts patchText", () => {
expect(parse(ApplyPatch, { patchText: "*** Begin Patch\n*** End Patch" })).toEqual({
@@ -254,13 +290,21 @@ describe("tool parameters", () => {
describe("webfetch", () => {
test("accepts url-only", () => {
expect(parse(WebFetch, { url: "https://example.com" }).url).toBe("https://example.com")
expect(parse(WebFetch, { url: "https://example.com" })).toEqual({
url: "https://example.com",
format: "markdown",
})
})
})
describe("websearch", () => {
test("accepts query", () => {
expect(parse(WebSearch, { query: "opencode" }).query).toBe("opencode")
expect(parse(WebSearch, { query: "opencode" })).toEqual({
query: "opencode",
numResults: 8,
livecrawl: "fallback",
type: "auto",
})
})
})

View File

@@ -2,7 +2,7 @@ import { describe, expect, test } from "bun:test"
import { Effect, Schema, SchemaGetter } from "effect"
import z from "zod"
import { zod, ZodOverride, ZodPreprocess } from "../../src/util/effect-zod"
import { toJsonSchema, zod, ZodOverride, ZodPreprocess } from "../../src/util/effect-zod"
function json(schema: z.ZodTypeAny) {
const { $schema: _, ...rest } = z.toJSONSchema(schema)
@@ -750,6 +750,36 @@ describe("util.effect-zod", () => {
expect(schema.parse({})).toEqual({})
expect(schema.parse({ foo: "hi" })).toEqual({ foo: "hi" })
})
test("key defaults fill in missing struct keys", () => {
const schema = zod(
Schema.Struct({
mode: Schema.String.pipe(Schema.withDecodingDefaultTypeKey(Effect.succeed("ctrl-x"))),
}),
)
expect(schema.parse({})).toEqual({ mode: "ctrl-x" })
})
test("key defaults still accept explicit values", () => {
const schema = zod(
Schema.Struct({
mode: Schema.String.pipe(Schema.withDecodingDefaultTypeKey(Effect.succeed("ctrl-x"))),
}),
)
expect(schema.parse({ mode: "ctrl-c" })).toEqual({ mode: "ctrl-c" })
})
test("JSON Schema output includes the default key for key defaults", () => {
const shape = toJsonSchema(
Schema.Struct({
mode: Schema.String.pipe(Schema.withDecodingDefaultTypeKey(Effect.succeed("ctrl-x"))),
}),
) as any
expect(shape.properties.mode.default).toBe("ctrl-x")
expect(shape.required).toBeUndefined()
})
})
describe("ZodPreprocess annotation", () => {