diff --git a/routstr/upstream/anthropic.py b/routstr/upstream/anthropic.py index 5e48f058..7944cbd5 100644 --- a/routstr/upstream/anthropic.py +++ b/routstr/upstream/anthropic.py @@ -24,7 +24,7 @@ class AnthropicUpstreamProvider(BaseUpstreamProvider): ) @classmethod - def from_db_row( + def _build_from_row( cls, provider_row: "UpstreamProviderRow" ) -> "AnthropicUpstreamProvider": return cls( diff --git a/routstr/upstream/auto_topup.py b/routstr/upstream/auto_topup.py index 31397882..3582b3fc 100644 --- a/routstr/upstream/auto_topup.py +++ b/routstr/upstream/auto_topup.py @@ -97,6 +97,8 @@ async def _check_and_topup(row: UpstreamProviderRow) -> None: # Instantiate provider and check balance provider = RoutstrUpstreamProvider.from_db_row(row) + if provider is None: + return balance = await provider.get_balance() if balance is None: diff --git a/routstr/upstream/azure.py b/routstr/upstream/azure.py index 11cee17c..a693b763 100644 --- a/routstr/upstream/azure.py +++ b/routstr/upstream/azure.py @@ -38,7 +38,7 @@ class AzureUpstreamProvider(BaseUpstreamProvider): self.api_version = api_version @classmethod - def from_db_row( + def _build_from_row( cls, provider_row: "UpstreamProviderRow" ) -> "AzureUpstreamProvider | None": if not provider_row.api_version: diff --git a/routstr/upstream/base.py b/routstr/upstream/base.py index 2467044c..51afb080 100644 --- a/routstr/upstream/base.py +++ b/routstr/upstream/base.py @@ -6,13 +6,12 @@ import math import traceback import uuid from collections.abc import AsyncGenerator, AsyncIterator, Iterator -from typing import Any, Mapping, cast +from typing import Any, Mapping, Self, cast import httpx from fastapi import BackgroundTasks, HTTPException, Request from fastapi.responses import Response, StreamingResponse from pydantic.v1 import BaseModel -from sqlmodel import select from ..auth import adjust_payment_for_tokens from ..core import get_logger @@ -92,6 +91,11 @@ class BaseUpstreamProvider: base_url: str api_key: str provider_fee: float = 1.05 + # Primary key of the ``upstream_providers`` row this instance was built + # from. Set by ``from_db_row`` so a live provider can re-find its own row by + # stable identity instead of its rotatable ``api_key``. ``None`` for + # instances not sourced from a row. + db_id: int | None = None _models_cache: list[Model] = [] _models_by_id: dict[str, Model] = {} @@ -106,6 +110,7 @@ class BaseUpstreamProvider: self.base_url = base_url self.api_key = api_key self.provider_fee = provider_fee + self.db_id = None self._models_cache = [] self._models_by_id = {} @@ -123,10 +128,13 @@ class BaseUpstreamProvider: return detect_litellm_prefix(self.base_url) @classmethod - def from_db_row( - cls, provider_row: "UpstreamProviderRow" - ) -> "BaseUpstreamProvider | None": - """Factory method to instantiate provider from database row. + def from_db_row(cls, provider_row: "UpstreamProviderRow") -> "Self | None": + """Instantiate a provider from a database row, carrying its identity. + + Construction itself is delegated to the ``_build_from_row`` hook (which + subclasses override to match their constructor); this wrapper stamps the + row's primary key onto the instance as ``db_id`` so the provider can + later re-find its own row by identity rather than by its ``api_key``. Args: provider_row: Database row containing provider configuration @@ -134,6 +142,19 @@ class BaseUpstreamProvider: Returns: Instantiated provider or None if instantiation fails """ + provider = cls._build_from_row(provider_row) + if provider is not None: + provider.db_id = provider_row.id + return provider + + @classmethod + def _build_from_row(cls, provider_row: "UpstreamProviderRow") -> "Self | None": + """Construct the provider instance from a row (no identity stamping). + + Overridden by subclasses whose constructors differ from the base + ``(base_url, api_key, provider_fee)`` shape. Callers should use + ``from_db_row`` instead, which also attaches ``db_id``. + """ return cls( base_url=provider_row.base_url, api_key=provider_row.api_key, @@ -4878,14 +4899,11 @@ class BaseUpstreamProvider: """Refresh the in-memory models cache from upstream API.""" try: async with create_session() as session: - stmt = select(UpstreamProviderRow).where( - UpstreamProviderRow.base_url == self.base_url, - UpstreamProviderRow.api_key == self.api_key, + provider = ( + await session.get(UpstreamProviderRow, self.db_id) + if self.db_id is not None + else None ) - result = await session.exec(stmt) - - # .first() returns the object or None if not found - provider = result.first() if not provider or not provider.id: raise HTTPException(status_code=404, detail="Provider not found") diff --git a/routstr/upstream/fireworks.py b/routstr/upstream/fireworks.py index e0bb1fec..92b67838 100644 --- a/routstr/upstream/fireworks.py +++ b/routstr/upstream/fireworks.py @@ -20,7 +20,7 @@ class FireworksUpstreamProvider(BaseUpstreamProvider): ) @classmethod - def from_db_row( + def _build_from_row( cls, provider_row: "UpstreamProviderRow" ) -> "FireworksUpstreamProvider": return cls( diff --git a/routstr/upstream/gemini.py b/routstr/upstream/gemini.py index 54bf849b..54de41a3 100644 --- a/routstr/upstream/gemini.py +++ b/routstr/upstream/gemini.py @@ -51,7 +51,7 @@ class GeminiUpstreamProvider(BaseUpstreamProvider): return self._client @classmethod - def from_db_row( + def _build_from_row( cls, provider_row: "UpstreamProviderRow" ) -> "GeminiUpstreamProvider": return cls( diff --git a/routstr/upstream/generic.py b/routstr/upstream/generic.py index 390c8372..83b9f565 100644 --- a/routstr/upstream/generic.py +++ b/routstr/upstream/generic.py @@ -45,7 +45,7 @@ class GenericUpstreamProvider(BaseUpstreamProvider): ) @classmethod - def from_db_row( + def _build_from_row( cls, provider_row: "UpstreamProviderRow" ) -> "GenericUpstreamProvider": return cls( diff --git a/routstr/upstream/groq.py b/routstr/upstream/groq.py index 4020b2df..17103c35 100644 --- a/routstr/upstream/groq.py +++ b/routstr/upstream/groq.py @@ -20,7 +20,7 @@ class GroqUpstreamProvider(BaseUpstreamProvider): ) @classmethod - def from_db_row(cls, provider_row: "UpstreamProviderRow") -> "GroqUpstreamProvider": + def _build_from_row(cls, provider_row: "UpstreamProviderRow") -> "GroqUpstreamProvider": return cls( api_key=provider_row.api_key, provider_fee=provider_row.provider_fee, diff --git a/routstr/upstream/helpers.py b/routstr/upstream/helpers.py index 9ad4d73f..f40a3a2a 100644 --- a/routstr/upstream/helpers.py +++ b/routstr/upstream/helpers.py @@ -215,9 +215,6 @@ async def init_upstreams() -> list[BaseUpstreamProvider]: provider = _instantiate_provider(provider_row) if provider: - # Keep provider DB id on runtime instance so model mapping can - # bind DB overrides to the correct upstream. - setattr(provider, "db_id", provider_row.id) await provider.refresh_models_cache() logger.debug( f"Initialized {provider_row.provider_type} provider", @@ -391,9 +388,7 @@ def _instantiate_provider( return provider if provider_row.provider_type == "custom": - return BaseUpstreamProvider( - provider_row.base_url, provider_row.api_key, provider_row.provider_fee - ) + return BaseUpstreamProvider.from_db_row(provider_row) logger.error( f"Unknown provider type: {provider_row.provider_type}", diff --git a/routstr/upstream/ollama.py b/routstr/upstream/ollama.py index 74327d0b..9fed0154 100644 --- a/routstr/upstream/ollama.py +++ b/routstr/upstream/ollama.py @@ -43,7 +43,7 @@ class OllamaUpstreamProvider(BaseUpstreamProvider): ) @classmethod - def from_db_row( + def _build_from_row( cls, provider_row: "UpstreamProviderRow" ) -> "OllamaUpstreamProvider": return cls( diff --git a/routstr/upstream/openai.py b/routstr/upstream/openai.py index 11cc4336..f2f03cc5 100644 --- a/routstr/upstream/openai.py +++ b/routstr/upstream/openai.py @@ -20,7 +20,7 @@ class OpenAIUpstreamProvider(BaseUpstreamProvider): ) @classmethod - def from_db_row( + def _build_from_row( cls, provider_row: "UpstreamProviderRow" ) -> "OpenAIUpstreamProvider": return cls( diff --git a/routstr/upstream/openrouter.py b/routstr/upstream/openrouter.py index 0c69d2cc..63995295 100644 --- a/routstr/upstream/openrouter.py +++ b/routstr/upstream/openrouter.py @@ -90,7 +90,7 @@ class OpenRouterUpstreamProvider(BaseUpstreamProvider): ) @classmethod - def from_db_row( + def _build_from_row( cls, provider_row: "UpstreamProviderRow" ) -> "OpenRouterUpstreamProvider": return cls( diff --git a/routstr/upstream/perplexity.py b/routstr/upstream/perplexity.py index 55a9116e..a088218f 100644 --- a/routstr/upstream/perplexity.py +++ b/routstr/upstream/perplexity.py @@ -23,7 +23,7 @@ class PerplexityUpstreamProvider(BaseUpstreamProvider): ) @classmethod - def from_db_row( + def _build_from_row( cls, provider_row: "UpstreamProviderRow" ) -> "PerplexityUpstreamProvider": return cls( diff --git a/routstr/upstream/ppqai.py b/routstr/upstream/ppqai.py index 48ca9600..0824a08a 100644 --- a/routstr/upstream/ppqai.py +++ b/routstr/upstream/ppqai.py @@ -46,7 +46,7 @@ class PPQAIUpstreamProvider(BaseUpstreamProvider): ) @classmethod - def from_db_row( + def _build_from_row( cls, provider_row: "UpstreamProviderRow" ) -> "PPQAIUpstreamProvider": return cls( @@ -229,17 +229,14 @@ class PPQAIUpstreamProvider(BaseUpstreamProvider): f"Disabling PPQ.AI provider ({self.base_url}) due to insufficient balance", extra={"error": error_message}, ) - from sqlmodel import select - from ..core.db import UpstreamProviderRow, create_session async with create_session() as session: - statement = select(UpstreamProviderRow).where( - UpstreamProviderRow.base_url == self.base_url, - UpstreamProviderRow.api_key == self.api_key, + provider = ( + await session.get(UpstreamProviderRow, self.db_id) + if self.db_id is not None + else None ) - result = await session.exec(statement) - provider = result.first() if provider: provider.enabled = False diff --git a/routstr/upstream/routstr.py b/routstr/upstream/routstr.py index c9962301..0371946a 100644 --- a/routstr/upstream/routstr.py +++ b/routstr/upstream/routstr.py @@ -55,7 +55,7 @@ class RoutstrUpstreamProvider(BaseUpstreamProvider): return path.lstrip("/") @classmethod - def from_db_row( + def _build_from_row( cls, provider_row: "UpstreamProviderRow" ) -> "RoutstrUpstreamProvider": import json diff --git a/routstr/upstream/xai.py b/routstr/upstream/xai.py index b46676cb..58caaba0 100644 --- a/routstr/upstream/xai.py +++ b/routstr/upstream/xai.py @@ -21,7 +21,7 @@ class XAIUpstreamProvider(BaseUpstreamProvider): ) @classmethod - def from_db_row(cls, provider_row: "UpstreamProviderRow") -> "XAIUpstreamProvider": + def _build_from_row(cls, provider_row: "UpstreamProviderRow") -> "XAIUpstreamProvider": return cls( api_key=provider_row.api_key, provider_fee=provider_row.provider_fee, diff --git a/tests/integration/test_provider_self_lookup.py b/tests/integration/test_provider_self_lookup.py new file mode 100644 index 00000000..5745ab65 --- /dev/null +++ b/tests/integration/test_provider_self_lookup.py @@ -0,0 +1,128 @@ +"""A live upstream provider resolves its OWN database row by stable identity +(its primary key), not by its mutable/secret ``api_key``. + +Today ``from_db_row`` drops ``provider_row.id`` and the two self-referential +paths — PPQ.AI's insufficient-balance self-disable and the base +``refresh_models_cache`` — re-find their own row with +``WHERE base_url == self.base_url AND api_key == self.api_key``. That uses a +rotatable secret as a self-handle: the moment the row's key changes underneath a +live object (a rotation racing an in-flight request), the object can no longer +find itself. These tests pin the invariant that a provider looks itself up by +identity, so the lookup survives a key change (and, later, key encryption). +""" + +from unittest.mock import AsyncMock, patch + +import pytest + +from routstr.core.db import UpstreamProviderRow +from routstr.upstream.ppqai import PPQAIUpstreamProvider + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test_provider_object_carries_its_persistent_identity( + integration_session: object, + patched_db_engine: None, +) -> None: + """``from_db_row`` gives the in-memory object its row's identity (``db_id``).""" + row = UpstreamProviderRow( + provider_type="ppqai", + base_url="https://api.ppq.ai", + api_key="sk-original", + enabled=True, + provider_fee=1.0, + ) + integration_session.add(row) # type: ignore[attr-defined] + await integration_session.commit() # type: ignore[attr-defined] + await integration_session.refresh(row) # type: ignore[attr-defined] + + provider = PPQAIUpstreamProvider.from_db_row(row) + assert provider is not None + + assert provider.db_id == row.id + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test_self_disable_targets_own_row_after_key_rotation( + integration_session: object, + patched_db_engine: None, +) -> None: + """PPQ.AI self-disable must disable *its* row even after the key rotated. + + RED (current): the object holds the pre-rotation key, so the + ``(base_url, api_key)`` lookup misses the row → the provider is never + disabled. GREEN: lookup by ``id`` finds it and disables it. + """ + row = UpstreamProviderRow( + provider_type="ppqai", + base_url="https://api.ppq.ai", + api_key="sk-original", + enabled=True, + provider_fee=1.0, + ) + integration_session.add(row) # type: ignore[attr-defined] + await integration_session.commit() # type: ignore[attr-defined] + await integration_session.refresh(row) # type: ignore[attr-defined] + + provider = PPQAIUpstreamProvider.from_db_row(row) # captures sk-original + assert provider is not None + + # Key is rotated in the DB while `provider` is still live. + row.api_key = "sk-rotated" + integration_session.add(row) # type: ignore[attr-defined] + await integration_session.commit() # type: ignore[attr-defined] + + with patch("routstr.proxy.reinitialize_upstreams", new=AsyncMock()): + await provider.on_upstream_error_redirect(402, "Insufficient balance") + + await integration_session.refresh(row) # type: ignore[attr-defined] + assert row.enabled is False + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test_refresh_models_cache_finds_own_row_after_key_rotation( + integration_session: object, + patched_db_engine: None, +) -> None: + """``refresh_models_cache`` must resolve its own row after a key rotation. + + ``refresh_models_cache`` swallows every exception (it only logs), so the + observable proof it found its row is that it reaches ``list_models`` — which + is called with the row's ``id`` only *after* the row is resolved. RED + (current): the stale-key ``(base_url, api_key)`` lookup returns nothing, the + method raises ``404`` internally and returns before ``list_models`` is ever + called. GREEN: lookup by ``id`` finds the row and ``list_models`` runs for + that ``id``. + """ + row = UpstreamProviderRow( + provider_type="ppqai", + base_url="https://api.ppq.ai", + api_key="sk-original", + enabled=True, + provider_fee=1.0, + ) + integration_session.add(row) # type: ignore[attr-defined] + await integration_session.commit() # type: ignore[attr-defined] + await integration_session.refresh(row) # type: ignore[attr-defined] + row_id = row.id + + provider = PPQAIUpstreamProvider.from_db_row(row) # captures sk-original + assert provider is not None + + row.api_key = "sk-rotated" + integration_session.add(row) # type: ignore[attr-defined] + await integration_session.commit() # type: ignore[attr-defined] + + list_models_mock = AsyncMock(return_value=[]) + with ( + patch.object(provider, "fetch_models", new=AsyncMock(return_value=[])), + patch("routstr.upstream.base.list_models", new=list_models_mock), + ): + await provider.refresh_models_cache() + + list_models_mock.assert_awaited_once() + assert list_models_mock.await_args is not None + assert list_models_mock.await_args.kwargs["upstream_id"] == row_id