feat(s2ui): create chat add token show

packages/genui/a2ui-playground/src/pages/AIChatPage.css

@@ -55,6 +55,39 @@
   white-space: nowrap;
 }
 
+.chatTokenUsageBadge {
+  display: inline-flex;
+  align-items: center;
+  gap: 8px;
+  min-height: 26px;
+  margin-left: auto;
+  padding: 0 10px;
+  border: 1px solid var(--geist-border);
+  border-radius: 999px;
+  background: var(--geist-surface);
+  color: var(--geist-secondary);
+  font-family: var(--geist-mono);
+  font-size: 12px;
+  font-weight: 500;
+  line-height: 1;
+  white-space: nowrap;
+}
+
+.chatTokenUsageItem {
+  display: inline-flex;
+  align-items: center;
+}
+
+.chatTokenUsageTotal {
+  padding: 3px 7px;
+  border: 1px solid var(--geist-border);
+  border-radius: 999px;
+  background: var(--geist-background);
+  color: var(--geist-foreground);
+  font-size: 12px;
+  font-weight: 600;
+}
+
 .chatHeaderSub {
   font-size: 12px;
   color: var(--geist-secondary);

packages/genui/a2ui-playground/src/pages/AIChatPage.tsx

@@ -44,6 +44,50 @@ interface A2UIDonePayload {
   };
   error?: unknown;
   message?: unknown;
+  usage?: unknown;
+}
+
+interface TokenUsage {
+  promptTokens: number;
+  completionTokens: number;
+  totalTokens: number;
+}
+
+function parseUsage(value: unknown): TokenUsage | null {
+  if (!value || typeof value !== 'object') return null;
+  const record = value as Record<string, unknown>;
+  const pickNumber = (...keys: string[]): number => {
+    for (const key of keys) {
+      const v = record[key];
+      if (typeof v === 'number' && Number.isFinite(v)) return v;
+    }
+    return 0;
+  };
+  // Support both legacy (promptTokens/completionTokens) and new
+  // (inputTokens/outputTokens) AI SDK usage shapes.
+  const promptTokens = pickNumber(
+    'promptTokens',
+    'inputTokens',
+    'prompt_tokens',
+  );
+  const completionTokens = pickNumber(
+    'completionTokens',
+    'outputTokens',
+    'completion_tokens',
+  );
+  const totalTokens = pickNumber('totalTokens', 'total_tokens')
+    || promptTokens + completionTokens;
+  if (promptTokens === 0 && completionTokens === 0 && totalTokens === 0) {
+    return null;
+  }
+  return { promptTokens, completionTokens, totalTokens };
+}
+
+function formatTokenCount(value: number): string {
+  if (value < 1000) return String(value);
+  if (value < 10_000) return `${(value / 1000).toFixed(2)}k`;
+  if (value < 1_000_000) return `${(value / 1000).toFixed(1)}k`;
+  return `${(value / 1_000_000).toFixed(2)}M`;
 }
 
 interface BrowserResponse {
@@ -294,6 +338,7 @@ async function readA2UIResponse(
   onText: (text: string) => void,
   onMessages: (messages: unknown[]) => void,
   onStart?: (threadId: string) => void,
+  onUsage?: (usage: TokenUsage) => void,
 ): Promise<unknown[]> {
   const contentType = response.headers.get('content-type') ?? '';
   if (!contentType.includes('text/event-stream')) {
@@ -302,6 +347,10 @@ async function readA2UIResponse(
     if (messages.length === 0) {
       throw new Error(normalizeErrorPayload(payload));
     }
+    if (payload && typeof payload === 'object') {
+      const usage = parseUsage((payload as A2UIDonePayload).usage);
+      if (usage) onUsage?.(usage);
+    }
     onMessages(messages);
     return messages;
   }
@@ -348,6 +397,10 @@ async function readA2UIResponse(
 
       if (parsed.event === 'done') {
         const doneMessages = normalizeA2UIMessages(parsed.data);
+        if (parsed.data && typeof parsed.data === 'object') {
+          const usage = parseUsage((parsed.data as A2UIDonePayload).usage);
+          if (usage) onUsage?.(usage);
+        }
         if (doneMessages.length > 0) {
           latestMessages = doneMessages;
           onMessages(latestMessages);
@@ -372,14 +425,14 @@ async function readA2UIResponse(
 
 const SUGGESTED_PROMPTS: Array<{ label: string; text: string }> = [
   {
-    label: '⚡ Quiz card with actions',
+    label: '🛍️ Product card with Buy',
     text:
-      'Create a trivia quiz card. Show a question "What is the capital of Japan?" with 4 answer buttons: Tokyo, Beijing, Seoul, Bangkok. When the user taps an answer, show whether it is correct with a brief explanation.',
+      'Create a product card for a limited-edition sneaker. Include name, price ($189), a short description, and a "Buy Now" button. When tapped, show an order confirmation with a fake order number and estimated delivery.',
   },
   {
-    label: '🛍️ Product card with Buy',
+    label: '⚡ Quiz card with actions',
     text:
-      'Create a product card for a limited-edition sneaker. Include name, price ($189), a short description, and a "Buy Now" button. When tapped, show an order confirmation with a fake order number and estimated delivery.',
+      'Create a trivia quiz card. Show a question "Which shape has three sides?" with 4 answer buttons: Triangle, Square, Circle, Hexagon. When the user taps an answer, show whether it is correct with a brief explanation.',
   },
   {
     label: '🌤️ Weather with Refresh',
@@ -398,6 +451,11 @@ export function AIChatPage(props: { protocol: Protocol }) {
     null,
   );
   const [isGenerating, setIsGenerating] = useState<boolean>(false);
+  const [tokenUsage, setTokenUsage] = useState<TokenUsage>({
+    promptTokens: 0,
+    completionTokens: 0,
+    totalTokens: 0,
+  });
   const messagesEndRef = useRef<HTMLDivElement>(null);
   const chatMessagesRef = useRef<HTMLDivElement>(null);
   const followBottomRef = useRef<boolean>(true);
@@ -564,6 +622,11 @@ export function AIChatPage(props: { protocol: Protocol }) {
     setPreviewMessages(null);
     latestPreviewMessagesRef.current = [];
     threadIdRef.current = null;
+    setTokenUsage({
+      promptTokens: 0,
+      completionTokens: 0,
+      totalTokens: 0,
+    });
     setIsGenerating(true);
 
     void (async () => {
@@ -598,6 +661,14 @@ export function AIChatPage(props: { protocol: Protocol }) {
           (threadId) => {
             threadIdRef.current = threadId;
           },
+          (usage) => {
+            if (controller.signal.aborted) return;
+            setTokenUsage((prev) => ({
+              promptTokens: prev.promptTokens + usage.promptTokens,
+              completionTokens: prev.completionTokens + usage.completionTokens,
+              totalTokens: prev.totalTokens + usage.totalTokens,
+            }));
+          },
         );
 
         if (finalMessages.length === 0) {
@@ -773,6 +844,16 @@ export function AIChatPage(props: { protocol: Protocol }) {
               if (signal.aborted) return;
               responseMessages = msgs;
             },
+            undefined,
+            (usage) => {
+              if (signal.aborted) return;
+              setTokenUsage((prev) => ({
+                promptTokens: prev.promptTokens + usage.promptTokens,
+                completionTokens: prev.completionTokens
+                  + usage.completionTokens,
+                totalTokens: prev.totalTokens + usage.totalTokens,
+              }));
+            },
           );
 
           if (signal.aborted) return;
@@ -896,6 +977,24 @@ export function AIChatPage(props: { protocol: Protocol }) {
           <div className='chatHeaderTitleRow'>
             <h2 className='chatHeaderTitle'>Create</h2>
             <span className='constructionBadge'>Online Agent</span>
+            {tokenUsage.totalTokens > 0
+              ? (
+                <span
+                  className='chatTokenUsageBadge'
+                  title={`Prompt: ${tokenUsage.promptTokens} · Completion: ${tokenUsage.completionTokens} · Total: ${tokenUsage.totalTokens}`}
+                >
+                  <span className='chatTokenUsageItem'>
+                    Prompt {formatTokenCount(tokenUsage.promptTokens)}
+                  </span>
+                  <span className='chatTokenUsageItem'>
+                    Output {formatTokenCount(tokenUsage.completionTokens)}
+                  </span>
+                  <span className='chatTokenUsageItem chatTokenUsageTotal'>
+                    Total {formatTokenCount(tokenUsage.totalTokens)}
+                  </span>
+                </span>
+              )
+              : null}
           </div>
           <p className='chatHeaderSub'>Describe the UI you want to build</p>
         </div>

packages/genui/server/tsconfig.json

@@ -1,10 +1,14 @@
 {
   "compilerOptions": {
     "target": "ES2022",
-    "lib": ["dom", "dom.iterable", "esnext"],
+    "lib": [
+      "dom",
+      "dom.iterable",
+      "esnext",
+    ],
     "module": "esnext",
     "moduleResolution": "bundler",
-    "jsx": "preserve",
+    "jsx": "react-jsx",
     "allowJs": true,
     "skipLibCheck": true,
     "strict": true,
@@ -23,7 +27,9 @@
       },
     ],
     "paths": {
-      "@/*": ["./*"],
+      "@/*": [
+        "./*",
+      ],
     },
   },
   "include": [
@@ -34,5 +40,7 @@
     ".next/dev/types/**/*.ts",
     "**/*.mts",
   ],
-  "exclude": ["node_modules"],
+  "exclude": [
+    "node_modules",
+  ],
 }