refactor(protocol): harden simulation wire contract (#42628)

packages/protocol/src/simulation.ts

@@ -22,13 +22,24 @@ export namespace JsonRpc {
     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 Response = Schema.Union(
+    [
+      Schema.Struct({
+        jsonrpc: Schema.Literal("2.0"),
+        id: JsonRpcID,
+        result: Schema.Json,
+        error: Schema.optionalKey(Schema.Never),
+      }),
+      Schema.Struct({
+        jsonrpc: Schema.Literal("2.0"),
+        id: JsonRpcID,
+        result: Schema.optionalKey(Schema.Never),
+        error: ErrorObject,
+      }),
+    ],
+    { mode: "oneOf" },
+  )
+  export type Response = Schema.Schema.Type<typeof Response>
 
   export const decodeRequest = Schema.decodeUnknownSync(Request)
 
@@ -49,6 +60,28 @@ export namespace JsonRpc {
   }
 }
 
+export class SimulationRequestError extends Schema.TaggedErrorClass<SimulationRequestError>()(
+  "SimulationRequestError",
+  {
+    method: Schema.String,
+    code: Schema.Number,
+    message: Schema.String,
+    data: Schema.optionalKey(Schema.Json),
+  },
+) {}
+
+const request = <
+  const Tag extends string,
+  Payload extends Schema.Top | Schema.Struct.Fields = typeof Schema.Void,
+  Success extends Schema.Top = typeof Schema.Void,
+>(
+  tag: Tag,
+  options?: {
+    readonly payload?: Payload
+    readonly success?: Success
+  },
+) => Rpc.make(tag, { ...options, error: SimulationRequestError })
+
 export namespace Handshake {
   export const ProtocolVersion = Schema.Literal(1)
   export type ProtocolVersion = Schema.Schema.Type<typeof ProtocolVersion>
@@ -81,7 +114,7 @@ export namespace Handshake {
     protocolVersion: ProtocolVersion,
     role: EndpointRole,
     server: Identity,
-    capabilities: Schema.Array(Capability),
+    capabilities: Schema.Array(Capability).check(Schema.isUnique()),
   })
   export interface Response extends Schema.Schema.Type<typeof Response> {}
 
@@ -564,29 +597,28 @@ export namespace Backend {
     matched: Schema.Boolean,
   })
   export interface NetworkLogEntry extends Schema.Schema.Type<typeof NetworkLogEntry> {}
-}
 
-export class SimulationRequestError extends Schema.TaggedErrorClass<SimulationRequestError>()(
-  "SimulationRequestError",
-  {
-    method: Schema.String,
-    code: Schema.Number,
-    message: Schema.String,
-    data: Schema.optionalKey(Schema.Json),
-  },
-) {}
-
-const request = <
-  const Tag extends string,
-  Payload extends Schema.Top | Schema.Struct.Fields = typeof Schema.Void,
-  Success extends Schema.Top = typeof Schema.Void,
->(
-  tag: Tag,
-  options?: {
-    readonly payload?: Payload
-    readonly success?: Success
-  },
-) => Rpc.make(tag, { ...options, error: SimulationRequestError })
+  export const Notification = Schema.Union([
+    Schema.Struct({
+      jsonrpc: Schema.Literal("2.0"),
+      method: Schema.Literal("llm.request"),
+      params: ProviderInvocation,
+    }),
+    Schema.Struct({
+      jsonrpc: Schema.Literal("2.0"),
+      method: Schema.Literal("tool.invocation"),
+      params: ToolInvocation,
+    }),
+    Schema.Struct({
+      jsonrpc: Schema.Literal("2.0"),
+      method: Schema.Literal("tool.cancel"),
+      params: ToolCancellation,
+    }),
+  ])
+  export type Notification = Schema.Schema.Type<typeof Notification>
+  export const decodeNotification = Schema.decodeUnknownSync(Notification)
+  export const decodeNotificationEffect = Schema.decodeUnknownEffect(Schema.fromJsonString(Notification))
+}
 
 export const UiRpcs = RpcGroup.make(
   request("simulation.handshake", { payload: Handshake.Params, success: Handshake.Response }),

packages/simulation/test/protocol.test.ts

@@ -1,6 +1,64 @@
 import { describe, expect, test } from "bun:test"
 import { Effect, Schema } from "effect"
-import { Backend, Frontend, Handshake } from "../src/protocol"
+import { Backend, Frontend, Handshake, JsonRpc } from "../src/protocol"
+
+const successResponse: Schema.Schema.Type<typeof JsonRpc.Response> = { jsonrpc: "2.0", id: 1, result: null }
+// @ts-expect-error responses require one outcome
+const missingResponse: Schema.Schema.Type<typeof JsonRpc.Response> = { jsonrpc: "2.0", id: 1 }
+// @ts-expect-error responses cannot contain both outcomes
+const invalidResponse: Schema.Schema.Type<typeof JsonRpc.Response> = {
+  jsonrpc: "2.0",
+  id: 1,
+  result: null,
+  error: { code: -32600, message: "Invalid request" },
+}
+void [successResponse, missingResponse, invalidResponse]
+
+test("normalizes an omitted finish reason", () => {
+  expect(Backend.decodeRequest({ jsonrpc: "2.0", id: 1, method: "llm.finish", params: { id: "inv_1" } })).toMatchObject(
+    { params: { id: "inv_1", reason: "stop" } },
+  )
+})
+
+test("decodes typed backend notifications", () => {
+  expect(
+    Backend.decodeNotification({
+      jsonrpc: "2.0",
+      method: "tool.cancel",
+      params: { id: "tool_1", reason: "interrupted" },
+    }),
+  ).toEqual({
+    jsonrpc: "2.0",
+    method: "tool.cancel",
+    params: { id: "tool_1", reason: "interrupted" },
+  })
+  expect(() =>
+    Backend.decodeNotification({
+      jsonrpc: "2.0",
+      method: "tool.cancel",
+      params: { id: "tool_1", reason: "unknown" },
+    }),
+  ).toThrow()
+})
+
+test("requires exactly one JSON-RPC response outcome", () => {
+  const decode = Schema.decodeUnknownSync(JsonRpc.Response)
+  expect(decode({ jsonrpc: "2.0", id: 1, result: null })).toEqual({ jsonrpc: "2.0", id: 1, result: null })
+  expect(decode({ jsonrpc: "2.0", id: 1, error: { code: -32600, message: "Invalid request" } })).toEqual({
+    jsonrpc: "2.0",
+    id: 1,
+    error: { code: -32600, message: "Invalid request" },
+  })
+  expect(() => decode({ jsonrpc: "2.0", id: 1 })).toThrow()
+  expect(() =>
+    decode({
+      jsonrpc: "2.0",
+      id: 1,
+      result: null,
+      error: { code: -32600, message: "Invalid request" },
+    }),
+  ).toThrow()
+})
 
 test("decodes ui.matches text params", () => {
   expect(

packages/simulation/test/simulated-provider.test.ts

@@ -83,7 +83,7 @@ test("streams a Drive-controlled provider response and removes the finished invo
           jsonrpc: "2.0",
           id: 3,
           method: "llm.finish",
-          params: { id: params.id, reason: "stop" },
+          params: { id: params.id },
         }),
       )
       expect(yield* Queue.take(messages)).toMatchObject({ id: 3, result: { ok: true } })