bifroest/tests/test_workspace_mcp.py

"""MCP-Verbindungen je Workspace und Datenbanken über Datenbank-Treiber (DB-Dienst aus services/db-mcp).

Getestet gegen echte Streamable-HTTP-Server im selben Prozess: den Testserver aus test_mcp und den
DB-Dienst mit dem PostgreSQL-Treiber gegen eine eigene Datenbank im Test-Postgres.
"""

import socket
import sys
import threading
import time
import uuid
from pathlib import Path

import psycopg
import pytest
import uvicorn
from sqlalchemy import select
from test_mcp import Remote, ask, last_output, register

from app.audit.models import AuditEvent
from app.core.crypto import Vault
from app.identity.models import User
from app.mcp import service as mcp
from app.mcp.models import APPROVED, PENDING, McpServer
from app.notify.models import Notification
from app.tasks import states
from app.tasks.models import TaskRun
from app.tasks.worker import Worker
from app.tools import service as tools
from app.tools.models import ALLOW, REQUIRE_APPROVAL, ToolCall, ToolDefinition
from helpers import connect_repo, create_workspace

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "services" / "db-mcp"))
from dbmcp.core.config import Settings as ServiceSettings
from dbmcp.core.server import app as service_app
from dbmcp.postgres.driver import PostgresDriver

ADMIN_URL = "postgresql://aiw:aiw@127.0.0.1:55433/postgres"
SERVICE_KEY = "dienst-schluessel-0123456789"
PASSWORD = "lesen-nur-123"


@pytest.fixture
def settings(settings):
    return settings.model_copy(update={"mcp_allow_local": True, "mcp_allowed_hosts": "mcp.intern",
                                       "mcp_timeout_seconds": 5, "db_allow_local": True})


@pytest.fixture
def admin(browser, people, db):
    with db.session() as session:
        anna = session.get(User, people["anna"])
        anna.is_admin, anna.email = True, "anna@example.org"
    return browser("anna")


@pytest.fixture
def principal(people):
    from app.identity import Principal
    return Principal(user_id=people["anna"], username="anna", display_name="Anna", is_admin=True,
                     team_ids=frozenset({people["team_demo"]}))


@pytest.fixture
def workspace(admin, people, git_repo) -> int:
    workspace_id = create_workspace(admin, people["team_demo"])
    connect_repo(admin, workspace_id, git_repo)
    return workspace_id


@pytest.fixture
def other(browser, people, git_repo) -> tuple[object, int]:
    """Workspace eines anderen Teams; dora ist dort Admin."""
    dora = browser("dora")
    workspace_id = create_workspace(dora, people["team_other"], "Anderes")
    return dora, workspace_id


def new_chat(b, workspace_id: int) -> int:
    return int(b.post(f"/workspaces/{workspace_id}/chats").headers["location"].rsplit("/", 1)[1])


def add_connection(b, workspace_id: int, url: str, name: str = "eigen", **extra):
    return b.post(f"/workspaces/{workspace_id}/connections",
                  {"name": name, "label": "Eigener Dienst", "url": url, "exposure_confirmed": "1", **extra})


def server(db, name: str) -> McpServer:
    with db.session() as session:
        found = session.scalar(select(McpServer).where(McpServer.name == name))
        session.expunge(found)
        return found


def offered(fake_provider, index: int = -1) -> set[str]:
    return {t["name"] for t in fake_provider.calls[index]["tools"] or []}


# --- Eigene Verbindungen ---

def test_neue_adresse_wartet_auf_freigabe(admin, workspace, db, fake_provider, settings, principal):
    with Remote() as remote:
        response = add_connection(admin, workspace, remote.url, token="eigenes-token-9876")
        assert response.status_code == 303, response.text
        found = server(db, "eigen")
        assert found.approval == PENDING and found.workspace_id == workspace and found.secret_hint == "9876"
        assert "wartet auf Freigabe" in admin.get(f"/workspaces/{workspace}/settings?tab=mcp").text
        with db.session() as session:
            note = session.scalar(select(Notification).where(Notification.event == "mcp.connection_requested"))
            assert note is not None and remote.url in note.body
        # Vor der Freigabe kein Abruf, keine Tools
        assert remote.tokens == []
        page = admin.get("/admin/mcp").text
        assert "Warten auf Freigabe" in page and "Projekt X" in page

        assert admin.post(f"/admin/mcp/{found.id}/approve", {"note": "passt"}).status_code == 303
        found = server(db, "eigen")
        assert found.approval == APPROVED and found.status == "ok" and "3 Tools" in found.status_message
        assert "Bearer eigenes-token-9876" in remote.tokens
        # Neue Tools sind gesperrt, bis sie eingestuft sind
        chat = new_chat(admin, workspace)
        ask(admin, chat)
        Worker(db, settings).run_once()
        assert not any(n.startswith("eigen__") for n in offered(fake_provider))
        with db.session() as session:
            tools.sync(session)
            tools.classify(session, principal, "eigen__lookup", "read", ALLOW)
        fake_provider.script = [[("eigen__lookup", {"term": "Port"})], "Port 8443."]
        ask(admin, chat, "Und jetzt?")
        Worker(db, settings).run_once()
        assert "eigen__lookup" in offered(fake_provider, -2)
        assert "Treffer zu Port" in last_output(fake_provider)


def test_ablehnen(admin, workspace, db):
    with Remote() as remote:
        add_connection(admin, workspace, remote.url)
        found = server(db, "eigen")
        admin.post(f"/admin/mcp/{found.id}/reject", {"note": "nicht im Firmennetz"})
        found = server(db, "eigen")
        assert found.approval == "rejected" and "nicht im Firmennetz" in found.status_message
        assert remote.tokens == []
        # Abgelehnt bleibt abgelehnt; ein zweites Mal entscheiden geht nicht
        assert admin.post(f"/admin/mcp/{found.id}/approve").status_code == 400


def test_bekannte_adresse_sofort_aktiv_und_einstufung_wird_geerbt(admin, workspace, db, principal, fake_provider,
                                                                  settings):
    with Remote() as remote:
        register(admin, remote)
        with db.session() as session:
            tools.sync(session)
            tools.classify(session, principal, "munin__lookup", "read", ALLOW)
            tools.classify(session, principal, "munin__remember", "write_internal", REQUIRE_APPROVAL)
        assert add_connection(admin, workspace, remote.url, name="meins", token="anderes-token-1111").status_code == 303
        found = server(db, "meins")
        assert found.approval == APPROVED and found.status == "ok"
        with db.session() as session:
            definitions = {d.name: d for d in tools.sync(session)}
            events = list(session.scalars(select(AuditEvent).where(
                AuditEvent.action == "tool.classification_inherited")))
        lookup, remember = definitions["meins__lookup"], definitions["meins__remember"]
        assert lookup.classified and lookup.decision == ALLOW and lookup.inherited_from == "munin__lookup"
        assert remember.decision == REQUIRE_APPROVAL
        assert not definitions["meins__boom"].classified
        assert {e.target_id for e in events} >= {"meins__lookup", "meins__remember"}
        # Eine neue Einstufung des Originals gilt für alle gleichen Tools
        with db.session() as session:
            tools.classify(session, principal, "munin__lookup", "read", REQUIRE_APPROVAL)
            twin = session.scalar(select(ToolDefinition).where(ToolDefinition.name == "meins__lookup"))
            assert twin.decision == REQUIRE_APPROVAL
        # Das eigene Token der Verbindung geht an den Server
        remote.tokens.clear()
        fake_provider.script = [[("meins__remember", {"text": "x"})], "Ok."]
        ask(admin, new_chat(admin, workspace))
        Worker(db, settings).run_once()
        with db.session() as session:
            assert session.scalar(select(ToolCall).where(ToolCall.tool == "meins__remember")).status == "pending"


def test_header_werden_verschluesselt_und_gesendet(admin, workspace, db, principal, fake_provider, settings):
    with Remote() as remote:
        register(admin, remote)
        response = add_connection(admin, workspace, remote.url, name="kopf",
                                  headers="X-Api-Key: geheim-123\n\nX-Tenant: blau")
        assert response.status_code == 303, response.text
        found = server(db, "kopf")
        assert found.header_names == ["X-Api-Key", "X-Tenant"] and "geheim-123" not in (found.headers or "")
        page = admin.get(f"/workspaces/{workspace}/settings?tab=mcp").text
        assert "X-Api-Key, X-Tenant" in page and "geheim-123" not in page
        with db.session() as session:
            target = mcp.targets(session, Vault(settings.secret_key), settings, workspace)["kopf"]
        assert target.headers == {"X-Api-Key": "geheim-123", "X-Tenant": "blau"}
        for bad in ("Authorization: Bearer x", "kein header", "X-A: " + "y" * 5000):
            response = add_connection(admin, workspace, remote.url, name="kopf2", headers=bad)
            assert response.status_code == 400, bad
            assert "geheim" not in response.text


def test_kurzname_global_eindeutig_und_nur_workspace_admin(admin, workspace, other, browser, db):
    dora, other_id = other
    with Remote() as remote:
        assert add_connection(admin, workspace, remote.url, name="dienst").status_code == 303
        response = add_connection(dora, other_id, remote.url, name="dienst")
        assert response.status_code == 400 and "gibt es schon" in response.text
        assert add_connection(admin, workspace, remote.url, name="ohne", exposure_confirmed="").status_code == 400
        ben = browser("ben")
        assert add_connection(ben, workspace, remote.url, name="ben1").status_code in (403, 404)
        found = server(db, "dienst")
        # Fremde sehen die Verbindung nicht
        assert dora.get(f"/connections/{found.id}/edit").status_code == 404
        assert dora.post(f"/connections/{found.id}/delete").status_code == 404
        assert ben.get(f"/connections/{found.id}/edit").status_code in (403, 404)


def test_verbindung_nur_im_eigenen_workspace(admin, workspace, other, db, principal, fake_provider, settings,
                                             people):
    dora, other_id = other
    with Remote() as remote:
        register(admin, remote)
        add_connection(admin, workspace, remote.url, name="nurwir")
        with db.session() as session:
            tools.sync(session)
            tools.classify(session, principal, "nurwir__lookup", "read", ALLOW)
        # Im anderen Workspace weder angeboten noch in der Tool-Übersicht
        assert "nurwir__lookup" not in dora.get(f"/workspaces/{other_id}/settings?tab=tools").text
        assert "nurwir__lookup" in admin.get(f"/workspaces/{workspace}/settings?tab=tools").text
        assert dora.post(f"/workspaces/{other_id}/tools/nurwir__lookup", {"decision": "deny"}).status_code == 404
        with db.session() as session:
            assert "mcp:nurwir" not in mcp.kinds(session, other_id)
            assert "mcp:nurwir" in mcp.kinds(session, workspace)


def test_bearbeiten_und_entfernen(admin, workspace, db):
    with Remote() as remote:
        register(admin, remote)
        add_connection(admin, workspace, remote.url, name="aendern", token="erstes-token-1111")
        found = server(db, "aendern")
        assert admin.get(f"/connections/{found.id}/edit").status_code == 200
        assert admin.post(f"/connections/{found.id}", {"label": "Neu", "token": "zweites-token-2222",
                                                       "active": "1"}).status_code == 303
        found = server(db, "aendern")
        assert found.label == "Neu" and found.secret_hint == "2222" and found.status == "ok"
        assert "Bearer zweites-token-2222" in remote.tokens
        assert admin.post(f"/connections/{found.id}/delete").status_code == 303
        with db.session() as session:
            assert session.get(McpServer, found.id) is None


# --- Datenbanken ---

@pytest.fixture(scope="module")
def shop() -> dict[str, str]:
    suffix = uuid.uuid4().hex[:8]
    name, reader = f"wsmcp_{suffix}", f"wsmcp_ro_{suffix}"
    admin = psycopg.connect(ADMIN_URL, autocommit=True)
    admin.execute(f"CREATE DATABASE {name}")
    admin.execute(f"CREATE ROLE {reader} LOGIN PASSWORD '{PASSWORD}'")
    with psycopg.connect(ADMIN_URL.rsplit("/", 1)[0] + f"/{name}", autocommit=True) as conn:
        conn.execute("CREATE TABLE kunden (id int PRIMARY KEY, name text, umsatz numeric(10,2)); "
                     "INSERT INTO kunden SELECT i, 'Kunde ' || i, i * 10 FROM generate_series(1, 30) i;")
        conn.execute(f"GRANT SELECT ON kunden TO {reader}")
    yield {"database": name, "user": reader}
    admin.execute(f"DROP DATABASE {name} WITH (FORCE)")
    admin.execute(f"DROP ROLE {reader}")
    admin.close()


class DbService:
    """DB-Dienst (PostgreSQL) über HTTP im selben Prozess."""

    def __init__(self):
        settings = ServiceSettings(api_keys={SERVICE_KEY: "bifroest"}, allow_local=True, max_rows=100,
                                   query_timeout=5)
        with socket.socket() as s:
            s.bind(("127.0.0.1", 0))
            self.port = s.getsockname()[1]
        self.server = uvicorn.Server(uvicorn.Config(service_app(PostgresDriver(settings), settings),
                                                    host="127.0.0.1", port=self.port, log_level="warning",
                                                    lifespan="on"))
        self.thread = threading.Thread(target=self.server.run, daemon=True)

    @property
    def url(self) -> str:
        return f"http://localhost:{self.port}/mcp"

    def __enter__(self) -> "DbService":
        self.thread.start()
        for _ in range(100):
            if self.server.started:
                return self
            time.sleep(0.05)
        raise RuntimeError("DB-Dienst startet nicht")

    def __exit__(self, *exc) -> None:
        self.server.should_exit = True
        self.thread.join(5)


@pytest.fixture
def db_service():
    with DbService() as running:
        yield running


def register_driver(admin, db, principal, db_service, classify: bool = True) -> int:
    response = admin.post("/admin/databases", {"name": "pgdienst", "label": "PostgreSQL", "url": db_service.url,
                                               "token": SERVICE_KEY})
    assert response.status_code == 303, response.text
    driver_id = int(response.headers["location"].rsplit("/", 1)[1])
    assert admin.post(f"/admin/mcp/{driver_id}/discover").status_code == 200
    if classify:
        with db.session() as session:
            tools.sync(session)
            for name in ("list_tables", "describe_table", "query", "check_connection"):
                tools.classify(session, principal, f"pgdienst__{name}", "read", ALLOW)
    return driver_id


def add_database(b, workspace_id: int, driver_id: int, shop, name: str = "shop", **extra):
    return b.post(f"/workspaces/{workspace_id}/databases", {
        "driver_id": str(driver_id), "name": name, "label": "Shop", "host": "127.0.0.1", "port": "55433",
        "database": shop["database"], "user": shop["user"], "password": PASSWORD, "sslmode": "disable",
        "exposure_confirmed": "1", **extra})


def test_datenbank_anlegen_testen_und_abfragen(admin, workspace, db, principal, db_service, shop, fake_provider,
                                               settings):
    driver_id = register_driver(admin, db, principal, db_service)
    # Der Treiber lässt sich nicht freischalten und erscheint nicht unter den zentralen Servern
    page = admin.get(f"/workspaces/{workspace}/settings?tab=mcp").text
    assert "Datenbank hinzufügen" in page and "pgdienst" not in page
    assert "pgdienst" not in admin.get("/admin/mcp").text and "PostgreSQL" in admin.get("/admin/databases").text
    assert admin.post(f"/workspaces/{workspace}/mcp/{driver_id}").status_code == 400
    assert admin.get(f"/workspaces/{workspace}/databases/new").status_code == 200

    response = add_database(admin, workspace, driver_id, shop)
    assert response.status_code == 303, response.text
    found = server(db, "shop")
    assert found.template_id == driver_id and found.status == "ok", found.status_message
    assert "1 Tabelle" in found.status_message or "Tabellen" in found.status_message
    assert PASSWORD not in (found.headers or "") and found.config["user"] == shop["user"]
    with db.session() as session:
        definitions = {d.name: d for d in tools.sync(session)}
    # Die Einstufung des Treibers gilt für die Datenbank; einzeln einstufen geht nicht
    assert definitions["shop__query"].classified and definitions["shop__query"].inherited_from == "pgdienst__query"
    tools_page = admin.get("/admin/tools").text
    assert "pgdienst__query" in tools_page and "shop__query" not in tools_page
    assert admin.post("/admin/tools/shop__query", {"risk_class": "read", "decision": "deny"}).status_code == 400
    assert "shop" in admin.get("/admin/databases").text

    fake_provider.script = [[("shop__query", {"sql": "select id, umsatz from kunden order by id limit 3"})],
                            "Drei Kunden."]
    run_id = ask(admin, new_chat(admin, workspace), "Umsatz?")
    Worker(db, settings).run_once()
    first = fake_provider.calls[0]
    assert {"shop__list_tables", "shop__describe_table", "shop__query"} <= offered(fake_provider, 0)
    assert not any(n.startswith("pgdienst__") for n in offered(fake_provider, 0))
    assert "SQL-Dialekt: PostgreSQL" in first["instructions"]
    assert PASSWORD not in first["instructions"] and shop["user"] not in first["instructions"]
    with db.session() as session:
        assert session.get(TaskRun, run_id).status == states.COMPLETED
        call = session.scalar(select(ToolCall).where(ToolCall.run_id == run_id))
        assert call.status == "ok" and call.result["columns"] == ["id", "umsatz"]
        assert len(call.result["rows"]) == 3


def test_workspace_regel_gilt_je_datenbank(admin, workspace, db, principal, db_service, shop, fake_provider,
                                           settings):
    driver_id = register_driver(admin, db, principal, db_service)
    add_database(admin, workspace, driver_id, shop, name="prod")
    add_database(admin, workspace, driver_id, shop, name="test")
    admin.get(f"/workspaces/{workspace}/settings?tab=tools")
    assert admin.post(f"/workspaces/{workspace}/tools/prod__query",
                      {"decision": "require_approval"}).status_code in (200, 303)
    fake_provider.script = [[("prod__query", {"sql": "select 1"}), ("test__query", {"sql": "select 1"})], "Ok."]
    ask(admin, new_chat(admin, workspace))
    Worker(db, settings).run_once()
    with db.session() as session:
        status = {c.tool: c.status for c in session.scalars(select(ToolCall))}
    assert status["prod__query"] == "pending" and status["test__query"] == "ok"


def test_datenbank_regeln(admin, workspace, db, principal, db_service, shop):
    driver_id = register_driver(admin, db, principal, db_service)
    response = add_database(admin, workspace, driver_id, shop, host="db.anderswo.example", sslmode="require")
    assert response.status_code == 400 and "nicht freigegeben" in response.text
    response = add_database(admin, workspace, driver_id, shop, name="falsch", password="falsch-falsch")
    assert response.status_code == 303
    found = server(db, "falsch")
    assert found.status == "error" and "falsch-falsch" not in found.status_message
    # Passwort korrigieren: Bearbeiten testet neu
    page = admin.get(f"/connections/{found.id}/edit")
    assert page.status_code == 200 and "falsch-falsch" not in page.text and shop["user"] in page.text
    response = admin.post(f"/connections/{found.id}/database", {"label": "Shop", "user": shop["user"],
                                                                "password": PASSWORD, "sslmode": "disable",
                                                                "active": "1"})
    assert response.status_code == 303
    assert server(db, "falsch").status == "ok"
    # Der Treiber lässt sich nicht löschen, solange Datenbanken ihn nutzen; die Art bleibt
    assert admin.post(f"/admin/mcp/{driver_id}", {"label": "X", "url": db_service.url, "role": "",
                                                  "active": "1"}).status_code == 400
    response = admin.post(f"/admin/mcp/{driver_id}/delete")
    assert response.status_code == 400