Merge pull request #158 from aliyun/codex/fix-step1-handle-off

fix: harden A2A recovery e2e flows

Author: guima-why

scripts/a2a/e2e/fixtures/text-images/manifest.json

@@ -36,9 +36,9 @@
   },
   "rollback-interrupt": {
     "filename": "rollback-interrupt.png",
-    "text": "回退到 intent_parsing，选择一个已有vpc，创建一个安全组",
+    "text": "我改需求了：使用已有 VPC 创建一个安全组，不创建 VSwitch。请基于这个新需求重新规划。",
     "mediaType": "image/png",
-    "byteSize": 20967,
-    "sha256": "1dfa25bba58757704b27a7ee8f44a42f4e69045730bba5320986194c873a5937"
+    "byteSize": 33436,
+    "sha256": "99399b761853dd9a9188349549933de6cdd074051697ecbb23cd687769222ca3"
   }
 }

scripts/a2a/e2e/fixtures/text-images/rollback-interrupt.png

@@ -39,7 +39,11 @@
 from iac_code.a2a.pipeline_executor import IacCodeA2APipelineExecutor, recoverable_task_id_from_sidecar
 from iac_code.a2a.pipeline_journal import A2APipelineJournal
 from iac_code.a2a.pipeline_paths import existing_a2a_pipeline_dir_for_session
-from iac_code.a2a.pipeline_snapshot import A2APipelineSnapshotStore, reduce_pipeline_events
+from iac_code.a2a.pipeline_snapshot import (
+    A2APipelineSnapshotStore,
+    reduce_pipeline_events,
+    snapshot_needs_backup_commit_repair,
+)
 from iac_code.a2a.pipeline_stream import BACKUP_COMMITTED_EVENT_TYPE, PipelineA2AEventPublisher
 from iac_code.a2a.runtime_overrides import (
     a2a_request_context,
@@ -89,6 +93,7 @@
 
 logger = logging.getLogger(__name__)
 _CONTEXT_LOCK_ACQUIRE_TIMEOUT_SECONDS = 1
+_CANCEL_ACTIVE_TASK_DRAIN_TIMEOUT_SECONDS = 30
 _ERROR_TEXT_MAX_CHARS = 1000
 _DEFERRED_CLEANUP_PROMPTS_FILENAME = "cleanup-deferred-prompts.json"
 _A2A_SAFE_MODE_ENV = "IAC_CODE_A2A_SAFE_MODE"
@@ -434,7 +439,12 @@ def _a2a_pipeline_state_for_session(
         (_a2a_pipeline_sequence_number(event.get("sequence")) for event in journal_events if isinstance(event, dict)),
         default=0,
     )
-    if journal_events and (not isinstance(snapshot, dict) or journal_sequence != snapshot_sequence):
+    needs_backup_commit_repair = (
+        isinstance(snapshot, dict) and snapshot_needs_backup_commit_repair(snapshot, journal_events)
+    )
+    if journal_events and (
+        not isinstance(snapshot, dict) or journal_sequence != snapshot_sequence or needs_backup_commit_repair
+    ):
         snapshot = reduce_pipeline_events(journal_events)
         if not isinstance(snapshot, dict):
             return None
@@ -1336,7 +1346,10 @@ def runtime_factory(session_id: str) -> Any:
     async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None:
         task_id = context.task_id
         context_id = context.context_id or "unknown"
-        if task_id and await self._task_store.cancel_task(task_id):
+        if task_id and await self._task_store.cancel_task_and_wait(
+            task_id,
+            timeout=_CANCEL_ACTIVE_TASK_DRAIN_TIMEOUT_SECONDS,
+        ):
             await self._publish_status(
                 event_queue,
                 task_id=task_id,

scripts/a2a/e2e/run_recovery_scenarios.py

@@ -90,6 +90,13 @@
 _PENDING_QUESTION_NOT_ROUTED = "not_routed"
 _PENDING_QUESTION_ANSWERED = "answered"
 _PENDING_QUESTION_STALE_FINISHED = "stale_finished"
+_ACTIVE_INTERRUPT_TERMINAL_WAIT_TIMEOUT_SECONDS = 30.0
+
+
+def _new_set_asyncio_event() -> asyncio.Event:
+    event = asyncio.Event()
+    event.set()
+    return event
 
 
 class WaitingInputCancelResult(str, Enum):
@@ -143,6 +150,7 @@ class A2APipelineRuntime:
     restart_after_interrupt: bool = False
     pause_after_interrupt: bool = False
     restart_requested: asyncio.Event = field(default_factory=asyncio.Event)
+    interrupt_settled: asyncio.Event = field(default_factory=_new_set_asyncio_event)
 
 
 @dataclass(frozen=True)
@@ -574,7 +582,7 @@ def fresh_pipeline_factory() -> Any:
                     ctx.active_task_id = task.task_id
                 elif task.active_task is owner_task:
                     task.active_task = None
-                    if ctx.active_task_id == task.task_id and getattr(pipeline, "sidecar_status", None) != "running":
+                    if ctx.active_task_id == task.task_id:
                         ctx.active_task_id = None
                     if replacement_owner is owner_task and hasattr(ctx.runtime, "active_owner_task"):
                         ctx.runtime.active_owner_task = None
@@ -725,6 +733,8 @@ async def _route_active_pipeline_interrupt(
         paused = False
         verdict: Any | None = None
         interrupt_received_published = False
+        interrupt_settled = _interrupt_settled_event(runtime)
+        interrupt_settled.clear()
         try:
             publish_interrupt_received = getattr(publisher, "publish_interrupt_received", None)
             if callable(publish_interrupt_received):
@@ -743,9 +753,8 @@ async def _route_active_pipeline_interrupt(
                 apply_hard_interrupt = getattr(pipeline, "apply_hard_interrupt", None)
                 if callable(apply_hard_interrupt):
                     parameters = inspect.signature(apply_hard_interrupt).parameters
-                    if pipeline_input.has_images and (
-                        "source_input" in parameters
-                        or any(parameter.kind == inspect.Parameter.VAR_KEYWORD for parameter in parameters.values())
+                    if "source_input" in parameters or any(
+                        parameter.kind == inspect.Parameter.VAR_KEYWORD for parameter in parameters.values()
                     ):
                         applied = apply_hard_interrupt(verdict, source_input=runner_input)
                     else:
@@ -818,6 +827,7 @@ async def _route_active_pipeline_interrupt(
                 await self._complete_backup_blocked_transition(task=task, ctx=ctx)
             return True
         finally:
+            interrupt_settled.set()
             if paused and not bool(getattr(verdict, "paused", False)):
                 resume_agent_loops = getattr(pipeline, "resume_agent_loops", None)
                 if callable(resume_agent_loops):
@@ -1148,6 +1158,22 @@ async def _consume_stream_until_restart(
                 finally:
                     next_task = None
 
+                interrupt_action = await _consume_active_interrupt_action_before_terminal(runtime, event)
+                if interrupt_action == "restart":
+                    close_stream_on_exit = True
+                    return _StreamConsumeResult(
+                        had_events=had_events,
+                        restart_requested=True,
+                        terminal_handoff_unavailable=terminal_handoff_unavailable,
+                    )
+                if interrupt_action == "pause":
+                    close_stream_on_exit = True
+                    return _StreamConsumeResult(
+                        had_events=had_events,
+                        restart_requested=False,
+                        terminal_handoff_unavailable=terminal_handoff_unavailable,
+                    )
+
                 had_events = True
                 await publish_mcp_warnings(
                     publisher.event_queue,
@@ -3410,6 +3436,72 @@ def _restart_requested_event(runtime: Any) -> asyncio.Event:
     return restart_requested
 
 
+def _interrupt_settled_event(runtime: Any) -> asyncio.Event:
+    interrupt_settled = getattr(runtime, "interrupt_settled", None)
+    if isinstance(interrupt_settled, asyncio.Event):
+        return interrupt_settled
+    interrupt_settled = _new_set_asyncio_event()
+    runtime.interrupt_settled = interrupt_settled
+    return interrupt_settled
+
+
+async def _consume_active_interrupt_action_before_terminal(runtime: Any, event: Any) -> str | None:
+    if not _is_pipeline_terminal_stream_event(event):
+        return None
+
+    action = _consume_requested_interrupt_action(runtime)
+    if action is not None:
+        return action
+
+    interrupt_settled = _interrupt_settled_event(runtime)
+    if interrupt_settled.is_set():
+        return None
+
+    restart_event = _restart_requested_event(runtime)
+    settled_task = asyncio.create_task(interrupt_settled.wait())
+    restart_task = asyncio.create_task(restart_event.wait())
+    try:
+        done, _pending = await asyncio.wait(
+            {settled_task, restart_task},
+            timeout=_ACTIVE_INTERRUPT_TERMINAL_WAIT_TIMEOUT_SECONDS,
+            return_when=asyncio.FIRST_COMPLETED,
+        )
+    finally:
+        await _cancel_task_safely(settled_task)
+        await _cancel_task_safely(restart_task)
+
+    action = _consume_requested_interrupt_action(runtime)
+    if action is not None:
+        return action
+
+    if not done and not interrupt_settled.is_set():
+        logger.warning(
+            "Timed out waiting for active A2A pipeline interrupt before publishing terminal pipeline event"
+        )
+    return None
+
+
+def _consume_requested_interrupt_action(runtime: Any) -> str | None:
+    restart_event = _restart_requested_event(runtime)
+    if bool(getattr(runtime, "restart_after_interrupt", False)) and restart_event.is_set():
+        restart_event.clear()
+        runtime.restart_after_interrupt = False
+        return "restart"
+    if bool(getattr(runtime, "pause_after_interrupt", False)) and restart_event.is_set():
+        restart_event.clear()
+        runtime.pause_after_interrupt = False
+        return "pause"
+    return None
+
+
+def _is_pipeline_terminal_stream_event(event: Any) -> bool:
+    return isinstance(event, PipelineEvent) and event.type in {
+        PipelineEventType.PIPELINE_COMPLETED,
+        PipelineEventType.PIPELINE_ERROR,
+        PipelineEventType.BACKUP_BLOCKED,
+    }
+
+
 def _is_active_task_record(task: Any, active_task_id: str | None) -> bool:
     return active_task_id is not None and getattr(task, "task_id", None) == active_task_id
 

src/iac_code/a2a/executor.py

@@ -13,6 +13,7 @@
     reduce_pipeline_events,
     sanitize_pipeline_artifact_uris,
     sanitize_pipeline_cleanup_private_fields,
+    snapshot_needs_backup_commit_repair,
 )
 from iac_code.i18n import _
 
@@ -71,6 +72,11 @@ async def get_state(
         else:
             replay_events = _events_for_task(events, task_id=recovery_task_id, context_id=context_id)
 
+        if snapshot_needs_backup_commit_repair(snapshot, replay_events):
+            snapshot = reduce_pipeline_events(replay_events)
+            snapshot_store.save(snapshot)
+            snapshot = snapshot_store.load() or snapshot
+
         if (
             snapshot is not None
             and _snapshot_schema_is_stale(snapshot)

src/iac_code/a2a/pipeline_executor.py

@@ -1122,6 +1122,78 @@ def _backup_ack_authoritative_events(events: list[dict[str, Any]]) -> list[dict[
     return authoritative
 
 
+def snapshot_needs_backup_commit_repair(snapshot: dict[str, Any] | None, events: list[dict[str, Any]]) -> bool:
+    if not isinstance(snapshot, dict):
+        return False
+    pending_publications = _pending_backup_publications_from_snapshot(snapshot)
+    if not pending_publications:
+        return False
+
+    matching_events = [
+        event
+        for event in events
+        if isinstance(event, dict) and _event_matches_snapshot_task_context(event, snapshot)
+    ]
+    committed_by_id: dict[str, dict[str, Any]] = {}
+    committed_by_sequence_type: dict[tuple[int, str], dict[str, Any]] = {}
+    for event in matching_events:
+        if not _is_committed_backup_publication(event):
+            continue
+        event_id = _string_or_none(event.get("eventId"))
+        event_type = _string_or_none(event.get("eventType")) or ""
+        sequence = _sequence_value(event)
+        if event_id:
+            committed_by_id[event_id] = event
+        committed_by_sequence_type[(sequence, event_type)] = event
+
+    if not committed_by_id and not committed_by_sequence_type:
+        return False
+
+    for event in matching_events:
+        if event.get("eventType") != _BACKUP_COMMITTED_EVENT_TYPE:
+            continue
+        data = _dict_or_empty(event.get("data"))
+        committed_event = None
+        committed_event_id = _string_or_none(data.get("committedEventId"))
+        if committed_event_id:
+            committed_event = committed_by_id.get(committed_event_id)
+        if committed_event is None:
+            committed_event_type = _string_or_none(data.get("committedEventType")) or ""
+            committed_event = committed_by_sequence_type.get(
+                (_sequence_number(data.get("committedSequence")), committed_event_type),
+            )
+        if committed_event is None:
+            continue
+        committed_event_type = _string_or_none(committed_event.get("eventType"))
+        committed_sequence = _sequence_value(committed_event)
+        for pending in pending_publications:
+            if (
+                committed_event_type == _string_or_none(pending.get("eventType"))
+                and committed_sequence >= _sequence_value(pending)
+            ):
+                return True
+    return False
+
+
+def _pending_backup_publications_from_snapshot(snapshot: dict[str, Any]) -> list[dict[str, Any]]:
+    publications: list[dict[str, Any]] = []
+    for key in ("pendingTerminal", "pendingNormalHandoff"):
+        publication = snapshot.get(key)
+        if isinstance(publication, dict) and _publication_visibility(publication) == _PENDING_BACKUP_VISIBILITY:
+            publications.append(publication)
+    return publications
+
+
+def _event_matches_snapshot_task_context(event: dict[str, Any], snapshot: dict[str, Any]) -> bool:
+    context_id = _string_or_none(snapshot.get("contextId"))
+    if context_id is not None and event.get("contextId") != context_id:
+        return False
+    task_id = _string_or_none(snapshot.get("taskId"))
+    if task_id is None:
+        return True
+    return event.get("taskId") == task_id or event.get("deliveryTaskId") == task_id
+
+
 def _publication_visibility(event: dict[str, Any]) -> str | None:
     value = _string_or_none(event.get("visibility"))
     if value is not None:
@@ -1685,4 +1757,5 @@ def _utc_now() -> str:
     "SNAPSHOT_SCHEMA_VERSION",
     "reduce_pipeline_events",
     "sanitize_pipeline_cleanup_private_fields",
+    "snapshot_needs_backup_commit_repair",
 ]

src/iac_code/a2a/pipeline_recovery.py

@@ -431,6 +431,28 @@ async def cancel_task(self, task_id: str) -> bool:
             record.active_task.cancel()
             return True
 
+    async def cancel_task_and_wait(self, task_id: str, *, timeout: float | None = None) -> bool:
+        async with self._mutation_lock:
+            record = self._tasks.get(validate_protocol_id(task_id))
+            if record is None or record.active_task is None or record.active_task.done():
+                return False
+            active_task = record.active_task
+            active_task.cancel()
+
+        if active_task is asyncio.current_task():
+            return True
+        try:
+            if timeout is None:
+                await asyncio.shield(active_task)
+            else:
+                await asyncio.wait_for(asyncio.shield(active_task), timeout=timeout)
+        except asyncio.CancelledError:
+            if not active_task.done():
+                raise
+        except TimeoutError:
+            logger.warning("Timed out waiting for canceled A2A task %s to finish", task_id)
+        return True
+
     async def is_task_active(self, task_id: str) -> bool:
         async with self._mutation_lock:
             record = self._tasks.get(validate_protocol_id(task_id))

src/iac_code/a2a/pipeline_snapshot.py

@@ -3087,6 +3087,7 @@ msgid "The user-facing question to ask."
 msgstr "Die dem Benutzer zu stellende Frage."
 
 #: src/iac_code/pipeline/engine/ask_user_question_tool.py
+#: src/iac_code/ui/renderer.py
 msgid "ask_user_question validation failed."
 msgstr "Die Validierung von ask_user_question ist fehlgeschlagen."
 
@@ -3626,6 +3627,11 @@ msgstr ""
 msgid "Step {step_id} completed. Conclusion submitted."
 msgstr "Schritt {step_id} abgeschlossen. Schlussfolgerung übermittelt."
 
+#: src/iac_code/pipeline/engine/pipeline_runner.py
+#, python-brace-format
+msgid "用户反馈：{}"
+msgstr "Benutzerfeedback: {}"
+
 #: src/iac_code/pipeline/engine/pipeline_runner.py
 msgid "Pipeline state persistence failed."
 msgstr "Persistenz des Pipeline-Zustands fehlgeschlagen."

src/iac_code/a2a/task_store.py

@@ -3088,6 +3088,7 @@ msgid "The user-facing question to ask."
 msgstr "La pregunta que se mostrará al usuario."
 
 #: src/iac_code/pipeline/engine/ask_user_question_tool.py
+#: src/iac_code/ui/renderer.py
 msgid "ask_user_question validation failed."
 msgstr "La validación de ask_user_question falló."
 
@@ -3609,6 +3610,11 @@ msgstr ""
 msgid "Step {step_id} completed. Conclusion submitted."
 msgstr "Paso {step_id} completado. Conclusión enviada."
 
+#: src/iac_code/pipeline/engine/pipeline_runner.py
+#, python-brace-format
+msgid "用户反馈：{}"
+msgstr "Comentarios del usuario: {}"
+
 #: src/iac_code/pipeline/engine/pipeline_runner.py
 msgid "Pipeline state persistence failed."
 msgstr "Error al persistir el estado de la pipeline."