fix(tui): propagate aclose cancellation; dedupe modal suspend/restore

Address CodeRabbit review on #219:
- PromptLifecycle.aclose no longer swallows CancelledError of aclose()
  itself: a wait_for-timed-out shutdown now propagates cancellation
  instead of continuing to close remaining resources past the deadline.
  Individual closers still own their internal cancellation.
- Extract the duplicated modal suspend/restore transition shared by
  ModalAttached/ModalDetached into _suspend_restore_effects (behavior
  preserving).

Author: elkaix

src/pythinker_code/ui/shell/prompting/lifecycle.py

@@ -60,7 +60,13 @@ async def aclose(self) -> None:
                 try:
                     await closer()
                 except asyncio.CancelledError:
-                    continue
+                    # Cancellation here means aclose() itself was cancelled (e.g.
+                    # a wait_for timeout), not a closer's own internal cancel —
+                    # propagate it instead of swallowing and closing on regardless.
+                    logger.warning(
+                        "Prompt lifecycle aclose cancelled while closing resource={}", name
+                    )
+                    raise
                 except Exception as exc:
                     logger.warning(
                         "Prompt lifecycle resource failed during shutdown: resource={} error={!r}",

src/pythinker_code/ui/shell/prompting/state.py

@@ -182,6 +182,29 @@ def _mode_effects(mode: PromptMode) -> tuple[PromptEffect, ...]:
     return SelectCompleter(mode), SetEraseWhenDone(mode is PromptMode.AGENT), Invalidate()
 
 
+def _suspend_restore_effects(
+    old_active: ModalState | None,
+    new_active: ModalState | None,
+    suspended: Document | None,
+    document: Document,
+) -> tuple[Document | None, tuple[PromptEffect, ...]]:
+    """Compute the suspended-document/effects transition shared by modal
+    attach and detach: suspend the live input when a hides-input modal takes
+    over, restore it when the last such modal leaves."""
+    old_hides_input = old_active is not None and old_active.hides_input
+    new_hides_input = new_active is not None and new_active.hides_input
+    if not old_hides_input and new_hides_input and document.text:
+        if suspended is None:
+            return document, (SuspendDocument(document),)
+        return suspended, ()
+    if old_hides_input and not new_hides_input and suspended is not None:
+        effects: tuple[PromptEffect, ...] = (
+            (RestoreDocument(suspended),) if not document.text else ()
+        )
+        return None, effects
+    return suspended, ()
+
+
 def transition(state: PromptState, event: PromptEvent) -> PromptTransition:
     """Return the next prompt state and ordered facade effects without doing I/O."""
     if isinstance(event, TurnStarting):
@@ -231,18 +254,9 @@ def transition(state: PromptState, event: PromptEvent) -> PromptTransition:
             ModalState(event.delegate, event.priority, event.hides_input),
         )
         new_active = _active_modal(stack)
-        suspended = state.suspended_document
-        effects: tuple[PromptEffect, ...] = ()
-        old_hides_input = old_active is not None and old_active.hides_input
-        new_hides_input = new_active is not None and new_active.hides_input
-        if not old_hides_input and new_hides_input and event.document.text:
-            if suspended is None:
-                suspended = event.document
-                effects = (SuspendDocument(event.document),)
-        elif old_hides_input and not new_hides_input and suspended is not None:
-            if not event.document.text:
-                effects = (RestoreDocument(suspended),)
-            suspended = None
+        suspended, effects = _suspend_restore_effects(
+            old_active, new_active, state.suspended_document, event.document
+        )
         next_state = replace(
             state,
             modal_stack=stack,
@@ -257,18 +271,9 @@ def transition(state: PromptState, event: PromptEvent) -> PromptTransition:
         old_active = _active_modal(state.modal_stack)
         stack = tuple(modal for modal in state.modal_stack if modal.delegate is not event.delegate)
         new_active = _active_modal(stack)
-        suspended = state.suspended_document
-        effects: tuple[PromptEffect, ...] = ()
-        old_hides_input = old_active is not None and old_active.hides_input
-        new_hides_input = new_active is not None and new_active.hides_input
-        if not old_hides_input and new_hides_input and event.document.text:
-            if suspended is None:
-                suspended = event.document
-                effects = (SuspendDocument(event.document),)
-        elif old_hides_input and not new_hides_input and suspended is not None:
-            if not event.document.text:
-                effects = (RestoreDocument(suspended),)
-            suspended = None
+        suspended, effects = _suspend_restore_effects(
+            old_active, new_active, state.suspended_document, event.document
+        )
         next_state = replace(state, modal_stack=stack, suspended_document=suspended)
         return PromptTransition(next_state, (*effects, Invalidate()))
 

tests/ui_and_conv/test_prompt_lifecycle.py

@@ -97,6 +97,31 @@ async def refused() -> None:
         lifecycle.create_task(refused())
 
 
+@pytest.mark.asyncio
+async def test_aclose_cancellation_propagates_and_halts_remaining_closers() -> None:
+    # If aclose() itself is cancelled (e.g. a wait_for timeout) while awaiting a
+    # closer, the cancellation must propagate — not be swallowed so shutdown
+    # keeps closing the remaining resources past the caller's deadline.
+    lifecycle = PromptLifecycle()
+    earlier_closer_ran = False
+
+    async def blocking_closer() -> None:
+        await asyncio.Event().wait()
+
+    async def earlier_closer() -> None:
+        nonlocal earlier_closer_ran
+        earlier_closer_ran = True
+
+    # Closers run in reverse registration order, so "blocking" runs first.
+    lifecycle.register_closer("earlier", earlier_closer)
+    lifecycle.register_closer("blocking", blocking_closer)
+
+    with pytest.raises(asyncio.TimeoutError):
+        await asyncio.wait_for(lifecycle.aclose(), timeout=0.05)
+
+    assert earlier_closer_ran is False
+
+
 class _BlockingStdout:
     def __init__(self) -> None:
         self.read_started = asyncio.Event()