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

Aperta
#2,326 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub
bug
Lingua principale
Python
Stelle
15.9k
Fork
1.4k
Merge medio
1g 17h
PR unite (30g)
163

Descrizione

### 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

Guida per i contributori

Apri la guida per i contributori

Valutazione

Questa issue non è ancora stata valutata.

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.