From 47abfef23877300e0ae6ac2501595a803baae196 Mon Sep 17 00:00:00 2001
From: cnbugs <cnbugs@cnbugs.com>
Date: Sat, 25 Jul 2026 12:50:38 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E5=85=81=E8=AE=B8=20Agent=20=E6=8C=81?=
 =?UTF-8?q?=E7=BB=AD=E8=87=AA=E4=B8=BB=E8=AF=8A=E6=96=AD?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 backend/ai_agent.py            | 31 +++++++++++++++++++---
 tests/test_diagnostic_agent.py | 47 ++++++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+), 4 deletions(-)

diff --git a/backend/ai_agent.py b/backend/ai_agent.py
index dee5add..c8a42ef 100644
--- a/backend/ai_agent.py
+++ b/backend/ai_agent.py
@@ -2,6 +2,7 @@
 import json
 import re
 import uuid
+from collections import Counter
 import httpx
 from backend.config import AI_API_BASE, AI_API_KEY, AI_MODEL
 from backend.agent_tools import classify_kubectl_command, execute_kubectl
@@ -96,7 +97,7 @@ async def _complete_chat(messages: list[dict]) -> str:
 async def run_diagnostic_agent(
     report: str,
     question: str = "",
-    max_steps: int = 6,
+    max_steps: int = 30,
     execution_mode: str = "ai",
 ):
     """自主诊断循环：AI 模式执行只读命令，手动模式等待用户执行。"""
@@ -112,8 +113,11 @@ async def run_diagnostic_agent(
         {"role": "user", "content": user_content},
     ]
 
+    command_counts = Counter()
+
     if execution_mode == "ai":
         initial_command = "kubectl get nodes -o wide"
+        command_counts[initial_command] += 1
         yield {
             "type": "command", "step": 0, "command": initial_command,
             "reason": "建立集群节点状态基线，确保 AI 执行模式实际采集现场证据",
@@ -141,6 +145,23 @@ async def run_diagnostic_agent(
 
         command = action["command"]
         reason = action.get("reason", "AI 需要更多集群证据")
+        command_counts[command] += 1
+        if command_counts[command] > 2:
+            messages.extend([
+                {"role": "assistant", "content": content},
+                {
+                    "role": "user",
+                    "content": (
+                        f"命令 {command!r} 已重复执行两次，不能再次执行。"
+                        "请基于已有输出选择不同的只读诊断命令，或在证据充分时返回 final。"
+                    ),
+                },
+            ])
+            yield {
+                "type": "command_rejected", "command": command,
+                "reason": "相同命令已执行两次，请更换诊断方向",
+            }
+            continue
         try:
             decision = classify_kubectl_command(command)
         except ValueError as exc:
@@ -185,9 +206,11 @@ async def run_diagnostic_agent(
         ])
 
     yield {
-        "type": "final",
-        "content": f"已达到自主诊断最大轮次（{max_steps}），请根据上方命令输出继续人工排查。",
-        "model": AI_MODEL,
+        "type": "error",
+        "message": (
+            f"自主诊断已达到安全上限（{max_steps} 轮），但 Agent 仍未形成有效结论。"
+            "系统已停止继续执行以避免无限循环，请检查是否重复执行相同命令或 AI 服务响应异常。"
+        ),
     }
 
 
diff --git a/tests/test_diagnostic_agent.py b/tests/test_diagnostic_agent.py
index ea6b622..8a73862 100644
--- a/tests/test_diagnostic_agent.py
+++ b/tests/test_diagnostic_agent.py
@@ -121,6 +121,53 @@ class DiagnosticAgentTest(unittest.TestCase):
         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":"已完成根因分析"}',
+        ])
+        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_kubectl", 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_kubectl", 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_stops_when_model_returns_invalid_action(self):
         async def fake_completion(messages):
             return "not-json"