diff --git a/src/mistralai/client/_hooks/registration.py b/src/mistralai/client/_hooks/registration.py index 4da6eb7c..f262aee1 100644 --- a/src/mistralai/client/_hooks/registration.py +++ b/src/mistralai/client/_hooks/registration.py @@ -3,6 +3,7 @@ from .traceparent import TraceparentInjectionHook from .tracing import TracingHook from .types import Hooks +from .stream_error_hook import WorkflowStreamErrorHook from .workflow_encoding_hook import WorkflowEncodingHook # This file is only ever generated once on the first generation and then is free to be modified. @@ -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/client/_hooks/stream_error_hook.py b/src/mistralai/client/_hooks/stream_error_hook.py new file mode 100644 index 00000000..71e6c7bb --- /dev/null +++ b/src/mistralai/client/_hooks/stream_error_hook.py @@ -0,0 +1,188 @@ +import json +import re +from typing import Any, AsyncIterator, Dict, Iterator, Optional, Tuple, Union, get_args + +import httpx +from httpx._types import AsyncByteStream, SyncByteStream + +from .types import AfterSuccessContext, AfterSuccessHook +from mistralai.extra.exceptions import ( + StreamDisconnectReason, + StreamDisconnectedError, +) + +# Operation IDs of the SSE-backed workflow stream endpoints that can emit a +# terminal ``event: error`` frame (event, execution, and logs streams). +STREAM_OPERATIONS_WITH_ERROR_EVENT = { + "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 = get_args(StreamDisconnectReason) +_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 _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: + # 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: + 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 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_WITH_ERROR_EVENT: + 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 + + # 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=response.headers, + stream=wrapped, + request=response.request, + extensions=response.extensions, + ) 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 new file mode 100644 index 00000000..c6762e08 --- /dev/null +++ b/src/mistralai/extra/tests/test_stream_error_hook.py @@ -0,0 +1,257 @@ +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, +) + +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 + ) + assert isinstance(result, httpx.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_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_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])) + 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 == "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 + ) + assert isinstance(result, httpx.Response) + + body = b"".join(result.iter_bytes()) + assert body.count(b"workflow.event") == 2 + + +@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) + 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" + + +@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) + 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