src/free_claude_code/runtime/application.py

"""Single owner for application startup, shutdown, and runtime operations."""

import asyncio
import inspect
import logging
import os
import traceback
from collections.abc import Awaitable, Callable, Mapping
from dataclasses import asdict, replace
from typing import Any

from loguru import logger

import free_claude_code.cli.managed as cli_managed
import free_claude_code.messaging.session as messaging_session
import free_claude_code.messaging.workflow as messaging_workflow_module
from free_claude_code.application.connected_accounts import (
    ConnectedAccountLoginMode,
    ConnectedAccountPort,
    ConnectedAccountStatus,
)
from free_claude_code.application.errors import ApplicationUnavailableError
from free_claude_code.application.model_metadata import ProviderModelRefreshResult
from free_claude_code.application.ports import StopResult
from free_claude_code.config.admin.persistence import (
    PreparedAdminUpdate,
    commit_prepared_admin_update,
    prepare_admin_update,
)
from free_claude_code.config.admin.status import provider_config_status
from free_claude_code.config.admin.values import load_value_state
from free_claude_code.config.env_files import (
    ANTHROPIC_AUTH_TOKEN_ENV,
    process_env_key_is_effective,
)
from free_claude_code.config.model_refs import parse_provider_type
from free_claude_code.config.paths import messaging_state_dir_path, usage_path
from free_claude_code.config.server_urls import local_admin_url, local_proxy_root_url
from free_claude_code.config.settings import Settings, get_settings
from free_claude_code.core.usage_tracking import get_buffer, init_buffer
from free_claude_code.messaging.platforms import factory as messaging_platform_factory
from free_claude_code.messaging.platforms.factory import MessagingPlatformOptions
from free_claude_code.messaging.platforms.ports import (
    MessagingPlatformComponents,
    MessagingRuntime,
)
from free_claude_code.messaging.voice import Transcriber

from .provider_manager import ProviderRuntimeManager

RestartCallback = Callable[[], Awaitable[None] | None]


async def best_effort(
    name: str,
    awaitable: Awaitable[Any],
    *,
    log_verbose_errors: bool = False,
) -> bool:
    """Run one cleanup step and report whether it completed.

    The lifecycle owner intentionally applies no generic timeout here. Cancelling
    an arbitrary cleanup at a deadline can abandon a half-closed SDK, thread, or
    provider resource; resource-specific cleanup or the process supervisor owns
    any force-termination deadline.
    """
    try:
        await awaitable
    except Exception as exc:
        if log_verbose_errors:
            logger.warning(
                "Shutdown step failed: {}: {}: {}",
                name,
                type(exc).__name__,
                exc,
            )
        else:
            logger.warning(
                "Shutdown step failed: {}: exc_type={}",
                name,
                type(exc).__name__,
            )
        return False
    return True


def warn_if_process_auth_token(settings: Settings) -> None:
    """Warn when server auth was implicitly inherited from the shell."""
    model_config = getattr(settings, "model_config", Settings.model_config)
    if process_env_key_is_effective(model_config, ANTHROPIC_AUTH_TOKEN_ENV):
        logger.warning(
            "ANTHROPIC_AUTH_TOKEN is set in the process environment but not in "
            "a configured .env file. The proxy will require that token. Add "
            "ANTHROPIC_AUTH_TOKEN= to .env to disable proxy auth, or set the "
            "same token in .env to make server auth explicit."
        )


def startup_failure_message(settings: Settings, exc: Exception) -> str:
    """Return the existing concise ASGI startup failure message."""
    if isinstance(exc, ApplicationUnavailableError):
        return exc.message.strip() or "Server startup failed."
    if settings.log_api_error_tracebacks:
        return f"{type(exc).__name__}: {exc}"
    return f"Server startup failed: exc_type={type(exc).__name__}"


class ApplicationRuntime:
    """Own every process-lifetime resource used by one server instance."""

    def __init__(
        self,
        provider_manager: ProviderRuntimeManager,
        *,
        transcriber: Transcriber | None,
        restart_callback: RestartCallback | None = None,
        connected_accounts: Mapping[str, ConnectedAccountPort] | None = None,
    ) -> None:
        self.provider_manager = provider_manager
        self._transcriber = transcriber
        self._restart_callback = restart_callback
        self._connected_accounts = dict(connected_accounts or {})
        self._connected_account_revisions = {
            provider_id: manager.status().revision
            for provider_id, manager in self._connected_accounts.items()
        }
        self._config_lock = asyncio.Lock()
        self._pending_fields: list[str] = []
        self._messaging_runtime: MessagingRuntime | None = None
        self._messaging_workflow: messaging_workflow_module.MessagingWorkflow | None = (
            None
        )
        self._cli_manager: cli_managed.ManagedClaudeSessionManager | None = None
        self._usage_buffer = init_buffer(usage_path())
        self._started = False
        self._closed = False
        self._provider_manager_closed = False
        self._connected_accounts_closed = False
        self._close_lock = asyncio.Lock()

    @property
    def settings(self) -> Settings:
        return self.provider_manager.current_settings()

    @property
    def is_closed(self) -> bool:
        """Whether this runtime released its complete ownership graph."""
        return self._closed

    async def start(self) -> None:
        if self._started:
            return
        logger.info("Starting Claude Code Proxy...")
        try:
            warn_if_process_auth_token(self.settings)
            await self.provider_manager.warm_referenced_model_cache()
            self.provider_manager.start_model_list_refresh()
            await self._start_messaging_if_configured()
            logging.getLogger("uvicorn.error").info(
                "Admin UI: %s (local-only)",
                local_admin_url(self.settings),
            )
            self._started = True
        except asyncio.CancelledError:
            await self.close()
            raise
        except Exception as exc:
            logger.error(
                "Startup failed:\n{}",
                startup_failure_message(self.settings, exc),
            )
            await self.close()
            raise

    async def close(self) -> bool:
        async with self._close_lock:
            if self._closed:
                return True
            logger.info("Shutdown requested, cleaning up...")
            self._closed = await self._close_owned_resources()
            if self._closed:
                self._started = False
                logger.info("Server shut down cleanly")
            else:
                logger.warning(
                    "Server shutdown incomplete; owned resources remain for retry"
                )
            return self._closed

    async def apply_admin_config(
        self,
        updates: Mapping[str, Any],
    ) -> dict[str, Any]:
        """Apply one validated config update without splitting runtime ownership."""
        async with self._config_lock:
            prepared = prepare_admin_update(updates)
            if not prepared.valid:
                return prepared.applied_response()
            assert prepared.settings is not None

            if prepared.pending_fields:
                result = self._commit_admin_update(prepared)
                restart = self._restart_metadata(
                    prepared.pending_fields,
                    prepared.settings,
                )
                result["restart"] = restart
                self._pending_fields = (
                    [] if restart["automatic"] else list(prepared.pending_fields)
                )
                return result

            result: dict[str, Any] = {}

            def commit() -> None:
                result.update(self._commit_admin_update(prepared))

            await self.provider_manager.replace(
                prepared.settings,
                commit=commit,
                reason="admin_apply",
            )
            self._pending_fields = []
            result["restart"] = self._restart_metadata((), prepared.settings)
            return result

    def admin_status(self) -> dict[str, Any]:
        settings = self.settings
        return {
            "status": "running",
            "host": settings.host,
            "port": settings.port,
            "model": settings.model,
            "provider": parse_provider_type(settings.model),
            "pending_fields": list(self._pending_fields),
            "provider_status": provider_config_status(load_value_state()),
            "cached_models": {
                provider_id: sorted(model_ids)
                for provider_id, model_ids in self.provider_manager.cached_model_ids().items()
            },
        }

    async def test_provider(self, provider_id: str) -> dict[str, Any]:
        lease = await self.provider_manager.acquire()
        try:
            provider = lease.resolve_provider(provider_id)
            infos = await provider.list_model_infos()
        except Exception as exc:
            return {
                "provider_id": provider_id,
                "ok": False,
                "error_type": type(exc).__name__,
            }
        finally:
            await lease.release()
        self.provider_manager.cache_model_infos(provider_id, infos)
        return {
            "provider_id": provider_id,
            "ok": True,
            "models": sorted(info.model_id for info in infos),
        }

    async def refresh_models(self) -> ProviderModelRefreshResult:
        return await self.provider_manager.refresh_model_list_cache()

    def usage_stats(self) -> dict[str, Any]:
        """Return usage aggregates plus the reverse-chronological request log."""
        buffer = self._usage_buffer
        if buffer is None:
            buffer = get_buffer()
        if buffer is None:
            return {
                "stats": {
                    "total_requests": 0,
                    "total_input_tokens": 0,
                    "total_output_tokens": 0,
                    "total_cache_creation_tokens": 0,
                    "total_cache_read_tokens": 0,
                    "total_reasoning_tokens": 0,
                    "errors": 0,
                    "cancelled": 0,
                    "tpm": 0.0,
                    "tps": 0.0,
                },
                "records": [],
            }
        records = buffer.query()
        return {
            "stats": buffer.stats(),
            "records": [asdict(record) for record in records],
        }

    async def connected_account_status(
        self, provider_id: str
    ) -> ConnectedAccountStatus:
        """Return safe account state and synchronize model availability."""

        manager = self._connected_account(provider_id)
        status = manager.status()
        previous_revision = self._connected_account_revisions.get(provider_id)
        if status.revision != previous_revision:
            await self.provider_manager.connected_provider_changed(
                provider_id, connected=status.connected
            )
            self._connected_account_revisions[provider_id] = status.revision
        model_count = len(self.provider_manager.cached_model_ids().get(provider_id, ()))
        return replace(status, model_count=model_count)

    async def start_connected_account_login(
        self,
        provider_id: str,
        mode: ConnectedAccountLoginMode,
    ) -> ConnectedAccountStatus:
        """Start one provider-owned interactive login."""

        return await self._connected_account(provider_id).start_login(mode)

    async def cancel_connected_account_login(
        self, provider_id: str
    ) -> ConnectedAccountStatus:
        """Cancel one pending provider login."""

        return await self._connected_account(provider_id).cancel_login()

    async def disconnect_connected_account(
        self, provider_id: str
    ) -> ConnectedAccountStatus:
        """Disconnect an account and evict only that provider's models."""

        status = await self._connected_account(provider_id).disconnect()
        await self.provider_manager.connected_provider_changed(
            provider_id, connected=False
        )
        self._connected_account_revisions[provider_id] = status.revision
        return status

    async def request_restart(self) -> None:
        callback = self._restart_callback
        if callback is None:
            return
        result = callback()
        if inspect.isawaitable(result):
            await result

    async def stop_all(self) -> StopResult | None:
        if self._messaging_workflow is not None:
            outcome = await self._messaging_workflow.stop_all_tasks()
            return StopResult(cancelled_count=outcome.cancelled_count)
        if self._cli_manager is not None:
            await self._cli_manager.stop_all()
            return StopResult(source="cli_manager")
        return None

    def _commit_admin_update(
        self,
        prepared: PreparedAdminUpdate,
    ) -> dict[str, Any]:
        result = commit_prepared_admin_update(prepared)
        get_settings.cache_clear()
        return result

    def _restart_metadata(
        self,
        fields: tuple[str, ...],
        settings: Settings,
    ) -> dict[str, Any]:
        automatic = bool(fields and self._restart_callback is not None)
        return {
            "required": bool(fields),
            "automatic": automatic,
            "admin_url": local_admin_url(settings) if automatic else None,
            "fields": list(fields),
        }

    async def _start_messaging_if_configured(self) -> None:
        try:
            components = messaging_platform_factory.create_messaging_components(
                self.settings.messaging_platform,
                self._messaging_options(),
            )
            if components is not None:
                await self._start_messaging_workflow(components)
        except ImportError as exc:
            cleaned = await self._cleanup_messaging()
            if self.settings.log_api_error_tracebacks:
                logger.warning("Messaging module import error: {}", exc)
            else:
                logger.warning(
                    "Messaging module import error: exc_type={}",
                    type(exc).__name__,
                )
            if not cleaned:
                raise RuntimeError("Messaging startup cleanup incomplete") from exc
        except Exception as exc:
            cleaned = await self._cleanup_messaging()
            if self.settings.log_api_error_tracebacks:
                logger.error("Failed to start messaging platform: {}", exc)
                logger.error(traceback.format_exc())
            else:
                logger.error(
                    "Failed to start messaging platform: exc_type={}",
                    type(exc).__name__,
                )
            if not cleaned:
                raise RuntimeError("Messaging startup cleanup incomplete") from exc

    def _messaging_options(self) -> MessagingPlatformOptions:
        settings = self.settings
        return MessagingPlatformOptions(
            telegram_bot_token=settings.telegram_bot_token,
            allowed_telegram_user_id=settings.allowed_telegram_user_id,
            telegram_proxy_url=settings.telegram_proxy_url,
            discord_bot_token=settings.discord_bot_token,
            allowed_discord_channels=settings.allowed_discord_channels,
            transcriber=self._transcriber,
            messaging_rate_limit=settings.messaging_rate_limit,
            messaging_rate_window=settings.messaging_rate_window,
            log_raw_messaging_content=settings.log_raw_messaging_content,
            log_messaging_error_details=settings.log_messaging_error_details,
            log_api_error_tracebacks=settings.log_api_error_tracebacks,
        )

    async def _start_messaging_workflow(
        self,
        components: MessagingPlatformComponents,
    ) -> None:
        settings = self.settings
        self._messaging_runtime = components.runtime
        workspace = (
            os.path.abspath(settings.allowed_dir)
            if settings.allowed_dir
            else os.getcwd()
        )
        os.makedirs(workspace, exist_ok=True)
        data_path = os.path.abspath(messaging_state_dir_path())
        os.makedirs(data_path, exist_ok=True)
        allowed_dirs = [workspace] if settings.allowed_dir else []

        self._cli_manager = cli_managed.ManagedClaudeSessionManager(
            workspace_path=workspace,
            proxy_root_url=local_proxy_root_url(settings),
            allowed_dirs=allowed_dirs,
            auth_token=settings.anthropic_auth_token,
            log_raw_cli_diagnostics=settings.log_raw_cli_diagnostics,
            log_messaging_error_details=settings.log_messaging_error_details,
        )
        session_store = messaging_session.SessionStore(
            storage_path=os.path.join(data_path, "sessions.json"),
            managed_message_cap=settings.max_message_log_entries_per_chat,
        )
        workflow = messaging_workflow_module.MessagingWorkflow(
            platform_name=components.name,
            outbound=components.outbound,
            voice_cancellation=components.voice_cancellation,
            cli_manager=self._cli_manager,
            session_store=session_store,
            debug_platform_edits=settings.debug_platform_edits,
            debug_subagent_stack=settings.debug_subagent_stack,
            log_raw_cli_diagnostics=settings.log_raw_cli_diagnostics,
            log_messaging_error_details=settings.log_messaging_error_details,
        )
        self._messaging_workflow = workflow
        workflow.restore()
        components.runtime.on_message(workflow.handle_message)
        await components.runtime.start()
        await workflow.repair_restored_statuses()
        if components.startup_notice is not None:
            await workflow.publish_startup_notice(components.startup_notice)
        logger.info("{} platform started with messaging workflow", components.name)

    async def _close_owned_resources(self) -> bool:
        if not await self._cleanup_messaging():
            return False
        if not await self._cleanup_transcriber():
            return False
        verbose = self.settings.log_api_error_tracebacks
        if not self._provider_manager_closed:
            self._provider_manager_closed = await best_effort(
                "provider_manager.close",
                self.provider_manager.close(),
                log_verbose_errors=verbose,
            )
            if not self._provider_manager_closed:
                return False
        if self._connected_accounts_closed:
            return True
        results = await asyncio.gather(
            *(
                best_effort(
                    f"connected_account.{provider_id}.close",
                    manager.close(),
                    log_verbose_errors=verbose,
                )
                for provider_id, manager in self._connected_accounts.items()
            )
        )
        self._connected_accounts_closed = all(results)
        return self._connected_accounts_closed

    def _connected_account(self, provider_id: str) -> ConnectedAccountPort:
        manager = self._connected_accounts.get(provider_id)
        if manager is None:
            raise ApplicationUnavailableError(
                f"Provider {provider_id!r} does not support connected-account login."
            )
        return manager

    async def _cleanup_messaging(self) -> bool:
        verbose = self.settings.log_api_error_tracebacks
        workflow = self._messaging_workflow
        runtime = self._messaging_runtime
        cli_manager = self._cli_manager

        if runtime is not None:
            quiesced = await best_effort(
                "messaging_runtime.quiesce",
                runtime.quiesce(),
                log_verbose_errors=verbose,
            )
            if not quiesced:
                # Delivery must remain available until ingress is known stopped.
                # Retaining the graph lets the next close retry this exact gate.
                return False

        if workflow is not None:
            closed = await best_effort(
                "messaging_workflow.close",
                workflow.close(),
                log_verbose_errors=verbose,
            )
            if not closed:
                # Active workflow tasks may still need delivery, transcription,
                # CLI sessions, and providers while a later close retries drain.
                return False
            if self._messaging_workflow is workflow:
                self._messaging_workflow = None
            if self._cli_manager is cli_manager:
                self._cli_manager = None
        elif cli_manager is not None:
            drained = await best_effort(
                "cli_manager.stop_all",
                cli_manager.stop_all(),
                log_verbose_errors=verbose,
            )
            if not drained:
                return False
            if self._cli_manager is cli_manager:
                self._cli_manager = None

        if runtime is not None:
            closed = await best_effort(
                "messaging_runtime.close",
                runtime.close(),
                log_verbose_errors=verbose,
            )
            if not closed:
                return False
            if self._messaging_runtime is runtime:
                self._messaging_runtime = None
        return True

    async def _cleanup_transcriber(self) -> bool:
        transcriber = self._transcriber
        if transcriber is None:
            return True
        closed = await best_effort(
            "transcriber.close",
            transcriber.close(),
            log_verbose_errors=self.settings.log_api_error_tracebacks,
        )
        if closed and self._transcriber is transcriber:
            self._transcriber = None
        return closed