Python: Drop post-limit streamed function calls

python/packages/core/agent_framework/_tools.py

@@ -1923,7 +1923,14 @@ def _response_has_visible_content(response: ChatResponse[Any]) -> bool:
     return False
 
 
+def _drop_function_call_contents_from_response(response: ChatResponse[Any]) -> None:
+    for message in response.messages:
+        if any(content.type == "function_call" for content in message.contents):
+            message.contents = [content for content in message.contents if content.type != "function_call"]
+
+
 def _ensure_function_invocation_limit_fallback_response(response: ChatResponse[Any]) -> ChatResponse[Any]:
+    _drop_function_call_contents_from_response(response)
     if _response_has_visible_content(response):
         return response
 
@@ -1948,6 +1955,31 @@ def _function_invocation_limit_fallback_update() -> ChatResponseUpdate:
     )
 
 
+def _drop_function_call_contents_from_update(update: ChatResponseUpdate) -> ChatResponseUpdate | None:
+    if not any(content.type == "function_call" for content in update.contents):
+        return update
+
+    update.contents = [content for content in update.contents if content.type != "function_call"]
+    if update.contents or _update_has_meaningful_metadata(update):
+        return update
+    return None
+
+
+def _update_has_meaningful_metadata(update: ChatResponseUpdate) -> bool:
+    return any((
+        update.author_name is not None,
+        update.response_id is not None,
+        update.message_id is not None,
+        update.conversation_id is not None,
+        update.model is not None,
+        update.created_at is not None,
+        update.finish_reason is not None,
+        update.continuation_token is not None,
+        bool(update.additional_properties),
+        update.raw_representation is not None,
+    ))
+
+
 def _extract_tools(
     options: dict[str, Any] | None,
 ) -> ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None:
@@ -2823,20 +2855,31 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]:
                 await inner_stream
                 # Collect result hooks from the inner stream to run later
                 stream_result_hooks[:] = _get_result_hooks_from_stream(inner_stream)
+                dropping_post_limit_function_calls = (
+                    mutable_options.get("tool_choice") == "none"
+                    and max_function_calls is not None
+                    and total_function_calls >= max_function_calls
+                )
 
                 # Yield updates from the inner stream, letting it collect them
                 async for update in inner_stream:
+                    if dropping_post_limit_function_calls:
+                        update = _drop_function_call_contents_from_update(update)
+                        if update is None:
+                            continue
                     yield update
 
                 # Get the finalized response from the inner stream
                 # This triggers the inner stream's finalizer and result hooks
                 response = await inner_stream.get_final_response()
-                response_had_visible_content = _response_has_visible_content(response)
                 function_call_limit_reached = (
                     mutable_options.get("tool_choice") == "none"
                     and max_function_calls is not None
                     and total_function_calls >= max_function_calls
                 )
+                if function_call_limit_reached:
+                    _drop_function_call_contents_from_response(response)
+                response_had_visible_content = _response_has_visible_content(response)
                 if function_call_limit_reached:
                     response = _ensure_function_invocation_limit_fallback_response(response)
                 _update_continuation_state(

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

@@ -2,7 +2,7 @@
 
 import asyncio
 from collections.abc import AsyncIterable, Awaitable, Callable, Sequence
-from typing import Any
+from typing import Any, Literal
 
 import pytest
 
@@ -96,6 +96,47 @@ def _finalize(updates: Sequence[ChatResponseUpdate]) -> ChatResponse:
     chat_client_base._get_streaming_response = _get_streaming_response
 
 
+def _force_function_call_tool_choice_none_stream(
+    chat_client_base: Any,
+    *,
+    call_id: str,
+    name: str,
+    arguments: str,
+    additional_properties: dict[str, Any] | None = None,
+    finish_reason: Literal["stop", "length", "tool_calls", "content_filter"] | None = None,
+) -> None:
+    original_get_streaming_response = chat_client_base._get_streaming_response
+
+    def _get_streaming_response(
+        *,
+        messages: Sequence[Message],
+        options: dict[str, Any],
+        **kwargs: Any,
+    ) -> ResponseStream[ChatResponseUpdate, ChatResponse]:
+        if options.get("tool_choice") != "none":
+            return original_get_streaming_response(messages=messages, options=options, **kwargs)
+
+        updates = (
+            ChatResponseUpdate(
+                contents=[Content.from_function_call(call_id=call_id, name=name, arguments=arguments)],
+                role="assistant",
+                additional_properties=additional_properties,
+                finish_reason=finish_reason,
+            ),
+        )
+
+        async def _stream() -> AsyncIterable[ChatResponseUpdate]:
+            for update in updates:
+                yield update
+
+        def _finalize(updates: Sequence[ChatResponseUpdate]) -> ChatResponse:
+            return ChatResponse.from_updates(updates, output_format_type=options.get("response_format"))
+
+        return ResponseStream(_stream(), finalizer=_finalize)
+
+    chat_client_base._get_streaming_response = _get_streaming_response
+
+
 async def test_base_client_with_function_calling(chat_client_base: SupportsChatGetResponse):
     exec_counter = 0
 
@@ -2374,6 +2415,7 @@ def test_replace_approval_contents_with_results_allows_reused_call_id_after_comp
         ("call_reused", "second output"),
     ]
 
+
 def test_replace_approval_contents_with_results_uses_result_call_ids_for_placeholders() -> None:
     from agent_framework._tools import _collect_approval_responses, _replace_approval_contents_with_results
 
@@ -3284,6 +3326,109 @@ def lookup_func(key: str) -> str:
     assert updates[-1].text == _EXPECTED_FUNCTION_INVOCATION_LIMIT_FALLBACK_TEXT
 
 
+@pytest.mark.parametrize("max_iterations", [10])
+async def test_streaming_max_function_calls_drops_post_limit_function_call_update(
+    chat_client_base: SupportsChatGetResponse,
+):
+    """Function calls streamed after max_function_calls is reached should not be forwarded."""
+    _force_function_call_tool_choice_none_stream(
+        chat_client_base,
+        call_id="call_2",
+        name="lookup",
+        arguments='{"key": "b"}',
+    )
+    exec_counter = 0
+
+    @tool(name="lookup", approval_mode="never_require")
+    def lookup_func(key: str) -> str:
+        nonlocal exec_counter
+        exec_counter += 1
+        return f"Value for {key}"
+
+    chat_client_base.streaming_responses = [  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+        [
+            ChatResponseUpdate(
+                contents=[Content.from_function_call(call_id="call_1", name="lookup", arguments='{"key": "a"}')],
+                role="assistant",
+            ),
+        ],
+    ]
+    chat_client_base.function_invocation_configuration["max_function_calls"] = 1  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+
+    updates = []
+    async for update in chat_client_base.get_response(  # type: ignore[call-overload]  # pyrefly: ignore[no-matching-overload]
+        [Message(role="user", contents=["look up key"])],
+        options={"tool_choice": "auto", "tools": [lookup_func]},
+        stream=True,
+    ):
+        updates.append(update)
+
+    function_call_ids = {
+        content.call_id for update in updates for content in update.contents if content.type == "function_call"
+    }
+    function_result_ids = {
+        content.call_id for update in updates for content in update.contents if content.type == "function_result"
+    }
+
+    assert exec_counter == 1
+    assert "call_1" in function_call_ids
+    assert "call_1" in function_result_ids
+    assert "call_2" not in function_call_ids
+    assert "call_2" not in function_result_ids
+    assert updates[-1].role == "assistant"
+    assert updates[-1].text == _EXPECTED_FUNCTION_INVOCATION_LIMIT_FALLBACK_TEXT
+
+
+@pytest.mark.parametrize("max_iterations", [10])
+async def test_streaming_max_function_calls_preserves_post_limit_update_metadata(
+    chat_client_base: SupportsChatGetResponse,
+):
+    """Post-limit function call chunks should retain metadata after call content is stripped."""
+    _force_function_call_tool_choice_none_stream(
+        chat_client_base,
+        call_id="call_2",
+        name="lookup",
+        arguments='{"key": "b"}',
+        additional_properties={"provider_metadata": "keep"},
+        finish_reason="stop",
+    )
+    exec_counter = 0
+
+    @tool(name="lookup", approval_mode="never_require")
+    def lookup_func(key: str) -> str:
+        nonlocal exec_counter
+        exec_counter += 1
+        return f"Value for {key}"
+
+    chat_client_base.streaming_responses = [  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+        [
+            ChatResponseUpdate(
+                contents=[Content.from_function_call(call_id="call_1", name="lookup", arguments='{"key": "a"}')],
+                role="assistant",
+            ),
+        ],
+    ]
+    chat_client_base.function_invocation_configuration["max_function_calls"] = 1  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+
+    updates = []
+    async for update in chat_client_base.get_response(  # type: ignore[call-overload]  # pyrefly: ignore[no-matching-overload]
+        [Message(role="user", contents=["look up key"])],
+        options={"tool_choice": "auto", "tools": [lookup_func]},
+        stream=True,
+    ):
+        updates.append(update)
+
+    metadata_updates = [
+        update
+        for update in updates
+        if update.additional_properties == {"provider_metadata": "keep"} and update.finish_reason == "stop"
+    ]
+
+    assert exec_counter == 1
+    assert len(metadata_updates) == 1
+    assert metadata_updates[0].contents == []
+
+
 async def test_streaming_function_invocation_config_enabled_false(chat_client_base: SupportsChatGetResponse):
     """Test that setting enabled=False disables function invocation in streaming mode."""
     exec_counter = 0