Address PR review feedback

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: eavanvalkenburg

python/packages/core/agent_framework/_mcp.py

@@ -13,7 +13,7 @@
 from contextlib import AsyncExitStack, _AsyncGeneratorContextManager  # type: ignore
 from datetime import timedelta
 from functools import partial
-from typing import TYPE_CHECKING, Any, Literal, TypedDict
+from typing import TYPE_CHECKING, Any, Literal, TypedDict, cast
 
 import httpx
 from anyio import ClosedResourceError
@@ -733,9 +733,9 @@ async def sampling_callback(
         for msg in params.messages:
             messages.append(_parse_message_from_mcp(msg))
         try:
-            chat_client: Any = self.client
-            response: Any = await chat_client.get_response(
-                messages,
+            chat_client = cast("SupportsChatGetResponse[Any]", self.client)
+            response = await chat_client.get_response(
+                messages=messages,
                 options={
                     "temperature": params.temperature,
                     "max_tokens": params.maxTokens,

python/packages/core/agent_framework/_middleware.py

@@ -1072,12 +1072,12 @@ def get_response(
         """Execute the chat pipeline if middleware is configured."""
         super_get_response = super().get_response  # type: ignore[misc]
         effective_client_kwargs = dict(client_kwargs) if client_kwargs is not None else {}
+        call_middleware = effective_client_kwargs.pop("middleware", [])
         context_kwargs = dict(effective_client_kwargs)
         if compaction_strategy is not None:
             context_kwargs["compaction_strategy"] = compaction_strategy
         if tokenizer is not None:
             context_kwargs["tokenizer"] = tokenizer
-        call_middleware = effective_client_kwargs.pop("middleware", [])
         pipeline = self._get_chat_middleware_pipeline(call_middleware)  # type: ignore[reportUnknownArgumentType]
         if not pipeline.has_middlewares:
             return super_get_response(  # type: ignore[no-any-return]

python/packages/core/agent_framework/_tools.py

@@ -2159,6 +2159,7 @@ def get_response(
                 options=mutable_options,
                 compaction_strategy=compaction_strategy,
                 tokenizer=tokenizer,
+                function_invocation_kwargs=function_invocation_kwargs,
                 client_kwargs=filtered_kwargs,
             )
         if not stream:

python/packages/core/tests/core/test_function_invocation_logic.py

@@ -13,6 +13,7 @@
     Content,
     Message,
     SupportsChatGetResponse,
+    chat_middleware,
     tool,
 )
 from agent_framework._compaction import (
@@ -1429,6 +1430,36 @@ def ai_func(arg1: str) -> str:
     assert len(response.messages) > 0
 
 
+async def test_function_invocation_config_enabled_false_preserves_invocation_kwargs(
+    chat_client_base: SupportsChatGetResponse,
+):
+    """Test disabled function invocation still forwards invocation kwargs downstream."""
+    captured_kwargs: dict[str, Any] = {}
+
+    @tool(name="test_function")
+    def ai_func(arg1: str) -> str:
+        return f"Processed {arg1}"
+
+    @chat_middleware
+    async def capture_middleware(context, call_next):
+        captured_kwargs.update(context.function_invocation_kwargs or {})
+        await call_next()
+
+    chat_client_base.chat_middleware = [capture_middleware]
+    chat_client_base.run_responses = [
+        ChatResponse(messages=Message(role="assistant", text="response without function calling")),
+    ]
+    chat_client_base.function_invocation_configuration["enabled"] = False
+
+    await chat_client_base.get_response(
+        [Message(role="user", text="hello")],
+        options={"tool_choice": "auto", "tools": [ai_func]},
+        function_invocation_kwargs={"tool_request_id": "tool-123"},
+    )
+
+    assert captured_kwargs == {"tool_request_id": "tool-123"}
+
+
 @pytest.mark.skip(reason="Error handling and failsafe behavior needs investigation in unified API")
 async def test_function_invocation_config_max_consecutive_errors(chat_client_base: SupportsChatGetResponse):
     """Test that max_consecutive_errors_per_request limits error retries."""

python/packages/core/tests/core/test_mcp.py

@@ -1573,6 +1573,13 @@ async def test_mcp_tool_sampling_callback_chat_client_exception():
     assert isinstance(result, types.ErrorData)
     assert result.code == types.INTERNAL_ERROR
     assert "Failed to get chat message content: Chat client error" in result.message
+    mock_chat_client.get_response.assert_awaited_once()
+    _, kwargs = mock_chat_client.get_response.await_args
+    assert kwargs["options"] == {
+        "temperature": None,
+        "max_tokens": None,
+        "stop": None,
+    }
 
 
 async def test_mcp_tool_sampling_callback_no_valid_content():
@@ -1616,6 +1623,13 @@ async def test_mcp_tool_sampling_callback_no_valid_content():
     assert isinstance(result, types.ErrorData)
     assert result.code == types.INTERNAL_ERROR
     assert "Failed to get right content types from the response." in result.message
+    mock_chat_client.get_response.assert_awaited_once()
+    _, kwargs = mock_chat_client.get_response.await_args
+    assert kwargs["options"] == {
+        "temperature": None,
+        "max_tokens": None,
+        "stop": None,
+    }
 
 
 # Test error handling in connect() method

python/packages/core/tests/core/test_middleware_with_chat.py

@@ -2,6 +2,7 @@
 
 from collections.abc import Awaitable, Callable
 from typing import Any
+from unittest.mock import patch
 
 from agent_framework import (
     Agent,
@@ -296,6 +297,35 @@ async def counting_middleware(context: ChatContext, call_next: Callable[[], Awai
         assert response3 is not None
         assert execution_count["count"] == 2  # Should be 2 now
 
+    async def test_run_level_middleware_is_not_forwarded_to_inner_client(
+        self, chat_client_base: "MockBaseChatClient"
+    ) -> None:
+        """Test that run-level middleware stays in the middleware pipeline only."""
+        observed_context_kwargs: dict[str, Any] = {}
+
+        @chat_middleware
+        async def inspecting_middleware(context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None:
+            observed_context_kwargs.update(context.kwargs)
+            await call_next()
+
+        async def fake_inner_get_response(**kwargs: Any) -> ChatResponse:
+            assert "middleware" not in kwargs
+            return ChatResponse(messages=[Message(role="assistant", text="ok")])
+
+        with patch.object(
+            chat_client_base,
+            "_inner_get_response",
+            side_effect=fake_inner_get_response,
+        ) as mock_inner_get_response:
+            response = await chat_client_base.get_response(
+                [Message(role="user", text="hello")],
+                client_kwargs={"middleware": [inspecting_middleware], "trace_id": "trace-123"},
+            )
+
+        assert response.messages[0].text == "ok"
+        assert observed_context_kwargs == {"trace_id": "trace-123"}
+        mock_inner_get_response.assert_called_once()
+
     async def test_chat_client_middleware_can_access_and_override_options(
         self, chat_client_base: "MockBaseChatClient"
     ) -> None: