From 19d7d86030beb7211d5ac36719f83ab7ff743f4b Mon Sep 17 00:00:00 2001
From: haidang1810 <99536085+haidang1810@users.noreply.github.com>
Date: Mon, 23 Mar 2026 15:17:31 +0700
Subject: [PATCH] fix(ui): use server timestamps for chat messages (#332)
 (#367)

---
 internal/providers/types.go                       | 6 ++++++
 internal/store/pg/sessions.go                     | 6 ++++++
 ui/web/src/pages/chat/hooks/use-chat-messages.ts  | 5 ++++-
 ui/web/src/pages/sessions/session-detail-page.tsx | 5 ++++-
 ui/web/src/types/session.ts                       | 1 +
 5 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/internal/providers/types.go b/internal/providers/types.go
index 0b3d204d..526982a7 100644
--- a/internal/providers/types.go
+++ b/internal/providers/types.go
@@ -3,6 +3,7 @@ package providers
 import (
 	"context"
 	"encoding/json"
+	"time"
 )
 
 // Options keys used in ChatRequest.Options across providers.
@@ -111,6 +112,11 @@ type Message struct {
 	// RawAssistantContent carries raw provider content blocks through tool loop iterations.
 	// Anthropic requires thinking blocks to be passed back exactly as received.
 	RawAssistantContent json.RawMessage `json:"-"`
+
+	// CreatedAt records when this message was added to the session.
+	// Pointer type so that older messages (stored before this field existed) deserialize as nil,
+	// allowing the frontend to fall back to synthetic timestamps.
+	CreatedAt *time.Time `json:"created_at,omitempty"`
 }
 
 // ToolCall represents a tool invocation requested by the LLM.
diff --git a/internal/store/pg/sessions.go b/internal/store/pg/sessions.go
index 3ea7c243..f26a6937 100644
--- a/internal/store/pg/sessions.go
+++ b/internal/store/pg/sessions.go
@@ -140,6 +140,12 @@ func (s *PGSessionStore) AddMessage(ctx context.Context, key string, msg provide
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
+	// Stamp message creation time if not already set.
+	if msg.CreatedAt == nil {
+		now := time.Now()
+		msg.CreatedAt = &now
+	}
+
 	data := s.getOrInit(ctx, key)
 	data.Messages = append(data.Messages, msg)
 	data.Updated = time.Now()
diff --git a/ui/web/src/pages/chat/hooks/use-chat-messages.ts b/ui/web/src/pages/chat/hooks/use-chat-messages.ts
index bb3f36a1..5a4518ba 100644
--- a/ui/web/src/pages/chat/hooks/use-chat-messages.ts
+++ b/ui/web/src/pages/chat/hooks/use-chat-messages.ts
@@ -90,7 +90,10 @@ export function useChatMessages(sessionKey: string, agentId: string) {
       const msgs: ChatMessage[] = allMsgs.map((m: Message, i: number) => {
         const chatMsg: ChatMessage = {
           ...m,
-          timestamp: Date.now() - (allMsgs.length - i) * 1000,
+          // Use server-provided created_at; fall back to synthetic spacing for older messages.
+          timestamp: m.created_at
+            ? new Date(m.created_at).getTime()
+            : Date.now() - (allMsgs.length - i) * 1000,
         };
         // Convert persisted media_refs to mediaItems for gallery display
         if (m.role === "assistant" && m.media_refs && m.media_refs.length > 0) {
diff --git a/ui/web/src/pages/sessions/session-detail-page.tsx b/ui/web/src/pages/sessions/session-detail-page.tsx
index 540b1c8c..e44a31a1 100644
--- a/ui/web/src/pages/sessions/session-detail-page.tsx
+++ b/ui/web/src/pages/sessions/session-detail-page.tsx
@@ -73,7 +73,10 @@ export function SessionDetailPage({
             allMsgs.map((m, i) => {
               const chatMsg: ChatMessage = {
                 ...m,
-                timestamp: Date.now() - (allMsgs.length - i) * 1000,
+                // Use server-provided created_at; fall back to synthetic spacing for older messages.
+                timestamp: m.created_at
+                  ? new Date(m.created_at).getTime()
+                  : Date.now() - (allMsgs.length - i) * 1000,
               };
               // Reconstruct toolDetails for assistant messages with tool_calls
               if (m.role === "assistant" && m.tool_calls && m.tool_calls.length > 0) {
diff --git a/ui/web/src/types/session.ts b/ui/web/src/types/session.ts
index 51958dd7..f043ab22 100644
--- a/ui/web/src/types/session.ts
+++ b/ui/web/src/types/session.ts
@@ -34,6 +34,7 @@ export interface Message {
   tool_call_id?: string;
   is_error?: boolean;
   media_refs?: { id: string; mime_type: string; kind: string; path?: string }[];
+  created_at?: string; // ISO 8601 timestamp from server; absent for older messages
 }
 
 export interface ToolCall {