mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-12 23:12:23 -06:00
Patrick Buckley
5378b33641
* feat: output guard — evaluate tool results before they enter context Add turnstone/core/output_guard.py — a time-budgeted heuristic that evaluates tool execution results after execution but before they enter the conversation context window. Priority-ordered detection (5s budget, highest priority first): 1. Prompt injection: override phrases, role injection, instruction override markers, meta-injection patterns 2. Credential leakage: API keys (OpenAI/GitHub/AWS/Google), PEM private key blocks, connection strings, .env secret format 3. Encoded payloads: script data URIs, hex shellcode sequences 4. Adversarial URLs: cloud metadata endpoints, credential query params 5. System info disclosure: private IPs, sensitive file paths Annotates and optionally redacts (credentials → [REDACTED:<type>]). Does NOT gate — surfaces warnings via on_output_warning callback. Integration: - Wired into session.py tool result loop via _evaluate_output() - JudgeConfig gains output_guard + redact_secrets fields (both default true) - SessionUI protocol gains on_output_warning callback - 25 compiled regex patterns, pure function, no I/O 29 tests covering all detection categories, benign output false positive checks, credential redaction, and time budget behavior. * fix: address PR #109 review — protocol, config, and guard fixes Copilot review feedback: - Replace _CLEAN singleton with _clean() factory to prevent mutable shared state (OutputAssessment has list fields) - Remove redundant second _CREDENTIAL_PATTERNS loop in _check_credentials - Evaluate text parts of list outputs (images) not just string outputs - Wire output_guard + redact_secrets through ConfigStore settings registry and _build_judge_config() so operators can configure via admin Settings tab - Remove --no-output-guard CLI flag claim from docs (use Settings tab) Typecheck fix: - Add on_output_warning to all SessionUI implementations: NullUI (eval, 5 test files), WebUI (server — emits SSE event), TerminalUI (CLI — ANSI colored warning), RecordingUI, FakeUI
972 lines
33 KiB
Python
972 lines
33 KiB
Python
"""Tests for turnstone.core.workstream — WorkstreamManager, state management, and UI adapters."""
|
|
|
|
import threading
|
|
import time
|
|
|
|
import pytest
|
|
|
|
from turnstone.core.workstream import Workstream, WorkstreamManager, WorkstreamState
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class FakeSession:
|
|
"""Minimal stand-in for ChatSession in workstream tests."""
|
|
|
|
def __init__(self):
|
|
self.model = "test-model"
|
|
self.messages = []
|
|
|
|
|
|
def _fake_factory(ui, model_alias=None, ws_id=None, **kwargs):
|
|
return FakeSession()
|
|
|
|
|
|
class FakeUI:
|
|
"""Minimal SessionUI that tracks state changes."""
|
|
|
|
def __init__(self, ws_id=""):
|
|
self.ws_id = ws_id
|
|
self.state_changes = []
|
|
self.auto_approve = False
|
|
|
|
def on_state_change(self, state):
|
|
self.state_changes.append(state)
|
|
|
|
# Stubs for the rest of the protocol
|
|
def on_thinking_start(self):
|
|
pass
|
|
|
|
def on_thinking_stop(self):
|
|
pass
|
|
|
|
def on_reasoning_token(self, text):
|
|
pass
|
|
|
|
def on_content_token(self, text):
|
|
pass
|
|
|
|
def on_stream_end(self):
|
|
pass
|
|
|
|
def approve_tools(self, items):
|
|
return True, None
|
|
|
|
def on_tool_result(self, call_id, name, output):
|
|
pass
|
|
|
|
def on_tool_output_chunk(self, call_id, chunk):
|
|
pass
|
|
|
|
def on_status(self, usage, context_window, effort):
|
|
pass
|
|
|
|
def on_plan_review(self, content):
|
|
return ""
|
|
|
|
def on_info(self, message):
|
|
pass
|
|
|
|
def on_error(self, message):
|
|
pass
|
|
|
|
def on_rename(self, name):
|
|
pass
|
|
|
|
def on_output_warning(self, call_id, assessment):
|
|
pass
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# WorkstreamState enum
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestWorkstreamState:
|
|
def test_values(self):
|
|
assert WorkstreamState.IDLE.value == "idle"
|
|
assert WorkstreamState.THINKING.value == "thinking"
|
|
assert WorkstreamState.RUNNING.value == "running"
|
|
assert WorkstreamState.ATTENTION.value == "attention"
|
|
assert WorkstreamState.ERROR.value == "error"
|
|
|
|
def test_from_string(self):
|
|
assert WorkstreamState("idle") == WorkstreamState.IDLE
|
|
assert WorkstreamState("attention") == WorkstreamState.ATTENTION
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Workstream dataclass
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestWorkstream:
|
|
def test_default_name(self):
|
|
ws = Workstream()
|
|
assert ws.name.startswith("ws-")
|
|
assert len(ws.name) == 7 # "ws-" + 4 hex chars
|
|
|
|
def test_custom_name(self):
|
|
ws = Workstream(name="my-stream")
|
|
assert ws.name == "my-stream"
|
|
|
|
def test_default_state(self):
|
|
ws = Workstream()
|
|
assert ws.state == WorkstreamState.IDLE
|
|
|
|
def test_id_uniqueness(self):
|
|
ws1 = Workstream()
|
|
ws2 = Workstream()
|
|
assert ws1.id != ws2.id
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# WorkstreamManager — creation and lookup
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestManagerCreation:
|
|
def test_create_first_sets_active(self):
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
ws = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
assert mgr.active_id == ws.id
|
|
assert mgr.get_active() is ws
|
|
|
|
def test_create_second_does_not_change_active(self):
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
ws1 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
_ws2 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
assert mgr.active_id == ws1.id
|
|
|
|
def test_create_assigns_session(self):
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
ws = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
assert isinstance(ws.session, FakeSession)
|
|
|
|
def test_create_assigns_ui(self):
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
ws = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
assert isinstance(ws.ui, FakeUI)
|
|
assert ws.ui.ws_id == ws.id
|
|
|
|
def test_create_custom_name(self):
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
ws = mgr.create(name="research", ui_factory=lambda wid: FakeUI(wid))
|
|
assert ws.name == "research"
|
|
|
|
def test_create_default_name(self):
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
ws = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
assert ws.name.startswith("ws-")
|
|
|
|
def test_create_max_workstreams_all_active(self):
|
|
mgr = WorkstreamManager(_fake_factory, max_workstreams=3)
|
|
ws1 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
ws2 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
ws3 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
# Mark all as non-idle so eviction cannot help
|
|
mgr.set_state(ws1.id, WorkstreamState.THINKING)
|
|
mgr.set_state(ws2.id, WorkstreamState.RUNNING)
|
|
mgr.set_state(ws3.id, WorkstreamState.ATTENTION)
|
|
with pytest.raises(RuntimeError, match="All 3 workstreams are active"):
|
|
mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
|
|
|
|
class TestManagerLookup:
|
|
def test_get_existing(self):
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
ws = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
assert mgr.get(ws.id) is ws
|
|
|
|
def test_get_nonexistent(self):
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
assert mgr.get("no-such-id") is None
|
|
|
|
def test_list_all_creation_order(self):
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
_ws1 = mgr.create(name="a", ui_factory=lambda wid: FakeUI(wid))
|
|
_ws2 = mgr.create(name="b", ui_factory=lambda wid: FakeUI(wid))
|
|
_ws3 = mgr.create(name="c", ui_factory=lambda wid: FakeUI(wid))
|
|
result = mgr.list_all()
|
|
assert [w.name for w in result] == ["a", "b", "c"]
|
|
|
|
def test_index_of(self):
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
ws1 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
ws2 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
assert mgr.index_of(ws1.id) == 1
|
|
assert mgr.index_of(ws2.id) == 2
|
|
assert mgr.index_of("nonexistent") == 0
|
|
|
|
def test_count(self):
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
assert mgr.count == 0
|
|
mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
assert mgr.count == 1
|
|
mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
assert mgr.count == 2
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# WorkstreamManager — switching
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestManagerSwitching:
|
|
def test_switch_by_id(self):
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
ws1 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
ws2 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
assert mgr.active_id == ws1.id
|
|
|
|
result = mgr.switch(ws2.id)
|
|
assert result is ws2
|
|
assert mgr.active_id == ws2.id
|
|
|
|
def test_switch_nonexistent_returns_none(self):
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
assert mgr.switch("bad-id") is None
|
|
|
|
def test_switch_by_index(self):
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
_ws1 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
ws2 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
|
|
result = mgr.switch_by_index(2)
|
|
assert result is ws2
|
|
assert mgr.active_id == ws2.id
|
|
|
|
def test_switch_by_index_out_of_range(self):
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
assert mgr.switch_by_index(0) is None
|
|
assert mgr.switch_by_index(5) is None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# WorkstreamManager — closing
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestManagerClose:
|
|
def test_close_removes_workstream(self):
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
_ws1 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
ws2 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
|
|
assert mgr.close(ws2.id) is True
|
|
assert mgr.count == 1
|
|
assert mgr.get(ws2.id) is None
|
|
|
|
def test_close_last_returns_false(self):
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
ws = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
assert mgr.close(ws.id) is False
|
|
assert mgr.count == 1
|
|
|
|
def test_close_nonexistent_returns_false(self):
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
assert mgr.close("nonexistent") is False
|
|
|
|
def test_close_active_switches_to_first(self):
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
ws1 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
ws2 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
mgr.switch(ws2.id)
|
|
|
|
mgr.close(ws2.id)
|
|
assert mgr.active_id == ws1.id
|
|
|
|
def test_close_updates_order(self):
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
_ws1 = mgr.create(name="a", ui_factory=lambda wid: FakeUI(wid))
|
|
ws2 = mgr.create(name="b", ui_factory=lambda wid: FakeUI(wid))
|
|
_ws3 = mgr.create(name="c", ui_factory=lambda wid: FakeUI(wid))
|
|
|
|
mgr.close(ws2.id)
|
|
names = [w.name for w in mgr.list_all()]
|
|
assert names == ["a", "c"]
|
|
|
|
def test_close_unblocks_approval_event(self):
|
|
"""Closing a workstream whose UI has a pending approval should unblock it."""
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
_ws1 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
|
|
# Create a workstream with a WebUI-like approval mechanism
|
|
from turnstone.server import WebUI
|
|
|
|
ws2 = mgr.create(ui_factory=lambda wid: WebUI(ws_id=wid))
|
|
ws2.ui._approval_event.clear() # simulate pending approval
|
|
|
|
mgr.close(ws2.id)
|
|
# The approval event should be set (unblocked)
|
|
assert ws2.ui._approval_event.is_set()
|
|
|
|
def test_close_unblocks_plan_event(self):
|
|
"""Closing a workstream with pending plan review should unblock it."""
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
_ws1 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
|
|
from turnstone.server import WebUI
|
|
|
|
ws2 = mgr.create(ui_factory=lambda wid: WebUI(ws_id=wid))
|
|
ws2.ui._plan_event.clear()
|
|
|
|
mgr.close(ws2.id)
|
|
assert ws2.ui._plan_event.is_set()
|
|
assert ws2.ui._plan_result == "reject"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# WorkstreamManager — auto-eviction
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestManagerEviction:
|
|
def test_evict_oldest_idle_on_create(self):
|
|
"""At capacity with idle workstreams, create() succeeds by evicting the oldest idle."""
|
|
mgr = WorkstreamManager(_fake_factory, max_workstreams=3)
|
|
ws1 = mgr.create(name="oldest", ui_factory=lambda wid: FakeUI(wid))
|
|
ws2 = mgr.create(name="middle", ui_factory=lambda wid: FakeUI(wid))
|
|
_ws3 = mgr.create(name="newest", ui_factory=lambda wid: FakeUI(wid))
|
|
# All three are IDLE. Mark ws2 as RUNNING so it won't be evicted.
|
|
mgr.set_state(ws2.id, WorkstreamState.RUNNING)
|
|
# ws1 is oldest idle, ws3 is newer idle. Creating should evict ws1.
|
|
ws4 = mgr.create(name="four", ui_factory=lambda wid: FakeUI(wid))
|
|
assert mgr.count == 3
|
|
assert mgr.get(ws1.id) is None, "oldest idle should have been evicted"
|
|
assert mgr.get(ws4.id) is ws4
|
|
# Creation order should reflect the eviction
|
|
names = [w.name for w in mgr.list_all()]
|
|
assert "oldest" not in names
|
|
assert names == ["middle", "newest", "four"]
|
|
|
|
def test_create_fails_when_all_active(self):
|
|
"""At capacity with ALL non-idle workstreams, create() raises RuntimeError."""
|
|
mgr = WorkstreamManager(_fake_factory, max_workstreams=2)
|
|
ws1 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
ws2 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
mgr.set_state(ws1.id, WorkstreamState.THINKING)
|
|
mgr.set_state(ws2.id, WorkstreamState.RUNNING)
|
|
with pytest.raises(RuntimeError, match="All 2 workstreams are active"):
|
|
mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
|
|
def test_configurable_max(self):
|
|
"""Constructor accepts max_workstreams param and respects it."""
|
|
mgr = WorkstreamManager(_fake_factory, max_workstreams=2)
|
|
ws1 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
ws2 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
mgr.set_state(ws1.id, WorkstreamState.RUNNING)
|
|
mgr.set_state(ws2.id, WorkstreamState.RUNNING)
|
|
with pytest.raises(RuntimeError):
|
|
mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
assert mgr.count == 2
|
|
|
|
def test_eviction_counter(self):
|
|
"""eviction_count increments on each auto-eviction."""
|
|
mgr = WorkstreamManager(_fake_factory, max_workstreams=2)
|
|
assert mgr.eviction_count == 0
|
|
mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
# Both IDLE — create should evict the oldest
|
|
mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
assert mgr.eviction_count == 1
|
|
# Again — evict another idle one
|
|
mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
assert mgr.eviction_count == 2
|
|
assert mgr.count == 2
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# WorkstreamManager — state management
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestManagerState:
|
|
def test_set_state(self):
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
ws = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
assert ws.state == WorkstreamState.IDLE
|
|
|
|
mgr.set_state(ws.id, WorkstreamState.THINKING)
|
|
assert ws.state == WorkstreamState.THINKING
|
|
|
|
def test_set_state_with_error(self):
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
ws = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
|
|
mgr.set_state(ws.id, WorkstreamState.ERROR, error_msg="API timeout")
|
|
assert ws.state == WorkstreamState.ERROR
|
|
assert ws.error_message == "API timeout"
|
|
|
|
def test_set_state_nonexistent_is_noop(self):
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
mgr.set_state("no-such-id", WorkstreamState.THINKING) # should not raise
|
|
|
|
def test_on_state_change_callback(self):
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
ws = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
|
|
changes = []
|
|
mgr._on_state_change = lambda wid, state: changes.append((wid, state))
|
|
|
|
mgr.set_state(ws.id, WorkstreamState.RUNNING)
|
|
assert changes == [(ws.id, WorkstreamState.RUNNING)]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# WorkstreamManager — thread safety
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestManagerThreadSafety:
|
|
def test_concurrent_create_respects_max(self):
|
|
"""Multiple threads creating workstreams should not exceed max."""
|
|
mgr = WorkstreamManager(_fake_factory, max_workstreams=5)
|
|
errors = []
|
|
created = []
|
|
|
|
def do_create():
|
|
try:
|
|
ws = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
# Mark as non-idle immediately so auto-eviction cannot reclaim it
|
|
mgr.set_state(ws.id, WorkstreamState.RUNNING)
|
|
created.append(ws.id)
|
|
except RuntimeError:
|
|
errors.append(True)
|
|
|
|
threads = [threading.Thread(target=do_create) for _ in range(10)]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
# All threads resolved (created or rejected)
|
|
assert len(created) + len(errors) == 10
|
|
# Never exceeded capacity
|
|
assert mgr.count <= 5
|
|
|
|
def test_concurrent_switch(self):
|
|
"""Concurrent switches should not corrupt state."""
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
ids = []
|
|
for _ in range(5):
|
|
ws = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
ids.append(ws.id)
|
|
|
|
def do_switch(wid):
|
|
for _ in range(20):
|
|
mgr.switch(wid)
|
|
|
|
threads = [threading.Thread(target=do_switch, args=(wid,)) for wid in ids]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
# active_id should be one of the valid ids
|
|
assert mgr.active_id in ids
|
|
|
|
def test_concurrent_close_and_list(self):
|
|
"""close() and list_all() running concurrently should not crash."""
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
# Keep one alive to prevent closing the last
|
|
anchor = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
targets = []
|
|
for _ in range(5):
|
|
ws = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
targets.append(ws.id)
|
|
|
|
def do_close():
|
|
for wid in targets:
|
|
mgr.close(wid)
|
|
|
|
def do_list():
|
|
for _ in range(50):
|
|
mgr.list_all()
|
|
|
|
t1 = threading.Thread(target=do_close)
|
|
t2 = threading.Thread(target=do_list)
|
|
t1.start()
|
|
t2.start()
|
|
t1.join()
|
|
t2.join()
|
|
|
|
assert mgr.count == 1
|
|
assert mgr.get(anchor.id) is not None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# WorkstreamTerminalUI
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestWorkstreamTerminalUI:
|
|
def test_foreground_detection(self):
|
|
from turnstone.cli import WorkstreamTerminalUI
|
|
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
ws = mgr.create(ui_factory=lambda wid: WorkstreamTerminalUI(wid, mgr))
|
|
assert ws.ui.is_foreground is True
|
|
|
|
ws2 = mgr.create(ui_factory=lambda wid: WorkstreamTerminalUI(wid, mgr))
|
|
# ws1 is still active, so ws2 is not foreground
|
|
assert ws2.ui.is_foreground is False
|
|
|
|
def test_state_change_updates_manager(self):
|
|
from turnstone.cli import WorkstreamTerminalUI
|
|
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
ws = mgr.create(ui_factory=lambda wid: WorkstreamTerminalUI(wid, mgr))
|
|
|
|
ws.ui.on_state_change("thinking")
|
|
assert ws.state == WorkstreamState.THINKING
|
|
|
|
ws.ui.on_state_change("idle")
|
|
assert ws.state == WorkstreamState.IDLE
|
|
|
|
def test_invalid_state_change_ignored(self):
|
|
from turnstone.cli import WorkstreamTerminalUI
|
|
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
ws = mgr.create(ui_factory=lambda wid: WorkstreamTerminalUI(wid, mgr))
|
|
ws.ui.on_state_change("not_a_real_state") # should not raise
|
|
assert ws.state == WorkstreamState.IDLE # unchanged
|
|
|
|
def _make_background_ws(self):
|
|
"""Create a manager with two workstreams; switch to the second so the first is background."""
|
|
from turnstone.cli import WorkstreamTerminalUI
|
|
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
bg = mgr.create(ui_factory=lambda wid: WorkstreamTerminalUI(wid, mgr))
|
|
fg = mgr.create(ui_factory=lambda wid: WorkstreamTerminalUI(wid, mgr))
|
|
mgr.switch(fg.id)
|
|
bg.ui.set_foreground(False)
|
|
fg.ui.set_foreground(True)
|
|
return mgr, bg, fg
|
|
|
|
def test_background_buffers_content(self):
|
|
mgr, bg, fg = self._make_background_ws()
|
|
assert bg.ui.is_foreground is False
|
|
|
|
bg.ui.on_content_token("hello ")
|
|
bg.ui.on_content_token("world")
|
|
bg.ui.on_stream_end()
|
|
|
|
assert len(bg.ui._output_buffer) == 3
|
|
assert bg.ui._output_buffer[0] == ("content", "hello ")
|
|
assert bg.ui._output_buffer[1] == ("content", "world")
|
|
assert bg.ui._output_buffer[2] == ("stream_end", "")
|
|
|
|
def test_flush_buffer_clears(self):
|
|
mgr, bg, fg = self._make_background_ws()
|
|
|
|
bg.ui.on_content_token("test")
|
|
bg.ui.on_stream_end()
|
|
assert len(bg.ui._output_buffer) == 2
|
|
|
|
mgr.switch(bg.id)
|
|
bg.ui.set_foreground(True)
|
|
bg.ui.flush_buffer()
|
|
assert len(bg.ui._output_buffer) == 0
|
|
|
|
def test_background_buffers_info_and_error(self):
|
|
mgr, bg, fg = self._make_background_ws()
|
|
|
|
bg.ui.on_info("info msg")
|
|
bg.ui.on_error("error msg")
|
|
|
|
assert ("info", "info msg") in bg.ui._output_buffer
|
|
assert ("error", "error msg") in bg.ui._output_buffer
|
|
|
|
def test_fg_event_blocks_approval_in_background(self):
|
|
"""approve_tools should block until foregrounded."""
|
|
mgr, bg, fg = self._make_background_ws()
|
|
bg.ui.auto_approve = True # so we don't need actual input()
|
|
|
|
result = [None]
|
|
|
|
def call_approve():
|
|
result[0] = bg.ui.approve_tools(
|
|
[{"needs_approval": True, "header": "test", "func_name": "bash"}]
|
|
)
|
|
|
|
t = threading.Thread(target=call_approve)
|
|
t.start()
|
|
time.sleep(0.1)
|
|
assert t.is_alive(), "approve_tools should be blocking"
|
|
|
|
# Bring to foreground — should unblock
|
|
mgr.switch(bg.id)
|
|
bg.ui.set_foreground(True)
|
|
t.join(timeout=2)
|
|
assert not t.is_alive()
|
|
assert result[0] == (True, None) # auto-approved
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# WebUI workstream support
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestWebUI:
|
|
def test_ws_id_assigned(self):
|
|
from turnstone.server import WebUI
|
|
|
|
ui = WebUI(ws_id="test-123")
|
|
assert ui.ws_id == "test-123"
|
|
|
|
def test_on_state_change_broadcasts(self):
|
|
"""on_state_change should put an event on the global queue."""
|
|
import queue
|
|
|
|
from turnstone.server import WebUI
|
|
|
|
gq = queue.Queue()
|
|
old = WebUI._global_queue
|
|
WebUI._global_queue = gq
|
|
try:
|
|
ui = WebUI(ws_id="abc")
|
|
ui.on_state_change("thinking")
|
|
|
|
event = gq.get_nowait()
|
|
assert event["type"] == "ws_state"
|
|
assert event["ws_id"] == "abc"
|
|
assert event["state"] == "thinking"
|
|
finally:
|
|
WebUI._global_queue = old
|
|
|
|
def test_on_state_change_no_global_queue(self):
|
|
"""on_state_change should not crash if no global queue is set."""
|
|
from turnstone.server import WebUI
|
|
|
|
old = WebUI._global_queue
|
|
WebUI._global_queue = None
|
|
try:
|
|
ui = WebUI(ws_id="xyz")
|
|
ui.on_state_change("running") # should not raise
|
|
finally:
|
|
WebUI._global_queue = old
|
|
|
|
def test_resolve_approval(self):
|
|
from turnstone.server import WebUI
|
|
|
|
ui = WebUI(ws_id="test")
|
|
ui._approval_event.clear()
|
|
|
|
# Resolve in a thread
|
|
def resolve():
|
|
time.sleep(0.05)
|
|
ui.resolve_approval(True, "looks good")
|
|
|
|
t = threading.Thread(target=resolve)
|
|
t.start()
|
|
|
|
ui._approval_event.wait(timeout=2)
|
|
assert ui._approval_result == (True, "looks good")
|
|
t.join()
|
|
|
|
def test_resolve_approval_emits_event(self):
|
|
"""resolve_approval should enqueue an approval_resolved SSE event."""
|
|
from turnstone.server import WebUI
|
|
|
|
ui = WebUI(ws_id="test-emit")
|
|
listener = ui._register_listener()
|
|
|
|
# Drain any init events
|
|
while not listener.empty():
|
|
listener.get_nowait()
|
|
|
|
ui.resolve_approval(False, "Approval timed out")
|
|
|
|
# Collect events from the listener
|
|
events = []
|
|
while not listener.empty():
|
|
events.append(listener.get_nowait())
|
|
|
|
ui._unregister_listener(listener)
|
|
|
|
resolved = [e for e in events if e.get("type") == "approval_resolved"]
|
|
assert len(resolved) == 1
|
|
assert resolved[0]["approved"] is False
|
|
assert resolved[0]["feedback"] == "Approval timed out"
|
|
|
|
def test_resolve_plan(self):
|
|
from turnstone.server import WebUI
|
|
|
|
ui = WebUI(ws_id="test")
|
|
ui._plan_event.clear()
|
|
|
|
def resolve():
|
|
time.sleep(0.05)
|
|
ui.resolve_plan("approved")
|
|
|
|
t = threading.Thread(target=resolve)
|
|
t.start()
|
|
|
|
ui._plan_event.wait(timeout=2)
|
|
assert ui._plan_result == "approved"
|
|
t.join()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# WebUI SSE fan-out
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestWebUIFanOut:
|
|
"""Verify per-client SSE fan-out on WebUI._enqueue / _register_listener."""
|
|
|
|
def test_enqueue_no_listeners(self):
|
|
"""Events silently dropped when no listeners are registered."""
|
|
from turnstone.server import WebUI
|
|
|
|
ui = WebUI(ws_id="test")
|
|
ui._enqueue({"type": "content", "text": "hello"}) # should not raise
|
|
|
|
def test_enqueue_single_listener(self):
|
|
"""Single listener receives the event."""
|
|
from turnstone.server import WebUI
|
|
|
|
ui = WebUI(ws_id="test")
|
|
q = ui._register_listener()
|
|
ui._enqueue({"type": "content", "text": "hello"})
|
|
assert q.get_nowait() == {"type": "content", "text": "hello"}
|
|
|
|
def test_enqueue_multiple_listeners(self):
|
|
"""All registered listeners receive an identical copy."""
|
|
from turnstone.server import WebUI
|
|
|
|
ui = WebUI(ws_id="test")
|
|
q1 = ui._register_listener()
|
|
q2 = ui._register_listener()
|
|
q3 = ui._register_listener()
|
|
|
|
event = {"type": "content", "text": "world"}
|
|
ui._enqueue(event)
|
|
|
|
assert q1.get_nowait() == event
|
|
assert q2.get_nowait() == event
|
|
assert q3.get_nowait() == event
|
|
|
|
def test_unregister_stops_delivery(self):
|
|
"""After unregister, the queue receives no further events."""
|
|
import queue as queue_mod
|
|
|
|
from turnstone.server import WebUI
|
|
|
|
ui = WebUI(ws_id="test")
|
|
q = ui._register_listener()
|
|
ui._unregister_listener(q)
|
|
ui._enqueue({"type": "content", "text": "gone"})
|
|
|
|
with pytest.raises(queue_mod.Empty):
|
|
q.get_nowait()
|
|
|
|
def test_slow_consumer_does_not_block(self):
|
|
"""A full queue doesn't block the producer or starve other listeners."""
|
|
from turnstone.server import WebUI
|
|
|
|
ui = WebUI(ws_id="test")
|
|
slow = ui._register_listener()
|
|
fast = ui._register_listener()
|
|
|
|
# Fill only the slow consumer's queue directly to capacity
|
|
for i in range(500):
|
|
slow.put_nowait({"type": "content", "text": f"fill-{i}"})
|
|
|
|
assert slow.qsize() == 500
|
|
assert fast.qsize() == 0
|
|
|
|
# Enqueue via fan-out — slow drops (full), fast receives
|
|
event = {"type": "content", "text": "overflow"}
|
|
ui._enqueue(event)
|
|
assert slow.qsize() == 500 # still full, overflow dropped
|
|
assert fast.qsize() == 1
|
|
assert fast.get_nowait() == event
|
|
|
|
def test_unregister_idempotent(self):
|
|
"""Double unregister does not raise."""
|
|
|
|
from turnstone.server import WebUI
|
|
|
|
ui = WebUI(ws_id="test")
|
|
q = ui._register_listener()
|
|
ui._unregister_listener(q)
|
|
ui._unregister_listener(q) # should not raise
|
|
|
|
def test_concurrent_enqueue_and_register(self):
|
|
"""Concurrent register/unregister and enqueue should not crash."""
|
|
from turnstone.server import WebUI
|
|
|
|
ui = WebUI(ws_id="test")
|
|
stop = threading.Event()
|
|
|
|
def register_loop():
|
|
while not stop.is_set():
|
|
q = ui._register_listener()
|
|
ui._unregister_listener(q)
|
|
|
|
def enqueue_loop():
|
|
for i in range(500):
|
|
ui._enqueue({"type": "content", "text": f"tok-{i}"})
|
|
|
|
t1 = threading.Thread(target=register_loop)
|
|
t2 = threading.Thread(target=enqueue_loop)
|
|
t1.start()
|
|
t2.start()
|
|
t2.join()
|
|
stop.set()
|
|
t1.join()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Integration: WorkstreamManager + session state transitions
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestStateTransitions:
|
|
def test_full_lifecycle(self):
|
|
"""Verify the expected state transition sequence."""
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
ws = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
|
|
# Simulate the state transitions that ChatSession.send() would emit
|
|
mgr.set_state(ws.id, WorkstreamState.THINKING)
|
|
assert ws.state == WorkstreamState.THINKING
|
|
|
|
mgr.set_state(ws.id, WorkstreamState.RUNNING)
|
|
assert ws.state == WorkstreamState.RUNNING
|
|
|
|
mgr.set_state(ws.id, WorkstreamState.ATTENTION)
|
|
assert ws.state == WorkstreamState.ATTENTION
|
|
|
|
mgr.set_state(ws.id, WorkstreamState.RUNNING)
|
|
assert ws.state == WorkstreamState.RUNNING
|
|
|
|
mgr.set_state(ws.id, WorkstreamState.IDLE)
|
|
assert ws.state == WorkstreamState.IDLE
|
|
|
|
def test_error_recovery(self):
|
|
"""After an error, sending again should transition back to thinking."""
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
ws = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
|
|
|
mgr.set_state(ws.id, WorkstreamState.ERROR, "API failed")
|
|
assert ws.state == WorkstreamState.ERROR
|
|
assert ws.error_message == "API failed"
|
|
|
|
mgr.set_state(ws.id, WorkstreamState.THINKING)
|
|
assert ws.state == WorkstreamState.THINKING
|
|
assert ws.error_message == ""
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Design polish: thread-safe buffer, approval context, NO_COLOR
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestBufferThreadSafety:
|
|
"""Verify that _buffer() uses the lock and flush_buffer copies under lock."""
|
|
|
|
def test_concurrent_buffer_and_flush(self):
|
|
"""Simultaneous buffering and flushing should not lose or corrupt events."""
|
|
from turnstone.cli import WorkstreamTerminalUI
|
|
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
bg = mgr.create(ui_factory=lambda wid: WorkstreamTerminalUI(wid, mgr))
|
|
fg = mgr.create(ui_factory=lambda wid: WorkstreamTerminalUI(wid, mgr))
|
|
mgr.switch(fg.id)
|
|
bg.ui.set_foreground(False)
|
|
|
|
n_events = 200
|
|
done = threading.Event()
|
|
|
|
def do_buffer():
|
|
for i in range(n_events):
|
|
bg.ui._buffer("content", f"token-{i}")
|
|
done.set()
|
|
|
|
t = threading.Thread(target=do_buffer)
|
|
t.start()
|
|
done.wait()
|
|
|
|
# All events should be in the buffer
|
|
with bg.ui._print_lock:
|
|
count = len(bg.ui._output_buffer)
|
|
assert count == n_events
|
|
t.join()
|
|
|
|
|
|
class TestApprovalContextMessage:
|
|
"""Verify approval in background buffers a context message."""
|
|
|
|
def test_approval_buffers_tool_names(self):
|
|
from turnstone.cli import WorkstreamTerminalUI
|
|
|
|
mgr = WorkstreamManager(_fake_factory)
|
|
bg = mgr.create(ui_factory=lambda wid: WorkstreamTerminalUI(wid, mgr))
|
|
fg = mgr.create(ui_factory=lambda wid: WorkstreamTerminalUI(wid, mgr))
|
|
mgr.switch(fg.id)
|
|
bg.ui.set_foreground(False)
|
|
bg.ui.auto_approve = True
|
|
|
|
result = [None]
|
|
|
|
def call_approve():
|
|
result[0] = bg.ui.approve_tools(
|
|
[
|
|
{
|
|
"needs_approval": True,
|
|
"header": "test",
|
|
"func_name": "bash",
|
|
"approval_label": "bash: ls",
|
|
},
|
|
]
|
|
)
|
|
|
|
t = threading.Thread(target=call_approve)
|
|
t.start()
|
|
time.sleep(0.1)
|
|
|
|
# Should have a waiting-for-approval message in the buffer
|
|
with bg.ui._print_lock:
|
|
info_msgs = [text for ev, text in bg.ui._output_buffer if ev == "info"]
|
|
assert any("bash: ls" in msg for msg in info_msgs)
|
|
|
|
# Unblock
|
|
mgr.switch(bg.id)
|
|
bg.ui.set_foreground(True)
|
|
t.join(timeout=2)
|
|
assert result[0] == (True, None)
|
|
|
|
|
|
class TestNoColor:
|
|
"""Verify NO_COLOR support in colors module."""
|
|
|
|
def test_no_color_env_disables_ansi(self):
|
|
import importlib
|
|
import os
|
|
|
|
import turnstone.ui.colors as colors_mod
|
|
|
|
old_env = os.environ.get("NO_COLOR")
|
|
try:
|
|
os.environ["NO_COLOR"] = "1"
|
|
importlib.reload(colors_mod)
|
|
assert colors_mod.RESET == ""
|
|
assert colors_mod.BOLD == ""
|
|
assert colors_mod.RED == ""
|
|
assert colors_mod.red("test") == "test"
|
|
finally:
|
|
if old_env is None:
|
|
os.environ.pop("NO_COLOR", None)
|
|
else:
|
|
os.environ["NO_COLOR"] = old_env
|
|
importlib.reload(colors_mod)
|