Opt-in tool_call_id injection into tool execution

README.md

@@ -295,6 +295,10 @@ We will define a single tool that prints a story. Tools may optionally take a fi
 exposed to the agent as a parameter but will be injected by the environment
 (if part of the function signature).
 
+A tool may also declare a `tool_call_id: str` argument. The environment injects the ID of
+the tool call being executed, so a tool can key its work by call. Like `state`, this
+argument stays hidden from the agent.
+
 ```py
 def print_story(story: str, state: ExampleState):
     """Print a story.

src/aviary/env.py

@@ -238,18 +238,26 @@ async def _exec_tool_call(tool_call: ToolCall) -> ToolResponseMessage:
                     f" { {t.info.name for t in self.tools} }."
                 ) from exc
 
-            # we do a special convenience to make
-            # state be optional in the function signature
+            # we do a special convenience to make state and tool_call_id be
+            # optional in the function signature: state is dropped when undeclared,
+            # tool_call_id is injected when declared. Tool.from_function keeps both
+            # out of the LLM-facing schema, and our id wins over an LLM-emitted one
+            fn_parameters = inspect.signature(tool._tool_fn).parameters
             need_to_filter = (
                 "state" in function_kwargs
-                and "state" not in inspect.signature(tool._tool_fn).parameters
+                and "state" not in fn_parameters
                 and not hasattr(tool._tool_fn, "requires_state")
             )
             filtered_kwargs = (
                 {k: v for k, v in function_kwargs.items() if k != "state"}
                 if need_to_filter
                 else function_kwargs
             )
+            tool_call_args = tool_call.function.arguments | (
+                {"tool_call_id": tool_call.id}
+                if "tool_call_id" in fn_parameters
+                else {}
+            )
 
             concurrency_context = (
                 concurrency_lock.read_lock()
@@ -262,18 +270,14 @@ async def _exec_tool_call(tool_call: ToolCall) -> ToolResponseMessage:
                 async with concurrency_context:
                     if is_coroutine_callable(tool._tool_fn):
                         content = await maybe_wait_for(
-                            tool._tool_fn(
-                                **tool_call.function.arguments, **filtered_kwargs
-                            ),
+                            tool._tool_fn(**tool_call_args, **filtered_kwargs),
                             exec_timeout,
                         )
                     else:
                         # If the function is synchronous, run on a thread
                         content = await maybe_wait_for(
                             asyncio.to_thread(
-                                tool._tool_fn,
-                                **tool_call.function.arguments,
-                                **filtered_kwargs,
+                                tool._tool_fn, **tool_call_args, **filtered_kwargs
                             ),
                             exec_timeout,
                         )

src/aviary/tools/argref.py

@@ -118,7 +118,7 @@ def argref_by_name(  # noqa: PLR0915
         >>> # Equivalent to my_func(state.refs["a"], state.refs["b"])
         >>> wrapped_fxn("a", "b", state=state)  # doctest: +SKIP
     """
-    args_to_skip = (args_to_skip or set()) | {"state", "return"}
+    args_to_skip = (args_to_skip or set()) | {"state", "tool_call_id", "return"}
 
     def decorator(func):  # noqa: PLR0915
         def get_call_args(*args, **kwargs):

src/aviary/tools/base.py

@@ -415,8 +415,8 @@ def from_function(
         required: dict[str, bool] = {}
         annotations = function.__annotations__
         for pname, parameter in inspect.signature(function).parameters.items():
-            if pname == "state":
-                # NOTE: ToolRequestMessage passes state for us, not the LLM
+            if pname in {"state", "tool_call_id"}:
+                # NOTE: exec_tool_calls injects state and tool_call_id for us, not the LLM
                 continue
             d = next(
                 (

tests/test_tools.py

@@ -19,6 +19,7 @@
 from aviary.core import (
     INVALID_TOOL_NAME,
     DummyEnv,
+    DummyEnvState,
     Environment,
     FunctionInfo,
     Message,
@@ -702,6 +703,106 @@ def get_todo_list_no_args():
             new_messages = await dummy_env.exec_tool_calls(action)
             assert new_messages[0].content == "Go for a walk"
 
+    @pytest.mark.asyncio
+    async def test_tool_call_id_injection(
+        self, dummy_env: DummyEnv, subtests: pytest.Subtests
+    ) -> None:
+        # NOTE: tool_call_id is left out of every docstring below, to confirm
+        # from_function doesn't demand a description for an injected parameter
+        async def remember(x: int, tool_call_id: str) -> str:  # noqa: D417
+            """Remember a number.
+
+            Args:
+                x: Number to remember.
+            """
+            await asyncio.sleep(0.01)  # Force concurrent calls to overlap
+            return f"{tool_call_id}:{x}"
+
+        def remember_sync(x: int, tool_call_id: str) -> str:  # noqa: D417
+            """Remember a number.
+
+            Args:
+                x: Number to remember.
+            """
+            return f"{tool_call_id}:{x}"
+
+        def remember_in_state(  # noqa: D417
+            x: int, state: DummyEnvState, tool_call_id: str
+        ) -> str:
+            """Remember a number in the state.
+
+            Args:
+                x: Number to remember.
+            """
+            state.reward = x
+            return f"{tool_call_id}:{x}"
+
+        class Rememberer:
+            async def remember_method(self, x: int, tool_call_id: str) -> str:  # noqa: D417
+                """Remember a number.
+
+                Args:
+                    x: Number to remember.
+                """
+                return f"{tool_call_id}:{x}"
+
+        tools = {
+            fn.__name__: Tool.from_function(fn)
+            for fn in (
+                remember,
+                remember_sync,
+                remember_in_state,
+                Rememberer().remember_method,
+            )
+        }
+
+        with subtests.test("injected parameters are absent from the schema"):
+            for name, tool in tools.items():
+                params = tool.info.parameters
+                assert params is not None
+                # tool_call_id is hidden just like state, so the LLM only sees x
+                assert set(params.properties) == {"x"}, name
+                assert params.required == ["x"], name
+
+        # state is passed for every call below, exercising the filtering of it out of
+        # the signatures that don't declare it
+        await dummy_env.reset()
+        for name, tool in tools.items():
+            with subtests.test(f"injected into {name}"):
+                dummy_env.tools = [tool]
+                action = ToolRequestMessage(
+                    tool_calls=[ToolCall.from_name(name, id=f"{name}-id", x=1)]
+                )
+                (response,) = await dummy_env.exec_tool_calls(
+                    action, state=dummy_env.state
+                )
+                assert response.content == f"{name}-id:1"
+        # remember_in_state ran above, so state arrived alongside tool_call_id
+        assert dummy_env.state.reward == 1
+
+        dummy_env.tools = [tools["remember"]]
+
+        with subtests.test("concurrent calls each see their own id"):
+            action = ToolRequestMessage(
+                tool_calls=[
+                    ToolCall.from_name("remember", id="first", x=1),
+                    ToolCall.from_name("remember", id="second", x=2),
+                ]
+            )
+            responses = await dummy_env.exec_tool_calls(action, concurrency=True)
+            assert [r.content for r in responses] == ["first:1", "second:2"]
+
+        with subtests.test("injected id beats an LLM-emitted one"):
+            action = ToolRequestMessage(
+                tool_calls=[
+                    ToolCall.from_name(
+                        "remember", id="authoritative", x=1, tool_call_id="spoofed"
+                    )
+                ]
+            )
+            (response,) = await dummy_env.exec_tool_calls(action)
+            assert response.content == "authoritative:1"
+
     @pytest.mark.asyncio
     async def test_tool_timing(self) -> None:
         sleep_time = 0.1