From 01b024f4560ad7ca96ad299a3b7d9824e322d788 Mon Sep 17 00:00:00 2001 From: CorgiBoyG Date: Thu, 17 Sep 2026 11:43:37 +0800 Subject: [PATCH] fix: preserve mixed-batch recovery invariants --- ...1-handle-provider-invalidated-responses.md | 8 +- .../specs/004-python-function-calling-loop.md | 17 +- python/packages/core/AGENTS.md | 4 + .../packages/core/agent_framework/_tools.py | 346 +++++++- .../core/test_function_invocation_logic.py | 740 +++++++++++++++++- 5 files changed, 1059 insertions(+), 56 deletions(-) diff --git a/docs/decisions/0041-handle-provider-invalidated-responses.md b/docs/decisions/0041-handle-provider-invalidated-responses.md index fdefd5a6762..ed9673525a6 100644 --- a/docs/decisions/0041-handle-provider-invalidated-responses.md +++ b/docs/decisions/0041-handle-provider-invalidated-responses.md @@ -54,8 +54,12 @@ usability is the portable default while invalidation is provider-specific protoc `FunctionInvocationLayer` remains independent of `finish_reason`: a newly completed actionable call proceeds when argument preparation and schema validation succeed. A provider that knows partial response output was invalidated raises `ResponseInvalidatedException`; any local function calls from that response must not execute. The layer -abandons that current iteration, clears request budget state, restores the last valid continuation, avoids successful -response persistence and local function side effects, and re-raises the same exception. +abandons that current iteration, ordinarily clears request budget state, restores the last valid continuation, avoids +successful response persistence and local function side effects from the invalid response, and re-raises the same exception. +If the invalidated call was instead delivering results from an already completed mixed approval/Host batch, the layer +retains that serializable provider outbox and its charged budget. A retry replays the stored Host and local results +without recovering approval authority or executing the local side effect again; only a successful provider response +clears the outbox. Anthropic applies the signal only to local actionable `tool_use` blocks. A valid stream has closed local blocks, a terminal `stop_reason` of `tool_use`, and `message_stop`. Non-tool terminal reasons, an open block at `message_stop`, diff --git a/docs/specs/004-python-function-calling-loop.md b/docs/specs/004-python-function-calling-loop.md index 71dc8f59140..e78251be270 100644 --- a/docs/specs/004-python-function-calling-loop.md +++ b/docs/specs/004-python-function-calling-loop.md @@ -374,11 +374,14 @@ that manually replay messages own the equivalent rule: do not resend an approval report another reason, including `length`, while still returning a complete and service-usable call. - A provider adapter that knows its protocol invalidated or cancelled partial response output raises `ResponseInvalidatedException`; any local function calls already exposed from that response must not execute. The - function invocation layer then abandons only that current model/tool iteration, clears its request budget state, - restores the last valid service continuation, and re-raises - the same exception. It emits no new approval request or function result, invokes no function middleware or tool - body, makes no second model request or dangling-call settlement request, and does not success-persist that service - response. Approval decisions resolved before the invalidated provider call retain their existing semantics. + function invocation layer then abandons only that current model/tool iteration, restores the last valid service + continuation, and re-raises the same exception. It emits no new approval request or function result, invokes no + function middleware or tool body for calls from the invalid response, makes no second model request or dangling-call + settlement request, and does not success-persist that service response. Ordinarily it clears the request budget. + When the failed request was delivering a completed mixed approval/Host batch, however, the serializable pending + batch, approval authority, ordered Host and local results, and already-charged budget remain in a provider outbox. + A later run replays that outbox without executing the approved function again, and clears it only after a provider + response succeeds. Streaming providers may yield local function-call deltas before discovering invalidation; direct stream consumers must treat the exception as an explicit instruction to discard those calls and never execute them. Caller cancellation remains cancellation rather than becoming this provider signal. @@ -624,7 +627,7 @@ that manually replay messages own the equivalent rule: do not resend an approval | Scenario | Required invariant | Primary regression test | |---|---|---| | Fatal call mixed with pauses | Complete-batch classification raises before approval or execution, independent of call order. | `packages/core/tests/core/test_function_invocation_logic.py::test_mixed_batch_fatal_unknown_precedes_every_pause` | -| Approval and Host-owned calls | Both pause types are returned in model order; a session-backed partial response remains pending across serialization; occurrence-identified Host results reserve their slots before id-less results use the unique unanswered occurrence, while authoritative session state recognizes equivalent duplicates and conflicting duplicates fail closed; stateless replay does not treat an ambiguous id-less result as a Host duplicate without occurrence provenance; a stateless zero-response or partial response fails closed across non-user message roles and cannot be hidden by a newer standalone request; standalone pauses separated by an unrelated user turn remain independent and response-order invariant when call IDs are reused; exact Host occurrence identity outranks newer call-ID-only candidates; each stateless response belongs to the nearest compatible request batch even when it cannot be assigned to one item, response ownership is discovered in linear time, and Host-result exclusions close only their Host occurrences without keeping historical completed calls open or reinterpreting later local terminal results during approval normalization; completed mixed batches remain inert even when approval and Host requests reuse a call ID; historical Host calls do not participate; a complete response executes the exact approved arguments once. | `test_mixed_batch_returns_approval_and_host_pause_in_model_order`, `test_mixed_batch_requires_complete_responses_before_execution`, `test_stateful_mixed_batch_accepts_equivalent_idless_host_result_replay`, `test_stateful_mixed_batch_assigns_idless_equal_result_to_unanswered_occurrence`, `test_stateless_mixed_batch_rejects_conflicting_identified_host_results`, `test_stateless_split_mixed_batch_rejects_incomplete_replay_before_execution`, `test_stateless_mixed_batch_across_non_user_message_roles_requires_complete_responses`, `test_stateless_abandoned_approval_does_not_join_later_host_request`, `test_stateless_separated_pauses_with_reused_call_id_are_order_independent`, `test_exact_older_host_result_does_not_consume_newer_reused_call_approval`, `test_later_standalone_request_does_not_hide_incomplete_stateless_mixed_batch`, `test_later_idless_host_result_does_not_complete_older_stateless_mixed_batch`, `test_completed_approval_result_is_not_claimed_by_older_stateless_host_request`, `test_historical_stateless_host_result_does_not_capture_later_reused_call_approval`, `test_excluded_host_result_closes_own_occurrence_before_reused_call_approval`, `test_ambiguous_later_host_result_does_not_complete_older_stateless_mixed_batch`, `test_id_bearing_result_for_idless_host_request_does_not_consume_approval`, `test_stateless_pause_response_ownership_scans_contents_linearly`, `test_completed_split_stateless_mixed_batch_is_inert_on_later_turn`, `test_completed_stateless_mixed_batch_with_reused_call_id_is_inert`, `test_equal_idless_terminal_result_does_not_reexecute_completed_stateless_mixed_approval`, `test_active_mixed_pause_ignores_historical_host_requests` | +| Approval and Host-owned calls | Both pause types are returned in model order; a session-backed partial response remains pending across serialization; occurrence-identified Host results reserve their slots before id-less results use the unique unanswered occurrence, while authoritative session state recognizes equivalent duplicates and conflicting duplicates fail closed; stateless replay does not treat an ambiguous id-less result as a Host duplicate without occurrence provenance; a stateless zero-response or partial response fails closed across non-user message roles and cannot be hidden by a newer standalone request; standalone pauses separated by an unrelated user turn remain independent and response-order invariant when call IDs are reused; exact Host occurrence identity outranks newer call-ID-only candidates; each stateless response belongs to the nearest compatible request batch even when it cannot be assigned to one item, response ownership is discovered in linear time, and Host-result exclusions close only their Host occurrences without keeping historical completed calls open or reinterpreting later local terminal results during approval normalization; completed mixed batches remain inert even when approval and Host requests reuse a call ID; historical Host calls do not participate; a complete response executes the exact approved arguments once. A provider-invalidated submission retains a serializable two-phase outbox, approval authority, ordered Host and local results, and charged budget for execution-free retry; restored duration metadata is validated and rebased without comparing unrelated monotonic-clock epochs. | `test_mixed_batch_returns_approval_and_host_pause_in_model_order`, `test_mixed_batch_requires_complete_responses_before_execution`, `test_stateful_mixed_batch_accepts_equivalent_idless_host_result_replay`, `test_stateful_mixed_batch_assigns_idless_equal_result_to_unanswered_occurrence`, `test_stateless_mixed_batch_rejects_conflicting_identified_host_results`, `test_stateless_split_mixed_batch_rejects_incomplete_replay_before_execution`, `test_stateless_mixed_batch_across_non_user_message_roles_requires_complete_responses`, `test_stateless_abandoned_approval_does_not_join_later_host_request`, `test_stateless_separated_pauses_with_reused_call_id_are_order_independent`, `test_exact_older_host_result_does_not_consume_newer_reused_call_approval`, `test_later_standalone_request_does_not_hide_incomplete_stateless_mixed_batch`, `test_later_idless_host_result_does_not_complete_older_stateless_mixed_batch`, `test_completed_approval_result_is_not_claimed_by_older_stateless_host_request`, `test_historical_stateless_host_result_does_not_capture_later_reused_call_approval`, `test_excluded_host_result_closes_own_occurrence_before_reused_call_approval`, `test_ambiguous_later_host_result_does_not_complete_older_stateless_mixed_batch`, `test_id_bearing_result_for_idless_host_request_does_not_consume_approval`, `test_stateless_pause_response_ownership_scans_contents_linearly`, `test_completed_split_stateless_mixed_batch_is_inert_on_later_turn`, `test_completed_stateless_mixed_batch_with_reused_call_id_is_inert`, `test_equal_idless_terminal_result_does_not_reexecute_completed_stateless_mixed_approval`, `test_completed_stateless_mixed_batch_is_inert_when_later_turn_reuses_host_identity`, `test_active_mixed_pause_ignores_historical_host_requests`, `test_completed_mixed_batch_replays_serialized_outbox_after_provider_invalidation`, `test_serialized_provider_outbox_rebases_duration_budget_across_monotonic_epochs`, `test_untrusted_provider_outbox_duration_budget_fails_closed_after_restore`, `test_malformed_provider_outbox_budget_is_rejected` | | Safe and approval-required calls in one batch | Hidden safe calls replay only with the matching visible approval. | `packages/core/tests/core/test_harness_tool_approval.py::test_mixed_batch_hides_already_approved_request_until_approval_replay` | | Restored approval state | Serialized `ToolApprovalState` restores mixed-batch behavior. | `test_mixed_batch_accepts_restored_tool_approval_state` | | Unrelated turn before approval | Hidden calls do not execute on an unrelated turn. | `test_hidden_mixed_batch_requests_do_not_replay_on_unrelated_turn` | @@ -666,7 +669,7 @@ that manually replay messages own the equivalent rule: do not resend an approval | Middleware failure batch cancellation | A fatal signal fails the whole parallel batch: in-flight sibling tool invocations are cancelled and awaited before the failure propagates. Cancellation is cooperative — an async sibling stops at its next suspension point; a synchronous tool body already executing in a worker thread cannot be interrupted and may complete its side effects, but its result is discarded and never reaches the transcript, the model, or history, and failure propagation is not delayed behind it. | `TestMiddlewareFailure::test_failure_cancels_concurrent_sibling_tool`, `test_failure_with_sync_sibling_discards_late_result` | | Middleware failure on a service-managed conversation | The continuation state is already persisted when the batch fails, so before propagating, the loop settles the hosted thread: one error `function_result` per dangling call, sent with `tool_choice="none"` in one extra request; the persisted continuation advances to the settlement response (required for response-ID continuations, a no-op for conversation-object ids) and the settlement response is otherwise discarded; a settlement failure never masks the abort. Without a service-managed conversation no extra request is made. | `TestMiddlewareFailure::test_failure_settles_dangling_calls_on_service_conversation`, `test_failure_settles_service_conversation_streaming`, `test_failure_settlement_advances_response_id_continuation`, `test_failure_without_service_conversation_makes_no_settlement_request` | | Middleware failure during approved-tool replay | A fatal abort while the approval-resolution phase replays an approved tool escapes loudly (never absorbed into a rejection result), the tool's original — already service-persisted — call is settled the same way, and the continuation advances; both response modes. | `TestMiddlewareFailure::test_failure_during_approved_replay_settles_and_escapes`, `test_failure_during_approved_replay_streaming` | -| Provider-invalidated partial response output | `ResponseInvalidatedException` propagates unchanged in both response modes after clearing request budget state and restoring the last valid continuation. Partial streamed call deltas may remain visible, but no new approval, function middleware, tool body, result, follow-up model call, or settlement occurs. Anthropic requires every local call block to close, terminal `stop_reason="tool_use"`, and `message_stop`; non-tool terminal reasons, open blocks, missing `message_stop`, and non-cancellation stream errors after local call start invalidate the call. Hosted/server-only calls and caller cancellation remain unaffected. | `packages/core/tests/core/test_function_invocation_logic.py::test_response_invalidation_short_circuits_current_iteration`, `test_invalidated_final_no_tool_response_preserves_prior_continuation`, `packages/anthropic/tests/test_anthropic_client.py::test_non_streaming_local_tool_call_with_invalidating_stop_reason_raises`, `test_valid_streaming_local_tool_call_executes_and_continues`, `test_streaming_local_tool_call_invalid_terminal_sequences_raise`, `test_streaming_provider_error_after_local_call_is_wrapped_as_invalidation`, `test_streaming_cancellation_after_local_call_is_not_wrapped`, `test_streaming_hosted_tool_pause_turn_is_not_invalidated` | +| Provider-invalidated partial response output | `ResponseInvalidatedException` propagates unchanged in both response modes after restoring the last valid continuation. It clears ordinary request budget state, but preserves a completed mixed-batch provider outbox and its charged budget for execution-free retry. Partial streamed call deltas may remain visible, but no new approval, function middleware, tool body, result, follow-up model call, or settlement occurs for calls from the invalid response. Anthropic requires every local call block to close, terminal `stop_reason="tool_use"`, and `message_stop`; non-tool terminal reasons, open blocks, missing `message_stop`, and non-cancellation stream errors after local call start invalidate the call. Hosted/server-only calls and caller cancellation remain unaffected. | `packages/core/tests/core/test_function_invocation_logic.py::test_response_invalidation_short_circuits_current_iteration`, `test_invalidated_final_no_tool_response_preserves_prior_continuation`, `test_completed_mixed_batch_replays_serialized_outbox_after_provider_invalidation`, `packages/anthropic/tests/test_anthropic_client.py::test_non_streaming_local_tool_call_with_invalidating_stop_reason_raises`, `test_valid_streaming_local_tool_call_executes_and_continues`, `test_streaming_local_tool_call_invalid_terminal_sequences_raise`, `test_streaming_provider_error_after_local_call_is_wrapped_as_invalidation`, `test_streaming_cancellation_after_local_call_is_not_wrapped`, `test_streaming_hosted_tool_pause_turn_is_not_invalidated` | | Maximum iterations | No orphan calls; a final no-tool response or deterministic fallback is returned. | `test_max_iterations_limit`, `test_max_iterations_no_orphaned_function_calls`, `test_max_iterations_makes_final_toolchoice_none_call`, `test_max_iterations_blank_final_fallback_synthesizes_message`, streaming equivalents | | Maximum function calls | Parallel overshoot is bounded after the batch; every executed result group counts even without a `function_result`; blank final responses get fallback content. | `test_max_function_calls_limits_parallel_invocations`, `test_max_function_calls_single_calls_per_iteration`, `test_user_input_request_multiple_contents_propagate`, `test_approval_resume_user_input_counts_toward_function_call_budget`, `test_max_function_calls_blank_final_fallback_synthesizes_message`, streaming equivalent | | Provider tool content after an active limit | Locally actionable calls and local approval requests returned despite `tool_choice="none"` are removed in both response modes. Provider-executed informational call/result pairs, hosted approval requests, and metadata-only streaming updates remain visible; fallback text never replaces retained transcript content. | `test_function_invocation_limit_drops_unexecutable_tool_content`, `test_streaming_function_invocation_limit_drops_unexecutable_tool_content`, `test_streaming_function_invocation_limit_preserves_metadata_after_tool_content_is_dropped`, `test_function_invocation_limit_preserves_provider_executed_tool_pair`, `test_streaming_function_invocation_limit_preserves_provider_executed_tool_pair`, `test_function_invocation_limit_appends_fallback_after_provider_executed_tool_pair`, `test_streaming_function_invocation_limit_appends_fallback_after_provider_executed_tool_pair`, `test_function_invocation_limit_preserves_hosted_approval_request`, `test_streaming_function_invocation_limit_preserves_hosted_approval_request` | diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index d32843f4204..6e01c21b280 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -232,6 +232,10 @@ The vector store API is experimental under the shared `VECTOR_STORES` feature ID available, approval requests for known non-approval-required tools are treated as already approved, hidden, stored in session state keyed to the visible approval request ids from that batch, and reinjected only when that visible approval flow resumes. +- Once a mixed approval/Host pause batch is complete, its ordered provider input and locally produced results remain + in a serializable two-phase outbox until the provider accepts them. Provider invalidation preserves the approval + authority, outbox, service continuation, and charged budget; retry replays the stored results without re-executing + local side effects. - Approval resume is an immutable response boundary: the function invocation layer normalizes a private copy of caller messages, returns approved and rejected terminal results in the resumed response (and stream) before any final assistant message, and does not mutate the caller's approval `Message` or the earlier approval-request diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index ff2756cc60d..d107a2eadab 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -24,6 +24,7 @@ from contextlib import suppress from dataclasses import dataclass from functools import partial, wraps +from math import isfinite from time import perf_counter, time_ns from typing import ( TYPE_CHECKING, @@ -118,6 +119,10 @@ def _has_authoritative_approval_session(invocation_session: AgentSession | None) _FUNCTION_CALL_ORDER_KEY: Final[str] = "function_call_order" _PENDING_APPROVAL_REQUESTS_KEY: Final[str] = "pending_approval_requests" _PENDING_MIXED_PAUSE_BATCH_KEY: Final[str] = "pending_mixed_pause_batch" +_PENDING_PROVIDER_OUTBOX_KEY: Final[str] = "pending_provider_outbox" +_PROVIDER_OUTBOX_BUDGET_ELAPSED_SECONDS_KEY: Final[str] = "elapsed_duration_seconds" +_PROVIDER_OUTBOX_BUDGET_SAVED_AT_NS_KEY: Final[str] = "saved_at_time_ns" +_FUNCTION_INVOCATION_DURATION_EXHAUSTED_KEY: Final[str] = "duration_exhausted" _APPROVAL_REQUEST_ID_KEY: Final[str] = "_approval_request_id" _FUNCTION_INVOCATION_BUDGET_STATE_KEY: Final[str] = "_function_invocation_budget_state" _FUNCTION_RESULT_CARRIER_CONTEXT_KEY: Final[str] = "_function_result_carrier" @@ -2778,6 +2783,17 @@ def _content_from_state(value: Any) -> Content | None: return None +def _message_from_state(value: Any) -> Message | None: + """Restore a Message item stored in session state.""" + from ._types import Message + + if isinstance(value, Message): + return Message.from_dict(value.to_dict()) + if isinstance(value, Mapping): + return Message.from_dict(dict(cast(Mapping[str, Any], value))) + return None + + def _load_pending_approval_requests(invocation_session: AgentSession | None) -> dict[str, Content]: """Load immutable approval-request snapshots keyed by request ID.""" state = _get_tool_approval_state(invocation_session, create=False) @@ -3061,9 +3077,7 @@ def _stage_approval_batch_responses( if any(request_id not in stored_responses for request_id in group_ids): updated_group = dict(group) updated_group[_APPROVAL_RESPONSES_KEY] = [ - stored_responses[request_id].to_dict() - for request_id in group_ids - if request_id in stored_responses + stored_responses[request_id].to_dict() for request_id in group_ids if request_id in stored_responses ] remaining_groups.append(updated_group) missing_request_ids = [request_id for request_id in group_ids if request_id not in stored_responses] @@ -3144,7 +3158,8 @@ def _match_mixed_pause_responses( *, approval_response_binder: Callable[[Content], Content | None] | None = None, allow_idless_host_duplicates: bool = True, -) -> tuple[set[int], bool, list[Content], set[int]]: + allowed_idless_host_duplicate_ids: set[int] | None = None, +) -> tuple[set[int], bool, list[Content], set[int], set[str]]: """Match one complete mixed pause batch without depending on its storage source.""" approval_items: dict[str, int | None] = {} host_items_by_occurrence: dict[str, int] = {} @@ -3170,6 +3185,9 @@ def _match_mixed_pause_responses( host_items_by_occurrence[request.id] = index matched_content_ids: set[int] = set() + preexisting_response_indexes = {index for index, item in enumerate(items) if item.get("response") is not None} + observed_preexisting_indexes: set[int] = set() + conflicting_host_response_errors: set[str] = set() for match_idless_host_results in (False, True): for response in responses: is_idless_host_result = response.type == "function_result" and response.id is None @@ -3204,14 +3222,29 @@ def _match_mixed_pause_responses( if request is None or request.call_id != response.call_id: item_index = None else: + replay_indexes = [ + pending_index + for pending_index in host_items_by_call[response.call_id] + if pending_index in preexisting_response_indexes + and pending_index not in observed_preexisting_indexes + and isinstance(items[pending_index].get("response"), Mapping) + and cast(Mapping[str, Any], items[pending_index]["response"]).get("id") is None + and _same_mixed_pause_response( + cast(Mapping[str, Any], items[pending_index]["response"]), + response.to_dict(), + ) + ] unanswered_indexes = [ pending_index for pending_index in host_items_by_call[response.call_id] if items[pending_index].get("response") is None ] - if len(unanswered_indexes) == 1: + if len(replay_indexes) == 1: + item_index = replay_indexes[0] + observed_preexisting_indexes.add(item_index) + elif len(unanswered_indexes) == 1: item_index = unanswered_indexes[0] - elif not unanswered_indexes and allow_idless_host_duplicates: + elif not unanswered_indexes: duplicate_indexes = [ pending_index for pending_index in host_items_by_call[response.call_id] @@ -3221,8 +3254,19 @@ def _match_mixed_pause_responses( response.to_dict(), ) ] - if len(duplicate_indexes) == 1: + if duplicate_indexes and ( + allow_idless_host_duplicates + or ( + allowed_idless_host_duplicate_ids is not None + and id(response) in allowed_idless_host_duplicate_ids + ) + ): item_index = duplicate_indexes[0] + elif not duplicate_indexes: + conflicting_host_response_errors.add( + f"Conflicting Host response for mixed pause call_id {response.call_id!r}." + ) + continue else: continue @@ -3234,23 +3278,28 @@ def _match_mixed_pause_responses( not isinstance(stored_response, Mapping) or not _same_mixed_pause_response(cast(Mapping[str, Any], stored_response), candidate_state) ): + if items[item_index].get("kind") == "host": + conflicting_host_response_errors.add( + f"Conflicting Host response for mixed pause occurrence {candidate.id!r}." + ) + continue raise RuntimeError(f"Conflicting response for mixed pause occurrence {candidate.id!r}.") items[item_index]["response"] = candidate_state matched_content_ids.add(id(response)) if any(item.get("response") is None for item in items): - return matched_content_ids, True, [], set() + return matched_content_ids, True, [], set(), conflicting_host_response_errors ordered_responses: list[Content] = [] host_result_ids: set[int] = set() for item in items: response = _content_from_state(item.get("response")) if response is None: - return matched_content_ids, True, [], set() + return matched_content_ids, True, [], set(), conflicting_host_response_errors ordered_responses.append(response) if item.get("kind") == "host": host_result_ids.add(id(response)) - return matched_content_ids, False, ordered_responses, host_result_ids + return matched_content_ids, False, ordered_responses, host_result_ids, conflicting_host_response_errors def _stage_pending_mixed_pause_responses( @@ -3289,11 +3338,15 @@ def bind_approval_response(response: Content) -> Content | None: for content in message.contents if content.type in {"function_approval_response", "function_result"} ] - matched_content_ids, incomplete, ordered_responses, host_result_ids = _match_mixed_pause_responses( - items, - responses, - approval_response_binder=bind_approval_response, + matched_content_ids, incomplete, ordered_responses, host_result_ids, conflicting_host_response_errors = ( + _match_mixed_pause_responses( + items, + responses, + approval_response_binder=bind_approval_response, + ) ) + if conflicting_host_response_errors: + raise RuntimeError(" ".join(sorted(conflicting_host_response_errors))) if matched_content_ids: filtered_messages: list[Message] = [] @@ -3478,7 +3531,7 @@ def register_request_batch(batch_index: int) -> None: reserved_host_result_ids.add(id(content)) host_result_content_indices[id(content)] = response_index - batch_matches: list[tuple[set[int], bool, list[Content], set[int]]] = [] + batch_matches: list[tuple[set[int], bool, list[Content], set[int], set[str]]] = [] matched_host_result_ids = set(reserved_host_result_ids) for batch_index, (_, items, _) in enumerate(request_batches): responses = responses_by_batch[batch_index] @@ -3486,6 +3539,7 @@ def register_request_batch(batch_index: int) -> None: items, responses, allow_idless_host_duplicates=False, + allowed_idless_host_duplicate_ids=reserved_host_result_ids, ) batch_matches.append(match) matched_response_ids = match[0] @@ -3518,7 +3572,13 @@ def register_request_batch(batch_index: int) -> None: } for batch_index in range(len(request_batches) - 1, -1, -1): _, _, kinds = request_batches[batch_index] - matched_response_ids, incomplete, ordered_responses, ordered_host_result_ids = batch_matches[batch_index] + ( + matched_response_ids, + incomplete, + ordered_responses, + ordered_host_result_ids, + conflicting_host_response_errors, + ) = batch_matches[batch_index] if kinds != {"approval", "host"}: continue if incomplete: @@ -3526,6 +3586,8 @@ def register_request_batch(batch_index: int) -> None: if matched_response_ids.isdisjoint(pending_approval_response_ids): continue + if conflicting_host_response_errors: + raise RuntimeError(" ".join(sorted(conflicting_host_response_errors))) original_host_result_ids = matched_response_ids & active_host_result_ids active_host_result_ids.difference_update(original_host_result_ids) @@ -4142,11 +4204,16 @@ def _response_invalidation_cleanup( def cleanup(error: ResponseInvalidatedException) -> None: if stream_error is not None: stream_error[:] = [error] - budget_state.clear() if invocation_session is None: + budget_state.clear() return - invocation_session.state.pop(_FUNCTION_INVOCATION_BUDGET_STATE_KEY, None) invocation_session.service_session_id = service_session_id + if _has_pending_provider_outbox(invocation_session): + _restore_pending_provider_outbox_budget(invocation_session, budget_state) + invocation_session.state[_FUNCTION_INVOCATION_BUDGET_STATE_KEY] = budget_state + return + budget_state.clear() + invocation_session.state.pop(_FUNCTION_INVOCATION_BUDGET_STATE_KEY, None) return cleanup @@ -4192,7 +4259,7 @@ def _apply_batch_limit_decision( return if max_duration_seconds is not None: elapsed = perf_counter() - budget_state["start_time"] - if elapsed >= max_duration_seconds: + if budget_state.get(_FUNCTION_INVOCATION_DURATION_EXHAUSTED_KEY) is True or elapsed >= max_duration_seconds: logger.info( "Maximum duration reached (%.2fs / %.2fs). Stopping further function calls for this request.", elapsed, @@ -4244,6 +4311,178 @@ class _FunctionProcessingResult: streaming_updates: tuple[ChatResponseUpdate, ...] = () +def _restore_pending_provider_outbox( + prepared_messages: list[Message], + invocation_session: AgentSession | None, +) -> _FunctionProcessingResult | None: + """Restore a completed mixed batch without re-authorizing or re-executing local calls.""" + state = _get_tool_approval_state(invocation_session, create=False) + if state is None: + return None + raw_outbox = state.get(_PENDING_PROVIDER_OUTBOX_KEY) + if not isinstance(raw_outbox, Mapping): + return None + outbox = cast(Mapping[str, Any], raw_outbox) + errors_in_a_row = outbox.get("errors_in_a_row", 0) + if not isinstance(errors_in_a_row, int) or isinstance(errors_in_a_row, bool) or errors_in_a_row < 0: + raise RuntimeError("The pending mixed-batch provider outbox has an invalid error counter.") + raw_provider_messages = outbox.get("provider_messages") + raw_response_messages = outbox.get("response_messages") + if not isinstance(raw_provider_messages, list) or not isinstance(raw_response_messages, list): + raise RuntimeError("The pending mixed-batch provider outbox is malformed.") + + restored_groups: list[list[Message]] = [[], []] + for raw_messages, restored_messages in zip( + (cast(list[Any], raw_provider_messages), cast(list[Any], raw_response_messages)), + restored_groups, + strict=True, + ): + for item in raw_messages: + message = _message_from_state(item) + if message is None: + raise RuntimeError("The pending mixed-batch provider outbox contains invalid messages.") + restored_messages.append(message) + provider_messages, response_message_items = restored_groups + response_messages = tuple(response_message_items) + prepared_messages[:] = provider_messages + + streaming_updates: tuple[ChatResponseUpdate, ...] = () + streaming_updates_published = outbox.get("streaming_updates_published", False) + if not isinstance(streaming_updates_published, bool): + raise RuntimeError("The pending mixed-batch provider outbox has an invalid publication marker.") + if not streaming_updates_published: + terminal_contents = [content for message in response_messages for content in message.contents] + _, streaming_updates = _messages_and_updates_for_terminal_contents(terminal_contents) + action = outbox.get("action") + if action not in {"continue", "stop"}: + raise RuntimeError("The pending mixed-batch provider outbox contains an invalid action.") + return _FunctionProcessingResult( + errors_in_a_row=errors_in_a_row, + action=cast("Literal['continue', 'stop']", action), + function_call_count=0, + response_messages=response_messages, + streaming_updates=streaming_updates, + ) + + +def _store_pending_provider_outbox( + invocation_session: AgentSession | None, + *, + prepared_messages: Sequence[Message], + processing_result: _FunctionProcessingResult, +) -> None: + """Persist completed mixed-batch results until provider delivery succeeds.""" + state = _get_tool_approval_state(invocation_session) + if state is None: + return + state[_PENDING_PROVIDER_OUTBOX_KEY] = { + "provider_messages": [message.to_dict() for message in prepared_messages], + "response_messages": [message.to_dict() for message in processing_result.response_messages], + "errors_in_a_row": processing_result.errors_in_a_row, + "action": processing_result.action, + "streaming_updates_published": False, + } + + +def _mark_pending_provider_outbox_streaming_updates_published( + invocation_session: AgentSession | None, +) -> None: + """Record that terminal streaming updates were already emitted to the caller.""" + state = _get_tool_approval_state(invocation_session, create=False) + if state is None: + return + raw_outbox = state.get(_PENDING_PROVIDER_OUTBOX_KEY) + if isinstance(raw_outbox, dict): + raw_outbox["streaming_updates_published"] = True + + +def _persist_pending_provider_outbox_budget( + invocation_session: AgentSession | None, + budget_state: dict[str, Any], +) -> None: + """Persist the charged budget beside an outbox before provider delivery.""" + state = _get_tool_approval_state(invocation_session, create=False) + if state is None or invocation_session is None: + return + raw_outbox = state.get(_PENDING_PROVIDER_OUTBOX_KEY) + if not isinstance(raw_outbox, dict): + return + serialized_budget = copy.deepcopy(budget_state) + serialized_budget.setdefault("attempt_count", 0) + serialized_budget.setdefault("total_function_calls", 0) + start_time = serialized_budget.pop("start_time", None) + if isinstance(start_time, int | float) and not isinstance(start_time, bool) and isfinite(float(start_time)): + serialized_budget[_PROVIDER_OUTBOX_BUDGET_ELAPSED_SECONDS_KEY] = max( + 0.0, + perf_counter() - float(start_time), + ) + serialized_budget[_PROVIDER_OUTBOX_BUDGET_SAVED_AT_NS_KEY] = time_ns() + raw_outbox["budget_state"] = serialized_budget + invocation_session.state[_FUNCTION_INVOCATION_BUDGET_STATE_KEY] = budget_state + + +def _restore_pending_provider_outbox_budget( + invocation_session: AgentSession | None, + budget_state: dict[str, Any], +) -> None: + """Restore an outbox's already-charged budget into the current run.""" + state = _get_tool_approval_state(invocation_session, create=False) + if state is None: + return + raw_outbox = state.get(_PENDING_PROVIDER_OUTBOX_KEY) + if not isinstance(raw_outbox, Mapping): + return + raw_budget = cast(Mapping[str, Any], raw_outbox).get("budget_state") + if not isinstance(raw_budget, Mapping): + raise RuntimeError("The pending mixed-batch provider outbox contains an invalid budget.") + restored_budget = copy.deepcopy(dict(cast(Mapping[str, Any], raw_budget))) + for key in ("attempt_count", "total_function_calls"): + value = restored_budget.get(key) + if not isinstance(value, int) or isinstance(value, bool) or value < 0: + raise RuntimeError("The pending mixed-batch provider outbox contains an invalid budget.") + for key in ("truncated", _FUNCTION_INVOCATION_DURATION_EXHAUSTED_KEY): + if key in restored_budget and not isinstance(restored_budget[key], bool): + raise RuntimeError("The pending mixed-batch provider outbox contains an invalid budget.") + + elapsed = restored_budget.pop(_PROVIDER_OUTBOX_BUDGET_ELAPSED_SECONDS_KEY, None) + saved_at_ns = restored_budget.pop(_PROVIDER_OUTBOX_BUDGET_SAVED_AT_NS_KEY, None) + current_time_ns = time_ns() + if ( + isinstance(elapsed, int | float) + and not isinstance(elapsed, bool) + and isfinite(float(elapsed)) + and float(elapsed) >= 0 + and isinstance(saved_at_ns, int) + and not isinstance(saved_at_ns, bool) + and 0 <= saved_at_ns <= current_time_ns + ): + elapsed_since_persist_ns = current_time_ns - saved_at_ns + restored_budget["start_time"] = perf_counter() + restored_budget["start_time"] -= float(elapsed) + elapsed_since_persist_ns / 1_000_000_000 + else: + restored_budget.pop("start_time", None) + restored_budget["start_time"] = perf_counter() + restored_budget[_FUNCTION_INVOCATION_DURATION_EXHAUSTED_KEY] = True + budget_state.clear() + budget_state.update(restored_budget) + + +def _complete_pending_provider_outbox(invocation_session: AgentSession | None) -> None: + """Clear completed mixed-batch state only after provider delivery succeeds.""" + state = _get_tool_approval_state(invocation_session, create=False) + if state is None or not isinstance(state.get(_PENDING_PROVIDER_OUTBOX_KEY), Mapping): + return + state.pop(_PENDING_PROVIDER_OUTBOX_KEY, None) + state.pop(_PENDING_MIXED_PAUSE_BATCH_KEY, None) + state.pop(_PENDING_APPROVAL_REQUESTS_KEY, None) + state.pop(_ALREADY_APPROVED_APPROVAL_REQUEST_GROUPS_KEY, None) + + +def _has_pending_provider_outbox(invocation_session: AgentSession | None) -> bool: + state = _get_tool_approval_state(invocation_session, create=False) + return state is not None and isinstance(state.get(_PENDING_PROVIDER_OUTBOX_KEY), Mapping) + + _FunctionCallExecutor: TypeAlias = Callable[..., Awaitable[_FunctionExecutionBatch]] @@ -4405,6 +4644,8 @@ async def _resolve_approval_responses( from ._types import Message approval_session = invocation_session if approval_session_is_authoritative else None + if outbox_replay := _restore_pending_provider_outbox(prepared_messages, approval_session): + return outbox_replay completed_mixed_batch = False if _has_authoritative_approval_session(approval_session): incomplete_mixed_batch, completed_mixed_batch, host_result_ids = _stage_pending_mixed_pause_responses( @@ -4423,6 +4664,12 @@ async def _resolve_approval_responses( pending_responses_before_binding = list( _collect_approval_responses(prepared_messages, non_approval_result_ids=host_result_ids).values() ) + pending_requests_before_staging = ( + _load_pending_approval_requests(approval_session) + if _has_authoritative_approval_session(approval_session) + else {} + ) + active_pending_ids = set(pending_requests_before_staging) if pending_requests_before_staging else None staged_responses, function_call_order, waiting_requests = _stage_approval_batch_responses( approval_session, pending_responses_before_binding, @@ -4436,16 +4683,8 @@ async def _resolve_approval_responses( streaming_updates=streaming_updates, ) - _bind_approval_responses_to_pending_requests(prepared_messages, approval_session) - active_pending_ids = ( - set(_load_pending_approval_requests(approval_session)) - if _has_authoritative_approval_session(approval_session) - else None - ) - if completed_mixed_batch: - state = _get_tool_approval_state(approval_session, create=False) - if state is not None: - state.pop(_PENDING_MIXED_PAUSE_BATCH_KEY, None) + if not completed_mixed_batch: + _bind_approval_responses_to_pending_requests(prepared_messages, approval_session) # 1. Restore safe siblings hidden with a prior mixed approval batch when its visible decision arrives. if staged_responses: @@ -4599,13 +4838,26 @@ async def _resolve_approval_responses( action = "return" elif reached_error_limit: action = "stop" - return _FunctionProcessingResult( + processing_result = _FunctionProcessingResult( errors_in_a_row=errors_in_a_row, action=action, function_call_count=executed_function_count, response_messages=response_messages, streaming_updates=streaming_updates, ) + if completed_mixed_batch and action in {"continue", "stop"}: + _save_pending_approval_requests(approval_session, pending_requests_before_staging) + _store_pending_provider_outbox( + approval_session, + prepared_messages=prepared_messages, + processing_result=processing_result, + ) + elif completed_mixed_batch: + state = _get_tool_approval_state(approval_session, create=False) + if state is not None: + state.pop(_PENDING_MIXED_PAUSE_BATCH_KEY, None) + state.pop(_PENDING_PROVIDER_OUTBOX_KEY, None) + return processing_result async def _process_model_function_calls( @@ -4832,6 +5084,7 @@ async def _get_response_with_function_invocation( from ._middleware import MiddlewareFailure from ._types import ChatResponse, add_usage_details + approval_session = invocation_session if approval_session_is_authoritative else None 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") @@ -4883,6 +5136,7 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non total_function_calls, approval_processing.function_call_count, ) + _persist_pending_provider_outbox_budget(approval_session, budget_state) if approval_processing.action == "return": response = ChatResponse(messages=list(function_call_messages)) response.usage_details = aggregated_usage @@ -4935,6 +5189,7 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non session=invocation_session, options=options, ) + _complete_pending_provider_outbox(approval_session) terminal_prefix_length = len(function_call_messages) try: @@ -5022,6 +5277,7 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non session=invocation_session, options=options, ) + _complete_pending_provider_outbox(approval_session) response.usage_details = aggregated_usage _prepend_function_call_messages(response, function_call_messages) _clear_budget_state_from_session(invocation_session) @@ -5047,6 +5303,7 @@ async def _stream_response_with_function_invocation( """Run the streaming function invocation loop.""" from ._middleware import MiddlewareFailure + approval_session = invocation_session if approval_session_is_authoritative else None 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") @@ -5095,6 +5352,9 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non total_function_calls, approval_processing.function_call_count, ) + _persist_pending_provider_outbox_budget(approval_session, budget_state) + if approval_processing.streaming_updates: + _mark_pending_provider_outbox_streaming_updates_published(approval_session) for update in approval_processing.streaming_updates: yield update if approval_processing.action == "return": @@ -5199,6 +5459,7 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non session=invocation_session, options=options, ) + _complete_pending_provider_outbox(approval_session) if not any( item.type == "function_approval_request" or _is_actionable_function_call(item) @@ -5292,6 +5553,7 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non session=invocation_session, options=options, ) + _complete_pending_provider_outbox(approval_session) if fallback_added: yield _function_invocation_limit_fallback_update() _clear_budget_state_from_session(invocation_session) @@ -5385,10 +5647,23 @@ 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 + approval_session_is_authoritative = ( + request_kwargs.pop(_APPROVAL_SESSION_IS_AUTHORITATIVE_KEY, True) is not False + ) + if invocation_session is None and requires_session_state: + invocation_session = _AgentSession() + approval_session_is_authoritative = False + 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 {} ) + approval_session = invocation_session if approval_session_is_authoritative else None + _restore_pending_provider_outbox_budget(approval_session, budget_state) # 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. @@ -5408,17 +5683,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 - approval_session_is_authoritative = ( - request_kwargs.pop(_APPROVAL_SESSION_IS_AUTHORITATIVE_KEY, True) is not False - ) - if invocation_session is None and requires_session_state: - invocation_session = _AgentSession() - approval_session_is_authoritative = False - # Bind one executor with the run's custom arguments, middleware, configuration, and session. execute_function_calls = partial( _execute_function_calls, 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 f97e80ea0f8..e5337470ac1 100644 --- a/python/packages/core/tests/core/test_function_invocation_logic.py +++ b/python/packages/core/tests/core/test_function_invocation_logic.py @@ -1186,16 +1186,23 @@ async def invalid_response(**kwargs: Any) -> ChatResponse: @pytest.mark.parametrize("streaming", [False, True], ids=["non_streaming", "streaming"]) +@pytest.mark.parametrize("authoritative", [True, False], ids=["authoritative", "temporary"]) async def test_response_invalidation_restores_structured_continuation_snapshot( chat_client_base: SupportsChatGetResponse, streaming: bool, + authoritative: bool, ) -> None: """Cleanup restores a copy when a provider mutates structured continuation state in place.""" from agent_framework._sessions import AgentSession + from agent_framework._tools import _APPROVAL_SESSION_IS_AUTHORITATIVE_KEY original_service_session_id = {"conversation_id": "valid", "metadata": {"generation": 1}} session = AgentSession(service_session_id=original_service_session_id) invalidated = ResponseInvalidatedException("provider invalidated partial response output") + client_kwargs = { + "session": session, + _APPROVAL_SESSION_IS_AUTHORITATIVE_KEY: authoritative, + } def mutate_continuation() -> None: service_session_id = session.service_session_id @@ -1222,7 +1229,7 @@ async def updates() -> AsyncIterable[ChatResponseUpdate]: [Message(role="user", contents=["run"])], options={"tools": []}, stream=True, - client_kwargs={"session": session}, + client_kwargs=client_kwargs, ) with pytest.raises(ResponseInvalidatedException): async for _ in stream: @@ -1239,7 +1246,7 @@ async def invalid_response(**kwargs: Any) -> ChatResponse: await chat_client_base.get_response( [Message(role="user", contents=["run"])], options={"tools": []}, - client_kwargs={"session": session}, + client_kwargs=client_kwargs, ) assert session.service_session_id == {"conversation_id": "valid", "metadata": {"generation": 1}} @@ -2327,9 +2334,10 @@ def second_write() -> str: for content in message.contents if content.type == "function_approval_request" ] - assert [ - request.function_call.name for request in approval_requests if request.function_call is not None - ] == ["first_write", "second_write"] + assert [request.function_call.name for request in approval_requests if request.function_call is not None] == [ + "first_write", + "second_write", + ] await chat_client_base.get_response( [ @@ -4708,10 +4716,66 @@ def test_stateless_mixed_batch_rejects_conflicting_identified_host_results() -> ), ] - with pytest.raises(RuntimeError, match="Conflicting response for mixed pause occurrence 'host-occurrence'"): + with pytest.raises(RuntimeError, match="Conflicting Host response.*host-occurrence"): _stateless_mixed_pause_batch_status(messages) +@pytest.mark.parametrize("identified", [False, True], ids=["idless", "identified"]) +def test_completed_stateless_mixed_batch_is_inert_when_later_turn_reuses_host_identity(identified: bool) -> None: + """A later id-less result with a reused call ID does not reopen a completed mixed batch.""" + from agent_framework._tools import _stateless_mixed_pause_batch_status + + approval_call = Content.from_function_call( + call_id="approval", + name="approval_func", + arguments={}, + id="approval-occurrence", + ) + approval_request = Content.from_function_approval_request( + id="approval-occurrence", + function_call=approval_call, + ) + host_request = Content.from_function_call( + call_id="reused-host", + name="host_func", + arguments={"round": 1}, + id="host-occurrence", + ) + host_request.user_input_request = True + host_result = Content.from_function_result(call_id="reused-host", result="first round") + host_result.id = "host-occurrence" + approval_result = Content.from_function_result(call_id="approval", result="approved") + later_host_request = Content.from_function_call( + call_id="reused-host", + name="host_func", + arguments={"round": 2}, + id="host-occurrence" if identified else "later-host-occurrence", + ) + later_host_request.user_input_request = True + later_host_result = Content.from_function_result(call_id="reused-host", result="second round") + if identified: + later_host_result.id = "host-occurrence" + messages = [ + Message(role="assistant", contents=[approval_request, host_request]), + Message( + role="user", + contents=[ + approval_request.to_function_approval_response(approved=True), + host_result, + ], + ), + Message(role="tool", contents=[approval_result]), + Message(role="assistant", contents=[later_host_request]), + Message(role="user", contents=[later_host_result]), + ] + + incomplete, host_result_ids = _stateless_mixed_pause_batch_status(messages) + + assert incomplete is False + assert host_result_ids == set() + assert messages[-1].contents == [later_host_result] + + def test_active_mixed_pause_ignores_historical_host_requests() -> None: """Only the session-recorded mixed batch participates in response correlation.""" from agent_framework._tools import ( @@ -4784,6 +4848,670 @@ def test_active_mixed_pause_ignores_historical_host_requests() -> None: assert host_result_ids == {id(messages[-1].contents[-1])} +@pytest.mark.parametrize("identified_first", [True, False], ids=["identified-first", "idless-first"]) +def test_stateful_mixed_pause_matches_equal_identified_and_idless_host_responses( + identified_first: bool, +) -> None: + """An identified Host result reserves its occurrence before an equal id-less sibling.""" + from agent_framework._tools import ( + _stage_pending_mixed_pause_responses, + _store_pending_approval_requests, + _store_pending_mixed_pause_batch, + ) + + session = AgentSession() + approval_call = Content.from_function_call( + call_id="approval-call", + name="approval_func", + arguments={}, + id="approval-occurrence", + ) + approval_request = Content.from_function_approval_request( + id="approval-occurrence", + function_call=approval_call, + ) + host_requests: list[Content] = [] + for occurrence in (1, 2): + request = Content.from_function_call( + call_id="reused-call", + name="host_func", + arguments={"value": occurrence}, + id=f"host-occurrence-{occurrence}", + ) + request.user_input_request = True + host_requests.append(request) + _store_pending_approval_requests(session, [approval_request]) + _store_pending_mixed_pause_batch(session, [[host_requests[0]], [approval_request], [host_requests[1]]]) + + identified = Content.from_function_result(call_id="reused-call", result={"same": True}) + identified.id = "host-occurrence-2" + idless = Content.from_function_result(call_id="reused-call", result={"same": True}) + host_results = [identified, idless] if identified_first else [idless, identified] + messages = [ + Message( + role="user", + contents=[approval_request.to_function_approval_response(approved=True), *host_results], + ) + ] + + incomplete, completed, host_result_ids = _stage_pending_mixed_pause_responses(messages, session) + + assert incomplete is False + assert completed is True + assert host_result_ids == {id(content) for content in messages[-1].contents if content.type == "function_result"} + assert [(content.type, content.id) for content in messages[-1].contents] == [ + ("function_result", None), + ("function_approval_response", "approval-occurrence"), + ("function_result", "host-occurrence-2"), + ] + + +def test_stateful_mixed_pause_matches_cross_turn_equal_identified_and_idless_host_responses() -> None: + """A later id-less result fills the sole unanswered occurrence instead of replaying an answered one.""" + from agent_framework._tools import ( + _stage_pending_mixed_pause_responses, + _store_pending_approval_requests, + _store_pending_mixed_pause_batch, + ) + + session = AgentSession() + approval_call = Content.from_function_call( + call_id="approval-call", + name="approval_func", + arguments={}, + id="approval-occurrence", + ) + approval_request = Content.from_function_approval_request( + id="approval-occurrence", + function_call=approval_call, + ) + host_requests: list[Content] = [] + for occurrence in (1, 2): + request = Content.from_function_call( + call_id="reused-call", + name="host_func", + arguments={"value": occurrence}, + id=f"host-occurrence-{occurrence}", + ) + request.user_input_request = True + host_requests.append(request) + _store_pending_approval_requests(session, [approval_request]) + _store_pending_mixed_pause_batch(session, [[host_requests[0]], [approval_request], [host_requests[1]]]) + + identified = Content.from_function_result(call_id="reused-call", result={"same": True}) + identified.id = "host-occurrence-2" + approval_response = approval_request.to_function_approval_response(approved=True) + first_turn = [Message(role="user", contents=[approval_response, identified])] + incomplete, completed, _ = _stage_pending_mixed_pause_responses(first_turn, session) + assert incomplete is True + assert completed is False + + idless = Content.from_function_result(call_id="reused-call", result={"same": True}) + second_turn = [ + Message(role="user", contents=[approval_response, identified]), + Message(role="user", contents=[idless]), + ] + incomplete, completed, host_result_ids = _stage_pending_mixed_pause_responses(second_turn, session) + + assert incomplete is False + assert completed is True + assert host_result_ids == {id(content) for content in second_turn[-1].contents if content.type == "function_result"} + assert [(content.type, content.id) for content in second_turn[-1].contents] == [ + ("function_result", None), + ("function_approval_response", "approval-occurrence"), + ("function_result", "host-occurrence-2"), + ] + + +def test_stateless_conflicting_host_results_for_same_occurrence_fail_closed() -> None: + """Conflicting identified Host results cannot both claim one stateless occurrence.""" + from agent_framework._tools import _stateless_mixed_pause_batch_status + + approval_call = Content.from_function_call( + call_id="approval-call", + name="approval_func", + arguments={}, + id="approval-occurrence", + ) + approval_request = Content.from_function_approval_request( + id="approval-occurrence", + function_call=approval_call, + ) + host_request = Content.from_function_call( + call_id="host-call", + name="host_func", + arguments={}, + id="host-occurrence", + ) + host_request.user_input_request = True + first = Content.from_function_result(call_id="host-call", result="first") + first.id = "host-occurrence" + conflicting = Content.from_function_result(call_id="host-call", result="second") + conflicting.id = "host-occurrence" + + with pytest.raises(RuntimeError, match="Conflicting Host response.*host-occurrence"): + _stateless_mixed_pause_batch_status([ + Message(role="assistant", contents=[approval_call, approval_request, host_request]), + Message( + role="user", + contents=[ + approval_request.to_function_approval_response(approved=True), + first, + conflicting, + ], + ), + ]) + + +@pytest.mark.parametrize("stateful", [False, True], ids=["stateless", "stateful"]) +def test_extra_idless_conflicting_host_result_fails_closed(stateful: bool) -> None: + """An extra id-less result cannot bypass conflict detection after all Host occurrences are answered.""" + from agent_framework._tools import ( + _stage_pending_mixed_pause_responses, + _stateless_mixed_pause_batch_status, + _store_pending_approval_requests, + _store_pending_mixed_pause_batch, + ) + + approval_call = Content.from_function_call( + call_id="approval-call", + name="approval_func", + arguments={}, + id="approval-occurrence", + ) + approval_request = Content.from_function_approval_request( + id="approval-occurrence", + function_call=approval_call, + ) + host_requests: list[Content] = [] + for occurrence in (1, 2): + request = Content.from_function_call( + call_id="reused-call", + name="host_func", + arguments={"value": occurrence}, + id=f"host-occurrence-{occurrence}", + ) + request.user_input_request = True + host_requests.append(request) + first = Content.from_function_result(call_id="reused-call", result="first") + second = Content.from_function_result(call_id="reused-call", result="second") + second.id = "host-occurrence-2" + conflicting = Content.from_function_result(call_id="reused-call", result="conflicting") + response_message = Message( + role="user", + contents=[ + approval_request.to_function_approval_response(approved=True), + first, + second, + conflicting, + ], + ) + + with pytest.raises(RuntimeError, match="Conflicting Host response.*reused-call"): + if stateful: + session = AgentSession() + _store_pending_approval_requests(session, [approval_request]) + _store_pending_mixed_pause_batch( + session, + [[host_requests[0]], [approval_request], [host_requests[1]]], + ) + _stage_pending_mixed_pause_responses([response_message], session) + else: + _stateless_mixed_pause_batch_status([ + Message(role="assistant", contents=[approval_request, *host_requests]), + response_message, + ]) + + +@pytest.mark.parametrize("stateful", [False, True], ids=["stateless", "stateful"]) +@pytest.mark.parametrize("equal_payloads", [False, True], ids=["distinct-payloads", "equal-payloads"]) +def test_extra_idless_equivalent_host_result_is_deduplicated(stateful: bool, equal_payloads: bool) -> None: + """An extra id-less replay is removed when it is equivalent to an accepted Host result.""" + from agent_framework._tools import ( + _stage_pending_mixed_pause_responses, + _stateless_mixed_pause_batch_status, + _store_pending_approval_requests, + _store_pending_mixed_pause_batch, + ) + + approval_call = Content.from_function_call( + call_id="approval-call", + name="approval_func", + arguments={}, + id="approval-occurrence", + ) + approval_request = Content.from_function_approval_request( + id="approval-occurrence", + function_call=approval_call, + ) + host_requests: list[Content] = [] + for occurrence in (1, 2): + request = Content.from_function_call( + call_id="reused-call", + name="host_func", + arguments={"value": occurrence}, + id=f"host-occurrence-{occurrence}", + ) + request.user_input_request = True + host_requests.append(request) + first = Content.from_function_result(call_id="reused-call", result="first") + second_result = "first" if equal_payloads else "second" + second = Content.from_function_result(call_id="reused-call", result=second_result) + second.id = "host-occurrence-2" + replay = Content.from_function_result(call_id="reused-call", result="first") + response_message = Message( + role="user", + contents=[ + approval_request.to_function_approval_response(approved=True), + first, + second, + replay, + ], + ) + messages = [response_message] + + if stateful: + session = AgentSession() + _store_pending_approval_requests(session, [approval_request]) + _store_pending_mixed_pause_batch( + session, + [[host_requests[0]], [approval_request], [host_requests[1]]], + ) + incomplete, completed, _ = _stage_pending_mixed_pause_responses(messages, session) + assert incomplete is False + assert completed is True + else: + messages.insert(0, Message(role="assistant", contents=[approval_request, *host_requests])) + incomplete, _ = _stateless_mixed_pause_batch_status(messages) + assert incomplete is False + + normalized_results = [ + content + for message in messages + for content in message.contents + if content.type == "function_result" and content.call_id == "reused-call" + ] + assert [(result.id, result.result) for result in normalized_results] == [ + (None, "first"), + ("host-occurrence-2", second_result), + ] + + +@pytest.mark.parametrize("streaming", [False, True], ids=["non_streaming", "streaming"]) +async def test_completed_mixed_batch_replays_serialized_outbox_after_provider_invalidation( + chat_client_base: SupportsChatGetResponse, + streaming: bool, +) -> None: + """Invalidated delivery replays persisted results without repeating the approved side effect.""" + from agent_framework import FunctionTool + from agent_framework._tools import ( + _FUNCTION_INVOCATION_BUDGET_STATE_KEY, + _PENDING_APPROVAL_REQUESTS_KEY, + _PENDING_MIXED_PAUSE_BATCH_KEY, + _PENDING_PROVIDER_OUTBOX_KEY, + _TOOL_APPROVAL_STATE_KEY, + ) + + approved_calls = 0 + provider_inputs: list[list[Message]] = [] + published_results: list[tuple[str | None, Any]] = [] + invalidated = ResponseInvalidatedException("provider invalidated result delivery") + + @tool(name="approval_func", approval_mode="always_require") + def approval_func() -> str: + nonlocal approved_calls + approved_calls += 1 + return "approved result" + + host_func = FunctionTool(name="host_func", func=None, description="A Host-owned function") + host_call = Content.from_function_call( + call_id="host-call", + name="host_func", + arguments={}, + id="host-occurrence", + ) + approval_call = Content.from_function_call( + call_id="approval-call", + name="approval_func", + arguments={}, + id="approval-occurrence", + ) + + def record_request(messages: Sequence[Message]) -> int: + provider_inputs.append([Message.from_dict(message.to_dict()) for message in messages]) + return len(provider_inputs) + + if streaming: + + def scripted_stream( + *, + messages: Sequence[Message], + **kwargs: Any, + ) -> ResponseStream[ChatResponseUpdate, ChatResponse]: + del kwargs + call_number = record_request(messages) + + async def updates() -> AsyncIterable[ChatResponseUpdate]: + if call_number == 1: + yield ChatResponseUpdate( + role="assistant", + contents=[host_call, approval_call], + finish_reason="tool_calls", + conversation_id="mixed-continuation", + ) + elif call_number == 2: + raise invalidated + yield # pragma: no cover + else: + yield ChatResponseUpdate( + role="assistant", + contents=[Content.from_text("done")], + finish_reason="stop", + conversation_id="completed-continuation", + ) + + return ResponseStream(updates(), finalizer=ChatResponse.from_updates) + + chat_client_base._get_streaming_response = scripted_stream # type: ignore[attr-defined, method-assign] # ty: ignore[unresolved-attribute] + else: + + async def scripted_response( + *, + messages: Sequence[Message], + **kwargs: Any, + ) -> ChatResponse: + del kwargs + call_number = record_request(messages) + if call_number == 1: + return ChatResponse( + messages=Message(role="assistant", contents=[host_call, approval_call]), + finish_reason="tool_calls", + conversation_id="mixed-continuation", + ) + if call_number == 2: + raise invalidated + return ChatResponse( + messages=Message(role="assistant", contents=["done"]), + finish_reason="stop", + conversation_id="completed-continuation", + ) + + chat_client_base._get_non_streaming_response = scripted_response # type: ignore[attr-defined, method-assign] # ty: ignore[unresolved-attribute] + + chat_client_base.function_invocation_configuration["max_function_calls"] = 1 # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + session = AgentSession() + + async def run(contents: list[Content]) -> ChatResponse: + options: ChatOptions = {"tool_choice": "auto", "tools": [approval_func, host_func]} + if isinstance(session.service_session_id, str): + options["conversation_id"] = session.service_session_id + if not streaming: + return await chat_client_base.get_response( + [Message(role="user", contents=contents)], + options=options, + client_kwargs={"session": session}, + ) + response_stream = chat_client_base.get_response( + [Message(role="user", contents=contents)], + stream=True, + options=options, + client_kwargs={"session": session}, + ) + async for update in response_stream: + published_results.extend( + (content.call_id, content.result) for content in update.contents if content.type == "function_result" + ) + return await response_stream.get_final_response() + + first_response = await run([Content.from_text("go")]) + approval_request = next( + content + for message in first_response.messages + for content in message.contents + if content.type == "function_approval_request" + ) + host_request = next( + content + for message in first_response.messages + for content in message.contents + if content.type == "function_call" and content.user_input_request + ) + host_result = Content.from_function_result(call_id="host-call", result="host result") + host_result.id = host_request.id + + with pytest.raises(ResponseInvalidatedException): + await run([approval_request.to_function_approval_response(approved=True), host_result]) + assert approved_calls == 1 + tool_state = session.state[_TOOL_APPROVAL_STATE_KEY] + assert isinstance(tool_state, dict) + assert _PENDING_MIXED_PAUSE_BATCH_KEY in tool_state + assert _PENDING_APPROVAL_REQUESTS_KEY in tool_state + assert _PENDING_PROVIDER_OUTBOX_KEY in tool_state + budget_state = session.state[_FUNCTION_INVOCATION_BUDGET_STATE_KEY] + assert isinstance(budget_state, dict) + assert budget_state["total_function_calls"] == 1 + + session = AgentSession.from_dict(json.loads(json.dumps(session.to_dict()))) + final_response = await run([Content.from_text("retry")]) + + assert final_response.text == "done" + assert approved_calls == 1 + assert len(provider_inputs) == 3 + delivered_results = [ + [ + (content.call_id, content.result) + for message in request_messages + for content in message.contents + if content.type == "function_result" + ] + for request_messages in provider_inputs[1:] + ] + assert delivered_results == [ + [("host-call", "host result"), ("approval-call", "approved result")], + [("host-call", "host result"), ("approval-call", "approved result")], + ] + assert published_results == ([("approval-call", "approved result")] if streaming else []) + final_tool_state = session.state[_TOOL_APPROVAL_STATE_KEY] + assert isinstance(final_tool_state, dict) + assert _PENDING_MIXED_PAUSE_BATCH_KEY not in final_tool_state + assert _PENDING_PROVIDER_OUTBOX_KEY not in final_tool_state + assert _FUNCTION_INVOCATION_BUDGET_STATE_KEY not in session.state + + +@pytest.mark.parametrize( + "restored_time_ns", + [1_005_000_000_000, 995_000_000_000], + ids=["wall-clock-advanced", "wall-clock-rewound"], +) +def test_serialized_provider_outbox_rebases_duration_budget_across_monotonic_epochs( + monkeypatch: pytest.MonkeyPatch, + restored_time_ns: int, +) -> None: + """A restored outbox preserves elapsed duration without comparing unrelated monotonic clocks.""" + from agent_framework import _tools + from agent_framework._tools import ( + _PENDING_PROVIDER_OUTBOX_KEY, + _TOOL_APPROVAL_STATE_KEY, + _apply_batch_limit_decision, + _persist_pending_provider_outbox_budget, + _restore_pending_provider_outbox_budget, + ) + + session = AgentSession() + session.state[_TOOL_APPROVAL_STATE_KEY] = { + _PENDING_PROVIDER_OUTBOX_KEY: {}, + } + budget_state: dict[str, Any] = { + "start_time": 100.0, + "attempt_count": 1, + "total_function_calls": 1, + } + monkeypatch.setattr(_tools, "perf_counter", lambda: 110.0) + monkeypatch.setattr(_tools, "time_ns", lambda: 1_000_000_000_000) + _persist_pending_provider_outbox_budget(session, budget_state) + + restored_budget: dict[str, Any] = {} + monkeypatch.setattr(_tools, "perf_counter", lambda: 5.0) + monkeypatch.setattr(_tools, "time_ns", lambda: restored_time_ns) + _restore_pending_provider_outbox_budget(session, restored_budget) + options: dict[str, Any] = {"tool_choice": "auto"} + _apply_batch_limit_decision( + "continue", + options, + restored_budget, + total_function_calls=1, + max_function_calls=None, + max_duration_seconds=12.0, + ) + + assert options["tool_choice"] == "none" + assert restored_budget["truncated"] is True + _persist_pending_provider_outbox_budget(session, restored_budget) + json.dumps(session.to_dict(), allow_nan=False) + + +@pytest.mark.parametrize( + "serialized_budget", + [ + {"start_time": 1_000_000.0, "attempt_count": 1, "total_function_calls": 1}, + { + "elapsed_duration_seconds": -1.0, + "saved_at_time_ns": 1_000, + "attempt_count": 1, + "total_function_calls": 1, + }, + { + "elapsed_duration_seconds": 1.0, + "saved_at_time_ns": 2_000, + "attempt_count": 1, + "total_function_calls": 1, + }, + ], + ids=["legacy-monotonic", "negative-elapsed", "future-wall-clock"], +) +def test_untrusted_provider_outbox_duration_budget_fails_closed_after_restore( + monkeypatch: pytest.MonkeyPatch, + serialized_budget: dict[str, Any], +) -> None: + """Untrusted duration metadata cannot relax the invocation budget.""" + from agent_framework import _tools + from agent_framework._tools import ( + _PENDING_PROVIDER_OUTBOX_KEY, + _TOOL_APPROVAL_STATE_KEY, + _apply_batch_limit_decision, + _restore_pending_provider_outbox_budget, + ) + + session = AgentSession() + session.state[_TOOL_APPROVAL_STATE_KEY] = { + _PENDING_PROVIDER_OUTBOX_KEY: { + "budget_state": serialized_budget, + }, + } + monkeypatch.setattr(_tools, "perf_counter", lambda: 5.0) + monkeypatch.setattr(_tools, "time_ns", lambda: 1_000) + restored_budget: dict[str, Any] = {} + _restore_pending_provider_outbox_budget(session, restored_budget) + options: dict[str, Any] = {"tool_choice": "auto"} + + _apply_batch_limit_decision( + "continue", + options, + restored_budget, + total_function_calls=1, + max_function_calls=None, + max_duration_seconds=12.0, + ) + + assert options["tool_choice"] == "none" + assert restored_budget["truncated"] is True + + +@pytest.mark.parametrize( + "serialized_budget", + [ + None, + [], + {"attempt_count": -1, "total_function_calls": 1}, + {"attempt_count": 1, "total_function_calls": False}, + {"attempt_count": "1", "total_function_calls": 1}, + ], + ids=["missing", "non-mapping", "negative-attempt", "boolean-total", "string-attempt"], +) +def test_malformed_provider_outbox_budget_is_rejected(serialized_budget: Any) -> None: + """A pending outbox cannot replay with missing or invalid charged counters.""" + from agent_framework._tools import ( + _PENDING_PROVIDER_OUTBOX_KEY, + _TOOL_APPROVAL_STATE_KEY, + _restore_pending_provider_outbox_budget, + ) + + session = AgentSession() + session.state[_TOOL_APPROVAL_STATE_KEY] = { + _PENDING_PROVIDER_OUTBOX_KEY: { + "budget_state": serialized_budget, + }, + } + + with pytest.raises(RuntimeError, match="provider outbox contains an invalid budget"): + _restore_pending_provider_outbox_budget(session, {}) + + +@pytest.mark.parametrize("errors_in_a_row", [-1, False, "1", 1.5]) +def test_malformed_provider_outbox_error_counter_is_rejected(errors_in_a_row: Any) -> None: + """A restored outbox cannot weaken the consecutive-error limit with an invalid counter.""" + from agent_framework._tools import ( + _PENDING_PROVIDER_OUTBOX_KEY, + _TOOL_APPROVAL_STATE_KEY, + _FunctionProcessingResult, + _restore_pending_provider_outbox, + _store_pending_provider_outbox, + ) + + session = AgentSession() + _store_pending_provider_outbox( + session, + prepared_messages=[Message(role="user", contents=["provider input"])], + processing_result=_FunctionProcessingResult( + errors_in_a_row=1, + response_messages=(Message(role="tool", contents=["tool result"]),), + ), + ) + approval_state = session.state[_TOOL_APPROVAL_STATE_KEY] + approval_state[_PENDING_PROVIDER_OUTBOX_KEY]["errors_in_a_row"] = errors_in_a_row + + with pytest.raises(RuntimeError, match="provider outbox has an invalid error counter"): + _restore_pending_provider_outbox([], session) + + +def test_missing_provider_outbox_error_counter_defaults_to_zero() -> None: + """An older outbox without an error counter resumes from the documented default.""" + from agent_framework._tools import ( + _PENDING_PROVIDER_OUTBOX_KEY, + _TOOL_APPROVAL_STATE_KEY, + _FunctionProcessingResult, + _restore_pending_provider_outbox, + _store_pending_provider_outbox, + ) + + session = AgentSession() + _store_pending_provider_outbox( + session, + prepared_messages=[Message(role="user", contents=["provider input"])], + processing_result=_FunctionProcessingResult( + errors_in_a_row=1, + response_messages=(Message(role="tool", contents=["tool result"]),), + ), + ) + approval_state = session.state[_TOOL_APPROVAL_STATE_KEY] + del approval_state[_PENDING_PROVIDER_OUTBOX_KEY]["errors_in_a_row"] + + restored = _restore_pending_provider_outbox([], session) + + assert restored is not None + assert restored.errors_in_a_row == 0 + + async def test_function_invocation_config_additional_tools(chat_client_base: SupportsChatGetResponse): """Test that additional_tools are available but treated as declaration_only.""" exec_counter_visible = 0