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)

Open
#2,653 1 comment 0 reactions 0 assignees View on GitHub
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

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.