patch_auto_continue_v6.py

"""V6: FIXED - Auto-continue with pause button for free models."""

import os
import re

EVENTS_FILE = "/source/frontend/src/types/events.ts"
SSE_TRANSPORT = "/source/frontend/src/lib/sse-chat-transport.ts"
HOOK_FILE = "/source/frontend/src/hooks/useAgentChat.ts"
CHAT_INPUT = "/source/frontend/src/components/Chat/ChatInput.tsx"
SESSION_CHAT = "/source/frontend/src/components/SessionChat.tsx"
AGENT_LOOP = "/app/agent/core/agent_loop.py"


def fix_events():
    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: events.ts")


def fix_sse():
    with open(SSE_TRANSPORT, "r") as f:
        content = f.read()

    # 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;"
        )

    # Add case handler  
    if "case 'task_incomplete'" not in content:
        case_code = """      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':", case_code + "\n        case 'turn_complete':")

    with open(SSE_TRANSPORT, "w") as f:
        f.write(content)
    print("OK: sse-chat-transport.ts")


def fix_hook():
    with open(HOOK_FILE, "r") as f:
        content = f.read()

    # 1. Import useState
    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 block
    state_block = """
  // Auto-continue for free models
  const _acTimer = useRef<ReturnType<typeof setTimeout> | null>(null);
  const [_showAc, _setShowAc] = useState(false);
  const _acCancel = useCallback(() => {
    _setShowAc(false);
    if (_acTimer.current) { clearTimeout(_acTimer.current); _acTimer.current = null; }
  }, []);"""
    
    if "_acTimer" 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 handler in sideChannel - uses chatActionsRef to avoid hoisting issues
    handler_code = """      onTaskIncomplete: (incompletePlan: Array<{ id: string; content: string; status: string }>) => {
        _setShowAc(true);
        if (_acTimer.current) clearTimeout(_acTimer.current);
        _acTimer.current = setTimeout(() => {
          _setShowAc(false);
          const planStr = incompletePlan.map(i => `- ${i.content}`).join('\\\\n');
          const msg = { text: `[TIẾP TỤC] Task chưa hoàn thành:\\\\n${planStr}`, metadata: { createdAt: new Date().toISOString() } };
          const setMsgs = chatActionsRef.current.setMessages;
          if (setMsgs) {
            chatActionsRef.current.messages = [...chatActionsRef.current.messages, { id: 'ac-continue', role: 'user', parts: [{ type: 'text', text: msg.text }], content: msg.text }];
          }
        }, 10000);
      },
      onInterrupted: () => { /* no-op - handled by stop() */ },"""
    
    if "onTaskIncomplete: (incompletePlan" not in content:
        old_handler = '      onInterrupted: () => { /* no-op \u2014 handled by stop() caller */ },'
        if old_handler in content:
            content = content.replace(old_handler, handler_code)
        else:
            old_handler2 = '      onInterrupted: () => { /* no-op */ },'
            content = content.replace(old_handler2, handler_code)

    # 5. Return values
    if "_showAc" not in content:
        # The _showAc and _acCancel need to be in the return
        return_marker = "refreshMessages,\n  };"
        if return_marker in content:
            content = content.replace(
                return_marker,
                "refreshMessages,\n    _showAc,\n    _acCancel,\n  };"
            )

    with open(HOOK_FILE, "w") as f:
        f.write(content)
    print("OK: useAgentChat.ts")


def fix_chat_input():
    with open(CHAT_INPUT, "r") as f:
        content = f.read()

    # Fix DatasetUploadResponse
    wrong_import = "import type { DatasetUploadResponse } from '@/types/agent';"
    if wrong_import in content:
        content = content.replace(wrong_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
    if "_showAc" not in content:
        iface_match = re.search(r'interface ChatInputProps\s*\{[^}]*\}', content, re.DOTALL)
        if iface_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;
  _showAc?: boolean;
  _acCancel?: () => void;
}"""
            content = content.replace(iface_match.group(0), new_iface)

    # Add pause button
    btn_code = """
      {/* Auto-continue pause for free models */}
      {_showAc && _acCancel && (
        <Stack direction="row" spacing={1} sx={{ mb: 1, px: 1 }}>
          <Button
            variant="outlined"
            color="secondary"
            size="small"
            onClick={_acCancel}
            sx={{ textTransform: 'none', fontSize: '0.7rem', opacity: 0.8 }}
          >
            T\u1ea1m d\u1eebng t\u1ef1 \u0111\u1ed9ng ti\u1ebfp t\u1ee5c (10s)
          </Button>
        </Stack>
      )}"""
    
    if "T\u1ea1m d\u1eebng" not in content and "free models" not in content:
        content = content.replace("<Box sx={{ flex: 1 ", btn_code + "\n        <Box sx={{ flex: 1 ")

    with open(CHAT_INPUT, "w") as f:
        f.write(content)
    print("OK: ChatInput.tsx")


def fix_session_chat():
    with open(SESSION_CHAT, "r") as f:
        content = f.read()
    
    # Add props to ChatInput
    if "_showAc" not in content:
        content = content.replace(
            "<ChatInput ",
            "<ChatInput _showAc={_showAc} _acCancel={_acCancel} "
        )

    with open(SESSION_CHAT, "w") as f:
        f.write(content)
    print("OK: SessionChat.tsx")


def fix_agent_loop():
    if not os.path.exists(AGENT_LOOP):
        print(f"SKIP: {AGENT_LOOP}")
        return
    with open(AGENT_LOOP, "r") as f:
        content = f.read()

    if "_unfinished_plan" not in content:
        fn_block = '''

def _unfinished_plan(s: Session) -> list[dict[str, str]]:
    p = getattr(s, "current_plan", None) or []
    return [it for it in p if it.get("status") in ("pending", "in_progress")]

'''
        content = content.replace("class Handlers:", fn_block + "\n\nclass Handlers:")

    if "task_incomplete" not in content:
        check_block = '''
            # Auto-continue detect
            if not llm_result.tool_calls_acc and llm_result.content:
                unfinished = _unfinished_plan(session)
                if unfinished:
                    await session.send_event(Event(event_type="task_incomplete", data={"incomplete_plan": unfinished}))
'''
        marker = "  # -- End of turn --"
        if marker in content:
            content = content.replace(marker, check_block + "\n" + marker)

    with open(AGENT_LOOP, "w") as f:
        f.write(content)
    print("OK: agent_loop.py")


if __name__ == "__main__":
    fix_events()
    fix_sse()
    fix_hook()
    fix_chat_input()
    fix_session_chat()
    fix_agent_loop()
    print("\nDONE: V6 - all fixed")