Skip to main content

streams.loop

WebSocket pub/sub infrastructure and monitoring snapshot builder.

This module provides the generic streaming loop used by the agent-runtimes WebSocket endpoints. It was extracted from routes/tool_approvals.py so that the common logic (subscriber management, snapshot assembly, stream loop) is decoupled from the tool-approval CRUD routes.

subscribe_stream​

def subscribe_stream(
        agent_id: str | None) -> asyncio.Queue[AgentStreamMessage]

Register a subscriber queue for agent_id (or global).

unsubscribe_stream​

def unsubscribe_stream(agent_id: str | None,
                       queue: asyncio.Queue[AgentStreamMessage]) -> None

Remove a previously registered subscriber queue.

enqueue_stream_message​

def enqueue_stream_message(agent_id: str | None,
                           message: AgentStreamMessage) -> None

Push message to every queue subscribed to agent_id.

build_context_snapshot​

def build_context_snapshot(agent_id: str) -> dict[str, Any] | None

Build the lightweight context snapshot dict for agent_id.

build_full_context​

def build_full_context(agent_id: str) -> dict[str, Any] | None

Build the full context snapshot (tools, messages, model config).

build_mcp_status​

def build_mcp_status() -> dict[str, Any] | None

Build MCP toolsets status dict.

get_codemode_mcp_tools_detailed​

def get_codemode_mcp_tools_detailed(
        agent_id: str | None = None) -> dict[str, list[dict[str, Any]]]

Return codemode-discovered MCP tools grouped by server for display.

Each tool dict carries name (bare tool name), description, enabled and inputSchema keys so listing commands and the /mcp/servers route can surface codemode-managed servers, which the MCP lifecycle manager does not track.

get_agent_mcp_enabled_tools_by_server​

def get_agent_mcp_enabled_tools_by_server(
        agent_id: str | None) -> dict[str, list[str]]

Public accessor for effective MCP enabled tools by server.

get_agent_approved_mcp_tool_names​

def get_agent_approved_mcp_tool_names(agent_id: str | None) -> set[str]

Return the set of approved MCP tool names for an agent.

get_agent_enabled_mcp_tool_names​

def get_agent_enabled_mcp_tool_names(agent_id: str | None) -> set[str]

Return the set of enabled MCP tool names for an agent.

has_agent_mcp_tool_selection​

def has_agent_mcp_tool_selection(agent_id: str | None) -> bool

Return True when an explicit MCP tool selection exists for an agent.

This indicates the user/client has intentionally sent a selection update (for example via builtinTools or monitoring websocket controls). When False, callers may seed a sensible default selection for first-turn flows.

get_known_mcp_tool_names​

def get_known_mcp_tool_names() -> set[str]

Return all known MCP tool names currently discovered by lifecycle manager.

get_agent_skills_snapshot​

def get_agent_skills_snapshot(agent_id: str | None) -> list[dict[str, Any]]

Return serialized skills state for an agent.

get_agent_enabled_skill_ids​

def get_agent_enabled_skill_ids(agent_id: str | None) -> set[str]

Return enabled (enabled/loaded) skill IDs for an agent.

get_agent_tracked_skill_ids​

def get_agent_tracked_skill_ids(agent_id: str | None) -> set[str]

Return all skill IDs tracked for an agent (available + enabled + loaded).

These are the skills declared by the agent's spec. The guardrail uses this as the "in-scope" set — anything in it is allowed to run pending user approval; anything outside is rejected as unknown.

set_agent_enabled_skills​

def set_agent_enabled_skills(agent_id: str | None,
                             skill_refs: list[str]) -> list[dict[str, Any]]

Set enabled skills for an agent (single source of truth).

set_agent_turn_enabled_skills​

def set_agent_turn_enabled_skills(
        agent_id: str | None, skill_refs: list[str]) -> list[dict[str, Any]]

Apply per-turn skill enablement without changing tracked skill scope.

get_agent_approved_skill_ids​

def get_agent_approved_skill_ids(agent_id: str | None) -> set[str]

Return approved skill IDs for an agent.

mark_agent_mcp_tool_approved​

def mark_agent_mcp_tool_approved(agent_id: str | None, tool_name: str) -> bool

Mark an MCP tool as approved for an agent.

Looks up which MCP server hosts the tool and adds it to the agent's approved set so the dropdown reflects the new state. Returns True if the tool was matched to a known MCP server.

Accepts both bare tool names (tavily_extract) and qualified names (tavily__tavily_extract).

set_agent_enabled_mcp_tool_names​

def set_agent_enabled_mcp_tool_names(
        agent_id: str | None, tool_names: list[str]) -> dict[str, list[str]]

Set enabled MCP tools from a flat per-turn tool-name list.

This is used by request transports that only receive a flat list of tool names for a turn (e.g. builtinTools in chat payloads). The mapping is projected back to per-server enabled tool sets.

enable_agent_skill​

def enable_agent_skill(agent_id: str | None, skill_ref: str) -> None

Enable one skill for an agent.

disable_agent_skill​

def disable_agent_skill(agent_id: str | None, skill_ref: str) -> None

Disable one skill for an agent.

approve_agent_skill​

def approve_agent_skill(agent_id: str | None, skill_ref: str) -> None

Mark one skill as approved for an agent.

unapprove_agent_skill​

def unapprove_agent_skill(agent_id: str | None, skill_ref: str) -> None

Mark one skill as unapproved for an agent.

purge_agent_stream_state​

def purge_agent_stream_state(agent_id: str | None) -> None

Purge per-agent in-memory stream state for skills/MCP/subscribers.

build_codemode_status​

def build_codemode_status(
        agent_id: str | None = None) -> dict[str, Any] | None

Build codemode status dict (sync — safe to call from async context).

Uses the stream-layer per-agent skills state. Each skill in the skills list carries a status field (available | enabled | loaded).

build_monitoring_snapshot_payload​

async def build_monitoring_snapshot_payload(
        agent_id: str | None,
        *,
        list_approvals: Any | None = None) -> AgentMonitoringSnapshotPayload

Assemble a full monitoring snapshot for agent_id.

Parameters​

list_approvals : callable, optional An async callable (agent_id, status) -> list[Record] that returns the current pending tool approvals. When None the approval fields are left empty (useful when the caller does not have access to the approval store).

publish_stream_event​

async def publish_stream_event(*,
                               event_type: str,
                               payload: dict[str, Any],
                               agent_id: str | None,
                               list_approvals: Any | None = None) -> None

Publish a stream event and follow it with a fresh snapshot.

stream_loop​

async def stream_loop(
        websocket: WebSocket,
        agent_id: str | None,
        *,
        list_approvals: Any | None = None,
        decide_approval: Callable[[str, bool, str | None, str | None],
                                  Awaitable[Any]]
    | None = None,
        delete_approval: Callable[[str], Awaitable[Any]] | None = None
) -> None

Run the main WebSocket stream loop.

Accepts the connection, streams an initial snapshot, then pushes periodic diffs and reactive messages until the client disconnects.