mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-22 23:52:23 +02:00
fix(langgraph): keep tool results out of v3 messages (#7838)
## Summary
- Filter `ToolMessage` from v3 `run.messages` streaming (handler +
`MessagesTransformer`) so tool results do not appear as chat text
deltas.
- Normalize `ToolCallStream.output` in `ToolCallTransformer` so live and
serialized `ToolMessage` payloads resolve to raw `content`.
- Add regression tests for message filtering and tool-output unwrapping.
<details>
<summary>Reproducible script</summary>
```python
"""Repro: v3 tool results must not leak through run.messages."""
from __future__ import annotations
import asyncio
from collections.abc import Callable, Sequence
from typing import Any
from uuid import uuid4
from deepagents import create_deep_agent
from langchain_core.callbacks import CallbackManagerForLLMRun
from langchain_core.language_models import BaseChatModel, LanguageModelInput
from langchain_core.messages import AIMessage, HumanMessage, ToolCall, ToolMessage
from langchain_core.outputs import ChatGeneration, ChatResult
from langchain_core.runnables import Runnable
from langchain_core.tools import BaseTool, tool
from pydantic import Field
TOOL_RESULT_SENTINEL = "[]"
class ScriptedChatModel(BaseChatModel):
responses: list[AIMessage] = Field(default_factory=list)
tools: Sequence[dict[str, Any] | type | Callable | BaseTool] = ()
_idx: int = 0
@property
def _llm_type(self) -> str:
return "scripted"
def _generate(
self,
messages: Sequence[Any],
stop: list[str] | None = None,
run_manager: CallbackManagerForLLMRun | None = None,
**kwargs: Any,
) -> ChatResult:
del messages, stop, run_manager, kwargs
idx = min(self._idx, len(self.responses) - 1)
self._idx += 1
return ChatResult(generations=[ChatGeneration(message=self.responses[idx])])
def bind_tools(
self,
tools: Sequence[dict[str, Any] | type | Callable | BaseTool],
*,
tool_choice: str | None = None,
**kwargs: Any,
) -> Runnable[LanguageModelInput, AIMessage]:
del tool_choice, kwargs
self.tools = tools
return self
@tool
def list_items() -> str:
"""List available items."""
return TOOL_RESULT_SENTINEL
def _tool_call_message() -> AIMessage:
return AIMessage(
content="",
tool_calls=[
ToolCall(id="call_list", name="list_items", args={}),
],
)
def _extract_text_delta(event: Any) -> str | None:
if isinstance(event, dict):
if event.get("event") != "content-block-delta":
return None
delta = event.get("delta")
if isinstance(delta, dict) and delta.get("type") == "text-delta":
text = delta.get("text")
return text if isinstance(text, str) else None
elif getattr(event, "event", None) == "content-block-delta":
delta = getattr(event, "delta", None)
if isinstance(delta, dict) and delta.get("type") == "text-delta":
text = delta.get("text")
return text if isinstance(text, str) else None
text = getattr(delta, "text", None)
return text if isinstance(text, str) else None
return None
async def _tool_output(tool_call: Any) -> Any:
output = getattr(tool_call, "output", None)
if callable(output):
return await output()
if hasattr(output, "__await__"):
return await output
return output
async def main() -> None:
model = ScriptedChatModel(
responses=[
_tool_call_message(),
AIMessage(content="No items found."),
]
)
agent = create_deep_agent(model=model, tools=[list_items])
run = await agent.astream_events(
{"messages": [HumanMessage(content="List items")]},
version="v3",
configurable={"thread_id": str(uuid4())},
recursion_limit=50,
)
async def collect_message_texts() -> list[str]:
texts: list[str] = []
async for message_stream in run.messages:
async for event in message_stream:
text = _extract_text_delta(event)
if text is not None:
texts.append(text)
return texts
async def collect_tool_outputs() -> list[Any]:
outputs: list[Any] = []
async for tool_call in run.tool_calls:
outputs.append(await _tool_output(tool_call))
return outputs
message_texts, tool_outputs, final_state = await asyncio.gather(
collect_message_texts(),
collect_tool_outputs(),
run.output(),
)
final_messages = final_state["messages"]
tool_message = next((m for m in final_messages if isinstance(m, ToolMessage)), None)
print("run.messages text deltas:", message_texts)
print("run.tool_calls outputs:", tool_outputs)
print("final state message roles:", [m.type for m in final_messages])
if TOOL_RESULT_SENTINEL in message_texts:
raise AssertionError("Tool result leaked through run.messages.")
if TOOL_RESULT_SENTINEL not in tool_outputs:
raise AssertionError("Tool output was not surfaced through run.tool_calls.")
if tool_message is None or tool_message.tool_call_id != "call_list":
raise AssertionError("Final state does not contain the expected ToolMessage.")
print("Reproduction passed: tool output stayed out of run.messages.")
if __name__ == "__main__":
asyncio.run(main())
```
</details>
<details>
<summary>Current behavior</summary>
```text
run.messages text deltas: ['[]', 'No items found.']
run.tool_calls outputs: [ToolMessage(content='[]', ...)]
final state message roles: ['human', 'ai', 'tool', 'ai']
AssertionError: Tool result leaked through run.messages.
```
</details>
<details>
<summary>Expected behavior</summary>
```text
run.messages text deltas: ['No items found.']
run.tool_calls outputs: ['[]']
final state message roles: ['human', 'ai', 'tool', 'ai']
Reproduction passed: tool output stayed out of run.messages.
```
</details>
## Test plan
- [ ] `uv run --project libs/langgraph pytest
libs/langgraph/tests/test_stream_messages_transformer.py`
- [ ] `uv run --project libs/prebuilt pytest
libs/prebuilt/tests/test_tool_call_transformer.py`
- [ ] `uv run --project libs/langgraph ruff check` (touched files)
- [ ] `uv run --project libs/prebuilt ruff check` (touched files)
Related:
[langchain-ai/langchainjs#10900](https://github.com/langchain-ai/langchainjs/pull/10900)
This commit is contained in:
@@ -10,7 +10,7 @@ from typing import (
|
||||
from uuid import UUID, uuid4
|
||||
|
||||
from langchain_core.callbacks import BaseCallbackHandler
|
||||
from langchain_core.messages import BaseMessage
|
||||
from langchain_core.messages import BaseMessage, ToolMessage
|
||||
from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, LLMResult
|
||||
from pydantic import BaseModel
|
||||
|
||||
@@ -303,6 +303,35 @@ class StreamMessagesHandlerV2(StreamMessagesHandler, _V2StreamingCallbackHandler
|
||||
super().__init__(stream, subgraphs, parent_ns=parent_ns)
|
||||
self._streamed_run_ids: set[UUID] = set()
|
||||
|
||||
def _find_and_emit_messages(self, meta: Meta, response: Any) -> None:
|
||||
"""Like the v1 handler, but skip ToolMessage from node outputs.
|
||||
|
||||
Tool results belong on the tools channel / state in v3; v2-flagged streams
|
||||
must not replay finalized ToolMessages as chat tokens (see MessagesTransformer).
|
||||
Legacy v1-only `stream_mode="messages"` still emits ToolMessages (see subgraph
|
||||
streaming tests).
|
||||
"""
|
||||
if isinstance(response, BaseMessage) and not isinstance(response, ToolMessage):
|
||||
self._emit(meta, response, dedupe=True)
|
||||
elif isinstance(response, Sequence):
|
||||
for value in response:
|
||||
if isinstance(value, BaseMessage) and not isinstance(
|
||||
value, ToolMessage
|
||||
):
|
||||
self._emit(meta, value, dedupe=True)
|
||||
else:
|
||||
for value in _state_values(response):
|
||||
if isinstance(value, BaseMessage) and not isinstance(
|
||||
value, ToolMessage
|
||||
):
|
||||
self._emit(meta, value, dedupe=True)
|
||||
elif isinstance(value, Sequence):
|
||||
for item in value:
|
||||
if isinstance(item, BaseMessage) and not isinstance(
|
||||
item, ToolMessage
|
||||
):
|
||||
self._emit(meta, item, dedupe=True)
|
||||
|
||||
def on_llm_end(
|
||||
self,
|
||||
response: LLMResult,
|
||||
|
||||
@@ -8,7 +8,7 @@ from langchain_core.language_models.chat_model_stream import (
|
||||
AsyncChatModelStream,
|
||||
ChatModelStream,
|
||||
)
|
||||
from langchain_core.messages import AIMessageChunk, BaseMessage
|
||||
from langchain_core.messages import AIMessageChunk, BaseMessage, ToolMessage
|
||||
from langchain_protocol.protocol import MessagesData
|
||||
from typing_extensions import NotRequired, TypedDict
|
||||
|
||||
@@ -203,6 +203,7 @@ class MessagesTransformer(StreamTransformer):
|
||||
# Correlate protocol events back to a ChatModelStream by run_id
|
||||
# (attached to the event's metadata by StreamMessagesHandler).
|
||||
self._by_run: dict[str, ChatModelStream] = {}
|
||||
self._ignored_runs: set[str] = set()
|
||||
self._pump_fn: Callable[[], bool] | None = None
|
||||
self._apump_fn: Callable[[], Awaitable[bool]] | None = None
|
||||
# Cached as a list once for cheap equality with the protocol
|
||||
@@ -276,8 +277,10 @@ class MessagesTransformer(StreamTransformer):
|
||||
self._route_protocol_event(
|
||||
cast("MessagesData", payload), run_id=run_id, node=node
|
||||
)
|
||||
elif isinstance(payload, BaseMessage) and not isinstance(
|
||||
payload, AIMessageChunk
|
||||
elif (
|
||||
isinstance(payload, BaseMessage)
|
||||
and not isinstance(payload, AIMessageChunk)
|
||||
and not isinstance(payload, ToolMessage)
|
||||
):
|
||||
self._route_whole_message(payload, node=node)
|
||||
# Legacy AIMessageChunk tuples (from on_llm_new_token) are ignored;
|
||||
@@ -295,6 +298,11 @@ class MessagesTransformer(StreamTransformer):
|
||||
) -> None:
|
||||
event_type = event.get("event")
|
||||
if event_type == "message-start":
|
||||
# Tool results are exposed on the tools projection and state
|
||||
# snapshots; run.messages is the chat-token projection.
|
||||
if event.get("role") == "tool":
|
||||
self._ignored_runs.add(run_id)
|
||||
return
|
||||
message_id = event.get("message_id")
|
||||
stream = self._make_stream(
|
||||
namespace=[],
|
||||
@@ -304,6 +312,9 @@ class MessagesTransformer(StreamTransformer):
|
||||
self._by_run[run_id] = stream
|
||||
self._log.push(stream)
|
||||
stream.dispatch(event)
|
||||
elif run_id in self._ignored_runs:
|
||||
if event_type == "message-finish":
|
||||
self._ignored_runs.discard(run_id)
|
||||
elif run_id in self._by_run:
|
||||
stream = self._by_run[run_id]
|
||||
stream.dispatch(event)
|
||||
@@ -319,12 +330,14 @@ class MessagesTransformer(StreamTransformer):
|
||||
def finalize(self) -> None:
|
||||
"""Clear any routing state — streams close themselves via `message-finish`."""
|
||||
self._by_run.clear()
|
||||
self._ignored_runs.clear()
|
||||
|
||||
def fail(self, err: BaseException) -> None:
|
||||
"""Propagate run error to any streams still open when the graph fails."""
|
||||
for stream in list(self._by_run.values()):
|
||||
stream.fail(err)
|
||||
self._by_run.clear()
|
||||
self._ignored_runs.clear()
|
||||
|
||||
|
||||
SubgraphStatus = Literal["started", "completed", "failed", "interrupted", "drained"]
|
||||
|
||||
@@ -12,7 +12,7 @@ from langchain_core.language_models.chat_model_stream import (
|
||||
AsyncChatModelStream,
|
||||
ChatModelStream,
|
||||
)
|
||||
from langchain_core.messages import AIMessage, AIMessageChunk
|
||||
from langchain_core.messages import AIMessage, AIMessageChunk, ToolMessage
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
@@ -213,6 +213,23 @@ class TestProtocolEventRouting:
|
||||
log.close()
|
||||
assert _unstamped(log._items) == []
|
||||
|
||||
def test_tool_role_protocol_events_are_ignored(self) -> None:
|
||||
t, log = _make_sync_transformer()
|
||||
for evt in [
|
||||
{"event": "message-start", "role": "tool", "message_id": "tool-msg-1"},
|
||||
{
|
||||
"event": "content-block-delta",
|
||||
"index": 0,
|
||||
"content_block": {"type": "text", "text": "[]"},
|
||||
},
|
||||
{"event": "message-finish", "reason": "stop"},
|
||||
]:
|
||||
t.process(_proto_event(evt, run_id="tool-run"))
|
||||
|
||||
log.close()
|
||||
assert _unstamped(log._items) == []
|
||||
assert t._ignored_runs == set()
|
||||
|
||||
def test_concurrent_streams_routed_by_run_id(self) -> None:
|
||||
t, log = _make_sync_transformer()
|
||||
life_a = _lifecycle(text="aaaa", message_id="run-a")
|
||||
@@ -273,6 +290,29 @@ class TestWholeMessageFallback:
|
||||
assert stream.done
|
||||
assert stream.output.text == "the full answer"
|
||||
|
||||
def test_whole_tool_message_is_ignored(self) -> None:
|
||||
t, log = _make_sync_transformer()
|
||||
t.process(
|
||||
{
|
||||
"type": "event",
|
||||
"method": "messages",
|
||||
"params": {
|
||||
"namespace": [],
|
||||
"timestamp": TS,
|
||||
"data": (
|
||||
ToolMessage(
|
||||
content="[]",
|
||||
id="tool-msg-1",
|
||||
tool_call_id="call_1",
|
||||
),
|
||||
{"langgraph_node": "tools"},
|
||||
),
|
||||
},
|
||||
}
|
||||
)
|
||||
log.close()
|
||||
assert _unstamped(log._items) == []
|
||||
|
||||
def test_whole_message_has_full_lifecycle(self) -> None:
|
||||
t, log = _make_sync_transformer()
|
||||
t.process(_whole_msg("full"))
|
||||
@@ -849,6 +889,23 @@ class TestStreamMessagesHandlerV2Unit:
|
||||
|
||||
assert emitted == []
|
||||
|
||||
def test_on_chain_end_does_not_emit_tool_messages(self) -> None:
|
||||
from uuid import uuid4
|
||||
|
||||
from langgraph.pregel._messages import StreamMessagesHandlerV2
|
||||
|
||||
emitted: list[Any] = []
|
||||
handler = StreamMessagesHandlerV2(emitted.append, subgraphs=False)
|
||||
run_id = uuid4()
|
||||
handler.metadata[run_id] = ((), {"langgraph_node": "tools"})
|
||||
|
||||
handler.on_chain_end(
|
||||
{"messages": [ToolMessage(content="[]", tool_call_id="call_1")]},
|
||||
run_id=run_id,
|
||||
)
|
||||
|
||||
assert emitted == []
|
||||
|
||||
def test_on_llm_end_dedupes_when_final_message_id_differs(self) -> None:
|
||||
"""A streamed v2 message should not be emitted again from the final
|
||||
AIMessage fallback when its final id does not match `message-start`."""
|
||||
|
||||
@@ -5,12 +5,42 @@ from __future__ import annotations
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Any
|
||||
|
||||
from langchain_core.messages import ToolMessage
|
||||
from langgraph.stream._types import ProtocolEvent, StreamTransformer
|
||||
from langgraph.stream.stream_channel import StreamChannel
|
||||
|
||||
from langgraph.prebuilt._tool_call_stream import ToolCallStream
|
||||
|
||||
|
||||
def _is_serialized_tool_message(value: Any) -> bool:
|
||||
"""Detect a serialized LangChain `ToolMessage` payload.
|
||||
|
||||
Example:
|
||||
{
|
||||
"lc": 1,
|
||||
"type": "constructor",
|
||||
"id": ["langchain_core", "messages", "ToolMessage"],
|
||||
"kwargs": {"content": "raw tool result", "tool_call_id": "call_1"},
|
||||
}
|
||||
"""
|
||||
return (
|
||||
isinstance(value, dict)
|
||||
and value.get("type") == "constructor"
|
||||
and isinstance(value.get("id"), list)
|
||||
and value["id"][-1] == "ToolMessage"
|
||||
)
|
||||
|
||||
|
||||
def _normalize_tool_output(output: Any) -> Any:
|
||||
if isinstance(output, ToolMessage):
|
||||
return output.content
|
||||
if _is_serialized_tool_message(output):
|
||||
kwargs = output.get("kwargs")
|
||||
if isinstance(kwargs, dict):
|
||||
return kwargs.get("content")
|
||||
return output
|
||||
|
||||
|
||||
class ToolCallTransformer(StreamTransformer):
|
||||
"""Project `tools` channel events into `ToolCallStream` handles.
|
||||
|
||||
@@ -109,7 +139,7 @@ class ToolCallTransformer(StreamTransformer):
|
||||
elif event_type == "tool-finished":
|
||||
stream = self._active.pop(tool_call_id, None)
|
||||
if stream is not None:
|
||||
stream._finish(data.get("output"))
|
||||
stream._finish(_normalize_tool_output(data.get("output")))
|
||||
elif event_type == "tool-error":
|
||||
stream = self._active.pop(tool_call_id, None)
|
||||
if stream is not None:
|
||||
|
||||
@@ -6,7 +6,7 @@ import time
|
||||
from typing import Annotated, Any
|
||||
|
||||
import pytest
|
||||
from langchain_core.messages import AIMessage
|
||||
from langchain_core.messages import AIMessage, ToolMessage
|
||||
from langchain_core.tools import tool
|
||||
from langgraph.constants import END, START
|
||||
from langgraph.graph import StateGraph
|
||||
@@ -128,6 +128,42 @@ class TestToolCallTransformerUnit:
|
||||
assert stream.error is None
|
||||
assert "tc1" not in transformer._active
|
||||
|
||||
def test_finish_unwraps_tool_message_output(self) -> None:
|
||||
mux, transformer = _mux()
|
||||
mux.push(_tool_event("tool-started", "tc1", tool_name="echo"))
|
||||
stream = transformer._active["tc1"]
|
||||
mux.push(
|
||||
_tool_event(
|
||||
"tool-finished",
|
||||
"tc1",
|
||||
output=ToolMessage(content="done", tool_call_id="tc1"),
|
||||
)
|
||||
)
|
||||
assert stream.completed is True
|
||||
assert stream.output == "done"
|
||||
|
||||
def test_finish_unwraps_serialized_tool_message_output(self) -> None:
|
||||
mux, transformer = _mux()
|
||||
mux.push(_tool_event("tool-started", "tc1", tool_name="echo"))
|
||||
stream = transformer._active["tc1"]
|
||||
mux.push(
|
||||
_tool_event(
|
||||
"tool-finished",
|
||||
"tc1",
|
||||
output={
|
||||
"lc": 1,
|
||||
"type": "constructor",
|
||||
"id": ["langchain_core", "messages", "ToolMessage"],
|
||||
"kwargs": {
|
||||
"content": "serialized done",
|
||||
"tool_call_id": "tc1",
|
||||
},
|
||||
},
|
||||
)
|
||||
)
|
||||
assert stream.completed is True
|
||||
assert stream.output == "serialized done"
|
||||
|
||||
def test_error_closes_stream(self) -> None:
|
||||
mux, transformer = _mux()
|
||||
mux.push(_tool_event("tool-started", "tc1", tool_name="boom"))
|
||||
|
||||
Reference in New Issue
Block a user