fix(aibox): batch G audit fixes - llm hardening, thread lifecycle, cancel propagation, db transactions
P-series (llm_client): truncate non-transient HTTP error body to 200 chars, cap Retry-After at 3600s, validate fallback profiles (missing model skipped with warning), dedupe fallback list, remove dead _cancelled flag, document FailoverClient concurrency contract. N-series (aibox/registry): batch-level drain thread with try/finally reclamation (N2), subagent cancel propagation via shared parent event + watcher thread (N5), /new reset moved off event loop (N3), streaming forward bounded by TOOL_MAX_RESULT_LENGTH (N8), persona path boundary check (N12), broadcast except narrowed + discard (N13), per-connection task set registration (N14). db/history: FTS5 quote escaping (N7), mark_as_triggered single BEGIN IMMEDIATE transaction (N10), LIMIT parameterization (L1), archive filename collision suffix (N16). Tests: +20 cases in test_g_fixes.py, full suite 375 passed.
This commit is contained in:
1 parent
aa15841d73
commit
cccc946660
6 files changed
+690
-184
No files matched your search
+200
-124
@@ -534,6 +534,14 @@ class ExpressEmotionTool(Tool):
|
||||
if not name:
|
||||
return ""
|
||||
path = os.path.join(constants.SKILLS_DIR, name + ".md")
|
||||
# N12:persona 为 LLM 可控输入,做路径边界校验——normcase 统一大小写与
|
||||
# 分隔符后要求解析结果仍在 SKILLS_DIR 内,防止 ../../ 越界读取并注入
|
||||
skills_root = os.path.normcase(os.path.abspath(constants.SKILLS_DIR))
|
||||
resolved = os.path.normcase(os.path.abspath(path))
|
||||
if not resolved.startswith(skills_root + os.sep):
|
||||
self.logger.warning(f"人格技能名非法(越界路径): {name}")
|
||||
return ""
|
||||
path = resolved
|
||||
if not os.path.isfile(path):
|
||||
self.logger.warning(f"未找到人格技能文件: {path}")
|
||||
return ""
|
||||
@@ -612,7 +620,8 @@ _SUBAGENT_TOOL_FAILED_MARKER = "[SUBAGENT_TOOL_FAILED]\n"
|
||||
|
||||
def run_full_subagent(role, task, config, logger=None, model=None,
|
||||
temperature=None, max_tokens=2048, max_turns=10,
|
||||
output_callback=None, outcome: dict = None) -> str:
|
||||
output_callback=None, outcome: dict = None,
|
||||
cancel_event=None) -> str:
|
||||
"""完整子代理:拥有全部工具,独立上下文,多轮对话直到任务完成。
|
||||
|
||||
子代理是完整的 DeepSeekChat 实例,拥有文件操作、命令执行、网页访问、
|
||||
@@ -623,8 +632,12 @@ def run_full_subagent(role, task, config, logger=None, model=None,
|
||||
output_callback: 可选回调(msg_type, content),用于转发子代理的实时思考过程
|
||||
(reasoning_start / reasoning / reasoning_end),不传则静默执行。
|
||||
|
||||
cancel_event: 父代理取消事件(N5)。传入时子代理共享该事件并起观察者
|
||||
线程转发深度中断——父级取消(用户中断 / --once 超时)可即时终止子代理。
|
||||
|
||||
供 create_subagent 工具使用。
|
||||
"""
|
||||
_watch_done = threading.Event() # N5:取消观察者线程退出信号(finally 置位)
|
||||
try:
|
||||
max_tokens = max(1, min(int(max_tokens or 2048), constants.DEFAULT_MAX_TOKENS))
|
||||
temp = temperature if temperature is not None else constants.DEFAULT_TEMPERATURE
|
||||
@@ -653,6 +666,23 @@ def run_full_subagent(role, task, config, logger=None, model=None,
|
||||
exit_callback=None,
|
||||
)
|
||||
|
||||
# N5(审计 2026-08-18):父代理取消传播。子代理此前持完全独立的
|
||||
# _cancel_event/executor/HTTP client,父级 cancel 或 --once 超时杀不到
|
||||
# 它,需等子代理自行结束。现共享父取消事件(子代理在 think 前置、
|
||||
# 工具调用间、流式事件间的既有检查点立即感知),并由观察者线程转发
|
||||
# 深度中断(sub_chat.cancel() 强关在途 HTTP 流并 kill 其子进程)。
|
||||
if cancel_event is not None:
|
||||
sub_chat._cancel_event = cancel_event
|
||||
|
||||
def _watch_parent_cancel():
|
||||
while not _watch_done.wait(timeout=0.2):
|
||||
if cancel_event.is_set():
|
||||
sub_chat.cancel()
|
||||
return
|
||||
|
||||
threading.Thread(target=_watch_parent_cancel, daemon=True,
|
||||
name="subagent-cancel-watch").start()
|
||||
|
||||
# 覆盖模型(如果指定了不同模型)
|
||||
# M1:经 set_client 统一替换——同步刷新 compressor.llm_client 与
|
||||
# api_key/name 快照,避免压缩链仍用默认模型发请求
|
||||
@@ -740,6 +770,9 @@ def run_full_subagent(role, task, config, logger=None, model=None,
|
||||
if logger:
|
||||
logger.error(f"子代理执行异常: {e}")
|
||||
return f"❌ 子代理执行失败: {str(e)}"
|
||||
finally:
|
||||
# N5:通知取消观察者线程退出(子代理已结束,无论成败)
|
||||
_watch_done.set()
|
||||
|
||||
|
||||
class SubAgentTool(Tool):
|
||||
@@ -750,6 +783,12 @@ class SubAgentTool(Tool):
|
||||
结束后由子代理基于完整执行过程生成最终报告,返回给父代理。
|
||||
"""
|
||||
|
||||
def __init__(self, logger=None, cancel_event=None):
|
||||
super().__init__(logger)
|
||||
# N5:父代理(DeepSeekChat)取消事件——父级被取消(用户中断/--once
|
||||
# 超时)时传播给 run_full_subagent,即时终止子代理
|
||||
self.cancel_event = cancel_event
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "create_subagent"
|
||||
|
||||
@@ -859,7 +898,8 @@ class SubAgentTool(Tool):
|
||||
model=model, temperature=temperature,
|
||||
max_tokens=max_tokens, max_turns=max_turns,
|
||||
output_callback=self._subagent_output,
|
||||
outcome=outcome)
|
||||
outcome=outcome,
|
||||
cancel_event=self.cancel_event)
|
||||
# 冲刷子代理回答缓冲(最后一段 chunk 可能尚未发出)
|
||||
self._flush_subagent_answer()
|
||||
# 子代理内部工具失败 → 打上标记,由 _handle_tool_calls 检测并置 orchestrator stats
|
||||
@@ -1103,7 +1143,8 @@ class DeepSeekChat:
|
||||
builtin_tools.append(GetProxyTool(self.logger))
|
||||
|
||||
# 子代理工具:主代理可派生独立上下文/角色的轻量子代理执行任务
|
||||
builtin_tools.append(SubAgentTool(self.logger))
|
||||
# cancel_event 注入(N5):父级取消传播到子代理
|
||||
builtin_tools.append(SubAgentTool(self.logger, cancel_event=self._cancel_event))
|
||||
|
||||
# 注册内置工具(标记为内置)
|
||||
self.tool_registry.register_many(builtin_tools, is_builtin=True)
|
||||
@@ -1243,122 +1284,140 @@ class DeepSeekChat:
|
||||
# 结果大小阈值(字符数),从配置获取
|
||||
MAX_RESULT_LENGTH = constants.TOOL_MAX_RESULT_LENGTH
|
||||
|
||||
for tool_call in tool_calls:
|
||||
# 每个工具调用前检查取消信号
|
||||
if self._is_cancelled():
|
||||
self.logger.info("工具调用被用户中断")
|
||||
break
|
||||
|
||||
function = tool_call.get("function", {})
|
||||
tool_name = function.get("name", "unknown")
|
||||
arguments = function.get("arguments", "{}")
|
||||
tc_id = tool_call.get("id") or f"call_fallback_{self.stats['tool_calls']}_{int(time.time()*1000)}"
|
||||
|
||||
try:
|
||||
args = json.loads(arguments) if arguments else {}
|
||||
except json.JSONDecodeError:
|
||||
args = {}
|
||||
|
||||
# 开始调用工具时发送提示
|
||||
if output_callback:
|
||||
output_callback("line", "\n\n")
|
||||
output_callback("line", f"🔧 调用工具: {tool_name}\n")
|
||||
|
||||
result_raw = None
|
||||
if self.tool_profile == TOOL_PROFILE_ORCHESTRATOR and tool_name in WORKER_ONLY_TOOLS:
|
||||
# 编排者不持有执行面工具:拦截并纠正,引导改用 create_subagent
|
||||
result_raw = (
|
||||
f"⛔ '{tool_name}' 调用被拦截:你是编排者,没有执行面工具(命令/文件/网页由 worker 子代理持有)。\n"
|
||||
"请改用 create_subagent 派生子代理执行该任务:在 role 中写明执行者身份,"
|
||||
"在 task 中写清目标、具体步骤与验收标准(任务描述须自包含,子代理无法向你提问)。"
|
||||
)
|
||||
self.logger.warning(f"编排者误调执行面工具 {tool_name},已拦截并引导改用 create_subagent")
|
||||
else:
|
||||
# N2 修复(审计 2026-08-18):drain 线程改为整个工具批次仅创建一个,
|
||||
# 由独立 stop 事件控制生命周期,批次结束(正常/异常/中断)在 finally
|
||||
# 中统一置位、join 并排空残余。此前每次工具调用都启动一个仅靠 cancel
|
||||
# 事件退出的常驻线程(注释声称"队列空自动退出",循环条件却从未检查
|
||||
# 队列是否为空),闭包长期持有 output_callback(含 websocket 引用)
|
||||
# 阻碍 GC,常驻服务下线程与内存持续累积。
|
||||
drain_stop = threading.Event()
|
||||
|
||||
def _drain_queue(q=output_queue, cb=output_callback, stop_event=drain_stop):
|
||||
"""持续排空 output_queue 并转发 line 输出(长时间工具执行期间防 WS 空闲断线)"""
|
||||
while not stop_event.is_set():
|
||||
try:
|
||||
# 启动后台线程持续排空 output_queue,将流式输出实时转发给 output_callback,
|
||||
# 避免 SubAgent 等长时间工具执行期间前端收不到数据导致 WS 空闲断线
|
||||
def _drain_queue(q=output_queue, cb=output_callback, stop_event=self._cancel_event):
|
||||
while not stop_event.is_set():
|
||||
try:
|
||||
msg_type, content = q.get(timeout=0.1)
|
||||
if cb and msg_type == "line":
|
||||
cb("line", str(content) if content is not None else "")
|
||||
except queue.Empty:
|
||||
continue
|
||||
msg_type, content = q.get(timeout=0.1)
|
||||
if cb and msg_type == "line":
|
||||
cb("line", str(content) if content is not None else "")
|
||||
except queue.Empty:
|
||||
continue
|
||||
|
||||
import threading as _th
|
||||
drain_thread = _th.Thread(target=_drain_queue, daemon=True)
|
||||
drain_thread.start()
|
||||
drain_thread = threading.Thread(target=_drain_queue, daemon=True,
|
||||
name="tool-output-drain")
|
||||
drain_thread.start()
|
||||
|
||||
result_raw = self.tool_registry.execute_tool(tool_name, **args)
|
||||
try:
|
||||
for tool_call in tool_calls:
|
||||
# 每个工具调用前检查取消信号
|
||||
if self._is_cancelled():
|
||||
self.logger.info("工具调用被用户中断")
|
||||
break
|
||||
|
||||
# 工具执行完毕,排空队列中剩余内容并等待 drain 线程退出
|
||||
while not output_queue.empty():
|
||||
try:
|
||||
msg_type, content = output_queue.get_nowait()
|
||||
if output_callback and msg_type == "line":
|
||||
output_callback("line", str(content) if content is not None else "")
|
||||
except queue.Empty:
|
||||
break
|
||||
# drain 线程会在下一次 get 超时后自动退出(队列已空)
|
||||
except Exception as e:
|
||||
result_raw = f"❌ 工具 {tool_name} 执行异常: {str(e)}"
|
||||
self.logger.error(f"工具 {tool_name} 执行异常: {e}")
|
||||
self.stats["tool_failed"] = True
|
||||
function = tool_call.get("function", {})
|
||||
tool_name = function.get("name", "unknown")
|
||||
arguments = function.get("arguments", "{}")
|
||||
tc_id = tool_call.get("id") or f"call_fallback_{self.stats['tool_calls']}_{int(time.time()*1000)}"
|
||||
|
||||
# 子代理内部工具失败透传:SubAgentTool 在结果前缀打标记,此处检测并置 stats
|
||||
if tool_name == "create_subagent" and result_raw and \
|
||||
str(result_raw).startswith(_SUBAGENT_TOOL_FAILED_MARKER):
|
||||
self.stats["tool_failed"] = True
|
||||
result_raw = str(result_raw)[len(_SUBAGENT_TOOL_FAILED_MARKER):]
|
||||
try:
|
||||
args = json.loads(arguments) if arguments else {}
|
||||
except json.JSONDecodeError:
|
||||
args = {}
|
||||
|
||||
# 安全转换并清理结果
|
||||
result = self._sanitize_tool_result(result_raw)
|
||||
|
||||
# drain 线程已排空队列中的 line 输出,这里只需判断是否有流式输出
|
||||
has_stream_output = not output_queue.empty()
|
||||
|
||||
# 普通(非流式)工具没有队列输出,直接发送结果
|
||||
if not has_stream_output and output_callback and result:
|
||||
output_callback("line", result + "\n")
|
||||
|
||||
# 检查结果大小,如果超过阈值,不保留截断内容,仅返回提示信息
|
||||
if len(result) > MAX_RESULT_LENGTH:
|
||||
truncated_note = f"⚠️ 工具调用结果长度超过 {MAX_RESULT_LENGTH} 字符,已丢弃。请使用更精确的查询获取所需信息。"
|
||||
result = truncated_note
|
||||
# 开始调用工具时发送提示
|
||||
if output_callback:
|
||||
output_callback("line", truncated_note + "\n")
|
||||
|
||||
# 结束调用工具时发送换行
|
||||
if output_callback:
|
||||
output_callback("line", "\n\n")
|
||||
|
||||
# 历史记录(安全截断)
|
||||
if self.history_manager:
|
||||
output_callback("line", "\n\n")
|
||||
output_callback("line", f"🔧 调用工具: {tool_name}\n")
|
||||
|
||||
result_raw = None
|
||||
if self.tool_profile == TOOL_PROFILE_ORCHESTRATOR and tool_name in WORKER_ONLY_TOOLS:
|
||||
# 编排者不持有执行面工具:拦截并纠正,引导改用 create_subagent
|
||||
result_raw = (
|
||||
f"⛔ '{tool_name}' 调用被拦截:你是编排者,没有执行面工具(命令/文件/网页由 worker 子代理持有)。\n"
|
||||
"请改用 create_subagent 派生子代理执行该任务:在 role 中写明执行者身份,"
|
||||
"在 task 中写清目标、具体步骤与验收标准(任务描述须自包含,子代理无法向你提问)。"
|
||||
)
|
||||
self.logger.warning(f"编排者误调执行面工具 {tool_name},已拦截并引导改用 create_subagent")
|
||||
else:
|
||||
try:
|
||||
result_raw = self.tool_registry.execute_tool(tool_name, **args)
|
||||
# 排空队列中剩余标记(complete 标记供 has_stream_output 判定、
|
||||
# 残余 line 补发)——与 drain 线程并发消费安全(每项仅被消费一次)
|
||||
while not output_queue.empty():
|
||||
try:
|
||||
msg_type, content = output_queue.get_nowait()
|
||||
if output_callback and msg_type == "line":
|
||||
output_callback("line", str(content) if content is not None else "")
|
||||
except queue.Empty:
|
||||
break
|
||||
except Exception as e:
|
||||
result_raw = f"❌ 工具 {tool_name} 执行异常: {str(e)}"
|
||||
self.logger.error(f"工具 {tool_name} 执行异常: {e}")
|
||||
self.stats["tool_failed"] = True
|
||||
|
||||
# 子代理内部工具失败透传:SubAgentTool 在结果前缀打标记,此处检测并置 stats
|
||||
if tool_name == "create_subagent" and result_raw and \
|
||||
str(result_raw).startswith(_SUBAGENT_TOOL_FAILED_MARKER):
|
||||
self.stats["tool_failed"] = True
|
||||
result_raw = str(result_raw)[len(_SUBAGENT_TOOL_FAILED_MARKER):]
|
||||
|
||||
# 安全转换并清理结果
|
||||
result = self._sanitize_tool_result(result_raw)
|
||||
|
||||
# drain 线程持续转发 line 输出,这里通过 complete 标记判断是否流式工具
|
||||
has_stream_output = not output_queue.empty()
|
||||
|
||||
# 普通(非流式)工具没有队列输出,直接发送结果
|
||||
if not has_stream_output and output_callback and result:
|
||||
output_callback("line", result + "\n")
|
||||
|
||||
# 检查结果大小,如果超过阈值,不保留截断内容,仅返回提示信息
|
||||
if len(result) > MAX_RESULT_LENGTH:
|
||||
truncated_note = f"⚠️ 工具调用结果长度超过 {MAX_RESULT_LENGTH} 字符,已丢弃。请使用更精确的查询获取所需信息。"
|
||||
result = truncated_note
|
||||
if output_callback:
|
||||
output_callback("line", truncated_note + "\n")
|
||||
|
||||
# 结束调用工具时发送换行
|
||||
if output_callback:
|
||||
output_callback("line", "\n\n")
|
||||
|
||||
# 历史记录(安全截断)
|
||||
if self.history_manager:
|
||||
try:
|
||||
history_preview = result[:200] + "..." if len(result) > 200 else result
|
||||
self.history_manager.add_tool_call(tool_name, args, history_preview)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# 检查是否结束对话
|
||||
end_conversation = (tool_name == "save_memory_and_end_conversation" and "✅" in result)
|
||||
if end_conversation:
|
||||
self.logger.info("检测到保存记忆并结束对话工具调用")
|
||||
self.conversation_active = False
|
||||
|
||||
tool_messages.append({
|
||||
"role": "tool",
|
||||
"tool_call_id": tc_id,
|
||||
"content": result
|
||||
})
|
||||
|
||||
if end_conversation:
|
||||
break
|
||||
finally:
|
||||
# 停止 drain 线程并回收(线程 0.1s 轮询粒度内感知置位)
|
||||
drain_stop.set()
|
||||
drain_thread.join(timeout=2)
|
||||
# 排空残余输出(drain 停止后无并发消费者,残余 line 补发)
|
||||
while True:
|
||||
try:
|
||||
history_preview = result[:200] + "..." if len(result) > 200 else result
|
||||
self.history_manager.add_tool_call(tool_name, args, history_preview)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# 检查是否结束对话
|
||||
end_conversation = (tool_name == "save_memory_and_end_conversation" and "✅" in result)
|
||||
if end_conversation:
|
||||
self.logger.info("检测到保存记忆并结束对话工具调用")
|
||||
self.conversation_active = False
|
||||
|
||||
tool_messages.append({
|
||||
"role": "tool",
|
||||
"tool_call_id": tc_id,
|
||||
"content": result
|
||||
})
|
||||
|
||||
if end_conversation:
|
||||
break
|
||||
|
||||
# 清除输出队列
|
||||
self.tool_registry.set_output_queue(None)
|
||||
|
||||
msg_type, content = output_queue.get_nowait()
|
||||
if output_callback and msg_type == "line":
|
||||
output_callback("line", str(content) if content is not None else "")
|
||||
except queue.Empty:
|
||||
break
|
||||
# 清除输出队列
|
||||
self.tool_registry.set_output_queue(None)
|
||||
|
||||
return tool_messages
|
||||
|
||||
def _should_search_memory(self, user_input: str) -> bool:
|
||||
@@ -1960,7 +2019,7 @@ class WebSocketHandler:
|
||||
self.logger = Logger("WebSocket")
|
||||
self.connected_clients = set()
|
||||
self.client_sessions = {}
|
||||
self._processing_tasks = {} # websocket id -> 当前正在处理的asyncio Task
|
||||
self._processing_tasks = {} # websocket -> {asyncio.Task}(N14:集合登记,支持并发多任务)
|
||||
self._shared_skill_manager = None # 共享技能管理器(懒加载,供所有连接复用)
|
||||
|
||||
def _get_shared_skill_manager(self):
|
||||
@@ -2020,7 +2079,9 @@ class WebSocketHandler:
|
||||
|
||||
async def unregister(self, websocket):
|
||||
"""注销客户端"""
|
||||
self.connected_clients.remove(websocket)
|
||||
# N13:discard 替代 remove——并发注销/重复注销时 remove 抛 KeyError
|
||||
# 会逃逸到调用方(如 broadcast 的 except 分支),discard 幂等
|
||||
self.connected_clients.discard(websocket)
|
||||
if websocket in self.client_sessions:
|
||||
session = self.client_sessions.pop(websocket)
|
||||
# 释放该连接独立的AI实例资源
|
||||
@@ -2128,7 +2189,8 @@ class WebSocketHandler:
|
||||
for websocket in self.connected_clients.copy():
|
||||
try:
|
||||
await self.send_message(websocket, message)
|
||||
except:
|
||||
except Exception:
|
||||
# N13:不再裸 except——吞掉 CancelledError 会破坏任务取消语义
|
||||
await self.unregister(websocket)
|
||||
|
||||
async def broadcast_task_status(self, task: Dict, status: str, message: str = None):
|
||||
@@ -2322,12 +2384,20 @@ class WebSocketHandler:
|
||||
if not ai.is_busy():
|
||||
break
|
||||
# 重置对话状态,重新加载系统提示词
|
||||
ai.reset_conversation(reload_prompts=True)
|
||||
# N3:reset_conversation 内部同步获取 _msg_lock(繁忙对话持锁
|
||||
# 可达数十秒),直接调用会冻结整个事件循环(所有连接的心跳/
|
||||
# 消息收发全部停摆)。移入线程池执行,事件循环仅等待结果。
|
||||
loop = asyncio.get_running_loop()
|
||||
|
||||
def _do_reset():
|
||||
ai.reset_conversation(reload_prompts=True)
|
||||
# 重新初始化提示词管理器(确保最新)
|
||||
self.ai_manager.prompt_manager.reload()
|
||||
|
||||
await loop.run_in_executor(None, _do_reset)
|
||||
# 重置该连接对话进行中标记
|
||||
if session:
|
||||
session["dialogue_in_progress"] = False
|
||||
# 重新初始化提示词管理器(确保最新)
|
||||
self.ai_manager.prompt_manager.reload()
|
||||
self.logger.info(i18n.get('new_conversation_reset'))
|
||||
|
||||
await self.send_message(websocket, {
|
||||
@@ -2785,20 +2855,26 @@ class EnhancedAIManager:
|
||||
task = asyncio.create_task(
|
||||
self.ws_handler.handle_message(websocket, message)
|
||||
)
|
||||
self.ws_handler._processing_tasks[id(websocket)] = task
|
||||
# N14:按连接登记任务集合。此前按 ws id 覆盖式登记,同连接并发
|
||||
# 多任务时先完成任务的 done_callback 会误 pop 掉后续任务的登记
|
||||
tasks = self.ws_handler._processing_tasks.setdefault(websocket, set())
|
||||
tasks.add(task)
|
||||
|
||||
def _cleanup_task(t, ws_id=id(websocket)):
|
||||
self.ws_handler._processing_tasks.pop(ws_id, None)
|
||||
def _cleanup_task(t, ws=websocket):
|
||||
pending = self.ws_handler._processing_tasks.get(ws)
|
||||
if pending is not None:
|
||||
pending.discard(t)
|
||||
task.add_done_callback(_cleanup_task)
|
||||
except websockets.exceptions.ConnectionClosed as e:
|
||||
self.logger.info(f"连接关闭: {e}")
|
||||
except Exception as e:
|
||||
self.logger.error(f"处理错误: {e}")
|
||||
finally:
|
||||
# 连接关闭时取消正在进行的处理任务
|
||||
task = self.ws_handler._processing_tasks.pop(id(websocket), None)
|
||||
if task and not task.done():
|
||||
task.cancel()
|
||||
# 连接关闭时取消该连接所有仍在进行的处理任务(N14:此前仅最后一个)
|
||||
pending = self.ws_handler._processing_tasks.pop(websocket, set())
|
||||
for t in pending:
|
||||
if not t.done():
|
||||
t.cancel()
|
||||
await self.ws_handler.unregister(websocket)
|
||||
|
||||
async def websocket_server(self):
|
||||
|
||||
+33
-12
@@ -47,9 +47,6 @@ class LLMClient(ABC):
|
||||
self.default_temperature = default_temperature
|
||||
# 当前活动的 HTTP 响应引用,供 cancel() 强制关闭底层连接以中断阻塞读取
|
||||
self._current_response = None
|
||||
# L4b:用户中断标志——cancel() 强关连接引发的异常与真实网络故障不可区分,
|
||||
# FailoverClient 据此标志跳过转移(否则用户越中断、系统越换厂商重试)
|
||||
self._cancelled = False
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
@@ -72,8 +69,11 @@ class LLMClient(ABC):
|
||||
pass
|
||||
|
||||
def cancel(self):
|
||||
"""中断当前流式请求。线程安全,可从任意线程调用。"""
|
||||
self._cancelled = True
|
||||
"""中断当前流式请求。线程安全,可从任意线程调用。
|
||||
|
||||
注:用户中断标志由 FailoverClient 维护(其 cancel() 广播整条链并置
|
||||
自身 _cancelled,据此跳过转移)——底层 client 只负责强关连接。
|
||||
"""
|
||||
resp = self._current_response
|
||||
self._current_response = None
|
||||
if resp is not None:
|
||||
@@ -213,7 +213,7 @@ class DeepSeekClient(LLMClient):
|
||||
f"LLM 流式请求瞬时失败: HTTP {response.status_code}: {response.text[:200]}",
|
||||
retry_after=_retry_after_of(response))
|
||||
raise RuntimeError(
|
||||
f"HTTP {response.status_code}: {response.text}")
|
||||
f"HTTP {response.status_code}: {response.text[:200]}")
|
||||
|
||||
try:
|
||||
for line in response.iter_lines():
|
||||
@@ -282,12 +282,16 @@ class TransientLLMError(RuntimeError):
|
||||
|
||||
|
||||
def _retry_after_of(response) -> Optional[float]:
|
||||
"""从 Retry-After 响应头解析冷却秒数(仅支持数字秒格式,其余忽略返回 None)。"""
|
||||
"""从 Retry-After 响应头解析冷却秒数(仅支持数字秒格式,其余忽略返回 None)。
|
||||
|
||||
P2:上限 cap 3600s——异常/恶意服务端返回超大值(如 999999)时,
|
||||
避免档案被单条响应头冷却数天。
|
||||
"""
|
||||
try:
|
||||
v = response.headers.get("Retry-After")
|
||||
if v is None:
|
||||
return None
|
||||
return max(1.0, float(v))
|
||||
return min(max(1.0, float(v)), 3600.0)
|
||||
except (TypeError, ValueError, AttributeError):
|
||||
return None
|
||||
|
||||
@@ -316,6 +320,11 @@ class FailoverClient:
|
||||
- timeout_budget 限制整条链尝试的总秒数(缺省不限),耗尽即停止尝试剩余档案,
|
||||
防止长链故障时超时叠加(3 档案最坏 3×timeout)
|
||||
- 429 响应带 Retry-After(数字秒)时按其延长该档案冷却期
|
||||
|
||||
并发契约(P4):实例非线程安全(_cancelled/_call_started/_stats 无锁)。
|
||||
当前调用面满足"单线程串行调用 + cancel() 可从其他线程"的模式(compressor
|
||||
在 _msg_lock 内同步调用、子代理各持独立实例);如需多线程共享调用,
|
||||
需由调用方加锁或为实例补锁。
|
||||
"""
|
||||
|
||||
def __init__(self, clients: List, logger=None, cooldown_seconds: float = 60.0,
|
||||
@@ -657,23 +666,35 @@ def create_llm_client(config, api_key: Optional[str] = None,
|
||||
# L4a:构造 fallback 链(一层,不递归——fallback 档案自身的 fallback 字段被忽略)
|
||||
models = config.get("models") or {}
|
||||
chain = [primary]
|
||||
seen = {profile_name} # P6:重复项去重(冷却键同名冲突/重复撞同端点)
|
||||
for fb_name in fallback_names:
|
||||
if fb_name == profile_name:
|
||||
raise ValueError(
|
||||
f"模型档案 '{profile_name}' 的 fallback 不能包含自身,"
|
||||
f"请检查 config.json 的 models.{profile_name}.fallback")
|
||||
if fb_name in seen:
|
||||
if logger is not None and hasattr(logger, "warning"):
|
||||
logger.warning(f"fallback 列表重复项 '{fb_name}',已忽略")
|
||||
continue
|
||||
seen.add(fb_name)
|
||||
fb_profile = models.get(fb_name)
|
||||
if fb_profile is None:
|
||||
raise ValueError(
|
||||
f"模型档案 '{profile_name}' 的 fallback 指向不存在的档案 '{fb_name}',"
|
||||
f"可用档案:{', '.join(models) or '(无)'}")
|
||||
try:
|
||||
# P3:fallback 档案与主档案同等校验(缺 api_url/model 时跳过并告警,
|
||||
# 而非构造出 model=None 的 client 在请求期表现为神秘失败)
|
||||
_validate_profile(fb_name, fb_profile)
|
||||
fb_model = fb_profile.get("model")
|
||||
_require_model(fb_name, fb_model)
|
||||
chain.append(_construct_profile_client(
|
||||
config, fb_profile, fb_name, fb_profile.get("model"),
|
||||
config, fb_profile, fb_name, fb_model,
|
||||
logger=logger))
|
||||
except RuntimeError as e:
|
||||
# 环境问题(如 fallback 档案缺 API Key):跳过该 fallback 并告警,
|
||||
# 不阻断主档案可用性(配置结构错误在上方已硬性拦截)
|
||||
except (RuntimeError, ValueError) as e:
|
||||
# 环境问题(如 fallback 档案缺 API Key)或配置缺陷(缺 model):
|
||||
# 跳过该 fallback 并告警,不阻断主档案可用性
|
||||
# (配置结构错误——fallback 指向不存在的档案——在上方已硬性拦截)
|
||||
if logger is not None and hasattr(logger, "warning"):
|
||||
logger.warning(f"fallback 档案 '{fb_name}' 构造失败,已忽略:{e}")
|
||||
return FailoverClient(chain, logger=logger,
|
||||
|
||||
@@ -61,7 +61,10 @@ class ConversationHistory:
|
||||
try:
|
||||
with open(self.history_file, 'r', encoding='utf-8') as f:
|
||||
return json.load(f)
|
||||
except:
|
||||
except Exception as e:
|
||||
# N16:收窄为 Exception——裸 except 连 KeyboardInterrupt/SystemExit
|
||||
# 也吞掉;损坏时重置为空索引并告警
|
||||
self.logger.warning(f"历史索引文件损坏,已重置: {e}")
|
||||
return {"conversations": [], "last_id": 0}
|
||||
return {"conversations": [], "last_id": 0}
|
||||
|
||||
@@ -132,8 +135,14 @@ class ConversationHistory:
|
||||
importance=2
|
||||
)
|
||||
|
||||
# N16:归档文件名加序号防同秒覆盖——DATE_FORMAT 精度为秒,
|
||||
# 同一秒内两次归档(如测试/快速重置)会静默覆盖前一份
|
||||
timestamp = datetime.now().strftime(constants.DATE_FORMAT)
|
||||
archive_filename = f"conversation_{timestamp}.json"
|
||||
counter = 1
|
||||
while (self.archive_dir / archive_filename).exists():
|
||||
archive_filename = f"conversation_{timestamp}_{counter}.json"
|
||||
counter += 1
|
||||
archive_path = self.archive_dir / archive_filename
|
||||
|
||||
with open(archive_path, 'w', encoding='utf-8') as f:
|
||||
|
||||
+65
-45
@@ -176,7 +176,10 @@ class MemoryDatabase:
|
||||
def _search_via_fts(self, cursor, keywords: List[str], match_all: bool, limit: int) -> List[Dict]:
|
||||
"""FTS5 全文检索(关键词含特殊字符导致 MATCH 语法错误时返回空,走 LIKE 兜底)"""
|
||||
joiner = " AND " if match_all else " OR "
|
||||
query = joiner.join([f'"{kw}"' for kw in keywords])
|
||||
# N7:短语内双引号翻倍转义——关键词含 " 时若不翻倍会提前闭合短语
|
||||
# 边界,破坏 MATCH 语法或改变匹配语义(FTS5 短语引号转义规则)
|
||||
terms = [f'"{kw.replace(chr(34), chr(34) * 2)}"' for kw in keywords]
|
||||
query = joiner.join(terms)
|
||||
|
||||
try:
|
||||
cursor.execute('''
|
||||
@@ -414,7 +417,8 @@ class MemoDatabase:
|
||||
params = [compare_time, compare_time]
|
||||
|
||||
if limit:
|
||||
query += f' LIMIT {limit}'
|
||||
query += ' LIMIT ?'
|
||||
params.append(int(limit))
|
||||
|
||||
cursor.execute(query, params)
|
||||
|
||||
@@ -455,10 +459,12 @@ class MemoDatabase:
|
||||
ORDER BY created_at ASC
|
||||
'''
|
||||
|
||||
params = []
|
||||
if limit:
|
||||
query += f' LIMIT {limit}'
|
||||
query += ' LIMIT ?'
|
||||
params.append(int(limit))
|
||||
|
||||
cursor.execute(query)
|
||||
cursor.execute(query, params)
|
||||
|
||||
results = []
|
||||
for row in cursor.fetchall():
|
||||
@@ -508,64 +514,78 @@ class MemoDatabase:
|
||||
|
||||
def mark_as_triggered(self, memo_id: int) -> Optional[str]:
|
||||
"""标记任务已被触发,并根据重复规则计算下一次提醒时间"""
|
||||
memo = self.get_memo(memo_id)
|
||||
if not memo:
|
||||
return None
|
||||
|
||||
# N10:读取与更新收进同一连接的 BEGIN IMMEDIATE 事务。此前先
|
||||
# get_memo(独立连接)后 UPDATE,check-then-act 之间存在竞态:
|
||||
# 调度器并发触发同一任务会双计数/双推进 repeat。IMMEDIATE 在事务
|
||||
# 开始即取写锁,串行化并发触发。
|
||||
now = datetime.now()
|
||||
next_reminder = None
|
||||
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("BEGIN IMMEDIATE")
|
||||
try:
|
||||
cursor.execute('SELECT * FROM memos WHERE id = ?', (memo_id,))
|
||||
row = cursor.fetchone()
|
||||
if row is None:
|
||||
conn.rollback()
|
||||
return None
|
||||
memo = dict(row)
|
||||
if memo.get('tags'):
|
||||
memo['tags'] = memo['tags'].split(',')
|
||||
|
||||
cursor.execute('''
|
||||
UPDATE memos
|
||||
SET last_triggered_at = ?, trigger_count = trigger_count + 1
|
||||
WHERE id = ?
|
||||
''', (now.strftime(constants.DATETIME_FORMAT), memo_id))
|
||||
cursor.execute('''
|
||||
UPDATE memos
|
||||
SET last_triggered_at = ?, trigger_count = trigger_count + 1
|
||||
WHERE id = ?
|
||||
''', (now.strftime(constants.DATETIME_FORMAT), memo_id))
|
||||
|
||||
# 处理重复任务
|
||||
if memo['repeat_type'] != 'none' and not self._should_end_repeat(memo, now):
|
||||
# 使用当前时间作为基准来计算下一次执行时间
|
||||
next_reminder = self._calculate_next_reminder(memo, now)
|
||||
# 处理重复任务
|
||||
if memo['repeat_type'] != 'none' and not self._should_end_repeat(memo, now):
|
||||
# 使用当前时间作为基准来计算下一次执行时间
|
||||
next_reminder = self._calculate_next_reminder(memo, now)
|
||||
|
||||
if next_reminder:
|
||||
cursor.execute('''
|
||||
UPDATE memos
|
||||
SET reminder_time = ?
|
||||
WHERE id = ?
|
||||
''', (next_reminder.strftime(constants.DATETIME_FORMAT), memo_id))
|
||||
if next_reminder:
|
||||
cursor.execute('''
|
||||
UPDATE memos
|
||||
SET reminder_time = ?
|
||||
WHERE id = ?
|
||||
''', (next_reminder.strftime(constants.DATETIME_FORMAT), memo_id))
|
||||
|
||||
self.logger.info(f"任务 #{memo_id} 已更新下一次执行时间: {next_reminder}")
|
||||
else:
|
||||
# 重复结束,标记为完成
|
||||
self.logger.info(f"任务 #{memo_id} 已更新下一次执行时间: {next_reminder}")
|
||||
else:
|
||||
# 重复结束,标记为完成
|
||||
cursor.execute('''
|
||||
UPDATE memos
|
||||
SET is_completed = 1, completed_at = ?
|
||||
WHERE id = ?
|
||||
''', (now.strftime(constants.DATETIME_FORMAT), memo_id))
|
||||
|
||||
self.logger.info(f"任务 #{memo_id} 重复结束,已标记为完成")
|
||||
elif memo['is_immediate']:
|
||||
# 即时任务执行后标记为完成
|
||||
cursor.execute('''
|
||||
UPDATE memos
|
||||
SET is_completed = 1, completed_at = ?
|
||||
WHERE id = ?
|
||||
''', (now.strftime(constants.DATETIME_FORMAT), memo_id))
|
||||
|
||||
self.logger.info(f"任务 #{memo_id} 重复结束,已标记为完成")
|
||||
elif memo['is_immediate']:
|
||||
# 即时任务执行后标记为完成
|
||||
cursor.execute('''
|
||||
UPDATE memos
|
||||
SET is_completed = 1, completed_at = ?
|
||||
WHERE id = ?
|
||||
''', (now.strftime(constants.DATETIME_FORMAT), memo_id))
|
||||
self.logger.info(f"即时任务 #{memo_id} 已完成")
|
||||
else:
|
||||
# 一次性任务,标记为完成
|
||||
cursor.execute('''
|
||||
UPDATE memos
|
||||
SET is_completed = 1, completed_at = ?
|
||||
WHERE id = ?
|
||||
''', (now.strftime(constants.DATETIME_FORMAT), memo_id))
|
||||
|
||||
self.logger.info(f"即时任务 #{memo_id} 已完成")
|
||||
else:
|
||||
# 一次性任务,标记为完成
|
||||
cursor.execute('''
|
||||
UPDATE memos
|
||||
SET is_completed = 1, completed_at = ?
|
||||
WHERE id = ?
|
||||
''', (now.strftime(constants.DATETIME_FORMAT), memo_id))
|
||||
self.logger.info(f"任务 #{memo_id} 已完成")
|
||||
|
||||
self.logger.info(f"任务 #{memo_id} 已完成")
|
||||
|
||||
conn.commit()
|
||||
conn.commit()
|
||||
except Exception:
|
||||
conn.rollback()
|
||||
raise
|
||||
|
||||
return next_reminder.strftime(constants.DATETIME_FORMAT) if next_reminder else None
|
||||
|
||||
|
||||
@@ -0,0 +1,359 @@
|
||||
"""aibox 第三轮审计 Batch G 修复回归测试(2026-08-18,审计报告 §9.4)
|
||||
|
||||
覆盖:
|
||||
- P1:非瞬时 HTTP 错误响应体截断(密钥/超大错误页不进异常栈)
|
||||
- P2:Retry-After 解析上限 cap 3600s
|
||||
- P3:fallback 档案缺 model → 跳过并告警(与缺 Key 行为对齐)
|
||||
- P5:LLMClient 死状态 _cancelled 已删除
|
||||
- P6:fallback 列表重复项去重
|
||||
- N2:drain 线程批次结束即回收(不再泄漏常驻线程)
|
||||
- N5:子代理共享父代理取消事件
|
||||
- N7:FTS5 关键词内双引号翻倍转义
|
||||
- N8:registry 流式转发受 TOOL_MAX_RESULT_LENGTH 约束
|
||||
- N10:mark_as_triggered 单事务(缺失 id 路径回归)
|
||||
- N12:persona 路径边界校验(拒绝越界路径)
|
||||
- N16:同秒归档不再覆盖(文件名加序号)
|
||||
|
||||
N3(/new 移线程池)/N13(broadcast)/N14(任务登记集合)为 async WebSocket
|
||||
路径,依赖事件循环与真实连接对象,由现有 WebSocket 测试面回归,不在此单测。
|
||||
"""
|
||||
|
||||
import os
|
||||
import queue
|
||||
import sys
|
||||
import threading
|
||||
import types
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
|
||||
# ---- 环境桩:websockets 未安装时保证 import aibox.py 不失败 ----
|
||||
if "websockets" not in sys.modules:
|
||||
sys.modules["websockets"] = types.ModuleType("websockets")
|
||||
|
||||
import aibox # noqa: E402
|
||||
from aibox import DeepSeekChat, run_full_subagent # noqa: E402
|
||||
|
||||
from llm_client import ( # noqa: E402
|
||||
DeepSeekClient,
|
||||
FailoverClient,
|
||||
_retry_after_of,
|
||||
create_llm_client,
|
||||
)
|
||||
from subsystems.db import MemoryDatabase, MemoDatabase # noqa: E402
|
||||
from subsystems.conversation_history import ConversationHistory # noqa: E402
|
||||
from tools.base import Tool # noqa: E402
|
||||
from tools.registry import ToolRegistry # noqa: E402
|
||||
from common import constants # noqa: E402
|
||||
|
||||
|
||||
# ==================== P1/P2:错误响应体截断 + Retry-After cap ====================
|
||||
|
||||
class TestP1P2LlmClient:
|
||||
def test_p1_non_transient_error_body_truncated(self):
|
||||
"""400 响应体 1000 字符 → 异常消息仅含前 200 字符"""
|
||||
c = DeepSeekClient(api_key="k", model="m", base_url="https://x")
|
||||
resp = mock.MagicMock()
|
||||
resp.status_code = 400
|
||||
resp.text = "A" * 1000
|
||||
resp.headers = {}
|
||||
with mock.patch("llm_client.requests.post", return_value=resp):
|
||||
with pytest.raises(RuntimeError) as ei:
|
||||
list(c.stream_chat([{"role": "user", "content": "hi"}]))
|
||||
msg = str(ei.value)
|
||||
assert "HTTP 400" in msg
|
||||
assert len(msg) < 300 # 200 字符正文 + 状态行前缀
|
||||
assert msg.count("A") <= 200
|
||||
|
||||
def test_p2_retry_after_capped(self):
|
||||
"""异常/恶意 Retry-After=999999 → cap 到 3600s"""
|
||||
resp = mock.MagicMock()
|
||||
resp.headers = {"Retry-After": "999999"}
|
||||
assert _retry_after_of(resp) == 3600.0
|
||||
|
||||
def test_p2_retry_after_normal_value_kept(self):
|
||||
resp = mock.MagicMock()
|
||||
resp.headers = {"Retry-After": "30"}
|
||||
assert _retry_after_of(resp) == 30.0
|
||||
|
||||
def test_p2_retry_after_invalid_returns_none(self):
|
||||
resp = mock.MagicMock()
|
||||
resp.headers = {"Retry-After": "not-a-number"}
|
||||
assert _retry_after_of(resp) is None
|
||||
|
||||
|
||||
# ==================== P3/P5/P6:fallback 链构造 ====================
|
||||
|
||||
def _g_cfg():
|
||||
"""主档案 main 的 fallback 链:缺 model 的 broken + 重复 dupe + 正常 ok2"""
|
||||
return {
|
||||
"models": {
|
||||
"main": {
|
||||
"api_url": "https://a.example/v1/chat/completions",
|
||||
"model": "m1", "api_key_env": "MAIN_KEY",
|
||||
"fallback": ["nomodel", "dupe", "dupe", "ok2"],
|
||||
},
|
||||
"nomodel": { # P3:缺 model 字段
|
||||
"api_url": "https://b.example/v1/chat/completions",
|
||||
"api_key_env": "B_KEY",
|
||||
},
|
||||
"dupe": {
|
||||
"api_url": "https://c.example/v1/chat/completions",
|
||||
"model": "c1", "api_key_env": "C_KEY",
|
||||
},
|
||||
"ok2": {
|
||||
"api_url": "https://d.example/v1/chat/completions",
|
||||
"model": "d1", "api_key_env": "D_KEY",
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
class TestP3P5P6FallbackChain:
|
||||
def test_p3_p6_fallback_missing_model_skipped_and_deduped(self):
|
||||
"""缺 model 的 fallback 跳过;重复项去重;链 = [main, dupe, ok2]"""
|
||||
env = {k: f"sk-{k.lower()}" for k in
|
||||
("MAIN_KEY", "B_KEY", "C_KEY", "D_KEY")}
|
||||
with mock.patch.dict(os.environ, env):
|
||||
c = create_llm_client(_g_cfg(), model="main")
|
||||
assert isinstance(c, FailoverClient)
|
||||
names = [cl.name for cl in c._clients]
|
||||
assert names == ["main", "dupe", "ok2"]
|
||||
|
||||
def test_p5_dead_cancel_flag_removed(self):
|
||||
"""LLMClient 不再持有只写不读的 _cancelled 死状态"""
|
||||
c = DeepSeekClient(api_key="k", model="m", base_url="https://x")
|
||||
assert not hasattr(c, "_cancelled")
|
||||
c.cancel() # 无在途响应时为安全 no-op
|
||||
|
||||
|
||||
# ==================== N2:drain 线程生命周期 ====================
|
||||
|
||||
class _EchoTool(Tool):
|
||||
def get_name(self):
|
||||
return "g_echo"
|
||||
|
||||
def get_description(self):
|
||||
return "测试用回显工具"
|
||||
|
||||
def get_parameters(self):
|
||||
return {"type": "object", "properties": {}}
|
||||
|
||||
def execute(self, **kwargs):
|
||||
return "ok"
|
||||
|
||||
|
||||
class TestN2DrainThread:
|
||||
def test_drain_thread_reclaimed_after_batch(self):
|
||||
"""_handle_tool_calls 返回后 drain 线程必须已退出(finally join)"""
|
||||
from test_once_mode import _make_ai
|
||||
ai = _make_ai()
|
||||
ai.tool_registry.register(_EchoTool(ai.logger))
|
||||
tool_calls = [{"id": "c1", "function": {"name": "g_echo", "arguments": "{}"}}]
|
||||
ai._handle_tool_calls(tool_calls, output_callback=lambda t, c: None)
|
||||
drains = [t for t in threading.enumerate() if t.name == "tool-output-drain"]
|
||||
assert drains == [] or all(not t.is_alive() for t in drains)
|
||||
|
||||
@pytest.mark.filterwarnings("ignore::pytest.PytestUnhandledThreadExceptionWarning")
|
||||
def test_drain_thread_reclaimed_on_tool_exception(self):
|
||||
"""工具执行抛异常时 drain 线程同样被回收(finally 路径)"""
|
||||
from test_once_mode import _make_ai
|
||||
|
||||
class _BoomTool(_EchoTool):
|
||||
def get_name(self):
|
||||
return "g_boom"
|
||||
|
||||
def execute(self, **kwargs):
|
||||
raise RuntimeError("boom")
|
||||
|
||||
ai = _make_ai()
|
||||
ai.tool_registry.register(_BoomTool(ai.logger))
|
||||
tool_calls = [{"id": "c1", "function": {"name": "g_boom", "arguments": "{}"}}]
|
||||
ai._handle_tool_calls(tool_calls, output_callback=lambda t, c: None)
|
||||
drains = [t for t in threading.enumerate() if t.name == "tool-output-drain"]
|
||||
assert drains == [] or all(not t.is_alive() for t in drains)
|
||||
|
||||
|
||||
# ==================== N5:子代理取消传播 ====================
|
||||
|
||||
def _fake_sub_chat():
|
||||
fake = mock.MagicMock()
|
||||
fake.think_and_respond = mock.MagicMock(return_value="任务完成报告")
|
||||
fake.stats = {}
|
||||
fake.messages = []
|
||||
fake._cancel_event = threading.Event()
|
||||
return fake
|
||||
|
||||
|
||||
class TestN5SubagentCancel:
|
||||
def test_subagent_shares_parent_cancel_event(self):
|
||||
"""cancel_event 传入 → 子代理取消事件被替换为父事件(检查点联动)"""
|
||||
fake = _fake_sub_chat()
|
||||
parent_event = threading.Event()
|
||||
with mock.patch("aibox.DeepSeekChat", return_value=fake), \
|
||||
mock.patch("aibox.create_llm_client"), \
|
||||
mock.patch("aibox.SubprocessCommandExecutor"):
|
||||
run_full_subagent("r", "t", None, cancel_event=parent_event)
|
||||
assert fake._cancel_event is parent_event
|
||||
|
||||
def test_subagent_keeps_own_event_when_not_provided(self):
|
||||
"""cancel_event=None → 兼容旧行为,子代理保留自身事件"""
|
||||
fake = _fake_sub_chat()
|
||||
own = fake._cancel_event
|
||||
with mock.patch("aibox.DeepSeekChat", return_value=fake), \
|
||||
mock.patch("aibox.create_llm_client"), \
|
||||
mock.patch("aibox.SubprocessCommandExecutor"):
|
||||
run_full_subagent("r", "t", None)
|
||||
assert fake._cancel_event is own
|
||||
|
||||
def test_subagent_tool_registered_with_parent_event(self):
|
||||
"""DeepSeekChat 注册 SubAgentTool 时注入自身 _cancel_event(N5 链路)"""
|
||||
from test_once_mode import _make_ai
|
||||
ai = _make_ai()
|
||||
tool = next((t for t in ai.tool_registry.get_all_tools()
|
||||
if t.get_name() == "create_subagent"), None)
|
||||
assert tool is not None
|
||||
assert tool.cancel_event is ai._cancel_event
|
||||
|
||||
|
||||
# ==================== N7:FTS5 引号转义 ====================
|
||||
|
||||
class TestN7FtsQuoteEscape:
|
||||
def test_keyword_with_quote_does_not_break_search(self, tmp_path):
|
||||
"""关键词含双引号 → MATCH 语法合法(翻倍转义),不抛异常"""
|
||||
db = MemoryDatabase(db_path=str(tmp_path / "mem.db"))
|
||||
db.add_memory(topic="引号话题", summary="摘要",
|
||||
content='含"引号"的记忆内容',
|
||||
start_time="2026-08-18 00:00:00",
|
||||
end_time="2026-08-18 00:00:01", importance=2)
|
||||
res = db.search_memories(['带"引号"的关键词'])
|
||||
assert isinstance(res, list) # 不抛 OperationalError 即通过
|
||||
|
||||
def test_normal_keyword_still_matches(self, tmp_path):
|
||||
"""常规关键词检索行为不回归"""
|
||||
db = MemoryDatabase(db_path=str(tmp_path / "mem.db"))
|
||||
db.add_memory(topic="普通话题", summary="天气摘要",
|
||||
content="今天天气晴朗",
|
||||
start_time="2026-08-18 00:00:00",
|
||||
end_time="2026-08-18 00:00:01", importance=2)
|
||||
res = db.search_memories(["天气"])
|
||||
assert len(res) >= 1
|
||||
|
||||
|
||||
# ==================== N8:流式转发长度上限 ====================
|
||||
|
||||
class _BigStreamTool(Tool):
|
||||
"""流式工具:发出 3 行 × 40000 字符(总 120000 > 100000 上限)"""
|
||||
|
||||
def get_name(self):
|
||||
return "g_bigstream"
|
||||
|
||||
def get_description(self):
|
||||
return "测试用超长流式工具"
|
||||
|
||||
def get_parameters(self):
|
||||
return {"type": "object", "properties": {}}
|
||||
|
||||
def execute(self, **kwargs):
|
||||
return "done"
|
||||
|
||||
def execute_streaming(self, output_queue, **kwargs):
|
||||
for _ in range(3):
|
||||
output_queue.put(("line", "x" * 40000))
|
||||
output_queue.put(("complete", "done"))
|
||||
|
||||
|
||||
class TestN8StreamingLimit:
|
||||
def test_streaming_forward_truncated_to_budget(self):
|
||||
registry = ToolRegistry()
|
||||
registry.register(_BigStreamTool())
|
||||
captured = queue.Queue()
|
||||
registry.set_output_queue(captured)
|
||||
|
||||
result = registry.execute_tool("g_bigstream")
|
||||
|
||||
lines = []
|
||||
while not captured.empty():
|
||||
msg_type, content = captured.get_nowait()
|
||||
if msg_type == "line":
|
||||
lines.append(content)
|
||||
budget = constants.TOOL_MAX_RESULT_LENGTH
|
||||
total = sum(len(ln) for ln in lines)
|
||||
# 截断标记行本身占少量额外字符(含提示文案)
|
||||
assert total <= budget + 200
|
||||
assert any("已停止转发" in ln or "已截断" in ln for ln in lines)
|
||||
assert result == "done"
|
||||
|
||||
|
||||
# ==================== N10:mark_as_triggered 单事务 ====================
|
||||
|
||||
class TestN10MarkTriggered:
|
||||
def test_missing_id_returns_none_via_transaction(self, tmp_path):
|
||||
"""事务内读取不存在行 → rollback + None(不再走独立 get_memo)"""
|
||||
db = MemoDatabase(db_path=str(tmp_path / "memo.db"))
|
||||
assert db.mark_as_triggered(99999) is None
|
||||
|
||||
def test_single_shot_memo_completes(self, tmp_path):
|
||||
"""一次性任务触发 → completed + 返回 None(无下次提醒)"""
|
||||
db = MemoDatabase(db_path=str(tmp_path / "memo.db"))
|
||||
mid = db.add_memo(title="一次性", content="c",
|
||||
reminder_time="2026-08-18 00:00:00")
|
||||
assert mid is not None
|
||||
assert db.mark_as_triggered(mid) is None
|
||||
memo = db.get_memo(mid)
|
||||
assert memo["is_completed"] == 1
|
||||
assert memo["trigger_count"] == 1
|
||||
|
||||
|
||||
# ==================== N12:persona 路径边界 ====================
|
||||
|
||||
class TestN12PersonaBoundary:
|
||||
def _tool(self):
|
||||
return aibox.ExpressEmotionTool()
|
||||
|
||||
def test_traversal_persona_rejected(self, tmp_path):
|
||||
"""persona 含 ../ 越界 → 拒绝(返回空串,不读文件)"""
|
||||
tool = self._tool()
|
||||
skills = tmp_path / "skills"
|
||||
skills.mkdir()
|
||||
with mock.patch.object(type(aibox.constants), "SKILLS_DIR",
|
||||
new_callable=mock.PropertyMock(return_value=str(skills))):
|
||||
assert tool._load_persona("../../etc/passwd", None) == ""
|
||||
assert tool._load_persona("../outside", None) == ""
|
||||
|
||||
def test_absolute_persona_rejected(self, tmp_path):
|
||||
"""persona 为绝对路径 → 解析结果不在技能目录内,拒绝"""
|
||||
tool = self._tool()
|
||||
skills = tmp_path / "skills"
|
||||
skills.mkdir()
|
||||
with mock.patch.object(type(aibox.constants), "SKILLS_DIR",
|
||||
new_callable=mock.PropertyMock(return_value=str(skills))):
|
||||
assert tool._load_persona(str(tmp_path / "evil"), None) == ""
|
||||
|
||||
def test_normal_persona_loads(self, tmp_path):
|
||||
"""合法 persona(技能目录内)正常加载"""
|
||||
tool = self._tool()
|
||||
skills = tmp_path / "skills"
|
||||
skills.mkdir()
|
||||
(skills / "ok.md").write_text("人格内容", encoding="utf-8")
|
||||
with mock.patch.object(type(aibox.constants), "SKILLS_DIR",
|
||||
new_callable=mock.PropertyMock(return_value=str(skills))):
|
||||
assert "人格内容" in tool._load_persona("ok", None)
|
||||
|
||||
|
||||
# ==================== N16:同秒归档防覆盖 ====================
|
||||
|
||||
class TestN16ArchiveCollision:
|
||||
def test_same_second_archive_no_overwrite(self, tmp_path):
|
||||
"""同一秒内两次归档 → 文件名加序号,互不覆盖"""
|
||||
hm = ConversationHistory(save_dir=str(tmp_path))
|
||||
hm.start_new_conversation(topic="t1")
|
||||
hm.add_message("user", "hi", speaker="u")
|
||||
p1 = hm.archive_current_conversation()
|
||||
|
||||
hm.start_new_conversation(topic="t2")
|
||||
hm.add_message("user", "hi2", speaker="u")
|
||||
p2 = hm.archive_current_conversation()
|
||||
|
||||
assert p1 != p2
|
||||
assert os.path.exists(p1) and os.path.exists(p2)
|
||||
+23
-2
@@ -11,7 +11,7 @@ import threading
|
||||
import time
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
from common import Logger, i18n
|
||||
from common import Logger, constants, i18n
|
||||
|
||||
from .base import Tool
|
||||
|
||||
@@ -110,6 +110,11 @@ class ToolRegistry:
|
||||
deadline = time.time() + tool_timeout
|
||||
final_result = None
|
||||
timed_out = False
|
||||
# N8:流式转发同样受 TOOL_MAX_RESULT_LENGTH 约束——此前仅非流式
|
||||
# 返回值受限,超长内容(如记忆全文)经流式通道可完整推到前端
|
||||
line_budget = constants.TOOL_MAX_RESULT_LENGTH
|
||||
forwarded = 0
|
||||
truncation_noted = False
|
||||
while thread.is_alive() or not output_queue.empty():
|
||||
if time.time() > deadline:
|
||||
self.logger.warning(
|
||||
@@ -119,9 +124,25 @@ class ToolRegistry:
|
||||
try:
|
||||
msg_type, content = output_queue.get(timeout=0.1)
|
||||
if msg_type == "line":
|
||||
piece = str(content) if content is not None else ""
|
||||
if forwarded >= line_budget:
|
||||
# 预算耗尽:停止转发,仅提示一次
|
||||
if not truncation_noted:
|
||||
truncation_noted = True
|
||||
if self.current_output_queue:
|
||||
self.current_output_queue.put(
|
||||
("line", f"\n⚠️ 流式输出超过 {line_budget} 字符上限,已停止转发剩余内容\n"))
|
||||
continue
|
||||
remaining = line_budget - forwarded
|
||||
if len(piece) > remaining:
|
||||
piece = piece[:remaining] + \
|
||||
f"\n⚠️ 流式输出超过 {line_budget} 字符上限,已截断"
|
||||
forwarded = line_budget
|
||||
else:
|
||||
forwarded += len(piece)
|
||||
# 发送到当前队列(会通过WebSocket发送)
|
||||
if self.current_output_queue:
|
||||
self.current_output_queue.put(("line", content))
|
||||
self.current_output_queue.put(("line", piece))
|
||||
elif msg_type == "complete":
|
||||
final_result = content
|
||||
# 也发送完整结果
|
||||
|
||||
Reference in new issue
Block a user