Consolidate Retry and Follow-up (#800)

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

@@ -1,113 +1,263 @@
+import {
+  attemptsApi,
+  type UpdateFollowUpDraftRequest,
+  type UpdateRetryFollowUpDraftRequest,
+} from '@/lib/api';
+import type { Draft, DraftResponse } from 'shared/types';
 import { useEffect, useRef, useState } from 'react';
-import { attemptsApi, type UpdateFollowUpDraftRequest } from '@/lib/api';
-import type { FollowUpDraft } from 'shared/types';
 
 export type SaveStatus = 'idle' | 'saving' | 'saved' | 'offline' | 'sent';
 
-type DraftData = Pick<FollowUpDraft, 'prompt' | 'variant' | 'image_ids'>;
-
-type Args = {
-  attemptId?: string;
-  serverDraft: FollowUpDraft | null;
-  current: DraftData;
-  isQueuedUI: boolean;
-  isDraftSending: boolean;
-  isQueuing: boolean;
-  isUnqueuing: boolean;
-  suppressNextSaveRef: React.MutableRefObject<boolean>;
-  lastServerVersionRef: React.MutableRefObject<number>;
-  forceNextApplyRef: React.MutableRefObject<boolean>;
+// Small helper to diff common draft fields
+type BaseCurrent = {
+  prompt: string;
+  variant: string | null | undefined;
+  image_ids: string[] | null | undefined;
+};
+type BaseServer = {
+  prompt?: string | null;
+  variant?: string | null;
+  image_ids?: string[] | null;
+} | null;
+type BasePayload = {
+  prompt?: string;
+  variant?: string | null;
+  image_ids?: string[];
 };
 
-export function useDraftAutosave({
+function diffBaseDraft(current: BaseCurrent, server: BaseServer): BasePayload {
+  const payload: BasePayload = {};
+  const serverPrompt = (server?.prompt ?? '') || '';
+  const serverVariant = server?.variant ?? null;
+  const serverIds = (server?.image_ids as string[] | undefined) ?? [];
+
+  if (current.prompt !== serverPrompt) payload.prompt = current.prompt || '';
+  if ((current.variant ?? null) !== serverVariant)
+    payload.variant = (current.variant ?? null) as string | null;
+
+  const currIds = (current.image_ids as string[] | null) ?? [];
+  const idsEqual =
+    currIds.length === serverIds.length &&
+    currIds.every((id, i) => id === serverIds[i]);
+  if (!idsEqual) payload.image_ids = currIds;
+
+  return payload;
+}
+
+function diffDraftPayload<
+  TExtra extends Record<string, unknown> = Record<string, never>,
+>(
+  current: BaseCurrent,
+  server: BaseServer,
+  extra?: TExtra,
+  requireBaseChange: boolean = true
+): (BasePayload & TExtra) | null {
+  const base = diffBaseDraft(current, server);
+  const baseChanged = Object.keys(base).length > 0;
+  if (!baseChanged && requireBaseChange) return null;
+  return { ...(extra as object), ...base } as BasePayload & TExtra;
+}
+
+// Private core
+function useDraftAutosaveCore<TServer, TCurrent, TPayload>({
   attemptId,
   serverDraft,
   current,
-  isQueuedUI,
   isDraftSending,
-  isQueuing,
-  isUnqueuing,
-  suppressNextSaveRef,
-  lastServerVersionRef,
-  forceNextApplyRef,
-}: Args) {
+  skipConditions = [],
+  buildPayload,
+  saveDraft,
+  fetchLatest,
+  debugLabel,
+}: {
+  attemptId?: string;
+  serverDraft: TServer | null;
+  current: TCurrent;
+  isDraftSending: boolean;
+  skipConditions?: boolean[];
+  buildPayload: (
+    current: TCurrent,
+    serverDraft: TServer | null
+  ) => TPayload | null;
+  saveDraft: (attemptId: string, payload: TPayload) => Promise<unknown>;
+  fetchLatest?: (attemptId: string) => Promise<unknown>;
+  debugLabel?: string;
+}) {
   const [isSaving, setIsSaving] = useState(false);
   const [saveStatus, setSaveStatus] = useState<SaveStatus>('idle');
-  // Presentation timers moved to FollowUpStatusRow; keep only raw status.
-
-  // debounced save
   const lastSentRef = useRef<string>('');
   const saveTimeoutRef = useRef<number | undefined>(undefined);
+
   useEffect(() => {
     if (!attemptId) return;
-    if (isDraftSending) return;
-    if (isQueuing || isUnqueuing) return;
-    if (suppressNextSaveRef.current) {
-      suppressNextSaveRef.current = false;
+    if (isDraftSending) {
+      if (import.meta.env.DEV)
+        console.debug(`[autosave:${debugLabel}] skip: draft is sending`, {
+          attemptId,
+        });
+      return;
+    }
+    if (skipConditions.some((c) => c)) {
+      if (import.meta.env.DEV)
+        console.debug(`[autosave:${debugLabel}] skip: skipConditions`, {
+          attemptId,
+          skipConditions,
+        });
       return;
     }
-    if (isQueuedUI) return;
 
-    const saveDraft = async () => {
-      const payload: Partial<UpdateFollowUpDraftRequest> = {};
-      if (serverDraft && current.prompt !== (serverDraft.prompt || ''))
-        payload.prompt = current.prompt || '';
-      if ((serverDraft?.variant ?? null) !== (current.variant ?? null))
-        payload.variant = (current.variant ?? null) as string | null;
-      const currentIds = (current.image_ids as string[] | null) ?? [];
-      const serverIds = (serverDraft?.image_ids as string[] | undefined) ?? [];
-      const idsEqual =
-        currentIds.length === serverIds.length &&
-        currentIds.every((id, i) => id === serverIds[i]);
-      if (!idsEqual) payload.image_ids = currentIds;
-      const keys = Object.keys(payload);
-      if (keys.length === 0) return;
+    const doSave = async () => {
+      const payload = buildPayload(current, serverDraft);
+      if (!payload) {
+        if (import.meta.env.DEV)
+          console.debug(`[autosave:${debugLabel}] no changes`, { attemptId });
+        return;
+      }
       const payloadKey = JSON.stringify(payload);
-      if (payloadKey === lastSentRef.current) return;
+      if (payloadKey === lastSentRef.current) {
+        if (import.meta.env.DEV)
+          console.debug(`[autosave:${debugLabel}] deduped identical payload`, {
+            attemptId,
+            payload,
+          });
+        return;
+      }
       lastSentRef.current = payloadKey;
+
       try {
         setIsSaving(true);
         setSaveStatus(navigator.onLine ? 'saving' : 'offline');
-        await attemptsApi.saveFollowUpDraft(
-          attemptId,
-          payload as UpdateFollowUpDraftRequest
-        );
+        if (import.meta.env.DEV)
+          console.debug(`[autosave:${debugLabel}] saving`, {
+            attemptId,
+            payload,
+          });
+        await saveDraft(attemptId, payload);
         setSaveStatus('saved');
+        if (import.meta.env.DEV)
+          console.debug(`[autosave:${debugLabel}] saved`, { attemptId });
       } catch {
-        try {
-          // Fetch latest server draft to ensure stream catches up,
-          // and force next apply to override local edits when it arrives.
-          await attemptsApi.getFollowUpDraft(attemptId);
-          suppressNextSaveRef.current = true;
-          forceNextApplyRef.current = true;
-        } catch {
-          /* ignore */
+        if (import.meta.env.DEV)
+          console.debug(`[autosave:${debugLabel}] error -> fetchLatest`, {
+            attemptId,
+          });
+        if (fetchLatest) {
+          try {
+            await fetchLatest(attemptId);
+          } catch {
+            /* empty */
+          }
         }
         setSaveStatus(navigator.onLine ? 'idle' : 'offline');
       } finally {
         setIsSaving(false);
       }
     };
+
     if (saveTimeoutRef.current) window.clearTimeout(saveTimeoutRef.current);
-    saveTimeoutRef.current = window.setTimeout(saveDraft, 400);
+    saveTimeoutRef.current = window.setTimeout(doSave, 400);
     return () => {
       if (saveTimeoutRef.current) window.clearTimeout(saveTimeoutRef.current);
     };
   }, [
     attemptId,
-    serverDraft?.prompt,
-    serverDraft?.variant,
-    serverDraft?.image_ids,
-    current.prompt,
-    current.variant,
-    current.image_ids,
-    isQueuedUI,
+    serverDraft,
+    current,
     isDraftSending,
-    isQueuing,
-    isUnqueuing,
-    suppressNextSaveRef,
-    lastServerVersionRef,
+    skipConditions,
+    buildPayload,
+    saveDraft,
+    fetchLatest,
+    debugLabel,
   ]);
 
   return { isSaving, saveStatus } as const;
 }
+
+type DraftData = Pick<Draft, 'prompt' | 'variant' | 'image_ids'>;
+
+type DraftArgs<TServer, TCurrent> = {
+  attemptId?: string;
+  serverDraft: TServer | null;
+  current: TCurrent;
+  isDraftSending: boolean;
+  // Queue-related flags (used for follow_up; not used for retry)
+  isQueuedUI?: boolean;
+  isQueuing?: boolean;
+  isUnqueuing?: boolean;
+  // Discriminant
+  draftType?: 'follow_up' | 'retry';
+};
+
+type FollowUpAutosaveArgs = DraftArgs<Draft, DraftData> & {
+  draftType?: 'follow_up';
+};
+type RetryAutosaveArgs = DraftArgs<RetryDraftResponse, RetryDraftData> & {
+  draftType: 'retry';
+};
+
+export function useDraftAutosave(
+  args: FollowUpAutosaveArgs | RetryAutosaveArgs
+) {
+  const skipConditions =
+    args.draftType === 'retry'
+      ? [!args.serverDraft]
+      : [!!args.isQueuing, !!args.isUnqueuing, !!args.isQueuedUI];
+
+  return useDraftAutosaveCore<
+    Draft | RetryDraftResponse,
+    DraftData | RetryDraftData,
+    UpdateFollowUpDraftRequest | UpdateRetryFollowUpDraftRequest
+  >({
+    attemptId: args.attemptId,
+    serverDraft: args.serverDraft as Draft | RetryDraftResponse | null,
+    current: args.current as DraftData | RetryDraftData,
+    isDraftSending: args.isDraftSending,
+    skipConditions,
+    debugLabel: (args.draftType ?? 'follow_up') as string,
+    buildPayload: (current, serverDraft) => {
+      if (args.draftType === 'retry') {
+        const c = current as RetryDraftData;
+        const s = serverDraft as RetryDraftResponse | null;
+        return diffDraftPayload(
+          c,
+          s,
+          { retry_process_id: c.retry_process_id },
+          true
+        ) as UpdateRetryFollowUpDraftRequest | null;
+      }
+      const c = current as DraftData;
+      const s = serverDraft as Draft | null;
+      return diffDraftPayload(c, s) as UpdateFollowUpDraftRequest | null;
+    },
+    saveDraft: (id, payload) => {
+      if (args.draftType === 'retry') {
+        return attemptsApi.saveDraft(
+          id,
+          'retry',
+          payload as UpdateRetryFollowUpDraftRequest
+        );
+      }
+      return attemptsApi.saveDraft(
+        id,
+        'follow_up',
+        payload as UpdateFollowUpDraftRequest
+      );
+    },
+    fetchLatest: (id) => {
+      if (args.draftType === 'retry') return attemptsApi.getDraft(id, 'retry');
+      return attemptsApi.getDraft(id, 'follow_up');
+    },
+  });
+}
+
+export type RetrySaveStatus = SaveStatus;
+
+type RetryDraftResponse = DraftResponse;
+
+type RetryDraftData = Pick<
+  DraftResponse,
+  'prompt' | 'variant' | 'image_ids'
+> & {
+  retry_process_id: string;
+};

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

@@ -0,0 +1,88 @@
+import { useEffect, useRef, useState, useCallback } from 'react';
+import type { Draft, ImageResponse } from 'shared/types';
+import { imagesApi } from '@/lib/api';
+import { useQuery } from '@tanstack/react-query';
+
+type PartialDraft = Pick<Draft, 'prompt' | 'image_ids'>;
+
+type Args = {
+  draft: PartialDraft | null;
+  taskId: string;
+};
+
+export function useDraftEditor({ draft, taskId }: Args) {
+  const [message, setMessageInner] = useState('');
+  const [localImages, setLocalImages] = useState<ImageResponse[]>([]);
+  const [newlyUploadedImageIds, setNewlyUploadedImageIds] = useState<string[]>(
+    []
+  );
+
+  const localDirtyRef = useRef<boolean>(false);
+  const imagesDirtyRef = useRef<boolean>(false);
+
+  // Sync message with server when not locally dirty
+  useEffect(() => {
+    if (!draft) return;
+    const serverPrompt = draft.prompt || '';
+    if (!localDirtyRef.current) {
+      setMessageInner(serverPrompt);
+    } else if (serverPrompt === message) {
+      // When server catches up to local text, clear dirty
+      localDirtyRef.current = false;
+    }
+  }, [draft, message]);
+
+  // Fetch images for task via react-query and map to the draft's image_ids
+  const serverIds = (draft?.image_ids ?? []).filter(Boolean);
+  const idsKey = serverIds.join(',');
+  const imagesQuery = useQuery({
+    queryKey: ['taskImagesForDraft', taskId, idsKey],
+    enabled: !!taskId,
+    queryFn: async () => {
+      const all = await imagesApi.getTaskImages(taskId);
+      const want = new Set(serverIds);
+      return all.filter((img) => want.has(img.id));
+    },
+    staleTime: 60_000,
+  });
+
+  const images = imagesDirtyRef.current
+    ? localImages
+    : (imagesQuery.data ?? []);
+
+  const setMessage = (v: React.SetStateAction<string>) => {
+    localDirtyRef.current = true;
+    if (typeof v === 'function') {
+      setMessageInner((prev) => v(prev));
+    } else {
+      setMessageInner(v);
+    }
+  };
+
+  const setImages = (next: ImageResponse[]) => {
+    imagesDirtyRef.current = true;
+    setLocalImages(next);
+  };
+
+  const handleImageUploaded = useCallback((image: ImageResponse) => {
+    imagesDirtyRef.current = true;
+    setLocalImages((prev) => [...prev, image]);
+    setNewlyUploadedImageIds((prev) => [...prev, image.id]);
+  }, []);
+
+  const clearImagesAndUploads = useCallback(() => {
+    imagesDirtyRef.current = false;
+    setLocalImages([]);
+    setNewlyUploadedImageIds([]);
+  }, []);
+
+  return {
+    message,
+    setMessage,
+    images,
+    setImages,
+    newlyUploadedImageIds,
+    handleImageUploaded,
+    clearImagesAndUploads,
+  } as const;
+}

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

@@ -1,104 +0,0 @@
-import { useEffect, useRef, useState, useCallback } from 'react';
-import type { FollowUpDraft, ImageResponse } from 'shared/types';
-import { imagesApi } from '@/lib/api';
-
-type Args = {
-  draft: FollowUpDraft | null;
-  lastServerVersionRef: React.MutableRefObject<number>;
-  suppressNextSaveRef: React.MutableRefObject<boolean>;
-  forceNextApplyRef: React.MutableRefObject<boolean>;
-  taskId: string;
-};
-
-export function useDraftEdits({
-  draft,
-  lastServerVersionRef,
-  suppressNextSaveRef,
-  forceNextApplyRef,
-  taskId,
-}: Args) {
-  const [message, setMessageInner] = useState('');
-  const [images, setImages] = useState<ImageResponse[]>([]);
-  const [newlyUploadedImageIds, setNewlyUploadedImageIds] = useState<string[]>(
-    []
-  );
-
-  const localDirtyRef = useRef<boolean>(false);
-  const imagesDirtyRef = useRef<boolean>(false);
-
-  useEffect(() => {
-    if (!draft) return;
-    const incomingVersion = Number(draft.version ?? 0n);
-
-    if (incomingVersion === lastServerVersionRef.current) return;
-    suppressNextSaveRef.current = true;
-    const isInitial = lastServerVersionRef.current === -1;
-    const shouldForce = forceNextApplyRef.current;
-    const allowApply = isInitial || shouldForce || !localDirtyRef.current;
-    if (allowApply && incomingVersion >= lastServerVersionRef.current) {
-      setMessageInner(draft.prompt || '');
-      localDirtyRef.current = false;
-      lastServerVersionRef.current = incomingVersion;
-      if (shouldForce) forceNextApplyRef.current = false;
-    } else if (incomingVersion > lastServerVersionRef.current) {
-      // Skip applying server changes while user is editing; still advance version to avoid loops
-      lastServerVersionRef.current = incomingVersion;
-    }
-  }, [draft]);
-
-  // Sync images from server when not locally dirty
-  useEffect(() => {
-    if (!draft) return;
-    const serverIds = (draft.image_ids || []) as string[];
-    const wantIds = new Set(serverIds);
-    const haveIds = new Set(images.map((img) => img.id));
-    const equal =
-      haveIds.size === wantIds.size &&
-      Array.from(haveIds).every((id) => wantIds.has(id));
-
-    if (equal) {
-      imagesDirtyRef.current = false;
-      return;
-    }
-
-    if (imagesDirtyRef.current) return;
-
-    imagesApi
-      .getTaskImages(taskId)
-      .then((all) => {
-        const next = all.filter((img) => wantIds.has(img.id));
-        setImages(next);
-        setNewlyUploadedImageIds([]);
-      })
-      .catch(() => void 0);
-  }, [draft?.image_ids, taskId, images]);
-
-  const handleImageUploaded = useCallback((image: ImageResponse) => {
-    imagesDirtyRef.current = true;
-    setImages((prev) => [...prev, image]);
-    setNewlyUploadedImageIds((prev) => [...prev, image.id]);
-  }, []);
-
-  const clearImagesAndUploads = useCallback(() => {
-    imagesDirtyRef.current = false;
-    setImages([]);
-    setNewlyUploadedImageIds([]);
-  }, []);
-
-  return {
-    message,
-    setMessage: (v: React.SetStateAction<string>) => {
-      localDirtyRef.current = true;
-      if (typeof v === 'function') {
-        setMessageInner((prev) => (v as (prev: string) => string)(prev));
-      } else {
-        setMessageInner(v);
-      }
-    },
-    images,
-    setImages,
-    newlyUploadedImageIds,
-    handleImageUploaded,
-    clearImagesAndUploads,
-  } as const;
-}

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

@@ -1,15 +1,13 @@
 import { useCallback } from 'react';
 import { attemptsApi, type UpdateFollowUpDraftRequest } from '@/lib/api';
-import type { FollowUpDraft, ImageResponse } from 'shared/types';
+import type { Draft, ImageResponse } from 'shared/types';
 
 type Args = {
   attemptId?: string;
-  draft: FollowUpDraft | null;
+  draft: Draft | null;
   message: string;
   selectedVariant: string | null;
   images: ImageResponse[];
-  suppressNextSaveRef: React.MutableRefObject<boolean>;
-  lastServerVersionRef: React.MutableRefObject<number>;
 };
 
 export function useDraftQueue({
@@ -18,8 +16,6 @@ export function useDraftQueue({
   message,
   selectedVariant,
   images,
-  suppressNextSaveRef,
-  lastServerVersionRef,
 }: Args) {
   const onQueue = useCallback(async (): Promise<boolean> => {
     if (!attemptId) return false;
@@ -37,26 +33,13 @@ export function useDraftQueue({
         currentIds.length === serverIds.length &&
         currentIds.every((id, i) => id === serverIds[i]);
       if (!idsEqual) immediatePayload.image_ids = currentIds;
-      suppressNextSaveRef.current = true;
-      await attemptsApi.saveFollowUpDraft(
+      await attemptsApi.saveDraft(
         attemptId,
+        'follow_up',
         immediatePayload as UpdateFollowUpDraftRequest
       );
-      try {
-        const resp = await attemptsApi.setFollowUpQueue(attemptId, true);
-        if (resp?.version !== undefined && resp?.version !== null) {
-          lastServerVersionRef.current = Number(resp.version ?? 0n);
-        }
-        return !!resp?.queued;
-      } catch {
-        /* adopt server on failure */
-        const latest = await attemptsApi.getFollowUpDraft(attemptId);
-        suppressNextSaveRef.current = true;
-        if (latest.version !== undefined && latest.version !== null) {
-          lastServerVersionRef.current = Number(latest.version ?? 0n);
-        }
-        return !!latest?.queued;
-      }
+      const resp = await attemptsApi.setDraftQueue(attemptId, true);
+      return !!resp?.queued;
     } finally {
       // presentation-only state handled by caller
     }
@@ -65,36 +48,22 @@ export function useDraftQueue({
     attemptId,
     draft?.variant,
     draft?.image_ids,
+    draft?.queued,
     images,
     message,
     selectedVariant,
-    suppressNextSaveRef,
-    lastServerVersionRef,
   ]);
 
   const onUnqueue = useCallback(async (): Promise<boolean> => {
     if (!attemptId) return false;
     try {
-      suppressNextSaveRef.current = true;
-      try {
-        const resp = await attemptsApi.setFollowUpQueue(attemptId, false);
-        if (resp?.version !== undefined && resp?.version !== null) {
-          lastServerVersionRef.current = Number(resp.version ?? 0n);
-        }
-        return !!resp && !resp.queued;
-      } catch {
-        const latest = await attemptsApi.getFollowUpDraft(attemptId);
-        suppressNextSaveRef.current = true;
-        if (latest.version !== undefined && latest.version !== null) {
-          lastServerVersionRef.current = Number(latest.version ?? 0n);
-        }
-        return !!latest && !latest.queued;
-      }
+      const resp = await attemptsApi.setDraftQueue(attemptId, false);
+      return !!resp && !resp.queued;
     } finally {
       // presentation-only state handled by caller
     }
     return false;
-  }, [attemptId, suppressNextSaveRef, lastServerVersionRef]);
+  }, [attemptId]);
 
   return { onQueue, onUnqueue } as const;
 }

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

@@ -1,143 +1,121 @@
-import { useCallback, useEffect, useRef, useState } from 'react';
-import { useJsonPatchWsStream } from '@/hooks/useJsonPatchWsStream';
-import { attemptsApi } from '@/lib/api';
-import type { FollowUpDraft } from 'shared/types';
-import { inIframe } from '@/vscode/bridge';
+import { useCallback, useEffect, useMemo } from 'react';
+import { useQuery, useQueryClient } from '@tanstack/react-query';
+import { applyPatch } from 'rfc6902';
+import type { Operation } from 'rfc6902';
+import useWebSocket from 'react-use-websocket';
+import type { Draft, DraftResponse } from 'shared/types';
+import { useProject } from '@/contexts/project-context';
 
-type DraftStreamState = { follow_up_draft: FollowUpDraft };
+interface Drafts {
+  [attemptId: string]: { follow_up: Draft; retry: DraftResponse | null };
+}
+
+type DraftsContainer = {
+  drafts: Drafts;
+};
+
+type WsJsonPatchMsg = { JsonPatch: Operation[] };
+type WsFinishedMsg = { finished: boolean };
+type WsMsg = WsJsonPatchMsg | WsFinishedMsg;
 
 export function useDraftStream(attemptId?: string) {
-  const [draft, setDraft] = useState<FollowUpDraft | null>(null);
-  const [isDraftLoaded, setIsDraftLoaded] = useState(false);
-  const lastServerVersionRef = useRef<number>(-1);
-  const suppressNextSaveRef = useRef<boolean>(false);
-  const forceNextApplyRef = useRef<boolean>(false);
+  const { projectId } = useProject();
+  const drafts = useDraftsStreamState(projectId);
 
-  const endpoint = attemptId
-    ? `/api/task-attempts/${attemptId}/follow-up-draft/stream/ws`
-    : undefined;
-
-  const makeInitial = useCallback(
-    (): DraftStreamState => ({
-      follow_up_draft: {
-        id: '',
-        task_attempt_id: attemptId || '',
-        prompt: '',
-        queued: false,
-        sending: false,
-        variant: null,
-        image_ids: [],
-        version: 0n,
-        created_at: new Date().toISOString(),
-        updated_at: new Date().toISOString(),
-      },
-    }),
-    [attemptId]
-  );
-
-  const { data, isConnected, error } = useJsonPatchWsStream<DraftStreamState>(
-    endpoint,
-    !!endpoint,
-    makeInitial
-  );
-
-  // Quick initial draft loading from REST
-  useEffect(() => {
-    let cancelled = false;
-    const hydrate = async () => {
-      if (!attemptId) return;
-      try {
-        const d = await attemptsApi.getFollowUpDraft(attemptId);
-        if (cancelled) return;
-        suppressNextSaveRef.current = true;
-        setDraft({
-          id: 'rest',
-          task_attempt_id: d.task_attempt_id,
-          prompt: d.prompt || '',
-          queued: !!d.queued,
-          sending: false,
-          variant: (d.variant ?? null) as string | null,
-          image_ids: (d.image_ids ?? []) as string[],
-          version: (d.version ?? 0n) as unknown as bigint,
-          created_at: new Date().toISOString(),
-          updated_at: new Date().toISOString(),
-        });
-        if (!isDraftLoaded) setIsDraftLoaded(true);
-      } catch {
-        // ignore, rely on stream
-      }
-    };
-    hydrate();
-    return () => {
-      cancelled = true;
-    };
-  }, [attemptId, isDraftLoaded]);
-
-  // Handle stream updates
-  useEffect(() => {
-    if (!data) return;
-    const d = data.follow_up_draft;
-    if (d.id === '') return;
-    const incomingVersion = Number(d.version ?? 0n);
-    if (incomingVersion === lastServerVersionRef.current) {
-      if (!isDraftLoaded) setIsDraftLoaded(true);
-      return;
-    }
-    suppressNextSaveRef.current = true;
-    // Let consumers decide whether to apply or ignore based on local dirty/forceApply.
-    setDraft(d);
-    if (!isDraftLoaded) setIsDraftLoaded(true);
-  }, [data, isDraftLoaded]);
-
-  // VSCode iframe poll fallback
-  const pollTimerRef = useRef<number | undefined>(undefined);
-  useEffect(() => {
-    if (!attemptId) return;
-    const shouldPoll = inIframe() && (!isConnected || !!error);
-    if (!shouldPoll) {
-      if (pollTimerRef.current) window.clearInterval(pollTimerRef.current);
-      pollTimerRef.current = undefined;
-      return;
-    }
-    const pollOnce = async () => {
-      try {
-        const d = await attemptsApi.getFollowUpDraft(attemptId);
-        const incomingVersion = Number((d as FollowUpDraft).version ?? 0n);
-        if (incomingVersion !== lastServerVersionRef.current) {
-          suppressNextSaveRef.current = true;
-          setDraft({
-            id: 'rest',
-            task_attempt_id: d.task_attempt_id,
-            prompt: d.prompt || '',
-            queued: !!d.queued,
-            sending: false,
-            variant: (d.variant ?? null) as string | null,
-            image_ids: (d.image_ids ?? []) as string[],
-            version: (d.version ?? 0n) as unknown as bigint,
-            created_at: new Date().toISOString(),
-            updated_at: new Date().toISOString(),
-          });
-          if (!isDraftLoaded) setIsDraftLoaded(true);
-        }
-      } catch {
-        // ignore
-      }
-    };
-    pollOnce();
-    pollTimerRef.current = window.setInterval(pollOnce, 1000);
-    return () => {
-      if (pollTimerRef.current) window.clearInterval(pollTimerRef.current);
-      pollTimerRef.current = undefined;
-    };
-  }, [attemptId, isConnected, error, isDraftLoaded]);
+  const attemptDrafts = useMemo(() => {
+    if (!attemptId || !drafts) return null;
+    return drafts[attemptId] ?? null;
+  }, [drafts, attemptId]);
 
   return {
-    draft,
-    isDraftLoaded,
-    isConnected,
-    error,
-    lastServerVersionRef,
-    suppressNextSaveRef,
-    forceNextApplyRef,
+    draft: attemptDrafts?.follow_up ?? null,
+    retryDraft: attemptDrafts?.retry ?? null,
+    isRetryLoaded: !!attemptDrafts,
+    isDraftLoaded: !!attemptDrafts,
   } as const;
 }
+
+function useDraftsStreamState(projectId?: string): Drafts | undefined {
+  const endpoint = useMemo(
+    () =>
+      projectId
+        ? `/api/drafts/stream/ws?project_id=${encodeURIComponent(projectId)}`
+        : undefined,
+    [projectId]
+  );
+  const wsUrl = useMemo(() => toWsUrl(endpoint), [endpoint]);
+  const isStreamEnabled = !!endpoint && !!wsUrl;
+
+  const queryClient = useQueryClient();
+  const initialData = useCallback((): DraftsContainer => ({ drafts: {} }), []);
+  const queryKey = useMemo(() => ['ws-json-patch', wsUrl], [wsUrl]);
+
+  const { data } = useQuery<DraftsContainer | undefined>({
+    queryKey,
+    enabled: isStreamEnabled,
+    staleTime: Infinity,
+    gcTime: 0,
+    initialData: undefined,
+  });
+
+  useEffect(() => {
+    if (!isStreamEnabled) return;
+    const current = queryClient.getQueryData<DraftsContainer | undefined>(
+      queryKey
+    );
+    if (current === undefined) {
+      queryClient.setQueryData<DraftsContainer>(queryKey, initialData());
+    }
+  }, [isStreamEnabled, queryClient, queryKey, initialData]);
+
+  const { getWebSocket } = useWebSocket(
+    wsUrl ?? 'ws://invalid',
+    {
+      share: true,
+      shouldReconnect: () => true,
+      reconnectInterval: (attempt) =>
+        Math.min(8000, 1000 * Math.pow(2, attempt)),
+      retryOnError: true,
+      onMessage: (event) => {
+        try {
+          const msg: WsMsg = JSON.parse(event.data);
+          if ('JsonPatch' in msg) {
+            const patches = msg.JsonPatch;
+            if (!patches.length) return;
+            queryClient.setQueryData<DraftsContainer | undefined>(
+              queryKey,
+              (prev) => {
+                const base = prev ?? initialData();
+                const next = structuredClone(base) as DraftsContainer;
+                applyPatch(next, patches);
+                return next;
+              }
+            );
+          } else if ('finished' in msg) {
+            try {
+              getWebSocket()?.close();
+            } catch {
+              /* noop */
+            }
+          }
+        } catch (e) {
+          console.error('Failed to process WebSocket message:', e);
+        }
+      },
+    },
+    isStreamEnabled
+  );
+
+  return isStreamEnabled ? data?.drafts : undefined;
+}
+
+function toWsUrl(endpoint?: string): string | undefined {
+  if (!endpoint) return undefined;
+  try {
+    const url = new URL(endpoint, window.location.origin);
+    url.protocol = url.protocol.replace('http', 'ws');
+    return url.toString();
+  } catch {
+    return undefined;
+  }
+}

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

@@ -52,7 +52,10 @@ export function useFollowUpSend({
         prompt: finalPrompt,
         variant: selectedVariant,
         image_ids,
-      });
+        retry_process_id: null,
+        force_when_dirty: null,
+        perform_git_reset: null,
+      } as any);
       setMessage('');
       clearComments();
       onAfterSendCleanup();

frontend/src/hooks/useProcessRetry.ts

@@ -2,29 +2,8 @@
 import { useCallback, useMemo, useState } from 'react';
 import { useAttemptExecution } from '@/hooks/useAttemptExecution';
 import { useBranchStatus } from '@/hooks/useBranchStatus';
-import { showModal } from '@/lib/modals';
-import {
-  shouldShowInLogs,
-  isCodingAgent,
-  PROCESS_RUN_REASONS,
-} from '@/constants/processes';
+import { attemptsApi, executionProcessesApi } from '@/lib/api';
 import type { ExecutionProcess, TaskAttempt } from 'shared/types';
-import type {
-  ExecutorActionType,
-  CodingAgentInitialRequest,
-  CodingAgentFollowUpRequest,
-} from 'shared/types';
-
-function isCodingAgentActionType(
-  t: ExecutorActionType
-): t is
-  | ({ type: 'CodingAgentInitialRequest' } & CodingAgentInitialRequest)
-  | ({ type: 'CodingAgentFollowUpRequest' } & CodingAgentFollowUpRequest) {
-  return (
-    t.type === 'CodingAgentInitialRequest' ||
-    t.type === 'CodingAgentFollowUpRequest'
-  );
-}
 
 /**
  * Reusable hook to retry a process given its executionProcessId and a new prompt.
@@ -38,10 +17,8 @@ export function useProcessRetry(attempt: TaskAttempt | undefined) {
   const attemptId = attempt?.id;
 
   // Fetch attempt + branch state the same way your component did
-  const { attemptData, refetch: refetchAttempt } =
-    useAttemptExecution(attemptId);
-  const { data: branchStatus, refetch: refetchBranch } =
-    useBranchStatus(attemptId);
+  const { attemptData } = useAttemptExecution(attemptId);
+  useBranchStatus(attemptId);
 
   const [busy, setBusy] = useState(false);
 
@@ -79,149 +56,49 @@ export function useProcessRetry(attempt: TaskAttempt | undefined) {
   /**
    * Primary entrypoint: retry a process with a new prompt.
    */
-  const retryProcess = useCallback(
+  // Initialize retry mode by creating a retry draft populated from the process
+  const startRetry = useCallback(
     async (executionProcessId: string, newPrompt: string) => {
       if (!attemptId) return;
-
       const proc = getProcessById(executionProcessId);
       if (!proc) return;
-
-      // Respect current disabled state
       const { disabled } = getRetryDisabledState(executionProcessId);
       if (disabled) return;
 
-      type WithBefore = { before_head_commit?: string | null };
-      const before =
-        (proc as WithBefore | undefined)?.before_head_commit || null;
-
-      // Try to gather comparison info (best-effort)
-      let targetSubject: string | null = null;
-      let commitsToReset: number | null = null;
-      let isLinear: boolean | null = null;
-
-      if (before) {
-        try {
-          const { commitsApi } = await import('@/lib/api');
-          const info = await commitsApi.getInfo(attemptId, before);
-          targetSubject = info.subject;
-          const cmp = await commitsApi.compareToHead(attemptId, before);
-          commitsToReset = cmp.is_linear ? cmp.ahead_from_head : null;
-          isLinear = cmp.is_linear;
-        } catch {
-          // ignore best-effort enrichments
-        }
-      }
-
-      const head = branchStatus?.head_oid || null;
-      const dirty = !!branchStatus?.has_uncommitted_changes;
-      const needReset = !!(before && (before !== head || dirty));
-      const canGitReset = needReset && !dirty;
-
-      // Compute “later processes” context for the dialog
-      const procs = (attemptData.processes || []).filter(
-        (p) => !p.dropped && shouldShowInLogs(p.run_reason)
-      );
-      const idx = procs.findIndex((p) => p.id === executionProcessId);
-      const later = idx >= 0 ? procs.slice(idx + 1) : [];
-      const laterCount = later.length;
-      const laterCoding = later.filter((p) =>
-        isCodingAgent(p.run_reason)
-      ).length;
-      const laterSetup = later.filter(
-        (p) => p.run_reason === PROCESS_RUN_REASONS.SETUP_SCRIPT
-      ).length;
-      const laterCleanup = later.filter(
-        (p) => p.run_reason === PROCESS_RUN_REASONS.CLEANUP_SCRIPT
-      ).length;
-
-      // Ask user for confirmation / reset options
-      let modalResult:
-        | {
-            action: 'confirmed' | 'canceled';
-            performGitReset?: boolean;
-            forceWhenDirty?: boolean;
-          }
-        | undefined;
-
-      try {
-        modalResult = await showModal<
-          typeof modalResult extends infer T
-            ? T extends object
-              ? T
-              : never
-            : never
-        >('restore-logs', {
-          targetSha: before,
-          targetSubject,
-          commitsToReset,
-          isLinear,
-          laterCount,
-          laterCoding,
-          laterSetup,
-          laterCleanup,
-          needGitReset: needReset,
-          canGitReset,
-          hasRisk: dirty,
-          uncommittedCount: branchStatus?.uncommitted_count ?? 0,
-          untrackedCount: branchStatus?.untracked_count ?? 0,
-          // Defaults
-          initialWorktreeResetOn: true,
-          initialForceReset: false,
-        });
-      } catch {
-        // user closed dialog
-        return;
-      }
-
-      if (!modalResult || modalResult.action !== 'confirmed') return;
-
       let variant: string | null = null;
-
-      const typ = proc?.executor_action?.typ; // type: ExecutorActionType
-
-      if (typ && isCodingAgentActionType(typ)) {
-        // executor_profile_id is ExecutorProfileId -> has `variant: string | null`
-        variant = typ.executor_profile_id.variant;
+      try {
+        const details =
+          await executionProcessesApi.getDetails(executionProcessId);
+        const typ: any = details?.executor_action?.typ as any;
+        if (
+          typ &&
+          (typ.type === 'CodingAgentInitialRequest' ||
+            typ.type === 'CodingAgentFollowUpRequest')
+        ) {
+          variant = (typ.executor_profile_id?.variant as string | null) ?? null;
+        }
+      } catch {
+        /* ignore */
       }
 
-      // Perform the replacement
+      setBusy(true);
       try {
-        setBusy(true);
-        const { attemptsApi } = await import('@/lib/api');
-        await attemptsApi.replaceProcess(attemptId, {
-          process_id: executionProcessId,
+        await attemptsApi.saveDraft(attemptId, 'retry', {
+          retry_process_id: executionProcessId,
           prompt: newPrompt,
           variant,
-          perform_git_reset: modalResult.performGitReset ?? true,
-          force_when_dirty: modalResult.forceWhenDirty ?? false,
+          image_ids: [],
+          version: null as any,
         });
-
-        // Refresh local caches
-        await refetchAttempt();
-        await refetchBranch();
       } finally {
         setBusy(false);
       }
     },
-    [
-      attemptId,
-      attemptData.processes,
-      branchStatus?.head_oid,
-      branchStatus?.has_uncommitted_changes,
-      branchStatus?.uncommitted_count,
-      branchStatus?.untracked_count,
-      getProcessById,
-      getRetryDisabledState,
-      refetchAttempt,
-      refetchBranch,
-    ]
+    [attemptId, getProcessById, getRetryDisabledState]
   );
 
   return {
-    retryProcess,
-    busy,
-    anyRunning,
-    /** Helpful for buttons/tooltips */
+    startRetry,
     getRetryDisabledState,
   };
 }