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

扩展 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
+8 -3
View File
@@ -15,10 +15,14 @@
- 根因定位、影响评估、修复命令、优先级排序、预防建议
- **AI 自主诊断(双执行模式)**:启动前可选择命令执行方式
- **手动执行命令**:Agent 只生成命令,用户复制到自己的终端执行,服务端不会运行
- **AI 执行命令**:自动执行 `get``describe``logs``top` 等只读 kubectl 命令
- **AI 执行命令**:自动执行只读诊断命令,覆盖 kubectl 与节点宿主机两个维度
- kubectl: `get``describe``logs``top` 等只读命令
- 系统诊断: `systemctl status``journalctl``crictl ps``df``free``ss``ip addr` 等只读命令
- 实时展示思考步骤、执行命令和真实输出
- 即使选择 AI 执行,`apply``delete``patch``scale``rollout``exec` 等命令仍必须人工确认
- 禁止 shell、管道、重定向和多命令拼接
- 即使选择 AI 执行,修改命令人工确认
- kubectl: `apply``delete``patch``scale``rollout``exec`
- 系统: `systemctl restart/stop/start`
- 命令白名单严格限制,禁止 shell、管道、重定向和多命令拼接
- **历史记录持久化**:使用 SQLite 保存全部诊断与处理轨迹
- 保存完整诊断结果、文本报告、AI 分析结果
- 保存 Agent 思考、命令、真实输出、审批和执行结果
@@ -90,6 +94,7 @@ k8smanager-cli/
│ ├── main.py # FastAPI 主应用
│ ├── config.py # 配置管理
│ ├── diagnosis.py # K8S 诊断引擎 (9 大检查项)
│ ├── agent_tools.py # Agent 命令工具与安全策略 (kubectl + 系统诊断白名单)
│ └── ai_agent.py # AI Agent 分析模块
├── frontend/
│ ├── src/App.vue # 仪表盘界面
+160 -7
View File
@@ -1,10 +1,19 @@
"""AI Agent 可用的 kubectl 命令工具和安全策略。"""
"""AI Agent 可用的命令工具和安全策略。
支持两类命令:
1. kubectl 命令 - 通过 classify_kubectl_command 分类
2. 系统诊断 shell 命令 - 通过 classify_shell_command 分类 (白名单)
classify_agent_command 统一分发,execute_command 自动选择执行方式。
所有命令均通过 subprocess_exec 执行 (无 shell),禁止管道、重定向和命令组合。
"""
import asyncio
import shlex
from dataclasses import dataclass
from backend.diagnosis import KUBECTL
# ── kubectl 命令分类 ──────────────────────────────────────
READ_ONLY_VERBS = {
"api-resources", "api-versions", "auth", "cluster-info", "config", "describe",
@@ -27,22 +36,27 @@ FORBIDDEN_TOKENS = {";", "|", "||", "&&", ">", ">>", "<", "<<", "`", "$"}
@dataclass(frozen=True)
class CommandDecision:
mode: str
mode: str # "read" (只读,自动执行) 或 "write" (修改,需审批)
requires_approval: bool
verb: str
verb: str # kubectl 子命令 或 shell 命令名
args: list[str]
def parse_kubectl_command(command: str) -> list[str]:
"""解析单条 kubectl 命令,不允许 shell、重定向命令组合。"""
def _parse_command(command: str) -> list[str]:
"""解析单条命令,禁止 shell 操作符、重定向命令组合。"""
if not isinstance(command, str) or not command.strip():
raise ValueError("命令不能为空")
if any(token in command for token in FORBIDDEN_TOKENS) or "\n" in command or "\r" in command:
raise ValueError("不允许 shell 操作符、重定向或多条命令")
try:
parts = shlex.split(command)
return shlex.split(command)
except ValueError as exc:
raise ValueError(f"命令格式错误: {exc}") from exc
def parse_kubectl_command(command: str) -> list[str]:
"""仅解析单条 kubectl 命令,不允许 shell、重定向或命令组合。"""
parts = _parse_command(command)
if not parts or parts[0] != "kubectl":
raise ValueError("只允许执行 kubectl 命令")
if len(parts) < 2:
@@ -91,8 +105,147 @@ def classify_kubectl_command(command: str) -> CommandDecision:
return CommandDecision("write", True, verb, args)
# ── 系统诊断 shell 命令分类 (白名单) ─────────────────────
# 只读诊断命令白名单:value 为 None 表示该命令本身只读,任意参数均可;
# 为集合时,第一个非 flag 参数必须属于该集合。
SHELL_READ_ONLY_COMMANDS: dict[str, set[str] | None] = {
"systemctl": {
"status", "is-active", "is-enabled", "is-failed", "list-units",
"list-sockets", "list-jobs", "show", "cat", "list-unit-files",
"list-dependencies",
},
"journalctl": None, # 只读,任意参数 (-u, -n, --no-pager, --since 等)
"crictl": {"ps", "inspect", "logs", "info", "version", "images", "stats"},
"docker": {"ps", "logs", "inspect", "stats", "version", "info", "images", "top"},
"df": None, "free": None, "uptime": None, "uname": None, "hostname": None,
"ip": {"addr", "address", "route", "link", "neighbor", "neigh", "rule"},
"ss": None, "netstat": None, "lsblk": None, "ls": None, "cat": None,
"head": None, "tail": None, "wc": None, "grep": None, "egrep": None, "fgrep": None,
"ps": None, "mount": None, "lsmod": None, "lscpu": None, "lsof": None,
"dmesg": None, "ping": None, "nslookup": None, "dig": None, "getent": None,
"timedatectl": None, "date": None, "stat": None, "file": None, "blkid": None,
}
# 系统修改命令白名单:需要人工审批
SHELL_WRITE_COMMANDS: dict[str, set[str] | None] = {
"systemctl": {"restart", "start", "stop", "reload", "enable", "disable", "daemon-reload"},
}
# ip 命令的写操作动词 (出现即拒绝自动执行)
_IP_WRITE_ACTIONS = {"set", "add", "del", "delete", "change", "replace", "append", "prepend"}
# cat 允许读取的安全路径
_SAFE_CAT_FILES = {"/etc/os-release", "/etc/hosts", "/etc/resolv.conf"}
_SAFE_CAT_DIRS = ("/proc/", "/sys/", "/etc/kubernetes/", "/etc/cni/", "/etc/systemd/system/")
def _is_safe_cat_path(path: str) -> bool:
if not path:
return False
if path in _SAFE_CAT_FILES:
return True
return any(path.startswith(d) for d in _SAFE_CAT_DIRS)
def _first_non_flag(args: list[str]) -> str | None:
for arg in args:
if not arg.startswith("-"):
return arg.lower()
return None
def parse_shell_command(command: str) -> list[str]:
"""解析单条 shell 诊断命令,禁止 shell 操作符和命令组合。"""
parts = _parse_command(command)
if not parts:
raise ValueError("命令不能为空")
return parts
def classify_shell_command(command: str) -> CommandDecision:
"""对白名单内的系统诊断命令进行分类。
systemctl 同时出现在只读和写白名单中:restart/stop 等走写表(需审批),
status/is-active 等走只读表(自动执行)。因此先检查写表,再检查只读表。
"""
parts = parse_shell_command(command)
binary = parts[0]
verb = _first_non_flag(parts[1:])
if binary in SHELL_WRITE_COMMANDS:
allowed = SHELL_WRITE_COMMANDS[binary]
if allowed is None or (verb is not None and verb in allowed):
return CommandDecision("write", True, binary, parts[1:])
if binary in SHELL_READ_ONLY_COMMANDS:
allowed = SHELL_READ_ONLY_COMMANDS[binary]
if allowed is not None and (verb is None or verb not in allowed):
raise ValueError(f"命令 {binary} 不支持该子命令: {verb}")
if binary == "ip" and any(arg in _IP_WRITE_ACTIONS for arg in parts[1:]):
raise ValueError("ip 命令包含写操作,请使用只读形式如 'ip addr show'")
if binary == "cat":
for arg in parts[1:]:
if not arg.startswith("-") and not _is_safe_cat_path(arg):
raise ValueError(
"cat 仅允许读取 /proc, /sys, /etc/kubernetes, /etc/cni, "
"/etc/systemd/system, /etc/os-release, /etc/hosts, /etc/resolv.conf"
)
return CommandDecision("read", False, binary, parts[1:])
raise ValueError(f"不允许执行该命令: {binary}")
def classify_agent_command(command: str) -> CommandDecision:
"""统一分类 Agent 命令:kubectl 或白名单系统诊断命令。"""
stripped = (command or "").strip()
if stripped.startswith("kubectl ") or stripped == "kubectl":
return classify_kubectl_command(stripped)
return classify_shell_command(stripped)
# ── 命令执行 ──────────────────────────────────────────────
async def execute_command(command: str, timeout: int = 30) -> dict:
"""执行单条 Agent 命令 (kubectl 或白名单系统诊断命令),无 shell。"""
decision = classify_agent_command(command)
parts = _parse_command(command)
if parts[0] == "kubectl":
executable = KUBECTL
args = decision.args
else:
executable = parts[0]
args = parts[1:]
proc = await asyncio.create_subprocess_exec(
executable,
*args,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
)
try:
stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=timeout)
return {
"command": command,
"mode": decision.mode,
"returncode": proc.returncode,
"stdout": stdout.decode(errors="replace")[:20000].strip(),
"stderr": stderr.decode(errors="replace")[:10000].strip(),
}
except asyncio.TimeoutError:
proc.kill()
await proc.communicate()
return {
"command": command,
"mode": decision.mode,
"returncode": -1,
"stdout": "",
"stderr": f"命令执行超时 ({timeout}s)",
}
async def execute_kubectl(command: str, timeout: int = 30) -> dict:
"""无 shell 执行 kubectl,并限制输出大小"""
"""执行单条 kubectl 命令 (兼容旧接口)"""
decision = classify_kubectl_command(command)
proc = await asyncio.create_subprocess_exec(
KUBECTL,
+48 -15
View File
@@ -6,7 +6,7 @@ from collections import Counter
from typing import Optional
import httpx
from backend.config import AI_API_BASE, AI_API_KEY, AI_MODEL
from backend.agent_tools import classify_kubectl_command, execute_kubectl
from backend.agent_tools import classify_agent_command, execute_command
SYSTEM_PROMPT = """你是一位资深的 Kubernetes 运维专家和 SRE 工程师。你的任务是分析 K8S 集群诊断报告,找出根本原因并给出修复建议。
@@ -25,19 +25,25 @@ _HEADERS = {
}
AGENT_SYSTEM_PROMPT = """你是 Kubernetes 自主诊断 Agent。你的首要职责是主动调用 kubectl 获取现场证据,而不是只给用户一段脚本或操作建议。
AGENT_SYSTEM_PROMPT = """你是 Kubernetes 自主诊断 Agent。你的首要职责是主动调用命令获取现场证据,而不是只给用户一段脚本或操作建议。
可用工具:
1. kubectl 命令 - 查询集群内资源状态。例如: kubectl get pods -A, kubectl describe node <name>, kubectl logs <pod> -n <ns>
2. 系统诊断 shell 命令 - 查询节点宿主机状态。例如: systemctl status kubelet, journalctl -u kubelet -n 100 --no-pager, crictl ps, df -h, free -m, ss -tlnp, ip addr
行为要求:
1. 如果证据不足,必须返回 command 动作并逐条调查;至少检查与故障资源直接相关的状态、describe、events 或 logs。
2. 禁止在 final 中仅提供请执行以下命令的脚本。只有在已经依据真实命令输出形成结论时才能返回 final。
3. 每次只提出一条 kubectl 命令,服务端会执行并把真实输出返回给你,然后你继续下一轮判断。
2. 禁止在 final 中仅提供"请执行以下命令"的脚本。只有在已经依据真实命令输出形成结论时才能返回 final。
3. 每次只提出一条命令,服务端会执行并把真实输出返回给你,然后你继续下一轮判断。
安全规则:
1. 只允许提出单条 kubectl 命令,不得使用 shell、管道、重定向或命令连接符。
2. get、describe、logs、top 等只读命令可以自动执行。
3. apply、delete、patch、scale、rollout、exec 等修改或交互命令必须暂停并请求人工批准
4. 不要猜测命令输出;必须依据实际工具结果继续判断
5. 最多执行有限轮次,优先使用 namespace 和资源名缩小范围
1. 每次只返回一条命令,不得使用 shell、管道、重定向或命令连接符 (; | && > < 等)
2. kubectl 的 get、describe、logs、top 等只读命令可以自动执行。
3. 系统诊断只读命令 (systemctl status、journalctl、crictl ps、df、free、ss、ip addr 等) 可以自动执行
4. apply、delete、patch、scale、rollout、exec 等修改或交互命令必须暂停并请求人工批准
5. systemctl restart/stop/start 等系统修改命令必须暂停并请求人工批准
6. 不要猜测命令输出;必须依据实际工具结果继续判断。
7. 最多执行有限轮次,优先使用 namespace 和资源名缩小范围。
每轮必须只返回一个 JSON 对象,不要使用 Markdown 代码块:
- 继续调查:{"action":"command","command":"kubectl ...","reason":"为什么执行"}
@@ -129,7 +135,7 @@ async def run_diagnostic_agent(
"reason": "建立集群节点状态基线,确保 AI 执行模式实际采集现场证据",
"mode": "read",
}
initial_result = await execute_kubectl(initial_command)
initial_result = await execute_command(initial_command)
yield {"type": "result", **initial_result}
messages.append({
"role": "user",
@@ -157,6 +163,7 @@ async def _run_agent_loop(
report: str,
question: str,
):
invalid_response_count = 0
for step in range(start_step, max_steps + 1):
yield {"type": "thinking", "step": step, "message": f"AI 正在规划第 {step} 轮诊断"}
try:
@@ -164,6 +171,7 @@ async def _run_agent_loop(
action = _parse_agent_action(content)
except ValueError as exc:
if "最终动作必须确认" in str(exc):
invalid_response_count = 0
messages.extend([
{"role": "assistant", "content": content},
{
@@ -177,13 +185,37 @@ async def _run_agent_loop(
])
yield {"type": "verification_required", "message": "尚未确认所有问题清零,Agent 继续执行诊断验证"}
continue
yield {"type": "error", "message": str(exc)}
return
invalid_response_count += 1
if invalid_response_count >= 3:
yield {
"type": "error",
"message": "AI 连续 3 次未返回合法 JSON 动作,已停止本次诊断以避免无效循环",
}
return
messages.extend([
{"role": "assistant", "content": content},
{
"role": "user",
"content": (
f"上一次响应无法解析:{exc}。请不要输出解释、脚本或多条命令,"
"严格只返回一个 JSON 对象:"
'{"action":"command","command":"kubectl ...","reason":"..."} '
"或在全量复检问题清零后返回 final。"
),
},
])
yield {
"type": "response_rejected",
"message": f"AI 响应格式无效,正在自动纠正并重试({invalid_response_count}/3",
"reason": str(exc),
}
continue
except Exception as exc:
yield {"type": "error", "message": str(exc)}
return
if action["action"] == "final":
invalid_response_count = 0
yield {
"type": "final", "content": action["analysis"], "model": AI_MODEL,
"verified": True, "remaining_issues": 0,
@@ -191,6 +223,7 @@ async def _run_agent_loop(
return
command = action["command"]
invalid_response_count = 0
reason = action.get("reason", "AI 需要更多集群证据")
command_counts[command] += 1
if command_counts[command] > 2:
@@ -210,11 +243,11 @@ async def _run_agent_loop(
}
continue
try:
decision = classify_kubectl_command(command)
decision = classify_agent_command(command)
except ValueError as exc:
messages.extend([
{"role": "assistant", "content": content},
{"role": "user", "content": f"命令被安全策略拒绝:{exc}。请改用单条合法 kubectl 命令。"},
{"role": "user", "content": f"命令被安全策略拒绝:{exc}。请改用单条合法 kubectl 或系统诊断命令。"},
])
yield {"type": "command_rejected", "command": command, "reason": str(exc)}
continue
@@ -252,7 +285,7 @@ async def _run_agent_loop(
}
return
result = await execute_kubectl(command)
result = await execute_command(command)
yield {"type": "result", **result}
messages.extend([
{"role": "assistant", "content": content},
+2 -2
View File
@@ -24,7 +24,7 @@ from backend.ai_agent import (
analyze_with_ai_stream, chat_with_ai_stream,
run_diagnostic_agent, resume_diagnostic_agent, take_pending_approval,
)
from backend.agent_tools import execute_kubectl
from backend.agent_tools import execute_command
from backend.history_store import HistoryStore
app = FastAPI(title="K8S 智能诊断平台", version="1.0.0")
@@ -370,7 +370,7 @@ async def execute_approved_agent_command(req: AgentApprovalRequest):
history_store.append_event(pending["run_id"], "approval_rejected", result)
history_store.finish_run(pending["run_id"], "rejected", result["message"])
return result
result = await execute_kubectl(pending["command"])
result = await execute_command(pending["command"])
approved_event = {"approved": True, "reason": pending["reason"], **result}
async def event_generator():
+3 -2
View File
@@ -280,8 +280,8 @@
</template>
<el-alert
:title="activeAgentMode === 'ai'
? 'AI 执行模式:Agent 自动执行只读 kubectl 命令,修改或交互命令仍需人工批准。'
: '手动执行模式:Agent 只生成安全的 kubectl 命令,不会在服务器上执行,请复制命令到终端手动运行。'"
? 'AI 执行模式:Agent 自动执行只读 kubectl 和系统诊断命令 (systemctl status、journalctl、crictl ps、df 等),修改或交互命令仍需人工批准。'
: '手动执行模式:Agent 只生成安全的诊断命令,不会在服务器上执行,请复制命令到终端手动运行。'"
type="info"
:closable="false"
show-icon
@@ -690,6 +690,7 @@ function agentEventTitle(event) {
thinking: 'Agent 思考', command: '准备执行命令', result: '命令执行结果',
manual_command: '请手动执行命令',
approval_required: '发现修改命令,等待审批', command_rejected: '命令已被安全策略拒绝',
verification_required: '需要继续复检', response_rejected: 'AI 响应格式已自动纠正',
final: '自主诊断结论', error: 'Agent 执行异常',
}
return titles[event.type] || event.type
+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]
+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()
+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()