brokkr-broker::ws::registry

Rust

Per-agent WebSocket connection registry.

Tracks which agents currently hold an open WebSocket connection to this broker process and routes outgoing messages to them on one of two priority lanes. Control-plane messages (work orders, target/stack changes, heartbeat) use the control lane; logs and kube events use the telemetry lane. The per-connection writer task drains control first (see handler).

Connections are keyed by agent_id. If a second connection arrives for the same agent (reconnect, duplicate process), it replaces the existing handle — dropping the prior senders causes the old writer task to wind down cleanly.

Structs

struct ConnectionHandle

pub

Derives: Clone

Sender-side handle for a single registered connection.

Owned by the registry. Cloned cheaply (just Arcs and mpsc senders). The corresponding receiver halves are held by the connection’s writer task — dropping all senders for a lane signals the writer to exit.

Fields

Name	Type	Description
agent_id	Uuid
connected_since	DateTime<Utc>
messages_in	Arc<AtomicU64>
messages_out	Arc<AtomicU64>

struct ConnectionInfo

pub

Derives: Debug, Clone

Snapshot view of one connection for diagnostics endpoints (WS-13).

Fields

Name	Type	Description
agent_id	Uuid
connected_since	DateTime<Utc>
messages_in	u64
messages_out	u64

struct ConnectionRegistry

pub

Derives: Default

Per-broker-process registry of live agent connections.

Wrap in Arc and share via axum::Extension so handlers and future broker-push code (WS-04) can route messages without depending on a specific app-state struct.

Methods

new

pub

fn new() -> Arc<Self>

Source

    pub fn new() -> Arc<Self> {
        Arc::new(Self::default())
    }

register

pub

fn register(&self, handle: ConnectionHandle)

Insert a new handle, evicting any prior connection for the same agent.

Source

    pub fn register(&self, handle: ConnectionHandle) {
        let agent_id = handle.agent_id;
        let mut map = self.inner.write().expect("registry poisoned");
        map.insert(agent_id, handle);
    }

unregister_​if_​matches

pub

fn unregister_if_matches(
    &self,
    agent_id: Uuid,
    connected_since: DateTime<Utc>,
)

Remove the handle iff it still matches the writer’s connected_since timestamp. This prevents a late-cleaning writer from unregistering a fresh reconnect that happens to share its agent id.

Source

    pub fn unregister_if_matches(&self, agent_id: Uuid, connected_since: DateTime<Utc>) {
        let mut map = self.inner.write().expect("registry poisoned");
        if let Some(existing) = map.get(&agent_id)
            && existing.connected_since == connected_since
        {
            map.remove(&agent_id);
        }
    }

is_​connected

pub

fn is_connected(&self, agent_id: Uuid) -> bool

True if any handle is registered for this agent.

Source

    pub fn is_connected(&self, agent_id: Uuid) -> bool {
        self.inner
            .read()
            .expect("registry poisoned")
            .contains_key(&agent_id)
    }

send_​control

pub

fn send_control(
    &self) -> Result<(,
    agent_id: Uuid) -> Result<(,
    msg: WsMessage) -> Result<(,
) -> Result<(), SendError>

Send a control-plane message to a specific agent. Non-blocking; returns immediately with SendError::LaneUnavailable if the lane queue is full (back-pressure). Use this for WorkOrder, TargetChanged, StackChanged, Heartbeat.

Source

    pub fn send_control(&self, agent_id: Uuid, msg: WsMessage) -> Result<(), SendError> {
        let handle = self
            .inner
            .read()
            .expect("registry poisoned")
            .get(&agent_id)
            .cloned()
            .ok_or(SendError::NotConnected(agent_id))?;
        handle
            .control_tx
            .try_send(msg)
            .map_err(|_| SendError::LaneUnavailable(agent_id))
    }

send_​telemetry

pub

fn send_telemetry(
    &self) -> Result<(,
    agent_id: Uuid) -> Result<(,
    msg: WsMessage) -> Result<(,
) -> Result<(), SendError>

Send a telemetry/log message to a specific agent on the low-priority lane. Use this for K8sEvent, PodLogLine, LogGap when the broker originates fan-out (which it does not in v1, but the lane exists for symmetry and for future broker→agent telemetry).

Source

    pub fn send_telemetry(&self, agent_id: Uuid, msg: WsMessage) -> Result<(), SendError> {
        let handle = self
            .inner
            .read()
            .expect("registry poisoned")
            .get(&agent_id)
            .cloned()
            .ok_or(SendError::NotConnected(agent_id))?;
        handle
            .telemetry_tx
            .try_send(msg)
            .map_err(|_| SendError::LaneUnavailable(agent_id))
    }

snapshot

pub

fn snapshot(&self) -> Vec<ConnectionInfo>

Snapshot every connection for diagnostics. O(n) clone — only call from admin/diagnostics paths.

Source

    pub fn snapshot(&self) -> Vec<ConnectionInfo> {
        self.inner
            .read()
            .expect("registry poisoned")
            .values()
            .map(|h| ConnectionInfo {
                agent_id: h.agent_id,
                connected_since: h.connected_since,
                messages_in: h.messages_in.load(Ordering::Relaxed),
                messages_out: h.messages_out.load(Ordering::Relaxed),
            })
            .collect()
    }

connected_​count

pub

fn connected_count(&self) -> usize

Number of connected agents (cheap; no clone).

Source

    pub fn connected_count(&self) -> usize {
        self.inner.read().expect("registry poisoned").len()
    }

close_​for_​agent

pub

fn close_for_agent(&self, agent_id: Uuid) -> usize

Forcibly close any live connection for agent_id, returning how many were closed (0 or 1 today — the map is one-connection-per-agent, since register evicts a prior handle for the same id).

Removing the handle drops its lane senders; the per-connection writer task then observes both lanes closed, closes the socket, and the connection’s run_connection loop unwinds (decrementing the connected gauge). This is the teardown path used by PAK revocation ([[BROKKR-T-0176]]): once an agent’s PAK is invalidated we must not leave its already-upgraded socket open until its TCP layer happens to notice. The agent’s reconnect will re-hit auth and be rejected.

Source

    pub fn close_for_agent(&self, agent_id: Uuid) -> usize {
        let mut map = self.inner.write().expect("registry poisoned");
        if map.remove(&agent_id).is_some() {
            1
        } else {
            0
        }
    }

Enums

enum SendError

pub

Errors returned when trying to push a message to a registered agent.

Variants

• NotConnected - No connection exists for the given agent id.
• LaneUnavailable - The connection exists but its lane is full or the writer has gone away. In practice this means the agent is either too slow or dying; callers should treat it the same as “not connected” and let the REST polling fallback surface the change on the next tick.