agui: TEXT_MESSAGE_CONTENT events are sometimes orphaned

[jcbwndsr]

Initial Checks

• I confirm that I'm using the latest version of Pydantic AI
• I confirm that I searched for my issue in https://github.com/pydantic/pydantic-ai/issues before opening this issue

Description

Sometimes, it seems that TEXT_MESSAGE_CONTENT events are sent after a TEXT_MESSAGE_END event. This causes problems when using the agui client as receive the error error: Cannot send 'TEXT_MESSAGE_CONTENT' event: No active text message found with ID '4e3fcf08-eff9-4594-843f-45dad275a301'. Start a text message with 'TEXT_MESSAGE_START' first.

I have no idea what is causing this to happen as it exhibits some race condition behaviour. I haven't been able to reliably get this to happen.

I have managed to fix this with the following patch on our end. At the very least I haven't seen the bug since 😆 but maybe there is still some race condition where it may happen

async def patch_agent_stream(run_stream: AsyncIterator[str]) -> AsyncIterator[str]:
    """
    Fix bug in pydantic_ai where TEXT_MESSAGE_CONTENT events orphaned.
    This fixes the bug by converting orphaned TEXT_MESSAGE_CONTENT events to TEXT_MESSAGE_CHUNK events.
    """
    last_text_message_start_event: TextMessageStartEvent | None = None
    last_text_message_end_event: TextMessageEndEvent | None = None
    encoder = EventEncoder()

    async for chunk in run_stream:
        chunk_without_data = chunk.replace("data: ", "").strip()
        if '"type":"TEXT_MESSAGE_START"' in chunk_without_data:
            last_text_message_start_event = TextMessageStartEvent.model_validate_json(
                chunk_without_data
            )
            last_text_message_end_event = None
            yield chunk
        elif '"type":"TEXT_MESSAGE_END"' in chunk_without_data:
            last_text_message_end_event = TextMessageEndEvent.model_validate_json(
                chunk_without_data
            )
            last_text_message_start_event = None
            yield chunk
        elif '"type":"TEXT_MESSAGE_CONTENT"' in chunk_without_data:
            # Check if this is an orphaned TEXT_MESSAGE_CONTENT
            # (comes without TEXT_MESSAGE_START or after TEXT_MESSAGE_END)
            is_inside_message = (
                last_text_message_start_event is not None and last_text_message_end_event is None
            )

            if not is_inside_message:
                # Convert orphaned TEXT_MESSAGE_CONTENT to TEXT_MESSAGE_CHUNK
                content_event = TextMessageContentEvent.model_validate_json(chunk_without_data)
                chunk_event = TextMessageChunkEvent(
                    message_id=content_event.message_id,
                    delta=content_event.delta,
                    timestamp=content_event.timestamp,
                )
                yield encoder.encode(chunk_event)
            else:
                # Inside a message, let it through normally
                yield chunk
        else:
            yield chunk

Example Code

    agent = await get_agent()
    event_stream = patch_agent_stream(
        run_agui(
            agent,
            run_input,
            accept=accept,
        )
    )

    return StreamingResponse(
        event_stream,
        media_type=accept,
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
    )

Python, Pydantic AI & LLM client version

Python 3.11.6
Pydantic AI 1.0.15

[DouweM]

@jcbwndsr Thanks for the report, that's strange! It'd be useful to see the full event stream for the scenario where this goes wrong. What comes before the orphan TEXT_MESSAGE_CONTENT could be enlightening.

[jcbwndsr]

@DouweM I will see if I can repeat it today. Not that easy to reproduce but it was a TEXT_MESSAGE_END event sometimes with a tool call in between as well

[jcbwndsr]

@DouweM Please find the events:

data: {"type":"RUN_STARTED","threadId":"8f8638fa-7e89-43cf-b036-e1188506ceca","runId":"644014ef-3b75-4e76-85f6-b09eed6dc682"}


---
data: {"type":"THINKING_START"}


---
data: {"type":"THINKING_END"}


---
data: {"type":"TEXT_MESSAGE_START","messageId":"a741ed5a-3183-4838-acdf-1e553d54123b","role":"assistant"}


---
data: {"type":"TEXT_MESSAGE_CONTENT","messageId":"a741ed5a-3183-4838-acdf-1e553d54123b","delta":"MESSAGE_CONTENT_HERE"}


---
data: {"type":"TEXT_MESSAGE_CONTENT","messageId":"a741ed5a-3183-4838-acdf-1e553d54123b","delta":"MESSAGE_CONTENT_HERE"}


---
data: {"type":"TEXT_MESSAGE_END","messageId":"a741ed5a-3183-4838-acdf-1e553d54123b"}


---
data: {"type":"TOOL_CALL_START","toolCallId":"pyd_ai_db677ba8098045d7b83b7a5fe44b3dfe","toolCallName":"axiomatic-documents_document_to_markdown","parentMessageId":"a741ed5a-3183-4838-acdf-1e553d54123b"}


---
data: {"type":"TOOL_CALL_ARGS","toolCallId":"pyd_ai_db677ba8098045d7b83b7a5fe44b3dfe","delta":"MESSAGE_CONTENT_HERE"}


---
data: {"type":"TOOL_CALL_END","toolCallId":"pyd_ai_db677ba8098045d7b83b7a5fe44b3dfe"}


---
data: {"type":"TOOL_CALL_RESULT","messageId":"517bc6f7-e6ea-49cd-bf78-278e522e9e61","toolCallId":"pyd_ai_db677ba8098045d7b83b7a5fe44b3dfe","content":"RESULT_HERE","role":"tool"}


---
data: {"type":"STATE_SNAPSHOT","snapshot":"SNAPSHOT_HERE"}


---
data: {"type":"TEXT_MESSAGE_START","messageId":"fd884f08-0494-4470-b46a-66aa52ef2a9c","role":"assistant"}


---
data: {"type":"TEXT_MESSAGE_CONTENT","messageId":"fd884f08-0494-4470-b46a-66aa52ef2a9c","delta":"MESSAGE_CONTENT_HERE"}


---
data: {"type":"TEXT_MESSAGE_END","messageId":"fd884f08-0494-4470-b46a-66aa52ef2a9c"}


---
data: {"type":"THINKING_START"}


---
data: {"type":"TEXT_MESSAGE_CONTENT","messageId":"fd884f08-0494-4470-b46a-66aa52ef2a9c","delta":"MESSAGE_CONTENT_HERE"}


---
data: {"type":"TEXT_MESSAGE_CONTENT","messageId":"fd884f08-0494-4470-b46a-66aa52ef2a9c","delta":"MESSAGE_CONTENT_HERE"}


---
data: {"type":"RUN_ERROR","message":"Content field missing from streaming Gemini response, body:\nsdk_http_response=HttpResponse(\n  headers=<dict len=11>\n) candidates=[Candidate(\n  finish_reason=<FinishReason.MALFORMED_FUNCTION_CALL: 'MALFORMED_FUNCTION_CALL'>,\n  index=0\n)] create_time=None model_version=None prompt_feedback=None response_id=None usage_metadata=GenerateContentResponseUsageMetadata(\n  candidates_token_count=77,\n  prompt_token_count=35796,\n  prompt_tokens_details=[\n    ModalityTokenCount(\n      modality=<MediaModality.TEXT: 'TEXT'>,\n      token_count=35796\n    ),\n  ],\n  total_token_count=35873\n) automatic_function_calling_history=None parsed=None"}


---

As you can see you have two TEXT_MESSAGE_CONTENT events directly after the corresponding TEXT_MESSAGE_END.

This is the first time I see the RUN_ERROR but I have left it in any case. We are using gemini-2.5-flash

[DouweM]

@jcbwndsr Interesting, that suggests the stream receives a PartStartEvent for a ThinkingPart followed by a PartDeltaEvent for a TextPartDelta, without a PartStartEvent for a TextPart in between.

Looking at our implementation for Google, I suppose this could happen if the API returns text parts, followed by thinking parts (or a text part with a thought_signature), followed by more text parts:

pydantic-ai/pydantic_ai_slim/pydantic_ai/models/google.py

Lines 653 to 667 in 9b1913e

	if part.thought_signature:
	signature = base64.b64encode(part.thought_signature).decode('utf-8')
	yield self._parts_manager.handle_thinking_delta(
	vendor_part_id='thinking',
	signature=signature,
	provider_name=self.provider_name,
	)
	if part.text is not None:
	if part.thought:
	yield self._parts_manager.handle_thinking_delta(vendor_part_id='thinking', content=part.text)
	else:
	maybe_event = self._parts_manager.handle_text_delta(vendor_part_id='content', content=part.text)
	if maybe_event is not None: # pragma: no branch
	yield maybe_event

To confirm, it would be helpful if you were able to catch the result.all_messages() from the on_complete hook and share it here.

[dsfaccini]

Seems related to #3007, right? I can test this using the #3206 branch

@jcbwndsr could you provide full example code including imports and get_agent? I just want to make sure I'm using the same provider etc as you.

[DouweM]

@dsfaccini I don't think it's related as thoughts from Google don't come in thinking tags, but in special parts from the API.