smart-org/apps/api/tests/test_mcp.py

"""MCP gateway: scoped tokens, tool calls, and what they may never do.

The invariants under test:

* a token from organization A can never read organization B — not by asking, not
  by passing an ``organization_id``, not by guessing an id;
* a scope is a floor: a token minted ``tasks:read`` cannot create a task, and a
  token minted ``memory:write`` cannot write a *rule* (only ``fact`` and
  ``preference``);
* a revoked or expired token is indistinguishable from an unknown one;
* the rate limit refuses the 61st call in a minute;
* every call leaves an audit row carrying the token prefix and the tool;
* no secret appears in any tool output — the reflection redacts, and a token
  never is part of a payload.

The tests drive the gateway directly (``McpGateway.call``), which is the layer
the SDK server is a thin adapter around: the transport is the SDK's problem, the
rules are ours.
"""

from __future__ import annotations

import json

import pytest
from httpx import AsyncClient
from sqlalchemy import select

from app.application.services.mcp_service import (
    PREFIX_DISPLAY_CHARS,
    TOKEN_PREFIX,
    McpService,
    TokenRateLimiter,
)
from app.core.config import settings
from app.core.errors import ValidationError
from app.domain.base import new_id, utcnow
from app.domain.identity.models import PlatformRole, Role
from app.infrastructure.db.models import (
    AuditEventRow,
    MembershipRow,
    OrganizationRow,
    ProjectRow,
    UserRow,
)
from app.infrastructure.db.models_delivery import (
    AppRuntimeRow,
    PipelineRunRow,
    ProductionDeployRow,
)
from app.infrastructure.db.models_org import McpTokenRow
from app.interfaces.mcp.gateway import McpGateway
from tests.conftest import auth_header, signup

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


async def _org(session_factory, slug: str, role: Role = Role.OWNER) -> tuple[str, str]:
    async with session_factory() as session:
        user = UserRow(
            id=new_id(),
            email=f"{slug}@example.com",
            password_hash="x",
            name=slug,
            platform_role=PlatformRole.MANAGER.value,
        )
        session.add(user)
        await session.flush()
        org = OrganizationRow(id=new_id(), name=slug, slug=slug, owner_user_id=user.id)
        session.add(org)
        await session.flush()
        session.add(
            MembershipRow(organization_id=org.id, user_id=user.id, role=role.value)
        )
        await session.commit()
        return org.id, user.id


async def _gateway(session_factory, *, limit: int | None = None) -> McpGateway:
    return McpGateway(
        session_factory=session_factory,
        rate_limiter=TokenRateLimiter(limit_per_minute=limit),
    )


async def _token(
    session_factory, gateway: McpGateway, org_id: str, *, scopes: list[str],
    user_id: str | None = None,
) -> str:
    async with session_factory() as session:
        issued = await McpService(session).issue(
            organization_id=org_id, name="agent", scopes=scopes,
            created_by_user_id=user_id,
        )
        await session.commit()
        return issued.token


async def _call(gateway: McpGateway, token: str, tool: str, **arguments):
    return await gateway.call(tool=tool, arguments=arguments, token=token)


# ---------------------------------------------------------------------------
# Token lifecycle
# ---------------------------------------------------------------------------


async def test_a_token_is_shown_once_and_stored_only_as_a_hash(session_factory):
    org_id, user_id = await _org(session_factory, "issue")
    async with session_factory() as session:
        issued = await McpService(session).issue(
            organization_id=org_id,
            name="agent-1",
            scopes=["tasks:read", "tasks:write"],
            created_by_user_id=user_id,
            ttl_days=30,
        )
        await session.commit()
        raw = issued.token
        rows = (await session.execute(select(McpTokenRow))).scalars().all()
        assert len(rows) == 1
        # Only the hash and the display prefix exist; the token itself is gone.
        assert all(raw not in (row.token_hash or "") for row in rows)
        assert all(row.prefix == raw[: len(TOKEN_PREFIX) + PREFIX_DISPLAY_CHARS] for row in rows)
        assert issued.record.scopes == ["tasks:read", "tasks:write"]


async def test_unknown_revoked_and_expired_are_all_the_same_answer(session_factory):
    org_id, _ = await _org(session_factory, "verify")
    gateway = await _gateway(session_factory)
    raw = await _token(session_factory, gateway, org_id, scopes=["tasks:read"])

    async with session_factory() as session:
        service = McpService(session)
        identity = await service.verify(raw)
        assert identity is not None and identity.organization_id == org_id
        # A token for another org id is not this org's business.
        assert (await service.verify("smo_" + "x" * 24)) is None

        row = (await session.execute(select(McpTokenRow))).scalar_one()
        row.revoked_at = utcnow()
        await session.commit()
        assert await service.verify(raw) is None
        row.revoked_at = None
        from datetime import UTC, datetime, timedelta

        row.expires_at = datetime.now(UTC) - timedelta(minutes=1)
        await session.commit()
        assert await service.verify(raw) is None


async def test_unknown_scopes_are_refused_at_mint_time(session_factory):
    org_id, _ = await _org(session_factory, "scopes")
    async with session_factory() as session:
        with pytest.raises(ValidationError):
            await McpService(session).issue(
                organization_id=org_id, name="x", scopes=["root:everything"]
            )
        with pytest.raises(ValidationError):
            await McpService(session).issue(organization_id=org_id, name="x", scopes=[])


# ---------------------------------------------------------------------------
# Tenant isolation and scope enforcement
# ---------------------------------------------------------------------------


async def test_one_organizations_token_cannot_read_anothers_app(session_factory, seeded):
    org_a, _ = await _org(session_factory, "mcp-a")
    org_b, user_b = await _org(session_factory, "mcp-b")
    gateway = await _gateway(session_factory)
    token_a = await _token(
        session_factory, gateway, org_a, scopes=["app:read", "tasks:read"]
    )

    # org A has no project yet, so the honest answer is a refusal, not org B's
    # data. The important part is what it is *not*: never org B's app.
    result = await _call(gateway, token_a, "app_get_profile")
    assert result.ok is False

    # Even with the id handed to it, the tool refuses: identity comes from the token.
    result = await _call(gateway, token_a, "tasks_list")
    assert result.ok is True
    assert result.payload == {"tasks": []}

    # A token minted for B does not exist for A either.
    other = await _token(session_factory, gateway, org_b, scopes=["app:read"], user_id=user_b)
    assert (await _call(gateway, other, "app_get_profile")).ok is False


async def test_scope_enforcement_is_per_tool(session_factory, seeded):
    org_id, user_id = await _org(session_factory, "scope-org")
    gateway = await _gateway(session_factory)
    reader = await _token(
        session_factory, gateway, org_id, scopes=["tasks:read"], user_id=user_id
    )

    # Reading is allowed.
    listed = await _call(gateway, reader, "tasks_list")
    assert listed.ok is True
    # Writing is not: the scope is missing, and the refusal names the reason.
    created = await _call(gateway, reader, "tasks_create", title="کار تازه")
    assert created.ok is False
    assert created.code == "scope_denied"
    assert "دسترسی" in (created.error or "")

    # A refused write leaves nothing behind.
    async with session_factory() as session:
        from app.infrastructure.db.models import TaskRow

        rows = (await session.execute(select(TaskRow))).scalars().all()
        assert rows == []


async def test_memory_write_through_a_token_is_only_fact_or_preference(
    session_factory, seeded
):
    org_id, user_id = await _org(session_factory, "mem-write")
    gateway = await _gateway(session_factory)
    token = await _token(
        session_factory, gateway, org_id, scopes=["memory:read", "memory:write"],
        user_id=user_id,
    )

    ok = await _call(
        gateway, token, "memory_write", kind="fact", content="پورت اپ ۸۰۰۰ است."
    )
    assert ok.ok is True
    preference = await _call(
        gateway, token, "memory_write", kind="preference", content="ترجیح: فارسی بنویس."
    )
    assert preference.ok is True

    # A rule is not an agent's to set, and a lesson is the engine's to learn.
    refused_rule = await _call(
        gateway, token, "memory_write", kind="rule", content="همیشه تست بنویس."
    )
    assert refused_rule.ok is False
    refused_lesson = await _call(
        gateway, token, "memory_write", kind="lesson", content="درس تازه."
    )
    assert refused_lesson.ok is False

    # What it wrote is searchable, and what it may not write is not there.
    found = await _call(gateway, token, "memory_search", query="پورت")
    assert found.ok is True
    contents = [m["content"] for m in found.payload["memories"]]
    assert any("۸۰۰۰" in content for content in contents)
    assert not any("همیشه تست بنویس" in content for content in contents)


async def test_a_secret_written_through_a_token_never_comes_back(session_factory, seeded):
    org_id, user_id = await _org(session_factory, "mem-secret")
    gateway = await _gateway(session_factory)
    token = await _token(
        session_factory, gateway, org_id, scopes=["memory:read", "memory:write"],
        user_id=user_id,
    )
    await _call(
        gateway, token, "memory_write", kind="fact",
        content="کلید CI استفاده‌شده ghp_ABCDEFGHIJKLMNOPQRSTUVWXyz123 است.",
    )
    found = await _call(gateway, token, "memory_search", query="کلید")
    assert found.ok is True
    text = found.to_json()
    assert "ghp_ABCDEFGHIJKLMNOPQRSTUVWXyz123" not in text


async def test_knowledge_and_pipeline_tools(session_factory, seeded):
    org_id, user_id = await _org(session_factory, "knowledge-org")
    gateway = await _gateway(session_factory)
    token = await _token(
        session_factory, gateway, org_id, scopes=["memory:read", "delivery:read"],
        user_id=user_id,
    )

    # No sources yet: an empty result, not an error.
    empty = await _call(gateway, token, "knowledge_search", query="هرچی")
    assert empty.ok is True
    assert empty.payload["results"] == []

    # No app yet: the honest answer is "not available", never an empty delivery state.
    status = await _call(gateway, token, "pipeline_status")
    assert status.ok is True
    assert status.payload["available"] is False


async def test_pipeline_status_reports_the_delivery_state_and_no_host_details(
    session_factory, seeded
):
    org_id, user_id = await _org(session_factory, "pipeline-org")
    project_id = new_id()
    async with session_factory() as session:
        session.add(
            ProjectRow(
                id=project_id, organization_id=org_id, name="اپ", slug="app",
                created_by_user_id=user_id,
            )
        )
        await session.flush()
        session.add(
            AppRuntimeRow(
                project_id=project_id, organization_id=org_id,
                demo_host="pipeline-org.preview.example", host_port=18080,
                demo_container_id="ctr-secret-42", status="live",
            )
        )
        session.add(
            PipelineRunRow(
                id=new_id(), organization_id=org_id, project_id=project_id,
                trigger="engine", branch="dev", sha="0123456789abcdef", status="succeeded",
            )
        )
        session.add(
            ProductionDeployRow(
                id=new_id(), organization_id=org_id, project_id=project_id,
                kind="deploy", host="deploy-host.internal", status="succeeded",
                log_tail="token=super-secret-log-line",
            )
        )
        await session.commit()
    gateway = await _gateway(session_factory)
    token = await _token(
        session_factory, gateway, org_id, scopes=["delivery:read"], user_id=user_id
    )

    status = await _call(gateway, token, "pipeline_status")

    assert status.ok is True
    payload = status.payload
    assert payload["available"] is True
    assert payload["demo_url"] == "https://pipeline-org.preview.example"
    assert payload["runtime_status"] == "live"
    assert [run["status"] for run in payload["runs"]] == ["succeeded"]
    assert payload["runs"][0]["sha"] == "0123456789ab"
    assert set(payload["last_production_deploy"]) == {"status", "created_at"}
    assert payload["last_production_deploy"]["status"] == "succeeded"
    # Rule 7: no port, container id, deploy-target host or log tail leaves the gateway.
    text = json.dumps(payload, ensure_ascii=False)
    for hidden in (
        "host_port", "18080", "ctr-secret-42", "deploy-host.internal", "super-secret-log-line",
    ):
        assert hidden not in text


async def test_a_task_created_through_a_token_is_attributed_to_its_minter(
    session_factory, seeded
):
    org_id, user_id = await _org(session_factory, "task-org")
    gateway = await _gateway(session_factory)
    token = await _token(
        session_factory, gateway, org_id, scopes=["tasks:write", "tasks:read"],
        user_id=user_id,
    )

    created = await _call(
        gateway, token, "tasks_create", title="پیاده‌سازی ورود", description="صفحهٔ ورود"
    )
    assert created.ok is True, created.error
    task_id = created.payload["task"]["id"]

    fetched = await _call(gateway, token, "tasks_get", task_id=task_id)
    assert fetched.ok is True
    assert fetched.payload["task"]["title"] == "پیاده‌سازی ورود"

    # The write is attributed to the administrator who minted the token: a
    # delegated key never erases who holds it.
    async with session_factory() as session:
        from app.infrastructure.db.models import TaskRow

        row = await session.get(TaskRow, task_id)
        assert row is not None
        assert row.created_by_user_id == user_id

    # …and a task from another organization is a 404, not data.
    other_org, _ = await _org(session_factory, "task-other")
    other_token = await _token(
        session_factory, gateway, other_org, scopes=["tasks:read", "tasks:write"]
    )
    foreign = await _call(gateway, other_token, "tasks_get", task_id=task_id)
    assert foreign.ok is False


async def test_a_bad_token_or_a_disabled_switch_is_refused_everywhere(session_factory, seeded):
    org_id, user_id = await _org(session_factory, "off-switch")
    gateway = await _gateway(session_factory)
    token = await _token(
        session_factory, gateway, org_id, scopes=["tasks:read"], user_id=user_id
    )

    assert (await _call(gateway, "smo_wrong", "tasks_list")).ok is False
    assert (await _call(gateway, "", "tasks_list")).ok is False
    assert (await _call(gateway, token, "no_such_tool")).ok is False

    original = settings.mcp_gateway_enabled
    settings.mcp_gateway_enabled = False
    try:
        result = await _call(gateway, token, "tasks_list")
        assert result.ok is False
        assert result.code == "mcp_disabled"
    finally:
        settings.mcp_gateway_enabled = original

    # The per-organization switch is the owner's, not the token's.
    async with session_factory() as session:
        from app.application.container_org import build_model_grant_service

        await build_model_grant_service(session).set_mcp_enabled(
            organization_id=org_id, enabled=False, actor_user_id=user_id
        )
        await session.commit()
    result = await _call(gateway, token, "tasks_list")
    assert result.ok is False
    assert result.code == "org_disabled"


async def test_the_rate_limit_refuses_the_sixty_first_call(session_factory, seeded):
    org_id, user_id = await _org(session_factory, "rate")
    gateway = await _gateway(session_factory, limit=5)
    token = await _token(
        session_factory, gateway, org_id, scopes=["tasks:read"], user_id=user_id
    )
    codes = [
        (await _call(gateway, token, "tasks_list")).ok for _ in range(7)
    ]
    assert codes == [True] * 5 + [False, False]
    limited = await _call(gateway, token, "tasks_list")
    assert limited.code == "rate_limited"


# ---------------------------------------------------------------------------
# Audit
# ---------------------------------------------------------------------------


async def test_every_call_leaves_an_audit_row_with_the_prefix(session_factory, seeded):
    org_id, user_id = await _org(session_factory, "audit")
    gateway = await _gateway(session_factory)
    token = await _token(
        session_factory, gateway, org_id, scopes=["tasks:read"], user_id=user_id
    )
    await _call(gateway, token, "tasks_list")
    await _call(gateway, token, "tasks_create", title="رد شده")

    async with session_factory() as session:
        rows = (
            (
                await session.execute(
                    select(AuditEventRow).where(
                        AuditEventRow.organization_id == org_id,
                        AuditEventRow.action == "mcp.tool_call",
                    )
                )
            )
            .scalars()
            .all()
        )
    assert len(rows) == 2
    by_tool = {row.target_id: row for row in rows}
    assert set(by_tool) == {"tasks_list", "tasks_create"}
    # The prefix identifies the token; the token itself never appears.
    prefix = token[: len(TOKEN_PREFIX) + PREFIX_DISPLAY_CHARS]
    assert all(row.metadata_json["prefix"] == prefix for row in rows)
    assert all(token not in str(row.metadata_json) for row in rows)
    assert by_tool["tasks_list"].metadata_json["ok"] is True
    assert by_tool["tasks_create"].metadata_json["ok"] is False
    assert by_tool["tasks_create"].metadata_json["error"]


# ---------------------------------------------------------------------------
# HTTP surface: tokens and the owner's switches
# ---------------------------------------------------------------------------


async def _owner_with_org(client: AsyncClient, slug: str) -> dict:
    body = await signup(client, f"{slug}-owner@example.com", name="Owner")
    org = await client.post(
        "/api/organizations", json={"name": slug}, headers=auth_header(body["access_token"])
    )
    assert org.status_code == 201, org.text
    return {"token": body["access_token"], "org_id": org.json()["id"]}


async def _member(client: AsyncClient, owner_token: str, org_id: str, email: str, role: str):
    invite = await client.post(
        f"/api/organizations/{org_id}/invites",
        json={"email": email, "role": role},
        headers=auth_header(owner_token),
    )
    assert invite.status_code == 201, invite.text
    account = await signup(client, email, name=role)
    accepted = await client.post(
        "/api/invites/accept",
        json={"token": invite.json()["invite_token"]},
        headers=auth_header(account["access_token"]),
    )
    assert accepted.status_code == 200, accepted.text
    return account["access_token"]


async def test_only_an_admin_mints_tokens_and_only_the_owner_switches_the_gateway(
    client: AsyncClient,
    manager_token: str,
    platform_owner_token: str,
):
    org = await _owner_with_org(client, "mcp-org")
    viewer = await _member(client, org["token"], org["org_id"], "v@example.com", "viewer")
    admin = await _member(client, org["token"], org["org_id"], "a@example.com", "admin")

    # A viewer reads the list but cannot mint.
    listing = await client.get(
        f"/api/organizations/{org['org_id']}/mcp-tokens", headers=auth_header(viewer)
    )
    assert listing.status_code == 200, listing.text
    assert listing.json() == []
    refused = await client.post(
        f"/api/organizations/{org['org_id']}/mcp-tokens",
        json={"name": "nope", "scopes": ["tasks:read"]},
        headers=auth_header(viewer),
    )
    assert refused.status_code == 403, refused.text

    # The owner's switches are the owner's.
    not_owner = await client.get("/api/admin/mcp", headers=auth_header(manager_token))
    assert not_owner.status_code == 403
    anonymous = await client.get("/api/admin/mcp")
    assert anonymous.status_code == 401

    # An admin mints one, sees the raw token once, and it never returns.
    created = await client.post(
        f"/api/organizations/{org['org_id']}/mcp-tokens",
        json={"name": "agent-1", "scopes": ["tasks:read", "tasks:write"]},
        headers=auth_header(admin),
    )
    assert created.status_code == 201, created.text
    body = created.json()
    assert body["token"].startswith("smo_")
    assert body["prefix"] == body["token"][: len(TOKEN_PREFIX) + PREFIX_DISPLAY_CHARS]
    assert body["scopes"] == ["tasks:read", "tasks:write"]
    assert body["status"] == "active"

    again = await client.get(
        f"/api/organizations/{org['org_id']}/mcp-tokens", headers=auth_header(admin)
    )
    assert again.status_code == 200, again.text
    assert "token" not in again.json()[0]
    assert again.json()[0]["prefix"] == body["prefix"]

    # An unknown scope is refused rather than silently dropped.
    bad_scope = await client.post(
        f"/api/organizations/{org['org_id']}/mcp-tokens",
        json={"name": "bad", "scopes": ["root:all"]},
        headers=auth_header(admin),
    )
    assert bad_scope.status_code == 422, bad_scope.text

    # Revoke, then revoking again is a 404 — the same idempotence rule as grants.
    revoked = await client.delete(
        f"/api/organizations/{org['org_id']}/mcp-tokens/{body['id']}", headers=auth_header(admin)
    )
    assert revoked.status_code == 200, revoked.text
    twice = await client.delete(
        f"/api/organizations/{org['org_id']}/mcp-tokens/{body['id']}", headers=auth_header(admin)
    )
    assert twice.status_code == 404

    # The owner sees the gateway state and may switch it and one organization.
    state = await client.get("/api/admin/mcp", headers=auth_header(platform_owner_token))
    assert state.status_code == 200, state.text
    assert state.json()["gateway_enabled"] is True
    assert any(
        org_row["organization_id"] == org["org_id"]
        for org_row in state.json()["organizations"]
    )

    switched = await client.put(
        "/api/admin/mcp",
        json={"enabled": False},
        headers=auth_header(platform_owner_token),
    )
    assert switched.status_code == 200, switched.text
    assert switched.json()["gateway_enabled"] is False
    org_switch = await client.put(
        f"/api/admin/mcp/organizations/{org['org_id']}",
        json={"enabled": False},
        headers=auth_header(platform_owner_token),
    )
    assert org_switch.status_code == 200, org_switch.text
    assert org_switch.json()["mcp_enabled"] is False
    settings_view = await client.get(
        f"/api/organizations/{org['org_id']}/mcp-tokens/settings",
        headers=auth_header(admin),
    )
    assert settings_view.json()["mcp_enabled"] is False
    assert settings_view.json()["gateway_enabled"] is False
    assert "tasks:read" in settings_view.json()["scopes"]