From 076edd1bd0aea3ea0efba9e1a639a3f17757fdae Mon Sep 17 00:00:00 2001 From: Amaan Javed <68250881+amaanJvd@users.noreply.github.com> Date: Wed, 15 Jul 2026 00:33:42 -0700 Subject: [PATCH] fix(quota): fail closed when quota verification errors (#331) (#523) * fix(quota): fail closed when quota verification errors (#331) Quota enforcement fell open on unexpected errors: the outer `except` in `authorize_workflow_run_start` returned `has_quota=True`, so a degraded database or a config-resolution bug let a billable run start unverified. Billing and abuse protection are control-plane functions, so this is the wrong default under exactly the degraded conditions that matter. - Fail closed by default: the outer handler now returns `has_quota=False` / `quota_check_failed`, reusing the existing message. - Add `QUOTA_FAIL_MODE=closed|open` (default `closed`) so OSS self-hosters can explicitly opt back into availability; the open path logs loudly. - Narrow the try-scope so `get_user_by_id` / `get_workflow_run` DB read failures surface as their specific `user_not_found` / `workflow_run_not_found` codes instead of the generic handler. - Tests cover the config-resolution and DB-read failure paths (denied, not `has_quota=True`) and the `QUOTA_FAIL_MODE=open` escape hatch. Fixes #331 Co-Authored-By: Claude Opus 4.8 * fix(quota): route DB read failures through the fail-mode policy gate Review (greptile) flagged that the narrowed get_user_by_id / get_workflow_run catches returned user_not_found / workflow_run_not_found before the outer QUOTA_FAIL_MODE handler ran, so QUOTA_FAIL_MODE=open never applied to a DB failure -- the exact "degraded database" case the escape hatch documents. Revert the two narrowed catches so DB read exceptions fall through to the single outer policy gate: closed -> quota_check_failed, open -> allow. The None checks still return the specific not_found codes for genuinely missing rows; an exception is a "cannot verify" condition, not a definitive absence. Add a regression test asserting QUOTA_FAIL_MODE=open allows a run when a DB read throws, and update the two DB-error tests to expect quota_check_failed. Co-Authored-By: Claude Opus 4.8 * docs(quota): scope the fail-mode comment to credit-verification failures (#331) Review (cubic) flagged the outer-handler comment as overclaiming: it said the handler is the single gate for "all cannot-verify errors", but the earlier workflow-load and org-membership catches always deny with workflow_not_found regardless of QUOTA_FAIL_MODE. That distinction is intentional (those are authorization/existence gates, not credit verification), so scope the comment accordingly. Comment-only, no behavior change. Co-Authored-By: Claude Opus 4.8 * fix(quota): fail open only when MPS is unreachable --------- Co-authored-by: Claude Opus 4.8 Co-authored-by: Abhishek Kumar --- api/services/quota_service.py | 45 ++++- api/tests/test_quota_service.py | 333 ++++++++++++++++++++++++++++++++ 2 files changed, 376 insertions(+), 2 deletions(-) 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"