From aeb5382f0a429f7b99a44e306c708286f766c87a Mon Sep 17 00:00:00 2001
From: shine1i <wasimysdev@gmail.com>
Date: Sun, 1 Feb 2026 09:08:23 +0100
Subject: [PATCH] feat: track and display reasoning duration, enhance runtime
 with adapter for copying during inference and edit and UI integration

---
 .../src/components/assistant-ui/reasoning.tsx | 10 ++++-
 .../src/components/assistant-ui/thread.tsx    | 44 ++++++++++++++++---
 frontend/src/features/chat/adapter.ts         | 18 +++++++-
 .../src/features/chat/runtime-provider.tsx    |  4 +-
 frontend/src/features/chat/types.ts           |  1 +
 5 files changed, 67 insertions(+), 10 deletions(-)

diff --git a/frontend/src/components/assistant-ui/reasoning.tsx b/frontend/src/components/assistant-ui/reasoning.tsx
index 3dd2585d2e..e43eb2ca4b 100644
--- a/frontend/src/components/assistant-ui/reasoning.tsx
+++ b/frontend/src/components/assistant-ui/reasoning.tsx
@@ -165,7 +165,7 @@ function ReasoningTrigger({
             Thinking...
           </AnimatedShinyText>
         ) : (
-          <span>Thought for {duration ?? 0}s</span>
+          <span>Thought for {duration ?? 0} seconds</span>
         )}
       </span>
       <ChevronDownIcon
@@ -276,6 +276,12 @@ const ReasoningGroupImpl: ReasoningGroupComponent = ({
     return lastIndex >= startIndex && lastIndex <= endIndex;
   });
 
+  const persistedDuration = useAuiState(({ message }) => {
+    const d = (message.metadata?.custom as Record<string, unknown>)
+      ?.reasoningDuration;
+    return typeof d === "number" ? d : 0;
+  });
+
   const [manualOpen, setManualOpen] = useState(false);
   const [duration, setDuration] = useState<number>(0);
   const startTimeRef = useRef<number | null>(null);
@@ -319,7 +325,7 @@ const ReasoningGroupImpl: ReasoningGroupComponent = ({
     >
       <ReasoningTrigger
         active={isReasoningStreaming}
-        duration={duration}
+        duration={duration || persistedDuration}
       />
       <ReasoningContent
         aria-busy={isReasoningStreaming}
diff --git a/frontend/src/components/assistant-ui/thread.tsx b/frontend/src/components/assistant-ui/thread.tsx
index 2ddf6fa65e..400fcb6c81 100644
--- a/frontend/src/components/assistant-ui/thread.tsx
+++ b/frontend/src/components/assistant-ui/thread.tsx
@@ -19,6 +19,9 @@ import {
   MessagePrimitive,
   SuggestionPrimitive,
   ThreadPrimitive,
+  useComposerRuntime,
+  useMessageRuntime,
+  useThreadRuntime,
 } from "@assistant-ui/react";
 import { motion } from "framer-motion";
 import {
@@ -333,12 +336,21 @@ const UserMessage: FC = () => {
 const UserActionBar: FC = () => {
   return (
     <ActionBarPrimitive.Root
-      hideWhenRunning={true}
       autohide="not-last"
-      className="aui-user-action-bar-root flex flex-col items-end"
+      className="aui-user-action-bar-root flex items-center"
     >
+      <ActionBarPrimitive.Copy asChild={true}>
+        <TooltipIconButton tooltip="Copy">
+          <AuiIf condition={({ message }) => message.isCopied}>
+            <CheckIcon />
+          </AuiIf>
+          <AuiIf condition={({ message }) => !message.isCopied}>
+            <CopyIcon />
+          </AuiIf>
+        </TooltipIconButton>
+      </ActionBarPrimitive.Copy>
       <ActionBarPrimitive.Edit asChild={true}>
-        <TooltipIconButton tooltip="Edit" className="aui-user-action-edit p-4">
+        <TooltipIconButton tooltip="Edit" className="aui-user-action-edit">
           <PencilIcon />
         </TooltipIconButton>
       </ActionBarPrimitive.Edit>
@@ -347,6 +359,10 @@ const UserActionBar: FC = () => {
 };
 
 const EditComposer: FC = () => {
+  const threadRuntime = useThreadRuntime();
+  const composerRuntime = useComposerRuntime();
+  const messageRuntime = useMessageRuntime();
+
   return (
     <MessagePrimitive.Root className="aui-edit-composer-wrapper mx-auto flex w-full max-w-(--thread-max-width) flex-col px-2 py-3">
       <ComposerPrimitive.Root className="aui-edit-composer-root ml-auto flex w-full max-w-[85%] flex-col rounded-2xl bg-muted">
@@ -360,9 +376,25 @@ const EditComposer: FC = () => {
               Cancel
             </Button>
           </ComposerPrimitive.Cancel>
-          <ComposerPrimitive.Send asChild={true}>
-            <Button size="sm">Update</Button>
-          </ComposerPrimitive.Send>
+          <Button
+            size="sm"
+            onClick={() => {
+              const newText = composerRuntime.getState().text;
+              const originalText = messageRuntime.unstable_getCopyText();
+
+              if (newText === originalText) {
+                composerRuntime.cancel();
+                return;
+              }
+
+              if (threadRuntime.getState().isRunning) {
+                threadRuntime.cancelRun();
+              }
+              composerRuntime.send();
+            }}
+          >
+            Update
+          </Button>
         </div>
       </ComposerPrimitive.Root>
     </MessagePrimitive.Root>
diff --git a/frontend/src/features/chat/adapter.ts b/frontend/src/features/chat/adapter.ts
index 3dcf881ad2..c4e103f62d 100644
--- a/frontend/src/features/chat/adapter.ts
+++ b/frontend/src/features/chat/adapter.ts
@@ -65,6 +65,9 @@ export function createStreamAdapter(apiUrl: string = API): ChatModelAdapter {
       }
       const decoder = new TextDecoder();
       let text = "";
+      let reasoningStart: number | null = null;
+      let reasoningDuration = 0;
+
       while (true) {
         const { done, value } = await reader.read();
         if (done) {
@@ -72,8 +75,21 @@ export function createStreamAdapter(apiUrl: string = API): ChatModelAdapter {
         }
         text += decoder.decode(value, { stream: true });
         const parts = parseThinkTags(text) ?? [];
+
+        if (parts.some((p) => p.type === "reasoning") && !reasoningStart) {
+          reasoningStart = Date.now();
+        }
+        if (text.includes("</think>") && reasoningStart && !reasoningDuration) {
+          reasoningDuration = Math.round(
+            (Date.now() - reasoningStart) / 1000,
+          );
+        }
+
         if (parts.length > 0) {
-          yield { content: parts };
+          yield {
+            content: parts,
+            metadata: { custom: { reasoningDuration } },
+          };
         }
       }
     },
diff --git a/frontend/src/features/chat/runtime-provider.tsx b/frontend/src/features/chat/runtime-provider.tsx
index d411d595f8..593fe42b31 100644
--- a/frontend/src/features/chat/runtime-provider.tsx
+++ b/frontend/src/features/chat/runtime-provider.tsx
@@ -52,7 +52,7 @@ function toThreadMessage(m: MessageRecord): ThreadMessage {
     role: "assistant" as const,
     status: { type: "complete" as const, reason: "unknown" as const },
     metadata: {
-      custom: {},
+      custom: (m.metadata as Record<string, unknown>) ?? {},
       steps: [],
       unstable_annotations: [],
       unstable_data: [],
@@ -183,11 +183,13 @@ function ThreadHistoryProvider({
         const content = Array.isArray(message.content)
           ? JSON.parse(JSON.stringify(message.content))
           : [];
+        const custom = message.metadata?.custom;
         await db.messages.put({
           id: message.id,
           threadId: remoteId,
           role: message.role,
           content,
+          ...(custom && Object.keys(custom).length > 0 && { metadata: custom }),
           createdAt: message.createdAt?.getTime() ?? Date.now(),
         });
       },
diff --git a/frontend/src/features/chat/types.ts b/frontend/src/features/chat/types.ts
index 655aba7095..2ef9b6fdc5 100644
--- a/frontend/src/features/chat/types.ts
+++ b/frontend/src/features/chat/types.ts
@@ -18,6 +18,7 @@ export interface MessageRecord {
   threadId: string;
   role: import("@assistant-ui/react").ThreadMessage["role"];
   content: import("@assistant-ui/react").ThreadMessage["content"];
+  metadata?: Record<string, unknown>;
   createdAt: number;
 }