diff --git a/packages/core/src/session/model-request.ts b/packages/core/src/session/model-request.ts index fab9e4c0e46..bae38339c00 100644 --- a/packages/core/src/session/model-request.ts +++ b/packages/core/src/session/model-request.ts @@ -189,34 +189,43 @@ export const layer = Layer.effect( .map(SystemPart.make) const history = toLLMMessages(input.context.messages, resolved.ref, providerMetadataKey) const messages = stepLimitReached ? [...history, Message.assistant(MAX_STEPS_PROMPT)] : history - const toolDefinitions = tools.definitions - const toolsByName = new Map(toolDefinitions.map((tool) => [tool.name, tool])) - // Hooks may reshape available definitions but cannot advertise tools omitted by permissions or the Step limit. - const contextEvent = yield* hooks.trigger("session", "context", { + const registry = new Map(tools.definitions.map((tool) => [tool.name, tool])) + // The definition objects we hand to hooks, mapped back to their tools. Hooks rename a + // tool by moving its definition to a new key; recognizing the object recovers the tool. + const given = new Map( + tools.definitions.map( + (tool) => [{ description: tool.description, input: { ...tool.inputSchema } }, tool] as const, + ), + ) + // Hooks mutate this record in place: edit descriptions and schemas, rename, or remove. + const context = yield* hooks.trigger("session", "context", { sessionID: session.id, agent: agent.id, model: resolved.ref, system, messages, - tools: Object.fromEntries( - toolDefinitions.map((tool) => [tool.name, { description: tool.description, input: { ...tool.inputSchema } }]), - ), - }) - const hookedTools = Object.entries(contextEvent.tools).flatMap(([name, tool]) => { - const registered = toolsByName.get(name) - return registered - ? [{ ...registered, description: tool.description, inputSchema: tool.input }] - : [] + tools: Object.fromEntries(Array.from(given, ([definition, tool]) => [tool.name, definition])), }) + // Match each surviving entry back to its tool, by recognizing a moved definition or + // by key. Identity wins so a definition moved onto another tool's name still executes + // the tool it describes. Entries matching neither were invented by a hook and dropped. + // `tool.name` stays canonical so execution can translate renamed calls back. + const hooked = new Map( + Object.entries(context.tools).flatMap(([name, definition]) => { + const tool = given.get(definition) ?? registry.get(name) + if (!tool) return [] + return [[name, { ...tool, description: definition.description, inputSchema: definition.input }] as const] + }), + ) const request = LLM.request({ model, http: { headers: SessionModelHeaders.make(session, app), }, providerOptions: { [providerMetadataKey]: { promptCacheKey } }, - system: contextEvent.system, - messages: boundImages(unsupportedParts(contextEvent.messages, resolved.capabilities)), - tools: hookedTools, + system: context.system, + messages: boundImages(unsupportedParts(context.messages, resolved.capabilities)), + tools: Array.from(hooked, ([name, tool]) => ({ ...tool, name })), toolChoice: stepLimitReached ? "none" : undefined, }) const options: StreamOptions = { @@ -256,13 +265,15 @@ export const layer = Layer.effect( }), ) } - const executeTool: Prepared["executeTool"] = (executeInput) => { + const executeTool: Prepared["executeTool"] = (input) => { if (stepLimitReached) return new Tool.Error({ message: "Tools are disabled after the maximum agent steps" }) - if (toolsByName.has(executeInput.call.name) && !Object.hasOwn(contextEvent.tools, executeInput.call.name)) - return new Tool.Error({ message: `Tool is not available for this request: ${executeInput.call.name}` }) + const tool = hooked.get(input.call.name) + // A registered tool absent from the hooked set was removed or renamed by a hook. + if (!tool && registry.has(input.call.name)) + return new Tool.Error({ message: `Tool is not available for this request: ${input.call.name}` }) return tools - .execute(executeInput) + .execute(tool ? { ...input, call: { ...input.call, name: tool.name } } : input) .pipe(Effect.catchCauseFilter(declineDefect, (decline) => Effect.fail(decline))) } return { diff --git a/packages/core/test/session-runner.test.ts b/packages/core/test/session-runner.test.ts index f7bca091004..99e8e023724 100644 --- a/packages/core/test/session-runner.test.ts +++ b/packages/core/test/session-runner.test.ts @@ -887,6 +887,27 @@ describe("SessionRunnerLLM", () => { }), ) + it.effect("executes a tool renamed by a session context hook", () => + Effect.gen(function* () { + const session = yield* setup + const hooks = yield* PluginHooks.Service + yield* hooks.register("session", "context", (event) => + Effect.sync(() => { + event.tools.renamed_echo = event.tools.echo! + delete event.tools.echo + }), + ) + yield* admit(session, "Use the renamed tool") + yield* TestLLM.push(TestLLM.tool("call-renamed", "renamed_echo", { text: "renamed" }), []) + + yield* session.resume(sessionID) + + expect(requests[0]?.tools.map((tool) => tool.name)).toContain("renamed_echo") + expect(requests[0]?.tools.map((tool) => tool.name)).not.toContain("echo") + expect(executions).toEqual(["renamed"]) + }), + ) + it.effect("advertises and executes a location registered tool", () => Effect.gen(function* () { const session = yield* setup