mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-08 02:37:52 +02:00
fix: tool node injection bug (#7391)
ensuring that injected args can only be injected by LC code, not LLMs :) this does not appear to be an actual security concern because it falls outside of documented usage, the tool node doesn't inject arbitrary args
This commit is contained in:
@@ -21,7 +21,7 @@ from langchain_core.messages import (
|
||||
ToolMessage,
|
||||
)
|
||||
from langchain_core.runnables.config import RunnableConfig
|
||||
from langchain_core.tools import BaseTool, ToolException
|
||||
from langchain_core.tools import BaseTool, InjectedToolArg, ToolException
|
||||
from langchain_core.tools import tool as dec_tool
|
||||
from langgraph.config import get_stream_writer
|
||||
from langgraph.errors import GraphBubbleUp, GraphInterrupt
|
||||
@@ -2008,3 +2008,98 @@ async def test_tool_node_inject_runtime_dynamic_tool_via_wrap_tool_call_async()
|
||||
tool_message = result["messages"][-1]
|
||||
assert tool_message.content == "dynamic: x=42, tool_call_id=call_dynamic_2"
|
||||
assert tool_message.tool_call_id == "call_dynamic_2"
|
||||
|
||||
|
||||
# --- InjectedToolArg security tests ---
|
||||
|
||||
|
||||
def test_tool_node_strips_plain_injected_tool_arg() -> None:
|
||||
"""Plain InjectedToolArg values supplied by the LLM should be stripped."""
|
||||
|
||||
@dec_tool
|
||||
def read_secret(
|
||||
query: str,
|
||||
auth: Annotated[dict, InjectedToolArg()],
|
||||
) -> str:
|
||||
"""Return secret data based on auth role."""
|
||||
if auth.get("role") == "admin":
|
||||
return "ADMIN_SECRET"
|
||||
return "PUBLIC_DATA"
|
||||
|
||||
node = ToolNode([read_secret], handle_tool_errors=True)
|
||||
|
||||
# LLM tries to supply the hidden 'auth' field
|
||||
tool_call = {
|
||||
"name": "read_secret",
|
||||
"args": {"query": "hello", "auth": {"role": "admin"}},
|
||||
"id": "call-1",
|
||||
"type": "tool_call",
|
||||
}
|
||||
msg = AIMessage("", tool_calls=[tool_call])
|
||||
result = node.invoke({"messages": [msg]}, config=_create_config_with_runtime())
|
||||
tool_message = result["messages"][-1]
|
||||
# auth should have been stripped, so tool should fail (missing required arg)
|
||||
assert "ADMIN_SECRET" not in tool_message.content
|
||||
|
||||
|
||||
def test_tool_node_strips_custom_injected_tool_arg_subclass() -> None:
|
||||
"""Custom InjectedToolArg subclasses should also be stripped."""
|
||||
|
||||
class InjectedAuth(InjectedToolArg):
|
||||
pass
|
||||
|
||||
@dec_tool
|
||||
def read_secret(
|
||||
query: str,
|
||||
auth: Annotated[dict, InjectedAuth()],
|
||||
) -> str:
|
||||
"""Return secret data based on auth role."""
|
||||
if auth.get("role") == "admin":
|
||||
return "ADMIN_SECRET"
|
||||
return "PUBLIC_DATA"
|
||||
|
||||
node = ToolNode([read_secret], handle_tool_errors=True)
|
||||
|
||||
tool_call = {
|
||||
"name": "read_secret",
|
||||
"args": {"query": "hello", "auth": {"role": "admin"}},
|
||||
"id": "call-1",
|
||||
"type": "tool_call",
|
||||
}
|
||||
msg = AIMessage("", tool_calls=[tool_call])
|
||||
result = node.invoke({"messages": [msg]}, config=_create_config_with_runtime())
|
||||
tool_message = result["messages"][-1]
|
||||
assert "ADMIN_SECRET" not in tool_message.content
|
||||
|
||||
|
||||
def test_tool_node_injected_state_overwrites_llm_value() -> None:
|
||||
"""InjectedState should use graph state, not LLM-supplied values."""
|
||||
|
||||
@dec_tool
|
||||
def read_secret(
|
||||
query: str,
|
||||
auth: Annotated[dict, InjectedState("auth")],
|
||||
) -> str:
|
||||
"""Return secret data based on auth from graph state."""
|
||||
if auth.get("role") == "admin":
|
||||
return "ADMIN_SECRET"
|
||||
return "PUBLIC_DATA"
|
||||
|
||||
node = ToolNode([read_secret])
|
||||
|
||||
# LLM tries to supply auth as admin
|
||||
tool_call = {
|
||||
"name": "read_secret",
|
||||
"args": {"query": "hello", "auth": {"role": "admin"}},
|
||||
"id": "call-1",
|
||||
"type": "tool_call",
|
||||
}
|
||||
msg = AIMessage("", tool_calls=[tool_call])
|
||||
|
||||
# Graph state has auth as viewer
|
||||
result = node.invoke(
|
||||
{"messages": [msg], "auth": {"role": "viewer"}},
|
||||
config=_create_config_with_runtime(),
|
||||
)
|
||||
tool_message = result["messages"][-1]
|
||||
assert tool_message.content == "PUBLIC_DATA"
|
||||
|
||||
Reference in New Issue
Block a user