by-framework 0.2.2.dev5__py3-none-any.whl → 0.2.2.dev7__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/common/__init__.py

@@ -28,7 +28,13 @@ from .config import (
 from .constants import RedisKeys
 from .emitter import (DataLayoutBuilder, DefaultSseLayoutBuilder, GatewayDataEmitter)
 from .logger import get_logger, logger, setup_logging
-from .redis_client import Redis, close_redis, get_redis, init_redis
+from .redis_client import (
+    Redis,
+    close_redis,
+    get_redis,
+    init_redis,
+    init_redis_from_url,
+)
 
 __all__ = [
     "RedisKeys",
@@ -37,6 +43,7 @@ __all__ = [
     "setup_logging",
     "get_redis",
     "init_redis",
+    "init_redis_from_url",
     "close_redis",
     "Redis",
     "DataLayoutBuilder",

by_framework/common/config.py

@@ -6,7 +6,7 @@ Provides typed configuration loaded from environment variables.
 
 import os
 from dataclasses import dataclass, field
-from typing import Optional
+from typing import Literal, Optional
 
 from by_framework.common.constants import RedisKeys
 
@@ -22,6 +22,10 @@ class RedisConfig:
     username: Optional[str] = None
     decode_responses: bool = True
     max_connections: Optional[int] = None
+    mode: Literal["standalone", "cluster"] = "standalone"
+    cluster_nodes: Optional[list[tuple[str, int]]] = None
+    socket_connect_timeout: int = 5
+    socket_timeout: int = 10
 
     @classmethod
     def from_env(cls) -> "RedisConfig":
@@ -29,6 +33,14 @@ class RedisConfig:
         password = os.environ.get("REDIS_PASSWORD", "")
         username = os.environ.get("REDIS_USERNAME") or None
         max_connections = os.environ.get("REDIS_MAX_CONNECTIONS")
+        mode = os.environ.get("REDIS_MODE", "standalone")
+        cluster_nodes_str = os.environ.get("REDIS_CLUSTER_NODES")
+        cluster_nodes = None
+        if cluster_nodes_str:
+            cluster_nodes = []
+            for node in cluster_nodes_str.split(","):
+                node_host, node_port = node.rsplit(":", 1)
+                cluster_nodes.append((node_host, int(node_port)))
         return cls(
             host=os.environ.get("REDIS_HOST", "localhost"),
             port=int(os.environ.get("REDIS_PORT", "6379")),
@@ -36,6 +48,8 @@ class RedisConfig:
             password=password,
             username=username,
             max_connections=int(max_connections) if max_connections else None,
+            mode=mode,
+            cluster_nodes=cluster_nodes,
         )