smart-org/apps/api/app/interfaces/mcp/gateway.py

"""The MCP tools themselves — framework-free, so they can be called directly.

``app.interfaces.mcp.server`` adapts these to the official SDK; this module
holds the actual work and is what the tests exercise. Four rules are enforced
here, once, for every tool:

1. **The organization comes from the token, never from the arguments.** A tool
   body receives an authenticated ``TokenIdentity`` and uses its
   ``organization_id``; an ``organization_id`` in the arguments is ignored, so an
   agent cannot ask about a tenant it was not issued for.
2. **A scope is checked before the work starts.** Not after.
3. **Every call is audited** (``mcp.tool_call`` with the token prefix, the tool,
   and whether it succeeded), rate-limited, and its output redacted.
4. **An error is an answer, not a crash**: the agent gets a plain message, and
   the exception never reaches the transport.
"""

from __future__ import annotations

import json
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from typing import Any

from app.application.container_org import (
    build_memory_service,
    build_model_grant_service,
)
from app.application.services.mcp_service import McpService, TokenIdentity
from app.core.config import settings
from app.core.errors import NotFoundError, PermissionDeniedError, ValidationError
from app.core.redact import redact
from app.domain.memory.models import MemoryKind
from app.infrastructure.error_log import record_error


@dataclass(frozen=True, slots=True)
class ToolSpec:
    """One tool: its name, the scope it needs, and its JSON schema."""
    name: str
    scope: str
    description: str
    parameters: dict[str, Any]


TOOL_SPECS: tuple[ToolSpec, ...] = (
    ToolSpec(
        name="app_get_profile",
        scope="app:read",
        description="پروفایل اپلیکیشن این سازمان: نام، نوع، زبان، فریم‌ورک، "
        "دستورهای ساخت و تست، پورت و مسیر سلامت.",
        parameters={
            "type": "object",
            "properties": {},
            "additionalProperties": False,
        },
    ),
    ToolSpec(
        name="tasks_list",
        scope="tasks:read",
        description="فهرست کارهای این سازمان، به ترتیب جدیدترین.",
        parameters={
            "type": "object",
            "properties": {
                "status": {
                    "type": "string",
                    "description": "وضعیت اختیاری برای فیلتر (new، running، done، …)",
                },
                "limit": {"type": "integer", "minimum": 1, "maximum": 100, "default": 25},
            },
            "additionalProperties": False,
        },
    ),
    ToolSpec(
        name="tasks_get",
        scope="tasks:read",
        description="یک کار بر اساس شناسهٔ آن، در همین سازمان.",
        parameters={
            "type": "object",
            "properties": {
                "task_id": {"type": "string", "description": "شناسهٔ کار"},
            },
            "required": ["task_id"],
            "additionalProperties": False,
        },
    ),
    ToolSpec(
        name="tasks_create",
        scope="tasks:write",
        description="ثبت یک کار تازه برای مدیر سازمان؛ مدیر خودش آن را تقسیم "
        "و به کارمندان می‌سپارد.",
        parameters={
            "type": "object",
            "properties": {
                "title": {"type": "string", "description": "عنوان کوتاه کار"},
                "description": {"type": "string", "description": "شرح کامل درخواست"},
            },
            "required": ["title"],
            "additionalProperties": False,
        },
    ),
    ToolSpec(
        name="memory_search",
        scope="memory:read",
        description="جست‌وجو در حافظهٔ سازمان: قوانین، ترجیحات، درس‌ها و نقشهٔ کد.",
        parameters={
            "type": "object",
            "properties": {
                "query": {"type": "string", "description": "واژه‌های جست‌وجو"},
                "kind": {
                    "type": "string",
                    "description": "نوع اختیاری: rule، preference، lesson، fact، code_map",
                },
                "limit": {"type": "integer", "minimum": 1, "maximum": 50, "default": 10},
            },
            "additionalProperties": False,
        },
    ),
    ToolSpec(
        name="memory_write",
        scope="memory:write",
        description="یادآوری یک واقعیت یا ترجیح به سازمان. گذرواژه و کلید "
        "هرگز ذخیره نمی‌شود.",
        parameters={
            "type": "object",
            "properties": {
                "kind": {
                    "type": "string",
                    "enum": ["fact", "preference"],
                    "description": "فقط «fact» یا «preference»",
                },
                "content": {"type": "string", "description": "یک جملهٔ کوتاه"},
            },
            "required": ["kind", "content"],
            "additionalProperties": False,
        },
    ),
    ToolSpec(
        name="knowledge_search",
        scope="memory:read",
        description="جست‌وجو در منابع دانش سازمان (اسناد بلند).",
        parameters={
            "type": "object",
            "properties": {
                "query": {"type": "string", "description": "واژه‌های جست‌وجو"},
                "limit": {"type": "integer", "minimum": 1, "maximum": 20, "default": 5},
            },
            "required": ["query"],
            "additionalProperties": False,
        },
    ),
    ToolSpec(
        name="pipeline_status",
        scope="delivery:read",
        description="وضعیت خط تحویل (استقرار و دمو). در این نسخه در دسترس نیست.",
        parameters={
            "type": "object",
            "properties": {},
            "additionalProperties": False,
        },
    ),
)

#: The only kinds an external agent may write. A rule is the organization's own
#: to set through the UI, and a lesson is the engine's to learn.
MCP_WRITABLE_KINDS = frozenset({MemoryKind.FACT, MemoryKind.PREFERENCE})


class McpError(Exception):
    """A refusal the agent can read. Never carries a secret."""

    def __init__(self, message: str, *, code: str = "mcp_error") -> None:
        super().__init__(message)
        self.code = code


@dataclass(slots=True)
class McpCallResult:
    """What one tool call produced: a redacted payload, or a readable refusal."""

    ok: bool
    payload: dict[str, Any]
    error: str | None = None
    code: str = "mcp_error"

    def as_dict(self) -> dict[str, Any]:
        if self.ok:
            return self.payload
        return {
            "error": {"code": self.code, "message": self.error or "عملیات ناموفق بود."}
        }

    def to_json(self) -> str:
        """The JSON text an agent receives. Always redacted, always short."""
        return json.dumps(redact(self.as_dict()), ensure_ascii=False, default=str)


class McpGateway:
    """Executes the tools for one authenticated token.

    ``session_factory`` is injectable so the tests can point the gateway at their
    in-memory database; production passes nothing and gets the application's own
    factory, exactly like the workers.
    """

    def __init__(self, *, session_factory=None, rate_limiter=None) -> None:
        from app.application.services.mcp_service import TokenRateLimiter

        self._session_factory = session_factory
        self._limiter = rate_limiter or TokenRateLimiter()

    # -- public API ----------------------------------------------------------

    async def call(
        self,
        *,
        tool: str,
        arguments: dict[str, Any] | None = None,
        token: str,
        identity: TokenIdentity | None = None,
    ) -> McpCallResult:
        """Run one tool for one token. Every failure mode is an answer."""
        spec = _SPECS_BY_NAME.get(tool)
        if spec is None:
            return McpCallResult(ok=False, payload={}, error=f"ابزار ناشناخته: {tool}")
        args = arguments or {}
        if not settings.mcp_gateway_enabled:
            # The owner's kill switch. Nothing is attributable yet, so this lands
            # in the error ledger rather than in one organization's audit trail.
            await _note("mcp.gateway_disabled", f"tool={spec.name}")
            return self._fail(
                "دسترسی MCP از سوی مالک سامانه غیرفعال شده است.",
                code="mcp_disabled",
            )

        factory = self._session_factory or _default_factory()
        async with factory() as session:
            service = McpService(session)
            identity = identity or await service.verify(token)
            if identity is None:
                # A token that does not exist, was revoked, or expired is one
                # answer — and it is unattributable, so it goes to the error
                # ledger where a human can see the pattern.
                await _note("mcp.invalid_token", f"tool={spec.name}")
                return self._fail(
                    "توکن نامعتبر یا منقضی است.", code="invalid_token"
                )

            async def refuse(message: str, code: str) -> McpCallResult:
                """Every refusal from here on is audited against this organization.

                The audit row is committed inside the refusal: a denial nobody can
                see afterwards is not an audit trail, it is a hope.
                """
                await self._audit(
                    session, identity=identity, spec=spec, ok=False, error=message
                )
                await session.commit()
                return McpCallResult(ok=False, payload={}, error=message, code=code)

            if not identity.allows(spec.scope):
                return await refuse(
                    "این توکن به این ابزار دسترسی ندارد.", "scope_denied"
                )
            policy = await build_model_grant_service(session).get_org_policy(
                identity.organization_id
            )
            if not policy.mcp_enabled:
                return await refuse(
                    "دسترسی MCP این سازمان غیرفعال است.", "org_disabled"
                )
            if not self._limiter.allow(identity.token_id):
                return await refuse(
                    "تعداد درخواست‌ها بیش از حد مجاز است؛ کمی بعد تلاش کنید.",
                    "rate_limited",
                )

            handler = _HANDLERS[spec.name]
            try:
                payload = await handler(self, session=session, identity=identity, **args)
            except (NotFoundError, PermissionDeniedError, ValidationError, McpError) as exc:
                message = str(getattr(exc, "message", exc))
                await self._audit(
                    session, identity=identity, spec=spec, ok=False, error=message
                )
                await session.commit()
                return McpCallResult(ok=False, payload={}, error=message, code="tool_failed")
            except TypeError:
                # An argument shape the tool did not expect: the agent's fault,
                # not the platform's.
                return await refuse(
                    "آرگومان‌های ابزار نامعتبر است.", "bad_arguments"
                )
            except Exception as exc:  # noqa: BLE001 - an outage is an answer too
                # Anything else is infrastructure, and its message may contain a
                # DSN or a heap address. The agent gets a plain sentence; the
                # operator gets the real error in the error console.
                await record_error(
                    kind="mcp.tool_crashed",
                    message=str(exc),
                    exc=exc,
                    organization_id=identity.organization_id,
                    metadata={"tool": spec.name, "prefix": identity.prefix},
                )
                await self._audit(
                    session, identity=identity, spec=spec, ok=False, error="internal"
                )
                await session.commit()
                return self._fail(
                    "خطای داخلی سامانه؛ این خطا ثبت شد.", code="internal_error"
                )
            await service.touch(identity.token_id)
            await self._audit(session, identity=identity, spec=spec, ok=True)
            await session.commit()
            return McpCallResult(ok=True, payload=payload)

    # -- helpers -------------------------------------------------------------

    @staticmethod
    def _fail(message: str, *, code: str) -> McpCallResult:
        """A refusal is a normal answer: the agent gets a message it can act on."""
        return McpCallResult(ok=False, payload={}, error=message, code=code)

    async def _audit(
        self,
        session,
        *,
        identity: TokenIdentity,
        spec: ToolSpec,
        ok: bool,
        error: str = "",
    ) -> None:
        """One audit row per call. Never allowed to fail the call itself."""
        from app.application.services.audit_service import record_audit

        try:
            await record_audit(
                session,
                action="mcp.tool_call",
                organization_id=identity.organization_id,
                actor_user_id=None,
                target_type="mcp_tool",
                target_id=spec.name,
                metadata={
                    "prefix": identity.prefix,
                    "scope": spec.scope,
                    "ok": ok,
                    "error": error[:200],
                },
                ip="mcp",
                user_agent=f"mcp/{identity.prefix}",
            )
        except Exception:  # noqa: BLE001 - auditing must not fail a tool call
            return


# ---------------------------------------------------------------------------
# Tool bodies
# ---------------------------------------------------------------------------


async def _app_profile(
    gateway: McpGateway, *, session, identity: TokenIdentity, **_ignored: Any
) -> dict[str, Any]:
    from app.infrastructure.repositories.app_profile import AppProfileRepository
    from app.infrastructure.repositories.projects import ProjectRepository

    profile = await AppProfileRepository(session).get_for_org(identity.organization_id)
    if profile is None:
        raise NotFoundError("اپلیکیشنی برای این سازمان یافت نشد.")
    project = await ProjectRepository(session).get(profile.project_id)
    return {
        "project_name": project.name if project else "",
        "project_slug": project.slug if project else "",
        "app_type": profile.app_type,
        "language": profile.language,
        "framework": profile.framework,
        "runtime_version": profile.runtime_version,
        "package_manager": profile.package_manager,
        "setup_command": profile.setup_command,
        "build_command": profile.build_command,
        "test_command": profile.test_command,
        "start_command": profile.start_command,
        "app_port": profile.app_port,
        "healthcheck_path": profile.healthcheck_path,
        "default_branch": profile.default_branch,
        "git_source": profile.git_source,
        "repo_url": profile.external_repo_url,
    }


async def _tasks_list(
    gateway: McpGateway,
    *,
    session,
    identity: TokenIdentity,
    status: str | None = None,
    limit: int = 25,
    **_ignored: Any,
) -> dict[str, Any]:
    from app.application.container import build_workforce_service
    from app.domain.workforce.models import TaskStatus

    service = build_workforce_service(session)
    parsed: TaskStatus | None = None
    if status:
        try:
            parsed = TaskStatus(str(status).strip().lower())
        except ValueError as exc:
            raise ValidationError("وضعیت کار نامعتبر است.") from exc
    tasks = await service.list_tasks(
        identity.organization_id, status=parsed, limit=max(1, min(int(limit), 100))
    )
    return {
        "tasks": [
            {
                "id": task.id,
                "title": task.title,
                "status": str(task.status),
                "created_at": task.created_at.isoformat() if task.created_at else None,
                "assigned_employee_id": task.assigned_employee_id,
            }
            for task in tasks
        ]
    }


async def _tasks_get(
    gateway: McpGateway, *, session, identity: TokenIdentity, task_id: str, **_ignored: Any
) -> dict[str, Any]:
    from app.application.container import build_workforce_service

    if not task_id:
        raise ValidationError("شناسهٔ کار لازم است.")
    service = build_workforce_service(session)
    task = await service.get_task_or_404(
        organization_id=identity.organization_id, task_id=task_id
    )
    return {
        "task": {
            "id": task.id,
            "title": task.title,
            "description": task.description,
            "status": str(task.status),
            "created_at": task.created_at.isoformat() if task.created_at else None,
            "project_id": task.project_id,
        }
    }


async def _tasks_create(
    gateway: McpGateway,
    *,
    session,
    identity: TokenIdentity,
    title: str,
    description: str = "",
    **_ignored: Any,
) -> dict[str, Any]:
    """A user order to the manager: the current public signature, untouched."""
    from app.application.container import build_workforce_service

    if not title or not title.strip():
        raise ValidationError("عنوان کار لازم است.")
    creator = await McpService(session).creator_of(identity.token_id)
    service = build_workforce_service(session)
    task = await service.create_task(
        organization_id=identity.organization_id,
        created_by_user_id=creator or "",
        title=title,
        description=description,
    )
    return {"task": {"id": task.id, "title": task.title, "status": str(task.status)}}


async def _memory_search(
    gateway: McpGateway,
    *,
    session,
    identity: TokenIdentity,
    query: str = "",
    kind: str | None = None,
    limit: int = 10,
    **_ignored: Any,
) -> dict[str, Any]:
    service = build_memory_service(session)
    items = await service.list(
        identity.organization_id, kind=kind, search=query or None
    )
    return {
        "memories": [
            {
                "id": memory.id,
                "kind": str(memory.kind),
                "key": memory.key,
                "content": memory.content,
                "pinned": memory.pinned,
                "confidence": memory.confidence,
            }
            for memory in items[: max(1, min(int(limit), 50))]
        ]
    }


async def _memory_write(
    gateway: McpGateway,
    *,
    session,
    identity: TokenIdentity,
    kind: str,
    content: str,
    **_ignored: Any,
) -> dict[str, Any]:
    from app.domain.memory.models import coerce_kind

    entry_kind = coerce_kind(kind)
    if entry_kind not in MCP_WRITABLE_KINDS:
        raise McpError("با توکن فقط «fact» یا «preference» قابل نوشتن است.")
    service = build_memory_service(session)
    memory_id = await service.record(
        organization_id=identity.organization_id,
        kind=str(entry_kind),
        content=content,
        source="user",
    )
    if memory_id is None:  # pragma: no cover - the task cap cannot apply here
        raise McpError("یادداشت ذخیره نشد.")
    return {"memory": {"id": memory_id, "kind": str(entry_kind)}}


async def _knowledge_search(
    gateway: McpGateway,
    *,
    session,
    identity: TokenIdentity,
    query: str,
    limit: int = 5,
    **_ignored: Any,
) -> dict[str, Any]:
    from app.application.container import build_knowledge_service

    if not query:
        raise ValidationError("عبارت جست‌وجو لازم است.")
    service = build_knowledge_service(session)
    hits = await service.search(
        organization_id=identity.organization_id,
        query=query,
        limit=max(1, min(int(limit), 20)),
    )
    return {
        "results": [
            {
                "content": hit.get("content", ""),
                "source_name": hit.get("source_name", ""),
                "document_title": hit.get("document_title", ""),
            }
            for hit in hits
        ]
    }


async def _pipeline_status(
    gateway: McpGateway, *, session, identity: TokenIdentity, **_ignored: Any
) -> dict[str, Any]:
    """Delivery is another workstream's table; the tool exists and says so."""
    return {"available": False, "message": "خط تحویل در این نسخه در دسترس نیست."}


_HANDLERS: dict[str, Callable[..., Awaitable[dict[str, Any]]]] = {
    "app_get_profile": _app_profile,
    "tasks_list": _tasks_list,
    "tasks_get": _tasks_get,
    "tasks_create": _tasks_create,
    "memory_search": _memory_search,
    "memory_write": _memory_write,
    "knowledge_search": _knowledge_search,
    "pipeline_status": _pipeline_status,
}

_SPECS_BY_NAME: dict[str, ToolSpec] = {spec.name: spec for spec in TOOL_SPECS}


def _default_factory():
    from app.core.database import SessionFactory

    return SessionFactory


async def _note(kind: str, detail: str) -> None:
    """Record an unattributable attempt in the error ledger, best-effort.

    A bad token and a disabled gateway happen before the caller is known, so they
    belong to nobody's audit trail — but they are exactly the patterns an operator
    needs to see, which is what the error console is for.
    """
    from app.infrastructure.error_log import record_error

    try:
        await record_error(kind=kind, message=detail, source="mcp")
    except Exception:  # noqa: BLE001 - telemetry must never fail a call
        return


def build_gateway(*, session_factory=None) -> McpGateway:
    """The gateway the server and the tests both use.

    Each gateway owns its own in-memory rate limiter: one process, one budget.
    """
    return McpGateway(session_factory=session_factory)


__all__ = [
    "MCP_WRITABLE_KINDS",
    "TOOL_SPECS",
    "McpCallResult",
    "McpError",
    "McpGateway",
    "build_gateway",
]