aidev-bkplugin 2.2.0__tar.gz → 2.2.0.post3__tar.gz

{aidev_bkplugin-2.2.0 → aidev_bkplugin-2.2.0.post3}/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.1
 Name: aidev-bkplugin
-Version: 2.2.0
+Version: 2.2.0.post3
 Summary: AiDev Agent Bkplugin builtin functions
 Requires-Python: <3.12,>=3.11
-Requires-Dist: aidev-agent>=2.2.0
+Requires-Dist: aidev-agent>=2.2.0post3
 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 → aidev_bkplugin-2.2.0.post3}/aidev_bkplugin/packages/checkpoint/bk_django_saver.py

@@ -21,10 +21,11 @@ from __future__ import annotations
 import json
 import random
 import threading
+import time
 from collections.abc import AsyncIterator, Iterator, Sequence
 from typing import Any, Optional, Tuple, cast
 
-from django.db import connections, router, transaction
+from django.db import OperationalError, close_old_connections, connections, router, transaction
 from langchain_core.runnables import RunnableConfig
 from langgraph.checkpoint.base import (
     WRITES_IDX_MAP,
@@ -40,6 +41,10 @@ from langgraph.checkpoint.base import (
 from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
 from langgraph.checkpoint.serde.types import ChannelProtocol
 
+RETRYABLE_DATABASE_ERROR_CODES = {1205, 1213}
+CHECKPOINT_WRITE_MAX_RETRIES = 3
+CHECKPOINT_WRITE_RETRY_DELAY_SECONDS = 0.05
+
 
 def bulk_upsert(model, objs, update_fields, unique_fields):
     """
@@ -605,17 +610,26 @@ class BKDjangoSaver(BaseCheckpointSaver[str]):
         serialized_metadata = json.loads(json.dumps(raw_metadata, default=str).replace("\\u0000", ""))
         # 使用Django的update_or_create方法
         with self.lock:
-            self.checkpoint_model.objects.update_or_create(
-                thread_id=thread_id,
-                checkpoint_ns=checkpoint_ns,
-                checkpoint_id=checkpoint["id"],
-                defaults={
-                    "parent_checkpoint_id": config["configurable"].get("checkpoint_id"),
-                    "type": type_,
-                    "checkpoint": serialized_checkpoint,
-                    "metadata": serialized_metadata,
-                },
-            )
+            for attempt in range(CHECKPOINT_WRITE_MAX_RETRIES):
+                try:
+                    self.checkpoint_model.objects.update_or_create(
+                        thread_id=thread_id,
+                        checkpoint_ns=checkpoint_ns,
+                        checkpoint_id=checkpoint["id"],
+                        defaults={
+                            "parent_checkpoint_id": config["configurable"].get("checkpoint_id"),
+                            "type": type_,
+                            "checkpoint": serialized_checkpoint,
+                            "metadata": serialized_metadata,
+                        },
+                    )
+                    break
+                except OperationalError as exc:
+                    error_code = exc.args[0] if exc.args else None
+                    if error_code not in RETRYABLE_DATABASE_ERROR_CODES or attempt == CHECKPOINT_WRITE_MAX_RETRIES - 1:
+                        raise
+                    close_old_connections()
+                    time.sleep(CHECKPOINT_WRITE_RETRY_DELAY_SECONDS * (2**attempt))
 
         return {
             "configurable": {

{aidev_bkplugin-2.2.0 → aidev_bkplugin-2.2.0.post3}/aidev_bkplugin/services/agent_bkplugin.py

@@ -16,12 +16,13 @@
 from __future__ import annotations
 
 import logging
+import threading
 import uuid
 from abc import ABC, abstractmethod
 from typing import Any, ClassVar
 
 from aidev_agent.config import settings as agent_settings
-from aidev_agent.enums import AgentBuildType, AgentType, ChannelType, PromptRole
+from aidev_agent.enums import AgentBuildType, AgentType, ChannelType, PromptRole, SessionsStatus
 from aidev_agent.packages.resource_manager.agent import AgentResourceManager
 from aidev_agent.pydantic_models import ExecuteKwargs
 from aidev_agent.services.agent import AgentInstanceFactory
@@ -212,11 +213,42 @@ class BkpluginAgentRunner(ABC):
         chat_context: list[dict] | None = None,
     ) -> None:
         """Celery worker 入口：捕获异常并写 session 失败状态。"""
+        logger.info(
+            "[Bkplugin] run_worker enter session_code=%s turn_id=%s thread=%s",
+            session_code,
+            execute_payload.get("turn_id") or "",
+            threading.current_thread().name,
+        )
         try:
             self.invoke_agent(session_code, execute_payload, chat_context=chat_context or [])
         except Exception as e:
             logger.exception("[Bkplugin] worker error session_code=%s", session_code)
-            SessionManager(self.username or "").save_stream_failure(
+            manager = SessionManager(self.username or "")
+            # 幂等保护：心跳超时误报时 producer 实际已完成（session=FINISHED），
+            # 不应再覆盖为 FAILED。仅当 session 仍处于非终态时才写失败，
+            # 避免覆盖已确定的终态（FINISHED/CANCELLED/FAILED）。
+            try:
+                current_status = str(manager.retrieve_session(session_code).get("status") or "")
+            except Exception:
+                logger.exception(
+                    "[Bkplugin] retrieve_session failed before save_stream_failure session_code=%s",
+                    session_code,
+                )
+                current_status = ""
+            if current_status in (
+                SessionsStatus.FINISHED.value,
+                SessionsStatus.CANCELLED.value,
+                SessionsStatus.FAILED.value,
+            ):
+                logger.warning(
+                    "[Bkplugin] skip save_stream_failure: session already %s "
+                    "(likely heartbeat-timeout false positive) session_code=%s exc=%r",
+                    current_status,
+                    session_code,
+                    e,
+                )
+                return
+            manager.save_stream_failure(
                 session_code,
                 f"Agent 执行异常: {e}",
                 turn_id=execute_payload.get("turn_id") or "",

{aidev_bkplugin-2.2.0 → aidev_bkplugin-2.2.0.post3}/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 → aidev_bkplugin-2.2.0.post3}/aidev_bkplugin/services/agent_session.py

@@ -238,24 +238,39 @@ class SessionManager:
         if status in (SessionsStatus.PENDING.value, SessionsStatus.RUNNING.value):
             return PluginPollTaskState.RUNNING, ""
         if status in (SessionsStatus.FAILED.value, SessionsStatus.CANCELLED.value):
-            return PluginPollTaskState.FAILED, "Agent 执行失败"
+            # 取回 save_stream_failure 写入的实际异常内容，避免 UI 只显示通用 "Agent 执行失败"
+            return PluginPollTaskState.FAILED, self._last_output(
+                self.list_session_contents(session_code),
+                turn_id=turn_id,
+                statuses=(ChatContentStatus.ERROR.value,),
+            ) or "Agent 执行失败"
         if status == SessionsStatus.FINISHED.value:
-            return PluginPollTaskState.SUCCESS, self._last_assistant_output(
+            return PluginPollTaskState.SUCCESS, self._last_output(
                 self.list_session_contents(session_code),
                 turn_id=turn_id,
+                statuses=(ChatContentStatus.COMPLETE.value, ChatContentStatus.SUCCESS.value),
             )
         return PluginPollTaskState.RUNNING, ""
 
     @staticmethod
-    def _last_assistant_output(items: list[dict], *, turn_id: str = "") -> str:
-        """取最后一条有内容的 assistant/ai；有 ``turn_id`` 时仅取同轮消息。"""
+    def _last_output(
+        items: list[dict],
+        *,
+        turn_id: str = "",
+        statuses: tuple[str, ...],
+    ) -> str:
+        """取最后一条指定 status 的 assistant/ai 内容；有 ``turn_id`` 时仅取同轮消息。
+
+        FINISHED 取 COMPLETE/SUCCESS 的正常回复；FAILED/CANCELLED 取 ERROR 的实际异常
+        （save_stream_failure 写入），避免上层只看到通用 "Agent 执行失败"。
+        """
         assistant_roles = (PromptRole.ASSISTANT.value, PromptRole.AI.value)
         for item in reversed(items):
             if item.get("role") not in assistant_roles:
                 continue
             if turn_id and (item.get("property") or {}).get("turn_id") != turn_id:
                 continue
-            if item.get("status") not in (ChatContentStatus.COMPLETE.value, ChatContentStatus.SUCCESS.value):
+            if item.get("status") not in statuses:
                 continue
             text = str(item.get("content") or "").strip()
             if text:

{aidev_bkplugin-2.2.0 → aidev_bkplugin-2.2.0.post3}/pyproject.toml

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

aidev_bkplugin-2.2.0.post3/tests/packages/test_bk_django_saver.py

@@ -0,0 +1,69 @@
+# -*- coding: utf-8 -*-
+
+import threading
+
+import pytest
+from aidev_bkplugin.packages.checkpoint.bk_django_saver import BKDjangoSaver
+from django.db import OperationalError
+
+
+@pytest.fixture
+def saver(mocker):
+    instance = object.__new__(BKDjangoSaver)
+    instance.lock = threading.Lock()
+    instance.checkpoint_model = mocker.Mock()
+    instance.serde = mocker.Mock()
+    instance.serde.dumps_typed.return_value = ("json", b"checkpoint")
+    return instance
+
+
+@pytest.fixture
+def checkpoint_args():
+    return (
+        {"configurable": {"thread_id": "thread-id"}},
+        {"id": "checkpoint-id"},
+        {},
+        {},
+    )
+
+
+@pytest.mark.parametrize("error_code", [1205, 1213])
+def test_put_retries_transient_database_lock_error(mocker, saver, checkpoint_args, error_code):
+    saver.checkpoint_model.objects.update_or_create.side_effect = [
+        OperationalError(error_code, "retryable"),
+        (mocker.Mock(), True),
+    ]
+    close_old_connections = mocker.patch("aidev_bkplugin.packages.checkpoint.bk_django_saver.close_old_connections")
+    sleep = mocker.patch("aidev_bkplugin.packages.checkpoint.bk_django_saver.time.sleep")
+
+    saver.put(*checkpoint_args)
+
+    assert saver.checkpoint_model.objects.update_or_create.call_count == 2
+    close_old_connections.assert_called_once_with()
+    sleep.assert_called_once_with(0.05)
+
+
+def test_put_does_not_retry_other_database_error(mocker, saver, checkpoint_args):
+    error = OperationalError(2006, "server has gone away")
+    saver.checkpoint_model.objects.update_or_create.side_effect = error
+    sleep = mocker.patch("aidev_bkplugin.packages.checkpoint.bk_django_saver.time.sleep")
+
+    with pytest.raises(OperationalError) as exc_info:
+        saver.put(*checkpoint_args)
+
+    assert exc_info.value is error
+    assert saver.checkpoint_model.objects.update_or_create.call_count == 1
+    sleep.assert_not_called()
+
+
+def test_put_raises_after_database_lock_retries_exhausted(mocker, saver, checkpoint_args):
+    error = OperationalError(1213, "deadlock")
+    saver.checkpoint_model.objects.update_or_create.side_effect = error
+    sleep = mocker.patch("aidev_bkplugin.packages.checkpoint.bk_django_saver.time.sleep")
+
+    with pytest.raises(OperationalError) as exc_info:
+        saver.put(*checkpoint_args)
+
+    assert exc_info.value is error
+    assert saver.checkpoint_model.objects.update_or_create.call_count == 3
+    assert [call.args[0] for call in sleep.call_args_list] == [0.05, 0.1]

aidev_bkplugin-2.2.0.post3/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 → aidev_bkplugin-2.2.0.post3}/tests/services/test_bkplugin.py

@@ -266,3 +266,41 @@ class TestBkpluginExecution:
         assert runner.execute_kwargs["session_code"] == "thread-x"
         assert runner.plugin_context == [{"k": "v"}]
         mock_resolve.assert_called_once_with("alice")
+
+    def test_run_worker_skips_save_stream_failure_when_session_finished(self, monkeypatch):
+        """心跳超时误报：producer 实际已完成（session=FINISHED），不应覆盖为 FAILED。"""
+        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("生产者心跳超时"))
+        )
+        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
+        )
+
+        agent.run_worker("sess-1", {"turn_id": "t1"})
+
+        mock_sm.retrieve_session.assert_called_once_with("sess-1")
+        mock_sm.save_stream_failure.assert_not_called()
+
+    def test_run_worker_calls_save_stream_failure_when_session_running(self, monkeypatch):
+        """producer 真崩溃：session 仍 RUNNING 时应写失败。"""
+        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("生产者心跳超时"))
+        )
+        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")
+