Python: Added custom args and thread object to ai_function kwargs (#2769)

* Added an example of using kwargs in ai_function

* Added thread object to ai_function kwargs

* Updated docs

* Small fix

* Added thread parameter filtering

python/packages/core/agent_framework/_agents.py

@@ -878,6 +878,9 @@ class ChatAgent(BaseAgent):
             user=user,
             additional_properties=merged_additional_options,  # type: ignore[arg-type]
         )
+
+        # Ensure thread is forwarded in kwargs for tool invocation
+        kwargs["thread"] = thread
         # Filter chat_options from kwargs to prevent duplicate keyword argument
         filtered_kwargs = {k: v for k, v in kwargs.items() if k != "chat_options"}
         response = await self.chat_client.get_response(
@@ -895,7 +898,12 @@ class ChatAgent(BaseAgent):
 
         # Only notify the thread of new messages if the chatResponse was successful
         # to avoid inconsistent messages state in the thread.
-        await self._notify_thread_of_new_messages(thread, input_messages, response.messages)
+        await self._notify_thread_of_new_messages(
+            thread,
+            input_messages,
+            response.messages,
+            **{k: v for k, v in kwargs.items() if k != "thread"},
+        )
         return AgentRunResponse(
             messages=response.messages,
             response_id=response.response_id,
@@ -1017,6 +1025,8 @@ class ChatAgent(BaseAgent):
             additional_properties=merged_additional_options,  # type: ignore[arg-type]
         )
 
+        # Ensure thread is forwarded in kwargs for tool invocation
+        kwargs["thread"] = thread
         # Filter chat_options from kwargs to prevent duplicate keyword argument
         filtered_kwargs = {k: v for k, v in kwargs.items() if k != "chat_options"}
         response_updates: list[ChatResponseUpdate] = []
@@ -1043,7 +1053,13 @@ class ChatAgent(BaseAgent):
 
         response = ChatResponse.from_chat_response_updates(response_updates, output_format_type=co.response_format)
         await self._update_thread_with_type_and_conversation_id(thread, response.conversation_id)
-        await self._notify_thread_of_new_messages(thread, input_messages, response.messages, **kwargs)
+
+        await self._notify_thread_of_new_messages(
+            thread,
+            input_messages,
+            response.messages,
+            **{k: v for k, v in kwargs.items() if k != "thread"},
+        )
 
     @override
     def get_new_thread(

python/packages/core/agent_framework/_tools.py

@@ -627,6 +627,12 @@ class AIFunction(BaseTool, Generic[ArgsT, ReturnT]):
         self._invocation_duration_histogram = _default_histogram()
         self.type: Literal["ai_function"] = "ai_function"
         self._forward_runtime_kwargs: bool = False
+        if self.func:
+            sig = inspect.signature(self.func)
+            for param in sig.parameters.values():
+                if param.kind == inspect.Parameter.VAR_KEYWORD:
+                    self._forward_runtime_kwargs = True
+                    break
 
     @property
     def declaration_only(self) -> bool:
@@ -915,6 +921,7 @@ def _create_input_model_from_func(func: Callable[..., Any], name: str) -> type[B
         )
         for pname, param in sig.parameters.items()
         if pname not in {"self", "cls"}
+        and param.kind not in {inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD}
     }
     return create_model(f"{name}_input", **fields)  # type: ignore[call-overload, no-any-return]
 
@@ -1744,7 +1751,9 @@ def _handle_function_calls_response(
                             break
                     _replace_approval_contents_with_results(prepped_messages, fcc_todo, approved_function_results)
 
-                response = await func(self, messages=prepped_messages, **kwargs)
+                # Filter out internal framework kwargs before passing to clients.
+                filtered_kwargs = {k: v for k, v in kwargs.items() if k != "thread"}
+                response = await func(self, messages=prepped_messages, **filtered_kwargs)
                 # if there are function calls, we will handle them first
                 function_results = {
                     it.call_id for it in response.messages[0].contents if isinstance(it, FunctionResultContent)
@@ -1833,7 +1842,10 @@ def _handle_function_calls_response(
 
             # Failsafe: give up on tools, ask model for plain answer
             kwargs["tool_choice"] = "none"
-            response = await func(self, messages=prepped_messages, **kwargs)
+
+            # Filter out internal framework kwargs before passing to clients.
+            filtered_kwargs = {k: v for k, v in kwargs.items() if k != "thread"}
+            response = await func(self, messages=prepped_messages, **filtered_kwargs)
             if fcc_messages:
                 for msg in reversed(fcc_messages):
                     response.messages.insert(0, msg)
@@ -1920,7 +1932,9 @@ def _handle_function_calls_streaming_response(
                     _replace_approval_contents_with_results(prepped_messages, fcc_todo, approved_function_results)
 
                 all_updates: list["ChatResponseUpdate"] = []
-                async for update in func(self, messages=prepped_messages, **kwargs):
+                # Filter out internal framework kwargs before passing to clients.
+                filtered_kwargs = {k: v for k, v in kwargs.items() if k != "thread"}
+                async for update in func(self, messages=prepped_messages, **filtered_kwargs):
                     all_updates.append(update)
                     yield update
 
@@ -2031,7 +2045,9 @@ def _handle_function_calls_streaming_response(
 
             # Failsafe: give up on tools, ask model for plain answer
             kwargs["tool_choice"] = "none"
-            async for update in func(self, messages=prepped_messages, **kwargs):
+            # Filter out internal framework kwargs before passing to clients.
+            filtered_kwargs = {k: v for k, v in kwargs.items() if k != "thread"}
+            async for update in func(self, messages=prepped_messages, **filtered_kwargs):
                 yield update
 
         return streaming_function_invocation_wrapper

python/packages/core/tests/core/test_agents.py

@@ -21,9 +21,11 @@ from agent_framework import (
     ChatResponse,
     Context,
     ContextProvider,
+    FunctionCallContent,
     HostedCodeInterpreterTool,
     Role,
     TextContent,
+    ai_function,
 )
 from agent_framework._mcp import MCPTool
 from agent_framework.exceptions import AgentExecutionException
@@ -595,3 +597,38 @@ async def test_chat_agent_with_local_mcp_tools(chat_client: ChatClientProtocol)
         # Test async context manager with MCP tools
         async with agent:
             pass
+
+
+async def test_agent_tool_receives_thread_in_kwargs(chat_client_base: Any) -> None:
+    """Verify tool execution receives 'thread' inside **kwargs when function is called by client."""
+
+    captured: dict[str, Any] = {}
+
+    @ai_function(name="echo_thread_info")
+    def echo_thread_info(text: str, **kwargs: Any) -> str:  # type: ignore[reportUnknownParameterType]
+        thread = kwargs.get("thread")
+        captured["has_thread"] = thread is not None
+        captured["has_message_store"] = thread.message_store is not None if isinstance(thread, AgentThread) else False
+        return f"echo: {text}"
+
+    # Make the base client emit a function call for our tool
+    chat_client_base.run_responses = [
+        ChatResponse(
+            messages=ChatMessage(
+                role="assistant",
+                contents=[FunctionCallContent(call_id="1", name="echo_thread_info", arguments='{"text": "hello"}')],
+            )
+        ),
+        ChatResponse(messages=ChatMessage(role="assistant", text="done")),
+    ]
+
+    agent = ChatAgent(
+        chat_client=chat_client_base, tools=[echo_thread_info], chat_message_store_factory=ChatMessageStore
+    )
+    thread = agent.get_new_thread()
+
+    result = await agent.run("hello", thread=thread)
+
+    assert result.text == "done"
+    assert captured.get("has_thread") is True
+    assert captured.get("has_message_store") is True

python/packages/core/tests/core/test_tools.py

@@ -1334,3 +1334,37 @@ async def test_streaming_two_functions_mixed_approval():
     assert updates[2].role == Role.ASSISTANT
     assert len(updates[2].contents) == 2
     assert all(isinstance(c, FunctionApprovalRequestContent) for c in updates[2].contents)
+
+
+async def test_ai_function_with_kwargs_injection():
+    """Test that ai_function correctly handles kwargs injection and hides them from schema."""
+
+    @ai_function
+    def tool_with_kwargs(x: int, **kwargs: Any) -> str:
+        """A tool that accepts kwargs."""
+        user_id = kwargs.get("user_id", "unknown")
+        return f"x={x}, user={user_id}"
+
+    # Verify schema does not include kwargs
+    assert tool_with_kwargs.parameters() == {
+        "properties": {"x": {"title": "X", "type": "integer"}},
+        "required": ["x"],
+        "title": "tool_with_kwargs_input",
+        "type": "object",
+    }
+
+    # Verify direct invocation works
+    assert tool_with_kwargs(1, user_id="user1") == "x=1, user=user1"
+
+    # Verify invoke works with injected args
+    result = await tool_with_kwargs.invoke(
+        arguments=tool_with_kwargs.input_model(x=5),
+        user_id="user2",
+    )
+    assert result == "x=5, user=user2"
+
+    # Verify invoke works without injected args (uses default)
+    result_default = await tool_with_kwargs.invoke(
+        arguments=tool_with_kwargs.input_model(x=10),
+    )
+    assert result_default == "x=10, user=unknown"