fix: forward callbacks/config through static-args tool wrapper

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
 

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)