Spaces:
Sleeping
Sleeping
| """FastAPI route handlers.""" | |
| import os | |
| import shutil | |
| import tempfile | |
| import time | |
| import traceback | |
| import uuid | |
| from pathlib import Path | |
| from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException, Request | |
| from fastapi.responses import HTMLResponse, StreamingResponse | |
| from loguru import logger | |
| from config.settings import Settings | |
| from providers.common import get_user_facing_error_message | |
| from providers.exceptions import InvalidRequestError, ProviderError | |
| from .dependencies import get_provider_for_type, get_settings, require_api_key | |
| from .models.anthropic import MessagesRequest, TokenCountRequest | |
| from .models.responses import TokenCountResponse | |
| from .optimization_handlers import try_optimizations | |
| from .request_utils import get_token_count | |
| router = APIRouter() | |
| def _home_page_html(status_payload: dict[str, str]) -> str: | |
| """Render the home page HTML with the factory reset button.""" | |
| return f""" | |
| <!doctype html> | |
| <html lang=\"en\"> | |
| <head> | |
| <meta charset=\"utf-8\" /> | |
| <meta name=\"viewport\" content=\"width=device-width, initial-scale=1\" /> | |
| <title>Claude Code Proxy</title> | |
| <style> | |
| :root {{ color-scheme: dark; }} | |
| body {{ margin: 0; min-height: 100vh; display: grid; place-items: center; font-family: Inter, Segoe UI, Arial, sans-serif; background: radial-gradient(1200px 600px at 20% 10%, #2e1b4a, #12131f 45%, #0b0c12); color: #e8ecff; }} | |
| .card {{ width: min(92vw, 560px); background: rgba(255, 255, 255, 0.06); border: 1px solid rgba(255, 255, 255, 0.14); border-radius: 20px; padding: 28px; box-shadow: 0 20px 50px rgba(0, 0, 0, 0.35); }} | |
| h1 {{ margin: 0 0 10px; font-size: 1.35rem; }} | |
| p {{ margin: 0 0 8px; color: #cfd7ff; line-height: 1.45; }} | |
| .meta {{ margin: 14px 0 18px; font-size: 0.95rem; color: #dfe6ff; }} | |
| .meta span {{ display: inline-block; margin-right: 12px; opacity: 0.95; }} | |
| button {{ border: none; border-radius: 14px; padding: 12px 18px; font-size: 1rem; font-weight: 700; color: white; cursor: pointer; background: linear-gradient(135deg, #ff507a, #7f5bff); box-shadow: 0 10px 20px rgba(127, 91, 255, 0.35); }} | |
| button:disabled {{ opacity: 0.65; cursor: wait; }} | |
| .status {{ margin-top: 14px; min-height: 24px; font-size: 0.95rem; color: #b8ffd8; }} | |
| </style> | |
| </head> | |
| <body> | |
| <div class=\"card\"> | |
| <h1>Claude Code Proxy</h1> | |
| <p>Server is running.</p> | |
| <div class=\"meta\"> | |
| <span><strong>Status:</strong> {status_payload['status']}</span> | |
| <span><strong>Provider:</strong> {status_payload['provider']}</span> | |
| <span><strong>Model:</strong> {status_payload['model']}</span> | |
| </div> | |
| <button id=\"resetBtn\">Factory Restart</button> | |
| <div class=\"status\" id=\"status\"></div> | |
| </div> | |
| <script> | |
| const btn = document.getElementById('resetBtn'); | |
| const status = document.getElementById('status'); | |
| btn.addEventListener('click', async () => {{ | |
| btn.disabled = true; | |
| status.textContent = 'Resetting cache/workspace and restarting...'; | |
| try {{ | |
| const response = await fetch('/admin/factory-reset' + window.location.search, {{ method: 'POST' }}); | |
| const data = await response.json(); | |
| if (!response.ok) {{ | |
| throw new Error(data.detail || 'Request failed'); | |
| }} | |
| status.textContent = 'Restart initiated. This page will disconnect briefly.'; | |
| }} catch (err) {{ | |
| status.textContent = 'Failed: ' + (err.message || String(err)); | |
| btn.disabled = false; | |
| }} | |
| }}); | |
| </script> | |
| </body> | |
| </html> | |
| """ | |
| def _clear_path(path: Path) -> int: | |
| """Best-effort removal of a file/directory path. Returns removed item count.""" | |
| if not path.exists(): | |
| return 0 | |
| try: | |
| if path.is_dir(): | |
| shutil.rmtree(path) | |
| else: | |
| path.unlink() | |
| return 1 | |
| except Exception as e: | |
| logger.warning("Failed to remove path {}: {}", path, e) | |
| return 0 | |
| def _clear_workspace_contents(workspace: Path) -> int: | |
| """Best-effort clear of workspace contents while preserving root directory.""" | |
| if not workspace.exists() or not workspace.is_dir(): | |
| return 0 | |
| removed = 0 | |
| for child in workspace.iterdir(): | |
| removed += _clear_path(child) | |
| return removed | |
| def _clear_runtime_state(settings: Settings) -> dict[str, int]: | |
| """Clear runtime caches/workspace data for a lightweight factory reset.""" | |
| removed = { | |
| "workspace_items": 0, | |
| "cache_dirs": 0, | |
| "pycache_dirs": 0, | |
| } | |
| workspace = Path(settings.claude_workspace).expanduser().resolve() | |
| removed["workspace_items"] = _clear_workspace_contents(workspace) | |
| cache_dirs = [ | |
| Path.home() / ".cache" / "huggingface", | |
| Path.home() / ".cache" / "uv", | |
| Path.home() / ".cache" / "pip", | |
| Path(tempfile.gettempdir()) / "huggingface", | |
| ] | |
| for cache_dir in cache_dirs: | |
| removed["cache_dirs"] += _clear_path(cache_dir) | |
| project_root = Path.cwd() | |
| for pycache_dir in project_root.rglob("__pycache__"): | |
| if ".venv" in pycache_dir.parts: | |
| continue | |
| removed["pycache_dirs"] += _clear_path(pycache_dir) | |
| return removed | |
| def _restart_process() -> None: | |
| """Terminate process so container orchestrator restarts the app.""" | |
| logger.warning("Factory reset requested: restarting process") | |
| time.sleep(1.0) | |
| os._exit(0) | |
| # ============================================================================= | |
| # Routes | |
| # ============================================================================= | |
| async def create_message( | |
| request_data: MessagesRequest, | |
| raw_request: Request, | |
| settings: Settings = Depends(get_settings), | |
| _auth=Depends(require_api_key), | |
| ): | |
| """Create a message (always streaming).""" | |
| try: | |
| if not request_data.messages: | |
| raise InvalidRequestError("messages cannot be empty") | |
| optimized = try_optimizations(request_data, settings) | |
| if optimized is not None: | |
| return optimized | |
| logger.debug("No optimization matched, routing to provider") | |
| # Resolve provider from the model-aware mapping | |
| provider_type = Settings.parse_provider_type( | |
| request_data.resolved_provider_model or settings.model | |
| ) | |
| provider = get_provider_for_type(provider_type) | |
| request_id = f"req_{uuid.uuid4().hex[:12]}" | |
| logger.info( | |
| "API_REQUEST: request_id={} model={} messages={}", | |
| request_id, | |
| request_data.model, | |
| len(request_data.messages), | |
| ) | |
| logger.debug("FULL_PAYLOAD [{}]: {}", request_id, request_data.model_dump()) | |
| input_tokens = get_token_count( | |
| request_data.messages, request_data.system, request_data.tools | |
| ) | |
| return StreamingResponse( | |
| provider.stream_response( | |
| request_data, | |
| input_tokens=input_tokens, | |
| request_id=request_id, | |
| ), | |
| media_type="text/event-stream", | |
| headers={ | |
| "X-Accel-Buffering": "no", | |
| "Cache-Control": "no-cache", | |
| "Connection": "keep-alive", | |
| }, | |
| ) | |
| except ProviderError: | |
| raise | |
| except Exception as e: | |
| logger.error(f"Error: {e!s}\n{traceback.format_exc()}") | |
| raise HTTPException( | |
| status_code=getattr(e, "status_code", 500), | |
| detail=get_user_facing_error_message(e), | |
| ) from e | |
| async def count_tokens(request_data: TokenCountRequest, _auth=Depends(require_api_key)): | |
| """Count tokens for a request.""" | |
| request_id = f"req_{uuid.uuid4().hex[:12]}" | |
| with logger.contextualize(request_id=request_id): | |
| try: | |
| tokens = get_token_count( | |
| request_data.messages, request_data.system, request_data.tools | |
| ) | |
| logger.info( | |
| "COUNT_TOKENS: request_id={} model={} messages={} input_tokens={}", | |
| request_id, | |
| getattr(request_data, "model", "unknown"), | |
| len(request_data.messages), | |
| tokens, | |
| ) | |
| return TokenCountResponse(input_tokens=tokens) | |
| except Exception as e: | |
| logger.error( | |
| "COUNT_TOKENS_ERROR: request_id={} error={}\n{}", | |
| request_id, | |
| get_user_facing_error_message(e), | |
| traceback.format_exc(), | |
| ) | |
| raise HTTPException( | |
| status_code=500, detail=get_user_facing_error_message(e) | |
| ) from e | |
| async def root( | |
| request: Request, | |
| settings: Settings = Depends(get_settings), | |
| _auth=Depends(require_api_key), | |
| ): | |
| """Root endpoint (JSON for API clients, HTML for browsers).""" | |
| payload = { | |
| "status": "ok", | |
| "provider": settings.provider_type, | |
| "model": settings.model, | |
| } | |
| accept = request.headers.get("accept", "") | |
| if "__sign" in request.query_params or "text/html" in accept.lower(): | |
| return HTMLResponse(content=_home_page_html(payload)) | |
| return payload | |
| async def health(): | |
| """Health check endpoint.""" | |
| return {"status": "healthy"} | |
| async def stop_cli(request: Request, _auth=Depends(require_api_key)): | |
| """Stop all CLI sessions and pending tasks.""" | |
| handler = getattr(request.app.state, "message_handler", None) | |
| if not handler: | |
| # Fallback if messaging not initialized | |
| cli_manager = getattr(request.app.state, "cli_manager", None) | |
| if cli_manager: | |
| await cli_manager.stop_all() | |
| logger.info("STOP_CLI: source=cli_manager cancelled_count=N/A") | |
| return {"status": "stopped", "source": "cli_manager"} | |
| raise HTTPException(status_code=503, detail="Messaging system not initialized") | |
| count = await handler.stop_all_tasks() | |
| logger.info("STOP_CLI: source=handler cancelled_count={}", count) | |
| return {"status": "stopped", "cancelled_count": count} | |
| async def factory_reset_page(request: Request, _auth=Depends(require_api_key)): | |
| """Simple admin UI for one-click factory reset and restart.""" | |
| return """ | |
| <!doctype html> | |
| <html lang=\"en\"> | |
| <head> | |
| <meta charset=\"utf-8\" /> | |
| <meta name=\"viewport\" content=\"width=device-width, initial-scale=1\" /> | |
| <title>Factory Reset</title> | |
| <style> | |
| :root { color-scheme: dark; } | |
| body { margin: 0; min-height: 100vh; display: grid; place-items: center; font-family: Inter, Segoe UI, Arial, sans-serif; background: radial-gradient(1200px 600px at 20% 10%, #2e1b4a, #12131f 45%, #0b0c12); color: #e8ecff; } | |
| .card { width: min(92vw, 520px); background: rgba(255, 255, 255, 0.06); border: 1px solid rgba(255, 255, 255, 0.14); border-radius: 20px; padding: 28px; box-shadow: 0 20px 50px rgba(0, 0, 0, 0.35); } | |
| h1 { margin: 0 0 10px; font-size: 1.35rem; } | |
| p { margin: 0 0 18px; color: #cfd7ff; line-height: 1.45; } | |
| button { border: none; border-radius: 14px; padding: 12px 18px; font-size: 1rem; font-weight: 700; color: white; cursor: pointer; background: linear-gradient(135deg, #ff507a, #7f5bff); box-shadow: 0 10px 20px rgba(127, 91, 255, 0.35); } | |
| button:disabled { opacity: 0.65; cursor: wait; } | |
| .status { margin-top: 14px; min-height: 24px; font-size: 0.95rem; color: #b8ffd8; } | |
| </style> | |
| </head> | |
| <body> | |
| <div class=\"card\"> | |
| <h1>Factory Reset & Restart</h1> | |
| <p>Clears runtime cache and workspace data, then restarts this server.</p> | |
| <button id=\"resetBtn\">Factory Restart</button> | |
| <div class=\"status\" id=\"status\"></div> | |
| </div> | |
| <script> | |
| const btn = document.getElementById('resetBtn'); | |
| const status = document.getElementById('status'); | |
| btn.addEventListener('click', async () => { | |
| btn.disabled = true; | |
| status.textContent = 'Resetting cache/workspace and restarting...'; | |
| try { | |
| const response = await fetch('/admin/factory-reset' + window.location.search, { method: 'POST' }); | |
| const data = await response.json(); | |
| if (!response.ok) { | |
| throw new Error(data.detail || 'Request failed'); | |
| } | |
| status.textContent = 'Restart initiated. This page will disconnect briefly.'; | |
| } catch (err) { | |
| status.textContent = 'Failed: ' + (err.message || String(err)); | |
| btn.disabled = false; | |
| } | |
| }); | |
| </script> | |
| </body> | |
| </html> | |
| """ | |
| async def factory_reset( | |
| background_tasks: BackgroundTasks, | |
| settings: Settings = Depends(get_settings), | |
| _auth=Depends(require_api_key), | |
| ): | |
| """Clear runtime state and restart process (for Space maintenance).""" | |
| cleared = _clear_runtime_state(settings) | |
| background_tasks.add_task(_restart_process) | |
| return { | |
| "status": "restarting", | |
| "cleared": cleared, | |
| } | |