Python: propagate as_tool() kwargs. Add sample for runtime context with as_tool kwargs and middleware. (#2311)

* as tool kwargs

* simplify

python/packages/core/agent_framework/_agents.py

@@ -454,13 +454,16 @@ class BaseAgent(SerializationMixin):
             # Extract the input from kwargs using the specified arg_name
             input_text = kwargs.get(arg_name, "")
 
+            # Forward all kwargs except the arg_name to support runtime context propagation
+            forwarded_kwargs = {k: v for k, v in kwargs.items() if k != arg_name}
+
             if stream_callback is None:
                 # Use non-streaming mode
-                return (await self.run(input_text)).text
+                return (await self.run(input_text, **forwarded_kwargs)).text
 
             # Use streaming mode - accumulate updates and create final response
             response_updates: list[AgentRunResponseUpdate] = []
-            async for update in self.run_stream(input_text):
+            async for update in self.run_stream(input_text, **forwarded_kwargs):
                 response_updates.append(update)
                 if is_async_callback:
                     await stream_callback(update)  # type: ignore[misc]
@@ -470,12 +473,14 @@ class BaseAgent(SerializationMixin):
             # Create final text from accumulated updates
             return AgentRunResponse.from_agent_run_response_updates(response_updates).text
 
-        return AIFunction(
+        agent_tool: AIFunction[BaseModel, str] = AIFunction(
             name=tool_name,
             description=tool_description,
             func=agent_wrapper,
             input_model=input_model,  # type: ignore
         )
+        agent_tool._forward_runtime_kwargs = True  # type: ignore
+        return agent_tool
 
     def _normalize_messages(
         self,
@@ -868,7 +873,9 @@ class ChatAgent(BaseAgent):
             user=user,
             **(additional_chat_options or {}),
         )
-        response = await self.chat_client.get_response(messages=thread_messages, chat_options=co, **kwargs)
+        # 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(messages=thread_messages, chat_options=co, **filtered_kwargs)
 
         await self._update_thread_with_type_and_conversation_id(thread, response.conversation_id)
 
@@ -1000,9 +1007,11 @@ class ChatAgent(BaseAgent):
             **(additional_chat_options or {}),
         )
 
+        # 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] = []
         async for update in self.chat_client.get_streaming_response(
-            messages=thread_messages, chat_options=co, **kwargs
+            messages=thread_messages, chat_options=co, **filtered_kwargs
         ):
             response_updates.append(update)
 

python/packages/core/agent_framework/_tools.py

@@ -627,6 +627,7 @@ class AIFunction(BaseTool, Generic[ArgsT, ReturnT]):
         self.invocation_exception_count = 0
         self._invocation_duration_histogram = _default_histogram()
         self.type: Literal["ai_function"] = "ai_function"
+        self._forward_runtime_kwargs: bool = False
 
     @property
     def declaration_only(self) -> bool:
@@ -728,11 +729,16 @@ class AIFunction(BaseTool, Generic[ArgsT, ReturnT]):
         global OBSERVABILITY_SETTINGS
         from .observability import OBSERVABILITY_SETTINGS
 
-        tool_call_id = kwargs.pop("tool_call_id", None)
+        original_kwargs = dict(kwargs)
+        tool_call_id = original_kwargs.pop("tool_call_id", None)
         if arguments is not None:
             if not isinstance(arguments, self.input_model):
                 raise TypeError(f"Expected {self.input_model.__name__}, got {type(arguments).__name__}")
             kwargs = arguments.model_dump(exclude_none=True)
+            if getattr(self, "_forward_runtime_kwargs", False) and original_kwargs:
+                kwargs.update(original_kwargs)
+        else:
+            kwargs = original_kwargs
         if not OBSERVABILITY_SETTINGS.ENABLED:  # type: ignore[name-defined]
             logger.info(f"Function name: {self.name}")
             logger.debug(f"Function arguments: {kwargs}")
@@ -1272,15 +1278,20 @@ async def _auto_invoke_function(
 
     parsed_args: dict[str, Any] = dict(function_call_content.parse_arguments() or {})
 
-    # Merge with user-supplied args; right-hand side dominates, so parsed args win on conflicts.
-    merged_args: dict[str, Any] = (custom_args or {}) | parsed_args
+    # Filter out internal framework kwargs before passing to tools.
+    runtime_kwargs: dict[str, Any] = {
+        key: value
+        for key, value in (custom_args or {}).items()
+        if key not in {"_function_middleware_pipeline", "middleware"}
+    }
     try:
-        args = tool.input_model.model_validate(merged_args)
+        args = tool.input_model.model_validate(parsed_args)
     except ValidationError as exc:
         message = "Error: Argument parsing failed."
         if config.include_detailed_errors:
             message = f"{message} Exception: {exc}"
         return FunctionResultContent(call_id=function_call_content.call_id, result=message, exception=exc)
+
     if not middleware_pipeline or (
         not hasattr(middleware_pipeline, "has_middlewares") and not middleware_pipeline.has_middlewares
     ):
@@ -1289,7 +1300,8 @@ async def _auto_invoke_function(
             function_result = await tool.invoke(
                 arguments=args,
                 tool_call_id=function_call_content.call_id,
-            )  # type: ignore[arg-type]
+                **runtime_kwargs if getattr(tool, "_forward_runtime_kwargs", False) else {},
+            )
             return FunctionResultContent(
                 call_id=function_call_content.call_id,
                 result=function_result,
@@ -1305,13 +1317,14 @@ async def _auto_invoke_function(
     middleware_context = FunctionInvocationContext(
         function=tool,
         arguments=args,
-        kwargs=custom_args or {},
+        kwargs=runtime_kwargs.copy(),
     )
 
     async def final_function_handler(context_obj: Any) -> Any:
         return await tool.invoke(
             arguments=context_obj.arguments,
             tool_call_id=function_call_content.call_id,
+            **context_obj.kwargs if getattr(tool, "_forward_runtime_kwargs", False) else {},
         )
 
     try:

python/packages/core/agent_framework/observability.py

@@ -1104,6 +1104,7 @@ def _trace_agent_run(
         if not OBSERVABILITY_SETTINGS.ENABLED:
             # If model diagnostics are not enabled, just return the completion
             return await run_func(self, messages=messages, thread=thread, **kwargs)
+        filtered_kwargs = {k: v for k, v in kwargs.items() if k != "chat_options"}
         attributes = _get_span_attributes(
             operation_name=OtelAttr.AGENT_INVOKE_OPERATION,
             provider_name=provider_name,
@@ -1112,7 +1113,7 @@ def _trace_agent_run(
             agent_description=self.description,
             thread_id=thread.service_thread_id if thread else None,
             chat_options=getattr(self, "chat_options", None),
-            **kwargs,
+            **filtered_kwargs,
         )
         with _get_span(attributes=attributes, span_name_attribute=OtelAttr.AGENT_NAME) as span:
             if OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED and messages:
@@ -1173,6 +1174,7 @@ def _trace_agent_run_stream(
 
         all_updates: list["AgentRunResponseUpdate"] = []
 
+        filtered_kwargs = {k: v for k, v in kwargs.items() if k != "chat_options"}
         attributes = _get_span_attributes(
             operation_name=OtelAttr.AGENT_INVOKE_OPERATION,
             provider_name=provider_name,
@@ -1181,7 +1183,7 @@ def _trace_agent_run_stream(
             agent_description=self.description,
             thread_id=thread.service_thread_id if thread else None,
             chat_options=getattr(self, "chat_options", None),
-            **kwargs,
+            **filtered_kwargs,
         )
         with _get_span(attributes=attributes, span_name_attribute=OtelAttr.AGENT_NAME) as span:
             if OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED and messages:

python/packages/core/tests/core/test_as_tool_kwargs_propagation.py

@@ -0,0 +1,315 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for kwargs propagation through as_tool() method."""
+
+from collections.abc import Awaitable, Callable
+from typing import Any
+
+from agent_framework import ChatAgent, ChatMessage, ChatResponse, FunctionCallContent, agent_middleware
+from agent_framework._middleware import AgentRunContext
+
+from .conftest import MockChatClient
+
+
+class TestAsToolKwargsPropagation:
+    """Test cases for kwargs propagation through as_tool() delegation."""
+
+    async def test_as_tool_forwards_runtime_kwargs(self, chat_client: MockChatClient) -> None:
+        """Test that runtime kwargs are forwarded through as_tool() to sub-agent."""
+        captured_kwargs: dict[str, Any] = {}
+
+        @agent_middleware
+        async def capture_middleware(
+            context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
+        ) -> None:
+            # Capture kwargs passed to the sub-agent
+            captured_kwargs.update(context.kwargs)
+            await next(context)
+
+        # Setup mock response
+        chat_client.responses = [
+            ChatResponse(messages=[ChatMessage(role="assistant", text="Response from sub-agent")]),
+        ]
+
+        # Create sub-agent with middleware
+        sub_agent = ChatAgent(
+            chat_client=chat_client,
+            name="sub_agent",
+            middleware=[capture_middleware],
+        )
+
+        # Create tool from sub-agent
+        tool = sub_agent.as_tool(name="delegate", arg_name="task")
+
+        # Directly invoke the tool with kwargs (simulating what happens during agent execution)
+        _ = await tool.invoke(
+            arguments=tool.input_model(task="Test delegation"),
+            api_token="secret-xyz-123",
+            user_id="user-456",
+            session_id="session-789",
+        )
+
+        # Verify kwargs were forwarded to sub-agent
+        assert "api_token" in captured_kwargs, f"Expected 'api_token' in {captured_kwargs}"
+        assert captured_kwargs["api_token"] == "secret-xyz-123"
+        assert "user_id" in captured_kwargs
+        assert captured_kwargs["user_id"] == "user-456"
+        assert "session_id" in captured_kwargs
+        assert captured_kwargs["session_id"] == "session-789"
+
+    async def test_as_tool_excludes_arg_name_from_forwarded_kwargs(self, chat_client: MockChatClient) -> None:
+        """Test that the arg_name parameter is not forwarded as a kwarg."""
+        captured_kwargs: dict[str, Any] = {}
+
+        @agent_middleware
+        async def capture_middleware(
+            context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
+        ) -> None:
+            captured_kwargs.update(context.kwargs)
+            await next(context)
+
+        # Setup mock response
+        chat_client.responses = [
+            ChatResponse(messages=[ChatMessage(role="assistant", text="Response from sub-agent")]),
+        ]
+
+        sub_agent = ChatAgent(
+            chat_client=chat_client,
+            name="sub_agent",
+            middleware=[capture_middleware],
+        )
+
+        tool = sub_agent.as_tool(arg_name="custom_task")
+
+        # Invoke tool with both the arg_name field and additional kwargs
+        await tool.invoke(
+            arguments=tool.input_model(custom_task="Test task"),
+            api_token="token-123",
+            custom_task="should_be_excluded",  # This should be filtered out
+        )
+
+        # The arg_name ("custom_task") should NOT be in the forwarded kwargs
+        assert "custom_task" not in captured_kwargs
+        # But other kwargs should be present
+        assert "api_token" in captured_kwargs
+        assert captured_kwargs["api_token"] == "token-123"
+
+    async def test_as_tool_nested_delegation_propagates_kwargs(self, chat_client: MockChatClient) -> None:
+        """Test that kwargs propagate through multiple levels of delegation (A → B → C)."""
+        captured_kwargs_list: list[dict[str, Any]] = []
+
+        @agent_middleware
+        async def capture_middleware(
+            context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
+        ) -> None:
+            # Capture kwargs at each level
+            captured_kwargs_list.append(dict(context.kwargs))
+            await next(context)
+
+        # Setup mock responses to trigger nested tool invocation: B calls tool C, then completes.
+        chat_client.responses = [
+            ChatResponse(
+                messages=[
+                    ChatMessage(
+                        role="assistant",
+                        contents=[
+                            FunctionCallContent(
+                                call_id="call_c_1",
+                                name="call_c",
+                                arguments='{"task": "Please execute agent_c"}',
+                            )
+                        ],
+                    )
+                ]
+            ),
+            ChatResponse(messages=[ChatMessage(role="assistant", text="Response from agent_c")]),
+            ChatResponse(messages=[ChatMessage(role="assistant", text="Response from agent_b")]),
+        ]
+
+        # Create agent C (bottom level)
+        agent_c = ChatAgent(
+            chat_client=chat_client,
+            name="agent_c",
+            middleware=[capture_middleware],
+        )
+
+        # Create agent B (middle level) - delegates to C
+        agent_b = ChatAgent(
+            chat_client=chat_client,
+            name="agent_b",
+            tools=[agent_c.as_tool(name="call_c")],
+            middleware=[capture_middleware],
+        )
+
+        # Create tool from B for direct invocation
+        tool_b = agent_b.as_tool(name="call_b")
+
+        # Invoke tool B with kwargs - should propagate to both B and C
+        await tool_b.invoke(
+            arguments=tool_b.input_model(task="Test cascade"),
+            trace_id="trace-abc-123",
+            tenant_id="tenant-xyz",
+        )
+
+        # Verify both levels received the kwargs
+        # We should have 2 captures: one from B, one from C
+        assert len(captured_kwargs_list) >= 2
+        for kwargs_dict in captured_kwargs_list:
+            assert kwargs_dict.get("trace_id") == "trace-abc-123"
+            assert kwargs_dict.get("tenant_id") == "tenant-xyz"
+
+    async def test_as_tool_streaming_mode_forwards_kwargs(self, chat_client: MockChatClient) -> None:
+        """Test that kwargs are forwarded in streaming mode."""
+        captured_kwargs: dict[str, Any] = {}
+
+        @agent_middleware
+        async def capture_middleware(
+            context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
+        ) -> None:
+            captured_kwargs.update(context.kwargs)
+            await next(context)
+
+        # Setup mock streaming responses
+        from agent_framework import ChatResponseUpdate, TextContent
+
+        chat_client.streaming_responses = [
+            [ChatResponseUpdate(text=TextContent(text="Streaming response"), role="assistant")],
+        ]
+
+        sub_agent = ChatAgent(
+            chat_client=chat_client,
+            name="sub_agent",
+            middleware=[capture_middleware],
+        )
+
+        captured_updates: list[Any] = []
+
+        async def stream_callback(update: Any) -> None:
+            captured_updates.append(update)
+
+        tool = sub_agent.as_tool(stream_callback=stream_callback)
+
+        # Invoke tool with kwargs while streaming callback is active
+        await tool.invoke(
+            arguments=tool.input_model(task="Test streaming"),
+            api_key="streaming-key-999",
+        )
+
+        # Verify kwargs were forwarded even in streaming mode
+        assert "api_key" in captured_kwargs
+        assert captured_kwargs["api_key"] == "streaming-key-999"
+        assert len(captured_updates) == 1
+
+    async def test_as_tool_empty_kwargs_still_works(self, chat_client: MockChatClient) -> None:
+        """Test that as_tool works correctly when no extra kwargs are provided."""
+        # Setup mock response
+        chat_client.responses = [
+            ChatResponse(messages=[ChatMessage(role="assistant", text="Response from agent")]),
+        ]
+
+        sub_agent = ChatAgent(
+            chat_client=chat_client,
+            name="sub_agent",
+        )
+
+        tool = sub_agent.as_tool()
+
+        # Invoke without any extra kwargs - should work without errors
+        result = await tool.invoke(arguments=tool.input_model(task="Simple task"))
+
+        # Verify tool executed successfully
+        assert result is not None
+
+    async def test_as_tool_kwargs_with_chat_options(self, chat_client: MockChatClient) -> None:
+        """Test that kwargs including chat_options are properly forwarded."""
+        captured_kwargs: dict[str, Any] = {}
+
+        @agent_middleware
+        async def capture_middleware(
+            context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
+        ) -> None:
+            captured_kwargs.update(context.kwargs)
+            await next(context)
+
+        # Setup mock response
+        chat_client.responses = [
+            ChatResponse(messages=[ChatMessage(role="assistant", text="Response with options")]),
+        ]
+
+        sub_agent = ChatAgent(
+            chat_client=chat_client,
+            name="sub_agent",
+            middleware=[capture_middleware],
+        )
+
+        tool = sub_agent.as_tool()
+
+        # Invoke with various kwargs
+        await tool.invoke(
+            arguments=tool.input_model(task="Test with options"),
+            temperature=0.8,
+            max_tokens=500,
+            custom_param="custom_value",
+        )
+
+        # Verify all kwargs were forwarded
+        assert "temperature" in captured_kwargs
+        assert captured_kwargs["temperature"] == 0.8
+        assert "max_tokens" in captured_kwargs
+        assert captured_kwargs["max_tokens"] == 500
+        assert "custom_param" in captured_kwargs
+        assert captured_kwargs["custom_param"] == "custom_value"
+
+    async def test_as_tool_kwargs_isolated_per_invocation(self, chat_client: MockChatClient) -> None:
+        """Test that kwargs are isolated per invocation and don't leak between calls."""
+        first_call_kwargs: dict[str, Any] = {}
+        second_call_kwargs: dict[str, Any] = {}
+        call_count = 0
+
+        @agent_middleware
+        async def capture_middleware(
+            context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
+        ) -> None:
+            nonlocal call_count
+            call_count += 1
+            if call_count == 1:
+                first_call_kwargs.update(context.kwargs)
+            elif call_count == 2:
+                second_call_kwargs.update(context.kwargs)
+            await next(context)
+
+        # Setup mock responses for both calls
+        chat_client.responses = [
+            ChatResponse(messages=[ChatMessage(role="assistant", text="First response")]),
+            ChatResponse(messages=[ChatMessage(role="assistant", text="Second response")]),
+        ]
+
+        sub_agent = ChatAgent(
+            chat_client=chat_client,
+            name="sub_agent",
+            middleware=[capture_middleware],
+        )
+
+        tool = sub_agent.as_tool()
+
+        # First call with specific kwargs
+        await tool.invoke(
+            arguments=tool.input_model(task="First task"),
+            session_id="session-1",
+            api_token="token-1",
+        )
+
+        # Second call with different kwargs
+        await tool.invoke(
+            arguments=tool.input_model(task="Second task"),
+            session_id="session-2",
+            api_token="token-2",
+        )
+
+        # Verify first call had its own kwargs
+        assert first_call_kwargs.get("session_id") == "session-1"
+        assert first_call_kwargs.get("api_token") == "token-1"
+
+        # Verify second call had its own kwargs (not leaked from first)
+        assert second_call_kwargs.get("session_id") == "session-2"
+        assert second_call_kwargs.get("api_token") == "token-2"

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

@@ -321,6 +321,26 @@ async def test_ai_function_invoke_telemetry_sensitive_disabled(span_exporter: In
     assert attributes[OtelAttr.TOOL_CALL_ID] == "test_call_id"
 
 
+async def test_ai_function_invoke_ignores_additional_kwargs() -> None:
+    """Ensure ai_function tools drop unknown kwargs when invoked with validated arguments."""
+
+    @ai_function
+    async def simple_tool(message: str) -> str:
+        """Echo tool."""
+        return message.upper()
+
+    args = simple_tool.input_model(message="hello world")
+
+    # These kwargs simulate runtime context passed through function invocation.
+    result = await simple_tool.invoke(
+        arguments=args,
+        api_token="secret-token",
+        chat_options={"model_id": "dummy"},
+    )
+
+    assert result == "HELLO WORLD"
+
+
 async def test_ai_function_invoke_telemetry_with_pydantic_args(span_exporter: InMemorySpanExporter):
     """Test the ai_function invoke method with Pydantic model arguments."""