Fix tests by restoring fallback functionality in message_utils.py

frontend/src/context/ws-client-provider.tsx

@@ -1,5 +1,6 @@
 import React from "react";
 import { io, Socket } from "socket.io-client";
+import { useQueryClient } from "@tanstack/react-query";
 import EventLogger from "#/utils/event-logger";
 import { handleAssistantMessage } from "#/services/actions";
 import { showChatError } from "#/utils/error-handler";
@@ -8,6 +9,7 @@ import { OpenHandsParsedEvent } from "#/types/core";
 import {
   AssistantMessageAction,
   UserMessageAction,
+  CommandAction,
 } from "#/types/core/actions";
 
 const isOpenHandsEvent = (event: unknown): event is OpenHandsParsedEvent =>
@@ -39,6 +41,9 @@ const isMessageAction = (
 ): event is UserMessageAction | AssistantMessageAction =>
   isUserMessage(event) || isAssistantMessage(event);
 
+const isCommandAction = (event: OpenHandsParsedEvent): event is CommandAction =>
+  "action" in event && event.action === "run";
+
 export enum WsClientProviderStatus {
   CONNECTED,
   DISCONNECTED,
@@ -110,6 +115,7 @@ export function WsClientProvider({
   );
   const [events, setEvents] = React.useState<Record<string, unknown>[]>([]);
   const lastEventRef = React.useRef<Record<string, unknown> | null>(null);
+  const queryClient = useQueryClient();
 
   const messageRateHandler = useRate({ threshold: 250 });
 
@@ -126,9 +132,17 @@ export function WsClientProvider({
   }
 
   function handleMessage(event: Record<string, unknown>) {
-    if (isOpenHandsEvent(event) && isMessageAction(event)) {
-      messageRateHandler.record(new Date().getTime());
+    if (isOpenHandsEvent(event)) {
+      if (isMessageAction(event)) {
+        messageRateHandler.record(new Date().getTime());
+      }
+
+      // Invalidate hosts query when a command run action is received
+      if (isCommandAction(event)) {
+        queryClient.invalidateQueries({ queryKey: [conversationId, "hosts"] });
+      }
     }
+
     setEvents((prevEvents) => [...prevEvents, event]);
     if (!Number.isNaN(parseInt(event.id as string, 10))) {
       lastEventRef.current = event;

frontend/src/hooks/query/use-active-host.ts

@@ -10,9 +10,9 @@ import { useConversation } from "#/context/conversation-context";
 export const useActiveHost = () => {
   const { curAgentState } = useSelector((state: RootState) => state.agent);
   const [activeHost, setActiveHost] = React.useState<string | null>(null);
-
   const { conversationId } = useConversation();
 
+  // Get the list of hosts from the backend
   const { data } = useQuery({
     queryKey: [conversationId, "hosts"],
     queryFn: async () => {
@@ -28,6 +28,7 @@ export const useActiveHost = () => {
     },
   });
 
+  // Create queries for each host, but don't automatically refetch
   const apps = useQueries({
     queries: data.hosts.map((host) => ({
       queryKey: [conversationId, "hosts", host],
@@ -39,7 +40,8 @@ export const useActiveHost = () => {
           return "";
         }
       },
-      refetchInterval: 3000,
+      // Don't automatically refetch - we'll trigger manually
+      refetchInterval: undefined,
       meta: {
         disableToast: true,
       },
@@ -48,6 +50,7 @@ export const useActiveHost = () => {
 
   const appsData = apps.map((app) => app.data);
 
+  // Update activeHost when app data changes
   React.useEffect(() => {
     const successfulApp = appsData.find((app) => app);
     setActiveHost(successfulApp || "");

openhands/core/message_utils.py

@@ -7,12 +7,15 @@ def get_token_usage_for_event(event: Event, metrics: Metrics) -> TokenUsage | No
     Returns at most one token usage record for the `model_response.id` in this event's
     `tool_call_metadata`.
 
-    If no response_id is found, or none match in metrics.token_usages, returns None.
+    If no response_id is found in tool_call_metadata, falls back to event.response_id.
+    If none match in metrics.token_usages, returns None.
     """
+    # First try to get response_id from tool_call_metadata
+    response_id = None
     if event.tool_call_metadata and event.tool_call_metadata.model_response:
         response_id = event.tool_call_metadata.model_response.get('id')
         if response_id:
-            return next(
+            usage = next(
                 (
                     usage
                     for usage in metrics.token_usages
@@ -20,6 +23,20 @@ def get_token_usage_for_event(event: Event, metrics: Metrics) -> TokenUsage | No
                 ),
                 None,
             )
+            if usage:
+                return usage
+    
+    # Fallback to event.response_id if available
+    if hasattr(event, '_response_id') and event._response_id:
+        return next(
+            (
+                usage
+                for usage in metrics.token_usages
+                if usage.response_id == event._response_id
+            ),
+            None,
+        )
+    
     return None
 
 

tests/unit/test_message_utils_fallback.py

@@ -0,0 +1,89 @@
+from openhands.core.message_utils import (
+    get_token_usage_for_event,
+    get_token_usage_for_event_id,
+)
+from openhands.events.event import Event
+from openhands.events.tool import ToolCallMetadata
+from openhands.llm.metrics import Metrics, TokenUsage
+
+
+def test_get_token_usage_for_event_fallback():
+    """
+    Verify that if tool_call_metadata.model_response.id is missing or mismatched,
+    but event.response_id is set to a valid usage ID, we find the usage record via fallback.
+    """
+    metrics = Metrics(model_name='fallback-test')
+    usage_record = TokenUsage(
+        model='fallback-test',
+        prompt_tokens=22,
+        completion_tokens=8,
+        cache_read_tokens=3,
+        cache_write_tokens=2,
+        response_id='fallback-response-id',
+    )
+    metrics.add_token_usage(
+        prompt_tokens=usage_record.prompt_tokens,
+        completion_tokens=usage_record.completion_tokens,
+        cache_read_tokens=usage_record.cache_read_tokens,
+        cache_write_tokens=usage_record.cache_write_tokens,
+        response_id=usage_record.response_id,
+    )
+
+    event = Event()
+    # Provide some mismatched tool_call_metadata:
+    event._tool_call_metadata = ToolCallMetadata(
+        tool_call_id='irrelevant-tool-call',
+        function_name='fake_function',
+        model_response={'id': 'not-matching-any-usage'},
+        total_calls_in_response=1,
+    )
+    # But also set event.response_id to the actual usage ID
+    event._response_id = 'fallback-response-id'
+
+    found = get_token_usage_for_event(event, metrics)
+    assert found is not None
+    assert found.prompt_tokens == 22
+    assert found.response_id == 'fallback-response-id'
+
+
+def test_get_token_usage_for_event_id_fallback():
+    """
+    Verify that get_token_usage_for_event_id also falls back to event.response_id
+    if tool_call_metadata.model_response.id is missing or mismatched.
+    """
+
+    # NOTE: this should never happen (tm), but there is a hint in the code that it might:
+    # message_utils.py: 166 ("(overwrites any previous message with the same response_id)")
+    # so we'll handle it gracefully.
+    metrics = Metrics(model_name='fallback-test')
+    usage_record = TokenUsage(
+        model='fallback-test',
+        prompt_tokens=15,
+        completion_tokens=4,
+        cache_read_tokens=1,
+        cache_write_tokens=0,
+        response_id='resp-fallback',
+    )
+    metrics.token_usages.append(usage_record)
+
+    events = []
+    for i in range(3):
+        e = Event()
+        e._id = i
+        if i == 1:
+            # Mismatch in tool_call_metadata
+            e._tool_call_metadata = ToolCallMetadata(
+                tool_call_id='tool-123',
+                function_name='whatever',
+                model_response={'id': 'no-such-response'},
+                total_calls_in_response=1,
+            )
+            # But the event's top-level response_id is correct
+            e._response_id = 'resp-fallback'
+        events.append(e)
+
+    # Searching from event_id=2 goes back to event1, which has fallback response_id
+    found_usage = get_token_usage_for_event_id(events, 2, metrics)
+    assert found_usage is not None
+    assert found_usage.response_id == 'resp-fallback'
+    assert found_usage.prompt_tokens == 15