608 lines
22 KiB
Python
608 lines
22 KiB
Python
"""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
|
|
|
|
# Only shorten the final idle interval; scheduler latency must not break the first chain.
|
|
monkeypatch.setattr(websocket, "IDLE_SECONDS", 60)
|
|
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
|
|
monkeypatch.setattr(websocket, "IDLE_SECONDS", 0.01)
|
|
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
|