fix(ai/rsc): Remove extra reconciliation of streamUI (#1818)

Author: shuding

.changeset/late-years-share.md

@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+fix(ai/rsc): Remove extra reconcilation of streamUI

packages/core/rsc/stream-ui/__snapshots__/stream-ui.ui.test.tsx.snap

@@ -63,13 +63,7 @@ exports[`result.value > should render text function returned ui 1`] = `
                 "next": {
                   "done": false,
                   "next": {
-                    "done": false,
-                    "next": {
-                      "done": true,
-                      "value": <h1>
-                        { "content": "Hello, world!" }
-                      </h1>,
-                    },
+                    "done": true,
                     "value": <h1>
                       { "content": "Hello, world!" }
                     </h1>,
@@ -115,22 +109,11 @@ exports[`result.value > should render tool call results 1`] = `
     "props": {
       "c": undefined,
       "n": {
-        "done": false,
-        "next": {
-          "done": false,
-          "next": {
-            "done": true,
-            "value": <div>
-              tool1: 
-              value
-            </div>,
-          },
-          "value": <div>
-            tool1: 
-            value
-          </div>,
-        },
-        "value": "",
+        "done": true,
+        "value": <div>
+          tool1: 
+          value
+        </div>,
       },
     },
     "type": "",
@@ -151,22 +134,11 @@ exports[`result.value > should render tool call results with generator render fu
       "n": {
         "done": false,
         "next": {
-          "done": false,
-          "next": {
-            "done": false,
-            "next": {
-              "done": true,
-              "value": <div>
-                tool: 
-                value
-              </div>,
-            },
-            "value": <div>
-              tool: 
-              value
-            </div>,
-          },
-          "value": "",
+          "done": true,
+          "value": <div>
+            tool: 
+            value
+          </div>,
         },
         "value": <div>
           Loading...

packages/core/rsc/stream-ui/stream-ui.tsx

@@ -145,6 +145,7 @@ The tool choice strategy. Default: 'auto'.
     args: [payload: any] | [payload: any, options: any],
     renderer: undefined | Renderer<any>,
     res: ReturnType<typeof createStreamableUI>,
+    lastCall = false,
   ) {
     if (!renderer) return;
 
@@ -165,7 +166,13 @@ The tool choice strategy. Default: 'auto'.
         typeof value.then === 'function')
     ) {
       const node = await (value as Promise<React.ReactNode>);
-      res.update(node);
+
+      if (lastCall) {
+        res.done(node);
+      } else {
+        res.update(node);
+      }
+
       resolvable.resolve(void 0);
     } else if (
       value &&
@@ -179,20 +186,32 @@ The tool choice strategy. Default: 'auto'.
       >;
       while (true) {
         const { done, value } = await it.next();
-        res.update(value);
+        if (lastCall && done) {
+          res.done(value);
+        } else {
+          res.update(value);
+        }
         if (done) break;
       }
       resolvable.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();
-        res.update(value);
+        if (lastCall && done) {
+          res.done(value);
+        } else {
+          res.update(value);
+        }
         if (done) break;
       }
       resolvable.resolve(void 0);
     } else {
-      res.update(value);
+      if (lastCall) {
+        res.done(value);
+      } else {
+        res.update(value);
+      }
       resolvable.resolve(void 0);
     }
   }
@@ -257,6 +276,7 @@ The tool choice strategy. Default: 'auto'.
               });
             }
 
+            hasToolCall = true;
             const parseResult = safeParseJSON({
               text: value.args,
               schema: tool.parameters,
@@ -280,6 +300,7 @@ The tool choice strategy. Default: 'auto'.
               ],
               tool.generate,
               ui,
+              true,
             );
 
             break;
@@ -297,11 +318,9 @@ The tool choice strategy. Default: 'auto'.
 
       if (hasToolCall) {
         await finished;
-        ui.done();
       } else {
-        handleRender([{ content, done: true }], textRender, ui);
+        handleRender([{ content, done: true }], textRender, ui, true);
         await finished;
-        ui.done();
       }
     } catch (error) {
       // During the stream rendering, we don't want to throw the error to the