From 06daf9586c713584d37c7ac8da9de72a96235adc Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Mon, 10 Aug 2026 13:27:09 +0900 Subject: [PATCH 01/13] Route local approvals through lifecycle owner Key decisions: - Add an internal typed approval lifecycle with pending, claimed, executing, and settled states. - Keep authorization separate from execution; only LocalPendingToolTransitionOwner invokes approved local calls. - Register server-owned occurrences before canonical ResumeDecision claims and retain one replayable result under the original call identity. Files changed: - packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py - packages/ag-ui/agent_framework_ag_ui/_approval_state.py - packages/ag-ui/agent_framework_ag_ui/_agent_run.py - packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py - packages/ag-ui/tests/ag_ui/test_approval_result_event.py - packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py Verification: - 952 AG-UI tests passed. - Focused lifecycle/public tracer passed with warnings treated as errors. - Ruff format/check and AG-UI Pyright passed. - git diff --check passed. Notes for next iteration: - The function-calling-loop scenario mapping is inaccessible under the organization content-exclusion policy and could not be updated. - The workspace Poe package fan-out is blocked by the pre-existing missing packages/durabletask/pyproject.toml; equivalent package-local checks were run. --- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 183 +++++++++++++++--- .../_approval_lifecycle.py | 177 +++++++++++++++++ .../agent_framework_ag_ui/_approval_state.py | 40 ++++ .../ag_ui/test_agent_wrapper_comprehensive.py | 30 +-- .../tests/ag_ui/test_approval_lifecycle.py | 54 ++++++ .../tests/ag_ui/test_approval_result_event.py | 135 +++++++++---- 6 files changed, 541 insertions(+), 78 deletions(-) create mode 100644 python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py create mode 100644 python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index a19320c31e..3b604eacaa 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py @@ -55,6 +55,12 @@ _use_telemetry_conversation_id, # pyright: ignore[reportPrivateUsage] ) +from ._approval_lifecycle import ( + ApprovalLifecycle, + AuthorizedExecution, + LocalPendingToolTransitionOwner, + ResumeDecision, +) from ._approval_state import _APPROVAL_SCOPE_INPUT_KEY, InMemoryAGUIApprovalStateStore, approval_state_thread_id from ._message_adapters import normalize_agui_input_messages from ._predictive_state import PredictiveStateHandler @@ -915,6 +921,9 @@ def _pop_collected_tool_approval_response_messages( session: AgentSession, pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] | None, thread_id: str, + *, + lifecycle: ApprovalLifecycle | None = None, + authorized_executions: dict[str, AuthorizedExecution] | None = None, ) -> list[Message]: """Pop server-collected auto-approved responses into provider-visible messages.""" raw_state = session.state.get(_TOOL_APPROVAL_STATE_KEY) @@ -932,6 +941,33 @@ def _pop_collected_tool_approval_response_messages( if response is None or response.type != "function_approval_response": continue _register_server_generated_approval_response(response, pending_approvals, thread_id) + function_call = response.function_call + if ( + lifecycle is not None + and authorized_executions is not None + and response.approved + and function_call is not None + and function_call.call_id + and function_call.name + and not _function_call_server_label(function_call) + ): + arguments = canonical_function_arguments(function_call) or "{}" + lifecycle.register_local( + thread_id=thread_id, + interrupt_id=str(response.id or function_call.call_id), + call_id=str(function_call.call_id), + name=function_call.name, + arguments=arguments, + ) + intent = lifecycle.claim( + thread_id=thread_id, + decision=ResumeDecision( + interrupt_id=str(response.id or function_call.call_id), + accepted=True, + arguments=arguments, + ), + ) + authorized_executions[intent.identity.call_id] = intent responses.append(response) state[_COLLECTED_APPROVAL_RESPONSES_KEY] = [] @@ -1194,6 +1230,9 @@ def _canonical_approval_resume_messages( pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] | None, thread_id: str, expected_interrupt_ids: set[str] | None = None, + *, + lifecycle: ApprovalLifecycle | None = None, + authorized_executions: dict[str, AuthorizedExecution] | None = None, ) -> tuple[list[dict[str, Any]], set[str], set[str], RunErrorEvent | None]: """Translate canonical ResumeEntry approvals into existing approval response messages.""" expected_ids = set(expected_interrupt_ids or set()) @@ -1321,6 +1360,7 @@ def _canonical_approval_resume_messages( ) argument_updates: list[tuple[_PendingApprovalWithSiblings, str]] = [] + lifecycle_decisions: list[ResumeDecision] = [] restored_sibling_response_ids: set[str] = set() for entry in entries: interrupt_id = cast(str, entry["interrupt_id"]) @@ -1393,9 +1433,17 @@ def _canonical_approval_resume_messages( ) merged_arguments = {**original_arguments, **edited_arguments} + canonical_arguments = json.dumps(make_json_safe(merged_arguments), sort_keys=True, separators=(",", ":")) if not isinstance(pending_entry, str): - argument_updates.append( - (pending_entry, json.dumps(make_json_safe(merged_arguments), sort_keys=True, separators=(",", ":"))) + argument_updates.append((pending_entry, canonical_arguments)) + if lifecycle is not None and not _pending_approval_server_label(pending_entry): + lifecycle_decisions.append( + ResumeDecision( + interrupt_id=interrupt_id, + accepted=accepted, + arguments=canonical_arguments, + original_arguments=pending_arguments, + ) ) function_approvals = [ { @@ -1434,6 +1482,24 @@ def _canonical_approval_resume_messages( str(response_id), str(function_call.call_id) if function_call.call_id else None, ) + if lifecycle is not None and not _function_call_server_label(function_call): + sibling_interrupt_id = str(response_id) + sibling_call_id = str(function_call.call_id or response_id) + sibling_arguments = canonical_function_arguments(function_call) or "{}" + lifecycle.register_local( + thread_id=thread_id, + interrupt_id=sibling_interrupt_id, + call_id=sibling_call_id, + name=function_call.name, + arguments=sibling_arguments, + ) + lifecycle_decisions.append( + ResumeDecision( + interrupt_id=sibling_interrupt_id, + accepted=True, + arguments=sibling_arguments, + ) + ) function_approvals.append( { "id": str(response_id), @@ -1448,6 +1514,20 @@ def _canonical_approval_resume_messages( for pending_entry, arguments_json in argument_updates: pending_entry["arguments"] = arguments_json + if lifecycle is not None and authorized_executions is not None: + try: + for decision in lifecycle_decisions: + if decision.accepted: + intent = lifecycle.claim(thread_id=thread_id, decision=decision) + authorized_executions[intent.identity.call_id] = intent + except (KeyError, ValueError) as exc: + return ( + [], + handled_ids, + cancelled_ids, + RunErrorEvent(message=str(exc), code="APPROVAL_RESUME_INVALID"), + ) + return messages, handled_ids, cancelled_ids, None @@ -1472,6 +1552,9 @@ async def _resolve_approval_responses( pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] | None = None, thread_id: str = "", validated_approved_responses: list[Content] | None = None, + *, + lifecycle: ApprovalLifecycle | None = None, + authorized_executions: dict[str, AuthorizedExecution] | None = None, ) -> list[Content]: """Execute approved function calls and replace approval content with results. @@ -1672,27 +1755,43 @@ def matches_pending_entry(candidate: PendingApprovalEntry | None, expected: Pend if tool_name in tool_map and not _is_hosted_tool_approval(approval): static_approved.append(approval) - # Execute only statically-available approved tool calls - if static_approved and tools: + # Execute lifecycle-authorized local calls only through their transition owner. + if static_approved and tools and lifecycle is not None and authorized_executions is not None: client = getattr(agent, "client", None) config = normalize_function_invocation_configuration(getattr(client, "function_invocation_configuration", None)) middleware_pipeline = FunctionMiddlewarePipeline( *getattr(client, "function_middleware", ()), *run_kwargs.get("middleware", ()), ) - # Filter out AG-UI-specific kwargs that should not be passed to tool execution tool_kwargs = {k: v for k, v in run_kwargs.items() if k != "options"} - try: - approved_function_result_groups, _ = await _try_execute_function_call_groups( - custom_args=tool_kwargs, - function_calls=static_approved, - tools=tools, - middleware_pipeline=middleware_pipeline, - config=config, - ) - except Exception as e: - logger.exception("Failed to execute approved tool calls; injecting error results: %s", e) - approved_function_result_groups = [] + for approval in static_approved: + function_call = approval.function_call + call_id = (function_call.call_id if function_call else None) or approval.id or "" + intent = authorized_executions.get(call_id) + if intent is None: + logger.warning("Skipping local approval without lifecycle authority for call_id=%s.", call_id) + approved_function_result_groups.append([]) + continue + + async def execute_local_call(approval: Content = approval, call_id: str = call_id) -> list[Content]: + try: + result_groups, _ = await _try_execute_function_call_groups( + custom_args=tool_kwargs, + function_calls=[approval], + tools=tools, + middleware_pipeline=middleware_pipeline, + config=config, + ) + except Exception as exc: + logger.exception("Failed to execute approved tool call; injecting error result: %s", exc) + return [Content.from_function_result(call_id=call_id, result="Error: Tool call invocation failed.")] + if not result_groups or not result_groups[0]: + return [Content.from_function_result(call_id=call_id, result="Error: Tool call invocation failed.")] + return result_groups[0] + + owner = LocalPendingToolTransitionOwner(execute_local_call) + outcome = await owner.execute(intent, lifecycle=lifecycle) + approved_function_result_groups.append(list(outcome.result_group)) # Normalize one group per static approval and collect only terminal results for TOOL_CALL_RESULT events. # Deferred provider-injected approvals are left in messages for ToolApprovalMiddleware to process. @@ -2231,6 +2330,8 @@ async def run_agent_stream( pending_approvals, approval_thread_id, expected_interrupt_ids=stored_pending_approval_interrupt_ids or None, + lifecycle=approval_state_store.lifecycle if approval_state_store is not None else None, + authorized_executions=(authorized_executions := {}), ) ) if resume_error is not None: @@ -2331,7 +2432,15 @@ async def run_agent_stream( # Resolve approval responses (execute approved tools, replace approvals with results) # This must happen before running the agent so it sees the tool results tools_for_execution = tools if tools is not None else server_tools - messages.extend(_pop_collected_tool_approval_response_messages(session, pending_approvals, approval_thread_id)) + messages.extend( + _pop_collected_tool_approval_response_messages( + session, + pending_approvals, + approval_thread_id, + lifecycle=approval_state_store.lifecycle if approval_state_store is not None else None, + authorized_executions=authorized_executions, + ) + ) validated_approved_responses: list[Content] = [] resolved_approval_results = await _resolve_approval_responses( messages, @@ -2341,6 +2450,8 @@ async def run_agent_stream( pending_approvals, approval_thread_id, validated_approved_responses, + lifecycle=approval_state_store.lifecycle if approval_state_store is not None else None, + authorized_executions=authorized_executions, ) # Defense-in-depth: replace approval payloads in snapshot with actual tool results @@ -2467,20 +2578,32 @@ async def run_agent_stream( scope=approval_scope, thread_id=provider_thread_id or thread_id, ) - _register_pending_approval( - pending_approvals, - [approval_thread_id, provider_approval_thread_id], - content.function_call.name, - canonical_function_arguments(content.function_call), - request_id=str(content.id), - interrupt_id=str(canonical_interrupt_id), - already_approved_requests=_stored_already_approved_requests_for_visible_approval( - session, - str(content.id), - str(canonical_interrupt_id) if canonical_interrupt_id else None, - ), - server_label=_function_call_server_label(content.function_call), + server_label = _function_call_server_label(content.function_call) + already_approved_requests = _stored_already_approved_requests_for_visible_approval( + session, + str(content.id), + str(canonical_interrupt_id) if canonical_interrupt_id else None, ) + if approval_state_store is not None and not server_label: + approval_state_store.register_local( + thread_ids=[approval_thread_id, provider_approval_thread_id], + name=content.function_call.name, + arguments=canonical_function_arguments(content.function_call) or "{}", + request_id=str(content.id), + interrupt_id=str(canonical_interrupt_id), + already_approved_requests=already_approved_requests, + ) + else: + _register_pending_approval( + pending_approvals, + [approval_thread_id, provider_approval_thread_id], + content.function_call.name, + canonical_function_arguments(content.function_call), + request_id=str(content.id), + interrupt_id=str(canonical_interrupt_id), + already_approved_requests=already_approved_requests, + server_label=server_label, + ) # Evict oldest entries if the registry exceeds a safe bound (LRU) _evict_oldest_approvals(pending_approvals, max_size=10_000) else: diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py b/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py new file mode 100644 index 0000000000..5e0823d355 --- /dev/null +++ b/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py @@ -0,0 +1,177 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Server-owned lifecycle for AG-UI approval-gated tool calls.""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field +from enum import Enum + +from agent_framework import Content + + +class ApprovalStatus(str, Enum): + """Lifecycle state of one server-owned approval occurrence.""" + + PENDING = "pending" + CLAIMED = "claimed" + EXECUTING = "executing" + SETTLED = "settled" + + +@dataclass(frozen=True) +class ApprovalOccurrenceIdentity: + """Identity of one approval occurrence within a server-owned thread.""" + + thread_id: str + interrupt_id: str + call_id: str + + +@dataclass(frozen=True) +class ResumeDecision: + """Canonical client decision presented to the approval lifecycle.""" + + interrupt_id: str + accepted: bool + arguments: str + original_arguments: str | None = None + + +@dataclass +class ApprovalOccurrence: + """Server-owned state for one approval-gated call occurrence.""" + + identity: ApprovalOccurrenceIdentity + name: str + arguments: str + status: ApprovalStatus = ApprovalStatus.PENDING + replayable_results: list[ReplayableToolResult] = field(default_factory=list) + + +@dataclass(frozen=True) +class AuthorizedExecution: + """Authority for a Pending Tool Transition Owner to execute one local call.""" + + identity: ApprovalOccurrenceIdentity + name: str + arguments: str + + +@dataclass(frozen=True) +class ReplayableToolResult: + """A settled tool result retained under its original call identity.""" + + content: Content + + +@dataclass(frozen=True) +class ApprovalOutcome: + """Terminal outcome retained for one approval occurrence.""" + + identity: ApprovalOccurrenceIdentity + replayable_results: tuple[ReplayableToolResult, ...] + result_group: tuple[Content, ...] + + +class ApprovalLifecycle: + """Own registration, authority transitions, and settlement for approvals.""" + + def __init__(self) -> None: + self._occurrences: dict[ApprovalOccurrenceIdentity, ApprovalOccurrence] = {} + self._pending_by_interrupt: dict[tuple[str, str], ApprovalOccurrenceIdentity] = {} + + def register_local( + self, + *, + thread_id: str, + interrupt_id: str, + call_id: str, + name: str, + arguments: str, + ) -> ApprovalOccurrence: + """Register one server-generated local approval occurrence.""" + identity = ApprovalOccurrenceIdentity( + thread_id=thread_id, + interrupt_id=interrupt_id, + call_id=call_id, + ) + occurrence = ApprovalOccurrence(identity=identity, name=name, arguments=arguments) + self._occurrences[identity] = occurrence + self._pending_by_interrupt[(thread_id, interrupt_id)] = identity + return occurrence + + def get(self, identity: ApprovalOccurrenceIdentity) -> ApprovalOccurrence: + """Return server-owned state for a registered occurrence.""" + return self._occurrences[identity] + + def claim(self, *, thread_id: str, decision: ResumeDecision) -> AuthorizedExecution: + """Validate and reserve one accepted decision before execution.""" + identity = self._pending_by_interrupt[(thread_id, decision.interrupt_id)] + occurrence = self._occurrences[identity] + if occurrence.status is not ApprovalStatus.PENDING: + raise ValueError(f"Approval occurrence is not pending: {occurrence.status}.") + if not decision.accepted: + raise ValueError("A rejected decision cannot authorize execution.") + if (decision.original_arguments or decision.arguments) != occurrence.arguments: + raise ValueError("Approval decision arguments do not match the registered occurrence.") + occurrence.arguments = decision.arguments + occurrence.status = ApprovalStatus.CLAIMED + return AuthorizedExecution(identity=identity, name=occurrence.name, arguments=occurrence.arguments) + + def begin_execution(self, intent: AuthorizedExecution) -> None: + """Mark a claimed occurrence immediately before its owner may invoke a tool.""" + occurrence = self._occurrences[intent.identity] + if occurrence.status is not ApprovalStatus.CLAIMED: + raise ValueError(f"Approval occurrence is not claimed: {occurrence.status}.") + occurrence.status = ApprovalStatus.EXECUTING + + def settle(self, intent: AuthorizedExecution, results: list[Content]) -> ApprovalOutcome: + """Settle an executing occurrence with results under its original call identity.""" + occurrence = self._occurrences[intent.identity] + if occurrence.status is not ApprovalStatus.EXECUTING: + raise ValueError(f"Approval occurrence is not executing: {occurrence.status}.") + replayable_results = [ + ReplayableToolResult(content=result) + for result in results + if result.type == "function_result" and result.call_id == occurrence.identity.call_id + ] + if len(replayable_results) != 1: + raise ValueError("A settled local approval must produce exactly one result for its original call.") + occurrence.replayable_results = replayable_results + occurrence.status = ApprovalStatus.SETTLED + self._pending_by_interrupt.pop((intent.identity.thread_id, intent.identity.interrupt_id), None) + return ApprovalOutcome( + identity=occurrence.identity, + replayable_results=tuple(replayable_results), + result_group=tuple(results), + ) + + def defer(self, intent: AuthorizedExecution, results: list[Content]) -> ApprovalOutcome: + """Return an execution that yielded only follow-up requests to pending.""" + occurrence = self._occurrences[intent.identity] + if occurrence.status is not ApprovalStatus.EXECUTING: + raise ValueError(f"Approval occurrence is not executing: {occurrence.status}.") + occurrence.status = ApprovalStatus.PENDING + return ApprovalOutcome(identity=occurrence.identity, replayable_results=(), result_group=tuple(results)) + + +class LocalPendingToolTransitionOwner: + """Execute an authorized call through the process-local transition owner.""" + + def __init__(self, executor: Callable[[], Awaitable[list[Content]]]) -> None: + self._executor = executor + + async def execute( + self, + intent: AuthorizedExecution, + *, + lifecycle: ApprovalLifecycle, + ) -> ApprovalOutcome: + """Execute and settle one call after lifecycle authorization.""" + lifecycle.begin_execution(intent) + results = await self._executor() + if not any(result.type == "function_result" for result in results): + return lifecycle.defer(intent, results) + return lifecycle.settle(intent, results) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py b/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py index 00b472d9ee..e28c62fa62 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py @@ -7,6 +7,8 @@ from collections import OrderedDict from typing import Any +from ._approval_lifecycle import ApprovalLifecycle + ApprovalScope = str """Application-defined scope for server-side AG-UI Approval State.""" @@ -50,6 +52,44 @@ def __init__(self, *, max_entries: int = DEFAULT_MAX_APPROVAL_STATES) -> None: self.max_entries = max_entries self.pending_approvals: OrderedDict[tuple[str, str], Any] = OrderedDict() self.tool_approval_states: OrderedDict[str, dict[str, Any]] = OrderedDict() + self.lifecycle = ApprovalLifecycle() + + def register_local( + self, + *, + thread_ids: list[str], + name: str, + arguments: str, + request_id: str, + interrupt_id: str, + already_approved_requests: list[dict[str, Any]] | None = None, + ) -> None: + """Register one local occurrence and its trusted aliases.""" + entry: dict[str, Any] = { + "name": name, + "arguments": arguments, + "request_id": request_id, + "interrupt_id": interrupt_id, + } + if already_approved_requests: + entry["already_approved_requests"] = already_approved_requests + + for thread_id in dict.fromkeys(thread_ids): + aliases = {(thread_id, request_id), (thread_id, interrupt_id)} + replaced_entries = {id(existing) for key, existing in self.pending_approvals.items() if key in aliases} + for key, existing in list(self.pending_approvals.items()): + if key in aliases or id(existing) in replaced_entries: + self.pending_approvals.pop(key, None) + for key in aliases: + self.pending_approvals[key] = entry + self.lifecycle.register_local( + thread_id=thread_id, + interrupt_id=interrupt_id, + call_id=interrupt_id, + name=name, + arguments=arguments, + ) + self.evict_oldest() def evict_oldest(self) -> None: """Evict oldest pending approval entries until the store is within bounds.""" diff --git a/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py b/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py index 533026bca0..fd66081246 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py +++ b/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py @@ -819,21 +819,25 @@ async def stream_fn_turn2( assert len(run_started) == 1 assert len(run_finished) == 1 - # Verify that a FunctionResultContent was created and sent to the agent - tool_result_found = False - for msg in messages_received: - for content in msg.contents: - if content.type == "function_result": - tool_result_found = True - assert content.call_id == "call_get_datetime_123" - assert content.result == "2025/12/01 12:00:00" - break - - assert tool_result_found, ( - "FunctionResultContent should be included in messages sent to agent. " - "This is required for the model to see the approved tool execution result." + result_events = [event for event in events2 if event.type == "TOOL_CALL_RESULT"] + assert len(result_events) == 1 + assert result_events[0].tool_call_id == "call_get_datetime_123" + assert result_events[0].content == "2025/12/01 12:00:00" + assert not any( + event.type in {"TOOL_CALL_START", "TOOL_CALL_ARGS", "TOOL_CALL_END"} + and getattr(event, "tool_call_id", None) == "call_get_datetime_123" + for event in events2 ) + replayable_results = [ + content + for message in messages_received + for content in message.contents + if content.type == "function_result" and content.call_id == "call_get_datetime_123" + ] + assert len(replayable_results) == 1 + assert replayable_results[0].result == "2025/12/01 12:00:00" + async def test_function_approval_mode_rejection(streaming_chat_client_stub): """Test that function approval rejection creates a rejection response.""" diff --git a/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py b/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py new file mode 100644 index 0000000000..f01a70316f --- /dev/null +++ b/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py @@ -0,0 +1,54 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Behavior tests for the approval batch continuity lifecycle seam.""" + +from __future__ import annotations + +from agent_framework import Content + +from agent_framework_ag_ui._approval_lifecycle import ( + ApprovalLifecycle, + ApprovalStatus, + LocalPendingToolTransitionOwner, + ResumeDecision, +) + + +async def test_local_approval_crosses_lifecycle_before_execution_and_settlement() -> None: + """One accepted local occurrence is claimed, executed by its owner, and settled.""" + lifecycle = ApprovalLifecycle() + occurrence = lifecycle.register_local( + thread_id="thread-1", + interrupt_id="approval-1", + call_id="call-1", + name="get_weather", + arguments='{"city":"Seattle"}', + ) + observed_statuses = [occurrence.status] + + intent = lifecycle.claim( + thread_id="thread-1", + decision=ResumeDecision( + interrupt_id="approval-1", + accepted=True, + arguments='{"city":"Seattle"}', + ), + ) + observed_statuses.append(lifecycle.get(occurrence.identity).status) + + async def execute_authorized_call() -> list[Content]: + observed_statuses.append(lifecycle.get(occurrence.identity).status) + return [Content.from_function_result(call_id="call-1", result="Sunny")] + + owner = LocalPendingToolTransitionOwner(execute_authorized_call) + outcome = await owner.execute(intent, lifecycle=lifecycle) + observed_statuses.append(lifecycle.get(occurrence.identity).status) + + assert observed_statuses == [ + ApprovalStatus.PENDING, + ApprovalStatus.CLAIMED, + ApprovalStatus.EXECUTING, + ApprovalStatus.SETTLED, + ] + assert [result.content.call_id for result in outcome.replayable_results] == ["call-1"] + assert [result.content.result for result in outcome.replayable_results] == ["Sunny"] diff --git a/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py b/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py index 3e60017112..c02110760b 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py +++ b/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py @@ -12,6 +12,60 @@ from agent_framework_ag_ui._agent import AgentConfig from agent_framework_ag_ui._agent_run import PendingApprovalEntry, PendingApprovalKey, run_agent_stream +from agent_framework_ag_ui._approval_lifecycle import ResumeDecision +from agent_framework_ag_ui._approval_state import InMemoryAGUIApprovalStateStore + + +async def _run_with_registered_approval_state( + input_data: dict[str, Any], + agent: StubAgent, + config: AgentConfig, +) -> list[Any]: + """Cross the lifecycle seam with server-owned state for approval fixtures.""" + thread_id = str(input_data["thread_id"]) + calls: dict[str, dict[str, Any]] = {} + decisions: list[dict[str, Any]] = [] + messages: list[dict[str, Any]] = [] + for message in input_data["messages"]: + if message.get("role") == "assistant": + for tool_call in message.get("tool_calls", []): + calls[str(tool_call["id"])] = tool_call["function"] + if message.get("role") == "tool": + decision = json.loads(message["content"]) + if isinstance(decision, dict) and "accepted" in decision: + decisions.append( + { + "interruptId": str(message["toolCallId"]), + "status": "resolved", + "payload": decision, + } + ) + continue + messages.append(message) + + store = InMemoryAGUIApprovalStateStore() + for decision in decisions: + call_id = str(decision["interruptId"]) + function = calls[call_id] + arguments = json.dumps(json.loads(function["arguments"]), sort_keys=True, separators=(",", ":")) + store.register_local( + thread_ids=[thread_id], + name=str(function["name"]), + arguments=arguments, + request_id=call_id, + interrupt_id=call_id, + ) + + events: list[Any] = [] + async for event in run_agent_stream( + {**input_data, "messages": messages, "resume": decisions}, + agent, + config, + pending_approvals=store.pending_approvals, + approval_state_store=store, + ): + events.append(event) + return events def _make_weather_tool() -> FunctionTool: @@ -75,9 +129,7 @@ async def test_approval_resume_emits_tool_call_result() -> None: "messages": resume_messages, } - events: list[Any] = [] - async for event in run_agent_stream(input_data, agent, config): - events.append(event) + events = await _run_with_registered_approval_state(input_data, agent, config) event_types = [getattr(e, "type", None) for e in events] @@ -141,9 +193,7 @@ async def test_approval_resume_result_has_content() -> None: "messages": resume_messages, } - events: list[Any] = [] - async for event in run_agent_stream(input_data, agent, config): - events.append(event) + events = await _run_with_registered_approval_state(input_data, agent, config) tool_result_events = [e for e in events if getattr(e, "type", None) == "TOOL_CALL_RESULT"] assert len(tool_result_events) == 1 @@ -189,8 +239,7 @@ async def test_approval_resume_snapshot_replaces_approval_payload_with_tool_resu }, ] - events: list[Any] = [] - async for event in run_agent_stream( + events = await _run_with_registered_approval_state( { "thread_id": "thread-snapshot-replay", "run_id": "run-snapshot-replay", @@ -198,8 +247,7 @@ async def test_approval_resume_snapshot_replaces_approval_payload_with_tool_resu }, agent, config, - ): - events.append(event) + ) snapshots = [event.messages for event in events if getattr(event, "type", None) == "MESSAGES_SNAPSHOT"] assert snapshots @@ -289,9 +337,7 @@ async def test_rejection_does_not_emit_tool_call_result() -> None: "messages": resume_messages, } - events: list[Any] = [] - async for event in run_agent_stream(input_data, agent, config): - events.append(event) + events = await _run_with_registered_approval_state(input_data, agent, config) tool_result_events = [e for e in events if getattr(e, "type", None) == "TOOL_CALL_RESULT"] assert len(tool_result_events) == 0, ( @@ -368,9 +414,7 @@ async def test_mixed_approve_reject_emits_only_approved_tool_result() -> None: "messages": resume_messages, } - events: list[Any] = [] - async for event in run_agent_stream(input_data, agent, config): - events.append(event) + events = await _run_with_registered_approval_state(input_data, agent, config) tool_result_events = [e for e in events if getattr(e, "type", None) == "TOOL_CALL_RESULT"] @@ -423,9 +467,7 @@ async def test_approval_resume_zero_updates_emits_tool_result() -> None: "messages": resume_messages, } - events: list[Any] = [] - async for event in run_agent_stream(input_data, agent, config): - events.append(event) + events = await _run_with_registered_approval_state(input_data, agent, config) event_types = [getattr(e, "type", None) for e in events] assert "RUN_STARTED" in event_types @@ -552,8 +594,29 @@ def request_consent() -> str: Message(role="user", contents=[approval_request.to_function_approval_response(approved=True)]), ] agent = StubAgent(updates=[], default_options={"tools": [consent_tool]}) + store = InMemoryAGUIApprovalStateStore() + store.register_local( + thread_ids=["thread-consent"], + name="request_consent", + arguments="{}", + request_id="approval_consent", + interrupt_id="call_consent", + ) + intent = store.lifecycle.claim( + thread_id="thread-consent", + decision=ResumeDecision(interrupt_id="call_consent", accepted=True, arguments="{}"), + ) - results = await _resolve_approval_responses(messages, [consent_tool], agent, {}) + results = await _resolve_approval_responses( + messages, + [consent_tool], + agent, + {}, + store.pending_approvals, + "thread-consent", + lifecycle=store.lifecycle, + authorized_executions={"call_consent": intent}, + ) follow_up_requests = [content for message in messages for content in message.contents if content.user_input_request] assert results == [] @@ -612,11 +675,7 @@ async def test_resolve_approval_responses_keeps_fresh_occurrence_when_canonical_ """A completed occurrence cannot consume a later approval that reuses its canonical call id.""" from agent_framework import Message - from agent_framework_ag_ui._agent_run import ( - _make_pending_approval_entry, - _pending_approval_key, - _resolve_approval_responses, - ) + from agent_framework_ag_ui._agent_run import _resolve_approval_responses executions: list[str] = [] @@ -647,14 +706,18 @@ def guarded_write(value: str) -> str: ), ] thread_id = "thread-reused" - pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] = { - _pending_approval_key(thread_id, call_id): _make_pending_approval_entry( - "guarded_write", - '{"value":"same"}', - request_id=call_id, - interrupt_id=call_id, - ) - } + store = InMemoryAGUIApprovalStateStore() + store.register_local( + thread_ids=[thread_id], + name="guarded_write", + arguments='{"value":"same"}', + request_id=call_id, + interrupt_id=call_id, + ) + intent = store.lifecycle.claim( + thread_id=thread_id, + decision=ResumeDecision(interrupt_id=call_id, accepted=True, arguments='{"value":"same"}'), + ) agent = StubAgent(updates=[], default_options={"tools": [tool]}) results = await _resolve_approval_responses( @@ -662,13 +725,15 @@ def guarded_write(value: str) -> str: [tool], agent, {}, - pending_approvals, + store.pending_approvals, thread_id, + lifecycle=store.lifecycle, + authorized_executions={call_id: intent}, ) assert executions == ["same"] assert [result.result for result in results] == ["wrote:same"] - assert pending_approvals == {} + assert store.pending_approvals == {} assert not [ content for message in messages for content in message.contents if content.type == "function_approval_response" ] From 26ef56d71571c0da015a87c234e4aae34293d4fc Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Mon, 10 Aug 2026 13:44:25 +0900 Subject: [PATCH 02/13] Make approval batches occurrence-safe Key decisions: - Give each local approval a scoped logical occurrence identity and share one occurrence across trusted thread aliases. - Validate complete Resume Decision batches before applying claims, then account for accepted, rejected, and cancelled occurrences independently. - Preserve sibling authority and original result identity across failures, mixed decisions, and reused raw call IDs. Files changed: - packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py - packages/ag-ui/agent_framework_ag_ui/_approval_state.py - packages/ag-ui/agent_framework_ag_ui/_agent_run.py - packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py - packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py Verification: - 959 AG-UI tests passed with 90% lifecycle branch coverage. - 30 focused lifecycle/public tracer tests passed with warnings treated as errors. - Ruff format/check and AG-UI Pyright passed. - git diff --check passed. Notes for next iteration: - The function-calling-loop scenario mapping remains inaccessible under the organization content-exclusion policy. - Workspace typing fan-out remains blocked by the pre-existing missing packages/durabletask/pyproject.toml; package-local Pyright passed, while package-local MyPy retains three unrelated baseline errors. --- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 30 ++- .../_approval_lifecycle.py | 138 +++++++++-- .../agent_framework_ag_ui/_approval_state.py | 17 +- .../ag_ui/test_agent_wrapper_comprehensive.py | 9 +- .../tests/ag_ui/test_approval_lifecycle.py | 234 ++++++++++++++++++ 5 files changed, 397 insertions(+), 31 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index 3b604eacaa..512976484b 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py @@ -1344,6 +1344,22 @@ def _canonical_approval_resume_messages( ) if cancelled_ids: + if lifecycle is not None: + local_cancelled_ids = [ + interrupt_id + for interrupt_id in cancelled_ids + if (pending_entry := entries_by_interrupt_id.get(interrupt_id)) is not None + and not _pending_approval_server_label(pending_entry) + ] + try: + lifecycle.cancel_batch(thread_id=thread_id, interrupt_ids=local_cancelled_ids) + except (KeyError, ValueError) as exc: + return ( + [], + handled_ids, + cancelled_ids, + RunErrorEvent(message=str(exc), code="APPROVAL_RESUME_INVALID"), + ) for interrupt_id in cancelled_ids: pending_entry = entries_by_interrupt_id.get(interrupt_id) if pending_entry is not None: @@ -1516,10 +1532,9 @@ def _canonical_approval_resume_messages( if lifecycle is not None and authorized_executions is not None: try: - for decision in lifecycle_decisions: - if decision.accepted: - intent = lifecycle.claim(thread_id=thread_id, decision=decision) - authorized_executions[intent.identity.call_id] = intent + intents = lifecycle.claim_batch(thread_id=thread_id, decisions=lifecycle_decisions) + for intent in intents: + authorized_executions[intent.identity.call_id] = intent except (KeyError, ValueError) as exc: return ( [], @@ -1644,8 +1659,8 @@ def matches_pending_entry(candidate: PendingApprovalEntry | None, expected: Pend primary_response = responses[-1] response_content_ids_to_strip.update(id(response) for response in responses[:-1]) resp_id = primary_response.id - registry_key = _pending_approval_key(thread_id, resp_id) if resp_id is not None else None - id_entry = pending_approvals.get(registry_key) if registry_key is not None else None + primary_registry_key = _pending_approval_key(thread_id, resp_id) if resp_id is not None else None + id_entry = pending_approvals.get(primary_registry_key) if primary_registry_key is not None else None if not matches_pending_entry(id_entry, pending_entry): logger.warning( "Rejected approval response id=%s: no matching pending approval request", @@ -2324,6 +2339,7 @@ async def run_agent_stream( current_state=flow.current_state, ) + authorized_executions: dict[str, AuthorizedExecution] = {} approval_resume_messages, handled_resume_ids, cancelled_resume_ids, resume_error = ( _canonical_approval_resume_messages( resume_payload, @@ -2331,7 +2347,7 @@ async def run_agent_stream( approval_thread_id, expected_interrupt_ids=stored_pending_approval_interrupt_ids or None, lifecycle=approval_state_store.lifecycle if approval_state_store is not None else None, - authorized_executions=(authorized_executions := {}), + authorized_executions=authorized_executions, ) ) if resume_error is not None: diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py b/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py index 5e0823d355..59dabe8316 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py @@ -7,6 +7,7 @@ from collections.abc import Awaitable, Callable from dataclasses import dataclass, field from enum import Enum +from uuid import uuid4 from agent_framework import Content @@ -18,13 +19,16 @@ class ApprovalStatus(str, Enum): CLAIMED = "claimed" EXECUTING = "executing" SETTLED = "settled" + REJECTED = "rejected" + CANCELLED = "cancelled" @dataclass(frozen=True) class ApprovalOccurrenceIdentity: - """Identity of one approval occurrence within a server-owned thread.""" + """Identity of one occurrence within a scoped server-owned thread.""" thread_id: str + occurrence_id: str interrupt_id: str call_id: str @@ -44,6 +48,7 @@ class ApprovalOccurrence: """Server-owned state for one approval-gated call occurrence.""" identity: ApprovalOccurrenceIdentity + thread_ids: tuple[str, ...] name: str arguments: str status: ApprovalStatus = ApprovalStatus.PENDING @@ -92,14 +97,60 @@ def register_local( arguments: str, ) -> ApprovalOccurrence: """Register one server-generated local approval occurrence.""" + return self.register_local_aliases( + thread_ids=[thread_id], + interrupt_id=interrupt_id, + call_id=call_id, + name=name, + arguments=arguments, + ) + + def register_local_aliases( + self, + *, + thread_ids: list[str], + interrupt_id: str, + call_id: str, + name: str, + arguments: str, + ) -> ApprovalOccurrence: + """Register one occurrence under its trusted scoped-thread aliases.""" + unique_thread_ids = tuple(dict.fromkeys(thread_ids)) + if not unique_thread_ids: + raise ValueError("An approval occurrence requires at least one scoped thread identity.") + existing_identities = { + identity + for thread_id in unique_thread_ids + if (identity := self._pending_by_interrupt.get((thread_id, interrupt_id))) is not None + } + if len(existing_identities) > 1: + raise ValueError("Approval aliases resolve to different pending occurrences.") + if existing_identities: + occurrence = self._occurrences[next(iter(existing_identities))] + if occurrence.status is not ApprovalStatus.PENDING: + raise ValueError(f"Approval occurrence is not pending: {occurrence.status}.") + if occurrence.identity.call_id != call_id or occurrence.name != name or occurrence.arguments != arguments: + raise ValueError("Approval alias conflicts with an existing pending occurrence.") + occurrence.thread_ids = tuple(dict.fromkeys((*occurrence.thread_ids, *unique_thread_ids))) + for thread_id in occurrence.thread_ids: + self._pending_by_interrupt[(thread_id, interrupt_id)] = occurrence.identity + return occurrence + identity = ApprovalOccurrenceIdentity( - thread_id=thread_id, + thread_id=unique_thread_ids[0], + occurrence_id=str(uuid4()), interrupt_id=interrupt_id, call_id=call_id, ) - occurrence = ApprovalOccurrence(identity=identity, name=name, arguments=arguments) + occurrence = ApprovalOccurrence( + identity=identity, + thread_ids=unique_thread_ids, + name=name, + arguments=arguments, + ) self._occurrences[identity] = occurrence - self._pending_by_interrupt[(thread_id, interrupt_id)] = identity + for thread_id in unique_thread_ids: + self._pending_by_interrupt[(thread_id, interrupt_id)] = identity return occurrence def get(self, identity: ApprovalOccurrenceIdentity) -> ApprovalOccurrence: @@ -108,17 +159,76 @@ def get(self, identity: ApprovalOccurrenceIdentity) -> ApprovalOccurrence: def claim(self, *, thread_id: str, decision: ResumeDecision) -> AuthorizedExecution: """Validate and reserve one accepted decision before execution.""" - identity = self._pending_by_interrupt[(thread_id, decision.interrupt_id)] - occurrence = self._occurrences[identity] - if occurrence.status is not ApprovalStatus.PENDING: - raise ValueError(f"Approval occurrence is not pending: {occurrence.status}.") if not decision.accepted: raise ValueError("A rejected decision cannot authorize execution.") - if (decision.original_arguments or decision.arguments) != occurrence.arguments: - raise ValueError("Approval decision arguments do not match the registered occurrence.") - occurrence.arguments = decision.arguments - occurrence.status = ApprovalStatus.CLAIMED - return AuthorizedExecution(identity=identity, name=occurrence.name, arguments=occurrence.arguments) + return self.claim_batch(thread_id=thread_id, decisions=[decision])[0] + + def claim_batch( + self, + *, + thread_id: str, + decisions: list[ResumeDecision], + ) -> tuple[AuthorizedExecution, ...]: + """Validate a complete decision batch before reserving accepted occurrences.""" + resolved: list[tuple[ResumeDecision, ApprovalOccurrence]] = [] + seen_interrupt_ids: set[str] = set() + for decision in decisions: + if decision.interrupt_id in seen_interrupt_ids: + raise ValueError(f"Approval batch repeats interrupt: {decision.interrupt_id}.") + seen_interrupt_ids.add(decision.interrupt_id) + identity = self._pending_by_interrupt[(thread_id, decision.interrupt_id)] + occurrence = self._occurrences[identity] + if occurrence.status is not ApprovalStatus.PENDING: + raise ValueError(f"Approval occurrence is not pending: {occurrence.status}.") + if (decision.original_arguments or decision.arguments) != occurrence.arguments: + raise ValueError("Approval decision arguments do not match the registered occurrence.") + resolved.append((decision, occurrence)) + intents: list[AuthorizedExecution] = [] + for decision, occurrence in resolved: + if not decision.accepted: + occurrence.replayable_results = [ + ReplayableToolResult( + content=Content.from_function_result( + call_id=occurrence.identity.call_id, + result="Error: Tool call invocation was rejected by user.", + ) + ) + ] + occurrence.status = ApprovalStatus.REJECTED + self._remove_pending_aliases(occurrence) + continue + occurrence.arguments = decision.arguments + occurrence.status = ApprovalStatus.CLAIMED + intents.append( + AuthorizedExecution( + identity=occurrence.identity, + name=occurrence.name, + arguments=occurrence.arguments, + ) + ) + return tuple(intents) + + def cancel_batch(self, *, thread_id: str, interrupt_ids: list[str]) -> None: + """Validate and cancel selected occurrences without changing their siblings.""" + occurrences: list[ApprovalOccurrence] = [] + seen_interrupt_ids: set[str] = set() + for interrupt_id in interrupt_ids: + if interrupt_id in seen_interrupt_ids: + raise ValueError(f"Approval batch repeats interrupt: {interrupt_id}.") + seen_interrupt_ids.add(interrupt_id) + identity = self._pending_by_interrupt[(thread_id, interrupt_id)] + occurrence = self._occurrences[identity] + if occurrence.status is not ApprovalStatus.PENDING: + raise ValueError(f"Approval occurrence is not pending: {occurrence.status}.") + occurrences.append(occurrence) + + for occurrence in occurrences: + occurrence.status = ApprovalStatus.CANCELLED + self._remove_pending_aliases(occurrence) + + def _remove_pending_aliases(self, occurrence: ApprovalOccurrence) -> None: + for thread_id in occurrence.thread_ids: + self._pending_by_interrupt.pop((thread_id, occurrence.identity.interrupt_id), None) def begin_execution(self, intent: AuthorizedExecution) -> None: """Mark a claimed occurrence immediately before its owner may invoke a tool.""" @@ -141,7 +251,7 @@ def settle(self, intent: AuthorizedExecution, results: list[Content]) -> Approva raise ValueError("A settled local approval must produce exactly one result for its original call.") occurrence.replayable_results = replayable_results occurrence.status = ApprovalStatus.SETTLED - self._pending_by_interrupt.pop((intent.identity.thread_id, intent.identity.interrupt_id), None) + self._remove_pending_aliases(occurrence) return ApprovalOutcome( identity=occurrence.identity, replayable_results=tuple(replayable_results), diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py b/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py index e28c62fa62..540f90f8ed 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py @@ -74,7 +74,15 @@ def register_local( if already_approved_requests: entry["already_approved_requests"] = already_approved_requests - for thread_id in dict.fromkeys(thread_ids): + unique_thread_ids = list(dict.fromkeys(thread_ids)) + self.lifecycle.register_local_aliases( + thread_ids=unique_thread_ids, + interrupt_id=interrupt_id, + call_id=interrupt_id, + name=name, + arguments=arguments, + ) + for thread_id in unique_thread_ids: aliases = {(thread_id, request_id), (thread_id, interrupt_id)} replaced_entries = {id(existing) for key, existing in self.pending_approvals.items() if key in aliases} for key, existing in list(self.pending_approvals.items()): @@ -82,13 +90,6 @@ def register_local( self.pending_approvals.pop(key, None) for key in aliases: self.pending_approvals[key] = entry - self.lifecycle.register_local( - thread_id=thread_id, - interrupt_id=interrupt_id, - call_id=interrupt_id, - name=name, - arguments=arguments, - ) self.evict_oldest() def evict_oldest(self) -> None: diff --git a/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py b/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py index fd66081246..7b009aee47 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py +++ b/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py @@ -872,8 +872,13 @@ async def stream_fn( thread_id = "thread-rejection-test" - # Pre-populate the pending approval as if Turn 1 had emitted the request. - wrapper._pending_approvals[(thread_id, "call_delete_123")] = "delete_all_data" + wrapper._approval_state_store.register_local( + thread_ids=[thread_id], + name="delete_all_data", + arguments="{}", + request_id="call_delete_123", + interrupt_id="call_delete_123", + ) input_data: dict[str, Any] = { "thread_id": thread_id, diff --git a/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py b/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py index f01a70316f..2f8b4eebaf 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py +++ b/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py @@ -4,6 +4,7 @@ from __future__ import annotations +import pytest from agent_framework import Content from agent_framework_ag_ui._approval_lifecycle import ( @@ -52,3 +53,236 @@ async def execute_authorized_call() -> list[Content]: ] assert [result.content.call_id for result in outcome.replayable_results] == ["call-1"] assert [result.content.result for result in outcome.replayable_results] == ["Sunny"] + + +async def test_execution_failure_keeps_unexecuted_batch_sibling_claimed() -> None: + """A failed occurrence does not erase a claimed sibling that can still execute.""" + lifecycle = ApprovalLifecycle() + first = lifecycle.register_local( + thread_id="thread-1", + interrupt_id="approval-1", + call_id="call-1", + name="write_record", + arguments='{"value":"first"}', + ) + second = lifecycle.register_local( + thread_id="thread-1", + interrupt_id="approval-2", + call_id="call-2", + name="write_record", + arguments='{"value":"second"}', + ) + first_intent, second_intent = lifecycle.claim_batch( + thread_id="thread-1", + decisions=[ + ResumeDecision(interrupt_id="approval-1", accepted=True, arguments='{"value":"first"}'), + ResumeDecision(interrupt_id="approval-2", accepted=True, arguments='{"value":"second"}'), + ], + ) + + async def fail_first() -> list[Content]: + raise RuntimeError("side effect failed") + + with pytest.raises(RuntimeError, match="side effect failed"): + await LocalPendingToolTransitionOwner(fail_first).execute(first_intent, lifecycle=lifecycle) + + assert lifecycle.get(first.identity).status is ApprovalStatus.EXECUTING + assert lifecycle.get(second.identity).status is ApprovalStatus.CLAIMED + + async def execute_second() -> list[Content]: + return [Content.from_function_result(call_id="call-2", result="wrote second")] + + await LocalPendingToolTransitionOwner(execute_second).execute(second_intent, lifecycle=lifecycle) + assert lifecycle.get(second.identity).status is ApprovalStatus.SETTLED + + +def test_batch_validation_is_atomic_before_claiming_any_occurrence() -> None: + """One invalid decision leaves every occurrence pending and eligible for a corrected batch.""" + lifecycle = ApprovalLifecycle() + first = lifecycle.register_local( + thread_id="tenant-a\x1fthread-1", + interrupt_id="approval-1", + call_id="call-1", + name="write_record", + arguments='{"value":"first"}', + ) + second = lifecycle.register_local( + thread_id="tenant-a\x1fthread-1", + interrupt_id="approval-2", + call_id="call-2", + name="write_record", + arguments='{"value":"second"}', + ) + + with pytest.raises(ValueError, match="arguments do not match"): + lifecycle.claim_batch( + thread_id="tenant-a\x1fthread-1", + decisions=[ + ResumeDecision(interrupt_id="approval-1", accepted=True, arguments='{"value":"first"}'), + ResumeDecision(interrupt_id="approval-2", accepted=True, arguments='{"value":"forged"}'), + ], + ) + + assert lifecycle.get(first.identity).status is ApprovalStatus.PENDING + assert lifecycle.get(second.identity).status is ApprovalStatus.PENDING + + +def test_mixed_batch_accounts_for_rejection_under_original_call_identity() -> None: + """A rejected occurrence remains represented while its accepted sibling is claimed.""" + lifecycle = ApprovalLifecycle() + accepted = lifecycle.register_local( + thread_id="thread-1", + interrupt_id="approval-1", + call_id="call-1", + name="write_record", + arguments='{"value":"first"}', + ) + rejected = lifecycle.register_local( + thread_id="thread-1", + interrupt_id="approval-2", + call_id="call-2", + name="write_record", + arguments='{"value":"second"}', + ) + + intents = lifecycle.claim_batch( + thread_id="thread-1", + decisions=[ + ResumeDecision(interrupt_id="approval-1", accepted=True, arguments='{"value":"first"}'), + ResumeDecision(interrupt_id="approval-2", accepted=False, arguments='{"value":"second"}'), + ], + ) + + assert [intent.identity for intent in intents] == [accepted.identity] + assert lifecycle.get(rejected.identity).status is ApprovalStatus.REJECTED + assert [result.content.call_id for result in lifecycle.get(rejected.identity).replayable_results] == ["call-2"] + + +def test_batch_claims_preserve_order_and_scope_reused_raw_call_ids() -> None: + """Raw call ids reused in another scoped thread cannot correlate approval authority.""" + lifecycle = ApprovalLifecycle() + tenant_a = lifecycle.register_local( + thread_id="tenant-a\x1fthread-1", + interrupt_id="approval-shared", + call_id="call-shared", + name="write_record", + arguments='{"tenant":"a"}', + ) + tenant_b = lifecycle.register_local( + thread_id="tenant-b\x1fthread-1", + interrupt_id="approval-shared", + call_id="call-shared", + name="write_record", + arguments='{"tenant":"b"}', + ) + + intents = lifecycle.claim_batch( + thread_id="tenant-a\x1fthread-1", + decisions=[ + ResumeDecision( + interrupt_id="approval-shared", + accepted=True, + arguments='{"tenant":"a"}', + ) + ], + ) + + assert [intent.identity for intent in intents] == [tenant_a.identity] + assert tenant_a.identity != tenant_b.identity + assert lifecycle.get(tenant_a.identity).status is ApprovalStatus.CLAIMED + assert lifecycle.get(tenant_b.identity).status is ApprovalStatus.PENDING + + +def test_batch_cancellation_preserves_each_original_occurrence() -> None: + """Cancelling selected occurrences is terminal without consuming an unrelated sibling.""" + lifecycle = ApprovalLifecycle() + cancelled = lifecycle.register_local( + thread_id="tenant-a\x1fthread-1", + interrupt_id="approval-1", + call_id="call-1", + name="write_record", + arguments='{"value":"first"}', + ) + pending = lifecycle.register_local( + thread_id="tenant-a\x1fthread-1", + interrupt_id="approval-2", + call_id="call-2", + name="write_record", + arguments='{"value":"second"}', + ) + + lifecycle.cancel_batch( + thread_id="tenant-a\x1fthread-1", + interrupt_ids=["approval-1"], + ) + + assert lifecycle.get(cancelled.identity).status is ApprovalStatus.CANCELLED + assert lifecycle.get(pending.identity).status is ApprovalStatus.PENDING + + +def test_one_occurrence_can_be_claimed_through_a_trusted_thread_alias() -> None: + """Provider conversation aliases address one occurrence rather than duplicating authority.""" + lifecycle = ApprovalLifecycle() + occurrence = lifecycle.register_local_aliases( + thread_ids=["tenant-a\x1fag-ui-thread", "tenant-a\x1fprovider-thread"], + interrupt_id="approval-1", + call_id="call-1", + name="write_record", + arguments='{"value":"first"}', + ) + + intents = lifecycle.claim_batch( + thread_id="tenant-a\x1fprovider-thread", + decisions=[ResumeDecision(interrupt_id="approval-1", accepted=True, arguments='{"value":"first"}')], + ) + + assert [intent.identity for intent in intents] == [occurrence.identity] + with pytest.raises(ValueError, match="not pending"): + lifecycle.register_local_aliases( + thread_ids=["tenant-a\x1fag-ui-thread", "tenant-a\x1fprovider-thread"], + interrupt_id="approval-1", + call_id="call-1", + name="write_record", + arguments='{"value":"first"}', + ) + with pytest.raises(ValueError, match="not pending"): + lifecycle.claim_batch( + thread_id="tenant-a\x1fag-ui-thread", + decisions=[ResumeDecision(interrupt_id="approval-1", accepted=True, arguments='{"value":"first"}')], + ) + + +async def test_settled_raw_call_id_can_be_reused_for_a_new_occurrence() -> None: + """Sequential reuse creates a fresh logical occurrence instead of reviving settled authority.""" + lifecycle = ApprovalLifecycle() + first = lifecycle.register_local( + thread_id="thread-1", + interrupt_id="approval-shared", + call_id="call-shared", + name="write_record", + arguments='{"value":"first"}', + ) + first_intent = lifecycle.claim( + thread_id="thread-1", + decision=ResumeDecision( + interrupt_id="approval-shared", + accepted=True, + arguments='{"value":"first"}', + ), + ) + + async def execute_first() -> list[Content]: + return [Content.from_function_result(call_id="call-shared", result="wrote first")] + + await LocalPendingToolTransitionOwner(execute_first).execute(first_intent, lifecycle=lifecycle) + second = lifecycle.register_local( + thread_id="thread-1", + interrupt_id="approval-shared", + call_id="call-shared", + name="write_record", + arguments='{"value":"second"}', + ) + + assert first.identity != second.identity + assert lifecycle.get(first.identity).status is ApprovalStatus.SETTLED + assert lifecycle.get(second.identity).status is ApprovalStatus.PENDING From f24f57a45253a21a8832d10b838a2d533d4c9509 Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Mon, 10 Aug 2026 13:55:41 +0900 Subject: [PATCH 03/13] Make approval resume retries idempotent Key decisions: - Retain terminal decisions and outcomes by scoped occurrence so identical accepted and rejected retries reproject results without granting execution authority again. - Reject conflicting names, arguments, decisions, wrong-scope lookups, and expired authority before an execution intent can reach the local transition owner. - Keep protocol normalization in the runner while using server-owned lifecycle context to canonicalize retries and preserve existing AG-UI wire aliases. Files changed: - packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py - packages/ag-ui/agent_framework_ag_ui/_agent_run.py - packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py - packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py - packages/ag-ui/tests/ag_ui/test_endpoint.py Verification: - 965 AG-UI tests passed with 90% approval lifecycle coverage. - 18 focused lifecycle, hostile-resume, wrong-thread, and endpoint retry tests passed with runtime and deprecation warnings treated as errors. - Ruff format/check and AG-UI package-local Pyright passed. - git diff --check passed. Notes for next iteration: - Terminal retention is process-local and unbounded until the later bounded-retention issue adds its explicit policy. - The function-calling-loop scenario mapping remains inaccessible under the organization content-exclusion policy. --- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 71 ++++++++- .../_approval_lifecycle.py | 137 ++++++++++++++--- .../ag_ui/test_agent_wrapper_comprehensive.py | 22 ++- .../tests/ag_ui/test_approval_lifecycle.py | 144 ++++++++++++++++++ .../ag-ui/tests/ag_ui/test_endpoint.py | 13 +- 5 files changed, 358 insertions(+), 29 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index 512976484b..cd01bf0802 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py @@ -965,6 +965,7 @@ def _pop_collected_tool_approval_response_messages( interrupt_id=str(response.id or function_call.call_id), accepted=True, arguments=arguments, + name=function_call.name, ), ) authorized_executions[intent.identity.call_id] = intent @@ -1233,6 +1234,7 @@ def _canonical_approval_resume_messages( *, lifecycle: ApprovalLifecycle | None = None, authorized_executions: dict[str, AuthorizedExecution] | None = None, + retained_results: list[Content] | None = None, ) -> tuple[list[dict[str, Any]], set[str], set[str], RunErrorEvent | None]: """Translate canonical ResumeEntry approvals into existing approval response messages.""" expected_ids = set(expected_interrupt_ids or set()) @@ -1281,6 +1283,63 @@ def _canonical_approval_resume_messages( if _resume_payload_has_approval_decision(resume_payload): normalized_interrupts = _normalize_resume_interrupts(resume_payload) interrupt_id = normalized_interrupts[0]["id"] if normalized_interrupts else "unknown" + if lifecycle is not None and retained_results is not None: + decisions: list[ResumeDecision] = [] + for interrupt in normalized_interrupts: + if interrupt.get("status") != "resolved": + break + payload = _parse_json_object(interrupt.get("value")) + if payload is None: + break + accepted = payload.get("accepted", payload.get("approved")) + if not isinstance(accepted, bool): + break + interrupt_id = str(interrupt["id"]) + try: + name, retained_arguments = lifecycle.decision_context( + thread_id=thread_id, + interrupt_id=interrupt_id, + ) + except KeyError: + break + edited_arguments = { + key: value for key, value in payload.items() if key not in {"accepted", "approved"} + } + canonical_arguments: str | None = None + if edited_arguments: + retained_argument_values = _parse_json_object(retained_arguments) + if retained_argument_values is None: + break + canonical_arguments = json.dumps( + make_json_safe({**retained_argument_values, **edited_arguments}), + sort_keys=True, + separators=(",", ":"), + ) + decisions.append( + ResumeDecision( + interrupt_id=interrupt_id, + accepted=accepted, + arguments=canonical_arguments, + name=name, + ) + ) + if len(decisions) == len(normalized_interrupts) and decisions: + try: + batch = lifecycle.claim_batch(thread_id=thread_id, decisions=decisions) + except KeyError: + pass + except ValueError as exc: + return ( + [], + handled_ids, + cancelled_ids, + RunErrorEvent(message=str(exc), code="APPROVAL_RESUME_INVALID"), + ) + else: + for outcome in batch.retained_outcomes: + retained_results.extend(result.content for result in outcome.replayable_results) + handled_ids.update(decision.interrupt_id for decision in decisions) + return [], handled_ids, cancelled_ids, None return ( [], handled_ids, @@ -1458,6 +1517,7 @@ def _canonical_approval_resume_messages( interrupt_id=interrupt_id, accepted=accepted, arguments=canonical_arguments, + name=_pending_approval_name(pending_entry), original_arguments=pending_arguments, ) ) @@ -1514,6 +1574,7 @@ def _canonical_approval_resume_messages( interrupt_id=sibling_interrupt_id, accepted=True, arguments=sibling_arguments, + name=function_call.name, ) ) function_approvals.append( @@ -2340,6 +2401,7 @@ async def run_agent_stream( ) authorized_executions: dict[str, AuthorizedExecution] = {} + retained_approval_results: list[Content] = [] approval_resume_messages, handled_resume_ids, cancelled_resume_ids, resume_error = ( _canonical_approval_resume_messages( resume_payload, @@ -2348,6 +2410,7 @@ async def run_agent_stream( expected_interrupt_ids=stored_pending_approval_interrupt_ids or None, lifecycle=approval_state_store.lifecycle if approval_state_store is not None else None, authorized_executions=authorized_executions, + retained_results=retained_approval_results, ) ) if resume_error is not None: @@ -2374,6 +2437,12 @@ async def run_agent_stream( if resume_messages: logger.info(f"Appending {len(resume_messages)} synthesized resume message(s) to AG-UI input.") raw_messages.extend(resume_messages) + if retained_approval_results and not raw_messages: + yield RunStartedEvent(run_id=run_id, thread_id=thread_id) + for event in _make_approval_tool_result_events(retained_approval_results): + yield event + yield _build_run_finished_event(run_id=run_id, thread_id=thread_id) + return protected_tool_call_ids = _approval_state_tool_call_ids(pending_approvals, approval_state_store, approval_thread_id) messages, snapshot_messages = normalize_agui_input_messages( raw_messages, @@ -2458,7 +2527,7 @@ async def run_agent_stream( ) ) validated_approved_responses: list[Content] = [] - resolved_approval_results = await _resolve_approval_responses( + resolved_approval_results = retained_approval_results + await _resolve_approval_responses( messages, tools_for_execution, agent, diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py b/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py index 59dabe8316..7e13d806cc 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py @@ -4,7 +4,7 @@ from __future__ import annotations -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Iterator from dataclasses import dataclass, field from enum import Enum from uuid import uuid4 @@ -21,6 +21,7 @@ class ApprovalStatus(str, Enum): SETTLED = "settled" REJECTED = "rejected" CANCELLED = "cancelled" + EXPIRED = "expired" @dataclass(frozen=True) @@ -39,7 +40,8 @@ class ResumeDecision: interrupt_id: str accepted: bool - arguments: str + arguments: str | None + name: str | None = None original_arguments: str | None = None @@ -53,6 +55,8 @@ class ApprovalOccurrence: arguments: str status: ApprovalStatus = ApprovalStatus.PENDING replayable_results: list[ReplayableToolResult] = field(default_factory=list) + decision: ResumeDecision | None = None + outcome: ApprovalOutcome | None = None @dataclass(frozen=True) @@ -80,12 +84,25 @@ class ApprovalOutcome: result_group: tuple[Content, ...] +@dataclass(frozen=True) +class ApprovalBatchDecision: + """Validated batch result containing new authority and retained outcomes.""" + + authorized_executions: tuple[AuthorizedExecution, ...] + retained_outcomes: tuple[ApprovalOutcome, ...] = () + + def __iter__(self) -> Iterator[AuthorizedExecution]: + """Iterate newly authorized executions for compatibility with existing callers.""" + return iter(self.authorized_executions) + + class ApprovalLifecycle: """Own registration, authority transitions, and settlement for approvals.""" def __init__(self) -> None: self._occurrences: dict[ApprovalOccurrenceIdentity, ApprovalOccurrence] = {} self._pending_by_interrupt: dict[tuple[str, str], ApprovalOccurrenceIdentity] = {} + self._terminal_by_interrupt: dict[tuple[str, str], ApprovalOccurrenceIdentity] = {} def register_local( self, @@ -157,46 +174,95 @@ def get(self, identity: ApprovalOccurrenceIdentity) -> ApprovalOccurrence: """Return server-owned state for a registered occurrence.""" return self._occurrences[identity] + def decision_context(self, *, thread_id: str, interrupt_id: str) -> tuple[str, str]: + """Return canonical server-owned call data needed to normalize a typed retry.""" + key = (thread_id, interrupt_id) + identity = self._pending_by_interrupt.get(key) or self._terminal_by_interrupt[key] + occurrence = self._occurrences[identity] + return occurrence.name, occurrence.arguments + def claim(self, *, thread_id: str, decision: ResumeDecision) -> AuthorizedExecution: """Validate and reserve one accepted decision before execution.""" if not decision.accepted: raise ValueError("A rejected decision cannot authorize execution.") - return self.claim_batch(thread_id=thread_id, decisions=[decision])[0] + batch = self.claim_batch(thread_id=thread_id, decisions=[decision]) + if not batch.authorized_executions: + raise ValueError("A duplicate settled decision cannot authorize execution again.") + return batch.authorized_executions[0] def claim_batch( self, *, thread_id: str, decisions: list[ResumeDecision], - ) -> tuple[AuthorizedExecution, ...]: + ) -> ApprovalBatchDecision: """Validate a complete decision batch before reserving accepted occurrences.""" - resolved: list[tuple[ResumeDecision, ApprovalOccurrence]] = [] + resolved: list[tuple[ResumeDecision, ApprovalOccurrence, bool]] = [] seen_interrupt_ids: set[str] = set() for decision in decisions: if decision.interrupt_id in seen_interrupt_ids: raise ValueError(f"Approval batch repeats interrupt: {decision.interrupt_id}.") seen_interrupt_ids.add(decision.interrupt_id) - identity = self._pending_by_interrupt[(thread_id, decision.interrupt_id)] + key = (thread_id, decision.interrupt_id) + identity = self._pending_by_interrupt.get(key) + is_terminal = identity is None + if identity is None: + identity = self._terminal_by_interrupt[key] occurrence = self._occurrences[identity] + if is_terminal: + if occurrence.status is ApprovalStatus.EXPIRED: + raise ValueError("Approval authority has expired.") + retained_decision = occurrence.decision + if ( + retained_decision is None + or retained_decision.accepted != decision.accepted + or (decision.name is not None and retained_decision.name != decision.name) + or (decision.arguments is not None and retained_decision.arguments != decision.arguments) + or ( + decision.original_arguments is not None + and retained_decision.original_arguments != decision.original_arguments + ) + ): + raise ValueError("Approval decision conflicts with the retained terminal decision.") + if occurrence.outcome is None: + raise ValueError(f"Approval occurrence has no replayable terminal outcome: {occurrence.status}.") + resolved.append((decision, occurrence, True)) + continue if occurrence.status is not ApprovalStatus.PENDING: raise ValueError(f"Approval occurrence is not pending: {occurrence.status}.") - if (decision.original_arguments or decision.arguments) != occurrence.arguments: + if decision.name is not None and decision.name != occurrence.name: + raise ValueError("Approval decision tool name does not match the registered occurrence.") + canonical_arguments = decision.arguments + if canonical_arguments is None: + raise ValueError("A pending approval decision must include canonical arguments.") + if (decision.original_arguments or canonical_arguments) != occurrence.arguments: raise ValueError("Approval decision arguments do not match the registered occurrence.") - resolved.append((decision, occurrence)) + resolved.append((decision, occurrence, False)) intents: list[AuthorizedExecution] = [] - for decision, occurrence in resolved: + retained_outcomes: list[ApprovalOutcome] = [] + for decision, occurrence, is_terminal in resolved: + if is_terminal: + if occurrence.outcome is None: + raise RuntimeError("Validated terminal approval is missing its retained outcome.") + retained_outcomes.append(occurrence.outcome) + continue + occurrence.decision = decision if not decision.accepted: - occurrence.replayable_results = [ - ReplayableToolResult( - content=Content.from_function_result( - call_id=occurrence.identity.call_id, - result="Error: Tool call invocation was rejected by user.", - ) - ) - ] + result = Content.from_function_result( + call_id=occurrence.identity.call_id, + result="Error: Tool call invocation was rejected by user.", + ) + occurrence.replayable_results = [ReplayableToolResult(content=result)] + occurrence.outcome = ApprovalOutcome( + identity=occurrence.identity, + replayable_results=tuple(occurrence.replayable_results), + result_group=(result,), + ) occurrence.status = ApprovalStatus.REJECTED self._remove_pending_aliases(occurrence) continue + if decision.arguments is None: + raise RuntimeError("Validated pending approval is missing canonical arguments.") occurrence.arguments = decision.arguments occurrence.status = ApprovalStatus.CLAIMED intents.append( @@ -206,7 +272,10 @@ def claim_batch( arguments=occurrence.arguments, ) ) - return tuple(intents) + return ApprovalBatchDecision( + authorized_executions=tuple(intents), + retained_outcomes=tuple(retained_outcomes), + ) def cancel_batch(self, *, thread_id: str, interrupt_ids: list[str]) -> None: """Validate and cancel selected occurrences without changing their siblings.""" @@ -216,7 +285,14 @@ def cancel_batch(self, *, thread_id: str, interrupt_ids: list[str]) -> None: if interrupt_id in seen_interrupt_ids: raise ValueError(f"Approval batch repeats interrupt: {interrupt_id}.") seen_interrupt_ids.add(interrupt_id) - identity = self._pending_by_interrupt[(thread_id, interrupt_id)] + key = (thread_id, interrupt_id) + identity = self._pending_by_interrupt.get(key) + if identity is None: + identity = self._terminal_by_interrupt[key] + terminal = self._occurrences[identity] + if terminal.status is ApprovalStatus.CANCELLED: + continue + raise ValueError("Approval cancellation conflicts with the retained terminal decision.") occurrence = self._occurrences[identity] if occurrence.status is not ApprovalStatus.PENDING: raise ValueError(f"Approval occurrence is not pending: {occurrence.status}.") @@ -226,9 +302,28 @@ def cancel_batch(self, *, thread_id: str, interrupt_ids: list[str]) -> None: occurrence.status = ApprovalStatus.CANCELLED self._remove_pending_aliases(occurrence) + def expire_batch(self, *, thread_id: str, interrupt_ids: list[str]) -> None: + """Expire pending authority without permitting later execution.""" + occurrences: list[ApprovalOccurrence] = [] + seen_interrupt_ids: set[str] = set() + for interrupt_id in interrupt_ids: + if interrupt_id in seen_interrupt_ids: + raise ValueError(f"Approval batch repeats interrupt: {interrupt_id}.") + seen_interrupt_ids.add(interrupt_id) + identity = self._pending_by_interrupt[(thread_id, interrupt_id)] + occurrence = self._occurrences[identity] + if occurrence.status is not ApprovalStatus.PENDING: + raise ValueError(f"Approval occurrence is not pending: {occurrence.status}.") + occurrences.append(occurrence) + + for occurrence in occurrences: + occurrence.status = ApprovalStatus.EXPIRED + self._remove_pending_aliases(occurrence) + def _remove_pending_aliases(self, occurrence: ApprovalOccurrence) -> None: for thread_id in occurrence.thread_ids: self._pending_by_interrupt.pop((thread_id, occurrence.identity.interrupt_id), None) + self._terminal_by_interrupt[(thread_id, occurrence.identity.interrupt_id)] = occurrence.identity def begin_execution(self, intent: AuthorizedExecution) -> None: """Mark a claimed occurrence immediately before its owner may invoke a tool.""" @@ -252,11 +347,13 @@ def settle(self, intent: AuthorizedExecution, results: list[Content]) -> Approva occurrence.replayable_results = replayable_results occurrence.status = ApprovalStatus.SETTLED self._remove_pending_aliases(occurrence) - return ApprovalOutcome( + outcome = ApprovalOutcome( identity=occurrence.identity, replayable_results=tuple(replayable_results), result_group=tuple(results), ) + occurrence.outcome = outcome + return outcome def defer(self, intent: AuthorizedExecution, results: list[Content]) -> ApprovalOutcome: """Return an execution that yielded only follow-up requests to pending.""" diff --git a/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py b/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py index 7b009aee47..79899e1998 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py +++ b/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py @@ -1179,10 +1179,28 @@ def approval_input(thread_id: str) -> dict[str, Any]: assert ("provider-conversation", "call_sensitive") not in wrapper._pending_approvals replay_thread_id = "provider-conversation" if resume_thread_id == "client-thread" else "client-thread" - async for _ in wrapper.run(approval_input(replay_thread_id)): - pass + retry_events = [event async for event in wrapper.run(approval_input(replay_thread_id))] + + assert execution_count == 1 + retry_results = [event for event in retry_events if event.type == "TOOL_CALL_RESULT"] + assert len(retry_results) == 1 + assert retry_results[0].tool_call_id == "call_sensitive" + assert retry_results[0].content == "executed" + assert not any(event.type == "RUN_ERROR" for event in retry_events) + + conflicting_input = approval_input(replay_thread_id) + conflicting_input["resume"][0]["payload"]["accepted"] = False + conflicting_events = [event async for event in wrapper.run(conflicting_input)] + + assert execution_count == 1 + assert any(event.type == "RUN_ERROR" and event.code == "APPROVAL_RESUME_INVALID" for event in conflicting_events) + + changed_input = approval_input(replay_thread_id) + changed_input["resume"][0]["payload"]["forged"] = True + changed_events = [event async for event in wrapper.run(changed_input)] assert execution_count == 1 + assert any(event.type == "RUN_ERROR" and event.code == "APPROVAL_RESUME_INVALID" for event in changed_events) async def test_approval_function_name_mismatch_is_blocked(streaming_chat_client_stub): diff --git a/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py b/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py index 2f8b4eebaf..827e4376bb 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py +++ b/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py @@ -286,3 +286,147 @@ async def execute_first() -> list[Content]: assert first.identity != second.identity assert lifecycle.get(first.identity).status is ApprovalStatus.SETTLED assert lifecycle.get(second.identity).status is ApprovalStatus.PENDING + + +async def test_identical_accepted_retry_returns_retained_outcome_without_execution() -> None: + """A settled accepted decision reprojects its result instead of granting authority again.""" + lifecycle = ApprovalLifecycle() + occurrence = lifecycle.register_local( + thread_id="thread-1", + interrupt_id="approval-1", + call_id="call-1", + name="write_record", + arguments='{"value":"first"}', + ) + decision = ResumeDecision( + interrupt_id="approval-1", + accepted=True, + arguments='{"value":"first"}', + ) + intent = lifecycle.claim(thread_id="thread-1", decision=decision) + invocation_count = 0 + + async def execute_once() -> list[Content]: + nonlocal invocation_count + invocation_count += 1 + return [Content.from_function_result(call_id="call-1", result="wrote first")] + + first_outcome = await LocalPendingToolTransitionOwner(execute_once).execute(intent, lifecycle=lifecycle) + retry = lifecycle.claim_batch(thread_id="thread-1", decisions=[decision]) + + assert retry.authorized_executions == () + assert retry.retained_outcomes == (first_outcome,) + assert lifecycle.get(occurrence.identity).status is ApprovalStatus.SETTLED + assert invocation_count == 1 + + +def test_accepted_retry_after_rejection_fails_as_a_conflict() -> None: + """A terminal rejection cannot be changed into execution authority by a retry.""" + lifecycle = ApprovalLifecycle() + occurrence = lifecycle.register_local( + thread_id="thread-1", + interrupt_id="approval-1", + call_id="call-1", + name="write_record", + arguments='{"value":"first"}', + ) + lifecycle.claim_batch( + thread_id="thread-1", + decisions=[ResumeDecision(interrupt_id="approval-1", accepted=False, arguments='{"value":"first"}')], + ) + + with pytest.raises(ValueError, match="conflicts with the retained terminal decision"): + lifecycle.claim_batch( + thread_id="thread-1", + decisions=[ResumeDecision(interrupt_id="approval-1", accepted=True, arguments='{"value":"first"}')], + ) + + assert lifecycle.get(occurrence.identity).status is ApprovalStatus.REJECTED + + +def test_identical_rejection_retry_returns_retained_outcome() -> None: + """A repeated rejection preserves and returns the original rejection result.""" + lifecycle = ApprovalLifecycle() + occurrence = lifecycle.register_local( + thread_id="thread-1", + interrupt_id="approval-1", + call_id="call-1", + name="write_record", + arguments='{"value":"first"}', + ) + decision = ResumeDecision(interrupt_id="approval-1", accepted=False, arguments='{"value":"first"}') + + first = lifecycle.claim_batch(thread_id="thread-1", decisions=[decision]) + retry = lifecycle.claim_batch(thread_id="thread-1", decisions=[decision]) + + assert first.authorized_executions == () + assert first.retained_outcomes == () + assert retry.authorized_executions == () + assert retry.retained_outcomes == (lifecycle.get(occurrence.identity).outcome,) + assert lifecycle.get(occurrence.identity).status is ApprovalStatus.REJECTED + + +def test_changed_tool_name_fails_before_authority_is_claimed() -> None: + """A typed decision for another tool cannot claim the registered occurrence.""" + lifecycle = ApprovalLifecycle() + occurrence = lifecycle.register_local( + thread_id="thread-1", + interrupt_id="approval-1", + call_id="call-1", + name="safe_action", + arguments="{}", + ) + + with pytest.raises(ValueError, match="tool name does not match"): + lifecycle.claim_batch( + thread_id="thread-1", + decisions=[ + ResumeDecision( + interrupt_id="approval-1", + accepted=True, + name="dangerous_action", + arguments="{}", + ) + ], + ) + + assert lifecycle.get(occurrence.identity).status is ApprovalStatus.PENDING + + +def test_identical_cancel_retry_keeps_terminal_cancellation() -> None: + """Retrying an explicit cancellation is idempotent and cannot restore authority.""" + lifecycle = ApprovalLifecycle() + occurrence = lifecycle.register_local( + thread_id="thread-1", + interrupt_id="approval-1", + call_id="call-1", + name="write_record", + arguments="{}", + ) + + lifecycle.cancel_batch(thread_id="thread-1", interrupt_ids=["approval-1"]) + lifecycle.cancel_batch(thread_id="thread-1", interrupt_ids=["approval-1"]) + + assert lifecycle.get(occurrence.identity).status is ApprovalStatus.CANCELLED + + +def test_expired_authority_cannot_be_claimed() -> None: + """Expiration is terminal and an otherwise valid decision cannot revive it.""" + lifecycle = ApprovalLifecycle() + occurrence = lifecycle.register_local( + thread_id="thread-1", + interrupt_id="approval-1", + call_id="call-1", + name="write_record", + arguments="{}", + ) + + lifecycle.expire_batch(thread_id="thread-1", interrupt_ids=["approval-1"]) + + with pytest.raises(ValueError, match="expired"): + lifecycle.claim_batch( + thread_id="thread-1", + decisions=[ResumeDecision(interrupt_id="approval-1", accepted=True, arguments="{}")], + ) + + assert lifecycle.get(occurrence.identity).status is ApprovalStatus.EXPIRED diff --git a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py index ab7e609af1..4c73054d2d 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -2133,8 +2133,8 @@ async def test_endpoint_agent_approval_cancellation_does_not_release_already_app assert messages_received == [] -async def test_endpoint_agent_approval_replayed_resume_entry_emits_run_error(): - """A consumed server-side approval cannot be replayed to execute a tool again.""" +async def test_endpoint_agent_approval_replayed_resume_entry_reprojects_retained_result(): + """An identical retry reprojects the retained result without executing the tool again.""" client, agent, executed_cities = _build_weather_approval_endpoint() first_response = client.post( @@ -2162,11 +2162,12 @@ async def test_endpoint_agent_approval_replayed_resume_entry_emits_run_error(): assert replay_response.status_code == 200 replay_events = _decode_sse_events(replay_response) - run_errors = [event for event in replay_events if event.get("type") == "RUN_ERROR"] - assert len(run_errors) == 1 - assert run_errors[0]["code"] == "APPROVAL_RESUME_NOT_FOUND" assert executed_cities == ["Seattle"] - assert not [event for event in replay_events if event.get("type") == "TOOL_CALL_RESULT"] + assert not [event for event in replay_events if event.get("type") == "RUN_ERROR"] + result_events = [event for event in replay_events if event.get("type") == "TOOL_CALL_RESULT"] + assert len(result_events) == 1 + assert result_events[0]["toolCallId"] == "call_get_weather" + assert result_events[0]["content"] == "Sunny in Seattle" async def test_endpoint_agent_approval_resume_wrong_thread_emits_run_error(): From 5c21854ec475cb2189154af82f90c4f874934e11 Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Mon, 10 Aug 2026 14:09:57 +0900 Subject: [PATCH 04/13] Separate approval execution ownership Key decisions: - Carry explicit local, hosted, deferred in-run, or unavailable ownership on every approval occurrence and authorized intent. - Keep lifecycle authorization separate from execution; local calls execute only through the local adapter while hosted and setup-injected decisions forward through owner-specific adapters. - Leave declaration-only calls pending when no transition owner can act, and settle forwarded outcomes against the original occurrence without local fallback. Files changed: - packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py - packages/ag-ui/agent_framework_ag_ui/_approval_state.py - packages/ag-ui/agent_framework_ag_ui/_agent_run.py - packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py - packages/ag-ui/tests/ag_ui/test_endpoint.py Verification: - 967 AG-UI tests passed with 92% package coverage and 89% approval lifecycle coverage. - 94 focused lifecycle, hosted, deferred-owner, hostile-resume, and approval tests passed. - Ruff format/check and package-local Pyright passed. - git diff --check passed. Notes for next iteration: - Executing-without-outcome recovery remains for the indeterminate execution-window issue. - The function-calling-loop scenario mapping remains inaccessible under the organization content-exclusion policy. - Workspace Poe fan-out remains blocked by the pre-existing missing packages/durabletask/pyproject.toml; equivalent package-local checks passed. --- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 187 ++++++++++-- .../_approval_lifecycle.py | 265 +++++++++++++++++- .../agent_framework_ag_ui/_approval_state.py | 105 ++++++- .../tests/ag_ui/test_approval_lifecycle.py | 83 ++++++ .../ag-ui/tests/ag_ui/test_endpoint.py | 18 ++ 5 files changed, 630 insertions(+), 28 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index cd01bf0802..044848b134 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py @@ -56,8 +56,12 @@ ) from ._approval_lifecycle import ( + ApprovalExecutionOwner, ApprovalLifecycle, AuthorizedExecution, + DeferredPendingToolTransitionOwner, + ForwardedPendingToolTransitionOwner, + HostedPendingToolTransitionOwner, LocalPendingToolTransitionOwner, ResumeDecision, ) @@ -698,6 +702,7 @@ class _PendingApprovalWithSiblings(_PendingApproval, total=False): """Pending approval details including sibling calls and trusted hosted metadata.""" already_approved_requests: list[dict[str, Any]] + execution_owner: str server_label: str @@ -718,6 +723,7 @@ def _make_pending_approval_entry( interrupt_id: str | None = None, already_approved_requests: list[dict[str, Any]] | None = None, server_label: str | None = None, + execution_owner: ApprovalExecutionOwner | None = None, ) -> _PendingApprovalWithSiblings: entry: _PendingApprovalWithSiblings = { "name": name, @@ -729,6 +735,8 @@ def _make_pending_approval_entry( entry["already_approved_requests"] = already_approved_requests if server_label: entry["server_label"] = server_label + if execution_owner is not None: + entry["execution_owner"] = execution_owner.value return entry @@ -774,6 +782,23 @@ def _function_call_server_label(function_call: Content | None) -> str | None: return server_label if isinstance(server_label, str) and server_label else None +def _function_call_execution_owner( + function_call: Content, + tools: list[Any] | None, + *, + has_deferred_owner: bool = False, +) -> ApprovalExecutionOwner: + """Resolve execution ownership only after the call and available tools exist.""" + if _function_call_server_label(function_call): + return ApprovalExecutionOwner.HOSTED + tool = _get_tool_map(tools).get(function_call.name) if tools and function_call.name else None + if tool is not None and not getattr(tool, "declaration_only", False): + return ApprovalExecutionOwner.LOCAL + if has_deferred_owner: + return ApprovalExecutionOwner.DEFERRED + return ApprovalExecutionOwner.UNAVAILABLE + + def _stored_already_approved_requests_for_visible_approval( session: AgentSession, *approval_ids: str | None, @@ -894,6 +919,9 @@ def _register_server_generated_approval_response( response: Content, pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] | None, thread_id: str, + tools: list[Any] | None, + *, + has_deferred_owner: bool, ) -> None: """Register a server-owned approval response so normal validation can consume it.""" if pending_approvals is None or response.function_call is None or not response.function_call.name: @@ -901,12 +929,18 @@ def _register_server_generated_approval_response( response_id = response.id or response.function_call.call_id if not response_id: return + execution_owner = _function_call_execution_owner( + response.function_call, + tools, + has_deferred_owner=has_deferred_owner, + ) entry = _make_pending_approval_entry( response.function_call.name, canonical_function_arguments(response.function_call), request_id=str(response.id) if response.id else None, interrupt_id=str(response.function_call.call_id) if response.function_call.call_id else None, server_label=_function_call_server_label(response.function_call), + execution_owner=execution_owner, ) _register_pending_approval_entry( pending_approvals, @@ -921,6 +955,7 @@ def _pop_collected_tool_approval_response_messages( session: AgentSession, pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] | None, thread_id: str, + tools: list[Any] | None, *, lifecycle: ApprovalLifecycle | None = None, authorized_executions: dict[str, AuthorizedExecution] | None = None, @@ -940,7 +975,13 @@ def _pop_collected_tool_approval_response_messages( response = _content_from_approval_state(raw_response) if response is None or response.type != "function_approval_response": continue - _register_server_generated_approval_response(response, pending_approvals, thread_id) + _register_server_generated_approval_response( + response, + pending_approvals, + thread_id, + tools, + has_deferred_owner=True, + ) function_call = response.function_call if ( lifecycle is not None @@ -949,10 +990,22 @@ def _pop_collected_tool_approval_response_messages( and function_call is not None and function_call.call_id and function_call.name - and not _function_call_server_label(function_call) ): arguments = canonical_function_arguments(function_call) or "{}" - lifecycle.register_local( + execution_owner = _function_call_execution_owner( + function_call, + tools, + has_deferred_owner=True, + ) + if execution_owner is ApprovalExecutionOwner.HOSTED: + register = lifecycle.register_hosted + elif execution_owner is ApprovalExecutionOwner.DEFERRED: + register = lifecycle.register_deferred + elif execution_owner is ApprovalExecutionOwner.LOCAL: + register = lifecycle.register_local + else: + register = lifecycle.register_unowned + register( thread_id=thread_id, interrupt_id=str(response.id or function_call.call_id), call_id=str(function_call.call_id), @@ -1404,14 +1457,11 @@ def _canonical_approval_resume_messages( if cancelled_ids: if lifecycle is not None: - local_cancelled_ids = [ - interrupt_id - for interrupt_id in cancelled_ids - if (pending_entry := entries_by_interrupt_id.get(interrupt_id)) is not None - and not _pending_approval_server_label(pending_entry) + lifecycle_cancelled_ids = [ + interrupt_id for interrupt_id in cancelled_ids if entries_by_interrupt_id.get(interrupt_id) is not None ] try: - lifecycle.cancel_batch(thread_id=thread_id, interrupt_ids=local_cancelled_ids) + lifecycle.cancel_batch(thread_id=thread_id, interrupt_ids=lifecycle_cancelled_ids) except (KeyError, ValueError) as exc: return ( [], @@ -1511,7 +1561,7 @@ def _canonical_approval_resume_messages( canonical_arguments = json.dumps(make_json_safe(merged_arguments), sort_keys=True, separators=(",", ":")) if not isinstance(pending_entry, str): argument_updates.append((pending_entry, canonical_arguments)) - if lifecycle is not None and not _pending_approval_server_label(pending_entry): + if lifecycle is not None: lifecycle_decisions.append( ResumeDecision( interrupt_id=interrupt_id, @@ -1550,6 +1600,11 @@ def _canonical_approval_resume_messages( request_id=str(response.id) if response.id else None, interrupt_id=str(function_call.call_id) if function_call.call_id else None, server_label=_function_call_server_label(function_call), + execution_owner=( + ApprovalExecutionOwner.HOSTED + if _function_call_server_label(function_call) + else ApprovalExecutionOwner.LOCAL + ), ) _register_pending_approval_entry( pending_approvals, @@ -1558,11 +1613,16 @@ def _canonical_approval_resume_messages( str(response_id), str(function_call.call_id) if function_call.call_id else None, ) - if lifecycle is not None and not _function_call_server_label(function_call): + if lifecycle is not None: sibling_interrupt_id = str(response_id) sibling_call_id = str(function_call.call_id or response_id) sibling_arguments = canonical_function_arguments(function_call) or "{}" - lifecycle.register_local( + register = ( + lifecycle.register_hosted + if _function_call_server_label(function_call) + else lifecycle.register_local + ) + register( thread_id=thread_id, interrupt_id=sibling_interrupt_id, call_id=sibling_call_id, @@ -1631,6 +1691,9 @@ async def _resolve_approval_responses( *, lifecycle: ApprovalLifecycle | None = None, authorized_executions: dict[str, AuthorizedExecution] | None = None, + forwarded_executions: ( + dict[str, tuple[ForwardedPendingToolTransitionOwner, AuthorizedExecution, Content]] | None + ) = None, ) -> list[Content]: """Execute approved function calls and replace approval content with results. @@ -1667,6 +1730,7 @@ async def _resolve_approval_responses( valid_response_content_ids: set[int] | None = None pending_local_response_content_ids: set[int] | None = None + validated_forwarded_approvals: list[Content] = [] response_content_ids_to_strip: set[int] = set() if pending_approvals is not None: valid_response_content_ids = set() @@ -1760,12 +1824,33 @@ def matches_pending_entry(candidate: PendingApprovalEntry | None, expected: Pend continue server_label = _pending_approval_server_label(pending_entry) + intent: AuthorizedExecution | None = None if primary_response.function_call is not None: if server_label: primary_response.function_call.additional_properties["server_label"] = server_label else: primary_response.function_call.additional_properties.pop("server_label", None) + if ( + primary_response.approved + and lifecycle is not None + and authorized_executions is not None + and primary_response.function_call is not None + ): + call_id = primary_response.function_call.call_id or primary_response.id or "" + intent = authorized_executions.get(call_id) + if intent is None: + logger.warning( + "Approval remains pending because no transition owner can act for call_id=%s.", call_id + ) + response_content_ids_to_strip.add(id(primary_response)) + continue valid_response_content_ids.add(id(primary_response)) + if ( + primary_response.approved + and intent is not None + and intent.owner in {ApprovalExecutionOwner.HOSTED, ApprovalExecutionOwner.DEFERRED} + ): + validated_forwarded_approvals.append(primary_response) if not server_label: pending_local_response_content_ids.add(id(primary_response)) _consume_pending_approval_entry( @@ -1808,6 +1893,34 @@ def matches_pending_entry(candidate: PendingApprovalEntry | None, expected: Pend pending_response_content_ids=pending_local_response_content_ids, ) + if ( + validated_forwarded_approvals + and lifecycle is not None + and authorized_executions is not None + and forwarded_executions is not None + ): + for approval in validated_forwarded_approvals: + function_call = approval.function_call + call_id = (function_call.call_id if function_call else None) or approval.id or "" + intent = authorized_executions.get(call_id) + if intent is None: + logger.warning("Skipping hosted approval without lifecycle authority for call_id=%s.", call_id) + continue + + async def forward_hosted_decision(approval: Content = approval) -> list[Content]: + return [approval] + + if intent.owner is ApprovalExecutionOwner.HOSTED: + forwarded_owner: ForwardedPendingToolTransitionOwner = HostedPendingToolTransitionOwner( + forward_hosted_decision + ) + else: + forwarded_owner = DeferredPendingToolTransitionOwner(forward_hosted_decision) + forwarded = await forwarded_owner.forward(intent, lifecycle=lifecycle) + if len(forwarded) != 1: + raise RuntimeError("Hosted transition owner did not forward exactly one approval decision.") + forwarded_executions[call_id] = (forwarded_owner, intent, forwarded[0]) + fcc_todo = _collect_approval_responses(messages) if valid_response_content_ids is not None: fcc_todo = { @@ -1865,8 +1978,8 @@ async def execute_local_call(approval: Content = approval, call_id: str = call_i return [Content.from_function_result(call_id=call_id, result="Error: Tool call invocation failed.")] return result_groups[0] - owner = LocalPendingToolTransitionOwner(execute_local_call) - outcome = await owner.execute(intent, lifecycle=lifecycle) + local_owner = LocalPendingToolTransitionOwner(execute_local_call) + outcome = await local_owner.execute(intent, lifecycle=lifecycle) approved_function_result_groups.append(list(outcome.result_group)) # Normalize one group per static approval and collect only terminal results for TOOL_CALL_RESULT events. @@ -2401,6 +2514,7 @@ async def run_agent_stream( ) authorized_executions: dict[str, AuthorizedExecution] = {} + forwarded_executions: dict[str, tuple[ForwardedPendingToolTransitionOwner, AuthorizedExecution, Content]] = {} retained_approval_results: list[Content] = [] approval_resume_messages, handled_resume_ids, cancelled_resume_ids, resume_error = ( _canonical_approval_resume_messages( @@ -2522,6 +2636,7 @@ async def run_agent_stream( session, pending_approvals, approval_thread_id, + tools_for_execution, lifecycle=approval_state_store.lifecycle if approval_state_store is not None else None, authorized_executions=authorized_executions, ) @@ -2537,6 +2652,7 @@ async def run_agent_stream( validated_approved_responses, lifecycle=approval_state_store.lifecycle if approval_state_store is not None else None, authorized_executions=authorized_executions, + forwarded_executions=forwarded_executions, ) # Defense-in-depth: replace approval payloads in snapshot with actual tool results @@ -2655,6 +2771,15 @@ async def run_agent_stream( content_type = getattr(content, "type", None) logger.debug(f"Processing content type={content_type}, message_id={flow.message_id}") + if ( + content_type == "function_result" + and content.call_id + and (forwarded := forwarded_executions.pop(content.call_id, None)) is not None + and approval_state_store is not None + ): + owner, intent, _ = forwarded + owner.record_outcome(intent, [content], lifecycle=approval_state_store.lifecycle) + # Register pending approval requests so we can validate responses later if content_type == "function_approval_request" and pending_approvals is not None: if content.id and content.function_call and content.function_call.name: @@ -2669,15 +2794,28 @@ async def run_agent_stream( str(content.id), str(canonical_interrupt_id) if canonical_interrupt_id else None, ) - if approval_state_store is not None and not server_label: - approval_state_store.register_local( - thread_ids=[approval_thread_id, provider_approval_thread_id], - name=content.function_call.name, - arguments=canonical_function_arguments(content.function_call) or "{}", - request_id=str(content.id), - interrupt_id=str(canonical_interrupt_id), - already_approved_requests=already_approved_requests, + if approval_state_store is not None: + execution_owner = _function_call_execution_owner( + content.function_call, + tools, + has_deferred_owner=_TOOL_APPROVAL_STATE_KEY in session.state, ) + registration_kwargs = { + "thread_ids": [approval_thread_id, provider_approval_thread_id], + "name": content.function_call.name, + "arguments": canonical_function_arguments(content.function_call) or "{}", + "request_id": str(content.id), + "interrupt_id": str(canonical_interrupt_id), + "already_approved_requests": already_approved_requests, + } + if execution_owner is ApprovalExecutionOwner.HOSTED: + approval_state_store.register_hosted(server_label=server_label, **registration_kwargs) + elif execution_owner is ApprovalExecutionOwner.DEFERRED: + approval_state_store.register_deferred(**registration_kwargs) + elif execution_owner is ApprovalExecutionOwner.LOCAL: + approval_state_store.register_local(**registration_kwargs) + else: + approval_state_store.register_unowned(**registration_kwargs) else: _register_pending_approval( pending_approvals, @@ -2893,6 +3031,11 @@ async def run_agent_stream( # Always emit RunFinished - confirm_changes tool call is complete (Start -> Args -> End) # The UI will show confirmation dialog and send a new request when user responds + if approval_state_store is not None: + for owner, intent, forwarded_approval in forwarded_executions.values(): + owner.record_outcome(intent, [forwarded_approval], lifecycle=approval_state_store.lifecycle) + forwarded_executions.clear() + persisted_messages = latest_messages_snapshot if resume_payload is not None and not seeded_resume_from_snapshot: # Generic resume requests carry only the synthesized response, so prepend diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py b/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py index 7e13d806cc..b649a14a99 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py @@ -24,6 +24,15 @@ class ApprovalStatus(str, Enum): EXPIRED = "expired" +class ApprovalExecutionOwner(str, Enum): + """Runtime owner authorized to continue an approved occurrence.""" + + LOCAL = "local" + HOSTED = "hosted" + DEFERRED = "deferred" + UNAVAILABLE = "unavailable" + + @dataclass(frozen=True) class ApprovalOccurrenceIdentity: """Identity of one occurrence within a scoped server-owned thread.""" @@ -53,6 +62,7 @@ class ApprovalOccurrence: thread_ids: tuple[str, ...] name: str arguments: str + owner: ApprovalExecutionOwner status: ApprovalStatus = ApprovalStatus.PENDING replayable_results: list[ReplayableToolResult] = field(default_factory=list) decision: ResumeDecision | None = None @@ -61,11 +71,12 @@ class ApprovalOccurrence: @dataclass(frozen=True) class AuthorizedExecution: - """Authority for a Pending Tool Transition Owner to execute one local call.""" + """Authority for one Pending Tool Transition Owner to continue a call.""" identity: ApprovalOccurrenceIdentity name: str arguments: str + owner: ApprovalExecutionOwner @dataclass(frozen=True) @@ -132,6 +143,136 @@ def register_local_aliases( arguments: str, ) -> ApprovalOccurrence: """Register one occurrence under its trusted scoped-thread aliases.""" + return self._register_aliases( + thread_ids=thread_ids, + interrupt_id=interrupt_id, + call_id=call_id, + name=name, + arguments=arguments, + owner=ApprovalExecutionOwner.LOCAL, + ) + + def register_hosted( + self, + *, + thread_id: str, + interrupt_id: str, + call_id: str, + name: str, + arguments: str, + ) -> ApprovalOccurrence: + """Register one server-generated hosted approval occurrence.""" + return self.register_hosted_aliases( + thread_ids=[thread_id], + interrupt_id=interrupt_id, + call_id=call_id, + name=name, + arguments=arguments, + ) + + def register_hosted_aliases( + self, + *, + thread_ids: list[str], + interrupt_id: str, + call_id: str, + name: str, + arguments: str, + ) -> ApprovalOccurrence: + """Register one hosted occurrence under its trusted scoped-thread aliases.""" + return self._register_aliases( + thread_ids=thread_ids, + interrupt_id=interrupt_id, + call_id=call_id, + name=name, + arguments=arguments, + owner=ApprovalExecutionOwner.HOSTED, + ) + + def register_unowned( + self, + *, + thread_id: str, + interrupt_id: str, + call_id: str, + name: str, + arguments: str, + ) -> ApprovalOccurrence: + """Register an occurrence that has no executable transition owner.""" + return self.register_unowned_aliases( + thread_ids=[thread_id], + interrupt_id=interrupt_id, + call_id=call_id, + name=name, + arguments=arguments, + ) + + def register_deferred( + self, + *, + thread_id: str, + interrupt_id: str, + call_id: str, + name: str, + arguments: str, + ) -> ApprovalOccurrence: + """Register one occurrence owned by the in-run transition pipeline.""" + return self.register_deferred_aliases( + thread_ids=[thread_id], + interrupt_id=interrupt_id, + call_id=call_id, + name=name, + arguments=arguments, + ) + + def register_deferred_aliases( + self, + *, + thread_ids: list[str], + interrupt_id: str, + call_id: str, + name: str, + arguments: str, + ) -> ApprovalOccurrence: + """Register one deferred occurrence under its trusted scoped-thread aliases.""" + return self._register_aliases( + thread_ids=thread_ids, + interrupt_id=interrupt_id, + call_id=call_id, + name=name, + arguments=arguments, + owner=ApprovalExecutionOwner.DEFERRED, + ) + + def register_unowned_aliases( + self, + *, + thread_ids: list[str], + interrupt_id: str, + call_id: str, + name: str, + arguments: str, + ) -> ApprovalOccurrence: + """Register an unowned occurrence under its trusted scoped-thread aliases.""" + return self._register_aliases( + thread_ids=thread_ids, + interrupt_id=interrupt_id, + call_id=call_id, + name=name, + arguments=arguments, + owner=ApprovalExecutionOwner.UNAVAILABLE, + ) + + def _register_aliases( + self, + *, + thread_ids: list[str], + interrupt_id: str, + call_id: str, + name: str, + arguments: str, + owner: ApprovalExecutionOwner, + ) -> ApprovalOccurrence: unique_thread_ids = tuple(dict.fromkeys(thread_ids)) if not unique_thread_ids: raise ValueError("An approval occurrence requires at least one scoped thread identity.") @@ -146,7 +287,12 @@ def register_local_aliases( occurrence = self._occurrences[next(iter(existing_identities))] if occurrence.status is not ApprovalStatus.PENDING: raise ValueError(f"Approval occurrence is not pending: {occurrence.status}.") - if occurrence.identity.call_id != call_id or occurrence.name != name or occurrence.arguments != arguments: + if ( + occurrence.identity.call_id != call_id + or occurrence.name != name + or occurrence.arguments != arguments + or occurrence.owner is not owner + ): raise ValueError("Approval alias conflicts with an existing pending occurrence.") occurrence.thread_ids = tuple(dict.fromkeys((*occurrence.thread_ids, *unique_thread_ids))) for thread_id in occurrence.thread_ids: @@ -164,6 +310,7 @@ def register_local_aliases( thread_ids=unique_thread_ids, name=name, arguments=arguments, + owner=owner, ) self._occurrences[identity] = occurrence for thread_id in unique_thread_ids: @@ -264,12 +411,15 @@ def claim_batch( if decision.arguments is None: raise RuntimeError("Validated pending approval is missing canonical arguments.") occurrence.arguments = decision.arguments + if occurrence.owner is ApprovalExecutionOwner.UNAVAILABLE: + continue occurrence.status = ApprovalStatus.CLAIMED intents.append( AuthorizedExecution( identity=occurrence.identity, name=occurrence.name, arguments=occurrence.arguments, + owner=occurrence.owner, ) ) return ApprovalBatchDecision( @@ -325,9 +475,13 @@ def _remove_pending_aliases(self, occurrence: ApprovalOccurrence) -> None: self._pending_by_interrupt.pop((thread_id, occurrence.identity.interrupt_id), None) self._terminal_by_interrupt[(thread_id, occurrence.identity.interrupt_id)] = occurrence.identity - def begin_execution(self, intent: AuthorizedExecution) -> None: + def begin_execution(self, intent: AuthorizedExecution, *, owner: ApprovalExecutionOwner) -> None: """Mark a claimed occurrence immediately before its owner may invoke a tool.""" occurrence = self._occurrences[intent.identity] + if intent.owner is not owner or occurrence.owner is not owner: + raise ValueError( + f"Approval occurrence belongs to the {occurrence.owner.value} transition owner, not {owner.value}." + ) if occurrence.status is not ApprovalStatus.CLAIMED: raise ValueError(f"Approval occurrence is not claimed: {occurrence.status}.") occurrence.status = ApprovalStatus.EXECUTING @@ -335,6 +489,8 @@ def begin_execution(self, intent: AuthorizedExecution) -> None: def settle(self, intent: AuthorizedExecution, results: list[Content]) -> ApprovalOutcome: """Settle an executing occurrence with results under its original call identity.""" occurrence = self._occurrences[intent.identity] + if intent.owner is not ApprovalExecutionOwner.LOCAL or occurrence.owner is not ApprovalExecutionOwner.LOCAL: + raise ValueError("Only the local transition owner can settle a local execution result.") if occurrence.status is not ApprovalStatus.EXECUTING: raise ValueError(f"Approval occurrence is not executing: {occurrence.status}.") replayable_results = [ @@ -355,9 +511,54 @@ def settle(self, intent: AuthorizedExecution, results: list[Content]) -> Approva occurrence.outcome = outcome return outcome + def settle_forwarded( + self, + intent: AuthorizedExecution, + results: list[Content], + *, + owner: ApprovalExecutionOwner, + ) -> ApprovalOutcome: + """Record a forwarded owner's outcome under its original occurrence.""" + occurrence = self._occurrences[intent.identity] + if owner not in {ApprovalExecutionOwner.HOSTED, ApprovalExecutionOwner.DEFERRED}: + raise ValueError("Only a forwarding transition owner can record a forwarded outcome.") + if intent.owner is not owner or occurrence.owner is not owner: + raise ValueError( + f"Approval occurrence belongs to the {occurrence.owner.value} transition owner, not {owner.value}." + ) + if occurrence.status is not ApprovalStatus.EXECUTING: + raise ValueError(f"Approval occurrence is not executing: {occurrence.status}.") + replayable_results = [ + ReplayableToolResult(content=result) + for result in results + if result.type == "function_result" and result.call_id == occurrence.identity.call_id + ] + forwarded_responses = [ + result + for result in results + if result.type == "function_approval_response" + and result.approved + and result.function_call is not None + and result.function_call.call_id == occurrence.identity.call_id + ] + if len(replayable_results) + len(forwarded_responses) != 1: + raise ValueError("A hosted approval must record exactly one outcome for its original call.") + outcome = ApprovalOutcome( + identity=occurrence.identity, + replayable_results=tuple(replayable_results), + result_group=tuple(results), + ) + occurrence.replayable_results = replayable_results + occurrence.status = ApprovalStatus.SETTLED + occurrence.outcome = outcome + self._remove_pending_aliases(occurrence) + return outcome + def defer(self, intent: AuthorizedExecution, results: list[Content]) -> ApprovalOutcome: """Return an execution that yielded only follow-up requests to pending.""" occurrence = self._occurrences[intent.identity] + if intent.owner is not ApprovalExecutionOwner.LOCAL or occurrence.owner is not ApprovalExecutionOwner.LOCAL: + raise ValueError("Only the local transition owner can defer a local execution.") if occurrence.status is not ApprovalStatus.EXECUTING: raise ValueError(f"Approval occurrence is not executing: {occurrence.status}.") occurrence.status = ApprovalStatus.PENDING @@ -377,8 +578,64 @@ async def execute( lifecycle: ApprovalLifecycle, ) -> ApprovalOutcome: """Execute and settle one call after lifecycle authorization.""" - lifecycle.begin_execution(intent) + lifecycle.begin_execution(intent, owner=ApprovalExecutionOwner.LOCAL) results = await self._executor() if not any(result.type == "function_result" for result in results): return lifecycle.defer(intent, results) return lifecycle.settle(intent, results) + + +class ForwardedPendingToolTransitionOwner: + """Forward an authorized decision to a non-transport transition owner.""" + + def __init__( + self, + forwarder: Callable[[], Awaitable[list[Content]]], + *, + owner: ApprovalExecutionOwner, + ) -> None: + self._forwarder = forwarder + self._owner = owner + + async def forward( + self, + intent: AuthorizedExecution, + *, + lifecycle: ApprovalLifecycle, + ) -> list[Content]: + """Forward one hosted approval after lifecycle authorization.""" + lifecycle.begin_execution(intent, owner=self._owner) + return await self._forwarder() + + def record_outcome( + self, + intent: AuthorizedExecution, + results: list[Content], + *, + lifecycle: ApprovalLifecycle, + ) -> ApprovalOutcome: + """Record the hosted owner's outcome against the authorized occurrence.""" + return lifecycle.settle_forwarded(intent, results, owner=self._owner) + + async def execute( + self, + intent: AuthorizedExecution, + *, + lifecycle: ApprovalLifecycle, + ) -> ApprovalOutcome: + """Forward and immediately record an available hosted outcome.""" + return self.record_outcome(intent, await self.forward(intent, lifecycle=lifecycle), lifecycle=lifecycle) + + +class HostedPendingToolTransitionOwner(ForwardedPendingToolTransitionOwner): + """Forward an authorized decision through the hosted transition owner.""" + + def __init__(self, forwarder: Callable[[], Awaitable[list[Content]]]) -> None: + super().__init__(forwarder, owner=ApprovalExecutionOwner.HOSTED) + + +class DeferredPendingToolTransitionOwner(ForwardedPendingToolTransitionOwner): + """Forward an authorized decision to the in-run transition pipeline.""" + + def __init__(self, forwarder: Callable[[], Awaitable[list[Content]]]) -> None: + super().__init__(forwarder, owner=ApprovalExecutionOwner.DEFERRED) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py b/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py index 540f90f8ed..052389b201 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py @@ -7,7 +7,7 @@ from collections import OrderedDict from typing import Any -from ._approval_lifecycle import ApprovalLifecycle +from ._approval_lifecycle import ApprovalExecutionOwner, ApprovalLifecycle ApprovalScope = str """Application-defined scope for server-side AG-UI Approval State.""" @@ -65,6 +65,96 @@ def register_local( already_approved_requests: list[dict[str, Any]] | None = None, ) -> None: """Register one local occurrence and its trusted aliases.""" + self._register( + thread_ids=thread_ids, + name=name, + arguments=arguments, + request_id=request_id, + interrupt_id=interrupt_id, + already_approved_requests=already_approved_requests, + server_label=None, + owner=ApprovalExecutionOwner.LOCAL, + ) + + def register_hosted( + self, + *, + thread_ids: list[str], + name: str, + arguments: str, + request_id: str, + interrupt_id: str, + server_label: str | None, + already_approved_requests: list[dict[str, Any]] | None = None, + ) -> None: + """Register one hosted occurrence and its trusted aliases.""" + self._register( + thread_ids=thread_ids, + name=name, + arguments=arguments, + request_id=request_id, + interrupt_id=interrupt_id, + already_approved_requests=already_approved_requests, + server_label=server_label, + owner=ApprovalExecutionOwner.HOSTED, + ) + + def register_unowned( + self, + *, + thread_ids: list[str], + name: str, + arguments: str, + request_id: str, + interrupt_id: str, + already_approved_requests: list[dict[str, Any]] | None = None, + ) -> None: + """Register one occurrence that has no executable transition owner.""" + self._register( + thread_ids=thread_ids, + name=name, + arguments=arguments, + request_id=request_id, + interrupt_id=interrupt_id, + already_approved_requests=already_approved_requests, + server_label=None, + owner=ApprovalExecutionOwner.UNAVAILABLE, + ) + + def register_deferred( + self, + *, + thread_ids: list[str], + name: str, + arguments: str, + request_id: str, + interrupt_id: str, + already_approved_requests: list[dict[str, Any]] | None = None, + ) -> None: + """Register one occurrence owned by the in-run transition pipeline.""" + self._register( + thread_ids=thread_ids, + name=name, + arguments=arguments, + request_id=request_id, + interrupt_id=interrupt_id, + already_approved_requests=already_approved_requests, + server_label=None, + owner=ApprovalExecutionOwner.DEFERRED, + ) + + def _register( + self, + *, + thread_ids: list[str], + name: str, + arguments: str, + request_id: str, + interrupt_id: str, + already_approved_requests: list[dict[str, Any]] | None, + server_label: str | None, + owner: ApprovalExecutionOwner, + ) -> None: entry: dict[str, Any] = { "name": name, "arguments": arguments, @@ -73,9 +163,20 @@ def register_local( } if already_approved_requests: entry["already_approved_requests"] = already_approved_requests + if server_label: + entry["server_label"] = server_label + entry["execution_owner"] = owner.value unique_thread_ids = list(dict.fromkeys(thread_ids)) - self.lifecycle.register_local_aliases( + if owner is ApprovalExecutionOwner.HOSTED: + register_aliases = self.lifecycle.register_hosted_aliases + elif owner is ApprovalExecutionOwner.DEFERRED: + register_aliases = self.lifecycle.register_deferred_aliases + elif owner is ApprovalExecutionOwner.UNAVAILABLE: + register_aliases = self.lifecycle.register_unowned_aliases + else: + register_aliases = self.lifecycle.register_local_aliases + register_aliases( thread_ids=unique_thread_ids, interrupt_id=interrupt_id, call_id=interrupt_id, diff --git a/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py b/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py index 827e4376bb..85a2014c1c 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py +++ b/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py @@ -8,8 +8,10 @@ from agent_framework import Content from agent_framework_ag_ui._approval_lifecycle import ( + ApprovalExecutionOwner, ApprovalLifecycle, ApprovalStatus, + HostedPendingToolTransitionOwner, LocalPendingToolTransitionOwner, ResumeDecision, ) @@ -55,6 +57,67 @@ async def execute_authorized_call() -> list[Content]: assert [result.content.result for result in outcome.replayable_results] == ["Sunny"] +async def test_hosted_approval_is_forwarded_only_by_its_owner_and_settles_same_occurrence() -> None: + """Hosted authority cannot execute locally and records forwarding against its occurrence.""" + lifecycle = ApprovalLifecycle() + occurrence = lifecycle.register_hosted( + thread_id="thread-1", + interrupt_id="approval-1", + call_id="call-1", + name="hosted_search", + arguments='{"query":"azure"}', + ) + intent = lifecycle.claim( + thread_id="thread-1", + decision=ResumeDecision( + interrupt_id="approval-1", + accepted=True, + arguments='{"query":"azure"}', + ), + ) + local_invocations = 0 + + async def execute_locally() -> list[Content]: + nonlocal local_invocations + local_invocations += 1 + return [Content.from_function_result(call_id="call-1", result="local")] + + with pytest.raises(ValueError, match="hosted"): + await LocalPendingToolTransitionOwner(execute_locally).execute(intent, lifecycle=lifecycle) + + forwarded_response = Content.from_function_approval_response( + approved=True, + id="approval-1", + function_call=Content.from_function_call( + call_id="call-1", + name="hosted_search", + arguments={"query": "azure"}, + additional_properties={"server_label": "hosted"}, + ), + ) + hosted_forwards = 0 + + async def forward_to_hosted_owner() -> list[Content]: + nonlocal hosted_forwards + hosted_forwards += 1 + return [forwarded_response] + + owner = HostedPendingToolTransitionOwner(forward_to_hosted_owner) + forwarded = await owner.forward(intent, lifecycle=lifecycle) + assert lifecycle.get(occurrence.identity).status is ApprovalStatus.EXECUTING + remote_result = Content.from_function_result(call_id="call-1", result="hosted") + outcome = owner.record_outcome(intent, [remote_result], lifecycle=lifecycle) + + assert intent.owner is ApprovalExecutionOwner.HOSTED + assert local_invocations == 0 + assert hosted_forwards == 1 + assert forwarded == [forwarded_response] + assert outcome.identity == occurrence.identity + assert outcome.result_group == (remote_result,) + assert [result.content for result in outcome.replayable_results] == [remote_result] + assert lifecycle.get(occurrence.identity).status is ApprovalStatus.SETTLED + + async def test_execution_failure_keeps_unexecuted_batch_sibling_claimed() -> None: """A failed occurrence does not erase a claimed sibling that can still execute.""" lifecycle = ApprovalLifecycle() @@ -127,6 +190,26 @@ def test_batch_validation_is_atomic_before_claiming_any_occurrence() -> None: assert lifecycle.get(second.identity).status is ApprovalStatus.PENDING +def test_accepted_declaration_without_execution_owner_remains_pending() -> None: + """Approval alone does not grant local authority to a declaration-only call.""" + lifecycle = ApprovalLifecycle() + occurrence = lifecycle.register_unowned( + thread_id="thread-1", + interrupt_id="approval-1", + call_id="call-1", + name="client_action", + arguments="{}", + ) + + batch = lifecycle.claim_batch( + thread_id="thread-1", + decisions=[ResumeDecision(interrupt_id="approval-1", accepted=True, arguments="{}")], + ) + + assert batch.authorized_executions == () + assert lifecycle.get(occurrence.identity).status is ApprovalStatus.PENDING + + def test_mixed_batch_accounts_for_rejection_under_original_call_identity() -> None: """A rejected occurrence remains represented while its accepted sibling is claimed.""" lifecycle = ApprovalLifecycle() diff --git a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py index 4c73054d2d..6ab195608f 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -5573,6 +5573,7 @@ async def test_endpoint_canonical_resume_preserves_hosted_approval_for_provider( state = {"phase": "pause"} local_executions: list[str] = [] provider_messages: list[Message] = [] + provider_invocations = 0 hosted_call = Content.from_function_call( call_id=call_id, name="docs_search", @@ -5595,7 +5596,9 @@ async def stream_fn( options: dict[str, Any], **kwargs: Any, ) -> AsyncIterator[ChatResponseUpdate]: + nonlocal provider_invocations del options, kwargs + provider_invocations += 1 if state["phase"] == "pause": yield ChatResponseUpdate( contents=[Content.from_function_approval_request(id=call_id, function_call=hosted_call)], @@ -5657,6 +5660,21 @@ async def stream_fn( assert approval_responses[0].function_call is not None assert approval_responses[0].function_call.additional_properties["server_label"] == server_label + retry_response = client.post( + "/approval", + json={ + "runId": "run-retry", + "threadId": "thread-hosted-approval", + "messages": [], + "resume": [{"interruptId": call_id, "status": "resolved", "payload": {"accepted": True}}], + }, + ) + + assert retry_response.status_code == 200 + assert not [event for event in _decode_sse_events(retry_response) if event.get("type") == "RUN_ERROR"] + assert provider_invocations == 2 + assert local_executions == [] + async def test_endpoint_does_not_forward_resolved_local_approval_control_to_chat_client( streaming_chat_client_stub, From de6fce99d65690687e360b043e68e3b1fad3eeb6 Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Mon, 10 Aug 2026 14:18:21 +0900 Subject: [PATCH 05/13] Represent approval execution uncertainty Key decisions: - Distinguish reserved claims from execution windows that may have started an external side effect. - Recover non-idempotent execution failures as indeterminate and reject identical retries without another invocation. - Permit claim release only under an explicit safe policy and execution retry only with a predeclared idempotency key shared by local and forwarded owners. Files changed: - packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py - packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py - packages/ag-ui/tests/ag_ui/test_endpoint.py Verification: - 972 AG-UI tests passed with 92% line coverage and 89% package branch coverage. - 23 focused lifecycle, duplicate-resume, hosted-owner, and public settlement-window tests passed. - Package-local Ruff and Pyright passed; git diff --check passed. Notes for next iteration: - The function-calling-loop scenario mapping remains inaccessible under the organization content-exclusion policy. - Workspace Poe fan-out remains blocked by the pre-existing missing packages/durabletask/pyproject.toml; equivalent package-local checks passed. --- .../_approval_lifecycle.py | 107 +++++++++++++- .../tests/ag_ui/test_approval_lifecycle.py | 131 +++++++++++++++++- .../ag-ui/tests/ag_ui/test_endpoint.py | 46 ++++++ 3 files changed, 274 insertions(+), 10 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py b/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py index b649a14a99..e3957390a4 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py @@ -4,6 +4,7 @@ from __future__ import annotations +from asyncio import CancelledError from collections.abc import Awaitable, Callable, Iterator from dataclasses import dataclass, field from enum import Enum @@ -12,6 +13,10 @@ from agent_framework import Content +class ApprovalIndeterminateError(ValueError): + """An approval may have executed but has no retained terminal outcome.""" + + class ApprovalStatus(str, Enum): """Lifecycle state of one server-owned approval occurrence.""" @@ -22,6 +27,7 @@ class ApprovalStatus(str, Enum): REJECTED = "rejected" CANCELLED = "cancelled" EXPIRED = "expired" + INDETERMINATE = "indeterminate" class ApprovalExecutionOwner(str, Enum): @@ -33,6 +39,12 @@ class ApprovalExecutionOwner(str, Enum): UNAVAILABLE = "unavailable" +class ClaimRecoveryPolicy(str, Enum): + """Proof required to release authority before execution begins.""" + + SAFE_TO_RETRY = "safe_to_retry" + + @dataclass(frozen=True) class ApprovalOccurrenceIdentity: """Identity of one occurrence within a scoped server-owned thread.""" @@ -63,6 +75,7 @@ class ApprovalOccurrence: name: str arguments: str owner: ApprovalExecutionOwner + idempotency_key: str | None = None status: ApprovalStatus = ApprovalStatus.PENDING replayable_results: list[ReplayableToolResult] = field(default_factory=list) decision: ResumeDecision | None = None @@ -77,6 +90,7 @@ class AuthorizedExecution: name: str arguments: str owner: ApprovalExecutionOwner + idempotency_key: str | None = None @dataclass(frozen=True) @@ -123,6 +137,7 @@ def register_local( call_id: str, name: str, arguments: str, + idempotency_key: str | None = None, ) -> ApprovalOccurrence: """Register one server-generated local approval occurrence.""" return self.register_local_aliases( @@ -131,6 +146,7 @@ def register_local( call_id=call_id, name=name, arguments=arguments, + idempotency_key=idempotency_key, ) def register_local_aliases( @@ -141,6 +157,7 @@ def register_local_aliases( call_id: str, name: str, arguments: str, + idempotency_key: str | None = None, ) -> ApprovalOccurrence: """Register one occurrence under its trusted scoped-thread aliases.""" return self._register_aliases( @@ -150,6 +167,7 @@ def register_local_aliases( name=name, arguments=arguments, owner=ApprovalExecutionOwner.LOCAL, + idempotency_key=idempotency_key, ) def register_hosted( @@ -160,6 +178,7 @@ def register_hosted( call_id: str, name: str, arguments: str, + idempotency_key: str | None = None, ) -> ApprovalOccurrence: """Register one server-generated hosted approval occurrence.""" return self.register_hosted_aliases( @@ -168,6 +187,7 @@ def register_hosted( call_id=call_id, name=name, arguments=arguments, + idempotency_key=idempotency_key, ) def register_hosted_aliases( @@ -178,6 +198,7 @@ def register_hosted_aliases( call_id: str, name: str, arguments: str, + idempotency_key: str | None = None, ) -> ApprovalOccurrence: """Register one hosted occurrence under its trusted scoped-thread aliases.""" return self._register_aliases( @@ -187,6 +208,7 @@ def register_hosted_aliases( name=name, arguments=arguments, owner=ApprovalExecutionOwner.HOSTED, + idempotency_key=idempotency_key, ) def register_unowned( @@ -215,6 +237,7 @@ def register_deferred( call_id: str, name: str, arguments: str, + idempotency_key: str | None = None, ) -> ApprovalOccurrence: """Register one occurrence owned by the in-run transition pipeline.""" return self.register_deferred_aliases( @@ -223,6 +246,7 @@ def register_deferred( call_id=call_id, name=name, arguments=arguments, + idempotency_key=idempotency_key, ) def register_deferred_aliases( @@ -233,6 +257,7 @@ def register_deferred_aliases( call_id: str, name: str, arguments: str, + idempotency_key: str | None = None, ) -> ApprovalOccurrence: """Register one deferred occurrence under its trusted scoped-thread aliases.""" return self._register_aliases( @@ -242,6 +267,7 @@ def register_deferred_aliases( name=name, arguments=arguments, owner=ApprovalExecutionOwner.DEFERRED, + idempotency_key=idempotency_key, ) def register_unowned_aliases( @@ -272,7 +298,10 @@ def _register_aliases( name: str, arguments: str, owner: ApprovalExecutionOwner, + idempotency_key: str | None = None, ) -> ApprovalOccurrence: + if idempotency_key == "": + raise ValueError("An execution idempotency key cannot be empty.") unique_thread_ids = tuple(dict.fromkeys(thread_ids)) if not unique_thread_ids: raise ValueError("An approval occurrence requires at least one scoped thread identity.") @@ -292,6 +321,7 @@ def _register_aliases( or occurrence.name != name or occurrence.arguments != arguments or occurrence.owner is not owner + or occurrence.idempotency_key != idempotency_key ): raise ValueError("Approval alias conflicts with an existing pending occurrence.") occurrence.thread_ids = tuple(dict.fromkeys((*occurrence.thread_ids, *unique_thread_ids))) @@ -311,6 +341,7 @@ def _register_aliases( name=name, arguments=arguments, owner=owner, + idempotency_key=idempotency_key, ) self._occurrences[identity] = occurrence for thread_id in unique_thread_ids: @@ -357,6 +388,10 @@ def claim_batch( identity = self._terminal_by_interrupt[key] occurrence = self._occurrences[identity] if is_terminal: + if occurrence.status is ApprovalStatus.INDETERMINATE: + raise ApprovalIndeterminateError( + "Approval execution outcome is indeterminate; automatic retry is unsafe." + ) if occurrence.status is ApprovalStatus.EXPIRED: raise ValueError("Approval authority has expired.") retained_decision = occurrence.decision @@ -420,6 +455,7 @@ def claim_batch( name=occurrence.name, arguments=occurrence.arguments, owner=occurrence.owner, + idempotency_key=occurrence.idempotency_key, ) ) return ApprovalBatchDecision( @@ -476,7 +512,11 @@ def _remove_pending_aliases(self, occurrence: ApprovalOccurrence) -> None: self._terminal_by_interrupt[(thread_id, occurrence.identity.interrupt_id)] = occurrence.identity def begin_execution(self, intent: AuthorizedExecution, *, owner: ApprovalExecutionOwner) -> None: - """Mark a claimed occurrence immediately before its owner may invoke a tool.""" + """Mark that an external side effect may begin. + + A claimed occurrence only reserves authority. Once this transition succeeds, + arbitrary execution cannot be assumed safe to retry without idempotency proof. + """ occurrence = self._occurrences[intent.identity] if intent.owner is not owner or occurrence.owner is not owner: raise ValueError( @@ -486,6 +526,47 @@ def begin_execution(self, intent: AuthorizedExecution, *, owner: ApprovalExecuti raise ValueError(f"Approval occurrence is not claimed: {occurrence.status}.") occurrence.status = ApprovalStatus.EXECUTING + def release_claim(self, intent: AuthorizedExecution, *, policy: ClaimRecoveryPolicy) -> None: + """Release reserved authority when execution is known not to have begun.""" + occurrence = self._occurrences[intent.identity] + if policy is not ClaimRecoveryPolicy.SAFE_TO_RETRY: + raise ValueError("Claim recovery policy does not permit retry.") + if occurrence.status is not ApprovalStatus.CLAIMED: + raise ValueError(f"Approval occurrence is not claimed: {occurrence.status}.") + occurrence.status = ApprovalStatus.PENDING + + def mark_indeterminate(self, intent: AuthorizedExecution, *, owner: ApprovalExecutionOwner) -> None: + """Record that execution may have begun but no result was settled.""" + occurrence = self._occurrences[intent.identity] + if intent.owner is not owner or occurrence.owner is not owner: + raise ValueError( + f"Approval occurrence belongs to the {occurrence.owner.value} transition owner, not {owner.value}." + ) + if occurrence.status is not ApprovalStatus.EXECUTING: + raise ValueError(f"Approval occurrence is not executing: {occurrence.status}.") + occurrence.status = ApprovalStatus.INDETERMINATE + self._remove_pending_aliases(occurrence) + + def recover_execution( + self, + intent: AuthorizedExecution, + *, + owner: ApprovalExecutionOwner, + ) -> AuthorizedExecution | None: + """Recover an execution that has no settled result.""" + occurrence = self._occurrences[intent.identity] + if intent.owner is not owner or occurrence.owner is not owner: + raise ValueError( + f"Approval occurrence belongs to the {occurrence.owner.value} transition owner, not {owner.value}." + ) + if occurrence.status is not ApprovalStatus.EXECUTING: + raise ValueError(f"Approval occurrence is not executing: {occurrence.status}.") + if intent.idempotency_key is not None and intent.idempotency_key == occurrence.idempotency_key: + occurrence.status = ApprovalStatus.CLAIMED + return intent + self.mark_indeterminate(intent, owner=owner) + return None + def settle(self, intent: AuthorizedExecution, results: list[Content]) -> ApprovalOutcome: """Settle an executing occurrence with results under its original call identity.""" occurrence = self._occurrences[intent.identity] @@ -579,10 +660,14 @@ async def execute( ) -> ApprovalOutcome: """Execute and settle one call after lifecycle authorization.""" lifecycle.begin_execution(intent, owner=ApprovalExecutionOwner.LOCAL) - results = await self._executor() - if not any(result.type == "function_result" for result in results): - return lifecycle.defer(intent, results) - return lifecycle.settle(intent, results) + try: + results = await self._executor() + if not any(result.type == "function_result" for result in results): + return lifecycle.defer(intent, results) + return lifecycle.settle(intent, results) + except (Exception, CancelledError): + lifecycle.recover_execution(intent, owner=ApprovalExecutionOwner.LOCAL) + raise class ForwardedPendingToolTransitionOwner: @@ -605,7 +690,11 @@ async def forward( ) -> list[Content]: """Forward one hosted approval after lifecycle authorization.""" lifecycle.begin_execution(intent, owner=self._owner) - return await self._forwarder() + try: + return await self._forwarder() + except (Exception, CancelledError): + lifecycle.recover_execution(intent, owner=self._owner) + raise def record_outcome( self, @@ -615,7 +704,11 @@ def record_outcome( lifecycle: ApprovalLifecycle, ) -> ApprovalOutcome: """Record the hosted owner's outcome against the authorized occurrence.""" - return lifecycle.settle_forwarded(intent, results, owner=self._owner) + try: + return lifecycle.settle_forwarded(intent, results, owner=self._owner) + except (Exception, CancelledError): + lifecycle.recover_execution(intent, owner=self._owner) + raise async def execute( self, diff --git a/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py b/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py index 85a2014c1c..7f2a83e80e 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py +++ b/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py @@ -9,8 +9,10 @@ from agent_framework_ag_ui._approval_lifecycle import ( ApprovalExecutionOwner, + ApprovalIndeterminateError, ApprovalLifecycle, ApprovalStatus, + ClaimRecoveryPolicy, HostedPendingToolTransitionOwner, LocalPendingToolTransitionOwner, ResumeDecision, @@ -118,8 +120,8 @@ async def forward_to_hosted_owner() -> list[Content]: assert lifecycle.get(occurrence.identity).status is ApprovalStatus.SETTLED -async def test_execution_failure_keeps_unexecuted_batch_sibling_claimed() -> None: - """A failed occurrence does not erase a claimed sibling that can still execute.""" +async def test_execution_failure_becomes_indeterminate_and_keeps_unexecuted_sibling_claimed() -> None: + """A possibly started side effect is not retried and does not erase a claimed sibling.""" lifecycle = ApprovalLifecycle() first = lifecycle.register_local( thread_id="thread-1", @@ -143,14 +145,26 @@ async def test_execution_failure_keeps_unexecuted_batch_sibling_claimed() -> Non ], ) + invocation_count = 0 + async def fail_first() -> list[Content]: + nonlocal invocation_count + invocation_count += 1 raise RuntimeError("side effect failed") with pytest.raises(RuntimeError, match="side effect failed"): await LocalPendingToolTransitionOwner(fail_first).execute(first_intent, lifecycle=lifecycle) - assert lifecycle.get(first.identity).status is ApprovalStatus.EXECUTING + assert lifecycle.get(first.identity).status is ApprovalStatus.INDETERMINATE assert lifecycle.get(second.identity).status is ApprovalStatus.CLAIMED + with pytest.raises(ApprovalIndeterminateError) as error: + lifecycle.claim_batch( + thread_id="thread-1", + decisions=[ResumeDecision(interrupt_id="approval-1", accepted=True, arguments='{"value":"first"}')], + ) + assert str(error.value) == "Approval execution outcome is indeterminate; automatic retry is unsafe." + assert "first" not in str(error.value) + assert invocation_count == 1 async def execute_second() -> list[Content]: return [Content.from_function_result(call_id="call-2", result="wrote second")] @@ -159,6 +173,117 @@ async def execute_second() -> list[Content]: assert lifecycle.get(second.identity).status is ApprovalStatus.SETTLED +def test_claim_can_be_released_before_execution_only_with_explicit_safe_policy() -> None: + """Reserved authority can be reclaimed when the owner proves execution never began.""" + lifecycle = ApprovalLifecycle() + occurrence = lifecycle.register_local( + thread_id="thread-1", + interrupt_id="approval-1", + call_id="call-1", + name="write_record", + arguments="{}", + ) + decision = ResumeDecision(interrupt_id="approval-1", accepted=True, arguments="{}") + intent = lifecycle.claim(thread_id="thread-1", decision=decision) + + lifecycle.release_claim(intent, policy=ClaimRecoveryPolicy.SAFE_TO_RETRY) + retry = lifecycle.claim(thread_id="thread-1", decision=decision) + + assert retry.identity == occurrence.identity + assert lifecycle.get(occurrence.identity).status is ApprovalStatus.CLAIMED + + +def test_recovering_execution_without_a_result_becomes_indeterminate() -> None: + """Recovery preserves an uncertain occurrence instead of granting authority again.""" + lifecycle = ApprovalLifecycle() + occurrence = lifecycle.register_local( + thread_id="thread-1", + interrupt_id="approval-1", + call_id="call-1", + name="write_record", + arguments="{}", + ) + decision = ResumeDecision(interrupt_id="approval-1", accepted=True, arguments="{}") + intent = lifecycle.claim(thread_id="thread-1", decision=decision) + lifecycle.begin_execution(intent, owner=ApprovalExecutionOwner.LOCAL) + + recovered = lifecycle.recover_execution(intent, owner=ApprovalExecutionOwner.LOCAL) + + assert recovered is None + assert lifecycle.get(occurrence.identity).status is ApprovalStatus.INDETERMINATE + with pytest.raises(ApprovalIndeterminateError): + lifecycle.claim_batch(thread_id="thread-1", decisions=[decision]) + + +async def test_explicit_idempotency_key_allows_retry_after_execution_interruption() -> None: + """A predeclared idempotency key permits retrying a potentially started side effect.""" + lifecycle = ApprovalLifecycle() + occurrence = lifecycle.register_local( + thread_id="thread-1", + interrupt_id="approval-1", + call_id="call-1", + name="write_record", + arguments="{}", + idempotency_key="operation-1", + ) + intent = lifecycle.claim( + thread_id="thread-1", + decision=ResumeDecision(interrupt_id="approval-1", accepted=True, arguments="{}"), + ) + invocation_count = 0 + + async def execute_idempotently() -> list[Content]: + nonlocal invocation_count + invocation_count += 1 + if invocation_count == 1: + raise RuntimeError("connection lost") + return [Content.from_function_result(call_id="call-1", result="recorded")] + + owner = LocalPendingToolTransitionOwner(execute_idempotently) + with pytest.raises(RuntimeError, match="connection lost"): + await owner.execute(intent, lifecycle=lifecycle) + + assert lifecycle.get(occurrence.identity).status is ApprovalStatus.CLAIMED + outcome = await owner.execute(intent, lifecycle=lifecycle) + assert lifecycle.get(occurrence.identity).status is ApprovalStatus.SETTLED + assert outcome.replayable_results[0].content.result == "recorded" + assert invocation_count == 2 + + +async def test_hosted_idempotency_key_allows_retry_after_forwarding_interruption() -> None: + """A hosted owner uses the same explicit recovery rule as the local owner.""" + lifecycle = ApprovalLifecycle() + occurrence = lifecycle.register_hosted( + thread_id="thread-1", + interrupt_id="approval-1", + call_id="call-1", + name="hosted_write", + arguments="{}", + idempotency_key="hosted-operation-1", + ) + intent = lifecycle.claim( + thread_id="thread-1", + decision=ResumeDecision(interrupt_id="approval-1", accepted=True, arguments="{}"), + ) + forwarding_count = 0 + + async def forward_idempotently() -> list[Content]: + nonlocal forwarding_count + forwarding_count += 1 + if forwarding_count == 1: + raise RuntimeError("host disconnected") + return [Content.from_function_result(call_id="call-1", result="recorded")] + + owner = HostedPendingToolTransitionOwner(forward_idempotently) + with pytest.raises(RuntimeError, match="host disconnected"): + await owner.execute(intent, lifecycle=lifecycle) + + assert lifecycle.get(occurrence.identity).status is ApprovalStatus.CLAIMED + await owner.execute(intent, lifecycle=lifecycle) + assert lifecycle.get(occurrence.identity).status is ApprovalStatus.SETTLED + assert forwarding_count == 2 + + def test_batch_validation_is_atomic_before_claiming_any_occurrence() -> None: """One invalid decision leaves every occurrence pending and eligible for a corrected batch.""" lifecycle = ApprovalLifecycle() diff --git a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py index 6ab195608f..8785b3a055 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -47,6 +47,7 @@ add_agent_framework_fastapi_endpoint, ) from agent_framework_ag_ui._agent import AgentFrameworkAgent +from agent_framework_ag_ui._approval_lifecycle import ApprovalLifecycle from agent_framework_ag_ui._workflow import AgentFrameworkWorkflow @@ -2170,6 +2171,51 @@ async def test_endpoint_agent_approval_replayed_resume_entry_reprojects_retained assert result_events[0]["content"] == "Sunny in Seattle" +async def test_endpoint_agent_approval_settlement_failure_prevents_automatic_reexecution(monkeypatch): + """A lost settlement becomes indeterminate and an identical resume cannot execute again.""" + client, _, executed_cities = _build_weather_approval_endpoint() + original_settle = ApprovalLifecycle.settle + + def fail_settlement(self, intent, results): + del self, intent, results + raise RuntimeError("settlement unavailable") + + monkeypatch.setattr(ApprovalLifecycle, "settle", fail_settlement) + first_response = client.post( + "/approval", + json={ + "runId": "run-resume", + "threadId": "thread-weather", + "messages": [], + "resume": [{"interruptId": "call_get_weather", "status": "resolved", "payload": {"accepted": True}}], + }, + ) + monkeypatch.setattr(ApprovalLifecycle, "settle", original_settle) + + assert first_response.status_code == 200 + assert executed_cities == ["Seattle"] + assert [event for event in _decode_sse_events(first_response) if event.get("type") == "RUN_ERROR"] + + retry_response = client.post( + "/approval", + json={ + "runId": "run-retry", + "threadId": "thread-weather", + "messages": [], + "resume": [{"interruptId": "call_get_weather", "status": "resolved", "payload": {"accepted": True}}], + }, + ) + + retry_events = _decode_sse_events(retry_response) + run_errors = [event for event in retry_events if event.get("type") == "RUN_ERROR"] + assert retry_response.status_code == 200 + assert executed_cities == ["Seattle"] + assert len(run_errors) == 1 + assert run_errors[0]["code"] == "APPROVAL_RESUME_INVALID" + assert "indeterminate" in run_errors[0]["message"] + assert not [event for event in retry_events if event.get("type") == "TOOL_CALL_RESULT"] + + async def test_endpoint_agent_approval_resume_wrong_thread_emits_run_error(): """A valid approval id on a different AG-UI thread cannot execute the pending tool.""" client, agent, executed_cities = _build_weather_approval_endpoint() From d7fcd066ca4101fac800b5fba29a5c67e112ea5d Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Mon, 10 Aug 2026 14:26:34 +0900 Subject: [PATCH 06/13] Reconcile approval snapshots with lifecycle state Key decisions: - Keep Approval State authoritative and emit typed snapshot reconciliation keyed by logical occurrence identity. - Retire settled, rejected, cancelled, expired, indeterminate, and missing controls while preserving nonterminal authority. - Reconcile stale snapshots before hydration or resume, and retain lifecycle deduplication when snapshot saves fail. Files changed: - packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py - packages/ag-ui/agent_framework_ag_ui/_agent_run.py - packages/ag-ui/agent_framework_ag_ui/_snapshot_session.py - packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py - packages/ag-ui/tests/ag_ui/test_endpoint.py Verification: - 975 AG-UI tests passed with 92% package coverage and 89% approval lifecycle coverage. - Package-local Ruff and Pyright passed; git diff --check passed. Notes for next iteration: - The function-calling-loop scenario mapping remains inaccessible under the organization content-exclusion policy. - Workspace Poe fan-out remains blocked by the pre-existing missing packages/durabletask/pyproject.toml; equivalent package-local checks passed. --- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 50 ++++++- .../_approval_lifecycle.py | 119 +++++++++++++++-- .../_snapshot_session.py | 10 ++ .../tests/ag_ui/test_approval_lifecycle.py | 48 ++++++- .../ag-ui/tests/ag_ui/test_endpoint.py | 122 +++++++++++++++++- 5 files changed, 331 insertions(+), 18 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index 044848b134..f9c819f8fb 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py @@ -58,6 +58,7 @@ from ._approval_lifecycle import ( ApprovalExecutionOwner, ApprovalLifecycle, + ApprovalSnapshotReconciliation, AuthorizedExecution, DeferredPendingToolTransitionOwner, ForwardedPendingToolTransitionOwner, @@ -1288,6 +1289,7 @@ def _canonical_approval_resume_messages( lifecycle: ApprovalLifecycle | None = None, authorized_executions: dict[str, AuthorizedExecution] | None = None, retained_results: list[Content] | None = None, + snapshot_reconciliations: list[ApprovalSnapshotReconciliation] | None = None, ) -> tuple[list[dict[str, Any]], set[str], set[str], RunErrorEvent | None]: """Translate canonical ResumeEntry approvals into existing approval response messages.""" expected_ids = set(expected_interrupt_ids or set()) @@ -1391,6 +1393,8 @@ def _canonical_approval_resume_messages( else: for outcome in batch.retained_outcomes: retained_results.extend(result.content for result in outcome.replayable_results) + if snapshot_reconciliations is not None: + snapshot_reconciliations.extend(batch.snapshot_reconciliations) handled_ids.update(decision.interrupt_id for decision in decisions) return [], handled_ids, cancelled_ids, None return ( @@ -1461,7 +1465,12 @@ def _canonical_approval_resume_messages( interrupt_id for interrupt_id in cancelled_ids if entries_by_interrupt_id.get(interrupt_id) is not None ] try: - lifecycle.cancel_batch(thread_id=thread_id, interrupt_ids=lifecycle_cancelled_ids) + reconciliations = lifecycle.cancel_batch( + thread_id=thread_id, + interrupt_ids=lifecycle_cancelled_ids, + ) + if snapshot_reconciliations is not None: + snapshot_reconciliations.extend(reconciliations) except (KeyError, ValueError) as exc: return ( [], @@ -1654,6 +1663,8 @@ def _canonical_approval_resume_messages( if lifecycle is not None and authorized_executions is not None: try: intents = lifecycle.claim_batch(thread_id=thread_id, decisions=lifecycle_decisions) + if snapshot_reconciliations is not None: + snapshot_reconciliations.extend(intents.snapshot_reconciliations) for intent in intents: authorized_executions[intent.identity.call_id] = intent except (KeyError, ValueError) as exc: @@ -2467,16 +2478,34 @@ async def run_agent_stream( scope=snapshot_scope, thread_id=thread_id, ) - if snapshot_session.enabled and not raw_messages and resume_payload is None: - async for event in snapshot_session.hydrate_events(run_id=run_id): - yield event - return stored_snapshot = snapshot_session.stored stored_pending_approval_interrupt_ids: set[str] = set() seeded_resume_from_snapshot = False if stored_snapshot is not None: stored_pending_approval_interrupt_ids = _stored_pending_approval_interrupt_ids(stored_snapshot.interrupt) + if approval_state_store is not None and stored_pending_approval_interrupt_ids: + reconciliations = approval_state_store.lifecycle.reconcile_snapshot( + thread_id=approval_thread_id, + interrupt_ids=list(stored_pending_approval_interrupt_ids), + ) + retired_interrupt_ids = { + reconciliation.identity.interrupt_id + if reconciliation.identity is not None + else reconciliation.interrupt_id + for reconciliation in reconciliations + if reconciliation.retire_interrupt + } + if retired_interrupt_ids: + await snapshot_session.clear_interrupts(interrupt_ids=retired_interrupt_ids) + stored_snapshot = snapshot_session.stored + stored_pending_approval_interrupt_ids.difference_update(retired_interrupt_ids) + if snapshot_session.enabled and not raw_messages and resume_payload is None: + async for event in snapshot_session.hydrate_events(run_id=run_id): + yield event + return + + if stored_snapshot is not None: if resume_payload is not None and stored_pending_approval_interrupt_ids: raw_messages = snapshot_session.resume_seeded_messages(raw_messages) seeded_resume_from_snapshot = True @@ -2516,6 +2545,7 @@ async def run_agent_stream( authorized_executions: dict[str, AuthorizedExecution] = {} forwarded_executions: dict[str, tuple[ForwardedPendingToolTransitionOwner, AuthorizedExecution, Content]] = {} retained_approval_results: list[Content] = [] + approval_snapshot_reconciliations: list[ApprovalSnapshotReconciliation] = [] approval_resume_messages, handled_resume_ids, cancelled_resume_ids, resume_error = ( _canonical_approval_resume_messages( resume_payload, @@ -2525,6 +2555,7 @@ async def run_agent_stream( lifecycle=approval_state_store.lifecycle if approval_state_store is not None else None, authorized_executions=authorized_executions, retained_results=retained_approval_results, + snapshot_reconciliations=approval_snapshot_reconciliations, ) ) if resume_error is not None: @@ -2539,7 +2570,14 @@ async def run_agent_stream( if should_clear_tool_approval_state: _clear_tool_approval_state(approval_state_store, approval_thread_id) if resume_error_code == "APPROVAL_RESUME_CANCELLED": - await snapshot_session.clear_interrupts(interrupt_ids=cancelled_resume_ids or None) + retired_interrupt_ids = { + reconciliation.identity.interrupt_id + if reconciliation.identity is not None + else reconciliation.interrupt_id + for reconciliation in approval_snapshot_reconciliations + if reconciliation.retire_interrupt + } + await snapshot_session.clear_interrupts(interrupt_ids=retired_interrupt_ids or cancelled_resume_ids or None) yield resume_error return resume_messages = _resume_to_tool_messages(resume_payload, exclude_interrupt_ids=handled_resume_ids) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py b/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py index e3957390a4..9167a8448d 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py @@ -30,6 +30,20 @@ class ApprovalStatus(str, Enum): INDETERMINATE = "indeterminate" +class ApprovalSnapshotStatus(str, Enum): + """Approval authority status projected during snapshot reconciliation.""" + + PENDING = "pending" + CLAIMED = "claimed" + EXECUTING = "executing" + SETTLED = "settled" + REJECTED = "rejected" + CANCELLED = "cancelled" + EXPIRED = "expired" + INDETERMINATE = "indeterminate" + MISSING = "missing" + + class ApprovalExecutionOwner(str, Enum): """Runtime owner authorized to continue an approved occurrence.""" @@ -55,6 +69,16 @@ class ApprovalOccurrenceIdentity: call_id: str +@dataclass(frozen=True) +class ApprovalSnapshotReconciliation: + """Semantic lifecycle result used to retire or retain one snapshot control.""" + + interrupt_id: str + identity: ApprovalOccurrenceIdentity | None + status: ApprovalSnapshotStatus + retire_interrupt: bool + + @dataclass(frozen=True) class ResumeDecision: """Canonical client decision presented to the approval lifecycle.""" @@ -107,6 +131,7 @@ class ApprovalOutcome: identity: ApprovalOccurrenceIdentity replayable_results: tuple[ReplayableToolResult, ...] result_group: tuple[Content, ...] + snapshot_reconciliation: ApprovalSnapshotReconciliation @dataclass(frozen=True) @@ -115,6 +140,7 @@ class ApprovalBatchDecision: authorized_executions: tuple[AuthorizedExecution, ...] retained_outcomes: tuple[ApprovalOutcome, ...] = () + snapshot_reconciliations: tuple[ApprovalSnapshotReconciliation, ...] = () def __iter__(self) -> Iterator[AuthorizedExecution]: """Iterate newly authorized executions for compatibility with existing callers.""" @@ -359,6 +385,48 @@ def decision_context(self, *, thread_id: str, interrupt_id: str) -> tuple[str, s occurrence = self._occurrences[identity] return occurrence.name, occurrence.arguments + def reconcile_snapshot( + self, + *, + thread_id: str, + interrupt_ids: list[str], + ) -> tuple[ApprovalSnapshotReconciliation, ...]: + """Describe which stored approval controls remain actionable.""" + reconciliations: list[ApprovalSnapshotReconciliation] = [] + for interrupt_id in interrupt_ids: + key = (thread_id, interrupt_id) + identity = self._pending_by_interrupt.get(key) or self._terminal_by_interrupt.get(key) + if identity is None: + reconciliations.append( + ApprovalSnapshotReconciliation( + interrupt_id=interrupt_id, + identity=None, + status=ApprovalSnapshotStatus.MISSING, + retire_interrupt=True, + ) + ) + continue + occurrence = self._occurrences[identity] + reconciliations.append(self._snapshot_reconciliation(occurrence)) + return tuple(reconciliations) + + @staticmethod + def _snapshot_reconciliation(occurrence: ApprovalOccurrence) -> ApprovalSnapshotReconciliation: + status = ApprovalSnapshotStatus(occurrence.status.value) + terminal_statuses = { + ApprovalSnapshotStatus.SETTLED, + ApprovalSnapshotStatus.REJECTED, + ApprovalSnapshotStatus.CANCELLED, + ApprovalSnapshotStatus.EXPIRED, + ApprovalSnapshotStatus.INDETERMINATE, + } + return ApprovalSnapshotReconciliation( + interrupt_id=occurrence.identity.interrupt_id, + identity=occurrence.identity, + status=status, + retire_interrupt=status in terminal_statuses, + ) + def claim(self, *, thread_id: str, decision: ResumeDecision) -> AuthorizedExecution: """Validate and reserve one accepted decision before execution.""" if not decision.accepted: @@ -422,11 +490,13 @@ def claim_batch( resolved.append((decision, occurrence, False)) intents: list[AuthorizedExecution] = [] retained_outcomes: list[ApprovalOutcome] = [] + snapshot_reconciliations: list[ApprovalSnapshotReconciliation] = [] for decision, occurrence, is_terminal in resolved: if is_terminal: if occurrence.outcome is None: raise RuntimeError("Validated terminal approval is missing its retained outcome.") retained_outcomes.append(occurrence.outcome) + snapshot_reconciliations.append(occurrence.outcome.snapshot_reconciliation) continue occurrence.decision = decision if not decision.accepted: @@ -435,13 +505,15 @@ def claim_batch( result="Error: Tool call invocation was rejected by user.", ) occurrence.replayable_results = [ReplayableToolResult(content=result)] + occurrence.status = ApprovalStatus.REJECTED + self._remove_pending_aliases(occurrence) occurrence.outcome = ApprovalOutcome( identity=occurrence.identity, replayable_results=tuple(occurrence.replayable_results), result_group=(result,), + snapshot_reconciliation=self._snapshot_reconciliation(occurrence), ) - occurrence.status = ApprovalStatus.REJECTED - self._remove_pending_aliases(occurrence) + snapshot_reconciliations.append(occurrence.outcome.snapshot_reconciliation) continue if decision.arguments is None: raise RuntimeError("Validated pending approval is missing canonical arguments.") @@ -461,11 +533,18 @@ def claim_batch( return ApprovalBatchDecision( authorized_executions=tuple(intents), retained_outcomes=tuple(retained_outcomes), + snapshot_reconciliations=tuple(snapshot_reconciliations), ) - def cancel_batch(self, *, thread_id: str, interrupt_ids: list[str]) -> None: + def cancel_batch( + self, + *, + thread_id: str, + interrupt_ids: list[str], + ) -> tuple[ApprovalSnapshotReconciliation, ...]: """Validate and cancel selected occurrences without changing their siblings.""" occurrences: list[ApprovalOccurrence] = [] + reconciliations: list[ApprovalSnapshotReconciliation] = [] seen_interrupt_ids: set[str] = set() for interrupt_id in interrupt_ids: if interrupt_id in seen_interrupt_ids: @@ -477,6 +556,7 @@ def cancel_batch(self, *, thread_id: str, interrupt_ids: list[str]) -> None: identity = self._terminal_by_interrupt[key] terminal = self._occurrences[identity] if terminal.status is ApprovalStatus.CANCELLED: + reconciliations.append(self._snapshot_reconciliation(terminal)) continue raise ValueError("Approval cancellation conflicts with the retained terminal decision.") occurrence = self._occurrences[identity] @@ -487,8 +567,15 @@ def cancel_batch(self, *, thread_id: str, interrupt_ids: list[str]) -> None: for occurrence in occurrences: occurrence.status = ApprovalStatus.CANCELLED self._remove_pending_aliases(occurrence) + reconciliations.append(self._snapshot_reconciliation(occurrence)) + return tuple(reconciliations) - def expire_batch(self, *, thread_id: str, interrupt_ids: list[str]) -> None: + def expire_batch( + self, + *, + thread_id: str, + interrupt_ids: list[str], + ) -> tuple[ApprovalSnapshotReconciliation, ...]: """Expire pending authority without permitting later execution.""" occurrences: list[ApprovalOccurrence] = [] seen_interrupt_ids: set[str] = set() @@ -505,6 +592,7 @@ def expire_batch(self, *, thread_id: str, interrupt_ids: list[str]) -> None: for occurrence in occurrences: occurrence.status = ApprovalStatus.EXPIRED self._remove_pending_aliases(occurrence) + return tuple(self._snapshot_reconciliation(occurrence) for occurrence in occurrences) def _remove_pending_aliases(self, occurrence: ApprovalOccurrence) -> None: for thread_id in occurrence.thread_ids: @@ -535,7 +623,12 @@ def release_claim(self, intent: AuthorizedExecution, *, policy: ClaimRecoveryPol raise ValueError(f"Approval occurrence is not claimed: {occurrence.status}.") occurrence.status = ApprovalStatus.PENDING - def mark_indeterminate(self, intent: AuthorizedExecution, *, owner: ApprovalExecutionOwner) -> None: + def mark_indeterminate( + self, + intent: AuthorizedExecution, + *, + owner: ApprovalExecutionOwner, + ) -> ApprovalSnapshotReconciliation: """Record that execution may have begun but no result was settled.""" occurrence = self._occurrences[intent.identity] if intent.owner is not owner or occurrence.owner is not owner: @@ -546,6 +639,7 @@ def mark_indeterminate(self, intent: AuthorizedExecution, *, owner: ApprovalExec raise ValueError(f"Approval occurrence is not executing: {occurrence.status}.") occurrence.status = ApprovalStatus.INDETERMINATE self._remove_pending_aliases(occurrence) + return self._snapshot_reconciliation(occurrence) def recover_execution( self, @@ -588,6 +682,7 @@ def settle(self, intent: AuthorizedExecution, results: list[Content]) -> Approva identity=occurrence.identity, replayable_results=tuple(replayable_results), result_group=tuple(results), + snapshot_reconciliation=self._snapshot_reconciliation(occurrence), ) occurrence.outcome = outcome return outcome @@ -624,15 +719,16 @@ def settle_forwarded( ] if len(replayable_results) + len(forwarded_responses) != 1: raise ValueError("A hosted approval must record exactly one outcome for its original call.") + occurrence.replayable_results = replayable_results + occurrence.status = ApprovalStatus.SETTLED + self._remove_pending_aliases(occurrence) outcome = ApprovalOutcome( identity=occurrence.identity, replayable_results=tuple(replayable_results), result_group=tuple(results), + snapshot_reconciliation=self._snapshot_reconciliation(occurrence), ) - occurrence.replayable_results = replayable_results - occurrence.status = ApprovalStatus.SETTLED occurrence.outcome = outcome - self._remove_pending_aliases(occurrence) return outcome def defer(self, intent: AuthorizedExecution, results: list[Content]) -> ApprovalOutcome: @@ -643,7 +739,12 @@ def defer(self, intent: AuthorizedExecution, results: list[Content]) -> Approval if occurrence.status is not ApprovalStatus.EXECUTING: raise ValueError(f"Approval occurrence is not executing: {occurrence.status}.") occurrence.status = ApprovalStatus.PENDING - return ApprovalOutcome(identity=occurrence.identity, replayable_results=(), result_group=tuple(results)) + return ApprovalOutcome( + identity=occurrence.identity, + replayable_results=(), + result_group=tuple(results), + snapshot_reconciliation=self._snapshot_reconciliation(occurrence), + ) class LocalPendingToolTransitionOwner: diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_snapshot_session.py b/python/packages/ag-ui/agent_framework_ag_ui/_snapshot_session.py index ea291853b0..b0694596e7 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_snapshot_session.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_snapshot_session.py @@ -181,6 +181,16 @@ async def clear_interrupts(self, *, interrupt_ids: set[str] | None = None) -> No Clears all interrupts when ``interrupt_ids`` is omitted. Failures are logged and swallowed for the same reason as :meth:`save`. """ + if self._stored is not None and self._stored.interrupt is not None: + if interrupt_ids is None: + self._stored.interrupt = None + else: + remaining_interrupts = [ + interrupt + for interrupt in self._stored.interrupt + if str(interrupt.get("id") or interrupt.get("interruptId")) not in interrupt_ids + ] + self._stored.interrupt = remaining_interrupts or None if self._store is None or self._scope is None: return await _clear_thread_snapshot_interrupt( diff --git a/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py b/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py index 7f2a83e80e..7cd68b6a62 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py +++ b/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py @@ -11,6 +11,7 @@ ApprovalExecutionOwner, ApprovalIndeterminateError, ApprovalLifecycle, + ApprovalSnapshotStatus, ApprovalStatus, ClaimRecoveryPolicy, HostedPendingToolTransitionOwner, @@ -117,6 +118,8 @@ async def forward_to_hosted_owner() -> list[Content]: assert outcome.identity == occurrence.identity assert outcome.result_group == (remote_result,) assert [result.content for result in outcome.replayable_results] == [remote_result] + assert outcome.snapshot_reconciliation.status is ApprovalSnapshotStatus.SETTLED + assert outcome.snapshot_reconciliation.retire_interrupt is True assert lifecycle.get(occurrence.identity).status is ApprovalStatus.SETTLED @@ -419,13 +422,56 @@ def test_batch_cancellation_preserves_each_original_occurrence() -> None: arguments='{"value":"second"}', ) - lifecycle.cancel_batch( + reconciliations = lifecycle.cancel_batch( thread_id="tenant-a\x1fthread-1", interrupt_ids=["approval-1"], ) assert lifecycle.get(cancelled.identity).status is ApprovalStatus.CANCELLED assert lifecycle.get(pending.identity).status is ApprovalStatus.PENDING + assert [(item.identity, item.status, item.retire_interrupt) for item in reconciliations] == [ + (cancelled.identity, ApprovalSnapshotStatus.CANCELLED, True) + ] + + +def test_snapshot_reconciliation_reports_terminal_pending_and_missing_occurrences() -> None: + """Snapshot projection receives lifecycle semantics without recreating authority.""" + lifecycle = ApprovalLifecycle() + settled = lifecycle.register_local( + thread_id="thread-1", + interrupt_id="approval-settled", + call_id="call-settled", + name="write_record", + arguments="{}", + ) + pending = lifecycle.register_local( + thread_id="thread-1", + interrupt_id="approval-pending", + call_id="call-pending", + name="write_record", + arguments="{}", + ) + intent = lifecycle.claim( + thread_id="thread-1", + decision=ResumeDecision(interrupt_id="approval-settled", accepted=True, arguments="{}"), + ) + lifecycle.begin_execution(intent, owner=ApprovalExecutionOwner.LOCAL) + outcome = lifecycle.settle( + intent, + [Content.from_function_result(call_id="call-settled", result="done")], + ) + + reconciliations = lifecycle.reconcile_snapshot( + thread_id="thread-1", + interrupt_ids=["approval-settled", "approval-pending", "approval-missing"], + ) + + assert outcome.snapshot_reconciliation == reconciliations[0] + assert [(item.identity, item.status, item.retire_interrupt) for item in reconciliations] == [ + (settled.identity, ApprovalSnapshotStatus.SETTLED, True), + (pending.identity, ApprovalSnapshotStatus.PENDING, False), + (None, ApprovalSnapshotStatus.MISSING, True), + ] def test_one_occurrence_can_be_claimed_through_a_trusted_thread_alias() -> None: diff --git a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py index 8785b3a055..a10599456c 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -1432,7 +1432,10 @@ async def test_endpoint_agent_approval_pause_emits_canonical_interrupt_outcome() } -def _build_weather_approval_endpoint() -> tuple[TestClient, StubAgent, list[str]]: +def _build_weather_approval_endpoint( + *, + snapshot_store: InMemoryAGUIThreadSnapshotStore | None = None, +) -> tuple[TestClient, StubAgent, list[str]]: executed_cities: list[str] = [] def get_weather(city: str) -> str: @@ -1460,7 +1463,13 @@ def get_weather(city: str) -> str: ) wrapped_agent = AgentFrameworkAgent(agent=agent, require_confirmation=False) app = FastAPI() - add_agent_framework_fastapi_endpoint(app, wrapped_agent, path="/approval") + add_agent_framework_fastapi_endpoint( + app, + wrapped_agent, + path="/approval", + snapshot_store=snapshot_store, + snapshot_scope_resolver=(lambda _request: "tenant-a") if snapshot_store is not None else None, + ) client = TestClient(app) pause_response = client.post( @@ -4968,6 +4977,85 @@ def get_weather(city: str) -> str: assert "outcome" not in hydrate_events[-1] +async def test_agent_endpoint_stale_approval_snapshot_cannot_recreate_missing_authority(): + """A snapshot from a prior process cannot advertise approval authority the new process does not own.""" + executed_cities: list[str] = [] + + def get_weather(city: str) -> str: + executed_cities.append(city) + return f"Sunny in {city}" + + weather_tool = FunctionTool( + name="get_weather", + description="Get the weather for a city", + func=get_weather, + approval_mode="always_require", + ) + approval_request = Content.from_function_approval_request( + id="call_get_weather", + function_call=Content.from_function_call( + call_id="call_get_weather", + name="get_weather", + arguments={"city": "Seattle"}, + ), + ) + agent = StubAgent( + updates=[AgentResponseUpdate(contents=[approval_request], role="assistant")], + default_options={"tools": [weather_tool]}, + ) + store = InMemoryAGUIThreadSnapshotStore() + first_app = FastAPI() + add_agent_framework_fastapi_endpoint( + first_app, + AgentFrameworkAgent(agent=agent, require_confirmation=False), + path="/approval-snapshots", + snapshot_store=store, + snapshot_scope_resolver=lambda _request: "tenant-a", + ) + pause_response = TestClient(first_app).post( + "/approval-snapshots", + json={ + "thread_id": "agent-approval-thread", + "messages": [{"role": "user", "content": "What is the weather?"}], + }, + ) + assert pause_response.status_code == 200 + assert _run_finished_interrupts(_decode_sse_events(pause_response)[-1]) + + restarted_app = FastAPI() + add_agent_framework_fastapi_endpoint( + restarted_app, + AgentFrameworkAgent(agent=agent, require_confirmation=False), + path="/approval-snapshots", + snapshot_store=store, + snapshot_scope_resolver=lambda _request: "tenant-a", + ) + restarted_client = TestClient(restarted_app) + stale_resume = restarted_client.post( + "/approval-snapshots", + json={ + "runId": "run-stale-resume", + "thread_id": "agent-approval-thread", + "messages": [], + "resume": [{"interruptId": "call_get_weather", "status": "resolved", "payload": {"accepted": True}}], + }, + ) + stale_events = _decode_sse_events(stale_resume) + assert [event["code"] for event in stale_events if event.get("type") == "RUN_ERROR"] == [ + "APPROVAL_RESUME_NOT_FOUND" + ] + assert executed_cities == [] + + hydrate_response = restarted_client.post( + "/approval-snapshots", + json={"thread_id": "agent-approval-thread", "messages": []}, + ) + + assert hydrate_response.status_code == 200 + hydrate_events = _decode_sse_events(hydrate_response) + assert "outcome" not in hydrate_events[-1] + + async def test_agent_endpoint_ignores_forged_suffix_messages(streaming_chat_client_stub): """Client-forged assistant/tool messages after the stored prefix never become history.""" app = FastAPI() @@ -5170,6 +5258,36 @@ async def save(self, *, scope: str, thread_id: str, snapshot: Any) -> None: await super().save(scope=scope, thread_id=thread_id, snapshot=snapshot) +async def test_agent_endpoint_approval_snapshot_save_failure_does_not_duplicate_execution(): + """A stale interrupt left by a failed save is retired from terminal Approval State before a retry.""" + store = _FailNextSaveStore() + client, _, executed_cities = _build_weather_approval_endpoint(snapshot_store=store) + store.fail_next_save = True + + resume_payload = { + "threadId": "thread-weather", + "messages": [], + "resume": [{"interruptId": "call_get_weather", "status": "resolved", "payload": {"accepted": True}}], + } + first_resume = client.post("/approval", json={"runId": "run-resume", **resume_payload}) + retry = client.post("/approval", json={"runId": "run-retry", **resume_payload}) + + assert first_resume.status_code == 200 + assert retry.status_code == 200 + assert executed_cities == ["Seattle"] + retry_events = _decode_sse_events(retry) + assert not [event for event in retry_events if event.get("type") == "RUN_ERROR"] + assert [ + (event["toolCallId"], event["content"]) for event in retry_events if event.get("type") == "TOOL_CALL_RESULT" + ] == [("call_get_weather", "Sunny in Seattle")] + + hydrate_response = client.post( + "/approval", + json={"runId": "run-hydrate", "threadId": "thread-weather", "messages": []}, + ) + assert "outcome" not in _decode_sse_events(hydrate_response)[-1] + + async def test_agent_endpoint_snapshot_save_failure_does_not_fail_run(streaming_chat_client_stub): """A failing snapshot save must not turn a completed agent run into RUN_ERROR.""" app = FastAPI() From 42b3d05b47efaaecb519a776c381ef79b9f9594c Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Mon, 10 Aug 2026 14:41:53 +0900 Subject: [PATCH 07/13] Bound process-local approval lifecycle state Key decisions: - Protect pending, claimed, executing, and indeterminate occurrences from eviction while retaining terminal outcomes for a configurable 15-minute process-local deduplication window. - Serialize complete approval batches by logical occurrence locks so aliases share atomic decisions and independent batches can progress concurrently. - Fail capacity, claim, and settlement conflicts explicitly, and emit redacted structured lifecycle telemetry without tool names, arguments, or approval payloads. - Remove legacy LRU eviction paths so active Approval State and middleware state are never silently discarded. Files changed: - packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py - packages/ag-ui/agent_framework_ag_ui/_approval_state.py - packages/ag-ui/agent_framework_ag_ui/_agent_run.py - packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py - packages/ag-ui/tests/ag_ui/test_approval_state.py Verification: - 982 AG-UI tests passed with 92% package coverage and 91% approval lifecycle coverage. - 34 focused lifecycle and storage tests passed with RuntimeWarning and DeprecationWarning treated as errors. - Package-local Ruff and Pyright passed; git diff --check passed. Notes for next iteration: - The function-calling-loop scenario mapping remains inaccessible under the organization content-exclusion policy. - Workspace Poe fan-out remains blocked by the pre-existing missing packages/durabletask/pyproject.toml; equivalent package-local checks passed. --- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 20 +- .../_approval_lifecycle.py | 202 ++++++++++++- .../agent_framework_ag_ui/_approval_state.py | 39 ++- .../tests/ag_ui/test_approval_lifecycle.py | 278 ++++++++++++++++++ .../ag-ui/tests/ag_ui/test_approval_state.py | 36 ++- 5 files changed, 522 insertions(+), 53 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index f9c819f8fb..c96444c334 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py @@ -8,7 +8,6 @@ import json import logging import uuid -from collections import OrderedDict from collections.abc import AsyncIterable, Awaitable, Mapping, Sequence from dataclasses import dataclass, field from functools import partial @@ -901,9 +900,7 @@ def _save_tool_approval_state( serialized_state = _serialized_tool_approval_state(raw_state) if serialized_state is None: return - approval_state_store.tool_approval_states[thread_id] = serialized_state - approval_state_store.tool_approval_states.move_to_end(thread_id) - approval_state_store.evict_oldest() + approval_state_store.set_tool_approval_state(thread_id, serialized_state) def _clear_tool_approval_state( @@ -1678,19 +1675,6 @@ def _canonical_approval_resume_messages( return messages, handled_ids, cancelled_ids, None -def _evict_oldest_approvals(registry: dict[PendingApprovalKey, PendingApprovalEntry], max_size: int = 10_000) -> None: - """Evict the oldest entries from the pending-approvals registry (LRU). - - Only effective when *registry* is an ``OrderedDict``; plain dicts are - left untouched because insertion-order eviction is unreliable for them. - """ - if len(registry) <= max_size or not isinstance(registry, OrderedDict): - return - while len(registry) > max_size: - oldest_key = next(iter(registry)) - _remove_pending_approval(registry, oldest_key) - - async def _resolve_approval_responses( messages: list[Any], tools: list[Any], @@ -2865,8 +2849,6 @@ async def run_agent_stream( already_approved_requests=already_approved_requests, server_label=server_label, ) - # Evict oldest entries if the registry exceeds a safe bound (LRU) - _evict_oldest_approvals(pending_approvals, max_size=10_000) else: logger.warning( "Approval request not registered: missing id=%s, function_call=%s, or function name", diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py b/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py index 9167a8448d..a48366669d 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py @@ -4,19 +4,90 @@ from __future__ import annotations +import logging from asyncio import CancelledError from collections.abc import Awaitable, Callable, Iterator +from contextlib import ExitStack from dataclasses import dataclass, field from enum import Enum +from functools import wraps +from threading import RLock +from time import monotonic +from typing import Any, TypeVar from uuid import uuid4 from agent_framework import Content +logger = logging.getLogger(__name__) +_ReturnT = TypeVar("_ReturnT") + + +def _serialized_registration(method: Callable[..., _ReturnT]) -> Callable[..., _ReturnT]: + @wraps(method) + def wrapper(self: ApprovalLifecycle, *args: Any, **kwargs: Any) -> _ReturnT: + with self._index_lock: + return method(self, *args, **kwargs) + + return wrapper + + +def _serialized_by_batch(method: Callable[..., _ReturnT]) -> Callable[..., _ReturnT]: + @wraps(method) + def wrapper(self: ApprovalLifecycle, *args: Any, **kwargs: Any) -> _ReturnT: + thread_id = kwargs.get("thread_id") + if not isinstance(thread_id, str): + raise TypeError("Serialized approval transitions require a thread_id.") + decisions = kwargs.get("decisions") + interrupt_ids = ( + [decision.interrupt_id for decision in decisions] + if isinstance(decisions, list) and all(isinstance(decision, ResumeDecision) for decision in decisions) + else kwargs.get("interrupt_ids") + ) + if not isinstance(interrupt_ids, list) or not all( + isinstance(interrupt_id, str) for interrupt_id in interrupt_ids + ): + raise TypeError("Serialized approval transitions require interrupt identities.") + locks = self._locks_for_batch(thread_id=thread_id, interrupt_ids=interrupt_ids) + with ExitStack() as stack: + for lock in locks: + stack.enter_context(lock) + try: + return method(self, *args, **kwargs) + except (KeyError, ValueError) as exc: + self._emit_event("authority_failure", failure_type=type(exc).__name__) + raise + + return wrapper + + +def _serialized_by_occurrence(method: Callable[..., _ReturnT]) -> Callable[..., _ReturnT]: + @wraps(method) + def wrapper(self: ApprovalLifecycle, *args: Any, **kwargs: Any) -> _ReturnT: + intent = args[0] if args else kwargs.get("intent") + if not isinstance(intent, AuthorizedExecution): + raise TypeError("Serialized approval transitions require an authorized execution.") + with self._lock_for_identity(intent.identity): + return method(self, *args, **kwargs) + + return wrapper + class ApprovalIndeterminateError(ValueError): """An approval may have executed but has no retained terminal outcome.""" +class ApprovalCapacityError(RuntimeError): + """Approval state capacity is exhausted by protected occurrences.""" + + +class ApprovalClaimConflictError(ValueError): + """Approval authority cannot be claimed in its current state.""" + + +class ApprovalSettlementConflictError(ValueError): + """An approval outcome cannot settle in its current state.""" + + class ApprovalStatus(str, Enum): """Lifecycle state of one server-owned approval occurrence.""" @@ -104,6 +175,7 @@ class ApprovalOccurrence: replayable_results: list[ReplayableToolResult] = field(default_factory=list) decision: ResumeDecision | None = None outcome: ApprovalOutcome | None = None + terminal_at: float | None = None @dataclass(frozen=True) @@ -148,9 +220,24 @@ def __iter__(self) -> Iterator[AuthorizedExecution]: class ApprovalLifecycle: - """Own registration, authority transitions, and settlement for approvals.""" + """Own process-local registration, authority transitions, and settlement for approvals.""" - def __init__(self) -> None: + def __init__( + self, + *, + max_entries: int = 10_000, + terminal_retention_seconds: float = 900, + clock: Callable[[], float] = monotonic, + ) -> None: + if max_entries < 1: + raise ValueError("max_entries must be greater than 0.") + if terminal_retention_seconds <= 0: + raise ValueError("terminal_retention_seconds must be greater than 0.") + self._max_entries = max_entries + self._terminal_retention_seconds = terminal_retention_seconds + self._clock = clock + self._index_lock = RLock() + self._locks_by_identity: dict[ApprovalOccurrenceIdentity, RLock] = {} self._occurrences: dict[ApprovalOccurrenceIdentity, ApprovalOccurrence] = {} self._pending_by_interrupt: dict[tuple[str, str], ApprovalOccurrenceIdentity] = {} self._terminal_by_interrupt: dict[tuple[str, str], ApprovalOccurrenceIdentity] = {} @@ -315,6 +402,7 @@ def register_unowned_aliases( owner=ApprovalExecutionOwner.UNAVAILABLE, ) + @_serialized_registration def _register_aliases( self, *, @@ -326,6 +414,7 @@ def _register_aliases( owner: ApprovalExecutionOwner, idempotency_key: str | None = None, ) -> ApprovalOccurrence: + self._purge_expired_terminal() if idempotency_key == "": raise ValueError("An execution idempotency key cannot be empty.") unique_thread_ids = tuple(dict.fromkeys(thread_ids)) @@ -341,7 +430,7 @@ def _register_aliases( if existing_identities: occurrence = self._occurrences[next(iter(existing_identities))] if occurrence.status is not ApprovalStatus.PENDING: - raise ValueError(f"Approval occurrence is not pending: {occurrence.status}.") + raise ApprovalClaimConflictError(f"Approval occurrence is not pending: {occurrence.status}.") if ( occurrence.identity.call_id != call_id or occurrence.name != name @@ -355,6 +444,9 @@ def _register_aliases( self._pending_by_interrupt[(thread_id, interrupt_id)] = occurrence.identity return occurrence + if len(self._occurrences) >= self._max_entries: + self._emit_event("capacity_failure") + raise ApprovalCapacityError("Approval state capacity is exhausted by protected occurrences.") identity = ApprovalOccurrenceIdentity( thread_id=unique_thread_ids[0], occurrence_id=str(uuid4()), @@ -372,14 +464,18 @@ def _register_aliases( self._occurrences[identity] = occurrence for thread_id in unique_thread_ids: self._pending_by_interrupt[(thread_id, interrupt_id)] = identity + self._locks_by_identity[identity] = RLock() + self._emit_event("registration", occurrence) return occurrence def get(self, identity: ApprovalOccurrenceIdentity) -> ApprovalOccurrence: """Return server-owned state for a registered occurrence.""" + self._purge_expired_terminal() return self._occurrences[identity] def decision_context(self, *, thread_id: str, interrupt_id: str) -> tuple[str, str]: """Return canonical server-owned call data needed to normalize a typed retry.""" + self._purge_expired_terminal() key = (thread_id, interrupt_id) identity = self._pending_by_interrupt.get(key) or self._terminal_by_interrupt[key] occurrence = self._occurrences[identity] @@ -392,6 +488,7 @@ def reconcile_snapshot( interrupt_ids: list[str], ) -> tuple[ApprovalSnapshotReconciliation, ...]: """Describe which stored approval controls remain actionable.""" + self._purge_expired_terminal() reconciliations: list[ApprovalSnapshotReconciliation] = [] for interrupt_id in interrupt_ids: key = (thread_id, interrupt_id) @@ -436,6 +533,7 @@ def claim(self, *, thread_id: str, decision: ResumeDecision) -> AuthorizedExecut raise ValueError("A duplicate settled decision cannot authorize execution again.") return batch.authorized_executions[0] + @_serialized_by_batch def claim_batch( self, *, @@ -443,6 +541,7 @@ def claim_batch( decisions: list[ResumeDecision], ) -> ApprovalBatchDecision: """Validate a complete decision batch before reserving accepted occurrences.""" + self._purge_expired_terminal() resolved: list[tuple[ResumeDecision, ApprovalOccurrence, bool]] = [] seen_interrupt_ids: set[str] = set() for decision in decisions: @@ -479,7 +578,7 @@ def claim_batch( resolved.append((decision, occurrence, True)) continue if occurrence.status is not ApprovalStatus.PENDING: - raise ValueError(f"Approval occurrence is not pending: {occurrence.status}.") + raise ApprovalClaimConflictError(f"Approval occurrence is not pending: {occurrence.status}.") if decision.name is not None and decision.name != occurrence.name: raise ValueError("Approval decision tool name does not match the registered occurrence.") canonical_arguments = decision.arguments @@ -497,6 +596,7 @@ def claim_batch( raise RuntimeError("Validated terminal approval is missing its retained outcome.") retained_outcomes.append(occurrence.outcome) snapshot_reconciliations.append(occurrence.outcome.snapshot_reconciliation) + self._emit_event("duplicate", occurrence) continue occurrence.decision = decision if not decision.accepted: @@ -514,6 +614,7 @@ def claim_batch( snapshot_reconciliation=self._snapshot_reconciliation(occurrence), ) snapshot_reconciliations.append(occurrence.outcome.snapshot_reconciliation) + self._emit_event("rejection", occurrence) continue if decision.arguments is None: raise RuntimeError("Validated pending approval is missing canonical arguments.") @@ -521,6 +622,7 @@ def claim_batch( if occurrence.owner is ApprovalExecutionOwner.UNAVAILABLE: continue occurrence.status = ApprovalStatus.CLAIMED + self._emit_event("claim", occurrence) intents.append( AuthorizedExecution( identity=occurrence.identity, @@ -536,6 +638,7 @@ def claim_batch( snapshot_reconciliations=tuple(snapshot_reconciliations), ) + @_serialized_by_batch def cancel_batch( self, *, @@ -543,6 +646,7 @@ def cancel_batch( interrupt_ids: list[str], ) -> tuple[ApprovalSnapshotReconciliation, ...]: """Validate and cancel selected occurrences without changing their siblings.""" + self._purge_expired_terminal() occurrences: list[ApprovalOccurrence] = [] reconciliations: list[ApprovalSnapshotReconciliation] = [] seen_interrupt_ids: set[str] = set() @@ -557,6 +661,7 @@ def cancel_batch( terminal = self._occurrences[identity] if terminal.status is ApprovalStatus.CANCELLED: reconciliations.append(self._snapshot_reconciliation(terminal)) + self._emit_event("duplicate", terminal) continue raise ValueError("Approval cancellation conflicts with the retained terminal decision.") occurrence = self._occurrences[identity] @@ -568,8 +673,10 @@ def cancel_batch( occurrence.status = ApprovalStatus.CANCELLED self._remove_pending_aliases(occurrence) reconciliations.append(self._snapshot_reconciliation(occurrence)) + self._emit_event("cancellation", occurrence) return tuple(reconciliations) + @_serialized_by_batch def expire_batch( self, *, @@ -577,6 +684,7 @@ def expire_batch( interrupt_ids: list[str], ) -> tuple[ApprovalSnapshotReconciliation, ...]: """Expire pending authority without permitting later execution.""" + self._purge_expired_terminal() occurrences: list[ApprovalOccurrence] = [] seen_interrupt_ids: set[str] = set() for interrupt_id in interrupt_ids: @@ -592,13 +700,73 @@ def expire_batch( for occurrence in occurrences: occurrence.status = ApprovalStatus.EXPIRED self._remove_pending_aliases(occurrence) + self._emit_event("expiration", occurrence) return tuple(self._snapshot_reconciliation(occurrence) for occurrence in occurrences) def _remove_pending_aliases(self, occurrence: ApprovalOccurrence) -> None: - for thread_id in occurrence.thread_ids: - self._pending_by_interrupt.pop((thread_id, occurrence.identity.interrupt_id), None) - self._terminal_by_interrupt[(thread_id, occurrence.identity.interrupt_id)] = occurrence.identity + if occurrence.status is not ApprovalStatus.INDETERMINATE: + occurrence.terminal_at = self._clock() + with self._index_lock: + for thread_id in occurrence.thread_ids: + self._pending_by_interrupt.pop((thread_id, occurrence.identity.interrupt_id), None) + self._terminal_by_interrupt[(thread_id, occurrence.identity.interrupt_id)] = occurrence.identity + + def _purge_expired_terminal(self) -> None: + with self._index_lock: + cutoff = self._clock() - self._terminal_retention_seconds + expired = [ + occurrence + for occurrence in self._occurrences.values() + if occurrence.terminal_at is not None and occurrence.terminal_at <= cutoff + ] + for occurrence in expired: + self._emit_event("retention_purge", occurrence) + self._occurrences.pop(occurrence.identity, None) + self._locks_by_identity.pop(occurrence.identity, None) + for thread_id in occurrence.thread_ids: + key = (thread_id, occurrence.identity.interrupt_id) + if self._terminal_by_interrupt.get(key) == occurrence.identity: + self._terminal_by_interrupt.pop(key, None) + + def _locks_for_batch(self, *, thread_id: str, interrupt_ids: list[str]) -> tuple[RLock, ...]: + with self._index_lock: + identities = { + identity + for interrupt_id in interrupt_ids + if ( + identity := self._pending_by_interrupt.get((thread_id, interrupt_id)) + or self._terminal_by_interrupt.get((thread_id, interrupt_id)) + ) + is not None + } + return tuple( + self._locks_by_identity[identity] + for identity in sorted(identities, key=lambda item: item.occurrence_id) + ) + + def _lock_for_identity(self, identity: ApprovalOccurrenceIdentity) -> RLock: + with self._index_lock: + return self._locks_by_identity.setdefault(identity, RLock()) + @staticmethod + def _emit_event( + event: str, + occurrence: ApprovalOccurrence | None = None, + *, + failure_type: str | None = None, + ) -> None: + extra: dict[str, str] = {"approval_event": event} + if occurrence is not None: + extra.update( + approval_occurrence_id=occurrence.identity.occurrence_id, + approval_status=occurrence.status.value, + approval_owner=occurrence.owner.value, + ) + if failure_type is not None: + extra["approval_failure_type"] = failure_type + logger.info("AG-UI approval lifecycle transition", extra=extra) + + @_serialized_by_occurrence def begin_execution(self, intent: AuthorizedExecution, *, owner: ApprovalExecutionOwner) -> None: """Mark that an external side effect may begin. @@ -611,9 +779,11 @@ def begin_execution(self, intent: AuthorizedExecution, *, owner: ApprovalExecuti f"Approval occurrence belongs to the {occurrence.owner.value} transition owner, not {owner.value}." ) if occurrence.status is not ApprovalStatus.CLAIMED: - raise ValueError(f"Approval occurrence is not claimed: {occurrence.status}.") + raise ApprovalClaimConflictError(f"Approval occurrence is not claimed: {occurrence.status}.") occurrence.status = ApprovalStatus.EXECUTING + self._emit_event("execution_start", occurrence) + @_serialized_by_occurrence def release_claim(self, intent: AuthorizedExecution, *, policy: ClaimRecoveryPolicy) -> None: """Release reserved authority when execution is known not to have begun.""" occurrence = self._occurrences[intent.identity] @@ -623,6 +793,7 @@ def release_claim(self, intent: AuthorizedExecution, *, policy: ClaimRecoveryPol raise ValueError(f"Approval occurrence is not claimed: {occurrence.status}.") occurrence.status = ApprovalStatus.PENDING + @_serialized_by_occurrence def mark_indeterminate( self, intent: AuthorizedExecution, @@ -636,11 +807,13 @@ def mark_indeterminate( f"Approval occurrence belongs to the {occurrence.owner.value} transition owner, not {owner.value}." ) if occurrence.status is not ApprovalStatus.EXECUTING: - raise ValueError(f"Approval occurrence is not executing: {occurrence.status}.") + raise ApprovalSettlementConflictError(f"Approval occurrence is not executing: {occurrence.status}.") occurrence.status = ApprovalStatus.INDETERMINATE self._remove_pending_aliases(occurrence) + self._emit_event("indeterminate_recovery", occurrence) return self._snapshot_reconciliation(occurrence) + @_serialized_by_occurrence def recover_execution( self, intent: AuthorizedExecution, @@ -654,20 +827,21 @@ def recover_execution( f"Approval occurrence belongs to the {occurrence.owner.value} transition owner, not {owner.value}." ) if occurrence.status is not ApprovalStatus.EXECUTING: - raise ValueError(f"Approval occurrence is not executing: {occurrence.status}.") + raise ApprovalSettlementConflictError(f"Approval occurrence is not executing: {occurrence.status}.") if intent.idempotency_key is not None and intent.idempotency_key == occurrence.idempotency_key: occurrence.status = ApprovalStatus.CLAIMED return intent self.mark_indeterminate(intent, owner=owner) return None + @_serialized_by_occurrence def settle(self, intent: AuthorizedExecution, results: list[Content]) -> ApprovalOutcome: """Settle an executing occurrence with results under its original call identity.""" occurrence = self._occurrences[intent.identity] if intent.owner is not ApprovalExecutionOwner.LOCAL or occurrence.owner is not ApprovalExecutionOwner.LOCAL: raise ValueError("Only the local transition owner can settle a local execution result.") if occurrence.status is not ApprovalStatus.EXECUTING: - raise ValueError(f"Approval occurrence is not executing: {occurrence.status}.") + raise ApprovalSettlementConflictError(f"Approval occurrence is not executing: {occurrence.status}.") replayable_results = [ ReplayableToolResult(content=result) for result in results @@ -685,8 +859,10 @@ def settle(self, intent: AuthorizedExecution, results: list[Content]) -> Approva snapshot_reconciliation=self._snapshot_reconciliation(occurrence), ) occurrence.outcome = outcome + self._emit_event("settlement", occurrence) return outcome + @_serialized_by_occurrence def settle_forwarded( self, intent: AuthorizedExecution, @@ -703,7 +879,7 @@ def settle_forwarded( f"Approval occurrence belongs to the {occurrence.owner.value} transition owner, not {owner.value}." ) if occurrence.status is not ApprovalStatus.EXECUTING: - raise ValueError(f"Approval occurrence is not executing: {occurrence.status}.") + raise ApprovalSettlementConflictError(f"Approval occurrence is not executing: {occurrence.status}.") replayable_results = [ ReplayableToolResult(content=result) for result in results @@ -729,8 +905,10 @@ def settle_forwarded( snapshot_reconciliation=self._snapshot_reconciliation(occurrence), ) occurrence.outcome = outcome + self._emit_event("settlement", occurrence) return outcome + @_serialized_by_occurrence def defer(self, intent: AuthorizedExecution, results: list[Content]) -> ApprovalOutcome: """Return an execution that yielded only follow-up requests to pending.""" occurrence = self._occurrences[intent.identity] diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py b/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py index 052389b201..913f129911 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py @@ -7,12 +7,13 @@ from collections import OrderedDict from typing import Any -from ._approval_lifecycle import ApprovalExecutionOwner, ApprovalLifecycle +from ._approval_lifecycle import ApprovalCapacityError, ApprovalExecutionOwner, ApprovalLifecycle ApprovalScope = str """Application-defined scope for server-side AG-UI Approval State.""" DEFAULT_MAX_APPROVAL_STATES = 10_000 +DEFAULT_TERMINAL_RETENTION_SECONDS = 900 _APPROVAL_SCOPE_INPUT_KEY = "__ag_ui_approval_scope" _APPROVAL_THREAD_SEPARATOR = "\x1f" @@ -34,15 +35,23 @@ def approval_state_thread_id(*, scope: object | None, thread_id: str) -> str: class InMemoryAGUIApprovalStateStore: """Bounded process-local server-side store for AG-UI Approval State. - The default store keeps only pending approval entries. It does not store - general ``AgentSession.state`` or AG-UI Thread Snapshots. + State is local to one process and is not durable across restarts or replicas. + Active and indeterminate occurrences are protected from eviction. Terminal + outcomes guarantee duplicate-execution protection for the configured + retention interval. """ - def __init__(self, *, max_entries: int = DEFAULT_MAX_APPROVAL_STATES) -> None: + def __init__( + self, + *, + max_entries: int = DEFAULT_MAX_APPROVAL_STATES, + terminal_retention_seconds: float = DEFAULT_TERMINAL_RETENTION_SECONDS, + ) -> None: """Initialize the process-local Approval State store. Keyword Args: - max_entries: Maximum pending approval entries to retain. + max_entries: Maximum approval occurrences or middleware state entries to retain. + terminal_retention_seconds: Process-local duplicate-execution protection window. Raises: ValueError: If ``max_entries`` is less than 1. @@ -52,7 +61,10 @@ def __init__(self, *, max_entries: int = DEFAULT_MAX_APPROVAL_STATES) -> None: self.max_entries = max_entries self.pending_approvals: OrderedDict[tuple[str, str], Any] = OrderedDict() self.tool_approval_states: OrderedDict[str, dict[str, Any]] = OrderedDict() - self.lifecycle = ApprovalLifecycle() + self.lifecycle = ApprovalLifecycle( + max_entries=max_entries, + terminal_retention_seconds=terminal_retention_seconds, + ) def register_local( self, @@ -191,11 +203,10 @@ def _register( self.pending_approvals.pop(key, None) for key in aliases: self.pending_approvals[key] = entry - self.evict_oldest() - - def evict_oldest(self) -> None: - """Evict oldest pending approval entries until the store is within bounds.""" - while len(self.pending_approvals) > self.max_entries: - self.pending_approvals.popitem(last=False) - while len(self.tool_approval_states) > self.max_entries: - self.tool_approval_states.popitem(last=False) + + def set_tool_approval_state(self, thread_id: str, state: dict[str, Any]) -> None: + """Store approval middleware state without evicting another active thread.""" + if thread_id not in self.tool_approval_states and len(self.tool_approval_states) >= self.max_entries: + raise ApprovalCapacityError("Approval state capacity is exhausted by protected occurrences.") + self.tool_approval_states[thread_id] = state + self.tool_approval_states.move_to_end(thread_id) diff --git a/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py b/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py index 7cd68b6a62..7d89d40a8f 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py +++ b/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py @@ -4,13 +4,19 @@ from __future__ import annotations +from concurrent.futures import ThreadPoolExecutor +from threading import Event + import pytest from agent_framework import Content from agent_framework_ag_ui._approval_lifecycle import ( + ApprovalCapacityError, + ApprovalClaimConflictError, ApprovalExecutionOwner, ApprovalIndeterminateError, ApprovalLifecycle, + ApprovalSettlementConflictError, ApprovalSnapshotStatus, ApprovalStatus, ClaimRecoveryPolicy, @@ -60,6 +66,278 @@ async def execute_authorized_call() -> list[Content]: assert [result.content.result for result in outcome.replayable_results] == ["Sunny"] +def test_active_occurrence_is_not_evicted_when_capacity_is_exhausted() -> None: + """Storage pressure fails explicitly instead of discarding pending authority.""" + lifecycle = ApprovalLifecycle(max_entries=1) + occurrence = lifecycle.register_local( + thread_id="thread-1", + interrupt_id="approval-1", + call_id="call-1", + name="write_record", + arguments='{"secret":"first"}', + ) + + with pytest.raises(ApprovalCapacityError): + lifecycle.register_local( + thread_id="thread-2", + interrupt_id="approval-2", + call_id="call-2", + name="write_record", + arguments='{"secret":"second"}', + ) + + assert lifecycle.get(occurrence.identity).status is ApprovalStatus.PENDING + + +async def test_terminal_outcome_expires_only_after_configured_retention_window() -> None: + """Duplicate execution protection lasts for the configured terminal retention window.""" + now = 100.0 + lifecycle = ApprovalLifecycle(max_entries=1, terminal_retention_seconds=30, clock=lambda: now) + occurrence = lifecycle.register_local( + thread_id="thread-1", + interrupt_id="approval-1", + call_id="call-1", + name="write_record", + arguments="{}", + ) + decision = ResumeDecision(interrupt_id="approval-1", accepted=True, arguments="{}") + intent = lifecycle.claim(thread_id="thread-1", decision=decision) + + async def execute() -> list[Content]: + return [Content.from_function_result(call_id="call-1", result="done")] + + outcome = await LocalPendingToolTransitionOwner(execute).execute(intent, lifecycle=lifecycle) + now = 129.0 + assert lifecycle.claim_batch(thread_id="thread-1", decisions=[decision]).retained_outcomes == (outcome,) + + now = 131.0 + replacement = lifecycle.register_local( + thread_id="thread-2", + interrupt_id="approval-2", + call_id="call-2", + name="write_record", + arguments="{}", + ) + + assert replacement.status is ApprovalStatus.PENDING + with pytest.raises(KeyError): + lifecycle.claim_batch(thread_id="thread-1", decisions=[decision]) + with pytest.raises(KeyError): + lifecycle.get(occurrence.identity) + + +def test_indeterminate_occurrence_remains_protected_after_terminal_retention_window() -> None: + """Uncertain execution is never aged out as a retryable terminal tombstone.""" + now = 100.0 + lifecycle = ApprovalLifecycle(max_entries=1, terminal_retention_seconds=30, clock=lambda: now) + occurrence = lifecycle.register_local( + thread_id="thread-1", + interrupt_id="approval-1", + call_id="call-1", + name="write_record", + arguments="{}", + ) + intent = lifecycle.claim( + thread_id="thread-1", + decision=ResumeDecision(interrupt_id="approval-1", accepted=True, arguments="{}"), + ) + lifecycle.begin_execution(intent, owner=ApprovalExecutionOwner.LOCAL) + lifecycle.recover_execution(intent, owner=ApprovalExecutionOwner.LOCAL) + now = 1_000.0 + + with pytest.raises(ApprovalCapacityError): + lifecycle.register_local( + thread_id="thread-2", + interrupt_id="approval-2", + call_id="call-2", + name="write_record", + arguments="{}", + ) + + assert lifecycle.get(occurrence.identity).status is ApprovalStatus.INDETERMINATE + + +async def test_transition_telemetry_covers_lifecycle_without_sensitive_payloads( + caplog: pytest.LogCaptureFixture, +) -> None: + """Operators can distinguish lifecycle transitions without logging tool inputs.""" + caplog.set_level("INFO", logger="agent_framework_ag_ui._approval_lifecycle") + lifecycle = ApprovalLifecycle() + secret = "sensitive-value" + settled = lifecycle.register_local( + thread_id="thread-settled", + interrupt_id="approval-settled", + call_id="call-settled", + name="write_secret", + arguments=f'{{"value":"{secret}"}}', + ) + decision = ResumeDecision( + interrupt_id="approval-settled", + accepted=True, + arguments=f'{{"value":"{secret}"}}', + ) + intent = lifecycle.claim(thread_id="thread-settled", decision=decision) + + async def execute() -> list[Content]: + return [Content.from_function_result(call_id="call-settled", result="done")] + + await LocalPendingToolTransitionOwner(execute).execute(intent, lifecycle=lifecycle) + lifecycle.claim_batch(thread_id="thread-settled", decisions=[decision]) + + lifecycle.register_local( + thread_id="thread-rejected", + interrupt_id="approval-rejected", + call_id="call-rejected", + name="reject_secret", + arguments="{}", + ) + lifecycle.claim_batch( + thread_id="thread-rejected", + decisions=[ResumeDecision(interrupt_id="approval-rejected", accepted=False, arguments="{}")], + ) + lifecycle.register_local( + thread_id="thread-cancelled", + interrupt_id="approval-cancelled", + call_id="call-cancelled", + name="cancel_secret", + arguments="{}", + ) + lifecycle.cancel_batch(thread_id="thread-cancelled", interrupt_ids=["approval-cancelled"]) + lifecycle.register_local( + thread_id="thread-expired", + interrupt_id="approval-expired", + call_id="call-expired", + name="expire_secret", + arguments="{}", + ) + lifecycle.expire_batch(thread_id="thread-expired", interrupt_ids=["approval-expired"]) + uncertain = lifecycle.register_local( + thread_id="thread-uncertain", + interrupt_id="approval-uncertain", + call_id="call-uncertain", + name="uncertain_secret", + arguments="{}", + ) + uncertain_intent = lifecycle.claim( + thread_id="thread-uncertain", + decision=ResumeDecision(interrupt_id="approval-uncertain", accepted=True, arguments="{}"), + ) + lifecycle.begin_execution(uncertain_intent, owner=ApprovalExecutionOwner.LOCAL) + lifecycle.recover_execution(uncertain_intent, owner=ApprovalExecutionOwner.LOCAL) + with pytest.raises(KeyError): + lifecycle.claim_batch( + thread_id="thread-missing", + decisions=[ResumeDecision(interrupt_id="approval-missing", accepted=True, arguments="{}")], + ) + capacity_lifecycle = ApprovalLifecycle(max_entries=1) + capacity_lifecycle.register_local( + thread_id="thread-capacity-1", + interrupt_id="approval-capacity-1", + call_id="call-capacity-1", + name="capacity_secret", + arguments="{}", + ) + with pytest.raises(ApprovalCapacityError): + capacity_lifecycle.register_local( + thread_id="thread-capacity-2", + interrupt_id="approval-capacity-2", + call_id="call-capacity-2", + name="capacity_secret", + arguments="{}", + ) + + events = {getattr(record, "approval_event", None) for record in caplog.records} + assert { + "registration", + "claim", + "execution_start", + "settlement", + "rejection", + "cancellation", + "duplicate", + "expiration", + "indeterminate_recovery", + "authority_failure", + "capacity_failure", + } <= events + assert any( + getattr(record, "approval_occurrence_id", None) == settled.identity.occurrence_id for record in caplog.records + ) + assert lifecycle.get(uncertain.identity).status is ApprovalStatus.INDETERMINATE + assert secret not in caplog.text + assert all(secret not in repr(record.__dict__) for record in caplog.records) + + +def test_claim_and_settlement_conflicts_have_typed_outcomes() -> None: + """Adapters can distinguish transition conflicts without parsing error messages.""" + lifecycle = ApprovalLifecycle() + lifecycle.register_local( + thread_id="thread-1", + interrupt_id="approval-1", + call_id="call-1", + name="write_record", + arguments="{}", + ) + decision = ResumeDecision(interrupt_id="approval-1", accepted=True, arguments="{}") + intent = lifecycle.claim(thread_id="thread-1", decision=decision) + + with pytest.raises(ApprovalClaimConflictError): + lifecycle.claim_batch(thread_id="thread-1", decisions=[decision]) + with pytest.raises(ApprovalSettlementConflictError): + lifecycle.settle(intent, [Content.from_function_result(call_id="call-1", result="done")]) + + +def test_same_thread_transitions_serialize_without_blocking_an_independent_thread() -> None: + """One scoped thread is serialized while another can claim concurrently.""" + lifecycle = ApprovalLifecycle() + first = lifecycle.register_local( + thread_id="thread-1", + interrupt_id="approval-1", + call_id="call-1", + name="write_record", + arguments="{}", + ) + lifecycle.register_local( + thread_id="thread-2", + interrupt_id="approval-2", + call_id="call-2", + name="write_record", + arguments="{}", + ) + first_decision = ResumeDecision(interrupt_id="approval-1", accepted=True, arguments="{}") + second_decision = ResumeDecision(interrupt_id="approval-2", accepted=True, arguments="{}") + first_claim_entered = Event() + release_first_claim = Event() + second_same_thread_started = Event() + original_emit = lifecycle._emit_event + + def blocking_emit(event: str, occurrence=None, *, failure_type: str | None = None) -> None: + if event == "claim" and occurrence is not None and occurrence.identity == first.identity: + first_claim_entered.set() + assert release_first_claim.wait(timeout=2) + original_emit(event, occurrence, failure_type=failure_type) + + lifecycle._emit_event = blocking_emit # type: ignore[method-assign] + + def repeat_first_claim(): + second_same_thread_started.set() + return lifecycle.claim_batch(thread_id="thread-1", decisions=[first_decision]) + + with ThreadPoolExecutor(max_workers=3) as executor: + first_claim = executor.submit(lifecycle.claim, thread_id="thread-1", decision=first_decision) + assert first_claim_entered.wait(timeout=2) + conflicting_claim = executor.submit(repeat_first_claim) + assert second_same_thread_started.wait(timeout=2) + independent_claim = executor.submit(lifecycle.claim, thread_id="thread-2", decision=second_decision) + + assert independent_claim.result(timeout=2).identity.call_id == "call-2" + assert not conflicting_claim.done() + release_first_claim.set() + assert first_claim.result(timeout=2).identity == first.identity + with pytest.raises(ApprovalClaimConflictError): + conflicting_claim.result(timeout=2) + + async def test_hosted_approval_is_forwarded_only_by_its_owner_and_settles_same_occurrence() -> None: """Hosted authority cannot execute locally and records forwarding against its occurrence.""" lifecycle = ApprovalLifecycle() diff --git a/python/packages/ag-ui/tests/ag_ui/test_approval_state.py b/python/packages/ag-ui/tests/ag_ui/test_approval_state.py index 99695c361e..f64400aa5e 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_approval_state.py +++ b/python/packages/ag-ui/tests/ag_ui/test_approval_state.py @@ -4,6 +4,7 @@ import pytest +from agent_framework_ag_ui._approval_lifecycle import ApprovalCapacityError from agent_framework_ag_ui._approval_state import InMemoryAGUIApprovalStateStore, approval_state_thread_id @@ -30,14 +31,33 @@ def test_approval_state_store_rejects_invalid_max_entries() -> None: InMemoryAGUIApprovalStateStore(max_entries=0) -def test_approval_state_store_evicts_oldest_entries() -> None: +def test_approval_state_store_does_not_evict_active_entries() -> None: store = InMemoryAGUIApprovalStateStore(max_entries=1) - store.pending_approvals[("thread-1", "call-1")] = "first" - store.pending_approvals[("thread-2", "call-2")] = "second" - store.tool_approval_states["thread-1"] = {"call_id": "call-1"} - store.tool_approval_states["thread-2"] = {"call_id": "call-2"} + store.register_local( + thread_ids=["thread-1"], + name="write_record", + arguments="{}", + request_id="request-1", + interrupt_id="approval-1", + ) + + with pytest.raises(ApprovalCapacityError): + store.register_local( + thread_ids=["thread-2"], + name="write_record", + arguments="{}", + request_id="request-2", + interrupt_id="approval-2", + ) + + assert {entry["interrupt_id"] for entry in store.pending_approvals.values()} == {"approval-1"} + + +def test_approval_state_store_does_not_evict_active_middleware_state() -> None: + store = InMemoryAGUIApprovalStateStore(max_entries=1) + store.set_tool_approval_state("thread-1", {"call_id": "call-1"}) - store.evict_oldest() + with pytest.raises(ApprovalCapacityError): + store.set_tool_approval_state("thread-2", {"call_id": "call-2"}) - assert list(store.pending_approvals.items()) == [(("thread-2", "call-2"), "second")] - assert list(store.tool_approval_states.items()) == [("thread-2", {"call_id": "call-2"})] + assert store.tool_approval_states == {"thread-1": {"call_id": "call-1"}} From 31c483e65fc69c3f01d87eb6c0986e0d853cfab0 Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Mon, 10 Aug 2026 14:53:00 +0900 Subject: [PATCH 08/13] Complete approval lifecycle cutover Key decisions: - Make ApprovalLifecycle the sole owner of trusted aliases, occurrence metadata, authority transitions, and retained outcomes. - Remove the parallel mutable pending-approval registry and route local, hosted, deferred, cancellation, replay, and snapshot reconciliation through lifecycle occurrences. - Encapsulate middleware Approval State behind copy-isolated store methods while keeping AG-UI protocol normalization and event projection in the runner. Files changed: - packages/ag-ui/AGENTS.md - packages/ag-ui/agent_framework_ag_ui/_agent.py - packages/ag-ui/agent_framework_ag_ui/_agent_run.py - packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py - packages/ag-ui/agent_framework_ag_ui/_approval_state.py - packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py - packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py - packages/ag-ui/tests/ag_ui/test_approval_result_event.py - packages/ag-ui/tests/ag_ui/test_approval_state.py - packages/ag-ui/tests/ag_ui/test_endpoint.py - packages/ag-ui/tests/ag_ui/test_run.py Verification: - 964 package-local AG-UI tests passed with 92% coverage and 90% approval lifecycle coverage. - 85 warning-strict focused approval tests passed. - Package-local Ruff and Pyright passed; git diff --check passed. Notes for next iteration: - The function-calling-loop scenario mapping remains inaccessible under the organization content-exclusion policy. - Workspace Poe fan-out remains blocked by the pre-existing missing packages/durabletask/pyproject.toml; equivalent package-local checks passed. --- python/packages/ag-ui/AGENTS.md | 3 + .../ag-ui/agent_framework_ag_ui/_agent.py | 7 +- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 812 +++------- .../_approval_lifecycle.py | 110 +- .../agent_framework_ag_ui/_approval_state.py | 46 +- .../ag_ui/test_agent_wrapper_comprehensive.py | 42 +- .../tests/ag_ui/test_approval_lifecycle.py | 30 + .../tests/ag_ui/test_approval_result_event.py | 1386 +---------------- .../ag-ui/tests/ag_ui/test_approval_state.py | 5 +- .../ag-ui/tests/ag_ui/test_endpoint.py | 21 +- python/packages/ag-ui/tests/ag_ui/test_run.py | 75 +- 11 files changed, 532 insertions(+), 2005 deletions(-) diff --git a/python/packages/ag-ui/AGENTS.md b/python/packages/ag-ui/AGENTS.md index ac4b52f0ae..bee4a0f0e6 100644 --- a/python/packages/ag-ui/AGENTS.md +++ b/python/packages/ag-ui/AGENTS.md @@ -33,6 +33,9 @@ AG-UI protocol integration for building agent UIs with the AG-UI standard. resumed messages, while `TOOL_CALL_RESULT` events are emitted only for terminal `function_result` contents. - Approval responses for tools injected during `before_run` are deferred to the in-run approval middleware rather than executed or rejected by the transport before those tools exist. +- `_approval_lifecycle.py` is the sole owner of approval occurrence registration, trusted aliases, authority + validation, claims, terminal outcomes, and retry deduplication. Runner code normalizes AG-UI protocol values and + projects lifecycle outcomes but must not maintain a parallel pending-approval registry. - `confirm_changes` snapshot cleanup resolves the synthetic confirmation back to its original `function_call_id`; it must never concatenate unrelated tool results or record accepted changes without a matching real result. - SSE keepalive is endpoint-owned transport behavior configured through diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent.py index 09766b7e1a..3a380f11e2 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent.py @@ -9,7 +9,7 @@ from agent_framework import SupportsAgentRun from agent_framework._telemetry import mark_feature_used -from ._agent_run import PendingApprovalEntry, PendingApprovalKey, run_agent_stream +from ._agent_run import run_agent_stream from ._approval_state import InMemoryAGUIApprovalStateStore from ._feature_usage import FeatureIndex from ._snapshots import AGUIThreadSnapshotStore @@ -122,10 +122,6 @@ def __init__( # Server-side Approval State. Populated when approval requests are emitted # and consumed when resume decisions arrive. self._approval_state_store = InMemoryAGUIApprovalStateStore() - self._pending_approvals = cast( - dict[PendingApprovalKey, PendingApprovalEntry], - self._approval_state_store.pending_approvals, - ) @property def snapshot_store(self) -> AGUIThreadSnapshotStore | None: @@ -149,7 +145,6 @@ async def run( input_data, self.agent, self.config, - pending_approvals=self._pending_approvals, approval_state_store=self._approval_state_store, ): yield event diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index c96444c334..dc227171aa 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py @@ -11,7 +11,7 @@ from collections.abc import AsyncIterable, Awaitable, Mapping, Sequence from dataclasses import dataclass, field from functools import partial -from typing import TYPE_CHECKING, Any, TypedDict, cast +from typing import TYPE_CHECKING, Any, cast from ag_ui.core import ( BaseEvent, @@ -57,6 +57,7 @@ from ._approval_lifecycle import ( ApprovalExecutionOwner, ApprovalLifecycle, + ApprovalOccurrence, ApprovalSnapshotReconciliation, AuthorizedExecution, DeferredPendingToolTransitionOwner, @@ -689,90 +690,22 @@ def _make_approval_tool_result_events(resolved_approval_results: list[Content]) return events -class _PendingApproval(TypedDict): - """Pending approval details for a requested function call.""" - - name: str - arguments: str | None - request_id: str | None - interrupt_id: str | None - - -class _PendingApprovalWithSiblings(_PendingApproval, total=False): - """Pending approval details including sibling calls and trusted hosted metadata.""" - - already_approved_requests: list[dict[str, Any]] - execution_owner: str - server_label: str - - -PendingApprovalEntry = _PendingApprovalWithSiblings | str -PendingApprovalKey = tuple[str, str] - - -def _pending_approval_key(thread_id: str, interrupt_id: str) -> PendingApprovalKey: - """Build a structured pending-approval key scoped by thread and interrupt id.""" - return (thread_id, interrupt_id) - - -def _make_pending_approval_entry( - name: str, - arguments: str | None, - *, - request_id: str | None = None, - interrupt_id: str | None = None, - already_approved_requests: list[dict[str, Any]] | None = None, - server_label: str | None = None, - execution_owner: ApprovalExecutionOwner | None = None, -) -> _PendingApprovalWithSiblings: - entry: _PendingApprovalWithSiblings = { - "name": name, - "arguments": arguments, - "request_id": request_id, - "interrupt_id": interrupt_id, - } - if already_approved_requests: - entry["already_approved_requests"] = already_approved_requests - if server_label: - entry["server_label"] = server_label - if execution_owner is not None: - entry["execution_owner"] = execution_owner.value - return entry +def _pending_approval_name(entry: ApprovalOccurrence) -> str: + return entry.name -def _register_pending_approval_entry( - pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry], - thread_id: str, - entry: PendingApprovalEntry, - *ids: str | None, -) -> None: - """Register all server-owned aliases for one pending approval entry.""" - for alias_key in _pending_approval_alias_keys(thread_id, entry, *ids): - pending_approvals[alias_key] = entry - - -def _pending_approval_name(entry: PendingApprovalEntry) -> str | None: - if isinstance(entry, str): - return entry - return entry["name"] - - -def _pending_approval_arguments(entry: PendingApprovalEntry) -> str | None: - if isinstance(entry, str): - return None - return entry["arguments"] +def _pending_approval_arguments(entry: ApprovalOccurrence) -> str: + return entry.arguments -def _pending_approval_already_approved_requests(entry: PendingApprovalEntry) -> list[dict[str, Any]]: - if isinstance(entry, str): - return [] - return list(entry.get("already_approved_requests", [])) +def _pending_approval_already_approved_requests( + entry: ApprovalOccurrence, +) -> list[dict[str, Any]]: + return list(entry.already_approved_requests) -def _pending_approval_server_label(entry: PendingApprovalEntry) -> str | None: - if isinstance(entry, str): - return None - return entry.get("server_label") +def _pending_approval_server_label(entry: ApprovalOccurrence) -> str | None: + return entry.server_label def _function_call_server_label(function_call: Content | None) -> str | None: @@ -878,11 +811,10 @@ def _restore_tool_approval_state( session.state.pop(_TOOL_APPROVAL_STATE_KEY, None) if approval_state_store is None: return - stored_state = approval_state_store.tool_approval_states.get(thread_id) + stored_state = approval_state_store.get_tool_approval_state(thread_id) if stored_state is None: return - approval_state_store.tool_approval_states.move_to_end(thread_id) - session.state[_TOOL_APPROVAL_STATE_KEY] = copy.deepcopy(stored_state) + session.state[_TOOL_APPROVAL_STATE_KEY] = stored_state def _save_tool_approval_state( @@ -895,7 +827,7 @@ def _save_tool_approval_state( return raw_state = session.state.get(_TOOL_APPROVAL_STATE_KEY) if raw_state is None: - approval_state_store.tool_approval_states.pop(thread_id, None) + approval_state_store.delete_tool_approval_state(thread_id) return serialized_state = _serialized_tool_approval_state(raw_state) if serialized_state is None: @@ -910,52 +842,65 @@ def _clear_tool_approval_state( """Discard queued ToolApprovalMiddleware state for a cancelled approval flow.""" if approval_state_store is None: return - approval_state_store.tool_approval_states.pop(thread_id, None) + approval_state_store.delete_tool_approval_state(thread_id) def _register_server_generated_approval_response( response: Content, - pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] | None, thread_id: str, tools: list[Any] | None, *, + lifecycle: ApprovalLifecycle, has_deferred_owner: bool, -) -> None: +) -> AuthorizedExecution | None: """Register a server-owned approval response so normal validation can consume it.""" - if pending_approvals is None or response.function_call is None or not response.function_call.name: - return + if response.function_call is None or not response.function_call.name: + return None response_id = response.id or response.function_call.call_id if not response_id: - return + return None execution_owner = _function_call_execution_owner( response.function_call, tools, has_deferred_owner=has_deferred_owner, ) - entry = _make_pending_approval_entry( - response.function_call.name, - canonical_function_arguments(response.function_call), - request_id=str(response.id) if response.id else None, - interrupt_id=str(response.function_call.call_id) if response.function_call.call_id else None, + if execution_owner is ApprovalExecutionOwner.HOSTED: + register = lifecycle.register_hosted + elif execution_owner is ApprovalExecutionOwner.DEFERRED: + register = lifecycle.register_deferred + elif execution_owner is ApprovalExecutionOwner.LOCAL: + register = lifecycle.register_local + else: + register = lifecycle.register_unowned + arguments = canonical_function_arguments(response.function_call) or "{}" + register( + thread_id=thread_id, + interrupt_id=str(response_id), + call_id=str(response.function_call.call_id or response_id), + name=response.function_call.name, + arguments=arguments, + aliases=[str(response.function_call.call_id)] if response.function_call.call_id else None, server_label=_function_call_server_label(response.function_call), - execution_owner=execution_owner, ) - _register_pending_approval_entry( - pending_approvals, - thread_id, - entry, - str(response_id), - str(response.function_call.call_id) if response.function_call.call_id else None, + if not response.approved or execution_owner is ApprovalExecutionOwner.UNAVAILABLE: + return None + return lifecycle.claim( + thread_id=thread_id, + decision=ResumeDecision( + interrupt_id=str(response_id), + accepted=True, + arguments=arguments, + name=response.function_call.name, + ), ) def _pop_collected_tool_approval_response_messages( session: AgentSession, - pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] | None, thread_id: str, tools: list[Any] | None, *, - lifecycle: ApprovalLifecycle | None = None, + lifecycle: ApprovalLifecycle, authorized_executions: dict[str, AuthorizedExecution] | None = None, ) -> list[Message]: """Pop server-collected auto-approved responses into provider-visible messages.""" @@ -973,52 +918,14 @@ def _pop_collected_tool_approval_response_messages( response = _content_from_approval_state(raw_response) if response is None or response.type != "function_approval_response": continue - _register_server_generated_approval_response( + intent = _register_server_generated_approval_response( response, - pending_approvals, thread_id, tools, + lifecycle=lifecycle, has_deferred_owner=True, ) - function_call = response.function_call - if ( - lifecycle is not None - and authorized_executions is not None - and response.approved - and function_call is not None - and function_call.call_id - and function_call.name - ): - arguments = canonical_function_arguments(function_call) or "{}" - execution_owner = _function_call_execution_owner( - function_call, - tools, - has_deferred_owner=True, - ) - if execution_owner is ApprovalExecutionOwner.HOSTED: - register = lifecycle.register_hosted - elif execution_owner is ApprovalExecutionOwner.DEFERRED: - register = lifecycle.register_deferred - elif execution_owner is ApprovalExecutionOwner.LOCAL: - register = lifecycle.register_local - else: - register = lifecycle.register_unowned - register( - thread_id=thread_id, - interrupt_id=str(response.id or function_call.call_id), - call_id=str(function_call.call_id), - name=function_call.name, - arguments=arguments, - ) - intent = lifecycle.claim( - thread_id=thread_id, - decision=ResumeDecision( - interrupt_id=str(response.id or function_call.call_id), - accepted=True, - arguments=arguments, - name=function_call.name, - ), - ) + if intent is not None and authorized_executions is not None: authorized_executions[intent.identity.call_id] = intent responses.append(response) @@ -1041,33 +948,6 @@ def _parse_json_object(value: Any) -> dict[str, Any] | None: return cast(dict[str, Any], parsed) if isinstance(parsed, dict) else None -def _thread_has_pending_approvals( - pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] | None, - thread_id: str, -) -> bool: - if not pending_approvals: - return False - return any(key[0] == thread_id for key in pending_approvals) - - -def _pending_approval_interrupt_ids( - pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] | None, - thread_id: str, -) -> set[str]: - if not pending_approvals: - return set() - interrupt_ids: set[str] = set() - for key, entry in pending_approvals.items(): - if key[0] != thread_id: - continue - if isinstance(entry, str): - interrupt_ids.add(key[1]) - continue - interrupt_id = entry.get("interrupt_id") or entry.get("request_id") or key[1] - interrupt_ids.add(str(interrupt_id)) - return interrupt_ids - - def _content_tool_call_ids(value: Any) -> set[str]: """Collect tool call ids from serialized approval-specific state values.""" call_ids: set[str] = set() @@ -1099,24 +979,17 @@ def _content_tool_call_ids(value: Any) -> set[str]: def _approval_state_tool_call_ids( - pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] | None, - approval_state_store: InMemoryAGUIApprovalStateStore | None, + approval_state_store: InMemoryAGUIApprovalStateStore, thread_id: str, ) -> set[str]: """Return server-owned approval call ids that are not abandoned.""" - call_ids = _pending_approval_interrupt_ids(pending_approvals, thread_id) - if pending_approvals: - for key, entry in pending_approvals.items(): - if key[0] != thread_id: - continue - call_ids.add(key[1]) - if isinstance(entry, str): - continue - call_ids.update(_content_tool_call_ids(entry.get("already_approved_requests", []))) - - if approval_state_store is None: - return call_ids - stored_state = approval_state_store.tool_approval_states.get(thread_id) + call_ids: set[str] = set() + for occurrence in approval_state_store.lifecycle.occurrences_for_thread(thread_id=thread_id): + call_ids.add(occurrence.identity.call_id) + call_ids.add(occurrence.identity.interrupt_id) + call_ids.update(occurrence.aliases) + call_ids.update(_content_tool_call_ids(list(occurrence.already_approved_requests))) + stored_state = approval_state_store.get_tool_approval_state(thread_id) if stored_state is not None: call_ids.update(_content_tool_call_ids(stored_state)) return call_ids @@ -1133,7 +1006,7 @@ def _tool_approval_state_exists_for_cancelled_resume( cancelled_ids = _cancelled_resume_interrupt_ids(resume_payload) if not cancelled_ids: return False - return thread_id in approval_state_store.tool_approval_states + return approval_state_store.has_tool_approval_state(thread_id) def _stored_pending_approval_interrupt_ids(interrupts: list[dict[str, Any]] | None) -> set[str]: @@ -1169,94 +1042,6 @@ def _resume_payload_has_approval_decision(resume_payload: Any) -> bool: return False -def _find_pending_approval_entry( - pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] | None, - thread_id: str, - interrupt_id: str, -) -> PendingApprovalEntry | None: - if pending_approvals is None: - return None - return pending_approvals.get(_pending_approval_key(thread_id, interrupt_id)) - - -def _pending_approval_alias_keys( - thread_id: str, - entry: PendingApprovalEntry, - *ids: str | None, -) -> set[PendingApprovalKey]: - aliases = {item for item in ids if item} - if not isinstance(entry, str): - request_id = entry.get("request_id") - interrupt_id = entry.get("interrupt_id") - if request_id: - aliases.add(request_id) - if interrupt_id: - aliases.add(interrupt_id) - return {_pending_approval_key(thread_id, alias) for alias in aliases} - - -def _remove_pending_approval(registry: dict[PendingApprovalKey, PendingApprovalEntry], key: PendingApprovalKey) -> None: - """Remove one pending approval and every alias that references the same entry.""" - entry = registry.pop(key, None) - if entry is None or isinstance(entry, str): - return - - for alias_key, alias_entry in list(registry.items()): - if alias_entry is entry: - registry.pop(alias_key, None) - - -def _register_pending_approval( - registry: dict[PendingApprovalKey, PendingApprovalEntry], - thread_ids: list[str], - name: str, - arguments: str | None, - *, - request_id: str, - interrupt_id: str | None, - already_approved_requests: list[dict[str, Any]] | None = None, - server_label: str | None = None, -) -> None: - """Register one pending approval under each distinct thread identity.""" - keys = list( - dict.fromkeys( - _pending_approval_key(thread_id, approval_id) - for thread_id in thread_ids - for approval_id in (request_id, interrupt_id) - if approval_id - ) - ) - for key in keys: - _remove_pending_approval(registry, key) - - entry = _make_pending_approval_entry( - name, - arguments, - request_id=request_id, - interrupt_id=interrupt_id, - already_approved_requests=already_approved_requests, - server_label=server_label, - ) - for key in keys: - registry[key] = entry - - -def _consume_pending_approval_entry( - pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry], - thread_id: str, - entry: PendingApprovalEntry, - *ids: str | None, -) -> None: - if not isinstance(entry, str): - for key, candidate in list(pending_approvals.items()): - if candidate is entry: - pending_approvals.pop(key, None) - return - - for alias_key in _pending_approval_alias_keys(thread_id, entry, *ids): - pending_approvals.pop(alias_key, None) - - def _approval_arguments_match_pending(pending_arguments: str | None, response_arguments: str | None) -> bool: return pending_arguments is None or response_arguments == pending_arguments @@ -1279,63 +1064,26 @@ def _json_schema_value_matches(original_value: Any, edited_value: Any) -> bool: def _canonical_approval_resume_messages( resume_payload: Any, - pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] | None, thread_id: str, expected_interrupt_ids: set[str] | None = None, *, - lifecycle: ApprovalLifecycle | None = None, + lifecycle: ApprovalLifecycle, authorized_executions: dict[str, AuthorizedExecution] | None = None, retained_results: list[Content] | None = None, snapshot_reconciliations: list[ApprovalSnapshotReconciliation] | None = None, ) -> tuple[list[dict[str, Any]], set[str], set[str], RunErrorEvent | None]: """Translate canonical ResumeEntry approvals into existing approval response messages.""" expected_ids = set(expected_interrupt_ids or set()) - if pending_approvals is None: - if not expected_ids: - return [], set(), set(), None - entries, contract_error, contract_code = _resume_contract_error( - resume_payload, - expected_ids, - required_code="APPROVAL_RESUME_REQUIRED", - invalid_code="APPROVAL_RESUME_INVALID", - unknown_code="APPROVAL_RESUME_NOT_FOUND", - missing_code="APPROVAL_RESUME_MISSING_INTERRUPT", - ) - if contract_error is not None and contract_code is not None: - return [], set(), set(), RunErrorEvent(message=contract_error, code=contract_code) - cancelled_expected_ids = {str(entry["interrupt_id"]) for entry in entries if entry.get("status") == "cancelled"} - if cancelled_expected_ids: - interrupt_id = next(str(entry["interrupt_id"]) for entry in entries if entry.get("status") == "cancelled") - return ( - [], - cancelled_expected_ids, - cancelled_expected_ids, - RunErrorEvent( - message=f"Approval resume for interruptId '{interrupt_id}' was cancelled.", - code="APPROVAL_RESUME_CANCELLED", - ), - ) - interrupt_id = entries[0]["interrupt_id"] if entries else sorted(expected_ids)[0] - return ( - [], - set(), - set(), - RunErrorEvent( - message=f"No pending approval interrupt found for resume interruptId '{interrupt_id}'.", - code="APPROVAL_RESUME_NOT_FOUND", - ), - ) - messages: list[dict[str, Any]] = [] handled_ids: set[str] = set() cancelled_ids: set[str] = set() - pending_interrupt_ids = _pending_approval_interrupt_ids(pending_approvals, thread_id) + pending_interrupt_ids = lifecycle.pending_interrupt_ids(thread_id=thread_id) contract_interrupt_ids = expected_ids | pending_interrupt_ids if not contract_interrupt_ids: if _resume_payload_has_approval_decision(resume_payload): normalized_interrupts = _normalize_resume_interrupts(resume_payload) interrupt_id = normalized_interrupts[0]["id"] if normalized_interrupts else "unknown" - if lifecycle is not None and retained_results is not None: + if retained_results is not None: decisions: list[ResumeDecision] = [] for interrupt in normalized_interrupts: if interrupt.get("status") != "resolved": @@ -1416,12 +1164,12 @@ def _canonical_approval_resume_messages( if contract_error is not None and contract_code is not None: return [], handled_ids, cancelled_ids, RunErrorEvent(message=contract_error, code=contract_code) - has_pending_for_thread = _thread_has_pending_approvals(pending_approvals, thread_id) - entries_by_interrupt_id: dict[str, PendingApprovalEntry | None] = {} + has_pending_for_thread = bool(pending_interrupt_ids) + entries_by_interrupt_id: dict[str, ApprovalOccurrence | None] = {} for entry in entries: interrupt_id = cast(str, entry["interrupt_id"]) status = entry["status"] - pending_entry = _find_pending_approval_entry(pending_approvals, thread_id, interrupt_id) + pending_entry = lifecycle.pending_occurrence(thread_id=thread_id, interrupt_id=interrupt_id) if pending_entry is None: if status == "cancelled" and interrupt_id in expected_ids: handled_ids.add(interrupt_id) @@ -1457,28 +1205,23 @@ def _canonical_approval_resume_messages( ) if cancelled_ids: - if lifecycle is not None: - lifecycle_cancelled_ids = [ - interrupt_id for interrupt_id in cancelled_ids if entries_by_interrupt_id.get(interrupt_id) is not None - ] - try: - reconciliations = lifecycle.cancel_batch( - thread_id=thread_id, - interrupt_ids=lifecycle_cancelled_ids, - ) - if snapshot_reconciliations is not None: - snapshot_reconciliations.extend(reconciliations) - except (KeyError, ValueError) as exc: - return ( - [], - handled_ids, - cancelled_ids, - RunErrorEvent(message=str(exc), code="APPROVAL_RESUME_INVALID"), - ) - for interrupt_id in cancelled_ids: - pending_entry = entries_by_interrupt_id.get(interrupt_id) - if pending_entry is not None: - _consume_pending_approval_entry(pending_approvals, thread_id, pending_entry, interrupt_id) + lifecycle_cancelled_ids = [ + interrupt_id for interrupt_id in cancelled_ids if entries_by_interrupt_id.get(interrupt_id) is not None + ] + try: + reconciliations = lifecycle.cancel_batch( + thread_id=thread_id, + interrupt_ids=lifecycle_cancelled_ids, + ) + if snapshot_reconciliations is not None: + snapshot_reconciliations.extend(reconciliations) + except (KeyError, ValueError) as exc: + return ( + [], + handled_ids, + cancelled_ids, + RunErrorEvent(message=str(exc), code="APPROVAL_RESUME_INVALID"), + ) interrupt_id = next(str(entry["interrupt_id"]) for entry in entries if entry.get("status") == "cancelled") return ( [], @@ -1490,7 +1233,6 @@ def _canonical_approval_resume_messages( ), ) - argument_updates: list[tuple[_PendingApprovalWithSiblings, str]] = [] lifecycle_decisions: list[ResumeDecision] = [] restored_sibling_response_ids: set[str] = set() for entry in entries: @@ -1565,22 +1307,19 @@ def _canonical_approval_resume_messages( merged_arguments = {**original_arguments, **edited_arguments} canonical_arguments = json.dumps(make_json_safe(merged_arguments), sort_keys=True, separators=(",", ":")) - if not isinstance(pending_entry, str): - argument_updates.append((pending_entry, canonical_arguments)) - if lifecycle is not None: - lifecycle_decisions.append( - ResumeDecision( - interrupt_id=interrupt_id, - accepted=accepted, - arguments=canonical_arguments, - name=_pending_approval_name(pending_entry), - original_arguments=pending_arguments, - ) + lifecycle_decisions.append( + ResumeDecision( + interrupt_id=interrupt_id, + accepted=accepted, + arguments=canonical_arguments, + name=_pending_approval_name(pending_entry), + original_arguments=pending_arguments, ) + ) function_approvals = [ { "id": interrupt_id, - "call_id": interrupt_id, + "call_id": pending_entry.identity.call_id, "name": _pending_approval_name(pending_entry) or "", "approved": accepted, "arguments": merged_arguments, @@ -1600,49 +1339,29 @@ def _canonical_approval_resume_messages( if str(response_id) in restored_sibling_response_ids: continue restored_sibling_response_ids.add(str(response_id)) - sibling_entry = _make_pending_approval_entry( - function_call.name, - canonical_function_arguments(function_call), - request_id=str(response.id) if response.id else None, - interrupt_id=str(function_call.call_id) if function_call.call_id else None, - server_label=_function_call_server_label(function_call), - execution_owner=( - ApprovalExecutionOwner.HOSTED - if _function_call_server_label(function_call) - else ApprovalExecutionOwner.LOCAL - ), + sibling_interrupt_id = str(response_id) + sibling_call_id = str(function_call.call_id or response_id) + sibling_arguments = canonical_function_arguments(function_call) or "{}" + register = ( + lifecycle.register_hosted if _function_call_server_label(function_call) else lifecycle.register_local ) - _register_pending_approval_entry( - pending_approvals, - thread_id, - sibling_entry, - str(response_id), - str(function_call.call_id) if function_call.call_id else None, + register( + thread_id=thread_id, + interrupt_id=sibling_interrupt_id, + call_id=sibling_call_id, + name=function_call.name, + arguments=sibling_arguments, + aliases=[sibling_call_id], + server_label=_function_call_server_label(function_call), ) - if lifecycle is not None: - sibling_interrupt_id = str(response_id) - sibling_call_id = str(function_call.call_id or response_id) - sibling_arguments = canonical_function_arguments(function_call) or "{}" - register = ( - lifecycle.register_hosted - if _function_call_server_label(function_call) - else lifecycle.register_local - ) - register( - thread_id=thread_id, + lifecycle_decisions.append( + ResumeDecision( interrupt_id=sibling_interrupt_id, - call_id=sibling_call_id, - name=function_call.name, + accepted=True, arguments=sibling_arguments, + name=function_call.name, ) - lifecycle_decisions.append( - ResumeDecision( - interrupt_id=sibling_interrupt_id, - accepted=True, - arguments=sibling_arguments, - name=function_call.name, - ) - ) + ) function_approvals.append( { "id": str(response_id), @@ -1654,10 +1373,7 @@ def _canonical_approval_resume_messages( ) messages.append({"role": "user", "function_approvals": function_approvals}) - for pending_entry, arguments_json in argument_updates: - pending_entry["arguments"] = arguments_json - - if lifecycle is not None and authorized_executions is not None: + if authorized_executions is not None: try: intents = lifecycle.claim_batch(thread_id=thread_id, decisions=lifecycle_decisions) if snapshot_reconciliations is not None: @@ -1680,11 +1396,10 @@ async def _resolve_approval_responses( tools: list[Any], agent: SupportsAgentRun, run_kwargs: dict[str, Any], - pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] | None = None, thread_id: str = "", validated_approved_responses: list[Content] | None = None, *, - lifecycle: ApprovalLifecycle | None = None, + lifecycle: ApprovalLifecycle, authorized_executions: dict[str, AuthorizedExecution] | None = None, forwarded_executions: ( dict[str, tuple[ForwardedPendingToolTransitionOwner, AuthorizedExecution, Content]] | None @@ -1700,10 +1415,6 @@ async def _resolve_approval_responses( tools: List of available tools agent: The agent instance (to get client and config) run_kwargs: Kwargs for tool execution - pending_approvals: Server-side registry of pending approval requests. - Keys are ``(thread_id, request_id)``, values are function names. - When provided, every approval response is validated against this - registry to prevent bypass, function name spoofing, and replay. thread_id: The conversation thread ID used to scope registry keys. validated_approved_responses: Optional collector for validated local approval responses, including controls removed because the matching @@ -1727,142 +1438,107 @@ async def _resolve_approval_responses( pending_local_response_content_ids: set[int] | None = None validated_forwarded_approvals: list[Content] = [] response_content_ids_to_strip: set[int] = set() - if pending_approvals is not None: - valid_response_content_ids = set() - pending_local_response_content_ids = set() - - def matches_pending_entry(candidate: PendingApprovalEntry | None, expected: PendingApprovalEntry) -> bool: - if isinstance(candidate, str) and isinstance(expected, str): - return candidate == expected - return candidate is expected - - pending_response_groups: dict[tuple[str, object], tuple[PendingApprovalEntry, list[Content]]] = {} - for response in approval_responses: - resp_id = response.id - if resp_id is None: - continue - registry_key = _pending_approval_key(thread_id, resp_id) - id_entry = pending_approvals.get(registry_key) - function_call_id = response.function_call.call_id if response.function_call else None - call_registry_key = ( - _pending_approval_key(thread_id, function_call_id) if function_call_id is not None else None + valid_response_content_ids = set() + pending_local_response_content_ids = set() + pending_response_groups: dict[object, tuple[ApprovalOccurrence, list[Content]]] = {} + for response in approval_responses: + resp_id = response.id + function_call_id = response.function_call.call_id if response.function_call else None + pending_entry = None + for alias in (resp_id, function_call_id): + if alias is not None: + pending_entry = lifecycle.occurrence_for_alias(thread_id=thread_id, interrupt_id=str(alias)) + if pending_entry is not None: + break + if pending_entry is None: + if not _is_hosted_tool_approval(response): + logger.warning("Rejected approval response id=%s: no matching approval occurrence", resp_id) + response_content_ids_to_strip.add(id(response)) + continue + group = pending_response_groups.get(pending_entry.identity) + if group is None: + pending_response_groups[pending_entry.identity] = (pending_entry, [response]) + else: + group[1].append(response) + + for pending_entry, responses in pending_response_groups.values(): + pending_name = pending_entry.name + # The canonical AG-UI approval id may be the provider call id, which can + # be reused by a later call occurrence, while provider request ids may + # alias that same pending entry. Only the latest response across every + # trusted alias can answer the current entry; earlier responses are + # stale replay controls and must not authorize a malformed fresh one. + primary_response = responses[-1] + response_content_ids_to_strip.update(id(response) for response in responses[:-1]) + resp_id = primary_response.id + id_entry = ( + lifecycle.occurrence_for_alias(thread_id=thread_id, interrupt_id=str(resp_id)) + if resp_id is not None + else None + ) + if id_entry is not pending_entry: + logger.warning( + "Rejected approval response id=%s: no matching pending approval request", + resp_id, ) - call_entry = pending_approvals.get(call_registry_key) if call_registry_key is not None else None - pending_entry = id_entry or call_entry - if pending_entry is None: - if not _is_hosted_tool_approval(response): - logger.warning( - "Rejected approval response id=%s: no matching pending approval request", - resp_id, - ) - response_content_ids_to_strip.add(id(response)) - continue + response_content_ids_to_strip.add(id(primary_response)) + continue + function_call_id = primary_response.function_call.call_id if primary_response.function_call else None + if function_call_id != pending_entry.identity.call_id: + logger.warning( + "Rejected approval response id=%s: function call id mismatch (response=%s)", + resp_id, + function_call_id, + ) + response_content_ids_to_strip.add(id(primary_response)) + continue + response_name = primary_response.function_call.name if primary_response.function_call else None + if response_name != pending_name: + logger.warning( + "Rejected approval response id=%s: function name mismatch (response=%s, pending=%s)", + resp_id, + response_name, + pending_name, + ) + response_content_ids_to_strip.add(id(primary_response)) + continue + pending_arguments = pending_entry.arguments + response_arguments = canonical_function_arguments(primary_response.function_call) + if not _approval_arguments_match_pending(pending_arguments, response_arguments): + logger.warning("Rejected approval response id=%s: function arguments mismatch", resp_id) + response_content_ids_to_strip.add(id(primary_response)) + continue - group_key: tuple[str, object] - if isinstance(pending_entry, str): - group_key = ("legacy", call_registry_key if call_entry is not None else registry_key) + server_label = pending_entry.server_label + intent: AuthorizedExecution | None = None + if primary_response.function_call is not None: + if server_label: + primary_response.function_call.additional_properties["server_label"] = server_label else: - group_key = ("entry", id(pending_entry)) - group = pending_response_groups.get(group_key) - if group is None: - pending_response_groups[group_key] = (pending_entry, [response]) - else: - group[1].append(response) - - for pending_entry, responses in pending_response_groups.values(): - pending_name = _pending_approval_name(pending_entry) - # The canonical AG-UI approval id may be the provider call id, which can - # be reused by a later call occurrence, while provider request ids may - # alias that same pending entry. Only the latest response across every - # trusted alias can answer the current entry; earlier responses are - # stale replay controls and must not authorize a malformed fresh one. - primary_response = responses[-1] - response_content_ids_to_strip.update(id(response) for response in responses[:-1]) - resp_id = primary_response.id - primary_registry_key = _pending_approval_key(thread_id, resp_id) if resp_id is not None else None - id_entry = pending_approvals.get(primary_registry_key) if primary_registry_key is not None else None - if not matches_pending_entry(id_entry, pending_entry): - logger.warning( - "Rejected approval response id=%s: no matching pending approval request", - resp_id, - ) - response_content_ids_to_strip.add(id(primary_response)) - continue - function_call_id = primary_response.function_call.call_id if primary_response.function_call else None - call_registry_key = ( - _pending_approval_key(thread_id, function_call_id) if function_call_id is not None else None - ) - call_entry = pending_approvals.get(call_registry_key) if call_registry_key is not None else None - if not isinstance(pending_entry, str) and not matches_pending_entry(call_entry, pending_entry): - logger.warning( - "Rejected approval response id=%s: function call id mismatch (response=%s)", - resp_id, - function_call_id, - ) - response_content_ids_to_strip.add(id(primary_response)) - continue - response_name = primary_response.function_call.name if primary_response.function_call else None - if response_name != pending_name: - logger.warning( - "Rejected approval response id=%s: function name mismatch (response=%s, pending=%s)", - resp_id, - response_name, - pending_name, - ) - response_content_ids_to_strip.add(id(primary_response)) - continue - pending_arguments = _pending_approval_arguments(pending_entry) - response_arguments = canonical_function_arguments(primary_response.function_call) - if not _approval_arguments_match_pending(pending_arguments, response_arguments): - logger.warning("Rejected approval response id=%s: function arguments mismatch", resp_id) + primary_response.function_call.additional_properties.pop("server_label", None) + if ( + primary_response.approved + and lifecycle is not None + and authorized_executions is not None + and primary_response.function_call is not None + ): + call_id = primary_response.function_call.call_id or primary_response.id or "" + intent = authorized_executions.get(call_id) + if intent is None: + logger.warning("Approval remains pending because no transition owner can act for call_id=%s.", call_id) response_content_ids_to_strip.add(id(primary_response)) continue - - server_label = _pending_approval_server_label(pending_entry) - intent: AuthorizedExecution | None = None - if primary_response.function_call is not None: - if server_label: - primary_response.function_call.additional_properties["server_label"] = server_label - else: - primary_response.function_call.additional_properties.pop("server_label", None) - if ( - primary_response.approved - and lifecycle is not None - and authorized_executions is not None - and primary_response.function_call is not None - ): - call_id = primary_response.function_call.call_id or primary_response.id or "" - intent = authorized_executions.get(call_id) - if intent is None: - logger.warning( - "Approval remains pending because no transition owner can act for call_id=%s.", call_id - ) - response_content_ids_to_strip.add(id(primary_response)) - continue - valid_response_content_ids.add(id(primary_response)) - if ( - primary_response.approved - and intent is not None - and intent.owner in {ApprovalExecutionOwner.HOSTED, ApprovalExecutionOwner.DEFERRED} - ): - validated_forwarded_approvals.append(primary_response) - if not server_label: - pending_local_response_content_ids.add(id(primary_response)) - _consume_pending_approval_entry( - pending_approvals, - thread_id, - pending_entry, - resp_id, - primary_response.function_call.call_id if primary_response.function_call else None, - ) - if validated_approved_responses is not None and primary_response.approved and not server_label: - validated_approved_responses.append(primary_response) - elif validated_approved_responses is not None: - validated_approved_responses.extend( - responses[-1] - for responses in responses_by_id.values() - if responses[-1].approved and not _is_hosted_tool_approval(responses[-1]) - ) + valid_response_content_ids.add(id(primary_response)) + if ( + primary_response.approved + and intent is not None + and intent.owner in {ApprovalExecutionOwner.HOSTED, ApprovalExecutionOwner.DEFERRED} + ): + validated_forwarded_approvals.append(primary_response) + if not server_label: + pending_local_response_content_ids.add(id(primary_response)) + if validated_approved_responses is not None and primary_response.approved and not server_label: + validated_approved_responses.append(primary_response) if response_content_ids_to_strip: filtered_messages: list[Message] = [] @@ -2419,7 +2095,6 @@ async def run_agent_stream( input_data: dict[str, Any], agent: SupportsAgentRun, config: AgentConfig, - pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] | None = None, approval_state_store: InMemoryAGUIApprovalStateStore | None = None, ) -> AsyncGenerator[BaseEvent]: """Run agent and yield AG-UI events. @@ -2431,10 +2106,6 @@ async def run_agent_stream( input_data: AG-UI request data with messages, state, tools, etc. agent: The Agent Framework agent to run config: Agent configuration - pending_approvals: Optional server-side registry of pending approval - requests. Keys are ``(thread_id, request_id)``, values are - function names. When provided, approval responses are validated - against this registry to prevent bypass, spoofing, and replay. approval_state_store: Optional server-side Approval State store used to preserve approval-only middleware state across AG-UI requests. @@ -2449,6 +2120,8 @@ async def run_agent_stream( snapshot_scope = cast(str | None, input_data.get(_SNAPSHOT_SCOPE_INPUT_KEY)) approval_scope = cast(str | None, input_data.get(_APPROVAL_SCOPE_INPUT_KEY)) approval_thread_id = approval_state_thread_id(scope=approval_scope, thread_id=thread_id) + if approval_state_store is None: + approval_state_store = InMemoryAGUIApprovalStateStore() state_schema = cast(dict[str, Any], getattr(config, "state_schema", {}) or {}) predict_state_config = cast(dict[str, dict[str, str]], getattr(config, "predict_state_config", {}) or {}) @@ -2533,10 +2206,9 @@ async def run_agent_stream( approval_resume_messages, handled_resume_ids, cancelled_resume_ids, resume_error = ( _canonical_approval_resume_messages( resume_payload, - pending_approvals, approval_thread_id, expected_interrupt_ids=stored_pending_approval_interrupt_ids or None, - lifecycle=approval_state_store.lifecycle if approval_state_store is not None else None, + lifecycle=approval_state_store.lifecycle, authorized_executions=authorized_executions, retained_results=retained_approval_results, snapshot_reconciliations=approval_snapshot_reconciliations, @@ -2579,7 +2251,7 @@ async def run_agent_stream( yield event yield _build_run_finished_event(run_id=run_id, thread_id=thread_id) return - protected_tool_call_ids = _approval_state_tool_call_ids(pending_approvals, approval_state_store, approval_thread_id) + protected_tool_call_ids = _approval_state_tool_call_ids(approval_state_store, approval_thread_id) messages, snapshot_messages = normalize_agui_input_messages( raw_messages, protected_tool_call_ids=protected_tool_call_ids, @@ -2656,10 +2328,9 @@ async def run_agent_stream( messages.extend( _pop_collected_tool_approval_response_messages( session, - pending_approvals, approval_thread_id, tools_for_execution, - lifecycle=approval_state_store.lifecycle if approval_state_store is not None else None, + lifecycle=approval_state_store.lifecycle, authorized_executions=authorized_executions, ) ) @@ -2669,10 +2340,9 @@ async def run_agent_stream( tools_for_execution, agent, run_kwargs, - pending_approvals, approval_thread_id, validated_approved_responses, - lifecycle=approval_state_store.lifecycle if approval_state_store is not None else None, + lifecycle=approval_state_store.lifecycle, authorized_executions=authorized_executions, forwarded_executions=forwarded_executions, ) @@ -2803,7 +2473,7 @@ async def run_agent_stream( owner.record_outcome(intent, [content], lifecycle=approval_state_store.lifecycle) # Register pending approval requests so we can validate responses later - if content_type == "function_approval_request" and pending_approvals is not None: + if content_type == "function_approval_request": if content.id and content.function_call and content.function_call.name: canonical_interrupt_id = content.function_call.call_id or content.id provider_approval_thread_id = approval_state_thread_id( @@ -2816,39 +2486,27 @@ async def run_agent_stream( str(content.id), str(canonical_interrupt_id) if canonical_interrupt_id else None, ) - if approval_state_store is not None: - execution_owner = _function_call_execution_owner( - content.function_call, - tools, - has_deferred_owner=_TOOL_APPROVAL_STATE_KEY in session.state, - ) - registration_kwargs = { - "thread_ids": [approval_thread_id, provider_approval_thread_id], - "name": content.function_call.name, - "arguments": canonical_function_arguments(content.function_call) or "{}", - "request_id": str(content.id), - "interrupt_id": str(canonical_interrupt_id), - "already_approved_requests": already_approved_requests, - } - if execution_owner is ApprovalExecutionOwner.HOSTED: - approval_state_store.register_hosted(server_label=server_label, **registration_kwargs) - elif execution_owner is ApprovalExecutionOwner.DEFERRED: - approval_state_store.register_deferred(**registration_kwargs) - elif execution_owner is ApprovalExecutionOwner.LOCAL: - approval_state_store.register_local(**registration_kwargs) - else: - approval_state_store.register_unowned(**registration_kwargs) + execution_owner = _function_call_execution_owner( + content.function_call, + tools, + has_deferred_owner=_TOOL_APPROVAL_STATE_KEY in session.state, + ) + registration_kwargs = { + "thread_ids": [approval_thread_id, provider_approval_thread_id], + "name": content.function_call.name, + "arguments": canonical_function_arguments(content.function_call) or "{}", + "request_id": str(content.id), + "interrupt_id": str(canonical_interrupt_id), + "already_approved_requests": already_approved_requests, + } + if execution_owner is ApprovalExecutionOwner.HOSTED: + approval_state_store.register_hosted(server_label=server_label, **registration_kwargs) + elif execution_owner is ApprovalExecutionOwner.DEFERRED: + approval_state_store.register_deferred(**registration_kwargs) + elif execution_owner is ApprovalExecutionOwner.LOCAL: + approval_state_store.register_local(**registration_kwargs) else: - _register_pending_approval( - pending_approvals, - [approval_thread_id, provider_approval_thread_id], - content.function_call.name, - canonical_function_arguments(content.function_call), - request_id=str(content.id), - interrupt_id=str(canonical_interrupt_id), - already_approved_requests=already_approved_requests, - server_label=server_label, - ) + approval_state_store.register_unowned(**registration_kwargs) else: logger.warning( "Approval request not registered: missing id=%s, function_call=%s, or function name", diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py b/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py index a48366669d..7d8681878a 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py @@ -170,6 +170,9 @@ class ApprovalOccurrence: name: str arguments: str owner: ApprovalExecutionOwner + aliases: tuple[str, ...] = () + already_approved_requests: tuple[dict[str, Any], ...] = () + server_label: str | None = None idempotency_key: str | None = None status: ApprovalStatus = ApprovalStatus.PENDING replayable_results: list[ReplayableToolResult] = field(default_factory=list) @@ -250,6 +253,9 @@ def register_local( call_id: str, name: str, arguments: str, + aliases: list[str] | None = None, + already_approved_requests: list[dict[str, Any]] | None = None, + server_label: str | None = None, idempotency_key: str | None = None, ) -> ApprovalOccurrence: """Register one server-generated local approval occurrence.""" @@ -259,6 +265,9 @@ def register_local( call_id=call_id, name=name, arguments=arguments, + aliases=aliases, + already_approved_requests=already_approved_requests, + server_label=server_label, idempotency_key=idempotency_key, ) @@ -270,6 +279,9 @@ def register_local_aliases( call_id: str, name: str, arguments: str, + aliases: list[str] | None = None, + already_approved_requests: list[dict[str, Any]] | None = None, + server_label: str | None = None, idempotency_key: str | None = None, ) -> ApprovalOccurrence: """Register one occurrence under its trusted scoped-thread aliases.""" @@ -280,6 +292,9 @@ def register_local_aliases( name=name, arguments=arguments, owner=ApprovalExecutionOwner.LOCAL, + aliases=aliases, + already_approved_requests=already_approved_requests, + server_label=server_label, idempotency_key=idempotency_key, ) @@ -291,6 +306,9 @@ def register_hosted( call_id: str, name: str, arguments: str, + aliases: list[str] | None = None, + already_approved_requests: list[dict[str, Any]] | None = None, + server_label: str | None = None, idempotency_key: str | None = None, ) -> ApprovalOccurrence: """Register one server-generated hosted approval occurrence.""" @@ -300,6 +318,9 @@ def register_hosted( call_id=call_id, name=name, arguments=arguments, + aliases=aliases, + already_approved_requests=already_approved_requests, + server_label=server_label, idempotency_key=idempotency_key, ) @@ -311,6 +332,9 @@ def register_hosted_aliases( call_id: str, name: str, arguments: str, + aliases: list[str] | None = None, + already_approved_requests: list[dict[str, Any]] | None = None, + server_label: str | None = None, idempotency_key: str | None = None, ) -> ApprovalOccurrence: """Register one hosted occurrence under its trusted scoped-thread aliases.""" @@ -321,6 +345,9 @@ def register_hosted_aliases( name=name, arguments=arguments, owner=ApprovalExecutionOwner.HOSTED, + aliases=aliases, + already_approved_requests=already_approved_requests, + server_label=server_label, idempotency_key=idempotency_key, ) @@ -332,6 +359,9 @@ def register_unowned( call_id: str, name: str, arguments: str, + aliases: list[str] | None = None, + already_approved_requests: list[dict[str, Any]] | None = None, + server_label: str | None = None, ) -> ApprovalOccurrence: """Register an occurrence that has no executable transition owner.""" return self.register_unowned_aliases( @@ -340,6 +370,9 @@ def register_unowned( call_id=call_id, name=name, arguments=arguments, + aliases=aliases, + already_approved_requests=already_approved_requests, + server_label=server_label, ) def register_deferred( @@ -350,6 +383,9 @@ def register_deferred( call_id: str, name: str, arguments: str, + aliases: list[str] | None = None, + already_approved_requests: list[dict[str, Any]] | None = None, + server_label: str | None = None, idempotency_key: str | None = None, ) -> ApprovalOccurrence: """Register one occurrence owned by the in-run transition pipeline.""" @@ -359,6 +395,9 @@ def register_deferred( call_id=call_id, name=name, arguments=arguments, + aliases=aliases, + already_approved_requests=already_approved_requests, + server_label=server_label, idempotency_key=idempotency_key, ) @@ -370,6 +409,9 @@ def register_deferred_aliases( call_id: str, name: str, arguments: str, + aliases: list[str] | None = None, + already_approved_requests: list[dict[str, Any]] | None = None, + server_label: str | None = None, idempotency_key: str | None = None, ) -> ApprovalOccurrence: """Register one deferred occurrence under its trusted scoped-thread aliases.""" @@ -380,6 +422,9 @@ def register_deferred_aliases( name=name, arguments=arguments, owner=ApprovalExecutionOwner.DEFERRED, + aliases=aliases, + already_approved_requests=already_approved_requests, + server_label=server_label, idempotency_key=idempotency_key, ) @@ -391,6 +436,9 @@ def register_unowned_aliases( call_id: str, name: str, arguments: str, + aliases: list[str] | None = None, + already_approved_requests: list[dict[str, Any]] | None = None, + server_label: str | None = None, ) -> ApprovalOccurrence: """Register an unowned occurrence under its trusted scoped-thread aliases.""" return self._register_aliases( @@ -400,6 +448,9 @@ def register_unowned_aliases( name=name, arguments=arguments, owner=ApprovalExecutionOwner.UNAVAILABLE, + aliases=aliases, + already_approved_requests=already_approved_requests, + server_label=server_label, ) @_serialized_registration @@ -412,18 +463,24 @@ def _register_aliases( name: str, arguments: str, owner: ApprovalExecutionOwner, + aliases: list[str] | None = None, + already_approved_requests: list[dict[str, Any]] | None = None, + server_label: str | None = None, idempotency_key: str | None = None, ) -> ApprovalOccurrence: self._purge_expired_terminal() if idempotency_key == "": raise ValueError("An execution idempotency key cannot be empty.") unique_thread_ids = tuple(dict.fromkeys(thread_ids)) + alias_values: list[str] = [interrupt_id, *(aliases or [])] + occurrence_aliases = tuple(dict.fromkeys(alias_values)) if not unique_thread_ids: raise ValueError("An approval occurrence requires at least one scoped thread identity.") existing_identities = { identity for thread_id in unique_thread_ids - if (identity := self._pending_by_interrupt.get((thread_id, interrupt_id))) is not None + for alias in occurrence_aliases + if (identity := self._pending_by_interrupt.get((thread_id, alias))) is not None } if len(existing_identities) > 1: raise ValueError("Approval aliases resolve to different pending occurrences.") @@ -437,11 +494,15 @@ def _register_aliases( or occurrence.arguments != arguments or occurrence.owner is not owner or occurrence.idempotency_key != idempotency_key + or occurrence.already_approved_requests != tuple(already_approved_requests or ()) + or occurrence.server_label != server_label ): raise ValueError("Approval alias conflicts with an existing pending occurrence.") occurrence.thread_ids = tuple(dict.fromkeys((*occurrence.thread_ids, *unique_thread_ids))) + occurrence.aliases = tuple(dict.fromkeys((*occurrence.aliases, *occurrence_aliases))) for thread_id in occurrence.thread_ids: - self._pending_by_interrupt[(thread_id, interrupt_id)] = occurrence.identity + for alias in occurrence.aliases: + self._pending_by_interrupt[(thread_id, alias)] = occurrence.identity return occurrence if len(self._occurrences) >= self._max_entries: @@ -459,11 +520,15 @@ def _register_aliases( name=name, arguments=arguments, owner=owner, + aliases=occurrence_aliases, + already_approved_requests=tuple(already_approved_requests or ()), + server_label=server_label, idempotency_key=idempotency_key, ) self._occurrences[identity] = occurrence for thread_id in unique_thread_ids: - self._pending_by_interrupt[(thread_id, interrupt_id)] = identity + for alias in occurrence.aliases: + self._pending_by_interrupt[(thread_id, alias)] = identity self._locks_by_identity[identity] = RLock() self._emit_event("registration", occurrence) return occurrence @@ -481,6 +546,33 @@ def decision_context(self, *, thread_id: str, interrupt_id: str) -> tuple[str, s occurrence = self._occurrences[identity] return occurrence.name, occurrence.arguments + def pending_occurrence(self, *, thread_id: str, interrupt_id: str) -> ApprovalOccurrence | None: + """Return pending server-owned state for one trusted interrupt alias.""" + self._purge_expired_terminal() + identity = self._pending_by_interrupt.get((thread_id, interrupt_id)) + return self._occurrences[identity] if identity is not None else None + + def occurrence_for_alias(self, *, thread_id: str, interrupt_id: str) -> ApprovalOccurrence | None: + """Return retained server-owned state for one trusted interrupt alias.""" + self._purge_expired_terminal() + key = (thread_id, interrupt_id) + identity = self._pending_by_interrupt.get(key) or self._terminal_by_interrupt.get(key) + return self._occurrences[identity] if identity is not None else None + + def occurrences_for_thread(self, *, thread_id: str) -> tuple[ApprovalOccurrence, ...]: + """Return retained occurrences owned by one scoped thread.""" + self._purge_expired_terminal() + return tuple(occurrence for occurrence in self._occurrences.values() if thread_id in occurrence.thread_ids) + + def pending_interrupt_ids(self, *, thread_id: str) -> set[str]: + """Return canonical interrupt identities with pending authority for one thread.""" + self._purge_expired_terminal() + return { + occurrence.identity.interrupt_id + for occurrence in self._occurrences.values() + if thread_id in occurrence.thread_ids and occurrence.status is ApprovalStatus.PENDING + } + def reconcile_snapshot( self, *, @@ -708,8 +800,9 @@ def _remove_pending_aliases(self, occurrence: ApprovalOccurrence) -> None: occurrence.terminal_at = self._clock() with self._index_lock: for thread_id in occurrence.thread_ids: - self._pending_by_interrupt.pop((thread_id, occurrence.identity.interrupt_id), None) - self._terminal_by_interrupt[(thread_id, occurrence.identity.interrupt_id)] = occurrence.identity + for alias in occurrence.aliases: + self._pending_by_interrupt.pop((thread_id, alias), None) + self._terminal_by_interrupt[(thread_id, alias)] = occurrence.identity def _purge_expired_terminal(self) -> None: with self._index_lock: @@ -724,9 +817,10 @@ def _purge_expired_terminal(self) -> None: self._occurrences.pop(occurrence.identity, None) self._locks_by_identity.pop(occurrence.identity, None) for thread_id in occurrence.thread_ids: - key = (thread_id, occurrence.identity.interrupt_id) - if self._terminal_by_interrupt.get(key) == occurrence.identity: - self._terminal_by_interrupt.pop(key, None) + for alias in occurrence.aliases: + key = (thread_id, alias) + if self._terminal_by_interrupt.get(key) == occurrence.identity: + self._terminal_by_interrupt.pop(key, None) def _locks_for_batch(self, *, thread_id: str, interrupt_ids: list[str]) -> tuple[RLock, ...]: with self._index_lock: diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py b/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py index 913f129911..7efa39a657 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py @@ -4,7 +4,7 @@ from __future__ import annotations -from collections import OrderedDict +import copy from typing import Any from ._approval_lifecycle import ApprovalCapacityError, ApprovalExecutionOwner, ApprovalLifecycle @@ -59,8 +59,7 @@ def __init__( if max_entries < 1: raise ValueError("max_entries must be greater than 0.") self.max_entries = max_entries - self.pending_approvals: OrderedDict[tuple[str, str], Any] = OrderedDict() - self.tool_approval_states: OrderedDict[str, dict[str, Any]] = OrderedDict() + self._tool_approval_states: dict[str, dict[str, Any]] = {} self.lifecycle = ApprovalLifecycle( max_entries=max_entries, terminal_retention_seconds=terminal_retention_seconds, @@ -167,18 +166,6 @@ def _register( server_label: str | None, owner: ApprovalExecutionOwner, ) -> None: - entry: dict[str, Any] = { - "name": name, - "arguments": arguments, - "request_id": request_id, - "interrupt_id": interrupt_id, - } - if already_approved_requests: - entry["already_approved_requests"] = already_approved_requests - if server_label: - entry["server_label"] = server_label - entry["execution_owner"] = owner.value - unique_thread_ids = list(dict.fromkeys(thread_ids)) if owner is ApprovalExecutionOwner.HOSTED: register_aliases = self.lifecycle.register_hosted_aliases @@ -194,19 +181,26 @@ def _register( call_id=interrupt_id, name=name, arguments=arguments, + aliases=[request_id], + already_approved_requests=already_approved_requests, + server_label=server_label, ) - for thread_id in unique_thread_ids: - aliases = {(thread_id, request_id), (thread_id, interrupt_id)} - replaced_entries = {id(existing) for key, existing in self.pending_approvals.items() if key in aliases} - for key, existing in list(self.pending_approvals.items()): - if key in aliases or id(existing) in replaced_entries: - self.pending_approvals.pop(key, None) - for key in aliases: - self.pending_approvals[key] = entry def set_tool_approval_state(self, thread_id: str, state: dict[str, Any]) -> None: """Store approval middleware state without evicting another active thread.""" - if thread_id not in self.tool_approval_states and len(self.tool_approval_states) >= self.max_entries: + if thread_id not in self._tool_approval_states and len(self._tool_approval_states) >= self.max_entries: raise ApprovalCapacityError("Approval state capacity is exhausted by protected occurrences.") - self.tool_approval_states[thread_id] = state - self.tool_approval_states.move_to_end(thread_id) + self._tool_approval_states[thread_id] = copy.deepcopy(state) + + def get_tool_approval_state(self, thread_id: str) -> dict[str, Any] | None: + """Return an isolated copy of server-owned middleware approval state.""" + state = self._tool_approval_states.get(thread_id) + return copy.deepcopy(state) if state is not None else None + + def delete_tool_approval_state(self, thread_id: str) -> None: + """Delete server-owned middleware approval state for one scoped thread.""" + self._tool_approval_states.pop(thread_id, None) + + def has_tool_approval_state(self, thread_id: str) -> bool: + """Return whether middleware approval state exists for one scoped thread.""" + return thread_id in self._tool_approval_states diff --git a/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py b/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py index 79899e1998..f6992752e8 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py +++ b/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py @@ -1050,7 +1050,7 @@ async def stream_fn_approval( if getattr(e, "type", None) == "CUSTOM" and getattr(e, "name", None) == "function_approval_request" ] assert len(approval_events) == 1, "Expected one approval request event" - assert any("call_sens_001" in k for k in wrapper._pending_approvals) + assert wrapper._approval_state_store.lifecycle.pending_occurrence(thread_id=thread_id, interrupt_id="call_sens_001") # --- Turn 2: legitimate approval --- async def stream_fn_post_approval( @@ -1064,7 +1064,7 @@ async def stream_fn_post_approval( instructions="Test", tools=[sensitive_action], ) - # Reuse the same wrapper (same _pending_approvals) with a new agent for Turn 2 + # Reuse the same wrapper with its server-owned Approval State for Turn 2. wrapper.agent = agent2 turn2_input: dict[str, Any] = { @@ -1078,7 +1078,9 @@ async def stream_fn_post_approval( events2.append(event) assert call_count == 1, "Tool should have been executed once" - assert not any("call_sens_001" in k for k in wrapper._pending_approvals), "Pending approval should be consumed" + assert not wrapper._approval_state_store.lifecycle.pending_occurrence( + thread_id=thread_id, interrupt_id="call_sens_001" + ) # --- Turn 3: replay attempt with the same approval ID --- call_count = 0 # reset @@ -1149,8 +1151,12 @@ async def approval_stream( async for _ in wrapper.run({"thread_id": "client-thread", "messages": [{"role": "user", "content": "do it"}]}): pass - assert ("client-thread", "call_sensitive") in wrapper._pending_approvals - assert ("provider-conversation", "call_sensitive") in wrapper._pending_approvals + assert wrapper._approval_state_store.lifecycle.pending_occurrence( + thread_id="client-thread", interrupt_id="call_sensitive" + ) + assert wrapper._approval_state_store.lifecycle.pending_occurrence( + thread_id="provider-conversation", interrupt_id="call_sensitive" + ) async def completion_stream( messages: MutableSequence[Message], options: ChatOptions, **kwargs: Any @@ -1175,8 +1181,12 @@ def approval_input(thread_id: str) -> dict[str, Any]: pass assert execution_count == 1 - assert ("client-thread", "call_sensitive") not in wrapper._pending_approvals - assert ("provider-conversation", "call_sensitive") not in wrapper._pending_approvals + assert not wrapper._approval_state_store.lifecycle.pending_occurrence( + thread_id="client-thread", interrupt_id="call_sensitive" + ) + assert not wrapper._approval_state_store.lifecycle.pending_occurrence( + thread_id="provider-conversation", interrupt_id="call_sensitive" + ) replay_thread_id = "provider-conversation" if resume_thread_id == "client-thread" else "client-thread" retry_events = [event async for event in wrapper.run(approval_input(replay_thread_id))] @@ -1258,7 +1268,7 @@ async def stream_fn_approval( async for event in wrapper.run({"thread_id": thread_id, "messages": [{"role": "user", "content": "do safe"}]}): events1.append(event) - assert any("call_safe_001" in k for k in wrapper._pending_approvals) + assert wrapper._approval_state_store.lifecycle.pending_occurrence(thread_id=thread_id, interrupt_id="call_safe_001") # Turn 2: try to approve with a different function name (function name spoofing) async def stream_fn_post( @@ -1297,9 +1307,9 @@ async def stream_fn_post( events2.append(event) assert not tool_executed, "Function name spoofing should be blocked" - assert any("call_safe_001" in k for k in wrapper._pending_approvals), ( - "Pending approval should be preserved after mismatch for legitimate retry" - ) + assert wrapper._approval_state_store.lifecycle.pending_occurrence( + thread_id=thread_id, interrupt_id="call_safe_001" + ), "Pending approval should be preserved after mismatch for legitimate retry" async def test_approval_bypass_via_fabricated_tool_result_is_blocked(streaming_chat_client_stub): @@ -1496,7 +1506,9 @@ async def stream_fn_approval( async for event in wrapper.run({"thread_id": thread_id, "messages": [{"role": "user", "content": "update"}]}): events1.append(event) - assert any("call_update_001" in k for k in wrapper._pending_approvals) + assert wrapper._approval_state_store.lifecycle.pending_occurrence( + thread_id=thread_id, interrupt_id="call_update_001" + ) async def stream_fn_post( messages: MutableSequence[Message], options: ChatOptions, **kwargs: Any @@ -1534,9 +1546,9 @@ async def stream_fn_post( events2.append(event) assert executed_args == [] - assert any("call_update_001" in k for k in wrapper._pending_approvals), ( - "Pending approval should be preserved after argument mismatch for legitimate retry" - ) + assert wrapper._approval_state_store.lifecycle.pending_occurrence( + thread_id=thread_id, interrupt_id="call_update_001" + ), "Pending approval should be preserved after argument mismatch for legitimate retry" async def test_state_update_end_to_end_via_real_tool_invocation(streaming_chat_client_stub): diff --git a/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py b/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py index 7d89d40a8f..d7254b1b51 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py +++ b/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py @@ -66,6 +66,36 @@ async def execute_authorized_call() -> list[Content]: assert [result.content.result for result in outcome.replayable_results] == ["Sunny"] +def test_registration_keeps_trusted_aliases_and_projection_metadata_in_lifecycle() -> None: + lifecycle = ApprovalLifecycle() + sibling = { + "type": "function_approval_request", + "id": "sibling-request", + "function_call": { + "type": "function_call", + "call_id": "sibling-call", + "name": "write_record", + "arguments": '{"value":"second"}', + }, + } + + occurrence = lifecycle.register_hosted( + thread_id="thread-1", + interrupt_id="call-1", + call_id="call-1", + name="write_record", + arguments='{"value":"first"}', + aliases=["request-1"], + already_approved_requests=[sibling], + server_label="hosted-tools", + ) + + assert lifecycle.pending_occurrence(thread_id="thread-1", interrupt_id="request-1") is occurrence + assert lifecycle.pending_interrupt_ids(thread_id="thread-1") == {"call-1"} + assert occurrence.already_approved_requests == (sibling,) + assert occurrence.server_label == "hosted-tools" + + def test_active_occurrence_is_not_evicted_when_capacity_is_exhausted() -> None: """Storage pressure fails explicitly instead of discarding pending authority.""" lifecycle = ApprovalLifecycle(max_entries=1) diff --git a/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py b/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py index c02110760b..032dbbc21b 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py +++ b/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py @@ -1,6 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. -"""Tests for TOOL_CALL_RESULT event emission on approval resume flows.""" +"""Public event-stream tests for approval result projection.""" from __future__ import annotations @@ -11,67 +11,13 @@ from conftest import StubAgent # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports] from agent_framework_ag_ui._agent import AgentConfig -from agent_framework_ag_ui._agent_run import PendingApprovalEntry, PendingApprovalKey, run_agent_stream -from agent_framework_ag_ui._approval_lifecycle import ResumeDecision +from agent_framework_ag_ui._agent_run import run_agent_stream from agent_framework_ag_ui._approval_state import InMemoryAGUIApprovalStateStore -async def _run_with_registered_approval_state( - input_data: dict[str, Any], - agent: StubAgent, - config: AgentConfig, -) -> list[Any]: - """Cross the lifecycle seam with server-owned state for approval fixtures.""" - thread_id = str(input_data["thread_id"]) - calls: dict[str, dict[str, Any]] = {} - decisions: list[dict[str, Any]] = [] - messages: list[dict[str, Any]] = [] - for message in input_data["messages"]: - if message.get("role") == "assistant": - for tool_call in message.get("tool_calls", []): - calls[str(tool_call["id"])] = tool_call["function"] - if message.get("role") == "tool": - decision = json.loads(message["content"]) - if isinstance(decision, dict) and "accepted" in decision: - decisions.append( - { - "interruptId": str(message["toolCallId"]), - "status": "resolved", - "payload": decision, - } - ) - continue - messages.append(message) - - store = InMemoryAGUIApprovalStateStore() - for decision in decisions: - call_id = str(decision["interruptId"]) - function = calls[call_id] - arguments = json.dumps(json.loads(function["arguments"]), sort_keys=True, separators=(",", ":")) - store.register_local( - thread_ids=[thread_id], - name=str(function["name"]), - arguments=arguments, - request_id=call_id, - interrupt_id=call_id, - ) - - events: list[Any] = [] - async for event in run_agent_stream( - {**input_data, "messages": messages, "resume": decisions}, - agent, - config, - pending_approvals=store.pending_approvals, - approval_state_store=store, - ): - events.append(event) - return events - - -def _make_weather_tool() -> FunctionTool: - """Create a real executable weather tool with approval_mode='always_require'.""" - +def _weather_tool(executions: list[str]) -> FunctionTool: def get_weather(city: str) -> str: + executions.append(city) return f"Sunny in {city}" return FunctionTool( @@ -82,89 +28,29 @@ def get_weather(city: str) -> str: ) -async def test_approval_resume_emits_tool_call_result() -> None: - """After approving a tool call, the resume stream should contain a TOOL_CALL_RESULT event. - - The message format follows the AG-UI approval pattern: - - assistant message with tool_calls - - tool message with {"accepted": true} content and toolCallId - """ - tool_name = "get_weather" - call_id = "call_abc123" - weather_tool = _make_weather_tool() - - agent = StubAgent( - updates=[AgentResponseUpdate(contents=[Content.from_text(text="The weather is sunny.")], role="assistant")], - default_options={"tools": [weather_tool]}, - ) - config = AgentConfig() - - # Build resume messages: user query, assistant tool call, approval response - resume_messages: list[dict[str, Any]] = [ - {"role": "user", "content": "What's the weather in Seattle?"}, - { - "role": "assistant", - "content": "", - "tool_calls": [ - { - "id": call_id, - "type": "function", - "function": { - "name": tool_name, - "arguments": json.dumps({"city": "Seattle"}), - }, - } - ], - }, - { - "role": "tool", - "content": json.dumps({"accepted": True}), - "toolCallId": call_id, - }, - ] - - input_data: dict[str, Any] = { - "thread_id": "thread-approval-result", - "run_id": "run-resume", - "messages": resume_messages, - } - - events = await _run_with_registered_approval_state(input_data, agent, config) - - event_types = [getattr(e, "type", None) for e in events] - - assert "RUN_STARTED" in event_types, f"Expected RUN_STARTED, got types: {event_types}" - assert "RUN_FINISHED" in event_types, f"Expected RUN_FINISHED, got types: {event_types}" - - # TOOL_CALL_RESULT must be present for the approved tool - tool_result_events = [e for e in events if getattr(e, "type", None) == "TOOL_CALL_RESULT"] - - assert len(tool_result_events) > 0, ( - f"Expected at least one TOOL_CALL_RESULT event for the approved tool, " - f"but found none. Event types in stream: {event_types}" - ) - - result_event = tool_result_events[0] - assert result_event.tool_call_id == call_id, ( - f"Expected TOOL_CALL_RESULT with tool_call_id={call_id}, got tool_call_id={result_event.tool_call_id}" - ) - # Verify the result contains the actual tool execution output - assert result_event.content == "Sunny in Seattle" - - -async def test_approval_resume_result_has_content() -> None: - """TOOL_CALL_RESULT event from an approved tool should contain the execution result.""" - tool_name = "get_weather" - call_id = "call_content_check" - weather_tool = _make_weather_tool() - +async def _run_resume( + *, + thread_id: str, + calls: list[tuple[str, str]], + decisions: list[tuple[str, bool]], + executions: list[str], +) -> list[Any]: + tool = _weather_tool(executions) agent = StubAgent( updates=[AgentResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant")], - default_options={"tools": [weather_tool]}, + default_options={"tools": [tool]}, ) - config = AgentConfig() + store = InMemoryAGUIApprovalStateStore() + for call_id, city in calls: + store.register_local( + thread_ids=[thread_id], + name="get_weather", + arguments=json.dumps({"city": city}, sort_keys=True, separators=(",", ":")), + request_id=call_id, + interrupt_id=call_id, + ) - resume_messages: list[dict[str, Any]] = [ + messages: list[dict[str, Any]] = [ {"role": "user", "content": "Check the weather"}, { "role": "assistant", @@ -173,1218 +59,70 @@ async def test_approval_resume_result_has_content() -> None: { "id": call_id, "type": "function", - "function": { - "name": tool_name, - "arguments": json.dumps({"city": "Portland"}), - }, - } - ], - }, - { - "role": "tool", - "content": json.dumps({"accepted": True}), - "toolCallId": call_id, - }, - ] - - input_data: dict[str, Any] = { - "thread_id": "thread-result-content", - "run_id": "run-resume-2", - "messages": resume_messages, - } - - events = await _run_with_registered_approval_state(input_data, agent, config) - - tool_result_events = [e for e in events if getattr(e, "type", None) == "TOOL_CALL_RESULT"] - assert len(tool_result_events) == 1 - - result_event = tool_result_events[0] - assert result_event.tool_call_id == call_id - assert result_event.role == "tool" - # Verify the result contains the actual tool execution output (string returned directly) - assert result_event.content == "Sunny in Portland" - - -async def test_approval_resume_snapshot_replaces_approval_payload_with_tool_result() -> None: - """Approved HITL tools persist their executed result in MESSAGES_SNAPSHOT for replay.""" - from agent_framework_ag_ui._message_adapters import normalize_agui_input_messages - - call_id = "call_snapshot_replay" - weather_tool = _make_weather_tool() - agent = StubAgent( - updates=[AgentResponseUpdate(contents=[Content.from_text(text="The weather is sunny.")], role="assistant")], - default_options={"tools": [weather_tool]}, - ) - config = AgentConfig() - resume_messages: list[dict[str, Any]] = [ - {"role": "user", "content": "What's the weather in Seattle?"}, - { - "role": "assistant", - "content": "", - "tool_calls": [ - { - "id": call_id, - "type": "function", - "function": { - "name": "get_weather", - "arguments": json.dumps({"city": "Seattle"}), - }, + "function": {"name": "get_weather", "arguments": json.dumps({"city": city})}, } + for call_id, city in calls ], }, - { - "role": "tool", - "content": json.dumps({"accepted": True}), - "toolCallId": call_id, - }, - ] - - events = await _run_with_registered_approval_state( - { - "thread_id": "thread-snapshot-replay", - "run_id": "run-snapshot-replay", - "messages": resume_messages, - }, - agent, - config, - ) - - snapshots = [event.messages for event in events if getattr(event, "type", None) == "MESSAGES_SNAPSHOT"] - assert snapshots - snapshot_messages = [ - message.model_dump(by_alias=True, exclude_none=True) if hasattr(message, "model_dump") else message - for message in snapshots[-1] ] - tool_messages = [message for message in snapshot_messages if message.get("role") == "tool"] - assert any( - message.get("toolCallId") == call_id and message.get("content") == "Sunny in Seattle" - for message in tool_messages - ) - assert not any(message.get("content") == json.dumps({"accepted": True}) for message in tool_messages) - - replay_messages = snapshot_messages + [{"role": "user", "content": "What is the weather now?"}] - provider_messages, _ = normalize_agui_input_messages(replay_messages) - - assert not any( - content.type == "function_approval_response" - for message in provider_messages - for content in message.contents or [] - ) - assert any( - content.type == "function_result" and content.call_id == call_id and content.result == "Sunny in Seattle" - for message in provider_messages - for content in message.contents or [] - ) - - -async def test_no_approval_no_extra_tool_result() -> None: - """When no approval response is present, no extra TOOL_CALL_RESULT events should be emitted.""" - agent = StubAgent(updates=[AgentResponseUpdate(contents=[Content.from_text(text="Hello.")], role="assistant")]) - config = AgentConfig() - - input_data: dict[str, Any] = { - "thread_id": "thread-no-approval", - "run_id": "run-normal", - "messages": [{"role": "user", "content": "Hi"}], - } - events: list[Any] = [] - async for event in run_agent_stream(input_data, agent, config): - events.append(event) - - tool_result_events = [e for e in events if getattr(e, "type", None) == "TOOL_CALL_RESULT"] - assert len(tool_result_events) == 0, f"Unexpected TOOL_CALL_RESULT events: {tool_result_events}" - - -async def test_rejection_does_not_emit_tool_call_result() -> None: - """Rejected tool calls should not produce TOOL_CALL_RESULT events.""" - tool_name = "get_weather" - call_id = "call_rejected" - weather_tool = _make_weather_tool() - - agent = StubAgent( - updates=[AgentResponseUpdate(contents=[Content.from_text(text="OK, I won't check.")], role="assistant")], - default_options={"tools": [weather_tool]}, - ) - config = AgentConfig() - - resume_messages: list[dict[str, Any]] = [ - {"role": "user", "content": "What's the weather?"}, - { - "role": "assistant", - "content": "", - "tool_calls": [ - { - "id": call_id, - "type": "function", - "function": { - "name": tool_name, - "arguments": json.dumps({"city": "Denver"}), - }, - } - ], - }, - { - "role": "tool", - "content": json.dumps({"accepted": False}), - "toolCallId": call_id, - }, - ] - - input_data: dict[str, Any] = { - "thread_id": "thread-rejection", - "run_id": "run-rejected", - "messages": resume_messages, - } - - events = await _run_with_registered_approval_state(input_data, agent, config) - - tool_result_events = [e for e in events if getattr(e, "type", None) == "TOOL_CALL_RESULT"] - assert len(tool_result_events) == 0, ( - f"Expected no TOOL_CALL_RESULT for rejected tool, got {len(tool_result_events)}" - ) - - -def _make_temperature_tool() -> FunctionTool: - """Create a real executable temperature tool with approval_mode='always_require'.""" - - def get_temperature(city: str) -> str: - return f"72F in {city}" - - return FunctionTool( - name="get_temperature", - description="Get the temperature for a city", - func=get_temperature, - approval_mode="always_require", - ) - - -async def test_mixed_approve_reject_emits_only_approved_tool_result() -> None: - """When one tool call is approved and another rejected, only the approved one produces a TOOL_CALL_RESULT event.""" - weather_tool = _make_weather_tool() - temperature_tool = _make_temperature_tool() - approved_call_id = "call_approved" - rejected_call_id = "call_rejected" - - agent = StubAgent( - updates=[AgentResponseUpdate(contents=[Content.from_text(text="Here are the results.")], role="assistant")], - default_options={"tools": [weather_tool, temperature_tool]}, - ) - config = AgentConfig() - - resume_messages: list[dict[str, Any]] = [ - {"role": "user", "content": "Weather and temperature in Seattle?"}, - { - "role": "assistant", - "content": "", - "tool_calls": [ - { - "id": approved_call_id, - "type": "function", - "function": { - "name": "get_weather", - "arguments": json.dumps({"city": "Seattle"}), - }, - }, - { - "id": rejected_call_id, - "type": "function", - "function": { - "name": "get_temperature", - "arguments": json.dumps({"city": "Seattle"}), - }, - }, - ], - }, - { - "role": "tool", - "content": json.dumps({"accepted": True}), - "toolCallId": approved_call_id, - }, - { - "role": "tool", - "content": json.dumps({"accepted": False}), - "toolCallId": rejected_call_id, - }, - ] - - input_data: dict[str, Any] = { - "thread_id": "thread-mixed", - "run_id": "run-mixed", - "messages": resume_messages, - } - - events = await _run_with_registered_approval_state(input_data, agent, config) - - tool_result_events = [e for e in events if getattr(e, "type", None) == "TOOL_CALL_RESULT"] - - # Only the approved tool call should produce a TOOL_CALL_RESULT event - assert len(tool_result_events) == 1, ( - f"Expected exactly 1 TOOL_CALL_RESULT (approved only), got {len(tool_result_events)}" - ) - assert tool_result_events[0].tool_call_id == approved_call_id - assert tool_result_events[0].content == "Sunny in Seattle" - - -async def test_approval_resume_zero_updates_emits_tool_result() -> None: - """When the agent produces zero updates, TOOL_CALL_RESULT events should still be emitted via the fallback path.""" - tool_name = "get_weather" - call_id = "call_zero_updates" - weather_tool = _make_weather_tool() - - agent = StubAgent( - updates=[], - default_options={"tools": [weather_tool]}, - ) - config = AgentConfig() - - resume_messages: list[dict[str, Any]] = [ - {"role": "user", "content": "What's the weather?"}, + async for event in run_agent_stream( { - "role": "assistant", - "content": "", - "tool_calls": [ - { - "id": call_id, - "type": "function", - "function": { - "name": tool_name, - "arguments": json.dumps({"city": "Boston"}), - }, - } + "thread_id": thread_id, + "run_id": "resume-run", + "messages": messages, + "resume": [ + {"interruptId": call_id, "status": "resolved", "payload": {"accepted": accepted}} + for call_id, accepted in decisions ], }, - { - "role": "tool", - "content": json.dumps({"accepted": True}), - "toolCallId": call_id, - }, - ] - - input_data: dict[str, Any] = { - "thread_id": "thread-zero-updates", - "run_id": "run-zero-updates", - "messages": resume_messages, - } - - events = await _run_with_registered_approval_state(input_data, agent, config) - - event_types = [getattr(e, "type", None) for e in events] - assert "RUN_STARTED" in event_types - - tool_result_events = [e for e in events if getattr(e, "type", None) == "TOOL_CALL_RESULT"] - assert len(tool_result_events) == 1, ( - f"Expected 1 TOOL_CALL_RESULT in zero-updates fallback path, got {len(tool_result_events)}" - ) - assert tool_result_events[0].tool_call_id == call_id - assert tool_result_events[0].content == "Sunny in Boston" - - -async def test_resolve_approval_responses_returns_only_approved() -> None: - """_resolve_approval_responses should return only approved results; rejection results go into messages only.""" - from agent_framework import Message - - from agent_framework_ag_ui._agent_run import _resolve_approval_responses - - weather_tool = _make_weather_tool() - temperature_tool = _make_temperature_tool() - approved_call_id = "call_a" - rejected_call_id = "call_r" - - messages: list[Any] = [ - Message(role="user", contents=[Content.from_text(text="Hi")]), - Message( - role="assistant", - contents=[ - Content( - type="function_approval_request", - id=approved_call_id, - function_call=Content( - type="function_call", - name="get_weather", - call_id=approved_call_id, - arguments='{"city": "NYC"}', - ), - ), - Content( - type="function_approval_request", - id=rejected_call_id, - function_call=Content( - type="function_call", - name="get_temperature", - call_id=rejected_call_id, - arguments='{"city": "NYC"}', - ), - ), - ], - ), - Message( - role="user", - contents=[ - Content( - type="function_approval_response", - id=approved_call_id, - approved=True, - function_call=Content( - type="function_call", - name="get_weather", - call_id=approved_call_id, - arguments='{"city": "NYC"}', - ), - ), - Content( - type="function_approval_response", - id=rejected_call_id, - approved=False, - function_call=Content( - type="function_call", - name="get_temperature", - call_id=rejected_call_id, - arguments='{"city": "NYC"}', - ), - ), - ], - ), - ] - - agent = StubAgent( - updates=[], - default_options={"tools": [weather_tool, temperature_tool]}, - ) - - results = await _resolve_approval_responses(messages, [weather_tool, temperature_tool], agent, {}) - - # Return value should only contain approved results - assert len(results) == 1 - assert results[0].call_id == approved_call_id - assert results[0].type == "function_result" - - # Rejection result should be written into messages (by _replace_approval_contents_with_results) - all_contents = [c for msg in messages for c in msg.contents] - rejection_results = [c for c in all_contents if c.type == "function_result" and c.call_id == rejected_call_id] - assert len(rejection_results) == 1 - assert "rejected" in str(rejection_results[0].result).lower() - - -async def test_resolve_approval_responses_preserves_follow_up_user_input_group() -> None: - """Approval-time follow-up requests stay grouped and do not emit a synthetic tool result.""" - from agent_framework import Message - from agent_framework.exceptions import UserInputRequiredException - - from agent_framework_ag_ui._agent_run import _resolve_approval_responses - - def request_consent() -> str: - raise UserInputRequiredException( - contents=[ - Content.from_oauth_consent_request(consent_link="https://example.com/consent-1"), - Content.from_oauth_consent_request(consent_link="https://example.com/consent-2"), - ] - ) - - consent_tool = FunctionTool( - name="request_consent", - description="Request two consent steps", - func=request_consent, - approval_mode="always_require", - ) - function_call = Content.from_function_call(call_id="call_consent", name="request_consent", arguments="{}") - approval_request = Content.from_function_approval_request(id="approval_consent", function_call=function_call) - messages: list[Any] = [ - Message(role="assistant", contents=[approval_request]), - Message(role="user", contents=[approval_request.to_function_approval_response(approved=True)]), - ] - agent = StubAgent(updates=[], default_options={"tools": [consent_tool]}) - store = InMemoryAGUIApprovalStateStore() - store.register_local( - thread_ids=["thread-consent"], - name="request_consent", - arguments="{}", - request_id="approval_consent", - interrupt_id="call_consent", - ) - intent = store.lifecycle.claim( - thread_id="thread-consent", - decision=ResumeDecision(interrupt_id="call_consent", accepted=True, arguments="{}"), - ) - - results = await _resolve_approval_responses( - messages, - [consent_tool], agent, - {}, - store.pending_approvals, - "thread-consent", - lifecycle=store.lifecycle, - authorized_executions={"call_consent": intent}, - ) - - follow_up_requests = [content for message in messages for content in message.contents if content.user_input_request] - assert results == [] - assert [request.consent_link for request in follow_up_requests] == [ - "https://example.com/consent-1", - "https://example.com/consent-2", - ] - assert not [content for message in messages for content in message.contents if content.type == "function_result"] - assert any(message.role == "assistant" and message.contents == follow_up_requests for message in messages) - - -async def test_resolve_approval_responses_returns_failure_when_grouped_execution_raises( - monkeypatch: Any, -) -> None: - """A grouped-execution failure produces one deterministic result for the approved call.""" - from agent_framework import Message - - from agent_framework_ag_ui._agent_run import _resolve_approval_responses - - async def fail_grouped_execution(**kwargs: Any) -> tuple[list[list[Content]], bool]: - del kwargs - raise RuntimeError("execution failed") - - monkeypatch.setattr( - "agent_framework_ag_ui._agent_run._try_execute_function_call_groups", - fail_grouped_execution, - ) - weather_tool = _make_weather_tool() - function_call = Content.from_function_call( - call_id="call_execution_failure", - name="get_weather", - arguments='{"city": "Seattle"}', - ) - approval_request = Content.from_function_approval_request( - id="approval_execution_failure", - function_call=function_call, - ) - messages: list[Any] = [ - Message(role="assistant", contents=[approval_request]), - Message(role="user", contents=[approval_request.to_function_approval_response(approved=True)]), - ] - agent = StubAgent(updates=[], default_options={"tools": [weather_tool]}) - - results = await _resolve_approval_responses(messages, [weather_tool], agent, {}) - - assert len(results) == 1 - assert results[0].type == "function_result" - assert results[0].call_id == "call_execution_failure" - assert results[0].result == "Error: Tool call invocation failed." - assert [ - content.result for message in messages for content in message.contents if content.type == "function_result" - ] == ["Error: Tool call invocation failed."] - - -async def test_resolve_approval_responses_keeps_fresh_occurrence_when_canonical_id_is_reused() -> None: - """A completed occurrence cannot consume a later approval that reuses its canonical call id.""" - from agent_framework import Message - - from agent_framework_ag_ui._agent_run import _resolve_approval_responses - - executions: list[str] = [] - - def guarded_write(value: str) -> str: - executions.append(value) - return f"wrote:{value}" - - tool = FunctionTool( - name="guarded_write", - description="Write a value", - func=guarded_write, - approval_mode="always_require", - ) - call_id = "call_reused" - first_call = Content.from_function_call(call_id=call_id, name="guarded_write", arguments={"value": "same"}) - second_call = Content.from_function_call(call_id=call_id, name="guarded_write", arguments={"value": "same"}) - messages = [ - Message(role="assistant", contents=[first_call]), - Message(role="tool", contents=[Content.from_function_result(call_id=call_id, result="already wrote")]), - Message( - role="user", - contents=[Content.from_function_approval_response(approved=True, id=call_id, function_call=first_call)], - ), - Message(role="assistant", contents=[second_call]), - Message( - role="user", - contents=[Content.from_function_approval_response(approved=True, id=call_id, function_call=second_call)], - ), - ] - thread_id = "thread-reused" - store = InMemoryAGUIApprovalStateStore() - store.register_local( - thread_ids=[thread_id], - name="guarded_write", - arguments='{"value":"same"}', - request_id=call_id, - interrupt_id=call_id, - ) - intent = store.lifecycle.claim( - thread_id=thread_id, - decision=ResumeDecision(interrupt_id=call_id, accepted=True, arguments='{"value":"same"}'), - ) - agent = StubAgent(updates=[], default_options={"tools": [tool]}) - - results = await _resolve_approval_responses( - messages, - [tool], - agent, - {}, - store.pending_approvals, - thread_id, - lifecycle=store.lifecycle, - authorized_executions={call_id: intent}, - ) - - assert executions == ["same"] - assert [result.result for result in results] == ["wrote:same"] - assert store.pending_approvals == {} - assert not [ - content for message in messages for content in message.contents if content.type == "function_approval_response" - ] - - -async def test_resolve_approval_responses_uses_fresh_decision_when_canonical_id_is_reused() -> None: - """A historical approval does not conflict with a fresh rejection for a reused call id.""" - from agent_framework import Message - - from agent_framework_ag_ui._agent_run import ( - _make_pending_approval_entry, - _pending_approval_key, - _resolve_approval_responses, - ) - - executions: list[str] = [] - - def guarded_write(value: str) -> str: - executions.append(value) - return f"wrote:{value}" - - tool = FunctionTool( - name="guarded_write", - description="Write a value", - func=guarded_write, - approval_mode="always_require", - ) - call_id = "call_reused_decision" - first_call = Content.from_function_call(call_id=call_id, name="guarded_write", arguments={"value": "same"}) - second_call = Content.from_function_call(call_id=call_id, name="guarded_write", arguments={"value": "same"}) - messages = [ - Message(role="assistant", contents=[first_call]), - Message(role="tool", contents=[Content.from_function_result(call_id=call_id, result="already wrote")]), - Message( - role="user", - contents=[Content.from_function_approval_response(approved=True, id=call_id, function_call=first_call)], - ), - Message(role="assistant", contents=[second_call]), - Message( - role="user", - contents=[Content.from_function_approval_response(approved=False, id=call_id, function_call=second_call)], - ), - ] - thread_id = "thread-reused-decision" - pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] = { - _pending_approval_key(thread_id, call_id): _make_pending_approval_entry( - "guarded_write", - '{"value":"same"}', - request_id=call_id, - interrupt_id=call_id, - ) - } - agent = StubAgent(updates=[], default_options={"tools": [tool]}) - - results = await _resolve_approval_responses( - messages, - [tool], - agent, - {}, - pending_approvals, - thread_id, - ) - - assert executions == [] - assert results == [] - assert pending_approvals == {} - assert [ - content.result for message in messages for content in message.contents if content.type == "function_result" - ] == ["already wrote", "Error: Tool call invocation was rejected by user."] - assert not [ - content for message in messages for content in message.contents if content.type == "function_approval_response" - ] - - -async def test_resolve_approval_responses_does_not_fall_back_when_fresh_reused_response_is_invalid() -> None: - """An invalid fresh response cannot consume pending state through a valid historical response.""" - from agent_framework import Message - - from agent_framework_ag_ui._agent_run import ( - _make_pending_approval_entry, - _pending_approval_key, - _resolve_approval_responses, - ) - - executions: list[str] = [] - - def guarded_write(value: str) -> str: - executions.append(value) - return f"wrote:{value}" - - tool = FunctionTool( - name="guarded_write", - description="Write a value", - func=guarded_write, - approval_mode="always_require", - ) - call_id = "call_reused_invalid" - first_call = Content.from_function_call(call_id=call_id, name="guarded_write", arguments={"value": "same"}) - edited_second_call = Content.from_function_call( - call_id=call_id, - name="guarded_write", - arguments={"value": "tampered"}, - ) - messages = [ - Message(role="assistant", contents=[first_call]), - Message(role="tool", contents=[Content.from_function_result(call_id=call_id, result="already wrote")]), - Message( - role="user", - contents=[Content.from_function_approval_response(approved=True, id=call_id, function_call=first_call)], - ), - Message( - role="assistant", - contents=[ - Content.from_function_call( - call_id=call_id, - name="guarded_write", - arguments={"value": "same"}, - ) - ], - ), - Message( - role="user", - contents=[ - Content.from_function_approval_response( - approved=True, - id=call_id, - function_call=edited_second_call, - ) - ], - ), - ] - thread_id = "thread-reused-invalid" - pending_entry = _make_pending_approval_entry( - "guarded_write", - '{"value":"same"}', - request_id=call_id, - interrupt_id=call_id, - ) - pending_key = _pending_approval_key(thread_id, call_id) - pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] = {pending_key: pending_entry} - agent = StubAgent(updates=[], default_options={"tools": [tool]}) - - results = await _resolve_approval_responses( - messages, - [tool], - agent, - {}, - pending_approvals, - thread_id, - ) - - assert executions == [] - assert results == [] - assert pending_approvals == {pending_key: pending_entry} - assert not [ - content for message in messages for content in message.contents if content.type == "function_approval_response" - ] - - -async def test_resolve_approval_responses_consumes_trusted_hosted_pending_entry() -> None: - """A server-collected hosted response remains provider-bound but cannot be replayed.""" - from agent_framework import Message - - from agent_framework_ag_ui._agent_run import ( - _make_pending_approval_entry, - _pending_approval_key, - _resolve_approval_responses, - ) - - call_id = "mcpr_hosted" - server_label = "hosted-mcp" - hosted_call = Content.from_function_call( - call_id=call_id, - name="hosted_write", - arguments={"value": "same"}, - additional_properties={"server_label": server_label}, - ) - hosted_response = Content.from_function_approval_response( - approved=True, - id=call_id, - function_call=hosted_call, - ) - messages = [Message(role="user", contents=[hosted_response])] - thread_id = "thread-hosted-collected" - pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] = { - _pending_approval_key(thread_id, call_id): _make_pending_approval_entry( - "hosted_write", - '{"value":"same"}', - request_id=call_id, - interrupt_id=call_id, - server_label=server_label, - ) - } - - results = await _resolve_approval_responses( - messages, - [], - StubAgent(updates=[]), - {}, - pending_approvals, - thread_id, - ) - - assert results == [] - assert pending_approvals == {} - assert len(messages) == 1 - assert messages[0].role == "user" - assert messages[0].contents == [hosted_response] - - -async def test_resolve_approval_responses_uses_fresh_response_across_pending_aliases() -> None: - """The interrupt-id response wins over historical replay under the request-id alias.""" - from agent_framework import Message - - from agent_framework_ag_ui._agent_run import ( - _make_pending_approval_entry, - _pending_approval_key, - _resolve_approval_responses, - ) - - executions: list[str] = [] - - def guarded_write(value: str) -> str: - executions.append(value) - return f"wrote:{value}" - - tool = FunctionTool(name="guarded_write", description="Write a value", func=guarded_write) - request_id = "approval_alias" - interrupt_id = "call_alias" - function_call = Content.from_function_call( - call_id=interrupt_id, - name="guarded_write", - arguments={"value": "same"}, - ) - messages = [ - Message(role="assistant", contents=[function_call]), - Message( - role="user", - contents=[ - Content.from_function_approval_response( - approved=True, - id=request_id, - function_call=function_call, - ) - ], - ), - Message( - role="user", - contents=[ - Content.from_function_approval_response( - approved=False, - id=interrupt_id, - function_call=function_call, - ) - ], - ), - ] - thread_id = "thread-alias" - pending_entry = _make_pending_approval_entry( - "guarded_write", - '{"value":"same"}', - request_id=request_id, - interrupt_id=interrupt_id, - ) - pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] = { - _pending_approval_key(thread_id, request_id): pending_entry, - _pending_approval_key(thread_id, interrupt_id): pending_entry, - } - - results = await _resolve_approval_responses( - messages, - [tool], - StubAgent(updates=[], default_options={"tools": [tool]}), - {}, - pending_approvals, - thread_id, - ) - - assert executions == [] - assert results == [] - assert pending_approvals == {} - assert [ - content.result for message in messages for content in message.contents if content.type == "function_result" - ] == ["Error: Tool call invocation was rejected by user."] - - -async def test_resolve_approval_responses_rejects_fresh_unknown_alias_without_historical_fallback() -> None: - """An unknown fresh response id cannot consume pending state through a trusted historical alias.""" - from agent_framework import Message - - from agent_framework_ag_ui._agent_run import ( - _make_pending_approval_entry, - _pending_approval_key, - _resolve_approval_responses, - ) - - executions: list[str] = [] - - def guarded_write(value: str) -> str: - executions.append(value) - return f"wrote:{value}" - - tool = FunctionTool(name="guarded_write", description="Write a value", func=guarded_write) - request_id = "approval_known" - interrupt_id = "call_known" - function_call = Content.from_function_call( - call_id=interrupt_id, - name="guarded_write", - arguments={"value": "same"}, - ) - messages = [ - Message(role="assistant", contents=[function_call]), - Message( - role="user", - contents=[ - Content.from_function_approval_response( - approved=True, - id=request_id, - function_call=function_call, - ), - Content.from_function_approval_response( - approved=True, - id="approval_unknown", - function_call=function_call, - ), - ], - ), - ] - thread_id = "thread-unknown-alias" - pending_entry = _make_pending_approval_entry( - "guarded_write", - '{"value":"same"}', - request_id=request_id, - interrupt_id=interrupt_id, - ) - request_key = _pending_approval_key(thread_id, request_id) - interrupt_key = _pending_approval_key(thread_id, interrupt_id) - pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] = { - request_key: pending_entry, - interrupt_key: pending_entry, - } - - results = await _resolve_approval_responses( - messages, - [tool], - StubAgent(updates=[], default_options={"tools": [tool]}), - {}, - pending_approvals, - thread_id, - ) - - assert executions == [] - assert results == [] - assert pending_approvals == {request_key: pending_entry, interrupt_key: pending_entry} - assert not [ - content for message in messages for content in message.contents if content.type == "function_approval_response" - ] - - -async def test_resolve_approval_responses_rejects_forged_call_id_for_valid_response_alias() -> None: - """A trusted response id cannot authorize a result under an unknown function-call id.""" - from agent_framework import Message - - from agent_framework_ag_ui._agent_run import ( - _make_pending_approval_entry, - _pending_approval_key, - _resolve_approval_responses, - ) - - executions: list[str] = [] - - def guarded_write(value: str) -> str: - executions.append(value) - return f"wrote:{value}" - - tool = FunctionTool(name="guarded_write", description="Write a value", func=guarded_write) - request_id = "approval_valid" - interrupt_id = "call_valid" - actual_call = Content.from_function_call( - call_id=interrupt_id, - name="guarded_write", - arguments={"value": "same"}, - ) - forged_call = Content.from_function_call( - call_id="call_forged", - name="guarded_write", - arguments={"value": "same"}, - ) - messages = [ - Message(role="assistant", contents=[actual_call]), - Message( - role="user", - contents=[ - Content.from_function_approval_response( - approved=True, - id=request_id, - function_call=forged_call, - ) - ], - ), - ] - thread_id = "thread-forged-call" - pending_entry = _make_pending_approval_entry( - "guarded_write", - '{"value":"same"}', - request_id=request_id, - interrupt_id=interrupt_id, - ) - request_key = _pending_approval_key(thread_id, request_id) - interrupt_key = _pending_approval_key(thread_id, interrupt_id) - pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] = { - request_key: pending_entry, - interrupt_key: pending_entry, - } - - results = await _resolve_approval_responses( - messages, - [tool], - StubAgent(updates=[], default_options={"tools": [tool]}), - {}, - pending_approvals, - thread_id, - ) - - assert executions == [] - assert results == [] - assert pending_approvals == {request_key: pending_entry, interrupt_key: pending_entry} - assert not [ - content - for message in messages - for content in message.contents - if content.type in {"function_approval_response", "function_result"} - ] - - -async def test_resolve_approval_responses_rejects_call_id_from_different_pending_entry() -> None: - """A response id from one approval cannot be paired with another approval's call id.""" - from agent_framework import Message + AgentConfig(), + approval_state_store=store, + ): + events.append(event) + return events - from agent_framework_ag_ui._agent_run import ( - _make_pending_approval_entry, - _pending_approval_key, - _resolve_approval_responses, - ) +async def test_approved_call_emits_one_live_result_under_original_identity() -> None: executions: list[str] = [] - def guarded_write(value: str) -> str: - executions.append(value) - return f"wrote:{value}" - - tool = FunctionTool(name="guarded_write", description="Write a value", func=guarded_write) - request_a = "approval_a" - call_a = Content.from_function_call( - call_id="call_a", - name="guarded_write", - arguments={"value": "same"}, + events = await _run_resume( + thread_id="thread-approved", + calls=[("call-weather", "Seattle")], + decisions=[("call-weather", True)], + executions=executions, ) - request_b = "approval_b" - call_b = Content.from_function_call( - call_id="call_b", - name="guarded_write", - arguments={"value": "same"}, - ) - messages = [ - Message(role="assistant", contents=[call_a]), - Message( - role="user", - contents=[ - Content.from_function_approval_response(approved=True, id=request_a, function_call=call_a), - Content.from_function_approval_response(approved=False, id=request_a, function_call=call_b), - ], - ), - ] - thread_id = "thread-crossed-aliases" - pending_a = _make_pending_approval_entry( - "guarded_write", - '{"value":"same"}', - request_id=request_a, - interrupt_id="call_a", - ) - pending_b = _make_pending_approval_entry( - "guarded_write", - '{"value":"same"}', - request_id=request_b, - interrupt_id="call_b", - ) - pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] = { - _pending_approval_key(thread_id, request_a): pending_a, - _pending_approval_key(thread_id, "call_a"): pending_a, - _pending_approval_key(thread_id, request_b): pending_b, - _pending_approval_key(thread_id, "call_b"): pending_b, - } - expected_pending = dict(pending_approvals) - - results = await _resolve_approval_responses( - messages, - [tool], - StubAgent(updates=[], default_options={"tools": [tool]}), - {}, - pending_approvals, - thread_id, - ) - - assert executions == [] - assert results == [] - assert pending_approvals == expected_pending - assert not [ - content - for message in messages - for content in message.contents - if content.type in {"function_approval_response", "function_result"} - ] - -async def test_resolve_approval_responses_without_registry_uses_latest_duplicate_decision() -> None: - """The optional no-registry path preserves the established last-response-wins behavior.""" - from agent_framework import Message + results = [event for event in events if getattr(event, "type", None) == "TOOL_CALL_RESULT"] + assert executions == ["Seattle"] + assert [(event.tool_call_id, event.content) for event in results] == [("call-weather", "Sunny in Seattle")] - from agent_framework_ag_ui._agent_run import _resolve_approval_responses +async def test_rejected_call_does_not_execute_or_emit_live_result() -> None: executions: list[str] = [] - def guarded_write(value: str) -> str: - executions.append(value) - return f"wrote:{value}" - - tool = FunctionTool(name="guarded_write", description="Write a value", func=guarded_write) - call_id = "call_no_registry" - function_call = Content.from_function_call( - call_id=call_id, - name="guarded_write", - arguments={"value": "same"}, - ) - messages = [ - Message(role="assistant", contents=[function_call]), - Message( - role="user", - contents=[ - Content.from_function_approval_response(approved=True, id=call_id, function_call=function_call), - Content.from_function_approval_response(approved=False, id=call_id, function_call=function_call), - ], - ), - ] - - results = await _resolve_approval_responses( - messages, - [tool], - StubAgent(updates=[], default_options={"tools": [tool]}), - {}, + events = await _run_resume( + thread_id="thread-rejected", + calls=[("call-weather", "Seattle")], + decisions=[("call-weather", False)], + executions=executions, ) assert executions == [] - assert results == [] - assert [ - content.result for message in messages for content in message.contents if content.type == "function_result" - ] == ["Error: Tool call invocation was rejected by user."] - - -async def test_resolve_approval_responses_legacy_registry_uses_latest_duplicate_decision() -> None: - """Legacy string entries group duplicate decisions by their matched response id.""" - from agent_framework import Message + assert not [event for event in events if getattr(event, "type", None) == "TOOL_CALL_RESULT"] - from agent_framework_ag_ui._agent_run import _pending_approval_key, _resolve_approval_responses +async def test_mixed_batch_preserves_approved_result_identity_and_order() -> None: executions: list[str] = [] - def guarded_write(value: str) -> str: - executions.append(value) - return f"wrote:{value}" - - tool = FunctionTool(name="guarded_write", description="Write a value", func=guarded_write) - approval_id = "approval_legacy" - first_call = Content.from_function_call( - call_id="call_legacy_old", - name="guarded_write", - arguments={"value": "same"}, - ) - latest_call = Content.from_function_call( - call_id="call_legacy_new", - name="guarded_write", - arguments={"value": "same"}, - ) - messages = [ - Message( - role="user", - contents=[ - Content.from_function_approval_response( - approved=True, - id=approval_id, - function_call=first_call, - ), - Content.from_function_approval_response( - approved=False, - id=approval_id, - function_call=latest_call, - ), - ], - ) - ] - thread_id = "thread-legacy" - pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] = { - _pending_approval_key(thread_id, approval_id): "guarded_write" - } - - results = await _resolve_approval_responses( - messages, - [tool], - StubAgent(updates=[], default_options={"tools": [tool]}), - {}, - pending_approvals, - thread_id, + events = await _run_resume( + thread_id="thread-mixed", + calls=[("call-seattle", "Seattle"), ("call-portland", "Portland")], + decisions=[("call-seattle", True), ("call-portland", False)], + executions=executions, ) - assert executions == [] - assert results == [] - assert pending_approvals == {} - assert [ - content.result for message in messages for content in message.contents if content.type == "function_result" - ] == ["Error: Tool call invocation was rejected by user."] - - -class TestApprovalToolResultDisplayChannel: - """Approved tools using ``state_update(..., tool_result=...)`` must route the - display payload to the UI event while ``flow.tool_results`` still receives - the LLM-bound text. The HITL approval emitter is separate from the standard - streaming emitter, so it gets its own coverage. - """ - - def test_approval_emits_display_payload_when_marker_present(self) -> None: - from agent_framework_ag_ui import state_update - from agent_framework_ag_ui._agent_run import _make_approval_tool_result_events - - display_payload = {"city": "Seattle", "temp": 14, "conditions": "foggy"} - inner = state_update(text="14°C, foggy", tool_result=display_payload) - resolved = Content.from_function_result(call_id="call_disp", result=[inner]) - - events = _make_approval_tool_result_events([resolved]) - - assert len(events) == 1 - # UI event must carry the serialized display payload, NOT the LLM text. - assert json.loads(events[0].content) == display_payload - assert events[0].content != "14°C, foggy" - - def test_approval_falls_back_to_text_when_no_marker(self) -> None: - """Backward compat: without a display marker, behaviour is unchanged.""" - from agent_framework_ag_ui._agent_run import _make_approval_tool_result_events - - resolved = Content.from_function_result(call_id="call_plain", result="Sunny in Seattle") - - events = _make_approval_tool_result_events([resolved]) - - assert len(events) == 1 - assert events[0].content == "Sunny in Seattle" + results = [event for event in events if getattr(event, "type", None) == "TOOL_CALL_RESULT"] + assert executions == ["Seattle"] + assert [(event.tool_call_id, event.content) for event in results] == [("call-seattle", "Sunny in Seattle")] diff --git a/python/packages/ag-ui/tests/ag_ui/test_approval_state.py b/python/packages/ag-ui/tests/ag_ui/test_approval_state.py index f64400aa5e..bf1a515be0 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_approval_state.py +++ b/python/packages/ag-ui/tests/ag_ui/test_approval_state.py @@ -50,7 +50,7 @@ def test_approval_state_store_does_not_evict_active_entries() -> None: interrupt_id="approval-2", ) - assert {entry["interrupt_id"] for entry in store.pending_approvals.values()} == {"approval-1"} + assert store.lifecycle.pending_interrupt_ids(thread_id="thread-1") == {"approval-1"} def test_approval_state_store_does_not_evict_active_middleware_state() -> None: @@ -60,4 +60,5 @@ def test_approval_state_store_does_not_evict_active_middleware_state() -> None: with pytest.raises(ApprovalCapacityError): store.set_tool_approval_state("thread-2", {"call_id": "call-2"}) - assert store.tool_approval_states == {"thread-1": {"call_id": "call-1"}} + assert store.get_tool_approval_state("thread-1") == {"call_id": "call-1"} + assert store.get_tool_approval_state("thread-2") is None diff --git a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py index a10599456c..5be6655a50 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -48,6 +48,7 @@ ) from agent_framework_ag_ui._agent import AgentFrameworkAgent from agent_framework_ag_ui._approval_lifecycle import ApprovalLifecycle +from agent_framework_ag_ui._approval_state import InMemoryAGUIApprovalStateStore, approval_state_thread_id from agent_framework_ag_ui._workflow import AgentFrameworkWorkflow @@ -1977,10 +1978,11 @@ async def test_endpoint_agent_approval_cancel_clears_queued_state_when_visible_e assert pause_response.status_code == 200 pause_finished = [event for event in _decode_sse_events(pause_response) if event.get("type") == "RUN_FINISHED"] assert [interrupt["id"] for interrupt in _run_finished_interrupts(pause_finished[-1])] == ["call_first"] - stored_state = wrapped_agent._approval_state_store.tool_approval_states["thread-queued-cancel-evicted"] + stored_state = wrapped_agent._approval_state_store.get_tool_approval_state("thread-queued-cancel-evicted") + assert stored_state is not None assert "call_second" in json.dumps(stored_state) - wrapped_agent._pending_approvals.clear() + wrapped_agent._approval_state_store = InMemoryAGUIApprovalStateStore() state["phase"] = "resume" cancel_response = client.post( "/approval", @@ -2680,7 +2682,7 @@ def get_weather(city: str) -> str: pause_finished = [event for event in _decode_sse_events(pause_response) if event.get("type") == "RUN_FINISHED"] assert _run_finished_interrupts(pause_finished[-1])[0]["id"] == "call_get_weather" - wrapped_agent._pending_approvals.clear() + wrapped_agent._approval_state_store = InMemoryAGUIApprovalStateStore() agent.updates = [AgentResponseUpdate(contents=[Content.from_text(text="Should not run")], role="assistant")] response = client.post( @@ -2928,8 +2930,10 @@ def record_city(city: str) -> str: assert len(run_errors) == 1 assert run_errors[0]["code"] == "APPROVAL_RESUME_CANCELLED" assert executed == [] - assert not any("call_seattle" in key for key in wrapped_agent._pending_approvals) - assert any("call_portland" in key for key in wrapped_agent._pending_approvals) + approval_thread_id = approval_state_thread_id(scope="tenant-a", thread_id="thread-two-approvals") + pending_ids = wrapped_agent._approval_state_store.lifecycle.pending_interrupt_ids(thread_id=approval_thread_id) + assert "call_seattle" not in pending_ids + assert "call_portland" in pending_ids hydrate_response = client.post( "/approval-snapshots", @@ -4965,7 +4969,8 @@ def get_weather(city: str) -> str: "code" ] == "APPROVAL_RESUME_CANCELLED" assert executed_cities == [] - assert not wrapped_agent._pending_approvals + approval_thread_id = approval_state_thread_id(scope="tenant-a", thread_id="agent-approval-thread") + assert not wrapped_agent._approval_state_store.lifecycle.pending_interrupt_ids(thread_id=approval_thread_id) hydrate_response = client.post( "/approval-snapshots", @@ -5811,7 +5816,7 @@ async def stream_fn( assert resume_response.status_code == 200 assert not [event for event in _decode_sse_events(resume_response) if event.get("type") == "RUN_ERROR"] assert local_executions == [] - assert not wrapped_agent._pending_approvals # pyright: ignore[reportPrivateUsage] + assert not wrapped_agent._approval_state_store.lifecycle.pending_interrupt_ids(thread_id="thread-hosted-approval") approval_responses = [ content for message in provider_messages @@ -5956,7 +5961,7 @@ async def stream_fn( assert resume_response.status_code == 200 assert local_executions == ["Approved draft"] - assert not wrapped_agent._pending_approvals # pyright: ignore[reportPrivateUsage] + assert not wrapped_agent._approval_state_store.lifecycle.pending_interrupt_ids(thread_id="thread-local-approval") state_snapshots = [ event["snapshot"] for event in _decode_sse_events(resume_response) if event.get("type") == "STATE_SNAPSHOT" ] diff --git a/python/packages/ag-ui/tests/ag_ui/test_run.py b/python/packages/ag-ui/tests/ag_ui/test_run.py index 95d9ec9c48..2e22067613 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_run.py @@ -24,21 +24,18 @@ from agent_framework_ag_ui._agent import AgentConfig from agent_framework_ag_ui._agent_run import ( - PendingApprovalEntry, - PendingApprovalKey, _build_messages_snapshot, _build_safe_metadata, _canonical_approval_resume_messages, _create_state_context_message, _filter_local_approval_responses_for_provider, _inject_state_context, - _make_pending_approval_entry, _normalize_response_stream, - _pending_approval_key, _resume_to_tool_messages, _should_suppress_intermediate_snapshot, run_agent_stream, ) +from agent_framework_ag_ui._approval_lifecycle import ApprovalLifecycle from agent_framework_ag_ui._run_common import ( FlowState, _build_run_finished_event, @@ -1041,29 +1038,29 @@ def test_resume_to_tool_messages_skips_cancelled_entries(): def test_canonical_approval_resume_does_not_mutate_arguments_until_batch_validates(): """Edited approval arguments are committed only after every resume entry validates.""" - pending_entry = _make_pending_approval_entry( - "get_weather", - '{"city":"Seattle"}', - request_id="call_a", + lifecycle = ApprovalLifecycle() + pending_entry = lifecycle.register_local( + thread_id="thread-weather", interrupt_id="call_a", + call_id="call_a", + name="get_weather", + arguments='{"city":"Seattle"}', + ) + lifecycle.register_local( + thread_id="thread-weather", + interrupt_id="call_b", + call_id="call_b", + name="get_weather", + arguments='{"city":"Portland"}', ) - pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] = { - _pending_approval_key("thread-weather", "call_a"): pending_entry, - _pending_approval_key("thread-weather", "call_b"): _make_pending_approval_entry( - "get_weather", - '{"city":"Portland"}', - request_id="call_b", - interrupt_id="call_b", - ), - } messages, handled_ids, cancelled_ids, error = _canonical_approval_resume_messages( [ {"interruptId": "call_a", "status": "resolved", "payload": {"accepted": True, "city": "Portland"}}, {"interruptId": "call_b", "status": "resolved", "payload": "not an object"}, ], - pending_approvals, "thread-weather", + lifecycle=lifecycle, ) assert messages == [] @@ -1071,20 +1068,20 @@ def test_canonical_approval_resume_does_not_mutate_arguments_until_batch_validat assert cancelled_ids == set() assert error is not None assert error.code == "APPROVAL_RESUME_INVALID" - assert pending_entry["arguments"] == '{"city":"Seattle"}' + assert pending_entry.arguments == '{"city":"Seattle"}' def test_canonical_hosted_approval_resume_rejects_edited_arguments_without_mutating_pending() -> None: """Hosted approvals accept a decision only because providers ignore edited arguments.""" - pending_entry = _make_pending_approval_entry( - "docs_search", - '{"query":"azure"}', - request_id="mcpr_docs", + lifecycle = ApprovalLifecycle() + pending_entry = lifecycle.register_hosted( + thread_id="thread-hosted", interrupt_id="mcpr_docs", + call_id="mcpr_docs", + name="docs_search", + arguments='{"query":"azure"}', server_label="Microsoft_Learn_MCP", ) - key = _pending_approval_key("thread-hosted", "mcpr_docs") - pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] = {key: pending_entry} messages, handled_ids, cancelled_ids, error = _canonical_approval_resume_messages( [ @@ -1094,8 +1091,8 @@ def test_canonical_hosted_approval_resume_rejects_edited_arguments_without_mutat "payload": {"accepted": True, "query": "untrusted edit"}, } ], - pending_approvals, "thread-hosted", + lifecycle=lifecycle, ) assert messages == [] @@ -1103,23 +1100,23 @@ def test_canonical_hosted_approval_resume_rejects_edited_arguments_without_mutat assert cancelled_ids == set() assert error is not None assert error.code == "APPROVAL_RESUME_INVALID_RESPONSE" - assert pending_entry["arguments"] == '{"query":"azure"}' - assert pending_approvals[key] is pending_entry + assert pending_entry.arguments == '{"query":"azure"}' + assert lifecycle.pending_occurrence(thread_id="thread-hosted", interrupt_id="mcpr_docs") is pending_entry -def test_pending_approval_registry_scans_exact_thread_keys_with_colons(): +def test_approval_lifecycle_scans_exact_thread_keys_with_colons(): """A thread id that prefixes another thread id must not inherit its pending approval contract.""" - pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] = { - _pending_approval_key("tenant:thread", "call_1"): _make_pending_approval_entry( - "get_weather", - '{"city":"Seattle"}', - request_id="call_1", - interrupt_id="call_1", - ) - } + lifecycle = ApprovalLifecycle() + lifecycle.register_local( + thread_id="tenant:thread", + interrupt_id="call_1", + call_id="call_1", + name="get_weather", + arguments='{"city":"Seattle"}', + ) - _, _, _, unrelated_error = _canonical_approval_resume_messages(None, pending_approvals, "tenant") - _, _, _, owning_error = _canonical_approval_resume_messages(None, pending_approvals, "tenant:thread") + _, _, _, unrelated_error = _canonical_approval_resume_messages(None, "tenant", lifecycle=lifecycle) + _, _, _, owning_error = _canonical_approval_resume_messages(None, "tenant:thread", lifecycle=lifecycle) assert unrelated_error is None assert owning_error is not None From 99feeaf529c0bc801279eabad3961b6dd1751b72 Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Mon, 10 Aug 2026 19:04:40 +0900 Subject: [PATCH 09/13] Align AG-UI approval resumes with protocol --- .../specs/004-python-function-calling-loop.md | 8 + python/packages/ag-ui/AGENTS.md | 3 + python/packages/ag-ui/README.md | 21 +- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 401 ++++++++++-------- .../_approval_lifecycle.py | 245 ++--------- .../agent_framework_ag_ui/_approval_state.py | 107 +---- .../agent_framework_ag_ui/_run_common.py | 21 +- .../agent_framework_ag_ui/_workflow_run.py | 21 +- .../ag_ui/test_agent_wrapper_comprehensive.py | 5 +- .../tests/ag_ui/test_approval_lifecycle.py | 201 ++++++--- .../tests/ag_ui/test_approval_result_event.py | 4 +- .../ag-ui/tests/ag_ui/test_approval_state.py | 27 +- .../ag-ui/tests/ag_ui/test_endpoint.py | 275 +++++++++--- python/packages/ag-ui/tests/ag_ui/test_run.py | 23 +- 14 files changed, 753 insertions(+), 609 deletions(-) diff --git a/docs/specs/004-python-function-calling-loop.md b/docs/specs/004-python-function-calling-loop.md index 8462dc76ac..a6e7534362 100644 --- a/docs/specs/004-python-function-calling-loop.md +++ b/docs/specs/004-python-function-calling-loop.md @@ -361,6 +361,10 @@ that manually replay messages own the equivalent rule: do not resend an approval response for local execution, and leaves hosted approval responses as provider protocol data. - Hosted AG-UI approval interrupts expose an accept/reject decision only; argument edits are rejected because the hosted provider executes the server-owned request rather than client-edited arguments. +- AG-UI tool approval resumes accept the standard `approved` decision and full-replacement `editedArgs` payload. + Existing MAF clients remain compatible through the `accepted` decision alias and direct partial argument edits. +- An AG-UI `cancelled` resume is a valid terminal decision, not a run error. In a resume covering parallel open + interrupts, resolved siblings still execute and cancelled calls do not. - A server-issued approval request must not be replayed inline during service-side continuation. - History providers may retain approval control contents in their backing store for audit, but base history replay filters them before later model calls. @@ -449,6 +453,10 @@ that manually replay messages own the equivalent rule: do not resend an approval | Argument-scoped rule | Exact arguments are required; empty arguments are not tool-wide. | `test_tool_approval_middleware_always_approve_tool_with_arguments_rule`, `test_tool_approval_middleware_empty_arguments_rule_is_not_tool_wide` | | Provider-injected approval tool | A tool added during `before_run` defers to in-run resolution, executes once, and emits one result. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_agent_approval_deferred_provider_tool_executes` | | AG-UI provider boundary | Completed local approval controls from AG-UI request and snapshot replay are absent from raw chat-client input while deferred and hosted approvals keep their respective in-run/provider paths. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_does_not_forward_resolved_local_approval_control_to_chat_client`, `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_agent_approval_deferred_provider_tool_executes`, `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_canonical_resume_preserves_hosted_approval_for_provider`, `packages/ag-ui/tests/ag_ui/test_run.py::test_filter_local_approval_responses_for_provider_removes_duplicate_completed_controls`, `packages/ag-ui/tests/ag_ui/test_run.py::test_filter_local_approval_responses_for_provider_pairs_reused_call_ids_by_occurrence`, `packages/ag-ui/tests/ag_ui/test_run.py::test_canonical_hosted_approval_resume_rejects_edited_arguments_without_mutating_pending` | +| AG-UI standard approval payload | `approved` plus full-replacement `editedArgs` executes once and replays idempotently, while legacy `accepted` plus direct partial edits remains supported. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_agent_approval_resume_entry_applies_standard_full_replacement_edited_args`, `test_endpoint_agent_approval_replayed_standard_edited_resume_is_idempotent`, `test_endpoint_agent_approval_resume_entry_applies_edited_arguments` | +| AG-UI cancellation | A cancelled interrupt executes zero times and completes normally; resolved siblings in the same complete resume still execute once. Workflow `request_info` cancellation follows the same terminal lifecycle. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_agent_approval_cancelled_resume_entry_completes_without_execution`, `test_endpoint_agent_approval_mixed_cancelled_and_resolved_resume_executes_resolved_tool`, `test_endpoint_workflow_request_info_cancelled_resume_completes_normally` | +| AG-UI local executor unavailable on resume | A claimed local occurrence whose executor disappeared releases its unstarted claim, reports temporary unavailability, and remains safely retryable. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_agent_approval_resume_remains_retryable_when_local_tool_is_temporarily_unavailable` | +| AG-UI forwarded execution interruption | A provider failure, cancellation, or stream close after forwarding an approval recovers the open occurrence as indeterminate when no idempotency key proves retry safety. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_hosted_approval_becomes_indeterminate_when_provider_stream_fails` | ### Errors, control flow, and limits diff --git a/python/packages/ag-ui/AGENTS.md b/python/packages/ag-ui/AGENTS.md index bee4a0f0e6..86a178f322 100644 --- a/python/packages/ag-ui/AGENTS.md +++ b/python/packages/ag-ui/AGENTS.md @@ -29,6 +29,9 @@ AG-UI protocol integration for building agent UIs with the AG-UI standard. - Multimodal user inputs support both legacy (`text`, `binary`) and draft-style (`image`, `audio`, `video`, `document`) shapes. - Interrupted runs complete with `RUN_FINISHED.outcome.type == "interrupt"` and canonical `outcome.interrupts`; do not document or add new flows that depend on the legacy top-level `RUN_FINISHED.interrupt` field. - `Interrupt` and `ResumeEntry` come from the `ag-ui-protocol` package (`ag_ui.core`), not from an Agent Framework-specific interrupt model. +- Tool approval interrupts advertise standard `approved` and full-replacement `editedArgs` responses while retaining + the existing `accepted` alias and direct partial edits for MAF client compatibility. A `cancelled` resume completes + normally without executing that call; resolved siblings in the same complete resume still proceed. - Approval-time execution preserves each call's complete result group. Follow-up user-input requests remain in the resumed messages, while `TOOL_CALL_RESULT` events are emitted only for terminal `function_result` contents. - Approval responses for tools injected during `before_run` are deferred to the in-run approval middleware rather diff --git a/python/packages/ag-ui/README.md b/python/packages/ag-ui/README.md index 9302f42b2f..b9769233f8 100644 --- a/python/packages/ag-ui/README.md +++ b/python/packages/ag-ui/README.md @@ -167,10 +167,23 @@ Interrupted terminal event shape: "responseSchema": { "type": "object", "properties": { + "approved": { "type": "boolean" }, "accepted": { "type": "boolean" }, - "arguments": { "type": "object" } + "city": { "type": "string" }, + "editedArgs": { + "type": "object", + "description": "Full replacement of the tool arguments. Not merged.", + "properties": { + "city": { "type": "string" } + }, + "required": ["city"], + "additionalProperties": false + } }, - "required": ["accepted"] + "anyOf": [ + { "required": ["approved"] }, + { "required": ["accepted"] } + ] }, "metadata": { "agent_framework": { @@ -192,6 +205,10 @@ Interrupted terminal event shape: Resume the paused thread with a canonical `resume` array. Each entry addresses exactly one open interrupt by `interruptId`; `status` is `resolved` or `cancelled`; resolved entries carry the approval or workflow response payload. +Tool approvals use the standard `approved` field and may provide `editedArgs` as a full replacement of the tool +arguments. For compatibility with existing MAF clients, `accepted` remains an alias for `approved`, and direct +argument fields remain supported as partial edits. Cancellation is a normal terminal decision: cancelled calls do +not execute, while resolved siblings in the same complete resume continue normally. ```json { diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index dc227171aa..c8176b9d58 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py @@ -60,6 +60,7 @@ ApprovalOccurrence, ApprovalSnapshotReconciliation, AuthorizedExecution, + ClaimRecoveryPolicy, DeferredPendingToolTransitionOwner, ForwardedPendingToolTransitionOwner, HostedPendingToolTransitionOwner, @@ -864,17 +865,10 @@ def _register_server_generated_approval_response( tools, has_deferred_owner=has_deferred_owner, ) - if execution_owner is ApprovalExecutionOwner.HOSTED: - register = lifecycle.register_hosted - elif execution_owner is ApprovalExecutionOwner.DEFERRED: - register = lifecycle.register_deferred - elif execution_owner is ApprovalExecutionOwner.LOCAL: - register = lifecycle.register_local - else: - register = lifecycle.register_unowned arguments = canonical_function_arguments(response.function_call) or "{}" - register( - thread_id=thread_id, + lifecycle.register( + owner=execution_owner, + thread_ids=[thread_id], interrupt_id=str(response_id), call_id=str(response.function_call.call_id or response_id), name=response.function_call.name, @@ -1068,6 +1062,8 @@ def _canonical_approval_resume_messages( expected_interrupt_ids: set[str] | None = None, *, lifecycle: ApprovalLifecycle, + tools: list[Any] | None = None, + has_deferred_owner: bool = False, authorized_executions: dict[str, AuthorizedExecution] | None = None, retained_results: list[Content] | None = None, snapshot_reconciliations: list[ApprovalSnapshotReconciliation] | None = None, @@ -1091,7 +1087,7 @@ def _canonical_approval_resume_messages( payload = _parse_json_object(interrupt.get("value")) if payload is None: break - accepted = payload.get("accepted", payload.get("approved")) + accepted = payload.get("approved", payload.get("accepted")) if not isinstance(accepted, bool): break interrupt_id = str(interrupt["id"]) @@ -1103,10 +1099,19 @@ def _canonical_approval_resume_messages( except KeyError: break edited_arguments = { - key: value for key, value in payload.items() if key not in {"accepted", "approved"} + key: value + for key, value in payload.items() + if key not in {"accepted", "approved", "editedArgs"} } + standard_edited_arguments = payload.get("editedArgs") canonical_arguments: str | None = None - if edited_arguments: + if isinstance(standard_edited_arguments, dict) and not edited_arguments: + canonical_arguments = json.dumps( + make_json_safe(standard_edited_arguments), + sort_keys=True, + separators=(",", ":"), + ) + elif edited_arguments: retained_argument_values = _parse_json_object(retained_arguments) if retained_argument_values is None: break @@ -1222,23 +1227,13 @@ def _canonical_approval_resume_messages( cancelled_ids, RunErrorEvent(message=str(exc), code="APPROVAL_RESUME_INVALID"), ) - interrupt_id = next(str(entry["interrupt_id"]) for entry in entries if entry.get("status") == "cancelled") - return ( - [], - handled_ids, - cancelled_ids, - RunErrorEvent( - message=f"Approval resume for interruptId '{interrupt_id}' was cancelled.", - code="APPROVAL_RESUME_CANCELLED", - ), - ) lifecycle_decisions: list[ResumeDecision] = [] restored_sibling_response_ids: set[str] = set() for entry in entries: interrupt_id = cast(str, entry["interrupt_id"]) pending_entry = entries_by_interrupt_id.get(interrupt_id) - if pending_entry is None: + if pending_entry is None or entry["status"] == "cancelled": continue payload = _parse_json_object(entry.get("payload")) @@ -1252,7 +1247,7 @@ def _canonical_approval_resume_messages( code="APPROVAL_RESUME_INVALID", ), ) - accepted = payload.get("accepted", payload.get("approved")) + accepted = payload.get("approved", payload.get("accepted")) if not isinstance(accepted, bool): return ( [], @@ -1266,7 +1261,40 @@ def _canonical_approval_resume_messages( pending_arguments = _pending_approval_arguments(pending_entry) original_arguments = _parse_json_object(pending_arguments) or {} - edited_arguments = {key: value for key, value in payload.items() if key not in {"accepted", "approved"}} + direct_edited_arguments = { + key: value for key, value in payload.items() if key not in {"accepted", "approved", "editedArgs"} + } + standard_edited_arguments = payload.get("editedArgs") + if standard_edited_arguments is not None: + if not isinstance(standard_edited_arguments, dict) or direct_edited_arguments: + return ( + [], + handled_ids, + cancelled_ids, + RunErrorEvent( + message=( + f"Approval resume for interruptId '{interrupt_id}' must provide editedArgs as the " + "only edited-argument representation." + ), + code="APPROVAL_RESUME_INVALID_RESPONSE", + ), + ) + edited_arguments = cast(dict[str, Any], standard_edited_arguments) + if set(edited_arguments) != set(original_arguments): + return ( + [], + handled_ids, + cancelled_ids, + RunErrorEvent( + message=( + f"Approval resume for interruptId '{interrupt_id}' must provide editedArgs as a full " + "replacement of the pending tool arguments." + ), + code="APPROVAL_RESUME_INVALID_RESPONSE", + ), + ) + else: + edited_arguments = direct_edited_arguments if edited_arguments and _pending_approval_server_label(pending_entry): return ( [], @@ -1305,7 +1333,11 @@ def _canonical_approval_resume_messages( ), ) - merged_arguments = {**original_arguments, **edited_arguments} + merged_arguments = ( + dict(edited_arguments) + if standard_edited_arguments is not None + else {**original_arguments, **edited_arguments} + ) canonical_arguments = json.dumps(make_json_safe(merged_arguments), sort_keys=True, separators=(",", ":")) lifecycle_decisions.append( ResumeDecision( @@ -1342,11 +1374,14 @@ def _canonical_approval_resume_messages( sibling_interrupt_id = str(response_id) sibling_call_id = str(function_call.call_id or response_id) sibling_arguments = canonical_function_arguments(function_call) or "{}" - register = ( - lifecycle.register_hosted if _function_call_server_label(function_call) else lifecycle.register_local + execution_owner = _function_call_execution_owner( + function_call, + tools, + has_deferred_owner=has_deferred_owner, ) - register( - thread_id=thread_id, + lifecycle.register( + owner=execution_owner, + thread_ids=[thread_id], interrupt_id=sibling_interrupt_id, call_id=sibling_call_id, name=function_call.name, @@ -1607,13 +1642,15 @@ async def forward_hosted_decision(approval: Content = approval) -> list[Content] approved_function_result_groups: list[list[Content]] = [] # Partition approved responses into static (execute now) and deferred (execute during run) - tool_map = _get_tool_map(tools) if tools else {} static_approved: list[Content] = [] for approval in approved_responses: - tool_name = approval.function_call.name if approval.function_call else None - if tool_name in tool_map and not _is_hosted_tool_approval(approval): - static_approved.append(approval) + function_call = approval.function_call + call_id = (function_call.call_id if function_call else None) or approval.id or "" + intent = authorized_executions.get(call_id) if authorized_executions is not None else None + if intent is None or intent.owner is not ApprovalExecutionOwner.LOCAL: + continue + static_approved.append(approval) # Execute lifecycle-authorized local calls only through their transition owner. if static_approved and tools and lifecycle is not None and authorized_executions is not None: @@ -2203,12 +2240,18 @@ async def run_agent_stream( forwarded_executions: dict[str, tuple[ForwardedPendingToolTransitionOwner, AuthorizedExecution, Content]] = {} retained_approval_results: list[Content] = [] approval_snapshot_reconciliations: list[ApprovalSnapshotReconciliation] = [] + client_tools = convert_agui_tools_to_agent_framework(input_data.get("tools")) + server_tools = collect_server_tools(agent) + register_additional_client_tools(agent, client_tools) + tools = merge_tools(server_tools, client_tools) approval_resume_messages, handled_resume_ids, cancelled_resume_ids, resume_error = ( _canonical_approval_resume_messages( resume_payload, approval_thread_id, expected_interrupt_ids=stored_pending_approval_interrupt_ids or None, lifecycle=approval_state_store.lifecycle, + tools=tools, + has_deferred_owner=approval_state_store.has_tool_approval_state(approval_thread_id), authorized_executions=authorized_executions, retained_results=retained_approval_results, snapshot_reconciliations=approval_snapshot_reconciliations, @@ -2236,6 +2279,17 @@ async def run_agent_stream( await snapshot_session.clear_interrupts(interrupt_ids=retired_interrupt_ids or cancelled_resume_ids or None) yield resume_error return + if cancelled_resume_ids and handled_resume_ids == cancelled_resume_ids: + yield RunStartedEvent(run_id=run_id, thread_id=thread_id) + _clear_tool_approval_state(approval_state_store, approval_thread_id) + retired_interrupt_ids = { + reconciliation.identity.interrupt_id if reconciliation.identity is not None else reconciliation.interrupt_id + for reconciliation in approval_snapshot_reconciliations + if reconciliation.retire_interrupt + } + await snapshot_session.clear_interrupts(interrupt_ids=retired_interrupt_ids or cancelled_resume_ids) + yield _build_run_finished_event(run_id=run_id, thread_id=thread_id) + return resume_messages = _resume_to_tool_messages(resume_payload, exclude_interrupt_ids=handled_resume_ids) if available_interrupts: logger.debug("Received available interrupts metadata: %s", available_interrupts) @@ -2273,12 +2327,6 @@ async def run_agent_stream( yield _build_run_finished_event(run_id=run_id, thread_id=thread_id) return - # Prepare tools - client_tools = convert_agui_tools_to_agent_framework(input_data.get("tools")) - server_tools = collect_server_tools(agent) - register_additional_client_tools(agent, client_tools) - tools = merge_tools(server_tools, client_tools) - # Create session (with service session support) if config.use_service_session: session = AgentSession(session_id=thread_id, service_session_id=supplied_thread_id) @@ -2334,8 +2382,27 @@ async def run_agent_stream( authorized_executions=authorized_executions, ) ) + execution_tool_map = _get_tool_map(tools_for_execution) if tools_for_execution else {} + local_intents = [ + intent for intent in authorized_executions.values() if intent.owner is ApprovalExecutionOwner.LOCAL + ] + unavailable_local_intents = [ + intent + for intent in local_intents + if (tool := execution_tool_map.get(intent.name)) is None or getattr(tool, "declaration_only", False) + ] + if unavailable_local_intents: + for intent in local_intents: + approval_state_store.lifecycle.release_claim(intent, policy=ClaimRecoveryPolicy.SAFE_TO_RETRY) + unavailable_names = ", ".join(sorted({intent.name for intent in unavailable_local_intents})) + yield RunStartedEvent(run_id=run_id, thread_id=thread_id) + yield RunErrorEvent( + message=f"Approved tool(s) {unavailable_names} are temporarily unavailable; retry the approval later.", + code="APPROVAL_TOOL_UNAVAILABLE", + ) + return validated_approved_responses: list[Content] = [] - resolved_approval_results = retained_approval_results + await _resolve_approval_responses( + newly_resolved_approval_results = await _resolve_approval_responses( messages, tools_for_execution, agent, @@ -2346,6 +2413,7 @@ async def run_agent_stream( authorized_executions=authorized_executions, forwarded_executions=forwarded_executions, ) + resolved_approval_results = retained_approval_results + newly_resolved_approval_results # Defense-in-depth: replace approval payloads in snapshot with actual tool results # so CopilotKit does not re-send stale approval content on subsequent turns. @@ -2409,125 +2477,133 @@ async def run_agent_stream( # telemetry override must cover construction, stream resolution, and every pull. telemetry_conversation_id = str(supplied_thread_id) if supplied_thread_id is not None else None telemetry_context = partial(_use_telemetry_conversation_id, telemetry_conversation_id) - with telemetry_context(): - response_stream = agent.run(messages, stream=True, **run_kwargs) - stream = await _normalize_response_stream(response_stream) - - async for update in _iterate_with_context(stream, telemetry_context): - # Collect updates for structured output processing - if response_format is not None: - all_updates.append(update) - - # Use service-generated IDs only when the AG-UI request omitted them. Client-supplied - # IDs remain authoritative for lifecycle correlation and thread-scoped persistence. - if not run_started_emitted: - conv_id = get_conversation_id_from_update(update) - if conv_id: - provider_thread_id = conv_id - if supplied_thread_id is None and conv_id: - thread_id = conv_id - snapshot_session.rebind_thread_id(thread_id) - if supplied_run_id is None and update.response_id: - run_id = update.response_id - # NOW emit RunStarted with proper IDs - yield RunStartedEvent(run_id=run_id, thread_id=thread_id) - # Emit PredictState custom event if configured - if predict_state_config: - predict_state_value = [ - { - "state_key": state_key, - "tool": cfg["tool"], - "tool_argument": cfg["tool_argument"], - } - for state_key, cfg in predict_state_config.items() - ] - yield CustomEvent(name="PredictState", value=predict_state_value) - # Emit initial state snapshot only if we have both state_schema and state - if state_schema and flow.current_state: - latest_state_snapshot = cast(dict[str, Any], make_json_safe(flow.current_state)) - yield StateSnapshotEvent(snapshot=flow.current_state) - run_started_emitted = True - - for event in _make_approval_tool_result_events(resolved_approval_results): - yield event - - # Feature #4: Detect tool-only messages (no text content) - # Emit TextMessageStartEvent to create message context for tool calls - if not flow.message_id and _has_only_tool_calls(update.contents): - flow.message_id = generate_event_id() - logger.info(f"Tool-only response detected, creating message_id={flow.message_id}") - yield TextMessageStartEvent(message_id=flow.message_id, role="assistant") - - # Emit events for each content item - for content in update.contents: - content_type = getattr(content, "type", None) - logger.debug(f"Processing content type={content_type}, message_id={flow.message_id}") - - if ( - content_type == "function_result" - and content.call_id - and (forwarded := forwarded_executions.pop(content.call_id, None)) is not None - and approval_state_store is not None - ): - owner, intent, _ = forwarded - owner.record_outcome(intent, [content], lifecycle=approval_state_store.lifecycle) - - # Register pending approval requests so we can validate responses later - if content_type == "function_approval_request": - if content.id and content.function_call and content.function_call.name: - canonical_interrupt_id = content.function_call.call_id or content.id - provider_approval_thread_id = approval_state_thread_id( - scope=approval_scope, - thread_id=provider_thread_id or thread_id, - ) - server_label = _function_call_server_label(content.function_call) - already_approved_requests = _stored_already_approved_requests_for_visible_approval( - session, - str(content.id), - str(canonical_interrupt_id) if canonical_interrupt_id else None, - ) - execution_owner = _function_call_execution_owner( - content.function_call, - tools, - has_deferred_owner=_TOOL_APPROVAL_STATE_KEY in session.state, - ) - registration_kwargs = { - "thread_ids": [approval_thread_id, provider_approval_thread_id], - "name": content.function_call.name, - "arguments": canonical_function_arguments(content.function_call) or "{}", - "request_id": str(content.id), - "interrupt_id": str(canonical_interrupt_id), - "already_approved_requests": already_approved_requests, - } - if execution_owner is ApprovalExecutionOwner.HOSTED: - approval_state_store.register_hosted(server_label=server_label, **registration_kwargs) - elif execution_owner is ApprovalExecutionOwner.DEFERRED: - approval_state_store.register_deferred(**registration_kwargs) - elif execution_owner is ApprovalExecutionOwner.LOCAL: - approval_state_store.register_local(**registration_kwargs) + stream_completed = False + try: + with telemetry_context(): + response_stream = agent.run(messages, stream=True, **run_kwargs) + stream = await _normalize_response_stream(response_stream) + + async for update in _iterate_with_context(stream, telemetry_context): + # Collect updates for structured output processing + if response_format is not None: + all_updates.append(update) + + # Use service-generated IDs only when the AG-UI request omitted them. Client-supplied + # IDs remain authoritative for lifecycle correlation and thread-scoped persistence. + if not run_started_emitted: + conv_id = get_conversation_id_from_update(update) + if conv_id: + provider_thread_id = conv_id + if supplied_thread_id is None and conv_id: + thread_id = conv_id + snapshot_session.rebind_thread_id(thread_id) + if supplied_run_id is None and update.response_id: + run_id = update.response_id + # NOW emit RunStarted with proper IDs + yield RunStartedEvent(run_id=run_id, thread_id=thread_id) + # Emit PredictState custom event if configured + if predict_state_config: + predict_state_value = [ + { + "state_key": state_key, + "tool": cfg["tool"], + "tool_argument": cfg["tool_argument"], + } + for state_key, cfg in predict_state_config.items() + ] + yield CustomEvent(name="PredictState", value=predict_state_value) + # Emit initial state snapshot only if we have both state_schema and state + if state_schema and flow.current_state: + latest_state_snapshot = cast(dict[str, Any], make_json_safe(flow.current_state)) + yield StateSnapshotEvent(snapshot=flow.current_state) + run_started_emitted = True + + for event in _make_approval_tool_result_events(resolved_approval_results): + yield event + + # Feature #4: Detect tool-only messages (no text content) + # Emit TextMessageStartEvent to create message context for tool calls + if not flow.message_id and _has_only_tool_calls(update.contents): + flow.message_id = generate_event_id() + logger.info(f"Tool-only response detected, creating message_id={flow.message_id}") + yield TextMessageStartEvent(message_id=flow.message_id, role="assistant") + + # Emit events for each content item + for content in update.contents: + content_type = getattr(content, "type", None) + logger.debug(f"Processing content type={content_type}, message_id={flow.message_id}") + + if ( + content_type == "function_result" + and content.call_id + and (forwarded := forwarded_executions.pop(content.call_id, None)) is not None + and approval_state_store is not None + ): + owner, intent, _ = forwarded + owner.record_outcome(intent, [content], lifecycle=approval_state_store.lifecycle) + + # Register pending approval requests so we can validate responses later + if content_type == "function_approval_request": + if content.id and content.function_call and content.function_call.name: + canonical_interrupt_id = content.function_call.call_id or content.id + provider_approval_thread_id = approval_state_thread_id( + scope=approval_scope, + thread_id=provider_thread_id or thread_id, + ) + server_label = _function_call_server_label(content.function_call) + already_approved_requests = _stored_already_approved_requests_for_visible_approval( + session, + str(content.id), + str(canonical_interrupt_id) if canonical_interrupt_id else None, + ) + execution_owner = _function_call_execution_owner( + content.function_call, + tools, + has_deferred_owner=_TOOL_APPROVAL_STATE_KEY in session.state, + ) + registration_kwargs = { + "thread_ids": [approval_thread_id, provider_approval_thread_id], + "name": content.function_call.name, + "arguments": canonical_function_arguments(content.function_call) or "{}", + "request_id": str(content.id), + "interrupt_id": str(canonical_interrupt_id), + "already_approved_requests": already_approved_requests, + } + approval_state_store.register( + owner=execution_owner, + server_label=server_label, + **registration_kwargs, + ) else: - approval_state_store.register_unowned(**registration_kwargs) - else: - logger.warning( - "Approval request not registered: missing id=%s, function_call=%s, or function name", - getattr(content, "id", None), - getattr(content, "function_call", None), - ) + logger.warning( + "Approval request not registered: missing id=%s, function_call=%s, or function name", + getattr(content, "id", None), + getattr(content, "function_call", None), + ) - for event in _emit_content( - content, - flow, - predictive_handler, - skip_text, - config.require_confirmation, - ): - if isinstance(event, StateSnapshotEvent): - latest_state_snapshot = cast(dict[str, Any], make_json_safe(event.snapshot)) - yield event - - # Stop if waiting for approval - if flow.waiting_for_approval: - break + for event in _emit_content( + content, + flow, + predictive_handler, + skip_text, + config.require_confirmation, + ): + if isinstance(event, StateSnapshotEvent): + latest_state_snapshot = cast(dict[str, Any], make_json_safe(event.snapshot)) + yield event + + # Stop if waiting for approval + if flow.waiting_for_approval: + break + stream_completed = True + finally: + if approval_state_store is not None: + for owner, intent, forwarded_approval in forwarded_executions.values(): + if stream_completed: + owner.record_outcome(intent, [forwarded_approval], lifecycle=approval_state_store.lifecycle) + else: + approval_state_store.lifecycle.recover_execution(intent, owner=intent.owner) + forwarded_executions.clear() if flow.waiting_for_approval and isinstance(stream, ResponseStream): await stream.get_final_response() @@ -2707,13 +2783,6 @@ async def run_agent_stream( ): yield snapshot_event - # Always emit RunFinished - confirm_changes tool call is complete (Start -> Args -> End) - # The UI will show confirmation dialog and send a new request when user responds - if approval_state_store is not None: - for owner, intent, forwarded_approval in forwarded_executions.values(): - owner.record_outcome(intent, [forwarded_approval], lifecycle=approval_state_store.lifecycle) - forwarded_executions.clear() - persisted_messages = latest_messages_snapshot if resume_payload is not None and not seeded_resume_from_snapshot: # Generic resume requests carry only the synthesized response, so prepend diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py b/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py index 7d8681878a..cd2317dcec 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py @@ -100,19 +100,21 @@ class ApprovalStatus(str, Enum): EXPIRED = "expired" INDETERMINATE = "indeterminate" + @property + def is_terminal(self) -> bool: + """Whether this state has ended its approval authority.""" + return self in { + ApprovalStatus.SETTLED, + ApprovalStatus.REJECTED, + ApprovalStatus.CANCELLED, + ApprovalStatus.EXPIRED, + ApprovalStatus.INDETERMINATE, + } -class ApprovalSnapshotStatus(str, Enum): - """Approval authority status projected during snapshot reconciliation.""" - - PENDING = "pending" - CLAIMED = "claimed" - EXECUTING = "executing" - SETTLED = "settled" - REJECTED = "rejected" - CANCELLED = "cancelled" - EXPIRED = "expired" - INDETERMINATE = "indeterminate" - MISSING = "missing" + @property + def is_purgeable(self) -> bool: + """Whether this terminal state may expire after the retention window.""" + return self.is_terminal and self is not ApprovalStatus.INDETERMINATE class ApprovalExecutionOwner(str, Enum): @@ -146,9 +148,14 @@ class ApprovalSnapshotReconciliation: interrupt_id: str identity: ApprovalOccurrenceIdentity | None - status: ApprovalSnapshotStatus + status: ApprovalStatus | None retire_interrupt: bool + @property + def is_missing(self) -> bool: + """Whether no lifecycle occurrence exists for the snapshot interrupt.""" + return self.status is None + @dataclass(frozen=True) class ResumeDecision: @@ -245,166 +252,12 @@ def __init__( self._pending_by_interrupt: dict[tuple[str, str], ApprovalOccurrenceIdentity] = {} self._terminal_by_interrupt: dict[tuple[str, str], ApprovalOccurrenceIdentity] = {} - def register_local( - self, - *, - thread_id: str, - interrupt_id: str, - call_id: str, - name: str, - arguments: str, - aliases: list[str] | None = None, - already_approved_requests: list[dict[str, Any]] | None = None, - server_label: str | None = None, - idempotency_key: str | None = None, - ) -> ApprovalOccurrence: - """Register one server-generated local approval occurrence.""" - return self.register_local_aliases( - thread_ids=[thread_id], - interrupt_id=interrupt_id, - call_id=call_id, - name=name, - arguments=arguments, - aliases=aliases, - already_approved_requests=already_approved_requests, - server_label=server_label, - idempotency_key=idempotency_key, - ) - - def register_local_aliases( - self, - *, - thread_ids: list[str], - interrupt_id: str, - call_id: str, - name: str, - arguments: str, - aliases: list[str] | None = None, - already_approved_requests: list[dict[str, Any]] | None = None, - server_label: str | None = None, - idempotency_key: str | None = None, - ) -> ApprovalOccurrence: - """Register one occurrence under its trusted scoped-thread aliases.""" - return self._register_aliases( - thread_ids=thread_ids, - interrupt_id=interrupt_id, - call_id=call_id, - name=name, - arguments=arguments, - owner=ApprovalExecutionOwner.LOCAL, - aliases=aliases, - already_approved_requests=already_approved_requests, - server_label=server_label, - idempotency_key=idempotency_key, - ) - - def register_hosted( - self, - *, - thread_id: str, - interrupt_id: str, - call_id: str, - name: str, - arguments: str, - aliases: list[str] | None = None, - already_approved_requests: list[dict[str, Any]] | None = None, - server_label: str | None = None, - idempotency_key: str | None = None, - ) -> ApprovalOccurrence: - """Register one server-generated hosted approval occurrence.""" - return self.register_hosted_aliases( - thread_ids=[thread_id], - interrupt_id=interrupt_id, - call_id=call_id, - name=name, - arguments=arguments, - aliases=aliases, - already_approved_requests=already_approved_requests, - server_label=server_label, - idempotency_key=idempotency_key, - ) - - def register_hosted_aliases( - self, - *, - thread_ids: list[str], - interrupt_id: str, - call_id: str, - name: str, - arguments: str, - aliases: list[str] | None = None, - already_approved_requests: list[dict[str, Any]] | None = None, - server_label: str | None = None, - idempotency_key: str | None = None, - ) -> ApprovalOccurrence: - """Register one hosted occurrence under its trusted scoped-thread aliases.""" - return self._register_aliases( - thread_ids=thread_ids, - interrupt_id=interrupt_id, - call_id=call_id, - name=name, - arguments=arguments, - owner=ApprovalExecutionOwner.HOSTED, - aliases=aliases, - already_approved_requests=already_approved_requests, - server_label=server_label, - idempotency_key=idempotency_key, - ) - - def register_unowned( - self, - *, - thread_id: str, - interrupt_id: str, - call_id: str, - name: str, - arguments: str, - aliases: list[str] | None = None, - already_approved_requests: list[dict[str, Any]] | None = None, - server_label: str | None = None, - ) -> ApprovalOccurrence: - """Register an occurrence that has no executable transition owner.""" - return self.register_unowned_aliases( - thread_ids=[thread_id], - interrupt_id=interrupt_id, - call_id=call_id, - name=name, - arguments=arguments, - aliases=aliases, - already_approved_requests=already_approved_requests, - server_label=server_label, - ) - - def register_deferred( - self, - *, - thread_id: str, - interrupt_id: str, - call_id: str, - name: str, - arguments: str, - aliases: list[str] | None = None, - already_approved_requests: list[dict[str, Any]] | None = None, - server_label: str | None = None, - idempotency_key: str | None = None, - ) -> ApprovalOccurrence: - """Register one occurrence owned by the in-run transition pipeline.""" - return self.register_deferred_aliases( - thread_ids=[thread_id], - interrupt_id=interrupt_id, - call_id=call_id, - name=name, - arguments=arguments, - aliases=aliases, - already_approved_requests=already_approved_requests, - server_label=server_label, - idempotency_key=idempotency_key, - ) - - def register_deferred_aliases( + def register( self, *, - thread_ids: list[str], + owner: ApprovalExecutionOwner, + thread_ids: list[str] | None = None, + thread_id: str | None = None, interrupt_id: str, call_id: str, name: str, @@ -414,45 +267,23 @@ def register_deferred_aliases( server_label: str | None = None, idempotency_key: str | None = None, ) -> ApprovalOccurrence: - """Register one deferred occurrence under its trusted scoped-thread aliases.""" + """Register one approval occurrence with its pending transition owner.""" + if (thread_ids is None) == (thread_id is None): + raise ValueError("Provide exactly one of thread_id or thread_ids when registering an approval.") + resolved_thread_ids = thread_ids if thread_ids is not None else [thread_id] return self._register_aliases( - thread_ids=thread_ids, + thread_ids=[value for value in resolved_thread_ids if value is not None], interrupt_id=interrupt_id, call_id=call_id, name=name, arguments=arguments, - owner=ApprovalExecutionOwner.DEFERRED, + owner=owner, aliases=aliases, already_approved_requests=already_approved_requests, server_label=server_label, idempotency_key=idempotency_key, ) - def register_unowned_aliases( - self, - *, - thread_ids: list[str], - interrupt_id: str, - call_id: str, - name: str, - arguments: str, - aliases: list[str] | None = None, - already_approved_requests: list[dict[str, Any]] | None = None, - server_label: str | None = None, - ) -> ApprovalOccurrence: - """Register an unowned occurrence under its trusted scoped-thread aliases.""" - return self._register_aliases( - thread_ids=thread_ids, - interrupt_id=interrupt_id, - call_id=call_id, - name=name, - arguments=arguments, - owner=ApprovalExecutionOwner.UNAVAILABLE, - aliases=aliases, - already_approved_requests=already_approved_requests, - server_label=server_label, - ) - @_serialized_registration def _register_aliases( self, @@ -590,7 +421,7 @@ def reconcile_snapshot( ApprovalSnapshotReconciliation( interrupt_id=interrupt_id, identity=None, - status=ApprovalSnapshotStatus.MISSING, + status=None, retire_interrupt=True, ) ) @@ -601,19 +432,11 @@ def reconcile_snapshot( @staticmethod def _snapshot_reconciliation(occurrence: ApprovalOccurrence) -> ApprovalSnapshotReconciliation: - status = ApprovalSnapshotStatus(occurrence.status.value) - terminal_statuses = { - ApprovalSnapshotStatus.SETTLED, - ApprovalSnapshotStatus.REJECTED, - ApprovalSnapshotStatus.CANCELLED, - ApprovalSnapshotStatus.EXPIRED, - ApprovalSnapshotStatus.INDETERMINATE, - } return ApprovalSnapshotReconciliation( interrupt_id=occurrence.identity.interrupt_id, identity=occurrence.identity, - status=status, - retire_interrupt=status in terminal_statuses, + status=occurrence.status, + retire_interrupt=occurrence.status.is_terminal, ) def claim(self, *, thread_id: str, decision: ResumeDecision) -> AuthorizedExecution: @@ -796,7 +619,7 @@ def expire_batch( return tuple(self._snapshot_reconciliation(occurrence) for occurrence in occurrences) def _remove_pending_aliases(self, occurrence: ApprovalOccurrence) -> None: - if occurrence.status is not ApprovalStatus.INDETERMINATE: + if occurrence.status.is_purgeable: occurrence.terminal_at = self._clock() with self._index_lock: for thread_id in occurrence.thread_ids: diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py b/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py index 7efa39a657..a09e18d5fc 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py @@ -65,7 +65,7 @@ def __init__( terminal_retention_seconds=terminal_retention_seconds, ) - def register_local( + def register( self, *, thread_ids: list[str], @@ -73,109 +73,14 @@ def register_local( arguments: str, request_id: str, interrupt_id: str, - already_approved_requests: list[dict[str, Any]] | None = None, - ) -> None: - """Register one local occurrence and its trusted aliases.""" - self._register( - thread_ids=thread_ids, - name=name, - arguments=arguments, - request_id=request_id, - interrupt_id=interrupt_id, - already_approved_requests=already_approved_requests, - server_label=None, - owner=ApprovalExecutionOwner.LOCAL, - ) - - def register_hosted( - self, - *, - thread_ids: list[str], - name: str, - arguments: str, - request_id: str, - interrupt_id: str, - server_label: str | None, - already_approved_requests: list[dict[str, Any]] | None = None, - ) -> None: - """Register one hosted occurrence and its trusted aliases.""" - self._register( - thread_ids=thread_ids, - name=name, - arguments=arguments, - request_id=request_id, - interrupt_id=interrupt_id, - already_approved_requests=already_approved_requests, - server_label=server_label, - owner=ApprovalExecutionOwner.HOSTED, - ) - - def register_unowned( - self, - *, - thread_ids: list[str], - name: str, - arguments: str, - request_id: str, - interrupt_id: str, - already_approved_requests: list[dict[str, Any]] | None = None, - ) -> None: - """Register one occurrence that has no executable transition owner.""" - self._register( - thread_ids=thread_ids, - name=name, - arguments=arguments, - request_id=request_id, - interrupt_id=interrupt_id, - already_approved_requests=already_approved_requests, - server_label=None, - owner=ApprovalExecutionOwner.UNAVAILABLE, - ) - - def register_deferred( - self, - *, - thread_ids: list[str], - name: str, - arguments: str, - request_id: str, - interrupt_id: str, - already_approved_requests: list[dict[str, Any]] | None = None, - ) -> None: - """Register one occurrence owned by the in-run transition pipeline.""" - self._register( - thread_ids=thread_ids, - name=name, - arguments=arguments, - request_id=request_id, - interrupt_id=interrupt_id, - already_approved_requests=already_approved_requests, - server_label=None, - owner=ApprovalExecutionOwner.DEFERRED, - ) - - def _register( - self, - *, - thread_ids: list[str], - name: str, - arguments: str, - request_id: str, - interrupt_id: str, - already_approved_requests: list[dict[str, Any]] | None, - server_label: str | None, owner: ApprovalExecutionOwner, + already_approved_requests: list[dict[str, Any]] | None = None, + server_label: str | None = None, ) -> None: + """Register one occurrence with its pending transition owner.""" unique_thread_ids = list(dict.fromkeys(thread_ids)) - if owner is ApprovalExecutionOwner.HOSTED: - register_aliases = self.lifecycle.register_hosted_aliases - elif owner is ApprovalExecutionOwner.DEFERRED: - register_aliases = self.lifecycle.register_deferred_aliases - elif owner is ApprovalExecutionOwner.UNAVAILABLE: - register_aliases = self.lifecycle.register_unowned_aliases - else: - register_aliases = self.lifecycle.register_local_aliases - register_aliases( + self.lifecycle.register( + owner=owner, thread_ids=unique_thread_ids, interrupt_id=interrupt_id, call_id=interrupt_id, diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py index 777450a686..bf1658c6c0 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py @@ -375,21 +375,34 @@ def _json_schema_for_value(value: Any) -> dict[str, Any]: def _approval_response_schema(arguments: Mapping[str, Any] | None = None) -> dict[str, Any]: """Build the response schema generic AG-UI clients use to render approval input.""" properties: dict[str, Any] = { - "accepted": { + "approved": { "type": "boolean", "description": "Whether the requested tool call is approved.", - } + }, + "accepted": { + "type": "boolean", + "description": "Legacy alias for approved.", + }, } - if arguments: + if arguments is not None: + edited_argument_properties: dict[str, Any] = {} for name, value in arguments.items(): argument_schema = _json_schema_for_value(value) argument_schema["description"] = f"Optional edited value for the '{name}' tool argument." properties[str(name)] = argument_schema + edited_argument_properties[str(name)] = _json_schema_for_value(value) + properties["editedArgs"] = { + "type": "object", + "description": "Full replacement of the tool arguments. Not merged.", + "properties": edited_argument_properties, + "required": list(edited_argument_properties), + "additionalProperties": False, + } return { "type": "object", "properties": properties, - "required": ["accepted"], + "anyOf": [{"required": ["approved"]}, {"required": ["accepted"]}], "additionalProperties": False, } diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py index bd3089775e..8b99d67c60 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py @@ -270,16 +270,11 @@ def _pending_workflow_interrupt_ids(pending_events: dict[str, Any]) -> set[str]: def _resume_error_for_pending_workflow_requests( resume_entries: list[dict[str, Any]], ) -> RunErrorEvent | None: - """Return a workflow resume error for explicit non-resolved canonical resume entries.""" + """Return a workflow resume error for unsupported canonical resume entries.""" for entry in resume_entries: interrupt_id = str(entry["interrupt_id"]) status = entry.get("status") - if status == "cancelled": - return RunErrorEvent( - message=f"Workflow resume for interruptId '{interrupt_id}' was cancelled.", - code="WORKFLOW_RESUME_CANCELLED", - ) - if status not in {None, "resolved"}: + if status not in {None, "resolved", "cancelled"}: return RunErrorEvent( message=f"Unsupported workflow resume status '{status}' for interruptId '{interrupt_id}'.", code="WORKFLOW_RESUME_INVALID", @@ -818,11 +813,19 @@ async def run_workflow_stream( return resume_error = _resume_error_for_pending_workflow_requests(resume_entries) if resume_error is not None: - if getattr(resume_error, "code", None) == "WORKFLOW_RESUME_CANCELLED": - _consume_cancelled_workflow_requests(workflow, resume_entries) yield RunStartedEvent(run_id=run_id, thread_id=thread_id) yield resume_error return + cancelled_request_ids = { + str(entry["interrupt_id"]) for entry in resume_entries if entry.get("status") == "cancelled" + } + if cancelled_request_ids: + _consume_cancelled_workflow_requests(workflow, resume_entries) + pending_before_run = { + request_id: request_event + for request_id, request_event in pending_before_run.items() + if str(getattr(request_event, "request_id", None) or request_id) not in cancelled_request_ids + } resume_responses = ( _resume_entries_to_workflow_responses(resume_entries) diff --git a/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py b/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py index f6992752e8..5116282925 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py +++ b/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py @@ -10,6 +10,8 @@ from agent_framework import Agent, ChatOptions, ChatResponseUpdate, Content, Message from pydantic import BaseModel +from agent_framework_ag_ui._approval_lifecycle import ApprovalExecutionOwner + async def test_agent_initialization_basic(streaming_chat_client_stub): """Test basic agent initialization without state schema.""" @@ -872,7 +874,8 @@ async def stream_fn( thread_id = "thread-rejection-test" - wrapper._approval_state_store.register_local( + wrapper._approval_state_store.register( + owner=ApprovalExecutionOwner.LOCAL, thread_ids=[thread_id], name="delete_all_data", arguments="{}", diff --git a/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py b/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py index d7254b1b51..9b81454a62 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py +++ b/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py @@ -17,7 +17,6 @@ ApprovalIndeterminateError, ApprovalLifecycle, ApprovalSettlementConflictError, - ApprovalSnapshotStatus, ApprovalStatus, ClaimRecoveryPolicy, HostedPendingToolTransitionOwner, @@ -26,10 +25,33 @@ ) +@pytest.mark.parametrize( + ("status", "is_terminal", "is_purgeable"), + [ + (ApprovalStatus.PENDING, False, False), + (ApprovalStatus.CLAIMED, False, False), + (ApprovalStatus.EXECUTING, False, False), + (ApprovalStatus.SETTLED, True, True), + (ApprovalStatus.REJECTED, True, True), + (ApprovalStatus.CANCELLED, True, True), + (ApprovalStatus.EXPIRED, True, True), + (ApprovalStatus.INDETERMINATE, True, False), + ], +) +def test_approval_status_owns_terminal_and_retention_semantics( + status: ApprovalStatus, + is_terminal: bool, + is_purgeable: bool, +) -> None: + assert status.is_terminal is is_terminal + assert status.is_purgeable is is_purgeable + + async def test_local_approval_crosses_lifecycle_before_execution_and_settlement() -> None: """One accepted local occurrence is claimed, executed by its owner, and settled.""" lifecycle = ApprovalLifecycle() - occurrence = lifecycle.register_local( + occurrence = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-1", interrupt_id="approval-1", call_id="call-1", @@ -79,7 +101,8 @@ def test_registration_keeps_trusted_aliases_and_projection_metadata_in_lifecycle }, } - occurrence = lifecycle.register_hosted( + occurrence = lifecycle.register( + owner=ApprovalExecutionOwner.HOSTED, thread_id="thread-1", interrupt_id="call-1", call_id="call-1", @@ -96,10 +119,40 @@ def test_registration_keeps_trusted_aliases_and_projection_metadata_in_lifecycle assert occurrence.server_label == "hosted-tools" +@pytest.mark.parametrize( + "owner", + [ + ApprovalExecutionOwner.LOCAL, + ApprovalExecutionOwner.HOSTED, + ApprovalExecutionOwner.DEFERRED, + ApprovalExecutionOwner.UNAVAILABLE, + ], +) +def test_registration_uses_one_owner_based_lifecycle_operation(owner: ApprovalExecutionOwner) -> None: + """Callers select lifecycle ownership explicitly through one registration operation.""" + lifecycle = ApprovalLifecycle() + + occurrence = lifecycle.register( + owner=owner, + thread_ids=["thread-primary", "thread-provider"], + interrupt_id="approval-1", + call_id="call-1", + name="write_record", + arguments='{"value":"first"}', + aliases=["request-1"], + server_label="hosted-tools" if owner is ApprovalExecutionOwner.HOSTED else None, + ) + + assert occurrence.owner is owner + assert lifecycle.pending_occurrence(thread_id="thread-primary", interrupt_id="request-1") is occurrence + assert lifecycle.pending_occurrence(thread_id="thread-provider", interrupt_id="approval-1") is occurrence + + def test_active_occurrence_is_not_evicted_when_capacity_is_exhausted() -> None: """Storage pressure fails explicitly instead of discarding pending authority.""" lifecycle = ApprovalLifecycle(max_entries=1) - occurrence = lifecycle.register_local( + occurrence = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-1", interrupt_id="approval-1", call_id="call-1", @@ -108,7 +161,8 @@ def test_active_occurrence_is_not_evicted_when_capacity_is_exhausted() -> None: ) with pytest.raises(ApprovalCapacityError): - lifecycle.register_local( + lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-2", interrupt_id="approval-2", call_id="call-2", @@ -123,7 +177,8 @@ async def test_terminal_outcome_expires_only_after_configured_retention_window() """Duplicate execution protection lasts for the configured terminal retention window.""" now = 100.0 lifecycle = ApprovalLifecycle(max_entries=1, terminal_retention_seconds=30, clock=lambda: now) - occurrence = lifecycle.register_local( + occurrence = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-1", interrupt_id="approval-1", call_id="call-1", @@ -141,7 +196,8 @@ async def execute() -> list[Content]: assert lifecycle.claim_batch(thread_id="thread-1", decisions=[decision]).retained_outcomes == (outcome,) now = 131.0 - replacement = lifecycle.register_local( + replacement = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-2", interrupt_id="approval-2", call_id="call-2", @@ -160,7 +216,8 @@ def test_indeterminate_occurrence_remains_protected_after_terminal_retention_win """Uncertain execution is never aged out as a retryable terminal tombstone.""" now = 100.0 lifecycle = ApprovalLifecycle(max_entries=1, terminal_retention_seconds=30, clock=lambda: now) - occurrence = lifecycle.register_local( + occurrence = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-1", interrupt_id="approval-1", call_id="call-1", @@ -176,7 +233,8 @@ def test_indeterminate_occurrence_remains_protected_after_terminal_retention_win now = 1_000.0 with pytest.raises(ApprovalCapacityError): - lifecycle.register_local( + lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-2", interrupt_id="approval-2", call_id="call-2", @@ -194,7 +252,8 @@ async def test_transition_telemetry_covers_lifecycle_without_sensitive_payloads( caplog.set_level("INFO", logger="agent_framework_ag_ui._approval_lifecycle") lifecycle = ApprovalLifecycle() secret = "sensitive-value" - settled = lifecycle.register_local( + settled = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-settled", interrupt_id="approval-settled", call_id="call-settled", @@ -214,7 +273,8 @@ async def execute() -> list[Content]: await LocalPendingToolTransitionOwner(execute).execute(intent, lifecycle=lifecycle) lifecycle.claim_batch(thread_id="thread-settled", decisions=[decision]) - lifecycle.register_local( + lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-rejected", interrupt_id="approval-rejected", call_id="call-rejected", @@ -225,7 +285,8 @@ async def execute() -> list[Content]: thread_id="thread-rejected", decisions=[ResumeDecision(interrupt_id="approval-rejected", accepted=False, arguments="{}")], ) - lifecycle.register_local( + lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-cancelled", interrupt_id="approval-cancelled", call_id="call-cancelled", @@ -233,7 +294,8 @@ async def execute() -> list[Content]: arguments="{}", ) lifecycle.cancel_batch(thread_id="thread-cancelled", interrupt_ids=["approval-cancelled"]) - lifecycle.register_local( + lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-expired", interrupt_id="approval-expired", call_id="call-expired", @@ -241,7 +303,8 @@ async def execute() -> list[Content]: arguments="{}", ) lifecycle.expire_batch(thread_id="thread-expired", interrupt_ids=["approval-expired"]) - uncertain = lifecycle.register_local( + uncertain = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-uncertain", interrupt_id="approval-uncertain", call_id="call-uncertain", @@ -260,7 +323,8 @@ async def execute() -> list[Content]: decisions=[ResumeDecision(interrupt_id="approval-missing", accepted=True, arguments="{}")], ) capacity_lifecycle = ApprovalLifecycle(max_entries=1) - capacity_lifecycle.register_local( + capacity_lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-capacity-1", interrupt_id="approval-capacity-1", call_id="call-capacity-1", @@ -268,7 +332,8 @@ async def execute() -> list[Content]: arguments="{}", ) with pytest.raises(ApprovalCapacityError): - capacity_lifecycle.register_local( + capacity_lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-capacity-2", interrupt_id="approval-capacity-2", call_id="call-capacity-2", @@ -301,7 +366,8 @@ async def execute() -> list[Content]: def test_claim_and_settlement_conflicts_have_typed_outcomes() -> None: """Adapters can distinguish transition conflicts without parsing error messages.""" lifecycle = ApprovalLifecycle() - lifecycle.register_local( + lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-1", interrupt_id="approval-1", call_id="call-1", @@ -320,14 +386,16 @@ def test_claim_and_settlement_conflicts_have_typed_outcomes() -> None: def test_same_thread_transitions_serialize_without_blocking_an_independent_thread() -> None: """One scoped thread is serialized while another can claim concurrently.""" lifecycle = ApprovalLifecycle() - first = lifecycle.register_local( + first = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-1", interrupt_id="approval-1", call_id="call-1", name="write_record", arguments="{}", ) - lifecycle.register_local( + lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-2", interrupt_id="approval-2", call_id="call-2", @@ -371,7 +439,8 @@ def repeat_first_claim(): async def test_hosted_approval_is_forwarded_only_by_its_owner_and_settles_same_occurrence() -> None: """Hosted authority cannot execute locally and records forwarding against its occurrence.""" lifecycle = ApprovalLifecycle() - occurrence = lifecycle.register_hosted( + occurrence = lifecycle.register( + owner=ApprovalExecutionOwner.HOSTED, thread_id="thread-1", interrupt_id="approval-1", call_id="call-1", @@ -426,7 +495,7 @@ async def forward_to_hosted_owner() -> list[Content]: assert outcome.identity == occurrence.identity assert outcome.result_group == (remote_result,) assert [result.content for result in outcome.replayable_results] == [remote_result] - assert outcome.snapshot_reconciliation.status is ApprovalSnapshotStatus.SETTLED + assert outcome.snapshot_reconciliation.status is ApprovalStatus.SETTLED assert outcome.snapshot_reconciliation.retire_interrupt is True assert lifecycle.get(occurrence.identity).status is ApprovalStatus.SETTLED @@ -434,14 +503,16 @@ async def forward_to_hosted_owner() -> list[Content]: async def test_execution_failure_becomes_indeterminate_and_keeps_unexecuted_sibling_claimed() -> None: """A possibly started side effect is not retried and does not erase a claimed sibling.""" lifecycle = ApprovalLifecycle() - first = lifecycle.register_local( + first = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-1", interrupt_id="approval-1", call_id="call-1", name="write_record", arguments='{"value":"first"}', ) - second = lifecycle.register_local( + second = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-1", interrupt_id="approval-2", call_id="call-2", @@ -487,7 +558,8 @@ async def execute_second() -> list[Content]: def test_claim_can_be_released_before_execution_only_with_explicit_safe_policy() -> None: """Reserved authority can be reclaimed when the owner proves execution never began.""" lifecycle = ApprovalLifecycle() - occurrence = lifecycle.register_local( + occurrence = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-1", interrupt_id="approval-1", call_id="call-1", @@ -507,7 +579,8 @@ def test_claim_can_be_released_before_execution_only_with_explicit_safe_policy() def test_recovering_execution_without_a_result_becomes_indeterminate() -> None: """Recovery preserves an uncertain occurrence instead of granting authority again.""" lifecycle = ApprovalLifecycle() - occurrence = lifecycle.register_local( + occurrence = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-1", interrupt_id="approval-1", call_id="call-1", @@ -529,7 +602,8 @@ def test_recovering_execution_without_a_result_becomes_indeterminate() -> None: async def test_explicit_idempotency_key_allows_retry_after_execution_interruption() -> None: """A predeclared idempotency key permits retrying a potentially started side effect.""" lifecycle = ApprovalLifecycle() - occurrence = lifecycle.register_local( + occurrence = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-1", interrupt_id="approval-1", call_id="call-1", @@ -564,7 +638,8 @@ async def execute_idempotently() -> list[Content]: async def test_hosted_idempotency_key_allows_retry_after_forwarding_interruption() -> None: """A hosted owner uses the same explicit recovery rule as the local owner.""" lifecycle = ApprovalLifecycle() - occurrence = lifecycle.register_hosted( + occurrence = lifecycle.register( + owner=ApprovalExecutionOwner.HOSTED, thread_id="thread-1", interrupt_id="approval-1", call_id="call-1", @@ -598,14 +673,16 @@ async def forward_idempotently() -> list[Content]: def test_batch_validation_is_atomic_before_claiming_any_occurrence() -> None: """One invalid decision leaves every occurrence pending and eligible for a corrected batch.""" lifecycle = ApprovalLifecycle() - first = lifecycle.register_local( + first = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="tenant-a\x1fthread-1", interrupt_id="approval-1", call_id="call-1", name="write_record", arguments='{"value":"first"}', ) - second = lifecycle.register_local( + second = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="tenant-a\x1fthread-1", interrupt_id="approval-2", call_id="call-2", @@ -629,7 +706,8 @@ def test_batch_validation_is_atomic_before_claiming_any_occurrence() -> None: def test_accepted_declaration_without_execution_owner_remains_pending() -> None: """Approval alone does not grant local authority to a declaration-only call.""" lifecycle = ApprovalLifecycle() - occurrence = lifecycle.register_unowned( + occurrence = lifecycle.register( + owner=ApprovalExecutionOwner.UNAVAILABLE, thread_id="thread-1", interrupt_id="approval-1", call_id="call-1", @@ -649,14 +727,16 @@ def test_accepted_declaration_without_execution_owner_remains_pending() -> None: def test_mixed_batch_accounts_for_rejection_under_original_call_identity() -> None: """A rejected occurrence remains represented while its accepted sibling is claimed.""" lifecycle = ApprovalLifecycle() - accepted = lifecycle.register_local( + accepted = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-1", interrupt_id="approval-1", call_id="call-1", name="write_record", arguments='{"value":"first"}', ) - rejected = lifecycle.register_local( + rejected = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-1", interrupt_id="approval-2", call_id="call-2", @@ -680,14 +760,16 @@ def test_mixed_batch_accounts_for_rejection_under_original_call_identity() -> No def test_batch_claims_preserve_order_and_scope_reused_raw_call_ids() -> None: """Raw call ids reused in another scoped thread cannot correlate approval authority.""" lifecycle = ApprovalLifecycle() - tenant_a = lifecycle.register_local( + tenant_a = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="tenant-a\x1fthread-1", interrupt_id="approval-shared", call_id="call-shared", name="write_record", arguments='{"tenant":"a"}', ) - tenant_b = lifecycle.register_local( + tenant_b = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="tenant-b\x1fthread-1", interrupt_id="approval-shared", call_id="call-shared", @@ -715,14 +797,16 @@ def test_batch_claims_preserve_order_and_scope_reused_raw_call_ids() -> None: def test_batch_cancellation_preserves_each_original_occurrence() -> None: """Cancelling selected occurrences is terminal without consuming an unrelated sibling.""" lifecycle = ApprovalLifecycle() - cancelled = lifecycle.register_local( + cancelled = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="tenant-a\x1fthread-1", interrupt_id="approval-1", call_id="call-1", name="write_record", arguments='{"value":"first"}', ) - pending = lifecycle.register_local( + pending = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="tenant-a\x1fthread-1", interrupt_id="approval-2", call_id="call-2", @@ -738,21 +822,23 @@ def test_batch_cancellation_preserves_each_original_occurrence() -> None: assert lifecycle.get(cancelled.identity).status is ApprovalStatus.CANCELLED assert lifecycle.get(pending.identity).status is ApprovalStatus.PENDING assert [(item.identity, item.status, item.retire_interrupt) for item in reconciliations] == [ - (cancelled.identity, ApprovalSnapshotStatus.CANCELLED, True) + (cancelled.identity, ApprovalStatus.CANCELLED, True) ] def test_snapshot_reconciliation_reports_terminal_pending_and_missing_occurrences() -> None: """Snapshot projection receives lifecycle semantics without recreating authority.""" lifecycle = ApprovalLifecycle() - settled = lifecycle.register_local( + settled = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-1", interrupt_id="approval-settled", call_id="call-settled", name="write_record", arguments="{}", ) - pending = lifecycle.register_local( + pending = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-1", interrupt_id="approval-pending", call_id="call-pending", @@ -776,16 +862,18 @@ def test_snapshot_reconciliation_reports_terminal_pending_and_missing_occurrence assert outcome.snapshot_reconciliation == reconciliations[0] assert [(item.identity, item.status, item.retire_interrupt) for item in reconciliations] == [ - (settled.identity, ApprovalSnapshotStatus.SETTLED, True), - (pending.identity, ApprovalSnapshotStatus.PENDING, False), - (None, ApprovalSnapshotStatus.MISSING, True), + (settled.identity, ApprovalStatus.SETTLED, True), + (pending.identity, ApprovalStatus.PENDING, False), + (None, None, True), ] + assert reconciliations[-1].is_missing def test_one_occurrence_can_be_claimed_through_a_trusted_thread_alias() -> None: """Provider conversation aliases address one occurrence rather than duplicating authority.""" lifecycle = ApprovalLifecycle() - occurrence = lifecycle.register_local_aliases( + occurrence = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_ids=["tenant-a\x1fag-ui-thread", "tenant-a\x1fprovider-thread"], interrupt_id="approval-1", call_id="call-1", @@ -800,7 +888,8 @@ def test_one_occurrence_can_be_claimed_through_a_trusted_thread_alias() -> None: assert [intent.identity for intent in intents] == [occurrence.identity] with pytest.raises(ValueError, match="not pending"): - lifecycle.register_local_aliases( + lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_ids=["tenant-a\x1fag-ui-thread", "tenant-a\x1fprovider-thread"], interrupt_id="approval-1", call_id="call-1", @@ -817,7 +906,8 @@ def test_one_occurrence_can_be_claimed_through_a_trusted_thread_alias() -> None: async def test_settled_raw_call_id_can_be_reused_for_a_new_occurrence() -> None: """Sequential reuse creates a fresh logical occurrence instead of reviving settled authority.""" lifecycle = ApprovalLifecycle() - first = lifecycle.register_local( + first = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-1", interrupt_id="approval-shared", call_id="call-shared", @@ -837,7 +927,8 @@ async def execute_first() -> list[Content]: return [Content.from_function_result(call_id="call-shared", result="wrote first")] await LocalPendingToolTransitionOwner(execute_first).execute(first_intent, lifecycle=lifecycle) - second = lifecycle.register_local( + second = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-1", interrupt_id="approval-shared", call_id="call-shared", @@ -853,7 +944,8 @@ async def execute_first() -> list[Content]: async def test_identical_accepted_retry_returns_retained_outcome_without_execution() -> None: """A settled accepted decision reprojects its result instead of granting authority again.""" lifecycle = ApprovalLifecycle() - occurrence = lifecycle.register_local( + occurrence = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-1", interrupt_id="approval-1", call_id="call-1", @@ -885,7 +977,8 @@ async def execute_once() -> list[Content]: def test_accepted_retry_after_rejection_fails_as_a_conflict() -> None: """A terminal rejection cannot be changed into execution authority by a retry.""" lifecycle = ApprovalLifecycle() - occurrence = lifecycle.register_local( + occurrence = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-1", interrupt_id="approval-1", call_id="call-1", @@ -909,7 +1002,8 @@ def test_accepted_retry_after_rejection_fails_as_a_conflict() -> None: def test_identical_rejection_retry_returns_retained_outcome() -> None: """A repeated rejection preserves and returns the original rejection result.""" lifecycle = ApprovalLifecycle() - occurrence = lifecycle.register_local( + occurrence = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-1", interrupt_id="approval-1", call_id="call-1", @@ -931,7 +1025,8 @@ def test_identical_rejection_retry_returns_retained_outcome() -> None: def test_changed_tool_name_fails_before_authority_is_claimed() -> None: """A typed decision for another tool cannot claim the registered occurrence.""" lifecycle = ApprovalLifecycle() - occurrence = lifecycle.register_local( + occurrence = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-1", interrupt_id="approval-1", call_id="call-1", @@ -958,7 +1053,8 @@ def test_changed_tool_name_fails_before_authority_is_claimed() -> None: def test_identical_cancel_retry_keeps_terminal_cancellation() -> None: """Retrying an explicit cancellation is idempotent and cannot restore authority.""" lifecycle = ApprovalLifecycle() - occurrence = lifecycle.register_local( + occurrence = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-1", interrupt_id="approval-1", call_id="call-1", @@ -975,7 +1071,8 @@ def test_identical_cancel_retry_keeps_terminal_cancellation() -> None: def test_expired_authority_cannot_be_claimed() -> None: """Expiration is terminal and an otherwise valid decision cannot revive it.""" lifecycle = ApprovalLifecycle() - occurrence = lifecycle.register_local( + occurrence = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-1", interrupt_id="approval-1", call_id="call-1", diff --git a/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py b/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py index 032dbbc21b..9c58277cc4 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py +++ b/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py @@ -12,6 +12,7 @@ from agent_framework_ag_ui._agent import AgentConfig from agent_framework_ag_ui._agent_run import run_agent_stream +from agent_framework_ag_ui._approval_lifecycle import ApprovalExecutionOwner from agent_framework_ag_ui._approval_state import InMemoryAGUIApprovalStateStore @@ -42,7 +43,8 @@ async def _run_resume( ) store = InMemoryAGUIApprovalStateStore() for call_id, city in calls: - store.register_local( + store.register( + owner=ApprovalExecutionOwner.LOCAL, thread_ids=[thread_id], name="get_weather", arguments=json.dumps({"city": city}, sort_keys=True, separators=(",", ":")), diff --git a/python/packages/ag-ui/tests/ag_ui/test_approval_state.py b/python/packages/ag-ui/tests/ag_ui/test_approval_state.py index bf1a515be0..58a61391f3 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_approval_state.py +++ b/python/packages/ag-ui/tests/ag_ui/test_approval_state.py @@ -4,7 +4,7 @@ import pytest -from agent_framework_ag_ui._approval_lifecycle import ApprovalCapacityError +from agent_framework_ag_ui._approval_lifecycle import ApprovalCapacityError, ApprovalExecutionOwner from agent_framework_ag_ui._approval_state import InMemoryAGUIApprovalStateStore, approval_state_thread_id @@ -31,9 +31,29 @@ def test_approval_state_store_rejects_invalid_max_entries() -> None: InMemoryAGUIApprovalStateStore(max_entries=0) +def test_approval_state_store_registers_explicit_execution_owner() -> None: + store = InMemoryAGUIApprovalStateStore() + + store.register( + owner=ApprovalExecutionOwner.DEFERRED, + thread_ids=["thread-1", "provider-thread-1"], + name="write_record", + arguments="{}", + request_id="request-1", + interrupt_id="approval-1", + server_label=None, + ) + + occurrence = store.lifecycle.pending_occurrence(thread_id="thread-1", interrupt_id="approval-1") + assert occurrence is not None + assert occurrence.owner is ApprovalExecutionOwner.DEFERRED + assert store.lifecycle.pending_occurrence(thread_id="provider-thread-1", interrupt_id="request-1") is occurrence + + def test_approval_state_store_does_not_evict_active_entries() -> None: store = InMemoryAGUIApprovalStateStore(max_entries=1) - store.register_local( + store.register( + owner=ApprovalExecutionOwner.LOCAL, thread_ids=["thread-1"], name="write_record", arguments="{}", @@ -42,7 +62,8 @@ def test_approval_state_store_does_not_evict_active_entries() -> None: ) with pytest.raises(ApprovalCapacityError): - store.register_local( + store.register( + owner=ApprovalExecutionOwner.LOCAL, thread_ids=["thread-2"], name="write_record", arguments="{}", diff --git a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py index 5be6655a50..b1b78c32f3 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -1421,9 +1421,18 @@ async def test_endpoint_agent_approval_pause_emits_canonical_interrupt_outcome() assert interrupt["reason"] == "tool_call" assert interrupt["toolCallId"] == "call_write_doc" assert interrupt["message"] == "Approve running write_doc?" - assert interrupt["responseSchema"]["required"] == ["accepted"] - assert interrupt["responseSchema"]["properties"]["accepted"]["type"] == "boolean" - assert interrupt["responseSchema"]["properties"]["content"]["type"] == "string" + response_schema = interrupt["responseSchema"] + assert response_schema["anyOf"] == [{"required": ["approved"]}, {"required": ["accepted"]}] + assert response_schema["properties"]["approved"]["type"] == "boolean" + assert response_schema["properties"]["accepted"]["type"] == "boolean" + assert response_schema["properties"]["content"]["type"] == "string" + assert response_schema["properties"]["editedArgs"] == { + "type": "object", + "description": "Full replacement of the tool arguments. Not merged.", + "properties": {"content": {"type": "string"}}, + "required": ["content"], + "additionalProperties": False, + } metadata_value = interrupt["metadata"]["agent_framework"] assert metadata_value["type"] == "function_approval_request" assert metadata_value["function_call"] == { @@ -1705,6 +1714,50 @@ async def test_endpoint_agent_approval_resume_entry_executes_approved_tool(): assert "outcome" not in [event for event in events if event.get("type") == "RUN_FINISHED"][-1] +async def test_endpoint_agent_approval_resume_remains_retryable_when_local_tool_is_temporarily_unavailable(): + """A local approval can be retried after its executor disappears before resume.""" + client, agent, executed_cities = _build_weather_approval_endpoint(snapshot_store=InMemoryAGUIThreadSnapshotStore()) + weather_tool = agent.default_options["tools"][0] + agent.default_options["tools"] = [] + + unavailable_response = client.post( + "/approval", + json={ + "runId": "run-unavailable", + "threadId": "thread-weather", + "messages": [], + "resume": [{"interruptId": "call_get_weather", "status": "resolved", "payload": {"accepted": True}}], + }, + ) + + assert unavailable_response.status_code == 200 + unavailable_events = _decode_sse_events(unavailable_response) + run_errors = [event for event in unavailable_events if event.get("type") == "RUN_ERROR"] + assert len(run_errors) == 1 + assert run_errors[0]["code"] == "APPROVAL_TOOL_UNAVAILABLE" + assert "temporarily unavailable" in run_errors[0]["message"] + assert executed_cities == [] + + agent.default_options["tools"] = [weather_tool] + retry_response = client.post( + "/approval", + json={ + "runId": "run-retry", + "threadId": "thread-weather", + "messages": [], + "resume": [{"interruptId": "call_get_weather", "status": "resolved", "payload": {"accepted": True}}], + }, + ) + + assert retry_response.status_code == 200 + retry_events = _decode_sse_events(retry_response) + assert not [event for event in retry_events if event.get("type") == "RUN_ERROR"] + assert [ + (event["toolCallId"], event["content"]) for event in retry_events if event.get("type") == "TOOL_CALL_RESULT" + ] == [("call_get_weather", "Sunny in Seattle")] + assert executed_cities == ["Seattle"] + + async def test_endpoint_agent_approval_resume_releases_already_approved_sibling(streaming_chat_client_stub): """Resuming a visible approval should also complete never-require siblings from the same batch.""" client, executed, messages_received, state = _build_mixed_approval_batch_endpoint(streaming_chat_client_stub) @@ -1932,9 +1985,7 @@ async def test_endpoint_agent_approval_cancel_discards_queued_tool_approval(stre assert cancel_response.status_code == 200 cancel_events = _decode_sse_events(cancel_response) - run_errors = [event for event in cancel_events if event.get("type") == "RUN_ERROR"] - assert len(run_errors) == 1 - assert run_errors[0]["code"] == "APPROVAL_RESUME_CANCELLED" + assert [event.get("type") for event in cancel_events] == ["RUN_STARTED", "RUN_FINISHED"] assert executed == [] assert messages_received == [] @@ -2110,7 +2161,7 @@ async def test_endpoint_agent_approval_rejection_releases_already_approved_sibli async def test_endpoint_agent_approval_cancellation_does_not_release_already_approved_sibling( streaming_chat_client_stub, ): - """Cancelling a visible approval remains fail-closed and emits no sibling result.""" + """Cancelling a visible approval completes normally without releasing a hidden sibling.""" client, executed, messages_received, state = _build_mixed_approval_batch_endpoint(streaming_chat_client_stub) pause_response = client.post( "/approval", @@ -2137,9 +2188,7 @@ async def test_endpoint_agent_approval_cancellation_does_not_release_already_app assert cancel_response.status_code == 200 cancel_events = _decode_sse_events(cancel_response) - run_errors = [event for event in cancel_events if event.get("type") == "RUN_ERROR"] - assert len(run_errors) == 1 - assert run_errors[0]["code"] == "APPROVAL_RESUME_CANCELLED" + assert [event.get("type") for event in cancel_events] == ["RUN_STARTED", "RUN_FINISHED"] assert not [event for event in cancel_events if event.get("type") == "TOOL_CALL_RESULT"] assert executed == [] assert messages_received == [] @@ -2570,8 +2619,69 @@ async def test_endpoint_agent_approval_resume_entry_applies_edited_arguments(): assert executed_cities == ["Portland"] -async def test_endpoint_agent_approval_cancelled_resume_entry_emits_run_error(): - """A cancelled canonical approval resume should fail safely instead of proceeding.""" +async def test_endpoint_agent_approval_resume_entry_applies_standard_full_replacement_edited_args(): + """The standard approved/editedArgs payload replaces the complete pending tool arguments.""" + client, _, executed_cities = _build_weather_approval_endpoint() + + response = client.post( + "/approval", + json={ + "runId": "run-standard-edit", + "threadId": "thread-weather", + "messages": [], + "resume": [ + { + "interruptId": "call_get_weather", + "status": "resolved", + "payload": {"approved": True, "editedArgs": {"city": "Portland"}}, + } + ], + }, + ) + + assert response.status_code == 200 + events = _decode_sse_events(response) + assert not [event for event in events if event.get("type") == "RUN_ERROR"] + tool_results = [event for event in events if event.get("type") == "TOOL_CALL_RESULT"] + assert [(event["toolCallId"], event["content"]) for event in tool_results] == [ + ("call_get_weather", "Sunny in Portland") + ] + assert executed_cities == ["Portland"] + + +async def test_endpoint_agent_approval_replayed_standard_edited_resume_is_idempotent(): + """Replaying a standard edited approval returns its retained result without executing again.""" + client, _, executed_cities = _build_weather_approval_endpoint() + resume = [ + { + "interruptId": "call_get_weather", + "status": "resolved", + "payload": {"approved": True, "editedArgs": {"city": "Portland"}}, + } + ] + + first_response = client.post( + "/approval", + json={"runId": "run-standard-edit", "threadId": "thread-weather", "messages": [], "resume": resume}, + ) + retry_response = client.post( + "/approval", + json={"runId": "run-standard-retry", "threadId": "thread-weather", "messages": [], "resume": resume}, + ) + + assert first_response.status_code == 200 + assert retry_response.status_code == 200 + retry_events = _decode_sse_events(retry_response) + assert not [event for event in retry_events if event.get("type") == "RUN_ERROR"] + retry_results = [event for event in retry_events if event.get("type") == "TOOL_CALL_RESULT"] + assert [(event["toolCallId"], event["content"]) for event in retry_results] == [ + ("call_get_weather", "Sunny in Portland") + ] + assert executed_cities == ["Portland"] + + +async def test_endpoint_agent_approval_cancelled_resume_entry_completes_without_execution(): + """A cancelled canonical approval resume should complete without executing the pending tool.""" client, _, executed_cities = _build_weather_approval_endpoint() response = client.post( @@ -2586,9 +2696,7 @@ async def test_endpoint_agent_approval_cancelled_resume_entry_emits_run_error(): assert response.status_code == 200 events = _decode_sse_events(response) - run_errors = [event for event in events if event.get("type") == "RUN_ERROR"] - assert len(run_errors) == 1 - assert run_errors[0]["code"] == "APPROVAL_RESUME_CANCELLED" + assert [event.get("type") for event in events] == ["RUN_STARTED", "RUN_FINISHED"] assert executed_cities == [] assert not [event for event in events if event.get("type") == "TOOL_CALL_RESULT"] @@ -2850,8 +2958,8 @@ def record_city(city: str) -> str: assert executed == [] -async def test_endpoint_agent_approval_cancelled_resume_preserves_uncancelled_interrupt(): - """Cancelling one approval clears only that interrupt and leaves others resumable.""" +async def test_endpoint_agent_approval_mixed_cancelled_and_resolved_resume_executes_resolved_tool(): + """A mixed resume executes resolved approvals and treats cancelled calls as not executed.""" executed: list[str] = [] def record_city(city: str) -> str: @@ -2911,6 +3019,7 @@ def record_city(city: str) -> str: "call_portland", } + agent.updates = [AgentResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant")] cancel_response = client.post( "/approval-snapshots", json={ @@ -2926,14 +3035,16 @@ def record_city(city: str) -> str: assert cancel_response.status_code == 200 cancel_events = _decode_sse_events(cancel_response) - run_errors = [event for event in cancel_events if event.get("type") == "RUN_ERROR"] - assert len(run_errors) == 1 - assert run_errors[0]["code"] == "APPROVAL_RESUME_CANCELLED" - assert executed == [] + assert not [event for event in cancel_events if event.get("type") == "RUN_ERROR"] + tool_results = [event for event in cancel_events if event.get("type") == "TOOL_CALL_RESULT"] + assert [(event["toolCallId"], event["content"]) for event in tool_results] == [ + ("call_portland", "Recorded Portland") + ] + assert executed == ["Portland"] approval_thread_id = approval_state_thread_id(scope="tenant-a", thread_id="thread-two-approvals") pending_ids = wrapped_agent._approval_state_store.lifecycle.pending_interrupt_ids(thread_id=approval_thread_id) assert "call_seattle" not in pending_ids - assert "call_portland" in pending_ids + assert "call_portland" not in pending_ids hydrate_response = client.post( "/approval-snapshots", @@ -2941,23 +3052,8 @@ def record_city(city: str) -> str: ) assert hydrate_response.status_code == 200 hydrate_events = _decode_sse_events(hydrate_response) - assert [interrupt["id"] for interrupt in _run_finished_interrupts(hydrate_events[-1])] == ["call_portland"] - - agent.updates = [AgentResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant")] - resume_response = client.post( - "/approval-snapshots", - json={ - "runId": "run-resume-remaining", - "threadId": "thread-two-approvals", - "messages": [], - "resume": [{"interruptId": "call_portland", "status": "resolved", "payload": {"accepted": True}}], - }, - ) - - assert resume_response.status_code == 200 - assert executed == ["Portland"] - resume_events = _decode_sse_events(resume_response) - assert not [event for event in resume_events if event.get("type") == "RUN_ERROR"] + assert hydrate_events[-1]["type"] == "RUN_FINISHED" + assert "outcome" not in hydrate_events[-1] def _build_workflow_request_info_app() -> FastAPI: @@ -3037,8 +3133,8 @@ async def test_endpoint_workflow_request_info_emits_canonical_interrupt_and_resu assert "outcome" not in [event for event in resume_events if event.get("type") == "RUN_FINISHED"][-1] -async def test_endpoint_workflow_request_info_cancelled_resume_emits_run_error(): - """Cancelled workflow resumes fail explicitly and do not wedge the next turn.""" +async def test_endpoint_workflow_request_info_cancelled_resume_completes_normally(): + """Cancelled workflow resumes complete without output and do not wedge the next turn.""" app = _build_workflow_request_info_app() with TestClient(app) as client: @@ -3064,9 +3160,7 @@ async def test_endpoint_workflow_request_info_cancelled_resume_emits_run_error() assert resume_response.status_code == 200 events = _decode_sse_events(resume_response) - run_errors = [event for event in events if event.get("type") == "RUN_ERROR"] - assert len(run_errors) == 1 - assert run_errors[0]["code"] == "WORKFLOW_RESUME_CANCELLED" + assert [event.get("type") for event in events] == ["RUN_STARTED", "RUN_FINISHED"] assert not [event for event in events if event.get("type") == "TEXT_MESSAGE_CONTENT"] next_response = client.post( @@ -4965,9 +5059,8 @@ def get_weather(city: str) -> str: ) assert cancel_response.status_code == 200 cancel_events = _decode_sse_events(cancel_response) - assert [event for event in cancel_events if event.get("type") == "RUN_ERROR"][0][ - "code" - ] == "APPROVAL_RESUME_CANCELLED" + assert [event.get("type") for event in cancel_events][-1] == "RUN_FINISHED" + assert not [event for event in cancel_events if event.get("type") == "RUN_ERROR"] assert executed_cities == [] approval_thread_id = approval_state_thread_id(scope="tenant-a", thread_id="agent-approval-thread") assert not wrapped_agent._approval_state_store.lifecycle.pending_interrupt_ids(thread_id=approval_thread_id) @@ -5229,9 +5322,7 @@ async def requester(message: Any, ctx: WorkflowContext[Any, Any]) -> None: ) assert cancel_response.status_code == 200 cancel_events = _decode_sse_events(cancel_response) - assert [event for event in cancel_events if event.get("type") == "RUN_ERROR"][0][ - "code" - ] == "WORKFLOW_RESUME_CANCELLED" + assert [event.get("type") for event in cancel_events] == ["RUN_STARTED", "RUN_FINISHED"] hydrate_response = client.post( "/workflow-snapshots", @@ -5800,7 +5891,9 @@ async def stream_fn( pause_finished = [event for event in _decode_sse_events(pause_response) if event.get("type") == "RUN_FINISHED"] pause_interrupts = _run_finished_interrupts(pause_finished[-1]) assert [interrupt["id"] for interrupt in pause_interrupts] == [call_id] - assert set(pause_interrupts[0]["responseSchema"]["properties"]) == {"accepted"} + hosted_response_schema = pause_interrupts[0]["responseSchema"] + assert set(hosted_response_schema["properties"]) == {"approved", "accepted"} + assert hosted_response_schema["anyOf"] == [{"required": ["approved"]}, {"required": ["accepted"]}] state["phase"] = "resume" resume_response = client.post( @@ -5845,6 +5938,86 @@ async def stream_fn( assert local_executions == [] +async def test_endpoint_hosted_approval_becomes_indeterminate_when_provider_stream_fails( + streaming_chat_client_stub, +) -> None: + """A forwarded approval with an interrupted provider stream cannot be retried automatically.""" + call_id = "mcpr_docs_failure" + state = {"phase": "pause"} + hosted_call = Content.from_function_call( + call_id=call_id, + name="docs_search", + arguments={"query": "azure"}, + additional_properties={"server_label": "Microsoft_Learn_MCP"}, + ) + + async def stream_fn( + messages: list[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + del messages, options, kwargs + if state["phase"] == "pause": + yield ChatResponseUpdate( + contents=[Content.from_function_approval_request(id=call_id, function_call=hosted_call)], + role="assistant", + ) + return + raise RuntimeError("provider stream failed") + + agent = Agent( + name="test_agent", + instructions="Test", + client=streaming_chat_client_stub(stream_fn), + ) + app = FastAPI() + add_agent_framework_fastapi_endpoint( + app, + AgentFrameworkAgent(agent=agent, require_confirmation=False), + path="/approval", + ) + client = TestClient(app) + + pause_response = client.post( + "/approval", + json={ + "runId": "run-pause", + "threadId": "thread-hosted-failure", + "messages": [{"role": "user", "content": "Search the hosted docs"}], + }, + ) + assert pause_response.status_code == 200 + state["phase"] = "resume" + + failed_response = client.post( + "/approval", + json={ + "runId": "run-failed", + "threadId": "thread-hosted-failure", + "messages": [], + "resume": [{"interruptId": call_id, "status": "resolved", "payload": {"accepted": True}}], + }, + ) + + assert failed_response.status_code == 200 + assert [event for event in _decode_sse_events(failed_response) if event.get("type") == "RUN_ERROR"] + + retry_response = client.post( + "/approval", + json={ + "runId": "run-retry", + "threadId": "thread-hosted-failure", + "messages": [], + "resume": [{"interruptId": call_id, "status": "resolved", "payload": {"accepted": True}}], + }, + ) + + retry_errors = [event for event in _decode_sse_events(retry_response) if event.get("type") == "RUN_ERROR"] + assert len(retry_errors) == 1 + assert retry_errors[0]["code"] == "APPROVAL_RESUME_INVALID" + assert "indeterminate" in retry_errors[0]["message"] + + async def test_endpoint_does_not_forward_resolved_local_approval_control_to_chat_client( streaming_chat_client_stub, ) -> None: diff --git a/python/packages/ag-ui/tests/ag_ui/test_run.py b/python/packages/ag-ui/tests/ag_ui/test_run.py index 2e22067613..d986725806 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_run.py @@ -35,7 +35,7 @@ _should_suppress_intermediate_snapshot, run_agent_stream, ) -from agent_framework_ag_ui._approval_lifecycle import ApprovalLifecycle +from agent_framework_ag_ui._approval_lifecycle import ApprovalExecutionOwner, ApprovalLifecycle from agent_framework_ag_ui._run_common import ( FlowState, _build_run_finished_event, @@ -923,9 +923,12 @@ def test_emit_approval_request_populates_interrupt_metadata(): assert flow.interrupts[0]["reason"] == "tool_call" assert flow.interrupts[0]["toolCallId"] == "call_123" assert flow.interrupts[0]["message"] == "Approve running write_doc?" - assert flow.interrupts[0]["responseSchema"]["required"] == ["accepted"] - assert flow.interrupts[0]["responseSchema"]["properties"]["accepted"]["type"] == "boolean" - assert flow.interrupts[0]["responseSchema"]["properties"]["content"]["type"] == "string" + response_schema = flow.interrupts[0]["responseSchema"] + assert response_schema["anyOf"] == [{"required": ["approved"]}, {"required": ["accepted"]}] + assert response_schema["properties"]["approved"]["type"] == "boolean" + assert response_schema["properties"]["accepted"]["type"] == "boolean" + assert response_schema["properties"]["content"]["type"] == "string" + assert response_schema["properties"]["editedArgs"]["required"] == ["content"] assert flow.interrupts[0]["metadata"]["agent_framework"]["type"] == "function_approval_request" assert flow.interrupts[0]["metadata"]["agent_framework"]["function_call"] == { "call_id": "call_123", @@ -1039,14 +1042,16 @@ def test_resume_to_tool_messages_skips_cancelled_entries(): def test_canonical_approval_resume_does_not_mutate_arguments_until_batch_validates(): """Edited approval arguments are committed only after every resume entry validates.""" lifecycle = ApprovalLifecycle() - pending_entry = lifecycle.register_local( + pending_entry = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-weather", interrupt_id="call_a", call_id="call_a", name="get_weather", arguments='{"city":"Seattle"}', ) - lifecycle.register_local( + lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="thread-weather", interrupt_id="call_b", call_id="call_b", @@ -1074,7 +1079,8 @@ def test_canonical_approval_resume_does_not_mutate_arguments_until_batch_validat def test_canonical_hosted_approval_resume_rejects_edited_arguments_without_mutating_pending() -> None: """Hosted approvals accept a decision only because providers ignore edited arguments.""" lifecycle = ApprovalLifecycle() - pending_entry = lifecycle.register_hosted( + pending_entry = lifecycle.register( + owner=ApprovalExecutionOwner.HOSTED, thread_id="thread-hosted", interrupt_id="mcpr_docs", call_id="mcpr_docs", @@ -1107,7 +1113,8 @@ def test_canonical_hosted_approval_resume_rejects_edited_arguments_without_mutat def test_approval_lifecycle_scans_exact_thread_keys_with_colons(): """A thread id that prefixes another thread id must not inherit its pending approval contract.""" lifecycle = ApprovalLifecycle() - lifecycle.register_local( + lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, thread_id="tenant:thread", interrupt_id="call_1", call_id="call_1", From e4404e0c64685213adb7233201edbeac775b4be5 Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Mon, 10 Aug 2026 19:24:39 +0900 Subject: [PATCH 10/13] Align workflow approvals with AG-UI resumes --- .../specs/004-python-function-calling-loop.md | 4 +- python/packages/ag-ui/AGENTS.md | 7 +- python/packages/ag-ui/README.md | 3 +- .../agent_framework_ag_ui/_workflow_run.py | 153 ++++++- .../ag-ui/tests/ag_ui/test_endpoint.py | 372 ++++++++++++++++++ 5 files changed, 527 insertions(+), 12 deletions(-) diff --git a/docs/specs/004-python-function-calling-loop.md b/docs/specs/004-python-function-calling-loop.md index a6e7534362..d19a2b63ab 100644 --- a/docs/specs/004-python-function-calling-loop.md +++ b/docs/specs/004-python-function-calling-loop.md @@ -453,8 +453,8 @@ that manually replay messages own the equivalent rule: do not resend an approval | Argument-scoped rule | Exact arguments are required; empty arguments are not tool-wide. | `test_tool_approval_middleware_always_approve_tool_with_arguments_rule`, `test_tool_approval_middleware_empty_arguments_rule_is_not_tool_wide` | | Provider-injected approval tool | A tool added during `before_run` defers to in-run resolution, executes once, and emits one result. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_agent_approval_deferred_provider_tool_executes` | | AG-UI provider boundary | Completed local approval controls from AG-UI request and snapshot replay are absent from raw chat-client input while deferred and hosted approvals keep their respective in-run/provider paths. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_does_not_forward_resolved_local_approval_control_to_chat_client`, `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_agent_approval_deferred_provider_tool_executes`, `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_canonical_resume_preserves_hosted_approval_for_provider`, `packages/ag-ui/tests/ag_ui/test_run.py::test_filter_local_approval_responses_for_provider_removes_duplicate_completed_controls`, `packages/ag-ui/tests/ag_ui/test_run.py::test_filter_local_approval_responses_for_provider_pairs_reused_call_ids_by_occurrence`, `packages/ag-ui/tests/ag_ui/test_run.py::test_canonical_hosted_approval_resume_rejects_edited_arguments_without_mutating_pending` | -| AG-UI standard approval payload | `approved` plus full-replacement `editedArgs` executes once and replays idempotently, while legacy `accepted` plus direct partial edits remains supported. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_agent_approval_resume_entry_applies_standard_full_replacement_edited_args`, `test_endpoint_agent_approval_replayed_standard_edited_resume_is_idempotent`, `test_endpoint_agent_approval_resume_entry_applies_edited_arguments` | -| AG-UI cancellation | A cancelled interrupt executes zero times and completes normally; resolved siblings in the same complete resume still execute once. Workflow `request_info` cancellation follows the same terminal lifecycle. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_agent_approval_cancelled_resume_entry_completes_without_execution`, `test_endpoint_agent_approval_mixed_cancelled_and_resolved_resume_executes_resolved_tool`, `test_endpoint_workflow_request_info_cancelled_resume_completes_normally` | +| AG-UI standard approval payload | Agent and workflow tool approvals emit canonical `tool_call` interrupts. `approved` plus full-replacement `editedArgs` executes once and replays idempotently, while legacy `accepted` plus direct partial edits remains supported. Hosted approvals remain decision-only. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_agent_approval_resume_entry_applies_standard_full_replacement_edited_args`, `test_endpoint_agent_approval_replayed_standard_edited_resume_is_idempotent`, `test_endpoint_agent_approval_resume_entry_applies_edited_arguments`, `test_workflow_endpoint_emits_canonical_tool_approval_interrupt`, `test_workflow_endpoint_accepts_canonical_tool_approval_resume`, `test_workflow_endpoint_applies_canonical_approval_edited_args`, `test_workflow_endpoint_accepts_legacy_partial_approval_edits`, `test_workflow_endpoint_hosted_approval_rejects_argument_edits` | +| AG-UI cancellation | A cancelled interrupt executes zero times and completes normally; resolved siblings in the same complete resume still execute once. Workflow cancellation clears both runner correlation and the owning agent executor's pending request so later approvals remain resumable. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_agent_approval_cancelled_resume_entry_completes_without_execution`, `test_endpoint_agent_approval_mixed_cancelled_and_resolved_resume_executes_resolved_tool`, `test_endpoint_workflow_request_info_cancelled_resume_completes_normally`, `test_workflow_endpoint_cancelled_agent_approval_does_not_block_next_approval` | | AG-UI local executor unavailable on resume | A claimed local occurrence whose executor disappeared releases its unstarted claim, reports temporary unavailability, and remains safely retryable. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_agent_approval_resume_remains_retryable_when_local_tool_is_temporarily_unavailable` | | AG-UI forwarded execution interruption | A provider failure, cancellation, or stream close after forwarding an approval recovers the open occurrence as indeterminate when no idempotency key proves retry safety. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_hosted_approval_becomes_indeterminate_when_provider_stream_fails` | diff --git a/python/packages/ag-ui/AGENTS.md b/python/packages/ag-ui/AGENTS.md index 86a178f322..c1b6c32896 100644 --- a/python/packages/ag-ui/AGENTS.md +++ b/python/packages/ag-ui/AGENTS.md @@ -29,9 +29,10 @@ AG-UI protocol integration for building agent UIs with the AG-UI standard. - Multimodal user inputs support both legacy (`text`, `binary`) and draft-style (`image`, `audio`, `video`, `document`) shapes. - Interrupted runs complete with `RUN_FINISHED.outcome.type == "interrupt"` and canonical `outcome.interrupts`; do not document or add new flows that depend on the legacy top-level `RUN_FINISHED.interrupt` field. - `Interrupt` and `ResumeEntry` come from the `ag-ui-protocol` package (`ag_ui.core`), not from an Agent Framework-specific interrupt model. -- Tool approval interrupts advertise standard `approved` and full-replacement `editedArgs` responses while retaining - the existing `accepted` alias and direct partial edits for MAF client compatibility. A `cancelled` resume completes - normally without executing that call; resolved siblings in the same complete resume still proceed. +- Tool approval interrupts, including approvals surfaced through workflow `request_info`, advertise standard + `approved` and full-replacement `editedArgs` responses while retaining the existing `accepted` alias and direct + partial edits for MAF client compatibility. A `cancelled` resume completes normally without executing that call; + resolved siblings in the same complete resume still proceed. - Approval-time execution preserves each call's complete result group. Follow-up user-input requests remain in the resumed messages, while `TOOL_CALL_RESULT` events are emitted only for terminal `function_result` contents. - Approval responses for tools injected during `before_run` are deferred to the in-run approval middleware rather diff --git a/python/packages/ag-ui/README.md b/python/packages/ag-ui/README.md index b9769233f8..a7688e59c4 100644 --- a/python/packages/ag-ui/README.md +++ b/python/packages/ag-ui/README.md @@ -208,7 +208,8 @@ Resume the paused thread with a canonical `resume` array. Each entry addresses e Tool approvals use the standard `approved` field and may provide `editedArgs` as a full replacement of the tool arguments. For compatibility with existing MAF clients, `accepted` remains an alias for `approved`, and direct argument fields remain supported as partial edits. Cancellation is a normal terminal decision: cancelled calls do -not execute, while resolved siblings in the same complete resume continue normally. +not execute, while resolved siblings in the same complete resume continue normally. The same tool-approval shape and +resume payloads apply when an agent approval is surfaced through a workflow `request_info` event. ```json { diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py index 8b99d67c60..4128420862 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py @@ -33,6 +33,8 @@ from ._message_adapters import normalize_agui_input_messages from ._run_common import ( FlowState, + _approval_interrupt_for_function_call, # pyright: ignore[reportPrivateUsage] + _approval_response_schema, # pyright: ignore[reportPrivateUsage] _build_run_finished_event, _close_reasoning_block, _emit_content, @@ -158,6 +160,26 @@ def _interrupt_entry_for_request_event(request_event: Any) -> dict[str, Any] | N return None value = _workflow_interrupt_value(request_payload.get("data")) + request_data = getattr(request_event, "data", None) + if ( + isinstance(request_data, Content) + and request_data.type == "function_approval_request" + and request_data.function_call is not None + ): + workflow_metadata = _workflow_interrupt_metadata(request_payload, value)["agent_framework"] + workflow_metadata.pop("type", None) + response_schema = ( + _approval_response_schema() + if request_data.function_call.additional_properties.get("server_label") + else None + ) + return _approval_interrupt_for_function_call( + interrupt_id=str(request_payload["request_id"]), + function_call=request_data.function_call, + metadata=workflow_metadata, + response_schema=response_schema, + ) + entry: dict[str, Any] = { "id": str(request_payload["request_id"]), "reason": "input_required", @@ -283,7 +305,7 @@ def _resume_error_for_pending_workflow_requests( def _consume_cancelled_workflow_requests(workflow: Workflow, resume_entries: list[dict[str, Any]]) -> None: - """Remove cancelled workflow request_info events from the runner context.""" + """Remove cancelled workflow requests from runner and owning agent-executor state.""" cancelled_ids = {str(entry["interrupt_id"]) for entry in resume_entries if entry.get("status") == "cancelled"} if not cancelled_ids: return @@ -292,9 +314,15 @@ def _consume_cancelled_workflow_requests(workflow: Workflow, resume_entries: lis pending_events = getattr(runner_context, "_pending_request_info_events", None) if not isinstance(pending_events, dict): return + pending_events = cast(dict[str, Any], pending_events) for interrupt_id in cancelled_ids: - pending_events.pop(interrupt_id, None) + request_event = pending_events.pop(interrupt_id, None) + source_executor_id = getattr(request_event, "source_executor_id", None) + executor = workflow.executors.get(source_executor_id) if source_executor_id else None + pending_agent_requests = getattr(executor, "_pending_agent_requests", None) + if isinstance(pending_agent_requests, dict): + cast(dict[str, Any], pending_agent_requests).pop(interrupt_id, None) def _coerce_json_value(value: Any) -> Any: @@ -401,6 +429,79 @@ def _coerce_message(value: Any) -> Message | None: ) +def _approval_argument_value_matches(original_value: Any, edited_value: Any) -> bool: + """Return whether an edited approval argument preserves its JSON value type.""" + if isinstance(original_value, bool): + return isinstance(edited_value, bool) + if isinstance(original_value, int) and not isinstance(original_value, bool): + return isinstance(edited_value, int) and not isinstance(edited_value, bool) + if isinstance(original_value, float): + return isinstance(edited_value, (int, float)) and not isinstance(edited_value, bool) + if isinstance(original_value, str): + return isinstance(edited_value, str) + if isinstance(original_value, list): + return isinstance(edited_value, list) + if isinstance(original_value, dict): + return isinstance(edited_value, dict) + return True + + +def _coerce_compact_approval_response(request_data: Content, candidate: dict[str, Any]) -> Content | None: + """Reconstruct a workflow approval response from client-owned decision fields.""" + if {"type", "id", "function_call"}.intersection(candidate): + return None + + approved = candidate.get("approved", candidate.get("accepted")) + if not isinstance(approved, bool): + return None + + direct_edited_arguments = { + key: value for key, value in candidate.items() if key not in {"approved", "accepted", "editedArgs"} + } + standard_edited_arguments = candidate.get("editedArgs") + if request_data.function_call is None: + return None + if ( + direct_edited_arguments or standard_edited_arguments is not None + ) and request_data.function_call.additional_properties.get("server_label"): + return None + + original_arguments = request_data.function_call.parse_arguments() or {} + if standard_edited_arguments is not None: + if not isinstance(standard_edited_arguments, dict) or direct_edited_arguments: + return None + edited_arguments = cast(dict[str, Any], standard_edited_arguments) + if set(edited_arguments) != set(original_arguments): + return None + final_arguments = dict(edited_arguments) + else: + edited_arguments = direct_edited_arguments + if not set(edited_arguments).issubset(original_arguments): + return None + final_arguments = {**original_arguments, **edited_arguments} + + if any( + not _approval_argument_value_matches(original_arguments[name], edited_arguments[name]) + for name in edited_arguments + ): + return None + if not edited_arguments: + return request_data.to_function_approval_response(approved) + + edited_function_call = Content.from_function_call( + call_id=request_data.function_call.call_id or "", + name=request_data.function_call.name or "", + arguments=final_arguments, + informational_only=request_data.function_call.informational_only, + annotations=request_data.function_call.annotations, + additional_properties=request_data.function_call.additional_properties, + raw_representation=request_data.function_call.raw_representation, + ) + response = request_data.to_function_approval_response(approved) + response.function_call = edited_function_call + return response + + def _coerce_response_for_request(request_event: Any, value: Any) -> Any | None: """Coerce a candidate value into the request's expected response type.""" response_type = getattr(request_event, "response_type", None) @@ -447,6 +548,15 @@ def _coerce_response_for_request(request_event: Any, value: Any) -> Any | None: if target_type is Message: return _coerce_message(candidate) if target_type is Content: + request_data = getattr(request_event, "data", None) + if ( + isinstance(request_data, Content) + and request_data.type == "function_approval_request" + and isinstance(candidate, dict) + ): + compact_response = _coerce_compact_approval_response(request_data, cast(dict[str, Any], candidate)) + if compact_response is not None: + return compact_response return _coerce_content(candidate) if target_type is bool: return candidate if isinstance(candidate, bool) else None @@ -461,7 +571,21 @@ def _coerce_response_for_request(request_event: Any, value: Any) -> Any | None: return candidate -def _approval_response_matches_request(request_id: str, request_event: Any, response: Any) -> bool: +def _is_compact_approval_response_payload(value: Any) -> bool: + """Return whether a value contains only client-owned approval decision fields.""" + candidate = _coerce_json_value(value) + return isinstance(candidate, dict) and not {"type", "id", "function_call"}.intersection( + cast(dict[str, Any], candidate) + ) + + +def _approval_response_matches_request( + request_id: str, + request_event: Any, + response: Any, + *, + allow_edited_arguments: bool = False, +) -> bool: """Check whether an approval response matches the pending approval request.""" request_data = getattr(request_event, "data", None) if not isinstance(request_data, Content) or request_data.type != "function_approval_request": @@ -481,6 +605,8 @@ def _approval_response_matches_request(request_id: str, request_event: Any, resp if getattr(response_call, "name", None) != getattr(request_call, "name", None): return False + if allow_edited_arguments: + return True return canonical_function_arguments(response_call) == canonical_function_arguments(request_call) @@ -503,7 +629,12 @@ def _single_pending_response_from_value(pending_events: dict[str, Any], value: A ) return {} - if not _approval_response_matches_request(str(request_id), request_event, coerced_value): + if not _approval_response_matches_request( + str(request_id), + request_event, + coerced_value, + allow_edited_arguments=_is_compact_approval_response_payload(value), + ): logger.info( "Ignoring pending request response for request_id=%s: approval response does not match pending request", request_id, @@ -544,7 +675,12 @@ def _coerce_responses_for_pending_requests( _response_type_name(request_event), ) continue - if not _approval_response_matches_request(request_key, request_event, coerced_value): + if not _approval_response_matches_request( + request_key, + request_event, + coerced_value, + allow_edited_arguments=_is_compact_approval_response_payload(value), + ): logger.info( "Ignoring resume response for request_id=%s: approval response does not match pending request", request_key, @@ -589,7 +725,12 @@ def _coerce_responses_for_pending_requests_strict( code="WORKFLOW_RESUME_INVALID_RESPONSE", ), ) - if not _approval_response_matches_request(request_key, request_event, coerced_value): + if not _approval_response_matches_request( + request_key, + request_event, + coerced_value, + allow_edited_arguments=_is_compact_approval_response_payload(value), + ): return ( {}, RunErrorEvent( diff --git a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py index b1b78c32f3..8ed55d80f5 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -153,6 +153,315 @@ async def start(message: Any, ctx: WorkflowContext[Any, Any]) -> None: assert "RUN_FINISHED" in event_types +async def test_workflow_endpoint_emits_canonical_tool_approval_interrupt() -> None: + """Workflow agent approvals use the standard AG-UI tool-approval interrupt contract.""" + + @executor(id="approval") + async def approval(message: Any, ctx: WorkflowContext[Any, Any]) -> None: + del message + function_call = Content.from_function_call( + call_id="refund-call", + name="submit_refund", + arguments={"order_id": "12345", "amount": 89.99}, + ) + await ctx.request_info( + Content.from_function_approval_request(id="approval-1", function_call=function_call), + Content, + request_id="approval-1", + ) + + app = FastAPI() + workflow = WorkflowBuilder(start_executor=approval).build() + add_agent_framework_fastapi_endpoint(app, workflow, path="/workflow-approval") + + response = TestClient(app).post( + "/workflow-approval", + json={"messages": [{"role": "user", "content": "Refund the order"}]}, + ) + + assert response.status_code == 200 + finished = [event for event in _decode_sse_events(response) if event.get("type") == "RUN_FINISHED"] + interrupt = _run_finished_interrupts(finished[-1])[0] + assert interrupt["id"] == "approval-1" + assert interrupt["reason"] == "tool_call" + assert interrupt["toolCallId"] == "refund-call" + assert interrupt["responseSchema"] == { + "type": "object", + "properties": { + "approved": {"type": "boolean", "description": "Whether the requested tool call is approved."}, + "accepted": {"type": "boolean", "description": "Legacy alias for approved."}, + "order_id": {"type": "string", "description": "Optional edited value for the 'order_id' tool argument."}, + "amount": {"type": "number", "description": "Optional edited value for the 'amount' tool argument."}, + "editedArgs": { + "type": "object", + "description": "Full replacement of the tool arguments. Not merged.", + "properties": {"order_id": {"type": "string"}, "amount": {"type": "number"}}, + "required": ["order_id", "amount"], + "additionalProperties": False, + }, + }, + "anyOf": [{"required": ["approved"]}, {"required": ["accepted"]}], + "additionalProperties": False, + } + assert interrupt["metadata"]["agent_framework"]["type"] == "function_approval_request" + + +async def test_workflow_endpoint_accepts_canonical_tool_approval_resume() -> None: + """Workflow approvals reconstruct server-owned response identity from a canonical resume decision.""" + + class ApprovalExecutor(Executor): + def __init__(self) -> None: + super().__init__(id="approval") + + @handler + async def start(self, message: Any, ctx: WorkflowContext[Any, Any]) -> None: + del message + function_call = Content.from_function_call( + call_id="refund-call", + name="submit_refund", + arguments={"order_id": "12345", "amount": 89.99}, + ) + await ctx.request_info( + Content.from_function_approval_request(id="approval-1", function_call=function_call), + Content, + request_id="approval-1", + ) + + @response_handler + async def approve( + self, + original_request: Content, + response: Content, + ctx: WorkflowContext[Any, Any], + ) -> None: + del original_request + status = "approved" if response.approved else "rejected" + await ctx.yield_output(f"Refund {status}.") # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] + + app = FastAPI() + workflow = WorkflowBuilder(start_executor=ApprovalExecutor()).build() + add_agent_framework_fastapi_endpoint(app, workflow, path="/workflow-approval") + client = TestClient(app) + + pause_response = client.post( + "/workflow-approval", + json={"messages": [{"role": "user", "content": "Refund the order"}]}, + ) + assert pause_response.status_code == 200 + + resume_response = client.post( + "/workflow-approval", + json={ + "messages": [], + "resume": [{"interruptId": "approval-1", "status": "resolved", "payload": {"approved": True}}], + }, + ) + + assert resume_response.status_code == 200 + events = _decode_sse_events(resume_response) + assert not [event for event in events if event.get("type") == "RUN_ERROR"] + assert "Refund approved." == "".join( + str(event.get("delta", "")) for event in events if event.get("type") == "TEXT_MESSAGE_CONTENT" + ) + + +async def test_workflow_endpoint_applies_canonical_approval_edited_args() -> None: + """Workflow approvals apply standard editedArgs as a full replacement.""" + + class ApprovalExecutor(Executor): + def __init__(self) -> None: + super().__init__(id="approval") + + @handler + async def start(self, message: Any, ctx: WorkflowContext[Any, Any]) -> None: + del message + function_call = Content.from_function_call( + call_id="refund-call", + name="submit_refund", + arguments={"order_id": "12345", "amount": 89.99}, + ) + await ctx.request_info( + Content.from_function_approval_request(id="approval-1", function_call=function_call), + Content, + request_id="approval-1", + ) + + @response_handler + async def approve( + self, + original_request: Content, + response: Content, + ctx: WorkflowContext[Any, Any], + ) -> None: + del original_request + arguments = response.function_call.parse_arguments() if response.function_call is not None else None + await ctx.yield_output(json.dumps(arguments, sort_keys=True)) # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] + + app = FastAPI() + workflow = WorkflowBuilder(start_executor=ApprovalExecutor()).build() + add_agent_framework_fastapi_endpoint(app, workflow, path="/workflow-approval") + client = TestClient(app) + pause_response = client.post( + "/workflow-approval", + json={"messages": [{"role": "user", "content": "Refund the order"}]}, + ) + assert pause_response.status_code == 200 + + resume_response = client.post( + "/workflow-approval", + json={ + "messages": [], + "resume": [ + { + "interruptId": "approval-1", + "status": "resolved", + "payload": { + "approved": True, + "editedArgs": {"order_id": "54321", "amount": 49.5}, + }, + } + ], + }, + ) + + assert resume_response.status_code == 200 + events = _decode_sse_events(resume_response) + assert not [event for event in events if event.get("type") == "RUN_ERROR"] + assert '{"amount": 49.5, "order_id": "54321"}' == "".join( + str(event.get("delta", "")) for event in events if event.get("type") == "TEXT_MESSAGE_CONTENT" + ) + + +async def test_workflow_endpoint_accepts_legacy_partial_approval_edits() -> None: + """Workflow approvals retain the MAF accepted alias and direct partial argument edits.""" + + class ApprovalExecutor(Executor): + def __init__(self) -> None: + super().__init__(id="approval") + + @handler + async def start(self, message: Any, ctx: WorkflowContext[Any, Any]) -> None: + del message + function_call = Content.from_function_call( + call_id="refund-call", + name="submit_refund", + arguments={"order_id": "12345", "amount": 89.99}, + ) + await ctx.request_info( + Content.from_function_approval_request(id="approval-1", function_call=function_call), + Content, + request_id="approval-1", + ) + + @response_handler + async def approve( + self, + original_request: Content, + response: Content, + ctx: WorkflowContext[Any, Any], + ) -> None: + del original_request + arguments = response.function_call.parse_arguments() if response.function_call is not None else None + await ctx.yield_output(json.dumps(arguments, sort_keys=True)) # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] + + app = FastAPI() + workflow = WorkflowBuilder(start_executor=ApprovalExecutor()).build() + add_agent_framework_fastapi_endpoint(app, workflow, path="/workflow-approval") + client = TestClient(app) + pause_response = client.post( + "/workflow-approval", + json={"messages": [{"role": "user", "content": "Refund the order"}]}, + ) + assert pause_response.status_code == 200 + + resume_response = client.post( + "/workflow-approval", + json={ + "messages": [], + "resume": [ + { + "interruptId": "approval-1", + "status": "resolved", + "payload": {"accepted": True, "amount": 49.5}, + } + ], + }, + ) + + assert resume_response.status_code == 200 + events = _decode_sse_events(resume_response) + assert not [event for event in events if event.get("type") == "RUN_ERROR"] + assert '{"amount": 49.5, "order_id": "12345"}' == "".join( + str(event.get("delta", "")) for event in events if event.get("type") == "TEXT_MESSAGE_CONTENT" + ) + + +async def test_workflow_endpoint_hosted_approval_rejects_argument_edits() -> None: + """Workflow-hosted approvals remain decision-only because the remote owner controls arguments.""" + + class ApprovalExecutor(Executor): + def __init__(self) -> None: + super().__init__(id="approval") + + @handler + async def start(self, message: Any, ctx: WorkflowContext[Any, Any]) -> None: + del message + function_call = Content.from_function_call( + call_id="hosted-call", + name="hosted_refund", + arguments={"order_id": "12345"}, + additional_properties={"server_label": "refund-server"}, + ) + await ctx.request_info( + Content.from_function_approval_request(id="approval-1", function_call=function_call), + Content, + request_id="approval-1", + ) + + @response_handler + async def approve( + self, + original_request: Content, + response: Content, + ctx: WorkflowContext[Any, Any], + ) -> None: + del original_request, response + await ctx.yield_output("Hosted approval handled.") # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] + + app = FastAPI() + workflow = WorkflowBuilder(start_executor=ApprovalExecutor()).build() + add_agent_framework_fastapi_endpoint(app, workflow, path="/workflow-approval") + client = TestClient(app) + pause_response = client.post( + "/workflow-approval", + json={"messages": [{"role": "user", "content": "Refund the order"}]}, + ) + + assert pause_response.status_code == 200 + finished = [event for event in _decode_sse_events(pause_response) if event.get("type") == "RUN_FINISHED"] + interrupt = _run_finished_interrupts(finished[-1])[0] + assert set(interrupt["responseSchema"]["properties"]) == {"approved", "accepted"} + + resume_response = client.post( + "/workflow-approval", + json={ + "messages": [], + "resume": [ + { + "interruptId": "approval-1", + "status": "resolved", + "payload": {"approved": True, "editedArgs": {"order_id": "54321"}}, + } + ], + }, + ) + + assert resume_response.status_code == 200 + errors = [event for event in _decode_sse_events(resume_response) if event.get("type") == "RUN_ERROR"] + assert len(errors) == 1 + assert errors[0]["code"] == "WORKFLOW_RESUME_INVALID_RESPONSE" + + async def test_add_endpoint_accepts_keepalive_option_for_supported_runners(build_chat_client): """Keepalive configuration is accepted at the endpoint seam for every supported runner shape.""" @@ -5335,6 +5644,69 @@ async def requester(message: Any, ctx: WorkflowContext[Any, Any]) -> None: assert call_count == 1 +async def test_workflow_endpoint_cancelled_agent_approval_does_not_block_next_approval() -> None: + """Cancelling one workflow-agent approval leaves a later approval resumable.""" + app = FastAPI() + + def approval_update(approval_id: str, call_id: str) -> AgentResponseUpdate: + function_call = Content.from_function_call( + call_id=call_id, + name="submit_refund", + arguments={"order_id": approval_id}, + ) + approval_request = Content.from_function_approval_request(id=approval_id, function_call=function_call) + return AgentResponseUpdate(contents=[approval_request], role="assistant") + + agent = StubAgent(updates=[approval_update("approval-1", "refund-call-1")]) + workflow = WorkflowBuilder(start_executor=agent).build() + add_agent_framework_fastapi_endpoint(app, workflow, path="/workflow-agent-approval") + client = TestClient(app) + + first_pause = client.post( + "/workflow-agent-approval", + json={"messages": [{"role": "user", "content": "First refund"}]}, + ) + assert first_pause.status_code == 200 + first_finished = [event for event in _decode_sse_events(first_pause) if event.get("type") == "RUN_FINISHED"] + assert _run_finished_interrupts(first_finished[-1])[0]["id"] == "approval-1" + + cancelled = client.post( + "/workflow-agent-approval", + json={ + "messages": [], + "resume": [{"interruptId": "approval-1", "status": "cancelled"}], + }, + ) + assert cancelled.status_code == 200 + + agent.updates = [approval_update("approval-2", "refund-call-2")] + second_pause = client.post( + "/workflow-agent-approval", + json={"messages": [{"role": "user", "content": "Second refund"}]}, + ) + assert second_pause.status_code == 200 + second_finished = [event for event in _decode_sse_events(second_pause) if event.get("type") == "RUN_FINISHED"] + assert _run_finished_interrupts(second_finished[-1])[0]["id"] == "approval-2" + + agent.updates = [ + AgentResponseUpdate(contents=[Content.from_text(text="Second refund completed.")], role="assistant") + ] + resumed = client.post( + "/workflow-agent-approval", + json={ + "messages": [], + "resume": [{"interruptId": "approval-2", "status": "resolved", "payload": {"approved": True}}], + }, + ) + + assert resumed.status_code == 200 + events = _decode_sse_events(resumed) + assert not [event for event in events if event.get("type") == "RUN_ERROR"] + assert "Second refund completed." == "".join( + str(event.get("delta", "")) for event in events if event.get("type") == "TEXT_MESSAGE_CONTENT" + ) + + class _FailingSaveStore(InMemoryAGUIThreadSnapshotStore): """Store whose save always fails, simulating a transient backend outage.""" From d6969a262ce8bd07dc79bc24eb0b284cfd4f2bff Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Tue, 11 Aug 2026 17:28:00 +0900 Subject: [PATCH 11/13] Address AG-UI approval review findings --- .../specs/004-python-function-calling-loop.md | 10 +- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 379 ++++++++++-------- .../_approval_lifecycle.py | 138 +++++-- .../agent_framework_ag_ui/_approval_state.py | 20 +- .../agent_framework_ag_ui/_run_common.py | 4 +- .../agent_framework_ag_ui/_workflow_run.py | 15 +- .../tests/ag_ui/test_approval_lifecycle.py | 59 +++ .../tests/ag_ui/test_approval_result_event.py | 100 +++++ .../ag-ui/tests/ag_ui/test_approval_state.py | 30 ++ .../ag-ui/tests/ag_ui/test_endpoint.py | 116 +++++- python/packages/ag-ui/tests/ag_ui/test_run.py | 197 ++++++++- .../ag-ui/tests/ag_ui/test_workflow_run.py | 59 +++ 12 files changed, 892 insertions(+), 235 deletions(-) diff --git a/docs/specs/004-python-function-calling-loop.md b/docs/specs/004-python-function-calling-loop.md index d19a2b63ab..e682fd6213 100644 --- a/docs/specs/004-python-function-calling-loop.md +++ b/docs/specs/004-python-function-calling-loop.md @@ -491,11 +491,11 @@ that manually replay messages own the equivalent rule: do not resend an approval | Foundry encrypted reasoning opt-in | Foundry clients omit `reasoning.encrypted_content` by default and preserve an explicit caller opt-in. | `packages/foundry/tests/foundry/test_foundry_chat_client.py::test_get_response_does_not_request_encrypted_reasoning_by_default`, `test_get_response_preserves_explicit_encrypted_reasoning_opt_in`, `packages/foundry/tests/foundry/test_foundry_agent.py::test_foundry_agent_basic_call_does_not_request_unsupported_encrypted_reasoning`, `test_foundry_agent_preserves_caller_requested_encrypted_reasoning`, `packages/foundry_hosting/tests/test_responses_int.py::TestReasoningHostedMcpReplay::test_second_turn_replays_mcp_call_with_encrypted_reasoning` | | Opaque reasoning signature replay | Provider-specific opaque reasoning metadata is captured and restored on reconstructed calls. | `packages/gemini/tests/test_gemini_client.py::test_function_call_part_captures_thought_signature_as_reasoning_content`, `test_reconstructed_function_call_replays_thought_signature_from_reasoning_content` | | Chat Completions approval wrappers | Framework approval wrappers are not sent as chat messages. | `packages/openai/tests/openai/test_openai_chat_completion_client.py` approval serialization tests | -| AG-UI approval result event | Approved result emits once with content and persists in snapshot. | `packages/ag-ui/tests/ag_ui/test_approval_result_event.py::test_approval_resume_emits_tool_call_result`, `test_approval_resume_result_has_content`, `test_approval_resume_snapshot_replaces_approval_payload_with_tool_result`, `test_approval_resume_zero_updates_emits_tool_result` | -| AG-UI rejection/mixed decision | Transport emits only the events defined for approved and rejected calls without duplicates. | `test_rejection_does_not_emit_tool_call_result`, `test_mixed_approve_reject_emits_only_approved_tool_result`, `test_resolve_approval_responses_returns_only_approved` | -| AG-UI approval-time follow-up | The full grouped user-input pause remains in message history and emits no synthetic `TOOL_CALL_RESULT`. | `test_resolve_approval_responses_preserves_follow_up_user_input_group` | -| AG-UI approval execution failure | A grouped executor failure becomes one deterministic terminal error result for the approved call. | `test_resolve_approval_responses_returns_failure_when_grouped_execution_raises` | -| AG-UI no-approval path | Ordinary tool results do not gain an extra approval result event. | `test_no_approval_no_extra_tool_result` | +| AG-UI approval result event | Approved result emits once with content and persists in snapshot. | `packages/ag-ui/tests/ag_ui/test_approval_result_event.py::test_approved_call_emits_one_live_result_under_original_identity`, `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_agent_approval_resume_persists_replayable_tool_results`, `test_endpoint_agent_approval_replayed_resume_entry_reprojects_retained_result` | +| AG-UI rejection/mixed decision | Transport emits only the events defined for approved and rejected calls without duplicates. | `packages/ag-ui/tests/ag_ui/test_approval_result_event.py::test_rejected_call_does_not_execute_or_emit_live_result`, `test_mixed_batch_preserves_approved_result_identity_and_order`, `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_agent_approval_rejection_releases_already_approved_sibling` | +| AG-UI approval-time follow-up | The full grouped user-input pause remains in message history and emits no synthetic `TOOL_CALL_RESULT`. | `packages/ag-ui/tests/ag_ui/test_approval_result_event.py::test_approval_follow_up_group_remains_in_history_without_live_tool_result` | +| AG-UI approval execution failure | A grouped executor failure becomes one deterministic terminal error result for the approved call. | `packages/ag-ui/tests/ag_ui/test_approval_result_event.py::test_approval_execution_failure_emits_one_terminal_error_result` | +| AG-UI no-approval path | Ordinary tool results do not gain an extra approval result event. | `packages/ag-ui/tests/ag_ui/test_approval_result_event.py::test_no_approval_path_emits_no_approval_specific_duplicate_result` | | AG-UI `confirm_changes` snapshot | An accepted synthetic confirmation is replaced only when its original function call has a real result; rejection is cleaned explicitly, and missing accepted results remain inert. | `packages/ag-ui/tests/ag_ui/test_confirm_changes_snapshot.py` | | AG-UI malformed `confirm_changes` metadata | Non-list tool-call metadata and malformed argument JSON are ignored without guessing a target call. | `test_confirm_changes_target_ignores_non_list_tool_calls`, `test_confirm_changes_target_rejects_malformed_arguments_json` | | Compaction pair integrity | Adjacent and non-adjacent pairs, including assistant-embedded results and completed reused-id occurrences, remain atomic without pairing ambiguous or out-of-order ids. | `packages/core/tests/core/test_compaction.py::test_group_annotations_keep_tool_call_and_tool_result_atomic`, `test_group_annotations_include_reasoning_in_tool_call_group`, `test_group_annotations_pair_nonadjacent_function_result_by_call_id`, `test_group_annotations_pair_multiple_nonadjacent_results_with_declaration`, `test_group_annotations_pair_completed_reused_call_id_occurrences`, `test_group_annotations_close_assistant_embedded_result_before_reused_call_id`, `test_sliding_window_does_not_retain_orphan_result_after_assistant_embedded_result`, `test_sliding_window_keeps_reused_call_id_occurrences_atomic`, `test_group_annotations_do_not_pair_ambiguous_duplicate_call_ids` | diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index c8176b9d58..9114dccbcd 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py @@ -58,7 +58,9 @@ ApprovalExecutionOwner, ApprovalLifecycle, ApprovalOccurrence, + ApprovalOccurrenceIdentity, ApprovalSnapshotReconciliation, + ApprovalStatus, AuthorizedExecution, ClaimRecoveryPolicy, DeferredPendingToolTransitionOwner, @@ -876,7 +878,20 @@ def _register_server_generated_approval_response( aliases=[str(response.function_call.call_id)] if response.function_call.call_id else None, server_label=_function_call_server_label(response.function_call), ) - if not response.approved or execution_owner is ApprovalExecutionOwner.UNAVAILABLE: + if not response.approved: + lifecycle.claim_batch( + thread_id=thread_id, + decisions=[ + ResumeDecision( + interrupt_id=str(response_id), + accepted=False, + arguments=arguments, + name=response.function_call.name, + ) + ], + ) + return None + if execution_owner is ApprovalExecutionOwner.UNAVAILABLE: return None return lifecycle.claim( thread_id=thread_id, @@ -895,7 +910,7 @@ def _pop_collected_tool_approval_response_messages( tools: list[Any] | None, *, lifecycle: ApprovalLifecycle, - authorized_executions: dict[str, AuthorizedExecution] | None = None, + authorized_executions: dict[ApprovalOccurrenceIdentity, AuthorizedExecution] | None = None, ) -> list[Message]: """Pop server-collected auto-approved responses into provider-visible messages.""" raw_state = session.state.get(_TOOL_APPROVAL_STATE_KEY) @@ -920,7 +935,7 @@ def _pop_collected_tool_approval_response_messages( has_deferred_owner=True, ) if intent is not None and authorized_executions is not None: - authorized_executions[intent.identity.call_id] = intent + authorized_executions[intent.identity] = intent responses.append(response) state[_COLLECTED_APPROVAL_RESPONSES_KEY] = [] @@ -1056,6 +1071,114 @@ def _json_schema_value_matches(original_value: Any, edited_value: Any) -> bool: return True +def _canonical_approval_decision( + payload_value: Any, + *, + interrupt_id: str, + original_arguments_text: str, + server_label: str | None, +) -> tuple[bool | None, str | None, dict[str, Any] | None, RunErrorEvent | None]: + """Validate one approval payload and return its canonical full argument replacement.""" + payload = _parse_json_object(payload_value) + if payload is None: + return ( + None, + None, + None, + RunErrorEvent( + message=f"Approval resume for interruptId '{interrupt_id}' must include an object payload.", + code="APPROVAL_RESUME_INVALID", + ), + ) + accepted = payload.get("approved", payload.get("accepted")) + if not isinstance(accepted, bool): + return ( + None, + None, + None, + RunErrorEvent( + message=f"Approval resume for interruptId '{interrupt_id}' must include a boolean accepted value.", + code="APPROVAL_RESUME_INVALID", + ), + ) + + original_arguments = _parse_json_object(original_arguments_text) or {} + direct_edited_arguments = { + key: value for key, value in payload.items() if key not in {"accepted", "approved", "editedArgs"} + } + standard_edited_arguments = payload.get("editedArgs") + if standard_edited_arguments is not None: + if not isinstance(standard_edited_arguments, dict) or direct_edited_arguments: + return ( + None, + None, + None, + RunErrorEvent( + message=( + f"Approval resume for interruptId '{interrupt_id}' must provide editedArgs as the " + "only edited-argument representation." + ), + code="APPROVAL_RESUME_INVALID_RESPONSE", + ), + ) + edited_arguments = cast(dict[str, Any], standard_edited_arguments) + if set(edited_arguments) != set(original_arguments): + return ( + None, + None, + None, + RunErrorEvent( + message=( + f"Approval resume for interruptId '{interrupt_id}' must provide editedArgs as a full " + "replacement of the pending tool arguments." + ), + code="APPROVAL_RESUME_INVALID_RESPONSE", + ), + ) + else: + edited_arguments = direct_edited_arguments + if edited_arguments and server_label: + return ( + None, + None, + None, + RunErrorEvent( + message=f"Hosted approval resume for interruptId '{interrupt_id}' does not support edited arguments.", + code="APPROVAL_RESUME_INVALID_RESPONSE", + ), + ) + if not set(edited_arguments).issubset(set(original_arguments)): + return ( + None, + None, + None, + RunErrorEvent( + message=f"Approval resume for interruptId '{interrupt_id}' includes unsupported edited arguments.", + code="APPROVAL_RESUME_INVALID", + ), + ) + for name, edited_value in edited_arguments.items(): + if not _json_schema_value_matches(original_arguments[name], edited_value): + return ( + None, + None, + None, + RunErrorEvent( + message=( + f"Approval resume for interruptId '{interrupt_id}' has invalid type for edited argument " + f"'{name}'." + ), + code="APPROVAL_RESUME_INVALID_RESPONSE", + ), + ) + + merged_arguments = ( + dict(edited_arguments) if standard_edited_arguments is not None else {**original_arguments, **edited_arguments} + ) + canonical_arguments = json.dumps(make_json_safe(merged_arguments), sort_keys=True, separators=(",", ":")) + return accepted, canonical_arguments, merged_arguments, None + + def _canonical_approval_resume_messages( resume_payload: Any, thread_id: str, @@ -1064,7 +1187,7 @@ def _canonical_approval_resume_messages( lifecycle: ApprovalLifecycle, tools: list[Any] | None = None, has_deferred_owner: bool = False, - authorized_executions: dict[str, AuthorizedExecution] | None = None, + authorized_executions: dict[ApprovalOccurrenceIdentity, AuthorizedExecution] | None = None, retained_results: list[Content] | None = None, snapshot_reconciliations: list[ApprovalSnapshotReconciliation] | None = None, ) -> tuple[list[dict[str, Any]], set[str], set[str], RunErrorEvent | None]: @@ -1084,12 +1207,6 @@ def _canonical_approval_resume_messages( for interrupt in normalized_interrupts: if interrupt.get("status") != "resolved": break - payload = _parse_json_object(interrupt.get("value")) - if payload is None: - break - accepted = payload.get("approved", payload.get("accepted")) - if not isinstance(accepted, bool): - break interrupt_id = str(interrupt["id"]) try: name, retained_arguments = lifecycle.decision_context( @@ -1098,34 +1215,31 @@ def _canonical_approval_resume_messages( ) except KeyError: break - edited_arguments = { - key: value - for key, value in payload.items() - if key not in {"accepted", "approved", "editedArgs"} - } - standard_edited_arguments = payload.get("editedArgs") - canonical_arguments: str | None = None - if isinstance(standard_edited_arguments, dict) and not edited_arguments: - canonical_arguments = json.dumps( - make_json_safe(standard_edited_arguments), - sort_keys=True, - separators=(",", ":"), - ) - elif edited_arguments: - retained_argument_values = _parse_json_object(retained_arguments) - if retained_argument_values is None: - break - canonical_arguments = json.dumps( - make_json_safe({**retained_argument_values, **edited_arguments}), - sort_keys=True, - separators=(",", ":"), - ) + occurrence = lifecycle.occurrence_for_alias(thread_id=thread_id, interrupt_id=interrupt_id) + original_arguments_text = ( + occurrence.decision.original_arguments + if occurrence is not None + and occurrence.decision is not None + and occurrence.decision.original_arguments is not None + else retained_arguments + ) + accepted, canonical_arguments, _, validation_error = _canonical_approval_decision( + interrupt.get("value"), + interrupt_id=interrupt_id, + original_arguments_text=original_arguments_text, + server_label=occurrence.server_label if occurrence is not None else None, + ) + if validation_error is not None: + return [], handled_ids, cancelled_ids, validation_error + if accepted is None or canonical_arguments is None: + break decisions.append( ResumeDecision( interrupt_id=interrupt_id, accepted=accepted, arguments=canonical_arguments, name=name, + original_arguments=original_arguments_text, ) ) if len(decisions) == len(normalized_interrupts) and decisions: @@ -1142,7 +1256,8 @@ def _canonical_approval_resume_messages( ) else: for outcome in batch.retained_outcomes: - retained_results.extend(result.content for result in outcome.replayable_results) + if outcome.snapshot_reconciliation.status is ApprovalStatus.SETTLED: + retained_results.extend(result.content for result in outcome.replayable_results) if snapshot_reconciliations is not None: snapshot_reconciliations.extend(batch.snapshot_reconciliations) handled_ids.update(decision.interrupt_id for decision in decisions) @@ -1209,25 +1324,6 @@ def _canonical_approval_resume_messages( ), ) - if cancelled_ids: - lifecycle_cancelled_ids = [ - interrupt_id for interrupt_id in cancelled_ids if entries_by_interrupt_id.get(interrupt_id) is not None - ] - try: - reconciliations = lifecycle.cancel_batch( - thread_id=thread_id, - interrupt_ids=lifecycle_cancelled_ids, - ) - if snapshot_reconciliations is not None: - snapshot_reconciliations.extend(reconciliations) - except (KeyError, ValueError) as exc: - return ( - [], - handled_ids, - cancelled_ids, - RunErrorEvent(message=str(exc), code="APPROVAL_RESUME_INVALID"), - ) - lifecycle_decisions: list[ResumeDecision] = [] restored_sibling_response_ids: set[str] = set() for entry in entries: @@ -1236,109 +1332,17 @@ def _canonical_approval_resume_messages( if pending_entry is None or entry["status"] == "cancelled": continue - payload = _parse_json_object(entry.get("payload")) - if payload is None: - return ( - [], - handled_ids, - cancelled_ids, - RunErrorEvent( - message=f"Approval resume for interruptId '{interrupt_id}' must include an object payload.", - code="APPROVAL_RESUME_INVALID", - ), - ) - accepted = payload.get("approved", payload.get("accepted")) - if not isinstance(accepted, bool): - return ( - [], - handled_ids, - cancelled_ids, - RunErrorEvent( - message=f"Approval resume for interruptId '{interrupt_id}' must include a boolean accepted value.", - code="APPROVAL_RESUME_INVALID", - ), - ) - pending_arguments = _pending_approval_arguments(pending_entry) - original_arguments = _parse_json_object(pending_arguments) or {} - direct_edited_arguments = { - key: value for key, value in payload.items() if key not in {"accepted", "approved", "editedArgs"} - } - standard_edited_arguments = payload.get("editedArgs") - if standard_edited_arguments is not None: - if not isinstance(standard_edited_arguments, dict) or direct_edited_arguments: - return ( - [], - handled_ids, - cancelled_ids, - RunErrorEvent( - message=( - f"Approval resume for interruptId '{interrupt_id}' must provide editedArgs as the " - "only edited-argument representation." - ), - code="APPROVAL_RESUME_INVALID_RESPONSE", - ), - ) - edited_arguments = cast(dict[str, Any], standard_edited_arguments) - if set(edited_arguments) != set(original_arguments): - return ( - [], - handled_ids, - cancelled_ids, - RunErrorEvent( - message=( - f"Approval resume for interruptId '{interrupt_id}' must provide editedArgs as a full " - "replacement of the pending tool arguments." - ), - code="APPROVAL_RESUME_INVALID_RESPONSE", - ), - ) - else: - edited_arguments = direct_edited_arguments - if edited_arguments and _pending_approval_server_label(pending_entry): - return ( - [], - handled_ids, - cancelled_ids, - RunErrorEvent( - message=( - f"Hosted approval resume for interruptId '{interrupt_id}' does not support edited arguments." - ), - code="APPROVAL_RESUME_INVALID_RESPONSE", - ), - ) - if not set(edited_arguments).issubset(set(original_arguments)): - return ( - [], - handled_ids, - cancelled_ids, - RunErrorEvent( - message=f"Approval resume for interruptId '{interrupt_id}' includes unsupported edited arguments.", - code="APPROVAL_RESUME_INVALID", - ), - ) - for name, edited_value in edited_arguments.items(): - original_value = original_arguments[name] - if not _json_schema_value_matches(original_value, edited_value): - return ( - [], - handled_ids, - cancelled_ids, - RunErrorEvent( - message=( - f"Approval resume for interruptId '{interrupt_id}' has invalid type for edited " - f"argument '{name}'." - ), - code="APPROVAL_RESUME_INVALID_RESPONSE", - ), - ) - - merged_arguments = ( - dict(edited_arguments) - if standard_edited_arguments is not None - else {**original_arguments, **edited_arguments} + accepted, canonical_arguments, merged_arguments, validation_error = _canonical_approval_decision( + entry.get("payload"), + interrupt_id=interrupt_id, + original_arguments_text=pending_arguments, + server_label=_pending_approval_server_label(pending_entry), ) - canonical_arguments = json.dumps(make_json_safe(merged_arguments), sort_keys=True, separators=(",", ":")) + if validation_error is not None: + return [], handled_ids, cancelled_ids, validation_error + if accepted is None or canonical_arguments is None or merged_arguments is None: + raise RuntimeError("Validated approval decision is missing canonical values.") lifecycle_decisions.append( ResumeDecision( interrupt_id=interrupt_id, @@ -1406,15 +1410,40 @@ def _canonical_approval_resume_messages( "arguments": make_json_safe(function_call.parse_arguments() or {}), } ) - messages.append({"role": "user", "function_approvals": function_approvals}) + messages.append( + { + "id": f"approval-response-{interrupt_id}", + "role": "user", + "function_approvals": function_approvals, + } + ) + lifecycle_cancelled_ids = [ + interrupt_id for interrupt_id in cancelled_ids if entries_by_interrupt_id.get(interrupt_id) is not None + ] if authorized_executions is not None: try: - intents = lifecycle.claim_batch(thread_id=thread_id, decisions=lifecycle_decisions) + intents = lifecycle.resolve_batch( + thread_id=thread_id, + decisions=lifecycle_decisions, + cancelled_interrupt_ids=lifecycle_cancelled_ids, + ) if snapshot_reconciliations is not None: snapshot_reconciliations.extend(intents.snapshot_reconciliations) for intent in intents: - authorized_executions[intent.identity.call_id] = intent + authorized_executions[intent.identity] = intent + except (KeyError, ValueError) as exc: + return ( + [], + handled_ids, + cancelled_ids, + RunErrorEvent(message=str(exc), code="APPROVAL_RESUME_INVALID"), + ) + elif lifecycle_cancelled_ids: + try: + reconciliations = lifecycle.cancel_batch(thread_id=thread_id, interrupt_ids=lifecycle_cancelled_ids) + if snapshot_reconciliations is not None: + snapshot_reconciliations.extend(reconciliations) except (KeyError, ValueError) as exc: return ( [], @@ -1435,9 +1464,9 @@ async def _resolve_approval_responses( validated_approved_responses: list[Content] | None = None, *, lifecycle: ApprovalLifecycle, - authorized_executions: dict[str, AuthorizedExecution] | None = None, + authorized_executions: dict[ApprovalOccurrenceIdentity, AuthorizedExecution] | None = None, forwarded_executions: ( - dict[str, tuple[ForwardedPendingToolTransitionOwner, AuthorizedExecution, Content]] | None + dict[str, list[tuple[ForwardedPendingToolTransitionOwner, AuthorizedExecution, Content]]] | None ) = None, ) -> list[Content]: """Execute approved function calls and replace approval content with results. @@ -1476,6 +1505,7 @@ async def _resolve_approval_responses( valid_response_content_ids = set() pending_local_response_content_ids = set() pending_response_groups: dict[object, tuple[ApprovalOccurrence, list[Content]]] = {} + intents_by_response_content_id: dict[int, AuthorizedExecution] = {} for response in approval_responses: resp_id = response.id function_call_id = response.function_call.call_id if response.function_call else None @@ -1558,11 +1588,12 @@ async def _resolve_approval_responses( and primary_response.function_call is not None ): call_id = primary_response.function_call.call_id or primary_response.id or "" - intent = authorized_executions.get(call_id) + intent = authorized_executions.get(pending_entry.identity) if intent is None: logger.warning("Approval remains pending because no transition owner can act for call_id=%s.", call_id) response_content_ids_to_strip.add(id(primary_response)) continue + intents_by_response_content_id[id(primary_response)] = intent valid_response_content_ids.add(id(primary_response)) if ( primary_response.approved @@ -1608,7 +1639,7 @@ async def _resolve_approval_responses( for approval in validated_forwarded_approvals: function_call = approval.function_call call_id = (function_call.call_id if function_call else None) or approval.id or "" - intent = authorized_executions.get(call_id) + intent = intents_by_response_content_id.get(id(approval)) if intent is None: logger.warning("Skipping hosted approval without lifecycle authority for call_id=%s.", call_id) continue @@ -1625,7 +1656,7 @@ async def forward_hosted_decision(approval: Content = approval) -> list[Content] forwarded = await forwarded_owner.forward(intent, lifecycle=lifecycle) if len(forwarded) != 1: raise RuntimeError("Hosted transition owner did not forward exactly one approval decision.") - forwarded_executions[call_id] = (forwarded_owner, intent, forwarded[0]) + forwarded_executions.setdefault(call_id, []).append((forwarded_owner, intent, forwarded[0])) fcc_todo = _collect_approval_responses(messages) if valid_response_content_ids is not None: @@ -1647,7 +1678,7 @@ async def forward_hosted_decision(approval: Content = approval) -> list[Content] for approval in approved_responses: function_call = approval.function_call call_id = (function_call.call_id if function_call else None) or approval.id or "" - intent = authorized_executions.get(call_id) if authorized_executions is not None else None + intent = intents_by_response_content_id.get(id(approval)) if intent is None or intent.owner is not ApprovalExecutionOwner.LOCAL: continue static_approved.append(approval) @@ -1664,7 +1695,7 @@ async def forward_hosted_decision(approval: Content = approval) -> list[Content] for approval in static_approved: function_call = approval.function_call call_id = (function_call.call_id if function_call else None) or approval.id or "" - intent = authorized_executions.get(call_id) + intent = intents_by_response_content_id.get(id(approval)) if intent is None: logger.warning("Skipping local approval without lifecycle authority for call_id=%s.", call_id) approved_function_result_groups.append([]) @@ -2236,8 +2267,8 @@ async def run_agent_stream( current_state=flow.current_state, ) - authorized_executions: dict[str, AuthorizedExecution] = {} - forwarded_executions: dict[str, tuple[ForwardedPendingToolTransitionOwner, AuthorizedExecution, Content]] = {} + authorized_executions: dict[ApprovalOccurrenceIdentity, AuthorizedExecution] = {} + forwarded_executions: dict[str, list[tuple[ForwardedPendingToolTransitionOwner, AuthorizedExecution, Content]]] = {} retained_approval_results: list[Content] = [] approval_snapshot_reconciliations: list[ApprovalSnapshotReconciliation] = [] client_tools = convert_agui_tools_to_agent_framework(input_data.get("tools")) @@ -2392,7 +2423,7 @@ async def run_agent_stream( if (tool := execution_tool_map.get(intent.name)) is None or getattr(tool, "declaration_only", False) ] if unavailable_local_intents: - for intent in local_intents: + for intent in authorized_executions.values(): approval_state_store.lifecycle.release_claim(intent, policy=ClaimRecoveryPolicy.SAFE_TO_RETRY) unavailable_names = ", ".join(sorted({intent.name for intent in unavailable_local_intents})) yield RunStartedEvent(run_id=run_id, thread_id=thread_id) @@ -2536,9 +2567,12 @@ async def run_agent_stream( if ( content_type == "function_result" and content.call_id - and (forwarded := forwarded_executions.pop(content.call_id, None)) is not None + and (forwarded_queue := forwarded_executions.get(content.call_id)) and approval_state_store is not None ): + forwarded = forwarded_queue.pop(0) + if not forwarded_queue: + forwarded_executions.pop(content.call_id, None) owner, intent, _ = forwarded owner.record_outcome(intent, [content], lifecycle=approval_state_store.lifecycle) @@ -2598,11 +2632,12 @@ async def run_agent_stream( stream_completed = True finally: if approval_state_store is not None: - for owner, intent, forwarded_approval in forwarded_executions.values(): - if stream_completed: - owner.record_outcome(intent, [forwarded_approval], lifecycle=approval_state_store.lifecycle) - else: - approval_state_store.lifecycle.recover_execution(intent, owner=intent.owner) + for queued_executions in forwarded_executions.values(): + for owner, intent, forwarded_approval in queued_executions: + if stream_completed: + owner.record_outcome(intent, [forwarded_approval], lifecycle=approval_state_store.lifecycle) + else: + approval_state_store.lifecycle.recover_execution(intent, owner=intent.owner) forwarded_executions.clear() if flow.waiting_for_approval and isinstance(stream, ResponseStream): diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py b/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py index cd2317dcec..c42400bda6 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py @@ -38,11 +38,13 @@ def wrapper(self: ApprovalLifecycle, *args: Any, **kwargs: Any) -> _ReturnT: if not isinstance(thread_id, str): raise TypeError("Serialized approval transitions require a thread_id.") decisions = kwargs.get("decisions") - interrupt_ids = ( + decision_interrupt_ids = ( [decision.interrupt_id for decision in decisions] if isinstance(decisions, list) and all(isinstance(decision, ResumeDecision) for decision in decisions) - else kwargs.get("interrupt_ids") + else [] ) + cancellation_interrupt_ids = kwargs.get("cancelled_interrupt_ids", kwargs.get("interrupt_ids", [])) + interrupt_ids = [*decision_interrupt_ids, *cancellation_interrupt_ids] if not isinstance(interrupt_ids, list) or not all( isinstance(interrupt_id, str) for interrupt_id in interrupt_ids ): @@ -366,43 +368,49 @@ def _register_aliases( def get(self, identity: ApprovalOccurrenceIdentity) -> ApprovalOccurrence: """Return server-owned state for a registered occurrence.""" - self._purge_expired_terminal() - return self._occurrences[identity] + with self._index_lock: + self._purge_expired_terminal() + return self._occurrences[identity] def decision_context(self, *, thread_id: str, interrupt_id: str) -> tuple[str, str]: """Return canonical server-owned call data needed to normalize a typed retry.""" - self._purge_expired_terminal() - key = (thread_id, interrupt_id) - identity = self._pending_by_interrupt.get(key) or self._terminal_by_interrupt[key] - occurrence = self._occurrences[identity] - return occurrence.name, occurrence.arguments + with self._index_lock: + self._purge_expired_terminal() + key = (thread_id, interrupt_id) + identity = self._pending_by_interrupt.get(key) or self._terminal_by_interrupt[key] + occurrence = self._occurrences[identity] + return occurrence.name, occurrence.arguments def pending_occurrence(self, *, thread_id: str, interrupt_id: str) -> ApprovalOccurrence | None: """Return pending server-owned state for one trusted interrupt alias.""" - self._purge_expired_terminal() - identity = self._pending_by_interrupt.get((thread_id, interrupt_id)) - return self._occurrences[identity] if identity is not None else None + with self._index_lock: + self._purge_expired_terminal() + identity = self._pending_by_interrupt.get((thread_id, interrupt_id)) + return self._occurrences[identity] if identity is not None else None def occurrence_for_alias(self, *, thread_id: str, interrupt_id: str) -> ApprovalOccurrence | None: """Return retained server-owned state for one trusted interrupt alias.""" - self._purge_expired_terminal() - key = (thread_id, interrupt_id) - identity = self._pending_by_interrupt.get(key) or self._terminal_by_interrupt.get(key) - return self._occurrences[identity] if identity is not None else None + with self._index_lock: + self._purge_expired_terminal() + key = (thread_id, interrupt_id) + identity = self._pending_by_interrupt.get(key) or self._terminal_by_interrupt.get(key) + return self._occurrences[identity] if identity is not None else None def occurrences_for_thread(self, *, thread_id: str) -> tuple[ApprovalOccurrence, ...]: """Return retained occurrences owned by one scoped thread.""" - self._purge_expired_terminal() - return tuple(occurrence for occurrence in self._occurrences.values() if thread_id in occurrence.thread_ids) + with self._index_lock: + self._purge_expired_terminal() + return tuple(occurrence for occurrence in self._occurrences.values() if thread_id in occurrence.thread_ids) def pending_interrupt_ids(self, *, thread_id: str) -> set[str]: """Return canonical interrupt identities with pending authority for one thread.""" - self._purge_expired_terminal() - return { - occurrence.identity.interrupt_id - for occurrence in self._occurrences.values() - if thread_id in occurrence.thread_ids and occurrence.status is ApprovalStatus.PENDING - } + with self._index_lock: + self._purge_expired_terminal() + return { + occurrence.identity.interrupt_id + for occurrence in self._occurrences.values() + if thread_id in occurrence.thread_ids and occurrence.status is ApprovalStatus.PENDING + } def reconcile_snapshot( self, @@ -411,24 +419,25 @@ def reconcile_snapshot( interrupt_ids: list[str], ) -> tuple[ApprovalSnapshotReconciliation, ...]: """Describe which stored approval controls remain actionable.""" - self._purge_expired_terminal() - reconciliations: list[ApprovalSnapshotReconciliation] = [] - for interrupt_id in interrupt_ids: - key = (thread_id, interrupt_id) - identity = self._pending_by_interrupt.get(key) or self._terminal_by_interrupt.get(key) - if identity is None: - reconciliations.append( - ApprovalSnapshotReconciliation( - interrupt_id=interrupt_id, - identity=None, - status=None, - retire_interrupt=True, + with self._index_lock: + self._purge_expired_terminal() + reconciliations: list[ApprovalSnapshotReconciliation] = [] + for interrupt_id in interrupt_ids: + key = (thread_id, interrupt_id) + identity = self._pending_by_interrupt.get(key) or self._terminal_by_interrupt.get(key) + if identity is None: + reconciliations.append( + ApprovalSnapshotReconciliation( + interrupt_id=interrupt_id, + identity=None, + status=None, + retire_interrupt=True, + ) ) - ) - continue - occurrence = self._occurrences[identity] - reconciliations.append(self._snapshot_reconciliation(occurrence)) - return tuple(reconciliations) + continue + occurrence = self._occurrences[identity] + reconciliations.append(self._snapshot_reconciliation(occurrence)) + return tuple(reconciliations) @staticmethod def _snapshot_reconciliation(occurrence: ApprovalOccurrence) -> ApprovalSnapshotReconciliation: @@ -553,6 +562,32 @@ def claim_batch( snapshot_reconciliations=tuple(snapshot_reconciliations), ) + @_serialized_by_batch + def resolve_batch( + self, + *, + thread_id: str, + decisions: list[ResumeDecision], + cancelled_interrupt_ids: list[str], + ) -> ApprovalBatchDecision: + """Atomically validate and apply resolved and cancelled decisions for one resume batch.""" + self._purge_expired_terminal() + cancelled_occurrences, retained_cancellations = self._validate_cancellations( + thread_id=thread_id, + interrupt_ids=cancelled_interrupt_ids, + ) + decision_batch = self.claim_batch(thread_id=thread_id, decisions=decisions) + cancellation_reconciliations = self._cancel_occurrences(cancelled_occurrences) + return ApprovalBatchDecision( + authorized_executions=decision_batch.authorized_executions, + retained_outcomes=decision_batch.retained_outcomes, + snapshot_reconciliations=( + *decision_batch.snapshot_reconciliations, + *retained_cancellations, + *cancellation_reconciliations, + ), + ) + @_serialized_by_batch def cancel_batch( self, @@ -562,6 +597,19 @@ def cancel_batch( ) -> tuple[ApprovalSnapshotReconciliation, ...]: """Validate and cancel selected occurrences without changing their siblings.""" self._purge_expired_terminal() + occurrences, reconciliations = self._validate_cancellations( + thread_id=thread_id, + interrupt_ids=interrupt_ids, + ) + return (*reconciliations, *self._cancel_occurrences(occurrences)) + + def _validate_cancellations( + self, + *, + thread_id: str, + interrupt_ids: list[str], + ) -> tuple[list[ApprovalOccurrence], tuple[ApprovalSnapshotReconciliation, ...]]: + """Validate cancellations without mutating lifecycle state.""" occurrences: list[ApprovalOccurrence] = [] reconciliations: list[ApprovalSnapshotReconciliation] = [] seen_interrupt_ids: set[str] = set() @@ -584,6 +632,14 @@ def cancel_batch( raise ValueError(f"Approval occurrence is not pending: {occurrence.status}.") occurrences.append(occurrence) + return occurrences, tuple(reconciliations) + + def _cancel_occurrences( + self, + occurrences: list[ApprovalOccurrence], + ) -> tuple[ApprovalSnapshotReconciliation, ...]: + """Commit previously validated cancellations while their batch locks are held.""" + reconciliations: list[ApprovalSnapshotReconciliation] = [] for occurrence in occurrences: occurrence.status = ApprovalStatus.CANCELLED self._remove_pending_aliases(occurrence) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py b/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py index a09e18d5fc..5bae76017b 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py @@ -5,6 +5,7 @@ from __future__ import annotations import copy +from threading import RLock from typing import Any from ._approval_lifecycle import ApprovalCapacityError, ApprovalExecutionOwner, ApprovalLifecycle @@ -59,6 +60,7 @@ def __init__( if max_entries < 1: raise ValueError("max_entries must be greater than 0.") self.max_entries = max_entries + self._lock = RLock() self._tool_approval_states: dict[str, dict[str, Any]] = {} self.lifecycle = ApprovalLifecycle( max_entries=max_entries, @@ -93,19 +95,23 @@ def register( def set_tool_approval_state(self, thread_id: str, state: dict[str, Any]) -> None: """Store approval middleware state without evicting another active thread.""" - if thread_id not in self._tool_approval_states and len(self._tool_approval_states) >= self.max_entries: - raise ApprovalCapacityError("Approval state capacity is exhausted by protected occurrences.") - self._tool_approval_states[thread_id] = copy.deepcopy(state) + with self._lock: + if thread_id not in self._tool_approval_states and len(self._tool_approval_states) >= self.max_entries: + raise ApprovalCapacityError("Approval state capacity is exhausted by protected occurrences.") + self._tool_approval_states[thread_id] = copy.deepcopy(state) def get_tool_approval_state(self, thread_id: str) -> dict[str, Any] | None: """Return an isolated copy of server-owned middleware approval state.""" - state = self._tool_approval_states.get(thread_id) - return copy.deepcopy(state) if state is not None else None + with self._lock: + state = self._tool_approval_states.get(thread_id) + return copy.deepcopy(state) if state is not None else None def delete_tool_approval_state(self, thread_id: str) -> None: """Delete server-owned middleware approval state for one scoped thread.""" - self._tool_approval_states.pop(thread_id, None) + with self._lock: + self._tool_approval_states.pop(thread_id, None) def has_tool_approval_state(self, thread_id: str) -> bool: """Return whether middleware approval state exists for one scoped thread.""" - return thread_id in self._tool_approval_states + with self._lock: + return thread_id in self._tool_approval_states diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py index bf1658c6c0..3606428f60 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py @@ -374,6 +374,7 @@ def _json_schema_for_value(value: Any) -> dict[str, Any]: def _approval_response_schema(arguments: Mapping[str, Any] | None = None) -> dict[str, Any]: """Build the response schema generic AG-UI clients use to render approval input.""" + reserved_properties = {"approved", "accepted", "editedArgs"} properties: dict[str, Any] = { "approved": { "type": "boolean", @@ -389,7 +390,8 @@ def _approval_response_schema(arguments: Mapping[str, Any] | None = None) -> dic for name, value in arguments.items(): argument_schema = _json_schema_for_value(value) argument_schema["description"] = f"Optional edited value for the '{name}' tool argument." - properties[str(name)] = argument_schema + if str(name) not in reserved_properties: + properties[str(name)] = argument_schema edited_argument_properties[str(name)] = _json_schema_for_value(value) properties["editedArgs"] = { "type": "object", diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py index 4128420862..26c93deae5 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py @@ -939,6 +939,7 @@ async def run_workflow_stream( pending_before_run = await _pending_request_events(workflow) pending_interrupt_ids = _pending_workflow_interrupt_ids(pending_before_run) resume_entries: list[dict[str, Any]] = [] + cancelled_request_ids: set[str] = set() if pending_interrupt_ids: resume_entries, contract_error, contract_code = _resume_contract_error( resume_payload, @@ -960,13 +961,6 @@ async def run_workflow_stream( cancelled_request_ids = { str(entry["interrupt_id"]) for entry in resume_entries if entry.get("status") == "cancelled" } - if cancelled_request_ids: - _consume_cancelled_workflow_requests(workflow, resume_entries) - pending_before_run = { - request_id: request_event - for request_id, request_event in pending_before_run.items() - if str(getattr(request_event, "request_id", None) or request_id) not in cancelled_request_ids - } resume_responses = ( _resume_entries_to_workflow_responses(resume_entries) @@ -984,6 +978,13 @@ async def run_workflow_stream( yield RunStartedEvent(run_id=run_id, thread_id=thread_id) yield response_error return + if cancelled_request_ids: + _consume_cancelled_workflow_requests(workflow, resume_entries) + pending_before_run = { + request_id: request_event + for request_id, request_event in pending_before_run.items() + if str(getattr(request_event, "request_id", None) or request_id) not in cancelled_request_ids + } pending_interrupts = _interrupts_from_pending_requests(pending_before_run) if not responses and pending_before_run: diff --git a/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py b/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py index 9b81454a62..7e87e6e70f 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py +++ b/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py @@ -436,6 +436,65 @@ def repeat_first_claim(): conflicting_claim.result(timeout=2) +def test_terminal_purge_cannot_mutate_indexes_during_thread_snapshot_read() -> None: + """Thread-scoped lifecycle reads remain consistent while another operation purges terminal state.""" + now = 0.0 + lifecycle = ApprovalLifecycle(terminal_retention_seconds=1, clock=lambda: now) + terminal = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, + thread_id="terminal-thread", + interrupt_id="terminal-approval", + call_id="terminal-call", + name="write_record", + arguments="{}", + ) + lifecycle.cancel_batch(thread_id="terminal-thread", interrupt_ids=["terminal-approval"]) + lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, + thread_id="pending-thread", + interrupt_id="pending-approval", + call_id="pending-call", + name="write_record", + arguments="{}", + ) + read_entered = Event() + release_read = Event() + + class SlowThreadId(str): + def __eq__(self, other: object) -> bool: + read_entered.set() + assert release_read.wait(timeout=2) + return super().__eq__(other) + + __hash__ = str.__hash__ + + def read_occurrences() -> tuple[object, ...]: + return lifecycle.occurrences_for_thread(thread_id=SlowThreadId("unrelated-thread")) + + def purge_during_read() -> None: + nonlocal now + now = 2.0 + lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, + thread_id="new-thread", + interrupt_id="new-approval", + call_id="new-call", + name="write_record", + arguments="{}", + ) + + with ThreadPoolExecutor(max_workers=2) as executor: + read = executor.submit(read_occurrences) + assert read_entered.wait(timeout=2) + purge = executor.submit(purge_during_read) + release_read.set() + assert read.result(timeout=2) == () + purge.result(timeout=2) + + with pytest.raises(KeyError): + lifecycle.get(terminal.identity) + + async def test_hosted_approval_is_forwarded_only_by_its_owner_and_settles_same_occurrence() -> None: """Hosted authority cannot execute locally and records forwarding against its occurrence.""" lifecycle = ApprovalLifecycle() diff --git a/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py b/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py index 9c58277cc4..4652d028c8 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py +++ b/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py @@ -8,6 +8,7 @@ from typing import Any from agent_framework import AgentResponseUpdate, Content, FunctionTool +from agent_framework.exceptions import UserInputRequiredException from conftest import StubAgent # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports] from agent_framework_ag_ui._agent import AgentConfig @@ -86,6 +87,37 @@ async def _run_resume( return events +async def _run_custom_approval(tool: FunctionTool) -> tuple[list[Any], StubAgent]: + agent = StubAgent( + updates=[AgentResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant")], + default_options={"tools": [tool]}, + ) + store = InMemoryAGUIApprovalStateStore() + store.register( + owner=ApprovalExecutionOwner.LOCAL, + thread_ids=["thread-custom"], + name=tool.name, + arguments="{}", + request_id="approval-custom", + interrupt_id="call-custom", + ) + events = [ + event + async for event in run_agent_stream( + { + "thread_id": "thread-custom", + "run_id": "resume-custom", + "messages": [{"role": "user", "content": "Continue"}], + "resume": [{"interruptId": "call-custom", "status": "resolved", "payload": {"approved": True}}], + }, + agent, + AgentConfig(), + approval_state_store=store, + ) + ] + return events, agent + + async def test_approved_call_emits_one_live_result_under_original_identity() -> None: executions: list[str] = [] @@ -128,3 +160,71 @@ async def test_mixed_batch_preserves_approved_result_identity_and_order() -> Non results = [event for event in events if getattr(event, "type", None) == "TOOL_CALL_RESULT"] assert executions == ["Seattle"] assert [(event.tool_call_id, event.content) for event in results] == [("call-seattle", "Sunny in Seattle")] + + +async def test_approval_follow_up_group_remains_in_history_without_live_tool_result() -> None: + """Approval-time user-input requests remain grouped without becoming terminal tool results.""" + + def request_consent() -> str: + raise UserInputRequiredException( + contents=[ + Content.from_oauth_consent_request(consent_link="https://example.com/first"), + Content.from_oauth_consent_request(consent_link="https://example.com/second"), + ] + ) + + events, agent = await _run_custom_approval( + FunctionTool(name="request_consent", description="Request consent", func=request_consent) + ) + + assert not [event for event in events if getattr(event, "type", None) == "TOOL_CALL_RESULT"] + requests = [ + content + for message in agent.messages_received + for content in message.contents + if isinstance(content, Content) and content.user_input_request + ] + assert [request.consent_link for request in requests] == [ + "https://example.com/first", + "https://example.com/second", + ] + + +async def test_approval_execution_failure_emits_one_terminal_error_result() -> None: + """An approved tool failure produces one deterministic terminal result.""" + + def fail() -> str: + raise RuntimeError("secret failure detail") + + events, _ = await _run_custom_approval(FunctionTool(name="fail", description="Fail", func=fail)) + + results = [event for event in events if getattr(event, "type", None) == "TOOL_CALL_RESULT"] + assert [(event.tool_call_id, event.content) for event in results] == [("call-custom", "Error: Function failed.")] + + +async def test_no_approval_path_emits_no_approval_specific_duplicate_result() -> None: + """An ordinary provider tool result is projected once without approval bookkeeping.""" + agent = StubAgent( + updates=[ + AgentResponseUpdate( + contents=[Content.from_function_result(call_id="call-ordinary", result="ordinary result")], + role="tool", + ) + ] + ) + + events = [ + event + async for event in run_agent_stream( + { + "thread_id": "thread-ordinary", + "run_id": "run-ordinary", + "messages": [{"role": "user", "content": "Run it"}], + }, + agent, + AgentConfig(), + ) + ] + + results = [event for event in events if getattr(event, "type", None) == "TOOL_CALL_RESULT"] + assert [(event.tool_call_id, event.content) for event in results] == [("call-ordinary", "ordinary result")] diff --git a/python/packages/ag-ui/tests/ag_ui/test_approval_state.py b/python/packages/ag-ui/tests/ag_ui/test_approval_state.py index 58a61391f3..9a4a738584 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_approval_state.py +++ b/python/packages/ag-ui/tests/ag_ui/test_approval_state.py @@ -2,6 +2,10 @@ """Tests for server-side AG-UI approval state storage.""" +from concurrent.futures import ThreadPoolExecutor +from threading import Barrier +from time import sleep + import pytest from agent_framework_ag_ui._approval_lifecycle import ApprovalCapacityError, ApprovalExecutionOwner @@ -83,3 +87,29 @@ def test_approval_state_store_does_not_evict_active_middleware_state() -> None: assert store.get_tool_approval_state("thread-1") == {"call_id": "call-1"} assert store.get_tool_approval_state("thread-2") is None + + +def test_approval_state_store_enforces_capacity_across_concurrent_first_writes() -> None: + """Concurrent first writes cannot reserve more middleware slots than configured.""" + store = InMemoryAGUIApprovalStateStore(max_entries=1) + start = Barrier(2) + + class SlowCopy: + def __deepcopy__(self, memo: dict[int, object]) -> "SlowCopy": + del memo + sleep(0.05) + return self + + def write(thread_id: str) -> str: + start.wait(timeout=2) + try: + store.set_tool_approval_state(thread_id, {"call_id": thread_id, "slow": SlowCopy()}) + except ApprovalCapacityError: + return "rejected" + return "stored" + + with ThreadPoolExecutor(max_workers=2) as executor: + outcomes = list(executor.map(write, ["thread-1", "thread-2"])) + + assert sorted(outcomes) == ["rejected", "stored"] + assert sum(store.has_tool_approval_state(thread_id) for thread_id in ["thread-1", "thread-2"]) == 1 diff --git a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py index 8ed55d80f5..67d834162a 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -47,7 +47,7 @@ add_agent_framework_fastapi_endpoint, ) from agent_framework_ag_ui._agent import AgentFrameworkAgent -from agent_framework_ag_ui._approval_lifecycle import ApprovalLifecycle +from agent_framework_ag_ui._approval_lifecycle import ApprovalExecutionOwner, ApprovalLifecycle, ApprovalStatus from agent_framework_ag_ui._approval_state import InMemoryAGUIApprovalStateStore, approval_state_thread_id from agent_framework_ag_ui._workflow import AgentFrameworkWorkflow @@ -1809,6 +1809,120 @@ def get_weather(city: str) -> str: return client, agent, executed_cities +async def test_endpoint_agent_approval_batch_keeps_distinct_occurrences_for_reused_call_id() -> None: + """Unique interrupts sharing a provider call ID each execute and settle exactly once.""" + executed: list[str] = [] + + def first_tool() -> str: + executed.append("first") + return "first result" + + def second_tool() -> str: + executed.append("second") + return "second result" + + tools = [ + FunctionTool(name="first_tool", description="First", func=first_tool), + FunctionTool(name="second_tool", description="Second", func=second_tool), + ] + agent = StubAgent( + updates=[AgentResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant")], + default_options={"tools": tools}, + ) + wrapped_agent = AgentFrameworkAgent(agent=agent, require_confirmation=False) + lifecycle = wrapped_agent._approval_state_store.lifecycle + first = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, + thread_id="thread-shared-call", + interrupt_id="approval-first", + call_id="call-shared", + name="first_tool", + arguments="{}", + ) + second = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, + thread_id="thread-shared-call", + interrupt_id="approval-second", + call_id="call-shared", + name="second_tool", + arguments="{}", + ) + app = FastAPI() + add_agent_framework_fastapi_endpoint(app, wrapped_agent, path="/approval") + + response = TestClient(app).post( + "/approval", + json={ + "runId": "run-shared-call", + "threadId": "thread-shared-call", + "messages": [], + "resume": [ + {"interruptId": "approval-first", "status": "resolved", "payload": {"approved": True}}, + {"interruptId": "approval-second", "status": "resolved", "payload": {"approved": True}}, + ], + }, + ) + + assert response.status_code == 200 + events = _decode_sse_events(response) + assert not [event for event in events if event.get("type") == "RUN_ERROR"] + assert lifecycle.get(first.identity).status is ApprovalStatus.SETTLED + assert lifecycle.get(second.identity).status is ApprovalStatus.SETTLED + assert executed == ["first", "second"] + assert [event["content"] for event in events if event.get("type") == "TOOL_CALL_RESULT"] == [ + "first result", + "second result", + ] + + +async def test_endpoint_agent_unavailable_local_executor_releases_every_unstarted_batch_intent() -> None: + """A local executor failure leaves local and hosted siblings retryable.""" + agent = StubAgent( + updates=[AgentResponseUpdate(contents=[Content.from_text(text="Should not run.")], role="assistant")], + default_options={"tools": []}, + ) + wrapped_agent = AgentFrameworkAgent(agent=agent, require_confirmation=False) + lifecycle = wrapped_agent._approval_state_store.lifecycle + local = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, + thread_id="thread-unavailable-batch", + interrupt_id="approval-local", + call_id="call-local", + name="missing_local_tool", + arguments="{}", + ) + hosted = lifecycle.register( + owner=ApprovalExecutionOwner.HOSTED, + thread_id="thread-unavailable-batch", + interrupt_id="approval-hosted", + call_id="call-hosted", + name="hosted_tool", + arguments="{}", + server_label="hosted-server", + ) + app = FastAPI() + add_agent_framework_fastapi_endpoint(app, wrapped_agent, path="/approval") + + response = TestClient(app).post( + "/approval", + json={ + "runId": "run-unavailable-batch", + "threadId": "thread-unavailable-batch", + "messages": [], + "resume": [ + {"interruptId": "approval-local", "status": "resolved", "payload": {"approved": True}}, + {"interruptId": "approval-hosted", "status": "resolved", "payload": {"approved": True}}, + ], + }, + ) + + assert response.status_code == 200 + errors = [event for event in _decode_sse_events(response) if event.get("type") == "RUN_ERROR"] + assert [error["code"] for error in errors] == ["APPROVAL_TOOL_UNAVAILABLE"] + assert lifecycle.get(local.identity).status is ApprovalStatus.PENDING + assert lifecycle.get(hosted.identity).status is ApprovalStatus.PENDING + + def _build_mixed_approval_batch_endpoint( streaming_chat_client_stub: Any, *, diff --git a/python/packages/ag-ui/tests/ag_ui/test_run.py b/python/packages/ag-ui/tests/ag_ui/test_run.py index d986725806..b8075b2e48 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_run.py @@ -35,7 +35,13 @@ _should_suppress_intermediate_snapshot, run_agent_stream, ) -from agent_framework_ag_ui._approval_lifecycle import ApprovalExecutionOwner, ApprovalLifecycle +from agent_framework_ag_ui._approval_lifecycle import ( + ApprovalExecutionOwner, + ApprovalLifecycle, + ApprovalStatus, + ResumeDecision, +) +from agent_framework_ag_ui._approval_state import InMemoryAGUIApprovalStateStore from agent_framework_ag_ui._run_common import ( FlowState, _build_run_finished_event, @@ -937,6 +943,34 @@ def test_emit_approval_request_populates_interrupt_metadata(): } +def test_emit_approval_request_keeps_protocol_fields_when_tool_arguments_use_reserved_names() -> None: + """Reserved protocol fields remain controls while editedArgs carries colliding tool arguments.""" + flow = FlowState(message_id="msg-1") + function_call = Content.from_function_call( + call_id="call_reserved", + name="write_doc", + arguments={"approved": "draft", "accepted": 1, "editedArgs": {"value": True}}, + ) + approval_content = Content.from_function_approval_request(id="approval_reserved", function_call=function_call) + + _emit_approval_request(approval_content, flow) + + properties = flow.interrupts[0]["responseSchema"]["properties"] + assert properties["approved"]["type"] == "boolean" + assert properties["accepted"]["type"] == "boolean" + assert properties["editedArgs"] == { + "type": "object", + "description": "Full replacement of the tool arguments. Not merged.", + "properties": { + "approved": {"type": "string"}, + "accepted": {"type": "integer"}, + "editedArgs": {"type": "object", "additionalProperties": True}, + }, + "required": ["approved", "accepted", "editedArgs"], + "additionalProperties": False, + } + + def test_emit_approval_request_accumulates_multiple_interrupts(): """Multiple approval requests in the same turn should accumulate in flow.interrupts.""" flow = FlowState(message_id="msg-1") @@ -1076,6 +1110,167 @@ def test_canonical_approval_resume_does_not_mutate_arguments_until_batch_validat assert pending_entry.arguments == '{"city":"Seattle"}' +def test_canonical_approval_resume_does_not_cancel_until_resolved_siblings_validate() -> None: + """A malformed resolved sibling leaves every approval in the batch pending.""" + lifecycle = ApprovalLifecycle() + cancelled = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, + thread_id="thread-weather", + interrupt_id="call_a", + call_id="call_a", + name="get_weather", + arguments='{"city":"Seattle"}', + ) + resolved = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, + thread_id="thread-weather", + interrupt_id="call_b", + call_id="call_b", + name="get_weather", + arguments='{"city":"Portland"}', + ) + + _, _, _, error = _canonical_approval_resume_messages( + [ + {"interruptId": "call_a", "status": "cancelled"}, + {"interruptId": "call_b", "status": "resolved", "payload": "not an object"}, + ], + "thread-weather", + lifecycle=lifecycle, + ) + + assert error is not None + assert error.code == "APPROVAL_RESUME_INVALID" + assert lifecycle.get(cancelled.identity).status is ApprovalStatus.PENDING + assert lifecycle.get(resolved.identity).status is ApprovalStatus.PENDING + + +def test_terminal_approval_retry_validates_standard_edited_arguments() -> None: + """A terminal retry cannot bypass the pending path's editedArgs contract.""" + lifecycle = ApprovalLifecycle() + lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, + thread_id="thread-weather", + interrupt_id="call_a", + call_id="call_a", + name="get_weather", + arguments='{"city":"Seattle"}', + ) + lifecycle.claim_batch( + thread_id="thread-weather", + decisions=[ + ResumeDecision( + interrupt_id="call_a", + accepted=False, + arguments='{"city":"Seattle"}', + name="get_weather", + original_arguments='{"city":"Seattle"}', + ) + ], + ) + retained_results: list[Content] = [] + + _, handled_ids, _, error = _canonical_approval_resume_messages( + [ + { + "interruptId": "call_a", + "status": "resolved", + "payload": {"approved": False, "editedArgs": "not an object"}, + } + ], + "thread-weather", + lifecycle=lifecycle, + retained_results=retained_results, + ) + + assert handled_ids == set() + assert error is not None + assert error.code == "APPROVAL_RESUME_INVALID_RESPONSE" + assert retained_results == [] + + +def test_terminal_rejection_retry_does_not_project_a_live_tool_result() -> None: + """An identical rejected retry has the same no-result projection as the original rejection.""" + lifecycle = ApprovalLifecycle() + lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, + thread_id="thread-weather", + interrupt_id="call_a", + call_id="call_a", + name="get_weather", + arguments='{"city":"Seattle"}', + ) + decision = ResumeDecision( + interrupt_id="call_a", + accepted=False, + arguments='{"city":"Seattle"}', + name="get_weather", + original_arguments='{"city":"Seattle"}', + ) + lifecycle.claim_batch(thread_id="thread-weather", decisions=[decision]) + retained_results: list[Content] = [] + + _, handled_ids, _, error = _canonical_approval_resume_messages( + [ + { + "interruptId": "call_a", + "status": "resolved", + "payload": {"approved": False}, + } + ], + "thread-weather", + lifecycle=lifecycle, + retained_results=retained_results, + ) + + assert error is None + assert handled_ids == {"call_a"} + assert retained_results == [] + + +async def test_run_settles_server_collected_rejection_in_lifecycle() -> None: + """A rejection restored from approval middleware state no longer remains pending.""" + function_call = Content.from_function_call( + call_id="call_rejected", + name="write_record", + arguments={"value": "draft"}, + ) + response = Content.from_function_approval_response( + approved=False, + id="approval_rejected", + function_call=function_call, + ) + store = InMemoryAGUIApprovalStateStore() + store.set_tool_approval_state( + "thread-server-rejection", + {"collected_approval_responses": [response.to_dict()]}, + ) + agent = StubAgent(updates=[AgentResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant")]) + + events = [ + event + async for event in run_agent_stream( + { + "runId": "run-server-rejection", + "threadId": "thread-server-rejection", + "messages": [{"role": "user", "content": "Continue"}], + }, + agent, + AgentConfig(), + approval_state_store=store, + ) + ] + + assert not [event for event in events if event.type == "RUN_ERROR"] + occurrence = store.lifecycle.occurrence_for_alias( + thread_id="thread-server-rejection", + interrupt_id="approval_rejected", + ) + assert occurrence is not None + assert occurrence.status is ApprovalStatus.REJECTED + assert store.lifecycle.pending_interrupt_ids(thread_id="thread-server-rejection") == set() + + def test_canonical_hosted_approval_resume_rejects_edited_arguments_without_mutating_pending() -> None: """Hosted approvals accept a decision only because providers ignore edited arguments.""" lifecycle = ApprovalLifecycle() diff --git a/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py b/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py index 6eac56864c..59c1fce90b 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py @@ -848,6 +848,65 @@ async def requester(message: Any, ctx: WorkflowContext) -> None: assert getattr(run_error, "code") == "WORKFLOW_RESUME_REQUIRED" +async def test_workflow_run_does_not_cancel_until_resolved_siblings_validate() -> None: + """A malformed resolved workflow sibling leaves a cancelled request pending.""" + cancelled_call = Content.from_function_call( + call_id="call-cancelled", + name="write_record", + arguments={"value": "cancelled"}, + ) + resolved_call = Content.from_function_call( + call_id="call-resolved", + name="write_record", + arguments={"value": "resolved"}, + ) + pending = { + "approval-cancelled": SimpleNamespace( + request_id="approval-cancelled", + data=Content.from_function_approval_request(id="approval-cancelled", function_call=cancelled_call), + response_type=Content, + ), + "approval-resolved": SimpleNamespace( + request_id="approval-resolved", + data=Content.from_function_approval_request(id="approval-resolved", function_call=resolved_call), + response_type=Content, + ), + } + + async def get_pending_request_info_events() -> dict[str, Any]: + return dict(pending) + + runner_context = SimpleNamespace( + get_pending_request_info_events=get_pending_request_info_events, + _pending_request_info_events=pending, + ) + workflow = SimpleNamespace(_runner_context=runner_context) + + events = [ + event + async for event in run_workflow_stream( + { + "runId": "run-mixed-invalid", + "threadId": "thread-mixed-invalid", + "messages": [], + "resume": [ + {"interruptId": "approval-cancelled", "status": "cancelled"}, + { + "interruptId": "approval-resolved", + "status": "resolved", + "payload": {"approved": True, "editedArgs": "not an object"}, + }, + ], + }, + cast(Any, workflow), + ) + ] + + assert [event.type for event in events] == ["RUN_STARTED", "RUN_ERROR"] + assert getattr(events[-1], "code") == "WORKFLOW_RESUME_INVALID_RESPONSE" + assert set(runner_context._pending_request_info_events) == {"approval-cancelled", "approval-resolved"} + + async def test_workflow_run_agent_response_output_uses_latest_assistant_message_only() -> None: """Conversation payload outputs should not flatten full history into one assistant message.""" From 0b17f980288e6be8177871c848126b8d2d1169a0 Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Wed, 12 Aug 2026 11:42:25 +0900 Subject: [PATCH 12/13] fix AG-UI test typing checks --- .../tests/ag_ui/test_agent_wrapper_comprehensive.py | 9 ++++++--- .../ag-ui/tests/ag_ui/test_approval_lifecycle.py | 2 +- .../ag-ui/tests/ag_ui/test_approval_result_event.py | 9 +++++---- python/packages/ag-ui/tests/ag_ui/test_approval_state.py | 3 ++- 4 files changed, 14 insertions(+), 9 deletions(-) diff --git a/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py b/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py index 5116282925..692b280963 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py +++ b/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py @@ -7,6 +7,7 @@ from typing import Any import pytest +from ag_ui.core import RunErrorEvent, ToolCallResultEvent from agent_framework import Agent, ChatOptions, ChatResponseUpdate, Content, Message from pydantic import BaseModel @@ -1195,7 +1196,7 @@ def approval_input(thread_id: str) -> dict[str, Any]: retry_events = [event async for event in wrapper.run(approval_input(replay_thread_id))] assert execution_count == 1 - retry_results = [event for event in retry_events if event.type == "TOOL_CALL_RESULT"] + retry_results = [event for event in retry_events if isinstance(event, ToolCallResultEvent)] assert len(retry_results) == 1 assert retry_results[0].tool_call_id == "call_sensitive" assert retry_results[0].content == "executed" @@ -1206,14 +1207,16 @@ def approval_input(thread_id: str) -> dict[str, Any]: conflicting_events = [event async for event in wrapper.run(conflicting_input)] assert execution_count == 1 - assert any(event.type == "RUN_ERROR" and event.code == "APPROVAL_RESUME_INVALID" for event in conflicting_events) + assert any( + isinstance(event, RunErrorEvent) and event.code == "APPROVAL_RESUME_INVALID" for event in conflicting_events + ) changed_input = approval_input(replay_thread_id) changed_input["resume"][0]["payload"]["forged"] = True changed_events = [event async for event in wrapper.run(changed_input)] assert execution_count == 1 - assert any(event.type == "RUN_ERROR" and event.code == "APPROVAL_RESUME_INVALID" for event in changed_events) + assert any(isinstance(event, RunErrorEvent) and event.code == "APPROVAL_RESUME_INVALID" for event in changed_events) async def test_approval_function_name_mismatch_is_blocked(streaming_chat_client_stub): diff --git a/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py b/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py index 7e87e6e70f..8aa414dd33 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py +++ b/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py @@ -415,7 +415,7 @@ def blocking_emit(event: str, occurrence=None, *, failure_type: str | None = Non assert release_first_claim.wait(timeout=2) original_emit(event, occurrence, failure_type=failure_type) - lifecycle._emit_event = blocking_emit # type: ignore[method-assign] + lifecycle._emit_event = blocking_emit # type: ignore[method-assign] # ty: ignore[invalid-assignment] def repeat_first_claim(): second_same_thread_started.set() diff --git a/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py b/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py index 4652d028c8..78edb604e5 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py +++ b/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py @@ -7,6 +7,7 @@ import json from typing import Any +from ag_ui.core import ToolCallResultEvent from agent_framework import AgentResponseUpdate, Content, FunctionTool from agent_framework.exceptions import UserInputRequiredException from conftest import StubAgent # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports] @@ -128,7 +129,7 @@ async def test_approved_call_emits_one_live_result_under_original_identity() -> executions=executions, ) - results = [event for event in events if getattr(event, "type", None) == "TOOL_CALL_RESULT"] + results = [event for event in events if isinstance(event, ToolCallResultEvent)] assert executions == ["Seattle"] assert [(event.tool_call_id, event.content) for event in results] == [("call-weather", "Sunny in Seattle")] @@ -157,7 +158,7 @@ async def test_mixed_batch_preserves_approved_result_identity_and_order() -> Non executions=executions, ) - results = [event for event in events if getattr(event, "type", None) == "TOOL_CALL_RESULT"] + results = [event for event in events if isinstance(event, ToolCallResultEvent)] assert executions == ["Seattle"] assert [(event.tool_call_id, event.content) for event in results] == [("call-seattle", "Sunny in Seattle")] @@ -198,7 +199,7 @@ def fail() -> str: events, _ = await _run_custom_approval(FunctionTool(name="fail", description="Fail", func=fail)) - results = [event for event in events if getattr(event, "type", None) == "TOOL_CALL_RESULT"] + results = [event for event in events if isinstance(event, ToolCallResultEvent)] assert [(event.tool_call_id, event.content) for event in results] == [("call-custom", "Error: Function failed.")] @@ -226,5 +227,5 @@ async def test_no_approval_path_emits_no_approval_specific_duplicate_result() -> ) ] - results = [event for event in events if getattr(event, "type", None) == "TOOL_CALL_RESULT"] + results = [event for event in events if isinstance(event, ToolCallResultEvent)] assert [(event.tool_call_id, event.content) for event in results] == [("call-ordinary", "ordinary result")] diff --git a/python/packages/ag-ui/tests/ag_ui/test_approval_state.py b/python/packages/ag-ui/tests/ag_ui/test_approval_state.py index 9a4a738584..e7ef9f8c27 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_approval_state.py +++ b/python/packages/ag-ui/tests/ag_ui/test_approval_state.py @@ -7,6 +7,7 @@ from time import sleep import pytest +from typing_extensions import Self from agent_framework_ag_ui._approval_lifecycle import ApprovalCapacityError, ApprovalExecutionOwner from agent_framework_ag_ui._approval_state import InMemoryAGUIApprovalStateStore, approval_state_thread_id @@ -95,7 +96,7 @@ def test_approval_state_store_enforces_capacity_across_concurrent_first_writes() start = Barrier(2) class SlowCopy: - def __deepcopy__(self, memo: dict[int, object]) -> "SlowCopy": + def __deepcopy__(self, memo: dict[int, object]) -> Self: del memo sleep(0.05) return self From 6aa820af4fda1f652d372f32d0141013a8f7297b Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Wed, 12 Aug 2026 12:12:56 +0900 Subject: [PATCH 13/13] fix AG-UI approval retention and cancellation retries --- .../specs/004-python-function-calling-loop.md | 9 +- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 33 ++++ .../_approval_lifecycle.py | 48 +++++- .../agent_framework_ag_ui/_approval_state.py | 10 ++ .../tests/ag_ui/test_approval_lifecycle.py | 157 ++++++++++++++++++ .../ag-ui/tests/ag_ui/test_endpoint.py | 20 +++ 6 files changed, 271 insertions(+), 6 deletions(-) diff --git a/docs/specs/004-python-function-calling-loop.md b/docs/specs/004-python-function-calling-loop.md index e682fd6213..5caf0b6509 100644 --- a/docs/specs/004-python-function-calling-loop.md +++ b/docs/specs/004-python-function-calling-loop.md @@ -364,7 +364,11 @@ that manually replay messages own the equivalent rule: do not resend an approval - AG-UI tool approval resumes accept the standard `approved` decision and full-replacement `editedArgs` payload. Existing MAF clients remain compatible through the `accepted` decision alias and direct partial argument edits. - An AG-UI `cancelled` resume is a valid terminal decision, not a run error. In a resume covering parallel open - interrupts, resolved siblings still execute and cancelled calls do not. + interrupts, resolved siblings still execute and cancelled calls do not. An identical cancellation retry during + the retained terminal window also completes normally without restoring authority. +- AG-UI Approval State capacity is enforced independently for each trusted application scope. Abandoned pending + authority expires after its configured window, and indeterminate execution records remain non-retryable until + their separate safety window permits reclamation. Reclamation never recreates approval authority. - A server-issued approval request must not be replayed inline during service-side continuation. - History providers may retain approval control contents in their backing store for audit, but base history replay filters them before later model calls. @@ -454,7 +458,8 @@ that manually replay messages own the equivalent rule: do not resend an approval | Provider-injected approval tool | A tool added during `before_run` defers to in-run resolution, executes once, and emits one result. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_agent_approval_deferred_provider_tool_executes` | | AG-UI provider boundary | Completed local approval controls from AG-UI request and snapshot replay are absent from raw chat-client input while deferred and hosted approvals keep their respective in-run/provider paths. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_does_not_forward_resolved_local_approval_control_to_chat_client`, `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_agent_approval_deferred_provider_tool_executes`, `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_canonical_resume_preserves_hosted_approval_for_provider`, `packages/ag-ui/tests/ag_ui/test_run.py::test_filter_local_approval_responses_for_provider_removes_duplicate_completed_controls`, `packages/ag-ui/tests/ag_ui/test_run.py::test_filter_local_approval_responses_for_provider_pairs_reused_call_ids_by_occurrence`, `packages/ag-ui/tests/ag_ui/test_run.py::test_canonical_hosted_approval_resume_rejects_edited_arguments_without_mutating_pending` | | AG-UI standard approval payload | Agent and workflow tool approvals emit canonical `tool_call` interrupts. `approved` plus full-replacement `editedArgs` executes once and replays idempotently, while legacy `accepted` plus direct partial edits remains supported. Hosted approvals remain decision-only. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_agent_approval_resume_entry_applies_standard_full_replacement_edited_args`, `test_endpoint_agent_approval_replayed_standard_edited_resume_is_idempotent`, `test_endpoint_agent_approval_resume_entry_applies_edited_arguments`, `test_workflow_endpoint_emits_canonical_tool_approval_interrupt`, `test_workflow_endpoint_accepts_canonical_tool_approval_resume`, `test_workflow_endpoint_applies_canonical_approval_edited_args`, `test_workflow_endpoint_accepts_legacy_partial_approval_edits`, `test_workflow_endpoint_hosted_approval_rejects_argument_edits` | -| AG-UI cancellation | A cancelled interrupt executes zero times and completes normally; resolved siblings in the same complete resume still execute once. Workflow cancellation clears both runner correlation and the owning agent executor's pending request so later approvals remain resumable. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_agent_approval_cancelled_resume_entry_completes_without_execution`, `test_endpoint_agent_approval_mixed_cancelled_and_resolved_resume_executes_resolved_tool`, `test_endpoint_workflow_request_info_cancelled_resume_completes_normally`, `test_workflow_endpoint_cancelled_agent_approval_does_not_block_next_approval` | +| AG-UI cancellation | A cancelled interrupt executes zero times and completes normally, including an identical retry during retained cancellation state; resolved siblings in the same complete resume still execute once. Workflow cancellation clears both runner correlation and the owning agent executor's pending request so later approvals remain resumable. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_agent_approval_cancelled_resume_entry_completes_without_execution`, `test_endpoint_agent_approval_replayed_cancellation_completes_idempotently`, `test_endpoint_agent_approval_mixed_cancelled_and_resolved_resume_executes_resolved_tool`, `test_endpoint_workflow_request_info_cancelled_resume_completes_normally`, `test_workflow_endpoint_cancelled_agent_approval_does_not_block_next_approval` | +| AG-UI approval retention and capacity | Pending authority expires automatically, indeterminate outcomes remain non-retryable until their safety window permits reclamation, and one trusted scope cannot consume another scope's occurrence quota. | `packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py::test_abandoned_pending_occurrence_expires_and_releases_capacity`, `test_indeterminate_occurrence_is_reclaimed_after_its_safety_window`, `test_capacity_is_enforced_per_trusted_scope` | | AG-UI local executor unavailable on resume | A claimed local occurrence whose executor disappeared releases its unstarted claim, reports temporary unavailability, and remains safely retryable. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_agent_approval_resume_remains_retryable_when_local_tool_is_temporarily_unavailable` | | AG-UI forwarded execution interruption | A provider failure, cancellation, or stream close after forwarding an approval recovers the open occurrence as indeterminate when no idempotency key proves retry safety. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_hosted_approval_becomes_indeterminate_when_provider_stream_fails` | diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index 9114dccbcd..a5d461a932 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py @@ -855,6 +855,7 @@ def _register_server_generated_approval_response( *, lifecycle: ApprovalLifecycle, has_deferred_owner: bool, + approval_scope: str | None = None, ) -> AuthorizedExecution | None: """Register a server-owned approval response so normal validation can consume it.""" if response.function_call is None or not response.function_call.name: @@ -870,6 +871,7 @@ def _register_server_generated_approval_response( arguments = canonical_function_arguments(response.function_call) or "{}" lifecycle.register( owner=execution_owner, + scope=approval_scope, thread_ids=[thread_id], interrupt_id=str(response_id), call_id=str(response.function_call.call_id or response_id), @@ -911,6 +913,7 @@ def _pop_collected_tool_approval_response_messages( *, lifecycle: ApprovalLifecycle, authorized_executions: dict[ApprovalOccurrenceIdentity, AuthorizedExecution] | None = None, + approval_scope: str | None = None, ) -> list[Message]: """Pop server-collected auto-approved responses into provider-visible messages.""" raw_state = session.state.get(_TOOL_APPROVAL_STATE_KEY) @@ -933,6 +936,7 @@ def _pop_collected_tool_approval_response_messages( tools, lifecycle=lifecycle, has_deferred_owner=True, + approval_scope=approval_scope, ) if intent is not None and authorized_executions is not None: authorized_executions[intent.identity] = intent @@ -1190,6 +1194,7 @@ def _canonical_approval_resume_messages( authorized_executions: dict[ApprovalOccurrenceIdentity, AuthorizedExecution] | None = None, retained_results: list[Content] | None = None, snapshot_reconciliations: list[ApprovalSnapshotReconciliation] | None = None, + approval_scope: str | None = None, ) -> tuple[list[dict[str, Any]], set[str], set[str], RunErrorEvent | None]: """Translate canonical ResumeEntry approvals into existing approval response messages.""" expected_ids = set(expected_interrupt_ids or set()) @@ -1202,6 +1207,30 @@ def _canonical_approval_resume_messages( if _resume_payload_has_approval_decision(resume_payload): normalized_interrupts = _normalize_resume_interrupts(resume_payload) interrupt_id = normalized_interrupts[0]["id"] if normalized_interrupts else "unknown" + cancelled_retry_ids = [ + str(interrupt["id"]) for interrupt in normalized_interrupts if interrupt.get("status") == "cancelled" + ] + if len(cancelled_retry_ids) == len(normalized_interrupts) and cancelled_retry_ids: + try: + reconciliations = lifecycle.cancel_batch( + thread_id=thread_id, + interrupt_ids=cancelled_retry_ids, + ) + except KeyError: + pass + except ValueError as exc: + return ( + [], + handled_ids, + cancelled_ids, + RunErrorEvent(message=str(exc), code="APPROVAL_RESUME_INVALID"), + ) + else: + if snapshot_reconciliations is not None: + snapshot_reconciliations.extend(reconciliations) + handled_ids.update(cancelled_retry_ids) + cancelled_ids.update(cancelled_retry_ids) + return [], handled_ids, cancelled_ids, None if retained_results is not None: decisions: list[ResumeDecision] = [] for interrupt in normalized_interrupts: @@ -1385,6 +1414,7 @@ def _canonical_approval_resume_messages( ) lifecycle.register( owner=execution_owner, + scope=approval_scope, thread_ids=[thread_id], interrupt_id=sibling_interrupt_id, call_id=sibling_call_id, @@ -2286,6 +2316,7 @@ async def run_agent_stream( authorized_executions=authorized_executions, retained_results=retained_approval_results, snapshot_reconciliations=approval_snapshot_reconciliations, + approval_scope=approval_scope, ) ) if resume_error is not None: @@ -2411,6 +2442,7 @@ async def run_agent_stream( tools_for_execution, lifecycle=approval_state_store.lifecycle, authorized_executions=authorized_executions, + approval_scope=approval_scope, ) ) execution_tool_map = _get_tool_map(tools_for_execution) if tools_for_execution else {} @@ -2605,6 +2637,7 @@ async def run_agent_stream( } approval_state_store.register( owner=execution_owner, + scope=approval_scope, server_label=server_label, **registration_kwargs, ) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py b/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py index c42400bda6..eb74310a21 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py @@ -179,11 +179,13 @@ class ApprovalOccurrence: name: str arguments: str owner: ApprovalExecutionOwner + scope: str | None = None aliases: tuple[str, ...] = () already_approved_requests: tuple[dict[str, Any], ...] = () server_label: str | None = None idempotency_key: str | None = None status: ApprovalStatus = ApprovalStatus.PENDING + pending_since: float = 0 replayable_results: list[ReplayableToolResult] = field(default_factory=list) decision: ResumeDecision | None = None outcome: ApprovalOutcome | None = None @@ -238,14 +240,22 @@ def __init__( self, *, max_entries: int = 10_000, + pending_retention_seconds: float = 86_400, + indeterminate_retention_seconds: float = 604_800, terminal_retention_seconds: float = 900, clock: Callable[[], float] = monotonic, ) -> None: if max_entries < 1: raise ValueError("max_entries must be greater than 0.") + if pending_retention_seconds <= 0: + raise ValueError("pending_retention_seconds must be greater than 0.") + if indeterminate_retention_seconds <= 0: + raise ValueError("indeterminate_retention_seconds must be greater than 0.") if terminal_retention_seconds <= 0: raise ValueError("terminal_retention_seconds must be greater than 0.") self._max_entries = max_entries + self._pending_retention_seconds = pending_retention_seconds + self._indeterminate_retention_seconds = indeterminate_retention_seconds self._terminal_retention_seconds = terminal_retention_seconds self._clock = clock self._index_lock = RLock() @@ -258,6 +268,7 @@ def register( self, *, owner: ApprovalExecutionOwner, + scope: str | None = None, thread_ids: list[str] | None = None, thread_id: str | None = None, interrupt_id: str, @@ -280,6 +291,7 @@ def register( name=name, arguments=arguments, owner=owner, + scope=scope, aliases=aliases, already_approved_requests=already_approved_requests, server_label=server_label, @@ -296,12 +308,15 @@ def _register_aliases( name: str, arguments: str, owner: ApprovalExecutionOwner, + scope: str | None = None, aliases: list[str] | None = None, already_approved_requests: list[dict[str, Any]] | None = None, server_label: str | None = None, idempotency_key: str | None = None, ) -> ApprovalOccurrence: self._purge_expired_terminal() + if scope == "": + raise ValueError("An approval scope cannot be empty.") if idempotency_key == "": raise ValueError("An execution idempotency key cannot be empty.") unique_thread_ids = tuple(dict.fromkeys(thread_ids)) @@ -326,6 +341,7 @@ def _register_aliases( or occurrence.name != name or occurrence.arguments != arguments or occurrence.owner is not owner + or occurrence.scope != scope or occurrence.idempotency_key != idempotency_key or occurrence.already_approved_requests != tuple(already_approved_requests or ()) or occurrence.server_label != server_label @@ -338,7 +354,7 @@ def _register_aliases( self._pending_by_interrupt[(thread_id, alias)] = occurrence.identity return occurrence - if len(self._occurrences) >= self._max_entries: + if sum(occurrence.scope == scope for occurrence in self._occurrences.values()) >= self._max_entries: self._emit_event("capacity_failure") raise ApprovalCapacityError("Approval state capacity is exhausted by protected occurrences.") identity = ApprovalOccurrenceIdentity( @@ -353,10 +369,12 @@ def _register_aliases( name=name, arguments=arguments, owner=owner, + scope=scope, aliases=occurrence_aliases, already_approved_requests=tuple(already_approved_requests or ()), server_label=server_label, idempotency_key=idempotency_key, + pending_since=self._clock(), ) self._occurrences[identity] = occurrence for thread_id in unique_thread_ids: @@ -675,7 +693,7 @@ def expire_batch( return tuple(self._snapshot_reconciliation(occurrence) for occurrence in occurrences) def _remove_pending_aliases(self, occurrence: ApprovalOccurrence) -> None: - if occurrence.status.is_purgeable: + if occurrence.status.is_terminal: occurrence.terminal_at = self._clock() with self._index_lock: for thread_id in occurrence.thread_ids: @@ -685,11 +703,31 @@ def _remove_pending_aliases(self, occurrence: ApprovalOccurrence) -> None: def _purge_expired_terminal(self) -> None: with self._index_lock: - cutoff = self._clock() - self._terminal_retention_seconds + now = self._clock() + pending_cutoff = now - self._pending_retention_seconds + abandoned = [ + occurrence + for occurrence in self._occurrences.values() + if occurrence.status is ApprovalStatus.PENDING and occurrence.pending_since <= pending_cutoff + ] + for occurrence in abandoned: + occurrence.status = ApprovalStatus.EXPIRED + self._remove_pending_aliases(occurrence) + occurrence.terminal_at = occurrence.pending_since + self._pending_retention_seconds + self._emit_event("expiration", occurrence) + cutoff = now - self._terminal_retention_seconds + indeterminate_cutoff = now - self._indeterminate_retention_seconds expired = [ occurrence for occurrence in self._occurrences.values() - if occurrence.terminal_at is not None and occurrence.terminal_at <= cutoff + if occurrence.terminal_at is not None + and ( + (occurrence.status.is_purgeable and occurrence.terminal_at <= cutoff) + or ( + occurrence.status is ApprovalStatus.INDETERMINATE + and occurrence.terminal_at <= indeterminate_cutoff + ) + ) ] for occurrence in expired: self._emit_event("retention_purge", occurrence) @@ -765,6 +803,7 @@ def release_claim(self, intent: AuthorizedExecution, *, policy: ClaimRecoveryPol if occurrence.status is not ApprovalStatus.CLAIMED: raise ValueError(f"Approval occurrence is not claimed: {occurrence.status}.") occurrence.status = ApprovalStatus.PENDING + occurrence.pending_since = self._clock() @_serialized_by_occurrence def mark_indeterminate( @@ -890,6 +929,7 @@ def defer(self, intent: AuthorizedExecution, results: list[Content]) -> Approval if occurrence.status is not ApprovalStatus.EXECUTING: raise ValueError(f"Approval occurrence is not executing: {occurrence.status}.") occurrence.status = ApprovalStatus.PENDING + occurrence.pending_since = self._clock() return ApprovalOutcome( identity=occurrence.identity, replayable_results=(), diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py b/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py index 5bae76017b..6ed72c9c04 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py @@ -14,6 +14,8 @@ """Application-defined scope for server-side AG-UI Approval State.""" DEFAULT_MAX_APPROVAL_STATES = 10_000 +DEFAULT_PENDING_RETENTION_SECONDS = 86_400 +DEFAULT_INDETERMINATE_RETENTION_SECONDS = 604_800 DEFAULT_TERMINAL_RETENTION_SECONDS = 900 _APPROVAL_SCOPE_INPUT_KEY = "__ag_ui_approval_scope" _APPROVAL_THREAD_SEPARATOR = "\x1f" @@ -46,12 +48,16 @@ def __init__( self, *, max_entries: int = DEFAULT_MAX_APPROVAL_STATES, + pending_retention_seconds: float = DEFAULT_PENDING_RETENTION_SECONDS, + indeterminate_retention_seconds: float = DEFAULT_INDETERMINATE_RETENTION_SECONDS, terminal_retention_seconds: float = DEFAULT_TERMINAL_RETENTION_SECONDS, ) -> None: """Initialize the process-local Approval State store. Keyword Args: max_entries: Maximum approval occurrences or middleware state entries to retain. + pending_retention_seconds: Maximum time to retain abandoned pending approval authority. + indeterminate_retention_seconds: Safety window for uncertain execution records before reclamation. terminal_retention_seconds: Process-local duplicate-execution protection window. Raises: @@ -64,6 +70,8 @@ def __init__( self._tool_approval_states: dict[str, dict[str, Any]] = {} self.lifecycle = ApprovalLifecycle( max_entries=max_entries, + pending_retention_seconds=pending_retention_seconds, + indeterminate_retention_seconds=indeterminate_retention_seconds, terminal_retention_seconds=terminal_retention_seconds, ) @@ -76,6 +84,7 @@ def register( request_id: str, interrupt_id: str, owner: ApprovalExecutionOwner, + scope: ApprovalScope | None = None, already_approved_requests: list[dict[str, Any]] | None = None, server_label: str | None = None, ) -> None: @@ -83,6 +92,7 @@ def register( unique_thread_ids = list(dict.fromkeys(thread_ids)) self.lifecycle.register( owner=owner, + scope=scope, thread_ids=unique_thread_ids, interrupt_id=interrupt_id, call_id=interrupt_id, diff --git a/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py b/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py index 8aa414dd33..feeee69420 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py +++ b/python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py @@ -173,6 +173,105 @@ def test_active_occurrence_is_not_evicted_when_capacity_is_exhausted() -> None: assert lifecycle.get(occurrence.identity).status is ApprovalStatus.PENDING +def test_abandoned_pending_occurrence_expires_and_releases_capacity() -> None: + """Abandoned pending authority is reclaimed after its configured safety windows.""" + now = 100.0 + lifecycle = ApprovalLifecycle( + max_entries=1, + pending_retention_seconds=20, + terminal_retention_seconds=10, + clock=lambda: now, + ) + abandoned = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, + thread_id="thread-abandoned", + interrupt_id="approval-abandoned", + call_id="call-abandoned", + name="write_record", + arguments="{}", + ) + + now = 131.0 + replacement = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, + thread_id="thread-replacement", + interrupt_id="approval-replacement", + call_id="call-replacement", + name="write_record", + arguments="{}", + ) + + assert replacement.status is ApprovalStatus.PENDING + with pytest.raises(KeyError): + lifecycle.get(abandoned.identity) + with pytest.raises(KeyError): + lifecycle.claim_batch( + thread_id="thread-abandoned", + decisions=[ResumeDecision(interrupt_id="approval-abandoned", accepted=True, arguments="{}")], + ) + + +def test_safe_claim_release_restarts_pending_retention_window() -> None: + """A safely released claim receives a fresh pending decision window.""" + now = 100.0 + lifecycle = ApprovalLifecycle(max_entries=1, pending_retention_seconds=20, clock=lambda: now) + occurrence = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, + thread_id="thread-1", + interrupt_id="approval-1", + call_id="call-1", + name="write_record", + arguments="{}", + ) + intent = lifecycle.claim( + thread_id="thread-1", + decision=ResumeDecision(interrupt_id="approval-1", accepted=True, arguments="{}"), + ) + + now = 119.0 + lifecycle.release_claim(intent, policy=ClaimRecoveryPolicy.SAFE_TO_RETRY) + now = 121.0 + + assert lifecycle.pending_occurrence(thread_id="thread-1", interrupt_id="approval-1") is occurrence + + +def test_capacity_is_enforced_per_trusted_scope() -> None: + """One trusted application scope cannot exhaust another scope's approval quota.""" + lifecycle = ApprovalLifecycle(max_entries=1) + lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, + scope="tenant-a", + thread_id="tenant-a\x1fthread-1", + interrupt_id="approval-a-1", + call_id="call-a-1", + name="write_record", + arguments="{}", + ) + + with pytest.raises(ApprovalCapacityError): + lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, + scope="tenant-a", + thread_id="tenant-a\x1fthread-2", + interrupt_id="approval-a-2", + call_id="call-a-2", + name="write_record", + arguments="{}", + ) + + tenant_b = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, + scope="tenant-b", + thread_id="tenant-b\x1fthread-1", + interrupt_id="approval-b-1", + call_id="call-b-1", + name="write_record", + arguments="{}", + ) + + assert tenant_b.status is ApprovalStatus.PENDING + + async def test_terminal_outcome_expires_only_after_configured_retention_window() -> None: """Duplicate execution protection lasts for the configured terminal retention window.""" now = 100.0 @@ -245,6 +344,64 @@ def test_indeterminate_occurrence_remains_protected_after_terminal_retention_win assert lifecycle.get(occurrence.identity).status is ApprovalStatus.INDETERMINATE +def test_indeterminate_occurrence_is_reclaimed_after_its_safety_window() -> None: + """An uncertain execution is eventually reclaimed without restoring approval authority.""" + now = 100.0 + lifecycle = ApprovalLifecycle( + max_entries=1, + indeterminate_retention_seconds=30, + terminal_retention_seconds=10, + clock=lambda: now, + ) + uncertain = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, + scope="tenant-a", + thread_id="tenant-a\x1fthread-1", + interrupt_id="approval-uncertain", + call_id="call-uncertain", + name="write_record", + arguments="{}", + ) + intent = lifecycle.claim( + thread_id="tenant-a\x1fthread-1", + decision=ResumeDecision(interrupt_id="approval-uncertain", accepted=True, arguments="{}"), + ) + lifecycle.begin_execution(intent, owner=ApprovalExecutionOwner.LOCAL) + lifecycle.recover_execution(intent, owner=ApprovalExecutionOwner.LOCAL) + + now = 129.0 + with pytest.raises(ApprovalCapacityError): + lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, + scope="tenant-a", + thread_id="tenant-a\x1fthread-2", + interrupt_id="approval-blocked", + call_id="call-blocked", + name="write_record", + arguments="{}", + ) + + now = 131.0 + replacement = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, + scope="tenant-a", + thread_id="tenant-a\x1fthread-2", + interrupt_id="approval-replacement", + call_id="call-replacement", + name="write_record", + arguments="{}", + ) + + assert replacement.status is ApprovalStatus.PENDING + with pytest.raises(KeyError): + lifecycle.get(uncertain.identity) + with pytest.raises(KeyError): + lifecycle.claim_batch( + thread_id="tenant-a\x1fthread-1", + decisions=[ResumeDecision(interrupt_id="approval-uncertain", accepted=True, arguments="{}")], + ) + + async def test_transition_telemetry_covers_lifecycle_without_sensitive_payloads( caplog: pytest.LogCaptureFixture, ) -> None: diff --git a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py index 67d834162a..63e502effc 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -3124,6 +3124,26 @@ async def test_endpoint_agent_approval_cancelled_resume_entry_completes_without_ assert not [event for event in events if event.get("type") == "TOOL_CALL_RESULT"] +async def test_endpoint_agent_approval_replayed_cancellation_completes_idempotently() -> None: + """Retrying a cancellation after a lost response completes normally without execution.""" + client, _, executed_cities = _build_weather_approval_endpoint() + resume = [{"interruptId": "call_get_weather", "status": "cancelled"}] + + first_response = client.post( + "/approval", + json={"runId": "run-cancel-first", "threadId": "thread-weather", "messages": [], "resume": resume}, + ) + retry_response = client.post( + "/approval", + json={"runId": "run-cancel-retry", "threadId": "thread-weather", "messages": [], "resume": resume}, + ) + + assert first_response.status_code == 200 + assert retry_response.status_code == 200 + assert [event.get("type") for event in _decode_sse_events(retry_response)] == ["RUN_STARTED", "RUN_FINISHED"] + assert executed_cities == [] + + async def test_endpoint_agent_approval_unknown_resume_entry_emits_run_error(): """A canonical approval resume for an unknown pending interrupt should fail safely.""" client, _, executed_cities = _build_weather_approval_endpoint()