Import AITURK IDE 1.0.0-beta.1 from Hermes 63279301; preserve MIT license

tests/agent/test_relay_session_segments.py

@@ -0,0 +1,413 @@
+"""Session-span segmentation for continuous sessions.
+
+Continuous gateway sessions keep the Relay session scope open indefinitely;
+close-driven export means the session root span (and out-of-turn marks) are
+unexported until /new or idle-end, and a crash loses the whole segment.
+
+Segmentation closes the current session scope at a TURN BOUNDARY and pushes
+a fresh one, chaining segments via metadata:
+
+  gateway.telemetry.session_segments.on_compaction  (default False)
+  gateway.telemetry.session_segments.max_turns      (default 0 = unlimited)
+
+Both defaults off => behavior identical to today (no rotation, ever).
+Rotation never happens mid-turn: compaction only sets rotate_pending,
+consumed at the next begin_turn before the turn scope pushes.
+"""
+
+from __future__ import annotations
+
+import threading
+from typing import Any
+
+import pytest
+
+from agent import relay_runtime
+from agent.relay_runtime import (
+    RelayRuntime,
+    RelaySessionCoordinator,
+)
+
+
+class _ScopeHandle:
+    def __init__(self, name: str, seq: int) -> None:
+        self.name = name
+        self.seq = seq
+
+
+class _FakeScopeModule:
+    def __init__(self, wedge_pop: threading.Event | None = None) -> None:
+        self._wedge = wedge_pop
+        self._seq = 0
+        self.pushes: list[dict[str, Any]] = []  # {name, metadata, handle}
+        self.pops: list[_ScopeHandle] = []
+
+    def push(self, name: str, scope_type: Any, **kwargs: Any) -> _ScopeHandle:
+        self._seq += 1
+        self.pushes.append(
+            {
+                "name": name,
+                "metadata": dict(kwargs.get("metadata") or {}),
+                "parent": kwargs.get("handle"),
+                "seq": self._seq,
+            }
+        )
+        return _ScopeHandle(name, self._seq)
+
+    def pop(self, handle: _ScopeHandle, **kwargs: Any) -> None:
+        if self._wedge is not None:
+            self._wedge.wait()
+        self.pops.append(handle)
+
+    def event(self, *args: Any, **kwargs: Any) -> None:
+        return None
+
+
+class _FakeSubscribers:
+    def __init__(self) -> None:
+        self.flushed = 0
+
+    def flush(self) -> None:
+        self.flushed += 1
+
+
+class _FakeScopeType:
+    Function = "function"
+    Agent = "agent"
+
+
+class _FakeRelay:
+    def __init__(self, wedge_pop: threading.Event | None = None) -> None:
+        self.scope = _FakeScopeModule(wedge_pop)
+        self.subscribers = _FakeSubscribers()
+        self.ScopeType = _FakeScopeType()
+
+    def get_scope_stack(self) -> None:
+        return None
+
+
+_LIVE: list[tuple[RelayRuntime, _FakeRelay]] = []
+
+
+def _make_runtime(fake: _FakeRelay) -> RelayRuntime:
+    runtime = RelayRuntime(relay=fake, profile_key="/tmp/test-profile")
+    _LIVE.append((runtime, fake))
+    return runtime
+
+
+@pytest.fixture(autouse=True)
+def _teardown_runtimes():
+    """Unwedge and drain every runtime so exit paths never replay wedged ops."""
+    yield
+    for runtime, fake in _LIVE:
+        if fake.scope._wedge is not None:
+            fake.scope._wedge.set()
+        runtime.shutdown()
+    _LIVE.clear()
+
+
+@pytest.fixture(autouse=True)
+def _fast_scope_timeout(monkeypatch):
+    monkeypatch.setattr(relay_runtime, "_SCOPE_OP_TIMEOUT", 1.0)
+
+
+@pytest.fixture(autouse=True)
+def _default_config(monkeypatch):
+    """No config on disk by default; tests override _segments_config directly."""
+    monkeypatch.setattr(
+        "gateway.run._load_gateway_config", lambda: {}, raising=False
+    )
+    relay_runtime._reset_segments_config_for_tests()
+
+
+def _set_segments(monkeypatch, *, on_compaction=False, max_turns=0):
+    monkeypatch.setattr(
+        "gateway.run._load_gateway_config",
+        lambda: {
+            "gateway": {
+                "telemetry": {
+                    "session_segments": {
+                        "on_compaction": on_compaction,
+                        "max_turns": max_turns,
+                    }
+                }
+            }
+        },
+        raising=False,
+    )
+    relay_runtime._reset_segments_config_for_tests()
+
+
+@pytest.fixture()
+def coordinator() -> RelaySessionCoordinator:
+    return RelaySessionCoordinator()
+
+
+def _acquire(coordinator, runtime, session_id="sess-1"):
+    class _Registry:
+        def for_profile(self, key):
+            return runtime
+
+    coordinator.registry = _Registry()
+    coordinator._prepare_session = lambda host, ctx: None
+    return coordinator.acquire_conversation(
+        profile_key=runtime.profile_key,
+        session_id=session_id,
+        platform="test",
+    )
+
+
+def _session_pushes(fake):
+    return [p for p in fake.scope.pushes if p["name"] == relay_runtime.SESSION_SCOPE]
+
+
+def _run_turn(coordinator, lease, turn_id):
+    turn = coordinator.begin_turn(lease, turn_id=turn_id, task_id=f"task-{turn_id}")
+    coordinator.end_turn(turn, outcome="success")
+    return turn
+
+
+class TestDefaultsNeverRotate:
+    def test_no_rotation_across_many_turns_and_compactions(self, coordinator):
+        fake = _FakeRelay()
+        runtime = _make_runtime(fake)
+        lease = _acquire(coordinator, runtime)
+        assert lease.session is not None
+
+        coordinator.notify_session_compacted(
+            profile_key=runtime.profile_key, session_id="sess-1"
+        )
+        for i in range(5):
+            _run_turn(coordinator, lease, f"t{i}")
+
+        assert len(_session_pushes(fake)) == 1, (
+            "defaults off must never rotate the session scope — "
+            "today's behavior is the contract"
+        )
+
+
+class TestCompactionRotation:
+    def test_compaction_rotates_at_next_begin_turn_not_immediately(
+        self, coordinator, monkeypatch
+    ):
+        _set_segments(monkeypatch, on_compaction=True)
+        fake = _FakeRelay()
+        runtime = _make_runtime(fake)
+        lease = _acquire(coordinator, runtime)
+        original_handle = lease.session.handle
+
+        coordinator.notify_session_compacted(
+            profile_key=runtime.profile_key, session_id="sess-1"
+        )
+        # No rotation yet — compaction only flags; scope stack untouched.
+        assert len(_session_pushes(fake)) == 1
+        assert not fake.scope.pops
+
+        turn = coordinator.begin_turn(lease, turn_id="t1", task_id="task1")
+        sessions = _session_pushes(fake)
+        assert len(sessions) == 2, "rotation must happen at the next begin_turn"
+        # Old session scope was popped before the new push.
+        assert any(p.seq == 1 for p in fake.scope.pops), "old segment scope popped"
+        assert lease.session.handle is not original_handle
+        # The turn scope parents to the NEW segment handle.
+        turn_push = [p for p in fake.scope.pushes if p["name"] == relay_runtime.TURN_SCOPE][-1]
+        assert turn_push["parent"] is lease.session.handle
+        coordinator.end_turn(turn, outcome="success")
+
+    def test_segment_metadata_on_rotated_scope(self, coordinator, monkeypatch):
+        _set_segments(monkeypatch, on_compaction=True)
+        fake = _FakeRelay()
+        runtime = _make_runtime(fake)
+        lease = _acquire(coordinator, runtime)
+
+        coordinator.notify_session_compacted(
+            profile_key=runtime.profile_key, session_id="sess-1"
+        )
+        _run_turn(coordinator, lease, "t1")
+
+        new_seg = _session_pushes(fake)[-1]["metadata"]
+        assert new_seg.get("hermes.session.segment") == 1
+        assert new_seg.get("hermes.session.segment_reason") == "compaction"
+
+    def test_unknown_session_compaction_is_noop(self, coordinator, monkeypatch):
+        _set_segments(monkeypatch, on_compaction=True)
+        fake = _FakeRelay()
+        runtime = _make_runtime(fake)
+        _acquire(coordinator, runtime)
+        # Must not raise, must not rotate anything.
+        coordinator.notify_session_compacted(
+            profile_key=runtime.profile_key, session_id="never-seen"
+        )
+        assert len(_session_pushes(fake)) == 1
+
+    def test_rotating_compaction_closes_old_session_scope(
+        self, coordinator, monkeypatch
+    ):
+        """Legacy compaction rotates to a child session id: the OLD session's
+        scope must close (export) instead of orphaning unexported forever."""
+        _set_segments(monkeypatch, on_compaction=True)
+        fake = _FakeRelay()
+        runtime = _make_runtime(fake)
+        _acquire(coordinator, runtime, session_id="parent-1")
+        assert not fake.scope.pops
+
+        coordinator.notify_session_compacted(
+            profile_key=runtime.profile_key,
+            session_id="child-1",
+            old_session_id="parent-1",
+        )
+        assert len(fake.scope.pops) == 1, (
+            "rotating compaction must close the old session scope"
+        )
+        # Subscriber flushing is process-wide and happens once at final plugin
+        # teardown, after all sessions have drained. Flushing on this per-session
+        # close can block an active asyncio loop owned by another session.
+        assert fake.subscribers.flushed == 0
+
+    def test_rotating_compaction_mid_turn_defers_close_to_end_turn(
+        self, coordinator, monkeypatch
+    ):
+        """A rotating compaction completing while a turn is LIVE on the old
+        session must NOT close the session scope immediately — that would pop
+        it under the live turn scope (LIFO violation). The close defers to
+        that turn's end_turn."""
+        _set_segments(monkeypatch, on_compaction=True)
+        fake = _FakeRelay()
+        runtime = _make_runtime(fake)
+        lease = _acquire(coordinator, runtime, session_id="parent-1")
+
+        turn = coordinator.begin_turn(lease, turn_id="t1", task_id="task1")
+        coordinator.notify_session_compacted(
+            profile_key=runtime.profile_key,
+            session_id="child-1",
+            old_session_id="parent-1",
+        )
+        # No pops yet: neither the turn scope nor the session scope closed.
+        assert not fake.scope.pops, (
+            "old-session close must defer while its turn is live"
+        )
+
+        coordinator.end_turn(turn, outcome="success")
+        # Turn scope popped first, then the deferred session close popped
+        # the session scope — LIFO order preserved.
+        assert len(fake.scope.pops) == 2, "end_turn must consume deferred close"
+        assert fake.scope.pops[0].name == relay_runtime.TURN_SCOPE, (
+            "turn scope must pop before the session scope"
+        )
+        assert fake.scope.pops[-1].name == relay_runtime.SESSION_SCOPE
+        assert runtime.get_session("parent-1") is None
+
+    def test_rotating_compaction_noop_when_disabled(self, coordinator, monkeypatch):
+        fake = _FakeRelay()
+        runtime = _make_runtime(fake)
+        _acquire(coordinator, runtime, session_id="parent-1")
+        coordinator.notify_session_compacted(
+            profile_key=runtime.profile_key,
+            session_id="child-1",
+            old_session_id="parent-1",
+        )
+        assert not fake.scope.pops, "defaults off: rotating compaction is a no-op"
+
+
+class TestMaxTurnsRotation:
+    def test_rotates_after_cap(self, coordinator, monkeypatch):
+        _set_segments(monkeypatch, max_turns=2)
+        fake = _FakeRelay()
+        runtime = _make_runtime(fake)
+        lease = _acquire(coordinator, runtime)
+
+        for i in range(5):
+            _run_turn(coordinator, lease, f"t{i}")
+
+        # turns 0,1 in segment 0; rotation before turn 2; turns 2,3 in
+        # segment 1; rotation before turn 4.
+        sessions = _session_pushes(fake)
+        assert len(sessions) == 3, "cap of 2 over 5 turns => 2 rotations"
+        assert sessions[-1]["metadata"].get("hermes.session.segment_reason") == "max_turns"
+
+    def test_zero_cap_means_unlimited(self, coordinator, monkeypatch):
+        _set_segments(monkeypatch, max_turns=0)
+        fake = _FakeRelay()
+        runtime = _make_runtime(fake)
+        lease = _acquire(coordinator, runtime)
+        for i in range(4):
+            _run_turn(coordinator, lease, f"t{i}")
+        assert len(_session_pushes(fake)) == 1
+
+
+class TestRotationSafety:
+    def test_never_rotates_mid_turn(self, coordinator, monkeypatch):
+        _set_segments(monkeypatch, on_compaction=True)
+        fake = _FakeRelay()
+        runtime = _make_runtime(fake)
+        lease = _acquire(coordinator, runtime)
+
+        turn = coordinator.begin_turn(lease, turn_id="t1", task_id="task1")
+        # Compaction lands while the turn is LIVE.
+        coordinator.notify_session_compacted(
+            profile_key=runtime.profile_key, session_id="sess-1"
+        )
+        assert len(_session_pushes(fake)) == 1, "no rotation while a turn is live"
+        coordinator.end_turn(turn, outcome="success")
+        assert len(_session_pushes(fake)) == 1, "end_turn does not rotate either"
+
+        # The NEXT turn consumes the pending rotation.
+        turn2 = coordinator.begin_turn(lease, turn_id="t2", task_id="task2")
+        assert len(_session_pushes(fake)) == 2
+        coordinator.end_turn(turn2, outcome="success")
+
+    def test_wedged_rotation_is_bounded_and_agent_continues(
+        self, coordinator, monkeypatch
+    ):
+        _set_segments(monkeypatch, on_compaction=True)
+        wedge = threading.Event()  # never set until teardown
+        fake = _FakeRelay(wedge_pop=wedge)
+        runtime = _make_runtime(fake)
+        lease = _acquire(coordinator, runtime)
+        coordinator.notify_session_compacted(
+            profile_key=runtime.profile_key, session_id="sess-1"
+        )
+
+        result: list[Any] = []
+
+        def _begin():
+            result.append(
+                coordinator.begin_turn(lease, turn_id="t1", task_id="task1")
+            )
+
+        worker = threading.Thread(target=_begin, daemon=True)
+        worker.start()
+        worker.join(5.0)
+        assert not worker.is_alive(), (
+            "begin_turn must return even when the rotation pop wedges — "
+            "a wedged pipeline costs one segment span, never the agent"
+        )
+        turn = result[0]
+        coordinator.end_turn(turn, outcome="success")
+
+    def test_subagent_children_parent_to_new_segment_after_rotation(
+        self, coordinator, monkeypatch
+    ):
+        _set_segments(monkeypatch, on_compaction=True)
+        fake = _FakeRelay()
+        runtime = _make_runtime(fake)
+        lease = _acquire(coordinator, runtime)
+        coordinator.notify_session_compacted(
+            profile_key=runtime.profile_key, session_id="sess-1"
+        )
+        _run_turn(coordinator, lease, "t1")  # consumes rotation
+        new_handle = lease.session.handle
+
+        child = runtime.register_subagent(
+            {"parent_session_id": "sess-1", "child_session_id": "child-1"}
+        )
+        assert child is not None
+        child_push = [
+            p
+            for p in fake.scope.pushes
+            if p["name"] == relay_runtime.SESSION_SCOPE
+            and p["parent"] is not None
+        ][-1]
+        assert child_push["parent"] is new_handle, (
+            "post-rotation children must parent to the new segment handle"
+        )