From f1e5df29ce669568c1913dd10b8111699d66c5cf Mon Sep 17 00:00:00 2001 From: mjalmoudy94 Date: Tue, 13 Oct 2026 11:41:01 +0330 Subject: [PATCH] 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. --- apps/api/app/application/ports/realtime.py | 504 ++++++++++++++++++ .../app/application/services/audit_service.py | 12 +- .../application/services/billing_service.py | 13 +- .../application/services/sandbox_service.py | 29 + .../application/services/workforce_service.py | 223 +++++++- apps/api/app/core/config.py | 9 + apps/api/app/infrastructure/error_log.py | 13 + .../api/app/interfaces/http/routers/health.py | 17 + apps/api/app/interfaces/realtime/__init__.py | 20 + apps/api/app/interfaces/realtime/authz.py | 112 ++++ apps/api/app/interfaces/realtime/bus.py | 260 +++++++++ apps/api/app/interfaces/realtime/manager.py | 217 ++++++++ apps/api/app/interfaces/realtime/router.py | 175 ++++++ apps/api/app/main.py | 17 + apps/api/app/workers/__main__.py | 5 + apps/api/app/workers/ai_runtime.py | 8 + apps/api/app/workers/sandbox_gateway.py | 8 + apps/api/tests/conftest.py | 12 + apps/api/tests/test_realtime.py | 416 +++++++++++++++ apps/web/src/app/(app)/admin/errors/page.tsx | 58 ++ apps/web/src/app/(app)/billing/page.tsx | 39 ++ apps/web/src/app/(app)/dashboard/page.tsx | 27 + apps/web/src/app/(app)/manager/page.tsx | 60 ++- apps/web/src/app/(app)/tasks/[id]/page.tsx | 109 +++- apps/web/src/app/(app)/traces/page.tsx | 68 ++- apps/web/src/app/layout.tsx | 6 +- apps/web/src/components/AppShell.tsx | 36 ++ apps/web/src/components/sandbox-panel.tsx | 48 ++ apps/web/src/lib/realtime.tsx | 135 +++++ apps/web/src/lib/ws.ts | 270 ++++++++++ packages/contracts/src/index.ts | 83 +++ 31 files changed, 2981 insertions(+), 28 deletions(-) create mode 100644 apps/api/app/application/ports/realtime.py create mode 100644 apps/api/app/interfaces/realtime/__init__.py create mode 100644 apps/api/app/interfaces/realtime/authz.py create mode 100644 apps/api/app/interfaces/realtime/bus.py create mode 100644 apps/api/app/interfaces/realtime/manager.py create mode 100644 apps/api/app/interfaces/realtime/router.py create mode 100644 apps/api/tests/test_realtime.py create mode 100644 apps/web/src/lib/realtime.tsx create mode 100644 apps/web/src/lib/ws.ts diff --git a/apps/api/app/application/ports/realtime.py b/apps/api/app/application/ports/realtime.py new file mode 100644 index 0000000..49513a1 --- /dev/null +++ b/apps/api/app/application/ports/realtime.py @@ -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)}, + ) diff --git a/apps/api/app/application/services/audit_service.py b/apps/api/app/application/services/audit_service.py index 40ebff4..77af591 100644 --- a/apps/api/app/application/services/audit_service.py +++ b/apps/api/app/application/services/audit_service.py @@ -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 diff --git a/apps/api/app/application/services/billing_service.py b/apps/api/app/application/services/billing_service.py index 6700eb7..b9a9e21 100644 --- a/apps/api/app/application/services/billing_service.py +++ b/apps/api/app/application/services/billing_service.py @@ -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 diff --git a/apps/api/app/application/services/sandbox_service.py b/apps/api/app/application/services/sandbox_service.py index 8e42185..dd3f41e 100644 --- a/apps/api/app/application/services/sandbox_service.py +++ b/apps/api/app/application/services/sandbox_service.py @@ -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: diff --git a/apps/api/app/application/services/workforce_service.py b/apps/api/app/application/services/workforce_service.py index e30ff7d..613025c 100644 --- a/apps/api/app/application/services/workforce_service.py +++ b/apps/api/app/application/services/workforce_service.py @@ -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): diff --git a/apps/api/app/core/config.py b/apps/api/app/core/config.py index b41b612..5de0487 100644 --- a/apps/api/app/core/config.py +++ b/apps/api/app/core/config.py @@ -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" diff --git a/apps/api/app/infrastructure/error_log.py b/apps/api/app/infrastructure/error_log.py index 9505c9a..423b802 100644 --- a/apps/api/app/infrastructure/error_log.py +++ b/apps/api/app/infrastructure/error_log.py @@ -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 diff --git a/apps/api/app/interfaces/http/routers/health.py b/apps/api/app/interfaces/http/routers/health.py index 06d9af3..8ef16a0 100644 --- a/apps/api/app/interfaces/http/routers/health.py +++ b/apps/api/app/interfaces/http/routers/health.py @@ -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} diff --git a/apps/api/app/interfaces/realtime/__init__.py b/apps/api/app/interfaces/realtime/__init__.py new file mode 100644 index 0000000..903b29d --- /dev/null +++ b/apps/api/app/interfaces/realtime/__init__.py @@ -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", +] diff --git a/apps/api/app/interfaces/realtime/authz.py b/apps/api/app/interfaces/realtime/authz.py new file mode 100644 index 0000000..1783d9a --- /dev/null +++ b/apps/api/app/interfaces/realtime/authz.py @@ -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 + ) diff --git a/apps/api/app/interfaces/realtime/bus.py b/apps/api/app/interfaces/realtime/bus.py new file mode 100644 index 0000000..968ab2f --- /dev/null +++ b/apps/api/app/interfaces/realtime/bus.py @@ -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() diff --git a/apps/api/app/interfaces/realtime/manager.py b/apps/api/app/interfaces/realtime/manager.py new file mode 100644 index 0000000..fb9fc54 --- /dev/null +++ b/apps/api/app/interfaces/realtime/manager.py @@ -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) diff --git a/apps/api/app/interfaces/realtime/router.py b/apps/api/app/interfaces/realtime/router.py new file mode 100644 index 0000000..6a4d91b --- /dev/null +++ b/apps/api/app/interfaces/realtime/router.py @@ -0,0 +1,175 @@ +"""The WebSocket endpoint — one socket per browser, every live feed on it. + +``GET /api/ws?token=`` + +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)) diff --git a/apps/api/app/main.py b/apps/api/app/main.py index 94fe661..ba4bc41 100644 --- a/apps/api/app/main.py +++ b/apps/api/app/main.py @@ -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: diff --git a/apps/api/app/workers/__main__.py b/apps/api/app/workers/__main__.py index 612c02f..301e5fb 100644 --- a/apps/api/app/workers/__main__.py +++ b/apps/api/app/workers/__main__.py @@ -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)) diff --git a/apps/api/app/workers/ai_runtime.py b/apps/api/app/workers/ai_runtime.py index a5e5589..110728c 100644 --- a/apps/api/app/workers/ai_runtime.py +++ b/apps/api/app/workers/ai_runtime.py @@ -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 diff --git a/apps/api/app/workers/sandbox_gateway.py b/apps/api/app/workers/sandbox_gateway.py index d6d218e..06a6b19 100644 --- a/apps/api/app/workers/sandbox_gateway.py +++ b/apps/api/app/workers/sandbox_gateway.py @@ -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 diff --git a/apps/api/tests/conftest.py b/apps/api/tests/conftest.py index f8b54c6..358471f 100644 --- a/apps/api/tests/conftest.py +++ b/apps/api/tests/conftest.py @@ -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 diff --git a/apps/api/tests/test_realtime.py b/apps/api/tests/test_realtime.py new file mode 100644 index 0000000..aaa19a9 --- /dev/null +++ b/apps/api/tests/test_realtime.py @@ -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 diff --git a/apps/web/src/app/(app)/admin/errors/page.tsx b/apps/web/src/app/(app)/admin/errors/page.tsx index fc91351..aabbff5 100644 --- a/apps/web/src/app/(app)/admin/errors/page.tsx +++ b/apps/web/src/app/(app)/admin/errors/page.tsx @@ -25,6 +25,7 @@ import { } from "@agy/ui"; import type { TableColumn } from "@agy/ui"; import { ApiError, api } from "@/lib/api"; +import { useLiveChannel } from "@/lib/realtime"; /** * Platform-owner error console. @@ -101,6 +102,63 @@ export default function AdminErrorsPage() { const [busy, setBusy] = React.useState(null); const [notice, setNotice] = React.useState(null); + // Live feed: new errors appear in the console as they happen. + useLiveChannel(["admin"], (event) => { + if (event.type === "error.recorded") { + const data = event.data as { + error_id: string | null; + severity: ErrorSeverity; + source: ErrorSource; + kind: string; + message: string; + organization_id: string | null; + }; + const id = data.error_id ?? event.id; + setEvents((current) => { + if (current.some((e) => e.id === id)) return current; + return [ + { + id, + kind: data.kind, + message: data.message, + stack: "", + severity: data.severity, + source: data.source, + error_type: data.kind, + fingerprint: id, + status_code: null, + method: "", + path: "", + organization_id: data.organization_id, + user_id: null, + run_id: null, + request_id: "", + ip: "", + user_agent: "", + metadata_json: {}, + resolved_at: null, + resolved_note: "", + created_at: new Date().toISOString(), + }, + ...current, + ]; + }); + setTotal((current) => current + 1); + setSummary((current) => + current + ? { + ...current, + total_open: current.total_open + 1, + open_by_severity: { + ...current.open_by_severity, + [data.severity]: (current.open_by_severity[data.severity] ?? 0) + 1, + }, + } + : current, + ); + } + }); + const load = React.useCallback(async () => { setLoading(true); setError(null); diff --git a/apps/web/src/app/(app)/billing/page.tsx b/apps/web/src/app/(app)/billing/page.tsx index adbcaee..ea0e52c 100644 --- a/apps/web/src/app/(app)/billing/page.tsx +++ b/apps/web/src/app/(app)/billing/page.tsx @@ -23,6 +23,7 @@ import { } from "@agy/ui"; import { ApiError, api } from "@/lib/api"; import { CAN_MANAGE_BILLING, roleHasCapability } from "@/lib/permissions"; +import { useLiveChannel } from "@/lib/realtime"; import { useSession } from "@/lib/session"; const UNIT_LABELS: Record = { @@ -69,6 +70,44 @@ export default function BillingPage() { const [actionError, setActionError] = React.useState(null); const [checkoutUrl, setCheckoutUrl] = React.useState(null); + // Live feed: wallet balance and usage events update as they happen. + useLiveChannel( + orgId ? [`org:${orgId}`] : [], + (event) => { + switch (event.type) { + case "billing.wallet": { + const data = event.data as { organization_id: string; balance: number }; + setSummary((current) => + current ? { ...current, balance: data.balance, low_balance: Boolean(data.balance <= current.low_balance_threshold) } : current, + ); + break; + } + case "usage.recorded": { + const data = event.data as { + run_id: string | null; + unit: string; + quantity: number; + }; + setEvents((current) => [ + { + id: event.id, + run_id: data.run_id, + unit: data.unit as UsageEvent["unit"], + quantity: data.quantity, + unit_price: 0, + total_price: 0, + created_at: new Date().toISOString(), + }, + ...current, + ]); + break; + } + default: + break; + } + }, + ); + const load = React.useCallback(async () => { if (!orgId) { setLoading(false); diff --git a/apps/web/src/app/(app)/dashboard/page.tsx b/apps/web/src/app/(app)/dashboard/page.tsx index 1a8d0a8..692cc5e 100644 --- a/apps/web/src/app/(app)/dashboard/page.tsx +++ b/apps/web/src/app/(app)/dashboard/page.tsx @@ -15,6 +15,7 @@ import { UsageBar, } from "@agy/ui"; import { ApiError, api } from "@/lib/api"; +import { useLiveChannel } from "@/lib/realtime"; import { useSession } from "@/lib/session"; export default function DashboardPage() { @@ -24,6 +25,32 @@ export default function DashboardPage() { const [loading, setLoading] = React.useState(true); const [error, setError] = React.useState(null); + // Live feed: audit events stream into the timeline; usage updates the bar. + useLiveChannel( + activeOrganization ? [`org:${activeOrganization.id}`] : [], + (event) => { + switch (event.type) { + case "audit.recorded": { + const data = event.data as { action: string; target_type: string; target_id: string }; + setAudit((current) => [ + { + id: event.id, + action: data.action, + actor_user_id: null, + target_type: data.target_type, + target_id: data.target_id, + created_at: new Date().toISOString(), + }, + ...current, + ]); + break; + } + default: + break; + } + }, + ); + const load = React.useCallback(async () => { setLoading(true); setError(null); diff --git a/apps/web/src/app/(app)/manager/page.tsx b/apps/web/src/app/(app)/manager/page.tsx index 703265e..54f764f 100644 --- a/apps/web/src/app/(app)/manager/page.tsx +++ b/apps/web/src/app/(app)/manager/page.tsx @@ -2,7 +2,7 @@ import Link from "next/link"; import * as React from "react"; -import type { Employee, Message, Task, Thread } from "@agy/contracts"; +import type { Employee, Message, Task, TaskStatus, Thread } from "@agy/contracts"; import { Badge, Button, @@ -17,6 +17,7 @@ import { } from "@agy/ui"; import { ApiError, api } from "@/lib/api"; import { CAN_RUN_TASKS, roleHasCapability } from "@/lib/permissions"; +import { useLiveChannel } from "@/lib/realtime"; import { useSession } from "@/lib/session"; /** @@ -48,6 +49,63 @@ export default function ManagerPage() { const bottomRef = React.useRef(null); + // Live feed: new messages in the active thread, and tasks created from chat. + const liveChannels = React.useMemo(() => { + const base: string[] = []; + if (orgId) base.push(`org:${orgId}`); + if (activeThreadId) base.push(`thread:${activeThreadId}`); + return base; + }, [orgId, activeThreadId]); + + useLiveChannel(liveChannels, (event) => { + switch (event.type) { + case "message.created": { + const data = event.data as { + thread_id: string; + message: { id: string; role: string; content: string; employee_id?: string | null }; + }; + if (data.thread_id === activeThreadId) { + setMessages((current) => [ + ...current, + { + id: data.message.id, + thread_id: data.thread_id, + role: data.message.role as Message["role"], + content: data.message.content, + employee_id: data.message.employee_id ?? null, + created_at: new Date().toISOString(), + }, + ]); + } + break; + } + case "task.created": { + const data = event.data as { task_id: string; title: string; status: string }; + setTasks((current) => { + if (current.some((t) => t.id === data.task_id)) return current; + return [ + { + id: data.task_id, + organization_id: orgId ?? "", + project_id: null, + thread_id: activeThreadId, + title: data.title, + description: "", + status: data.status as TaskStatus, + assigned_employee_id: null, + created_at: new Date().toISOString(), + updated_at: new Date().toISOString(), + }, + ...current, + ]; + }); + break; + } + default: + break; + } + }); + const load = React.useCallback(async () => { if (!orgId) { setLoading(false); diff --git a/apps/web/src/app/(app)/tasks/[id]/page.tsx b/apps/web/src/app/(app)/tasks/[id]/page.tsx index d5970f3..9eb86bd 100644 --- a/apps/web/src/app/(app)/tasks/[id]/page.tsx +++ b/apps/web/src/app/(app)/tasks/[id]/page.tsx @@ -2,7 +2,7 @@ import Link from "next/link"; import * as React from "react"; -import type { Employee, RunDetail, TaskDetail, TaskStatus, ToolCall } from "@agy/contracts"; +import type { Employee, RunDetail, RunStatus, TaskDetail, TaskStatus, ToolCall } from "@agy/contracts"; import { Badge, Button, @@ -17,6 +17,7 @@ import { } from "@agy/ui"; import { ApiError, api } from "@/lib/api"; import { CAN_RUN_TASKS, roleHasCapability } from "@/lib/permissions"; +import { useLiveChannel } from "@/lib/realtime"; import { useSession } from "@/lib/session"; const STATUS_LABELS: Record = { @@ -64,6 +65,112 @@ export default function TaskDetailPage({ const [running, setRunning] = React.useState(false); const [actionError, setActionError] = React.useState(null); + // Live feed: task status, run steps, approvals, artifacts — all push. + const channels = React.useMemo(() => { + const base: string[] = []; + if (orgId) base.push(`org:${orgId}`); + if (taskId) base.push(`task:${taskId}`); + if (selectedRun?.id) base.push(`run:${selectedRun.id}`); + return base; + }, [orgId, taskId, selectedRun?.id]); + + useLiveChannel(channels, (event) => { + switch (event.type) { + case "task.status": { + const data = event.data as { task_id: string; status: string }; + if (data.task_id === taskId) { + setTask((current) => (current ? { ...current, status: data.status as TaskStatus } : current)); + } + break; + } + case "run.step": { + const data = event.data as { + run_id: string; + task_id: string; + step: { id: string; index: number; kind: string; title: string; payload: Record }; + }; + setSelectedRun((current) => { + if (!current || current.id !== data.run_id) return current; + return { + ...current, + steps: [ + ...current.steps, + { + id: data.step.id, + index: data.step.index, + kind: data.step.kind, + title: data.step.title, + payload_json: data.step.payload, + created_at: new Date().toISOString(), + }, + ], + }; + }); + break; + } + case "run.status": { + const data = event.data as { run_id: string; task_id: string; status: string }; + setSelectedRun((current) => + current && current.id === data.run_id + ? { ...current, status: data.status as RunDetail["status"] } + : current, + ); + setTask((current) => { + if (!current) return current; + return { + ...current, + runs: current.runs.map((r) => + r.id === data.run_id ? { ...r, status: data.status as RunStatus } : r, + ), + }; + }); + break; + } + case "approval.pending": { + const data = event.data as { tool_call_id: string; tool_key: string; run_id: string }; + setApprovals((current) => { + if (current.some((a) => a.id === data.tool_call_id)) return current; + return [ + ...current, + { + id: data.tool_call_id, + tool_key: data.tool_key, + status: "awaiting_approval", + requires_approval: true, + arguments_json: {}, + result_json: {}, + created_at: new Date().toISOString(), + }, + ]; + }); + break; + } + case "artifact.created": { + const data = event.data as { artifact_id: string; run_id: string; name: string }; + setSelectedRun((current) => { + if (!current || current.id !== data.run_id) return current; + return { + ...current, + artifacts: [ + ...current.artifacts, + { + id: data.artifact_id, + name: data.name, + kind: "text", + size_bytes: 0, + content_ref: "", + created_at: new Date().toISOString(), + }, + ], + }; + }); + break; + } + default: + break; + } + }); + const load = React.useCallback(async () => { if (!orgId) { setLoading(false); diff --git a/apps/web/src/app/(app)/traces/page.tsx b/apps/web/src/app/(app)/traces/page.tsx index 4deb9d9..2ade898 100644 --- a/apps/web/src/app/(app)/traces/page.tsx +++ b/apps/web/src/app/(app)/traces/page.tsx @@ -18,7 +18,7 @@ import Link from "next/link"; import * as React from "react"; -import type { AuditEvent, Run, RunDetail } from "@agy/contracts"; +import type { AuditEvent, Run, RunDetail, RunStatus } from "@agy/contracts"; import { Badge, Button, @@ -32,6 +32,7 @@ import { Timeline, } from "@agy/ui"; import { ApiError, api } from "@/lib/api"; +import { useLiveChannel } from "@/lib/realtime"; import { useSession } from "@/lib/session"; /** How many tasks to fan out into. Bounds the request count. */ @@ -108,6 +109,71 @@ export default function TracesPage() { const [error, setError] = React.useState(null); const [detailLoading, setDetailLoading] = React.useState(false); + // Live feed: run steps and status changes update both the list and the detail. + useLiveChannel( + orgId ? [`org:${orgId}`] : [], + (event) => { + switch (event.type) { + case "run.step": { + const data = event.data as { + run_id: string; + task_id: string; + step: { id: string; index: number; kind: string; title: string; payload: Record }; + }; + setSelected((current) => { + if (!current || current.id !== data.run_id) return current; + return { + ...current, + steps: [ + ...current.steps, + { + id: data.step.id, + index: data.step.index, + kind: data.step.kind, + title: data.step.title, + payload_json: data.step.payload, + created_at: new Date().toISOString(), + }, + ], + }; + }); + break; + } + case "run.status": { + const data = event.data as { run_id: string; task_id: string; status: string }; + setRows((current) => + current.map((row) => + row.run.id === data.run_id ? { ...row, run: { ...row.run, status: data.status as RunStatus } } : row, + ), + ); + setSelected((current) => + current && current.id === data.run_id + ? { ...current, status: data.status as RunDetail["status"] } + : current, + ); + break; + } + case "audit.recorded": { + const data = event.data as { action: string; target_type: string; target_id: string }; + setAudit((current) => [ + { + id: event.id, + action: data.action, + actor_user_id: null, + target_type: data.target_type, + target_id: data.target_id, + created_at: new Date().toISOString(), + }, + ...current, + ]); + break; + } + default: + break; + } + }, + ); + const load = React.useCallback(async () => { if (!orgId) { setLoading(false); diff --git a/apps/web/src/app/layout.tsx b/apps/web/src/app/layout.tsx index 8f60dd9..dba5ee1 100644 --- a/apps/web/src/app/layout.tsx +++ b/apps/web/src/app/layout.tsx @@ -1,5 +1,6 @@ import type { Metadata } from "next"; import "./globals.css"; +import { RealtimeProvider } from "@/lib/realtime"; import { SessionProvider } from "@/lib/session"; export const metadata: Metadata = { @@ -11,7 +12,10 @@ export default function RootLayout({ children }: { children: React.ReactNode }) return ( - {children} + + {/* One socket for the whole app: every live feed rides on it. */} + {children} + ); diff --git a/apps/web/src/components/AppShell.tsx b/apps/web/src/components/AppShell.tsx index 1d7d32b..0fb1342 100644 --- a/apps/web/src/components/AppShell.tsx +++ b/apps/web/src/components/AppShell.tsx @@ -8,6 +8,7 @@ import { api } from "@/lib/api"; import { CommandPalette } from "@/components/CommandPalette"; import { NAV_ITEMS } from "@/lib/nav"; import { APP_VERSION_LABEL } from "@/lib/version"; +import { useRealtime } from "@/lib/realtime"; import { useSession } from "@/lib/session"; /** @@ -38,6 +39,38 @@ function VersionBadge() { ); } +/** A tiny dot that tells the user whether the live feed is connected. */ +function LiveIndicator() { + const { status } = useRealtime(); + const tone = + status === "live" + ? "success" + : status === "connecting" || status === "reconnecting" + ? "warning" + : "neutral"; + const label = + status === "live" + ? "زنده" + : status === "connecting" + ? "در حال اتصال…" + : status === "reconnecting" + ? "در حال اتصال مجدد…" + : "آفلاین"; + return ( + + + {label} + + ); +} + /** Credit badge — reads the organization's real balance, never a placeholder. */ function CreditBadge({ organizationId }: { organizationId: string | undefined }) { const [balance, setBalance] = React.useState(null); @@ -195,6 +228,9 @@ export function AppShell({ children }: { children: React.ReactNode }) { + + + diff --git a/apps/web/src/components/sandbox-panel.tsx b/apps/web/src/components/sandbox-panel.tsx index 1b6b599..5158d16 100644 --- a/apps/web/src/components/sandbox-panel.tsx +++ b/apps/web/src/components/sandbox-panel.tsx @@ -35,6 +35,7 @@ import { Stat, } from "@agy/ui"; import { ApiError, api } from "@/lib/api"; +import { useLiveChannel } from "@/lib/realtime"; const STATUS_TONES = { ready: "success", @@ -316,6 +317,53 @@ function SessionPanel({ const [snapshotName, setSnapshotName] = React.useState(""); const [snapshotBusy, setSnapshotBusy] = React.useState(false); + // Live feed: command output and session lifecycle events push live. + useLiveChannel( + [`session:${session.id}`], + (event) => { + switch (event.type) { + case "sandbox.command": { + const data = event.data as { + run_id: string; + command: string; + status: string; + exit_code: number | null; + stdout: string; + stderr: string; + duration_ms: number; + truncated: boolean; + }; + setRuns((current) => { + if (current.some((r) => r.id === data.run_id)) return current; + return [ + ...current, + { + id: data.run_id, + command: data.command, + exit_code: data.exit_code, + stdout: data.stdout, + stderr: data.stderr, + duration_ms: data.duration_ms, + truncated: data.truncated, + created_at: new Date().toISOString(), + }, + ]; + }); + break; + } + case "sandbox.session": { + const data = event.data as { session_id: string; status: string }; + if (data.session_id === session.id) { + onSessionChange({ ...session, status: data.status }); + } + break; + } + default: + break; + } + }, + ); + const [previewPort, setPreviewPort] = React.useState("8080"); const [previewUrl, setPreviewUrl] = React.useState(null); diff --git a/apps/web/src/lib/realtime.tsx b/apps/web/src/lib/realtime.tsx new file mode 100644 index 0000000..d435e1b --- /dev/null +++ b/apps/web/src/lib/realtime.tsx @@ -0,0 +1,135 @@ +"use client"; + +/** + * Realtime session state. + * + * The provider owns the one socket for the whole app and exposes two hooks: + * + * * ``useRealtime()`` — the connection status (for the live indicator) and + * the subscribe primitives, for the rare component that needs raw access. + * * ``useLiveChannel(channels, onEvent)`` — what every page actually uses: + * subscribe for as long as the component is mounted or the channels change, + * and keep the callback stable so an inline arrow does not resubscribe on + * every render. + * + * The socket follows the session: it connects once a token exists and drops + * when the user logs out, so no page has to think about auth. + */ + +import * as React from "react"; +import type { WsEvent, WsEventType } from "@agy/contracts"; +import { getRealtimeClient, type RealtimeStatus } from "@/lib/ws"; +import { getStoredToken } from "@/lib/api"; + +interface RealtimeContextValue { + status: RealtimeStatus; + subscribe: (channels: string[]) => void; + unsubscribe: (channels: string[]) => void; +} + +const RealtimeContext = React.createContext(null); + +export function RealtimeProvider({ children }: { children: React.ReactNode }) { + const client = React.useMemo(() => getRealtimeClient(), []); + const [status] = useRealtimeState(client); + + // The connection follows the token: present -> live, gone -> closed. + React.useEffect(() => { + const connect = () => { + if (getStoredToken()) client.connect(); + else client.disconnect(); + }; + connect(); + // ``storage`` fires when another tab logs in or out; polling covers the + // same-tab case where the token is written and read within one tick. + const interval = setInterval(connect, 2_000); + window.addEventListener("storage", connect); + return () => { + clearInterval(interval); + window.removeEventListener("storage", connect); + client.disconnect(); + }; + }, [client]); + + const value = React.useMemo( + () => ({ + status, + subscribe: (channels) => client.subscribe(channels), + unsubscribe: (channels) => client.unsubscribe(channels), + }), + [status, client], + ); + + return {children}; +} + +function useRealtimeState(client: ReturnType): [RealtimeStatus] { + const [status, setStatus] = React.useState(() => client.getStatus()); + React.useEffect(() => client.onStatus(setStatus), [client]); + return [status]; +} + +export function useRealtime(): RealtimeContextValue { + const value = React.useContext(RealtimeContext); + // A missing provider is a wiring bug, not a runtime condition: fail loudly + // in development rather than silently showing a page that never updates. + if (!value) { + throw new Error("useRealtime must be used inside "); + } + return value; +} + +/** + * Subscribe to channels for the lifetime of the component. + * + * ``onEvent`` is kept in a ref, so the subscription is not re-established + * because the callback identity changed — the classic cause of a live feed + * that resubscribes on every render and drops events in between. + */ +export function useLiveChannel( + channels: string[], + onEvent: (event: WsEvent) => void, +): void { + const client = React.useMemo(() => getRealtimeClient(), []); + const handler = React.useRef(onEvent); + React.useEffect(() => { + handler.current = onEvent; + }, [onEvent]); + + // A stable key: the effect re-runs when the channel *set* changes, not when + // a new array with the same contents is passed. + const channelKey = channels.join("\u0000"); + + React.useEffect(() => { + const active = channelKey ? channelKey.split("\u0000") : []; + if (active.length === 0) return; + client.subscribe(active); + const off = client.onEvent((event) => handler.current(event)); + return () => { + off(); + client.unsubscribe(active); + }; + }, [client, channelKey]); +} + +/** + * Subscribe to one event type only — the common case for a page that renders + * exactly one live surface (a timeline, a terminal, a counter). + */ +export function useLiveEvent( + channels: string[], + type: WsEventType, + onEvent: (event: WsEvent) => void, +): void { + const handler = React.useRef(onEvent); + React.useEffect(() => { + handler.current = onEvent; + }, [onEvent]); + + useLiveChannel( + channels, + React.useCallback((event: WsEvent) => { + if (event.type === type) handler.current(event); + }, [type]), + ); +} diff --git a/apps/web/src/lib/ws.ts b/apps/web/src/lib/ws.ts new file mode 100644 index 0000000..3ec63ce --- /dev/null +++ b/apps/web/src/lib/ws.ts @@ -0,0 +1,270 @@ +/** + * The realtime client — one WebSocket for the whole app. + * + * Design decisions worth knowing before editing: + * + * * **One socket, shared.** Pages subscribe and unsubscribe through + * ``RealtimeProvider``; they never own a connection. A dashboard and a task + * page both being live must not mean two sockets and two auth round-trips. + * * **The token is a function, not a value.** ``getToken`` is read at connect + * time, so a login/logout cycle in another tab reconnects correctly instead + * of pushing on a dead credential. + * * **Reconnect with backoff, capped.** Networks flap; a reconnect storm is + * how a flapping network becomes an outage. The cap is deliberately low + * enough that "live" feels live again within seconds. + * * **Heartbeat.** A silent middlebox will happily hold a dead socket open + * forever. A ping every ``HEARTBEAT_INTERVAL_MS`` plus a pong timeout moves + * that failure from "the page quietly stops updating" to "the socket + * reconnects and updates resume". + */ + +import type { WsClientFrame, WsEvent, WsServerFrame } from "@agy/contracts"; +import { getStoredToken } from "@/lib/api"; + +/** Where the socket lives: the API origin, upgraded from http(s) to ws(s). */ +export function websocketUrl(path = "/api/ws"): string { + const base = + process.env.NEXT_PUBLIC_API_URL ?? + (typeof window !== "undefined" ? window.location.origin : "http://localhost:8000"); + const url = new URL(base); + url.protocol = url.protocol === "https:" ? "wss:" : "ws:"; + url.pathname = path; + return url.toString(); +} + +export type RealtimeStatus = "idle" | "connecting" | "live" | "reconnecting" | "offline"; + +const MIN_BACKOFF_MS = 500; +const MAX_BACKOFF_MS = 8_000; +const HEARTBEAT_INTERVAL_MS = 25_000; +/** If no pong arrives within this window the socket is considered dead. */ +const PONG_TIMEOUT_MS = 10_000; + +type EventListener = (event: WsEvent) => void; +type StatusListener = (status: RealtimeStatus) => void; + +export class RealtimeClient { + private socket: WebSocket | null = null; + private status: RealtimeStatus = "idle"; + private readonly listeners = new Set(); + private readonly statusListeners = new Set(); + private readonly subscriptions = new Set(); + private heartbeatTimer: ReturnType | null = null; + private pongTimer: ReturnType | null = null; + private reconnectAttempts = 0; + private reconnectTimer: ReturnType | null = null; + private closedByUs = false; + + constructor(private readonly getToken: () => string | null = getStoredToken) {} + + // -- lifecycle ----------------------------------------------------------- + + connect(): void { + const token = this.getToken(); + if (!token) { + // No session yet: nothing to push. The provider reconnects on login. + this.setStatus("idle"); + return; + } + if ( + this.socket && + (this.socket.readyState === WebSocket.OPEN || this.socket.readyState === WebSocket.CONNECTING) + ) { + return; + } + + this.closedByUs = false; + this.setStatus(this.reconnectAttempts > 0 ? "reconnecting" : "connecting"); + + const socket = new WebSocket(`${websocketUrl()}?token=${encodeURIComponent(token)}`); + this.socket = socket; + + socket.onopen = () => { + this.reconnectAttempts = 0; + this.setStatus("live"); + this.startHeartbeat(); + // Re-subscribe: a new socket knows nothing about what this client cared + // about, and pages do not (and should not) re-run their effects to fix it. + if (this.subscriptions.size > 0) { + this.send({ action: "subscribe", channels: [...this.subscriptions] }); + } + }; + + socket.onmessage = (message) => this.handleFrame(message.data); + + socket.onclose = () => { + this.stopHeartbeat(); + this.socket = null; + if (this.closedByUs) return; + this.scheduleReconnect(); + }; + + socket.onerror = () => { + // A close always follows an error; the close handler does the work. + // This only stops an unhandled rejection in stricter runtimes. + }; + } + + disconnect(): void { + this.closedByUs = true; + if (this.reconnectTimer) { + clearTimeout(this.reconnectTimer); + this.reconnectTimer = null; + } + this.stopHeartbeat(); + this.reconnectAttempts = 0; + if (this.socket) { + this.socket.close(); + this.socket = null; + } + this.setStatus("idle"); + } + + private scheduleReconnect(): void { + this.setStatus("reconnecting"); + const delay = Math.min( + MAX_BACKOFF_MS, + MIN_BACKOFF_MS * 2 ** this.reconnectAttempts, + ); + this.reconnectAttempts += 1; + this.reconnectTimer = setTimeout(() => this.connect(), delay); + } + + // -- subscriptions --------------------------------------------------------- + + subscribe(channels: string[]): void { + const fresh = channels.filter((channel) => !this.subscriptions.has(channel)); + if (fresh.length === 0) return; + fresh.forEach((channel) => this.subscriptions.add(channel)); + this.send({ action: "subscribe", channels: fresh }); + } + + unsubscribe(channels: string[]): void { + const active = channels.filter((channel) => this.subscriptions.has(channel)); + if (active.length === 0) return; + active.forEach((channel) => this.subscriptions.delete(channel)); + this.send({ action: "unsubscribe", channels: active }); + } + + activeSubscriptions(): string[] { + return [...this.subscriptions]; + } + + // -- listeners -------------------------------------------------------------- + + onEvent(listener: EventListener): () => void { + this.listeners.add(listener); + return () => this.listeners.delete(listener); + } + + onStatus(listener: StatusListener): () => void { + this.statusListeners.add(listener); + return () => this.statusListeners.delete(listener); + } + + getStatus(): RealtimeStatus { + return this.status; + } + + // -- frames ----------------------------------------------------------------- + + private send(frame: WsClientFrame): void { + if (this.socket && this.socket.readyState === WebSocket.OPEN) { + this.socket.send(JSON.stringify(frame)); + } + // A frame sent before the socket is open is not an error: on open, the + // full subscription set is replayed anyway. + } + + private handleFrame(raw: unknown): void { + if (typeof raw !== "string") return; + let frame: WsServerFrame; + try { + frame = JSON.parse(raw) as WsServerFrame; + } catch { + return; + } + + switch (frame.type) { + case "event": + this.emitEvent(frame.event); + break; + case "pong": + this.stopPongTimer(); + break; + case "subscribed": + // A rejected channel is a server answer, not a client error: the page + // keeps its initial data and the console shows why it is not live. + if (frame.rejected.length > 0) { + console.warn("realtime: channels rejected by the server", frame.rejected); + } + break; + case "error": + console.warn("realtime: server error frame", frame.message); + break; + default: + break; + } + } + + private emitEvent(event: WsEvent): void { + for (const listener of [...this.listeners]) { + try { + listener(event); + } catch (error) { + // One broken page must not silence every other page's feed. + console.error("realtime: event listener failed", error); + } + } + } + + private setStatus(status: RealtimeStatus): void { + if (this.status === status) return; + this.status = status; + for (const listener of [...this.statusListeners]) { + try { + listener(status); + } catch (error) { + console.error("realtime: status listener failed", error); + } + } + } + + // -- heartbeat ----------------------------------------------------------------- + + private startHeartbeat(): void { + this.stopHeartbeat(); + this.heartbeatTimer = setInterval(() => { + this.send({ action: "ping" }); + this.pongTimer = setTimeout(() => { + // No pong: the socket is a zombie. Close it so the reconnect path runs. + this.socket?.close(); + }, PONG_TIMEOUT_MS); + }, HEARTBEAT_INTERVAL_MS); + } + + private stopHeartbeat(): void { + if (this.heartbeatTimer) { + clearInterval(this.heartbeatTimer); + this.heartbeatTimer = null; + } + this.stopPongTimer(); + } + + private stopPongTimer(): void { + if (this.pongTimer) { + clearTimeout(this.pongTimer); + this.pongTimer = null; + } + } +} + +let sharedClient: RealtimeClient | null = null; + +/** The process-wide client the provider owns. */ +export function getRealtimeClient(): RealtimeClient { + if (!sharedClient) { + sharedClient = new RealtimeClient(); + } + return sharedClient; +} diff --git a/packages/contracts/src/index.ts b/packages/contracts/src/index.ts index fdad3a0..ec837f3 100644 --- a/packages/contracts/src/index.ts +++ b/packages/contracts/src/index.ts @@ -666,3 +666,86 @@ export interface ErrorQuery { limit?: number; offset?: number; } + +// --------------------------------------------------------------------------- +// Realtime (WebSocket) +// +// The live feed is the same product surface as the REST API: one envelope +// shape, mirrored key-for-key from ``app/application/ports/realtime.py``. +// A page that switches on ``WsEventType`` never needs to know which part of +// the backend produced the event. +// +// Channel naming is what a page subscribes to — ``org:{id}`` for the tenant +// feed, ``task:{id}`` for one task, and so on. The server authorizes every +// subscribe against the caller's memberships, so a channel the user may not +// see simply comes back rejected. +// --------------------------------------------------------------------------- + +export type WsEventType = + | "run.step" + | "run.status" + | "task.status" + | "task.created" + | "chat.pending" + | "message.created" + | "approval.pending" + | "artifact.created" + | "sandbox.command" + | "sandbox.session" + | "sandbox.reaped" + | "usage.recorded" + | "billing.wallet" + | "error.recorded" + | "audit.recorded" + | "worker.status"; + +/** One pushed event, as the browser receives it. */ +export interface WsEvent> { + id: string; + type: WsEventType; + data: T; + organization_id: string | null; + /** ISO-8601 UTC; render with ``toLocaleString`` like every other timestamp. */ + at: string; +} + +/** Every frame the server sends. Control frames and events share ``type``. */ +export type WsServerFrame = + | { type: "hello"; user_id: string; server_time: string } + | { type: "subscribed"; accepted: string[]; rejected: string[] } + | { type: "unsubscribed"; channels: string[] } + | { type: "pong" } + | { type: "error"; message: string } + | { type: "event"; event: WsEvent }; + +/** Every frame the client may send. */ +export type WsClientFrame = + | { action: "subscribe"; channels: string[] } + | { action: "unsubscribe"; channels: string[] } + | { action: "ping" }; + +/** Live run-step payload — what the task timeline renders as it arrives. */ +export interface WsRunStepData { + run_id: string; + task_id: string; + step: { + id: string; + index: number; + kind: string; + title: string; + payload: Record; + }; +} + +/** Live sandbox command payload — the terminal panel's stream. */ +export interface WsSandboxCommandData { + session_id: string; + run_id: string; + command: string; + status: string; + exit_code: number | null; + stdout: string; + stderr: string; + duration_ms: number; + truncated: boolean; +}