|
#!/usr/bin/env -S uv run --script |
|
# /// script |
|
# requires-python = ">=3.11" |
|
# dependencies = [ |
|
# "aiohttp>=3.10,<4", |
|
# "aws-bedrock-token-generator>=1.0,<2", |
|
# "boto3>=1.35,<2", |
|
# ] |
|
# /// |
|
""" |
|
bedrock-mantle-proxy — local OpenAI-compatible proxy that forwards to AWS Bedrock |
|
Mantle's OpenAI surface, doing region routing by model name and rotating Bedrock |
|
short-lived bearer tokens transparently. |
|
|
|
The point of this proxy is so a downstream OpenAI client (most commonly the |
|
Codex CLI) can be pointed at a stable localhost URL with a dummy API key, and |
|
have its requests cleanly served against multi-region Bedrock without ever |
|
needing to restart when AWS credentials rotate or Bedrock bearer tokens expire. |
|
|
|
Routing: |
|
model.startswith("gpt-5.5") → bedrock-mantle.us-east-2.api.aws/openai/v1 |
|
model.startswith("gpt-5.4") → bedrock-mantle.us-west-2.api.aws/openai/v1 |
|
(Extend MODEL_REGION_MAP below to add more model families.) |
|
|
|
Token caching: |
|
Per-region, generated lazily on first request, refreshed when within |
|
REFRESH_BEFORE_EXPIRY of expiry. The token generator reads the current |
|
process's AWS credentials chain (so EC2 instance role, AWS_PROFILE, |
|
static AWS_ACCESS_KEY_ID, etc. all work the standard way). |
|
|
|
Pass-through behaviour: |
|
Streaming (SSE) bodies stream back unchanged byte-for-byte. Non-streaming |
|
bodies are read in full. The proxy strips the inbound Authorization header |
|
(the dummy `Bearer sk-...` from the codex client) and replaces it with the |
|
fresh Bedrock bearer token. Other headers pass through; `Host` is |
|
rewritten to the upstream host. |
|
|
|
Run: |
|
bedrock-mantle-proxy # listens on 127.0.0.1:18181 |
|
bedrock-mantle-proxy --port 9000 |
|
|
|
Health check: |
|
curl http://127.0.0.1:18181/healthz |
|
""" |
|
|
|
from __future__ import annotations |
|
|
|
import argparse |
|
import asyncio |
|
import json |
|
import logging |
|
import os |
|
import re |
|
import signal |
|
import sys |
|
import time |
|
from pathlib import Path |
|
from typing import Optional |
|
|
|
import aiohttp |
|
from aiohttp import web |
|
|
|
# `aws_bedrock_token_generator` is the official package that wraps the |
|
# `bedrock-runtime` SigV4 token-exchange flow; no need to reimplement. |
|
from aws_bedrock_token_generator import provide_token # type: ignore[import] |
|
|
|
|
|
LOG = logging.getLogger("bedrock-mantle-proxy") |
|
|
|
|
|
# --------------------------------------------------------------------------- |
|
# Routing config |
|
# --------------------------------------------------------------------------- |
|
|
|
# Each entry maps a model-name prefix to an ORDERED region list. The first |
|
# region is the **primary** — sticky preference for hot prompt cache locality |
|
# (Bedrock's prompt cache is per-region; bouncing between regions destroys |
|
# hit rate). Subsequent regions are fallbacks, tried only when the primary |
|
# returns a retryable failure (5xx / 429 / connection error). |
|
# |
|
# Once a request fails over to a fallback, the proxy "sticks" on that fallback |
|
# for the rest of its lifetime (see ProxyState.preferred_region). Region-level |
|
# outages don't recover in a few minutes; flapping back to the primary on |
|
# every request would shred the secondary's prompt cache without buying real |
|
# availability. Restart the proxy to reset the preference. |
|
# |
|
# The first matching prefix wins. Longest prefix should come first if you |
|
# ever need overlap. Both `gpt-5.X` (codex's default in `~/.codex/config.toml`) |
|
# and the qualified `openai.gpt-5.X` (Bedrock's actual model id, in case a |
|
# user configs codex with the qualified form) are accepted. |
|
MODEL_REGION_MAP: list[tuple[str, list[str]]] = [ |
|
("gpt-5.5", ["us-east-2"]), |
|
("openai.gpt-5.5", ["us-east-2"]), |
|
("gpt-5.4", ["us-west-2", "us-east-2"]), |
|
("openai.gpt-5.4", ["us-west-2", "us-east-2"]), |
|
] |
|
|
|
|
|
def regions_for_model(model: str) -> Optional[list[str]]: |
|
"""Return the ordered region preference list, or None.""" |
|
for prefix, regions in MODEL_REGION_MAP: |
|
if model.startswith(prefix): |
|
return regions |
|
return None |
|
|
|
|
|
def upstream_for_region(region: str) -> str: |
|
return f"https://bedrock-mantle.{region}.api.aws/openai/v1" |
|
|
|
|
|
# --------------------------------------------------------------------------- |
|
# Bearer token cache |
|
# --------------------------------------------------------------------------- |
|
|
|
# Refresh when fewer than this many seconds remain on the token's *actual* |
|
# claimed expiry (parsed from the embedded SigV4 presigned URL — see |
|
# `parse_token_expiry`). The buffer keeps long streaming requests safely |
|
# inside a still-valid token even if they straddle the refresh point. |
|
REFRESH_BEFORE_EXPIRY = 30 * 60 # 30 minutes |
|
|
|
# Fallback refresh interval if the embedded SigV4 expiry can't be parsed |
|
# (defensive — the token format is stable in `aws-bedrock-token-generator>=1.0`, |
|
# but if it ever changes we don't want a parse failure to lock the cache |
|
# forever or yield a 0s expiry). |
|
FALLBACK_REFRESH_INTERVAL = 30 * 60 # 30 minutes |
|
|
|
# Boto3 STS-derived credentials carry their own `expiry_time`. If the |
|
# underlying credentials are short-lived (e.g. AssumeRole giving 1h), the |
|
# presigned URL inside the token claims 12h but goes 403 when the |
|
# underlying creds rotate. Cap the cached token's effective expiry by the |
|
# credentials' own expiry minus this safety margin. |
|
CREDENTIAL_EXPIRY_SAFETY = 60 # 1 minute |
|
|
|
|
|
# Token wire format (verified empirically on aws-bedrock-token-generator 1.0.x): |
|
# |
|
# bedrock-api-key-<base64url(presigned-url-bytes)> |
|
# |
|
# The base64url payload is the canonical-request body of a SigV4-signed |
|
# `bedrock.amazonaws.com/?Action=CallWithBearerToken&X-Amz-Algorithm=AWS4-HMAC-SHA256 |
|
# &X-Amz-Credential=...&X-Amz-Date=YYYYMMDDTHHMMSSZ&X-Amz-Expires=<seconds>&...` |
|
# request. The actual expiry is `parse(X-Amz-Date) + X-Amz-Expires`. |
|
_BEARER_PREFIX = "bedrock-api-key-" |
|
_X_AMZ_DATE_RE = re.compile(r"X-Amz-Date=(\d{8}T\d{6}Z)") |
|
_X_AMZ_EXPIRES_RE = re.compile(r"X-Amz-Expires=(\d+)") |
|
|
|
|
|
def parse_token_expiry(token: str) -> Optional[float]: |
|
"""Return the token's claimed expiry as a UNIX timestamp, or None if |
|
the token doesn't match the expected `bedrock-api-key-<b64>` shape.""" |
|
if not token.startswith(_BEARER_PREFIX): |
|
return None |
|
encoded = token[len(_BEARER_PREFIX):] |
|
try: |
|
import base64 |
|
decoded = base64.urlsafe_b64decode( |
|
encoded + "=" * (-len(encoded) % 4) |
|
) |
|
text = decoded.decode("utf-8", errors="replace") |
|
except (ValueError, UnicodeDecodeError): |
|
return None |
|
date_match = _X_AMZ_DATE_RE.search(text) |
|
expires_match = _X_AMZ_EXPIRES_RE.search(text) |
|
if not date_match or not expires_match: |
|
return None |
|
try: |
|
from datetime import datetime, timezone |
|
signed_at = datetime.strptime( |
|
date_match.group(1), "%Y%m%dT%H%M%SZ" |
|
).replace(tzinfo=timezone.utc) |
|
return signed_at.timestamp() + int(expires_match.group(1)) |
|
except (ValueError, OverflowError): |
|
return None |
|
|
|
|
|
def _credential_expiry() -> Optional[float]: |
|
"""Return the underlying boto3 credentials' own expiry as a UNIX |
|
timestamp, or None for static / non-expiring credentials. Catches the |
|
"assume-role gives 1h, but the bedrock token claims 12h" footgun.""" |
|
try: |
|
import boto3 # type: ignore[import] |
|
creds = boto3.Session().get_credentials() |
|
if creds is None: |
|
return None |
|
frozen = creds.get_frozen_credentials() |
|
# Static credentials don't have an expiry. |
|
token = getattr(frozen, "token", None) |
|
if not token: |
|
return None |
|
# `creds._expiry_time` is the documented attribute on the |
|
# botocore RefreshableCredentials base class but it's private. |
|
# Fall back gracefully if botocore renames it. |
|
expiry = getattr(creds, "_expiry_time", None) |
|
if expiry is None: |
|
return None |
|
return expiry.timestamp() |
|
except Exception: # pragma: no cover — anything unexpected → no cap |
|
return None |
|
|
|
|
|
class TokenCache: |
|
"""Per-region bearer token cache with concurrent-refresh coalescing. |
|
|
|
Cache TTL is the minimum of: |
|
- Token's own SigV4 expiry (`X-Amz-Date + X-Amz-Expires`) |
|
- Underlying boto3 credentials' expiry (caps the 12h claim when STS |
|
session is shorter, e.g. assume-role giving 1h) |
|
- 30-min safety margin subtracted from each |
|
""" |
|
|
|
def __init__(self) -> None: |
|
# region -> (token, refresh_after_unix_ts) |
|
self._tokens: dict[str, tuple[str, float]] = {} |
|
self._locks: dict[str, asyncio.Lock] = {} |
|
|
|
def _lock_for(self, region: str) -> asyncio.Lock: |
|
lock = self._locks.get(region) |
|
if lock is None: |
|
lock = asyncio.Lock() |
|
self._locks[region] = lock |
|
return lock |
|
|
|
def invalidate(self, region: str) -> None: |
|
"""Drop any cached token for `region`. Called from the request |
|
handler on upstream 401/403 — the next `get()` issues a fresh |
|
token. Safe to call concurrently with `get()`.""" |
|
self._tokens.pop(region, None) |
|
|
|
def _compute_refresh_at(self, token: str) -> float: |
|
now = time.time() |
|
token_expiry = parse_token_expiry(token) |
|
creds_expiry = _credential_expiry() |
|
candidates: list[float] = [] |
|
if token_expiry is not None: |
|
candidates.append(token_expiry - REFRESH_BEFORE_EXPIRY) |
|
if creds_expiry is not None: |
|
candidates.append(creds_expiry - CREDENTIAL_EXPIRY_SAFETY) |
|
if not candidates: |
|
return now + FALLBACK_REFRESH_INTERVAL |
|
return max(now + 60, min(candidates)) # never less than 1 min ahead |
|
|
|
async def get(self, region: str) -> str: |
|
cached = self._tokens.get(region) |
|
now = time.time() |
|
if cached is not None and now < cached[1]: |
|
return cached[0] |
|
async with self._lock_for(region): |
|
# Re-check inside the lock — another task may have refreshed |
|
# while we were waiting. |
|
cached = self._tokens.get(region) |
|
if cached is not None and time.time() < cached[1]: |
|
return cached[0] |
|
token = await asyncio.to_thread(provide_token, region) |
|
refresh_at = self._compute_refresh_at(token) |
|
self._tokens[region] = (token, refresh_at) |
|
LOG.info("issued new bedrock bearer token for region=%s " |
|
"(refresh after %.0fs)", |
|
region, refresh_at - time.time()) |
|
return token |
|
|
|
|
|
# --------------------------------------------------------------------------- |
|
# Sticky-region preference state |
|
# --------------------------------------------------------------------------- |
|
|
|
class ProxyState: |
|
"""Process-wide sticky-region preference. |
|
|
|
Once any model's request fails over from its primary region to a |
|
fallback, all subsequent requests (for any model whose preference list |
|
contains that region) start with the fallback region first. This keeps |
|
prompt cache locality intact during a region-level outage instead of |
|
flapping back to the broken primary on every request. |
|
|
|
Reset by restarting the proxy. There is intentionally no auto-recovery |
|
timer — region outages are typically minutes to hours, and a proxy |
|
restart is cheap; auto-recovery would just trade a measurable failure |
|
surface (test by stopping the proxy) for an invisible one (test by |
|
waiting an arbitrary duration). |
|
""" |
|
|
|
def __init__(self) -> None: |
|
self.preferred_region: Optional[str] = None |
|
|
|
|
|
# --------------------------------------------------------------------------- |
|
# Body inspection — find `model` field without disturbing the body bytes |
|
# --------------------------------------------------------------------------- |
|
|
|
_MODEL_FIELD_RE = re.compile(rb'"model"\s*:\s*"([^"]+)"') |
|
|
|
# Captures the model field's value for in-place rewrite to the Bedrock-qualified |
|
# id. Group 1 = `"model":"` prefix, group 2 = closing quote. |
|
_MODEL_REWRITE_RE = re.compile(rb'("model"\s*:\s*")[^"]*(")') |
|
|
|
|
|
def bedrock_model_id(model: str) -> str: |
|
"""Bedrock Mantle requires the `openai.`-qualified model id; codex configs |
|
use the bare name (e.g. `gpt-5.4` → `openai.gpt-5.4`).""" |
|
return model if model.startswith("openai.") else f"openai.{model}" |
|
|
|
|
|
def parse_model_from_body(body: bytes) -> Optional[str]: |
|
"""Extract the `model` field. Tries fast-path regex first; falls back to |
|
full JSON parse on failure (handles edge-case escaping inside the value).""" |
|
m = _MODEL_FIELD_RE.search(body) |
|
if m is not None: |
|
return m.group(1).decode("utf-8", errors="replace") |
|
try: |
|
obj = json.loads(body) |
|
except (ValueError, UnicodeDecodeError): |
|
return None |
|
if isinstance(obj, dict): |
|
v = obj.get("model") |
|
if isinstance(v, str): |
|
return v |
|
return None |
|
|
|
|
|
# --------------------------------------------------------------------------- |
|
# Request sanitisation — drop tools Bedrock Mantle's OpenAI surface rejects |
|
# --------------------------------------------------------------------------- |
|
|
|
# Bedrock Mantle's OpenAI surface only accepts these tool `type`s and 400s on |
|
# anything else (e.g. codex's built-in `web_search`). codex can't know the |
|
# upstream is Mantle, so the proxy — which does — strips unsupported tools from |
|
# the request `tools` array before forwarding. Allow-list, so any future |
|
# unsupported type is dropped automatically. |
|
MANTLE_SUPPORTED_TOOL_TYPES = { |
|
"function", "mcp", "custom", "namespace", "tool_search", |
|
} |
|
|
|
# Fast-path guard: only pay the JSON round-trip when the body has a `tools` |
|
# array at all. Cheap substring scan over the raw bytes. |
|
_TOOLS_HINT_RE = re.compile(rb'"tools"\s*:\s*\[') |
|
|
|
|
|
def strip_unsupported_tools(body: bytes) -> bytes: |
|
"""Return `body` with any `tools[]` entry whose `type` is unsupported by |
|
Bedrock Mantle removed. Returns the original bytes unchanged when there is |
|
nothing to strip or the body isn't the expected JSON object shape.""" |
|
if not body or not _TOOLS_HINT_RE.search(body): |
|
return body |
|
try: |
|
obj = json.loads(body) |
|
except (ValueError, UnicodeDecodeError): |
|
return body |
|
if not isinstance(obj, dict) or not isinstance(obj.get("tools"), list): |
|
return body |
|
kept: list = [] |
|
dropped_types: list = [] |
|
for t in obj["tools"]: |
|
if isinstance(t, dict) and t.get("type") not in MANTLE_SUPPORTED_TOOL_TYPES: |
|
dropped_types.append(t.get("type")) |
|
else: |
|
kept.append(t) |
|
if not dropped_types: |
|
return body |
|
LOG.info("stripped %d unsupported tool(s) for Mantle: %s", |
|
len(dropped_types), dropped_types) |
|
obj["tools"] = kept |
|
return json.dumps(obj, ensure_ascii=False).encode("utf-8") |
|
|
|
|
|
# --------------------------------------------------------------------------- |
|
# HTTP handlers |
|
# --------------------------------------------------------------------------- |
|
|
|
# Headers we never want to pass through to upstream — `Authorization` is |
|
# replaced with our bearer; `Host` is rewritten by the new URL; `Content-Length` |
|
# is recomputed by the upstream client; hop-by-hop headers per RFC 7230 §6.1. |
|
_STRIP_REQUEST_HEADERS = { |
|
"authorization", |
|
"host", |
|
"content-length", |
|
"connection", |
|
"keep-alive", |
|
"proxy-authenticate", |
|
"proxy-authorization", |
|
"te", |
|
"trailer", |
|
"transfer-encoding", |
|
"upgrade", |
|
} |
|
|
|
# Mirror set on the response side — some additionally don't make sense to |
|
# forward back (Server, etc., but those are harmless so we keep them). |
|
_STRIP_RESPONSE_HEADERS = { |
|
"connection", |
|
"keep-alive", |
|
"proxy-authenticate", |
|
"proxy-authorization", |
|
"te", |
|
"trailer", |
|
"transfer-encoding", |
|
"upgrade", |
|
"content-encoding", # aiohttp will re-encode if needed |
|
"content-length", |
|
} |
|
|
|
|
|
async def health_handler(request: web.Request) -> web.Response: |
|
state: ProxyState = request.app["state"] |
|
prefixes: list[str] = [] |
|
seen: set[str] = set() |
|
for prefix, _ in MODEL_REGION_MAP: |
|
if prefix not in seen: |
|
prefixes.append(prefix) |
|
seen.add(prefix) |
|
return web.json_response({ |
|
"ok": True, |
|
"regions": sorted({r for _, regs in MODEL_REGION_MAP for r in regs}), |
|
"models": prefixes, |
|
"preferred_region": state.preferred_region, |
|
}) |
|
|
|
|
|
# Status codes that should trigger a fallback to the next region in the |
|
# model's preference list. 4xx other than 429 are client-side (bad request, |
|
# auth, etc.) — fanning out wouldn't help and would just waste tokens. |
|
_RETRYABLE_STATUSES = {429, 500, 502, 503, 504} |
|
|
|
|
|
async def proxy_handler(request: web.Request) -> web.StreamResponse: |
|
body = await request.read() |
|
# Drop tools Bedrock Mantle rejects (e.g. codex's `web_search`) before |
|
# forwarding; downstream Content-Length is derived from this `body`. |
|
body = strip_unsupported_tools(body) |
|
|
|
model = parse_model_from_body(body) |
|
if model is None: |
|
return web.json_response( |
|
{"error": {"message": "request body missing `model` field", |
|
"type": "invalid_request_error", |
|
"code": "model_required"}}, |
|
status=400, |
|
) |
|
|
|
region_pref = regions_for_model(model) |
|
if not region_pref: |
|
# De-duplicate prefixes for the error message. |
|
seen: set[str] = set() |
|
ordered: list[str] = [] |
|
for prefix, _ in MODEL_REGION_MAP: |
|
if prefix not in seen: |
|
ordered.append(prefix) |
|
seen.add(prefix) |
|
return web.json_response( |
|
{"error": {"message": f"unknown model `{model}` " |
|
f"(supported prefixes: {', '.join(ordered)})", |
|
"type": "invalid_request_error", |
|
"code": "model_not_found"}}, |
|
status=404, |
|
) |
|
|
|
state: ProxyState = request.app["state"] |
|
# Sticky preference: if a previous request failed over to a non-primary |
|
# region, keep using it for as long as the model's preference list |
|
# supports that region. Region-level outages don't recover in seconds; |
|
# bouncing back to the primary on every request would shred prompt cache |
|
# locality with no real availability win. |
|
candidates = list(region_pref) |
|
if state.preferred_region in candidates and candidates[0] != state.preferred_region: |
|
candidates.remove(state.preferred_region) |
|
candidates.insert(0, state.preferred_region) |
|
LOG.debug("sticky-pref: trying %s first for model=%s", |
|
state.preferred_region, model) |
|
|
|
# The path on this server starts with /v1; the upstream base ALSO ends |
|
# with /openai/v1, so we strip the leading /v1 to avoid duplicating it. |
|
rel_path = request.path |
|
if rel_path.startswith("/v1/"): |
|
rel_path = rel_path[len("/v1"):] |
|
elif rel_path == "/v1": |
|
rel_path = "/" |
|
|
|
# Bedrock Mantle wants the `openai.`-qualified model id; codex sends the |
|
# bare name. Routing above used the bare name (MODEL_REGION_MAP has both); |
|
# rewrite the forwarded body's model value (first/top-level occurrence). |
|
bedrock_model = bedrock_model_id(model) |
|
if bedrock_model != model: |
|
body = _MODEL_REWRITE_RE.sub( |
|
lambda m: m.group(1) + bedrock_model.encode() + m.group(2), |
|
body, count=1) |
|
LOG.info("mapped model %s -> %s", model, bedrock_model) |
|
|
|
token_cache: TokenCache = request.app["token_cache"] |
|
session: aiohttp.ClientSession = request.app["session"] |
|
|
|
last_exc: Optional[BaseException] = None |
|
last_status: Optional[int] = None |
|
last_reason: Optional[str] = None |
|
|
|
for attempt_idx, region in enumerate(candidates): |
|
is_fallback = attempt_idx > 0 |
|
# Reactive auth retry: if upstream returns 401/403, the cached token |
|
# may be valid SigV4-wise but rejected by Bedrock (e.g. underlying |
|
# STS credentials rotated mid-session, or IAM policy changed). Drop |
|
# the cache, request a fresh token, retry ONCE before falling over |
|
# to the next region. Bounded loop — never more than 2 attempts |
|
# per region. |
|
auth_retried = False |
|
while True: |
|
try: |
|
token = await token_cache.get(region) |
|
except Exception as exc: # pragma: no cover — credential/network errors |
|
LOG.exception("failed to obtain bedrock bearer token (region=%s)", |
|
region) |
|
last_exc = exc |
|
break |
|
|
|
upstream_url = upstream_for_region(region) + rel_path |
|
if request.query_string: |
|
upstream_url = f"{upstream_url}?{request.query_string}" |
|
|
|
headers = {k: v for k, v in request.headers.items() |
|
if k.lower() not in _STRIP_REQUEST_HEADERS} |
|
headers["Authorization"] = f"Bearer {token}" |
|
if body: |
|
headers["Content-Length"] = str(len(body)) |
|
|
|
LOG.info("→ %s %s model=%s region=%s%s%s body=%dB", |
|
request.method, rel_path, model, region, |
|
" [fallback]" if is_fallback else "", |
|
" [auth-retry]" if auth_retried else "", |
|
len(body)) |
|
|
|
try: |
|
upstream = await session.request( |
|
request.method, |
|
upstream_url, |
|
data=body if body else None, |
|
headers=headers, |
|
allow_redirects=False, |
|
) |
|
except aiohttp.ClientError as exc: |
|
LOG.warning("connection error to region=%s: %s — " |
|
"%s", |
|
region, exc, |
|
"trying fallback" if region != candidates[-1] |
|
else "no fallback left") |
|
last_exc = exc |
|
break # leave the auth-retry loop, advance to next region |
|
|
|
# Auth failure → invalidate cache, refresh token, retry once. |
|
if upstream.status in (401, 403) and not auth_retried: |
|
LOG.warning("upstream %d from region=%s — invalidating cached " |
|
"token and retrying once", |
|
upstream.status, region) |
|
await upstream.release() |
|
token_cache.invalidate(region) |
|
auth_retried = True |
|
continue # re-enter the inner loop with a fresh token |
|
|
|
# Headers received. If the status is retryable AND we haven't yet |
|
# started streaming bytes back to the client, drain + close and try |
|
# the next region. Once status code + headers are forwarded to the |
|
# downstream client we cannot retry — partial output already on |
|
# the wire. |
|
if (upstream.status in _RETRYABLE_STATUSES |
|
and region != candidates[-1]): |
|
LOG.warning("upstream %d from region=%s, falling back to next", |
|
upstream.status, region) |
|
last_status = upstream.status |
|
last_reason = upstream.reason |
|
await upstream.release() |
|
break # advance to next region |
|
|
|
# Commit to this upstream. From here on, no retry. |
|
if is_fallback and state.preferred_region != region: |
|
LOG.info("sticky-pref: switching preferred region " |
|
"from %s to %s (fallback succeeded)", |
|
state.preferred_region, region) |
|
state.preferred_region = region |
|
|
|
try: |
|
resp_headers = {k: v for k, v in upstream.headers.items() |
|
if k.lower() not in _STRIP_RESPONSE_HEADERS} |
|
response = web.StreamResponse( |
|
status=upstream.status, |
|
reason=upstream.reason, |
|
headers=resp_headers, |
|
) |
|
await response.prepare(request) |
|
async for chunk in upstream.content.iter_any(): |
|
await response.write(chunk) |
|
await response.write_eof() |
|
LOG.info("← %d %s model=%s region=%s", |
|
upstream.status, upstream.reason or "", |
|
model, region) |
|
return response |
|
finally: |
|
upstream.release() |
|
|
|
# All regions exhausted. |
|
if last_exc is not None: |
|
LOG.error("all candidate regions failed for model=%s: %s", |
|
model, last_exc) |
|
return web.json_response( |
|
{"error": {"message": f"all candidate regions failed: {last_exc}", |
|
"type": "upstream_error", |
|
"code": "bad_gateway"}}, |
|
status=502, |
|
) |
|
LOG.error("all candidate regions returned retryable status for model=%s " |
|
"(last %d %s)", model, last_status, last_reason) |
|
return web.json_response( |
|
{"error": {"message": f"all candidate regions returned retryable " |
|
f"status (last {last_status} {last_reason})", |
|
"type": "upstream_error", |
|
"code": "bad_gateway"}}, |
|
status=502, |
|
) |
|
|
|
|
|
# --------------------------------------------------------------------------- |
|
# Lifecycle / pidfile |
|
# --------------------------------------------------------------------------- |
|
|
|
def cache_dir() -> Path: |
|
base = os.environ.get("XDG_CACHE_HOME") or str(Path.home() / ".cache") |
|
p = Path(base) / "bedrock-mantle-proxy" |
|
p.mkdir(parents=True, exist_ok=True) |
|
return p |
|
|
|
|
|
def write_pidfile(path: Path) -> None: |
|
path.write_text(str(os.getpid())) |
|
|
|
|
|
def cleanup_pidfile(path: Path) -> None: |
|
try: |
|
if path.exists() and path.read_text().strip() == str(os.getpid()): |
|
path.unlink() |
|
except OSError: |
|
pass |
|
|
|
|
|
async def on_startup(app: web.Application) -> None: |
|
app["token_cache"] = TokenCache() |
|
app["session"] = aiohttp.ClientSession() |
|
app["state"] = ProxyState() |
|
|
|
|
|
async def on_cleanup(app: web.Application) -> None: |
|
session: aiohttp.ClientSession = app["session"] |
|
await session.close() |
|
|
|
|
|
def main() -> int: |
|
ap = argparse.ArgumentParser(prog="bedrock-mantle-proxy") |
|
ap.add_argument("--port", type=int, default=18181, |
|
help="local port (default 18181)") |
|
ap.add_argument("--bind", default="127.0.0.1", |
|
help="bind address (default 127.0.0.1)") |
|
ap.add_argument("--log-level", default="INFO", |
|
choices=["DEBUG", "INFO", "WARNING", "ERROR"]) |
|
args = ap.parse_args() |
|
|
|
logging.basicConfig( |
|
level=getattr(logging, args.log_level), |
|
format="%(asctime)s %(levelname)s %(name)s: %(message)s", |
|
stream=sys.stderr, |
|
) |
|
|
|
app = web.Application(client_max_size=64 * 1024 * 1024) |
|
app.on_startup.append(on_startup) |
|
app.on_cleanup.append(on_cleanup) |
|
app.router.add_get("/healthz", health_handler) |
|
# Catch-all under /v1 — proxy any method, any subpath. |
|
app.router.add_route("*", "/v1", proxy_handler) |
|
app.router.add_route("*", "/v1/{tail:.*}", proxy_handler) |
|
|
|
# Pidfile name matches the bash wrapper in the shell wrapper (codex.sh) |
|
# (the `_BM_CACHE/pid` path used by `_bm_proxy_alive` / `_bm_proxy_start`). |
|
pid_path = cache_dir() / "pid" |
|
write_pidfile(pid_path) |
|
|
|
def _on_signal(*_: object) -> None: |
|
cleanup_pidfile(pid_path) |
|
# Let aiohttp's own signal handling proceed. |
|
raise SystemExit(0) |
|
|
|
signal.signal(signal.SIGTERM, _on_signal) |
|
signal.signal(signal.SIGINT, _on_signal) |
|
|
|
print(f"bedrock-mantle-proxy listening on http://{args.bind}:{args.port}", |
|
flush=True) |
|
try: |
|
web.run_app(app, host=args.bind, port=args.port, |
|
print=lambda *a, **kw: None, |
|
handle_signals=False) |
|
finally: |
|
cleanup_pidfile(pid_path) |
|
return 0 |
|
|
|
|
|
if __name__ == "__main__": |
|
sys.exit(main()) |
