369 lines
18 KiB
Python
369 lines
18 KiB
Python
"""Project-local HTTP/1.1 transport for Responses-compatible model gateways."""
|
|
from contextlib import contextmanager
|
|
from collections import deque
|
|
import http.server
|
|
import http.client
|
|
import json
|
|
from pathlib import Path
|
|
import re
|
|
import socket
|
|
import threading
|
|
import time
|
|
import urllib.error
|
|
import urllib.request
|
|
from urllib.parse import urlsplit
|
|
|
|
|
|
class NoRedirect(urllib.request.HTTPRedirectHandler):
|
|
def redirect_request(self, *_args, **_kwargs):
|
|
return None
|
|
|
|
|
|
class RequestCancelled(OSError):
|
|
pass
|
|
|
|
|
|
class ModelTransport:
|
|
def __init__(self, provider, upstream, *, timeout=60, retries=1, retry_delay=0.5, log_path=None,
|
|
max_request_bytes=512 * 1024):
|
|
parsed = urlsplit(upstream)
|
|
local_http = parsed.scheme == "http" and parsed.hostname in {"127.0.0.1", "localhost"}
|
|
if (not isinstance(provider, str) or not re.fullmatch(r"[A-Za-z0-9_-]+", provider)
|
|
or not parsed.hostname or parsed.scheme != "https" and not local_http
|
|
or parsed.username or parsed.password or parsed.query or parsed.fragment):
|
|
raise ValueError("model transport requires a named provider and a credential-free HTTPS upstream")
|
|
self.provider = provider
|
|
self.upstream = upstream.rstrip("/")
|
|
self.timeout = float(timeout)
|
|
if not 1 <= self.timeout <= 300:
|
|
raise ValueError("model transport timeout must be between 1 and 300 seconds")
|
|
self.retries = int(retries)
|
|
if not 0 <= self.retries <= 2:
|
|
raise ValueError("model transport retries must be between 0 and 2")
|
|
self.retry_delay = float(retry_delay)
|
|
if not 0 <= self.retry_delay <= 5:
|
|
raise ValueError("model transport retry delay must be between 0 and 5 seconds")
|
|
self.log_path = Path(log_path) if log_path else None
|
|
self.max_request_bytes = int(max_request_bytes)
|
|
if not 1024 <= self.max_request_bytes <= 64 * 1024 * 1024:
|
|
raise ValueError("model request byte budget must be between 1024 and 67108864")
|
|
self.started = time.monotonic()
|
|
self.lock = threading.Lock()
|
|
self.slots = threading.BoundedSemaphore(4)
|
|
self.requests = 0
|
|
self.request_usage = deque(maxlen=256)
|
|
self.active = {}
|
|
self.finished = deque(maxlen=256)
|
|
self.server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), self.handler())
|
|
self.server.daemon_threads = True
|
|
self.thread = threading.Thread(target=self.server.serve_forever, daemon=True)
|
|
self.thread.start()
|
|
|
|
def event(self, event, **fields):
|
|
if self.log_path is None:
|
|
return
|
|
record = {"event": event, "seconds": round(time.monotonic() - self.started, 3), **fields}
|
|
with self.lock:
|
|
self.log_path.parent.mkdir(parents=True, exist_ok=True)
|
|
with self.log_path.open("a", encoding="utf-8") as stream:
|
|
stream.write(json.dumps(record) + "\n")
|
|
|
|
@property
|
|
def url(self):
|
|
return f"http://127.0.0.1:{self.server.server_port}"
|
|
|
|
def usage_checkpoint(self):
|
|
with self.lock:
|
|
return self.requests
|
|
|
|
def begin_request(self):
|
|
with self.lock:
|
|
self.requests += 1
|
|
number = self.requests
|
|
now = time.monotonic()
|
|
self.active[number] = {"request": number, "started": now, "lastData": now,
|
|
"phase": "request", "bytes": 0, "cancel": threading.Event()}
|
|
return number
|
|
|
|
def update_request(self, number, **values):
|
|
with self.lock:
|
|
if number in self.active:
|
|
self.active[number].update(values)
|
|
|
|
def ensure_active(self, number):
|
|
with self.lock:
|
|
state = self.active.get(number)
|
|
if state and state["cancel"].is_set():
|
|
raise RequestCancelled("model request was cancelled")
|
|
|
|
def finish_request(self, number):
|
|
with self.lock:
|
|
state = self.active.pop(number, None)
|
|
if state:
|
|
if state["cancel"].is_set():
|
|
state["phase"] = "cancelled"
|
|
self.finished.append({key: value for key, value in state.items()
|
|
if key not in {"cancel", "connection", "response"}})
|
|
|
|
def activity_since(self, checkpoint):
|
|
now = time.monotonic()
|
|
with self.lock:
|
|
rows = [row for row in [*self.finished, *self.active.values()] if row["request"] > checkpoint]
|
|
if not rows:
|
|
return {}
|
|
latest = max(rows, key=lambda row: row["request"])
|
|
elapsed, idle = now - latest["started"], now - latest["lastData"]
|
|
phase = latest["phase"]
|
|
# Heartbeats cannot extend the absolute request budget. A small
|
|
# completed-response grace lets the CLI dispatch buffered events.
|
|
if phase == "request":
|
|
waiting = elapsed < self.response_budget
|
|
elif phase in {"headers", "streaming"}:
|
|
waiting = elapsed < self.response_budget and idle < self.timeout
|
|
else:
|
|
waiting = phase == "completed" and idle < 5
|
|
return {"request": latest["request"], "phase": phase, "bytes": latest["bytes"],
|
|
"requestBytes": latest.get("requestBytes"), "errorType": latest.get("errorType"),
|
|
"elapsedSeconds": round(elapsed, 3), "idleSeconds": round(idle, 3),
|
|
"responseBudgetSeconds": self.response_budget, "canWait": bool(waiting),
|
|
"activeRequests": sum(row["request"] > checkpoint for row in self.active.values())}
|
|
|
|
def cancel_since(self, checkpoint):
|
|
with self.lock:
|
|
rows = [row.copy() for row in self.active.values() if row["request"] > checkpoint]
|
|
for row in rows:
|
|
row["cancel"].set()
|
|
self.active[row["request"]]["phase"] = "cancelled"
|
|
for row in rows:
|
|
connection, response = row.get("connection"), row.get("response")
|
|
sockets = [getattr(connection, "sock", None), getattr(getattr(getattr(response, "fp", None), "raw", None), "_sock", None)]
|
|
for stream in sockets:
|
|
if stream is not None:
|
|
try:
|
|
stream.shutdown(socket.SHUT_RDWR)
|
|
except OSError:
|
|
pass
|
|
if connection is not None:
|
|
connection.close()
|
|
self.event("cancelled", request=row["request"])
|
|
return len(rows)
|
|
|
|
def record_usage(self, number, usage):
|
|
if not isinstance(usage, dict):
|
|
return
|
|
tokens = usage.get("input_tokens")
|
|
if type(tokens) is not int or tokens < 0:
|
|
return
|
|
with self.lock:
|
|
self.request_usage.append((number, tokens))
|
|
self.event("usage", request=number, inputTokens=tokens)
|
|
|
|
def usage_since(self, checkpoint):
|
|
with self.lock:
|
|
rows = sorted((number, tokens) for number, tokens in self.request_usage if number > checkpoint)
|
|
requests = self.requests - checkpoint
|
|
return {"requests": requests, "measuredRequests": len(rows),
|
|
"lastInputTokens": rows[-1][1] if rows else None,
|
|
"peakInputTokens": max((tokens for _, tokens in rows), default=None),
|
|
"complete": requests > 0 and len(rows) == requests}
|
|
|
|
@property
|
|
def response_budget(self):
|
|
return (self.retries + 1) * self.timeout + self.retry_delay * self.retries * (self.retries + 1) / 2 + 10
|
|
|
|
def cli_args(self):
|
|
return ["-c", f"model_provider={json.dumps(self.provider)}",
|
|
"-c", f"model_providers.{self.provider}.base_url={json.dumps(self.url)}",
|
|
"-c", "features.enable_request_compression=false"]
|
|
|
|
def handler(self):
|
|
transport = self
|
|
|
|
class Handler(http.server.BaseHTTPRequestHandler):
|
|
protocol_version = "HTTP/1.1"
|
|
|
|
def log_message(self, *_):
|
|
pass
|
|
|
|
def respond(self, status, body, content_type="application/json", retry_after=None):
|
|
self.send_response(status)
|
|
self.send_header("Content-Type", content_type)
|
|
self.send_header("Content-Length", str(len(body)))
|
|
self.send_header("Connection", "close")
|
|
if retry_after:
|
|
self.send_header("Retry-After", retry_after)
|
|
self.end_headers()
|
|
self.wfile.write(body)
|
|
self.wfile.flush()
|
|
|
|
def do_POST(self):
|
|
if self.path not in {"/responses", "/responses/compact"}:
|
|
self.send_error(404)
|
|
return
|
|
if not transport.slots.acquire(blocking=False):
|
|
self.respond(503, b'{"error":{"message":"Local model transport is busy"}}')
|
|
return
|
|
try:
|
|
self.forward()
|
|
finally:
|
|
self.close_connection = True
|
|
transport.slots.release()
|
|
|
|
def forward(self):
|
|
sent_headers = False
|
|
number = transport.begin_request()
|
|
try:
|
|
self.connection.settimeout(30)
|
|
length = int(self.headers.get("Content-Length", "0"))
|
|
transport.update_request(number, requestBytes=length)
|
|
if not 0 < length <= 64 * 1024 * 1024 or self.headers.get("Transfer-Encoding"):
|
|
self.respond(413, b'{"error":{"message":"Invalid model request size or framing"}}')
|
|
return
|
|
if length > transport.max_request_bytes:
|
|
transport.update_request(number, phase="context_rejected")
|
|
transport.event("context_rejected", request=number, bytes=length, limit=transport.max_request_bytes)
|
|
self.respond(413, b'{"error":{"message":"Model request exceeds local context byte budget; checkpoint and start a fresh bounded turn","code":"context_budget_exceeded"}}')
|
|
return
|
|
body = self.rfile.read(length)
|
|
if len(body) != length:
|
|
self.respond(400, b'{"error":{"message":"Incomplete model request"}}')
|
|
return
|
|
excluded = {"host", "connection", "content-length", "transfer-encoding", "keep-alive", "te", "trailer",
|
|
"upgrade", "proxy-authorization", "proxy-connection", "accept-encoding"}
|
|
headers = {k: v for k, v in self.headers.items() if k.lower() not in excluded}
|
|
headers.update({"Accept-Encoding": "identity", "Connection": "close"})
|
|
transport.event("request", request=number, path=self.path, bytes=length)
|
|
if not self.headers.get("Content-Encoding"):
|
|
try:
|
|
payload = json.loads(body)
|
|
inputs = payload.get("input", [])
|
|
tools = payload.get("tools", [])
|
|
tool_outputs = [row.get("output", "") for row in inputs
|
|
if isinstance(row, dict) and row.get("type") in {"function_call_output", "custom_tool_call_output"}] if isinstance(inputs, list) else []
|
|
transport.event("context", request=number, inputItems=len(inputs) if isinstance(inputs, list) else 1,
|
|
tools=len(tools) if isinstance(tools, list) else 0,
|
|
toolOutputBytes=sum(len((value if isinstance(value, str) else json.dumps(value)).encode("utf-8")) for value in tool_outputs))
|
|
except (ValueError, AttributeError):
|
|
pass
|
|
request = urllib.request.Request(transport.upstream + self.path, data=body, headers=headers)
|
|
with transport.open_request(request, number) as response:
|
|
transport.ensure_active(number)
|
|
transport.update_request(number, phase="headers", response=response, lastData=time.monotonic())
|
|
transport.event("headers", request=number, status=response.status)
|
|
self.send_response(response.status)
|
|
self.send_header("Content-Type", response.headers.get("Content-Type", "text/event-stream"))
|
|
self.send_header("Connection", "close")
|
|
self.end_headers()
|
|
sent_headers = True
|
|
size, terminal = 0, None
|
|
for line in response:
|
|
if size == 0:
|
|
transport.event("first_body", request=number)
|
|
size += len(line)
|
|
if line.startswith(b"data:"):
|
|
try:
|
|
item = json.loads(line[5:])
|
|
transport.update_request(number, phase="streaming", lastData=time.monotonic())
|
|
if item.get("type") in {"response.completed", "response.failed", "error"}:
|
|
terminal = item["type"]
|
|
if item.get("type") == "response.completed" and self.path == "/responses":
|
|
transport.record_usage(number, (item.get("response") or {}).get("usage"))
|
|
except (ValueError, AttributeError):
|
|
pass
|
|
self.wfile.write(line)
|
|
self.wfile.flush()
|
|
transport.ensure_active(number)
|
|
transport.update_request(number, bytes=size)
|
|
transport.event("end", request=number, bytes=size, terminal=terminal)
|
|
transport.update_request(number, phase="completed" if terminal == "response.completed" else "incomplete",
|
|
lastData=time.monotonic(), terminal=terminal)
|
|
except urllib.error.HTTPError as error:
|
|
transport.update_request(number, phase="http_error")
|
|
transport.event("http_error", request=number, status=error.code)
|
|
try:
|
|
self.respond(error.code, error.read(1024 * 1024), error.headers.get("Content-Type", "application/json"),
|
|
error.headers.get("Retry-After"))
|
|
except OSError:
|
|
pass
|
|
except (OSError, ValueError, http.client.HTTPException) as error:
|
|
transport.update_request(number, phase="cancelled" if isinstance(error, RequestCancelled) else "error", errorType=type(error).__name__)
|
|
transport.event("error", request=number, errorType=type(error).__name__, afterHeaders=sent_headers)
|
|
if not sent_headers:
|
|
try:
|
|
self.respond(502, b'{"error":{"message":"Model upstream request failed","type":"upstream_error"}}')
|
|
except OSError:
|
|
pass
|
|
finally:
|
|
transport.finish_request(number)
|
|
|
|
return Handler
|
|
|
|
def open_request(self, request, number):
|
|
transport = self
|
|
def connection_factory(kind):
|
|
def create(*args, **kwargs):
|
|
transport.ensure_active(number)
|
|
connection = kind(*args, **kwargs)
|
|
transport.update_request(number, connection=connection)
|
|
return connection
|
|
return create
|
|
class HttpHandler(urllib.request.HTTPHandler):
|
|
def http_open(self, request):
|
|
return self.do_open(connection_factory(http.client.HTTPConnection), request)
|
|
class HttpsHandler(urllib.request.HTTPSHandler):
|
|
def https_open(self, request):
|
|
return self.do_open(connection_factory(http.client.HTTPSConnection), request, context=self._context)
|
|
opener = urllib.request.build_opener(NoRedirect(), HttpHandler(), HttpsHandler())
|
|
for attempt in range(self.retries + 1):
|
|
self.ensure_active(number)
|
|
self.event("attempt", request=number, attempt=attempt + 1)
|
|
try:
|
|
return opener.open(request, timeout=self.timeout)
|
|
except urllib.error.HTTPError as error:
|
|
self.ensure_active(number)
|
|
if error.code not in {408, 502, 503, 504} or error.headers.get("Retry-After") or attempt == self.retries:
|
|
raise
|
|
error.close()
|
|
reason = f"http_{error.code}"
|
|
except (urllib.error.URLError, OSError, http.client.HTTPException) as error:
|
|
self.ensure_active(number)
|
|
if attempt == self.retries:
|
|
raise
|
|
reason = type(error).__name__
|
|
# No response bytes have reached the client, so no tool call
|
|
# from this attempt can have been dispatched locally.
|
|
self.event("retry", request=number, attempt=attempt + 1, reason=reason)
|
|
with self.lock:
|
|
state = self.active.get(number)
|
|
if state:
|
|
if state["cancel"].wait(self.retry_delay * (attempt + 1)):
|
|
raise RequestCancelled("cancelled before retry")
|
|
else:
|
|
time.sleep(self.retry_delay * (attempt + 1))
|
|
|
|
def close(self):
|
|
self.cancel_since(0)
|
|
self.server.shutdown()
|
|
self.server.server_close()
|
|
self.thread.join(timeout=2)
|
|
|
|
|
|
@contextmanager
|
|
def configured_transport(root, *, process_id):
|
|
path = Path(root) / "docs/current/scheduler-transport.json"
|
|
if not path.is_file():
|
|
yield None
|
|
return
|
|
config = json.loads(path.read_text(encoding="utf-8"))
|
|
if not config.get("enabled", False):
|
|
yield None
|
|
return
|
|
transport = ModelTransport(config["provider"], config["upstreamBaseUrl"], timeout=config.get("requestTimeoutSeconds", 60),
|
|
retries=config.get("retriesBeforeHeaders", 1),
|
|
max_request_bytes=config.get("maxRequestBytes", 512 * 1024),
|
|
log_path=Path(root) / f"docs/current/continuation-runs/transport-{process_id}.jsonl")
|
|
try:
|
|
yield transport
|
|
finally:
|
|
transport.close()
|