Fix render() generator cases and use onText (#1021)

Co-authored-by: Max Leiter <max.leiter@vercel.com>

Author: shuding

.changeset/late-pugs-cheer.md

@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+ai/rsc: fix text response and async generator

packages/core/rsc/streamable.tsx

@@ -227,6 +227,12 @@ export function render<
       )
     : undefined;
 
+  if (functions && tools) {
+    throw new Error(
+      "You can't have both functions and tools defined. Please choose one or the other.",
+    );
+  }
+
   let finished: ReturnType<typeof createResolvablePromise> | undefined;
 
   async function handleRender(
@@ -254,20 +260,23 @@ export function render<
       typeof value === 'object' &&
       Symbol.asyncIterator in value
     ) {
-      for await (const node of value as AsyncGenerator<
+      const it = value as AsyncGenerator<
         React.ReactNode,
         React.ReactNode,
         void
-      >) {
-        res.update(node);
+      >;
+      while (true) {
+        const { done, value } = await it.next();
+        res.update(value);
+        if (done) break;
       }
       finished?.resolve(void 0);
     } else if (value && typeof value === 'object' && Symbol.iterator in value) {
       const it = value as Generator<React.ReactNode, React.ReactNode, void>;
       while (true) {
         const { done, value } = it.next();
-        if (done) break;
         res.update(value);
+        if (done) break;
       }
       finished?.resolve(void 0);
     } else {
@@ -299,14 +308,19 @@ export function render<
             : {}),
         })) as any,
         {
-          async experimental_onFunctionCall(functionCallPayload) {
-            hasFunction = true;
-            handleRender(
-              functionCallPayload.arguments,
-              options.functions?.[functionCallPayload.name as any]?.render,
-              ui,
-            );
-          },
+          ...(functions
+            ? {
+                async experimental_onFunctionCall(functionCallPayload) {
+                  hasFunction = true;
+                  handleRender(
+                    functionCallPayload.arguments,
+                    options.functions?.[functionCallPayload.name as any]
+                      ?.render,
+                    ui,
+                  );
+                },
+              }
+            : {}),
           ...(tools
             ? {
                 async experimental_onToolCall(toolCallPayload: any) {
@@ -323,15 +337,18 @@ export function render<
                 },
               }
             : {}),
-          onToken(token) {
-            text += token;
-            if (hasFunction) return;
+          onText(chunk) {
+            text += chunk;
             handleRender({ content: text, done: false }, options.text, ui);
           },
           async onFinal() {
-            if (hasFunction) return;
-            handleRender({ content: text, done: true }, options.text, ui);
+            if (hasFunction) {
+              await finished?.promise;
+              ui.done();
+              return;
+            }
 
+            handleRender({ content: text, done: true }, options.text, ui);
             await finished?.promise;
             ui.done();
           },