Type-safe react stub streaming calls

ruifigueira · ruifigueira · commit 76fa487b41a8 · 2025-05-25T00:42:07.000+01:00
diff --git a/packages/agents/src/client.ts b/packages/agents/src/client.ts
@@ -27,11 +27,14 @@ export type AgentClientOptions<State = unknown> = Omit<
 /**
  * Options for streaming RPC calls
  */
-export type StreamOptions = {
+export type StreamOptions<
+  OnChunkT extends unknown | SerializableValue = unknown,
+  OnDoneT extends unknown | SerializableValue = unknown,
+> = {
   /** Called when a chunk of data is received */
-  onChunk?: (chunk: unknown) => void;
+  onChunk?: (chunk: OnChunkT) => void;
   /** Called when the stream ends */
-  onDone?: (finalChunk: unknown) => void;
+  onDone?: (finalChunk: OnDoneT) => void;
   /** Called when an error occurs */
   onError?: (error: string) => void;
 };
diff --git a/packages/agents/src/index.ts b/packages/agents/src/index.ts
@@ -25,7 +25,7 @@ import type { Client } from "@modelcontextprotocol/sdk/client/index.js";
 import type { SSEClientTransportOptions } from "@modelcontextprotocol/sdk/client/sse.js";
 
 import { camelCaseToKebabCase } from "./client";
-import type { MCPClientConnection } from "./mcp/client-connection";
+import type { SerializableValue } from "./serializable";
 
 export type { Connection, ConnectionContext, WSMessage } from "partyserver";
 
@@ -1205,7 +1205,10 @@ export async function getAgentByName<Env, T extends Agent<Env>>(
 /**
  * A wrapper for streaming responses in callable methods
  */
-export class StreamingResponse {
+export class StreamingResponse<
+  OnChunkT extends SerializableValue | unknown = unknown,
+  OnDoneT extends SerializableValue | unknown = unknown,
+> {
   private _connection: Connection;
   private _id: string;
   private _closed = false;
@@ -1219,7 +1222,7 @@ export class StreamingResponse {
    * Send a chunk of data to the client
    * @param chunk The data to send
    */
-  send(chunk: unknown) {
+  send(chunk: OnChunkT) {
     if (this._closed) {
       throw new Error("StreamingResponse is already closed");
     }
@@ -1237,7 +1240,7 @@ export class StreamingResponse {
    * End the stream and send the final chunk (if any)
    * @param finalChunk Optional final chunk of data to send
    */
-  end(finalChunk?: unknown) {
+  end(finalChunk?: OnDoneT) {
     if (this._closed) {
       throw new Error("StreamingResponse is already closed");
     }
diff --git a/packages/agents/src/react.tsx b/packages/agents/src/react.tsx
@@ -3,7 +3,11 @@ import { usePartySocket } from "partysocket/react";
 import { useCallback, useRef } from "react";
 import type { MCPServersState, RPCRequest, RPCResponse, Agent } from "./";
 import type { StreamOptions } from "./client";
-import type { Method, RPCMethod } from "./serializable";
+import type {
+  AllSerializableValues,
+  SerializableReturnValue,
+  SerializableValue,
+} from "./serializable";
 
 /**
  * Convert a camelCase string to a kebab-case string
@@ -44,17 +48,68 @@ export type UseAgentOptions<State = unknown> = Omit<
   onMcpUpdate?: (mcpServers: MCPServersState) => void;
 };
 
+// biome-ignore lint: suppressions/parse
+type Method = (...args: any[]) => any;
+
+type NonStreamingRPCMethod<T extends Method> =
+  AllSerializableValues<Parameters<T>> extends true
+    ? ReturnType<T> extends SerializableReturnValue
+      ? T
+      : never
+    : never;
+
+interface StreamingResponse<
+  OnChunkT extends SerializableValue | unknown = unknown,
+  OnDoneT extends SerializableValue | unknown = unknown,
+> {
+  send(chunk: OnChunkT): void;
+  end(finalChunk?: OnDoneT): void;
+}
+
+type StreamingRPCMethod<T extends Method> = T extends (
+  arg: infer A,
+  ...rest: infer R
+) => void | Promise<void>
+  ? A extends StreamingResponse<SerializableValue, SerializableValue>
+    ? AllSerializableValues<R> extends true
+      ? T
+      : never
+    : never
+  : never;
+
+type RPCMethod<T extends Method> =
+  T extends NonStreamingRPCMethod<T>
+    ? NonStreamingRPCMethod<T>
+    : T extends StreamingRPCMethod<T>
+      ? StreamingRPCMethod<T>
+      : never;
+
+type RPCMethods<T> = {
+  [K in keyof T as T[K] extends Method ? K : never]: T[K] extends Method
+    ? RPCMethod<T[K]>
+    : never;
+};
+
 type AllOptional<T> = T extends [infer A, ...infer R]
   ? undefined extends A
     ? AllOptional<R>
     : false
   : true; // no params means optional by default
 
-type RPCMethods<T> = {
-  [K in keyof T as T[K] extends RPCMethod<T[K]> ? K : never]: RPCMethod<T[K]>;
-};
+type StreamOptionsFrom<StreamingResponseT> =
+  StreamingResponseT extends StreamingResponse<
+    infer T extends SerializableValue,
+    infer U extends SerializableValue
+  >
+    ? StreamOptions<T, U>
+    : never;
 
-type OptionalParametersMethod<T extends RPCMethod> =
+type RestParameters<T extends Method> =
+  Parameters<StreamingRPCMethod<T>> extends [unknown, ...infer Rest]
+    ? Rest
+    : never;
+
+type OptionalParametersMethod<T extends Method> =
   AllOptional<Parameters<T>> extends true ? T : never;
 
 // all methods of the Agent, excluding the ones that are declared in the base Agent class
@@ -74,6 +129,14 @@ type RequiredAgentMethods<T> = Omit<
   keyof OptionalAgentMethods<T>
 >;
 
+type StreamingAgentMethods<T> = {
+  [K in keyof AgentMethods<T> as AgentMethods<T>[K] extends StreamingRPCMethod<
+    AgentMethods<T>[K]
+  >
+    ? K
+    : never]: StreamingRPCMethod<AgentMethods<T>[K]>;
+};
+
 type AgentPromiseReturnType<T, K extends keyof AgentMethods<T>> =
   // biome-ignore lint: suppressions/parse
   ReturnType<AgentMethods<T>[K]> extends Promise<any>
@@ -96,7 +159,18 @@ type RequiredArgsAgentMethodCall<AgentT> = <
   streamOptions?: StreamOptions
 ) => AgentPromiseReturnType<AgentT, K>;
 
-type AgentMethodCall<AgentT> = OptionalArgsAgentMethodCall<AgentT> &
+type StreamingAgentMethodCall<AgentT> = <
+  K extends keyof StreamingAgentMethods<AgentT>,
+>(
+  method: K,
+  args: RestParameters<StreamingAgentMethods<AgentT>[K]>,
+  streamOptions: StreamOptionsFrom<
+    Parameters<StreamingAgentMethods<AgentT>[K]>[0]
+  >
+) => void;
+
+type AgentMethodCall<AgentT> = StreamingAgentMethodCall<AgentT> &
+  OptionalArgsAgentMethodCall<AgentT> &
   RequiredArgsAgentMethodCall<AgentT>;
 
 type UntypedAgentMethodCall = <T = unknown>(
@@ -106,9 +180,16 @@ type UntypedAgentMethodCall = <T = unknown>(
 ) => Promise<T>;
 
 type AgentStub<T> = {
-  [K in keyof AgentMethods<T>]: (
-    ...args: Parameters<AgentMethods<T>[K]>
-  ) => AgentPromiseReturnType<AgentMethods<T>, K>;
+  [K in keyof AgentMethods<T>]: AgentMethods<T>[K] extends StreamingRPCMethod<
+    AgentMethods<T>[K]
+  >
+    ? (
+        options: StreamOptionsFrom<Parameters<AgentMethods<T>[K]>[0]>,
+        ...args: RestParameters<AgentMethods<T>[K]>
+      ) => void
+    : (
+        ...args: Parameters<AgentMethods<T>[K]>
+      ) => AgentPromiseReturnType<AgentMethods<T>, K>;
 };
 
 // we neet to use Method instead of RPCMethod here for retro-compatibility
@@ -150,7 +231,7 @@ export function useAgent<State>(
   agent: string;
   name: string;
   setState: (state: State) => void;
-  call: UntypedAgentMethodCall | AgentMethodCall<unknown>;
+  call: UntypedAgentMethodCall;
   stub: UntypedAgentStub;
 } {
   const agentNamespace = camelCaseToKebabCase(options.agent);
diff --git a/packages/agents/src/serializable.ts b/packages/agents/src/serializable.ts
@@ -13,21 +13,8 @@ export type SerializableReturnValue =
   | Promise<SerializableValue>
   | Promise<void>;
 
-type AllSerializableValues<A> = A extends [infer First, ...infer Rest]
+export type AllSerializableValues<A> = A extends [infer First, ...infer Rest]
   ? First extends SerializableValue
     ? AllSerializableValues<Rest>
     : false
   : true; // no params means serializable by default
-
-// biome-ignore lint: suspicious/noExplicitAny
-export type Method = (...args: any[]) => any;
-
-export type RPCMethod<T = Method> = T extends Method
-  ? T extends (...arg: infer A) => infer R
-    ? AllSerializableValues<A> extends true
-      ? R extends SerializableReturnValue
-        ? T
-        : never
-      : never
-    : never
-  : never;
diff --git a/packages/agents/src/tests-d/example-stub.test-d.ts b/packages/agents/src/tests-d/example-stub.test-d.ts
@@ -1,6 +1,11 @@
 import type { env } from "cloudflare:workers";
-import { unstable_callable as callable, Agent } from "..";
+import {
+  unstable_callable as callable,
+  Agent,
+  type StreamingResponse,
+} from "..";
 import { useAgent } from "../react.tsx";
+import type { StreamOptions } from "../client.ts";
 
 class MyAgent extends Agent<typeof env, {}> {
   @callable()
@@ -17,6 +22,42 @@ class MyAgent extends Agent<typeof env, {}> {
   nonRpc(): void {
     // do something
   }
+
+  @callable({ streaming: true })
+  performStream(
+    options: StreamingResponse<number, boolean>,
+    other: string
+  ): void {
+    // do something
+  }
+
+  // TODO should fail, first argument is not a streamOptions
+  @callable({ streaming: true })
+  performStreamFirstArgNotStreamOptions(
+    other: string,
+    options: StreamingResponse<number, boolean>
+  ): void {
+    // do something
+  }
+
+  // TODO should fail, should be marked as streaming
+  @callable()
+  performStreamFail(options: StreamingResponse): void {
+    // do something
+  }
+
+  // TODO should fail, has no streamOptions
+  @callable({ streaming: true })
+  async performFail(task: string): Promise<string> {
+    // do something
+    return "";
+  }
+
+  @callable({ streaming: true })
+  performStreamUnserializable(options: StreamingResponse<Date>): void {
+    // @ts-expect-error parameter is not serializable
+    options.onDone(new Date());
+  }
 }
 
 const { stub } = useAgent<MyAgent, {}>({ agent: "my-agent" });
@@ -38,9 +79,26 @@ await stub.nonRpc();
 // @ts-expect-error nonSerializable is not serializable
 await stub.nonSerializable("hello", new Date());
 
+const streamOptions: StreamOptions<number, boolean> = {};
+
+// biome-ignore lint: suspicious/noConfusingVoidType
+stub.performStream(streamOptions, "hello") satisfies void;
+
+// @ts-expect-error there's no 2nd argument
+stub.performStream(streamOptions, "hello", 1);
+
+const invalidStreamOptions: StreamOptions<string, boolean> = {};
+
+// @ts-expect-error streamOptions must be of type StreamOptions<number, boolean>
+stub.performStream(invalidStreamOptions, "hello");
+
+// @ts-expect-error first argument is not a streamOptions
+stub.performStreamFirstArgNotStreamOptions("hello", streamOptions);
+
 const { stub: stub2 } = useAgent<Omit<MyAgent, "nonRpc">, {}>({
   agent: "my-agent",
 });
+
 stub2.sayHello();
 // @ts-expect-error nonRpc excluded from useAgent
 stub2.nonRpc();
diff --git a/packages/agents/src/tests-d/example.test-d.ts b/packages/agents/src/tests-d/example.test-d.ts
@@ -1,6 +1,11 @@
 import type { env } from "cloudflare:workers";
-import { unstable_callable as callable, Agent } from "..";
+import {
+  unstable_callable as callable,
+  Agent,
+  type StreamingResponse,
+} from "..";
 import { useAgent } from "../react.tsx";
+import type { StreamOptions } from "../client.ts";
 
 class MyAgent extends Agent<typeof env, {}> {
   @callable()
@@ -17,6 +22,16 @@ class MyAgent extends Agent<typeof env, {}> {
   nonRpc(): void {
     // do something
   }
+
+  @callable({ streaming: true })
+  performStream(
+    response: StreamingResponse<number, boolean>,
+    other: string
+  ): void {
+    response.send(1);
+    response.send(2);
+    response.end(true);
+  }
 }
 
 const agent = useAgent<MyAgent, {}>({ agent: "my-agent" });
@@ -38,6 +53,18 @@ await agent.call("nonRpc");
 // @ts-expect-error nonSerializable is not serializable
 await agent.call("nonSerializable", ["hello", new Date()]);
 
+const streamOptions: StreamOptions<number, boolean> = {};
+
+agent.call("performStream", ["hello"], streamOptions);
+
+// @ts-expect-error there's no second parameter
+agent.call("performStream", ["a", 1], streamOptions);
+
+const invalidStreamOptions: StreamOptions<string, boolean> = {};
+
+// @ts-expect-error streamOptions must be of type StreamOptions<number, boolean>
+agent.call("performStream", ["a", 1], invalidStreamOptions);
+
 const agent2 = useAgent<Omit<MyAgent, "nonRpc">, {}>({ agent: "my-agent" });
 agent2.call("sayHello");
 // @ts-expect-error nonRpc excluded from useAgent
