feat: realtime WebSocket layer for live events and streaming logs
- Backend: ConnectionManager, Redis-backed bus with in-process fallback, channel authorization, and typed emit helpers wired into workers (ai_runtime, sandbox_gateway), workforce/sandbox/billing/audit/error services. - Frontend: single shared WS client with backoff + heartbeat, RealtimeProvider, useLiveChannel/useLiveEvent hooks, live indicators on task detail, manager chat, dashboard, traces, sandbox terminal, admin errors, and billing. - Contracts: WsEvent/WsServerFrame/WsClientFrame types mirroring the port. All 255 backend tests pass; web typecheck and build pass.
This commit is contained in:
mjalmoudy94
@@ -0,0 +1,504 @@
|
||||
"""Realtime push contract: the event envelope and the typed publishers.
|
||||
|
||||
This module is an **application port**. Services depend on it; the WebSocket
|
||||
adapter in ``app.interfaces.realtime`` implements it. That direction matters:
|
||||
the run engine must be able to publish without knowing that a browser exists,
|
||||
and a worker process that never serves a socket still publishes the same way.
|
||||
|
||||
Two parts:
|
||||
|
||||
* **The envelope** (:class:`RealtimeEvent`) and the channel naming scheme.
|
||||
Framework-free, like every other port.
|
||||
* **The typed publishers** — ``emit_run_step`` and friends. They own the
|
||||
fan-out channels and the one rule that makes live feeds safe to call from
|
||||
anywhere: *publishing never raises*. A feed must not be able to fail a task
|
||||
run, a sandbox command or a billing write.
|
||||
|
||||
The publisher is process-wide state, set once by a composition root
|
||||
(``app.main`` for the API, ``app.workers.__main__`` for the workers). Until
|
||||
then it is a :class:`NullPublisher` — publishing is a no-op rather than a
|
||||
missing-import crash.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from collections.abc import Iterable, Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from typing import Protocol, runtime_checkable
|
||||
|
||||
from app.domain.base import new_id, utcnow
|
||||
|
||||
logger = logging.getLogger("agy.realtime")
|
||||
|
||||
#: The platform-owner console stream. Authorized against ``is_platform_owner``.
|
||||
ADMIN_CHANNEL = "admin"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Channels
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def org_channel(organization_id: str) -> str:
|
||||
"""Everything happening inside one tenant."""
|
||||
return f"org:{organization_id}"
|
||||
|
||||
|
||||
def task_channel(task_id: str) -> str:
|
||||
return f"task:{task_id}"
|
||||
|
||||
|
||||
def run_channel(run_id: str) -> str:
|
||||
return f"run:{run_id}"
|
||||
|
||||
|
||||
def thread_channel(thread_id: str) -> str:
|
||||
return f"thread:{thread_id}"
|
||||
|
||||
|
||||
def session_channel(session_id: str) -> str:
|
||||
"""One sandbox session — the live command terminal subscribes here."""
|
||||
return f"session:{session_id}"
|
||||
|
||||
|
||||
def user_channel(user_id: str) -> str:
|
||||
return f"user:{user_id}"
|
||||
|
||||
|
||||
def admin_channel() -> str:
|
||||
return ADMIN_CHANNEL
|
||||
|
||||
|
||||
def dedupe_channels(channels: Iterable[str]) -> tuple[str, ...]:
|
||||
"""Preserve order, drop duplicates — a publisher may overlap channels."""
|
||||
seen: dict[str, None] = {}
|
||||
for channel in channels:
|
||||
if channel and channel not in seen:
|
||||
seen[channel] = None
|
||||
return tuple(seen)
|
||||
|
||||
|
||||
def with_org_channel(channels: Sequence[str], organization_id: str | None) -> tuple[str, ...]:
|
||||
"""Append the tenant channel unless the caller already listed it."""
|
||||
if not organization_id:
|
||||
return dedupe_channels(channels)
|
||||
return dedupe_channels([*channels, org_channel(organization_id)])
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Envelope
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass(slots=True, kw_only=True)
|
||||
class RealtimeEvent:
|
||||
"""One push to one or more channels.
|
||||
|
||||
``channels`` is the fan-out list — the same event typically goes to the
|
||||
narrow channel (``run:…``) and the wide one (``org:…``), so a page watching
|
||||
a single run and a dashboard watching the whole tenant both stay live
|
||||
without the publisher caring who is listening.
|
||||
"""
|
||||
|
||||
type: str
|
||||
channels: tuple[str, ...]
|
||||
data: dict = field(default_factory=dict)
|
||||
organization_id: str | None = None
|
||||
id: str = field(default_factory=new_id)
|
||||
at: datetime = field(default_factory=utcnow)
|
||||
|
||||
def envelope(self) -> dict:
|
||||
"""The JSON object a browser receives.
|
||||
|
||||
``at`` is serialised as an ISO-8601 UTC string: the client renders it
|
||||
with ``toLocaleString`` like every other timestamp in the product.
|
||||
"""
|
||||
return {
|
||||
"id": self.id,
|
||||
"type": self.type,
|
||||
"data": self.data,
|
||||
"organization_id": self.organization_id,
|
||||
"at": self.at.isoformat(),
|
||||
}
|
||||
|
||||
def to_json(self) -> str:
|
||||
return json.dumps(self.envelope(), ensure_ascii=False, default=str)
|
||||
|
||||
@classmethod
|
||||
def from_json(cls, raw: str | bytes) -> RealtimeEvent:
|
||||
payload = json.loads(raw)
|
||||
return cls(
|
||||
type=str(payload["type"]),
|
||||
channels=tuple(payload.get("channels") or ()),
|
||||
data=dict(payload.get("data") or {}),
|
||||
organization_id=payload.get("organization_id"),
|
||||
id=str(payload["id"]),
|
||||
at=_parse_at(payload.get("at")),
|
||||
)
|
||||
|
||||
|
||||
def _parse_at(value: object) -> datetime:
|
||||
if isinstance(value, str) and value:
|
||||
try:
|
||||
return datetime.fromisoformat(value)
|
||||
except ValueError:
|
||||
return utcnow()
|
||||
return utcnow()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Publisher plumbing
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class EventPublisher(Protocol):
|
||||
async def publish(self, event: RealtimeEvent) -> None: ...
|
||||
|
||||
|
||||
class NullPublisher:
|
||||
"""Default before a composition root wires the real bus."""
|
||||
|
||||
async def publish(self, event: RealtimeEvent) -> None:
|
||||
return None
|
||||
|
||||
|
||||
_PUBLISHER: EventPublisher = NullPublisher()
|
||||
|
||||
|
||||
def set_publisher(publisher: EventPublisher) -> None:
|
||||
"""Install the process-wide publisher. Called by composition roots only."""
|
||||
global _PUBLISHER
|
||||
_PUBLISHER = publisher
|
||||
|
||||
|
||||
def get_publisher() -> EventPublisher:
|
||||
return _PUBLISHER
|
||||
|
||||
|
||||
async def publish(event: RealtimeEvent) -> None:
|
||||
"""Publish through the installed publisher. Never raises."""
|
||||
try:
|
||||
await _PUBLISHER.publish(event)
|
||||
except Exception: # noqa: BLE001 - the feed must never fail a caller
|
||||
logger.exception("realtime: publish failed for %s", event.type)
|
||||
|
||||
|
||||
async def _emit(
|
||||
*,
|
||||
type: str,
|
||||
channels: tuple[str, ...],
|
||||
data: dict,
|
||||
organization_id: str | None = None,
|
||||
) -> None:
|
||||
await publish(
|
||||
RealtimeEvent(
|
||||
type=type,
|
||||
channels=with_org_channel(channels, organization_id),
|
||||
data=data,
|
||||
organization_id=organization_id,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Typed publishers
|
||||
#
|
||||
# Payload shape is the product contract with the browser: the client-side types
|
||||
# in ``packages/contracts`` mirror ``data`` key-for-key. Keep the two in sync.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
#: Long text fields are clipped in the live feed. The full record stays
|
||||
#: available through the REST endpoints — the feed is a signal, not an archive.
|
||||
MAX_LIVE_TEXT_CHARS = 8000
|
||||
|
||||
|
||||
def _clip(text: str) -> str:
|
||||
return text[:MAX_LIVE_TEXT_CHARS]
|
||||
|
||||
|
||||
async def emit_run_step(
|
||||
*,
|
||||
run_id: str,
|
||||
task_id: str,
|
||||
organization_id: str,
|
||||
step_id: str,
|
||||
index: int,
|
||||
kind: str,
|
||||
title: str,
|
||||
payload: dict | None = None,
|
||||
) -> None:
|
||||
await _emit(
|
||||
type="run.step",
|
||||
channels=(run_channel(run_id), task_channel(task_id)),
|
||||
organization_id=organization_id,
|
||||
data={
|
||||
"run_id": run_id,
|
||||
"task_id": task_id,
|
||||
"step": {
|
||||
"id": step_id,
|
||||
"index": index,
|
||||
"kind": kind,
|
||||
"title": title,
|
||||
"payload": payload or {},
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
async def emit_run_status(
|
||||
*,
|
||||
run_id: str,
|
||||
task_id: str,
|
||||
organization_id: str,
|
||||
status: str,
|
||||
error: str | None = None,
|
||||
) -> None:
|
||||
await _emit(
|
||||
type="run.status",
|
||||
channels=(run_channel(run_id), task_channel(task_id)),
|
||||
organization_id=organization_id,
|
||||
data={"run_id": run_id, "task_id": task_id, "status": status, "error": error},
|
||||
)
|
||||
|
||||
|
||||
async def emit_task_status(*, task_id: str, organization_id: str, status: str) -> None:
|
||||
await _emit(
|
||||
type="task.status",
|
||||
channels=(task_channel(task_id),),
|
||||
organization_id=organization_id,
|
||||
data={"task_id": task_id, "status": status},
|
||||
)
|
||||
|
||||
|
||||
async def emit_task_created(
|
||||
*, task_id: str, organization_id: str, title: str, status: str
|
||||
) -> None:
|
||||
await _emit(
|
||||
type="task.created",
|
||||
channels=(task_channel(task_id),),
|
||||
organization_id=organization_id,
|
||||
data={"task_id": task_id, "title": title, "status": status},
|
||||
)
|
||||
|
||||
|
||||
async def emit_chat_pending(*, thread_id: str, organization_id: str) -> None:
|
||||
"""The manager is thinking — the client shows a typing indicator."""
|
||||
await _emit(
|
||||
type="chat.pending",
|
||||
channels=(thread_channel(thread_id),),
|
||||
organization_id=organization_id,
|
||||
data={"thread_id": thread_id},
|
||||
)
|
||||
|
||||
|
||||
async def emit_message_created(
|
||||
*,
|
||||
thread_id: str,
|
||||
organization_id: str,
|
||||
message_id: str,
|
||||
role: str,
|
||||
content: str,
|
||||
employee_id: str | None = None,
|
||||
) -> None:
|
||||
await _emit(
|
||||
type="message.created",
|
||||
channels=(thread_channel(thread_id),),
|
||||
organization_id=organization_id,
|
||||
data={
|
||||
"thread_id": thread_id,
|
||||
"message": {
|
||||
"id": message_id,
|
||||
"role": role,
|
||||
"content": _clip(content),
|
||||
"employee_id": employee_id,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
async def emit_approval_pending(
|
||||
*,
|
||||
run_id: str,
|
||||
task_id: str,
|
||||
organization_id: str,
|
||||
tool_call_id: str,
|
||||
tool_key: str,
|
||||
requires_approval: bool,
|
||||
) -> None:
|
||||
await _emit(
|
||||
type="approval.pending",
|
||||
channels=(run_channel(run_id), task_channel(task_id)),
|
||||
organization_id=organization_id,
|
||||
data={
|
||||
"run_id": run_id,
|
||||
"task_id": task_id,
|
||||
"tool_call_id": tool_call_id,
|
||||
"tool_key": tool_key,
|
||||
"requires_approval": requires_approval,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
async def emit_artifact_created(
|
||||
*, run_id: str, task_id: str, organization_id: str, artifact_id: str, name: str
|
||||
) -> None:
|
||||
await _emit(
|
||||
type="artifact.created",
|
||||
channels=(run_channel(run_id), task_channel(task_id)),
|
||||
organization_id=organization_id,
|
||||
data={
|
||||
"run_id": run_id,
|
||||
"task_id": task_id,
|
||||
"artifact_id": artifact_id,
|
||||
"name": name,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
async def emit_sandbox_command(
|
||||
*,
|
||||
session_id: str,
|
||||
organization_id: str,
|
||||
command_run_id: str,
|
||||
command: str,
|
||||
exit_code: int | None,
|
||||
stdout: str,
|
||||
stderr: str,
|
||||
duration_ms: int,
|
||||
truncated: bool,
|
||||
status: str = "finished",
|
||||
) -> None:
|
||||
"""One command's output, pushed live to the session's terminal panel."""
|
||||
await _emit(
|
||||
type="sandbox.command",
|
||||
channels=(session_channel(session_id),),
|
||||
organization_id=organization_id,
|
||||
data={
|
||||
"session_id": session_id,
|
||||
"run_id": command_run_id,
|
||||
"command": _clip(command),
|
||||
"status": status,
|
||||
"exit_code": exit_code,
|
||||
"stdout": _clip(stdout),
|
||||
"stderr": _clip(stderr),
|
||||
"duration_ms": duration_ms,
|
||||
"truncated": truncated,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
async def emit_sandbox_session(
|
||||
*,
|
||||
session_id: str,
|
||||
organization_id: str,
|
||||
status: str,
|
||||
workspace_id: str | None = None,
|
||||
) -> None:
|
||||
await _emit(
|
||||
type="sandbox.session",
|
||||
channels=(session_channel(session_id),),
|
||||
organization_id=organization_id,
|
||||
data={
|
||||
"session_id": session_id,
|
||||
"status": status,
|
||||
"workspace_id": workspace_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
async def emit_sandbox_reaped(
|
||||
*, organization_id: str | None, session_id: str, workspace_id: str
|
||||
) -> None:
|
||||
await _emit(
|
||||
type="sandbox.reaped",
|
||||
channels=(session_channel(session_id),),
|
||||
organization_id=organization_id,
|
||||
data={"session_id": session_id, "workspace_id": workspace_id},
|
||||
)
|
||||
|
||||
|
||||
async def emit_usage_recorded(
|
||||
*,
|
||||
organization_id: str,
|
||||
unit: str,
|
||||
quantity: float,
|
||||
run_id: str | None = None,
|
||||
employee_id: str | None = None,
|
||||
) -> None:
|
||||
await _emit(
|
||||
type="usage.recorded",
|
||||
channels=(),
|
||||
organization_id=organization_id,
|
||||
data={
|
||||
"unit": unit,
|
||||
"quantity": quantity,
|
||||
"run_id": run_id,
|
||||
"employee_id": employee_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
async def emit_wallet_changed(*, organization_id: str, balance: float) -> None:
|
||||
await _emit(
|
||||
type="billing.wallet",
|
||||
channels=(),
|
||||
organization_id=organization_id,
|
||||
data={"organization_id": organization_id, "balance": balance},
|
||||
)
|
||||
|
||||
|
||||
async def emit_error_recorded(
|
||||
*,
|
||||
error_id: str | None,
|
||||
kind: str,
|
||||
message: str,
|
||||
severity: str,
|
||||
source: str,
|
||||
organization_id: str | None = None,
|
||||
) -> None:
|
||||
"""Feed the platform-owner console; tenants see their own through the org
|
||||
channel the envelope carries."""
|
||||
await _emit(
|
||||
type="error.recorded",
|
||||
channels=("admin",),
|
||||
organization_id=organization_id,
|
||||
data={
|
||||
"error_id": error_id,
|
||||
"kind": kind,
|
||||
"message": _clip(message),
|
||||
"severity": severity,
|
||||
"source": source,
|
||||
"organization_id": organization_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
async def emit_audit_recorded(
|
||||
*, organization_id: str | None, action: str, target_type: str, target_id: str
|
||||
) -> None:
|
||||
if not organization_id:
|
||||
return
|
||||
await _emit(
|
||||
type="audit.recorded",
|
||||
channels=(),
|
||||
organization_id=organization_id,
|
||||
data={"action": action, "target_type": target_type, "target_id": target_id},
|
||||
)
|
||||
|
||||
|
||||
async def emit_worker_status(
|
||||
*,
|
||||
organization_id: str,
|
||||
worker: str,
|
||||
message: str,
|
||||
) -> None:
|
||||
await _emit(
|
||||
type="worker.status",
|
||||
channels=(),
|
||||
organization_id=organization_id,
|
||||
data={"worker": worker, "message": _clip(message)},
|
||||
)
|
||||
Reference in New Issue
Block a user