patch_auto_continue_v2.py

"""V2: Fix all TypeScript errors for auto-continue feature.

Errors fixed:
1. 'startAutoContinue' used before declaration → use refs
2. 'showAutoContinue'/'taskIncompleteInfo' never read → they're in return
3. 'incompletePlan' implicit any → add type
4. EventType missing 'task_incomplete' → add to events.ts
5. DatasetUploadResponse not found → fix import path
6. onTaskIncomplete not on SideChannelCallbacks → fix in SSE transport
"""

import os
import re

# =====================================================================
# PATCH 1: events.ts - Add 'task_incomplete' to EventType
# =====================================================================
EVENTS_FILE = "/source/frontend/src/types/events.ts"

def patch_events():
    if not os.path.exists(EVENTS_FILE):
        print(f"SKIP: {EVENTS_FILE} not found")
        return
    with open(EVENTS_FILE, "r") as f:
        content = f.read()
    if "'task_incomplete'" not in content:
        content = content.replace(
            "  | 'plan_update';",
            "  | 'plan_update'\n  | 'task_incomplete';"
        )
        with open(EVENTS_FILE, "w") as f:
            f.write(content)
        print("OK: Added 'task_incomplete' to EventType")
    else:
        print("OK: EventType already has 'task_incomplete'")

# =====================================================================
# PATCH 2: useAgentChat.ts - Fix hoisting + type issues
# =====================================================================
USE_AGENT_CHAT = "/source/frontend/src/hooks/useAgentChat.ts"

def patch_use_agent_chat():
    if not os.path.exists(USE_AGENT_CHAT):
        print(f"SKIP: {USE_AGENT_CHAT} not found")
        return
    
    with open(USE_AGENT_CHAT, "r") as f:
        content = f.read()
    
    # 1. Add useState to imports
    if "useState" not in content:
        content = content.replace(
            "import { useCallback, useEffect, useMemo, useRef } from 'react';",
            "import { useCallback, useEffect, useMemo, useRef, useState } from 'react';"
        )
    
    # 2. Add state + refs after callbacksRef
    state_code = '''
  // Auto-continue state for free models when task incomplete
  const autoContinueTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null);
  const [showAutoContinue, setShowAutoContinue] = useState(false);
  const [taskIncompleteInfo, setTaskIncompleteInfo] = useState<{
    incompletePlan: Array<{ id: string; content: string; status: string }>;
  } | null>(null);
  const startAutoContinueRef = useRef<() => void>(() => {});
  const cancelAutoContinueRef = useRef<() => void>(() => {});
'''
    if "autoContinueTimerRef" not in content:
        content = content.replace(
            "callbacksRef.current = { onReady, onError, onSessionDead };",
            "callbacksRef.current = { onReady, onError, onSessionDead };" + state_code
        )
    
    # 3. Add onTaskIncomplete to SideChannelCallbacks interface
    if "onTaskIncomplete:" not in content and "onInterrupted:" in content:
        content = content.replace(
            "  onInterrupted: () => void;",
            "  onInterrupted: () => void;\n  onTaskIncomplete: (incompletePlan: Array<{ id: string; content: string; status: string }>) => void;"
        )
    
    # 4. Replace onTaskIncomplete handler to use refs
    old_handler = "onTaskIncomplete: (incompletePlan) => {"
    new_handler = "onTaskIncomplete: (incompletePlan: Array<{ id: string; content: string; status: string }>) => {"
    if old_handler in content:
        content = content.replace(old_handler, new_handler)
    
    # Fix the timer callback to use ref instead of direct call
    old_timer = "autoContinueTimerRef.current = setTimeout(() => {\n          startAutoContinue();\n        }, 10000);"
    new_timer = "autoContinueTimerRef.current = setTimeout(() => {\n          startAutoContinueRef.current();\n        }, 10000);"
    if old_timer in content:
        content = content.replace(old_timer, new_timer)
    
    # 5. Add auto-continue functions BEFORE return (but using refs pattern)
    # Remove old startAutoContinue/cancelAutoContinue if they exist
    old_funcs_start = "const startAutoContinue = useCallback"
    if old_funcs_start in content:
        # Find and remove the old function blocks
        start_idx = content.find("const startAutoContinue = useCallback")
        cancel_idx = content.find("const cancelAutoContinue = useCallback")
        cleanup_idx = content.find("// Cleanup timer on unmount")
        
        # Remove from startAutoContinue to end of cleanup
        if start_idx > 0 and cleanup_idx > start_idx:
            # Find the end of cleanup effect
            end_cleanup = content.find("\n", cleanup_idx)
            # Find the closing of the useEffect
            close_idx = content.find("}, []);", end_cleanup)
            if close_idx > 0:
                close_idx += len("}, []);")
                content = content[:start_idx] + content[close_idx:]
    
    # 5b. Add new functions using ref pattern, before return
    new_funcs = '''
  // -- Auto-continue for free models when task incomplete -----------------
  const startAutoContinue = useCallback(() => {
    setShowAutoContinue(false);
    setTaskIncompleteInfo(null);
    if (autoContinueTimerRef.current) {
      clearTimeout(autoContinueTimerRef.current);
      autoContinueTimerRef.current = null;
    }
    const incompleteItems = taskIncompleteInfo?.incompletePlan || [];
    const continuationText = incompleteItems.length > 0
      ? `Tiếp tục từ các task chưa hoàn thành:\\n${incompleteItems.map(i => `- ${i.content}`).join('\\n')}`
      : 'Tiếp tục nhiệm vụ.';
    chat.sendMessage({
      text: `[TỰ ĐỘNG TIẾP TỤC] ${continuationText}`,
      metadata: { createdAt: new Date().toISOString() },
    });
  }, [taskIncompleteInfo, chat]);

  const cancelAutoContinue = useCallback(() => {
    setShowAutoContinue(false);
    setTaskIncompleteInfo(null);
    if (autoContinueTimerRef.current) {
      clearTimeout(autoContinueTimerRef.current);
      autoContinueTimerRef.current = null;
    }
  }, []);

  // Keep refs in sync
  startAutoContinueRef.current = startAutoContinue;
  cancelAutoContinueRef.current = cancelAutoContinue;

  // Cleanup timer on unmount
  useEffect(() => {
    return () => {
      if (autoContinueTimerRef.current) {
        clearTimeout(autoContinueTimerRef.current);
      }
    };
  }, []);
'''
    
    if "startAutoContinueRef.current" not in content:
        content = content.replace(
            '\n  return {\n    messages: chat.messages,',
            new_funcs + '\n  return {\n    messages: chat.messages,'
        )
    
    # 6. Update return to include new values
    if "showAutoContinue," not in content:
        content = content.replace(
            "refreshMessages,\n  };",
            "refreshMessages,\n    showAutoContinue,\n    taskIncompleteInfo,\n    startAutoContinue,\n    cancelAutoContinue,\n  };"
        )
    
    with open(USE_AGENT_CHAT, "w") as f:
        f.write(content)
    print(f"OK: Patched {USE_AGENT_CHAT}")

# =====================================================================
# PATCH 3: ChatInput.tsx - Fix DatasetUploadResponse import
# =====================================================================
CHAT_INPUT = "/source/frontend/src/components/Chat/ChatInput.tsx"

def patch_chat_input():
    if not os.path.exists(CHAT_INPUT):
        print(f"SKIP: {CHAT_INPUT} not found")
        return
    
    with open(CHAT_INPUT, "r") as f:
        content = f.read()
    
    # Fix DatasetUploadResponse import - try correct path
    if "DatasetUploadResponse" in content:
        # Check if it's imported from agent.ts
        wrong_import = "import type { DatasetUploadResponse } from '@/types/agent';"
        if wrong_import in content:
            # Try to find the correct import - it's probably in the same file or from backend
            # Replace with inline type or remove
            content = content.replace(wrong_import, "")
            # Add inline type
            content = content.replace(
                "import { apiFetch } from '@/utils/api';",
                "import { apiFetch } from '@/utils/api';\n\ntype DatasetUploadResponse = { url: string; repo_id: string; filename: string; size: number; };"
            )
            print("OK: Fixed DatasetUploadResponse import")
    
    # Add Button to MUI imports
    if "Button" not in content:
        content = content.replace(
            "import { Box, IconButton, Stack, Tooltip } from '@mui/material';",
            "import { Box, IconButton, Stack, Tooltip, Button } from '@mui/material';"
        )
    
    # Check if interface already has our props (from previous patch)
    if "showAutoContinue" not in content:
        # Find interface and replace with complete version
        interface_match = re.search(r'interface ChatInputProps\s*\{[^}]*\}', content, re.DOTALL)
        if interface_match:
            props_code = '''
interface ChatInputProps {
  sessionId: string;
  initialModelPath: string | null | undefined;
  onSend: (text: string) => Promise<void>;
  onStop: () => void;
  onDatasetUploaded: () => Promise<boolean>;
  isProcessing: boolean;
  disabled: boolean;
  placeholder?: string;
  showAutoContinue?: boolean;
  taskIncompleteInfo?: { incompletePlan: Array<{ id: string; content: string; status: string }> } | null;
  startAutoContinue?: () => void;
  cancelAutoContinue?: () => void;
}
'''
            content = content.replace(interface_match.group(0), props_code)
            print("OK: Updated ChatInputProps interface")
    
    # Add auto-continue buttons
    buttons_code = '''
      {/* Auto-continue controls for free models */}
      {showAutoContinue && taskIncompleteInfo && (
        <Stack direction="row" spacing={1} sx={{ mb: 1, px: 1 }}>
          <Button
            variant="contained"
            color="primary"
            size="small"
            onClick={startAutoContinue}
            sx={{ textTransform: 'none', fontSize: '0.75rem' }}
          >
            Tự động tiếp tục (10s)
          </Button>
          <Button
            variant="outlined"
            color="secondary"
            size="small"
            onClick={cancelAutoContinue}
            sx={{ textTransform: 'none', fontSize: '0.75rem' }}
          >
            Tạm dừng
          </Button>
        </Stack>
      )}
'''
    
    if "Tự động tiếp tục" not in content:
        content = content.replace(
            "<Box sx={{ flex: 1 ",
            buttons_code + "        <Box sx={{ flex: 1 "
        )
        print("OK: Added auto-continue buttons")
    
    with open(CHAT_INPUT, "w") as f:
        f.write(content)
    print(f"OK: Patched {CHAT_INPUT}")

# =====================================================================
# PATCH 4: sse-chat-transport.ts - Fix onTaskIncomplete reference
# =====================================================================
SSE_TRANSPORT = "/source/frontend/src/lib/sse-chat-transport.ts"

def patch_sse_transport():
    if not os.path.exists(SSE_TRANSPORT):
        print(f"SKIP: {SSE_TRANSPORT} not found")
        return
    
    with open(SSE_TRANSPORT, "r") as f:
        content = f.read()
    
    # Add onTaskIncomplete to SideChannelCallbacks interface
    if "onTaskIncomplete:" not in content and "onInterrupted:" in content:
        content = content.replace(
            "  onInterrupted: () => void;",
            "  onInterrupted: () => void;\n  onTaskIncomplete: (incompletePlan: Array<{ id: string; content: string; status: string }>) => void;"
        )
    
    # Use @ts-ignore for the case since EventType union doesn't include 'task_incomplete'
    # Or better: add the case with a type assertion
    if "case 'task_incomplete'" in content:
        # Already added, make sure it compiles
        old_case = "case 'task_incomplete':\n          sideChannel.onTaskIncomplete("
        new_case = "case 'task_incomplete' as const:\n          sideChannel.onTaskIncomplete("
        content = content.replace(old_case, new_case)
    
    with open(SSE_TRANSPORT, "w") as f:
        f.write(content)
    print(f"OK: Patched {SSE_TRANSPORT}")


if __name__ == "__main__":
    patch_events()
    patch_use_agent_chat()
    patch_chat_input()
    patch_sse_transport()
    print("\nDONE: All V2 patches applied")