apunkt/dograh
1
0
Fork 0
mirror of https://github.com/dograh-hq/dograh.git synced 2026-07-10 11:12:13 +02:00
Code Issues Projects Releases Packages Wiki Activity

chore: cleaup mps v1 billing (#507)

Browse source 
* chore: cleaup mps v1 billing

* chore: remove legacy file upload path

* chore: implement review comments
This commit is contained in:
Abhishek 2026-07-07 18:38:29 +05:30 • committed by GitHub
parent ac01f7775e
commit fdb7f92fcc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
36 changed files with 268 additions and 1319 deletions
Download patch file Download diff file Expand all files Collapse all files

4
api/routes/knowledge_base.py
View file

@ -402,7 +402,7 @@ async def search_chunks(
)
# Manual search runs outside any workflow run, so resolve the MPS
# correlation id here (mint only for orgs already on v2; never create one).
# correlation id here.
embedding_service = await build_embedding_service(
db_client=db_client,
provider=embeddings_provider,
@ -411,8 +411,6 @@ async def search_chunks(
base_url=embeddings_base_url,
endpoint=embeddings_endpoint,
api_version=embeddings_api_version,
organization_id=user.selected_organization_id,
created_by=str(user.provider_id),
resolve_correlation=True,
)

16
api/routes/organization.py
View file

@ -42,7 +42,6 @@ from api.schemas.telephony_phone_number import (
ProviderSyncStatus,
)
from api.services.auth.depends import (
_sync_posthog_organization_mps_billing_v2_status,
get_user,
get_user_with_selected_organization,
)
@ -391,22 +390,21 @@ async def migrate_model_configuration_v2(
except ValueError as exc:
raise HTTPException(status_code=422, detail=exc.args[0])
billing_account_status = None
if DEPLOYMENT_MODE != "oss":
try:
billing_account_status = await ensure_hosted_mps_billing_account_v2(
await ensure_hosted_mps_billing_account_v2(
organization_id,
created_by=str(user.provider_id),
)
except Exception as exc:
logger.error(
"Failed to initialize MPS billing v2 account for organization {}: {}",
"Failed to initialize MPS billing account for organization {}: {}",
organization_id,
exc,
)
raise HTTPException(
status_code=502,
detail="Failed to initialize MPS billing v2 account",
detail="Failed to initialize MPS billing account",
)
await upsert_organization_ai_model_configuration_v2(
@ -417,14 +415,6 @@ async def migrate_model_configuration_v2(
organization_id=organization_id,
fallback_user_config=legacy,
)
if DEPLOYMENT_MODE != "oss":
_sync_posthog_organization_mps_billing_v2_status(
organization_id,
uses_mps_billing_v2=bool(
billing_account_status
and billing_account_status.get("billing_mode") == "v2"
),
)
return await _model_configuration_v2_response(
user=user,
configuration=configuration,

95
api/routes/organization_usage.py
View file

@ -1,6 +1,6 @@
import json
from datetime import datetime, timedelta
from typing import Any, Dict, List, Literal, Optional
from typing import Any, Dict, List, Optional
from fastapi import APIRouter, Depends, HTTPException, Query
from fastapi.responses import StreamingResponse
@ -29,12 +29,6 @@ class CurrentUsageResponse(BaseModel):
price_per_second_usd: Optional[float] = None
class MPSCreditsResponse(BaseModel):
total_credits_used: float
remaining_credits: float
total_quota: float
class MPSCreditPurchaseUrlResponse(BaseModel):
checkout_url: str
@ -69,7 +63,6 @@ class MPSCreditLedgerEntryResponse(BaseModel):
class MPSBillingCreditsResponse(BaseModel):
billing_version: Literal["legacy", "v2"]
total_credits_used: float = 0.0
remaining_credits: float = 0.0
total_quota: float = 0.0
@ -162,69 +155,13 @@ async def get_current_period_usage(user: UserModel = Depends(get_user)):
raise HTTPException(status_code=500, detail=str(e))
@router.get("/usage/mps-credits", response_model=MPSCreditsResponse)
async def get_mps_credits(user: UserModel = Depends(get_user)):
"""Get aggregated usage and quota from MPS.
OSS users: queries by provider_id (created_by).
Hosted users: queries by organization_id.
"""
try:
if DEPLOYMENT_MODE == "oss":
usage = await mps_service_key_client.get_usage_by_created_by(
str(user.provider_id)
)
else:
if not user.selected_organization_id:
raise HTTPException(status_code=400, detail="No organization selected")
usage = await mps_service_key_client.get_usage_by_organization(
user.selected_organization_id
)
total_used = usage.get("total_credits_used", 0.0)
total_remaining = usage.get("remaining_credits", 0.0)
return MPSCreditsResponse(
total_credits_used=total_used,
remaining_credits=total_remaining,
total_quota=total_used + total_remaining,
)
except HTTPException:
raise
except Exception as e:
logger.error(f"Failed to fetch MPS credits: {e}")
raise HTTPException(status_code=500, detail=str(e))
async def _get_mps_billing_account_status(
user: UserModel, organization_id: int
) -> Optional[dict]:
return await mps_service_key_client.get_billing_account_status(
organization_id=organization_id,
created_by=str(user.provider_id),
)
def _is_mps_billing_v2(account: Optional[dict]) -> bool:
return bool(account and account.get("billing_mode") == "v2")
async def _legacy_mps_credits_response(user: UserModel) -> MPSBillingCreditsResponse:
if DEPLOYMENT_MODE == "oss":
usage = await mps_service_key_client.get_usage_by_created_by(
str(user.provider_id)
)
else:
if not user.selected_organization_id:
raise HTTPException(status_code=400, detail="No organization selected")
usage = await mps_service_key_client.get_usage_by_organization(
user.selected_organization_id
)
async def _oss_mps_credits_response(user: UserModel) -> MPSBillingCreditsResponse:
"""Aggregate per-key MPS credits for OSS deployments (no billing account)."""
usage = await mps_service_key_client.get_usage_by_created_by(str(user.provider_id))
total_used = float(usage.get("total_credits_used", 0.0))
total_remaining = float(usage.get("remaining_credits", 0.0))
return MPSBillingCreditsResponse(
billing_version="legacy",
total_credits_used=total_used,
remaining_credits=total_remaining,
total_quota=total_used + total_remaining,
@ -237,16 +174,15 @@ async def get_billing_credits(
limit: int = Query(50, ge=1, le=100),
user: UserModel = Depends(get_user),
):
"""Return legacy MPS credits or paginated v2 billing ledger details for the org."""
"""Return per-key MPS credits (OSS) or the org's paginated billing ledger."""
try:
if DEPLOYMENT_MODE == "oss" or not user.selected_organization_id:
return await _legacy_mps_credits_response(user)
if DEPLOYMENT_MODE == "oss":
return await _oss_mps_credits_response(user)
if not user.selected_organization_id:
raise HTTPException(status_code=400, detail="No organization selected")
organization_id = user.selected_organization_id
account_status = await _get_mps_billing_account_status(user, organization_id)
if not _is_mps_billing_v2(account_status):
return await _legacy_mps_credits_response(user)
ledger = await mps_service_key_client.get_credit_ledger(
organization_id=organization_id,
page=page,
@ -287,7 +223,6 @@ async def get_billing_credits(
total_debits = float(ledger["total_debits_credits"])
return MPSBillingCreditsResponse(
billing_version="v2",
total_credits_used=total_debits,
remaining_credits=balance,
total_quota=balance + total_debits,
@ -350,7 +285,7 @@ async def get_billing_credits(
async def create_mps_credit_purchase_url(
user: UserModel = Depends(get_user_with_selected_organization),
):
"""Create a checkout URL for organizations using Dograh-managed MPS v2."""
"""Create a checkout URL for purchasing organization credits."""
if DEPLOYMENT_MODE == "oss":
raise HTTPException(
status_code=404,
@ -359,14 +294,6 @@ async def create_mps_credit_purchase_url(
organization_id = user.selected_organization_id
assert organization_id is not None
account_status = await _get_mps_billing_account_status(user, organization_id)
if not _is_mps_billing_v2(account_status):
raise HTTPException(
status_code=403,
detail=(
"Credit purchases are available only for organizations using billing v2"
),
)
try:
session = await mps_service_key_client.create_credit_purchase_url(

50
api/services/auth/depends.py
View file

@ -34,7 +34,6 @@ async def require_local_auth() -> None:
POSTHOG_ORGANIZATION_GROUP_TYPE = "organization"
POSTHOG_ORGANIZATION_USES_MPS_BILLING_V2_PROPERTY = "uses_mps_billing_v2"
async def get_user(
@ -193,7 +192,6 @@ def _sync_created_organization_to_posthog(
organization,
stack_user: dict | None = None,
created_by_provider_id: str | None = None,
uses_mps_billing_v2: bool | None = None,
) -> None:
"""Create/update the PostHog organization group for a newly-created org."""
try:
@ -209,10 +207,6 @@ def _sync_created_organization_to_posthog(
}
if created_by:
properties["created_by_provider_id"] = created_by
if uses_mps_billing_v2 is not None:
properties[POSTHOG_ORGANIZATION_USES_MPS_BILLING_V2_PROPERTY] = (
uses_mps_billing_v2
)
group_identify(
POSTHOG_ORGANIZATION_GROUP_TYPE,
@ -231,50 +225,6 @@ def _sync_created_organization_to_posthog(
logger.exception("Failed to sync created organization to PostHog")
def _sync_posthog_organization_group_properties(
*,
organization,
uses_mps_billing_v2: bool | None = None,
) -> None:
"""Update PostHog organization group properties without creating a person."""
try:
organization_id = int(organization.id)
properties = {
"organization_id": organization_id,
"organization_provider_id": getattr(organization, "provider_id", None),
"auth_provider": "stack",
}
if uses_mps_billing_v2 is not None:
properties[POSTHOG_ORGANIZATION_USES_MPS_BILLING_V2_PROPERTY] = (
uses_mps_billing_v2
)
group_identify(
POSTHOG_ORGANIZATION_GROUP_TYPE,
str(organization_id),
properties,
)
except Exception:
logger.exception("Failed to sync organization group properties to PostHog")
def _sync_posthog_organization_mps_billing_v2_status(
organization_id: int,
*,
uses_mps_billing_v2: bool,
) -> None:
"""Update the PostHog organization group with current MPS billing status."""
try:
organization_id = int(organization_id)
group_identify(
POSTHOG_ORGANIZATION_GROUP_TYPE,
str(organization_id),
{POSTHOG_ORGANIZATION_USES_MPS_BILLING_V2_PROPERTY: uses_mps_billing_v2},
)
except Exception:
logger.exception("Failed to sync organization billing status to PostHog")
def _associate_user_with_posthog_organization(
*,
user: UserModel,

48
api/services/gen_ai/embedding/factory.py
View file

@ -2,9 +2,8 @@
Centralizes the provider branching (Azure BYOK / Dograh-managed / OpenAI-compatible
BYOK) that was previously duplicated across document ingestion, the search route,
and the RAG tool, and resolves the MPS billing v2 protocol the same way the voice
path does: attach it only for orgs already on v2, and never create a billing
account to do so.
and the RAG tool, and resolves the MPS correlation id the same way the voice
path does.
"""
from typing import Optional
@ -24,46 +23,23 @@ DEFAULT_AZURE_API_VERSION = "2024-02-15-preview"
async def resolve_embedding_correlation_id(
*,
organization_id: Optional[int],
service_key: Optional[str],
created_by: Optional[str] = None,
) -> Optional[str]:
"""Resolve an MPS correlation id for a managed embedding call made outside a run.
"""Mint an MPS correlation id for a managed embedding call made outside a run.
Mirrors the voice path's gating:
- OSS deployments use a pasted hosted v2 key (v2 by definition), so mint
directly via the bearer endpoint — matching ``_authorize_oss_managed_v2_correlation``.
- Hosted/SaaS: read the org's billing mode (no side effects) and mint only when
it is already v2. Minting for an already-v2 org is a no-op on the account.
Returns ``None`` when the call should be sent without the protocol; MPS accepts
un-gated embedding calls from v1 orgs. Never creates a v2 billing account.
Matches the voice path's ``_authorize_oss_managed_v2_correlation``: the
correlation is minted via the bearer service-key endpoint, so it works for
hosted orgs and OSS keys alike. Returns ``None`` when minting fails; MPS
accepts un-correlated embedding calls.
"""
if not service_key:
return None
# Imported lazily to avoid import-time cycles between the gen_ai and service
# layers (matches the inline-import convention used elsewhere in the app).
from api.constants import DEPLOYMENT_MODE
from api.services.mps_service_key_client import mps_service_key_client
try:
if DEPLOYMENT_MODE == "oss":
minted = await mps_service_key_client.create_correlation_id(
service_key=service_key
)
return minted.get("correlation_id")
if organization_id is None:
return None
status = await mps_service_key_client.get_billing_account_status(
organization_id, created_by=created_by
)
if not status or status.get("billing_mode") != "v2":
return None
minted = await mps_service_key_client.create_correlation_id(
service_key=service_key
)
@ -71,7 +47,7 @@ async def resolve_embedding_correlation_id(
except Exception as e:
logger.warning(
"Could not resolve MPS correlation id for managed embeddings; "
"sending without v2 protocol: {}",
"sending without it: {}",
e,
)
return None
@ -87,8 +63,6 @@ async def build_embedding_service(
endpoint: Optional[str] = None,
api_version: Optional[str] = None,
correlation_id: Optional[str] = None,
organization_id: Optional[int] = None,
created_by: Optional[str] = None,
resolve_correlation: bool = False,
) -> BaseEmbeddingService:
"""Construct the right embedding service for a provider/config.
@ -116,11 +90,7 @@ async def build_embedding_service(
if provider == ServiceProviders.DOGRAH.value:
cid = correlation_id
if cid is None and resolve_correlation:
cid = await resolve_embedding_correlation_id(
organization_id=organization_id,
service_key=api_key,
created_by=created_by,
)
cid = await resolve_embedding_correlation_id(service_key=api_key)
return DograhEmbeddingService(
db_client=db_client,
api_key=api_key,

70
api/services/mps_service_key_client.py
View file

@ -241,19 +241,12 @@ class MPSServiceKeyClient:
)
return False
async def check_service_key_usage(
self,
service_key: str,
organization_id: Optional[int] = None,
created_by: Optional[str] = None,
) -> dict:
async def check_service_key_usage(self, service_key: str) -> dict:
"""
Check the usage and quota of a service key.
Args:
service_key: The service key to check usage for
organization_id: Organization ID (for authenticated mode)
created_by: User provider ID (for OSS mode)
Returns:
Dictionary containing:
@ -321,39 +314,6 @@ class MPSServiceKeyClient:
response=response,
)
async def get_usage_by_organization(self, organization_id: int) -> dict:
"""
Get aggregated usage for all service keys belonging to an organization (hosted mode).
Args:
organization_id: The organization's ID
Returns:
Dictionary containing total_credits_used and remaining_credits
"""
async with httpx.AsyncClient(timeout=self.timeout) as client:
response = await client.post(
f"{self.base_url}/api/v1/service-keys/usage/organization",
json={"organization_id": organization_id},
headers=self._get_headers(organization_id=organization_id),
)
if response.status_code == 200:
data = response.json()
return {
"total_credits_used": data.get("total_credits_used", 0.0),
"remaining_credits": data.get("remaining_credits", 0.0),
}
else:
logger.error(
f"Failed to get usage by organization: {response.status_code} - {response.text}"
)
raise httpx.HTTPStatusError(
f"Failed to get usage by organization: {response.text}",
request=response.request,
response=response,
)
async def create_credit_purchase_url(
self,
organization_id: int,
@ -422,34 +382,6 @@ class MPSServiceKeyClient:
response=response,
)
async def get_billing_account_status(
self,
organization_id: int,
created_by: Optional[str] = None,
) -> Optional[dict]:
"""Get an existing MPS v2 billing account without creating one."""
async with httpx.AsyncClient(timeout=self.timeout) as client:
response = await client.get(
f"{self.base_url}/api/v1/billing/accounts/{organization_id}/status",
headers=self._get_headers(
organization_id=organization_id,
created_by=created_by,
),
)
if response.status_code == 200:
return response.json()
logger.error(
"Failed to get MPS billing account status: "
f"{response.status_code} - {response.text}"
)
raise httpx.HTTPStatusError(
f"Failed to get MPS billing account status: {response.text}",
request=response.request,
response=response,
)
async def ensure_billing_account_v2(
self,
organization_id: int,

136
api/services/quota_service.py
View file

@ -25,13 +25,13 @@ from api.services.mps_service_key_client import mps_service_key_client
MINIMUM_DOGRAH_CREDITS_FOR_CALL = 0.10
LEGACY_QUOTA_EXCEEDED_MESSAGE = (
OSS_QUOTA_EXCEEDED_MESSAGE = (
"You have exhausted your trial credits. "
"Please email founders@dograh.com for additional Dograh credits "
"or change providers in Models configurations."
"Please sign up on app.dograh.com to create a "
"new service key and set up in your model configurations."
)
BILLING_V2_QUOTA_EXCEEDED_MESSAGE = (
HOSTED_QUOTA_EXCEEDED_MESSAGE = (
"You have exhausted your Dograh credits. "
"Please purchase more credits from /billing "
"or change providers in Models configurations."
@ -60,19 +60,19 @@ def _safe_float(value: Any, default: float = 0.0) -> float:
return default
def _insufficient_billing_v2_quota_result() -> QuotaCheckResult:
def _insufficient_hosted_quota_result() -> QuotaCheckResult:
return QuotaCheckResult(
has_quota=False,
error_code="insufficient_credits",
error_message=BILLING_V2_QUOTA_EXCEEDED_MESSAGE,
error_message=HOSTED_QUOTA_EXCEEDED_MESSAGE,
)
def _insufficient_legacy_quota_result() -> QuotaCheckResult:
def _insufficient_oss_quota_result() -> QuotaCheckResult:
return QuotaCheckResult(
has_quota=False,
error_code="quota_exceeded",
error_message=LEGACY_QUOTA_EXCEEDED_MESSAGE,
error_message=OSS_QUOTA_EXCEEDED_MESSAGE,
)
@ -157,10 +157,10 @@ async def _authorize_hosted_workflow_run_start(
workflow_id: int | None,
workflow_run_id: int | None,
user_config: Any,
) -> tuple[QuotaCheckResult, bool]:
"""Authorize hosted v2 billing and return whether MPS handled enforcement."""
if DEPLOYMENT_MODE == "oss" or organization_id is None:
return QuotaCheckResult(has_quota=True), False
) -> QuotaCheckResult:
"""Authorize a hosted workflow run against the org's MPS billing account."""
if organization_id is None:
return QuotaCheckResult(has_quota=True)
requires_correlation = bool(
workflow_run_id and uses_managed_model_services_v2(user_config)
@ -169,16 +169,13 @@ async def _authorize_hosted_workflow_run_start(
get_dograh_service_api_key(user_config) if requires_correlation else None
)
if requires_correlation and not service_key:
return (
QuotaCheckResult(
has_quota=False,
error_code="invalid_service_key",
error_message=(
"You have invalid keys in your model configuration. "
"Please validate the service keys."
),
return QuotaCheckResult(
has_quota=False,
error_code="invalid_service_key",
error_message=(
"You have invalid keys in your model configuration. "
"Please validate the service keys."
),
True,
)
try:
@ -205,38 +202,28 @@ async def _authorize_hosted_workflow_run_start(
e,
)
if _is_service_key_org_mismatch_error(e):
return (
QuotaCheckResult(
has_quota=False,
error_code="service_key_org_mismatch",
error_message=SERVICE_TOKEN_ORG_MISMATCH_MESSAGE,
),
True,
)
return (
QuotaCheckResult(
return QuotaCheckResult(
has_quota=False,
error_code="quota_check_failed",
error_message="Could not verify Dograh credits. Please try again.",
),
True,
error_code="service_key_org_mismatch",
error_message=SERVICE_TOKEN_ORG_MISMATCH_MESSAGE,
)
return QuotaCheckResult(
has_quota=False,
error_code="quota_check_failed",
error_message="Could not verify Dograh credits. Please try again.",
)
billing_mode = authorization.get("billing_mode")
if billing_mode != "v2":
return QuotaCheckResult(has_quota=True), False
remaining = _safe_float(authorization.get("remaining_credits"))
if (
not authorization.get("allowed", False)
or remaining < MINIMUM_DOGRAH_CREDITS_FOR_CALL
):
logger.warning(
"Insufficient Dograh billing v2 credits for org {}: {:.2f} credits remaining",
"Insufficient Dograh credits for org {}: {:.2f} credits remaining",
organization_id,
remaining,
)
return _insufficient_billing_v2_quota_result(), True
return _insufficient_hosted_quota_result()
try:
await _store_run_correlation_id(
@ -249,35 +236,27 @@ async def _authorize_hosted_workflow_run_start(
workflow_run_id,
e,
)
return (
QuotaCheckResult(
has_quota=False,
error_code="quota_check_failed",
error_message="Could not verify Dograh credits. Please try again.",
),
True,
return QuotaCheckResult(
has_quota=False,
error_code="quota_check_failed",
error_message="Could not verify Dograh credits. Please try again.",
)
logger.info(
"Dograh billing v2 run authorization passed for org {}: {:.2f} credits remaining",
"Dograh run authorization passed for org {}: {:.2f} credits remaining",
organization_id,
remaining,
)
return QuotaCheckResult(has_quota=True), True
return QuotaCheckResult(has_quota=True)
async def _authorize_legacy_dograh_keys(
async def _authorize_oss_dograh_keys(
*,
dograh_api_keys: set[str],
organization_id: int | None,
workflow_owner: UserModel,
) -> QuotaCheckResult:
"""Check per-key MPS credits for OSS deployments before a run starts."""
for api_key in dograh_api_keys:
try:
usage = await mps_service_key_client.check_service_key_usage(
api_key,
organization_id=organization_id,
created_by=workflow_owner.provider_id,
)
usage = await mps_service_key_client.check_service_key_usage(api_key)
remaining = usage.get("remaining_credits", 0.0)
# Require at least $0.10 for a short call
@ -286,7 +265,7 @@ async def _authorize_legacy_dograh_keys(
f"Insufficient Dograh credits for key ...{api_key[-8:]}: "
f"${remaining:.2f} remaining"
)
return _insufficient_legacy_quota_result()
return _insufficient_oss_quota_result()
logger.info(
f"Dograh quota check passed for key ...{api_key[-8:]}: "
@ -363,9 +342,9 @@ async def authorize_workflow_run_start(
) -> QuotaCheckResult:
"""Authorize a workflow run before any billable call/text runtime starts.
The workflow organization is the billing subject for hosted v2. The workflow
owner is used only to resolve the effective model configuration and legacy
service-key metadata.
The workflow organization is the billing subject for hosted deployments.
OSS deployments are billed per service key instead. The workflow owner is
used only to resolve the effective model configuration.
"""
try:
workflow = await db_client.get_workflow_by_id(workflow_id)
@ -405,38 +384,27 @@ async def authorize_workflow_run_start(
)
if DEPLOYMENT_MODE != "oss":
hosted_result, hosted_enforced = await _authorize_hosted_workflow_run_start(
return await _authorize_hosted_workflow_run_start(
workflow_owner=workflow_owner,
organization_id=workflow.organization_id,
workflow_id=workflow.id,
workflow_run_id=workflow_run_id,
user_config=user_config,
)
if hosted_enforced or not hosted_result.has_quota:
return hosted_result
dograh_api_keys = _dograh_api_keys(user_config)
if not dograh_api_keys:
return QuotaCheckResult(has_quota=True)
legacy_result = await _authorize_legacy_dograh_keys(
dograh_api_keys=dograh_api_keys,
organization_id=(
None if DEPLOYMENT_MODE == "oss" else workflow.organization_id
),
workflow_owner=workflow_owner,
)
if not legacy_result.has_quota:
return legacy_result
if DEPLOYMENT_MODE == "oss":
return await _authorize_oss_managed_v2_correlation(
workflow_id=workflow.id,
workflow_run_id=workflow_run_id,
user_config=user_config,
if dograh_api_keys:
oss_result = await _authorize_oss_dograh_keys(
dograh_api_keys=dograh_api_keys,
)
if not oss_result.has_quota:
return oss_result
return QuotaCheckResult(has_quota=True)
return await _authorize_oss_managed_v2_correlation(
workflow_id=workflow.id,
workflow_run_id=workflow_run_id,
user_config=user_config,
)
except Exception as e:
logger.error(f"Error during quota check: {str(e)}")

6
api/services/workflow/audit.py
View file

@ -1,10 +1,8 @@
"""Rule-based audit of a workflow definition's nodes + edges.
Pure, dependency-free helpers derived from `NodeSpec.graph_constraints`.
Lives in tracked code so the regression tests in
`test_workflow_graph_constraints.py` can pin it; the admin cleanup
script in `api/services/admin_utils/local_exec.py` is the production
consumer.
Lives in tracked code so `test_workflow_graph_constraints.py` can pin the
verdicts that one-off cleanup tooling needs to share with runtime validation.
"""
from collections import Counter

3
api/services/workflow/tools/knowledge_base.py
View file

@ -266,8 +266,7 @@ async def _perform_retrieval(
)
# Search runs inside a workflow run: reuse the run's MPS correlation
# id (present only for v2 orgs; None otherwise → sent without the
# protocol). The Dograh-managed path forwards it via request metadata.
# id. The Dograh-managed path forwards it via request metadata.
embedding_service = await build_embedding_service(
db_client=db_client,
provider=embeddings_provider,

13
api/services/workflow_run_billing.py
View file

@ -32,13 +32,6 @@ def _duration_seconds_from_usage_info(workflow_run) -> float | None:
return duration_seconds if duration_seconds > 0 else None
async def _organization_uses_mps_billing_v2(organization_id: int) -> bool:
account = await mps_service_key_client.get_billing_account_status(
organization_id=organization_id
)
return bool(account and account.get("billing_mode") == "v2")
def _is_usage_not_ready_error(exc: Exception) -> bool:
response = getattr(exc, "response", None)
if getattr(response, "status_code", None) != 409:
@ -79,12 +72,6 @@ async def report_workflow_run_platform_usage(workflow_run) -> None:
return
try:
if not await _organization_uses_mps_billing_v2(organization_id):
logger.debug(
"Not reporting platform usage since org not using mps billing v2"
)
return
result = await mps_service_key_client.report_platform_usage(
organization_id=organization_id,
correlation_id=correlation_id,

4
api/tasks/knowledge_base_processing.py
View file

@ -215,7 +215,7 @@ async def process_knowledge_base_document(
return
# Ingestion runs outside any workflow run, so resolve the MPS correlation
# id here (mint only for orgs already on v2; never create an account).
# id here.
embedding_service = await build_embedding_service(
db_client=db_client,
provider=embeddings_provider,
@ -224,8 +224,6 @@ async def process_knowledge_base_document(
base_url=embeddings_base_url,
endpoint=embeddings_endpoint,
api_version=embeddings_api_version,
organization_id=organization_id,
created_by=created_by_provider_id,
resolve_correlation=True,
)

82
api/tasks/workflow_completion.py
View file

@ -1,111 +1,31 @@
import asyncio
import os
from loguru import logger
from pipecat.utils.run_context import set_current_run_id
from api.services.workflow_run_artifacts import upload_workflow_run_artifacts
from api.services.workflow_run_billing import (
report_completed_workflow_run_platform_usage,
)
from api.tasks.run_integrations import run_integrations_post_workflow_run
def _read_and_remove_temp_file(temp_file_path: str | None, label: str) -> bytes | None:
if not temp_file_path:
return None
try:
if not os.path.exists(temp_file_path):
logger.warning(f"{label} temp file not found: {temp_file_path}")
return None
with open(temp_file_path, "rb") as f:
data = f.read()
os.remove(temp_file_path)
return data
except Exception as e:
logger.error(f"Error reading legacy {label} temp file {temp_file_path}: {e}")
return None
async def _upload_legacy_temp_artifacts(
workflow_run_id: int,
audio_temp_path: str | None,
transcript_temp_path: str | None,
user_audio_temp_path: str | None,
bot_audio_temp_path: str | None,
) -> None:
"""Handle jobs enqueued before uploads moved into the pipeline process.
Pre-refactor web workers passed local temp-file paths; upload them if this
worker can still see the files (same host / shared volume).
Deprecated: remove once no pre-refactor jobs can remain in the queue.
"""
logger.info(
f"Processing legacy temp-file artifacts for workflow run {workflow_run_id}"
)
transcript_bytes = await asyncio.to_thread(
_read_and_remove_temp_file, transcript_temp_path, "transcript"
)
await upload_workflow_run_artifacts(
workflow_run_id,
mixed_audio_wav=await asyncio.to_thread(
_read_and_remove_temp_file, audio_temp_path, "mixed audio"
),
user_audio_wav=await asyncio.to_thread(
_read_and_remove_temp_file, user_audio_temp_path, "user audio"
),
bot_audio_wav=await asyncio.to_thread(
_read_and_remove_temp_file, bot_audio_temp_path, "bot audio"
),
transcript_text=(
transcript_bytes.decode("utf-8") if transcript_bytes else None
),
)
async def process_workflow_completion(
_ctx,
workflow_run_id: int,
audio_temp_path: str | None = None,
transcript_temp_path: str | None = None,
user_audio_temp_path: str | None = None,
bot_audio_temp_path: str | None = None,
):
"""Process workflow completion: run integrations and report billing.
Recording/transcript uploads happen in the pipeline process itself
(api/services/workflow_run_artifacts.py) before this job is enqueued,
so this task needs no shared filesystem with the web tier. The temp-path
arguments only exist for jobs enqueued by pre-refactor web workers.
so this task needs no shared filesystem with the web tier.
Args:
_ctx: ARQ context (unused)
workflow_run_id: The workflow run ID
audio_temp_path: Deprecated, pre-refactor jobs only
transcript_temp_path: Deprecated, pre-refactor jobs only
user_audio_temp_path: Deprecated, pre-refactor jobs only
bot_audio_temp_path: Deprecated, pre-refactor jobs only
"""
run_id = str(workflow_run_id)
set_current_run_id(run_id)
logger.info(f"Processing workflow completion for run {workflow_run_id}")
if (
audio_temp_path
or transcript_temp_path
or user_audio_temp_path
or bot_audio_temp_path
):
await _upload_legacy_temp_artifacts(
workflow_run_id,
audio_temp_path,
transcript_temp_path,
user_audio_temp_path,
bot_audio_temp_path,
)
# Run integrations including QA analysis (after uploads are complete)
try:
await run_integrations_post_workflow_run(_ctx, workflow_run_id)

9
api/tests/test_ai_model_configuration_v2.py
View file

@ -1,5 +1,5 @@
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
from unittest.mock import AsyncMock
import pytest
from pydantic import ValidationError
@ -442,7 +442,6 @@ async def test_migrate_model_configuration_v2_initializes_hosted_mps_billing(
ensure_billing = AsyncMock(return_value={"billing_mode": "v2"})
upsert = AsyncMock()
migrate_workflows = AsyncMock()
sync_posthog_billing = Mock()
monkeypatch.setattr(organization_routes, "DEPLOYMENT_MODE", "saas")
monkeypatch.setattr(
@ -480,11 +479,6 @@ async def test_migrate_model_configuration_v2_initializes_hosted_mps_billing(
"_model_configuration_v2_response",
AsyncMock(return_value=expected_response),
)
monkeypatch.setattr(
organization_routes,
"_sync_posthog_organization_mps_billing_v2_status",
sync_posthog_billing,
)
user = SimpleNamespace(
id=7,
@ -503,5 +497,4 @@ async def test_migrate_model_configuration_v2_initializes_hosted_mps_billing(
organization_id=42,
fallback_user_config=legacy,
)
sync_posthog_billing.assert_called_once_with(42, uses_mps_billing_v2=True)
assert response == expected_response

66
api/tests/test_auth_depends.py
View file

@ -209,7 +209,7 @@ def test_associate_user_with_posthog_org_supports_backfill_arguments(monkeypatch
assert "backfilled" not in capture_kwargs["properties"]
def test_sync_created_organization_to_posthog_supports_billing_flag(monkeypatch):
def test_sync_created_organization_to_posthog_with_provider_id(monkeypatch):
organization = SimpleNamespace(id=42, provider_id="team-1")
group_calls = []
capture_calls = []
@ -228,11 +228,9 @@ def test_sync_created_organization_to_posthog_supports_billing_flag(monkeypatch)
auth_depends._sync_created_organization_to_posthog(
organization=organization,
created_by_provider_id="stack-user-1",
uses_mps_billing_v2=True,
)
_, group_kwargs = group_calls[0]
group_args, _ = group_calls[0]
group_args, group_kwargs = group_calls[0]
assert group_args == (
"organization",
"42",
@ -241,69 +239,9 @@ def test_sync_created_organization_to_posthog_supports_billing_flag(monkeypatch)
"organization_provider_id": "team-1",
"auth_provider": "stack",
"created_by_provider_id": "stack-user-1",
"uses_mps_billing_v2": True,
},
)
assert group_kwargs == {"distinct_id": "stack-user-1"}
_, capture_kwargs = capture_calls[0]
assert capture_kwargs["distinct_id"] == "stack-user-1"
assert capture_kwargs["properties"]["uses_mps_billing_v2"] is True
def test_sync_posthog_organization_group_properties_has_no_distinct_id(monkeypatch):
organization = SimpleNamespace(id=42, provider_id="team-1")
group_calls = []
monkeypatch.setattr(
auth_depends,
"group_identify",
lambda *args, **kwargs: group_calls.append((args, kwargs)),
)
auth_depends._sync_posthog_organization_group_properties(
organization=organization,
uses_mps_billing_v2=True,
)
assert group_calls == [
(
(
"organization",
"42",
{
"organization_id": 42,
"organization_provider_id": "team-1",
"auth_provider": "stack",
"uses_mps_billing_v2": True,
},
),
{},
)
]
def test_sync_posthog_organization_mps_billing_v2_status(monkeypatch):
group_calls = []
monkeypatch.setattr(
auth_depends,
"group_identify",
lambda *args, **kwargs: group_calls.append((args, kwargs)),
)
auth_depends._sync_posthog_organization_mps_billing_v2_status(
42,
uses_mps_billing_v2=True,
)
assert group_calls == [
(
(
"organization",
"42",
{"uses_mps_billing_v2": True},
),
{},
)
]

80
api/tests/test_dograh_embedding_service.py
View file

@ -49,114 +49,52 @@ async def test_dograh_embedding_sends_plain_without_correlation():
await service.embed_texts(["hello"])
create.assert_awaited_once()
# No correlation id (e.g. a v1 org) → no MPS metadata; MPS accepts plain calls.
# No correlation id → no MPS metadata; MPS accepts plain calls.
assert "extra_body" not in create.await_args.kwargs
def _fake_mps_client(*, status_return=None, minted="minted"):
def _fake_mps_client(*, minted="minted"):
return SimpleNamespace(
get_billing_account_status=AsyncMock(return_value=status_return),
create_correlation_id=AsyncMock(return_value={"correlation_id": minted}),
)
@pytest.mark.asyncio
async def test_resolve_correlation_oss_mints_directly(monkeypatch):
async def test_resolve_correlation_mints_via_service_key(monkeypatch):
fake = _fake_mps_client()
monkeypatch.setattr(
"api.services.mps_service_key_client.mps_service_key_client", fake
)
monkeypatch.setattr("api.constants.DEPLOYMENT_MODE", "oss")
result = await resolve_embedding_correlation_id(
organization_id=None, service_key="sk-mps"
)
result = await resolve_embedding_correlation_id(service_key="sk-mps")
assert result == "minted"
fake.create_correlation_id.assert_awaited_once_with(service_key="sk-mps")
fake.get_billing_account_status.assert_not_awaited()
@pytest.mark.asyncio
async def test_resolve_correlation_hosted_v2_mints(monkeypatch):
fake = _fake_mps_client(status_return={"billing_mode": "v2"})
monkeypatch.setattr(
"api.services.mps_service_key_client.mps_service_key_client", fake
)
monkeypatch.setattr("api.constants.DEPLOYMENT_MODE", "hosted")
result = await resolve_embedding_correlation_id(
organization_id=42, service_key="sk-mps", created_by="user-1"
)
assert result == "minted"
fake.get_billing_account_status.assert_awaited_once_with(42, created_by="user-1")
fake.create_correlation_id.assert_awaited_once_with(service_key="sk-mps")
@pytest.mark.asyncio
async def test_resolve_correlation_hosted_v1_returns_none_without_minting(monkeypatch):
fake = _fake_mps_client(status_return={"billing_mode": "v1"})
monkeypatch.setattr(
"api.services.mps_service_key_client.mps_service_key_client", fake
)
monkeypatch.setattr("api.constants.DEPLOYMENT_MODE", "hosted")
result = await resolve_embedding_correlation_id(
organization_id=42, service_key="sk-mps"
)
assert result is None
fake.create_correlation_id.assert_not_awaited()
@pytest.mark.asyncio
async def test_resolve_correlation_hosted_no_account_returns_none(monkeypatch):
fake = _fake_mps_client(status_return=None)
monkeypatch.setattr(
"api.services.mps_service_key_client.mps_service_key_client", fake
)
monkeypatch.setattr("api.constants.DEPLOYMENT_MODE", "hosted")
result = await resolve_embedding_correlation_id(
organization_id=42, service_key="sk-mps"
)
assert result is None
fake.create_correlation_id.assert_not_awaited()
@pytest.mark.asyncio
async def test_resolve_correlation_no_service_key_returns_none(monkeypatch):
fake = _fake_mps_client(status_return={"billing_mode": "v2"})
fake = _fake_mps_client()
monkeypatch.setattr(
"api.services.mps_service_key_client.mps_service_key_client", fake
)
monkeypatch.setattr("api.constants.DEPLOYMENT_MODE", "hosted")
result = await resolve_embedding_correlation_id(
organization_id=42, service_key=None
)
result = await resolve_embedding_correlation_id(service_key=None)
assert result is None
fake.get_billing_account_status.assert_not_awaited()
fake.create_correlation_id.assert_not_awaited()
@pytest.mark.asyncio
async def test_resolve_correlation_swallows_errors(monkeypatch):
fake = SimpleNamespace(
get_billing_account_status=AsyncMock(side_effect=RuntimeError("mps down")),
create_correlation_id=AsyncMock(),
create_correlation_id=AsyncMock(side_effect=RuntimeError("mps down")),
)
monkeypatch.setattr(
"api.services.mps_service_key_client.mps_service_key_client", fake
)
monkeypatch.setattr("api.constants.DEPLOYMENT_MODE", "hosted")
# A transient MPS failure must not break embeddings — fall back to no protocol.
result = await resolve_embedding_correlation_id(
organization_id=42, service_key="sk-mps"
)
# A transient MPS failure must not break embeddings — send without it.
result = await resolve_embedding_correlation_id(service_key="sk-mps")
assert result is None

45
api/tests/test_mps_service_key_client.py
View file

@ -130,51 +130,6 @@ async def test_create_correlation_id_uses_bearer_auth(monkeypatch):
]
@pytest.mark.asyncio
async def test_get_billing_account_status_uses_hosted_org_auth(monkeypatch):
calls = []
class FakeAsyncClient:
def __init__(self, timeout):
self.timeout = timeout
async def __aenter__(self):
return self
async def __aexit__(self, exc_type, exc, tb):
return None
async def get(self, url, headers):
calls.append(("GET", url, headers))
return _Response(200, {"organization_id": 42, "billing_mode": "v2"})
monkeypatch.setattr(
"api.services.mps_service_key_client.httpx.AsyncClient", FakeAsyncClient
)
monkeypatch.setattr("api.services.mps_service_key_client.DEPLOYMENT_MODE", "saas")
monkeypatch.setattr(
"api.services.mps_service_key_client.DOGRAH_MPS_SECRET_KEY", "mps-secret"
)
client = MPSServiceKeyClient()
assert await client.get_billing_account_status(organization_id=42) == {
"organization_id": 42,
"billing_mode": "v2",
}
assert calls == [
(
"GET",
f"{client.base_url}/api/v1/billing/accounts/42/status",
{
"Content-Type": "application/json",
"X-Secret-Key": "mps-secret",
"X-Organization-Id": "42",
},
)
]
@pytest.mark.asyncio
async def test_authorize_workflow_run_start_uses_hosted_org_auth(monkeypatch):
calls = []

44
api/tests/test_organization_usage_billing.py
View file

@ -6,41 +6,36 @@ import pytest
from api.routes import organization_usage
def test_is_mps_billing_v2_depends_only_on_account_mode():
assert organization_usage._is_mps_billing_v2({"billing_mode": "v2"}) is True
assert organization_usage._is_mps_billing_v2({"billing_mode": "v1"}) is False
assert organization_usage._is_mps_billing_v2({"billing_mode": "shadow"}) is False
assert organization_usage._is_mps_billing_v2(None) is False
@pytest.mark.asyncio
async def test_get_mps_billing_account_status_uses_user_provider_id(monkeypatch):
get_status = AsyncMock(return_value={"billing_mode": "v2"})
async def test_get_billing_credits_oss_aggregates_by_created_by(monkeypatch):
monkeypatch.setattr(organization_usage, "DEPLOYMENT_MODE", "oss")
get_usage = AsyncMock(
return_value={"total_credits_used": 12.5, "remaining_credits": 487.5}
)
monkeypatch.setattr(
organization_usage.mps_service_key_client,
"get_billing_account_status",
get_status,
"get_usage_by_created_by",
get_usage,
)
user = SimpleNamespace(provider_id="provider-123")
user = SimpleNamespace(provider_id="provider-123", selected_organization_id=None)
assert await organization_usage._get_mps_billing_account_status(user, 42) == {
"billing_mode": "v2"
}
get_status.assert_awaited_once_with(
organization_id=42,
created_by="provider-123",
response = await organization_usage.get_billing_credits(
page=1,
limit=50,
user=user,
)
get_usage.assert_awaited_once_with("provider-123")
assert response.total_credits_used == 12.5
assert response.remaining_credits == 487.5
assert response.total_quota == 500.0
assert response.ledger_entries == []
@pytest.mark.asyncio
async def test_get_billing_credits_pages_v2_ledger(monkeypatch):
async def test_get_billing_credits_pages_hosted_ledger(monkeypatch):
monkeypatch.setattr(organization_usage, "DEPLOYMENT_MODE", "saas")
monkeypatch.setattr(
organization_usage,
"_get_mps_billing_account_status",
AsyncMock(return_value={"billing_mode": "v2"}),
)
get_ledger = AsyncMock(
return_value={
"account": {
@ -90,7 +85,6 @@ async def test_get_billing_credits_pages_v2_ledger(monkeypatch):
limit=25,
created_by="provider-123",
)
assert response.billing_version == "v2"
assert response.total_credits_used == 75
assert response.total_count == 101
assert response.page == 3

29
api/tests/test_quota_service.py
View file

@ -166,34 +166,22 @@ async def test_authorize_workflow_run_v2_insufficient_credits_prompts_billing(
@pytest.mark.asyncio
async def test_authorize_workflow_run_v1_uses_legacy_key_usage(
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))
authorize = AsyncMock(
return_value={
"allowed": True,
"billing_mode": "v1",
"remaining_credits": "0.0000",
}
)
check_usage = AsyncMock(
return_value={"total_credits_used": 500.0, "remaining_credits": 0.0}
)
monkeypatch.setattr(quota_service, "DEPLOYMENT_MODE", "saas")
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,
"authorize_workflow_run_start",
authorize,
)
monkeypatch.setattr(
quota_service.mps_service_key_client,
"check_service_key_usage",
@ -206,12 +194,7 @@ async def test_authorize_workflow_run_v1_uses_legacy_key_usage(
assert result.error_code == "quota_exceeded"
assert "founders@dograh.com" in result.error_message
assert "/billing" not in result.error_message
authorize.assert_awaited_once()
check_usage.assert_awaited_once_with(
api_key,
organization_id=42,
created_by="provider-123",
)
check_usage.assert_awaited_once_with(api_key)
@pytest.mark.asyncio
@ -404,11 +387,7 @@ async def test_authorize_workflow_run_oss_uses_key_paths_not_workflow_org(
assert result.has_quota is True
hosted_authorize.assert_not_awaited()
check_usage.assert_awaited_once_with(
api_key,
organization_id=None,
created_by="provider-123",
)
check_usage.assert_awaited_once_with(api_key)
create_correlation.assert_awaited_once_with(
service_key=api_key,
workflow_run_id=88,

4
api/tests/test_workflow_graph_constraints.py
View file

@ -9,9 +9,9 @@ category of violation we found in production. We pin two layers:
layer ever stops rejecting one of these fixtures, the production
write paths will quietly start accepting bad workflows again.
2. audit_definition (api.services.workflow.audit) — read-only sweep
over persisted rows used by the admin cleanup script to find
over persisted rows for one-off cleanup tooling that finds
legacy/imported breakage. Pinned so refactors of the rule set
don't silently change the verdicts the migration relies on.
don't silently change those cleanup verdicts.
DTO-level shape validation is covered by `test_dto.py` and isn't
re-pinned here.

49
api/tests/test_workflow_run_billing.py
View file

@ -40,15 +40,9 @@ async def test_report_workflow_run_platform_usage_reports_hosted_completion(
monkeypatch,
):
workflow_run = _make_workflow_run()
get_status = AsyncMock(return_value={"billing_mode": "v2"})
report_usage = AsyncMock(return_value={"metered": True})
monkeypatch.setattr(workflow_run_billing_mod, "DEPLOYMENT_MODE", "saas")
monkeypatch.setattr(
workflow_run_billing_mod.mps_service_key_client,
"get_billing_account_status",
get_status,
)
monkeypatch.setattr(
workflow_run_billing_mod.mps_service_key_client,
"report_platform_usage",
@ -76,15 +70,9 @@ async def test_report_workflow_run_platform_usage_reports_duration_without_corre
):
workflow_run = _make_workflow_run()
workflow_run.initial_context = {}
get_status = AsyncMock(return_value={"billing_mode": "v2"})
report_usage = AsyncMock(return_value={"metered": True})
monkeypatch.setattr(workflow_run_billing_mod, "DEPLOYMENT_MODE", "saas")
monkeypatch.setattr(
workflow_run_billing_mod.mps_service_key_client,
"get_billing_account_status",
get_status,
)
monkeypatch.setattr(
workflow_run_billing_mod.mps_service_key_client,
"report_platform_usage",
@ -106,30 +94,6 @@ async def test_report_workflow_run_platform_usage_reports_duration_without_corre
)
@pytest.mark.asyncio
async def test_report_workflow_run_platform_usage_skips_non_v2_account(monkeypatch):
workflow_run = _make_workflow_run()
get_status = AsyncMock(return_value={"billing_mode": "v1"})
report_usage = AsyncMock()
monkeypatch.setattr(workflow_run_billing_mod, "DEPLOYMENT_MODE", "saas")
monkeypatch.setattr(
workflow_run_billing_mod.mps_service_key_client,
"get_billing_account_status",
get_status,
)
monkeypatch.setattr(
workflow_run_billing_mod.mps_service_key_client,
"report_platform_usage",
report_usage,
)
await report_workflow_run_platform_usage(workflow_run)
get_status.assert_awaited_once_with(organization_id=42)
report_usage.assert_not_awaited()
@pytest.mark.asyncio
async def test_report_workflow_run_platform_usage_skips_missing_duration_without_correlation(
monkeypatch,
@ -137,15 +101,9 @@ async def test_report_workflow_run_platform_usage_skips_missing_duration_without
workflow_run = _make_workflow_run()
workflow_run.initial_context = {}
workflow_run.usage_info = {}
get_status = AsyncMock(return_value={"billing_mode": "v2"})
report_usage = AsyncMock()
monkeypatch.setattr(workflow_run_billing_mod, "DEPLOYMENT_MODE", "saas")
monkeypatch.setattr(
workflow_run_billing_mod.mps_service_key_client,
"get_billing_account_status",
get_status,
)
monkeypatch.setattr(
workflow_run_billing_mod.mps_service_key_client,
"report_platform_usage",
@ -154,7 +112,6 @@ async def test_report_workflow_run_platform_usage_skips_missing_duration_without
await report_workflow_run_platform_usage(workflow_run)
get_status.assert_not_awaited()
report_usage.assert_not_awaited()
@ -197,7 +154,6 @@ async def test_report_workflow_run_platform_usage_skips_incomplete(monkeypatch):
async def test_report_completed_workflow_run_platform_usage_loads_run(monkeypatch):
workflow_run = _make_workflow_run()
get_run = AsyncMock(return_value=workflow_run)
get_status = AsyncMock(return_value={"billing_mode": "v2"})
report_usage = AsyncMock(return_value={"metered": True})
monkeypatch.setattr(workflow_run_billing_mod, "DEPLOYMENT_MODE", "saas")
@ -206,11 +162,6 @@ async def test_report_completed_workflow_run_platform_usage_loads_run(monkeypatc
"get_workflow_run_by_id",
get_run,
)
monkeypatch.setattr(
workflow_run_billing_mod.mps_service_key_client,
"get_billing_account_status",
get_status,
)
monkeypatch.setattr(
workflow_run_billing_mod.mps_service_key_client,
"report_platform_usage",