feat(simulation): share control protocol schemas (#35230)

packages/simulation/package.json

@@ -9,7 +9,8 @@
     "./backend": "./src/backend/index.ts",
     "./backend/*": "./src/backend/*.ts",
     "./frontend": "./src/frontend/simulation.ts",
-    "./frontend/*": "./src/frontend/*.ts"
+    "./frontend/*": "./src/frontend/*.ts",
+    "./protocol": "./src/protocol/index.ts"
   },
   "scripts": {
     "typecheck": "tsgo --noEmit"

packages/simulation/src/backend/control.ts

@@ -1,4 +1,5 @@
-import { Effect, Schema } from "effect"
+import { Effect } from "effect"
+import { SimulationProtocol } from "../protocol"
 import { SimulationLLMExchange } from "./llm-exchange"
 import { SimulationNetwork } from "./network"
 
@@ -23,43 +24,13 @@ import { SimulationNetwork } from "./network"
 const DefaultPort = 40950
 const MaxPortAttempts = 100
 
-const ChunkItem = Schema.Union([
-  Schema.Struct({ type: Schema.Literal("textDelta"), text: Schema.String }),
-  Schema.Struct({ type: Schema.Literal("reasoningDelta"), text: Schema.String }),
-  Schema.Struct({ type: Schema.Literal("toolCall"), id: Schema.String, name: Schema.String, input: Schema.Unknown }),
-  Schema.Struct({ type: Schema.Literal("raw"), chunk: Schema.Unknown }),
-])
-
-const ChunkParams = Schema.Struct({ id: Schema.String, items: Schema.Array(ChunkItem) })
-
-const FinishParams = Schema.Struct({
-  id: Schema.String,
-  reason: Schema.Literals(["stop", "tool-calls", "length", "content-filter"]).pipe(
-    Schema.withDecodingDefault(Effect.succeed("stop" as const)),
-  ),
-})
-
-const decodeChunkParams = Schema.decodeUnknownPromise(ChunkParams)
-const decodeFinishParams = Schema.decodeUnknownPromise(FinishParams)
-
-type JsonRpcRequest = {
-  readonly jsonrpc: "2.0"
-  readonly id?: string | number | null
-  readonly method: string
-  readonly params?: unknown
-}
-
 type ControlSocket = Bun.ServerWebSocket<{ unsubscribe?: () => void }>
 
-function parseRequest(input: string | Buffer): JsonRpcRequest {
-  const value = JSON.parse(typeof input === "string" ? input : input.toString()) as unknown
-  if (typeof value !== "object" || value === null) throw new Error("Invalid JSON-RPC request")
-  if (!("jsonrpc" in value) || value.jsonrpc !== "2.0") throw new Error("Invalid JSON-RPC version")
-  if (!("method" in value) || typeof value.method !== "string") throw new Error("Invalid JSON-RPC method")
-  return value as JsonRpcRequest
+function parseRequest(input: string | Buffer) {
+  return SimulationProtocol.JsonRpc.decodeRequest(JSON.parse(typeof input === "string" ? input : input.toString()))
 }
 
-async function handle(socket: ControlSocket, request: JsonRpcRequest): Promise<unknown> {
+async function handle(socket: ControlSocket, request: SimulationProtocol.JsonRpc.Request): Promise<unknown> {
   switch (request.method) {
     case "llm.attach": {
       socket.data.unsubscribe?.()
@@ -69,7 +40,7 @@ async function handle(socket: ControlSocket, request: JsonRpcRequest): Promise<u
       return { attached: true }
     }
     case "llm.chunk": {
-      const params = await decodeChunkParams(request.params)
+      const params = await SimulationProtocol.Backend.decodeChunkParams(request.params)
       await Effect.runPromise(
         SimulationLLMExchange.push(
           params.id,
@@ -79,7 +50,7 @@ async function handle(socket: ControlSocket, request: JsonRpcRequest): Promise<u
       return { ok: true }
     }
     case "llm.finish": {
-      const params = await decodeFinishParams(request.params)
+      const params = await SimulationProtocol.Backend.decodeFinishParams(request.params)
       await Effect.runPromise(SimulationLLMExchange.push(params.id, [{ type: "finish", reason: params.reason }]))
       return { ok: true }
     }
@@ -105,19 +76,14 @@ function serve(port = DefaultPort, attempts = MaxPortAttempts): Bun.Server<{ uns
           socket.data.unsubscribe?.()
         },
         async message(socket, message) {
-          let request: JsonRpcRequest | undefined
+          let request: SimulationProtocol.JsonRpc.Request | undefined
           try {
             request = parseRequest(message)
             const result = await handle(socket, request)
-            if (request.id !== undefined) socket.send(JSON.stringify({ jsonrpc: "2.0", id: request.id, result }))
+            const response = SimulationProtocol.JsonRpc.success(request.id, result)
+            if (response) socket.send(JSON.stringify(response))
           } catch (error) {
-            socket.send(
-              JSON.stringify({
-                jsonrpc: "2.0",
-                id: request?.id ?? null,
-                error: { code: -32000, message: error instanceof Error ? error.message : String(error) },
-              }),
-            )
+            socket.send(JSON.stringify(SimulationProtocol.JsonRpc.failure(request?.id, error)))
           }
         },
       },

packages/simulation/src/frontend/actions.ts

@@ -1,36 +1,11 @@
 import type { CliRenderer, Renderable } from "@opentui/core"
 import { createMockKeys, createMockMouse, type MockInput, type MockMouse } from "@opentui/core/testing"
+import type { SimulationProtocol } from "../protocol"
 import { SimulationRenderer } from "./renderer"
 import { SimulationTrace } from "./trace"
 
-export interface KeyModifiers {
-  readonly ctrl?: boolean
-  readonly shift?: boolean
-  readonly meta?: boolean
-  readonly super?: boolean
-  readonly hyper?: boolean
-}
-
-export type Action =
-  | { readonly type: "typeText"; readonly text: string }
-  | { readonly type: "pressKey"; readonly key: string; readonly modifiers?: KeyModifiers }
-  | { readonly type: "pressEnter" }
-  | { readonly type: "pressArrow"; readonly direction: "up" | "down" | "left" | "right" }
-  | { readonly type: "focus"; readonly target: number }
-  | { readonly type: "click"; readonly target: number; readonly x: number; readonly y: number }
-
-export interface Element {
-  readonly id: string
-  readonly num: number
-  readonly x: number
-  readonly y: number
-  readonly width: number
-  readonly height: number
-  readonly focusable: boolean
-  readonly focused: boolean
-  readonly clickable: boolean
-  readonly editor: boolean
-}
+export type Action = SimulationProtocol.Frontend.Action
+export type Element = SimulationProtocol.Frontend.Element
 
 export interface Harness {
   readonly renderer: CliRenderer

packages/simulation/src/frontend/server.ts

@@ -1,27 +1,10 @@
-import { SimulationActions, type Action, type Harness } from "./actions"
+import { SimulationProtocol } from "../protocol"
+import { SimulationActions, type Harness } from "./actions"
 import { SimulationTrace } from "./trace"
 
 const DefaultPort = 40900
 const MaxPortAttempts = 100
 
-type JsonRpcRequest = {
-  readonly jsonrpc: "2.0"
-  readonly id?: string | number | null
-  readonly method: string
-  readonly params?: unknown
-}
-
-type JsonRpcResponse = {
-  readonly jsonrpc: "2.0"
-  readonly id: string | number | null
-  readonly result?: unknown
-  readonly error?: {
-    readonly code: number
-    readonly message: string
-    readonly data?: unknown
-  }
-}
-
 export interface Server {
   readonly url: string
   readonly stop: () => void
@@ -36,63 +19,15 @@ function isPortUnavailable(error: unknown) {
   return message.includes("eaddrinuse") || message.includes("address already in use") || message.includes(" in use")
 }
 
-function parseRequest(input: string | Buffer): JsonRpcRequest {
-  const value = JSON.parse(typeof input === "string" ? input : input.toString()) as unknown
-  if (typeof value !== "object" || value === null) throw new Error("Invalid JSON-RPC request")
-  if (!("jsonrpc" in value) || value.jsonrpc !== "2.0") throw new Error("Invalid JSON-RPC version")
-  if (!("method" in value) || typeof value.method !== "string") throw new Error("Invalid JSON-RPC method")
-  return value as JsonRpcRequest
-}
-
-function isAction(input: unknown): input is Action {
-  if (typeof input !== "object" || input === null || !("type" in input)) return false
-  switch (input.type) {
-    case "typeText":
-      return "text" in input && typeof input.text === "string"
-    case "pressKey":
-      return "key" in input && typeof input.key === "string"
-    case "pressEnter":
-      return true
-    case "pressArrow":
-      return "direction" in input && ["up", "down", "left", "right"].includes(String(input.direction))
-    case "focus":
-      return "target" in input && typeof input.target === "number"
-    case "click":
-      return (
-        "target" in input &&
-        typeof input.target === "number" &&
-        "x" in input &&
-        typeof input.x === "number" &&
-        "y" in input &&
-        typeof input.y === "number"
-      )
-  }
-  return false
-}
-
 function actionParam(params: unknown) {
-  if (typeof params !== "object" || params === null || !("action" in params)) throw new Error("Missing action")
-  if (!isAction(params.action)) throw new Error("Invalid action")
-  return params.action
+  return SimulationProtocol.Frontend.decodeActionParams(params).action
 }
 
-function response(id: JsonRpcRequest["id"], result: unknown): JsonRpcResponse | undefined {
-  if (id === undefined) return undefined
-  return { jsonrpc: "2.0", id, result }
+function parseRequest(input: string | Buffer) {
+  return SimulationProtocol.JsonRpc.decodeRequest(JSON.parse(typeof input === "string" ? input : input.toString()))
 }
 
-function errorResponse(id: JsonRpcRequest["id"], error: unknown): JsonRpcResponse {
-  return {
-    jsonrpc: "2.0",
-    id: id ?? null,
-    error: {
-      code: -32000,
-      message: error instanceof Error ? error.message : String(error),
-    },
-  }
-}
-
-async function handle(harness: Harness, request: JsonRpcRequest) {
+async function handle(harness: Harness, request: SimulationProtocol.JsonRpc.Request) {
   switch (request.method) {
     case "ui.state": {
       const result = SimulationActions.state(harness)
@@ -139,14 +74,14 @@ function serve(
           SimulationTrace.add("control.disconnect")
         },
         async message(socket, message) {
-          let request: JsonRpcRequest | undefined
+          let request: SimulationProtocol.JsonRpc.Request | undefined
           try {
             request = parseRequest(message)
             const result = await handle(harness, request)
-            const next = response(request.id, result)
+            const next = SimulationProtocol.JsonRpc.success(request.id, result)
             if (next) socket.send(JSON.stringify(next))
           } catch (error) {
-            socket.send(JSON.stringify(errorResponse(request?.id, error)))
+            socket.send(JSON.stringify(SimulationProtocol.JsonRpc.failure(request?.id, error)))
           }
         },
       },

packages/simulation/src/protocol/index.ts

@@ -0,0 +1,145 @@
+import { Effect, Schema } from "effect"
+
+const JsonRpcID = Schema.Union([Schema.String, Schema.Number, Schema.Null])
+type Json = Schema.Schema.Type<typeof Schema.Json>
+
+export namespace JsonRpc {
+  export const Request = Schema.Struct({
+    jsonrpc: Schema.Literal("2.0"),
+    id: Schema.optional(JsonRpcID),
+    method: Schema.String,
+    params: Schema.optional(Schema.Json),
+  })
+  export interface Request extends Schema.Schema.Type<typeof Request> {}
+
+  export const ErrorObject = Schema.Struct({
+    code: Schema.Number,
+    message: Schema.String,
+    data: Schema.optional(Schema.Json),
+  })
+
+  export const Response = Schema.Struct({
+    jsonrpc: Schema.Literal("2.0"),
+    id: JsonRpcID,
+    result: Schema.optional(Schema.Json),
+    error: Schema.optional(ErrorObject),
+  })
+  export interface Response extends Schema.Schema.Type<typeof Response> {}
+
+  export const decodeRequest = Schema.decodeUnknownSync(Request)
+
+  export function success(id: Request["id"], result: unknown): Response | undefined {
+    if (id === undefined) return undefined
+    return { jsonrpc: "2.0", id, result: result as Json }
+  }
+
+  export function failure(id: Request["id"], error: unknown): Response {
+    return {
+      jsonrpc: "2.0",
+      id: id ?? null,
+      error: {
+        code: -32000,
+        message: error instanceof Error ? error.message : String(error),
+      },
+    }
+  }
+}
+
+export namespace Frontend {
+  export const KeyModifiers = Schema.Struct({
+    ctrl: Schema.optional(Schema.Boolean),
+    shift: Schema.optional(Schema.Boolean),
+    meta: Schema.optional(Schema.Boolean),
+    super: Schema.optional(Schema.Boolean),
+    hyper: Schema.optional(Schema.Boolean),
+  })
+  export interface KeyModifiers extends Schema.Schema.Type<typeof KeyModifiers> {}
+
+  export const Action = Schema.Union([
+    Schema.Struct({ type: Schema.Literal("typeText"), text: Schema.String }),
+    Schema.Struct({ type: Schema.Literal("pressKey"), key: Schema.String, modifiers: Schema.optional(KeyModifiers) }),
+    Schema.Struct({ type: Schema.Literal("pressEnter") }),
+    Schema.Struct({ type: Schema.Literal("pressArrow"), direction: Schema.Literals(["up", "down", "left", "right"]) }),
+    Schema.Struct({ type: Schema.Literal("focus"), target: Schema.Number }),
+    Schema.Struct({ type: Schema.Literal("click"), target: Schema.Number, x: Schema.Number, y: Schema.Number }),
+  ])
+  export type Action = Schema.Schema.Type<typeof Action>
+
+  export const Element = Schema.Struct({
+    id: Schema.String,
+    num: Schema.Number,
+    x: Schema.Number,
+    y: Schema.Number,
+    width: Schema.Number,
+    height: Schema.Number,
+    focusable: Schema.Boolean,
+    focused: Schema.Boolean,
+    clickable: Schema.Boolean,
+    editor: Schema.Boolean,
+  })
+  export interface Element extends Schema.Schema.Type<typeof Element> {}
+
+  export const State = Schema.Struct({
+    screen: Schema.String,
+    focused: Schema.Struct({
+      renderable: Schema.optional(Schema.Number),
+      editor: Schema.Boolean,
+    }),
+    elements: Schema.Array(Element),
+    actions: Schema.Array(Action),
+  })
+  export interface State extends Schema.Schema.Type<typeof State> {}
+
+  export const ActionParams = Schema.Struct({ action: Action })
+  export interface ActionParams extends Schema.Schema.Type<typeof ActionParams> {}
+  export const decodeActionParams = Schema.decodeUnknownSync(ActionParams)
+
+  export const TraceRecord = Schema.Struct({
+    id: Schema.Number,
+    time: Schema.String,
+    type: Schema.String,
+    data: Schema.optional(Schema.Json),
+  })
+  export interface TraceRecord extends Schema.Schema.Type<typeof TraceRecord> {}
+
+  export const TraceList = Schema.Struct({ records: Schema.Array(TraceRecord) })
+  export interface TraceList extends Schema.Schema.Type<typeof TraceList> {}
+}
+
+export namespace Backend {
+  export const Item = Schema.Union([
+    Schema.Struct({ type: Schema.Literal("textDelta"), text: Schema.String }),
+    Schema.Struct({ type: Schema.Literal("reasoningDelta"), text: Schema.String }),
+    Schema.Struct({ type: Schema.Literal("toolCall"), id: Schema.String, name: Schema.String, input: Schema.Json }),
+    Schema.Struct({ type: Schema.Literal("raw"), chunk: Schema.Json }),
+  ])
+  export type Item = Schema.Schema.Type<typeof Item>
+
+  export const FinishReason = Schema.Literals(["stop", "tool-calls", "length", "content-filter"])
+  export type FinishReason = Schema.Schema.Type<typeof FinishReason>
+
+  export const ChunkParams = Schema.Struct({ id: Schema.String, items: Schema.Array(Item) })
+  export interface ChunkParams extends Schema.Schema.Type<typeof ChunkParams> {}
+
+  export const FinishParams = Schema.Struct({
+    id: Schema.String,
+    reason: FinishReason.pipe(Schema.withDecodingDefault(Effect.succeed("stop" as const))),
+  })
+  export interface FinishParams extends Schema.Schema.Type<typeof FinishParams> {}
+
+  export const OpenedExchange = Schema.Struct({ id: Schema.String, url: Schema.String, body: Schema.Json })
+  export interface OpenedExchange extends Schema.Schema.Type<typeof OpenedExchange> {}
+
+  export const NetworkLogEntry = Schema.Struct({
+    time: Schema.Number,
+    method: Schema.String,
+    url: Schema.String,
+    matched: Schema.Boolean,
+  })
+  export interface NetworkLogEntry extends Schema.Schema.Type<typeof NetworkLogEntry> {}
+
+  export const decodeChunkParams = Schema.decodeUnknownPromise(ChunkParams)
+  export const decodeFinishParams = Schema.decodeUnknownPromise(FinishParams)
+}
+
+export * as SimulationProtocol from "./index"