diff --git a/api/ee/src/core/organizations/service.py b/api/ee/src/core/organizations/service.py index e04aada4636..f8dcfd84631 100644 --- a/api/ee/src/core/organizations/service.py +++ b/api/ee/src/core/organizations/service.py @@ -963,6 +963,9 @@ async def _to_response( scope_from, Gauge, ) +from ee.src.core.starter_credits_bridge.service import ( # noqa: E402 + seed_starter_credits_bridge_safely, +) _subscription_service = SubscriptionsService( @@ -1001,6 +1004,14 @@ async def provision_signup_subscription( scope=scope_from(organization_id=organization.id), ) + # Bounded (10s) and swallow-all: a failure degrades to "no starter credits" + # and must never raise here — the signup path deletes the new user when + # setup fails. Nothing retries a failed seed. + await seed_starter_credits_bridge_safely( + organization_id=str(organization.id), + organization_email=organization_email, + ) + async def provision_user_subscription(organization: OrganizationDB) -> None: """Start the default plan + seed the user gauge for an explicitly-created org. diff --git a/api/ee/src/core/starter_credits_bridge/__init__.py b/api/ee/src/core/starter_credits_bridge/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/api/ee/src/core/starter_credits_bridge/__init__.py @@ -0,0 +1 @@ + diff --git a/api/ee/src/core/starter_credits_bridge/client.py b/api/ee/src/core/starter_credits_bridge/client.py new file mode 100644 index 00000000000..bf1e0c4746f --- /dev/null +++ b/api/ee/src/core/starter_credits_bridge/client.py @@ -0,0 +1,136 @@ +import re +from typing import Any, Optional + +import httpx + +from ee.src.core.starter_credits_bridge.types import ( + KeyAliasExistsError, + MintedKey, + ProxyRequestError, +) + +_REQUEST_TIMEOUT_SECONDS = 10.0 + +_KEY_PATTERN = re.compile(r"sk-[A-Za-z0-9_\-]+") + + +class StarterCreditsProxyClient: + """Admin client for the starter-credits proxy (mint / block keys, team info). + + Authenticates with the master key, which must never leave the backend. + """ + + def __init__( + self, + *, + base_url: str, + master_key: str, + transport: Optional[httpx.AsyncBaseTransport] = None, + ): + self._base_url = base_url.rstrip("/") + self._master_key = master_key + self._transport = transport + + def _http_client(self) -> httpx.AsyncClient: + return httpx.AsyncClient( + headers={"Authorization": f"Bearer {self._master_key}"}, + timeout=_REQUEST_TIMEOUT_SECONDS, + transport=self._transport, + ) + + async def generate_key( + self, + *, + key_alias: str, + max_budget: float, + models: list[str], + metadata: dict[str, Any], + team_id: str, + max_parallel_requests: Optional[int] = None, + rpm_limit: Optional[int] = None, + tpm_limit: Optional[int] = None, + ) -> MintedKey: + body: dict[str, Any] = { + "key_alias": key_alias, + "max_budget": max_budget, + # An explicit list always; an omitted list would mean "any model". + "models": models, + "metadata": metadata, + # Always under the program team so its ceiling bounds total exposure. + "team_id": team_id, + } + if max_parallel_requests is not None: + body["max_parallel_requests"] = max_parallel_requests + if rpm_limit is not None: + body["rpm_limit"] = rpm_limit + if tpm_limit is not None: + body["tpm_limit"] = tpm_limit + + payload = await self._request("POST", "/key/generate", json=body) + + key = payload.get("key") if isinstance(payload, dict) else None + if not isinstance(key, str) or not key: + raise ProxyRequestError( + status_code=200, + detail="key generation response carried no key", + ) + + return MintedKey(key=key, key_alias=key_alias) + + async def get_team_info(self, *, team_id: str) -> dict[str, Any]: + payload = await self._request("GET", "/team/info", params={"team_id": team_id}) + return payload if isinstance(payload, dict) else {} + + async def block_key(self, *, key: str) -> None: + await self._request("POST", "/key/block", json={"key": key}) + + async def _request( + self, + method: str, + path: str, + *, + json: Optional[dict[str, Any]] = None, + params: Optional[dict[str, Any]] = None, + ) -> Any: + try: + async with self._http_client() as client: + response = await client.request( + method, + f"{self._base_url}{path}", + json=json, + params=params, + ) + except httpx.HTTPError as exc: + raise ProxyRequestError( + status_code=None, + detail=f"request to {path} failed: {type(exc).__name__}", + ) from exc + + if response.status_code >= 400: + # Redact key material before the body can reach logs (a proxy error + # may echo the failing request, which can carry a virtual key). + detail = _KEY_PATTERN.sub("sk-[redacted]", response.text)[:300] + # Only the measured conflict wording may read as "this organization + # already holds its key"; a validation 400 that merely mentions the + # alias field must not. + if ( + response.status_code == 400 + and "alias" in detail.lower() + and "already exists" in detail.lower() + ): + raise KeyAliasExistsError( + status_code=response.status_code, + detail=detail, + ) + raise ProxyRequestError( + status_code=response.status_code, + detail=detail, + ) + + try: + return response.json() + except ValueError as exc: + raise ProxyRequestError( + status_code=response.status_code, + detail=f"non-JSON response from {path}", + ) from exc diff --git a/api/ee/src/core/starter_credits_bridge/service.py b/api/ee/src/core/starter_credits_bridge/service.py new file mode 100644 index 00000000000..a940703f22c --- /dev/null +++ b/api/ee/src/core/starter_credits_bridge/service.py @@ -0,0 +1,696 @@ +import asyncio +import json +import math +import time +from datetime import datetime, timezone +from typing import Any, Optional +from uuid import UUID + +import httpx + +from oss.src.services import db_manager +from oss.src.utils.caching import get_cache, set_cache +from oss.src.utils.env import env, StarterCreditsBridgeConfig +from oss.src.utils.lazy import _load_posthog +from oss.src.utils.logging import get_module_logger +from oss.src.dbs.redis.shared.engine import get_cache_engine +from oss.src.core.secrets.dtos import ( + CreateSecretDTO, + CustomModelSettingsDTO, + CustomProviderDTO, + CustomProviderSettingsDTO, + SecretDTO, + SecretKind, + SecretResponseDTO, +) +from oss.src.core.secrets.enums import CustomProviderKind +from oss.src.core.secrets.managed import ( + SecretManagementDTO, + SecretManagementPolicy, + SecretManager, +) +from oss.src.core.secrets.services import VaultService +from oss.src.dbs.postgres.secrets.dao import SecretsDAO +from oss.src.core.shared.dtos import Header + +from ee.src.core.starter_credits_bridge.client import StarterCreditsProxyClient +from ee.src.core.starter_credits_bridge.types import ( + DEVELOPMENT_POLICY_VALUES, + KeyAliasExistsError, + MintedKey, + MintPolicy, + ProxyRequestError, +) + +log = get_module_logger(__name__) + +# Slug of the seeded vault connection; with the project+slug unique index it makes +# the vault write idempotent. NOT proof of ownership: a user can delete and +# recreate the slug, so nothing ownership-sensitive reads it. +STARTER_CREDITS_SLUG = "starter-credits" + +# The connection's display name, and therefore the namespace of the stored model +# keys ("/custom/"). It is permanent PER ROW: an organization keeps +# the name and the keys it was seeded with, so changing this renames nothing that +# already exists — it sets what the next organization gets. Changing it after a +# real rollout would leave two namespaces in the field, each correct for its own +# organizations, so treat it as settled once seeding runs anywhere real. +STARTER_CREDITS_NAME = "Agenta" + +# Origin carried in the proxy key's metadata (with the org id). Key metadata is +# master-key-only mutable, so it supports later operator-side inspection. +PROXY_ORIGIN = "starter-credits-bridge" + +# Human-facing description stored on the connection. +STARTER_CREDITS_DESCRIPTION = "Provided and managed by Agenta." + +# The grant invariant: we mint at most one key per organization, guarded by the +# alias (one key per org id on the proxy) and by the slug's unique index in the +# vault. A failed seed is never retried by a repair — the organization simply +# stays unseeded and meets the connect-your-key wall as before. +_SEED_TIMEOUT_SECONDS = 10.0 + +# The mint is the one call worth retrying inside the seeding bound: a dropped +# connection or a proxy restart is common enough to cost real signups, and a +# retry cannot double-grant (the alias conflict below stops the second key). +_MINT_ATTEMPTS = 3 +_MINT_RETRY_BACKOFF_SECONDS = 0.2 + +_VELOCITY_COUNTER_TTL_SECONDS = 2 * 24 * 3600 + +# Re-verify the team ceiling after this long, so removing or loosening the +# ceiling is caught without a process restart. +_TEAM_VERIFY_TTL_SECONDS = 600.0 + +_verified_teams: dict[str, float] = {} + +# asyncio only weak-refs scheduled tasks; hold strong refs until each finishes. +_background_tasks: set[asyncio.Task] = set() + + +async def seed_starter_credits_bridge_safely( + *, + organization_id: str, + organization_email: str, +) -> None: + """Bounded, swallow-all wrapper for the signup hook: a seeding failure or a + slow dependency degrades to "no starter credits", never to a broken signup + (the signup path deletes the new user when setup raises). The alert is + fire-and-forget so it cannot extend the bound.""" + try: + async with asyncio.timeout(_SEED_TIMEOUT_SECONDS): + await seed_starter_credits_bridge( + organization_id=organization_id, + organization_email=organization_email, + ) + except Exception: + log.warning( + "[starter_credits_bridge] seed_failed; organization stays unseeded", + organization_id=organization_id, + exc_info=True, + ) + _send_alert_background( + f"starter-credits-bridge: seed_failed for organization {organization_id}; " + "it stays unseeded" + ) + + +async def seed_starter_credits_bridge( + *, + organization_id: str, + organization_email: str, +) -> None: + """Mint a budget-capped virtual key and seed it into the org's default-project + vault as a ready-to-use provider connection.""" + config = env.starter_credits_bridge + # Two kill switches, both program-wide: AGENTA_STARTER_CREDITS_BRIDGE_ENABLED + # (needs a redeploy) and the PostHog policy payload below — emptying or deleting + # it leaves the policy unresolved and stops seeding without one. `armed` also + # requires both proxy addresses (public for the seeded row, admin for minting). + if not config.armed: + return + + # The row this seeds is write-only, and a write-only secret is readable only through + # a credential the platform runtime is issued — which the API issues only to a caller + # holding the runtime key. Without it, seeding would mint a funded key and store it + # where no run can reach it: money spent on a connection that cannot work. Refuse + # instead, and say which variable is missing. + if not _platform_runtime_key_configured(): + log.error( + "[starter_credits_bridge] AGENTA_SERVICES_INTERNAL_KEY is not configured; " + "a seeded connection would be unreadable by any run. Refusing to seed.", + organization_id=organization_id, + ) + _alert_once_about_the_runtime_key() + return + + policy = await _resolve_mint_policy() + if policy is None: + return + + client = _proxy_client(config) + if not await _team_ceiling_verified(client, config): + return + + project = await db_manager.get_default_project_by_organization_id(organization_id) + if project is None: + log.warning( + "[starter_credits_bridge] no default project; skipping seed", + organization_id=organization_id, + ) + return + + vault_service = _vault_service() + row = await vault_service.get_secret_by_slug( + STARTER_CREDITS_SLUG, + project_id=project.id, + ) + if row is not None: + return + + if not await _mint_policy_allows(organization_email, policy): + return + + seeded = False + try: + seeded = await _provision( + client=client, + config=config, + policy=policy, + vault_service=vault_service, + project_id=project.id, + organization_id=organization_id, + ) + finally: + if not seeded: + # The consumed velocity slot funded no mint; hand it back best-effort. + await _release_velocity_slots(organization_email, policy) + + +async def _provision( + *, + client: StarterCreditsProxyClient, + config: StarterCreditsBridgeConfig, + policy: MintPolicy, + vault_service: VaultService, + project_id: UUID, + organization_id: str, +) -> bool: + try: + minted = await _mint_key( + client=client, + config=config, + policy=policy, + organization_id=organization_id, + ) + except KeyAliasExistsError: + # One alias per organization, so a conflict means this org already holds + # its key (a duplicate signup race). Never re-mint. + log.info( + "[starter_credits_bridge] alias already minted; organization is seeded", + organization_id=organization_id, + ) + return False + + try: + row = await _create_row( + vault_service=vault_service, + project_id=project_id, + config=config, + virtual_key=minted.key, + ) + except Exception: + # The key is live but nothing references it; block it so an orphaned + # grant can never be spent. + await _block_key(client, minted.key) + raise + + log.info( + "[starter_credits_bridge] seeded starter credits", + organization_id=organization_id, + project_id=str(project_id), + secret_id=str(row.id), + ) + return True + + +async def _mint_key( + *, + client: StarterCreditsProxyClient, + config: StarterCreditsBridgeConfig, + policy: MintPolicy, + organization_id: str, +) -> MintedKey: + """Mint the org's one key, retrying only transport failures and proxy 5xx. + A 4xx is the proxy's verdict on this request (an alias conflict, a rejected + body) and repeating it can only waste the seeding bound.""" + for attempt in range(_MINT_ATTEMPTS): + try: + # The proxy caps what a key may ask for (`upperbound_key_generate_params`: + # max_budget 5, max_parallel_requests 2, rpm 30, tpm 200000, duration 90d) + # and fills an OMITTED duration with its cap, so every funded key expires + # after 90 days. Send no duration rather than a longer one. Raising + # `grant_usd` or a per-key limit in the policy payload ABOVE a cap makes + # every mint fail with HTTP 400 until the proxy config is bumped and + # redeployed; lowering a value is a live payload edit. + return await client.generate_key( + key_alias=organization_id, + max_budget=policy.grant_usd, + models=[config.model_id], + metadata={ + "organization_id": organization_id, + "origin": PROXY_ORIGIN, + }, + team_id=config.team_id, + max_parallel_requests=policy.key_max_parallel_requests, + rpm_limit=policy.key_rpm_limit, + tpm_limit=policy.key_tpm_limit, + ) + except ProxyRequestError as exc: + if not _is_transient(exc) or attempt == _MINT_ATTEMPTS - 1: + raise + log.warning( + "[starter_credits_bridge] transient mint failure; retrying", + organization_id=organization_id, + attempt=attempt + 1, + status_code=exc.status_code, + ) + await asyncio.sleep(_MINT_RETRY_BACKOFF_SECONDS * (2**attempt)) + + raise ProxyRequestError(status_code=None, detail="mint retries exhausted") + + +def _is_transient(error: ProxyRequestError) -> bool: + return error.status_code is None or error.status_code >= 500 + + +async def _block_key(client: StarterCreditsProxyClient, key: str) -> None: + """Best-effort block; a failed block leaves a key nothing can reach through + the product, still bounded by its own budget and the team ceiling.""" + try: + await client.block_key(key=key) + except Exception: + log.warning( + "[starter_credits_bridge] could not block the orphaned key", + exc_info=True, + ) + + +async def _create_row( + *, + vault_service: VaultService, + project_id: UUID, + config: StarterCreditsBridgeConfig, + virtual_key: str, +) -> SecretResponseDTO: + data = CustomProviderDTO( + kind=CustomProviderKind.CUSTOM, + provider=CustomProviderSettingsDTO( + # The PUBLIC address: this is what a run dials, and a sandboxed run + # is outside the proxy's network. + url=config.proxy_public_url, + key=virtual_key, + ), + models=[CustomModelSettingsDTO(slug=config.model_id)], + # The namespace half of every model key this connection publishes + # (`//`), and it must equal the display name: + # that is the namespace the resolver rebuilds keys under, and a model key is + # permanent once a config references it. `CreateSecretDTO` fills this in from + # the header only when the payload arrives as plain dicts, which is not how + # this call builds it, so it is set here explicitly. + provider_slug=STARTER_CREDITS_NAME, + ).model_dump() + + return await vault_service.create_managed_secret( + project_id=project_id, + create_secret_dto=CreateSecretDTO( + slug=STARTER_CREDITS_SLUG, + header=Header( + name=STARTER_CREDITS_NAME, + description=STARTER_CREDITS_DESCRIPTION, + ), + secret=SecretDTO( + kind=SecretKind.CUSTOM_PROVIDER, + data=data, + ), + write_only=True, + ), + management=SecretManagementDTO( + manager=SecretManager.STARTER_CREDITS_BRIDGE, + policy=SecretManagementPolicy.MANAGER_ONLY, + ), + ) + + +# --- gates ----------------------------------------------------------------- + + +def _proxy_client(config: StarterCreditsBridgeConfig) -> StarterCreditsProxyClient: + # The ADMIN address, never the public one: the proxy publishes only its + # inference paths publicly, so /key/generate, /key/block and /team/info are + # reachable on the internal network alone. + return StarterCreditsProxyClient( + base_url=config.proxy_admin_url, + master_key=config.master_key, + ) + + +def _vault_service() -> VaultService: + return VaultService(SecretsDAO()) + + +async def _team_ceiling_verified( + client: StarterCreditsProxyClient, + config: StarterCreditsBridgeConfig, +) -> bool: + """Refuse to mint unless the program team stands with a numeric, finite, + non-resetting budget ceiling — the always-on bound on total exposure. Fail + closed; a positive result is trusted only for a short TTL so a removed or + loosened ceiling is caught without a restart.""" + team_id = config.team_id + verified_at = _verified_teams.get(team_id) + if ( + verified_at is not None + and _monotonic() - verified_at < _TEAM_VERIFY_TTL_SECONDS + ): + return True + + try: + payload = await client.get_team_info(team_id=team_id) + except Exception as exc: + log.error( + "[starter_credits_bridge] team ceiling unverifiable; refusing to seed", + team_id=team_id, + reason=str(exc), + ) + _send_alert_background( + f"starter-credits-bridge: team '{team_id}' unverifiable; seeding refused" + ) + return False + + team_info = payload.get("team_info") or payload + max_budget = team_info.get("max_budget") + budget_duration = team_info.get("budget_duration") + + if ( + not isinstance(max_budget, (int, float)) + or isinstance(max_budget, bool) + or not _is_finite_positive(max_budget) + ): + log.error( + "[starter_credits_bridge] team has no sound max_budget; refusing to seed", + team_id=team_id, + ) + _send_alert_background( + f"starter-credits-bridge: team '{team_id}' has no budget ceiling; seeding refused" + ) + return False + + if budget_duration: + # A duration makes the ceiling reset periodically — that is a rate, not + # the program's total-exposure bound. + log.error( + "[starter_credits_bridge] team budget resets; refusing to seed", + team_id=team_id, + budget_duration=budget_duration, + ) + _send_alert_background( + f"starter-credits-bridge: team '{team_id}' budget resets " + f"(duration {budget_duration}); seeding refused" + ) + return False + + _verified_teams[team_id] = _monotonic() + return True + + +def _is_finite_positive(value: float) -> bool: + return math.isfinite(value) and value > 0 + + +def _monotonic() -> float: + return time.monotonic() + + +# The development-policy notice is worth seeing, and worth seeing once: it would +# otherwise repeat on every signup for the whole life of a dev process. +_development_policy_announced = False + + +_UNCONFIGURED_RUNTIME_KEY = "replace-me" +_runtime_key_alerted = False + + +def _platform_runtime_key_configured() -> bool: + """Whether this deployment can issue a credential that reads a write-only secret.""" + runtime_key = (env.agenta.services_internal_key or "").strip() + + return bool(runtime_key) and runtime_key != _UNCONFIGURED_RUNTIME_KEY + + +def _alert_once_about_the_runtime_key() -> None: + """Page the operator the first time, then stay quiet: every signup would repeat it.""" + global _runtime_key_alerted + + if _runtime_key_alerted: + return + + _runtime_key_alerted = True + _send_alert_background( + "starter-credits-bridge: AGENTA_SERVICES_INTERNAL_KEY is not configured; " + "seeding is refused because the seeded connection would be unreadable by any run" + ) + + +def _development_policy() -> MintPolicy: + """The policy a deployment without PostHog runs on. Values live in `types.py`.""" + global _development_policy_announced + + if not _development_policy_announced: + _development_policy_announced = True + log.warning( + "starter credits: PostHog not configured, using the built-in development policy" + ) + + return MintPolicy(**DEVELOPMENT_POLICY_VALUES) + + +async def _resolve_mint_policy() -> Optional[MintPolicy]: + """Resolve the mint policy. Live-first; a MALFORMED live payload fails + closed with an alert (a bad rollout must never silently keep old caps via + the cache) — only a transport failure may fall back to the Redis-cached + payload. The payload is the only source of policy values, except on a + deployment that configured no PostHog of its own, which runs on the built-in + development policy. No resolvable policy means no seeding.""" + flag = env.starter_credits_bridge.policy_flag + # Deliberately global: the mint policy is one program-wide payload (caps, domain + # rules), identical for every organization. + cache_key = {"ff": flag} + + # Deliberately NOT `env.posthog.enabled`: the PostHog config falls back to a + # built-in project key, so `enabled` is true in every checkout and could never + # tell a local stack from a deployment that runs its own PostHog. A deployment + # that supplied no key of its own has no way to publish a policy, so failing + # closed would simply block it; one that did keeps failing closed below on a + # missing or malformed payload, which is what protects production. + if not env.posthog.api_key_configured: + return _development_policy() + + payload: Optional[dict] = None + live_malformed = False + posthog = _load_posthog() + if posthog is not None: + try: + raw = await asyncio.to_thread( + posthog.get_feature_flag_payload, flag, "starter-credits-bridge" + ) + except Exception as exc: + log.warning( + "[starter_credits_bridge] live policy lookup failed; using cached value", + reason=str(exc), + ) + else: + if raw is None: + # A reachable PostHog with no payload is a real "no policy" signal: + # nothing else can supply one, so the resolve below fails closed. + payload = None + else: + payload = _parse_policy_payload(raw) + if payload is None: + live_malformed = True + + if payload is None and not live_malformed: + cached = await get_cache( + namespace="starter_credits_bridge:policy", + key=cache_key, + retry=False, + ) + if isinstance(cached, dict): + payload = cached + + try: + policy = MintPolicy(**({} if live_malformed else (payload or {}))) + except Exception: + log.error( + "[starter_credits_bridge] mint policy missing or invalid; seeding disabled (fail closed)", + malformed_live_payload=live_malformed, + ) + if live_malformed: + _send_alert_background( + "starter-credits-bridge: live policy payload is malformed; seeding disabled" + ) + return None + + if payload is not None and not live_malformed: + await set_cache( + namespace="starter_credits_bridge:policy", + key=cache_key, + value=payload, + ) + return policy + + +def _parse_policy_payload(raw: Any) -> Optional[dict]: + if isinstance(raw, dict): + return raw + if isinstance(raw, str) and raw.strip(): + try: + parsed = json.loads(raw) + except ValueError: + return None + return parsed if isinstance(parsed, dict) else None + return None + + +def _velocity_counters( + organization_email: str, + policy: MintPolicy, +) -> list[tuple[str, int, int, str]]: + day = datetime.now(timezone.utc).strftime("%Y%m%d") + hour = datetime.now(timezone.utc).strftime("%Y%m%d%H") + domain = _email_domain(organization_email) + + counters = [ + ( + f"starter_credits_bridge:mints:{day}", + policy.global_daily, + _VELOCITY_COUNTER_TTL_SECONDS, + "global_daily", + ), + ( + f"starter_credits_bridge:mints:hour:{hour}", + policy.global_hourly, + 2 * 3600, + "global_hourly", + ), + ] + if not policy.is_freemail(domain): + counters.append( + ( + f"starter_credits_bridge:mints:{day}:domain:{domain}", + policy.work_domain_daily, + _VELOCITY_COUNTER_TTL_SECONDS, + "work_domain_daily", + ) + ) + return counters + + +def _email_domain(organization_email: str) -> str: + _, _, domain = organization_email.rpartition("@") + return domain.lower() or "unknown" + + +async def _mint_policy_allows( + organization_email: str, + policy: MintPolicy, +) -> bool: + """Apply the mint policy: eligibility rules, then Redis counters (global + daily + hourly; per-domain daily on non-free-mail domains only). Fail closed + on Redis errors: an unverifiable mint is a skipped mint.""" + local_part, _, _ = organization_email.rpartition("@") + domain = _email_domain(organization_email) + freemail = policy.is_freemail(domain) + + if ( + not freemail + and policy.block_digit_locals + and any(character.isdigit() for character in local_part) + ): + log.warning( + "[starter_credits_bridge] policy refused mint; skipping seed", + rule="digit_local_part", + # The domain, never the address: it is what makes a refusal diagnosable + # (which provider, which rule) without logging who signed up. + domain=domain, + ) + return False + + engine = get_cache_engine() + try: + for counter_key, cap, ttl, scope in _velocity_counters( + organization_email, policy + ): + count = await engine.incr(counter_key) + if count == 1: + await engine.expire(counter_key, ttl) + if count > cap: + log.warning( + "[starter_credits_bridge] velocity cap reached; skipping seed", + rule=scope, + domain=domain, + ) + return False + except Exception as exc: + log.warning( + "[starter_credits_bridge] velocity counters unavailable; skipping seed (fail closed)", + rule="velocity_counters_unavailable", + domain=domain, + reason=str(exc), + ) + return False + + return True + + +async def _release_velocity_slots( + organization_email: str, + policy: MintPolicy, +) -> None: + """Best-effort: hand consumed counter slots back when the attempt funded no + mint, so an outage does not eat the day's allowance.""" + engine = get_cache_engine() + for counter_key, _cap, _ttl, _scope in _velocity_counters( + organization_email, policy + ): + try: + await engine.decr(counter_key) + except Exception: + return + + +def _send_alert_background(text: str) -> None: + """Schedule the operator alert without blocking the caller; never raises.""" + try: + task = asyncio.create_task(_send_alert(text)) + task.add_done_callback(_background_tasks.discard) + _background_tasks.add(task) + except Exception: + log.warning("[starter_credits_bridge] could not schedule alert", exc_info=True) + + +async def _send_alert(text: str) -> None: + """Best-effort operator alert; failures only log.""" + webhook = env.starter_credits_bridge.alert_webhook + if not webhook: + return + + try: + async with httpx.AsyncClient(timeout=5.0) as client: + response = await client.post(webhook, json={"text": text}) + response.raise_for_status() + except Exception: + log.warning( + "[starter_credits_bridge] alert webhook failed", + exc_info=True, + ) diff --git a/api/ee/src/core/starter_credits_bridge/types.py b/api/ee/src/core/starter_credits_bridge/types.py new file mode 100644 index 00000000000..7fdf774fc6d --- /dev/null +++ b/api/ee/src/core/starter_credits_bridge/types.py @@ -0,0 +1,194 @@ +import math + +from pydantic import BaseModel, ConfigDict, field_validator + + +# Consumer mail providers. Classification drives the per-domain daily cap, which only +# means anything on a domain one company controls: unrecognized free mail counts every +# signup from it against a single "company" and, with `block_digit_locals`, judges a +# personal address by rules meant for a work one. A plain constant rather than config +# because it tracks the mail industry, not an operator decision, and because the policy +# payload UNIONS with it (see `_normalize_domains`) — extending the list must never +# silently drop gmail.com. +DEFAULT_FREEMAIL_DOMAINS: tuple[str, ...] = ( + "gmail.com", + "googlemail.com", + "yahoo.com", + "yahoo.co.uk", + "yahoo.co.in", + "yahoo.fr", + "yahoo.de", + "ymail.com", + "rocketmail.com", + "hotmail.com", + "hotmail.co.uk", + "hotmail.fr", + "hotmail.de", + "hotmail.it", + "outlook.com", + "outlook.de", + "outlook.fr", + "outlook.es", + "outlook.in", + "live.com", + "live.co.uk", + "msn.com", + "aol.com", + "aim.com", + "icloud.com", + "me.com", + "mac.com", + "proton.me", + "protonmail.com", + "pm.me", + "tutanota.com", + "tuta.com", + "fastmail.com", + "hey.com", + "zoho.com", + "mail.com", + "gmx.com", + "gmx.de", + "gmx.net", + "web.de", + "t-online.de", + "freenet.de", + "orange.fr", + "wanadoo.fr", + "free.fr", + "laposte.net", + "libero.it", + "virgilio.it", + "seznam.cz", + "wp.pl", + "onet.pl", + "interia.pl", + "mail.ru", + "yandex.ru", + "yandex.com", + "rambler.ru", + "qq.com", + "163.com", + "126.com", + "sina.com", + "naver.com", + "daum.net", + "hanmail.net", + "rediffmail.com", + "comcast.net", + "verizon.net", + "att.net", + "sbcglobal.net", + "cox.net", + "btinternet.com", + "sky.com", + "virginmedia.com", + "bigpond.com", + "uol.com.br", + "bol.com.br", +) + + +class MintPolicy(BaseModel): + """Mint policy: velocity caps, domain classification, eligibility rules, and + the money values (grant, per-key limits). Everything ships via the PostHog + policy flag payload so no real value lives in source; an unresolved or invalid + policy means no seeding. Unknown payload fields are rejected so a malformed + rollout fails closed instead of half-applying. The one exception is a + deployment with no PostHog at all, which gets DEVELOPMENT_POLICY_VALUES + below.""" + + # `validate_default` so an absent `freemail_domains` still picks up the built-in + # defaults through the validator below. + model_config = ConfigDict(extra="forbid", validate_default=True) + + global_daily: int + global_hourly: int + work_domain_daily: int + freemail_domains: list[str] = [] + block_digit_locals: bool + grant_usd: float + key_max_parallel_requests: int + key_rpm_limit: int + key_tpm_limit: int + + @field_validator("freemail_domains") + @classmethod + def _normalize_domains(cls, domains: list[str]) -> list[str]: + """Whatever the payload names, UNIONED with the built-in defaults. + + A configured list adds providers the defaults miss; it never replaces them, so a + rollout that names three domains cannot quietly reclassify gmail.com as a company + domain. + """ + configured = {domain.strip().lower() for domain in domains if domain.strip()} + + return sorted(configured | set(DEFAULT_FREEMAIL_DOMAINS)) + + @field_validator("grant_usd") + @classmethod + def _finite_positive_grant(cls, value: float) -> float: + if not math.isfinite(value) or value <= 0: + raise ValueError("grant_usd must be a finite positive number") + return value + + @field_validator( + "global_daily", + "global_hourly", + "work_domain_daily", + "key_max_parallel_requests", + "key_rpm_limit", + "key_tpm_limit", + ) + @classmethod + def _positive_ints(cls, value: int) -> int: + if value <= 0: + raise ValueError("policy limits must be positive") + return value + + def is_freemail(self, domain: str) -> bool: + return domain.lower() in self.freemail_domains + + +class StarterCreditsBridgeError(Exception): + """Base exception for starter-credits-bridge errors.""" + + +class ProxyRequestError(StarterCreditsBridgeError): + """The proxy admin API refused or failed a request. + + Never carries key material; `detail` is a short, log-safe diagnosis. + """ + + def __init__(self, *, status_code: int | None, detail: str): + self.status_code = status_code + self.detail = detail + super().__init__(f"proxy request failed ({status_code}): {detail}") + + +class KeyAliasExistsError(ProxyRequestError): + """A key with this alias was already minted (idempotency signal).""" + + +class MintedKey(BaseModel): + key: str + key_alias: str + + +# What a deployment with no PostHog runs on: local development and live QA, which +# would otherwise be blocked by the fail-closed rule with no way to unblock them. +# The numbers are deliberately generic — a small grant and caps loose enough not to +# get in the way of testing — so they say nothing about the real program, whose +# values live only in the PostHog payload. A deployment that HAS PostHog never +# reaches these: a missing or malformed payload there still fails closed. +DEVELOPMENT_POLICY_VALUES: dict = { + "global_daily": 1000, + "global_hourly": 1000, + "work_domain_daily": 1000, + "freemail_domains": [], + "block_digit_locals": False, + "grant_usd": 5.0, + "key_max_parallel_requests": 2, + "key_rpm_limit": 30, + "key_tpm_limit": 200_000, +} diff --git a/api/ee/src/main.py b/api/ee/src/main.py index 3261c2ebd0a..cb33c8b292e 100644 --- a/api/ee/src/main.py +++ b/api/ee/src/main.py @@ -125,7 +125,6 @@ records_retention_service=records_retention_service, ) - log = get_module_logger(__name__) diff --git a/api/ee/tests/pytest/unit/test_starter_credits_bridge_client.py b/api/ee/tests/pytest/unit/test_starter_credits_bridge_client.py new file mode 100644 index 00000000000..de9f33d4e49 --- /dev/null +++ b/api/ee/tests/pytest/unit/test_starter_credits_bridge_client.py @@ -0,0 +1,232 @@ +"""Unit tests for the starter-credits proxy admin client +(``ee.src.core.starter_credits_bridge.client``) against a mocked proxy. + +The base URL here is the proxy's INTERNAL address on purpose: admin routes are +not publicly served, so this client is always pointed at `proxy_admin_url`.""" + +import json + +import httpx +import pytest + +from ee.src.core.starter_credits_bridge.client import StarterCreditsProxyClient +from ee.src.core.starter_credits_bridge.types import ( + KeyAliasExistsError, + ProxyRequestError, +) + + +def _client_with_handler(handler) -> StarterCreditsProxyClient: + return StarterCreditsProxyClient( + base_url="https://proxy.internal.test", + master_key="sk-master-test", + transport=httpx.MockTransport(handler), + ) + + +def _read_body(request: httpx.Request) -> dict: + return json.loads(httpx.Request.read(request)) + + +class TestGenerateKey: + async def test_sends_master_key_and_required_body(self): + seen = {} + + def handler(request: httpx.Request) -> httpx.Response: + seen["url"] = str(request.url) + seen["auth"] = request.headers.get("Authorization") + seen["body"] = _read_body(request) + return httpx.Response(200, json={"key": "sk-virtual-abc"}) + + client = _client_with_handler(handler) + minted = await client.generate_key( + key_alias="org-123", + max_budget=10.0, + models=["some-model"], + metadata={"organization_id": "org-123", "origin": "starter-credits-bridge"}, + team_id="team-1", + ) + + assert minted.key == "sk-virtual-abc" + assert minted.key_alias == "org-123" + assert seen["url"] == "https://proxy.internal.test/key/generate" + assert seen["auth"] == "Bearer sk-master-test" + + body = seen["body"] + assert body["key_alias"] == "org-123" + assert body["max_budget"] == 10.0 + assert body["models"] == ["some-model"] + assert body["metadata"]["origin"] == "starter-credits-bridge" + assert body["team_id"] == "team-1" + assert "max_parallel_requests" not in body + assert "rpm_limit" not in body + assert "tpm_limit" not in body + + async def test_per_key_limits_are_sent(self): + seen = {} + + def handler(request: httpx.Request) -> httpx.Response: + seen["body"] = _read_body(request) + return httpx.Response(200, json={"key": "sk-virtual-abc"}) + + client = _client_with_handler(handler) + await client.generate_key( + key_alias="org-123", + max_budget=10.0, + models=["some-model"], + metadata={}, + team_id="team-1", + max_parallel_requests=2, + rpm_limit=30, + tpm_limit=200_000, + ) + + assert seen["body"]["max_parallel_requests"] == 2 + assert seen["body"]["rpm_limit"] == 30 + assert seen["body"]["tpm_limit"] == 200_000 + + async def test_server_error_raises_proxy_request_error(self): + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(500, text="internal error") + + client = _client_with_handler(handler) + with pytest.raises(ProxyRequestError) as excinfo: + await client.generate_key( + key_alias="org-123", + max_budget=1.0, + models=["some-model"], + metadata={}, + team_id="team-1", + ) + assert excinfo.value.status_code == 500 + + async def test_alias_conflict_raises_key_alias_exists(self): + # The measured v1.97.0 wording; only this shape may trigger the + # destructive delete-and-remint compensation. + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 400, + json={ + "error": { + "message": "Key with alias 'org-123' already exists. " + "Unique key aliases across all keys are required." + } + }, + ) + + client = _client_with_handler(handler) + with pytest.raises(KeyAliasExistsError): + await client.generate_key( + key_alias="org-123", + max_budget=1.0, + models=["some-model"], + metadata={}, + team_id="team-1", + ) + + async def test_validation_400_mentioning_alias_is_not_a_conflict(self): + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 400, + json={"error": {"message": "key_alias: value is not a valid string"}}, + ) + + client = _client_with_handler(handler) + with pytest.raises(ProxyRequestError) as excinfo: + await client.generate_key( + key_alias="org-123", + max_budget=1.0, + models=["some-model"], + metadata={}, + team_id="team-1", + ) + assert not isinstance(excinfo.value, KeyAliasExistsError) + + async def test_error_detail_redacts_key_material(self): + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 500, + text='update failed for {"key": "sk-virtual-secret-123"}', + ) + + client = _client_with_handler(handler) + with pytest.raises(ProxyRequestError) as excinfo: + await client.generate_key( + key_alias="org-123", + max_budget=1.0, + models=["some-model"], + metadata={}, + team_id="team-1", + ) + assert "sk-virtual-secret-123" not in excinfo.value.detail + assert "sk-[redacted]" in excinfo.value.detail + + async def test_missing_key_in_response_raises(self): + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"expires": None}) + + client = _client_with_handler(handler) + with pytest.raises(ProxyRequestError): + await client.generate_key( + key_alias="org-123", + max_budget=1.0, + models=["some-model"], + metadata={}, + team_id="team-1", + ) + + async def test_connection_failure_raises_proxy_request_error(self): + def handler(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("refused") + + client = _client_with_handler(handler) + with pytest.raises(ProxyRequestError) as excinfo: + await client.generate_key( + key_alias="org-123", + max_budget=1.0, + models=["some-model"], + metadata={}, + team_id="team-1", + ) + assert excinfo.value.status_code is None + + +class TestKeyLifecycle: + async def test_block_key_posts_to_key_block(self): + seen = {} + + def handler(request: httpx.Request) -> httpx.Response: + seen["url"] = str(request.url) + seen["body"] = _read_body(request) + return httpx.Response(200, json={}) + + client = _client_with_handler(handler) + await client.block_key(key="sk-virtual-abc") + + assert seen["url"] == "https://proxy.internal.test/key/block" + assert seen["body"] == {"key": "sk-virtual-abc"} + + +class TestTeamInfo: + async def test_get_team_info_queries_team_id(self): + def handler(request: httpx.Request) -> httpx.Response: + assert str(request.url).startswith("https://proxy.internal.test/team/info") + assert request.url.params["team_id"] == "team-1" + return httpx.Response( + 200, + json={"team_id": "team-1", "team_info": {"max_budget": 500.0}}, + ) + + client = _client_with_handler(handler) + payload = await client.get_team_info(team_id="team-1") + + assert payload["team_info"]["max_budget"] == 500.0 + + async def test_missing_team_raises(self): + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(404, json={"error": {"message": "team not found"}}) + + client = _client_with_handler(handler) + with pytest.raises(ProxyRequestError) as excinfo: + await client.get_team_info(team_id="team-1") + assert excinfo.value.status_code == 404 diff --git a/api/ee/tests/pytest/unit/test_starter_credits_bridge_seeding.py b/api/ee/tests/pytest/unit/test_starter_credits_bridge_seeding.py new file mode 100644 index 00000000000..b53e912464a --- /dev/null +++ b/api/ee/tests/pytest/unit/test_starter_credits_bridge_seeding.py @@ -0,0 +1,1141 @@ +"""Unit tests for the starter-credits-bridge seeding service +(``ee.src.core.starter_credits_bridge.service``): gating, team-ceiling refusal, +the mint-then-write seed and its failure behavior (transient retry, orphan +blocking, alias conflicts), policy resolution, and velocity caps, with every +external dependency stubbed.""" + +import asyncio +from types import SimpleNamespace +from uuid import uuid4 + +import pytest + +from oss.src.utils.env import env, PostHogConfig, StarterCreditsBridgeConfig +from oss.src.core.secrets.dtos import SecretResponseDTO +from oss.src.core.secrets.enums import CustomProviderKind +from oss.src.core.secrets.managed import ( + SecretManagementDTO, + SecretManagementPolicy, + SecretManager, +) + +from ee.src.core.starter_credits_bridge import service +from ee.src.core.starter_credits_bridge.types import ( + KeyAliasExistsError, + MintedKey, + MintPolicy, + ProxyRequestError, +) + + +# Captured before any fixture replaces it, so the one end-to-end policy test below +# can run the real resolver. +_REAL_RESOLVE_MINT_POLICY = service._resolve_mint_policy + +ORGANIZATION_ID = "9f0e0f39-0000-4000-8000-000000000001" +ORGANIZATION_EMAIL = "someone@example.com" + + +def _armed_config(**overrides) -> StarterCreditsBridgeConfig: + values = dict( + enabled=True, + proxy_public_url="https://credits-proxy.example.test", + proxy_admin_url="http://litellm-proxy:4000", + master_key="sk-master-test", + team_id="team-starter", + model_id="vertex_ai/some-model", + ) + values.update(overrides) + return StarterCreditsBridgeConfig(**values) + + +# Synthetic policy values for tests only; the real values ship via the PostHog +# payload and are deliberately absent from source. +def _policy(**overrides) -> MintPolicy: + values = dict( + global_daily=4, + global_hourly=3, + work_domain_daily=1, + freemail_domains=["freemail.test"], + block_digit_locals=True, + grant_usd=10.0, + key_max_parallel_requests=2, + key_rpm_limit=30, + key_tpm_limit=200_000, + ) + values.update(overrides) + return MintPolicy(**values) + + +class FakeVaultService: + """Stores real validated DTOs and enforces the project+slug unique index.""" + + def __init__(self): + self.row = None + self.created_count = 0 + self.create_dto = None + self.management = None + self.create_error = None + + async def get_secret_by_slug(self, secret_slug, project_id=None, **kwargs): + assert secret_slug == service.STARTER_CREDITS_SLUG + return self.row + + async def create_managed_secret(self, *, project_id, create_secret_dto, management): + await asyncio.sleep(0) + if self.create_error is not None: + raise self.create_error + if self.row is not None: + raise RuntimeError("duplicate slug (unique index)") + self.create_dto = create_secret_dto + self.management = management + self.row = SimpleNamespace( + id=uuid4(), + header=create_secret_dto.header, + data=create_secret_dto.secret.data, + management=management, + write_only=bool(create_secret_dto.write_only), + ) + self.created_count += 1 + return self.row + + +class FakeProxyClient: + """Stands in for StarterCreditsProxyClient with a per-alias key registry, so + the alias conflict that guards the one-key-per-organization invariant is + exercised against realistic proxy state.""" + + records: dict = {} + generate_failures: list = [] + instances: list = [] + + def __init__(self, *, base_url, master_key): + self.base_url = base_url + self.master_key = master_key + self.generate_calls = [] + self.block_calls = [] + FakeProxyClient.instances.append(self) + + @classmethod + def fail_generate(cls, times: int, *, status_code: int = 500): + """Queue `times` mint failures before the next success.""" + cls.generate_failures = [ + ProxyRequestError(status_code=status_code, detail="mint failed") + for _ in range(times) + ] + + async def generate_key( + self, + *, + key_alias, + max_budget, + models, + metadata, + team_id, + max_parallel_requests=None, + rpm_limit=None, + tpm_limit=None, + ): + await asyncio.sleep(0) + self.generate_calls.append( + dict( + key_alias=key_alias, + max_budget=max_budget, + models=models, + metadata=metadata, + team_id=team_id, + max_parallel_requests=max_parallel_requests, + rpm_limit=rpm_limit, + tpm_limit=tpm_limit, + ) + ) + if FakeProxyClient.generate_failures: + raise FakeProxyClient.generate_failures.pop(0) + if key_alias in FakeProxyClient.records: + raise KeyAliasExistsError(status_code=400, detail="already exists: alias") + value = f"sk-virtual-{uuid4().hex[:8]}" + FakeProxyClient.records[key_alias] = { + "key_alias": key_alias, + "key": value, + "max_budget": max_budget, + "spend": 0.0, + "metadata": dict(metadata), + } + return MintedKey(key=value, key_alias=key_alias) + + async def block_key(self, *, key): + self.block_calls.append(key) + + +def _all_generate_calls(): + return [ + call + for instance in FakeProxyClient.instances + for call in instance.generate_calls + ] + + +def _all_block_calls(): + return [ + call for instance in FakeProxyClient.instances for call in instance.block_calls + ] + + +@pytest.fixture +def seeding_env(monkeypatch): + """Arm the config and stub every dependency; returns the mutable stubs.""" + FakeProxyClient.records = {} + FakeProxyClient.generate_failures = [] + FakeProxyClient.instances = [] + service._verified_teams.clear() + + config = _armed_config() + monkeypatch.setattr(env, "starter_credits_bridge", config) + # A deployment that can issue a credential able to read the write-only row it seeds; + # the refusal when it cannot is its own test below. + monkeypatch.setattr(env.agenta, "services_internal_key", "runtime-key-for-tests") + monkeypatch.setattr(service, "_runtime_key_alerted", False) + # The retry backoff is behavior under test only for its shape, not its wall time. + monkeypatch.setattr(service, "_MINT_RETRY_BACKOFF_SECONDS", 0.0) + + project = SimpleNamespace(id=uuid4()) + vault = FakeVaultService() + alerts = [] + released = [] + policy = _policy() + + async def fake_get_default_project(organization_id): + assert organization_id == ORGANIZATION_ID + return project + + async def fake_resolve_policy(): + return policy + + async def fake_team_verified(client, config): + return True + + async def fake_policy_allows(organization_email, policy): + return True + + def fake_send_alert_background(text): + alerts.append(text) + + async def fake_release(organization_email, policy): + released.append(organization_email) + + monkeypatch.setattr( + service.db_manager, + "get_default_project_by_organization_id", + fake_get_default_project, + ) + monkeypatch.setattr(service, "_vault_service", lambda: vault) + monkeypatch.setattr(service, "_resolve_mint_policy", fake_resolve_policy) + monkeypatch.setattr(service, "_team_ceiling_verified", fake_team_verified) + monkeypatch.setattr(service, "_mint_policy_allows", fake_policy_allows) + monkeypatch.setattr(service, "_send_alert_background", fake_send_alert_background) + monkeypatch.setattr(service, "_release_velocity_slots", fake_release) + monkeypatch.setattr(service, "StarterCreditsProxyClient", FakeProxyClient) + + return SimpleNamespace( + config=config, + policy=policy, + project=project, + vault=vault, + alerts=alerts, + released=released, + monkeypatch=monkeypatch, + ) + + +async def _seed(): + await service.seed_starter_credits_bridge( + organization_id=ORGANIZATION_ID, + organization_email=ORGANIZATION_EMAIL, + ) + + +async def _seed_safely(): + await service.seed_starter_credits_bridge_safely( + organization_id=ORGANIZATION_ID, + organization_email=ORGANIZATION_EMAIL, + ) + + +class TestSeeding: + async def test_happy_path_mints_then_writes_the_row(self, seeding_env): + await _seed() + + row = seeding_env.vault.row + assert row is not None + (call,) = _all_generate_calls() + assert call["key_alias"] == ORGANIZATION_ID + assert call["max_budget"] == 10.0 + assert call["models"] == ["vertex_ai/some-model"] + assert call["team_id"] == "team-starter" + assert call["max_parallel_requests"] == 2 + assert call["rpm_limit"] == 30 + assert call["tpm_limit"] == 200_000 + assert call["metadata"] == { + "organization_id": ORGANIZATION_ID, + "origin": service.PROXY_ORIGIN, + } + + assert row.data.kind == CustomProviderKind.CUSTOM + # The row carries the PUBLIC address: a sandboxed run dials it from + # outside, while minting used the internal admin one. + assert row.data.provider.url == "https://credits-proxy.example.test" + (client,) = FakeProxyClient.instances + assert client.base_url == "http://litellm-proxy:4000" + assert row.data.provider.key.startswith("sk-virtual-") + assert row.data.provider.key == FakeProxyClient.records[ORGANIZATION_ID]["key"] + assert [model.slug for model in row.data.models] == ["vertex_ai/some-model"] + assert row.header.name == service.STARTER_CREDITS_NAME + assert row.header.description == service.STARTER_CREDITS_DESCRIPTION + assert service.PROXY_ORIGIN not in row.header.description + assert seeding_env.config.proxy_public_url not in row.header.description + assert seeding_env.released == [] + + async def test_the_seeded_row_keys_its_models_under_the_display_name( + self, seeding_env + ): + # The display name is the namespace half of every model key this connection + # publishes, and a key is permanent once a config references it: a row seeded + # with no `provider_slug` published "None/custom/" forever. + await _seed() + + created = seeding_env.vault.create_dto + # `mode="json"` is the shape the row is stored and read back in. + response = SecretResponseDTO( + id=uuid4(), + slug=service.STARTER_CREDITS_SLUG, + kind=created.secret.kind, + data=created.secret.data.model_dump(mode="json"), + header=created.header, + ) + + assert response.data.model_keys == [ + f"{service.STARTER_CREDITS_NAME}/custom/vertex_ai/some-model" + ] + + async def test_the_resolver_reads_the_seeded_keys_and_strips_the_namespace( + self, seeding_env + ): + # The resolver takes the STORED keys as they are (it only rebuilds when a row + # has none), and strips the namespace off the chosen one to get the model id + # the harness runs. So the namespace the row was seeded with is the namespace + # the runtime sees, and it has to be the display name. + from agenta.sdk.agents.connections import ModelRef + from agenta.sdk.agents.platform import connections + + await _seed() + + created = seeding_env.vault.create_dto + stored = SecretResponseDTO( + id=uuid4(), + slug=service.STARTER_CREDITS_SLUG, + kind=created.secret.kind, + data=created.secret.data.model_dump(mode="json"), + header=created.header, + ) + # The row as the vault routes return it, model keys included. + (candidate,) = connections._catalog([stored.model_dump(mode="json")]) + + assert candidate.model_keys == { + f"{service.STARTER_CREDITS_NAME}/custom/vertex_ai/some-model" + } + + model = ModelRef( + model=f"{service.STARTER_CREDITS_NAME}/custom/vertex_ai/some-model", + connection={"mode": "agenta", "slug": service.STARTER_CREDITS_SLUG}, + ) + assert candidate.matches_model(model) is True + assert candidate.selected_model_id(model) == "vertex_ai/some-model" + + async def test_a_deployment_without_posthog_still_seeds(self, seeding_env): + # The development-policy path end to end: a local stack has no way to publish + # a payload, and a signup there must still get its funded connection. + seeding_env.monkeypatch.setattr( + service, "_resolve_mint_policy", _REAL_RESOLVE_MINT_POLICY + ) + seeding_env.monkeypatch.setattr( + env, "posthog", PostHogConfig(api_key_configured=False) + ) + seeding_env.monkeypatch.setattr(service, "_development_policy_announced", False) + + await _seed() + + assert seeding_env.vault.row is not None + (client,) = FakeProxyClient.instances + (call,) = client.generate_calls + assert call["max_budget"] == 5.0 + assert call["rpm_limit"] == 30 + + @pytest.mark.parametrize("runtime_key", ["", "replace-me"]) + async def test_an_unreadable_deployment_refuses_to_seed( + self, seeding_env, runtime_key + ): + # Without a runtime key the API cannot issue a credential that reads a write-only + # secret, so a seeded row would be a funded key nothing can spend. Refuse rather + # than mint into a connection that cannot work. + seeding_env.monkeypatch.setattr( + env.agenta, "services_internal_key", runtime_key + ) + + await _seed() + + assert FakeProxyClient.instances == [] + assert seeding_env.vault.row is None + assert len(seeding_env.alerts) == 1 + assert "AGENTA_SERVICES_INTERNAL_KEY" in seeding_env.alerts[0] + + async def test_the_unreadable_deployment_alert_is_sent_once(self, seeding_env): + # One alert, not one per signup. + seeding_env.monkeypatch.setattr(env.agenta, "services_internal_key", "") + + await _seed() + await _seed() + + assert len(seeding_env.alerts) == 1 + + async def test_disarmed_config_is_a_noop(self, seeding_env): + seeding_env.monkeypatch.setattr( + env, "starter_credits_bridge", _armed_config(enabled=False) + ) + + await _seed() + + assert FakeProxyClient.instances == [] + assert seeding_env.vault.row is None + + async def test_missing_team_id_disarms(self, seeding_env): + seeding_env.monkeypatch.setattr( + env, "starter_credits_bridge", _armed_config(team_id=None) + ) + + await _seed() + + assert FakeProxyClient.instances == [] + + async def test_missing_admin_url_disarms(self, seeding_env): + # Every admin route lives on the internal address; without it a mint + # would dial a path the proxy does not serve publicly. + seeding_env.monkeypatch.setattr( + env, "starter_credits_bridge", _armed_config(proxy_admin_url=None) + ) + + await _seed() + + assert FakeProxyClient.instances == [] + assert seeding_env.vault.row is None + + async def test_unresolved_policy_skips_seed(self, seeding_env): + async def no_policy(): + return None + + seeding_env.monkeypatch.setattr(service, "_resolve_mint_policy", no_policy) + + await _seed() + + assert _all_generate_calls() == [] + assert seeding_env.vault.row is None + + async def test_unverified_team_refuses_to_seed(self, seeding_env): + async def team_unverified(client, config): + return False + + seeding_env.monkeypatch.setattr( + service, "_team_ceiling_verified", team_unverified + ) + + await _seed() + + assert _all_generate_calls() == [] + assert seeding_env.vault.row is None + + async def test_existing_row_is_idempotent_no_second_mint(self, seeding_env): + await _seed() + assert len(_all_generate_calls()) == 1 + + await _seed() + + assert len(_all_generate_calls()) == 1 + + async def test_policy_refusal_skips_mint(self, seeding_env): + async def refused(organization_email, policy): + return False + + seeding_env.monkeypatch.setattr(service, "_mint_policy_allows", refused) + + await _seed() + + assert _all_generate_calls() == [] + assert seeding_env.vault.row is None + + async def test_missing_default_project_skips(self, seeding_env): + async def no_project(organization_id): + return None + + seeding_env.monkeypatch.setattr( + service.db_manager, + "get_default_project_by_organization_id", + no_project, + ) + + await _seed() + + assert _all_generate_calls() == [] + + +class TestOneKeyPerOrganization: + """We mint at most one key per organization; a failed seed is never retried + by a repair.""" + + async def test_alias_conflict_reads_as_already_seeded(self, seeding_env): + # A key already holds this org's alias (a duplicate signup race, or an + # earlier seed whose row the user deleted). + FakeProxyClient.records[ORGANIZATION_ID] = { + "key_alias": ORGANIZATION_ID, + "key": "sk-preexisting", + "max_budget": 10.0, + "spend": 4.0, + "metadata": {}, + } + + await _seed() + + # Exactly one mint attempt, no second key, no row, no failure alert. + assert len(_all_generate_calls()) == 1 + assert FakeProxyClient.records[ORGANIZATION_ID]["key"] == "sk-preexisting" + assert seeding_env.vault.row is None + assert seeding_env.alerts == [] + assert seeding_env.released == [ORGANIZATION_EMAIL] + + async def test_concurrent_seeds_produce_one_key_and_one_row(self, seeding_env): + await asyncio.gather(_seed(), _seed(), return_exceptions=True) + + assert len(FakeProxyClient.records) == 1 + assert seeding_env.vault.created_count == 1 + row = seeding_env.vault.row + assert row.data.provider.key == FakeProxyClient.records[ORGANIZATION_ID]["key"] + + async def test_failed_seed_is_not_retried_on_the_next_signup_hook( + self, seeding_env + ): + FakeProxyClient.fail_generate(service._MINT_ATTEMPTS) + + await _seed_safely() + + assert seeding_env.vault.row is None + # The org stays unseeded; nothing in the module converges it later. + assert not hasattr(service, "reconcile_starter_credits_bridge") + + +class TestMintRetry: + async def test_transient_failure_is_retried_within_the_bound(self, seeding_env): + FakeProxyClient.fail_generate(1, status_code=500) + + await _seed() + + assert len(_all_generate_calls()) == 2 + assert seeding_env.vault.row is not None + assert seeding_env.alerts == [] + + async def test_connection_failure_is_retried(self, seeding_env): + FakeProxyClient.generate_failures = [ + ProxyRequestError(status_code=None, detail="connection failed") + ] + + await _seed() + + assert len(_all_generate_calls()) == 2 + assert seeding_env.vault.row is not None + + async def test_retries_are_bounded(self, seeding_env): + FakeProxyClient.fail_generate(service._MINT_ATTEMPTS) + + await _seed_safely() + + assert len(_all_generate_calls()) == service._MINT_ATTEMPTS + assert seeding_env.vault.row is None + assert len(seeding_env.alerts) == 1 + assert "seed_failed" in seeding_env.alerts[0] + assert seeding_env.released == [ORGANIZATION_EMAIL] + + async def test_client_error_is_never_retried(self, seeding_env): + FakeProxyClient.fail_generate(1, status_code=400) + + await _seed_safely() + + assert len(_all_generate_calls()) == 1 + assert seeding_env.vault.row is None + assert len(seeding_env.alerts) == 1 + + async def test_alias_conflict_is_never_retried(self, seeding_env): + FakeProxyClient.generate_failures = [ + KeyAliasExistsError(status_code=400, detail="already exists: alias") + ] + + await _seed() + + assert len(_all_generate_calls()) == 1 + assert seeding_env.vault.row is None + + +class TestRowWriteFailure: + async def test_row_write_failure_blocks_the_orphaned_key(self, seeding_env): + seeding_env.vault.create_error = RuntimeError("vault write failed") + + await _seed_safely() + + assert seeding_env.vault.row is None + minted = FakeProxyClient.records[ORGANIZATION_ID]["key"] + assert _all_block_calls() == [minted] + assert len(seeding_env.alerts) == 1 + assert "seed_failed" in seeding_env.alerts[0] + + async def test_a_failing_block_still_degrades_quietly(self, seeding_env): + seeding_env.vault.create_error = RuntimeError("vault write failed") + + async def block_fails(self, *, key): + raise ProxyRequestError(status_code=500, detail="block failed") + + seeding_env.monkeypatch.setattr(FakeProxyClient, "block_key", block_fails) + + await _seed_safely() + + assert seeding_env.vault.row is None + assert len(seeding_env.alerts) == 1 + + +class TestTimeoutBound: + async def test_slow_seed_is_bounded_and_alerts(self, seeding_env): + async def slow_seed(**kwargs): + await asyncio.sleep(1.0) + + seeding_env.monkeypatch.setattr(service, "_SEED_TIMEOUT_SECONDS", 0.05) + seeding_env.monkeypatch.setattr( + service, "seed_starter_credits_bridge", slow_seed + ) + + await _seed_safely() + + assert len(seeding_env.alerts) == 1 + + +class TestTeamCeilingGate: + @pytest.fixture(autouse=True) + def armed(self, monkeypatch): + service._verified_teams.clear() + monkeypatch.setattr(env, "starter_credits_bridge", _armed_config()) + self.alerts = [] + self.clock = [0.0] + + monkeypatch.setattr( + service, "_send_alert_background", lambda text: self.alerts.append(text) + ) + monkeypatch.setattr(service, "_monotonic", lambda: self.clock[0]) + self.monkeypatch = monkeypatch + + def _client(self, payloads): + """payloads: list consumed one per lookup (last one repeats).""" + calls = [] + + class Client: + async def get_team_info(self, *, team_id): + calls.append(team_id) + payload = payloads[min(len(calls), len(payloads)) - 1] + if isinstance(payload, Exception): + raise payload + return payload + + client = Client() + client.calls = calls + return client + + async def test_sound_ceiling_verifies_and_caches_within_ttl(self): + client = self._client( + [ + { + "team_id": "team-starter", + "team_info": {"max_budget": 500.0, "budget_duration": None}, + } + ] + ) + config = env.starter_credits_bridge + + assert await service._team_ceiling_verified(client, config) is True + self.clock[0] += 100.0 + assert await service._team_ceiling_verified(client, config) is True + assert client.calls == ["team-starter"] + + async def test_reverifies_after_ttl_and_catches_removed_ceiling(self): + client = self._client( + [ + {"team_info": {"max_budget": 500.0, "budget_duration": None}}, + {"team_info": {"max_budget": None, "budget_duration": None}}, + ] + ) + config = env.starter_credits_bridge + + assert await service._team_ceiling_verified(client, config) is True + self.clock[0] += service._TEAM_VERIFY_TTL_SECONDS + 1 + assert await service._team_ceiling_verified(client, config) is False + assert len(client.calls) == 2 + assert len(self.alerts) == 1 + + async def test_top_level_budget_fields_also_accepted(self): + client = self._client([{"team_id": "team-starter", "max_budget": 500.0}]) + + assert ( + await service._team_ceiling_verified(client, env.starter_credits_bridge) + is True + ) + + async def test_unreachable_team_refuses(self): + client = self._client( + [ProxyRequestError(status_code=404, detail="team not found")] + ) + + assert ( + await service._team_ceiling_verified(client, env.starter_credits_bridge) + is False + ) + assert len(self.alerts) == 1 + + async def test_non_finite_budget_refuses(self): + client = self._client([{"team_info": {"max_budget": float("inf")}}]) + + assert ( + await service._team_ceiling_verified(client, env.starter_credits_bridge) + is False + ) + + async def test_resetting_budget_refuses(self): + client = self._client( + [{"team_info": {"max_budget": 500.0, "budget_duration": "30d"}}] + ) + + assert ( + await service._team_ceiling_verified(client, env.starter_credits_bridge) + is False + ) + assert len(self.alerts) == 1 + + +class TestMintPolicyResolution: + _PAYLOAD = { + "global_daily": 4, + "global_hourly": 3, + "work_domain_daily": 1, + "freemail_domains": ["freemail.test"], + "block_digit_locals": True, + "grant_usd": 10.0, + "key_max_parallel_requests": 2, + "key_rpm_limit": 30, + "key_tpm_limit": 200_000, + } + + @pytest.fixture(autouse=True) + def armed(self, monkeypatch): + monkeypatch.setattr(env, "starter_credits_bridge", _armed_config()) + # These cases are about a deployment that runs its own PostHog; the + # development-policy cases below say so explicitly. + monkeypatch.setattr(env, "posthog", PostHogConfig(api_key_configured=True)) + self.cache: dict = {} + self.alerts = [] + + async def fake_get_cache(*, namespace, key, retry=False, **kwargs): + return self.cache.get((namespace, tuple(sorted(key.items())))) + + async def fake_set_cache(*, namespace, key, value, **kwargs): + self.cache[(namespace, tuple(sorted(key.items())))] = value + + monkeypatch.setattr(service, "get_cache", fake_get_cache) + monkeypatch.setattr(service, "set_cache", fake_set_cache) + monkeypatch.setattr( + service, "_send_alert_background", lambda text: self.alerts.append(text) + ) + self.monkeypatch = monkeypatch + + def _posthog_with(self, payload): + return SimpleNamespace( + get_feature_flag_payload=lambda flag, distinct_id: payload + ) + + def _broken_posthog(self): + def boom(flag, distinct_id): + raise RuntimeError("posthog down") + + return SimpleNamespace(get_feature_flag_payload=boom) + + async def test_payload_resolves_and_caches(self): + self.monkeypatch.setattr( + service, "_load_posthog", lambda: self._posthog_with(dict(self._PAYLOAD)) + ) + + policy = await service._resolve_mint_policy() + + assert policy is not None + assert policy.global_daily == 4 + assert policy.grant_usd == 10.0 + assert policy.key_rpm_limit == 30 + assert len(self.cache) == 1 + + async def test_json_string_payload_parses(self): + import json + + self.monkeypatch.setattr( + service, + "_load_posthog", + lambda: self._posthog_with(json.dumps(self._PAYLOAD)), + ) + + policy = await service._resolve_mint_policy() + + assert policy is not None + assert policy.global_hourly == 3 + + async def test_the_payload_is_the_only_source_of_policy_values(self): + # No deployment can raise a cap or a grant on its own: the payload is it. + payload = dict(self._PAYLOAD, grant_usd=7.5) + self.monkeypatch.setattr( + service, "_load_posthog", lambda: self._posthog_with(payload) + ) + + policy = await service._resolve_mint_policy() + + assert policy is not None + assert policy.grant_usd == 7.5 + # The payload's list adds to the built-in defaults rather than replacing them. + assert policy.is_freemail("freemail.test") is True + assert policy.is_freemail("gmail.com") is True + + def _unconfigured_posthog(self): + """A deployment that supplied no PostHog key of its own. + + Driven through the config the bridge actually reads, not by stubbing the + loader: `PostHogConfig` falls back to a built-in project key, so a stubbed + loader would hide the very case this branch exists for. + """ + self.monkeypatch.setattr( + env, "posthog", PostHogConfig(api_key_configured=False) + ) + + def _must_not_be_consulted(): + raise AssertionError( + "PostHog must not be consulted when the deployment configured none" + ) + + self.monkeypatch.setattr(service, "_load_posthog", _must_not_be_consulted) + + def test_enabled_cannot_tell_a_stock_checkout_from_a_configured_one(self): + # The trap that made an earlier version of this fallback unreachable: the + # PostHog config falls back to a built-in project key, so `enabled` is true + # even where no operator ever configured PostHog. + unconfigured = PostHogConfig(api_key_configured=False) + + assert unconfigured.enabled is True + assert unconfigured.api_key_configured is False + + async def test_unconfigured_posthog_uses_the_development_policy(self): + # Local dev and QA stacks cannot publish a payload, so they run on the + # built-in one instead of being blocked by the fail-closed rule. + self._unconfigured_posthog() + self.monkeypatch.setattr(service, "_development_policy_announced", False) + + policy = await service._resolve_mint_policy() + + assert policy is not None + assert policy.grant_usd == 5.0 + assert policy.global_daily == 1000 + assert policy.key_tpm_limit == 200_000 + assert policy.block_digit_locals is False + # The built-in domain list still applies through the union. + assert policy.is_freemail("gmail.com") is True + # Nothing was cached: the development policy never becomes a stored payload. + assert self.cache == {} + + async def test_the_development_policy_announces_itself_once(self, caplog): + self._unconfigured_posthog() + self.monkeypatch.setattr(service, "_development_policy_announced", False) + + with caplog.at_level("WARNING"): + await service._resolve_mint_policy() + await service._resolve_mint_policy() + + notices = [ + record + for record in caplog.records + if "built-in development policy" in record.getMessage() + ] + assert len(notices) == 1 + + async def test_configured_posthog_without_a_payload_still_fails_closed(self): + # The fallback is only for a deployment that configured no PostHog. Here one + # is configured, and a missing payload is a real "no policy" signal. + self.monkeypatch.setattr( + service, "_load_posthog", lambda: self._posthog_with(None) + ) + + assert await service._resolve_mint_policy() is None + + async def test_a_configured_posthog_that_will_not_load_still_fails_closed(self): + # Configured, but the client cannot be built: on such a deployment that is a + # fault to fail closed on, never a reason to grant credits on dev values. + self.monkeypatch.setattr(service, "_load_posthog", lambda: None) + + assert await service._resolve_mint_policy() is None + + async def test_incomplete_payload_fails_closed(self): + partial = dict(self._PAYLOAD) + del partial["grant_usd"] + self.monkeypatch.setattr( + service, "_load_posthog", lambda: self._posthog_with(partial) + ) + + assert await service._resolve_mint_policy() is None + + async def test_unknown_payload_field_fails_closed(self): + payload = dict(self._PAYLOAD) + payload["surprise"] = 1 + self.monkeypatch.setattr( + service, "_load_posthog", lambda: self._posthog_with(payload) + ) + + assert await service._resolve_mint_policy() is None + + async def test_non_finite_grant_fails_closed(self): + payload = dict(self._PAYLOAD) + payload["grant_usd"] = float("inf") + self.monkeypatch.setattr( + service, "_load_posthog", lambda: self._posthog_with(payload) + ) + + assert await service._resolve_mint_policy() is None + + async def test_outage_falls_back_to_cached_payload(self): + self.monkeypatch.setattr( + service, "_load_posthog", lambda: self._posthog_with(dict(self._PAYLOAD)) + ) + assert await service._resolve_mint_policy() is not None + + self.monkeypatch.setattr(service, "_load_posthog", self._broken_posthog) + policy = await service._resolve_mint_policy() + + assert policy is not None + assert policy.global_daily == 4 + + async def test_malformed_live_payload_fails_closed_despite_cache(self): + # A cached valid payload may stand in for an OUTAGE only, never for a + # malformed live response (a bad rollout must fail closed, loudly). + self.monkeypatch.setattr( + service, "_load_posthog", lambda: self._posthog_with(dict(self._PAYLOAD)) + ) + assert await service._resolve_mint_policy() is not None + + self.monkeypatch.setattr( + service, "_load_posthog", lambda: self._posthog_with("{not json") + ) + + assert await service._resolve_mint_policy() is None + assert len(self.alerts) == 1 + + +class FakeCacheEngine: + def __init__(self, counts=None, error=None): + self.counts = counts or {} + self.error = error + self.expired = [] + + async def incr(self, key): + if self.error is not None: + raise self.error + self.counts[key] = self.counts.get(key, 0) + 1 + return self.counts[key] + + async def decr(self, key): + self.counts[key] = self.counts.get(key, 0) - 1 + return self.counts[key] + + async def expire(self, key, ttl): + self.expired.append((key, ttl)) + + +class TestMintPolicyAllows: + @pytest.fixture(autouse=True) + def armed(self, monkeypatch): + monkeypatch.setattr(env, "starter_credits_bridge", _armed_config()) + self.engine = FakeCacheEngine() + monkeypatch.setattr(service, "get_cache_engine", lambda: self.engine) + self.monkeypatch = monkeypatch + + async def test_freemail_skips_domain_counter_and_digit_rule(self): + assert ( + await service._mint_policy_allows("john99@freemail.test", _policy()) is True + ) + assert len(self.engine.counts) == 2 + assert not any("domain" in key for key in self.engine.counts) + + async def test_work_domain_digit_local_refused_before_counters(self): + assert await service._mint_policy_allows("john99@acme.test", _policy()) is False + assert self.engine.counts == {} + + async def test_work_domain_daily_cap_blocks(self): + assert await service._mint_policy_allows("alice@acme.test", _policy()) is True + assert await service._mint_policy_allows("bob@acme.test", _policy()) is False + + async def test_digit_rule_can_be_disabled_by_policy(self): + policy = _policy(block_digit_locals=False) + + assert await service._mint_policy_allows("john99@acme.test", policy) is True + + async def test_global_hourly_cap_blocks(self): + policy = _policy(global_hourly=2) + + assert await service._mint_policy_allows("a@freemail.test", policy) is True + assert await service._mint_policy_allows("b@freemail.test", policy) is True + assert await service._mint_policy_allows("c@freemail.test", policy) is False + + async def test_global_daily_cap_blocks(self): + policy = _policy(global_daily=1, global_hourly=10) + + assert await service._mint_policy_allows("a@freemail.test", policy) is True + assert await service._mint_policy_allows("b@freemail.test", policy) is False + + async def test_redis_error_fails_closed(self): + self.engine.error = ConnectionError("redis down") + + assert await service._mint_policy_allows("a@freemail.test", _policy()) is False + + async def test_release_hands_slots_back(self): + policy = _policy() + assert await service._mint_policy_allows("a@freemail.test", policy) is True + + await service._release_velocity_slots("a@freemail.test", policy) + + assert all(count == 0 for count in self.engine.counts.values()) + + +class TestManagedAndWriteOnlyRow: + """Management and value visibility are explicit, separate creation policies.""" + + async def test_the_row_uses_typed_management_and_explicit_write_only( + self, seeding_env + ): + await _seed() + + create_dto = seeding_env.vault.create_dto + assert create_dto is not None + assert create_dto.write_only is True + assert seeding_env.vault.management == SecretManagementDTO( + manager=SecretManager.STARTER_CREDITS_BRIDGE, + policy=SecretManagementPolicy.MANAGER_ONLY, + ) + + assert seeding_env.vault.row.management == seeding_env.vault.management + assert seeding_env.vault.row.write_only is True + + +class TestFreemailDefaults: + """Consumer mail providers are recognized without an operator listing them. + + A domain the policy does not classify as free mail is treated as a company domain: + every signup from it shares one per-domain daily cap, and a digit in the local part + refuses the mint. Applying that to proton.me or icloud.com refuses ordinary personal + signups, so the defaults ship in code and a configured list only adds to them. + """ + + @pytest.mark.parametrize( + "domain", + ["proton.me", "icloud.com", "outlook.de", "gmx.net", "yandex.ru", "qq.com"], + ) + def test_a_payload_without_a_list_still_classifies_consumer_mail(self, domain): + policy = MintPolicy( + global_daily=4, + global_hourly=3, + work_domain_daily=1, + block_digit_locals=True, + grant_usd=10.0, + key_max_parallel_requests=2, + key_rpm_limit=30, + key_tpm_limit=200_000, + ) + + assert policy.is_freemail(domain) is True + + def test_a_configured_list_is_added_to_the_defaults_not_swapped_for_them(self): + policy = _policy(freemail_domains=["freemail.test"]) + + assert policy.is_freemail("freemail.test") is True + assert policy.is_freemail("gmail.com") is True + assert policy.is_freemail("acme.test") is False + + def test_classification_ignores_case_and_padding(self): + policy = _policy(freemail_domains=[" Other.TEST "]) + + assert policy.is_freemail("PROTON.ME") is True + assert policy.is_freemail("other.test") is True + + +class TestRefusalLogging: + """A refusal names the rule and the DOMAIN. Never the address: the domain is what + makes a refusal diagnosable, and the local part identifies the person.""" + + @pytest.fixture(autouse=True) + def armed(self, monkeypatch): + monkeypatch.setattr(env, "starter_credits_bridge", _armed_config()) + self.engine = FakeCacheEngine() + monkeypatch.setattr(service, "get_cache_engine", lambda: self.engine) + self.records: list = [] + + class _RecordingLog: + def __init__(self, records): + self._records = records + + def warning(self, message, **kwargs): + self._records.append((message, kwargs)) + + def __getattr__(self, _name): + return lambda *args, **kwargs: None + + monkeypatch.setattr(service, "log", _RecordingLog(self.records)) + + async def test_a_digit_local_refusal_names_the_rule_and_the_domain(self): + assert await service._mint_policy_allows("john99@acme.test", _policy()) is False + + message, fields = self.records[-1] + assert fields["rule"] == "digit_local_part" + assert fields["domain"] == "acme.test" + assert "john99" not in repr((message, fields)) + + async def test_a_velocity_refusal_names_the_rule_and_the_domain(self): + policy = _policy(work_domain_daily=1) + + assert await service._mint_policy_allows("alice@acme.test", policy) is True + assert await service._mint_policy_allows("bob@acme.test", policy) is False + + message, fields = self.records[-1] + assert fields["rule"] == "work_domain_daily" + assert fields["domain"] == "acme.test" + assert "bob" not in repr((message, fields)) + + async def test_an_unverifiable_refusal_names_the_rule_and_the_domain(self): + self.engine.error = ConnectionError("redis down") + + assert await service._mint_policy_allows("carol@acme.test", _policy()) is False + + message, fields = self.records[-1] + assert fields["rule"] == "velocity_counters_unavailable" + assert fields["domain"] == "acme.test" + assert "carol" not in repr((message, fields)) + + +def test_a_bridge_deployment_without_a_runtime_key_fails_at_startup(monkeypatch): + # The bridge seeds write-only rows, so the deployment needs the dedicated runtime key. + from oss.src.utils import helpers + + monkeypatch.setattr(env.agenta, "services_internal_key", "replace-me") + monkeypatch.setattr(env, "starter_credits_bridge", _armed_config()) + + with pytest.raises(RuntimeError, match="AGENTA_SERVICES_INTERNAL_KEY"): + helpers.validate_platform_runtime_key() diff --git a/api/oss/src/utils/env.py b/api/oss/src/utils/env.py index d164427f184..b94b950192a 100644 --- a/api/oss/src/utils/env.py +++ b/api/oss/src/utils/env.py @@ -1356,6 +1356,11 @@ class PostHogConfig(BaseModel): os.getenv("POSTHOG_API_KEY") or "phc_hmVSxIjTW1REBHXgj2aw4HW9X6CXb6FzerBgP9XenC7" ) + # Whether THIS deployment supplied the key, as opposed to falling back to the + # built-in project key above. `enabled` cannot answer that — the fallback makes it + # true in every checkout — and a consumer that must tell "this operator runs + # PostHog" from "this is a stock local stack" needs the difference. + api_key_configured: bool = bool(os.getenv("POSTHOG_API_KEY")) model_config = ConfigDict(extra="ignore") @@ -1519,6 +1524,85 @@ def enabled(self) -> bool: return bool(self.api_key and self.from_email) +# --------------------------------------------------------------------------- +# starter_credits_bridge — starter-credit seeding at signup (EE, temporary bridge). +# --------------------------------------------------------------------------- + + +class StarterCreditsBridgeConfig(BaseModel): + """Starter-credits bridge (EE): mint a budget-capped proxy key at signup and + seed it into the new organization's vault as a ready-to-use connection. + + Inert unless `enabled` is true AND `proxy_admin_url` + `proxy_public_url` + + `master_key` + `team_id` are present (the team's budget ceiling is the + program's total-exposure bound, so seeding refuses to run without it). The + redeploy-free runtime switch is the PostHog policy payload: clearing or + emptying it leaves the policy unresolved, and seeding fails closed. + """ + + enabled: bool = _parse_bool_env("AGENTA_STARTER_CREDITS_BRIDGE_ENABLED", False) + + # Two base URLs for the same proxy, because it is reachable two ways. The + # public one is what the SEEDED CONNECTION stores, so a sandboxed run can + # reach the proxy's inference paths from outside; only those paths are + # publicly routed. The admin routes (/key/generate, /key/block, /team/info) + # are not, so the master-keyed admin client dials the proxy's internal + # address instead — and the master key never leaves the private network. + proxy_public_url: str | None = ( + os.getenv("AGENTA_STARTER_CREDITS_BRIDGE_PROXY_PUBLIC_URL") or None + ) + proxy_admin_url: str | None = ( + os.getenv("AGENTA_STARTER_CREDITS_BRIDGE_PROXY_ADMIN_URL") or None + ) + master_key: str | None = ( + os.getenv("AGENTA_STARTER_CREDITS_BRIDGE_MASTER_KEY") or None + ) + # Every minted key joins this team; its max_budget is the program ceiling. + team_id: str | None = os.getenv("AGENTA_STARTER_CREDITS_BRIDGE_TEAM_ID") or None + + model_id: str = ( + os.getenv("AGENTA_STARTER_CREDITS_BRIDGE_MODEL_ID") + or "vertex_ai/gemini-3.6-flash" + ) + + # The mint policy (velocity caps, domain classification, eligibility rules, + # grant size, per-key limits) ships via the PostHog policy flag's payload: + # the payload is the only source, so no real value lives in source and no + # money value can be changed one field at a time on a single deployment. + # Without a resolvable payload the policy is unresolved and seeding fails + # closed. Only the flag's NAME is configurable here. + # + # One exception, for developer convenience: a deployment with NO PostHog + # configured at all (local dev, a QA stack) runs on the built-in development + # policy in `starter_credits_bridge/types.py` instead of being blocked, since + # it has no way to publish a payload. Cloud always has PostHog, so it always + # takes the payload path and still fails closed on a missing or bad one. + policy_flag: str = ( + os.getenv("AGENTA_STARTER_CREDITS_BRIDGE_POLICY_FLAG") + or "starter-credits-bridge-policy" + ) + # Optional operator webhook ({"text": ...} POST) for refusals and failures. + alert_webhook: str | None = ( + os.getenv("AGENTA_STARTER_CREDITS_BRIDGE_ALERT_WEBHOOK") or None + ) + + model_config = ConfigDict(extra="ignore") + + @property + def armed(self) -> bool: + """True when the deployment opted in and holds both proxy addresses and + the proxy credentials, plus the program team whose ceiling bounds total + exposure. Without the admin address every mint would dial a route that is + not publicly served, so a missing one keeps the bridge inert.""" + return bool( + self.enabled + and self.proxy_public_url + and self.proxy_admin_url + and self.master_key + and self.team_id + ) + + # --------------------------------------------------------------------------- # stripe # --------------------------------------------------------------------------- @@ -1672,6 +1756,7 @@ class EnvironSettings(BaseModel): sessions: SessionsRedisConfig = SessionsRedisConfig() smtp: SmtpConfig = SmtpConfig() sendgrid: SendgridConfig = SendgridConfig() + starter_credits_bridge: StarterCreditsBridgeConfig = StarterCreditsBridgeConfig() store: StoreConfig = StoreConfig() stripe: StripeConfig = StripeConfig() supertokens: SuperTokensConfig = SuperTokensConfig() diff --git a/docs/design/starter-credits-seeding/README.md b/docs/design/starter-credits-seeding/README.md new file mode 100644 index 00000000000..30e2c58308c --- /dev/null +++ b/docs/design/starter-credits-seeding/README.md @@ -0,0 +1,60 @@ +# Starter credits seeding + +This document records the production contract for the starter-credits connection created for a +new organization's default project. + +## Seeding flow + +The EE signup hook calls the starter-credits bridge. The bridge: + +1. Confirms the bridge configuration and `AGENTA_SERVICES_INTERNAL_KEY` are configured. +2. Resolves the mint policy and verifies the proxy team budget ceiling. +3. Loads the organization's default project and checks the `starter-credits` Vault slug. +4. Applies the signup and velocity policy. +5. Mints one budget-capped proxy key for the organization. +6. Creates one managed Vault connection through `VaultService.create_managed_secret`. + +The proxy alias and the project-plus-slug unique constraint make duplicate signup attempts +idempotent. If the Vault write fails after minting, the bridge blocks the orphaned proxy key. + +## Vault connection contract + +The bridge creates a custom-provider connection with these independent properties: + +| Property | Value | Purpose | +| --- | --- | --- | +| Slug | `starter-credits` | Stable Vault identity | +| Name | `Agenta` | User-facing connection name and model-key namespace | +| Description | `Provided and managed by Agenta.` | User-facing explanation | +| Provider URL | `AGENTA_STARTER_CREDITS_BRIDGE_PROXY_PUBLIC_URL` | Endpoint used by model runs | +| `write_only` | `true` | Prevents ordinary API and UI callers from reading the virtual key | +| Manager | `SecretManager.STARTER_CREDITS_BRIDGE` | Identifies the trusted lifecycle owner internally | +| Management policy | `SecretManagementPolicy.MANAGER_ONLY` | Prevents general update and delete operations | + +Management and value visibility are separate policies. The bridge selects both explicitly. The +general Vault layer does not infer `write_only` from management. + +The public Vault response exposes the management policy, not the internal manager identity. + +## Names with different roles + +The bridge keeps these values separate even when two serialized values happen to match: + +- `PROXY_ORIGIN` is proxy audit metadata attached to the minted key. +- `AGENTA_STARTER_CREDITS_BRIDGE_PROXY_ADMIN_URL` routes administrative mint and block calls. +- `AGENTA_STARTER_CREDITS_BRIDGE_PROXY_PUBLIC_URL` is stored as the connection endpoint. +- `STARTER_CREDITS_DESCRIPTION` is human-facing copy stored in the Vault header. + +Do not reuse the proxy origin or either URL as the user-facing description. + +## Cache ownership + +`VaultService.create_managed_secret` owns Vault list-cache invalidation. The bridge does not call +the cache helper directly. This keeps HTTP and in-process Vault writers on the same invalidation +path. + +## Runtime key requirement + +The bridge refuses to mint when `AGENTA_SERVICES_INTERNAL_KEY` is absent, blank, or the +`replace-me` placeholder. It does not fall back to `AGENTA_AUTH_KEY`. Without the dedicated key, +the runtime cannot obtain the grant required to resolve the write-only connection.