aidev-bkplugin 2.2.0.post2__tar.gz → 2.2.0.post4__tar.gz

{aidev_bkplugin-2.2.0.post2 → aidev_bkplugin-2.2.0.post4}/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.1
 Name: aidev-bkplugin
-Version: 2.2.0.post2
+Version: 2.2.0.post4
 Summary: AiDev Agent Bkplugin builtin functions
 Requires-Python: <3.12,>=3.11
-Requires-Dist: aidev-agent>=2.2.0post2
+Requires-Dist: aidev-agent>=2.2.0post4
 Requires-Dist: blueapps<5.0,>=4.15.1
 Requires-Dist: django<5.0,>=4.2.28
 Description-Content-Type: text/markdown

{aidev_bkplugin-2.2.0.post2 → aidev_bkplugin-2.2.0.post4}/aidev_bkplugin/services/agent_bkplugin.py

@@ -28,6 +28,7 @@ from aidev_agent.pydantic_models import ExecuteKwargs
 from aidev_agent.services.agent import AgentInstanceFactory
 from aidev_agent.services.common_agent import common_agent_factory
 from aidev_agent.services.event_handlers.agui_writer import AGUISessionWriter
+from aidev_agent.services.messages_handler import RetryableHeartbeatTimeoutError
 from aidev_agent.utils.local import request_local
 from django.contrib.auth import get_user_model
 
@@ -173,6 +174,13 @@ class BkpluginAgentRunner(ABC):
         """同步执行 Agent，返回最终 AI 回复字符串；异常时写回失败状态 (D-01)。"""
         try:
             return self._do_execute()
+        except RetryableHeartbeatTimeoutError:
+            logger.warning(
+                "[Bkplugin] skip session terminal update for retryable consumer heartbeat timeout session_code=%s",
+                self.session_code or "",
+                exc_info=True,
+            )
+            raise
         except Exception as e:
             logger.exception("[Bkplugin] execute error: %s", e)
             session_code = self.session_code or ""
@@ -212,7 +220,7 @@ class BkpluginAgentRunner(ABC):
         *,
         chat_context: list[dict] | None = None,
     ) -> None:
-        """Celery worker 入口：捕获异常并写 session 失败状态。"""
+        """Celery worker 入口：业务异常写失败；可重试消费异常保持 session 原状态。"""
         logger.info(
             "[Bkplugin] run_worker enter session_code=%s turn_id=%s thread=%s",
             session_code,
@@ -221,6 +229,14 @@ class BkpluginAgentRunner(ABC):
         )
         try:
             self.invoke_agent(session_code, execute_payload, chat_context=chat_context or [])
+        except RetryableHeartbeatTimeoutError:
+            logger.warning(
+                "[Bkplugin] skip session terminal update for retryable consumer heartbeat timeout "
+                "session_code=%s turn_id=%s",
+                session_code,
+                execute_payload.get("turn_id") or "",
+                exc_info=True,
+            )
         except Exception as e:
             logger.exception("[Bkplugin] worker error session_code=%s", session_code)
             manager = SessionManager(self.username or "")

{aidev_bkplugin-2.2.0.post2 → aidev_bkplugin-2.2.0.post4}/aidev_bkplugin/services/agent_execution.py

@@ -89,7 +89,7 @@ class AgentExecutor:
         *,
         turn_id: str = "",
     ):
-        """执行 agent 直至结束；AG-UI 流式负责事件落库并收尾会话状态。"""
+        """执行 agent 直至结束；后台消费者完整 drain 后统一收尾会话状态。"""
         # 后台 drain（for _ in out: pass，无 SSE 下游）：标记为 background_only，
         # 使消费者读到 EOD 时不立即清理队列，保留 DLQ 历史供前端在清理窗口内接管续流。
         execute_kwargs.background_only = True
@@ -99,24 +99,18 @@ class AgentExecutor:
                 handler.turn_id = turn_id
             if hasattr(handler, "set_streaming_started"):
                 handler.set_streaming_started()
-        try:
-            out = cls(session_manager).execute_with_save(
-                agent_instance,
-                execute_kwargs,
-                session_code,
-                turn_id=turn_id,
-            )
-            if execute_kwargs.stream:
-                for _ in out:
-                    pass
-            return out
-        finally:
-            if (
-                execute_kwargs.stream
-                and isinstance(handler, BaseSessionWriter)
-                and hasattr(handler, "set_streaming_finished")
-            ):
+        out = cls(session_manager).execute_with_save(
+            agent_instance,
+            execute_kwargs,
+            session_code,
+            turn_id=turn_id,
+        )
+        if execute_kwargs.stream:
+            for _ in out:
+                pass
+            if isinstance(handler, BaseSessionWriter) and hasattr(handler, "set_streaming_finished"):
                 handler.set_streaming_finished()
+        return out
 
     def wrap_generator(self, generator, session_code: str, *, turn_id: str = ""):
         """SSE 数据格式约定：

{aidev_bkplugin-2.2.0.post2 → aidev_bkplugin-2.2.0.post4}/aidev_bkplugin/views/chat.py

@@ -397,9 +397,11 @@ class ChatCompletionViewSet(PluginViewSet):
         _preempted = False
         _client_disconnected = False
         _cancelled = False
+        _stream_completed = False
         try:
             for chunk in generator:
                 yield chunk
+            _stream_completed = True
         except GeneratorExit:
             _client_disconnected = True
             logger.info(f"[WRAP_STATUS] Client disconnected (GeneratorExit): session_code={session_code}")
@@ -421,7 +423,7 @@ class ChatCompletionViewSet(PluginViewSet):
                 if GeneratorStreamingHelper.is_cancelled(session_code):
                     _cancelled = True
                     logger.info(f"[WRAP_STATUS] Generator结束后检测到取消标志: session_code={session_code}")
-            if not _preempted and not _client_disconnected and not _cancelled:
+            if _stream_completed and not _preempted and not _client_disconnected and not _cancelled:
                 logger.info(f"[WRAP_STATUS] 流正常结束, 调用 set_streaming_finished: session_code={session_code}")
                 event_handler.set_streaming_finished()
 

{aidev_bkplugin-2.2.0.post2 → aidev_bkplugin-2.2.0.post4}/pyproject.toml

@@ -1,11 +1,11 @@
 [project]
 name = "aidev-bkplugin"
-version = "2.2.0post2"
+version = "2.2.0post4"
 description = "AiDev Agent Bkplugin builtin functions"
 readme = "readme.md"
 requires-python = "<3.12,>=3.11"
 dependencies = [
-    "aidev-agent>=2.2.0post2",
+    "aidev-agent>=2.2.0post4",
     "blueapps>=4.15.1, <5.0",
     "django>=4.2.28,<5.0",
 ]

aidev_bkplugin-2.2.0.post4/tests/services/test_agent_completion.py

@@ -0,0 +1,58 @@
+# -*- coding: utf-8 -*-
+"""后台 Agent 流式终态写入回归测试。"""
+
+import sys
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+if "pkg_resources" not in sys.modules:
+    sys.modules["pkg_resources"] = MagicMock()
+
+sys.modules.setdefault("aidev_bkplugin.models", MagicMock())
+_bk_plugin_framework = MagicMock()
+_bk_plugin_framework.kit.decorators.inject_user_token = lambda func: func
+sys.modules.setdefault("bk_plugin_framework", _bk_plugin_framework)
+sys.modules.setdefault("bk_plugin_framework.kit", _bk_plugin_framework.kit)
+sys.modules.setdefault("bk_plugin_framework.kit.decorators", _bk_plugin_framework.kit.decorators)
+
+
+class TestAgentExecutorCompletion:
+    def test_sets_finished_after_stream_is_drained(self):
+        from aidev_agent.services.event_handlers.base import BaseSessionWriter
+        from aidev_bkplugin.services.agent_execution import AgentExecutor
+
+        lifecycle = []
+
+        def stream():
+            yield "chunk"
+            lifecycle.append("drained")
+
+        handler = MagicMock(spec=BaseSessionWriter)
+        handler.set_streaming_finished.side_effect = lambda: lifecycle.append("finished")
+        agent = MagicMock(event_handler=handler)
+        execute_kwargs = MagicMock(stream=True)
+        with patch.object(AgentExecutor, "execute_with_save", return_value=stream()):
+            AgentExecutor.run_agent_to_completion(agent, execute_kwargs, "session-abc", MagicMock())
+
+        assert execute_kwargs.background_only is True
+        assert lifecycle == ["drained", "finished"]
+
+    def test_does_not_set_finished_when_stream_drain_fails(self):
+        from aidev_agent.services.event_handlers.base import BaseSessionWriter
+        from aidev_bkplugin.services.agent_execution import AgentExecutor
+
+        def failing_stream():
+            yield "chunk"
+            raise RuntimeError("producer heartbeat timeout")
+
+        handler = MagicMock(spec=BaseSessionWriter)
+        agent = MagicMock(event_handler=handler)
+        execute_kwargs = MagicMock(stream=True)
+        with (
+            patch.object(AgentExecutor, "execute_with_save", return_value=failing_stream()),
+            pytest.raises(RuntimeError, match="heartbeat timeout"),
+        ):
+            AgentExecutor.run_agent_to_completion(agent, execute_kwargs, "session-abc", MagicMock())
+
+        handler.set_streaming_finished.assert_not_called()

{aidev_bkplugin-2.2.0.post2 → aidev_bkplugin-2.2.0.post4}/tests/services/test_bkplugin.py

@@ -23,6 +23,7 @@ sys.modules.setdefault("bk_plugin_framework.kit", _bk_plugin_framework.kit)
 sys.modules.setdefault("bk_plugin_framework.kit.decorators", _bk_plugin_framework.kit.decorators)
 
 from aidev_agent.enums import AgentType, ChatContentStatus, PromptRole, SessionsStatus  # noqa: E402
+from aidev_agent.services.messages_handler import RetryableHeartbeatTimeoutError  # noqa: E402
 from aidev_bkplugin.enums import PluginPollTaskState  # noqa: E402
 from aidev_bkplugin.services.agent_session import SessionManager  # noqa: E402
 
@@ -272,14 +273,10 @@ class TestBkpluginExecution:
         from aidev_bkplugin.services.agent_bkplugin import BkpluginChat
 
         agent = BkpluginChat(chat_history=[], execute_kwargs={}, username="alice")
-        monkeypatch.setattr(
-            agent, "invoke_agent", MagicMock(side_effect=RuntimeError("生产者心跳超时"))
-        )
+        monkeypatch.setattr(agent, "invoke_agent", MagicMock(side_effect=RuntimeError("生产者心跳超时")))
         mock_sm = MagicMock()
         mock_sm.retrieve_session.return_value = {"status": SessionsStatus.FINISHED.value}
-        monkeypatch.setattr(
-            "aidev_bkplugin.services.agent_bkplugin.SessionManager", lambda username: mock_sm
-        )
+        monkeypatch.setattr("aidev_bkplugin.services.agent_bkplugin.SessionManager", lambda username: mock_sm)
 
         agent.run_worker("sess-1", {"turn_id": "t1"})
 
@@ -291,16 +288,46 @@ class TestBkpluginExecution:
         from aidev_bkplugin.services.agent_bkplugin import BkpluginChat
 
         agent = BkpluginChat(chat_history=[], execute_kwargs={}, username="alice")
-        monkeypatch.setattr(
-            agent, "invoke_agent", MagicMock(side_effect=RuntimeError("生产者心跳超时"))
-        )
+        monkeypatch.setattr(agent, "invoke_agent", MagicMock(side_effect=RuntimeError("producer crashed")))
         mock_sm = MagicMock()
         mock_sm.retrieve_session.return_value = {"status": SessionsStatus.RUNNING.value}
+        monkeypatch.setattr("aidev_bkplugin.services.agent_bkplugin.SessionManager", lambda username: mock_sm)
+
+        agent.run_worker("sess-1", {"turn_id": "t1"})
+
+        mock_sm.save_stream_failure.assert_called_once_with("sess-1", ANY, turn_id="t1")
+
+    def test_run_worker_skips_session_update_for_retryable_heartbeat_timeout(self, monkeypatch):
+        """消费者心跳超时可重试，不应更新 session 终态。"""
+        from aidev_bkplugin.services.agent_bkplugin import BkpluginChat
+
+        agent = BkpluginChat(chat_history=[], execute_kwargs={}, username="alice")
         monkeypatch.setattr(
-            "aidev_bkplugin.services.agent_bkplugin.SessionManager", lambda username: mock_sm
+            agent,
+            "invoke_agent",
+            MagicMock(side_effect=RetryableHeartbeatTimeoutError("生产者心跳超时")),
         )
+        manager_factory = MagicMock()
+        monkeypatch.setattr("aidev_bkplugin.services.agent_bkplugin.SessionManager", manager_factory)
 
         agent.run_worker("sess-1", {"turn_id": "t1"})
 
-        mock_sm.save_stream_failure.assert_called_once_with("sess-1", ANY, turn_id="t1")
+        manager_factory.assert_not_called()
+
+    def test_execute_skips_session_update_for_retryable_heartbeat_timeout(self, monkeypatch):
+        """同步消费遇到可重试心跳超时，也不应覆盖 session 状态。"""
+        from aidev_bkplugin.services.agent_bkplugin import BkpluginChat
+
+        agent = BkpluginChat(chat_history=[], execute_kwargs={"session_code": "sess-1"}, username="alice")
+        monkeypatch.setattr(
+            agent,
+            "_do_execute",
+            MagicMock(side_effect=RetryableHeartbeatTimeoutError("生产者心跳超时")),
+        )
+        manager_factory = MagicMock()
+        monkeypatch.setattr("aidev_bkplugin.services.agent_bkplugin.SessionManager", manager_factory)
+
+        with pytest.raises(RetryableHeartbeatTimeoutError):
+            agent.execute()
 
+        manager_factory.assert_not_called()

aidev_bkplugin-2.2.0.post4/tests/views/test_chat_stream_status.py

@@ -0,0 +1,40 @@
+# -*- coding: utf-8 -*-
+"""Chat SSE 异常时的会话终态回归测试。"""
+
+from unittest.mock import MagicMock
+
+import pytest
+
+
+@pytest.mark.parametrize("error", [RuntimeError("生产者心跳超时"), SystemExit("worker abort")])
+def test_unexpected_stream_error_does_not_mark_session_finished(monkeypatch, error):
+    from aidev_agent.services.messages_handler.streaming_helper import GeneratorStreamingHelper
+    from aidev_bkplugin.views.chat import ChatCompletionViewSet
+
+    def failing_stream():
+        yield "chunk"
+        raise error
+
+    monkeypatch.setattr(GeneratorStreamingHelper, "is_cancelled", lambda _session_code: False)
+    event_handler = MagicMock()
+    stream = ChatCompletionViewSet()._wrap_streaming_with_status(
+        failing_stream(), event_handler, session_code="session-1"
+    )
+
+    assert next(stream) == "chunk"
+    with pytest.raises(type(error)):
+        next(stream)
+    event_handler.set_streaming_finished.assert_not_called()
+
+
+def test_normal_stream_marks_session_finished(monkeypatch):
+    from aidev_agent.services.messages_handler.streaming_helper import GeneratorStreamingHelper
+    from aidev_bkplugin.views.chat import ChatCompletionViewSet
+
+    monkeypatch.setattr(GeneratorStreamingHelper, "is_cancelled", lambda _session_code: False)
+    event_handler = MagicMock()
+
+    assert list(
+        ChatCompletionViewSet()._wrap_streaming_with_status(iter(["chunk"]), event_handler, session_code="session-1")
+    ) == ["chunk"]
+    event_handler.set_streaming_finished.assert_called_once_with()