Jack1808/Claude_Code
0
1"""FastAPI route handlers."""2 3import os4import shutil5import tempfile6import time7import traceback8import uuid9from pathlib import Path10 11from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException, Request12from fastapi.responses import HTMLResponse, StreamingResponse13from loguru import logger14 15from config.settings import Settings16from providers.common import get_user_facing_error_message17from providers.exceptions import InvalidRequestError, ProviderError18 19from .dependencies import get_provider_for_type, get_settings, require_api_key20from .models.anthropic import MessagesRequest, TokenCountRequest21from .models.responses import TokenCountResponse22from .optimization_handlers import try_optimizations23from .request_utils import get_token_count24 25router = APIRouter()26 27 28def _home_page_html(status_payload: dict[str, str]) -> str:29 """Render the home page HTML with the factory reset button."""30 return f"""31<!doctype html>32<html lang=\"en\">33<head>34 <meta charset=\"utf-8\" />35 <meta name=\"viewport\" content=\"width=device-width, initial-scale=1\" />36 <title>Claude Code Proxy</title>37 <style>38 :root {{ color-scheme: dark; }}39 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; }}40 .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); }}41 h1 {{ margin: 0 0 10px; font-size: 1.35rem; }}42 p {{ margin: 0 0 8px; color: #cfd7ff; line-height: 1.45; }}43 .meta {{ margin: 14px 0 18px; font-size: 0.95rem; color: #dfe6ff; }}44 .meta span {{ display: inline-block; margin-right: 12px; opacity: 0.95; }}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); }}46 button:disabled {{ opacity: 0.65; cursor: wait; }}47 .status {{ margin-top: 14px; min-height: 24px; font-size: 0.95rem; color: #b8ffd8; }}48 </style>49</head>50<body>51 <div class=\"card\">52 <h1>Claude Code Proxy</h1>53 <p>Server is running.</p>54 <div class=\"meta\">55 <span><strong>Status:</strong> {status_payload['status']}</span>56 <span><strong>Provider:</strong> {status_payload['provider']}</span>57 <span><strong>Model:</strong> {status_payload['model']}</span>58 </div>59 <button id=\"resetBtn\">Factory Restart</button>60 <div class=\"status\" id=\"status\"></div>61 </div>62 <script>63 const btn = document.getElementById('resetBtn');64 const status = document.getElementById('status');65 btn.addEventListener('click', async () => {{66 btn.disabled = true;67 status.textContent = 'Resetting cache/workspace and restarting...';68 try {{69 const response = await fetch('/admin/factory-reset' + window.location.search, {{ method: 'POST' }});70 const data = await response.json();71 if (!response.ok) {{72 throw new Error(data.detail || 'Request failed');73 }}74 status.textContent = 'Restart initiated. This page will disconnect briefly.';75 }} catch (err) {{76 status.textContent = 'Failed: ' + (err.message || String(err));77 btn.disabled = false;78 }}79 }});80 </script>81</body>82</html>83"""84 85 86def _clear_path(path: Path) -> int:87 """Best-effort removal of a file/directory path. Returns removed item count."""88 if not path.exists():89 return 090 try:91 if path.is_dir():92 shutil.rmtree(path)93 else:94 path.unlink()95 return 196 except Exception as e:97 logger.warning("Failed to remove path {}: {}", path, e)98 return 099 100 101def _clear_workspace_contents(workspace: Path) -> int:102 """Best-effort clear of workspace contents while preserving root directory."""103 if not workspace.exists() or not workspace.is_dir():104 return 0105 removed = 0106 for child in workspace.iterdir():107 removed += _clear_path(child)108 return removed109 110 111def _clear_runtime_state(settings: Settings) -> dict[str, int]:112 """Clear runtime caches/workspace data for a lightweight factory reset."""113 removed = {114 "workspace_items": 0,115 "cache_dirs": 0,116 "pycache_dirs": 0,117 }118 119 workspace = Path(settings.claude_workspace).expanduser().resolve()120 removed["workspace_items"] = _clear_workspace_contents(workspace)121 122 cache_dirs = [123 Path.home() / ".cache" / "huggingface",124 Path.home() / ".cache" / "uv",125 Path.home() / ".cache" / "pip",126 Path(tempfile.gettempdir()) / "huggingface",127 ]128 for cache_dir in cache_dirs:129 removed["cache_dirs"] += _clear_path(cache_dir)130 131 project_root = Path.cwd()132 for pycache_dir in project_root.rglob("__pycache__"):133 if ".venv" in pycache_dir.parts:134 continue135 removed["pycache_dirs"] += _clear_path(pycache_dir)136 137 return removed138 139 140def _restart_process() -> None:141 """Terminate process so container orchestrator restarts the app."""142 logger.warning("Factory reset requested: restarting process")143 time.sleep(1.0)144 os._exit(0)145 146 147# =============================================================================148# Routes149# =============================================================================150@router.post("/v1/messages")151async def create_message(152 request_data: MessagesRequest,153 raw_request: Request,154 settings: Settings = Depends(get_settings),155 _auth=Depends(require_api_key),156):157 """Create a message (always streaming)."""158 159 try:160 if not request_data.messages:161 raise InvalidRequestError("messages cannot be empty")162 163 optimized = try_optimizations(request_data, settings)164 if optimized is not None:165 return optimized166 logger.debug("No optimization matched, routing to provider")167 168 # Resolve provider from the model-aware mapping169 provider_type = Settings.parse_provider_type(170 request_data.resolved_provider_model or settings.model171 )172 provider = get_provider_for_type(provider_type)173 174 request_id = f"req_{uuid.uuid4().hex[:12]}"175 logger.info(176 "API_REQUEST: request_id={} model={} messages={}",177 request_id,178 request_data.model,179 len(request_data.messages),180 )181 logger.debug("FULL_PAYLOAD [{}]: {}", request_id, request_data.model_dump())182 183 input_tokens = get_token_count(184 request_data.messages, request_data.system, request_data.tools185 )186 return StreamingResponse(187 provider.stream_response(188 request_data,189 input_tokens=input_tokens,190 request_id=request_id,191 ),192 media_type="text/event-stream",193 headers={194 "X-Accel-Buffering": "no",195 "Cache-Control": "no-cache",196 "Connection": "keep-alive",197 },198 )199 200 except ProviderError:201 raise202 except Exception as e:203 logger.error(f"Error: {e!s}\n{traceback.format_exc()}")204 raise HTTPException(205 status_code=getattr(e, "status_code", 500),206 detail=get_user_facing_error_message(e),207 ) from e208 209 210@router.post("/v1/messages/count_tokens")211async def count_tokens(request_data: TokenCountRequest, _auth=Depends(require_api_key)):212 """Count tokens for a request."""213 request_id = f"req_{uuid.uuid4().hex[:12]}"214 with logger.contextualize(request_id=request_id):215 try:216 tokens = get_token_count(217 request_data.messages, request_data.system, request_data.tools218 )219 logger.info(220 "COUNT_TOKENS: request_id={} model={} messages={} input_tokens={}",221 request_id,222 getattr(request_data, "model", "unknown"),223 len(request_data.messages),224 tokens,225 )226 return TokenCountResponse(input_tokens=tokens)227 except Exception as e:228 logger.error(229 "COUNT_TOKENS_ERROR: request_id={} error={}\n{}",230 request_id,231 get_user_facing_error_message(e),232 traceback.format_exc(),233 )234 raise HTTPException(235 status_code=500, detail=get_user_facing_error_message(e)236 ) from e237 238 239@router.get("/")240async def root(241 request: Request,242 settings: Settings = Depends(get_settings),243 _auth=Depends(require_api_key),244):245 """Root endpoint (JSON for API clients, HTML for browsers)."""246 payload = {247 "status": "ok",248 "provider": settings.provider_type,249 "model": settings.model,250 }251 accept = request.headers.get("accept", "")252 if "__sign" in request.query_params or "text/html" in accept.lower():253 return HTMLResponse(content=_home_page_html(payload))254 return payload255 256 257@router.get("/health")258async def health():259 """Health check endpoint."""260 return {"status": "healthy"}261 262 263@router.post("/stop")264async def stop_cli(request: Request, _auth=Depends(require_api_key)):265 """Stop all CLI sessions and pending tasks."""266 handler = getattr(request.app.state, "message_handler", None)267 if not handler:268 # Fallback if messaging not initialized269 cli_manager = getattr(request.app.state, "cli_manager", None)270 if cli_manager:271 await cli_manager.stop_all()272 logger.info("STOP_CLI: source=cli_manager cancelled_count=N/A")273 return {"status": "stopped", "source": "cli_manager"}274 raise HTTPException(status_code=503, detail="Messaging system not initialized")275 276 count = await handler.stop_all_tasks()277 logger.info("STOP_CLI: source=handler cancelled_count={}", count)278 return {"status": "stopped", "cancelled_count": count}279 280 281@router.get("/admin/factory-reset", response_class=HTMLResponse)282async def factory_reset_page(request: Request, _auth=Depends(require_api_key)):283 """Simple admin UI for one-click factory reset and restart."""284 return """285<!doctype html>286<html lang=\"en\">287<head>288 <meta charset=\"utf-8\" />289 <meta name=\"viewport\" content=\"width=device-width, initial-scale=1\" />290 <title>Factory Reset</title>291 <style>292 :root { color-scheme: dark; }293 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; }294 .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); }295 h1 { margin: 0 0 10px; font-size: 1.35rem; }296 p { margin: 0 0 18px; color: #cfd7ff; line-height: 1.45; }297 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); }298 button:disabled { opacity: 0.65; cursor: wait; }299 .status { margin-top: 14px; min-height: 24px; font-size: 0.95rem; color: #b8ffd8; }300 </style>301</head>302<body>303 <div class=\"card\">304 <h1>Factory Reset & Restart</h1>305 <p>Clears runtime cache and workspace data, then restarts this server.</p>306 <button id=\"resetBtn\">Factory Restart</button>307 <div class=\"status\" id=\"status\"></div>308 </div>309 <script>310 const btn = document.getElementById('resetBtn');311 const status = document.getElementById('status');312 btn.addEventListener('click', async () => {313 btn.disabled = true;314 status.textContent = 'Resetting cache/workspace and restarting...';315 try {316 const response = await fetch('/admin/factory-reset' + window.location.search, { method: 'POST' });317 const data = await response.json();318 if (!response.ok) {319 throw new Error(data.detail || 'Request failed');320 }321 status.textContent = 'Restart initiated. This page will disconnect briefly.';322 } catch (err) {323 status.textContent = 'Failed: ' + (err.message || String(err));324 btn.disabled = false;325 }326 });327 </script>328</body>329</html>330"""331 332 333@router.post("/admin/factory-reset")334async def factory_reset(335 background_tasks: BackgroundTasks,336 settings: Settings = Depends(get_settings),337 _auth=Depends(require_api_key),338):339 """Clear runtime state and restart process (for Space maintenance)."""340 cleared = _clear_runtime_state(settings)341 background_tasks.add_task(_restart_process)342 return {343 "status": "restarting",344 "cleared": cleared,345 }346