From 02587758c9085c6e51a6a84312b7340c9d860373 Mon Sep 17 00:00:00 2001 From: Eduard Stanculet Date: Tue, 23 Jun 2026 10:41:38 +0200 Subject: [PATCH] fix: forward callbacks/config through static-args tool wrapper The static-args wrapper replaced a tool's coroutine/func with a **kwargs-only callable, which hid the original signature from StructuredTool._arun/_run. Those methods inspect the callable signature to decide whether to forward the child `callbacks` and the runnable `config`, so a wrapped tool that declares either stopped receiving them. Use functools.wraps so inspect.signature sees through to the original callable (restoring detection), and forward the injected callbacks/config as-is, bypassing apply_static_args whose serialization sanitization would mangle the callback manager / config objects. Static-arg injection still applies only to the model-supplied args. Adds a test invoking a config/callbacks-declaring coroutine through ainvoke and asserting both arrive alongside the injected static value. Co-Authored-By: Claude Opus 4.8 --- .../agent/tools/static_args.py | 42 +++++++++++++++++-- tests/agent/tools/test_wrap_static_args.py | 39 +++++++++++++++++ 2 files changed, 77 insertions(+), 4 deletions(-) diff --git a/src/uipath_langchain/agent/tools/static_args.py b/src/uipath_langchain/agent/tools/static_args.py index 3b2a8a47f..1312633ee 100644 --- a/src/uipath_langchain/agent/tools/static_args.py +++ b/src/uipath_langchain/agent/tools/static_args.py @@ -1,13 +1,16 @@ """Handles static arguments for tool calls.""" import copy +import functools import logging import re -from typing import Any, Iterator, Mapping, Sequence, TypeVar +from inspect import signature +from typing import Any, Callable, Iterator, Mapping, Sequence, TypeVar from jsonpath_ng import parse # type: ignore[import-untyped] from langchain_core.messages import ToolCall from langchain_core.tools import BaseTool, StructuredTool +from langchain_core.tools.base import _get_runnable_config_param from pydantic import BaseModel from typing_extensions import deprecated from uipath.agent.models.agent import ( @@ -346,6 +349,24 @@ def apply_to_response(self, tool_calls: list[ToolCall]) -> None: tool_call["args"] = apply_static_args(static_values, tool_call["args"]) +def _runtime_passthrough_keys(target: Callable[..., Any]) -> set[str]: + """Keyword names LangChain injects by inspecting a tool callable's signature. + + ``StructuredTool._arun`` / ``_run`` forward the child ``callbacks`` and the + runnable ``config`` into the callable only when its signature declares them. + These are runtime objects (a callback manager, a ``RunnableConfig``) that + must be forwarded untouched — never routed through ``apply_static_args``, + whose serialization sanitization would mangle them. + """ + keys: set[str] = set() + if "callbacks" in signature(target).parameters: + keys.add("callbacks") + config_param = _get_runnable_config_param(target) + if config_param: + keys.add(config_param) + return keys + + def wrap_tool_with_static_args(tool: BaseTool) -> BaseTool: """Wrap a tool so its *static* parameters are hidden from the model and injected just before the underlying tool runs. @@ -443,18 +464,31 @@ def wrap_tool_with_static_args(tool: BaseTool) -> BaseTool: coroutine = tool.coroutine if coroutine is not None: _coroutine = coroutine - + _co_passthrough = _runtime_passthrough_keys(_coroutine) + + # ``functools.wraps`` lets ``inspect.signature`` (used by StructuredTool + # to detect ``callbacks`` / ``config``) see through to the original + # callable, so LangChain's runtime injection keeps working after + # wrapping. Those injected kwargs are forwarded as-is, bypassing + # ``apply_static_args`` which only applies to the model-supplied args. + @functools.wraps(_coroutine) async def _acall(**kwargs: Any) -> Any: - return await _coroutine(**apply_static_args(inject_values, kwargs)) + passthrough = {k: kwargs.pop(k) for k in _co_passthrough if k in kwargs} + return await _coroutine( + **apply_static_args(inject_values, kwargs), **passthrough + ) update["coroutine"] = _acall func = tool.func if func is not None: _func = func + _fn_passthrough = _runtime_passthrough_keys(_func) + @functools.wraps(_func) def _call(**kwargs: Any) -> Any: - return _func(**apply_static_args(inject_values, kwargs)) + passthrough = {k: kwargs.pop(k) for k in _fn_passthrough if k in kwargs} + return _func(**apply_static_args(inject_values, kwargs), **passthrough) update["func"] = _call diff --git a/tests/agent/tools/test_wrap_static_args.py b/tests/agent/tools/test_wrap_static_args.py index f6a288677..5efb8a03a 100644 --- a/tests/agent/tools/test_wrap_static_args.py +++ b/tests/agent/tools/test_wrap_static_args.py @@ -8,6 +8,7 @@ from typing import Any +from langchain_core.runnables import RunnableConfig from pydantic import BaseModel, Field from uipath.agent.models.agent import ( AgentToolArgumentArgumentProperties, @@ -60,6 +61,44 @@ def _argument(path: str) -> AgentToolArgumentArgumentProperties: return AgentToolArgumentArgumentProperties(argument_path=path, is_sensitive=False) +async def test_wrapper_forwards_config_and_callbacks_to_coroutine() -> None: + """LangChain forwards ``config``/``callbacks`` into a tool coroutine only + when it can detect them in the signature. The wrapper stays transparent + (via functools.wraps) so a coroutine that declares them still receives them, + untouched by the static-arg merge.""" + received: dict[str, Any] = {} + + async def tool_fn( + *, + host: str, + port: int, + api_key: str, + config: RunnableConfig, + callbacks: Any = None, + ) -> str: + received.update( + host=host, port=port, api_key=api_key, config=config, callbacks=callbacks + ) + return "ok" + + tool = StructuredToolWithArgumentProperties( + name="rec", + description="A recording tool", + args_schema=_ToolInput, + coroutine=tool_fn, + output_type=None, + argument_properties={"$['host']": _static("localhost")}, + ) + wrapped = wrap_tool_with_static_args(tool) + + await wrapped.ainvoke({"port": 9090, "api_key": "k"}) + + assert received["host"] == "localhost" # static arg still injected + assert received["port"] == 9090 + assert received["config"] is not None # runnable config forwarded as-is + assert received["callbacks"] is not None # child callbacks forwarded + + async def test_static_arg_hidden_from_schema_and_injected_on_call() -> None: props: dict[str, AgentToolArgumentProperties] = {"$['host']": _static("localhost")} tool, received = _recording_tool(props)