polyrob 0.5.0__py3-none-any.whl

api/jwt_middleware.py

@@ -0,0 +1,83 @@
+"""JWT authentication middleware."""
+
+import logging
+import jwt
+from typing import Callable
+
+from fastapi import Request
+from fastapi.responses import JSONResponse
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.responses import Response
+
+from api.auth_constants import is_admin
+
+logger = logging.getLogger(__name__)
+
+
+class JWTAuthMiddleware(BaseHTTPMiddleware):
+    """Decode JWT and populate request.state."""
+
+    PUBLIC_PATHS = {
+        "/", "/health", "/docs", "/openapi.json", "/redoc",
+        "/api/auth/nonce", "/api/auth/verify",
+        "/api/payments/pricing", "/api/x402/pricing"
+    }
+
+    def __init__(self, app, jwt_secret: str):
+        super().__init__(app)
+        self.jwt_secret = jwt_secret
+
+    async def dispatch(self, request: Request, call_next: Callable) -> Response:
+        """Process request with JWT authentication."""
+
+        # Skip public endpoints
+        if (request.url.path in self.PUBLIC_PATHS or
+            request.url.path.startswith("/static")):
+            return await call_next(request)
+
+        # Get JWT
+        auth_header = request.headers.get("Authorization", "")
+        if not auth_header.startswith("Bearer "):
+            return await call_next(request)  # x402 flow
+
+        token = auth_header[7:]
+
+        try:
+            # Decode JWT
+            payload = jwt.decode(token, self.jwt_secret, algorithms=["HS256"])
+
+            # Populate request.state via the canonical C4 contract.
+            from api.auth_state import set_auth_state
+            set_auth_state(
+                request.state,
+                user_id=payload.get("user_id"),
+                tier=payload.get("tier", "free"),
+                role=payload.get("role", "user"),
+                payment_method=None,
+                authenticated=True,
+            )
+            # NOTE: "sub" contains wallet_address (JWT standard)
+            request.state.wallet_address = payload.get("sub")
+            request.state.chain = payload.get("chain", "ethereum")
+
+            # Admin flag (simplified: only 'admin' role or admin wallet)
+            request.state.is_admin = is_admin(
+                role=request.state.role,
+                wallet_address=request.state.wallet_address
+            )
+
+            wallet_display = payload.get('sub', '')[:10] if payload.get('sub') else 'unknown'
+            logger.debug(f"JWT valid for {wallet_display}...")
+
+        except jwt.ExpiredSignatureError:
+            return JSONResponse(
+                status_code=401,
+                content={"error": "Token expired", "code": "TOKEN_EXPIRED"}
+            )
+        except jwt.InvalidTokenError:
+            return JSONResponse(
+                status_code=401,
+                content={"error": "Invalid token", "code": "INVALID_TOKEN"}
+            )
+
+        return await call_next(request)

api/kb/__init__.py

@@ -0,0 +1 @@
+# api/kb package

api/kb/endpoints.py

@@ -0,0 +1,314 @@
+"""KB (knowledge-base) HTTP API — gated, multi-tenant.
+
+Mount in ``api/app.py`` behind ``KB_API_ENABLED`` (default OFF), mirroring the
+OpenAI-compat block:
+
+    from api.kb.endpoints import router as kb_router, kb_api_enabled
+    if kb_api_enabled():
+        app.include_router(kb_router)
+
+Auth
+----
+``user_id`` is ALWAYS resolved from the authenticated request state via the shared
+``get_user_id`` dependency (same seam the task HTTP API and other routers use).
+A caller-supplied ``user_id`` in the body is never trusted.
+
+Provider registration
+---------------------
+The SQLite memory backend (``MEMORY_BACKEND=sqlite``, default) is registered lazily
+on the first ingest/search request via ``_ensure_backend``.  On the full server path
+``build_server_bot`` does NOT call ``maybe_register_memory_backend`` — that call lives
+in ``construction.py`` (per-session agent build).  Calling it here on first request is
+idempotent: the factory's active-check short-circuits if a provider is already live
+(a sibling session already registered it).
+
+Path confinement (path-body variant)
+-------------------------------------
+``POST /api/kb/ingest`` with a path in the body: the path must be **relative** and
+must not escape the authenticated tenant's session workspace.  Absolute paths and
+``..``-escapes are rejected with HTTP 400 *before* the engine is called (defence in
+depth — the engine also confines, but we fail-fast here).
+
+File-upload variant
+-------------------
+``POST /api/kb/ingest/upload`` accepts a multipart file upload.  The bytes are
+written to a unique temp file UNDER the authenticated tenant's session workspace
+(``<workspace>/.kb_uploads/<unique>``), so the shared ``kb_ingest`` engine's
+confinement check passes naturally.  The upload is ingested through the SAME engine
+as the path-body variant (no duplicated chunking/hashing), with ``source_name`` set
+to the stripped original filename so a re-upload of the same file dedups by its
+logical name rather than the volatile temp path.  The temp file is removed in a
+``finally`` block even on error.
+"""
+from __future__ import annotations
+
+import logging
+import os
+import tempfile
+from pathlib import Path
+from typing import Optional
+
+from fastapi import APIRouter, Depends, HTTPException, UploadFile, File
+from fastapi.responses import JSONResponse
+from pydantic import BaseModel, Field
+
+from api.dependencies import get_user_id
+from core.path_safety import is_within_root
+
+logger = logging.getLogger(__name__)
+
+router = APIRouter(prefix="/api/kb", tags=["knowledge-base"])
+
+# ---------------------------------------------------------------------------
+# Gate helper
+# ---------------------------------------------------------------------------
+
+
+def kb_api_enabled() -> bool:
+    """Whether the KB HTTP API is mounted (default OFF).
+
+    Uses the project's shared ``_bool_env`` parser (falsey-set semantics:
+    none/off/false/0/no/'' = off), so the flag behaves consistently with the
+    rest of POLYROB's env flags.
+    """
+    from agents.task.constants import _bool_env
+    return _bool_env("KB_API_ENABLED", False)
+
+
+# ---------------------------------------------------------------------------
+# Lazy backend registration (idempotent)
+# ---------------------------------------------------------------------------
+
+def _ensure_backend() -> None:
+    """Register the memory backend if it isn't already active.
+
+    Called on the first ingest/search request; idempotent — the factory
+    short-circuits if a provider is already registered.
+    """
+    try:
+        from modules.memory.backend_factory import maybe_register_memory_backend
+        from core.config import BotConfig
+        try:
+            cfg = BotConfig()
+            data_dir = getattr(cfg, "data_dir", "data") or "data"
+        except Exception:
+            data_dir = "data"
+        maybe_register_memory_backend(data_dir=data_dir)
+    except Exception as e:
+        logger.debug("KB API: memory backend registration skipped: %s", e)
+
+
+# ---------------------------------------------------------------------------
+# Request models
+# ---------------------------------------------------------------------------
+
+class KbIngestBody(BaseModel):
+    """Body for ``POST /api/kb/ingest`` when providing a workspace-relative path."""
+    path: str = Field(..., description="Path relative to the tenant's session workspace.")
+    session_id: str = Field(..., description="Session ID whose workspace root to use.")
+    collection: str = Field("default", description="KB collection name.")
+    recursive: bool = Field(True, description="Recurse into subdirectories.")
+    globs: Optional[list[str]] = Field(None, description="Optional glob patterns.")
+
+
+class KbSearchBody(BaseModel):
+    """Body for ``POST /api/kb/search``."""
+    query: str = Field(..., description="Search query.")
+    collection: str = Field("default", description="KB collection to search.")
+    limit: int = Field(8, ge=1, le=50, description="Maximum number of results.")
+
+
+# ---------------------------------------------------------------------------
+# Path-safety helper (endpoint-level early rejection)
+# ---------------------------------------------------------------------------
+
+def _resolve_and_guard(
+    path: str,
+    session_id: str,
+    user_id: str,
+) -> Path:
+    """Resolve *path* against the tenant workspace and return the resolved Path.
+
+    Raises ``HTTPException(400)`` for:
+    - Absolute paths.
+    - Paths that escape the workspace root after resolution.
+    """
+    if os.path.isabs(path):
+        raise HTTPException(
+            status_code=400,
+            detail="Absolute paths are not allowed. Provide a path relative to your session workspace.",
+        )
+
+    from agents.task.path import pm
+    try:
+        workspace_root = pm().get_workspace_dir(session_id, user_id)
+    except Exception as exc:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Could not resolve workspace for session '{session_id}': {exc}",
+        )
+
+    resolved = (workspace_root / path).resolve()
+
+    if not is_within_root(str(resolved), str(workspace_root.resolve())):
+        raise HTTPException(
+            status_code=400,
+            detail="Path escapes the session workspace. Only paths within the workspace are allowed.",
+        )
+
+    return resolved
+
+
+# ---------------------------------------------------------------------------
+# POST /api/kb/ingest  (path body variant)
+# ---------------------------------------------------------------------------
+
+@router.post("/ingest")
+async def kb_ingest_path(
+    body: KbIngestBody,
+    user_id: str = Depends(get_user_id),
+) -> JSONResponse:
+    """Ingest a file or directory from the tenant's session workspace into the KB.
+
+    The ``path`` in the body must be **relative** to the authenticated tenant's
+    session workspace.  Absolute paths and ``..``-escaping paths are rejected.
+    ``user_id`` is derived from the auth dependency — the body never supplies it.
+    """
+    _ensure_backend()
+
+    # Early path safety — reject absolute / escaping paths before calling the engine.
+    _resolve_and_guard(body.path, body.session_id, user_id)
+
+    try:
+        from tools.knowledge_ingest import kb_ingest
+        result = await kb_ingest(
+            body.path,
+            collection=body.collection,
+            recursive=body.recursive,
+            globs=body.globs,
+            user_id=user_id,
+            session_id=body.session_id,
+        )
+    except Exception as exc:
+        logger.warning("KB ingest failed for user=%s: %s", user_id, exc)
+        raise HTTPException(status_code=500, detail=f"KB ingest error: {exc}")
+
+    if "error" in result:
+        raise HTTPException(status_code=400, detail=result["error"])
+
+    return JSONResponse(content=result)
+
+
+# ---------------------------------------------------------------------------
+# POST /api/kb/ingest/upload  (file-upload variant)
+# ---------------------------------------------------------------------------
+
+@router.post("/ingest/upload")
+async def kb_ingest_upload(
+    file: UploadFile = File(...),
+    collection: str = "default",
+    session_id: str = "_kb_upload_",
+    user_id: str = Depends(get_user_id),
+) -> JSONResponse:
+    """Ingest an uploaded file into the tenant's KB.
+
+    The bytes are written to a unique temp file UNDER the authenticated tenant's
+    session workspace, then ingested through the SAME ``kb_ingest`` engine as the
+    path-body variant — so size guards, MIME detection, dedup and confinement all
+    apply for free.  ``source_name`` is the stripped original filename, so a
+    re-upload of the same file dedups by logical name (not the volatile temp path).
+    ``user_id`` comes from the auth dependency, NEVER from the request body.
+    """
+    _ensure_backend()
+
+    # Resolve the tenant's confined workspace root and write the upload UNDER it,
+    # so the engine's is_within_root confinement check passes naturally.
+    from agents.task.path import pm
+    try:
+        workspace_root = pm().get_workspace_dir(session_id, user_id)
+    except Exception as exc:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Could not resolve workspace for session '{session_id}': {exc}",
+        )
+
+    upload_dir = Path(workspace_root) / ".kb_uploads"
+    upload_dir.mkdir(parents=True, exist_ok=True)
+
+    # Strip path components from the client filename (no traversal); keep it as the
+    # STABLE source identity so a re-upload dedups by its original name.
+    source_name = Path(file.filename or "upload.bin").name
+    suffix = Path(source_name).suffix or ".bin"
+
+    tmp_file: Optional[Path] = None
+    try:
+        fd, tmp_path_str = tempfile.mkstemp(suffix=suffix, dir=str(upload_dir))
+        tmp_file = Path(tmp_path_str)
+        try:
+            with os.fdopen(fd, "wb") as fh:
+                content = await file.read()
+                fh.write(content)
+        except Exception:
+            os.close(fd)
+            raise
+
+        # Engine path is relative to the workspace root (it resolves
+        # allowed_root / path); pass the temp file's path relative to the root.
+        rel_path = str(tmp_file.relative_to(Path(workspace_root)))
+
+        from tools.knowledge_ingest import kb_ingest
+        result = await kb_ingest(
+            rel_path,
+            collection=collection,
+            recursive=False,
+            user_id=user_id,
+            session_id=session_id,
+            source_name=source_name,
+        )
+    except HTTPException:
+        raise
+    except Exception as exc:
+        logger.warning("KB upload ingest failed for user=%s: %s", user_id, exc)
+        raise HTTPException(status_code=500, detail=f"KB ingest error: {exc}")
+    finally:
+        if tmp_file is not None:
+            try:
+                tmp_file.unlink(missing_ok=True)
+            except OSError:
+                pass
+
+    if "error" in result:
+        raise HTTPException(status_code=400, detail=result["error"])
+
+    return JSONResponse(content=result)
+
+
+# ---------------------------------------------------------------------------
+# POST /api/kb/search
+# ---------------------------------------------------------------------------
+
+@router.post("/search")
+async def kb_search_endpoint(
+    body: KbSearchBody,
+    user_id: str = Depends(get_user_id),
+) -> JSONResponse:
+    """Search the tenant's KB.
+
+    Results are scoped to the authenticated ``user_id``; a caller cannot search
+    another tenant's data.
+    """
+    _ensure_backend()
+
+    try:
+        from modules.memory.registry import kb_search as _kb_search
+        result = await _kb_search(
+            body.query,
+            user_id=user_id,
+            collection=body.collection,
+            limit=body.limit,
+        )
+    except Exception as exc:
+        logger.warning("KB search failed for user=%s: %s", user_id, exc)
+        raise HTTPException(status_code=500, detail=f"KB search error: {exc}")
+
+    return JSONResponse(content={"results": result or ""})

api/mcp_models.py

@@ -0,0 +1,279 @@
+"""
+Pydantic models for MCP server management API.
+"""
+
+from typing import Optional, List, Dict, Any
+from pydantic import BaseModel, Field, validator
+import re
+
+
+class AddServerRequest(BaseModel):
+    """Request to add a new MCP server."""
+
+    server_name: str = Field(
+        ...,
+        min_length=3,
+        max_length=50,
+        description="Unique server name (lowercase, alphanumeric, hyphens only)"
+    )
+    server_url: str = Field(
+        ...,
+        description="Server endpoint URL (must be HTTPS)"
+    )
+    server_type: str = Field(
+        default="sse",
+        description="Server type: 'sse' or 'http' (STDIO not allowed)"
+    )
+    auth_method: str = Field(
+        default="api_key",
+        description="Authentication method: 'api_key', 'bearer', or 'none'"
+    )
+    api_key: Optional[str] = Field(
+        default=None,
+        description="API key for authentication (will be encrypted)"
+    )
+    headers: Optional[Dict[str, str]] = Field(
+        default=None,
+        description="Custom headers (will be encrypted)"
+    )
+    display_name: Optional[str] = Field(
+        default=None,
+        max_length=100,
+        description="Friendly display name"
+    )
+    enabled: bool = Field(
+        default=True,
+        description="Whether server is enabled"
+    )
+    timeout: int = Field(
+        default=30,
+        ge=5,
+        le=300,
+        description="Connection timeout in seconds"
+    )
+    retry_attempts: int = Field(
+        default=3,
+        ge=0,
+        le=10,
+        description="Number of retry attempts"
+    )
+    retry_delay: int = Field(
+        default=5,
+        ge=1,
+        le=60,
+        description="Delay between retry attempts in seconds"
+    )
+    auto_reconnect: bool = Field(
+        default=True,
+        description="Whether to automatically reconnect on failure"
+    )
+    max_concurrent_requests: int = Field(
+        default=5,
+        ge=1,
+        le=50,
+        description="Maximum concurrent requests to this server"
+    )
+    verify_connection: bool = Field(
+        default=False,
+        description="If true, verify MCP protocol before saving (tests connection and discovers tools)"
+    )
+    message_endpoint: Optional[str] = Field(
+        default=None,
+        description="For SSE servers: explicit POST endpoint for sending messages (if different from main URL)"
+    )
+
+    @validator('server_name')
+    def validate_server_name(cls, v):
+        """Validate server name format."""
+        if not re.match(r'^[a-z0-9][a-z0-9-]*[a-z0-9]$', v) and len(v) > 2:
+            raise ValueError('Server name must be lowercase alphanumeric with hyphens')
+        if len(v) <= 2 and not re.match(r'^[a-z0-9]+$', v):
+            raise ValueError('Server name must be lowercase alphanumeric')
+        return v
+
+    @validator('server_type')
+    def validate_server_type(cls, v):
+        """Validate server type."""
+        if v not in ('sse', 'http', 'streamable_http'):
+            raise ValueError('Server type must be "sse", "http", or "streamable_http" (STDIO not allowed)')
+        return v
+
+    @validator('auth_method')
+    def validate_auth_method(cls, v):
+        """Validate auth method."""
+        if v not in ('api_key', 'bearer', 'none'):
+            raise ValueError('Auth method must be "api_key", "bearer", or "none"')
+        return v
+
+
+class UpdateServerRequest(BaseModel):
+    """Request to update an MCP server."""
+
+    server_url: Optional[str] = Field(
+        default=None,
+        description="Server endpoint URL (must be HTTPS)"
+    )
+    auth_method: Optional[str] = Field(
+        default=None,
+        description="Authentication method"
+    )
+    api_key: Optional[str] = Field(
+        default=None,
+        description="API key for authentication"
+    )
+    headers: Optional[Dict[str, str]] = Field(
+        default=None,
+        description="Custom headers"
+    )
+    display_name: Optional[str] = Field(
+        default=None,
+        max_length=100,
+        description="Friendly display name"
+    )
+    enabled: Optional[bool] = Field(
+        default=None,
+        description="Whether server is enabled"
+    )
+    timeout: Optional[int] = Field(
+        default=None,
+        ge=5,
+        le=300,
+        description="Connection timeout"
+    )
+    retry_attempts: Optional[int] = Field(
+        default=None,
+        ge=0,
+        le=10,
+        description="Retry attempts"
+    )
+    retry_delay: Optional[int] = Field(
+        default=None,
+        ge=1,
+        le=60,
+        description="Delay between retry attempts in seconds"
+    )
+    auto_reconnect: Optional[bool] = Field(
+        default=None,
+        description="Whether to automatically reconnect on failure"
+    )
+    max_concurrent_requests: Optional[int] = Field(
+        default=None,
+        ge=1,
+        le=50,
+        description="Maximum concurrent requests to this server"
+    )
+    message_endpoint: Optional[str] = Field(
+        default=None,
+        description="For SSE servers: explicit POST endpoint for sending messages"
+    )
+
+
+class ServerResponse(BaseModel):
+    """Response with server details."""
+
+    server_name: str
+    display_name: Optional[str]
+    server_url: str
+    server_type: str
+    auth_method: str
+    enabled: bool
+    timeout: int
+    retry_attempts: int
+    retry_delay: int
+    auto_reconnect: bool
+    max_concurrent_requests: int
+    message_endpoint: Optional[str]
+    auth_status: str
+    connection_count: int
+    tools_discovered: int
+    last_connected_at: Optional[str]
+    last_error: Optional[str]
+    created_at: str
+    updated_at: str
+
+    # Tool ID for use in session creation
+    tool_id: str
+
+
+class ServerListResponse(BaseModel):
+    """Response with list of servers."""
+
+    servers: List[ServerResponse]
+    count: int
+    max_servers: int
+
+
+class TestConnectionResponse(BaseModel):
+    """Response from connection test."""
+
+    success: bool
+    latency_ms: Optional[float]
+    error: Optional[str]
+    tools_discovered: Optional[int]
+
+
+class UserSettingsRequest(BaseModel):
+    """Request to update user MCP settings."""
+
+    mcp_enabled: Optional[bool] = Field(
+        default=None,
+        description="Enable/disable MCP for user"
+    )
+    include_global_servers: Optional[bool] = Field(
+        default=None,
+        description="Include global platform servers"
+    )
+    max_servers: Optional[int] = Field(
+        default=None,
+        ge=1,
+        le=50,
+        description="Maximum number of custom servers"
+    )
+    default_timeout: Optional[int] = Field(
+        default=None,
+        ge=5,
+        le=300,
+        description="Default timeout for new servers"
+    )
+
+
+class UserSettingsResponse(BaseModel):
+    """Response with user MCP settings."""
+
+    mcp_enabled: bool
+    include_global_servers: bool
+    max_servers: int
+    default_timeout: int
+
+
+class AvailableServersResponse(BaseModel):
+    """Response with all available MCP servers for session creation."""
+
+    global_servers: List[Dict[str, Any]]
+    user_servers: List[Dict[str, Any]]
+    include_global: bool
+
+
+class AuditLogEntry(BaseModel):
+    """Single audit log entry."""
+
+    id: int
+    action: str
+    server_name: Optional[str]
+    details: Optional[Dict[str, Any]]
+    ip_address: Optional[str]
+    timestamp: str
+
+
+class AuditLogResponse(BaseModel):
+    """Response with audit log entries."""
+
+    entries: List[AuditLogEntry]
+    count: int
+
+
+class ErrorResponse(BaseModel):
+    """Standard error response."""
+
+    error: str
+    detail: Optional[str] = None