diff --git a/src/youtube_extension/backend/services/api_cost_monitor.py b/src/youtube_extension/backend/services/api_cost_monitor.py index 3b875a232..bfbb2ba5c 100644 --- a/src/youtube_extension/backend/services/api_cost_monitor.py +++ b/src/youtube_extension/backend/services/api_cost_monitor.py @@ -8,12 +8,15 @@ """ import asyncio +import contextvars import json import logging import os +import random import re import threading import time +import uuid from collections import defaultdict, deque from collections.abc import Iterator from contextlib import contextmanager @@ -31,8 +34,11 @@ delete, func, inspect, + or_, text, ) +from sqlalchemy.dialects.postgresql import insert as postgresql_insert +from sqlalchemy.dialects.sqlite import insert as sqlite_insert from sqlalchemy.engine import URL, make_url from sqlalchemy.orm import Session, sessionmaker from sqlalchemy.pool import StaticPool @@ -57,6 +63,11 @@ # Configure logging logger = logging.getLogger(__name__) +# Keep the public webhook helper's one-argument signature for existing callers and +# tests while attaching an outbox event identifier to each delivery attempt. +_WEBHOOK_EVENT_ID: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar( + "api_cost_webhook_event_id", default=None +) _PRODUCTION_NAMES = {"staging", "prod", "production"} _RUNTIME_ROLE_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]{0,62}$") _REQUIRED_API_COST_COLUMNS = { @@ -566,6 +577,8 @@ class APICostMonitor: "claude-3-haiku-20240307": {"input": 0.00025, "output": 0.00125}, }, "google": { + # Standard paid-tier prices, normalized from per-million to per-1K USD. + "gemini-3.5-flash": {"input": 0.0015, "output": 0.009}, "gemini-3-pro": {"input": 0.000875, "output": 0.0035}, "gemini-3-flash": {"input": 0.000052, "output": 0.00021}, "gemini-1.5-pro": {"input": 0.00125, "output": 0.005}, @@ -603,6 +616,22 @@ def __init__( # Webhook notification settings self.webhook_url = os.getenv("API_COST_WEBHOOK_URL") + self.webhook_max_attempts = 5 + self.webhook_retry_base_seconds = max( + 0.0, float(os.getenv("API_COST_WEBHOOK_RETRY_BASE_SECONDS", "5")) + ) + self.webhook_retry_max_seconds = max( + self.webhook_retry_base_seconds, + float(os.getenv("API_COST_WEBHOOK_RETRY_MAX_SECONDS", "300")), + ) + self.webhook_poll_interval_seconds = max( + 0.01, float(os.getenv("API_COST_WEBHOOK_POLL_SECONDS", "1")) + ) + self.webhook_stale_timeout_seconds = max( + 1, int(os.getenv("API_COST_WEBHOOK_STALE_SECONDS", "30")) + ) + self._worker_task: Optional[asyncio.Task[None]] = None + self._worker_wake_event: Optional[asyncio.Event] = None # Rate limiters for different services self.rate_limiters = { @@ -723,7 +752,6 @@ def _create_database_engine(database_url: URL): if not database_url.database: kwargs["poolclass"] = StaticPool return create_engine(database_url, **kwargs) - pool_size = max(1, int(os.getenv("API_COST_DB_POOL_SIZE", "2"))) max_overflow = max(0, int(os.getenv("API_COST_DB_MAX_OVERFLOW", "0"))) pool_timeout = max(1, int(os.getenv("API_COST_DB_POOL_TIMEOUT", "10"))) @@ -762,8 +790,92 @@ def _init_database(self) -> None: self.engine, tables=[APIUsage.__table__, DailyBudget.__table__, WebhookOutbox.__table__], ) + self._upgrade_sqlite_outbox_schema() self._validate_database_schema() + def _upgrade_sqlite_outbox_schema(self) -> None: + """Add compatible SQLite outbox columns and indexes without data loss.""" + if self.engine is None or self.engine.dialect.name != "sqlite": + return + + with self.engine.connect() as connection: + connection.exec_driver_sql("BEGIN EXCLUSIVE") + try: + columns = { + row[1] + for row in connection.exec_driver_sql( + "PRAGMA table_info(webhook_outbox)" + ) + } + if not columns: + connection.commit() + return + + unique_indexes = [ + row[1] + for row in connection.exec_driver_sql( + "PRAGMA index_list(webhook_outbox)" + ) + if row[2] + ] + unique_index_columns = [] + for index_name in unique_indexes: + unique_index_columns.append( + [ + row[2] + for row in connection.exec_driver_sql( + f"PRAGMA index_info({index_name!r})" + ) + ] + ) + if ["utc_date", "alert_type"] not in unique_index_columns: + raise RuntimeError( + "Legacy API-cost SQLite schema is incompatible; " + "back up and recreate the local database" + ) + + for column_name, column_type in ( + ("next_attempt_at", "DATETIME"), + ("claimed_at", "DATETIME"), + ("claim_token", "VARCHAR(64)"), + ("last_recovered_at", "DATETIME"), + ("sent_at", "DATETIME"), + ): + if column_name not in columns: + connection.exec_driver_sql( + "ALTER TABLE webhook_outbox " + f"ADD COLUMN {column_name} {column_type}" + ) + + index_columns = [ + row[2] + for row in connection.exec_driver_sql( + "PRAGMA index_info(ix_webhook_outbox_due)" + ) + ] + expected_due_index_columns = [ + "status", + "next_attempt_at", + "retry_count", + "id", + ] + if index_columns and index_columns != expected_due_index_columns: + connection.exec_driver_sql( + "DROP INDEX IF EXISTS ix_webhook_outbox_due" + ) + connection.exec_driver_sql( + "CREATE INDEX IF NOT EXISTS ix_webhook_outbox_due " + "ON webhook_outbox (status, next_attempt_at, retry_count, id)" + ) + connection.exec_driver_sql( + "CREATE INDEX IF NOT EXISTS ix_webhook_outbox_stale_claims " + "ON webhook_outbox (status, claimed_at, id)" + ) + connection.commit() + except BaseException: + connection.rollback() + raise + def _validate_database_schema(self) -> None: if self.engine is None: raise RuntimeError("API-cost persistence is not configured") @@ -1143,8 +1255,13 @@ def calculate_cost( service_costs = self.COST_MODELS[service] if model not in service_costs: - # Use average cost for unknown models - model = list(service_costs.keys())[0] + if model == "default": + model = next(iter(service_costs)) + else: + raise ValueError( + f"Unknown pricing model for {service}: {model!r}; " + "refusing to apply an unrelated fallback price" + ) if service == "youtube": # YouTube uses quota units, not token pricing @@ -1200,28 +1317,33 @@ async def record_usage( self.session_costs[service] += cost self.session_requests[service] += 1 - stored = False + claimed_alerts: list[tuple[str, float]] = [] if self.Session is None: logger.warning( "API usage was not persisted because persistence is disabled" ) else: try: - await asyncio.to_thread(self._record_usage_sync, record) - stored = True + claimed_alerts = await asyncio.to_thread( + self._record_usage_sync, record + ) except Exception as exc: # The provider operation has already completed. Telemetry is # best effort and must never make that paid result retry/fail. logger.error("Failed to record API usage: %s", exc) - if stored: - await self._check_budget_alerts() + # This only wakes the explicitly managed worker; network I/O remains + # outside the accounting path. + for alert_type, current_cost in claimed_alerts: + await self._send_budget_alert(current_cost, alert_type) logger.debug("API usage: %s - $%.4f (%s tokens)", service, cost, tokens_used) return record - def _record_usage_sync(self, record: APIUsageRecord) -> None: - """Persist one usage record on a worker thread.""" + def _record_usage_sync( + self, record: APIUsageRecord + ) -> list[tuple[str, float]]: + """Persist usage and any newly crossed alert in one transaction.""" with self._session_scope(commit=True) as session: session.add( APIUsage( @@ -1237,6 +1359,83 @@ def _record_usage_sync(self, record: APIUsageRecord) -> None: error_message=record.error_message, ) ) + session.flush() + return self._stage_budget_alerts(session, record.timestamp) + + def _stage_budget_alerts( + self, session: Session, timestamp: datetime + ) -> list[tuple[str, float]]: + """Aggregate the UTC day and enqueue crossed alerts transactionally.""" + utc_date = timestamp.astimezone(timezone.utc).date().isoformat() + start_at, end_at = self._utc_day_bounds(utc_date) + insert_factory = postgresql_insert if self._is_postgres else sqlite_insert + + session.execute( + insert_factory(DailyBudget) + .values( + date=utc_date, + total_cost=0.0, + alert_sent=False, + budget_exceeded=False, + ) + .on_conflict_do_nothing(index_elements=["date"]) + ) + budget = ( + session.query(DailyBudget) + .filter_by(date=utc_date) + .with_for_update() + .one() + ) + total = ( + session.query(func.sum(APIUsage.cost)) + .filter( + APIUsage.timestamp >= start_at, + APIUsage.timestamp < end_at, + ) + .scalar() + ) + current_cost = float(total) if total is not None else 0.0 + budget.total_cost = current_cost + + claimed: list[tuple[str, float]] = [] + alert_specs = ( + ("threshold", self.alert_threshold, "alert_sent"), + ("exceeded", self.daily_budget, "budget_exceeded"), + ) + for alert_type, limit, flag_name in alert_specs: + if current_cost < limit or getattr(budget, flag_name): + continue + + if alert_type == "threshold": + payload = ( + f"🚨 API Budget Alert: ${current_cost:.2f} " + f"(Alert threshold: ${self.alert_threshold})" + ) + else: + payload = ( + f"🚨 API Budget Alert: ${current_cost:.2f} " + f"EXCEEDED daily budget of ${self.daily_budget}" + ) + + inserted = session.execute( + insert_factory(WebhookOutbox) + .values( + utc_date=utc_date, + alert_type=alert_type, + status="pending", + retry_count=0, + current_cost=current_cost, + payload=payload, + ) + .on_conflict_do_nothing( + index_elements=["utc_date", "alert_type"] + ) + ) + setattr(budget, flag_name, True) + if inserted.rowcount == 1: + claimed.append((alert_type, current_cost)) + + return claimed async def _check_budget_alerts(self) -> None: """Check and enqueue budget alerts if thresholds are exceeded.""" @@ -1329,112 +1528,351 @@ def _claim_alert( return False def _trigger_delivery(self): - """Leave delivery to the dedicated worker's polling loop.""" - logger.debug("API-cost alert queued for the dedicated worker") + """Wake the explicitly managed worker without spawning per-alert tasks.""" + if self._worker_wake_event is not None: + self._worker_wake_event.set() + + async def start(self) -> asyncio.Task[None]: + """Start the monitor's single managed outbox worker.""" + if self._worker_task is not None and not self._worker_task.done(): + return self._worker_task + + self._worker_wake_event = asyncio.Event() + self._worker_task = asyncio.create_task( + self._outbox_worker(), name="api-cost-webhook-outbox" + ) + return self._worker_task + + async def close(self) -> None: + """Stop the managed worker and wait for any claim cleanup to finish.""" + task = self._worker_task + if task is None: + return + + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + finally: + if self._worker_task is task: + self._worker_task = None + self._worker_wake_event = None + + async def _outbox_worker(self) -> None: + """Continuously deliver due outbox items until explicitly closed.""" + while True: + wake_event = self._worker_wake_event + if wake_event is None: + return + wake_event.clear() + try: + await self.process_outbox() + except asyncio.CancelledError: + raise + except Exception: + logger.exception("Unhandled error in API-cost webhook outbox worker") + + try: + await asyncio.wait_for( + wake_event.wait(), timeout=self.webhook_poll_interval_seconds + ) + except asyncio.TimeoutError: + pass + + def _retry_at(self, attempt: int, now: datetime) -> datetime: + """Return a bounded exponential equal-jitter retry timestamp.""" + exponential_cap = min( + self.webhook_retry_max_seconds, + self.webhook_retry_base_seconds * (2 ** max(0, attempt - 1)), + ) + half_cap = exponential_cap / 2 + delay = half_cap + random.uniform(0, half_cap) + return now + timedelta(seconds=delay) + + def _retry_state( + self, attempt: int, now: datetime, error_message: str + ) -> tuple[Optional[datetime], str]: + """Return persisted scheduling and error state for a failed attempt.""" + if attempt >= self.webhook_max_attempts: + return ( + None, + f"Retry exhausted after {self.webhook_max_attempts} attempts: " + f"{error_message}", + ) + return self._retry_at(attempt, now), error_message - async def recover_stale_deliveries(self, stale_timeout_seconds: int = 30): - """Recover items left processing after a crash or cancellation.""" + async def recover_stale_deliveries( + self, stale_timeout_seconds: Optional[int] = None + ) -> None: + """Recover abandoned claims without blocking the application event loop.""" await asyncio.to_thread( self._recover_stale_deliveries_sync, stale_timeout_seconds ) - def _recover_stale_deliveries_sync(self, stale_timeout_seconds: int) -> None: + def _recover_stale_deliveries_sync( + self, stale_timeout_seconds: Optional[int] = None + ) -> None: + """Recover processing claims in a worker thread.""" + if stale_timeout_seconds is None: + stale_timeout_seconds = self.webhook_stale_timeout_seconds + try: - cutoff = datetime.now(timezone.utc) - timedelta( - seconds=stale_timeout_seconds - ) with self._session_scope(commit=True) as session: + now = datetime.now(timezone.utc) + cutoff = now - timedelta(seconds=stale_timeout_seconds) stale_items = ( session.query(WebhookOutbox) .filter( WebhookOutbox.status == "processing", - WebhookOutbox.last_attempt < cutoff, + or_( + WebhookOutbox.claimed_at.is_(None), + WebhookOutbox.claimed_at < cutoff, + ), ) .all() ) + for item in stale_items: - item.status = "failed" - item.last_recovered_at = datetime.now(timezone.utc) - item.error_message = ( - "Recovery: Stale/Crashed delivery task recovered" + next_attempt_at, recovery_error = self._retry_state( + max(1, item.retry_count), + now, + "Recovery: Stale/Crashed delivery task recovered", ) - logger.info( - "Recovered stale webhook delivery %s for %s (%s)", - item.id, - item.utc_date, - item.alert_type, + filters = [ + WebhookOutbox.id == item.id, + WebhookOutbox.status == "processing", + ] + if item.claimed_at is None: + filters.append(WebhookOutbox.claimed_at.is_(None)) + else: + filters.append(WebhookOutbox.claimed_at == item.claimed_at) + if item.claim_token is None: + filters.append(WebhookOutbox.claim_token.is_(None)) + else: + filters.append(WebhookOutbox.claim_token == item.claim_token) + + recovered = ( + session.query(WebhookOutbox) + .filter(*filters) + .update( + { + WebhookOutbox.status: "failed", + WebhookOutbox.next_attempt_at: next_attempt_at, + WebhookOutbox.error_message: recovery_error, + WebhookOutbox.last_recovered_at: now, + WebhookOutbox.claimed_at: None, + WebhookOutbox.claim_token: None, + }, + synchronize_session=False, + ) ) - except Exception as exc: - logger.error("Error during stale webhook delivery recovery: %s", exc) + if recovered: + logger.info( + "Recovered stale webhook delivery %s for %s (%s)", + item.id, + item.utc_date, + item.alert_type, + ) + except Exception as e: + logger.error("Error during stale webhook delivery recovery: %s", e) + + def _try_claim_outbox_item( + self, + item_id: int, + claim_time: datetime, + respect_schedule: bool = True, + ) -> Optional[dict[str, Any]]: + """Claim one due item with a single compare-and-swap UPDATE.""" + try: + with self._session_scope(commit=True) as session: + filters = [ + WebhookOutbox.id == item_id, + WebhookOutbox.status.in_(["pending", "failed"]), + WebhookOutbox.retry_count < self.webhook_max_attempts, + ] + if respect_schedule: + filters.append( + or_( + WebhookOutbox.next_attempt_at.is_(None), + WebhookOutbox.next_attempt_at <= claim_time, + ) + ) + + claim_token = uuid.uuid4().hex + claimed = ( + session.query(WebhookOutbox) + .filter(*filters) + .update( + { + WebhookOutbox.status: "processing", + WebhookOutbox.retry_count: WebhookOutbox.retry_count + 1, + WebhookOutbox.last_attempt: claim_time, + WebhookOutbox.claimed_at: claim_time, + WebhookOutbox.claim_token: claim_token, + WebhookOutbox.next_attempt_at: None, + }, + synchronize_session=False, + ) + ) + if claimed != 1: + return None - async def process_outbox(self, max_items: Optional[int] = None): - """Process a bounded set of pending or failed outbox deliveries.""" + item = session.query(WebhookOutbox).filter_by(id=item_id).one() + return { + "id": item.id, + "payload": item.payload, + "utc_date": item.utc_date, + "alert_type": item.alert_type, + "retry_count": item.retry_count, + "last_attempt": item.last_attempt, + "claim_token": item.claim_token, + } + except Exception as e: + logger.debug("Could not claim webhook outbox item %s: %s", item_id, e) + return None + + def _complete_outbox_claim( + self, + claim: dict[str, Any], + *, + success: bool, + error_message: Optional[str] = None, + ) -> bool: + """Conditionally complete exactly the represented delivery attempt.""" + try: + with self._session_scope(commit=True) as session: + values: dict[Any, Any] + if success: + values = { + WebhookOutbox.status: "sent", + WebhookOutbox.next_attempt_at: None, + WebhookOutbox.error_message: None, + WebhookOutbox.sent_at: datetime.now(timezone.utc), + WebhookOutbox.claimed_at: None, + WebhookOutbox.claim_token: None, + } + else: + next_attempt_at, persisted_error = self._retry_state( + claim["retry_count"], + datetime.now(timezone.utc), + error_message or "Delivery failed", + ) + values = { + WebhookOutbox.status: "failed", + WebhookOutbox.next_attempt_at: next_attempt_at, + WebhookOutbox.error_message: persisted_error, + WebhookOutbox.claimed_at: None, + WebhookOutbox.claim_token: None, + } + + completed = ( + session.query(WebhookOutbox) + .filter( + WebhookOutbox.id == claim["id"], + WebhookOutbox.status == "processing", + WebhookOutbox.retry_count == claim["retry_count"], + WebhookOutbox.last_attempt == claim["last_attempt"], + WebhookOutbox.claim_token == claim["claim_token"], + ) + .update(values, synchronize_session=False) + ) + if completed != 1: + return False + return True + except Exception as e: + logger.error("Error completing outbox item %s: %s", claim["id"], e) + return False + + def _select_outbox_item_ids( + self, *, now: datetime, force: bool, max_items: Optional[int] + ) -> list[int]: + """Return due outbox IDs using a short worker-thread transaction.""" + try: + with self._session_scope() as session: + filters = [ + WebhookOutbox.status.in_(["pending", "failed"]), + WebhookOutbox.retry_count < self.webhook_max_attempts, + ] + if not force: + filters.append( + or_( + WebhookOutbox.next_attempt_at.is_(None), + WebhookOutbox.next_attempt_at <= now, + ) + ) + query = ( + session.query(WebhookOutbox.id) + .filter(*filters) + .order_by(WebhookOutbox.next_attempt_at, WebhookOutbox.id) + ) + if max_items is not None: + query = query.limit(max(0, max_items)) + return [row[0] for row in query.all()] + except Exception as e: + logger.error("Error selecting webhook outbox items: %s", e) + return [] + + async def process_outbox( + self, max_items: Optional[int] = None, *, force: bool = False + ) -> int: + """Deliver eligible items, honoring persisted due times by default. + + ``force=True`` is an explicit operational/test escape hatch that ignores + only the due timestamp; compare-and-swap claims and retry bounds remain. + """ if not self.delivery_enabled: logger.debug("API-cost outbox delivery is disabled") - return - if not self.webhook_url: - logger.warning( - "API-cost outbox delivery skipped because no webhook is configured" - ) - return + return 0 + await self.recover_stale_deliveries() - item_ids = await asyncio.to_thread(self._list_outbox_candidates, max_items) + if not self.webhook_url: + return 0 + + item_ids = await asyncio.to_thread( + self._select_outbox_item_ids, + now=datetime.now(timezone.utc), + force=force, + max_items=max_items, + ) + + completed = 0 for item_id in item_ids: - payload = await asyncio.to_thread(self._claim_outbox_item, item_id) - if payload is None: + claim = await asyncio.to_thread( + self._try_claim_outbox_item, + item_id, + datetime.now(timezone.utc), + respect_schedule=not force, + ) + if claim is None: continue - success = await self._send_webhook_notification(payload) - await asyncio.to_thread(self._complete_outbox_item, item_id, success) - def _list_outbox_candidates(self, max_items: Optional[int]) -> list[int]: - with self._session_scope() as session: - query = ( - session.query(WebhookOutbox.id) - .filter( - WebhookOutbox.status.in_(["pending", "failed"]), - WebhookOutbox.retry_count < 5, + event_id = f"api-cost:{claim['utc_date']}:{claim['alert_type']}" + token = _WEBHOOK_EVENT_ID.set(event_id) + try: + success = await self._send_webhook_notification(claim["payload"]) + except asyncio.CancelledError: + await asyncio.to_thread( + self._complete_outbox_claim, + claim, + success=False, + error_message="Delivery cancelled during worker shutdown", ) - .order_by(WebhookOutbox.id) + raise + except Exception as e: + logger.error("Webhook outbox delivery %s raised: %s", item_id, e) + success = False + finally: + _WEBHOOK_EVENT_ID.reset(token) + + claim_completed = await asyncio.to_thread( + self._complete_outbox_claim, claim, success=success ) - if max_items is not None: - query = query.limit(max(0, max_items)) - return [row[0] for row in query.all()] + if success and claim_completed: + completed += 1 - def _claim_outbox_item(self, item_id: int) -> Optional[str]: - try: - with self._session_scope(commit=True) as session: - item = session.query(WebhookOutbox).filter_by(id=item_id).first() - if ( - item is None - or item.status not in {"pending", "failed"} - or item.retry_count >= 5 - ): - return None - item.status = "processing" - item.retry_count += 1 - item.last_attempt = datetime.now(timezone.utc) - item.claimed_at = item.last_attempt - return item.payload - except Exception as exc: - logger.error("Error claiming outbox item %s: %s", item_id, exc) - return None - - def _complete_outbox_item(self, item_id: int, success: bool) -> None: - try: - with self._session_scope(commit=True) as session: - item = session.query(WebhookOutbox).filter_by(id=item_id).first() - if item is None: - return - if success: - item.status = "sent" - item.sent_at = datetime.now(timezone.utc) - item.error_message = None - else: - item.status = "failed" - item.error_message = "Delivery failed" - except Exception as exc: - logger.error("Error updating outbox item %s: %s", item_id, exc) + return completed async def _send_webhook_notification(self, message: str) -> bool: """Send an async webhook notification if URL is configured. @@ -1446,15 +1884,26 @@ async def _send_webhook_notification(self, message: str) -> bool: True if the POST completed with a successful 2xx status; False otherwise. """ if not self.webhook_url: - return True # Behave as successful delivery if no webhook is configured + return False try: payload = {"text": message, "content": message} + event_id = _WEBHOOK_EVENT_ID.get() + headers = ( + {"Idempotency-Key": event_id, "X-Event-ID": event_id} + if event_id + else None + ) + request_kwargs: dict[str, Any] = { + "json": payload, + "timeout": aiohttp.ClientTimeout(total=5), + } + if headers is not None: + request_kwargs["headers"] = headers async with aiohttp.ClientSession() as session: async with session.post( self.webhook_url, - json=payload, - timeout=aiohttp.ClientTimeout(total=5), + **request_kwargs, ) as response: if response.status >= 200 and response.status < 300: return True diff --git a/src/youtube_extension/services/ai/gemini_service.py b/src/youtube_extension/services/ai/gemini_service.py index d474fc179..c5f1f0ede 100644 --- a/src/youtube_extension/services/ai/gemini_service.py +++ b/src/youtube_extension/services/ai/gemini_service.py @@ -305,6 +305,7 @@ class GeminiResult: model_name: str backend: str # "api" or "vertex" error: Optional[str] = None + usage_metadata: Optional[Any] = None class GeminiService: @@ -588,7 +589,8 @@ async def process_image( response=response.text, latency=latency, model_name=self.config.model_name, - backend="vertex" if self._use_vertex else "api" + backend="vertex" if self._use_vertex else "api", + usage_metadata=getattr(response, "usage_metadata", None), ) except Exception as e: @@ -676,6 +678,7 @@ async def process_text( latency=time.time() - start_time, model_name=self.config.model_name, backend=self._backend_kind, + usage_metadata=getattr(response, "usage_metadata", None), ) except Exception as exc: @@ -825,7 +828,8 @@ async def process_video( response=response.text, latency=latency, model_name=self.config.model_name, - backend="vertex" if self._use_vertex else "api" + backend="vertex" if self._use_vertex else "api", + usage_metadata=getattr(response, "usage_metadata", None), ) except Exception as e: @@ -892,6 +896,7 @@ async def process_audio( latency=latency, model_name=self.config.model_name, backend="vertex" if self._use_vertex else "api", + usage_metadata=getattr(response, "usage_metadata", None), ) except Exception as e: @@ -1173,7 +1178,8 @@ async def process_youtube( response=response.text, latency=latency, model_name=self.config.model_name, - backend="api" + backend="api", + usage_metadata=getattr(response, "usage_metadata", None), ) except Exception as e: diff --git a/src/youtube_extension/services/ai/hybrid_processor_service.py b/src/youtube_extension/services/ai/hybrid_processor_service.py index 3d0b9b3bb..3afed7606 100644 --- a/src/youtube_extension/services/ai/hybrid_processor_service.py +++ b/src/youtube_extension/services/ai/hybrid_processor_service.py @@ -26,6 +26,13 @@ from .gemini_service import GeminiConfig, GeminiResult, GeminiService +async def _record_api_usage(*args: Any, **kwargs: Any) -> Any: + """Load cost tracking only when provider usage is actually available.""" + from youtube_extension.backend.services.api_cost_monitor import track_api_call + + return await track_api_call(*args, **kwargs) + + class ProcessingMode(Enum): """Processing mode roadmap retained for compatibility.""" @@ -261,6 +268,11 @@ async def process( **kwargs, ) + await self._track_gemini_usage( + cloud_result, + routing_decision.task_type, + ) + hybrid_result = HybridResult( success=cloud_result.success, response=cloud_result.response, @@ -290,6 +302,46 @@ async def process( error=str(exc), ) + async def _track_gemini_usage( + self, + result: GeminiResult, + task_type: TaskType, + ) -> None: + """Persist provider-reported usage without delaying a paid result.""" + if not result.success or result.backend not in {"api", "vertex", "gemini"}: + return + + usage = result.usage_metadata + if usage is None: + self.logger.warning( + "Gemini response omitted usage metadata; cost record skipped" + ) + return + + input_tokens = int(getattr(usage, "prompt_token_count", 0) or 0) + output_tokens = int(getattr(usage, "candidates_token_count", 0) or 0) + output_tokens += int(getattr(usage, "thoughts_token_count", 0) or 0) + if input_tokens <= 0 and output_tokens <= 0: + self.logger.warning( + "Gemini usage metadata contained no billable token counts" + ) + return + + try: + await _record_api_usage( + "google", + "hybrid/process", + input_tokens, + model=result.model_name, + output_tokens=output_tokens, + request_type=task_type.value, + success=True, + ) + except Exception: + self.logger.exception( + "Gemini usage tracking failed after provider completion" + ) + async def _call_gemini( self, input_data: str | Path | Image.Image, diff --git a/tests/unit/test_api_cost_database_substrate.py b/tests/unit/test_api_cost_database_substrate.py index 6baefd69c..48dd5510b 100644 --- a/tests/unit/test_api_cost_database_substrate.py +++ b/tests/unit/test_api_cost_database_substrate.py @@ -630,8 +630,9 @@ async def tracking_to_thread( await monitor.record_usage("openai", "/chat", 100, model="gpt-4o") - assert "_record_usage_sync" in calls - assert "_get_daily_cost_sync" in calls + # Usage persistence and UTC-day aggregation now share one worker-thread + # transaction; a second daily-cost query would reopen the crash boundary. + assert calls == ["_record_usage_sync"] async def test_telemetry_database_failure_does_not_fail_paid_api_result( diff --git a/tests/unit/test_api_cost_monitor.py b/tests/unit/test_api_cost_monitor.py index f56055167..adedcd82b 100644 --- a/tests/unit/test_api_cost_monitor.py +++ b/tests/unit/test_api_cost_monitor.py @@ -89,12 +89,24 @@ def test_youtube_quota_cost(self, monitor): def test_unknown_service_returns_zero(self, monitor): assert monitor.calculate_cost("nonexistent", "model", 1000) == 0.0 - def test_unknown_model_falls_back_to_first_model(self, monitor): + def test_unknown_model_fails_closed(self, monitor): + with pytest.raises(ValueError, match="Unknown pricing model"): + monitor.calculate_cost( + "anthropic", "unknown-model", input_tokens=1000, output_tokens=0 + ) + + def test_default_model_preserves_legacy_service_costing(self, monitor): cost = monitor.calculate_cost( - "anthropic", "unknown-model", input_tokens=1000, output_tokens=0 + "openai", "default", input_tokens=1000, output_tokens=0 ) assert cost > 0.0 + def test_google_gemini_35_flash_cost(self, monitor): + cost = monitor.calculate_cost( + "google", "gemini-3.5-flash", input_tokens=1000, output_tokens=1000 + ) + assert pytest.approx(cost, rel=1e-6) == 0.0015 + 0.009 + def test_zero_tokens_returns_zero_cost(self, monitor): cost = monitor.calculate_cost( "anthropic", "claude-opus-4-8", input_tokens=0, output_tokens=0 @@ -345,6 +357,56 @@ async def test_record_failure_usage(self, monitor): assert record.success is False assert record.error_message == "rate limited" + async def test_usage_and_crossed_alert_commit_atomically(self, monitor): + from youtube_extension.backend.models.api_cost import ( + APIUsage, + DailyBudget, + WebhookOutbox, + ) + + monitor.alert_threshold = 0.001 + monitor.daily_budget = 100.0 + record = await monitor.record_usage( + service="anthropic", + endpoint="/messages", + tokens_used=1000, + model="claude-opus-4-8", + ) + + with monitor._session_scope() as session: + assert session.query(APIUsage).count() == 1 + budget = session.query(DailyBudget).one() + alert = session.query(WebhookOutbox).one() + + assert budget.total_cost == pytest.approx(record.cost) + assert budget.alert_sent is True + assert alert.alert_type == "threshold" + assert alert.current_cost == pytest.approx(record.cost) + + async def test_alert_staging_failure_rolls_back_usage(self, monitor, monkeypatch): + from youtube_extension.backend.models.api_cost import ( + APIUsage, + DailyBudget, + WebhookOutbox, + ) + + def fail_staging(session, timestamp): + raise RuntimeError("simulated crash boundary") + + monkeypatch.setattr(monitor, "_stage_budget_alerts", fail_staging) + record = await monitor.record_usage( + service="anthropic", + endpoint="/messages", + tokens_used=1000, + model="claude-opus-4-8", + ) + + assert record is not None + with monitor._session_scope() as session: + assert session.query(APIUsage).count() == 0 + assert session.query(DailyBudget).count() == 0 + assert session.query(WebhookOutbox).count() == 0 + # =========================================================================== # APICostMonitor — get_daily_cost @@ -740,7 +802,7 @@ async def fake_notification(message): # Attempt 2, 3, 4, 5 for expected_retry in [2, 3, 4, 5]: - await monitor.process_outbox() + await monitor.process_outbox(force=True) session = monitor.Session() try: item = ( @@ -754,7 +816,7 @@ async def fake_notification(message): session.close() # Attempt 6 (should not be retried because retry count reached 5) - await monitor.process_outbox() + await monitor.process_outbox(force=True) session = monitor.Session() try: item = ( diff --git a/tests/unit/test_api_cost_outbox_worker.py b/tests/unit/test_api_cost_outbox_worker.py new file mode 100644 index 000000000..f4603e8d6 --- /dev/null +++ b/tests/unit/test_api_cost_outbox_worker.py @@ -0,0 +1,504 @@ +"""Focused durability and lifecycle tests for the API-cost webhook outbox.""" + +from __future__ import annotations + +import asyncio +import sqlite3 +import sys +import threading +from datetime import datetime, timedelta, timezone +from pathlib import Path + +import pytest +from sqlalchemy import text + +sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "src")) + +from youtube_extension.backend.services import api_cost_monitor as monitor_module +from youtube_extension.backend.services.api_cost_monitor import ( + APICostMonitor, + WebhookOutbox, +) + + +def _get_item(monitor: APICostMonitor, utc_date: str) -> WebhookOutbox: + session = monitor.Session() + try: + item = ( + session.query(WebhookOutbox) + .filter_by(utc_date=utc_date, alert_type="threshold") + .one() + ) + session.expunge(item) + return item + finally: + session.close() + + +async def _wait_until(predicate, timeout: float = 1.0) -> None: + deadline = asyncio.get_running_loop().time() + timeout + while not predicate(): + if asyncio.get_running_loop().time() >= deadline: + raise AssertionError("condition was not reached before timeout") + await asyncio.sleep(0.005) + + +def test_additive_schema_upgrade_preserves_rows_and_adds_due_index(tmp_path): + db_path = tmp_path / "legacy.db" + connection = sqlite3.connect(db_path) + try: + connection.executescript(""" + CREATE TABLE webhook_outbox ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + utc_date VARCHAR NOT NULL, + alert_type VARCHAR NOT NULL, + status VARCHAR NOT NULL, + retry_count INTEGER NOT NULL, + last_attempt DATETIME, + error_message VARCHAR, + current_cost FLOAT NOT NULL, + payload VARCHAR, + CONSTRAINT uq_utc_date_alert_type UNIQUE (utc_date, alert_type) + ); + INSERT INTO webhook_outbox ( + utc_date, alert_type, status, retry_count, current_cost, payload + ) VALUES ( + '2026-07-17', 'threshold', 'pending', 0, 8.5, 'keep me' + ); + """) + connection.commit() + finally: + connection.close() + + monitor = APICostMonitor(db_path=str(db_path)) + + with monitor.engine.connect() as connection: + columns = { + row[1] + for row in connection.execute(text("PRAGMA table_info(webhook_outbox)")) + } + indexes = { + row[1] + for row in connection.execute(text("PRAGMA index_list(webhook_outbox)")) + } + due_index_columns = [ + row[2] + for row in connection.execute( + text("PRAGMA index_info(ix_webhook_outbox_due)") + ) + ] + + assert "next_attempt_at" in columns + assert "ix_webhook_outbox_due" in indexes + assert due_index_columns == ["status", "next_attempt_at", "retry_count", "id"] + assert _get_item(monitor, "2026-07-17").payload == "keep me" + + +def test_schema_initialization_failure_is_not_suppressed(tmp_path, monkeypatch): + def fail_upgrade(self): + raise sqlite3.OperationalError("migration failed") + + monkeypatch.setattr(APICostMonitor, "_upgrade_sqlite_outbox_schema", fail_upgrade) + + with pytest.raises(sqlite3.OperationalError, match="migration failed"): + APICostMonitor(db_path=str(tmp_path / "broken.db")) + + +async def test_constructor_and_alert_do_not_spawn_background_work( + tmp_path, monkeypatch +): + recovery_started = asyncio.Event() + + async def blocking_recovery(self, stale_timeout_seconds=None): + recovery_started.set() + await asyncio.Event().wait() + + monkeypatch.setattr(APICostMonitor, "recover_stale_deliveries", blocking_recovery) + existing_tasks = asyncio.all_tasks() + monitor = APICostMonitor(db_path=str(tmp_path / "no_implicit_tasks.db")) + await asyncio.sleep(0) + spawned_tasks = asyncio.all_tasks() - existing_tasks + + try: + assert spawned_tasks == set() + assert not recovery_started.is_set() + + loop = asyncio.get_running_loop() + created = [] + original_create_task = loop.create_task + + def record_create_task(coro, *args, **kwargs): + created.append(coro) + coro.close() + return None + + with monkeypatch.context() as context: + context.setattr(loop, "create_task", record_create_task) + await monitor._send_budget_alert(8.5, "threshold") + + assert created == [] + assert loop.create_task == original_create_task + finally: + for task in spawned_tasks: + task.cancel() + if spawned_tasks: + await asyncio.gather(*spawned_tasks, return_exceptions=True) + + +async def test_start_is_idempotent_and_close_stops_the_single_worker(tmp_path): + monitor = APICostMonitor(db_path=str(tmp_path / "lifecycle.db")) + monitor.webhook_poll_interval_seconds = 60 + + first = await monitor.start() + second = await monitor.start() + + assert first is second + assert first is monitor._worker_task + assert not first.done() + + await monitor.close() + + assert first.done() + assert monitor._worker_task is None + await monitor.close() + + +async def test_missing_webhook_url_leaves_item_unattempted(tmp_path): + monitor = APICostMonitor(db_path=str(tmp_path / "missing_url.db")) + monitor.webhook_url = None + assert monitor._claim_alert("2026-07-18", "threshold", 8.5) + + await monitor.process_outbox() + + item = _get_item(monitor, "2026-07-18") + assert item.status == "pending" + assert item.retry_count == 0 + assert item.last_attempt is None + assert item.next_attempt_at is None + + +async def test_in_memory_outbox_is_shared_with_worker_threads(monkeypatch): + monitor = APICostMonitor(db_path=":memory:") + monitor.webhook_url = "https://example.test/hook" + + async def succeed(message): + return True + + monkeypatch.setattr(monitor, "_send_webhook_notification", succeed) + assert monitor._claim_alert("2026-07-28", "threshold", 8.5) + + assert await monitor.process_outbox(force=True) == 1 + assert _get_item(monitor, "2026-07-28").status == "sent" + + +async def test_failed_delivery_is_not_counted_as_completed(tmp_path, monkeypatch): + monitor = APICostMonitor(db_path=str(tmp_path / "failed-count.db")) + monitor.webhook_url = "https://example.test/hook" + + async def fail(message): + return False + + monkeypatch.setattr(monitor, "_send_webhook_notification", fail) + assert monitor._claim_alert("2026-07-29", "threshold", 8.5) + + assert await monitor.process_outbox(force=True) == 0 + item = _get_item(monitor, "2026-07-29") + assert item.status == "failed" + assert item.retry_count == 1 + + +async def test_claim_is_compare_and_swap_across_monitor_instances(tmp_path): + db_path = str(tmp_path / "shared.db") + first = APICostMonitor(db_path=db_path) + second = APICostMonitor(db_path=db_path) + assert first._claim_alert("2026-07-19", "threshold", 8.5) + item_id = _get_item(first, "2026-07-19").id + claim_time = datetime.now(timezone.utc) + + claims = await asyncio.gather( + asyncio.to_thread(first._try_claim_outbox_item, item_id, claim_time, True), + asyncio.to_thread(second._try_claim_outbox_item, item_id, claim_time, True), + ) + + assert sum(claim is not None for claim in claims) == 1 + item = _get_item(first, "2026-07-19") + assert item.status == "processing" + assert item.retry_count == 1 + + +async def test_completion_is_conditional_on_the_original_claim(tmp_path): + db_path = str(tmp_path / "conditional-completion.db") + first = APICostMonitor(db_path=db_path) + second = APICostMonitor(db_path=db_path) + assert first._claim_alert("2026-07-25", "threshold", 8.5) + item_id = _get_item(first, "2026-07-25").id + + old_claim = first._try_claim_outbox_item(item_id, datetime.now(timezone.utc), False) + assert old_claim is not None + + session = second.Session() + try: + item = session.query(WebhookOutbox).filter_by(id=item_id).one() + item.status = "failed" + session.commit() + finally: + session.close() + + new_claim = second._try_claim_outbox_item( + item_id, datetime.now(timezone.utc) + timedelta(seconds=1), False + ) + assert new_claim is not None + + assert first._complete_outbox_claim(old_claim, success=True) is False + item = _get_item(first, "2026-07-25") + assert item.status == "processing" + assert item.retry_count == 2 + + assert second._complete_outbox_claim(new_claim, success=True) is True + assert _get_item(first, "2026-07-25").status == "sent" + + +async def test_failure_persists_equal_jitter_backoff_and_respects_due_time( + tmp_path, monkeypatch +): + monitor = APICostMonitor(db_path=str(tmp_path / "backoff.db")) + monitor.webhook_url = "https://example.test/hook" + monitor.webhook_retry_base_seconds = 10 + monitor.webhook_retry_max_seconds = 25 + monkeypatch.setattr(monitor_module.random, "uniform", lambda low, high: high) + + attempts = 0 + + async def fail(message): + nonlocal attempts + attempts += 1 + return False + + monkeypatch.setattr(monitor, "_send_webhook_notification", fail) + assert monitor._claim_alert("2026-07-20", "threshold", 8.5) + + expected_delays = [10, 20, 25, 25] + for expected_attempt, expected_delay in enumerate(expected_delays, start=1): + before = datetime.now(timezone.utc).replace(tzinfo=None) + await monitor.process_outbox(force=True) + item = _get_item(monitor, "2026-07-20") + assert item.retry_count == expected_attempt + assert item.status == "failed" + assert item.next_attempt_at is not None + actual_delay = (item.next_attempt_at - before).total_seconds() + assert expected_delay - 0.5 <= actual_delay <= expected_delay + 0.5 + + await monitor.process_outbox() + assert _get_item(monitor, "2026-07-20").retry_count == expected_attempt + + await monitor.process_outbox(force=True) + item = _get_item(monitor, "2026-07-20") + assert item.retry_count == 5 + assert item.next_attempt_at is None + assert item.error_message.startswith("Retry exhausted") + + await monitor.process_outbox(force=True) + assert _get_item(monitor, "2026-07-20").retry_count == 5 + assert attempts == 5 + + +async def test_worker_automatically_retries_due_delivery(tmp_path, monkeypatch): + monitor = APICostMonitor(db_path=str(tmp_path / "automatic.db")) + monitor.webhook_url = "https://example.test/hook" + monitor.webhook_retry_base_seconds = 0.01 + monitor.webhook_retry_max_seconds = 0.01 + monitor.webhook_poll_interval_seconds = 0.005 + monkeypatch.setattr(monitor_module.random, "uniform", lambda low, high: high) + attempts = 0 + + async def fail_once(message): + nonlocal attempts + attempts += 1 + return attempts > 1 + + monkeypatch.setattr(monitor, "_send_webhook_notification", fail_once) + assert monitor._claim_alert("2026-07-21", "threshold", 8.5) + + try: + await monitor.start() + await _wait_until(lambda: _get_item(monitor, "2026-07-21").status == "sent") + finally: + await monitor.close() + + assert attempts == 2 + assert _get_item(monitor, "2026-07-21").retry_count == 2 + + +@pytest.mark.parametrize("claimed_at", [None, datetime(2020, 1, 1)]) +async def test_stale_processing_recovery_handles_null_and_old_timestamps( + tmp_path, claimed_at +): + suffix = "null" if claimed_at is None else "old" + monitor = APICostMonitor(db_path=str(tmp_path / f"stale-{suffix}.db")) + assert monitor._claim_alert("2026-07-22", "threshold", 8.5) + session = monitor.Session() + try: + item = session.query(WebhookOutbox).one() + item.status = "processing" + item.retry_count = 1 + item.claimed_at = claimed_at + session.commit() + finally: + session.close() + + await monitor.recover_stale_deliveries(stale_timeout_seconds=30) + + item = _get_item(monitor, "2026-07-22") + assert item.status == "failed" + assert item.next_attempt_at is not None + assert "Recovery:" in item.error_message + + +async def test_stale_processing_at_max_attempts_is_terminal(tmp_path): + monitor = APICostMonitor(db_path=str(tmp_path / "stale-exhausted.db")) + assert monitor._claim_alert("2026-07-26", "threshold", 8.5) + session = monitor.Session() + try: + item = session.query(WebhookOutbox).one() + item.status = "processing" + item.retry_count = 5 + item.claimed_at = datetime(2020, 1, 1) + session.commit() + finally: + session.close() + + await monitor.recover_stale_deliveries(stale_timeout_seconds=30) + + item = _get_item(monitor, "2026-07-26") + assert item.status == "failed" + assert item.next_attempt_at is None + assert item.error_message.startswith("Retry exhausted") + + +async def test_cancellation_releases_claim_and_schedules_retry(tmp_path, monkeypatch): + monitor = APICostMonitor(db_path=str(tmp_path / "cancel.db")) + monitor.webhook_url = "https://example.test/hook" + monitor.webhook_retry_base_seconds = 0.01 + monitor.webhook_poll_interval_seconds = 60 + delivery_started = asyncio.Event() + + async def block(message): + delivery_started.set() + await asyncio.Event().wait() + + monkeypatch.setattr(monitor, "_send_webhook_notification", block) + assert monitor._claim_alert("2026-07-23", "threshold", 8.5) + + await monitor.start() + await asyncio.wait_for(delivery_started.wait(), timeout=1) + await monitor.close() + + item = _get_item(monitor, "2026-07-23") + assert item.status == "failed" + assert item.retry_count == 1 + assert item.next_attempt_at is not None + assert "cancel" in item.error_message.lower() + + +async def test_every_attempt_uses_stable_idempotency_headers_and_sent_is_terminal( + tmp_path, monkeypatch +): + monitor = APICostMonitor(db_path=str(tmp_path / "headers.db")) + monitor.webhook_url = "https://example.test/hook" + responses = iter([500, 204]) + captured_headers: list[dict[str, str]] = [] + + class FakeResponse: + def __init__(self, status): + self.status = status + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + return False + + class FakeSession: + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + return False + + def post(self, url, json=None, timeout=None, headers=None): + captured_headers.append(headers) + return FakeResponse(next(responses)) + + monkeypatch.setattr(monitor_module.aiohttp, "ClientSession", FakeSession) + assert monitor._claim_alert("2026-07-24", "threshold", 8.5) + + await monitor.process_outbox(force=True) + await monitor.process_outbox(force=True) + await monitor.process_outbox() + + expected_event_id = "api-cost:2026-07-24:threshold" + assert captured_headers == [ + {"Idempotency-Key": expected_event_id, "X-Event-ID": expected_event_id}, + {"Idempotency-Key": expected_event_id, "X-Event-ID": expected_event_id}, + ] + item = _get_item(monitor, "2026-07-24") + assert item.status == "sent" + assert item.retry_count == 2 + + +async def test_worker_database_transactions_run_off_event_loop(tmp_path, monkeypatch): + monitor = APICostMonitor(db_path=str(tmp_path / "off-loop.db")) + monitor.webhook_url = "https://example.test/hook" + assert monitor._claim_alert("2026-07-27", "threshold", 8.5) + + event_loop_thread = threading.get_ident() + observed_threads: list[tuple[str, int]] = [] + helper_names = ( + "_recover_stale_deliveries_sync", + "_select_outbox_item_ids", + "_try_claim_outbox_item", + "_complete_outbox_claim", + ) + + for helper_name in helper_names: + original = getattr(monitor, helper_name) + + def record_thread(*args, _name=helper_name, _original=original, **kwargs): + observed_threads.append((_name, threading.get_ident())) + return _original(*args, **kwargs) + + monkeypatch.setattr(monitor, helper_name, record_thread) + + async def succeed(message): + return True + + monkeypatch.setattr(monitor, "_send_webhook_notification", succeed) + + assert await monitor.process_outbox(force=True) == 1 + assert _get_item(monitor, "2026-07-27").status == "sent" + assert {name for name, _ in observed_threads} == set(helper_names) + assert all(thread_id != event_loop_thread for _, thread_id in observed_threads) + + +def test_claim_token_fences_completion_and_is_cleared(tmp_path): + monitor = APICostMonitor(db_path=str(tmp_path / "claim-token.db")) + assert monitor._claim_alert("2026-07-28", "threshold", 8.5) + + item = _get_item(monitor, "2026-07-28") + claim = monitor._try_claim_outbox_item( + item.id, datetime.now(timezone.utc), respect_schedule=False + ) + + assert claim is not None + assert claim["claim_token"] + assert _get_item(monitor, "2026-07-28").claim_token == claim["claim_token"] + + stale_claim = dict(claim, claim_token="not-the-owner") + assert monitor._complete_outbox_claim(stale_claim, success=True) is False + assert _get_item(monitor, "2026-07-28").status == "processing" + + assert monitor._complete_outbox_claim(claim, success=True) is True + completed = _get_item(monitor, "2026-07-28") + assert completed.status == "sent" + assert completed.claim_token is None + assert completed.claimed_at is None diff --git a/tests/unit/test_gemini_service.py b/tests/unit/test_gemini_service.py index 0ced34295..e36e34645 100644 --- a/tests/unit/test_gemini_service.py +++ b/tests/unit/test_gemini_service.py @@ -49,6 +49,7 @@ def _make_service(api_key: str = "fake_key", model_name: str | None = None, **ex def _mock_response(text: str = "test response") -> MagicMock: resp = MagicMock() resp.text = text + resp.usage_metadata = None return resp @@ -687,6 +688,21 @@ async def test_process_text_success(self): assert result.success is True assert result.response == "text response here" + async def test_process_text_preserves_usage_metadata(self): + svc, mock_model, m = self._make_initialized_service() + usage = SimpleNamespace( + prompt_token_count=12, + candidates_token_count=7, + total_token_count=19, + ) + mock_response = _mock_response("tracked response") + mock_response.usage_metadata = usage + mock_model.generate_content.return_value = mock_response + + result = await svc.process_text("hello") + + assert result.usage_metadata is usage + async def test_process_text_with_input_text(self): svc, mock_model, m = self._make_initialized_service() mock_response = _mock_response("expanded response") diff --git a/tests/unit/test_hybrid_processor_service.py b/tests/unit/test_hybrid_processor_service.py index c125cff3a..3bf43770e 100644 --- a/tests/unit/test_hybrid_processor_service.py +++ b/tests/unit/test_hybrid_processor_service.py @@ -6,6 +6,7 @@ import sys import types as _types from pathlib import Path +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -102,6 +103,7 @@ def _make_gemini_result( model_name: str = "gemini-2.0-flash", backend: str = "api", error: str | None = None, + usage_metadata: object | None = None, ) -> GeminiResult: return GeminiResult( success=success, @@ -110,6 +112,7 @@ def _make_gemini_result( model_name=model_name, backend=backend, error=error, + usage_metadata=usage_metadata, ) @@ -619,6 +622,56 @@ async def test_process_routes_youtube_url(self): await svc.process("https://www.youtube.com/watch?v=abc", "summarize") svc.gemini.process_youtube.assert_awaited_once() + async def test_process_tracks_provider_reported_usage(self): + usage = SimpleNamespace( + prompt_token_count=125, + candidates_token_count=40, + thoughts_token_count=5, + total_token_count=170, + ) + svc = self._svc( + _make_gemini_result(usage_metadata=usage) + ) + + with patch( + "youtube_extension.services.ai.hybrid_processor_service._record_api_usage", + new=AsyncMock(), + ) as track: + result = await svc.process( + "video.mp4", + "describe", + task_type=TaskType.VIDEO_UNDERSTANDING, + ) + + assert result.success is True + track.assert_awaited_once_with( + "google", + "hybrid/process", + 125, + model="gemini-2.0-flash", + output_tokens=45, + request_type="video_understanding", + success=True, + ) + + async def test_usage_tracking_failure_does_not_discard_paid_result(self): + usage = SimpleNamespace( + prompt_token_count=25, + candidates_token_count=10, + ) + svc = self._svc( + _make_gemini_result(usage_metadata=usage) + ) + + with patch( + "youtube_extension.services.ai.hybrid_processor_service._record_api_usage", + new=AsyncMock(side_effect=RuntimeError("database unavailable")), + ): + result = await svc.process("video.mp4", "describe") + + assert result.success is True + assert result.response == "ok" + async def test_process_routes_mp4_video(self): svc = self._svc() await svc.process("/data/video.mp4", "describe")