from datetime import datetime, timedelta from typing import Dict, List, Optional from app.models import SessionState, WorkerState, DevicePolicy from app.security import generate_session_id, generate_session_secret from app.storage import save_session_event _sessions_by_id: Dict[str, SessionState] = {} _workers_by_session: Dict[str, Dict[str, WorkerState]] = {} def create_session(owner_label: Optional[str] = None) -> SessionState: from app.config import Settings session_id = generate_session_id() secret = generate_session_secret() ttl = Settings.get_session_ttl_seconds() now = datetime.utcnow() session = SessionState( session_id=session_id, owner_label=owner_label, secret=secret, created_at=now, expires_at=now + timedelta(seconds=ttl), policy=DevicePolicy(session_id=session_id), ) _sessions_by_id[session_id] = session _workers_by_session[session_id] = {} save_session_event({ "event": "created", "session_id": session_id, "owner_label": owner_label, "timestamp": now.isoformat(), }) return session def get_session(session_id: str) -> Optional[SessionState]: expire_old_sessions() return _sessions_by_id.get(session_id) def delete_session(session_id: str) -> bool: if session_id in _sessions_by_id: del _sessions_by_id[session_id] _workers_by_session.pop(session_id, None) return True return False def expire_old_sessions() -> None: now = datetime.utcnow() expired = [sid for sid, s in _sessions_by_id.items() if s.expires_at < now] for sid in expired: delete_session(sid) def attach_worker(session_id: str, worker: WorkerState) -> bool: session = get_session(session_id) if not session: return False if session_id not in _workers_by_session: _workers_by_session[session_id] = {} _workers_by_session[session_id][worker.worker_id] = worker session.workers[worker.worker_id] = worker return True def detach_worker(session_id: str, worker_id: str) -> bool: session = get_session(session_id) if session and worker_id in session.workers: session.workers.pop(worker_id, None) ws = _workers_by_session.get(session_id) if ws and worker_id in ws: del ws[worker_id] return True return False def list_workers(session_id: str) -> List[WorkerState]: return list(_workers_by_session.get(session_id, {}).values()) def update_worker_heartbeat(session_id: str, worker_id: str) -> bool: ws = _workers_by_session.get(session_id, {}) worker = ws.get(worker_id) if worker: worker.last_heartbeat = datetime.utcnow() return True return False def set_session_policy(session_id: str, policy: DevicePolicy) -> bool: session = get_session(session_id) if session: session.policy = policy return True return False def get_session_policy(session_id: str) -> Optional[DevicePolicy]: session = get_session(session_id) return session.policy if session else None