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

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