patch_agent_loop.py

"""Patch agent_loop.py - detect task incomplete and send event."""

AGENT_LOOP = "/app/agent/core/agent_loop.py"

def patch():
    import os
    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_task_incomplete function before class Handlers
    check_func = '''

def _check_task_incomplete(session: Session, llm_result: LLMResult) -> bool:
    """Check if model stopped streaming but task is incomplete."""
    if llm_result.tool_calls_acc:
        return False  # Has tool calls, task continues via tool execution
    
    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


def _unfinished_plan_items(session: Session) -> list[dict[str, str]]:
    """Helper to get unfinished plan items."""
    plan = getattr(session, "current_plan", None) or []
    return [item for item in plan if item.get("status") in ("pending", "in_progress")]

'''
    
    if "_check_task_incomplete" not in content:
        content = content.replace(
            "class Handlers:",
            check_func + "\n\nclass Handlers:"
        )
    
    # Add the check after LLM response processing
    check_call = '''
            # === Check for incomplete task after LLM response ===
            if not llm_result.tool_calls_acc and llm_result.content:
                unfinished = _unfinished_plan_items(session)
                if unfinished:
                    await session.send_event(
                        Event(
                            event_type="task_incomplete",
                            data={
                                "incomplete_plan": unfinished,
                                "message": "Model phản hồi đã dừng nhưng nhiệm vụ chưa hoàn thành."
                            }
                        )
                    )
'''
    
    if "task_incomplete" not in content:
        # Find a good insertion point - after tool calls processing
        insert_marker = "  # -- End of turn --"
        if insert_marker in content:
            content = content.replace(insert_marker, check_call + "\n" + insert_marker)
        else:
            # Try another location
            result_marker = "final_response = llm_result.content or None"
            if result_marker in content:
                content = content.replace(result_marker, result_marker + "\n" + check_call)
    
    with open(AGENT_LOOP, "w", encoding="utf-8") as f:
        f.write(content)
    print(f"OK: Patched {AGENT_LOOP}")

if __name__ == "__main__":
    patch()