From 2a3f43fb6080ff801d3f4a6dae0966d0547f2bab Mon Sep 17 00:00:00 2001 From: chenmingxuan Date: Wed, 4 Nov 2026 17:15:32 +0800 Subject: [PATCH] feat(llm): support incremental Responses WebSocket transport --- .env.example | 2 + README.md | 2 + .../services/model_settings_service.py | 16 +- .../model-settings.component.spec.ts | 25 + .../model-settings.component.ts | 11 +- artemis/config/llm.py | 3 + artemis/config/settings.py | 2 + artemis/llm/diagnostics.py | 9 +- artemis/llm/router.py | 31 +- artemis/llm/websocket.py | 449 +++++++++++++ artemis/resources/config/artemis.jsonc | 1 + artemis/services/llm.py | 1 + config/artemis.jsonc | 1 + docs/ai-automation.md | 12 + mcp_server/catalog.json | 2 +- mcp_server/tools/automation.py | 3 +- pyproject.toml | 2 +- skills/artemis-use/references/modes.md | 2 + .../unit/admin_console/test_model_settings.py | 5 + tests/unit/mcp/test_automation_tools.py | 2 + tests/unit/test_responses_websocket.py | 605 ++++++++++++++++++ uv.lock | 2 +- 22 files changed, 1176 insertions(+), 12 deletions(-) create mode 100644 artemis/llm/websocket.py create mode 100644 tests/unit/test_responses_websocket.py diff --git a/.env.example b/.env.example index 2ee467c..5ccaafa 100644 --- a/.env.example +++ b/.env.example @@ -9,6 +9,8 @@ OPENAI_TIMEOUT_SECONDS=300 # 模型请求和模型凭据检查默认直连；设为 true 才读取 Windows/HTTP(S)_PROXY/ALL_PROXY/NO_PROXY。 # 仅环境配置，不显示在网页中；修改后重启服务生效。TLS 校验及自定义 CA 不受影响。 OPENAI_USE_SYSTEM_PROXY=false +# http 保持兼容；websocket 复用连接并按已验证的响应链增量发送，要求网关支持 Responses WebSocket。 +OPENAI_TRANSPORT=http # 可选的后台步骤截图摘要单次超时，独立于上面的主模型请求超时。 ARTEMIS_STEP_SUMMARIZER_TIMEOUT_SECONDS=120 # 失败后最多重试 5 次，由 Artemis 统一管理，SDK 不叠加重试。 diff --git a/README.md b/README.md index 72df2e5..96dfe6d 100644 --- a/README.md +++ b/README.md @@ -61,6 +61,8 @@ cd Artemis-Eapil 也可在项目 `.env` 中设置 `OPENAI_API_KEY`、`OPENAI_BASE_URL` 和 `OPENAI_MODEL`，其他参数见 [.env.example](.env.example)。新安装默认接口地址为 `https://api-n-cd.playones.com/v1`；已保存的地址不变。直接编辑环境文件后需重启相应服务。不要提交自己的密钥。 +模型配置的“传输协议”可选择 **WebSocket（增量传输）**，或设置 `OPENAI_TRANSPORT=websocket`。Flash、Pro、探索及辅助模型的异步请求复用 Responses 连接，历史一致时只上传新增输入；网关需支持 `/v1/responses` 的 WebSocket 升级与响应链。默认 HTTP 保持兼容，详细行为与流量诊断见[传输说明](docs/ai-automation.md#responses-websocket-增量传输)。 + “功能开关”默认将任务截图长边限制为 2048 像素；关闭后保留设备原始分辨率。可在 [config/artemis.jsonc](config/artemis.jsonc) 的 `agent.screenshot` 中覆盖 `enabled` 和 `max_long_edge`，重启后对新任务生效。 首页“功能配置”提供“Flash 最终 AI 判定与分析”，默认关闭。切换后自动保存并提示结果，失败恢复原值；无需重启服务。开启后，Flash 在执行结束时根据用例需求和执行步骤进行独立复核；关闭后显示“执行完成 · 未做 AI 判定”，保留执行步骤、截图和报告，不自动记为测试通过。此设置适用于所有 Flash 任务，包括批量、精选、MCP/SDK 和探索提炼用例的 Flash 复跑；保存后对尚未开始的任务生效，运行中的任务保持原设置。Pro 和主动探索不受影响。其余功能配置仍需点击底部保存按钮。 diff --git a/apps/admin_console/services/model_settings_service.py b/apps/admin_console/services/model_settings_service.py index f241740..5f28e9d 100644 --- a/apps/admin_console/services/model_settings_service.py +++ b/apps/admin_console/services/model_settings_service.py @@ -22,6 +22,7 @@ class ModelSettingsUpdate(BaseModel): api_key: SecretStr | None = None # Omitted/blank keeps the existing key. timeout_seconds: float = Field(default=300, gt=0, le=3600) max_retries: int = Field(default=5, ge=0, le=10) + transport: Literal["http", "websocket"] | None = None retry_interval_seconds: float = Field(default=10, ge=0, le=300) rate_limit_retry_seconds: float | None = Field(default=None, ge=0, le=3600) tool_fallback_model: str | None = Field(default=None, min_length=1, max_length=200) @@ -72,6 +73,7 @@ def get_model_settings() -> dict: "api_key_configured": bool(settings.OPENAI_API_KEY), "timeout_seconds": settings.OPENAI_TIMEOUT_SECONDS, "max_retries": settings.OPENAI_MAX_RETRIES, + "transport": settings.OPENAI_TRANSPORT, "retry_interval_seconds": settings.OPENAI_RETRY_INTERVAL_SECONDS, "rate_limit_retry_seconds": settings.OPENAI_RATE_LIMIT_RETRY_SECONDS, "tool_fallback_model": settings.OPENAI_TOOL_FALLBACK_MODEL, @@ -102,6 +104,7 @@ def save_model_settings(update: ModelSettingsUpdate) -> dict: "OPENAI_REASONING_EFFORT": update.reasoning_effort, } optional_settings = { + "transport": ("OPENAI_TRANSPORT", str), "rate_limit_retry_seconds": ("OPENAI_RATE_LIMIT_RETRY_SECONDS", float), "tool_fallback_model": ("OPENAI_TOOL_FALLBACK_MODEL", str), "tool_fallback_reasoning_effort": ("OPENAI_TOOL_FALLBACK_REASONING_EFFORT", str), @@ -261,10 +264,19 @@ async def test_model_settings(update: ModelSettingsUpdate) -> dict: api_base=update.base_url, timeout_seconds=update.timeout_seconds, max_retries=0, + transport=update.transport or settings.OPENAI_TRANSPORT, reasoning_effort=None if update.reasoning_effort == "default" else update.reasoning_effort, ) model = ModelFactory.create_model(endpoint) - result = await model.ainvoke([HumanMessage(content="Reply with OK only.")]) + try: + result = await model.ainvoke([HumanMessage(content="Reply with OK only.")]) + finally: + await model.http_async_client.aclose() + model.http_client.close() if not result.content: raise ValueError("Responses 接口返回空内容，请检查模型配置") - return {"success": True, "message": "Responses 连接测试成功", "model": update.model} + return { + "success": True, + "message": f"Responses {endpoint.transport.upper()} 连接测试成功", + "model": update.model, + } diff --git a/apps/showcase_ui/src/app/components/model-settings/model-settings.component.spec.ts b/apps/showcase_ui/src/app/components/model-settings/model-settings.component.spec.ts index 15eb29f..24c26fd 100644 --- a/apps/showcase_ui/src/app/components/model-settings/model-settings.component.spec.ts +++ b/apps/showcase_ui/src/app/components/model-settings/model-settings.component.spec.ts @@ -30,6 +30,31 @@ describe('持久化模型与功能设置', () => { http.expectNone('/api/system/model-settings'); })); + it('加载、测试和保存 WebSocket 选择，旧配置默认 HTTP', fakeAsync(() => { + const fixture = TestBed.createComponent(ModelSettingsComponent); + fixture.detectChanges(); + http.expectOne('/api/system/model-settings').flush(config); + tick(); + fixture.detectChanges(); + const component = fixture.componentInstance; + expect(component.transport).toBe('http'); + const selector: HTMLSelectElement = fixture.nativeElement.querySelector('#model-transport'); + selector.value = 'websocket'; + selector.dispatchEvent(new Event('change')); + tick(); + fixture.detectChanges(); + component.test(); + const check = http.expectOne('/api/system/model-settings/test'); + expect(check.request.body.transport).toBe('websocket'); + check.flush({message: 'Responses WEBSOCKET 连接测试成功'}); + component.save(); + const save = http.expectOne('/api/system/model-settings'); + expect(save.request.body.transport).toBe('websocket'); + expect(save.request.body.api_key).toBeUndefined(); + save.flush({...config, transport: 'websocket'}); + expect(component.transport).toBe('websocket'); + })); + it('加载配置后保留密钥，保存新模型且不发送空密钥', () => { const fixture = TestBed.createComponent(ModelSettingsComponent); fixture.detectChanges(); diff --git a/apps/showcase_ui/src/app/components/model-settings/model-settings.component.ts b/apps/showcase_ui/src/app/components/model-settings/model-settings.component.ts index 6f6d54a..ea25125 100644 --- a/apps/showcase_ui/src/app/components/model-settings/model-settings.component.ts +++ b/apps/showcase_ui/src/app/components/model-settings/model-settings.component.ts @@ -8,6 +8,7 @@ import { ModelSelectorComponent } from '../model-selector/model-selector.compone import { MODEL_PRESETS } from '../../core/model-presets'; interface ModelSettings { + transport?: 'http' | 'websocket'; base_url: string; model: string; api_key_configured: boolean; @@ -36,6 +37,12 @@ interface ModelSettings { 填写接口根地址，通常包含 /v1，无需添加 /responses。 + + + WebSocket 可减少重复历史与截图的上传，需要接口支持。保存后对新任务生效。 选择预设模型，或选择“自定义模型”填写名称。 @@ -102,6 +109,7 @@ interface ModelSettings { `] }) export class ModelSettingsComponent implements OnInit { + transport: 'http' | 'websocket' = 'http'; readonly reasoningValues = REASONING_VALUES; readonly reasoningLabels = REASONING_LABELS; private http = inject(HttpClient); @@ -128,7 +136,7 @@ export class ModelSettingsComponent implements OnInit { && Number.isInteger(this.toolFallbackNoProgress) && this.toolFallbackNoProgress >= 1 && this.toolFallbackNoProgress <= 10; } private payload() { - return { base_url: this.baseUrl.trim(), model: this.model.trim(), api_key: this.apiKey.trim() || undefined, + return { base_url: this.baseUrl.trim(), transport: this.transport, model: this.model.trim(), api_key: this.apiKey.trim() || undefined, timeout_seconds: this.timeout, max_retries: this.retries, retry_interval_seconds: this.retryInterval, reasoning_effort: this.reasoning, rate_limit_retry_seconds: this.rateLimitRetry, tool_fallback_model: this.toolFallbackModel.trim(), tool_fallback_reasoning_effort: this.toolFallbackReasoning, tool_fallback_max_retries: this.toolFallbackRetries, @@ -136,6 +144,7 @@ export class ModelSettingsComponent implements OnInit { } private apply(config: ModelSettings): void { this.baseUrl = config.base_url; this.model = config.model; this.timeout = config.timeout_seconds; + this.transport = config.transport ?? 'http'; this.modelPresets = config.model_presets; this.retries = config.max_retries; this.retryInterval = config.retry_interval_seconds ?? 10; this.reasoning = config.reasoning_effort ?? 'high'; this.keyConfigured.set(config.api_key_configured); this.envPath.set(config.env_path); this.rateLimitRetry = config.rate_limit_retry_seconds ?? 60; diff --git a/artemis/config/llm.py b/artemis/config/llm.py index 47c7293..fbd53bf 100644 --- a/artemis/config/llm.py +++ b/artemis/config/llm.py @@ -62,6 +62,9 @@ class LLM(BaseModel): provider: LLMProvider = "openai" model: str = Field(default_factory=lambda: settings.OPENAI_MODEL) api_base: str | None = None + transport: Literal["http", "websocket"] = Field( + default_factory=lambda: settings.OPENAI_TRANSPORT + ) timeout: float = Field(default_factory=lambda: settings.OPENAI_TIMEOUT_SECONDS, gt=0) max_retries: int = Field(default_factory=lambda: settings.OPENAI_MAX_RETRIES, ge=0) max_tokens: int | None = Field(default=None, gt=0) diff --git a/artemis/config/settings.py b/artemis/config/settings.py index 093cff0..cab4cdc 100644 --- a/artemis/config/settings.py +++ b/artemis/config/settings.py @@ -16,6 +16,7 @@ import os from pathlib import Path +from typing import Literal from dotenv import load_dotenv from pydantic import Field, SecretStr, model_validator @@ -109,6 +110,7 @@ class Settings(BaseSettings): OPENAI_BASE_URL: str | None = "https://api-n-cd.playones.com/v1" OPENAI_MODEL: str = Field(default=DEFAULT_MODEL, min_length=1) OPENAI_TIMEOUT_SECONDS: float = Field(default=300.0, gt=0) + OPENAI_TRANSPORT: Literal["http", "websocket"] = "http" # Environment only; deliberately absent from the saved UI/task configuration. OPENAI_USE_SYSTEM_PROXY: bool = False OPENAI_MAX_RETRIES: int = Field(default=5, ge=0, le=10) diff --git a/artemis/llm/diagnostics.py b/artemis/llm/diagnostics.py index c11e627..c3f9a39 100644 --- a/artemis/llm/diagnostics.py +++ b/artemis/llm/diagnostics.py @@ -390,7 +390,12 @@ async def observe_response_headers(response): """Correlate allowlisted headers and passively observe bounded SSE frames.""" attempt = ACTIVE_ATTEMPT.get() if attempt is not None: - attempt.headers(response) + websocket = response.extensions.get("artemis_websocket", False) + if websocket: + # 200/SSE is an in-process SDK facade, not a wire HTTP response. + attempt.data.update(headers_at=time.time(), phase="waiting_output") + else: + attempt.headers(response) content_type = response.headers.get("content-type", "").split(";", 1)[0].lower() if ( response.status_code == 200 @@ -407,5 +412,7 @@ async def observe_response_headers(response): ) response.stream = stream attempt.wire_streams.append(stream) + if websocket: + stream.observation.data["observation_boundary"] = "websocket_event_sse_adapter" else: attempt.data["wire_streams_truncated"] = True diff --git a/artemis/llm/router.py b/artemis/llm/router.py index 3684dab..66ef7b7 100644 --- a/artemis/llm/router.py +++ b/artemis/llm/router.py @@ -17,7 +17,7 @@ from enum import StrEnum import hashlib import os -from typing import Any +from typing import Any, Literal from langchain_core.language_models.chat_models import BaseChatModel from pydantic import BaseModel, Field @@ -40,6 +40,9 @@ class ModelEndpoint(BaseModel): model_name: str = Field(default_factory=lambda: settings.OPENAI_MODEL) api_key: str | None = Field(default=None, repr=False) api_base: str | None = None + transport: Literal["http", "websocket"] = Field( + default_factory=lambda: settings.OPENAI_TRANSPORT + ) temperature: float | None = None max_tokens: int | None = Field(default=None, gt=0) timeout_seconds: float = Field(default_factory=lambda: settings.OPENAI_TIMEOUT_SECONDS, gt=0) @@ -63,6 +66,7 @@ class ModelEndpoint(BaseModel): self.enable_grounding, self.api_base or settings.OPENAI_BASE_URL, settings.OPENAI_USE_SYSTEM_PROXY, + self.transport, hashlib.sha256(key.encode()).hexdigest(), ) @@ -119,13 +123,32 @@ class ModelFactory: "http_client": DefaultHttpxClient( timeout=endpoint.timeout_seconds, **model_http_options() ), - "http_async_client": DefaultAsyncHttpxClient( + } + if endpoint.transport == "websocket": + import httpx + + from artemis.llm.websocket import ResponsesWebSocketTransport + + # Disable HTTP proxy mounts on this facade: they would bypass the + # WebSocket transport. Each actual transport owns its proxy policy. + kwargs["http_async_client"] = httpx.AsyncClient( + timeout=endpoint.timeout_seconds, + trust_env=False, + transport=ResponsesWebSocketTransport( + timeout=endpoint.timeout_seconds, use_proxy=settings.OPENAI_USE_SYSTEM_PROXY + ), + event_hooks={ + "request": [observe_request_headers], + "response": [observe_response_headers], + }, + ) + else: + kwargs["http_async_client"] = DefaultAsyncHttpxClient( timeout=endpoint.timeout_seconds, **model_http_options(), event_hooks={ "request": [observe_request_headers], "response": [observe_response_headers], }, - ), - } + ) return ChatOpenAI(**{k: v for k, v in kwargs.items() if v is not None}) diff --git a/artemis/llm/websocket.py b/artemis/llm/websocket.py new file mode 100644 index 0000000..29c3655 --- /dev/null +++ b/artemis/llm/websocket.py @@ -0,0 +1,449 @@ +"""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() diff --git a/artemis/resources/config/artemis.jsonc b/artemis/resources/config/artemis.jsonc index cd43142..2803aca 100644 --- a/artemis/resources/config/artemis.jsonc +++ b/artemis/resources/config/artemis.jsonc @@ -2,6 +2,7 @@ { "default": { "provider": "openai", + "transport": "${env:OPENAI_TRANSPORT:-http}", "model": "${env:OPENAI_MODEL:-gpt-6-luna}", "timeout": "${env:OPENAI_TIMEOUT_SECONDS:-300}", "max_retries": "${env:OPENAI_MAX_RETRIES:-5}", diff --git a/artemis/services/llm.py b/artemis/services/llm.py index c47865e..e853e9f 100644 --- a/artemis/services/llm.py +++ b/artemis/services/llm.py @@ -1174,6 +1174,7 @@ def _resolve_endpoint( else settings.OPENAI_MAX_RETRIES, max_tokens=_get_val(cfg, "max_tokens", int), api_base=_get_val(cfg, "api_base", str), + transport=_get_val(cfg, "transport", str) or settings.OPENAI_TRANSPORT, reasoning_effort=_get_val(cfg, "reasoning_effort", str), enable_grounding=_get_val(cfg, "enable_grounding", bool) or False, ) diff --git a/config/artemis.jsonc b/config/artemis.jsonc index cd43142..2803aca 100644 --- a/config/artemis.jsonc +++ b/config/artemis.jsonc @@ -2,6 +2,7 @@ { "default": { "provider": "openai", + "transport": "${env:OPENAI_TRANSPORT:-http}", "model": "${env:OPENAI_MODEL:-gpt-6-luna}", "timeout": "${env:OPENAI_TIMEOUT_SECONDS:-300}", "max_retries": "${env:OPENAI_MAX_RETRIES:-5}", diff --git a/docs/ai-automation.md b/docs/ai-automation.md index 2f06001..eb3a0c4 100644 --- a/docs/ai-automation.md +++ b/docs/ai-automation.md @@ -169,6 +169,18 @@ Flash、Pro（启用 transcript）和探索共用分阶段上下文。一个阶 辅助摘要继续独立处理各自片段，避免为了提高命中比例累积不相关的历史。视觉摘要中完全重复的插件执行结果与控件树只发送一份；独立错误、状态、不同观察及原始 trace 均保留。 +## Responses WebSocket 增量传输 + +网页“模型配置 → 传输协议”选择 `WebSocket（增量传输）`，或在 `.env` 设置 `OPENAI_TRANSPORT=websocket`。网页保存对新任务生效；直接编辑 `.env` 后重启相应服务。默认 `http`，旧配置不变。JSON 模型/节点配置可用 `transport` 单独覆盖；MCP 使用 `mobile_config(section="model", changes={"transport":"websocket"}, persist_defaults=true)` 保存，仍须用户明确要求修改默认值。 + +异步任务请求在现有 SDK 序列化与完整性保护之下接入 `/v1/responses` WebSocket。首轮发送完整上下文；同一任务（或探索）、角色、模型、路由头及凭据范围内复用连接。只在当前输入严格包含上一轮输入和完整输出、生成选项一致时，发送 `previous_response_id` 与新增输入。仅忽略消息 ID、状态、输出概率等不影响历史语义的响应字段，不按文本相似度截取。工具结果保留 `call_id`，文本、图片、结构化输出与 usage 继续走原有解析和安全检查。同步底层模型调用仍使用 HTTP；Flash、Pro、探索和辅助模型运行路径均为异步。 + +本地历史、截图替换或压缩后重新发送当前完整上下文；不会为了续传保留已经移除的旧界面索引。连接重建或服务端明确返回 `previous_response_not_found` 时，从本地完整历史恢复，后者最多补发一次且必须尚未出现输出。连接在输出中途断开时，不在传输层盲目重发，交给现有模型重试预算；取消关闭未完成连接，不采纳半截工具参数或结论。不同会话并行、同一连接串行；连接池最多保留 32 个会话，空闲 60 秒释放，连接满 55 分钟主动重建。`store=false` 保持不变。 + +网关必须支持 WebSocket Upgrade、双向长连接、`response.create` 与响应 ID 续接；只透传会话头不能证明支持。握手传递现有授权和路由头，保持 TLS/自定义 CA 验证与 `OPENAI_USE_SYSTEM_PROXY` 策略，拒绝携带凭据跨地址重定向。不支持升级时明确报错，可在网页改回 HTTP，不会静默改成 HTTP 全量请求。 + +请求诊断的 `transport` 记录 `full_request_bytes`、`sent_payload_bytes`、`received_payload_bytes`、`payload_bytes_saved`、`incremental`、`reused_input_items`、`connection_reused`、`create_events` 和重建原因。字节数衡量 JSON 应用载荷，不包含 TLS/WebSocket 帧和压缩开销；SDK 内部兼容 SSE 的字节不冒充网络流量或 HTTP 200。握手上的请求标识属于建连尝试，后续轮次以本地请求诊断及响应 ID 关联。减少上传不等于减少模型上下文 token，也不等于缓存命中，usage 仍以服务端实报为准。协议见 [OpenAI 官方说明](https://developers.openai.com/api/docs/guides/websocket-mode)。 + ## 指定任务诊断、复盘与管理 流式请求中的 `click`、`long_press`、`swipe`、`press_key`、`manage_app`、`wait_for_delay` 每次调用参数超过 16,384 字符时，会立即关闭该请求并丢弃未完成输出，避免持续生成异常参数直到超时。该保护只针对小型设备动作，不限制输入文本、笔记、报告或插件参数，也不改变模型思考强度与输出 token 配置。异常最多额外尝试一次（配置为不重试时不重试）；仍异常时进入现有备用模型或暂停恢复流程。诊断记录保留 `invalid_output` 分类及 `stream_guard` 中的工具名、实际字符数和保护上限，不保存参数正文；无法据此断定异常来自模型生成还是 API 中转。 diff --git a/mcp_server/catalog.json b/mcp_server/catalog.json index 8792315..f93c364 100644 --- a/mcp_server/catalog.json +++ b/mcp_server/catalog.json @@ -437,7 +437,7 @@ }, { "name": "mobile_config", - "description": "Read configuration or explicitly save non-secret defaults for FUTURE tasks.\n\nOnly set persist_defaults=true when the user requests saved/global configuration.\nmodel changes: model (preset or custom), base_url, reasoning_effort\n(default/none/minimal/low/medium/high/xhigh/max), timeout_seconds, max_retries,\nretry_interval_seconds, rate_limit_retry_seconds, tool_fallback_model,\ntool_fallback_reasoning_effort, tool_fallback_max_retries,\ntool_fallback_no_progress_attempts. 'default' omits reasoning;\nunsupported values fail visibly.\ncapabilities: video_enabled and recording_enabled are independent; outputter_enabled,\nflash_final_check_enabled (false; Flash exit-only AI review),\nstep_summarizer_enabled, screenshot_limit_enabled, task_timeout_seconds (default 1800),\nobservation_timeout_seconds (60), observation_max_attempts (3),\ndevice_command_timeout_seconds (10), action_failure_attempts (3),\nworker_cleanup_timeout_seconds (30), adb_auto_connect (true),\nadb_discovery_interval_seconds (10), adb_discovery_timeout_seconds (5).\nreports: formats=['docx','xlsx'] (at least one).\nexploration: model (default gpt-6.1-terra), reasoning_effort (high), timeout_seconds\n(1800 active seconds), demonstration_timeout_seconds (900), demonstration_warning_seconds (60).\nAPI keys are never accepted or returned: use the Web UI or .env for credentials.\nSaving settings does not change an already running task. Flash/Pro are execution\nprofiles, not model names. Model retries do not repeat device test cases.\n", + "description": "Read configuration or explicitly save non-secret defaults for FUTURE tasks.\n\nOnly set persist_defaults=true when the user requests saved/global configuration.\nmodel changes: model (preset or custom), base_url, transport (http/websocket), reasoning_effort\n(default/none/minimal/low/medium/high/xhigh/max), timeout_seconds, max_retries,\nretry_interval_seconds, rate_limit_retry_seconds, tool_fallback_model,\ntool_fallback_reasoning_effort, tool_fallback_max_retries,\ntool_fallback_no_progress_attempts. 'default' omits reasoning;\nunsupported values fail visibly.\ncapabilities: video_enabled and recording_enabled are independent; outputter_enabled,\nflash_final_check_enabled (false; Flash exit-only AI review),\nstep_summarizer_enabled, screenshot_limit_enabled, task_timeout_seconds (default 1800),\nobservation_timeout_seconds (60), observation_max_attempts (3),\ndevice_command_timeout_seconds (10), action_failure_attempts (3),\nworker_cleanup_timeout_seconds (30), adb_auto_connect (true),\nadb_discovery_interval_seconds (10), adb_discovery_timeout_seconds (5).\nreports: formats=['docx','xlsx'] (at least one).\nexploration: model (default gpt-6.1-terra), reasoning_effort (high), timeout_seconds\n(1800 active seconds), demonstration_timeout_seconds (900), demonstration_warning_seconds (60).\nAPI keys are never accepted or returned: use the Web UI or .env for credentials.\nSaving settings does not change an already running task. Flash/Pro are execution\nprofiles, not model names. Model retries do not repeat device test cases.\n", "inputSchema": { "properties": { "section": { diff --git a/mcp_server/tools/automation.py b/mcp_server/tools/automation.py index ae45dae..ef67993 100644 --- a/mcp_server/tools/automation.py +++ b/mcp_server/tools/automation.py @@ -14,6 +14,7 @@ from mcp_server.base import mcp EXPORT_DIR = ROOT_DIR / "artifacts" / "mcp_exports" MODEL_FIELDS = { + "transport", "base_url", "model", "timeout_seconds", @@ -186,7 +187,7 @@ async def mobile_config( """Read configuration or explicitly save non-secret defaults for FUTURE tasks. Only set persist_defaults=true when the user requests saved/global configuration. - model changes: model (preset or custom), base_url, reasoning_effort + model changes: model (preset or custom), base_url, transport (http/websocket), reasoning_effort (default/none/minimal/low/medium/high/xhigh/max), timeout_seconds, max_retries, retry_interval_seconds, rate_limit_retry_seconds, tool_fallback_model, tool_fallback_reasoning_effort, tool_fallback_max_retries, diff --git a/pyproject.toml b/pyproject.toml index 1df42e2..9cf1697 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -57,7 +57,7 @@ dependencies = [ "uvicorn[standard]>=0.30.1", "httpx>=0.28.1", "requests>=2.31.0", - "websockets>=12.0", + "websockets>=15.0", "sseclient-py>=1.8.0", # Image, Video & Multimodal Perception "pillow>=13.0.0", diff --git a/skills/artemis-use/references/modes.md b/skills/artemis-use/references/modes.md index 017d592..a249dc9 100644 --- a/skills/artemis-use/references/modes.md +++ b/skills/artemis-use/references/modes.md @@ -16,6 +16,8 @@ ## 模式、模型与配置 +网页模型配置和 `mobile_config(section="model")` 提供 `transport=http/websocket`。用户明确要求保存默认值时可设置 `changes={"transport":"websocket"}, persist_defaults=true`；新任务的 Flash、Pro、探索和辅助请求共用此选择。WebSocket 在历史一致时仅上传新增输入，网关须支持 Responses WebSocket；压缩、重连会发送完整上下文，上传减少不能当作 token 减少或缓存命中。不支持协议时应查看错误并确认网关能力，不自动改写用户配置。 + - `mobile_run_task.model` 取 `Flash` / `Pro`。API 模型名放 `runtime_overrides.model`；思考强度放 `runtime_overrides.reasoning_effort`。批次使用 `task_config.profile` 选择模式，`task_config.model` 选择 API 模型。 - 探索独立配置，默认 `gpt-6.1-terra/high`，支持 `gpt-6-astra` 及已配置网关支持的自定义模型。执行、示范分析、提炼共用该探索配置；不改变 Flash/Pro 默认值。 - 默认 high、单次模型请求时限 300 秒。`default` 表示不发送思考强度，`none` 是显式关闭，不能混为一谈；接口不支持参数时如实报告，不擅自降级。 diff --git a/tests/unit/admin_console/test_model_settings.py b/tests/unit/admin_console/test_model_settings.py index f35a384..8641cee 100644 --- a/tests/unit/admin_console/test_model_settings.py +++ b/tests/unit/admin_console/test_model_settings.py @@ -21,6 +21,7 @@ def isolated_settings(monkeypatch, tmp_path): "OPENAI_MODEL", "OPENAI_BASE_URL", "OPENAI_TIMEOUT_SECONDS", + "OPENAI_TRANSPORT", "OPENAI_MAX_RETRIES", "OPENAI_RETRY_INTERVAL_SECONDS", "OPENAI_RATE_LIMIT_RETRY_SECONDS", @@ -50,6 +51,7 @@ def isolated_settings(monkeypatch, tmp_path): def test_save_reload_and_blank_key_preserves_secret(isolated_settings): update = service.ModelSettingsUpdate( + transport="websocket", base_url="https://gateway.example/v1/", model="my-model", api_key="new-test-secret", @@ -63,6 +65,9 @@ def test_save_reload_and_blank_key_preserves_secret(isolated_settings): ) result = service.save_model_settings(update) assert result["api_key_configured"] + assert result["transport"] == "websocket" + assert dotenv_values(isolated_settings)["OPENAI_TRANSPORT"] == "websocket" + assert service.ModelSettingsUpdate(model="test").transport is None assert "new-test-secret" not in repr(result) saved = dotenv_values(isolated_settings) assert saved["OPENAI_API_KEY"] == "new-test-secret" diff --git a/tests/unit/mcp/test_automation_tools.py b/tests/unit/mcp/test_automation_tools.py index 05be9af..e02b004 100644 --- a/tests/unit/mcp/test_automation_tools.py +++ b/tests/unit/mcp/test_automation_tools.py @@ -72,6 +72,8 @@ async def test_config_requires_explicit_persistence_and_excludes_secrets(monkeyp payload = request.call_args.kwargs["json"] assert payload["rate_limit_retry_seconds"] == 90 assert payload["tool_fallback_model"] == "gpt-6-luna" + await tools.mobile_config(changes={"transport": "websocket"}, persist_defaults=True) + assert request.call_args.kwargs["json"]["transport"] == "websocket" @pytest.mark.asyncio diff --git a/tests/unit/test_responses_websocket.py b/tests/unit/test_responses_websocket.py new file mode 100644 index 0000000..d06e11e --- /dev/null +++ b/tests/unit/test_responses_websocket.py @@ -0,0 +1,605 @@ +"""Real local WebSocket/SDK round trips; no device, credential, or paid endpoint.""" + +import asyncio +from contextlib import asynccontextmanager +from copy import deepcopy +import json + +import httpx +from langchain_core.messages import HumanMessage, ToolMessage +import pytest +import pytest_asyncio +from websockets.asyncio.server import serve + +from artemis.llm.diagnostics import ACTIVE_ATTEMPT +from artemis.llm.response_guard import ResponseIntegrityError, ensure_response_complete +from artemis.llm.router import ModelEndpoint, ModelFactory +from artemis.services.llm import RobustChatModelWrapper + + +def text_output(text="OK"): + return [ + { + "id": "msg_test", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": text, "annotations": [], "logprobs": []}], + } + ] + + +class Wire: + def __init__(self): + self.requests = [] + self.connections = [] + self.outputs = [] + self.reject_previous = False + self.drop = False + self.stall = False + self.status = "completed" + self.sent = asyncio.Event() + self.gate = None + + async def handle(self, socket): + self.connections.append(socket) + async for raw in socket: + body = json.loads(raw) + self.requests.append(body) + self.sent.set() + if self.reject_previous and body.get("previous_response_id"): + await socket.send( + json.dumps( + { + "type": "error", + "status": 400, + "error": {"code": "previous_response_not_found"}, + } + ) + ) + continue + response = { + "id": f"resp_{len(self.requests)}", + "object": "response", + "created_at": 1, + "status": "in_progress", + "model": "test-model", + "output": [], + } + await socket.send( + json.dumps({"type": "response.created", "sequence_number": 0, "response": response}) + ) + if self.drop: + await socket.close() + continue + if self.stall: + await socket.wait_closed() + return + if self.gate: + await self.gate.wait() + output = self.outputs.pop(0) if self.outputs else text_output() + seq = 0 + for index, item in enumerate(output): + seq += 1 + await socket.send( + json.dumps( + { + "type": "response.output_item.added", + "sequence_number": seq, + "output_index": index, + "item": {**item, "content": [], "arguments": ""}, + } + ) + ) + for content_index, block in enumerate(item.get("content", [])): + if block.get("type") != "output_text": + continue + seq += 1 + await socket.send( + json.dumps( + { + "type": "response.output_text.delta", + "sequence_number": seq, + "output_index": index, + "content_index": content_index, + "item_id": item["id"], + "delta": block["text"], + } + ) + ) + seq += 1 + await socket.send( + json.dumps( + { + "type": "response.output_text.done", + "sequence_number": seq, + "output_index": index, + "content_index": content_index, + "item_id": item["id"], + "text": block["text"], + } + ) + ) + seq += 1 + await socket.send( + json.dumps( + { + "type": "response.output_item.done", + "sequence_number": seq, + "output_index": index, + "item": item, + } + ) + ) + response.update( + status=self.status, + output=output, + usage={ + "input_tokens": 20, + "output_tokens": 3, + "total_tokens": 23, + "input_tokens_details": {"cached_tokens": 10}, + }, + ) + await socket.send( + json.dumps( + { + "type": f"response.{self.status}", + "sequence_number": seq + 1, + "response": response, + } + ) + ) + + +@pytest_asyncio.fixture +async def wire(monkeypatch): + monkeypatch.delenv("ARTEMIS_FAKE_LLM", raising=False) + instance = Wire() + async with serve(instance.handle, "127.0.0.1", 0) as server: + instance.base = f"http://127.0.0.1:{server.sockets[0].getsockname()[1]}/v1" + yield instance + + +@asynccontextmanager +async def model(wire, **kwargs): + endpoint = ModelEndpoint( + model_name="test-model", + api_key="test-secret", + api_base=wire.base, + transport="websocket", + max_retries=0, + timeout_seconds=2, + **kwargs, + ) + value = ModelFactory.create_model(endpoint) + try: + yield value + finally: + await value.http_async_client.aclose() + value.http_client.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +async def test_sdk_roundtrip_incremental_history_and_usage(wire, stream): + async with model(wire) as raw: + messages = [ + HumanMessage( + content=[ + {"type": "text", "text": "Inspect"}, + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64," + "A" * 80000}, + }, + ] + ) + ] + + async def call(history): + if not stream: + return await raw.ainvoke(history, prompt_cache_key="task:operator") + chunks = [ + chunk async for chunk in raw.astream(history, prompt_cache_key="task:operator") + ] + return sum(chunks[1:], chunks[0]) + + first = await call(messages) + ensure_response_complete(first) + second = await call([*messages, first, HumanMessage(content="Continue")]) + ensure_response_complete(second) + assert second.usage_metadata["input_token_details"]["cache_read"] == 10 + assert wire.requests[1].get("previous_response_id") == "resp_1", { + "first": first.content, + "tail": wire.requests[1]["input"][1:], + } + assert wire.requests[1]["input"] == [ + {"type": "message", "role": "user", "content": "Continue"} + ] + assert "stream" not in wire.requests[0] + assert wire.requests[0]["type"] == "response.create" + assert wire.requests[0]["store"] is False + assert len(json.dumps(wire.requests[1])) < len(json.dumps(wire.requests[0])) / 100 + assert len(wire.connections) == 1 + assert wire.connections[0].request.path == "/v1/responses" + assert wire.connections[0].request.headers["Authorization"] == "Bearer test-secret" + + +@pytest.mark.asyncio +async def test_tool_call_chain_preserves_call_id(wire): + wire.outputs.append( + [ + { + "id": "fc_test", + "type": "function_call", + "call_id": "call_test", + "name": "lookup", + "arguments": '{"x":1}', + "status": "completed", + } + ] + ) + async with model(wire) as raw: + messages = [HumanMessage(content="Look up")] + first = await raw.ainvoke(messages, prompt_cache_key="task:operator") + await raw.ainvoke( + [*messages, first, ToolMessage(content="found", tool_call_id="call_test")], + prompt_cache_key="task:operator", + ) + assert wire.requests[1]["previous_response_id"] == "resp_1" + assert wire.requests[1]["input"] == [ + {"type": "function_call_output", "call_id": "call_test", "output": "found"} + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "change", ["history", "options", "scope", "new_connection", "server_closed"] +) +async def test_chain_resets_on_context_or_connection_change(wire, change): + async with model(wire) as raw: + messages = [HumanMessage(content="Original")] + first = await raw.ainvoke(messages, prompt_cache_key="task:operator") + history = [*messages, first, HumanMessage(content="Next")] + options = {"prompt_cache_key": "task:operator"} + if change == "history": + history[0] = HumanMessage(content="Compacted") + if change == "options": + options["temperature"] = 0.3 + if change == "scope": + options["prompt_cache_key"] = "task:checker" + if change == "new_connection": + await raw.http_async_client._transport.sessions[ + next(iter(raw.http_async_client._transport.sessions)) + ].close() + if change == "server_closed": + await wire.connections[0].close() + await wire.connections[0].wait_closed() + await raw.ainvoke(history, **options) + assert "previous_response_id" not in wire.requests[1] + assert len(wire.requests[1]["input"]) == 3 + + +@pytest.mark.asyncio +async def test_uncached_id_recovers_once_with_full_context(wire): + async with model(wire) as raw: + history = [HumanMessage(content="First")] + first = await raw.ainvoke(history, prompt_cache_key="task:operator") + wire.reject_previous = True + result = await raw.ainvoke( + [*history, first, HumanMessage(content="Next")], prompt_cache_key="task:operator" + ) + ensure_response_complete(result) + assert len(wire.requests) == 3 + assert wire.requests[1]["previous_response_id"] == "resp_1" + assert "previous_response_id" not in wire.requests[2] + assert len(wire.requests[2]["input"]) == 3 + + +@pytest.mark.asyncio +async def test_missing_affinity_never_shares_connections(wire): + async with model(wire) as raw: + history = [HumanMessage(content="First")] + first = await raw.ainvoke(history) + await raw.ainvoke([*history, first, HumanMessage(content="Next")]) + assert len(wire.connections) == 2 + assert "previous_response_id" not in wire.requests[1] + + +@pytest.mark.asyncio +async def test_cancel_closes_socket_and_next_call_starts_full(wire): + async with model(wire) as raw: + wire.stall = True + task = asyncio.create_task( + raw.ainvoke([HumanMessage(content="Wait")], prompt_cache_key="task:operator") + ) + await asyncio.wait_for(wire.sent.wait(), 2) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + await asyncio.wait_for(wire.connections[0].wait_closed(), 2) + wire.stall = False + await raw.ainvoke([HumanMessage(content="Next")], prompt_cache_key="task:operator") + assert len(wire.connections) == 2 + assert "previous_response_id" not in wire.requests[1] + + +@pytest.mark.asyncio +async def test_disconnect_never_silently_resends_model_request(wire): + async with model(wire) as raw: + wire.drop = True + with pytest.raises(Exception, match="connection|Connection"): + await raw.ainvoke([HumanMessage(content="Test")], prompt_cache_key="task:operator") + assert len(wire.requests) == 1 + + +@pytest.mark.asyncio +async def test_incomplete_result_does_not_advance_chain(wire): + async with model(wire) as raw: + wire.status = "incomplete" + first = await raw.ainvoke([HumanMessage(content="First")], prompt_cache_key="task:operator") + with pytest.raises(ResponseIntegrityError): + ensure_response_complete(first) + wire.status = "completed" + await raw.ainvoke( + [HumanMessage(content="First"), first, HumanMessage(content="Next")], + prompt_cache_key="task:operator", + ) + assert "previous_response_id" not in wire.requests[1] + + +@pytest.mark.asyncio +async def test_diagnostics_measure_actual_wire_payload_without_logging_content(wire): + class Attempt: + data = {} + + def publish_if_due(self): + pass + + attempt = Attempt() + async with model(wire) as raw: + # Disable the unrelated HTTP trace hook: this checks the transport's + # own byte counters, alongside the real SDK request serialization. + raw.http_async_client.event_hooks = {"request": [], "response": []} + first = await raw.ainvoke( + [HumanMessage(content="A" * 5000)], prompt_cache_key="task:operator" + ) + token = ACTIVE_ATTEMPT.set(attempt) + try: + await raw.ainvoke( + [HumanMessage(content="A" * 5000), first, HumanMessage(content="Next")], + prompt_cache_key="task:operator", + ) + finally: + ACTIVE_ATTEMPT.reset(token) + stats = deepcopy(attempt.data["transport"]) + assert stats["incremental"] is True + assert stats["sent_payload_bytes"] == len( + json.dumps( + wire.requests[1], ensure_ascii=False, separators=(",", ":"), sort_keys=True + ).encode() + ) + assert stats["payload_bytes_saved"] > 4000 + assert "AAAA" not in json.dumps(stats) + + +@pytest.mark.asyncio +async def test_gateway_rejects_upgrade_with_explicit_error(wire): + import openai + from websockets.http11 import Response + from websockets.datastructures import Headers + + async with serve( + wire.handle, + "127.0.0.1", + 0, + process_request=lambda *_: Response(404, "Not Found", Headers(), b""), + ) as server: + wire.base = f"http://127.0.0.1:{server.sockets[0].getsockname()[1]}/v1" + async with model(wire) as raw: + with pytest.raises(openai.NotFoundError, match="WebSocket"): + await raw.ainvoke([HumanMessage(content="Test")]) + + +@pytest.mark.asyncio +async def test_wrapper_retains_stream_guards_and_diagnostic_events(wire, monkeypatch): + from artemis.services import llm + from tests.unit.test_responses_wire_lifecycle import Engine + + engine = Engine() + monkeypatch.setattr(llm, "_get_current_data_engine", lambda: engine) + async with model(wire) as raw: + wrapper = RobustChatModelWrapper( + raw, + endpoint=ModelEndpoint(model_name="test-model", transport="websocket", max_retries=0), + ) + result = await wrapper.ainvoke( + [HumanMessage(content="Test")], prompt_cache_key="task:operator" + ) + ensure_response_complete(result) + assert engine.data["transport"]["mode"] == "websocket" + assert engine.data["transport"]["received_payload_bytes"] > 0 + assert engine.data.get("http_status") is None + + +@pytest.mark.asyncio +async def test_parallel_scopes_keep_independent_response_chains(wire): + async with model(wire) as raw: + + async def conversation(scope): + history = [HumanMessage(content=f"First {scope}")] + first = await raw.ainvoke(history, prompt_cache_key=scope) + await raw.ainvoke( + [*history, first, HumanMessage(content=f"Next {scope}")], prompt_cache_key=scope + ) + return scope, first.response_metadata["id"] + + results = await asyncio.gather( + conversation("task-a:operator"), conversation("task-b:operator") + ) + assert len(wire.connections) == 2 + for scope, response_id in results: + calls = [item for item in wire.requests if item["prompt_cache_key"] == scope] + assert len(calls) == 2 + assert calls[1]["previous_response_id"] == response_id + assert len(calls[1]["input"]) == 1 + + +@pytest.mark.asyncio +async def test_idle_connections_are_closed_and_removed(wire, monkeypatch): + from artemis.llm import websocket + + monkeypatch.setattr(websocket, "IDLE_SECONDS", 0.01) + async with model(wire) as raw: + await raw.ainvoke([HumanMessage(content="Test")], prompt_cache_key="task:operator") + await asyncio.wait_for(wire.connections[0].wait_closed(), 2) + assert not raw.http_async_client._transport.sessions + + +@pytest.mark.asyncio +async def test_cancel_waiting_turn_keeps_active_chain_and_idle_cleanup(wire, monkeypatch): + from artemis.llm import websocket + + monkeypatch.setattr(websocket, "IDLE_SECONDS", 0.05) + wire.gate = asyncio.Event() + async with model(wire) as raw: + history = [HumanMessage(content="First")] + first = asyncio.create_task(raw.ainvoke(history, prompt_cache_key="task:operator")) + await asyncio.wait_for(wire.sent.wait(), 2) + queued = asyncio.create_task(raw.ainvoke(history, prompt_cache_key="task:operator")) + async with asyncio.timeout(2): + while next(iter(raw.http_async_client._transport.sessions.values())).users != 2: + await asyncio.sleep(0) + queued.cancel() + with pytest.raises(asyncio.CancelledError): + await queued + wire.gate.set() + result = await first + await raw.ainvoke( + [*history, result, HumanMessage(content="Next")], prompt_cache_key="task:operator" + ) + assert len(wire.connections) == 1 + assert len(wire.requests) == 2 + assert wire.requests[1]["previous_response_id"] == result.response_metadata["id"] + await asyncio.wait_for(wire.connections[0].wait_closed(), 2) + assert not raw.http_async_client._transport.sessions + + +@pytest.mark.asyncio +async def test_pool_eviction_does_not_split_concurrent_scope(wire, monkeypatch): + from artemis.llm import websocket + + monkeypatch.setattr(websocket, "MAX_SESSIONS", 1) + closing, resume = asyncio.Event(), asyncio.Event() + async with model(wire) as raw: + await raw.ainvoke([HumanMessage(content="Old")], prompt_cache_key="old") + old = next(iter(raw.http_async_client._transport.sessions.values())) + close = old.close + + async def delayed_close(): + closing.set() + await resume.wait() + await close() + + monkeypatch.setattr(old, "close", delayed_close) + first = asyncio.create_task( + raw.ainvoke([HumanMessage(content="New first")], prompt_cache_key="new") + ) + try: + await asyncio.wait_for(closing.wait(), 2) + await asyncio.wait_for( + asyncio.gather( + first, + raw.ainvoke([HumanMessage(content="New second")], prompt_cache_key="new"), + ), + 2, + ) + assert len(wire.connections) == 2 + assert len(wire.requests) == 3 + finally: + resume.set() + + +@pytest.mark.asyncio +async def test_structured_output_uses_same_websocket_path(wire): + from pydantic import BaseModel + + class Result(BaseModel): + ok: bool + + wire.outputs.append( + [ + { + "id": "fc_test", + "type": "function_call", + "call_id": "call_test", + "name": "Result", + "arguments": '{"ok":true}', + "status": "completed", + } + ] + ) + async with model(wire) as raw: + result = ( + await RobustChatModelWrapper( + raw, endpoint=ModelEndpoint(transport="websocket", max_retries=0) + ) + .with_structured_output(Result) + .ainvoke([HumanMessage(content="Test")]) + ) + assert result.ok is True + assert len(wire.requests) == 1 + + +@pytest.mark.asyncio +async def test_invalid_tool_arguments_are_rejected_before_use(wire, monkeypatch): + from artemis.services import llm + from artemis.llm.reliability import LLMExhaustedError + + monkeypatch.setattr(llm, "_get_current_data_engine", lambda: None) + wire.outputs.append( + [ + { + "id": "fc_test", + "type": "function_call", + "call_id": "call_test", + "name": "lookup", + "arguments": '{"x":', + "status": "completed", + } + ] + ) + async with model(wire) as raw: + wrapper = RobustChatModelWrapper( + raw, endpoint=ModelEndpoint(transport="websocket", max_retries=0) + ) + with pytest.raises(LLMExhaustedError): + await wrapper.ainvoke([HumanMessage(content="Test")], prompt_cache_key="task:operator") + assert len(wire.requests) == 1 + + +@pytest.mark.asyncio +async def test_redirect_does_not_forward_authorization(wire): + from websockets.http11 import Response + from websockets.datastructures import Headers + + async with serve( + wire.handle, + "127.0.0.1", + 0, + process_request=lambda *_: Response( + 307, + "Redirect", + Headers({"Location": wire.base.replace("http://", "ws://") + "/responses"}), + b"", + ), + ) as server: + redirect_base = f"http://127.0.0.1:{server.sockets[0].getsockname()[1]}/v1" + original_base, wire.base = wire.base, redirect_base + try: + async with model(wire) as raw: + with pytest.raises(Exception, match="redirects are disabled"): + await raw.ainvoke([HumanMessage(content="Test")]) + assert not wire.connections + finally: + wire.base = original_base diff --git a/uv.lock b/uv.lock index 4a65727..7535c0e 100644 --- a/uv.lock +++ b/uv.lock @@ -369,7 +369,7 @@ requires-dist = [ { name = "uiautomator2", specifier = ">=3.5.0" }, { name = "uuid-utils", specifier = ">=0.9.0" }, { name = "uvicorn", extras = ["standard"], specifier = ">=0.30.1" }, - { name = "websockets", specifier = ">=12.0" }, + { name = "websockets", specifier = ">=15.0" }, { name = "windows-capture", marker = "platform_machine == 'AMD64' and sys_platform == 'win32'", specifier = ">=2.0.1,<3" }, ] provides-extras = ["dev"]