Can I still check messages even if the API is interrupted due to an error?

[phi-friday]

Question

I am using ag-ui, and on the frontend, I am building a chat UI with @ag-ui/client.
I have several questions about this process, but one of them concerns error handling.

When the API finishes normally, on_complete is called, and since this function takes AgentRunResult as a parameter, it’s easy to get the user’s input and the newly generated LLM messages. I use this callback to store messages in the database.

However, if an error occurs during API execution (especially when client.abortAgent is executed from the frontend), this callback is not triggered. I considered implementing a subclass of AGUIEventStream to handle errors, but I noticed that not all event objects have a result, only certain events have one, and these events appear to be generated from the graph.

Is there a way to collect the ModelRequest and ModelResponse up until just before the error occurs, even in cases of errors? If not, would I need to collect all events and manually construct ModelRequest and ModelResponse from them to achieve this?

Since a function that converts ag-ui messages into ModelMessage is already implemented(AGUIAdapter.load_messages), this isn’t a difficult task. However, I want to avoid handling event branching and collection twic, once on the frontend and again on the backend. I’m asking because there might be a better way that I’m unaware of.

Additional Context

No response

[phi-friday]

At the moment, it seems that the simplest way is to define a class that inherits from Agent and override the run method to manipulate the event_stream_handler. However, I'm not sure if this is the correct approach. (I thought it would be possible because RunContext has messages.)

like this

event_stream_handler = event_stream_handler or self.event_stream_handler
async def _event_stream_handler(
    context: RunContext[AgentDepsT], events: AsyncIterable[_messages.AgentStreamEvent]
) -> None:
    await queue.put(context.messages)
    if event_stream_handler is not None:
        await event_stream_handler(context, events)

If it were possible to pass event_stream_handler to dispatch_request, that would be best, but I don't think it's an easy task.

[phi-friday]

I created a pr in the simplest form. I hope this will be sufficient.