From 8e28c53252d4b70730751415b36c8d00e70958a6 Mon Sep 17 00:00:00 2001 From: Tristan Simas Date: Fri, 23 Oct 2026 03:35:58 -0400 Subject: [PATCH 1/4] Give scheduled turn absence an owner and use child writes --- src/agent_comms/backend.py | 62 +++++++++++++---------------- src/agent_comms/input_drain.py | 4 +- src/agent_comms/messages.py | 21 +++++++++- src/agent_comms/prompt_state.py | 2 +- src/agent_comms/routing.py | 65 +++++++++++++++++++++++++------ src/agent_comms/schedule_rules.py | 2 +- src/agent_comms/send_outcome.py | 2 +- src/agent_comms/turn_inputs.py | 10 ++--- 8 files changed, 111 insertions(+), 57 deletions(-) diff --git a/src/agent_comms/backend.py b/src/agent_comms/backend.py index 6437186226..beb5adf2ee 100644 --- a/src/agent_comms/backend.py +++ b/src/agent_comms/backend.py @@ -395,9 +395,7 @@ def committable_message(self, text): return self.started_input and self.accepts_output and self.output.matches_message(text) def ensure_input_forwarding(self): - if self.native.proc.stdin is not None: - self.stdin = self.native.proc.stdin - self.native.forward_input(partial(self.inputs.forward, self)) + self.native.forward_input(partial(self.inputs.forward, self)) def observe_input_during_abort(self, event): if self.accepts_output: @@ -413,32 +411,30 @@ def permits_live_status(self) -> bool: async def abort_stalled_rpc(self) -> None: await self.native.finish_input() - if self.native.proc.stdin is not None and self.native.proc.returncode is None: - abort_deadline = self.loop.time() + max(0.0, self.rpc_abort_grace) - try: - self.native.proc.stdin.write( - self.native.reader.encode(commands.Abort(id=RequestId("agent-comms-watchdog"))) + abort_deadline = self.loop.time() + max(0.0, self.rpc_abort_grace) + try: + remaining = max(0.0, abort_deadline - self.loop.time()) + await asyncio.wait_for(self.native.proc.write( + self.native.reader.encode(commands.Abort(id=RequestId("agent-comms-watchdog"))) + ), timeout=remaining) + while remaining := max(0.0, abort_deadline - self.loop.time()): + response = await asyncio.wait_for( + self.native.reader.readline(), timeout=remaining ) - remaining = max(0.0, abort_deadline - self.loop.time()) - await asyncio.wait_for(self.native.proc.stdin.drain(), timeout=remaining) - while remaining := max(0.0, abort_deadline - self.loop.time()): - response = await asyncio.wait_for( - self.native.reader.readline(), timeout=remaining - ) - if not response: - break - try: - abort_payload = PiRpcChannel.decode_record(response) - except (json.JSONDecodeError, UnicodeDecodeError, MalformedRecord): - continue - abort_payload.observe_abort(self) - if ( - isinstance(abort_payload, pi.Response) - and abort_payload.command is commands.Abort - ): - break - except TimeoutError: - pass + if not response: + break + try: + abort_payload = PiRpcChannel.decode_record(response) + except (json.JSONDecodeError, UnicodeDecodeError, MalformedRecord): + continue + abort_payload.observe_abort(self) + if ( + isinstance(abort_payload, pi.Response) + and abort_payload.command is commands.Abort + ): + break + except (TimeoutError, BrokenPipeError): + pass await self.native.shutdown() def context_info(self) -> events.ContextInfo: @@ -490,12 +486,10 @@ async def run(self) -> AsyncGenerator[events.AgentEvent, None]: ) async with self.native.failures(): self.output.sensitive |= self.native.sensitive_diagnostics - if self.native.proc.stdin is not None: - try: - self.native.proc.stdin.write(self.native.reader.encode(self.native.attestation.request)) - await self.native.proc.stdin.drain() - except (BrokenPipeError, ConnectionResetError): - pass + try: + await self.native.proc.write(self.native.reader.encode(self.native.attestation.request)) + except (BrokenPipeError, ConnectionResetError): + pass await self.initialize_rpc() while True: self.skip = False diff --git a/src/agent_comms/input_drain.py b/src/agent_comms/input_drain.py index 2f9863fa66..34ad4a916d 100644 --- a/src/agent_comms/input_drain.py +++ b/src/agent_comms/input_drain.py @@ -638,9 +638,9 @@ async def delivery_scope(self, session_id: str) -> InputDeliveryScope: return UNOBSERVED_DELIVERY_SCOPE keys = set(self.input_keys(session_id)) keys.update( - self.dispositions.bus_key(turn.origin, owner) + key for turn in self.state(session_id).pending_turns - if turn.origin is not None + for key in turn.origin.bus_keys(self.dispositions, owner) ) return OwnerQueueInputDeliveryScope(frozenset(keys)) diff --git a/src/agent_comms/messages.py b/src/agent_comms/messages.py index b7722c205d..d781cd574f 100644 --- a/src/agent_comms/messages.py +++ b/src/agent_comms/messages.py @@ -37,12 +37,14 @@ from .mentions import ThreadMention from .message_reference import MessageReference from .page_traversal import LiveCursor, PageCursor +from .routing import TurnOrigin from .task_sources import Decision, NoTaskAttachment, TaskAttachment from .thread_identity import AgentThreadRole, ThreadRole from .ui_model.message_category import InboundCategory, MessageCategory, OtherCategory from .wake_policy import BoundedTriageWake, PassiveWake if TYPE_CHECKING: + from .input_disposition import InputDispositions from .historical_views import HistorySource from .retained_task_facts import ExactTaskFact from .threads import Thread @@ -361,7 +363,7 @@ class HistoricalMembershipNoticeRole(HistoricalRole, MembershipNoticeRole, decla @dataclass(frozen=True, slots=True) -class Message(Addressed, BusCapability): +class Message(Addressed, BusCapability, TurnOrigin): """Declares one inter-thread message. The bus assigns ``seq`` (a monotonically increasing log position) at send @@ -560,6 +562,23 @@ def exact_reply_target(self) -> str: """The exact response destination owned by this opening message.""" return CommandTarget.named(self.target).reply_target(self.sender) + autonomous: ClassVar[bool] = False + + @property + def separate_turn(self) -> bool: + return self.role.response_policy(self).separate_turn + + @property + def batch_target(self) -> str: + return self.exact_reply_target + + @property + def messages(self) -> tuple[Message, ...]: + return (self,) + + def bus_keys(self, dispositions: InputDispositions, owner: Thread) -> tuple[str, ...]: + return (dispositions.bus_key(self, owner),) + @property def reply_target(self) -> str | None: """An external sender may receive an automatic live reply.""" diff --git a/src/agent_comms/prompt_state.py b/src/agent_comms/prompt_state.py index 1d5e517b95..329043cb30 100644 --- a/src/agent_comms/prompt_state.py +++ b/src/agent_comms/prompt_state.py @@ -67,7 +67,7 @@ def dispatch(self): async def deliver_to(self, inputs, session): def write(): session.grant_prompt() - session.stdin.write(session.native.reader.encode(session.command)) + session.native.proc.stdin.write(session.native.reader.encode(session.command)) async with MaintenanceBarrier(RegistryFile.path_in(session.startup.root)).ingress_async() as status: with status.phase.send_new(session.send_boundary, None, session.command.input_id, session.command.message) as outcome: completion = outcome.dispatch(inputs, session, self, write) diff --git a/src/agent_comms/routing.py b/src/agent_comms/routing.py index cd9d40cd19..65ab7915ea 100644 --- a/src/agent_comms/routing.py +++ b/src/agent_comms/routing.py @@ -2,21 +2,67 @@ from __future__ import annotations +from abc import ABC, abstractmethod from collections.abc import Callable, Iterator, Mapping from dataclasses import dataclass from typing import Annotated, TYPE_CHECKING from .channel_targets import CommandTarget -from .field_codec import FieldCodec, WireOmitDefault +from .field_codec import FieldCodec, NullMember, WireOmitDefault from .message_reference import MessageReference -from .messages import Message from .read_basis import DMDisplayScope, MessageDisplayScope, ReadCursor if TYPE_CHECKING: from .catalog_document import CatalogDocument from .registry_provenance import RegistryProvenance from .thread_identity import ThreadIncarnation + from .input_disposition import InputDispositions + from .messages import Message + from .threads import Thread + + +class TurnOrigin(ABC): + """The source of a scheduled turn owns batching and delivery identity.""" + + __slots__ = () + + @property + @abstractmethod + def separate_turn(self) -> bool: ... + + @property + @abstractmethod + def batch_target(self) -> object: ... + + @property + @abstractmethod + def autonomous(self) -> bool: ... + + @property + @abstractmethod + def messages(self) -> tuple[Message, ...]: ... + + @abstractmethod + def bus_keys(self, dispositions: InputDispositions, owner: Thread) -> tuple[str, ...]: ... + + +class NoTurnOrigin(NullMember, TurnOrigin): + """An autonomous turn has no message, delivery key or reply destination.""" + + separate_turn = False + autonomous = True + messages = () + + @property + def batch_target(self) -> NoTurnOrigin: + return self + + def bus_keys(self, dispositions: InputDispositions, owner: Thread) -> tuple[str, ...]: + return () + + +NO_TURN_ORIGIN = NoTurnOrigin() @dataclass(frozen=True, slots=True) @@ -88,13 +134,13 @@ def require_publications(self, originals, owner_lookup: str) -> None: @dataclass(frozen=True, slots=True) class ScheduledTurn: prompt: str - origin: Message | None = None + origin: TurnOrigin = NO_TURN_ORIGIN goal_id: str | None = None goal_wait_id: str | None = None @property def autonomous_goal(self) -> bool: - return self.goal_id is not None and self.origin is None + return self.goal_id is not None and self.origin.autonomous @classmethod def resuming_in_project(cls, project: str) -> ScheduledTurn: @@ -110,10 +156,6 @@ def require_dependency_wake(self, error: Exception) -> None: if self.goal_wait_id is None: raise error - @property - def origins(self) -> tuple[Message, ...]: - return (self.origin,) if self.origin is not None else () - @classmethod def incoming( cls, message: Message, *, aliases: Mapping[str, str] | None = None @@ -131,15 +173,14 @@ def take_batch(pending: list[ScheduledTurn]) -> tuple[list[ScheduledTurn], list[ """Combine channel turns sharing a route, keeping each direct input separate.""" if not pending: return [], [] - if pending[0].origin and pending[0].origin.role.response_policy(pending[0].origin).separate_turn: + if pending[0].origin.separate_turn: return pending[:1], pending[1:] boundary = next( ( index for index, turn in enumerate(pending) - if (turn.origin.exact_reply_target if turn.origin else None) - != (pending[0].origin.exact_reply_target if pending[0].origin else None) - or (turn.origin and turn.origin.role.response_policy(turn.origin).separate_turn) + if (turn.origin.batch_target != pending[0].origin.batch_target + or turn.origin.separate_turn) ), len(pending), ) diff --git a/src/agent_comms/schedule_rules.py b/src/agent_comms/schedule_rules.py index f8e84ea247..a131c67f4d 100644 --- a/src/agent_comms/schedule_rules.py +++ b/src/agent_comms/schedule_rules.py @@ -105,7 +105,7 @@ async def run(self) -> None: session_id, inputs.sessions.require(session_id), "\n\n".join(turn.prompt for turn in pending), - origins=tuple(origin for turn in pending for origin in turn.origins), + origins=tuple(origin for turn in pending for origin in turn.origin.messages), autonomous_goal=len(pending) == 1 and pending[0].autonomous_goal, dependency_wait_id=pending[0].goal_wait_id, ) diff --git a/src/agent_comms/send_outcome.py b/src/agent_comms/send_outcome.py index f3b35905f7..d9c8a5011b 100644 --- a/src/agent_comms/send_outcome.py +++ b/src/agent_comms/send_outcome.py @@ -24,7 +24,7 @@ def dispatch(self, inputs, session, item, write): return self._drain(session) async def _drain(self, session): - await session.stdin.drain() + await session.native.proc.stdin.drain() return True diff --git a/src/agent_comms/turn_inputs.py b/src/agent_comms/turn_inputs.py index 31c32485c7..3012286eed 100644 --- a/src/agent_comms/turn_inputs.py +++ b/src/agent_comms/turn_inputs.py @@ -96,8 +96,8 @@ async def forward(self, session: TurnSession) -> None: async def write(self, session: TurnSession, *written: commands.PiCommand) -> None: for command in written: - session.stdin.write(session.native.reader.encode(command)) - await session.stdin.drain() + session.native.proc.stdin.write(session.native.reader.encode(command)) + await session.native.proc.stdin.drain() async def queued(self, input_ids: tuple[InputId, ...]) -> tuple[str, ...]: """A native prompt acknowledgement proves it reached the child's queue.""" @@ -117,7 +117,7 @@ async def clear_queue(self, session: TurnSession, item: ClearQueue) -> bool: command = commands.ClearQueue(input_ids=await self.queued(item.input_ids)) if not command.input_ids: return True - async with command.pending_response(session.native.reader, session.stdin) as future: + async with command.pending_response(session.native.reader, session.native.proc.stdin) as future: response = await future if not response.succeeded: raise RuntimeError( @@ -151,7 +151,7 @@ async def send_now(self, session: TurnSession, item: SendNow) -> bool: with status.phase.send_running( session.interrupt_boundary, first.require_id(), candidates[0], first.message, ) as outcome: - completion = outcome.dispatch(self, session, item, lambda: session.stdin.write( + completion = outcome.dispatch(self, session, item, lambda: session.native.proc.stdin.write( session.native.reader.encode(commands.InterruptSteering(input_ids=list(candidates))) )) return await completion @@ -172,7 +172,7 @@ async def place(self, session: TurnSession, item: InputRequest) -> bool: def write(): if command.images: session.output.sensitive = True - session.stdin.write(session.native.reader.encode(command)) + session.native.proc.stdin.write(session.native.reader.encode(command)) self.pending[native_id] = command self.generation += 1 From a32bd915e1d0b6ee0f71c8d90abf63a31bc14bf8 Mon Sep 17 00:00:00 2001 From: Tristan Simas Date: Fri, 23 Oct 2026 03:38:39 -0400 Subject: [PATCH 2/4] Check autonomous origin boundaries and delivery keys --- tests/test_response_policy.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/tests/test_response_policy.py b/tests/test_response_policy.py index fa657f329e..7e736f847d 100644 --- a/tests/test_response_policy.py +++ b/tests/test_response_policy.py @@ -40,3 +40,25 @@ def response_policy(self, message): finally: del ResponsePolicy.__registry__["review"] + + +def test_autonomous_turns_keep_message_batches_and_delivery_keys_separate(tmp_path, thread_registration): + from agent_comms.input_disposition import InputDispositions + + goal = ScheduledTurn("Continue the goal", goal_id="goal") + resume = ScheduledTurn.resuming_in_project(str(tmp_path)) + first = ScheduledTurn.incoming(Message("peer", "#team", "first", InfoKind, seq=1)) + second = ScheduledTurn.incoming(Message("peer", "#team", "second", InfoKind, seq=2)) + direct = ScheduledTurn.incoming(Message("peer", "owner", "direct", InfoKind, seq=3)) + + assert goal.autonomous_goal and not resume.autonomous_goal + assert ScheduledTurn.take_batch([goal, resume, first]) == ([goal, resume], [first]) + assert ScheduledTurn.take_batch([first, second, goal]) == ([first, second], [goal]) + assert ScheduledTurn.take_batch([direct, direct]) == ([direct], [direct]) + + owner = thread_registration.new_thread("owner", frozenset(), str(tmp_path)) + dispositions = InputDispositions(tmp_path) + assert goal.origin.messages == () + assert goal.origin.bus_keys(dispositions, owner) == () + assert first.origin.messages == (first.origin,) + assert first.origin.bus_keys(dispositions, owner) == (dispositions.bus_key(first.origin, owner),) From 325dd9557aec21712d2efdcd3fa15e543ff884d6 Mon Sep 17 00:00:00 2001 From: Tristan Simas Date: Fri, 23 Oct 2026 03:44:40 -0400 Subject: [PATCH 3/4] Centralize Pi sends on the child and preserve message enumeration --- src/agent_comms/backend.py | 8 ++++---- src/agent_comms/child_process.py | 17 +++++++++++++++-- src/agent_comms/messages.py | 4 ---- src/agent_comms/native_child.py | 16 +++++++++++++--- src/agent_comms/pi_commands.py | 25 ++++++++++++++----------- src/agent_comms/pi_rpc.py | 6 +++--- src/agent_comms/private_send_claim.py | 8 +++----- src/agent_comms/prompt_state.py | 2 +- src/agent_comms/routing.py | 5 +++-- src/agent_comms/schedule_rules.py | 2 +- src/agent_comms/send_outcome.py | 2 +- src/agent_comms/tracked_turn.py | 9 ++++----- src/agent_comms/turn_inputs.py | 12 ++++++------ src/agent_comms/turn_stats.py | 6 +++--- tests/test_response_policy.py | 4 ++-- 15 files changed, 73 insertions(+), 53 deletions(-) diff --git a/src/agent_comms/backend.py b/src/agent_comms/backend.py index beb5adf2ee..e4d932e2f0 100644 --- a/src/agent_comms/backend.py +++ b/src/agent_comms/backend.py @@ -414,9 +414,8 @@ async def abort_stalled_rpc(self) -> None: abort_deadline = self.loop.time() + max(0.0, self.rpc_abort_grace) try: remaining = max(0.0, abort_deadline - self.loop.time()) - await asyncio.wait_for(self.native.proc.write( - self.native.reader.encode(commands.Abort(id=RequestId("agent-comms-watchdog"))) - ), timeout=remaining) + self.native.send(commands.Abort(id=RequestId("agent-comms-watchdog"))) + await asyncio.wait_for(self.native.drain(), timeout=remaining) while remaining := max(0.0, abort_deadline - self.loop.time()): response = await asyncio.wait_for( self.native.reader.readline(), timeout=remaining @@ -487,7 +486,8 @@ async def run(self) -> AsyncGenerator[events.AgentEvent, None]: async with self.native.failures(): self.output.sensitive |= self.native.sensitive_diagnostics try: - await self.native.proc.write(self.native.reader.encode(self.native.attestation.request)) + self.native.send(self.native.attestation.request) + await self.native.drain() except (BrokenPipeError, ConnectionResetError): pass await self.initialize_rpc() diff --git a/src/agent_comms/child_process.py b/src/agent_comms/child_process.py index 0ce62d27d0..4597e0d852 100644 --- a/src/agent_comms/child_process.py +++ b/src/agent_comms/child_process.py @@ -1525,13 +1525,26 @@ async def start( raise return child - async def write(self, data: bytes) -> None: - """Write to this still-running child's input and observe transport backpressure.""" + def send(self, data: bytes) -> None: + """Write synchronously while this child's input still exists and it runs.""" if self.stdin is None or self.returncode is not None: raise BrokenPipeError self.stdin.write(data) + + async def drain(self) -> None: await self.stdin.drain() + async def write(self, data: bytes) -> None: + """Write to this still-running child's input and observe transport backpressure.""" + self.send(data) + await self.drain() + + async def send_fenced(self, data: bytes, boundary, *, measurements) -> None: + """Keep the independent raw writer on this child's original input.""" + from .native_prompt_send import send_fenced_prompt + + await send_fenced_prompt(self.stdin, data, boundary, measurements=measurements) + @property def returncode(self) -> int | None: return self.process.returncode diff --git a/src/agent_comms/messages.py b/src/agent_comms/messages.py index d781cd574f..756c1c90e0 100644 --- a/src/agent_comms/messages.py +++ b/src/agent_comms/messages.py @@ -572,10 +572,6 @@ def separate_turn(self) -> bool: def batch_target(self) -> str: return self.exact_reply_target - @property - def messages(self) -> tuple[Message, ...]: - return (self,) - def bus_keys(self, dispositions: InputDispositions, owner: Thread) -> tuple[str, ...]: return (dispositions.bus_key(self, owner),) diff --git a/src/agent_comms/native_child.py b/src/agent_comms/native_child.py index 109d9950fe..d1622a51d4 100644 --- a/src/agent_comms/native_child.py +++ b/src/agent_comms/native_child.py @@ -91,8 +91,18 @@ class PiChild: forwarding: InputForwarder = field(default_factory=NoInputForwarder, init=False, repr=False) _close_job: asyncio.Task[ChildOutcome | BaseException] | None = field(default=None, init=False, repr=False) + def send(self, command) -> None: + self.proc.send(self.reader.encode(command)) + + async def drain(self) -> None: + await self.proc.drain() + + async def send_fenced(self, command, boundary, *, measurements) -> None: + await self.proc.send_fenced(self.reader.encode(command), boundary, measurements=measurements) + async def reply_ui(self, response) -> None: - await asyncio.wait_for(self.proc.write(self.reader.encode(response)), timeout=2) + self.send(response) + await asyncio.wait_for(self.drain(), timeout=2) @classmethod async def start_command(cls, argv, *, cwd, env=INHERITED_ENVIRONMENT, pass_fds=(), @@ -104,7 +114,7 @@ async def request(self, command): if self.retired: raise RuntimeError('Pi child has retired') async with self.failures(): - response = await self.reader.request(command, self.proc.stdin) + response = await self.reader.request(command, self) if not response.succeeded: raise RuntimeError(response.error.value_or("Pi gave no error text")) return response.data.require_payload() @@ -302,7 +312,7 @@ class BorrowedNative(SuccessorNative): async def inspect(self, persistent, prepare, request): # The active TurnSession owns receive/correlation. Borrow its original # pending response instead of starting a competing reader. - async with request.pending_response(self.child.reader, self.child.proc.stdin) as response: + async with request.pending_response(self.child) as response: return (await response).data.require_payload() @dataclass diff --git a/src/agent_comms/pi_commands.py b/src/agent_comms/pi_commands.py index 867b3f2d0b..8748e8026e 100644 --- a/src/agent_comms/pi_commands.py +++ b/src/agent_comms/pi_commands.py @@ -86,19 +86,21 @@ def invalidates_identity(cls, response: Response, session: TurnSession) -> bool: return False @asynccontextmanager - async def pending_response(self, channel, writer): + async def pending_response(self, child): + channel = child.reader future = channel.track(self) try: - writer.write(channel.command_bytes(self)) - await writer.drain() + child.send(self) + await child.drain() yield future finally: self.id.release_from(channel.pending, type(self)) future.cancel() - async def exchange(self, channel, writer, *, strict=False, max_bytes=None, on_event=lambda event: None): + async def exchange(self, child, *, strict=False, max_bytes=None, on_event=lambda event: None): from .pi_events import Response - async with self.pending_response(channel, writer) as future: + channel = child.reader + async with self.pending_response(child) as future: while not future.done(): event = await channel.receive(strict=strict, max_bytes=max_bytes) if event is None: @@ -261,6 +263,7 @@ async def catalog( from .child_process import BoundedRun from .native_pi import NativePiRpcLaunch from .pi_rpc import PiRpcChannel + from .native_child import PiChild try: launch = await asyncio.to_thread( @@ -274,7 +277,7 @@ async def catalog( async with BoundedRun.session( launch.argv, cwd=launch.cwd, env=launch.env, timeout=CATALOG_TIMEOUT ) as child: - yield PiRpcChannel(child.stdout), child.stdin + yield PiChild(child, PiRpcChannel(child.stdout)) except TimeoutError as error: raise NativePiUnavailable( f"Native catalog discovery timed out: no catalog reply from {agent_bin!r} " @@ -287,8 +290,8 @@ async def discover( self, agent_bin: str, arguments: Sequence[str], root: Path | str | None = None, *, package: VerifiedNativePackage, ) -> PiResponseData: - async with self.catalog(agent_bin, arguments, root=root, package=package) as (channel, writer): - response = await self.exchange(channel, writer) + async with self.catalog(agent_bin, arguments, root=root, package=package) as child: + response = await self.exchange(child) if response.succeeded: return response.data.require_payload() raise NativePiUnavailable(response.error.text_or("Native catalog query was refused")) @@ -342,11 +345,11 @@ async def configure( semantics cannot change a running turn, a saved transcript or its budget. """ catalog = self.catalog(agent_bin, arguments, package=package, worktree=worktree, root=root) - async with catalog as (channel, writer): - response = await self.exchange(channel, writer) + async with catalog as child: + response = await self.exchange(child) if not response.succeeded: raise NativePiUnavailable(response.error.text_or(self.error_message)) - state = await GetState().exchange(channel, writer) + state = await GetState().exchange(child) if not state.succeeded: raise NativePiUnavailable( state.error.text_or("Native configuration was not reported") diff --git a/src/agent_comms/pi_rpc.py b/src/agent_comms/pi_rpc.py index 0d24a2f4a2..8484fd44fb 100644 --- a/src/agent_comms/pi_rpc.py +++ b/src/agent_comms/pi_rpc.py @@ -71,10 +71,10 @@ async def next_event(self): except asyncio.QueueShutDown: return reading.result() - async def request(self, command, writer): + async def request(self, command, child): if self._read_task is None: - return await command.exchange(self, writer, strict=True, on_event=self.event_type.unowned_event) - async with command.pending_response(self, writer) as future: + return await command.exchange(child, strict=True, on_event=self.event_type.unowned_event) + async with command.pending_response(child) as future: await asyncio.wait((future, self._read_task), return_when=asyncio.FIRST_COMPLETED) response = future.result() if future.done() else self._read_task.result() if response is None: diff --git a/src/agent_comms/private_send_claim.py b/src/agent_comms/private_send_claim.py index 0adc2232e9..c137121385 100644 --- a/src/agent_comms/private_send_claim.py +++ b/src/agent_comms/private_send_claim.py @@ -30,7 +30,7 @@ from .native_entries import NativeEvidenceScope from .native_pi import NativeContextProof, NativePiTerminalFailure, NativeTurnResult from .native_prompt_binding import bind_expected_prompt -from .native_prompt_send import PromptBoundary, send_fenced_prompt +from .native_prompt_send import PromptBoundary from .native_session_reopen import NativeSessionIdentity from .native_startup import NativeStartup from .native_tool_call import NativeToolMode @@ -271,11 +271,9 @@ def _start_once(self, identity: NativeSessionIdentity) -> None: # all of these operations have finished successfully. async def send_prompt(self, turn: TrackedTurnSession) -> None: - stdin = turn.native.proc.stdin - assert stdin is not None with turn.startup.measurements.operation("prompt_writer_join"): - await send_fenced_prompt( - stdin, turn.native.reader.encode(turn.command), + await turn.native.send_fenced( + turn.command, partial(self.prompt_boundary, turn), measurements=turn.startup.measurements, ) diff --git a/src/agent_comms/prompt_state.py b/src/agent_comms/prompt_state.py index 329043cb30..4248556a65 100644 --- a/src/agent_comms/prompt_state.py +++ b/src/agent_comms/prompt_state.py @@ -67,7 +67,7 @@ def dispatch(self): async def deliver_to(self, inputs, session): def write(): session.grant_prompt() - session.native.proc.stdin.write(session.native.reader.encode(session.command)) + session.native.send(session.command) async with MaintenanceBarrier(RegistryFile.path_in(session.startup.root)).ingress_async() as status: with status.phase.send_new(session.send_boundary, None, session.command.input_id, session.command.message) as outcome: completion = outcome.dispatch(inputs, session, self, write) diff --git a/src/agent_comms/routing.py b/src/agent_comms/routing.py index 65ab7915ea..88ac5758fe 100644 --- a/src/agent_comms/routing.py +++ b/src/agent_comms/routing.py @@ -39,7 +39,6 @@ def batch_target(self) -> object: ... @abstractmethod def autonomous(self) -> bool: ... - @property @abstractmethod def messages(self) -> tuple[Message, ...]: ... @@ -52,7 +51,9 @@ class NoTurnOrigin(NullMember, TurnOrigin): separate_turn = False autonomous = True - messages = () + + def messages(self) -> tuple[Message, ...]: + return () @property def batch_target(self) -> NoTurnOrigin: diff --git a/src/agent_comms/schedule_rules.py b/src/agent_comms/schedule_rules.py index a131c67f4d..4a995f4ba5 100644 --- a/src/agent_comms/schedule_rules.py +++ b/src/agent_comms/schedule_rules.py @@ -105,7 +105,7 @@ async def run(self) -> None: session_id, inputs.sessions.require(session_id), "\n\n".join(turn.prompt for turn in pending), - origins=tuple(origin for turn in pending for origin in turn.origin.messages), + origins=tuple(origin for turn in pending for origin in turn.origin.messages()), autonomous_goal=len(pending) == 1 and pending[0].autonomous_goal, dependency_wait_id=pending[0].goal_wait_id, ) diff --git a/src/agent_comms/send_outcome.py b/src/agent_comms/send_outcome.py index d9c8a5011b..0f7c3b076d 100644 --- a/src/agent_comms/send_outcome.py +++ b/src/agent_comms/send_outcome.py @@ -24,7 +24,7 @@ def dispatch(self, inputs, session, item, write): return self._drain(session) async def _drain(self, session): - await session.native.proc.stdin.drain() + await session.native.drain() return True diff --git a/src/agent_comms/tracked_turn.py b/src/agent_comms/tracked_turn.py index bac058f7d1..eaf534c2f1 100644 --- a/src/agent_comms/tracked_turn.py +++ b/src/agent_comms/tracked_turn.py @@ -407,16 +407,15 @@ async def next_event(self) -> pi.PiEvent: return event async def send(self, command: commands.PiCommand) -> None: - payload = self.native.reader.encode(command) try: async with self.ingress.ingress_async() as status: status.require_open() - self.write(command, payload) + self.write(command) except RelationViolationError as error: raise NativePiUnavailable("Maintenance closed before Pi capability preflight") from error wait = self.watchdog.wait_for(self) with wait.charge(self.watchdog.clock): - await asyncio.wait_for(self.native.proc.stdin.drain(), timeout=wait.timeout()) + await asyncio.wait_for(self.native.drain(), timeout=wait.timeout()) async def attest(self) -> None: self.native.attestation = self.launch.session.tracked_attestation( @@ -449,10 +448,10 @@ async def attest(self) -> None: with self.startup.measurements.operation("attestation_publication"): await self.observe_event(observed) - def write(self, command: commands.PiCommand, payload: bytes) -> None: + def write(self, command: commands.PiCommand) -> None: if command is self.command: self.grant_prompt() - self.native.proc.stdin.write(payload) + self.native.send(command) async def send_prompt(self) -> None: await self.prompt_send_boundary.send_prompt(self) diff --git a/src/agent_comms/turn_inputs.py b/src/agent_comms/turn_inputs.py index 3012286eed..3d35c99710 100644 --- a/src/agent_comms/turn_inputs.py +++ b/src/agent_comms/turn_inputs.py @@ -96,8 +96,8 @@ async def forward(self, session: TurnSession) -> None: async def write(self, session: TurnSession, *written: commands.PiCommand) -> None: for command in written: - session.native.proc.stdin.write(session.native.reader.encode(command)) - await session.native.proc.stdin.drain() + session.native.send(command) + await session.native.drain() async def queued(self, input_ids: tuple[InputId, ...]) -> tuple[str, ...]: """A native prompt acknowledgement proves it reached the child's queue.""" @@ -117,7 +117,7 @@ async def clear_queue(self, session: TurnSession, item: ClearQueue) -> bool: command = commands.ClearQueue(input_ids=await self.queued(item.input_ids)) if not command.input_ids: return True - async with command.pending_response(session.native.reader, session.native.proc.stdin) as future: + async with command.pending_response(session.native) as future: response = await future if not response.succeeded: raise RuntimeError( @@ -151,8 +151,8 @@ async def send_now(self, session: TurnSession, item: SendNow) -> bool: with status.phase.send_running( session.interrupt_boundary, first.require_id(), candidates[0], first.message, ) as outcome: - completion = outcome.dispatch(self, session, item, lambda: session.native.proc.stdin.write( - session.native.reader.encode(commands.InterruptSteering(input_ids=list(candidates))) + completion = outcome.dispatch(self, session, item, lambda: session.native.send( + commands.InterruptSteering(input_ids=list(candidates)) )) return await completion @@ -172,7 +172,7 @@ async def place(self, session: TurnSession, item: InputRequest) -> bool: def write(): if command.images: session.output.sensitive = True - session.native.proc.stdin.write(session.native.reader.encode(command)) + session.native.send(command) self.pending[native_id] = command self.generation += 1 diff --git a/src/agent_comms/turn_stats.py b/src/agent_comms/turn_stats.py index d1e28bab64..18f55ff7e8 100644 --- a/src/agent_comms/turn_stats.py +++ b/src/agent_comms/turn_stats.py @@ -53,7 +53,7 @@ def busy(self) -> bool: ) async def request(self, session: NativeSession) -> None: - if session.native.proc.stdin is None or self.requested: + if self.requested: return self.generation = session.inputs.generation self.settlement_count = session.settlement_count @@ -65,8 +65,8 @@ async def request(self, session: NativeSession) -> None: ) self._pending = tuple(session.native.reader.track(command) for command in requests) for command in requests: - session.native.proc.stdin.write(session.native.reader.command_bytes(command)) - await session.native.proc.stdin.drain() + session.native.send(command) + await session.native.drain() async def settle(self, session: NativeSession) -> AsyncIterator[events.AgentEvent]: if not (self.complete or self.failed): diff --git a/tests/test_response_policy.py b/tests/test_response_policy.py index 7e736f847d..15406dcf8d 100644 --- a/tests/test_response_policy.py +++ b/tests/test_response_policy.py @@ -58,7 +58,7 @@ def test_autonomous_turns_keep_message_batches_and_delivery_keys_separate(tmp_pa owner = thread_registration.new_thread("owner", frozenset(), str(tmp_path)) dispositions = InputDispositions(tmp_path) - assert goal.origin.messages == () + assert goal.origin.messages() == () assert goal.origin.bus_keys(dispositions, owner) == () - assert first.origin.messages == (first.origin,) + assert first.origin.messages() == (first.origin,) assert first.origin.bus_keys(dispositions, owner) == (dispositions.bus_key(first.origin, owner),) From 595cb6ac394852698c2054c08179144d46b3907f Mon Sep 17 00:00:00 2001 From: Tristan Simas Date: Fri, 23 Oct 2026 03:54:55 -0400 Subject: [PATCH 4/4] Register correlated Pi commands once at their child send --- src/agent_comms/native_child.py | 9 ++++++--- src/agent_comms/pi_commands.py | 3 +-- src/agent_comms/pi_rpc.py | 5 ----- src/agent_comms/turn_stats.py | 4 +--- tests/test_pi_rpc_nominal.py | 33 +++++++++++++++++++++++++++++---- 5 files changed, 37 insertions(+), 17 deletions(-) diff --git a/src/agent_comms/native_child.py b/src/agent_comms/native_child.py index d1622a51d4..b04f6c88e5 100644 --- a/src/agent_comms/native_child.py +++ b/src/agent_comms/native_child.py @@ -91,14 +91,17 @@ class PiChild: forwarding: InputForwarder = field(default_factory=NoInputForwarder, init=False, repr=False) _close_job: asyncio.Task[ChildOutcome | BaseException] | None = field(default=None, init=False, repr=False) - def send(self, command) -> None: - self.proc.send(self.reader.encode(command)) + def send(self, command) -> asyncio.Future: + future = self.reader.track(command) + self.proc.send(self.reader.command_bytes(command)) + return future async def drain(self) -> None: await self.proc.drain() async def send_fenced(self, command, boundary, *, measurements) -> None: - await self.proc.send_fenced(self.reader.encode(command), boundary, measurements=measurements) + self.reader.track(command) + await self.proc.send_fenced(self.reader.command_bytes(command), boundary, measurements=measurements) async def reply_ui(self, response) -> None: self.send(response) diff --git a/src/agent_comms/pi_commands.py b/src/agent_comms/pi_commands.py index 8748e8026e..79186e3e6a 100644 --- a/src/agent_comms/pi_commands.py +++ b/src/agent_comms/pi_commands.py @@ -88,9 +88,8 @@ def invalidates_identity(cls, response: Response, session: TurnSession) -> bool: @asynccontextmanager async def pending_response(self, child): channel = child.reader - future = channel.track(self) + future = child.send(self) try: - child.send(self) await child.drain() yield future finally: diff --git a/src/agent_comms/pi_rpc.py b/src/agent_comms/pi_rpc.py index 8484fd44fb..34c987022b 100644 --- a/src/agent_comms/pi_rpc.py +++ b/src/agent_comms/pi_rpc.py @@ -106,11 +106,6 @@ async def receive( return None return self.decode_record(raw, strict=strict, max_bytes=max_bytes, event_type=self.event_type) - def encode(self, command: PiCommand) -> bytes: - """Register before writing; correlated proof still requires its native input ID.""" - self.track(command) - return self.command_bytes(command) - def track(self, command: PiCommand) -> asyncio.Future[Response]: """Expose the same pending response to the request's lifecycle owner.""" return command.id.track_in(self.pending, type(command), command) diff --git a/src/agent_comms/turn_stats.py b/src/agent_comms/turn_stats.py index 18f55ff7e8..45d2863224 100644 --- a/src/agent_comms/turn_stats.py +++ b/src/agent_comms/turn_stats.py @@ -63,9 +63,7 @@ async def request(self, session: NativeSession) -> None: id=RequestId(f"agent-comms-stats-usage-{secrets.token_hex(16)}") ), ) - self._pending = tuple(session.native.reader.track(command) for command in requests) - for command in requests: - session.native.send(command) + self._pending = tuple(session.native.send(command) for command in requests) await session.native.drain() async def settle(self, session: NativeSession) -> AsyncIterator[events.AgentEvent]: diff --git a/tests/test_pi_rpc_nominal.py b/tests/test_pi_rpc_nominal.py index 001ffee0f8..cc2a08e51d 100644 --- a/tests/test_pi_rpc_nominal.py +++ b/tests/test_pi_rpc_nominal.py @@ -157,8 +157,8 @@ async def test_correlation_handles_out_of_order_ids_and_anonymous_commands(): channel = PiRpcChannel(asyncio.StreamReader()) first = commands.SetModel(id=RequestId("a"), provider="p", model_id="one") second = commands.SetModel(id=RequestId("b"), provider="p", model_id="two") - channel.encode(first) - channel.encode(second) + channel.track(first) + channel.track(second) def response(identity): return PiRpcChannel.decode_record( @@ -173,14 +173,39 @@ def response(identity): commands.InterruptSteering(id=AnonymousRequest()), commands.InterruptSteering(id=AnonymousRequest()), ) - channel.encode(one) - channel.encode(two) + channel.track(one) + channel.track(two) reply = PiRpcChannel.decode_record(b'{"type":"response","command":"interrupt_steering"}') assert channel.correlate(reply) is one assert channel.correlate(reply) is two channel.pending.cancel_all() +@pytest.mark.parametrize("anonymous", (False, True)) +async def test_pending_response_registers_each_child_exchange_once(tmp_path, anonymous): + import sys + + from agent_comms.child_process import BoundedRun + from agent_comms.native_child import PiChild + + responder = """import json, sys +for line in sys.stdin: + request = json.loads(line) + response = {"type": "response", "command": request["type"], "success": True} + if "id" in request: + response["id"] = request["id"] + print(json.dumps(response), flush=True) +""" + async with BoundedRun.session((sys.executable, "-u", "-c", responder), cwd=tmp_path, timeout=2) as proc: + child = PiChild(proc, PiRpcChannel(proc.stdout)) + for index in range(2): + request = commands.SetSteeringMode( + id=AnonymousRequest() if anonymous else RequestId(f"exchange-{index}"), + mode="all", + ) + assert (await request.exchange(child, strict=True)).succeeded + + async def test_fragmented_cancelled_line_remains_intact_in_the_only_reader(): stream = asyncio.StreamReader(limit=8) channel = PiRpcChannel(stream)