patch_auto_continue_v7.py

"""V7: PRECISE patch - exact string matching with source code.

Adds auto-continue pause button for free models. All patches use exact 
string content from the original source files for reliable replacement.
"""

import os, sys

# ========== PATCH 1: events.ts - add 'task_incomplete' to EventType ==========
events_ts = "/source/frontend/src/types/events.ts"
def patch_events():
    with open(events_ts) as f:
        c = f.read()
    if "'task_incomplete'" not in c:
        c = c.replace(
            "  | 'interrupted'",
            "  | 'interrupted'\n  | 'task_incomplete'"
        )
        with open(events_ts, 'w') as f: f.write(c)
    print("OK: events.ts")

# ========== PATCH 2: SSE transport - add onTaskIncomplete to interface + handler ==========
sse_ts = "/source/frontend/src/lib/sse-chat-transport.ts"
def patch_sse():
    with open(sse_ts) as f:
        c = f.read()
    
    # Add 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 before default
    if "case 'task_incomplete'" not in c:
        case_handler = """      case 'task_incomplete':
          sideChannel.onTaskIncomplete(
            (event.data?.incomplete_plan as Array<{ id: string; content: string; status: string }>) || [],
          );
          break;

        default:"""
        c = c.replace("\n        default:", "\n" + case_handler)
    
    with open(sse_ts, 'w') as f: f.write(c)
    print("OK: sse-chat-transport.ts")

# ========== PATCH 3: useAgentChat.ts - add state + handler + return ==========
hook_ts = "/source/frontend/src/hooks/useAgentChat.ts"
def patch_hook():
    with open(hook_ts) as f:
        c = f.read()
    
    # 1. Add useState to imports
    if "useState" not in c:
        c = c.replace(
            "import { useCallback, useEffect, useMemo, useRef } from 'react';",
            "import { useCallback, useEffect, useMemo, useRef, useState } from 'react';"
        )
    
    # 2. Add state variables after callbacksRef
    state_block = """
  // Auto-continue state for free models
  const _acTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null);
  const [_showAc, _setShowAc] = useState(false);"""
    if "_acTimerRef" not in c:
        c = c.replace(
            "callbacksRef.current = { onReady, onError, onSessionDead };",
            "callbacksRef.current = { onReady, onError, onSessionDead };" + state_block
        )
    
    # 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 in sideChannel object
    # This inserts BEFORE the existing onInterrupted handler
    handler = """      onTaskIncomplete: (incompletePlan: Array<{ id: string; content: string; status: string }>) => {
        _setShowAc(true);
        if (_acTimerRef.current) clearTimeout(_acTimerRef.current);
        _acTimerRef.current = setTimeout(() => {
          _setShowAc(false);
          const planStr = incompletePlan.map(i => `- ${i.content}`).join('\\\\n');
          const setMsgs = chatActionsRef.current.setMessages;
          if (setMsgs) {
            const msgs = chatActionsRef.current.messages;
            setMsgs([...msgs, { id: 'ac-'+Date.now(), role: 'user' as const, parts: [{ type: 'text' as const, text: `[TIẾP TỤC] Task chưa hoàn thành:\\\\n${planStr}` }], content: '' }]);
          }
        }, 10000);
      },
      onInterrupted: () => { /* no-op */ },"""
    
    if "onTaskIncomplete: (incompletePlan" not in c:
        old_interrupted = '      onInterrupted: () => { /* no-op \u2014 handled by stop() caller */ },'
        if old_interrupted in c:
            c = c.replace(old_interrupted, handler)
        else:
            # Try simpler version
            old2 = '      onInterrupted: () => { /* no-op */ },'
            c = c.replace(old2, handler)
    
    # 5. Add _acCancel function + return values
    cancel_func = """
  const _acCancel = useCallback(() => {
    _setShowAc(false);
    if (_acTimerRef.current) { clearTimeout(_acTimerRef.current); _acTimerRef.current = null; }
  }, []);"""
    if "_acCancel" not in c:
        c = c.replace("\n  return {\n    messages: chat.messages,", cancel_func + "\n  return {\n    messages: chat.messages,")
    
    # 6. Return values
    if "_showAc," not in c:
        c = c.replace(
            "refreshMessages,\n  };",
            "refreshMessages,\n    _showAc,\n    _acCancel,\n  };"
        )
    
    with open(hook_ts, 'w') as f: f.write(c)
    print("OK: useAgentChat.ts")

# ========== PATCH 4: ChatInput.tsx - add pause button ==========
input_tsx = "/source/frontend/src/components/Chat/ChatInput.tsx"
def patch_chat_input():
    with open(input_tsx) as f:
        c = f.read()
    
    # 1. Add Button import (it's missing from the original imports)
    if "Button" not in c:
        c = c.replace(
            "  Tooltip,",
            "  Tooltip,\n  Button,"
        )
    
    # 2. Add props to interface
    if "_showAc" not in c:
        old_interface_end = """interface ChatInputProps {
  sessionId?: string;
  initialModelPath?: string | null;
  onSend: (text: string) => void;
  onStop?: () => void;
  onDatasetUploaded?: () => Promise<boolean> | boolean;
  isProcessing?: boolean;
  disabled?: boolean;
  placeholder?: string;
}"""
        new_interface = """interface ChatInputProps {
  sessionId?: string;
  initialModelPath?: string | null;
  onSend: (text: string) => void;
  onStop?: () => void;
  onDatasetUploaded?: () => Promise<boolean> | boolean;
  isProcessing?: boolean;
  disabled?: boolean;
  placeholder?: string;
  _showAc?: boolean;
  _acCancel?: () => void;
}"""
        c = c.replace(old_interface_end, new_interface)
    
    # 3. Add destructured props in function signature
    if "_showAc" not in c:
        old_props = "placeholder = 'Ask anything...' }: ChatInputProps) {"
        new_props = "placeholder = 'Ask anything...', _showAc = false, _acCancel }: ChatInputProps) {"
        if old_props not in c:
            old_props2 = "placeholder = 'Ask anything...' }: ChatInputProps) {"
            new_props2 = "placeholder = 'Ask anything...', _showAc = false, _acCancel }: ChatInputProps) {"
            c = c.replace(old_props2, new_props2)
        else:
            c = c.replace(old_props, new_props)
    
    # 4. Add pause button after the model badge area, before <JobsUpgradeDialog>
    if "T\u1ea1m d\u1eebng" not in c:
        btn = """
        {/* Auto-continue pause for free models */}
        {_showAc && _acCancel && (
          <Box sx={{ display: 'flex', justifyContent: 'center', mt: 1 }}>
            <Button
              variant="outlined"
              size="small"
              onClick={_acCancel}
              sx={{ textTransform: 'none', fontSize: '0.7rem', opacity: 0.7, '&:hover': { opacity: 1 } }}
            >
              T\u1ea1m d\u1eebng t\u1ef1 \u0111\u1ed9ng ti\u1ebfp t\u1ee5c (10s)
            </Button>
          </Box>
        )}
"""
        if "<JobsUpgradeDialog" in c:
            c = c.replace("<JobsUpgradeDialog", btn + "\n        <JobsUpgradeDialog")
    
    with open(input_tsx, 'w') as f: f.write(c)
    print("OK: ChatInput.tsx")

# ========== PATCH 5: SessionChat.tsx - pass props down ==========
session_tsx = "/source/frontend/src/components/SessionChat.tsx"
def patch_session():
    if not os.path.exists(session_tsx):
        print(f"SKIP: {session_tsx}")
        return
    with open(session_tsx) as f:
        c = f.read()
    if "_showAc" not in c:
        old = "<ChatInput "
        new = "<ChatInput _showAc={_showAc} _acCancel={_acCancel} "
        c = c.replace(old, new)
    with open(session_tsx, 'w') as f: f.write(c)
    print("OK: SessionChat.tsx")

# ========== PATCH 6: agent_loop.py - detection ==========
agent_py = "/app/agent/core/agent_loop.py"
def patch_agent():
    if not os.path.exists(agent_py):
        print(f"SKIP: {agent_py}")
        return
    with open(agent_py) as f:
        c = f.read()
    
    # Add helper function
    if "_unfinished_plan" not in c:
        fn = """
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")]

"""
        c = c.replace("class Handlers:", fn + "\n\nclass Handlers:")
    
    # Add detection
    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_py, 'w') as f: f.write(c)
    print("OK: agent_loop.py")

if __name__ == "__main__":
    patch_events()
    patch_sse()
    patch_hook()
    patch_chat_input()
    patch_session()
    patch_agent()
    print("\nDONE: V7 - precise patches applied")