From 8dd03235a205e66bc8d6ebb3dc26be10ff50a2e7 Mon Sep 17 00:00:00 2001 From: Sam Schillace Date: Fri, 7 Aug 2026 12:34:34 -0700 Subject: [PATCH] fix(session): persist interrupted child sessions --- .gitignore | 1 + amplifier_app_cli/session_spawner.py | 529 ++++++++++----- tests/test_session_spawner_interruption.py | 708 +++++++++++++++++++++ 3 files changed, 1090 insertions(+), 148 deletions(-) create mode 100644 tests/test_session_spawner_interruption.py diff --git a/.gitignore b/.gitignore index 6802ccb9..0cb734ff 100644 --- a/.gitignore +++ b/.gitignore @@ -71,3 +71,4 @@ next-steps.md # Working folders ai_working/tmp tests/recipes/DECISIONS.md +.next/ diff --git a/amplifier_app_cli/session_spawner.py b/amplifier_app_cli/session_spawner.py index 4d1b0fa5..fdb26d4c 100644 --- a/amplifier_app_cli/session_spawner.py +++ b/amplifier_app_cli/session_spawner.py @@ -1,22 +1,172 @@ +# pyright: reportMissingImports=false """Session spawning for agent delegation. Implements sub-session creation with configuration inheritance and overlays. """ +import asyncio import copy import logging import sys +from collections.abc import Awaitable, Callable from pathlib import Path +from typing import Any, TypeVar, cast from amplifier_core import AmplifierSession -from amplifier_foundation import generate_sub_session_id -from amplifier_foundation import bridge_child_cost -from amplifier_foundation import RUNTIME_SKILL_OVERLAY_CAPABILITY +from amplifier_foundation import ( + RUNTIME_SKILL_OVERLAY_CAPABILITY, + bridge_child_cost, + generate_sub_session_id, +) from .agent_config import merge_configs logger = logging.getLogger(__name__) +_T = TypeVar("_T") + + +class _SubSessionLifecycle: + """Track and reliably release resources acquired by an in-process child.""" + + def __init__(self) -> None: + self.child_session: Any | None = None + self.display_system: Any | None = None + self.nesting_pushed = False + self.unregister_hook: Callable[[], Any] | None = None + self.parent_cancellation: Any | None = None + self.child_cancellation: Any | None = None + self._persist_interruption: Callable[[], Awaitable[None]] | None = None + self._finalization_started = False + self._teardown_started = False + + def push_nesting(self, display_system: Any) -> None: + self.display_system = display_system + if hasattr(display_system, "push_nesting"): + # Mark first so even a partially successful push is balanced. + self.nesting_pushed = True + display_system.push_nesting() + + def register_cancellation( + self, parent_cancellation: Any, child_cancellation: Any + ) -> None: + parent_cancellation.register_child(child_cancellation) + self.parent_cancellation = parent_cancellation + self.child_cancellation = child_cancellation + + def persist_on_interruption(self, callback: Callable[[], Awaitable[None]]) -> None: + """Register persistence to run only when child execution was cancelled.""" + self._persist_interruption = callback + + async def finalize(self, *, preserve_error: bool) -> None: + """Persist an interrupted execution, then release every resource once.""" + if self._finalization_started: + return + self._finalization_started = True + + if self._persist_interruption is not None: + persist_interruption = self._persist_interruption + self._persist_interruption = None + try: + await persist_interruption() + except BaseException: + # Persistence is best effort and must never replace the original + # execution cancellation. + logger.exception("Failed to persist interrupted sub-session") + + await self.teardown(preserve_error=preserve_error) + + async def teardown(self, *, preserve_error: bool) -> None: + """Attempt every acquired teardown action exactly once.""" + if self._teardown_started: + return + self._teardown_started = True + errors: list[BaseException] = [] + + if self.unregister_hook is not None: + unregister_hook = self.unregister_hook + self.unregister_hook = None + try: + unregister_hook() + except BaseException as error: + errors.append(error) + logger.exception("Failed to unregister sub-session completion hook") + + if self.parent_cancellation is not None and self.child_cancellation is not None: + parent_cancellation = self.parent_cancellation + child_cancellation = self.child_cancellation + self.parent_cancellation = None + self.child_cancellation = None + try: + parent_cancellation.unregister_child(child_cancellation) + except BaseException as error: + errors.append(error) + logger.exception("Failed to unregister child cancellation token") + + if self.nesting_pushed and self.display_system is not None: + display_system = self.display_system + self.nesting_pushed = False + try: + if hasattr(display_system, "pop_nesting"): + display_system.pop_nesting() + except BaseException as error: + errors.append(error) + logger.exception("Failed to pop sub-session display nesting") + + if self.child_session is not None: + child_session = self.child_session + self.child_session = None + try: + await child_session.cleanup() + except BaseException as error: + errors.append(error) + logger.exception("Failed to clean up sub-session") + + if errors and not preserve_error: + raise errors[0] + + +async def _run_with_finalizer( + operation: Awaitable[_T], lifecycle: _SubSessionLifecycle +) -> _T: + """Run an operation, then shield one separately scheduled finalizer task.""" + result = cast(_T, None) + primary_error: BaseException | None = None + + try: + result = await operation + except BaseException as error: # noqa: BLE001 - teardown must follow any exit + primary_error = error + + finalizer = asyncio.create_task( + lifecycle.finalize(preserve_error=primary_error is not None) + ) + while True: + try: + await asyncio.shield(finalizer) + break + except asyncio.CancelledError as error: + if primary_error is None: + primary_error = error + # Re-await the same shielded task. A new task would duplicate work, + # while awaiting this task directly would let cancellation reach it. + if not finalizer.done(): + continue + if finalizer.cancelled(): + break + except BaseException as finalization_error: + if primary_error is None: + raise + logger.error( + "Sub-session finalization failed after primary error", + exc_info=finalization_error, + ) + break + + if primary_error is not None: + raise primary_error + return result + # Capture default sys.path entries at import time. # Used to filter out bundle-added paths when forwarding sys_paths to subprocess children. @@ -228,7 +378,7 @@ def _find_redacted_values(value: object, path: str = "") -> list[str]: return found -async def spawn_sub_session( +async def _spawn_sub_session( agent_name: str, instruction: str, parent_session: AmplifierSession, @@ -242,6 +392,8 @@ async def spawn_sub_session( self_delegation_depth: int = 0, session_metadata: dict | None = None, use_subprocess: bool = False, + *, + lifecycle: _SubSessionLifecycle, ) -> dict: """ Spawn sub-session with agent configuration overlay. @@ -539,10 +691,10 @@ async def spawn_sub_session( approval_system=parent_session.coordinator.approval_system, # Inherit from parent display_system=display_system, # Inherit from parent ) + lifecycle.child_session = child_session # Notify display system we're entering a nested session (for indentation) - if hasattr(display_system, "push_nesting"): - display_system.push_nesting() + lifecycle.push_nesting(display_system) # NOTE: Parent message injection moved to AFTER initialize() because # the context module is only mounted during initialize(). @@ -574,8 +726,9 @@ async def spawn_sub_session( paths_to_share: list[str] = [] # Source 1: Module paths from parent loader - if hasattr(parent_session, "loader") and parent_session.loader is not None: - parent_added_paths = getattr(parent_session.loader, "_added_paths", []) + parent_loader = getattr(parent_session, "loader", None) + if parent_loader is not None: + parent_added_paths = getattr(parent_loader, "_added_paths", []) paths_to_share.extend(parent_added_paths) # Source 2: Bundle package paths (src/ directories from bundles like python-dev) @@ -651,7 +804,7 @@ async def spawn_sub_session( # This enables graceful Ctrl+C handling for nested agent sessions parent_cancellation = parent_session.coordinator.cancellation child_cancellation = child_session.coordinator.cancellation - parent_cancellation.register_child(child_cancellation) + lifecycle.register_cancellation(parent_cancellation, child_cancellation) logger.debug( f"Registered child cancellation token for sub-session {sub_session_id}" ) @@ -819,6 +972,7 @@ async def _capture_completion(event: str, data: dict) -> HookResult: priority=999, name="_spawn_capture", ) + lifecycle.unregister_hook = unregister_hook # Expand @-mentions in delegation instruction before executing. # Content lands inline as XML blocks prepended to the instruction. @@ -839,75 +993,82 @@ async def _capture_completion(event: str, data: dict) -> HookResult: relative_to=_instr_rel, ) - # Execute instruction in child session; cleanup MUST run even on CancelledError - try: - try: - response = await child_session.execute(instruction) - finally: - if unregister_hook: - unregister_hook() + # Prepare canonical reconstruction state before execution so a cancelled + # child can be resumed from whatever transcript it produced. + from datetime import UTC, datetime - # Persist state for multi-turn resumption - from datetime import UTC - from datetime import datetime - - from .session_store import SessionStore - - context = child_session.coordinator.get("context") - transcript = await context.get_messages() if context else [] - - # Extract or generate trace_id for W3C Trace Context pattern - # Root session ID is the trace_id, propagate it to all children - parent_trace_id = getattr(parent_session, "trace_id", parent_session.session_id) + from .session_store import SessionStore - # Extract child_span from sub_session_id for short_id resolution - # Format: {parent_id}-{child_span}_{agent_name} - child_span: str | None = None - if sub_session_id and "_" in sub_session_id and "-" in sub_session_id: - base = sub_session_id.rsplit("_", 1)[0] # Remove agent name - child_span = base.rsplit("-", 1)[-1] # Get child_span (16 hex chars) + context = child_session.coordinator.get("context") + parent_trace_id = getattr(parent_session, "trace_id", parent_session.session_id) - metadata = { - "session_id": sub_session_id, - "parent_id": parent_session.session_id, - "trace_id": parent_trace_id, # W3C Trace Context: trace entire conversation - "agent_name": agent_name, - "child_span": child_span, # For short_id resolution (first 8 chars = short_id) - "created": datetime.now(UTC).isoformat(), - "config": merged_config, - "agent_overlay": agent_config, - "turn_count": 1, - "bundle_context": _extract_bundle_context(parent_session), - "self_delegation_depth": self_delegation_depth, # For recursion limit tracking - # Store working_dir for session sync between CLI and web - "working_dir": str(Path.cwd().resolve()), - } + child_span: str | None = None + if sub_session_id and "_" in sub_session_id and "-" in sub_session_id: + base = sub_session_id.rsplit("_", 1)[0] + child_span = base.rsplit("-", 1)[-1] - store = SessionStore() - store.save(sub_session_id, transcript, metadata) - logger.debug(f"Sub-session {sub_session_id} state persisted") + metadata = { + "session_id": sub_session_id, + "parent_id": parent_session.session_id, + "trace_id": parent_trace_id, + "agent_name": agent_name, + "child_span": child_span, + "created": datetime.now(UTC).isoformat(), + "config": merged_config, + "agent_overlay": agent_config, + "turn_count": 1, + "bundle_context": _extract_bundle_context(parent_session), + "self_delegation_depth": self_delegation_depth, + "working_dir": str(Path.cwd().resolve()), + } + store = SessionStore() - # Bridge child session costs to parent coordinator (bridge_child_cost never raises) - await bridge_child_cost( - child_coordinator=child_session.coordinator, - parent_coordinator=parent_session.coordinator, - child_session_id=sub_session_id, - ) + # Execute and collect the transcript for normal persistence. If either await + # is cancelled, the wrapper persists interruption state and tears down in + # its shielded finalizer. + try: + response = await child_session.execute(instruction) + transcript = await context.get_messages() if context else [] + except asyncio.CancelledError: + + async def _persist_interrupted_spawn() -> None: + transcript = [] + if context: + try: + transcript = await context.get_messages() + except BaseException: + logger.exception( + "Failed to read interrupted sub-session %s transcript; " + "saving fallback transcript", + sub_session_id, + ) + try: + store.save( + sub_session_id, + transcript, + {**metadata, "status": "interrupted"}, + ) + logger.debug( + "Interrupted sub-session %s state persisted", sub_session_id + ) + except BaseException: + logger.exception( + "Failed to persist interrupted sub-session %s", sub_session_id + ) - finally: - # Unregister child cancellation token before cleanup - # MUST run even if execution was cancelled (CancelledError) or failed - parent_cancellation.unregister_child(child_cancellation) - logger.debug( - f"Unregistered child cancellation token for sub-session {sub_session_id}" - ) + lifecycle.persist_on_interruption(_persist_interrupted_spawn) + raise - # Notify display system we're exiting the nested session (for indentation) - if hasattr(display_system, "pop_nesting"): - display_system.pop_nesting() + # Persist state for multi-turn resumption + store.save(sub_session_id, transcript, metadata) + logger.debug(f"Sub-session {sub_session_id} state persisted") - # Cleanup child session - await child_session.cleanup() + # Bridge child session costs to parent coordinator (bridge_child_cost never raises) + await bridge_child_cost( + child_coordinator=child_session.coordinator, + parent_coordinator=parent_session.coordinator, + child_session_id=sub_session_id, + ) # Return response and session ID for potential multi-turn # Include enriched fields from orchestrator:complete hook @@ -920,10 +1081,49 @@ async def _capture_completion(event: str, data: dict) -> HookResult: } -async def resume_sub_session( +async def spawn_sub_session( + agent_name: str, + instruction: str, + parent_session: AmplifierSession, + agent_configs: dict[str, dict], + sub_session_id: str | None = None, + tool_inheritance: dict[str, list[str]] | None = None, + hook_inheritance: dict[str, list[str]] | None = None, + orchestrator_config: dict | None = None, + parent_messages: list[dict] | None = None, + provider_preferences: list | None = None, + self_delegation_depth: int = 0, + session_metadata: dict | None = None, + use_subprocess: bool = False, +) -> dict: + """Run a spawned child under lifecycle management from construction onward.""" + lifecycle = _SubSessionLifecycle() + operation = _spawn_sub_session( + agent_name=agent_name, + instruction=instruction, + parent_session=parent_session, + agent_configs=agent_configs, + sub_session_id=sub_session_id, + tool_inheritance=tool_inheritance, + hook_inheritance=hook_inheritance, + orchestrator_config=orchestrator_config, + parent_messages=parent_messages, + provider_preferences=provider_preferences, + self_delegation_depth=self_delegation_depth, + session_metadata=session_metadata, + use_subprocess=use_subprocess, + lifecycle=lifecycle, + ) + return await _run_with_finalizer(operation, lifecycle) + + +async def _resume_sub_session( sub_session_id: str, instruction: str, parent_session: AmplifierSession | None = None, + *, + lifecycle: _SubSessionLifecycle, + child_spawn_capability, ) -> dict: """Resume existing sub-session for multi-turn engagement. @@ -942,8 +1142,7 @@ async def resume_sub_session( RuntimeError: If session metadata corrupted or incomplete ValueError: If session_id is invalid """ - from datetime import UTC - from datetime import datetime + from datetime import UTC, datetime from .session_store import SessionStore @@ -959,7 +1158,7 @@ async def resume_sub_session( transcript, metadata = store.load(sub_session_id) except Exception as e: raise RuntimeError( - f"Failed to load sub-session '{sub_session_id}': {str(e)}" + f"Failed to load sub-session '{sub_session_id}': {e!s}" ) from e # Extract reconstruction data @@ -1098,8 +1297,7 @@ async def resume_sub_session( # 2. Serializing full UX state would add significant complexity # 3. The parent session may no longer be running when sub-session resumes # 4. Approval decisions are contextual to the current execution state - from amplifier_app_cli.ui import CLIApprovalSystem - from amplifier_app_cli.ui import CLIDisplaySystem + from amplifier_app_cli.ui import CLIApprovalSystem, CLIDisplaySystem logger.debug( "Resuming sub-session %s (agent=%s, parent=%s, trace=%s). " @@ -1121,6 +1319,8 @@ async def resume_sub_session( approval_system=approval_system, display_system=display_system, ) + lifecycle.child_session = child_session + lifecycle.push_nesting(display_system) # Register app-layer capabilities for resumed child session BEFORE initialization # Must be mounted before initialize() so modules with source: directives can be resolved @@ -1216,40 +1416,9 @@ async def resume_sub_session( "self_delegation_depth", self_delegation_depth ) - # Register session spawning capabilities on resumed child session - # This enables nested agent delegation (child can spawn grandchildren) - # The capabilities are closures that reference the spawn/resume functions - async def child_spawn_capability( - agent_name: str, - instruction: str, - parent_session: "AmplifierSession", - agent_configs: dict[str, dict], - sub_session_id: str | None = None, - tool_inheritance: dict[str, list[str]] | None = None, - hook_inheritance: dict[str, list[str]] | None = None, - orchestrator_config: dict | None = None, - parent_messages: list[dict] | None = None, - provider_preferences: list | None = None, - self_delegation_depth: int = 0, - session_metadata: dict | None = None, - use_subprocess: bool = False, - ) -> dict: - return await spawn_sub_session( - agent_name=agent_name, - instruction=instruction, - parent_session=parent_session, - agent_configs=agent_configs, - sub_session_id=sub_session_id, - tool_inheritance=tool_inheritance, - hook_inheritance=hook_inheritance, - orchestrator_config=orchestrator_config, - parent_messages=parent_messages, - provider_preferences=provider_preferences, - self_delegation_depth=self_delegation_depth, - session_metadata=session_metadata, - use_subprocess=use_subprocess, - ) - + # Register session spawning capabilities on resumed child session. + # child_spawn_capability is defined by the public resume function so the + # established source layout and public introspection behavior remain intact. async def child_resume_capability(sub_session_id: str, instruction: str) -> dict: return await resume_sub_session( sub_session_id=sub_session_id, @@ -1321,19 +1490,19 @@ async def _capture_completion(event: str, data: dict) -> HookResult: priority=999, name="_spawn_capture", ) + lifecycle.unregister_hook = unregister_hook # Wire up cancellation propagation if parent session provided # Enables graceful Ctrl+C to stop the child after its current tool call if parent_session is not None: resume_parent_cancellation = parent_session.coordinator.cancellation resume_child_cancellation = child_session.coordinator.cancellation - resume_parent_cancellation.register_child(resume_child_cancellation) + lifecycle.register_cancellation( + resume_parent_cancellation, resume_child_cancellation + ) logger.debug( f"Registered child cancellation token for resumed sub-session {sub_session_id}" ) - else: - resume_parent_cancellation = None - resume_child_cancellation = None # Expand @-mentions in the resumed instruction (consistent with spawn path). # Content lands inline as XML blocks prepended to the instruction. @@ -1354,46 +1523,58 @@ async def _capture_completion(event: str, data: dict) -> HookResult: relative_to=_resume_rel, ) - # Execute new instruction with full context; cleanup MUST run even on CancelledError + # Execute and collect the transcript for normal persistence. If either await + # is cancelled, the wrapper persists interruption state and tears down in + # its shielded finalizer. try: - try: - response = await child_session.execute(instruction) - finally: - if unregister_hook: - unregister_hook() - - # Update state for next resumption + response = await child_session.execute(instruction) updated_transcript = await context.get_messages() if context else [] - metadata["turn_count"] = len(updated_transcript) - metadata["last_updated"] = datetime.now(UTC).isoformat() + except asyncio.CancelledError: + + async def _persist_interrupted_resume() -> None: + updated_transcript = transcript + if context: + try: + updated_transcript = await context.get_messages() + except BaseException: + logger.exception( + "Failed to read interrupted sub-session %s transcript; " + "saving loaded transcript as fallback", + sub_session_id, + ) + metadata["turn_count"] = len(updated_transcript) + metadata["last_updated"] = datetime.now(UTC).isoformat() + metadata["status"] = "interrupted" + try: + store.save(sub_session_id, updated_transcript, metadata) + logger.debug( + "Interrupted sub-session %s state persisted", sub_session_id + ) + except BaseException: + logger.exception( + "Failed to persist interrupted sub-session %s", sub_session_id + ) - store.save(sub_session_id, updated_transcript, metadata) - logger.debug( - f"Sub-session {sub_session_id} state updated (turn {metadata['turn_count']})" - ) + lifecycle.persist_on_interruption(_persist_interrupted_resume) + raise - # Bridge child session costs to parent coordinator (bridge_child_cost never raises) - if parent_session is not None: - await bridge_child_cost( - child_coordinator=child_session.coordinator, - parent_coordinator=parent_session.coordinator, - child_session_id=sub_session_id, - ) + # Update state for next resumption and clear any prior interruption marker. + metadata["turn_count"] = len(updated_transcript) + metadata["last_updated"] = datetime.now(UTC).isoformat() + metadata.pop("status", None) - finally: - # Unregister child cancellation token before cleanup - # MUST run even if execution was cancelled (CancelledError) or failed - if ( - resume_parent_cancellation is not None - and resume_child_cancellation is not None - ): - resume_parent_cancellation.unregister_child(resume_child_cancellation) - logger.debug( - f"Unregistered child cancellation token for resumed sub-session {sub_session_id}" - ) + store.save(sub_session_id, updated_transcript, metadata) + logger.debug( + f"Sub-session {sub_session_id} state updated (turn {metadata['turn_count']})" + ) - # Cleanup child session - await child_session.cleanup() + # Bridge child session costs to parent coordinator (bridge_child_cost never raises) + if parent_session is not None: + await bridge_child_cost( + child_coordinator=child_session.coordinator, + parent_coordinator=parent_session.coordinator, + child_session_id=sub_session_id, + ) # Return response and same session ID # Include enriched fields from orchestrator:complete hook @@ -1404,3 +1585,55 @@ async def _capture_completion(event: str, data: dict) -> HookResult: "turn_count": completion_data.get("turn_count", 1), "metadata": completion_data.get("metadata", {}), } + + +async def resume_sub_session( + sub_session_id: str, + instruction: str, + parent_session: AmplifierSession | None = None, +) -> dict: + """Run a resumed child under lifecycle management from construction onward.""" + + # Keep this capability closure in the public resume implementation. Besides + # preserving source-level compatibility, it is the callable mounted on the + # resumed child and enables subprocess grandchildren. + async def child_spawn_capability( + agent_name: str, + instruction: str, + parent_session: "AmplifierSession", + agent_configs: dict[str, dict], + sub_session_id: str | None = None, + tool_inheritance: dict[str, list[str]] | None = None, + hook_inheritance: dict[str, list[str]] | None = None, + orchestrator_config: dict | None = None, + parent_messages: list[dict] | None = None, + provider_preferences: list | None = None, + self_delegation_depth: int = 0, + session_metadata: dict | None = None, + use_subprocess: bool = False, + ) -> dict: + return await spawn_sub_session( + agent_name=agent_name, + instruction=instruction, + parent_session=parent_session, + agent_configs=agent_configs, + sub_session_id=sub_session_id, + tool_inheritance=tool_inheritance, + hook_inheritance=hook_inheritance, + orchestrator_config=orchestrator_config, + parent_messages=parent_messages, + provider_preferences=provider_preferences, + self_delegation_depth=self_delegation_depth, + session_metadata=session_metadata, + use_subprocess=use_subprocess, + ) + + lifecycle = _SubSessionLifecycle() + operation = _resume_sub_session( + sub_session_id=sub_session_id, + instruction=instruction, + parent_session=parent_session, + lifecycle=lifecycle, + child_spawn_capability=child_spawn_capability, + ) + return await _run_with_finalizer(operation, lifecycle) diff --git a/tests/test_session_spawner_interruption.py b/tests/test_session_spawner_interruption.py new file mode 100644 index 00000000..d1e02623 --- /dev/null +++ b/tests/test_session_spawner_interruption.py @@ -0,0 +1,708 @@ +"""Deterministic cancellation coverage for in-process sub-sessions.""" + +import asyncio +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from amplifier_app_cli.session_spawner import resume_sub_session, spawn_sub_session + +pytestmark = pytest.mark.anyio + + +@pytest.fixture(scope="module") +def anyio_backend(): + return "asyncio" + + +class FakeHooks: + def __init__(self): + self.unregister = MagicMock() + self.emit = AsyncMock() + + def register(self, event, handler, priority=0, name=None): + return self.unregister + + +def make_child(transcript): + hooks = FakeHooks() + context = MagicMock() + context.get_messages = AsyncMock(return_value=transcript) + context.add_message = AsyncMock() + + coordinator = MagicMock() + coordinator.config = {} + coordinator.get_capability.return_value = None + coordinator.mount = AsyncMock() + coordinator.collect_contributions = AsyncMock(return_value=[]) + coordinator.cancellation = MagicMock() + + def get_component(name): + return {"hooks": hooks, "context": context}.get(name) + + coordinator.get = get_component + + child = MagicMock() + child.coordinator = coordinator + child.initialize = AsyncMock() + child.cleanup = AsyncMock() + return child, hooks, context + + +def make_parent(display=None): + coordinator = MagicMock() + coordinator.config = {} + coordinator.get.return_value = None + coordinator.get_capability.return_value = None + coordinator.display_system = display or MagicMock() + coordinator.cancellation = MagicMock() + + parent = MagicMock() + parent.coordinator = coordinator + parent.config = { + "session": {"orchestrator": "loop-basic", "context": "context-simple"} + } + parent.session_id = "parent-123" + parent.trace_id = "trace-abc" + parent.loader = None + return parent + + +async def test_spawn_cancellation_persists_partial_transcript_and_unwinds_once(): + partial = [ + {"role": "user", "content": "work"}, + {"role": "assistant", "content": "partial"}, + ] + child, hooks, context = make_child(partial) + cancellation = asyncio.CancelledError("cancel spawn") + child.execute = AsyncMock(side_effect=cancellation) + parent = make_parent() + store = MagicMock() + + with ( + patch("amplifier_app_cli.session_spawner.AmplifierSession", return_value=child), + patch("amplifier_app_cli.paths.create_foundation_resolver"), + patch("amplifier_app_cli.session_store.SessionStore", return_value=store), + pytest.raises(asyncio.CancelledError) as raised, + ): + await spawn_sub_session( + agent_name="test-agent", + instruction="do work", + parent_session=parent, + agent_configs={"test-agent": {"description": "test"}}, + sub_session_id="parent-abcdef_test-agent", + ) + + assert raised.value is cancellation + store.save.assert_called_once() + session_id, saved_transcript, saved_metadata = store.save.call_args.args + assert session_id == "parent-abcdef_test-agent" + assert saved_transcript == partial + assert saved_metadata["status"] == "interrupted" + assert saved_metadata["session_id"] == session_id + assert saved_metadata["parent_id"] == "parent-123" + assert saved_metadata["trace_id"] == "trace-abc" + assert saved_metadata["agent_name"] == "test-agent" + assert saved_metadata["config"]["session"] == parent.config["session"] + assert saved_metadata["agent_overlay"] == {"description": "test"} + context.get_messages.assert_awaited_once() + hooks.unregister.assert_called_once_with() + parent.coordinator.cancellation.unregister_child.assert_called_once_with( + child.coordinator.cancellation + ) + parent.coordinator.display_system.pop_nesting.assert_called_once_with() + child.cleanup.assert_awaited_once_with() + + +async def test_resume_cancellation_preserves_metadata_and_unwinds_once(): + original = [{"role": "user", "content": "first"}] + partial = original + [{"role": "assistant", "content": "partial follow-up"}] + child, hooks, context = make_child(partial) + cancellation = asyncio.CancelledError("cancel resume") + child.execute = AsyncMock(side_effect=cancellation) + display = MagicMock() + parent = make_parent() + store = MagicMock() + metadata = { + "session_id": "resumed-child", + "parent_id": "parent-123", + "agent_name": "test-agent", + "trace_id": "trace-abc", + "config": { + "session": {"orchestrator": "loop-basic", "context": "context-simple"} + }, + "working_dir": "/fixed/project", + "resume_marker": "keep-me", + } + store.exists.return_value = True + store.load.return_value = (original, metadata) + + with ( + patch("amplifier_app_cli.session_spawner.AmplifierSession", return_value=child), + patch("amplifier_app_cli.ui.CLIApprovalSystem"), + patch("amplifier_app_cli.ui.CLIDisplaySystem", return_value=display), + patch("amplifier_app_cli.paths.create_foundation_resolver"), + patch("amplifier_app_cli.session_store.SessionStore", return_value=store), + pytest.raises(asyncio.CancelledError) as raised, + ): + await resume_sub_session("resumed-child", "continue", parent) + + assert raised.value is cancellation + store.save.assert_called_once() + session_id, saved_transcript, saved_metadata = store.save.call_args.args + assert session_id == "resumed-child" + assert saved_transcript == partial + assert saved_metadata["status"] == "interrupted" + assert saved_metadata["resume_marker"] == "keep-me" + assert saved_metadata["config"] == metadata["config"] + assert saved_metadata["turn_count"] == len(partial) + assert "last_updated" in saved_metadata + context.get_messages.assert_awaited_once() + hooks.unregister.assert_called_once_with() + parent.coordinator.cancellation.unregister_child.assert_called_once_with( + child.coordinator.cancellation + ) + display.push_nesting.assert_called_once_with() + display.pop_nesting.assert_called_once_with() + child.cleanup.assert_awaited_once_with() + + +async def test_spawn_cancel_during_post_execute_get_messages_persists_interrupted_once(): + partial = [{"role": "assistant", "content": "completed before persistence"}] + child, hooks, context = make_child(partial) + child.execute = AsyncMock(return_value="done") + parent = make_parent() + store = MagicMock() + transcript_started = asyncio.Event() + transcript_cancellations = [] + transcript_reads = 0 + + async def get_messages(): + nonlocal transcript_reads + transcript_reads += 1 + if transcript_reads == 1: + transcript_started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError as error: + transcript_cancellations.append(error) + raise + return partial + + context.get_messages = AsyncMock(side_effect=get_messages) + + with ( + patch("amplifier_app_cli.session_spawner.AmplifierSession", return_value=child), + patch("amplifier_app_cli.paths.create_foundation_resolver"), + patch("amplifier_app_cli.session_store.SessionStore", return_value=store), + ): + task = asyncio.create_task( + spawn_sub_session( + agent_name="test-agent", + instruction="do work", + parent_session=parent, + agent_configs={"test-agent": {"description": "test"}}, + sub_session_id="post-execute-cancel-spawn", + ) + ) + await transcript_started.wait() + task.cancel("cancel spawn transcript read") + + with pytest.raises(asyncio.CancelledError) as raised: + await task + + assert raised.value is transcript_cancellations[0] + assert context.get_messages.await_count == 2 + store.save.assert_called_once() + assert store.save.call_args.args[0] == "post-execute-cancel-spawn" + assert store.save.call_args.args[1] == partial + assert store.save.call_args.args[2]["status"] == "interrupted" + hooks.unregister.assert_called_once_with() + parent.coordinator.cancellation.unregister_child.assert_called_once_with( + child.coordinator.cancellation + ) + parent.coordinator.display_system.pop_nesting.assert_called_once_with() + child.cleanup.assert_awaited_once_with() + + +async def test_resume_cancel_during_post_execute_get_messages_persists_interrupted_once(): + original = [{"role": "user", "content": "first turn"}] + partial = original + [{"role": "assistant", "content": "completed follow-up"}] + child, hooks, context = make_child(partial) + child.execute = AsyncMock(return_value="done") + parent = make_parent() + display = MagicMock() + store = MagicMock() + metadata = { + "session_id": "post-execute-cancel-resume", + "parent_id": "parent-123", + "agent_name": "test-agent", + "config": { + "session": {"orchestrator": "loop-basic", "context": "context-simple"} + }, + "working_dir": "/fixed/project", + "resume_marker": "preserved", + } + store.exists.return_value = True + store.load.return_value = (original, metadata) + transcript_started = asyncio.Event() + transcript_cancellations = [] + transcript_reads = 0 + + async def get_messages(): + nonlocal transcript_reads + transcript_reads += 1 + if transcript_reads == 1: + transcript_started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError as error: + transcript_cancellations.append(error) + raise + return partial + + context.get_messages = AsyncMock(side_effect=get_messages) + + with ( + patch("amplifier_app_cli.session_spawner.AmplifierSession", return_value=child), + patch("amplifier_app_cli.ui.CLIApprovalSystem"), + patch("amplifier_app_cli.ui.CLIDisplaySystem", return_value=display), + patch("amplifier_app_cli.paths.create_foundation_resolver"), + patch("amplifier_app_cli.session_store.SessionStore", return_value=store), + ): + task = asyncio.create_task( + resume_sub_session("post-execute-cancel-resume", "continue", parent) + ) + await transcript_started.wait() + task.cancel("cancel resume transcript read") + + with pytest.raises(asyncio.CancelledError) as raised: + await task + + assert raised.value is transcript_cancellations[0] + assert context.get_messages.await_count == 2 + store.save.assert_called_once() + session_id, saved_transcript, saved_metadata = store.save.call_args.args + assert session_id == "post-execute-cancel-resume" + assert saved_transcript == partial + assert saved_metadata["status"] == "interrupted" + assert saved_metadata["resume_marker"] == "preserved" + assert saved_metadata["turn_count"] == len(partial) + assert "last_updated" in saved_metadata + hooks.unregister.assert_called_once_with() + parent.coordinator.cancellation.unregister_child.assert_called_once_with( + child.coordinator.cancellation + ) + display.push_nesting.assert_called_once_with() + display.pop_nesting.assert_called_once_with() + child.cleanup.assert_awaited_once_with() + + +async def test_spawn_setup_cancellation_unwinds_resources_once(): + child, hooks, _ = make_child([]) + child.execute = AsyncMock() + parent = make_parent() + cancellation = asyncio.CancelledError("cancel mention expansion") + + def get_capability(name): + return object() if name == "mention_resolver" else None + + child.coordinator.get_capability.side_effect = get_capability + + with ( + patch("amplifier_app_cli.session_spawner.AmplifierSession", return_value=child), + patch("amplifier_app_cli.paths.create_foundation_resolver"), + patch( + "amplifier_foundation.mentions.expand_mentions_in_instruction", + new=AsyncMock(side_effect=cancellation), + ), + pytest.raises(asyncio.CancelledError) as raised, + ): + await spawn_sub_session( + agent_name="test-agent", + instruction="expand @mention", + parent_session=parent, + agent_configs={"test-agent": {"description": "test"}}, + sub_session_id="setup-cancelled-child", + ) + + assert raised.value is cancellation + child.execute.assert_not_awaited() + hooks.unregister.assert_called_once_with() + parent.coordinator.cancellation.unregister_child.assert_called_once_with( + child.coordinator.cancellation + ) + parent.coordinator.display_system.pop_nesting.assert_called_once_with() + child.cleanup.assert_awaited_once_with() + + +async def test_resume_setup_cancellation_unwinds_resources_once(): + child, hooks, _ = make_child([]) + child.execute = AsyncMock() + parent = make_parent() + display = MagicMock() + cancellation = asyncio.CancelledError("cancel resume mention expansion") + store = MagicMock() + store.exists.return_value = True + store.load.return_value = ( + [], + { + "session_id": "resumed-child", + "parent_id": "parent-123", + "agent_name": "test-agent", + "config": { + "session": { + "orchestrator": "loop-basic", + "context": "context-simple", + } + }, + "working_dir": "/fixed/project", + }, + ) + + def get_capability(name): + return object() if name == "mention_resolver" else None + + child.coordinator.get_capability.side_effect = get_capability + + with ( + patch("amplifier_app_cli.session_spawner.AmplifierSession", return_value=child), + patch("amplifier_app_cli.ui.CLIApprovalSystem"), + patch("amplifier_app_cli.ui.CLIDisplaySystem", return_value=display), + patch("amplifier_app_cli.paths.create_foundation_resolver"), + patch("amplifier_app_cli.session_store.SessionStore", return_value=store), + patch( + "amplifier_foundation.mentions.expand_mentions_in_instruction", + new=AsyncMock(side_effect=cancellation), + ), + pytest.raises(asyncio.CancelledError) as raised, + ): + await resume_sub_session("resumed-child", "expand @mention", parent) + + assert raised.value is cancellation + child.execute.assert_not_awaited() + hooks.unregister.assert_called_once_with() + parent.coordinator.cancellation.unregister_child.assert_called_once_with( + child.coordinator.cancellation + ) + display.push_nesting.assert_called_once_with() + display.pop_nesting.assert_called_once_with() + child.cleanup.assert_awaited_once_with() + + +async def test_spawn_preserves_cancellation_when_persistence_and_teardown_fail(): + child, hooks, context = make_child([{"role": "assistant", "content": "partial"}]) + cancellation = asyncio.CancelledError("original cancellation") + child.execute = AsyncMock(side_effect=cancellation) + child.cleanup = AsyncMock(side_effect=RuntimeError("cleanup failed")) + hooks.unregister.side_effect = RuntimeError("hook unregister failed") + parent = make_parent() + parent.coordinator.cancellation.unregister_child.side_effect = RuntimeError( + "token unregister failed" + ) + parent.coordinator.display_system.pop_nesting.side_effect = RuntimeError( + "display pop failed" + ) + store = MagicMock() + store.save.side_effect = OSError("save failed") + + with ( + patch("amplifier_app_cli.session_spawner.AmplifierSession", return_value=child), + patch("amplifier_app_cli.paths.create_foundation_resolver"), + patch("amplifier_app_cli.session_store.SessionStore", return_value=store), + pytest.raises(asyncio.CancelledError) as raised, + ): + await spawn_sub_session( + agent_name="test-agent", + instruction="do work", + parent_session=parent, + agent_configs={"test-agent": {"description": "test"}}, + sub_session_id="failure-child", + ) + + assert raised.value is cancellation + context.get_messages.assert_awaited_once_with() + store.save.assert_called_once() + hooks.unregister.assert_called_once_with() + parent.coordinator.cancellation.unregister_child.assert_called_once_with( + child.coordinator.cancellation + ) + parent.coordinator.display_system.pop_nesting.assert_called_once_with() + child.cleanup.assert_awaited_once_with() + + +async def test_successful_resume_clears_interrupted_status(): + child, _, context = make_child([{"role": "assistant", "content": "complete"}]) + child.execute = AsyncMock(return_value="done") + store = MagicMock() + metadata = { + "session_id": "resumed-child", + "parent_id": "parent-123", + "agent_name": "test-agent", + "config": { + "session": {"orchestrator": "loop-basic", "context": "context-simple"} + }, + "working_dir": "/fixed/project", + "status": "interrupted", + } + store.exists.return_value = True + store.load.return_value = ([], metadata) + + with ( + patch("amplifier_app_cli.session_spawner.AmplifierSession", return_value=child), + patch("amplifier_app_cli.ui.CLIApprovalSystem"), + patch("amplifier_app_cli.ui.CLIDisplaySystem"), + patch("amplifier_app_cli.paths.create_foundation_resolver"), + patch("amplifier_app_cli.session_store.SessionStore", return_value=store), + ): + result = await resume_sub_session("resumed-child", "continue") + + assert set(result) == {"output", "session_id", "status", "turn_count", "metadata"} + assert result["output"] == "done" + store.save.assert_called_once() + assert "status" not in store.save.call_args.args[2] + context.get_messages.assert_awaited_once_with() + + +async def test_spawn_repeated_cancellation_cannot_interrupt_save_or_cleanup(): + partial = [{"role": "assistant", "content": "partial spawn"}] + child, hooks, context = make_child(partial) + parent = make_parent() + store = MagicMock() + execute_started = asyncio.Event() + transcript_started = asyncio.Event() + release_transcript = asyncio.Event() + cleanup_started = asyncio.Event() + release_cleanup = asyncio.Event() + execution_cancellations = [] + + async def execute(_instruction): + execute_started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError as error: + execution_cancellations.append(error) + raise + + async def get_messages(): + transcript_started.set() + await release_transcript.wait() + return partial + + async def cleanup(): + cleanup_started.set() + await release_cleanup.wait() + + child.execute = AsyncMock(side_effect=execute) + context.get_messages = AsyncMock(side_effect=get_messages) + child.cleanup = AsyncMock(side_effect=cleanup) + + with ( + patch("amplifier_app_cli.session_spawner.AmplifierSession", return_value=child), + patch("amplifier_app_cli.paths.create_foundation_resolver"), + patch("amplifier_app_cli.session_store.SessionStore", return_value=store), + ): + task = asyncio.create_task( + spawn_sub_session( + agent_name="test-agent", + instruction="do work", + parent_session=parent, + agent_configs={"test-agent": {"description": "test"}}, + sub_session_id="repeated-cancel-spawn", + ) + ) + await execute_started.wait() + task.cancel("original spawn cancellation") + await transcript_started.wait() + task.cancel("cancel while reading transcript") + release_transcript.set() + await cleanup_started.wait() + task.cancel("cancel while cleaning up") + release_cleanup.set() + + with pytest.raises(asyncio.CancelledError) as raised: + await task + + assert raised.value is execution_cancellations[0] + store.save.assert_called_once() + assert store.save.call_args.args[1] == partial + assert store.save.call_args.args[2]["status"] == "interrupted" + context.get_messages.assert_awaited_once_with() + hooks.unregister.assert_called_once_with() + parent.coordinator.cancellation.unregister_child.assert_called_once_with( + child.coordinator.cancellation + ) + parent.coordinator.display_system.pop_nesting.assert_called_once_with() + child.cleanup.assert_awaited_once_with() + + +async def test_resume_repeated_cancellation_cannot_interrupt_save_or_cleanup(): + original = [{"role": "user", "content": "first turn"}] + partial = original + [{"role": "assistant", "content": "partial resume"}] + child, hooks, context = make_child(partial) + parent = make_parent() + display = MagicMock() + store = MagicMock() + store.exists.return_value = True + store.load.return_value = ( + original, + { + "session_id": "repeated-cancel-resume", + "parent_id": "parent-123", + "agent_name": "test-agent", + "config": { + "session": { + "orchestrator": "loop-basic", + "context": "context-simple", + } + }, + "working_dir": "/fixed/project", + "resume_marker": "preserved", + }, + ) + execute_started = asyncio.Event() + transcript_started = asyncio.Event() + release_transcript = asyncio.Event() + cleanup_started = asyncio.Event() + release_cleanup = asyncio.Event() + execution_cancellations = [] + + async def execute(_instruction): + execute_started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError as error: + execution_cancellations.append(error) + raise + + async def get_messages(): + transcript_started.set() + await release_transcript.wait() + return partial + + async def cleanup(): + cleanup_started.set() + await release_cleanup.wait() + + child.execute = AsyncMock(side_effect=execute) + context.get_messages = AsyncMock(side_effect=get_messages) + child.cleanup = AsyncMock(side_effect=cleanup) + + with ( + patch("amplifier_app_cli.session_spawner.AmplifierSession", return_value=child), + patch("amplifier_app_cli.ui.CLIApprovalSystem"), + patch("amplifier_app_cli.ui.CLIDisplaySystem", return_value=display), + patch("amplifier_app_cli.paths.create_foundation_resolver"), + patch("amplifier_app_cli.session_store.SessionStore", return_value=store), + ): + task = asyncio.create_task( + resume_sub_session("repeated-cancel-resume", "continue", parent) + ) + await execute_started.wait() + task.cancel("original resume cancellation") + await transcript_started.wait() + task.cancel("cancel while reading resumed transcript") + release_transcript.set() + await cleanup_started.wait() + task.cancel("cancel while cleaning up resumed child") + release_cleanup.set() + + with pytest.raises(asyncio.CancelledError) as raised: + await task + + assert raised.value is execution_cancellations[0] + store.save.assert_called_once() + saved_metadata = store.save.call_args.args[2] + assert store.save.call_args.args[1] == partial + assert saved_metadata["status"] == "interrupted" + assert saved_metadata["resume_marker"] == "preserved" + context.get_messages.assert_awaited_once_with() + hooks.unregister.assert_called_once_with() + parent.coordinator.cancellation.unregister_child.assert_called_once_with( + child.coordinator.cancellation + ) + display.push_nesting.assert_called_once_with() + display.pop_nesting.assert_called_once_with() + child.cleanup.assert_awaited_once_with() + + +async def test_resume_get_messages_failure_saves_loaded_transcript_once(): + original = [{"role": "user", "content": "recover me"}] + child, _, context = make_child([]) + cancellation = asyncio.CancelledError("cancel resume") + child.execute = AsyncMock(side_effect=cancellation) + context.get_messages = AsyncMock(side_effect=RuntimeError("context unavailable")) + store = MagicMock() + store.exists.return_value = True + store.load.return_value = ( + original, + { + "session_id": "fallback-resume", + "parent_id": "parent-123", + "agent_name": "test-agent", + "config": { + "session": { + "orchestrator": "loop-basic", + "context": "context-simple", + } + }, + "working_dir": "/fixed/project", + }, + ) + + with ( + patch("amplifier_app_cli.session_spawner.AmplifierSession", return_value=child), + patch("amplifier_app_cli.ui.CLIApprovalSystem"), + patch("amplifier_app_cli.ui.CLIDisplaySystem"), + patch("amplifier_app_cli.paths.create_foundation_resolver"), + patch("amplifier_app_cli.session_store.SessionStore", return_value=store), + pytest.raises(asyncio.CancelledError) as raised, + ): + await resume_sub_session("fallback-resume", "continue") + + assert raised.value is cancellation + context.get_messages.assert_awaited_once_with() + store.save.assert_called_once() + assert store.save.call_args.args[1] == original + assert store.save.call_args.args[2]["status"] == "interrupted" + child.cleanup.assert_awaited_once_with() + + +async def test_spawn_cancellation_after_normal_save_does_not_save_twice(): + child, _, context = make_child([{"role": "assistant", "content": "complete"}]) + child.execute = AsyncMock(return_value="done") + parent = make_parent() + store = MagicMock() + bridge_started = asyncio.Event() + + async def blocking_bridge(**_kwargs): + bridge_started.set() + await asyncio.Event().wait() + + with ( + patch("amplifier_app_cli.session_spawner.AmplifierSession", return_value=child), + patch("amplifier_app_cli.paths.create_foundation_resolver"), + patch("amplifier_app_cli.session_store.SessionStore", return_value=store), + patch( + "amplifier_app_cli.session_spawner.bridge_child_cost", + side_effect=blocking_bridge, + ), + ): + task = asyncio.create_task( + spawn_sub_session( + agent_name="test-agent", + instruction="do work", + parent_session=parent, + agent_configs={"test-agent": {"description": "test"}}, + sub_session_id="cancel-after-save", + ) + ) + await bridge_started.wait() + task.cancel("cancel after normal save") + with pytest.raises(asyncio.CancelledError): + await task + + context.get_messages.assert_awaited_once_with() + store.save.assert_called_once() + assert "status" not in store.save.call_args.args[2] + child.cleanup.assert_awaited_once_with()