patch_auto_continue_v3.py

"""V3: FINAL - Fix remaining TS errors for auto-continue feature.

Errors fixed:
1. TS6133: 'showAutoContinue'/'taskIncompleteInfo' never read → prefix with _
2. TS2741: Property 'onTaskIncomplete' missing in sideChannel object → add implementation
"""

import os
import re

USE_AGENT_CHAT = "/source/frontend/src/hooks/useAgentChat.ts"
EVENTS_FILE = "/source/frontend/src/types/events.ts"
CHAT_INPUT = "/source/frontend/src/components/Chat/ChatInput.tsx"
SSE_TRANSPORT = "/source/frontend/src/lib/sse-chat-transport.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")

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 import
    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 after callbacksRef
    state_block = '''
  // 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);
'''
    if "autoContinueTimerRef" not in content:
        content = content.replace(
            "callbacksRef.current = { onReady, onError, onSessionDead };",
            "callbacksRef.current = { onReady, onError, onSessionDead };" + state_block
        )
    
    # 3. Add onTaskIncomplete to interface
    if "onTaskIncomplete:" not in content and "onInterrupted: () => void;" in content:
        content = content.replace(
            "  onInterrupted: () => void;",
            "  onInterrupted: () => void;\n  onTaskIncomplete: (incompletePlan: Array<{ id: string; content: string; status: string }>) => void;"
        )
    
    # 4. Add onTaskIncomplete handler in sideChannel (BEFORE onInterrupted)
    handler_block = '''      onTaskIncomplete: (incompletePlan: Array<{ id: string; content: string; status: string }>) => {
        _setTaskIncompleteInfo({ incompletePlan });
        _setShowAutoContinue(true);
        if (autoContinueTimerRef.current) {
          clearTimeout(autoContinueTimerRef.current);
        }
        autoContinueTimerRef.current = setTimeout(() => {
          _setShowAutoContinue(false);
          _setTaskIncompleteInfo(null);
          const plan = incompletePlan.map(i => `- ${i.content}`).join('\\n');
          chat.sendMessage({
            text: `[TỰ ĐỘNG TIẾP TỤC] Tiếp tục từ các task chưa hoàn thành:\\n${plan}`,
            metadata: { createdAt: new Date().toISOString() },
          });
        }, 10000);
      },
      onInterrupted: () => { /* no-op */ },'''
    
    if "onTaskIncomplete: (incompletePlan" not in content:
        old_interrupted = '      onInterrupted: () => { /* no-op — handled by stop() caller */ },'
        if old_interrupted in content:
            content = content.replace(old_interrupted, handler_block)
        else:
            # Try the simpler version
            simple_interrupted = '      onInterrupted: () => { /* no-op */ },'
            if simple_interrupted in content:
                content = content.replace(simple_interrupted, handler_block)
    
    # 5. Add cleanup timer before return
    cleanup = '''
  // Cleanup auto-continue timer
  useEffect(() => {
    return () => {
      if (autoContinueTimerRef.current) {
        clearTimeout(autoContinueTimerRef.current);
      }
    };
  }, []);
'''
    if "autoContinueTimerRef.current" in content and "Cleanup timer" not in content:
        content = content.replace(
            '\n  return {\n    messages: chat.messages,',
            cleanup + '\n  return {\n    messages: chat.messages,'
        )
    
    with open(USE_AGENT_CHAT, "w") as f:
        f.write(content)
    print(f"OK: Patched {USE_AGENT_CHAT}")

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 - define inline type
    if "DatasetUploadResponse" in content and "type DatasetUploadResponse" not in content:
        wrong = "import type { DatasetUploadResponse } from '@/types/agent';"
        if wrong in content:
            content = content.replace(wrong, "")
            # Add inline type after apiFetch import
            content = content.replace(
                "import { apiFetch } from '@/utils/api';",
                "import { apiFetch } from '@/utils/api';\ntype DatasetUploadResponse = { url: string; repo_id: string; filename: string; size: number; };"
            )
    
    # Add Button import
    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';"
        )
    
    # Fix interface - replace with complete version
    if "_showAutoContinue" not in content:
        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;
  _onCancelAutoContinue?: () => void;
}
'''
            content = content.replace(interface_match.group(0), props_code)
    
    # Add cancel button
    buttons = '''
      {/* Auto-continue cancel button for free models */}
      {_showAutoContinue && _taskIncompleteInfo && _onCancelAutoContinue && (
        <Stack direction="row" spacing={1} sx={{ mb: 1, px: 1 }}>
          <Button
            variant="outlined"
            color="secondary"
            size="small"
            onClick={_onCancelAutoContinue}
            sx={{ textTransform: 'none', fontSize: '0.75rem' }}
          >
            Tạm dừng (để nhập nội dung khác)
          </Button>
        </Stack>
      )}
'''
    if "Tạm dừng" not in content:
        content = content.replace(
            "<Box sx={{ flex: 1 ",
            buttons + "\n        <Box sx={{ flex: 1 "
        )
    
    with open(CHAT_INPUT, "w") as f:
        f.write(content)
    print(f"OK: Patched {CHAT_INPUT}")

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: () => void;" in content:
        content = content.replace(
            "  onInterrupted: () => void;",
            "  onInterrupted: () => void;\n  onTaskIncomplete: (incompletePlan: Array<{ id: string; content: string; status: string }>) => void;"
        )
    
    # Add case 'task_incomplete'
    if "case 'task_incomplete'" not in content:
        task_case = '''      case 'task_incomplete' as const:
          sideChannel.onTaskIncomplete(
            (event.data?.incomplete_plan as Array<{ id: string; content: string; status: string }>) || [],
          );
          break;
'''
        content = content.replace(
            "case 'turn_complete':",
            task_case + "\n        case 'turn_complete':"
        )
    
    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: V3 patches applied")