"""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()