WebSocket Handler

November 11, 2025 · View on GitHub

The UFOWebSocketHandler is the central nervous system of the server, implementing the Agent Interaction Protocol (AIP) to manage structured, reliable communication between the server and all connected clients.

For context on how this component fits into the server architecture, see the Server Overview.

🎯 Overview

The WebSocket Handler acts as the protocol orchestrator, managing all aspects of client communication:

Responsibility	Description	Protocol Used
Client Registration	Validate and register new device/constellation connections	AIP Registration Protocol
Task Dispatch	Route task requests to appropriate devices	AIP Task Execution Protocol
Heartbeat Monitoring	Maintain connection health via periodic pings	AIP Heartbeat Protocol
Device Info Exchange	Query and share device capabilities	AIP Device Info Protocol
Command Results	Relay execution results from devices to requesters	AIP Message Transport
Error Handling	Gracefully handle communication failures	AIP Error Protocol
Connection Lifecycle	Manage registration → active → cleanup flow	WebSocket + AIP

Architecture Position

graph TB
    subgraph "Clients"
        DC[Device Clients]
        CC[Constellation Clients]
    end
    
    subgraph "Server - WebSocket Handler"
        WH[UFOWebSocketHandler]
        
        subgraph "AIP Protocols"
            REG[Registration Protocol]
            HB[Heartbeat Protocol]
            DI[Device Info Protocol]
            TE[Task Execution Protocol]
        end
        
        subgraph "Message Router"
            MH[handle_message]
        end
    end
    
    subgraph "Server Components"
        WSM[Client Connection Manager]
        SM[Session Manager]
    end
    
    DC -->|WebSocket /ws| WH
    CC -->|WebSocket /ws| WH
    
    WH --> REG
    WH --> HB
    WH --> DI
    WH --> TE
    
    WH --> MH
    MH -->|"handle_task_request"| TE
    MH -->|"handle_heartbeat"| HB
    MH -->|"handle_device_info_request"| DI
    MH -->|"handle_command_result"| SM
    
    WH -->|"add_client / get_client"| WSM
    WH -->|"execute_task_async"| SM
    
    style WH fill:#ffecb3
    style MH fill:#bbdefb
    style SM fill:#c8e6c9
    style WSM fill:#f8bbd0

🔌 AIP Protocol Integration

The handler uses four specialized AIP protocols, each handling a specific aspect of communication. This separation of concerns makes the code maintainable and testable. For detailed protocol specifications, see the AIP Protocol Documentation.

def __init__(self, client_manager, session_manager, local=False):
    # Initialize per-connection protocols
    self.transport = None
    self.registration_protocol = None
    self.heartbeat_protocol = None
    self.device_info_protocol = None
    self.task_protocol = None

Protocol	Purpose	Key Methods	Message Types
Registration Protocol	Client identity and validation	`send_registration_confirmation()`, `send_registration_error()`	`REGISTER`, `REGISTER_CONFIRM`
Heartbeat Protocol	Connection health monitoring	`send_heartbeat_ack()`	`HEARTBEAT`, `HEARTBEAT_ACK`
Device Info Protocol	Capability exchange	`send_device_info_response()`, `send_device_info_request()`	`DEVICE_INFO_REQUEST`, `DEVICE_INFO_RESPONSE`
Task Execution Protocol	Task lifecycle management	`send_task_assignment()`, `send_ack()`, `send_error()`	`TASK`, `TASK_ASSIGNMENT`, `TASK_END`

Protocol Initialization per Connection:

async def connect(self, websocket: WebSocket) -> str:
    await websocket.accept()
    
    # Initialize AIP protocols for this connection
    self.transport = WebSocketTransport(websocket)
    self.registration_protocol = RegistrationProtocol(self.transport)
    self.heartbeat_protocol = HeartbeatProtocol(self.transport)
    self.device_info_protocol = DeviceInfoProtocol(self.transport)
    self.task_protocol = TaskExecutionProtocol(self.transport)
    
    # ... registration flow ...

Per-Connection Protocol Instances:

Each WebSocket connection gets its own set of protocol instances, ensuring message routing and state management are isolated between clients.

📝 Client Registration

Registration is the critical first step when a client connects. The handler validates client identity, checks permissions, and establishes the communication session.

Registration Flow

sequenceDiagram
    participant C as Client<br/>(Device/Constellation)
    participant WS as WebSocket Handler
    participant RP as Registration Protocol
    participant WSM as Client Connection Manager
    
    Note over C,WS: 1️⃣ Connection Establishment
    C->>WS: WebSocket CONNECT /ws
    WS->>WS: await websocket.accept()
    WS->>WS: Initialize AIP protocols
    
    Note over WS,RP: 2️⃣ Registration Message
    WS->>C: (AIP Transport ready)
    C->>RP: REGISTER<br/>{client_id, client_type, platform, metadata}
    RP->>RP: Parse & validate JSON
    
    Note over RP,WS: 3️⃣ Validation
    RP->>WS: ClientMessage object
    WS->>WS: Validate client_id exists
    WS->>WS: Validate message type = REGISTER
    
    alt Client Type = Constellation
        WS->>WSM: is_device_connected(target_id)?
        alt Target device offline
            WSM-->>WS: False
            WS->>RP: send_registration_error()
            RP->>C: ERROR: "Target device not connected"
            WS->>C: WebSocket close()
        else Target device online
            WSM-->>WS: True
            WS->>WSM: add_client(client_id, ...)
        end
    else Client Type = Device
        WS->>WSM: add_client(client_id, platform, ...)
    end
    
    Note over WS,C: 4️⃣ Confirmation
    WS->>RP: send_registration_confirmation()
    RP->>C: REGISTER_CONFIRM<br/>{status: "success"}
    
    Note over C: Client registered, ready for tasks
    C->>C: Start message listening loop

Registration Steps (Code Walkthrough)

Step 1: WebSocket Connection Accepted

async def connect(self, websocket: WebSocket) -> str:
    # Accept WebSocket connection
    await websocket.accept()
    
    # Initialize AIP protocols for this connection
    self.transport = WebSocketTransport(websocket)
    self.registration_protocol = RegistrationProtocol(self.transport)
    self.heartbeat_protocol = HeartbeatProtocol(self.transport)
    self.device_info_protocol = DeviceInfoProtocol(self.transport)
    self.task_protocol = TaskExecutionProtocol(self.transport)

Step 2: Receive Registration Message

async def _parse_registration_message(self) -> ClientMessage:
    """Parse and validate registration message using AIP Transport."""
    self.logger.info("[WS] [AIP] Waiting for registration message...")
    
    # Receive via AIP Transport
    reg_data = await self.transport.receive()
    if isinstance(reg_data, bytes):
        reg_data = reg_data.decode("utf-8")
    
    # Parse using Pydantic model
    reg_info = ClientMessage.model_validate_json(reg_data)
    
    self.logger.info(
        f"[WS] [AIP] Received registration from {reg_info.client_id}, "
        f"type={reg_info.client_type}"
    )
    
    return reg_info

Expected Registration Message:

{
  "type": "REGISTER",
  "client_id": "device_windows_001",
  "client_type": "DEVICE",
  "platform": "windows",
  "metadata": {
    "hostname": "DESKTOP-ABC123",
    "os_version": "Windows 11 Pro",
    "screen_resolution": "1920x1080"
  }
}

Step 3: Validation

# Basic validation
if not reg_info.client_id:
    raise ValueError("Client ID is required for WebSocket registration")
if reg_info.type != ClientMessageType.REGISTER:
    raise ValueError("First message must be a registration message")

# Constellation-specific validation
if client_type == ClientType.CONSTELLATION:
    await self._validate_constellation_client(reg_info)

Constellation Validation:

async def _validate_constellation_client(self, reg_info: ClientMessage) -> None:
    """Validate constellation's claimed target_id."""
    claimed_device_id = reg_info.target_id
    
    if not claimed_device_id:
        return  # No device_id to validate
    
    # Check if target device is connected
    if not self.client_manager.is_device_connected(claimed_device_id):
        error_msg = f"Target device '{claimed_device_id}' is not connected"
        self.logger.warning(f"[WS] Constellation registration failed: {error_msg}")
        
        # Send error via AIP protocol
        await self._send_error_response(error_msg)
        await self.transport.close()
        raise ValueError(error_msg)

Step 4: Register Client in ClientConnectionManager

client_type = reg_info.client_type
platform = reg_info.metadata.get("platform", "windows") if reg_info.metadata else "windows"

# Register in Client Connection Manager
self.client_manager.add_client(
    client_id,
    platform,
    websocket,
    client_type,
    reg_info.metadata,
    transport=self.transport,
    task_protocol=self.task_protocol,
)

Step 5: Send Confirmation

async def _send_registration_confirmation(self) -> None:
    """Send successful registration confirmation using AIP RegistrationProtocol."""
    self.logger.info("[WS] [AIP] Sending registration confirmation...")
    await self.registration_protocol.send_registration_confirmation()
    self.logger.info("[WS] [AIP] Registration confirmation sent")

Confirmation Message:

{
  "type": "REGISTER_CONFIRM",
  "status": "success",
  "timestamp": "2024-11-04T14:30:22.123456+00:00",
  "response_id": "uuid-v4"
}

Step 6: Log Success

def _log_client_connection(self, client_id: str, client_type: ClientType) -> None:
    """Log successful client connection with appropriate emoji."""
    if client_type == ClientType.DEVICE:
        self.logger.info(f"[WS] Registered device client: {client_id}")
    elif client_type == ClientType.CONSTELLATION:
        self.logger.info(f"[WS] 🌟 Registered constellation client: {client_id}")

Validation Rules

Validation	Check	Error Message	Action
Client ID Presence	`client_id` field exists and not empty	`"Client ID is required"`	Reject connection
Message Type	First message type == `REGISTER`	`"First message must be a registration message"`	Reject connection
Target Device (Constellation)	If `target_id` specified, device must be online	`"Target device '<id>' is not connected"`	Send error + close
Client ID Uniqueness	No existing client with same ID	Handled by ClientConnectionManager	Disconnect old connection

Constellation Dependency:

Constellations must specify a valid target_id that refers to an already-connected device. If the device is offline or doesn't exist, registration fails immediately.

Workaround: Connect devices first, then constellations.

Security Consideration: The current implementation does not authenticate clients. Any client can register with any client_id. For production deployments:

- Implement authentication tokens in `metadata`
- Validate client certificates (TLS client auth)
- Use API keys or OAuth tokens
- Whitelist allowed `client_id` patterns

📨 Message Handling

After registration, the handler enters a message loop, routing incoming client messages to specialized handlers based on message type.

Message Dispatcher

graph TB
    WS[WebSocket receive_text] --> Parse[Parse ClientMessage JSON]
    Parse --> Router{Message Type?}
    
    Router -->|TASK| HT[handle_task_request]
    Router -->|COMMAND_RESULTS| HC[handle_command_result]
    Router -->|HEARTBEAT| HH[handle_heartbeat]
    Router -->|ERROR| HE[handle_error]
    Router -->|DEVICE_INFO_REQUEST| HD[handle_device_info_request]
    Router -->|DEVICE_INFO_RESPONSE| HDR[handle_device_info_response]
    Router -->|Unknown| HU[handle_unknown]
    
    HT --> SM[Session Manager]
    HC --> CD[Command Dispatcher]
    HH --> HP[Heartbeat Protocol]
    HE --> Log[Error Logging]
    HD --> DIP[Device Info Protocol]
    
    style Router fill:#ffe0b2
    style SM fill:#c8e6c9
    style HP fill:#bbdefb

Dispatcher Implementation:

async def handle_message(self, msg: str, websocket: WebSocket) -> None:
    """Dispatch incoming messages to specific handlers."""
    try:
        # Parse message using Pydantic model
        data = ClientMessage.model_validate_json(msg)
        
        client_id = data.client_id
        client_type = data.client_type
        msg_type = data.type
        
        # Route to appropriate handler
        if msg_type == ClientMessageType.TASK:
            await self.handle_task_request(data, websocket)
        elif msg_type == ClientMessageType.COMMAND_RESULTS:
            await self.handle_command_result(data)
        elif msg_type == ClientMessageType.HEARTBEAT:
            await self.handle_heartbeat(data, websocket)
        elif msg_type == ClientMessageType.ERROR:
            await self.handle_error(data, websocket)
        elif msg_type == ClientMessageType.DEVICE_INFO_REQUEST:
            await self.handle_device_info_request(data, websocket)
        elif msg_type == ClientMessageType.DEVICE_INFO_RESPONSE:
            await self.handle_device_info_response(data, websocket)
        else:
            await self.handle_unknown(data, websocket)
    
    except Exception as e:
        self.logger.error(f"Error handling message: {e}")
        try:
            await self.task_protocol.send_error(str(e))
        except (ConnectionError, IOError):
            pass  # Connection already closed

Message Type Handlers:

Handler	Triggered By	Purpose	Response
`handle_task_request`	`TASK`	Client requests task execution	`TASK_ASSIGNMENT` device
`handle_command_result`	`COMMAND_RESULTS`	Device reports command execution result	Unblock command dispatcher
`handle_heartbeat`	`HEARTBEAT`	Connection health ping	`HEARTBEAT_ACK`
`handle_error`	`ERROR`	Client reports error	Log + send error acknowledgment
`handle_device_info_request`	`DEVICE_INFO_REQUEST`	Constellation queries device capabilities	`DEVICE_INFO_RESPONSE`
`handle_device_info_response`	`DEVICE_INFO_RESPONSE`	Device provides info (pull model)	Store in ClientConnectionManager
`handle_unknown`	Any other type	Unknown/unsupported message	Log warning + send error

Task Request Handling

The handler supports task requests from both device clients (self-execution) and constellation clients (orchestrated execution on target devices).

Task Request Flow:

sequenceDiagram
    participant C as Constellation
    participant WH as WebSocket Handler
    participant WSM as Client Connection Manager
    participant SM as Session Manager
    participant D as Device
    
    Note over C,WH: 1️⃣ Task Request
    C->>WH: TASK<br/>{request, target_id, session_id}
    WH->>WH: Validate target_id
    
    Note over WH,WSM: 2️⃣ Resolve Target Device
    WH->>WSM: get_client(target_id)
    WSM-->>WH: Device WebSocket
    WH->>WSM: get_client_info(target_id)
    WSM-->>WH: {platform: "windows"}
    
    Note over WH,SM: 3️⃣ Create Session
    WH->>SM: execute_task_async(<br/>  session_id,<br/>  request,<br/>  target_ws,<br/>  platform,<br/>  callback<br/>)
    SM-->>WH: session_id (non-blocking!)
    
    Note over WH,C: 4️⃣ Immediate Acknowledgment
    WH->>C: ACK<br/>{session_id, status: "dispatched"}
    
    Note over SM,D: 5️⃣ Background Execution
    SM->>D: TASK_ASSIGNMENT<br/>{request, session_id}
    D->>D: Execute task (LLM + actions)
    
    Note over D,WH: 6️⃣ Result Callback
    D-->>SM: Task complete
    SM->>WH: callback(session_id, result_msg)
    WH->>C: TASK_END<br/>{status, result}

Device Client Self-Execution:

When a device requests a task for itself:

async def handle_task_request(self, data: ClientMessage, websocket: WebSocket):
    client_id = data.client_id
    client_type = data.client_type
    
    if client_type == ClientType.DEVICE:
        # Device executing task on itself
        target_ws = websocket  # Use requesting client's WebSocket
        platform = self.client_manager.get_client_info(client_id).platform
        target_device_id = client_id
    # ...

Constellation Orchestrated Execution:

When a constellation dispatches a task to a target device:

async def handle_task_request(self, data: ClientMessage, websocket: WebSocket):
    client_id = data.client_id
    client_type = data.client_type
    
    if client_type == ClientType.CONSTELLATION:
        # Constellation dispatching to target device
        target_device_id = data.target_id
        target_ws = self.client_manager.get_client(target_device_id)
        platform = self.client_manager.get_client_info(target_device_id).platform
        
        # Validate target device exists
            if not target_ws:
                raise ValueError(f"Target device '{target_device_id}' not connected")
            
            # Track session mappings
            session_id = data.session_id or str(uuid.uuid4())
            self.client_manager.add_constellation_session(client_id, session_id)
            self.client_manager.add_device_session(target_device_id, session_id)
        # ...
    ```

**Background Task Execution:**

```python
# Define callback for result delivery
async def send_result(sid: str, result_msg: ServerMessage):
    """Send result back to requester when task completes."""
    # Send to constellation client
    if client_type == ClientType.CONSTELLATION:
        if websocket.client_state == WebSocketState.CONNECTED:
            await websocket.send_text(result_msg.model_dump_json())
        
        # Also send to target device (optional)
        if target_ws and target_ws.client_state == WebSocketState.CONNECTED:
            await target_ws.send_text(result_msg.model_dump_json())
    else:
        # Send to device client
        if websocket.client_state == WebSocketState.CONNECTED:
            await websocket.send_text(result_msg.model_dump_json())

# Execute in background via SessionManager
await self.session_manager.execute_task_async(
    session_id=session_id,
    task_name=task_name,
    request=data.request,
    websocket=target_ws,  # Device WebSocket for command dispatcher
    platform_override=platform,
    callback=send_result  # Called when task completes
)

# Send immediate acknowledgment (non-blocking)
await self.task_protocol.send_ack(session_id=session_id)

Why Immediate ACK?

The handler sends an immediate ACK after dispatching the task to the SessionManager. This confirms:

Task was received and validated
Session was created successfully
Task is now executing in background

The actual task result is delivered later via the send_result callback.

Session Tracking:

Client Type	Session Tracking	Purpose
Device	`client_manager.add_device_session(device_id, session_id)`	Track which device is executing the session
Constellation	`client_manager.add_constellation_session(constellation_id, session_id)`	Track which constellation requested the session
Both	Session Manager stores session `BaseSession` object	Execute and manage task lifecycle

Command Result Handling

When a device executes a command (e.g., "click button", "type text"), it sends results back to the server for processing by the session's command dispatcher.

Command Result Flow:

sequenceDiagram
    participant S as Session (on server)
    participant CD as Command Dispatcher
    participant D as Device
    participant WH as WebSocket Handler
    
    Note over S,D: Session is running on server
    S->>CD: execute_command("open_notepad")
    CD->>D: Send command via WebSocket<br/>(response_id="cmd_123")
    CD->>CD: await response (blocking)
    
    Note over D: Device executes command
    D->>D: Open Notepad application
    D->>D: Take screenshot
    
    Note over D,WH: Send result back
    D->>WH: COMMAND_RESULTS<br/>{response_id="cmd_123", result, screenshot}
    WH->>WH: handle_command_result()
    WH->>CD: set_result(response_id, data)
    
    Note over CD: Unblocks await!
    CD-->>S: Return command result
    S->>S: Continue session execution

Handler Implementation:

async def handle_command_result(self, data: ClientMessage):
    """
    Handle command execution results from devices.
    Unblocks the command dispatcher waiting for this response.
    """
    response_id = data.prev_response_id  # ID of the command request
    session_id = data.session_id
    
    self.logger.debug(
        f"[WS] Received command result for response_id={response_id}, "
        f"session_id={session_id}"
    )
    
    # Get session's command dispatcher
    session = self.session_manager.get_or_create_session(session_id)
    command_dispatcher = session.context.command_dispatcher
    
    # Set result (unblocks waiting dispatcher)
    await command_dispatcher.set_result(response_id, data)
    
    self.logger.debug(
        f"[WS] Command result set for response_id={response_id}"
    )

Critical for Session Execution:

Without proper command result handling, sessions would hang indefinitely waiting for device responses. The set_result() call is what unblocks the await in the command dispatcher.

Heartbeat Handling

Heartbeats are lightweight ping/pong messages that ensure the WebSocket connection is alive and healthy.

async def handle_heartbeat(self, data: ClientMessage, websocket: WebSocket) -> None:
    """Handle heartbeat messages using AIP HeartbeatProtocol."""
    self.logger.debug(f"[WS] [AIP] Heartbeat from {data.client_id}")
    
    try:
        # Send acknowledgment via AIP protocol
        await self.heartbeat_protocol.send_heartbeat_ack()
        self.logger.debug(f"[WS] [AIP] Heartbeat response sent to {data.client_id}")
    except (ConnectionError, IOError) as e:
        # Connection closed - log but don't fail
        self.logger.debug(f"[WS] [AIP] Could not send heartbeat ack: {e}")

Heartbeat Message:

{
  "type": "HEARTBEAT",
  "client_id": "device_windows_001",
  "timestamp": "2024-11-04T14:30:22.123456+00:00"
}

Heartbeat ACK:

{
  "type": "HEARTBEAT_ACK",
  "timestamp": "2024-11-04T14:30:22.234567+00:00",
  "response_id": "uuid-v4"
}

Heartbeat Best Practices:

Frequency: Clients should send heartbeats every 30-60 seconds
Timeout: Server should consider connection dead after 2-3 missed heartbeats
Lightweight: Heartbeat messages are small and processed quickly
Non-blocking: Heartbeat handling doesn't block task execution

Device Info Handling

Constellations can query device capabilities (screen resolution, installed apps, OS version) to make intelligent task routing decisions.

Device Info Request Flow:

sequenceDiagram
    participant C as Constellation
    participant WH as WebSocket Handler
    participant WSM as Client Connection Manager
    participant DIP as Device Info Protocol
    participant D as Device
    
    Note over C,WH: 1️⃣ Request Device Info
    C->>WH: DEVICE_INFO_REQUEST<br/>{target_id, request_id}
    
    Note over WH,WSM: 2️⃣ Resolve Device
    WH->>WSM: get_client(target_id)
    WSM-->>WH: Device WebSocket
    
    Note over WH,D: 3️⃣ Forward Request
    WH->>DIP: send_device_info_request()
    DIP->>D: DEVICE_INFO_REQUEST
    
    Note over D: 4️⃣ Collect Info
    D->>D: Gather system info<br/>(screen, OS, apps)
    
    Note over D,WH: 5️⃣ Response
    D->>DIP: DEVICE_INFO_RESPONSE<br/>{screen_res, os_version, ...}
    DIP->>WH: Parse response
    
    Note over WH,C: 6️⃣ Forward to Constellation
    WH->>C: DEVICE_INFO_RESPONSE<br/>{device_info, request_id}

async def handle_device_info_request(
    self, data: ClientMessage, websocket: WebSocket
) -> None:
    """Handle device info requests from constellations."""
    device_id = data.target_id
    request_id = data.request_id
    
    self.logger.info(
        f"[WS] Constellation {data.client_id} requesting info for device {device_id}"
    )
    
    # Get device info (may involve querying the device)
    device_info = await self.get_device_info(device_id)
    
    # Send via AIP protocol
    await self.device_info_protocol.send_device_info_response(
        device_info=device_info,
        request_id=request_id
    )

Device Info Structure:

{
  "device_id": "device_windows_001",
  "platform": "windows",
  "os_version": "Windows 11 Pro 22H2",
  "screen_resolution": "1920x1080",
  "installed_applications": ["Chrome", "Excel", "Notepad", "..."],
  "capabilities": ["ui_automation", "file_operations", "web_browsing"],
  "cpu_cores": 8,
  "memory_gb": 16
}

🔌 Client Disconnection

Critical Cleanup Process:

When a client disconnects (gracefully or abruptly), the handler must clean up sessions, remove registry entries, and prevent resource leaks.

Disconnection Detection

async def handler(self, websocket: WebSocket) -> None:
    """FastAPI WebSocket entry point."""
    client_id = None
    
    try:
        # Registration
        client_id = await self.connect(websocket)
        
        # Message loop
        while True:
            msg = await websocket.receive_text()
            asyncio.create_task(self.handle_message(msg, websocket))
    
    except WebSocketDisconnect as e:
        # Normal disconnection
        self.logger.warning(
            f"[WS] {client_id} disconnected code={e.code}, reason={e.reason}"
        )
        if client_id:
            await self.disconnect(client_id)
    
    except Exception as e:
        # Unexpected error
        self.logger.error(f"[WS] Error with client {client_id}: {e}")
        if client_id:
            await self.disconnect(client_id)

Cleanup Process

graph TD
    A[Client Disconnects] --> B{Get Client Info}
    B --> C{Client Type?}
    
    C -->|Device| D[Get Device Sessions]
    C -->|Constellation| E[Get Constellation Sessions]
    
    D --> F[Cancel Each Session<br/>reason='device_disconnected']
    E --> G[Cancel Each Session<br/>reason='constellation_disconnected']
    
    F --> H[Remove Device Session Mappings]
    G --> I[Remove Constellation Session Mappings]
    
    H --> J[Remove Client from ClientConnectionManager]
    I --> J
    
    J --> K[Log Disconnection]
    K --> L[Cleanup Complete]
    
    style F fill:#ffcdd2
    style G fill:#ffcdd2
    style J fill:#c8e6c9

Device Client Cleanup:

async def disconnect(self, client_id: str) -> None:
    """Handle client disconnection and cleanup."""
    client_info = self.client_manager.get_client_info(client_id)
    
    if client_info and client_info.client_type == ClientType.DEVICE:
        # Get all sessions running on this device
        session_ids = self.client_manager.get_device_sessions(client_id)
        
        if session_ids:
            self.logger.info(
                f"[WS] 📱 Device {client_id} disconnected, "
                f"cancelling {len(session_ids)} active session(s)"
            )
        
        # Cancel all sessions
        for session_id in session_ids:
            try:
                await self.session_manager.cancel_task(
                    session_id,
                    reason="device_disconnected"  # Send callback to constellation
                )
            except Exception as e:
                self.logger.error(f"Error cancelling session {session_id}: {e}")
        
        # Clean up mappings
        self.client_manager.remove_device_sessions(client_id)

Constellation Client Cleanup:

if client_info and client_info.client_type == ClientType.CONSTELLATION:
    # Get all sessions initiated by constellation
    session_ids = self.client_manager.get_constellation_sessions(client_id)
    
    if session_ids:
        self.logger.info(
            f"[WS] 🌟 Constellation {client_id} disconnected, "
            f"cancelling {len(session_ids)} active session(s)"
        )
    
    # Cancel all associated sessions
    for session_id in session_ids:
        try:
            await self.session_manager.cancel_task(
                session_id,
                reason="constellation_disconnected"  # Don't send callback
            )
        except Exception as e:
            self.logger.error(f"Error cancelling session {session_id}: {e}")
    
    # Clean up mappings
    self.client_manager.remove_constellation_sessions(client_id)

Final Registry Cleanup:

# Remove client from registry
self.client_manager.remove_client(client_id)
self.logger.info(f"[WS] {client_id} disconnected")

Cancellation Behavior Comparison

Scenario	Cancellation Reason	Callback Sent?	Why?
Device Disconnects	`device_disconnected`	Yes Constellation	Notify orchestrator to reassign task
Constellation Disconnects	`constellation_disconnected`	No	Requester is gone, no one to notify

Proper Cleanup is Critical:

Failing to clean up disconnected clients leads to:

Orphaned sessions consuming server memory
Stale WebSocket references causing errors
Registry pollution with non-existent clients
Resource leaks (file handles, memory)

🚨 Error Handling

The handler implements comprehensive error handling to prevent failures from cascading and breaking the entire server.

Error Categories

Error Type	Handler Location	Recovery Strategy
Connection Errors	`send_*` methods	Log and skip (connection already closed)
Message Parsing Errors	`handle_message`	Send error response via AIP
Task Execution Errors	`handle_task_request`	Log + send error via task protocol
Validation Errors	`_validate_*` methods	Send error + close connection
Callback Errors	Session Manager	Log but don't fail session

Connection Error Handling

async def handle_heartbeat(self, data: ClientMessage, websocket: WebSocket):
    try:
        await self.heartbeat_protocol.send_heartbeat_ack()
    except (ConnectionError, IOError) as e:
        # Connection closed - log but don't fail
        self.logger.debug(f"Could not send heartbeat ack: {e}")
        # Don't raise - connection is already closed

Why Catch and Ignore?

When a connection is abruptly closed, attempts to send messages will raise ConnectionError. Since the client is already gone, there's no point in propagating the error—just log it and continue cleanup.

Message Parsing Errors

async def handle_message(self, msg: str, websocket: WebSocket):
    try:
        data = ClientMessage.model_validate_json(msg)
        # ... route to handlers ...
    
    except Exception as e:
        import traceback
        traceback.print_exc()
        self.logger.error(f"Error handling message: {e}")
        
        # Try to send error response
        try:
            await self.task_protocol.send_error(str(e))
        except (ConnectionError, IOError) as send_error:
            self.logger.debug(f"Could not send error response: {send_error}")

Error Message Format:

{
  "type": "ERROR",
  "error": "Invalid message format: missing required field 'client_id'",
  "timestamp": "2024-11-04T14:30:22.123456+00:00",
  "response_id": "uuid-v4"
}

Task Execution Errors

async def handle_task_request(self, data: ClientMessage, websocket: WebSocket):
    try:
        # Validate target device
        if client_type == ClientType.CONSTELLATION:
            target_ws = self.client_manager.get_client(target_device_id)
            if not target_ws:
                raise ValueError(f"Target device '{target_device_id}' not connected")
        
        # Execute task
        await self.session_manager.execute_task_async(...)
    
    except Exception as e:
        self.logger.error(f"Error handling task: {e}")
        await self.task_protocol.send_error(str(e))

Validation Errors with Connection Closure

async def _validate_constellation_client(self, reg_info: ClientMessage) -> None:
    """Validate constellation's target device."""
    claimed_device_id = reg_info.target_id
    
    if not self.client_manager.is_device_connected(claimed_device_id):
        error_msg = f"Target device '{claimed_device_id}' is not connected"
        self.logger.warning(f"Constellation registration failed: {error_msg}")
        
        # Send error via AIP protocol
        await self._send_error_response(error_msg)
        
        # Close connection immediately
        await self.transport.close()
        
        # Raise to prevent further processing
        raise ValueError(error_msg)

When to Close Connections:

Close connections immediately for:

Invalid registration (missing client_id, wrong message type)
Authorization failures (target device not connected for constellations)
Protocol violations (sending TASK before REGISTER)

For other errors, log and send error messages but keep connection alive.

Best Practices

1. Validate Early and Thoroughly

# Good: Validate immediately after parsing
async def handle_task_request(self, data: ClientMessage, websocket: WebSocket):
    if not data.request:
        raise ValueError("Task request cannot be empty")
    if not data.client_id:
        raise ValueError("Client ID required")
    if data.client_type == ClientType.CONSTELLATION and not data.target_id:
        raise ValueError("Constellation must specify target_id")
    # ... proceed with validated data ...

2. Always Check Connection State Before Sending

from starlette.websockets import WebSocketState

# Good: Check state before sending
async def send_result(sid: str, result_msg: ServerMessage):
    if websocket.client_state == WebSocketState.CONNECTED:
        await websocket.send_text(result_msg.model_dump_json())
    else:
        self.logger.debug(f"Cannot send result, connection closed for {sid}")

WebSocket States:

State	Description	Can Send?
`CONNECTING`	Handshake in progress	No
`CONNECTED`	Active connection	Yes
`DISCONNECTED`	Connection closed	No

3. Handle Cancellation Gracefully with Context

# Good: Different reasons need different handling
async def disconnect(self, client_id: str):
    client_info = self.client_manager.get_client_info(client_id)
    
    if client_info.client_type == ClientType.CONSTELLATION:
        reason = "constellation_disconnected"  # Don't send callback
    else:
        reason = "device_disconnected"  # Send callback to constellation
    
    for session_id in session_ids:
        await self.session_manager.cancel_task(session_id, reason=reason)

4. Use Structured Logging with Context

# Good: Include client type and context
if client_type == ClientType.CONSTELLATION:
    self.logger.info(
        f"[WS] 🌟 Constellation {client_id} requesting task on {target_id}"
    )
else:
    self.logger.debug(
        f"[WS] 📱 Received device message from {client_id}, type: {data.type}"
    )

Logging Levels:

DEBUG: Heartbeats, message routing, low-level protocol details
INFO: Registration, disconnection, task dispatch, major lifecycle events
WARNING: Validation failures, connection issues, recoverable errors
ERROR: Unexpected exceptions, critical failures

5. Implement Async Message Handling

# Good: Process messages in background tasks
async def handler(self, websocket: WebSocket):
    while True:
        msg = await websocket.receive_text()
        asyncio.create_task(self.handle_message(msg, websocket))
        # Loop continues immediately, doesn't wait for handler to finish

Why asyncio.create_task?

Without create_task, the handler would process messages sequentially, blocking new messages while handling the current one. This is problematic for:

Long-running task dispatches
Command result processing
Device info queries

Background tasks allow concurrent message processing while keeping the receive loop responsive.

Explore related components to understand the full server architecture:

Component	Purpose	Link
Server Overview	High-level architecture and capabilities	Overview
Quick Start	Start server and dispatch first task	Quick Start
Session Manager	Session lifecycle and background execution	Session Manager
Client Connection Manager	Connection registry and session tracking	Client Connection Manager
HTTP API	RESTful API endpoints	API Reference
AIP Protocol	Agent Interaction Protocol details	AIP Overview

🎓 What You Learned

After reading this guide, you should understand:

AIP Protocol Integration - Four specialized protocols handle different communication aspects
Registration Flow - Validation → Registration → Confirmation
Message Routing - Central dispatcher routes messages to specialized handlers
Dual Client Support - Devices (self-execution) vs. Constellations (orchestration)
Background Task Dispatch - Immediate ACK + async execution
Command Result Handling - Unblocks command dispatcher waiting for device responses
Heartbeat Monitoring - Lightweight connection health checks
Disconnection Cleanup - Context-aware session cancellation and registry cleanup
Error Handling - Graceful degradation without cascading failures

Next Steps:

Explore Session Manager to understand background execution internals
Learn about Client Connection Manager for client registry management
Review AIP Protocol Documentation for message format specifications