Merge pull request #3 from Zipstack/fix-token-usage-endpoint-OSS-mode

Fix skip token-usage API calls in OSS mode — fetch only in cloud

Author: wicky-zipstack

backend/backend/core/routers/chat_message/views/message_views.py

@@ -52,7 +52,7 @@ def persist_prompt(self, request: Request, *args, **kwargs) -> Response:
     def get_token_usage(self, request, project_id=None, chat_id=None, chat_message_id=None, *args, **kwargs) -> Response:
         """
         Get token usage data for a specific chat message.
-        
+
         Returns:
             - remaining_balance: Current token balance
             - total_consumed: Total tokens consumed
@@ -67,18 +67,18 @@ def get_token_usage(self, request, project_id=None, chat_id=None, chat_message_i
                     {"error": "Organization header is required"},
                     status=status.HTTP_400_BAD_REQUEST
                 )
-            
+
             # Get token usage data using the existing function
             token_data = get_token_usage_data(org_id, str(chat_message_id), str(chat_id))
-            
+
             if token_data is None:
                 return Response(
                     {"error": "Failed to retrieve token usage data"},
                     status=status.HTTP_500_INTERNAL_SERVER_ERROR
                 )
-            
+
             return Response(token_data, status=status.HTTP_200_OK)
-            
+
         except Exception as e:
             return Response(
                 {"error": f"Error retrieving token usage: {str(e)}"},

frontend/src/ide/chat-ai/ExistingChat.jsx

@@ -16,7 +16,17 @@ import { TodoGuide } from "./TodoGuide";
 import { OnboardingGuide } from "./OnboardingGuide";
 import { OnboardingCompletionPopup } from "./OnboardingCompletionPopup";
 import { useNotificationService } from "../../service/notification-service";
+import { useAxiosPrivate } from "../../service/axios-service";
 import { SpinnerLoader } from "../../widgets/spinner_loader";
+import { useSessionStore } from "../../store/session-store";
+
+// Cloud-only: fetch per-message token usage (unavailable in OSS — import fails gracefully)
+let getTokenUsage = null;
+try {
+  ({ getTokenUsage } = require("../../plugins/token-management/token-usage"));
+} catch {
+  // OSS: token usage API not available
+}
 
 const ExistingChat = memo(function ExistingChat({
   selectedChatId,
@@ -61,8 +71,9 @@ const ExistingChat = memo(function ExistingChat({
   onSkipCurrentTask,
   onSendButtonClick,
 }) {
-  const { getChatMessagesByChatId, getTokenUsage, updateChatName } =
-    useChatAIService();
+  const { getChatMessagesByChatId, updateChatName } = useChatAIService();
+  const axiosPrivate = useAxiosPrivate();
+  const isCloud = useSessionStore((state) => state.sessionDetails?.is_cloud);
   const chatContainerRef = useRef(null);
 
   const [isLoadingChats, setIsLoadingChats] = useState(false);
@@ -262,10 +273,15 @@ const ExistingChat = memo(function ExistingChat({
           : [msg.response].filter(Boolean),
       }));
 
-      // Fetch token usage for all messages to display in historical conversations
-      if (updatedData.length > 0) {
+      // Fetch token usage for all messages to display in historical conversations.
+      // Only available in cloud mode via the token-usage plugin.
+      if (getTokenUsage && isCloud && updatedData.length > 0) {
         const tokenUsagePromises = updatedData.map((msg) =>
-          getTokenUsage(selectedChatId, msg.chat_message_id).catch(() => null)
+          getTokenUsage(
+            axiosPrivate,
+            selectedChatId,
+            msg.chat_message_id
+          ).catch(() => null)
         );
         const tokenUsageResults = await Promise.all(tokenUsagePromises);
 

frontend/src/ide/chat-ai/services.js

@@ -134,22 +134,6 @@ export function useChatAIService() {
     return response.data;
   };
 
-  const getTokenUsage = async (chatId, chatMessageId) => {
-    const url = `/api/v1/visitran/${orgId}/project/${projectId}/chat/${chatId}/chat-message/${chatMessageId}/token-usage/`;
-    try {
-      const response = await axiosPrivate.get(url, {
-        headers: {
-          "Content-Type": "application/json",
-          "X-Organization": orgId,
-        },
-      });
-      return response.data || null;
-    } catch (error) {
-      console.warn("Failed to fetch token usage:", error);
-      return null;
-    }
-  };
-
   return {
     getAllChats,
     deleteChatById,
@@ -161,6 +145,5 @@ export function useChatAIService() {
     getChatLlmModels,
     completeOnboardingTask,
     getOnboardingStatus,
-    getTokenUsage,
   };
 }