AutoBrowser/server.py at main · exponentialXP/AutoBrowser · GitHub

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
import asyncio
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
import json
from agent import WebAgent

app = FastAPI()

# Allow all CORS for local development
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# Global agent instance
agent = None
# Active WebSocket connections
active_connections = []

@app.on_event("startup")
async def startup_event():
    global agent
    agent = WebAgent(logger=broadcast_log)
    # Start the browser immediately
    asyncio.create_task(agent.start_browser())

async def broadcast_log(message, type="info"):
    dead_connections = []
    payload = json.dumps({"type": "log", "message": message, "logType": type})
    for websocket in active_connections:
        try:
            await websocket.send_text(payload)
        except:
            dead_connections.append(websocket)

    for dead in dead_connections:
        if dead in active_connections:
            active_connections.remove(dead)

async def broadcast_status(status):
    """Broadcast agent status (idle/running) to all clients."""
    dead_connections = []
    payload = json.dumps({"type": "status", "status": status})
    for websocket in active_connections:
        try:
            await websocket.send_text(payload)
        except:
            dead_connections.append(websocket)

    for dead in dead_connections:
        if dead in active_connections:
            active_connections.remove(dead)

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    active_connections.append(websocket)
    try:
        while True:
            data = await websocket.receive_text()
            # Handle incoming messages if needed
            message = json.loads(data)
            if message.get("type") == "start_task":
                task = message.get("task")
                api_key = message.get("api_key")
                asyncio.create_task(run_agent_task(task, api_key))
            elif message.get("type") == "pause":
                if agent:
                    agent.paused = True
                    await broadcast_log("Agent paused.", "warning")
            elif message.get("type") == "resume":
                if agent:
                    agent.paused = False
                    await broadcast_log("Agent resumed.", "info")
            elif message.get("type") == "stop":
                if agent:
                    agent.stopped = True
                    await broadcast_log("Agent stopping...", "warning")
            elif message.get("type") == "reset":
                if agent:
                    agent.stopped = True
                    agent.history = []
                    await broadcast_log("Agent reset. History cleared.", "warning")
    except WebSocketDisconnect:
        if websocket in active_connections:
            active_connections.remove(websocket)

async def run_agent_task(task, api_key=None):
    global agent
    await broadcast_status("running")
    await broadcast_log(f"Starting task: {task}", "info")
    try:
        if not agent:
            agent = WebAgent(logger=broadcast_log)
        await agent.run(task, api_key=api_key)
    except Exception as e:
        await broadcast_log(f"Agent error: {str(e)}", "error")
    finally:
        await broadcast_status("idle")

@app.get("/health")
async def health():
    return {"status": "ok"}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="127.0.0.1", port=8000)