fix(core): route Vertex multi-regions through REP

packages/core/src/plugin/provider/google-vertex.ts

@@ -26,6 +26,7 @@ function resolveLocation(options: Record<string, any>) {
 
 function vertexEndpoint(location: string) {
   if (location === "global") return "aiplatform.googleapis.com"
+  if (location === "eu" || location === "us") return `aiplatform.${location}.rep.googleapis.com`
   return `${location}-aiplatform.googleapis.com`
 }
 
@@ -111,13 +112,18 @@ export const GoogleVertexPlugin = define({
         if (evt.package !== "@ai-sdk/google-vertex") return
         const mod = yield* Effect.promise(() => import("@ai-sdk/google-vertex"))
         const project = resolveProject(evt.options)
-        const location = resolveLocation(evt.options)
+        const location = String(resolveLocation(evt.options))
         const options = { ...evt.options }
         delete options.fetch
         evt.sdk = mod.createVertex({
           ...options,
           project,
           location,
+          ...((location === "eu" || location === "us") && project && !options.apiKey && !options.baseURL
+            ? {
+                baseURL: `https://${vertexEndpoint(location)}/v1beta1/projects/${project}/locations/${location}/publishers/google`,
+              }
+            : {}),
         })
       }),
     )

packages/core/test/plugin/provider-google-vertex.test.ts

@@ -230,7 +230,7 @@ describe("GoogleVertexPlugin", () => {
     ),
   )
 
-  it.effect("keeps OpenAI-compatible Vertex endpoint templates regional for eu", () =>
+  it.effect("uses the REP endpoint for OpenAI-compatible Vertex multi-regions", () =>
     Effect.gen(function* () {
       const catalog = yield* Catalog.Service
       yield* catalog.transform((catalog) =>
@@ -248,11 +248,65 @@ describe("GoogleVertexPlugin", () => {
       const provider = required(yield* catalog.provider.get(Provider.ID.make("google-vertex")))
       expect(provider).toMatchObject({
         package: "aisdk:@ai-sdk/openai-compatible",
-        settings: { baseURL: "https://eu-aiplatform.googleapis.com/v1/projects/config-project/locations/eu" },
+        settings: {
+          baseURL: "https://aiplatform.eu.rep.googleapis.com/v1/projects/config-project/locations/eu",
+        },
       })
     }),
   )
 
+  it.effect("uses the REP endpoint for the native Vertex SDK in multi-regions", () =>
+    withEnv(
+      {
+        GOOGLE_CLOUD_PROJECT: "env-project",
+        GOOGLE_VERTEX_LOCATION: "us",
+      },
+      () =>
+        Effect.gen(function* () {
+          vertexOptions.length = 0
+          const aisdk = yield* AISDK.Service
+          yield* addPlugin()
+          yield* aisdk.runSDK({
+            model: Model.Info.make({
+              ...Model.Info.default(Provider.ID.make("google-vertex"), Model.ID.make("gemini")),
+              modelID: Model.ID.make("gemini"),
+              package: "aisdk:@ai-sdk/google-vertex",
+            }),
+            package: "@ai-sdk/google-vertex",
+            options: { name: "google-vertex" },
+          })
+          expect(vertexOptions).toHaveLength(1)
+          expect(vertexOptions[0].baseURL).toBe(
+            "https://aiplatform.us.rep.googleapis.com/v1beta1/projects/env-project/locations/us/publishers/google",
+          )
+        }),
+    ),
+  )
+
+  it.effect("preserves custom native Vertex base URLs in multi-regions", () =>
+    Effect.gen(function* () {
+      vertexOptions.length = 0
+      const aisdk = yield* AISDK.Service
+      yield* addPlugin()
+      yield* aisdk.runSDK({
+        model: Model.Info.make({
+          ...Model.Info.default(Provider.ID.make("google-vertex"), Model.ID.make("gemini")),
+          modelID: Model.ID.make("gemini"),
+          package: "aisdk:@ai-sdk/google-vertex",
+        }),
+        package: "@ai-sdk/google-vertex",
+        options: {
+          name: "google-vertex",
+          project: "project",
+          location: "eu",
+          baseURL: "https://vertex.example/v1",
+        },
+      })
+      expect(vertexOptions).toHaveLength(1)
+      expect(vertexOptions[0].baseURL).toBe("https://vertex.example/v1")
+    }),
+  )
+
   it.effect("defaults location to us-central1 when only project is configured", () =>
     withEnv(
       {

packages/protocol/src/simulation.ts

@@ -4,10 +4,6 @@ import { Rpc, RpcGroup } from "effect/unstable/rpc"
 const JsonRpcID = Schema.Union([Schema.String, Schema.Number, Schema.Null])
 const decodeJson = Schema.decodeUnknownSync(Schema.Json)
 
-// Generated schema unions lose mapped tuple types even though both come from the same operation tuple.
-// eslint-disable-next-line typescript-eslint/no-unsafe-type-assertion, typescript-eslint/no-unnecessary-type-parameters
-const decoded = <Type>(value: unknown) => value as Type
-
 export namespace JsonRpc {
   export const RequestFields = {
     jsonrpc: Schema.Literal("2.0"),
@@ -26,33 +22,13 @@ export namespace JsonRpc {
     data: Schema.optional(Schema.Json),
   })
 
-  export type Response =
-    | {
-        readonly jsonrpc: "2.0"
-        readonly id: string | number | null
-        readonly result: Schema.Schema.Type<typeof Schema.Json>
-        readonly error?: never
-      }
-    | {
-        readonly jsonrpc: "2.0"
-        readonly id: string | number | null
-        readonly error: Schema.Schema.Type<typeof ErrorObject>
-        readonly result?: never
-      }
-  export const Response = decoded<Schema.Decoder<Response>>(
-    Schema.Struct({
-      jsonrpc: Schema.Literal("2.0"),
-      id: JsonRpcID,
-      result: Schema.optionalKey(Schema.Json),
-      error: Schema.optionalKey(ErrorObject),
-    }).check(
-      Schema.makeFilter((response) =>
-        "result" in response === "error" in response
-          ? "JSON-RPC responses must contain exactly one of result or error"
-          : undefined,
-      ),
-    ),
-  )
+  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)
 
@@ -73,161 +49,6 @@ 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 })
-
-function operation<const Tag extends string, Success extends Schema.Decoder<unknown>>(method: Tag, success: Success) {
-  return {
-    method,
-    success,
-    request: Schema.Struct({ ...JsonRpc.RequestFields, method: Schema.Literal(method) }),
-    rpc: request(method, { success }),
-  }
-}
-
-function operationWithPayload<
-  const Tag extends string,
-  Payload extends Schema.Decoder<unknown>,
-  Success extends Schema.Decoder<unknown>,
->(method: Tag, payload: Payload, success: Success) {
-  return {
-    method,
-    payload,
-    success,
-    request: Schema.Struct({ ...JsonRpc.RequestFields, method: Schema.Literal(method), params: payload }),
-    rpc: request(method, { payload, success }),
-  }
-}
-
-function operationWithRpcPayload<
-  const Tag extends string,
-  Payload extends Schema.Decoder<unknown>,
-  RpcPayload extends Schema.Decoder<unknown>,
-  Success extends Schema.Decoder<unknown>,
->(method: Tag, payload: Payload, rpcPayload: RpcPayload, success: Success) {
-  return {
-    method,
-    payload,
-    success,
-    request: Schema.Struct({ ...JsonRpc.RequestFields, method: Schema.Literal(method), params: payload }),
-    rpc: request(method, { payload: rpcPayload, success }),
-  }
-}
-
-function notification<const Method extends string, Payload extends Schema.Decoder<unknown>>(
-  method: Method,
-  payload: Payload,
-) {
-  return {
-    method,
-    payload,
-    schema: Schema.Struct({
-      jsonrpc: Schema.Literal("2.0"),
-      method: Schema.Literal(method),
-      params: payload,
-    }),
-  }
-}
-
-type OperationRequest<Operation> = Operation extends {
-  readonly method: infer Method extends string
-  readonly payload: infer Payload extends Schema.Top
-}
-  ? Omit<JsonRpc.Request, "method" | "params"> & {
-      readonly method: Method
-      readonly params: Schema.Schema.Type<Payload>
-    }
-  : Operation extends { readonly method: infer Method extends string }
-    ? Omit<JsonRpc.Request, "method" | "params"> & { readonly method: Method }
-    : never
-
-type EndpointRequest<Operations extends ReadonlyArray<{ readonly method: string }>> =
-  | Handshake.Request
-  | OperationRequest<Operations[number]>
-
-type EndpointNotification<Notifications extends ReadonlyArray<{ readonly schema: Schema.Top }>> = Schema.Schema.Type<
-  Notifications[number]["schema"]
->
-
-function endpoint<
-  const Operations extends ReadonlyArray<{
-    readonly method: string
-    readonly success: Schema.Top
-    readonly rpc: Rpc.Any
-    readonly request: Schema.Decoder<unknown>
-    readonly payload?: Schema.Decoder<unknown>
-  }>,
-  const Notifications extends ReadonlyArray<{
-    readonly method: string
-    readonly payload: Schema.Decoder<unknown>
-    readonly schema: Schema.Decoder<unknown>
-  }>,
-  const Capabilities extends ReadonlyArray<Handshake.Capability>,
->(
-  operations: Operations,
-  notifications: Notifications,
-  features: ReadonlyArray<Handshake.Capability>,
-  Capabilities: Capabilities,
-) {
-  const requests: ReadonlyArray<Schema.Decoder<unknown>> = [
-    Handshake.Request,
-    ...operations.map((operation) => operation.request),
-  ]
-  const Request = Schema.Union(requests)
-  const Notification =
-    notifications.length === 0 ? Schema.Never : Schema.Union(notifications.map((notification) => notification.schema))
-  const decodeRequest = Schema.decodeUnknownSync(Request)
-  const decodeRequestEffect = Schema.decodeUnknownEffect(Schema.fromJsonString(Request))
-  const decodeNotification = Schema.decodeUnknownSync(Notification)
-  const decodeNotificationEffect = Schema.decodeUnknownEffect(Schema.fromJsonString(Notification))
-  const handshake = request("simulation.handshake", { payload: Handshake.Params, success: Handshake.Response })
-  const rpcs = RpcGroup.make(handshake, ...operations.map((operation) => operation.rpc)) as RpcGroup.RpcGroup<
-    typeof handshake | Operations[number]["rpc"]
-  >
-  const derived = [
-    ...operations.map((operation) => operation.method),
-    ...notifications.map((notification) => notification.method),
-    ...features,
-  ]
-  if (
-    new Set(Capabilities).size !== Capabilities.length ||
-    derived.length !== Capabilities.length ||
-    derived.some((capability) => !Capabilities.includes(capability))
-  )
-    throw new Error("Simulation capabilities must exactly match endpoint operations")
-  return {
-    Capabilities,
-    Request: decoded<Schema.Decoder<EndpointRequest<Operations>>>(Request),
-    Notification: decoded<Schema.Decoder<EndpointNotification<Notifications>>>(Notification),
-    decodeRequest: (input: unknown) => decoded<EndpointRequest<Operations>>(decodeRequest(input)),
-    decodeRequestEffect: (input: string) =>
-      decodeRequestEffect(input).pipe(Effect.map(decoded<EndpointRequest<Operations>>)),
-    decodeNotification: (input: unknown) => decoded<EndpointNotification<Notifications>>(decodeNotification(input)),
-    decodeNotificationEffect: (input: string) =>
-      decodeNotificationEffect(input).pipe(Effect.map(decoded<EndpointNotification<Notifications>>)),
-    rpcs,
-  }
-}
-
 export namespace Handshake {
   export const ProtocolVersion = Schema.Literal(1)
   export type ProtocolVersion = Schema.Schema.Type<typeof ProtocolVersion>
@@ -260,7 +81,7 @@ export namespace Handshake {
     protocolVersion: ProtocolVersion,
     role: EndpointRole,
     server: Identity,
-    capabilities: Schema.Array(Capability).check(Schema.isUnique()),
+    capabilities: Schema.Array(Capability),
   })
   export interface Response extends Schema.Schema.Type<typeof Response> {}
 
@@ -344,6 +165,23 @@ export namespace Handshake {
 }
 
 export namespace Frontend {
+  export const Capabilities = [
+    "ui.type",
+    "ui.press",
+    "ui.enter",
+    "ui.arrow",
+    "ui.focus",
+    "ui.click",
+    "ui.click.semantic",
+    "ui.resize",
+    "ui.matches",
+    "ui.state",
+    "ui.snapshot",
+    "ui.capture",
+    "ui.recording.finish",
+  ] as const satisfies ReadonlyArray<Handshake.Capability>
+  export type Capability = (typeof Capabilities)[number]
+
   export const KeyModifiers = Schema.Struct({
     ctrl: Schema.optional(Schema.Boolean),
     shift: Schema.optional(Schema.Boolean),
@@ -499,51 +337,44 @@ export namespace Frontend {
 
   export const ResizeParams = Schema.Struct({ cols: Schema.Number, rows: Schema.Number })
   export interface ResizeParams extends Schema.Schema.Type<typeof ResizeParams> {}
-}
 
-const FrontendOperations = [
-  operation("ui.state", Frontend.State),
-  operation("ui.snapshot", Frontend.SemanticSnapshot),
-  operation("ui.capture", Frontend.CapturedFrame),
-  operationWithPayload("ui.matches", Frontend.MatchesParams, Frontend.Matches),
-  operation("ui.recording.finish", Frontend.RecordingFinish),
-  operationWithPayload("ui.type", Frontend.TypeParams, Frontend.State),
-  operationWithPayload("ui.press", Frontend.PressParams, Frontend.State),
-  operation("ui.enter", Frontend.State),
-  operationWithPayload("ui.arrow", Frontend.ArrowParams, Frontend.State),
-  operationWithPayload("ui.focus", Frontend.FocusParams, Frontend.State),
-  operationWithPayload("ui.click", Frontend.ClickParams, Frontend.State),
-  operationWithPayload("ui.resize", Frontend.ResizeParams, Frontend.State),
-] as const
-const FrontendCapabilities = [
-  "ui.type",
-  "ui.press",
-  "ui.enter",
-  "ui.arrow",
-  "ui.focus",
-  "ui.click",
-  "ui.click.semantic",
-  "ui.resize",
-  "ui.matches",
-  "ui.state",
-  "ui.snapshot",
-  "ui.capture",
-  "ui.recording.finish",
-] as const
-const FrontendEndpoint = endpoint(FrontendOperations, [], ["ui.click.semantic"], FrontendCapabilities)
-type FrontendRequest = EndpointRequest<typeof FrontendOperations>
-
-export namespace Frontend {
-  export const Capabilities: typeof FrontendCapabilities = FrontendEndpoint.Capabilities
-  export type Capability = (typeof Capabilities)[number]
-  export const Request: Schema.Decoder<FrontendRequest> = FrontendEndpoint.Request
-  export type Request = FrontendRequest
-  export const decodeRequest: (input: unknown) => Request = FrontendEndpoint.decodeRequest
-  export const decodeRequestEffect: (input: string) => Effect.Effect<Request, Schema.SchemaError> =
-    FrontendEndpoint.decodeRequestEffect
+  export const Request = Schema.Union([
+    Handshake.Request,
+    Schema.Struct({ ...JsonRpc.RequestFields, method: Schema.Literal("ui.type"), params: TypeParams }),
+    Schema.Struct({ ...JsonRpc.RequestFields, method: Schema.Literal("ui.press"), params: PressParams }),
+    Schema.Struct({ ...JsonRpc.RequestFields, method: Schema.Literal("ui.arrow"), params: ArrowParams }),
+    Schema.Struct({ ...JsonRpc.RequestFields, method: Schema.Literal("ui.focus"), params: FocusParams }),
+    Schema.Struct({ ...JsonRpc.RequestFields, method: Schema.Literal("ui.click"), params: ClickParams }),
+    Schema.Struct({ ...JsonRpc.RequestFields, method: Schema.Literal("ui.resize"), params: ResizeParams }),
+    Schema.Struct({ ...JsonRpc.RequestFields, method: Schema.Literal("ui.matches"), params: MatchesParams }),
+    Schema.Struct({
+      ...JsonRpc.RequestFields,
+      method: Schema.Literals(["ui.enter", "ui.state", "ui.snapshot", "ui.recording.finish"]),
+    }),
+    Schema.Struct({ ...JsonRpc.RequestFields, method: Schema.Literal("ui.capture") }),
+  ])
+  export type Request = Schema.Schema.Type<typeof Request>
+  export const decodeRequest = Schema.decodeUnknownSync(Request)
+  export const decodeRequestEffect = Schema.decodeUnknownEffect(Schema.fromJsonString(Request))
 }
 
 export namespace Backend {
+  export const Capabilities = [
+    "llm.attach",
+    "llm.chunk",
+    "llm.finish",
+    "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<Handshake.Capability>
+
   export const Item = Schema.Union([
     Schema.Struct({ type: Schema.Literal("textDelta"), text: Schema.String }),
     Schema.Struct({ type: Schema.Literal("reasoningDelta"), text: Schema.String }),
@@ -702,6 +533,24 @@ export namespace Backend {
   export const DisconnectParams = Schema.Struct({ id: Schema.String })
   export interface DisconnectParams extends Schema.Schema.Type<typeof DisconnectParams> {}
 
+  export const Request = Schema.Union([
+    Handshake.Request,
+    Schema.Struct({ ...JsonRpc.RequestFields, method: Schema.Literal("llm.chunk"), params: ChunkParams }),
+    Schema.Struct({ ...JsonRpc.RequestFields, method: Schema.Literal("llm.finish"), params: FinishParams }),
+    Schema.Struct({ ...JsonRpc.RequestFields, 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"]),
+    }),
+  ])
+  export type Request = Schema.Schema.Type<typeof Request>
+  export const decodeRequest = Schema.decodeUnknownSync(Request)
+  export const decodeRequestEffect = Schema.decodeUnknownEffect(Schema.fromJsonString(Request))
+
   export const ProviderInvocation = Schema.Struct({ id: Schema.String, url: Schema.String, body: Schema.Json })
   export interface ProviderInvocation extends Schema.Schema.Type<typeof ProviderInvocation> {}
 
@@ -717,54 +566,53 @@ export namespace Backend {
   export interface NetworkLogEntry extends Schema.Schema.Type<typeof NetworkLogEntry> {}
 }
 
-const BackendOperations = [
-  operation("llm.attach", Backend.Attached),
-  operation("llm.pending", Backend.Pending),
-  operationWithPayload("llm.chunk", Backend.ChunkParams, Backend.Ok),
-  operationWithRpcPayload("llm.finish", Backend.FinishParams, Backend.FinishPayload, Backend.Ok),
-  operationWithPayload("llm.disconnect", Backend.DisconnectParams, Backend.Ok),
-  operationWithPayload("tool.attach", Backend.ToolAttachParams, Backend.Attached),
-  operationWithPayload("tool.update", Backend.ToolUpdateParams, Backend.Ok),
-  operationWithPayload("tool.finish", Backend.ToolFinishParams, Backend.Ok),
-  operationWithPayload("tool.fail", Backend.ToolFailParams, Backend.Ok),
-] as const
-const BackendNotifications = [
-  notification("llm.request", Backend.ProviderInvocation),
-  notification("tool.invocation", Backend.ToolInvocation),
-  notification("tool.cancel", Backend.ToolCancellation),
-] as const
-const BackendCapabilities = [
-  "llm.attach",
-  "llm.chunk",
-  "llm.finish",
-  "llm.disconnect",
-  "llm.pending",
-  "llm.request",
-  "llm.tool-input-delta",
-  "tool.attach",
-  "tool.update",
-  "tool.finish",
-  "tool.fail",
-  "tool.invocation",
-  "tool.cancel",
-] as const
-const BackendEndpoint = endpoint(BackendOperations, BackendNotifications, ["llm.tool-input-delta"], BackendCapabilities)
-type BackendRequest = EndpointRequest<typeof BackendOperations>
-type BackendNotification = EndpointNotification<typeof BackendNotifications>
+export class SimulationRequestError extends Schema.TaggedErrorClass<SimulationRequestError>()(
+  "SimulationRequestError",
+  {
+    method: Schema.String,
+    code: Schema.Number,
+    message: Schema.String,
+    data: Schema.optionalKey(Schema.Json),
+  },
+) {}
 
-export namespace Backend {
-  export const Capabilities: typeof BackendCapabilities = BackendEndpoint.Capabilities
-  export const Request: Schema.Decoder<BackendRequest> = BackendEndpoint.Request
-  export type Request = BackendRequest
-  export const decodeRequest: (input: unknown) => Request = BackendEndpoint.decodeRequest
-  export const decodeRequestEffect: (input: string) => Effect.Effect<Request, Schema.SchemaError> =
-    BackendEndpoint.decodeRequestEffect
-  export const Notification: Schema.Decoder<BackendNotification> = BackendEndpoint.Notification
-  export type Notification = BackendNotification
-  export const decodeNotification: (input: unknown) => Notification = BackendEndpoint.decodeNotification
-  export const decodeNotificationEffect: (input: string) => Effect.Effect<Notification, Schema.SchemaError> =
-    BackendEndpoint.decodeNotificationEffect
-}
+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 = FrontendEndpoint.rpcs
-export const BackendRpcs = BackendEndpoint.rpcs
+export const UiRpcs = RpcGroup.make(
+  request("simulation.handshake", { payload: Handshake.Params, success: Handshake.Response }),
+  request("ui.state", { success: Frontend.State }),
+  request("ui.snapshot", { success: Frontend.SemanticSnapshot }),
+  request("ui.capture", { success: Frontend.CapturedFrame }),
+  request("ui.matches", { payload: Frontend.MatchesParams, success: Frontend.Matches }),
+  request("ui.recording.finish", { success: Frontend.RecordingFinish }),
+  request("ui.type", { payload: Frontend.TypeParams, success: Frontend.State }),
+  request("ui.press", { payload: Frontend.PressParams, success: Frontend.State }),
+  request("ui.enter", { success: Frontend.State }),
+  request("ui.arrow", { payload: Frontend.ArrowParams, success: Frontend.State }),
+  request("ui.focus", { payload: Frontend.FocusParams, success: Frontend.State }),
+  request("ui.click", { payload: Frontend.ClickParams, success: Frontend.State }),
+  request("ui.resize", { payload: Frontend.ResizeParams, success: Frontend.State }),
+)
+
+export const BackendRpcs = RpcGroup.make(
+  request("simulation.handshake", { payload: Handshake.Params, success: Handshake.Response }),
+  request("llm.attach", { success: Backend.Attached }),
+  request("llm.pending", { success: Backend.Pending }),
+  request("llm.chunk", { payload: Backend.ChunkParams, success: Backend.Ok }),
+  request("llm.finish", { payload: Backend.FinishPayload, success: Backend.Ok }),
+  request("llm.disconnect", { 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 }),
+)

packages/simulation/test/protocol.test.ts

@@ -1,110 +1,6 @@
 import { describe, expect, test } from "bun:test"
 import { Effect, Schema } from "effect"
-import { Backend, BackendRpcs, Frontend, Handshake, JsonRpc, UiRpcs } from "../src/protocol"
-
-const uiCapability: Frontend.Capability = "ui.state"
-// @ts-expect-error capability literals must remain narrow for consumers
-const invalidUiCapability: Frontend.Capability = "ui.future"
-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 [uiCapability, invalidUiCapability, successResponse, missingResponse, invalidResponse]
-
-test("preserves capability order and keeps request capabilities aligned with RPC groups", () => {
-  expect(Frontend.Capabilities).toEqual([
-    "ui.type",
-    "ui.press",
-    "ui.enter",
-    "ui.arrow",
-    "ui.focus",
-    "ui.click",
-    "ui.click.semantic",
-    "ui.resize",
-    "ui.matches",
-    "ui.state",
-    "ui.snapshot",
-    "ui.capture",
-    "ui.recording.finish",
-  ])
-  expect(Backend.Capabilities).toEqual([
-    "llm.attach",
-    "llm.chunk",
-    "llm.finish",
-    "llm.disconnect",
-    "llm.pending",
-    "llm.request",
-    "llm.tool-input-delta",
-    "tool.attach",
-    "tool.update",
-    "tool.finish",
-    "tool.fail",
-    "tool.invocation",
-    "tool.cancel",
-  ])
-  expect(new Set<string>(Frontend.Capabilities.filter((capability) => capability !== "ui.click.semantic"))).toEqual(
-    new Set(Array.from(UiRpcs.requests.keys()).filter((method) => method !== "simulation.handshake")),
-  )
-  expect(
-    new Set<string>(
-      Backend.Capabilities.filter(
-        (capability) => !["llm.request", "llm.tool-input-delta", "tool.invocation", "tool.cancel"].includes(capability),
-      ),
-    ),
-  ).toEqual(new Set(Array.from(BackendRpcs.requests.keys()).filter((method) => method !== "simulation.handshake")))
-})
-
-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()
-})
+import { Backend, Frontend, Handshake } from "../src/protocol"
 
 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 },
+          params: { id: params.id, reason: "stop" },
         }),
       )
       expect(yield* Queue.take(messages)).toMatchObject({ id: 3, result: { ok: true } })