feat(platform): close AI expense value loop
Add tenant-safe value, telemetry, connector, commercial, and production-readiness foundations.
This commit is contained in:
caoxiaozhu
301
server/tests/test_agent_run_tenant_security.py
Normal file
301
server/tests/test_agent_run_tenant_security.py
Normal file
@@ -0,0 +1,301 @@
|
||||
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",
|
||||
)
|
||||
Reference in New Issue
Block a user