beadhub 0.1.0__py3-none-any.whl

beadhub/presence.py

@@ -0,0 +1,495 @@
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from typing import Dict, List, Optional
+from urllib.parse import quote
+
+from redis.asyncio import Redis
+
+from .roles import ROLE_MAX_LENGTH, is_valid_role, normalize_role
+
+DEFAULT_PRESENCE_TTL_SECONDS = 1800  # 30 minutes
+
+
+def _safe_key_component(value: str) -> str:
+    """URL-encode a value for safe use in Redis keys.
+
+    Prevents key collision attacks where values containing colons could
+    create ambiguous key boundaries. For example, without encoding:
+      project_id="abc", alias="xyz:def" -> "idx:alias:abc:xyz:def"
+      project_id="abc:xyz", alias="def" -> "idx:alias:abc:xyz:def" (COLLISION!)
+
+    With encoding:
+      project_id="abc", alias="xyz:def" -> "idx:alias:abc:xyz%3Adef"
+      project_id="abc:xyz", alias="def" -> "idx:alias:abc%3Axyz:def" (DISTINCT)
+    """
+    return quote(value, safe="")
+
+
+def _presence_key(workspace_id: str) -> str:
+    """Presence key: one agent per workspace."""
+    return f"presence:{workspace_id}"
+
+
+def _project_workspaces_index_key(project_id: str) -> str:
+    """Secondary index: workspace_ids by project_id."""
+    return f"idx:project_workspaces:{project_id}"
+
+
+def _repo_workspaces_index_key(repo_id: str) -> str:
+    """Secondary index: workspace_ids by repo_id."""
+    return f"idx:repo_workspaces:{repo_id}"
+
+
+def _branch_workspaces_index_key(repo_id: str, branch: str) -> str:
+    """Secondary index: workspace_ids by repo_id and branch."""
+    return f"idx:branch_workspaces:{repo_id}:{_safe_key_component(branch)}"
+
+
+def _project_slug_workspaces_index_key(project_slug: str) -> str:
+    """Secondary index: workspace_ids by project_slug."""
+    return f"idx:project_slug_workspaces:{_safe_key_component(project_slug)}"
+
+
+def _all_workspaces_index_key() -> str:
+    """Global index: all workspace_ids with active presence."""
+    return "idx:all_workspaces"
+
+
+def _alias_index_key(project_id: str, alias: str) -> str:
+    """Secondary index: workspace_id by (project_id, alias).
+
+    Enables O(1) alias collision checking instead of SCAN.
+    Key maps to a single workspace_id (aliases are unique per project).
+    """
+    return f"idx:alias:{project_id}:{_safe_key_component(alias)}"
+
+
+async def update_agent_presence(
+    redis: Redis,
+    workspace_id: str,
+    alias: str,
+    program: Optional[str],
+    model: Optional[str],
+    human_name: Optional[str] = None,
+    project_id: Optional[str] = None,
+    project_slug: Optional[str] = None,
+    repo_id: Optional[str] = None,
+    member_email: str = "",
+    status: str = "active",
+    current_branch: Optional[str] = None,
+    role: Optional[str] = None,
+    ttl_seconds: int = DEFAULT_PRESENCE_TTL_SECONDS,
+) -> str:
+    """
+    Upsert agent presence in Redis and return the ISO timestamp used.
+
+    Args:
+        workspace_id: UUID identifying the workspace.
+        alias: Human-friendly workspace identifier for addressing.
+        human_name: Name of the human who owns this workspace.
+        project_id: UUID of the project (for secondary index).
+        project_slug: Human-readable project slug.
+        repo_id: UUID of the repo (for secondary index).
+        current_branch: Optional branch name.
+        role: Brief description of workspace purpose (max 50 chars).
+        ttl_seconds: How long until presence expires if not refreshed. Default 5 minutes.
+    """
+    key = _presence_key(workspace_id)
+    now = datetime.now(timezone.utc).isoformat()
+
+    fields = {
+        "workspace_id": workspace_id,
+        "alias": alias,
+        "human_name": human_name or "",
+        "project_id": project_id or "",
+        "project_slug": project_slug or "",
+        "repo_id": repo_id or "",
+        "member_email": member_email,
+        "program": program or "",
+        "model": model or "",
+        "status": status,
+        "current_branch": current_branch or "",
+        "last_seen": now,
+    }
+    if role is not None and len(role) <= ROLE_MAX_LENGTH and is_valid_role(role):
+        fields["role"] = normalize_role(role)
+
+    await redis.hset(key, mapping=fields)
+    await redis.expire(key, ttl_seconds)
+
+    # Update secondary indexes
+    # Index TTL is 2x presence TTL to ensure index entries outlive presence keys,
+    # allowing lazy cleanup to detect stale entries via EXISTS checks.
+    # Note: workspace → project is immutable (see architecture docs), so project_id
+    # and project_slug don't change for a given workspace. Branch and repo indexes
+    # may have transient staleness (up to TTL*2) when workspaces switch branches.
+
+    # Global all_workspaces index (always maintained)
+    all_idx_key = _all_workspaces_index_key()
+    await redis.sadd(all_idx_key, workspace_id)
+    await redis.expire(all_idx_key, ttl_seconds * 2)
+
+    if project_id:
+        idx_key = _project_workspaces_index_key(project_id)
+        await redis.sadd(idx_key, workspace_id)
+        await redis.expire(idx_key, ttl_seconds * 2)
+
+        # Alias index for O(1) collision checking (1:1 mapping, not a set)
+        alias_idx_key = _alias_index_key(project_id, alias)
+        await redis.set(alias_idx_key, workspace_id, ex=ttl_seconds * 2)
+
+    if project_slug:
+        idx_key = _project_slug_workspaces_index_key(project_slug)
+        await redis.sadd(idx_key, workspace_id)
+        await redis.expire(idx_key, ttl_seconds * 2)
+
+    if repo_id:
+        idx_key = _repo_workspaces_index_key(repo_id)
+        await redis.sadd(idx_key, workspace_id)
+        await redis.expire(idx_key, ttl_seconds * 2)
+
+        if current_branch:
+            idx_key = _branch_workspaces_index_key(repo_id, current_branch)
+            await redis.sadd(idx_key, workspace_id)
+            await redis.expire(idx_key, ttl_seconds * 2)
+
+    return now
+
+
+async def get_agent_presence(
+    redis: Redis,
+    workspace_id: str,
+) -> Optional[Dict[str, str]]:
+    """
+    Fetch an agent's presence record from Redis.
+
+    One agent per workspace architecture: workspace_id is the only key.
+
+    Args:
+        workspace_id: UUID identifying the workspace.
+    """
+    key = _presence_key(workspace_id)
+    data: Dict[str, str] = await redis.hgetall(key)
+    if not data:
+        return None
+    return data
+
+
+async def list_agent_presences(
+    redis: Redis,
+    workspace_id: Optional[str] = None,
+) -> List[Dict[str, str]]:
+    """
+    List agent presence records.
+
+    Uses the all_workspaces secondary index for O(M) lookup where M is the
+    number of active workspaces, instead of O(N) SCAN over all Redis keys.
+
+    Args:
+        workspace_id: If provided, return the presence for this workspace only.
+                      Otherwise, list all presences.
+    """
+    if workspace_id:
+        # One agent per workspace - direct lookup
+        key = _presence_key(workspace_id)
+        data: Dict[str, str] = await redis.hgetall(key)
+        return [data] if data else []
+
+    # List all presences using secondary index (avoids SCAN)
+    idx_key = _all_workspaces_index_key()
+    workspace_ids = await _filter_valid_workspace_ids(redis, idx_key)
+
+    if not workspace_ids:
+        return []
+
+    # Batch fetch all presence hashes with pipeline (N round-trips → 1)
+    pipe = redis.pipeline()
+    for ws_id in workspace_ids:
+        pipe.hgetall(_presence_key(ws_id))
+    presence_data = await pipe.execute()
+
+    results: List[Dict[str, str]] = []
+    for data in presence_data:
+        if data:
+            results.append(data)
+
+    return results
+
+
+async def list_agent_presences_by_workspace_ids(
+    redis: Redis,
+    workspace_ids: List[str],
+) -> List[Dict[str, str]]:
+    """
+    Fetch presence records for specific workspace IDs.
+
+    This avoids scanning global indexes when callers already know which
+    workspaces they need to enrich with presence.
+    """
+    if not workspace_ids:
+        return []
+
+    pipe = redis.pipeline()
+    for ws_id in workspace_ids:
+        pipe.hgetall(_presence_key(ws_id))
+    presence_data = await pipe.execute()
+
+    results: List[Dict[str, str]] = []
+    for data in presence_data:
+        if data:
+            results.append(data)
+
+    return results
+
+
+async def _filter_valid_workspace_ids(
+    redis: Redis,
+    idx_key: str,
+) -> List[str]:
+    """
+    Filter workspace_ids from a secondary index, removing stale entries.
+
+    Uses Redis pipeline to batch EXISTS checks (N+1 round-trips → 2:
+    one SMEMBERS + one pipeline for all EXISTS checks).
+
+    Stale entries (presence expired but index entry remains) are lazily
+    removed from the index. There's a theoretical race where a workspace
+    could be removed from the index just as it refreshes presence, but
+    the entry gets re-added on next presence update.
+
+    Args:
+        redis: Redis client.
+        idx_key: Key for the secondary index (Set).
+
+    Returns:
+        List of workspace_ids with active presence.
+    """
+    members = await redis.smembers(idx_key)
+    if not members:
+        return []
+
+    # Normalize to strings
+    workspace_ids = [
+        ws_id.decode("utf-8") if isinstance(ws_id, bytes) else ws_id for ws_id in members
+    ]
+
+    # Batch EXISTS checks with pipeline (N round-trips → 1)
+    pipe = redis.pipeline()
+    for ws_id in workspace_ids:
+        pipe.exists(_presence_key(ws_id))
+    exists_results = await pipe.execute()
+
+    # Separate valid and stale workspace_ids
+    valid_workspace_ids: List[str] = []
+    stale_workspace_ids: List[str] = []
+    for ws_id, exists in zip(workspace_ids, exists_results):
+        if exists:
+            valid_workspace_ids.append(ws_id)
+        else:
+            stale_workspace_ids.append(ws_id)
+
+    # Lazy cleanup: remove stale entries from index
+    if stale_workspace_ids:
+        cleanup_pipe = redis.pipeline()
+        for ws_id in stale_workspace_ids:
+            cleanup_pipe.srem(idx_key, ws_id)
+        await cleanup_pipe.execute()
+
+    return valid_workspace_ids
+
+
+async def get_workspace_ids_by_project_slug(
+    redis: Redis,
+    project_slug: str,
+) -> List[str]:
+    """
+    Get all workspace_ids that belong to a project by slug.
+
+    Uses secondary index for O(M) lookup where M is workspaces in project.
+    Stale entries (presence expired but index entry remains) are filtered
+    out and lazily removed from the index.
+
+    Args:
+        project_slug: Human-readable project slug.
+
+    Returns:
+        List of workspace_ids with matching project_slug.
+    """
+    idx_key = _project_slug_workspaces_index_key(project_slug)
+    return await _filter_valid_workspace_ids(redis, idx_key)
+
+
+async def get_workspace_ids_by_project_id(
+    redis: Redis,
+    project_id: str,
+) -> List[str]:
+    """
+    Get all workspace_ids that belong to a project by ID.
+
+    Uses secondary index for O(1) lookup. Stale entries (presence expired but
+    index entry remains) are filtered out and lazily removed from the index.
+
+    Args:
+        project_id: Project UUID.
+
+    Returns:
+        List of workspace_ids in the project.
+    """
+    idx_key = _project_workspaces_index_key(project_id)
+    return await _filter_valid_workspace_ids(redis, idx_key)
+
+
+async def get_workspace_ids_by_repo_id(
+    redis: Redis,
+    repo_id: str,
+) -> List[str]:
+    """
+    Get all workspace_ids that belong to a repo by ID.
+
+    Uses secondary index for O(1) lookup. Stale entries (presence expired but
+    index entry remains) are filtered out and lazily removed from the index.
+
+    Args:
+        repo_id: Repo UUID.
+
+    Returns:
+        List of workspace_ids in the repo.
+    """
+    idx_key = _repo_workspaces_index_key(repo_id)
+    return await _filter_valid_workspace_ids(redis, idx_key)
+
+
+async def get_workspace_ids_by_branch(
+    redis: Redis,
+    repo_id: str,
+    branch: str,
+) -> List[str]:
+    """
+    Get all workspace_ids working on a specific branch of a repo.
+
+    Uses secondary index for O(1) lookup. Stale entries (presence expired but
+    index entry remains) are filtered out and lazily removed from the index.
+
+    Args:
+        repo_id: Repo UUID.
+        branch: Branch name.
+
+    Returns:
+        List of workspace_ids on the branch.
+    """
+    idx_key = _branch_workspaces_index_key(repo_id, branch)
+    return await _filter_valid_workspace_ids(redis, idx_key)
+
+
+async def get_all_workspace_ids(
+    redis: Redis,
+) -> List[str]:
+    """
+    Get all workspace_ids with active presence.
+
+    Uses the global all_workspaces index. Stale entries (presence expired but
+    index entry remains) are filtered out and lazily removed from the index.
+
+    Returns:
+        List of all active workspace_ids.
+    """
+    idx_key = _all_workspaces_index_key()
+    return await _filter_valid_workspace_ids(redis, idx_key)
+
+
+async def get_workspace_id_by_alias(
+    redis: Redis,
+    project_id: str,
+    alias: str,
+) -> Optional[str]:
+    """
+    Get the workspace_id using a specific alias within a project.
+
+    Uses the alias secondary index for O(1) lookup. Returns the workspace_id
+    if the alias is in use and the workspace has active presence, None otherwise.
+
+    Note: This is for presence-based collision checking only. The database
+    (workspaces table) is the authoritative source for alias ownership.
+
+    Args:
+        project_id: Project UUID.
+        alias: The alias to look up.
+
+    Returns:
+        workspace_id if alias is in use with active presence, None otherwise.
+    """
+    idx_key = _alias_index_key(project_id, alias)
+    workspace_id = await redis.get(idx_key)
+
+    if not workspace_id:
+        return None
+
+    # Normalize to string (may be bytes depending on Redis client config)
+    ws_id = workspace_id.decode("utf-8") if isinstance(workspace_id, bytes) else workspace_id
+
+    # Verify presence is still active (index may outlive presence due to TTL difference)
+    presence_key = _presence_key(ws_id)
+    if not await redis.exists(presence_key):
+        # Stale index entry - lazy cleanup
+        await redis.delete(idx_key)
+        return None
+
+    return ws_id
+
+
+async def get_workspace_project_slug(
+    redis: Redis,
+    workspace_id: str,
+) -> Optional[str]:
+    """Get the project_slug for a workspace from its presence data.
+
+    Args:
+        redis: Redis client.
+        workspace_id: Workspace UUID.
+
+    Returns:
+        project_slug if available, None otherwise.
+    """
+    data = await redis.hget(_presence_key(workspace_id), "project_slug")
+    if not data:
+        return None
+    slug = data.decode("utf-8") if isinstance(data, bytes) else data
+    return slug if slug else None
+
+
+async def clear_workspace_presence(
+    redis: Redis,
+    workspace_ids: List[str],
+) -> int:
+    """
+    Clear presence for a list of workspaces.
+
+    Deletes presence keys and removes from all secondary indexes.
+    Used when soft-deleting repos or projects.
+
+    Args:
+        workspace_ids: List of workspace UUIDs to clear.
+
+    Returns:
+        Number of presence keys deleted.
+    """
+    if not workspace_ids:
+        return 0
+
+    # Delete presence keys
+    pipe = redis.pipeline()
+    for ws_id in workspace_ids:
+        pipe.delete(_presence_key(ws_id))
+    results = await pipe.execute()
+    deleted_count = sum(1 for r in results if r)
+
+    # Remove from all secondary indexes (lazy cleanup handles misses)
+    # We remove from all possible indexes to ensure cleanup
+    all_idx_key = _all_workspaces_index_key()
+    pipe = redis.pipeline()
+    for ws_id in workspace_ids:
+        pipe.srem(all_idx_key, ws_id)
+    await pipe.execute()
+
+    return deleted_count

beadhub/rate_limit.py

@@ -0,0 +1,152 @@
+"""Rate limiting for BeadHub endpoints.
+
+Provides Redis-based rate limiting with fixed window counters.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import Optional
+
+from fastapi import HTTPException, Request
+from redis.asyncio import Redis
+from redis.exceptions import RedisError
+
+logger = logging.getLogger(__name__)
+
+# Rate limit configuration for /v1/init endpoint
+# Can be overridden via environment for testing
+INIT_RATE_LIMIT = int(os.getenv("BEADHUB_INIT_RATE_LIMIT", "10"))  # requests per window
+INIT_RATE_WINDOW = int(os.getenv("BEADHUB_INIT_RATE_WINDOW", "60"))  # seconds
+
+# Lua script for atomic increment with expiration
+# This prevents race condition between INCR and EXPIRE
+RATE_LIMIT_SCRIPT = """
+local current = redis.call('INCR', KEYS[1])
+if current == 1 then
+    redis.call('EXPIRE', KEYS[1], ARGV[1])
+end
+return current
+"""
+
+
+def get_client_ip(request: Request) -> str:
+    """Extract client IP address from request.
+
+    Uses request.client.host which is set by the ASGI server based on
+    the actual TCP connection. This is secure against X-Forwarded-For
+    spoofing attacks.
+
+    For deployments behind a reverse proxy, configure the proxy to set
+    the client IP correctly, or use FastAPI's TrustedHostMiddleware.
+    Do NOT trust X-Forwarded-For without proper proxy configuration.
+
+    Args:
+        request: FastAPI request object
+
+    Returns:
+        Client IP address string
+    """
+    # request.client can be None in some test scenarios
+    if request.client is None:
+        return "unknown"
+    return request.client.host
+
+
+async def check_rate_limit(
+    request: Request,
+    redis: Redis,
+    key_prefix: str,
+    limit: int,
+    window_seconds: int,
+) -> Optional[int]:
+    """Check if request is within rate limit.
+
+    Uses a Lua script for atomic increment with expiration. The window
+    is fixed from the first request (not sliding).
+
+    Note: Fixed window allows burst at window boundaries (up to 2x limit
+    in a short period). This is acceptable for /v1/init.
+
+    Args:
+        request: FastAPI request object
+        redis: Async Redis client
+        key_prefix: Prefix for the rate limit key (e.g., "ratelimit:init")
+        limit: Maximum requests per window
+        window_seconds: Window duration in seconds
+
+    Returns:
+        None if within limit, otherwise seconds until limit resets
+
+    Raises:
+        RedisError: If Redis is unavailable (caller should handle)
+    """
+    client_ip = get_client_ip(request)
+    key = f"{key_prefix}:{client_ip}"
+
+    # Atomic increment with expiration using Lua script
+    # This prevents race condition where INCR succeeds but EXPIRE fails
+    current = await redis.eval(RATE_LIMIT_SCRIPT, 1, key, window_seconds)
+
+    if current > limit:
+        # Over limit - get TTL for Retry-After header
+        ttl = await redis.ttl(key)
+        if ttl < 0:
+            # Key has no expiry (shouldn't happen with Lua script)
+            # Delete and allow retry on next request
+            await redis.delete(key)
+            ttl = window_seconds
+            logger.error(
+                "Rate limit key without TTL detected and deleted: ip=%s key=%s",
+                client_ip,
+                key,
+            )
+        logger.warning(
+            "Rate limit exceeded: ip=%s key=%s count=%d limit=%d",
+            client_ip,
+            key,
+            current,
+            limit,
+        )
+        return ttl
+
+    return None
+
+
+async def enforce_init_rate_limit(request: Request, redis: Redis) -> None:
+    """Enforce rate limit for /v1/init endpoint.
+
+    Fails closed: If Redis is unavailable, returns 503 Service Unavailable
+    rather than allowing unlimited requests.
+
+    Args:
+        request: FastAPI request object
+        redis: Async Redis client
+
+    Raises:
+        HTTPException: 429 Too Many Requests if rate limit exceeded
+        HTTPException: 503 Service Unavailable if Redis is unavailable
+    """
+    try:
+        retry_after = await check_rate_limit(
+            request=request,
+            redis=redis,
+            key_prefix="ratelimit:init",
+            limit=INIT_RATE_LIMIT,
+            window_seconds=INIT_RATE_WINDOW,
+        )
+    except RedisError as e:
+        # Fail closed: deny request if we can't check rate limit
+        logger.error("Redis error during rate limit check: %s", e)
+        raise HTTPException(
+            status_code=503,
+            detail="Service temporarily unavailable. Please retry.",
+        )
+
+    if retry_after is not None:
+        raise HTTPException(
+            status_code=429,
+            detail="Rate limit exceeded. Too many initialization requests.",
+            headers={"Retry-After": str(retry_after)},
+        )

beadhub/redis_client.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+from fastapi import Request
+from redis.asyncio import Redis
+
+
+def get_redis(request: Request) -> Redis:
+    """
+    FastAPI dependency that returns the shared async Redis client from app state.
+    """
+    return request.app.state.redis

beadhub/roles.py

@@ -0,0 +1,35 @@
+"""Role validation and normalization helpers."""
+
+from __future__ import annotations
+
+import re
+
+ROLE_MAX_LENGTH = 50
+ROLE_MAX_WORDS = 2
+ROLE_WORD_PATTERN = re.compile(r"^[a-zA-Z0-9][a-zA-Z0-9_-]*$")
+ROLE_ERROR_MESSAGE = (
+    "Invalid role: use 1-2 words (letters/numbers) with hyphens/underscores allowed; max 50 chars"
+)
+
+
+def normalize_role(role: str) -> str:
+    """Normalize role string by trimming, collapsing spaces, and lowercasing."""
+    return " ".join(role.strip().split()).lower()
+
+
+def is_valid_role(role: str) -> bool:
+    """Check if role is 1-2 words with allowed characters."""
+    if not role or not isinstance(role, str):
+        return False
+    normalized = normalize_role(role)
+    if not normalized or len(normalized) > ROLE_MAX_LENGTH:
+        return False
+    words = normalized.split(" ")
+    if len(words) > ROLE_MAX_WORDS:
+        return False
+    return all(ROLE_WORD_PATTERN.match(word) for word in words)
+
+
+def role_to_alias_prefix(role: str) -> str:
+    """Convert role to an alias-friendly prefix."""
+    return normalize_role(role).replace(" ", "-")

beadhub/routes/__init__.py

@@ -0,0 +1 @@
+