fix(db): 使用asyncio.shield确保数据库连接正确关闭

1. 简化对话状态重置逻辑并移除冗余代码
2. 重构聊天路由的数据库会话管理

server/routers/chat_router.py

@@ -549,112 +549,124 @@ async def chat_agent(
             thread_id = str(uuid.uuid4())
             logger.warning(f"No thread_id provided, generated new thread_id: {thread_id}")
 
-        # Initialize conversation manager
-        conv_manager = ConversationManager(db)
-
-        # Save user message
         try:
-            await conv_manager.add_message_by_thread_id(
+            async with db_manager.get_async_session_context() as db:
-                thread_id=thread_id,
+                # Initialize conversation manager
-                role="user",
+                conv_manager = ConversationManager(db)
-                content=query,
-                message_type=message_type,
-                image_content=image_content,
-                extra_metadata={"raw_message": human_message.model_dump()},
-            )
-        except Exception as e:
-            logger.error(f"Error saving user message: {e}")
 
-        try:
+                # Save user message
-            assert thread_id, "thread_id is required"
+                try:
-            attachments = await conv_manager.get_attachments_by_thread_id(thread_id)
+                    await conv_manager.add_message_by_thread_id(
-            input_context["attachments"] = attachments
+                        thread_id=thread_id,
-            logger.debug(f"Loaded {len(attachments)} attachments for thread_id={thread_id}")
+                        role="user",
-        except Exception as e:
+                        content=query,
-            logger.error(f"Error loading attachments for thread_id={thread_id}: {e}")
+                        message_type=message_type,
-            input_context["attachments"] = []
+                        image_content=image_content,
+                        extra_metadata={"raw_message": human_message.model_dump()},
+                    )
+                except Exception as e:
+                    logger.error(f"Error saving user message: {e}")
 
-        try:
+                try:
-            full_msg = None
+                    assert thread_id, "thread_id is required"
-            langgraph_config = {"configurable": input_context}
+                    attachments = await conv_manager.get_attachments_by_thread_id(thread_id)
-            async for msg, metadata in agent.stream_messages(messages, input_context=input_context):
+                    input_context["attachments"] = attachments
-                if isinstance(msg, AIMessageChunk):
+                    logger.debug(f"Loaded {len(attachments)} attachments for thread_id={thread_id}")
-                    full_msg = msg if not full_msg else full_msg + msg
+                except Exception as e:
-                    if conf.enable_content_guard and await content_guard.check_with_keywords(full_msg.content[-20:]):
+                    logger.error(f"Error loading attachments for thread_id={thread_id}: {e}")
-                        logger.warning("Sensitive content detected in stream")
+                    input_context["attachments"] = []
-                        await save_partial_message(conv_manager, thread_id, full_msg, "content_guard_blocked")
-                        meta["time_cost"] = asyncio.get_event_loop().time() - start_time
-                        yield make_chunk(status="interrupted", message="检测到敏感内容，已中断输出", meta=meta)
-                        return
 
-                    yield make_chunk(content=msg.content, msg=msg.model_dump(), metadata=metadata, status="loading")
+                full_msg = None
+                langgraph_config = {"configurable": input_context}
+                async for msg, metadata in agent.stream_messages(messages, input_context=input_context):
+                    if isinstance(msg, AIMessageChunk):
+                        full_msg = msg if not full_msg else full_msg + msg
+                        if conf.enable_content_guard and await content_guard.check_with_keywords(full_msg.content[-20:]):
+                            logger.warning("Sensitive content detected in stream")
+                            await save_partial_message(conv_manager, thread_id, full_msg, "content_guard_blocked")
+                            meta["time_cost"] = asyncio.get_event_loop().time() - start_time
+                            yield make_chunk(status="interrupted", message="检测到敏感内容，已中断输出", meta=meta)
+                            return
 
-                else:
+                        yield make_chunk(content=msg.content, msg=msg.model_dump(), metadata=metadata, status="loading")
-                    msg_dict = msg.model_dump()
-                    yield make_chunk(msg=msg_dict, metadata=metadata, status="loading")
 
-                    try:
+                    else:
-                        if msg_dict.get("type") == "tool":
+                        msg_dict = msg.model_dump()
-                            graph = await agent.get_graph()
+                        yield make_chunk(msg=msg_dict, metadata=metadata, status="loading")
-                            state = await graph.aget_state(langgraph_config)
+
-                            agent_state = _extract_agent_state(getattr(state, "values", {})) if state else {}
+                        try:
-                            if agent_state:
+                            if msg_dict.get("type") == "tool":
-                                yield make_chunk(status="agent_state", agent_state=agent_state, meta=meta)
+                                graph = await agent.get_graph()
-                    except Exception:
+                                state = await graph.aget_state(langgraph_config)
-                        pass
+                                agent_state = _extract_agent_state(getattr(state, "values", {})) if state else {}
+                                if agent_state:
+                                    yield make_chunk(status="agent_state", agent_state=agent_state, meta=meta)
+                        except Exception:
+                            pass
+
+                if (
+                    conf.enable_content_guard
+                    and hasattr(full_msg, "content")
+                    and await content_guard.check(full_msg.content)
+                ):
+                    logger.warning("Sensitive content detected in final message")
+                    await save_partial_message(conv_manager, thread_id, full_msg, "content_guard_blocked")
+                    meta["time_cost"] = asyncio.get_event_loop().time() - start_time
+                    yield make_chunk(status="interrupted", message="检测到敏感内容，已中断输出", meta=meta)
+                    return
+
+                # After streaming finished, check for interrupts and save messages
+
+                # Check for human approval interrupts
+                async for chunk in check_and_handle_interrupts(agent, langgraph_config, make_chunk, meta, thread_id):
+                    yield chunk
 
-            if (
-                conf.enable_content_guard
-                and hasattr(full_msg, "content")
-                and await content_guard.check(full_msg.content)
-            ):
-                logger.warning("Sensitive content detected in final message")
-                await save_partial_message(conv_manager, thread_id, full_msg, "content_guard_blocked")
                 meta["time_cost"] = asyncio.get_event_loop().time() - start_time
-                yield make_chunk(status="interrupted", message="检测到敏感内容，已中断输出", meta=meta)
+                try:
-                return
+                    graph = await agent.get_graph()
+                    state = await graph.aget_state(langgraph_config)
+                    agent_state = _extract_agent_state(getattr(state, "values", {})) if state else {}
+                except Exception:
+                    agent_state = {}
 
-            # After streaming finished, check for interrupts and save messages
+                if agent_state:
+                    yield make_chunk(status="agent_state", agent_state=agent_state, meta=meta)
 
-            # Check for human approval interrupts
+                yield make_chunk(status="finished", meta=meta)
-            async for chunk in check_and_handle_interrupts(agent, langgraph_config, make_chunk, meta, thread_id):
-                yield chunk
 
-            meta["time_cost"] = asyncio.get_event_loop().time() - start_time
+                # Save all messages from LangGraph state
-            try:
+                await save_messages_from_langgraph_state(
-                graph = await agent.get_graph()
+                    agent_instance=agent,
-                state = await graph.aget_state(langgraph_config)
+                    thread_id=thread_id,
-                agent_state = _extract_agent_state(getattr(state, "values", {})) if state else {}
+                    conv_mgr=conv_manager,
-            except Exception:
+                    config_dict=langgraph_config,
-                agent_state = {}
+                )
-
-            if agent_state:
-                yield make_chunk(status="agent_state", agent_state=agent_state, meta=meta)
-
-            yield make_chunk(status="finished", meta=meta)
-
-            # Save all messages from LangGraph state
-            await save_messages_from_langgraph_state(
-                agent_instance=agent,
-                thread_id=thread_id,
-                conv_mgr=conv_manager,
-                config_dict=langgraph_config,
-            )
 
         except (asyncio.CancelledError, ConnectionError) as e:
             # 客户端主动中断连接，检查中断并保存已生成的部分内容
             logger.warning(f"Client disconnected, cancelling stream: {e}")
 
-            # 保存中断消息到数据库
+            # Run save in a separate task to avoid cancellation
-            async with db_manager.get_async_session_context() as new_db:
+            async def save_cleanup():
-                new_conv_manager = ConversationManager(new_db)
+                async with db_manager.get_async_session_context() as new_db:
-                await save_partial_message(
+                    new_conv_manager = ConversationManager(new_db)
-                    new_conv_manager,
+                    await save_partial_message(
-                    thread_id,
+                        new_conv_manager,
-                    full_msg=full_msg,
+                        thread_id,
-                    error_message="对话已中断" if not full_msg else None,
+                        full_msg=full_msg,
-                    error_type="interrupted",
+                        error_message="对话已中断" if not full_msg else None,
-                )
+                        error_type="interrupted",
+                    )
+
+            # Create a task and await it, shielding it from cancellation
+            # ensuring the DB operation completes even if the stream is cancelled
+            cleanup_task = asyncio.create_task(save_cleanup())
+            try:
+                await asyncio.shield(cleanup_task)
+            except asyncio.CancelledError:
+                pass
+            except Exception as exc:
+                logger.error(f"Error during cleanup save: {exc}")
 
             # 通知前端中断（可能发送不到，但用于一致性）
             yield make_chunk(status="interrupted", message="对话已中断", meta=meta)
@@ -780,29 +792,30 @@ async def resume_agent_chat(
         )
 
         try:
-            async for msg, metadata in stream_source:
+            async with db_manager.get_async_session_context() as db:
-                # 确保msg有正确的ID结构
+                async for msg, metadata in stream_source:
-                msg_dict = msg.model_dump()
+                    # 确保msg有正确的ID结构
-                if "id" not in msg_dict:
+                    msg_dict = msg.model_dump()
-                    msg_dict["id"] = str(uuid.uuid4())
+                    if "id" not in msg_dict:
+                        msg_dict["id"] = str(uuid.uuid4())
 
-                yield make_resume_chunk(
+                    yield make_resume_chunk(
-                    content=getattr(msg, "content", ""), msg=msg_dict, metadata=metadata, status="loading"
+                        content=getattr(msg, "content", ""), msg=msg_dict, metadata=metadata, status="loading"
+                    )
+
+                meta["time_cost"] = asyncio.get_event_loop().time() - start_time
+                yield make_resume_chunk(status="finished", meta=meta)
+
+                # 保存消息到数据库
+                langgraph_config = {"configurable": input_context}
+                conv_manager = ConversationManager(db)
+                await save_messages_from_langgraph_state(
+                    agent_instance=agent,
+                    thread_id=thread_id,
+                    conv_mgr=conv_manager,
+                    config_dict=langgraph_config,
                 )
 
-            meta["time_cost"] = asyncio.get_event_loop().time() - start_time
-            yield make_resume_chunk(status="finished", meta=meta)
-
-            # 保存消息到数据库
-            langgraph_config = {"configurable": input_context}
-            conv_manager = ConversationManager(db)
-            await save_messages_from_langgraph_state(
-                agent_instance=agent,
-                thread_id=thread_id,
-                conv_mgr=conv_manager,
-                config_dict=langgraph_config,
-            )
-
         except (asyncio.CancelledError, ConnectionError) as e:
             # 客户端主动中断连接
             logger.warning(f"Client disconnected during resume: {e}")

src/storage/db/manager.py

@@ -1,3 +1,4 @@
+import asyncio
 import json
 import os
 import pathlib
@@ -128,7 +129,9 @@ class DBManager(metaclass=SingletonMeta):
             logger.error(f"Async database operation failed: {e}")
             raise
         finally:
-            await session.close()
+            # Shield close operation to ensure connection is properly closed even if task is cancelled
+            # This prevents aiosqlite from raising errors during cancellation
+            await asyncio.shield(session.close())
 
     def check_first_run(self):
         """检查是否首次运行（同步版本）"""

web/src/components/AgentChatComponent.vue

@@ -233,10 +233,8 @@
 
 <script setup>
 import { ref, reactive, onMounted, watch, nextTick, computed, onUnmounted } from 'vue';
-import { LoadingOutlined } from '@ant-design/icons-vue';
 import { message } from 'ant-design-vue';
 import MessageInputComponent from '@/components/MessageInputComponent.vue'
-import AttachmentInputPanel from '@/components/AttachmentInputPanel.vue'
 import AttachmentOptionsComponent from '@/components/AttachmentOptionsComponent.vue'
 import AttachmentStatusIndicator from '@/components/AttachmentStatusIndicator.vue'
 import AgentMessageComponent from '@/components/AgentMessageComponent.vue'
@@ -431,27 +429,16 @@ const conversations = computed(() => {
   const threadState = currentThreadState.value;
 
   // 如果有进行中的消息且线程状态显示正在流式处理，添加进行中的对话
-  if (onGoingConvMessages.value.length > 0 && threadState?.isStreaming) {
+  if (onGoingConvMessages.value.length > 0) {
     const onGoingConv = {
       messages: onGoingConvMessages.value,
       status: 'streaming'
     };
     return [...historyConvs, onGoingConv];
   }
-
-  // 即使流式结束，如果历史记录为空但还有消息没有完全同步，也保持显示
-  if (historyConvs.length === 0 && onGoingConvMessages.value.length > 0 && !threadState?.isStreaming) {
-    const finalConv = {
-      messages: onGoingConvMessages.value,
-      status: 'finished'
-    };
-    return [finalConv];
-  }
-
   return historyConvs;
 });
 
-const isLoadingThreads = computed(() => chatUIStore.isLoadingThreads);
 const isLoadingMessages = computed(() => chatUIStore.isLoadingMessages);
 const isStreaming = computed(() => {
   const threadState = currentThreadState.value;
@@ -522,54 +509,28 @@ const cleanupThreadState = (threadId) => {
 };
 
 // ==================== STREAM HANDLING LOGIC ====================
-const resetOnGoingConv = (threadId = null, preserveMessages = false) => {
+const resetOnGoingConv = (threadId = null) => {
-  console.log('🔄 [RESET] Resetting on going conversation:', threadId, preserveMessages);
+  console.log(`🔄 [RESET] Resetting on going conversation: ${new Date().toLocaleTimeString()}.${new Date().getMilliseconds()}`, threadId);
-  if (threadId) {
+
+  const targetThreadId = threadId || currentChatId.value;
+
+  if (targetThreadId) {
     // 清理指定线程的状态
-    const threadState = getThreadState(threadId);
+    const threadState = getThreadState(targetThreadId);
     if (threadState) {
       if (threadState.streamAbortController) {
         threadState.streamAbortController.abort();
         threadState.streamAbortController = null;
       }
-      // 如果指定要保留消息，则延迟清空
+
-      if (preserveMessages) {
+      // 直接重置对话状态
-        // 延迟清空消息，给历史记录加载足够时间
-        setTimeout(() => {
-          if (threadState.onGoingConv) {
       threadState.onGoingConv = createOnGoingConvState();
-    }
-        }, 100);
-      } else {
-        threadState.onGoingConv = createOnGoingConvState();
-      }
     }
   } else {
-    // 清理当前线程或所有线程的状态
+    // 如果没有当前线程，清理所有线程状态
-    const targetThreadId = currentChatId.value;
+    Object.keys(chatState.threadStates).forEach(tid => {
-    if (targetThreadId) {
+      cleanupThreadState(tid);
-      const threadState = getThreadState(targetThreadId);
+    });
-      if (threadState) {
-        if (threadState.streamAbortController) {
-          threadState.streamAbortController.abort();
-          threadState.streamAbortController = null;
-        }
-        if (preserveMessages) {
-          setTimeout(() => {
-            if (threadState.onGoingConv) {
-              threadState.onGoingConv = createOnGoingConvState();
-            }
-          }, 100);
-        } else {
-          threadState.onGoingConv = createOnGoingConvState();
-        }
-      }
-    } else {
-      // 如果没有当前线程，清理所有线程状态
-      Object.keys(chatState.threadStates).forEach(tid => {
-        cleanupThreadState(tid);
-      });
-    }
   }
 };
 
@@ -604,10 +565,6 @@ const _processStreamChunk = (chunk, threadId) => {
           threadState.streamAbortController = null;
         }
       }
-
-      // Reload messages to show any partial content saved by the backend
-      fetchThreadMessages({ agentId: currentAgentId.value, threadId: threadId, delay: 500 });
-      resetOnGoingConv(threadId);
       return true;
     case 'human_approval_required':
       // 使用审批 composable 处理审批请求
@@ -627,22 +584,17 @@ const _processStreamChunk = (chunk, threadId) => {
       if (threadState) {
         threadState.isStreaming = false;
         if ((supportsTodo.value || supportsFiles.value) && threadState.agentState) {
-          console.log('[AgentState|Final]', {
+          console.log(`[AgentState|Final] ${new Date().toLocaleTimeString()}.${new Date().getMilliseconds()}`, {
             threadId,
             todos: threadState.agentState?.todos || [],
             files: threadState.agentState?.files || []
           });
         }
       }
-      // 异步加载历史记录，保持当前消息显示直到历史记录加载完成
-      fetchThreadMessages({ agentId: currentAgentId.value, threadId: threadId, delay: 500 })
-      .finally(() => {
-        // 历史记录加载完成后，安全地清空当前进行中的对话
-        resetOnGoingConv(threadId, true);
-      });
       return true;
     case 'interrupted':
       // 中断状态，刷新消息历史
+      console.warn("[Interrupted] case");
       if (threadState) {
         threadState.isStreaming = false;
       }
@@ -650,10 +602,6 @@ const _processStreamChunk = (chunk, threadId) => {
       if (chunkMessage) {
         message.info(chunkMessage);
       }
-      fetchThreadMessages({ agentId: currentAgentId.value, threadId: threadId, delay: 1000 })
-      .finally(() => {
-        resetOnGoingConv(threadId, true);
-      });
       return true;
   }
 
@@ -761,7 +709,7 @@ const fetchThreadMessages = async ({ agentId, threadId, delay = 0 }) => {
 
   try {
     const response = await agentApi.getAgentHistory(agentId, threadId);
-    console.log('🔄 [FETCH] Thread messages:', response);
+    console.log(`🔄 [FETCH] Thread messages: ${new Date().toLocaleTimeString()}.${new Date().getMilliseconds()}`, response);
     threadMessages.value[threadId] = response.history || [];
   } catch (error) {
     handleChatError(error, 'load');
@@ -1089,12 +1037,21 @@ const handleSendMessage = async () => {
     }
   } catch (error) {
     if (error.name !== 'AbortError') {
+      console.error('Stream error:', error);
       handleChatError(error, 'send');
+    } else {
+      console.warn("[Interrupted] Catch");
     }
-  } finally {
     threadState.isStreaming = false;
+  } finally {
     threadState.streamAbortController = null;
-    resetOnGoingConv(threadId);
+    // 异步加载历史记录，保持当前消息显示直到历史记录加载完成
+    fetchThreadMessages({ agentId: currentAgentId.value, threadId: threadId, delay: 500 })
+    .finally(() => {
+      // 历史记录加载完成后，安全地清空当前进行中的对话
+      resetOnGoingConv(threadId);
+      scrollController.scrollToBottom();
+    });
   }
 };
 
@@ -1208,6 +1165,14 @@ const handleApprovalWithStream = async (approved) => {
       threadState.isStreaming = false;
       threadState.streamAbortController = null;
     }
+
+    // 异步加载历史记录，保持当前消息显示直到历史记录加载完成
+    fetchThreadMessages({ agentId: currentAgentId.value, threadId: threadId, delay: 500 })
+    .finally(() => {
+      // 历史记录加载完成后，安全地清空当前进行中的对话
+      resetOnGoingConv(threadId);
+      scrollController.scrollToBottom();
+    });
   }
 };
 

web/src/components/RefsComponent.vue

@@ -153,23 +153,6 @@ const getModelName = (msg) => {
   }
   return null;
 }
-
-// Load existing feedback on mount
-onMounted(async () => {
-  if (msg.value?.id) {
-    try {
-      const response = await agentApi.getMessageFeedback(msg.value.id)
-      if (response.has_feedback) {
-        feedbackState.hasSubmitted = true
-        feedbackState.rating = response.feedback.rating
-        feedbackState.reason = response.feedback.reason
-      }
-    } catch (error) {
-      console.error('Failed to load feedback:', error)
-    }
-  }
-})
-
 // Handle like action
 const likeThisResponse = async (msg) => {
   if (feedbackState.hasSubmitted) {