Retire held backends and join original session observations

src/agent_comms/goal_scheduler.py

@@ -41,7 +41,7 @@
 if TYPE_CHECKING:
     from .threads import Thread
     from .comms import Comms
-    from .input_drain import InputDrain
+    from .input_drain import InputDrain, OwnerSession
     from .session_lifecycle import SessionLifecycle
     from .turn_effects import TurnEffects
 
@@ -263,16 +263,23 @@ async def sync_goal_execution(
             )
 
     def observe_root(
-        self, document: RegistryDocument, session_ids: Iterable[AcpSessionId], paths: frozenset[Path],
-    ) -> None:
+        self, document: RegistryDocument, owners: Iterable[OwnerSession], paths: frozenset[Path],
+    ) -> tuple[asyncio.Task[None], ...]:
         """Request independent publications without holding up the root's inbox wake.
 
-        RuntimeServer owns each task's retirement and original failure; these
-        projections are not work on a nonexistent inbox input.
+        The original OwnerSession admits and joins publication before retirement;
+        RuntimeServer reports its original failure. No inbox input is invented.
         """
-        execution_changed = not paths.isdisjoint(GoalWakeFile.execution_paths_in(self.comms.root))
-        for session_id in session_ids:
-            thread_name = self.sessions.require(session_id)
-            request = (GoalRefresh() if execution_changed else
-                       DeclaredGoalRefresh(document.entry(thread_name).thread.goal))
-            self.effects._runtime.background(self.sync_goal_execution(session_id, request))
+        return tuple(self.effects._runtime.background(owner.observe_live(partial(
+            self.publish_observed_goal, owner.identifier, document, paths,
+        ))) for owner in owners)
+
+    async def publish_observed_goal(
+        self, session_id: AcpSessionId, document: RegistryDocument, paths: frozenset[Path],
+    ) -> None:
+        """Resolve captured source data only after its original owner admits observation."""
+        thread_name = self.sessions.require(session_id)
+        request = (GoalRefresh()
+                   if not paths.isdisjoint(GoalWakeFile.execution_paths_in(self.comms.root))
+                   else DeclaredGoalRefresh(document.entry(thread_name).thread.goal))
+        await self.sync_goal_execution(session_id, request)

src/agent_comms/input_drain.py

@@ -254,7 +254,7 @@ async def observe_root(self) -> None:
                             state.observation_ready.set()
                     else:
                         self.effects.turns.goals.observe_root(
-                            document, (owner.identifier for state in targets.values()
+                            document, (owner for state in targets.values()
                                        for owner in state.executors()), watcher.observed_paths,
                         )
                     self.observed_registry = read
@@ -737,6 +737,7 @@ class OwnerSession(DrainedSession):
     auth_observer: asyncio.Task[None] = field(kw_only=True)
     retirement_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
     retired: asyncio.Future[None] = field(default_factory=asyncio.Future)
+    _observations: set[asyncio.Future[None]] = field(default_factory=set, init=False, repr=False)
 
     def owned(self) -> OwnerSession:
         return self
@@ -752,6 +753,27 @@ def configuration_for(self, ingress: ConfigOptions) -> ConfigOptions:
     def executors(self) -> tuple[OwnerSession, ...]:
         return (self,)
 
+    async def observe_live(self, observe: Callable[[], Awaitable[object]]) -> None:
+        """Join an admitted observation before retirement drops session state.
+
+        A shared source may retain this owner across a storage await. Admission
+        belongs to the original owner's stop, not a later session-ID lookup that
+        could create fresh unowned input state after retirement.
+        """
+        if self.stop.done():
+            return
+        completion = asyncio.get_running_loop().create_future()
+        self._observations.add(completion)
+        try:
+            await observe()
+        finally:
+            completion.set_result(None)
+            self._observations.remove(completion)
+
+    async def join_observations(self) -> None:
+        """Join admitted readers after stop closes admission; readers stay parallel."""
+        await join_retirement(asyncio.gather(*self._observations))
+
     def request_stop(self) -> None:
         self.stop.request()
 

src/agent_comms/session_config_options.py

@@ -1,3 +1,3 @@
 """ACP session publication consumes the shared configuration catalogs."""
 
 from __future__ import annotations
@@ -33,6 +33,7 @@
     from .runtime import RuntimeServer
     from .session_effects import SessionEffects
     from .session_lifecycle import SessionLifecycle
+    from .input_drain import OwnerSession
     from .threads import Thread
 
 
@@ -49,7 +50,7 @@
         self.comms, self.runtime, self.sessions, self.effects = comms, runtime, sessions, effects
         self.catalog_publish_lock = asyncio.Lock()
         self.publications: dict[AcpSessionId, ConfigOptionUpdate] = {}
-        self._publication_sessions: set[AcpSessionId] = set()
+        self._observed_owners: dict[AcpSessionId, OwnerSession] = {}
         self._tasks: set[asyncio.Task[None]] = set()
         self.observer = self.runtime.background(self._observe_configuration())
         # The stored settings each session's owner last saw, to notice a change.
@@ -114,7 +115,7 @@
             await self.effects.inputs.state(session_id).owned().configuration.catalog_for(member).change(self, thread, value)
         except (NativePiUnavailable, InvalidConfigChoice) as error:
             raise Refused.invalid_params({"reason": str(error)}) from error
-        await self.effects.turns.close_idle_backend(session_id, thread)
+        await self.effects.turns.close_idle_backend(session_id)
         await self.use_stored_settings(session_id, name)
         return SetSessionConfigOptionResponse(
             config_options=await self.publish_configuration(session_id),
@@ -167,8 +168,9 @@
 
     def start_menu(self, session_id: AcpSessionId, client: SessionUpdateSink) -> asyncio.Task[None]:
         """Retain initial publication until it finishes or its consumer retires."""
-        self._publication_sessions.add(session_id)
-        task = self.runtime.background(self.send_menu(client, session_id))
+        owner = self.effects.inputs.state(session_id).owned()
+        self._observed_owners[session_id] = owner
+        task = self.runtime.background(owner.observe_live(partial(self.send_menu, client, session_id)))
         self._tasks.add(task)
         task.add_done_callback(self._tasks.discard)
         return task
@@ -209,19 +211,23 @@
         )
 
     def stop_observing_configuration(self, session_id: AcpSessionId) -> None:
-        self._publication_sessions.discard(session_id)
+        self._observed_owners.pop(session_id, None)
         self.publications.pop(session_id, None)
 
     async def _observe_configuration(self) -> None:
         async with aclosing(Bus.watch_sources(FileSource(self.comms.registry.store.path))) as observations:
             async for _ in observations:
-                for session_id in tuple(self._publication_sessions):
-                    await self.sessions.sync_identity(session_id)
-                    await self.publish_configuration(session_id)
+                for owner in tuple(self._observed_owners.values()):
+                    await owner.observe_live(partial(self.refresh_configuration, owner.identifier))
+
+    async def refresh_configuration(self, session_id: AcpSessionId) -> None:
+        """Refresh within the original owner's admitted observation lifetime."""
+        await self.sessions.sync_identity(session_id)
+        await self.publish_configuration(session_id)
 
     async def close(self) -> None:
         await self.retire_observations((*self._tasks, self.observer))
-        self._publication_sessions.clear()
+        self._observed_owners.clear()
 
     async def publish_configuration(
         self, session_id: AcpSessionId,

src/agent_comms/session_lifecycle.py

@@ -305,6 +305,7 @@ async def retire_owner(self, state: OwnerSession, expected: OwnerIdentity, claim
             state.request_stop()
             self.runtime.retire_session(state.identifier)
             await self.effects.inputs.close_sessions((state,))
+            await state.join_observations()
             await self.effects.turns.close_sessions((state,))
             await Coordination.run_worker(partial(
                 self.comms.owners.complete_session, state.owner, expected,
@@ -382,7 +383,7 @@ async def sync_identity(self, session_id: AcpSessionId) -> str:
                 or not thread.process_identity.runs_here()
                 or not entry.status.running
             ):
-                await self.effects.turns.close_idle_backend(session_id, thread)
+                await self.effects.turns.close_idle_backend(session_id)
             current = PublishedSession.of(thread)
             if self.published[session_id] != current:
                 await self.publish_metadata(session_id, title=thread.title)

src/agent_comms/turn_runner.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import asyncio
-from contextlib import AsyncExitStack, asynccontextmanager
+from contextlib import AsyncExitStack
 from collections.abc import Awaitable, Callable
 from functools import partial
 from typing import TYPE_CHECKING
@@ -391,30 +391,22 @@ async def active_backend_turn(self, session_id: AcpSessionId) -> DrainTurn:
         thread = await Coordination.run_worker(partial(self.comms.registry.require, name))
         return self.inputs.state(session_id).turn.for_owner(thread)
 
-    async def close_idle_backend(self, session_id: AcpSessionId, thread: Thread) -> None:
-        async with self.idle_backend(session_id, thread) as persistent:
-            if persistent is not None:
-                await persistent.close()
+    async def close_idle_backend(self, session_id: AcpSessionId) -> None:
+        """Retire only the backend already held by this idle session.
 
-    @asynccontextmanager
-    async def idle_backend(self, session_id: AcpSessionId, thread: Thread):
-        """Hold this session's turn and child acquisition until an idle effect ends.
-
-        Configuration retirement remains conditional on an idle owner. Source
-        selection uses this same resource and refuses a busy owner. Neither
-        consumer can close a child after a new turn acquired it.
+        The turn lock excludes new input while the backend's own close_idle
+        serializes child retirement. Cleanup never acquires execution or creates
+        a backend, and needs no thread declaration to close its original child.
         """
+        persistent = self.persistent_backends.get(session_id)
+        if persistent is None:
+            return
         lock = self.inputs.state(session_id).turn_lock
         if lock.locked():
-            yield None
             return
         async with lock:
-            if await Coordination.run_worker(partial(self.session_busy, session_id)):
-                yield None
-                return
-            persistent = self.backend_for(session_id, thread)
-            async with persistent.lock:
-                yield persistent
+            if not await Coordination.run_worker(partial(self.session_busy, session_id)):
+                await persistent.close_idle()
 
     async def run_selected(self, session_id: AcpSessionId, execution: SelectedExecution):
         """Selected work and fresh ACP input share this session's one turn lifetime.

tests/test_goal_publication_sources.py

@@ -2,7 +2,9 @@
 
 import asyncio
 from dataclasses import dataclass, field, replace
+from functools import partial
 import sqlite3
+import sys
 
 from acp.schema import SessionInfoUpdate
 from pydantic import BaseModel
@@ -21,6 +23,7 @@
 from agent_comms.goal_waits import GoalWaits
 from agent_comms.locked_store import GoalWakeFile, RegistryFile
 from agent_comms.repo import Checkout
+from agent_comms.registry_document import RegistryDocument
 from agent_comms.thread_execution import ExternalThreadExecution
 from agent_comms.thread_text import GivenThreadText
 from agent_comms.threads import Thread
@@ -115,7 +118,10 @@ async def test_registry_routing_uses_only_the_shared_document_and_published_goal
     document = comms.registry.store.read()
     original = scheduler.publication(session).snapshot
     reads = comms.registry.store.cache.read_counts().requests
-    await scheduler.observe_root(document, (session,), frozenset((RegistryFile.path_in(comms.root),)))
+    await asyncio.gather(*scheduler.observe_root(
+        document, (agent.inputs.state(session).owned(),),
+        frozenset((RegistryFile.path_in(comms.root),)),
+    ))
     assert comms.registry.store.cache.read_counts().requests == reads
     assert scheduler.publication(session).snapshot is original
     assert len(client.published) == 1
@@ -164,7 +170,10 @@ async def test_attempt_signal_refreshes_complete_projection_but_does_not_publish
     document = comms.registry.store.read()
     before = scheduler.publication(session).snapshot
     reads = comms.registry.store.cache.read_counts().requests
-    await scheduler.observe_root(document, (session,), frozenset(store.source_paths(store.path.parent)))
+    await asyncio.gather(*scheduler.observe_root(
+        document, (agent.inputs.state(session).owned(),),
+        frozenset(store.source_paths(store.path.parent)),
+    ))
     assert comms.registry.store.cache.read_counts().requests > reads
     assert scheduler.publication(session).snapshot is before
     assert len(client.published) == 1
@@ -183,9 +192,117 @@ async def test_failed_goal_read_keeps_cause_and_never_creates_input_diagnostic(p
         db.execute(f'DROP TABLE "{Generation.declared_name}"')
     with pytest.raises(StorageUncertainError):
         await scheduler.sync_goal_execution(session)
-    await scheduler.observe_root(document, (session,), frozenset(store.source_paths(path.parent)))
+    outcomes = await asyncio.gather(*scheduler.observe_root(
+        document, (agent.inputs.state(session).owned(),),
+        frozenset(store.source_paths(path.parent)),
+    ), return_exceptions=True)
+    assert len(outcomes) == 1 and isinstance(outcomes[0], StorageUncertainError)
     assert any(record.exc_info and isinstance(record.exc_info[1], StorageUncertainError)
                for record in caplog.records)
     assert scheduler.publication(session).snapshot is before
     assert len(client.published) == 1
     assert tuple((comms.root / "diagnostics").glob("drain-*.json")) == diagnostics
+
+
+async def test_idle_retirement_never_acquires_an_unloaded_session(publication):
+    _comms, agent, _session, _client = publication
+    missing = AcpSessionId("never-loaded")
+    before = dict(agent.inputs.drained)
+    await agent.turns.close_idle_backend(missing)
+    assert agent.inputs.drained == before
+    assert missing not in agent.turns.persistent_backends
+
+
+async def test_idle_retirement_uses_the_backend_contract_and_respects_a_held_turn(publication):
+    from agent_comms.native_child import PiChild
+    from agent_comms.pi_durable_backend import PiDurableBackend
+
+    comms, agent, session, _client = publication
+    backend = PiDurableBackend()
+    child = await PiChild.start_command(
+        (sys.executable, "-c", "import sys; sys.stdin.buffer.read()"), cwd=str(comms.root),
+    )
+    backend.child = child
+    agent.turns.persistent_backends[session] = backend
+    try:
+        async with agent.inputs.state(session).turn_lock:
+            await agent.turns.close_idle_backend(session)
+            assert backend.child is child
+            assert child.proc.alive()
+        await agent.turns.close_idle_backend(session)
+        assert backend.child is None
+        assert child.retired
+    finally:
+        await backend.close_idle()
+
+
+async def test_observation_joins_before_retirement_and_captured_sources_cannot_reopen(publication):
+    comms, agent, session, client = publication
+    owner = agent.inputs.state(session).owned()
+    document = comms.registry.store.read()
+    entered, release = (asyncio.Event(), asyncio.Event()), asyncio.Event()
+
+    async def observation(started):
+        started.set()
+        await release.wait()
+        assert agent.sessions.require(session) == session
+        assert agent.inputs.state(session) is owner
+
+    observing = tuple(asyncio.create_task(owner.observe_live(partial(observation, started)))
+                      for started in entered)
+    stopping = None
+    try:
+        await asyncio.wait_for(asyncio.gather(*(started.wait() for started in entered)), 3)
+        fenced = comms.owners.fence_stop(session)
+        request, = comms.owners.stop_requests(fenced)
+        stopping = asyncio.create_task(request.retire(agent.sessions, owner))
+        await asyncio.sleep(0)
+        assert not stopping.done()
+        release.set()
+        await asyncio.wait_for(asyncio.gather(*observing, stopping), 8)
+        assert owner.retired.done()
+        assert session not in agent.sessions.published
+        assert session not in agent.inputs.drained
+        before = len(client.published)
+        await owner.observe_live(lambda: agent.sessions.config.refresh_configuration(session))
+        for source in (document, RegistryDocument()):
+            await asyncio.gather(*agent.turns.goals.observe_root(
+                source, (owner,), frozenset((RegistryFile.path_in(comms.root),)),
+            ))
+        assert len(client.published) == before
+        assert session not in agent.inputs.drained
+        assert session not in agent.turns.persistent_backends
+    finally:
+        release.set()
+        await asyncio.gather(*(task for task in (*observing, stopping) if task is not None),
+                             return_exceptions=True)
+
+
+async def test_queued_initial_menu_obeys_the_original_owner_stop(publication):
+    _comms, agent, session, client = publication
+    owner = agent.inputs.state(session).owned()
+    menu = agent.sessions.config.start_menu(session, client)
+    owner.request_stop()
+    await asyncio.wait_for(menu, 3)
+    assert not agent.turns.persistent_backends
+
+
+async def test_cancelled_observation_releases_retirement_admission(publication):
+    _comms, agent, session, _client = publication
+    owner = agent.inputs.state(session).owned()
+    entered = asyncio.Event()
+
+    async def observation():
+        entered.set()
+        await asyncio.Future()
+
+    observing = asyncio.create_task(owner.observe_live(observation))
+    try:
+        await asyncio.wait_for(entered.wait(), 3)
+        observing.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await observing
+        await asyncio.wait_for(owner.join_observations(), 3)
+    finally:
+        observing.cancel()
+        await asyncio.gather(observing, return_exceptions=True)