diff --git a/backend/packages/harness/deerflow/client.py b/backend/packages/harness/deerflow/client.py index de89771bf..8a61fcb57 100644 --- a/backend/packages/harness/deerflow/client.py +++ b/backend/packages/harness/deerflow/client.py @@ -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 diff --git a/backend/tests/test_client.py b/backend/tests/test_client.py index 2a8460b1d..1a2a73365 100644 --- a/backend/tests/test_client.py +++ b/backend/tests/test_client.py @@ -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(