Something went wrong. Try again.
A declarative, hermetic harness for Nix-defined agents.
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532import asyncioimport shleximport subprocessfrom collections.abc import AsyncIteratorfrom typing import cast
from tartarus.agent_loop import AgentLoop, ToolFinished, ToolStartedfrom tartarus.broker import Brokerfrom tartarus.context import ContextLedger, ContextLimits, ContextManagerfrom tartarus.jail import ExecResult, JailBuilderfrom tartarus.models import ( AssistantTurn, StreamEvent, TextDelta, ToolOutputDelta, ToolCall, ToolResult, TurnComplete,)from tartarus.policy import PolicyEnginefrom tests.manifest_fixtures import echo_manifest
class LocalJail: """A jail double that runs the command directly (no bwrap), so the loop's broker integration can be exercised without Linux namespaces."""
def build(self, grant): return grant
async def exec(self, spec, command, timeout=None, output_callback=None): completed = subprocess.run(shlex.split(command), capture_output=True, text=True) if output_callback is not None and completed.stdout: output_callback(completed.stdout) return ExecResult(completed.returncode, completed.stdout, completed.stderr)
class ScriptedProvider: """A provider that replays a fixed list of turns as streams, recording the tool results it is handed back. Stands in for any real backend to prove the loop is provider-neutral and brokers tool calls correctly.
Each turn's text is emitted as a single TextDelta followed by TurnComplete, so the streaming loop drives identically to a real SSE backend. """
def __init__(self, turns: list[AssistantTurn]): self._turns = list(turns) self.received_results: list = [] self.received_messages: list[list[dict]] = []
async def complete( self, system: str, messages: list[dict], tools: list[dict] ) -> AssistantTurn: return self._turns.pop(0)
async def stream( self, system: str, messages: list[dict], tools: list[dict] ) -> AsyncIterator[StreamEvent]: self.received_messages.append(list(messages)) turn = self._turns.pop(0) if turn.text: yield TextDelta(turn.text) yield TurnComplete(turn)
def assistant_message(self, turn: AssistantTurn) -> dict: return {"role": "assistant", "content": turn.text or ""}
def tool_result_messages(self, results: list) -> list[dict]: self.received_results.extend(results) return [ {"role": "tool", "tool_call_id": r.call_id, "content": r.output} for r in results ]
def adapt_tools(self, tools: list[dict]) -> list[dict]: return tools
async def _drain(loop, messages): """Collect every UI event run_turn yields for one human turn.""" events = [] async for event in loop.run_turn(messages): events.append(event) return events
def _text(events): return "".join(e.text for e in events if isinstance(e, TextDelta))
def test_loop_brokers_tool_calls_then_streams_final_text(): manifest = echo_manifest() provider = ScriptedProvider( [ AssistantTurn( text=None, tool_calls=[ToolCall("call-1", "echo", {"message": "ping"})], raw={"role": "assistant"}, stop_reason="tool_calls", ), AssistantTurn( text="The tool said ping.", tool_calls=[], raw={"role": "assistant"}, stop_reason="end", ), ] ) loop = AgentLoop( provider, Broker(manifest, cast(JailBuilder, LocalJail()), PolicyEngine()), manifest, "system", )
messages = [{"role": "user", "content": "use echo"}] events = asyncio.run(_drain(loop, messages))
assert _text(events) == "The tool said ping." # A ToolStarted/ToolFinished pair surfaced the brokered call. assert any(isinstance(e, ToolStarted) for e in events) finished = [e for e in events if isinstance(e, ToolFinished)] assert finished and finished[0].result.output == "ping" assert len(provider.received_results) == 1 # transcript: user, assistant(tool_calls), tool result, assistant(final) assert len(messages) == 4
def test_loop_streams_tool_output_before_tool_finished(): manifest = echo_manifest() provider = ScriptedProvider( [ AssistantTurn( text=None, tool_calls=[ToolCall("call-1", "echo", {"message": "ping"})], raw={"role": "assistant"}, stop_reason="tool_calls", ), AssistantTurn( text="done", tool_calls=[], raw={"role": "assistant"}, stop_reason="end", ), ] ) loop = AgentLoop( provider, Broker(manifest, cast(JailBuilder, LocalJail()), PolicyEngine()), manifest, "system", )
events = asyncio.run(_drain(loop, [{"role": "user", "content": "use echo"}]))
delta_index = next( index for index, event in enumerate(events) if isinstance(event, ToolOutputDelta) ) finished_index = next( index for index, event in enumerate(events) if isinstance(event, ToolFinished) ) assert events[delta_index].text == "ping\n" assert delta_index < finished_index
def test_loop_stops_immediately_when_no_tools_requested(): manifest = echo_manifest() provider = ScriptedProvider( [ AssistantTurn( text="Hello, no tools needed.", tool_calls=[], raw={"role": "assistant"}, stop_reason="end", ) ] ) loop = AgentLoop( provider, Broker(manifest, cast(JailBuilder, LocalJail()), PolicyEngine()), manifest, "system", )
messages = [{"role": "user", "content": "hi"}] events = asyncio.run(_drain(loop, messages))
assert _text(events) == "Hello, no tools needed." assert provider.received_results == [] # transcript: user, assistant(final) assert len(messages) == 2
def test_cancel_mid_turn_leaves_a_valid_transcript(): """Cancelling while the first round-trip is still streaming must not commit a partial assistant message: the transcript stays exactly the user message.""" manifest = echo_manifest()
class HangingProvider(ScriptedProvider): async def stream(self, system, messages, tools): yield TextDelta("thinking") # Never reaches TurnComplete; the consumer cancels mid-stream. await asyncio.Event().wait()
provider = HangingProvider([]) loop = AgentLoop( provider, Broker(manifest, cast(JailBuilder, LocalJail()), PolicyEngine()), manifest, "system", ) messages = [{"role": "user", "content": "hi"}]
async def run_and_cancel(): agen = loop.run_turn(messages) # Pull the first delta, then abort the generator mid-turn. first = await agen.__anext__() assert isinstance(first, TextDelta) await agen.aclose()
asyncio.run(run_and_cancel())
# Nothing partial was committed; only the original user message remains. assert messages == [{"role": "user", "content": "hi"}]
def test_aclose_mid_tool_cancels_worker_and_leaves_transcript(): """Closing the turn generator while a tool is running must cancel the worker so the jail terminates, instead of leaking it or committing a result.""" manifest = echo_manifest()
class CancellingBroker: def __init__(self): self.cancelled = False
async def handle(self, call, output_callback=None): # Emit a chunk so the caller can advance past ToolStarted, then block # until the worker is cancelled by the closing turn. if output_callback is not None: output_callback("partial output") try: await asyncio.Event().wait() except asyncio.CancelledError: self.cancelled = True raise return ToolResult(call.id, "should-not-appear", is_error=False)
provider = ScriptedProvider( [ AssistantTurn( text=None, tool_calls=[ToolCall("call-1", "echo", {"message": "ping"})], raw={"role": "assistant"}, stop_reason="tool_calls", ), ] ) broker = CancellingBroker() loop = AgentLoop(provider, cast(Broker, broker), manifest, "system") messages = [{"role": "user", "content": "use echo"}]
async def run_and_cancel(): agen = loop.run_turn(messages) async for event in agen: if isinstance(event, ToolOutputDelta): await agen.aclose() break
asyncio.run(run_and_cancel())
# The closing handshake cancelled the worker and did not commit anything. assert broker.cancelled assert messages == [{"role": "user", "content": "use echo"}]
def test_task_cancel_mid_tool_terminates_worker_synchronously(): """The CLI cancels an in-flight turn with ``task.cancel()`` (not ``aclose``). That CancelledError chains through ``__anext__`` into ``_run_tool``'s own ``await``, which cancels the worker and awaits its teardown before the cancelled task returns.""" manifest = echo_manifest()
class CancellingBroker: def __init__(self): self.cancelled = False
async def handle(self, call, output_callback=None): if output_callback is not None: output_callback("partial output") try: await asyncio.Event().wait() except asyncio.CancelledError: self.cancelled = True raise return ToolResult(call.id, "should-not-appear", is_error=False)
provider = ScriptedProvider( [ AssistantTurn( text=None, tool_calls=[ToolCall("call-1", "echo", {"message": "ping"})], raw={"role": "assistant"}, stop_reason="tool_calls", ), ] ) broker = CancellingBroker() loop = AgentLoop(provider, cast(Broker, broker), manifest, "system") messages = [{"role": "user", "content": "use echo"}]
async def consume(agen, seen): async for event in agen: seen.append(event)
async def run_and_cancel(): agen = loop.run_turn(messages) seen: list = [] task = asyncio.create_task(consume(agen, seen)) # Wait until output is flowing, then cancel the task as the CLI does. while not any(isinstance(e, ToolOutputDelta) for e in seen): await asyncio.sleep(0) task.cancel() try: await task except asyncio.CancelledError: pass # Teardown is synchronous on this path: no extra loop pumping needed. return broker.cancelled
cancelled_on_return = asyncio.run(run_and_cancel())
assert cancelled_on_return assert messages == [{"role": "user", "content": "use echo"}]
def test_loop_survives_unknown_tool_call(): """When the model calls a tool not in the manifest, the broker returns an error result and the loop continues normally (no crash, tool result recorded).""" manifest = echo_manifest() provider = ScriptedProvider( [ AssistantTurn( text=None, tool_calls=[ToolCall("call-1", "nonexistent", {})], raw={"role": "assistant"}, stop_reason="tool_calls", ), AssistantTurn( text="I tried but the tool was not found.", tool_calls=[], raw={"role": "assistant"}, stop_reason="end", ), ] ) loop = AgentLoop( provider, Broker(manifest, cast(JailBuilder, LocalJail()), PolicyEngine()), manifest, "system", )
messages = [{"role": "user", "content": "use bad tool"}] events = asyncio.run(_drain(loop, messages))
assert _text(events) == "I tried but the tool was not found." finished = [e for e in events if isinstance(e, ToolFinished)] assert len(finished) == 1 assert finished[0].result.is_error assert "unknown tool" in finished[0].result.output
def test_loop_sends_effective_messages_without_mutating_raw_transcript(tmp_path): manifest = echo_manifest() ledger = ContextLedger(str(tmp_path), "s1") ledger.append_event( { "type": "context_summary", "covered": {"start": 0, "end": 2}, "summary": "Earlier work was completed.", "source": "deterministic-local", "estimated_chars": 29, } ) provider = ScriptedProvider( [ AssistantTurn( text="done", tool_calls=[], raw={"role": "assistant"}, stop_reason="end", ) ] ) loop = AgentLoop( provider, Broker(manifest, cast(JailBuilder, LocalJail()), PolicyEngine()), manifest, "system", ContextManager(ledger, ContextLimits(max_chars=10_000, recent_turns=1)), ) messages = [ {"role": "user", "content": "old"}, {"role": "assistant", "content": "old reply"}, {"role": "user", "content": "new"}, ]
asyncio.run(_drain(loop, messages))
assert provider.received_messages[0][0]["role"] == "system" assert "Earlier work" in provider.received_messages[0][0]["content"] assert provider.received_messages[0][1:] == [{"role": "user", "content": "new"}] assert messages == [ {"role": "user", "content": "old"}, {"role": "assistant", "content": "old reply"}, {"role": "user", "content": "new"}, {"role": "assistant", "content": "done"}, ]
def test_loop_does_not_auto_compact_before_provider_call(tmp_path): manifest = echo_manifest() ledger = ContextLedger(str(tmp_path), "s1") provider = ScriptedProvider( [ AssistantTurn( text="done", tool_calls=[], raw={"role": "assistant"}, stop_reason="end", ) ] ) loop = AgentLoop( provider, Broker(manifest, cast(JailBuilder, LocalJail()), PolicyEngine()), manifest, "system", ContextManager( ledger, ContextLimits(max_chars=100, recent_turns=1), auto_compact=True, ), ) messages = [ {"role": "user", "content": "old " + "x" * 200}, {"role": "assistant", "content": "old reply " + "y" * 200}, {"role": "user", "content": "new " + "z" * 200}, ]
asyncio.run(_drain(loop, messages))
assert [event["type"] for event in ledger.load_events()] == []
def test_loop_handles_context_status_as_internal_tool(tmp_path): manifest = echo_manifest() provider = ScriptedProvider( [ AssistantTurn( text=None, tool_calls=[ToolCall("call-1", "context_status", {})], raw={"role": "assistant"}, stop_reason="tool_calls", ), AssistantTurn( text="I checked context.", tool_calls=[], raw={"role": "assistant"}, stop_reason="end", ), ] ) loop = AgentLoop( provider, Broker(manifest, cast(JailBuilder, LocalJail()), PolicyEngine()), manifest, "system", ContextManager(ContextLedger(str(tmp_path), "s1")), ) messages = [{"role": "user", "content": "status"}]
events = asyncio.run(_drain(loop, messages))
assert _text(events) == "I checked context." assert len(provider.received_results) == 1 assert '"message_count": 1' in provider.received_results[0].output
def test_loop_brokers_multiple_parallel_tool_calls(): """One assistant turn with two tool calls brokers both, aggregates results.""" manifest = echo_manifest() provider = ScriptedProvider( [ AssistantTurn( text=None, tool_calls=[ ToolCall("call-1", "echo", {"message": "first"}), ToolCall("call-2", "echo", {"message": "second"}), ], raw={"role": "assistant"}, stop_reason="tool_calls", ), AssistantTurn( text="Both tools ran.", tool_calls=[], raw={"role": "assistant"}, stop_reason="end", ), ] ) loop = AgentLoop( provider, Broker(manifest, cast(JailBuilder, LocalJail()), PolicyEngine()), manifest, "system", )
messages = [{"role": "user", "content": "run both"}] events = asyncio.run(_drain(loop, messages))
assert _text(events) == "Both tools ran." finished = [e for e in events if isinstance(e, ToolFinished)] assert len(finished) == 2 results = {f.call.id: f.result.output for f in finished} assert results["call-1"] == "first" assert results["call-2"] == "second" assert len(provider.received_results) == 2