From 32bea0d4eaa291505a4223255b2193d12c978f82 Mon Sep 17 00:00:00 2001 From: cnbugs Date: Sat, 25 Jul 2026 13:38:48 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20AI=20=E8=AF=8A=E6=96=AD=20Agent=20?= =?UTF-8?q?=E6=94=AF=E6=8C=81=E6=89=A7=E8=A1=8C=E7=B3=BB=E7=BB=9F=E8=AF=8A?= =?UTF-8?q?=E6=96=AD=20shell=20=E5=91=BD=E4=BB=A4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 扩展 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 等) 直接拒绝。 --- README.md | 11 ++- backend/agent_tools.py | 167 +++++++++++++++++++++++++++++++-- backend/ai_agent.py | 63 ++++++++++--- backend/main.py | 4 +- frontend/src/App.vue | 5 +- tests/test_agent_api.py | 4 +- tests/test_agent_tools.py | 96 ++++++++++++++++++- tests/test_diagnostic_agent.py | 119 +++++++++++++++++++++-- 8 files changed, 430 insertions(+), 39 deletions(-) diff --git a/README.md b/README.md index 260ebdf..4d9ff15 100644 --- a/README.md +++ b/README.md @@ -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 # 仪表盘界面 diff --git a/backend/agent_tools.py b/backend/agent_tools.py index 0bf995f..04a0a79 100644 --- a/backend/agent_tools.py +++ b/backend/agent_tools.py @@ -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, diff --git a/backend/ai_agent.py b/backend/ai_agent.py index 02e8843..3404ce3 100644 --- a/backend/ai_agent.py +++ b/backend/ai_agent.py @@ -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 , kubectl logs -n +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}, diff --git a/backend/main.py b/backend/main.py index a7c423e..e4f3a34 100644 --- a/backend/main.py +++ b/backend/main.py @@ -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(): diff --git a/frontend/src/App.vue b/frontend/src/App.vue index 9e78eeb..4911ccd 100644 --- a/frontend/src/App.vue +++ b/frontend/src/App.vue @@ -280,8 +280,8 @@ /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() diff --git a/tests/test_diagnostic_agent.py b/tests/test_diagnostic_agent.py index 2b0fa17..8d51065 100644 --- a/tests/test_diagnostic_agent.py +++ b/tests/test_diagnostic_agent.py @@ -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()