agent-Specialization/modules/memory_debug.py

"""内存调试工具：在关键路径记录内存变化，帮助定位暴涨点。

使用方式：
    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/<mode>/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