wehub-resource-sync
352 lines
14 KiB
Python
352 lines
14 KiB
Python
"""In-process index of currently outstanding elicitation requests.
|
|
|
|
Lets the Omnigent server answer "which sessions have a pending approval
|
|
prompt?" without scanning per-session state or persisting elicitation
|
|
rows. The index is the *sidebar's* view of pending state — it lives
|
|
alongside the underlying parked awaiter (a runner-side Future or a
|
|
server-side ``_harness_elicitation_registry`` Future) and shares its
|
|
lifecycle: when the Omnigent process dies, both the index and every parked
|
|
awaiter die together, so the index cannot diverge from the underlying
|
|
state into "phantom" pending rows.
|
|
|
|
The index is populated automatically by
|
|
:func:`omnigent.runtime.session_stream.publish` whenever a
|
|
``response.elicitation_request`` event passes through (server-emitted
|
|
policy elicitations, the claude-native PermissionRequest hook, and
|
|
runner-originated elicitations relayed by ``_relay_runner_stream``
|
|
all funnel through that single chokepoint). It is decremented either
|
|
by the approval-dispatch path on ``POST /v1/sessions/{id}/events``
|
|
when an ``approval`` verdict arrives, or by a
|
|
``response.elicitation_resolved`` event from the runner when its own
|
|
Future timed out or was cancelled without a UI verdict.
|
|
|
|
The index also stores the full event payload (not just the id) so
|
|
``GET /v1/sessions/{id}`` can replay outstanding prompts into the
|
|
UI's chat blocks on cold load — the SSE stream itself has no replay,
|
|
so an elicitation emitted before the user opened the chat would
|
|
otherwise render as nothing.
|
|
|
|
Limitations:
|
|
|
|
* In-memory only; multi-replica Omnigent deploys would each see their own
|
|
slice. This matches the existing ``_harness_elicitation_registry``
|
|
constraint — when a shared backplane is added for the registry,
|
|
this index should be wired through the same backplane.
|
|
* Events emitted before the Omnigent server starts (e.g. between turns,
|
|
with the session_stream having dropped them) are not tracked,
|
|
same as every other AP-server-side in-memory state.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
import threading
|
|
from collections.abc import Callable
|
|
from typing import Any
|
|
|
|
# Per-conversation mapping of outstanding elicitation_id → original
|
|
# event payload. Storing the full event (not just the id) lets
|
|
# ``GET /v1/sessions/{id}`` replay the prompt into the UI on cold
|
|
# load. Populated by ``record_publish`` on
|
|
# ``response.elicitation_request``; drained by ``resolve`` (called
|
|
# directly from the approval-dispatch path, or via ``record_publish``
|
|
# when a ``response.elicitation_resolved`` event flows through the
|
|
# SSE chokepoint). Empty inner dicts are popped eagerly so
|
|
# :func:`count_for` doesn't see stale keys.
|
|
_pending: dict[str, dict[str, dict[str, Any]]] = {}
|
|
_lock = threading.Lock()
|
|
|
|
# Optional observer (``subagent_block_notifier``) run synchronously on every
|
|
# tracked event — must be cheap + non-blocking. ``None`` (runner, tests) skips it.
|
|
_observer: Callable[[str, dict[str, Any]], None] | None = None
|
|
|
|
|
|
def set_elicitation_observer(
|
|
observer: Callable[[str, dict[str, Any]], None] | None,
|
|
) -> None:
|
|
"""
|
|
Register (or clear) the elicitation-event observer.
|
|
|
|
:param observer: Callback invoked as ``observer(conversation_id,
|
|
event)`` for every ``response.elicitation_request`` and
|
|
``response.elicitation_resolved`` event passing through
|
|
:func:`record_publish`. Pass ``None`` to clear (e.g. test
|
|
teardown). Replaces any previously registered observer.
|
|
:returns: None.
|
|
"""
|
|
global _observer
|
|
_observer = observer
|
|
|
|
|
|
def record_publish(conversation_id: str, event: dict[str, Any]) -> None:
|
|
"""
|
|
Update the index when an SSE event is published.
|
|
|
|
Acts on two event types and silently ignores every other type —
|
|
the function sits on the hot publish path so unrelated events
|
|
must pay only one dict-key lookup.
|
|
|
|
* ``response.elicitation_request`` — add the elicitation id +
|
|
full event payload to the index. The payload is what
|
|
:func:`snapshot_for` replays into ``GET /v1/sessions/{id}``
|
|
so the UI can render the ApprovalCard on cold load.
|
|
* ``response.elicitation_resolved`` — drop the elicitation id
|
|
from the index. Used by the runner to clear an entry when
|
|
its own ``_pending_approvals`` Future timed out or was
|
|
cancelled without a UI verdict; same effect as
|
|
:func:`resolve` but routed through the SSE chokepoint so the
|
|
Omnigent server picks it up via ``_relay_runner_stream`` without
|
|
a separate out-of-band signal.
|
|
|
|
Idempotent on both event types — adds use a dict assignment
|
|
(re-publishing the same id overwrites the payload) and drops are
|
|
no-ops for unknown ids.
|
|
|
|
After updating the index, a registered observer (if any — see
|
|
:func:`set_elicitation_observer`) is notified with the same
|
|
``(conversation_id, event)`` so cross-session consumers (the
|
|
parent-wake notifier) can react without re-deriving the event type.
|
|
|
|
:param conversation_id: Conversation/session id the event was
|
|
published on, e.g. ``"conv_abc123"``.
|
|
:param event: The event dict as passed to
|
|
:func:`omnigent.runtime.session_stream.publish`. Reads
|
|
``event["type"]`` to dispatch and ``event["elicitation_id"]``
|
|
for both event types.
|
|
"""
|
|
event_type = event.get("type")
|
|
if event_type == "response.elicitation_request":
|
|
elicitation_id = event.get("elicitation_id")
|
|
if not isinstance(elicitation_id, str) or not elicitation_id:
|
|
return
|
|
with _lock:
|
|
_pending.setdefault(conversation_id, {})[elicitation_id] = event
|
|
_notify_observer(conversation_id, event)
|
|
return
|
|
if event_type == "response.elicitation_resolved":
|
|
elicitation_id = event.get("elicitation_id")
|
|
if not isinstance(elicitation_id, str) or not elicitation_id:
|
|
return
|
|
resolve(conversation_id, elicitation_id)
|
|
_notify_observer(conversation_id, event)
|
|
|
|
|
|
def _notify_observer(conversation_id: str, event: dict[str, Any]) -> None:
|
|
"""
|
|
Invoke the registered elicitation observer, if any.
|
|
|
|
Read the module-global once into a local so a concurrent
|
|
:func:`set_elicitation_observer` clearing it cannot turn the call
|
|
into a ``None`` dereference between the check and the call.
|
|
|
|
:param conversation_id: Conversation/session id the event was
|
|
published on, e.g. ``"conv_abc123"``.
|
|
:param event: The elicitation request/resolved event dict.
|
|
:returns: None.
|
|
"""
|
|
observer = _observer
|
|
if observer is not None:
|
|
observer(conversation_id, event)
|
|
|
|
|
|
def resolve(conversation_id: str, elicitation_id: str) -> None:
|
|
"""
|
|
Drop an outstanding elicitation from the index.
|
|
|
|
Called by the approval-dispatch path once a verdict has been
|
|
accepted (regardless of which downstream awaiter — runner
|
|
workflow or server-side ``_harness_elicitation_registry``
|
|
Future — receives the verdict). Also called by
|
|
:func:`record_publish` when a
|
|
``response.elicitation_resolved`` event passes through. The
|
|
decrement here is the sidebar's signal that the session no
|
|
longer needs attention.
|
|
|
|
Idempotent: removing an id that isn't tracked is a no-op (the
|
|
approval dispatch path may resolve an id whose publish landed
|
|
on a different replica, or whose tracking failed validation).
|
|
Empty conversation entries are popped so :func:`count_for`
|
|
returns ``0`` cleanly without leaving stale keys.
|
|
|
|
:param conversation_id: Conversation/session id the verdict
|
|
was dispatched against, e.g. ``"conv_abc123"``.
|
|
:param elicitation_id: The elicitation correlation id from the
|
|
approval payload, e.g. ``"elicit_abc123"``.
|
|
"""
|
|
with _lock:
|
|
ids = _pending.get(conversation_id)
|
|
if ids is None:
|
|
return
|
|
ids.pop(elicitation_id, None)
|
|
if not ids:
|
|
_pending.pop(conversation_id, None)
|
|
|
|
|
|
def count_for(conversation_id: str) -> int:
|
|
"""
|
|
Return the number of outstanding elicitations for one session.
|
|
|
|
:param conversation_id: Conversation/session id to query,
|
|
e.g. ``"conv_abc123"``.
|
|
:returns: Count of outstanding elicitations; ``0`` when the
|
|
session has none tracked.
|
|
"""
|
|
with _lock:
|
|
ids = _pending.get(conversation_id)
|
|
return len(ids) if ids is not None else 0
|
|
|
|
|
|
def counts_for(conversation_ids: list[str]) -> dict[str, int]:
|
|
"""
|
|
Batch lookup of pending counts for a list of session ids.
|
|
|
|
Used by ``GET /v1/sessions`` to populate the
|
|
``pending_elicitations_count`` field on each
|
|
:class:`omnigent.server.schemas.SessionListItem` in one
|
|
pass without re-acquiring the lock per session.
|
|
|
|
:param conversation_ids: Conversation/session ids to query,
|
|
e.g. ``["conv_abc123", "conv_def456"]``.
|
|
:returns: Mapping from each id in the input to its outstanding
|
|
count. Ids not tracked in the index map to ``0``.
|
|
"""
|
|
with _lock:
|
|
return {
|
|
conv_id: len(_pending[conv_id]) if conv_id in _pending else 0
|
|
for conv_id in conversation_ids
|
|
}
|
|
|
|
|
|
def pending_session_ids() -> list[str]:
|
|
"""
|
|
Return ids of every session with at least one outstanding elicitation.
|
|
|
|
Used by ``GET /v1/sessions/{id}`` as a cheap pre-check before the
|
|
descendant walk that mirrors child pending prompts into an ancestor
|
|
snapshot — that walk costs one ``list_conversations`` query per
|
|
session in the tree, so it should run only when some session other
|
|
than the one being snapshotted actually has an outstanding prompt
|
|
(the rare, transient case).
|
|
|
|
:returns: Session ids with outstanding elicitations, e.g.
|
|
``["conv_child123"]``. Empty when nothing is pending anywhere.
|
|
"""
|
|
with _lock:
|
|
return list(_pending.keys())
|
|
|
|
|
|
def snapshot_for(conversation_id: str) -> list[dict[str, Any]]:
|
|
"""
|
|
Return outstanding elicitation event payloads for one session.
|
|
|
|
Used by ``GET /v1/sessions/{id}`` to replay outstanding
|
|
``response.elicitation_request`` events into the UI on cold
|
|
load. Without replay, an elicitation emitted before the user
|
|
opened the chat would never render — the SSE stream has no
|
|
replay buffer.
|
|
|
|
Returns deep copies of the stored event dicts so callers
|
|
can mutate at any depth without poisoning the index. The
|
|
elicitation event carries a nested ``params`` block; a
|
|
shallow copy here would leak nested-dict mutations back into
|
|
the index for subsequent reads.
|
|
|
|
:param conversation_id: Conversation/session id to query,
|
|
e.g. ``"conv_abc123"``.
|
|
:returns: List of event dicts (each shaped like the original
|
|
``response.elicitation_request`` payload). Order is
|
|
insertion order. Empty list when the session has no
|
|
outstanding prompts.
|
|
"""
|
|
with _lock:
|
|
ids = _pending.get(conversation_id)
|
|
if ids is None:
|
|
return []
|
|
return [copy.deepcopy(event) for event in ids.values()]
|
|
|
|
|
|
def project_for_peek(event: dict[str, Any]) -> dict[str, Any]:
|
|
"""
|
|
Project a stored elicitation event into a compact peek item.
|
|
|
|
``sys_session_get_history`` returns a tail of compact conversation
|
|
items so a parent agent can read a sub-agent's recent activity. A
|
|
parked elicitation never lands in the conversation store (it lives
|
|
only in this index — see the module docstring), so get_history must
|
|
synthesize an item from the stored ``response.elicitation_request``
|
|
payload. This projector produces that item, shaped to match the
|
|
other compact items: a ``type`` discriminator plus the human-facing
|
|
prompt and (form mode only) the fields being requested.
|
|
|
|
Used by both read paths — the in-process
|
|
:class:`omnigent.tools.builtins.spawn.SysSessionGetHistoryTool`
|
|
(reads :func:`snapshot_for` directly) and the runner's REST
|
|
get_history (reads the same payloads off the
|
|
``GET /v1/sessions/{id}`` snapshot). Both receive the identical
|
|
event dict, so a single projector keeps the two outputs consistent.
|
|
|
|
:param event: A stored ``response.elicitation_request`` event dict,
|
|
as returned by :func:`snapshot_for`. Reads ``elicitation_id``
|
|
and the nested ``params`` block (``message`` and, for form
|
|
mode, ``requestedSchema.properties``).
|
|
:returns: A compact dict, e.g.
|
|
``{"type": "pending_elicitation", "elicitation_id":
|
|
"elicit_abc123", "prompt": "Approve running 'rm -rf'?",
|
|
"fields": ["approve"]}``. ``prompt`` is ``None`` when the
|
|
payload carried no message; ``fields`` is omitted when the
|
|
elicitation is not a form (or has no declared properties).
|
|
"""
|
|
params = event.get("params")
|
|
params = params if isinstance(params, dict) else {}
|
|
item: dict[str, Any] = {
|
|
"type": "pending_elicitation",
|
|
"elicitation_id": event.get("elicitation_id"),
|
|
"prompt": params.get("message"),
|
|
}
|
|
schema = params.get("requestedSchema")
|
|
if isinstance(schema, dict):
|
|
properties = schema.get("properties")
|
|
if isinstance(properties, dict) and properties:
|
|
item["fields"] = list(properties.keys())
|
|
return item
|
|
|
|
|
|
def lookup(elicitation_id: str) -> tuple[str, dict[str, Any]] | None:
|
|
"""
|
|
Look up a single outstanding elicitation by its correlation id.
|
|
|
|
Returns the ``(conversation_id, event_payload)`` pair if the
|
|
elicitation is still pending, or ``None`` if it has already been
|
|
resolved, timed out, or was never tracked. Used by the standalone
|
|
approval page route to render the elicitation prompt without
|
|
requiring a database round-trip — the in-memory index already
|
|
holds the full event payload.
|
|
|
|
:param elicitation_id: The elicitation correlation id to look up,
|
|
e.g. ``"elicit_abc123"``.
|
|
:returns: ``(conversation_id, event_dict)`` when found, ``None``
|
|
otherwise. The event dict is a deep copy so callers cannot
|
|
mutate the index.
|
|
"""
|
|
with _lock:
|
|
for conv_id, ids in _pending.items():
|
|
if elicitation_id in ids:
|
|
return conv_id, copy.deepcopy(ids[elicitation_id])
|
|
return None
|
|
|
|
|
|
def reset_for_tests() -> None:
|
|
"""
|
|
Clear the entire index. For test isolation only.
|
|
|
|
Tests that exercise the publish or dispatch paths can mutate
|
|
the module-global state; this resets between tests so leak
|
|
from one test doesn't change the behavior of another. Not for
|
|
production callers — there is no legitimate use case for
|
|
wiping the index at runtime.
|
|
"""
|
|
global _observer
|
|
with _lock:
|
|
_pending.clear()
|
|
_observer = None
|