1 files changed, 244 insertions, 0 deletions
diff --git a/makima/frontend/src/hooks/useWebSocket.ts b/makima/frontend/src/hooks/useWebSocket.ts
new file mode 100644
index 0000000..de6c1a6
--- /dev/null
+++ b/makima/frontend/src/hooks/useWebSocket.ts
@@ -0,0 +1,244 @@
+import { useState, useCallback, useRef, useEffect } from "react";
+import { LISTEN_ENDPOINT } from "../lib/api";
+import type {
+  ClientMessage,
+  ServerMessage,
+  TranscriptEntry,
+} from "../types/messages";
+
+export type ConnectionStatus =
+  | "disconnected"
+  | "connecting"
+  | "connected"
+  | "error";
+
+export interface WebSocketState {
+  status: ConnectionStatus;
+  sessionId: string | null;
+  error: string | null;
+  transcripts: TranscriptEntry[];
+}
+
+interface UseWebSocketOptions {
+  onReady?: (sessionId: string) => void;
+  onTranscript?: (transcript: TranscriptEntry) => void;
+  onError?: (code: string, message: string) => void;
+  onStopped?: (reason: string) => void;
+}
+
+export function useWebSocket(options: UseWebSocketOptions = {}) {
+  const { onReady, onTranscript, onError, onStopped } = options;
+
+  const [state, setState] = useState<WebSocketState>({
+    status: "disconnected",
+    sessionId: null,
+    error: null,
+    transcripts: [],
+  });
+
+  const wsRef = useRef<WebSocket | null>(null);
+  const transcriptIdRef = useRef(0);
+
+  // Store callbacks in refs to avoid recreating handlers
+  const callbacksRef = useRef({ onReady, onTranscript, onError, onStopped });
+  useEffect(() => {
+    callbacksRef.current = { onReady, onTranscript, onError, onStopped };
+  }, [onReady, onTranscript, onError, onStopped]);
+
+  const connect = useCallback((): Promise<boolean> => {
+    return new Promise((resolve) => {
+      if (wsRef.current?.readyState === WebSocket.OPEN) {
+        resolve(true);
+        return;
+      }
+
+      // Close any existing connection
+      if (wsRef.current) {
+        wsRef.current.close();
+        wsRef.current = null;
+      }
+
+      setState((s) => ({ ...s, status: "connecting", error: null }));
+
+      try {
+        const ws = new WebSocket(LISTEN_ENDPOINT);
+        wsRef.current = ws;
+
+        ws.onopen = () => {
+          setState((s) => ({ ...s, status: "connected", error: null }));
+          resolve(true);
+        };
+
+        ws.onmessage = (event) => {
+          try {
+            const message: ServerMessage = JSON.parse(event.data);
+
+            switch (message.type) {
+              case "ready":
+                setState((s) => ({ ...s, sessionId: message.sessionId }));
+                callbacksRef.current.onReady?.(message.sessionId);
+                break;
+
+              case "transcript": {
+                const entry: TranscriptEntry = {
+                  id: `t-${++transcriptIdRef.current}`,
+                  speaker: message.speaker,
+                  start: message.start,
+                  end: message.end,
+                  text: message.text,
+                  isFinal: message.isFinal,
+                };
+
+                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] };
+                  } 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] };
+                  }
+                });
+
+                callbacksRef.current.onTranscript?.(entry);
+                break;
+              }
+
+              case "error":
+                setState((s) => ({ ...s, error: message.message }));
+                callbacksRef.current.onError?.(message.code, message.message);
+                break;
+
+              case "stopped":
+                setState((s) => ({ ...s, status: "disconnected" }));
+                callbacksRef.current.onStopped?.(message.reason);
+                break;
+            }
+          } catch {
+            console.error("Failed to parse WebSocket message:", event.data);
+          }
+        };
+
+        ws.onerror = () => {
+          setState((s) => ({
+            ...s,
+            status: "error",
+            error: "Failed to connect to server",
+          }));
+          resolve(false);
+        };
+
+        ws.onclose = (event) => {
+          // Check for specific close codes
+          let errorMessage: string | null = null;
+          if (event.code === 1006) {
+            errorMessage = "Connection failed - server may be unavailable";
+          } else if (event.code !== 1000 && event.code !== 1001) {
+            errorMessage = `Connection closed unexpectedly (code: ${event.code})`;
+          }
+
+          setState((s) => ({
+            ...s,
+            status: "disconnected",
+            sessionId: null,
+            error: errorMessage || s.error,
+          }));
+          wsRef.current = null;
+        };
+      } catch (err) {
+        const message = err instanceof Error ? err.message : "Failed to create WebSocket connection";
+        setState((s) => ({
+          ...s,
+          status: "error",
+          error: message,
+        }));
+        resolve(false);
+      }
+    });
+  }, []);
+
+  const disconnect = useCallback(() => {
+    if (wsRef.current) {
+      wsRef.current.close(1000, "User disconnected");
+      wsRef.current = null;
+    }
+    setState((s) => ({ ...s, status: "disconnected", sessionId: null }));
+  }, []);
+
+  const sendMessage = useCallback((message: ClientMessage) => {
+    if (wsRef.current?.readyState === WebSocket.OPEN) {
+      wsRef.current.send(JSON.stringify(message));
+    }
+  }, []);
+
+  const sendAudio = useCallback((samples: Float32Array) => {
+    if (wsRef.current?.readyState === WebSocket.OPEN) {
+      // Convert Float32Array to bytes (little-endian)
+      const bytes = new Uint8Array(samples.length * 4);
+      const view = new DataView(bytes.buffer);
+      for (let i = 0; i < samples.length; i++) {
+        view.setFloat32(i * 4, samples[i], true);
+      }
+      wsRef.current.send(bytes);
+    }
+  }, []);
+
+  const startSession = useCallback(
+    (sampleRate: number, channels: number = 1) => {
+      sendMessage({
+        type: "start",
+        sampleRate,
+        channels,
+        encoding: "pcm32f",
+      });
+    },
+    [sendMessage]
+  );
+
+  const stopSession = useCallback(
+    (reason?: string) => {
+      sendMessage({
+        type: "stop",
+        reason,
+      });
+    },
+    [sendMessage]
+  );
+
+  const clearTranscripts = useCallback(() => {
+    setState((s) => ({ ...s, transcripts: [], error: null }));
+  }, []);
+
+  // Cleanup on unmount
+  useEffect(() => {
+    return () => {
+      if (wsRef.current) {
+        wsRef.current.close();
+      }
+    };
+  }, []);
+
+  return {
+    ...state,
+    connect,
+    disconnect,
+    sendAudio,
+    startSession,
+    stopSession,
+    clearTranscripts,
+    isConnected: state.status === "connected",
+  };
+}