router-maestro 0.1.2__py3-none-any.whl

router_maestro/server/oauth_sessions.py

@@ -0,0 +1,159 @@
+"""OAuth session management for remote OAuth flows."""
+
+import secrets
+import time
+from dataclasses import dataclass
+from threading import Lock
+
+
+@dataclass
+class OAuthSession:
+    """An OAuth session for device flow authentication."""
+
+    session_id: str
+    provider: str
+    device_code: str
+    user_code: str
+    verification_uri: str
+    expires_at: float
+    interval: int
+    status: str = "pending"  # pending, complete, expired, error
+    error: str | None = None
+    # Result data when complete
+    access_token: str | None = None
+    refresh_token: str | None = None
+
+
+class OAuthSessionManager:
+    """Manages OAuth sessions for remote device flow authentication."""
+
+    def __init__(self, session_timeout: int = 900) -> None:
+        """Initialize the session manager.
+
+        Args:
+            session_timeout: Default session timeout in seconds (default: 15 minutes)
+        """
+        self._sessions: dict[str, OAuthSession] = {}
+        self._lock = Lock()
+        self._session_timeout = session_timeout
+
+    def create_session(
+        self,
+        provider: str,
+        device_code: str,
+        user_code: str,
+        verification_uri: str,
+        expires_in: int,
+        interval: int = 5,
+    ) -> OAuthSession:
+        """Create a new OAuth session.
+
+        Args:
+            provider: Provider name (e.g., 'github-copilot')
+            device_code: Device code from OAuth provider
+            user_code: User code to display
+            verification_uri: URL for user to visit
+            expires_in: Seconds until expiration
+            interval: Polling interval in seconds
+
+        Returns:
+            The created OAuth session
+        """
+        session_id = secrets.token_urlsafe(16)
+        session = OAuthSession(
+            session_id=session_id,
+            provider=provider,
+            device_code=device_code,
+            user_code=user_code,
+            verification_uri=verification_uri,
+            expires_at=time.time() + expires_in,
+            interval=interval,
+        )
+
+        with self._lock:
+            self._sessions[session_id] = session
+            self._cleanup_expired()
+
+        return session
+
+    def get_session(self, session_id: str) -> OAuthSession | None:
+        """Get a session by ID.
+
+        Args:
+            session_id: Session ID to look up
+
+        Returns:
+            The session if found and not expired, None otherwise
+        """
+        with self._lock:
+            session = self._sessions.get(session_id)
+            if session is None:
+                return None
+
+            # Check if expired
+            if session.status == "pending" and time.time() > session.expires_at:
+                session.status = "expired"
+
+            return session
+
+    def update_session_status(
+        self,
+        session_id: str,
+        status: str,
+        error: str | None = None,
+        access_token: str | None = None,
+        refresh_token: str | None = None,
+    ) -> bool:
+        """Update session status.
+
+        Args:
+            session_id: Session ID to update
+            status: New status
+            error: Error message if status is 'error'
+            access_token: Access token if status is 'complete'
+            refresh_token: Refresh token if status is 'complete'
+
+        Returns:
+            True if session was updated, False if not found
+        """
+        with self._lock:
+            session = self._sessions.get(session_id)
+            if session is None:
+                return False
+
+            session.status = status
+            session.error = error
+            session.access_token = access_token
+            session.refresh_token = refresh_token
+            return True
+
+    def remove_session(self, session_id: str) -> bool:
+        """Remove a session.
+
+        Args:
+            session_id: Session ID to remove
+
+        Returns:
+            True if session was removed, False if not found
+        """
+        with self._lock:
+            if session_id in self._sessions:
+                del self._sessions[session_id]
+                return True
+            return False
+
+    def _cleanup_expired(self) -> None:
+        """Remove expired sessions. Must be called with lock held."""
+        now = time.time()
+        expired = [
+            sid
+            for sid, session in self._sessions.items()
+            if session.status in ("complete", "error", "expired")
+            or now > session.expires_at + 60  # Keep for 1 minute after expiry
+        ]
+        for sid in expired:
+            del self._sessions[sid]
+
+
+# Global session manager instance
+oauth_sessions = OAuthSessionManager()

router_maestro/server/routes/__init__.py

@@ -0,0 +1,8 @@
+"""Server routes module."""
+
+from router_maestro.server.routes.admin import router as admin_router
+from router_maestro.server.routes.anthropic import router as anthropic_router
+from router_maestro.server.routes.chat import router as chat_router
+from router_maestro.server.routes.models import router as models_router
+
+__all__ = ["admin_router", "anthropic_router", "chat_router", "models_router"]

router_maestro/server/routes/admin.py

@@ -0,0 +1,358 @@
+"""Admin API routes for remote management."""
+
+from typing import Annotated
+
+import httpx
+from fastapi import APIRouter, BackgroundTasks, HTTPException, Query
+
+from router_maestro.auth import AuthManager, AuthType
+from router_maestro.auth.github_oauth import (
+    GitHubOAuthError,
+    get_copilot_token,
+    poll_access_token,
+    request_device_code,
+)
+from router_maestro.auth.storage import OAuthCredential
+from router_maestro.config import (
+    load_priorities_config,
+    save_priorities_config,
+)
+from router_maestro.routing import get_router, reset_router
+from router_maestro.server.oauth_sessions import oauth_sessions
+from router_maestro.server.schemas.admin import (
+    AuthListResponse,
+    AuthProviderInfo,
+    LoginRequest,
+    ModelInfo,
+    ModelsResponse,
+    OAuthInitResponse,
+    OAuthStatusResponse,
+    PrioritiesResponse,
+    PrioritiesUpdateRequest,
+    StatsResponse,
+)
+from router_maestro.stats import StatsStorage
+
+router = APIRouter(prefix="/api/admin", tags=["admin"])
+
+
+# ============================================================================
+# Auth endpoints
+# ============================================================================
+
+
+@router.get("/auth", response_model=AuthListResponse)
+async def list_auth() -> AuthListResponse:
+    """List all authenticated providers."""
+    manager = AuthManager()
+    providers = []
+
+    for provider_name in manager.list_authenticated():
+        cred = manager.get_credential(provider_name)
+        if cred:
+            auth_type = "oauth" if cred.type == AuthType.OAUTH else "api"
+            # For OAuth, check if token might be expired
+            status = "active"
+            if isinstance(cred, OAuthCredential) and cred.expires > 0:
+                import time
+
+                if cred.expires < time.time():
+                    status = "expired"
+
+            providers.append(
+                AuthProviderInfo(
+                    provider=provider_name,
+                    auth_type=auth_type,
+                    status=status,
+                )
+            )
+
+    return AuthListResponse(providers=providers)
+
+
+@router.post("/auth/login")
+async def login(
+    request: LoginRequest,
+    background_tasks: BackgroundTasks,
+) -> OAuthInitResponse | dict:
+    """Initiate login for a provider.
+
+    For OAuth providers (github-copilot): Returns session info for device flow polling.
+    For API key providers: Saves the key and returns success.
+    """
+    manager = AuthManager()
+
+    if request.provider == "github-copilot":
+        # OAuth device flow
+        async with httpx.AsyncClient() as client:
+            try:
+                device_code = await request_device_code(client)
+            except httpx.HTTPError as e:
+                raise HTTPException(status_code=502, detail=f"Failed to get device code: {e}")
+
+        # Create session for polling
+        session = oauth_sessions.create_session(
+            provider=request.provider,
+            device_code=device_code.device_code,
+            user_code=device_code.user_code,
+            verification_uri=device_code.verification_uri,
+            expires_in=device_code.expires_in,
+            interval=device_code.interval,
+        )
+
+        # Start background task to poll for token
+        background_tasks.add_task(
+            _poll_oauth_completion,
+            session.session_id,
+            device_code.device_code,
+            device_code.interval,
+        )
+
+        return OAuthInitResponse(
+            session_id=session.session_id,
+            user_code=device_code.user_code,
+            verification_uri=device_code.verification_uri,
+            expires_in=device_code.expires_in,
+        )
+
+    elif request.api_key:
+        # API key auth
+        manager.login_api_key(request.provider, request.api_key)
+        # Reset router to pick up new authentication
+        reset_router()
+        return {"success": True, "provider": request.provider}
+
+    else:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Provider '{request.provider}' requires an API key",
+        )
+
+
+async def _poll_oauth_completion(
+    session_id: str,
+    device_code: str,
+    interval: int,
+) -> None:
+    """Background task to poll for OAuth completion and save credentials."""
+    manager = AuthManager()
+
+    async with httpx.AsyncClient() as client:
+        try:
+            # Poll for access token
+            access_token = await poll_access_token(
+                client,
+                device_code,
+                interval=interval,
+                timeout=900,  # 15 minutes
+            )
+
+            # Get Copilot token
+            copilot_token = await get_copilot_token(client, access_token.access_token)
+
+            # Save credentials
+            manager.storage.set(
+                "github-copilot",
+                OAuthCredential(
+                    refresh=access_token.access_token,
+                    access=copilot_token.token,
+                    expires=copilot_token.expires_at,
+                ),
+            )
+            manager.save()
+
+            # Update session status
+            oauth_sessions.update_session_status(
+                session_id,
+                status="complete",
+                access_token=copilot_token.token,
+                refresh_token=access_token.access_token,
+            )
+
+            # Reset router to pick up new authentication
+            reset_router()
+
+        except GitHubOAuthError as e:
+            oauth_sessions.update_session_status(
+                session_id,
+                status="error",
+                error=str(e),
+            )
+        except Exception as e:
+            oauth_sessions.update_session_status(
+                session_id,
+                status="error",
+                error=f"Unexpected error: {e}",
+            )
+
+
+@router.get("/auth/oauth/status/{session_id}", response_model=OAuthStatusResponse)
+async def get_oauth_status(session_id: str) -> OAuthStatusResponse:
+    """Get OAuth session status for polling."""
+    session = oauth_sessions.get_session(session_id)
+    if not session:
+        raise HTTPException(status_code=404, detail="Session not found")
+
+    return OAuthStatusResponse(
+        status=session.status,
+        error=session.error,
+    )
+
+
+@router.delete("/auth/{provider}")
+async def logout(provider: str) -> dict:
+    """Log out from a provider."""
+    manager = AuthManager()
+
+    if manager.logout(provider):
+        # Reset router to reflect authentication change
+        reset_router()
+        return {"success": True, "provider": provider}
+    else:
+        raise HTTPException(status_code=404, detail=f"Not authenticated with {provider}")
+
+
+# ============================================================================
+# Model endpoints
+# ============================================================================
+
+
+@router.get("/models", response_model=ModelsResponse)
+async def list_models() -> ModelsResponse:
+    """List all available models from authenticated providers."""
+    router_instance = get_router()
+
+    try:
+        models = await router_instance.list_models()
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Failed to list models: {e}")
+
+    model_list = [
+        ModelInfo(
+            provider=model.provider,
+            id=model.id,
+            name=model.name,
+        )
+        for model in models
+    ]
+
+    return ModelsResponse(models=model_list)
+
+
+@router.post("/models/refresh")
+async def refresh_models() -> dict:
+    """Force refresh the models cache."""
+    router_instance = get_router()
+    router_instance.invalidate_cache()
+    # Trigger re-population
+    try:
+        models = await router_instance.list_models()
+        return {"success": True, "models_count": len(models)}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Failed to refresh models: {e}")
+
+
+# ============================================================================
+# Priority endpoints
+# ============================================================================
+
+
+@router.get("/priorities", response_model=PrioritiesResponse)
+async def get_priorities() -> PrioritiesResponse:
+    """Get current priority configuration."""
+    config = load_priorities_config()
+
+    return PrioritiesResponse(
+        priorities=config.priorities,
+        fallback=config.fallback.model_dump(),
+    )
+
+
+@router.put("/priorities", response_model=PrioritiesResponse)
+async def update_priorities(request: PrioritiesUpdateRequest) -> PrioritiesResponse:
+    """Update priority configuration."""
+    config = load_priorities_config()
+
+    # Update priorities
+    config.priorities = request.priorities
+
+    # Update fallback if provided
+    if request.fallback is not None:
+        from router_maestro.config import FallbackConfig
+
+        config.fallback = FallbackConfig.model_validate(request.fallback)
+
+    save_priorities_config(config)
+
+    return PrioritiesResponse(
+        priorities=config.priorities,
+        fallback=config.fallback.model_dump(),
+    )
+
+
+# ============================================================================
+# Stats endpoints
+# ============================================================================
+
+
+@router.get("/stats", response_model=StatsResponse)
+async def get_stats(
+    days: Annotated[int, Query(ge=1, le=365)] = 7,
+    provider: str | None = None,
+    model: str | None = None,
+) -> StatsResponse:
+    """Get usage statistics."""
+    storage = StatsStorage()
+
+    # Get total stats
+    total = storage.get_total_usage(days=days)
+
+    # Get stats by model (which includes provider info)
+    by_model_raw = storage.get_usage_by_model(days=days)
+
+    # Aggregate by provider
+    by_provider: dict[str, dict] = {}
+    by_model: dict[str, dict] = {}
+
+    for record in by_model_raw:
+        provider_name = record["provider"]
+        model_name = record["model"]
+        model_key = f"{provider_name}/{model_name}"
+
+        # Filter if requested
+        if provider and provider_name != provider:
+            continue
+        if model and model_name != model:
+            continue
+
+        # Aggregate by provider
+        if provider_name not in by_provider:
+            by_provider[provider_name] = {
+                "total_tokens": 0,
+                "prompt_tokens": 0,
+                "completion_tokens": 0,
+                "request_count": 0,
+            }
+
+        by_provider[provider_name]["total_tokens"] += record.get("total_tokens", 0) or 0
+        by_provider[provider_name]["prompt_tokens"] += record.get("prompt_tokens", 0) or 0
+        by_provider[provider_name]["completion_tokens"] += record.get("completion_tokens", 0) or 0
+        by_provider[provider_name]["request_count"] += record.get("request_count", 0) or 0
+
+        # Store by model
+        by_model[model_key] = {
+            "total_tokens": record.get("total_tokens", 0) or 0,
+            "prompt_tokens": record.get("prompt_tokens", 0) or 0,
+            "completion_tokens": record.get("completion_tokens", 0) or 0,
+            "request_count": record.get("request_count", 0) or 0,
+            "avg_latency_ms": record.get("avg_latency_ms", 0) or 0,
+        }
+
+    return StatsResponse(
+        total_requests=total.get("request_count", 0) or 0,
+        total_tokens=total.get("total_tokens", 0) or 0,
+        prompt_tokens=total.get("prompt_tokens", 0) or 0,
+        completion_tokens=total.get("completion_tokens", 0) or 0,
+        by_provider=by_provider,
+        by_model=by_model,
+    )