dograh/api/services/tool_management.py

"""Service layer for reusable tool management.

Routes and MCP tools both use this module so validation, credential
scoping, MCP discovery, and analytics stay consistent.
"""

from __future__ import annotations

import asyncio
from typing import Any, Optional

from loguru import logger

from api.db import db_client
from api.db.models import UserModel
from api.enums import PostHogEvent, ToolCategory
from api.schemas.tool import (
    CreatedByResponse,
    CreateToolRequest,
    McpRefreshResponse,
    ToolResponse,
)
from api.services.posthog_client import capture_event
from api.services.workflow.mcp_tool_session import discover_mcp_tools
from api.services.workflow.tools.mcp_tool import (
    McpDefinitionError,
    validate_mcp_definition,
)


class ToolManagementError(ValueError):
    """Recoverable tool-management error with an MCP/HTTP friendly code."""

    def __init__(self, error_code: str, message: str, *, status_code: int = 400):
        super().__init__(message)
        self.error_code = error_code
        self.message = message
        self.status_code = status_code


def build_tool_response(tool: Any, include_created_by: bool = False) -> ToolResponse:
    """Build a public response from a ToolModel-like object."""
    created_by = None
    if include_created_by and tool.created_by_user:
        created_by = CreatedByResponse(
            id=tool.created_by_user.id,
            provider_id=tool.created_by_user.provider_id,
        )

    return ToolResponse(
        id=tool.id,
        tool_uuid=tool.tool_uuid,
        name=tool.name,
        description=tool.description,
        category=tool.category,
        icon=tool.icon,
        icon_color=tool.icon_color,
        status=tool.status,
        definition=tool.definition,
        created_at=tool.created_at,
        updated_at=tool.updated_at,
        created_by=created_by,
    )


def _credential_uuid_from_definition(definition: dict[str, Any]) -> Optional[str]:
    config = definition.get("config")
    if not isinstance(config, dict):
        return None
    credential_uuid = config.get("credential_uuid")
    return credential_uuid if isinstance(credential_uuid, str) else None


def _credential_uuids_from_definition(definition: dict[str, Any]) -> list[str]:
    credential_uuids: list[str] = []
    top_level = _credential_uuid_from_definition(definition)
    if top_level:
        credential_uuids.append(top_level)

    config = definition.get("config")
    if isinstance(config, dict):
        resolver = config.get("resolver")
        if isinstance(resolver, dict):
            resolver_credential_uuid = resolver.get("credential_uuid")
            if isinstance(resolver_credential_uuid, str) and resolver_credential_uuid:
                credential_uuids.append(resolver_credential_uuid)

    return list(dict.fromkeys(credential_uuids))


async def fetch_credential(credential_uuid: Optional[str], organization_id: int):
    """Best-effort credential lookup for MCP auth/discovery."""
    if not credential_uuid:
        return None
    try:
        return await db_client.get_credential_by_uuid(credential_uuid, organization_id)
    except Exception as e:  # noqa: BLE001
        logger.warning(f"Tool credential fetch failed: {e}")
        return None


async def validate_tool_credential_references(
    definition: dict[str, Any], *, organization_id: int
) -> None:
    """Ensure credential UUID references belong to the caller's organization."""
    for credential_uuid in _credential_uuids_from_definition(definition):
        credential = await db_client.get_credential_by_uuid(
            credential_uuid, organization_id
        )
        if not credential:
            raise ToolManagementError(
                "credential_not_found",
                (
                    f"Credential '{credential_uuid}' was not found in this "
                    "organization. Create it in the UI first, then retry with its "
                    "credential_uuid."
                ),
                status_code=404,
            )


async def populate_discovered_tools(
    definition: dict[str, Any], *, organization_id: int
) -> dict[str, Any]:
    """Best-effort MCP discovery before saving a tool definition.

    Non-MCP definitions pass through untouched. For MCP definitions, a dead
    server yields ``discovered_tools: []`` and does not block creation.
    """
    if not isinstance(definition, dict) or definition.get("type") != "mcp":
        return definition
    try:
        cfg = validate_mcp_definition(definition)
    except McpDefinitionError:
        return definition

    credential = await fetch_credential(cfg.get("credential_uuid"), organization_id)

    async def _run() -> list:
        try:
            return await discover_mcp_tools(
                url=cfg["url"],
                credential=credential,
                timeout_secs=cfg["timeout_secs"],
                sse_read_timeout_secs=cfg["sse_read_timeout_secs"],
            )
        except BaseException as e:  # noqa: BLE001
            logger.warning(f"MCP discovery failed; caching empty list: {e}")
            return []

    discovered = await asyncio.ensure_future(_run())
    definition["config"]["discovered_tools"] = discovered
    return definition


async def create_tool_for_user(
    request: CreateToolRequest,
    user: UserModel,
    *,
    source: str = "api",
) -> ToolResponse:
    """Create a reusable tool for the authenticated user's selected org."""
    if not user.selected_organization_id:
        raise ToolManagementError(
            "organization_required",
            "No organization selected for the user",
            status_code=400,
        )

    definition = request.definition.model_dump()
    await validate_tool_credential_references(
        definition, organization_id=user.selected_organization_id
    )
    definition = await populate_discovered_tools(
        definition,
        organization_id=user.selected_organization_id,
    )

    tool = await db_client.create_tool(
        organization_id=user.selected_organization_id,
        user_id=user.id,
        name=request.name,
        definition=definition,
        category=request.category,
        description=request.description,
        icon=request.icon,
        icon_color=request.icon_color,
    )

    capture_event(
        distinct_id=str(user.provider_id),
        event=PostHogEvent.TOOL_CREATED,
        properties={
            "tool_name": request.name,
            "tool_category": request.category,
            "source": source,
            "organization_id": user.selected_organization_id,
        },
    )

    return build_tool_response(tool)


async def refresh_mcp_tool_for_user(
    tool_uuid: str,
    user: UserModel,
) -> McpRefreshResponse:
    """Refresh cached MCP catalog for a tool owned by the user's org."""
    if not user.selected_organization_id:
        raise ToolManagementError(
            "organization_required",
            "No organization selected for the user",
            status_code=400,
        )

    tool = await db_client.get_tool_by_uuid(
        tool_uuid, user.selected_organization_id, include_archived=True
    )
    if not tool:
        raise ToolManagementError("tool_not_found", "Tool not found", status_code=404)
    if tool.category != ToolCategory.MCP.value:
        raise ToolManagementError(
            "not_mcp_tool", "Tool is not an MCP tool", status_code=400
        )

    try:
        cfg = validate_mcp_definition(tool.definition)
    except McpDefinitionError as e:
        raise ToolManagementError(
            "invalid_mcp_definition",
            f"Invalid MCP definition: {e}",
            status_code=400,
        ) from e

    credential = await fetch_credential(
        cfg.get("credential_uuid"), user.selected_organization_id
    )

    try:
        discovered = await discover_mcp_tools(
            url=cfg["url"],
            credential=credential,
            timeout_secs=cfg["timeout_secs"],
            sse_read_timeout_secs=cfg["sse_read_timeout_secs"],
        )
    except Exception as e:  # noqa: BLE001
        logger.warning(f"MCP refresh discovery failed: {e}")
        discovered = []

    if not discovered:
        error = (
            f"Could not reach the MCP server at {cfg['url']} "
            f"(or it exposes no tools). Previously cached list retained."
        )
        return McpRefreshResponse(tool_uuid=tool_uuid, discovered_tools=[], error=error)

    new_def = dict(tool.definition or {})
    new_def["config"] = {**new_def.get("config", {}), "discovered_tools": discovered}
    await db_client.update_tool(
        tool_uuid=tool_uuid,
        organization_id=user.selected_organization_id,
        definition=new_def,
    )
    return McpRefreshResponse(
        tool_uuid=tool_uuid, discovered_tools=discovered, error=None
    )