From e1c19f2e48209774bb439762cb46d613a8d20c95 Mon Sep 17 00:00:00 2001 From: Kit Langton Date: Sun, 19 Jul 2026 21:02:42 -0400 Subject: [PATCH 1/3] feat(drive): control arbitrary tool lifecycles --- .changeset/provider-tool-lifecycle.md | 5 + packages/drive/README.md | 60 ++ packages/drive/compile/public-api.ts | 9 + packages/drive/docs/open-code-driver-api.md | 52 +- .../docs/open-code-driver-architecture.md | 44 +- packages/drive/src/cli/script.ts | 2 +- packages/drive/src/driver/llm-responder.ts | 62 +- packages/drive/src/driver/prepared.ts | 13 +- packages/drive/src/driver/server.ts | 65 +- packages/drive/src/instance/runtime.ts | 4 +- packages/drive/src/simulation/connector.ts | 135 +++- packages/drive/src/simulation/protocol.ts | 158 ++++ packages/drive/src/simulation/rpc.ts | 16 + packages/drive/src/tool/controller.ts | 10 +- packages/drive/src/tool/producer.ts | 681 ++++++++++++++++++ packages/drive/src/tool/types.ts | 49 +- packages/drive/test/driver/index.test.ts | 81 +++ .../drive/test/driver/llm-controller.test.ts | 130 +++- packages/drive/test/fixtures/fake-opencode.ts | 42 +- .../drive/test/simulation/backend.test.ts | 100 ++- .../drive/test/simulation/connector.test.ts | 36 + .../drive/test/simulation/protocol.test.ts | 98 +++ .../drive/test/simulation/transport-peer.ts | 12 +- packages/drive/test/tool/producer.test.ts | 470 ++++++++++++ skills/opencode-drive/SKILL.md | 36 + 25 files changed, 2277 insertions(+), 93 deletions(-) create mode 100644 .changeset/provider-tool-lifecycle.md create mode 100644 packages/drive/src/tool/producer.ts create mode 100644 packages/drive/test/simulation/protocol.test.ts create mode 100644 packages/drive/test/tool/producer.test.ts diff --git a/.changeset/provider-tool-lifecycle.md b/.changeset/provider-tool-lifecycle.md new file mode 100644 index 0000000..8160cd7 --- /dev/null +++ b/.changeset/provider-tool-lifecycle.md @@ -0,0 +1,5 @@ +--- +"opencode-drive": minor +--- + +Control arbitrary provider-backed tool lifecycles with dynamic registration, structured progress, success, failure, cancellation, and reconnect-safe replay. diff --git a/packages/drive/README.md b/packages/drive/README.md index 01f6047..f438097 100644 --- a/packages/drive/README.md +++ b/packages/drive/README.md @@ -286,6 +286,66 @@ Declared `config` and `tuiConfig` values are deeply merged over fixture `.opencode/opencode.jsonc` and `.opencode/tui.jsonc` files. Arrays replace instead of merging, and mutations made in `setup` take final precedence. +Attach arbitrary provider-backed tools at runtime with their JSON schemas, then +take and settle native OpenCode invocations by model call ID. `attach` replaces +the complete dynamic set atomically; it does not affect the built-in adapters +configured through the driver or script `tools` option. + +```ts +import { Effect } from "effect" +import { Llm, OpenCodeDriver } from "opencode-drive" + +export default OpenCodeDriver.use(({ tools, llm, ui }) => + Effect.gen(function* () { + yield* tools.attach({ + tools: [ + { + name: "lookup", + description: "Look up a value", + inputSchema: { + type: "object", + properties: { query: { type: "string" } }, + required: ["query"], + }, + outputSchema: { + type: "object", + properties: { answer: { type: "number" } }, + required: ["answer"], + }, + options: { codemode: false }, + }, + ], + }) + yield* llm.queue( + Llm.toolCall({ + index: 0, + id: "call_lookup", + name: "lookup", + input: { query: "meaning" }, + }), + Llm.finish("tool-calls"), + ) + yield* ui.submit("Look up the meaning") + + const lookup = yield* tools.take("call_lookup") + yield* lookup.progress({ + structured: { phase: "searching" }, + content: [{ type: "text", text: "Searching" }], + }) + yield* lookup.finish({ + structured: { answer: 42 }, + content: [{ type: "text", text: "42" }], + }) + }), +) +``` + +Drive owns progress sequence numbers and retries uncertain operations without +rerunning a claimed call. `awaitCancelled()` completes when OpenCode interrupts +the native invocation before `finish` or `fail`. Dynamic registrations survive +the tool-only controller reconnecting; an intentional server generation change +cancels unresolved calls and reapplies the desired set after launch. + Declare which built-in tools Drive should intercept with `tools`, then control their invocations inside `run`. Each tool controller accepts calls in arrival order or by the stable call ID chosen in `Llm.toolCall`: diff --git a/packages/drive/compile/public-api.ts b/packages/drive/compile/public-api.ts index 45ffe4e..c79d3eb 100644 --- a/packages/drive/compile/public-api.ts +++ b/packages/drive/compile/public-api.ts @@ -55,6 +55,9 @@ export type ZeroConfigUseIsRunnable = Assert< const controlledOptions: OpenCodeDriver.Options = { tools: ["shell"] } declare const controls: Tool.Controls const shellCalls = controls.control("shell") +declare const dynamicTools: Tool.AttachParams +const attached = controls.attach(dynamicTools) +const dynamicCall = controls.take("call_lookup") declare const toolName: Tool.Name controls.control(toolName) export type ShellControlIsTyped = Assert< @@ -63,6 +66,12 @@ export type ShellControlIsTyped = Assert< Tool.ControlledCalls > > +export type DynamicAttachIsTyped = Assert< + Equal, void> +> +export type DynamicCallIsTyped = Assert< + Equal, Tool.Invocation> +> const controlled = OpenCodeDriver.use(controlledOptions, ({ tools }) => tools.control("shell").pipe(Effect.asVoid), ) diff --git a/packages/drive/docs/open-code-driver-api.md b/packages/drive/docs/open-code-driver-api.md index 3715dc0..b62ce86 100644 --- a/packages/drive/docs/open-code-driver-api.md +++ b/packages/drive/docs/open-code-driver-api.md @@ -243,6 +243,53 @@ unresolved calls, and waits for transport cleanup. The callback-style for fixed handlers; callback-controlled tools are not also available through the runtime `tools.control` capability. +## Arbitrary tools use the provider-backed lifecycle + +`tools.attach({ tools })` atomically replaces the complete dynamic registration +set for the current run. Registrations use OpenCode's canonical JSON Schema, +permission, namespace, and CodeMode options. Static `shell`, `webfetch`, and +`websearch` adapters remain installed separately. + +```ts +yield* tools.attach({ + tools: [ + { + name: "lookup", + description: "Look up a value", + inputSchema: { + type: "object", + properties: { query: { type: "string" } }, + required: ["query"], + }, + options: { codemode: false }, + }, + ], +}) + +const invocation = yield* tools.take("call_lookup") +yield* invocation.progress({ + structured: { phase: "searching" }, + content: [{ type: "text", text: "Searching" }], +}) +yield* invocation.finish({ + structured: { answer: 42 }, + content: [{ type: "text", text: "42" }], +}) +``` + +`take(callID)` matches `context.callID`, the model call ID supplied to +`Llm.toolCall`; `invocation.id` is the producer's transport identity. Drive +deduplicates invocation replay after a controller reconnect and retries +progress or terminal operations with the same producer identity and progress +sequence. `awaitCancelled()` observes OpenCode's native interruption. There is +no public cancel operation because cancellation flows from OpenCode to Drive. + +Attaching a dynamic effective name that collides with a configured static +adapter fails locally. Calling `attach({ tools: [] })` clears the dynamic set. +Older OpenCode revisions remain compatible with static adapters and ordinary +LLM control; dynamic attachment fails with `Tool.LifecycleError` when the six +tool lifecycle capabilities are unavailable. + ## LLM response description is separate from live LLM control `Llm` is a pure data module. `llm` is the live capability that queues, sends, and serves responses. @@ -262,8 +309,9 @@ yield* llm.queue( Each constructor returns an ordinary serializable value. Raw values with the same schema remain accepted. Tool calls remain atomic when options are omitted. Supplying stream options -serializes the input to JSON and emits it incrementally through the simulated -provider, producing OpenCode's normal tool-input lifecycle: +serializes the input to JSON and emits provider-neutral partial tool input when +the endpoint advertises that capability. Older endpoints retain the existing +OpenAI-compatible fallback: ```ts Llm.toolCall( diff --git a/packages/drive/docs/open-code-driver-architecture.md b/packages/drive/docs/open-code-driver-architecture.md index f351158..cd0a661 100644 --- a/packages/drive/docs/open-code-driver-architecture.md +++ b/packages/drive/docs/open-code-driver-architecture.md @@ -15,7 +15,7 @@ OpenCodeDriver tuis additional frontend process factory ui convenience alias for tui.ui llm shared simulated-model control - tools runtime control for declared tool adapters + tools runtime control for static adapters and arbitrary tools ``` The names distinguish the two kinds of client involved: @@ -25,8 +25,10 @@ The names distinguish the two kinds of client involved: optional `recording`. - `Tuis` launches and supervises additional frontend processes connected to the same server. -- `tools` accepts independently controlled invocations for adapters declared - before OpenCode starts. +- `tools.control` accepts independently controlled invocations for adapters + declared before OpenCode starts. +- `tools.attach` and `tools.take` control arbitrary native tools through the + canonical provider-backed lifecycle. - Transport-level JSON-RPC clients remain private implementation details. `defineScript` consumes these exact capabilities. It adds a branded module @@ -48,7 +50,9 @@ Effect Scope controlled invocation exchanges OpenCodeServer backend simulation connection + reconnecting tool-only backend connection LLM controller + dynamic ToolProducer generated OpenCode SDK connection TUI supervisor primary TUI scope @@ -59,9 +63,10 @@ Effect Scope Library drivers create their ToolController before project preparation and pass that controller into `OpenCodeInstance`. CLI scripts create the controller -inside `OpenCodeInstance`. Prepared drivers and script contexts derive controls -from their instance, so the controller that wrote the plugin configuration is -structurally the controller exposed at runtime. +inside `OpenCodeInstance`. Prepared drivers and script contexts combine the +instance's static controller with the server's dynamic producer. The static +controller that wrote plugin configuration remains the one exposed through +`tools.control`. `OpenCodeDriver.make(options)` requires `Scope.Scope`. It returns once the server, generated SDK client, primary TUI, and simulation connections are @@ -73,11 +78,12 @@ settlement even when the user program fails. Settlement is one shared terminal operation. It runs in this order: 1. Validate that queued LLM work was consumed. -2. Shut down the LLM controller. -3. Finish active recording timelines. -4. Close all TUI scopes and processes. -5. Export completed recordings. -6. Decode the schema-validated `RunReport`. +2. Validate that native dynamic-tool invocations were settled. +3. Shut down the LLM controller. +4. Finish active recording timelines. +5. Close all TUI scopes and processes. +6. Export completed recordings. +7. Decode the schema-validated `RunReport`. `driver.settle()` is shared and idempotent. Once settlement starts, `tuis` rejects new launches and `llm` rejects new responses. `OpenCodeDriver.use` @@ -98,6 +104,15 @@ Terminal commitment uses a synchronous first-writer-wins Deferred completion; Drive guarantees exactly-once acceptance inside the controller, not delivery across a transport disconnect. +`ToolProducer` owns a separate backend socket because LLM chunks are not +idempotent and an LLM socket closure is terminal to `LlmController`. Dynamic +tool progress and terminal RPCs are idempotent by producer invocation ID and +sequence, so the tool-only connection may reconnect and replay pending +invocations safely. One ordered event stream preserves invocation-before- +cancellation order. The desired registration set survives reconnects and +manual server relaunches; invocation records are scoped to one server +generation because producer IDs may be reused by a new process. + ## TUI Lifecycle `Tuis.launch(options)` generates an internal identity. `Tuis.launch(name, @@ -153,9 +168,10 @@ schema validation, request correlation, interruption, and connection failure. The driver receives the connector through an Effect service and does not expose it in userland. -The UI connection is request-response JSON-RPC. The backend additionally -receives unsolicited `llm.request` notifications and exposes them as a -validated stream to the LLM controller. +The UI connection is request-response JSON-RPC. The LLM backend additionally +receives unsolicited `llm.request` notifications. The tool-only backend keeps +ordered `tool.invocation` and `tool.cancel` notifications on one validated +stream and does not call `llm.attach`. ## Project Setup diff --git a/packages/drive/src/cli/script.ts b/packages/drive/src/cli/script.ts index 974da50..a2de6e3 100644 --- a/packages/drive/src/cli/script.ts +++ b/packages/drive/src/cli/script.ts @@ -130,7 +130,7 @@ export const runScript = Effect.fn("DriveCli.runScript")(function* ( kill: prepared.server.kill, }, llm: prepared.llm, - tools: instance.tools, + tools: prepared.tools, artifacts: instance.artifacts, } const primaryTui = prepared.primary diff --git a/packages/drive/src/driver/llm-responder.ts b/packages/drive/src/driver/llm-responder.ts index 59826a3..6f6beab 100644 --- a/packages/drive/src/driver/llm-responder.ts +++ b/packages/drive/src/driver/llm-responder.ts @@ -116,20 +116,11 @@ export const make = ({ requestTimeout }: Options): Responder => { const delay = toolCall.options?.delay ?? 2 const chunkSize = toolCall.options?.chunkSize ?? 15 const chunks = [...chunkText(JSON.stringify(toolCall.input), chunkSize)] - for (let index = 0; index < chunks.length; index++) { - const text = chunks[index] - if (text === undefined) continue - const callDelta = - index === 0 - ? { - index: toolCall.index, - id: toolCall.id, - function: { name: toolCall.name, arguments: text }, - } - : { - index: toolCall.index, - function: { arguments: text }, - } + const providerNeutral = supportsCapability( + backend.compatibility, + "llm.tool-input-delta", + ) + if (providerNeutral) yield* call( backend, "llm.chunk", @@ -138,12 +129,51 @@ export const make = ({ requestTimeout }: Options): Responder => { id: requestId, items: [ { - type: "raw", - chunk: { choices: [{ delta: { tool_calls: [callDelta] } }] }, + type: "toolInputStart", + index: toolCall.index, + id: toolCall.id, + name: toolCall.name, }, ], }), ) + for (let index = 0; index < chunks.length; index++) { + const text = chunks[index] + if (text === undefined) continue + const item = providerNeutral + ? { type: "toolInputDelta" as const, index: toolCall.index, text } + : { + type: "raw" as const, + chunk: { + choices: [ + { + delta: { + tool_calls: [ + index === 0 + ? { + index: toolCall.index, + id: toolCall.id, + function: { name: toolCall.name, arguments: text }, + } + : { + index: toolCall.index, + function: { arguments: text }, + }, + ], + }, + }, + ], + }, + } + yield* call( + backend, + "llm.chunk", + requestId, + backend.rpc["llm.chunk"]({ + id: requestId, + items: [item], + }), + ) if (index < chunks.length - 1 && delay > 0) yield* Effect.sleep(delay) } }) diff --git a/packages/drive/src/driver/prepared.ts b/packages/drive/src/driver/prepared.ts index 1b3ecf6..46fe172 100644 --- a/packages/drive/src/driver/prepared.ts +++ b/packages/drive/src/driver/prepared.ts @@ -31,6 +31,7 @@ export interface Prepared { readonly driver: Driver | undefined readonly primary: OpenCodeTui.Tui | undefined readonly llm: Llm + readonly tools: Driver["tools"] readonly tuis: OpenCodeTui.Tuis readonly server: Pick readonly artifacts: string @@ -67,6 +68,11 @@ export const makeWithServices = Effect.fn("OpenCodeDriver.makePreparedWithServic ) => Effect.gen(function* () { const llm = yield* Effect.exit(server.llm.settle()) + const tools = yield* Effect.exit( + server.settleTools.pipe( + Effect.mapError((cause) => error("tools.settle", cause)), + ), + ) const shutdown = yield* Effect.exit(server.llm.shutdown()) const tuiExit = yield* Effect.exit(tuis) let failure: Cause.Cause< @@ -76,6 +82,10 @@ export const makeWithServices = Effect.fn("OpenCodeDriver.makePreparedWithServic | OpenCodeUi.OperationError > | undefined if (Exit.isFailure(llm)) failure = llm.cause + if (Exit.isFailure(tools)) + failure = failure === undefined + ? tools.cause + : Cause.combine(failure, tools.cause) if (Exit.isFailure(shutdown)) failure = failure === undefined ? shutdown.cause @@ -110,7 +120,7 @@ export const makeWithServices = Effect.fn("OpenCodeDriver.makePreparedWithServic tui: primary, ui: primary.ui, llm, - tools: instance.tools, + tools: server.tools, tuis: server.tuis, artifacts: instance.artifacts, settle: () => settle, @@ -119,6 +129,7 @@ export const makeWithServices = Effect.fn("OpenCodeDriver.makePreparedWithServic driver, primary, llm, + tools: server.tools, tuis: server.tuis, server, artifacts: instance.artifacts, diff --git a/packages/drive/src/driver/server.ts b/packages/drive/src/driver/server.ts index 52179ef..edcec81 100644 --- a/packages/drive/src/driver/server.ts +++ b/packages/drive/src/driver/server.ts @@ -10,6 +10,8 @@ import * as OpenCodeTui from "./client.js" import * as OpenCodeSdk from "./opencode.js" import { error, type OpenCodeDriverError } from "./error.js" import * as LlmController from "./llm-controller.js" +import * as ToolProducer from "../tool/producer.js" +import type * as Tool from "../tool/index.js" export interface Target { readonly command?: ReadonlyArray @@ -26,6 +28,8 @@ export interface Options { export interface Server { readonly llm: LlmController.Controller + readonly tools: Tool.Controls + readonly settleTools: ToolProducer.Controller["settle"] readonly tuis: OpenCodeTui.Control readonly launch: () => Effect.Effect< OpenCodeSdk.OpenCode, @@ -47,6 +51,11 @@ export const make = Effect.fn("OpenCodeServer.make")(function* ( const target = options.target ?? {} const instance = options.instance const llm = yield* LlmController.make() + const toolProducer = yield* ToolProducer.make(instance.toolNames) + const tools: Tool.Controls = { + ...instance.tools, + ...toolProducer.controls, + } const tuis = yield* OpenCodeTui.makeTuis( instance, target.visible ?? false, @@ -89,6 +98,45 @@ export const make = Effect.fn("OpenCodeServer.make")(function* ( ), ), ) + const connectTools = Effect.fn("OpenCodeServer.connectTools")(function* () { + const connectionScope = yield* Scope.fork(scope) + const backend = yield* connector.backend(launched.endpoint, { + attach: false, + compatibility: target.compatibility, + }).pipe( + Scope.provide(connectionScope), + Effect.mapError((cause) => error("tools.connect", cause)), + Effect.onError(() => Scope.close(connectionScope, Exit.void)), + ) + const attachment = yield* toolProducer.connect(backend).pipe( + Effect.mapError((cause) => error("tools.connect", cause)), + Effect.onError(() => Scope.close(connectionScope, Exit.void)), + ) + return { attachment, backend, scope: connectionScope } + }) + function reconnectTools(): Effect.Effect { + return connectTools().pipe( + Effect.flatMap(superviseTools), + Effect.catch(() => Effect.sleep(25).pipe(Effect.andThen(reconnectTools()))), + ) + } + function superviseTools( + connection: Effect.Success>, + ): Effect.Effect { + return connection.backend.closed.pipe( + Effect.ensuring( + connection.attachment.detach().pipe( + Effect.andThen(Scope.close(connection.scope, Exit.void)), + ), + ), + Effect.andThen(Effect.sleep(25)), + Effect.andThen(reconnectTools()), + ) + } + const toolConnection = yield* connectTools() + yield* superviseTools(toolConnection).pipe( + Effect.forkIn(scope), + ) const attachment = yield* llm.attach(backend).pipe( Effect.onError(() => instance.killServer.pipe( @@ -107,6 +155,12 @@ export const make = Effect.fn("OpenCodeServer.make")(function* ( ) const process = yield* instance.primary.pipe( Effect.mapError((cause) => error("server.launch", cause)), + Effect.onError(() => + instance.killServer.pipe( + Effect.ignore, + Effect.andThen(Scope.close(scope, Exit.void)), + ), + ), ) yield* Ref.set(generation, { scope, attachment, process }) compatibility = [...compatibility, backend.compatibility] @@ -138,12 +192,13 @@ export const make = Effect.fn("OpenCodeServer.make")(function* ( ) yield* Ref.set(generation, undefined) yield* active.attachment.detach() + yield* Scope.close(active.scope, Exit.void) + yield* toolProducer.endGeneration const stopped = yield* Effect.exit( instance.killServer.pipe( Effect.mapError((cause) => error("server.kill", cause)), ), ) - yield* Scope.close(active.scope, Exit.void) if (Exit.isFailure(stopped)) return yield* Effect.failCause(stopped.cause) return undefined }) @@ -152,12 +207,16 @@ export const make = Effect.fn("OpenCodeServer.make")(function* ( return { llm, + tools, + settleTools: toolProducer.settle, tuis, launch, kill, failure: Effect.raceFirst( - llm.failure, - Deferred.await(unexpectedExit), + toolProducer.failure.pipe( + Effect.mapError((cause) => error("tools", cause)), + ), + Effect.raceFirst(llm.failure, Deferred.await(unexpectedExit)), ), compatibility: Effect.sync(() => compatibility), } satisfies Server diff --git a/packages/drive/src/instance/runtime.ts b/packages/drive/src/instance/runtime.ts index 4951531..dc66a45 100644 --- a/packages/drive/src/instance/runtime.ts +++ b/packages/drive/src/instance/runtime.ts @@ -65,7 +65,8 @@ export interface Instance { readonly ui: string readonly backend: string } - readonly tools: Tool.Controls + readonly tools: Tool.StaticControls + readonly toolNames: ReadonlySet readonly recording: Effect.Effect readonly primary: Effect.Effect readonly launchServer: Effect.Effect< @@ -564,6 +565,7 @@ export const make = Effect.fn("OpenCodeInstance.make")(function* ( visible: options.visible ?? false, endpoints, tools: toolController.controls, + toolNames: toolController.names, recording: Ref.get(state).pipe(Effect.map((current) => current.recording)), primary: Ref.get(state).pipe( Effect.flatMap((current) => diff --git a/packages/drive/src/simulation/connector.ts b/packages/drive/src/simulation/connector.ts index ceb9037..2c92da5 100644 --- a/packages/drive/src/simulation/connector.ts +++ b/packages/drive/src/simulation/connector.ts @@ -84,6 +84,10 @@ export interface UiConnection { readonly compatibility: EndpointCompatibility } +export type ToolEvent = + | { readonly type: "invocation"; readonly invocation: BackendProtocol.ToolInvocation } + | { readonly type: "cancellation"; readonly cancellation: BackendProtocol.ToolCancellation } + export interface BackendConnection { readonly endpoint: string readonly rpc: BackendClient @@ -92,6 +96,7 @@ export interface BackendConnection { BackendProtocol.OpenedExchange, Schema.SchemaError > + readonly toolEvents: Stream.Stream readonly closed: Effect.Effect readonly attach: () => Effect.Effect< { readonly attached: true }, @@ -99,8 +104,51 @@ export interface BackendConnection { | RpcClientError.RpcClientError | SimulationRequestError > + readonly attachTools: ( + tools: ReadonlyArray, + ) => Effect.Effect< + { readonly attached: true }, + | SimulationCompatibilityError + | SimulationConnectionError + | RpcClientError.RpcClientError + | SimulationRequestError + > + readonly updateTool: ( + params: BackendProtocol.ToolUpdateParams, + ) => Effect.Effect< + { readonly ok: true }, + | SimulationConnectionError + | RpcClientError.RpcClientError + | SimulationRequestError + > + readonly finishTool: ( + params: BackendProtocol.ToolFinishParams, + ) => Effect.Effect< + { readonly ok: true }, + | SimulationConnectionError + | RpcClientError.RpcClientError + | SimulationRequestError + > + readonly failTool: ( + params: BackendProtocol.ToolFailParams, + ) => Effect.Effect< + { readonly ok: true }, + | SimulationConnectionError + | RpcClientError.RpcClientError + | SimulationRequestError + > } +const toolCapabilities = [ + "tool.attach", + "tool.update", + "tool.finish", + "tool.fail", + "tool.invocation", + "tool.cancel", +] as const satisfies ReadonlyArray +const toolCapabilitySet: ReadonlySet = new Set(toolCapabilities) + export const ui = Effect.fn("SimulationConnector.ui")(function* ( endpoint: string, options?: Options, @@ -140,10 +188,12 @@ export const backend = Effect.fn("SimulationConnector.backend")(function* ( BackendProtocol.OpenedExchange, Schema.SchemaError >() + const toolEvents = yield* Queue.unbounded() const closed = yield* Deferred.make() const close = Effect.all([ Deferred.succeed(closed, undefined), Queue.shutdown(requests), + Queue.shutdown(toolEvents), ], { discard: true }) yield* Effect.addFinalizer(() => close) const protocol = yield* OpenCodeRpcProtocol.make(endpoint, { @@ -151,22 +201,46 @@ export const backend = Effect.fn("SimulationConnector.backend")(function* ( firstWireId: 0, onClose: () => close, onNotification: ({ method, params }) => { - if (method !== "llm.request") return Effect.void - return Effect.matchEffect( - Schema.decodeUnknownEffect(BackendProtocol.OpenedExchange)(params), - { - onFailure: (error) => Queue.fail(requests, error).pipe(Effect.asVoid), - onSuccess: (request) => - Queue.offer(requests, request).pipe(Effect.asVoid), - }, - ) + switch (method) { + case "llm.request": + return Effect.matchEffect( + Schema.decodeUnknownEffect(BackendProtocol.OpenedExchange)(params), + { + onFailure: (error) => Queue.fail(requests, error).pipe(Effect.asVoid), + onSuccess: (request) => Queue.offer(requests, request).pipe(Effect.asVoid), + }, + ) + case "tool.invocation": + return Effect.matchEffect( + Schema.decodeUnknownEffect(BackendProtocol.ToolInvocation)(params), + { + onFailure: (error) => Queue.fail(toolEvents, error).pipe(Effect.asVoid), + onSuccess: (invocation) => + Queue.offer(toolEvents, { type: "invocation", invocation }).pipe(Effect.asVoid), + }, + ) + case "tool.cancel": + return Effect.matchEffect( + Schema.decodeUnknownEffect(BackendProtocol.ToolCancellation)(params), + { + onFailure: (error) => Queue.fail(toolEvents, error).pipe(Effect.asVoid), + onSuccess: (cancellation) => + Queue.offer(toolEvents, { type: "cancellation", cancellation }).pipe(Effect.asVoid), + }, + ) + default: + return Effect.void + } }, }) const rpc = yield* RpcClient.make(BackendRpcs).pipe( Effect.provideService(RpcClient.Protocol, protocol), ) const required = BackendProtocol.Capabilities.filter( - (capability) => capability !== "llm.pending", + (capability) => + capability !== "llm.pending" && + capability !== "llm.tool-input-delta" && + !toolCapabilitySet.has(capability), ) const compatibility = yield* negotiate( endpoint, @@ -177,7 +251,7 @@ export const backend = Effect.fn("SimulationConnector.backend")(function* ( expectedRole: "backend", offeredVersions: [1], requiredCapabilities: required, - optionalCapabilities: ["llm.pending"], + optionalCapabilities: ["llm.pending", "llm.tool-input-delta", ...toolCapabilities], }), options?.compatibility, ) @@ -197,13 +271,52 @@ export const backend = Effect.fn("SimulationConnector.backend")(function* ( ), ) if (options?.attach !== false) yield* attach() + const withTimeout = (operation: string, effect: Effect.Effect) => + effect.pipe( + Effect.timeoutOrElse({ + duration: options?.requestTimeout ?? 30_000, + orElse: () => + Effect.fail( + new SimulationConnectionError({ + endpoint, + operation, + message: `${operation} timed out after ${options?.requestTimeout ?? 30_000}ms`, + }), + ), + }), + ) + const attachTools: BackendConnection["attachTools"] = (tools) => { + const missing = toolCapabilities.filter( + (capability) => !supportsCapability(compatibility, capability), + ) + if (missing.length > 0) + return Effect.fail( + new SimulationCompatibilityError({ + endpoint, + role: "backend", + message: `Simulation endpoint is missing required capabilities: ${missing.join(", ")}`, + }), + ) + return withTimeout("tool.attach", rpc["tool.attach"]({ tools })) + } + const updateTool: BackendConnection["updateTool"] = (params) => + withTimeout("tool.update", rpc["tool.update"](params)) + const finishTool: BackendConnection["finishTool"] = (params) => + withTimeout("tool.finish", rpc["tool.finish"](params)) + const failTool: BackendConnection["failTool"] = (params) => + withTimeout("tool.fail", rpc["tool.fail"](params)) return { endpoint, rpc, compatibility, requests: Stream.fromQueue(requests), + toolEvents: Stream.fromQueue(toolEvents), closed: Deferred.await(closed), attach, + attachTools, + updateTool, + finishTool, + failTool, } satisfies BackendConnection }) diff --git a/packages/drive/src/simulation/protocol.ts b/packages/drive/src/simulation/protocol.ts index 66b8b70..270fbdc 100644 --- a/packages/drive/src/simulation/protocol.ts +++ b/packages/drive/src/simulation/protocol.ts @@ -380,6 +380,13 @@ 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([ @@ -388,6 +395,17 @@ export namespace Backend { 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, @@ -402,6 +420,126 @@ export namespace Backend { export const FinishReason = Llm.FinishReason 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 Attached = Schema.Struct({ attached: Schema.Literal(true) }) export interface Attached extends Schema.Schema.Type {} @@ -453,6 +591,26 @@ export namespace Backend { 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/drive/src/simulation/rpc.ts b/packages/drive/src/simulation/rpc.ts index f544311..0215a84 100644 --- a/packages/drive/src/simulation/rpc.ts +++ b/packages/drive/src/simulation/rpc.ts @@ -93,4 +93,20 @@ export const BackendRpcs = RpcGroup.make( payload: Backend.DisconnectParams, success: Backend.Ok, }), + request("tool.attach", { + payload: Backend.ToolAttachParams, + success: Backend.Attached, + }), + request("tool.update", { + payload: Backend.ToolUpdateParams, + success: Backend.Ok, + }), + request("tool.finish", { + payload: Backend.ToolFinishParams, + success: Backend.Ok, + }), + request("tool.fail", { + payload: Backend.ToolFailParams, + success: Backend.Ok, + }), ) diff --git a/packages/drive/src/tool/controller.ts b/packages/drive/src/tool/controller.ts index 3546bdc..246547a 100644 --- a/packages/drive/src/tool/controller.ts +++ b/packages/drive/src/tool/controller.ts @@ -24,7 +24,7 @@ import { type ControlledCall, type ControlledCalls, type ControlledCallsFor, - type Controls, + type StaticControls, type Handler, type Name, type Registration, @@ -92,8 +92,9 @@ function controlError( export interface Controller { readonly enabled: boolean - readonly controls: Controls + readonly controls: StaticControls readonly configure: (config: OpenCodeConfig) => void + readonly names: ReadonlySet } export function composeSetup( @@ -231,9 +232,9 @@ export const make = Effect.fn("ToolController.make")(function* (configuration?: ? Effect.fail(controlError("control", "not-controlled", name)) : Effect.succeed(calls) } - const controls: Controls = { control } + const controls: StaticControls = { control } if (definitions.size === 0) - return { enabled: false, controls, configure() {} } satisfies Controller + return { enabled: false, controls, configure() {}, names: new Set() } satisfies Controller const token = crypto.randomUUID() const indexes = new Map() @@ -299,6 +300,7 @@ export const make = Effect.fn("ToolController.make")(function* (configuration?: return { enabled: true, controls, + names: new Set(definitions.keys()), configure(config) { const current = config.plugins if (current !== undefined && !Array.isArray(current)) diff --git a/packages/drive/src/tool/producer.ts b/packages/drive/src/tool/producer.ts new file mode 100644 index 0000000..e23640f --- /dev/null +++ b/packages/drive/src/tool/producer.ts @@ -0,0 +1,681 @@ +import { createHash } from "node:crypto" +import * as Cause from "effect/Cause" +import * as Deferred from "effect/Deferred" +import * as Effect from "effect/Effect" +import * as Exit from "effect/Exit" +import * as Queue from "effect/Queue" +import * as Ref from "effect/Ref" +import * as Schema from "effect/Schema" +import * as Semaphore from "effect/Semaphore" +import * as Scope from "effect/Scope" +import * as Stream from "effect/Stream" +import { RpcClientError } from "effect/unstable/rpc" +import { + SimulationCompatibilityError, + SimulationConnectionError, + type BackendConnection, +} from "../simulation/connector.js" +import { Backend } from "../simulation/protocol.js" +import { SimulationRequestError } from "../simulation/rpc.js" +import { + LifecycleError, + type AttachParams, + type Cancellation, + type DynamicControls, + type Invocation, + type Output, + type Progress, +} from "./types.js" + +export interface BackendAttachment { + readonly detach: () => Effect.Effect +} + +export interface Controller { + readonly controls: DynamicControls + readonly connect: ( + backend: BackendConnection, + ) => Effect.Effect + readonly endGeneration: Effect.Effect + readonly settle: Effect.Effect + readonly shutdown: Effect.Effect + readonly failure: Effect.Effect +} + +type Waiter = { + readonly callID?: string + readonly resume: (effect: Effect.Effect) => void + delivered?: Active +} + +type Active = { + readonly fingerprint: string + readonly call: Invocation + readonly cancelled: Deferred.Deferred + readonly operations: Semaphore.Semaphore + claimed: boolean + sequence: number + state: "pending" | "settled" | "cancelled" +} + +const decodeAttach = Schema.decodeUnknownEffect(Backend.ToolAttachParams) +const decodeProgress = Schema.decodeUnknownEffect(Backend.ToolProgress) +const decodeOutput = Schema.decodeUnknownEffect(Backend.ToolOutput) +const decodeFailure = Schema.decodeUnknownEffect(Schema.String) + +export const make = Effect.fn("ToolProducer.make")(function* ( + staticNames: ReadonlySet, +) { + const parentScope = yield* Scope.Scope + const lifecycle = yield* Semaphore.make(1) + const attachmentLock = yield* Semaphore.make(1) + const records = new Map() + const waiters: Waiter[] = [] + const completed = new Map() + const backendChanges = yield* Queue.sliding(1) + const failure = yield* Deferred.make() + const current = yield* Ref.make< + | { readonly backend: BackendConnection; readonly scope: Scope.Scope } + | undefined + >(undefined) + let desired: AttachParams | undefined + let closed = false + + const lifecycleError = ( + operation: LifecycleError["operation"], + reason: LifecycleError["reason"], + message: string, + callID?: string, + ) => + new LifecycleError({ + operation, + reason, + message, + ...(callID === undefined ? {} : { callID }), + }) + + const notifyBackend = Queue.offer(backendChanges, undefined).pipe( + Effect.asVoid, + ) + + const remember = (id: string, fingerprint: string) => { + completed.set(id, fingerprint) + if (completed.size > 256) { + const oldest = completed.keys().next().value + if (oldest !== undefined) completed.delete(oldest) + } + } + + const deliver = (record: Active) => { + const callID = record.call.context.callID + const exact = waiters.findIndex((waiter) => waiter.callID === callID) + const index = + exact >= 0 + ? exact + : waiters.findIndex((waiter) => waiter.callID === undefined) + if (index < 0) { + record.claimed = false + return + } + const waiter = waiters.splice(index, 1)[0] + if (waiter === undefined) + throw new Error(`missing dynamic tool waiter for ${record.call.id}`) + record.claimed = true + waiter.delivered = record + waiter.resume(Effect.succeed(record.call)) + if (record.state === "cancelled") records.delete(record.call.id) + } + + const awaitBackend = ( + operation: LifecycleError["operation"], + callID: string | undefined, + ): Effect.Effect => + Effect.suspend(() => + closed + ? Effect.fail( + lifecycleError( + operation, + "controller-closed", + "dynamic tool controller is closed", + callID, + ), + ) + : Ref.get(current).pipe( + Effect.flatMap((attached) => + attached === undefined + ? Queue.take(backendChanges).pipe( + Effect.andThen(awaitBackend(operation, callID)), + ) + : Effect.succeed(attached.backend), + ), + ), + ) + + const whileConnected = ( + backend: BackendConnection, + operation: LifecycleError["operation"], + effect: Effect.Effect, + ) => + Effect.raceFirst( + effect, + backend.closed.pipe( + Effect.andThen( + Effect.fail( + new SimulationConnectionError({ + endpoint: backend.endpoint, + operation, + message: `dynamic tool ${operation} connection closed`, + }), + ), + ), + ), + ) + + const request = ( + operation: LifecycleError["operation"], + callID: string | undefined, + send: (backend: BackendConnection) => Effect.Effect, + ): Effect.Effect => + Effect.gen(function* () { + if (closed) + return yield* Effect.fail( + lifecycleError( + operation, + "controller-closed", + "dynamic tool controller is closed", + callID, + ), + ) + const backend = yield* awaitBackend(operation, callID) + const result = yield* Effect.exit( + whileConnected(backend, operation, send(backend)), + ) + if (Exit.isSuccess(result)) return result.value + if (Cause.hasInterrupts(result.cause)) return yield* Effect.interrupt + const found = Cause.findErrorOption(result.cause) + if (found._tag === "None") + return yield* Effect.die(Cause.squash(result.cause)) + if ( + found.value instanceof SimulationRequestError || + found.value instanceof SimulationCompatibilityError + ) + return yield* Effect.fail( + lifecycleError(operation, "rejected", found.value.message, callID), + ) + if ( + !(found.value instanceof SimulationConnectionError) && + !(found.value instanceof RpcClientError.RpcClientError) + ) + return yield* Effect.fail( + lifecycleError(operation, "rejected", String(found.value), callID), + ) + yield* Effect.sleep(25) + return yield* request(operation, callID, send) + }) + + const unavailable = ( + record: Active, + operation: LifecycleError["operation"], + ) => + record.state === "pending" + ? undefined + : lifecycleError( + operation, + record.state === "settled" ? "already-settled" : "cancelled", + record.state === "settled" + ? `dynamic tool call ${record.call.context.callID} is settled` + : `dynamic tool call ${record.call.context.callID} was cancelled`, + record.call.context.callID, + ) + + const operate = ( + record: Active, + operation: "progress" | "finish" | "fail", + effect: Effect.Effect, + ) => + record.operations.withPermit( + Effect.suspend(() => { + const error = unavailable(record, operation) + if (error !== undefined) return Effect.fail(error) + return Effect.raceFirst( + effect, + Deferred.await(record.cancelled).pipe( + Effect.andThen( + Effect.fail( + lifecycleError( + operation, + "cancelled", + `dynamic tool call ${record.call.context.callID} was cancelled`, + record.call.context.callID, + ), + ), + ), + ), + ) + }), + ) + + const commit = ( + record: Active, + operation: "finish" | "fail", + effect: Effect.Effect, + ) => + operate( + record, + operation, + effect.pipe( + Effect.andThen( + Effect.sync(() => { + record.state = "settled" + records.delete(record.call.id) + remember(record.call.id, record.fingerprint) + }), + ), + ), + ) + + const makeInvocation = ( + invocation: Backend.ToolInvocation, + fingerprint: string, + ): Active => { + const cancelled = Deferred.makeUnsafe() + const operations = Semaphore.makeUnsafe(1) + let record: Active + const call: Invocation = { + id: invocation.id, + name: invocation.name, + input: invocation.input, + context: invocation.context, + progress: (update: Progress) => + operate( + record, + "progress", + decodeProgress(update).pipe( + Effect.mapError((error) => + lifecycleError( + "progress", + "rejected", + error.message, + invocation.context.callID, + ), + ), + Effect.flatMap((decoded) => + request("progress", invocation.context.callID, (backend) => + backend.updateTool({ + id: invocation.id, + sequence: record.sequence, + update: decoded, + }), + ), + ), + Effect.andThen( + Effect.sync(() => { + record.sequence++ + }), + ), + ), + ), + finish: (output: Output) => + commit( + record, + "finish", + decodeOutput(output).pipe( + Effect.mapError((error) => + lifecycleError( + "finish", + "rejected", + error.message, + invocation.context.callID, + ), + ), + Effect.flatMap((decoded) => + request("finish", invocation.context.callID, (backend) => + backend.finishTool({ id: invocation.id, output: decoded }), + ), + ), + ), + ), + fail: (message) => + commit( + record, + "fail", + decodeFailure(message).pipe( + Effect.mapError((error) => + lifecycleError( + "fail", + "rejected", + error.message, + invocation.context.callID, + ), + ), + Effect.flatMap((decoded) => + request("fail", invocation.context.callID, (backend) => + backend.failTool({ id: invocation.id, message: decoded }), + ), + ), + ), + ), + awaitCancelled: () => Deferred.await(cancelled), + } + record = { + fingerprint, + call, + cancelled, + operations, + claimed: false, + sequence: 0, + state: "pending", + } + return record + } + + const receiveInvocation = (invocation: Backend.ToolInvocation) => + lifecycle.withPermit( + Effect.gen(function* () { + const fingerprint = fingerprintJson(invocation) + const existing = records.get(invocation.id) + if (existing !== undefined) { + if (existing.fingerprint !== fingerprint) + yield* Deferred.fail( + failure, + lifecycleError( + "take", + "rejected", + `dynamic tool invocation ${invocation.id} was replayed with different input`, + invocation.context.callID, + ), + ) + return + } + const settled = completed.get(invocation.id) + if (settled !== undefined) { + if (settled !== fingerprint) + yield* Deferred.fail( + failure, + lifecycleError( + "take", + "rejected", + `dynamic tool invocation ${invocation.id} was reused with different input`, + invocation.context.callID, + ), + ) + return + } + const record = makeInvocation(invocation, fingerprint) + records.set(invocation.id, record) + deliver(record) + }), + ) + + const receiveCancellation = (cancellation: Cancellation) => + lifecycle.withPermit( + Effect.sync(() => { + const record = records.get(cancellation.id) + if (record === undefined || record.state !== "pending") return + record.state = "cancelled" + remember(cancellation.id, record.fingerprint) + Deferred.doneUnsafe(record.cancelled, Effect.succeed(cancellation)) + if (record.claimed) records.delete(cancellation.id) + else deliver(record) + }), + ) + + const connect: Controller["connect"] = (backend) => + lifecycle.withPermit( + Effect.gen(function* () { + if ((yield* Ref.get(current)) !== undefined) + return yield* Effect.fail( + lifecycleError( + "attach", + "rejected", + "dynamic tool backend is already connected", + ), + ) + const scope = yield* Scope.fork(parentScope) + yield* backend.toolEvents.pipe( + Stream.runForEach((event) => + event.type === "invocation" + ? receiveInvocation(event.invocation) + : receiveCancellation(event.cancellation), + ), + Effect.matchCauseEffect({ + onFailure: (cause) => + Deferred.fail( + failure, + lifecycleError( + "take", + "rejected", + `dynamic tool event stream failed: ${Cause.pretty(cause)}`, + ), + ).pipe(Effect.asVoid), + onSuccess: () => Effect.void, + }), + Effect.forkIn(scope), + ) + const applied = yield* Effect.exit( + attachmentLock.withPermit( + Effect.suspend(() => + desired === undefined + ? Effect.void + : whileConnected( + backend, + "attach", + backend.attachTools(desired.tools), + ).pipe( + Effect.mapError((error) => + lifecycleError("attach", "rejected", String(error)), + ), + ), + ), + ), + ) + if (Exit.isFailure(applied)) { + yield* Scope.close(scope, Exit.void) + return yield* Effect.failCause(applied.cause) + } + yield* Ref.set(current, { backend, scope }) + yield* notifyBackend + const detach = Effect.fn("ToolProducer.detach")(function* () { + const shouldClose = yield* lifecycle.withPermit( + Effect.gen(function* () { + const attached = yield* Ref.get(current) + if (attached?.backend !== backend) return false + yield* Ref.set(current, undefined) + yield* notifyBackend + return true + }), + ) + if (shouldClose) yield* Scope.close(scope, Exit.void) + }) + return { detach } + }), + ) + + const attach: DynamicControls["attach"] = (params) => + Effect.gen(function* () { + const decoded = yield* decodeAttach(params).pipe( + Effect.mapError((error) => + lifecycleError("attach", "rejected", error.message), + ), + ) + const collision = decoded.tools + .map(Backend.exposedToolName) + .find((name) => staticNames.has(name)) + if (collision !== undefined) + yield* Effect.fail( + lifecycleError( + "attach", + "rejected", + `dynamic tool conflicts with configured static adapter: ${collision}`, + ), + ) + yield* request("attach", undefined, (backend) => + attachmentLock.withPermit( + backend.attachTools(decoded.tools).pipe( + Effect.andThen( + Effect.sync(() => { + desired = decoded + }), + ), + ), + ), + ) + return undefined + }) + + function take(): Effect.Effect + function take(callID: string): Effect.Effect + function take(callID?: string): Effect.Effect { + return Effect.callback((resume) => { + if (closed) { + resume( + Effect.fail( + lifecycleError( + "take", + "controller-closed", + "dynamic tool controller is closed", + callID, + ), + ), + ) + return undefined + } + if ( + callID !== undefined && + (Array.from(records.values()).some( + (record) => record.call.context.callID === callID && record.claimed, + ) || + waiters.some((waiter) => waiter.callID === callID)) + ) { + resume( + Effect.fail( + lifecycleError( + "take", + "already-claimed", + `dynamic tool call is already claimed: ${callID}`, + callID, + ), + ), + ) + return undefined + } + let record = Array.from(records.values()).find( + (candidate) => + !candidate.claimed && + (callID === undefined || candidate.call.context.callID === callID), + ) + if (record !== undefined) { + record.claimed = true + resume(Effect.succeed(record.call)) + if (record.state === "cancelled") records.delete(record.call.id) + return Effect.sync(() => { + if (record !== undefined && records.get(record.call.id) === record) { + record.claimed = false + deliver(record) + } + }) + } + const waiter: Waiter = { callID, resume } + waiters.push(waiter) + return Effect.sync(() => { + const index = waiters.indexOf(waiter) + if (index >= 0) waiters.splice(index, 1) + else if (waiter.delivered !== undefined) { + record = waiter.delivered + if (records.get(record.call.id) === record) { + record.claimed = false + deliver(record) + } + } + }) + }) + } + + const endGeneration = lifecycle.withPermit( + Effect.gen(function* () { + const attached = yield* Ref.get(current) + if (attached !== undefined) { + yield* Ref.set(current, undefined) + yield* Scope.close(attached.scope, Exit.void) + } + for (const record of records.values()) { + if (record.state !== "pending") continue + record.state = "cancelled" + Deferred.doneUnsafe( + record.cancelled, + Effect.succeed({ id: record.call.id, reason: "interrupted" }), + ) + } + records.clear() + completed.clear() + yield* notifyBackend + }), + ) + + const settle = lifecycle.withPermit( + Effect.suspend(() => { + const pending = Array.from(records.values()).filter( + (record) => record.state === "pending", + ).length + return pending === 0 + ? Effect.void + : Effect.fail( + lifecycleError( + "take", + "rejected", + `${pending} dynamic tool invocation(s) remain unsettled`, + ), + ) + }), + ) + + const shutdown = Effect.suspend(() => { + if (closed) return Effect.void + return Effect.sync(() => { + closed = true + for (const waiter of waiters) + waiter.resume( + Effect.fail( + lifecycleError( + "take", + "controller-closed", + "dynamic tool controller is closed", + waiter.callID, + ), + ), + ) + waiters.length = 0 + }).pipe( + Effect.andThen(endGeneration), + Effect.ensuring(Queue.shutdown(backendChanges)), + ) + }) + + yield* Effect.addFinalizer(() => shutdown) + + return { + controls: { attach, take }, + connect, + endGeneration, + settle, + shutdown, + failure: Deferred.await(failure), + } satisfies Controller +}) + +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") +} + +export * as ToolProducer from "./producer.js" diff --git a/packages/drive/src/tool/types.ts b/packages/drive/src/tool/types.ts index 189e792..feea129 100644 --- a/packages/drive/src/tool/types.ts +++ b/packages/drive/src/tool/types.ts @@ -1,5 +1,6 @@ import * as Effect from "effect/Effect" import * as Schema from "effect/Schema" +import type { Backend } from "../simulation/protocol.js" const PositiveInt = Schema.Number.check(Schema.isInt(), Schema.isGreaterThan(0)) @@ -147,6 +148,52 @@ export type ControlledCallsFor = ControlledCalls< ToolTypes[Tool]["result"] > -export interface Controls { +export interface StaticControls { control(name: Tool): Effect.Effect, ControlError> } + +export type Content = Backend.ToolContent +export type DynamicRegistration = Backend.ToolRegistration +export type AttachParams = Backend.ToolAttachParams +export type Progress = Backend.ToolProgress +export type Output = Backend.ToolOutput +export type Cancellation = Backend.ToolCancellation + +export class LifecycleError extends Schema.TaggedErrorClass()( + "OpenCodeDrive.ToolLifecycleError", + { + operation: Schema.Literals(["attach", "take", "progress", "finish", "fail"]), + reason: Schema.Literals([ + "controller-closed", + "already-claimed", + "already-settled", + "cancelled", + "rejected", + ]), + callID: Schema.optional(Schema.String), + message: Schema.String, + }, +) {} + +export interface Invocation { + /** Producer invocation ID, stable across controller reconnects. */ + readonly id: string + /** Effective provider-visible name, including any flattened namespace. */ + readonly name: string + readonly input: Backend.ToolInvocation["input"] + readonly context: Backend.ToolInvocation["context"] + readonly progress: (update: Progress) => Effect.Effect + readonly finish: (output: Output) => Effect.Effect + readonly fail: (message: string) => Effect.Effect + /** Completes only when OpenCode cancels before `finish` or `fail`. */ + readonly awaitCancelled: () => Effect.Effect +} + +export interface DynamicControls { + /** Atomically replaces the complete provider-backed dynamic tool set. */ + readonly attach: (params: AttachParams) => Effect.Effect + take(): Effect.Effect + take(callID: string): Effect.Effect +} + +export interface Controls extends StaticControls, DynamicControls {} diff --git a/packages/drive/test/driver/index.test.ts b/packages/drive/test/driver/index.test.ts index 40d5525..ed5f261 100644 --- a/packages/drive/test/driver/index.test.ts +++ b/packages/drive/test/driver/index.test.ts @@ -252,6 +252,87 @@ it.live("exposes runtime controls for tools declared by name", () => }), ) +it.live("controls arbitrary provider-backed tools through the runtime lifecycle", () => + Effect.gen(function* () { + const artifacts = yield* OpenCodeDriver.use( + { + keepArtifacts: true, + opencode: { + command: [...fakeOpenCode, "dynamic-tool", "reconnect-tool"], + }, + }, + (driver) => + Effect.gen(function* () { + yield* driver.tools.attach({ + tools: [ + { + name: "lookup", + description: "Look up a value", + inputSchema: { + type: "object", + properties: { query: { type: "string" } }, + required: ["query"], + }, + outputSchema: { + type: "object", + properties: { answer: { type: "number" } }, + required: ["answer"], + }, + options: { codemode: false }, + }, + ], + }) + const call = yield* driver.tools.take("call_lookup") + expect(call).toMatchObject({ + name: "lookup", + input: { query: "meaning" }, + context: { sessionID: "ses_dynamic", callID: "call_lookup" }, + }) + yield* call.progress({ + structured: { phase: "searching" }, + content: [{ type: "text", text: "Searching" }], + }) + yield* call.finish({ + structured: { answer: 42 }, + content: [{ type: "text", text: "42" }], + }) + return driver.artifacts + }), + ) + yield* Effect.addFinalizer(() => + Effect.promise(() => rm(artifacts, { recursive: true, force: true })), + ) + const events = (yield* Effect.promise(() => + readFile(`${artifacts}/tool-events.jsonl`, "utf8"), + )).trim().split("\n").map((line) => JSON.parse(line)) + expect(events).toEqual([ + expect.objectContaining({ method: "tool.attach" }), + expect.objectContaining({ method: "tool.attach" }), + { + method: "tool.update", + params: { + id: "tool_1", + sequence: 0, + update: { + structured: { phase: "searching" }, + content: [{ type: "text", text: "Searching" }], + }, + }, + }, + { + method: "tool.finish", + params: { + id: "tool_1", + output: { + structured: { answer: 42 }, + content: [{ type: "text", text: "42" }], + }, + }, + }, + ]) + }), +) + it.live("returns structured evidence from the safe lifecycle boundary", () => Effect.gen(function* () { const result = yield* OpenCodeDriver.useReport( diff --git a/packages/drive/test/driver/llm-controller.test.ts b/packages/drive/test/driver/llm-controller.test.ts index c222a09..733c5e6 100644 --- a/packages/drive/test/driver/llm-controller.test.ts +++ b/packages/drive/test/driver/llm-controller.test.ts @@ -102,7 +102,7 @@ describe("LlmController", () => { }) }) - it.live("streams tool call input as OpenAI argument deltas", () => { + it.live("streams tool call input as provider-neutral deltas", () => { const random = vi.spyOn(Math, "random").mockReturnValue(0.5) const peer = startTransportPeer(({ request: frame, socket }) => { if (frame.method === "llm.attach") @@ -144,25 +144,10 @@ describe("LlmController", () => { id: "exchange-1", items: [ { - type: "raw", - chunk: { - choices: [ - { - delta: { - tool_calls: [ - { - index: 0, - id: "call_1", - function: { - name: "lookup", - arguments: '{"query":"', - }, - }, - ], - }, - }, - ], - }, + type: "toolInputStart", + index: 0, + id: "call_1", + name: "lookup", }, ], }, @@ -172,21 +157,21 @@ describe("LlmController", () => { id: "exchange-1", items: [ { - type: "raw", - chunk: { - choices: [ - { - delta: { - tool_calls: [ - { - index: 0, - function: { arguments: 'weather"}' }, - }, - ], - }, - }, - ], - }, + type: "toolInputDelta", + index: 0, + text: '{"query":"', + }, + ], + }, + }), + expect.objectContaining({ + params: { + id: "exchange-1", + items: [ + { + type: "toolInputDelta", + index: 0, + text: 'weather"}', }, ], }, @@ -195,6 +180,81 @@ describe("LlmController", () => { }) }) + it.live("falls back to OpenAI deltas when partial tool input is unavailable", () => { + const peer = startTransportPeer( + ({ request: frame, socket }) => { + if (frame.method === "llm.attach") + socket.send( + JSON.stringify({ + jsonrpc: "2.0", + method: "llm.request", + params: request, + }), + ) + sendResult( + socket, + frame, + frame.method === "llm.attach" ? { attached: true } : { ok: true }, + ) + }, + { + capabilities: (offered) => + offered.filter((capability) => capability !== "llm.tool-input-delta"), + }, + ) + + return Effect.gen(function* () { + yield* Effect.addFinalizer(() => Effect.promise(() => peer.stop())) + const backend = yield* SimulationConnector.backend(peer.url) + const llm = yield* LlmController.make(backend) + yield* llm.queue( + Llm.toolCall( + { + index: 0, + id: "call_legacy", + name: "lookup", + input: { query: "weather" }, + }, + { delay: 0, chunkSize: 100 }, + ), + Llm.finish("tool-calls"), + ) + yield* llm.settle() + + const chunk = peer.received.find( + ({ request: frame }) => frame.method === "llm.chunk", + )?.request + expect(chunk).toMatchObject({ + params: { + id: "exchange-1", + items: [ + { + type: "raw", + chunk: { + choices: [ + { + delta: { + tool_calls: [ + { + index: 0, + id: "call_legacy", + function: { + name: "lookup", + arguments: '{"query":"weather"}', + }, + }, + ], + }, + }, + ], + }, + }, + ], + }, + }) + }) + }) + it.live("serves requests and keeps titles outside normal sequencing", () => { const random = vi.spyOn(Math, "random").mockReturnValue(0.5) const titleRequest = { diff --git a/packages/drive/test/fixtures/fake-opencode.ts b/packages/drive/test/fixtures/fake-opencode.ts index 9438a24..5c9c7b7 100644 --- a/packages/drive/test/fixtures/fake-opencode.ts +++ b/packages/drive/test/fixtures/fake-opencode.ts @@ -22,6 +22,7 @@ const screen = { value: `Fake OpenCode${role === "client" ? ` ${process.env.OPEN const drive = await resolveDrive() const recordingStarted = performance.now() const endpoints = drive.endpoints +let toolAttachments = 0 const servicePassword = "drive-test-password" const api = role === "service" ? Bun.serve({ @@ -138,14 +139,53 @@ const backend = role === "client" ? undefined : Bun.serve({ `${JSON.stringify({ method: request.method, params: request.params })}\n`, ) } + if (request.method.startsWith("tool.") && process.env.OPENCODE_TEST_HOME) { + const events = `${process.env.OPENCODE_TEST_HOME}/tool-events.jsonl` + await appendFile( + events, + `${JSON.stringify({ method: request.method, params: request.params })}\n`, + ) + } + if ( + request.method === "tool.attach" && + process.argv.includes("dynamic-tool") + ) + if ( + (++toolAttachments === 1 && + !process.argv.includes("reconnect-tool")) || + toolAttachments === 2 + ) + socket.send( + JSON.stringify({ + jsonrpc: "2.0", + method: "tool.invocation", + params: { + id: "tool_1", + name: "lookup", + input: { query: "meaning" }, + context: { + sessionID: "ses_dynamic", + agent: "build", + messageID: "msg_dynamic", + callID: "call_lookup", + }, + }, + }), + ) const result = request.method === "simulation.handshake" ? handshake(request.params, "backend") - : request.method === "llm.attach" + : request.method === "llm.attach" || request.method === "tool.attach" ? { attached: true } : { ok: true } if (request.id !== undefined) socket.send(JSON.stringify({ jsonrpc: "2.0", id: request.id, result })) + if ( + request.method === "tool.attach" && + process.argv.includes("reconnect-tool") && + toolAttachments === 1 + ) + setTimeout(() => socket.close(), 10) if (request.method === "llm.attach") { const requestDelay = Number(process.argv[3]) if (Number.isFinite(requestDelay)) await Bun.sleep(requestDelay) diff --git a/packages/drive/test/simulation/backend.test.ts b/packages/drive/test/simulation/backend.test.ts index f2b78dc..8092f89 100644 --- a/packages/drive/test/simulation/backend.test.ts +++ b/packages/drive/test/simulation/backend.test.ts @@ -1,5 +1,5 @@ import { describe, expect, it } from "@effect/vitest" -import { Effect, Option, Stream } from "effect" +import { Effect, Fiber, Option, Stream } from "effect" import { Backend } from "../../src/client/index.js" import * as SimulationConnector from "../../src/simulation/connector.js" import { sendResult, startTransportPeer } from "./transport-peer.js" @@ -192,4 +192,102 @@ describe("OpenCode backend simulation transport", () => { expect([...received]).toEqual([exchanges.first, exchanges.second]) }), ) + + it.live("mirrors tool lifecycle RPCs and preserves invocation-cancellation order", () => + Effect.gen(function* () { + const invocation = { + id: "tool_1", + name: "lookup", + input: { query: "meaning" }, + context: { + sessionID: "ses_tools", + agent: "build", + messageID: "msg_tools", + callID: "call_lookup", + }, + } + const peer = startTransportPeer(({ request, socket }) => { + if (request.method === "tool.attach") { + sendNotification(socket, "tool.invocation", invocation) + sendNotification(socket, "tool.cancel", { + id: "tool_1", + reason: "interrupted", + }) + sendResult(socket, request, { attached: true }) + return + } + sendResult(socket, request, { ok: true }) + }) + yield* Effect.addFinalizer(() => Effect.promise(() => peer.stop())) + const connection = yield* SimulationConnector.backend(peer.url, { + attach: false, + }) + const events = yield* connection.toolEvents.pipe( + Stream.take(2), + Stream.runCollect, + Effect.forkScoped, + ) + const tools = [ + { + name: "lookup", + description: "Look up a value", + inputSchema: { type: "object" }, + options: { codemode: false }, + }, + ] + + yield* connection.attachTools(tools) + yield* connection.updateTool({ + id: "tool_1", + sequence: 0, + update: { structured: { phase: "searching" } }, + }) + yield* connection.finishTool({ + id: "tool_1", + output: { structured: { answer: 42 }, content: [] }, + }) + yield* connection.failTool({ id: "tool_2", message: "failed" }) + + expect([...(yield* Fiber.join(events))]).toEqual([ + { type: "invocation", invocation }, + { + type: "cancellation", + cancellation: { id: "tool_1", reason: "interrupted" }, + }, + ]) + expect(peer.received.map(({ request }) => request)).toEqual([ + { + jsonrpc: "2.0", + id: 1, + method: "tool.attach", + params: { tools }, + }, + { + jsonrpc: "2.0", + id: 2, + method: "tool.update", + params: { + id: "tool_1", + sequence: 0, + update: { structured: { phase: "searching" } }, + }, + }, + { + jsonrpc: "2.0", + id: 3, + method: "tool.finish", + params: { + id: "tool_1", + output: { structured: { answer: 42 }, content: [] }, + }, + }, + { + jsonrpc: "2.0", + id: 4, + method: "tool.fail", + params: { id: "tool_2", message: "failed" }, + }, + ]) + }), + ) }) diff --git a/packages/drive/test/simulation/connector.test.ts b/packages/drive/test/simulation/connector.test.ts index f5cc4ba..db7e983 100644 --- a/packages/drive/test/simulation/connector.test.ts +++ b/packages/drive/test/simulation/connector.test.ts @@ -109,4 +109,40 @@ describe("SimulationConnector", () => { expect(Schema.isSchemaError(error)).toBe(true) }).pipe(Effect.provide(SimulationConnector.layer)), ) + + it.live("keeps tool capabilities optional until dynamic tools are attached", () => + Effect.gen(function* () { + const peer = startTransportPeer( + ({ request, socket }) => + sendResult( + socket, + request, + request.method === "llm.attach" + ? { attached: true } + : { ok: true }, + ), + { + capabilities: (offered) => + offered.filter((capability) => !capability.startsWith("tool.")), + }, + ) + yield* Effect.addFinalizer(() => Effect.promise(() => peer.stop())) + + const connection = yield* SimulationConnector.backend(peer.url) + expect(connection.compatibility).toMatchObject({ + _tag: "Negotiated", + role: "backend", + }) + expect( + yield* connection.attachTools([]).pipe(Effect.flip), + ).toMatchObject({ + _tag: "SimulationCompatibilityError", + role: "backend", + message: expect.stringContaining("tool.attach"), + }) + expect(peer.received.map(({ request }) => request.method)).toEqual([ + "llm.attach", + ]) + }).pipe(Effect.provide(SimulationConnector.layer)), + ) }) diff --git a/packages/drive/test/simulation/protocol.test.ts b/packages/drive/test/simulation/protocol.test.ts new file mode 100644 index 0000000..2b39cc8 --- /dev/null +++ b/packages/drive/test/simulation/protocol.test.ts @@ -0,0 +1,98 @@ +import { expect, it } from "@effect/vitest" +import { Backend } from "../../src/simulation/protocol.js" + +it("decodes the canonical dynamic tool lifecycle", () => { + expect( + Backend.decodeRequest({ + jsonrpc: "2.0", + id: 1, + method: "tool.attach", + params: { + tools: [ + { + name: "search", + description: "Search GitHub", + inputSchema: { type: "object" }, + outputSchema: { type: "object" }, + permission: "search", + options: { namespace: "github.api", codemode: false }, + }, + ], + }, + }), + ).toMatchObject({ method: "tool.attach" }) + 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" }) +}) + +it("rejects unsafe, colliding, and reserved exposed names", () => { + const attach = (tools: ReadonlyArray) => + Backend.decodeRequest({ + jsonrpc: "2.0", + id: 1, + method: "tool.attach", + params: { tools }, + }) + const tool = { + name: "search", + description: "Search", + inputSchema: { type: "object" }, + } + + expect(() => attach([{ ...tool, name: "unsafe.name" }])).toThrow() + expect(() => + attach([ + { ...tool, options: { namespace: "a.b" } }, + { ...tool, options: { namespace: "a_b" } }, + ]), + ).toThrow() + expect(() => + attach([ + { + ...tool, + name: "execute", + options: { codemode: false }, + }, + ]), + ).toThrow() +}) + +it("decodes provider-neutral partial tool input chunks", () => { + expect( + Backend.Item.make({ + type: "toolInputStart", + index: 0, + id: "call_lookup", + name: "lookup", + }), + ).toEqual({ + type: "toolInputStart", + index: 0, + id: "call_lookup", + name: "lookup", + }) + expect( + Backend.Item.make({ + type: "toolInputDelta", + index: 0, + text: '{"query":"meaning"}', + }), + ).toEqual({ + type: "toolInputDelta", + index: 0, + text: '{"query":"meaning"}', + }) +}) diff --git a/packages/drive/test/simulation/transport-peer.ts b/packages/drive/test/simulation/transport-peer.ts index c912158..019266b 100644 --- a/packages/drive/test/simulation/transport-peer.ts +++ b/packages/drive/test/simulation/transport-peer.ts @@ -13,7 +13,12 @@ export interface ReceivedRequest { export function startTransportPeer( onRequest: (received: ReceivedRequest) => void, - options?: { readonly handshake?: boolean }, + options?: { + readonly handshake?: boolean + readonly capabilities?: ( + offered: ReadonlyArray, + ) => ReadonlyArray + }, ) { const received: ReceivedRequest[] = [] const server = Bun.serve({ @@ -40,7 +45,10 @@ export function startTransportPeer( protocolVersion: 1, role: params.expectedRole, server: { name: "opencode", version: "test" }, - capabilities: [ + capabilities: options?.capabilities?.([ + ...params.requiredCapabilities, + ...params.optionalCapabilities, + ]) ?? [ ...params.requiredCapabilities, ...params.optionalCapabilities, ], diff --git a/packages/drive/test/tool/producer.test.ts b/packages/drive/test/tool/producer.test.ts new file mode 100644 index 0000000..f22165f --- /dev/null +++ b/packages/drive/test/tool/producer.test.ts @@ -0,0 +1,470 @@ +import { expect, it } from "@effect/vitest" +import { Deferred, Effect, Exit, Fiber, Scope } from "effect" +import * as SimulationConnector from "../../src/simulation/connector.js" +import * as ToolProducer from "../../src/tool/producer.js" +import type { Progress } from "../../src/tool/types.js" +import { sendResult, startTransportPeer } from "../simulation/transport-peer.js" + +const registration = { + name: "lookup", + description: "Look up a value", + inputSchema: { + type: "object", + properties: { query: { type: "string" } }, + required: ["query"], + additionalProperties: false, + }, + outputSchema: { + type: "object", + properties: { answer: { type: "number" } }, + required: ["answer"], + }, + options: { codemode: false }, +} as const + +const invocation = { + id: "tool_1", + name: "lookup", + input: { query: "meaning" }, + context: { + sessionID: "ses_tools", + agent: "build", + messageID: "msg_tools", + callID: "call_lookup", + }, +} as const + +function notify(socket: Bun.ServerWebSocket, method: "tool.invocation" | "tool.cancel", params: unknown) { + socket.send(JSON.stringify({ jsonrpc: "2.0", method, params })) +} + +it.live("attaches, sequences progress, and settles one dynamic invocation", () => + Effect.scoped( + Effect.gen(function* () { + const peer = startTransportPeer(({ request, socket }) => { + sendResult(socket, request, request.method === "tool.attach" ? { attached: true } : { ok: true }) + }) + yield* Effect.addFinalizer(() => Effect.promise(() => peer.stop())) + const backend = yield* SimulationConnector.backend(peer.url, { + attach: false, + }) + const controller = yield* ToolProducer.make(new Set()) + yield* controller.connect(backend) + + yield* controller.controls.attach({ tools: [registration] }) + const socket = peer.received[0].socket + notify(socket, "tool.invocation", invocation) + const call = yield* controller.controls.take("call_lookup") + + expect(call).toMatchObject(invocation) + yield* call.progress({ + structured: { phase: "searching" }, + content: [{ type: "text", text: "Searching" }], + }) + yield* call.progress({ structured: { phase: "done" } }) + yield* call.finish({ + structured: { answer: 42 }, + content: [{ type: "text", text: "42" }], + }) + + expect(peer.received.map(({ request }) => request)).toEqual([ + { + jsonrpc: "2.0", + id: 1, + method: "tool.attach", + params: { tools: [registration] }, + }, + { + jsonrpc: "2.0", + id: 2, + method: "tool.update", + params: { + id: "tool_1", + sequence: 0, + update: { + structured: { phase: "searching" }, + content: [{ type: "text", text: "Searching" }], + }, + }, + }, + { + jsonrpc: "2.0", + id: 3, + method: "tool.update", + params: { + id: "tool_1", + sequence: 1, + update: { structured: { phase: "done" } }, + }, + }, + { + jsonrpc: "2.0", + id: 4, + method: "tool.finish", + params: { + id: "tool_1", + output: { + structured: { answer: 42 }, + content: [{ type: "text", text: "42" }], + }, + }, + }, + ]) + expect(yield* call.fail("late").pipe(Effect.flip)).toMatchObject({ + _tag: "OpenCodeDrive.ToolLifecycleError", + operation: "fail", + reason: "already-settled", + callID: "call_lookup", + }) + yield* controller.settle + }), + ), +) + +it.live("observes cancellation and rejects terminal work after it wins", () => + Effect.scoped( + Effect.gen(function* () { + const peer = startTransportPeer(({ request, socket }) => + sendResult(socket, request, request.method === "tool.attach" ? { attached: true } : { ok: true }), + ) + yield* Effect.addFinalizer(() => Effect.promise(() => peer.stop())) + const backend = yield* SimulationConnector.backend(peer.url, { + attach: false, + }) + const controller = yield* ToolProducer.make(new Set()) + yield* controller.connect(backend) + yield* controller.controls.attach({ tools: [registration] }) + const socket = peer.received[0].socket + notify(socket, "tool.invocation", invocation) + notify(socket, "tool.cancel", { id: "tool_1", reason: "interrupted" }) + yield* Effect.sleep(10) + const call = yield* controller.controls.take("call_lookup") + const cancelled = yield* call.awaitCancelled().pipe(Effect.forkScoped) + + expect(yield* Fiber.join(cancelled)).toEqual({ + id: "tool_1", + reason: "interrupted", + }) + expect(yield* call.finish({ structured: 42, content: [] }).pipe(Effect.flip)).toMatchObject({ + operation: "finish", + reason: "cancelled", + callID: "call_lookup", + }) + yield* controller.settle + }), + ), +) + +it.live("settles concurrent invocations in controlled reverse order", () => + Effect.scoped( + Effect.gen(function* () { + const peer = startTransportPeer(({ request, socket }) => + sendResult(socket, request, request.method === "tool.attach" ? { attached: true } : { ok: true }), + ) + yield* Effect.addFinalizer(() => Effect.promise(() => peer.stop())) + const backend = yield* SimulationConnector.backend(peer.url, { attach: false }) + const controller = yield* ToolProducer.make(new Set()) + yield* controller.connect(backend) + yield* controller.controls.attach({ tools: [registration] }) + const socket = peer.received[0].socket + notify(socket, "tool.invocation", invocation) + notify(socket, "tool.invocation", { + ...invocation, + id: "tool_2", + input: { query: "second" }, + context: { ...invocation.context, callID: "call_second" }, + }) + + const second = yield* controller.controls.take("call_second") + const first = yield* controller.controls.take("call_lookup") + const start = yield* Deferred.make() + const attempts = yield* Effect.all( + [ + Deferred.await(start).pipe( + Effect.andThen(second.finish({ structured: 2, content: [] })), + Effect.exit, + ), + Deferred.await(start).pipe( + Effect.andThen(second.fail("second failed")), + Effect.exit, + ), + ], + { concurrency: "unbounded" }, + ).pipe(Effect.forkScoped) + yield* Deferred.succeed(start, undefined) + const exits = yield* Fiber.join(attempts) + expect(exits.filter(Exit.isSuccess)).toHaveLength(1) + expect(exits.filter(Exit.isFailure)).toHaveLength(1) + yield* first.fail("first failed") + + const terminals = peer.received.filter(({ request }) => + request.method === "tool.finish" || request.method === "tool.fail", + ) + expect(terminals).toHaveLength(2) + expect(terminals[0]?.request.params).toMatchObject({ id: "tool_2" }) + expect(terminals[1]?.request).toMatchObject({ + method: "tool.fail", + params: { id: "tool_1", message: "first failed" }, + }) + yield* controller.settle + }), + ), +) + +it.live("reattaches the desired set and deduplicates replay after reconnect", () => + Effect.scoped( + Effect.gen(function* () { + let attaches = 0 + const peer = startTransportPeer(({ request, socket }) => { + if (request.method === "tool.attach") { + attaches++ + if (attaches === 2) notify(socket, "tool.invocation", invocation) + sendResult(socket, request, { attached: true }) + return + } + sendResult(socket, request, { ok: true }) + }) + yield* Effect.addFinalizer(() => Effect.promise(() => peer.stop())) + const controller = yield* ToolProducer.make(new Set()) + const firstScope = yield* Scope.make() + const first = yield* SimulationConnector.backend(peer.url, { + attach: false, + }).pipe(Scope.provide(firstScope)) + const firstAttachment = yield* controller.connect(first) + yield* controller.controls.attach({ tools: [registration] }) + notify(peer.received[0].socket, "tool.invocation", invocation) + const call = yield* controller.controls.take("call_lookup") + + yield* firstAttachment.detach() + yield* Scope.close(firstScope, Exit.void) + const secondScope = yield* Scope.make() + const second = yield* SimulationConnector.backend(peer.url, { + attach: false, + }).pipe(Scope.provide(secondScope)) + const secondAttachment = yield* controller.connect(second) + + expect(yield* controller.controls.take("call_lookup").pipe(Effect.flip)).toMatchObject({ + reason: "already-claimed", + }) + yield* call.finish({ + structured: { answer: 42 }, + content: [{ type: "text", text: "42" }], + }) + expect(attaches).toBe(2) + expect(peer.received.filter(({ request }) => request.method === "tool.finish")).toHaveLength(1) + yield* controller.settle + yield* secondAttachment.detach() + yield* Scope.close(secondScope, Exit.void) + yield* controller.endGeneration + yield* controller.shutdown + }), + ), +) + +it.live("replays the current set before applying an attachment started while disconnected", () => + Effect.scoped( + Effect.gen(function* () { + const peer = startTransportPeer(({ request, socket }) => + sendResult(socket, request, { attached: true }), + ) + yield* Effect.addFinalizer(() => Effect.promise(() => peer.stop())) + const controller = yield* ToolProducer.make(new Set()) + const firstScope = yield* Scope.make() + const first = yield* SimulationConnector.backend(peer.url, { + attach: false, + }).pipe(Scope.provide(firstScope)) + const firstAttachment = yield* controller.connect(first) + yield* controller.controls.attach({ tools: [registration] }) + yield* firstAttachment.detach() + yield* Scope.close(firstScope, Exit.void) + + const replacement = { + ...registration, + name: "search", + description: "Search for a value", + } + const replacing = yield* controller.controls + .attach({ tools: [replacement] }) + .pipe(Effect.forkScoped) + yield* Effect.yieldNow + const secondScope = yield* Scope.make() + const second = yield* SimulationConnector.backend(peer.url, { + attach: false, + }).pipe(Scope.provide(secondScope)) + const secondAttachment = yield* controller.connect(second) + yield* Fiber.join(replacing) + + expect( + peer.received + .filter(({ request }) => request.method === "tool.attach") + .map(({ request }) => request.params), + ).toEqual([ + { tools: [registration] }, + { tools: [registration] }, + { tools: [replacement] }, + ]) + yield* secondAttachment.detach() + yield* Scope.close(secondScope, Exit.void) + }), + ), +) + +it.live("retries in-flight progress after reconnect without advancing its sequence", () => + Effect.scoped( + Effect.gen(function* () { + let updates = 0 + const firstUpdate = yield* Deferred.make() + const peer = startTransportPeer(({ request, socket }) => { + if (request.method === "tool.update" && ++updates === 1) { + Deferred.doneUnsafe(firstUpdate, Effect.void) + socket.close() + return + } + sendResult(socket, request, request.method === "tool.attach" ? { attached: true } : { ok: true }) + }) + yield* Effect.addFinalizer(() => + Effect.sync(() => { + // Bun does not resolve server.stop after this peer initiates a WebSocket close. + void peer.stop() + }), + ) + const controller = yield* ToolProducer.make(new Set()) + const firstScope = yield* Scope.make() + const first = yield* SimulationConnector.backend(peer.url, { + attach: false, + }).pipe(Scope.provide(firstScope)) + const firstAttachment = yield* controller.connect(first) + yield* controller.controls.attach({ tools: [registration] }) + notify(peer.received[0].socket, "tool.invocation", invocation) + const call = yield* controller.controls.take("call_lookup") + + const progress = yield* call + .progress({ structured: { phase: "searching" } }) + .pipe(Effect.forkScoped) + yield* Deferred.await(firstUpdate) + yield* first.closed + yield* firstAttachment.detach() + yield* Scope.close(firstScope, Exit.void) + const secondScope = yield* Scope.make() + const second = yield* SimulationConnector.backend(peer.url, { + attach: false, + }).pipe(Scope.provide(secondScope)) + const secondAttachment = yield* controller.connect(second) + yield* Fiber.join(progress) + yield* call.fail("finished") + + expect( + peer.received.filter(({ request }) => request.method === "tool.update").map(({ request }) => request.params), + ).toEqual([ + { + id: "tool_1", + sequence: 0, + update: { structured: { phase: "searching" } }, + }, + { + id: "tool_1", + sequence: 0, + update: { structured: { phase: "searching" } }, + }, + ]) + yield* secondAttachment.detach() + yield* Scope.close(secondScope, Exit.void) + yield* controller.endGeneration + yield* controller.shutdown + }), + ), +) + +it.live("does not retry progress interrupted by its caller", () => + Effect.scoped( + Effect.gen(function* () { + const updateReceived = yield* Deferred.make() + const peer = startTransportPeer(({ request, socket }) => { + if (request.method === "tool.update") { + Deferred.doneUnsafe(updateReceived, Effect.void) + return + } + sendResult(socket, request, request.method === "tool.attach" ? { attached: true } : { ok: true }) + }) + yield* Effect.addFinalizer(() => Effect.promise(() => peer.stop())) + const backend = yield* SimulationConnector.backend(peer.url, { attach: false }) + const controller = yield* ToolProducer.make(new Set()) + yield* controller.connect(backend) + yield* controller.controls.attach({ tools: [registration] }) + notify(peer.received[0].socket, "tool.invocation", invocation) + const call = yield* controller.controls.take("call_lookup") + + const progress = yield* call + .progress({ structured: { phase: "searching" } }) + .pipe(Effect.forkScoped) + yield* Deferred.await(updateReceived) + yield* Fiber.interrupt(progress) + yield* call.fail("finished") + + expect(peer.received.filter(({ request }) => request.method === "tool.update")).toHaveLength(1) + }), + ), +) + +it.live("rejects malformed progress without advancing its sequence", () => + Effect.scoped( + Effect.gen(function* () { + const peer = startTransportPeer(({ request, socket }) => + sendResult(socket, request, request.method === "tool.attach" ? { attached: true } : { ok: true }), + ) + yield* Effect.addFinalizer(() => Effect.promise(() => peer.stop())) + const backend = yield* SimulationConnector.backend(peer.url, { attach: false }) + const controller = yield* ToolProducer.make(new Set()) + yield* controller.connect(backend) + yield* controller.controls.attach({ tools: [registration] }) + notify(peer.received[0].socket, "tool.invocation", invocation) + const call = yield* controller.controls.take("call_lookup") + + const malformed = { structured: [] } as unknown as Progress + expect(yield* call.progress(malformed).pipe(Effect.flip)).toMatchObject({ + operation: "progress", + reason: "rejected", + callID: "call_lookup", + }) + yield* call.progress({ structured: { phase: "valid" } }) + yield* call.fail("finished") + + expect( + peer.received.filter(({ request }) => request.method === "tool.update").map(({ request }) => request.params), + ).toEqual([ + { + id: "tool_1", + sequence: 0, + update: { structured: { phase: "valid" } }, + }, + ]) + }), + ), +) + +it.live("rejects dynamic names that collide with static adapters", () => + Effect.scoped( + Effect.gen(function* () { + const peer = startTransportPeer(({ request, socket }) => sendResult(socket, request, { attached: true })) + yield* Effect.addFinalizer(() => Effect.promise(() => peer.stop())) + const backend = yield* SimulationConnector.backend(peer.url, { + attach: false, + }) + const controller = yield* ToolProducer.make(new Set(["shell"])) + yield* controller.connect(backend) + + expect( + yield* controller.controls + .attach({ + tools: [{ ...registration, name: "shell" }], + }) + .pipe(Effect.flip), + ).toMatchObject({ + operation: "attach", + reason: "rejected", + message: expect.stringContaining("static adapter: shell"), + }) + expect(peer.received).toEqual([]) + }), + ), +) diff --git a/skills/opencode-drive/SKILL.md b/skills/opencode-drive/SKILL.md index d549bc9..c11835f 100644 --- a/skills/opencode-drive/SKILL.md +++ b/skills/opencode-drive/SKILL.md @@ -183,6 +183,42 @@ yield* secondary.ui.screenshot("secondary") Drive prefers protocol negotiation and reports explicit legacy fallback. Set `opencode.compatibility` to `"required"` when protocol skew must fail before the program runs. Additional built-in tool adapters remain follow-ups. +### Arbitrary Dynamic Tools + +Use runtime `tools.attach` for tools that do not have a shipped Drive adapter. +Attachment replaces the complete dynamic set without affecting configured +static adapters. Take invocations by the model call ID, while Drive owns +producer IDs, progress sequences, reconnect replay, and exactly-one terminal +commitment. + +```ts +yield* tools.attach({ + tools: [ + { + name: "lookup", + description: "Look up a value", + inputSchema: { + type: "object", + properties: { query: { type: "string" } }, + required: ["query"], + }, + options: { codemode: false }, + }, + ], +}) + +const lookup = yield* tools.take("call_lookup") +yield* lookup.progress({ structured: { phase: "searching" } }) +yield* lookup.finish({ + structured: { answer: 42 }, + content: [{ type: "text", text: "42" }], +}) +``` + +Use `awaitCancelled()` to observe native interruption. Do not synthesize +cancellation or expose transport sequence numbers. Dynamic effective names may +not collide with configured `shell`, `webfetch`, or `websearch` adapters. + ### Simulated Shell Execution Declare the built-in tools Drive should intercept, then control each invocation From a45bb7d460ab9be65e2dea4d5edf4902bb11f8da Mon Sep 17 00:00:00 2001 From: Kit Langton Date: Sun, 19 Jul 2026 21:41:34 -0400 Subject: [PATCH 2/3] fix(drive): harden dynamic tool lifecycle --- .../docs/open-code-driver-architecture.md | 4 + packages/drive/src/driver/server.ts | 89 +++-- packages/drive/src/simulation/connector.ts | 8 + packages/drive/src/tool/producer.ts | 306 ++++++++++++------ packages/drive/src/tool/types.ts | 1 + packages/drive/test/driver/index.test.ts | 81 +++++ packages/drive/test/fixtures/fake-opencode.ts | 31 ++ packages/drive/test/tool/producer.test.ts | 196 ++++++++++- 8 files changed, 581 insertions(+), 135 deletions(-) diff --git a/packages/drive/docs/open-code-driver-architecture.md b/packages/drive/docs/open-code-driver-architecture.md index cd0a661..279bab9 100644 --- a/packages/drive/docs/open-code-driver-architecture.md +++ b/packages/drive/docs/open-code-driver-architecture.md @@ -112,6 +112,10 @@ invocations safely. One ordered event stream preserves invocation-before- cancellation order. The desired registration set survives reconnects and manual server relaunches; invocation records are scoped to one server generation because producer IDs may be reused by a new process. +Settlement first clears the dynamic registration set on OpenCode, then drains +the ordered local event stream before checking for unresolved invocations. The +clear acts as the server-side barrier that prevents a native invocation from +appearing after a successful settlement snapshot. ## TUI Lifecycle diff --git a/packages/drive/src/driver/server.ts b/packages/drive/src/driver/server.ts index edcec81..3d92b3e 100644 --- a/packages/drive/src/driver/server.ts +++ b/packages/drive/src/driver/server.ts @@ -1,9 +1,11 @@ import * as Effect from "effect/Effect" +import * as Cause from "effect/Cause" import * as Deferred from "effect/Deferred" import * as Ref from "effect/Ref" import * as Semaphore from "effect/Semaphore" import * as Scope from "effect/Scope" import * as Exit from "effect/Exit" +import { RpcClientError } from "effect/unstable/rpc" import * as OpenCodeInstance from "../instance/runtime.js" import * as SimulationConnector from "../simulation/connector.js" import * as OpenCodeTui from "./client.js" @@ -12,6 +14,7 @@ import { error, type OpenCodeDriverError } from "./error.js" import * as LlmController from "./llm-controller.js" import * as ToolProducer from "../tool/producer.js" import type * as Tool from "../tool/index.js" +import { LifecycleError } from "../tool/types.js" export interface Target { readonly command?: ReadonlyArray @@ -71,6 +74,7 @@ export const make = Effect.fn("OpenCodeServer.make")(function* ( } | undefined >(undefined) const unexpectedExit = yield* Deferred.make() + const toolConnectionFailure = yield* Deferred.make() const lifecycle = yield* Semaphore.make(1) let compatibility: ReadonlyArray< SimulationConnector.EndpointCompatibility @@ -86,17 +90,16 @@ export const make = Effect.fn("OpenCodeServer.make")(function* ( Effect.mapError((cause) => error("server.launch", cause)), Effect.onError(() => Scope.close(scope, Exit.void)), ) + const rollbackLaunch = Scope.close(scope, Exit.void).pipe( + Effect.andThen(toolProducer.endGeneration), + Effect.andThen(instance.killServer.pipe(Effect.ignore)), + ) const backend = yield* connector.backend(launched.endpoint, { compatibility: target.compatibility, }).pipe( Scope.provide(scope), Effect.mapError((cause) => error("server.connect", cause)), - Effect.onError(() => - instance.killServer.pipe( - Effect.ignore, - Effect.andThen(Scope.close(scope, Exit.void)), - ), - ), + Effect.onError(() => rollbackLaunch), ) const connectTools = Effect.fn("OpenCodeServer.connectTools")(function* () { const connectionScope = yield* Scope.fork(scope) @@ -105,19 +108,35 @@ export const make = Effect.fn("OpenCodeServer.make")(function* ( compatibility: target.compatibility, }).pipe( Scope.provide(connectionScope), - Effect.mapError((cause) => error("tools.connect", cause)), Effect.onError(() => Scope.close(connectionScope, Exit.void)), ) const attachment = yield* toolProducer.connect(backend).pipe( - Effect.mapError((cause) => error("tools.connect", cause)), Effect.onError(() => Scope.close(connectionScope, Exit.void)), ) return { attachment, backend, scope: connectionScope } }) + const failToolConnection = (cause: unknown) => + toolProducer.shutdown.pipe( + Effect.andThen( + Deferred.fail( + toolConnectionFailure, + error("tools.connect", cause), + ), + ), + Effect.asVoid, + ) function reconnectTools(): Effect.Effect { return connectTools().pipe( Effect.flatMap(superviseTools), - Effect.catch(() => Effect.sleep(25).pipe(Effect.andThen(reconnectTools()))), + Effect.catchIf( + isRetryableToolConnectionError, + () => Effect.sleep(25).pipe(Effect.andThen(reconnectTools())), + ), + Effect.catch(failToolConnection), + Effect.catchCauseIf( + (cause) => !Cause.hasInterrupts(cause), + (cause) => failToolConnection(Cause.pretty(cause)), + ), ) } function superviseTools( @@ -133,34 +152,22 @@ export const make = Effect.fn("OpenCodeServer.make")(function* ( Effect.andThen(reconnectTools()), ) } - const toolConnection = yield* connectTools() + const toolConnection = yield* connectTools().pipe( + Effect.mapError((cause) => error("tools.connect", cause)), + Effect.onError(() => rollbackLaunch), + ) yield* superviseTools(toolConnection).pipe( Effect.forkIn(scope), ) const attachment = yield* llm.attach(backend).pipe( - Effect.onError(() => - instance.killServer.pipe( - Effect.ignore, - Effect.andThen(Scope.close(scope, Exit.void)), - ), - ), + Effect.onError(() => rollbackLaunch), ) const opencode = yield* OpenCodeSdk.make(instance.artifacts).pipe( - Effect.onError(() => - instance.killServer.pipe( - Effect.ignore, - Effect.andThen(Scope.close(scope, Exit.void)), - ), - ), + Effect.onError(() => rollbackLaunch), ) const process = yield* instance.primary.pipe( Effect.mapError((cause) => error("server.launch", cause)), - Effect.onError(() => - instance.killServer.pipe( - Effect.ignore, - Effect.andThen(Scope.close(scope, Exit.void)), - ), - ), + Effect.onError(() => rollbackLaunch), ) yield* Ref.set(generation, { scope, attachment, process }) compatibility = [...compatibility, backend.compatibility] @@ -216,10 +223,34 @@ export const make = Effect.fn("OpenCodeServer.make")(function* ( toolProducer.failure.pipe( Effect.mapError((cause) => error("tools", cause)), ), - Effect.raceFirst(llm.failure, Deferred.await(unexpectedExit)), + Effect.raceFirst( + Deferred.await(toolConnectionFailure), + Effect.raceFirst(llm.failure, Deferred.await(unexpectedExit)), + ), ), compatibility: Effect.sync(() => compatibility), } satisfies Server }) +function isRetryableToolConnectionError(cause: unknown) { + return ( + cause instanceof SimulationConnector.SimulationConnectionError || + (cause instanceof RpcClientError.RpcClientError && + isTransientRpcClientError(cause)) || + (cause instanceof LifecycleError && cause.reason === "transport-interrupted") + ) +} + +function isTransientRpcClientError(error: RpcClientError.RpcClientError) { + if (error.reason._tag !== "RpcClientDefect") return true + const message = error.reason.message + return ( + message.startsWith("cannot connect") || + message === "connection closed" || + message === "connection error" || + message === "connection is not open" || + message === "failed to send request" + ) +} + export * as OpenCodeServer from "./server.js" diff --git a/packages/drive/src/simulation/connector.ts b/packages/drive/src/simulation/connector.ts index 2c92da5..59f4015 100644 --- a/packages/drive/src/simulation/connector.ts +++ b/packages/drive/src/simulation/connector.ts @@ -87,6 +87,7 @@ export interface UiConnection { export type ToolEvent = | { readonly type: "invocation"; readonly invocation: BackendProtocol.ToolInvocation } | { readonly type: "cancellation"; readonly cancellation: BackendProtocol.ToolCancellation } + | { readonly type: "barrier"; readonly completed: Deferred.Deferred } export interface BackendConnection { readonly endpoint: string @@ -97,6 +98,7 @@ export interface BackendConnection { Schema.SchemaError > readonly toolEvents: Stream.Stream + readonly flushToolEvents: () => Effect.Effect readonly closed: Effect.Effect readonly attach: () => Effect.Effect< { readonly attached: true }, @@ -305,12 +307,18 @@ export const backend = Effect.fn("SimulationConnector.backend")(function* ( withTimeout("tool.finish", rpc["tool.finish"](params)) const failTool: BackendConnection["failTool"] = (params) => withTimeout("tool.fail", rpc["tool.fail"](params)) + const flushToolEvents = Effect.fn("SimulationConnector.flushToolEvents")(function* () { + const completed = yield* Deferred.make() + yield* Queue.offer(toolEvents, { type: "barrier", completed }) + yield* Deferred.await(completed) + }) return { endpoint, rpc, compatibility, requests: Stream.fromQueue(requests), toolEvents: Stream.fromQueue(toolEvents), + flushToolEvents, closed: Deferred.await(closed), attach, attachTools, diff --git a/packages/drive/src/tool/producer.ts b/packages/drive/src/tool/producer.ts index e23640f..8eb1580 100644 --- a/packages/drive/src/tool/producer.ts +++ b/packages/drive/src/tool/producer.ts @@ -58,16 +58,23 @@ type Active = { state: "pending" | "settled" | "cancelled" } +type AttachmentIntent = { + readonly params: AttachParams + attempted: boolean +} + const decodeAttach = Schema.decodeUnknownEffect(Backend.ToolAttachParams) const decodeProgress = Schema.decodeUnknownEffect(Backend.ToolProgress) const decodeOutput = Schema.decodeUnknownEffect(Backend.ToolOutput) const decodeFailure = Schema.decodeUnknownEffect(Schema.String) +const decodeCallID = Schema.decodeUnknownEffect(Schema.String) export const make = Effect.fn("ToolProducer.make")(function* ( staticNames: ReadonlySet, ) { const parentScope = yield* Scope.Scope const lifecycle = yield* Semaphore.make(1) + const attachmentCalls = yield* Semaphore.make(1) const attachmentLock = yield* Semaphore.make(1) const records = new Map() const waiters: Waiter[] = [] @@ -78,7 +85,7 @@ export const make = Effect.fn("ToolProducer.make")(function* ( | { readonly backend: BackendConnection; readonly scope: Scope.Scope } | undefined >(undefined) - let desired: AttachParams | undefined + let desired: AttachmentIntent | undefined let closed = false const lifecycleError = ( @@ -193,8 +200,14 @@ export const make = Effect.fn("ToolProducer.make")(function* ( if (Exit.isSuccess(result)) return result.value if (Cause.hasInterrupts(result.cause)) return yield* Effect.interrupt const found = Cause.findErrorOption(result.cause) - if (found._tag === "None") - return yield* Effect.die(Cause.squash(result.cause)) + if (found._tag === "None") { + const defect = Cause.squash(result.cause) + if (Schema.isSchemaError(defect)) + return yield* Effect.fail( + lifecycleError(operation, "rejected", defect.message, callID), + ) + return yield* Effect.die(defect) + } if ( found.value instanceof SimulationRequestError || found.value instanceof SimulationCompatibilityError @@ -204,7 +217,10 @@ export const make = Effect.fn("ToolProducer.make")(function* ( ) if ( !(found.value instanceof SimulationConnectionError) && - !(found.value instanceof RpcClientError.RpcClientError) + !( + found.value instanceof RpcClientError.RpcClientError && + isTransientRpcError(found.value) + ) ) return yield* Effect.fail( lifecycleError(operation, "rejected", String(found.value), callID), @@ -433,11 +449,15 @@ export const make = Effect.fn("ToolProducer.make")(function* ( ) const scope = yield* Scope.fork(parentScope) yield* backend.toolEvents.pipe( - Stream.runForEach((event) => - event.type === "invocation" - ? receiveInvocation(event.invocation) - : receiveCancellation(event.cancellation), - ), + Stream.runForEach((event) => { + if (event.type === "invocation") + return receiveInvocation(event.invocation) + if (event.type === "cancellation") + return receiveCancellation(event.cancellation) + return Deferred.succeed(event.completed, undefined).pipe( + Effect.asVoid, + ) + }), Effect.matchCauseEffect({ onFailure: (cause) => Deferred.fail( @@ -454,23 +474,35 @@ export const make = Effect.fn("ToolProducer.make")(function* ( ) const applied = yield* Effect.exit( attachmentLock.withPermit( - Effect.suspend(() => - desired === undefined - ? Effect.void - : whileConnected( - backend, + Effect.suspend(() => { + const intent = desired + if (intent === undefined) return Effect.void + intent.attempted = true + return whileConnected( + backend, + "attach", + backend.attachTools(intent.params.tools), + ).pipe( + Effect.mapError((error) => + lifecycleError( "attach", - backend.attachTools(desired.tools), - ).pipe( - Effect.mapError((error) => - lifecycleError("attach", "rejected", String(error)), - ), + isTransientConnectionError(error) + ? "transport-interrupted" + : "rejected", + String(error), ), - ), + ), + ) + }), ), ) if (Exit.isFailure(applied)) { yield* Scope.close(scope, Exit.void) + const defect = Cause.squash(applied.cause) + if (Schema.isSchemaError(defect)) + return yield* Effect.fail( + lifecycleError("attach", "rejected", defect.message), + ) return yield* Effect.failCause(applied.cause) } yield* Ref.set(current, { backend, scope }) @@ -509,16 +541,39 @@ export const make = Effect.fn("ToolProducer.make")(function* ( `dynamic tool conflicts with configured static adapter: ${collision}`, ), ) - yield* request("attach", undefined, (backend) => - attachmentLock.withPermit( - backend.attachTools(decoded.tools).pipe( - Effect.andThen( - Effect.sync(() => { - desired = decoded + yield* attachmentCalls.withPermit( + Effect.gen(function* () { + const previous = desired + const intent: AttachmentIntent = { params: decoded, attempted: false } + desired = intent + const rollback = Effect.sync(() => { + if (desired === intent) desired = previous + }) + yield* request("attach", undefined, (backend) => + attachmentLock.withPermit( + Effect.suspend(() => { + if (desired !== intent) + return Effect.succeed({ attached: true as const }) + intent.attempted = true + return backend + .attachTools(decoded.tools) + .pipe( + Effect.tapError((error) => + error instanceof SimulationRequestError || + error instanceof SimulationCompatibilityError + ? rollback + : Effect.void, + ), + ) }), ), - ), - ), + ).pipe( + Effect.tapError(() => (intent.attempted ? Effect.void : rollback)), + Effect.onInterrupt(() => + intent.attempted ? Effect.void : rollback, + ), + ) + }), ) return undefined }) @@ -526,69 +581,86 @@ export const make = Effect.fn("ToolProducer.make")(function* ( function take(): Effect.Effect function take(callID: string): Effect.Effect function take(callID?: string): Effect.Effect { - return Effect.callback((resume) => { - if (closed) { - resume( - Effect.fail( - lifecycleError( - "take", - "controller-closed", - "dynamic tool controller is closed", - callID, - ), - ), - ) - return undefined - } - if ( - callID !== undefined && - (Array.from(records.values()).some( - (record) => record.call.context.callID === callID && record.claimed, - ) || - waiters.some((waiter) => waiter.callID === callID)) - ) { - resume( - Effect.fail( - lifecycleError( - "take", - "already-claimed", - `dynamic tool call is already claimed: ${callID}`, - callID, + const decoded: Effect.Effect = + callID === undefined + ? Effect.succeed(undefined) + : decodeCallID(callID).pipe( + Effect.mapError((error) => + lifecycleError("take", "rejected", error.message), ), - ), - ) - return undefined - } - let record = Array.from(records.values()).find( - (candidate) => - !candidate.claimed && - (callID === undefined || candidate.call.context.callID === callID), - ) - if (record !== undefined) { - record.claimed = true - resume(Effect.succeed(record.call)) - if (record.state === "cancelled") records.delete(record.call.id) - return Effect.sync(() => { - if (record !== undefined && records.get(record.call.id) === record) { - record.claimed = false - deliver(record) + ) + return decoded.pipe( + Effect.flatMap((callID) => + Effect.callback((resume) => { + if (closed) { + resume( + Effect.fail( + lifecycleError( + "take", + "controller-closed", + "dynamic tool controller is closed", + callID, + ), + ), + ) + return undefined } - }) - } - const waiter: Waiter = { callID, resume } - waiters.push(waiter) - return Effect.sync(() => { - const index = waiters.indexOf(waiter) - if (index >= 0) waiters.splice(index, 1) - else if (waiter.delivered !== undefined) { - record = waiter.delivered - if (records.get(record.call.id) === record) { - record.claimed = false - deliver(record) + if ( + callID !== undefined && + (Array.from(records.values()).some( + (record) => + record.call.context.callID === callID && record.claimed, + ) || + waiters.some((waiter) => waiter.callID === callID)) + ) { + resume( + Effect.fail( + lifecycleError( + "take", + "already-claimed", + `dynamic tool call is already claimed: ${callID}`, + callID, + ), + ), + ) + return undefined } - } - }) - }) + let record = Array.from(records.values()).find( + (candidate) => + !candidate.claimed && + (callID === undefined || + candidate.call.context.callID === callID), + ) + if (record !== undefined) { + record.claimed = true + resume(Effect.succeed(record.call)) + if (record.state === "cancelled") records.delete(record.call.id) + return Effect.sync(() => { + if ( + record !== undefined && + records.get(record.call.id) === record + ) { + record.claimed = false + deliver(record) + } + }) + } + const waiter: Waiter = { callID, resume } + waiters.push(waiter) + return Effect.sync(() => { + const index = waiters.indexOf(waiter) + if (index >= 0) waiters.splice(index, 1) + else if (waiter.delivered !== undefined) { + record = waiter.delivered + if (records.get(record.call.id) === record) { + record.claimed = false + deliver(record) + } + } + }) + }), + ), + ) } const endGeneration = lifecycle.withPermit( @@ -612,22 +684,28 @@ export const make = Effect.fn("ToolProducer.make")(function* ( }), ) - const settle = lifecycle.withPermit( - Effect.suspend(() => { - const pending = Array.from(records.values()).filter( - (record) => record.state === "pending", - ).length - return pending === 0 - ? Effect.void - : Effect.fail( - lifecycleError( - "take", - "rejected", - `${pending} dynamic tool invocation(s) remain unsettled`, - ), - ) - }), - ) + const settle = Effect.gen(function* () { + if (desired !== undefined) { + if (desired.params.tools.length > 0) yield* attach({ tools: [] }) + yield* request("take", undefined, (backend) => backend.flushToolEvents()) + } + yield* lifecycle.withPermit( + Effect.suspend(() => { + const pending = Array.from(records.values()).filter( + (record) => record.state === "pending", + ).length + return pending === 0 + ? Effect.void + : Effect.fail( + lifecycleError( + "take", + "rejected", + `${pending} dynamic tool invocation(s) remain unsettled`, + ), + ) + }), + ) + }) const shutdown = Effect.suspend(() => { if (closed) return Effect.void @@ -678,4 +756,24 @@ function fingerprintJson(value: unknown) { return createHash("sha256").update(canonicalJson(value)).digest("base64url") } +function isTransientConnectionError(error: unknown) { + return ( + error instanceof SimulationConnectionError || + (error instanceof RpcClientError.RpcClientError && + isTransientRpcError(error)) + ) +} + +function isTransientRpcError(error: RpcClientError.RpcClientError) { + if (error.reason._tag !== "RpcClientDefect") return true + const message = error.reason.message + return ( + message.startsWith("cannot connect") || + message === "connection closed" || + message === "connection error" || + message === "connection is not open" || + message === "failed to send request" + ) +} + export * as ToolProducer from "./producer.js" diff --git a/packages/drive/src/tool/types.ts b/packages/drive/src/tool/types.ts index feea129..87ea299 100644 --- a/packages/drive/src/tool/types.ts +++ b/packages/drive/src/tool/types.ts @@ -169,6 +169,7 @@ export class LifecycleError extends Schema.TaggedErrorClass()( "already-settled", "cancelled", "rejected", + "transport-interrupted", ]), callID: Schema.optional(Schema.String), message: Schema.String, diff --git a/packages/drive/test/driver/index.test.ts b/packages/drive/test/driver/index.test.ts index ed5f261..e352e3b 100644 --- a/packages/drive/test/driver/index.test.ts +++ b/packages/drive/test/driver/index.test.ts @@ -196,6 +196,25 @@ it.live("supports explicit terminal settlement with make", () => }), ) +it.live("rolls back a failed initial tool connection", () => + Effect.gen(function* () { + const failure = yield* OpenCodeDriver.make({ + opencode: { command: [...fakeOpenCode, "reject-tool-handshake"] }, + }).pipe( + Effect.scoped, + Effect.timeoutOrElse({ + duration: 10_000, + orElse: () => Effect.dieMessage("failed tool connection did not roll back"), + }), + Effect.flip, + ) + expect(failure).toMatchObject({ + _tag: "OpenCodeDriverError", + operation: "tools.connect", + }) + }), +) + it.live("injects declared tool handlers for library drivers", () => Effect.gen(function* () { const artifacts = yield* OpenCodeDriver.use( @@ -329,10 +348,72 @@ it.live("controls arbitrary provider-backed tools through the runtime lifecycle" }, }, }, + { + method: "tool.attach", + params: { tools: [] }, + }, ]) }), ) +it.live("surfaces a permanent tool reconnect failure", () => + Effect.gen(function* () { + let artifacts = "" + const result = yield* Effect.exit( + OpenCodeDriver.use( + { + keepArtifacts: true, + opencode: { + command: [ + ...fakeOpenCode, + "dynamic-tool", + "reconnect-tool", + "reject-tool-reconnect", + ], + }, + }, + (driver) => + Effect.gen(function* () { + artifacts = driver.artifacts + yield* driver.tools.attach({ + tools: [ + { + name: "lookup", + description: "Look up a value", + inputSchema: { type: "object" }, + options: { codemode: false }, + }, + ], + }) + return yield* Effect.never + }), + ).pipe( + Effect.timeoutOrElse({ + duration: 10_000, + orElse: () => Effect.dieMessage("permanent tool reconnect failure was not observed"), + }), + ), + ) + yield* Effect.addFinalizer(() => + artifacts === "" + ? Effect.void + : Effect.promise(() => rm(artifacts, { recursive: true, force: true })), + ) + if (Exit.isSuccess(result)) + return yield* Effect.dieMessage("driver unexpectedly succeeded") + expect( + result.cause.reasons + .filter(Cause.isFailReason) + .map((reason) => reason.error), + ).toContainEqual( + expect.objectContaining({ + _tag: "OpenCodeDriverError", + operation: "tools.connect", + }), + ) + }), +) + it.live("returns structured evidence from the safe lifecycle boundary", () => Effect.gen(function* () { const result = yield* OpenCodeDriver.useReport( diff --git a/packages/drive/test/fixtures/fake-opencode.ts b/packages/drive/test/fixtures/fake-opencode.ts index 5c9c7b7..295b4f3 100644 --- a/packages/drive/test/fixtures/fake-opencode.ts +++ b/packages/drive/test/fixtures/fake-opencode.ts @@ -23,6 +23,7 @@ const drive = await resolveDrive() const recordingStarted = performance.now() const endpoints = drive.endpoints let toolAttachments = 0 +let backendHandshakes = 0 const servicePassword = "drive-test-password" const api = role === "service" ? Bun.serve({ @@ -126,6 +127,21 @@ const backend = role === "client" ? undefined : Bun.serve({ readonly method: string readonly params?: unknown } + if ( + request.method === "simulation.handshake" && + process.argv.includes("reject-tool-handshake") && + ++backendHandshakes === 2 + ) { + if (request.id !== undefined) + socket.send( + JSON.stringify({ + jsonrpc: "2.0", + id: request.id, + error: { code: -32000, message: "tool handshake rejected" }, + }), + ) + return + } if (request.method === "llm.chunk" && process.env.OPENCODE_TEST_HOME) { await Bun.write( `${process.env.OPENCODE_TEST_HOME}/mock-response.json`, @@ -172,6 +188,21 @@ const backend = role === "client" ? undefined : Bun.serve({ }, }), ) + if ( + request.method === "tool.attach" && + process.argv.includes("reject-tool-reconnect") && + toolAttachments === 2 + ) { + if (request.id !== undefined) + socket.send( + JSON.stringify({ + jsonrpc: "2.0", + id: request.id, + error: { code: -32000, message: "tool reconnect rejected" }, + }), + ) + return + } const result = request.method === "simulation.handshake" ? handshake(request.params, "backend") diff --git a/packages/drive/test/tool/producer.test.ts b/packages/drive/test/tool/producer.test.ts index f22165f..e2fdce9 100644 --- a/packages/drive/test/tool/producer.test.ts +++ b/packages/drive/test/tool/producer.test.ts @@ -3,7 +3,7 @@ import { Deferred, Effect, Exit, Fiber, Scope } from "effect" import * as SimulationConnector from "../../src/simulation/connector.js" import * as ToolProducer from "../../src/tool/producer.js" import type { Progress } from "../../src/tool/types.js" -import { sendResult, startTransportPeer } from "../simulation/transport-peer.js" +import { sendError, sendResult, startTransportPeer } from "../simulation/transport-peer.js" const registration = { name: "lookup", @@ -261,7 +261,7 @@ it.live("reattaches the desired set and deduplicates replay after reconnect", () ), ) -it.live("replays the current set before applying an attachment started while disconnected", () => +it.live("replays a replacement intent started while disconnected", () => Effect.scoped( Effect.gen(function* () { const peer = startTransportPeer(({ request, socket }) => @@ -300,11 +300,203 @@ it.live("replays the current set before applying an attachment started while dis .map(({ request }) => request.params), ).toEqual([ { tools: [registration] }, + { tools: [replacement] }, + { tools: [replacement] }, + ]) + yield* secondAttachment.detach() + yield* Scope.close(secondScope, Exit.void) + }), + ), +) + +it.live("preserves a replacement intent when its acknowledgement is lost", () => + Effect.scoped( + Effect.gen(function* () { + let attaches = 0 + const replacement = { + ...registration, + name: "search", + description: "Search for a value", + } + const replacementInvocation = { + ...invocation, + id: "tool_2", + name: "search", + context: { ...invocation.context, callID: "call_search" }, + } + const firstReplacement = yield* Deferred.make() + const peer = startTransportPeer(({ request, socket }) => { + if (request.method === "tool.attach") { + attaches++ + if (attaches === 2) { + notify(socket, "tool.invocation", replacementInvocation) + Deferred.doneUnsafe(firstReplacement, Effect.void) + socket.close() + return + } + if (attaches === 3) notify(socket, "tool.invocation", replacementInvocation) + sendResult(socket, request, { attached: true }) + return + } + sendResult(socket, request, { ok: true }) + }) + yield* Effect.addFinalizer(() => + Effect.sync(() => { + void peer.stop() + }), + ) + const controller = yield* ToolProducer.make(new Set()) + const firstScope = yield* Scope.make() + const first = yield* SimulationConnector.backend(peer.url, { + attach: false, + }).pipe(Scope.provide(firstScope)) + const firstAttachment = yield* controller.connect(first) + yield* controller.controls.attach({ tools: [registration] }) + + const replacing = yield* controller.controls + .attach({ tools: [replacement] }) + .pipe(Effect.forkScoped) + yield* Deferred.await(firstReplacement) + yield* first.closed + yield* firstAttachment.detach() + yield* Scope.close(firstScope, Exit.void) + const secondScope = yield* Scope.make() + const second = yield* SimulationConnector.backend(peer.url, { + attach: false, + }).pipe(Scope.provide(secondScope)) + const secondAttachment = yield* controller.connect(second) + yield* Fiber.join(replacing) + const call = yield* controller.controls.take("call_search") + yield* call.fail("finished") + + expect( + peer.received + .filter(({ request }) => request.method === "tool.attach") + .map(({ request }) => request.params), + ).toEqual([ { tools: [registration] }, { tools: [replacement] }, + { tools: [replacement] }, + { tools: [replacement] }, ]) yield* secondAttachment.detach() yield* Scope.close(secondScope, Exit.void) + yield* controller.endGeneration + yield* controller.shutdown + }), + ), +) + +it.live("restores the acknowledged set after a rejected replacement", () => + Effect.scoped( + Effect.gen(function* () { + let attaches = 0 + const replacement = { + ...registration, + name: "search", + description: "Search for a value", + } + const peer = startTransportPeer(({ request, socket }) => { + if (request.method === "tool.attach" && ++attaches === 2) { + sendError(socket, request, "replacement rejected") + return + } + sendResult(socket, request, { attached: true }) + }) + yield* Effect.addFinalizer(() => Effect.promise(() => peer.stop())) + const controller = yield* ToolProducer.make(new Set()) + const firstScope = yield* Scope.make() + const first = yield* SimulationConnector.backend(peer.url, { + attach: false, + }).pipe(Scope.provide(firstScope)) + const firstAttachment = yield* controller.connect(first) + yield* controller.controls.attach({ tools: [registration] }) + expect( + yield* controller.controls.attach({ tools: [replacement] }).pipe(Effect.flip), + ).toMatchObject({ operation: "attach", reason: "rejected" }) + yield* firstAttachment.detach() + yield* Scope.close(firstScope, Exit.void) + const secondScope = yield* Scope.make() + const second = yield* SimulationConnector.backend(peer.url, { + attach: false, + }).pipe(Scope.provide(secondScope)) + const secondAttachment = yield* controller.connect(second) + + expect( + peer.received + .filter(({ request }) => request.method === "tool.attach") + .map(({ request }) => request.params), + ).toEqual([ + { tools: [registration] }, + { tools: [replacement] }, + { tools: [registration] }, + ]) + yield* secondAttachment.detach() + yield* Scope.close(secondScope, Exit.void) + }), + ), +) + +it.live("rejects malformed lifecycle acknowledgements without retrying", () => + Effect.scoped( + Effect.gen(function* () { + const peer = startTransportPeer(({ request, socket }) => + sendResult(socket, request, { attached: "invalid" }), + ) + yield* Effect.addFinalizer(() => Effect.promise(() => peer.stop())) + const backend = yield* SimulationConnector.backend(peer.url, { attach: false }) + const controller = yield* ToolProducer.make(new Set()) + yield* controller.connect(backend) + + expect( + yield* controller.controls.attach({ tools: [registration] }).pipe(Effect.flip), + ).toMatchObject({ operation: "attach", reason: "rejected" }) + expect(peer.received.filter(({ request }) => request.method === "tool.attach")).toHaveLength(1) + }), + ), +) + +it.live("drains queued invocations before reporting settlement", () => + Effect.scoped( + Effect.gen(function* () { + let attaches = 0 + const peer = startTransportPeer(({ request, socket }) => { + if (request.method === "tool.attach") { + attaches++ + if (attaches === 2) notify(socket, "tool.invocation", invocation) + sendResult(socket, request, { attached: true }) + return + } + sendResult(socket, request, { ok: true }) + }) + yield* Effect.addFinalizer(() => Effect.promise(() => peer.stop())) + const backend = yield* SimulationConnector.backend(peer.url, { attach: false }) + const controller = yield* ToolProducer.make(new Set()) + yield* controller.connect(backend) + yield* controller.controls.attach({ tools: [registration] }) + + expect(yield* controller.settle.pipe(Effect.flip)).toMatchObject({ + operation: "take", + reason: "rejected", + message: expect.stringContaining("1 dynamic tool invocation"), + }) + const call = yield* controller.controls.take("call_lookup") + yield* call.fail("finished") + yield* controller.settle + expect( + peer.received.filter(({ request }) => request.method === "tool.attach").map(({ request }) => request.params), + ).toEqual([{ tools: [registration] }, { tools: [] }]) + }), + ), +) + +it.effect("rejects a malformed take call ID", () => + Effect.scoped( + Effect.gen(function* () { + const controller = yield* ToolProducer.make(new Set()) + expect( + yield* controller.controls.take(null as unknown as string).pipe(Effect.flip), + ).toMatchObject({ operation: "take", reason: "rejected" }) }), ), ) From 96f5411465224f240f782309f279420eb146961e Mon Sep 17 00:00:00 2001 From: Kit Langton Date: Sun, 19 Jul 2026 23:33:52 -0400 Subject: [PATCH 3/3] fix(drive): close dynamic tool teardown races --- .../docs/open-code-driver-architecture.md | 7 +- packages/drive/src/driver/server.ts | 50 +- packages/drive/src/simulation/connector.ts | 20 +- packages/drive/src/tool/producer.ts | 389 +++++-- packages/drive/test/cli/integration.test.ts | 26 + packages/drive/test/driver/index.test.ts | 10 + .../fixtures/failed-launch-retry-script.ts | 52 + packages/drive/test/fixtures/fake-opencode.ts | 6 +- .../drive/test/fixtures/kill-server-script.ts | 12 +- packages/drive/test/tool/producer.test.ts | 967 ++++++++++++++++-- 10 files changed, 1330 insertions(+), 209 deletions(-) create mode 100644 packages/drive/test/fixtures/failed-launch-retry-script.ts diff --git a/packages/drive/docs/open-code-driver-architecture.md b/packages/drive/docs/open-code-driver-architecture.md index 279bab9..06a6f93 100644 --- a/packages/drive/docs/open-code-driver-architecture.md +++ b/packages/drive/docs/open-code-driver-architecture.md @@ -115,7 +115,12 @@ generation because producer IDs may be reused by a new process. Settlement first clears the dynamic registration set on OpenCode, then drains the ordered local event stream before checking for unresolved invocations. The clear acts as the server-side barrier that prevents a native invocation from -appearing after a successful settlement snapshot. +appearing after a successful settlement snapshot. Settlement is terminal for +dynamic attachment. Reconnects remain available while the clear is in flight; +the final connection gate drains any reconnect that landed during settlement +before preventing further backend creation. If the server generation has +already ended, its teardown has cleared the generation-scoped invocation +records, so settlement does not wait for a replacement backend. ## TUI Lifecycle diff --git a/packages/drive/src/driver/server.ts b/packages/drive/src/driver/server.ts index 3d92b3e..3b1e287 100644 --- a/packages/drive/src/driver/server.ts +++ b/packages/drive/src/driver/server.ts @@ -90,9 +90,13 @@ export const make = Effect.fn("OpenCodeServer.make")(function* ( Effect.mapError((cause) => error("server.launch", cause)), Effect.onError(() => Scope.close(scope, Exit.void)), ) - const rollbackLaunch = Scope.close(scope, Exit.void).pipe( - Effect.andThen(toolProducer.endGeneration), - Effect.andThen(instance.killServer.pipe(Effect.ignore)), + let llmAttachment: LlmController.Attachment | undefined + const rollbackLaunch = Effect.suspend(() => + (llmAttachment?.detach() ?? Effect.void).pipe( + Effect.andThen(toolProducer.endGeneration), + Effect.andThen(Scope.close(scope, Exit.void)), + Effect.andThen(instance.killServer.pipe(Effect.ignore)), + ), ) const backend = yield* connector.backend(launched.endpoint, { compatibility: target.compatibility, @@ -103,17 +107,15 @@ export const make = Effect.fn("OpenCodeServer.make")(function* ( ) const connectTools = Effect.fn("OpenCodeServer.connectTools")(function* () { const connectionScope = yield* Scope.fork(scope) - const backend = yield* connector.backend(launched.endpoint, { - attach: false, - compatibility: target.compatibility, - }).pipe( - Scope.provide(connectionScope), - Effect.onError(() => Scope.close(connectionScope, Exit.void)), - ) - const attachment = yield* toolProducer.connect(backend).pipe( - Effect.onError(() => Scope.close(connectionScope, Exit.void)), - ) - return { attachment, backend, scope: connectionScope } + const connection = yield* toolProducer + .connectFrom( + connector.backend(launched.endpoint, { + attach: false, + compatibility: target.compatibility, + }).pipe(Scope.provide(connectionScope)), + ) + .pipe(Effect.onError(() => Scope.close(connectionScope, Exit.void))) + return { ...connection, scope: connectionScope } }) const failToolConnection = (cause: unknown) => toolProducer.shutdown.pipe( @@ -132,6 +134,7 @@ export const make = Effect.fn("OpenCodeServer.make")(function* ( isRetryableToolConnectionError, () => Effect.sleep(25).pipe(Effect.andThen(reconnectTools())), ), + Effect.catchIf(isClosedToolConnectionError, () => Effect.void), Effect.catch(failToolConnection), Effect.catchCauseIf( (cause) => !Cause.hasInterrupts(cause), @@ -162,6 +165,7 @@ export const make = Effect.fn("OpenCodeServer.make")(function* ( const attachment = yield* llm.attach(backend).pipe( Effect.onError(() => rollbackLaunch), ) + llmAttachment = attachment const opencode = yield* OpenCodeSdk.make(instance.artifacts).pipe( Effect.onError(() => rollbackLaunch), ) @@ -177,10 +181,12 @@ export const make = Effect.fn("OpenCodeServer.make")(function* ( Ref.get(generation).pipe( Effect.flatMap((active) => active?.process === process - ? Deferred.fail( - unexpectedExit, - error("server.exit", `OpenCode server exited with status ${status}`), - ).pipe(Effect.asVoid) + ? toolProducer.endGeneration.pipe( + Effect.andThen( + Deferred.fail(unexpectedExit, error("server.exit", `OpenCode server exited with status ${status}`)), + ), + Effect.asVoid, + ) : Effect.void, ), ), @@ -199,8 +205,8 @@ export const make = Effect.fn("OpenCodeServer.make")(function* ( ) yield* Ref.set(generation, undefined) yield* active.attachment.detach() - yield* Scope.close(active.scope, Exit.void) yield* toolProducer.endGeneration + yield* Scope.close(active.scope, Exit.void) const stopped = yield* Effect.exit( instance.killServer.pipe( Effect.mapError((cause) => error("server.kill", cause)), @@ -227,7 +233,7 @@ export const make = Effect.fn("OpenCodeServer.make")(function* ( Deferred.await(toolConnectionFailure), Effect.raceFirst(llm.failure, Deferred.await(unexpectedExit)), ), - ), + ).pipe(Effect.tapError(() => toolProducer.endGeneration)), compatibility: Effect.sync(() => compatibility), } satisfies Server }) @@ -241,6 +247,10 @@ function isRetryableToolConnectionError(cause: unknown) { ) } +function isClosedToolConnectionError(cause: unknown) { + return cause instanceof LifecycleError && cause.reason === "controller-closed" +} + function isTransientRpcClientError(error: RpcClientError.RpcClientError) { if (error.reason._tag !== "RpcClientDefect") return true const message = error.reason.message diff --git a/packages/drive/src/simulation/connector.ts b/packages/drive/src/simulation/connector.ts index 59f4015..af02b6d 100644 --- a/packages/drive/src/simulation/connector.ts +++ b/packages/drive/src/simulation/connector.ts @@ -35,6 +35,14 @@ export class SimulationCompatibilityError extends Schema.TaggedErrorClass()( + "SimulationEventStreamError", + { + endpoint: Schema.String, + message: Schema.String, + }, +) {} + export const EndpointCompatibility = Schema.TaggedUnion({ Negotiated: { endpoint: Schema.String, @@ -98,7 +106,7 @@ export interface BackendConnection { Schema.SchemaError > readonly toolEvents: Stream.Stream - readonly flushToolEvents: () => Effect.Effect + readonly flushToolEvents: () => Effect.Effect readonly closed: Effect.Effect readonly attach: () => Effect.Effect< { readonly attached: true }, @@ -309,8 +317,16 @@ export const backend = Effect.fn("SimulationConnector.backend")(function* ( withTimeout("tool.fail", rpc["tool.fail"](params)) const flushToolEvents = Effect.fn("SimulationConnector.flushToolEvents")(function* () { const completed = yield* Deferred.make() - yield* Queue.offer(toolEvents, { type: "barrier", completed }) + const offered = yield* Queue.offer(toolEvents, { type: "barrier", completed }) + if (!offered) + return yield* Effect.fail( + new SimulationEventStreamError({ + endpoint, + message: "Dynamic tool event stream is unavailable", + }), + ) yield* Deferred.await(completed) + return undefined }) return { endpoint, diff --git a/packages/drive/src/tool/producer.ts b/packages/drive/src/tool/producer.ts index 8eb1580..968b61f 100644 --- a/packages/drive/src/tool/producer.ts +++ b/packages/drive/src/tool/producer.ts @@ -33,9 +33,14 @@ export interface BackendAttachment { export interface Controller { readonly controls: DynamicControls - readonly connect: ( - backend: BackendConnection, - ) => Effect.Effect + readonly connect: (backend: BackendConnection) => Effect.Effect + readonly connectFrom: ( + backend: Effect.Effect, + ) => Effect.Effect< + { readonly backend: BackendConnection; readonly attachment: BackendAttachment }, + E | LifecycleError, + R + > readonly endGeneration: Effect.Effect readonly settle: Effect.Effect readonly shutdown: Effect.Effect @@ -60,6 +65,8 @@ type Active = { type AttachmentIntent = { readonly params: AttachParams + previous: AttachmentIntent | undefined + readonly rejection: Deferred.Deferred attempted: boolean } @@ -76,16 +83,29 @@ export const make = Effect.fn("ToolProducer.make")(function* ( const lifecycle = yield* Semaphore.make(1) const attachmentCalls = yield* Semaphore.make(1) const attachmentLock = yield* Semaphore.make(1) + const connectionCalls = yield* Semaphore.make(1) const records = new Map() const waiters: Waiter[] = [] const completed = new Map() + const unclaimedCancellations: Active[] = [] const backendChanges = yield* Queue.sliding(1) const failure = yield* Deferred.make() + const settlementStarted = yield* Deferred.make() const current = yield* Ref.make< - | { readonly backend: BackendConnection; readonly scope: Scope.Scope } + | { + readonly backend: BackendConnection + readonly scope: Scope.Scope + readonly disconnected: Deferred.Deferred + } | undefined >(undefined) let desired: AttachmentIntent | undefined + let acknowledged: AttachmentIntent | undefined + let terminalFailure: LifecycleError | undefined + let generationActive = false + let generationEnded = Deferred.makeUnsafe() + let settling = false + let settled = false let closed = false const lifecycleError = ( @@ -101,9 +121,21 @@ export const make = Effect.fn("ToolProducer.make")(function* ( ...(callID === undefined ? {} : { callID }), }) - const notifyBackend = Queue.offer(backendChanges, undefined).pipe( - Effect.asVoid, - ) + const rejectIntent = (intent: AttachmentIntent, error: LifecycleError) => + Effect.sync(() => { + if (desired === intent) desired = intent.previous + if (acknowledged === intent) acknowledged = intent.previous + Deferred.doneUnsafe(intent.rejection, Effect.fail(error)) + }) + + const publishFailure = (error: LifecycleError) => + Effect.sync(() => { + if (terminalFailure !== undefined) return + terminalFailure = error + Deferred.doneUnsafe(failure, Effect.fail(error)) + }) + + const notifyBackend = Queue.offer(backendChanges, undefined).pipe(Effect.asVoid) const remember = (id: string, fingerprint: string) => { completed.set(id, fingerprint) @@ -113,6 +145,19 @@ export const make = Effect.fn("ToolProducer.make")(function* ( } } + const retainCancellation = (record: Active) => { + unclaimedCancellations.push(record) + if (unclaimedCancellations.length <= 256) return + const oldest = unclaimedCancellations.shift() + if ( + oldest !== undefined && + records.get(oldest.call.id) === oldest && + oldest.state === "cancelled" && + !oldest.claimed + ) + records.delete(oldest.call.id) + } + const deliver = (record: Active) => { const callID = record.call.context.callID const exact = waiters.findIndex((waiter) => waiter.callID === callID) @@ -392,8 +437,7 @@ export const make = Effect.fn("ToolProducer.make")(function* ( const existing = records.get(invocation.id) if (existing !== undefined) { if (existing.fingerprint !== fingerprint) - yield* Deferred.fail( - failure, + yield* publishFailure( lifecycleError( "take", "rejected", @@ -406,8 +450,7 @@ export const make = Effect.fn("ToolProducer.make")(function* ( const settled = completed.get(invocation.id) if (settled !== undefined) { if (settled !== fingerprint) - yield* Deferred.fail( - failure, + yield* publishFailure( lifecycleError( "take", "rejected", @@ -432,13 +475,24 @@ export const make = Effect.fn("ToolProducer.make")(function* ( remember(cancellation.id, record.fingerprint) Deferred.doneUnsafe(record.cancelled, Effect.succeed(cancellation)) if (record.claimed) records.delete(cancellation.id) - else deliver(record) + else { + deliver(record) + if (!record.claimed) retainCancellation(record) + } }), ) - const connect: Controller["connect"] = (backend) => + const connectBackend: Controller["connect"] = (backend) => lifecycle.withPermit( Effect.gen(function* () { + if (closed) + return yield* Effect.fail( + lifecycleError( + "attach", + "controller-closed", + "dynamic tool controller is closed", + ), + ) if ((yield* Ref.get(current)) !== undefined) return yield* Effect.fail( lifecycleError( @@ -459,39 +513,46 @@ export const make = Effect.fn("ToolProducer.make")(function* ( ) }), Effect.matchCauseEffect({ - onFailure: (cause) => - Deferred.fail( - failure, - lifecycleError( - "take", - "rejected", + onFailure: (cause) => { + const error = lifecycleError( + "take", + "rejected", `dynamic tool event stream failed: ${Cause.pretty(cause)}`, - ), - ).pipe(Effect.asVoid), + ) + return publishFailure(error) + }, onSuccess: () => Effect.void, }), Effect.forkIn(scope), ) + let replayed: AttachmentIntent | undefined const applied = yield* Effect.exit( attachmentLock.withPermit( Effect.suspend(() => { const intent = desired if (intent === undefined) return Effect.void + replayed = intent intent.attempted = true + const reject = (error: unknown) => + lifecycleError( + "attach", + isTransientConnectionError(error) + ? "transport-interrupted" + : "rejected", + String(error), + ) return whileConnected( backend, "attach", backend.attachTools(intent.params.tools), ).pipe( - Effect.mapError((error) => - lifecycleError( - "attach", - isTransientConnectionError(error) - ? "transport-interrupted" - : "rejected", - String(error), - ), + Effect.tapError((error) => + error instanceof SimulationRequestError || + error instanceof SimulationCompatibilityError + ? rejectIntent(intent, reject(error)) + : Effect.void, ), + Effect.mapError(reject), ) }), ), @@ -499,13 +560,27 @@ export const make = Effect.fn("ToolProducer.make")(function* ( if (Exit.isFailure(applied)) { yield* Scope.close(scope, Exit.void) const defect = Cause.squash(applied.cause) - if (Schema.isSchemaError(defect)) - return yield* Effect.fail( - lifecycleError("attach", "rejected", defect.message), - ) + if (Schema.isSchemaError(defect)) { + const error = lifecycleError("attach", "rejected", defect.message) + if (replayed !== undefined) yield* rejectIntent(replayed, error) + return yield* Effect.fail(error) + } + const found = Cause.findErrorOption(applied.cause) + if ( + replayed !== undefined && + found._tag === "Some" && + found.value instanceof LifecycleError && + found.value.reason === "rejected" + ) + yield* rejectIntent(replayed, found.value) return yield* Effect.failCause(applied.cause) } - yield* Ref.set(current, { backend, scope }) + yield* Effect.sync(() => { + if (!generationActive) generationEnded = Deferred.makeUnsafe() + generationActive = true + }) + const disconnected = Deferred.makeUnsafe() + yield* Ref.set(current, { backend, scope, disconnected }) yield* notifyBackend const detach = Effect.fn("ToolProducer.detach")(function* () { const shouldClose = yield* lifecycle.withPermit( @@ -513,6 +588,7 @@ export const make = Effect.fn("ToolProducer.make")(function* ( const attached = yield* Ref.get(current) if (attached?.backend !== backend) return false yield* Ref.set(current, undefined) + Deferred.doneUnsafe(disconnected, Effect.void) yield* notifyBackend return true }), @@ -523,6 +599,86 @@ export const make = Effect.fn("ToolProducer.make")(function* ( }), ) + const connectionClosed = () => + lifecycleError( + "attach", + "controller-closed", + "dynamic tool controller is settled", + ) + + const connect: Controller["connect"] = (backend) => + connectionCalls.withPermit( + Effect.suspend(() => + settled ? Effect.fail(connectionClosed()) : connectBackend(backend), + ), + ) + + const connectFrom: Controller["connectFrom"] = (backend) => + connectionCalls.withPermit( + Effect.suspend(() => + settled + ? Effect.fail(connectionClosed()) + : backend.pipe( + Effect.flatMap((backend) => + connectBackend(backend).pipe( + Effect.map((attachment) => ({ backend, attachment })), + ), + ), + ), + ), + ) + + const replace = (params: AttachParams, duringSettlement = false) => + Effect.gen(function* () { + const previous = acknowledged + const intent: AttachmentIntent = { + params, + previous, + rejection: Deferred.makeUnsafe(), + attempted: false, + } + desired = intent + const rollback = Effect.sync(() => { + if (desired === intent) desired = previous + }) + const apply = Effect.raceFirst( + request("attach", undefined, (backend) => + attachmentLock.withPermit( + Effect.suspend(() => { + if (desired !== intent) return Effect.succeed({ attached: true as const }) + intent.attempted = true + return backend.attachTools(params.tools).pipe( + Effect.tapError((error) => { + if (!(error instanceof SimulationRequestError) && !(error instanceof SimulationCompatibilityError)) + return Effect.void + return rejectIntent(intent, lifecycleError("attach", "rejected", error.message)) + }), + ) + }), + ), + ).pipe( + Effect.tapError((error) => (error.reason === "rejected" ? rejectIntent(intent, error) : Effect.void)), + Effect.tapError(() => (intent.attempted ? Effect.void : rollback)), + Effect.onInterrupt(() => (intent.attempted ? Effect.void : rollback)), + ), + Deferred.await(intent.rejection), + ) + yield* duringSettlement + ? apply + : Effect.raceFirst( + apply, + Deferred.await(settlementStarted).pipe( + Effect.andThen( + Effect.fail(lifecycleError("attach", "controller-closed", "dynamic tool controller is settling")), + ), + ), + ) + if (desired === intent) { + intent.previous = undefined + acknowledged = intent + } + }) + const attach: DynamicControls["attach"] = (params) => Effect.gen(function* () { const decoded = yield* decodeAttach(params).pipe( @@ -542,38 +698,11 @@ export const make = Effect.fn("ToolProducer.make")(function* ( ), ) yield* attachmentCalls.withPermit( - Effect.gen(function* () { - const previous = desired - const intent: AttachmentIntent = { params: decoded, attempted: false } - desired = intent - const rollback = Effect.sync(() => { - if (desired === intent) desired = previous - }) - yield* request("attach", undefined, (backend) => - attachmentLock.withPermit( - Effect.suspend(() => { - if (desired !== intent) - return Effect.succeed({ attached: true as const }) - intent.attempted = true - return backend - .attachTools(decoded.tools) - .pipe( - Effect.tapError((error) => - error instanceof SimulationRequestError || - error instanceof SimulationCompatibilityError - ? rollback - : Effect.void, - ), - ) - }), - ), - ).pipe( - Effect.tapError(() => (intent.attempted ? Effect.void : rollback)), - Effect.onInterrupt(() => - intent.attempted ? Effect.void : rollback, - ), - ) - }), + Effect.suspend(() => + settling + ? Effect.fail(lifecycleError("attach", "controller-closed", "dynamic tool controller is settling")) + : replace(decoded), + ), ) return undefined }) @@ -592,13 +721,13 @@ export const make = Effect.fn("ToolProducer.make")(function* ( return decoded.pipe( Effect.flatMap((callID) => Effect.callback((resume) => { - if (closed) { + if (closed || settled) { resume( Effect.fail( lifecycleError( "take", "controller-closed", - "dynamic tool controller is closed", + settled ? "dynamic tool controller is settled" : "dynamic tool controller is closed", callID, ), ), @@ -665,9 +794,12 @@ export const make = Effect.fn("ToolProducer.make")(function* ( const endGeneration = lifecycle.withPermit( Effect.gen(function* () { + generationActive = false + Deferred.doneUnsafe(generationEnded, Effect.void) const attached = yield* Ref.get(current) if (attached !== undefined) { yield* Ref.set(current, undefined) + Deferred.doneUnsafe(attached.disconnected, Effect.void) yield* Scope.close(attached.scope, Exit.void) } for (const record of records.values()) { @@ -680,32 +812,112 @@ export const make = Effect.fn("ToolProducer.make")(function* ( } records.clear() completed.clear() + unclaimedCancellations.length = 0 yield* notifyBackend }), ) - const settle = Effect.gen(function* () { - if (desired !== undefined) { - if (desired.params.tools.length > 0) yield* attach({ tools: [] }) - yield* request("take", undefined, (backend) => backend.flushToolEvents()) - } - yield* lifecycle.withPermit( - Effect.suspend(() => { - const pending = Array.from(records.values()).filter( - (record) => record.state === "pending", - ).length - return pending === 0 - ? Effect.void - : Effect.fail( - lifecycleError( - "take", - "rejected", - `${pending} dynamic tool invocation(s) remain unsettled`, + function commitSettlement(): Effect.Effect { + return Effect.gen(function* () { + const retryAfter = yield* connectionCalls.withPermit( + Effect.gen(function* () { + const attached = yield* Ref.get(current) + if (generationActive && attached === undefined) + return generationEnded + if (generationActive && attached !== undefined) { + const drained = yield* Effect.raceFirst( + attached.backend.flushToolEvents().pipe( + Effect.mapError((error) => + lifecycleError("take", "rejected", error.message), + ), + Effect.as(true), ), + Deferred.await(attached.disconnected).pipe(Effect.as(false)), ) - }), - ) - }) + if (!drained) return generationEnded + } + if (terminalFailure !== undefined) + return yield* Effect.fail(terminalFailure) + yield* lifecycle.withPermit( + Effect.suspend(() => { + if (terminalFailure !== undefined) + return Effect.fail(terminalFailure) + const pending = Array.from(records.values()).filter( + (record) => record.state === "pending", + ).length + if (pending > 0) + return Effect.fail( + lifecycleError( + "take", + "rejected", + `${pending} dynamic tool invocation(s) remain unsettled`, + ), + ) + return Effect.sync(() => { + settled = true + for (const waiter of waiters) + waiter.resume( + Effect.fail( + lifecycleError( + "take", + "controller-closed", + "dynamic tool controller is settled", + waiter.callID, + ), + ), + ) + waiters.length = 0 + }) + }), + ) + return undefined + }), + ) + if (retryAfter === undefined) return + yield* Effect.raceFirst( + Deferred.await(retryAfter), + Effect.raceFirst( + awaitBackend("take", undefined).pipe(Effect.asVoid), + Deferred.await(failure), + ), + ) + yield* commitSettlement() + }) + } + + const settle = Effect.sync(() => { + settling = true + Deferred.doneUnsafe(settlementStarted, Effect.void) + }).pipe( + Effect.andThen( + attachmentCalls.withPermit( + Effect.gen(function* () { + yield* connectionCalls.withPermit(Effect.void) + if (terminalFailure !== undefined) + return yield* Effect.fail(terminalFailure) + if (generationActive && desired !== undefined) { + const ended = generationEnded + if (desired.params.tools.length > 0) + yield* Effect.raceFirst( + replace({ tools: [] }, true), + Deferred.await(ended), + ) + if (generationActive) + yield* Effect.raceFirst( + request("take", undefined, (backend) => + backend.flushToolEvents(), + ), + Deferred.await(ended), + ) + } + if (terminalFailure !== undefined) + return yield* Effect.fail(terminalFailure) + yield* commitSettlement() + return undefined + }), + ), + ), + ) const shutdown = Effect.suspend(() => { if (closed) return Effect.void @@ -734,6 +946,7 @@ export const make = Effect.fn("ToolProducer.make")(function* ( return { controls: { attach, take }, connect, + connectFrom, endGeneration, settle, shutdown, diff --git a/packages/drive/test/cli/integration.test.ts b/packages/drive/test/cli/integration.test.ts index 14ce542..03c4987 100644 --- a/packages/drive/test/cli/integration.test.ts +++ b/packages/drive/test/cli/integration.test.ts @@ -1013,6 +1013,32 @@ describe("opencode-drive", () => { expect((await Bun.file(join(artifacts, "launches.txt")).text()).trim().split("\n")).toHaveLength(3) }, 60_000) + test("retries a launch that fails after LLM attachment", async () => { + const root = await temporary() + const child = spawn( + [ + "start", + "--name", + "failed-launch-retry-test", + "--script", + fixture("failed-launch-retry-script.ts"), + "--", + process.execPath, + fixture("fake-opencode.ts"), + "omit-service-registration", + ], + root, + ) + const stderr = new Response(child.stderr).text() + expect(await child.exited).toBe(0) + const artifacts = artifactPath(await stderr) + roots.push(artifacts) + const result = await Bun.file(join(artifacts, "failed-launch-retry.json")).json() + expect(Number.isInteger(result.firstPid)).toBe(true) + expect(Number.isInteger(result.secondPid)).toBe(true) + expect(result.secondPid).not.toBe(result.firstPid) + }, 60_000) + test("rejects the removed callback script shape", async () => { const root = await temporary() const child = spawn(["start", "--name", "callback-script-test", "--script", fixture("callback-script.ts")], root) diff --git a/packages/drive/test/driver/index.test.ts b/packages/drive/test/driver/index.test.ts index e352e3b..624b6f9 100644 --- a/packages/drive/test/driver/index.test.ts +++ b/packages/drive/test/driver/index.test.ts @@ -544,6 +544,16 @@ it.live("interrupts use when the backend disconnects", () => (driver) => Effect.gen(function* () { artifacts = driver.artifacts + yield* driver.tools.attach({ + tools: [ + { + name: "lookup", + description: "Look up a value", + inputSchema: { type: "object" }, + options: { codemode: false }, + }, + ], + }) const pid = Number( yield* Effect.promise(() => readFile(`${artifacts}/service.pid`, "utf8"), diff --git a/packages/drive/test/fixtures/failed-launch-retry-script.ts b/packages/drive/test/fixtures/failed-launch-retry-script.ts new file mode 100644 index 0000000..50726c7 --- /dev/null +++ b/packages/drive/test/fixtures/failed-launch-retry-script.ts @@ -0,0 +1,52 @@ +import { defineScript } from "../../src/index.js" +import * as Effect from "effect/Effect" + +export default defineScript({ + launch: "manual", + run: ({ server, artifacts }) => + Effect.gen(function* () { + const first = yield* server.launch().pipe(Effect.flip) + const firstPid = yield* servicePid(artifacts) + yield* waitUntilStopped(firstPid) + + const second = yield* server.launch().pipe(Effect.flip) + const secondPid = yield* servicePid(artifacts) + yield* waitUntilStopped(secondPid) + + if (first.operation !== "opencode.connect" || second.operation !== "opencode.connect") + return yield* Effect.fail( + new Error(`unexpected launch failures: ${first.operation}, ${second.operation}`), + ) + if (firstPid === secondPid) + return yield* Effect.fail(new Error("failed launch did not start a fresh server process")) + + yield* Effect.tryPromise(() => + Bun.write( + `${artifacts}/failed-launch-retry.json`, + JSON.stringify({ firstPid, secondPid }), + ), + ) + }), +}) + +const servicePid = (artifacts: string) => + Effect.tryPromise(() => + Bun.file(`${artifacts}/service.pid`).text().then(Number), + ) + +const waitUntilStopped = (pid: number) => + Effect.gen(function* () { + for (let attempt = 0; attempt < 100 && running(pid); attempt++) + yield* Effect.sleep(10) + if (running(pid)) + return yield* Effect.fail(new Error(`server process ${pid} is still running`)) + }) + +function running(pid: number) { + try { + process.kill(pid, 0) + return true + } catch { + return false + } +} diff --git a/packages/drive/test/fixtures/fake-opencode.ts b/packages/drive/test/fixtures/fake-opencode.ts index 295b4f3..cbd9e02 100644 --- a/packages/drive/test/fixtures/fake-opencode.ts +++ b/packages/drive/test/fixtures/fake-opencode.ts @@ -50,7 +50,11 @@ const api = role === "service" }, }) : undefined -if (api !== undefined && process.env.XDG_STATE_HOME) { +if ( + api !== undefined && + process.env.XDG_STATE_HOME && + !process.argv.includes("omit-service-registration") +) { const directory = `${process.env.XDG_STATE_HOME}/opencode` await mkdir(directory, { recursive: true }) await Bun.write( diff --git a/packages/drive/test/fixtures/kill-server-script.ts b/packages/drive/test/fixtures/kill-server-script.ts index 5b72ff3..1b30661 100644 --- a/packages/drive/test/fixtures/kill-server-script.ts +++ b/packages/drive/test/fixtures/kill-server-script.ts @@ -3,7 +3,7 @@ import * as Effect from "effect/Effect" export default defineScript({ launch: "manual", - run: ({ server, tuis, artifacts }) => + run: ({ server, tools, tuis, artifacts }) => Effect.gen(function* () { yield* server.launch() const firstServer = Number( @@ -41,6 +41,16 @@ export default defineScript({ yield* alice.close() const relaunchedAlice = yield* tuis.launch("alice") yield* relaunchedAlice.close() + yield* tools.attach({ + tools: [ + { + name: "lookup", + description: "Look up a value", + inputSchema: { type: "object" }, + options: { codemode: false }, + }, + ], + }) yield* server.kill() yield* Effect.tryPromise(() => diff --git a/packages/drive/test/tool/producer.test.ts b/packages/drive/test/tool/producer.test.ts index e2fdce9..7b2feac 100644 --- a/packages/drive/test/tool/producer.test.ts +++ b/packages/drive/test/tool/producer.test.ts @@ -117,6 +117,42 @@ it.live("attaches, sequences progress, and settles one dynamic invocation", () = callID: "call_lookup", }) yield* controller.settle + expect(yield* controller.controls.attach({ tools: [registration] }).pipe(Effect.flip)).toMatchObject({ + operation: "attach", + reason: "controller-closed", + }) + }), + ), +) + +it.live("fails settlement after a completed invocation ID is reused", () => + Effect.scoped( + Effect.gen(function* () { + const peer = startTransportPeer(({ request, socket }) => + sendResult( + socket, + request, + request.method === "tool.attach" ? { attached: true } : { ok: true }, + ), + ) + yield* Effect.addFinalizer(() => Effect.promise(() => peer.stop())) + const backend = yield* SimulationConnector.backend(peer.url, { + attach: false, + }) + const controller = yield* ToolProducer.make(new Set()) + yield* controller.connect(backend) + yield* controller.controls.attach({ tools: [registration] }) + const socket = peer.received[0].socket + notify(socket, "tool.invocation", invocation) + const call = yield* controller.controls.take("call_lookup") + yield* call.finish({ structured: { answer: 42 }, content: [] }) + notify(socket, "tool.invocation", { + ...invocation, + input: { query: "different" }, + }) + const failure = yield* controller.failure.pipe(Effect.flip) + + expect(yield* controller.settle.pipe(Effect.flip)).toEqual(failure) }), ), ) @@ -155,6 +191,47 @@ it.live("observes cancellation and rejects terminal work after it wins", () => ), ) +it.live("bounds retained unclaimed cancellations", () => + Effect.scoped( + Effect.gen(function* () { + const peer = startTransportPeer(({ request, socket }) => + sendResult(socket, request, { attached: true }), + ) + yield* Effect.addFinalizer(() => Effect.promise(() => peer.stop())) + const backend = yield* SimulationConnector.backend(peer.url, { + attach: false, + }) + const controller = yield* ToolProducer.make(new Set()) + yield* controller.connect(backend) + yield* controller.controls.attach({ tools: [registration] }) + const latest = yield* controller.controls.take("call_257").pipe(Effect.forkScoped) + const socket = peer.received[0].socket + + for (let index = 0; index < 258; index++) { + notify(socket, "tool.invocation", { + ...invocation, + id: `tool_${index}`, + context: { ...invocation.context, callID: `call_${index}` }, + }) + notify(socket, "tool.cancel", { + id: `tool_${index}`, + reason: "interrupted", + }) + } + + const latestCall = yield* Fiber.join(latest) + yield* latestCall.awaitCancelled() + const retained = yield* controller.controls.take("call_256") + expect(yield* retained.awaitCancelled()).toMatchObject({ id: "tool_256" }) + const evicted = yield* controller.controls.take("call_0").pipe(Effect.forkScoped) + yield* Effect.yieldNow + expect(evicted.pollUnsafe()).toBeUndefined() + yield* Fiber.interrupt(evicted) + yield* controller.settle + }), + ), +) + it.live("settles concurrent invocations in controlled reverse order", () => Effect.scoped( Effect.gen(function* () { @@ -162,7 +239,9 @@ it.live("settles concurrent invocations in controlled reverse order", () => sendResult(socket, request, request.method === "tool.attach" ? { attached: true } : { ok: true }), ) yield* Effect.addFinalizer(() => Effect.promise(() => peer.stop())) - const backend = yield* SimulationConnector.backend(peer.url, { attach: false }) + const backend = yield* SimulationConnector.backend(peer.url, { + attach: false, + }) const controller = yield* ToolProducer.make(new Set()) yield* controller.connect(backend) yield* controller.controls.attach({ tools: [registration] }) @@ -437,6 +516,63 @@ it.live("restores the acknowledged set after a rejected replacement", () => ), ) +it.live("rejects a disconnected replacement without poisoning relaunch", () => + Effect.scoped( + Effect.gen(function* () { + let attaches = 0 + const replacement = { + ...registration, + name: "search", + description: "Search for a value", + } + const peer = startTransportPeer(({ request, socket }) => { + if (request.method === "tool.attach" && ++attaches === 2) { + sendError(socket, request, "replacement rejected") + return + } + sendResult(socket, request, { attached: true }) + }) + yield* Effect.addFinalizer(() => Effect.promise(() => peer.stop())) + const controller = yield* ToolProducer.make(new Set()) + const firstScope = yield* Scope.make() + const first = yield* SimulationConnector.backend(peer.url, { + attach: false, + }).pipe(Scope.provide(firstScope)) + const firstAttachment = yield* controller.connect(first) + yield* controller.controls.attach({ tools: [registration] }) + yield* firstAttachment.detach() + yield* Scope.close(firstScope, Exit.void) + + const replacing = yield* controller.controls.attach({ tools: [replacement] }).pipe(Effect.forkScoped) + yield* Effect.yieldNow + const rejectedScope = yield* Scope.make() + const rejectedBackend = yield* SimulationConnector.backend(peer.url, { + attach: false, + }).pipe(Scope.provide(rejectedScope)) + expect(yield* controller.connect(rejectedBackend).pipe(Effect.flip)).toMatchObject({ + operation: "attach", + reason: "rejected", + }) + expect(yield* Fiber.join(replacing).pipe(Effect.flip)).toMatchObject({ + operation: "attach", + reason: "rejected", + }) + yield* Scope.close(rejectedScope, Exit.void) + + const recoveredScope = yield* Scope.make() + const recovered = yield* SimulationConnector.backend(peer.url, { + attach: false, + }).pipe(Scope.provide(recoveredScope)) + const recoveredAttachment = yield* controller.connect(recovered) + expect( + peer.received.filter(({ request }) => request.method === "tool.attach").map(({ request }) => request.params), + ).toEqual([{ tools: [registration] }, { tools: [replacement] }, { tools: [registration] }]) + yield* recoveredAttachment.detach() + yield* Scope.close(recoveredScope, Exit.void) + }), + ), +) + it.live("rejects malformed lifecycle acknowledgements without retrying", () => Effect.scoped( Effect.gen(function* () { @@ -444,82 +580,99 @@ it.live("rejects malformed lifecycle acknowledgements without retrying", () => sendResult(socket, request, { attached: "invalid" }), ) yield* Effect.addFinalizer(() => Effect.promise(() => peer.stop())) - const backend = yield* SimulationConnector.backend(peer.url, { attach: false }) + const backend = yield* SimulationConnector.backend(peer.url, { + attach: false, + }) const controller = yield* ToolProducer.make(new Set()) - yield* controller.connect(backend) + const attachment = yield* controller.connect(backend) expect( yield* controller.controls.attach({ tools: [registration] }).pipe(Effect.flip), ).toMatchObject({ operation: "attach", reason: "rejected" }) expect(peer.received.filter(({ request }) => request.method === "tool.attach")).toHaveLength(1) + yield* attachment.detach() + const recovered = yield* SimulationConnector.backend(peer.url, { + attach: false, + }) + yield* controller.connect(recovered) + expect(peer.received.filter(({ request }) => request.method === "tool.attach")).toHaveLength(1) }), ), ) -it.live("drains queued invocations before reporting settlement", () => +it.live("rejects a malformed disconnected replacement and restores the acknowledged set", () => Effect.scoped( Effect.gen(function* () { let attaches = 0 + const replacement = { + ...registration, + name: "search", + description: "Search for a value", + } const peer = startTransportPeer(({ request, socket }) => { - if (request.method === "tool.attach") { - attaches++ - if (attaches === 2) notify(socket, "tool.invocation", invocation) - sendResult(socket, request, { attached: true }) + if (request.method === "tool.attach" && ++attaches === 2) { + sendResult(socket, request, { attached: "invalid" }) return } - sendResult(socket, request, { ok: true }) + sendResult(socket, request, { attached: true }) }) yield* Effect.addFinalizer(() => Effect.promise(() => peer.stop())) - const backend = yield* SimulationConnector.backend(peer.url, { attach: false }) const controller = yield* ToolProducer.make(new Set()) - yield* controller.connect(backend) + const firstScope = yield* Scope.make() + const first = yield* SimulationConnector.backend(peer.url, { + attach: false, + }).pipe(Scope.provide(firstScope)) + const firstAttachment = yield* controller.connect(first) yield* controller.controls.attach({ tools: [registration] }) + yield* firstAttachment.detach() + yield* Scope.close(firstScope, Exit.void) - expect(yield* controller.settle.pipe(Effect.flip)).toMatchObject({ - operation: "take", + const replacing = yield* controller.controls.attach({ tools: [replacement] }).pipe(Effect.forkScoped) + const rejectedScope = yield* Scope.make() + const rejected = yield* SimulationConnector.backend(peer.url, { + attach: false, + }).pipe(Scope.provide(rejectedScope)) + expect(yield* controller.connect(rejected).pipe(Effect.flip)).toMatchObject({ + operation: "attach", reason: "rejected", - message: expect.stringContaining("1 dynamic tool invocation"), }) - const call = yield* controller.controls.take("call_lookup") - yield* call.fail("finished") - yield* controller.settle - expect( - peer.received.filter(({ request }) => request.method === "tool.attach").map(({ request }) => request.params), - ).toEqual([{ tools: [registration] }, { tools: [] }]) - }), - ), -) + expect(yield* Fiber.join(replacing).pipe(Effect.flip)).toMatchObject({ + operation: "attach", + reason: "rejected", + }) + yield* Scope.close(rejectedScope, Exit.void) -it.effect("rejects a malformed take call ID", () => - Effect.scoped( - Effect.gen(function* () { - const controller = yield* ToolProducer.make(new Set()) + const recoveredScope = yield* Scope.make() + const recovered = yield* SimulationConnector.backend(peer.url, { + attach: false, + }).pipe(Scope.provide(recoveredScope)) + const recoveredAttachment = yield* controller.connect(recovered) expect( - yield* controller.controls.take(null as unknown as string).pipe(Effect.flip), - ).toMatchObject({ operation: "take", reason: "rejected" }) + peer.received.filter(({ request }) => request.method === "tool.attach").map(({ request }) => request.params), + ).toEqual([{ tools: [registration] }, { tools: [replacement] }, { tools: [registration] }]) + yield* recoveredAttachment.detach() + yield* Scope.close(recoveredScope, Exit.void) }), ), ) -it.live("retries in-flight progress after reconnect without advancing its sequence", () => +it.live("rejects a disconnected replacement after an invalid protocol response", () => Effect.scoped( Effect.gen(function* () { - let updates = 0 - const firstUpdate = yield* Deferred.make() + let attaches = 0 + const replacement = { + ...registration, + name: "search", + description: "Search for a value", + } const peer = startTransportPeer(({ request, socket }) => { - if (request.method === "tool.update" && ++updates === 1) { - Deferred.doneUnsafe(firstUpdate, Effect.void) - socket.close() + if (request.method === "tool.attach" && ++attaches === 2) { + socket.send("not-json") return } - sendResult(socket, request, request.method === "tool.attach" ? { attached: true } : { ok: true }) + sendResult(socket, request, { attached: true }) }) - yield* Effect.addFinalizer(() => - Effect.sync(() => { - // Bun does not resolve server.stop after this peer initiates a WebSocket close. - void peer.stop() - }), - ) + yield* Effect.addFinalizer(() => Effect.promise(() => peer.stop())) const controller = yield* ToolProducer.make(new Set()) const firstScope = yield* Scope.make() const first = yield* SimulationConnector.backend(peer.url, { @@ -527,85 +680,707 @@ it.live("retries in-flight progress after reconnect without advancing its sequen }).pipe(Scope.provide(firstScope)) const firstAttachment = yield* controller.connect(first) yield* controller.controls.attach({ tools: [registration] }) - notify(peer.received[0].socket, "tool.invocation", invocation) - const call = yield* controller.controls.take("call_lookup") - - const progress = yield* call - .progress({ structured: { phase: "searching" } }) - .pipe(Effect.forkScoped) - yield* Deferred.await(firstUpdate) - yield* first.closed yield* firstAttachment.detach() yield* Scope.close(firstScope, Exit.void) - const secondScope = yield* Scope.make() - const second = yield* SimulationConnector.backend(peer.url, { - attach: false, - }).pipe(Scope.provide(secondScope)) - const secondAttachment = yield* controller.connect(second) - yield* Fiber.join(progress) - yield* call.fail("finished") - expect( - peer.received.filter(({ request }) => request.method === "tool.update").map(({ request }) => request.params), - ).toEqual([ - { - id: "tool_1", - sequence: 0, - update: { structured: { phase: "searching" } }, - }, - { - id: "tool_1", - sequence: 0, - update: { structured: { phase: "searching" } }, - }, - ]) - yield* secondAttachment.detach() - yield* Scope.close(secondScope, Exit.void) - yield* controller.endGeneration - yield* controller.shutdown + const replacing = yield* controller.controls + .attach({ tools: [replacement] }) + .pipe(Effect.forkScoped) + const rejectedScope = yield* Scope.make() + const rejected = yield* SimulationConnector.backend(peer.url, { + attach: false, + }).pipe(Scope.provide(rejectedScope)) + expect(yield* controller.connect(rejected).pipe(Effect.flip)).toMatchObject({ + operation: "attach", + reason: "rejected", + }) + expect(yield* Fiber.join(replacing).pipe(Effect.flip)).toMatchObject({ + operation: "attach", + reason: "rejected", + }) + yield* Scope.close(rejectedScope, Exit.void) }), ), ) -it.live("does not retry progress interrupted by its caller", () => +it.live("drains queued invocations before reporting settlement", () => Effect.scoped( Effect.gen(function* () { - const updateReceived = yield* Deferred.make() + let attaches = 0 const peer = startTransportPeer(({ request, socket }) => { - if (request.method === "tool.update") { - Deferred.doneUnsafe(updateReceived, Effect.void) + if (request.method === "tool.attach") { + attaches++ + if (attaches === 2) notify(socket, "tool.invocation", invocation) + sendResult(socket, request, { attached: true }) return } - sendResult(socket, request, request.method === "tool.attach" ? { attached: true } : { ok: true }) + sendResult(socket, request, { ok: true }) }) yield* Effect.addFinalizer(() => Effect.promise(() => peer.stop())) - const backend = yield* SimulationConnector.backend(peer.url, { attach: false }) + const backend = yield* SimulationConnector.backend(peer.url, { + attach: false, + }) const controller = yield* ToolProducer.make(new Set()) yield* controller.connect(backend) yield* controller.controls.attach({ tools: [registration] }) - notify(peer.received[0].socket, "tool.invocation", invocation) - const call = yield* controller.controls.take("call_lookup") - const progress = yield* call - .progress({ structured: { phase: "searching" } }) - .pipe(Effect.forkScoped) - yield* Deferred.await(updateReceived) - yield* Fiber.interrupt(progress) + expect(yield* controller.settle.pipe(Effect.flip)).toMatchObject({ + operation: "take", + reason: "rejected", + message: expect.stringContaining("1 dynamic tool invocation"), + }) + const call = yield* controller.controls.take("call_lookup") yield* call.fail("finished") - - expect(peer.received.filter(({ request }) => request.method === "tool.update")).toHaveLength(1) + yield* controller.settle + expect( + peer.received.filter(({ request }) => request.method === "tool.attach").map(({ request }) => request.params), + ).toEqual([{ tools: [registration] }, { tools: [] }]) }), ), ) -it.live("rejects malformed progress without advancing its sequence", () => +it.live("rolls interrupted replacement history back to the acknowledged set", () => Effect.scoped( Effect.gen(function* () { - const peer = startTransportPeer(({ request, socket }) => - sendResult(socket, request, request.method === "tool.attach" ? { attached: true } : { ok: true }), - ) - yield* Effect.addFinalizer(() => Effect.promise(() => peer.stop())) - const backend = yield* SimulationConnector.backend(peer.url, { attach: false }) + let attaches = 0 + const interruptedStarted = yield* Deferred.make() + const interrupted = { + ...registration, + name: "search", + description: "Search for a value", + } + const rejected = { + ...registration, + name: "fetch", + description: "Fetch a value", + } + const peer = startTransportPeer(({ request, socket }) => { + if (request.method === "tool.attach") { + attaches++ + if (attaches === 2) { + Deferred.doneUnsafe(interruptedStarted, Effect.void) + return + } + if (attaches === 3) { + sendError(socket, request, "replacement rejected") + return + } + } + sendResult(socket, request, { attached: true }) + }) + yield* Effect.addFinalizer(() => Effect.promise(() => peer.stop())) + const controller = yield* ToolProducer.make(new Set()) + const firstScope = yield* Scope.make() + const first = yield* SimulationConnector.backend(peer.url, { + attach: false, + }).pipe(Scope.provide(firstScope)) + const firstAttachment = yield* controller.connect(first) + yield* controller.controls.attach({ tools: [registration] }) + + const replacing = yield* controller.controls.attach({ tools: [interrupted] }).pipe(Effect.forkScoped) + yield* Deferred.await(interruptedStarted) + yield* Fiber.interrupt(replacing) + expect(yield* controller.controls.attach({ tools: [rejected] }).pipe(Effect.flip)).toMatchObject({ + operation: "attach", + reason: "rejected", + }) + yield* firstAttachment.detach() + yield* Scope.close(firstScope, Exit.void) + + const recoveredScope = yield* Scope.make() + const recovered = yield* SimulationConnector.backend(peer.url, { + attach: false, + }).pipe(Scope.provide(recoveredScope)) + const recoveredAttachment = yield* controller.connect(recovered) + expect( + peer.received.filter(({ request }) => request.method === "tool.attach").map(({ request }) => request.params), + ).toEqual([{ tools: [registration] }, { tools: [interrupted] }, { tools: [rejected] }, { tools: [registration] }]) + yield* recoveredAttachment.detach() + yield* Scope.close(recoveredScope, Exit.void) + }), + ), +) + +it.live("fails settlement after the inbound event stream fails", () => + Effect.scoped( + Effect.gen(function* () { + const peer = startTransportPeer(({ request, socket }) => sendResult(socket, request, { attached: true })) + yield* Effect.addFinalizer(() => Effect.promise(() => peer.stop())) + const backend = yield* SimulationConnector.backend(peer.url, { + attach: false, + }) + const controller = yield* ToolProducer.make(new Set()) + yield* controller.connect(backend) + yield* controller.controls.attach({ tools: [registration] }) + notify(peer.received[0].socket, "tool.invocation", { + ...invocation, + id: 1, + }) + const streamFailure = yield* controller.failure.pipe(Effect.flip) + + expect(yield* controller.settle.pipe(Effect.flip)).toEqual(streamFailure) + expect(peer.received.filter(({ request }) => request.method === "tool.attach")).toHaveLength(1) + }), + ), +) + +it.live("settles without a backend after its generation ends", () => + Effect.scoped( + Effect.gen(function* () { + const peer = startTransportPeer(({ request, socket }) => sendResult(socket, request, { attached: true })) + yield* Effect.addFinalizer(() => Effect.promise(() => peer.stop())) + const backendScope = yield* Scope.make() + const backend = yield* SimulationConnector.backend(peer.url, { + attach: false, + }).pipe(Scope.provide(backendScope)) + const controller = yield* ToolProducer.make(new Set()) + const attachment = yield* controller.connect(backend) + yield* controller.controls.attach({ tools: [registration] }) + yield* attachment.detach() + yield* Scope.close(backendScope, Exit.void) + yield* controller.endGeneration + + yield* controller.settle + }), + ), +) + +it.live("settles when the generation ends during its clear request", () => + Effect.scoped( + Effect.gen(function* () { + let attaches = 0 + const clearStarted = yield* Deferred.make() + const peer = startTransportPeer(({ request, socket }) => { + if (request.method === "tool.attach" && ++attaches === 2) { + Deferred.doneUnsafe(clearStarted, Effect.void) + return + } + sendResult(socket, request, { attached: true }) + }) + yield* Effect.addFinalizer(() => Effect.promise(() => peer.stop())) + const backendScope = yield* Scope.make() + const backend = yield* SimulationConnector.backend(peer.url, { + attach: false, + }).pipe(Scope.provide(backendScope)) + const controller = yield* ToolProducer.make(new Set()) + const attachment = yield* controller.connect(backend) + yield* controller.controls.attach({ tools: [registration] }) + + const settling = yield* controller.settle.pipe(Effect.forkScoped) + yield* Deferred.await(clearStarted) + yield* attachment.detach() + yield* controller.endGeneration + yield* Scope.close(backendScope, Exit.void) + + yield* Fiber.join(settling) + }), + ), +) + +it.live("reconnects when the backend disconnects during settlement", () => + Effect.scoped( + Effect.gen(function* () { + let attaches = 0 + const clearInterrupted = yield* Deferred.make() + const peer = startTransportPeer(({ request, socket }) => { + if (request.method === "tool.attach" && ++attaches === 2) { + Deferred.doneUnsafe(clearInterrupted, Effect.void) + socket.close() + return + } + sendResult(socket, request, { attached: true }) + }) + yield* Effect.addFinalizer(() => + Effect.sync(() => { + void peer.stop() + }), + ) + const controller = yield* ToolProducer.make(new Set()) + const firstScope = yield* Scope.make() + const first = yield* SimulationConnector.backend(peer.url, { + attach: false, + }).pipe(Scope.provide(firstScope)) + const firstAttachment = yield* controller.connect(first) + yield* controller.controls.attach({ tools: [registration] }) + + const settling = yield* controller.settle.pipe(Effect.forkScoped) + yield* Deferred.await(clearInterrupted) + yield* first.closed + yield* firstAttachment.detach() + yield* Scope.close(firstScope, Exit.void) + const recoveredScope = yield* Scope.make() + const recovered = yield* SimulationConnector.backend(peer.url, { + attach: false, + }).pipe(Scope.provide(recoveredScope)) + const recoveredAttachment = yield* controller.connect(recovered) + + yield* Fiber.join(settling) + expect( + peer.received + .filter(({ request }) => request.method === "tool.attach") + .map(({ request }) => request.params), + ).toEqual([ + { tools: [registration] }, + { tools: [] }, + { tools: [] }, + { tools: [] }, + ]) + yield* recoveredAttachment.detach() + yield* Scope.close(recoveredScope, Exit.void) + }), + ), +) + +it.live("retries the final event barrier after a disconnect", () => + Effect.scoped( + Effect.gen(function* () { + let flushes = 0 + const finalBarrierStarted = yield* Deferred.make() + const peer = startTransportPeer(({ request, socket }) => + sendResult(socket, request, { attached: true }), + ) + yield* Effect.addFinalizer(() => Effect.promise(() => peer.stop())) + const controller = yield* ToolProducer.make(new Set()) + const firstScope = yield* Scope.make() + const first = yield* SimulationConnector.backend(peer.url, { + attach: false, + }).pipe(Scope.provide(firstScope)) + const blocked = { + ...first, + flushToolEvents: () => { + flushes++ + return flushes === 2 + ? Deferred.succeed(finalBarrierStarted, undefined).pipe( + Effect.andThen(Effect.never), + ) + : first.flushToolEvents() + }, + } + const firstAttachment = yield* controller.connect(blocked) + yield* controller.controls.attach({ tools: [registration] }) + + const settling = yield* controller.settle.pipe(Effect.forkScoped) + yield* Deferred.await(finalBarrierStarted) + yield* firstAttachment.detach() + yield* Scope.close(firstScope, Exit.void) + yield* Effect.yieldNow + expect(settling.pollUnsafe()).toBeUndefined() + + const recoveredScope = yield* Scope.make() + const recovered = yield* SimulationConnector.backend(peer.url, { + attach: false, + }).pipe(Scope.provide(recoveredScope)) + const recoveredAttachment = yield* controller.connect(recovered) + yield* Fiber.join(settling) + expect( + peer.received + .filter(({ request }) => request.method === "tool.attach") + .map(({ request }) => request.params), + ).toEqual([ + { tools: [registration] }, + { tools: [] }, + { tools: [] }, + ]) + yield* recoveredAttachment.detach() + yield* Scope.close(recoveredScope, Exit.void) + }), + ), +) + +it.live("rejects a backend connection after terminal settlement", () => + Effect.scoped( + Effect.gen(function* () { + const peer = startTransportPeer(({ request, socket }) => sendResult(socket, request, { attached: true })) + yield* Effect.addFinalizer(() => Effect.promise(() => peer.stop())) + const firstScope = yield* Scope.make() + const first = yield* SimulationConnector.backend(peer.url, { + attach: false, + }).pipe(Scope.provide(firstScope)) + const controller = yield* ToolProducer.make(new Set()) + const firstAttachment = yield* controller.connect(first) + yield* controller.controls.attach({ tools: [registration] }) + yield* firstAttachment.detach() + yield* Scope.close(firstScope, Exit.void) + yield* controller.endGeneration + yield* controller.settle + + const nextScope = yield* Scope.make() + const next = yield* SimulationConnector.backend(peer.url, { + attach: false, + }).pipe(Scope.provide(nextScope)) + expect(yield* controller.connect(next).pipe(Effect.flip)).toMatchObject({ + operation: "attach", + reason: "controller-closed", + }) + expect(peer.received.filter(({ request }) => request.method === "tool.attach")).toHaveLength(1) + yield* Scope.close(nextScope, Exit.void) + }), + ), +) + +it.effect("closes take waiters and backend creation after settlement", () => + Effect.scoped( + Effect.gen(function* () { + const controller = yield* ToolProducer.make(new Set()) + const waiting = yield* controller.controls.take().pipe(Effect.forkScoped) + yield* controller.settle + + expect(yield* Fiber.join(waiting).pipe(Effect.flip)).toMatchObject({ + operation: "take", + reason: "controller-closed", + }) + expect(yield* controller.controls.take().pipe(Effect.flip)).toMatchObject({ + operation: "take", + reason: "controller-closed", + }) + let evaluated = false + const backend = Effect.sync(() => { + evaluated = true + throw new Error("backend factory was evaluated") + }) + expect(yield* controller.connectFrom(backend).pipe(Effect.flip)).toMatchObject({ + operation: "attach", + reason: "controller-closed", + }) + expect(evaluated).toBe(false) + }), + ), +) + +it.live("settles after a backend connection that was already in progress", () => + Effect.scoped( + Effect.gen(function* () { + let attaches = 0 + const replayStarted = yield* Deferred.make() + const clearStarted = yield* Deferred.make() + const peer = startTransportPeer(({ request, socket }) => { + if (request.method === "tool.attach") { + attaches++ + if (attaches === 2) { + Deferred.doneUnsafe(replayStarted, Effect.void) + return + } + if (attaches === 3) Deferred.doneUnsafe(clearStarted, Effect.void) + } + sendResult(socket, request, { attached: true }) + }) + yield* Effect.addFinalizer(() => Effect.promise(() => peer.stop())) + const controller = yield* ToolProducer.make(new Set()) + const firstScope = yield* Scope.make() + const first = yield* SimulationConnector.backend(peer.url, { + attach: false, + }).pipe(Scope.provide(firstScope)) + const firstAttachment = yield* controller.connect(first) + yield* controller.controls.attach({ tools: [registration] }) + yield* firstAttachment.detach() + yield* controller.endGeneration + yield* Scope.close(firstScope, Exit.void) + + const nextScope = yield* Scope.make() + const next = yield* SimulationConnector.backend(peer.url, { + attach: false, + }).pipe(Scope.provide(nextScope)) + const connecting = yield* controller.connect(next).pipe(Effect.forkScoped) + yield* Deferred.await(replayStarted) + const settling = yield* controller.settle.pipe(Effect.forkScoped) + const replay = peer.received.filter(({ request }) => request.method === "tool.attach")[1] + if (replay === undefined) return yield* Effect.die(new Error("missing replay attachment request")) + sendResult(replay.socket, replay.request, { attached: true }) + yield* Deferred.await(clearStarted) + + const nextAttachment = yield* Fiber.join(connecting) + yield* Fiber.join(settling) + expect( + peer.received.filter(({ request }) => request.method === "tool.attach").map(({ request }) => request.params), + ).toEqual([{ tools: [registration] }, { tools: [registration] }, { tools: [] }]) + yield* nextAttachment.detach() + yield* Scope.close(nextScope, Exit.void) + }), + ), +) + +it.live("drains a reconnect that finishes during settlement commit", () => + Effect.scoped( + Effect.gen(function* () { + let attaches = 0 + let flushes = 0 + const firstDrained = yield* Deferred.make() + const releaseDrain = yield* Deferred.make() + const staleInvocation = { + ...invocation, + id: "tool_stale", + context: { ...invocation.context, callID: "call_stale" }, + } + const peer = startTransportPeer(({ request, socket }) => { + if (request.method === "tool.attach") { + attaches++ + if (attaches === 3) notify(socket, "tool.invocation", staleInvocation) + } + sendResult(socket, request, request.method === "tool.attach" ? { attached: true } : { ok: true }) + }) + yield* Effect.addFinalizer(() => Effect.promise(() => peer.stop())) + const controller = yield* ToolProducer.make(new Set()) + const firstScope = yield* Scope.make() + const first = yield* SimulationConnector.backend(peer.url, { + attach: false, + }).pipe(Scope.provide(firstScope)) + const blocked = { + ...first, + flushToolEvents: () => { + flushes++ + return first.flushToolEvents().pipe( + Effect.andThen( + flushes === 1 + ? Deferred.succeed(firstDrained, undefined).pipe( + Effect.andThen(Deferred.await(releaseDrain)), + ) + : Effect.void, + ), + ) + }, + } + const firstAttachment = yield* controller.connect(blocked) + yield* controller.controls.attach({ tools: [registration] }) + + const settling = yield* controller.settle.pipe(Effect.forkScoped) + yield* Deferred.await(firstDrained) + yield* firstAttachment.detach() + yield* Scope.close(firstScope, Exit.void) + const nextScope = yield* Scope.make() + const next = yield* SimulationConnector.backend(peer.url, { + attach: false, + }).pipe(Scope.provide(nextScope)) + const nextAttachment = yield* controller.connect(next) + yield* Deferred.succeed(releaseDrain, undefined) + + expect(yield* Fiber.join(settling).pipe(Effect.flip)).toMatchObject({ + operation: "take", + reason: "rejected", + message: expect.stringContaining("1 dynamic tool invocation"), + }) + const stale = yield* controller.controls.take("call_stale") + yield* stale.fail("finished") + yield* controller.settle + yield* nextAttachment.detach() + yield* Scope.close(nextScope, Exit.void) + }), + ), +) + +it.live("serializes settlement after interrupting an in-flight replacement", () => + Effect.scoped( + Effect.gen(function* () { + let attaches = 0 + const replacementStarted = yield* Deferred.make() + const clearStarted = yield* Deferred.make() + const replacement = { + ...registration, + name: "search", + description: "Search for a value", + } + const peer = startTransportPeer(({ request, socket }) => { + if (request.method === "tool.attach") { + attaches++ + if (attaches === 2) { + Deferred.doneUnsafe(replacementStarted, Effect.void) + return + } + if (attaches === 3) Deferred.doneUnsafe(clearStarted, Effect.void) + } + sendResult(socket, request, { attached: true }) + }) + yield* Effect.addFinalizer(() => Effect.promise(() => peer.stop())) + const backend = yield* SimulationConnector.backend(peer.url, { + attach: false, + }) + const controller = yield* ToolProducer.make(new Set()) + yield* controller.connect(backend) + yield* controller.controls.attach({ tools: [registration] }) + + const replacing = yield* controller.controls.attach({ tools: [replacement] }).pipe(Effect.forkScoped) + yield* Deferred.await(replacementStarted) + const settling = yield* controller.settle.pipe(Effect.forkScoped) + yield* Deferred.await(clearStarted) + + expect(yield* Fiber.join(replacing).pipe(Effect.flip)).toMatchObject({ + operation: "attach", + reason: "controller-closed", + }) + yield* Fiber.join(settling) + expect( + peer.received.filter(({ request }) => request.method === "tool.attach").map(({ request }) => request.params), + ).toEqual([{ tools: [registration] }, { tools: [replacement] }, { tools: [] }]) + }), + ), +) + +it.live("releases a disconnected replacement when its generation ends", () => + Effect.scoped( + Effect.gen(function* () { + let attaches = 0 + const replacementStarted = yield* Deferred.make() + const replacement = { + ...registration, + name: "search", + description: "Search for a value", + } + const peer = startTransportPeer(({ request, socket }) => { + if (request.method === "tool.attach" && ++attaches === 2) { + Deferred.doneUnsafe(replacementStarted, Effect.void) + socket.close() + return + } + sendResult(socket, request, { attached: true }) + }) + yield* Effect.addFinalizer(() => + Effect.sync(() => { + void peer.stop() + }), + ) + const backendScope = yield* Scope.make() + const backend = yield* SimulationConnector.backend(peer.url, { + attach: false, + }).pipe(Scope.provide(backendScope)) + const controller = yield* ToolProducer.make(new Set()) + const attachment = yield* controller.connect(backend) + yield* controller.controls.attach({ tools: [registration] }) + + const replacing = yield* controller.controls.attach({ tools: [replacement] }).pipe(Effect.forkScoped) + yield* Deferred.await(replacementStarted) + yield* backend.closed + yield* attachment.detach() + yield* Scope.close(backendScope, Exit.void) + yield* controller.endGeneration + yield* controller.settle + + expect(yield* Fiber.join(replacing).pipe(Effect.flip)).toMatchObject({ + operation: "attach", + reason: "controller-closed", + }) + }), + ), +) + +it.effect("rejects a malformed take call ID", () => + Effect.scoped( + Effect.gen(function* () { + const controller = yield* ToolProducer.make(new Set()) + expect( + yield* controller.controls.take(null as unknown as string).pipe(Effect.flip), + ).toMatchObject({ operation: "take", reason: "rejected" }) + }), + ), +) + +it.live("retries in-flight progress after reconnect without advancing its sequence", () => + Effect.scoped( + Effect.gen(function* () { + let updates = 0 + const firstUpdate = yield* Deferred.make() + const peer = startTransportPeer(({ request, socket }) => { + if (request.method === "tool.update" && ++updates === 1) { + Deferred.doneUnsafe(firstUpdate, Effect.void) + socket.close() + return + } + sendResult(socket, request, request.method === "tool.attach" ? { attached: true } : { ok: true }) + }) + yield* Effect.addFinalizer(() => + Effect.sync(() => { + // Bun does not resolve server.stop after this peer initiates a WebSocket close. + void peer.stop() + }), + ) + const controller = yield* ToolProducer.make(new Set()) + const firstScope = yield* Scope.make() + const first = yield* SimulationConnector.backend(peer.url, { + attach: false, + }).pipe(Scope.provide(firstScope)) + const firstAttachment = yield* controller.connect(first) + yield* controller.controls.attach({ tools: [registration] }) + notify(peer.received[0].socket, "tool.invocation", invocation) + const call = yield* controller.controls.take("call_lookup") + + const progress = yield* call + .progress({ structured: { phase: "searching" } }) + .pipe(Effect.forkScoped) + yield* Deferred.await(firstUpdate) + yield* first.closed + yield* firstAttachment.detach() + yield* Scope.close(firstScope, Exit.void) + const secondScope = yield* Scope.make() + const second = yield* SimulationConnector.backend(peer.url, { + attach: false, + }).pipe(Scope.provide(secondScope)) + const secondAttachment = yield* controller.connect(second) + yield* Fiber.join(progress) + yield* call.fail("finished") + + expect( + peer.received.filter(({ request }) => request.method === "tool.update").map(({ request }) => request.params), + ).toEqual([ + { + id: "tool_1", + sequence: 0, + update: { structured: { phase: "searching" } }, + }, + { + id: "tool_1", + sequence: 0, + update: { structured: { phase: "searching" } }, + }, + ]) + yield* secondAttachment.detach() + yield* Scope.close(secondScope, Exit.void) + yield* controller.endGeneration + yield* controller.shutdown + }), + ), +) + +it.live("does not retry progress interrupted by its caller", () => + Effect.scoped( + Effect.gen(function* () { + const updateReceived = yield* Deferred.make() + const peer = startTransportPeer(({ request, socket }) => { + if (request.method === "tool.update") { + Deferred.doneUnsafe(updateReceived, Effect.void) + return + } + sendResult(socket, request, request.method === "tool.attach" ? { attached: true } : { ok: true }) + }) + yield* Effect.addFinalizer(() => Effect.promise(() => peer.stop())) + const backend = yield* SimulationConnector.backend(peer.url, { + attach: false, + }) + const controller = yield* ToolProducer.make(new Set()) + yield* controller.connect(backend) + yield* controller.controls.attach({ tools: [registration] }) + notify(peer.received[0].socket, "tool.invocation", invocation) + const call = yield* controller.controls.take("call_lookup") + + const progress = yield* call + .progress({ structured: { phase: "searching" } }) + .pipe(Effect.forkScoped) + yield* Deferred.await(updateReceived) + yield* Fiber.interrupt(progress) + yield* call.fail("finished") + + expect(peer.received.filter(({ request }) => request.method === "tool.update")).toHaveLength(1) + }), + ), +) + +it.live("rejects malformed progress without advancing its sequence", () => + Effect.scoped( + Effect.gen(function* () { + const peer = startTransportPeer(({ request, socket }) => + sendResult(socket, request, request.method === "tool.attach" ? { attached: true } : { ok: true }), + ) + yield* Effect.addFinalizer(() => Effect.promise(() => peer.stop())) + const backend = yield* SimulationConnector.backend(peer.url, { + attach: false, + }) const controller = yield* ToolProducer.make(new Set()) yield* controller.connect(backend) yield* controller.controls.attach({ tools: [registration] })