diff --git a/packages/genui/a2ui-playground/src/pages/AIChatPage.css b/packages/genui/a2ui-playground/src/pages/AIChatPage.css index c0e5cf16f1..65df65a835 100644 --- a/packages/genui/a2ui-playground/src/pages/AIChatPage.css +++ b/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); diff --git a/packages/genui/a2ui-playground/src/pages/AIChatPage.tsx b/packages/genui/a2ui-playground/src/pages/AIChatPage.tsx index 95b6e1bded..2c7db9d720 100644 --- a/packages/genui/a2ui-playground/src/pages/AIChatPage.tsx +++ b/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; + 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 { 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(false); + const [tokenUsage, setTokenUsage] = useState({ + promptTokens: 0, + completionTokens: 0, + totalTokens: 0, + }); const messagesEndRef = useRef(null); const chatMessagesRef = useRef(null); const followBottomRef = useRef(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 }) {

Create

Online Agent + {tokenUsage.totalTokens > 0 + ? ( + + + Prompt {formatTokenCount(tokenUsage.promptTokens)} + + + Output {formatTokenCount(tokenUsage.completionTokens)} + + + Total {formatTokenCount(tokenUsage.totalTokens)} + + + ) + : null}

Describe the UI you want to build

diff --git a/packages/genui/server/tsconfig.json b/packages/genui/server/tsconfig.json index bc01542b8d..324d6bb847 100644 --- a/packages/genui/server/tsconfig.json +++ b/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", + ], }