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:
Sydney Runkle
2026-04-02 11:19:24 -04:00
committed by GitHub
parent ee0566d31d
commit 57b314c5d7
2 changed files with 113 additions and 2 deletions
+96 -1
View File
@@ -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"