X-Financial/server/tests/test_agent_run_tenant_security.py

from __future__ import annotations

import json
from collections.abc import Generator
from datetime import UTC, datetime, timedelta

import pytest
from auth_helpers import install_legacy_header_auth_override
from fastapi import FastAPI
from fastapi.testclient import TestClient
from sqlalchemy import create_engine
from sqlalchemy.orm import Session, sessionmaker
from sqlalchemy.pool import StaticPool

from app.api.deps import get_db
from app.api.v1.endpoints.agent_runs import router as agent_runs_router
from app.db.base import Base
from app.models.agent_run import AgentRun, AgentToolCall, SemanticParseLog
from app.services.agent_runs import AgentRunService
from app.services.finance_dashboard_scope import (
    FINANCE_DASHBOARD_TASK_TYPE,
    resolve_finance_dashboard_data_scope,
)


def _session_factory() -> sessionmaker[Session]:
    engine = create_engine(
        "sqlite+pysqlite:///:memory:",
        connect_args={"check_same_thread": False},
        poolclass=StaticPool,
    )
    Base.metadata.create_all(bind=engine)
    return sessionmaker(bind=engine, autoflush=False, autocommit=False)


def _run(
    *,
    run_id: str,
    started_at: datetime,
    route_tenant_id: str | None,
    ontology_tenant_id: str | None,
    status: str = "succeeded",
    error_message: str | None = None,
) -> AgentRun:
    route_json: dict[str, object] = {"stage": "tenant-security-test"}
    ontology_json: dict[str, object] = {
        "scenario": "expense",
        "intent": "query",
    }
    if route_tenant_id is not None:
        route_json["tenant_id"] = route_tenant_id
    if ontology_tenant_id is not None:
        ontology_json["tenant_id"] = ontology_tenant_id
    return AgentRun(
        run_id=run_id,
        agent="orchestrator",
        source="user_message",
        user_id=f"user-{run_id}",
        ontology_json=ontology_json,
        route_json=route_json,
        permission_level="read",
        status=status,
        result_summary=f"summary-{run_id}",
        error_message=error_message,
        started_at=started_at,
        finished_at=started_at + timedelta(seconds=1),
    )


def _seed_runs(db: Session) -> None:
    now = datetime.now(UTC)
    tenant_a = _run(
        run_id="run-normal-tenant-a",
        started_at=now - timedelta(minutes=4),
        route_tenant_id="tenant-a",
        ontology_tenant_id="tenant-a",
    )
    tenant_b = _run(
        run_id="run-normal-tenant-b",
        started_at=now - timedelta(minutes=3),
        route_tenant_id="tenant-b",
        ontology_tenant_id="tenant-b",
        status="failed",
        error_message="tenant-b-run-secret",
    )
    legacy_unscoped = _run(
        run_id="run-normal-unscoped",
        started_at=now - timedelta(minutes=2),
        route_tenant_id=None,
        ontology_tenant_id=None,
        status="failed",
        error_message="legacy-unscoped-secret",
    )
    partial_scope = _run(
        run_id="run-normal-partial-scope",
        started_at=now - timedelta(seconds=90),
        route_tenant_id="tenant-a",
        ontology_tenant_id=None,
        status="failed",
        error_message="partial-scope-secret",
    )
    conflicting = _run(
        run_id="run-normal-conflicting",
        started_at=now - timedelta(minutes=1),
        route_tenant_id="tenant-a",
        ontology_tenant_id="tenant-b",
        status="failed",
        error_message="conflicting-scope-secret",
    )
    finance_scope = resolve_finance_dashboard_data_scope("tenant-a")
    finance_snapshot = _run(
        run_id="run-finance-tenant-a",
        started_at=now,
        route_tenant_id="tenant-a",
        ontology_tenant_id="tenant-a",
        status="failed",
        error_message="finance-snapshot-secret",
    )
    finance_snapshot.route_json.update(
        {
            "task_type": FINANCE_DASHBOARD_TASK_TYPE,
            "data_scope": finance_scope,
        }
    )
    finance_snapshot.ontology_json["data_scope"] = finance_scope
    db.add_all(
        [
            tenant_a,
            tenant_b,
            legacy_unscoped,
            partial_scope,
            conflicting,
            finance_snapshot,
        ]
    )
    db.flush()
    db.add_all(
        [
            AgentToolCall(
                run_id=tenant_a.run_id,
                tool_type="database",
                tool_name="tenant-a.tool",
                request_json={"private": "tenant-a-request"},
                response_json={"private": "tenant-a-response"},
                status="succeeded",
                duration_ms=3,
            ),
            AgentToolCall(
                run_id=tenant_b.run_id,
                tool_type="database",
                tool_name="tenant-b.tool",
                request_json={"private": "tenant-b-request-secret"},
                response_json={"private": "tenant-b-response-secret"},
                status="failed",
                duration_ms=5,
                error_message="tenant-b-tool-secret",
            ),
            SemanticParseLog(
                run_id=tenant_a.run_id,
                user_id="tenant-a-user",
                raw_query="tenant-a-raw-query",
                scenario="expense",
                intent="query",
                confidence=0.9,
            ),
            SemanticParseLog(
                run_id=tenant_b.run_id,
                user_id="tenant-b-user",
                raw_query="tenant-b-raw-query-secret",
                scenario="expense",
                intent="query",
                confidence=0.9,
            ),
        ]
    )
    db.commit()


def _client(
    session_factory: sessionmaker[Session],
) -> TestClient:
    app = FastAPI()
    app.include_router(agent_runs_router)
    install_legacy_header_auth_override(app)

    def override_db() -> Generator[Session, None, None]:
        with session_factory() as db:
            yield db

    app.dependency_overrides[get_db] = override_db
    return TestClient(app)


def _headers(tenant_id: str, *, is_admin: bool = False) -> dict[str, str]:
    return {
        "X-Auth-Username": f"user-{tenant_id}",
        "X-Auth-Role-Codes": "user",
        "X-Auth-Is-Admin": "true" if is_admin else "false",
        "X-Auth-Tenant-Id": tenant_id,
    }


def test_agent_run_endpoints_fail_closed_for_normal_cross_tenant_runs() -> None:
    session_factory = _session_factory()
    with session_factory() as db:
        _seed_runs(db)
    client = _client(session_factory)

    tenant_a_list = client.get(
        "/agent-runs",
        params={"limit": 1},
        headers=_headers("tenant-a"),
    )
    tenant_a_summary = client.get(
        "/agent-runs/summary",
        params={"limit": 1},
        headers=_headers("tenant-a"),
    )
    tenant_b_list = client.get(
        "/agent-runs",
        headers=_headers("tenant-b"),
    )

    assert tenant_a_list.status_code == 200
    assert [item["run_id"] for item in tenant_a_list.json()] == ["run-normal-tenant-a"]
    assert tenant_a_summary.status_code == 200
    assert tenant_a_summary.json()["total_runs"] == 1
    assert tenant_a_summary.json()["failed_runs"] == 0
    assert tenant_a_summary.json()["recent_errors"] == []
    assert tenant_b_list.status_code == 200
    assert [item["run_id"] for item in tenant_b_list.json()] == ["run-normal-tenant-b"]

    tenant_a_payload = json.dumps(
        [tenant_a_list.json(), tenant_a_summary.json()],
        ensure_ascii=False,
    )
    assert "tenant-b" not in tenant_a_payload
    assert "legacy-unscoped-secret" not in tenant_a_payload
    assert "partial-scope-secret" not in tenant_a_payload
    assert "conflicting-scope-secret" not in tenant_a_payload
    assert "finance-snapshot-secret" not in tenant_a_payload

    own_detail = client.get(
        "/agent-runs/run-normal-tenant-a",
        headers=_headers("tenant-a"),
    )
    assert own_detail.status_code == 200
    assert own_detail.json()["tool_calls"][0]["request_json"] == {"private": "tenant-a-request"}
    assert own_detail.json()["semantic_parse"]["raw_query"] == "tenant-a-raw-query"

    hidden_run_ids = [
        "run-normal-tenant-b",
        "run-normal-unscoped",
        "run-normal-partial-scope",
        "run-normal-conflicting",
    ]
    for run_id in hidden_run_ids:
        response = client.get(
            f"/agent-runs/{run_id}",
            headers=_headers("tenant-a", is_admin=True),
        )
        assert response.status_code == 404
        response_body = json.dumps(response.json(), ensure_ascii=False)
        assert "tenant-b-request-secret" not in response_body
        assert "tenant-b-response-secret" not in response_body
        assert "tenant-b-raw-query-secret" not in response_body


def test_agent_run_creation_stamps_both_payloads_and_rejects_conflicts() -> None:
    session_factory = _session_factory()
    with session_factory() as db:
        service = AgentRunService(db)
        created = service.create_run(
            agent="orchestrator",
            source="user_message",
            tenant_id="tenant-a",
            ontology_json={"scenario": "expense"},
            route_json={"stage": "created"},
            status="running",
        )

        assert created.ontology_json["tenant_id"] == "tenant-a"
        assert created.route_json["tenant_id"] == "tenant-a"

        updated = service.update_run(
            created.run_id,
            ontology_json={"scenario": "expense", "intent": "query"},
            route_json={"stage": "finished"},
            status="succeeded",
        )
        assert updated.ontology_json["tenant_id"] == "tenant-a"
        assert updated.route_json["tenant_id"] == "tenant-a"

        with pytest.raises(ValueError, match="tenant_id 与业务上下文冲突"):
            service.create_run(
                agent="orchestrator",
                source="user_message",
                tenant_id="tenant-a",
                route_json={"tenant_id": "tenant-b"},
                status="running",
            )