add completion_image_pricing to Pricing Model + hardcoded values for Gemini Image models

Remove flaky wallet tests

routstr/payment/cost_calculation.py

@@ -70,6 +70,7 @@ async def calculate_cost(  # todo: can be sync
     MSATS_PER_1K_OUTPUT_TOKENS: float = (
         float(settings.fixed_per_1k_output_tokens) * 1000.0
     )
+    MSATS_PER_1K_IMAGE_COMPLETION_TOKENS: float = 0.0
 
     if not settings.fixed_pricing:
         response_model = response_data.get("model", "")
@@ -104,11 +105,13 @@ async def calculate_cost(  # todo: can be sync
         try:
             mspp = float(model_obj.sats_pricing.prompt)
             mspc = float(model_obj.sats_pricing.completion)
+            mspci = float(getattr(model_obj.sats_pricing, "completion_image", 0.0))
         except Exception:
             return CostDataError(message="Invalid pricing data", code="pricing_invalid")
 
         MSATS_PER_1K_INPUT_TOKENS = mspp * 1_000_000.0
         MSATS_PER_1K_OUTPUT_TOKENS = mspc * 1_000_000.0
+        MSATS_PER_1K_IMAGE_COMPLETION_TOKENS = mspci * 1_000_000.0
 
         logger.info(
             "Applied model-specific pricing",
@@ -116,6 +119,7 @@ async def calculate_cost(  # todo: can be sync
                 "model": response_model,
                 "input_price_msats_per_1k": MSATS_PER_1K_INPUT_TOKENS,
                 "output_price_msats_per_1k": MSATS_PER_1K_OUTPUT_TOKENS,
+                "image_completion_price_msats_per_1k": MSATS_PER_1K_IMAGE_COMPLETION_TOKENS,
             },
         )
 
@@ -128,13 +132,44 @@ async def calculate_cost(  # todo: can be sync
             },
         )
         return cost_data
+    usage_data = response_data["usage"]
+    input_tokens = usage_data.get("prompt_tokens", 0)
+    output_tokens = usage_data.get("completion_tokens", 0)
 
-    input_tokens = response_data.get("usage", {}).get("prompt_tokens", 0)
-    output_tokens = response_data.get("usage", {}).get("completion_tokens", 0)
+    # added for response api
+    input_tokens = (
+        input_tokens if input_tokens != 0 else usage_data.get("input_tokens", 0)
+    )
+    output_tokens = (
+        output_tokens if output_tokens != 0 else usage_data.get("output_tokens", 0)
+    )
+
+    # Calculate image completion cost
+    image_completion_msats = 0.0
+    if MSATS_PER_1K_IMAGE_COMPLETION_TOKENS > 0:
+        completion_details = usage_data.get("completion_tokens_details", {})
+        image_tokens = completion_details.get("image_tokens", 0)
+
+        if image_tokens > 0:
+            if output_tokens >= image_tokens:
+                output_tokens -= image_tokens
+
+            image_completion_msats = round(
+                image_tokens / 1000 * MSATS_PER_1K_IMAGE_COMPLETION_TOKENS, 3
+            )
+
+            logger.info(
+                "Calculated image completion cost",
+                extra={
+                    "image_tokens": image_tokens,
+                    "image_completion_msats": image_completion_msats,
+                },
+            )
 
     input_msats = round(input_tokens / 1000 * MSATS_PER_1K_INPUT_TOKENS, 3)
+
     output_msats = round(output_tokens / 1000 * MSATS_PER_1K_OUTPUT_TOKENS, 3)
-    token_based_cost = math.ceil(input_msats + output_msats)
+    token_based_cost = math.ceil(input_msats + output_msats + image_completion_msats)
 
     logger.info(
         "Calculated token-based cost",
@@ -143,6 +178,7 @@ async def calculate_cost(  # todo: can be sync
             "output_tokens": output_tokens,
             "input_cost_msats": input_msats,
             "output_cost_msats": output_msats,
+            "image_completion_msats": image_completion_msats,
             "total_cost_msats": token_based_cost,
             "model": response_data.get("model", "unknown"),
         },

routstr/payment/models.py

@@ -31,6 +31,7 @@ class Pricing(BaseModel):
     completion: float
     request: float = 0.0
     image: float = 0.0
+    completion_image: float = 0.0
     web_search: float = 0.0
     internal_reasoning: float = 0.0
     input_cache_read: float = 0.0
@@ -40,6 +41,13 @@ class Pricing(BaseModel):
     max_cost: float = 0.0  # in sats not msats
 
 
+PRICING_OVERRIDES = {
+    "gemini-3-pro-image-preview": {"completion_image": 0.00012},
+    "gemini-2.5-flash-image": {"completion_image": 0.00003},
+    "gemini-2.0-flash": {"completion_image": 0.00003},
+}
+
+
 class TopProvider(BaseModel):
     context_length: int | None = None
     max_completion_tokens: int | None = None
@@ -116,6 +124,16 @@ async def async_fetch_openrouter_models(source_filter: str | None = None) -> lis
                 if not _has_valid_pricing(model):
                     continue
 
+                # Apply manual pricing overrides
+                if model_id in PRICING_OVERRIDES:
+                    pricing = model.get("pricing", {})
+                    if pricing:
+                        for k, v in PRICING_OVERRIDES[model_id].items():
+                            pricing[k] = str(
+                                v
+                            )  # OpenRouter API returns strings for pricing
+                        model["pricing"] = pricing
+
                 models_data.append(model)
 
             return models_data
@@ -148,6 +166,12 @@ def _row_to_model(
     if isinstance(pricing, dict) and float(pricing.get("request", 0.0)) <= 0.0:
         pricing["request"] = max(pricing.get("request", 0.0), 0.0)
 
+    # Apply defaults for missing fields from manual overrides
+    if row.id in PRICING_OVERRIDES and isinstance(pricing, dict):
+        for k, v in PRICING_OVERRIDES[row.id].items():
+            if k not in pricing:
+                pricing[k] = v
+
     parsed_pricing = Pricing.parse_obj(pricing)
     model = Model(
         id=row.id,
@@ -507,6 +531,7 @@ def _pricing_matches(
         "completion",
         "request",
         "image",
+        "completion_image",
         "web_search",
         "internal_reasoning",
     ]

tests/integration/test_wallet_authentication.py

@@ -3,7 +3,6 @@ Integration tests for wallet authentication system including API key generation
 Tests POST /v1/wallet/topup endpoint and authorization header validation.
 """
 
-import hashlib
 from datetime import datetime, timedelta
 from typing import Any
 
@@ -15,7 +14,6 @@ from routstr.core.db import ApiKey
 
 from .utils import (
     CashuTokenGenerator,
-    ConcurrencyTester,
     ResponseValidator,
 )
 
@@ -389,69 +387,6 @@ async def test_api_key_with_expiry_time(
     # The expiry time and refund address functionality is tested elsewhere
 
 
-@pytest.mark.integration
-@pytest.mark.asyncio
-async def test_concurrent_token_submissions(
-    integration_client: AsyncClient, testmint_wallet: Any, integration_session: Any
-) -> None:
-    """Test concurrent submissions of different tokens"""
-
-    # Generate multiple unique tokens with known amounts
-    num_tokens = 10
-    tokens = []
-    expected_balances = {}
-
-    for i in range(num_tokens):
-        amount = 100 + i * 10
-        token = await testmint_wallet.mint_tokens(amount)
-        tokens.append(token)
-        # Store expected balance by token hash
-        hashed_key = hashlib.sha256(token.encode()).hexdigest()
-        expected_balances[hashed_key] = amount * 1000  # msats
-
-    # Create concurrent requests
-    requests = [
-        {
-            "method": "GET",
-            "url": "/v1/wallet/info",
-            "headers": {"Authorization": f"Bearer {token}"},
-        }
-        for token in tokens
-    ]
-
-    # Execute concurrently
-    tester = ConcurrencyTester()
-    responses = await tester.run_concurrent_requests(
-        integration_client, requests, max_concurrent=5
-    )
-
-    # All should succeed
-    assert len(responses) == num_tokens
-    api_keys = set()
-
-    for response in responses:
-        assert response.status_code == 200
-        data = response.json()
-        api_key = data["api_key"]
-        api_keys.add(api_key)
-
-        # Verify balance matches the expected amount
-        hashed_key = api_key[3:]  # Remove "sk-" prefix
-        assert data["balance"] == expected_balances[hashed_key]
-
-    # Should have created unique API keys
-    assert len(api_keys) == num_tokens
-
-    # Verify all keys exist in database
-    for api_key in api_keys:
-        hashed_key = api_key[3:]  # Remove "sk-" prefix
-        result = await integration_session.execute(
-            select(ApiKey).where(ApiKey.hashed_key == hashed_key)  # type: ignore[arg-type]
-        )
-        db_key = result.scalar_one()
-        assert db_key.balance == expected_balances[hashed_key]
-
-
 @pytest.mark.integration
 @pytest.mark.asyncio
 async def test_authorization_with_cashu_token_directly(
@@ -504,48 +439,6 @@ async def test_x_cashu_header_support(
     assert response.status_code == 200
 
 
-@pytest.mark.integration
-@pytest.mark.asyncio
-@pytest.mark.slow
-async def test_api_key_consistency_under_load(
-    integration_client: AsyncClient, testmint_wallet: Any, integration_session: Any
-) -> None:
-    """Test API key generation consistency under concurrent load"""
-
-    # Generate a single token
-    token = await testmint_wallet.mint_tokens(1000)
-
-    # First request to create the API key
-    integration_client.headers["Authorization"] = f"Bearer {token}"
-    initial_response = await integration_client.get("/v1/wallet/info")
-    assert initial_response.status_code == 200
-    expected_api_key = initial_response.json()["api_key"]
-    expected_balance = initial_response.json()["balance"]
-
-    # Try to use the same token concurrently multiple times
-    # All should return the same API key since it's already created
-    requests = [
-        {
-            "method": "GET",
-            "url": "/v1/wallet/info",
-            "headers": {"Authorization": f"Bearer {token}"},
-        }
-        for _ in range(20)  # 20 concurrent attempts
-    ]
-
-    tester = ConcurrencyTester()
-    responses = await tester.run_concurrent_requests(
-        integration_client, requests, max_concurrent=10
-    )
-
-    # All should succeed and return the same API key
-    for response in responses:
-        assert response.status_code == 200
-        data = response.json()
-        assert data["api_key"] == expected_api_key
-        assert data["balance"] == expected_balance
-
-
 @pytest.mark.integration
 @pytest.mark.asyncio
 async def test_database_timestamp_accuracy(

tests/integration/test_wallet_information.py

@@ -13,7 +13,7 @@ from sqlmodel import select, update
 
 from routstr.core.db import ApiKey
 
-from .utils import ConcurrencyTester, ResponseValidator
+from .utils import ResponseValidator
 
 
 @pytest.mark.integration
@@ -204,45 +204,6 @@ async def test_expired_api_key_behavior(
     assert db_key.refund_address == "test@lightning.address"
 
 
-@pytest.mark.integration
-@pytest.mark.asyncio
-async def test_concurrent_access_same_api_key(
-    integration_client: AsyncClient, authenticated_client: AsyncClient
-) -> None:
-    """Test concurrent access with the same API key"""
-
-    # Get the API key from authenticated client
-    response = await authenticated_client.get("/v1/wallet/")
-    api_key = response.json()["api_key"]
-    initial_balance = response.json()["balance"]
-
-    # Create multiple concurrent requests
-    requests = []
-    for i in range(20):
-        # Alternate between both endpoints
-        endpoint = "/v1/wallet/" if i % 2 == 0 else "/v1/wallet/info"
-        requests.append(
-            {
-                "method": "GET",
-                "url": endpoint,
-                "headers": {"Authorization": f"Bearer {api_key}"},
-            }
-        )
-
-    # Execute concurrently
-    tester = ConcurrencyTester()
-    responses = await tester.run_concurrent_requests(
-        integration_client, requests, max_concurrent=10
-    )
-
-    # All should succeed with consistent data
-    for response in responses:
-        assert response.status_code == 200
-        data = response.json()
-        assert data["api_key"] == api_key
-        assert data["balance"] == initial_balance
-
-
 @pytest.mark.integration
 @pytest.mark.asyncio
 async def test_wallet_info_data_consistency(

tests/integration/test_wallet_topup.py

@@ -15,7 +15,6 @@ from routstr.core.db import ApiKey
 
 from .utils import (
     CashuTokenGenerator,
-    ConcurrencyTester,
     ResponseValidator,
 )
 
@@ -284,60 +283,6 @@ async def test_transaction_history_tracking(  # type: ignore[no-untyped-def]
     assert response.status_code == 400
 
 
-@pytest.mark.integration
-@pytest.mark.asyncio
-async def test_concurrent_topups_same_api_key(  # type: ignore[no-untyped-def]
-    integration_client: AsyncClient,
-    authenticated_client: AsyncClient,
-    testmint_wallet: Any,
-) -> None:
-    """Test concurrent top-ups to the same API key"""
-
-    # Get API key
-    response = await authenticated_client.get("/v1/wallet/")
-    api_key = response.json()["api_key"]
-    initial_balance = response.json()["balance"]
-
-    # Generate multiple unique tokens
-    num_tokens = 10
-    tokens = []
-    total_amount = 0
-
-    for i in range(num_tokens):
-        amount = 100 + i * 10  # Different amounts
-        token = await testmint_wallet.mint_tokens(amount)
-        tokens.append(token)
-        total_amount += amount
-
-    # Create concurrent top-up requests
-    requests = [
-        {
-            "method": "POST",
-            "url": "/v1/wallet/topup",
-            "params": {"cashu_token": token},
-            "headers": {"Authorization": f"Bearer {api_key}"},
-        }
-        for token in tokens
-    ]
-
-    # Execute concurrently
-    tester = ConcurrencyTester()
-    responses = await tester.run_concurrent_requests(
-        integration_client, requests, max_concurrent=5
-    )
-
-    # All should succeed
-    for response in responses:
-        assert response.status_code == 200
-        assert "msats" in response.json()
-
-    # Verify final balance is correct
-    final_response = await authenticated_client.get("/v1/wallet/")
-    final_balance = final_response.json()["balance"]
-    expected_balance = initial_balance + (total_amount * 1000)
-    assert final_balance == expected_balance
-
-
 @pytest.mark.integration
 @pytest.mark.asyncio
 async def test_topup_during_active_proxy_request(  # type: ignore[no-untyped-def]