File size: 5,086 Bytes
4975526
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
"""Patch script: Add auto-continue feature for free models in ml-intern.

Apply this script during Docker build to enable:
1. Detection when streaming ends but task incomplete
2. "Tự động tiếp tục (10s)" button - auto continues after 10s
3. "Tạm dừng" button - lets user send different content

Usage: Run this during Docker build after cloning source.
"""

import os
import re

# Paths - adjust if needed
USE_AGENT_CHAT = "/source/frontend/src/hooks/useAgentChat.ts"
SSE_TRANSPORT = "/source/frontend/src/lib/sse-chat-transport.ts"
AGENT_LOOP = "/app/agent/core/agent_loop.py"

# =============================================================================
# Patch 1: useAgentChat.ts
# =============================================================================
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", encoding="utf-8") as f:
        content = f.read()
    
    # Add useState to imports if missing
    if "import { useCallback, useEffect, useMemo, useRef, useState } from 'react';" not in content:
        content = content.replace(
            "import { useCallback, useEffect, useMemo, useRef } from 'react';",
            "import { useCallback, useEffect, useMemo, useRef, useState } from 'react';"
        )
    
    # Add auto-continue state after callbacksRef
    state_code = '''
  // 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 and "callbacksRef.current = { onReady, onError, onSessionDead }" in content:
        content = content.replace(
            "callbacksRef.current = { onReady, onError, onSessionDead };",
            "callbacksRef.current = { onReady, onError, onSessionDead };" + state_code
        )
    
    # Add onTaskIncomplete to SideChannelCallbacks interface
    if "onTaskIncomplete" not in content:
        content = content.replace(
            "onInterrupted: () => void;",
            "onInterrupted: () => void;\n  onTaskIncomplete: (incompletePlan: Array<{ id: string; content: string; status: string }>) => void;"
        )
    
    with open(USE_AGENT_CHAT, "w", encoding="utf-8") as f:
        f.write(content)
    print(f"OK: Patched {USE_AGENT_CHAT}")

# =============================================================================
# Patch 2: sse-chat-transport.ts  
# =============================================================================
def patch_sse_transport():
    if not os.path.exists(SSE_TRANSPORT):
        print(f"SKIP: {SSE_TRANSPORT} not found")
        return
    
    with open(SSE_TRANSPORT, "r", encoding="utf-8") as f:
        content = f.read()
    
    # Add task_incomplete case
    if "case 'task_incomplete'" not in content:
        task_incomplete_case = '''      case 'task_incomplete':
          sideChannel.onTaskIncomplete(
            (event.data?.incomplete_plan as Array<{ id: string; content: string; status: string }>) || [],
          );
          break;
'''
        # Insert before turn_complete case
        content = content.replace(
            "case 'turn_complete':",
            task_incomplete_case + "\n        case 'turn_complete':"
        )
    
    with open(SSE_TRANSPORT, "w", encoding="utf-8") as f:
        f.write(content)
    print(f"OK: Patched {SSE_TRANSPORT}")

# =============================================================================
# Patch 3: agent_loop.py
# =============================================================================
def patch_agent_loop():
    if not os.path.exists(AGENT_LOOP):
        print(f"SKIP: {AGENT_LOOP} not found")
        return
    
    with open(AGENT_LOOP, "r", encoding="utf-8") as f:
        content = f.read()
    
    # Add check function if not exists
    check_func = '''
def _check_task_incomplete(session: Session, llm_result: LLMResult) -> bool:
    """Check if model stopped but task is incomplete (has unfinished plan items)."""
    if llm_result.tool_calls_acc:
        return False  # Has tool calls, not incomplete
    plan = getattr(session, "current_plan", None) or []
    unfinished = [item for item in plan if item.get("status") in ("pending", "in_progress")]
    return len(unfinished) > 0
'''
    if "_check_task_incomplete" not in content:
        # Insert after LLMResult dataclass
        content = content.replace(
            "class Handlers:",
            check_func + "\n\nclass Handlers:"
        )
    
    with open(AGENT_LOOP, "w", encoding="utf-8") as f:
        f.write(content)
    print(f"OK: Patched {AGENT_LOOP}")


if __name__ == "__main__":
    patch_use_agent_chat()
    patch_sse_transport()
    patch_agent_loop()
    print("\nDONE: Auto-continue patches applied")