Skip to content

batch settlement for tee gateway #131

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
dixitaniket wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

batch settlement for tee gateway #131

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ dependencies = [
"setuptools>=21.0.0",
"Flask>=3.0.0",
"gunicorn>=23.0.0",
"og-x402[evm]==0.0.2.dev8",
"og-x402[evm,flask,extensions]==2.15.0",
"fastapi>=0.128.0",
"uvicorn[standard]>=0.40.0",
"pydantic>=2.12.5",
Expand Down
201 changes: 96 additions & 105 deletions tee_gateway/__main__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,22 +24,28 @@
from tee_gateway.llm_backend import get_provider_config, set_provider_config
from tee_gateway.heartbeat import create_heartbeat_service
from tee_gateway.controllers.ohttp_controller import (
OHTTP_BATCH_SETTLEMENT_BOUNDARY,
build_ohttp_batch_settlement_frame,
create_anonymous_chat_completion,
get_hpke_config,
)

from x402.http import FacilitatorConfig, HTTPFacilitatorClientSync, PaymentOption
from x402.http.middleware.flask import payment_middleware
from x402.http.types import RouteConfig
from x402.mechanisms.evm.exact import ExactEvmServerScheme
from x402.mechanisms.evm.upto import UptoEvmServerScheme
from x402.mechanisms.evm.batch_settlement.server import (
AutoSettlementConfig,
BatchSettlementChannelManagerSync,
BatchSettlementEvmScheme,
BatchSettlementEvmSchemeServerConfig,
ChannelManagerConfigSync,
FileChannelStorage,
)
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 x402.http.middleware.flask as x402_flask

from .model_registry import get_model_config
from .price_feed import OPGPriceFeed, set_price_feed
Expand Down Expand Up @@ -89,6 +95,18 @@
# ---------------------------------------------------------------------------
_heartbeat_service = None

# The gateway has to retain the latest signed voucher for every channel in
# order to claim it later. This directory must be mounted on persistent
# storage in production; an in-memory store would lose claimable vouchers on
# a gateway restart.
_batch_settlement_storage_dir = os.path.expanduser(
os.getenv(
"X402_BATCH_SETTLEMENT_STORAGE_DIR",
"~/.tee-gateway/x402-batch-channels",
)
)
_batch_settlement_manager = None


def _init_heartbeat(heartbeat_config: HeartbeatConfig | None):
"""Create and start the heartbeat service if a HeartbeatConfig is provided."""
Expand All @@ -115,6 +133,21 @@ def _shutdown_heartbeat():
atexit.register(_shutdown_heartbeat)


def _shutdown_batch_settlement_manager():
"""Flush claimable vouchers before the gateway process exits."""
manager = _batch_settlement_manager
if manager is None:
return
try:
logger.info("Shutting down batch-settlement manager; flushing claims")
manager.stop(flush=True)
except Exception:
logger.exception("Failed to flush batch-settlement claims during shutdown")


atexit.register(_shutdown_batch_settlement_manager)


# ---------------------------------------------------------------------------
# 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.
Expand All @@ -136,102 +169,16 @@ def _gateway_version() -> str:
_active_facilitator_url: str | None = None


# ---------------------------------------------------------------------------
# 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 _patched_stream_session_response(
self,
environ,
start_response,
context,
session_id,
payment_payload,
payment_requirements,
):
"""Expose x402's per-request cost context to Flask route handlers.

OHTTP requests arrive as ciphertext and return ciphertext, so the x402
middleware cannot parse request_json/response_json from the outer HTTP
bodies. The OHTTP controller decrypts the inner request and plaintext
response inside the enclave; this patch gives it a request-local dict where
it can attach those inner JSON objects for dynamic settlement.
"""
self._start_reaper()

request_body_bytes = x402_flask._read_body_bytes(environ)
request_json = x402_flask._try_parse_json(request_body_bytes)
parsed_request_json = (
request_json if isinstance(request_json, (dict, list)) else None
)

x402_flask.g.payment_payload = payment_payload
x402_flask.g.payment_requirements = payment_requirements
x402_flask.g.x402_session_id = session_id

cost_context = {
"method": context.method,
"path": context.path,
"request_body_bytes": request_body_bytes,
"request_json": parsed_request_json,
"payment_payload": payment_payload,
"payment_requirements": payment_requirements,
}
environ["x402.cost_context"] = cost_context

status_capture = x402_flask.StatusCapture(start_response)
status_capture.add_header(x402_flask.UPTO_SESSION_HEADER, session_id)

upstream_iter = self._original_wsgi(environ, status_capture)

return x402_flask.StreamingSessionResponse(
upstream_iter,
middleware=self,
session_id=session_id,
cost_context=cost_context,
status_ref=status_capture,
)


setattr(
x402_flask.PaymentMiddleware,
"_stream_session_response",
_patched_stream_session_response,
)


def _session_cost_calculator(ctx: dict) -> int:
def _batch_settlement_cost_calculator(ctx: dict) -> int:
# The chat/completions controllers compute cost in-band and embed it on
# the response as a SessionCost model BEFORE returning. We parse it back
# here — single source of truth for cost lives in the controller, not
# split between the controller (which serves clients) and x402's callback
# (which charges them).
#
# If the controller couldn't compute cost (e.g. missing usage from the
# provider), the block is absent — SessionCost.model_validate raises and
# x402's close() swallows it so the client is not charged. The controller
# has already logged CRITICAL in that case.
# provider), the block is absent and settlement must fail rather than
# commit a voucher for an unknown amount.
from .pricing import SessionCost

if ctx.get("path") == "/v1/ohttp":
Expand Down Expand Up @@ -267,18 +214,61 @@ def _init_payment_middleware(facilitator_url: str) -> None:
Called once from set_provider_keys() after the facilitator URL is known.
Swaps application.wsgi_app so all subsequent requests flow through it.
"""
global _batch_settlement_manager

facilitator = HTTPFacilitatorClientSync(FacilitatorConfig(url=facilitator_url))
scheme = BatchSettlementEvmScheme(
EVM_PAYMENT_ADDRESS,
BatchSettlementEvmSchemeServerConfig(
storage=FileChannelStorage(_batch_settlement_storage_dir),
),
)
server = x402ResourceServerSync(facilitator) # type: ignore[arg-type]
store = SessionStore()

server.register(BASE_MAINNET_NETWORK, ExactEvmServerScheme())
server.register(BASE_MAINNET_NETWORK, UptoEvmServerScheme())
server.register(BASE_MAINNET_NETWORK, scheme)

# Claims are asynchronous with respect to inference. Voucher requests
# only update the durable channel record; this background manager submits
# batched claim and settle transactions through the facilitator.
_batch_settlement_manager = BatchSettlementChannelManagerSync(
ChannelManagerConfigSync(
scheme=scheme,
facilitator=facilitator,
receiver=EVM_PAYMENT_ADDRESS,
token=BASE_MAINNET_OPG_ADDRESS,
network=BASE_MAINNET_NETWORK,
)
)
_batch_settlement_manager.start(
AutoSettlementConfig(
claim_interval_secs=60,
settle_interval_secs=120,
max_claims_per_batch=100,
on_claim=lambda result: logger.info(
"Claimed %s batch-settlement voucher(s): %s",
result.vouchers,
result.transaction,
),
on_settle=lambda result: logger.info(
"Settled batch-settlement claims: %s",
result.transaction,
),
on_error=lambda error: logger.error(
"Batch-settlement channel-manager error: %s", error
),
)
)
logger.info(
"Batch-settlement channel manager started "
"(storage=%s, token=%s, claim_interval=60s, settle_interval=120s)",
_batch_settlement_storage_dir,
BASE_MAINNET_OPG_ADDRESS,
)

routes = {
"POST /v1/chat/completions": RouteConfig(
accepts=[
PaymentOption(
scheme="upto",
scheme="batch-settlement",
pay_to=EVM_PAYMENT_ADDRESS,
price=AssetAmount(
amount=CHAT_COMPLETIONS_OPG_SESSION_MAX_SPEND,
Expand All @@ -301,7 +291,7 @@ def _init_payment_middleware(facilitator_url: str) -> None:
"POST /v1/completions": RouteConfig(
accepts=[
PaymentOption(
scheme="upto",
scheme="batch-settlement",
pay_to=EVM_PAYMENT_ADDRESS,
price=AssetAmount(
amount=COMPLETIONS_OPG_SESSION_MAX_SPEND,
Expand All @@ -324,7 +314,7 @@ def _init_payment_middleware(facilitator_url: str) -> None:
"POST /v1/ohttp": RouteConfig(
accepts=[
PaymentOption(
scheme="upto",
scheme="batch-settlement",
pay_to=EVM_PAYMENT_ADDRESS,
price=AssetAmount(
amount=OHTTP_OPG_SESSION_MAX_SPEND,
Expand Down Expand Up @@ -356,13 +346,14 @@ def _init_payment_middleware(facilitator_url: str) -> None:
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,
streaming_cost_calculator=_batch_settlement_cost_calculator,
settlement_data_enabled=True,
streaming_settlement_boundary=OHTTP_BATCH_SETTLEMENT_BOUNDARY,
streaming_receipt_encoder=build_ohttp_batch_settlement_frame,
)
logger.info(
"x402 payment middleware initialized with facilitator: %s", facilitator_url
"x402 batch-settlement middleware initialized with facilitator: %s",
facilitator_url,
)


Expand Down
59 changes: 43 additions & 16 deletions tee_gateway/controllers/ohttp_controller.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,11 +24,9 @@
model name, no token counts: those would be a fingerprint of the
inner request and have no role in billing.
* stream=true: outer response headers are flushed before any body chunk
is yielded, so cost isn't known at header-write time. The relay reads
the settled amount from x402: either by querying the facilitator with
its X-Upto-Session, or via X-Payment-Response on its next request. The
client still sees cost in the final SSE event inside the encrypted
stream (the ``opengradient`` block written by the chat controller).
is yielded, so final cost and the batch voucher receipt are carried in a
private plaintext billing frame immediately before the final sealed OHTTP
chunk. The relay strips that frame before forwarding bytes to the client.

Trust / payment model:
* The CLIENT encrypts only an LLM chat-completion request. It does not see,
Expand Down Expand Up @@ -85,6 +83,10 @@
OHTTP_RESPONSE_MEDIA_TYPE = "message/ohttp-res"
OHTTP_CHUNKED_RESPONSE_MEDIA_TYPE = "message/ohttp-chunked-res"
OHTTP_BILLING_FRAME_MAGIC = b"\n--opengradient-ohttp-billing-v1--\n"
# This marker is never sent on the wire. The x402 Flask middleware consumes it
# after the final cost is known, settles the batch voucher, and replaces it
# with the billing frame above before the final sealed OHTTP chunk is yielded.
OHTTP_BATCH_SETTLEMENT_BOUNDARY = b"\n--opengradient-ohttp-batch-settlement-boundary-v1--\n"
_SSE_CONTENT_TYPE = "text/event-stream"

# Cap on the encapsulated request size. The inner payload is a chat-completion
Expand Down Expand Up @@ -282,9 +284,12 @@ def _stream() -> Iterator[bytes]:
response_json=response_json,
status_code=status,
)
billing_frame = _build_billing_frame(response_json)
if billing_frame:
yield billing_frame
# The outer x402 middleware consumes this internal boundary,
# settles the batch voucher using the cost context above, and
# injects the private billing/receipt frame in its place. This
# must remain before the final OHTTP chunk so the relay can strip
# it while preserving valid chunked-OHTTP framing for the browser.
yield OHTTP_BATCH_SETTLEMENT_BOUNDARY

# Always emit exactly one final chunk so the AAD=b"final"
# marker is present — that's what protects clients from
Expand Down Expand Up @@ -351,27 +356,49 @@ def _extract_cost_headers(body_bytes: bytes) -> dict[str, str]:
}


def _build_billing_frame(response_json: dict[str, Any] | None) -> bytes:
"""Build the private gateway-to-relay billing frame for streaming OHTTP.
def build_ohttp_batch_settlement_frame(
receipt: dict[str, Any],
streaming_cost_context: dict[str, Any] | None,
) -> bytes:
"""Build the private gateway-to-relay batch settlement frame.

This plaintext frame carries only the same billing fields projected as
outer headers for non-streaming OHTTP. The relay strips it before forwarding
bytes to the browser.
This plaintext frame carries the same billing fields projected as outer
headers for non-streaming OHTTP plus the encoded x402 payment response.
The relay strips it before forwarding bytes to the browser and persists
the voucher receipt in its SDK channel storage.
"""
if not isinstance(response_json, dict):
return b""
response_json = (
streaming_cost_context.get("inner_response_json")
if isinstance(streaming_cost_context, dict)
else None
)
payload: dict[str, str] = {}
try:
if not isinstance(response_json, dict):
raise ValueError("streaming response cost context is missing")
cost = SessionCost.model_validate(response_json.get("opengradient"))
payload.update(cost.model_dump(mode="json"))
except Exception as exc:
logger.warning(
"OHTTP billing frame skipped — cost block missing/invalid on "
"streamed response: %s: %s",
type(exc).__name__,
exc,
)
if receipt.get("success"):
payment_response = receipt.get("paymentResponse")
if isinstance(payment_response, str) and payment_response:
payload["payment_response"] = payment_response
else:
payload["payment_error"] = "batch settlement response is missing"
else:
payload["payment_error"] = str(
receipt.get("error") or "batch settlement failed"
)
if not payload:
return b""
payload = json.dumps(
cost.model_dump(mode="json"),
payload,
separators=(",", ":"),
).encode("utf-8")
return OHTTP_BILLING_FRAME_MAGIC + len(payload).to_bytes(4, "big") + payload
Expand Down
Loading
Loading