From a5b5112abc5bfeaad376b9f8552c718ee17dd938 Mon Sep 17 00:00:00 2001 From: Tristan Simas Date: Sun, 1 Nov 2026 01:17:30 -0500 Subject: [PATCH 1/2] Keep session observations inside their original owner lifetime --- src/agent_comms/goal_scheduler.py | 29 +++-- src/agent_comms/input_drain.py | 13 ++- src/agent_comms/session_config_options.py | 24 +++-- src/agent_comms/session_lifecycle.py | 2 +- src/agent_comms/turn_runner.py | 30 ++---- tests/test_goal_publication_sources.py | 124 +++++++++++++++++++++- 6 files changed, 178 insertions(+), 44 deletions(-) diff --git a/src/agent_comms/goal_scheduler.py b/src/agent_comms/goal_scheduler.py index 55059d7922..73cc8d3f37 100644 --- a/src/agent_comms/goal_scheduler.py +++ b/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) diff --git a/src/agent_comms/input_drain.py b/src/agent_comms/input_drain.py index 03b996c086..5118a214b4 100644 --- a/src/agent_comms/input_drain.py +++ b/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 @@ -752,6 +752,17 @@ 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. + """ + async with self.retirement_lock: + if not self.stop.done(): + await observe() + def request_stop(self) -> None: self.stop.request() diff --git a/src/agent_comms/session_config_options.py b/src/agent_comms/session_config_options.py index 37ca1f8180..f781b884dd 100644 --- a/src/agent_comms/session_config_options.py +++ b/src/agent_comms/session_config_options.py @@ -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 @@ def __init__( 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 @@ async def set_option( 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 @@ async def apply_stored_settings(self, session_id: AcpSessionId) -> None: 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 @@ async def options_for_session( ) 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, diff --git a/src/agent_comms/session_lifecycle.py b/src/agent_comms/session_lifecycle.py index 33fcbbbeeb..2e30f92e87 100644 --- a/src/agent_comms/session_lifecycle.py +++ b/src/agent_comms/session_lifecycle.py @@ -382,7 +382,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) diff --git a/src/agent_comms/turn_runner.py b/src/agent_comms/turn_runner.py index a057f3cf7c..da9b7706d6 100644 --- a/src/agent_comms/turn_runner.py +++ b/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. diff --git a/tests/test_goal_publication_sources.py b/tests/test_goal_publication_sources.py index e3134dd7e9..ac399b68de 100644 --- a/tests/test_goal_publication_sources.py +++ b/tests/test_goal_publication_sources.py @@ -3,6 +3,7 @@ import asyncio from dataclasses import dataclass, field, replace import sqlite3 +import sys from acp.schema import SessionInfoUpdate from pydantic import BaseModel @@ -21,6 +22,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 +117,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 +169,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 +191,119 @@ 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() + + async def observation(): + entered.set() + await release.wait() + assert agent.sessions.require(session) == session + assert agent.inputs.state(session) is owner + + observing = asyncio.create_task(owner.observe_live(observation)) + stopping = None + try: + await asyncio.wait_for(entered.wait(), 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() + async with owner.retirement_lock: + menu = agent.sessions.config.start_menu(session, client) + owner.request_stop() + await asyncio.sleep(0) + assert not menu.done() + 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 + assert not owner.retirement_lock.locked() + finally: + observing.cancel() + await asyncio.gather(observing, return_exceptions=True) From 1a462d05a78f719262c8a40cc5ff9ab4528a7a9c Mon Sep 17 00:00:00 2001 From: Tristan Simas Date: Sun, 1 Nov 2026 01:37:32 -0500 Subject: [PATCH 2/2] Join admitted observations without serializing independent readers --- src/agent_comms/input_drain.py | 17 ++++++++++++++--- src/agent_comms/session_lifecycle.py | 1 + tests/test_goal_publication_sources.py | 25 ++++++++++++------------- 3 files changed, 27 insertions(+), 16 deletions(-) diff --git a/src/agent_comms/input_drain.py b/src/agent_comms/input_drain.py index 5118a214b4..00dec10e16 100644 --- a/src/agent_comms/input_drain.py +++ b/src/agent_comms/input_drain.py @@ -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 @@ -759,9 +760,19 @@ async def observe_live(self, observe: Callable[[], Awaitable[object]]) -> None: belongs to the original owner's stop, not a later session-ID lookup that could create fresh unowned input state after retirement. """ - async with self.retirement_lock: - if not self.stop.done(): - await observe() + 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() diff --git a/src/agent_comms/session_lifecycle.py b/src/agent_comms/session_lifecycle.py index 2e30f92e87..6872cfd112 100644 --- a/src/agent_comms/session_lifecycle.py +++ b/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, diff --git a/tests/test_goal_publication_sources.py b/tests/test_goal_publication_sources.py index ac399b68de..d7fe6fcc67 100644 --- a/tests/test_goal_publication_sources.py +++ b/tests/test_goal_publication_sources.py @@ -2,6 +2,7 @@ import asyncio from dataclasses import dataclass, field, replace +from functools import partial import sqlite3 import sys @@ -239,25 +240,26 @@ async def test_observation_joins_before_retirement_and_captured_sources_cannot_r comms, agent, session, client = publication owner = agent.inputs.state(session).owned() document = comms.registry.store.read() - entered, release = asyncio.Event(), asyncio.Event() + entered, release = (asyncio.Event(), asyncio.Event()), asyncio.Event() - async def observation(): - entered.set() + async def observation(started): + started.set() await release.wait() assert agent.sessions.require(session) == session assert agent.inputs.state(session) is owner - observing = asyncio.create_task(owner.observe_live(observation)) + observing = tuple(asyncio.create_task(owner.observe_live(partial(observation, started))) + for started in entered) stopping = None try: - await asyncio.wait_for(entered.wait(), 3) + 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) + 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 @@ -272,18 +274,15 @@ async def observation(): 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), + 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() - async with owner.retirement_lock: - menu = agent.sessions.config.start_menu(session, client) - owner.request_stop() - await asyncio.sleep(0) - assert not menu.done() + menu = agent.sessions.config.start_menu(session, client) + owner.request_stop() await asyncio.wait_for(menu, 3) assert not agent.turns.persistent_backends @@ -303,7 +302,7 @@ async def observation(): observing.cancel() with pytest.raises(asyncio.CancelledError): await observing - assert not owner.retirement_lock.locked() + await asyncio.wait_for(owner.join_observations(), 3) finally: observing.cancel() await asyncio.gather(observing, return_exceptions=True)