Migrate followup draft SSE to WebSockets (#776)

frontend/src/hooks/follow-up/useDraftStream.ts

@@ -1,5 +1,5 @@
 import { useCallback, useEffect, useRef, useState } from 'react';
-import { useJsonPatchStream } from '@/hooks/useJsonPatchStream';
+import { useJsonPatchWsStream } from '@/hooks/useJsonPatchWsStream';
 import { attemptsApi } from '@/lib/api';
 import type { FollowUpDraft } from 'shared/types';
 import { inIframe } from '@/vscode/bridge';
@@ -14,7 +14,7 @@ export function useDraftStream(attemptId?: string) {
   const forceNextApplyRef = useRef<boolean>(false);
 
   const endpoint = attemptId
-    ? `/api/task-attempts/${attemptId}/follow-up-draft/stream`
+    ? `/api/task-attempts/${attemptId}/follow-up-draft/stream/ws`
     : undefined;
 
   const makeInitial = useCallback(
@@ -35,7 +35,7 @@ export function useDraftStream(attemptId?: string) {
     [attemptId]
   );
 
-  const { data, isConnected, error } = useJsonPatchStream<DraftStreamState>(
+  const { data, isConnected, error } = useJsonPatchWsStream<DraftStreamState>(
     endpoint,
     !!endpoint,
     makeInitial
@@ -64,7 +64,7 @@ export function useDraftStream(attemptId?: string) {
         });
         if (!isDraftLoaded) setIsDraftLoaded(true);
       } catch {
-        // ignore, rely on SSE
+        // ignore, rely on stream
       }
     };
     hydrate();
@@ -73,7 +73,7 @@ export function useDraftStream(attemptId?: string) {
     };
   }, [attemptId, isDraftLoaded]);
 
-  // Handle SSE stream
+  // Handle stream updates
   useEffect(() => {
     if (!data) return;
     const d = data.follow_up_draft;

frontend/src/hooks/useJsonPatchWsStream.ts

@@ -2,6 +2,10 @@ import { useEffect, useState, useRef } from 'react';
 import { applyPatch } from 'rfc6902';
 import type { Operation } from 'rfc6902';
 
+type WsJsonPatchMsg = { JsonPatch: Operation[] };
+type WsFinishedMsg = { finished: boolean };
+type WsMsg = WsJsonPatchMsg | WsFinishedMsg;
+
 interface UseJsonPatchStreamOptions<T> {
   /**
    * Called once when the stream starts to inject initial data
@@ -98,10 +102,10 @@ export const useJsonPatchWsStream = <T>(
 
       ws.onmessage = (event) => {
         try {
-          const msg = JSON.parse(event.data);
+          const msg: WsMsg = JSON.parse(event.data);
 
           // Handle JsonPatch messages (same as SSE json_patch event)
-          if (msg.JsonPatch) {
+          if ('JsonPatch' in msg) {
             const patches: Operation[] = msg.JsonPatch;
             const filtered = options.deduplicatePatches
               ? options.deduplicatePatches(patches)
@@ -119,8 +123,8 @@ export const useJsonPatchWsStream = <T>(
             setData(dataRef.current);
           }
 
-          // Handle Finished messages (same as SSE finished event)
-          if (msg.Finished !== undefined) {
+          // Handle finished messages ({finished: true})
+          if ('finished' in msg) {
             ws.close();
             wsRef.current = null;
             setIsConnected(false);