feat: 增加 AI Agent 自主诊断执行模式
This commit is contained in:
@@ -0,0 +1,48 @@
|
||||
import asyncio
|
||||
import json
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
import backend.main as main
|
||||
|
||||
|
||||
def parse_sse(event):
|
||||
return json.loads(event[6:].strip())
|
||||
|
||||
|
||||
class AgentApiTest(unittest.TestCase):
|
||||
def test_agent_stream_forwards_agent_events(self):
|
||||
async def fake_agent(report, question="", max_steps=6, execution_mode="ai"):
|
||||
self.assertEqual(execution_mode, "manual")
|
||||
yield {"type": "thinking", "step": 1, "message": "规划"}
|
||||
yield {"type": "final", "content": "完成"}
|
||||
|
||||
async def run():
|
||||
main._last_diagnosis = {"score": 60}
|
||||
main._last_report = "report"
|
||||
with patch.object(main, "run_diagnostic_agent", fake_agent):
|
||||
response = await main.agent_diagnose_stream(question="检查 Pod", execution_mode="manual")
|
||||
return [parse_sse(event) async for event in response.body_iterator]
|
||||
|
||||
events = asyncio.run(run())
|
||||
self.assertEqual(events[0]["type"], "start")
|
||||
self.assertEqual(events[-1]["type"], "done")
|
||||
self.assertEqual(events[-2]["type"], "final")
|
||||
|
||||
def test_execute_approved_command_once(self):
|
||||
approval = {"command": "kubectl scale deployment/api --replicas=2", "reason": "恢复副本"}
|
||||
|
||||
async def fake_execute(command):
|
||||
return {"command": command, "mode": "write", "returncode": 0, "stdout": "scaled", "stderr": ""}
|
||||
|
||||
async def run():
|
||||
with patch.object(main, "take_pending_approval", return_value=approval), patch.object(main, "execute_kubectl", fake_execute):
|
||||
return await main.execute_approved_agent_command(main.AgentApprovalRequest(approval_id="abc", approved=True))
|
||||
|
||||
result = asyncio.run(run())
|
||||
self.assertEqual(result["returncode"], 0)
|
||||
self.assertEqual(result["stdout"], "scaled")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,38 @@
|
||||
import unittest
|
||||
|
||||
from backend.agent_tools import classify_kubectl_command, parse_kubectl_command
|
||||
|
||||
|
||||
class KubectlCommandPolicyTest(unittest.TestCase):
|
||||
def test_allows_read_only_kubectl_command_for_autonomous_execution(self):
|
||||
decision = classify_kubectl_command("kubectl get pods -n production -o wide")
|
||||
|
||||
self.assertEqual(decision.mode, "read")
|
||||
self.assertFalse(decision.requires_approval)
|
||||
|
||||
def test_requires_approval_for_mutating_kubectl_command(self):
|
||||
decision = classify_kubectl_command("kubectl rollout restart deployment/api -n production")
|
||||
|
||||
self.assertEqual(decision.mode, "write")
|
||||
self.assertTrue(decision.requires_approval)
|
||||
|
||||
def test_rejects_shell_operators_and_non_kubectl_commands(self):
|
||||
for command in (
|
||||
"kubectl get pods; rm -rf /",
|
||||
"kubectl get pods | sh",
|
||||
"bash -c kubectl get pods",
|
||||
"sudo kubectl get pods",
|
||||
):
|
||||
with self.subTest(command=command):
|
||||
with self.assertRaises(ValueError):
|
||||
parse_kubectl_command(command)
|
||||
|
||||
def test_rejects_exec_even_when_command_looks_read_only(self):
|
||||
decision = classify_kubectl_command("kubectl exec api-0 -- cat /etc/hosts")
|
||||
|
||||
self.assertEqual(decision.mode, "write")
|
||||
self.assertTrue(decision.requires_approval)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,97 @@
|
||||
import asyncio
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
from backend.ai_agent import run_diagnostic_agent
|
||||
|
||||
|
||||
async def collect_events(stream):
|
||||
return [event async for event in stream]
|
||||
|
||||
|
||||
class DiagnosticAgentTest(unittest.TestCase):
|
||||
def test_executes_read_only_tool_and_returns_final_answer(self):
|
||||
responses = iter([
|
||||
'{"action":"command","command":"kubectl get pods -A","reason":"查看异常 Pod"}',
|
||||
'{"action":"final","analysis":"发现一个异常 Pod"}',
|
||||
])
|
||||
|
||||
async def fake_completion(messages):
|
||||
return next(responses)
|
||||
|
||||
async def fake_execute(command):
|
||||
return {"command": command, "mode": "read", "returncode": 0, "stdout": "pod-a Running", "stderr": ""}
|
||||
|
||||
async def run():
|
||||
with patch("backend.ai_agent._complete_chat", fake_completion), patch("backend.ai_agent.execute_kubectl", fake_execute):
|
||||
return await collect_events(run_diagnostic_agent("report", max_steps=3))
|
||||
|
||||
events = asyncio.run(run())
|
||||
self.assertEqual([event["type"] for event in events], ["thinking", "command", "result", "thinking", "final"])
|
||||
self.assertEqual(events[1]["mode"], "read")
|
||||
self.assertEqual(events[-1]["content"], "发现一个异常 Pod")
|
||||
|
||||
def test_pauses_mutating_command_for_human_approval(self):
|
||||
async def fake_completion(messages):
|
||||
return '{"action":"command","command":"kubectl rollout restart deployment/api -n prod","reason":"重启故障工作负载"}'
|
||||
|
||||
async def run():
|
||||
with patch("backend.ai_agent._complete_chat", fake_completion):
|
||||
return await collect_events(run_diagnostic_agent("report", max_steps=3))
|
||||
|
||||
events = asyncio.run(run())
|
||||
self.assertEqual(events[-1]["type"], "approval_required")
|
||||
self.assertEqual(events[-1]["command"], "kubectl rollout restart deployment/api -n prod")
|
||||
self.assertTrue(events[-1]["approval_id"])
|
||||
|
||||
def test_manual_mode_never_executes_command_and_waits_for_result(self):
|
||||
async def fake_completion(messages):
|
||||
return '{"action":"command","command":"kubectl get pods -A","reason":"查看 Pod"}'
|
||||
|
||||
async def run():
|
||||
with patch("backend.ai_agent._complete_chat", fake_completion), patch("backend.ai_agent.execute_kubectl") as execute:
|
||||
events = await collect_events(run_diagnostic_agent("report", max_steps=3, execution_mode="manual"))
|
||||
return events, execute
|
||||
|
||||
events, execute = asyncio.run(run())
|
||||
execute.assert_not_called()
|
||||
self.assertEqual(events[-1]["type"], "manual_command")
|
||||
self.assertEqual(events[-1]["command"], "kubectl get pods -A")
|
||||
self.assertTrue(events[-1]["command_id"])
|
||||
|
||||
def test_ai_mode_keeps_write_command_behind_approval(self):
|
||||
async def fake_completion(messages):
|
||||
return '{"action":"command","command":"kubectl delete pod api-0 -n prod","reason":"删除故障 Pod"}'
|
||||
|
||||
async def run():
|
||||
with patch("backend.ai_agent._complete_chat", fake_completion), patch("backend.ai_agent.execute_kubectl") as execute:
|
||||
events = await collect_events(run_diagnostic_agent("report", max_steps=3, execution_mode="ai"))
|
||||
return events, execute
|
||||
|
||||
events, execute = asyncio.run(run())
|
||||
execute.assert_not_called()
|
||||
self.assertEqual(events[-1]["type"], "approval_required")
|
||||
|
||||
def test_rejects_unknown_execution_mode(self):
|
||||
async def run():
|
||||
return await collect_events(run_diagnostic_agent("report", execution_mode="invalid"))
|
||||
|
||||
events = asyncio.run(run())
|
||||
self.assertEqual(events[-1]["type"], "error")
|
||||
self.assertIn("执行模式", events[-1]["message"])
|
||||
|
||||
def test_stops_when_model_returns_invalid_action(self):
|
||||
async def fake_completion(messages):
|
||||
return "not-json"
|
||||
|
||||
async def run():
|
||||
with patch("backend.ai_agent._complete_chat", fake_completion):
|
||||
return await collect_events(run_diagnostic_agent("report", max_steps=3))
|
||||
|
||||
events = asyncio.run(run())
|
||||
self.assertEqual(events[-1]["type"], "error")
|
||||
self.assertIn("JSON", events[-1]["message"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user