embed-client 4.0.22__tar.gz → 4.0.23__tar.gz

{embed_client-4.0.22/embed_client.egg-info → embed_client-4.0.23}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: embed-client
-Version: 4.0.22
+Version: 4.0.23
 Summary: Async client for the Embedding Service API. Canonical class: EmbeddingClient (JSON-RPC via mcp-proxy-adapter >=8.10).
 Author-email: Vasiliy Zdanovskiy <vasilyvz@gmail.com>
 License-Expression: MIT

{embed_client-4.0.22 → embed_client-4.0.23}/embed_client/__init__.py

@@ -27,7 +27,7 @@ try:
     __version__ = version("embed-client")
 except PackageNotFoundError:
     # Editable checkout without install metadata (e.g. some unit-test sandboxes).
-    __version__ = "4.0.22"
+    __version__ = "4.0.23"
 
 __all__ = [
     "EmbeddingClient",

{embed_client-4.0.22 → embed_client-4.0.23}/embed_client/embedding_client.py

@@ -31,6 +31,31 @@ from embed_client.exceptions import EmbedAPIError, EmbedClientError, EmbedError
 
 logger = logging.getLogger(__name__)
 
+try:  # mcp-proxy-adapter >= 8.10 raises this when /ws closes mid-wait
+    from mcp_proxy_adapter.client.jsonrpc_client.exceptions import (
+        ClientWebSocketClosedError as _AdapterWsClosedError,
+    )
+except ImportError:  # pragma: no cover - older adapter builds
+    _AdapterWsClosedError = None  # type: ignore[assignment,misc]
+
+# Terminal queue states reported by queue_get_job_status / embed_job_status.
+_POLL_COMPLETED_STATES = ("completed", "success")
+_POLL_FAILED_STATES = ("failed", "error")
+_POLL_STOPPED_STATES = ("stopped", "cancelled", "canceled")
+
+
+def _is_ws_closed_before_terminal(exc: BaseException) -> bool:
+    """Whether *exc* means the /ws stream closed before a terminal job event.
+
+    The job itself is unaffected by a dropped stream: the server keeps
+    executing it and ``queue_get_job_status`` still reports its terminal
+    state. A heartbeat miss while the server is busy closes the stream
+    without any job having failed.
+    """
+    if _AdapterWsClosedError is not None and isinstance(exc, _AdapterWsClosedError):
+        return True
+    return "websocket closed" in str(exc).lower()
+
 
 def _client_has_async_method(client: Any, name: str) -> bool:
     fn = getattr(client, name, None)
@@ -757,6 +782,7 @@ class EmbeddingClient:
         """
         if _client_has_async_method(self.client, "wait_for_job"):
             t_arg = float(timeout) if timeout > 0 else 0.0
+            started = asyncio.get_running_loop().time()
             try:
                 raw = await self.client.wait_for_job(job_id, timeout=t_arg)
             except asyncio.TimeoutError as e:
@@ -765,13 +791,78 @@ class EmbeddingClient:
                 else:
                     msg = f"Job {job_id} WebSocket wait timed out"
                 raise TimeoutError(msg) from e
-            except RuntimeError as e:
-                raise EmbedClientError(f"WebSocket error: {e}") from e
+            except Exception as e:  # noqa: BLE001 - classified below
+                if not _is_ws_closed_before_terminal(e):
+                    if isinstance(e, RuntimeError):
+                        raise EmbedClientError(f"WebSocket error: {e}") from e
+                    raise
+                return await self._poll_terminal_after_ws_close(
+                    job_id, t_arg, started, e, poll_interval=poll_interval
+                )
             return _normalize_terminal_job_payload(raw)
         return await wait_for_job_poll(
             self.job_status, job_id, timeout=timeout, poll_interval=poll_interval
         )
 
+    async def _poll_terminal_after_ws_close(
+        self,
+        job_id: str,
+        timeout: float,
+        started: float,
+        cause: BaseException,
+        poll_interval: float = 1.0,
+    ) -> Dict[str, Any]:
+        """Finish a wait whose /ws stream closed before the terminal event.
+
+        Polls :meth:`job_status` (``queue_get_job_status``) for the rest of the
+        caller's *timeout* and returns exactly what the WebSocket path would
+        have returned: the terminal status is wrapped as the equivalent
+        ``job_completed`` / ``job_failed`` event and passed through the same
+        normalizer, so a failed command inside a completed envelope is still
+        reported as a failure. ``timeout`` 0 means no deadline, as on the WS path.
+        """
+        loop = asyncio.get_running_loop()
+        deadline = None if timeout <= 0 else started + timeout
+        logger.warning(
+            "WebSocket closed before a terminal event for job %s (%s); "
+            "polling queue_get_job_status for its terminal state",
+            job_id,
+            cause,
+        )
+        interval = poll_interval if poll_interval and poll_interval > 0 else 1.0
+        while True:
+            status = await self.job_status(job_id)
+            if isinstance(status, dict):
+                if status.get("exists") is False:
+                    raise EmbedClientError(f"Job {job_id} not found") from cause
+                state = str(status.get("status") or "").lower()
+                if state in _POLL_COMPLETED_STATES:
+                    return _normalize_terminal_job_payload(
+                        {
+                            "event": "job_completed",
+                            "job_id": job_id,
+                            "result": status.get("result") or {},
+                        }
+                    )
+                if state in _POLL_FAILED_STATES:
+                    return _normalize_terminal_job_payload(
+                        {
+                            "event": "job_failed",
+                            "job_id": job_id,
+                            "error": status.get("error")
+                            or status.get("description")
+                            or status.get("message")
+                            or "Unknown error",
+                        }
+                    )
+                if state in _POLL_STOPPED_STATES:
+                    raise EmbedClientError(f"Job {job_id} was stopped") from cause
+            if deadline is not None and loop.time() >= deadline:
+                raise TimeoutError(
+                    f"Job {job_id} did not complete within {timeout} seconds"
+                ) from cause
+            await asyncio.sleep(interval)
+
     async def _embed_via_execute_async(
         self, params: Dict[str, Any], timeout: int = 60
     ) -> Dict[str, Any]:
@@ -833,6 +924,7 @@ class EmbeddingClient:
         """
         if _client_has_async_method(self.client, "wait_for_job"):
             t_arg = float(timeout) if timeout > 0 else 0.0
+            started = asyncio.get_running_loop().time()
 
             async def _on_event(msg: Dict[str, Any]) -> None:
                 if on_ws_event is not None:
@@ -850,8 +942,14 @@ class EmbeddingClient:
                 else:
                     msg = f"Job {job_id} WebSocket wait timed out"
                 raise TimeoutError(msg) from e
-            except RuntimeError as e:
-                raise EmbedClientError(f"WebSocket error: {e}") from e
+            except Exception as e:  # noqa: BLE001 - classified below
+                if not _is_ws_closed_before_terminal(e):
+                    if isinstance(e, RuntimeError):
+                        raise EmbedClientError(f"WebSocket error: {e}") from e
+                    raise
+                return await self._poll_terminal_after_ws_close(
+                    job_id, t_arg, started, e
+                )
             return _normalize_terminal_job_payload(payload)
 
         timeout_sec = float(timeout) if timeout > 0 else 60.0

{embed_client-4.0.22 → embed_client-4.0.23/embed_client.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: embed-client
-Version: 4.0.22
+Version: 4.0.23
 Summary: Async client for the Embedding Service API. Canonical class: EmbeddingClient (JSON-RPC via mcp-proxy-adapter >=8.10).
 Author-email: Vasiliy Zdanovskiy <vasilyvz@gmail.com>
 License-Expression: MIT

{embed_client-4.0.22 → embed_client-4.0.23}/embed_client.egg-info/SOURCES.txt

@@ -227,4 +227,5 @@ tests/test_local_json_embed_file.py
 tests/test_md_embed_file_prep.py
 tests/test_security_matrix.py
 tests/test_ssl_manager.py
+tests/test_wait_for_job_ws_fallback.py
 tests/test_with_auth.py

{embed_client-4.0.22 → embed_client-4.0.23}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "embed-client"
-version = "4.0.22"
+version = "4.0.23"
 description = """Async client for the Embedding Service API. Canonical class: EmbeddingClient (JSON-RPC via mcp-proxy-adapter >=8.10).
 
 Public API highlights:

embed_client-4.0.23/tests/test_wait_for_job_ws_fallback.py

@@ -0,0 +1,113 @@
+"""
+wait_for_job: a /ws stream that closes before the terminal event falls back to
+queue_get_job_status polling and returns the same normalized result.
+
+Live evidence (embed 4.0.22, 2026-09-24): jobs completed on the server while the
+client's WebSocket wait raised "WebSocket closed before a job terminal event"
+because the busy server missed the 30 s heartbeat.
+
+Author: Vasiliy Zdanovskiy
+email: vasilyvz@gmail.com
+"""
+
+import asyncio
+from typing import Any, Dict, List
+
+import pytest
+from mcp_proxy_adapter.client.jsonrpc_client.exceptions import (
+    ClientWebSocketClosedError,
+)
+
+from embed_client.embedding_client import EmbeddingClient
+from embed_client.exceptions import EmbedClientError
+
+_ROW = {"body": "x", "embedding": [0.1, 0.2], "tokens": ["x"], "bm25_tokens": ["x"], "error": None}
+
+
+def _completed_status(success: bool = True) -> Dict[str, Any]:
+    inner: Dict[str, Any]
+    if success:
+        inner = {"success": True, "data": {"results": [_ROW], "model": "m", "dimension": 2}}
+    else:
+        inner = {"success": False, "error": {"message": "Worker unreachable", "code": -32603}}
+    return {
+        "job_id": "j1",
+        "status": "completed",
+        "result": {"job_id": "j1", "command": "embed_thin", "result": inner},
+    }
+
+
+class _FakeTransport:
+    """Adapter client double: WS wait closes, status polls follow a script."""
+
+    def __init__(self, statuses: List[Dict[str, Any]], ws_error: Exception) -> None:
+        self._statuses = list(statuses)
+        self._ws_error = ws_error
+        self.status_calls = 0
+
+    async def wait_for_job(self, job_id: str, timeout: float = 0.0, on_event: Any = None) -> Dict[str, Any]:
+        raise self._ws_error
+
+    async def jsonrpc_call(self, method: str, params: Dict[str, Any]) -> Dict[str, Any]:
+        assert method == "queue_get_job_status"
+        self.status_calls += 1
+        status = self._statuses.pop(0) if len(self._statuses) > 1 else self._statuses[0]
+        return {"jsonrpc": "2.0", "id": 1, "result": {"success": True, "data": status}}
+
+
+def _client(transport: _FakeTransport) -> EmbeddingClient:
+    client = EmbeddingClient.__new__(EmbeddingClient)
+    client.client = transport  # type: ignore[attr-defined]
+    return client
+
+
+def _ws_closed() -> ClientWebSocketClosedError:
+    return ClientWebSocketClosedError(
+        "WebSocket closed before a job terminal event",
+        operation="websocket.wait_for_job",
+        url="wss://example/ws",
+    )
+
+
+def test_ws_close_falls_back_to_polling_and_returns_normalized_result() -> None:
+    transport = _FakeTransport(
+        [{"job_id": "j1", "status": "running"}, _completed_status()], _ws_closed()
+    )
+    result = asyncio.run(_client(transport).wait_for_job("j1", timeout=10, poll_interval=0.01))
+    assert result["results"] == [_ROW]
+    assert result["model"] == "m"
+    assert result["dimension"] == 2
+    assert transport.status_calls == 2
+
+
+def test_ws_close_then_failed_command_inside_completed_envelope_raises() -> None:
+    transport = _FakeTransport([_completed_status(success=False)], _ws_closed())
+    with pytest.raises(EmbedClientError, match="Worker unreachable"):
+        asyncio.run(_client(transport).wait_for_job("j1", timeout=10, poll_interval=0.01))
+
+
+def test_ws_close_then_failed_job_raises_job_failed() -> None:
+    transport = _FakeTransport(
+        [{"job_id": "j1", "status": "failed", "error": "boom"}], _ws_closed()
+    )
+    with pytest.raises(EmbedClientError, match="Job failed: boom"):
+        asyncio.run(_client(transport).wait_for_job("j1", timeout=10, poll_interval=0.01))
+
+
+def test_ws_close_poll_respects_the_callers_timeout() -> None:
+    transport = _FakeTransport([{"job_id": "j1", "status": "running"}], _ws_closed())
+    with pytest.raises(TimeoutError, match="did not complete within"):
+        asyncio.run(_client(transport).wait_for_job("j1", timeout=0.2, poll_interval=0.05))
+
+
+def test_other_ws_errors_are_not_masked_by_the_fallback() -> None:
+    transport = _FakeTransport([_completed_status()], RuntimeError("TLS handshake failed"))
+    with pytest.raises(EmbedClientError, match="WebSocket error: TLS handshake failed"):
+        asyncio.run(_client(transport).wait_for_job("j1", timeout=10, poll_interval=0.01))
+    assert transport.status_calls == 0
+
+
+def test_adapter_ws_path_used_by_embed_file_also_falls_back() -> None:
+    transport = _FakeTransport([_completed_status()], _ws_closed())
+    result = asyncio.run(_client(transport)._wait_for_job_via_adapter_ws("j1", timeout=10))
+    assert result["results"] == [_ROW]