From 77bcbe76c0b44e06512e3ea6b9942b4f7d16689d Mon Sep 17 00:00:00 2001 From: Nicolas Faurie Date: Wed, 22 Jul 2026 14:40:44 +0200 Subject: [PATCH 1/7] 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. --- src/mistralai/client/_hooks/registration.py | 2 + .../extra/tests/test_stream_error_hook.py | 154 ++++++++++++++ src/mistralai/extra/workflows/__init__.py | 6 + src/mistralai/extra/workflows/errors.py | 23 +++ .../extra/workflows/stream_error_hook.py | 189 ++++++++++++++++++ 5 files changed, 374 insertions(+) create mode 100644 src/mistralai/extra/tests/test_stream_error_hook.py create mode 100644 src/mistralai/extra/workflows/errors.py create mode 100644 src/mistralai/extra/workflows/stream_error_hook.py diff --git a/src/mistralai/client/_hooks/registration.py b/src/mistralai/client/_hooks/registration.py index 4da6eb7c..c9539eb5 100644 --- a/src/mistralai/client/_hooks/registration.py +++ b/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()) diff --git a/src/mistralai/extra/tests/test_stream_error_hook.py b/src/mistralai/extra/tests/test_stream_error_hook.py new file mode 100644 index 00000000..5277f41a --- /dev/null +++ b/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 diff --git a/src/mistralai/extra/workflows/__init__.py b/src/mistralai/extra/workflows/__init__.py index cc806268..ea0a8523 100644 --- a/src/mistralai/extra/workflows/__init__.py +++ b/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", ] diff --git a/src/mistralai/extra/workflows/errors.py b/src/mistralai/extra/workflows/errors.py new file mode 100644 index 00000000..98043838 --- /dev/null +++ b/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") diff --git a/src/mistralai/extra/workflows/stream_error_hook.py b/src/mistralai/extra/workflows/stream_error_hook.py new file mode 100644 index 00000000..2a348b99 --- /dev/null +++ b/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, + ) From 874a5939f6347ee5ebfcdabb025abfdf62579f12 Mon Sep 17 00:00:00 2001 From: Nicolas Faurie Date: Wed, 22 Jul 2026 14:46:37 +0200 Subject: [PATCH 2/7] refactor(workflows): move StreamDisconnectedError into extra/exceptions.py --- src/mistralai/extra/exceptions.py | 21 ++++++++++++++++- .../extra/tests/test_stream_error_hook.py | 2 +- src/mistralai/extra/workflows/__init__.py | 6 ----- src/mistralai/extra/workflows/errors.py | 23 ------------------- .../extra/workflows/stream_error_hook.py | 2 +- 5 files changed, 22 insertions(+), 32 deletions(-) delete mode 100644 src/mistralai/extra/workflows/errors.py diff --git a/src/mistralai/extra/exceptions.py b/src/mistralai/extra/exceptions.py index a51d5dc3..82507f2c 100644 --- a/src/mistralai/extra/exceptions.py +++ b/src/mistralai/extra/exceptions.py @@ -2,7 +2,7 @@ from dataclasses import dataclass from enum import Enum -from typing import Any, Optional, Union, TYPE_CHECKING +from typing import Any, Literal, Optional, Union, TYPE_CHECKING import typing from mistralai.client.models import ( @@ -32,6 +32,25 @@ class WorkflowPayloadCompressionException(MistralClientException): """Workflow payload compression exception""" +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") + + class RunException(MistralClientException): """Conversation run errors.""" diff --git a/src/mistralai/extra/tests/test_stream_error_hook.py b/src/mistralai/extra/tests/test_stream_error_hook.py index 5277f41a..8dc5c9fb 100644 --- a/src/mistralai/extra/tests/test_stream_error_hook.py +++ b/src/mistralai/extra/tests/test_stream_error_hook.py @@ -4,7 +4,7 @@ from mistralai.client import Mistral from mistralai.client._hooks.types import AfterSuccessContext, HookContext -from mistralai.extra.workflows.errors import StreamDisconnectedError +from mistralai.extra.exceptions import StreamDisconnectedError from mistralai.extra.workflows.stream_error_hook import WorkflowStreamErrorHook STREAM_OPERATION_ID = "get_stream_events_v1_workflows_events_stream_get" diff --git a/src/mistralai/extra/workflows/__init__.py b/src/mistralai/extra/workflows/__init__.py index ea0a8523..cc806268 100644 --- a/src/mistralai/extra/workflows/__init__.py +++ b/src/mistralai/extra/workflows/__init__.py @@ -22,10 +22,6 @@ configure_workflow_encoding, generate_two_part_id, ) -from .errors import ( - StreamDisconnectReason, - StreamDisconnectedError, -) __all__ = [ "ConnectorAuthTaskState", @@ -46,6 +42,4 @@ "EncryptedStrField", "configure_workflow_encoding", "generate_two_part_id", - "StreamDisconnectedError", - "StreamDisconnectReason", ] diff --git a/src/mistralai/extra/workflows/errors.py b/src/mistralai/extra/workflows/errors.py deleted file mode 100644 index 98043838..00000000 --- a/src/mistralai/extra/workflows/errors.py +++ /dev/null @@ -1,23 +0,0 @@ -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") diff --git a/src/mistralai/extra/workflows/stream_error_hook.py b/src/mistralai/extra/workflows/stream_error_hook.py index 2a348b99..0cf6d0b5 100644 --- a/src/mistralai/extra/workflows/stream_error_hook.py +++ b/src/mistralai/extra/workflows/stream_error_hook.py @@ -8,7 +8,7 @@ from httpx._types import AsyncByteStream, SyncByteStream from mistralai.client._hooks.types import AfterSuccessContext, AfterSuccessHook -from mistralai.extra.workflows.errors import ( +from mistralai.extra.exceptions import ( StreamDisconnectReason, StreamDisconnectedError, ) From a7d9e33217f4b6a1fdc478a72c9ffac450f043db Mon Sep 17 00:00:00 2001 From: Nicolas Faurie Date: Wed, 22 Jul 2026 14:50:50 +0200 Subject: [PATCH 3/7] test: narrow hook result to Response before iterating (pyright) --- src/mistralai/extra/tests/test_stream_error_hook.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/mistralai/extra/tests/test_stream_error_hook.py b/src/mistralai/extra/tests/test_stream_error_hook.py index 8dc5c9fb..035015e1 100644 --- a/src/mistralai/extra/tests/test_stream_error_hook.py +++ b/src/mistralai/extra/tests/test_stream_error_hook.py @@ -110,6 +110,7 @@ def test_hook_detects_error_frame_split_across_chunks(): result = WorkflowStreamErrorHook().after_success( _hook_ctx(STREAM_OPERATION_ID), response ) + assert isinstance(result, httpx.Response) with pytest.raises(StreamDisconnectedError) as exc_info: list(result.iter_bytes()) @@ -124,6 +125,7 @@ def test_hook_defaults_reason_when_missing_or_invalid(): result = WorkflowStreamErrorHook().after_success( _hook_ctx(STREAM_OPERATION_ID), response ) + assert isinstance(result, httpx.Response) with pytest.raises(StreamDisconnectedError) as exc_info: list(result.iter_bytes()) @@ -137,6 +139,7 @@ def test_hook_passes_normal_stream_through_without_raising(): result = WorkflowStreamErrorHook().after_success( _hook_ctx(STREAM_OPERATION_ID), response ) + assert isinstance(result, httpx.Response) body = b"".join(result.iter_bytes()) assert body.count(b"workflow.event") == 2 From 472daf225097ee2d0123e3d4ae42800fec2fef00 Mon Sep 17 00:00:00 2001 From: Nicolas Faurie Date: Wed, 22 Jul 2026 17:20:21 +0200 Subject: [PATCH 4/7] refactor(workflows): move stream error hook into _hooks and cover logs streams - Move WorkflowStreamErrorHook to client/_hooks alongside the other hooks - Cover deployment + execution logs SSE streams (same error-frame contract) - Derive valid reasons from the StreamDisconnectReason Literal - Add tests: error frame without trailing boundary; all stream operations --- src/mistralai/client/_hooks/registration.py | 2 +- .../_hooks}/stream_error_hook.py | 17 +++++----- .../extra/tests/test_stream_error_hook.py | 33 ++++++++++++++++++- 3 files changed, 42 insertions(+), 10 deletions(-) rename src/mistralai/{extra/workflows => client/_hooks}/stream_error_hook.py (91%) diff --git a/src/mistralai/client/_hooks/registration.py b/src/mistralai/client/_hooks/registration.py index c9539eb5..f262aee1 100644 --- a/src/mistralai/client/_hooks/registration.py +++ b/src/mistralai/client/_hooks/registration.py @@ -3,8 +3,8 @@ from .traceparent import TraceparentInjectionHook from .tracing import TracingHook from .types import Hooks +from .stream_error_hook import WorkflowStreamErrorHook 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 diff --git a/src/mistralai/extra/workflows/stream_error_hook.py b/src/mistralai/client/_hooks/stream_error_hook.py similarity index 91% rename from src/mistralai/extra/workflows/stream_error_hook.py rename to src/mistralai/client/_hooks/stream_error_hook.py index 0cf6d0b5..aa6dc5ef 100644 --- a/src/mistralai/extra/workflows/stream_error_hook.py +++ b/src/mistralai/client/_hooks/stream_error_hook.py @@ -1,26 +1,27 @@ -from __future__ import annotations - import json import re -from typing import Any, AsyncIterator, Iterator, Optional, Tuple, Union +from typing import Any, AsyncIterator, Dict, Iterator, Optional, Tuple, Union, get_args import httpx from httpx._types import AsyncByteStream, SyncByteStream -from mistralai.client._hooks.types import AfterSuccessContext, AfterSuccessHook +from .types import AfterSuccessContext, AfterSuccessHook from mistralai.extra.exceptions import ( StreamDisconnectReason, StreamDisconnectedError, ) -# Operation IDs of the two SSE-backed workflow stream endpoints. +# Operation IDs of the SSE-backed workflow stream endpoints that can emit a +# terminal ``event: error`` frame (event, execution, and logs streams). STREAM_OPERATIONS = { "get_stream_events_v1_workflows_events_stream_get", "stream_v1_workflows_executions__execution_id__stream_get", + "stream_deployment_logs", + "stream_workflow_execution_logs", } _ERROR_EVENT = "error" -_VALID_REASONS = ("read_error", "stream_error", "internal_error") +_VALID_REASONS = get_args(StreamDisconnectReason) _DEFAULT_REASON: StreamDisconnectReason = "stream_error" # SSE frame boundaries (blank line), longest first so the full separator is consumed. @@ -59,7 +60,7 @@ def _find_boundary(buffer: bytearray) -> Optional[Tuple[int, int]]: def _parse_error_payload(data: str) -> Tuple[str, StreamDisconnectReason]: - payload: dict[str, Any] = {} + payload: Dict[str, Any] = {} try: parsed = json.loads(data.strip()) if isinstance(parsed, dict): @@ -156,7 +157,7 @@ async def aclose(self) -> None: 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 + Wraps the response byte stream for the 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. """ diff --git a/src/mistralai/extra/tests/test_stream_error_hook.py b/src/mistralai/extra/tests/test_stream_error_hook.py index 035015e1..0d10d3f9 100644 --- a/src/mistralai/extra/tests/test_stream_error_hook.py +++ b/src/mistralai/extra/tests/test_stream_error_hook.py @@ -5,7 +5,10 @@ from mistralai.client import Mistral from mistralai.client._hooks.types import AfterSuccessContext, HookContext from mistralai.extra.exceptions import StreamDisconnectedError -from mistralai.extra.workflows.stream_error_hook import WorkflowStreamErrorHook +from mistralai.client._hooks.stream_error_hook import ( + STREAM_OPERATIONS, + WorkflowStreamErrorHook, +) STREAM_OPERATION_ID = "get_stream_events_v1_workflows_events_stream_get" NON_STREAM_OPERATION_ID = "chat_completion_v1_chat_completions_post" @@ -119,6 +122,21 @@ def test_hook_detects_error_frame_split_across_chunks(): assert exc_info.value.error == "x" +def test_hook_raises_on_error_frame_without_trailing_boundary(): + frame = b'event: error\ndata: {"error": "boom", "reason": "read_error"}' + response = _sse_response(_SyncSource([frame])) + result = WorkflowStreamErrorHook().after_success( + _hook_ctx(STREAM_OPERATION_ID), response + ) + assert isinstance(result, httpx.Response) + + with pytest.raises(StreamDisconnectedError) as exc_info: + list(result.iter_bytes()) + + assert exc_info.value.reason == "read_error" + assert exc_info.value.error == "boom" + + 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])) @@ -145,6 +163,19 @@ def test_hook_passes_normal_stream_through_without_raising(): assert body.count(b"workflow.event") == 2 +@pytest.mark.parametrize("operation_id", sorted(STREAM_OPERATIONS)) +def test_hook_raises_for_every_workflow_stream_operation(operation_id: str): + response = _sse_response(_SyncSource([ERROR_FRAME])) + result = WorkflowStreamErrorHook().after_success(_hook_ctx(operation_id), response) + assert isinstance(result, httpx.Response) + + with pytest.raises(StreamDisconnectedError) as exc_info: + list(result.iter_bytes()) + + assert exc_info.value.reason == "read_error" + assert exc_info.value.error == "boom" + + def test_hook_ignores_non_stream_operations(): source = _SyncSource([ERROR_FRAME]) response = _sse_response(source) From 6a711e1d11c289cb881e0a6bd526071c2fe45f09 Mon Sep 17 00:00:00 2001 From: Nicolas Faurie Date: Wed, 22 Jul 2026 17:23:29 +0200 Subject: [PATCH 5/7] refactor: rename to STREAM_OPERATIONS_WITH_ERROR_EVENT for clarity --- src/mistralai/client/_hooks/stream_error_hook.py | 4 ++-- src/mistralai/extra/tests/test_stream_error_hook.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/mistralai/client/_hooks/stream_error_hook.py b/src/mistralai/client/_hooks/stream_error_hook.py index aa6dc5ef..b6c10768 100644 --- a/src/mistralai/client/_hooks/stream_error_hook.py +++ b/src/mistralai/client/_hooks/stream_error_hook.py @@ -13,7 +13,7 @@ # Operation IDs of the SSE-backed workflow stream endpoints that can emit a # terminal ``event: error`` frame (event, execution, and logs streams). -STREAM_OPERATIONS = { +STREAM_OPERATIONS_WITH_ERROR_EVENT = { "get_stream_events_v1_workflows_events_stream_get", "stream_v1_workflows_executions__execution_id__stream_get", "stream_deployment_logs", @@ -167,7 +167,7 @@ def after_success( hook_ctx: AfterSuccessContext, response: httpx.Response, ) -> Union[httpx.Response, Exception]: - if hook_ctx.operation_id not in STREAM_OPERATIONS: + if hook_ctx.operation_id not in STREAM_OPERATIONS_WITH_ERROR_EVENT: return response if "text/event-stream" not in response.headers.get("content-type", ""): return response diff --git a/src/mistralai/extra/tests/test_stream_error_hook.py b/src/mistralai/extra/tests/test_stream_error_hook.py index 0d10d3f9..66e6963a 100644 --- a/src/mistralai/extra/tests/test_stream_error_hook.py +++ b/src/mistralai/extra/tests/test_stream_error_hook.py @@ -6,7 +6,7 @@ from mistralai.client._hooks.types import AfterSuccessContext, HookContext from mistralai.extra.exceptions import StreamDisconnectedError from mistralai.client._hooks.stream_error_hook import ( - STREAM_OPERATIONS, + STREAM_OPERATIONS_WITH_ERROR_EVENT, WorkflowStreamErrorHook, ) @@ -163,7 +163,7 @@ def test_hook_passes_normal_stream_through_without_raising(): assert body.count(b"workflow.event") == 2 -@pytest.mark.parametrize("operation_id", sorted(STREAM_OPERATIONS)) +@pytest.mark.parametrize("operation_id", sorted(STREAM_OPERATIONS_WITH_ERROR_EVENT)) def test_hook_raises_for_every_workflow_stream_operation(operation_id: str): response = _sse_response(_SyncSource([ERROR_FRAME])) result = WorkflowStreamErrorHook().after_success(_hook_ctx(operation_id), response) From 584fab7914ef9fc1ae353ab4f80bf347ff6e0bbc Mon Sep 17 00:00:00 2001 From: Nicolas Faurie Date: Wed, 22 Jul 2026 17:28:21 +0200 Subject: [PATCH 6/7] refactor: drop redundant content-encoding strip; test hook composition - Error hook forwards the raw stream unchanged, so it must keep Content-Encoding (httpx decodes downstream); removes duplication with workflow_encoding_hook - Add test that the encoding hook + stream error hook compose correctly --- .../client/_hooks/stream_error_hook.py | 10 ++-- .../extra/tests/test_stream_error_hook.py | 50 +++++++++++++++++++ 2 files changed, 53 insertions(+), 7 deletions(-) diff --git a/src/mistralai/client/_hooks/stream_error_hook.py b/src/mistralai/client/_hooks/stream_error_hook.py index b6c10768..deb125b0 100644 --- a/src/mistralai/client/_hooks/stream_error_hook.py +++ b/src/mistralai/client/_hooks/stream_error_hook.py @@ -37,12 +37,6 @@ ] -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 @@ -181,9 +175,11 @@ def after_success( else: return response + # Keep the original headers: this hook forwards the raw stream unchanged, + # so httpx still applies any Content-Encoding when the consumer iterates. return httpx.Response( status_code=response.status_code, - headers=_strip_content_encoding_header(response.headers), + headers=response.headers, stream=wrapped, request=response.request, extensions=response.extensions, diff --git a/src/mistralai/extra/tests/test_stream_error_hook.py b/src/mistralai/extra/tests/test_stream_error_hook.py index 66e6963a..61e58286 100644 --- a/src/mistralai/extra/tests/test_stream_error_hook.py +++ b/src/mistralai/extra/tests/test_stream_error_hook.py @@ -1,10 +1,20 @@ import httpx import pytest from httpx._types import AsyncByteStream, SyncByteStream +from pydantic import SecretStr from mistralai.client import Mistral from mistralai.client._hooks.types import AfterSuccessContext, HookContext +from mistralai.client._hooks.workflow_encoding_hook import ( + WorkflowEncodingHook, + configure_workflow_encoding, +) from mistralai.extra.exceptions import StreamDisconnectedError +from mistralai.extra.workflows import ( + PayloadEncryptionConfig, + PayloadEncryptionMode, + WorkflowEncodingConfig, +) from mistralai.client._hooks.stream_error_hook import ( STREAM_OPERATIONS_WITH_ERROR_EVENT, WorkflowStreamErrorHook, @@ -176,6 +186,46 @@ def test_hook_raises_for_every_workflow_stream_operation(operation_id: str): assert exc_info.value.error == "boom" +@pytest.mark.asyncio +async def test_encoding_and_error_hooks_compose(): + client = Mistral(api_key="test-key") + configure_workflow_encoding( + WorkflowEncodingConfig( + payload_encryption=PayloadEncryptionConfig( + mode=PayloadEncryptionMode.FULL, main_key=SecretStr("0" * 64) + ) + ), + namespace="demo", + sdk_config=client.sdk_configuration, + ) + ctx = AfterSuccessContext( + HookContext( + config=client.sdk_configuration, + base_url="https://api.example.com", + operation_id=STREAM_OPERATION_ID, + oauth2_scopes=[], + security_source=None, + ) + ) + benign = b'event: message\ndata: {"hello": "world"}\n\n' + response = _sse_response(_AsyncSource([benign, ERROR_FRAME])) + + # Encoding hook wraps first (decryption), then the error hook wraps that stream. + decrypted = WorkflowEncodingHook().after_success(ctx, response) + assert isinstance(decrypted, httpx.Response) + final = WorkflowStreamErrorHook().after_success(ctx, decrypted) + assert isinstance(final, httpx.Response) + + collected: list[bytes] = [] + with pytest.raises(StreamDisconnectedError) as exc_info: + async for chunk in final.aiter_bytes(): + collected.append(chunk) + + assert exc_info.value.reason == "read_error" + assert exc_info.value.error == "boom" + assert b'"hello": "world"' in b"".join(collected) + + def test_hook_ignores_non_stream_operations(): source = _SyncSource([ERROR_FRAME]) response = _sse_response(source) From e49af8e9de3995537e010abb742da01ae78fcc13 Mon Sep 17 00:00:00 2001 From: Nicolas Faurie Date: Wed, 22 Jul 2026 17:32:42 +0200 Subject: [PATCH 7/7] fix: parse error frames with data split across multiple SSE data lines Use json.loads(strict=False) so a value spanning several data: lines (rejoined with a literal newline per the SSE spec) still yields reason/error instead of silently falling back. --- .../client/_hooks/stream_error_hook.py | 4 +++- .../extra/tests/test_stream_error_hook.py | 19 +++++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/src/mistralai/client/_hooks/stream_error_hook.py b/src/mistralai/client/_hooks/stream_error_hook.py index deb125b0..71e6c7bb 100644 --- a/src/mistralai/client/_hooks/stream_error_hook.py +++ b/src/mistralai/client/_hooks/stream_error_hook.py @@ -56,7 +56,9 @@ def _find_boundary(buffer: bytearray) -> Optional[Tuple[int, int]]: def _parse_error_payload(data: str) -> Tuple[str, StreamDisconnectReason]: payload: Dict[str, Any] = {} try: - parsed = json.loads(data.strip()) + # strict=False: SSE joins multi-line data with "\n", so a value spanning + # several data: lines contains literal newlines that strict JSON rejects. + parsed = json.loads(data.strip(), strict=False) if isinstance(parsed, dict): payload = parsed except json.JSONDecodeError: diff --git a/src/mistralai/extra/tests/test_stream_error_hook.py b/src/mistralai/extra/tests/test_stream_error_hook.py index 61e58286..c6762e08 100644 --- a/src/mistralai/extra/tests/test_stream_error_hook.py +++ b/src/mistralai/extra/tests/test_stream_error_hook.py @@ -147,6 +147,25 @@ def test_hook_raises_on_error_frame_without_trailing_boundary(): assert exc_info.value.error == "boom" +def test_hook_handles_error_frame_split_across_data_lines(): + frame = ( + b"event: error\r\n" + b'data: {"error": "connection\r\n' + b'data: lost", "reason": "read_error"}\r\n\r\n' + ) + response = _sse_response(_SyncSource([frame])) + result = WorkflowStreamErrorHook().after_success( + _hook_ctx(STREAM_OPERATION_ID), response + ) + assert isinstance(result, httpx.Response) + + with pytest.raises(StreamDisconnectedError) as exc_info: + list(result.iter_bytes()) + + assert exc_info.value.reason == "read_error" + assert exc_info.value.error == "connection\nlost" + + 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]))