From d6d512b6cdd60944849e3ca77bfbd17da1655159 Mon Sep 17 00:00:00 2001 From: kylexqian Date: Mon, 20 Apr 2026 12:50:15 -0700 Subject: [PATCH 1/6] Defer x402 middleware setup to injection time; unify facilitator_url MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously the x402 payment middleware (facilitator, server, routes, session store) was created at module load using a hardcoded FACILITATOR_URL, which meant the URL couldn't be injected at runtime and the middleware was never actually wired into Flask's request chain (application.run() bypassed _payment_mw entirely). Changes: - Move all x402 setup into _init_payment_middleware(facilitator_url), called once from set_provider_keys(). Uses application.wsgi_app = mw (the standard Flask WSGI middleware pattern) so all requests flow through payment checking after injection. - Accept a single `facilitator_url` field in POST /v1/keys, used for both x402 payment verification and the heartbeat relay. Removes the separate `heartbeat_facilitator_url` field. - Fallback chain: injection payload → FACILITATOR_URL env var → definitions.py hardcoded default. - Update run-enclave.sh: HEARTBEAT_FACILITATOR_URL → FACILITATOR_URL, heartbeat_facilitator_url JSON key → facilitator_url. - Update definitions.py comment to document the full precedence chain. Co-Authored-By: Claude Sonnet 4.6 --- scripts/run-enclave.sh | 14 +- tee_gateway/__main__.py | 255 ++++++++++++++++++++----------------- tee_gateway/definitions.py | 4 + 3 files changed, 148 insertions(+), 125 deletions(-) diff --git a/scripts/run-enclave.sh b/scripts/run-enclave.sh index b966368..d4aac35 100755 --- a/scripts/run-enclave.sh +++ b/scripts/run-enclave.sh @@ -90,9 +90,11 @@ if [ -f "$ENV_FILE" ]; then ANTHROPIC_API_KEY="$(grep -E '^ANTHROPIC_API_KEY=' "$ENV_FILE" | cut -d'=' -f2-)" XAI_API_KEY="$(grep -E '^XAI_API_KEY=' "$ENV_FILE" | cut -d'=' -f2-)" - # Heartbeat configuration (optional — wallet key is generated inside the TEE) + # FACILITATOR_URL is used for both x402 payment verification and the heartbeat relay. + # HEARTBEAT_CONTRACT_ADDRESS and TEE_HEARTBEAT_INTERVAL are optional heartbeat parameters. + # The TEE wallet key is generated inside the enclave and never injected. HEARTBEAT_CONTRACT_ADDRESS="$(grep -E '^HEARTBEAT_CONTRACT_ADDRESS=' "$ENV_FILE" | cut -d'=' -f2-)" - HEARTBEAT_FACILITATOR_URL="$(grep -E '^HEARTBEAT_FACILITATOR_URL=' "$ENV_FILE" | cut -d'=' -f2-)" + FACILITATOR_URL="$(grep -E '^FACILITATOR_URL=' "$ENV_FILE" | cut -d'=' -f2-)" TEE_HEARTBEAT_INTERVAL="$(grep -E '^TEE_HEARTBEAT_INTERVAL=' "$ENV_FILE" | cut -d'=' -f2-)" # Build the JSON payload using jq for safe escaping @@ -103,7 +105,7 @@ if [ -f "$ENV_FILE" ]; then --arg anthropic "$ANTHROPIC_API_KEY" \ --arg xai "$XAI_API_KEY" \ --arg hb_contract "$HEARTBEAT_CONTRACT_ADDRESS" \ - --arg hb_facilitator "$HEARTBEAT_FACILITATOR_URL" \ + --arg facilitator "$FACILITATOR_URL" \ --arg hb_interval "$TEE_HEARTBEAT_INTERVAL" \ '{ openai_api_key: $openai, @@ -112,7 +114,7 @@ if [ -f "$ENV_FILE" ]; then xai_api_key: $xai } + if $hb_contract != "" then {heartbeat_contract_address: $hb_contract} else {} end - + if $hb_facilitator != "" then {heartbeat_facilitator_url: $hb_facilitator} else {} end + + if $facilitator != "" then {facilitator_url: $facilitator} else {} end + if $hb_interval != "" then {tee_heartbeat_interval: $hb_interval} else {} end ') @@ -125,7 +127,7 @@ if [ -f "$ENV_FILE" ]; then if [ "$http_status" = "200" ]; then echo "[ec2] API keys injected successfully." - if [ -n "$HEARTBEAT_CONTRACT_ADDRESS" ] && [ -n "$HEARTBEAT_FACILITATOR_URL" ]; then + if [ -n "$HEARTBEAT_CONTRACT_ADDRESS" ] && [ -n "$FACILITATOR_URL" ]; then echo "[ec2] Heartbeat service configured via facilitator relay (contract: ${HEARTBEAT_CONTRACT_ADDRESS})" else echo "[ec2] Heartbeat service not configured (missing env vars)." @@ -136,7 +138,7 @@ if [ -f "$ENV_FILE" ]; then # Clear key variables from this shell immediately after use unset OPENAI_API_KEY GOOGLE_API_KEY ANTHROPIC_API_KEY XAI_API_KEY - unset HEARTBEAT_CONTRACT_ADDRESS HEARTBEAT_FACILITATOR_URL TEE_HEARTBEAT_INTERVAL + unset HEARTBEAT_CONTRACT_ADDRESS FACILITATOR_URL TEE_HEARTBEAT_INTERVAL fi else echo "[ec2] No .env file found at $ENV_FILE" diff --git a/tee_gateway/__main__.py b/tee_gateway/__main__.py index 7ad75d4..7692949 100644 --- a/tee_gateway/__main__.py +++ b/tee_gateway/__main__.py @@ -107,6 +107,7 @@ def _shutdown_heartbeat(): atexit.register(_shutdown_heartbeat) + # --------------------------------------------------------------------------- # OPG price feed — start before x402 middleware so the first request can be # priced correctly. Runs as a daemon thread; no cleanup needed on exit. @@ -114,76 +115,142 @@ def _shutdown_heartbeat(): _price_feed = OPGPriceFeed() _price_feed.start() -facilitator = HTTPFacilitatorClientSync(FacilitatorConfig(url=FACILITATOR_URL)) -server = x402ResourceServerSync(facilitator) -store = SessionStore() - -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( - accepts=[ - PaymentOption( - scheme="upto", - pay_to=EVM_PAYMENT_ADDRESS, - price=AssetAmount( - amount=CHAT_COMPLETIONS_OPG_SESSION_MAX_SPEND, - asset=BASE_MAINNET_OPG_ADDRESS, - extra={ - "name": "OpenGradient", - "version": "1", - "assetTransferMethod": "permit2", - }, - ), - network=BASE_MAINNET_NETWORK, - ), - ], - extensions={ - **declare_erc20_approval_gas_sponsoring_extension(), - }, - mime_type="application/json", - description="Chat completion", - ), - "POST /v1/completions": RouteConfig( - accepts=[ - PaymentOption( - scheme="upto", - pay_to=EVM_PAYMENT_ADDRESS, - price=AssetAmount( - amount=CHAT_COMPLETIONS_OPG_SESSION_MAX_SPEND, - asset=BASE_MAINNET_OPG_ADDRESS, - extra={ - "name": "OpenGradient", - "version": "1", - "assetTransferMethod": "permit2", - }, - ), - network=BASE_MAINNET_NETWORK, - ), - ], - extensions={ - **declare_erc20_approval_gas_sponsoring_extension(), - }, - mime_type="application/json", - description="Completion", - ), -} + +# --------------------------------------------------------------------------- +# x402 read-body patch +# +# Ensures that non-payment 0-length requests can bypass the middleware without +# errors. Applied at module load so it is in place before the middleware +# instance is created at injection time. +# --------------------------------------------------------------------------- +_original_read_body_bytes = x402_flask._read_body_bytes + + +def _patched_read_body_bytes(environ): + try: + content_length = int(environ.get("CONTENT_LENGTH") or 0) + except (ValueError, TypeError): + content_length = 0 + + if content_length <= 0: + return b"" + + return _original_read_body_bytes(environ) + + +x402_flask._read_body_bytes = _patched_read_body_bytes + + +def _session_cost_calculator(ctx: dict) -> int: + # Post-inference cost calculation — response already sent to client. + # Predictable failures (unknown price, unknown model) are blocked by the + # pre-inference gate; any exception here indicates a provider-side error + # (e.g. missing usage field in the LLM response). The x402 middleware + # swallows the exception in close(), so the client is not charged. + # Log CRITICAL so provider errors are never silently missed. + try: + return calculate_session_cost(ctx, _price_feed.get_price) + except Exception as exc: + logger.critical( + "Post-inference cost calculation failed (provider error) — " + "client was NOT charged: %s", + exc, + exc_info=True, + ) + raise # --------------------------------------------------------------------------- -# One-time provider key injection +# One-time runtime configuration injection # --------------------------------------------------------------------------- _keys_initialized: bool = False _keys_lock = threading.Lock() +def _init_payment_middleware(facilitator_url: str) -> None: + """Build and attach the x402 payment middleware to the running Flask app. + + Called once from set_provider_keys() after the facilitator URL is known. + Swaps application.wsgi_app so all subsequent requests flow through it. + """ + facilitator = HTTPFacilitatorClientSync(FacilitatorConfig(url=facilitator_url)) + server = x402ResourceServerSync(facilitator) + store = SessionStore() + + server.register(BASE_MAINNET_NETWORK, ExactEvmServerScheme()) + server.register(BASE_MAINNET_NETWORK, UptoEvmServerScheme()) + + routes = { + "POST /v1/chat/completions": RouteConfig( + accepts=[ + PaymentOption( + scheme="upto", + pay_to=EVM_PAYMENT_ADDRESS, + price=AssetAmount( + amount=CHAT_COMPLETIONS_OPG_SESSION_MAX_SPEND, + asset=BASE_MAINNET_OPG_ADDRESS, + extra={ + "name": "OpenGradient", + "version": "1", + "assetTransferMethod": "permit2", + }, + ), + network=BASE_MAINNET_NETWORK, + ), + ], + extensions={ + **declare_erc20_approval_gas_sponsoring_extension(), + }, + mime_type="application/json", + description="Chat completion", + ), + "POST /v1/completions": RouteConfig( + accepts=[ + PaymentOption( + scheme="upto", + pay_to=EVM_PAYMENT_ADDRESS, + price=AssetAmount( + amount=CHAT_COMPLETIONS_OPG_SESSION_MAX_SPEND, + asset=BASE_MAINNET_OPG_ADDRESS, + extra={ + "name": "OpenGradient", + "version": "1", + "assetTransferMethod": "permit2", + }, + ), + network=BASE_MAINNET_NETWORK, + ), + ], + extensions={ + **declare_erc20_approval_gas_sponsoring_extension(), + }, + mime_type="application/json", + description="Completion", + ), + } + + mw = payment_middleware( + application.wsgi_app, + routes=routes, + server=server, + session_store=store, + cost_per_request=100000000000000, # static precheck/fallback estimate + session_idle_timeout=100, + session_cost_calculator=_session_cost_calculator, + ) + application.wsgi_app = mw + logger.info( + "x402 payment middleware initialized with facilitator: %s", facilitator_url + ) + + def set_provider_keys(): """ - POST /v1/keys — inject LLM provider API keys into the enclave. - Can only be called once; subsequent calls return HTTP 409. + POST /v1/keys — inject runtime configuration into the enclave. + + Accepts LLM provider API keys, a shared facilitator_url (used for both + x402 payment verification and the heartbeat relay), and optional heartbeat + parameters. Can only be called once; subsequent calls return HTTP 409. """ global _keys_initialized @@ -206,15 +273,16 @@ def set_provider_keys(): ) set_provider_config(provider_config) - # Build heartbeat config from request body (optional) - contract_address = body.get("heartbeat_contract_address") facilitator_url = ( - body.get("heartbeat_facilitator_url") + body.get("facilitator_url") or os.getenv("FACILITATOR_URL") or FACILITATOR_URL ) + + # Build heartbeat config from request body (optional) + contract_address = body.get("heartbeat_contract_address") heartbeat_config: HeartbeatConfig | None = None - if contract_address and facilitator_url: + if contract_address: interval_raw = body.get( "tee_heartbeat_interval", DEFAULT_HEARTBEAT_INTERVAL ) @@ -261,14 +329,11 @@ def _set(val: str | None) -> str: logger.info( " xai_api_key : %s", _set(provider_config.xai_api_key) ) + logger.info(" facilitator_url : %s", facilitator_url) logger.info( " heartbeat_contract_address : %s", _set(heartbeat_config.contract_address if heartbeat_config else None), ) - logger.info( - " heartbeat_facilitator_url : %s", - _set(heartbeat_config.facilitator_url if heartbeat_config else None), - ) logger.info( " tee_heartbeat_interval : %s", heartbeat_config.interval if heartbeat_config else "900 (default)", @@ -283,6 +348,8 @@ def _set(val: str | None) -> str: except Exception as e: logger.warning(f"Heartbeat initialization failed: {e}") + _init_payment_middleware(facilitator_url) + _keys_initialized = True providers_set = [ @@ -358,59 +425,11 @@ def create_app(): # --------------------------------------------------------------------------- -# WSGI application + x402 payment middleware +# WSGI application # --------------------------------------------------------------------------- -# Create the WSGI application application = create_app() -# This patch ensures that non-payment 0-length requests can still bypass the middleware -_original_read_body_bytes = x402_flask._read_body_bytes - - -def _patched_read_body_bytes(environ): - try: - content_length = int(environ.get("CONTENT_LENGTH") or 0) - except (ValueError, TypeError): - content_length = 0 - - if content_length <= 0: - return b"" - - return _original_read_body_bytes(environ) - - -x402_flask._read_body_bytes = _patched_read_body_bytes - - -def _session_cost_calculator(ctx: dict) -> int: - # Post-inference cost calculation — response already sent to client. - # Predictable failures (unknown price, unknown model) are blocked by the - # pre-inference gate; any exception here indicates a provider-side error - # (e.g. missing usage field in the LLM response). The x402 middleware - # swallows the exception in close(), so the client is not charged. - # Log CRITICAL so provider errors are never silently missed. - try: - return calculate_session_cost(ctx, _price_feed.get_price) - except Exception as exc: - logger.critical( - "Post-inference cost calculation failed (provider error) — " - "client was NOT charged: %s", - exc, - exc_info=True, - ) - raise - - -_payment_mw = payment_middleware( - application, - routes=routes, - server=server, - session_store=store, - cost_per_request=100000000000000, # static precheck/fallback estimate - session_idle_timeout=100, - session_cost_calculator=_session_cost_calculator, -) # --------------------------------------------------------------------------- # Pre-inference pricing gate @@ -443,8 +462,6 @@ def _check_pricing_ready(): return jsonify({"error": f"Model '{model}' is not supported"}), 400 -logger.info("x402 payment middleware initialized") - if __name__ == "__main__": port = int(os.getenv("API_SERVER_PORT", "8000")) host = os.getenv("API_SERVER_HOST", "0.0.0.0") diff --git a/tee_gateway/definitions.py b/tee_gateway/definitions.py index ca316bd..932f9eb 100644 --- a/tee_gateway/definitions.py +++ b/tee_gateway/definitions.py @@ -13,6 +13,10 @@ # --------------------------------------------------------------------------- # X402 Facilitator # --------------------------------------------------------------------------- +# Default fallback only. The live value is injected at runtime via POST /v1/keys +# (facilitator_url field) and used for both x402 payment verification and the +# heartbeat service. Override at the OS level with the FACILITATOR_URL env var, +# or supply it directly in the injection payload. FACILITATOR_URL = os.getenv("FACILITATOR_URL", "https://facilitator.memchat.io") # --------------------------------------------------------------------------- From 9be95176bd3cd0a8cdb36b28130a5ba7d5395f86 Mon Sep 17 00:00:00 2001 From: kylexqian Date: Mon, 20 Apr 2026 14:39:28 -0700 Subject: [PATCH 2/6] Remove redundant os.getenv in facilitator_url fallback FACILITATOR_URL in definitions.py already reads from os.getenv, so the middle step was always superseded by the final fallback. Co-Authored-By: Claude Sonnet 4.6 --- tee_gateway/__main__.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/tee_gateway/__main__.py b/tee_gateway/__main__.py index 7692949..f179c26 100644 --- a/tee_gateway/__main__.py +++ b/tee_gateway/__main__.py @@ -273,11 +273,7 @@ def set_provider_keys(): ) set_provider_config(provider_config) - facilitator_url = ( - body.get("facilitator_url") - or os.getenv("FACILITATOR_URL") - or FACILITATOR_URL - ) + facilitator_url = body.get("facilitator_url") or FACILITATOR_URL # Build heartbeat config from request body (optional) contract_address = body.get("heartbeat_contract_address") From db6784705c70f5321bf8acedd6d1378f7ccae9df Mon Sep 17 00:00:00 2001 From: kylexqian Date: Mon, 20 Apr 2026 15:46:12 -0700 Subject: [PATCH 3/6] Fix payment_middleware receiving wsgi_app function instead of Flask app PaymentMiddleware.__init__ does app.wsgi_app internally, so it needs the full Flask application object. Passing application.wsgi_app (a plain function) caused AttributeError. payment_middleware captures the inner wsgi_app by value at creation time, so setting application.wsgi_app = mw afterwards is safe and does not create a circular reference. Co-Authored-By: Claude Sonnet 4.6 --- tee_gateway/__main__.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tee_gateway/__main__.py b/tee_gateway/__main__.py index f179c26..cb6d1a6 100644 --- a/tee_gateway/__main__.py +++ b/tee_gateway/__main__.py @@ -230,7 +230,7 @@ def _init_payment_middleware(facilitator_url: str) -> None: } mw = payment_middleware( - application.wsgi_app, + application, routes=routes, server=server, session_store=store, @@ -238,6 +238,8 @@ def _init_payment_middleware(facilitator_url: str) -> None: session_idle_timeout=100, session_cost_calculator=_session_cost_calculator, ) + # payment_middleware captures application.wsgi_app by value at creation time, + # so assigning mw back here does not create a circular reference. application.wsgi_app = mw logger.info( "x402 payment middleware initialized with facilitator: %s", facilitator_url From b7058ea06b31dbe7a6c27df6ec65e334ba5a84b6 Mon Sep 17 00:00:00 2001 From: kylexqian Date: Mon, 20 Apr 2026 15:56:51 -0700 Subject: [PATCH 4/6] =?UTF-8?q?Remove=20manual=20wsgi=5Fapp=20assignment?= =?UTF-8?q?=20=E2=80=94=20PaymentMiddleware=20self-wires=20on=20init?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PaymentMiddleware.__init__ already does app.wsgi_app = self._wsgi_middleware internally. Our manual application.wsgi_app = mw was overwriting that bound method with the bare PaymentMiddleware instance (which has no __call__), causing TypeError on every request. Co-Authored-By: Claude Sonnet 4.6 --- tee_gateway/__main__.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/tee_gateway/__main__.py b/tee_gateway/__main__.py index cb6d1a6..1e8782d 100644 --- a/tee_gateway/__main__.py +++ b/tee_gateway/__main__.py @@ -229,7 +229,7 @@ def _init_payment_middleware(facilitator_url: str) -> None: ), } - mw = payment_middleware( + payment_middleware( application, routes=routes, server=server, @@ -238,9 +238,6 @@ def _init_payment_middleware(facilitator_url: str) -> None: session_idle_timeout=100, session_cost_calculator=_session_cost_calculator, ) - # payment_middleware captures application.wsgi_app by value at creation time, - # so assigning mw back here does not create a circular reference. - application.wsgi_app = mw logger.info( "x402 payment middleware initialized with facilitator: %s", facilitator_url ) From 8e5914cb342f0f3fd5b14eea7ba77826170acef7 Mon Sep 17 00:00:00 2001 From: kylexqian Date: Mon, 20 Apr 2026 16:59:53 -0700 Subject: [PATCH 5/6] Address Copilot review feedback - Add comment explaining why payment_middleware return value is discarded (PaymentMiddleware self-wires via app.wsgi_app in __init__) - Fix run-enclave.sh heartbeat status message: condition on HEARTBEAT_CONTRACT_ADDRESS alone, with a nested check for FACILITATOR_URL to distinguish injected vs enclave-default URL Co-Authored-By: Claude Sonnet 4.6 --- scripts/run-enclave.sh | 10 +++++++--- tee_gateway/__main__.py | 2 ++ 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/scripts/run-enclave.sh b/scripts/run-enclave.sh index d4aac35..4a05030 100755 --- a/scripts/run-enclave.sh +++ b/scripts/run-enclave.sh @@ -127,10 +127,14 @@ if [ -f "$ENV_FILE" ]; then if [ "$http_status" = "200" ]; then echo "[ec2] API keys injected successfully." - if [ -n "$HEARTBEAT_CONTRACT_ADDRESS" ] && [ -n "$FACILITATOR_URL" ]; then - echo "[ec2] Heartbeat service configured via facilitator relay (contract: ${HEARTBEAT_CONTRACT_ADDRESS})" + if [ -n "$HEARTBEAT_CONTRACT_ADDRESS" ]; then + if [ -n "$FACILITATOR_URL" ]; then + echo "[ec2] Heartbeat service configured via facilitator relay (contract: ${HEARTBEAT_CONTRACT_ADDRESS})" + else + echo "[ec2] Heartbeat service configured using enclave default facilitator URL (contract: ${HEARTBEAT_CONTRACT_ADDRESS})" + fi else - echo "[ec2] Heartbeat service not configured (missing env vars)." + echo "[ec2] Heartbeat service not configured (missing HEARTBEAT_CONTRACT_ADDRESS)." fi else echo "[ec2] Warning: Key injection returned HTTP $http_status. Check enclave logs." diff --git a/tee_gateway/__main__.py b/tee_gateway/__main__.py index 1e8782d..a980fac 100644 --- a/tee_gateway/__main__.py +++ b/tee_gateway/__main__.py @@ -229,6 +229,8 @@ def _init_payment_middleware(facilitator_url: str) -> None: ), } + # Return value intentionally discarded — PaymentMiddleware.__init__ self-wires + # by setting application.wsgi_app = self._wsgi_middleware internally. payment_middleware( application, routes=routes, From 2789a8195c5dbd48e99679729090e6a8e7e6b025 Mon Sep 17 00:00:00 2001 From: kylexqian Date: Mon, 20 Apr 2026 17:02:51 -0700 Subject: [PATCH 6/6] Change .env.example --- .env.example | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.env.example b/.env.example index f29ee6d..a372744 100644 --- a/.env.example +++ b/.env.example @@ -13,5 +13,5 @@ XAI_API_KEY= # If HEARTBEAT_CONTRACT_ADDRESS and HEARTBEAT_FACILITATOR_URL are set, the enclave # signs heartbeat payloads and the facilitator relays on-chain txs. HEARTBEAT_CONTRACT_ADDRESS= -HEARTBEAT_FACILITATOR_URL= +FACILITATOR_URL= TEE_HEARTBEAT_INTERVAL=900