dashboard-nanobot/backend/services/bot_management_service.py

import os
import re
import shutil
from typing import Any, Dict, List, Optional

import httpx
from fastapi import HTTPException
from sqlmodel import Session, select

from core.cache import cache
from core.docker_instance import docker_manager
from core.settings import BOTS_WORKSPACE_ROOT
from models.bot import BotInstance, NanobotImage
from schemas.bot import BotCreateRequest, BotUpdateRequest
from services.bot_service import (
    _normalize_env_params,
    _normalize_initial_channels,
    _normalize_resource_limits,
    _normalize_system_timezone,
    _provider_defaults,
    _resolve_bot_env_params,
    _serialize_bot,
    _serialize_bot_list_item,
    _sync_workspace_channels,
)
from services.bot_storage_service import _write_env_store
from services.cache_service import _cache_key_bot_detail, _cache_key_bots_list, _invalidate_bot_detail_cache
from services.platform_service import record_activity_event
from services.template_service import get_agent_md_templates

BOT_ID_PATTERN = re.compile(r"^[A-Za-z0-9_]+$")
MANAGED_WORKSPACE_FILENAMES = ("AGENTS.md", "SOUL.md", "USER.md", "TOOLS.md", "IDENTITY.md")


def _managed_bot_file_paths(bot_id: str) -> Dict[str, str]:
    bot_root = os.path.join(BOTS_WORKSPACE_ROOT, bot_id, ".nanobot")
    workspace_root = os.path.join(bot_root, "workspace")
    paths = {
        "config": os.path.join(bot_root, "config.json"),
        "env": os.path.join(bot_root, "env.json"),
        "resources": os.path.join(bot_root, "resources.json"),
    }
    for filename in MANAGED_WORKSPACE_FILENAMES:
        paths[f"workspace:{filename}"] = os.path.join(workspace_root, filename)
    return paths


def _snapshot_managed_bot_files(bot_id: str) -> Dict[str, Optional[bytes]]:
    snapshot: Dict[str, Optional[bytes]] = {}
    for key, path in _managed_bot_file_paths(bot_id).items():
        if os.path.isfile(path):
            with open(path, "rb") as file:
                snapshot[key] = file.read()
        else:
            snapshot[key] = None
    return snapshot


def _restore_managed_bot_files(bot_id: str, snapshot: Dict[str, Optional[bytes]]) -> None:
    for key, path in _managed_bot_file_paths(bot_id).items():
        payload = snapshot.get(key)
        if payload is None:
            if os.path.exists(path):
                os.remove(path)
            continue
        os.makedirs(os.path.dirname(path), exist_ok=True)
        tmp_path = f"{path}.tmp"
        with open(tmp_path, "wb") as file:
            file.write(payload)
        os.replace(tmp_path, path)


def _cleanup_bot_workspace_root(bot_id: str) -> None:
    bot_root = os.path.join(BOTS_WORKSPACE_ROOT, bot_id)
    if os.path.isdir(bot_root):
        shutil.rmtree(bot_root, ignore_errors=True)


async def test_provider_connection(payload: Dict[str, Any]) -> Dict[str, Any]:
    provider = (payload.get("provider") or "").strip()
    api_key = (payload.get("api_key") or "").strip()
    model = (payload.get("model") or "").strip()
    api_base = (payload.get("api_base") or "").strip()

    if not provider or not api_key:
        raise HTTPException(status_code=400, detail="provider and api_key are required")

    normalized_provider, default_base = _provider_defaults(provider)
    base = (api_base or default_base).rstrip("/")
    if normalized_provider not in {"openrouter", "dashscope", "kimi", "minimax", "openai", "deepseek"}:
        raise HTTPException(status_code=400, detail=f"provider not supported for test: {provider}")
    if not base:
        raise HTTPException(status_code=400, detail=f"api_base is required for provider: {provider}")

    headers = {"Authorization": f"Bearer {api_key}"}
    timeout = httpx.Timeout(20.0, connect=10.0)
    url = f"{base}/models"
    try:
        async with httpx.AsyncClient(timeout=timeout) as client:
            response = await client.get(url, headers=headers)
        if response.status_code >= 400:
            return {
                "ok": False,
                "provider": normalized_provider,
                "status_code": response.status_code,
                "detail": response.text[:500],
            }
        data = response.json()
        models_raw = data.get("data", []) if isinstance(data, dict) else []
        model_ids: List[str] = [
            str(item["id"]) for item in models_raw[:20] if isinstance(item, dict) and item.get("id")
        ]
        return {
            "ok": True,
            "provider": normalized_provider,
            "endpoint": url,
            "models_preview": model_ids[:8],
            "model_hint": (
                "model_found"
                if model and any(model in item for item in model_ids)
                else ("model_not_listed" if model else "")
            ),
        }
    except Exception as exc:
        return {
            "ok": False,
            "provider": normalized_provider,
            "endpoint": url,
            "detail": str(exc),
        }


def create_bot_record(session: Session, *, payload: BotCreateRequest) -> Dict[str, Any]:
    normalized_bot_id = str(payload.id or "").strip()
    if not normalized_bot_id:
        raise HTTPException(status_code=400, detail="Bot ID is required")
    if not BOT_ID_PATTERN.fullmatch(normalized_bot_id):
        raise HTTPException(status_code=400, detail="Bot ID can only contain letters, numbers, and underscores")
    if session.get(BotInstance, normalized_bot_id):
        raise HTTPException(status_code=409, detail=f"Bot ID already exists: {normalized_bot_id}")

    image_row = session.get(NanobotImage, payload.image_tag)
    if not image_row:
        raise HTTPException(status_code=400, detail=f"Image not registered in DB: {payload.image_tag}")
    if image_row.status != "READY":
        raise HTTPException(status_code=400, detail=f"Image status is not READY: {payload.image_tag} ({image_row.status})")
    if not docker_manager.has_image(payload.image_tag):
        raise HTTPException(status_code=400, detail=f"Docker image not found locally: {payload.image_tag}")

    normalized_env_params = _normalize_env_params(payload.env_params)
    try:
        normalized_env_params["TZ"] = _normalize_system_timezone(payload.system_timezone)
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc

    bot = BotInstance(
        id=normalized_bot_id,
        name=payload.name,
        enabled=bool(payload.enabled) if payload.enabled is not None else True,
        access_password=str(payload.access_password or ""),
        image_tag=payload.image_tag,
        workspace_dir=os.path.join(BOTS_WORKSPACE_ROOT, normalized_bot_id),
    )
    template_defaults = get_agent_md_templates()
    resource_limits = _normalize_resource_limits(payload.cpu_cores, payload.memory_mb, payload.storage_gb)
    try:
        session.add(bot)
        session.flush()
        _write_env_store(normalized_bot_id, normalized_env_params)
        _sync_workspace_channels(
            session,
            normalized_bot_id,
            channels_override=_normalize_initial_channels(normalized_bot_id, payload.channels),
            global_delivery_override={
                "sendProgress": bool(payload.send_progress) if payload.send_progress is not None else False,
                "sendToolHints": bool(payload.send_tool_hints) if payload.send_tool_hints is not None else False,
            },
            runtime_overrides={
                "llm_provider": payload.llm_provider,
                "llm_model": payload.llm_model,
                "api_key": payload.api_key,
                "api_base": payload.api_base or "",
                "temperature": payload.temperature,
                "top_p": payload.top_p,
                "max_tokens": payload.max_tokens,
                "cpu_cores": resource_limits["cpu_cores"],
                "memory_mb": resource_limits["memory_mb"],
                "storage_gb": resource_limits["storage_gb"],
                "system_prompt": payload.system_prompt or payload.soul_md or template_defaults.get("soul_md", ""),
                "soul_md": payload.soul_md or payload.system_prompt or template_defaults.get("soul_md", ""),
                "agents_md": payload.agents_md or template_defaults.get("agents_md", ""),
                "user_md": payload.user_md or template_defaults.get("user_md", ""),
                "tools_md": payload.tools_md or template_defaults.get("tools_md", ""),
                "identity_md": payload.identity_md or template_defaults.get("identity_md", ""),
                "send_progress": bool(payload.send_progress) if payload.send_progress is not None else False,
                "send_tool_hints": bool(payload.send_tool_hints) if payload.send_tool_hints is not None else False,
            },
        )
        record_activity_event(
            session,
            normalized_bot_id,
            "bot_created",
            channel="system",
            detail=f"Bot {normalized_bot_id} created",
            metadata={"image_tag": payload.image_tag},
        )
        session.commit()
        session.refresh(bot)
    except Exception:
        session.rollback()
        _cleanup_bot_workspace_root(normalized_bot_id)
        raise
    _invalidate_bot_detail_cache(normalized_bot_id)
    return _serialize_bot(bot)


def list_bots_with_cache(session: Session) -> List[Dict[str, Any]]:
    cached = cache.get_json(_cache_key_bots_list())
    if isinstance(cached, list):
        return cached
    bots = session.exec(
        select(BotInstance).order_by(BotInstance.created_at.desc(), BotInstance.id.asc())
    ).all()
    dirty = False
    for bot in bots:
        actual_status = docker_manager.get_bot_status(bot.id)
        if bot.docker_status != actual_status:
            bot.docker_status = actual_status
            if actual_status != "RUNNING" and str(bot.current_state or "").upper() not in {"ERROR"}:
                bot.current_state = "IDLE"
            session.add(bot)
            dirty = True
    if dirty:
        session.commit()
        for bot in bots:
            session.refresh(bot)
    rows = [_serialize_bot_list_item(bot) for bot in bots]
    cache.set_json(_cache_key_bots_list(), rows, ttl=30)
    return rows


def get_bot_detail_cached(session: Session, *, bot_id: str) -> Dict[str, Any]:
    cached = cache.get_json(_cache_key_bot_detail(bot_id))
    if isinstance(cached, dict):
        return cached
    bot = session.get(BotInstance, bot_id)
    if not bot:
        raise HTTPException(status_code=404, detail="Bot not found")
    row = _serialize_bot(bot)
    cache.set_json(_cache_key_bot_detail(bot_id), row, ttl=30)
    return row


def authenticate_bot_page_access(session: Session, *, bot_id: str, password: str) -> Dict[str, Any]:
    bot = session.get(BotInstance, bot_id)
    if not bot:
        raise HTTPException(status_code=404, detail="Bot not found")

    configured = str(bot.access_password or "").strip()
    if not configured:
        return {"ok": True, "enabled": False, "bot_id": bot_id}

    candidate = str(password or "").strip()
    if not candidate:
        raise HTTPException(status_code=401, detail="Bot access password required")
    if candidate != configured:
        raise HTTPException(status_code=401, detail="Invalid bot access password")
    return {"ok": True, "enabled": True, "bot_id": bot_id}


def update_bot_record(session: Session, *, bot_id: str, payload: BotUpdateRequest) -> Dict[str, Any]:
    bot = session.get(BotInstance, bot_id)
    if not bot:
        raise HTTPException(status_code=404, detail="Bot not found")
    managed_file_snapshot = _snapshot_managed_bot_files(bot_id)

    update_data = payload.model_dump(exclude_unset=True)
    if "image_tag" in update_data and update_data["image_tag"]:
        image_tag = str(update_data["image_tag"]).strip()
        image_row = session.get(NanobotImage, image_tag)
        if not image_row:
            raise HTTPException(status_code=400, detail=f"Image not registered in DB: {image_tag}")
        if image_row.status != "READY":
            raise HTTPException(status_code=400, detail=f"Image status is not READY: {image_tag} ({image_row.status})")
        if not docker_manager.has_image(image_tag):
            raise HTTPException(status_code=400, detail=f"Docker image not found locally: {image_tag}")

    env_params = update_data.pop("env_params", None) if isinstance(update_data, dict) else None
    system_timezone = update_data.pop("system_timezone", None) if isinstance(update_data, dict) else None
    normalized_system_timezone: Optional[str] = None
    if system_timezone is not None:
        try:
            normalized_system_timezone = _normalize_system_timezone(system_timezone)
        except ValueError as exc:
            raise HTTPException(status_code=400, detail=str(exc)) from exc

    runtime_fields = {
        "llm_provider",
        "llm_model",
        "api_key",
        "api_base",
        "temperature",
        "top_p",
        "max_tokens",
        "cpu_cores",
        "memory_mb",
        "storage_gb",
        "soul_md",
        "agents_md",
        "user_md",
        "tools_md",
        "identity_md",
        "send_progress",
        "send_tool_hints",
        "system_prompt",
    }
    runtime_overrides: Dict[str, Any] = {}
    update_data.pop("tools_config", None) if isinstance(update_data, dict) else None
    for field in runtime_fields:
        if field in update_data:
            runtime_overrides[field] = update_data.pop(field)

    for text_field in ("llm_provider", "llm_model", "api_key"):
        if text_field in runtime_overrides:
            text = str(runtime_overrides.get(text_field) or "").strip()
            if not text:
                runtime_overrides.pop(text_field, None)
            else:
                runtime_overrides[text_field] = text
    if "api_base" in runtime_overrides:
        runtime_overrides["api_base"] = str(runtime_overrides.get("api_base") or "").strip()
    if "system_prompt" in runtime_overrides and "soul_md" not in runtime_overrides:
        runtime_overrides["soul_md"] = runtime_overrides["system_prompt"]
    if "soul_md" in runtime_overrides and "system_prompt" not in runtime_overrides:
        runtime_overrides["system_prompt"] = runtime_overrides["soul_md"]
    if {"cpu_cores", "memory_mb", "storage_gb"} & set(runtime_overrides.keys()):
        runtime_overrides.update(
            _normalize_resource_limits(
                runtime_overrides.get("cpu_cores"),
                runtime_overrides.get("memory_mb"),
                runtime_overrides.get("storage_gb"),
            )
        )

    for key, value in update_data.items():
        if key in {"name", "image_tag", "access_password", "enabled"}:
            setattr(bot, key, value)
    try:
        session.add(bot)
        session.flush()

        if env_params is not None or normalized_system_timezone is not None:
            next_env_params = _resolve_bot_env_params(bot_id)
            if env_params is not None:
                next_env_params = _normalize_env_params(env_params)
            if normalized_system_timezone is not None:
                next_env_params["TZ"] = normalized_system_timezone
            _write_env_store(bot_id, next_env_params)

        global_delivery_override: Optional[Dict[str, Any]] = None
        if "send_progress" in runtime_overrides or "send_tool_hints" in runtime_overrides:
            global_delivery_override = {}
            if "send_progress" in runtime_overrides:
                global_delivery_override["sendProgress"] = bool(runtime_overrides.get("send_progress"))
            if "send_tool_hints" in runtime_overrides:
                global_delivery_override["sendToolHints"] = bool(runtime_overrides.get("send_tool_hints"))

        _sync_workspace_channels(
            session,
            bot_id,
            runtime_overrides=runtime_overrides if runtime_overrides else None,
            global_delivery_override=global_delivery_override,
        )
        session.commit()
        session.refresh(bot)
    except Exception:
        session.rollback()
        _restore_managed_bot_files(bot_id, managed_file_snapshot)
        refreshed_bot = session.get(BotInstance, bot_id)
        if refreshed_bot:
            session.refresh(refreshed_bot)
            bot = refreshed_bot
        raise
    _invalidate_bot_detail_cache(bot_id)
    return _serialize_bot(bot)