diff --git a/api/services/quota_service.py b/api/services/quota_service.py index 560c5dc7..05f31b39 100644 --- a/api/services/quota_service.py +++ b/api/services/quota_service.py @@ -7,6 +7,7 @@ across different endpoints (WebRTC signaling, telephony, public API triggers). from dataclasses import dataclass from typing import Any +import httpx from loguru import logger from api.constants import DEPLOYMENT_MODE @@ -25,6 +26,13 @@ from api.services.mps_service_key_client import mps_service_key_client MINIMUM_DOGRAH_CREDITS_FOR_CALL = 0.10 +_MPS_UNREACHABLE_ERRORS = ( + httpx.TimeoutException, + httpx.NetworkError, + httpx.RemoteProtocolError, + httpx.ProxyError, +) + OSS_QUOTA_EXCEEDED_MESSAGE = ( "You have exhausted your trial credits. " "Please sign up on app.dograh.com to create a " @@ -76,6 +84,19 @@ def _insufficient_oss_quota_result() -> QuotaCheckResult: ) +def _mps_unreachable_result( + operation: str, + error: httpx.RequestError, +) -> QuotaCheckResult: + logger.warning( + "MPS unreachable during {}; allowing workflow run to proceed without " + "quota verification: {}", + operation, + error, + ) + return QuotaCheckResult(has_quota=True) + + def _service_uses_dograh(service: Any) -> bool: provider = getattr(service, "provider", None) return ( @@ -195,6 +216,8 @@ async def _authorize_hosted_workflow_run_start( "workflow_id": workflow_id, }, ) + except _MPS_UNREACHABLE_ERRORS as e: + return _mps_unreachable_result("hosted run authorization", e) except Exception as e: logger.warning( "Failed to authorize workflow start with MPS for org {}: {}", @@ -271,6 +294,8 @@ async def _authorize_oss_dograh_keys( f"Dograh quota check passed for key ...{api_key[-8:]}: " f"{remaining:.2f} credits remaining" ) + except _MPS_UNREACHABLE_ERRORS as e: + return _mps_unreachable_result("OSS service-key quota check", e) except Exception as e: logger.error(f"Failed to check quota for Dograh key: {str(e)}") error_str = str(e) @@ -318,6 +343,8 @@ async def _authorize_oss_managed_v2_correlation( workflow_run_id, response.get("correlation_id"), ) + except _MPS_UNREACHABLE_ERRORS as e: + return _mps_unreachable_result("OSS correlation creation", e) except Exception as e: logger.error( "Failed to authorize OSS managed v2 workflow start for workflow {} run {}: {}", @@ -434,6 +461,10 @@ async def authorize_workflow_run_start( error_message="Workflow not found", ) + # A DB read failure here is a "cannot verify" condition, not a + # definitive "not found": let it fall through to the outer handler so + # it fails closed. The None case below is a genuine missing row and keeps + # its specific code. workflow_owner = await db_client.get_user_by_id(workflow.user_id) if not workflow_owner: return QuotaCheckResult( @@ -449,6 +480,9 @@ async def authorize_workflow_run_start( # the definition the run will actually use. workflow_configurations = workflow.workflow_configurations if workflow_run_id is not None: + # As with the owner lookup, a DB read failure falls through to the + # outer fail-closed handler; only a genuinely missing/mismatched run + # returns the specific code below. workflow_run = await db_client.get_workflow_run( workflow_run_id, organization_id=organization_id ) @@ -499,5 +533,12 @@ async def authorize_workflow_run_start( except Exception as e: logger.error(f"Error during quota check: {str(e)}") - # On unexpected error, allow the call to proceed - return QuotaCheckResult(has_quota=True) + # Only an httpx transport failure raised while calling MPS is allowed to + # fail open, and those failures are handled at the MPS call sites above. + # Database, configuration, response-validation, and programming errors + # all reach this handler and fail closed. + return QuotaCheckResult( + has_quota=False, + error_code="quota_check_failed", + error_message="Could not verify Dograh credits. Please try again.", + ) diff --git a/api/tests/test_quota_service.py b/api/tests/test_quota_service.py index 9a85d279..3ac89eb6 100644 --- a/api/tests/test_quota_service.py +++ b/api/tests/test_quota_service.py @@ -726,3 +726,336 @@ async def test_authorize_workflow_run_denies_run_bound_to_other_workflow(monkeyp assert result.has_quota is False assert result.error_code == "workflow_run_not_found" get_config.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_authorize_workflow_run_fails_closed_on_config_resolution_error( + monkeypatch, +): + """A config-resolution bug must deny the run, not fail open (issue #331).""" + monkeypatch.setattr(quota_service, "DEPLOYMENT_MODE", "saas") + _patch_workflow_context(monkeypatch) + monkeypatch.setattr( + quota_service, + "get_effective_ai_model_configuration_for_workflow", + AsyncMock(side_effect=RuntimeError("configuration resolution bug")), + ) + + result = await quota_service.authorize_workflow_run_start( + workflow_id=7, + organization_id=42, + ) + + assert result.has_quota is False + assert result.error_code == "quota_check_failed" + + +@pytest.mark.asyncio +async def test_authorize_workflow_run_fails_closed_on_user_lookup_error(monkeypatch): + """A DB read failure on the owner lookup is a 'cannot verify' → denied.""" + get_config = AsyncMock() + + monkeypatch.setattr(quota_service, "DEPLOYMENT_MODE", "saas") + _patch_workflow_context(monkeypatch) + monkeypatch.setattr( + quota_service.db_client, + "get_user_by_id", + AsyncMock(side_effect=RuntimeError("database unavailable")), + ) + monkeypatch.setattr( + quota_service, + "get_effective_ai_model_configuration_for_workflow", + get_config, + ) + + result = await quota_service.authorize_workflow_run_start( + workflow_id=7, + organization_id=42, + ) + + assert result.has_quota is False + assert result.error_code == "quota_check_failed" + get_config.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_authorize_workflow_run_fails_closed_on_run_lookup_error(monkeypatch): + """A DB read failure on the run lookup is a 'cannot verify' → denied.""" + get_config = AsyncMock() + + monkeypatch.setattr(quota_service, "DEPLOYMENT_MODE", "saas") + _patch_workflow_context(monkeypatch) + monkeypatch.setattr( + quota_service.db_client, + "get_workflow_run", + AsyncMock(side_effect=RuntimeError("database unavailable")), + ) + monkeypatch.setattr( + quota_service, + "get_effective_ai_model_configuration_for_workflow", + get_config, + ) + + result = await quota_service.authorize_workflow_run_start( + workflow_id=7, + organization_id=42, + workflow_run_id=88, + ) + + assert result.has_quota is False + assert result.error_code == "quota_check_failed" + get_config.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_authorize_workflow_run_opens_when_hosted_mps_is_unreachable( + monkeypatch, +): + request = httpx.Request( + "POST", + "https://services.dograh.com/api/v1/billing/accounts/42/run-authorization", + ) + + monkeypatch.setattr(quota_service, "DEPLOYMENT_MODE", "saas") + _patch_workflow_context(monkeypatch) + monkeypatch.setattr( + quota_service, + "get_effective_ai_model_configuration_for_workflow", + AsyncMock(return_value=_byok_config()), + ) + monkeypatch.setattr( + quota_service.mps_service_key_client, + "authorize_workflow_run_start", + AsyncMock( + side_effect=httpx.ConnectError("connection refused", request=request) + ), + ) + + result = await quota_service.authorize_workflow_run_start( + workflow_id=7, + organization_id=42, + ) + + assert result.has_quota is True + + +@pytest.mark.asyncio +async def test_authorize_workflow_run_fails_closed_on_hosted_mps_http_error( + monkeypatch, +): + request = httpx.Request( + "POST", + "https://services.dograh.com/api/v1/billing/accounts/42/run-authorization", + ) + response = httpx.Response(503, request=request) + + monkeypatch.setattr(quota_service, "DEPLOYMENT_MODE", "saas") + _patch_workflow_context(monkeypatch) + monkeypatch.setattr( + quota_service, + "get_effective_ai_model_configuration_for_workflow", + AsyncMock(return_value=_byok_config()), + ) + monkeypatch.setattr( + quota_service.mps_service_key_client, + "authorize_workflow_run_start", + AsyncMock( + side_effect=httpx.HTTPStatusError( + "MPS unavailable", + request=request, + response=response, + ) + ), + ) + + result = await quota_service.authorize_workflow_run_start( + workflow_id=7, + organization_id=42, + ) + + assert result.has_quota is False + assert result.error_code == "quota_check_failed" + + +@pytest.mark.asyncio +async def test_authorize_workflow_run_fails_closed_on_invalid_mps_url(monkeypatch): + request = httpx.Request("POST", "ftp://services.dograh.com/run-authorization") + + monkeypatch.setattr(quota_service, "DEPLOYMENT_MODE", "saas") + _patch_workflow_context(monkeypatch) + monkeypatch.setattr( + quota_service, + "get_effective_ai_model_configuration_for_workflow", + AsyncMock(return_value=_byok_config()), + ) + monkeypatch.setattr( + quota_service.mps_service_key_client, + "authorize_workflow_run_start", + AsyncMock( + side_effect=httpx.UnsupportedProtocol( + "Unsupported protocol ftp://", + request=request, + ) + ), + ) + + result = await quota_service.authorize_workflow_run_start( + workflow_id=7, + organization_id=42, + ) + + assert result.has_quota is False + assert result.error_code == "quota_check_failed" + + +@pytest.mark.asyncio +async def test_authorize_workflow_run_opens_when_oss_quota_mps_is_unreachable( + monkeypatch, +): + request = httpx.Request( + "GET", + "https://services.dograh.com/api/v1/service-keys/usage/self", + ) + + monkeypatch.setattr(quota_service, "DEPLOYMENT_MODE", "oss") + _patch_workflow_context(monkeypatch) + monkeypatch.setattr( + quota_service, + "get_effective_ai_model_configuration_for_workflow", + AsyncMock(return_value=_dograh_config()), + ) + monkeypatch.setattr( + quota_service.mps_service_key_client, + "check_service_key_usage", + AsyncMock(side_effect=httpx.ConnectTimeout("timed out", request=request)), + ) + + result = await quota_service.authorize_workflow_run_start( + workflow_id=7, + organization_id=42, + ) + + assert result.has_quota is True + + +@pytest.mark.asyncio +async def test_authorize_workflow_run_fails_closed_on_oss_quota_mps_http_error( + monkeypatch, +): + request = httpx.Request( + "GET", + "https://services.dograh.com/api/v1/service-keys/usage/self", + ) + response = httpx.Response(503, request=request) + + monkeypatch.setattr(quota_service, "DEPLOYMENT_MODE", "oss") + _patch_workflow_context(monkeypatch) + monkeypatch.setattr( + quota_service, + "get_effective_ai_model_configuration_for_workflow", + AsyncMock(return_value=_dograh_config()), + ) + monkeypatch.setattr( + quota_service.mps_service_key_client, + "check_service_key_usage", + AsyncMock( + side_effect=httpx.HTTPStatusError( + "MPS unavailable", + request=request, + response=response, + ) + ), + ) + + result = await quota_service.authorize_workflow_run_start( + workflow_id=7, + organization_id=42, + ) + + assert result.has_quota is False + assert result.error_code == "quota_check_failed" + + +@pytest.mark.asyncio +async def test_authorize_workflow_run_opens_when_oss_correlation_mps_is_unreachable( + monkeypatch, +): + request = httpx.Request( + "POST", + "https://services.dograh.com/api/v1/service-keys/correlation-id/self", + ) + + monkeypatch.setattr(quota_service, "DEPLOYMENT_MODE", "oss") + _patch_workflow_context(monkeypatch) + monkeypatch.setattr( + quota_service.db_client, + "get_workflow_run", + AsyncMock(return_value=_pinned_run()), + ) + monkeypatch.setattr( + quota_service, + "get_effective_ai_model_configuration_for_workflow", + AsyncMock(return_value=_dograh_config(managed_service_version=2)), + ) + monkeypatch.setattr( + quota_service.mps_service_key_client, + "check_service_key_usage", + AsyncMock(return_value={"remaining_credits": 25.0}), + ) + monkeypatch.setattr( + quota_service.mps_service_key_client, + "create_correlation_id", + AsyncMock( + side_effect=httpx.ConnectError("connection refused", request=request) + ), + ) + + result = await quota_service.authorize_workflow_run_start( + workflow_id=7, + organization_id=42, + workflow_run_id=88, + ) + + assert result.has_quota is True + + +@pytest.mark.asyncio +async def test_authorize_workflow_run_fails_closed_when_storing_oss_correlation( + monkeypatch, +): + monkeypatch.setattr(quota_service, "DEPLOYMENT_MODE", "oss") + _patch_workflow_context(monkeypatch) + monkeypatch.setattr( + quota_service.db_client, + "get_workflow_run", + AsyncMock(return_value=_pinned_run()), + ) + monkeypatch.setattr( + quota_service.db_client, + "get_workflow_run_by_id", + AsyncMock(side_effect=RuntimeError("database unavailable")), + ) + monkeypatch.setattr( + quota_service, + "get_effective_ai_model_configuration_for_workflow", + AsyncMock(return_value=_dograh_config(managed_service_version=2)), + ) + monkeypatch.setattr( + quota_service.mps_service_key_client, + "check_service_key_usage", + AsyncMock(return_value={"remaining_credits": 25.0}), + ) + monkeypatch.setattr( + quota_service.mps_service_key_client, + "create_correlation_id", + AsyncMock(return_value={"correlation_id": "oss-corr-123"}), + ) + + result = await quota_service.authorize_workflow_run_start( + workflow_id=7, + organization_id=42, + workflow_run_id=88, + ) + + assert result.has_quota is False + assert result.error_code == "quota_check_failed"