Skip to content

fix(agui): defer RunErrorEvent emission to prevent premature disconnect #292

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
weimch wants to merge 1 commit into from
Closed

fix(agui): defer RunErrorEvent emission to prevent premature disconnect #292

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
23 changes: 23 additions & 0 deletions trpc_agent_sdk/server/ag_ui/_core/_agui_agent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1241,10 +1241,18 @@ async def _run_trpc_in_background(self,
return

# Run TRPC agent
# Track whether the last observed trpc_event was a fatal (non-recoverable) error, so
# that once the stream ends we know whether the run actually terminated due to an
# error or completed normally. We can't emit RunErrorEvent as soon as we see an error
# event mid-stream because most agent types (GraphAgent, ChainAgent, etc.) continue
# running after a sub-agent yields an error event, and RunErrorEvent is terminal per
# the AG-UI protocol - compliant clients close the connection upon receiving it.
last_event_is_error = False
async for trpc_event in runner.run_async(user_id=user_id,
session_id=input.thread_id,
new_message=new_message,
run_config=run_config):
last_event_is_error = trpc_event.is_error()
if not isinstance(trpc_event, LongRunningEvent):
# Check if custom translator should handle this event
if self._custom_event_translator and self._custom_event_translator.need_translate(trpc_event):
Expand Down Expand Up @@ -1281,6 +1289,21 @@ async def _run_trpc_in_background(self,
current_timestamp = datetime.now().timestamp()
ag_ui_event = event_translator._create_state_snapshot_event(final_state, current_timestamp)
await event_queue.put(ag_ui_event)

# The run ended - decide whether it finished normally or terminated due to an error
# based on whether the last trpc_event observed was an error. `trpc_event` still
# refers to the last event yielded by the loop above.
if last_event_is_error:
error_msg = (trpc_event.error_message or (trpc_event.custom_metadata or {}).get("error")
or "Unknown error")
logger.error("Run for thread %s ended with a fatal error as its last event: %s", input.thread_id,
error_msg)
await event_queue.put(
RunErrorEvent(
type=EventType.RUN_ERROR,
message=error_msg,
code=trpc_event.error_code or "MODEL_ERROR",
))
# Signal completion - TRPC execution is done
logger.debug("Background task sending completion signal for thread %s", input.thread_id)
await event_queue.put(None)
Expand Down
25 changes: 17 additions & 8 deletions trpc_agent_sdk/server/ag_ui/_core/_event_translator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,6 @@
from ag_ui.core import BaseEvent
from ag_ui.core import CustomEvent
from ag_ui.core import EventType
from ag_ui.core import RunErrorEvent
from ag_ui.core import StateDeltaEvent
from ag_ui.core import StateSnapshotEvent
from ag_ui.core import TextMessageContentEvent
Expand Down Expand Up @@ -202,22 +201,32 @@ async def translate(self, trpc_event: TRPCEvent, thread_id: str, run_id: str) ->
# Tool execution errors (with function_response) are recoverable: the error is already
# passed back to the LLM as a tool result, so the LLM can retry or adjust its approach.
# Only fatal errors (LLM failures, system errors) without function_response should
# emit RunErrorEvent to terminate the run.
# be surfaced as an error to the client.
#
# NOTE (deferred decision): We deliberately do NOT emit RunErrorEvent here. Per the
# AG-UI protocol, RunErrorEvent is terminal and compliant clients close the connection
# upon receiving it. However, most agent types (GraphAgent, ChainAgent, etc.) continue
# running after a sub-agent yields an error event, so terminating the connection here
# would desynchronize the client from the still-running backend. Instead, we surface the
# error as a CustomEvent, and the caller (who observes the full event stream) decides
# whether to emit RunErrorEvent or RunFinishedEvent once the run actually ends.
if trpc_event.is_error() and not function_responses:
# Fatal system/LLM error - emit RunErrorEvent to terminate the run
# Fatal system/LLM error - surface it via CustomEvent instead of terminating the run
logger.error("Fatal error (non-recoverable), error_code=%s, error_message=%s", trpc_event.error_code,
trpc_event.error_message)
# Force close any streaming message before emitting error
async for close_event in self.force_close_streaming_message():
yield close_event
error_msg = (trpc_event.error_message or (trpc_event.custom_metadata or {}).get("error")
or "Unknown error")
yield RunErrorEvent(
type=EventType.RUN_ERROR,
message=error_msg,
code=trpc_event.error_code or "MODEL_ERROR",
yield CustomEvent(
type=EventType.CUSTOM,
name="trpc_error",
value={
"code": trpc_event.error_code or "MODEL_ERROR",
"message": error_msg,
},
)
return

# Handle custom events or metadata
if trpc_event.custom_metadata:
Expand Down
Loading