"""Bounded Manager status channel. Runtime state and command ownership stay in Analytics.""" import asyncio from collections import Counter import contextlib import json import ssl import time import uuid import subprocess from fastapi import WebSocket, WebSocketDisconnect from control_sidecar.protocol import BoundaryError, bearer, loads class ControlSocket: interval = 30 max_missed = 3 def __init__(self, cfg, transport, auth): self.cfg, self.tx, self.auth = cfg, transport, auth self.active = None self.admitting = False self.started = int(time.time()*1000) self.stats = Counter() self.actions = Counter() self.last_error = None self.stopping = False def snapshot(self): return {'connected': self.active is not None, 'counts': dict(self.stats), 'actions': dict(self.actions), 'lastError': self.last_error} async def result(self, action): body = self.cfg.login_body if action == 'tokens.flushCache': return {'evicted': 0} if action == 'ai.connectionInfo.get': return {'serverId': self.cfg.server_id, **{k: body[k] for k in ('advertisedIp', 'controlPort', 'bboxStreamKey', 'connectionConfigVersion')}, 'redis': {k: body['redis'][k] for k in ('host', 'port', 'db', 'username', 'tls')}} if action == 'control.server.status': cert = None if self.cfg.certificate_file: output = subprocess.run(['openssl', 'x509', '-in', self.cfg.certificate_file, '-noout', '-subject', '-enddate'], check=True, capture_output=True, text=True, timeout=2).stdout fields = dict(line.split('=', 1) for line in output.splitlines()) cert = {'subject': fields['subject'], 'notAfterWallMs': int(ssl.cert_time_to_seconds(fields['notAfter'])*1000)} return {'enabled': True, 'running': True, 'controlBindIp': body['advertisedIp'], 'boundPort': body['controlPort'], 'startedAtWallMs': self.started, 'clientCount': int(self.active is not None), 'heartbeat': {'intervalSeconds': self.interval, 'maxMissedPongs': self.max_missed}, 'certificate': cert, 'lastError': None} if action == 'ai.runtimeStatus.get': try: code, raw = await self.tx.request('GET', self.cfg.runtime + '/internal/secure-control/runtime-status', token=self.cfg.runtime_token, seconds=1) if code != 200: raise BoundaryError('RUNTIME_UNAVAILABLE') return loads(raw) except BoundaryError: return {'serverId': self.cfg.server_id, 'status': 'NOT_READY', 'modelName': None, 'dependencies': {'runtime': {'ok': False, 'error': 'RUNTIME_UNAVAILABLE'}}, 'streamCounts': dict(running=0, starting=0, stopped=0, error=0, total=0), 'streams': []} raise BoundaryError('UNKNOWN_ACTION', 422) async def handle(self, socket: WebSocket): if self.stopping or self.active is not None or self.admitting: await socket.close(code=1013) return self.admitting = True try: incoming = bearer(socket.headers) async with asyncio.timeout(3): await self.auth.authorize(incoming) await socket.accept() self.active = socket self.stats['connections'] += 1 self.last_error = None except (BoundaryError, TimeoutError): self.stats['auth_rejected'] += 1 await socket.close(code=1008) return finally: self.admitting = False pending = set() send_lock = asyncio.Lock() async def send(value): async with send_lock: async with asyncio.timeout(3): await socket.send_json(value) async def heartbeat(): while True: await asyncio.sleep(self.interval) if len(pending) >= self.max_missed: self.last_error = 'HEARTBEAT_TIMEOUT' await socket.close(code=1001) return # Revoked Manager sessions cannot live forever on a previously accepted socket. try: await self.auth.authorize(incoming) except BoundaryError: self.last_error = 'AUTHORITY_UNAVAILABLE' await socket.close(code=1008) return ping = uuid.uuid4().hex pending.add(ping) await send({'type': 'ping', 'pingId': ping, 'serverTimeMs': int(time.time()*1000)}) self.stats['pings'] += 1 async def heartbeat_guard(): try: await heartbeat() except (BoundaryError, TimeoutError, OSError, RuntimeError): self.last_error = 'HEARTBEAT_TRANSPORT' with contextlib.suppress(Exception): await socket.close(code=1001) ticker = asyncio.create_task(heartbeat_guard(), name='bridge-heartbeat') start, count = time.monotonic(), 0 try: while not self.stopping: raw = await socket.receive_text() if time.monotonic()-start >= 1: start, count = time.monotonic(), 0 count += 1 if count > 32 or len(raw.encode()) > 65536: raise BoundaryError('WSS_LIMIT') doc = loads(raw.encode()) if doc.get('type') == 'pong': ping = doc.get('pingId') if isinstance(ping, str) and ping in pending: pending.remove(ping) self.stats['pongs'] += 1 continue rid, action = doc.get('requestId'), doc.get('action') if (not isinstance(rid, str) or not 1 <= len(rid) <= 128 or not isinstance(action, str) or not 1 <= len(action) <= 100 or doc.get('params', {}) != {} or set(doc)-{'requestId','action','params'}): raise BoundaryError('INVALID_WSS_REQUEST') try: async with asyncio.timeout(2): await self.auth.authorize(incoming) result = await self.result(action) self.actions[action] += 1 response = {'requestId': rid, 'ok': True, 'result': result} except BoundaryError as exc: response = {'requestId': rid, 'ok': False, 'error': {'code': exc.code, 'message': 'Action unavailable.', 'retryable': False}} except TimeoutError: response = {'requestId': rid, 'ok': False, 'error': {'code': 'ACTION_TIMEOUT', 'message': 'Action timed out.', 'retryable': False}} response['serverTimeMs'] = int(time.time()*1000) await send(response) except WebSocketDisconnect: pass except (BoundaryError, TimeoutError, OSError, ValueError, RuntimeError, KeyError, subprocess.SubprocessError): self.last_error = self.last_error or 'WSS_PROTOCOL_OR_TRANSPORT' with contextlib.suppress(Exception): await socket.close(code=1008) finally: ticker.cancel() with contextlib.suppress(asyncio.CancelledError, Exception): await ticker if self.active is socket: self.active = None self.stats['disconnects'] += 1 async def close(self): self.stopping = True if self.active is not None: with contextlib.suppress(Exception): async with asyncio.timeout(3): await self.active.close(code=1001)