From f1b6fdfea4fb26976532d83f606aeb8c6942ee9e Mon Sep 17 00:00:00 2001
From: Adrian Gavrila <agavrila@microsoft.com>
Date: Wed, 18 Mar 2026 11:45:25 -0400
Subject: [PATCH] Fix conversation switching during in-flight requests and sort
 ordering

- Replace broad sendingConvIdsRef guard in ChatWindow useEffect with a
  forceLoadRef mechanism. User-initiated panel clicks set forceLoadRef
  to bypass the guard, allowing switches to sending conversations.
  Internal activeConversationId changes from handleSend (e.g., after
  attack creation) leave forceLoadRef unset so the guard protects
  optimistic messages from being overwritten by loadConversation.

- Normalize naive timestamps from SQLite to UTC before sorting conversations.
  Raw SQL queries return timezone-naive datetimes while the rest of the
  codebase uses UTC-aware datetimes, causing TypeError on comparison.

- In-flight conversations (no stored messages, created_at=None) now sort
  to the end using datetime.now(timezone.utc) as fallback, keeping them
  at the bottom of the list consistently.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 frontend/src/components/Chat/ChatWindow.tsx | 14 +++++++++++---
 pyrit/backend/services/attack_service.py    | 13 +++++++++----
 2 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/frontend/src/components/Chat/ChatWindow.tsx b/frontend/src/components/Chat/ChatWindow.tsx
index 92249fa02..b4a216c87 100644
--- a/frontend/src/components/Chat/ChatWindow.tsx
+++ b/frontend/src/components/Chat/ChatWindow.tsx
@@ -74,6 +74,10 @@ export default function ChatWindow({
       setIsPanelOpen(true)
     }
   }, [attackResultId, relatedConversationCount])
+  // Set by panel click to bypass the in-flight guard on the next useEffect cycle.
+  // This lets users switch to a sending conversation while still protecting
+  // optimistic messages when handleSend internally updates activeConversationId.
+  const forceLoadRef = useRef(false)
   // Always-current ref of the conversation being viewed so async callbacks can
   // check whether the user navigated away while a request was in-flight.
   const viewedConvRef = useRef(activeConversationId ?? conversationId)
@@ -126,9 +130,12 @@ export default function ChatWindow({
   // Reload messages when activeConversationId changes
   useEffect(() => {
     if (!attackResultId || !activeConversationId) { return }
-    // Skip loading if a send is already in-flight for this conversation —
-    // the send handler will update messages when it completes.
-    if (sendingConvIdsRef.current.has(activeConversationId)) { return }
+    // Allow user-initiated switches (forceLoadRef), but skip re-loading when
+    // handleSend internally updated activeConversationId during an in-flight
+    // send — the optimistic messages are already displayed.
+    const force = forceLoadRef.current
+    forceLoadRef.current = false
+    if (!force && sendingConvIdsRef.current.has(activeConversationId)) { return }
     loadConversation(attackResultId, activeConversationId)
   }, [activeConversationId, attackResultId, loadConversation])
 
@@ -143,6 +150,7 @@ export default function ChatWindow({
   // Handle conversation selection from the panel
   // For a different ID the useEffect handles loading; for same ID force a refresh
   const handlePanelSelectConversation = useCallback((convId: string) => {
+    forceLoadRef.current = true
     onSelectConversation(convId)
     if (convId === activeConversationId && attackResultId) {
       loadConversation(attackResultId, convId)
diff --git a/pyrit/backend/services/attack_service.py b/pyrit/backend/services/attack_service.py
index 5bb4f6bbc..16a8fb704 100644
--- a/pyrit/backend/services/attack_service.py
+++ b/pyrit/backend/services/attack_service.py
@@ -390,6 +390,9 @@ async def get_conversations_async(self, *, attack_result_id: str) -> Optional[At
         for conv_id in active_conv_ids:
             stats = stats_map.get(conv_id)
             created_at = stats.created_at if stats else None
+            # SQLite returns naive datetimes — normalize to UTC (same pattern as _ensure_utc)
+            if created_at is not None and created_at.tzinfo is None:
+                created_at = created_at.replace(tzinfo=timezone.utc)
             conversations.append(
                 ConversationSummary(
                     conversation_id=conv_id,
@@ -399,10 +402,12 @@ async def get_conversations_async(self, *, attack_result_id: str) -> Optional[At
                 )
             )
 
-        # Sort all conversations by created_at (earliest first, None last)
-        conversations.sort(
-            key=lambda c: (c.created_at is None, c.created_at or datetime.min.replace(tzinfo=timezone.utc))
-        )
+        # Sort conversations by created_at (earliest first). In-flight conversations
+        # have no stored messages yet so created_at is None — treat them as the most
+        # recent (they were just created) so they sort after older conversations
+        # instead of jumping to an arbitrary position.
+        now = datetime.now(timezone.utc)
+        conversations.sort(key=lambda c: c.created_at or now)
 
         return AttackConversationsResponse(
             attack_result_id=attack_result_id,