diff --git a/src/openjarvis/agents/errors.py b/src/openjarvis/agents/errors.py index d130d859..6ef1e9d0 100644 --- a/src/openjarvis/agents/errors.py +++ b/src/openjarvis/agents/errors.py @@ -2,6 +2,8 @@ from __future__ import annotations +from openjarvis.engine._base import looks_like_context_length_error + class AgentTickError(Exception): """Base class for agent tick errors.""" @@ -64,6 +66,14 @@ def classify_error(exc: Exception) -> AgentTickError: msg = str(exc).lower() + # A context-window overflow is deterministic: retrying the identical + # over-length request can never succeed, so fail fast instead of burning + # the retry budget on it. + if getattr(exc, "is_context_length_error", False) or ( + looks_like_context_length_error(msg) + ): + return FatalError(str(exc)) + # Check fatal patterns first (more specific) if isinstance(exc, PermissionError): return FatalError(str(exc)) @@ -90,6 +100,11 @@ def retry_delay(attempt: int) -> int: def suggest_action(error: AgentTickError) -> str: """Return a human-readable suggested action for the given error.""" msg = str(error).lower() + if looks_like_context_length_error(msg): + return ( + "Conversation too long for the model's context window \u2014 " + "start a new chat or shorten the conversation" + ) if any(p in msg for p in ("rate limit", "rate_limit", "429", "too many requests")): return "Rate limited \u2014 agent will auto-retry on next tick" if any(p in msg for p in ("timeout", "timed out", "connection", "unavailable")): diff --git a/src/openjarvis/agents/hybrid/mini_swe_agent.py b/src/openjarvis/agents/hybrid/mini_swe_agent.py index 1fe563ce..e4023e02 100644 --- a/src/openjarvis/agents/hybrid/mini_swe_agent.py +++ b/src/openjarvis/agents/hybrid/mini_swe_agent.py @@ -1403,12 +1403,10 @@ def _loop_local( # server walled the call. Compact aggressively (keep_last=1) # and retry once. Re-raise on anything else or on a second # failure — the runner records the row as errored. + from openjarvis.engine._base import looks_like_context_length_error + msg = str(exc) - is_ctx = ( - "maximum context length" in msg - or "context length" in msg.lower() - and "exceed" in msg.lower() - ) + is_ctx = looks_like_context_length_error(msg) if not is_ctx: raise _record_event( diff --git a/src/openjarvis/cli/ask.py b/src/openjarvis/cli/ask.py index 17b9ccbd..b10fc88a 100644 --- a/src/openjarvis/cli/ask.py +++ b/src/openjarvis/cli/ask.py @@ -20,6 +20,7 @@ from openjarvis.core.events import EventBus, EventType from openjarvis.core.types import Message, Role from openjarvis.engine import ( EngineConnectionError, + EngineContextLengthError, discover_engines, discover_models, get_engine, @@ -881,6 +882,11 @@ def ask( capability_policy=sec.capability_policy, memory_files_config=effective_mf, ) + except EngineContextLengthError as exc: + # Not a reachability problem — pointing the user at server/host + # config (hint_no_engine) would be misleading here. + console.print(f"[red]{exc}[/red]") + sys.exit(1) except EngineConnectionError as exc: console.print(f"[red]Engine error:[/red] {exc}") console.print(hint_no_engine()) @@ -990,6 +996,11 @@ def ask( temperature=temperature, max_tokens=max_tokens, ) + except EngineContextLengthError as exc: + # Not a reachability problem — pointing the user at server/host + # config (hint_no_engine) would be misleading here. + console.print(f"[red]{exc}[/red]") + sys.exit(1) except EngineConnectionError as exc: console.print(f"[red]Engine error:[/red] {exc}") console.print(hint_no_engine()) diff --git a/src/openjarvis/engine/__init__.py b/src/openjarvis/engine/__init__.py index 90486fe6..d575da8f 100644 --- a/src/openjarvis/engine/__init__.py +++ b/src/openjarvis/engine/__init__.py @@ -9,7 +9,9 @@ import openjarvis.engine.ollama # noqa: F401 import openjarvis.engine.openai_compat_engines # noqa: F401 from openjarvis.engine._base import ( EngineConnectionError, + EngineContextLengthError, InferenceEngine, + looks_like_context_length_error, messages_to_dicts, ) from openjarvis.engine._discovery import discover_engines, discover_models, get_engine @@ -23,9 +25,11 @@ for _optional in ("cloud", "litellm", "gemma_cpp"): __all__ = [ "EngineConnectionError", + "EngineContextLengthError", "InferenceEngine", "discover_engines", "discover_models", "get_engine", + "looks_like_context_length_error", "messages_to_dicts", ] diff --git a/src/openjarvis/engine/_base.py b/src/openjarvis/engine/_base.py index 6ecaf81a..102936ba 100644 --- a/src/openjarvis/engine/_base.py +++ b/src/openjarvis/engine/_base.py @@ -13,6 +13,46 @@ class EngineConnectionError(Exception): """Raised when an engine is unreachable.""" +class EngineContextLengthError(EngineConnectionError): + """The prompt exceeds the served model's maximum context window. + + Subclasses ``EngineConnectionError`` so existing ``except + EngineConnectionError`` handlers keep catching it, while callers that want a + distinct, user-facing "conversation too long" message can branch on this type + (or the ``is_context_length_error`` marker) instead of surfacing a generic + engine failure. + """ + + is_context_length_error: bool = True + + +# Substrings that identify an error body as a context-window overflow (vLLM, +# SGLang, and OpenAI-compatible servers phrase this a few different ways). +# Every marker is anchored on "context" on purpose: generic phrases like +# "please reduce" or "too many tokens" also appear in unrelated 400 bodies +# (max_tokens validation, rate limiting, oversized images) and would +# misclassify those as "conversation too long". +CONTEXT_LENGTH_MARKERS = ( + "context length", + "maximum context", + "context window", + "maximum_context", + "context_length_exceeded", +) + + +def looks_like_context_length_error(text: str) -> bool: + """True when *text* reads like a context-window overflow error. + + The single shared heuristic for recognizing vendor context-overflow + phrasings — used by the engine layer (typing upstream 400s), agent error + classification, and the server stream bridge, so a new vendor phrasing + only ever needs to be added here. + """ + low = (text or "").lower() + return any(marker in low for marker in CONTEXT_LENGTH_MARKERS) + + _REASONING_METADATA_KEYS = ("reasoning_content", "thinking") @@ -80,8 +120,11 @@ def estimate_prompt_tokens(messages: Sequence[Message]) -> int: __all__ = [ + "CONTEXT_LENGTH_MARKERS", "EngineConnectionError", + "EngineContextLengthError", "InferenceEngine", "estimate_prompt_tokens", + "looks_like_context_length_error", "messages_to_dicts", ] diff --git a/src/openjarvis/engine/_http_async.py b/src/openjarvis/engine/_http_async.py new file mode 100644 index 00000000..cc1239a2 --- /dev/null +++ b/src/openjarvis/engine/_http_async.py @@ -0,0 +1,128 @@ +"""Shared async-HTTP plumbing for engines that stream over httpx. + +Home of the pieces the OpenAI-compat and Ollama engines were each hand-rolling: +the async-client factory (with the configured timeout applied), a cached +long-lived client so consecutive streams reuse pooled connections instead of +paying a fresh TCP/TLS handshake per turn, the transport-error set that maps to +``EngineConnectionError``, and the non-2xx → engine-error translation. +""" + +from __future__ import annotations + +import asyncio +import logging +from typing import NoReturn + +import httpx + +from openjarvis.engine._base import ( + EngineConnectionError, + EngineContextLengthError, + looks_like_context_length_error, +) + +logger = logging.getLogger(__name__) + +# Transport failures that map to EngineConnectionError on the streaming paths. +# ``RemoteProtocolError``/``ReadError`` cover a server dying MID-STREAM (peer +# closed between tokens); a wedged read trips the configured timeout +# (TimeoutException). Kept exactly this narrow on purpose: +# ``asyncio.CancelledError``/``GeneratorExit`` are NOT ``httpx.TransportError`` +# subclasses and must keep propagating for correct cancellation. +STREAM_TRANSPORT_ERRORS = ( + httpx.ConnectError, + httpx.TimeoutException, + httpx.RemoteProtocolError, + httpx.ReadError, +) + +_CONTEXT_LENGTH_USER_MESSAGE = ( + "The conversation is too long for the model's context window. " + "Start a new chat or shorten the conversation, then try again." +) + + +class AsyncHTTPEngineMixin: + """Async streaming plumbing shared by httpx-backed engines. + + Expects the engine to provide ``engine_id``, ``_host``, ``_timeout``, an + ``_async_transport`` test seam (``httpx.MockTransport`` in tests, ``None`` + in production), and optionally ``_headers``. + """ + + engine_id: str + _host: str + _timeout: float + _async_transport: httpx.AsyncBaseTransport | None + + # Set True by engines whose upstream reports context-window overflows in + # 400 bodies (OpenAI-compat servers). Ollama has no such signal. + _stream_400_signals_context_length: bool = False + + # Lazily-created shared client (and the loop it belongs to). Class-level + # ``None`` defaults keep engine ``__init__``s free of mixin bookkeeping. + _async_client: httpx.AsyncClient | None = None + _async_client_loop: asyncio.AbstractEventLoop | None = None + + def _make_async_client(self) -> httpx.AsyncClient: + """Build an async client that honours the configured timeout.""" + return httpx.AsyncClient( + base_url=self._host, + timeout=self._timeout, + headers=getattr(self, "_headers", None), + transport=self._async_transport, + ) + + def _get_async_client(self) -> httpx.AsyncClient: + """Return the shared async client for the running event loop. + + Reusing one client across calls preserves connection pooling — without + it every conversation turn pays a fresh TCP (and TLS) handshake. The + client is cached per event loop: pooled connections die with their + loop, so CLI flows that run ``asyncio.run()`` per turn transparently + get a fresh client while a long-lived server loop keeps one pool. + """ + loop = asyncio.get_running_loop() + client = self._async_client + if client is None or client.is_closed or self._async_client_loop is not loop: + # Any previous client belonged to a finished loop; its pooled + # connections are already dead, so just drop the reference. + client = self._make_async_client() + self._async_client = client + self._async_client_loop = loop + return client + + def _close_async_client(self) -> None: + """Best-effort close of the shared async client (for ``close()``).""" + client = self._async_client + loop = self._async_client_loop + self._async_client = None + self._async_client_loop = None + if client is None or client.is_closed: + return + try: + if loop is not None and not loop.is_closed(): + if loop.is_running(): + loop.create_task(client.aclose()) + else: + loop.run_until_complete(client.aclose()) + except Exception: # noqa: BLE001 — cleanup must never mask the close + logger.debug("Async client did not close cleanly", exc_info=True) + + def _raise_stream_http_error(self, status: int, detail: str) -> NoReturn: + """Map a non-success streaming HTTP response to a clean engine error.""" + detail = (detail or "").strip() + if ( + status == 400 + and self._stream_400_signals_context_length + and looks_like_context_length_error(detail) + ): + raise EngineContextLengthError(_CONTEXT_LENGTH_USER_MESSAGE) + detail_suffix = f": {detail}" if detail else "" + raise EngineConnectionError( + f"{self.engine_id} engine at {self._host} returned HTTP " + f"{status}{detail_suffix}" + ) + + +__all__ = ["AsyncHTTPEngineMixin", "STREAM_TRANSPORT_ERRORS"] diff --git a/src/openjarvis/engine/_openai_compat.py b/src/openjarvis/engine/_openai_compat.py index 5f767478..507b6d83 100644 --- a/src/openjarvis/engine/_openai_compat.py +++ b/src/openjarvis/engine/_openai_compat.py @@ -12,18 +12,27 @@ import httpx from openjarvis.core.types import Message from openjarvis.engine._base import ( EngineConnectionError, + EngineContextLengthError, InferenceEngine, estimate_prompt_tokens, messages_to_dicts, ) +from openjarvis.engine._http_async import ( + STREAM_TRANSPORT_ERRORS, + AsyncHTTPEngineMixin, +) from openjarvis.engine._stubs import StreamChunk logger = logging.getLogger(__name__) -class _OpenAICompatibleEngine(InferenceEngine): +class _OpenAICompatibleEngine(AsyncHTTPEngineMixin, InferenceEngine): """Base for engines that serve the OpenAI ``/v1/chat/completions`` API.""" + # vLLM/SGLang report context-window overflows in 400 bodies; the shared + # ``_raise_stream_http_error`` types those as ``EngineContextLengthError``. + _stream_400_signals_context_length = True + engine_id: str = "" _default_host: str = "http://localhost:8000" _api_prefix: str = "/v1" @@ -50,6 +59,16 @@ class _OpenAICompatibleEngine(InferenceEngine): headers = ( {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None ) + # Used by the shared async streaming plumbing (AsyncHTTPEngineMixin) so + # the bounded request timeout is applied to streaming reads, not just + # the synchronous methods (a wedged token read fails at ``timeout`` + # rather than hanging the caller for the httpx default). + self._timeout = timeout + self._headers = headers + # Injection seam for tests: an ``httpx.MockTransport`` swapped in here lets + # the async stream path be exercised with a mocked transport and no real + # server. ``None`` in production so httpx uses its default networking. + self._async_transport: httpx.AsyncBaseTransport | None = None self._client = httpx.Client( base_url=self._host, timeout=timeout, headers=headers ) @@ -168,11 +187,26 @@ class _OpenAICompatibleEngine(InferenceEngine): # Default to tool_choice=auto when tools are provided if "tools" in payload and "tool_choice" not in payload: payload["tool_choice"] = "auto" + url = f"{self._api_prefix}/chat/completions" try: - url = f"{self._api_prefix}/chat/completions" - with self._client.stream("POST", url, json=payload) as resp: - resp.raise_for_status() - for line in resp.iter_lines(): + # ASYNC streaming: ``httpx.AsyncClient`` + ``aiter_lines`` never + # blocks the event loop between tokens (the previous SYNC + # ``httpx.Client`` + ``iter_lines`` inside this ``async def`` blocked + # the single uvicorn worker on every inter-token wait, serializing all + # concurrent chats and letting one wedged read freeze the whole API). + # The shared client keeps pooled connections across turns. + client = self._get_async_client() + async with client.stream("POST", url, json=payload) as resp: + # ``not is_success`` covers 3xx as well as 4xx/5xx. With + # ``follow_redirects`` off (the default) an unexpected redirect + # would otherwise fall through to ``aiter_lines`` and surface as + # a silent EMPTY stream instead of a clean engine error. + if not resp.is_success: + # Load the (short) error body before touching ``.text``: + # a streaming response is otherwise unread. + await resp.aread() + self._raise_stream_http_error(resp.status_code, resp.text) + async for line in resp.aiter_lines(): if not line.startswith("data:"): continue data_str = line[len("data:") :].strip() @@ -186,7 +220,11 @@ class _OpenAICompatibleEngine(InferenceEngine): content = delta.get("content") if content: yield content - except (httpx.ConnectError, httpx.TimeoutException) as exc: + except STREAM_TRANSPORT_ERRORS as exc: + # A wedged upstream read trips ``timeout`` (ReadTimeout) and is mapped + # here, so the request fails cleanly at the configured bound instead + # of hanging indefinitely (see STREAM_TRANSPORT_ERRORS for why the + # set is exactly this narrow). raise EngineConnectionError( f"{self.engine_id} engine not reachable at {self._host}" ) from exc @@ -212,11 +250,20 @@ class _OpenAICompatibleEngine(InferenceEngine): } if "tools" in payload and "tool_choice" not in payload: payload["tool_choice"] = "auto" + url = f"{self._api_prefix}/chat/completions" try: - url = f"{self._api_prefix}/chat/completions" - with self._client.stream("POST", url, json=payload) as resp: - resp.raise_for_status() - for line in resp.iter_lines(): + # ASYNC streaming (see ``stream``): non-blocking shared client so + # rich streaming never stalls the event loop and honours ``timeout``. + client = self._get_async_client() + async with client.stream("POST", url, json=payload) as resp: + # ``not is_success`` covers 3xx as well as 4xx/5xx. With + # ``follow_redirects`` off (the default) an unexpected redirect + # would otherwise fall through to ``aiter_lines`` and surface as + # a silent EMPTY stream instead of a clean engine error. + if not resp.is_success: + await resp.aread() + self._raise_stream_http_error(resp.status_code, resp.text) + async for line in resp.aiter_lines(): if not line.startswith("data:"): continue data_str = line[len("data:") :].strip() @@ -240,7 +287,10 @@ class _OpenAICompatibleEngine(InferenceEngine): finish_reason=finish, usage=usage, ) - except (httpx.ConnectError, httpx.TimeoutException) as exc: + except STREAM_TRANSPORT_ERRORS as exc: + # See ``stream``: transport failures (incl. a mid-stream server + # disconnect) map to a clean error; the set is kept narrow so + # cancellation still propagates. raise EngineConnectionError( f"{self.engine_id} engine not reachable at {self._host}" ) from exc @@ -279,6 +329,9 @@ class _OpenAICompatibleEngine(InferenceEngine): def close(self) -> None: self._client.close() + self._close_async_client() -__all__ = ["_OpenAICompatibleEngine"] +# ``EngineContextLengthError`` moved to ``openjarvis.engine._base``; re-exported +# here for callers/tests that import it from this module. +__all__ = ["_OpenAICompatibleEngine", "EngineContextLengthError"] diff --git a/src/openjarvis/engine/litellm.py b/src/openjarvis/engine/litellm.py index a90ec474..6e9c114b 100644 --- a/src/openjarvis/engine/litellm.py +++ b/src/openjarvis/engine/litellm.py @@ -123,8 +123,12 @@ class LiteLLMEngine(InferenceEngine): call_kwargs["api_base"] = self._api_base call_kwargs.update(kwargs) - resp = litellm.completion(**call_kwargs) - for chunk in resp: + # ``acompletion`` + ``async for``: the sync ``litellm.completion`` used + # before made a blocking network call (and blocking per-chunk reads) + # inside this ``async def``, stalling the whole event loop between + # tokens — the same bug the httpx engines' streaming paths had. + resp = await litellm.acompletion(**call_kwargs) + async for chunk in resp: delta = chunk.choices[0].delta if chunk.choices else None if delta and delta.content: yield delta.content diff --git a/src/openjarvis/engine/ollama.py b/src/openjarvis/engine/ollama.py index c959d44f..ae07e3cc 100644 --- a/src/openjarvis/engine/ollama.py +++ b/src/openjarvis/engine/ollama.py @@ -18,6 +18,10 @@ from openjarvis.engine._base import ( estimate_prompt_tokens, messages_to_dicts, ) +from openjarvis.engine._http_async import ( + STREAM_TRANSPORT_ERRORS, + AsyncHTTPEngineMixin, +) from openjarvis.engine._stubs import StreamChunk logger = logging.getLogger(__name__) @@ -80,11 +84,15 @@ def _default_num_ctx() -> int: @EngineRegistry.register("ollama") -class OllamaEngine(InferenceEngine): +class OllamaEngine(AsyncHTTPEngineMixin, InferenceEngine): """Ollama backend via its native HTTP API.""" engine_id = "ollama" + # Ollama has no context-length overflow signal in its 400 bodies, so the + # shared ``_raise_stream_http_error`` keeps its default (no + # ``EngineContextLengthError`` branch, unlike the OpenAI-compat engines). + _DEFAULT_HOST = "http://localhost:11434" def __init__( @@ -98,6 +106,14 @@ class OllamaEngine(InferenceEngine): env_host = os.environ.get("OLLAMA_HOST") host = env_host or self._DEFAULT_HOST self._host = host.rstrip("/") + # Used by the shared async streaming plumbing (AsyncHTTPEngineMixin) so a + # wedged token read is bounded by ``timeout`` instead of hanging the + # single event loop for the httpx default. + self._timeout = timeout + # Injection seam for tests: an ``httpx.MockTransport`` swapped in here drives + # the async stream path with no real Ollama server. ``None`` in production so + # httpx uses its default networking. + self._async_transport: httpx.AsyncBaseTransport | None = None self._client = httpx.Client(base_url=self._host, timeout=timeout) # Last stream usage — captured from Ollama's final chunk self._last_stream_usage: Dict[str, int] = {} @@ -263,9 +279,26 @@ class OllamaEngine(InferenceEngine): elif kwargs["think"] is not None: payload["think"] = kwargs["think"] try: - with self._client.stream("POST", "/api/chat", json=payload) as resp: - resp.raise_for_status() - for line in resp.iter_lines(): + # ASYNC streaming: ``httpx.AsyncClient`` + ``aiter_lines`` never + # blocks the event loop between tokens (the previous SYNC + # ``self._client`` + ``iter_lines`` inside this ``async def`` blocked + # the single uvicorn worker on every inter-token wait, serializing all + # concurrent chats and letting one wedged read freeze the whole API). + # The shared client keeps pooled connections across turns. + client = self._get_async_client() + async with client.stream("POST", "/api/chat", json=payload) as resp: + # ``not is_success`` covers 3xx as well as 4xx/5xx and maps + # to ``EngineConnectionError`` (matching the OpenAI-compat + # path) instead of leaking a raw ``httpx.HTTPStatusError``. + # With redirects off (the default) an unexpected 3xx would + # otherwise fall through to ``aiter_lines`` and surface as a + # silent EMPTY stream rather than a clean engine error. + if not resp.is_success: + # Read the (short) error body before touching ``.text``: + # a streaming response is otherwise unread. + await resp.aread() + self._raise_stream_http_error(resp.status_code, resp.text) + async for line in resp.aiter_lines(): if not line.strip(): continue try: @@ -290,7 +323,10 @@ class OllamaEngine(InferenceEngine): "total_tokens": full_prompt + comp, } break - except (httpx.ConnectError, httpx.TimeoutException) as exc: + except STREAM_TRANSPORT_ERRORS as exc: + # Transport failures (incl. a mid-stream server disconnect) map to a + # clean error; the set is kept narrow (see STREAM_TRANSPORT_ERRORS) + # so cancellation still propagates. raise EngineConnectionError( f"Ollama not reachable at {self._host}" ) from exc @@ -356,19 +392,34 @@ class OllamaEngine(InferenceEngine): ) -> AsyncIterator[StreamChunk]: """Execute the streaming request and yield parsed StreamChunks.""" try: - with self._client.stream("POST", "/api/chat", json=payload) as resp: + # ASYNC streaming (see ``stream``): shared ``AsyncClient`` + + # ``aiter_lines`` so rich streaming never stalls the event loop and + # honours ``timeout``. + client = self._get_async_client() + async with client.stream("POST", "/api/chat", json=payload) as resp: if resp.status_code == 400 and retry_without_tools: # Model doesn't support tools — retry without them. + # PRESERVED: this specific 400 path must still trigger the + # tools-less retry; only OTHER non-2xx responses map to + # EngineConnectionError below. payload.pop("tools", None) async for c in self._run_stream( payload, messages, retry_without_tools=False ): yield c return - resp.raise_for_status() + # ``not is_success`` covers 3xx as well as 4xx/5xx and maps + # to ``EngineConnectionError`` (matching the OpenAI-compat + # path) instead of leaking a raw ``httpx.HTTPStatusError``. + # With redirects off (the default) an unexpected 3xx would + # otherwise fall through to ``aiter_lines`` and surface as a + # silent EMPTY stream rather than a clean engine error. + if not resp.is_success: + await resp.aread() + self._raise_stream_http_error(resp.status_code, resp.text) finish_reason: str | None = None - for line in resp.iter_lines(): + async for line in resp.aiter_lines(): if not line.strip(): continue try: @@ -441,7 +492,10 @@ class OllamaEngine(InferenceEngine): usage=dict(self._last_stream_usage), ) break - except (httpx.ConnectError, httpx.TimeoutException) as exc: + except STREAM_TRANSPORT_ERRORS as exc: + # See ``stream``: transport failures (incl. a mid-stream server + # disconnect) map to a clean error; the set is kept narrow so + # cancellation still propagates. raise EngineConnectionError( f"Ollama not reachable at {self._host}" ) from exc @@ -474,6 +528,7 @@ class OllamaEngine(InferenceEngine): def close(self) -> None: self._client.close() + self._close_async_client() __all__ = ["OllamaEngine"] diff --git a/src/openjarvis/server/api_routes.py b/src/openjarvis/server/api_routes.py index cac2623c..b70b4f5e 100644 --- a/src/openjarvis/server/api_routes.py +++ b/src/openjarvis/server/api_routes.py @@ -741,8 +741,11 @@ async def websocket_chat_stream(websocket: WebSocket): ) except TypeError: # stream() didn't return an iterable; fall back to - # generate() - result = engine.generate(messages, model=model) + # generate(). It makes a blocking upstream call, so run + # it in a worker thread to keep the event loop free. + result = await asyncio.to_thread( + engine.generate, messages, model=model + ) content = ( result.get("content", "") if isinstance( @@ -767,8 +770,11 @@ async def websocket_chat_stream(websocket: WebSocket): ended_at=_time.time(), ) else: - # No stream method — single-shot generate - result = engine.generate(messages, model=model) + # No stream method — single-shot generate. Blocking upstream + # call, so run in a worker thread to keep the event loop free. + result = await asyncio.to_thread( + engine.generate, messages, model=model + ) content = ( result.get("content", "") if isinstance( diff --git a/src/openjarvis/server/routes.py b/src/openjarvis/server/routes.py index cda60f56..ec8f187e 100644 --- a/src/openjarvis/server/routes.py +++ b/src/openjarvis/server/routes.py @@ -231,8 +231,13 @@ async def chat_completions(request_body: ChatCompletionRequest, request: Request # tools (e.g. injecting MCP tools through this endpoint and wanting # the agent to execute them), add an explicit opt-in header rather # than removing this guard — silent re-routing is what produced #414. + # ``_handle_agent`` (sync ``agent.run()``) and ``_handle_direct`` (sync + # ``engine.generate()``) both make blocking upstream calls; run them in a + # worker thread so a slow/wedged non-streaming request can't stall the + # event loop and every other concurrent request with it. if agent is not None and not request_body.tools: - response = _handle_agent( + response = await asyncio.to_thread( + _handle_agent, agent, model, request_body, @@ -242,7 +247,8 @@ async def chat_completions(request_body: ChatCompletionRequest, request: Request ) else: bus = getattr(request.app.state, "bus", None) - response = _handle_direct( + response = await asyncio.to_thread( + _handle_direct, engine, model, request_body, diff --git a/src/openjarvis/server/stream_bridge.py b/src/openjarvis/server/stream_bridge.py index 386a677f..6d331021 100644 --- a/src/openjarvis/server/stream_bridge.py +++ b/src/openjarvis/server/stream_bridge.py @@ -16,6 +16,7 @@ from fastapi.responses import StreamingResponse from openjarvis.agents._stubs import AgentContext, BaseAgent from openjarvis.core.events import Event, EventBus, EventType +from openjarvis.engine._base import looks_like_context_length_error from openjarvis.server.models import ( ChatCompletionChunk, ChatCompletionRequest, @@ -194,8 +195,10 @@ class AgentStreamBridge: logger.error("Agent stream error: %s", exc, exc_info=True) error_str = str(exc) - if "context length" in error_str.lower() or ( - "400" in error_str and "too long" in error_str.lower() + if ( + getattr(exc, "is_context_length_error", False) + or looks_like_context_length_error(error_str) + or ("400" in error_str and "too long" in error_str.lower()) ): error_content = ( "The input is too long for the model's context window. " diff --git a/tests/agents/test_errors.py b/tests/agents/test_errors.py index 72b7ae46..94341912 100644 --- a/tests/agents/test_errors.py +++ b/tests/agents/test_errors.py @@ -56,3 +56,30 @@ class TestErrorClassification: assert retry_delay(2) == 40 # Capped at 300 seconds assert retry_delay(10) == 300 + + def test_classify_context_length_is_fatal(self): + # A context-window overflow is deterministic — retrying the identical + # over-length request can never succeed, so it must NOT be classified + # retryable (which would burn ~30s of backoff on guaranteed failures). + from openjarvis.agents.errors import classify_error + from openjarvis.engine._base import EngineContextLengthError + + typed = classify_error( + EngineContextLengthError( + "The conversation is too long for the model's context window." + ) + ) + assert typed.retryable is False + + # Same for untyped errors whose message reads like a context overflow + # (e.g. raw vendor errors from engines without the typed mapping). + untyped = classify_error( + Exception("This model's maximum context length is 4096 tokens.") + ) + assert untyped.retryable is False + + def test_suggest_action_context_length(self): + from openjarvis.agents.errors import FatalError, suggest_action + + action = suggest_action(FatalError("prompt exceeds the model's context window")) + assert "context window" in action or "too long" in action.lower() diff --git a/tests/engine/test_litellm.py b/tests/engine/test_litellm.py index b46b5b61..ed8793ed 100644 --- a/tests/engine/test_litellm.py +++ b/tests/engine/test_litellm.py @@ -158,6 +158,9 @@ class TestLiteLLMEngineGenerate: class TestLiteLLMEngineStream: def test_stream(self) -> None: + # stream() must use the ASYNC litellm entry point (acompletion): the + # sync litellm.completion makes blocking network reads inside an + # ``async def``, stalling the event loop between tokens. chunk1 = SimpleNamespace( choices=[SimpleNamespace(delta=SimpleNamespace(content="Hel"))] ) @@ -168,8 +171,12 @@ class TestLiteLLMEngineStream: choices=[SimpleNamespace(delta=SimpleNamespace(content=None))] ) + async def _chunks(): + for c in (chunk1, chunk2, chunk3): + yield c + fake_litellm = mock.MagicMock() - fake_litellm.completion.return_value = iter([chunk1, chunk2, chunk3]) + fake_litellm.acompletion = mock.AsyncMock(return_value=_chunks()) with mock.patch.dict("sys.modules", {"litellm": fake_litellm}): engine = LiteLLMEngine() @@ -186,6 +193,8 @@ class TestLiteLLMEngineStream: tokens = asyncio.run(collect()) assert tokens == ["Hel", "lo!"] + fake_litellm.acompletion.assert_awaited_once() + assert fake_litellm.completion.call_count == 0 class TestLiteLLMEngineListModels: diff --git a/tests/engine/test_ollama.py b/tests/engine/test_ollama.py index 0e62a343..30317481 100644 --- a/tests/engine/test_ollama.py +++ b/tests/engine/test_ollama.py @@ -6,13 +6,27 @@ import json import httpx import pytest -import respx + +try: + import respx + + _HAS_RESPX = True +except ImportError: # respx is an optional test-only dep; the async MockTransport + respx = None # type: ignore[assignment] # pins below run without it. + _HAS_RESPX = False from openjarvis.core.registry import EngineRegistry from openjarvis.core.types import Message, Role from openjarvis.engine._base import EngineConnectionError from openjarvis.engine.ollama import OllamaEngine, _is_control_token_only_args +# respx-backed tests exercise the SYNC client paths (generate/list_models/health) +# and the respx-driven stream tests; they skip cleanly when respx is absent. The +# async pins in TestOllamaStreamIsAsyncAndBounded use httpx.MockTransport directly. +requires_respx = pytest.mark.skipif( + not _HAS_RESPX, reason="respx not installed (optional test-only dependency)" +) + @pytest.fixture() def engine() -> OllamaEngine: @@ -20,6 +34,7 @@ def engine() -> OllamaEngine: return OllamaEngine(host="http://testhost:11434") +@requires_respx class TestOllamaGenerate: def test_generate_returns_content(self, engine: OllamaEngine) -> None: with respx.mock: @@ -53,6 +68,7 @@ class TestOllamaGenerate: ) +@requires_respx class TestOllamaListModels: def test_list_models(self, engine: OllamaEngine) -> None: with respx.mock: @@ -66,6 +82,7 @@ class TestOllamaListModels: assert models == ["qwen3:8b", "llama3.2:3b"] +@requires_respx class TestOllamaHealth: def test_health_true(self, engine: OllamaEngine) -> None: with respx.mock: @@ -120,6 +137,7 @@ class TestControlTokenFilter: assert _is_control_token_only_args(raw_args) is False +@requires_respx class TestOllamaGenerateControlToken: def test_generate_drops_control_token_tool_call(self, engine: OllamaEngine) -> None: with respx.mock: @@ -219,6 +237,7 @@ class TestOllamaGenerateControlToken: assert json.loads(result["tool_calls"][0]["arguments"]) == {"command": "date"} +@requires_respx class TestOllamaStreamFullControlToken: @pytest.mark.asyncio async def test_stream_full_drops_control_token_tool_call( @@ -257,6 +276,7 @@ class TestOllamaStreamFullControlToken: assert all(not c.tool_calls for c in chunks) +@requires_respx class TestOllamaStream: @pytest.mark.asyncio async def test_stream_yields_content(self, engine: OllamaEngine) -> None: @@ -275,3 +295,214 @@ class TestOllamaStream: ): tokens.append(tok) assert "Hello" in tokens + + +def _ndjson_transport(lines: list[str]) -> httpx.MockTransport: + body = "\n".join(lines) + "\n" + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, text=body) + + return httpx.MockTransport(handler) + + +class TestOllamaStreamIsAsyncAndBounded: + """Regression pins: the Ollama stream paths are async (never iterate the SYNC + httpx client between tokens, which blocked the single uvicorn worker on every + inter-token wait) and a mid-stream disconnect is bounded and mapped to a clean + error. Uses httpx.MockTransport directly, so it runs without respx.""" + + @pytest.mark.asyncio + async def test_stream_does_not_use_blocking_sync_client(self) -> None: + # PIN: the old code iterated ``self._client`` (a SYNC httpx.Client) via + # ``iter_lines`` inside this ``async def``. The async path must not touch the + # sync client at all — swap in a bomb that explodes if ``.stream`` is used. + engine = OllamaEngine(host="http://localhost:11434") + engine._async_transport = _ndjson_transport( + [ + json.dumps({"message": {"content": "Hi"}, "done": False}), + json.dumps({"message": {"content": " there"}, "done": True}), + ] + ) + + class _Boom: + def stream(self, *a, **k): # pragma: no cover - must never run + raise AssertionError("streaming used the blocking sync client") + + engine._client = _Boom() # type: ignore[assignment] + tokens = [ + tok + async for tok in engine.stream( + [Message(role=Role.USER, content="Hi")], model="qwen3:8b" + ) + ] + assert tokens == ["Hi", " there"] + + @pytest.mark.asyncio + async def test_stream_full_does_not_use_blocking_sync_client(self) -> None: + # stream_full delegates to _run_stream; prove that path is async too. + engine = OllamaEngine(host="http://localhost:11434") + engine._async_transport = _ndjson_transport( + [ + json.dumps({"message": {"content": "Hi"}, "done": False}), + json.dumps( + {"message": {"content": "", "tool_calls": []}, "done": True} + ), + ] + ) + + class _Boom: + def stream(self, *a, **k): # pragma: no cover - must never run + raise AssertionError("stream_full used the blocking sync client") + + engine._client = _Boom() # type: ignore[assignment] + chunks = [ + c + async for c in engine.stream_full( + [Message(role=Role.USER, content="Hi")], model="qwen3:8b" + ) + ] + assert any(c.content == "Hi" for c in chunks) + + @pytest.mark.asyncio + async def test_timeout_is_applied_to_async_stream_client(self) -> None: + # The configured timeout must be APPLIED to the async stream client, so a + # wedged read is actually bounded (not just stored on the engine). + engine = OllamaEngine(host="http://localhost:11434", timeout=0.05) + assert engine._timeout == 0.05 + client = engine._make_async_client() + try: + assert client.timeout == httpx.Timeout(0.05) + finally: + await client.aclose() + + @pytest.mark.asyncio + async def test_mid_stream_disconnect_maps_to_connection_error(self) -> None: + # PIN: a server dying MID-STREAM raises httpx.RemoteProtocolError from + # aiter_lines; it must surface as a clean EngineConnectionError, not raw. + class _MidStreamCrashStream(httpx.AsyncByteStream): + def __init__(self, request: httpx.Request) -> None: + self._request = request + + async def __aiter__(self): + yield b'{"message": {"content": "Hi"}, "done": false}\n' + raise httpx.RemoteProtocolError( + "peer closed connection mid-stream", request=self._request + ) + + async def aclose(self) -> None: # pragma: no cover - trivial + pass + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, stream=_MidStreamCrashStream(request)) + + engine = OllamaEngine(host="http://localhost:11434") + engine._async_transport = httpx.MockTransport(handler) + tokens: list[str] = [] + with pytest.raises(EngineConnectionError): + async for tok in engine.stream( + [Message(role=Role.USER, content="Hi")], model="qwen3:8b" + ): + tokens.append(tok) + # The disconnect happened AFTER the first token was delivered (mid-stream). + assert tokens == ["Hi"] + + +class TestOllamaStreamHttpErrorMapping: + """Regression pins: Ollama streaming non-2xx responses map to the same + ``EngineConnectionError`` as the OpenAI-compat engine (via + ``_raise_stream_http_error``), instead of leaking a raw + ``httpx.HTTPStatusError`` from ``raise_for_status()``. A 3xx must NOT fall + through to a silent empty stream. Uses ``httpx.MockTransport`` directly, so + it runs without respx.""" + + @staticmethod + def _status_transport( + status: int, *, text: str = "", headers: dict | None = None + ) -> httpx.MockTransport: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(status, text=text, headers=headers or {}) + + return httpx.MockTransport(handler) + + @pytest.mark.asyncio + async def test_stream_500_maps_to_connection_error(self) -> None: + # PIN: the old code called ``resp.raise_for_status()`` and leaked a raw + # httpx.HTTPStatusError on a streaming 500. It must now be a clean + # EngineConnectionError carrying the status + body, like the compat path. + engine = OllamaEngine(host="http://localhost:11434") + engine._async_transport = self._status_transport(500, text="internal boom") + tokens: list[str] = [] + with pytest.raises(EngineConnectionError) as excinfo: + async for tok in engine.stream( + [Message(role=Role.USER, content="Hi")], model="qwen3:8b" + ): + tokens.append(tok) + assert not isinstance(excinfo.value, httpx.HTTPStatusError) + assert "500" in str(excinfo.value) + assert "internal boom" in str(excinfo.value) + assert tokens == [] + + @pytest.mark.asyncio + async def test_stream_full_500_maps_to_connection_error(self) -> None: + # Same pin for the rich (_run_stream-backed) path. + engine = OllamaEngine(host="http://localhost:11434") + engine._async_transport = self._status_transport(500, text="internal boom") + with pytest.raises(EngineConnectionError) as excinfo: + async for _ in engine.stream_full( + [Message(role=Role.USER, content="Hi")], model="qwen3:8b" + ): + pass + assert not isinstance(excinfo.value, httpx.HTTPStatusError) + assert "500" in str(excinfo.value) + assert "internal boom" in str(excinfo.value) + + @pytest.mark.asyncio + async def test_stream_3xx_maps_to_connection_error_not_silent(self) -> None: + # PIN: with redirects off, a 3xx must map to EngineConnectionError, NOT + # fall through to ``aiter_lines`` as a silent empty stream. + engine = OllamaEngine(host="http://localhost:11434") + engine._async_transport = self._status_transport( + 302, headers={"location": "http://elsewhere/api/chat"} + ) + tokens: list[str] = [] + with pytest.raises(EngineConnectionError) as excinfo: + async for tok in engine.stream( + [Message(role=Role.USER, content="Hi")], model="qwen3:8b" + ): + tokens.append(tok) + assert "302" in str(excinfo.value) + assert tokens == [] + + @pytest.mark.asyncio + async def test_400_tools_retry_still_fires(self) -> None: + # REGRESSION GUARD: the tools-retry (400 WITH tools -> retry WITHOUT + # tools) must keep working; only OTHER non-2xx map to + # EngineConnectionError. A 400 carrying tools must NOT be treated as a + # generic connection error. + calls: list[bool] = [] # whether each request carried "tools" + + def handler(request: httpx.Request) -> httpx.Response: + payload = json.loads(request.content) + had_tools = "tools" in payload + calls.append(had_tools) + if had_tools: + return httpx.Response(400, text="model does not support tools") + body = ( + json.dumps({"message": {"content": "recovered"}, "done": True}) + "\n" + ) + return httpx.Response(200, text=body) + + engine = OllamaEngine(host="http://localhost:11434") + engine._async_transport = httpx.MockTransport(handler) + chunks = [ + c + async for c in engine.stream_full( + [Message(role=Role.USER, content="Hi")], + model="qwen3:8b", + tools=[{"type": "function", "function": {"name": "shell_exec"}}], + ) + ] + # First request had tools (400), second retried without them (200). + assert calls == [True, False] + assert any(c.content == "recovered" for c in chunks) diff --git a/tests/engine/test_openai_compat.py b/tests/engine/test_openai_compat.py index 31cc2af8..3df1a283 100644 --- a/tests/engine/test_openai_compat.py +++ b/tests/engine/test_openai_compat.py @@ -4,13 +4,37 @@ from __future__ import annotations import httpx import pytest -import respx + +try: + import respx + + _HAS_RESPX = True +except ImportError: # respx is an optional test-only dep; MockTransport tests + respx = None # type: ignore[assignment] # still run without it. + _HAS_RESPX = False from openjarvis.core.registry import EngineRegistry from openjarvis.core.types import Message, Role from openjarvis.engine._base import EngineConnectionError +from openjarvis.engine._openai_compat import EngineContextLengthError from openjarvis.engine.openai_compat_engines import VLLMEngine +# respx-backed tests exercise the SYNC client paths (generate/list_models/health) +# and skip cleanly when respx is absent; the async stream/timeout/disconnect tests +# below use httpx.MockTransport directly and never need respx. +requires_respx = pytest.mark.skipif( + not _HAS_RESPX, reason="respx not installed (optional test-only dependency)" +) + + +def _sse_transport(sse_lines: list[str]) -> httpx.MockTransport: + body = "\n".join(sse_lines) + "\n" + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, text=body) + + return httpx.MockTransport(handler) + @pytest.fixture() def engine() -> VLLMEngine: @@ -18,6 +42,7 @@ def engine() -> VLLMEngine: return VLLMEngine(host="http://testhost:8000") +@requires_respx class TestOpenAICompatGenerate: def test_generate_returns_content(self, engine: VLLMEngine) -> None: with respx.mock: @@ -83,6 +108,7 @@ class TestOpenAICompatGenerate: ) +@requires_respx class TestOpenAICompatListModels: def test_list_models(self, engine: VLLMEngine) -> None: with respx.mock: @@ -95,6 +121,7 @@ class TestOpenAICompatListModels: assert engine.list_models() == ["model-a", "model-b"] +@requires_respx class TestOpenAICompatHealth: def test_health_true(self, engine: VLLMEngine) -> None: with respx.mock: @@ -111,6 +138,7 @@ class TestOpenAICompatHealth: assert engine.health() is False +@requires_respx class TestOpenAICompatStream: @pytest.mark.asyncio async def test_stream_sse(self, engine: VLLMEngine) -> None: @@ -129,3 +157,199 @@ class TestOpenAICompatStream: ): tokens.append(tok) assert tokens == ["Hi", " there"] + + +class TestStreamIsAsyncAndBounded: + """Regression pins for BC2: the stream path is async (never blocks the event + loop on the SYNC httpx client) and a wedged/oversized upstream is bounded and + mapped to a clean error instead of hanging or surfacing raw HTTP.""" + + @pytest.mark.asyncio + async def test_timeout_is_applied_to_async_stream_client(self) -> None: + # HARDENED: assert the configured timeout is actually APPLIED to the async + # stream client, not merely stored on the engine. Deleting + # ``timeout=self._timeout`` from ``_make_async_client`` drops the client to + # httpx's 5s default and fails this (a stored-only assertion would not). + engine = VLLMEngine(host="http://testhost:8000", timeout=180.0) + assert engine._timeout == 180.0 + client = engine._make_async_client() + try: + assert client.timeout == httpx.Timeout(180.0) + finally: + await client.aclose() + + @pytest.mark.asyncio + async def test_stream_does_not_use_blocking_sync_client(self) -> None: + # PIN: the old code iterated ``self._client`` (a SYNC httpx.Client) inside + # this ``async def``, blocking the single uvicorn worker between tokens. + # The async path must not touch the sync client at all. + engine = VLLMEngine(host="http://localhost:8000") + engine._async_transport = _sse_transport( + [ + 'data: {"choices":[{"delta":{"content":"Hi"}}]}', + 'data: {"choices":[{"delta":{"content":" there"}}]}', + "data: [DONE]", + ] + ) + + class _Boom: + def stream(self, *a, **k): # pragma: no cover - must never run + raise AssertionError("streaming used the blocking sync client") + + engine._client = _Boom() # type: ignore[assignment] + tokens = [ + tok + async for tok in engine.stream( + [Message(role=Role.USER, content="Hello")], model="m" + ) + ] + assert tokens == ["Hi", " there"] + + @pytest.mark.asyncio + async def test_wedged_read_is_bounded_by_timeout(self) -> None: + # A wedged upstream read trips the (small, honoured) timeout and maps to a + # clean EngineConnectionError rather than hanging the caller. + seen: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + # httpx populates request.extensions["timeout"] with the per-op timeouts + # that were actually applied to THIS request; capturing it here proves + # the configured 0.05s reached the wire, not just the engine attribute. + seen["timeout"] = request.extensions["timeout"] + raise httpx.ReadTimeout("wedged upstream", request=request) + + engine = VLLMEngine(host="http://localhost:8000", timeout=0.05) + engine._async_transport = httpx.MockTransport(handler) + with pytest.raises(EngineConnectionError): + async for _ in engine.stream( + [Message(role=Role.USER, content="Hello")], model="m" + ): + pass + # HARDENED: the configured timeout was APPLIED to the request. Deleting + # ``timeout=self._timeout`` from ``_make_async_client`` drops this to httpx's + # 5.0s default and fails the assertion. + assert seen["timeout"]["read"] == 0.05 + + @pytest.mark.asyncio + async def test_context_length_400_maps_to_context_error(self) -> None: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 400, + text=( + "This model's maximum context length is 4096 tokens. " + "However, you requested 5200 tokens. Please reduce the length." + ), + ) + + engine = VLLMEngine(host="http://localhost:8000") + engine._async_transport = httpx.MockTransport(handler) + with pytest.raises(EngineContextLengthError) as excinfo: + async for _ in engine.stream( + [Message(role=Role.USER, content="Hello")], model="m" + ): + pass + assert getattr(excinfo.value, "is_context_length_error", False) is True + + @pytest.mark.asyncio + async def test_other_upstream_error_maps_to_connection_error(self) -> None: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(500, text="internal server error") + + engine = VLLMEngine(host="http://localhost:8000") + engine._async_transport = httpx.MockTransport(handler) + with pytest.raises(EngineConnectionError) as excinfo: + async for _ in engine.stream( + [Message(role=Role.USER, content="Hello")], model="m" + ): + pass + # A generic upstream failure is NOT reported as a context-length problem. + assert not isinstance(excinfo.value, EngineContextLengthError) + + @pytest.mark.asyncio + async def test_mid_stream_disconnect_maps_to_connection_error(self) -> None: + # PIN: a server dying MID-STREAM raises httpx.RemoteProtocolError from + # aiter_lines. Before the except tuple was widened it propagated raw; it + # must now surface as a clean EngineConnectionError. + class _MidStreamCrashStream(httpx.AsyncByteStream): + def __init__(self, request: httpx.Request) -> None: + self._request = request + + async def __aiter__(self): + yield b'data: {"choices":[{"delta":{"content":"Hi"}}]}\n\n' + raise httpx.RemoteProtocolError( + "peer closed connection mid-stream", request=self._request + ) + + async def aclose(self) -> None: # pragma: no cover - trivial + pass + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, stream=_MidStreamCrashStream(request)) + + engine = VLLMEngine(host="http://localhost:8000") + engine._async_transport = httpx.MockTransport(handler) + tokens: list[str] = [] + with pytest.raises(EngineConnectionError) as excinfo: + async for tok in engine.stream( + [Message(role=Role.USER, content="Hello")], model="m" + ): + tokens.append(tok) + # The disconnect happened AFTER the first token was delivered (mid-stream), + # and did not masquerade as a context-length error. + assert tokens == ["Hi"] + assert not isinstance(excinfo.value, EngineContextLengthError) + + @pytest.mark.asyncio + async def test_unrelated_400_is_not_context_error(self) -> None: + # PIN: the context-length markers are anchored on "context" so generic + # 400 bodies containing phrases like "please reduce" (max_tokens + # validation, rate limiting, oversized images) are NOT misclassified + # as "conversation too long" — that message would send the user off to + # shorten a conversation that isn't the problem. + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 400, + text=( + "Invalid max_tokens: the maximum number of tokens you can " + "request is 4096; please reduce max_tokens and retry." + ), + ) + + engine = VLLMEngine(host="http://localhost:8000") + engine._async_transport = httpx.MockTransport(handler) + with pytest.raises(EngineConnectionError) as excinfo: + async for _ in engine.stream( + [Message(role=Role.USER, content="Hello")], model="m" + ): + pass + assert not isinstance(excinfo.value, EngineContextLengthError) + + @pytest.mark.asyncio + async def test_async_client_is_reused_across_streams(self) -> None: + # PIN: consecutive streams on the same event loop share one AsyncClient + # (connection pooling). A per-call client would pay a fresh TCP/TLS + # handshake on every conversation turn. + engine = VLLMEngine(host="http://localhost:8000") + engine._async_transport = _sse_transport( + [ + 'data: {"choices":[{"delta":{"content":"Hi"}}]}', + "data: [DONE]", + ] + ) + + async def one_turn() -> list[str]: + return [ + tok + async for tok in engine.stream( + [Message(role=Role.USER, content="Hello")], model="m" + ) + ] + + assert await one_turn() == ["Hi"] + first_client = engine._async_client + assert first_client is not None and not first_client.is_closed + assert await one_turn() == ["Hi"] + assert engine._async_client is first_client + # close() tears the shared client down with the sync one. + engine.close() + assert engine._async_client is None diff --git a/tests/engine/test_stream_full.py b/tests/engine/test_stream_full.py index 4003738c..94a8979a 100644 --- a/tests/engine/test_stream_full.py +++ b/tests/engine/test_stream_full.py @@ -5,12 +5,27 @@ from __future__ import annotations import json from collections.abc import AsyncIterator from typing import Any, Dict, List -from unittest.mock import MagicMock +import httpx import pytest from openjarvis.core.types import Message, Role from openjarvis.engine._stubs import InferenceEngine, StreamChunk +from openjarvis.engine.openai_compat_engines import VLLMEngine + + +def _sse_transport(sse_lines: list[str]) -> httpx.MockTransport: + """A MockTransport that replies to /v1/chat/completions with SSE ``sse_lines``. + + Exercises the REAL async httpx streaming path (aiter_lines) with no server. + """ + body = "\n".join(sse_lines) + "\n" + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, text=body) + + return httpx.MockTransport(handler) + # --------------------------------------------------------------------------- # StreamChunk dataclass tests @@ -147,8 +162,6 @@ class TestOpenAICompatStreamFull: @pytest.mark.asyncio async def test_parses_sse_with_content_and_finish(self): - from openjarvis.engine._openai_compat import _OpenAICompatibleEngine - # Build mock SSE lines sse_lines = [] for token in ["Hello", " world"]: @@ -164,22 +177,8 @@ class TestOpenAICompatStreamFull: sse_lines.append(f"data: {json.dumps(final)}") sse_lines.append("data: [DONE]") - # Mock the httpx client stream context manager - mock_resp = MagicMock() - mock_resp.raise_for_status = MagicMock() - mock_resp.iter_lines.return_value = iter(sse_lines) - - engine = _OpenAICompatibleEngine.__new__(_OpenAICompatibleEngine) - engine.engine_id = "test" - engine._host = "http://localhost:8000" - engine._api_prefix = "/v1" - - mock_client = MagicMock() - mock_stream_ctx = MagicMock() - mock_stream_ctx.__enter__ = MagicMock(return_value=mock_resp) - mock_stream_ctx.__exit__ = MagicMock(return_value=False) - mock_client.stream.return_value = mock_stream_ctx - engine._client = mock_client + engine = VLLMEngine(host="http://localhost:8000") + engine._async_transport = _sse_transport(sse_lines) chunks = [] async for chunk in engine.stream_full( @@ -198,8 +197,6 @@ class TestOpenAICompatStreamFull: @pytest.mark.asyncio async def test_parses_tool_call_fragments(self): - from openjarvis.engine._openai_compat import _OpenAICompatibleEngine - # Simulate streamed tool_call fragments _tc1 = ( '{"choices": [{"delta": {"tool_calls": [{"index": 0, "id": "call_1",' @@ -218,21 +215,8 @@ class TestOpenAICompatStreamFull: "data: [DONE]", ] - mock_resp = MagicMock() - mock_resp.raise_for_status = MagicMock() - mock_resp.iter_lines.return_value = iter(sse_lines) - - engine = _OpenAICompatibleEngine.__new__(_OpenAICompatibleEngine) - engine.engine_id = "test" - engine._host = "http://localhost:8000" - engine._api_prefix = "/v1" - - mock_client = MagicMock() - mock_stream_ctx = MagicMock() - mock_stream_ctx.__enter__ = MagicMock(return_value=mock_resp) - mock_stream_ctx.__exit__ = MagicMock(return_value=False) - mock_client.stream.return_value = mock_stream_ctx - engine._client = mock_client + engine = VLLMEngine(host="http://localhost:8000") + engine._async_transport = _sse_transport(sse_lines) chunks = [] async for chunk in engine.stream_full(