Python: Harden AG-UI approval lifecycle and resume semantics - #7594
Python: Harden AG-UI approval lifecycle and resume semantics#7594Evan Mattson (moonbox3) wants to merge 11 commits into
Conversation
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.
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.
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.
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.
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.
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.
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.
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.
There was a problem hiding this comment.
Pull request overview
Centralizes AG-UI approval authority in an occurrence-aware lifecycle while updating resume, cancellation, replay, and snapshot behavior.
Changes:
- Adds lifecycle ownership, retention, recovery, and replay deduplication.
- Supports canonical
approved/editedArgspayloads and normal cancellation. - Expands regression coverage and documentation.
Reviewed changes
Copilot reviewed 16 out of 16 changed files in this pull request and generated 10 comments.
Show a summary per file
| File | Description |
|---|---|
python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py |
Adds approval lifecycle state machine. |
python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py |
Integrates lifecycle-backed state storage. |
python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py |
Routes approval execution and replay through lifecycle. |
python/packages/ag-ui/agent_framework_ag_ui/_agent.py |
Removes the legacy approval registry. |
python/packages/ag-ui/agent_framework_ag_ui/_run_common.py |
Extends approval response schemas. |
python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py |
Makes workflow cancellation terminal and normal. |
python/packages/ag-ui/agent_framework_ag_ui/_snapshot_session.py |
Reconciles in-memory snapshot interrupts. |
python/packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py |
Adds lifecycle behavior and concurrency tests. |
python/packages/ag-ui/tests/ag_ui/test_approval_state.py |
Tests bounded state behavior. |
python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py |
Refactors result projection tests. |
python/packages/ag-ui/tests/ag_ui/test_endpoint.py |
Expands endpoint lifecycle regressions. |
python/packages/ag-ui/tests/ag_ui/test_run.py |
Updates runner tests for lifecycle state. |
python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py |
Updates wrapper approval assertions. |
python/packages/ag-ui/README.md |
Documents canonical approval payloads. |
python/packages/ag-ui/AGENTS.md |
Records lifecycle ownership conventions. |
docs/specs/004-python-function-calling-loop.md |
Updates function-loop approval semantics. |
Suppressed comments (1)
python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py:405
- This second iteration has the same unlocked-read race: concurrent registration/purge can mutate
_occurrenceswhile the set comprehension runs, causing a runtime error or an inconsistent pending-ID snapshot. Keep the purge and iteration under_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
}
💡 Add a code-review agent skill for context-aware, tailored reviews. Learn more in the docs.
| if snapshot_reconciliations is not None: | ||
| snapshot_reconciliations.extend(intents.snapshot_reconciliations) | ||
| for intent in intents: | ||
| authorized_executions[intent.identity.call_id] = intent |
| for intent in local_intents: | ||
| approval_state_store.lifecycle.release_claim(intent, policy=ClaimRecoveryPolicy.SAFE_TO_RETRY) |
| 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, | ||
| ), | ||
| ) |
| reconciliations = lifecycle.cancel_batch( | ||
| thread_id=thread_id, | ||
| interrupt_ids=lifecycle_cancelled_ids, | ||
| ) |
| | 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` | |
| for outcome in batch.retained_outcomes: | ||
| retained_results.extend(result.content for result in outcome.replayable_results) |
| 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: |
| self._purge_expired_terminal() | ||
| return tuple(occurrence for occurrence in self._occurrences.values() if thread_id in occurrence.thread_ids) |
| 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) |
| 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) |
There was a problem hiding this comment.
Agent Framework Review — Iteration 1
Completed passes: 5 | Result: Findings reported
Scope: full PR (9 commit(s)): 06daf9586c71, 26ef56d71571, f24f57a45253, 5c21854ec475, de6fce99d656, d7fcd066ca41, 42b3d05b47ef, 31c483e65fc6, 99feeaf529c0
Review passes
- Correctness (
gpt-5.6-sol) — The unavailable-local-tool recovery path can strand sibling hosted or deferred approvals in the claimed state, breaking safe retry and batch continuity. Two medium-severity correctness issues were found: reserved approval fields can be overwritten by colliding tool argument names, and lifecycle dictionary iteration is vulnerable to concurrent mutation. - Security Reliability (
claude-opus-4.8) — No issues found in this pass. - Test Coverage (
gpt-5.6-sol) — Two regressions remain uncovered: duplicate call IDs within a single approval batch can cause incorrect settlement, and tool arguments namedapprovedcan corrupt the approval response schema. - Failure Modes (
claude-opus-4.8) — No issues found in this pass. - Design Approach (
claude-opus-4.8) — No issues found in this pass.
| 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: |
There was a problem hiding this comment.
When a local tool is unavailable, release every unstarted authorized intent before returning the temporary-unavailability error. Releasing only local_intents leaves sibling hosted or deferred intents permanently CLAIMED.
| if unavailable_local_intents: | |
| await self._release_claimed_intents(authorized_intents) |
| 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( |
There was a problem hiding this comment.
This test reuses the call ID only across separate threads and claims one occurrence, so it misses the promised same-batch case. The runner indexes AuthorizedExecution values by raw call_id; two distinct interrupts sharing a call ID overwrite each other, causing one intent to settle the wrong occurrence and the other to fail. Add an endpoint regression with unique interrupt IDs but the same call ID, asserting both execute and settle exactly once.
| tenant_a = lifecycle.register( | |
| # Add a regression test with two unique interrupt IDs sharing one call ID, | |
| # and assert both execute and settle exactly once. |
Motivation & Context
AG-UI approval-gated tool calls need occurrence-scoped authority across interruption, resume, execution, retry, and snapshot restoration. The previous runner-owned paths could leave approvals claimed or executing when tools disappeared or streams ended unexpectedly, confuse repeated call identifiers, and diverge from current AG-UI resume semantics.
This change centralizes approval lifecycle ownership so each pending call has one authoritative transition path while retaining the existing MAF endpoint and response compatibility.
Description & Review Guide
approvedwith full-replacementeditedArgs.acceptedand direct partial edits, keeps hosted approvals decision-only, and clears workflow-agent pending state when an approval is cancelled.acceptedalias and direct partial argument edits for compatibility; no exported API surface changes.Related Issue
N/A — approval lifecycle hardening identified through local adversarial review.
Contribution Checklist
breaking changelabel (or add "[BREAKING]" to the title prefix, before or after any language prefix) — a workflow keeps the label and title prefix in sync automatically.