From 927cdd0837f68f057c24c7aef2bca5fb56ad17a7 Mon Sep 17 00:00:00 2001
From: Mike Fortman <michael.fortman@datastax.com>
Date: Tue, 24 Jun 2025 18:07:19 -0500
Subject: [PATCH] Fix: Update display time for very short operations and fix
 outputs and titles for tool calling executions (#8718)

* update display time for very short operations and fix outputs and titles

* [autofix.ci] apply automated fixes

* [autofix.ci] apply automated fixes (attempt 2/3)

* cleanup

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: Edwin Jose <edwin.jose@datastax.com>
---
 .../base/langflow/base/agents/events.py       | 40 +++++++++++++++----
 .../core/chatComponents/DurationDisplay.tsx   |  3 +-
 2 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/src/backend/base/langflow/base/agents/events.py b/src/backend/base/langflow/base/agents/events.py
index 42ffa448d..62c130503 100644
--- a/src/backend/base/langflow/base/agents/events.py
+++ b/src/backend/base/langflow/base/agents/events.py
@@ -161,7 +161,7 @@ async def handle_on_tool_start(
     tool_content = ToolContent(
         type="tool_use",
         name=tool_name,
-        input=tool_input,
+        tool_input=tool_input,
         output=None,
         error=None,
         header={"title": f"Accessing **{tool_name}**", "icon": "Hammer"},
@@ -191,13 +191,39 @@ async def handle_on_tool_end(
     tool_content = tool_blocks_map.get(tool_key)
 
     if tool_content and isinstance(tool_content, ToolContent):
-        tool_content.output = event["data"].get("output")
-        duration = _calculate_duration(start_time)
-        tool_content.duration = duration
-        tool_content.header = {"title": f"Executed **{tool_content.name}**", "icon": "Hammer"}
-
+        # Call send_message_method first to get the updated message structure
         agent_message = await send_message_method(message=agent_message)
-        new_start_time = perf_counter()  # Get new start time for next operation
+        new_start_time = perf_counter()
+
+        # Now find and update the tool content in the current message
+        duration = _calculate_duration(start_time)
+        tool_key = f"{tool_name}_{run_id}"
+
+        # Find the corresponding tool content in the updated message
+        updated_tool_content = None
+        if agent_message.content_blocks and agent_message.content_blocks[0].contents:
+            for content in agent_message.content_blocks[0].contents:
+                if (
+                    isinstance(content, ToolContent)
+                    and content.name == tool_name
+                    and content.tool_input == tool_content.tool_input
+                ):
+                    updated_tool_content = content
+                    break
+
+        # Update the tool content that's actually in the message
+        if updated_tool_content:
+            updated_tool_content.duration = duration
+            updated_tool_content.header = {"title": f"Executed **{updated_tool_content.name}**", "icon": "Hammer"}
+            updated_tool_content.output = event["data"].get("output")
+
+            # Update the map reference
+            tool_blocks_map[tool_key] = updated_tool_content
+
+            for content in agent_message.content_blocks[0].contents:
+                if isinstance(content, ToolContent):
+                    header_title = content.header.get("title", "N/A") if content.header else "None"
+
         return agent_message, new_start_time
     return agent_message, start_time
 
diff --git a/src/frontend/src/components/core/chatComponents/DurationDisplay.tsx b/src/frontend/src/components/core/chatComponents/DurationDisplay.tsx
index 5facf421d..f8b4f11d4 100644
--- a/src/frontend/src/components/core/chatComponents/DurationDisplay.tsx
+++ b/src/frontend/src/components/core/chatComponents/DurationDisplay.tsx
@@ -41,7 +41,8 @@ export default function DurationDisplay({
 
   const displayTime = duration ?? durations[chatId] ?? 0;
   const secondsValue = displayTime / 1000;
-  const humanizedTime = `${secondsValue.toFixed(1)}s`;
+  const humanizedTime =
+    secondsValue < 0.05 ? "< 0.1s" : `${secondsValue.toFixed(1)}s`;
 
   return (
     <div