mirror of
https://github.com/dograh-hq/dograh.git
synced 2026-07-13 11:22:14 +02:00
f32395f859
* fix(auth): allow invited org members to start workflow runs Users invited to an org could not start workflows belonging to that org because the authorization check compared actor.selected_organization_id directly against workflow.organization_id. An invited user's selected org correctly reflects the invited org, but if the Stack Auth token resolves to a different org id than expected the strict equality fails. Per api/AGENTS.md: "Whenever you read or write an organization-scoped field, you must filter or validate by organization_id." The correct policy is org membership, not selected-org identity. - Add is_user_member_of_organization() to OrganizationClient; queries the organization_users association table directly (no lazy-load risk). - Replace the identity check in authorize_workflow_run_start() with a membership lookup. Deny when actor_user.id is not in the org's member set; error_code stays workflow_not_found to avoid leaking existence. - Update test: rename rejects_actor_from_another_org to rejects_actor_not_a_member (reflects actual policy), add positive test allows_invited_member that seeds membership and asserts has_quota=True. Closes #491 * fix(auth): skip membership check for personal workflows (organization_id=None) When workflow.organization_id is None (personal or legacy workflow with no org), the membership lookup was still called, producing a SQL IS NULL comparison that matched nothing and denied the run. Guard the check so it only runs when the workflow is org-scoped. Adds a regression test confirming that an actor with a known id can start a personal workflow without triggering is_user_member_of_organization. * fix(auth): fail closed on workflow membership lookup errors --------- Co-authored-by: Abhishek Kumar <abhishek@a6k.me>
507 lines
15 KiB
Python
507 lines
15 KiB
Python
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
from api.services import quota_service
|
|
from api.services.configuration.registry import ServiceProviders
|
|
from api.services.managed_model_services import MPS_CORRELATION_ID_CONTEXT_KEY
|
|
from api.services.quota_service import QuotaCheckResult
|
|
|
|
|
|
def _dograh_config(
|
|
api_key: str = "mps_sk_12345678",
|
|
*,
|
|
managed_service_version: int = 1,
|
|
):
|
|
return SimpleNamespace(
|
|
managed_service_version=managed_service_version,
|
|
llm=SimpleNamespace(provider=ServiceProviders.DOGRAH, api_key=api_key),
|
|
stt=None,
|
|
tts=None,
|
|
embeddings=None,
|
|
)
|
|
|
|
|
|
def _byok_config():
|
|
return SimpleNamespace(
|
|
managed_service_version=2,
|
|
llm=SimpleNamespace(provider="openai", api_key="sk-openai"),
|
|
stt=None,
|
|
tts=None,
|
|
embeddings=None,
|
|
)
|
|
|
|
|
|
def _workflow():
|
|
return SimpleNamespace(
|
|
id=7,
|
|
user_id=123,
|
|
organization_id=42,
|
|
workflow_configurations={"model_overrides": {}},
|
|
)
|
|
|
|
|
|
def _workflow_owner():
|
|
return SimpleNamespace(
|
|
id=123,
|
|
provider_id="provider-123",
|
|
)
|
|
|
|
|
|
def _actor():
|
|
return SimpleNamespace(
|
|
id=456,
|
|
provider_id="actor-456",
|
|
selected_organization_id=42,
|
|
)
|
|
|
|
|
|
def _patch_workflow_context(monkeypatch, *, workflow=None, owner=None):
|
|
monkeypatch.setattr(
|
|
quota_service.db_client,
|
|
"get_workflow_by_id",
|
|
AsyncMock(return_value=workflow or _workflow()),
|
|
)
|
|
monkeypatch.setattr(
|
|
quota_service.db_client,
|
|
"get_user_by_id",
|
|
AsyncMock(return_value=owner or _workflow_owner()),
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_authorize_workflow_run_uses_workflow_org_for_hosted_v2(
|
|
monkeypatch,
|
|
):
|
|
get_config = AsyncMock(return_value=_dograh_config())
|
|
authorize = AsyncMock(
|
|
return_value={
|
|
"allowed": True,
|
|
"billing_mode": "v2",
|
|
"remaining_credits": "25.0000",
|
|
}
|
|
)
|
|
check_usage = AsyncMock()
|
|
|
|
monkeypatch.setattr(quota_service, "DEPLOYMENT_MODE", "saas")
|
|
_patch_workflow_context(monkeypatch)
|
|
monkeypatch.setattr(
|
|
quota_service,
|
|
"get_effective_ai_model_configuration_for_workflow",
|
|
get_config,
|
|
)
|
|
monkeypatch.setattr(
|
|
quota_service.mps_service_key_client,
|
|
"authorize_workflow_run_start",
|
|
authorize,
|
|
)
|
|
monkeypatch.setattr(
|
|
quota_service.mps_service_key_client,
|
|
"check_service_key_usage",
|
|
check_usage,
|
|
)
|
|
|
|
result = await quota_service.authorize_workflow_run_start(workflow_id=7)
|
|
|
|
assert result.has_quota is True
|
|
get_config.assert_awaited_once_with(
|
|
user_id=123,
|
|
organization_id=42,
|
|
workflow_configurations={"model_overrides": {}},
|
|
)
|
|
authorize.assert_awaited_once_with(
|
|
organization_id=42,
|
|
workflow_run_id=None,
|
|
service_key=None,
|
|
require_correlation_id=False,
|
|
minimum_credits=quota_service.MINIMUM_DOGRAH_CREDITS_FOR_CALL,
|
|
created_by="provider-123",
|
|
metadata={"dograh_user_id": "123", "workflow_id": 7},
|
|
)
|
|
check_usage.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_authorize_workflow_run_v2_insufficient_credits_prompts_billing(
|
|
monkeypatch,
|
|
):
|
|
get_config = AsyncMock(return_value=_byok_config())
|
|
authorize = AsyncMock(
|
|
return_value={
|
|
"allowed": False,
|
|
"billing_mode": "v2",
|
|
"remaining_credits": "0.0000",
|
|
"error": "insufficient_credits",
|
|
}
|
|
)
|
|
check_usage = AsyncMock()
|
|
|
|
monkeypatch.setattr(quota_service, "DEPLOYMENT_MODE", "saas")
|
|
_patch_workflow_context(monkeypatch)
|
|
monkeypatch.setattr(
|
|
quota_service,
|
|
"get_effective_ai_model_configuration_for_workflow",
|
|
get_config,
|
|
)
|
|
monkeypatch.setattr(
|
|
quota_service.mps_service_key_client,
|
|
"authorize_workflow_run_start",
|
|
authorize,
|
|
)
|
|
monkeypatch.setattr(
|
|
quota_service.mps_service_key_client,
|
|
"check_service_key_usage",
|
|
check_usage,
|
|
)
|
|
|
|
result = await quota_service.authorize_workflow_run_start(workflow_id=7)
|
|
|
|
assert result.has_quota is False
|
|
assert result.error_code == "insufficient_credits"
|
|
assert "/billing" in result.error_message
|
|
assert "founders@dograh.com" not in result.error_message
|
|
authorize.assert_awaited_once()
|
|
check_usage.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_authorize_workflow_run_oss_exhausted_key_blocks_run(
|
|
monkeypatch,
|
|
):
|
|
api_key = "mps_sk_12345678"
|
|
get_config = AsyncMock(return_value=_dograh_config(api_key))
|
|
check_usage = AsyncMock(
|
|
return_value={"total_credits_used": 500.0, "remaining_credits": 0.0}
|
|
)
|
|
|
|
monkeypatch.setattr(quota_service, "DEPLOYMENT_MODE", "oss")
|
|
_patch_workflow_context(monkeypatch)
|
|
monkeypatch.setattr(
|
|
quota_service,
|
|
"get_effective_ai_model_configuration_for_workflow",
|
|
get_config,
|
|
)
|
|
monkeypatch.setattr(
|
|
quota_service.mps_service_key_client,
|
|
"check_service_key_usage",
|
|
check_usage,
|
|
)
|
|
|
|
result = await quota_service.authorize_workflow_run_start(workflow_id=7)
|
|
|
|
assert result.has_quota is False
|
|
assert result.error_code == "quota_exceeded"
|
|
assert "founders@dograh.com" in result.error_message
|
|
assert "/billing" not in result.error_message
|
|
check_usage.assert_awaited_once_with(api_key)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_authorize_workflow_run_managed_v2_stores_hosted_correlation(
|
|
monkeypatch,
|
|
):
|
|
api_key = "mps_sk_12345678"
|
|
workflow_run = SimpleNamespace(initial_context={"existing": "value"})
|
|
get_config = AsyncMock(
|
|
return_value=_dograh_config(api_key, managed_service_version=2)
|
|
)
|
|
authorize = AsyncMock(
|
|
return_value={
|
|
"allowed": True,
|
|
"billing_mode": "v2",
|
|
"remaining_credits": "25.0000",
|
|
"correlation_id": "mps-corr-123",
|
|
}
|
|
)
|
|
update_workflow_run = AsyncMock()
|
|
|
|
monkeypatch.setattr(quota_service, "DEPLOYMENT_MODE", "saas")
|
|
_patch_workflow_context(monkeypatch)
|
|
monkeypatch.setattr(
|
|
quota_service.db_client,
|
|
"get_workflow_run_by_id",
|
|
AsyncMock(return_value=workflow_run),
|
|
)
|
|
monkeypatch.setattr(
|
|
quota_service.db_client,
|
|
"update_workflow_run",
|
|
update_workflow_run,
|
|
)
|
|
monkeypatch.setattr(
|
|
quota_service,
|
|
"get_effective_ai_model_configuration_for_workflow",
|
|
get_config,
|
|
)
|
|
monkeypatch.setattr(
|
|
quota_service.mps_service_key_client,
|
|
"authorize_workflow_run_start",
|
|
authorize,
|
|
)
|
|
monkeypatch.setattr(
|
|
quota_service.mps_service_key_client,
|
|
"check_service_key_usage",
|
|
AsyncMock(),
|
|
)
|
|
|
|
result = await quota_service.authorize_workflow_run_start(
|
|
workflow_id=7,
|
|
workflow_run_id=88,
|
|
)
|
|
|
|
assert result.has_quota is True
|
|
authorize.assert_awaited_once_with(
|
|
organization_id=42,
|
|
workflow_run_id=88,
|
|
service_key=api_key,
|
|
require_correlation_id=True,
|
|
minimum_credits=quota_service.MINIMUM_DOGRAH_CREDITS_FOR_CALL,
|
|
created_by="provider-123",
|
|
metadata={"dograh_user_id": "123", "workflow_id": 7},
|
|
)
|
|
update_workflow_run.assert_awaited_once_with(
|
|
88,
|
|
initial_context={
|
|
"existing": "value",
|
|
MPS_CORRELATION_ID_CONTEXT_KEY: "mps-corr-123",
|
|
},
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_authorize_workflow_run_service_token_from_wrong_org_prompts_new_token(
|
|
monkeypatch,
|
|
):
|
|
api_key = "mps_sk_12345678"
|
|
get_config = AsyncMock(
|
|
return_value=_dograh_config(api_key, managed_service_version=2)
|
|
)
|
|
request = httpx.Request(
|
|
"POST",
|
|
"http://localhost:8004/api/v1/billing/accounts/42/run-authorization",
|
|
)
|
|
response = httpx.Response(
|
|
403,
|
|
json={"detail": "Service key organization mismatch"},
|
|
request=request,
|
|
)
|
|
authorize = AsyncMock(
|
|
side_effect=httpx.HTTPStatusError(
|
|
"Failed to authorize MPS workflow run start",
|
|
request=request,
|
|
response=response,
|
|
)
|
|
)
|
|
|
|
monkeypatch.setattr(quota_service, "DEPLOYMENT_MODE", "saas")
|
|
_patch_workflow_context(monkeypatch)
|
|
monkeypatch.setattr(
|
|
quota_service,
|
|
"get_effective_ai_model_configuration_for_workflow",
|
|
get_config,
|
|
)
|
|
monkeypatch.setattr(
|
|
quota_service.mps_service_key_client,
|
|
"authorize_workflow_run_start",
|
|
authorize,
|
|
)
|
|
monkeypatch.setattr(
|
|
quota_service.mps_service_key_client,
|
|
"check_service_key_usage",
|
|
AsyncMock(),
|
|
)
|
|
|
|
result = await quota_service.authorize_workflow_run_start(
|
|
workflow_id=7,
|
|
workflow_run_id=88,
|
|
)
|
|
|
|
assert result.has_quota is False
|
|
assert result.error_code == "service_key_org_mismatch"
|
|
assert result.error_message == quota_service.SERVICE_TOKEN_ORG_MISMATCH_MESSAGE
|
|
assert "new service token from the Developers tab" in result.error_message
|
|
authorize.assert_awaited_once_with(
|
|
organization_id=42,
|
|
workflow_run_id=88,
|
|
service_key=api_key,
|
|
require_correlation_id=True,
|
|
minimum_credits=quota_service.MINIMUM_DOGRAH_CREDITS_FOR_CALL,
|
|
created_by="provider-123",
|
|
metadata={"dograh_user_id": "123", "workflow_id": 7},
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_authorize_workflow_run_oss_uses_key_paths_not_workflow_org(
|
|
monkeypatch,
|
|
):
|
|
api_key = "mps_sk_12345678"
|
|
workflow_run = SimpleNamespace(initial_context={})
|
|
get_config = AsyncMock(
|
|
return_value=_dograh_config(api_key, managed_service_version=2)
|
|
)
|
|
hosted_authorize = AsyncMock()
|
|
check_usage = AsyncMock(
|
|
return_value={"total_credits_used": 1.0, "remaining_credits": 499.0}
|
|
)
|
|
create_correlation = AsyncMock(return_value={"correlation_id": "oss-corr-123"})
|
|
update_workflow_run = AsyncMock()
|
|
|
|
monkeypatch.setattr(quota_service, "DEPLOYMENT_MODE", "oss")
|
|
_patch_workflow_context(monkeypatch)
|
|
monkeypatch.setattr(
|
|
quota_service.db_client,
|
|
"get_workflow_run_by_id",
|
|
AsyncMock(return_value=workflow_run),
|
|
)
|
|
monkeypatch.setattr(
|
|
quota_service.db_client,
|
|
"update_workflow_run",
|
|
update_workflow_run,
|
|
)
|
|
monkeypatch.setattr(
|
|
quota_service,
|
|
"get_effective_ai_model_configuration_for_workflow",
|
|
get_config,
|
|
)
|
|
monkeypatch.setattr(
|
|
quota_service.mps_service_key_client,
|
|
"authorize_workflow_run_start",
|
|
hosted_authorize,
|
|
)
|
|
monkeypatch.setattr(
|
|
quota_service.mps_service_key_client,
|
|
"check_service_key_usage",
|
|
check_usage,
|
|
)
|
|
monkeypatch.setattr(
|
|
quota_service.mps_service_key_client,
|
|
"create_correlation_id",
|
|
create_correlation,
|
|
)
|
|
|
|
result = await quota_service.authorize_workflow_run_start(
|
|
workflow_id=7,
|
|
workflow_run_id=88,
|
|
)
|
|
|
|
assert result.has_quota is True
|
|
hosted_authorize.assert_not_awaited()
|
|
check_usage.assert_awaited_once_with(api_key)
|
|
create_correlation.assert_awaited_once_with(
|
|
service_key=api_key,
|
|
workflow_run_id=88,
|
|
)
|
|
update_workflow_run.assert_awaited_once_with(
|
|
88,
|
|
initial_context={MPS_CORRELATION_ID_CONTEXT_KEY: "oss-corr-123"},
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_authorize_workflow_run_rejects_actor_not_a_member(monkeypatch):
|
|
monkeypatch.setattr(quota_service, "DEPLOYMENT_MODE", "saas")
|
|
_patch_workflow_context(monkeypatch)
|
|
monkeypatch.setattr(
|
|
quota_service.db_client,
|
|
"is_user_member_of_organization",
|
|
AsyncMock(return_value=False),
|
|
)
|
|
|
|
result = await quota_service.authorize_workflow_run_start(
|
|
workflow_id=7,
|
|
actor_user=SimpleNamespace(id=456, selected_organization_id=999),
|
|
)
|
|
|
|
assert result.has_quota is False
|
|
assert result.error_code == "workflow_not_found"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_authorize_workflow_run_membership_lookup_error_fails_closed(monkeypatch):
|
|
monkeypatch.setattr(quota_service, "DEPLOYMENT_MODE", "saas")
|
|
_patch_workflow_context(monkeypatch)
|
|
monkeypatch.setattr(
|
|
quota_service.db_client,
|
|
"is_user_member_of_organization",
|
|
AsyncMock(side_effect=RuntimeError("database unavailable")),
|
|
)
|
|
|
|
result = await quota_service.authorize_workflow_run_start(
|
|
workflow_id=7,
|
|
actor_user=SimpleNamespace(id=456, selected_organization_id=42),
|
|
)
|
|
|
|
assert result.has_quota is False
|
|
assert result.error_code == "workflow_not_found"
|
|
quota_service.db_client.get_user_by_id.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_authorize_workflow_run_allows_invited_member(monkeypatch):
|
|
"""User invited to an org can start workflows belonging to that org."""
|
|
monkeypatch.setattr(quota_service, "DEPLOYMENT_MODE", "saas")
|
|
_patch_workflow_context(monkeypatch)
|
|
monkeypatch.setattr(
|
|
quota_service.db_client,
|
|
"is_user_member_of_organization",
|
|
AsyncMock(return_value=True),
|
|
)
|
|
monkeypatch.setattr(
|
|
quota_service,
|
|
"get_effective_ai_model_configuration_for_workflow",
|
|
AsyncMock(return_value=_byok_config()),
|
|
)
|
|
hosted_authorize = AsyncMock(return_value=QuotaCheckResult(has_quota=True))
|
|
monkeypatch.setattr(
|
|
quota_service,
|
|
"_authorize_hosted_workflow_run_start",
|
|
hosted_authorize,
|
|
)
|
|
|
|
# actor_user.selected_organization_id=999 differs from workflow.organization_id=42,
|
|
# but is_user_member_of_organization returns True so the run should be allowed.
|
|
result = await quota_service.authorize_workflow_run_start(
|
|
workflow_id=7,
|
|
actor_user=SimpleNamespace(id=456, selected_organization_id=999),
|
|
)
|
|
|
|
assert result.has_quota is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_authorize_workflow_run_allows_personal_workflow_with_actor(monkeypatch):
|
|
"""Personal/legacy workflows (organization_id=None) bypass membership check."""
|
|
personal_workflow = SimpleNamespace(
|
|
id=7,
|
|
user_id=123,
|
|
organization_id=None,
|
|
workflow_configurations={"model_overrides": {}},
|
|
)
|
|
monkeypatch.setattr(quota_service, "DEPLOYMENT_MODE", "saas")
|
|
_patch_workflow_context(monkeypatch, workflow=personal_workflow)
|
|
is_member_mock = AsyncMock()
|
|
monkeypatch.setattr(
|
|
quota_service.db_client,
|
|
"is_user_member_of_organization",
|
|
is_member_mock,
|
|
)
|
|
monkeypatch.setattr(
|
|
quota_service,
|
|
"get_effective_ai_model_configuration_for_workflow",
|
|
AsyncMock(return_value=_byok_config()),
|
|
)
|
|
monkeypatch.setattr(
|
|
quota_service,
|
|
"_authorize_hosted_workflow_run_start",
|
|
AsyncMock(return_value=QuotaCheckResult(has_quota=True)),
|
|
)
|
|
|
|
result = await quota_service.authorize_workflow_run_start(
|
|
workflow_id=7,
|
|
actor_user=SimpleNamespace(id=456),
|
|
)
|
|
|
|
assert result.has_quota is True
|
|
is_member_mock.assert_not_awaited()
|