feat: add live data pull for opengradient token

.github/workflows/test.yml

@@ -15,4 +15,12 @@ jobs:
       - uses: actions/checkout@v4
       - uses: astral-sh/setup-uv@v5
       - name: Run unit tests
-        run: uv run --group test pytest tee_gateway/test/test_tool_forwarding.py tee_gateway/test/test_tee_core.py tests/test_pricing.py -v --import-mode=importlib
+        run: uv run --group test pytest tee_gateway/test/test_tool_forwarding.py tee_gateway/test/test_tee_core.py tee_gateway/test/test_util.py tests/test_pricing.py -v --import-mode=importlib
+
+  integration-tests:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v5
+      - name: Run integration tests (live CoinGecko API)
+        run: uv run --group test pytest tests/test_integration.py -v -m integration --import-mode=importlib

pyproject.toml

@@ -11,7 +11,7 @@ dependencies = [
     "setuptools>=21.0.0",
     "Flask>=3.0.0",
     "gunicorn>=23.0.0",
-    "og-x402[evm]==0.0.1.dev6",
+    "og-x402[evm]>=0.0.1.dev9",
     "fastapi>=0.128.0",
     "uvicorn[standard]>=0.40.0",
     "pydantic>=2.12.5",
@@ -58,6 +58,12 @@ exclude = [
     "**/site-packages",
 ]
 
+[tool.pytest.ini_options]
+pythonpath = ["."]
+markers = [
+    "integration: tests that require live network access (deselect with '-m not integration')",
+]
+
 [tool.uv]
 # Pre-release needed for og-test-v2-x402==0.0.11.dev5
 prerelease = "allow"

tee_gateway/__main__.py

@@ -27,19 +27,20 @@
 from x402.http.types import RouteConfig
 from x402.mechanisms.evm.exact import ExactEvmServerScheme
 from x402.mechanisms.evm.upto import UptoEvmServerScheme
+from x402.extensions.erc20_approval_gas_sponsoring import (
+    declare_erc20_approval_gas_sponsoring_extension,
+)
 from x402.schemas import AssetAmount
 from x402.server import x402ResourceServerSync
 from x402.session import SessionStore
-import types as _types
 import x402.http.middleware.flask as x402_flask
 
 from .util import dynamic_session_cost_calculator
 from .definitions import (
-    BASE_TESTNET_NETWORK,
     EVM_PAYMENT_ADDRESS,
-    BASE_OPG_ADDRESS,
+    BASE_MAINNET_NETWORK,
+    BASE_MAINNET_OPG_ADDRESS,
     CHAT_COMPLETIONS_OPG_SESSION_MAX_SPEND,
-    COMPLETIONS_OPG_SESSION_MAX_SPEND,
     FACILITATOR_URL,
 )
 
@@ -108,8 +109,10 @@ def _shutdown_heartbeat():
 server = x402ResourceServerSync(facilitator)
 store = SessionStore()
 
-server.register(BASE_TESTNET_NETWORK, ExactEvmServerScheme())
-server.register(BASE_TESTNET_NETWORK, UptoEvmServerScheme())
+server.register(BASE_MAINNET_NETWORK, ExactEvmServerScheme())
+
+# Upto scheme registrations (permit2-based, variable settlement)
+server.register(BASE_MAINNET_NETWORK, UptoEvmServerScheme())
 
 routes = {
     "POST /v1/chat/completions": RouteConfig(
@@ -119,16 +122,19 @@ def _shutdown_heartbeat():
                 pay_to=EVM_PAYMENT_ADDRESS,
                 price=AssetAmount(
                     amount=CHAT_COMPLETIONS_OPG_SESSION_MAX_SPEND,
-                    asset=BASE_OPG_ADDRESS,
+                    asset=BASE_MAINNET_OPG_ADDRESS,
                     extra={
-                        "name": "OPG",
-                        "version": "2",
+                        "name": "OpenGradient",
+                        "version": "1",
                         "assetTransferMethod": "permit2",
                     },
                 ),
-                network=BASE_TESTNET_NETWORK,
+                network=BASE_MAINNET_NETWORK,
             ),
         ],
+        extensions={
+            **declare_erc20_approval_gas_sponsoring_extension(),
+        },
         mime_type="application/json",
         description="Chat completion",
     ),
@@ -138,17 +144,20 @@ def _shutdown_heartbeat():
                 scheme="upto",
                 pay_to=EVM_PAYMENT_ADDRESS,
                 price=AssetAmount(
-                    amount=COMPLETIONS_OPG_SESSION_MAX_SPEND,
-                    asset=BASE_OPG_ADDRESS,
+                    amount=CHAT_COMPLETIONS_OPG_SESSION_MAX_SPEND,
+                    asset=BASE_MAINNET_OPG_ADDRESS,
                     extra={
-                        "name": "OPG",
-                        "version": "2",
+                        "name": "OpenGradient",
+                        "version": "1",
                         "assetTransferMethod": "permit2",
                     },
                 ),
-                network=BASE_TESTNET_NETWORK,
+                network=BASE_MAINNET_NETWORK,
             ),
         ],
+        extensions={
+            **declare_erc20_approval_gas_sponsoring_extension(),
+        },
         mime_type="application/json",
         description="Completion",
     ),
@@ -374,99 +383,21 @@ def _patched_read_body_bytes(environ):
 )
 
 # ---------------------------------------------------------------------------
-# Strict cost-resolution patch
+# Cost-resolution behaviour note
 #
-# Why this exists
-# ---------------
-# The upstream x402 PaymentMiddleware._resolve_session_request_cost wraps the
-# call to the session_cost_calculator in a broad try/except.  If the calculator
-# raises (e.g. ValueError for an unrecognised model name, KeyError for missing
-# usage data), the exception is swallowed and the middleware silently falls back
-# to the static session maximum (CHAT_COMPLETIONS_OPG_SESSION_MAX_SPEND /
-# CHAT_COMPLETIONS_USDC_AMOUNT).  That silent fallback means:
-#   • The client is charged the full pre-check cap instead of actual usage.
-#   • The server has no visible indication that pricing failed.
+# dynamic_session_cost_calculator() is invoked by PaymentMiddleware via
+# _accumulate_session_cost(), which is itself called from
+# StreamingSessionResponse.close().  That close() wraps the call in a broad
+# try/except, so any exception raised by the calculator (e.g. ValueError for
+# an unrecognised model or missing usage data) is logged but otherwise
+# swallowed — the session cost is simply not incremented for that request.
 #
-# The fix
-# -------
-# We replace _resolve_session_request_cost with our own implementation that is
-# identical to upstream, except the cost-calculator call is NOT wrapped in a
-# try/except.  Any exception from dynamic_session_cost_calculator() therefore
-# propagates up through the middleware and Flask, producing a proper HTTP 500
-# response to the client instead of an incorrect silent charge.
+# A previous version of this file monkey-patched _resolve_session_request_cost
+# to let exceptions propagate, but that method no longer exists in the library.
+# If stricter error propagation is needed in future, patch _accumulate_session_cost
+# or StreamingSessionResponse.close() instead.
 # ---------------------------------------------------------------------------
 
-
-def _strict_resolve_session_request_cost(
-    self,
-    *,
-    method: str,
-    path: str,
-    request_body_bytes: bytes,
-    response_body_bytes: bytes,
-    payment_payload: object,
-    payment_requirements: object,
-    status_code: int | None,
-    output_object: object = None,
-    is_streaming: bool = False,
-) -> int:
-    """Replacement for PaymentMiddleware._resolve_session_request_cost.
-
-    Identical to the upstream implementation except that exceptions raised by
-    the dynamic cost calculator are NOT caught.  This means a request whose
-    cost cannot be determined (unknown model, missing usage data, etc.) will
-    result in a 500 error rather than silently falling back to the static cap
-    amount and charging the user an incorrect amount.
-    """
-    from x402.http.middleware.flask import _parse_json_bytes as _x402_parse_json  # noqa: PLC0415
-
-    default_cost = self._get_session_cost(payment_requirements)
-    if not self._should_charge_response(status_code):
-        return default_cost
-    if not callable(self._session_cost_calculator):
-        return default_cost
-
-    request_object = _x402_parse_json(request_body_bytes)
-    response_object = (
-        output_object
-        if output_object is not None
-        else _x402_parse_json(response_body_bytes)
-    )
-
-    callback_context = {
-        "method": method,
-        "path": path,
-        "status_code": status_code,
-        "is_streaming": is_streaming,
-        "request_body_bytes": request_body_bytes,
-        "response_body_bytes": response_body_bytes,
-        "request_json": request_object
-        if isinstance(request_object, (dict, list))
-        else None,
-        "response_json": response_object
-        if isinstance(response_object, (dict, list))
-        else None,
-        "response_object": response_object,
-        "payment_payload": payment_payload,
-        "payment_requirements": payment_requirements,
-        "default_cost": default_cost,
-    }
-
-    # Do NOT catch exceptions here — let them propagate so the request fails
-    # with a 500 rather than silently charging the static fallback amount.
-    dynamic_cost = self._session_cost_calculator(callback_context)
-    if dynamic_cost is None:
-        raise ValueError(
-            f"dynamic_session_cost_calculator returned None for {method} {path}; "
-            "cannot determine request cost"
-        )
-    return self._coerce_non_negative_int(dynamic_cost)
-
-
-_payment_mw._resolve_session_request_cost = _types.MethodType(  # type: ignore[method-assign]
-    _strict_resolve_session_request_cost, _payment_mw
-)
-
 logger.info("x402 payment middleware initialized")
 
 if __name__ == "__main__":

tee_gateway/config.py

@@ -9,7 +9,34 @@
 from dataclasses import dataclass
 from typing import Optional
 
+# ---------------------------------------------------------------------------
+# OPG / token price feed
+# ---------------------------------------------------------------------------
+
+# How long (seconds) to reuse a cached price before fetching a fresh one.
+# At 120 s the gateway makes at most 30 CoinGecko calls/hour — well within
+# the free-tier limit (30/min).
+OPG_PRICE_CACHE_TTL_SECONDS: int = 120
+
+# Number of times to retry a failed CoinGecko fetch before giving up.
+# Each attempt uses the same 5-second timeout; retries are immediate (no backoff).
+OPG_PRICE_FETCH_RETRIES: int = 3
+
+# CoinGecko coin ID for the OPG token.
+# https://www.coingecko.com/en/coins/opengradient
+OPG_PRICE_COINGECKO_ID: str = "opengradient"
+
+# Sanity bounds for the fetched token price.
+# Used in integration tests to catch obviously wrong API responses
+# (wrong currency, implausibly large value).
+# Update when OPG establishes a trading range.
+OPG_PRICE_SANITY_MAX_USD: str = (
+    "1000000"  # $1 000 000 — rules out obviously corrupt data
+)
+
+# ---------------------------------------------------------------------------
 # Heartbeat defaults
+# ---------------------------------------------------------------------------
 DEFAULT_HEARTBEAT_INTERVAL = 900  # 15 minutes
 DEFAULT_HEARTBEAT_BUFFER = (
     300  # 5 minutes — subtracted from time.time() to compensate for enclave clock drift

tee_gateway/controllers/chat_controller.py

@@ -30,6 +30,7 @@
     convert_messages,
     extract_usage,
 )
+from tee_gateway.util import validate_pricing_preflight
 
 logger = logging.getLogger(__name__)
 
@@ -46,6 +47,12 @@ def create_chat_completion(body):
         connexion.request.get_json()
     )
 
+    try:
+        validate_pricing_preflight(chat_request.model)
+    except ValueError as exc:
+        logger.error("Pricing preflight failed for model %r: %s", chat_request.model, exc)
+        return {"error": "Bad Request", "message": str(exc)}, 400
+
     if chat_request.stream:
         return _create_streaming_response(chat_request)
     else:

tee_gateway/controllers/completions_controller.py

@@ -10,6 +10,7 @@
 
 from tee_gateway.tee_manager import get_tee_keys, compute_tee_msg_hash
 from tee_gateway.llm_backend import get_chat_model_cached, extract_usage
+from tee_gateway.util import validate_pricing_preflight
 
 logger = logging.getLogger(__name__)
 
@@ -21,6 +22,12 @@ def create_completion(body):
     else:
         return {"error": "Request must be application/json"}, 415
 
+    try:
+        validate_pricing_preflight(body.model)
+    except ValueError as exc:
+        logger.error("Pricing preflight failed for model %r: %s", body.model, exc)
+        return {"error": "Bad Request", "message": str(exc)}, 400
+
     try:
         request_dict = {
             "model": body.model,

tee_gateway/definitions.py

@@ -19,8 +19,9 @@
 # Network IDs (EIP-155 chain identifiers)
 # ---------------------------------------------------------------------------
 
-# Base Testnet — where OPG payments are accepted
-BASE_TESTNET_NETWORK: str = "eip155:84532"
+
+# Base Mainnet — where OPG payments are accepted
+BASE_MAINNET_NETWORK: str = "eip155:8453"
 
 # ---------------------------------------------------------------------------
 # Payment recipient
@@ -31,23 +32,23 @@
 # your own instance.
 EVM_PAYMENT_ADDRESS: str = os.getenv(
     "EVM_PAYMENT_ADDRESS",
-    "0x40eFb45552EDfB2502D90A657a8ab41F03ec460d",
+    "0x9deEBB5D1b22e4a6e027977CeAd13893A7E4cC1a",
 )
 
 # ---------------------------------------------------------------------------
 # ERC-20 token contract addresses
 # ---------------------------------------------------------------------------
 
-# OpenGradient token (OPG) on Base Testnet
-BASE_OPG_ADDRESS: str = "0x240b09731D96979f50B2C649C9CE10FcF9C7987F"
+# OpenGradient token (OPG) on Base Mainnet
+BASE_MAINNET_OPG_ADDRESS: str = "0xFbC2051AE2265686a469421b2C5A2D5462FbF5eB"
 
 # ---------------------------------------------------------------------------
 # Token decimal places
 # ---------------------------------------------------------------------------
 
 # Maps lowercase contract address → number of decimals for unit conversion.
 ASSET_DECIMALS_BY_ADDRESS: dict[str, int] = {
-    BASE_OPG_ADDRESS.lower(): 18,  # OPG: 18 decimals (ERC-20 standard)
+    BASE_MAINNET_OPG_ADDRESS.lower(): 18,  # OPG: 18 decimals (ERC-20 standard)
 }
 
 # Fallback for any asset not explicitly listed above

tee_gateway/test/test_tool_forwarding.py

@@ -1,4 +1,5 @@
 import unittest
+from decimal import Decimal
 from unittest.mock import patch, Mock
 
 from tee_gateway.controllers.chat_controller import (
@@ -11,6 +12,21 @@
     ChatCompletionRequestFunctionMessage,
 )
 
+# Pin the token price for all integration-style tests in this file so they
+# never hit the real CoinGecko API.  The price feed is tested separately in
+# test_util.py and test_integration.py.
+_price_patcher = patch(
+    "tee_gateway.util.get_token_a_price_usd", return_value=Decimal("1")
+)
+
+
+def setUpModule():
+    _price_patcher.start()
+
+
+def tearDownModule():
+    _price_patcher.stop()
+
 
 # ---------------------------------------------------------------------------
 # Shared helpers