feat (rsc): add streamUI onFinish callback

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

@@ -13,10 +13,14 @@ import { getValidatedPrompt } from '../../core/prompt/get-validated-prompt';
 import { prepareCallSettings } from '../../core/prompt/prepare-call-settings';
 import { prepareToolsAndToolChoice } from '../../core/prompt/prepare-tools-and-tool-choice';
 import { Prompt } from '../../core/prompt/prompt';
-import { CoreToolChoice } from '../../core/types';
+import { CallWarning, CoreToolChoice, FinishReason } from '../../core/types';
 import { retryWithExponentialBackoff } from '../../core/util/retry-with-exponential-backoff';
 import { createStreamableUI } from '../streamable';
 import { createResolvablePromise } from '../utils';
+import {
+  TokenUsage,
+  calculateTokenUsage,
+} from '../../core/generate-text/token-usage';
 
 type Streamable = ReactNode | Promise<ReactNode>;
 
@@ -84,6 +88,7 @@ export async function streamUI<
   abortSignal,
   initial,
   text,
+  onFinish,
   ...settings
 }: CallSettings &
   Prompt & {
@@ -100,12 +105,42 @@ export async function streamUI<
     };
 
     /**
-The tool choice strategy. Default: 'auto'.
+     * The tool choice strategy. Default: 'auto'.
      */
     toolChoice?: CoreToolChoice<TOOLS>;
 
     text?: RenderText;
     initial?: ReactNode;
+    /**
+     * Callback that is called when the LLM response and the final object validation are finished.
+     */
+    onFinish?: (event: {
+      /**
+       * The reason why the generation finished.
+       */
+      finishReason: FinishReason;
+      /**
+       * The token usage of the generated response.
+       */
+      usage: TokenUsage;
+      /**
+       * The final ui node that was generated.
+       */
+      value: ReactNode;
+      /**
+       * Warnings from the model provider (e.g. unsupported settings)
+       */
+      warnings?: CallWarning[];
+      /**
+       * Optional raw response data.
+       */
+      rawResponse?: {
+        /**
+         * Response headers.
+         */
+        headers?: Record<string, string>;
+      };
+    }) => Promise<void> | void;
   }): Promise<RenderResult> {
   // TODO: Remove these errors after the experimental phase.
   if (typeof model === 'string') {
@@ -311,7 +346,13 @@ The tool choice strategy. Default: 'auto'.
           }
 
           case 'finish': {
-            // Nothing to do here.
+            onFinish?.({
+              finishReason: value.finishReason,
+              usage: calculateTokenUsage(value.usage),
+              value: ui.value,
+              warnings: result.warnings,
+              rawResponse: result.rawResponse,
+            });
           }
         }
       }