WebSocket Real-time System

Relevant source files

The following files were used as context for generating this wiki page:

• crates/infra/src/websocket/manager.rs
• crates/infra/src/websocket/handler.rs
• crates/core-service/src/service/ws_message.rs
• apps/api/src/api/ws/mod.rs
• crates/infra/src/websocket/mod.rs

The WebSocket Real-time System in ATMOS provides a high-performance, asynchronous communication layer designed to handle persistent connections between clients and the server. It facilitates bi-directional data exchange, allowing the server to push updates to clients instantly without the overhead of traditional HTTP polling. The system is built on top of the Axum web framework and utilizes the Tokio runtime for efficient task scheduling and I/O operations.

The architecture is centered around a centralized connection management strategy that tracks active sessions, handles message routing, and manages the lifecycle of each socket. This system is critical for features requiring low-latency feedback, such as terminal emulation, file system monitoring, and collaborative workspace updates.

Architecture and Connection Lifecycle

The system employs a layered architecture where the API layer handles the initial HTTP upgrade request, which is then handed off to infrastructure-level managers.

High-Level Component Flow

The following diagram illustrates the flow from an initial client request to the establishment of a managed WebSocket session.

The entry point for WebSocket connections is located in the API module, which defines the /ws route. When a client initiates an upgrade, the system extracts necessary session metadata and passes the underlying stream to the infrastructure layer.
Sources: apps/api/src/api/ws/mod.rs, crates/infra/src/websocket/handler.rs

Session Management

The WsManager serves as the central registry for all active WebSocket connections. It utilizes thread-safe data structures to track sessions by a unique identifier, typically a session_id or client_id.

Component	Responsibility
WsManager	Maintains a map of active connections and provides methods for broadcasting or targeted messaging.
WsHandler	Handles the low-level read/write loops for an individual socket.
WsMessageService	Orchestrates business logic by processing incoming messages and dispatching responses.

Sources: crates/infra/src/websocket/manager.rs, crates/core-service/src/service/ws_message.rs

Message Processing and Routing

Incoming messages are serialized and deserialized using JSON. The system uses a specific message format that identifies the "action" to be performed and contains the associated payload.

Message Sequence Diagram

The following sequence diagram demonstrates how an incoming message is received, processed by the core service, and how a response is sent back to the client.

The WsMessageService acts as the primary router for WebSocket-based business logic. It matches the action field in the incoming JSON against a set of handlers for file system operations, git commands, and terminal interactions.
Sources: crates/core-service/src/service/ws_message.rs, crates/infra/src/websocket/handler.rs

Data Structures

The system relies on structured message types to ensure type safety across the asynchronous boundary.

// Representative structure based on message handling logic
pub struct WsRequest {
    pub action: String,
    pub payload: Value,
    pub request_id: Option<String>,
}

pub struct WsResponse {
    pub status: String,
    pub data: Value,
    pub request_id: Option<String>,
}

Citations: crates/core-service/src/service/ws_message.rs:10-50

Infrastructure Implementation

The infrastructure layer provides the heavy lifting for socket I/O and session persistence.

Connection Registry Logic

The manager.rs file implements the WsManager struct, which often wraps a HashMap protected by an Arc<RwLock<...>> or uses channels to communicate with per-connection tasks. This allows the system to scale to many concurrent users while providing a simple API for other services to "push" data to specific users.
Sources: crates/infra/src/websocket/manager.rs

Error Handling and Disconnection

The system monitors connection health through heartbeats and handles abrupt disconnections by cleaning up the registry. If a socket closes, the WsHandler notifies the WsManager to remove the session ID, ensuring that resources are not leaked.
Sources: crates/infra/src/websocket/handler.rs, crates/infra/src/websocket/manager.rs

Summary of Key Source Files

File	Purpose
apps/api/src/api/ws/mod.rs	Defines the API routes and handles the initial HTTP to WebSocket upgrade.
crates/infra/src/websocket/handler.rs	Contains the low-level logic for reading from and writing to the WebSocket stream.
crates/infra/src/websocket/manager.rs	Manages the global state of active connections and session tracking.
crates/core-service/src/service/ws_message.rs	Implements the business logic dispatcher for incoming WebSocket requests.
crates/infra/src/websocket/mod.rs	Orchestrates the websocket module exports and initialization.

The WebSocket Real-time System is a foundational component of the ATMOS project, providing the necessary infrastructure for interactive, low-latency features. By decoupling connection management from business logic through the WsManager and WsMessageService, the system maintains high modularity and scalability.