From f2f5eb6f16260f50f611f75858ed92a154b60b25 Mon Sep 17 00:00:00 2001 From: Dax Raad Date: Wed, 15 Jul 2026 15:58:14 -0400 Subject: [PATCH] feat(plugin): restore ai request hook --- bun.lock | 1 + packages/core/src/plugin/hooks.ts | 2 + packages/core/src/plugin/host.ts | 3 ++ packages/core/src/plugin/promise.ts | 4 ++ packages/core/src/session/runner/llm.ts | 36 ++++++++++++++++-- packages/core/src/tool/patch.ts | 13 +++++++ packages/core/src/tool/subagent.ts | 27 ++++++++++++++ packages/core/test/lib/tool.ts | 3 ++ packages/core/test/plugin-hooks.test.ts | 45 +++++++++++++++++++++++ packages/core/test/plugin/host.ts | 3 ++ packages/core/test/plugin/promise.test.ts | 37 +++++++++++++++++++ packages/core/test/session-runner.test.ts | 30 +++++++++++++++ packages/plugin/package.json | 1 + packages/plugin/src/v2/effect/README.md | 12 +++--- packages/plugin/src/v2/effect/ai.ts | 23 ++++++++++++ packages/plugin/src/v2/effect/plugin.ts | 2 + packages/plugin/src/v2/promise/README.md | 4 +- packages/plugin/src/v2/promise/ai.ts | 23 ++++++++++++ packages/plugin/src/v2/promise/plugin.ts | 2 + 19 files changed, 261 insertions(+), 10 deletions(-) create mode 100644 packages/core/test/plugin-hooks.test.ts create mode 100644 packages/plugin/src/v2/effect/ai.ts create mode 100644 packages/plugin/src/v2/promise/ai.ts diff --git a/bun.lock b/bun.lock index 260f3a16ee..87354a2d82 100644 --- a/bun.lock +++ b/bun.lock @@ -720,6 +720,7 @@ "version": "1.17.20", "dependencies": { "@ai-sdk/provider": "3.0.8", + "@opencode-ai/ai": "workspace:*", "@opencode-ai/client": "workspace:*", "@opencode-ai/schema": "workspace:*", "@opencode-ai/sdk": "workspace:*", diff --git a/packages/core/src/plugin/hooks.ts b/packages/core/src/plugin/hooks.ts index 0023e62f6c..5c36938612 100644 --- a/packages/core/src/plugin/hooks.ts +++ b/packages/core/src/plugin/hooks.ts @@ -1,5 +1,6 @@ export * as PluginHooks from "./hooks" +import type { AIHooks } from "@opencode-ai/plugin/v2/effect/ai" import type { AISDKHooks } from "@opencode-ai/plugin/v2/effect/aisdk" import type { ToolHooks } from "@opencode-ai/plugin/v2/effect/tool" import { Context, Effect, Layer, Scope } from "effect" @@ -7,6 +8,7 @@ import { makeLocationNode } from "../effect/app-node" import { State } from "../state" export interface Domains { + readonly ai: AIHooks readonly aisdk: AISDKHooks readonly tool: ToolHooks } diff --git a/packages/core/src/plugin/host.ts b/packages/core/src/plugin/host.ts index 9b175a80b4..207a0cb1e1 100644 --- a/packages/core/src/plugin/host.ts +++ b/packages/core/src/plugin/host.ts @@ -80,6 +80,9 @@ export const make = Effect.fn("PluginHost.make")(function* (plugin: PluginV2.Int }) }), }, + ai: { + hook: (name, callback) => hooks.register("ai", name, callback), + }, aisdk: { hook: (name, callback) => { if (name === "sdk") { diff --git a/packages/core/src/plugin/promise.ts b/packages/core/src/plugin/promise.ts index 53caee351c..670964b3ca 100644 --- a/packages/core/src/plugin/promise.ts +++ b/packages/core/src/plugin/promise.ts @@ -63,6 +63,10 @@ export function fromPromise(plugin: Plugin) { transform: transform(host.agent), reload: () => run(host.agent.reload()), }, + ai: { + hook: (name, callback) => + register(host.ai.hook(name, (event) => Effect.promise(() => Promise.resolve(callback(event))))), + }, aisdk: { hook: (name, callback) => register(host.aisdk.hook(name, (event) => Effect.promise(() => Promise.resolve(callback(event))))), diff --git a/packages/core/src/session/runner/llm.ts b/packages/core/src/session/runner/llm.ts index 6801559fa8..3a7ab625bc 100644 --- a/packages/core/src/session/runner/llm.ts +++ b/packages/core/src/session/runner/llm.ts @@ -10,6 +10,7 @@ import { isContextOverflowFailure, type ProviderErrorEvent, } from "@opencode-ai/ai" +import type { AIHooks } from "@opencode-ai/plugin/v2/effect/ai" import { SessionError } from "@opencode-ai/schema/session-error" import { Money } from "@opencode-ai/schema/money" import { Cause, Effect, Exit, Fiber, FiberSet, Layer, Option, Semaphore, Stream } from "effect" @@ -51,6 +52,7 @@ import { AgentNotFoundError, StepFailedError } from "../error" import { toSessionError } from "../to-session-error" import { SessionRunnerRetry } from "./retry" import { PluginSupervisor } from "../../plugin/supervisor" +import { PluginHooks } from "../../plugin/hooks" import { SessionModelHeaders } from "../model-headers" type StepTokens = { @@ -89,6 +91,7 @@ const layer = Layer.effect( const llm = yield* LLMClient.Service const agents = yield* AgentV2.Service const tools = yield* ToolRegistry.Service + const hooks = yield* PluginHooks.Service const models = yield* SessionRunnerModel.Service const store = yield* SessionStore.Service const location = yield* Location.Service @@ -212,6 +215,30 @@ const layer = Layer.effect( tools: toolMaterialization?.definitions ?? [], toolChoice: isLastStep ? "none" : undefined, }) + const availableTools = new Map(request.tools.map((tool) => [tool.name, tool])) + const requestEvent: AIHooks["request"] = { + sessionID: session.id, + agent: agent.id, + model: resolved.ref, + system: [...request.system], + messages: [...request.messages], + tools: Object.fromEntries( + request.tools.map((tool) => [tool.name, { description: tool.description, input: { ...tool.inputSchema } }]), + ), + } + // Plugins may reshape the draft but cannot advertise tools excluded by + // permissions, registration state, or the selected agent's step limit. + yield* hooks.trigger("ai", "request", requestEvent) + const hookedRequest = LLM.updateRequest(request, { + system: requestEvent.system, + messages: requestEvent.messages, + tools: Object.entries(requestEvent.tools).flatMap(([name, tool]) => { + const registered = availableTools.get(name) + if (!registered) return [] + return [{ ...registered, description: tool.description, inputSchema: tool.input }] + }), + }) + const advertisedTools = new Set(hookedRequest.tools.map((tool) => tool.name)) const toolFibers = yield* FiberSet.make() const ownedToolFibers: Array> = [] let needsContinuation = false @@ -232,7 +259,7 @@ const layer = Layer.effect( const serialized = (effect: Effect.Effect) => publication.withPermit(effect) const publish = (event: LLMEvent, error?: SessionError.Error) => serialized(publisher.publish(event, error)) let overflowFailure: ProviderErrorEvent | undefined - const providerStream = llm.stream(request).pipe( + const providerStream = llm.stream(hookedRequest).pipe( Stream.runForEach((event) => Effect.gen(function* () { if (overflowFailure || publisher.hasProviderError()) return @@ -244,11 +271,13 @@ const layer = Layer.effect( } yield* publish(event) if (event.type !== "tool-call" || event.providerExecuted) return - if (!toolMaterialization) { + if (!toolMaterialization || (availableTools.has(event.name) && !advertisedTools.has(event.name))) { yield* serialized( publisher.failUnsettledTools({ type: "tool.execution", - message: "Tools are disabled after the maximum agent steps", + message: toolMaterialization + ? `Tool is not available for this request: ${event.name}` + : "Tools are disabled after the maximum agent steps", }), ) return @@ -585,6 +614,7 @@ export const node = makeLocationNode({ llmClient, AgentV2.node, ToolRegistry.node, + PluginHooks.node, SessionRunnerModel.node, SessionStore.node, Location.node, diff --git a/packages/core/src/tool/patch.ts b/packages/core/src/tool/patch.ts index fba773b735..57e8c0ef2a 100644 --- a/packages/core/src/tool/patch.ts +++ b/packages/core/src/tool/patch.ts @@ -195,6 +195,19 @@ export const Plugin = { ), ) .pipe(Effect.orDie) + + yield* ctx.ai.hook("request", (event) => + Effect.sync(() => { + const usePatch = + event.model.providerID.toLowerCase() === "openai" || event.model.id.toLowerCase().includes("gpt") + if (usePatch) { + delete event.tools.edit + delete event.tools.write + return + } + delete event.tools.patch + }), + ) }), } diff --git a/packages/core/src/tool/subagent.ts b/packages/core/src/tool/subagent.ts index 7e08a3a273..0cc2de9db5 100644 --- a/packages/core/src/tool/subagent.ts +++ b/packages/core/src/tool/subagent.ts @@ -215,5 +215,32 @@ export const Plugin = { ), ) .pipe(Effect.orDie) + + yield* ctx.ai.hook("request", (event) => + Effect.gen(function* () { + const tool = event.tools[name] + if (!tool) return + const selected = yield* agents.resolve(event.agent) + if (!selected) return + const available = (yield* agents.list()) + .filter( + (agent) => + agent.mode !== "primary" && + !agent.hidden && + PermissionV2.evaluate(name, agent.id, selected.permissions).effect !== "deny", + ) + .toSorted((a, b) => a.id.localeCompare(b.id)) + if (available.length === 0) return + tool.description = [ + tool.description, + "", + "Available subagents:", + ...available.map( + (agent) => + `- ${agent.id}: ${agent.description ?? "This subagent should only be called when explicitly requested."}`, + ), + ].join("\n") + }), + ) }), } diff --git a/packages/core/test/lib/tool.ts b/packages/core/test/lib/tool.ts index 468f9e3d9c..9adf3c240b 100644 --- a/packages/core/test/lib/tool.ts +++ b/packages/core/test/lib/tool.ts @@ -44,6 +44,9 @@ export const registerToolPlugin = (plugin: { Effect.gen(function* () { const tools = yield* Tools.Service const context = host({ + ai: { + hook: () => Effect.succeed({ dispose: Effect.void }), + }, tool: { transform: (callback) => Effect.gen(function* () { diff --git a/packages/core/test/plugin-hooks.test.ts b/packages/core/test/plugin-hooks.test.ts new file mode 100644 index 0000000000..1223aa0c94 --- /dev/null +++ b/packages/core/test/plugin-hooks.test.ts @@ -0,0 +1,45 @@ +import { describe, expect } from "bun:test" +import { Message, SystemPart } from "@opencode-ai/ai" +import { Agent } from "@opencode-ai/schema/agent" +import { Model } from "@opencode-ai/schema/model" +import { Provider } from "@opencode-ai/schema/provider" +import { Session } from "@opencode-ai/schema/session" +import { Effect, Layer } from "effect" +import { PluginHooks } from "../src/plugin/hooks" +import { testEffect } from "./lib/effect" + +const layer = PluginHooks.node.implementation as Layer.Layer +const it = testEffect(layer) + +describe("PluginHooks", () => { + it.effect("registers scoped AI hooks and triggers them sequentially", () => + Effect.gen(function* () { + const hooks = yield* PluginHooks.Service + const seen: string[] = [] + yield* hooks.register("ai", "request", (event) => + Effect.sync(() => { + seen.push("first") + event.system.push(SystemPart.make("second")) + }), + ) + yield* hooks.register("ai", "request", (event) => + Effect.sync(() => { + seen.push(event.system[1]?.text ?? "missing") + event.messages = [Message.user("changed")] + }), + ) + const event = { + sessionID: Session.ID.make("ses_hooks"), + agent: Agent.ID.make("build"), + model: Model.Ref.make({ providerID: Provider.ID.make("test"), id: Model.ID.make("model") }), + system: [SystemPart.make("first")], + messages: [Message.user("original")], + tools: {}, + } + + expect(yield* hooks.trigger("ai", "request", event)).toBe(event) + expect(seen).toEqual(["first", "second"]) + expect(event.messages).toEqual([Message.user("changed")]) + }), + ) +}) diff --git a/packages/core/test/plugin/host.ts b/packages/core/test/plugin/host.ts index 679624ee57..ae7c66ec83 100644 --- a/packages/core/test/plugin/host.ts +++ b/packages/core/test/plugin/host.ts @@ -18,6 +18,9 @@ export function host(overrides: Overrides = {}): PluginContext { transform: () => Effect.die("unused agent.transform"), reload: () => Effect.die("unused agent.reload"), }, + ai: overrides.ai ?? { + hook: () => Effect.die("unused ai.hook"), + }, aisdk: overrides.aisdk ?? { hook: () => Effect.die("unused aisdk.hook"), }, diff --git a/packages/core/test/plugin/promise.test.ts b/packages/core/test/plugin/promise.test.ts index 2c6fccc892..e7577faa6a 100644 --- a/packages/core/test/plugin/promise.test.ts +++ b/packages/core/test/plugin/promise.test.ts @@ -1,13 +1,18 @@ import { describe, expect } from "bun:test" +import { Message, SystemPart } from "@opencode-ai/ai" import { Effect, Schema } from "effect" import { AgentV2 } from "@opencode-ai/core/agent" import { PluginV2 } from "@opencode-ai/core/plugin" +import { PluginHooks } from "@opencode-ai/core/plugin/hooks" import { PluginHost } from "@opencode-ai/core/plugin/host" import { PluginPromise } from "@opencode-ai/core/plugin/promise" import { SessionV2 } from "@opencode-ai/core/session" import { SessionMessage } from "@opencode-ai/core/session/message" import { ToolRegistry } from "@opencode-ai/core/tool/registry" import { Plugin } from "@opencode-ai/plugin/v2" +import type { AIHooks } from "@opencode-ai/plugin/v2/effect/ai" +import { Model } from "@opencode-ai/schema/model" +import { Provider } from "@opencode-ai/schema/provider" import { testEffect } from "../lib/effect" import { PluginTestLayer } from "./fixture" @@ -72,6 +77,38 @@ describe("fromPromise", () => { }), ) + it.effect("forwards AI request hooks", () => + Effect.gen(function* () { + const plugin = yield* PluginV2.Service + const hooks = yield* PluginHooks.Service + const host = yield* PluginHost.make(plugin) + yield* PluginPromise.fromPromise( + Plugin.define({ + id: "promise-ai-request", + setup: async (ctx) => { + await ctx.ai.hook("request", (event) => { + event.system.push(SystemPart.make("Promise hook")) + delete event.tools.echo + }) + }, + }), + ).effect(host) + const event: AIHooks["request"] = { + sessionID: SessionV2.ID.make("ses_promise_ai_request"), + agent: AgentV2.ID.make("build"), + model: Model.Ref.make({ providerID: Provider.ID.make("test"), id: Model.ID.make("model") }), + system: [SystemPart.make("Initial")], + messages: [Message.user("Hello")], + tools: { echo: { description: "Echo", input: { type: "object" } } }, + } + + yield* hooks.trigger("ai", "request", event) + + expect(event.system.map((part) => part.text)).toEqual(["Initial", "Promise hook"]) + expect(event.tools).toEqual({}) + }), + ) + it.effect("disposes a hook registration on request", () => Effect.gen(function* () { const agents = yield* AgentV2.Service diff --git a/packages/core/test/session-runner.test.ts b/packages/core/test/session-runner.test.ts index fc72e17e25..2585355fe7 100644 --- a/packages/core/test/session-runner.test.ts +++ b/packages/core/test/session-runner.test.ts @@ -3,7 +3,9 @@ import { LLMClient, LLMError, LLMEvent, + Message, Model, + SystemPart, ToolFailure, TransportReason, InvalidProviderOutputReason, @@ -42,6 +44,7 @@ import { SessionRunnerModel } from "@opencode-ai/core/session/runner/model" import { SessionRunnerSystemPrompt } from "@opencode-ai/core/session/runner/system-prompt" import { ToolRegistry } from "@opencode-ai/core/tool/registry" import { PluginSupervisor } from "@opencode-ai/core/plugin/supervisor" +import { PluginHooks } from "@opencode-ai/core/plugin/hooks" import { QuestionTool } from "@opencode-ai/core/tool/question" import { ToolOutputStore } from "@opencode-ai/core/tool-output-store" import { AgentV2 } from "@opencode-ai/core/agent" @@ -397,6 +400,7 @@ const it = testEffect( AgentV2.node, ToolRegistry.node, ToolRegistry.toolsNode, + PluginHooks.node, echoNode, SessionRunnerModel.node, InstructionBuiltIns.node, @@ -773,6 +777,32 @@ const verifyPartialFlushOnInterruption = (kind: FragmentKind) => }) describe("SessionRunnerLLM", () => { + it.effect("applies AI request hooks without exposing unavailable tools", () => + Effect.gen(function* () { + const session = yield* setup + const hooks = yield* PluginHooks.Service + yield* hooks.register("ai", "request", (event) => + Effect.sync(() => { + event.system = [SystemPart.make("Hooked system")] + event.messages = [Message.user("Hooked message")] + delete event.tools.echo + event.tools.unregistered = { description: "Unavailable", input: { type: "object" } } + }), + ) + yield* admit(session, "Original message") + responses = [reply.tool("call-removed", "echo", { text: "blocked" })] + + yield* session.resume(sessionID) + + expect(requests).toHaveLength(1) + expect(requests[0]?.system.map((part) => part.text)).toEqual(["Hooked system"]) + expect(requests[0]?.messages).toEqual([Message.user("Hooked message")]) + expect(requests[0]?.tools.map((tool) => tool.name)).not.toContain("echo") + expect(requests[0]?.tools.map((tool) => tool.name)).not.toContain("unregistered") + expect(executions).toEqual([]) + }), + ) + it.effect("advertises and executes a location registered tool", () => Effect.gen(function* () { const session = yield* setup diff --git a/packages/plugin/package.json b/packages/plugin/package.json index bae8482a36..c3a0c6ec99 100644 --- a/packages/plugin/package.json +++ b/packages/plugin/package.json @@ -25,6 +25,7 @@ ], "dependencies": { "@ai-sdk/provider": "3.0.8", + "@opencode-ai/ai": "workspace:*", "@opencode-ai/client": "workspace:*", "@opencode-ai/schema": "workspace:*", "@opencode-ai/sdk": "workspace:*", diff --git a/packages/plugin/src/v2/effect/README.md b/packages/plugin/src/v2/effect/README.md index f2526ec860..50ee83615b 100644 --- a/packages/plugin/src/v2/effect/README.md +++ b/packages/plugin/src/v2/effect/README.md @@ -80,14 +80,16 @@ yield * Hooks run sequentially in registration order. Later hooks observe mutations made by earlier hooks. -Session request context is mutable immediately before provider dispatch: +AI request context is mutable immediately before provider dispatch: ```ts yield * - ctx.session.hook("request", (event) => { - event.tools.read.description = "Read a file using narrow line ranges." - delete event.tools.write - }) + ctx.ai.hook("request", (event) => + Effect.sync(() => { + event.tools.read.description = "Read a file using narrow line ranges." + delete event.tools.write + }), + ) ``` ## Reloading A Domain diff --git a/packages/plugin/src/v2/effect/ai.ts b/packages/plugin/src/v2/effect/ai.ts new file mode 100644 index 0000000000..306d53cd17 --- /dev/null +++ b/packages/plugin/src/v2/effect/ai.ts @@ -0,0 +1,23 @@ +import type { Message, SystemPart } from "@opencode-ai/ai" +import type { Agent } from "@opencode-ai/schema/agent" +import type { Model } from "@opencode-ai/schema/model" +import type { Session } from "@opencode-ai/schema/session" +import type { JsonSchema } from "effect" +import type { Hooks } from "./registration.js" + +export interface AIRequest { + readonly sessionID: Session.ID + readonly agent: Agent.ID + readonly model: Model.Ref + system: Array + messages: Array + tools: Record +} + +export interface AIHooks { + readonly request: AIRequest +} + +export interface AIDomain { + readonly hook: Hooks +} diff --git a/packages/plugin/src/v2/effect/plugin.ts b/packages/plugin/src/v2/effect/plugin.ts index 3b0b6eb918..7be0b228e1 100644 --- a/packages/plugin/src/v2/effect/plugin.ts +++ b/packages/plugin/src/v2/effect/plugin.ts @@ -2,6 +2,7 @@ import type { PluginApi } from "@opencode-ai/client/effect/api" import type { Effect, Scope } from "effect" import type { PluginOptions } from "../options.js" import type { AgentDomain } from "./agent.js" +import type { AIDomain } from "./ai.js" import type { AISDKDomain } from "./aisdk.js" import type { CatalogDomain } from "./catalog.js" import type { CommandDomain } from "./command.js" @@ -15,6 +16,7 @@ import type { ToolDomain } from "./tool.js" export interface Context { readonly options: PluginOptions readonly agent: AgentDomain + readonly ai: AIDomain readonly aisdk: AISDKDomain readonly catalog: CatalogDomain readonly command: CommandDomain diff --git a/packages/plugin/src/v2/promise/README.md b/packages/plugin/src/v2/promise/README.md index f94f18d060..268a98a94b 100644 --- a/packages/plugin/src/v2/promise/README.md +++ b/packages/plugin/src/v2/promise/README.md @@ -85,10 +85,10 @@ await ctx.aisdk.hook("language", (event) => { }) ``` -Session request context is mutable immediately before provider dispatch: +AI request context is mutable immediately before provider dispatch: ```ts -await ctx.session.hook("request", (event) => { +await ctx.ai.hook("request", (event) => { event.tools.read.description = "Read a file using narrow line ranges." delete event.tools.write }) diff --git a/packages/plugin/src/v2/promise/ai.ts b/packages/plugin/src/v2/promise/ai.ts new file mode 100644 index 0000000000..306d53cd17 --- /dev/null +++ b/packages/plugin/src/v2/promise/ai.ts @@ -0,0 +1,23 @@ +import type { Message, SystemPart } from "@opencode-ai/ai" +import type { Agent } from "@opencode-ai/schema/agent" +import type { Model } from "@opencode-ai/schema/model" +import type { Session } from "@opencode-ai/schema/session" +import type { JsonSchema } from "effect" +import type { Hooks } from "./registration.js" + +export interface AIRequest { + readonly sessionID: Session.ID + readonly agent: Agent.ID + readonly model: Model.Ref + system: Array + messages: Array + tools: Record +} + +export interface AIHooks { + readonly request: AIRequest +} + +export interface AIDomain { + readonly hook: Hooks +} diff --git a/packages/plugin/src/v2/promise/plugin.ts b/packages/plugin/src/v2/promise/plugin.ts index 2e05b1fded..1928ca3672 100644 --- a/packages/plugin/src/v2/promise/plugin.ts +++ b/packages/plugin/src/v2/promise/plugin.ts @@ -1,6 +1,7 @@ import type { PluginApi } from "@opencode-ai/client/promise/api" import type { PluginOptions } from "../options.js" import type { AgentDomain } from "./agent.js" +import type { AIDomain } from "./ai.js" import type { AISDKDomain } from "./aisdk.js" import type { CatalogDomain } from "./catalog.js" import type { CommandDomain } from "./command.js" @@ -14,6 +15,7 @@ import type { ToolDomain } from "./tool.js" export interface Context { readonly options: PluginOptions readonly agent: AgentDomain + readonly ai: AIDomain readonly aisdk: AISDKDomain readonly catalog: CatalogDomain readonly command: CommandDomain