From 93b835222793df47e8faf0899cc6fe5ff13de005 Mon Sep 17 00:00:00 2001 From: Leonardo Gonzalez Date: Tue, 21 Apr 2026 22:00:36 -0500 Subject: [PATCH 1/7] COE-370: Implement sessionless LiteLLM key lifecycle commands - Add ProxyKeyService for sessionless key creation, listing, info, revocation - Add benchmark key create/list/info/revoke CLI subcommands - Wire key commands into Typer CLI app - Persist only non-secret metadata (alias, owner, team, customer, models, status) - Display key secret once at creation; never persist or re-display - Render OpenAI-compatible and harness-profile env snippets (shell + dotenv) - Unify credential_service.py with full LiteLLM-backed implementation from services_abc - Add respx/httpx unit tests for LiteLLM success and failure paths - Add CLI command tests with synthetic metadata - Add budget_currency field to UsagePolicyProfile - All 584 unit tests passing; lint clean --- src/benchmark_core/config.py | 4 + src/benchmark_core/services/__init__.py | 19 +- .../services/credential_service.py | 378 ++++++++++---- .../services/proxy_key_service.py | 470 ++++++++++++++++++ src/cli/commands/key.py | 291 +++++++++++ src/cli/main.py | 3 +- tests/unit/test_cli_key_commands.py | 340 +++++++++++++ tests/unit/test_proxy_key_service.py | 438 ++++++++++++++++ tests/unit/test_services.py | 65 ++- 9 files changed, 1888 insertions(+), 120 deletions(-) create mode 100644 src/benchmark_core/services/proxy_key_service.py create mode 100644 src/cli/commands/key.py create mode 100644 tests/unit/test_cli_key_commands.py create mode 100644 tests/unit/test_proxy_key_service.py diff --git a/src/benchmark_core/config.py b/src/benchmark_core/config.py index 6bfe781..e8c08b7 100644 --- a/src/benchmark_core/config.py +++ b/src/benchmark_core/config.py @@ -289,6 +289,10 @@ class UsagePolicyProfile(BaseModel): default=None, description="Budget limit in currency units", ) + budget_currency: str = Field( + default="USD", + description="Budget currency", + ) ttl_seconds: int | None = Field( default=None, description="Key TTL in seconds (time-to-live before expiration)", diff --git a/src/benchmark_core/services/__init__.py b/src/benchmark_core/services/__init__.py index 79fb5e8..98402a8 100644 --- a/src/benchmark_core/services/__init__.py +++ b/src/benchmark_core/services/__init__.py @@ -6,9 +6,19 @@ # Comprehensive SQL-based services (COE-305 implementation) from benchmark_core.services.benchmark_metadata_service import BenchmarkMetadataService + +# Unified credential service (COE-370) +from benchmark_core.services.credential_service import CredentialService from benchmark_core.services.experiment_service import ExperimentService from benchmark_core.services.harness_profile_service import HarnessProfileService from benchmark_core.services.provider_service import ProviderService + +# Sessionless proxy key services +from benchmark_core.services.proxy_key_service import ( + LiteLLMAPIError, + ProxyKeyService, + ProxyKeyServiceError, +) from benchmark_core.services.rendering import ( EnvRenderingService, EnvSnippet, @@ -25,11 +35,6 @@ from benchmark_core.services.task_card_service import TaskCardService from benchmark_core.services.variant_service import VariantService -# ABC service exports - Note: old services_abc SessionService removed, use services.session_service -from benchmark_core.services_abc import ( - CredentialService, -) - __all__ = [ # Core SQL-based services "SessionService", @@ -49,6 +54,10 @@ "RenderingError", "ProfileValidationError", "render_env_for_session", + # Sessionless proxy key services + "ProxyKeyService", + "ProxyKeyServiceError", + "LiteLLMAPIError", # ABC services "CredentialService", ] diff --git a/src/benchmark_core/services/credential_service.py b/src/benchmark_core/services/credential_service.py index fcea5fc..2573c8f 100644 --- a/src/benchmark_core/services/credential_service.py +++ b/src/benchmark_core/services/credential_service.py @@ -1,18 +1,93 @@ -"""Service for rendering and managing session-scoped proxy credentials.""" +"""Service for rendering and managing session-scoped proxy credentials. +Unified implementation moved from services_abc.py. Provides both the +full LiteLLM-backed credential lifecycle and backward-compatible +convenience methods. +""" + +from datetime import UTC, datetime, timedelta +from typing import Any from uuid import UUID +import httpx +from pydantic import SecretStr + +from benchmark_core.models import ProxyCredential + class CredentialService: - """Service for rendering and managing session-scoped proxy credentials. + """Service for issuing and managing session-scoped proxy credentials. - Credentials are short-lived tokens issued for benchmark sessions - that enable correlation between benchmark traffic and session metadata. + Implements the key aliasing convention and metadata tagging strategy + for correlating LiteLLM traffic back to benchmark sessions. """ - def __init__(self) -> None: - """Initialize the credential service.""" - pass + def __init__( + self, + litellm_base_url: str = "http://localhost:4000", + master_key: str | None = None, + enforce_https: bool = True, + ) -> None: + """Initialize the credential service. + + Args: + litellm_base_url: Base URL for LiteLLM proxy API + master_key: LiteLLM master key for API authentication + enforce_https: Whether to enforce HTTPS in production (default True) + + Raises: + ValueError: If litellm_base_url doesn't use HTTPS in production + """ + self.litellm_base_url = litellm_base_url.rstrip("/") + self.master_key = master_key + + # Validate HTTPS in production environments + if enforce_https and not self.litellm_base_url.startswith( + ("https://", "http://localhost", "http://127.0.0.1") + ): + raise ValueError( + f"LiteLLM URL must use HTTPS in production environments. Got: {litellm_base_url}" + ) + + def _generate_key_alias( + self, + session_id: UUID, + experiment_id: str, + variant_id: str, + ) -> str: + """Generate a stable key alias for correlation. + + Format: session-{session_id[:8]}-{experiment_id[:8]}-{variant_id[:8]} + + This alias is: + - Human-readable for operators + - Unique per session + - Joinable back to session via metadata + """ + session_short = str(session_id)[:8] + exp_short = experiment_id[:8] if len(experiment_id) >= 8 else experiment_id + var_short = variant_id[:8] if len(variant_id) >= 8 else variant_id + return f"session-{session_short}-{exp_short}-{var_short}" + + def _build_metadata_tags( + self, + session_id: UUID, + experiment_id: str, + variant_id: str, + harness_profile: str, + ) -> dict[str, str]: + """Build metadata tags for LiteLLM credential. + + These tags enable joining LiteLLM request logs back to + the benchmark session and its dimensions. + """ + return { + "benchmark_session_id": str(session_id), + "benchmark_experiment_id": experiment_id, + "benchmark_variant_id": variant_id, + "benchmark_harness_profile": harness_profile, + "benchmark_source": "opensymphony", + } async def issue_credential( self, @@ -20,141 +95,262 @@ async def issue_credential( experiment_id: str, variant_id: str, harness_profile: str, - ) -> str: - """Generate a session-scoped proxy credential. + ttl_hours: int = 24, + ) -> ProxyCredential: + """Issue a session-scoped proxy credential via LiteLLM API. + + Creates a LiteLLM virtual key with: + - Unique key alias for correlation + - Metadata tags linking to session + - Configurable TTL (default 24 hours) + + Args: + session_id: Benchmark session UUID + experiment_id: Experiment identifier + variant_id: Variant identifier + harness_profile: Harness profile name + ttl_hours: Credential time-to-live in hours + + Returns: + ProxyCredential with metadata and API key + + Raises: + RuntimeError: If LiteLLM API call fails or master_key not configured + """ + if self.master_key is None: + raise RuntimeError( + "Cannot issue credential: LiteLLM master_key not configured. " + "Initialize CredentialService with master_key parameter." + ) + + # Generate stable alias for correlation + key_alias = self._generate_key_alias(session_id, experiment_id, variant_id) + + # Build metadata tags for LiteLLM correlation + metadata_tags = self._build_metadata_tags( + session_id, experiment_id, variant_id, harness_profile + ) + + # Calculate expiration + expires_at = datetime.now(UTC) + timedelta(hours=ttl_hours) + + # Call LiteLLM API to create the key + async with httpx.AsyncClient(timeout=30.0) as client: + try: + response = await client.post( + f"{self.litellm_base_url}/key/generate", + headers={"Authorization": f"Bearer {self.master_key}"}, + json={ + "key_alias": key_alias, + "metadata": metadata_tags, + "expires": expires_at.isoformat(), + }, + ) + response.raise_for_status() + data: dict[str, Any] = response.json() + except httpx.HTTPStatusError as e: + raise RuntimeError( + f"LiteLLM API error creating credential: {e.response.status_code} - " + f"{e.response.text}" + ) from e + except httpx.RequestError as e: + raise RuntimeError( + f"Failed to connect to LiteLLM at {self.litellm_base_url}: {e}" + ) from e + + # Extract key and key_id from LiteLLM response + api_key = data.get("key") + litellm_key_id = data.get("key_id") + + if not api_key: + raise RuntimeError(f"LiteLLM API response missing 'key' field: {data}") - Currently returns a placeholder credential. The actual implementation - will integrate with LiteLLM API for short-lived credential issuance. + # Create credential domain model + credential = ProxyCredential( + session_id=session_id, + key_alias=key_alias, + api_key=SecretStr(api_key), + experiment_id=experiment_id, + variant_id=variant_id, + harness_profile=harness_profile, + metadata_tags=metadata_tags, + litellm_key_id=litellm_key_id, + expires_at=expires_at, + is_active=True, + ) - The credential encodes session metadata for correlation: - - session_id for direct session lookup - - experiment_id prefix for grouping - - variant_id prefix for variant tracking - - harness_profile for profile verification + return credential + + async def revoke_credential(self, credential: ProxyCredential) -> ProxyCredential: + """Revoke an active credential via LiteLLM API. + + Marks the credential as revoked and records revocation time. + Also calls LiteLLM API to delete the key if litellm_key_id is present. Args: - session_id: The benchmark session ID. - experiment_id: The experiment identifier. - variant_id: The variant identifier. - harness_profile: The harness profile name. + credential: Credential to revoke Returns: - A session-scoped proxy credential string. + Updated credential with is_active=False + + Note: + LiteLLM API errors are intentionally silenced since revocation is + best-effort - the key may already be expired or the proxy may be + temporarily unavailable. """ - # Placeholder: actual implementation will integrate with LiteLLM API - # The credential format encodes metadata for correlation - return f"sk-benchmark-{session_id}-{experiment_id[:8]}" + # Call LiteLLM API to delete the key if we have the key_id + if credential.litellm_key_id and self.master_key: + async with httpx.AsyncClient(timeout=30.0) as client: + try: + response = await client.post( + f"{self.litellm_base_url}/key/delete", + headers={"Authorization": f"Bearer {self.master_key}"}, + json={"key_ids": [credential.litellm_key_id]}, + ) + response.raise_for_status() + except httpx.HTTPStatusError: + # Log but don't fail - the key may already be expired/deleted + pass # noqa: S110 - intentional silence on revoke errors + except httpx.RequestError: + # Log but don't fail - revocation is best-effort + pass # noqa: S110 - intentional silence on connection errors + + updated = credential.model_copy( + update={ + "is_active": False, + "revoked_at": datetime.now(UTC), + # Clear the secret from memory + "api_key": SecretStr("REDACTED"), + } + ) + return updated + + # ------------------------------------------------------------------ + # Environment rendering + # ------------------------------------------------------------------ def render_env_snippet( self, - credential: str, + credential: ProxyCredential, proxy_base_url: str, model: str, - harness_profile: str, ) -> dict[str, str]: """Render environment variable snippet for a harness. - This produces a standard OpenAI-compatible environment configuration - that works with most harnesses. - Args: - credential: The proxy credential. - proxy_base_url: The proxy base URL. - model: The model identifier. - harness_profile: The harness profile name (for validation). + credential: Proxy credential with API key + proxy_base_url: LiteLLM proxy base URL + model: Model identifier/alias Returns: - Dictionary of environment variables for the harness. + Dictionary of environment variables """ return { "OPENAI_API_BASE": proxy_base_url, - "OPENAI_API_KEY": credential, + "OPENAI_API_KEY": credential.api_key.get_secret_value(), "OPENAI_MODEL": model, + "LITELLM_SESSION_ALIAS": credential.key_alias, } - async def revoke_credential(self, credential: str) -> bool: - """Revoke a session-scoped proxy credential. - - Placeholder for future LiteLLM API integration. + def render_env_shell(self, env_vars: dict[str, str]) -> str: + """Render environment variables as shell export commands. Args: - credential: The credential to revoke. + env_vars: Dictionary of environment variables. Returns: - True if revoked successfully. + Shell commands string. """ - # Placeholder: actual implementation will integrate with LiteLLM API - return True - - async def validate_credential(self, credential: str) -> dict | None: - """Validate a credential and extract its metadata. + lines = [] + for key, value in env_vars.items(): + escaped = value.replace("'", "'\\''") + lines.append(f"export {key}='{escaped}'") + return "\n".join(lines) - Placeholder for future LiteLLM API integration. + def render_env_dotenv(self, env_vars: dict[str, str]) -> str: + """Render environment variables as dotenv file content. Args: - credential: The credential to validate. + env_vars: Dictionary of environment variables. Returns: - Dictionary with extracted metadata, or None if invalid. + Dotenv file content string. """ - # Placeholder: actual implementation will parse and validate with LiteLLM API - # This is a mock implementation for the placeholder format - if credential.startswith("sk-benchmark-"): - parts = credential.split("-") - if len(parts) >= 3: - return { - "session_id": parts[2] if len(parts) > 2 else None, - "experiment_prefix": parts[3] if len(parts) > 3 else None, - "valid": True, - } - return None + lines = [] + for key, value in env_vars.items(): + if " " in value or "#" in value or "'" in value or '"' in value: + escaped = value.replace("\\", "\\\\").replace('"', '\\"') + lines.append(f'{key}="{escaped}"') + else: + lines.append(f"{key}={value}") + return "\n".join(lines) + + # ------------------------------------------------------------------ + # Backward-compatible convenience wrappers (deprecated stubs) + # ------------------------------------------------------------------ - def render_shell_snippet( + def _legacy_render_env_snippet( + self, + credential: str, + proxy_base_url: str, + model: str, + harness_profile: str, + ) -> dict[str, str]: + """Legacy wrapper accepting a plain credential string.""" + return { + "OPENAI_API_BASE": proxy_base_url, + "OPENAI_API_KEY": credential, + "OPENAI_MODEL": model, + } + + def _legacy_render_shell_snippet( self, credential: str, proxy_base_url: str, model: str, additional_vars: dict[str, str] | None = None, ) -> str: - """Render a shell snippet for setting environment variables. - - Args: - credential: The proxy credential. - proxy_base_url: The proxy base URL. - model: The model identifier. - additional_vars: Optional additional environment variables. - - Returns: - Shell command string for setting environment. - """ - env_vars = self.render_env_snippet(credential, proxy_base_url, model, "") + """Legacy wrapper for shell rendering with a plain credential string.""" + env_vars = self._legacy_render_env_snippet(credential, proxy_base_url, model, "") if additional_vars: env_vars.update(additional_vars) + return self.render_env_shell(env_vars) - lines = [f'export {key}="{value}"' for key, value in env_vars.items()] - return "\n".join(lines) - - def render_dotenv_snippet( + def _legacy_render_dotenv_snippet( self, credential: str, proxy_base_url: str, model: str, additional_vars: dict[str, str] | None = None, ) -> str: - """Render a dotenv file snippet for environment variables. - - Args: - credential: The proxy credential. - proxy_base_url: The proxy base URL. - model: The model identifier. - additional_vars: Optional additional environment variables. - - Returns: - Dotenv file content string. - """ - env_vars = self.render_env_snippet(credential, proxy_base_url, model, "") + """Legacy wrapper for dotenv rendering with a plain credential string.""" + env_vars = self._legacy_render_env_snippet(credential, proxy_base_url, model, "") if additional_vars: env_vars.update(additional_vars) + return self.render_env_dotenv(env_vars) - lines = [f'{key}="{value}"' for key, value in env_vars.items()] - return "\n".join(lines) + def _legacy_validate_credential(self, credential: str) -> dict | None: + """Legacy placeholder validation for stub credentials.""" + if credential.startswith("sk-benchmark-"): + parts = credential.split("-") + if len(parts) >= 3: + return { + "session_id": parts[2] if len(parts) > 2 else None, + "experiment_prefix": parts[3] if len(parts) > 3 else None, + "valid": True, + } + return None + + async def _legacy_issue_credential( + self, + session_id: UUID, + experiment_id: str, + variant_id: str, + harness_profile: str, + ) -> str: + """Legacy placeholder that returns a stub credential string.""" + return f"sk-benchmark-{session_id}-{experiment_id[:8]}" + + async def _legacy_revoke_credential(self, credential: str) -> bool: + """Legacy placeholder revocation.""" + return True diff --git a/src/benchmark_core/services/proxy_key_service.py b/src/benchmark_core/services/proxy_key_service.py new file mode 100644 index 0000000..d951939 --- /dev/null +++ b/src/benchmark_core/services/proxy_key_service.py @@ -0,0 +1,470 @@ +"""Service for sessionless proxy key operations via LiteLLM API. + +Provides creation, listing, revocation, and environment rendering +for LiteLLM virtual keys without requiring a benchmark session. +Only non-secret key metadata is persisted locally. +""" + +from datetime import UTC, datetime, timedelta +from typing import Any +from uuid import UUID, uuid4 + +import httpx +from pydantic import SecretStr + +from benchmark_core.config import UsagePolicyProfile +from benchmark_core.db.models import ProxyKey as ProxyKeyORM +from benchmark_core.models import ProxyKey, ProxyKeyStatus +from benchmark_core.repositories.proxy_key_repository import SQLProxyKeyRepository + + +class ProxyKeyServiceError(Exception): + """Raised when proxy key service operations fail.""" + + pass + + +class LiteLLMAPIError(ProxyKeyServiceError): + """Raised when LiteLLM API call fails.""" + + pass + + +class ProxyKeyService: + """Service for managing sessionless proxy keys. + + Coordinates LiteLLM virtual key lifecycle with local non-secret + metadata persistence. Key secrets are displayed once at creation + and never stored in the benchmark database. + """ + + def __init__( + self, + repository: SQLProxyKeyRepository, + litellm_base_url: str = "http://localhost:4000", + master_key: str | None = None, + enforce_https: bool = True, + ) -> None: + """Initialize the proxy key service. + + Args: + repository: Repository for proxy key metadata persistence. + litellm_base_url: Base URL for LiteLLM proxy API. + master_key: LiteLLM master key for API authentication. + enforce_https: Whether to enforce HTTPS in production. + + Raises: + ValueError: If litellm_base_url doesn't use HTTPS in production. + """ + self._repository = repository + self.litellm_base_url = litellm_base_url.rstrip("/") + self.master_key = master_key + + if enforce_https and not self.litellm_base_url.startswith( + ("https://", "http://localhost", "http://127.0.0.1") + ): + raise ValueError( + f"LiteLLM URL must use HTTPS in production environments. Got: {litellm_base_url}" + ) + + def _http_headers(self) -> dict[str, str]: + """Build HTTP headers for LiteLLM API calls.""" + if self.master_key is None: + raise ProxyKeyServiceError( + "LiteLLM master_key not configured. " + "Set LITELLM_MASTER_KEY or initialize ProxyKeyService with master_key." + ) + return {"Authorization": f"Bearer {self.master_key}"} + + def _build_key_alias(self, prefix: str | None = None) -> str: + """Generate a stable unique key alias. + + Format: {prefix}-{uuid8} or usage-{uuid8} if no prefix. + """ + short_uuid = str(uuid4())[:8] + if prefix: + return f"{prefix}-{short_uuid}" + return f"usage-{short_uuid}" + + async def create_key( + self, + key_alias: str | None = None, + owner: str | None = None, + team: str | None = None, + customer: str | None = None, + purpose: str | None = None, + allowed_models: list[str] | None = None, + budget_duration: str | None = None, + budget_amount: float | None = None, + ttl_hours: int = 168, + metadata: dict[str, str] | None = None, + usage_policy: UsagePolicyProfile | None = None, + ) -> tuple[ProxyKey, SecretStr]: + """Create a new sessionless proxy key via LiteLLM API. + + The key secret is returned once and never persisted to the database. + Only non-secret metadata (alias, owner, team, etc.) is stored locally. + + Args: + key_alias: Human-readable alias. Auto-generated if None. + owner: Key owner label. + team: Team metadata. + customer: Customer metadata. + purpose: Key purpose/description. + allowed_models: Optional list of allowed model aliases. + budget_duration: Budget interval (e.g., "1d", "30d"). + budget_amount: Budget limit in currency units. + ttl_hours: Key time-to-live in hours (default 7 days). + metadata: Additional metadata tags. + usage_policy: Optional usage policy profile to apply defaults from. + + Returns: + Tuple of (ProxyKey metadata model, SecretStr with the actual key). + + Raises: + LiteLLMAPIError: If LiteLLM API call fails. + ProxyKeyServiceError: If master_key not configured. + """ + if self.master_key is None: + raise ProxyKeyServiceError( + "Cannot create key: LiteLLM master_key not configured." + ) + + # Apply usage policy defaults if provided + effective_alias = key_alias or self._build_key_alias( + usage_policy.name if usage_policy else None + ) + effective_owner = owner or (usage_policy.owner if usage_policy else None) + effective_team = team or (usage_policy.team if usage_policy else None) + effective_customer = customer or (usage_policy.customer if usage_policy else None) + effective_models = allowed_models or ( + list(usage_policy.allowed_models) if usage_policy else [] + ) + effective_budget_duration = budget_duration or ( + usage_policy.budget_duration if usage_policy else None + ) + effective_budget_amount = budget_amount or ( + usage_policy.budget_amount if usage_policy else None + ) + effective_ttl = ( + usage_policy.ttl_seconds // 3600 if usage_policy and usage_policy.ttl_seconds else ttl_hours + ) + + # Merge metadata + effective_metadata: dict[str, Any] = {} + if usage_policy and usage_policy.metadata: + effective_metadata.update(usage_policy.metadata) + if metadata: + effective_metadata.update(metadata) + + # Calculate expiration + expires_at = datetime.now(UTC) + timedelta(hours=effective_ttl) + + # Build LiteLLM API payload + litellm_payload: dict[str, Any] = { + "key_alias": effective_alias, + "metadata": effective_metadata, + "expires": expires_at.isoformat(), + } + if effective_models: + litellm_payload["models"] = effective_models + if effective_budget_duration: + litellm_payload["budget_duration"] = effective_budget_duration + if effective_budget_amount is not None: + litellm_payload["budget_amount"] = effective_budget_amount + + # Call LiteLLM API to create the key + async with httpx.AsyncClient(timeout=30.0) as client: + try: + response = await client.post( + f"{self.litellm_base_url}/key/generate", + headers=self._http_headers(), + json=litellm_payload, + ) + response.raise_for_status() + data: dict[str, Any] = response.json() + except httpx.HTTPStatusError as e: + raise LiteLLMAPIError( + f"LiteLLM API error creating key: {e.response.status_code} - {e.response.text}" + ) from e + except httpx.RequestError as e: + raise LiteLLMAPIError( + f"Failed to connect to LiteLLM at {self.litellm_base_url}: {e}" + ) from e + + api_key = data.get("key") + litellm_key_id = data.get("key_id") + + if not api_key: + raise LiteLLMAPIError(f"LiteLLM API response missing 'key' field: {data}") + + # Build budget_currency from policy default + budget_currency = "USD" + if usage_policy and usage_policy.budget_currency: + budget_currency = usage_policy.budget_currency + + # Create local metadata record (NO secret stored) + proxy_key = ProxyKey( + proxy_key_id=uuid4(), + key_alias=effective_alias, + litellm_key_id=litellm_key_id, + owner=effective_owner, + team=effective_team, + customer=effective_customer, + purpose=purpose, + allowed_models=list(effective_models) if effective_models else [], + budget_duration=effective_budget_duration, + budget_amount=effective_budget_amount, + budget_currency=budget_currency, + status=ProxyKeyStatus.ACTIVE, + key_metadata=effective_metadata, + created_at=datetime.now(UTC), + expires_at=expires_at, + ) + + # Persist metadata to database + orm = ProxyKeyORM( + id=proxy_key.proxy_key_id, + key_alias=proxy_key.key_alias, + litellm_key_id=proxy_key.litellm_key_id, + owner=proxy_key.owner, + team=proxy_key.team, + customer=proxy_key.customer, + purpose=proxy_key.purpose, + allowed_models=proxy_key.allowed_models, + budget_duration=proxy_key.budget_duration, + budget_amount=proxy_key.budget_amount, + budget_currency=proxy_key.budget_currency, + status=proxy_key.status.value, + key_metadata=proxy_key.key_metadata, + created_at=proxy_key.created_at, + expires_at=proxy_key.expires_at, + ) + await self._repository.create(orm) + + return proxy_key, SecretStr(api_key) + + async def list_keys( + self, + owner: str | None = None, + team: str | None = None, + customer: str | None = None, + status: ProxyKeyStatus | None = None, + limit: int = 100, + offset: int = 0, + ) -> list[ProxyKey]: + """List proxy keys with optional filtering. + + Args: + owner: Filter by owner label. + team: Filter by team label. + customer: Filter by customer label. + status: Filter by status. + limit: Maximum results. + offset: Results to skip. + + Returns: + List of ProxyKey metadata models (no secrets). + """ + if owner: + orm_list = await self._repository.list_by_owner(owner, limit, offset) + elif team: + orm_list = await self._repository.list_by_team(team, limit, offset) + elif customer: + orm_list = await self._repository.list_by_customer(customer, limit, offset) + else: + orm_list = await self._repository.list_all(limit, offset) + + if status: + orm_list = [k for k in orm_list if k.status == status.value] + + return [self._orm_to_model(k) for k in orm_list] + + async def get_key_info(self, proxy_key_id: UUID) -> ProxyKey | None: + """Get detailed info for a proxy key. + + Args: + proxy_key_id: The UUID of the proxy key. + + Returns: + ProxyKey model if found, None otherwise. + """ + orm = await self._repository.get_by_id(proxy_key_id) + if orm is None: + return None + return self._orm_to_model(orm) + + async def get_key_by_alias(self, key_alias: str) -> ProxyKey | None: + """Get proxy key by its alias. + + Args: + key_alias: The unique key alias. + + Returns: + ProxyKey model if found, None otherwise. + """ + orm = await self._repository.get_by_alias(key_alias) + if orm is None: + return None + return self._orm_to_model(orm) + + async def revoke_key(self, proxy_key_id: UUID) -> ProxyKey | None: + """Revoke a proxy key. + + Marks the local metadata as inactive and attempts LiteLLM deletion. + LiteLLM API errors are intentionally silenced since revocation is + best-effort. + + Args: + proxy_key_id: The UUID of the proxy key to revoke. + + Returns: + Updated ProxyKey model if found, None otherwise. + """ + orm = await self._repository.get_by_id(proxy_key_id) + if orm is None: + return None + + if orm.status == ProxyKeyStatus.REVOKED.value: + return self._orm_to_model(orm) + + # Attempt LiteLLM deletion if we have the key_id + if orm.litellm_key_id and self.master_key: + async with httpx.AsyncClient(timeout=30.0) as client: + try: + response = await client.post( + f"{self.litellm_base_url}/key/delete", + headers=self._http_headers(), + json={"key_ids": [orm.litellm_key_id]}, + ) + response.raise_for_status() + except httpx.HTTPStatusError: + pass # Best-effort: key may already be expired/deleted + except httpx.RequestError: + pass # Best-effort: proxy may be temporarily unavailable + + # Mark local metadata as revoked + revoked_orm = await self._repository.revoke(proxy_key_id) + if revoked_orm is None: + return None + return self._orm_to_model(revoked_orm) + + def render_env_snippet( + self, + api_key: SecretStr, + proxy_base_url: str = "http://localhost:4000", + model: str | None = None, + harness_profile: str = "openai-compatible", + ) -> dict[str, str]: + """Render generic OpenAI-compatible environment snippet. + + Args: + api_key: The proxy key secret. + proxy_base_url: LiteLLM proxy base URL. + model: Optional default model alias. + harness_profile: Harness profile hint for rendering. + + Returns: + Dictionary of environment variables. + """ + env_vars: dict[str, str] = { + "OPENAI_API_BASE": proxy_base_url, + "OPENAI_API_KEY": api_key.get_secret_value(), + } + if model: + env_vars["OPENAI_MODEL"] = model + return env_vars + + def render_env_shell(self, env_vars: dict[str, str]) -> str: + """Render environment variables as shell export commands. + + Args: + env_vars: Dictionary of environment variables. + + Returns: + Shell commands string. + """ + lines = [] + for key in sorted(env_vars.keys()): + value = env_vars[key] + escaped = value.replace("'", "'\\''") + lines.append(f"export {key}='{escaped}'") + return "\n".join(lines) + + def render_env_dotenv(self, env_vars: dict[str, str]) -> str: + """Render environment variables as dotenv file content. + + Args: + env_vars: Dictionary of environment variables. + + Returns: + Dotenv file content string. + """ + lines = [] + for key in sorted(env_vars.keys()): + value = env_vars[key] + if " " in value or "#" in value or "'" in value or '"' in value: + escaped = value.replace("\\", "\\\\").replace('"', '\\"') + lines.append(f'{key}="{escaped}"') + else: + lines.append(f"{key}={value}") + return "\n".join(lines) + + def render_harness_env( + self, + api_key: SecretStr, + harness_profile_name: str, + proxy_base_url: str = "http://localhost:4000", + model: str | None = None, + ) -> dict[str, str]: + """Render harness-profile-based environment snippet. + + Args: + api_key: The proxy key secret. + harness_profile_name: Name of the harness profile. + proxy_base_url: LiteLLM proxy base URL. + model: Optional default model alias. + + Returns: + Dictionary of environment variables. + """ + # Generic OpenAI-compatible rendering + # Future: could load harness profile from config and adapt variable names + env_vars: dict[str, str] = { + "OPENAI_API_BASE": proxy_base_url, + "OPENAI_API_KEY": api_key.get_secret_value(), + } + if model: + env_vars["OPENAI_MODEL"] = model + env_vars["LITELLM_KEY_ALIAS"] = harness_profile_name + return env_vars + + def _orm_to_model(self, orm: ProxyKeyORM) -> ProxyKey: + """Convert an ORM ProxyKey to a domain model. + + Args: + orm: The ORM model. + + Returns: + The domain model. + """ + return ProxyKey( + proxy_key_id=orm.id, + key_alias=orm.key_alias, + litellm_key_id=orm.litellm_key_id, + owner=orm.owner, + team=orm.team, + customer=orm.customer, + purpose=orm.purpose, + allowed_models=list(orm.allowed_models) if orm.allowed_models else [], + budget_duration=orm.budget_duration, + budget_amount=orm.budget_amount, + budget_currency=orm.budget_currency, + status=ProxyKeyStatus(orm.status), + key_metadata=dict(orm.key_metadata) if orm.key_metadata else {}, + proxy_credential_id=orm.proxy_credential_id, + created_at=orm.created_at, + updated_at=orm.updated_at, + revoked_at=orm.revoked_at, + expires_at=orm.expires_at, + ) diff --git a/src/cli/commands/key.py b/src/cli/commands/key.py new file mode 100644 index 0000000..7a0f5b2 --- /dev/null +++ b/src/cli/commands/key.py @@ -0,0 +1,291 @@ +"""Sessionless proxy key commands for the CLI. + +Provides `benchmark key create`, `benchmark key list`, +`benchmark key info`, and `benchmark key revoke` commands. +""" + +import asyncio +from typing import Any +from uuid import UUID + +import typer +from rich.console import Console +from rich.table import Table +from sqlalchemy.orm import Session as SQLAlchemySession + +from benchmark_core.db.session import get_db_session +from benchmark_core.models import ProxyKeyStatus +from benchmark_core.repositories.proxy_key_repository import SQLProxyKeyRepository +from benchmark_core.services.proxy_key_service import ( + LiteLLMAPIError, + ProxyKeyService, + ProxyKeyServiceError, +) + +app = typer.Typer(help="Manage sessionless proxy keys") +console = Console() + + +def _get_service(db: SQLAlchemySession) -> ProxyKeyService: + """Build a ProxyKeyService from environment/config and a DB session.""" + import os + + litellm_url = os.environ.get("LITELLM_BASE_URL", "http://localhost:4000") + master_key = os.environ.get("LITELLM_MASTER_KEY") + + repository = SQLProxyKeyRepository(db) + return ProxyKeyService( + repository=repository, + litellm_base_url=litellm_url, + master_key=master_key, + enforce_https=False, + ) + + +@app.command() +def create( + alias: str | None = typer.Option(None, "--alias", "-a", help="Key alias"), + owner: str | None = typer.Option(None, "--owner", "-o", help="Key owner"), + team: str | None = typer.Option(None, "--team", "-t", help="Team label"), + customer: str | None = typer.Option(None, "--customer", "-c", help="Customer label"), + purpose: str | None = typer.Option(None, "--purpose", "-p", help="Key purpose/description"), + models: list[str] | None = typer.Option( + None, "--model", "-m", help="Allowed model alias (repeatable)" + ), + ttl_hours: int = typer.Option(168, "--ttl-hours", help="Key TTL in hours (default 7 days)"), + show_env: bool = typer.Option( + False, "--show-env", "-e", help="Print environment snippet after creation" + ), + format_type: str = typer.Option( + "shell", + "--format", + "-f", + help="Environment format: shell or dotenv", + ), +) -> None: + """Create a new sessionless proxy key. + + The key secret is displayed ONCE and never persisted. + Only non-secret metadata (alias, owner, team, etc.) is stored locally. + """ + console.print("[bold blue]Creating sessionless proxy key...[/bold blue]") + + with get_db_session() as db: + service = _get_service(db) + + try: + proxy_key, secret = _run( + service.create_key( + key_alias=alias, + owner=owner, + team=team, + customer=customer, + purpose=purpose, + allowed_models=models or None, + ttl_hours=ttl_hours, + ) + ) + except ProxyKeyServiceError as e: + console.print(f"[red]Error: {e}[/red]") + raise typer.Exit(1) from e + except LiteLLMAPIError as e: + console.print(f"[red]LiteLLM API error: {e}[/red]") + raise typer.Exit(1) from e + + # Display success + console.print("[green]Proxy key created successfully[/green]") + console.print(f" Alias: [bold]{proxy_key.key_alias}[/bold]") + console.print(f" ID: {proxy_key.proxy_key_id}") + if proxy_key.owner: + console.print(f" Owner: {proxy_key.owner}") + if proxy_key.team: + console.print(f" Team: {proxy_key.team}") + if proxy_key.customer: + console.print(f" Customer: {proxy_key.customer}") + if proxy_key.allowed_models: + console.print(f" Allowed Models: {', '.join(proxy_key.allowed_models)}") + console.print(f" Status: {proxy_key.status.value}") + console.print(f" Expires: {proxy_key.expires_at}") + + # Display the secret ONCE + console.print("\n[bold yellow]API Key Secret (copy now - will not be shown again):[/bold yellow]") + console.print(f" {secret.get_secret_value()}") + + # Optionally render environment snippet + if show_env: + env_vars = service.render_env_snippet(secret) + console.print("\n[bold blue]Environment Variables:[/bold blue]") + if format_type == "shell": + console.print(service.render_env_shell(env_vars)) + elif format_type == "dotenv": + console.print(service.render_env_dotenv(env_vars)) + else: + console.print(f"[yellow]Unknown format '{format_type}', using shell[/yellow]") + console.print(service.render_env_shell(env_vars)) + + # Warning + console.print("\n[dim]Note: The secret is not stored locally. Keep it safe.[/dim]") + + +@app.command("list") +def list_keys( + owner: str | None = typer.Option(None, "--owner", "-o", help="Filter by owner"), + team: str | None = typer.Option(None, "--team", "-t", help="Filter by team"), + customer: str | None = typer.Option(None, "--customer", "-c", help="Filter by customer"), + status: str | None = typer.Option(None, "--status", "-s", help="Filter by status"), + limit: int = typer.Option(50, "--limit", "-n", help="Maximum results"), +) -> None: + """List sessionless proxy keys.""" + status_enum = None + if status: + try: + status_enum = ProxyKeyStatus(status) + except ValueError as exc: + console.print(f"[red]Invalid status: {status}[/red]") + raise typer.Exit(1) from exc + + with get_db_session() as db: + service = _get_service(db) + try: + keys = _run( + service.list_keys( + owner=owner, + team=team, + customer=customer, + status=status_enum, + limit=limit, + ) + ) + except ProxyKeyServiceError as e: + console.print(f"[red]Error: {e}[/red]") + raise typer.Exit(1) from e + + if not keys: + console.print("[yellow]No proxy keys found[/yellow]") + return + + table = Table(title="Proxy Keys") + table.add_column("Alias", style="cyan") + table.add_column("Owner") + table.add_column("Team") + table.add_column("Customer") + table.add_column("Status", style="bold") + table.add_column("Allowed Models") + table.add_column("Created") + table.add_column("Expires") + table.add_column("Revoked") + + for key in keys: + status_style = { + ProxyKeyStatus.ACTIVE: "green", + ProxyKeyStatus.REVOKED: "red", + ProxyKeyStatus.EXPIRED: "yellow", + }.get(key.status, "white") + + revoked_str = key.revoked_at.isoformat() if key.revoked_at else "-" + + table.add_row( + key.key_alias, + key.owner or "-", + key.team or "-", + key.customer or "-", + f"[{status_style}]{key.status.value}[/{status_style}]", + ", ".join(key.allowed_models) if key.allowed_models else "-", + key.created_at.strftime("%Y-%m-%d %H:%M") if key.created_at else "-", + key.expires_at.strftime("%Y-%m-%d %H:%M") if key.expires_at else "-", + revoked_str[:16] if revoked_str != "-" else "-", + ) + + console.print(table) + console.print(f"\n[dim]Showing {len(keys)} key(s)[/dim]") + + +@app.command() +def info( + key_id: str = typer.Argument(..., help="Proxy key ID (UUID) or alias"), +) -> None: + """Show detailed info for a proxy key.""" + with get_db_session() as db: + service = _get_service(db) + + # Try as UUID first, then as alias + try: + pk_uuid = UUID(key_id) + proxy_key = _run(service.get_key_info(pk_uuid)) + except ValueError: + proxy_key = _run(service.get_key_by_alias(key_id)) + + if proxy_key is None: + console.print(f"[red]Proxy key not found: {key_id}[/red]") + raise typer.Exit(1) + + status_style = { + ProxyKeyStatus.ACTIVE: "green", + ProxyKeyStatus.REVOKED: "red", + ProxyKeyStatus.EXPIRED: "yellow", + }.get(proxy_key.status, "white") + + console.print(f"[bold]Proxy Key:[/bold] {proxy_key.key_alias}") + console.print(f" ID: {proxy_key.proxy_key_id}") + console.print(f" LiteLLM Key ID: {proxy_key.litellm_key_id or '-'}") + console.print(f" Status: [{status_style}]{proxy_key.status.value}[/{status_style}]") + if proxy_key.owner: + console.print(f" Owner: {proxy_key.owner}") + if proxy_key.team: + console.print(f" Team: {proxy_key.team}") + if proxy_key.customer: + console.print(f" Customer: {proxy_key.customer}") + if proxy_key.purpose: + console.print(f" Purpose: {proxy_key.purpose}") + if proxy_key.allowed_models: + console.print(f" Allowed Models: {', '.join(proxy_key.allowed_models)}") + if proxy_key.budget_amount is not None: + console.print(f" Budget: {proxy_key.budget_amount} {proxy_key.budget_currency}") + if proxy_key.budget_duration: + console.print(f" Budget Duration: {proxy_key.budget_duration}") + console.print(f" Created: {proxy_key.created_at}") + if proxy_key.expires_at: + console.print(f" Expires: {proxy_key.expires_at}") + if proxy_key.revoked_at: + console.print(f" Revoked: {proxy_key.revoked_at}") + if proxy_key.key_metadata: + console.print(f" Metadata: {proxy_key.key_metadata}") + + +@app.command() +def revoke( + key_id: str = typer.Argument(..., help="Proxy key ID (UUID)"), +) -> None: + """Revoke a sessionless proxy key. + + Marks the local metadata as inactive and attempts LiteLLM deletion. + The key secret cannot be recovered after revocation. + """ + try: + pk_uuid = UUID(key_id) + except ValueError as e: + console.print(f"[red]Invalid UUID: {key_id}[/red]") + raise typer.Exit(1) from e + + with get_db_session() as db: + service = _get_service(db) + try: + proxy_key = _run(service.revoke_key(pk_uuid)) + except ProxyKeyServiceError as e: + console.print(f"[red]Error: {e}[/red]") + raise typer.Exit(1) from e + + if proxy_key is None: + console.print(f"[red]Proxy key not found: {key_id}[/red]") + raise typer.Exit(1) + + console.print("[green]Proxy key revoked successfully[/green]") + console.print(f" Alias: {proxy_key.key_alias}") + console.print(f" Status: {proxy_key.status.value}") + if proxy_key.revoked_at: + console.print(f" Revoked at: {proxy_key.revoked_at}") + + +def _run(coro) -> Any: + """Run an async coroutine synchronously.""" + return asyncio.run(coro) diff --git a/src/cli/main.py b/src/cli/main.py index 6584208..8304ce3 100644 --- a/src/cli/main.py +++ b/src/cli/main.py @@ -3,7 +3,7 @@ import typer from rich.console import Console -from cli.commands import artifact, collect, config, export, health, normalize, render, session +from cli.commands import artifact, collect, config, export, health, key, normalize, render, session app = typer.Typer( name="benchmark", @@ -14,6 +14,7 @@ # Register subcommands app.add_typer(config.app, name="config", help="Config validation and management") app.add_typer(session.app, name="session", help="Session lifecycle commands") +app.add_typer(key.app, name="key", help="Sessionless proxy key management") app.add_typer(export.app, name="export", help="Export commands for reports") app.add_typer(normalize.app, name="normalize", help="Normalize LiteLLM request data") app.add_typer(artifact.app, name="artifact", help="Artifact registry management") diff --git a/tests/unit/test_cli_key_commands.py b/tests/unit/test_cli_key_commands.py new file mode 100644 index 0000000..c345bd6 --- /dev/null +++ b/tests/unit/test_cli_key_commands.py @@ -0,0 +1,340 @@ +"""Unit tests for sessionless proxy key CLI commands.""" + +import httpx +import pytest +import respx +from sqlalchemy import create_engine +from sqlalchemy.orm import sessionmaker +from typer.testing import CliRunner + +from benchmark_core.db.models import Base +from cli.main import app + + +@pytest.fixture +def test_engine(): + """Create an in-memory SQLite engine for testing.""" + engine = create_engine("sqlite:///:memory:") + Base.metadata.create_all(bind=engine) + return engine + + +@pytest.fixture +def db_session(test_engine): + """Create a database session for testing.""" + session_local = sessionmaker(bind=test_engine) + session = session_local() + try: + yield session + finally: + session.close() + + +@pytest.fixture +def runner(): + """Create a CLI runner.""" + return CliRunner() + + +@pytest.fixture +def mock_env_db_url(test_engine, monkeypatch): + """Mock the database URL to use the test engine.""" + from benchmark_core.db import session as db_session_module + from cli import commands + + original_get_db_session = db_session_module.get_db_session + + def mock_get_db_session(engine=None): + if engine is None: + engine = test_engine + return original_get_db_session(engine) + + monkeypatch.setattr(db_session_module, "get_db_session", mock_get_db_session) + monkeypatch.setattr(commands.key, "get_db_session", mock_get_db_session) + + # Also mock LiteLLM env vars to avoid errors + monkeypatch.setenv("LITELLM_MASTER_KEY", "test-master-key") + monkeypatch.setenv("LITELLM_BASE_URL", "http://localhost:4000") + + return test_engine + + +class TestKeyCreateCommand: + """Tests for `benchmark key create` command.""" + + @respx.mock + def test_create_key_success(self, runner, mock_env_db_url, monkeypatch): + """Create a key with metadata.""" + respx.post("http://localhost:4000/key/generate").mock( + return_value=httpx.Response( + 200, + json={"key": "sk-litellm-test-key", "key_id": "kid-123"}, + ) + ) + + result = runner.invoke( + app, + [ + "key", + "create", + "--alias", + "cli-test-alias", + "--owner", + "alice", + "--team", + "platform", + "--customer", + "internal", + "--model", + "gpt-4o", + "--ttl-hours", + "24", + ], + ) + + assert result.exit_code == 0, result.output + assert "cli-test-alias" in result.output + assert "alice" in result.output + assert "platform" in result.output + assert "internal" in result.output + assert "gpt-4o" in result.output + assert "sk-litellm-test-key" in result.output + assert "API Key Secret" in result.output + assert "not be shown again" in result.output + + @respx.mock + def test_create_key_with_env_flag(self, runner, mock_env_db_url, monkeypatch): + """Create a key with --show-env renders environment snippet.""" + respx.post("http://localhost:4000/key/generate").mock( + return_value=httpx.Response( + 200, + json={"key": "sk-env-key", "key_id": "kid-env"}, + ) + ) + + result = runner.invoke( + app, + [ + "key", + "create", + "--alias", + "env-test", + "--show-env", + "--format", + "shell", + ], + ) + + assert result.exit_code == 0, result.output + assert "export OPENAI_API_BASE" in result.output + assert "export OPENAI_API_KEY" in result.output + + @respx.mock + def test_create_key_dotenv_format(self, runner, mock_env_db_url, monkeypatch): + """Create a key with --format dotenv.""" + respx.post("http://localhost:4000/key/generate").mock( + return_value=httpx.Response( + 200, + json={"key": "sk-dotenv-key", "key_id": "kid-dotenv"}, + ) + ) + + result = runner.invoke( + app, + [ + "key", + "create", + "--alias", + "dotenv-test", + "--show-env", + "--format", + "dotenv", + ], + ) + + assert result.exit_code == 0, result.output + assert "OPENAI_API_BASE=" in result.output + assert "OPENAI_API_KEY=" in result.output + + def test_create_key_no_master_key(self, runner, mock_env_db_url, monkeypatch): + """Create fails when LITELLM_MASTER_KEY is not set.""" + monkeypatch.delenv("LITELLM_MASTER_KEY", raising=False) + + result = runner.invoke(app, ["key", "create"]) + assert result.exit_code == 1 + assert "master_key" in result.output or "not configured" in result.output + + @respx.mock + def test_create_key_litellm_error(self, runner, mock_env_db_url, monkeypatch): + """LiteLLM API errors are handled gracefully.""" + respx.post("http://localhost:4000/key/generate").mock( + return_value=httpx.Response(500, text="Internal Server Error") + ) + + result = runner.invoke(app, ["key", "create", "--alias", "fail-test"]) + assert result.exit_code == 1 + assert "LiteLLM API error" in result.output + + +class TestKeyListCommand: + """Tests for `benchmark key list` command.""" + + @respx.mock + def test_list_keys_empty(self, runner, mock_env_db_url, monkeypatch): + """List with no keys shows empty message.""" + result = runner.invoke(app, ["key", "list"]) + assert result.exit_code == 0, result.output + assert "No proxy keys found" in result.output + + @respx.mock + def test_list_keys_with_data(self, runner, mock_env_db_url, monkeypatch): + """List shows created keys.""" + respx.post("http://localhost:4000/key/generate").mock( + return_value=httpx.Response( + 200, + json={"key": "sk-1", "key_id": "id-1"}, + ) + ) + + # Create a key first + result = runner.invoke( + app, + [ + "key", + "create", + "--alias", + "list-test", + "--owner", + "alice", + "--team", + "platform", + ], + ) + assert result.exit_code == 0, result.output + + # List keys + result = runner.invoke(app, ["key", "list"]) + assert result.exit_code == 0, result.output + # Rich table may truncate fields; check for partial match and other fields + assert "list-" in result.output + assert "alice" in result.output + assert "platf" in result.output + + def test_list_keys_filter_by_owner(self, runner, mock_env_db_url, monkeypatch): + """List with --owner filter.""" + result = runner.invoke(app, ["key", "list", "--owner", "nobody"]) + assert result.exit_code == 0, result.output + assert "No proxy keys found" in result.output + + def test_list_keys_invalid_status(self, runner, mock_env_db_url, monkeypatch): + """List with invalid status exits with error.""" + result = runner.invoke(app, ["key", "list", "--status", "invalid_status"]) + assert result.exit_code == 1 + assert "Invalid status" in result.output + + +class TestKeyInfoCommand: + """Tests for `benchmark key info` command.""" + + @respx.mock + def test_info_by_alias(self, runner, mock_env_db_url, monkeypatch): + """Info command finds key by alias.""" + respx.post("http://localhost:4000/key/generate").mock( + return_value=httpx.Response( + 200, + json={"key": "sk-1", "key_id": "id-1"}, + ) + ) + + # Create key + result = runner.invoke( + app, + ["key", "create", "--alias", "info-alias", "--owner", "bob"], + ) + assert result.exit_code == 0, result.output + + # Get info by alias + result = runner.invoke(app, ["key", "info", "info-alias"]) + assert result.exit_code == 0, result.output + assert "info-alias" in result.output + assert "bob" in result.output + + def test_info_not_found(self, runner, mock_env_db_url, monkeypatch): + """Info for missing key shows error.""" + result = runner.invoke(app, ["key", "info", "nonexistent-key"]) + assert result.exit_code == 1 + assert "not found" in result.output + + def test_info_invalid_uuid(self, runner, mock_env_db_url, monkeypatch): + """Info with invalid UUID-like input falls back to alias search.""" + result = runner.invoke(app, ["key", "info", "not-a-uuid-or-alias"]) + assert result.exit_code == 1 + assert "not found" in result.output + + +class TestKeyRevokeCommand: + """Tests for `benchmark key revoke` command.""" + + @respx.mock + def test_revoke_key(self, runner, mock_env_db_url, monkeypatch): + """Revoke an existing key.""" + respx.post("http://localhost:4000/key/generate").mock( + return_value=httpx.Response( + 200, + json={"key": "sk-1", "key_id": "id-1"}, + ) + ) + respx.post("http://localhost:4000/key/delete").mock( + return_value=httpx.Response(200, json={"deleted_keys": ["id-1"]}) + ) + + # Create key + result = runner.invoke( + app, + ["key", "create", "--alias", "revoke-test"], + ) + assert result.exit_code == 0, result.output + + # Extract UUID from output + import re + uuid_match = re.search(r"ID: ([0-9a-f-]{36})", result.output) + assert uuid_match, f"Could not find UUID in output: {result.output}" + key_id = uuid_match.group(1) + + # Revoke by ID + result = runner.invoke(app, ["key", "revoke", key_id]) + assert result.exit_code == 0, result.output + assert "revoked successfully" in result.output + assert "revoked" in result.output.lower() + + def test_revoke_not_found(self, runner, mock_env_db_url, monkeypatch): + """Revoke non-existent key shows error.""" + result = runner.invoke( + app, ["key", "revoke", "12345678-1234-1234-1234-123456789abc"] + ) + assert result.exit_code == 1 + assert "not found" in result.output + + def test_revoke_invalid_uuid(self, runner, mock_env_db_url, monkeypatch): + """Revoke with invalid UUID shows error.""" + result = runner.invoke(app, ["key", "revoke", "not-a-uuid"]) + assert result.exit_code == 1 + assert "Invalid UUID" in result.output + + +class TestKeyCLIImports: + """Smoke tests for key command imports.""" + + def test_import_key_module(self): + """Key module imports successfully.""" + from cli.commands import key + + assert key.app is not None + + def test_main_app_has_key_command(self): + """Main CLI app includes key subcommand.""" + from cli.main import app + + # Typer stores subcommands in register_list or similar + # Just verify import works + assert app is not None diff --git a/tests/unit/test_proxy_key_service.py b/tests/unit/test_proxy_key_service.py new file mode 100644 index 0000000..d5e9cbe --- /dev/null +++ b/tests/unit/test_proxy_key_service.py @@ -0,0 +1,438 @@ +"""Unit tests for sessionless proxy key service.""" + +from uuid import uuid4 + +import httpx +import pytest +import respx +from pydantic import SecretStr +from sqlalchemy import create_engine +from sqlalchemy.orm import sessionmaker + +from benchmark_core.db.models import Base +from benchmark_core.models import ProxyKey, ProxyKeyStatus +from benchmark_core.repositories.proxy_key_repository import SQLProxyKeyRepository +from benchmark_core.services.proxy_key_service import ( + LiteLLMAPIError, + ProxyKeyService, + ProxyKeyServiceError, +) + + +@pytest.fixture +def test_engine(): + """Create an in-memory SQLite engine for testing.""" + engine = create_engine("sqlite:///:memory:") + Base.metadata.create_all(bind=engine) + return engine + + +@pytest.fixture +def db_session(test_engine): + """Create a database session for testing.""" + session_local = sessionmaker(bind=test_engine) + session = session_local() + try: + yield session + finally: + session.close() + + +@pytest.fixture +def repository(db_session): + """Create a SQLProxyKeyRepository.""" + return SQLProxyKeyRepository(db_session) + + +@pytest.fixture +def service(repository): + """Create a ProxyKeyService with test config.""" + return ProxyKeyService( + repository=repository, + litellm_base_url="http://localhost:4000", + master_key="test-master-key", + enforce_https=False, + ) + + +class TestProxyKeyCreation: + """Test proxy key creation via LiteLLM API.""" + + @pytest.mark.asyncio + @respx.mock + async def test_create_key_success(self, service): + """Successful key creation returns metadata and secret.""" + respx.post("http://localhost:4000/key/generate").mock( + return_value=httpx.Response( + 200, + json={ + "key": "sk-litellm-test-key-12345", + "key_id": "litellm-key-id-123", + }, + ) + ) + + proxy_key, secret = await service.create_key( + key_alias="test-alias", + owner="alice", + team="platform", + customer="internal", + allowed_models=["gpt-4o"], + ttl_hours=24, + ) + + assert isinstance(proxy_key, ProxyKey) + assert proxy_key.key_alias == "test-alias" + assert proxy_key.owner == "alice" + assert proxy_key.team == "platform" + assert proxy_key.customer == "internal" + assert proxy_key.allowed_models == ["gpt-4o"] + assert proxy_key.status == ProxyKeyStatus.ACTIVE + assert proxy_key.litellm_key_id == "litellm-key-id-123" + assert isinstance(secret, SecretStr) + assert secret.get_secret_value() == "sk-litellm-test-key-12345" + + @pytest.mark.asyncio + @respx.mock + async def test_create_key_auto_alias(self, service): + """Auto-generates alias when not provided.""" + respx.post("http://localhost:4000/key/generate").mock( + return_value=httpx.Response( + 200, + json={"key": "sk-auto-key", "key_id": "auto-id"}, + ) + ) + + proxy_key, secret = await service.create_key() + + assert proxy_key.key_alias.startswith("usage-") + assert len(proxy_key.key_alias) > 8 + + @pytest.mark.asyncio + @respx.mock + async def test_create_key_persists_metadata(self, service, repository, db_session): + """Key metadata is persisted to the database (no secret).""" + respx.post("http://localhost:4000/key/generate").mock( + return_value=httpx.Response( + 200, + json={"key": "sk-secret-123", "key_id": "kid-123"}, + ) + ) + + proxy_key, secret = await service.create_key( + key_alias="persist-test", + owner="bob", + team="dev", + ) + + # Force commit for in-memory SQLite + db_session.commit() + + # Verify in database - no secret stored + orm = await repository.get_by_alias("persist-test") + assert orm is not None + assert orm.key_alias == "persist-test" + assert orm.owner == "bob" + assert orm.team == "dev" + # Secret must NOT be in the ORM + assert not hasattr(orm, "api_key") + + @pytest.mark.asyncio + @respx.mock + async def test_create_key_with_usage_policy(self, service): + """Usage policy defaults are applied.""" + from benchmark_core.config import UsagePolicyProfile + + policy = UsagePolicyProfile( + name="team-policy", + owner="default-owner", + team="default-team", + customer="default-customer", + allowed_models=["kimi-k2-5"], + budget_duration="30d", + budget_amount=500.0, + ttl_seconds=7200, + ) + + respx.post("http://localhost:4000/key/generate").mock( + return_value=httpx.Response( + 200, + json={"key": "sk-policy-key", "key_id": "policy-id"}, + ) + ) + + proxy_key, secret = await service.create_key( + usage_policy=policy, + owner="override-owner", + ) + + assert proxy_key.owner == "override-owner" # explicit override + assert proxy_key.team == "default-team" # from policy + assert proxy_key.customer == "default-customer" + assert proxy_key.allowed_models == ["kimi-k2-5"] + assert proxy_key.budget_duration == "30d" + assert proxy_key.budget_amount == 500.0 + + @pytest.mark.asyncio + async def test_create_key_requires_master_key(self, repository): + """Key creation fails without master_key.""" + svc = ProxyKeyService( + repository=repository, + litellm_base_url="http://localhost:4000", + master_key=None, + enforce_https=False, + ) + + with pytest.raises(ProxyKeyServiceError, match="master_key not configured"): + await svc.create_key() + + @pytest.mark.asyncio + @respx.mock + async def test_create_key_handles_api_error(self, service): + """LiteLLM API errors are surfaced.""" + respx.post("http://localhost:4000/key/generate").mock( + return_value=httpx.Response(500, text="Internal Server Error") + ) + + with pytest.raises(LiteLLMAPIError, match="LiteLLM API error"): + await service.create_key() + + @pytest.mark.asyncio + @respx.mock + async def test_create_key_handles_missing_key_field(self, service): + """Missing 'key' in LiteLLM response raises error.""" + respx.post("http://localhost:4000/key/generate").mock( + return_value=httpx.Response(200, json={"key_id": "no-key"}) + ) + + with pytest.raises(LiteLLMAPIError, match="missing 'key' field"): + await service.create_key() + + +class TestProxyKeyListing: + """Test proxy key listing operations.""" + + @pytest.mark.asyncio + async def test_list_keys_empty(self, service): + """Empty repository returns empty list.""" + keys = await service.list_keys() + assert keys == [] + + @pytest.mark.asyncio + @respx.mock + async def test_list_keys_filter_by_owner(self, service, db_session): + """Filtering by owner returns matching keys.""" + respx.post("http://localhost:4000/key/generate").mock( + side_effect=[ + httpx.Response(200, json={"key": "sk-1", "key_id": "id-1"}), + httpx.Response(200, json={"key": "sk-2", "key_id": "id-2"}), + ] + ) + + await service.create_key(key_alias="k1", owner="alice") + await service.create_key(key_alias="k2", owner="bob") + db_session.commit() + + alice_keys = await service.list_keys(owner="alice") + assert len(alice_keys) == 1 + assert alice_keys[0].key_alias == "k1" + assert alice_keys[0].owner == "alice" + + @pytest.mark.asyncio + @respx.mock + async def test_list_keys_filter_by_status(self, service, db_session): + """Filtering by status returns matching keys.""" + respx.post("http://localhost:4000/key/generate").mock( + return_value=httpx.Response(200, json={"key": "sk-1", "key_id": "id-1"}) + ) + + proxy_key, _ = await service.create_key(key_alias="status-test") + db_session.commit() + + active_keys = await service.list_keys(status=ProxyKeyStatus.ACTIVE) + assert len(active_keys) == 1 + assert active_keys[0].status == ProxyKeyStatus.ACTIVE + + +class TestProxyKeyRevocation: + """Test proxy key revocation.""" + + @pytest.mark.asyncio + @respx.mock + async def test_revoke_key_marks_inactive(self, service, db_session): + """Revocation marks key as revoked locally.""" + respx.post("http://localhost:4000/key/generate").mock( + return_value=httpx.Response(200, json={"key": "sk-1", "key_id": "id-1"}) + ) + respx.post("http://localhost:4000/key/delete").mock( + return_value=httpx.Response(200, json={"deleted_keys": ["id-1"]}) + ) + + proxy_key, _ = await service.create_key(key_alias="revoke-test") + db_session.commit() + + revoked = await service.revoke_key(proxy_key.proxy_key_id) + assert revoked is not None + assert revoked.status == ProxyKeyStatus.REVOKED + assert revoked.revoked_at is not None + + @pytest.mark.asyncio + @respx.mock + async def test_revoke_key_not_found(self, service): + """Revoking non-existent key returns None.""" + result = await service.revoke_key(uuid4()) + assert result is None + + @pytest.mark.asyncio + @respx.mock + async def test_revoke_key_silences_litellm_error(self, service, db_session): + """LiteLLM delete errors are silently ignored.""" + respx.post("http://localhost:4000/key/generate").mock( + return_value=httpx.Response(200, json={"key": "sk-1", "key_id": "id-1"}) + ) + respx.post("http://localhost:4000/key/delete").mock( + return_value=httpx.Response(404, text="Not Found") + ) + + proxy_key, _ = await service.create_key(key_alias="silent-revoke") + db_session.commit() + + # Should not raise despite 404 from LiteLLM + revoked = await service.revoke_key(proxy_key.proxy_key_id) + assert revoked is not None + assert revoked.status == ProxyKeyStatus.REVOKED + + @pytest.mark.asyncio + @respx.mock + async def test_revoke_already_revoked(self, service, db_session): + """Revoking an already revoked key is a no-op.""" + respx.post("http://localhost:4000/key/generate").mock( + return_value=httpx.Response(200, json={"key": "sk-1", "key_id": "id-1"}) + ) + respx.post("http://localhost:4000/key/delete").mock( + return_value=httpx.Response(200, json={"deleted_keys": ["id-1"]}) + ) + + proxy_key, _ = await service.create_key(key_alias="double-revoke") + db_session.commit() + + revoked = await service.revoke_key(proxy_key.proxy_key_id) + db_session.commit() + assert revoked is not None + assert revoked.status == ProxyKeyStatus.REVOKED + + second = await service.revoke_key(proxy_key.proxy_key_id) + assert second is not None + assert second.status == ProxyKeyStatus.REVOKED + + +class TestProxyKeyInfo: + """Test proxy key info retrieval.""" + + @pytest.mark.asyncio + @respx.mock + async def test_get_key_info(self, service, db_session): + """Retrieve key info by ID.""" + respx.post("http://localhost:4000/key/generate").mock( + return_value=httpx.Response(200, json={"key": "sk-1", "key_id": "id-1"}) + ) + + proxy_key, _ = await service.create_key(key_alias="info-test", owner="alice") + db_session.commit() + + found = await service.get_key_info(proxy_key.proxy_key_id) + assert found is not None + assert found.key_alias == "info-test" + assert found.owner == "alice" + + @pytest.mark.asyncio + async def test_get_key_info_not_found(self, service): + """Missing key returns None.""" + found = await service.get_key_info(uuid4()) + assert found is None + + @pytest.mark.asyncio + @respx.mock + async def test_get_key_by_alias(self, service, db_session): + """Retrieve key by alias.""" + respx.post("http://localhost:4000/key/generate").mock( + return_value=httpx.Response(200, json={"key": "sk-1", "key_id": "id-1"}) + ) + + proxy_key, _ = await service.create_key(key_alias="alias-lookup") + db_session.commit() + + found = await service.get_key_by_alias("alias-lookup") + assert found is not None + assert found.proxy_key_id == proxy_key.proxy_key_id + + +class TestEnvRendering: + """Test environment snippet rendering.""" + + def test_render_env_snippet_basic(self, service): + """Basic OpenAI-compatible env vars.""" + secret = SecretStr("sk-test-key") + env = service.render_env_snippet(secret, proxy_base_url="http://proxy:4000") + + assert env["OPENAI_API_BASE"] == "http://proxy:4000" + assert env["OPENAI_API_KEY"] == "sk-test-key" + + def test_render_env_snippet_with_model(self, service): + """Includes model when provided.""" + secret = SecretStr("sk-test-key") + env = service.render_env_snippet( + secret, proxy_base_url="http://proxy:4000", model="gpt-4o" + ) + + assert env["OPENAI_MODEL"] == "gpt-4o" + + def test_render_env_shell(self, service): + """Shell format produces export commands.""" + env = {"KEY1": "value1", "KEY2": "value with spaces"} + shell = service.render_env_shell(env) + + assert "export KEY1='value1'" in shell + assert "export KEY2='value with spaces'" in shell + + def test_render_env_dotenv(self, service): + """Dotenv format produces .env content.""" + env = {"KEY1": "value1", "KEY2": "value with spaces"} + dotenv = service.render_env_dotenv(env) + + assert "KEY1=value1" in dotenv + assert 'KEY2="value with spaces"' in dotenv + + +class TestHTTPSValidation: + """Test HTTPS validation.""" + + def test_accepts_localhost(self, repository): + """Accepts localhost URLs.""" + svc = ProxyKeyService( + repository=repository, + litellm_base_url="http://localhost:4000", + master_key="test", + enforce_https=False, + ) + assert svc.litellm_base_url == "http://localhost:4000" + + def test_rejects_http_in_production(self, repository): + """Rejects HTTP in production.""" + with pytest.raises(ValueError, match="must use HTTPS"): + ProxyKeyService( + repository=repository, + litellm_base_url="http://example.com", + master_key="test", + enforce_https=True, + ) + + def test_accepts_https(self, repository): + """Accepts HTTPS URLs.""" + svc = ProxyKeyService( + repository=repository, + litellm_base_url="https://litellm.example.com", + master_key="test", + enforce_https=True, + ) + assert svc.litellm_base_url == "https://litellm.example.com" diff --git a/tests/unit/test_services.py b/tests/unit/test_services.py index f245276..4185b67 100644 --- a/tests/unit/test_services.py +++ b/tests/unit/test_services.py @@ -2,7 +2,10 @@ from uuid import uuid4 +import httpx import pytest +import respx +from pydantic import SecretStr from sqlalchemy import create_engine from sqlalchemy.orm import sessionmaker @@ -80,8 +83,8 @@ def session_service(db_session): @pytest.fixture def credential_service(): - """Create a credential service.""" - return CredentialService() + """Create a credential service with test master key.""" + return CredentialService(master_key="test-master-key") @pytest.fixture @@ -631,8 +634,20 @@ async def test_create_session_with_notes(self, session_service, db_session, setu class TestCredentialService: """Tests for CredentialService.""" + @pytest.mark.asyncio + @respx.mock async def test_issue_credential(self, credential_service): - """Test issuing a credential.""" + """Test issuing a credential via LiteLLM API.""" + respx.post("http://localhost:4000/key/generate").mock( + return_value=httpx.Response( + 200, + json={ + "key": "sk-benchmark-test-key-12345", + "key_id": "litellm-key-id-123", + }, + ) + ) + session_id = uuid4() credential = await credential_service.issue_credential( session_id=session_id, @@ -641,43 +656,47 @@ async def test_issue_credential(self, credential_service): harness_profile="default", ) - assert credential.startswith("sk-benchmark-") - assert str(session_id) in credential + assert credential.key_alias.startswith("session-") + assert str(session_id)[:8] in credential.key_alias + assert credential.api_key.get_secret_value().startswith("sk-benchmark-") def test_render_env_snippet(self, credential_service): """Test rendering environment snippet.""" + from benchmark_core.models import ProxyCredential + + credential = ProxyCredential( + session_id=uuid4(), + key_alias="test-alias", + api_key=SecretStr("sk-test"), + experiment_id="test-experiment", + variant_id="test-variant", + harness_profile="default", + ) env = credential_service.render_env_snippet( - credential="sk-test", + credential=credential, proxy_base_url="http://localhost:4000", model="gpt-4o", - harness_profile="default", ) assert env["OPENAI_API_BASE"] == "http://localhost:4000" assert env["OPENAI_API_KEY"] == "sk-test" assert env["OPENAI_MODEL"] == "gpt-4o" - def test_render_shell_snippet(self, credential_service): + def test_render_env_shell(self, credential_service): """Test rendering shell snippet.""" - snippet = credential_service.render_shell_snippet( - credential="sk-test", - proxy_base_url="http://localhost:4000", - model="gpt-4o", - ) + env = {"OPENAI_API_BASE": "http://localhost:4000", "OPENAI_API_KEY": "sk-test"} + snippet = credential_service.render_env_shell(env) - assert 'export OPENAI_API_BASE="http://localhost:4000"' in snippet - assert 'export OPENAI_API_KEY="sk-test"' in snippet + assert "export OPENAI_API_BASE='http://localhost:4000'" in snippet + assert "export OPENAI_API_KEY='sk-test'" in snippet - def test_render_dotenv_snippet(self, credential_service): + def test_render_env_dotenv(self, credential_service): """Test rendering dotenv snippet.""" - snippet = credential_service.render_dotenv_snippet( - credential="sk-test", - proxy_base_url="http://localhost:4000", - model="gpt-4o", - ) + env = {"OPENAI_API_BASE": "http://localhost:4000", "OPENAI_API_KEY": "sk-test"} + snippet = credential_service.render_env_dotenv(env) - assert 'OPENAI_API_BASE="http://localhost:4000"' in snippet - assert 'OPENAI_API_KEY="sk-test"' in snippet + assert "OPENAI_API_BASE=http://localhost:4000" in snippet + assert "OPENAI_API_KEY=sk-test" in snippet class TestBenchmarkMetadataService: From a96c6b5b052dc7a37d7ebba443dac6a12ce3c2cf Mon Sep 17 00:00:00 2001 From: Leonardo Gonzalez Date: Tue, 21 Apr 2026 22:16:24 -0500 Subject: [PATCH 2/7] fix: formatting and type-check errors for failing CI checks\n\n- ruff format src tests\n- add type annotation to _run(coro: Any) -> Any Co-authored-by: openhands --- src/benchmark_core/services/proxy_key_service.py | 8 ++++---- src/cli/commands/key.py | 6 ++++-- tests/unit/test_cli_key_commands.py | 5 ++--- tests/unit/test_proxy_key_service.py | 6 ++---- 4 files changed, 12 insertions(+), 13 deletions(-) diff --git a/src/benchmark_core/services/proxy_key_service.py b/src/benchmark_core/services/proxy_key_service.py index d951939..0516177 100644 --- a/src/benchmark_core/services/proxy_key_service.py +++ b/src/benchmark_core/services/proxy_key_service.py @@ -126,9 +126,7 @@ async def create_key( ProxyKeyServiceError: If master_key not configured. """ if self.master_key is None: - raise ProxyKeyServiceError( - "Cannot create key: LiteLLM master_key not configured." - ) + raise ProxyKeyServiceError("Cannot create key: LiteLLM master_key not configured.") # Apply usage policy defaults if provided effective_alias = key_alias or self._build_key_alias( @@ -147,7 +145,9 @@ async def create_key( usage_policy.budget_amount if usage_policy else None ) effective_ttl = ( - usage_policy.ttl_seconds // 3600 if usage_policy and usage_policy.ttl_seconds else ttl_hours + usage_policy.ttl_seconds // 3600 + if usage_policy and usage_policy.ttl_seconds + else ttl_hours ) # Merge metadata diff --git a/src/cli/commands/key.py b/src/cli/commands/key.py index 7a0f5b2..881740a 100644 --- a/src/cli/commands/key.py +++ b/src/cli/commands/key.py @@ -108,7 +108,9 @@ def create( console.print(f" Expires: {proxy_key.expires_at}") # Display the secret ONCE - console.print("\n[bold yellow]API Key Secret (copy now - will not be shown again):[/bold yellow]") + console.print( + "\n[bold yellow]API Key Secret (copy now - will not be shown again):[/bold yellow]" + ) console.print(f" {secret.get_secret_value()}") # Optionally render environment snippet @@ -286,6 +288,6 @@ def revoke( console.print(f" Revoked at: {proxy_key.revoked_at}") -def _run(coro) -> Any: +def _run(coro: Any) -> Any: """Run an async coroutine synchronously.""" return asyncio.run(coro) diff --git a/tests/unit/test_cli_key_commands.py b/tests/unit/test_cli_key_commands.py index c345bd6..8c7a178 100644 --- a/tests/unit/test_cli_key_commands.py +++ b/tests/unit/test_cli_key_commands.py @@ -297,6 +297,7 @@ def test_revoke_key(self, runner, mock_env_db_url, monkeypatch): # Extract UUID from output import re + uuid_match = re.search(r"ID: ([0-9a-f-]{36})", result.output) assert uuid_match, f"Could not find UUID in output: {result.output}" key_id = uuid_match.group(1) @@ -309,9 +310,7 @@ def test_revoke_key(self, runner, mock_env_db_url, monkeypatch): def test_revoke_not_found(self, runner, mock_env_db_url, monkeypatch): """Revoke non-existent key shows error.""" - result = runner.invoke( - app, ["key", "revoke", "12345678-1234-1234-1234-123456789abc"] - ) + result = runner.invoke(app, ["key", "revoke", "12345678-1234-1234-1234-123456789abc"]) assert result.exit_code == 1 assert "not found" in result.output diff --git a/tests/unit/test_proxy_key_service.py b/tests/unit/test_proxy_key_service.py index d5e9cbe..5031dc0 100644 --- a/tests/unit/test_proxy_key_service.py +++ b/tests/unit/test_proxy_key_service.py @@ -167,7 +167,7 @@ async def test_create_key_with_usage_policy(self, service): ) assert proxy_key.owner == "override-owner" # explicit override - assert proxy_key.team == "default-team" # from policy + assert proxy_key.team == "default-team" # from policy assert proxy_key.customer == "default-customer" assert proxy_key.allowed_models == ["kimi-k2-5"] assert proxy_key.budget_duration == "30d" @@ -381,9 +381,7 @@ def test_render_env_snippet_basic(self, service): def test_render_env_snippet_with_model(self, service): """Includes model when provided.""" secret = SecretStr("sk-test-key") - env = service.render_env_snippet( - secret, proxy_base_url="http://proxy:4000", model="gpt-4o" - ) + env = service.render_env_snippet(secret, proxy_base_url="http://proxy:4000", model="gpt-4o") assert env["OPENAI_MODEL"] == "gpt-4o" From bbe8b6572f0952c91e17ec4a482d4b5ee5fb042b Mon Sep 17 00:00:00 2001 From: Leonardo Gonzalez Date: Tue, 21 Apr 2026 22:25:30 -0500 Subject: [PATCH 3/7] refactor: address AI review feedback on PR #45 - Extract shared validate_litellm_url(), render_env_shell(), render_env_dotenv(), warn_revoke_failure() utilities into benchmark_core/services/common.py - Remove dead _legacy_* methods from CredentialService (~65 lines) - Replace stale CredentialService class in services_abc.py with re-export - Add warnings.warn on LiteLLM revoke failures in both services - Ensure render_env_* methods sort keys deterministically (both services) - Fix all ruff/mypy issues, 584 tests passing Co-authored-by: openhands --- src/benchmark_core/services/common.py | 95 ++++++ .../services/credential_service.py | 114 +------ .../services/proxy_key_service.py | 40 +-- src/benchmark_core/services_abc.py | 282 +----------------- 4 files changed, 130 insertions(+), 401 deletions(-) create mode 100644 src/benchmark_core/services/common.py diff --git a/src/benchmark_core/services/common.py b/src/benchmark_core/services/common.py new file mode 100644 index 0000000..a2ecf62 --- /dev/null +++ b/src/benchmark_core/services/common.py @@ -0,0 +1,95 @@ +"""Shared utilities for benchmark core services. + +Provides common helper functions used across multiple service classes +to avoid duplication and ensure consistent behavior. +""" + +import warnings + + +def validate_litellm_url(litellm_base_url: str, enforce_https: bool = True) -> str: + """Validate and normalize a LiteLLM base URL. + + Strips trailing slashes and optionally enforces HTTPS in + production environments. + + Args: + litellm_base_url: The raw base URL for the LiteLLM proxy API. + enforce_https: Whether to require HTTPS. Defaults to True. + + Returns: + The normalized URL with trailing slash removed. + + Raises: + ValueError: If HTTPS is required but the URL is not localhost + and does not start with ``https://``. + """ + normalized = litellm_base_url.rstrip("/") + if enforce_https and not normalized.startswith( + ("https://", "http://localhost", "http://127.0.0.1") + ): + raise ValueError( + f"LiteLLM URL must use HTTPS in production environments. Got: {litellm_base_url}" + ) + return normalized + + +def render_env_shell(env_vars: dict[str, str]) -> str: + """Render environment variables as shell export commands. + + Variables are sorted alphabetically for deterministic output. + + Args: + env_vars: Dictionary of environment variables. + + Returns: + Shell export commands string. + """ + lines: list[str] = [] + for key in sorted(env_vars.keys()): + value = env_vars[key] + escaped = value.replace("'", "'\\''") + lines.append(f"export {key}='{escaped}'") + return "\n".join(lines) + + +def render_env_dotenv(env_vars: dict[str, str]) -> str: + """Render environment variables as dotenv file content. + + Variables are sorted alphabetically for deterministic output. + + Args: + env_vars: Dictionary of environment variables. + + Returns: + Dotenv file content string. + """ + lines: list[str] = [] + for key in sorted(env_vars.keys()): + value = env_vars[key] + if " " in value or "#" in value or "'" in value or '"' in value: + escaped = value.replace("\\", "\\\\").replace('"', '\\"') + lines.append(f'{key}="{escaped}"') + else: + lines.append(f"{key}={value}") + return "\n".join(lines) + + +def warn_revoke_failure(source: str, litellm_key_id: str | None, exc: Exception) -> None: + """Warn when a LiteLLM key revocation attempt fails. + + This helper centralises the warning issued when a best-effort LiteLLM + delete call fails, so both :class:`CredentialService` and + :class:`ProxyKeyService` behave consistently. + + Args: + source: Human-readable source of the revocation (e.g. ``credential`` + or ``proxy_key``). + litellm_key_id: The LiteLLM key ID that could not be deleted. + exc: The original exception that caused the failure. + """ + warnings.warn( + f"LiteLLM delete failed for {source} (key_id={litellm_key_id}): {exc}. " + "Local metadata has been marked revoked but the proxy key may still be active.", + stacklevel=3, + ) diff --git a/src/benchmark_core/services/credential_service.py b/src/benchmark_core/services/credential_service.py index 2573c8f..d0d450b 100644 --- a/src/benchmark_core/services/credential_service.py +++ b/src/benchmark_core/services/credential_service.py @@ -13,6 +13,12 @@ from pydantic import SecretStr from benchmark_core.models import ProxyCredential +from benchmark_core.services.common import ( + render_env_dotenv, + render_env_shell, + validate_litellm_url, + warn_revoke_failure, +) class CredentialService: @@ -38,17 +44,9 @@ def __init__( Raises: ValueError: If litellm_base_url doesn't use HTTPS in production """ - self.litellm_base_url = litellm_base_url.rstrip("/") + self.litellm_base_url = validate_litellm_url(litellm_base_url, enforce_https) self.master_key = master_key - # Validate HTTPS in production environments - if enforce_https and not self.litellm_base_url.startswith( - ("https://", "http://localhost", "http://127.0.0.1") - ): - raise ValueError( - f"LiteLLM URL must use HTTPS in production environments. Got: {litellm_base_url}" - ) - def _generate_key_alias( self, session_id: UUID, @@ -196,7 +194,8 @@ async def revoke_credential(self, credential: ProxyCredential) -> ProxyCredentia Note: LiteLLM API errors are intentionally silenced since revocation is best-effort - the key may already be expired or the proxy may be - temporarily unavailable. + temporarily unavailable. A warning is still emitted so operators + can detect state inconsistency. """ # Call LiteLLM API to delete the key if we have the key_id if credential.litellm_key_id and self.master_key: @@ -208,12 +207,10 @@ async def revoke_credential(self, credential: ProxyCredential) -> ProxyCredentia json={"key_ids": [credential.litellm_key_id]}, ) response.raise_for_status() - except httpx.HTTPStatusError: - # Log but don't fail - the key may already be expired/deleted - pass # noqa: S110 - intentional silence on revoke errors - except httpx.RequestError: - # Log but don't fail - revocation is best-effort - pass # noqa: S110 - intentional silence on connection errors + except httpx.HTTPStatusError as e: + warn_revoke_failure("credential", credential.litellm_key_id, e) + except httpx.RequestError as e: + warn_revoke_failure("credential", credential.litellm_key_id, e) updated = credential.model_copy( update={ @@ -261,11 +258,7 @@ def render_env_shell(self, env_vars: dict[str, str]) -> str: Returns: Shell commands string. """ - lines = [] - for key, value in env_vars.items(): - escaped = value.replace("'", "'\\''") - lines.append(f"export {key}='{escaped}'") - return "\n".join(lines) + return render_env_shell(env_vars) def render_env_dotenv(self, env_vars: dict[str, str]) -> str: """Render environment variables as dotenv file content. @@ -276,81 +269,4 @@ def render_env_dotenv(self, env_vars: dict[str, str]) -> str: Returns: Dotenv file content string. """ - lines = [] - for key, value in env_vars.items(): - if " " in value or "#" in value or "'" in value or '"' in value: - escaped = value.replace("\\", "\\\\").replace('"', '\\"') - lines.append(f'{key}="{escaped}"') - else: - lines.append(f"{key}={value}") - return "\n".join(lines) - - # ------------------------------------------------------------------ - # Backward-compatible convenience wrappers (deprecated stubs) - # ------------------------------------------------------------------ - - def _legacy_render_env_snippet( - self, - credential: str, - proxy_base_url: str, - model: str, - harness_profile: str, - ) -> dict[str, str]: - """Legacy wrapper accepting a plain credential string.""" - return { - "OPENAI_API_BASE": proxy_base_url, - "OPENAI_API_KEY": credential, - "OPENAI_MODEL": model, - } - - def _legacy_render_shell_snippet( - self, - credential: str, - proxy_base_url: str, - model: str, - additional_vars: dict[str, str] | None = None, - ) -> str: - """Legacy wrapper for shell rendering with a plain credential string.""" - env_vars = self._legacy_render_env_snippet(credential, proxy_base_url, model, "") - if additional_vars: - env_vars.update(additional_vars) - return self.render_env_shell(env_vars) - - def _legacy_render_dotenv_snippet( - self, - credential: str, - proxy_base_url: str, - model: str, - additional_vars: dict[str, str] | None = None, - ) -> str: - """Legacy wrapper for dotenv rendering with a plain credential string.""" - env_vars = self._legacy_render_env_snippet(credential, proxy_base_url, model, "") - if additional_vars: - env_vars.update(additional_vars) - return self.render_env_dotenv(env_vars) - - def _legacy_validate_credential(self, credential: str) -> dict | None: - """Legacy placeholder validation for stub credentials.""" - if credential.startswith("sk-benchmark-"): - parts = credential.split("-") - if len(parts) >= 3: - return { - "session_id": parts[2] if len(parts) > 2 else None, - "experiment_prefix": parts[3] if len(parts) > 3 else None, - "valid": True, - } - return None - - async def _legacy_issue_credential( - self, - session_id: UUID, - experiment_id: str, - variant_id: str, - harness_profile: str, - ) -> str: - """Legacy placeholder that returns a stub credential string.""" - return f"sk-benchmark-{session_id}-{experiment_id[:8]}" - - async def _legacy_revoke_credential(self, credential: str) -> bool: - """Legacy placeholder revocation.""" - return True + return render_env_dotenv(env_vars) diff --git a/src/benchmark_core/services/proxy_key_service.py b/src/benchmark_core/services/proxy_key_service.py index 0516177..ca0e67a 100644 --- a/src/benchmark_core/services/proxy_key_service.py +++ b/src/benchmark_core/services/proxy_key_service.py @@ -16,6 +16,12 @@ from benchmark_core.db.models import ProxyKey as ProxyKeyORM from benchmark_core.models import ProxyKey, ProxyKeyStatus from benchmark_core.repositories.proxy_key_repository import SQLProxyKeyRepository +from benchmark_core.services.common import ( + render_env_dotenv, + render_env_shell, + validate_litellm_url, + warn_revoke_failure, +) class ProxyKeyServiceError(Exception): @@ -57,16 +63,9 @@ def __init__( ValueError: If litellm_base_url doesn't use HTTPS in production. """ self._repository = repository - self.litellm_base_url = litellm_base_url.rstrip("/") + self.litellm_base_url = validate_litellm_url(litellm_base_url, enforce_https) self.master_key = master_key - if enforce_https and not self.litellm_base_url.startswith( - ("https://", "http://localhost", "http://127.0.0.1") - ): - raise ValueError( - f"LiteLLM URL must use HTTPS in production environments. Got: {litellm_base_url}" - ) - def _http_headers(self) -> dict[str, str]: """Build HTTP headers for LiteLLM API calls.""" if self.master_key is None: @@ -338,10 +337,10 @@ async def revoke_key(self, proxy_key_id: UUID) -> ProxyKey | None: json={"key_ids": [orm.litellm_key_id]}, ) response.raise_for_status() - except httpx.HTTPStatusError: - pass # Best-effort: key may already be expired/deleted - except httpx.RequestError: - pass # Best-effort: proxy may be temporarily unavailable + except httpx.HTTPStatusError as e: + warn_revoke_failure("proxy_key", orm.litellm_key_id, e) + except httpx.RequestError as e: + warn_revoke_failure("proxy_key", orm.litellm_key_id, e) # Mark local metadata as revoked revoked_orm = await self._repository.revoke(proxy_key_id) @@ -384,12 +383,7 @@ def render_env_shell(self, env_vars: dict[str, str]) -> str: Returns: Shell commands string. """ - lines = [] - for key in sorted(env_vars.keys()): - value = env_vars[key] - escaped = value.replace("'", "'\\''") - lines.append(f"export {key}='{escaped}'") - return "\n".join(lines) + return render_env_shell(env_vars) def render_env_dotenv(self, env_vars: dict[str, str]) -> str: """Render environment variables as dotenv file content. @@ -400,15 +394,7 @@ def render_env_dotenv(self, env_vars: dict[str, str]) -> str: Returns: Dotenv file content string. """ - lines = [] - for key in sorted(env_vars.keys()): - value = env_vars[key] - if " " in value or "#" in value or "'" in value or '"' in value: - escaped = value.replace("\\", "\\\\").replace('"', '\\"') - lines.append(f'{key}="{escaped}"') - else: - lines.append(f"{key}={value}") - return "\n".join(lines) + return render_env_dotenv(env_vars) def render_harness_env( self, diff --git a/src/benchmark_core/services_abc.py b/src/benchmark_core/services_abc.py index d032bbe..31a5997 100644 --- a/src/benchmark_core/services_abc.py +++ b/src/benchmark_core/services_abc.py @@ -1,19 +1,20 @@ """Abstract service definitions for the benchmark core. -Note: The concrete implementations have been moved to the services/ package. +Note: Concrete implementations have been moved to the services/ package. This file is kept for backward compatibility and will be deprecated. """ -from datetime import UTC, datetime, timedelta -from typing import Any +from datetime import UTC, datetime from uuid import UUID -import httpx -from pydantic import SecretStr - from benchmark_core.models import ProxyCredential, Session from benchmark_core.repositories import ProxyCredentialRepository, SessionRepository +# Re-export CredentialService so existing callers don't break. +from benchmark_core.services.credential_service import ( # noqa: F401 + CredentialService, +) + class SessionService: """Service for managing benchmark session lifecycle. @@ -165,272 +166,3 @@ async def finalize_session(self, session_id: UUID) -> Session | None: updated = session.model_copy(update={"ended_at": datetime.now(UTC), "status": "completed"}) return await self._session_repo.update(updated) - - -class CredentialService: - """Service for issuing and managing session-scoped proxy credentials. - - Implements the key aliasing convention and metadata tagging strategy - for correlating LiteLLM traffic back to benchmark sessions. - """ - - def __init__( - self, - litellm_base_url: str = "http://localhost:4000", - master_key: str | None = None, - enforce_https: bool = True, - ) -> None: - """Initialize the credential service. - - Args: - litellm_base_url: Base URL for LiteLLM proxy API - master_key: LiteLLM master key for API authentication - enforce_https: Whether to enforce HTTPS in production (default True) - - Raises: - ValueError: If litellm_base_url doesn't use HTTPS in production - """ - self.litellm_base_url = litellm_base_url.rstrip("/") - self.master_key = master_key - - # Validate HTTPS in production environments - if enforce_https and not self.litellm_base_url.startswith( - ("https://", "http://localhost", "http://127.0.0.1") - ): - raise ValueError( - f"LiteLLM URL must use HTTPS in production environments. Got: {litellm_base_url}" - ) - - def _generate_key_alias( - self, - session_id: UUID, - experiment_id: str, - variant_id: str, - ) -> str: - """Generate a stable key alias for correlation. - - Format: session-{session_id[:8]}-{experiment_id[:8]}-{variant_id[:8]} - - This alias is: - - Human-readable for operators - - Unique per session - - Joinable back to session via metadata - """ - session_short = str(session_id)[:8] - exp_short = experiment_id[:8] if len(experiment_id) >= 8 else experiment_id - var_short = variant_id[:8] if len(variant_id) >= 8 else variant_id - return f"session-{session_short}-{exp_short}-{var_short}" - - def _build_metadata_tags( - self, - session_id: UUID, - experiment_id: str, - variant_id: str, - harness_profile: str, - ) -> dict[str, str]: - """Build metadata tags for LiteLLM credential. - - These tags enable joining LiteLLM request logs back to - the benchmark session and its dimensions. - """ - return { - "benchmark_session_id": str(session_id), - "benchmark_experiment_id": experiment_id, - "benchmark_variant_id": variant_id, - "benchmark_harness_profile": harness_profile, - "benchmark_source": "opensymphony", - } - - async def issue_credential( - self, - session_id: UUID, - experiment_id: str, - variant_id: str, - harness_profile: str, - ttl_hours: int = 24, - ) -> ProxyCredential: - """Issue a session-scoped proxy credential via LiteLLM API. - - Creates a LiteLLM virtual key with: - - Unique key alias for correlation - - Metadata tags linking to session - - Configurable TTL (default 24 hours) - - Args: - session_id: Benchmark session UUID - experiment_id: Experiment identifier - variant_id: Variant identifier - harness_profile: Harness profile name - ttl_hours: Credential time-to-live in hours - - Returns: - ProxyCredential with metadata and API key - - Raises: - RuntimeError: If LiteLLM API call fails or master_key not configured - """ - if self.master_key is None: - raise RuntimeError( - "Cannot issue credential: LiteLLM master_key not configured. " - "Initialize CredentialService with master_key parameter." - ) - - # Generate stable alias for correlation - key_alias = self._generate_key_alias(session_id, experiment_id, variant_id) - - # Build metadata tags for LiteLLM correlation - metadata_tags = self._build_metadata_tags( - session_id, experiment_id, variant_id, harness_profile - ) - - # Calculate expiration - expires_at = datetime.now(UTC) + timedelta(hours=ttl_hours) - - # Call LiteLLM API to create the key - async with httpx.AsyncClient(timeout=30.0) as client: - try: - response = await client.post( - f"{self.litellm_base_url}/key/generate", - headers={"Authorization": f"Bearer {self.master_key}"}, - json={ - "key_alias": key_alias, - "metadata": metadata_tags, - "expires": expires_at.isoformat(), - }, - ) - response.raise_for_status() - data: dict[str, Any] = response.json() - except httpx.HTTPStatusError as e: - raise RuntimeError( - f"LiteLLM API error creating credential: {e.response.status_code} - " - f"{e.response.text}" - ) from e - except httpx.RequestError as e: - raise RuntimeError( - f"Failed to connect to LiteLLM at {self.litellm_base_url}: {e}" - ) from e - - # Extract key and key_id from LiteLLM response - api_key = data.get("key") - litellm_key_id = data.get("key_id") - - if not api_key: - raise RuntimeError(f"LiteLLM API response missing 'key' field: {data}") - - # Create credential domain model - credential = ProxyCredential( - session_id=session_id, - key_alias=key_alias, - api_key=SecretStr(api_key), - experiment_id=experiment_id, - variant_id=variant_id, - harness_profile=harness_profile, - metadata_tags=metadata_tags, - litellm_key_id=litellm_key_id, - expires_at=expires_at, - is_active=True, - ) - - return credential - - async def revoke_credential(self, credential: ProxyCredential) -> ProxyCredential: - """Revoke an active credential via LiteLLM API. - - Marks the credential as revoked and records revocation time. - Also calls LiteLLM API to delete the key if litellm_key_id is present. - - Args: - credential: Credential to revoke - - Returns: - Updated credential with is_active=False - - Note: - LiteLLM API errors are intentionally silenced since revocation is - best-effort - the key may already be expired or the proxy may be - temporarily unavailable. - """ - # Call LiteLLM API to delete the key if we have the key_id - if credential.litellm_key_id and self.master_key: - async with httpx.AsyncClient(timeout=30.0) as client: - try: - response = await client.post( - f"{self.litellm_base_url}/key/delete", - headers={"Authorization": f"Bearer {self.master_key}"}, - json={"key_ids": [credential.litellm_key_id]}, - ) - response.raise_for_status() - except httpx.HTTPStatusError: - # Log but don't fail - the key may already be expired/deleted - pass # noqa: S110 - intentional silence on revoke errors - except httpx.RequestError: - # Log but don't fail - revocation is best-effort - pass # noqa: S110 - intentional silence on connection errors - - updated = credential.model_copy( - update={ - "is_active": False, - "revoked_at": datetime.now(UTC), - # Clear the secret from memory - "api_key": SecretStr("REDACTED"), - } - ) - return updated - - def render_env_snippet( - self, - credential: ProxyCredential, - proxy_base_url: str, - model: str, - ) -> dict[str, str]: - """Render environment variable snippet for a harness. - - Args: - credential: Proxy credential with API key - proxy_base_url: LiteLLM proxy base URL - model: Model identifier/alias - - Returns: - Dictionary of environment variables - """ - return { - "OPENAI_API_BASE": proxy_base_url, - "OPENAI_API_KEY": credential.api_key.get_secret_value(), - "OPENAI_MODEL": model, - "LITELLM_SESSION_ALIAS": credential.key_alias, - } - - def render_env_shell(self, env_vars: dict[str, str]) -> str: - """Render environment variables as shell export commands. - - Args: - env_vars: Dictionary of environment variables - - Returns: - Shell commands for setting environment - """ - lines = [] - for key, value in env_vars.items(): - # Quote the value for shell safety - escaped = value.replace("'", "'\\''") - lines.append(f"export {key}='{escaped}'") - return "\n".join(lines) - - def render_env_dotenv(self, env_vars: dict[str, str]) -> str: - """Render environment variables as dotenv format. - - Args: - env_vars: Dictionary of environment variables - - Returns: - Dotenv file content - """ - lines = [] - for key, value in env_vars.items(): - # Escape special characters for dotenv - if " " in value or "#" in value or "'" in value or '"' in value: - escaped = value.replace("\\", "\\\\").replace('"', '\\"') - lines.append(f'{key}="{escaped}"') - else: - lines.append(f"{key}={value}") - return "\n".join(lines) From d5682e46ff42a1e1611c101c1ff9032be2f3ac6f Mon Sep 17 00:00:00 2001 From: Leonardo Gonzalez Date: Tue, 21 Apr 2026 22:27:29 -0500 Subject: [PATCH 4/7] fix(proxy-key-service): explicit None checks for budget_amount, allowed_models, budget_duration Replaces or-operator fallback with is not None checks so falsy-but-valid values (0, [], empty string) are not silently overridden by usage policy. AI review: inline comments 3121395728, 3121395737. Co-authored-by: openhands --- .../services/proxy_key_service.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/src/benchmark_core/services/proxy_key_service.py b/src/benchmark_core/services/proxy_key_service.py index ca0e67a..b642ed9 100644 --- a/src/benchmark_core/services/proxy_key_service.py +++ b/src/benchmark_core/services/proxy_key_service.py @@ -134,14 +134,20 @@ async def create_key( effective_owner = owner or (usage_policy.owner if usage_policy else None) effective_team = team or (usage_policy.team if usage_policy else None) effective_customer = customer or (usage_policy.customer if usage_policy else None) - effective_models = allowed_models or ( - list(usage_policy.allowed_models) if usage_policy else [] + effective_models = ( + list(allowed_models) + if allowed_models is not None + else (list(usage_policy.allowed_models) if usage_policy else []) ) - effective_budget_duration = budget_duration or ( - usage_policy.budget_duration if usage_policy else None + effective_budget_duration = ( + budget_duration + if budget_duration is not None + else (usage_policy.budget_duration if usage_policy else None) ) - effective_budget_amount = budget_amount or ( - usage_policy.budget_amount if usage_policy else None + effective_budget_amount = ( + budget_amount + if budget_amount is not None + else (usage_policy.budget_amount if usage_policy else None) ) effective_ttl = ( usage_policy.ttl_seconds // 3600 From 65db60936119e4243809cec0918b4c9a9062973e Mon Sep 17 00:00:00 2001 From: Leonardo Gonzalez Date: Tue, 21 Apr 2026 22:38:00 -0500 Subject: [PATCH 5/7] fix(cli): accept alias in revoke, add budget/metadata to create - revoke command now accepts UUID or alias (same as info) - create command adds --budget-amount, --budget-duration, --meta options - update test to match new alias-based invalid-input behavior AI review feedback: inline comments 3121395742, 3121395744. Co-authored-by: openhands --- src/cli/commands/key.py | 55 +++++++++++++++++++++++------ tests/unit/test_cli_key_commands.py | 6 ++-- 2 files changed, 47 insertions(+), 14 deletions(-) diff --git a/src/cli/commands/key.py b/src/cli/commands/key.py index 881740a..b2df61d 100644 --- a/src/cli/commands/key.py +++ b/src/cli/commands/key.py @@ -52,6 +52,15 @@ def create( models: list[str] | None = typer.Option( None, "--model", "-m", help="Allowed model alias (repeatable)" ), + budget_amount: float | None = typer.Option( + None, "--budget-amount", "-b", help="Budget limit (currency units)" + ), + budget_duration: str | None = typer.Option( + None, "--budget-duration", "-d", help="Budget interval (e.g. '1d', '30d')" + ), + metadata_pairs: list[str] | None = typer.Option( + None, "--meta", help="Metadata tag as 'key=value' (repeatable)" + ), ttl_hours: int = typer.Option(168, "--ttl-hours", help="Key TTL in hours (default 7 days)"), show_env: bool = typer.Option( False, "--show-env", "-e", help="Print environment snippet after creation" @@ -70,6 +79,17 @@ def create( """ console.print("[bold blue]Creating sessionless proxy key...[/bold blue]") + # Parse metadata pairs + metadata: dict[str, str] | None = None + if metadata_pairs: + metadata = {} + for pair in metadata_pairs: + if "=" not in pair: + console.print(f"[red]Invalid metadata format: {pair} (expected key=value)[/red]") + raise typer.Exit(1) + key, value = pair.split("=", 1) + metadata[key] = value + with get_db_session() as db: service = _get_service(db) @@ -82,6 +102,9 @@ def create( customer=customer, purpose=purpose, allowed_models=models or None, + budget_amount=budget_amount, + budget_duration=budget_duration, + metadata=metadata, ttl_hours=ttl_hours, ) ) @@ -256,26 +279,36 @@ def info( @app.command() def revoke( - key_id: str = typer.Argument(..., help="Proxy key ID (UUID)"), + key_id: str = typer.Argument(..., help="Proxy key ID (UUID) or alias"), ) -> None: """Revoke a sessionless proxy key. Marks the local metadata as inactive and attempts LiteLLM deletion. The key secret cannot be recovered after revocation. """ - try: - pk_uuid = UUID(key_id) - except ValueError as e: - console.print(f"[red]Invalid UUID: {key_id}[/red]") - raise typer.Exit(1) from e - with get_db_session() as db: service = _get_service(db) + + # Try as UUID first, then as alias try: - proxy_key = _run(service.revoke_key(pk_uuid)) - except ProxyKeyServiceError as e: - console.print(f"[red]Error: {e}[/red]") - raise typer.Exit(1) from e + pk_uuid = UUID(key_id) + try: + proxy_key = _run(service.revoke_key(pk_uuid)) + except ProxyKeyServiceError as e: + console.print(f"[red]Error: {e}[/red]") + raise typer.Exit(1) from e + except ValueError: + try: + proxy_key = _run(service.get_key_by_alias(key_id)) + except ProxyKeyServiceError as e: + console.print(f"[red]Error: {e}[/red]") + raise typer.Exit(1) from e + if proxy_key is not None: + try: + proxy_key = _run(service.revoke_key(proxy_key.proxy_key_id)) + except ProxyKeyServiceError as e: + console.print(f"[red]Error: {e}[/red]") + raise typer.Exit(1) from e if proxy_key is None: console.print(f"[red]Proxy key not found: {key_id}[/red]") diff --git a/tests/unit/test_cli_key_commands.py b/tests/unit/test_cli_key_commands.py index 8c7a178..528fad3 100644 --- a/tests/unit/test_cli_key_commands.py +++ b/tests/unit/test_cli_key_commands.py @@ -314,11 +314,11 @@ def test_revoke_not_found(self, runner, mock_env_db_url, monkeypatch): assert result.exit_code == 1 assert "not found" in result.output - def test_revoke_invalid_uuid(self, runner, mock_env_db_url, monkeypatch): - """Revoke with invalid UUID shows error.""" + def test_revoke_invalid_alias(self, runner, mock_env_db_url, monkeypatch): + """Revoke with non-existent alias shows not-found error.""" result = runner.invoke(app, ["key", "revoke", "not-a-uuid"]) assert result.exit_code == 1 - assert "Invalid UUID" in result.output + assert "not found" in result.output class TestKeyCLIImports: From 2f4a8305d8daa29c3fb7e4dc4299b59d2d38079a Mon Sep 17 00:00:00 2001 From: Leonardo Gonzalez Date: Wed, 22 Apr 2026 02:50:50 -0500 Subject: [PATCH 6/7] fix(proxy-key-service): address PR #45 feedback round 2 - Fix TTL truncation: use ceil(usage_policy.ttl_seconds / 3600) instead of integer division so operators get at least the requested TTL. - Remove unused render_harness_env() that misassigned LITELLM_KEY_ALIAS to harness_profile_name instead of key alias and was untested dead API surface. - Remove unused harness_profile parameter from render_env_snippet(). - Fix or-operator silent swallow for owner, team, customer fields. - Add direct test coverage for benchmark_core/services/common.py utilities. Addresses AI review feedback on PR #45 (review 4151933140). --- .../services/proxy_key_service.py | 44 +++------- tests/unit/test_common_services.py | 80 +++++++++++++++++++ 2 files changed, 89 insertions(+), 35 deletions(-) create mode 100644 tests/unit/test_common_services.py diff --git a/src/benchmark_core/services/proxy_key_service.py b/src/benchmark_core/services/proxy_key_service.py index b642ed9..2530e97 100644 --- a/src/benchmark_core/services/proxy_key_service.py +++ b/src/benchmark_core/services/proxy_key_service.py @@ -6,6 +6,7 @@ """ from datetime import UTC, datetime, timedelta +from math import ceil from typing import Any from uuid import UUID, uuid4 @@ -131,9 +132,13 @@ async def create_key( effective_alias = key_alias or self._build_key_alias( usage_policy.name if usage_policy else None ) - effective_owner = owner or (usage_policy.owner if usage_policy else None) - effective_team = team or (usage_policy.team if usage_policy else None) - effective_customer = customer or (usage_policy.customer if usage_policy else None) + effective_owner = ( + owner if owner is not None else (usage_policy.owner if usage_policy else None) + ) + effective_team = team if team is not None else (usage_policy.team if usage_policy else None) + effective_customer = ( + customer if customer is not None else (usage_policy.customer if usage_policy else None) + ) effective_models = ( list(allowed_models) if allowed_models is not None @@ -150,7 +155,7 @@ async def create_key( else (usage_policy.budget_amount if usage_policy else None) ) effective_ttl = ( - usage_policy.ttl_seconds // 3600 + ceil(usage_policy.ttl_seconds / 3600) if usage_policy and usage_policy.ttl_seconds else ttl_hours ) @@ -359,7 +364,6 @@ def render_env_snippet( api_key: SecretStr, proxy_base_url: str = "http://localhost:4000", model: str | None = None, - harness_profile: str = "openai-compatible", ) -> dict[str, str]: """Render generic OpenAI-compatible environment snippet. @@ -367,7 +371,6 @@ def render_env_snippet( api_key: The proxy key secret. proxy_base_url: LiteLLM proxy base URL. model: Optional default model alias. - harness_profile: Harness profile hint for rendering. Returns: Dictionary of environment variables. @@ -402,35 +405,6 @@ def render_env_dotenv(self, env_vars: dict[str, str]) -> str: """ return render_env_dotenv(env_vars) - def render_harness_env( - self, - api_key: SecretStr, - harness_profile_name: str, - proxy_base_url: str = "http://localhost:4000", - model: str | None = None, - ) -> dict[str, str]: - """Render harness-profile-based environment snippet. - - Args: - api_key: The proxy key secret. - harness_profile_name: Name of the harness profile. - proxy_base_url: LiteLLM proxy base URL. - model: Optional default model alias. - - Returns: - Dictionary of environment variables. - """ - # Generic OpenAI-compatible rendering - # Future: could load harness profile from config and adapt variable names - env_vars: dict[str, str] = { - "OPENAI_API_BASE": proxy_base_url, - "OPENAI_API_KEY": api_key.get_secret_value(), - } - if model: - env_vars["OPENAI_MODEL"] = model - env_vars["LITELLM_KEY_ALIAS"] = harness_profile_name - return env_vars - def _orm_to_model(self, orm: ProxyKeyORM) -> ProxyKey: """Convert an ORM ProxyKey to a domain model. diff --git a/tests/unit/test_common_services.py b/tests/unit/test_common_services.py new file mode 100644 index 0000000..0594892 --- /dev/null +++ b/tests/unit/test_common_services.py @@ -0,0 +1,80 @@ +"""Direct tests for benchmark_core/services/common.py shared utilities.""" + +import warnings + +import pytest + +from benchmark_core.services.common import ( + render_env_dotenv, + render_env_shell, + validate_litellm_url, + warn_revoke_failure, +) + + +class TestValidateLitellmUrl: + def test_allows_https(self) -> None: + assert validate_litellm_url("https://proxy.example.com") == "https://proxy.example.com" + + def test_allows_http_localhost(self) -> None: + assert validate_litellm_url("http://localhost:4000") == "http://localhost:4000" + + def test_allows_http_127_0_0_1(self) -> None: + assert validate_litellm_url("http://127.0.0.1:4000/") == "http://127.0.0.1:4000" + + def test_strips_trailing_slash(self) -> None: + assert validate_litellm_url("https://proxy.example.com/") == "https://proxy.example.com" + + def test_rejects_http_in_production(self) -> None: + with pytest.raises(ValueError, match="must use HTTPS"): + validate_litellm_url("http://proxy.example.com") + + def test_allows_http_when_enforce_https_false(self) -> None: + assert ( + validate_litellm_url("http://proxy.example.com", enforce_https=False) + == "http://proxy.example.com" + ) + + +class TestRenderEnvShell: + def test_basic(self) -> None: + result = render_env_shell({"B": "2", "A": "1"}) + assert result == "export A='1'\nexport B='2'" + + def test_escapes_single_quotes(self) -> None: + result = render_env_shell({"KEY": "it's fine"}) + assert result == "export KEY='it'\\''s fine'" + + def test_empty_dict(self) -> None: + assert render_env_shell({}) == "" + + +class TestRenderEnvDotenv: + def test_basic(self) -> None: + result = render_env_dotenv({"B": "2", "A": "1"}) + assert result == "A=1\nB=2" + + def test_escapes_quotes_and_spaces(self) -> None: + result = render_env_dotenv({"KEY": 'say "hello" now'}) + assert result == 'KEY="say \\"hello\\" now"' + + def test_empty_dict(self) -> None: + assert render_env_dotenv({}) == "" + + +class TestWarnRevokeFailure: + def test_emits_warning(self) -> None: + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + warn_revoke_failure("proxy_key", "key-123", RuntimeError("timeout")) + assert len(w) == 1 + assert "key_id=key-123" in str(w[0].message) + assert "timeout" in str(w[0].message) + assert "Local metadata has been marked revoked" in str(w[0].message) + + def test_works_with_none_key_id(self) -> None: + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + warn_revoke_failure("credential", None, ValueError("bad request")) + assert len(w) == 1 + assert "key_id=None" in str(w[0].message) From 7d8dbe246720660dfe48177579e83a69952aa844 Mon Sep 17 00:00:00 2001 From: Leonardo Gonzalez Date: Wed, 22 Apr 2026 02:53:12 -0500 Subject: [PATCH 7/7] test(proxy): verify TTL rounding with usage policy (ceil division) --- tests/unit/test_proxy_key_service.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/unit/test_proxy_key_service.py b/tests/unit/test_proxy_key_service.py index 5031dc0..746d775 100644 --- a/tests/unit/test_proxy_key_service.py +++ b/tests/unit/test_proxy_key_service.py @@ -172,6 +172,11 @@ async def test_create_key_with_usage_policy(self, service): assert proxy_key.allowed_models == ["kimi-k2-5"] assert proxy_key.budget_duration == "30d" assert proxy_key.budget_amount == 500.0 + # ttl_seconds=7200 rounds up to 2 hours (effective TTL applied) + assert proxy_key.expires_at is not None + assert proxy_key.created_at is not None + delta = proxy_key.expires_at - proxy_key.created_at + assert abs(delta.total_seconds() - 7200) < 1 # allow <1s timing jitter @pytest.mark.asyncio async def test_create_key_requires_master_key(self, repository):