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 <noreply@anthropic.com>

* 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 <noreply@anthropic.com>

* 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 <noreply@anthropic.com>

* fix(quota): fail open only when MPS is unreachable

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: Abhishek Kumar <abhishek@a6k.me>
This commit is contained in:
Amaan Javed 2026-07-15 00:33:42 -07:00 committed by GitHub
parent e4d2bc8e69
commit 076edd1bd0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 376 additions and 2 deletions

View file

@ -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.",
)

View file

@ -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"