-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy path176.py
More file actions
189 lines (158 loc) · 5.96 KB
/
Copy path176.py
File metadata and controls
189 lines (158 loc) · 5.96 KB
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
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
import asyncio
import logging
from contextlib import asynccontextmanager
import uvicorn
from autogpt_libs.auth import parse_jwt_token
from fastapi import Depends, FastAPI, WebSocket, WebSocketDisconnect
from starlette.middleware.cors import CORSMiddleware
from backend.data import redis
from backend.data.execution import AsyncRedisExecutionEventBus
from backend.data.user import DEFAULT_USER_ID
from backend.server.conn_manager import ConnectionManager
from backend.server.model import ExecutionSubscription, Methods, WsMessage
from backend.util.service import AppProcess
from backend.util.settings import AppEnvironment, Config, Settings
logger = logging.getLogger(__name__)
settings = Settings()
@asynccontextmanager
async def lifespan(app: FastAPI):
manager = get_connection_manager()
fut = asyncio.create_task(event_broadcaster(manager))
fut.add_done_callback(lambda _: logger.info("Event broadcaster stopped"))
yield
docs_url = "/docs" if settings.config.app_env == AppEnvironment.LOCAL else None
app = FastAPI(lifespan=lifespan, docs_url=docs_url)
_connection_manager = None
def get_connection_manager():
global _connection_manager
if _connection_manager is None:
_connection_manager = ConnectionManager()
return _connection_manager
async def event_broadcaster(manager: ConnectionManager):
try:
redis.connect()
event_queue = AsyncRedisExecutionEventBus()
async for event in event_queue.listen():
await manager.send_execution_result(event)
except Exception as e:
logger.exception(f"Event broadcaster error: {e}")
raise
finally:
redis.disconnect()
async def authenticate_websocket(websocket: WebSocket) -> str:
if not settings.config.enable_auth:
return DEFAULT_USER_ID
token = websocket.query_params.get("token")
if not token:
await websocket.close(code=4001, reason="Missing authentication token")
return ""
try:
payload = parse_jwt_token(token)
user_id = payload.get("sub")
if not user_id:
await websocket.close(code=4002, reason="Invalid token")
return ""
return user_id
except ValueError:
await websocket.close(code=4003, reason="Invalid token")
return ""
async def handle_subscribe(
websocket: WebSocket, manager: ConnectionManager, message: WsMessage
):
if not message.data:
await websocket.send_text(
WsMessage(
method=Methods.ERROR,
success=False,
error="Subscription data missing",
).model_dump_json()
)
else:
ex_sub = ExecutionSubscription.model_validate(message.data)
await manager.subscribe(ex_sub.graph_id, websocket)
logger.debug(f"New execution subscription for graph {ex_sub.graph_id}")
await websocket.send_text(
WsMessage(
method=Methods.SUBSCRIBE,
success=True,
channel=ex_sub.graph_id,
).model_dump_json()
)
async def handle_unsubscribe(
websocket: WebSocket, manager: ConnectionManager, message: WsMessage
):
if not message.data:
await websocket.send_text(
WsMessage(
method=Methods.ERROR,
success=False,
error="Subscription data missing",
).model_dump_json()
)
else:
ex_sub = ExecutionSubscription.model_validate(message.data)
await manager.unsubscribe(ex_sub.graph_id, websocket)
logger.debug(f"Removed execution subscription for graph {ex_sub.graph_id}")
await websocket.send_text(
WsMessage(
method=Methods.UNSUBSCRIBE,
success=True,
channel=ex_sub.graph_id,
).model_dump_json()
)
@app.get("/")
async def health():
return {"status": "healthy"}
@app.websocket("/ws")
async def websocket_router(
websocket: WebSocket, manager: ConnectionManager = Depends(get_connection_manager)
):
user_id = await authenticate_websocket(websocket)
if not user_id:
return
await manager.connect(websocket)
try:
while True:
data = await websocket.receive_text()
message = WsMessage.model_validate_json(data)
if message.method == Methods.HEARTBEAT:
await websocket.send_json(
{"method": Methods.HEARTBEAT.value, "data": "pong", "success": True}
)
continue
if message.method == Methods.SUBSCRIBE:
await handle_subscribe(websocket, manager, message)
elif message.method == Methods.UNSUBSCRIBE:
await handle_unsubscribe(websocket, manager, message)
elif message.method == Methods.ERROR:
logger.error(f"WebSocket Error message received: {message.data}")
else:
logger.warning(
f"Unknown WebSocket message type {message.method} received: "
f"{message.data}"
)
await websocket.send_text(
WsMessage(
method=Methods.ERROR,
success=False,
error="Message type is not processed by the server",
).model_dump_json()
)
except WebSocketDisconnect:
manager.disconnect(websocket)
logger.debug("WebSocket client disconnected")
class WebsocketServer(AppProcess):
def run(self):
logger.info(f"CORS allow origins: {settings.config.backend_cors_allow_origins}")
server_app = CORSMiddleware(
app=app,
allow_origins=settings.config.backend_cors_allow_origins,
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
uvicorn.run(
server_app,
host=Config().websocket_server_host,
port=Config().websocket_server_port,
)