wehub-resource-sync
459 lines
19 KiB
Python
459 lines
19 KiB
Python
"""Wake a parent agent when one of its descendants blocks on an approval.
|
|
|
|
A blocked sub-agent (e.g. a codex worker asking to run ``git fetch``) never
|
|
reaches a terminal status, so the completion-only parent-wake in
|
|
``omnigent/runner/app.py`` leaves the orchestrator idle and unaware. This
|
|
module observes every elicitation at the
|
|
``pending_elicitations.record_publish`` chokepoint. The approval prompt itself
|
|
is mirrored into ancestor chats for the human, so for a child session the
|
|
parent *agent* is woken only after an escalation grace passes with the block
|
|
still unanswered (once per block, re-arming when it clears); a delivered wake
|
|
is then paired with a resolution notice when the block clears so the parent is
|
|
never left waiting on a stale notice. Delivery is an injected ``wake_dispatch``
|
|
coroutine (kept free of HTTP/server knowledge) wired up by the Omnigent server
|
|
in ``configure_subagent_block_notifier``. See
|
|
``designs/SUBAGENT_ELICITATION_VISIBILITY_V2.md``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import concurrent.futures
|
|
import contextlib
|
|
import enum
|
|
import logging
|
|
import threading
|
|
from collections.abc import Awaitable, Callable
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
if TYPE_CHECKING:
|
|
from omnigent.entities.conversation import Conversation
|
|
from omnigent.stores import ConversationStore
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
_REQUEST_TYPE = "response.elicitation_request"
|
|
_RESOLVED_TYPE = "response.elicitation_resolved"
|
|
|
|
# Max length of the elicitation prompt echoed into the parent's notice.
|
|
_REASON_MAX_CHARS = 200
|
|
|
|
# Bounded in-handler retry for a transient delivery miss (parent runner briefly
|
|
# unroutable during a reconnect/relaunch). Total attempts = 1 + _WAKE_RETRIES.
|
|
_WAKE_RETRIES = 2
|
|
# Backoff between retry attempts; short because the window we cover (a runner
|
|
# rebinding) is on the order of a second, not minutes.
|
|
_WAKE_RETRY_BACKOFF_S = 0.5
|
|
|
|
# How long a block may sit unanswered before the parent agent is woken.
|
|
# The approval card is already mirrored into the parent chat, so an
|
|
# attended session resolves silently; the wake is the unattended fallback.
|
|
_BLOCK_WAKE_ESCALATION_DELAY_S = 120.0
|
|
|
|
# Injected delivery: post ``notice`` to the parent and wake it. Returns True when
|
|
# the wake was confirmed delivered, False when it could not be (no runner bound /
|
|
# caught transport error) so the notifier can release the debounce and retry.
|
|
WakeDispatch = Callable[[str, "Conversation", str], Awaitable[bool]]
|
|
|
|
|
|
class _WakeOutcome(enum.Enum):
|
|
"""
|
|
Result of one bounded-retry wake delivery.
|
|
|
|
``DELIVERED`` — the dispatch confirmed delivery; ``MOOT`` — the block
|
|
resolved mid-flight so nothing (stale) was sent and the cleared
|
|
debounce slot must be left alone; ``FAILED`` — every attempt failed,
|
|
so the caller releases the arm and a re-publish can retry.
|
|
"""
|
|
|
|
DELIVERED = "delivered"
|
|
MOOT = "moot"
|
|
FAILED = "failed"
|
|
|
|
|
|
async def _sleep(seconds: float) -> None:
|
|
"""
|
|
Indirection over :func:`asyncio.sleep` so tests can stub the retry
|
|
backoff without clobbering the process-global ``asyncio.sleep``.
|
|
|
|
:param seconds: Seconds to sleep, e.g. ``0.5``.
|
|
:returns: None.
|
|
"""
|
|
await asyncio.sleep(seconds)
|
|
|
|
|
|
async def _escalation_sleep(seconds: float) -> None:
|
|
"""
|
|
Indirection over :func:`asyncio.sleep` for the escalation grace, kept
|
|
separate from :func:`_sleep` so tests can gate the grace and the retry
|
|
backoff independently.
|
|
|
|
:param seconds: Seconds to sleep, e.g. ``120.0``.
|
|
:returns: None.
|
|
"""
|
|
await asyncio.sleep(seconds)
|
|
|
|
|
|
class SubagentBlockNotifier:
|
|
"""
|
|
Observe elicitation events and wake a blocked child's parent.
|
|
|
|
One instance is constructed and registered per Omnigent server process.
|
|
Its :meth:`observe` method is registered with
|
|
:func:`omnigent.runtime.pending_elicitations.set_elicitation_observer`;
|
|
every tracked elicitation event flows through it.
|
|
|
|
:param conversation_store: Store used to resolve a child session's
|
|
immediate parent (``parent_conversation_id``).
|
|
:param wake_dispatch: Coroutine that delivers a notice to a parent
|
|
session and wakes it, returning ``True`` on a confirmed delivery
|
|
and ``False`` when it could not deliver. See :data:`WakeDispatch`.
|
|
:param loop: The Omnigent server's event loop, captured at registration
|
|
time so :meth:`observe` (which runs synchronously on the publish
|
|
path, possibly off the loop) can schedule async handling onto it.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
conversation_store: ConversationStore,
|
|
wake_dispatch: WakeDispatch,
|
|
loop: asyncio.AbstractEventLoop,
|
|
) -> None:
|
|
self._conversation_store = conversation_store
|
|
self._wake_dispatch = wake_dispatch
|
|
self._loop = loop
|
|
self._lock = threading.Lock()
|
|
# Debounce: armed ids (one wake per block), cleared + re-armed on resolve.
|
|
self._notified: set[str] = set()
|
|
# Per-block resolve signal, registered by the handler before it
|
|
# dispatches a wake so a resolve racing the in-flight wake is never
|
|
# missed; observe() sets it on the resolved event.
|
|
self._resolution_signals: dict[str, asyncio.Event] = {}
|
|
# Strong refs so scheduled wake futures aren't GC'd mid-flight; close() cancels them.
|
|
self._inflight: set[concurrent.futures.Future[None]] = set()
|
|
|
|
def observe(self, conversation_id: str, event: dict[str, Any]) -> None:
|
|
"""
|
|
React to one tracked elicitation event (the registered observer).
|
|
|
|
Runs synchronously on the publish path (no I/O): a ``resolved`` event
|
|
clears the debounce and signals any handler waiting to send the
|
|
resolution notice; a ``request`` event arms the debounce and schedules
|
|
:meth:`handle_request` onto the captured loop for the off-hot-path
|
|
grace + parent lookup. Arm and clear both happen here, so they order
|
|
deterministically. Other event types are ignored.
|
|
|
|
:param conversation_id: Session the event was published on,
|
|
e.g. ``"conv_child123"``.
|
|
:param event: The elicitation request/resolved event dict.
|
|
:returns: None.
|
|
"""
|
|
elicitation_id = event.get("elicitation_id")
|
|
# No id → can't be matched by its resolve, so never arm.
|
|
if not isinstance(elicitation_id, str) or not elicitation_id:
|
|
return
|
|
event_type = event.get("type")
|
|
if event_type == _RESOLVED_TYPE:
|
|
with self._lock:
|
|
self._notified.discard(elicitation_id)
|
|
signal = self._resolution_signals.get(elicitation_id)
|
|
if signal is not None:
|
|
# Event.set is not thread-safe; hop onto the loop. A gone
|
|
# loop (shutdown) means the waiting handler dies with it.
|
|
with contextlib.suppress(RuntimeError):
|
|
self._loop.call_soon_threadsafe(signal.set)
|
|
return
|
|
if event_type != _REQUEST_TYPE:
|
|
return
|
|
with self._lock:
|
|
if elicitation_id in self._notified:
|
|
return # already armed for this block — debounce
|
|
self._notified.add(elicitation_id)
|
|
try:
|
|
future = asyncio.run_coroutine_threadsafe(
|
|
self.handle_request(conversation_id, event),
|
|
self._loop,
|
|
)
|
|
except RuntimeError:
|
|
# Loop gone (shutdown): release the arm so a later block isn't wrongly debounced.
|
|
with self._lock:
|
|
self._notified.discard(elicitation_id)
|
|
_logger.debug(
|
|
"subagent block notifier: loop unavailable, dropping wake for %s",
|
|
conversation_id,
|
|
)
|
|
return
|
|
with self._lock:
|
|
self._inflight.add(future)
|
|
future.add_done_callback(self._discard_inflight)
|
|
|
|
def _discard_inflight(self, future: concurrent.futures.Future[None]) -> None:
|
|
"""
|
|
Drop a completed wake future from the in-flight set.
|
|
|
|
Registered as the done-callback on every future scheduled by
|
|
:meth:`observe`, and locked against a concurrent observe-add or
|
|
:meth:`close`. Delivery swallows its own transport errors, so a future
|
|
completing with an exception is unexpected and gets logged rather than
|
|
silently discarded.
|
|
|
|
:param future: The just-completed (or cancelled) wake future.
|
|
:returns: None.
|
|
"""
|
|
with self._lock:
|
|
self._inflight.discard(future)
|
|
if future.cancelled():
|
|
return # teardown via close() — expected, not an error
|
|
exc = future.exception()
|
|
if exc is not None:
|
|
_logger.warning(
|
|
"subagent block notifier: wake handling raised unexpectedly",
|
|
exc_info=exc,
|
|
)
|
|
|
|
def close(self) -> None:
|
|
"""
|
|
Cancel any in-flight wake futures. For lifespan teardown.
|
|
|
|
Idempotent. After ``close`` the notifier should also be
|
|
unregistered via
|
|
:func:`omnigent.runtime.pending_elicitations.set_elicitation_observer`
|
|
so no new work is scheduled.
|
|
|
|
:returns: None.
|
|
"""
|
|
with self._lock:
|
|
inflight = list(self._inflight)
|
|
self._inflight.clear()
|
|
for future in inflight:
|
|
future.cancel()
|
|
|
|
async def handle_request(
|
|
self,
|
|
conversation_id: str,
|
|
event: dict[str, Any],
|
|
) -> None:
|
|
"""
|
|
Escalate one unanswered block to the parent, then close the loop.
|
|
|
|
The approval prompt is mirrored into ancestor chats the moment it
|
|
publishes, so the human's surface is immediate; this handler waits
|
|
out ``_BLOCK_WAKE_ESCALATION_DELAY_S`` and only wakes the parent
|
|
*agent* if the block is still unanswered (the debounce arm doubles
|
|
as the cancellation token — :meth:`observe` clears it on resolve).
|
|
For a top-level session (no parent) this is a no-op.
|
|
|
|
Delivery is attempted up to ``1 + _WAKE_RETRIES`` times (the parent's
|
|
runner may be briefly unroutable during a reconnect/relaunch). An
|
|
exhausted failure *releases* the arm so a re-publish of the same
|
|
elicitation re-fires; a confirmed delivery keeps it held (one wake
|
|
per block) and the handler then waits for the resolve signal to send
|
|
the parent a resolution notice, so it is never left acting on a
|
|
stale block notice.
|
|
|
|
:param conversation_id: The blocked child session id,
|
|
e.g. ``"conv_child123"``.
|
|
:param event: The ``response.elicitation_request`` event dict.
|
|
:returns: None.
|
|
"""
|
|
elicitation_id = event.get("elicitation_id")
|
|
if not isinstance(elicitation_id, str) or not elicitation_id:
|
|
return
|
|
await _escalation_sleep(_BLOCK_WAKE_ESCALATION_DELAY_S)
|
|
with self._lock:
|
|
if elicitation_id not in self._notified:
|
|
return # answered within the grace — the parent never needs to know
|
|
child = await asyncio.to_thread(
|
|
self._conversation_store.get_conversation,
|
|
conversation_id,
|
|
)
|
|
if child is None or child.parent_conversation_id is None:
|
|
# Top-level session: no parent; the resolve event clears the arm.
|
|
return
|
|
parent_id = child.parent_conversation_id
|
|
notice = _format_block_notice(child, event)
|
|
# Register the resolve signal before dispatching so a resolve that
|
|
# lands while the wake is in flight is never missed; the arm check
|
|
# and registration share one critical section, so a resolve landing
|
|
# earlier makes this return instead of registering.
|
|
with self._lock:
|
|
if elicitation_id not in self._notified:
|
|
return
|
|
signal = asyncio.Event()
|
|
self._resolution_signals[elicitation_id] = signal
|
|
try:
|
|
outcome = await self._deliver_with_retry(
|
|
parent_id, child, notice, armed_id=elicitation_id
|
|
)
|
|
if outcome is _WakeOutcome.FAILED:
|
|
# Release the arm so a later publish of this id can retry.
|
|
with self._lock:
|
|
self._notified.discard(elicitation_id)
|
|
return
|
|
if outcome is _WakeOutcome.MOOT:
|
|
return
|
|
await signal.wait()
|
|
await self._deliver_with_retry(
|
|
parent_id, child, _format_resolution_notice(child), armed_id=None
|
|
)
|
|
finally:
|
|
with self._lock:
|
|
self._resolution_signals.pop(elicitation_id, None)
|
|
|
|
async def _deliver_with_retry(
|
|
self,
|
|
parent_id: str,
|
|
child: Conversation,
|
|
notice: str,
|
|
*,
|
|
armed_id: str | None,
|
|
) -> _WakeOutcome:
|
|
"""
|
|
Attempt one notice dispatch with a bounded retry.
|
|
|
|
With ``armed_id`` set (block notices) the arm is re-checked before
|
|
each attempt: a block that resolved while an earlier attempt was in
|
|
flight clears the arm and this reports ``MOOT`` so the caller leaves
|
|
the cleared slot alone. ``None`` (resolution notices — the block is
|
|
already resolved) skips that gate. A dispatch that returns ``False``
|
|
or raises is retried after a short backoff; an exhausted run reports
|
|
``FAILED``. A raising dispatch is caught broadly (its error taxonomy
|
|
is unknown to this module) so a delivery failure never crashes the
|
|
publish path; ``CancelledError`` (a ``BaseException`` raised by
|
|
:meth:`close`) is not caught and still tears the handler down.
|
|
|
|
:param parent_id: Parent session id to wake, e.g. ``"conv_parent123"``.
|
|
:param child: The blocked child :class:`Conversation`.
|
|
:param notice: Pre-formatted ``[System: …]`` notice text.
|
|
:param armed_id: Correlation id whose arm gates the dispatch, or
|
|
``None`` to dispatch unconditionally.
|
|
:returns: The :class:`_WakeOutcome` of the delivery.
|
|
"""
|
|
for attempt in range(1 + _WAKE_RETRIES):
|
|
if armed_id is not None:
|
|
with self._lock:
|
|
if armed_id not in self._notified:
|
|
# Resolved during the grace/backoff — block's gone, the
|
|
# slot is already clear; nothing (stale) to wake.
|
|
return _WakeOutcome.MOOT
|
|
try:
|
|
if await self._wake_dispatch(parent_id, child, notice):
|
|
return _WakeOutcome.DELIVERED
|
|
except Exception:
|
|
# Broad: injected dispatch error types are unknown; don't crash the publish path.
|
|
_logger.warning(
|
|
"subagent block notifier: wake dispatch raised for parent=%s child=%s "
|
|
"(attempt %d/%d)",
|
|
parent_id,
|
|
child.id,
|
|
attempt + 1,
|
|
1 + _WAKE_RETRIES,
|
|
exc_info=True,
|
|
)
|
|
if attempt < _WAKE_RETRIES:
|
|
await _sleep(_WAKE_RETRY_BACKOFF_S)
|
|
_logger.warning(
|
|
"subagent block notifier: notice undelivered after %d attempt(s) for "
|
|
"parent=%s child=%s",
|
|
1 + _WAKE_RETRIES,
|
|
parent_id,
|
|
child.id,
|
|
)
|
|
return _WakeOutcome.FAILED
|
|
|
|
|
|
def _format_block_notice(child: Conversation, event: dict[str, Any]) -> str:
|
|
"""
|
|
Build the ``[System: …]`` notice posted into the parent session.
|
|
|
|
Mirrors the shape of the runner's terminal-completion wake notice
|
|
(``_format_subagent_wake_notice``). Describes the situation and asks
|
|
the parent to surface it — it does not prescribe a specific tool.
|
|
|
|
:param child: The blocked child :class:`Conversation`, used for its
|
|
``<agent>:<title>`` label.
|
|
:param event: The ``response.elicitation_request`` event dict; its
|
|
``params.message`` (when present) is echoed as the reason.
|
|
:returns: A one-line ``[System: …]`` notice, e.g. ``"[System:
|
|
sub-agent codex/auth-refactor is blocked awaiting human
|
|
approval: Codex wants to run 'git fetch'. Its approval prompt
|
|
is mirrored into this conversation but has gone unanswered —
|
|
surface the situation to the human and do not wait silently. It
|
|
cannot continue until the request is resolved.]"``.
|
|
"""
|
|
label = _child_label(child)
|
|
reason = _block_reason(event)
|
|
detail = f": {reason}" if reason else ""
|
|
return (
|
|
f"[System: sub-agent {label} is blocked awaiting human approval{detail}. "
|
|
"Its approval prompt is mirrored into this conversation but has gone "
|
|
"unanswered — surface the situation to the human and do not wait "
|
|
"silently. It cannot continue until the request is resolved.]"
|
|
)
|
|
|
|
|
|
def _format_resolution_notice(child: Conversation) -> str:
|
|
"""
|
|
Build the follow-up notice sent after a woken block resolves.
|
|
|
|
Sent only when a block notice was actually delivered, so the parent
|
|
stops acting on it (offering to relay answers, polling the child)
|
|
once the human has dealt with the prompt directly.
|
|
|
|
:param child: The previously blocked child :class:`Conversation`,
|
|
used for its ``<agent>:<title>`` label.
|
|
:returns: A one-line ``[System: …]`` notice, e.g. ``"[System:
|
|
sub-agent codex/auth-refactor's pending approval has been
|
|
resolved and it is continuing. …]"``.
|
|
"""
|
|
label = _child_label(child)
|
|
return (
|
|
f"[System: sub-agent {label}'s pending approval has been resolved and it "
|
|
"is continuing. No action is needed on the earlier block notice — if you "
|
|
"are waiting on its output, the result will arrive in your inbox as "
|
|
"usual.]"
|
|
)
|
|
|
|
|
|
def _child_label(child: Conversation) -> str:
|
|
"""
|
|
Derive a human label from a sub-agent conversation title.
|
|
|
|
Named sub-agents store ``"<agent>:<title>"`` in
|
|
:attr:`Conversation.title`. Returns ``"<agent>/<title>"`` for that
|
|
form, falling back to the raw title or, last, the conversation id so
|
|
the notice always names *something*.
|
|
|
|
:param child: The child :class:`Conversation`.
|
|
:returns: A label like ``"codex/auth-refactor"``.
|
|
"""
|
|
title = child.title or ""
|
|
if ":" in title:
|
|
agent, _, remainder = title.partition(":")
|
|
sa_title = remainder.partition(":closed:")[0]
|
|
return f"{agent}/{sa_title}" if sa_title else agent
|
|
return title or child.id
|
|
|
|
|
|
def _block_reason(event: dict[str, Any]) -> str | None:
|
|
"""
|
|
Extract a short human reason from an elicitation request event.
|
|
|
|
:param event: The ``response.elicitation_request`` event dict; reads
|
|
the nested ``params.message``.
|
|
:returns: The trimmed/truncated prompt text, or ``None`` when the
|
|
event carried no message.
|
|
"""
|
|
params = event.get("params")
|
|
if not isinstance(params, dict):
|
|
return None
|
|
message = params.get("message")
|
|
if not isinstance(message, str):
|
|
return None
|
|
message = message.strip()
|
|
if not message:
|
|
return None
|
|
if len(message) > _REASON_MAX_CHARS:
|
|
return message[: _REASON_MAX_CHARS - 1].rstrip() + "…"
|
|
return message
|