wehub-resource-sync
320 lines
9.7 KiB
Python
320 lines
9.7 KiB
Python
"""P2 gate tests — turn engine + event bus (scripted provider, no network)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import threading
|
|
import time
|
|
|
|
import aisuite as ai
|
|
from coworker.engine import ApprovalOutcome, PermissionRequest, TurnEngine
|
|
from coworker.events import EventType
|
|
from coworker.permissions import PermissionEngine
|
|
from coworker.providers import (
|
|
AssistantTurn,
|
|
ModelCapabilities,
|
|
ProviderClient,
|
|
StreamChunk,
|
|
ToolCall,
|
|
)
|
|
from coworker.tools import ToolRegistry
|
|
|
|
|
|
def _text_turn(text):
|
|
return AssistantTurn(text=text, finish_reason="stop")
|
|
|
|
|
|
def _tool_turn(name, args, call_id="call_1"):
|
|
return AssistantTurn(
|
|
tool_calls=[ToolCall(id=call_id, name=name, arguments=args)],
|
|
finish_reason="tool_calls",
|
|
)
|
|
|
|
|
|
class ScriptedProvider(ProviderClient):
|
|
"""Returns queued AssistantTurns; streams via the base default (one final chunk)."""
|
|
|
|
def __init__(self, turns, *, loop=False):
|
|
self._turns = list(turns)
|
|
self._loop = loop
|
|
self.calls = 0
|
|
|
|
def complete(self, *, model, messages, tools=None, **settings):
|
|
self.calls += 1
|
|
return self._turns[0] if self._loop else self._turns.pop(0)
|
|
|
|
def capabilities(self, model):
|
|
return ModelCapabilities()
|
|
|
|
|
|
def _engine(tmp_path, turns, *, approver=None, loop=False, max_iterations=12):
|
|
provider = ScriptedProvider(turns, loop=loop)
|
|
registry = ToolRegistry()
|
|
registry.register_all(ai.toolkits.files(root=str(tmp_path), allow_write=True))
|
|
permissions = PermissionEngine(workspace_root=tmp_path)
|
|
engine = TurnEngine(
|
|
provider=provider,
|
|
registry=registry,
|
|
permissions=permissions,
|
|
model="gpt-5.5",
|
|
approver=approver,
|
|
max_iterations=max_iterations,
|
|
)
|
|
return engine, provider
|
|
|
|
|
|
def _collect(engine, user_input):
|
|
async def _run():
|
|
return [ev async for ev in engine.run(user_input)]
|
|
|
|
return asyncio.run(_run())
|
|
|
|
|
|
def _types(events):
|
|
return [ev.type for ev in events]
|
|
|
|
|
|
# -- tests ----------------------------------------------------------------------
|
|
|
|
|
|
def test_no_tool_turn(tmp_path):
|
|
engine, _ = _engine(tmp_path, [_text_turn("all done")])
|
|
events = _collect(engine, "hi")
|
|
assert _types(events) == [
|
|
EventType.TURN_START,
|
|
EventType.ASSISTANT_MESSAGE,
|
|
EventType.TURN_END,
|
|
]
|
|
assert events[1].data["text"] == "all done"
|
|
assert events[-1].data["status"] == "completed"
|
|
|
|
|
|
def test_tool_turn_order_and_execution(tmp_path):
|
|
(tmp_path / "a.txt").write_text("hello", encoding="utf-8")
|
|
engine, _ = _engine(
|
|
tmp_path,
|
|
[_tool_turn("read_file", {"path": "a.txt"}), _text_turn("it says hello")],
|
|
)
|
|
events = _collect(engine, "read a.txt")
|
|
assert EventType.PERMISSION_REQUIRED not in _types(events)
|
|
assert _types(events) == [
|
|
EventType.TURN_START,
|
|
EventType.ASSISTANT_MESSAGE,
|
|
EventType.TOOL_PROPOSED,
|
|
EventType.TOOL_STARTED,
|
|
EventType.TOOL_FINISHED,
|
|
EventType.ITERATION_END,
|
|
EventType.ASSISTANT_MESSAGE,
|
|
EventType.TURN_END,
|
|
]
|
|
finished = next(e for e in events if e.type == EventType.TOOL_FINISHED)
|
|
assert finished.data["status"] == "ok"
|
|
assert any(
|
|
m.get("role") == "tool" and "hello" in m["content"] for m in engine.messages
|
|
)
|
|
|
|
|
|
def test_write_requires_approval_then_approved(tmp_path):
|
|
async def approve_once(_req: PermissionRequest):
|
|
return ApprovalOutcome.ONCE
|
|
|
|
engine, _ = _engine(
|
|
tmp_path,
|
|
[
|
|
_tool_turn("write_file", {"path": "new.py", "content": "print(1)\n"}),
|
|
_text_turn("wrote new.py"),
|
|
],
|
|
approver=approve_once,
|
|
)
|
|
events = _collect(engine, "create new.py")
|
|
assert EventType.PERMISSION_REQUIRED in _types(events)
|
|
assert (tmp_path / "new.py").read_text() == "print(1)\n"
|
|
|
|
|
|
def test_denied_tool_yields_error_and_continues(tmp_path):
|
|
async def deny(_req: PermissionRequest):
|
|
return ApprovalOutcome.DENY
|
|
|
|
engine, _ = _engine(
|
|
tmp_path,
|
|
[
|
|
_tool_turn("write_file", {"path": "new.py", "content": "x"}),
|
|
_text_turn("ok, skipped it"),
|
|
],
|
|
approver=deny,
|
|
)
|
|
events = _collect(engine, "create new.py")
|
|
assert not (tmp_path / "new.py").exists()
|
|
finished = next(e for e in events if e.type == EventType.TOOL_FINISHED)
|
|
assert finished.data["status"] == "denied"
|
|
assert _types(events)[-1] == EventType.TURN_END
|
|
assert any(
|
|
m.get("role") == "tool" and "not executed" in m["content"]
|
|
for m in engine.messages
|
|
)
|
|
|
|
|
|
def test_max_iterations_rail(tmp_path):
|
|
engine, provider = _engine(
|
|
tmp_path, [_tool_turn("list_files", {})], loop=True, max_iterations=3
|
|
)
|
|
events = _collect(engine, "loop forever")
|
|
end = events[-1]
|
|
assert end.type == EventType.TURN_END
|
|
assert end.data["status"] == "max_iterations_exceeded"
|
|
assert provider.calls == 3
|
|
|
|
|
|
def test_interrupt_between_iterations(tmp_path):
|
|
engine_holder = {}
|
|
|
|
async def approve_and_interrupt(_req: PermissionRequest):
|
|
engine_holder["engine"].request_interrupt()
|
|
return ApprovalOutcome.ONCE
|
|
|
|
engine, provider = _engine(
|
|
tmp_path,
|
|
[
|
|
_tool_turn("write_file", {"path": "x.py", "content": "x"}),
|
|
_text_turn("should not be reached"),
|
|
],
|
|
approver=approve_and_interrupt,
|
|
)
|
|
engine_holder["engine"] = engine
|
|
events = _collect(engine, "do a thing")
|
|
assert events[-1].type == EventType.INTERRUPTED
|
|
assert provider.calls == 1
|
|
|
|
|
|
def test_steering_injects_next_turn(tmp_path):
|
|
engine, provider = _engine(tmp_path, [_text_turn("first"), _text_turn("second")])
|
|
engine.queue_steering("actually, also do this")
|
|
events = _collect(engine, "do the first thing")
|
|
assert provider.calls == 2
|
|
assert any(
|
|
m.get("role") == "user" and m["content"] == "actually, also do this"
|
|
for m in engine.messages
|
|
)
|
|
assert events[-1].data["status"] == "completed"
|
|
|
|
|
|
# -- parallel tool execution ------------------------------------------------------
|
|
|
|
|
|
def _multi_tool_turn(calls):
|
|
return AssistantTurn(
|
|
tool_calls=[
|
|
ToolCall(id=f"call_{i}", name=name, arguments=args)
|
|
for i, (name, args) in enumerate(calls)
|
|
],
|
|
finish_reason="tool_calls",
|
|
)
|
|
|
|
|
|
def _bare_engine(tmp_path, turns):
|
|
provider = ScriptedProvider(turns)
|
|
registry = ToolRegistry()
|
|
permissions = PermissionEngine(workspace_root=tmp_path)
|
|
engine = TurnEngine(
|
|
provider=provider,
|
|
registry=registry,
|
|
permissions=permissions,
|
|
model="gpt-5.5",
|
|
)
|
|
return engine, registry
|
|
|
|
|
|
def test_low_risk_tool_calls_run_concurrently(tmp_path):
|
|
# Both tools block on a 2-party barrier: the turn only completes if the engine
|
|
# really runs them at the same time (sequential execution would trip the timeout
|
|
# and surface as an error result).
|
|
barrier = threading.Barrier(2, timeout=5)
|
|
low = ai.ToolMetadata(category="search", risk_level="low", requires_approval=False)
|
|
|
|
def side_a():
|
|
"""Wait for side_b."""
|
|
barrier.wait()
|
|
return {"side": "a"}
|
|
|
|
def side_b():
|
|
"""Wait for side_a."""
|
|
barrier.wait()
|
|
return {"side": "b"}
|
|
|
|
engine, registry = _bare_engine(
|
|
tmp_path,
|
|
[_multi_tool_turn([("side_a", {}), ("side_b", {})]), _text_turn("done")],
|
|
)
|
|
registry.register(side_a, metadata=low)
|
|
registry.register(side_b, metadata=low)
|
|
|
|
events = _collect(engine, "go")
|
|
finished = [e for e in events if e.type == EventType.TOOL_FINISHED]
|
|
assert len(finished) == 2
|
|
assert all(e.data["status"] == "ok" for e in finished)
|
|
# a tool result message exists for every call id
|
|
tool_ids = {
|
|
m.get("tool_call_id") for m in engine.messages if m.get("role") == "tool"
|
|
}
|
|
assert tool_ids == {"call_0", "call_1"}
|
|
|
|
|
|
def test_non_low_risk_tool_calls_stay_sequential(tmp_path):
|
|
order = []
|
|
medium = ai.ToolMetadata(
|
|
category="filesystem", risk_level="medium", requires_approval=False
|
|
)
|
|
|
|
def first():
|
|
"""Record start/end with a delay."""
|
|
order.append("first-start")
|
|
time.sleep(0.2)
|
|
order.append("first-end")
|
|
return "ok"
|
|
|
|
def second():
|
|
"""Record start/end."""
|
|
order.append("second-start")
|
|
order.append("second-end")
|
|
return "ok"
|
|
|
|
engine, registry = _bare_engine(
|
|
tmp_path,
|
|
[_multi_tool_turn([("first", {}), ("second", {})]), _text_turn("done")],
|
|
)
|
|
registry.register(first, metadata=medium)
|
|
registry.register(second, metadata=medium)
|
|
|
|
_collect(engine, "go")
|
|
assert order == ["first-start", "first-end", "second-start", "second-end"]
|
|
|
|
|
|
class StreamingProvider(ProviderClient):
|
|
def complete(self, **kwargs): # pragma: no cover - streamed instead
|
|
raise NotImplementedError
|
|
|
|
def capabilities(self, model):
|
|
return ModelCapabilities()
|
|
|
|
def stream(self, *, model, messages, tools=None, **settings):
|
|
for piece in ["Hel", "lo, ", "world"]:
|
|
yield StreamChunk(text_delta=piece)
|
|
yield StreamChunk(turn=AssistantTurn(text="Hello, world", finish_reason="stop"))
|
|
|
|
|
|
def test_streaming_emits_deltas(tmp_path):
|
|
registry = ToolRegistry()
|
|
permissions = PermissionEngine(workspace_root=tmp_path)
|
|
engine = TurnEngine(
|
|
provider=StreamingProvider(),
|
|
registry=registry,
|
|
permissions=permissions,
|
|
model="gpt-5.5",
|
|
)
|
|
events = _collect(engine, "say hi")
|
|
deltas = [e.data["text"] for e in events if e.type == EventType.ASSISTANT_DELTA]
|
|
assert deltas == ["Hel", "lo, ", "world"]
|
|
final = next(e for e in events if e.type == EventType.ASSISTANT_MESSAGE)
|
|
assert final.data["text"] == "Hello, world"
|
|
assert events[-1].type == EventType.TURN_END
|