ag-ui-protocol / ag-ui-protocol/ag-ui
[ag-ui-langgraph] TOOL_CALL_ARGS is never emitted when a provider returns a tool call's name and args in one chunk (all Gemini models)
- Dominant language
- Python
- Stars
- 15.9k
- Forks
- 1.4k
- Avg merge
- 1d 17h
- Merged PRs (30d)
- 163
Description
## Summary
`ag_ui_langgraph` never emits `TOOL_CALL_ARGS` when the model provider returns a tool call's name and its arguments in a **single** streaming chunk. The client receives `TOOL_CALL_START` → `TOOL_CALL_END` → `TOOL_CALL_RESULT` and never learns what the tool was called with.
This affects **every Gemini model** through `langchain-google-genai` (verified on `gemini-3.5-flash`, `3.6-flash`, `3.7-flash`, `3.8-flash` and `gemini-3.1-pro-preview` — all identical), and any other provider that returns a complete function call in one chunk.
## Reproduction
No API key required — two fake chat models reproduce the two chunk shapes.
```python
import asyncio, uuid
from collections import Counter
from typing import Iterator
from ag_ui.core import RunAgentInput
from ag_ui_langgraph import LangGraphAgent
from langchain.agents import create_agent
from langchain_core.language_models import BaseChatModel
from langchain_core.messages import AIMessageChunk
from langchain_core.outputs import ChatGenerationChunk
from langchain_core.tools import tool
from langgraph.checkpoint.memory import InMemorySaver
CALL_ID, ARGS = "call_1", '{"file_path": "/a.txt"}'
@tool
def read_file(file_path: str) -> str:
"""Read a file."""
return "file contents"
class SplitArgsModel(BaseChatModel):
"""How OpenAI/Anthropic stream: the name first, then the arguments."""
@property
def _llm_type(self) -> str:
return "fake-split"
def bind_tools(self, tools, **kwargs):
return self
def tool_call_chunks(self) -> list[list[dict]]:
return [
[{"name": "read_file", "args": "", "id": CALL_ID, "index": 0, "type": "tool_call_chunk"}],
[{"name": None, "args": ARGS, "id": None, "index": 0, "type": "tool_call_chunk"}],
]
def _stream(self, messages, stop=None, run_manager=None, **kwargs) -> Iterator[ChatGenerationChunk]:
if any(getattr(m, "type", None) == "tool" for m in messages):
yield ChatGenerationChunk(message=AIMessageChunk(content="done"))
return
for chunks in self.tool_call_chunks():
yield ChatGenerationChunk(message=AIMessageChunk(content="", tool_call_chunks=chunks))
def _generate(self, messages, stop=None, run_manager=None, **kwargs):
raise NotImplementedError
class CombinedArgsModel(SplitArgsModel):
"""How Gemini streams: the name and the whole argument string in one chunk."""
@property
def _llm_type(self) -> str:
return "fake-combined"
def tool_call_chunks(self) -> list[list[dict]]:
return [[{"name": "read_file", "args": ARGS, "id": CALL_ID, "index": 0, "type": "tool_call_chunk"}]]
async def events_for(model) -> Counter:
agent = LangGraphAgent(
name="repro", graph=create_agent(model, tools=[read_file], checkpointer=InMemorySaver())
)
run_input = RunAgentInput(
thread_id=f"t-{uuid.uuid4().hex[:8]}",
run_id=f"r-{uuid.uuid4().hex[:8]}",
messages=[{"id": "m1", "role": "user", "content": "read /a.txt"}],
tools=[], context=[], forwarded_props={}, state={},
)
seen: Counter = Counter()
async for event in agent.run(run_input):
kind = getattr(event, "type", None)
seen[str(getattr(kind, "value", kind))] += 1
return seen
async def main() -> None:
for label, model in (
("name and args in SEPARATE chunks (OpenAI shape)", SplitArgsModel()),
("name and args in ONE chunk (Gemini shape)", CombinedArgsModel()),
):
seen = await events_for(model)
print(f"\n{label}")
print(f" TOOL_CALL_START : {seen.get('TOOL_CALL_START', 0)}")
print(f" TOOL_CALL_ARGS : {seen.get('TOOL_CALL_ARGS', 0)} <-- expected 1")
print(f" TOOL_CALL_END : {seen.get('TOOL_CALL_END', 0)}")
asyncio.run(main())
```
### Actual
```
name and args in SEPARATE chunks (OpenAI shape)
TOOL_CALL_START : 1
TOOL_CALL_ARGS : 1 <-- expected 1
TOOL_CALL_END : 1
name and args in ONE chunk (Gemini shape)
TOOL_CALL_START : 1
TOOL_CALL_ARGS : 0 <-- expected 1
TOOL_CALL_END : 1
```
### Expected
Both shapes should produce one `TOOL_CALL_ARGS` carrying `{"file_path": "/a.txt"}`. The chunking is a provider transport detail and should not change the emitted event sequence.
## Cause
In `ag_ui_langgraph/agent.py`, the tool-call branches are evaluated per chunk:
```python
is_tool_call_start_event = not has_current_stream and tool_call_data and tool_call_data.get("name")
is_tool_call_args_event = has_current_stream and current_stream.get("tool_call_id") and tool_call_data and tool_call_data.get("args")
is_tool_call_end_event = has_current_stream and current_stream.get("tool_call_id") and not tool_call_data
```
and the start branch ends in a bare `return`:
```python
if is_tool_call_start_event or is_tool_call_end_event or is_tool_call_args_event:
...
if should_emit_tool_calls:
yield self._dispatch_event(ToolCallStartEvent(...))
self.set_message_in_progress(...)
return # <-- args carried by this same chunk are never emitted
if is_tool_call_args_event and should_emit_tool_calls:
yield self._dispatch_event(ToolCallArgsEvent(...))
return
```
When one chunk carries both `name` and `args`, the start branch matches first and returns, so the arguments in that chunk are dropped. The next chunk has no tool-call data, which trips the end branch — hence `START → END` with no `ARGS`.
## Impact
Any AG-UI client driven by a single-chunk provider cannot display tool arguments. A UI can show *"`execute` was called"* and its result, but never the command that ran or the file that was read. The arguments are absent from the event stream entirely — `STATE_SNAPSHOT` deltas carry one message each and do not include them either, so there is no client-side recovery path.
## Suggested fix
Do not `return` after emitting `TOOL_CALL_START`; fall through and emit `TOOL_CALL_ARGS` when the same chunk also carries `args`. Roughly:
```python
if is_tool_call_start_event:
...
yield self._dispatch_event(ToolCallStartEvent(...))
self.set_message_in_progress(...)
# A provider may deliver the whole call in one chunk (Gemini). Emit its
# arguments now rather than waiting for a follow-up chunk that never comes.
if tool_call_data.get("args") and should_emit_tool_calls:
yield self._dispatch_event(
ToolCallArgsEvent(
type=EventType.TOOL_CALL_ARGS,
tool_call_id=public_call_id,
delta=tool_call_data["args"],
raw_event=event,
)
)
return
```
Happy to open a PR if the approach looks right.
## Workaround
Splitting the provider's combined chunk into the two-chunk shape before it reaches the adapter restores correct events (verified: 42 tool calls → 42 `TOOL_CALL_ARGS`). That requires wrapping the chat model's `_stream`/`_astream` per provider, which every consumer has to reimplement.
## Environment
| | |
|---|---|
| `ag-ui-langgraph` | 0.0.44 (latest) |
| `ag-ui-protocol` | 0.1.22 |
| `langchain` | 1.4.0 |
| `langchain-core` | 1.6.2 |
| `langchain-google-genai` | 4.4.0 |
| Python | 3.12 |
## Note on the provider side
[langchain-ai/langchain-google#1752](https://github.com/langchain-ai/langchain-google/issues/1752) requests incremental argument streaming for Gemini, but that flag (`stream_function_call_arguments`) is Vertex AI only — the `google-genai` SDK states it is "not supported in Gemini API" — and it is not wired in any released `langchain-google-genai` or `langchain-google-vertexai`. So the single-chunk shape is the only shape available on the Gemini Developer API, and the adapter needs to handle it regardless.
Contributor guide
Assessment
This issue has not been assessed yet.