by-framework 0.2.2.dev4__py3-none-any.whl → 0.2.2.dev6__py3-none-any.whl

by_framework/__init__.py

@@ -5,6 +5,7 @@ Allows developers to quickly start a Redis-based agent node by inheriting
 from `GatewayWorker` and running `run_worker`.
 """
 
+from .admin import WorkerManager
 from .client.byai_client import ByaiGatewayClient
 from .client.client import (
     CancelTaskResponse,
@@ -59,9 +60,12 @@ from .core.protocol.commands import (
     AskAgentCommand,
     BaseCommand,
     CancelTaskCommand,
+    EvictWorkerCommand,
     GatewayCommand,
     ReloadPluginsCommand,
     ResumeCommand,
+    ResumeWorkerCommand,
+    SuspendWorkerCommand,
     command_from_dict,
     get_registered_command,
     register_command,
@@ -152,6 +156,7 @@ __all__ = [
     "close_redis",
     "Redis",
     "WorkerRegistry",
+    "WorkerManager",
     "check_agent_type_online",
     "check_worker_online",
     "RedisKeys",
@@ -168,6 +173,9 @@ __all__ = [
     "ResumeCommand",
     "CancelTaskCommand",
     "ReloadPluginsCommand",
+    "SuspendWorkerCommand",
+    "ResumeWorkerCommand",
+    "EvictWorkerCommand",
     "GatewayCommand",
     "command_from_dict",
     "register_command",

by_framework/admin/__init__.py

@@ -0,0 +1,5 @@
+"""Admin-side APIs for worker lifecycle and admission control."""
+
+from .worker_manager import WorkerManager
+
+__all__ = ["WorkerManager"]

by_framework/admin/cli.py

@@ -0,0 +1,344 @@
+"""by-admin — by-framework cluster admin CLI."""
+
+# pylint: disable=global-statement
+
+from __future__ import annotations
+
+import asyncio
+import json
+from typing import Optional
+
+import typer
+from rich.console import Console
+from rich.table import Table
+
+from by_framework.admin.worker_manager import WorkerManager
+from by_framework.common.redis_client import init_redis_from_url
+from by_framework.core.registry import WorkerRegistry
+from by_framework.metrics.snapshot import (
+    build_observability_snapshot,
+    load_history_from_redis,
+)
+
+# --------------------------------------------------------------------------- #
+#  App hierarchy
+# --------------------------------------------------------------------------- #
+
+app = typer.Typer(
+    name="by-admin",
+    no_args_is_help=True,
+    help="by-framework cluster admin CLI",
+)
+worker_app = typer.Typer(no_args_is_help=True)
+type_app = typer.Typer(no_args_is_help=True)
+metrics_app = typer.Typer(no_args_is_help=True)
+
+app.add_typer(worker_app, name="worker", help="Worker lifecycle management")
+app.add_typer(type_app, name="type", help="Agent-type admission control")
+app.add_typer(metrics_app, name="metrics", help="Metrics and observability")
+
+# --------------------------------------------------------------------------- #
+#  Global state + helpers
+# --------------------------------------------------------------------------- #
+
+console = Console()
+err_console = Console(stderr=True)
+_redis_url = "redis://localhost:6379/0"
+
+
+@app.callback()
+def _global(
+    redis_url: Optional[str] = typer.Option(
+        None,
+        "--redis-url",
+        envvar=["BYAI_REDIS_URL", "REDIS_URL"],
+        help="Redis connection URL  [env: BYAI_REDIS_URL]",
+    ),
+):
+    """by-framework cluster admin CLI."""
+    global _redis_url
+    if redis_url:
+        _redis_url = redis_url
+
+
+def _run(coro):  # type: ignore[no-untyped-def]
+    return asyncio.run(coro)
+
+
+def _die(msg: str, code: int = 1) -> None:
+    err_console.print(f"[red]Error:[/red] {msg}")
+    raise typer.Exit(code)
+
+
+def _get_redis():
+    return init_redis_from_url(_redis_url)
+
+
+# --------------------------------------------------------------------------- #
+#  worker subcommands
+# --------------------------------------------------------------------------- #
+
+_LIFECYCLE_COLORS = {"active": "green", "suspended": "yellow", "evicted": "red"}
+
+
+@worker_app.command("list")
+def worker_list(
+    agent_type: Optional[str] = typer.Option(
+        None, "--type", "-t", help="Filter by agent type"
+    ),
+    as_json: bool = typer.Option(False, "--json", help="Output as JSON"),
+):
+    """List all online workers."""
+
+    async def _inner():
+        registry = WorkerRegistry(_get_redis())
+        workers = await registry.get_all_workers()
+
+        rows = [{"worker_id": wid, **data} for wid, data in workers.items()]
+        if agent_type:
+            rows = [r for r in rows if agent_type in r.get("agent_types", [])]
+
+        if as_json:
+            typer.echo(json.dumps(rows, default=str))
+            return
+
+        if not rows:
+            console.print("[dim]No workers found.[/dim]")
+            return
+
+        table = Table(show_header=True, header_style="bold")
+        table.add_column("Worker ID")
+        table.add_column("Lifecycle")
+        table.add_column("Agent Types")
+        table.add_column("IP")
+        table.add_column("Last Seen (ms)")
+
+        for w in rows:
+            lc = w.get("lifecycle", "active")
+            color = _LIFECYCLE_COLORS.get(lc, "white")
+            table.add_row(
+                w["worker_id"],
+                f"[{color}]{lc}[/{color}]",
+                ", ".join(w.get("agent_types", [])),
+                w.get("ip_address", ""),
+                str(w.get("last_seen", "")),
+            )
+
+        console.print(table)
+        console.print(f"\n[dim]{len(rows)} worker(s)[/dim]")
+
+    _run(_inner())
+
+
+@worker_app.command("info")
+def worker_info(
+    worker_id: str = typer.Argument(..., help="Worker ID"),
+    as_json: bool = typer.Option(False, "--json", help="Output as JSON"),
+):
+    """Show detailed info for a single worker."""
+
+    async def _inner():
+        registry = WorkerRegistry(_get_redis())
+        workers = await registry.get_all_workers()
+        worker = workers.get(worker_id)
+
+        if worker is None:
+            _die(f"worker '{worker_id}' not found")
+            return
+
+        result = {"worker_id": worker_id, **worker}
+        if as_json:
+            typer.echo(json.dumps(result, default=str))
+            return
+
+        lc = result.get("lifecycle", "active")
+        color = _LIFECYCLE_COLORS.get(lc, "white")
+        lc_reason = result.get("lifecycle_reason", "")
+        agent_types_str = ", ".join(result.get("agent_types", []))
+        ip_address = result.get("ip_address", "")
+        last_seen = result.get("last_seen", "")
+        console.print(f"[bold]Worker:[/bold] {worker_id}")
+        console.print(f"  lifecycle:        [{color}]{lc}[/{color}]")
+        console.print(f"  lifecycle_reason: {lc_reason}")
+        console.print(f"  agent_types:      {agent_types_str}")
+        console.print(f"  ip_address:       {ip_address}")
+        console.print(f"  last_seen (ms):   {last_seen}")
+
+    _run(_inner())
+
+
+@worker_app.command("suspend")
+def worker_suspend(
+    worker_id: str = typer.Argument(..., help="Worker ID"),
+    reason: str = typer.Option("", "--reason", "-r", help="Reason for suspension"),
+):
+    """Suspend a worker (stops consuming, stays online)."""
+
+    async def _inner():
+        mgr = WorkerManager(_get_redis())
+        await mgr.suspend_worker(worker_id, reason=reason)
+        console.print(f"[green]✓[/green] Suspended {worker_id}")
+
+    _run(_inner())
+
+
+@worker_app.command("resume")
+def worker_resume(
+    worker_id: str = typer.Argument(..., help="Worker ID"),
+):
+    """Resume a suspended worker."""
+
+    async def _inner():
+        mgr = WorkerManager(_get_redis())
+        await mgr.resume_worker(worker_id)
+        console.print(f"[green]✓[/green] Resumed {worker_id}")
+
+    _run(_inner())
+
+
+@worker_app.command("evict")
+def worker_evict(
+    worker_id: str = typer.Argument(..., help="Worker ID"),
+    force: bool = typer.Option(
+        False, "--force", help="Cancel in-flight tasks immediately"
+    ),
+):
+    """Evict a worker (clears lease, removes from routing)."""
+
+    async def _inner():
+        mgr = WorkerManager(_get_redis())
+        await mgr.evict_worker(worker_id, force=force)
+        console.print(f"[green]✓[/green] Evicted {worker_id}")
+
+    _run(_inner())
+
+
+# --------------------------------------------------------------------------- #
+#  type subcommands
+# --------------------------------------------------------------------------- #
+
+
+@type_app.command("denylist")
+def type_denylist(
+    agent_type: str = typer.Argument(..., help="Agent type name"),
+    as_json: bool = typer.Option(False, "--json", help="Output as JSON"),
+):
+    """Show all workers denied for an agent type."""
+
+    async def _inner():
+        mgr = WorkerManager(_get_redis())
+        denied = await mgr.get_type_denylist(agent_type)
+        if as_json:
+            typer.echo(json.dumps(denied))
+            return
+        if not denied:
+            console.print(f"[dim]No denied workers for type '{agent_type}'[/dim]")
+            return
+        for wid in denied:
+            console.print(f"  {wid}")
+
+    _run(_inner())
+
+
+@type_app.command("deny")
+def type_deny(
+    agent_type: str = typer.Argument(..., help="Agent type name"),
+    worker_id: str = typer.Argument(..., help="Worker ID to deny"),
+):
+    """Deny a worker from consuming an agent type."""
+
+    async def _inner():
+        mgr = WorkerManager(_get_redis())
+        await mgr.deny_worker_for_type(agent_type, worker_id)
+        console.print(
+            f"[green]✓[/green] Denied {worker_id} from consuming '{agent_type}'"
+        )
+
+    _run(_inner())
+
+
+@type_app.command("allow")
+def type_allow(
+    agent_type: str = typer.Argument(..., help="Agent type name"),
+    worker_id: str = typer.Argument(..., help="Worker ID to allow"),
+):
+    """Remove denial for a worker on an agent type."""
+
+    async def _inner():
+        mgr = WorkerManager(_get_redis())
+        await mgr.allow_worker_for_type(agent_type, worker_id)
+        console.print(f"[green]✓[/green] Allowed {worker_id} on '{agent_type}'")
+
+    _run(_inner())
+
+
+# --------------------------------------------------------------------------- #
+#  metrics subcommands
+# --------------------------------------------------------------------------- #
+
+
+@metrics_app.command("snapshot")
+def metrics_snapshot(
+    as_json: bool = typer.Option(False, "--json", help="Output as JSON"),
+):
+    """Print current cluster observability snapshot."""
+
+    async def _inner():
+        snapshot = await build_observability_snapshot(_get_redis())
+        if as_json:
+            typer.echo(json.dumps(snapshot, default=str))
+            return
+        totals = snapshot.get("totals", {})
+        workers_online = totals.get("workers_online", 0)
+        agent_types_n = totals.get("agent_types", 0)
+        active_execs = totals.get("active_executions", 0)
+        queue_depth = snapshot.get("queue_depth_total", 0)
+        console.print("[bold]Cluster Snapshot[/bold]")
+        console.print(f"  Workers online:     {workers_online}")
+        console.print(f"  Agent types:        {agent_types_n}")
+        console.print(f"  Active executions:  {active_execs}")
+        console.print(f"  Queue depth total:  {queue_depth}")
+
+    _run(_inner())
+
+
+@metrics_app.command("history")
+def metrics_history(
+    limit: int = typer.Option(
+        20, "--limit", "-n", help="Number of history points to show"
+    ),
+    as_json: bool = typer.Option(False, "--json", help="Output as JSON"),
+):
+    """Print recent metrics history points."""
+
+    async def _inner():
+        points = await load_history_from_redis(_get_redis(), limit=limit)
+        if as_json:
+            typer.echo(json.dumps(points, default=str))
+            return
+        if not points:
+            console.print("[dim]No history points found.[/dim]")
+            return
+        table = Table(show_header=True, header_style="bold")
+        table.add_column("Timestamp (ms)")
+        table.add_column("Workers")
+        table.add_column("Active")
+        table.add_column("Queue Depth")
+        for p in points:
+            table.add_row(
+                str(p.get("generated_at", "")),
+                str(p.get("workers_online", "")),
+                str(p.get("active_executions", "")),
+                str(p.get("queue_depth_total", "")),
+            )
+        console.print(table)
+
+    _run(_inner())
+
+
+# --------------------------------------------------------------------------- #
+#  Entry point
+# --------------------------------------------------------------------------- #
+
+if __name__ == "__main__":
+    app()

by_framework/admin/worker_manager.py

@@ -0,0 +1,146 @@
+"""
+WorkerManager — admin-side API for controlling worker lifecycle and agent-type access.
+
+Responsibilities:
+  - Suspend / resume / evict individual workers (lifecycle control).
+  - Deny / allow workers from consuming a specific agent_type (admission control).
+
+Lifecycle commands are delivered via two channels:
+  1. Push: XADD to byai_gateway:ctrl:worker:{worker_id} (immediate delivery).
+  2. Pull: HSET to byai_gateway:registry:worker:admin:{worker_id} (durable fallback,
+     read by the worker's heartbeat loop every heartbeat interval).
+"""
+
+from __future__ import annotations
+
+import uuid
+from typing import Any, Optional
+
+from by_framework.common.constants import RedisKeys
+from by_framework.common.redis_client import Redis, get_redis
+from by_framework.core.protocol.commands import (
+    EvictWorkerCommand,
+    ResumeWorkerCommand,
+    SuspendWorkerCommand,
+)
+from by_framework.core.protocol.message_header import MessageHeader
+from by_framework.core.registry import WorkerRegistry
+
+
+def _admin_header() -> MessageHeader:
+    return MessageHeader(
+        session_id="admin",
+        trace_id=uuid.uuid4().hex,
+        message_id=uuid.uuid4().hex,
+    )
+
+
+class WorkerManager:
+    """Admin API for worker lifecycle and agent-type admission control."""
+
+    def __init__(
+        self,
+        redis_client: Optional[Redis] = None,
+        registry: Optional[WorkerRegistry] = None,
+    ):
+        self.redis = redis_client or get_redis()
+        self.registry = registry or WorkerRegistry(self.redis)
+
+    # ------------------------------------------------------------------
+    # Lifecycle control
+    # ------------------------------------------------------------------
+
+    async def suspend_worker(self, worker_id: str, reason: str = "") -> None:
+        """Pause a running worker from consuming new tasks.
+
+        The worker finishes in-flight tasks but stops accepting new ones.
+        Immediately removes the worker from all agent_type:members sets so that
+        routing skips it at once; the worker re-adds itself on resume.
+        """
+        await self.registry.set_worker_admin_state(worker_id, "suspended", reason)
+        await self.registry.remove_worker_from_type_members(worker_id)
+        command = SuspendWorkerCommand(header=_admin_header(), reason=reason)
+        await self.redis.xadd(
+            RedisKeys.worker_ctrl_stream(worker_id),
+            command.to_redis_payload(),
+        )
+
+    async def resume_worker(self, worker_id: str) -> None:
+        """Resume a previously suspended worker.
+
+        Re-adds the worker to all agent_type:members sets immediately (respecting
+        the denylist), so routing can reach it again without waiting for the next
+        heartbeat cycle.
+        """
+        await self.registry.set_worker_admin_state(worker_id, "active", "")
+        await self.registry.restore_worker_to_type_members(worker_id)
+        command = ResumeWorkerCommand(header=_admin_header())
+        await self.redis.xadd(
+            RedisKeys.worker_ctrl_stream(worker_id),
+            command.to_redis_payload(),
+        )
+
+    async def evict_worker(
+        self,
+        worker_id: str,
+        *,
+        force: bool = False,
+        reason: str = "",
+    ) -> None:
+        """Shut down a worker.
+
+        Args:
+            worker_id: Target worker ID.
+            force: When True, cancels in-flight tasks immediately instead of
+                waiting for them to finish.
+            reason: Human-readable eviction reason.
+
+        Immediately removes the worker from all agent_type:members sets so
+        routing stops sending new messages before the heartbeat TTL expires.
+        """
+        await self.registry.set_worker_admin_state(worker_id, "evicted", reason)
+        await self.registry.remove_worker_from_type_members(worker_id)
+        command = EvictWorkerCommand(header=_admin_header(), reason=reason, force=force)
+        await self.redis.xadd(
+            RedisKeys.worker_ctrl_stream(worker_id),
+            command.to_redis_payload(),
+        )
+
+    # ------------------------------------------------------------------
+    # Agent-type admission control
+    # ------------------------------------------------------------------
+
+    async def deny_worker_for_type(self, agent_type: str, worker_id: str) -> None:
+        """Prevent worker_id from consuming the agent_type stream.
+
+        Takes effect on the worker's next heartbeat cycle (or immediately if
+        the worker checks the denylist before each XREADGROUP call).
+        """
+        await self.registry.deny_worker_for_type(agent_type, worker_id)
+
+    async def allow_worker_for_type(self, agent_type: str, worker_id: str) -> None:
+        """Remove worker_id from the denylist for agent_type."""
+        await self.registry.allow_worker_for_type(agent_type, worker_id)
+
+    async def get_type_denylist(self, agent_type: str) -> list[str]:
+        """Return all worker_ids currently denied for agent_type."""
+        return await self.registry.get_agent_type_denylist(agent_type)
+
+    # ------------------------------------------------------------------
+    # Status queries
+    # ------------------------------------------------------------------
+
+    async def get_worker_admin_state(self, worker_id: str) -> dict[str, Any]:
+        """Return the admin-controlled state for a worker.
+
+        Returns an empty dict when no admin state has been set (default active).
+        """
+        return await self.registry.get_worker_admin_state(worker_id)
+
+    async def clear_worker_admin_state(self, worker_id: str) -> None:
+        """Remove all admin state for a worker, restoring default-active behaviour."""
+        await self.registry.clear_worker_admin_state(worker_id)
+
+    async def allow_worker_rejoin(self, worker_id: str) -> None:
+        """Allow a previously evicted offline worker ID to join on next startup."""
+        await self.clear_worker_admin_state(worker_id)

by_framework/client/client.py

@@ -43,7 +43,16 @@ from by_framework.core.protocol.responses import (
 )
 from by_framework.core.registry import WorkerRegistry
 from by_framework.errors import WorkerRegistryNotSetError
-from by_framework.trace.span_recorder import (SpanRecorder, TraceSpan, str_to_uint64)
+from by_framework.trace.span_recorder import (
+    ObservabilityConfig,
+    SpanRecorder,
+    TraceSpan,
+    build_observability_config,
+    sanitize_io_value,
+    str_to_uint64,
+)
+from by_framework.trace.trace_schema import TraceRecord
+from by_framework.trace.trace_writer import TraceWriteClient
 
 if TYPE_CHECKING:
     pass
@@ -89,13 +98,20 @@ class GatewayClient:
         redis_client: Optional[Redis] = None,
         interceptors: Optional[List[GatewayInterceptor]] = None,
         span_recorder: Optional[SpanRecorder] = None,
+        obs_config: Optional[ObservabilityConfig] = None,
     ):
         self.registry = registry
         self.redis = (
             redis_client or (registry.redis if registry else None) or get_redis()
         )
         self.interceptors = interceptors or []
-        self.span_recorder = span_recorder or SpanRecorder(self.redis)
+        self._obs_config = obs_config or build_observability_config()
+        self.span_recorder = span_recorder or SpanRecorder(
+            self.redis, config=self._obs_config
+        )
+        self._trace_writer = TraceWriteClient(
+            self.redis, ttl_seconds=self._obs_config.ttl_seconds
+        )
         self._langfuse_dispatch_fn = self._resolve_langfuse_dispatch_fn()
 
     @staticmethod
@@ -617,6 +633,20 @@ class GatewayClient:
             message_id = f"{MESSAGE_ID_PREFIX}{uuid.uuid4().hex[:8]}"
         if not trace_id:
             trace_id = uuid.uuid4().hex
+        trace_start_ts = int(time.time() * 1000)
+        is_root_dispatch = not params.get("parent_message_id", "")
+
+        # Write trace root start so the trace is visible even before the worker
+        # picks up the task.  Only written for root dispatches (no parent).
+        if is_root_dispatch:
+            await self._write_trace_root_start(
+                trace_id=trace_id,
+                message_id=message_id,
+                session_id=str(params.get("session_id", "")),
+                target_agent_type=str(params.get("target_agent_type", "")),
+                content=params.get("content"),
+                start_ts=trace_start_ts,
+            )
 
         metadata = dict(params.get("metadata", {}) or {})
         trace_parent_span_id = metadata.pop("trace_parent_span_id", "")
@@ -743,6 +773,13 @@ class GatewayClient:
                         output={"success": False, "error": availability.error},
                         error=availability.error,
                     )
+                    if is_root_dispatch:
+                        await self._write_trace_root_end(
+                            trace_id=trace_id,
+                            status="FAILED",
+                            end_ts=int(time.time() * 1000),
+                            error=availability.error or "",
+                        )
                     return response
                 if availability.status == AvailabilityStatus.QUEUE_PENDING:
                     should_dispatch_control = False
@@ -759,6 +796,13 @@ class GatewayClient:
                 output={"success": False, "error": str(err)},
                 error=str(err),
             )
+            if is_root_dispatch:
+                await self._write_trace_root_end(
+                    trace_id=trace_id,
+                    status="FAILED",
+                    end_ts=int(time.time() * 1000),
+                    error=str(err),
+                )
             return SendMessageResponse(
                 success=False,
                 status=ExecutionStatus.FAILED,
@@ -775,6 +819,13 @@ class GatewayClient:
                 output={"success": False, "error": str(err)},
                 error=str(err),
             )
+            if is_root_dispatch:
+                await self._write_trace_root_end(
+                    trace_id=trace_id,
+                    status="FAILED",
+                    end_ts=int(time.time() * 1000),
+                    error=str(err),
+                )
             return SendMessageResponse(
                 success=False,
                 status=ExecutionStatus.FAILED,
@@ -959,3 +1010,52 @@ class GatewayClient:
             logger.warning(
                 "Failed to record client dispatch span: %s", err, exc_info=True
             )
+
+    async def _write_trace_root_start(
+        self,
+        *,
+        trace_id: str,
+        message_id: str,
+        session_id: str,
+        target_agent_type: str,
+        content: Any,
+        start_ts: int,
+    ) -> None:
+        """Write trace root start — best-effort, never propagates errors."""
+        try:
+            await self._trace_writer.record_trace(
+                TraceRecord(
+                    trace_id=trace_id,
+                    name=target_agent_type,
+                    session_id=session_id,
+                    root_message_id=message_id,
+                    root_agent_type=target_agent_type,
+                    input=sanitize_io_value(content, self._obs_config),
+                    status="QUEUED",
+                    start_ts=start_ts,
+                )
+            )
+        except Exception:  # pylint: disable=broad-exception-caught
+            pass
+
+    async def _write_trace_root_end(
+        self,
+        *,
+        trace_id: str,
+        status: str,
+        end_ts: int,
+        error: str = "",
+    ) -> None:
+        """Write trace root end on routing failure — best-effort."""
+        try:
+            meta: dict = {"error": error} if error else {}
+            await self._trace_writer.record_trace(
+                TraceRecord(
+                    trace_id=trace_id,
+                    status=status,
+                    end_ts=end_ts,
+                    metadata=meta,
+                )
+            )
+        except Exception:  # pylint: disable=broad-exception-caught
+            pass