""" AirContext mode — V2 上下文管理器。 V2 改进：压缩质量校验、自适应 Token 估算、陈旧锁检测、三级降级压缩。 """ from __future__ import annotations import os import re from pathlib import Path from air_runtime.io import atomic_json_write, safe_json_load from air_runtime.paths import airplan_root, event_log_path from air_runtime.events import EventLog, CONTEXT_COMPACTED from air_runtime.utils import now_iso class CompressionLevel: """三级降级""" RETRY = "retry" # 1. 重试一次 FALLBACK_MODEL = "fallback_model" # 2. 换模型 TRUNCATE = "truncate" # 3. 激进截断 DEFAULT_TRUNCATION_KEEP = 10 # 保留最近 10 轮 CHARS_PER_TOKEN = { "chinese": 1.5, "english": 4.0, "code": 3.0, "markup": 5.0, } MUST_PRESERVE_PATTERNS = [ r"[A-Za-z0-9_\-/]+\.(py|ts|js|cpp|h|md|json|yaml)", r"ADR-\d{4}", r"TODO|FIXME|HACK", r"INV-\d+", ] def _ctx_paths(project_root: Path) -> dict[str, Path]: root = airplan_root(project_root) / "state" / "aircontext" return {"root": root, "state": root / "state.json", "lock": root / "compactor.lock"} def estimate_tokens(text: str) -> int: chinese = len(re.findall(r"[一-鿿]", text)) code = len(re.findall(r"[{}()\[\];=<>]", text)) markup = len(re.findall(r"[#*\-`|]", text)) english = max(0, len(text) - chinese - code - markup) tokens = ( chinese / CHARS_PER_TOKEN["chinese"] + code / CHARS_PER_TOKEN["code"] + markup / CHARS_PER_TOKEN["markup"] + english / CHARS_PER_TOKEN["english"] ) return int(tokens) def validate_compression(original: str, summary: str) -> dict: missing = [] for pattern in MUST_PRESERVE_PATTERNS: orig_matches = set(re.findall(pattern, original)) summary_matches = set(re.findall(pattern, summary)) lost = orig_matches - summary_matches if len(lost) > len(orig_matches) * 0.3 and len(orig_matches) > 3: missing.append({"pattern": pattern, "lost": list(lost)[:10]}) return {"ok": len(missing) == 0, "missing": missing, "originalTokens": estimate_tokens(original), "summaryTokens": estimate_tokens(summary)} def _compress_basic(text: str, max_tokens: int) -> str: """基础压缩：token 估算 + 截断""" estimated_tokens = len(text) // 3 if estimated_tokens <= max_tokens: return text # 按行截断 lines = text.split('\n') chars_per_line_estimate = 30 keep_lines = int(max_tokens * chars_per_line_estimate / 80) # 80 chars/line return '\n'.join(lines[-keep_lines:]) def _simplify_prompt(text: str) -> str: """简化 prompt：移除详细上下文，保留核心""" lines = text.split('\n') # 只保留前 3 行 + 包含 "def " / "class " / "#" 的行 kept = lines[:3] kept.extend([l for l in lines[3:] if 'def ' in l or 'class ' in l or l.startswith('#')]) return '\n'.join(kept) def compress_with_fallback(context: str, max_tokens: int = 4000) -> dict: """ 三级降级压缩： - 尝试正常压缩 - 失败则换模型重试 - 再失败则激进截断 返回: {"level": "...", "result": "...", "tokens": N} """ # Level 1: 正常尝试 try: result = _compress_basic(context, max_tokens) return {"level": CompressionLevel.RETRY, "result": result, "tokens": len(result.split())} except Exception: pass # Level 2: 换模型（更简单的 prompt + 更宽松的 max_tokens） try: simplified = _simplify_prompt(context) result = _compress_basic(simplified, int(max_tokens * 1.5)) return {"level": CompressionLevel.FALLBACK_MODEL, "result": result, "tokens": len(result.split())} except Exception: pass # Level 3: 激进截断 lines = context.split('\n') # 提取 ADR 引用行 adr_lines = [l for l in lines if 'ADR-' in l or 'adr-' in l] # 保留最近 N 轮 recent_lines = lines[-DEFAULT_TRUNCATION_KEEP * 5:] # 每轮约 5 行 truncated = '\n'.join(recent_lines + adr_lines) return { "level": CompressionLevel.TRUNCATE, "result": truncated, "tokens": len(truncated.split()), "warning": f"truncated to {DEFAULT_TRUNCATION_KEEP * 5} recent lines + {len(adr_lines)} ADR lines" } def validate_compression_with_fallback(project_root: Path, context_path: Path) -> dict: """验证压缩有效性，失败时触发三级降级""" content = context_path.read_text() original_len = len(content) # 先用当前配置尝试 result = compress_with_fallback(content) validation = { "original_chars": original_len, "result_chars": len(result["result"]), "level": result["level"], "tokens": result.get("tokens", 0), } if result["level"] == CompressionLevel.TRUNCATE: validation["warning"] = result.get("warning", "") return validation def acquire_compactor_lock(lock_path: Path) -> bool: try: fd = os.open(lock_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY) os.write(fd, str(os.getpid()).encode()) os.close(fd) return True except FileExistsError: try: pid = int(lock_path.read_text().strip()) os.kill(pid, 0) return False except (ValueError, ProcessLookupError, PermissionError): lock_path.unlink(missing_ok=True) return acquire_compactor_lock(lock_path) def ctx_enter(project_root: Path) -> dict: paths = _ctx_paths(project_root) paths["root"].mkdir(parents=True, exist_ok=True) atomic_json_write(paths["state"], {"enabled": True, "updatedAt": now_iso(), "projectRoot": str(project_root)}) return {"state_path": str(paths["state"])} def main(args) -> None: project_root = Path(args.project).expanduser().resolve() sub = args.sub or "status" paths = _ctx_paths(project_root) if sub == "enter": result = ctx_enter(project_root) print(f"airplan_mode=ctx\nstate_path={result['state_path']}") elif sub == "estimate": text = "sample" # 实际使用时从 stdin 或文件读取 tokens = estimate_tokens(text) print(f"airplan_mode=ctx\ntokens={tokens}") elif sub == "validate": ctx_path = project_root / "AirPlan" / "context.md" if ctx_path.exists(): validation = validate_compression_with_fallback(project_root, ctx_path) ok = validation["level"] != CompressionLevel.TRUNCATE print(f"airplan_mode=ctx\nvalidation_ok={ok}\nlevel={validation['level']}\ntokens={validation['tokens']}") if "warning" in validation: print(f"warning={validation['warning']}") else: print("airplan_mode=ctx\nvalidation_ok=true") elif sub == "compress": # 读取 context 文件 ctx_path = project_root / "AirPlan" / "context.md" if not ctx_path.exists(): print("error: context.md not found") return # 调用三级降级压缩 content = ctx_path.read_text() result = compress_with_fallback(content) log = EventLog(event_log_path(project_root)) log.emit(CONTEXT_COMPACTED, { "compressionLevel": result["level"], "originalChars": len(content), "resultChars": len(result["result"]), "tokens": result.get("tokens", 0), }) print(f"airplan_mode=ctx") print(f"compression_level={result['level']}") print(f"original_chars={len(content)}") print(f"result_chars={len(result['result'])}") if 'warning' in result: print(f"warning={result['warning']}") # 可选：写回压缩结果 if getattr(args, "write_back", False): compressed_path = project_root / "AirPlan" / "context.compressed.md" compressed_path.write_text(result['result']) print(f"written_to={compressed_path}") else: state = safe_json_load(paths["state"]) or {} print(f"airplan_mode=ctx\nenabled={state.get('enabled', False)}")