no advanced testing for now

Add cost usage to messages endpoint

routstr/payment/models.py

@@ -7,7 +7,7 @@ from fastapi import APIRouter, Depends
 from pydantic.v1 import BaseModel
 from sqlmodel.ext.asyncio.session import AsyncSession
 
-from ..core.db import ModelRow, get_session
+from ..core.db import ModelRow, UpstreamProviderRow, get_session
 from ..core.logging import get_logger
 from ..core.settings import settings
 from .price import sats_usd_price
@@ -405,6 +405,76 @@ async def update_sats_pricing() -> None:
             logger.error(f"Error updating sats pricing: {e}")
 
 
+class ModelTestRequest(BaseModel):
+    model_id: str
+    endpoint_type: str
+    request_data: dict
+
+
+@models_router.post("/api/models/test")
+async def test_model(
+    payload: ModelTestRequest,
+    session: AsyncSession = Depends(get_session),
+) -> dict:
+    """Test a model by sending a request through its configured upstream provider."""
+    from sqlmodel import select
+
+    result = await session.execute(
+        select(ModelRow).where(ModelRow.id == payload.model_id)
+    )
+    model_row = result.scalars().first()
+
+    if not model_row:
+        return {
+            "success": False,
+            "error": f"Model '{payload.model_id}' not found in database",
+            "status_code": 404,
+        }
+
+    provider = await session.get(UpstreamProviderRow, model_row.upstream_provider_id)
+    if not provider:
+        return {
+            "success": False,
+            "error": "Upstream provider not found",
+            "status_code": 404,
+        }
+
+    base_url = provider.base_url.rstrip("/")
+    if payload.endpoint_type == "chat-completions":
+        url = f"{base_url}/chat/completions"
+    else:
+        url = f"{base_url}/{payload.endpoint_type}"
+
+    actual_model_id = model_row.forwarded_model_id or model_row.id
+    request_data = dict(payload.request_data)
+    request_data["model"] = actual_model_id
+
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {provider.api_key}",
+    }
+
+    try:
+        async with httpx.AsyncClient(timeout=30.0) as client:
+            response = await client.post(url, json=request_data, headers=headers)
+            try:
+                response_data = response.json()
+            except Exception:
+                response_data = {"raw": response.text}
+
+            return {
+                "success": response.status_code < 400,
+                "data": response_data,
+                "status_code": response.status_code,
+            }
+    except Exception as e:
+        return {
+            "success": False,
+            "error": str(e),
+            "status_code": 500,
+        }
+
+
 @models_router.get("/v1/models")
 @models_router.get("/v1/models/", include_in_schema=False)
 @models_router.get("/models")

routstr/proxy.py

@@ -69,7 +69,25 @@ def get_upstreams() -> list[BaseUpstreamProvider]:
 
 def get_model_instance(model_id: str) -> Model | None:
     """Get Model instance by ID from global cache."""
-    return _model_instances.get(model_id.lower())
+    if not model_id:
+        return None
+
+    model_id_lower = model_id.lower()
+    # Try exact match first
+    if model := _model_instances.get(model_id_lower):
+        return model
+
+    # Try stripping common version suffixes (e.g., -20251222)
+    # This handles cases where upstream returns a specific version
+    # but we only track the base model name.
+    import re
+
+    base_model_id = re.sub(r"-\d{8}$", "", model_id_lower)
+    if base_model_id != model_id_lower:
+        if model := _model_instances.get(base_model_id):
+            return model
+
+    return None
 
 
 def get_provider_for_model(model_id: str) -> list[BaseUpstreamProvider] | None:

routstr/upstream/base.py

@@ -119,6 +119,51 @@ class BaseUpstreamProvider:
             "can_show_balance": False,
         }
 
+    def inject_cost_metadata(
+        self,
+        response_json: dict,
+        cost_data: CostData | MaxCostData | dict,
+        key: ApiKey,
+    ) -> None:
+        """Unifies the injection of cost and usage metadata across all completion types."""
+        if isinstance(cost_data, dict):
+            total_msats = cost_data.get("total_msats", 0)
+            total_usd = cost_data.get("total_usd", 0.0)
+            cost_dict = cost_data
+        else:
+            total_msats = cost_data.total_msats
+            total_usd = cost_data.total_usd
+            cost_dict = cost_data.dict()
+
+        sats_cost = total_msats // 1000
+
+        # Inject into top-level usage block (OpenAI/Anthropic style)
+        if "usage" in response_json:
+            response_json["usage"]["cost"] = total_usd
+            response_json["usage"]["cost_sats"] = sats_cost
+            response_json["usage"]["remaining_balance_msats"] = key.balance
+
+        # Inject into Anthropic nested usage block if present
+        if (
+            "message" in response_json
+            and isinstance(response_json["message"], dict)
+            and "usage" in response_json["message"]
+        ):
+            response_json["message"]["usage"]["sats_cost"] = sats_cost
+
+        # Unified Routstr metadata
+        response_json["metadata"] = response_json.get("metadata", {})
+        response_json["metadata"]["routstr"] = {
+            "cost": cost_dict,
+            "sats_cost": sats_cost,
+            "remaining_balance_msats": key.balance,
+        }
+
+        # Legacy/Compatibility fields
+        response_json["cost"] = cost_dict.copy()
+        response_json["cost"]["sats_cost"] = sats_cost
+        response_json["cost"]["remaining_balance_msats"] = key.balance
+
     def prepare_headers(self, request_headers: dict) -> dict:
         """Prepare headers for upstream request by removing proxy-specific headers and adding authentication.
 
@@ -518,24 +563,32 @@ class BaseUpstreamProvider:
                             continue
 
                         try:
-                            obj = json.loads(part)
-                            if isinstance(obj, dict):
-                                if obj.get("model"):
-                                    last_model_seen = str(obj.get("model"))
-
-                                if requested_model:
-                                    obj["model"] = requested_model
-
-                                if "id" not in obj or not isinstance(obj["id"], str):
-                                    obj["id"] = f"chatcmpl-{uuid.uuid4()}"
-
-                                if isinstance(obj.get("usage"), dict):
-                                    # Hold this chunk back to merge cost later
-                                    usage_chunk_data = obj
+                            # Only parse if it looks like a JSON object to avoid SSE control messages or partials
+                            if part.strip().startswith(b"{") and part.strip().endswith(
+                                b"}"
+                            ):
+                                obj = json.loads(part)
+                                if isinstance(obj, dict):
+                                    if obj.get("model"):
+                                        last_model_seen = str(obj.get("model"))
+                                    if requested_model:
+                                        obj["model"] = requested_model
+                                    if (
+                                        "id" not in obj
+                                        or not isinstance(obj["id"], str)
+                                        or obj["id"] == "existing-id"
+                                    ):
+                                        if not hasattr(self, "_current_stream_id"):
+                                            self._current_stream_id = (
+                                                f"chatcmpl-{uuid.uuid4()}"
+                                            )
+                                        obj["id"] = self._current_stream_id
+                                    if isinstance(obj.get("usage"), dict):
+                                        usage_chunk_data = obj
+                                        continue
+                                    yield b"data: " + json.dumps(obj).encode() + b"\n\n"
                                     continue
-                                yield b"data: " + json.dumps(obj).encode() + b"\n\n"
-                                continue
-                        except json.JSONDecodeError:
+                        except Exception:
                             pass
 
                         prefix = (
@@ -1131,7 +1184,11 @@ class BaseUpstreamProvider:
                 )
 
     async def handle_streaming_messages_completion(
-        self, response: httpx.Response, key: ApiKey, max_cost_for_model: int
+        self,
+        response: httpx.Response,
+        key: ApiKey,
+        max_cost_for_model: int,
+        requested_model: str | None = None,
     ) -> StreamingResponse:
         async def stream_with_cost(
             max_cost_for_model: int,
@@ -1170,6 +1227,8 @@ class BaseUpstreamProvider:
                     stored_chunks.append(chunk)
                     try:
                         decoded_chunk = chunk.decode("utf-8", errors="ignore")
+                        modified_lines = []
+                        changed = False
                         for line in decoded_chunk.split("\n"):
                             if line.startswith("data: "):
                                 try:
@@ -1179,6 +1238,20 @@ class BaseUpstreamProvider:
                                         if msg and msg.get("model"):
                                             last_model_seen = str(msg.get("model"))
 
+                                        if requested_model:
+                                            # Apply requested_model override
+                                            model_updated = False
+                                            if msg:
+                                                msg["model"] = requested_model
+                                                model_updated = True
+                                            if data.get("model"):
+                                                data["model"] = requested_model
+                                                model_updated = True
+
+                                            if model_updated:
+                                                line = "data: " + json.dumps(data)
+                                                changed = True
+
                                         if usage := msg.get("usage"):
                                             input_tokens += usage.get("input_tokens", 0)
                                             output_tokens += usage.get(
@@ -1192,10 +1265,14 @@ class BaseUpstreamProvider:
                                             )
                                 except json.JSONDecodeError:
                                     pass
-                    except Exception:
-                        pass
+                            modified_lines.append(line)
 
-                    yield chunk
+                        if changed:
+                            yield "\n".join(modified_lines).encode("utf-8")
+                        else:
+                            yield chunk
+                    except Exception:
+                        yield chunk
 
                 usage_data = {
                     "input_tokens": input_tokens,
@@ -1217,6 +1294,11 @@ class BaseUpstreamProvider:
                                     new_session,
                                     max_cost_for_model,
                                 )
+
+                                self.inject_cost_metadata(
+                                    combined_data, cost_data, fresh_key
+                                )
+
                                 usage_finalized = True
                                 yield f"event: cost\ndata: {json.dumps({'cost': cost_data})}\n\n".encode()
                             except Exception:
@@ -1256,11 +1338,22 @@ class BaseUpstreamProvider:
         session: AsyncSession,
         deducted_max_cost: int,
         path: str,
+        requested_model: str | None = None,
     ) -> Response:
         try:
             content = await response.aread()
             response_json = json.loads(content)
 
+            if requested_model:
+                if "model" in response_json:
+                    response_json["model"] = requested_model
+                if (
+                    "message" in response_json
+                    and isinstance(response_json["message"], dict)
+                    and "model" in response_json["message"]
+                ):
+                    response_json["message"]["model"] = requested_model
+
             if path.endswith("count_tokens") and "usage" not in response_json:
                 input_tokens = response_json.get("input_tokens", 0)
                 response_json["usage"] = {"input_tokens": input_tokens}
@@ -1268,7 +1361,8 @@ class BaseUpstreamProvider:
             cost_data = await adjust_payment_for_tokens(
                 key, response_json, session, deducted_max_cost
             )
-            response_json["cost"] = cost_data
+
+            self.inject_cost_metadata(response_json, cost_data, key)
 
             allowed_headers = {
                 "content-type",
@@ -1422,7 +1516,10 @@ class BaseUpstreamProvider:
 
                     if is_streaming and response.status_code == 200:
                         result = await self.handle_streaming_messages_completion(
-                            response, key, max_cost_for_model
+                            response,
+                            key,
+                            max_cost_for_model,
+                            requested_model=original_model_id,
                         )
                         background_tasks = BackgroundTasks()
                         background_tasks.add_task(response.aclose)
@@ -1433,7 +1530,12 @@ class BaseUpstreamProvider:
                     if response.status_code == 200:
                         try:
                             return await self.handle_non_streaming_messages_completion(
-                                response, key, session, max_cost_for_model, path
+                                response,
+                                key,
+                                session,
+                                max_cost_for_model,
+                                path,
+                                requested_model=original_model_id,
                             )
                         finally:
                             await response.aclose()
@@ -1443,7 +1545,12 @@ class BaseUpstreamProvider:
                     if response.status_code == 200:
                         try:
                             return await self.handle_non_streaming_messages_completion(
-                                response, key, session, max_cost_for_model, path
+                                response,
+                                key,
+                                session,
+                                max_cost_for_model,
+                                path,
+                                requested_model=original_model_id,
                             )
                         finally:
                             await response.aclose()

routstr/upstream/ppqai.py

@@ -1,9 +1,9 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
 
 import httpx
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 from ..core.logging import get_logger
 from ..payment.models import Architecture, Model, Pricing, async_fetch_openrouter_models
@@ -16,18 +16,20 @@ logger = get_logger(__name__)
 
 
 class PPQAIModelPricing(BaseModel):
-    ui: dict[str, float]
-    api: dict[str, float]
+    ui: Optional[dict[str, float]] = None
+    api: Optional[dict[str, float]] = None
+    input_per_1M_tokens: Optional[float] = Field(None, alias="input_per_1M_tokens")
+    output_per_1M_tokens: Optional[float] = Field(None, alias="output_per_1M_tokens")
 
 
 class PPQAIModel(BaseModel):
     id: str
-    provider: str
+    provider: Optional[str] = None
     name: str
     created_at: int
     context_length: int
     pricing: PPQAIModelPricing
-    popular: bool
+    popular: bool = False
 
 
 class PPQAIUpstreamProvider(BaseUpstreamProvider):
@@ -134,31 +136,54 @@ class PPQAIUpstreamProvider(BaseUpstreamProvider):
                         )
 
                         if or_model:
-                            if input_price := ppqai_model.pricing.api.get(
-                                "input_per_1M"
-                            ):
+                            input_price = None
+                            if ppqai_model.pricing.api:
+                                input_price = ppqai_model.pricing.api.get(
+                                    "input_per_1M"
+                                )
+                            elif ppqai_model.pricing.input_per_1M_tokens:
+                                input_price = ppqai_model.pricing.input_per_1M_tokens
+
+                            if input_price is not None:
                                 or_model.pricing.prompt = input_price / 1_000_000
-                            if output_price := ppqai_model.pricing.api.get(
-                                "output_per_1M"
-                            ):
+
+                            output_price = None
+                            if ppqai_model.pricing.api:
+                                output_price = ppqai_model.pricing.api.get(
+                                    "output_per_1M"
+                                )
+                            elif ppqai_model.pricing.output_per_1M_tokens:
+                                output_price = ppqai_model.pricing.output_per_1M_tokens
+
+                            if output_price is not None:
                                 or_model.pricing.completion = output_price / 1_000_000
+
                             if cl := ppqai_model.context_length:
                                 or_model.context_length = cl
                             models.append(or_model)
                         else:
-                            input_price = ppqai_model.pricing.api.get(
-                                "input_per_1M", 0.0
-                            )
-                            output_price = ppqai_model.pricing.api.get(
-                                "output_per_1M", 0.0
-                            )
+                            input_price = 0.0
+                            if ppqai_model.pricing.api:
+                                input_price = ppqai_model.pricing.api.get(
+                                    "input_per_1M", 0.0
+                                )
+                            elif ppqai_model.pricing.input_per_1M_tokens:
+                                input_price = ppqai_model.pricing.input_per_1M_tokens
+
+                            output_price = 0.0
+                            if ppqai_model.pricing.api:
+                                output_price = ppqai_model.pricing.api.get(
+                                    "output_per_1M", 0.0
+                                )
+                            elif ppqai_model.pricing.output_per_1M_tokens:
+                                output_price = ppqai_model.pricing.output_per_1M_tokens
 
                             models.append(
                                 Model(
                                     id=ppqai_model.id,
                                     name=ppqai_model.name,
                                     created=ppqai_model.created_at // 1000,
-                                    description=f"{ppqai_model.provider} model",
+                                    description=f"{ppqai_model.provider or 'PPQ.AI'} model",
                                     context_length=ppqai_model.context_length,
                                     architecture=Architecture(
                                         modality="text->text",

tests/unit/test_stream_id_injection.py

@@ -0,0 +1,109 @@
+import json
+from collections.abc import AsyncGenerator
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from routstr.core.db import ApiKey
+from routstr.upstream.base import BaseUpstreamProvider
+
+
+@pytest.mark.asyncio
+async def test_stream_with_id_injection() -> None:
+    """Test that stream_with_cost correctly injects IDs into complete JSON chunks but skips partials."""
+    provider = BaseUpstreamProvider(
+        base_url="https://api.example.com", api_key="test_key"
+    )
+
+    # Mock response with mixed chunks:
+    # 1. Complete JSON without ID
+    # 2. Partial JSON (should be passed through)
+    # 3. Complete JSON with ID (should be preserved or updated if requested_model is set)
+    # 4. [DONE] message
+    chunks = [
+        b'data: {"choices": [{"delta": {"content": "Hello"}}]}\n\n',
+        b'data: {"choices": [{"delta": {"content": "',  # Partial
+        b'world"}}]}\n\n',
+        b'data: {"id": "existing-id", "choices": [{"delta": {"content": "!"}}]}\n\n',
+        b"data: [DONE]\n\n",
+    ]
+
+    async def aiter_bytes() -> AsyncGenerator[bytes, None]:
+        for chunk in chunks:
+            yield chunk
+
+    mock_response = MagicMock()
+    mock_response.status_code = 200
+    mock_response.headers = {"content-type": "text/event-stream"}
+    mock_response.aiter_bytes = aiter_bytes
+
+    key = MagicMock(spec=ApiKey)
+    key.hashed_key = "test_hash"
+    key.balance = 1000
+
+    background_tasks = MagicMock()
+
+    # We need to mock adjust_payment_for_tokens since it's called at the end
+    with MagicMock():
+        from routstr.upstream import base
+
+        # Mocking the module-level function used in the generator
+        base.adjust_payment_for_tokens = AsyncMock(
+            return_value={"total_usd": 0.1, "total_msats": 100}
+        )
+        base.create_session = MagicMock()
+
+        streaming_response = await provider.handle_streaming_chat_completion(
+            response=mock_response,
+            key=key,
+            max_cost_for_model=100,
+            background_tasks=background_tasks,
+            requested_model="test-model",
+        )
+
+        results = []
+        async for chunk in streaming_response.body_iterator:
+            results.append(chunk)
+
+        # Parse results
+        parsed_results = []
+        for r in results:
+            if isinstance(r, bytes) and r.startswith(b"data: "):
+                data = r[6:].decode().strip()
+                if data == "[DONE]":
+                    parsed_results.append(data)
+                else:
+                    try:
+                        parsed_results.append(json.loads(data))
+                    except (json.JSONDecodeError, UnicodeDecodeError):
+                        parsed_results.append(
+                            data
+                        )  # Keep as string if it failed to parse
+
+        # Verifications
+        # 1. First chunk should have an injected ID and the requested model
+        assert isinstance(parsed_results[0], dict)
+        assert "id" in parsed_results[0]
+        assert parsed_results[0]["id"].startswith("chatcmpl-")
+        assert parsed_results[0]["model"] == "test-model"
+
+        # 2. Second chunk was partial, should be passed as-is
+        # In current implementation, re.split(b"data: ", b'data: {...') gives ['', '{...']
+        # The first empty part is skipped. The second part is processed.
+
+        # Check that we have results
+        assert len(parsed_results) >= 4
+
+        # Find the chunk that was "existing-id"
+        id_chunk = next(
+            r
+            for r in parsed_results
+            if isinstance(r, dict)
+            and "choices" in r
+            and r["choices"][0]["delta"].get("content") == "!"
+        )
+        assert id_chunk["id"] == parsed_results[0]["id"]
+        assert id_chunk["model"] == "test-model"
+
+        # 4. [DONE] should be there
+        assert "[DONE]" in parsed_results

ui/components/api-endpoint-tester.tsx

@@ -471,7 +471,11 @@ export function ApiEndpointTester({ models }: ApiEndpointTesterProps) {
     testEndpointMutation.mutate(requestData);
   };
 
-  const enabledModels = models.filter((model) => model.isEnabled);
+  const enabledModels = Array.from(
+    new Map(
+      models.filter((model) => model.isEnabled).map((m) => [m.id, m])
+    ).values()
+  );
   const credentials = selectedModel ? getModelCredentials(selectedModel) : null;
   const endpointUrl = credentials
     ? buildEndpointUrl(

ui/components/model-tester.tsx

@@ -197,7 +197,11 @@ export function ModelTester({ models }: ModelTesterProps) {
     testModelMutation.mutate(request);
   };
 
-  const enabledModels = models.filter((model) => model.isEnabled);
+  const enabledModels = Array.from(
+    new Map(
+      models.filter((model) => model.isEnabled).map((m) => [m.id, m])
+    ).values()
+  );
   const credentials = selectedModel ? getModelCredentials(selectedModel) : null;
 
   return (

ui/components/models-page.tsx

@@ -124,12 +124,14 @@ export function ModelsPage() {
             >
               Basic Testing
             </TabsTrigger>
+            {/*
             <TabsTrigger
               value='test-api'
               className='h-9 snap-start px-2 text-[13px] sm:h-10 sm:px-2.5 sm:text-sm'
             >
               API Endpoints
             </TabsTrigger>
+            */}
           </TabsList>
 
           <TabsContent value='manage' className='mt-0'>