mirror of
https://github.com/Routstr/routstr-core.git
synced 2026-10-05 12:28:22 +00:00
refactor(upstream): resolve a provider's own row by primary key
A live upstream provider re-finds its own database row in two places — PPQ.AI's insufficient-balance self-disable and the base refresh_models_cache — with WHERE base_url == self.base_url AND api_key == self.api_key. That uses a rotatable secret as a self-handle: if the row's key rotates under a live object, it can no longer find itself. Carry the row's primary key on the instance as db_id, stamped centrally by from_db_row via a _build_from_row construction hook that subclasses override, and look the row up with session.get(UpstreamProviderRow, db_id). This also closes a latent gap where providers built outside the init path (auto-topup) never received db_id. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 4.8
parent
782b233d40
commit
434283c58d
@@ -24,7 +24,7 @@ class AnthropicUpstreamProvider(BaseUpstreamProvider):
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_db_row(
|
||||
def _build_from_row(
|
||||
cls, provider_row: "UpstreamProviderRow"
|
||||
) -> "AnthropicUpstreamProvider":
|
||||
return cls(
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
+31
-13
@@ -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")
|
||||
|
||||
|
||||
@@ -20,7 +20,7 @@ class FireworksUpstreamProvider(BaseUpstreamProvider):
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_db_row(
|
||||
def _build_from_row(
|
||||
cls, provider_row: "UpstreamProviderRow"
|
||||
) -> "FireworksUpstreamProvider":
|
||||
return cls(
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -45,7 +45,7 @@ class GenericUpstreamProvider(BaseUpstreamProvider):
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_db_row(
|
||||
def _build_from_row(
|
||||
cls, provider_row: "UpstreamProviderRow"
|
||||
) -> "GenericUpstreamProvider":
|
||||
return cls(
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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}",
|
||||
|
||||
@@ -43,7 +43,7 @@ class OllamaUpstreamProvider(BaseUpstreamProvider):
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_db_row(
|
||||
def _build_from_row(
|
||||
cls, provider_row: "UpstreamProviderRow"
|
||||
) -> "OllamaUpstreamProvider":
|
||||
return cls(
|
||||
|
||||
@@ -20,7 +20,7 @@ class OpenAIUpstreamProvider(BaseUpstreamProvider):
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_db_row(
|
||||
def _build_from_row(
|
||||
cls, provider_row: "UpstreamProviderRow"
|
||||
) -> "OpenAIUpstreamProvider":
|
||||
return cls(
|
||||
|
||||
@@ -90,7 +90,7 @@ class OpenRouterUpstreamProvider(BaseUpstreamProvider):
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_db_row(
|
||||
def _build_from_row(
|
||||
cls, provider_row: "UpstreamProviderRow"
|
||||
) -> "OpenRouterUpstreamProvider":
|
||||
return cls(
|
||||
|
||||
@@ -23,7 +23,7 @@ class PerplexityUpstreamProvider(BaseUpstreamProvider):
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_db_row(
|
||||
def _build_from_row(
|
||||
cls, provider_row: "UpstreamProviderRow"
|
||||
) -> "PerplexityUpstreamProvider":
|
||||
return cls(
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user