wright 0.1.0__py3-none-any.whl

api/auth.py

@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+import os
+import secrets
+import sys
+from pathlib import Path
+
+import httpx
+from fastapi import HTTPException, Request, Security
+from fastapi.security import APIKeyHeader, HTTPAuthorizationCredentials, HTTPBearer
+from jose import JWTError, jwt
+
+_API_KEY_HEADER = APIKeyHeader(name="X-Wright-API-Key", auto_error=False)
+_BEARER = HTTPBearer(auto_error=False)
+_KEY_FILE = Path(os.getenv("WRIGHT_KEY_FILE", Path.home() / ".wright" / "api.key"))
+
+WORKOS_CLIENT_ID = os.getenv("WORKOS_CLIENT_ID", "")
+WORKOS_API_KEY = os.getenv("WORKOS_API_KEY", "")
+
+
+def _load_or_generate_key() -> str:
+    """
+    Loads or generates a WrightAI API key by checking the environment variable, an existing key file, or creating a new cryptographically secure random key.
+
+    Attempts to retrieve the API key in the following priority order: (1) from the WRIGHT_API_KEY environment variable, (2) from an existing key file at the path defined by _KEY_FILE, or (3) generates a new cryptographically secure random key using secrets.token_urlsafe(32), saves it to the key file with restricted permissions (0o600), creates parent directories with mode 0o700 if needed, and prints a notification message to stderr on first run.
+
+    Returns:
+        str: The WrightAI API key as a URL-safe string, sourced from the WRIGHT_API_KEY environment variable, an existing key file, or a newly generated secure random token.
+
+    Example:
+        ```
+        api_key = _load_or_generate_key()
+        print(api_key)  # e.g., 'aB3dEfGhIjKlMnOpQrStUvWxYz0123456789-_Ab'
+        ```
+    """
+    env_key = os.getenv("WRIGHT_API_KEY", "")
+    if env_key:
+        return env_key
+    if _KEY_FILE.exists():
+        return _KEY_FILE.read_text().strip()
+    _KEY_FILE.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
+    key = secrets.token_urlsafe(32)
+    _KEY_FILE.write_text(key)
+    _KEY_FILE.chmod(0o600)
+    print(
+        f"\n WrightAI API key generated (first run).\n"
+        f"  Saved to: {_KEY_FILE}\n"
+        f"  Read it with: cat {_KEY_FILE}\n",
+        file=sys.stderr,
+    )
+    return key
+
+
+_WRIGHT_API_KEY = _load_or_generate_key()
+
+
+_jwks_cache: dict | None = None
+
+
+def _get_jwks() -> dict:
+    """Fetch and cache the WorkOS JWKS public key set."""
+    global _jwks_cache
+    if _jwks_cache is None:
+        client_id = os.getenv("WORKOS_CLIENT_ID", "")
+        resp = httpx.get(
+            f"https://api.workos.com/user_management/jwks/{client_id}",
+            timeout=5,
+        )
+        if resp.status_code != 200:
+            raise HTTPException(status_code=503, detail="Could not fetch WorkOS JWKS")
+        _jwks_cache = resp.json()
+    return _jwks_cache
+
+
+def _verify_workos_token(token: str) -> dict:
+    """
+    Validates a WorkOS JWT access token by fetching the JWKS public keys and verifying the signature.
+
+    Fetches the WorkOS JWKS (JSON Web Key Set) for the configured client ID, then decodes and
+    verifies the provided JWT against those public keys. Raises HTTP 401 if the token is invalid,
+    expired, or the signature doesn't match.
+
+    Args:
+        token (str): The WorkOS JWT access token to validate.
+
+    Returns:
+        dict: The decoded JWT payload containing user identity claims.
+
+    Raises:
+        HTTPException: Raised with HTTP status code 401 when the token is invalid or expired.
+        HTTPException: Raised with HTTP status code 503 when the JWKS endpoint is unreachable.
+    """
+    try:
+        jwks = _get_jwks()
+        # Extract the key ID from the JWT header to pick the right key from the set
+        header = jwt.get_unverified_header(token)
+        kid = header.get("kid")
+        key = next(
+            (k for k in jwks.get("keys", []) if k.get("kid") == kid),
+            jwks.get("keys", [None])[0] if jwks.get("keys") else None,
+        )
+        if key is None:
+            raise HTTPException(status_code=401, detail="Invalid WorkOS token: no matching key")
+        claims = jwt.decode(
+            token,
+            key,
+            algorithms=["RS256"],
+            audience=os.getenv("WORKOS_CLIENT_ID", ""),
+        )
+        return claims
+    except JWTError as exc:
+        raise HTTPException(status_code=401, detail=f"Invalid WorkOS token: {exc}") from exc
+
+
+async def verify_api_key(
+    request: Request,
+    api_key: str | None = Security(_API_KEY_HEADER),
+    bearer: HTTPAuthorizationCredentials | None = Security(_BEARER),
+) -> None:
+    """
+    Verifies incoming API authentication credentials against user-issued API keys, a server-level static key, or WorkOS bearer tokens, raising HTTP 401 on any failure.
+
+    Authenticates requests via one of three ordered methods: (1) user-issued API keys prefixed with 'wai_' are looked up in the Supabase user store via get_user_by_api_key(); (2) a server-level static key is accepted directly for CLI, GitHub Actions, or MCP integrations; (3) a WorkOS bearer token is verified by calling _verify_workos_token(). The function returns silently on the first successful match. If no method succeeds or no credentials are supplied at all, an HTTP 401 exception is raised.
+
+    Args:
+        request (Request): The incoming FastAPI/Starlette HTTP request object, injected by the dependency injection system.
+        api_key (str | None): Optional API key extracted from the request headers. Accepts either a user-issued key (prefixed with 'wai_') validated against Supabase, or a server-level static key for CLI/GitHub Actions/MCP.
+        bearer (HTTPAuthorizationCredentials | None): Optional bearer token credentials extracted from the Authorization header, used for WorkOS token authentication.
+
+    Returns:
+        None: Returns None implicitly on successful authentication; no value is produced.
+
+    Raises:
+        HTTPException: Raised with status 401 when an API key with the 'wai_' prefix is not found in the Supabase user store.
+        HTTPException: Raised with status 401 when WorkOS bearer token verification fails via _verify_workos_token().
+        HTTPException: Raised with status 401 when no valid credentials are provided or all authentication methods fail.
+
+    Example:
+        ```
+        await verify_api_key(request, api_key='wai_abc123def456', bearer=None)
+        ```
+    """
+    # Accept user API keys issued via WorkOS + Supabase (wai_ prefix)
+    if api_key and api_key.startswith("wai_"):
+        from api.user_store import get_user_by_api_key
+
+        if get_user_by_api_key(api_key):
+            return
+        raise HTTPException(status_code=401, detail="Invalid API key")
+
+    # Accept server-level static key (CLI / GitHub Action / MCP)
+    if api_key and api_key == _WRIGHT_API_KEY:
+        return
+
+    # Accept WorkOS Bearer token directly
+    if bearer and bearer.credentials:
+        try:
+            _verify_workos_token(bearer.credentials)
+            return
+        except HTTPException:
+            raise
+        except Exception:
+            raise HTTPException(status_code=401, detail="Invalid WorkOS token")
+
+    raise HTTPException(status_code=401, detail="Invalid or missing credentials")

api/chroma_cache.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+import logging
+import shutil
+import threading
+import time
+from pathlib import Path
+
+_logger = logging.getLogger("wright.chroma_cache")
+
+_cache: dict[str, tuple[object, float]] = {}
+_lock = threading.Lock()
+_TTL = 300  # seconds — accept up to 5 min staleness from cross-container writes
+_restored_paths: set[str] = set()  # per persist_path, one-time GCS restore
+
+
+def _restore_chroma_from_gcs(persist_path: str, repo_root: str) -> None:
+    """Copy the per-user/repo GCS backup into persist_path on first access."""
+    from api.routes.repos import _parse_repo_root  # lazy — avoids circular import
+
+    dst = Path(persist_path)
+    if dst.exists():
+        return
+    parsed = _parse_repo_root(repo_root)
+    src = Path("/data/chroma") / parsed[0] / parsed[1] if parsed else Path("/data/chroma")
+    if src.exists():
+        try:
+            shutil.copytree(str(src), str(dst))
+        except Exception as e:
+            _logger.warning("ChromaDB GCS restore failed for %s: %s", persist_path, e)
+
+
+def get(persist_path: str, repo_root: str):
+    """Return a cached ChromaStore, creating one if missing or TTL-expired.
+
+    Opening a PersistentClient on GCS Fuse loads SQLite and the HNSW index from
+    the network filesystem — ~1-3 s per cold open. Caching the instance means
+    only the first request per (path, repo_root) per container pays that cost;
+    subsequent requests reuse the warm in-memory index (~10-50 ms).
+
+    If the resulting store is empty (fresh /tmp on a cold container, with no
+    /data/chroma backup either), it's repopulated from the Supabase pgvector
+    backup — the durable source of truth for bidirectional sync.
+    """
+    from core.embeddings.chroma_store import ChromaStore
+
+    key = f"{persist_path}::{repo_root}"
+    now = time.monotonic()
+    needs_rebuild = False
+    needs_restore = False
+    with _lock:
+        if persist_path not in _restored_paths:
+            _restored_paths.add(persist_path)
+            needs_restore = True
+    if needs_restore:
+        _restore_chroma_from_gcs(persist_path, repo_root)
+    with _lock:
+        if key in _cache:
+            store, ts = _cache[key]
+            if now - ts < _TTL:
+                return store
+        store = ChromaStore(persist_path=persist_path, repo_root=repo_root)
+        _cache[key] = (store, now)
+        needs_rebuild = store.count() == 0
+
+    if needs_rebuild:
+        _rebuild_from_pgvector(store, repo_root)
+    return store
+
+
+def _rebuild_from_pgvector(store, repo_root: str) -> None:
+    """Repopulate a freshly-created, empty ChromaStore from its Supabase
+    pgvector backup (e.g. on a cold container start with no /data/chroma)."""
+    from api.routes.repos import _parse_repo_root  # lazy import, avoids circular import
+
+    parsed = _parse_repo_root(repo_root)
+    if parsed is None:
+        return
+
+    from core.embeddings.pgvector_store import PgVectorStore
+
+    chunks, embeddings = PgVectorStore(*parsed).get_all_chunks()
+    if chunks:
+        store.upsert_chunks(chunks, embeddings)
+        _logger.info(
+            "Rebuilt local chroma for %s from pgvector (%d chunks)", repo_root, len(chunks)
+        )
+
+
+def invalidate(persist_path: str, repo_root: str) -> None:
+    """Drop the cached entry so the next request reloads from disk.
+
+    Call this after a successful upsert to force the cache to reflect newly
+    written embeddings on the next request.
+    """
+    key = f"{persist_path}::{repo_root}"
+    with _lock:
+        _cache.pop(key, None)

api/embedder.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+import os
+
+_embedder = None
+_gateway = None
+
+
+def get_embedder():
+    """Return the process-level VoyageEmbedder singleton.
+
+    Creating a new VoyageEmbedder on every request initialises a voyageai.Client
+    object each time. This singleton avoids that overhead while keeping the API
+    key resolution lazy (read at first call, not at import time).
+    """
+    global _embedder
+    if _embedder is None:
+        from core.embeddings.voyage_embeddings import VoyageEmbedder
+
+        _embedder = VoyageEmbedder(api_key=os.getenv("VOYAGE_API_KEY", ""))
+    return _embedder
+
+
+def get_gateway():
+    """Return the process-level LLMGateway singleton.
+
+    LLMGateway.__init__ creates an anthropic.AsyncAnthropic client on every
+    call. Sharing one instance across requests avoids that allocation and keeps
+    the underlying httpx connection pool alive between requests.
+    """
+    global _gateway
+    if _gateway is None:
+        from core.llm.gateway import LLMGateway
+
+        _gateway = LLMGateway(
+            anthropic_key=os.getenv("ANTHROPIC_API_KEY", ""),
+            gemini_key=os.getenv("GEMINI_API_KEY"),
+        )
+    return _gateway

api/main.py

@@ -0,0 +1,264 @@
+# WrightAI — AI-powered code documentation tool
+# Copyright (C) 2026 Suraj Sahoo
+# SPDX-License-Identifier: AGPL-3.0-or-later
+# https://github.com/surajs1999/WrightAI
+from __future__ import annotations
+
+import logging
+import os
+import time
+
+import uvicorn
+from dotenv import load_dotenv
+from fastapi import FastAPI, Request
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse
+
+load_dotenv()
+
+app = FastAPI(
+    title="Wright API",
+    version="1.0.0",
+    description="AI-powered code documentation API",
+)
+
+from api.observability import setup_observability  # noqa: E402
+from api.rate_limit import limiter  # noqa: E402
+from slowapi import _rate_limit_exceeded_handler  # noqa: E402
+from slowapi.errors import RateLimitExceeded  # noqa: E402
+
+setup_observability(app)
+app.state.limiter = limiter
+app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
+
+_ALLOWED_ORIGINS = [
+    o.strip()
+    for o in os.getenv(
+        "CORS_ORIGINS",
+        "http://localhost:3000,https://www.wrightai.live,https://wrightai.live,https://wrightai-web.netlify.app",
+    ).split(",")
+    if o.strip()
+]
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=_ALLOWED_ORIGINS,
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+_logger = logging.getLogger("wright.api")
+
+
+@app.middleware("http")
+async def log_requests(request: Request, call_next):
+    """
+    Logs HTTP request details including method, path, status code, and processing duration.
+
+    FastAPI middleware that intercepts HTTP requests, measures the time taken to process them, and logs the request method, URL path, response status code, and duration in milliseconds.
+
+    Args:
+        request (Request): The incoming HTTP request object containing request metadata and information.
+        call_next (Callable): Middleware chain callable that processes the request and returns the response.
+
+    Returns:
+        Response: The HTTP response object returned from the next middleware or route handler.
+
+    Example:
+        ```
+        # Registered automatically via the @app.middleware("http") decorator above.
+        # Logs one line per request, e.g.:
+        # 2026-01-01T00:00:00 INFO GET /health 200 1.2ms
+        ```
+
+    Complexity: O(1) time, O(1) space
+    """
+    start = time.perf_counter()
+    response = await call_next(request)
+    duration_ms = (time.perf_counter() - start) * 1000
+    _logger.info(
+        "http_request",
+        extra={
+            "http_method": request.method,
+            "http_path": request.url.path,
+            "http_status": response.status_code,
+            "duration_ms": round(duration_ms, 1),
+        },
+    )
+    return response
+
+
+from api.routes import (  # noqa: E402
+    auth,
+    billing,
+    chat,
+    coverage,
+    drift,
+    fix_pr,
+    generate,
+    internal,
+    llms_txt,
+    repos,
+    usage,
+    webhooks,
+)
+
+app.include_router(auth.router)
+app.include_router(billing.router)
+app.include_router(generate.router)
+app.include_router(coverage.router)
+app.include_router(drift.router)
+app.include_router(chat.router)
+app.include_router(repos.router)
+app.include_router(fix_pr.router)
+app.include_router(llms_txt.router)
+app.include_router(usage.router)
+app.include_router(webhooks.router)
+app.include_router(internal.router)
+
+
+@app.get("/health")
+async def health() -> dict:
+    """Liveness probe — returns immediately. Used by Cloud Run to detect crashes."""
+    return {"status": "ok", "version": "1.0.0"}
+
+
+@app.get("/ready")
+async def ready() -> JSONResponse:
+    """Readiness probe — checks live dependencies before accepting traffic."""
+    checks: dict[str, str] = {}
+
+    # Supabase reachability
+    try:
+        from api.user_store import _db
+
+        _db().table("plans").select("id").limit(1).execute()
+        checks["database"] = "ok"
+    except Exception:
+        checks["database"] = "error"
+
+    # Required LLM credentials
+    checks["anthropic_key"] = "ok" if os.getenv("ANTHROPIC_API_KEY") else "missing"
+    checks["voyage_key"] = "ok" if os.getenv("VOYAGE_API_KEY") else "missing"
+
+    all_ok = all(v == "ok" for v in checks.values())
+    return JSONResponse(
+        {"status": "ready" if all_ok else "degraded", "checks": checks},
+        status_code=200 if all_ok else 503,
+    )
+
+
+from fastapi import Depends  # noqa: E402
+from api.auth import verify_api_key  # noqa: E402
+
+
+@app.post("/user/key/rotate", dependencies=[Depends(verify_api_key)])
+async def rotate_key(request: Request) -> dict:
+    """
+    Rotates the API key for the authenticated user by invalidating the old key and generating a new one.
+
+    This FastAPI POST endpoint reads the current API key from the 'X-Wright-API-Key' request header, delegates key rotation to the user store, and returns the newly generated API key. Access is protected by the 'verify_api_key' dependency. If no user is associated with the provided key, a 404 HTTP exception is raised.
+
+    Args:
+        request (Request): The incoming FastAPI HTTP request object, used to extract the current API key from the 'X-Wright-API-Key' header.
+
+    Returns:
+        dict: A dictionary containing the newly generated API key under the key 'api_key', e.g., {'api_key': 'new-generated-key-string'}.
+
+    Raises:
+        HTTPException: Raised with status code 404 and detail 'User not found' when no user is associated with the provided API key.
+
+    Example:
+        ```
+        # Using httpx or requests in an async context:
+        import httpx
+
+        response = httpx.post(
+            'http://localhost:8000/user/key/rotate',
+            headers={'X-Wright-API-Key': 'existing-api-key-abc123'}
+        )
+        new_key = response.json()  # {'api_key': 'newly-rotated-key-xyz789'}
+        ```
+    """
+    from api.user_store import rotate_api_key
+    from fastapi import HTTPException
+
+    old_key = request.headers.get("X-Wright-API-Key", "")
+    user = rotate_api_key(old_key)
+    if not user:
+        raise HTTPException(status_code=404, detail="User not found")
+    return {"api_key": user.api_key}
+
+
+@app.get("/user/me", dependencies=[Depends(verify_api_key)])
+async def user_me(request: Request) -> dict:
+    """
+    Retrieves the authenticated user's profile information by looking up the API key from the request headers.
+
+    An async FastAPI endpoint protected by the verify_api_key dependency. It extracts the X-Wright-API-Key header from the incoming request, resolves the corresponding user via the user store, and returns the user's email, API key, and account creation timestamp. Raises a 401 Unauthorized exception if the key is absent or unrecognized.
+
+    Args:
+        request (Request): The FastAPI Request object used to access HTTP headers, specifically the X-Wright-API-Key header for user lookup.
+
+    Returns:
+        dict: A dictionary containing three keys: 'email' (str) — the user's email address, 'api_key' (str) — the user's API key, and 'created_at' (str or datetime) — the ISO 8601 timestamp of account creation.
+
+    Raises:
+        HTTPException: Raised with status_code=401 when the X-Wright-API-Key header is missing, empty, or does not correspond to any user in the user store.
+
+    Example:
+        ```
+        # GET /user/me with Header: X-Wright-API-Key: abc123xyz
+        # Response:
+        # {
+        #   "email": "user@example.com",
+        #   "api_key": "abc123xyz",
+        #   "created_at": "2024-01-01T00:00:00"
+        # }
+        ```
+    """
+    from api.user_store import get_user_by_api_key
+
+    api_key = request.headers.get("X-Wright-API-Key", "")
+    user = get_user_by_api_key(api_key)
+    if not user:
+        from fastapi import HTTPException
+
+        raise HTTPException(status_code=401, detail="Invalid API key")
+    return {
+        "email": user.email,
+        "api_key": user.api_key,
+        "created_at": user.created_at,
+    }
+
+
+def start() -> None:
+    """
+    Starts the Uvicorn ASGI server to run the Wright API application on the configured host and port.
+
+    Initializes and runs the FastAPI application using Uvicorn with the host bound to all network interfaces (0.0.0.0) and the port configured via the WRIGHT_API_PORT environment variable (defaulting to 8765). Auto-reload is disabled for production stability. This function blocks until the server is interrupted.
+
+    Returns:
+        None: Does not return a value; runs the server until interrupted.
+
+    Raises:
+        ValueError: When the WRIGHT_API_PORT environment variable is set to a non-integer value that cannot be converted via int().
+
+    Example:
+        ```
+        import os
+        os.environ['WRIGHT_API_PORT'] = '8765'
+        start()  # Starts the Wright API server on http://0.0.0.0:8765
+        ```
+    """
+    uvicorn.run(
+        "api.main:app",
+        host="0.0.0.0",
+        port=int(os.getenv("WRIGHT_API_PORT", "8765")),
+        reload=False,
+    )
+
+
+if __name__ == "__main__":
+    start()

api/observability.py

@@ -0,0 +1,74 @@
+"""
+Observability setup: structured JSON logging + Sentry error tracking.
+Call setup_observability(app) once, right after FastAPI() is created.
+
+Required env vars:
+  SENTRY_DSN        — enables Sentry (get from sentry.io → Project Settings → DSN)
+
+Optional env vars:
+  ENVIRONMENT       — "production" / "staging"  (default "production")
+  K_REVISION        — Cloud Run revision tag, used as Sentry release identifier
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+
+_logger = logging.getLogger("wright.observability")
+
+
+def configure_logging() -> None:
+    """Replace the plaintext formatter with JSON so Cloud Logging can parse fields."""
+    try:
+        from pythonjsonlogger import jsonlogger
+
+        handler = logging.StreamHandler()
+        handler.setFormatter(
+            jsonlogger.JsonFormatter(
+                fmt="%(asctime)s %(levelname)s %(name)s %(message)s",
+                datefmt="%Y-%m-%dT%H:%M:%SZ",
+                rename_fields={"levelname": "severity", "asctime": "timestamp"},
+            )
+        )
+        logging.root.setLevel(logging.INFO)
+        logging.root.handlers = []
+        logging.root.addHandler(handler)
+    except ImportError:
+        logging.basicConfig(
+            level=logging.INFO,
+            format="%(asctime)s %(levelname)s %(name)s %(message)s",
+            datefmt="%Y-%m-%dT%H:%M:%S",
+        )
+
+
+def _setup_sentry() -> None:
+    dsn = os.getenv("SENTRY_DSN", "")
+    if not dsn:
+        return
+    try:
+        import sentry_sdk
+        from sentry_sdk.integrations.fastapi import FastApiIntegration
+        from sentry_sdk.integrations.logging import LoggingIntegration
+
+        sentry_sdk.init(
+            dsn=dsn,
+            integrations=[
+                FastApiIntegration(),
+                LoggingIntegration(level=logging.ERROR, event_level=logging.ERROR),
+            ],
+            traces_sample_rate=0.0,  # no performance tracing — error tracking only
+            environment=os.getenv("ENVIRONMENT", "production"),
+            release=os.getenv("K_REVISION", "unknown"),
+        )
+        _logger.info(
+            "Sentry initialised", extra={"environment": os.getenv("ENVIRONMENT", "production")}
+        )
+    except ImportError:
+        _logger.warning("sentry-sdk not installed — error tracking disabled")
+
+
+def setup_observability(app) -> None:  # noqa: ARG001
+    """Wire up structured logging and Sentry. Call once after FastAPI() is created."""
+    configure_logging()
+    _setup_sentry()