fix(db): 使用asyncio.shield确保数据库连接正确关闭
1. 简化对话状态重置逻辑并移除冗余代码 2. 重构聊天路由的数据库会话管理
This commit is contained in:
Wenjie Zhang
parent
be3abc00b5
commit
ba2d0d0007
@ -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(
|
||||
thread_id=thread_id,
|
||||
role="user",
|
||||
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}")
|
||||
async with db_manager.get_async_session_context() as db:
|
||||
# Initialize conversation manager
|
||||
conv_manager = ConversationManager(db)
|
||||
|
||||
try:
|
||||
assert thread_id, "thread_id is required"
|
||||
attachments = await conv_manager.get_attachments_by_thread_id(thread_id)
|
||||
input_context["attachments"] = attachments
|
||||
logger.debug(f"Loaded {len(attachments)} attachments for thread_id={thread_id}")
|
||||
except Exception as e:
|
||||
logger.error(f"Error loading attachments for thread_id={thread_id}: {e}")
|
||||
input_context["attachments"] = []
|
||||
# Save user message
|
||||
try:
|
||||
await conv_manager.add_message_by_thread_id(
|
||||
thread_id=thread_id,
|
||||
role="user",
|
||||
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:
|
||||
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
|
||||
try:
|
||||
assert thread_id, "thread_id is required"
|
||||
attachments = await conv_manager.get_attachments_by_thread_id(thread_id)
|
||||
input_context["attachments"] = attachments
|
||||
logger.debug(f"Loaded {len(attachments)} attachments for thread_id={thread_id}")
|
||||
except Exception as e:
|
||||
logger.error(f"Error loading attachments for thread_id={thread_id}: {e}")
|
||||
input_context["attachments"] = []
|
||||
|
||||
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:
|
||||
msg_dict = msg.model_dump()
|
||||
yield make_chunk(msg=msg_dict, metadata=metadata, status="loading")
|
||||
yield make_chunk(content=msg.content, msg=msg.model_dump(), metadata=metadata, status="loading")
|
||||
|
||||
try:
|
||||
if msg_dict.get("type") == "tool":
|
||||
graph = await agent.get_graph()
|
||||
state = await graph.aget_state(langgraph_config)
|
||||
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
|
||||
else:
|
||||
msg_dict = msg.model_dump()
|
||||
yield make_chunk(msg=msg_dict, metadata=metadata, status="loading")
|
||||
|
||||
try:
|
||||
if msg_dict.get("type") == "tool":
|
||||
graph = await agent.get_graph()
|
||||
state = await graph.aget_state(langgraph_config)
|
||||
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)
|
||||
return
|
||||
try:
|
||||
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
|
||||
async for chunk in check_and_handle_interrupts(agent, langgraph_config, make_chunk, meta, thread_id):
|
||||
yield chunk
|
||||
yield make_chunk(status="finished", meta=meta)
|
||||
|
||||
meta["time_cost"] = asyncio.get_event_loop().time() - start_time
|
||||
try:
|
||||
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 = {}
|
||||
|
||||
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,
|
||||
)
|
||||
# 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}")
|
||||
|
||||
# 保存中断消息到数据库
|
||||
async with db_manager.get_async_session_context() as new_db:
|
||||
new_conv_manager = ConversationManager(new_db)
|
||||
await save_partial_message(
|
||||
new_conv_manager,
|
||||
thread_id,
|
||||
full_msg=full_msg,
|
||||
error_message="对话已中断" if not full_msg else None,
|
||||
error_type="interrupted",
|
||||
)
|
||||
# Run save in a separate task to avoid cancellation
|
||||
async def save_cleanup():
|
||||
async with db_manager.get_async_session_context() as new_db:
|
||||
new_conv_manager = ConversationManager(new_db)
|
||||
await save_partial_message(
|
||||
new_conv_manager,
|
||||
thread_id,
|
||||
full_msg=full_msg,
|
||||
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:
|
||||
# 确保msg有正确的ID结构
|
||||
msg_dict = msg.model_dump()
|
||||
if "id" not in msg_dict:
|
||||
msg_dict["id"] = str(uuid.uuid4())
|
||||
async with db_manager.get_async_session_context() as db:
|
||||
async for msg, metadata in stream_source:
|
||||
# 确保msg有正确的ID结构
|
||||
msg_dict = msg.model_dump()
|
||||
if "id" not in msg_dict:
|
||||
msg_dict["id"] = str(uuid.uuid4())
|
||||
|
||||
yield make_resume_chunk(
|
||||
content=getattr(msg, "content", ""), msg=msg_dict, metadata=metadata, status="loading"
|
||||
yield make_resume_chunk(
|
||||
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}")
|
||||
|
||||
@ -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):
|
||||
"""检查是否首次运行（同步版本）"""
|
||||
|
||||
@ -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) => {
|
||||
console.log('🔄 [RESET] Resetting on going conversation:', threadId, preserveMessages);
|
||||
if (threadId) {
|
||||
const resetOnGoingConv = (threadId = null) => {
|
||||
console.log(`🔄 [RESET] Resetting on going conversation: ${new Date().toLocaleTimeString()}.${new Date().getMilliseconds()}`, 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;
|
||||
if (targetThreadId) {
|
||||
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);
|
||||
});
|
||||
}
|
||||
// 如果没有当前线程，清理所有线程状态
|
||||
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();
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
@ -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) {
|
||||
|
||||
Loading…
Reference in New Issue
Block a user