""" core/agent/compaction.py ======================== 🌟 pi 上下文压缩算法的 Python 1:1 移植(harness 原版算法) 对照 pi-main 源码(逐函数对应): packages/agent/src/harness/compaction/compaction.ts DEFAULT_COMPACTION_SETTINGS -> CompactionSettings / DEFAULT_COMPACTION_SETTINGS shouldCompact -> 见 context.py: should_compact(同一公式) estimateTokens / estimateContextTokens -> context.py(usage 锚定,CJK 感知见说明) findValidCutPoints / findCutPoint -> find_valid_cut_points / find_cut_point findTurnStartIndex -> find_turn_start prepareCompaction -> prepare_compaction SUMMARIZATION_SYSTEM_PROMPT -> 同名(逐字移植) SUMMARIZATION_PROMPT -> 同名(逐字移植) UPDATE_SUMMARIZATION_PROMPT -> 同名(逐字移植,迭代更新用) TURN_PREFIX_SUMMARIZATION_PROMPT -> 同名(逐字移植,断轮前缀用) generateSummaryWithUsage -> generate_summary(maxTokens = 0.8×reserve) generateTurnPrefixSummary -> generate_turn_prefix_summary(0.5×reserve) compact -> compact_context(断轮双摘要 + 拼接格式 1:1) packages/agent/src/harness/compaction/utils.ts serializeConversation -> serialize_conversation(1:1,含 2000 字符截断) extractFileOpsFromMessage -> extract_file_ops_from_message computeFileLists -> compute_file_lists formatFileOperations -> format_file_operations TOOL_RESULT_MAX_CHARS = 2000 -> 同名常量 摘要 LLM 调用由上层注入: summarize_fn(prompt_text: str, system_prompt: str, max_tokens: int) -> str (pi 里是 models.completeSimple + retry;haocode 用 OpenAI 客户端非流式调用, 由 llm_engine.AgentWorker 实现并注入。) 已声明的偏差(仅 2 处,见 context.py 头注): 1. 单条消息 token 估算用 CJK 感知启发式(pi 是 chars/4)——对中文会话更安全 2. 摘要 LLM 调用失败时降级为机械摘录(pi 返回 CompactionError)——桌面应用优先不丢上下文 其余全部 1:1:触发公式、usage 锚定、token 预算切点、有效切点规则、断轮双摘要、 迭代式 previousSummary 更新、摘要提示词逐字、文件操作附录、拼接格式。 """ from __future__ import annotations import json from dataclasses import dataclass, field from typing import Callable, List, Optional, Set, Tuple from .types import AgentMessage, ModelConfig # ====================================================================== # 压缩设置 —— 1:1 对照 DEFAULT_COMPACTION_SETTINGS # ====================================================================== @dataclass class CompactionSettings: """pi: interface CompactionSettings { enabled; reserveTokens; keepRecentTokens }""" enabled: bool = True reserve_tokens: int = 16384 # pi: 16384(摘要提示词与输出预留) keep_recent_tokens: int = 20000 # pi: 20000(压缩后保留的近期上下文预算) DEFAULT_COMPACTION_SETTINGS = CompactionSettings() # pi utils.ts: const TOOL_RESULT_MAX_CHARS = 2000 TOOL_RESULT_MAX_CHARS = 2000 # ====================================================================== # 摘要提示词 —— 从 pi compaction.ts 逐字移植(不得改写,摘要质量依赖它) # ====================================================================== SUMMARIZATION_SYSTEM_PROMPT = \ "You are a context summarization assistant. Your task is to read a conversation " \ "between a user and an AI assistant, then produce a structured summary following " \ "the exact format specified.\n\n" \ "Do NOT continue the conversation. Do NOT respond to any questions in the " \ "conversation. ONLY output the structured summary." SUMMARIZATION_PROMPT = """The messages above are a conversation to summarize. Create a structured context checkpoint summary that another LLM will use to continue the work. Use this EXACT format: ## Goal [What is the user trying to accomplish? Can be multiple items if the session covers different tasks.] ## Constraints & Preferences - [Any constraints, preferences, or requirements mentioned by user] - [Or "(none)" if none were mentioned] ## Progress ### Done - [x] [Completed tasks/changes] ### In Progress - [ ] [Current work] ### Blocked - [Issues preventing progress, if any] ## Key Decisions - **[Decision]**: [Brief rationale] ## Next Steps 1. [Ordered list of what should happen next] ## Critical Context - [Any data, examples, or references needed to continue] - [Or "(none)" if not applicable] Keep each section concise. Preserve exact file paths, function names, and error messages.""" UPDATE_SUMMARIZATION_PROMPT = """The messages above are NEW conversation messages to incorporate into the existing summary provided in tags. Update the existing structured summary with new information. RULES: - PRESERVE all existing information from the previous summary - ADD new progress, decisions, and context from the new messages - UPDATE the Progress section: move items from "In Progress" to "Done" when completed - UPDATE "Next Steps" based on what was accomplished - PRESERVE exact file paths, function names, and error messages - If something is no longer relevant, you may remove it Use this EXACT format: ## Goal [Preserve existing goals, add new ones if the task expanded] ## Constraints & Preferences - [Preserve existing, add new ones discovered] ## Progress ### Done - [x] [Include previously done items AND newly completed items] ### In Progress - [ ] [Current work - update based on progress] ### Blocked - [Current blockers - remove if resolved] ## Key Decisions - **[Decision]**: [Brief rationale] (preserve all previous, add new) ## Next Steps 1. [Update based on current state] ## Critical Context - [Preserve important context, add new if needed] Keep each section concise. Preserve exact file paths, function names, and error messages.""" TURN_PREFIX_SUMMARIZATION_PROMPT = """This is the PREFIX of a turn that was too large to keep. The SUFFIX (recent work) is retained. Summarize the prefix to provide context for the retained suffix: ## Original Request [What did the user ask for in this turn?] ## Early Progress - [Key decisions and work done in the prefix] ## Context for Suffix - [Information needed to understand the retained recent work] Be concise. Focus on what's needed to understand the kept suffix.""" # ====================================================================== # 对话序列化 —— 1:1 对照 utils.ts serializeConversation # ====================================================================== def _safe_json(value) -> str: try: s = json.dumps(value, ensure_ascii=False) return s if s is not None else "undefined" except Exception: return "[unserializable]" def _content_text(content) -> str: """pi contentText: str 或 [{type:"text",text}] 列表取文本拼接""" if content is None: return "" if isinstance(content, str): return content if isinstance(content, list): parts = [] for b in content: if isinstance(b, dict) and b.get("type") == "text": parts.append(str(b.get("text", ""))) return "\n".join(p for p in parts if p) return str(content) def _truncate_for_summary(text: str, max_chars: int) -> str: """pi utils.ts truncateForSummary(逐字逻辑)""" if len(text) <= max_chars: return text truncated = len(text) - max_chars return f"{text[:max_chars]}\n\n[... {truncated} more characters truncated]" def serialize_conversation(messages: List[AgentMessage]) -> str: """ 1:1 对照 utils.ts serializeConversation(输出格式逐字一致): [User]: ... [Assistant thinking]: ... [Assistant]: ... [Assistant tool calls]: name(k=v, k2=v2); name2(...) [Tool result]: ...(超 2000 字符截断) """ parts: List[str] = [] for msg in messages: if msg.role == "user": content = _content_text(msg.content) if content: parts.append(f"[User]: {content}") elif msg.role == "assistant": thinking_parts = [] tool_calls = [] if msg.reasoning: thinking_parts.append(msg.reasoning) text = _content_text(msg.content) for tc in (msg.tool_calls or []): args_str = ", ".join(f"{k}={_safe_json(v)}" for k, v in (tc.arguments or {}).items()) tool_calls.append(f"{tc.name}({args_str})") if thinking_parts: parts.append(f"[Assistant thinking]: {chr(10).join(thinking_parts)}") if text: parts.append(f"[Assistant]: {text}") if tool_calls: parts.append(f"[Assistant tool calls]: {'; '.join(tool_calls)}") elif msg.role == "toolResult": content = _content_text(msg.content) if content: parts.append(f"[Tool result]: " f"{_truncate_for_summary(content, TOOL_RESULT_MAX_CHARS)}") return "\n\n".join(parts) # ====================================================================== # 文件操作提取 —— 1:1 对照 utils.ts extractFileOps*/computeFileLists/formatFileOperations # ====================================================================== class FileOperations: def __init__(self): self.read: Set[str] = set() self.written: Set[str] = set() self.edited: Set[str] = set() def extract_file_ops_from_message(message: AgentMessage, file_ops: FileOperations): """pi: assistant 的 toolCall 参数里 path 字段 → read/write/edit 归类""" if message.role != "assistant": return for tc in (message.tool_calls or []): args = tc.arguments or {} path = args.get("path") if not isinstance(path, str) or not path: continue if tc.name == "read": file_ops.read.add(path) elif tc.name == "write": file_ops.written.add(path) elif tc.name == "edit": file_ops.edited.add(path) def compute_file_lists(file_ops: FileOperations) -> Tuple[List[str], List[str]]: """pi computeFileLists: modified = edited|written;readOnly = read-modified;均排序""" modified = file_ops.edited | file_ops.written read_only = sorted(f for f in file_ops.read if f not in modified) return read_only, sorted(modified) def format_file_operations(read_files: List[str], modified_files: List[str]) -> str: """pi formatFileOperations: / 标签拼接""" sections = [] if read_files: sections.append("\n" + "\n".join(read_files) + "\n") if modified_files: sections.append("\n" + "\n".join(modified_files) + "\n") if not sections: return "" return "\n\n" + "\n\n".join(sections) # ====================================================================== # 切分点 —— 1:1 对照 findValidCutPoints / findTurnStartIndex / findCutPoint # ====================================================================== @dataclass class CutPointResult: """pi: interface CutPointResult""" first_kept_index: int # 保留段首条在 compactable 列表里的下标 turn_start_index: int = -1 # 断轮时:该轮起点(user 消息)下标;否则 -1 is_split_turn: bool = False def find_valid_cut_points(messages: List[AgentMessage], start_index: int, end_index: int) -> List[int]: """ pi findValidCutPoints:消息角色为 user/assistant 的位置是有效切点 (toolResult 不能做切点——它会与前面的 toolCall 分离)。 pi 里的 bashExecution/branchSummary/compactionSummary 等角色 在 haocode 消息模型中不存在,等价规则即 role in (user, assistant)。 """ cut_points = [] for i in range(start_index, end_index): if messages[i].role in ("user", "assistant"): cut_points.append(i) return cut_points def find_turn_start(messages: List[AgentMessage], entry_index: int, start_index: int) -> int: """pi findTurnStartIndex:向前找本轮起点(user 消息 / branch_summary)""" for i in range(entry_index, start_index - 1, -1): if messages[i].role == "user": return i return -1 def _estimate(msg: AgentMessage) -> int: from .context import estimate_message_tokens return estimate_message_tokens(msg) def find_cut_point(messages: List[AgentMessage], start_index: int, end_index: int, keep_recent_tokens: int) -> CutPointResult: """ 1:1 对照 pi findCutPoint: 1. 从尾部向前累计 token,直到累计 >= keep_recent_tokens 2. 取该位置(含)之后的第一个有效切点 3. 切点不是 user 消息 → 断轮:找本轮起点,前缀单独摘要 (pi 的「回退跳过状态条目」循环针对 session 状态条目; haocode 消息列表没有状态条目,等价省略。) """ cut_points = find_valid_cut_points(messages, start_index, end_index) if not cut_points: return CutPointResult(first_kept_index=start_index) accumulated = 0 cut_index = cut_points[0] for i in range(end_index - 1, start_index - 1, -1): accumulated += _estimate(messages[i]) if accumulated >= keep_recent_tokens: for c in cut_points: if c >= i: cut_index = c break break is_user = messages[cut_index].role == "user" turn_start = -1 if is_user else find_turn_start(messages, cut_index, start_index) is_split = (not is_user) and turn_start != -1 return CutPointResult(first_kept_index=cut_index, turn_start_index=turn_start, is_split_turn=is_split) # ====================================================================== # 压缩准备 —— 1:1 对照 prepareCompaction # ====================================================================== @dataclass class CompactionPreparation: """pi: interface CompactionPreparation""" messages_to_summarize: List[AgentMessage] = field(default_factory=list) turn_prefix_messages: List[AgentMessage] = field(default_factory=list) retained_tail: List[AgentMessage] = field(default_factory=list) is_split_turn: bool = False tokens_before: int = 0 previous_summary: Optional[str] = None file_ops: FileOperations = field(default_factory=FileOperations) settings: CompactionSettings = DEFAULT_COMPACTION_SETTINGS def _compaction_diag(line: str) -> None: """压缩自诊断日志(与 recovery.compact_diag_log 同一文件/规范: print + 追加 compaction_diag.log + 失败静默)。""" import os as _os import time as _time t = _time.time() stamp = (f"[{_time.strftime('%H:%M:%S', _time.localtime(t))}" f".{int(t * 1000) % 1000:03d}]") print(f"{stamp} {line}", flush=True) try: path = _os.path.join( _os.path.dirname(_os.path.dirname(_os.path.dirname(_os.path.abspath(__file__)))), "compaction_diag.log") with open(path, "a", encoding="utf-8") as f: f.write(f"{stamp} {line}\n") except Exception: pass def _nothing_to_summarize_diag(messages: List[AgentMessage], compactable: List[AgentMessage], cut: CutPointResult, settings: CompactionSettings) -> str: """prepare_compaction 返回 None(无可摘要内容)时的自诊断行: 记录分支子类型 + 关键量,下次失败可直接从日志定位原因。 子类型: no_valid_cut_points —— 可压缩范围内没有任何 user/assistant 条目 total_below_keep_recent —— 可压缩范围总量 < keep_recent(通常意味着 上下文主体是旧摘要本身,无新内容) cut_pinned_at_zero —— 预算从尾部累加只在 i=0 才达标,即首条 条目独占 ≥ (总量-keep_recent) 的 token (单条超长消息/巨型工具输出主导上下文) """ from .context import estimate_context_tokens, estimate_message_tokens try: total = estimate_context_tokens(messages).tokens n_user = sum(1 for m in compactable if m.role == "user") n_asst = sum(1 for m in compactable if m.role == "assistant") ests = [(estimate_message_tokens(m), i) for i, m in enumerate(compactable)] e0 = ests[0][0] if ests else 0 rest = sum(e for e, _ in ests[1:]) top3 = sorted(ests, reverse=True)[:3] roles = " ".join(("T" if m.role == "toolResult" else m.role[0].upper()) for m in compactable[:10]) if not find_valid_cut_points(compactable, 0, len(compactable)): sub = "no_valid_cut_points" elif sum(e for e, _ in ests) < settings.keep_recent_tokens: sub = "total_below_keep_recent" else: sub = "cut_pinned_at_zero" big = "; ".join(f"idx{i}={e}" for e, i in top3) return (f"[COMPACT_NONE] branch=nothing_to_summarize sub={sub} " f"msgs={len(messages)} est_total={total} compactable={len(compactable)} " f"keep_recent={settings.keep_recent_tokens} " f"first_kept={cut.first_kept_index} turn_start={cut.turn_start_index} " f"split={cut.is_split_turn} user={n_user} asst={n_asst} " f"est_first={e0} est_rest={rest} top3=[{big}] head_roles=[{roles}]") except Exception as ex: # 诊断本身失败不能影响主流程 return f"[COMPACT_NONE] branch=nothing_to_summarize (diag failed: {ex})" def prepare_compaction(messages: List[AgentMessage], settings: Optional[CompactionSettings] = None ) -> Optional[CompactionPreparation]: """ 1:1 对照 pi prepareCompaction: - 上一条压缩摘要(messages[0].kind == "compaction_summary")不重复摘要, 其内容作为 previousSummary 走迭代更新提示词 - 可压缩范围 = 摘要之后的全部消息(pi 里等价于「上次保留尾 + 新消息」) - 切点在可压缩范围内选;tokens_before 按完整上下文(含摘要消息)估算 不可压缩(空/无摘要对象)时返回 None(对照 pi 返回 ok(undefined))。 """ from .context import estimate_context_tokens settings = settings or DEFAULT_COMPACTION_SETTINGS if not messages: _compaction_diag("[COMPACT_NONE] branch=empty_messages") return None if messages[-1].kind == "compaction_summary": _compaction_diag( f"[COMPACT_NONE] branch=tail_is_summary msgs={len(messages)} " f"tail_head={_content_text(messages[-1].content)[:60]!r}") return None previous_summary = None if messages[0].kind == "compaction_summary": previous_summary = _content_text(messages[0].content) compactable = messages[1:] else: compactable = messages if not compactable: _compaction_diag("[COMPACT_NONE] branch=compactable_empty (上下文只剩旧摘要)") return None tokens_before = estimate_context_tokens(messages).tokens cut = find_cut_point(compactable, 0, len(compactable), settings.keep_recent_tokens) history_end = cut.turn_start_index if cut.is_split_turn else cut.first_kept_index messages_to_summarize = compactable[:history_end] turn_prefix_messages = [] if cut.is_split_turn: turn_prefix_messages = compactable[cut.turn_start_index:cut.first_kept_index] retained_tail = compactable[cut.first_kept_index:] if not messages_to_summarize and not turn_prefix_messages: # 🆕 自诊断:记录是哪种子条件导致无东西可摘要(见 _nothing_to_summarize_diag) _compaction_diag(_nothing_to_summarize_diag(messages, compactable, cut, settings)) return None # 没有可摘要内容 file_ops = FileOperations() for m in messages_to_summarize: extract_file_ops_from_message(m, file_ops) if cut.is_split_turn: for m in turn_prefix_messages: extract_file_ops_from_message(m, file_ops) return CompactionPreparation( messages_to_summarize=messages_to_summarize, turn_prefix_messages=turn_prefix_messages, retained_tail=retained_tail, is_split_turn=cut.is_split_turn, tokens_before=tokens_before, previous_summary=previous_summary, file_ops=file_ops, settings=settings, ) # ====================================================================== # 摘要生成 —— 1:1 对照 generateSummaryWithUsage / generateTurnPrefixSummary # ====================================================================== def _build_summary_prompt(conversation_text: str, previous_summary: Optional[str]) -> str: """pi generateSummaryWithUsage 的 prompt 组装(逐字结构)""" base = UPDATE_SUMMARIZATION_PROMPT if previous_summary else SUMMARIZATION_PROMPT prompt = f"\n{conversation_text}\n\n\n" if previous_summary: prompt += f"\n{previous_summary}\n\n\n" prompt += base return prompt def generate_summary(messages: List[AgentMessage], summarize_fn: Callable[[str, str, int], str], reserve_tokens: int, model_max_tokens: int, previous_summary: Optional[str] = None ) -> Tuple[Optional[str], Optional[str]]: """ 返回 (summary_text, error)。 maxTokens = min(0.8 × reserveTokens, model.maxTokens) —— 1:1 对照。 """ max_tokens = min( int(0.8 * reserve_tokens), model_max_tokens if model_max_tokens > 0 else (1 << 30), ) conversation = serialize_conversation(messages) prompt = _build_summary_prompt(conversation, previous_summary) try: text = summarize_fn(prompt, SUMMARIZATION_SYSTEM_PROMPT, max_tokens) except Exception as e: return None, f"Summarization failed: {e}" if not text or not text.strip(): return None, "Summarization failed: empty response" return text.strip(), None def generate_turn_prefix_summary(messages: List[AgentMessage], summarize_fn: Callable[[str, str, int], str], reserve_tokens: int, model_max_tokens: int ) -> Tuple[Optional[str], Optional[str]]: """maxTokens = min(0.5 × reserveTokens, model.maxTokens) —— 1:1 对照""" max_tokens = min( int(0.5 * reserve_tokens), model_max_tokens if model_max_tokens > 0 else (1 << 30), ) conversation = serialize_conversation(messages) prompt = f"\n{conversation}\n\n\n{TURN_PREFIX_SUMMARIZATION_PROMPT}" try: text = summarize_fn(prompt, SUMMARIZATION_SYSTEM_PROMPT, max_tokens) except Exception as e: return None, f"Turn prefix summarization failed: {e}" if not text or not text.strip(): return None, "Turn prefix summarization failed: empty response" return text.strip(), None # ====================================================================== # 主入口 —— 1:1 对照 compact() # ====================================================================== def compact_context(messages: List[AgentMessage], model: ModelConfig, summarize_fn: Callable[[str, str, int], str], settings: Optional[CompactionSettings] = None ) -> Optional[List[AgentMessage]]: """ 执行压缩。返回 [压缩摘要消息] + 保留尾巴;不可压缩时返回 None。 摘要消息: role="user", kind="compaction_summary",content 为纯摘要文本 (pi 的 compaction 条目;下次压缩时自动走迭代更新提示词)。 断轮(切点落在某轮中间)时 1:1 对照 pi compact(): 历史摘要 与 轮前缀摘要 分两次 LLM 调用,拼接为 {history}\n\n---\n\n**Turn Context (split turn):**\n\n{prefix} 最后追加文件操作附录(/)。 """ settings = settings or DEFAULT_COMPACTION_SETTINGS prep = prepare_compaction(messages, settings) if prep is None: return None history_text: Optional[str] = None prefix_error: Optional[str] = None history_error: Optional[str] = None if prep.is_split_turn and prep.turn_prefix_messages: if prep.messages_to_summarize: history_text, history_error = generate_summary( prep.messages_to_summarize, summarize_fn, prep.settings.reserve_tokens, model.max_tokens, prep.previous_summary) if history_error: return _degraded_compact(prep, history_error) else: history_text = "No prior history." prefix_text, prefix_error = generate_turn_prefix_summary( prep.turn_prefix_messages, summarize_fn, prep.settings.reserve_tokens, model.max_tokens) if prefix_error: return _degraded_compact(prep, prefix_error) summary = (f"{history_text}\n\n---\n\n" f"**Turn Context (split turn):**\n\n{prefix_text}") else: if not prep.messages_to_summarize: return None summary, history_error = generate_summary( prep.messages_to_summarize, summarize_fn, prep.settings.reserve_tokens, model.max_tokens, prep.previous_summary) if history_error: return _degraded_compact(prep, history_error) read_files, modified_files = compute_file_lists(prep.file_ops) summary += format_file_operations(read_files, modified_files) summary_msg = AgentMessage(role="user", content=summary, kind="compaction_summary") return [summary_msg] + prep.retained_tail def _degraded_compact(prep: CompactionPreparation, error: str) -> Optional[List[AgentMessage]]: """ 已声明偏差(对照 pi: 直接返回 CompactionError): 桌面应用优先「不丢上下文」——摘要失败时降级为机械摘录, 保留尾巴原样不动。 """ old = prep.messages_to_summarize + prep.turn_prefix_messages if not old: return None excerpt = serialize_conversation(old)[-500:] summary = (f"(自动压缩:摘要生成失败 [{error}],以下为旧对话尾部摘录)\n\n{excerpt}") summary_msg = AgentMessage(role="user", content=summary, kind="compaction_summary") return [summary_msg] + prep.retained_tail