From 05ce001e55f311c18d903ead919d7b639b8e7cd6 Mon Sep 17 00:00:00 2001 From: karthik-0306 Date: Thu, 20 Aug 2026 02:16:23 +0530 Subject: [PATCH 1/5] feat(core): add max_duration_seconds bound and stop_reason signal to tool loop (#7587) --- python/CHANGELOG.md | 3 + .../packages/core/agent_framework/_tools.py | 87 +++++- .../core/test_function_invocation_logic.py | 259 ++++++++++++++++++ 3 files changed, 341 insertions(+), 8 deletions(-) diff --git a/python/CHANGELOG.md b/python/CHANGELOG.md index f0669705f1..b9be8a3ec4 100644 --- a/python/CHANGELOG.md +++ b/python/CHANGELOG.md @@ -7,6 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added +- **agent-framework-core**: Add `max_duration_seconds` to `FunctionInvocationConfiguration` to bound total wall-clock time of the function-invocation loop; when exceeded, tools are disabled and the model is forced to produce a final text response via the same graceful-degradation path as `max_function_calls`. Add `_agent_framework_stop_reason` to `ChatResponse.additional_properties` (values: `"completed"`, `"max_iterations"`, `"max_duration_seconds"`) so callers can detect how a run ended ([#7587](https://github.com/microsoft/agent-framework/issues/7587)) + ## [1.14.0] - 2026-08-13 ### Added diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 45fd29864e..12873e5946 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -1344,6 +1344,17 @@ class FunctionInvocationConfiguration(TypedDict, total=False): parallel tool calls completes, not before. If the model requests 20 parallel calls in a single iteration and the limit is 10, all 20 will execute before the loop stops. + - ``max_duration_seconds``: Wall-clock time budget (in seconds) for the + entire function-invocation loop, measured cumulatively across approval + round-trips. When exceeded, the loop disables further tool calls and + forces the model to produce a text response, reusing the same + graceful-degradation path as ``max_function_calls``. Default is + ``None`` (no time limit). + + This is a **best-effort** limit: the clock is checked *after* each batch + of tool calls completes. The budget includes time spent waiting for human + approval responses — this is intentional so that long-running unattended + sessions are always bounded, even if individual approval steps are slow. - ``max_consecutive_errors_per_request``: How many consecutive errors before abandoning the tool loop for this request. - ``terminate_on_unknown_calls``: Whether to raise an error when the model @@ -1354,11 +1365,13 @@ class FunctionInvocationConfiguration(TypedDict, total=False): function result returned to the model. Note: - ``max_iterations`` and ``max_function_calls`` serve complementary purposes. - ``max_iterations`` caps the number of model round-trips regardless of how - many tools are called per trip. ``max_function_calls`` caps the cumulative - number of individual tool executions regardless of how they are distributed - across iterations. + ``max_iterations``, ``max_function_calls``, and ``max_duration_seconds`` + are complementary limits. ``max_iterations`` caps the number of model + round-trips, ``max_function_calls`` caps the cumulative number of + individual tool executions, and ``max_duration_seconds`` caps cumulative + elapsed wall time. When multiple limits trigger in the same batch, the + stop reason is assigned to whichever condition is detected first in + execution order (duration is checked before the call-count check). Example: .. code-block:: python @@ -1367,14 +1380,17 @@ class FunctionInvocationConfiguration(TypedDict, total=False): client = OpenAIChatClient(api_key="your_api_key") - # Limit to 5 LLM roundtrips and 20 total function executions + # Limit to 5 LLM roundtrips, 20 total function executions, + # and a 30-second wall-clock budget. client.function_invocation_configuration["max_iterations"] = 5 client.function_invocation_configuration["max_function_calls"] = 20 + client.function_invocation_configuration["max_duration_seconds"] = 30.0 """ enabled: bool max_iterations: int max_function_calls: int | None + max_duration_seconds: float | None max_consecutive_errors_per_request: int terminate_on_unknown_calls: bool additional_tools: Sequence[FunctionTool] @@ -1388,6 +1404,7 @@ def normalize_function_invocation_configuration( "enabled": True, "max_iterations": DEFAULT_MAX_ITERATIONS, "max_function_calls": None, + "max_duration_seconds": None, "max_consecutive_errors_per_request": DEFAULT_MAX_CONSECUTIVE_ERRORS_PER_REQUEST, "terminate_on_unknown_calls": False, "additional_tools": [], @@ -1399,6 +1416,8 @@ def normalize_function_invocation_configuration( raise ValueError("max_iterations must be at least 1.") if normalized["max_function_calls"] is not None and normalized["max_function_calls"] < 1: raise ValueError("max_function_calls must be at least 1 or None.") + if normalized["max_duration_seconds"] is not None and normalized["max_duration_seconds"] <= 0: + raise ValueError("max_duration_seconds must be greater than 0 or None.") if normalized["max_consecutive_errors_per_request"] < 0: raise ValueError("max_consecutive_errors_per_request must be 0 or more.") return normalized @@ -3187,6 +3206,7 @@ async def _get_response_with_function_invocation( errors_in_a_row = 0 total_function_calls = int(budget_state.get("total_function_calls", 0) or 0) max_function_calls = self.function_invocation_configuration.get("max_function_calls") + max_duration_seconds = self.function_invocation_configuration.get("max_duration_seconds") prepared_messages = _copy_messages_for_function_invocation(messages) function_call_messages: list[Message] = [] response: ChatResponse[Any] | None = None @@ -3225,6 +3245,7 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non if approval_processing.action == "return": response = ChatResponse(messages=list(function_call_messages)) response.usage_details = aggregated_usage + response.additional_properties["_agent_framework_stop_reason"] = "completed" return _clear_internal_conversation_id(response) if approval_processing.action == "stop": options["tool_choice"] = "none" @@ -3290,9 +3311,26 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non ) if function_processing.action == "return": response.usage_details = aggregated_usage + response.additional_properties["_agent_framework_stop_reason"] = budget_state.get( + "stop_reason", "completed" + ) return _clear_internal_conversation_id(response) + # Check duration limit before the call-count check so duration wins precedence + # when both limits are reached in the same batch. + if ( + max_duration_seconds is not None + and (perf_counter() - budget_state["start_time"]) >= max_duration_seconds + ): + logger.info( + "Maximum duration reached (%.2fs / %.2fs). Stopping further function calls for this request.", + perf_counter() - budget_state["start_time"], + max_duration_seconds, + ) + options["tool_choice"] = "none" + budget_state.setdefault("stop_reason", "max_duration_seconds") if function_processing.action == "stop": options["tool_choice"] = "none" + budget_state.setdefault("stop_reason", "stop") else: _disable_tools_at_function_call_limit(options, total_function_calls, max_function_calls) errors_in_a_row = function_processing.errors_in_a_row @@ -3300,6 +3338,7 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non _prepare_messages_for_next_iteration(prepared_messages, response) # Phase 3: the iteration budget is exhausted, so request one final response with tools disabled. + budget_state.setdefault("stop_reason", "max_iterations") if response is not None: logger.info( "Maximum iterations reached (%d). Requesting final response without tools.", @@ -3327,6 +3366,7 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non ) response.usage_details = aggregated_usage _prepend_function_call_messages(response, function_call_messages) + response.additional_properties["_agent_framework_stop_reason"] = budget_state.get("stop_reason", "completed") return _clear_internal_conversation_id(response) async def _stream_response_with_function_invocation( @@ -3349,6 +3389,7 @@ async def _stream_response_with_function_invocation( errors_in_a_row = 0 total_function_calls = int(budget_state.get("total_function_calls", 0) or 0) max_function_calls = self.function_invocation_configuration.get("max_function_calls") + max_duration_seconds = self.function_invocation_configuration.get("max_duration_seconds") prepared_messages = _copy_messages_for_function_invocation(messages) response: ChatResponse[Any] | None = None max_iterations = self.function_invocation_configuration.get("max_iterations", DEFAULT_MAX_ITERATIONS) @@ -3474,14 +3515,31 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non yield update if function_processing.action == "stop": options["tool_choice"] = "none" + budget_state.setdefault("stop_reason", "stop") elif function_processing.action != "continue": + # "return" action: model produced a terminal response. return else: - _disable_tools_at_function_call_limit(options, total_function_calls, max_function_calls) + # Check duration limit before the call-count check so duration wins precedence + # when both limits are reached in the same batch. + if ( + max_duration_seconds is not None + and (perf_counter() - budget_state["start_time"]) >= max_duration_seconds + ): + logger.info( + "Maximum duration reached (%.2fs / %.2fs). Stopping further function calls for this request.", + perf_counter() - budget_state["start_time"], + max_duration_seconds, + ) + options["tool_choice"] = "none" + budget_state.setdefault("stop_reason", "max_duration_seconds") + else: + _disable_tools_at_function_call_limit(options, total_function_calls, max_function_calls) _reset_required_tool_choice(options) _prepare_messages_for_next_iteration(prepared_messages, response) # Phase 3: the iteration budget is exhausted, so stream one final response with tools disabled. + budget_state.setdefault("stop_reason", "max_iterations") if response is not None: logger.info( "Maximum iterations reached (%d). Requesting final response without tools.", @@ -3606,9 +3664,14 @@ def get_response( budget_state: dict[str, Any] = ( cast(dict[str, Any], raw_budget_state) if isinstance(raw_budget_state, dict) else {} ) + # Record the start time once for the full logical run (including approval round-trips). + # setdefault preserves the original timestamp across approval re-entries so that + # max_duration_seconds measures cumulative elapsed time, not just the current segment. + budget_state.setdefault("start_time", perf_counter()) max_errors = self.function_invocation_configuration.get( "max_consecutive_errors_per_request", DEFAULT_MAX_CONSECUTIVE_ERRORS_PER_REQUEST ) + additional_function_arguments = ( dict(function_invocation_kwargs) if function_invocation_kwargs is not None else {} ) @@ -3668,6 +3731,14 @@ def get_response( ) response_format = mutable_options.get("response_format") + + def _finalize_with_stop_reason(updates: Sequence[ChatResponseUpdate]) -> ChatResponse[Any]: + response = ChatResponse.from_updates(updates, output_format_type=response_format) + response.additional_properties["_agent_framework_stop_reason"] = budget_state.get( + "stop_reason", "completed" + ) + return response + return ResponseStream( self._stream_response_with_function_invocation( super_get_response=super_get_response, @@ -3681,7 +3752,7 @@ def get_response( budget_state=budget_state, max_errors=max_errors, ), - finalizer=partial(ChatResponse.from_updates, output_format_type=response_format), + finalizer=_finalize_with_stop_reason, ) diff --git a/python/packages/core/tests/core/test_function_invocation_logic.py b/python/packages/core/tests/core/test_function_invocation_logic.py index 1d6c70fb39..891e17e97a 100644 --- a/python/packages/core/tests/core/test_function_invocation_logic.py +++ b/python/packages/core/tests/core/test_function_invocation_logic.py @@ -6384,3 +6384,262 @@ def a(x: int) -> int: # endregion + + +# region max_duration_seconds and stop_reason + + +@pytest.mark.parametrize("max_iterations", [10]) +async def test_max_duration_seconds_non_streaming_triggers_graceful_degradation( + chat_client_base: SupportsChatGetResponse, +): + """When max_duration_seconds is exceeded mid-loop, tools are disabled and the model + produces a final text response (same graceful-degradation path as max_function_calls).""" + exec_counter = 0 + + @tool(name="step", approval_mode="never_require") + def step_func() -> str: + nonlocal exec_counter + exec_counter += 1 + return "done" + + chat_client_base.function_invocation_configuration["max_duration_seconds"] = 5.0 # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + chat_client_base.run_responses = [ # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + ChatResponse( + messages=Message( + role="assistant", + contents=[Content.from_function_call(call_id="c1", name="step", arguments="{}")], + ) + ), + # model would call again, but duration check fires first + ChatResponse( + messages=Message( + role="assistant", + contents=[Content.from_function_call(call_id="c2", name="step", arguments="{}")], + ) + ), + ChatResponse(messages=Message(role="assistant", contents=["giving up"])), + ] + + call_count = 0 + + def fake_perf_counter() -> float: + nonlocal call_count + call_count += 1 + # First call (start_time setdefault): t=0. Second call (duration check after iteration 1): t=10. + return 0.0 if call_count == 1 else 10.0 + + from unittest.mock import patch + + with patch("agent_framework._tools.perf_counter", side_effect=fake_perf_counter): + response = await chat_client_base.get_response( + [Message(role="user", contents=["go"])], options={"tool_choice": "auto", "tools": [step_func]} + ) + + # Only first tool call executes; duration fires before second iteration's tool call + assert exec_counter == 1 + assert response.additional_properties.get("_agent_framework_stop_reason") == "max_duration_seconds" + + +@pytest.mark.parametrize("max_iterations", [10]) +async def test_max_duration_seconds_streaming_triggers_graceful_degradation( + chat_client_base: SupportsChatGetResponse, +): + """max_duration_seconds works in the streaming path: stop_reason is in the final ChatResponse.""" + exec_counter = 0 + + @tool(name="step", approval_mode="never_require") + def step_func() -> str: + nonlocal exec_counter + exec_counter += 1 + return "done" + + chat_client_base.function_invocation_configuration["max_duration_seconds"] = 5.0 # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + chat_client_base.streaming_responses = [ # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + [ + ChatResponseUpdate( + contents=[Content.from_function_call(call_id="s1", name="step", arguments="{}")], + role="assistant", + finish_reason="tool_calls", + ) + ], + # duration fires after iteration 1; second batch never executes tool + [ + ChatResponseUpdate( + contents=[Content.from_function_call(call_id="s2", name="step", arguments="{}")], + role="assistant", + finish_reason="tool_calls", + ) + ], + ] + + call_count = 0 + + def fake_perf_counter() -> float: + nonlocal call_count + call_count += 1 + return 0.0 if call_count == 1 else 10.0 + + from unittest.mock import patch + + with patch("agent_framework._tools.perf_counter", side_effect=fake_perf_counter): + stream = chat_client_base.get_response( + [Message(role="user", contents=["go"])], + stream=True, + options={"tool_choice": "auto", "tools": [step_func]}, + ) + async for _ in stream: + pass + final = await stream.get_final_response() + + assert exec_counter == 1 + assert final.additional_properties.get("_agent_framework_stop_reason") == "max_duration_seconds" + + +@pytest.mark.parametrize("max_iterations", [10]) +async def test_stop_reason_completed_on_normal_finish(chat_client_base: SupportsChatGetResponse): + """When the model finishes normally (no limit hit), stop_reason is 'completed'.""" + + @tool(name="q", approval_mode="never_require") + def q_func() -> str: + return "answer" + + chat_client_base.run_responses = [ # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + ChatResponse( + messages=Message( + role="assistant", + contents=[Content.from_function_call(call_id="q1", name="q", arguments="{}")], + ) + ), + ChatResponse(messages=Message(role="assistant", contents=["Final answer."])), + ] + + response = await chat_client_base.get_response( + [Message(role="user", contents=["question"])], options={"tool_choice": "auto", "tools": [q_func]} + ) + + assert response.additional_properties.get("_agent_framework_stop_reason") == "completed" + + +@pytest.mark.parametrize("max_iterations", [1]) +async def test_stop_reason_max_iterations(chat_client_base: SupportsChatGetResponse): + """When the iteration budget is exhausted, stop_reason is 'max_iterations'.""" + + @tool(name="loop", approval_mode="never_require") + def loop_func() -> str: + return "looping" + + chat_client_base.run_responses = [ # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + ChatResponse( + messages=Message( + role="assistant", + contents=[Content.from_function_call(call_id="l1", name="loop", arguments="{}")], + ) + ), + ChatResponse(messages=Message(role="assistant", contents=["Gave up."])), + ] + + response = await chat_client_base.get_response( + [Message(role="user", contents=["loop"])], options={"tool_choice": "auto", "tools": [loop_func]} + ) + + assert response.additional_properties.get("_agent_framework_stop_reason") == "max_iterations" + + +@pytest.mark.parametrize("max_iterations", [10]) +async def test_stop_reason_max_function_calls(chat_client_base: SupportsChatGetResponse): + """When the function-call budget is exhausted, stop_reason reflects the limit-triggered degradation.""" + + @tool(name="w", approval_mode="never_require") + def w_func() -> str: + return "work" + + chat_client_base.function_invocation_configuration["max_function_calls"] = 1 # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + chat_client_base.run_responses = [ # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + ChatResponse( + messages=Message( + role="assistant", + contents=[Content.from_function_call(call_id="w1", name="w", arguments="{}")], + ) + ), + ChatResponse( + messages=Message( + role="assistant", + contents=[Content.from_function_call(call_id="w2", name="w", arguments="{}")], + ) + ), + ChatResponse(messages=Message(role="assistant", contents=["Done."])), + ] + + response = await chat_client_base.get_response( + [Message(role="user", contents=["work"])], options={"tool_choice": "auto", "tools": [w_func]} + ) + + # max_function_calls degradation: stop_reason will be "completed" (the final phase-3 fallback + # path doesn't set a specific reason for max_function_calls; the limit disables tools and + # the model returns normally via action=="return"). This test documents the current behaviour. + assert response.additional_properties.get("_agent_framework_stop_reason") in ("completed", "max_iterations") + + +@pytest.mark.parametrize("max_iterations", [10]) +async def test_duration_wins_precedence_over_iterations(chat_client_base: SupportsChatGetResponse): + """When duration fires in the same batch as the last iteration, duration takes precedence + (setdefault on stop_reason is set by the duration check before Phase 3 stamps max_iterations).""" + exec_counter = 0 + + @tool(name="t", approval_mode="never_require") + def t_func() -> str: + nonlocal exec_counter + exec_counter += 1 + return "t" + + # Exactly max_iterations=1 so the loop falls into Phase 3 — but duration also fires. + chat_client_base.function_invocation_configuration["max_iterations"] = 1 # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + chat_client_base.function_invocation_configuration["max_duration_seconds"] = 0.001 # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + chat_client_base.run_responses = [ # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + ChatResponse( + messages=Message( + role="assistant", + contents=[Content.from_function_call(call_id="t1", name="t", arguments="{}")], + ) + ), + ChatResponse(messages=Message(role="assistant", contents=["done"])), + ] + + call_count = 0 + + def fake_perf_counter() -> float: + nonlocal call_count + call_count += 1 + return 0.0 if call_count == 1 else 100.0 # always exceeded after start + + from unittest.mock import patch + + with patch("agent_framework._tools.perf_counter", side_effect=fake_perf_counter): + response = await chat_client_base.get_response( + [Message(role="user", contents=["go"])], options={"tool_choice": "auto", "tools": [t_func]} + ) + + # Duration check runs inside the loop (before Phase 3 stamps max_iterations via setdefault). + assert response.additional_properties.get("_agent_framework_stop_reason") == "max_duration_seconds" + + +def test_normalize_rejects_non_positive_max_duration_seconds(): + """normalize_function_invocation_configuration raises ValueError for max_duration_seconds <= 0.""" + from agent_framework import normalize_function_invocation_configuration + + with pytest.raises(ValueError, match="max_duration_seconds"): + normalize_function_invocation_configuration({"max_duration_seconds": 0.0}) + + with pytest.raises(ValueError, match="max_duration_seconds"): + normalize_function_invocation_configuration({"max_duration_seconds": -1.0}) + + # None and positive values are accepted. + cfg = normalize_function_invocation_configuration({"max_duration_seconds": None}) + assert cfg["max_duration_seconds"] is None + + cfg = normalize_function_invocation_configuration({"max_duration_seconds": 30.0}) + assert cfg["max_duration_seconds"] == 30.0 + + +# endregion From 7483e81b4cd3bd8c5a6e01bd29120843cd6f63e2 Mon Sep 17 00:00:00 2001 From: karthik-0306 Date: Fri, 21 Aug 2026 22:06:39 +0530 Subject: [PATCH 2/5] fix(core): unify batch limit decision precedence and propagate stop_reason to AgentResponse --- python/CHANGELOG.md | 2 +- .../packages/core/agent_framework/_agents.py | 26 ++- .../_harness/_tool_approval.py | 3 +- .../packages/core/agent_framework/_tools.py | 195 +++++++++++++----- .../core/test_function_invocation_logic.py | 124 ++++++++++- 5 files changed, 291 insertions(+), 59 deletions(-) diff --git a/python/CHANGELOG.md b/python/CHANGELOG.md index b9be8a3ec4..ba1fb52c4c 100644 --- a/python/CHANGELOG.md +++ b/python/CHANGELOG.md @@ -8,7 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added -- **agent-framework-core**: Add `max_duration_seconds` to `FunctionInvocationConfiguration` to bound total wall-clock time of the function-invocation loop; when exceeded, tools are disabled and the model is forced to produce a final text response via the same graceful-degradation path as `max_function_calls`. Add `_agent_framework_stop_reason` to `ChatResponse.additional_properties` (values: `"completed"`, `"max_iterations"`, `"max_duration_seconds"`) so callers can detect how a run ended ([#7587](https://github.com/microsoft/agent-framework/issues/7587)) +- **agent-framework-core**: Add `max_duration_seconds` to `FunctionInvocationConfiguration` to bound total wall-clock time of the function-invocation loop; when exceeded, tools are disabled and the model is forced to produce a final text response via the same graceful-degradation path as `max_function_calls`. Rejects `float("nan")` during configuration validation. Add `_agent_framework_stop_reason` to `ChatResponse.additional_properties` with machine-readable enum values (`"completed"`, `"max_iterations"`, `"max_duration_seconds"`, `"max_function_calls"`, `"max_consecutive_errors"`) so callers can detect how a run ended ([#7587](https://github.com/microsoft/agent-framework/issues/7587), [#7772](https://github.com/microsoft/agent-framework/pull/7772)). Includes a small, additive update to `ToolApprovalMiddleware` to persist budget state across human approval round-trips via session state. ## [1.14.0] - 2026-08-13 diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py index 85615b111f..d5651a1aef 100644 --- a/python/packages/core/agent_framework/_agents.py +++ b/python/packages/core/agent_framework/_agents.py @@ -1268,10 +1268,34 @@ def _suppress_response_id(update: AgentResponseUpdate) -> AgentResponseUpdate: return update def _finalizer(updates: Sequence[AgentResponseUpdate]) -> AgentResponse[Any]: - return self._finalize_response_updates( + agent_response = self._finalize_response_updates( updates, response_format=context["chat_options"].get("response_format"), ) + # Propagate _agent_framework_stop_reason from the inner ChatResponse to the outer + # AgentResponse. _finalize_with_stop_reason sets it on the inner ChatResponse + # (a ResponseStream finalizer), but AgentResponse.from_updates only sees the + # AgentResponseUpdate sequence — none of which carries this key. The non-streaming + # path copies additional_properties directly via _build_agent_response_from_chat_response; + # for streaming we recover it here from the inner ChatResponse stored in the closure. + if inner_chat_responses: + inner = inner_chat_responses[0] + stop_reason = inner.additional_properties.get("_agent_framework_stop_reason") + if ( + stop_reason is not None + and "_agent_framework_stop_reason" not in agent_response.additional_properties + ): + agent_response.additional_properties["_agent_framework_stop_reason"] = stop_reason + return agent_response + + # Mutable container for the inner ChatResponse; populated by the result hook below before + # _finalizer is called (inner finalizer runs before outer finalizer per ResponseStream.map contract). + inner_chat_responses: list[ChatResponse[Any]] = [] + + def capture_inner_chat_response(inner: ChatResponse[Any]) -> None: + inner_chat_responses.append(inner) + + stream_response = stream_response.with_result_hook(capture_inner_chat_response) stream = stream_response.map( transform=partial( diff --git a/python/packages/core/agent_framework/_harness/_tool_approval.py b/python/packages/core/agent_framework/_harness/_tool_approval.py index 5c3bf7a2c9..f833258776 100644 --- a/python/packages/core/agent_framework/_harness/_tool_approval.py +++ b/python/packages/core/agent_framework/_harness/_tool_approval.py @@ -385,7 +385,8 @@ async def process(self, context: AgentContext, call_next: Callable[[], Awaitable raise RuntimeError("ToolApprovalMiddleware requires an AgentSession.") state = _get_state(context.session, source_id=self.source_id) - context.client_kwargs.setdefault(_FUNCTION_INVOCATION_BUDGET_STATE_KEY, {}) + session_budget = context.session.state.setdefault(_FUNCTION_INVOCATION_BUDGET_STATE_KEY, {}) + context.client_kwargs.setdefault(_FUNCTION_INVOCATION_BUDGET_STATE_KEY, session_budget) context.messages = self._prepare_inbound_messages(context.messages, state, context.session) await self._drain_auto_approvable_queue(state) if next_queued := self._pop_next_queued_request(state): diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 12873e5946..b344687f20 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -1416,7 +1416,7 @@ def normalize_function_invocation_configuration( raise ValueError("max_iterations must be at least 1.") if normalized["max_function_calls"] is not None and normalized["max_function_calls"] < 1: raise ValueError("max_function_calls must be at least 1 or None.") - if normalized["max_duration_seconds"] is not None and normalized["max_duration_seconds"] <= 0: + if normalized["max_duration_seconds"] is not None and not normalized["max_duration_seconds"] > 0: raise ValueError("max_duration_seconds must be greater than 0 or None.") if normalized["max_consecutive_errors_per_request"] < 0: raise ValueError("max_consecutive_errors_per_request must be 0 or more.") @@ -2725,6 +2725,26 @@ def _copy_messages_for_function_invocation(messages: Any) -> list[Message]: return copied_messages +def _is_duration_limit_reached(budget_state: dict[str, Any], max_duration_seconds: float | None) -> bool: + if max_duration_seconds is None: + return False + start_time = budget_state.get("start_time") + if start_time is None: + return False + return (perf_counter() - start_time) >= max_duration_seconds + + +def _is_tool_limit_reached( + total_function_calls: int, + max_function_calls: int | None, + budget_state: dict[str, Any], + max_duration_seconds: float | None, +) -> bool: + return _function_call_limit_reached(total_function_calls, max_function_calls) or _is_duration_limit_reached( + budget_state, max_duration_seconds + ) + + def _function_call_limit_reached(total_function_calls: int, max_function_calls: int | None) -> bool: return max_function_calls is not None and total_function_calls >= max_function_calls @@ -2751,6 +2771,7 @@ def _disable_tools_at_function_call_limit( options: dict[str, Any], total_function_calls: int, max_function_calls: int | None, + budget_state: dict[str, Any] | None = None, ) -> None: if not _function_call_limit_reached(total_function_calls, max_function_calls): return @@ -2760,6 +2781,8 @@ def _disable_tools_at_function_call_limit( max_function_calls, ) options["tool_choice"] = "none" + if budget_state is not None: + budget_state.setdefault("stop_reason", "max_function_calls") def _record_function_calls( @@ -2779,6 +2802,33 @@ def _reset_required_tool_choice(options: dict[str, Any]) -> None: options["tool_choice"] = None +def _apply_batch_limit_decision( + *, + options: dict[str, Any], + budget_state: dict[str, Any], + total_function_calls: int, + max_function_calls: int | None, + max_duration_seconds: float | None, +) -> None: + """Apply the post-batch tool-disable and stop-reason decision for the 'continue' path. + + Single source of truth for limit precedence shared by streaming and non-streaming loops: + duration is checked before call-count so it wins via ``setdefault`` when both fire in the + same batch. The ``action=='stop'`` (consecutive-error) path is handled identically in both + loops and is not routed through here. + """ + if _is_duration_limit_reached(budget_state, max_duration_seconds): + logger.info( + "Maximum duration reached (%.2fs / %.2fs). Stopping further function calls for this request.", + perf_counter() - budget_state["start_time"], + max_duration_seconds, + ) + options["tool_choice"] = "none" + budget_state.setdefault("stop_reason", "max_duration_seconds") + else: + _disable_tools_at_function_call_limit(options, total_function_calls, max_function_calls, budget_state) + + def _prepare_messages_for_next_iteration(prepared_messages: list[Message], response: ChatResponse[Any]) -> None: if response.conversation_id is None: prepared_messages.extend(response.messages) @@ -2932,7 +2982,11 @@ async def _resolve_approval_responses( execution_result_groups: list[list[Content]] = [] should_terminate = False reached_error_limit = False - if responses_to_execute: + # Skip execution when tool invocation is globally disabled for this request, regardless of why: + # duration limit, max_function_calls limit, or explicit caller-set "none" all mean the same thing — + # no more tool calls this turn. Narrowing to only the duration case would allow max_function_calls + # or explicit-none callers to still execute approved calls, which would be wrong. + if responses_to_execute and not (options and options.get("tool_choice") == "none"): try: execution = await execute_function_calls( function_calls=responses_to_execute, @@ -3226,6 +3280,14 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non ) # Phase 1: resolve inbound approvals before consuming another model iteration. + if _is_duration_limit_reached(budget_state, max_duration_seconds): + logger.info( + "Maximum duration reached (%.2fs / %.2fs). Stopping further function calls for this request.", + perf_counter() - budget_state["start_time"], + max_duration_seconds, + ) + options["tool_choice"] = "none" + budget_state.setdefault("stop_reason", "max_duration_seconds") approval_processing = await _resolve_approval_responses( prepared_messages=prepared_messages, options=options, @@ -3245,12 +3307,17 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non if approval_processing.action == "return": response = ChatResponse(messages=list(function_call_messages)) response.usage_details = aggregated_usage - response.additional_properties["_agent_framework_stop_reason"] = "completed" + response.additional_properties["_agent_framework_stop_reason"] = budget_state.get( + "stop_reason", "completed" + ) + if invocation_session is not None: + invocation_session.state.pop(_FUNCTION_INVOCATION_BUDGET_STATE_KEY, None) return _clear_internal_conversation_id(response) if approval_processing.action == "stop": options["tool_choice"] = "none" + budget_state.setdefault("stop_reason", "max_consecutive_errors") else: - _disable_tools_at_function_call_limit(options, total_function_calls, max_function_calls) + _disable_tools_at_function_call_limit(options, total_function_calls, max_function_calls, budget_state) # Phase 2: alternate model turns and local execution until a terminal response or safety limit is reached. for attempt_idx in range(attempt_start, max_iterations): @@ -3266,8 +3333,8 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non client_kwargs=request_kwargs, ), ) - if options.get("tool_choice") == "none" and _function_call_limit_reached( - total_function_calls, max_function_calls + if options.get("tool_choice") == "none" and _is_tool_limit_reached( + total_function_calls, max_function_calls, budget_state, max_duration_seconds ): _ensure_function_invocation_limit_fallback_response(response) aggregated_usage = add_usage_details(aggregated_usage, response.usage_details) @@ -3314,25 +3381,25 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non response.additional_properties["_agent_framework_stop_reason"] = budget_state.get( "stop_reason", "completed" ) - return _clear_internal_conversation_id(response) - # Check duration limit before the call-count check so duration wins precedence - # when both limits are reached in the same batch. - if ( - max_duration_seconds is not None - and (perf_counter() - budget_state["start_time"]) >= max_duration_seconds - ): - logger.info( - "Maximum duration reached (%.2fs / %.2fs). Stopping further function calls for this request.", - perf_counter() - budget_state["start_time"], - max_duration_seconds, + has_pending_approvals = any( + item.type == "function_approval_request" + for message in response.messages + for item in message.contents ) - options["tool_choice"] = "none" - budget_state.setdefault("stop_reason", "max_duration_seconds") + if invocation_session is not None and not has_pending_approvals: + invocation_session.state.pop(_FUNCTION_INVOCATION_BUDGET_STATE_KEY, None) + return _clear_internal_conversation_id(response) if function_processing.action == "stop": options["tool_choice"] = "none" - budget_state.setdefault("stop_reason", "stop") + budget_state.setdefault("stop_reason", "max_consecutive_errors") else: - _disable_tools_at_function_call_limit(options, total_function_calls, max_function_calls) + _apply_batch_limit_decision( + options=options, + budget_state=budget_state, + total_function_calls=total_function_calls, + max_function_calls=max_function_calls, + max_duration_seconds=max_duration_seconds, + ) errors_in_a_row = function_processing.errors_in_a_row _reset_required_tool_choice(options) _prepare_messages_for_next_iteration(prepared_messages, response) @@ -3366,6 +3433,8 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non ) response.usage_details = aggregated_usage _prepend_function_call_messages(response, function_call_messages) + if invocation_session is not None: + invocation_session.state.pop(_FUNCTION_INVOCATION_BUDGET_STATE_KEY, None) response.additional_properties["_agent_framework_stop_reason"] = budget_state.get("stop_reason", "completed") return _clear_internal_conversation_id(response) @@ -3407,6 +3476,14 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non ) # Phase 1: resolve and emit inbound approval outcomes before opening another provider stream. + if _is_duration_limit_reached(budget_state, max_duration_seconds): + logger.info( + "Maximum duration reached (%.2fs / %.2fs). Stopping further function calls for this request.", + perf_counter() - budget_state["start_time"], + max_duration_seconds, + ) + options["tool_choice"] = "none" + budget_state.setdefault("stop_reason", "max_duration_seconds") approval_processing = await _resolve_approval_responses( prepared_messages=prepared_messages, options=options, @@ -3425,11 +3502,14 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non for update in approval_processing.streaming_updates: yield update if approval_processing.action == "return": + if invocation_session is not None: + invocation_session.state.pop(_FUNCTION_INVOCATION_BUDGET_STATE_KEY, None) return if approval_processing.action == "stop": options["tool_choice"] = "none" + budget_state.setdefault("stop_reason", "max_consecutive_errors") else: - _disable_tools_at_function_call_limit(options, total_function_calls, max_function_calls) + _disable_tools_at_function_call_limit(options, total_function_calls, max_function_calls, budget_state) # Phase 2: stream each model turn, finalize it, execute its calls, then advance the transcript. for attempt_idx in range(attempt_start, max_iterations): @@ -3446,9 +3526,11 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non ), ) await inner_stream - drop_unexecutable_calls = options.get("tool_choice") == "none" and _function_call_limit_reached( + drop_unexecutable_calls = options.get("tool_choice") == "none" and _is_tool_limit_reached( total_function_calls, max_function_calls, + budget_state, + max_duration_seconds, ) async for update in inner_stream: if drop_unexecutable_calls: @@ -3458,8 +3540,8 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non yield update response = await inner_stream.get_final_response() - function_call_limit_reached = options.get("tool_choice") == "none" and _function_call_limit_reached( - total_function_calls, max_function_calls + function_call_limit_reached = options.get("tool_choice") == "none" and _is_tool_limit_reached( + total_function_calls, max_function_calls, budget_state, max_duration_seconds ) fallback_added = False if function_call_limit_reached: @@ -3478,6 +3560,8 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non ): if fallback_added: yield _function_invocation_limit_fallback_update() + if invocation_session is not None: + invocation_session.state.pop(_FUNCTION_INVOCATION_BUDGET_STATE_KEY, None) return try: @@ -3515,26 +3599,24 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non yield update if function_processing.action == "stop": options["tool_choice"] = "none" - budget_state.setdefault("stop_reason", "stop") + budget_state.setdefault("stop_reason", "max_consecutive_errors") elif function_processing.action != "continue": - # "return" action: model produced a terminal response. + # "return" action: the streamed response is terminal (no executable function calls + # remain). Unlike the non-streaming path, no has_pending_approvals guard is needed + # here: _process_model_function_calls stores any approval requests under + # _PENDING_APPROVAL_REQUESTS_KEY *before* returning action="return", so they are + # already persisted separately. Popping the budget key unconditionally is safe. + if invocation_session is not None: + invocation_session.state.pop(_FUNCTION_INVOCATION_BUDGET_STATE_KEY, None) return else: - # Check duration limit before the call-count check so duration wins precedence - # when both limits are reached in the same batch. - if ( - max_duration_seconds is not None - and (perf_counter() - budget_state["start_time"]) >= max_duration_seconds - ): - logger.info( - "Maximum duration reached (%.2fs / %.2fs). Stopping further function calls for this request.", - perf_counter() - budget_state["start_time"], - max_duration_seconds, - ) - options["tool_choice"] = "none" - budget_state.setdefault("stop_reason", "max_duration_seconds") - else: - _disable_tools_at_function_call_limit(options, total_function_calls, max_function_calls) + _apply_batch_limit_decision( + options=options, + budget_state=budget_state, + total_function_calls=total_function_calls, + max_function_calls=max_function_calls, + max_duration_seconds=max_duration_seconds, + ) _reset_required_tool_choice(options) _prepare_messages_for_next_iteration(prepared_messages, response) @@ -3558,6 +3640,8 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non ), ) await final_inner_stream + if invocation_session is not None: + invocation_session.state.pop(_FUNCTION_INVOCATION_BUDGET_STATE_KEY, None) async for update in final_inner_stream: update = _drop_unexecutable_tool_contents_from_update(update) if update is None: @@ -3660,14 +3744,31 @@ def get_response( ) if categorized_runtime_middleware["chat"]: request_kwargs["middleware"] = categorized_runtime_middleware["chat"] + from ._sessions import AgentSession as _AgentSession + + raw_session = request_kwargs.get("session") + invocation_session = raw_session if isinstance(raw_session, _AgentSession) else None + raw_budget_state = request_kwargs.pop(_FUNCTION_INVOCATION_BUDGET_STATE_KEY, None) budget_state: dict[str, Any] = ( cast(dict[str, Any], raw_budget_state) if isinstance(raw_budget_state, dict) else {} ) - # Record the start time once for the full logical run (including approval round-trips). - # setdefault preserves the original timestamp across approval re-entries so that - # max_duration_seconds measures cumulative elapsed time, not just the current segment. + if invocation_session is not None: + # Only merge with session budget when ToolApprovalMiddleware has already placed + # the key there (approval round-trip path). For plain session-based conversations + # without ToolApprovalMiddleware or max_duration_seconds, this is a no-op so that + # budget state remains isolated per agent.run() call, matching pre-PR behavior. + existing_session_budget = invocation_session.state.get(_FUNCTION_INVOCATION_BUDGET_STATE_KEY) + if isinstance(existing_session_budget, dict): + for k, v in budget_state.items(): + existing_session_budget[k] = v + budget_state = cast(dict[str, Any], existing_session_budget) + + # Record start time for this invocation. setdefault preserves the existing start_time when + # re-entering an approval round-trip (budget_state was merged from session.state above), + # so duration is measured from the *first* agent.run() call, not the resume call. budget_state.setdefault("start_time", perf_counter()) + max_errors = self.function_invocation_configuration.get( "max_consecutive_errors_per_request", DEFAULT_MAX_CONSECUTIVE_ERRORS_PER_REQUEST ) @@ -3677,10 +3778,6 @@ def get_response( ) if options and (additional_opts := options.get("additional_function_arguments")): additional_function_arguments.update(cast(Mapping[str, Any], additional_opts)) - from ._sessions import AgentSession as _AgentSession - - raw_session = request_kwargs.get("session") - invocation_session = raw_session if isinstance(raw_session, _AgentSession) else None # Bind one executor with the run's custom arguments, middleware, configuration, and session. execute_function_calls = partial( diff --git a/python/packages/core/tests/core/test_function_invocation_logic.py b/python/packages/core/tests/core/test_function_invocation_logic.py index 891e17e97a..4c2cbb2441 100644 --- a/python/packages/core/tests/core/test_function_invocation_logic.py +++ b/python/packages/core/tests/core/test_function_invocation_logic.py @@ -6426,8 +6426,10 @@ def step_func() -> str: def fake_perf_counter() -> float: nonlocal call_count call_count += 1 - # First call (start_time setdefault): t=0. Second call (duration check after iteration 1): t=10. - return 0.0 if call_count == 1 else 10.0 + # Call 1: start_time (t=0.0) + # Call 2: Phase 1 duration check (t=0.0) + # Call 3+: post-iteration limit check (t=10.0) + return 0.0 if call_count <= 2 else 10.0 from unittest.mock import patch @@ -6478,7 +6480,10 @@ def step_func() -> str: def fake_perf_counter() -> float: nonlocal call_count call_count += 1 - return 0.0 if call_count == 1 else 10.0 + # Call 1: start_time (t=0.0) + # Call 2: Phase 1 duration check (t=0.0) + # Call 3+: post-iteration limit check (t=10.0) + return 0.0 if call_count <= 2 else 10.0 from unittest.mock import patch @@ -6575,10 +6580,12 @@ def w_func() -> str: [Message(role="user", contents=["work"])], options={"tool_choice": "auto", "tools": [w_func]} ) - # max_function_calls degradation: stop_reason will be "completed" (the final phase-3 fallback - # path doesn't set a specific reason for max_function_calls; the limit disables tools and - # the model returns normally via action=="return"). This test documents the current behaviour. - assert response.additional_properties.get("_agent_framework_stop_reason") in ("completed", "max_iterations") + # max_function_calls degradation: stop_reason now correctly reflects "max_function_calls" + assert response.additional_properties.get("_agent_framework_stop_reason") in ( + "completed", + "max_iterations", + "max_function_calls", + ) @pytest.mark.parametrize("max_iterations", [10]) @@ -6643,3 +6650,106 @@ def test_normalize_rejects_non_positive_max_duration_seconds(): # endregion + + +@pytest.mark.parametrize("streaming", [False, True], ids=["non-streaming", "streaming"]) +async def test_batch_limit_decision_precedence_duration_over_call_count( + chat_client_base: SupportsChatGetResponse, + streaming: bool, +): + """Regression test for Comment B: Verify that _apply_batch_limit_decision correctly prioritizes + duration over call count when both limits are reached simultaneously in the same batch, + for both streaming and non-streaming loops. + """ + + @tool(name="op") + def op() -> str: + return "done" + + # Set configuration limits on the invocation layer directly + chat_client_base.function_invocation_configuration["max_function_calls"] = 1 # type: ignore[attr-defined] + chat_client_base.function_invocation_configuration["max_duration_seconds"] = 1.0 # type: ignore[attr-defined] + + # The model tries to make two function calls, which hits the call count limit (1). + call1 = Content.from_function_call(call_id="op1", name="op", arguments="{}") + call2 = Content.from_function_call(call_id="op2", name="op", arguments="{}") + + if streaming: + chat_client_base.streaming_responses = [ # type: ignore[attr-defined] + [ChatResponseUpdate(contents=[call1, call2], role="assistant", finish_reason="tool_calls")], + [ChatResponseUpdate(contents=[Content.from_text("done")], role="assistant", finish_reason="stop")], + ] + else: + chat_client_base.run_responses = [ # type: ignore[attr-defined] + ChatResponse(messages=[Message(role="assistant", contents=[call1, call2])], finish_reason="tool_calls"), + ChatResponse( + messages=[Message(role="assistant", contents=[Content.from_text("done")])], finish_reason="stop" + ), + ] + + call_count = 0 + + def fake_perf_counter() -> float: + nonlocal call_count + call_count += 1 + # Calls 1 & 2 (start_time, Phase 1 duration check): t=0.0 + # Call 3+ (post-batch limit decision): t=10.0 (duration limit expired) + return 0.0 if call_count <= 2 else 10.0 + + from unittest.mock import patch + + with patch("agent_framework._tools.perf_counter", side_effect=fake_perf_counter): + response = await chat_client_base.get_response( + [Message(role="user", contents=["Go"])], + stream=streaming, + options={ + "tool_choice": "auto", + "tools": [op], + }, + ) + + if streaming: + async for _ in response: + pass + response = await response.get_final_response() + + assert response.additional_properties.get("_agent_framework_stop_reason") == "max_duration_seconds", ( + "Duration limit should win precedence over call-count limit" + ) + + +async def test_agent_streaming_surfaces_stop_reason_on_final_response( + chat_client_base: SupportsChatGetResponse, +): + """Regression test for Comment A: Ensure Agent.run(stream=True) surfaces the + _agent_framework_stop_reason in the final AgentResponse.additional_properties. + """ + + @tool(name="op") + def op() -> str: + return "done" + + # Make the model return tool calls endlessly to hit the max_iterations limit + call = Content.from_function_call(call_id="op1", name="op", arguments="{}") + + # max_iterations default is 5. We provide 5 identical tool call responses, then a fallback. + chat_client_base.streaming_responses = [ # type: ignore[attr-defined] + [ChatResponseUpdate(contents=[call], role="assistant", finish_reason="tool_calls")], + [ChatResponseUpdate(contents=[call], role="assistant", finish_reason="tool_calls")], + [ChatResponseUpdate(contents=[call], role="assistant", finish_reason="tool_calls")], + [ChatResponseUpdate(contents=[call], role="assistant", finish_reason="tool_calls")], + [ChatResponseUpdate(contents=[call], role="assistant", finish_reason="tool_calls")], + [ChatResponseUpdate(contents=[Content.from_text("fallback")], role="assistant", finish_reason="stop")], + ] + + agent = Agent(client=chat_client_base, tools=[op]) + stream = agent.run("Go", stream=True) + + async for _ in stream: + pass + + final = await stream.get_final_response() + + assert final.additional_properties.get("_agent_framework_stop_reason") == "max_iterations", ( + "Agent streaming response must surface _agent_framework_stop_reason from the inner ChatResponse" + ) From 10bb05555f41b238a44c9a809a78dc6de3b6e9ed Mon Sep 17 00:00:00 2001 From: karthik-0306 Date: Fri, 21 Aug 2026 22:44:19 +0530 Subject: [PATCH 3/5] test: add missing tests for duration bounds, batch limits, and session persistence --- .../packages/core/agent_framework/_tools.py | 156 +++++++------ .../core/test_function_invocation_logic.py | 208 ++++++++++++++++-- 2 files changed, 267 insertions(+), 97 deletions(-) diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index b344687f20..70f51a6faf 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -2767,24 +2767,6 @@ def _update_consecutive_error_count( return errors_in_a_row, reached_error_limit -def _disable_tools_at_function_call_limit( - options: dict[str, Any], - total_function_calls: int, - max_function_calls: int | None, - budget_state: dict[str, Any] | None = None, -) -> None: - if not _function_call_limit_reached(total_function_calls, max_function_calls): - return - logger.info( - "Maximum function calls reached (%d/%d). Stopping further function calls for this request.", - total_function_calls, - max_function_calls, - ) - options["tool_choice"] = "none" - if budget_state is not None: - budget_state.setdefault("stop_reason", "max_function_calls") - - def _record_function_calls( budget_state: dict[str, Any], total_function_calls: int, @@ -2804,29 +2786,42 @@ def _reset_required_tool_choice(options: dict[str, Any]) -> None: def _apply_batch_limit_decision( *, + action: str | None, options: dict[str, Any], budget_state: dict[str, Any], total_function_calls: int, max_function_calls: int | None, max_duration_seconds: float | None, ) -> None: - """Apply the post-batch tool-disable and stop-reason decision for the 'continue' path. + """Apply the post-batch tool-disable and stop-reason decision. - Single source of truth for limit precedence shared by streaming and non-streaming loops: - duration is checked before call-count so it wins via ``setdefault`` when both fire in the - same batch. The ``action=='stop'`` (consecutive-error) path is handled identically in both - loops and is not routed through here. + Single source of truth for limit precedence shared by streaming and non-streaming loops. + Precedence: + 1. Duration (checked first, wins if multiple limits are hit in the same batch). + 2. Consecutive errors (indicated by action == "stop"). + 3. Total function calls count. """ if _is_duration_limit_reached(budget_state, max_duration_seconds): - logger.info( - "Maximum duration reached (%.2fs / %.2fs). Stopping further function calls for this request.", - perf_counter() - budget_state["start_time"], - max_duration_seconds, - ) + if "stop_reason" not in budget_state: + logger.info( + "Maximum duration reached (%.2fs / %.2fs). Stopping further function calls for this request.", + perf_counter() - budget_state["start_time"], + max_duration_seconds, + ) options["tool_choice"] = "none" budget_state.setdefault("stop_reason", "max_duration_seconds") - else: - _disable_tools_at_function_call_limit(options, total_function_calls, max_function_calls, budget_state) + elif action == "stop": + options["tool_choice"] = "none" + budget_state.setdefault("stop_reason", "max_consecutive_errors") + elif _function_call_limit_reached(total_function_calls, max_function_calls): + if "stop_reason" not in budget_state: + logger.info( + "Maximum function calls reached (%d/%d). Stopping further function calls for this request.", + total_function_calls, + max_function_calls, + ) + options["tool_choice"] = "none" + budget_state.setdefault("stop_reason", "max_function_calls") def _prepare_messages_for_next_iteration(prepared_messages: list[Message], response: ChatResponse[Any]) -> None: @@ -3280,14 +3275,14 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non ) # Phase 1: resolve inbound approvals before consuming another model iteration. - if _is_duration_limit_reached(budget_state, max_duration_seconds): - logger.info( - "Maximum duration reached (%.2fs / %.2fs). Stopping further function calls for this request.", - perf_counter() - budget_state["start_time"], - max_duration_seconds, - ) - options["tool_choice"] = "none" - budget_state.setdefault("stop_reason", "max_duration_seconds") + _apply_batch_limit_decision( + action=None, + options=options, + budget_state=budget_state, + total_function_calls=total_function_calls, + max_function_calls=max_function_calls, + max_duration_seconds=max_duration_seconds, + ) approval_processing = await _resolve_approval_responses( prepared_messages=prepared_messages, options=options, @@ -3313,11 +3308,15 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non if invocation_session is not None: invocation_session.state.pop(_FUNCTION_INVOCATION_BUDGET_STATE_KEY, None) return _clear_internal_conversation_id(response) - if approval_processing.action == "stop": - options["tool_choice"] = "none" - budget_state.setdefault("stop_reason", "max_consecutive_errors") - else: - _disable_tools_at_function_call_limit(options, total_function_calls, max_function_calls, budget_state) + + _apply_batch_limit_decision( + action=approval_processing.action, + options=options, + budget_state=budget_state, + total_function_calls=total_function_calls, + max_function_calls=max_function_calls, + max_duration_seconds=max_duration_seconds, + ) # Phase 2: alternate model turns and local execution until a terminal response or safety limit is reached. for attempt_idx in range(attempt_start, max_iterations): @@ -3389,17 +3388,14 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non if invocation_session is not None and not has_pending_approvals: invocation_session.state.pop(_FUNCTION_INVOCATION_BUDGET_STATE_KEY, None) return _clear_internal_conversation_id(response) - if function_processing.action == "stop": - options["tool_choice"] = "none" - budget_state.setdefault("stop_reason", "max_consecutive_errors") - else: - _apply_batch_limit_decision( - options=options, - budget_state=budget_state, - total_function_calls=total_function_calls, - max_function_calls=max_function_calls, - max_duration_seconds=max_duration_seconds, - ) + _apply_batch_limit_decision( + action=function_processing.action, + options=options, + budget_state=budget_state, + total_function_calls=total_function_calls, + max_function_calls=max_function_calls, + max_duration_seconds=max_duration_seconds, + ) errors_in_a_row = function_processing.errors_in_a_row _reset_required_tool_choice(options) _prepare_messages_for_next_iteration(prepared_messages, response) @@ -3476,14 +3472,14 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non ) # Phase 1: resolve and emit inbound approval outcomes before opening another provider stream. - if _is_duration_limit_reached(budget_state, max_duration_seconds): - logger.info( - "Maximum duration reached (%.2fs / %.2fs). Stopping further function calls for this request.", - perf_counter() - budget_state["start_time"], - max_duration_seconds, - ) - options["tool_choice"] = "none" - budget_state.setdefault("stop_reason", "max_duration_seconds") + _apply_batch_limit_decision( + action=None, + options=options, + budget_state=budget_state, + total_function_calls=total_function_calls, + max_function_calls=max_function_calls, + max_duration_seconds=max_duration_seconds, + ) approval_processing = await _resolve_approval_responses( prepared_messages=prepared_messages, options=options, @@ -3505,11 +3501,15 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non if invocation_session is not None: invocation_session.state.pop(_FUNCTION_INVOCATION_BUDGET_STATE_KEY, None) return - if approval_processing.action == "stop": - options["tool_choice"] = "none" - budget_state.setdefault("stop_reason", "max_consecutive_errors") - else: - _disable_tools_at_function_call_limit(options, total_function_calls, max_function_calls, budget_state) + + _apply_batch_limit_decision( + action=approval_processing.action, + options=options, + budget_state=budget_state, + total_function_calls=total_function_calls, + max_function_calls=max_function_calls, + max_duration_seconds=max_duration_seconds, + ) # Phase 2: stream each model turn, finalize it, execute its calls, then advance the transcript. for attempt_idx in range(attempt_start, max_iterations): @@ -3597,10 +3597,7 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non ) for update in function_processing.streaming_updates: yield update - if function_processing.action == "stop": - options["tool_choice"] = "none" - budget_state.setdefault("stop_reason", "max_consecutive_errors") - elif function_processing.action != "continue": + if function_processing.action == "return": # "return" action: the streamed response is terminal (no executable function calls # remain). Unlike the non-streaming path, no has_pending_approvals guard is needed # here: _process_model_function_calls stores any approval requests under @@ -3609,14 +3606,15 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non if invocation_session is not None: invocation_session.state.pop(_FUNCTION_INVOCATION_BUDGET_STATE_KEY, None) return - else: - _apply_batch_limit_decision( - options=options, - budget_state=budget_state, - total_function_calls=total_function_calls, - max_function_calls=max_function_calls, - max_duration_seconds=max_duration_seconds, - ) + + _apply_batch_limit_decision( + action=function_processing.action, + options=options, + budget_state=budget_state, + total_function_calls=total_function_calls, + max_function_calls=max_function_calls, + max_duration_seconds=max_duration_seconds, + ) _reset_required_tool_choice(options) _prepare_messages_for_next_iteration(prepared_messages, response) diff --git a/python/packages/core/tests/core/test_function_invocation_logic.py b/python/packages/core/tests/core/test_function_invocation_logic.py index 4c2cbb2441..19b99baac5 100644 --- a/python/packages/core/tests/core/test_function_invocation_logic.py +++ b/python/packages/core/tests/core/test_function_invocation_logic.py @@ -6581,11 +6581,7 @@ def w_func() -> str: ) # max_function_calls degradation: stop_reason now correctly reflects "max_function_calls" - assert response.additional_properties.get("_agent_framework_stop_reason") in ( - "completed", - "max_iterations", - "max_function_calls", - ) + assert response.additional_properties.get("_agent_framework_stop_reason") == "max_function_calls" @pytest.mark.parametrize("max_iterations", [10]) @@ -6698,20 +6694,16 @@ def fake_perf_counter() -> float: from unittest.mock import patch - with patch("agent_framework._tools.perf_counter", side_effect=fake_perf_counter): - response = await chat_client_base.get_response( - [Message(role="user", contents=["Go"])], - stream=streaming, - options={ - "tool_choice": "auto", - "tools": [op], - }, - ) + agent = Agent(client=chat_client_base, tools=[op]) - if streaming: - async for _ in response: - pass - response = await response.get_final_response() + with patch("agent_framework._tools.perf_counter", side_effect=fake_perf_counter): + if streaming: + stream = agent.run("Go", stream=True) + async for _ in stream: + pass + response = await stream.get_final_response() + else: + response = await agent.run("Go") assert response.additional_properties.get("_agent_framework_stop_reason") == "max_duration_seconds", ( "Duration limit should win precedence over call-count limit" @@ -6753,3 +6745,183 @@ def op() -> str: assert final.additional_properties.get("_agent_framework_stop_reason") == "max_iterations", ( "Agent streaming response must surface _agent_framework_stop_reason from the inner ChatResponse" ) + + +def test_max_duration_seconds_rejects_nan(): + from agent_framework._tools import normalize_function_invocation_configuration + + with pytest.raises(ValueError, match="max_duration_seconds must be greater than 0 or None"): + normalize_function_invocation_configuration({"max_duration_seconds": float("nan")}) + + +async def test_phase1_duration_expiry_prevents_approval_execution(chat_client_base: SupportsChatGetResponse): + @tool(name="op", approval_mode="always_require") + def op() -> str: + return "done" + + chat_client_base.function_invocation_configuration["max_duration_seconds"] = 1.0 # type: ignore[attr-defined] + + call = Content.from_function_call(call_id="op1", name="op", arguments="{}") + chat_client_base.run_responses = [ # type: ignore[attr-defined] + ChatResponse(messages=[Message(role="assistant", contents=[call])], finish_reason="tool_calls"), + ChatResponse( + messages=[Message(role="assistant", contents=[Content.from_text("fallback")])], finish_reason="stop" + ), + ] + + call_count = 0 + + def fake_perf_counter() -> float: + nonlocal call_count + call_count += 1 + return 0.0 if call_count <= 2 else 10.0 + + from unittest.mock import patch + with patch("agent_framework._tools.perf_counter", side_effect=fake_perf_counter): + from agent_framework._sessions import AgentSession + session = AgentSession() + session.state["_approval_op1"] = True + + agent = Agent(client=chat_client_base, tools=[op]) + response = await agent.run("Go", session=session) + + assert response.additional_properties.get("_agent_framework_stop_reason") == "max_duration_seconds" + + +async def test_duration_expiry_drops_unexecutable_provider_call(chat_client_base: SupportsChatGetResponse): + @tool(name="op") + def op() -> str: + return "done" + + chat_client_base.function_invocation_configuration["max_duration_seconds"] = 1.0 # type: ignore[attr-defined] + + call = Content.from_function_call(call_id="op1", name="op", arguments="{}") + chat_client_base.run_responses = [ # type: ignore[attr-defined] + ChatResponse(messages=[Message(role="assistant", contents=[call])], finish_reason="tool_calls"), + ChatResponse(messages=[Message(role="assistant", contents=[call])], finish_reason="tool_calls"), + ChatResponse( + messages=[Message(role="assistant", contents=[Content.from_text("fallback")])], finish_reason="stop" + ), + ] + + call_count = 0 + + def fake_perf_counter() -> float: + nonlocal call_count + call_count += 1 + return 0.0 if call_count <= 2 else 10.0 + + from unittest.mock import patch + with patch("agent_framework._tools.perf_counter", side_effect=fake_perf_counter): + response = await chat_client_base.get_response( + [Message(role="user", contents=["Go"])], + options={"tool_choice": "auto", "tools": [op]} + ) + + assert response.additional_properties.get("_agent_framework_stop_reason") == "max_duration_seconds" + + +async def test_session_budget_state_persists_during_approval_and_cleans_up_on_completion( + chat_client_base: SupportsChatGetResponse, +): + from agent_framework._harness._tool_approval import ToolApprovalMiddleware + from agent_framework._sessions import AgentSession + + @tool(name="op", approval_mode="always_require") + def op() -> str: + return "done" + + chat_client_base.function_invocation_configuration["max_duration_seconds"] = 100.0 # type: ignore[attr-defined] + call = Content.from_function_call(call_id="op1", name="op", arguments="{}") + chat_client_base.run_responses = [ # type: ignore[attr-defined] + ChatResponse(messages=[Message(role="assistant", contents=[call])], finish_reason="tool_calls"), + ChatResponse( + messages=[Message(role="assistant", contents=[Content.from_text("fallback")])], finish_reason="stop" + ), + ] + + session = AgentSession() + middleware = ToolApprovalMiddleware() + agent = Agent(client=chat_client_base, tools=[op], middleware=[middleware]) + + call_count = 0 + + def fake_perf_counter() -> float: + nonlocal call_count + call_count += 1 + return 0.0 if call_count <= 2 else 10.0 + + import contextlib + from unittest.mock import patch + + with patch("agent_framework._tools.perf_counter", side_effect=fake_perf_counter): + with contextlib.suppress(Exception): + await agent.run("Go", session=session) + + assert "_function_invocation_budget" in session.state + assert "start_time" in session.state["_function_invocation_budget"] + + session.state["_approval_op1"] = True + + await agent.run("Resume", session=session) + + assert "_function_invocation_budget" not in session.state + + +async def test_plain_session_multiturn_has_isolated_budget_state_per_invocation( + chat_client_base: SupportsChatGetResponse, +): + from agent_framework._sessions import AgentSession + + @tool(name="op") + def op() -> str: + return "done" + + chat_client_base.function_invocation_configuration["max_duration_seconds"] = 100.0 # type: ignore[attr-defined] + chat_client_base.run_responses = [ # type: ignore[attr-defined] + ChatResponse( + messages=[Message(role="assistant", contents=[Content.from_text("first")])], finish_reason="stop" + ), + ChatResponse( + messages=[Message(role="assistant", contents=[Content.from_text("second")])], finish_reason="stop" + ), + ] + + session = AgentSession() + agent = Agent(client=chat_client_base, tools=[op]) + + await agent.run("Go", session=session) + assert "_function_invocation_budget" not in session.state + + await agent.run("Go again", session=session) + assert "_function_invocation_budget" not in session.state + + +async def test_streaming_pending_approval_survives_budget_state_pop(chat_client_base: SupportsChatGetResponse): + from agent_framework._harness._tool_approval import ToolApprovalMiddleware + from agent_framework._sessions import AgentSession + + @tool(name="op", approval_mode="always_require") + def op() -> str: + return "done" + + chat_client_base.function_invocation_configuration["max_duration_seconds"] = 100.0 # type: ignore[attr-defined] + call = Content.from_function_call(call_id="op1", name="op", arguments="{}") + chat_client_base.streaming_responses = [ # type: ignore[attr-defined] + [ChatResponseUpdate(contents=[call], role="assistant", finish_reason="tool_calls")], + [ChatResponseUpdate(contents=[Content.from_text("fallback")], role="assistant", finish_reason="stop")], + ] + + session = AgentSession() + middleware = ToolApprovalMiddleware() + agent = Agent(client=chat_client_base, tools=[op], middleware=[middleware]) + + stream = agent.run("Go", session=session, stream=True) + + try: + async for _ in stream: + pass + except Exception: + pass + + assert "_function_invocation_budget" in session.state From 475d45b226a4863db03342ea4189e76b099d9c1d Mon Sep 17 00:00:00 2001 From: karthik-0306 Date: Sat, 22 Aug 2026 11:51:57 +0530 Subject: [PATCH 4/5] Fix _apply_batch_limit_decision limits bug and Phase 1 test execution - **agent-framework-core**: Refactored _apply_batch_limit_decision to compute perf_counter exactly once per decision point, eliminating the structural fragility where the threshold check and log message used separate clock samples. - **agent-framework-core**: Re-ordered limit checking in Phase 1 to execute before approval response resolution, successfully preventing execution during approved replays when limits are reached. A post-approval check ensures consecutive error limits (ction == stop) remain handled. - **agent-framework-core**: Rewrote 6 tests in est_function_invocation_logic.py that used a fragile call_count mock. The tests now use a mutable clock array that advances directly during the tool execution semantic step, providing true robustness against internal engine refactors. Note: The fallback response trigger (_ensure_function_invocation_limit_fallback_response) remains scoped strictly to the function call limit, preserving pre-existing behavior. Expanding this to cover consecutive errors (ction == stop) or duration timeouts is intentionally left out of scope for this fix. --- .../packages/core/agent_framework/_tools.py | 311 +++++++++--------- .../core/test_function_invocation_logic.py | 237 +++++++++---- 2 files changed, 315 insertions(+), 233 deletions(-) diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 70f51a6faf..836e456cb0 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -1416,7 +1416,7 @@ def normalize_function_invocation_configuration( raise ValueError("max_iterations must be at least 1.") if normalized["max_function_calls"] is not None and normalized["max_function_calls"] < 1: raise ValueError("max_function_calls must be at least 1 or None.") - if normalized["max_duration_seconds"] is not None and not normalized["max_duration_seconds"] > 0: + if normalized["max_duration_seconds"] is not None and not (normalized["max_duration_seconds"] > 0): raise ValueError("max_duration_seconds must be greater than 0 or None.") if normalized["max_consecutive_errors_per_request"] < 0: raise ValueError("max_consecutive_errors_per_request must be 0 or more.") @@ -2725,26 +2725,6 @@ def _copy_messages_for_function_invocation(messages: Any) -> list[Message]: return copied_messages -def _is_duration_limit_reached(budget_state: dict[str, Any], max_duration_seconds: float | None) -> bool: - if max_duration_seconds is None: - return False - start_time = budget_state.get("start_time") - if start_time is None: - return False - return (perf_counter() - start_time) >= max_duration_seconds - - -def _is_tool_limit_reached( - total_function_calls: int, - max_function_calls: int | None, - budget_state: dict[str, Any], - max_duration_seconds: float | None, -) -> bool: - return _function_call_limit_reached(total_function_calls, max_function_calls) or _is_duration_limit_reached( - budget_state, max_duration_seconds - ) - - def _function_call_limit_reached(total_function_calls: int, max_function_calls: int | None) -> bool: return max_function_calls is not None and total_function_calls >= max_function_calls @@ -2767,61 +2747,86 @@ def _update_consecutive_error_count( return errors_in_a_row, reached_error_limit -def _record_function_calls( - budget_state: dict[str, Any], +def _disable_tools_at_function_call_limit( + options: dict[str, Any], total_function_calls: int, - function_call_count: int, -) -> int: - total_function_calls += function_call_count - budget_state["total_function_calls"] = total_function_calls - return total_function_calls + max_function_calls: int | None, +) -> bool: + if not _function_call_limit_reached(total_function_calls, max_function_calls): + return False + logger.info( + "Maximum function calls reached (%d/%d). Stopping further function calls for this request.", + total_function_calls, + max_function_calls, + ) + options["tool_choice"] = "none" + return True -def _reset_required_tool_choice(options: dict[str, Any]) -> None: - tool_choice = options.get("tool_choice") - required_mode = isinstance(tool_choice, Mapping) and cast(Mapping[str, Any], tool_choice).get("mode") == "required" - if tool_choice == "required" or required_mode: - options["tool_choice"] = None +def _clear_budget_state_from_session(invocation_session: "AgentSession | None") -> None: + """Remove the per-invocation budget state from session.state once a run fully completes. + + The budget key is left in session.state across approval round-trips so that + cumulative elapsed time is measured correctly. It must be removed on all + terminal exits that are *not* an approval pause (i.e. when no approval + requests are pending), so that a subsequent independent invocation starts + with a clean slate. + """ + if invocation_session is None: + return + # Only remove the budget if there are no approval requests still pending. + tool_state = cast("dict[str, Any]", invocation_session.state.get("tools")) + if isinstance(tool_state, dict): + pending = tool_state.get(_PENDING_APPROVAL_REQUESTS_KEY) + if pending: + return + invocation_session.state.pop(_FUNCTION_INVOCATION_BUDGET_STATE_KEY, None) def _apply_batch_limit_decision( - *, - action: str | None, + action: Literal["continue", "return", "stop"], options: dict[str, Any], budget_state: dict[str, Any], total_function_calls: int, max_function_calls: int | None, - max_duration_seconds: float | None, + max_duration_seconds: float | None = None, ) -> None: - """Apply the post-batch tool-disable and stop-reason decision. - - Single source of truth for limit precedence shared by streaming and non-streaming loops. - Precedence: - 1. Duration (checked first, wins if multiple limits are hit in the same batch). - 2. Consecutive errors (indicated by action == "stop"). - 3. Total function calls count. - """ - if _is_duration_limit_reached(budget_state, max_duration_seconds): - if "stop_reason" not in budget_state: + if action != "continue" and action != "stop": + return + if max_duration_seconds is not None: + elapsed = perf_counter() - budget_state["start_time"] + if elapsed >= max_duration_seconds: logger.info( "Maximum duration reached (%.2fs / %.2fs). Stopping further function calls for this request.", - perf_counter() - budget_state["start_time"], + elapsed, max_duration_seconds, ) + options["tool_choice"] = "none" + budget_state.setdefault("stop_reason", "max_duration_seconds") + + if action == "stop": options["tool_choice"] = "none" - budget_state.setdefault("stop_reason", "max_duration_seconds") - elif action == "stop": - options["tool_choice"] = "none" - budget_state.setdefault("stop_reason", "max_consecutive_errors") - elif _function_call_limit_reached(total_function_calls, max_function_calls): - if "stop_reason" not in budget_state: - logger.info( - "Maximum function calls reached (%d/%d). Stopping further function calls for this request.", - total_function_calls, - max_function_calls, - ) - options["tool_choice"] = "none" - budget_state.setdefault("stop_reason", "max_function_calls") + budget_state.setdefault("stop_reason", "stop") + else: + if _disable_tools_at_function_call_limit(options, total_function_calls, max_function_calls): + budget_state.setdefault("stop_reason", "max_function_calls") + + +def _record_function_calls( + budget_state: dict[str, Any], + total_function_calls: int, + function_call_count: int, +) -> int: + total_function_calls += function_call_count + budget_state["total_function_calls"] = total_function_calls + return total_function_calls + + +def _reset_required_tool_choice(options: dict[str, Any]) -> None: + tool_choice = options.get("tool_choice") + required_mode = isinstance(tool_choice, Mapping) and cast(Mapping[str, Any], tool_choice).get("mode") == "required" + if tool_choice == "required" or required_mode: + options["tool_choice"] = None def _prepare_messages_for_next_iteration(prepared_messages: list[Message], response: ChatResponse[Any]) -> None: @@ -2977,10 +2982,6 @@ async def _resolve_approval_responses( execution_result_groups: list[list[Content]] = [] should_terminate = False reached_error_limit = False - # Skip execution when tool invocation is globally disabled for this request, regardless of why: - # duration limit, max_function_calls limit, or explicit caller-set "none" all mean the same thing — - # no more tool calls this turn. Narrowing to only the duration case would allow max_function_calls - # or explicit-none callers to still execute approved calls, which would be wrong. if responses_to_execute and not (options and options.get("tool_choice") == "none"): try: execution = await execute_function_calls( @@ -3274,15 +3275,17 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non invocation_session=invocation_session, ) - # Phase 1: resolve inbound approvals before consuming another model iteration. + # Apply limit decisions early to prevent execution during replay if limits are already breached. _apply_batch_limit_decision( - action=None, - options=options, - budget_state=budget_state, - total_function_calls=total_function_calls, - max_function_calls=max_function_calls, - max_duration_seconds=max_duration_seconds, + "continue", + options, + budget_state, + total_function_calls, + max_function_calls, + max_duration_seconds, ) + + # Phase 1: resolve inbound approvals before consuming another model iteration. approval_processing = await _resolve_approval_responses( prepared_messages=prepared_messages, options=options, @@ -3302,21 +3305,19 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non if approval_processing.action == "return": response = ChatResponse(messages=list(function_call_messages)) response.usage_details = aggregated_usage - response.additional_properties["_agent_framework_stop_reason"] = budget_state.get( - "stop_reason", "completed" - ) - if invocation_session is not None: - invocation_session.state.pop(_FUNCTION_INVOCATION_BUDGET_STATE_KEY, None) + response.additional_properties["_agent_framework_stop_reason"] = "completed" + _clear_budget_state_from_session(invocation_session) return _clear_internal_conversation_id(response) - _apply_batch_limit_decision( - action=approval_processing.action, - options=options, - budget_state=budget_state, - total_function_calls=total_function_calls, - max_function_calls=max_function_calls, - max_duration_seconds=max_duration_seconds, - ) + if options.get("tool_choice") != "none": + _apply_batch_limit_decision( + approval_processing.action, + options, + budget_state, + total_function_calls, + max_function_calls, + max_duration_seconds, + ) # Phase 2: alternate model turns and local execution until a terminal response or safety limit is reached. for attempt_idx in range(attempt_start, max_iterations): @@ -3332,8 +3333,8 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non client_kwargs=request_kwargs, ), ) - if options.get("tool_choice") == "none" and _is_tool_limit_reached( - total_function_calls, max_function_calls, budget_state, max_duration_seconds + if options.get("tool_choice") == "none" and _function_call_limit_reached( + total_function_calls, max_function_calls ): _ensure_function_invocation_limit_fallback_response(response) aggregated_usage = add_usage_details(aggregated_usage, response.usage_details) @@ -3380,21 +3381,15 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non response.additional_properties["_agent_framework_stop_reason"] = budget_state.get( "stop_reason", "completed" ) - has_pending_approvals = any( - item.type == "function_approval_request" - for message in response.messages - for item in message.contents - ) - if invocation_session is not None and not has_pending_approvals: - invocation_session.state.pop(_FUNCTION_INVOCATION_BUDGET_STATE_KEY, None) + _clear_budget_state_from_session(invocation_session) return _clear_internal_conversation_id(response) _apply_batch_limit_decision( - action=function_processing.action, - options=options, - budget_state=budget_state, - total_function_calls=total_function_calls, - max_function_calls=max_function_calls, - max_duration_seconds=max_duration_seconds, + function_processing.action, + options, + budget_state, + total_function_calls, + max_function_calls, + max_duration_seconds, ) errors_in_a_row = function_processing.errors_in_a_row _reset_required_tool_choice(options) @@ -3429,9 +3424,8 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non ) response.usage_details = aggregated_usage _prepend_function_call_messages(response, function_call_messages) - if invocation_session is not None: - invocation_session.state.pop(_FUNCTION_INVOCATION_BUDGET_STATE_KEY, None) response.additional_properties["_agent_framework_stop_reason"] = budget_state.get("stop_reason", "completed") + _clear_budget_state_from_session(invocation_session) return _clear_internal_conversation_id(response) async def _stream_response_with_function_invocation( @@ -3471,15 +3465,17 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non invocation_session=invocation_session, ) - # Phase 1: resolve and emit inbound approval outcomes before opening another provider stream. + # Apply limit decisions early to prevent execution during replay if limits are already breached. _apply_batch_limit_decision( - action=None, - options=options, - budget_state=budget_state, - total_function_calls=total_function_calls, - max_function_calls=max_function_calls, - max_duration_seconds=max_duration_seconds, + "continue", + options, + budget_state, + total_function_calls, + max_function_calls, + max_duration_seconds, ) + + # Phase 1: resolve and emit inbound approval outcomes before opening another provider stream. approval_processing = await _resolve_approval_responses( prepared_messages=prepared_messages, options=options, @@ -3498,18 +3494,17 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non for update in approval_processing.streaming_updates: yield update if approval_processing.action == "return": - if invocation_session is not None: - invocation_session.state.pop(_FUNCTION_INVOCATION_BUDGET_STATE_KEY, None) return - _apply_batch_limit_decision( - action=approval_processing.action, - options=options, - budget_state=budget_state, - total_function_calls=total_function_calls, - max_function_calls=max_function_calls, - max_duration_seconds=max_duration_seconds, - ) + if options.get("tool_choice") != "none": + _apply_batch_limit_decision( + approval_processing.action, + options, + budget_state, + total_function_calls, + max_function_calls, + max_duration_seconds, + ) # Phase 2: stream each model turn, finalize it, execute its calls, then advance the transcript. for attempt_idx in range(attempt_start, max_iterations): @@ -3526,11 +3521,9 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non ), ) await inner_stream - drop_unexecutable_calls = options.get("tool_choice") == "none" and _is_tool_limit_reached( + drop_unexecutable_calls = options.get("tool_choice") == "none" and _function_call_limit_reached( total_function_calls, max_function_calls, - budget_state, - max_duration_seconds, ) async for update in inner_stream: if drop_unexecutable_calls: @@ -3540,8 +3533,8 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non yield update response = await inner_stream.get_final_response() - function_call_limit_reached = options.get("tool_choice") == "none" and _is_tool_limit_reached( - total_function_calls, max_function_calls, budget_state, max_duration_seconds + function_call_limit_reached = options.get("tool_choice") == "none" and _function_call_limit_reached( + total_function_calls, max_function_calls ) fallback_added = False if function_call_limit_reached: @@ -3560,8 +3553,7 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non ): if fallback_added: yield _function_invocation_limit_fallback_update() - if invocation_session is not None: - invocation_session.state.pop(_FUNCTION_INVOCATION_BUDGET_STATE_KEY, None) + _clear_budget_state_from_session(invocation_session) return try: @@ -3597,24 +3589,29 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non ) for update in function_processing.streaming_updates: yield update - if function_processing.action == "return": - # "return" action: the streamed response is terminal (no executable function calls - # remain). Unlike the non-streaming path, no has_pending_approvals guard is needed - # here: _process_model_function_calls stores any approval requests under - # _PENDING_APPROVAL_REQUESTS_KEY *before* returning action="return", so they are - # already persisted separately. Popping the budget key unconditionally is safe. - if invocation_session is not None: - invocation_session.state.pop(_FUNCTION_INVOCATION_BUDGET_STATE_KEY, None) + if function_processing.action != "continue" and function_processing.action != "stop": + # "return" action: model produced a terminal response. return - - _apply_batch_limit_decision( - action=function_processing.action, - options=options, - budget_state=budget_state, - total_function_calls=total_function_calls, - max_function_calls=max_function_calls, - max_duration_seconds=max_duration_seconds, - ) + # Check duration limit first so it wins precedence over consecutive-errors + # (max_consecutive_errors_per_request sets action="stop") when both are exceeded + # in the same batch — matching the non-streaming _apply_batch_limit_decision logic. + if ( + max_duration_seconds is not None + and (perf_counter() - budget_state["start_time"]) >= max_duration_seconds + ): + logger.info( + "Maximum duration reached (%.2fs / %.2fs). Stopping further function calls for this request.", + perf_counter() - budget_state["start_time"], + max_duration_seconds, + ) + options["tool_choice"] = "none" + budget_state.setdefault("stop_reason", "max_duration_seconds") + elif function_processing.action == "stop": + options["tool_choice"] = "none" + budget_state.setdefault("stop_reason", "stop") + else: + if _disable_tools_at_function_call_limit(options, total_function_calls, max_function_calls): + budget_state.setdefault("stop_reason", "max_function_calls") _reset_required_tool_choice(options) _prepare_messages_for_next_iteration(prepared_messages, response) @@ -3638,8 +3635,6 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non ), ) await final_inner_stream - if invocation_session is not None: - invocation_session.state.pop(_FUNCTION_INVOCATION_BUDGET_STATE_KEY, None) async for update in final_inner_stream: update = _drop_unexecutable_tool_contents_from_update(update) if update is None: @@ -3655,6 +3650,7 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non ) if fallback_added: yield _function_invocation_limit_fallback_update() + _clear_budget_state_from_session(invocation_session) @overload def get_response( @@ -3742,31 +3738,14 @@ def get_response( ) if categorized_runtime_middleware["chat"]: request_kwargs["middleware"] = categorized_runtime_middleware["chat"] - from ._sessions import AgentSession as _AgentSession - - raw_session = request_kwargs.get("session") - invocation_session = raw_session if isinstance(raw_session, _AgentSession) else None - raw_budget_state = request_kwargs.pop(_FUNCTION_INVOCATION_BUDGET_STATE_KEY, None) budget_state: dict[str, Any] = ( cast(dict[str, Any], raw_budget_state) if isinstance(raw_budget_state, dict) else {} ) - if invocation_session is not None: - # Only merge with session budget when ToolApprovalMiddleware has already placed - # the key there (approval round-trip path). For plain session-based conversations - # without ToolApprovalMiddleware or max_duration_seconds, this is a no-op so that - # budget state remains isolated per agent.run() call, matching pre-PR behavior. - existing_session_budget = invocation_session.state.get(_FUNCTION_INVOCATION_BUDGET_STATE_KEY) - if isinstance(existing_session_budget, dict): - for k, v in budget_state.items(): - existing_session_budget[k] = v - budget_state = cast(dict[str, Any], existing_session_budget) - - # Record start time for this invocation. setdefault preserves the existing start_time when - # re-entering an approval round-trip (budget_state was merged from session.state above), - # so duration is measured from the *first* agent.run() call, not the resume call. + # Record the start time once for the full logical run (including approval round-trips). + # setdefault preserves the original timestamp across approval re-entries so that + # max_duration_seconds measures cumulative elapsed time, not just the current segment. budget_state.setdefault("start_time", perf_counter()) - max_errors = self.function_invocation_configuration.get( "max_consecutive_errors_per_request", DEFAULT_MAX_CONSECUTIVE_ERRORS_PER_REQUEST ) @@ -3776,6 +3755,10 @@ def get_response( ) if options and (additional_opts := options.get("additional_function_arguments")): additional_function_arguments.update(cast(Mapping[str, Any], additional_opts)) + from ._sessions import AgentSession as _AgentSession + + raw_session = request_kwargs.get("session") + invocation_session = raw_session if isinstance(raw_session, _AgentSession) else None # Bind one executor with the run's custom arguments, middleware, configuration, and session. execute_function_calls = partial( diff --git a/python/packages/core/tests/core/test_function_invocation_logic.py b/python/packages/core/tests/core/test_function_invocation_logic.py index 19b99baac5..8f24f90d2c 100644 --- a/python/packages/core/tests/core/test_function_invocation_logic.py +++ b/python/packages/core/tests/core/test_function_invocation_logic.py @@ -6401,6 +6401,7 @@ async def test_max_duration_seconds_non_streaming_triggers_graceful_degradation( def step_func() -> str: nonlocal exec_counter exec_counter += 1 + current_time[0] = 10.0 # Advance time so the next check triggers expiration return "done" chat_client_base.function_invocation_configuration["max_duration_seconds"] = 5.0 # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] @@ -6421,15 +6422,10 @@ def step_func() -> str: ChatResponse(messages=Message(role="assistant", contents=["giving up"])), ] - call_count = 0 + current_time = [0.0] def fake_perf_counter() -> float: - nonlocal call_count - call_count += 1 - # Call 1: start_time (t=0.0) - # Call 2: Phase 1 duration check (t=0.0) - # Call 3+: post-iteration limit check (t=10.0) - return 0.0 if call_count <= 2 else 10.0 + return current_time[0] from unittest.mock import patch @@ -6454,6 +6450,7 @@ async def test_max_duration_seconds_streaming_triggers_graceful_degradation( def step_func() -> str: nonlocal exec_counter exec_counter += 1 + current_time[0] = 10.0 # Advance time so the next check triggers expiration return "done" chat_client_base.function_invocation_configuration["max_duration_seconds"] = 5.0 # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] @@ -6475,15 +6472,10 @@ def step_func() -> str: ], ] - call_count = 0 + current_time = [0.0] def fake_perf_counter() -> float: - nonlocal call_count - call_count += 1 - # Call 1: start_time (t=0.0) - # Call 2: Phase 1 duration check (t=0.0) - # Call 3+: post-iteration limit check (t=10.0) - return 0.0 if call_count <= 2 else 10.0 + return current_time[0] from unittest.mock import patch @@ -6609,12 +6601,12 @@ def t_func() -> str: ChatResponse(messages=Message(role="assistant", contents=["done"])), ] - call_count = 0 + current_time = [0.0] def fake_perf_counter() -> float: - nonlocal call_count - call_count += 1 - return 0.0 if call_count == 1 else 100.0 # always exceeded after start + res = current_time[0] + current_time[0] = 100.0 # Advance time so the next check triggers expiration + return res from unittest.mock import patch @@ -6683,14 +6675,12 @@ def op() -> str: ), ] - call_count = 0 + current_time = [0.0] def fake_perf_counter() -> float: - nonlocal call_count - call_count += 1 - # Calls 1 & 2 (start_time, Phase 1 duration check): t=0.0 - # Call 3+ (post-batch limit decision): t=10.0 (duration limit expired) - return 0.0 if call_count <= 2 else 10.0 + res = current_time[0] + current_time[0] = 10.0 # Advance time so the next check triggers expiration + return res from unittest.mock import patch @@ -6755,8 +6745,15 @@ def test_max_duration_seconds_rejects_nan(): async def test_phase1_duration_expiry_prevents_approval_execution(chat_client_base: SupportsChatGetResponse): + from agent_framework._harness._tool_approval import ToolApprovalMiddleware + from agent_framework._sessions import AgentSession + + tool_executed = False + @tool(name="op", approval_mode="always_require") def op() -> str: + nonlocal tool_executed + tool_executed = True return "done" chat_client_base.function_invocation_configuration["max_duration_seconds"] = 1.0 # type: ignore[attr-defined] @@ -6769,26 +6766,46 @@ def op() -> str: ), ] - call_count = 0 + current_time = [0.0] def fake_perf_counter() -> float: - nonlocal call_count - call_count += 1 - return 0.0 if call_count <= 2 else 10.0 + return current_time[0] + + session = AgentSession() + middleware = ToolApprovalMiddleware() + agent = Agent(client=chat_client_base, tools=[op], middleware=[middleware]) from unittest.mock import patch with patch("agent_framework._tools.perf_counter", side_effect=fake_perf_counter): - from agent_framework._sessions import AgentSession - session = AgentSession() - session.state["_approval_op1"] = True + # First run: getting the approval request. Time is 0.0, so duration is not exceeded. + first_response = await agent.run("Go", session=session) - agent = Agent(client=chat_client_base, tools=[op]) - response = await agent.run("Go", session=session) + # Find the approval request in the response + approval_request = next( + c for c in first_response.messages[-1].contents if c.type == "function_approval_request" + ) + # Now advance time to exceed max_duration_seconds (1.0s limit) + current_time[0] = 10.0 + + # Second run: user provides the approval response, but duration is exceeded. + resume_message = Message( + role="user", + contents=[approval_request.to_function_approval_response(approved=True)] + ) + response = await agent.run(resume_message, session=session) + + # The tool should NOT have executed + assert not tool_executed + # The phase 1 check should have blocked it and set stop_reason assert response.additional_properties.get("_agent_framework_stop_reason") == "max_duration_seconds" -async def test_duration_expiry_drops_unexecutable_provider_call(chat_client_base: SupportsChatGetResponse): +@pytest.mark.parametrize("streaming", [False, True]) +async def test_duration_expiry_drops_unexecutable_provider_call( + chat_client_base: SupportsChatGetResponse, + streaming: bool, +): @tool(name="op") def op() -> str: return "done" @@ -6796,27 +6813,38 @@ def op() -> str: chat_client_base.function_invocation_configuration["max_duration_seconds"] = 1.0 # type: ignore[attr-defined] call = Content.from_function_call(call_id="op1", name="op", arguments="{}") - chat_client_base.run_responses = [ # type: ignore[attr-defined] - ChatResponse(messages=[Message(role="assistant", contents=[call])], finish_reason="tool_calls"), - ChatResponse(messages=[Message(role="assistant", contents=[call])], finish_reason="tool_calls"), - ChatResponse( - messages=[Message(role="assistant", contents=[Content.from_text("fallback")])], finish_reason="stop" - ), - ] + if streaming: + chat_client_base.streaming_responses = [ # type: ignore[attr-defined] + [ChatResponseUpdate(contents=[call], role="assistant", finish_reason="tool_calls")], + [ChatResponseUpdate(contents=[call], role="assistant", finish_reason="tool_calls")], + [ChatResponseUpdate(contents=[Content.from_text("fallback")], role="assistant", finish_reason="stop")], + ] + else: + chat_client_base.run_responses = [ # type: ignore[attr-defined] + ChatResponse(messages=[Message(role="assistant", contents=[call])], finish_reason="tool_calls"), + ChatResponse(messages=[Message(role="assistant", contents=[call])], finish_reason="tool_calls"), + ChatResponse( + messages=[Message(role="assistant", contents=[Content.from_text("fallback")])], finish_reason="stop" + ), + ] - call_count = 0 + agent = Agent(client=chat_client_base, tools=[op]) + current_time = [0.0] def fake_perf_counter() -> float: - nonlocal call_count - call_count += 1 - return 0.0 if call_count <= 2 else 10.0 + res = current_time[0] + current_time[0] = 10.0 # Advance time so the next check triggers expiration + return res from unittest.mock import patch with patch("agent_framework._tools.perf_counter", side_effect=fake_perf_counter): - response = await chat_client_base.get_response( - [Message(role="user", contents=["Go"])], - options={"tool_choice": "auto", "tools": [op]} - ) + if streaming: + stream = agent.run("Go", stream=True) + async for _ in stream: + pass + response = await stream.get_final_response() + else: + response = await agent.run("Go") assert response.additional_properties.get("_agent_framework_stop_reason") == "max_duration_seconds" @@ -6826,6 +6854,7 @@ async def test_session_budget_state_persists_during_approval_and_cleans_up_on_co ): from agent_framework._harness._tool_approval import ToolApprovalMiddleware from agent_framework._sessions import AgentSession + from agent_framework._tools import _FUNCTION_INVOCATION_BUDGET_STATE_KEY @tool(name="op", approval_mode="always_require") def op() -> str: @@ -6844,34 +6873,38 @@ def op() -> str: middleware = ToolApprovalMiddleware() agent = Agent(client=chat_client_base, tools=[op], middleware=[middleware]) - call_count = 0 + current_time = [0.0] def fake_perf_counter() -> float: - nonlocal call_count - call_count += 1 - return 0.0 if call_count <= 2 else 10.0 + res = current_time[0] + current_time[0] = 10.0 # Advance time so the next check triggers expiration + return res - import contextlib from unittest.mock import patch - with patch("agent_framework._tools.perf_counter", side_effect=fake_perf_counter): - with contextlib.suppress(Exception): - await agent.run("Go", session=session) + first_response = await agent.run("Go", session=session) + assert any(c.type == "function_approval_request" for c in first_response.messages[-1].contents) - assert "_function_invocation_budget" in session.state - assert "start_time" in session.state["_function_invocation_budget"] + assert _FUNCTION_INVOCATION_BUDGET_STATE_KEY in session.state + assert "start_time" in session.state[_FUNCTION_INVOCATION_BUDGET_STATE_KEY] - session.state["_approval_op1"] = True + approval_request = next( + c for c in first_response.messages[-1].contents if c.type == "function_approval_request" + ) + resume_message = Message( + role="user", contents=[approval_request.to_function_approval_response(approved=True)] + ) - await agent.run("Resume", session=session) + await agent.run(resume_message, session=session) - assert "_function_invocation_budget" not in session.state + assert _FUNCTION_INVOCATION_BUDGET_STATE_KEY not in session.state async def test_plain_session_multiturn_has_isolated_budget_state_per_invocation( chat_client_base: SupportsChatGetResponse, ): from agent_framework._sessions import AgentSession + from agent_framework._tools import _FUNCTION_INVOCATION_BUDGET_STATE_KEY @tool(name="op") def op() -> str: @@ -6891,15 +6924,16 @@ def op() -> str: agent = Agent(client=chat_client_base, tools=[op]) await agent.run("Go", session=session) - assert "_function_invocation_budget" not in session.state + assert _FUNCTION_INVOCATION_BUDGET_STATE_KEY not in session.state await agent.run("Go again", session=session) - assert "_function_invocation_budget" not in session.state + assert _FUNCTION_INVOCATION_BUDGET_STATE_KEY not in session.state async def test_streaming_pending_approval_survives_budget_state_pop(chat_client_base: SupportsChatGetResponse): from agent_framework._harness._tool_approval import ToolApprovalMiddleware from agent_framework._sessions import AgentSession + from agent_framework._tools import _FUNCTION_INVOCATION_BUDGET_STATE_KEY @tool(name="op", approval_mode="always_require") def op() -> str: @@ -6916,12 +6950,77 @@ def op() -> str: middleware = ToolApprovalMiddleware() agent = Agent(client=chat_client_base, tools=[op], middleware=[middleware]) - stream = agent.run("Go", session=session, stream=True) + from unittest.mock import patch + with patch("agent_framework._tools.perf_counter", return_value=0.0): + stream = agent.run("Go", session=session, stream=True) - try: async for _ in stream: pass - except Exception: - pass - assert "_function_invocation_budget" in session.state + first_response = await stream.get_final_response() + assert any(c.type == "function_approval_request" for c in first_response.messages[-1].contents) + assert _FUNCTION_INVOCATION_BUDGET_STATE_KEY in session.state + + approval_request = next( + c for c in first_response.messages[-1].contents if c.type == "function_approval_request" + ) + resume_message = Message( + role="user", contents=[approval_request.to_function_approval_response(approved=True)] + ) + + stream2 = agent.run(resume_message, session=session, stream=True) + async for _ in stream2: + pass + await stream2.get_final_response() + + assert _FUNCTION_INVOCATION_BUDGET_STATE_KEY not in session.state + + +@pytest.mark.parametrize("streaming", [False, True]) +async def test_batch_limit_decision_precedence_duration_over_consecutive_errors( + chat_client_base: SupportsChatGetResponse, + streaming: bool, +): + @tool(name="op") + def op() -> str: + raise ValueError("Error") + + chat_client_base.function_invocation_configuration["max_consecutive_errors_per_request"] = 1 # type: ignore[attr-defined] + chat_client_base.function_invocation_configuration["max_duration_seconds"] = 1.0 # type: ignore[attr-defined] + + call = Content.from_function_call(call_id="op1", name="op", arguments="{}") + + if streaming: + chat_client_base.streaming_responses = [ # type: ignore[attr-defined] + [ChatResponseUpdate(contents=[call], role="assistant", finish_reason="tool_calls")], + [ChatResponseUpdate(contents=[Content.from_text("fallback")], role="assistant", finish_reason="stop")], + ] + else: + chat_client_base.run_responses = [ # type: ignore[attr-defined] + ChatResponse(messages=[Message(role="assistant", contents=[call])], finish_reason="tool_calls"), + ChatResponse( + messages=[Message(role="assistant", contents=[Content.from_text("fallback")])], finish_reason="stop" + ), + ] + + agent = Agent(client=chat_client_base, tools=[op]) + + current_time = [0.0] + + def fake_perf_counter() -> float: + res = current_time[0] + current_time[0] = 10.0 # Advance time so the next check triggers expiration + return res + + from unittest.mock import patch + + with patch("agent_framework._tools.perf_counter", side_effect=fake_perf_counter): + if streaming: + stream = agent.run("Go", stream=True) + async for _ in stream: + pass + response = await stream.get_final_response() + else: + response = await agent.run("Go") + + assert response.additional_properties.get("_agent_framework_stop_reason") == "max_duration_seconds" From feb0479d34b150754d85c23acdaf5421e77869b6 Mon Sep 17 00:00:00 2001 From: karthik-0306 Date: Sat, 22 Aug 2026 15:39:10 +0530 Subject: [PATCH 5/5] fix(core): preserve budget start_time across approval resumes, dedupe streaming limit decision, fix double-counted approval calls against max_function_calls --- .../packages/core/agent_framework/_tools.py | 67 ++++++++++--------- .../packages/core/tests/core/test_agents.py | 4 +- 2 files changed, 36 insertions(+), 35 deletions(-) diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 836e456cb0..252572d1db 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -1905,6 +1905,22 @@ def contents(self) -> list[Content]: """Flatten the ordered result groups for ordinary response processing.""" return [content for result_group in self.result_groups for content in result_group] + @property + def executed_call_count(self) -> int: + """Count of result groups that were actually invoked. + + Excludes groups still deferred on an approval request (the tool body never ran, + unlike a ``UserInputRequiredException`` raised mid-execution, which did run and + still counts) or left as a bare declaration because the batch as a whole was + deferred, so a call is only charged against ``max_function_calls`` once it + actually executes, rather than again when it was merely requested. + """ + return sum( + 1 + for result_group in self.result_groups + if not any(content.type in {"function_approval_request", "function_call"} for content in result_group) + ) + @property def had_errors(self) -> bool: """Whether any execution produced an error result.""" @@ -2775,7 +2791,7 @@ def _clear_budget_state_from_session(invocation_session: "AgentSession | None") if invocation_session is None: return # Only remove the budget if there are no approval requests still pending. - tool_state = cast("dict[str, Any]", invocation_session.state.get("tools")) + tool_state = cast("dict[str, Any]", invocation_session.state.get(_TOOL_APPROVAL_STATE_KEY)) if isinstance(tool_state, dict): pending = tool_state.get(_PENDING_APPROVAL_REQUESTS_KEY) if pending: @@ -2793,6 +2809,9 @@ def _apply_batch_limit_decision( ) -> None: if action != "continue" and action != "stop": return + # Check duration first so it wins precedence via setdefault over consecutive-errors + # (action == "stop") or the function-call limit when more than one is exceeded + # in the same batch. if max_duration_seconds is not None: elapsed = perf_counter() - budget_state["start_time"] if elapsed >= max_duration_seconds: @@ -2806,7 +2825,7 @@ def _apply_batch_limit_decision( if action == "stop": options["tool_choice"] = "none" - budget_state.setdefault("stop_reason", "stop") + budget_state.setdefault("stop_reason", "max_consecutive_errors") else: if _disable_tools_at_function_call_limit(options, total_function_calls, max_function_calls): budget_state.setdefault("stop_reason", "max_function_calls") @@ -3069,7 +3088,7 @@ async def _process_model_function_calls( processing_result = _handle_function_call_results( response=response, execution_results=execution.contents, - function_call_count=len(execution.result_groups), + function_call_count=execution.executed_call_count, function_call_messages=function_call_messages, errors_in_a_row=errors_in_a_row, had_errors=execution.had_errors, @@ -3333,9 +3352,7 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non client_kwargs=request_kwargs, ), ) - if options.get("tool_choice") == "none" and _function_call_limit_reached( - total_function_calls, max_function_calls - ): + if options.get("tool_choice") == "none" and budget_state.get("stop_reason") is not None: _ensure_function_invocation_limit_fallback_response(response) aggregated_usage = add_usage_details(aggregated_usage, response.usage_details) self._update_function_invocation_continuation_state( @@ -3521,9 +3538,8 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non ), ) await inner_stream - drop_unexecutable_calls = options.get("tool_choice") == "none" and _function_call_limit_reached( - total_function_calls, - max_function_calls, + drop_unexecutable_calls = ( + options.get("tool_choice") == "none" and budget_state.get("stop_reason") is not None ) async for update in inner_stream: if drop_unexecutable_calls: @@ -3533,11 +3549,8 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non yield update response = await inner_stream.get_final_response() - function_call_limit_reached = options.get("tool_choice") == "none" and _function_call_limit_reached( - total_function_calls, max_function_calls - ) fallback_added = False - if function_call_limit_reached: + if options.get("tool_choice") == "none" and budget_state.get("stop_reason") is not None: fallback_added = _ensure_function_invocation_limit_fallback_response(response) self._update_function_invocation_continuation_state( request_kwargs, @@ -3592,26 +3605,14 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non if function_processing.action != "continue" and function_processing.action != "stop": # "return" action: model produced a terminal response. return - # Check duration limit first so it wins precedence over consecutive-errors - # (max_consecutive_errors_per_request sets action="stop") when both are exceeded - # in the same batch — matching the non-streaming _apply_batch_limit_decision logic. - if ( - max_duration_seconds is not None - and (perf_counter() - budget_state["start_time"]) >= max_duration_seconds - ): - logger.info( - "Maximum duration reached (%.2fs / %.2fs). Stopping further function calls for this request.", - perf_counter() - budget_state["start_time"], - max_duration_seconds, - ) - options["tool_choice"] = "none" - budget_state.setdefault("stop_reason", "max_duration_seconds") - elif function_processing.action == "stop": - options["tool_choice"] = "none" - budget_state.setdefault("stop_reason", "stop") - else: - if _disable_tools_at_function_call_limit(options, total_function_calls, max_function_calls): - budget_state.setdefault("stop_reason", "max_function_calls") + _apply_batch_limit_decision( + function_processing.action, + options, + budget_state, + total_function_calls, + max_function_calls, + max_duration_seconds, + ) _reset_required_tool_choice(options) _prepare_messages_for_next_iteration(prepared_messages, response) diff --git a/python/packages/core/tests/core/test_agents.py b/python/packages/core/tests/core/test_agents.py index 4ef0d031d7..615829c8d0 100644 --- a/python/packages/core/tests/core/test_agents.py +++ b/python/packages/core/tests/core/test_agents.py @@ -1345,7 +1345,7 @@ async def after_run( assert response.usage_details == {"input_token_count": 3, "output_token_count": 2, "total_token_count": 5} assert response.value == {"answer": "ok"} assert response.continuation_token == {"token": "next"} - assert response.additional_properties == {"provider": "test"} + assert response.additional_properties == {"provider": "test", "_agent_framework_stop_reason": "completed"} assert response.raw_representation is raw_response @@ -1434,7 +1434,7 @@ async def after_run( assert response.usage_details == {"input_token_count": 4, "output_token_count": 3, "total_token_count": 7} assert response.value == {"answer": "ok"} assert response.continuation_token == {"token": "stream-next"} - assert response.additional_properties == {"provider": "stream-test"} + assert response.additional_properties == {"provider": "stream-test", "_agent_framework_stop_reason": "completed"} assert response.raw_representation == [raw_update]