1 files changed, 39 insertions, 19 deletions
diff --git a/makima/frontend/src/hooks/useWebSocket.ts b/makima/frontend/src/hooks/useWebSocket.ts
index de6c1a6..961951f 100644
--- a/makima/frontend/src/hooks/useWebSocket.ts
+++ b/makima/frontend/src/hooks/useWebSocket.ts
@@ -38,6 +38,8 @@ export function useWebSocket(options: UseWebSocketOptions = {}) {
 
   const wsRef = useRef<WebSocket | null>(null);
   const transcriptIdRef = useRef(0);
+  const stoppingRef = useRef(false);
+  const pendingDisconnectRef = useRef(false);
 
   // Store callbacks in refs to avoid recreating handlers
   const callbacksRef = useRef({ onReady, onTranscript, onError, onStopped });
@@ -90,27 +92,28 @@ export function useWebSocket(options: UseWebSocketOptions = {}) {
                 };
 
                 setState((s) => {
-                  if (message.isFinal) {
-                    // Final transcript replaces all previous transcripts from this speaker
-                    const filtered = s.transcripts.filter(
-                      (t) => t.speaker !== message.speaker
-                    );
-                    return { ...s, transcripts: [...filtered, entry] };
+                  // Find existing transcript with same speaker and overlapping timestamp
+                  const existingIdx = s.transcripts.findIndex(
+                    (t) =>
+                      t.speaker === message.speaker &&
+                      Math.abs(t.start - message.start) < 0.1
+                  );
+
+                  let newTranscripts: TranscriptEntry[];
+
+                  if (existingIdx >= 0) {
+                    // Replace existing transcript (final replaces non-final, or update in place)
+                    newTranscripts = [...s.transcripts];
+                    newTranscripts[existingIdx] = entry;
                   } else {
-                    // Non-final: replace if same speaker and overlapping time, otherwise append
-                    const existingIdx = s.transcripts.findIndex(
-                      (t) =>
-                        !t.isFinal &&
-                        t.speaker === message.speaker &&
-                        Math.abs(t.start - message.start) < 0.1
-                    );
-                    if (existingIdx >= 0) {
-                      const newTranscripts = [...s.transcripts];
-                      newTranscripts[existingIdx] = entry;
-                      return { ...s, transcripts: newTranscripts };
-                    }
-                    return { ...s, transcripts: [...s.transcripts, entry] };
+                    // No overlap - insert in time order
+                    newTranscripts = [...s.transcripts, entry];
                   }
+
+                  // Sort by start time to maintain chronological order
+                  newTranscripts.sort((a, b) => a.start - b.start);
+
+                  return { ...s, transcripts: newTranscripts };
                 });
 
                 callbacksRef.current.onTranscript?.(entry);
@@ -123,8 +126,17 @@ export function useWebSocket(options: UseWebSocketOptions = {}) {
                 break;
 
               case "stopped":
+                stoppingRef.current = false;
                 setState((s) => ({ ...s, status: "disconnected" }));
                 callbacksRef.current.onStopped?.(message.reason);
+                // Execute pending disconnect if requested during stopping
+                if (pendingDisconnectRef.current) {
+                  pendingDisconnectRef.current = false;
+                  if (wsRef.current) {
+                    wsRef.current.close(1000, "User disconnected");
+                    wsRef.current = null;
+                  }
+                }
                 break;
             }
           } catch {
@@ -171,6 +183,11 @@ export function useWebSocket(options: UseWebSocketOptions = {}) {
   }, []);
 
   const disconnect = useCallback(() => {
+    if (stoppingRef.current) {
+      // Defer disconnect until "stopped" message is received
+      pendingDisconnectRef.current = true;
+      return;
+    }
     if (wsRef.current) {
       wsRef.current.close(1000, "User disconnected");
       wsRef.current = null;
@@ -210,6 +227,7 @@ export function useWebSocket(options: UseWebSocketOptions = {}) {
 
   const stopSession = useCallback(
     (reason?: string) => {
+      stoppingRef.current = true;
       sendMessage({
         type: "stop",
         reason,
@@ -219,6 +237,8 @@ export function useWebSocket(options: UseWebSocketOptions = {}) {
   );
 
   const clearTranscripts = useCallback(() => {
+    stoppingRef.current = false;
+    pendingDisconnectRef.current = false;
     setState((s) => ({ ...s, transcripts: [], error: null }));
   }, []);