diff --git a/backend/app/api/enterprise.py b/backend/app/api/enterprise.py index 86eef7afd..01eb77269 100644 --- a/backend/app/api/enterprise.py +++ b/backend/app/api/enterprise.py @@ -83,6 +83,28 @@ def _is_platform_admin_user(user: User) -> bool: return user.role == "platform_admin" or bool(getattr(getattr(user, "identity", None), "is_platform_admin", False)) +def _llm_management_tenant_id(current_user: User, requested_tenant_id: str | None = None) -> uuid.UUID | None: + """Resolve an LLM-management tenant without letting org admins switch tenants.""" + raw_tenant_id = requested_tenant_id or current_user.tenant_id + if raw_tenant_id is None: + return None + try: + tenant_id = uuid.UUID(str(raw_tenant_id)) + except ValueError as exc: + raise HTTPException(status_code=422, detail="Invalid tenant ID") from exc + if not _is_platform_admin_user(current_user) and tenant_id != current_user.tenant_id: + raise HTTPException(status_code=403, detail="Cannot manage another tenant's models") + return tenant_id + + +def _llm_model_scope(model_id: uuid.UUID, current_user: User): + """Build the tenant-scoped model lookup used by all mutable LLM routes.""" + conditions = [LLMModel.id == model_id, LLMModel.deleted_at.is_(None)] + if not _is_platform_admin_user(current_user): + conditions.append(LLMModel.tenant_id == current_user.tenant_id) + return select(LLMModel).where(*conditions) + + # ─── Public: Check Email Exists ──────────────────────── class CheckEmailRequest(BaseModel): @@ -368,12 +390,7 @@ async def list_llm_models( db: AsyncSession = Depends(get_db), ): """List LLM models scoped to the selected tenant.""" - # Authorization: non-platform admins can only see their own tenant's models - if tenant_id and current_user.role != "platform_admin": - if str(current_user.tenant_id) != tenant_id: - raise HTTPException(status_code=403, detail="Cannot access other tenant's models") - - tid = tenant_id or str(current_user.tenant_id) if current_user.tenant_id else None + tid = _llm_management_tenant_id(current_user, tenant_id) query = ( select(LLMModel) .where(LLMModel.deleted_at.is_(None)) @@ -400,7 +417,7 @@ async def add_llm_model( db: AsyncSession = Depends(get_db), ): """Add a new LLM model to the tenant's pool (admin).""" - tid = tenant_id or (str(current_user.tenant_id) if current_user.tenant_id else None) + tid = _llm_management_tenant_id(current_user, tenant_id) model = LLMModel( provider=data.provider, model=data.model, @@ -413,7 +430,7 @@ async def add_llm_model( supports_vision=data.supports_vision, max_output_tokens=data.max_output_tokens, request_timeout=data.request_timeout, - tenant_id=uuid.UUID(tid) if tid else None, + tenant_id=tid, ) db.add(model) await db.flush() @@ -437,12 +454,7 @@ async def set_default_llm_model( db: AsyncSession = Depends(get_db), ): """Mark this model as the tenant's default for new agents.""" - result = await db.execute( - select(LLMModel).where( - LLMModel.id == model_id, - LLMModel.deleted_at.is_(None), - ) - ) + result = await db.execute(_llm_model_scope(model_id, current_user)) model = result.scalar_one_or_none() if not model: raise HTTPException(status_code=404, detail="Model not found") @@ -526,12 +538,7 @@ async def update_llm_model( db: AsyncSession = Depends(get_db), ): """Update an existing LLM model in the pool (admin).""" - result = await db.execute( - select(LLMModel).where( - LLMModel.id == model_id, - LLMModel.deleted_at.is_(None), - ) - ) + result = await db.execute(_llm_model_scope(model_id, current_user)) model = result.scalar_one_or_none() if not model: raise HTTPException(status_code=404, detail="Model not found") diff --git a/backend/tests/test_llm_model_tenant_scope.py b/backend/tests/test_llm_model_tenant_scope.py new file mode 100644 index 000000000..0f606856c --- /dev/null +++ b/backend/tests/test_llm_model_tenant_scope.py @@ -0,0 +1,42 @@ +import uuid +from types import SimpleNamespace + +import pytest +from fastapi import HTTPException + +from app.api.enterprise import _llm_management_tenant_id, _llm_model_scope + + +def _user(tenant_id: uuid.UUID, role: str = "org_admin") -> SimpleNamespace: + return SimpleNamespace(tenant_id=tenant_id, role=role) + + +def test_org_admin_cannot_select_another_tenant_for_llm_models() -> None: + user = _user(uuid.uuid4()) + + with pytest.raises(HTTPException, match="Cannot manage another tenant's models") as error: + _llm_management_tenant_id(user, str(uuid.uuid4())) + + assert error.value.status_code == 403 + + +def test_platform_admin_can_select_another_tenant_for_llm_models() -> None: + target_tenant_id = uuid.uuid4() + + assert _llm_management_tenant_id(_user(uuid.uuid4(), "platform_admin"), str(target_tenant_id)) == target_tenant_id + + +def test_org_admin_model_mutation_query_is_tenant_scoped() -> None: + tenant_id = uuid.uuid4() + statement = _llm_model_scope(uuid.uuid4(), _user(tenant_id)) + where_clause = " ".join(str(criteria) for criteria in statement._where_criteria) + + assert "llm_models.tenant_id" in where_clause + assert tenant_id.hex in str(statement.compile(compile_kwargs={"literal_binds": True})) + + +def test_platform_admin_model_mutation_query_is_not_tenant_scoped() -> None: + statement = _llm_model_scope(uuid.uuid4(), _user(uuid.uuid4(), "platform_admin")) + where_clause = " ".join(str(criteria) for criteria in statement._where_criteria) + + assert "llm_models.tenant_id" not in where_clause