#!/usr/bin/env python3 """Tests for zubble agent loop orchestration. These tests mock only I/O (stream_chat HTTP calls, env.exec/verify), never the orchestration logic that wires detectors, protocol, ledger, and abort handling together. python3 -m unittest discover -s tests -v """ import os import sys import json import time import unittest from unittest.mock import MagicMock, patch sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from zubble import benchlib # noqa: E402 from zubble.agent import run_agent_loop # noqa: E402 from zubble import protocol # noqa: E402 class _MockEnv: """Minimal env stub matching scenariobench's interface.""" intro = "You have a root shell on a host." report_path = "/root/report.md" def __init__(self, exec_results=None, verify_result=(1, "unhealthy")): self._exec_results = exec_results or [] self._exec_idx = 0 self._verify_result = verify_result def exec(self, cmd, timeout=90): if self._exec_idx < len(self._exec_results): r = self._exec_results[self._exec_idx] self._exec_idx += 1 return r return (0, "") def verify(self): return self._verify_result def setup(self): return True, "ok" def teardown(self): pass def motd(self): return "welcome" def report(self): return "" def _make_stream_response(content): """Build a minimal perf dict + content for a mock stream_chat return.""" perf = dict(ttft=0.01, t_answer=0.01, wall=0.01, reasoning=0, visible=len(content), total=len(content), decode=0.0, user_tps=0.0, prompt_tokens=10, prefill=0.0, finish="stop") return content, "", perf class AgentLoopActionKillTest(unittest.TestCase): """Verify that repeated identical actions trigger kill → fresh context → reset cap → doom_loop.""" def test_action_repetition_leads_to_doom_loop(self): repeated_action = json.dumps({ "action": "run", "command": "kubectl get pods", "plan_step_id": 1, "reasoning": "checking pods", }) def mock_stream_chat(*args, **kwargs): return _make_stream_response(repeated_action) with patch.object(benchlib, "stream_chat", side_effect=mock_stream_chat): env = _MockEnv(exec_results=[(0, "NAME READY STATUS\npod-a Running")] * 100) deadline = time.monotonic() + 10 result = run_agent_loop( "http://x", env, "Fix the cluster", "Connected.", deadline, "off", None, max_turns=200, maxtok=4000, name="test") self.assertEqual(result["outcome"], "doom_loop") self.assertTrue(result["turns"] > 0) def test_plan_then_run_then_submit_passes(self): responses = [ json.dumps({"action": "plan", "steps": ["diagnose", "fix", "verify"], "reasoning": "3-step plan"}), json.dumps({"action": "run", "command": "kubectl get pods", "plan_step_id": 1, "reasoning": "step 1"}), json.dumps({"action": "run", "command": "kubectl apply -f fix.yaml", "plan_step_id": 2, "reasoning": "step 2"}), json.dumps({"action": "submit", "report": "Fixed the cert.", "reasoning": "done"}), ] idx = [0] def mock_stream_chat(*args, **kwargs): r = _make_stream_response(responses[idx[0]]) idx[0] += 1 return r with patch.object(benchlib, "stream_chat", side_effect=mock_stream_chat): env = _MockEnv( exec_results=[(0, "pods running"), (0, "applied")], verify_result=(0, "cluster healthy"), ) deadline = time.monotonic() + 10 result = run_agent_loop( "http://x", env, "Fix the cluster", "Connected.", deadline, "off", None, max_turns=200, maxtok=4000, name="test") self.assertTrue(result["passed"]) self.assertEqual(result["outcome"], "passed") self.assertEqual(result["submissions"], 1) self.assertEqual(result["commands"], 2) def test_agent_requests_constrained_json_schema(self): response = json.dumps({"action": "submit", "report": "fixed", "reasoning": "verified"}) with patch.object(benchlib, "stream_chat", return_value=_make_stream_response(response)) as stream_chat: env = _MockEnv(verify_result=(0, "cluster healthy")) result = run_agent_loop( "http://x", env, "Fix the cluster", "Connected.", time.monotonic() + 10, "off", None, max_turns=200, maxtok=4000, name="test") self.assertTrue(result["passed"]) self.assertEqual(stream_chat.call_args.kwargs["response_format"], protocol.RESPONSE_FORMAT) def test_rejected_submit_continues_until_later_pass(self): responses = [ json.dumps({"action": "plan", "steps": ["check"], "reasoning": "plan"}), json.dumps({"action": "run", "command": "kubectl get pods", "plan_step_id": 1, "reasoning": "check"}), json.dumps({"action": "submit", "report": "first attempt", "reasoning": "ready to verify"}), json.dumps({"action": "run", "command": "kubectl apply -f fix.yaml", "plan_step_id": 1, "reasoning": "apply correction"}), json.dumps({"action": "submit", "report": "fix applied", "reasoning": "ready to verify again"}), ] idx = [0] verify_calls = [0] def mock_stream_chat(*args, **kwargs): r = _make_stream_response(responses[idx[0]]) idx[0] += 1 return r def verify(): verify_calls[0] += 1 return ((1, "still unhealthy") if verify_calls[0] == 1 else (0, "cluster healthy")) with patch.object(benchlib, "stream_chat", side_effect=mock_stream_chat): env = _MockEnv(exec_results=[(0, "pods running"), (0, "applied")]) env.verify = verify deadline = time.monotonic() + 10 result = run_agent_loop( "http://x", env, "Fix the cluster", "Connected.", deadline, "off", None, max_turns=200, maxtok=4000, name="test") self.assertTrue(result["passed"]) self.assertEqual(result["outcome"], "passed") self.assertEqual(result["submissions"], 2) self.assertEqual(result["rejected"], 1) self.assertEqual(result["commands"], 2) def test_malformed_response_then_valid_repair_executes_action(self): responses = [ '{"action":"finished","reasoning":"done"}', json.dumps({"action": "submit", "report": "fixed", "reasoning": "ready to verify"}), ] idx = [0] def mock_stream_chat(*args, **kwargs): r = _make_stream_response(responses[idx[0]]) idx[0] += 1 return r with patch.object(benchlib, "stream_chat", side_effect=mock_stream_chat): env = _MockEnv(verify_result=(0, "cluster healthy")) result = run_agent_loop( "http://x", env, "Fix the cluster", "Connected.", time.monotonic() + 10, "off", None, max_turns=200, maxtok=4000, name="test") self.assertTrue(result["passed"]) self.assertEqual(result["submissions"], 1) self.assertEqual(result["parse_failures"], 1) def test_malformed_response_and_repair_end_as_parse_error(self): def mock_stream_chat(*args, **kwargs): return _make_stream_response('{"action":"submit","report":"bad\nreport"}') with patch.object(benchlib, "stream_chat", side_effect=mock_stream_chat): env = _MockEnv() result = run_agent_loop( "http://x", env, "Fix the cluster", "Connected.", time.monotonic() + 10, "off", None, max_turns=200, maxtok=4000, name="test") self.assertFalse(result["passed"]) self.assertEqual(result["outcome"], "parse_error") self.assertEqual(result["turns"], 1) self.assertEqual(result["parse_failures"], 2) self.assertEqual(result["submissions"], 0) class ProtocolTest(unittest.TestCase): def test_final_and_finished_are_rejected_with_submit_guidance(self): for action in ("final", "finished"): with self.subTest(action=action): with self.assertRaisesRegex(protocol.ProtocolError, "use 'submit'"): protocol.parse_response(json.dumps({ "action": action, "report": "done", "reasoning": "done", })) def test_trailing_content_is_rejected(self): with self.assertRaisesRegex(protocol.ProtocolError, "trailing content"): protocol.parse_response( '{"action":"submit","report":"done","reasoning":"done"} trailing') def test_long_escaped_report_is_valid(self): response = json.dumps({ "action": "submit", "report": "line one\nline two with {braces} and \\\"quotes\\\"", "reasoning": "verified", }) self.assertEqual(protocol.parse_response(response)["action"], "submit") if __name__ == "__main__": unittest.main()