diff --git a/src/components/ChatMessage.tsx b/src/components/ChatMessage.tsx index a959444..a57912b 100644 --- a/src/components/ChatMessage.tsx +++ b/src/components/ChatMessage.tsx @@ -539,6 +539,13 @@ export const ChatMessageComponent = memo(function ChatMessageComponent({ message )} {message.timestamp && } + {!isUser && message.generationTimeMs != null && ( + + · {message.generationTimeMs < 1000 + ? `${message.generationTimeMs}ms` + : `${(message.generationTimeMs / 1000).toFixed(1)}s`} + + )} {isUser && message.sendStatus === 'sending' && ( )} diff --git a/src/hooks/useGateway.ts b/src/hooks/useGateway.ts index a155323..c7a2848 100644 --- a/src/hooks/useGateway.ts +++ b/src/hooks/useGateway.ts @@ -54,6 +54,8 @@ export function useGateway() { const [activeSessions, setActiveSessions] = useState>(new Set()); const [unreadSessions, setUnreadSessions] = useState>(new Set()); const [agentIdentity, setAgentIdentity] = useState(null); + /** Map of runId → generation duration (ms), preserved across loadHistory reloads */ + const generationTimesRef = useRef>(new Map()); const handleAgentEvent = useCallback((payload: JsonPayload) => { if (payload?.stream !== 'tool') return; @@ -228,6 +230,18 @@ export function useGateway() { merged.push(msg); } } + // Apply stored generation time to the last assistant message if available + const genKey = sessionKey + ':latest'; + const genTime = generationTimesRef.current.get(genKey); + if (genTime) { + generationTimesRef.current.delete(genKey); + for (let i = merged.length - 1; i >= 0; i--) { + if (merged[i].role === 'assistant') { + merged[i] = { ...merged[i], generationTimeMs: genTime }; + break; + } + } + } setMessages(merged); } } catch { @@ -342,10 +356,16 @@ export function useGateway() { blocks, isStreaming: true, runId, + streamStartedAt: Date.now(), }; return [...prev, msg]; }); } else if (state === 'final') { + // Compute generation time from the streaming message before history reload replaces it + const lastMsg = messagesRef.current[messagesRef.current.length - 1]; + if (lastMsg?.role === 'assistant' && lastMsg.streamStartedAt) { + generationTimesRef.current.set(activeSessionRef.current + ':latest', Date.now() - lastMsg.streamStartedAt); + } currentRunIdRef.current = null; setIsGenerating(false); loadHistory(activeSessionRef.current); diff --git a/src/types/index.ts b/src/types/index.ts index 012b115..58f20c7 100644 --- a/src/types/index.ts +++ b/src/types/index.ts @@ -10,6 +10,10 @@ export interface ChatMessage { metadata?: Record; /** Optimistic send status for user messages */ sendStatus?: 'sending' | 'sent' | 'error'; + /** Timestamp (ms) when streaming started for this message */ + streamStartedAt?: number; + /** Total generation time in milliseconds (set when streaming ends) */ + generationTimeMs?: number; } export type MessageBlock =