diff --git a/bun.lock b/bun.lock index fa2d6a102b..4f08f1b31b 100644 --- a/bun.lock +++ b/bun.lock @@ -902,6 +902,7 @@ "@napi-rs/canvas": "1.0.2", "@opencode-ai/ai": "workspace:*", "@opencode-ai/core": "workspace:*", + "@opencode-ai/plugin": "workspace:*", "@opentui/core": "catalog:", "effect": "catalog:", }, diff --git a/packages/core/src/plugin/host.ts b/packages/core/src/plugin/host.ts index eb977dc811..69007e7833 100644 --- a/packages/core/src/plugin/host.ts +++ b/packages/core/src/plugin/host.ts @@ -365,11 +365,14 @@ export const make = Effect.fn("PluginHost.make")(function* (plugin: PluginV2.Int }, }), ) - yield* Effect.forEach( - registrations, - (registration) => tools.register({ [registration.name]: registration.tool }, registration.options), - { discard: true }, - ).pipe(Effect.orDie) + yield* tools + .registerBatch( + registrations.map((registration) => ({ + tools: { [registration.name]: registration.tool }, + ...(registration.options === undefined ? {} : { options: registration.options }), + })), + ) + .pipe(Effect.orDie) return { dispose: Effect.void } }), hook: (name, callback) => { diff --git a/packages/core/src/tool/registry.ts b/packages/core/src/tool/registry.ts index 844e4896db..ac003775f4 100644 --- a/packages/core/src/tool/registry.ts +++ b/packages/core/src/tool/registry.ts @@ -1,7 +1,7 @@ export * as ToolRegistry from "./registry" import { ToolOutput, type ToolCall, type ToolDefinition, type ToolResultValue } from "@opencode-ai/ai" -import { Context, Effect, Layer, Scope } from "effect" +import { Context, Effect, Layer, Scope, Semaphore } from "effect" import type { AgentV2 } from "../agent" import { Image } from "../image" import { PermissionV2 } from "../permission" @@ -45,6 +45,13 @@ export interface Interface { tools: Readonly>, options?: Tools.RegisterOptions, ) => Effect.Effect + /** Internal atomic registration capability used by plugin transforms. */ + readonly registerBatch: ( + registrations: ReadonlyArray<{ + readonly tools: Readonly> + readonly options?: Tools.RegisterOptions + }>, + ) => Effect.Effect } export interface Materialization { @@ -107,6 +114,7 @@ const registryLayer = Layer.effect( readonly codemode: boolean } const local = new Map>() + const registrationLock = Semaphore.makeUnsafe(1) const settleTool = Effect.fn("ToolRegistry.settleTool")(function* (input: ExecuteInput, tool: AnyTool) { // Hooks fire only for hosted/local tools; provider-executed calls never reach settleTool. @@ -192,82 +200,111 @@ const registryLayer = Layer.effect( } }) - return Service.of({ - register: Effect.fn("ToolRegistry.register")(function* (tools, options) { - if (options?.namespace !== undefined) yield* validateNamespace(options.namespace) - const entries = registrationEntries(tools, options?.namespace) - if (entries.length === 0) return - const codemode = options?.codemode ?? true - const reserved = codemode ? undefined : entries.find((entry) => entry.key === "execute") - if (reserved) - return yield* Effect.fail( - new RegistrationError({ name: reserved.key, message: 'Tool name "execute" is reserved for CodeMode' }), - ) - yield* Effect.uninterruptible( + const registerBatch: Interface["registerBatch"] = Effect.fn("ToolRegistry.registerBatch")( + function* (registrations) { + const planned = yield* Effect.forEach(registrations, ({ tools, options }) => Effect.gen(function* () { - const token = {} - for (const entry of entries) - local.set(entry.key, [ - ...(local.get(entry.key) ?? []), - { - token, - registration: { - tool: entry.tool, - name: entry.name, - namespace: entry.namespace, - codemode, - }, - }, - ]) - yield* Effect.addFinalizer(() => - Effect.sync(() => { - for (const entry of entries) { - const registrations = - local.get(entry.key)?.filter((registration) => registration.token !== token) ?? [] - if (registrations.length > 0) local.set(entry.key, registrations) - else local.delete(entry.key) - } - }), - ) + if (options?.namespace !== undefined) yield* validateNamespace(options.namespace) + const entries = registrationEntries(tools, options?.namespace) + const codemode = options?.codemode ?? true + const reserved = codemode ? undefined : entries.find((entry) => entry.key === "execute") + if (reserved) + return yield* Effect.fail( + new RegistrationError({ name: reserved.key, message: 'Tool name "execute" is reserved for CodeMode' }), + ) + return { entries, codemode } }), ) - }), - materialize: Effect.fn("ToolRegistry.materialize")(function* (permissions) { - const direct = new Map() - const codemode = new Map() - const rules = permissions ?? [] - for (const [name, entries] of local) { - const registration = entries.at(-1)?.registration - if (!registration) continue - if (whollyDisabled(permission(registration.tool, name), rules)) continue - if (registration.codemode) codemode.set(name, registration) - else direct.set(name, registration) - } - const execute = - codemode.size > 0 && !whollyDisabled("execute", rules) ? ExecuteTool.create(codemode) : undefined - return { - definitions: [ - ...Array.from(direct, ([name, registration]) => definition(name, registration.tool)), - ...(execute ? [definition("execute", execute)] : []), - ], - settle: (input) => { - if (input.call.name === "execute" && execute) return settleTool(input, execute) - const registration = direct.get(input.call.name) - if (registration) return settleTool(input, registration.tool) - return Effect.succeed({ - result: { type: "error", value: `Unknown tool: ${input.call.name}` }, - error: { type: "tool.unknown", message: `Unknown tool: ${input.call.name}` }, - }) + if (planned.every((plan) => plan.entries.length === 0)) return + yield* Effect.uninterruptible( + registrationLock.withPermit( + Effect.gen(function* () { + const token = {} + for (const { entries, codemode } of planned) + for (const entry of entries) + local.set(entry.key, [ + ...(local.get(entry.key) ?? []), + { + token, + registration: { + tool: entry.tool, + name: entry.name, + namespace: entry.namespace, + codemode, + }, + }, + ]) + yield* Effect.addFinalizer(() => + registrationLock.withPermit( + Effect.sync(() => { + for (const { entries } of planned) + for (const entry of entries) { + const registrations = + local.get(entry.key)?.filter((registration) => registration.token !== token) ?? [] + if (registrations.length > 0) local.set(entry.key, registrations) + else local.delete(entry.key) + } + }), + ), + ) + }), + ), + ) + }, + ) + + return Service.of({ + register: Effect.fn("ToolRegistry.register")((tools, options) => + registerBatch([ + { + tools, + ...(options === undefined ? {} : { options }), }, - } - }), + ]), + ), + registerBatch, + materialize: Effect.fn("ToolRegistry.materialize")((permissions) => + registrationLock.withPermit( + Effect.sync(() => { + const direct = new Map() + const codemode = new Map() + const rules = permissions ?? [] + for (const [name, entries] of local) { + const registration = entries.at(-1)?.registration + if (!registration) continue + if (whollyDisabled(permission(registration.tool, name), rules)) continue + if (registration.codemode) codemode.set(name, registration) + else direct.set(name, registration) + } + const execute = + codemode.size > 0 && !whollyDisabled("execute", rules) ? ExecuteTool.create(codemode) : undefined + return { + definitions: [ + ...Array.from(direct, ([name, registration]) => definition(name, registration.tool)), + ...(execute ? [definition("execute", execute)] : []), + ], + settle: (input: ExecuteInput) => { + if (input.call.name === "execute" && execute) return settleTool(input, execute) + const registration = direct.get(input.call.name) + if (registration) return settleTool(input, registration.tool) + return Effect.succeed({ + result: { type: "error", value: `Unknown tool: ${input.call.name}` }, + error: { type: "tool.unknown", message: `Unknown tool: ${input.call.name}` }, + }) + }, + } + }), + ), + ), }) }), ) const layer = Layer.effect( Tools.Service, - Service.use((registry) => Effect.succeed(Tools.Service.of({ register: registry.register }))), + Service.use((registry) => + Effect.succeed(Tools.Service.of({ register: registry.register, registerBatch: registry.registerBatch })), + ), ).pipe(Layer.provideMerge(registryLayer)) function whollyDisabled(action: string, rules: PermissionV2.Ruleset) { diff --git a/packages/core/src/tool/tools.ts b/packages/core/src/tool/tools.ts index 06b920a484..cfc4dd95dc 100644 --- a/packages/core/src/tool/tools.ts +++ b/packages/core/src/tool/tools.ts @@ -10,6 +10,13 @@ export interface Interface { tools: Readonly>, options?: Tool.RegisterOptions, ) => Effect.Effect + /** Internal atomic registration capability used by plugin transforms. */ + readonly registerBatch: ( + registrations: ReadonlyArray<{ + readonly tools: Readonly> + readonly options?: Tool.RegisterOptions + }>, + ) => Effect.Effect } /** Narrow registration-only Location capability. */ diff --git a/packages/core/test/session-runner-tool-registry.test.ts b/packages/core/test/session-runner-tool-registry.test.ts index 6f7b13f6e1..19ba17d5e6 100644 --- a/packages/core/test/session-runner-tool-registry.test.ts +++ b/packages/core/test/session-runner-tool-registry.test.ts @@ -95,6 +95,21 @@ describe("ToolRegistry", () => { }), ) + it.effect("validates a registration batch before installing any tools", () => + Effect.gen(function* () { + const service = yield* ToolRegistry.Service + const error = yield* service + .registerBatch([ + { tools: { first: make() }, options: { codemode: false } }, + { tools: { second: make() }, options: { namespace: "invalid..namespace", codemode: false } }, + ]) + .pipe(Effect.flip) + + expect(error).toBeInstanceOf(Tool.RegistrationError) + expect((yield* service.materialize()).definitions).toEqual([]) + }), + ) + it.effect("filters disabled tools with edit aliases and ordered wildcard precedence", () => Effect.gen(function* () { const service = yield* ToolRegistry.Service diff --git a/packages/simulation/package.json b/packages/simulation/package.json index 8775fc2567..010097c4bd 100644 --- a/packages/simulation/package.json +++ b/packages/simulation/package.json @@ -24,6 +24,7 @@ "@napi-rs/canvas": "1.0.2", "@opencode-ai/core": "workspace:*", "@opencode-ai/ai": "workspace:*", + "@opencode-ai/plugin": "workspace:*", "@opentui/core": "catalog:", "effect": "catalog:" }, diff --git a/packages/simulation/src/backend/index.ts b/packages/simulation/src/backend/index.ts index 56c2ed87c2..a1044058db 100644 --- a/packages/simulation/src/backend/index.ts +++ b/packages/simulation/src/backend/index.ts @@ -1,5 +1,7 @@ import { LayerNode } from "@opencode-ai/core/effect/layer-node" +import { makeGlobalNode } from "@opencode-ai/core/effect/app-node" import { httpClient } from "@opencode-ai/core/effect/app-node-platform" +import { SdkPlugins } from "@opencode-ai/core/plugin/sdk" import { Config, Effect, Layer } from "effect" import { HttpClient } from "effect/unstable/http" import { DriveManifest } from "../manifest" @@ -41,7 +43,12 @@ export const simulationReplacements = Effect.fn("Simulation.replacements")(funct }), ), ) - return [[httpClient, networkLayer]] satisfies LayerNode.Replacements + const networkNode = makeGlobalNode({ + service: HttpClient.HttpClient, + layer: networkLayer, + deps: [SdkPlugins.node], + }) + return [[httpClient, networkNode]] satisfies LayerNode.Replacements }) export * as Simulation from "./index" diff --git a/packages/simulation/src/backend/openai.ts b/packages/simulation/src/backend/openai.ts index 248866e1a1..3130f00feb 100644 --- a/packages/simulation/src/backend/openai.ts +++ b/packages/simulation/src/backend/openai.ts @@ -29,6 +29,32 @@ type FinishReason = Extract Effect.Effect } +type ControlSocket = SimulationControlServer.Socket + +interface ToolController { + readonly socket: ControlSocket +} + +type ToolCompletion = + | { readonly type: "success"; readonly output: SimulationProtocol.Backend.ToolOutput } + | { readonly type: "failure"; readonly message: string } + +interface PendingToolInvocation { + readonly id: string + readonly notification: SimulationProtocol.Backend.ToolInvocation + readonly progress: Tool.Context["progress"] + readonly completion: Deferred.Deferred + readonly operations: Semaphore.Semaphore + update?: { readonly sequence: number; readonly fingerprint: string } +} + +interface ToolReconciliation { + readonly generation: number + readonly result: Deferred.Deferred> +} + +interface ToolRegistrationUpdate { + readonly generation: number + readonly registrations: ReadonlyArray +} + +interface ToolState { + readonly counter: number + readonly generation: number + readonly appliedGeneration: number + readonly controller?: ToolController + readonly registrations: ReadonlyArray + readonly pending: ReadonlyMap + readonly completed: ReadonlyMap + readonly activeOverlays: ReadonlySet + readonly reconciliation: ReadonlyMap +} + +interface ToolDriver { + readonly attach: ( + socket: ControlSocket, + registrations: ReadonlyArray, + ) => Effect.Effect<{ readonly attached: true }, ToolControllerError> + readonly update: ( + socket: ControlSocket, + params: SimulationProtocol.Backend.ToolUpdateParams, + ) => Effect.Effect + readonly finish: ( + socket: ControlSocket, + params: SimulationProtocol.Backend.ToolFinishParams, + ) => Effect.Effect + readonly fail: ( + socket: ControlSocket, + params: SimulationProtocol.Backend.ToolFailParams, + ) => Effect.Effect + readonly release: (socket: ControlSocket) => Effect.Effect + readonly shutdown: Effect.Effect +} + class InvocationNotFoundError extends Schema.TaggedErrorClass()( "SimulatedProvider.InvocationNotFoundError", { id: Schema.String, message: Schema.String }, @@ -60,7 +142,15 @@ class ControllerDisconnectedError extends Schema.TaggedErrorClass()( + "SimulatedProvider.ToolInvocationNotFoundError", + { id: Schema.String, message: Schema.String }, +) {} + +class ToolControllerError extends Schema.TaggedErrorClass()( + "SimulatedProvider.ToolControllerError", + { message: Schema.String }, +) {} export const layerDrive = (options: { readonly endpoint: string }) => Layer.effect( @@ -180,13 +270,18 @@ export const layerDrive = (options: { readonly endpoint: string }) => const fibers = yield* FiberSet.make() const activeController = yield* Ref.make | undefined>(undefined) const controllerLock = yield* Semaphore.make(1) + const tools = yield* makeToolDriver() + yield* Effect.addFinalizer(() => tools.shutdown) yield* SimulationControlServer.start({ endpoint: options.endpoint, label: "opencode drive backend websocket", data: () => ({}), decode: SimulationProtocol.Backend.decodeRequestEffect, - handle: (socket, request) => handle(driver, fibers, activeController, controllerLock, socket, request), - close: (socket) => releaseController(activeController, controllerLock, socket), + handle: (socket, request) => handle(driver, tools, fibers, activeController, controllerLock, socket, request), + close: (socket) => + Effect.all([releaseController(activeController, controllerLock, socket), tools.release(socket)], { + discard: true, + }), }) yield* Effect.sync(() => process.stderr.write(`opencode drive backend websocket: ${options.endpoint}\n`)) @@ -205,6 +300,7 @@ export const layerDrive = (options: { readonly endpoint: string }) => function handle( driver: Driver, + tools: ToolDriver, fibers: FiberSet.FiberSet, activeController: Ref.Ref | undefined>, controllerLock: Semaphore.Semaphore, @@ -259,9 +355,539 @@ function handle( return driver.disconnect(request.params.id).pipe(Effect.as({ ok: true })) case "llm.pending": return driver.pending().pipe(Effect.map((invocations) => ({ invocations }))) + case "tool.attach": + return tools.attach(socket, request.params.tools) + case "tool.update": + return tools.update(socket, request.params).pipe(Effect.as({ ok: true })) + case "tool.finish": + return tools.finish(socket, request.params).pipe(Effect.as({ ok: true })) + case "tool.fail": + return tools.fail(socket, request.params).pipe(Effect.as({ ok: true })) } } +const makeToolDriver = Effect.fn("SimulatedProvider.makeToolDriver")(function* () { + const completedRetention = 256 + const plugins = yield* SdkPlugins.Service + const state = yield* Ref.make({ + counter: 0, + generation: 0, + appliedGeneration: 0, + registrations: [], + pending: new Map(), + completed: new Map(), + activeOverlays: new Set(), + reconciliation: new Map(), + }) + const registrationUpdates = yield* PubSub.unbounded() + const lock = yield* Semaphore.make(1) + const attachmentLock = yield* Semaphore.make(1) + + const remove = (current: ToolState, id: string) => { + const pending = new Map(current.pending) + pending.delete(id) + return { ...current, pending } + } + + const complete = (current: ToolState, id: string, completion: string) => { + const completed = new Map(current.completed) + completed.set(id, completion) + if (completed.size > completedRetention) { + const oldest = completed.keys().next().value + if (oldest !== undefined) completed.delete(oldest) + } + return { ...remove(current, id), completed } + } + + const notify = ( + socket: ControlSocket, + method: "tool.invocation" | "tool.cancel", + params: SimulationProtocol.Backend.ToolInvocation | SimulationProtocol.Backend.ToolCancellation, + ) => + Effect.sync(() => { + socket.send(JSON.stringify({ jsonrpc: "2.0", method, params })) + }) + + const requireController = (socket: ControlSocket) => + Effect.gen(function* () { + const current = yield* Ref.get(state) + if (current.controller?.socket === socket) return current + return yield* Effect.fail(new ToolControllerError({ message: "Drive tool controller is not attached" })) + }) + + const requireInvocation = (socket: ControlSocket, id: string) => + Effect.gen(function* () { + const current = yield* requireController(socket) + const invocation = current.pending.get(id) + if (invocation) return { current, invocation } + return yield* Effect.fail( + new ToolInvocationNotFoundError({ + id, + message: `Simulated tool invocation not found or already finished: ${id}`, + }), + ) + }) + + const cancel = (id: string) => + Effect.gen(function* () { + const invocation = yield* lock.withPermit(Ref.get(state).pipe(Effect.map((current) => current.pending.get(id)))) + if (!invocation) return + yield* invocation.operations.withPermit( + lock.withPermit( + Effect.gen(function* () { + const current = yield* Ref.get(state) + if (current.pending.get(id) !== invocation) return + yield* Ref.set(state, remove(current, id)) + if (current.controller && !current.controller.socket.data.closed) + yield* notify(current.controller.socket, "tool.cancel", { id, reason: "interrupted" }) + }), + ), + ) + }) + + const invoke = ( + registrationGeneration: number, + name: string, + input: unknown, + context: Tool.Context, + ): Effect.Effect => + Effect.gen(function* () { + const encoded = yield* Schema.decodeUnknownEffect(Schema.Json)(input).pipe( + Effect.mapError((error) => new Tool.Failure({ message: `Simulated tool input is not JSON: ${error.message}` })), + ) + const invocation = yield* Effect.uninterruptibleMask((restore) => + attachmentLock + .withPermit( + lock.withPermit( + Effect.gen(function* () { + const current = yield* Ref.get(state) + if (current.generation !== registrationGeneration) + yield* Effect.fail( + new Tool.Failure({ message: `Simulated tool registration is no longer active: ${name}` }), + ) + const id = `tool_${current.counter + 1}` + const completion = yield* Deferred.make() + const notification: SimulationProtocol.Backend.ToolInvocation = { + id, + name, + input: encoded, + context: { + sessionID: context.sessionID, + agent: context.agent, + messageID: context.messageID, + callID: context.callID, + }, + } + const pending: PendingToolInvocation = { + id, + notification, + progress: context.progress, + completion, + operations: Semaphore.makeUnsafe(1), + } + yield* Ref.set(state, { + counter: current.counter + 1, + generation: current.generation, + appliedGeneration: current.appliedGeneration, + ...(current.controller === undefined ? {} : { controller: current.controller }), + registrations: current.registrations, + pending: new Map(current.pending).set(id, pending), + completed: current.completed, + activeOverlays: current.activeOverlays, + reconciliation: current.reconciliation, + }) + if (current.controller && !current.controller.socket.data.closed) + yield* notify(current.controller.socket, "tool.invocation", notification) + return pending + }), + ), + ) + .pipe( + Effect.flatMap((pending) => + restore(Deferred.await(pending.completion)).pipe( + Effect.onInterrupt(() => cancel(pending.id)), + Effect.ensuring( + lock.withPermit( + Ref.update(state, (current) => + current.pending.get(pending.id) === pending ? remove(current, pending.id) : current, + ), + ), + ), + ), + ), + ), + ) + if (invocation.type === "success") return invocation.output + return yield* Effect.fail(new Tool.Failure({ message: invocation.message })) + }) + + yield* plugins.register( + Plugin.define({ + id: "opencode.simulation.tools", + effect: (ctx) => + Effect.gen(function* () { + const scope = yield* Scope.Scope + const token = {} + const registrationLock = Semaphore.makeUnsafe(1) + let currentScope: Scope.Closeable | undefined + const reconcile = ( + generation: number, + nextRegistrations: ReadonlyArray, + ) => + registrationLock.withPermit( + Effect.gen(function* () { + const nextScope = yield* Scope.fork(scope) + const applied = yield* Effect.exit( + ctx.tool + .transform((draft) => { + for (const registration of nextRegistrations) + draft.add( + registration.name, + Tool.make({ + description: registration.description, + jsonSchema: registration.inputSchema, + ...(registration.outputSchema === undefined + ? {} + : { outputSchema: registration.outputSchema }), + ...(registration.permission === undefined ? {} : { permission: registration.permission }), + execute: (input, context) => + invoke( + generation, + SimulationProtocol.Backend.exposedToolName(registration), + input, + context, + ), + }), + registration.options, + ) + }) + .pipe(Scope.provide(nextScope)), + ) + if (Exit.isFailure(applied)) { + yield* Scope.close(nextScope, applied) + yield* Effect.failCause(applied.cause) + } + const previousScope = currentScope + currentScope = nextScope + if (previousScope) yield* Scope.close(previousScope, Exit.void) + }), + ) + + const acknowledge = (generation: number, result: Exit.Exit) => + lock.withPermit( + Effect.gen(function* () { + const current = yield* Ref.get(state) + const reconciliation = current.reconciliation.get(token) + if (reconciliation?.generation !== generation) return + yield* Deferred.succeed(reconciliation.result, result) + }), + ) + + const initialized = yield* lock.withPermit( + Effect.gen(function* () { + const subscription = yield* PubSub.subscribe(registrationUpdates) + const current = yield* Ref.get(state) + const activeOverlays = new Set(current.activeOverlays).add(token) + yield* Ref.set(state, { ...current, activeOverlays }) + return { subscription, generation: current.generation, registrations: current.registrations } + }), + ) + yield* Effect.addFinalizer(() => + lock.withPermit( + Effect.gen(function* () { + const current = yield* Ref.get(state) + const activeOverlays = new Set(current.activeOverlays) + activeOverlays.delete(token) + const reconciliation = new Map(current.reconciliation) + const pending = reconciliation.get(token) + reconciliation.delete(token) + yield* Ref.set(state, { ...current, activeOverlays, reconciliation }) + if (pending) yield* Deferred.succeed(pending.result, Exit.void) + }), + ), + ) + yield* reconcile(initialized.generation, initialized.registrations) + yield* Stream.fromEffectRepeat(PubSub.take(initialized.subscription)).pipe( + Stream.runForEach((update) => + Effect.gen(function* () { + const result = yield* Effect.exit(reconcile(update.generation, update.registrations)) + yield* acknowledge(update.generation, result) + }), + ), + Effect.forkScoped({ startImmediately: true }), + ) + }), + }), + ) + + const reconcileRegistrations = ( + controller: ToolController | undefined, + registrations: ReadonlyArray, + targetGeneration?: number, + ) => + Effect.gen(function* () { + const reconciliation = yield* lock.withPermit( + Effect.gen(function* () { + const current = yield* Ref.get(state) + const generation = targetGeneration ?? current.generation + 1 + const reconciliation = new Map() + for (const token of current.activeOverlays) + reconciliation.set(token, { + generation, + result: Deferred.makeUnsafe>(), + }) + yield* Ref.set(state, { + counter: current.counter, + generation, + appliedGeneration: current.appliedGeneration, + ...(controller === undefined ? {} : { controller }), + registrations, + pending: current.pending, + completed: current.completed, + activeOverlays: current.activeOverlays, + reconciliation, + }) + yield* PubSub.publish(registrationUpdates, { generation, registrations }) + return { generation, previous: current, reconciliation } + }), + ) + const results = yield* Effect.forEach(reconciliation.reconciliation.values(), (item) => + Deferred.await(item.result), + ) + const failure = results.find(Exit.isFailure) + yield* lock.withPermit( + Ref.update(state, (current) => + current.generation === reconciliation.generation + ? { + ...current, + ...(failure === undefined ? { appliedGeneration: reconciliation.generation } : {}), + reconciliation: new Map(), + } + : current, + ), + ) + return { failure, previous: reconciliation.previous } + }) + + const attach: ToolDriver["attach"] = (socket, registrations) => + attachmentLock.withPermit( + Effect.gen(function* () { + if (socket.data.closed) + return yield* Effect.fail( + new ToolControllerError({ message: "Drive tool controller disconnected before attachment" }), + ) + const current = yield* Ref.get(state) + if (current.controller && current.controller.socket !== socket && !current.controller.socket.data.closed) + return yield* Effect.fail( + new ToolControllerError({ message: "Another Drive tool controller is already attached" }), + ) + const controller: ToolController = { socket } + const replay = current.controller?.socket !== socket + const sameRegistrations = + current.appliedGeneration === current.generation && + fingerprintJson(current.registrations) === fingerprintJson(registrations) + if (replay && current.pending.size > 0 && !sameRegistrations) + return yield* Effect.fail( + new ToolControllerError({ + message: "A reconnecting Drive tool controller must settle pending invocations before replacing tools", + }), + ) + if (sameRegistrations) { + yield* lock.withPermit( + Effect.gen(function* () { + const current = yield* Ref.get(state) + yield* Ref.set(state, { ...current, controller }) + if (replay) + yield* Effect.forEach( + current.pending.values(), + (invocation) => notify(socket, "tool.invocation", invocation.notification), + { discard: true }, + ) + }), + ) + return { attached: true as const } + } + const reconciled = yield* reconcileRegistrations(controller, registrations) + if (reconciled.failure) { + const rolledBack = yield* reconcileRegistrations( + reconciled.previous.controller, + reconciled.previous.registrations, + reconciled.previous.generation, + ) + return yield* Effect.fail( + new ToolControllerError({ + message: rolledBack.failure + ? `Failed to apply and restore simulated tools: ${Cause.pretty(reconciled.failure.cause)}; ${Cause.pretty(rolledBack.failure.cause)}` + : `Failed to apply simulated tools: ${Cause.pretty(reconciled.failure.cause)}`, + }), + ) + } + if (replay) + yield* lock.withPermit( + Effect.gen(function* () { + const current = yield* Ref.get(state) + if (current.controller?.socket !== socket) return + yield* Effect.forEach( + current.pending.values(), + (invocation) => notify(socket, "tool.invocation", invocation.notification), + { discard: true }, + ) + }), + ) + return { attached: true as const } + }), + ) + + const update: ToolDriver["update"] = (socket, params) => + Effect.gen(function* () { + const { invocation } = yield* lock.withPermit(requireInvocation(socket, params.id)) + yield* invocation.operations.withPermit( + Effect.gen(function* () { + const current = yield* lock.withPermit(Ref.get(state)) + if (current.pending.get(params.id) !== invocation) + yield* Effect.fail( + new ToolInvocationNotFoundError({ + id: params.id, + message: `Simulated tool invocation not found or already finished: ${params.id}`, + }), + ) + const fingerprint = fingerprintJson(params.update) + const applied = invocation.update + if (applied?.sequence === params.sequence && applied.fingerprint === fingerprint) return + if (applied?.sequence === params.sequence) + yield* Effect.fail( + new ToolControllerError({ + message: `Simulated tool update sequence ${params.sequence} was reused with different progress`, + }), + ) + const expected = applied === undefined ? 0 : applied.sequence + 1 + if (params.sequence !== expected) + yield* Effect.fail( + new ToolControllerError({ + message: `Expected simulated tool update sequence ${expected}, received ${params.sequence}`, + }), + ) + yield* invocation.progress(params.update) + invocation.update = { sequence: params.sequence, fingerprint } + }), + ) + }) + + const settle = (socket: ControlSocket, id: string, completion: ToolCompletion) => + Effect.gen(function* () { + const current = yield* lock.withPermit(requireController(socket)) + const fingerprint = fingerprintJson(completion) + const invocation = current.pending.get(id) + if (!invocation) { + if (current.completed.get(id) === fingerprint) return + yield* Effect.fail( + new ToolInvocationNotFoundError({ + id, + message: `Simulated tool invocation not found or already finished: ${id}`, + }), + ) + return + } + yield* invocation.operations.withPermit( + Effect.uninterruptible( + lock.withPermit( + Effect.gen(function* () { + const current = yield* Ref.get(state) + if (current.pending.get(id) !== invocation) { + if (current.completed.get(id) === fingerprint) return + yield* Effect.fail( + new ToolInvocationNotFoundError({ + id, + message: `Simulated tool invocation not found or already finished: ${id}`, + }), + ) + return + } + yield* Ref.set(state, complete(current, id, fingerprint)) + yield* Deferred.succeed(invocation.completion, completion) + }), + ), + ), + ) + }) + + const release: ToolDriver["release"] = (socket) => + attachmentLock.withPermit( + lock.withPermit( + Ref.update( + state, + (current): ToolState => + current.controller?.socket !== socket + ? current + : { + counter: current.counter, + generation: current.generation, + appliedGeneration: current.appliedGeneration, + registrations: current.registrations, + pending: current.pending, + completed: current.completed, + activeOverlays: current.activeOverlays, + reconciliation: current.reconciliation, + }, + ), + ), + ) + + const shutdown = attachmentLock.withPermit( + Effect.gen(function* () { + yield* reconcileRegistrations(undefined, []) + const current = yield* lock.withPermit( + Effect.gen(function* () { + const current = yield* Ref.get(state) + yield* Ref.set(state, { + counter: current.counter, + generation: current.generation, + appliedGeneration: current.appliedGeneration, + registrations: [], + pending: new Map(), + completed: current.completed, + activeOverlays: new Set(), + reconciliation: new Map(), + }) + return current + }), + ) + yield* Effect.forEach( + current.pending.values(), + (invocation) => + Deferred.succeed(invocation.completion, { + type: "failure", + message: "Simulated tool controller shut down", + }), + { discard: true }, + ) + yield* PubSub.shutdown(registrationUpdates) + }), + ) + + return { + attach, + update, + finish: (socket, params) => settle(socket, params.id, { type: "success", output: params.output }), + fail: (socket, params) => settle(socket, params.id, { type: "failure", message: params.message }), + release, + shutdown, + } satisfies ToolDriver +}) + +function canonicalJson(value: unknown): string { + if (Array.isArray(value)) return `[${value.map(canonicalJson).join(",")}]` + if (typeof value === "object" && value !== null) { + const entries = Object.entries(value).sort(([left], [right]) => left.localeCompare(right)) + return `{${entries.map(([key, item]) => `${JSON.stringify(key)}:${canonicalJson(item)}`).join(",")}}` + } + return JSON.stringify(value) ?? String(value) +} + +function fingerprintJson(value: unknown) { + return createHash("sha256").update(canonicalJson(value)).digest("base64url") +} + function releaseController( activeController: Ref.Ref | undefined>, controllerLock: Semaphore.Semaphore, diff --git a/packages/simulation/src/protocol/index.ts b/packages/simulation/src/protocol/index.ts index a373be5ecf..d97b808686 100644 --- a/packages/simulation/src/protocol/index.ts +++ b/packages/simulation/src/protocol/index.ts @@ -371,11 +371,29 @@ export namespace Backend { "llm.disconnect", "llm.pending", "llm.request", + "llm.tool-input-delta", + "tool.attach", + "tool.update", + "tool.finish", + "tool.fail", + "tool.invocation", + "tool.cancel", ] as const satisfies ReadonlyArray export const Item = Schema.Union([ Schema.Struct({ type: Schema.Literal("textDelta"), text: Schema.String }), Schema.Struct({ type: Schema.Literal("reasoningDelta"), text: Schema.String }), + Schema.Struct({ + type: Schema.Literal("toolInputStart"), + index: Schema.Number, + id: Schema.String, + name: Schema.String, + }), + Schema.Struct({ + type: Schema.Literal("toolInputDelta"), + index: Schema.Number, + text: Schema.String, + }), Schema.Struct({ type: Schema.Literal("toolCall"), index: Schema.Number, @@ -390,6 +408,115 @@ export namespace Backend { export const FinishReason = Schema.Literals(["stop", "tool-calls", "length", "content-filter"]) export type FinishReason = Schema.Schema.Type + export const ToolContent = Schema.Union([ + Schema.Struct({ type: Schema.Literal("text"), text: Schema.String }), + Schema.Struct({ + type: Schema.Literal("file"), + data: Schema.String, + mime: Schema.NonEmptyString, + name: Schema.optionalKey(Schema.String), + }), + ]) + export type ToolContent = Schema.Schema.Type + + const ToolName = Schema.NonEmptyString.check( + Schema.makeFilter((name) => + /^[A-Za-z][A-Za-z0-9_-]{0,63}$/.test(name) ? undefined : "simulated tool names must be provider-safe", + ), + ) + const ToolNamespace = Schema.NonEmptyString.check( + Schema.makeFilter((namespace) => + namespace.split(".").every((segment) => /^[A-Za-z][A-Za-z0-9_-]{0,63}$/.test(segment)) + ? undefined + : "simulated tool namespaces must contain provider-safe segments", + ), + ) + + export const ToolRegistration = Schema.Struct({ + name: ToolName, + description: Schema.String, + inputSchema: Schema.Record(Schema.String, Schema.Json), + outputSchema: Schema.optionalKey(Schema.Record(Schema.String, Schema.Json)), + permission: Schema.optionalKey(Schema.NonEmptyString), + options: Schema.optionalKey( + Schema.Struct({ + namespace: Schema.optionalKey(ToolNamespace), + codemode: Schema.optionalKey(Schema.Boolean), + }), + ), + }) + export interface ToolRegistration extends Schema.Schema.Type {} + + export const ToolAttachParams = Schema.Struct({ + tools: Schema.Array(ToolRegistration).check( + Schema.makeFilter((tools) => { + const names = tools.map(exposedToolName) + if (names.some((name) => !/^[A-Za-z][A-Za-z0-9_-]{0,63}$/.test(name))) + return "simulated tool names including namespaces must be provider-safe" + if (new Set(names).size !== names.length) return "simulated tool registrations must have unique exposed names" + if ( + tools.some( + (tool) => + tool.name === "execute" && tool.options?.namespace === undefined && tool.options?.codemode === false, + ) + ) + return 'direct simulated tool name "execute" is reserved' + return undefined + }), + ), + }) + export interface ToolAttachParams extends Schema.Schema.Type {} + + export function exposedToolName(registration: ToolRegistration) { + return registration.options?.namespace === undefined + ? registration.name + : `${registration.options.namespace.replaceAll(".", "_")}_${registration.name}` + } + + export const ToolProgress = Schema.Struct({ + structured: Schema.Record(Schema.String, Schema.Json), + content: Schema.optionalKey(Schema.Array(ToolContent)), + }) + export interface ToolProgress extends Schema.Schema.Type {} + + export const ToolOutput = Schema.Struct({ + structured: Schema.Json, + content: Schema.Array(ToolContent), + }) + export interface ToolOutput extends Schema.Schema.Type {} + + export const ToolUpdateParams = Schema.Struct({ + id: Schema.String, + sequence: Schema.Int.check(Schema.isGreaterThanOrEqualTo(0)), + update: ToolProgress, + }) + export interface ToolUpdateParams extends Schema.Schema.Type {} + + export const ToolFinishParams = Schema.Struct({ id: Schema.String, output: ToolOutput }) + export interface ToolFinishParams extends Schema.Schema.Type {} + + export const ToolFailParams = Schema.Struct({ id: Schema.String, message: Schema.String }) + export interface ToolFailParams extends Schema.Schema.Type {} + + export const ToolInvocation = Schema.Struct({ + id: Schema.String, + name: Schema.String, + input: Schema.Json, + context: Schema.Struct({ + sessionID: Schema.String, + agent: Schema.String, + messageID: Schema.String, + callID: Schema.String, + }), + }) + export interface ToolInvocation extends Schema.Schema.Type {} + + export const ToolCancellation = Schema.Struct({ + id: Schema.String, + reason: Schema.Literal("interrupted"), + }) + export interface ToolCancellation extends Schema.Schema.Type {} + export const ChunkParams = Schema.Struct({ id: Schema.String, items: Schema.Array(Item) }) export interface ChunkParams extends Schema.Schema.Type {} @@ -407,6 +534,10 @@ export namespace Backend { Schema.Struct({ ...JsonRpc.RequestFields, method: Schema.Literal("llm.chunk"), params: ChunkParams }), Schema.Struct({ ...JsonRpc.RequestFields, method: Schema.Literal("llm.finish"), params: FinishParams }), Schema.Struct({ ...JsonRpc.RequestFields, method: Schema.Literal("llm.disconnect"), params: DisconnectParams }), + Schema.Struct({ ...JsonRpc.RequestFields, method: Schema.Literal("tool.attach"), params: ToolAttachParams }), + Schema.Struct({ ...JsonRpc.RequestFields, method: Schema.Literal("tool.update"), params: ToolUpdateParams }), + Schema.Struct({ ...JsonRpc.RequestFields, method: Schema.Literal("tool.finish"), params: ToolFinishParams }), + Schema.Struct({ ...JsonRpc.RequestFields, method: Schema.Literal("tool.fail"), params: ToolFailParams }), Schema.Struct({ ...JsonRpc.RequestFields, method: Schema.Literals(["llm.attach", "llm.pending"]), diff --git a/packages/simulation/test/openai.test.ts b/packages/simulation/test/openai.test.ts index 069b4feed4..a4a5b16fab 100644 --- a/packages/simulation/test/openai.test.ts +++ b/packages/simulation/test/openai.test.ts @@ -33,6 +33,35 @@ test("encodes every simulated provider event as OpenAI SSE", async () => { ) }) +test("encodes provider-neutral partial tool input as OpenAI deltas", async () => { + const provider: SimulatedProvider.Interface = { + stream: () => + Stream.make( + { type: "toolInputStart", index: 0, id: "call_lookup", name: "lookup" }, + { type: "toolInputDelta", index: 0, text: '{"query":' }, + { type: "toolInputDelta", index: 0, text: '"OpenCode"}' }, + { type: "finish", reason: "tool-calls" }, + ), + } + const url = new URL(DEFAULT_BASE_URL + PATH) + const request = HttpClientRequest.post(url).pipe(HttpClientRequest.bodyJsonUnsafe({ model: "gpt-5" })) + const matched = SimulationOpenAI.route(provider).match(request, url) + if (!matched) throw new Error("The simulated OpenAI route did not match") + + const body = await Effect.runPromise(matched.pipe(Effect.flatMap((response) => response.text))) + + expect(body).toBe( + [ + 'data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_lookup","function":{"name":"lookup","arguments":""}}]}}]}', + 'data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\\"query\\":"}}]}}]}', + 'data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\\"OpenCode\\"}"}}]}}]}', + 'data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}]}', + "data: [DONE]", + "", + ].join("\n\n"), + ) +}) + test("rejects malformed intercepted OpenAI JSON as an HTTP client error", async () => { const provider: SimulatedProvider.Interface = { stream: () => Stream.empty } const url = new URL(DEFAULT_BASE_URL + PATH) diff --git a/packages/simulation/test/protocol.test.ts b/packages/simulation/test/protocol.test.ts index af30dde25e..b293f35439 100644 --- a/packages/simulation/test/protocol.test.ts +++ b/packages/simulation/test/protocol.test.ts @@ -97,6 +97,96 @@ test("decodes semantic UI snapshots", () => { expect(() => decode({ format: "opencode-ui-snapshot-v1", nodes })).toThrow() }) +test("decodes the simulated tool lifecycle", () => { + const registration = { + name: "lookup", + description: "Look up a value", + inputSchema: { + type: "object", + properties: { query: { type: "string" } }, + required: ["query"], + additionalProperties: false, + }, + outputSchema: { type: "object" }, + permission: "lookup", + options: { codemode: false }, + } + expect( + Backend.decodeRequest({ + jsonrpc: "2.0", + id: 1, + method: "tool.attach", + params: { tools: [registration] }, + }), + ).toMatchObject({ method: "tool.attach", params: { tools: [registration] } }) + expect( + Backend.decodeRequest({ + jsonrpc: "2.0", + id: 2, + method: "tool.update", + params: { + id: "tool_1", + sequence: 0, + update: { structured: { phase: "searching" }, content: [{ type: "text", text: "Searching" }] }, + }, + }), + ).toMatchObject({ method: "tool.update" }) + expect( + Backend.decodeRequest({ + jsonrpc: "2.0", + id: 3, + method: "tool.finish", + params: { + id: "tool_1", + output: { structured: { answer: 42 }, content: [{ type: "text", text: "42" }] }, + }, + }), + ).toMatchObject({ method: "tool.finish" }) + expect( + Backend.decodeRequest({ + jsonrpc: "2.0", + id: 4, + method: "tool.fail", + params: { id: "tool_2", message: "lookup failed" }, + }), + ).toMatchObject({ method: "tool.fail" }) + expect(() => + Backend.decodeRequest({ + jsonrpc: "2.0", + id: 5, + method: "tool.attach", + params: { tools: [registration, registration] }, + }), + ).toThrow() + for (const invalid of [ + { ...registration, name: "1lookup" }, + { ...registration, options: { namespace: "bad group", codemode: false } }, + { ...registration, name: "execute", options: { codemode: false } }, + { ...registration, options: { namespace: "a".repeat(64), codemode: false } }, + ]) + expect(() => + Backend.decodeRequest({ + jsonrpc: "2.0", + id: 6, + method: "tool.attach", + params: { tools: [invalid] }, + }), + ).toThrow() + expect(() => + Backend.decodeRequest({ + jsonrpc: "2.0", + id: 7, + method: "tool.attach", + params: { + tools: [ + { ...registration, name: "b_c", options: { namespace: "a", codemode: false } }, + { ...registration, name: "c", options: { namespace: "a.b", codemode: false } }, + ], + }, + }), + ).toThrow() +}) + const params: Handshake.Params = { client: { name: "opencode-drive", version: "test" }, expectedRole: "ui", diff --git a/packages/simulation/test/simulated-provider.test.ts b/packages/simulation/test/simulated-provider.test.ts index 650fd2d172..271ef58009 100644 --- a/packages/simulation/test/simulated-provider.test.ts +++ b/packages/simulation/test/simulated-provider.test.ts @@ -1,5 +1,24 @@ import { expect, test } from "bun:test" -import { Deferred, Effect, Fiber, Queue, Stream } from "effect" +import { mkdir, mkdtemp, rm } from "node:fs/promises" +import { tmpdir } from "node:os" +import { join } from "node:path" +import { AgentV2 } from "@opencode-ai/core/agent" +import { Config } from "@opencode-ai/core/config" +import { Database } from "@opencode-ai/core/database/database" +import { makeGlobalNode } from "@opencode-ai/core/effect/app-node" +import { AppNodeBuilder } from "@opencode-ai/core/effect/app-node-builder" +import { LayerNode } from "@opencode-ai/core/effect/layer-node" +import { EventV2 } from "@opencode-ai/core/event" +import { Location } from "@opencode-ai/core/location" +import { LocationServiceMap } from "@opencode-ai/core/location-services" +import { SdkPlugins } from "@opencode-ai/core/plugin/sdk" +import { PluginSupervisor } from "@opencode-ai/core/plugin/supervisor" +import { AbsolutePath } from "@opencode-ai/core/schema" +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/effect" +import { Deferred, Effect, Fiber, Layer, Queue, Stream } from "effect" import type { Scope } from "effect/Scope" import { SimulatedProvider } from "../src/backend/simulated-provider" import { availableEndpoint, connect } from "./fixture/websocket" @@ -125,7 +144,7 @@ test("replaces the previous attached controller", async () => { ) expect(yield* Queue.take(secondMessages)).toMatchObject({ id: 2, result: { ok: true } }) expect(Array.from(yield* Fiber.join(response))).toEqual([{ type: "finish", reason: "stop" }]) - }).pipe(Effect.provide(SimulatedProvider.layerDrive({ endpoint })), Effect.scoped), + }).pipe(Effect.provide(providerLayer(endpoint)), Effect.scoped), ) }) @@ -194,6 +213,454 @@ test("fails the provider stream when Drive disconnects the invocation", async () ) }) +test("controls arbitrary tools through scoped SDK overlays", async () => { + const endpoint = availableEndpoint() + const directory = await mkdtemp(join(tmpdir(), "opencode-simulated-tools-")) + const secondDirectory = join(directory, "second") + await mkdir(secondDirectory) + try { + await Effect.runPromise( + Effect.gen(function* () { + const socket = yield* connect(endpoint) + const messages = yield* messagesFrom(socket) + const plugins = yield* SdkPlugins.Service + let activations = 0 + yield* plugins.register( + Plugin.define({ + id: "opencode.simulation.test.activation-count", + effect: () => Effect.sync(() => void activations++), + }), + ) + const registration = { + name: "lookup", + description: "Look up a value", + inputSchema: { + type: "object", + properties: { query: { type: "string" } }, + required: ["query"], + additionalProperties: false, + }, + outputSchema: { type: "object" }, + options: { codemode: false }, + } + const locations = yield* LocationServiceMap.Service + const [primary, secondary] = yield* Effect.all([ + Layer.build(locations.get(Location.Ref.make({ directory: AbsolutePath.make(directory) }))), + Layer.build(locations.get(Location.Ref.make({ directory: AbsolutePath.make(secondDirectory) }))), + ]) + yield* Effect.forEach([primary, secondary], (context) => + PluginSupervisor.Service.use((supervisor) => supervisor.flush).pipe(Effect.provide(context)), + ) + expect(activations).toBe(2) + + yield* Effect.gen(function* () { + socket.send( + JSON.stringify({ + jsonrpc: "2.0", + id: 1, + method: "tool.attach", + params: { tools: [registration] }, + }), + ) + expect(yield* Queue.take(messages)).toMatchObject({ id: 1, result: { attached: true } }) + const registry = yield* ToolRegistry.Service + const materialized = yield* registry.materialize() + expect(materialized.definitions).toContainEqual( + expect.objectContaining({ name: "lookup", description: "Look up a value" }), + ) + const secondaryMaterialized = yield* ToolRegistry.Service.use((secondaryRegistry) => + secondaryRegistry.materialize(), + ).pipe(Effect.provide(secondary)) + expect(secondaryMaterialized.definitions).toContainEqual( + expect.objectContaining({ name: "lookup", description: "Look up a value" }), + ) + const progress: ToolRegistry.Progress[] = [] + const settle = (callID: string, query: string) => + materialized.settle({ + sessionID: SessionV2.ID.make("ses_simulated_tools"), + agent: AgentV2.ID.make("build"), + messageID: SessionMessage.ID.make("msg_simulated_tools"), + progress: (update) => Effect.sync(() => progress.push(update)), + call: { + type: "tool-call", + id: callID, + name: "lookup", + input: { query }, + }, + }) + + const successful = yield* settle("call_success", "answer").pipe(Effect.forkScoped) + const successInvocation = yield* takeToolInvocation(messages) + expect(successInvocation.params).toMatchObject({ + name: "lookup", + input: { query: "answer" }, + context: { + sessionID: "ses_simulated_tools", + agent: "build", + messageID: "msg_simulated_tools", + callID: "call_success", + }, + }) + const successID = requireString(requireRecord(successInvocation.params).id) + const update = JSON.stringify({ + jsonrpc: "2.0", + id: 20, + method: "tool.update", + params: { + id: successID, + sequence: 0, + update: { + structured: { phase: "searching" }, + content: [{ type: "text", text: "Searching" }], + }, + }, + }) + socket.send(update) + expect(yield* Queue.take(messages)).toMatchObject({ id: 20, result: { ok: true } }) + socket.send(update) + expect(yield* Queue.take(messages)).toMatchObject({ id: 20, result: { ok: true } }) + socket.send( + JSON.stringify({ + ...JSON.parse(update), + id: 21, + }), + ) + expect(yield* Queue.take(messages)).toMatchObject({ id: 21, result: { ok: true } }) + socket.send( + JSON.stringify({ + jsonrpc: "2.0", + id: 20, + method: "tool.update", + params: { + id: successID, + sequence: 0, + update: { structured: { phase: "different" } }, + }, + }), + ) + expect(yield* Queue.take(messages)).toMatchObject({ + id: 20, + error: { message: expect.stringContaining("reused with different progress") }, + }) + socket.send( + JSON.stringify({ + jsonrpc: "2.0", + id: 22, + method: "tool.update", + params: { + id: successID, + sequence: 2, + update: { structured: { phase: "skipped" } }, + }, + }), + ) + expect(yield* Queue.take(messages)).toMatchObject({ + id: 22, + error: { message: expect.stringContaining("Expected simulated tool update sequence 1") }, + }) + socket.send( + JSON.stringify({ + jsonrpc: "2.0", + id: 3, + method: "tool.finish", + params: { + id: successID, + output: { + structured: { answer: 42 }, + content: [{ type: "text", text: "42" }], + }, + }, + }), + ) + expect(yield* Queue.take(messages)).toMatchObject({ id: 3, result: { ok: true } }) + socket.send( + JSON.stringify({ + jsonrpc: "2.0", + id: 23, + method: "tool.finish", + params: { + id: successID, + output: { + structured: { answer: 42 }, + content: [{ type: "text", text: "42" }], + }, + }, + }), + ) + expect(yield* Queue.take(messages)).toMatchObject({ id: 23, result: { ok: true } }) + expect(yield* Fiber.join(successful)).toMatchObject({ + result: { type: "text", value: "42" }, + output: { + structured: { answer: 42 }, + content: [{ type: "text", text: "42" }], + }, + }) + expect(progress).toEqual([ + { + structured: { phase: "searching" }, + content: [{ type: "text", text: "Searching" }], + }, + ]) + + const failed = yield* settle("call_failure", "missing").pipe(Effect.forkScoped) + const failedInvocation = yield* takeToolInvocation(messages) + const failedID = requireString(requireRecord(failedInvocation.params).id) + socket.send( + JSON.stringify({ + jsonrpc: "2.0", + id: 4, + method: "tool.fail", + params: { id: failedID, message: "lookup failed" }, + }), + ) + expect(yield* Queue.take(messages)).toMatchObject({ id: 4, result: { ok: true } }) + expect(yield* Fiber.join(failed)).toMatchObject({ + result: { type: "error", value: "lookup failed" }, + }) + + const concurrent = [ + yield* settle("call_first", "first").pipe(Effect.forkScoped), + yield* settle("call_second", "second").pipe(Effect.forkScoped), + ] + const invocations = [yield* takeToolInvocation(messages), yield* takeToolInvocation(messages)] + const byCall = new Map( + invocations.map((invocation) => { + const params = requireRecord(invocation.params) + const context = requireRecord(params.context) + return [requireString(context.callID), requireString(params.id)] + }), + ) + for (const [id, callID, value] of [ + [5, "call_second", "second result"], + [6, "call_first", "first result"], + ] as const) { + socket.send( + JSON.stringify({ + jsonrpc: "2.0", + id, + method: "tool.finish", + params: { + id: byCall.get(callID), + output: { structured: value, content: [{ type: "text", text: value }] }, + }, + }), + ) + expect(yield* Queue.take(messages)).toMatchObject({ id, result: { ok: true } }) + } + expect((yield* Fiber.join(concurrent[0])).result).toEqual({ type: "text", value: "first result" }) + expect((yield* Fiber.join(concurrent[1])).result).toEqual({ type: "text", value: "second result" }) + + const cancelled = yield* settle("call_cancelled", "slow").pipe(Effect.forkScoped) + const cancelledInvocation = yield* takeToolInvocation(messages) + const cancelledID = requireString(requireRecord(cancelledInvocation.params).id) + yield* Fiber.interrupt(cancelled) + expect(yield* Queue.take(messages)).toMatchObject({ + method: "tool.cancel", + params: { id: cancelledID, reason: "interrupted" }, + }) + socket.send( + JSON.stringify({ + jsonrpc: "2.0", + id: 7, + method: "tool.finish", + params: { + id: cancelledID, + output: { structured: null, content: [] }, + }, + }), + ) + expect(yield* Queue.take(messages)).toMatchObject({ + id: 7, + error: { message: expect.stringContaining("not found or already finished") }, + }) + + const replayed = yield* settle("call_replayed", "reconnect").pipe(Effect.forkScoped) + const original = yield* takeToolInvocation(messages) + const originalID = requireString(requireRecord(original.params).id) + const replayedProgress = { + structured: { phase: "before-reconnect" }, + content: [{ type: "text", text: "Still running" }], + } + socket.send( + JSON.stringify({ + jsonrpc: "2.0", + id: 25, + method: "tool.update", + params: { id: originalID, sequence: 0, update: replayedProgress }, + }), + ) + expect(yield* Queue.take(messages)).toMatchObject({ id: 25, result: { ok: true } }) + const replacement = yield* connect(endpoint) + const replacementMessages = yield* messagesFrom(replacement) + replacement.send( + JSON.stringify({ + jsonrpc: "2.0", + id: 8, + method: "tool.attach", + params: { tools: [registration] }, + }), + ) + expect(yield* Queue.take(replacementMessages)).toMatchObject({ + id: 8, + error: { message: expect.stringContaining("already attached") }, + }) + yield* closeSocket(socket) + const disconnected = yield* registry.materialize() + expect(disconnected.definitions).toContainEqual(expect.objectContaining({ name: "lookup" })) + replacement.send( + JSON.stringify({ + jsonrpc: "2.0", + id: 90, + method: "tool.attach", + params: { tools: [{ ...registration, name: "replacement" }] }, + }), + ) + expect(yield* Queue.take(replacementMessages)).toMatchObject({ + id: 90, + error: { message: expect.stringContaining("must settle pending invocations") }, + }) + replacement.send( + JSON.stringify({ + jsonrpc: "2.0", + id: 9, + method: "tool.attach", + params: { tools: [registration] }, + }), + ) + const attached = [ + requireRecord(yield* Queue.take(replacementMessages)), + requireRecord(yield* Queue.take(replacementMessages)), + ] + expect(attached).toContainEqual(expect.objectContaining({ id: 9, result: { attached: true } })) + expect(attached).toContainEqual( + expect.objectContaining({ + method: "tool.invocation", + params: expect.objectContaining({ id: originalID }), + }), + ) + replacement.send( + JSON.stringify({ + jsonrpc: "2.0", + id: 26, + method: "tool.update", + params: { id: originalID, sequence: 0, update: replayedProgress }, + }), + ) + expect(yield* Queue.take(replacementMessages)).toMatchObject({ id: 26, result: { ok: true } }) + expect(progress.filter((update) => update.structured.phase === "before-reconnect")).toHaveLength(1) + replacement.send( + JSON.stringify({ + jsonrpc: "2.0", + id: 10, + method: "tool.finish", + params: { + id: originalID, + output: { + structured: "replayed result", + content: [{ type: "text", text: "replayed result" }], + }, + }, + }), + ) + expect(yield* Queue.take(replacementMessages)).toMatchObject({ id: 10, result: { ok: true } }) + expect((yield* Fiber.join(replayed)).result).toEqual({ + type: "text", + value: "replayed result", + }) + + const preserved = yield* settle("call_preserved", "same generation").pipe(Effect.forkScoped) + const preservedInvocation = yield* takeToolInvocation(replacementMessages) + const preservedID = requireString(requireRecord(preservedInvocation.params).id) + replacement.send( + JSON.stringify({ + jsonrpc: "2.0", + id: 27, + method: "tool.finish", + params: { + id: preservedID, + output: { structured: "preserved", content: [{ type: "text", text: "preserved" }] }, + }, + }), + ) + expect(yield* Queue.take(replacementMessages)).toMatchObject({ id: 27, result: { ok: true } }) + expect((yield* Fiber.join(preserved)).result).toEqual({ type: "text", value: "preserved" }) + + const namespaced = [ + { ...registration, name: "search", options: { namespace: "github", codemode: false } }, + { ...registration, name: "search", options: { namespace: "web", codemode: false } }, + ] + replacement.send( + JSON.stringify({ + jsonrpc: "2.0", + id: 11, + method: "tool.attach", + params: { tools: namespaced }, + }), + ) + expect(yield* Queue.take(replacementMessages)).toMatchObject({ id: 11, result: { attached: true } }) + const replaced = yield* registry.materialize() + const replacedNames = replaced.definitions.map((definition) => definition.name) + expect(replacedNames).toEqual(expect.arrayContaining(["github_search", "web_search"])) + expect(replacedNames).not.toContain("lookup") + const secondaryReplaced = yield* ToolRegistry.Service.use((secondaryRegistry) => + secondaryRegistry.materialize(), + ).pipe(Effect.provide(secondary)) + const secondaryNames = secondaryReplaced.definitions.map((definition) => definition.name) + expect(secondaryNames).toEqual(expect.arrayContaining(["github_search", "web_search"])) + expect(secondaryNames).not.toContain("lookup") + const routed = yield* replaced + .settle({ + sessionID: SessionV2.ID.make("ses_simulated_tools"), + agent: AgentV2.ID.make("build"), + messageID: SessionMessage.ID.make("msg_simulated_tools"), + call: { + type: "tool-call", + id: "call_namespaced", + name: "github_search", + input: { query: "routing" }, + }, + }) + .pipe(Effect.forkScoped) + const routedInvocation = yield* takeToolInvocation(replacementMessages) + expect(routedInvocation.params).toMatchObject({ name: "github_search" }) + const routedID = requireString(requireRecord(routedInvocation.params).id) + replacement.send( + JSON.stringify({ + jsonrpc: "2.0", + id: 12, + method: "tool.finish", + params: { + id: routedID, + output: { structured: "routed", content: [{ type: "text", text: "routed" }] }, + }, + }), + ) + expect(yield* Queue.take(replacementMessages)).toMatchObject({ id: 12, result: { ok: true } }) + expect((yield* Fiber.join(routed)).result).toEqual({ type: "text", value: "routed" }) + expect( + yield* materialized.settle({ + sessionID: SessionV2.ID.make("ses_simulated_tools"), + agent: AgentV2.ID.make("build"), + messageID: SessionMessage.ID.make("msg_simulated_tools"), + call: { + type: "tool-call", + id: "call_stale", + name: "lookup", + input: { query: "stale" }, + }, + }), + ).toMatchObject({ + result: { type: "error", value: expect.stringContaining("no longer active") }, + }) + expect(activations).toBe(2) + }).pipe(Effect.provide(primary)) + }).pipe(Effect.provide(toolLifecycleLayer(endpoint)), Effect.scoped), + ) + } finally { + await rm(directory, { recursive: true, force: true }) + } +}, 15_000) + const request: SimulatedProvider.ProviderRequest = { url: "https://api.openai.com/v1/chat/completions", body: { model: "gpt-5", messages: [{ role: "user", content: "Hello" }] }, @@ -213,7 +680,26 @@ function runProvider( const socket = yield* connect(endpoint) const messages = yield* messagesFrom(socket) yield* body(provider, socket, messages) - }).pipe(Effect.provide(SimulatedProvider.layerDrive({ endpoint })), Effect.scoped), + }).pipe(Effect.provide(providerLayer(endpoint)), Effect.scoped), + ) +} + +const providerLayer = (endpoint: string) => + SimulatedProvider.layerDrive({ endpoint }).pipe( + Layer.provide( + Layer.succeed(SdkPlugins.Service, SdkPlugins.Service.of({ register: () => Effect.void, all: () => [] })), + ), + ) + +const toolLifecycleLayer = (endpoint: string) => { + const provider = makeGlobalNode({ + service: SimulatedProvider.Service, + layer: SimulatedProvider.layerDrive({ endpoint }), + deps: [SdkPlugins.node], + }) + return AppNodeBuilder.build( + LayerNode.group([Database.node, EventV2.node, SdkPlugins.node, LocationServiceMap.node, provider]), + [[Config.node, Layer.succeed(Config.Service, Config.Service.of({ entries: () => Effect.succeed([]) }))]], ) } @@ -227,6 +713,19 @@ function messagesFrom(socket: WebSocket) { }) } +function closeSocket(socket: WebSocket) { + return Effect.callback((resume) => { + if (socket.readyState === WebSocket.CLOSED) { + resume(Effect.void) + return Effect.void + } + const closed = () => resume(Effect.void) + socket.addEventListener("close", closed, { once: true }) + socket.close() + return Effect.sync(() => socket.removeEventListener("close", closed)) + }) +} + function attach(socket: WebSocket, messages: Queue.Queue) { return Effect.gen(function* () { socket.send(JSON.stringify({ jsonrpc: "2.0", id: 1, method: "llm.attach" })) @@ -244,6 +743,21 @@ function takeInvocation(messages: Queue.Queue) { ) } +function takeToolInvocation(messages: Queue.Queue) { + return Queue.take(messages).pipe( + Effect.map((message) => { + const opened = requireRecord(message) + if (opened.method !== "tool.invocation") throw new Error("Expected a tool.invocation notification") + return opened + }), + ) +} + +function requireString(value: unknown) { + if (typeof value !== "string") throw new Error("Expected a string") + return value +} + function requireRecord(value: unknown): Record { if (!isRecord(value)) throw new Error("Expected an object") return value