- Prepaid credit system: org wallet, debit on usage, chat gate, admin credit console - Engine UI: /manager, /tasks, /tasks/[id], /workforce, /dashboard, /traces pages - Direct task creation (HTTP POST /tasks, MCP tasks_create) at pending_manager_review - Pause/resume/cancel/answer actions in engine API and UI - Employee cap (7 per org default), review pipeline guard (active reviewer required) - Scheduler: wide window prevents org starvation; archived employees excluded from cap - Tenant isolation: project-id cross-org 404, engine never uses another tenant's sandbox - Order route fix: manager built with pipeline states + employee profiles - Contract extension: ModelGrantPort.get_org_policy with DefaultOrgPolicy - Engine event types in @agy/contracts (task.step, task.gate, run.log, run.tool, team.changed) - api-engine.ts typed client; api-credit.ts client - Alembic migration c4d5e6f7a8b9: prepaid credit + engine UI tables - 859 tests passed, 1 skipped; ruff/mypy/tsc/eslint/vitest all green Refs: PLAN-v2.md §9 (T1–T43), todo-v2-final.md
403 lines
15 KiB
Python
403 lines
15 KiB
Python
"""Owner model grants, the default set and org policy.
|
||
|
||
The invariant under test: **an organization may only use the models the owner
|
||
granted it.** Two switches have to agree (model active + grant active), which is
|
||
what makes revocation real, and a grant belongs to exactly one organization.
|
||
|
||
Nothing here depends on the engine's code: the checks go through the HTTP
|
||
endpoints *and* through ``ModelGrantPort`` as the engine would call it, so a
|
||
break in either path fails the suite.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import pytest
|
||
from httpx import AsyncClient
|
||
from sqlalchemy import select
|
||
|
||
from app.application.container_org import build_model_grant_service, build_model_grants
|
||
from app.application.engine.manager import BOOTSTRAP_TIERS
|
||
from app.core.config import settings
|
||
from app.infrastructure.db.models import AiModelRow, AuditEventRow, OrgModelGrantRow
|
||
from app.infrastructure.db.models_org import DefaultModelGrantRow, OrgPolicyRow
|
||
from tests.conftest import auth_header
|
||
|
||
GRANTS = "/api/admin/model-grants"
|
||
DEFAULTS = "/api/admin/default-model-grants"
|
||
POLICIES = "/api/admin/org-policies"
|
||
|
||
|
||
def _payload(name: str, slug: str) -> dict:
|
||
return {
|
||
"org_name": name,
|
||
"slug": slug,
|
||
"app_name": f"app-{slug}",
|
||
"app_type": "web_app",
|
||
"language": "node",
|
||
"framework": "next",
|
||
"git_source": "platform",
|
||
}
|
||
|
||
|
||
async def _create_org(client: AsyncClient, token: str, name: str, slug: str) -> str:
|
||
response = await client.post(
|
||
"/api/setup/complete", json=_payload(name, slug), headers=auth_header(token)
|
||
)
|
||
assert response.status_code == 201, response.text
|
||
return response.json()["organization"]["id"]
|
||
|
||
|
||
async def _model_ids(session_factory) -> list[str]:
|
||
async with session_factory() as session:
|
||
rows = (
|
||
await session.execute(select(AiModelRow).order_by(AiModelRow.internal_name))
|
||
).scalars().all()
|
||
return [row.id for row in rows]
|
||
|
||
|
||
async def _audit_actions(session_factory, organization_id: str | None = None) -> list[str]:
|
||
async with session_factory() as session:
|
||
stmt = select(AuditEventRow.action)
|
||
if organization_id is not None:
|
||
stmt = stmt.where(AuditEventRow.organization_id == organization_id)
|
||
return list((await session.execute(stmt)).scalars().all())
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Authorization
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
("method", "path"),
|
||
[
|
||
("get", f"{GRANTS}/org-1"),
|
||
("get", f"{GRANTS}/org-1/active"),
|
||
("put", f"{GRANTS}/org-1/models/model-1"),
|
||
("delete", f"{GRANTS}/org-1/models/model-1"),
|
||
("get", DEFAULTS),
|
||
("put", f"{DEFAULTS}/model-1"),
|
||
("delete", f"{DEFAULTS}/model-1"),
|
||
("get", f"{POLICIES}/org-1"),
|
||
("put", f"{POLICIES}/org-1"),
|
||
],
|
||
)
|
||
async def test_non_owner_is_forbidden_everywhere(
|
||
client: AsyncClient, manager_token: str, method: str, path: str
|
||
):
|
||
call = getattr(client, method)
|
||
kwargs: dict = {"headers": auth_header(manager_token)}
|
||
if method == "put":
|
||
kwargs["json"] = {"daily_credit_cap": 5, "is_active": True, "max_employees": 3}
|
||
response = await call(path, **kwargs)
|
||
assert response.status_code == 403, response.text
|
||
|
||
|
||
async def test_anonymous_is_rejected(client: AsyncClient):
|
||
assert (await client.get(f"{GRANTS}/org-1")).status_code == 401
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Grant matrix and revocation
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
async def test_wizard_copies_the_default_set(
|
||
client: AsyncClient,
|
||
manager_token: str,
|
||
platform_owner_token: str,
|
||
session_factory,
|
||
seeded,
|
||
):
|
||
model_id = (await _model_ids(session_factory))[0]
|
||
set_default = await client.put(
|
||
f"{DEFAULTS}/{model_id}",
|
||
json={"is_active": True, "daily_credit_cap": 12.5},
|
||
headers=auth_header(platform_owner_token),
|
||
)
|
||
assert set_default.status_code == 200, set_default.text
|
||
assert set_default.json()["granted"] is True
|
||
|
||
org_id = await _create_org(client, manager_token, "سازمان الف", "alpha-org")
|
||
|
||
matrix = await client.get(f"{GRANTS}/{org_id}", headers=auth_header(platform_owner_token))
|
||
assert matrix.status_code == 200, matrix.text
|
||
rows = {m["ai_model_id"]: m for m in matrix.json()["models"]}
|
||
assert rows[model_id]["granted"] is True
|
||
assert rows[model_id]["daily_credit_cap"] == 12.5
|
||
# Every other model is listed but not granted.
|
||
assert all(m["granted"] is False for m_id, m in rows.items() if m_id != model_id)
|
||
|
||
# The same answer through the engine's port.
|
||
async with session_factory() as session:
|
||
service = build_model_grants(session)
|
||
granted = await service.granted_models(organization_id=org_id)
|
||
assert [g.ai_model_id for g in granted] == [model_id]
|
||
assert granted[0].tier in {"fast", "balanced", "strong"}
|
||
assert granted[0].daily_credit_cap == 12.5
|
||
assert await service.is_granted(organization_id=org_id, ai_model_id=model_id)
|
||
|
||
|
||
async def test_grant_edit_and_revoke_are_immediate_and_audited(
|
||
client: AsyncClient,
|
||
manager_token: str,
|
||
platform_owner_token: str,
|
||
session_factory,
|
||
seeded,
|
||
):
|
||
model_id = (await _model_ids(session_factory))[0]
|
||
org_id = await _create_org(client, manager_token, "سازمان ب", "beta-org")
|
||
owner = auth_header(platform_owner_token)
|
||
|
||
grant = await client.put(
|
||
f"{GRANTS}/{org_id}/models/{model_id}",
|
||
json={"daily_credit_cap": 3},
|
||
headers=owner,
|
||
)
|
||
assert grant.status_code == 200, grant.text
|
||
assert grant.json()["is_active"] is True
|
||
|
||
async with session_factory() as session:
|
||
service = build_model_grant_service(session)
|
||
assert await service.is_granted(organization_id=org_id, ai_model_id=model_id)
|
||
|
||
# Raising the cap keeps it granted.
|
||
raised = await client.put(
|
||
f"{GRANTS}/{org_id}/models/{model_id}",
|
||
json={"daily_credit_cap": 40},
|
||
headers=owner,
|
||
)
|
||
assert raised.json()["daily_credit_cap"] == 40
|
||
|
||
revoked = await client.delete(f"{GRANTS}/{org_id}/models/{model_id}", headers=owner)
|
||
assert revoked.status_code == 200, revoked.text
|
||
|
||
async with session_factory() as session:
|
||
service = build_model_grant_service(session)
|
||
assert not await service.is_granted(organization_id=org_id, ai_model_id=model_id)
|
||
assert await service.granted_models(organization_id=org_id) == []
|
||
# Revocation deactivates rather than deletes, so the history survives and
|
||
# re-granting is idempotent instead of racing a unique constraint.
|
||
row = (
|
||
await session.execute(
|
||
select(OrgModelGrantRow).where(
|
||
OrgModelGrantRow.organization_id == org_id,
|
||
OrgModelGrantRow.ai_model_id == model_id,
|
||
)
|
||
)
|
||
).scalar_one()
|
||
assert row.is_active is False
|
||
|
||
active = await client.get(f"{GRANTS}/{org_id}/active", headers=owner)
|
||
assert active.json() == []
|
||
|
||
actions = await _audit_actions(session_factory, org_id)
|
||
assert "model_grant.set" in actions
|
||
assert "model_grant.revoke" in actions
|
||
|
||
# Revoking something that is not granted is a 404, not a silent success.
|
||
again = await client.delete(f"{GRANTS}/{org_id}/models/{model_id}", headers=owner)
|
||
assert again.status_code == 404
|
||
|
||
|
||
async def test_inactive_model_is_never_granted(
|
||
client: AsyncClient,
|
||
manager_token: str,
|
||
platform_owner_token: str,
|
||
session_factory,
|
||
seeded,
|
||
):
|
||
model_id = (await _model_ids(session_factory))[0]
|
||
org_id = await _create_org(client, manager_token, "سازمان ج", "gamma-org")
|
||
owner = auth_header(platform_owner_token)
|
||
granted = await client.put(f"{GRANTS}/{org_id}/models/{model_id}", json={}, headers=owner)
|
||
assert granted.status_code == 200, granted.text
|
||
|
||
# The owner retires the model from the catalogue.
|
||
async with session_factory() as session:
|
||
row = await session.get(AiModelRow, model_id)
|
||
assert row is not None
|
||
row.is_active = False
|
||
await session.commit()
|
||
|
||
async with session_factory() as session:
|
||
service = build_model_grant_service(session)
|
||
assert not await service.is_granted(organization_id=org_id, ai_model_id=model_id)
|
||
assert await service.granted_models(organization_id=org_id) == []
|
||
|
||
# Granting a retired model is refused rather than quietly stored.
|
||
refused = await client.put(f"{GRANTS}/{org_id}/models/{model_id}", json={}, headers=owner)
|
||
assert refused.status_code == 404
|
||
listing = await client.get(f"{GRANTS}/{org_id}", headers=owner)
|
||
assert all(m["ai_model_id"] != model_id for m in listing.json()["models"])
|
||
|
||
|
||
async def test_grants_are_isolated_per_organization(
|
||
client: AsyncClient,
|
||
manager_token: str,
|
||
platform_owner_token: str,
|
||
session_factory,
|
||
seeded,
|
||
):
|
||
model_id = (await _model_ids(session_factory))[0]
|
||
org_a = await _create_org(client, manager_token, "شرکت آلفا", "alpha-two")
|
||
org_b = await _create_org(client, manager_token, "شرکت بتا", "beta-two")
|
||
owner = auth_header(platform_owner_token)
|
||
|
||
granted = await client.put(
|
||
f"{GRANTS}/{org_a}/models/{model_id}", json={"daily_credit_cap": 7}, headers=owner
|
||
)
|
||
assert granted.status_code == 200, granted.text
|
||
|
||
async with session_factory() as session:
|
||
service = build_model_grant_service(session)
|
||
assert await service.is_granted(organization_id=org_a, ai_model_id=model_id)
|
||
assert not await service.is_granted(organization_id=org_b, ai_model_id=model_id)
|
||
assert await service.granted_models(organization_id=org_b) == []
|
||
|
||
matrix_b = await client.get(f"{GRANTS}/{org_b}", headers=owner)
|
||
assert all(m["granted"] is False for m in matrix_b.json()["models"])
|
||
# …and org B's grant is never visible in org A's active list either.
|
||
active_a = await client.get(f"{GRANTS}/{org_a}/active", headers=owner)
|
||
assert [g["ai_model_id"] for g in active_a.json()] == [model_id]
|
||
|
||
|
||
async def test_default_set_can_be_removed_and_deactivated(
|
||
client: AsyncClient,
|
||
manager_token: str,
|
||
platform_owner_token: str,
|
||
session_factory,
|
||
seeded,
|
||
):
|
||
model_id = (await _model_ids(session_factory))[0]
|
||
owner = auth_header(platform_owner_token)
|
||
|
||
assert (await client.put(f"{DEFAULTS}/{model_id}", json={}, headers=owner)).status_code == 200
|
||
# Switched off: still in the set for the admin UI, but not copied to new orgs.
|
||
off = await client.put(
|
||
f"{DEFAULTS}/{model_id}", json={"is_active": False}, headers=owner
|
||
)
|
||
assert off.status_code == 200, off.text
|
||
|
||
org_off = await _create_org(client, manager_token, "سازمان خاموش", "off-org")
|
||
async with session_factory() as session:
|
||
service = build_model_grant_service(session)
|
||
assert await service.granted_models(organization_id=org_off) == []
|
||
|
||
removed = await client.delete(f"{DEFAULTS}/{model_id}", headers=owner)
|
||
assert removed.status_code == 200, removed.text
|
||
assert removed.json()["in_default_set"] is False
|
||
|
||
async with session_factory() as session:
|
||
assert (
|
||
await session.execute(
|
||
select(DefaultModelGrantRow).where(
|
||
DefaultModelGrantRow.ai_model_id == model_id
|
||
)
|
||
)
|
||
).scalar_one_or_none() is None
|
||
|
||
# Already-created organizations keep the grant they were given.
|
||
async with session_factory() as session:
|
||
assert (
|
||
await session.execute(
|
||
select(OrgModelGrantRow).where(
|
||
OrgModelGrantRow.organization_id == org_off
|
||
)
|
||
)
|
||
).scalar_one_or_none() is None
|
||
|
||
actions = await _audit_actions(session_factory)
|
||
assert "default_model_grant.set" in actions
|
||
assert "default_model_grant.remove" in actions
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Org policy
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
async def test_policy_defaults_caps_and_audit(
|
||
client: AsyncClient,
|
||
manager_token: str,
|
||
platform_owner_token: str,
|
||
session_factory,
|
||
seeded,
|
||
):
|
||
"""D-A: the default cap is 7 employees per organization (it was 5).
|
||
|
||
The starting team (``BOOTSTRAP_TIERS``, 4 employees) must fit under that default.
|
||
"""
|
||
org_id = await _create_org(client, manager_token, "سازمان د", "delta-org")
|
||
owner = auth_header(platform_owner_token)
|
||
|
||
fresh = await client.get(f"{POLICIES}/{org_id}", headers=owner)
|
||
assert fresh.status_code == 200, fresh.text
|
||
body = fresh.json()
|
||
assert body["max_employees"] == 7
|
||
assert len(BOOTSTRAP_TIERS) <= body["max_employees"]
|
||
assert body["automation_enabled"] is True
|
||
assert body["mcp_enabled"] is True
|
||
assert body["monthly_credit_budget"] is None
|
||
assert body["max_employees_limit"] == settings.engine_max_employees
|
||
|
||
too_many = await client.put(
|
||
f"{POLICIES}/{org_id}",
|
||
json={"max_employees": settings.engine_max_employees + 1},
|
||
headers=owner,
|
||
)
|
||
assert too_many.status_code == 422
|
||
|
||
negative = await client.put(
|
||
f"{POLICIES}/{org_id}", json={"monthly_credit_budget": -1}, headers=owner
|
||
)
|
||
assert negative.status_code == 422
|
||
|
||
empty = await client.put(f"{POLICIES}/{org_id}", json={}, headers=owner)
|
||
assert empty.status_code == 422
|
||
|
||
updated = await client.put(
|
||
f"{POLICIES}/{org_id}",
|
||
json={
|
||
"max_employees": 8,
|
||
"monthly_credit_budget": 250.5,
|
||
"automation_enabled": False,
|
||
"mcp_enabled": False,
|
||
},
|
||
headers=owner,
|
||
)
|
||
assert updated.status_code == 200, updated.text
|
||
assert updated.json()["max_employees"] == 8
|
||
assert updated.json()["automation_enabled"] is False
|
||
assert updated.json()["mcp_enabled"] is False
|
||
|
||
async with session_factory() as session:
|
||
row = await session.get(OrgPolicyRow, org_id)
|
||
assert row is not None
|
||
assert row.max_employees == 8
|
||
assert row.monthly_credit_budget == 250.5
|
||
# Agent 2 reads the policy through the same service.
|
||
service = build_model_grant_service(session)
|
||
policy = await service.get_org_policy(org_id)
|
||
assert policy.max_employees == 8
|
||
assert policy.automation_enabled is False
|
||
assert await service.mcp_enabled(org_id) is False
|
||
|
||
assert "org_policy.update" in await _audit_actions(session_factory, org_id)
|
||
|
||
|
||
async def test_policy_is_isolated_per_organization(
|
||
client: AsyncClient, manager_token: str, platform_owner_token: str, seeded
|
||
):
|
||
"""D-A: a cap written for one organization leaves the other at the default of 7."""
|
||
org_a = await _create_org(client, manager_token, "سازمان ه", "epsilon-org")
|
||
org_b = await _create_org(client, manager_token, "سازمان و", "zeta-org")
|
||
owner = auth_header(platform_owner_token)
|
||
|
||
await client.put(f"{POLICIES}/{org_a}", json={"max_employees": 2}, headers=owner)
|
||
|
||
a = await client.get(f"{POLICIES}/{org_a}", headers=owner)
|
||
b = await client.get(f"{POLICIES}/{org_b}", headers=owner)
|
||
assert a.json()["max_employees"] == 2
|
||
assert b.json()["max_employees"] == 7
|