ag-ui-protocol / ag-ui-protocol/ag-ui

[Bug]: Suppressed internal LLM calls (emit-messages=False) permanently break TEXT_MESSAGE_START for later messages in the same run

Abierto
#2,326 0 comentarios 0 reacciones 0 asignados Ver en GitHub
bug
Lenguaje dominante
Python
Estrellas
15.9k
Forks
1.4k
Merge medio
1 d 17 h
PR fusionados (30 d)
163

Descripción

### Pre-flight Checklist

- [x] I have searched [existing issues](https://github.com/ag-ui-protocol/ag-ui/issues) and this hasn't been reported yet.
- [x] I am using the **latest** version AG-UI.

### Describe the Bug

## Summary

When a LangGraph run contains an internal LLM call whose `TEXT_MESSAGE_*`/`TOOL_CALL_*` events are suppressed (e.g. via CopilotKit's `copilotkit_customize_config(emit_messages=False)`), any *later* organically-streamed assistant message in the **same run** silently loses its `TEXT_MESSAGE_START` event. AG-UI clients reject the resulting `TEXT_MESSAGE_CONTENT` events with:

```
Cannot send 'TEXT_MESSAGE_CONTENT' event: No active text message found with ID
'lc_run--...'. Start a text message with 'TEXT_MESSAGE_START' first.
```

This silently drops the assistant's real reply text from the user's perspective (tool-call/generative-UI rendering is unaffected — only plain text messages are lost).

### Steps to Reproduce

## Reproduction

1. Build a LangGraph graph where an early node makes an LLM call wrapped in `copilotkit_customize_config(config, emit_messages=False, emit_tool_calls=False)` (e.g. a structured-output classification/routing step).
2. A later node in the *same run* makes a normal LLM call with no suppression (organic streaming, e.g. LangChain's `create_agent` "model" node).
3. Drive the graph via `add_langgraph_fastapi_endpoint`/`LangGraphAGUIAgent` and observe the raw AG-UI SSE stream for a single request: the real assistant reply is missing its `TEXT_MESSAGE_START` event — only `TEXT_MESSAGE_CONTENT`/`TEXT_MESSAGE_END` are emitted, and AG-UI clients reject them.

Reproducible Code:

This script drives ``LangGraphAgent._handle_single_event`` directly with
synthetic LangGraph ``astream_events``-shaped events -- no LangGraph graph,
LLM, or API key required. It exercises exactly two calls in one run:

1. An internal call whose events are suppressed (simulating
``copilotkit:emit-messages=False``), which streams and then ends.
2. A real call, in a different graph node, that streams normally.

Expected (correct) behavior: step 2 emits TEXT_MESSAGE_START before its
TEXT_MESSAGE_CONTENT. Buggy behavior: TEXT_MESSAGE_START is missing.

Usage:
pip install ag-ui-langgraph fastapi
python repro.py

```
import asyncio

from ag_ui.core import EventType
from ag_ui_langgraph.agent import LangGraphAgent
from ag_ui_langgraph.types import LangGraphEventTypes

def dispatch_passthrough(ev):
"""Mimics CopilotKit's LangGraphAGUIAgent._dispatch_event: suppresses
TEXT_MESSAGE_*/TOOL_CALL_* events (returns None) when the originating
LangGraph run metadata carries ``copilotkit:emit-messages=False`` --
the pattern used to hide internal, non-user-facing LLM calls (e.g. a
structured-output classification/routing step) from the AG-UI stream."""
if ev.type in (
EventType.TEXT_MESSAGE_START,
EventType.TEXT_MESSAGE_CONTENT,
EventType.TEXT_MESSAGE_END,
):
raw_metadata = (ev.raw_event or {}).get("metadata", {})
if raw_metadata.get("copilotkit:emit-messages") is False:
return None
return ev

def make_agent():
from unittest.mock import MagicMock

agent = LangGraphAgent(name="repro", graph=MagicMock())
agent.active_run = {
"id": "run-1",
"thread_id": "t1",
"reasoning_process": None,
"node_name": "agent",
"has_function_streaming": False,
"model_made_tool_call": False,
"state_reliable": True,
"streamed_messages": [],
"manually_emitted_state": None,
"schema_keys": {"input": ["messages", "tools"], "output": ["messages", "tools"], "config": [], "context": []},
}
agent._dispatch_event = dispatch_passthrough
return agent

async def stream_chunk(agent, message_id: str, content: str, *, suppress: bool):
event = {
"event": LangGraphEventTypes.OnChatModelStream.value,
"data": {"chunk": {"id": message_id, "content": content, "tool_call_chunks": []}},
"metadata": {"copilotkit:emit-messages": not suppress},
}
return [ev async for ev in agent._handle_single_event(event, {}) if ev is not None]

async def end_chat_model(agent, *, suppress: bool):
event = {
"event": LangGraphEventTypes.OnChatModelEnd.value,
"data": {},
"metadata": {"copilotkit:emit-messages": not suppress},
}
return [ev async for ev in agent._handle_single_event(event, {}) if ev is not None]

async def main() -> int:
agent = make_agent()

print("Step 1: internal (suppressed) LLM call streams + ends, node 'classify'")
for _ in agent.handle_node_change("classify"):
pass
await stream_chunk(agent, "internal-msg", "internal output", suppress=True)
await end_chat_model(agent, suppress=True)
print(f" messages_in_process['run-1'] after suppressed call: "
f"{agent.get_message_in_progress('run-1')!r}")

print("Step 2: real LLM call streams normally, node 'model' (different node, same run)")
for _ in agent.handle_node_change("model"):
pass
events = await stream_chunk(agent, "real-msg", "Hello there", suppress=False)
event_types = [e.type for e in events]
print(f" Emitted event types: {event_types}")

ok = EventType.TEXT_MESSAGE_START in event_types and EventType.TEXT_MESSAGE_CONTENT in event_types
if ok:
print("\nPASS: TEXT_MESSAGE_START was emitted for the real message. Fix is present.")
return 0
else:
print(
"\nFAIL (bug reproduced): TEXT_MESSAGE_START is missing for the real message.\n"
"AG-UI clients would reject the TEXT_MESSAGE_CONTENT that follows with:\n"
" \"Cannot send 'TEXT_MESSAGE_CONTENT' event: No active text message found "
"with ID '...'. Start a text message with 'TEXT_MESSAGE_START' first.\"\n"
"Root cause: OnChatModelEnd only clears messages_in_process[run_id] when\n"
"_dispatch_event's return value is truthy; a suppressed (None-returning)\n"
"dispatch for the internal call in Step 1 leaves it permanently poisoned."
)
return 1

if __name__ == "__main__":
raise SystemExit(asyncio.run(main()))
```

### Expected Behavior

Expected (correct): emits TEXT_MESSAGE_START before its
TEXT_MESSAGE_CONTENT.

### Environment

```text
- `copilotkit` 0.1.94
- `ag-ui-langgraph` 0.0.41 (bug also present in 0.0.42, the latest release)
```

### Screenshots

_No response_

### Logs & Errors

```shell

```

### Additional Context

## Fix

Clear `self.messages_in_process[run_id] = None` **unconditionally** in `OnChatModelEnd`, for both the `ToolCallEndEvent` and `TextMessageEndEvent` cases, regardless of `_dispatch_event`'s return value. Every LLM call — suppressed or not — should clean up its own tracking slot when it ends, so a later organically-streamed message never inherits stale state from an earlier suppressed one.

I've written, tested, and verified this fix, including a regression test that fails without the fix and passes with it (full existing suite, 397 tests, passes with no regressions):

- Fork: https://github.com/parkerroan/ag-ui
- Branch (based on `main`): https://github.com/parkerroan/ag-ui/tree/fix/text-message-start-suppressed-messages
- New regression test: `integrations/langgraph/python/tests/test_suppressed_message_does_not_poison_stream.py`

Happy to open a PR

Guía de contribución

Abrir la guía de contribución

Evaluación

Este issue todavía no se ha evaluado.

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.