Jack1808/Claude_Code

1"""FastAPI application factory and configuration."""2 3import asyncio4import os5from contextlib import asynccontextmanager6 7from fastapi import FastAPI, HTTPException, Request8from fastapi.responses import JSONResponse9from loguru import logger10 11from config.logging_config import configure_logging12from config.settings import get_settings13from providers.exceptions import ProviderError14 15from .dependencies import cleanup_provider, validate_request_api_key16from .routes import router17 18# Opt-in to future behavior for python-telegram-bot19os.environ["PTB_TIMEDELTA"] = "1"20 21# Configure logging first (before any module logs)22_settings = get_settings()23configure_logging(_settings.log_file)24 25 26_SHUTDOWN_TIMEOUT_S = 5.027 28 29def _normalize_malformed_query_base_url_request(request: Request) -> None:30    """Normalize malformed request targets when base URL contains query auth.31 32    Some clients concatenate paths onto a base URL containing query params as plain33    strings, producing targets like:34      /?psw:token/v1/messages?beta=true35    This rewrites them to:36      /v1/messages?psw:token&beta=true37    """38    if request.scope.get("path") != "/":39        return40 41    raw_query_bytes = request.scope.get("query_string", b"")42    raw_query = raw_query_bytes.decode("utf-8", errors="ignore")43    if not raw_query or "/v1/" not in raw_query:44        return45 46    auth_part, _, remainder = raw_query.partition("/v1/")47    if not auth_part or not remainder:48        return49 50    if "?" in remainder:51        path_suffix, trailing_query = remainder.split("?", 1)52    else:53        path_suffix, trailing_query = remainder, ""54 55    new_path = f"/v1/{path_suffix}"56    new_query = auth_part if not trailing_query else f"{auth_part}&{trailing_query}"57 58    request.scope["path"] = new_path59    request.scope["raw_path"] = new_path.encode("utf-8")60    request.scope["query_string"] = new_query.encode("utf-8")61 62 63async def _best_effort(64    name: str, awaitable, timeout_s: float = _SHUTDOWN_TIMEOUT_S65) -> None:66    """Run a shutdown step with timeout; never raise to callers."""67    try:68        await asyncio.wait_for(awaitable, timeout=timeout_s)69    except TimeoutError:70        logger.warning(f"Shutdown step timed out: {name} ({timeout_s}s)")71    except Exception as e:72        logger.warning(f"Shutdown step failed: {name}: {type(e).__name__}: {e}")73 74 75@asynccontextmanager76async def lifespan(app: FastAPI):77    """Application lifespan manager."""78    settings = get_settings()79    logger.info("Starting Claude Code Proxy...")80 81    # Initialize messaging platform if configured82    messaging_platform = None83    message_handler = None84    cli_manager = None85 86    try:87        # Use the messaging factory to create the right platform88        from messaging.platforms.factory import create_messaging_platform89 90        messaging_platform = create_messaging_platform(91            platform_type=settings.messaging_platform,92            bot_token=settings.telegram_bot_token,93            allowed_user_id=settings.allowed_telegram_user_id,94            discord_bot_token=settings.discord_bot_token,95            allowed_discord_channels=settings.allowed_discord_channels,96        )97 98        if messaging_platform:99            from cli.manager import CLISessionManager100            from messaging.handler import ClaudeMessageHandler101            from messaging.session import SessionStore102 103            # Setup workspace - CLI runs in allowed_dir if set (e.g. project root)104            workspace = (105                os.path.abspath(settings.allowed_dir)106                if settings.allowed_dir107                else os.getcwd()108            )109            os.makedirs(workspace, exist_ok=True)110 111            # Session data stored in agent_workspace112            data_path = os.path.abspath(settings.claude_workspace)113            os.makedirs(data_path, exist_ok=True)114 115            api_url = f"http://{settings.host}:{settings.port}/v1"116            allowed_dirs = [workspace] if settings.allowed_dir else []117            plans_dir_abs = os.path.abspath(118                os.path.join(settings.claude_workspace, "plans")119            )120            plans_directory = os.path.relpath(plans_dir_abs, workspace)121            cli_manager = CLISessionManager(122                workspace_path=workspace,123                api_url=api_url,124                allowed_dirs=allowed_dirs,125                plans_directory=plans_directory,126            )127 128            # Initialize session store129            session_store = SessionStore(130                storage_path=os.path.join(data_path, "sessions.json")131            )132 133            # Create and register message handler134            message_handler = ClaudeMessageHandler(135                platform=messaging_platform,136                cli_manager=cli_manager,137                session_store=session_store,138            )139 140            # Restore tree state if available141            saved_trees = session_store.get_all_trees()142            if saved_trees:143                logger.info(f"Restoring {len(saved_trees)} conversation trees...")144                from messaging.trees.queue_manager import TreeQueueManager145 146                message_handler.replace_tree_queue(147                    TreeQueueManager.from_dict(148                        {149                            "trees": saved_trees,150                            "node_to_tree": session_store.get_node_mapping(),151                        },152                        queue_update_callback=message_handler.update_queue_positions,153                        node_started_callback=message_handler.mark_node_processing,154                    )155                )156                # Reconcile restored state - anything PENDING/IN_PROGRESS is lost across restart157                if message_handler.tree_queue.cleanup_stale_nodes() > 0:158                    # Sync back and save159                    tree_data = message_handler.tree_queue.to_dict()160                    session_store.sync_from_tree_data(161                        tree_data["trees"], tree_data["node_to_tree"]162                    )163 164            # Wire up the handler165            messaging_platform.on_message(message_handler.handle_message)166 167            # Start the platform168            await messaging_platform.start()169            logger.info(170                f"{messaging_platform.name} platform started with message handler"171            )172 173    except ImportError as e:174        logger.warning(f"Messaging module import error: {e}")175    except Exception as e:176        logger.error(f"Failed to start messaging platform: {e}")177        import traceback178 179        logger.error(traceback.format_exc())180 181    # Store in app state for access in routes182    app.state.messaging_platform = messaging_platform183    app.state.message_handler = message_handler184    app.state.cli_manager = cli_manager185 186    yield187 188    # Cleanup189    if message_handler and hasattr(message_handler, "session_store"):190        try:191            message_handler.session_store.flush_pending_save()192        except Exception as e:193            logger.warning(f"Session store flush on shutdown: {e}")194    logger.info("Shutdown requested, cleaning up...")195    if messaging_platform:196        await _best_effort("messaging_platform.stop", messaging_platform.stop())197    if cli_manager:198        await _best_effort("cli_manager.stop_all", cli_manager.stop_all())199    await _best_effort("cleanup_provider", cleanup_provider())200 201    # Ensure background limiter worker doesn't keep the loop alive.202    try:203        from messaging.limiter import MessagingRateLimiter204 205        await _best_effort(206            "MessagingRateLimiter.shutdown_instance",207            MessagingRateLimiter.shutdown_instance(),208            timeout_s=2.0,209        )210    except Exception:211        # Limiter may never have been imported/initialized.212        pass213 214    logger.info("Server shut down cleanly")215 216 217def create_app() -> FastAPI:218    """Create and configure the FastAPI application."""219    app = FastAPI(220        title="Claude Code Proxy",221        version="2.0.0",222        lifespan=lifespan,223    )224 225    @app.middleware("http")226    async def enforce_api_key(request: Request, call_next):227        """Enforce API key for every request before routing/method matching."""228        _normalize_malformed_query_base_url_request(request)229        try:230            validate_request_api_key(request, get_settings())231        except HTTPException as exc:232            return JSONResponse(233                status_code=exc.status_code,234                content={"detail": exc.detail},235            )236        return await call_next(request)237 238    # Register routes239    app.include_router(router)240 241    # Exception handlers242    @app.exception_handler(ProviderError)243    async def provider_error_handler(request: Request, exc: ProviderError):244        """Handle provider-specific errors and return Anthropic format."""245        logger.error(f"Provider Error: {exc.error_type} - {exc.message}")246        return JSONResponse(247            status_code=exc.status_code,248            content=exc.to_anthropic_format(),249        )250 251    @app.exception_handler(Exception)252    async def general_error_handler(request: Request, exc: Exception):253        """Handle general errors and return Anthropic format."""254        logger.error(f"General Error: {exc!s}")255        import traceback256 257        logger.error(traceback.format_exc())258        return JSONResponse(259            status_code=500,260            content={261                "type": "error",262                "error": {263                    "type": "api_error",264                    "message": "An unexpected error occurred.",265                },266            },267        )268 269    return app270 271 272# Default app instance for uvicorn273app = create_app()274