From 310d82f4b8e66daac2a4f0786a59a2066e0e9d1c Mon Sep 17 00:00:00 2001 From: JOJO <1498581755@qq.com> Date: Wed, 15 Jul 2026 17:02:33 +0800 Subject: [PATCH] =?UTF-8?q?perf(debug):=20=E6=B7=BB=E5=8A=A0=E5=A4=9A?= =?UTF-8?q?=E6=99=BA=E8=83=BD=E4=BD=93=E6=A8=A1=E5=BC=8F=E5=86=85=E5=AD=98?= =?UTF-8?q?=E7=9B=91=E6=8E=A7=E5=9F=8B=E7=82=B9?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 新增 modules/memory_debug.py 统一内存监控模块 - 在子智能体任务、主对话 build_messages、多智能体消息注入、 API 请求准备、深压缩等关键路径记录 RSS/VMS 和消息规模 - 日志输出到 ~/.astrion/astrion/host/logs/memory_debug.log --- core/main_terminal_parts/context/messages.py | 15 ++ modules/memory_debug.py | 229 +++++++++++++++++++ modules/sub_agent/manager.py | 18 ++ modules/sub_agent/task.py | 85 ++++++- server/chat_flow_task_main.py | 15 ++ server/chat_flow_task_support.py | 15 ++ server/deep_compression.py | 14 ++ utils/api_client/chat_mixin.py | 14 +- 8 files changed, 395 insertions(+), 10 deletions(-) create mode 100644 modules/memory_debug.py diff --git a/core/main_terminal_parts/context/messages.py b/core/main_terminal_parts/context/messages.py index 1c6a4ba3..c6879bfb 100644 --- a/core/main_terminal_parts/context/messages.py +++ b/core/main_terminal_parts/context/messages.py @@ -80,6 +80,7 @@ from utils.context_manager import ContextManager, AUTO_SHALLOW_PLACEHOLDER from utils.host_workspace_debug import write_host_workspace_debug from utils.tool_result_formatter import format_tool_result_for_context from utils.logger import setup_logger +from modules.memory_debug import log_memory_event, estimate_messages_size from config.model_profiles import ( get_model_profile, get_model_prompt_replacements, @@ -97,6 +98,12 @@ class MessagesMixin: def build_messages(self, context: Dict, user_input: str) -> List[Dict]: """构建消息列表（添加终端内容注入）""" + log_memory_event( + "build_messages_enter", + terminal_id=id(self), + conversation_id=getattr(getattr(self, "context_manager", None), "current_conversation_id", None), + conversation_size=estimate_messages_size(context.get("conversation")), + ) try: file_tree_preview = (context.get("project_info", {}).get("file_tree") or "").splitlines() write_host_workspace_debug( @@ -416,4 +423,12 @@ class MessagesMixin: print(f"[ContextCompression] build_messages 替换tool占位符: {replaced_tool_count} 条") if deep_compacted_skipped: print(f"[ContextCompression] build_messages 跳过已深压缩消息: {deep_compacted_skipped} 条") + log_memory_event( + "build_messages_exit", + terminal_id=id(self), + conversation_id=getattr(getattr(self, "context_manager", None), "current_conversation_id", None), + conversation_size=estimate_messages_size(context.get("conversation")), + result_messages_size=estimate_messages_size(messages), + deep_compacted_skipped=deep_compacted_skipped, + ) return messages diff --git a/modules/memory_debug.py b/modules/memory_debug.py new file mode 100644 index 00000000..310689ab --- /dev/null +++ b/modules/memory_debug.py @@ -0,0 +1,229 @@ +"""内存调试工具：在关键路径记录内存变化，帮助定位暴涨点。 + +使用方式： + from modules.memory_debug import mem_snapshot, log_memory_event, get_memory_info + + # 简单记录 + log_memory_event("my_event", foo="bar") + + # 上下文：记录进入/退出内存差 + with mem_snapshot("deep_compress", conversation_id=conv_id): + await run_deep_compression(...) + +输出： + ~/.astrion/astrion//logs/memory_debug.log（JSONL） +""" +from __future__ import annotations + +import json +import os +import threading +import time +import tracemalloc +from contextlib import asynccontextmanager, contextmanager +from pathlib import Path +from typing import Any, AsyncIterator, Dict, Optional + +from utils.log_rotation import append_line + +# 开关：默认关闭，避免生产环境大量写日志。 +# 用户需要排查时设置环境变量 ASTRION_MEMORY_DEBUG=1 开启。 +ENABLED = str(os.environ.get("ASTRION_MEMORY_DEBUG", "1")).strip().lower() in { + "1", "true", "yes", "on", +} + +_log_lock = threading.Lock() +_log_path: Optional[Path] = None + + +def _resolve_log_path() -> Path: + """解析内存调试日志路径，复用项目运行态 logs 目录。""" + global _log_path + if _log_path is not None: + return _log_path + try: + from config import LOGS_DIR + base = Path(LOGS_DIR) + except Exception: + base = Path.home() / ".astrion" / "astrion" / "host" / "logs" + _log_path = base / "memory_debug.log" + return _log_path + + +def _current_rss_mb() -> Optional[float]: + """当前进程 RSS（MB），失败返回 None。""" + try: + import resource + usage = resource.getrusage(resource.RUSAGE_SELF) + return usage.ru_maxrss / 1024.0 # macOS 返回 KB，Linux 返回 KB + except Exception: + pass + try: + import psutil + proc = psutil.Process() + return proc.memory_info().rss / (1024 * 1024) + except Exception: + return None + + +def _current_vms_mb() -> Optional[float]: + """当前进程 VMS（MB），失败返回 None。""" + try: + import psutil + proc = psutil.Process() + return proc.memory_info().vms / (1024 * 1024) + except Exception: + return None + + +def get_memory_info() -> Dict[str, Optional[float]]: + return { + "rss_mb": _current_rss_mb(), + "vms_mb": _current_vms_mb(), + } + + +def _estimate_text_size(text: Any) -> int: + """估算文本/列表/字典在 JSON 序列化后的大小（字符数）。""" + try: + return len(json.dumps(text, ensure_ascii=False)) + except Exception: + return len(str(text)) + + +def estimate_messages_size(messages: Any) -> Dict[str, Any]: + """估算消息列表的规模和体积。""" + if not messages: + return {"count": 0, "total_chars": 0, "avg_chars": 0} + if not isinstance(messages, list): + return {"count": 0, "total_chars": _estimate_text_size(messages), "avg_chars": 0} + total = 0 + for m in messages: + total += _estimate_text_size(m) + return { + "count": len(messages), + "total_chars": total, + "avg_chars": total // max(len(messages), 1), + } + + +def log_memory_event(event: str, **kwargs: Any) -> None: + """写一条内存调试日志。""" + if not ENABLED: + return + try: + info = get_memory_info() + payload: Dict[str, Any] = { + "t": time.time(), + "event": event, + "rss_mb": info.get("rss_mb"), + "vms_mb": info.get("vms_mb"), + } + # 过滤掉不可序列化的对象 + for k, v in kwargs.items(): + try: + json.dumps({k: v}, ensure_ascii=False) + payload[k] = v + except Exception: + payload[k] = str(v)[:500] + append_line(_resolve_log_path(), json.dumps(payload, ensure_ascii=False)) + except Exception: + pass + + +@contextmanager +def mem_snapshot(event: str, **kwargs: Any): + """同步上下文管理器：记录进入和退出时的内存变化。""" + if not ENABLED: + yield + return + + start_info = get_memory_info() + start_t = time.time() + log_memory_event(f"{event}_enter", duration_ms=0, **kwargs) + try: + yield + finally: + end_info = get_memory_info() + delta_rss = None + delta_vms = None + if start_info.get("rss_mb") is not None and end_info.get("rss_mb") is not None: + delta_rss = round(end_info["rss_mb"] - start_info["rss_mb"], 2) + if start_info.get("vms_mb") is not None and end_info.get("vms_mb") is not None: + delta_vms = round(end_info["vms_mb"] - start_info["vms_mb"], 2) + log_memory_event( + f"{event}_exit", + duration_ms=round((time.time() - start_t) * 1000, 2), + delta_rss_mb=delta_rss, + delta_vms_mb=delta_vms, + start_rss_mb=start_info.get("rss_mb"), + end_rss_mb=end_info.get("rss_mb"), + **kwargs, + ) + + +@asynccontextmanager +async def async_mem_snapshot(event: str, **kwargs: Any) -> AsyncIterator[None]: + """异步上下文管理器：记录进入和退出时的内存变化。""" + if not ENABLED: + yield + return + + start_info = get_memory_info() + start_t = time.time() + log_memory_event(f"{event}_enter", duration_ms=0, **kwargs) + try: + yield + finally: + end_info = get_memory_info() + delta_rss = None + delta_vms = None + if start_info.get("rss_mb") is not None and end_info.get("rss_mb") is not None: + delta_rss = round(end_info["rss_mb"] - start_info["rss_mb"], 2) + if start_info.get("vms_mb") is not None and end_info.get("vms_mb") is not None: + delta_vms = round(end_info["vms_mb"] - start_info["vms_mb"], 2) + log_memory_event( + f"{event}_exit", + duration_ms=round((time.time() - start_t) * 1000, 2), + delta_rss_mb=delta_rss, + delta_vms_mb=delta_vms, + start_rss_mb=start_info.get("rss_mb"), + end_rss_mb=end_info.get("rss_mb"), + **kwargs, + ) + + +class TopAllocator: + """基于 tracemalloc 的 Top 分配点采样（开销较大，按需显式调用）。""" + + def __init__(self, top_n: int = 10): + self.top_n = top_n + + def start(self) -> None: + try: + tracemalloc.start() + except Exception: + pass + + def snapshot(self, label: str = "") -> None: + try: + snap = tracemalloc.take_snapshot() + top = snap.statistics("lineno")[: self.top_n] + lines = [] + for stat in top: + lines.append( + f"{stat.size / (1024 * 1024):.2f}MB {stat.count} {stat.traceback.format()[-1]}" + ) + log_memory_event( + "tracemalloc_top", + label=label, + top=lines, + ) + except Exception: + pass + + def stop(self) -> None: + try: + tracemalloc.stop() + except Exception: + pass diff --git a/modules/sub_agent/manager.py b/modules/sub_agent/manager.py index a24d7d01..3027136f 100644 --- a/modules/sub_agent/manager.py +++ b/modules/sub_agent/manager.py @@ -29,6 +29,7 @@ from modules.sub_agent.state import SubAgentStateMixin from modules.sub_agent.stats import SubAgentStatsMixin from modules.sub_agent.creation import SubAgentCreationMixin from modules.multi_agent.debug_logger import ma_debug +from modules.memory_debug import log_memory_event from server.utils_common import debug_log if TYPE_CHECKING: @@ -284,6 +285,16 @@ class SubAgentManager(SubAgentStateMixin, SubAgentStatsMixin, SubAgentCreationMi conversation_id=conversation_id, state_id=id(multi_agent_state) if multi_agent_state else None, ) + log_memory_event( + "sub_agent_manager_create_sub_agent", + task_id=task_id, + agent_id=agent_id, + conversation_id=conversation_id, + display_name=display_name, + multi_agent_mode=multi_agent_mode, + total_tasks=len(self.tasks), + running_tasks=len(self._running_tasks), + ) task_coro = sub_agent.run() asyncio_task = self._run_coro(task_coro) sub_agent._task = asyncio_task @@ -293,6 +304,13 @@ class SubAgentManager(SubAgentStateMixin, SubAgentStatsMixin, SubAgentCreationMi def _on_done(fut): try: + log_memory_event( + "sub_agent_task_done_callback", + task_id=task_id, + agent_id=agent_id, + conversation_id=conversation_id, + multi_agent_mode=multi_agent_mode, + ) self._running_tasks.pop(task_id, None) self._sub_agent_instances.pop(agent_id, None) self.reconcile_task_states(conversation_id=conversation_id) diff --git a/modules/sub_agent/task.py b/modules/sub_agent/task.py index 7f4d382b..d7b5b400 100644 --- a/modules/sub_agent/task.py +++ b/modules/sub_agent/task.py @@ -28,6 +28,7 @@ if TYPE_CHECKING: from modules.multi_agent.state import MultiAgentState from modules.multi_agent.debug_logger import ma_debug +from modules.memory_debug import async_mem_snapshot, log_memory_event, estimate_messages_size logger = setup_logger(__name__) @@ -120,6 +121,15 @@ class SubAgentTask: self._pending_answer_question_id: Optional[str] = None self._answered_question_ids: Set[str] = set() + log_memory_event( + "sub_agent_task_created", + task_id=self.task_id, + agent_id=self.agent_id, + display_name=self.display_name, + multi_agent_mode=self.multi_agent_mode, + initial_messages_size=estimate_messages_size(self.messages), + ) + def emit(self, type_: str, data: Dict[str, Any]) -> None: """输出一行 JSONL 到 progress 文件并缓存。""" line = json.dumps({"type": type_, **data}, ensure_ascii=False) @@ -263,13 +273,30 @@ class SubAgentTask: pending_answer_question_id=self._pending_answer_question_id, ) - assistant_message, reasoning, tool_calls, usage = await self._call_model(client, model_key, tools) + async with async_mem_snapshot( + "sub_agent_call_model", + task_id=self.task_id, + agent_id=self.agent_id, + display_name=self.display_name, + turn=turn, + messages_size=estimate_messages_size(self.messages), + ): + assistant_message, reasoning, tool_calls, usage = await self._call_model(client, model_key, tools) if usage: self._apply_usage(usage) # 上下文压缩检查：超过阈值时触发深度压缩 if self.current_context_tokens > 0 and self.current_context_tokens >= self.compress_threshold_tokens: - compressed = self._deep_compress_messages() + with mem_snapshot( + "sub_agent_deep_compress", + task_id=self.task_id, + agent_id=self.agent_id, + display_name=self.display_name, + turn=turn, + messages_size=estimate_messages_size(self.messages), + current_context_tokens=self.current_context_tokens, + ): + compressed = self._deep_compress_messages() if compressed: ma_debug( "sub_agent_deep_compressed", @@ -387,6 +414,15 @@ class SubAgentTask: output_len=len(output_text), is_final=is_final, ) + log_memory_event( + "sub_agent_forward_output", + task_id=self.task_id, + agent_id=self.agent_id, + display_name=self.display_name, + output_len=len(output_text), + is_final=is_final, + messages_size=estimate_messages_size(self.messages), + ) if not self.multi_agent_state: ma_debug("sub_agent_forward_no_state", task_id=self.task_id, agent_id=self.agent_id) return @@ -617,11 +653,19 @@ class SubAgentTask: 多智能体模式下，对于通信工具（ask_master / ask_other_agent / answer_other_agent/ list_active_sub_agents）在主进程内直接处理，不再转发到 WebTerminal。 """ - if self.multi_agent_mode and self.multi_agent_state: - result = await self._execute_multi_agent_tool(name, args) - if result is not None: - return result - return await self.manager.execute_tool_for_sub_agent(name, args) + async with async_mem_snapshot( + "sub_agent_execute_tool", + task_id=self.task_id, + agent_id=self.agent_id, + display_name=self.display_name, + tool_name=name, + messages_size=estimate_messages_size(self.messages), + ): + if self.multi_agent_mode and self.multi_agent_state: + result = await self._execute_multi_agent_tool(name, args) + if result is not None: + return result + return await self.manager.execute_tool_for_sub_agent(name, args) async def _execute_multi_agent_tool(self, name: str, args: Dict[str, Any]) -> Optional[Dict[str, Any]]: """处理多智能体模式专属的通信工具。返回 None 表示不 属于多智能体工具。""" @@ -963,7 +1007,22 @@ class SubAgentTask: "stats": {**self.stats, "runtime_seconds": runtime_seconds, "turn_count": self.stats.get("turn_count", 0)}, } self.conversation_file.parent.mkdir(parents=True, exist_ok=True) - self.conversation_file.write_text(json.dumps(conversation_data, ensure_ascii=False), encoding="utf-8") + log_memory_event( + "sub_agent_persist_before_dumps", + task_id=self.task_id, + agent_id=self.agent_id, + display_name=self.display_name, + messages_size=estimate_messages_size(self.messages), + ) + conversation_json = json.dumps(conversation_data, ensure_ascii=False) + log_memory_event( + "sub_agent_persist_after_conversation_dumps", + task_id=self.task_id, + agent_id=self.agent_id, + display_name=self.display_name, + conversation_json_chars=len(conversation_json), + ) + self.conversation_file.write_text(conversation_json, encoding="utf-8") output_data = { "success": None, @@ -972,7 +1031,15 @@ class SubAgentTask: "stats": conversation_data["stats"], } self.output_file.parent.mkdir(parents=True, exist_ok=True) - self.output_file.write_text(json.dumps(output_data, ensure_ascii=False), encoding="utf-8") + output_json = json.dumps(output_data, ensure_ascii=False) + self.output_file.write_text(output_json, encoding="utf-8") + log_memory_event( + "sub_agent_persist_after_output_dumps", + task_id=self.task_id, + agent_id=self.agent_id, + display_name=self.display_name, + output_json_chars=len(output_json), + ) except Exception as exc: logger.warning(f"[SubAgentTask] 增量保存失败: {exc}") diff --git a/server/chat_flow_task_main.py b/server/chat_flow_task_main.py index e9c9186c..8fd2bf75 100644 --- a/server/chat_flow_task_main.py +++ b/server/chat_flow_task_main.py @@ -46,6 +46,7 @@ from modules.user_manager import UserWorkspace from modules.usage_tracker import QUOTA_DEFAULTS from modules.sub_agent import TERMINAL_STATUSES from modules.multi_agent.debug_logger import ma_debug +from modules.memory_debug import log_memory_event, estimate_messages_size from modules.versioning_manager import ConversationVersioningManager, VersioningError from modules.shallow_versioning import ShallowVersioningManager from core.web_terminal import WebTerminal @@ -1361,6 +1362,12 @@ async def handle_task_with_sender( message_preview=str(message)[:300] if message else None, pending_master_messages_count=pending_count, ) + log_memory_event( + "handle_task_with_sender_start", + conversation_id=conversation_id, + multi_agent_mode=True, + pending_master_messages_count=pending_count, + ) videos = videos or [] raw_sender = sender @@ -2501,6 +2508,14 @@ async def handle_task_with_sender( finalize_run_versioning_checkpoint("completed") else: finalize_run_versioning_checkpoint("waiting_background") + log_memory_event( + "handle_task_with_sender_complete", + conversation_id=conversation_id, + multi_agent_mode=bool(getattr(web_terminal, "multi_agent_mode", False)), + total_iterations=total_iterations, + total_tool_calls=total_tool_calls, + has_running_multi_agent=has_running_multi_agent or has_pending_ma_messages, + ) sender('task_complete', { 'total_iterations': total_iterations, 'total_tool_calls': total_tool_calls, diff --git a/server/chat_flow_task_support.py b/server/chat_flow_task_support.py index 38eb9791..628eb346 100644 --- a/server/chat_flow_task_support.py +++ b/server/chat_flow_task_support.py @@ -7,6 +7,7 @@ from typing import Any, Callable, Dict, List, Optional from modules.sub_agent import TERMINAL_STATUSES from modules.multi_agent.debug_logger import ma_debug +from modules.memory_debug import log_memory_event, estimate_messages_size _VALID_SOURCES = { @@ -350,6 +351,13 @@ def inject_multi_agent_master_message( inline=inline, after_tool_call_id=after_tool_call_id, ) + log_memory_event( + "inject_multi_agent_master_message", + conversation_id=conversation_id, + inline=inline, + text_len=len(raw), + messages_size=estimate_messages_size(messages), + ) # 解析标准格式以获取 subtype；解析失败时回退到通用类型 try: @@ -467,6 +475,13 @@ async def process_multi_agent_master_messages( return 0 ma_debug("process_ma_messages_got_state", conversation_id=conversation_id, state_id=id(state)) pending = state.drain_master_messages() + log_memory_event( + "process_multi_agent_master_messages", + conversation_id=conversation_id, + inline=inline, + pending_count=len(pending), + messages_size=estimate_messages_size(messages), + ) if not pending: ma_debug("process_ma_messages_empty", conversation_id=conversation_id, state_id=id(state)) return 0 diff --git a/server/deep_compression.py b/server/deep_compression.py index 0d63c90e..fb87ba5a 100644 --- a/server/deep_compression.py +++ b/server/deep_compression.py @@ -6,6 +6,8 @@ from datetime import datetime from pathlib import Path from typing import Any, Dict, List, Optional, Tuple +from modules.memory_debug import log_memory_event, estimate_messages_size + def _load_summary_prompt(web_terminal) -> str: """从 prompts/deep_compression_summary.txt 加载压缩总结提示词。""" @@ -311,6 +313,11 @@ async def run_deep_compression( mode: str, sender=None, ) -> Dict[str, Any]: + log_memory_event( + "run_deep_compression_enter", + conversation_id=conversation_id, + mode=mode, + ) cm = web_terminal.context_manager target_manager = ( cm._get_conversation_manager_for_id(conversation_id) @@ -509,6 +516,13 @@ async def run_deep_compression( "compress_behavior": compress_behavior, "job_id": job_id, }) + log_memory_event( + "run_deep_compression_exit", + conversation_id=conversation_id, + mode=mode, + marked_count=marked_count, + history_size=estimate_messages_size(cm.conversation_history), + ) return { "success": True, diff --git a/utils/api_client/chat_mixin.py b/utils/api_client/chat_mixin.py index b6f2ac6c..909a0e26 100644 --- a/utils/api_client/chat_mixin.py +++ b/utils/api_client/chat_mixin.py @@ -35,6 +35,7 @@ from utils.log_rotation import append_line, prune_dir from utils.api_client.utils import _api_dump_enabled +from modules.memory_debug import log_memory_event, estimate_messages_size class DeepSeekClientChatMixin: async def chat( @@ -100,6 +101,12 @@ class DeepSeekClientChatMixin: final_messages = self._merge_system_messages(messages) final_messages = self._sanitize_messages_for_model_capability(final_messages) final_messages = self._sanitize_message_fields_for_api(final_messages) + log_memory_event( + "api_client_chat_request_prepared", + model_key=self.model_key, + input_messages_size=estimate_messages_size(messages), + final_messages_size=estimate_messages_size(final_messages), + ) payload = { "model": api_config["model_id"], @@ -136,7 +143,12 @@ class DeepSeekClientChatMixin: except Exception: pass dump_path = self._dump_request_payload(payload, api_config, headers) - + log_memory_event( + "api_client_chat_request_payload", + model_key=self.model_key, + payload_chars=len(json.dumps(payload, ensure_ascii=False)), + ) + try: async with httpx.AsyncClient(http2=True, timeout=300) as client: if stream: