cnbugs
32bea0d4ea
扩展 Agent 命令工具，除 kubectl 外新增节点宿主机系统诊断命令支持：
- agent_tools.py: 新增 shell 命令白名单分类器 (classify_shell_command)
- 只读白名单: systemctl status/is-active、journalctl、crictl ps/logs、
docker ps/logs、df、free、ss、ip addr、ps、mount、lsmod、dmesg、
ping、nslookup、cat (限 /proc /sys /etc/kubernetes 等安全路径) 等 30+ 命令
- 写白名单: systemctl restart/stop/start/reload 等需人工审批
- 统一分发器 classify_agent_command 按命令前缀路由
- 统一执行入口 execute_command 走 subprocess_exec (无 shell 注入风险)
- ai_agent.py: 更新 AGENT_SYSTEM_PROMPT 告知 AI 两类可用工具及安全规则
- 三个调用点 (初始探测/分类/执行) 切换到统一接口
- main.py: 审批执行处改用 execute_command
- tests: 新增 19 个用例覆盖 shell 命令分类、混合执行、白名单拒绝
- 全部 44 个测试通过
- 前端/README: 文案与文档补充 shell 命令支持说明
安全策略不变: 禁止 shell 操作符/管道/重定向/多命令拼接，
非白名单命令 (rm/vi/apt 等) 直接拒绝。
329 lines
15 KiB
Python
329 lines
15 KiB
Python
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","verified":true,"remaining_issues":0}',
|
|
])
|
|
|
|
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_command", 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],
|
|
["command", "result", "thinking", "command", "result", "thinking", "final"],
|
|
)
|
|
self.assertEqual(events[3]["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_command") 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"}'
|
|
|
|
executed = []
|
|
|
|
async def run():
|
|
async def fake_execute(command):
|
|
executed.append(command)
|
|
return {"command": command, "mode": "read", "returncode": 0, "stdout": "nodes ready", "stderr": ""}
|
|
|
|
with patch("backend.ai_agent._complete_chat", fake_completion), patch("backend.ai_agent.execute_command", fake_execute):
|
|
return await collect_events(run_diagnostic_agent("report", max_steps=3, execution_mode="ai"))
|
|
|
|
events = asyncio.run(run())
|
|
self.assertEqual(executed, ["kubectl get nodes -o wide"])
|
|
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_ai_mode_forces_initial_read_only_probe_before_model_planning(self):
|
|
responses = iter([
|
|
'{"action":"final","analysis":"根据采集结果完成分析","verified":true,"remaining_issues":0}',
|
|
])
|
|
|
|
async def fake_completion(messages):
|
|
return next(responses)
|
|
|
|
executed = []
|
|
|
|
async def fake_execute(command):
|
|
executed.append(command)
|
|
return {"command": command, "mode": "read", "returncode": 0, "stdout": "node-a Ready", "stderr": ""}
|
|
|
|
async def run():
|
|
with patch("backend.ai_agent._complete_chat", fake_completion), patch("backend.ai_agent.execute_command", fake_execute):
|
|
return await collect_events(run_diagnostic_agent("report", max_steps=3, execution_mode="ai"))
|
|
|
|
events = asyncio.run(run())
|
|
self.assertEqual(executed, ["kubectl get nodes -o wide"])
|
|
self.assertEqual([event["type"] for event in events[:2]], ["command", "result"])
|
|
self.assertEqual(events[-1]["type"], "final")
|
|
|
|
def test_parse_agent_action_accepts_markdown_json_fence(self):
|
|
from backend.ai_agent import _parse_agent_action
|
|
|
|
action = _parse_agent_action(
|
|
'我将先检查 Pod。\n```json\n{"action":"command","command":"kubectl get pods -A","reason":"查看 Pod"}\n```'
|
|
)
|
|
|
|
self.assertEqual(action["action"], "command")
|
|
self.assertEqual(action["command"], "kubectl get pods -A")
|
|
|
|
def test_ai_mode_continues_beyond_six_commands_until_model_finishes(self):
|
|
command_count = 8
|
|
responses = iter([
|
|
*[
|
|
(
|
|
'{"action":"command","command":"kubectl get pods -A '
|
|
f'--field-selector metadata.name=pod-{index}","reason":"继续收集证据"}}'
|
|
)
|
|
for index in range(command_count)
|
|
],
|
|
'{"action":"final","analysis":"已完成根因分析","verified":true,"remaining_issues":0}',
|
|
])
|
|
executed = []
|
|
|
|
async def fake_completion(messages):
|
|
return next(responses)
|
|
|
|
async def fake_execute(command):
|
|
executed.append(command)
|
|
return {"command": command, "mode": "read", "returncode": 0, "stdout": "evidence", "stderr": ""}
|
|
|
|
async def run():
|
|
with patch("backend.ai_agent._complete_chat", fake_completion), patch("backend.ai_agent.execute_command", fake_execute):
|
|
return await collect_events(run_diagnostic_agent("report", execution_mode="ai", max_steps=20))
|
|
|
|
events = asyncio.run(run())
|
|
self.assertEqual(len(executed), command_count + 1)
|
|
self.assertEqual(events[-1]["type"], "final")
|
|
self.assertEqual(events[-1]["content"], "已完成根因分析")
|
|
self.assertNotIn("最大轮次", events[-1]["content"])
|
|
|
|
def test_safety_limit_reports_agent_error_instead_of_sending_user_to_manual_diagnosis(self):
|
|
async def fake_completion(messages):
|
|
return '{"action":"command","command":"kubectl get pods -A","reason":"继续调查"}'
|
|
|
|
async def fake_execute(command):
|
|
return {"command": command, "mode": "read", "returncode": 0, "stdout": "same result", "stderr": ""}
|
|
|
|
async def run():
|
|
with patch("backend.ai_agent._complete_chat", fake_completion), patch("backend.ai_agent.execute_command", fake_execute):
|
|
return await collect_events(run_diagnostic_agent("report", execution_mode="ai", max_steps=2))
|
|
|
|
events = asyncio.run(run())
|
|
self.assertEqual(events[-1]["type"], "error")
|
|
self.assertIn("安全上限", events[-1]["message"])
|
|
self.assertNotIn("人工排查", events[-1]["message"])
|
|
|
|
def test_agent_rejects_final_until_verification_reports_no_remaining_issues(self):
|
|
responses = iter([
|
|
'{"action":"final","analysis":"已经修复"}',
|
|
'{"action":"command","command":"kubectl get pods -A","reason":"执行全量复检"}',
|
|
'{"action":"final","analysis":"所有工作负载恢复正常","verified":true,"remaining_issues":0}',
|
|
])
|
|
executed = []
|
|
|
|
async def fake_completion(messages):
|
|
return next(responses)
|
|
|
|
async def fake_execute(command):
|
|
executed.append(command)
|
|
stdout = "node-a Ready" if command == "kubectl get nodes -o wide" else "all pods Running"
|
|
return {"command": command, "mode": "read", "returncode": 0, "stdout": stdout, "stderr": ""}
|
|
|
|
async def run():
|
|
with patch("backend.ai_agent._complete_chat", fake_completion), patch("backend.ai_agent.execute_command", fake_execute):
|
|
return await collect_events(run_diagnostic_agent("report", execution_mode="ai", max_steps=5))
|
|
|
|
events = asyncio.run(run())
|
|
self.assertEqual(executed, ["kubectl get nodes -o wide", "kubectl get pods -A"])
|
|
self.assertEqual(events[-1]["type"], "final")
|
|
self.assertTrue(events[-1]["verified"])
|
|
self.assertEqual(events[-1]["remaining_issues"], 0)
|
|
|
|
def test_parse_final_requires_explicit_verification_fields(self):
|
|
from backend.ai_agent import _parse_agent_action
|
|
|
|
with self.assertRaises(ValueError):
|
|
_parse_agent_action('{"action":"final","analysis":"看起来正常"}')
|
|
|
|
action = _parse_agent_action(
|
|
'{"action":"final","analysis":"全部正常","verified":true,"remaining_issues":0}'
|
|
)
|
|
self.assertTrue(action["verified"])
|
|
self.assertEqual(action["remaining_issues"], 0)
|
|
|
|
def test_invalid_json_response_is_retried_instead_of_stopping_agent(self):
|
|
responses = iter([
|
|
"我需要继续检查，但没有按格式输出",
|
|
'{"action":"command","command":"kubectl get pods -A","reason":"继续检查"}',
|
|
'{"action":"final","analysis":"检查完成","verified":true,"remaining_issues":0}',
|
|
])
|
|
executed = []
|
|
|
|
async def fake_completion(messages):
|
|
return next(responses)
|
|
|
|
async def fake_execute(command):
|
|
executed.append(command)
|
|
return {"command": command, "mode": "read", "returncode": 0, "stdout": "healthy", "stderr": ""}
|
|
|
|
async def run():
|
|
with patch("backend.ai_agent._complete_chat", fake_completion), patch("backend.ai_agent.execute_command", fake_execute):
|
|
return await collect_events(run_diagnostic_agent("report", max_steps=5))
|
|
|
|
events = asyncio.run(run())
|
|
self.assertTrue(any(event["type"] == "response_rejected" for event in events))
|
|
self.assertEqual(events[-1]["type"], "final")
|
|
self.assertEqual(executed, ["kubectl get nodes -o wide", "kubectl get pods -A"])
|
|
|
|
def test_stops_after_three_consecutive_invalid_json_responses(self):
|
|
async def fake_completion(messages):
|
|
return "still invalid"
|
|
|
|
async def fake_execute(command):
|
|
return {"command": command, "mode": "read", "returncode": 0, "stdout": "node Ready", "stderr": ""}
|
|
|
|
async def run():
|
|
with patch("backend.ai_agent._complete_chat", fake_completion), patch("backend.ai_agent.execute_command", fake_execute):
|
|
return await collect_events(run_diagnostic_agent("report", max_steps=10))
|
|
|
|
events = asyncio.run(run())
|
|
self.assertEqual(events[-1]["type"], "error")
|
|
self.assertIn("连续 3 次", 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"])
|
|
|
|
def test_agent_can_mix_kubectl_and_shell_commands(self):
|
|
"""Agent 应能在同一轮诊断中交替使用 kubectl 和系统诊断 shell 命令。"""
|
|
responses = iter([
|
|
'{"action":"command","command":"kubectl get nodes -o wide","reason":"查看节点状态"}',
|
|
'{"action":"command","command":"systemctl status kubelet","reason":"检查 kubelet 服务"}',
|
|
'{"action":"command","command":"journalctl -u kubelet -n 50 --no-pager","reason":"查看 kubelet 日志"}',
|
|
'{"action":"final","analysis":"kubelet 服务正常，节点就绪","verified":true,"remaining_issues":0}',
|
|
])
|
|
executed = []
|
|
|
|
async def fake_completion(messages):
|
|
return next(responses)
|
|
|
|
async def fake_execute(command):
|
|
executed.append(command)
|
|
if command.startswith("kubectl"):
|
|
stdout = "node-a Ready"
|
|
elif command.startswith("systemctl"):
|
|
stdout = "active (running)"
|
|
else:
|
|
stdout = "no errors"
|
|
return {"command": command, "mode": "read", "returncode": 0, "stdout": stdout, "stderr": ""}
|
|
|
|
async def run():
|
|
with patch("backend.ai_agent._complete_chat", fake_completion), patch("backend.ai_agent.execute_command", fake_execute):
|
|
return await collect_events(run_diagnostic_agent("report", execution_mode="ai", max_steps=10))
|
|
|
|
events = asyncio.run(run())
|
|
# 初始探测 + 3 条 Agent 命令 = 4 次执行
|
|
self.assertEqual(
|
|
executed,
|
|
["kubectl get nodes -o wide", "kubectl get nodes -o wide", "systemctl status kubelet", "journalctl -u kubelet -n 50 --no-pager"],
|
|
)
|
|
# 所有命令都是只读模式 (自动执行)
|
|
command_events = [e for e in events if e["type"] == "command"]
|
|
self.assertTrue(all(e["mode"] == "read" for e in command_events))
|
|
self.assertEqual(events[-1]["type"], "final")
|
|
|
|
def test_agent_rejects_non_whitelisted_shell_command(self):
|
|
"""Agent 应拒绝不在白名单内的 shell 命令 (如 rm)。"""
|
|
responses = iter([
|
|
'{"action":"command","command":"rm -rf /tmp/test","reason":"清理测试文件"}',
|
|
'{"action":"final","analysis":"已完成","verified":true,"remaining_issues":0}',
|
|
])
|
|
executed = []
|
|
|
|
async def fake_completion(messages):
|
|
return next(responses)
|
|
|
|
async def fake_execute(command):
|
|
executed.append(command)
|
|
return {"command": command, "mode": "read", "returncode": 0, "stdout": "", "stderr": ""}
|
|
|
|
async def run():
|
|
with patch("backend.ai_agent._complete_chat", fake_completion), patch("backend.ai_agent.execute_command", fake_execute):
|
|
return await collect_events(run_diagnostic_agent("report", execution_mode="ai", max_steps=5))
|
|
|
|
events = asyncio.run(run())
|
|
# rm 命令应被拒绝
|
|
rejected = [e for e in events if e["type"] == "command_rejected"]
|
|
self.assertTrue(rejected)
|
|
self.assertIn("不允许执行该命令", rejected[0]["reason"])
|
|
# rm 不应被执行
|
|
self.assertNotIn("rm -rf /tmp/test", executed)
|
|
self.assertEqual(events[-1]["type"], "final")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|