WIP: AsyncGenerator streaming stub

Author: ruifigueira

packages/agents/src/react.tsx

@@ -145,11 +145,11 @@ type NonStreamingRPCMethod<T extends Method> =
     : never;
 
 interface StreamingResponse<
-  OnChunkT extends SerializableValue | unknown = unknown,
-  OnDoneT extends SerializableValue | unknown = unknown
+  Chunk extends SerializableValue | unknown = unknown,
+  Done extends SerializableValue | unknown = unknown
 > {
-  send(chunk: OnChunkT): void;
-  end(finalChunk?: OnDoneT): void;
+  send(chunk: Chunk): void;
+  end(finalChunk?: Done): void;
 }
 
 type StreamingRPCMethod<T extends Method> = T extends (
@@ -190,6 +190,14 @@ type StreamOptionsFrom<StreamingResponseT> =
     ? StreamOptions<T, U>
     : never;
 
+type ReturnAndChunkTypesFrom<StreamingResponseT extends StreamingResponse> =
+  StreamingResponseT extends StreamingResponse<
+    infer Chunk extends SerializableValue,
+    infer Done extends SerializableValue
+  >
+    ? [Chunk, Done]
+    : never;
+
 type RestParameters<T extends Method> =
   Parameters<StreamingRPCMethod<T>> extends [unknown, ...infer Rest]
     ? Rest
@@ -266,20 +274,35 @@ type UntypedAgentMethodCall = <T = unknown>(
 ) => Promise<T>;
 
 type AgentStub<T> = {
+  [K in keyof AgentMethods<T>]: AgentMethods<T>[K] extends NonStreamingRPCMethod<
+    AgentMethods<T>[K]
+  >
+    ? (
+        ...args: Parameters<AgentMethods<T>[K]>
+      ) => AgentPromiseReturnType<AgentMethods<T>, K>
+    : never;
+};
+
+type AgentStreamingStub<T> = {
   [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>;
+      ) => AsyncGenerator<
+        ReturnAndChunkTypesFrom<
+          Parameters<StreamingRPCMethod<AgentMethods<T>[K]>>[0]
+        >[0],
+        ReturnAndChunkTypesFrom<
+          Parameters<StreamingRPCMethod<AgentMethods<T>[K]>>[0]
+        >[1]
+      >
+    : never;
 };
 
 // we neet to use Method instead of RPCMethod here for retro-compatibility
 type UntypedAgentStub = Record<string, Method>;
+type UntypedAgentStreamingStub = StreamingAgentMethods<unknown>;
 
 /**
  * React hook for connecting to an Agent
@@ -292,6 +315,7 @@ export function useAgent<State = unknown>(
   setState: (state: State) => void;
   call: UntypedAgentMethodCall;
   stub: UntypedAgentStub;
+  streamingStub: UntypedAgentStreamingStub;
 };
 export function useAgent<
   AgentT extends {
@@ -306,6 +330,7 @@ export function useAgent<
   setState: (state: State) => void;
   call: AgentMethodCall<AgentT>;
   stub: AgentStub<AgentT>;
+  streamingStub: AgentStreamingStub<AgentT>;
 };
 export function useAgent<State>(
   options: UseAgentOptions<unknown>
@@ -462,6 +487,7 @@ export function useAgent<State>(
     setState: (state: State) => void;
     call: UntypedAgentMethodCall;
     stub: UntypedAgentStub;
+    streamingStub: UntypedAgentStreamingStub;
   };
   // Create the call method
   const call = useCallback(
@@ -510,6 +536,55 @@ export function useAgent<State>(
       }
     }
   );
+  // biome-ignore lint: suppressions/parse
+  agent.streamingStub = new Proxy<any>(
+    {},
+    {
+      get: (_target, method) => {
+        return async function* (...args: unknown[]) {
+          let resolve: (value: unknown) => void;
+          let reject: (reason: unknown) => void;
+          let promise = new Promise((res, rej) => {
+            resolve = res;
+            reject = rej;
+          });
+
+          // 4. State flags
+          let isDone = false;
+
+          // 5. Callback implementation
+          const streamOptions: StreamOptions = {
+            onChunk: (chunk: unknown) => {
+              resolve(chunk);
+              promise = new Promise((res, rej) => {
+                resolve = res;
+                reject = rej;
+              });
+            },
+            onError: (error: unknown) => {
+              isDone = true;
+              reject(error);
+            },
+            onDone: (done: unknown) => {
+              isDone = true;
+              resolve(done);
+            }
+          };
+
+          call(method as string, args, streamOptions);
+
+          while (!isDone) {
+            const result = await promise;
+            if (isDone) {
+              return result;
+            } else {
+              yield result;
+            }
+          }
+        };
+      }
+    }
+  );
 
   // warn if agent isn't in lowercase
   if (agent.agent !== agent.agent.toLowerCase()) {

packages/agents/src/tests-d/example-stub.test-d.ts

@@ -2,7 +2,6 @@
 import type { env } from "cloudflare:workers";
 import { Agent, callable, type StreamingResponse } from "..";
 import { useAgent } from "../react.tsx";
-import type { StreamOptions } from "../client.ts";
 
 class MyAgent extends Agent<typeof env, {}> {
   @callable()
@@ -57,7 +56,7 @@ class MyAgent extends Agent<typeof env, {}> {
   }
 }
 
-const { stub } = useAgent<MyAgent, {}>({ agent: "my-agent" });
+const { stub, streamingStub } = useAgent<MyAgent, {}>({ agent: "my-agent" });
 // return type is promisified
 stub.sayHello() satisfies Promise<string>;
 
@@ -76,21 +75,16 @@ 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;
+const generator = streamingStub.performStream("hello") satisfies AsyncGenerator<
+  number,
+  boolean
+>;
+for await (const chunk of generator) {
+  chunk satisfies number;
+}
 
 // @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);
+streamingStub.performStream("hello", 1);
 
 const { stub: stub2 } = useAgent<Omit<MyAgent, "nonRpc">, {}>({
   agent: "my-agent"