450 lines
17 KiB
Python
450 lines
17 KiB
Python
"""Responses WebSocket transport beneath the existing SDK and safety guards.
|
|
|
|
The SDK still serializes complete, authoritative requests. Only this wire boundary
|
|
omits a verified prefix. No model messages, tool results, or screenshots are lost
|
|
from local history, so any new connection can start from the full request.
|
|
"""
|
|
|
|
import asyncio
|
|
from collections import OrderedDict
|
|
from contextlib import suppress
|
|
import hashlib
|
|
import json
|
|
import time
|
|
from uuid import uuid4
|
|
|
|
import httpx
|
|
from websockets.asyncio.client import connect
|
|
from websockets.exceptions import ConnectionClosed, InvalidStatus, SecurityError
|
|
from websockets.protocol import State
|
|
|
|
from artemis.llm.diagnostics import ACTIVE_ATTEMPT
|
|
from artemis.llm.http import model_http_options
|
|
from artemis.llm.responses_diagnostics import identifier
|
|
|
|
IDLE_SECONDS = 60
|
|
MAX_SESSIONS = 32
|
|
MAX_CONNECTION_SECONDS = 55 * 60
|
|
|
|
|
|
def _json(value):
|
|
return json.dumps(value, ensure_ascii=False, separators=(",", ":"), sort_keys=True)
|
|
|
|
|
|
def _fingerprint(item):
|
|
# Assistant message IDs/status are transport metadata; store=false strips
|
|
# message IDs on replay. Keep every semantic field, including call_id/phase.
|
|
if isinstance(item, dict):
|
|
item = dict(item)
|
|
item.pop("status", None)
|
|
if item.get("type") == "message" or item.get("role") == "assistant":
|
|
item.pop("id", None)
|
|
item.setdefault("type", "message")
|
|
# Completion events may include token probabilities; LangChain
|
|
# correctly omits those response-only statistics from replay input.
|
|
if isinstance(item.get("content"), list):
|
|
item["content"] = [
|
|
{k: v for k, v in block.items() if k != "logprobs"}
|
|
if isinstance(block, dict) and block.get("type") == "output_text"
|
|
else block
|
|
for block in item["content"]
|
|
]
|
|
if item.get("type") == "function_call" and isinstance(item.get("arguments"), str):
|
|
with suppress(ValueError):
|
|
item["arguments"] = json.loads(item["arguments"])
|
|
return hashlib.sha256(_json(item).encode()).digest()
|
|
|
|
|
|
class _Connect(connect):
|
|
def process_redirect(self, exc):
|
|
# Never forward the API credential to a redirected endpoint.
|
|
if not isinstance(exc, InvalidStatus) or exc.response.status_code not in {
|
|
301,
|
|
302,
|
|
303,
|
|
307,
|
|
308,
|
|
}:
|
|
return exc
|
|
return SecurityError("Responses WebSocket redirects are disabled")
|
|
|
|
|
|
class _Session:
|
|
def __init__(self):
|
|
self.lock = asyncio.Lock()
|
|
self.users = 0
|
|
self.connection = None
|
|
self.opened_at = 0.0
|
|
self.previous_id = None
|
|
self.history = ()
|
|
self.options = None
|
|
self.timer = None
|
|
self.handshake = {}
|
|
|
|
async def close(self):
|
|
if self.timer:
|
|
self.timer.cancel()
|
|
self.timer = None
|
|
self.previous_id = None
|
|
self.history = ()
|
|
self.options = None
|
|
connection, self.connection = self.connection, None
|
|
if connection is not None:
|
|
with suppress(ConnectionClosed, OSError, TimeoutError):
|
|
async with asyncio.timeout(1):
|
|
await connection.close()
|
|
|
|
|
|
class ResponsesWebSocketTransport(httpx.AsyncBaseTransport):
|
|
"""One serial socket per conversation/role, independent sockets for concurrency."""
|
|
|
|
def __init__(self, *, timeout, use_proxy):
|
|
self.timeout = timeout
|
|
self.use_proxy = use_proxy
|
|
self.http = httpx.AsyncClient(timeout=timeout, **model_http_options())
|
|
self.sessions = OrderedDict()
|
|
self.cleanup_tasks = set()
|
|
|
|
async def aclose(self):
|
|
sessions = list(self.sessions.values())
|
|
self.sessions.clear()
|
|
await asyncio.gather(*(session.close() for session in sessions))
|
|
if self.cleanup_tasks:
|
|
await asyncio.gather(*tuple(self.cleanup_tasks), return_exceptions=True)
|
|
await self.http.aclose()
|
|
|
|
def _close_background(self, session):
|
|
task = asyncio.create_task(session.close())
|
|
self.cleanup_tasks.add(task)
|
|
task.add_done_callback(self.cleanup_tasks.discard)
|
|
|
|
async def _session(self, request, body):
|
|
key = body.get("prompt_cache_key")
|
|
# Header identity prevents credentials, routing, and different gateways
|
|
# from sharing a socket. Per-request correlation headers aren't affinity.
|
|
headers = {
|
|
k: v
|
|
for k, v in request.headers.items()
|
|
if k not in {"content-length", "x-client-request-id", "x-client-attempt-id"}
|
|
}
|
|
identity = hashlib.sha256(_json([str(request.url), key, headers]).encode()).digest()
|
|
reusable = isinstance(key, str) and bool(key)
|
|
if reusable and identity in self.sessions:
|
|
session = self.sessions.pop(identity)
|
|
self.sessions[identity] = session
|
|
else:
|
|
if len(self.sessions) >= MAX_SESSIONS:
|
|
for old_key, old in list(self.sessions.items()):
|
|
if not old.users:
|
|
del self.sessions[old_key]
|
|
# Reserve the replacement before yielding: concurrent
|
|
# requests for its scope must find the same session.
|
|
self._close_background(old)
|
|
break
|
|
session = _Session()
|
|
reusable = reusable and len(self.sessions) < MAX_SESSIONS
|
|
if reusable:
|
|
self.sessions[identity] = session
|
|
acquired = False
|
|
session.users += 1
|
|
try:
|
|
await session.lock.acquire()
|
|
acquired = True
|
|
finally:
|
|
if not acquired:
|
|
session.users -= 1
|
|
self._idle(identity, session, reusable)
|
|
if session.timer:
|
|
session.timer.cancel()
|
|
session.timer = None
|
|
return identity, session, reusable
|
|
|
|
def release(self, identity, session, reusable):
|
|
session.lock.release()
|
|
session.users -= 1
|
|
self._idle(identity, session, reusable)
|
|
|
|
def _idle(self, identity, session, reusable):
|
|
if not reusable or session.users or self.sessions.get(identity) is not session:
|
|
return
|
|
|
|
def expire():
|
|
if session.users or self.sessions.get(identity) is not session:
|
|
return
|
|
del self.sessions[identity]
|
|
self._close_background(session)
|
|
|
|
session.timer = asyncio.get_running_loop().call_later(IDLE_SECONDS, expire)
|
|
|
|
async def handle_async_request(self, request):
|
|
if request.method != "POST" or not request.url.path.rstrip("/").endswith("/responses"):
|
|
return await self.http.send(request, stream=True)
|
|
body = json.loads(await request.aread())
|
|
identity, session, reusable = await self._session(request, body)
|
|
exchange = _Exchange(self, request, body, identity, session, reusable)
|
|
stream_handed_off = False
|
|
try:
|
|
await exchange.start()
|
|
first = await exchange.read()
|
|
if first.get("type") == "error":
|
|
status = first.get("status", 400)
|
|
status = status if type(status) is int and 400 <= status <= 599 else 400
|
|
return httpx.Response(status, json={"error": first.get("error", {})})
|
|
if not body.get("stream"):
|
|
event = first
|
|
while not exchange.done:
|
|
event = await exchange.read()
|
|
if event.get("type") == "error":
|
|
raise httpx.ReadError("Responses WebSocket request failed", request=request)
|
|
return httpx.Response(
|
|
200, json=event["response"], extensions={"artemis_websocket": True}
|
|
)
|
|
response = httpx.Response(
|
|
200,
|
|
headers={"content-type": "text/event-stream"},
|
|
stream=_EventStream(exchange, first),
|
|
extensions={"artemis_websocket": True},
|
|
)
|
|
stream_handed_off = True
|
|
return response
|
|
except InvalidStatus as error:
|
|
return httpx.Response(
|
|
error.response.status_code,
|
|
json={
|
|
"error": {
|
|
"message": "Responses WebSocket 升级失败，请检查网关 WebSocket 支持或选择 HTTP。"
|
|
}
|
|
},
|
|
headers={
|
|
k: v
|
|
for k, v in error.response.headers.raw_items()
|
|
if k.lower() in {"retry-after", "x-request-id"}
|
|
},
|
|
)
|
|
finally:
|
|
if not stream_handed_off:
|
|
await exchange.aclose()
|
|
|
|
|
|
class _Exchange:
|
|
def __init__(self, transport, request, body, identity, session, reusable):
|
|
self.transport = transport
|
|
self.request = request
|
|
self.body = body
|
|
self.identity = identity
|
|
self.session = session
|
|
self.reusable = reusable
|
|
self.done = False
|
|
self.closed = False
|
|
self.recovered = False
|
|
self.output_seen = False
|
|
self.input = body.get("input")
|
|
self.input_hashes = (
|
|
tuple(map(_fingerprint, self.input)) if isinstance(self.input, list) else ()
|
|
)
|
|
self.options = hashlib.sha256(
|
|
_json(
|
|
{
|
|
k: v
|
|
for k, v in body.items()
|
|
if k not in {"input", "stream", "background", "previous_response_id"}
|
|
}
|
|
).encode()
|
|
).digest()
|
|
self.explicit_state = bool(body.get("previous_response_id") or body.get("conversation"))
|
|
self.attempt = ACTIVE_ATTEMPT.get()
|
|
self.stats = {
|
|
"mode": "websocket",
|
|
"full_request_bytes": len(request.content),
|
|
"sent_payload_bytes": 0,
|
|
"received_payload_bytes": 0,
|
|
"create_events": 0,
|
|
"incremental": False,
|
|
"reused_input_items": 0,
|
|
"connection_reused": session.connection is not None,
|
|
}
|
|
if self.attempt is not None:
|
|
self.attempt.data["transport"] = self.stats
|
|
|
|
async def start(self):
|
|
session = self.session
|
|
if session.connection is not None and (
|
|
session.connection.state is not State.OPEN
|
|
or time.monotonic() - session.opened_at >= MAX_CONNECTION_SECONDS
|
|
):
|
|
await session.close()
|
|
self.stats["connection_reused"] = False
|
|
if session.connection is None:
|
|
url = self.request.url.copy_with(
|
|
scheme="wss" if self.request.url.scheme == "https" else "ws"
|
|
)
|
|
headers = {
|
|
k: v
|
|
for k, v in self.request.headers.items()
|
|
if k
|
|
not in {"host", "content-length", "content-type", "connection", "accept-encoding"}
|
|
}
|
|
options = {"ssl": model_http_options()["verify"]} if url.scheme == "wss" else {}
|
|
session.connection = await _Connect(
|
|
str(url),
|
|
additional_headers=headers,
|
|
user_agent_header=None,
|
|
proxy=True if self.transport.use_proxy else None,
|
|
open_timeout=self.transport.timeout,
|
|
close_timeout=1,
|
|
max_size=64 * 1024 * 1024,
|
|
**options,
|
|
)
|
|
session.opened_at = time.monotonic()
|
|
session.handshake = {
|
|
"client_request_id": identifier(headers.get("x-client-request-id")),
|
|
"client_attempt_id": identifier(headers.get("x-client-attempt-id")),
|
|
}
|
|
if self.attempt:
|
|
self.attempt.data.update(session.handshake, request_header_scope="websocket_handshake")
|
|
for field, local in (
|
|
("client_request_id", "request_id"),
|
|
("client_attempt_id", "attempt_id"),
|
|
):
|
|
self.attempt.data[field + "_matches"] = session.handshake.get(
|
|
field
|
|
) == self.attempt.data.get(local)
|
|
self.stats["handshake_http_status"] = 101
|
|
await self.send(incremental=True)
|
|
|
|
async def send(self, *, incremental):
|
|
session = self.session
|
|
payload = {k: v for k, v in self.body.items() if k not in {"stream", "background"}}
|
|
count = len(session.history)
|
|
can_continue = (
|
|
incremental
|
|
and not self.explicit_state
|
|
and session.previous_id
|
|
and session.options == self.options
|
|
and count
|
|
and len(self.input_hashes) > count
|
|
and self.input_hashes[:count] == session.history
|
|
)
|
|
if can_continue:
|
|
payload["previous_response_id"] = session.previous_id
|
|
payload["input"] = self.input[count:]
|
|
self.stats.update(incremental=True, reused_input_items=count)
|
|
else:
|
|
self.stats.update(incremental=False, reused_input_items=0)
|
|
self.stats["full_context_reason"] = (
|
|
"explicit_state"
|
|
if self.explicit_state
|
|
else "history_or_options_changed"
|
|
if session.previous_id
|
|
else "new_connection_or_chain"
|
|
)
|
|
payload["type"] = "response.create"
|
|
encoded = _json(payload)
|
|
await session.connection.send(encoded)
|
|
self.stats["sent_payload_bytes"] += len(encoded.encode())
|
|
self.stats["create_events"] += 1
|
|
self.stats["payload_bytes_saved"] = max(
|
|
0, self.stats["full_request_bytes"] - self.stats["sent_payload_bytes"]
|
|
)
|
|
if self.attempt:
|
|
self.attempt.data["request_sent_at"] = time.time()
|
|
self.attempt.publish_if_due()
|
|
|
|
async def read(self):
|
|
try:
|
|
async with asyncio.timeout(self.transport.timeout):
|
|
raw = await self.session.connection.recv()
|
|
self.stats["received_payload_bytes"] += len(
|
|
raw.encode() if isinstance(raw, str) else raw
|
|
)
|
|
event = json.loads(raw)
|
|
if not isinstance(event, dict):
|
|
raise ValueError("Responses WebSocket event is not an object")
|
|
except ConnectionClosed as error:
|
|
raise httpx.ReadError(
|
|
"Responses WebSocket connection closed before completion", request=self.request
|
|
) from error
|
|
except (ValueError, UnicodeError) as error:
|
|
raise httpx.ReadError(
|
|
"Invalid Responses WebSocket event", request=self.request
|
|
) from error
|
|
kind = event.get("type", "")
|
|
error = event.get("error") or {}
|
|
if (
|
|
kind == "error"
|
|
and error.get("code") == "previous_response_not_found"
|
|
and self.stats["incremental"]
|
|
and not self.recovered
|
|
and not self.output_seen
|
|
):
|
|
# An explicit rejection means no model output was generated. Recover
|
|
# once from the authoritative input; never replay an ambiguous drop.
|
|
self.recovered = True
|
|
self.session.previous_id = None
|
|
self.stats["chain_recovered"] = True
|
|
await self.send(incremental=False)
|
|
return await self.read()
|
|
if kind not in {"error", "response.created", "response.in_progress", "response.queued"}:
|
|
self.output_seen = True
|
|
if kind in {
|
|
"response.completed",
|
|
"response.failed",
|
|
"response.incomplete",
|
|
"response.cancelled",
|
|
}:
|
|
self.done = True
|
|
response = event.get("response") or {}
|
|
self.session.previous_id = None
|
|
if (
|
|
kind == "response.completed"
|
|
and response.get("status") == "completed"
|
|
and isinstance(response.get("id"), str)
|
|
and response["id"]
|
|
and isinstance(response.get("output"), list)
|
|
and self.input_hashes
|
|
and not self.explicit_state
|
|
):
|
|
output = [
|
|
item
|
|
for item in response["output"]
|
|
if not (
|
|
item.get("type") == "reasoning"
|
|
and not item.get("encrypted_content")
|
|
and self.body.get("store") is False
|
|
)
|
|
]
|
|
self.session.previous_id = response["id"]
|
|
self.session.history = self.input_hashes + tuple(map(_fingerprint, output))
|
|
self.session.options = self.options
|
|
return event
|
|
|
|
async def aclose(self):
|
|
if self.closed:
|
|
return
|
|
self.closed = True
|
|
try:
|
|
if not self.done or not self.reusable:
|
|
await self.session.close()
|
|
finally:
|
|
self.transport.release(self.identity, self.session, self.reusable)
|
|
|
|
|
|
class _EventStream(httpx.AsyncByteStream):
|
|
def __init__(self, exchange, first):
|
|
self.exchange = exchange
|
|
self.first = first
|
|
|
|
async def __aiter__(self):
|
|
try:
|
|
event = self.first
|
|
while True:
|
|
yield b"data: " + _json(event).encode() + b"\n\n"
|
|
if self.exchange.done or event.get("type") == "error":
|
|
return
|
|
event = await self.exchange.read()
|
|
finally:
|
|
await self.aclose()
|
|
|
|
async def aclose(self):
|
|
await self.exchange.aclose()
|