fix(client): emit streamed tool calls once with complete args (#5408)

* fix(client): emit streamed tool calls once with complete args

When a model streams a tool call as chunks (name and id first, then
argument fragments), DeerFlowClient.stream() emitted a tool_calls event
per chunk, each parsed from that fragment alone. Consumers got
`args={}` and then a call with no name or id, and never the complete
call, because the values snapshot skipped the already-streamed id.

Hold tool calls from AIMessageChunk.tool_call_chunks and emit them once
from the values snapshot, where the message is complete.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

* test(client): pin additional_kwargs on the deferred tool_calls event

Also note why a tool-call chunk without a message id keeps the per-chunk
event.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
alanhuangyoo 2026-09-14 09:50:14 +08:00 committed by GitHub
parent 1dd48d14d2
commit 529b0e05ec
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 65 additions and 3 deletions

View File

@ -30,7 +30,7 @@ from typing import Any, Literal
from langchain.agents import create_agent
from langchain.agents.middleware import AgentMiddleware
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage, SystemMessage, ToolMessage
from langchain_core.runnables import RunnableConfig
from deerflow.agents.lead_agent.agent import _authorize_model_name, build_middlewares
@ -949,6 +949,10 @@ class DeerFlowClient:
# Cross-mode handoff: ids already streamed via LangGraph ``messages``
# mode so the ``values`` path skips re-synthesis of the same message.
streamed_ids: set[str] = set()
# AI messages whose tool calls arrived as streamed fragments. The
# arguments only parse once the message is complete, so their
# tool_calls event is emitted from the values snapshot instead.
pending_tool_call_ids: set[str] = set()
# The same message id carries identical cumulative ``usage_metadata``
# in both the final ``messages`` chunk and the values snapshot —
# count it only on whichever arrives first.
@ -1041,7 +1045,12 @@ class DeerFlowClient:
)
sent_additional_kwargs = bool(additional_kwargs_delta)
if msg_chunk.tool_calls:
# A chunk without an id can't be matched to its values
# snapshot, so it keeps the per-chunk event below.
if isinstance(msg_chunk, AIMessageChunk) and msg_chunk.tool_call_chunks and msg_id:
streamed_ids.add(msg_id)
pending_tool_call_ids.add(msg_id)
elif msg_chunk.tool_calls:
if msg_id:
streamed_ids.add(msg_id)
additional_kwargs_delta = None if sent_additional_kwargs else _unsent_additional_kwargs(msg_id, additional_kwargs)
@ -1074,7 +1083,10 @@ class DeerFlowClient:
_account_usage(msg_id, getattr(msg, "usage_metadata", None))
additional_kwargs = self._serialize_additional_kwargs(msg)
additional_kwargs_delta = _unsent_additional_kwargs(msg_id, additional_kwargs)
if additional_kwargs_delta:
if msg_id in pending_tool_call_ids and msg.tool_calls:
pending_tool_call_ids.discard(msg_id)
yield self._ai_tool_calls_event(msg_id, msg.tool_calls, additional_kwargs_delta)
elif additional_kwargs_delta:
# Metadata-only follow-up: ``messages-tuple`` has no
# dedicated attribution event, so clients should
# merge this empty-content AI event by message id

View File

@ -621,6 +621,56 @@ class TestStream:
call_kwargs = agent.stream.call_args.kwargs
assert "messages" in call_kwargs["stream_mode"]
def test_stream_emits_streamed_tool_calls_once_with_complete_args(self, client):
"""Tool-call arguments streamed in fragments are emitted once, complete.
Each chunk only parses to a partial call (``args={}``, or no name/id), so
the tool_calls event comes from the values snapshot, not from the chunks.
"""
call = {"name": "bash", "args": {"command": "ls -la"}, "id": "call-1"}
attribution = {"version": 1, "kind": "tool_batch", "shared_attribution": False, "actions": []}
assembled = AIMessage(content="", id="ai-1", tool_calls=[call], additional_kwargs={"token_usage_attribution": attribution})
agent = MagicMock()
agent.stream.return_value = iter(
[
(
"messages",
(
AIMessageChunk(
content="",
id="ai-1",
tool_call_chunks=[{"name": "bash", "args": "", "id": "call-1", "index": 0}],
),
{},
),
),
(
"messages",
(
AIMessageChunk(
content="",
id="ai-1",
tool_call_chunks=[{"name": None, "args": '{"command": "ls -la"}', "id": None, "index": 0}],
),
{},
),
),
("values", {"messages": [HumanMessage(content="hi", id="h-1"), assembled]}),
]
)
with (
patch.object(client, "_ensure_agent"),
patch.object(client, "_agent", agent),
):
events = list(client.stream("hi", thread_id="t-stream-tools"))
tool_call_events = _tool_call_events(events)
assert len(tool_call_events) == 1
assert tool_call_events[0].data["id"] == "ai-1"
assert [(tc["name"], tc["args"], tc["id"]) for tc in tool_call_events[0].data["tool_calls"]] == [("bash", {"command": "ls -la"}, "call-1")]
assert tool_call_events[0].data["additional_kwargs"] == {"token_usage_attribution": attribution}
def test_stream_emits_additional_kwargs_updates_for_streamed_ai_messages(self, client):
"""stream() emits a follow-up AI event when attribution metadata arrives via values."""
assembled = AIMessage(