From 7484e32f62ebe5938a7a03b7941c53672e708079 Mon Sep 17 00:00:00 2001 From: Kit Langton Date: Fri, 14 Aug 2026 20:29:53 -0400 Subject: [PATCH] refactor(protocol): harden simulation wire contract (#42628) --- packages/protocol/src/simulation.ts | 92 +++++++++++++------ packages/simulation/test/protocol.test.ts | 60 +++++++++++- .../test/simulated-provider.test.ts | 2 +- 3 files changed, 122 insertions(+), 32 deletions(-) diff --git a/packages/protocol/src/simulation.ts b/packages/protocol/src/simulation.ts index 89b82c2f283..43b50fc8b64 100644 --- a/packages/protocol/src/simulation.ts +++ b/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 {} + 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 export const decodeRequest = Schema.decodeUnknownSync(Request) @@ -49,6 +60,28 @@ export namespace JsonRpc { } } +export class SimulationRequestError extends Schema.TaggedErrorClass()( + "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 @@ -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 {} @@ -564,30 +597,29 @@ export namespace Backend { matched: Schema.Boolean, }) export interface NetworkLogEntry extends Schema.Schema.Type {} + + 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 + export const decodeNotification = Schema.decodeUnknownSync(Notification) + export const decodeNotificationEffect = Schema.decodeUnknownEffect(Schema.fromJsonString(Notification)) } -export class SimulationRequestError extends Schema.TaggedErrorClass()( - "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 UiRpcs = RpcGroup.make( request("simulation.handshake", { payload: Handshake.Params, success: Handshake.Response }), request("ui.state", { success: Frontend.State }), diff --git a/packages/simulation/test/protocol.test.ts b/packages/simulation/test/protocol.test.ts index 34893914c6b..89ba0b0c588 100644 --- a/packages/simulation/test/protocol.test.ts +++ b/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 = { jsonrpc: "2.0", id: 1, result: null } +// @ts-expect-error responses require one outcome +const missingResponse: Schema.Schema.Type = { jsonrpc: "2.0", id: 1 } +// @ts-expect-error responses cannot contain both outcomes +const invalidResponse: Schema.Schema.Type = { + 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( diff --git a/packages/simulation/test/simulated-provider.test.ts b/packages/simulation/test/simulated-provider.test.ts index e3b0b07439e..d016f6d1a24 100644 --- a/packages/simulation/test/simulated-provider.test.ts +++ b/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 } })