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
2026-10-13 12:56:36 +03:30
parent afdef4b444
commit f1e5df29ce
31 changed files with 2981 additions and 28 deletions
+504
View File
@@ -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)},
)
@@ -8,6 +8,7 @@ from __future__ import annotations
from sqlalchemy.ext.asyncio import AsyncSession
from app.application.ports.realtime import emit_audit_recorded
from app.domain.audit.models import AuditEvent
from app.infrastructure.repositories.audit import AuditRepository
@@ -34,4 +35,13 @@ async def record_audit(
ip_address=ip,
user_agent=user_agent,
)
return await AuditRepository(session).record(event=event)
recorded = await AuditRepository(session).record(event=event)
# Live push: the audit timelines (dashboard, traces) update as actions
# happen instead of waiting for the next reload.
await emit_audit_recorded(
organization_id=organization_id,
action=action,
target_type=target_type,
target_id=target_id,
)
return recorded
@@ -14,6 +14,7 @@ from __future__ import annotations
from datetime import datetime
from app.application.ports import PaymentProviderPort
from app.application.ports.realtime import emit_wallet_changed
from app.core.errors import NotFoundError, PermissionDeniedError, ValidationError
from app.domain.billing.models import (
CreditWallet,
@@ -121,6 +122,11 @@ class BillingService:
# Refuse to leave a wallet in debt; undo and surface the conflict.
await self.wallets.adjust(organization_id=organization_id, delta=-amount)
raise ValidationError("اعتبار کافی برای این کسر وجود ندارد.")
# Live push: the sidebar credit badge and the billing page follow the
# wallet as it moves, whoever moved it.
await emit_wallet_changed(
organization_id=organization_id, balance=wallet.balance
)
return wallet
async def start_checkout(
@@ -177,8 +183,13 @@ class BillingService:
return payment
if succeeded:
await self.wallets.adjust(organization_id=organization_id, delta=payment.amount)
wallet = await self.wallets.adjust(
organization_id=organization_id, delta=payment.amount
)
final = await self.payments.set_status(payment_id, PaymentStatus.SUCCEEDED)
await emit_wallet_changed(
organization_id=organization_id, balance=wallet.balance
)
else:
final = await self.payments.set_status(payment_id, PaymentStatus.FAILED)
return final or payment
@@ -18,6 +18,10 @@ import contextlib
from datetime import datetime
from app.application.ports import CommandResult, SandboxProviderPort, UsageRecord
from app.application.ports.realtime import (
emit_sandbox_command,
emit_sandbox_session,
)
from app.core.config import settings
from app.core.errors import ConflictError, NotFoundError, ProviderUnavailableError, ValidationError
from app.domain.ai_control.models import RateUnit
@@ -210,6 +214,12 @@ class SandboxService:
)
)
await self.workspaces.set_status(workspace.id, WorkspaceStatus.READY)
await emit_sandbox_session(
session_id=session.id,
organization_id=organization_id,
status=SessionStatus.RUNNING.value,
workspace_id=workspace.id,
)
return session
async def get_session_or_404(
@@ -229,6 +239,11 @@ class SandboxService:
await self.provider.destroy(session_external_id=session.external_session_id)
await self.sessions.set_status(session.id, SessionStatus.STOPPED, stopped=True)
refreshed = await self.sessions.get(session.id)
await emit_sandbox_session(
session_id=session.id,
organization_id=organization_id,
status=SessionStatus.STOPPED.value,
)
return refreshed or session
# -- commands ------------------------------------------------------------
@@ -285,6 +300,20 @@ class SandboxService:
)
await self.sessions.set_status(session.id, SessionStatus.RUNNING)
# Live push: the terminal panel on the project page streams the command
# and its output as the row is written, not on the next poll.
await emit_sandbox_command(
session_id=session.id,
organization_id=organization_id,
command_run_id=run.id,
command=cleaned,
exit_code=result.exit_code,
stdout=run.stdout_ref,
stderr=run.stderr_ref,
duration_ms=result.duration_ms,
truncated=result.truncated,
)
# CPU seconds are billed. 1 second of wall time ~= 1 CPU second at MVP.
cpu_seconds = max(0.0, result.duration_ms / 1000.0)
if cpu_seconds > 0:
@@ -24,6 +24,17 @@ from app.application.ports import (
ModelProviderPort,
UsageRecord,
)
from app.application.ports.realtime import (
emit_approval_pending,
emit_artifact_created,
emit_chat_pending,
emit_message_created,
emit_run_status,
emit_run_step,
emit_task_created,
emit_task_status,
emit_usage_recorded,
)
from app.core.errors import (
ConflictError,
NotFoundError,
@@ -324,6 +335,10 @@ class WorkforceService:
else:
employee = await self._default_manager(organization_id)
# The manager is about to call a model. Tell the thread so the UI can
# show a thinking indicator immediately.
await emit_chat_pending(thread_id=thread.id, organization_id=organization_id)
history = await self.messages.list_for_thread(thread.id)
system_prompt = await self._system_prompt_for(employee)
reply = await self._complete(
@@ -347,6 +362,16 @@ class WorkforceService:
employee_id=employee.id if employee else None,
)
)
# Live push: any browser on this thread sees the reply arrive, and the
# pending indicator clears, without waiting for a reload.
await emit_message_created(
thread_id=thread.id,
organization_id=organization_id,
message_id=assistant_message.id,
role=MessageRole.ASSISTANT.value,
content=reply,
employee_id=employee.id if employee else None,
)
return user_message, assistant_message
# ------------------------------------------------------------------ #
@@ -371,7 +396,7 @@ class WorkforceService:
await self.get_employee_or_404(
organization_id=organization_id, employee_id=assigned_employee_id
)
return await self.tasks.create(
task = await self.tasks.create(
Task(
organization_id=organization_id,
title=title.strip(),
@@ -382,6 +407,15 @@ class WorkforceService:
created_by_user_id=created_by_user_id,
)
)
# Live push: the tasks list and the manager-chat sidebar update the
# moment the task exists, whoever created it and from where.
await emit_task_created(
task_id=task.id,
organization_id=organization_id,
title=task.title,
status=task.status.value,
)
return task
async def list_tasks(
self,
@@ -435,21 +469,39 @@ class WorkforceService:
Run(task_id=task.id, employee_id=employee.id if employee else None)
)
await self.tasks.set_status(task.id, TaskStatus.RUNNING)
await emit_task_status(
task_id=task.id, organization_id=organization_id, status=TaskStatus.RUNNING.value
)
await self.run_steps.append(
run_id=run.id, kind="queued", title="کار در صف اجرا قرار گرفت.",
payload={"task_id": task.id},
await self._publish_step(
organization_id=organization_id,
run_id=run.id,
task_id=task.id,
step=await self.run_steps.append(
run_id=run.id, kind="queued", title="کار در صف اجرا قرار گرفت.",
payload={"task_id": task.id},
),
)
try:
run = await self._advance(run, RunStatus.PLANNING)
await self.run_steps.append(
run_id=run.id, kind="planning", title="برنامه‌ریزی مرحلهٔ اجرا."
await self._publish_step(
organization_id=organization_id,
run_id=run.id,
task_id=task.id,
step=await self.run_steps.append(
run_id=run.id, kind="planning", title="برنامه‌ریزی مرحلهٔ اجرا."
),
)
run = await self._advance(run, RunStatus.RUNNING)
await self.run_steps.append(
run_id=run.id, kind="running", title="اجرای کار آغاز شد."
await self._publish_step(
organization_id=organization_id,
run_id=run.id,
task_id=task.id,
step=await self.run_steps.append(
run_id=run.id, kind="running", title="اجرای کار آغاز شد."
),
)
reply = await self._complete(
@@ -463,9 +515,14 @@ class WorkforceService:
ChatMessage(role="user", content=self._task_brief(task)),
],
)
await self.run_steps.append(
run_id=run.id, kind="output", title="خروجی آماده شد.",
payload={"chars": len(reply)},
await self._publish_step(
organization_id=organization_id,
run_id=run.id,
task_id=task.id,
step=await self.run_steps.append(
run_id=run.id, kind="output", title="خروجی آماده شد.",
payload={"chars": len(reply)},
),
)
artifact = await self.artifacts.create(
@@ -479,30 +536,106 @@ class WorkforceService:
size_bytes=len(reply.encode("utf-8")),
)
)
await self.run_steps.append(
run_id=run.id, kind="artifact", title="سند خروجی ساخته شد.",
payload={"artifact_id": artifact.id},
await self._publish_step(
organization_id=organization_id,
run_id=run.id,
task_id=task.id,
step=await self.run_steps.append(
run_id=run.id, kind="artifact", title="سند خروجی ساخته شد.",
payload={"artifact_id": artifact.id},
),
)
await emit_artifact_created(
run_id=run.id,
task_id=task.id,
organization_id=organization_id,
artifact_id=artifact.id,
name=artifact.name,
)
run = await self._finish(run, RunStatus.COMPLETED)
await self.tasks.set_status(task.id, TaskStatus.DONE)
await emit_run_status(
run_id=run.id,
task_id=task.id,
organization_id=organization_id,
status=RunStatus.COMPLETED.value,
)
await emit_task_status(
task_id=task.id, organization_id=organization_id, status=TaskStatus.DONE.value
)
except ProviderUnavailableError as exc:
run = await self._finish(run, RunStatus.FAILED, error=str(exc))
await self.tasks.set_status(task.id, TaskStatus.FAILED)
await self.run_steps.append(
run_id=run.id, kind="error", title="اجرا ناموفق بود.",
payload={"error": str(exc)},
await self._publish_step(
organization_id=organization_id,
run_id=run.id,
task_id=task.id,
step=await self.run_steps.append(
run_id=run.id, kind="error", title="اجرا ناموفق بود.",
payload={"error": str(exc)},
),
)
await emit_run_status(
run_id=run.id,
task_id=task.id,
organization_id=organization_id,
status=RunStatus.FAILED.value,
error=str(exc),
)
await emit_task_status(
task_id=task.id, organization_id=organization_id, status=TaskStatus.FAILED.value
)
except Exception as exc: # noqa: BLE001 - the run must always terminate
run = await self._finish(run, RunStatus.FAILED, error=str(exc))
await self.tasks.set_status(task.id, TaskStatus.FAILED)
await self.run_steps.append(
run_id=run.id, kind="error", title="خطای غیرمنتظره در اجرا.",
payload={"error": str(exc)},
await self._publish_step(
organization_id=organization_id,
run_id=run.id,
task_id=task.id,
step=await self.run_steps.append(
run_id=run.id, kind="error", title="خطای غیرمنتظره در اجرا.",
payload={"error": str(exc)},
),
)
await emit_run_status(
run_id=run.id,
task_id=task.id,
organization_id=organization_id,
status=RunStatus.FAILED.value,
error=str(exc),
)
await emit_task_status(
task_id=task.id, organization_id=organization_id, status=TaskStatus.FAILED.value
)
return run
async def _publish_step(
self,
*,
organization_id: str,
run_id: str,
task_id: str,
step: RunStep,
) -> None:
"""Push one timeline step to the run, the task and the tenant feed.
The step is already persisted by the repository; this only broadcasts
it. Kept as a helper so the run engine above reads as the plain
sequence of transitions it is.
"""
await emit_run_step(
run_id=run_id,
task_id=task_id,
organization_id=organization_id,
step_id=step.id,
index=step.index,
kind=step.kind,
title=step.title,
payload=step.payload_json,
)
async def get_run_or_404(
self, *, organization_id: str, run_id: str
) -> tuple[Run, list[RunStep], list[ToolCall], list[Artifact]]:
@@ -530,12 +663,28 @@ class WorkforceService:
raise ConflictError("از وضعیت فعلی امکان لغو وجود ندارد.")
run = await self._finish(run, RunStatus.CANCELLED)
await self.run_steps.append(
run_id=run.id, kind="cancelled", title="اجرا توسط کاربر لغو شد."
await self._publish_step(
organization_id=organization_id,
run_id=run.id,
task_id=run.task_id,
step=await self.run_steps.append(
run_id=run.id, kind="cancelled", title="اجرا توسط کاربر لغو شد."
),
)
task = await self.tasks.get(run.task_id)
if task is not None and task.status not in (TaskStatus.DONE, TaskStatus.CANCELLED):
await self.tasks.set_status(task.id, TaskStatus.CANCELLED)
await emit_task_status(
task_id=task.id,
organization_id=organization_id,
status=TaskStatus.CANCELLED.value,
)
await emit_run_status(
run_id=run.id,
task_id=run.task_id,
organization_id=organization_id,
status=RunStatus.CANCELLED.value,
)
return run
async def list_pending_approvals(self, organization_id: str) -> list[ToolCall]:
@@ -627,6 +776,14 @@ class WorkforceService:
),
payload={"tool_call_id": call.id, "requires_approval": needs_approval},
)
await emit_approval_pending(
run_id=run.id,
task_id=run.task_id,
organization_id=organization_id,
tool_call_id=call.id,
tool_key=tool_key,
requires_approval=needs_approval,
)
return call
async def decide_approval(
@@ -655,12 +812,25 @@ class WorkforceService:
result=result,
approved_by_user_id=decided_by_user_id,
)
run = await self.runs.get(call.run_id)
await self.run_steps.append(
run_id=call.run_id,
kind="approval",
title=("درخواست تأیید شد." if approve else "درخواست رد شد."),
payload={"tool_call_id": call.id, "tool_key": call.tool_key},
)
# Live push: the task page clears the pending card and the run timeline
# gains the decision without a reload.
await emit_run_step(
run_id=call.run_id,
task_id=run.task_id if run else "",
organization_id=organization_id,
step_id=call.id,
index=0,
kind="approval",
title=("درخواست تأیید شد." if approve else "درخواست رد شد."),
payload={"tool_call_id": call.id, "tool_key": call.tool_key},
)
return updated or call
# ------------------------------------------------------------------ #
@@ -776,6 +946,15 @@ class WorkforceService:
metadata_json={"priced": card is not None},
)
)
# Live push: the billing page and the sidebar credit badge react to
# consumption as it happens instead of at the next page load.
await emit_usage_recorded(
organization_id=organization_id,
unit=unit.value if hasattr(unit, "value") else str(unit),
quantity=quantity,
run_id=run_id,
employee_id=employee_id,
)
async def _advance(self, run: Run, target: RunStatus) -> Run:
if not can_transition_run(run.status, target):
+9
View File
@@ -53,6 +53,15 @@ class Settings(BaseSettings):
database_url: str = "sqlite+aiosqlite:///./agy_web.db"
redis_url: str = "redis://localhost:6379/0"
# --- realtime push -------------------------------------------------------
#: How live events travel between processes. ``auto`` mirrors published
#: events onto Redis pub/sub so the standalone workers' pushes reach the
#: API process's browser sockets; if no broker is reachable the bus falls
#: back to in-process fan-out and every API-driven feed stays live.
#: ``local`` pins the bus to in-process delivery (tests, single-process
#: deployments without a broker).
realtime_transport: str = "auto"
# --- security ------------------------------------------------------------
jwt_secret: str = "change-me"
jwt_algorithm: str = "HS256"
+13
View File
@@ -255,6 +255,19 @@ async def record_error(
request_id or "-",
fingerprint,
)
# Live push to the platform-owner console. Best-effort like everything here:
# a failed publish must not turn a recorded error into an unrecorded one.
from app.application.ports.realtime import emit_error_recorded
await emit_error_recorded(
error_id=event.id,
kind=event.kind,
message=event.message,
severity=str(severity),
source=str(source),
organization_id=event.organization_id,
)
return event.id
@@ -42,5 +42,22 @@ async def readiness(session: AsyncSession = Depends(get_session)) -> dict:
"sandbox": settings.sandbox_provider,
"payment": settings.payment_provider,
}
# Live-feed visibility: how many browser sockets this process is serving,
# and whether the cross-process relay (Redis) is connected. An operator
# wondering "is realtime working?" should not have to read the logs.
try:
from app.interfaces.realtime.bus import get_bus
from app.interfaces.realtime.router import realtime_manager
manager = realtime_manager()
checks["realtime"] = {
"connections": manager.connection_count,
"channels": manager.channel_count,
"redis_relay": "connected" if get_bus().redis_healthy else "in_process_only",
}
except Exception: # noqa: BLE001 - health must answer regardless
checks["realtime"] = "unavailable"
healthy = checks.get("database") == "ok"
return {"status": "ok" if healthy else "degraded", "checks": checks}
@@ -0,0 +1,20 @@
"""Realtime push layer — the WebSocket *interface* adapter.
The adapter is transport and plumbing only: who is connected to which channel
(:mod:`manager`), how an event travels between processes (:mod:`bus`), which
channels a socket may subscribe to (:mod:`authz`), and the wire protocol
(:mod:`router`).
The event contract and the typed publishers live one layer in, in
``app.application.ports.realtime`` — services depend on the port, this package
implements it. ``app.main`` is the composition root that installs the bus as
the process publisher and starts the cross-process relay.
"""
from app.interfaces.realtime.bus import get_bus, start_relay, stop_relay
__all__ = [
"get_bus",
"start_relay",
"stop_relay",
]
+112
View File
@@ -0,0 +1,112 @@
"""Channel authorization — the only gate on what a socket may receive.
Authorization happens once, at subscribe time. After that the connection holds
plain channel names and the manager routes on them, so there is no code path
that can deliver an event to a connection that was never checked.
Every rule resolves the channel's tenant through the repositories and then the
caller's membership in it — exactly like the HTTP layer's ``get_org_context``.
A tenant that does not own the object gets a rejection, not an empty stream:
the difference matters when debugging "why is my page not live?".
The caller's identity arrives as two primitives (id, platform-owner flag) —
what an authenticated socket actually knows — rather than a ``User`` row, so a
subscribe costs one fewer database read.
"""
from __future__ import annotations
from sqlalchemy.ext.asyncio import AsyncSession
from app.infrastructure.repositories.identity import MembershipRepository
from app.infrastructure.repositories.sandbox import (
SandboxSessionRepository,
WorkspaceRepository,
)
from app.infrastructure.repositories.workforce import (
RunRepository,
TaskRepository,
ThreadRepository,
)
async def authorize_channel(
session: AsyncSession,
*,
user_id: str,
is_platform_owner: bool,
channel: str,
) -> bool:
"""May this identity subscribe to ``channel``? Deny by default."""
kind, _, identifier = channel.partition(":")
if not identifier and kind != "admin":
return False
if kind == "admin":
return is_platform_owner
if kind == "user":
return identifier == user_id
if kind == "org":
return await _is_member(
session=session, user_id=user_id, organization_id=identifier
)
if kind == "task":
return await _task_is_visible(
session=session, user_id=user_id, task_id=identifier
)
if kind == "run":
return await _run_is_visible(
session=session, user_id=user_id, run_id=identifier
)
if kind == "thread":
return await _thread_is_visible(
session=session, user_id=user_id, thread_id=identifier
)
if kind == "session":
return await _session_is_visible(
session=session, user_id=user_id, session_id=identifier
)
return False
async def _is_member(*, session: AsyncSession, user_id: str, organization_id: str) -> bool:
membership = await MembershipRepository(session).get(organization_id, user_id)
return membership is not None and membership.is_active
async def _task_is_visible(*, session: AsyncSession, user_id: str, task_id: str) -> bool:
task = await TaskRepository(session).get(task_id)
if task is None:
return False
return await _is_member(
session=session, user_id=user_id, organization_id=task.organization_id
)
async def _run_is_visible(*, session: AsyncSession, user_id: str, run_id: str) -> bool:
run = await RunRepository(session).get(run_id)
if run is None:
return False
return await _task_is_visible(session=session, user_id=user_id, task_id=run.task_id)
async def _thread_is_visible(*, session: AsyncSession, user_id: str, thread_id: str) -> bool:
thread = await ThreadRepository(session).get(thread_id)
if thread is None:
return False
return await _is_member(
session=session, user_id=user_id, organization_id=thread.organization_id
)
async def _session_is_visible(*, session: AsyncSession, user_id: str, session_id: str) -> bool:
"""A sandbox session belongs to a workspace, which belongs to a tenant."""
sandbox_session = await SandboxSessionRepository(session).get(session_id)
if sandbox_session is None:
return False
workspace = await WorkspaceRepository(session).get(sandbox_session.workspace_id)
if workspace is None:
return False
return await _is_member(
session=session, user_id=user_id, organization_id=workspace.organization_id
)
+260
View File
@@ -0,0 +1,260 @@
"""The event bus: how a published event reaches connected browsers.
Two transports, in this order:
1. **In-process.** Always on. When the API itself records the event (a task run
driven by a button, a sandbox command, an error handler), the fan-out is a
direct call into the ``ConnectionManager`` — microseconds, no broker.
2. **Redis pub/sub.** Carries events published by *other* processes: the
standalone workers (``pnpm worker:ai-runtime``) run tasks in their own
process, and without a broker their events would die there. The API process
relays whatever arrives on the shared channel into the same manager.
Redis is a bonus, never a dependency:
* The bus tries it once per cooldown window. If nothing is listening, the
attempt fails quietly and in-process delivery continues — the product stays
fully live for every API-driven action, only cross-process events are lost.
* The relay task owns the subscription lifecycle; it starts with the app and
stops on shutdown, and reconnects on its own after a broker restart.
Nothing here ever raises. Publishing is called from request handlers, from
inside exception handlers and from worker passes; an exception in the feed
must not be able to fail any of those.
"""
from __future__ import annotations
import asyncio
import logging
import time
from collections.abc import Awaitable, Callable
from contextlib import suppress
from app.application.ports.realtime import RealtimeEvent, set_publisher
from app.core.config import settings
logger = logging.getLogger("agy.realtime")
#: Redis pub/sub channel shared by every process of one deployment.
DEFAULT_REDIS_CHANNEL = "agy:realtime"
#: After a failed broker attempt, wait this long before trying again. Keeps a
#: dev machine without Redis from adding a connection timeout to every event.
REDIS_RETRY_COOLDOWN_SECONDS = 10.0
#: How long a single connection attempt may take before it is abandoned.
REDIS_CONNECT_TIMEOUT_SECONDS = 1.0
#: Listener signature: ``await listener(event)``. The ConnectionManager
#: registers one; tests register fakes. The return value is ignored so a
#: listener may report how many connections it fanned out to.
EventListener = Callable[[RealtimeEvent], Awaitable[object]]
class RealtimeBus:
"""Fan-out hub shared by every publisher in this process."""
def __init__(
self,
*,
redis_url: str | None = None,
channel: str = DEFAULT_REDIS_CHANNEL,
) -> None:
self._listeners: set[EventListener] = set()
self._redis_url = redis_url if redis_url is not None else settings.redis_url
self._channel = channel
self._redis = None
self._pubsub = None
self._relay_task: asyncio.Task | None = None
self._next_redis_attempt = 0.0
self._redis_healthy = False
# -- listeners (in-process fan-out) --------------------------------------
def attach(self, listener: EventListener) -> None:
self._listeners.add(listener)
def detach(self, listener: EventListener) -> None:
self._listeners.discard(listener)
# -- publishing -----------------------------------------------------------
async def publish(self, event: RealtimeEvent) -> None:
"""Deliver locally, then mirror to other processes. Never raises."""
try:
await self._dispatch_local(event)
except Exception: # noqa: BLE001 - the feed must never fail a caller
logger.exception("realtime: local dispatch failed for %s", event.type)
if not self._redis_enabled():
return
try:
await self._publish_redis(event)
except Exception: # noqa: BLE001 - same rule
logger.exception("realtime: redis publish failed for %s", event.type)
async def _dispatch_local(self, event: RealtimeEvent) -> None:
for listener in list(self._listeners):
try:
await listener(event)
except Exception: # noqa: BLE001 - one bad listener is not fatal
logger.exception("realtime: listener failed for %s", event.type)
def _redis_enabled(self) -> bool:
# "local" pins the bus to in-process delivery — used by tests and by
# single-process deployments that have no broker on purpose.
return getattr(settings, "realtime_transport", "auto").lower() != "local"
async def _publish_redis(self, event: RealtimeEvent) -> None:
client = await self._redis_client()
if client is None:
return
await client.publish(self._channel, event.to_json())
async def _redis_client(self):
"""Lazily connect, with a cooldown after failure. ``None`` if down."""
if self._redis is not None:
return self._redis
now = time.monotonic()
if now < self._next_redis_attempt:
return None
self._next_redis_attempt = now + REDIS_RETRY_COOLDOWN_SECONDS
try:
import redis.asyncio as aioredis
client = aioredis.from_url(
self._redis_url,
encoding="utf-8",
decode_responses=True,
socket_connect_timeout=REDIS_CONNECT_TIMEOUT_SECONDS,
socket_timeout=REDIS_CONNECT_TIMEOUT_SECONDS,
)
await client.ping()
self._redis = client
self._redis_healthy = True
logger.info("realtime: redis pub/sub connected (%s)", self._channel)
except Exception as exc: # noqa: BLE001 - broker is optional
self._redis = None
self._redis_healthy = False
logger.debug("realtime: redis unavailable (%s); in-process only", exc)
return None
return self._redis
# -- relay (other processes -> this one's browsers) ----------------------
async def start_relay(self) -> None:
"""Subscribe to the shared channel and feed the local manager."""
if self._relay_task is not None and not self._relay_task.done():
return
self._relay_task = asyncio.create_task(self._relay_loop(), name="realtime-relay")
async def stop_relay(self) -> None:
task = self._relay_task
self._relay_task = None
if task is not None:
task.cancel()
with suppress(asyncio.CancelledError):
await task
await self._close_redis()
async def _relay_loop(self) -> None:
while True:
client = await self._redis_client()
if client is None:
# Broker absent (or down mid-flight): wait out the cooldown and
# try again — a Redis started after the API must still work.
await asyncio.sleep(REDIS_RETRY_COOLDOWN_SECONDS)
continue
try:
pubsub = client.pubsub()
await pubsub.subscribe(self._channel)
self._pubsub = pubsub
logger.info("realtime: relay subscribed to %s", self._channel)
async for message in pubsub.listen():
if message.get("type") != "message":
continue
event = _decode_message(message.get("data"))
if event is not None:
# Dedupe by id happens in the manager, so an event this
# process published itself is not delivered twice.
await self._dispatch_local(event)
except asyncio.CancelledError:
raise
except Exception as exc: # noqa: BLE001 - relay must survive
self._redis_healthy = False
logger.warning("realtime: relay interrupted (%s); reconnecting", exc)
finally:
if self._pubsub is not None:
with suppress(Exception):
await self._pubsub.unsubscribe(self._channel)
self._pubsub = None
self._redis = None
await asyncio.sleep(1.0)
async def _close_redis(self) -> None:
client, self._redis = self._redis, None
if client is not None:
with suppress(Exception):
await client.aclose()
self._redis_healthy = False
# -- introspection -------------------------------------------------------
@property
def redis_healthy(self) -> bool:
return self._redis_healthy
@property
def listener_count(self) -> int:
return len(self._listeners)
def _decode_message(raw: object) -> RealtimeEvent | None:
if not isinstance(raw, (str, bytes)):
return None
try:
return RealtimeEvent.from_json(raw)
except Exception: # noqa: BLE001 - a bad frame is skipped, not fatal
logger.warning("realtime: undecodable relay frame dropped")
return None
_bus: RealtimeBus | None = None
def get_bus() -> RealtimeBus:
"""Process-wide bus. Created lazily so importing never opens a socket."""
global _bus
if _bus is None:
_bus = RealtimeBus()
return _bus
def install() -> RealtimeBus:
"""Make this process's bus the application-layer publisher.
Called by the composition roots (``app.main`` for the API,
``app.workers.__main__`` for the workers). Idempotent, so a worker that
imports several modules pays nothing for calling it more than once.
"""
bus = get_bus()
set_publisher(bus)
return bus
async def publish(event: RealtimeEvent) -> None:
"""Publish through the shared bus. Never raises."""
bus = get_bus()
try:
await bus.publish(event)
except Exception: # noqa: BLE001 - belt and braces; publish is best-effort
logger.exception("realtime: publish failed for %s", event.type)
async def start_relay() -> None:
await get_bus().start_relay()
async def stop_relay() -> None:
await get_bus().stop_relay()
+217
View File
@@ -0,0 +1,217 @@
"""Connection bookkeeping: who is connected, and to which channels.
The manager is deliberately transport-agnostic — it does not know it is talking
to a WebSocket, only to an object with ``send_text``. That keeps it testable
without a socket and lets a future transport (SSE, push bridge) reuse it.
Two invariants:
* **A connection only ever receives channels it was authorized for.** The
authorizer runs at subscribe time; the routing table built afterwards can
only contain pre-approved channel names.
* **An event is delivered at most once.** ``dispatch`` de-duplicates by event
id, because the same event reaches a process twice in the normal case: once
in-process and once relayed from Redis.
"""
from __future__ import annotations
import asyncio
import logging
from collections import deque
from collections.abc import Awaitable, Callable
from contextlib import suppress
from dataclasses import dataclass, field
from starlette.websockets import WebSocket, WebSocketState
from app.application.ports.realtime import RealtimeEvent
logger = logging.getLogger("agy.realtime")
#: Bounded dedupe memory. Large enough that a relayed duplicate is always
#: caught, small enough that a chatty day costs a few kilobytes.
MAX_SEEN_EVENTS = 2048
#: Per-connection outbox. If a browser stops reading (a frozen tab), a
#: producer must not be able to grow this without bound. Overflow drops the
#: oldest frame — ordering of what follows stays meaningful.
MAX_QUEUE_DEPTH = 256
#: What a subscribe check looks like: ``(connection, channel) -> allowed``.
ChannelAuthorizer = Callable[["Connection", str], Awaitable[bool]]
#: ``eq=False`` on purpose: a connection is identified by object identity. The
#: default dataclass ``__eq__`` would also make it unhashable (a ``set`` field
#: cannot be compared safely), and the manager keys connections by identity in
#: both its connection set and its channel routing table.
@dataclass(slots=True, eq=False)
class Connection:
"""One authenticated browser socket."""
user_id: str
is_platform_owner: bool
websocket: WebSocket
channels: set[str] = field(default_factory=set)
_queue: asyncio.Queue[str] = field(
default_factory=lambda: asyncio.Queue(maxsize=MAX_QUEUE_DEPTH)
)
_sender: asyncio.Task | None = None
async def start(self) -> None:
"""Single writer task per connection.
One task owns the socket: concurrent ``send_text`` calls on one
WebSocket interleave frames and corrupt the stream, so everything goes
through this queue instead.
"""
if self._sender is None or self._sender.done():
self._sender = asyncio.create_task(self._drain(), name="realtime-sender")
async def _drain(self) -> None:
while True:
payload = await self._queue.get()
if payload is None: # shutdown sentinel
return
try:
if self.websocket.client_state == WebSocketState.CONNECTED:
await self.websocket.send_text(payload)
except Exception: # noqa: BLE001 - a dead socket ends the connection
logger.debug("realtime: send failed; closing connection", exc_info=True)
return
def offer(self, payload: str) -> None:
"""Queue a frame, never blocking the publisher.
A slow consumer loses the oldest frames rather than stalling every
other tenant's stream behind it.
"""
if self._queue.full():
with suppress(asyncio.QueueFull):
# Drop the oldest frame to make room: ordering of what follows
# stays meaningful, which is what the UI renders.
self._queue.get_nowait()
with suppress(asyncio.QueueFull): # pragma: no cover - defensive
self._queue.put_nowait(payload)
async def close(self) -> None:
if self._sender is not None:
self._sender.cancel()
self._sender = None
class ConnectionManager:
"""The routing table between events and open sockets."""
def __init__(self) -> None:
self._connections: set[Connection] = set()
self._by_channel: dict[str, set[Connection]] = {}
self._seen: deque[str] = deque(maxlen=MAX_SEEN_EVENTS)
self._seen_set: set[str] = set()
# -- lifecycle -----------------------------------------------------------
async def register(self, connection: Connection) -> None:
self._connections.add(connection)
await connection.start()
async def unregister(self, connection: Connection) -> None:
self._connections.discard(connection)
for channel in list(connection.channels):
subscribers = self._by_channel.get(channel)
if subscribers is not None:
subscribers.discard(connection)
if not subscribers:
self._by_channel.pop(channel, None)
connection.channels.clear()
await connection.close()
# -- subscriptions ---------------------------------------------------------
async def subscribe(
self,
connection: Connection,
channels: list[str],
authorizer: ChannelAuthorizer,
) -> tuple[list[str], list[str]]:
"""Authorize and subscribe. Returns ``(accepted, rejected)``."""
accepted: list[str] = []
rejected: list[str] = []
for channel in channels:
if channel in connection.channels:
accepted.append(channel)
continue
try:
allowed = await authorizer(connection, channel)
except Exception: # noqa: BLE001 - a broken check denies, never allows
logger.exception("realtime: channel authorization crashed")
allowed = False
if not allowed:
rejected.append(channel)
continue
connection.channels.add(channel)
self._by_channel.setdefault(channel, set()).add(connection)
accepted.append(channel)
return accepted, rejected
async def unsubscribe(self, connection: Connection, channels: list[str]) -> None:
for channel in channels:
connection.channels.discard(channel)
subscribers = self._by_channel.get(channel)
if subscribers is not None:
subscribers.discard(connection)
if not subscribers:
self._by_channel.pop(channel, None)
# -- fan-out ---------------------------------------------------------------
async def dispatch(self, event: RealtimeEvent) -> int:
"""Deliver to every subscriber of every channel. Returns fan-out count.
The id check is what makes "publish locally, then receive the same event
from Redis" deliver once instead of twice.
"""
if event.id in self._seen_set:
return 0
self._remember(event.id)
recipients: set[Connection] = set()
for channel in event.channels:
recipients.update(self._by_channel.get(channel, ()))
envelope = event.envelope()
# Control frames and event frames share one top-level discriminator
# (``type``), so the client switches on one field for everything.
payload = _dumps({"type": "event", "event": envelope})
delivered = 0
for connection in recipients:
connection.offer(payload)
delivered += 1
return delivered
def _remember(self, event_id: str) -> None:
if len(self._seen) == self._seen.maxlen:
evicted = self._seen[0]
self._seen_set.discard(evicted)
self._seen.append(event_id)
self._seen_set.add(event_id)
# -- introspection -------------------------------------------------------
@property
def connection_count(self) -> int:
return len(self._connections)
@property
def channel_count(self) -> int:
return len(self._by_channel)
def channel_load(self) -> dict[str, int]:
return {channel: len(subs) for channel, subs in self._by_channel.items()}
def _dumps(envelope: dict) -> str:
import json
return json.dumps(envelope, ensure_ascii=False, default=str)
+175
View File
@@ -0,0 +1,175 @@
"""The WebSocket endpoint — one socket per browser, every live feed on it.
``GET /api/ws?token=<jwt>``
The token travels as a query parameter because the browser WebSocket API cannot
set an ``Authorization`` header. That is a deliberate, contained exception to
the header-only rule: the token is short-lived, the socket is TLS-terminated in
any real deployment, and the alternative (a cookie) would drag CSRF into a
protocol that has no CSRF story.
Protocol — small, symmetric, JSON in both directions:
.. code-block:: text
client -> {"action": "subscribe", "channels": ["org:…", "task:…"]}
client -> {"action": "unsubscribe", "channels": ["task:…"]}
client -> {"action": "ping"}
server -> {"type": "hello", "user_id": "…"}
server -> {"type": "subscribed", "accepted": [...], "rejected": [...]}
server -> {"type": "event", "event": {…envelope…}}
server -> {"type": "pong"}
server -> {"type": "error", "message": "…"}
A rejected channel is reported, not hidden: "my task page is not live" must be
answerable from the browser console alone.
Authorization queries open their own short-lived session: a socket can stay
open for hours, and one SQLAlchemy session per open tab would quietly exhaust
the connection pool.
"""
from __future__ import annotations
import json
import logging
from fastapi import APIRouter, Query, WebSocket, WebSocketDisconnect
from app.application.ports.realtime import dedupe_channels
from app.core.database import SessionFactory
from app.core.errors import AuthenticationError
from app.core.security import decode_access_token
from app.infrastructure.repositories.identity import UserRepository
from app.interfaces.realtime.authz import authorize_channel
from app.interfaces.realtime.manager import Connection, ConnectionManager
logger = logging.getLogger("agy.realtime")
router = APIRouter()
#: The manager instance this process fans out with. Created at import: the
#: endpoint and the lifecycle hooks can share it without a wiring ceremony.
_manager = ConnectionManager()
def realtime_manager() -> ConnectionManager:
"""The process-wide manager — attached to the bus by ``app.main``."""
return _manager
@router.websocket("/ws")
async def realtime_socket(
websocket: WebSocket,
token: str | None = Query(default=None),
) -> None:
connection = await _authenticate(websocket, token)
if connection is None:
return # _authenticate already closed with a policy-violation code
await _manager.register(connection)
await _send(
websocket,
{"type": "hello", "user_id": connection.user_id, "server_time": _now_iso()},
)
try:
while True:
raw = await websocket.receive_text()
await _handle_message(websocket, connection, raw)
except WebSocketDisconnect:
pass
except Exception: # noqa: BLE001 - any read failure ends the socket cleanly
logger.debug("realtime: socket loop ended", exc_info=True)
finally:
await _manager.unregister(connection)
async def _authenticate(websocket: WebSocket, token: str | None) -> Connection | None:
"""Accept only a valid, active user. Closes 1008 otherwise."""
user = None
if token:
try:
payload = decode_access_token(token)
user_id = payload.get("sub")
if user_id:
# Own session: the socket must not depend on request scope.
async with SessionFactory() as session:
user = await UserRepository(session).get(str(user_id))
except AuthenticationError:
user = None
if user is None or not user.is_active:
# 1008 (policy violation) rather than a bare close: the client's
# reconnect logic can then tell "bad token" from "network dropped".
await websocket.close(code=1008, reason="unauthorized")
return None
await websocket.accept()
return Connection(
user_id=user.id,
is_platform_owner=user.is_platform_owner,
websocket=websocket,
)
async def _handle_message(websocket: WebSocket, connection: Connection, raw: str) -> None:
try:
message = json.loads(raw)
except (TypeError, ValueError):
await _send(websocket, {"type": "error", "message": "malformed JSON frame"})
return
if not isinstance(message, dict):
await _send(websocket, {"type": "error", "message": "frame must be an object"})
return
action = message.get("action")
if action == "ping":
await _send(websocket, {"type": "pong"})
return
channels = dedupe_channels(
str(channel) for channel in (message.get("channels") or []) if isinstance(channel, str)
)
if action == "subscribe":
accepted, rejected = await _subscribe(connection, channels)
await _send(
websocket, {"type": "subscribed", "accepted": accepted, "rejected": rejected}
)
return
if action == "unsubscribe":
await _manager.unsubscribe(connection, list(channels))
await _send(websocket, {"type": "unsubscribed", "channels": channels})
return
await _send(websocket, {"type": "error", "message": f"unknown action: {action!r}"})
async def _subscribe(
connection: Connection, channels: tuple[str, ...]
) -> tuple[list[str], list[str]]:
if not channels:
return [], []
async def _authorize(conn: Connection, channel: str) -> bool:
async with SessionFactory() as session:
return await authorize_channel(
session,
user_id=conn.user_id,
is_platform_owner=conn.is_platform_owner,
channel=channel,
)
return await _manager.subscribe(connection, list(channels), _authorize)
def _now_iso() -> str:
from app.domain.base import utcnow
return utcnow().isoformat()
async def _send(websocket: WebSocket, payload: dict) -> None:
await websocket.send_text(json.dumps(payload, ensure_ascii=False, default=str))
+17
View File
@@ -27,6 +27,10 @@ from app.interfaces.http.routers import (
sandbox,
workforce,
)
from app.interfaces.realtime.bus import install as install_realtime_bus
from app.interfaces.realtime.bus import start_relay, stop_relay
from app.interfaces.realtime.router import realtime_manager
from app.interfaces.realtime.router import router as realtime_router
logging.basicConfig(
level=logging.INFO,
@@ -62,8 +66,19 @@ async def lifespan(_: FastAPI):
# loudly — a half-started API must not accept traffic.
await _record_startup_failure(exc)
raise
# --- realtime wiring ---------------------------------------------------
# This process is the publisher for API-driven events and the relay for
# anything the standalone workers publish. Both are best-effort: if Redis
# is not reachable the in-process fan-out keeps every HTTP-driven feed live.
bus = install_realtime_bus()
bus.attach(realtime_manager().dispatch)
await start_relay()
yield
await stop_relay()
async def _record_startup_failure(exc: BaseException) -> None:
"""Persist a boot failure. Never raises; the caller re-raises regardless."""
@@ -123,6 +138,8 @@ def create_app() -> FastAPI:
application.include_router(sandbox.router, prefix="/api")
application.include_router(admin.router, prefix="/api")
application.include_router(agy.router, prefix="/api")
# The live feed: one WebSocket per browser, every push on it.
application.include_router(realtime_router, prefix="/api")
@application.get("/", tags=["meta"])
async def root() -> dict:
+5
View File
@@ -125,6 +125,11 @@ def main(argv: list[str] | None = None) -> int:
format="%(asctime)s %(levelname)-7s %(name)s: %(message)s",
)
args = _build_parser().parse_args(argv)
# A worker is its own process: publishing here is how a run it starts, or a
# session it reaps, reaches the browsers connected to the API process.
from app.interfaces.realtime.bus import install as install_realtime
install_realtime()
return asyncio.run(_main_async(args))
+8
View File
@@ -14,6 +14,7 @@ the service's business.
from __future__ import annotations
from app.application.container import build_workforce_service
from app.application.ports.realtime import emit_worker_status
from app.core.errors import AppError
from app.domain.audit.models import ErrorSeverity, ErrorSource
from app.domain.workforce.models import TaskStatus
@@ -87,6 +88,13 @@ class AiRuntimeWorker:
outcome.processed += 1
outcome.note(f"task {task.id} -> run {run.id} ({run.status})")
# This worker is its own process: without a push here, a run
# started by the queue is invisible until someone reloads.
await emit_worker_status(
organization_id=task.organization_id,
worker=self.name,
message=f"کار «{current.title}» در حال اجراست (اجرا {run.id}).",
)
return outcome
+8
View File
@@ -21,6 +21,7 @@ from __future__ import annotations
from datetime import timedelta
from app.application.container import build_sandbox_service
from app.application.ports.realtime import emit_sandbox_reaped
from app.core.config import settings
from app.domain.audit.models import AuditAction, AuditEvent
from app.domain.base import utcnow
@@ -72,6 +73,13 @@ class SandboxGatewayWorker:
)
outcome.processed += 1
outcome.note(f"stopped session {sandbox_session.id}")
# A reaped sandbox is a tenant-visible fact: push it so the
# project page's terminal stops showing a live session.
await emit_sandbox_reaped(
organization_id=workspace.organization_id if workspace else None,
session_id=sandbox_session.id,
workspace_id=sandbox_session.workspace_id,
)
if not stopped:
outcome.skipped += 1
+12
View File
@@ -39,6 +39,18 @@ def _production_shaped_jwt_secret():
settings.jwt_secret = "test-only-jwt-secret-with-more-than-32-bytes"
@pytest.fixture(scope="session", autouse=True)
def _hermetic_realtime_transport():
"""Pin the realtime bus to in-process delivery.
``auto`` would make every publish attempt a Redis connection to
``localhost:6379`` — the suite has no broker, so each event would add a
connection timeout and log noise. The fan-out behaviour under test does not
depend on the transport.
"""
settings.realtime_transport = "local"
@pytest_asyncio.fixture
async def engine():
# StaticPool keeps one connection alive so the in-memory schema survives
+416
View File
@@ -0,0 +1,416 @@
"""Realtime push layer: authorization, fan-out, and the live socket end-to-end.
Three levels, mirroring the package's structure:
* **Authorization** — a channel is only ever granted through
``authorize_channel``; every rule is tested for both the allow and the deny
side, because a rule that always denies is indistinguishable from a bug.
* **Fan-out** — the manager delivers to the subscribers of a channel, keeps
ordering, de-duplicates relayed events and cleans up on unsubscribe.
* **The socket** — a real WebSocket against the real app: subscribe, act over
HTTP, and assert the event arrives without a reload.
"""
from __future__ import annotations
import asyncio
from contextlib import suppress
from pathlib import Path
import pytest
from fastapi.testclient import TestClient
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from sqlalchemy.pool import NullPool
import app.infrastructure.db.models as _orm_models # noqa: F401 - registers metadata
from app.application.ports.realtime import RealtimeEvent, set_publisher
from app.core.config import settings
from app.core.database import Base, get_session
from app.core.security import create_access_token
from app.domain.base import new_id
from app.infrastructure.db.models import UserRow
from app.interfaces.realtime.authz import authorize_channel
from app.interfaces.realtime.bus import get_bus
from app.interfaces.realtime.manager import Connection, ConnectionManager
from app.main import app
API_ROOT = Path(__file__).resolve().parents[1]
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
class FakeSocket:
"""The manager only needs ``send_text`` and a connected ``client_state``.
Both matter: the sender task checks ``client_state`` before writing, so a
fake without it would silently swallow every frame and make the manager
look broken.
"""
def __init__(self) -> None:
self.sent: list[str] = []
from starlette.websockets import WebSocketState
self.client_state = WebSocketState.CONNECTED
async def send_text(self, payload: str) -> None:
self.sent.append(payload)
def _event(event_id: str, *channels: str) -> RealtimeEvent:
return RealtimeEvent(type="test.event", channels=channels, data={"id": event_id}, id=event_id)
async def _allow(_: Connection, __: str) -> bool:
return True
# ---------------------------------------------------------------------------
# Authorization
# ---------------------------------------------------------------------------
async def _user_with_org(session: AsyncSession, *, platform_owner: bool = False):
"""A user, and optionally an organization they belong to."""
user_id = new_id()
session.add(
UserRow(
id=user_id,
email=f"{user_id}@example.com",
password_hash="x",
is_platform_owner=platform_owner,
)
)
await session.flush()
return user_id
async def _org_with_member(session: AsyncSession, user_id: str) -> str:
from app.domain.identity.models import Membership, Organization, Role
from app.infrastructure.repositories.identity import (
MembershipRepository,
OrganizationRepository,
)
organization = await OrganizationRepository(session).create(
Organization(name="سازمان تست", slug=new_id()[:8], owner_user_id=user_id)
)
await MembershipRepository(session).add(
Membership(organization_id=organization.id, user_id=user_id, role=Role.ADMIN)
)
return organization.id
@pytest.mark.asyncio
async def test_member_may_subscribe_to_own_org(session_factory):
async with session_factory() as session:
user_id = await _user_with_org(session)
org_id = await _org_with_member(session, user_id)
allowed = await authorize_channel(
session, user_id=user_id, is_platform_owner=False, channel=f"org:{org_id}"
)
assert allowed is True
@pytest.mark.asyncio
async def test_outsider_may_not_subscribe_to_foreign_org(session_factory):
async with session_factory() as session:
insider = await _user_with_org(session)
outsider = await _user_with_org(session)
org_id = await _org_with_member(session, insider)
allowed = await authorize_channel(
session, user_id=outsider, is_platform_owner=False, channel=f"org:{org_id}"
)
assert allowed is False
@pytest.mark.asyncio
async def test_unknown_channel_kinds_are_denied(session_factory):
async with session_factory() as session:
user_id = await _user_with_org(session)
allowed = await authorize_channel(
session, user_id=user_id, is_platform_owner=False, channel="nonsense:whatever"
)
assert allowed is False
@pytest.mark.asyncio
async def test_admin_channel_needs_platform_owner(session_factory):
async with session_factory() as session:
owner = await _user_with_org(session, platform_owner=True)
member = await _user_with_org(session)
owner_ok = await authorize_channel(
session, user_id=owner, is_platform_owner=True, channel="admin"
)
member_ok = await authorize_channel(
session, user_id=member, is_platform_owner=False, channel="admin"
)
assert owner_ok is True
assert member_ok is False
@pytest.mark.asyncio
async def test_user_channel_is_private(session_factory):
async with session_factory() as session:
user_id = await _user_with_org(session)
other = await _user_with_org(session)
own = await authorize_channel(
session, user_id=user_id, is_platform_owner=False, channel=f"user:{user_id}"
)
foreign = await authorize_channel(
session, user_id=other, is_platform_owner=False, channel=f"user:{user_id}"
)
assert own is True
assert foreign is False
# ---------------------------------------------------------------------------
# Fan-out
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_manager_delivers_to_subscribers_only():
manager = ConnectionManager()
listening = FakeSocket()
other = FakeSocket()
subscribed = Connection(user_id="u1", is_platform_owner=False, websocket=listening)
stranger = Connection(user_id="u2", is_platform_owner=False, websocket=other)
for connection in (subscribed, stranger):
await manager.register(connection)
await manager.subscribe(subscribed, ["org:1"], _allow)
delivered = await manager.dispatch(_event("e1", "org:1"))
assert delivered == 1
# The write happens on the connection's own sender task; give the loop the
# turn it needs to flush the queue.
await asyncio.sleep(0.05)
assert len(listening.sent) == 1
assert other.sent == []
@pytest.mark.asyncio
async def test_manager_wraps_events_in_the_wire_envelope():
"""A browser sees one discriminator: ``{"type": "event", "event": {...}}``."""
import json
manager = ConnectionManager()
socket = FakeSocket()
connection = Connection(user_id="u1", is_platform_owner=False, websocket=socket)
await manager.register(connection)
await manager.subscribe(connection, ["org:1"], _allow)
await manager.dispatch(_event("e1", "org:1"))
await asyncio.sleep(0.05)
assert len(socket.sent) == 1
frame = json.loads(socket.sent[0])
assert frame["type"] == "event"
assert frame["event"]["id"] == "e1"
assert frame["event"]["type"] == "test.event"
@pytest.mark.asyncio
async def test_manager_drops_duplicate_events():
"""In-process publish plus a relayed copy deliver exactly once."""
manager = ConnectionManager()
socket = FakeSocket()
connection = Connection(user_id="u1", is_platform_owner=False, websocket=socket)
await manager.register(connection)
await manager.subscribe(connection, ["org:1"], _allow)
event = _event("dup", "org:1")
first = await manager.dispatch(event)
second = await manager.dispatch(event)
assert (first, second) == (1, 0)
@pytest.mark.asyncio
async def test_unsubscribe_stops_delivery():
manager = ConnectionManager()
socket = FakeSocket()
connection = Connection(user_id="u1", is_platform_owner=False, websocket=socket)
await manager.register(connection)
await manager.subscribe(connection, ["org:1", "org:2"], _allow)
await manager.unsubscribe(connection, ["org:1"])
assert await manager.dispatch(_event("e1", "org:1")) == 0
assert await manager.dispatch(_event("e2", "org:2")) == 1
@pytest.mark.asyncio
async def test_publisher_never_raises_on_bad_listener():
"""A listener that explodes must not break the publisher or its peers."""
async def exploding(event: RealtimeEvent) -> None:
raise RuntimeError("boom")
bus = get_bus()
bus.attach(exploding)
try:
set_publisher(bus)
# Must not raise despite the exploding listener.
from app.application.ports.realtime import publish
await publish(_event("e1", "org:1"))
finally:
bus.detach(exploding)
set_publisher(_NullPublisher())
class _NullPublisher:
async def publish(self, event: RealtimeEvent) -> None:
return None
# ---------------------------------------------------------------------------
# The socket, end to end
# ---------------------------------------------------------------------------
async def _create_schema(engine) -> None:
async with engine.begin() as conn:
await conn.run_sync(Base.metadata.create_all)
@pytest.fixture()
def ws_client(monkeypatch):
"""A TestClient whose API and socket both talk to one throwaway database.
File-based SQLite with ``NullPool`` on purpose: the app runs on the test
portal's event loop while the schema is created from the test's own, and a
pooled connection cannot cross that boundary. The database file lives in
the package directory — ``*.db`` is gitignored — so the fixture needs no
writable temp directory outside the workspace.
"""
db_path = API_ROOT / f"realtime-test-{new_id()}.db"
engine = create_async_engine(f"sqlite+aiosqlite:///{db_path}", poolclass=NullPool)
asyncio.run(_create_schema(engine))
factory = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
monkeypatch.setattr("app.main.engine", engine)
monkeypatch.setattr("app.interfaces.realtime.router.SessionFactory", factory)
monkeypatch.setattr(settings, "realtime_transport", "local")
async def _override_session():
async with factory() as session:
try:
yield session
await session.commit()
except Exception:
await session.rollback()
raise
app.dependency_overrides[get_session] = _override_session
# The manager is module state; start each run from a clean slate so a
# connection from another test cannot colour this one's fan-out.
from app.interfaces.realtime.router import realtime_manager
realtime_manager()._connections.clear()
realtime_manager()._by_channel.clear()
try:
with TestClient(app) as client:
yield client
finally:
app.dependency_overrides.clear()
asyncio.run(engine.dispose())
with suppress(OSError):
db_path.unlink()
def _frames_until(websocket, wanted_type: str, *, max_frames: int = 20) -> dict:
"""Read frames until ``wanted_type`` arrives — hello/subscribed come first."""
for _ in range(max_frames):
frame = websocket.receive_json()
if frame.get("type") == "event" and frame.get("event", {}).get("type") == wanted_type:
return frame["event"]
raise AssertionError(f"no {wanted_type} event arrived")
def test_socket_rejects_a_missing_token(ws_client):
with pytest.raises(Exception), ws_client.websocket_connect("/api/ws"): # noqa: B017, SIM117
# No token: the endpoint closes with 1008 before accepting.
pass
def test_task_creation_is_pushed_live(ws_client):
# Sign up, create an organization — the tenant the socket subscribes to.
signup = ws_client.post(
"/api/auth/signup",
json={"email": "live@example.com", "password": "password123", "name": "زنده"},
)
assert signup.status_code == 201, signup.text
token = signup.json()["access_token"]
org = ws_client.post(
"/api/organizations",
json={"name": "سازمان زنده"},
headers={"Authorization": f"Bearer {token}"},
)
assert org.status_code == 201, org.text
org_id = org.json()["id"]
with ws_client.websocket_connect(f"/api/ws?token={token}") as socket:
hello = socket.receive_json()
assert hello["type"] == "hello"
socket.send_json({"action": "subscribe", "channels": [f"org:{org_id}"]})
subscribed = socket.receive_json()
assert subscribed["type"] == "subscribed"
assert subscribed["accepted"] == [f"org:{org_id}"]
# A foreign channel is refused, not silently ignored.
socket.send_json({"action": "subscribe", "channels": ["org:does-not-exist"]})
refused = socket.receive_json()
assert refused["type"] == "subscribed"
assert refused["rejected"] == ["org:does-not-exist"]
created = ws_client.post(
f"/api/organizations/{org_id}/tasks",
json={"title": "کار زنده"},
headers={"Authorization": f"Bearer {token}"},
)
assert created.status_code == 201, created.text
event = _frames_until(socket, "task.created")
assert event["data"]["title"] == "کار زنده"
assert event["organization_id"] == org_id
def test_second_user_cannot_read_another_org(ws_client):
"""Two tenants, one socket each: an outsider's org is out of reach."""
tokens = []
org_ids = []
for index in range(2):
signup = ws_client.post(
"/api/auth/signup",
json={"email": f"user{index}@example.com", "password": "password123"},
)
assert signup.status_code == 201, signup.text
token = signup.json()["access_token"]
org = ws_client.post(
"/api/organizations",
json={"name": f"org {index}"},
headers={"Authorization": f"Bearer {token}"},
)
assert org.status_code == 201, org.text
tokens.append(token)
org_ids.append(org.json()["id"])
with ws_client.websocket_connect(f"/api/ws?token={tokens[0]}") as socket:
socket.receive_json() # hello
socket.send_json({"action": "subscribe", "channels": [f"org:{org_ids[1]}"]})
refused = socket.receive_json()
assert refused["rejected"] == [f"org:{org_ids[1]}"]
def test_socket_token_identifies_a_real_user(ws_client):
"""A well-formed token for a deleted user must not open a socket."""
user_id = new_id()
token = create_access_token(user_id)
with pytest.raises(Exception), ws_client.websocket_connect(f"/api/ws?token={token}"): # noqa: B017, SIM117
# Closed with 1008: the token is well-formed but names no user.
pass