refactor(simulation): scope Drive lifecycle with Effect (#36908)
This commit is contained in:
Kit Langton
GitHub
parent
a149a61a89
commit
947566f611
30 changed files with 1313 additions and 600 deletions
|
|
@ -1,97 +0,0 @@
|
|||
import { Effect } from "effect"
|
||||
import { SimulationProtocol } from "../protocol"
|
||||
import { SimulationLLMExchange } from "./llm-exchange"
|
||||
|
||||
/**
|
||||
* Backend-hosted simulation control WebSocket.
|
||||
*
|
||||
* JSON-RPC 2.0 over a loopback WebSocket, mirroring the protocol of the TUI
|
||||
* simulation server. Drivers connect directly (standalone topology; no
|
||||
* frontend proxy) to answer LLM exchanges and inspect the simulated network.
|
||||
* This is also the headless-simulation interface: it works with no TUI at
|
||||
* all.
|
||||
*
|
||||
* Methods:
|
||||
* - `llm.attach` -> subscribe; pending and future exchanges arrive
|
||||
* as `llm.request` notifications
|
||||
* - `llm.chunk` { id, items } append response items to an exchange
|
||||
* - `llm.finish` { id, reason? } finish an exchange
|
||||
* - `llm.disconnect` { id } abruptly terminate an exchange without a finish
|
||||
* - `llm.pending` list open exchanges
|
||||
*/
|
||||
|
||||
type ControlSocket = Bun.ServerWebSocket<{ unsubscribe?: () => void }>
|
||||
|
||||
function parseRequest(input: string | Buffer) {
|
||||
return SimulationProtocol.Backend.decodeRequest(JSON.parse(typeof input === "string" ? input : input.toString()))
|
||||
}
|
||||
|
||||
async function handle(socket: ControlSocket, request: SimulationProtocol.Backend.Request): Promise<unknown> {
|
||||
switch (request.method) {
|
||||
case "llm.attach": {
|
||||
socket.data.unsubscribe?.()
|
||||
socket.data.unsubscribe = SimulationLLMExchange.subscribe((exchange) => {
|
||||
socket.send(JSON.stringify({ jsonrpc: "2.0", method: "llm.request", params: exchange }))
|
||||
})
|
||||
return { attached: true }
|
||||
}
|
||||
case "llm.chunk": {
|
||||
await Effect.runPromise(
|
||||
SimulationLLMExchange.push(
|
||||
request.params.id,
|
||||
request.params.items.map((item) => ({ type: "item", item }) as const),
|
||||
),
|
||||
)
|
||||
return { ok: true }
|
||||
}
|
||||
case "llm.finish": {
|
||||
await Effect.runPromise(
|
||||
SimulationLLMExchange.push(request.params.id, [{ type: "finish", reason: request.params.reason }]),
|
||||
)
|
||||
return { ok: true }
|
||||
}
|
||||
case "llm.disconnect": {
|
||||
await Effect.runPromise(SimulationLLMExchange.disconnect(request.params.id))
|
||||
return { ok: true }
|
||||
}
|
||||
case "llm.pending":
|
||||
return { exchanges: SimulationLLMExchange.pending() }
|
||||
}
|
||||
}
|
||||
|
||||
export function start(endpoint: string) {
|
||||
const url = new URL(endpoint)
|
||||
const server = Bun.serve<{ unsubscribe?: () => void }>({
|
||||
hostname: url.hostname,
|
||||
port: Number(url.port),
|
||||
fetch(request, server) {
|
||||
if (server.upgrade(request, { data: {} })) return undefined
|
||||
return new Response("opencode drive backend websocket", { status: 426 })
|
||||
},
|
||||
websocket: {
|
||||
close(socket) {
|
||||
socket.data.unsubscribe?.()
|
||||
},
|
||||
async message(socket, message) {
|
||||
let request: SimulationProtocol.Backend.Request | undefined
|
||||
try {
|
||||
request = parseRequest(message)
|
||||
const result = await handle(socket, request)
|
||||
const response = SimulationProtocol.JsonRpc.success(request.id, result)
|
||||
if (response) socket.send(JSON.stringify(response))
|
||||
} catch (error) {
|
||||
socket.send(JSON.stringify(SimulationProtocol.JsonRpc.failure(request?.id, error)))
|
||||
}
|
||||
},
|
||||
},
|
||||
})
|
||||
process.stderr.write(`opencode drive backend websocket: ${endpoint}\n`)
|
||||
return {
|
||||
url: endpoint,
|
||||
stop: () => {
|
||||
server.stop(true)
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
export * as SimulationControl from "./control"
|
||||
|
|
@ -1,9 +1,11 @@
|
|||
import { LayerNode } from "@opencode-ai/core/effect/layer-node"
|
||||
import { httpClient } from "@opencode-ai/core/effect/app-node-platform"
|
||||
import { Config, Effect, Layer } from "effect"
|
||||
import { HttpClient } from "effect/unstable/http"
|
||||
import { DriveManifest } from "../manifest"
|
||||
import { SimulationControl } from "./control"
|
||||
import { SimulationNetwork } from "./network"
|
||||
import { SimulationOpenAI } from "./openai"
|
||||
import { SimulatedProvider } from "./simulated-provider"
|
||||
|
||||
/**
|
||||
* Layer replacements applied when the server is built in simulation mode.
|
||||
|
|
@ -17,17 +19,29 @@ import { SimulationOpenAI } from "./openai"
|
|||
*
|
||||
*/
|
||||
|
||||
SimulationNetwork.register(SimulationOpenAI.route)
|
||||
// ModelsDev dies when its catalog fetch fails, so simulation answers it with
|
||||
// an empty catalog; providers come from seeded config instead.
|
||||
SimulationNetwork.register(SimulationNetwork.json("GET", "https://models.dev/api.json", {}))
|
||||
export const simulationReplacements = Effect.fn("Simulation.replacements")(function* () {
|
||||
// ModelsDev dies when its catalog fetch fails, so simulation answers it with
|
||||
// an empty catalog; providers come from seeded config instead.
|
||||
const models = SimulationNetwork.json("GET", "https://models.dev/api.json", {})
|
||||
const drive = yield* Config.string("OPENCODE_DRIVE").pipe(Config.withDefault(undefined))
|
||||
if (!drive) return [[httpClient, SimulationNetwork.layer([models])]] satisfies LayerNode.Replacements
|
||||
|
||||
export function startDriveServer() {
|
||||
return SimulationControl.start(DriveManifest.resolve().endpoints.backend)
|
||||
}
|
||||
|
||||
export const simulationReplacements: LayerNode.Replacements = [
|
||||
[httpClient, SimulationNetwork.layer],
|
||||
]
|
||||
const manifest = yield* DriveManifest.resolve()
|
||||
const networkLayer = Layer.effect(
|
||||
HttpClient.HttpClient,
|
||||
Effect.gen(function* () {
|
||||
const provider = yield* SimulatedProvider.Service
|
||||
const network = yield* SimulationNetwork.make([SimulationOpenAI.route(provider), models])
|
||||
return network.client
|
||||
}),
|
||||
).pipe(
|
||||
Layer.provide(
|
||||
SimulatedProvider.layerDrive({
|
||||
endpoint: manifest.endpoints.backend,
|
||||
}),
|
||||
),
|
||||
)
|
||||
return [[httpClient, networkLayer]] satisfies LayerNode.Replacements
|
||||
})
|
||||
|
||||
export * as Simulation from "./index"
|
||||
|
|
|
|||
|
|
@ -1,119 +0,0 @@
|
|||
import { Effect, Queue } from "effect"
|
||||
|
||||
/**
|
||||
* Pending driver-answered LLM exchanges.
|
||||
*
|
||||
* When the simulated network receives a provider request it opens an
|
||||
* exchange: the parsed request body plus a queue of response chunks. The
|
||||
* simulation control WebSocket notifies the external driver, and the driver
|
||||
* pushes chunks back until it finishes the exchange. The driver is the
|
||||
* model; nothing is scripted or enqueued server-side.
|
||||
*
|
||||
* Process-global by design (plain module state, like the network route
|
||||
* table): the simulated network and the control server must observe the same
|
||||
* exchanges regardless of which layer instance touched them.
|
||||
*/
|
||||
|
||||
/** One response item the driver sends back. Compiled to provider wire chunks by the endpoint. */
|
||||
export type Item =
|
||||
| { readonly type: "textDelta"; readonly text: string }
|
||||
| { readonly type: "reasoningDelta"; readonly text: string }
|
||||
| {
|
||||
readonly type: "toolCall"
|
||||
readonly index: number
|
||||
readonly id: string
|
||||
readonly name: string
|
||||
readonly input: unknown
|
||||
}
|
||||
| { readonly type: "raw"; readonly chunk: unknown }
|
||||
|
||||
export type FinishReason = "stop" | "tool-calls" | "length" | "content-filter"
|
||||
|
||||
export type Chunk =
|
||||
| { readonly type: "item"; readonly item: Item }
|
||||
| { readonly type: "finish"; readonly reason: FinishReason }
|
||||
|
||||
export interface Exchange {
|
||||
readonly id: string
|
||||
readonly url: string
|
||||
readonly body: unknown
|
||||
readonly queue: Queue.Queue<Chunk>
|
||||
}
|
||||
|
||||
export interface OpenedExchange {
|
||||
readonly id: string
|
||||
readonly url: string
|
||||
readonly body: unknown
|
||||
}
|
||||
|
||||
const state = {
|
||||
counter: 0,
|
||||
exchanges: new Map<string, Exchange>(),
|
||||
listeners: new Set<(exchange: OpenedExchange) => void>(),
|
||||
}
|
||||
|
||||
export class ExchangeNotFoundError extends Error {
|
||||
constructor(id: string) {
|
||||
super(`Simulation LLM exchange not found or already finished: ${id}`)
|
||||
}
|
||||
}
|
||||
|
||||
/** Opens an exchange and notifies listeners. Called by the simulated provider endpoint. */
|
||||
export const open = (input: { readonly url: string; readonly body: unknown }) =>
|
||||
Effect.gen(function* () {
|
||||
const id = `ex_${++state.counter}`
|
||||
const queue = yield* Queue.unbounded<Chunk>()
|
||||
const exchange: Exchange = { id, url: input.url, body: input.body, queue }
|
||||
state.exchanges.set(id, exchange)
|
||||
for (const listener of state.listeners) listener({ id, url: input.url, body: input.body })
|
||||
return exchange
|
||||
})
|
||||
|
||||
/** Closes an exchange without consuming remaining chunks (response interrupted or finished). */
|
||||
export const close = (id: string) =>
|
||||
Effect.suspend(() => {
|
||||
const exchange = state.exchanges.get(id)
|
||||
state.exchanges.delete(id)
|
||||
if (!exchange) return Effect.void
|
||||
return Queue.shutdown(exchange.queue).pipe(Effect.asVoid)
|
||||
})
|
||||
|
||||
/** Appends response chunks to an open exchange. Driver-facing. */
|
||||
export const push = (id: string, chunks: readonly Chunk[]) =>
|
||||
Effect.gen(function* () {
|
||||
const exchange = state.exchanges.get(id)
|
||||
if (!exchange) return yield* Effect.fail(new ExchangeNotFoundError(id))
|
||||
yield* Queue.offerAll(exchange.queue, chunks)
|
||||
})
|
||||
|
||||
/** Abruptly ends the provider body without a finish chunk or SSE sentinel. */
|
||||
export const disconnect = (id: string) =>
|
||||
Effect.gen(function* () {
|
||||
const exchange = state.exchanges.get(id)
|
||||
if (!exchange) return yield* Effect.fail(new ExchangeNotFoundError(id))
|
||||
yield* Queue.shutdown(exchange.queue)
|
||||
})
|
||||
|
||||
/**
|
||||
* Registers a listener for newly opened exchanges and immediately replays
|
||||
* currently-pending ones, so a late-attaching driver observes requests that
|
||||
* arrived before it connected. Returns an unsubscribe function.
|
||||
*/
|
||||
export function subscribe(listener: (exchange: OpenedExchange) => void) {
|
||||
state.listeners.add(listener)
|
||||
for (const exchange of pending()) listener(exchange)
|
||||
return () => {
|
||||
state.listeners.delete(listener)
|
||||
}
|
||||
}
|
||||
|
||||
/** Snapshot of currently open exchanges, for control-surface inspection. */
|
||||
export function pending(): OpenedExchange[] {
|
||||
return [...state.exchanges.values()].map((exchange) => ({
|
||||
id: exchange.id,
|
||||
url: exchange.url,
|
||||
body: exchange.body,
|
||||
}))
|
||||
}
|
||||
|
||||
export * as SimulationLLMExchange from "./llm-exchange"
|
||||
|
|
@ -1,7 +1,8 @@
|
|||
import { Effect, Layer } from "effect"
|
||||
import { HttpClient, HttpClientResponse } from "effect/unstable/http"
|
||||
import { Clock, Effect, Layer, Ref } from "effect"
|
||||
import { HttpClient, HttpClientResponse, type HttpMethod } from "effect/unstable/http"
|
||||
import { HttpClientError, TransportError } from "effect/unstable/http/HttpClientError"
|
||||
import type { HttpClientRequest } from "effect/unstable/http"
|
||||
import { SimulationProtocol } from "../protocol"
|
||||
|
||||
/**
|
||||
* Simulated network.
|
||||
|
|
@ -12,8 +13,7 @@ import type { HttpClientRequest } from "effect/unstable/http"
|
|||
* silently reach the real network. The scripted LLM is one registered route,
|
||||
* not a separate mechanism.
|
||||
*
|
||||
* The route table is process-global module state so the control surface and
|
||||
* the client layer observe the same registrations.
|
||||
* Each acquired run owns its routes and request log.
|
||||
*/
|
||||
|
||||
export interface Route {
|
||||
|
|
@ -21,33 +21,15 @@ export interface Route {
|
|||
readonly match: (
|
||||
request: HttpClientRequest.HttpClientRequest,
|
||||
url: URL,
|
||||
) => Effect.Effect<HttpClientResponse.HttpClientResponse> | undefined
|
||||
) => Effect.Effect<HttpClientResponse.HttpClientResponse, HttpClientError> | undefined
|
||||
}
|
||||
|
||||
interface LogEntry {
|
||||
readonly time: number
|
||||
readonly method: string
|
||||
readonly url: string
|
||||
readonly matched: boolean
|
||||
}
|
||||
|
||||
const state = {
|
||||
routes: [] as Route[],
|
||||
log: [] as LogEntry[],
|
||||
}
|
||||
export type LogEntry = SimulationProtocol.Backend.NetworkLogEntry
|
||||
|
||||
const LOG_LIMIT = 1000
|
||||
|
||||
export function register(route: Route) {
|
||||
state.routes.push(route)
|
||||
return () => {
|
||||
const index = state.routes.indexOf(route)
|
||||
if (index >= 0) state.routes.splice(index, 1)
|
||||
}
|
||||
}
|
||||
|
||||
/** Static JSON route: exact method + origin/path match answered with a fixed body. */
|
||||
export function json(method: string, url: string, body: unknown): Route {
|
||||
export function json(method: HttpMethod.HttpMethod, url: string, body: unknown): Route {
|
||||
return {
|
||||
match: (request, requestUrl) => {
|
||||
if (request.method !== method) return undefined
|
||||
|
|
@ -62,24 +44,29 @@ export function json(method: string, url: string, body: unknown): Route {
|
|||
}
|
||||
}
|
||||
|
||||
export function log(): readonly LogEntry[] {
|
||||
return state.log
|
||||
export interface Run {
|
||||
readonly client: HttpClient.HttpClient
|
||||
readonly log: () => Effect.Effect<readonly LogEntry[]>
|
||||
}
|
||||
|
||||
function record(entry: LogEntry) {
|
||||
state.log.push(entry)
|
||||
if (state.log.length > LOG_LIMIT) state.log.splice(0, state.log.length - LOG_LIMIT)
|
||||
}
|
||||
|
||||
export const layer = Layer.sync(HttpClient.HttpClient)(() =>
|
||||
HttpClient.make((request, url) =>
|
||||
Effect.suspend(() => {
|
||||
const matched = state.routes
|
||||
.map((route) => route.match(request, url))
|
||||
.find((response) => response !== undefined)
|
||||
record({ time: Date.now(), method: request.method, url: url.toString(), matched: matched !== undefined })
|
||||
if (matched) return matched
|
||||
return Effect.fail(
|
||||
export const make = Effect.fn("SimulationNetwork.make")(function* (routes: readonly Route[] = []) {
|
||||
const log = yield* Ref.make<readonly LogEntry[]>([])
|
||||
const client = HttpClient.make((request, url) =>
|
||||
Effect.gen(function* () {
|
||||
let matched: Effect.Effect<HttpClientResponse.HttpClientResponse, HttpClientError> | undefined
|
||||
for (const route of routes) {
|
||||
matched = route.match(request, url)
|
||||
if (matched) break
|
||||
}
|
||||
const entry = {
|
||||
time: yield* Clock.currentTimeMillis,
|
||||
method: request.method,
|
||||
url: url.toString(),
|
||||
matched: matched !== undefined,
|
||||
}
|
||||
yield* Ref.update(log, (entries) => [...entries, entry].slice(-LOG_LIMIT))
|
||||
if (matched) return yield* matched
|
||||
return yield* Effect.fail(
|
||||
new HttpClientError({
|
||||
reason: new TransportError({
|
||||
request,
|
||||
|
|
@ -88,7 +75,11 @@ export const layer = Layer.sync(HttpClient.HttpClient)(() =>
|
|||
}),
|
||||
)
|
||||
}),
|
||||
),
|
||||
)
|
||||
)
|
||||
return { client, log: () => Ref.get(log) } satisfies Run
|
||||
})
|
||||
|
||||
export const layer = (routes: readonly Route[] = []) =>
|
||||
Layer.effect(HttpClient.HttpClient, make(routes).pipe(Effect.map((run) => run.client)))
|
||||
|
||||
export * as SimulationNetwork from "./network"
|
||||
|
|
|
|||
|
|
@ -1,14 +1,15 @@
|
|||
import { Effect, Schema, Stream } from "effect"
|
||||
import { HttpClientResponse } from "effect/unstable/http"
|
||||
import { HttpClientError, TransportError } from "effect/unstable/http/HttpClientError"
|
||||
import { OpenAIChatEvent, DEFAULT_BASE_URL, PATH } from "@opencode-ai/llm/protocols/openai-chat"
|
||||
import { SimulationLLMExchange } from "./llm-exchange"
|
||||
import { SimulationNetwork } from "./network"
|
||||
import { SimulatedProvider } from "./simulated-provider"
|
||||
|
||||
/**
|
||||
* Driver-answered OpenAI endpoint for the simulated network.
|
||||
*
|
||||
* Claims `POST {DEFAULT_BASE_URL}{PATH}` (the real openai-chat route
|
||||
* endpoint), opens an LLM exchange, and streams the driver's chunks back as
|
||||
* endpoint), invokes the simulated provider, and streams the driver's events back as
|
||||
* an OpenAI Chat SSE response terminated by `[DONE]`. Everything downstream
|
||||
* of the response bytes is the real pipeline: SSE framing, the OpenAIChat
|
||||
* event schema, the protocol state machine, and Lifecycle grammar.
|
||||
|
|
@ -17,11 +18,15 @@ import { SimulationNetwork } from "./network"
|
|||
const encodeChunk = Schema.encodeUnknownSync(OpenAIChatEvent)
|
||||
|
||||
const encoder = new TextEncoder()
|
||||
const decodeBody = Schema.decodeUnknownEffect(Schema.fromJsonString(Schema.Json))
|
||||
|
||||
// The simulated model id is echoed back only in non-schema fields; the
|
||||
// protocol event schema ignores unknown fields, so id/object/model are
|
||||
// decorative wire realism.
|
||||
function chunkOf(item: SimulationLLMExchange.Item): OpenAIChatEvent | unknown {
|
||||
type ProviderItem = Exclude<SimulatedProvider.ProviderResponseEvent, { readonly type: "finish" }>
|
||||
type FinishReason = Extract<SimulatedProvider.ProviderResponseEvent, { readonly type: "finish" }>["reason"]
|
||||
|
||||
function chunkOf(item: ProviderItem): OpenAIChatEvent | unknown {
|
||||
if (item.type === "textDelta") return { choices: [{ delta: { content: item.text } }] }
|
||||
if (item.type === "reasoningDelta") return { choices: [{ delta: { reasoning_content: item.text } }] }
|
||||
if (item.type === "toolCall")
|
||||
|
|
@ -43,7 +48,7 @@ function chunkOf(item: SimulationLLMExchange.Item): OpenAIChatEvent | unknown {
|
|||
return item.chunk
|
||||
}
|
||||
|
||||
const finishReasonWire: Record<SimulationLLMExchange.FinishReason, string> = {
|
||||
const finishReasonWire: Record<FinishReason, string> = {
|
||||
stop: "stop",
|
||||
"tool-calls": "tool_calls",
|
||||
length: "length",
|
||||
|
|
@ -54,40 +59,49 @@ function frame(payload: unknown): Uint8Array {
|
|||
return encoder.encode(`data: ${JSON.stringify(payload)}\n\n`)
|
||||
}
|
||||
|
||||
function sseBody(exchange: SimulationLLMExchange.Exchange): Stream.Stream<Uint8Array> {
|
||||
const chunks = Stream.fromQueue(exchange.queue).pipe(
|
||||
Stream.takeUntil((chunk) => chunk.type === "finish"),
|
||||
Stream.map((chunk) => {
|
||||
if (chunk.type === "finish")
|
||||
return frame(encodeChunk({ choices: [{ delta: {}, finish_reason: finishReasonWire[chunk.reason] }] }))
|
||||
if (chunk.item.type === "raw") return frame(chunk.item.chunk)
|
||||
return frame(encodeChunk(chunkOf(chunk.item)))
|
||||
function sseBody(
|
||||
events: Stream.Stream<SimulatedProvider.ProviderResponseEvent, SimulatedProvider.ProviderDisconnectedError>,
|
||||
): Stream.Stream<Uint8Array, SimulatedProvider.ProviderDisconnectedError> {
|
||||
return events.pipe(
|
||||
Stream.map((event) => {
|
||||
if (event.type === "finish")
|
||||
return frame(encodeChunk({ choices: [{ delta: {}, finish_reason: finishReasonWire[event.reason] }] }))
|
||||
if (event.type === "raw") return frame(event.chunk)
|
||||
return frame(encodeChunk(chunkOf(event)))
|
||||
}),
|
||||
)
|
||||
return chunks.pipe(
|
||||
Stream.concat(Stream.make(encoder.encode("data: [DONE]\n\n"))),
|
||||
// Close the exchange when the response body ends or is interrupted, so
|
||||
// late driver pushes fail with ExchangeNotFoundError instead of leaking.
|
||||
Stream.ensuring(SimulationLLMExchange.close(exchange.id)),
|
||||
)
|
||||
}
|
||||
|
||||
export const route: SimulationNetwork.Route = {
|
||||
export const route = (provider: SimulatedProvider.Interface): SimulationNetwork.Route => ({
|
||||
match: (request, url) => {
|
||||
if (request.method !== "POST") return undefined
|
||||
if (url.origin + url.pathname !== DEFAULT_BASE_URL + PATH) return undefined
|
||||
return Effect.gen(function* () {
|
||||
const body = request.body._tag === "Uint8Array" ? JSON.parse(new TextDecoder().decode(request.body.body)) : {}
|
||||
const exchange = yield* SimulationLLMExchange.open({ url: url.toString(), body })
|
||||
const body =
|
||||
request.body._tag === "Uint8Array"
|
||||
? yield* decodeBody(new TextDecoder().decode(request.body.body)).pipe(
|
||||
Effect.mapError(
|
||||
(cause) =>
|
||||
new HttpClientError({
|
||||
reason: new TransportError({
|
||||
request,
|
||||
cause,
|
||||
description: "Simulation received an invalid OpenAI request body",
|
||||
}),
|
||||
}),
|
||||
),
|
||||
)
|
||||
: {}
|
||||
return HttpClientResponse.fromWeb(
|
||||
request,
|
||||
new Response(Stream.toReadableStream(sseBody(exchange)), {
|
||||
new Response(Stream.toReadableStream(sseBody(provider.stream({ url: url.toString(), body }))), {
|
||||
status: 200,
|
||||
headers: { "content-type": "text/event-stream" },
|
||||
}),
|
||||
)
|
||||
})
|
||||
},
|
||||
}
|
||||
})
|
||||
|
||||
export * as SimulationOpenAI from "./openai"
|
||||
|
|
|
|||
270
packages/simulation/src/backend/simulated-provider.ts
Normal file
270
packages/simulation/src/backend/simulated-provider.ts
Normal file
|
|
@ -0,0 +1,270 @@
|
|||
import { Cause, Context, Effect, Fiber, FiberSet, Layer, PubSub, Queue, Ref, Schema, Semaphore, Stream } from "effect"
|
||||
import { SimulationControlServer } from "../control-server"
|
||||
import { SimulationProtocol } from "../protocol"
|
||||
|
||||
export interface ProviderRequest {
|
||||
readonly url: string
|
||||
readonly body: unknown
|
||||
}
|
||||
|
||||
export type ProviderResponseEvent =
|
||||
| SimulationProtocol.Backend.Item
|
||||
| { readonly type: "finish"; readonly reason: SimulationProtocol.Backend.FinishReason }
|
||||
|
||||
export class ProviderDisconnectedError extends Schema.TaggedErrorClass<ProviderDisconnectedError>()(
|
||||
"SimulatedProvider.ProviderDisconnectedError",
|
||||
{ message: Schema.String },
|
||||
) {}
|
||||
|
||||
export interface Interface {
|
||||
readonly stream: (request: ProviderRequest) => Stream.Stream<ProviderResponseEvent, ProviderDisconnectedError>
|
||||
}
|
||||
|
||||
export class Service extends Context.Service<Service, Interface>()("@opencode/simulation/SimulatedProvider") {}
|
||||
|
||||
interface ProviderInvocation extends ProviderRequest {
|
||||
readonly id: string
|
||||
}
|
||||
|
||||
interface PendingInvocation extends ProviderInvocation {
|
||||
readonly responses: Queue.Queue<ProviderResponseEvent, ProviderDisconnectedError | Cause.Done>
|
||||
}
|
||||
|
||||
interface State {
|
||||
readonly counter: number
|
||||
readonly pending: ReadonlyMap<string, PendingInvocation>
|
||||
}
|
||||
|
||||
interface Driver {
|
||||
readonly requests: Stream.Stream<ProviderInvocation>
|
||||
readonly push: (
|
||||
id: string,
|
||||
items: readonly SimulationProtocol.Backend.Item[],
|
||||
) => Effect.Effect<void, InvocationNotFoundError>
|
||||
readonly finish: (
|
||||
id: string,
|
||||
reason: SimulationProtocol.Backend.FinishReason,
|
||||
) => Effect.Effect<void, InvocationNotFoundError>
|
||||
readonly disconnect: (id: string) => Effect.Effect<void, InvocationNotFoundError>
|
||||
readonly pending: () => Effect.Effect<readonly ProviderInvocation[]>
|
||||
}
|
||||
|
||||
class InvocationNotFoundError extends Schema.TaggedErrorClass<InvocationNotFoundError>()(
|
||||
"SimulatedProvider.InvocationNotFoundError",
|
||||
{ id: Schema.String, message: Schema.String },
|
||||
) {}
|
||||
|
||||
class ControllerDisconnectedError extends Schema.TaggedErrorClass<ControllerDisconnectedError>()(
|
||||
"SimulatedProvider.ControllerDisconnectedError",
|
||||
{ message: Schema.String },
|
||||
) {}
|
||||
|
||||
type ControlSocket = SimulationControlServer.Socket
|
||||
|
||||
export const layerDrive = (options: { readonly endpoint: string }) =>
|
||||
Layer.effect(
|
||||
Service,
|
||||
Effect.gen(function* () {
|
||||
const state = yield* Ref.make<State>({ counter: 0, pending: new Map() })
|
||||
const opened = yield* PubSub.unbounded<ProviderInvocation>()
|
||||
const lock = yield* Semaphore.make(1)
|
||||
|
||||
const close = (invocation: PendingInvocation) =>
|
||||
Effect.gen(function* () {
|
||||
yield* Queue.shutdown(invocation.responses)
|
||||
yield* lock.withPermit(
|
||||
Ref.update(state, (current) =>
|
||||
current.pending.get(invocation.id) === invocation ? remove(current, invocation.id) : current,
|
||||
),
|
||||
)
|
||||
})
|
||||
|
||||
yield* Effect.addFinalizer(() =>
|
||||
Effect.gen(function* () {
|
||||
const current = yield* Ref.get(state)
|
||||
yield* Effect.forEach(current.pending.values(), (invocation) => Queue.shutdown(invocation.responses), {
|
||||
discard: true,
|
||||
})
|
||||
yield* PubSub.shutdown(opened)
|
||||
}),
|
||||
)
|
||||
|
||||
const open = (request: ProviderRequest) =>
|
||||
lock.withPermit(
|
||||
Effect.gen(function* () {
|
||||
const current = yield* Ref.get(state)
|
||||
const id = `inv_${current.counter + 1}`
|
||||
const responses = yield* Queue.bounded<ProviderResponseEvent, ProviderDisconnectedError | Cause.Done>(256)
|
||||
const invocation: PendingInvocation = { id, ...request, responses }
|
||||
yield* Ref.set(state, {
|
||||
counter: current.counter + 1,
|
||||
pending: new Map(current.pending).set(id, invocation),
|
||||
})
|
||||
yield* PubSub.publish(opened, { id, ...request })
|
||||
return invocation
|
||||
}),
|
||||
)
|
||||
|
||||
const requireInvocation = (id: string) =>
|
||||
Effect.gen(function* () {
|
||||
const current = yield* Ref.get(state)
|
||||
const invocation = current.pending.get(id)
|
||||
if (invocation) return invocation
|
||||
return yield* Effect.fail(
|
||||
new InvocationNotFoundError({
|
||||
id,
|
||||
message: `Simulated provider invocation not found or already finished: ${id}`,
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
const remove = (current: State, id: string) => {
|
||||
const pending = new Map(current.pending)
|
||||
pending.delete(id)
|
||||
return { ...current, pending }
|
||||
}
|
||||
|
||||
const driver: Driver = {
|
||||
requests: Stream.unwrap(
|
||||
lock.withPermit(
|
||||
Effect.gen(function* () {
|
||||
const subscription = yield* PubSub.subscribe(opened)
|
||||
const current = yield* Ref.get(state)
|
||||
const pending = Array.from(current.pending.values(), ({ id, url, body }) => ({ id, url, body }))
|
||||
return Stream.concat(Stream.fromIterable(pending), Stream.fromEffectRepeat(PubSub.take(subscription)))
|
||||
}),
|
||||
),
|
||||
),
|
||||
push: (id, items) =>
|
||||
Effect.gen(function* () {
|
||||
const invocation = yield* lock.withPermit(requireInvocation(id))
|
||||
yield* Queue.offerAll(invocation.responses, items)
|
||||
}),
|
||||
finish: (id, reason) =>
|
||||
Effect.gen(function* () {
|
||||
const invocation = yield* lock.withPermit(
|
||||
Effect.gen(function* () {
|
||||
const invocation = yield* requireInvocation(id)
|
||||
const current = yield* Ref.get(state)
|
||||
yield* Ref.set(state, remove(current, id))
|
||||
return invocation
|
||||
}),
|
||||
)
|
||||
yield* Queue.offer(invocation.responses, { type: "finish", reason })
|
||||
yield* Queue.end(invocation.responses)
|
||||
}),
|
||||
disconnect: (id) =>
|
||||
Effect.gen(function* () {
|
||||
const invocation = yield* lock.withPermit(
|
||||
Effect.gen(function* () {
|
||||
const invocation = yield* requireInvocation(id)
|
||||
const current = yield* Ref.get(state)
|
||||
yield* Ref.set(state, remove(current, id))
|
||||
return invocation
|
||||
}),
|
||||
)
|
||||
yield* Queue.fail(
|
||||
invocation.responses,
|
||||
new ProviderDisconnectedError({ message: "Simulated model provider disconnected" }),
|
||||
)
|
||||
}),
|
||||
pending: () =>
|
||||
lock.withPermit(
|
||||
Ref.get(state).pipe(
|
||||
Effect.map((current) => Array.from(current.pending.values(), ({ id, url, body }) => ({ id, url, body }))),
|
||||
),
|
||||
),
|
||||
}
|
||||
|
||||
const fibers = yield* FiberSet.make<void, unknown>()
|
||||
const activeController = yield* Ref.make<Fiber.Fiber<void> | undefined>(undefined)
|
||||
const controllerLock = yield* Semaphore.make(1)
|
||||
yield* SimulationControlServer.start({
|
||||
endpoint: options.endpoint,
|
||||
label: "opencode drive backend websocket",
|
||||
data: () => ({}),
|
||||
decode: SimulationProtocol.Backend.decodeRequestEffect,
|
||||
handle: (socket, request) => handle(driver, fibers, activeController, controllerLock, socket, request),
|
||||
close: (socket) => releaseController(activeController, controllerLock, socket),
|
||||
})
|
||||
yield* Effect.sync(() => process.stderr.write(`opencode drive backend websocket: ${options.endpoint}\n`))
|
||||
|
||||
return Service.of({
|
||||
stream: (request) =>
|
||||
Stream.unwrap(
|
||||
Effect.acquireRelease(open(request), close).pipe(
|
||||
Effect.map((invocation) =>
|
||||
Stream.fromQueue(invocation.responses).pipe(Stream.takeUntil((event) => event.type === "finish")),
|
||||
),
|
||||
),
|
||||
),
|
||||
})
|
||||
}),
|
||||
)
|
||||
|
||||
function handle(
|
||||
driver: Driver,
|
||||
fibers: FiberSet.FiberSet<void, unknown>,
|
||||
activeController: Ref.Ref<Fiber.Fiber<void> | undefined>,
|
||||
controllerLock: Semaphore.Semaphore,
|
||||
socket: ControlSocket,
|
||||
request: SimulationProtocol.Backend.Request,
|
||||
) {
|
||||
switch (request.method) {
|
||||
case "llm.attach":
|
||||
return controllerLock.withPermit(
|
||||
Effect.gen(function* () {
|
||||
if (socket.data.closed)
|
||||
return yield* Effect.fail(
|
||||
new ControllerDisconnectedError({ message: "Drive controller disconnected before attachment" }),
|
||||
)
|
||||
const previous = yield* Ref.get(activeController)
|
||||
if (previous) yield* Fiber.interrupt(previous)
|
||||
const attachment = yield* FiberSet.run(
|
||||
fibers,
|
||||
driver.requests.pipe(
|
||||
Stream.runForEach((invocation) =>
|
||||
Effect.sync(() => {
|
||||
socket.send(JSON.stringify({ jsonrpc: "2.0", method: "llm.request", params: invocation }))
|
||||
}),
|
||||
),
|
||||
),
|
||||
)
|
||||
if (socket.data.closed) {
|
||||
yield* Fiber.interrupt(attachment)
|
||||
return yield* Effect.fail(
|
||||
new ControllerDisconnectedError({ message: "Drive controller disconnected during attachment" }),
|
||||
)
|
||||
}
|
||||
socket.data.attachment = attachment
|
||||
yield* Ref.set(activeController, attachment)
|
||||
return { attached: true }
|
||||
}),
|
||||
)
|
||||
case "llm.chunk":
|
||||
return driver.push(request.params.id, request.params.items).pipe(Effect.as({ ok: true }))
|
||||
case "llm.finish":
|
||||
return driver.finish(request.params.id, request.params.reason).pipe(Effect.as({ ok: true }))
|
||||
case "llm.disconnect":
|
||||
return driver.disconnect(request.params.id).pipe(Effect.as({ ok: true }))
|
||||
case "llm.pending":
|
||||
return driver.pending().pipe(Effect.map((invocations) => ({ invocations })))
|
||||
}
|
||||
}
|
||||
|
||||
function releaseController(
|
||||
activeController: Ref.Ref<Fiber.Fiber<void> | undefined>,
|
||||
controllerLock: Semaphore.Semaphore,
|
||||
socket: ControlSocket,
|
||||
) {
|
||||
return controllerLock.withPermit(
|
||||
Effect.gen(function* () {
|
||||
const attachment = socket.data.attachment
|
||||
if (!attachment) return
|
||||
yield* Fiber.interrupt(attachment)
|
||||
yield* Ref.update(activeController, (active) => (active === attachment ? undefined : active))
|
||||
}),
|
||||
)
|
||||
}
|
||||
|
||||
export * as SimulatedProvider from "./simulated-provider"
|
||||
Loading…
Add table
Add a link
Reference in a new issue