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", )