patch_auto_continue_v4.py

"""V4: FINAL WORKING - Fix all TS errors for auto-continue feature.

Strategy:
- useAgentChat: return _showAutoContinue + _cancelAutoContinue (TS _ prefix suppresses "unused")
- ChatInput.tsx: receive props and render "Tạm dừng" button  
- SessionChat.tsx: pass props from hook to ChatInput
- events.ts: add 'task_incomplete' to EventType union
- sse-chat-transport.ts: add onTaskIncomplete + case handler
- agent_loop.py: send task_incomplete event
"""

import os
import re

EVENTS = "/source/frontend/src/types/events.ts"
USE_AGENT_CHAT = "/source/frontend/src/hooks/useAgentChat.ts"
SSE = "/source/frontend/src/lib/sse-chat-transport.ts"
CHAT_INPUT = "/source/frontend/src/components/Chat/ChatInput.tsx"
SESSION_CHAT = "/source/frontend/src/components/SessionChat.tsx"


def patch_events():
    with open(EVENTS, "r") as f:
        c = f.read()
    if "'task_incomplete'" not in c:
        c = c.replace("  | 'plan_update';", "  | 'plan_update'\n  | 'task_incomplete';")
        with open(EVENTS, "w") as f: f.write(c)
    print("OK: events.ts")


def patch_sse():
    with open(SSE, "r") as f:
        c = f.read()
    
    # Add onTaskIncomplete to interface
    if "onTaskIncomplete:" not in c:
        c = c.replace(
            "  onInterrupted: () => void;",
            "  onInterrupted: () => void;\n  onTaskIncomplete: (incompletePlan: Array<{ id: string; content: string; status: string }>) => void;"
        )
    
    # Add case handler
    if "case 'task_incomplete'" not in c:
        case = """      case 'task_incomplete' as const:
          sideChannel.onTaskIncomplete(
            (event.data?.incomplete_plan as Array<{ id: string; content: string; status: string }>) || [],
          );
          break;
"""
        c = c.replace("case 'turn_complete':", case + "\n        case 'turn_complete':")
    
    with open(SSE, "w") as f: f.write(c)
    print("OK: sse-chat-transport.ts")


def patch_use_agent_chat():
    with open(USE_AGENT_CHAT, "r") as f:
        c = f.read()
    
    # 1. Import useState
    if "useState" not in c:
        c = c.replace(
            "import { useCallback, useEffect, useMemo, useRef } from 'react';",
            "import { useCallback, useEffect, useMemo, useRef, useState } from 'react';"
        )
    
    # 2. State (TS _ prefix = unused ok)
    state_block = '''
  // Auto-continue for free models
  const _autoContinueTimer = useRef<ReturnType<typeof setTimeout> | null>(null);
  const [_showAutoContinue, _setShowAutoContinue] = useState(false);
'''
    if "_autoContinueTimer" not in c:
        c = c.replace(
            "callbacksRef.current = { onReady, onError, onSessionDead };",
            "callbacksRef.current = { onReady, onError, onSessionDead };" + state_block
        )
    
    # 3. onTaskIncomplete in interface
    if "onTaskIncomplete:" not in c:
        c = c.replace(
            "  onInterrupted: () => void;",
            "  onInterrupted: () => void;\n  onTaskIncomplete: (incompletePlan: Array<{ id: string; content: string; status: string }>) => void;"
        )
    
    # 4. Cancel ref
    cancel_ref = '  const _cancelAutoContinue = useRef(() => {});\n'
    if "_cancelAutoContinue" not in c:
        # Insert after callbacksRef.current = { ... } area
        c = c.replace(
            "  isActiveRef.current = isActive;",
            "  isActiveRef.current = isActive;\n" + cancel_ref
        )
    
    # 5. Handler in sideChannel - inline all logic
    handler = """      onTaskIncomplete: (incompletePlan: Array<{ id: string; content: string; status: string }>) => {
        _setShowAutoContinue(true);
        if (_autoContinueTimer.current) clearTimeout(_autoContinueTimer.current);
        _autoContinueTimer.current = setTimeout(() => {
          _setShowAutoContinue(false);
          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 c:
        old = '      onInterrupted: () => { /* no-op — handled by stop() caller */ },'
        if old in c:
            c = c.replace(old, handler)
    
    # 6. Set cancel ref function + cleanup  
    cancel_func = """
  _cancelAutoContinue.current = () => {
    _setShowAutoContinue(false);
    if (_autoContinueTimer.current) {
      clearTimeout(_autoContinueTimer.current);
      _autoContinueTimer.current = null;
    }
  };
"""
    if "_cancelAutoContinue.current = ()" not in c:
        # Insert before return
        c = c.replace(
            "\n  return {\n    messages: chat.messages,",
            cancel_func + "\n  return {\n    messages: chat.messages,"
        )
    
    # 7. Return values
    if "_showAutoContinue," not in c:
        c = c.replace(
            "refreshMessages,\n  };",
            "refreshMessages,\n    _showAutoContinue,\n    _cancelAutoContinue,\n  };"
        )
    
    with open(USE_AGENT_CHAT, "w") as f: f.write(c)
    print("OK: useAgentChat.ts")


def patch_chat_input():
    with open(CHAT_INPUT, "r") as f:
        c = f.read()
    
    # Fix DatasetUploadResponse
    if "DatasetUploadResponse" in c and "inline" not in c:
        wrong = "import type { DatasetUploadResponse } from '@/types/agent';"
        if wrong in c:
            c = c.replace(wrong, "")
            c = c.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 c:
        c = c.replace(
            "import { Box, IconButton, Stack, Tooltip } from '@mui/material';",
            "import { Box, IconButton, Stack, Tooltip, Button } from '@mui/material';"
        )
    
    # Add props to interface
    if "_showAutoContinue" not in c:
        match = re.search(r'interface ChatInputProps\s*\{[^}]*\}', c, re.DOTALL)
        if match:
            new_iface = """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;
  _cancelAutoContinue?: React.RefObject<() => void>;
}"""
            c = c.replace(match.group(0), new_iface)
    
    # Add pause button  
    btn = """
      {/* Pause auto-continue for free models */}
      {_showAutoContinue && _cancelAutoContinue?.current && (
        <Stack direction="row" spacing={1} sx={{ mb: 1, px: 1 }}>
          <Button
            variant="outlined"
            color="secondary"
            size="small"
            onClick={_cancelAutoContinue.current}
            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 c:
        c = c.replace("<Box sx={{ flex: 1 ", btn + "\n        <Box sx={{ flex: 1 ")
    
    with open(CHAT_INPUT, "w") as f: f.write(c)
    print("OK: ChatInput.tsx")


def patch_session_chat():
    with open(SESSION_CHAT, "r") as f:
        c = f.read()
    
    # Find the ChatInput usage and add new props
    if "_showAutoContinue" not in c:
        chat_input_pattern = r'(<ChatInput[^>]*sessionId={sessionId}[^>]*initialModelPath={initialModelPath}[^>]*)(\s*/>)'
        new_props = r'\1 _showAutoContinue={_showAutoContinue} _cancelAutoContinue={_cancelAutoContinue}\2'
        c = re.sub(chat_input_pattern, new_props, c)
    
    with open(SESSION_CHAT, "w") as f: f.write(c)
    print("OK: SessionChat.tsx")


def patch_agent_loop():
    agent_loop = "/app/agent/core/agent_loop.py"
    if not os.path.exists(agent_loop):
        print(f"SKIP: {agent_loop} not found (backend)")
        return
    
    with open(agent_loop, "r") as f:
        c = f.read()
    
    if "_check_task_incomplete" not in c:
        func = '''

def _unfinished_plan_items(session: Session) -> list[dict[str, str]]:
    plan = getattr(session, "current_plan", None) or []
    return [item for item in plan if item.get("status") in ("pending", "in_progress")]

'''
        c = c.replace("class Handlers:", func + "\n\nclass Handlers:")
    
    if "task_incomplete" not in c:
        check = '''
            # === Auto-continue: if model stopped but plan incomplete ===
            if not llm_result.tool_calls_acc and llm_result.content:
                unfinished = _unfinished_plan_items(session)
                if unfinished:
                    await session.send_event(
                        Event(event_type="task_incomplete", data={"incomplete_plan": unfinished})
                    )
'''
        marker = "  # -- End of turn --"
        if marker in c:
            c = c.replace(marker, check + "\n" + marker)
    
    with open(agent_loop, "w") as f: f.write(c)
    print("OK: agent_loop.py")


if __name__ == "__main__":
    patch_events()
    patch_sse()
    patch_use_agent_chat()
    patch_chat_input()
    patch_session_chat()
    patch_agent_loop()
    print("\nDONE: V4 - everything applied")