(Work In Progress) Exploration of how to show token usage

app/components/chat/AssistantMessage.tsx

@@ -1,14 +1,24 @@
 import { memo } from 'react';
 import { Markdown } from './Markdown';
+import { USAGE_REGEX } from '~/utils/constants';
 
 interface AssistantMessageProps {
   content: string;
 }
 
 export const AssistantMessage = memo(({ content }: AssistantMessageProps) => {
+  const match = content.match(USAGE_REGEX);
+  const usage = match ? JSON.parse(match[1]) : null;
+  const cleanContent = content.replace(USAGE_REGEX, '').trim();
+
   return (
     <div className="overflow-hidden w-full">
-      <Markdown html>{content}</Markdown>
+      {usage && (
+        <div className="text-sm text-bolt-elements-textSecondary mb-2">
+          Tokens: {usage.totalTokens} (prompt: {usage.promptTokens}, completion: {usage.completionTokens})
+        </div>
+      )}
+      <Markdown html>{cleanContent}</Markdown>
     </div>
   );
 });

app/components/chat/BaseChat.tsx

@@ -77,7 +77,6 @@ export const BaseChat = React.forwardRef<HTMLDivElement, BaseChatProps>(
       input = '',
       enhancingPrompt,
       handleInputChange,
-      promptEnhanced,
       enhancePrompt,
       sendMessage,
       handleStop,
@@ -490,10 +489,7 @@ export const BaseChat = React.forwardRef<HTMLDivElement, BaseChatProps>(
                       <IconButton
                         title="Enhance prompt"
                         disabled={input.length === 0 || enhancingPrompt}
-                        className={classNames(
-                          'transition-all',
-                          enhancingPrompt ? 'opacity-100' : '',
-                        )}
+                        className={classNames('transition-all', enhancingPrompt ? 'opacity-100' : '')}
                         onClick={() => {
                           enhancePrompt?.();
                           toast.success('Prompt enhanced!');

app/components/chat/Chat.client.tsx

@@ -116,13 +116,22 @@ export const ChatImpl = memo(
         apiKeys,
         files,
       },
+      sendExtraMessageFields: true,
       onError: (error) => {
         logger.error('Request failed\n\n', error);
         toast.error(
           'There was an error processing your request: ' + (error.message ? error.message : 'No details were returned'),
         );
       },
-      onFinish: () => {
+      onFinish: (message, response) => {
+        const usage = response.usage;
+
+        if (usage) {
+          console.log('Token usage:', usage);
+
+          // You can now use the usage data as needed
+        }
+
         logger.debug('Finished streaming');
       },
       initialMessages,

app/components/chat/UserMessage.tsx

@@ -12,42 +12,36 @@ interface UserMessageProps {
 export function UserMessage({ content }: UserMessageProps) {
   if (Array.isArray(content)) {
     const textItem = content.find((item) => item.type === 'text');
-    const textContent = sanitizeUserMessage(textItem?.text || '');
+    const textContent = stripMetadata(textItem?.text || '');
     const images = content.filter((item) => item.type === 'image' && item.image);
 
     return (
       <div className="overflow-hidden pt-[4px]">
-        <div className="flex items-start gap-4">
-          <div className="flex-1">
-            <Markdown limitedMarkdown>{textContent}</Markdown>
-          </div>
-          {images.length > 0 && (
-            <div className="flex-shrink-0 w-[160px]">
-              {images.map((item, index) => (
-                <div key={index} className="relative">
-                  <img
-                    src={item.image}
-                    alt={`Uploaded image ${index + 1}`}
-                    className="w-full h-[160px] rounded-lg object-cover border border-bolt-elements-borderColor"
-                  />
-                </div>
-              ))}
-            </div>
-          )}
+        <div className="flex flex-col gap-4">
+          {textContent && <Markdown html>{textContent}</Markdown>}
+          {images.map((item, index) => (
+            <img
+              key={index}
+              src={item.image}
+              alt={`Image ${index + 1}`}
+              className="max-w-full h-auto rounded-lg"
+              style={{ maxHeight: '512px', objectFit: 'contain' }}
+            />
+          ))}
         </div>
       </div>
     );
   }
 
-  const textContent = sanitizeUserMessage(content);
+  const textContent = stripMetadata(content);
 
   return (
     <div className="overflow-hidden pt-[4px]">
-      <Markdown limitedMarkdown>{textContent}</Markdown>
+      <Markdown html>{textContent}</Markdown>
     </div>
   );
 }
 
-function sanitizeUserMessage(content: string) {
+function stripMetadata(content: string) {
   return content.replace(MODEL_REGEX, '').replace(PROVIDER_REGEX, '');
 }

app/components/settings/chat-history/ChatHistoryTab.tsx

@@ -22,7 +22,8 @@ export default function ChatHistoryTab() {
   };
 
   const handleDeleteAllChats = async () => {
-    const confirmDelete = window.confirm("Are you sure you want to delete all chats? This action cannot be undone.");
+    const confirmDelete = window.confirm('Are you sure you want to delete all chats? This action cannot be undone.');
+
     if (!confirmDelete) {
       return; // Exit if the user cancels
     }
@@ -31,11 +32,13 @@ export default function ChatHistoryTab() {
       const error = new Error('Database is not available');
       logStore.logError('Failed to delete chats - DB unavailable', error);
       toast.error('Database is not available');
+
       return;
     }
 
     try {
       setIsDeleting(true);
+
       const allChats = await getAll(db);
       await Promise.all(allChats.map((chat) => deleteById(db!, chat.id)));
       logStore.logSystem('All chats deleted successfully', { count: allChats.length });
@@ -55,6 +58,7 @@ export default function ChatHistoryTab() {
       const error = new Error('Database is not available');
       logStore.logError('Failed to export chats - DB unavailable', error);
       toast.error('Database is not available');
+
       return;
     }
 

app/components/settings/debug/DebugTab.tsx

@@ -41,7 +41,8 @@ const versionHash = commit.commit;
 const GITHUB_URLS = {
   original: 'https://api.github.com/repos/stackblitz-labs/bolt.diy/commits/main',
   fork: 'https://api.github.com/repos/Stijnus/bolt.new-any-llm/commits/main',
-  commitJson: (branch: string) => `https://raw.githubusercontent.com/stackblitz-labs/bolt.diy/${branch}/app/commit.json`,
+  commitJson: (branch: string) =>
+    `https://raw.githubusercontent.com/stackblitz-labs/bolt.diy/${branch}/app/commit.json`,
 };
 
 function getSystemInfo(): SystemInfo {
@@ -227,19 +228,20 @@ export default function DebugTab() {
               provider.name.toLowerCase() === 'ollama'
                 ? 'OLLAMA_API_BASE_URL'
                 : provider.name.toLowerCase() === 'lmstudio'
-                ? 'LMSTUDIO_API_BASE_URL'
-                : `REACT_APP_${provider.name.toUpperCase()}_URL`;
+                  ? 'LMSTUDIO_API_BASE_URL'
+                  : `REACT_APP_${provider.name.toUpperCase()}_URL`;
 
             // Access environment variables through import.meta.env
             const url = import.meta.env[envVarName] || provider.settings.baseUrl || null; // Ensure baseUrl is used
             console.log(`[Debug] Using URL for ${provider.name}:`, url, `(from ${envVarName})`);
 
             const status = await checkProviderStatus(url, provider.name);
+
             return {
               ...status,
               enabled: provider.settings.enabled ?? false,
             };
-          })
+          }),
       );
 
       setActiveProviders(statuses);
@@ -269,19 +271,20 @@ export default function DebugTab() {
       console.log(`[Debug] Checking for updates against ${branchToCheck} branch`);
 
       const localCommitResponse = await fetch(GITHUB_URLS.commitJson(branchToCheck));
+
       if (!localCommitResponse.ok) {
         throw new Error('Failed to fetch local commit info');
       }
 
-      const localCommitData = await localCommitResponse.json() as CommitData;
+      const localCommitData = (await localCommitResponse.json()) as CommitData;
       const remoteCommitHash = localCommitData.commit;
       const currentCommitHash = versionHash;
 
       if (remoteCommitHash !== currentCommitHash) {
         setUpdateMessage(
           `Update available from ${branchToCheck} branch!\n` +
-          `Current: ${currentCommitHash.slice(0, 7)}\n` +
-          `Latest: ${remoteCommitHash.slice(0, 7)}`
+            `Current: ${currentCommitHash.slice(0, 7)}\n` +
+            `Latest: ${remoteCommitHash.slice(0, 7)}`,
         );
       } else {
         setUpdateMessage(`You are on the latest version from the ${branchToCheck} branch`);
@@ -309,7 +312,7 @@ export default function DebugTab() {
       })),
       Version: {
         hash: versionHash.slice(0, 7),
-        branch: useLatestBranch ? 'main' : 'stable'
+        branch: useLatestBranch ? 'main' : 'stable',
       },
       Timestamp: new Date().toISOString(),
     };

app/components/settings/features/FeaturesTab.tsx

@@ -3,7 +3,15 @@ import { Switch } from '~/components/ui/Switch';
 import { useSettings } from '~/lib/hooks/useSettings';
 
 export default function FeaturesTab() {
-  const { debug, enableDebugMode, isLocalModel, enableLocalModels, eventLogs, enableEventLogs, useLatestBranch, enableLatestBranch } = useSettings();
+  const {
+    debug,
+    enableDebugMode,
+    isLocalModel,
+    enableLocalModels,
+    enableEventLogs,
+    useLatestBranch,
+    enableLatestBranch,
+  } = useSettings();
 
   const handleToggle = (enabled: boolean) => {
     enableDebugMode(enabled);
@@ -22,7 +30,9 @@ export default function FeaturesTab() {
           <div className="flex items-center justify-between">
             <div>
               <span className="text-bolt-elements-textPrimary">Use Main Branch</span>
-              <p className="text-sm text-bolt-elements-textSecondary">Check for updates against the main branch instead of stable</p>
+              <p className="text-sm text-bolt-elements-textSecondary">
+                Check for updates against the main branch instead of stable
+              </p>
             </div>
             <Switch className="ml-auto" checked={useLatestBranch} onCheckedChange={enableLatestBranch} />
           </div>

app/components/settings/providers/ProvidersTab.tsx

@@ -56,7 +56,8 @@ export default function ProvidersTab() {
             <div className="flex items-center gap-2">
               <img
                 src={`/icons/${provider.name}.svg`} // Attempt to load the specific icon
-                onError={(e) => { // Fallback to default icon on error
+                onError={(e) => {
+                  // Fallback to default icon on error
                   e.currentTarget.src = DefaultIcon;
                 }}
                 alt={`${provider.name} icon`}

app/lib/.server/llm/switchable-stream.ts

@@ -1,5 +1,5 @@
 export default class SwitchableStream extends TransformStream {
-  private _controller: TransformStreamDefaultController | null = null;
+  _controller: TransformStreamDefaultController | null = null;
   private _currentReader: ReadableStreamDefaultReader | null = null;
   private _switches = 0;
 

app/lib/hooks/useSettings.tsx

@@ -28,12 +28,17 @@ export function useSettings() {
   // Function to check if we're on stable version
   const checkIsStableVersion = async () => {
     try {
-      const stableResponse = await fetch('https://raw.githubusercontent.com/stackblitz-labs/bolt.diy/stable/app/commit.json');
+      const stableResponse = await fetch(
+        'https://raw.githubusercontent.com/stackblitz-labs/bolt.diy/stable/app/commit.json',
+      );
+
       if (!stableResponse.ok) {
         console.warn('Failed to fetch stable commit info');
         return false;
       }
-      const stableData = await stableResponse.json() as CommitData;
+
+      const stableData = (await stableResponse.json()) as CommitData;
+
       return commit.commit === stableData.commit;
     } catch (error) {
       console.warn('Error checking stable version:', error);
@@ -86,9 +91,10 @@ export function useSettings() {
 
     // load latest branch setting from cookies or determine based on version
     const savedLatestBranch = Cookies.get('useLatestBranch');
+
     if (savedLatestBranch === undefined) {
       // If setting hasn't been set by user, check version
-      checkIsStableVersion().then(isStable => {
+      checkIsStableVersion().then((isStable) => {
         const shouldUseLatest = !isStable;
         latestBranch.set(shouldUseLatest);
         Cookies.set('useLatestBranch', String(shouldUseLatest));

app/routes/api.chat.ts

@@ -9,17 +9,15 @@ export async function action(args: ActionFunctionArgs) {
   return chatAction(args);
 }
 
-function parseCookies(cookieHeader: string) {
-  const cookies: any = {};
+function parseCookies(cookieHeader: string): Record<string, string> {
+  const cookies: Record<string, string> = {};
 
-  // Split the cookie string by semicolons and spaces
   const items = cookieHeader.split(';').map((cookie) => cookie.trim());
 
   items.forEach((item) => {
     const [name, ...rest] = item.split('=');
 
     if (name && rest) {
-      // Decode the name and value, and join value parts in case it contains '='
       const decodedName = decodeURIComponent(name.trim());
       const decodedValue = decodeURIComponent(rest.join('=').trim());
       cookies[decodedName] = decodedValue;
@@ -36,19 +34,43 @@ async function chatAction({ context, request }: ActionFunctionArgs) {
   }>();
 
   const cookieHeader = request.headers.get('Cookie');
-
-  // Parse the cookie's value (returns an object or null if no cookie exists)
   const apiKeys = JSON.parse(parseCookies(cookieHeader || '').apiKeys || '{}');
   const providerSettings: Record<string, IProviderSetting> = JSON.parse(
     parseCookies(cookieHeader || '').providers || '{}',
   );
 
   const stream = new SwitchableStream();
 
+  const cumulativeUsage = {
+    completionTokens: 0,
+    promptTokens: 0,
+    totalTokens: 0,
+  };
+
   try {
     const options: StreamingOptions = {
       toolChoice: 'none',
-      onFinish: async ({ text: content, finishReason }) => {
+      onFinish: async ({ text: content, finishReason, usage }) => {
+        console.log('usage', usage);
+
+        if (usage && stream._controller) {
+          cumulativeUsage.completionTokens += usage.completionTokens || 0;
+          cumulativeUsage.promptTokens += usage.promptTokens || 0;
+          cumulativeUsage.totalTokens += usage.totalTokens || 0;
+
+          // Send usage info in message metadata for assistant messages
+          const usageMetadata = `0:"[Usage: ${JSON.stringify({
+            completionTokens: cumulativeUsage.completionTokens,
+            promptTokens: cumulativeUsage.promptTokens,
+            totalTokens: cumulativeUsage.totalTokens,
+          })}\n]"`;
+
+          console.log(usageMetadata);
+
+          const encodedData = new TextEncoder().encode(usageMetadata);
+          stream._controller.enqueue(encodedData);
+        }
+
         if (finishReason !== 'length') {
           return stream.close();
         }
@@ -73,7 +95,7 @@ async function chatAction({ context, request }: ActionFunctionArgs) {
           providerSettings,
         });
 
-        return stream.switchSource(result.toAIStream());
+        return stream.switchSource(result.toDataStream());
       },
     };
 
@@ -86,7 +108,7 @@ async function chatAction({ context, request }: ActionFunctionArgs) {
       providerSettings,
     });
 
-    stream.switchSource(result.toAIStream());
+    stream.switchSource(result.toDataStream());
 
     return new Response(stream.readable, {
       status: 200,
@@ -95,7 +117,7 @@ async function chatAction({ context, request }: ActionFunctionArgs) {
       },
     });
   } catch (error: any) {
-    console.log(error);
+    console.error(error);
 
     if (error.message?.includes('API key')) {
       throw new Response('Invalid or missing API key', {