Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,31 @@ async for chunk in ai.ui.ai_sdk.to_sse(agent_stream):
yield chunk
```

If the agent stream raises, `to_sse` sends a safe `error` event and closes the
stream with `[DONE]`. Pass `on_error` to log the exception or return a more
specific client-facing message:

```python
import logging

logger = logging.getLogger(__name__)


def error_text(error: Exception) -> str:
logger.exception("Agent stream failed", exc_info=error)
return "The model is temporarily unavailable."


async for chunk in ai.ui.ai_sdk.to_sse(
agent_stream,
on_error=error_text,
):
yield chunk
```

Without `on_error`, the client receives `An error occurred.` so server-side
details are not exposed.

```python
async for part in ai.ui.ai_sdk.to_stream(agent_stream):
yield part
Expand Down
23 changes: 19 additions & 4 deletions src/ai/ui/ai_sdk/outbound_stream.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,12 @@
from .tool_utils import normalize_tool_input

if TYPE_CHECKING:
from collections.abc import AsyncGenerator, AsyncIterable
from collections.abc import AsyncGenerator, AsyncIterable, Callable


def _default_error_text(_: Exception) -> str:
"""Return a safe client-facing message for stream failures."""
return "An error occurred."


def _tool_error_text(part: messages_.ToolResultPart) -> str:
Expand Down Expand Up @@ -621,8 +626,18 @@ async def to_stream(

async def to_sse(
events: AsyncIterable[events_.AgentEvent],
*,
on_error: Callable[[Exception], str] = _default_error_text,
) -> AsyncGenerator[str]:
"""Convert an internal event stream into SSE strings."""
async for event in to_stream(events):
yield format_sse(event)
"""Convert an internal event stream into SSE strings.

Stream failures become AI SDK UI ``error`` events. The default message
avoids exposing server-side details; use ``on_error`` to log the exception
or return a more specific client-facing message.
"""
try:
async for event in to_stream(events):
yield format_sse(event)
except Exception as error:
yield format_sse(ui_events.UIErrorEvent(error_text=on_error(error)))
yield format_done_sse()
31 changes: 31 additions & 0 deletions tests/ui/ai_sdk/test_outbound_stream.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,11 @@ async def _gen(
yield event


async def _broken_gen() -> AsyncGenerator[agent_events_.AgentEvent]:
yield events_.TextStart(block_id="t1")
raise RuntimeError("sensitive provider detail")


async def _collect(
stream_events: list[agent_events_.AgentEvent],
) -> list[ui_events.UIMessageStreamEvent]:
Expand Down Expand Up @@ -102,6 +107,32 @@ async def test_to_sse_emits_data_prefixed_lines() -> None:
assert lines[-1] == "data: [DONE]\n\n"


async def test_to_sse_converts_stream_errors_to_safe_error_events() -> None:
lines = [line async for line in to_sse(_broken_gen())]

error = json.loads(lines[-2].removeprefix("data: ").rstrip())
assert error == {"type": "error", "errorText": "An error occurred."}
assert lines[-1] == "data: [DONE]\n\n"


async def test_to_sse_supports_custom_error_messages() -> None:
errors: list[Exception] = []

def on_error(error: Exception) -> str:
errors.append(error)
return "The model is temporarily unavailable."

lines = [line async for line in to_sse(_broken_gen(), on_error=on_error)]

error = json.loads(lines[-2].removeprefix("data: ").rstrip())
assert error == {
"type": "error",
"errorText": "The model is temporarily unavailable.",
}
assert len(errors) == 1
assert str(errors[0]) == "sensitive provider detail"


async def test_stream_start_uses_runtime_message_id() -> None:
assistant = messages_.Message(
id="assistant-runtime-id",
Expand Down