feat(workflows): raise StreamDisconnectedError on SSE error frames

Adds an AfterSuccess hook that converts a workflow SSE `event: error` frame into a raised StreamDisconnectedError (reason + error), so consumers use try/except around stream iteration instead of inspecting each event.

Author: mistralai-nfau

src/mistralai/client/_hooks/registration.py

@@ -4,6 +4,7 @@
 from .tracing import TracingHook
 from .types import Hooks
 from .workflow_encoding_hook import WorkflowEncodingHook
+from mistralai.extra.workflows.stream_error_hook import WorkflowStreamErrorHook
 
 # This file is only ever generated once on the first generation and then is free to be modified.
 # Any hooks you wish to add should be registered in the init_hooks function. Feel free to define them
@@ -26,3 +27,4 @@ def init_hooks(hooks: Hooks):
     hooks.register_after_error_hook(tracing_hook)
     hooks.register_before_request_hook(workflow_encoding_hook)
     hooks.register_after_success_hook(workflow_encoding_hook)
+    hooks.register_after_success_hook(WorkflowStreamErrorHook())

src/mistralai/extra/tests/test_stream_error_hook.py

@@ -0,0 +1,154 @@
+import httpx
+import pytest
+from httpx._types import AsyncByteStream, SyncByteStream
+
+from mistralai.client import Mistral
+from mistralai.client._hooks.types import AfterSuccessContext, HookContext
+from mistralai.extra.workflows.errors import StreamDisconnectedError
+from mistralai.extra.workflows.stream_error_hook import WorkflowStreamErrorHook
+
+STREAM_OPERATION_ID = "get_stream_events_v1_workflows_events_stream_get"
+NON_STREAM_OPERATION_ID = "chat_completion_v1_chat_completions_post"
+
+GOOD_FRAME = b'event: workflow.event\ndata: {"attributes": {}}\n\n'
+ERROR_FRAME = b'event: error\ndata: {"error": "boom", "reason": "read_error"}\n\n'
+
+
+class _SyncSource(SyncByteStream):
+    def __init__(self, chunks: list[bytes]) -> None:
+        self._chunks = chunks
+
+    def __iter__(self):
+        yield from self._chunks
+
+    def close(self) -> None:
+        pass
+
+
+class _AsyncSource(AsyncByteStream):
+    def __init__(self, chunks: list[bytes]) -> None:
+        self._iter = iter(chunks)
+
+    def __aiter__(self) -> "_AsyncSource":
+        return self
+
+    async def __anext__(self) -> bytes:
+        try:
+            return next(self._iter)
+        except StopIteration:
+            raise StopAsyncIteration
+
+    async def aclose(self) -> None:
+        pass
+
+
+def _hook_ctx(operation_id: str) -> AfterSuccessContext:
+    client = Mistral(api_key="test-key")
+    return AfterSuccessContext(
+        HookContext(
+            config=client.sdk_configuration,
+            base_url="https://api.example.com",
+            operation_id=operation_id,
+            oauth2_scopes=[],
+            security_source=None,
+        )
+    )
+
+
+def _sse_response(source, *, content_type: str = "text/event-stream") -> httpx.Response:
+    return httpx.Response(
+        status_code=200,
+        headers={"content-type": content_type},
+        stream=source,
+        request=httpx.Request(
+            "GET", "https://api.example.com/v1/workflows/events/stream"
+        ),
+    )
+
+
+def test_hook_raises_stream_disconnected_error_and_passes_prior_events_through():
+    response = _sse_response(_SyncSource([GOOD_FRAME, ERROR_FRAME]))
+    result = WorkflowStreamErrorHook().after_success(
+        _hook_ctx(STREAM_OPERATION_ID), response
+    )
+    assert isinstance(result, httpx.Response)
+
+    collected: list[bytes] = []
+    with pytest.raises(StreamDisconnectedError) as exc_info:
+        for chunk in result.iter_bytes():
+            collected.append(chunk)
+
+    assert exc_info.value.reason == "read_error"
+    assert exc_info.value.error == "boom"
+    assert b"workflow.event" in b"".join(collected)
+
+
+@pytest.mark.asyncio
+async def test_hook_raises_stream_disconnected_error_on_async_stream():
+    response = _sse_response(_AsyncSource([GOOD_FRAME, ERROR_FRAME]))
+    result = WorkflowStreamErrorHook().after_success(
+        _hook_ctx(STREAM_OPERATION_ID), response
+    )
+    assert isinstance(result, httpx.Response)
+
+    collected: list[bytes] = []
+    with pytest.raises(StreamDisconnectedError) as exc_info:
+        async for chunk in result.aiter_bytes():
+            collected.append(chunk)
+
+    assert exc_info.value.reason == "read_error"
+    assert exc_info.value.error == "boom"
+    assert b"workflow.event" in b"".join(collected)
+
+
+def test_hook_detects_error_frame_split_across_chunks():
+    chunks = [
+        b"event: er",
+        b'ror\ndata: {"error": "x", "reason": "internal_error"}\n\n',
+    ]
+    response = _sse_response(_SyncSource(chunks))
+    result = WorkflowStreamErrorHook().after_success(
+        _hook_ctx(STREAM_OPERATION_ID), response
+    )
+
+    with pytest.raises(StreamDisconnectedError) as exc_info:
+        list(result.iter_bytes())
+
+    assert exc_info.value.reason == "internal_error"
+    assert exc_info.value.error == "x"
+
+
+def test_hook_defaults_reason_when_missing_or_invalid():
+    frame = b'event: error\ndata: {"error": "no reason given"}\n\n'
+    response = _sse_response(_SyncSource([frame]))
+    result = WorkflowStreamErrorHook().after_success(
+        _hook_ctx(STREAM_OPERATION_ID), response
+    )
+
+    with pytest.raises(StreamDisconnectedError) as exc_info:
+        list(result.iter_bytes())
+
+    assert exc_info.value.reason == "stream_error"
+    assert exc_info.value.error == "no reason given"
+
+
+def test_hook_passes_normal_stream_through_without_raising():
+    response = _sse_response(_SyncSource([GOOD_FRAME, GOOD_FRAME]))
+    result = WorkflowStreamErrorHook().after_success(
+        _hook_ctx(STREAM_OPERATION_ID), response
+    )
+
+    body = b"".join(result.iter_bytes())
+    assert body.count(b"workflow.event") == 2
+
+
+def test_hook_ignores_non_stream_operations():
+    source = _SyncSource([ERROR_FRAME])
+    response = _sse_response(source)
+    result = WorkflowStreamErrorHook().after_success(
+        _hook_ctx(NON_STREAM_OPERATION_ID), response
+    )
+
+    # Response is returned untouched: same object, original stream not wrapped.
+    assert result is response
+    assert response.stream is source

src/mistralai/extra/workflows/__init__.py

@@ -22,6 +22,10 @@
     configure_workflow_encoding,
     generate_two_part_id,
 )
+from .errors import (
+    StreamDisconnectReason,
+    StreamDisconnectedError,
+)
 
 __all__ = [
     "ConnectorAuthTaskState",
@@ -42,4 +46,6 @@
     "EncryptedStrField",
     "configure_workflow_encoding",
     "generate_two_part_id",
+    "StreamDisconnectedError",
+    "StreamDisconnectReason",
 ]

src/mistralai/extra/workflows/errors.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+from typing import Literal
+
+from mistralai.extra.exceptions import MistralClientException
+
+StreamDisconnectReason = Literal["read_error", "stream_error", "internal_error"]
+
+
+class StreamDisconnectedError(MistralClientException):
+    """Raised when a workflow SSE stream is terminated by a server error frame.
+
+    The server ends a stream by emitting an ``event: error`` SSE frame. The SDK
+    surfaces this as a raised exception so consumers can wrap stream iteration in
+    ``try`` / ``except`` instead of inspecting each event for ``event == "error"``.
+
+    Both attributes are populated from the frame's ``data`` JSON payload.
+    """
+
+    def __init__(self, *, reason: StreamDisconnectReason, error: str) -> None:
+        self.reason: StreamDisconnectReason = reason
+        self.error = error
+        super().__init__("Workflow stream disconnected by server")

src/mistralai/extra/workflows/stream_error_hook.py

@@ -0,0 +1,189 @@
+from __future__ import annotations
+
+import json
+import re
+from typing import Any, AsyncIterator, Iterator, Optional, Tuple, Union
+
+import httpx
+from httpx._types import AsyncByteStream, SyncByteStream
+
+from mistralai.client._hooks.types import AfterSuccessContext, AfterSuccessHook
+from mistralai.extra.workflows.errors import (
+    StreamDisconnectReason,
+    StreamDisconnectedError,
+)
+
+# Operation IDs of the two SSE-backed workflow stream endpoints.
+STREAM_OPERATIONS = {
+    "get_stream_events_v1_workflows_events_stream_get",
+    "stream_v1_workflows_executions__execution_id__stream_get",
+}
+
+_ERROR_EVENT = "error"
+_VALID_REASONS = ("read_error", "stream_error", "internal_error")
+_DEFAULT_REASON: StreamDisconnectReason = "stream_error"
+
+# SSE frame boundaries (blank line), longest first so the full separator is consumed.
+_BOUNDARIES = [
+    b"\r\n\r\n",
+    b"\r\n\r",
+    b"\r\n\n",
+    b"\r\r\n",
+    b"\n\r\n",
+    b"\r\r",
+    b"\n\r",
+    b"\n\n",
+]
+
+
+def _strip_content_encoding_header(headers: httpx.Headers) -> httpx.Headers:
+    return httpx.Headers(
+        [(k, v) for k, v in headers.items() if k.lower() != "content-encoding"]
+    )
+
+
+def _find_boundary(buffer: bytearray) -> Optional[Tuple[int, int]]:
+    """Return (index, length) of the earliest frame boundary, or None if incomplete."""
+    best: Optional[Tuple[int, int]] = None
+    for boundary in _BOUNDARIES:
+        idx = buffer.find(boundary)
+        if idx == -1:
+            continue
+        if (
+            best is None
+            or idx < best[0]
+            or (idx == best[0] and len(boundary) > best[1])
+        ):
+            best = (idx, len(boundary))
+    return best
+
+
+def _parse_error_payload(data: str) -> Tuple[str, StreamDisconnectReason]:
+    payload: dict[str, Any] = {}
+    try:
+        parsed = json.loads(data.strip())
+        if isinstance(parsed, dict):
+            payload = parsed
+    except json.JSONDecodeError:
+        pass
+    error = str(payload.get("error", data.strip()))
+    reason = payload.get("reason", _DEFAULT_REASON)
+    if reason not in _VALID_REASONS:
+        reason = _DEFAULT_REASON
+    return error, reason
+
+
+def _raise_if_error_frame(block: bytes) -> None:
+    """Raise StreamDisconnectedError if the SSE frame is an ``event: error`` frame."""
+    event_name: Optional[str] = None
+    data = ""
+    for line in re.split(r"\r?\n|\r", block.decode("utf-8", errors="replace")):
+        if not line or line.startswith(":"):
+            continue
+        field, _, value = line.partition(":")
+        if value.startswith(" "):
+            value = value[1:]
+        if field == "event":
+            event_name = value
+        elif field == "data":
+            data += value + "\n"
+
+    if event_name != _ERROR_EVENT:
+        return
+
+    error, reason = _parse_error_payload(data)
+    raise StreamDisconnectedError(reason=reason, error=error)
+
+
+class _FrameScanner:
+    """Buffers raw SSE bytes, raising on error frames and passing others through."""
+
+    def __init__(self) -> None:
+        self._buffer = bytearray()
+
+    def feed(self, chunk: bytes) -> Iterator[bytes]:
+        self._buffer += chunk
+        while True:
+            found = _find_boundary(self._buffer)
+            if found is None:
+                return
+            idx, length = found
+            block = bytes(self._buffer[:idx])
+            frame = bytes(self._buffer[: idx + length])
+            del self._buffer[: idx + length]
+            _raise_if_error_frame(block)
+            yield frame
+
+    def flush(self) -> Iterator[bytes]:
+        if not self._buffer:
+            return
+        block = bytes(self._buffer)
+        self._buffer.clear()
+        _raise_if_error_frame(block)
+        yield block
+
+
+class _ErrorDetectingSyncByteStream(SyncByteStream):
+    def __init__(self, original: SyncByteStream) -> None:
+        self._original = original
+        self._scanner = _FrameScanner()
+
+    def __iter__(self) -> Iterator[bytes]:
+        for chunk in self._original:
+            yield from self._scanner.feed(chunk)
+        yield from self._scanner.flush()
+
+    def close(self) -> None:
+        self._original.close()
+
+
+class _ErrorDetectingAsyncByteStream(AsyncByteStream):
+    def __init__(self, original: AsyncByteStream) -> None:
+        self._original = original
+        self._scanner = _FrameScanner()
+
+    async def __aiter__(self) -> AsyncIterator[bytes]:
+        async for chunk in self._original:
+            for frame in self._scanner.feed(chunk):
+                yield frame
+        for frame in self._scanner.flush():
+            yield frame
+
+    async def aclose(self) -> None:
+        await self._original.aclose()
+
+
+class WorkflowStreamErrorHook(AfterSuccessHook):
+    """Raise StreamDisconnectedError when a workflow SSE stream sends an error frame.
+
+    Wraps the response byte stream for the two workflow SSE operations so that an
+    ``event: error`` frame raises during iteration, terminating the consumer's
+    ``for event in stream`` loop instead of yielding the error as a normal event.
+    """
+
+    def after_success(
+        self,
+        hook_ctx: AfterSuccessContext,
+        response: httpx.Response,
+    ) -> Union[httpx.Response, Exception]:
+        if hook_ctx.operation_id not in STREAM_OPERATIONS:
+            return response
+        if "text/event-stream" not in response.headers.get("content-type", ""):
+            return response
+
+        stream = response.stream
+        wrapped: Union[SyncByteStream, AsyncByteStream]
+        if isinstance(stream, AsyncByteStream):
+            wrapped = _ErrorDetectingAsyncByteStream(stream)
+        elif isinstance(stream, SyncByteStream):
+            wrapped = _ErrorDetectingSyncByteStream(stream)
+        else:
+            return response
+
+        return httpx.Response(
+            status_code=response.status_code,
+            headers=_strip_content_encoding_header(response.headers),
+            stream=wrapped,
+            request=response.request,
+            extensions=response.extensions,
+        )