md

www/app/(app)/transcripts/useTranscriptChat.ts

@@ -1,7 +1,9 @@
 "use client";
 
 import { useEffect, useState, useRef } from "react";
+import { getSession } from "next-auth/react";
 import { WEBSOCKET_URL } from "../../lib/apiClient";
+import { assertExtendedToken } from "../../lib/types";
 
 export type Message = {
   id: string;
@@ -27,59 +29,84 @@ export const useTranscriptChat = (transcriptId: string): UseTranscriptChat => {
 
   useEffect(() => {
     isMountedRef.current = true;
-    const url = `${WEBSOCKET_URL}/v1/transcripts/${transcriptId}/chat`;
-    const ws = new WebSocket(url);
-    wsRef.current = ws;
 
-    ws.onopen = () => {
-      console.log("Chat WebSocket connected");
-    };
+    const connectWebSocket = async () => {
+      const url = `${WEBSOCKET_URL}/v1/transcripts/${transcriptId}/chat`;
 
-    ws.onmessage = (event) => {
-      if (!isMountedRef.current) return;
-
-      const msg = JSON.parse(event.data);
-
-      switch (msg.type) {
-        case "token":
-          setIsStreaming(true);
-          streamingTextRef.current += msg.text;
-          setCurrentStreamingText(streamingTextRef.current);
-          break;
-
-        case "done":
-          setMessages((prev) => [
-            ...prev,
-            {
-              id: Date.now().toString(),
-              role: "assistant",
-              text: streamingTextRef.current,
-              timestamp: new Date(),
-            },
-          ]);
-          streamingTextRef.current = "";
-          setCurrentStreamingText("");
-          setIsStreaming(false);
-          break;
-
-        case "error":
-          console.error("Chat error:", msg.message);
-          setIsStreaming(false);
-          break;
+      // Get auth token for WebSocket subprotocol
+      let protocols: string[] | undefined;
+      try {
+        const session = await getSession();
+        if (session) {
+          const token = assertExtendedToken(session).accessToken;
+          // Pass token via subprotocol: ["bearer", token]
+          protocols = ["bearer", token];
+        }
+      } catch (error) {
+        console.warn("Failed to get auth token for WebSocket:", error);
       }
+
+      const ws = new WebSocket(url, protocols);
+      wsRef.current = ws;
+
+      ws.onopen = () => {
+        console.log("Chat WebSocket connected");
+      };
+
+      ws.onmessage = (event) => {
+        if (!isMountedRef.current) return;
+
+        const msg = JSON.parse(event.data);
+
+        switch (msg.type) {
+          case "token":
+            setIsStreaming(true);
+            streamingTextRef.current += msg.text;
+            setCurrentStreamingText(streamingTextRef.current);
+            break;
+
+          case "done":
+            // CRITICAL: Save the text BEFORE resetting the ref
+            // The setMessages callback may execute later, after ref is reset
+            const finalText = streamingTextRef.current;
+
+            setMessages((prev) => [
+              ...prev,
+              {
+                id: Date.now().toString(),
+                role: "assistant",
+                text: finalText,
+                timestamp: new Date(),
+              },
+            ]);
+            streamingTextRef.current = "";
+            setCurrentStreamingText("");
+            setIsStreaming(false);
+            break;
+
+          case "error":
+            console.error("Chat error:", msg.message);
+            setIsStreaming(false);
+            break;
+        }
+      };
+
+      ws.onerror = (error) => {
+        console.error("WebSocket error:", error);
+      };
+
+      ws.onclose = () => {
+        console.log("Chat WebSocket closed");
+      };
     };
 
-    ws.onerror = (error) => {
-      console.error("WebSocket error:", error);
-    };
-
-    ws.onclose = () => {
-      console.log("Chat WebSocket closed");
-    };
+    connectWebSocket();
 
     return () => {
       isMountedRef.current = false;
-      ws.close();
+      if (wsRef.current) {
+        wsRef.current.close();
+      }
     };
   }, [transcriptId]);