Agent/cyber_agent/core/dream.py

"""
Dream：记忆反思操作（Phase 3）

两阶段执行：
    per_trace_reflect    → 为每个有新消息的 trace 生成反思摘要，写 cognition_log
    cross_trace_integrate → 汇总各 trace 的反思摘要 + 当前记忆文件，
                             用 dream_prompt 指导 LLM 更新记忆文件

对外入口：
    run_dream(store, llm_call, memory_config, dream_scope, model=...)
"""

from __future__ import annotations

import json
import logging
import os
import re
import tempfile
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any, Awaitable, Callable, Dict, List, Optional, Tuple

from cyber_agent.core.memory import (
    MEMORY_IDENTITY_CONTEXT_KEY,
    MemoryConfig,
    MemoryPathError,
    compute_memory_identity,
    format_memory_injection,
    load_memory_files,
    normalize_memory_relative_path,
    resolve_memory_path,
)
from cyber_agent.trace.models import Trace
from cyber_agent.trace.store import FileSystemTraceStore

logger = logging.getLogger(__name__)


# ===== 默认 prompts =====

DEFAULT_REFLECT_PROMPT = """你正在回顾一次 Agent 执行中发生的事情，为你自己（作为长期身份）的记忆做反思。

请综合下面的执行过程和知识使用情况，回答：
1. 这次执行中有什么值得记住的经验？（品味、判断、策略）
2. 哪些知识的评估反映了我的判断需要调整？
3. 用户的反馈（如果有）说明了什么？

用简洁的第一人称段落写，不要逐条列点，不要重复执行细节 —— 你在沉淀"这对未来的我意味着什么"。
只输出反思内容本身，不要任何其它前缀或 markdown 标题。"""


DEFAULT_DREAM_PROMPT = """你正在整理自己的长期记忆。下面是你最近的反思摘要、以及当前各记忆文件的内容。

请决定哪些文件应该更新、内容怎么改。原则：
- 只更新真正有新见解的文件，没有变化的就不要动
- 在原有内容基础上演进，不是重写；保留仍然有效的旧内容
- 简洁、人类可读的 markdown 格式
- 新增文件必须是 MemoryConfig.files 已声明的路径（否则不会被下次加载）

**严格按以下 JSON 格式输出，不要任何其它文字**：

```json
{
  "updates": [
    {"path": "taste.md", "new_content": "完整的新文件内容"},
    {"path": "strategy.md", "new_content": "..."}
  ],
  "reasoning": "你为什么做这些更新（简短）"
}
```

如果没有任何文件需要更新，输出 `{"updates": [], "reasoning": "..."}`。"""


# ===== 数据结构 =====

@dataclass(frozen=True)
class DreamScope:
    """Dream 可见 Trace 的完整身份边界。

    uid 允许为 None，但仍按 None 做精确匹配；不会将其视为
    “不过滤”。
    """

    uid: Optional[str]
    agent_type: str
    memory_identity: str

    def __post_init__(self) -> None:
        if not self.agent_type:
            raise ValueError("DreamScope.agent_type is required")
        if not self.memory_identity:
            raise ValueError("DreamScope.memory_identity is required")

    def matches(self, trace: Optional[Trace]) -> bool:
        if trace is None:
            return False
        context = trace.context if isinstance(trace.context, dict) else {}
        return (
            trace.uid == self.uid
            and trace.agent_type == self.agent_type
            and context.get(MEMORY_IDENTITY_CONTEXT_KEY) == self.memory_identity
        )


class DreamScopeError(ValueError):
    """DreamScope 与 MemoryConfig 或 Trace 身份不一致。"""


class DreamPlanError(ValueError):
    """LLM 产生的记忆更新计划不安全或不完整。"""

@dataclass
class DreamReport:
    per_trace_summaries: Dict[str, str] = field(default_factory=dict)  # {trace_id: summary}
    updated_files: List[str] = field(default_factory=list)             # 实际写入的文件路径
    consumed_reflection_count: int = 0                                  # 本次消化了多少条 reflection
    reasoning: str = ""
    skipped_traces: List[str] = field(default_factory=list)


LLMCall = Callable[..., Awaitable[Dict[str, Any]]]


# ===== Per-trace 反思 =====

async def per_trace_reflect(
    store: FileSystemTraceStore,
    llm_call: LLMCall,
    trace_id: str,
    memory_config: MemoryConfig,
    dream_scope: DreamScope,
    model: str = "gpt-4o-mini",
) -> Optional[str]:
    """为单个 trace 生成反思摘要，写入 cognition_log，更新 reflected_at_sequence。

    Returns:
        反思摘要字符串；若 trace 没有新消息或 LLM 返回空，返回 None。
    """
    trace = await store.get_trace(trace_id)
    if not dream_scope.matches(trace):
        logger.warning(f"[Dream] trace 超出 DreamScope，拒绝反思: {trace_id}")
        return None

    start_seq = (trace.reflected_at_sequence or 0) + 1
    end_seq = trace.last_sequence
    if start_seq > end_seq:
        logger.debug(f"[Dream] trace {trace_id} 没有新消息（{start_seq} > {end_seq}）")
        return None

    all_msgs = await store.get_trace_messages(trace_id)
    new_msgs = [m for m in all_msgs if start_seq <= m.sequence <= end_seq]
    if not new_msgs:
        logger.debug(f"[Dream] trace {trace_id} 范围内无消息")
        return None

    log = await store.get_cognition_log(trace_id)
    events = log.get("events", log.get("entries", []))
    relevant_events = [
        e for e in events
        if e.get("sequence") is not None
        and start_seq <= e["sequence"] <= end_seq
        and e.get("type") in ("query", "evaluation", "extraction_pending", "extraction_committed")
    ]

    user_content = _build_reflect_input(new_msgs, relevant_events)
    prompt = memory_config.reflect_prompt or DEFAULT_REFLECT_PROMPT

    try:
        result = await llm_call(
            messages=[
                {"role": "system", "content": prompt},
                {"role": "user", "content": user_content},
            ],
            model=model,
            tools=None,
            temperature=0.5,
        )
    except Exception as e:
        logger.error(f"[Dream] per_trace_reflect LLM 调用失败 {trace_id}: {e}")
        return None

    summary = (result.get("content") or "").strip()
    if not summary:
        logger.info(f"[Dream] trace {trace_id} 反思 LLM 返回空，视为无值得记录的内容")
        # 仍然更新 reflected_at_sequence，避免下次重复扫描
        await store.update_trace(trace_id, reflected_at_sequence=end_seq)
        return None

    await store.append_cognition_event(
        trace_id=trace_id,
        event={
            "type": "reflection",
            "sequence_range": [start_seq, end_seq],
            "summary": summary,
        },
    )
    await store.update_trace(trace_id, reflected_at_sequence=end_seq)
    logger.info(f"[Dream] trace {trace_id} 反思完成，覆盖 sequence {start_seq}-{end_seq}")
    return summary


def _build_reflect_input(messages: List[Any], events: List[Dict[str, Any]]) -> str:
    """把消息和事件组织为 LLM 可读的反思输入。"""
    parts: List[str] = ["## 执行过程"]
    for m in messages:
        role = getattr(m, "role", "?")
        desc = getattr(m, "description", "") or ""
        seq = getattr(m, "sequence", "?")
        # 截断，防止单条过长
        parts.append(f"[{seq}] {role}: {desc[:500]}")

    if events:
        parts.append("\n## 知识使用与提取情况（来自 cognition_log）")
        for e in events:
            etype = e.get("type")
            if etype == "query":
                parts.append(
                    f"- [{e.get('sequence')}] query: {e.get('query', '')[:100]} → "
                    f"source_ids={e.get('source_ids', [])}"
                )
            elif etype == "evaluation":
                parts.append(
                    f"- evaluation: knowledge_id={e.get('knowledge_id')} "
                    f"result={e.get('eval_result')}"
                )
            elif etype == "extraction_pending":
                payload = e.get("payload", {})
                parts.append(
                    f"- extraction_pending ({e.get('extraction_id')}): "
                    f"{payload.get('task', '')[:80]}"
                )
            elif etype == "extraction_committed":
                parts.append(
                    f"- extraction_committed: extraction={e.get('extraction_id')} "
                    f"→ knowledge_id={e.get('knowledge_id')}"
                )
    return "\n".join(parts)


# ===== 跨 trace 整合 =====

async def cross_trace_integrate(
    store: FileSystemTraceStore,
    llm_call: LLMCall,
    memory_config: MemoryConfig,
    dream_scope: DreamScope,
    model: str = "gpt-4o",
) -> Tuple[int, List[str], str]:
    """汇总各 trace 未消化的 reflection 事件，用 LLM 更新记忆文件。

    Returns:
        (consumed_reflection_count, updated_file_paths, reasoning)
    """
    _validate_dream_scope(memory_config, dream_scope)
    all_traces = await store.list_traces(limit=1000)
    all_traces = [t for t in all_traces if dream_scope.matches(t)]

    # 收集所有未消化的 reflection 事件
    reflections: List[Tuple[str, Dict[str, Any]]] = []  # [(trace_id, event)]
    for t in all_traces:
        log = await store.get_cognition_log(t.trace_id)
        events = log.get("events", log.get("entries", []))
        for e in events:
            if e.get("type") == "reflection" and not e.get("consumed_at"):
                reflections.append((t.trace_id, e))

    if not reflections:
        logger.info("[Dream] 没有未消化的 reflection 事件")
        return 0, [], ""

    # 读当前记忆文件
    existing_files = load_memory_files(memory_config)
    user_content = _build_dream_input(reflections, existing_files, memory_config)
    prompt = memory_config.dream_prompt or DEFAULT_DREAM_PROMPT

    try:
        result = await llm_call(
            messages=[
                {"role": "system", "content": prompt},
                {"role": "user", "content": user_content},
            ],
            model=model,
            tools=None,
            temperature=0.3,
        )
    except Exception as e:
        logger.error(f"[Dream] cross_trace_integrate LLM 调用失败: {e}")
        return 0, [], ""

    raw = (result.get("content") or "").strip()
    plan = _parse_dream_output(raw)
    if plan is None:
        logger.error(f"[Dream] LLM 输出无法解析为 JSON 计划，原文: {raw[:500]}")
        return 0, [], ""

    try:
        prepared_updates = _prepare_memory_updates(memory_config, plan)
    except (DreamPlanError, MemoryPathError) as exc:
        logger.error(f"[Dream] 拒绝不安全的记忆更新计划: {exc}")
        return 0, [], ""

    # LLM 返回后再次检查 Trace 身份，防止在长耗时调用期间边界变化。
    for trace_id in {trace_id for trace_id, _ in reflections}:
        if not dream_scope.matches(await store.get_trace(trace_id)):
            raise DreamScopeError(
                f"trace left DreamScope before memory update: {trace_id}"
            )

    updated_paths = _transactional_write_memory_updates(
        memory_config,
        prepared_updates,
    )

    # 只有全部记忆文件原子替换成功后，才标记 reflection 已消化。
    consumed_at = datetime.now().isoformat()
    reflections_by_trace: Dict[str, List[Dict[str, Any]]] = {}
    for trace_id, event in reflections:
        reflections_by_trace.setdefault(trace_id, []).append(event)
    consumed_count = 0
    for trace_id, trace_reflections in reflections_by_trace.items():
        consumed_count += await store.mark_reflections_consumed(
            trace_id,
            trace_reflections,
            consumed_at,
        )

    reasoning = plan.get("reasoning", "")
    return consumed_count, updated_paths, reasoning


def _validate_dream_scope(
    memory_config: MemoryConfig,
    dream_scope: DreamScope,
) -> None:
    expected_identity = compute_memory_identity(memory_config)
    if dream_scope.memory_identity != expected_identity:
        raise DreamScopeError(
            "DreamScope.memory_identity does not match the active MemoryConfig"
        )


def _prepare_memory_updates(
    memory_config: MemoryConfig,
    plan: Dict[str, Any],
) -> List[Tuple[str, Path, str]]:
    updates = plan.get("updates")
    if not isinstance(updates, list):
        raise DreamPlanError("updates must be a list")
    prepared: List[Tuple[str, Path, str]] = []
    seen: set[str] = set()
    for index, update in enumerate(updates):
        if not isinstance(update, dict):
            raise DreamPlanError(f"updates[{index}] must be an object")
        rel_path = update.get("path")
        new_content = update.get("new_content")
        if not isinstance(rel_path, str) or not isinstance(new_content, str):
            raise DreamPlanError(
                f"updates[{index}] requires string path and new_content"
            )
        normalized = normalize_memory_relative_path(rel_path)
        if normalized != rel_path:
            raise DreamPlanError(f"path is not normalized: {rel_path!r}")
        if normalized in seen:
            raise DreamPlanError(f"duplicate update path: {normalized}")
        seen.add(normalized)
        prepared.append(
            (normalized, resolve_memory_path(memory_config, normalized), new_content)
        )
    return prepared


def _atomic_write_text(target: Path, content: str) -> None:
    """在同目录写临时文件、fsync，然后原子替换目标。"""
    tmp_path: Optional[Path] = None
    try:
        with tempfile.NamedTemporaryFile(
            mode="w",
            encoding="utf-8",
            dir=str(target.parent),
            prefix=f".{target.name}.",
            suffix=".tmp",
            delete=False,
        ) as tmp:
            tmp.write(content)
            tmp.flush()
            os.fsync(tmp.fileno())
            tmp_path = Path(tmp.name)
        os.replace(tmp_path, target)
        tmp_path = None
        try:
            directory_fd = os.open(str(target.parent), os.O_RDONLY)
            try:
                os.fsync(directory_fd)
            finally:
                os.close(directory_fd)
        except OSError:
            # 某些文件系统不支持对目录 fsync；文件本身仍已 fsync + replace。
            logger.debug(f"[Dream] 目录 fsync 不可用: {target.parent}")
    finally:
        if tmp_path is not None:
            try:
                tmp_path.unlink(missing_ok=True)
            except OSError:
                logger.warning(f"[Dream] 清理临时文件失败: {tmp_path}")


def _stage_text(target: Path, content: str) -> Path:
    """Durably stage one replacement next to its target without publishing it."""
    target.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.NamedTemporaryFile(
        mode="w",
        encoding="utf-8",
        dir=str(target.parent),
        prefix=f".{target.name}.",
        suffix=".dream-stage",
        delete=False,
    ) as staged:
        staged.write(content)
        staged.flush()
        os.fsync(staged.fileno())
        return Path(staged.name)


def _transactional_write_memory_updates(
    memory_config: MemoryConfig,
    prepared_updates: List[Tuple[str, Path, str]],
) -> List[str]:
    """Publish a multi-file Dream plan as one recoverable in-process unit.

    Filesystems do not provide a native multi-file rename transaction.  We
    therefore stage every new value first and keep every old value until all
    replacements succeed.  A later replacement failure rolls already-published
    files back before the exception escapes, so reflections remain retryable
    without exposing a partially committed plan.
    """
    staged_updates: List[Tuple[str, Path, str, Path]] = []
    old_values: Dict[Path, Optional[str]] = {}
    committed: List[Path] = []
    try:
        for rel_path, target, new_content in prepared_updates:
            target.parent.mkdir(parents=True, exist_ok=True)
            checked_target = resolve_memory_path(memory_config, rel_path)
            if checked_target != target:
                raise MemoryPathError(
                    f"memory target changed before staging: {rel_path}"
                )
            old_values[target] = (
                target.read_text(encoding="utf-8") if target.exists() else None
            )
            staged_updates.append(
                (rel_path, target, new_content, _stage_text(target, new_content))
            )

        for rel_path, target, new_content, staged_path in staged_updates:
            checked_target = resolve_memory_path(memory_config, rel_path)
            if checked_target != target:
                raise MemoryPathError(
                    f"memory target changed before commit: {rel_path}"
                )
            os.replace(staged_path, target)
            committed.append(target)
            try:
                directory_fd = os.open(str(target.parent), os.O_RDONLY)
                try:
                    os.fsync(directory_fd)
                finally:
                    os.close(directory_fd)
            except OSError:
                logger.debug("[Dream] 目录 fsync 不可用: %s", target.parent)
            logger.info(
                "[Dream] 已更新记忆文件: %s (%d chars)",
                rel_path,
                len(new_content),
            )
        return [rel_path for rel_path, *_ in staged_updates]
    except Exception as commit_error:
        rollback_errors: List[str] = []
        for target in reversed(committed):
            try:
                old_content = old_values[target]
                if old_content is None:
                    target.unlink(missing_ok=True)
                else:
                    _atomic_write_text(target, old_content)
            except Exception as rollback_error:  # pragma: no cover - fatal FS failure
                rollback_errors.append(f"{target}: {rollback_error}")
        if rollback_errors:
            raise RuntimeError(
                "Dream update failed and rollback was incomplete: "
                + "; ".join(rollback_errors)
            ) from commit_error
        raise
    finally:
        for _rel_path, _target, _content, staged_path in staged_updates:
            try:
                staged_path.unlink(missing_ok=True)
            except OSError:
                logger.warning("[Dream] 清理 staged 文件失败: %s", staged_path)


def _build_dream_input(
    reflections: List[Tuple[str, Dict[str, Any]]],
    existing_files: List[Tuple[str, str, str]],
    memory_config: MemoryConfig,
) -> str:
    """为 dream prompt 准备输入：反思摘要汇总 + 当前记忆文件 + 允许的文件路径。"""
    parts: List[str] = ["## 最近的反思摘要\n"]
    for trace_id, e in reflections:
        seq_range = e.get("sequence_range", [None, None])
        parts.append(
            f"### trace {trace_id} (messages {seq_range[0]}-{seq_range[1]})\n"
            f"{e.get('summary', '')}\n"
        )

    parts.append("\n## 当前记忆文件\n")
    if existing_files:
        parts.append(format_memory_injection(existing_files))
    else:
        parts.append("（暂无记忆文件）")

    if memory_config.files:
        parts.append("\n## 允许更新/新增的文件路径\n")
        for key, purpose in memory_config.files.items():
            parts.append(f"- `{key}`" + (f" — {purpose}" if purpose else ""))

    return "\n".join(parts)


def _parse_dream_output(raw: str) -> Optional[Dict[str, Any]]:
    """解析 LLM 的 JSON 计划输出。容忍 ```json ... ``` 包裹。"""
    stripped = raw.strip()
    # 去除 markdown 代码块包裹
    m = re.match(r"^```(?:json)?\s*(.*?)\s*```$", stripped, re.DOTALL)
    if m:
        stripped = m.group(1).strip()
    try:
        data = json.loads(stripped)
    except json.JSONDecodeError:
        return None
    if not isinstance(data, dict) or "updates" not in data:
        return None
    return data


# ===== 顶层入口 =====

async def run_dream(
    store: FileSystemTraceStore,
    llm_call: LLMCall,
    memory_config: MemoryConfig,
    dream_scope: DreamScope,
    reflect_model: str = "gpt-4o-mini",
    dream_model: str = "gpt-4o",
) -> DreamReport:
    """执行完整的 dream 流程：per_trace_reflect → cross_trace_integrate。

    Args:
        dream_scope: 必填的 uid + agent_type + memory_identity 精确边界
        reflect_model: per-trace 反思用的模型（轻量模型即可）
        dream_model:   跨 trace 整合用的模型（需要更强推理能力）
    """
    report = DreamReport()

    if not memory_config.base_path:
        logger.warning("[Dream] memory_config.base_path 未配置，跳过")
        return report
    _validate_dream_scope(memory_config, dream_scope)

    # Phase 1: per-trace reflect
    all_traces = await store.list_traces(limit=1000)
    all_traces = [t for t in all_traces if dream_scope.matches(t)]

    for t in all_traces:
        if (t.reflected_at_sequence or 0) >= t.last_sequence:
            continue
        try:
            summary = await per_trace_reflect(
                store,
                llm_call,
                t.trace_id,
                memory_config,
                dream_scope,
                model=reflect_model,
            )
            if summary:
                report.per_trace_summaries[t.trace_id] = summary
        except Exception as e:
            logger.error(f"[Dream] per_trace_reflect 异常 {t.trace_id}: {e}")
            report.skipped_traces.append(t.trace_id)

    # Phase 2: cross-trace integrate
    try:
        consumed, updated, reasoning = await cross_trace_integrate(
            store, llm_call, memory_config,
            dream_scope=dream_scope, model=dream_model,
        )
        report.consumed_reflection_count = consumed
        report.updated_files = updated
        report.reasoning = reasoning
    except Exception as e:
        logger.error(f"[Dream] cross_trace_integrate 异常: {e}")

    return report