WIN-JHFT4D3SIVT\caoxiaozhu
5667190abe
- Phase 1: Infrastructure (state, prompts, registry) - Phase 2: Execution engine (AI adapters, security classifier, executors) - Phase 3: Agent integration (graph nodes, routing) - Phase 4: Streaming interaction (PTY terminal, WebSocket) - Phase 5: Frontend integration (Vue components)
80 lines
2.6 KiB
Python
80 lines
2.6 KiB
Python
"""
|
|
Terminal WebSocket Router - 终端 WebSocket 端点
|
|
"""
|
|
|
|
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
|
|
|
|
from app.agents.tools.terminal_engine import pty_manager
|
|
|
|
router = APIRouter(prefix="/ws/terminal", tags=["terminal"])
|
|
|
|
|
|
class ConnectionManager:
|
|
"""WebSocket 连接管理器"""
|
|
|
|
def __init__(self):
|
|
self.active_connections: dict[str, WebSocket] = {}
|
|
|
|
async def connect(self, session_id: str, websocket: WebSocket):
|
|
await websocket.accept()
|
|
self.active_connections[session_id] = websocket
|
|
|
|
def disconnect(self, session_id: str):
|
|
if session_id in self.active_connections:
|
|
del self.active_connections[session_id]
|
|
|
|
async def send(self, session_id: str, data: str):
|
|
if session_id in self.active_connections:
|
|
await self.active_connections[session_id].send_text(data)
|
|
|
|
def is_connected(self, session_id: str) -> bool:
|
|
return session_id in self.active_connections
|
|
|
|
|
|
manager = ConnectionManager()
|
|
|
|
|
|
@router.websocket("/{session_id}")
|
|
async def terminal_websocket(websocket: WebSocket, session_id: str):
|
|
"""终端 WebSocket 端点"""
|
|
await manager.connect(session_id, websocket)
|
|
|
|
try:
|
|
# 获取该 session 的输出队列
|
|
queue = pty_manager._output_queues.get(session_id)
|
|
if queue:
|
|
# 异步任务：转发 PTY 输出到 WebSocket
|
|
async def forward_output():
|
|
while True:
|
|
try:
|
|
data = await asyncio.wait_for(queue.get(), timeout=0.1)
|
|
if data is None:
|
|
await manager.send(session_id, "[SESSION_END]")
|
|
break
|
|
await manager.send(session_id, data)
|
|
except asyncio.TimeoutError:
|
|
# 检查连接是否还活跃
|
|
if not manager.is_connected(session_id):
|
|
break
|
|
except Exception:
|
|
break
|
|
|
|
import asyncio
|
|
|
|
forward_task = asyncio.create_task(forward_output())
|
|
|
|
# 主循环：接收用户输入
|
|
try:
|
|
while True:
|
|
data = await websocket.receive_text()
|
|
# 接收用户输入，写入 PTY
|
|
await pty_manager.write(session_id, data + "\n")
|
|
except WebSocketDisconnect:
|
|
pass
|
|
finally:
|
|
forward_task.cancel()
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
manager.disconnect(session_id)
|