AI-Agent/k8smanager-cli
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 3 Wiki Activity

feat: AI 诊断 Agent 支持执行系统诊断 shell 命令

Browse Source 
扩展 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 等) 直接拒绝。
This commit is contained in:
cnbugs
2026-07-25 13:38:48 +08:00
parent f2254646c1
commit 32bea0d4ea
8 changed files with 430 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_agent_api.py
+2 -2
View File
@@ -62,7 +62,7 @@ class AgentApiTest(unittest.TestCase):
yield {"type": "final", "content": "故障已恢复", "model": "test"}
async def run():
with patch.object(main, "take_pending_approval", return_value=approval), patch.object(main, "execute_kubectl", fake_execute), patch.object(main, "resume_diagnostic_agent", fake_resume):
with patch.object(main, "take_pending_approval", return_value=approval), patch.object(main, "execute_command", fake_execute), patch.object(main, "resume_diagnostic_agent", fake_resume):
response = await main.execute_approved_agent_command(main.AgentApprovalRequest(approval_id="abc", approved=True))
return [parse_sse(event) async for event in response.body_iterator]
@@ -82,7 +82,7 @@ class AgentApiTest(unittest.TestCase):
yield {"type": "final", "content": "验证完成", "model": "test"}
async def run():
with patch.object(main, "take_pending_approval", return_value=approval), patch.object(main, "execute_kubectl", fake_execute), patch.object(main, "resume_diagnostic_agent", fake_resume):
with patch.object(main, "take_pending_approval", return_value=approval), patch.object(main, "execute_command", fake_execute), patch.object(main, "resume_diagnostic_agent", fake_resume):
response = await main.execute_approved_agent_command(main.AgentApprovalRequest(approval_id="abc", approved=True))
return [parse_sse(event) async for event in response.body_iterator]
tests/test_agent_tools.py
+95 -1
View File
@@ -1,6 +1,12 @@
import unittest
from backend.agent_tools import classify_kubectl_command, parse_kubectl_command
from backend.agent_tools import (
classify_agent_command,
classify_kubectl_command,
classify_shell_command,
parse_kubectl_command,
parse_shell_command,
)
class KubectlCommandPolicyTest(unittest.TestCase):
@@ -34,5 +40,93 @@ class KubectlCommandPolicyTest(unittest.TestCase):
self.assertTrue(decision.requires_approval)
class ShellCommandPolicyTest(unittest.TestCase):
def test_allows_read_only_systemctl_status_for_autonomous_execution(self):
decision = classify_shell_command("systemctl status kubelet")
self.assertEqual(decision.mode, "read")
self.assertFalse(decision.requires_approval)
def test_allows_journalctl_with_flags(self):
decision = classify_shell_command("journalctl -u kubelet -n 100 --no-pager")
self.assertEqual(decision.mode, "read")
self.assertFalse(decision.requires_approval)
def test_allows_crictl_ps(self):
decision = classify_shell_command("crictl ps -a")
self.assertEqual(decision.mode, "read")
self.assertFalse(decision.requires_approval)
def test_allows_df_and_free(self):
for command in ("df -h", "free -m", "uptime"):
with self.subTest(command=command):
decision = classify_shell_command(command)
self.assertEqual(decision.mode, "read")
self.assertFalse(decision.requires_approval)
def test_allows_ip_addr(self):
decision = classify_shell_command("ip addr show")
self.assertEqual(decision.mode, "read")
self.assertFalse(decision.requires_approval)
def test_requires_approval_for_systemctl_restart(self):
decision = classify_shell_command("systemctl restart kubelet")
self.assertEqual(decision.mode, "write")
self.assertTrue(decision.requires_approval)
def test_rejects_unknown_binary(self):
with self.assertRaises(ValueError):
classify_shell_command("rm -rf /")
def test_rejects_shell_operators_in_shell_command(self):
for command in (
"systemctl status kubelet; rm -rf /",
"df -h | grep vda",
"journalctl -u kubelet > /tmp/log",
):
with self.subTest(command=command):
with self.assertRaises(ValueError):
parse_shell_command(command)
def test_cat_only_allows_safe_paths(self):
decision = classify_shell_command("cat /etc/os-release")
self.assertEqual(decision.mode, "read")
self.assertFalse(decision.requires_approval)
def test_cat_rejects_unsafe_path(self):
with self.assertRaises(ValueError):
classify_shell_command("cat /etc/shadow")
class AgentCommandDispatchTest(unittest.TestCase):
def test_dispatches_kubectl_command_to_kubectl_classifier(self):
decision = classify_agent_command("kubectl get pods -A")
self.assertEqual(decision.mode, "read")
self.assertEqual(decision.verb, "get")
def test_dispatches_shell_command_to_shell_classifier(self):
decision = classify_agent_command("systemctl status kubelet")
self.assertEqual(decision.mode, "read")
self.assertEqual(decision.verb, "systemctl")
def test_dispatches_write_shell_command_to_shell_classifier(self):
decision = classify_agent_command("systemctl restart kubelet")
self.assertEqual(decision.mode, "write")
self.assertTrue(decision.requires_approval)
def test_dispatches_write_kubectl_command(self):
decision = classify_agent_command("kubectl delete pod api-0 -n prod")
self.assertEqual(decision.mode, "write")
self.assertTrue(decision.requires_approval)
if __name__ == "__main__":
unittest.main()
tests/test_diagnostic_agent.py
+112 -7
View File
@@ -23,7 +23,7 @@ class DiagnosticAgentTest(unittest.TestCase):
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):
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())
@@ -52,7 +52,7 @@ class DiagnosticAgentTest(unittest.TestCase):
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:
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
@@ -73,7 +73,7 @@ class DiagnosticAgentTest(unittest.TestCase):
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_kubectl", fake_execute):
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())
@@ -103,7 +103,7 @@ class DiagnosticAgentTest(unittest.TestCase):
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_kubectl", fake_execute):
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())
@@ -143,7 +143,7 @@ class DiagnosticAgentTest(unittest.TestCase):
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_kubectl", fake_execute):
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())
@@ -160,7 +160,7 @@ class DiagnosticAgentTest(unittest.TestCase):
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_kubectl", fake_execute):
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())
@@ -185,7 +185,7 @@ class DiagnosticAgentTest(unittest.TestCase):
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_kubectl", fake_execute):
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())
@@ -206,6 +206,45 @@ class DiagnosticAgentTest(unittest.TestCase):
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"
@@ -218,6 +257,72 @@ class DiagnosticAgentTest(unittest.TestCase):
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()