mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-07-21 16:18:55 +00:00
fix(gateway): deduplicate completion delivery
This commit is contained in:
parent
bd740f203b
commit
94a7705bdd
3 changed files with 694 additions and 61 deletions
183
gateway/run.py
183
gateway/run.py
|
|
@ -2979,6 +2979,15 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
|
|||
# cannot grow unbounded over a long-running gateway lifetime.
|
||||
self._session_sources: "OrderedDict[str, SessionSource]" = OrderedDict()
|
||||
self._session_sources_max = 512
|
||||
# Completion delivery is intentionally lifecycle-scoped. This closes
|
||||
# duplicate queue/watcher races inside one gateway without pretending
|
||||
# the adapter call and a persistence write can be exactly-once across
|
||||
# a process crash. Any durable async-delegation replay state remains
|
||||
# owned by tools.async_delegation, not a parallel gateway ledger.
|
||||
self._completion_delivery_lock = threading.Lock()
|
||||
self._completion_deliveries_inflight: set[tuple[str, str, object]] = set()
|
||||
self._completion_deliveries_delivered: "OrderedDict[tuple[str, str, object], None]" = OrderedDict()
|
||||
self._completion_delivery_retention = 2048
|
||||
|
||||
# Cache AIAgent instances per session to preserve prompt caching.
|
||||
# Without this, a new AIAgent is created per message, rebuilding the
|
||||
|
|
@ -15454,11 +15463,17 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
|
|||
user_name=str(evt.get("user_name") or "").strip() or None,
|
||||
)
|
||||
|
||||
async def _inject_watch_notification(self, synth_text: str, evt: dict) -> None:
|
||||
"""Inject a watch-pattern notification as a synthetic message event.
|
||||
async def _inject_watch_notification(
|
||||
self, synth_text: str, evt: dict,
|
||||
) -> Optional[bool]:
|
||||
"""Inject a watch/completion notification as a synthetic message event.
|
||||
|
||||
Routing must come from the queued watch event itself, not from whatever
|
||||
Routing must come from the queued event itself, not from whatever
|
||||
foreground message happened to be active when the queue was drained.
|
||||
Returns ``True`` after adapter acceptance, ``False`` after a retryable
|
||||
adapter failure, and ``None`` when the event has no gateway route. This
|
||||
is not a transactional boundary: a process crash after adapter
|
||||
acceptance can still cause durable at-least-once replay.
|
||||
"""
|
||||
source = self._build_process_event_source(evt)
|
||||
if not source:
|
||||
|
|
@ -15466,7 +15481,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
|
|||
"Dropping watch notification with no routing metadata for process %s",
|
||||
evt.get("session_id", "unknown"),
|
||||
)
|
||||
return
|
||||
return None
|
||||
platform_name = source.platform.value if hasattr(source.platform, "value") else str(source.platform)
|
||||
adapter = None
|
||||
for p, a in self.adapters.items():
|
||||
|
|
@ -15474,7 +15489,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
|
|||
adapter = a
|
||||
break
|
||||
if not adapter:
|
||||
return
|
||||
return None
|
||||
try:
|
||||
metadata = {}
|
||||
parent_session_id = str(evt.get("parent_session_id") or "").strip()
|
||||
|
|
@ -15495,8 +15510,96 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
|
|||
source.thread_id,
|
||||
)
|
||||
await adapter.handle_message(synth_event)
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error("Watch notification injection error: %s", e)
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _completion_delivery_identity(evt: dict) -> Optional[tuple[str, str, object]]:
|
||||
"""Return a producer-stable identity when one is available.
|
||||
|
||||
Delegation UUIDs identify one producer completion. Process session IDs
|
||||
are normally unique too, but include the persisted spawn epoch so an
|
||||
explicitly reused ID represents a distinct process incarnation. Legacy
|
||||
process events without ``started_at`` are delivered without deduplication
|
||||
rather than risking suppression of a real completion.
|
||||
"""
|
||||
evt_type = str(evt.get("type") or "")
|
||||
if evt_type == "async_delegation":
|
||||
producer_id = str(evt.get("delegation_id") or "")
|
||||
return (evt_type, producer_id, "") if producer_id else None
|
||||
if evt_type == "completion":
|
||||
producer_id = str(evt.get("session_id") or "")
|
||||
started_at = evt.get("started_at")
|
||||
if producer_id and started_at is not None:
|
||||
return (evt_type, producer_id, started_at)
|
||||
return None
|
||||
|
||||
async def _deliver_completion_notification(
|
||||
self, synth_text: str, evt: dict,
|
||||
) -> Optional[bool]:
|
||||
"""Deliver once per live gateway, or return False for a retry.
|
||||
|
||||
``True`` means this caller reached adapter acceptance, ``False`` means
|
||||
injection failed and the claim was released for retry, and ``None``
|
||||
means either another same-lifecycle caller owns/delivered the producer
|
||||
event or the event has no gateway route. No cross-process exactly-once
|
||||
guarantee is claimed.
|
||||
"""
|
||||
identity = self._completion_delivery_identity(evt)
|
||||
if identity is not None:
|
||||
with self._completion_delivery_lock:
|
||||
if (
|
||||
identity in self._completion_deliveries_inflight
|
||||
or identity in self._completion_deliveries_delivered
|
||||
):
|
||||
return None
|
||||
self._completion_deliveries_inflight.add(identity)
|
||||
|
||||
accepted = False
|
||||
try:
|
||||
injection_result = await self._inject_watch_notification(synth_text, evt)
|
||||
if injection_result is not True:
|
||||
return injection_result
|
||||
accepted = True
|
||||
|
||||
if identity is not None:
|
||||
with self._completion_delivery_lock:
|
||||
self._completion_deliveries_inflight.discard(identity)
|
||||
self._completion_deliveries_delivered[identity] = None
|
||||
while (
|
||||
len(self._completion_deliveries_delivered)
|
||||
> self._completion_delivery_retention
|
||||
):
|
||||
self._completion_deliveries_delivered.popitem(last=False)
|
||||
|
||||
# If the durable async-delegation producer branch is present, its
|
||||
# SQLite row remains the authoritative replay state. Acknowledge it
|
||||
# after adapter acceptance; this gateway keeps no parallel ledger.
|
||||
if evt.get("type") == "async_delegation":
|
||||
from tools import async_delegation
|
||||
|
||||
mark_completion_delivered = getattr(
|
||||
async_delegation, "mark_completion_delivered", None,
|
||||
)
|
||||
if mark_completion_delivered is not None:
|
||||
try:
|
||||
mark_completion_delivered(str(evt.get("delegation_id") or ""))
|
||||
except Exception as exc:
|
||||
# Adapter acceptance already happened. Retrying now
|
||||
# would duplicate the turn; SQLite may replay after a
|
||||
# restart, which is the honest at-least-once contract.
|
||||
logger.warning(
|
||||
"Could not acknowledge durable async completion %s: %s",
|
||||
evt.get("delegation_id", "unknown"),
|
||||
exc,
|
||||
)
|
||||
return True
|
||||
finally:
|
||||
if identity is not None and not accepted:
|
||||
with self._completion_delivery_lock:
|
||||
self._completion_deliveries_inflight.discard(identity)
|
||||
|
||||
def _enrich_async_delegation_routing(self, evt: dict) -> None:
|
||||
"""Fill platform/chat_id/thread_id/chat_type on an async-delegation event.
|
||||
|
|
@ -15559,8 +15662,11 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
|
|||
if not synth_text:
|
||||
continue
|
||||
try:
|
||||
await self._inject_watch_notification(synth_text, evt)
|
||||
delivered = await self._deliver_completion_notification(synth_text, evt)
|
||||
if delivered is False:
|
||||
_pr.completion_queue.put(evt)
|
||||
except Exception as e:
|
||||
_pr.completion_queue.put(evt)
|
||||
logger.error("Async delegation injection error: %s", e)
|
||||
except Exception as e:
|
||||
logger.debug("Async delegation watcher error: %s", e)
|
||||
|
|
@ -15626,8 +15732,12 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
|
|||
# (#10156) — a status check must not suppress this delivery turn.
|
||||
from tools.process_registry import format_process_notification, process_registry as _pr_check
|
||||
if agent_notify and not _pr_check.is_completion_consumed(session_id):
|
||||
from agent.redact import redact_terminal_output
|
||||
from tools.ansi_strip import strip_ansi
|
||||
_command = getattr(session, "command", "") or ""
|
||||
_raw = strip_ansi(session.output_buffer) if session.output_buffer else ""
|
||||
_raw = redact_terminal_output(_raw, _command)
|
||||
_command = _redact_gateway_user_facing_secrets(_command)
|
||||
# Truncate at line boundaries so notifications never start
|
||||
# mid-line (fixes #23284). Keep the last ~2000 chars but
|
||||
# snap to the nearest preceding newline, then prepend a
|
||||
|
|
@ -15640,57 +15750,34 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
|
|||
_out = f"[… output truncated — showing last {len(_tail)} chars]\n{_tail}"
|
||||
else:
|
||||
_out = _raw
|
||||
synth_text = format_process_notification({
|
||||
completion_evt = {
|
||||
"type": "completion",
|
||||
"session_id": session_id,
|
||||
"command": session.command,
|
||||
"exit_code": session.exit_code,
|
||||
"completion_reason": getattr(session, "completion_reason", "exited"),
|
||||
"termination_source": getattr(session, "termination_source", ""),
|
||||
"output": _out,
|
||||
})
|
||||
if not synth_text:
|
||||
break
|
||||
source = self._build_process_event_source({
|
||||
"session_id": session_id,
|
||||
"session_key": session_key,
|
||||
"platform": platform_name,
|
||||
"chat_type": watcher.get("chat_type", ""),
|
||||
"chat_id": chat_id,
|
||||
"thread_id": thread_id,
|
||||
"user_id": user_id,
|
||||
"user_name": user_name,
|
||||
})
|
||||
if not source:
|
||||
logger.warning(
|
||||
"Dropping completion notification with no routing metadata for process %s",
|
||||
session_id,
|
||||
)
|
||||
"message_id": message_id,
|
||||
"started_at": getattr(session, "started_at", None),
|
||||
"command": _command,
|
||||
"exit_code": session.exit_code,
|
||||
"completion_reason": getattr(session, "completion_reason", "exited"),
|
||||
"termination_source": getattr(session, "termination_source", ""),
|
||||
"output": _out,
|
||||
}
|
||||
synth_text = format_process_notification(completion_evt)
|
||||
if not synth_text:
|
||||
break
|
||||
|
||||
adapter = None
|
||||
for p, a in self.adapters.items():
|
||||
if p == source.platform:
|
||||
adapter = a
|
||||
break
|
||||
if adapter and source.chat_id:
|
||||
try:
|
||||
synth_event = MessageEvent(
|
||||
text=synth_text,
|
||||
message_type=MessageType.TEXT,
|
||||
source=source,
|
||||
internal=True,
|
||||
message_id=message_id,
|
||||
)
|
||||
logger.info(
|
||||
"Process %s finished — injecting agent notification for session %s chat=%s thread=%s",
|
||||
session_id,
|
||||
session_key,
|
||||
source.chat_id,
|
||||
source.thread_id,
|
||||
)
|
||||
await adapter.handle_message(synth_event)
|
||||
except Exception as e:
|
||||
logger.error("Agent notify injection error: %s", e)
|
||||
delivered = await self._deliver_completion_notification(
|
||||
synth_text, completion_evt,
|
||||
)
|
||||
if delivered is False:
|
||||
# The process remains terminal; retry after failed
|
||||
# adapter injection instead of suppressing the result.
|
||||
continue
|
||||
break
|
||||
|
||||
# --- Normal text-only notification ---
|
||||
|
|
|
|||
494
tests/gateway/test_completion_delivery.py
Normal file
494
tests/gateway/test_completion_delivery.py
Normal file
|
|
@ -0,0 +1,494 @@
|
|||
"""Lifecycle-scoped gateway delivery regressions for terminal completions.
|
||||
|
||||
The gateway contract here is deliberately narrower than exactly-once: one live
|
||||
GatewayRunner suppresses concurrent/replayed copies after successful adapter
|
||||
injection, failed injection remains retryable, and durable async-delegation
|
||||
state (when available) is acknowledged through its authoritative SQLite API.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import queue
|
||||
from collections import OrderedDict
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from gateway.config import Platform
|
||||
from gateway.run import GatewayRunner
|
||||
from gateway.session import SessionSource
|
||||
from tools.process_registry import ProcessRegistry, ProcessSession
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def isolated_registry(tmp_path, monkeypatch):
|
||||
"""Any current/future durable compatibility path must stay in tmp state."""
|
||||
monkeypatch.setenv("HERMES_HOME", str(tmp_path))
|
||||
import tools.process_registry as pr_module
|
||||
|
||||
monkeypatch.setattr(pr_module, "CHECKPOINT_PATH", tmp_path / "processes.json")
|
||||
registry = pr_module.ProcessRegistry()
|
||||
monkeypatch.setattr(pr_module, "process_registry", registry)
|
||||
return registry
|
||||
|
||||
|
||||
def _runner(adapter, *, origins=None):
|
||||
runner = object.__new__(GatewayRunner)
|
||||
runner._running = True
|
||||
runner.adapters = {Platform.TELEGRAM: adapter}
|
||||
runner.session_store = SimpleNamespace(
|
||||
_ensure_loaded=lambda: None,
|
||||
_entries=origins or {},
|
||||
)
|
||||
runner._session_source_cache = {}
|
||||
runner._completion_delivery_lock = __import__("threading").Lock()
|
||||
runner._completion_deliveries_inflight = set()
|
||||
runner._completion_deliveries_delivered = OrderedDict()
|
||||
runner._completion_delivery_retention = 2048
|
||||
return runner
|
||||
|
||||
|
||||
def _async_event(delegation_id="deleg_duplicate"):
|
||||
return {
|
||||
"type": "async_delegation",
|
||||
"delegation_id": delegation_id,
|
||||
"session_key": "agent:main:telegram:dm:12345:678",
|
||||
"goal": "Investigate flaky test",
|
||||
"status": "completed",
|
||||
"summary": "Found it",
|
||||
"api_calls": 1,
|
||||
"duration_seconds": 12.0,
|
||||
"dispatched_at": 1000.0,
|
||||
"completed_at": 1012.0,
|
||||
# PR #62479 stamps these on gateway-owned events. They must not
|
||||
# change the producer identity used for queue replay.
|
||||
"origin_profile": "default",
|
||||
"origin_hermes_home": "/tmp/hermes-default",
|
||||
}
|
||||
|
||||
|
||||
def _completion_event(*, started_at, session_id="proc_reused"):
|
||||
return {
|
||||
"type": "completion",
|
||||
"session_id": session_id,
|
||||
"session_key": "agent:main:telegram:dm:123",
|
||||
"platform": "telegram",
|
||||
"chat_type": "dm",
|
||||
"chat_id": "123",
|
||||
"started_at": started_at,
|
||||
"command": "echo done",
|
||||
"exit_code": 0,
|
||||
"completion_reason": "exited",
|
||||
"output": "done\n",
|
||||
}
|
||||
|
||||
|
||||
def _stop_after_sleeps(monkeypatch, runner, count):
|
||||
sleep_calls = 0
|
||||
|
||||
async def _bounded_sleep(_delay):
|
||||
nonlocal sleep_calls
|
||||
sleep_calls += 1
|
||||
if sleep_calls >= count:
|
||||
runner._running = False
|
||||
|
||||
monkeypatch.setattr(asyncio, "sleep", _bounded_sleep)
|
||||
|
||||
|
||||
def test_duplicate_async_queue_replay_injects_once(monkeypatch, isolated_registry):
|
||||
"""Byte-identical queue replays produce one turn in one gateway lifecycle."""
|
||||
isolated = queue.Queue()
|
||||
monkeypatch.setattr(isolated_registry, "completion_queue", isolated)
|
||||
isolated.put(dict(_async_event()))
|
||||
isolated.put(dict(_async_event()))
|
||||
|
||||
adapter = SimpleNamespace(handle_message=AsyncMock())
|
||||
runner = _runner(adapter)
|
||||
_stop_after_sleeps(monkeypatch, runner, count=2)
|
||||
|
||||
asyncio.run(runner._async_delegation_watcher(interval=0))
|
||||
|
||||
adapter.handle_message.assert_awaited_once()
|
||||
|
||||
|
||||
def test_unroutable_async_event_is_not_requeued_forever(
|
||||
monkeypatch, isolated_registry,
|
||||
):
|
||||
isolated = queue.Queue()
|
||||
monkeypatch.setattr(isolated_registry, "completion_queue", isolated)
|
||||
event = _async_event("deleg_desktop_or_cli")
|
||||
event["session_key"] = "20260711_unparseable_ui_session"
|
||||
isolated.put(event)
|
||||
|
||||
adapter = SimpleNamespace(handle_message=AsyncMock())
|
||||
runner = _runner(adapter)
|
||||
_stop_after_sleeps(monkeypatch, runner, count=2)
|
||||
|
||||
asyncio.run(runner._async_delegation_watcher(interval=0))
|
||||
|
||||
adapter.handle_message.assert_not_awaited()
|
||||
assert isolated.empty()
|
||||
|
||||
|
||||
def test_concurrent_claims_share_the_same_narrow_delivery_seam():
|
||||
"""Concurrent consumers in one runner cannot both enter the adapter."""
|
||||
entered = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
|
||||
async def _blocked_injection(_event):
|
||||
entered.set()
|
||||
await release.wait()
|
||||
|
||||
adapter = SimpleNamespace(handle_message=AsyncMock(side_effect=_blocked_injection))
|
||||
runner = _runner(adapter)
|
||||
event = _async_event()
|
||||
text = "completion"
|
||||
|
||||
async def _exercise():
|
||||
first = asyncio.create_task(runner._deliver_completion_notification(text, dict(event)))
|
||||
await entered.wait()
|
||||
second = asyncio.create_task(runner._deliver_completion_notification(text, dict(event)))
|
||||
await asyncio.sleep(0)
|
||||
release.set()
|
||||
return await asyncio.gather(first, second)
|
||||
|
||||
assert sorted(asyncio.run(_exercise()), key=str) == [None, True]
|
||||
adapter.handle_message.assert_awaited_once()
|
||||
|
||||
|
||||
def test_failed_async_injection_is_retried_and_only_success_is_acked(
|
||||
monkeypatch, isolated_registry,
|
||||
):
|
||||
isolated = queue.Queue()
|
||||
monkeypatch.setattr(isolated_registry, "completion_queue", isolated)
|
||||
isolated.put(_async_event())
|
||||
|
||||
adapter = SimpleNamespace(
|
||||
handle_message=AsyncMock(side_effect=[RuntimeError("temporary"), None])
|
||||
)
|
||||
runner = _runner(adapter)
|
||||
_stop_after_sleeps(monkeypatch, runner, count=3)
|
||||
|
||||
from tools import async_delegation
|
||||
|
||||
acknowledgements = []
|
||||
monkeypatch.setattr(
|
||||
async_delegation,
|
||||
"mark_completion_delivered",
|
||||
lambda delegation_id: acknowledgements.append(delegation_id) or True,
|
||||
raising=False,
|
||||
)
|
||||
|
||||
asyncio.run(runner._async_delegation_watcher(interval=0))
|
||||
|
||||
assert adapter.handle_message.await_count == 2
|
||||
assert acknowledgements == ["deleg_duplicate"]
|
||||
|
||||
|
||||
def test_distinct_process_incarnations_are_not_deduplicated():
|
||||
"""Producer spawn time distinguishes a reused process session ID."""
|
||||
adapter = SimpleNamespace(handle_message=AsyncMock())
|
||||
runner = _runner(adapter)
|
||||
|
||||
async def _exercise():
|
||||
first = await runner._deliver_completion_notification(
|
||||
"first", _completion_event(started_at=10.0)
|
||||
)
|
||||
second = await runner._deliver_completion_notification(
|
||||
"second", _completion_event(started_at=20.0)
|
||||
)
|
||||
return first, second
|
||||
|
||||
assert asyncio.run(_exercise()) == (True, True)
|
||||
|
||||
assert adapter.handle_message.await_count == 2
|
||||
|
||||
|
||||
def test_delivered_identity_retention_is_bounded():
|
||||
"""Lifecycle dedupe cannot grow without bound in a long-running gateway."""
|
||||
adapter = SimpleNamespace(handle_message=AsyncMock())
|
||||
runner = _runner(adapter)
|
||||
runner._completion_delivery_retention = 2
|
||||
runner._completion_deliveries_delivered = OrderedDict()
|
||||
|
||||
async def _exercise():
|
||||
for index in range(3):
|
||||
await runner._deliver_completion_notification(
|
||||
f"completion {index}",
|
||||
_async_event(f"deleg_retention_{index}"),
|
||||
)
|
||||
|
||||
asyncio.run(_exercise())
|
||||
|
||||
assert len(runner._completion_deliveries_delivered) == 2
|
||||
assert ("async_delegation", "deleg_retention_0", "") not in (
|
||||
runner._completion_deliveries_delivered
|
||||
)
|
||||
assert ("async_delegation", "deleg_retention_2", "") in (
|
||||
runner._completion_deliveries_delivered
|
||||
)
|
||||
|
||||
|
||||
def test_delivery_state_is_isolated_per_gateway_profile_lifecycle():
|
||||
"""A process-local claim in one profile never suppresses another runner."""
|
||||
default_adapter = SimpleNamespace(handle_message=AsyncMock())
|
||||
profile_adapter = SimpleNamespace(handle_message=AsyncMock())
|
||||
default_runner = _runner(default_adapter)
|
||||
profile_runner = _runner(profile_adapter)
|
||||
event = _async_event("deleg_same_producer_id")
|
||||
|
||||
async def _exercise():
|
||||
first = await default_runner._deliver_completion_notification(
|
||||
"default", dict(event),
|
||||
)
|
||||
second = await profile_runner._deliver_completion_notification(
|
||||
"profile", dict(event),
|
||||
)
|
||||
return first, second
|
||||
|
||||
assert asyncio.run(_exercise()) == (True, True)
|
||||
default_adapter.handle_message.assert_awaited_once()
|
||||
profile_adapter.handle_message.assert_awaited_once()
|
||||
|
||||
|
||||
def test_async_completion_uses_canonical_origin_routing(monkeypatch, isolated_registry):
|
||||
isolated = queue.Queue()
|
||||
monkeypatch.setattr(isolated_registry, "completion_queue", isolated)
|
||||
event = _async_event("deleg_routing")
|
||||
isolated.put(event)
|
||||
|
||||
canonical = SessionSource(
|
||||
platform=Platform.TELEGRAM,
|
||||
chat_id="canonical-chat",
|
||||
chat_type="group",
|
||||
thread_id="canonical-topic",
|
||||
)
|
||||
entry = SimpleNamespace(origin=canonical)
|
||||
adapter = SimpleNamespace(handle_message=AsyncMock())
|
||||
runner = _runner(adapter, origins={event["session_key"]: entry})
|
||||
_stop_after_sleeps(monkeypatch, runner, count=2)
|
||||
|
||||
asyncio.run(runner._async_delegation_watcher(interval=0))
|
||||
|
||||
delivered = adapter.handle_message.await_args.args[0]
|
||||
assert delivered.source == canonical
|
||||
|
||||
|
||||
def test_explicit_kill_returns_output_before_consuming_notification(monkeypatch):
|
||||
import tools.process_registry as pr_module
|
||||
|
||||
registry = ProcessRegistry()
|
||||
session = ProcessSession(
|
||||
id="proc_kill_consumed",
|
||||
command="sleep 999",
|
||||
task_id="task",
|
||||
started_at=1.0,
|
||||
output_buffer="important terminal output\n",
|
||||
notify_on_complete=True,
|
||||
)
|
||||
session.process = MagicMock()
|
||||
session.process.pid = 4242
|
||||
registry._running[session.id] = session
|
||||
monkeypatch.setattr(registry, "_terminate_host_pid", lambda *_a, **_kw: None)
|
||||
monkeypatch.setattr(registry, "_write_checkpoint", lambda: None)
|
||||
monkeypatch.setattr(pr_module, "process_registry", registry)
|
||||
|
||||
result = registry.kill_process(session.id)
|
||||
assert result["status"] == "killed"
|
||||
assert result["output"] == "important terminal output\n"
|
||||
assert registry.is_completion_consumed(session.id)
|
||||
|
||||
adapter = SimpleNamespace(handle_message=AsyncMock())
|
||||
runner = _runner(adapter)
|
||||
|
||||
async def _instant_sleep(*_a, **_kw):
|
||||
pass
|
||||
|
||||
monkeypatch.setattr(asyncio, "sleep", _instant_sleep)
|
||||
asyncio.run(runner._run_process_watcher({
|
||||
"session_id": session.id,
|
||||
"check_interval": 0,
|
||||
"session_key": "agent:main:telegram:dm:123",
|
||||
"platform": "telegram",
|
||||
"chat_type": "dm",
|
||||
"chat_id": "123",
|
||||
"notify_on_complete": True,
|
||||
}))
|
||||
|
||||
adapter.handle_message.assert_not_awaited()
|
||||
|
||||
|
||||
def test_process_tool_redacts_explicit_kill_output(monkeypatch):
|
||||
from tools import process_registry as pr_module
|
||||
|
||||
registry = ProcessRegistry()
|
||||
session = ProcessSession(
|
||||
id="proc_kill_redacted",
|
||||
command="printenv",
|
||||
task_id="task",
|
||||
started_at=1.0,
|
||||
output_buffer="PRIVATE_TOKEN=opaque-value\n",
|
||||
exited=True,
|
||||
exit_code=0,
|
||||
)
|
||||
registry._finished[session.id] = session
|
||||
monkeypatch.setattr(pr_module, "process_registry", registry)
|
||||
|
||||
def _redact(result):
|
||||
assert result["output"] == "PRIVATE_TOKEN=opaque-value\n"
|
||||
result["output"] = "PRIVATE_TOKEN=<redacted>\n"
|
||||
return result
|
||||
|
||||
monkeypatch.setattr(pr_module, "_redact_process_result", _redact)
|
||||
|
||||
result = json.loads(pr_module._handle_process({
|
||||
"action": "kill",
|
||||
"session_id": session.id,
|
||||
}))
|
||||
assert result["output"] == "PRIVATE_TOKEN=<redacted>\n"
|
||||
|
||||
|
||||
def test_kill_of_already_exited_process_returns_output_before_consuming():
|
||||
registry = ProcessRegistry()
|
||||
session = ProcessSession(
|
||||
id="proc_already_exited",
|
||||
command="echo complete",
|
||||
task_id="task",
|
||||
started_at=1.0,
|
||||
output_buffer="complete\n",
|
||||
exited=True,
|
||||
exit_code=0,
|
||||
)
|
||||
registry._finished[session.id] = session
|
||||
|
||||
result = registry.kill_process(session.id)
|
||||
|
||||
assert result["status"] == "already_exited"
|
||||
assert result["output"] == "complete\n"
|
||||
assert registry.is_completion_consumed(session.id)
|
||||
|
||||
|
||||
def test_read_log_only_consumes_when_terminal_output_page_is_observed():
|
||||
registry = ProcessRegistry()
|
||||
session = ProcessSession(
|
||||
id="proc_paged_log",
|
||||
command="printf lines",
|
||||
task_id="task",
|
||||
started_at=1.0,
|
||||
output_buffer="first\nsecond\nfinal\n",
|
||||
exited=True,
|
||||
exit_code=0,
|
||||
)
|
||||
registry._finished[session.id] = session
|
||||
|
||||
middle_page = registry.read_log(session.id, offset=1, limit=1)
|
||||
assert middle_page["output"] == "second"
|
||||
assert not registry.is_completion_consumed(session.id)
|
||||
|
||||
final_page = registry.read_log(session.id, offset=2, limit=1)
|
||||
assert final_page["output"] == "final"
|
||||
assert registry.is_completion_consumed(session.id)
|
||||
|
||||
|
||||
def test_bulk_kill_does_not_consume_discarded_completion_output(monkeypatch):
|
||||
registry = ProcessRegistry()
|
||||
session = ProcessSession(
|
||||
id="proc_bulk_kill",
|
||||
command="sleep 999",
|
||||
task_id="task",
|
||||
started_at=1.0,
|
||||
output_buffer="output bulk cleanup does not return\n",
|
||||
notify_on_complete=True,
|
||||
)
|
||||
session.process = MagicMock()
|
||||
session.process.pid = 4243
|
||||
registry._running[session.id] = session
|
||||
monkeypatch.setattr(registry, "_terminate_host_pid", lambda *_a, **_kw: None)
|
||||
monkeypatch.setattr(registry, "_write_checkpoint", lambda: None)
|
||||
|
||||
assert registry.kill_all() == 1
|
||||
assert not registry.is_completion_consumed(session.id)
|
||||
queued = registry.completion_queue.get_nowait()
|
||||
assert queued["session_id"] == session.id
|
||||
assert queued["started_at"] == session.started_at
|
||||
assert queued["output"] == "output bulk cleanup does not return\n"
|
||||
|
||||
|
||||
def test_unobserved_normal_completion_still_notifies(monkeypatch):
|
||||
import tools.process_registry as pr_module
|
||||
|
||||
class _Registry:
|
||||
def get(self, _session_id):
|
||||
return SimpleNamespace(
|
||||
output_buffer="done\n",
|
||||
exited=True,
|
||||
exit_code=0,
|
||||
command="echo done",
|
||||
started_at=1234.5,
|
||||
)
|
||||
|
||||
def is_completion_consumed(self, _session_id):
|
||||
return False
|
||||
|
||||
monkeypatch.setattr(pr_module, "process_registry", _Registry())
|
||||
adapter = SimpleNamespace(handle_message=AsyncMock())
|
||||
runner = _runner(adapter)
|
||||
|
||||
async def _instant_sleep(*_a, **_kw):
|
||||
pass
|
||||
|
||||
monkeypatch.setattr(asyncio, "sleep", _instant_sleep)
|
||||
asyncio.run(runner._run_process_watcher({
|
||||
"session_id": "proc_unobserved",
|
||||
"check_interval": 0,
|
||||
"session_key": "agent:main:telegram:dm:123",
|
||||
"platform": "telegram",
|
||||
"chat_type": "dm",
|
||||
"chat_id": "123",
|
||||
"notify_on_complete": True,
|
||||
}))
|
||||
|
||||
adapter.handle_message.assert_awaited_once()
|
||||
|
||||
|
||||
def test_autonomous_completion_redacts_real_command_and_output_secrets(monkeypatch):
|
||||
import agent.redact as redact_module
|
||||
import tools.process_registry as pr_module
|
||||
|
||||
secret = "abc123randomopaquetokenvalue999"
|
||||
registry = ProcessRegistry()
|
||||
session = ProcessSession(
|
||||
id="proc_autonomous_redaction",
|
||||
command=f"printenv MY_SERVICE_TOKEN={secret}",
|
||||
task_id="task",
|
||||
started_at=1234.5,
|
||||
output_buffer=f"MY_SERVICE_TOKEN={secret}\nHOME=/home/user\n",
|
||||
exited=True,
|
||||
exit_code=0,
|
||||
notify_on_complete=True,
|
||||
)
|
||||
registry._finished[session.id] = session
|
||||
monkeypatch.setattr(pr_module, "process_registry", registry)
|
||||
monkeypatch.setattr(redact_module, "_REDACT_ENABLED", True)
|
||||
|
||||
adapter = SimpleNamespace(handle_message=AsyncMock())
|
||||
runner = _runner(adapter)
|
||||
|
||||
async def _instant_sleep(*_a, **_kw):
|
||||
pass
|
||||
|
||||
monkeypatch.setattr(asyncio, "sleep", _instant_sleep)
|
||||
asyncio.run(runner._run_process_watcher({
|
||||
"session_id": session.id,
|
||||
"check_interval": 0,
|
||||
"session_key": "agent:main:telegram:dm:123",
|
||||
"platform": "telegram",
|
||||
"chat_type": "dm",
|
||||
"chat_id": "123",
|
||||
"notify_on_complete": True,
|
||||
}))
|
||||
|
||||
delivered = adapter.handle_message.await_args.args[0]
|
||||
assert secret not in delivered.text
|
||||
assert "HOME=/home/user" in delivered.text
|
||||
|
|
@ -1091,6 +1091,10 @@ class ProcessRegistry:
|
|||
"completion_reason": session.completion_reason,
|
||||
"termination_source": session.termination_source,
|
||||
"output": output_tail,
|
||||
# Stable producer identity across checkpoint recovery; unlike
|
||||
# a consumer-observed completion timestamp, this does not vary
|
||||
# based on which watcher notices exit first.
|
||||
"started_at": session.started_at,
|
||||
})
|
||||
|
||||
# ----- Query Methods -----
|
||||
|
|
@ -1347,8 +1351,13 @@ class ProcessRegistry:
|
|||
# Default: last N lines
|
||||
if offset == 0 and limit > 0:
|
||||
selected = lines[-limit:]
|
||||
observed_completion_output = bool(selected) or total_lines == 0
|
||||
else:
|
||||
selected = lines[offset:offset + limit]
|
||||
stop = slice(offset, offset + limit).indices(total_lines)[1]
|
||||
observed_completion_output = (
|
||||
total_lines == 0 or (bool(selected) and stop == total_lines)
|
||||
)
|
||||
|
||||
result = {
|
||||
"session_id": session.id,
|
||||
|
|
@ -1358,7 +1367,7 @@ class ProcessRegistry:
|
|||
"total_lines": total_lines,
|
||||
"showing": f"{len(selected)} lines",
|
||||
}
|
||||
if session.exited:
|
||||
if session.exited and observed_completion_output:
|
||||
self._completion_consumed.add(session_id)
|
||||
return result
|
||||
|
||||
|
|
@ -1449,17 +1458,41 @@ class ProcessRegistry:
|
|||
result["timeout_note"] = f"Waited {effective_timeout}s, process still running"
|
||||
return result
|
||||
|
||||
def kill_process(self, session_id: str, *, source: str = "process.kill") -> dict:
|
||||
"""Kill a background process."""
|
||||
def kill_process(
|
||||
self,
|
||||
session_id: str,
|
||||
*,
|
||||
source: str = "process.kill",
|
||||
consume_output: bool = True,
|
||||
) -> dict:
|
||||
"""Kill a background process and return its output snapshot.
|
||||
|
||||
``consume_output`` is true for explicit tool/RPC kills because their
|
||||
caller observes the returned output. Bulk cleanup passes false: it
|
||||
discards each result and therefore must not suppress an autonomous
|
||||
output-bearing completion notification.
|
||||
"""
|
||||
from tools.ansi_strip import strip_ansi
|
||||
|
||||
session = self.get(session_id)
|
||||
if session is None:
|
||||
return {"status": "not_found", "error": f"No process with ID {session_id}"}
|
||||
|
||||
if session.exited:
|
||||
return {
|
||||
"status": "already_exited",
|
||||
"exit_code": session.exit_code,
|
||||
}
|
||||
with session._lock:
|
||||
result = {
|
||||
"status": "already_exited",
|
||||
"command": session.command,
|
||||
"exit_code": session.exit_code,
|
||||
"completion_reason": session.completion_reason,
|
||||
"termination_source": session.termination_source,
|
||||
"output": strip_ansi(session.output_buffer[-2000:]),
|
||||
}
|
||||
# Only suppress the autonomous turn after its output is present in
|
||||
# the explicit kill result, matching wait/log consumption.
|
||||
if consume_output:
|
||||
self._completion_consumed.add(session_id)
|
||||
return result
|
||||
|
||||
# Kill via PTY, Popen (local), or env execute (non-local)
|
||||
try:
|
||||
|
|
@ -1486,10 +1519,14 @@ class ProcessRegistry:
|
|||
with session._lock:
|
||||
session.exited = True
|
||||
session.exit_code = None
|
||||
output = strip_ansi(session.output_buffer[-2000:])
|
||||
if consume_output:
|
||||
self._completion_consumed.add(session_id)
|
||||
self._move_to_finished(session)
|
||||
return {
|
||||
"status": "already_exited",
|
||||
"exit_code": session.exit_code,
|
||||
"output": output,
|
||||
}
|
||||
self._terminate_host_pid(session.pid, session.host_start_time)
|
||||
else:
|
||||
|
|
@ -1500,10 +1537,17 @@ class ProcessRegistry:
|
|||
"its original runtime handle is no longer available"
|
||||
),
|
||||
}
|
||||
session.exited = True
|
||||
session.exit_code = -15 # SIGTERM
|
||||
session.completion_reason = "killed"
|
||||
session.termination_source = source
|
||||
# Capture output before marking consumed, then mark consumed before
|
||||
# exposing ``exited`` to watcher tasks. This closes the delayed
|
||||
# notification race without discarding the terminal transcript.
|
||||
with session._lock:
|
||||
output = strip_ansi(session.output_buffer[-2000:])
|
||||
if consume_output:
|
||||
self._completion_consumed.add(session_id)
|
||||
session.exited = True
|
||||
session.exit_code = -15 # SIGTERM
|
||||
session.completion_reason = "killed"
|
||||
session.termination_source = source
|
||||
self._move_to_finished(session)
|
||||
self._write_checkpoint()
|
||||
return {
|
||||
|
|
@ -1511,6 +1555,7 @@ class ProcessRegistry:
|
|||
"session_id": session.id,
|
||||
"completion_reason": session.completion_reason,
|
||||
"termination_source": session.termination_source,
|
||||
"output": output,
|
||||
}
|
||||
except Exception as e:
|
||||
return {"status": "error", "error": str(e)}
|
||||
|
|
@ -1747,7 +1792,11 @@ class ProcessRegistry:
|
|||
|
||||
killed = 0
|
||||
for session in targets:
|
||||
result = self.kill_process(session.id, source="kill_all")
|
||||
result = self.kill_process(
|
||||
session.id,
|
||||
source="kill_all",
|
||||
consume_output=False,
|
||||
)
|
||||
if result.get("status") in {"killed", "already_exited"}:
|
||||
killed += 1
|
||||
return killed
|
||||
|
|
@ -2238,7 +2287,10 @@ def _handle_process(args, **kw):
|
|||
elif action == "wait":
|
||||
return json.dumps(_redact_process_result(process_registry.wait(session_id, timeout=args.get("timeout"))), ensure_ascii=False)
|
||||
elif action == "kill":
|
||||
return json.dumps(process_registry.kill_process(session_id), ensure_ascii=False)
|
||||
return json.dumps(
|
||||
_redact_process_result(process_registry.kill_process(session_id)),
|
||||
ensure_ascii=False,
|
||||
)
|
||||
elif action == "write":
|
||||
return json.dumps(process_registry.write_stdin(session_id, str(args.get("data", ""))), ensure_ascii=False)
|
||||
elif action == "submit":
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue