patch_auto_continue_v5.py

"""V5: SIMPLEST approach - just show pause button with 10s auto-continue.

Architecture:
1. Agent loop detects model stopped mid-task → sends 'task_incomplete' event
2. SSE transport receives it → calls sideChannel.onTaskIncomplete()
3. useAgentChat: onTaskIncomplete sets _showAutoContinue=true, sends continue msg
4. ChatInput shows "Tạm dừng" button if _showAutoContinue
"""

import os, re

EVENTS = "/source/frontend/src/types/events.ts"
SSE = "/source/frontend/src/lib/sse-chat-transport.ts"
USE_AGENT_CHAT = "/source/frontend/src/hooks/useAgentChat.ts"
CHAT_INPUT = "/source/frontend/src/components/Chat/ChatInput.tsx"
AGENT_LOOP = "/app/agent/core/agent_loop.py"


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()
    
    if "onTaskIncomplete:" not in c:
        c = c.replace(
            "  onInterrupted: () => void;",
            "  onInterrupted: () => void;\n  onTaskIncomplete: (incompletePlan: Array<{ id: string; content: string; status: string }>) => void;"
        )
    
    if "case 'task_incomplete'" not in c:
        case_block = """      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_block + "\n        case 'turn_complete':")
    
    with open(SSE, "w") as f:
        f.write(c)
    print("OK: sse.ts")


def patch_use_agent():
    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 (prefixed with _ = TS allows unused)
    state = '''
  // Auto-continue state for free models
  const _acTimer = useRef<ReturnType<typeof setTimeout> | null>(null);
  const [_showAc, _setShowAc] = useState(false);'''
    if "_acTimer" not in c:
        c = c.replace(
            "callbacksRef.current = { onReady, onError, onSessionDead };",
            "callbacksRef.current = { onReady, onError, onSessionDead };" + state
        )
    
    # 3. 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;"
        )
    
    # 4. Add onTaskIncomplete handler - uses chatActionsRef ref instead of chat directly
    handler = '''      onTaskIncomplete: (incompletePlan: Array<{ id: string; content: string; status: string }>) => {
        _setShowAc(true);
        if (_acTimer.current) clearTimeout(_acTimer.current);
        _acTimer.current = setTimeout(() => {
          _setShowAc(false);
          // Auto-continue after 10s via chatActionsRef
          const chatRef = chatActionsRef.current;
          const sendMsg = (chatRef as any)?.sendMessage;
          if (sendMsg) {
            const planStr = incompletePlan.map(i => `- ${i.content}`).join('\\\\n');
            sendMsg({ text: `[TIẾP TỤC] Tiếp tục task chưa hoàn thành:\\\\n${planStr}`, 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)
        else:
            old2 = '      onInterrupted: () => { /* no-op */ },'
            c = c.replace(old2, handler)
    
    # 5. Cancel ref + return values
    if "_showAc" not in content:
        # Insert cancel before return
        cancel = '''
  // Expose auto-continue state
  const _acCancel = useCallback(() => {
    _setShowAc(false);
    if (_acTimer.current) { clearTimeout(_acTimer.current); _acTimer.current = null; }
  }, []);
'''
        if "_acCancel" not in c:
            c = c.replace('\n  return {\n    messages: chat.messages,', cancel + '\n  return {\n    messages: chat.messages,')
        
        if "_showAc," not in c:
            c = c.replace(
                "refreshMessages,\n  };",
                "refreshMessages,\n    _showAc,\n    _acCancel,\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()
    
    # DatasetUploadResponse fix
    if "DatasetUploadResponse" in c and "type DatasetUploadResponse" 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
    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';"
        )
    
    # Update interface
    if "_showAc" not in c:
        m = re.search(r'interface ChatInputProps\s*\{[^}]*\}', c, re.DOTALL)
        if m:
            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;
}"""
            c = c.replace(m.group(0), iface)
    
    # Add button
    btn = """
      {/* Auto-continue pause button 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' }}
          >
            Tạm dừng tự động tiếp tục (10s)
          </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():
    sess = "/source/frontend/src/components/SessionChat.tsx"
    if not os.path.exists(sess):
        print(f"SKIP: {sess}")
        return
    with open(sess, "r") as f:
        c = f.read()
    if "_showAc" not in c:
        c = c.replace("<ChatInput ", "<ChatInput _showAc={_showAc} _acCancel={_acCancel} ")
    with open(sess, "w") as f:
        f.write(c)
    print("OK: SessionChat.tsx")


def patch_agent_loop():
    if not os.path.exists(AGENT_LOOP):
        print(f"SKIP: {AGENT_LOOP}")
        return
    with open(AGENT_LOOP, "r") as f:
        c = f.read()
    
    if "_unfinished_plan" not in c:
        fn = '''

def _unfinished_plan(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:", fn + "\n\nclass Handlers:")
    
    if "task_incomplete" not in c:
        check = '''
            # === Auto-continue detection ===
            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 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()
    patch_chat_input()
    patch_session_chat()
    patch_agent_loop()
    print("\nDONE: V5 - simple approach applied")