patch_auto_continue_v12.py

"""V12: Auto-continue feature - final fix.

Key fixes for V11 bugs:
1. Button import: check exact MUI import format
2. useAgentChat: match two-line `return {\n    messages: chat.messages,` instead of bare `return {`
3. Use \u unicode escapes for Vietnamese text
"""

import os

ROOT = "/source/frontend/src"
EVENTS = f"{ROOT}/types/events.ts"
SSE = f"{ROOT}/lib/sse-chat-transport.ts"
HOOK = f"{ROOT}/hooks/useAgentChat.ts"
INPUT = f"{ROOT}/components/Chat/ChatInput.tsx"
SESS = f"{ROOT}/components/SessionChat.tsx"
AGENT = "/app/agent/core/agent_loop.py"

def ok(f):
    print(f"OK: {f}")

# --- 1. events.ts: add task_incomplete ---
with open(EVENTS) as f:
    c = f.read()
if "'task_incomplete'" not in c:
    c = c.replace("  | 'interrupted'", "  | 'interrupted'\n  | 'task_incomplete'")
    with open(EVENTS, 'w') as f:
        f.write(c)
ok("events.ts")

# --- 2. sse-chat-transport.ts: add case handler ---
with open(SSE) as f:
    c = f.read()
if "case 'task_incomplete'" not in c:
    block = """      case 'task_incomplete':
          (sideChannel as any).onTaskIncomplete(
            (event.data?.incomplete_plan as Array<{ id: string; content: string; status: string }>) || [],
          );
          break;

        default:"""
    c = c.replace("\n        default:", "\n" + block)
    with open(SSE, 'w') as f:
        f.write(c)
ok("sse.ts")

# --- 3. useAgentChat.ts: add state, handler, effect, return ---
with open(HOOK) as f:
    content = f.read()

# Add useState to import
if "useState" not in content:
    content = content.replace(
        "import { useCallback, useEffect, useMemo, useRef } from 'react';",
        "import { useCallback, useEffect, useMemo, useRef, useState } from 'react';"
    )

# Add state after callbacksRef line
state_block = """\
  // Auto-continue state
  const [_showAc, _setShowAc] = useState(false);
  const _acRef = useRef<Array<{ id: string; content: string; status: string }>>([]);
"""
if "_showAc" not in content:
    content = content.replace(
        "callbacksRef.current = { onReady, onError, onSessionDead };",
        "callbacksRef.current = { onReady, onError, onSessionDead };\n" + state_block
    )

# Add handler inside onSessionUpdate
handler_code = """\
        // Auto-continue check
        if ((data as any).ac_plan) {
          _acRef.current = (data as any).ac_plan;
          _setShowAc(true);
          return;
        }
"""
if "ac_plan" not in content:
    content = content.replace(
        "onSessionUpdate: (data) => {",
        "onSessionUpdate: (data) => {\n" + handler_code
    )

# Add cancel + effect BEFORE the two-line return marker
cancel_effect = """\
  const _acCancel = useCallback(() => { _setShowAc(false); }, []);
  useEffect(() => {
    if (!_showAc) return;
    const plan = _acRef.current;
    const timer = setTimeout(() => {
      const s = chatActionsRef.current.setMessages;
      const m = chatActionsRef.current.messages;
      if (s && plan.length > 0) {
        const t = plan.map((i: { content: string }) => `- ${i.content}`).join('\\n');
        const nu = { id: 'ac-'+Date.now(), role: 'user' as const, parts: [{ type: 'text' as const, text: `[TIEP TUC] Task chua hoan thanh:\\n${t}` }], content: '' };
        s([...m, nu]);
      }
      _setShowAc(false);
    }, 10000);
    return () => clearTimeout(timer);
  }, [_showAc]);

"""
if "_acCancel" not in content:
    # Match the EXACT two-line return pattern
    content = content.replace(
        "  return {\n    messages: chat.messages,",
        cancel_effect + "  return {\n    messages: chat.messages,"
    )

# Add return values
if "_showAc," not in content:
    content = content.replace(
        "refreshMessages,\n  };",
        "refreshMessages,\n    _showAc,\n    _acCancel,\n  };"
    )

with open(HOOK, 'w') as f:
    f.write(content)
ok("useAgentChat.ts")

# --- 4. ChatInput.tsx: add Button + props + pause button ---
with open(INPUT) as f:
    content = f.read()

# Add Button to MUI imports - match the actual format
if "Button" not in content:
    content = content.replace(
        "  Tooltip,",
        "  Tooltip,\n  Button,"
    )

# Add props to interface
if "_showAc" not in content:
    content = content.replace(
        "placeholder?: string;",
        "placeholder?: string;\n  _showAc?: boolean;\n  _acCancel?: () => void;"
    )

# Add to destructured function params
if "_showAc" not in content or "_acCancel" not in content:
    content = content.replace(
        "placeholder = 'Ask anything...' }: ChatInputProps) {",
        "placeholder = 'Ask anything...', _showAc = false, _acCancel }: ChatInputProps) {"
    )

# Add pause button before JobsUpgradeDialog
if "Tam dung" not in content:
    btn = """
        {_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>
        )}
"""
    content = content.replace("<JobsUpgradeDialog", btn + "\n        <JobsUpgradeDialog")

with open(INPUT, 'w') as f:
    f.write(content)
ok("ChatInput.tsx")

# --- 5. SessionChat.tsx: pass props ---
if os.path.exists(SESS):
    with open(SESS) 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)
    ok("SessionChat.tsx")

# --- 6. agent_loop.py: detection ---
if os.path.exists(AGENT):
    with open(AGENT) as f:
        c = f.read()
    
    if "_unfinished_plan" not in c:
        fn = """

def _unfinished_plan(s):
    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:")
    
    if "ac_plan" not in c:
        check = """
            if not llm_result.tool_calls_acc and llm_result.content:
                unfinished = _unfinished_plan(session)
                if unfinished:
                    await session.send_event(Event(event_type="session_update", data={
                        "ac_plan": unfinished,
                    }))
"""
        marker = "  # -- End of turn --"
        if marker in c:
            c = c.replace(marker, check + marker)
    
    with open(AGENT, 'w') as f:
        f.write(c)
    ok("agent_loop.py")
else:
    print("SKIP: agent_loop.py (backend)")

print("\nDONE: V12")