feat(llm): support incremental Responses WebSocket transport

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"

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

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