Repository: pydantic/pydantic-ai
File: pydantic_ai_slim/pydantic_ai/ag_ui.py (Python)

Code before:
            usage_limits=usage_limits,
            usage=usage,
            infer_name=infer_name,
            toolsets=toolsets,
        ) as run:
            completed_run = run
            async for event in _agent_stream(run):
                yield encoder.encode(event)
    except _RunError as e:
        yield encoder.encode(
            RunErrorEvent(message=e.message, code=e.code),
        )
    except Exception as e:
        yield encoder.encode(
            RunErrorEvent(message=str(e)),
        )
        raise e
    else:
        yield encoder.encode(
            RunFinishedEvent(
                thread_id=run_input.thread_id,
                run_id=run_input.run_id,
            ),
        )
    finally:
        # Call the callback if provided and we have a completed run
        if on_complete is not None and completed_run is not None:
            import inspect

            if inspect.iscoroutinefunction(on_complete):
                await on_complete(completed_run)
            else:
                on_complete(completed_run)


async def _agent_stream(run: AgentRun[AgentDepsT, Any]) -> AsyncIterator[BaseEvent]:
    """Run the agent streaming responses using AG-UI protocol events.

    Args:
        run: The agent run to process.

    Yields:
        AG-UI Server-Sent Events (SSE).
    """
    async for node in run:
        stream_ctx = _RequestStreamContext()
        if isinstance(node, ModelRequestNode):
            async with node.stream(run.ctx) as request_stream:
                async for agent_event in request_stream:
                    async for msg in _handle_model_request_event(stream_ctx, agent_event):
                        yield msg

Review comment (suggestion):
Let's use `await _utils.run_in_executor(on_complete, completed_run)` like we do for all other user-registered functions.

Code after:
                await _utils.run_in_executor(on_complete, run.result)
    except _RunError as e:
        yield encoder.encode(
            RunErrorEvent(message=e.message, code=e.code),
        )
    except Exception as e:
        yield encoder.encode(
            RunErrorEvent(message=str(e)),
        )
        raise e
    else:
        yield encoder.encode(
            RunFinishedEvent(
                thread_id=run_input.thread_id,
                run_id=run_input.run_id,
            ),
        )


async def _agent_stream(run: AgentRun[AgentDepsT, Any]) -> AsyncIterator[BaseEvent]:
    """Run the agent streaming responses using AG-UI protocol events.

    Args:
        run: The agent run to process.

    Yields:
        AG-UI Server-Sent Events (SSE).
    """
    async for node in run:
        stream_ctx = _RequestStreamContext()
        if isinstance(node, ModelRequestNode):
            async with node.stream(run.ctx) as request_stream:
                async for agent_event in request_stream:
                    async for msg in _handle_model_request_event(stream_ctx, agent_event):
                        yield msg

                if stream_ctx.part_end:  # pragma: no branch
                    yield stream_ctx.part_end
                    stream_ctx.part_end = None
                if stream_ctx.thinking:
                    yield ThinkingEndEvent(
                        type=EventType.THINKING_END,
                    )
                    stream_ctx.thinking = False
        elif isinstance(node, CallToolsNode):
            async with node.stream(run.ctx) as handle_stream:
                async for event in handle_stream:
                    if isinstance(event, FunctionToolResultEvent):
                        async for msg in _handle_tool_result_event(stream_ctx, event):
                            yield msg