From e0b6039e1f26b753c07d690f6ea0ba7035e72c9d Mon Sep 17 00:00:00 2001 From: Aiden Cline <63023139+rekram1-node@users.noreply.github.com> Date: Sat, 24 Oct 2026 22:22:35 -0500 Subject: [PATCH] refactor(core): convert tool schemas once at registration (#58848) --- packages/core/src/session/model-request.ts | 6 +- packages/core/src/tool.ts | 68 +++++++++++++------ packages/core/src/tool/runtime.ts | 24 ++----- .../test/session-model-request-hooks.test.ts | 36 ++++++++++ packages/core/test/tool-execute.test.ts | 6 +- packages/core/test/tool-registry.test.ts | 30 ++++++++ packages/core/test/tool-schema.test.ts | 47 +------------ 7 files changed, 127 insertions(+), 90 deletions(-) diff --git a/packages/core/src/session/model-request.ts b/packages/core/src/session/model-request.ts index e50575a24e8..bf622dc4812 100644 --- a/packages/core/src/session/model-request.ts +++ b/packages/core/src/session/model-request.ts @@ -238,8 +238,12 @@ export const layer = Layer.effect( } // Remember which tool each definition object came from. Hooks rename a tool by moving // its definition to a new key, so after the hook we find the tool by object identity. + // Definitions are shared across steps and sessions, and hooks may edit schemas in place, + // so each hook gets a deep copy. const given = new Map( - tools.definitions.map((t) => [{ description: t.description, input: { ...t.inputSchema } }, t] as const), + tools.definitions.map( + (t) => [{ description: t.description, input: JSON.parse(JSON.stringify(t.inputSchema)) }, t] as const, + ), ) // Hooks see the default output limit and may change or remove it. Titles and generate keep the provider default, // because their reasoning is hard to budget. diff --git a/packages/core/src/tool.ts b/packages/core/src/tool.ts index 2ce6c3c1e63..50f5816f948 100644 --- a/packages/core/src/tool.ts +++ b/packages/core/src/tool.ts @@ -34,6 +34,8 @@ export interface Editor { type Data = { tools: Map + // Converted once at registration; schemas cannot change without re-registering the tool. + definitions: Map namespaces: Map errors: { kind: "tool" | "namespace"; name: string; namespace?: string; error: RegistrationError }[] } @@ -156,12 +158,19 @@ const layer = Layer.effect( }) let catalog: - | { data: Data; names: string; inventory: CodeModeTool.Inventory; value: CodeModeCatalog.Inventory } + | { + data: Data + names: string + inventory: CodeModeTool.Inventory + value: CodeModeCatalog.Inventory + definition: ToolDefinition + } | undefined const state = State.create({ name: "tool", initial: () => ({ tools: new Map(), + definitions: new Map(), namespaces: new Map(), errors: [], }), @@ -177,13 +186,19 @@ const layer = Layer.effect( editor.namespaces.set(namespace.name, { ...namespace }) }, add: (tool) => { - const error = registrationError(tool) - if (error) { - editor.errors.push({ kind: "tool", name: tool.name, namespace: tool.options?.namespace, error }) + const registered = register(tool) + if (Result.isFailure(registered)) { + editor.errors.push({ + kind: "tool", + name: tool.name, + namespace: tool.options?.namespace, + error: registered.failure, + }) return } const id = effectiveName(tool) editor.tools.set(id, { ...tool, id, options: tool.options && { ...tool.options } }) + editor.definitions.set(id, registered.success) }, update: (id, update) => { const current = editor.tools.get(id) @@ -194,15 +209,22 @@ const layer = Layer.effect( tool.id = id if (tool.options?.namespace !== current.options?.namespace) tool.options = { ...tool.options, namespace: current.options?.namespace } - const error = registrationError(tool) - if (error) { - editor.errors.push({ kind: "tool", name: tool.name, namespace: tool.options?.namespace, error }) + const registered = register(tool) + if (Result.isFailure(registered)) { + editor.errors.push({ + kind: "tool", + name: tool.name, + namespace: tool.options?.namespace, + error: registered.failure, + }) return } editor.tools.set(id, tool) + editor.definitions.set(id, registered.success) }, remove: (id) => { editor.tools.delete(id) + editor.definitions.delete(id) }, }), notify: (value) => { @@ -237,9 +259,9 @@ const layer = Layer.effect( const codeModeTools = new Map(Array.from(active).filter(([, tool]) => tool.options?.codemode !== false)) const names = Array.from(codeModeTools.keys()).join("\0") const codeModeEnabled = !whollyDisabled("execute", rules) - // Discovery is immutable for a registry revision and visible tool set. Keep request - // definitions/executors fresh, but share the inventory, whose prepared runtime and much - // larger rendered catalog are cached by identity, across steps. + // Discovery is immutable for a registry revision and visible tool set. Keep executors + // fresh, but share the inventory, whose prepared runtime and much larger rendered catalog + // are cached by identity, and the execute tool's definition across steps. const cached = catalog?.data === data && catalog.names === names ? catalog : undefined const codeModeInventory = cached?.inventory ?? { tools: codeModeTools, namespaces: data.namespaces } const codeModeTool = codeModeEnabled @@ -252,14 +274,22 @@ const layer = Layer.effect( const codeModeCatalog = !codeModeEnabled ? undefined : (cached?.value ?? CodeModeTool.catalog(codeModeInventory)) - if (codeModeCatalog) catalog = { data, names, inventory: codeModeInventory, value: codeModeCatalog } + const codeModeDefinition = codeModeTool && (cached?.definition ?? definition(codeModeTool)) + if (codeModeCatalog && codeModeDefinition) + catalog = { + data, + names, + inventory: codeModeInventory, + value: codeModeCatalog, + definition: codeModeDefinition, + } return { ...(codeModeCatalog === undefined ? {} : { codeModeCatalog }), definitions: [ ...Array.from(direct) .sort(([left], [right]) => (left < right ? -1 : left > right ? 1 : 0)) - .map(([, tool]) => definition(tool)), - ...(codeModeTool ? [definition(codeModeTool)] : []), + .map(([name]) => data.definitions.get(name)!), + ...(codeModeDefinition ? [codeModeDefinition] : []), ], execute: Effect.fnUntraced(function* (input: Parameters[0]) { const context: Tool.Context = { @@ -302,23 +332,23 @@ function schemaMakeError(error: unknown) { return error instanceof Error ? error.message : String(error) } -function registrationError(tool: Tool.Info) { +function register(tool: Tool.Info): Result.Result { const namespace = tool.options?.namespace if (namespace !== undefined) { const error = namespaceError(namespace) - if (error) return error + if (error) return Result.fail(error) } const name = normalizedName(tool) - if (!/^[A-Za-z0-9_-]+$/.test(name)) return new RegistrationError({ name, message: `Invalid tool name: ${name}` }) + if (!/^[A-Za-z0-9_-]+$/.test(name)) + return Result.fail(new RegistrationError({ name, message: `Invalid tool name: ${name}` })) const id = effectiveName(tool) if (tool.options?.codemode === false && id === "execute") - return new RegistrationError({ name: id, message: 'Tool name "execute" is reserved for CodeMode' }) - const result = Result.try({ + return Result.fail(new RegistrationError({ name: id, message: 'Tool name "execute" is reserved for CodeMode' })) + return Result.try({ try: () => ToolDefinition.make(definition(tool)), catch: (error) => new RegistrationError({ name: id, message: `Invalid tool definition ${id}: ${schemaMakeError(error)}` }), }) - return Result.isFailure(result) ? result.failure : undefined } function namespaceError(name: string) { diff --git a/packages/core/src/tool/runtime.ts b/packages/core/src/tool/runtime.ts index ace47d5deb6..16c0502f2a9 100644 --- a/packages/core/src/tool/runtime.ts +++ b/packages/core/src/tool/runtime.ts @@ -154,32 +154,16 @@ const validateStandard = ( }), ) -// Definitions are rebuilt on every step, but schemas are immutable, so convert each one once. -const inputJsonSchemas = new WeakMap() -const outputJsonSchemas = new WeakMap() - -const memoize = ( - cache: WeakMap, - schema: object, - build: () => JsonSchema.JsonSchema, -): JsonSchema.JsonSchema => { - const cached = cache.get(schema) - if (cached) return cached - const built = build() - cache.set(schema, built) - return built -} - const inputJsonSchema = (schema: Tool.ValueSchema): JsonSchema.JsonSchema => { if (schema === undefined || schema === null) return {} - if (isStandardSchema(schema)) return memoize(inputJsonSchemas, schema, () => standardJsonSchema(schema, "input")) + if (isStandardSchema(schema)) return standardJsonSchema(schema, "input") if (!Schema.isSchema(schema)) return schema - return memoize(inputJsonSchemas, schema, () => emptyInputJsonSchema(schema) ?? toJsonSchema(schema)) + return emptyInputJsonSchema(schema) ?? toJsonSchema(schema) } const outputJsonSchema = (schema: Tool.ValueSchema): JsonSchema.JsonSchema => { - if (isStandardSchema(schema)) return memoize(outputJsonSchemas, schema, () => standardJsonSchema(schema, "output")) - return Schema.isSchema(schema) ? memoize(outputJsonSchemas, schema, () => toJsonSchema(schema)) : schema + if (isStandardSchema(schema)) return standardJsonSchema(schema, "output") + return Schema.isSchema(schema) ? toJsonSchema(schema) : schema } const standardJsonSchema = (schema: StandardSchemaV1, io: "input" | "output"): JsonSchema.JsonSchema => { diff --git a/packages/core/test/session-model-request-hooks.test.ts b/packages/core/test/session-model-request-hooks.test.ts index 8e761206634..cda56d2b49e 100644 --- a/packages/core/test/session-model-request-hooks.test.ts +++ b/packages/core/test/session-model-request-hooks.test.ts @@ -1,4 +1,5 @@ import { describe, expect } from "bun:test" +import { ToolDefinition } from "@opencode/ai" import { OpenAIChat } from "@opencode/ai/protocols" import { Agent } from "@opencode/schema/agent" import { Money } from "@opencode/schema/money" @@ -80,6 +81,41 @@ describe("SessionModelRequest HTTP hooks", () => { }).pipe(Effect.provideService(SessionModelTransport.Service, transport)), ) + it.effect("keeps hook edits to a tool schema out of later requests", () => + Effect.gen(function* () { + const hooks = yield* PluginHooks.Service + yield* hooks.register("session", "context", (event) => + Effect.sync(() => { + const input = event.tools.lookup!.input as { required: string[] } + input.required.push("extra") + }), + ) + const requests = yield* SessionModelRequest.Service.pipe(Effect.provide(SessionModelRequest.layer)) + const inputSchema = { type: "object", properties: { id: { type: "string" } }, required: ["id"] } + const tools = { + definitions: [new ToolDefinition({ name: "lookup", description: "Lookup", inputSchema })], + execute: () => Effect.die("unused tool execution"), + } + + const required = yield* Effect.forEach([1, 2], () => + requests + .primary({ session, agent: Agent.ID.make("build"), model, system: [], messages: [], tools }) + .pipe( + Effect.map((prepared) => { + const tool = prepared.request.tools[0] + return tool?.type === "tool" ? tool.inputSchema.required : undefined + }), + ), + ) + + expect(required).toEqual([ + ["id", "extra"], + ["id", "extra"], + ]) + expect(inputSchema.required).toEqual(["id"]) + }).pipe(Effect.provideService(SessionModelTransport.Service, transport)), + ) + it.effect("offers the WebSocket executor alongside HTTP hooks and routes the WebSocket hooks", () => Effect.gen(function* () { const hooks = yield* PluginHooks.Service diff --git a/packages/core/test/tool-execute.test.ts b/packages/core/test/tool-execute.test.ts index 2013b972f89..d2f190f24d9 100644 --- a/packages/core/test/tool-execute.test.ts +++ b/packages/core/test/tool-execute.test.ts @@ -205,12 +205,8 @@ test("executions share the prepared runtime of an inventory while keeping per-ca ).toBe("call_execute:x") expect(conversions).toBe(1) - // Converted schemas are cached by identity, so a fresh schema shows the new inventory's runtime is rebuilt. await Effect.runPromise( - createCodeMode(new Map([["echo", { ...echo, input: { ...echo.input } }]])).execute( - { code: 'return await tools.echo({ value: "next" })' }, - context, - ), + createCodeMode(new Map(inventory.tools)).execute({ code: 'return await tools.echo({ value: "next" })' }, context), ) expect(conversions).toBe(2) }) diff --git a/packages/core/test/tool-registry.test.ts b/packages/core/test/tool-registry.test.ts index 22c7c0d77a0..dc446ca3bb5 100644 --- a/packages/core/test/tool-registry.test.ts +++ b/packages/core/test/tool-registry.test.ts @@ -182,6 +182,36 @@ describe("Tool", () => { }), ) + it.effect("converts schemas at registration and reuses them in every snapshot", () => + Effect.gen(function* () { + const service = yield* Tool.Service + let builds = 0 + const portable = { + "~standard": { + version: 1, + vendor: "test", + validate: (value: unknown) => ({ value }), + jsonSchema: { + input: () => { + builds++ + return { type: "object", properties: { text: { type: "string" } } } + }, + output: () => ({ type: "string" }), + }, + }, + } as const + yield* transform(service, { portable: { ...make(), input: portable } }, { codemode: false }) + const registered = builds + + const first = yield* service.snapshot() + const second = yield* service.snapshot() + + expect(builds).toBe(registered) + expect(second.definitions).toHaveLength(2) + second.definitions.forEach((item, index) => expect(item).toBe(first.definitions[index])) + }), + ) + it.effect("skips invalid updates without dropping the existing definition", () => Effect.gen(function* () { const service = yield* Tool.Service diff --git a/packages/core/test/tool-schema.test.ts b/packages/core/test/tool-schema.test.ts index 7d406fc6214..cf6682e5365 100644 --- a/packages/core/test/tool-schema.test.ts +++ b/packages/core/test/tool-schema.test.ts @@ -193,50 +193,7 @@ test("Effect tool schemas resolve escaped definition names", () => { expect(JSON.stringify(definition(tool).inputSchema)).not.toContain("$defs") }) -test("repeat steps reuse converted schemas instead of rebuilding them", () => { - const builds: Array = [] - const portable = { - "~standard": { - version: 1, - vendor: "test", - validate: (value: unknown) => ({ value }), - jsonSchema: { - input: () => { - builds.push("input") - return { type: "object" } - }, - output: () => { - builds.push("output") - return { type: "string" } - }, - }, - }, - } - const effect: Info = { - name: "effect", - description: "Effect tool", - input: Schema.Struct({ value: Schema.String }), - output: Schema.Struct({ ok: Schema.Boolean }), - execute: () => Effect.succeed({ output: { ok: true } }), - } - const step = () => [ - // Snapshots recreate some tools, such as CodeMode's, around the same schemas each step. - definition({ ...effect }), - definition({ ...effect, name: "portable", input: portable, output: portable }), - ] - - const first = step() - const before = builds.length - const second = step() - - expect(builds.length).toBe(before) - second.forEach((item, index) => { - expect(item.inputSchema).toBe(first[index].inputSchema) - expect(item.outputSchema).toBe(first[index].outputSchema) - }) -}) - -test("cached definitions serialize byte-identically to freshly built ones", () => { +test("definitions keep their pinned serialization", () => { const Item = Schema.Struct({ id: Schema.String }).annotate({ identifier: "Item" }) const tool: Info = { name: "pinned.tool", @@ -249,7 +206,7 @@ test("cached definitions serialize byte-identically to freshly built ones", () = output: Schema.Struct({}).annotate({ description: "Nothing" }), execute: () => Effect.succeed({ content: "unused" }), } - // Serialized by the uncached implementation; key order matters for provider prompt caches. + // Key order matters for provider prompt caches. const expected = '{"type":"tool","name":"mcp_server_pinned_tool","description":"Pinned","inputSchema":{"type":"object","properties":{"offset":{"type":"integer","minimum":0},"items":{"type":"array","items":{"type":"object","properties":{"id":{"type":"string"}},"required":["id"],"additionalProperties":false}}},"required":["items"],"additionalProperties":false},"outputSchema":{"not":{"type":"null"},"description":"Nothing"}}'