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

Vertex AI Gemini streaming does not work with LangGraphAgent - tokens buffered until completion

Open
#724 3 comments 0 reactions 0 assignees View on GitHub
bug Integration
Dominant language
Python
Stars
15.9k
Forks
1.4k
Avg merge
1d 17h
Merged PRs (30d)
163

Description

### Environment
- **ag-ui-langgraph**: 0.0.22
- **langchain-google-vertexai**: 3.0.0
- **langgraph**: 1.0.3+
- **Model**: Gemini 2.5 Pro via Vertex AI
- **Python**: 3.12

### Issue
Token-level streaming does not work when using `ChatVertexAI` (Gemini) with `LangGraphAgent`. All tokens are buffered and emitted only after the LLM completes generation, despite:
- Setting `streaming=True` on the LLM
- Using `await llm.ainvoke()` in the agent node
- Using `add_langgraph_fastapi_endpoint`

### Expected Behavior
Tokens should stream progressively as they are generated by the LLM, similar to how OpenAI models work with the same setup.

### Actual Behavior
- All tokens arrive in a single batch after ~3-6 seconds
- No `TEXT_MESSAGE_CHUNK` or `TEXT_MESSAGE_CONTENT` events are emitted during generation
- Only `MESSAGES_SNAPSHOT` is emitted after completion with the full response

### Reproduction

**Graph setup:**
```python
from langchain_google_vertexai import ChatVertexAI
from langgraph.graph import StateGraph, MessagesState

llm = ChatVertexAI(
model="gemini-2.5-pro",
streaming=True,
temperature=0.1
)

async def agent_node(state: MessagesState):
response = await llm.ainvoke(state["messages"])
return {"messages": [response]}

workflow = StateGraph(MessagesState)
workflow.add_node("agent", agent_node)
workflow.set_entry_point("agent")
workflow.add_edge("agent", END)

graph = workflow.compile(checkpointer=MemorySaver())
```

**Endpoint setup:**
```python
from ag_ui_langgraph import LangGraphAgent, add_langgraph_fastapi_endpoint

agent = LangGraphAgent(name="test_agent", graph=graph)
add_langgraph_fastapi_endpoint(app, agent, "/agent")
```

### Comparison with OpenAI
The **exact same code** works perfectly with OpenAI:
```python
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model="gpt-4o-mini", streaming=True)
```
With OpenAI, tokens stream progressively and `TEXT_MESSAGE_CHUNK` events are emitted in real-time.

### Investigation
1. Direct testing with `llm.astream()` confirms Gemini **does** stream tokens properly
2. The issue appears to be with how `LangGraphAgent` uses `astream_events()` internally
3. Related to LangGraph issue: https://github.com/langchain-ai/langgraph/issues/4718 (streaming from subgraphs)
4. Custom events emitted via `adispatch_custom_event()` inside the node ARE captured, but not converted to text message events

### Questions
1. Is Vertex AI Gemini streaming officially supported with AG-UI + LangGraph?
2. Should custom events be automatically converted to text message events?
3. Is there a recommended way to enable token-level streaming with Vertex AI?
4. Should we manually intercept custom events and convert them, or is there a better approach?

Thank you!

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.