288 lines
12 KiB
Python
288 lines
12 KiB
Python
"""Run real streaming lifecycle code at the two control-plane boundaries.
|
|
|
|
Only the external Agent/provider is replaced; barriers exercise the real worker,
|
|
Stop, Steer, registry transitions and leftover SSE, without provider/tool calls.
|
|
"""
|
|
import queue
|
|
import sys
|
|
import threading
|
|
import types
|
|
from collections import OrderedDict
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from api import config, models, streaming
|
|
from api.models import Session
|
|
from tests.test_real_steer import _captured_response, _make_handler
|
|
from tests.test_steer_cancel_ordering import ObservedLock
|
|
|
|
|
|
@pytest.fixture
|
|
def worker_scene(tmp_path, monkeypatch):
|
|
session_dir = tmp_path / "sessions"
|
|
session_dir.mkdir()
|
|
monkeypatch.setattr(models, "SESSION_DIR", session_dir)
|
|
monkeypatch.setattr(models, "SESSION_INDEX_FILE", session_dir / "_index.json")
|
|
monkeypatch.setattr(streaming, "SESSION_DIR", session_dir)
|
|
maps = {name: {} for name in (
|
|
"SESSIONS", "STREAMS", "CANCEL_FLAGS", "AGENT_INSTANCES",
|
|
"STREAM_PARTIAL_TEXT", "STREAM_REASONING_TEXT", "STREAM_LIVE_TOOL_CALLS",
|
|
"STREAM_SESSION_OWNERS", "ACTIVE_RUNS", "SESSION_AGENT_LOCKS",
|
|
)}
|
|
maps["SESSION_AGENT_CACHE"] = OrderedDict()
|
|
maps["SESSIONS"] = OrderedDict()
|
|
for name, value in maps.items():
|
|
for module in (config, models, streaming):
|
|
if hasattr(module, name):
|
|
monkeypatch.setattr(module, name, value)
|
|
lock = ObservedLock()
|
|
monkeypatch.setattr(config, "STREAMS_LOCK", lock)
|
|
monkeypatch.setattr(streaming, "STREAMS_LOCK", lock)
|
|
session = Session(session_id="original", title="Control boundary",
|
|
workspace=str(tmp_path), model="test-model",
|
|
messages=[], context_messages=[])
|
|
session.active_stream_id = "run"
|
|
session.pending_user_message = "Do the task."
|
|
session.pending_started_at = 1.0
|
|
session.save()
|
|
maps["SESSIONS"]["original"] = session
|
|
events = queue.Queue()
|
|
maps["STREAMS"]["run"] = events
|
|
maps["STREAM_SESSION_OWNERS"]["run"] = "original"
|
|
scene = SimpleNamespace(session=session, events=events, lock=lock, agent=None,
|
|
on_init=lambda: None, on_run=lambda: None,
|
|
on_drain=lambda: None, calls=[], drained=[], result=None)
|
|
|
|
class FakeAgent:
|
|
def __init__(self, **kwargs):
|
|
self.session_id = kwargs.get("session_id")
|
|
self.context_compressor = None
|
|
self.session_prompt_tokens = self.session_completion_tokens = 0
|
|
self.session_cache_read_tokens = self.session_cache_write_tokens = 0
|
|
self.session_estimated_cost_usd = None
|
|
self.reasoning_config = self.ephemeral_system_prompt = self._last_error = None
|
|
self.pending = []
|
|
self.pending_lock = threading.Lock()
|
|
scene.agent = self
|
|
scene.on_init()
|
|
|
|
def run_conversation(self, **kwargs):
|
|
scene.calls.append("run")
|
|
scene.on_run()
|
|
if scene.result is not None:
|
|
return scene.result
|
|
return {"messages": [
|
|
{"role": "user", "content": "Do the task."},
|
|
{"role": "assistant", "content": "Finished."},
|
|
]}
|
|
|
|
def interrupt(self, reason):
|
|
assert lock.owner != threading.get_ident(), "interrupt under stream lock"
|
|
scene.calls.append("interrupt")
|
|
|
|
def steer(self, text):
|
|
with self.pending_lock:
|
|
self.pending.append(text)
|
|
scene.calls.append("steer")
|
|
return True
|
|
|
|
def _drain_pending_steer(self):
|
|
with self.pending_lock:
|
|
text = "\n".join(self.pending)
|
|
self.pending.clear()
|
|
scene.drained.append(text)
|
|
scene.on_drain()
|
|
return text or None
|
|
|
|
monkeypatch.setattr(streaming, "get_session", lambda sid: session)
|
|
monkeypatch.setattr(streaming, "_get_ai_agent", lambda: FakeAgent)
|
|
monkeypatch.setattr(streaming, "resolve_model_provider", lambda *a, **kw: ("test-model", "openai", None))
|
|
monkeypatch.setattr(config, "get_config", lambda *a, **kw: {})
|
|
monkeypatch.setattr(config, "_resolve_cli_toolsets", lambda *a, **kw: [])
|
|
state = types.ModuleType("hermes_state")
|
|
state.SessionDB = lambda *a, **kw: object()
|
|
monkeypatch.setitem(sys.modules, "hermes_state", state)
|
|
scene.run = lambda: streaming._run_agent_streaming(
|
|
"original", "Do the task.", "test-model", str(tmp_path), "run")
|
|
return scene
|
|
|
|
|
|
def test_stop_between_stream_lookup_and_run_registration(worker_scene, monkeypatch):
|
|
scene = worker_scene
|
|
looked_up, release = threading.Event(), threading.Event()
|
|
peek, register = streaming.peek_stream, streaming.register_active_run
|
|
registrations = []
|
|
|
|
def paused_peek(stream_id):
|
|
q = peek(stream_id)
|
|
looked_up.set()
|
|
assert release.wait(5), "stream lookup was not released"
|
|
return q
|
|
|
|
def observed_register(*args, **kwargs):
|
|
registrations.append(kwargs)
|
|
return register(*args, **kwargs)
|
|
|
|
monkeypatch.setattr(streaming, "peek_stream", paused_peek)
|
|
monkeypatch.setattr(streaming, "register_active_run", observed_register)
|
|
with ThreadPoolExecutor(max_workers=1) as pool:
|
|
future = pool.submit(scene.run)
|
|
try:
|
|
assert looked_up.wait(5)
|
|
assert streaming.cancel_stream("run")
|
|
finally:
|
|
release.set()
|
|
future.result(timeout=10)
|
|
assert registrations == [], "cancelled worker republished active-run state"
|
|
assert scene.agent is None
|
|
assert "run" not in config.ACTIVE_RUNS
|
|
assert "run" not in config.CANCEL_FLAGS
|
|
|
|
|
|
def test_initial_run_registration_is_fenced_with_cancel_flag(worker_scene, monkeypatch):
|
|
scene = worker_scene
|
|
register = streaming.register_active_run
|
|
registrations = []
|
|
|
|
def observed_register(stream_id, **metadata):
|
|
assert scene.lock.owner == threading.get_ident()
|
|
assert stream_id in config.STREAMS
|
|
flag = config.CANCEL_FLAGS[stream_id]
|
|
assert not flag.is_set()
|
|
registrations.append(flag)
|
|
return register(stream_id, **metadata)
|
|
|
|
monkeypatch.setattr(streaming, "register_active_run", observed_register)
|
|
# Stop just after registration, before journal setup used to recreate flags.
|
|
def cancel_at_journal(*args, **kwargs):
|
|
assert scene.lock.owner != threading.get_ident()
|
|
assert streaming.cancel_stream("run")
|
|
return None
|
|
|
|
monkeypatch.setattr(streaming, "RunJournalWriter", cancel_at_journal)
|
|
scene.run()
|
|
assert len(registrations) == 1
|
|
assert registrations[0].is_set()
|
|
assert scene.agent is None, "cancel event was replaced before preflight"
|
|
assert "run" not in config.CANCEL_FLAGS
|
|
assert "run" not in config.ACTIVE_RUNS
|
|
|
|
|
|
@pytest.mark.parametrize("cancellation", ["stop", "event-only", "detached-only"])
|
|
def test_stop_during_agent_creation_prevents_provider_run(worker_scene, monkeypatch, cancellation):
|
|
scene = worker_scene
|
|
reached, release = threading.Event(), threading.Event()
|
|
|
|
def during_init():
|
|
reached.set()
|
|
assert release.wait(5), "initialization barrier timed out"
|
|
|
|
scene.on_init = during_init
|
|
finalize = streaming._finalize_cancelled_turn
|
|
finalized = []
|
|
|
|
def finalize_outside_lock(*args, **kwargs):
|
|
assert scene.lock.owner != threading.get_ident(), "finalization under stream lock"
|
|
finalized.append(True)
|
|
return finalize(*args, **kwargs)
|
|
|
|
monkeypatch.setattr(streaming, "_finalize_cancelled_turn", finalize_outside_lock)
|
|
with ThreadPoolExecutor(max_workers=1) as pool:
|
|
worker = pool.submit(scene.run)
|
|
try:
|
|
assert reached.wait(5), "worker never reached agent creation"
|
|
event = config.CANCEL_FLAGS["run"]
|
|
if cancellation == "stop":
|
|
assert streaming.cancel_stream("run") is True
|
|
assert event.is_set()
|
|
assert "run" not in config.CANCEL_FLAGS
|
|
assert "run" not in config.STREAMS
|
|
elif cancellation == "event-only":
|
|
with config.STREAMS_LOCK:
|
|
event.set()
|
|
config.CANCEL_FLAGS.pop("run")
|
|
else:
|
|
with config.STREAMS_LOCK:
|
|
config.STREAMS.pop("run")
|
|
assert not event.is_set()
|
|
finally:
|
|
release.set()
|
|
worker.result(timeout=10)
|
|
assert "run" not in scene.calls, "accepted Stop must prevent the provider/tool run"
|
|
cached = config.SESSION_AGENT_CACHE.get("original")
|
|
assert not cached or cached[0] is not scene.agent, "cancelled initial candidate remained reusable"
|
|
assert finalized, "cancelled turn was not finalized"
|
|
assert "run" not in config.AGENT_INSTANCES
|
|
assert "run" not in config.ACTIVE_RUNS
|
|
|
|
|
|
@pytest.mark.parametrize("registered,rotated", [(True, True), (True, False), (False, False)])
|
|
@pytest.mark.parametrize("first", ["steer", "drain"])
|
|
def test_final_drain_fences_steer(worker_scene, monkeypatch, registered, rotated, first):
|
|
scene = worker_scene
|
|
reached, release = threading.Event(), threading.Event()
|
|
request_thread = threading.get_ident()
|
|
# An already-resolved request can retain the pre-settlement session pointer.
|
|
# Registries must reject late delivery even if this projection is stale.
|
|
request_session = SimpleNamespace(active_stream_id="run")
|
|
monkeypatch.setattr(streaming, "get_session", lambda sid:
|
|
request_session if threading.get_ident() == request_thread else scene.session)
|
|
update = streaming.update_active_run
|
|
finalizing_edges = []
|
|
|
|
def observe_update(*args, **kwargs):
|
|
if kwargs.get("phase") == "finalizing":
|
|
finalizing_edges.append(scene.lock.owner == threading.get_ident())
|
|
return update(*args, **kwargs)
|
|
|
|
monkeypatch.setattr(streaming, "update_active_run", observe_update)
|
|
|
|
def pause():
|
|
reached.set()
|
|
assert release.wait(5), "drain/steer barrier timed out"
|
|
|
|
def running():
|
|
if rotated:
|
|
scene.agent.session_id = "compressed-child"
|
|
if not registered:
|
|
with config.STREAMS_LOCK:
|
|
config.AGENT_INSTANCES.pop("run")
|
|
if first == "steer":
|
|
pause()
|
|
|
|
scene.on_run = running
|
|
if first == "drain":
|
|
scene.on_drain = pause
|
|
with ThreadPoolExecutor(max_workers=1) as pool:
|
|
worker = pool.submit(scene.run)
|
|
try:
|
|
assert reached.wait(5), "worker never reached the requested boundary"
|
|
handler = _make_handler()
|
|
streaming._handle_chat_steer(handler, {"session_id": "original", "text": "guidance"})
|
|
result = _captured_response(handler)
|
|
if first == "steer":
|
|
assert result == {"accepted": True, "fallback": None, "stream_id": "run"}
|
|
else:
|
|
assert result == {"accepted": False, "fallback": "not_running", "stream_id": "run"}
|
|
assert scene.agent.pending == []
|
|
finally:
|
|
release.set()
|
|
worker.result(timeout=10)
|
|
emitted = list(scene.events.queue)
|
|
assert not [payload for event, payload in emitted if event == "apperror"]
|
|
leftovers = [payload["text"] for event, payload in emitted if event == "pending_steer_leftover"]
|
|
assert leftovers == (["guidance"] if first == "steer" else [])
|
|
assert scene.drained == (["guidance"] if first == "steer" else [""])
|
|
assert scene.agent.pending == []
|
|
assert finalizing_edges == [True], "final drain admission must close under the stream lock"
|
|
assert "run" not in config.ACTIVE_RUNS
|
|
|
|
|
|
def test_successfully_registered_new_agent_remains_cached(worker_scene):
|
|
scene = worker_scene
|
|
scene.run()
|
|
assert scene.calls.count("run") == 1
|
|
assert config.SESSION_AGENT_CACHE["original"][0] is scene.agent
|