"""Context writers — write learned patterns to agent-specific context files. Writers take Recommendations and write them to the appropriate context injection mechanism for each agent system (CLAUDE.md, .cursorrules, etc.). """ from __future__ import annotations import re from abc import ABC, abstractmethod from datetime import datetime, timezone from pathlib import Path from ._shared import claude_config_dir from .models import ( ProjectInfo, Recommendation, RecommendationTarget, ) # Marker delimiters for Headroom-managed sections _MARKER_START = "" _MARKER_END = "" _MARKER_PATTERN = re.compile( re.escape(_MARKER_START) + r".*?" + re.escape(_MARKER_END), re.DOTALL, ) def _read_text_tolerant(file_path: Path) -> str: """Read an existing context file that we are about to rewrite as UTF-8. These files are predominantly valid UTF-8 but may carry a stray legacy byte (e.g. a cp1252 em-dash ``0x97``). Strict UTF-8 decoding aborts the whole ``--apply`` on a single such byte, so fall back to UTF-8 with replacement: this preserves the valid UTF-8 content — a full-file cp1252 fallback would instead turn every genuine UTF-8 em-dash into mojibake — and the subsequent ``write_text(encoding="utf-8")`` self-heals the file. """ raw = file_path.read_bytes() try: return raw.decode("utf-8") except UnicodeDecodeError: return raw.decode("utf-8", errors="replace") # ============================================================================= # Abstract Writer # ============================================================================= class ContextWriter(ABC): """Base class for writing recommendations to context/memory files.""" @abstractmethod def write( self, recommendations: list[Recommendation], project: ProjectInfo, dry_run: bool = True, ) -> WriteResult: ... # ============================================================================= # Write Result # ============================================================================= class WriteResult: """Result of a write operation.""" def __init__(self) -> None: self.files_written: list[Path] = [] self.content_by_file: dict[Path, str] = {} self.dry_run: bool = True # Human-readable notices (e.g. legacy CLAUDE.md migration) surfaced by the CLI. self.warnings: list[str] = [] def add(self, path: Path, content: str) -> None: self.files_written.append(path) self.content_by_file[path] = content # ============================================================================= # Shared section builder # ============================================================================= def _build_section(recommendations: list[Recommendation]) -> str: """Build the marker-delimited section content from recommendations.""" now = datetime.now(timezone.utc).strftime("%Y-%m-%d") lines = [ _MARKER_START, "## Headroom Learned Patterns", f"*Auto-generated by `headroom learn` on {now} — do not edit manually*", "", ] for rec in recommendations: lines.append(f"### {rec.section}") if rec.estimated_tokens_saved > 0: lines.append(f"*~{rec.estimated_tokens_saved:,} tokens/session saved*") lines.append(rec.content) lines.append("") lines.append(_MARKER_END) return "\n".join(lines) # Matches the "*~N tokens/session saved*" annotation emitted by _build_section. _TOKENS_ANNOTATION_PATTERN = re.compile(r"\*~([\d,]+) tokens/session saved\*\n?") def extract_marker_block(file_content: str) -> str | None: """Return the raw text of the headroom:learn marker block, or None. Unlike _parse_prior_recommendations, this returns the block verbatim (including the start/end markers) so it can be fed back to an LLM as context without losing formatting. Returns None if no block is present. """ match = _MARKER_PATTERN.search(file_content) return match.group(0) if match else None def _parse_prior_recommendations(existing: str) -> list[Recommendation]: """Parse recommendations out of a prior marker block. Returns [] if no marker block is present or it contains no sections. The returned Recommendation objects are round-trip compatible with _build_section — target is set to a placeholder since the marker block itself doesn't record it (blocks are always per-file and per-target). """ match = _MARKER_PATTERN.search(existing) if not match: return [] inner = match.group(0)[len(_MARKER_START) : -len(_MARKER_END)] recs: list[Recommendation] = [] for part in re.split(r"\n### ", "\n" + inner)[1:]: heading_line, _, body = part.partition("\n") heading = heading_line.strip() if not heading: continue tokens_saved = 0 tokens_match = _TOKENS_ANNOTATION_PATTERN.match(body) if tokens_match: tokens_saved = int(tokens_match.group(1).replace(",", "")) body = body[tokens_match.end() :] recs.append( Recommendation( target=RecommendationTarget.CONTEXT_FILE, section=heading, content=body.rstrip(), estimated_tokens_saved=tokens_saved, ) ) return recs def _merge_recommendations( file_path: Path, new_recommendations: list[Recommendation], ) -> list[Recommendation]: """Union new recommendations with prior ones whose section is not re-surfaced. Sections produced by the current run take precedence over same-named prior sections — the latest analysis is authoritative. Prior sections whose headings do not reappear in the new run are carried forward so a re-run doesn't silently drop accumulated learnings. To fully rebuild the block, delete it manually and re-run. """ if not file_path.exists(): return new_recommendations prior = _parse_prior_recommendations(_read_text_tolerant(file_path)) if not prior: return new_recommendations new_sections = {r.section for r in new_recommendations} carried = [p for p in prior if p.section not in new_sections] return list(new_recommendations) + carried def _merge_into_file(file_path: Path, new_recommendations: list[Recommendation]) -> str: """Merge new recommendations with any existing marker block and rebuild the file.""" merged = _merge_recommendations(file_path, new_recommendations) section = _build_section(merged) if file_path.exists(): existing = _read_text_tolerant(file_path) if _MARKER_START in existing: return _MARKER_PATTERN.sub(lambda _match: section, existing) return existing.rstrip() + "\n\n" + section + "\n" return section + "\n" def _strip_marker_block(content: str) -> str: """Remove the headroom:learn marker block from text, tidying blank lines. Used when migrating a stale block out of the team-shared CLAUDE.md into the personal CLAUDE.local.md. Returns "" if nothing but the block remained. """ cleaned = _MARKER_PATTERN.sub("", content) cleaned = re.sub(r"\n{3,}", "\n\n", cleaned).strip() return cleaned + "\n" if cleaned else "" # ============================================================================= # Claude Code Writer # ============================================================================= class ClaudeCodeWriter(ContextWriter): """Writes learned patterns to CLAUDE.local.md and MEMORY.md for Claude Code. Project-level learnings default to ``CLAUDE.local.md`` rather than ``CLAUDE.md``: per Claude Code's memory convention ``CLAUDE.md`` is team-shared and checked into git, while ``CLAUDE.local.md`` is personal and gitignored by default. Learned patterns are personal-by-default (they hold machine-specific absolute paths and tool-discovery byproducts), so writing them to the shared file pollutes it for teammates (issue #1072). Pass an explicit target via :meth:`set_context_target` (CLI ``--target``) to override -- e.g. ``CLAUDE.md`` to opt back into the shared file. """ def __init__(self, context_target: str | None = None) -> None: # Explicit write target for CONTEXT_FILE recs (overrides the default). self._context_target = context_target def set_context_target(self, context_target: str | None) -> None: """Override where CONTEXT_FILE recommendations are written. Accepts a path relative to the project root or an absolute path. """ self._context_target = context_target def write( self, recommendations: list[Recommendation], project: ProjectInfo, dry_run: bool = True, ) -> WriteResult: result = WriteResult() result.dry_run = dry_run context_recs = [r for r in recommendations if r.target == RecommendationTarget.CONTEXT_FILE] memory_recs = [r for r in recommendations if r.target == RecommendationTarget.MEMORY_FILE] if context_recs: target_path = self._resolve_context_path(project) # Migrate any stale block left in the team-shared CLAUDE.md by older # headroom versions into the new target, then strip it from CLAUDE.md # so the shared file is no longer polluted. migrated = self._migrate_legacy_block(project, target_path, result, dry_run) new_sections = {r.section for r in context_recs} merged_recs = context_recs + [r for r in migrated if r.section not in new_sections] full_content = _merge_into_file(target_path, merged_recs) result.add(target_path, full_content) if not dry_run: target_path.parent.mkdir(parents=True, exist_ok=True) target_path.write_text(full_content, encoding="utf-8") if memory_recs: memory_path = self._resolve_memory_path(project) full_content = _merge_into_file(memory_path, memory_recs) result.add(memory_path, full_content) if not dry_run: memory_path.parent.mkdir(parents=True, exist_ok=True) memory_path.write_text(full_content, encoding="utf-8") return result def _resolve_context_path(self, project: ProjectInfo) -> Path: # Explicit --target wins over every default. if self._context_target is not None: target = Path(self._context_target).expanduser() return target if target.is_absolute() else project.project_path / target # The home directory's CLAUDE.md (~/.claude/CLAUDE.md, or # $CLAUDE_CONFIG_DIR/CLAUDE.md) is the user's personal global memory, # not a team-shared file, so keep writing there. if project.project_path == Path.home(): return claude_config_dir() / "CLAUDE.md" # Project level: default to the gitignored, personal CLAUDE.local.md so # we never pollute the team-shared CLAUDE.md (issue #1072). return project.project_path / "CLAUDE.local.md" def _migrate_legacy_block( self, project: ProjectInfo, target_path: Path, result: WriteResult, dry_run: bool, ) -> list[Recommendation]: """Move a stale headroom block out of CLAUDE.md into the new target. Only fires for the default project-level case (no explicit --target, not the home directory) when CLAUDE.md still carries a marker block and the new target doesn't yet own one. Returns the migrated recommendations so the caller can carry them forward; records the cleaned CLAUDE.md and a warning on ``result``. Honors ``dry_run`` (no writes, warning still set). """ legacy_path = project.project_path / "CLAUDE.md" if self._context_target is not None or project.project_path == Path.home(): return [] if target_path == legacy_path or not legacy_path.exists(): return [] legacy_text = _read_text_tolerant(legacy_path) if _MARKER_START not in legacy_text: return [] # If the target already owns a block, it is the source of truth -- don't # double-migrate or clobber accumulated learnings. if target_path.exists() and _MARKER_START in _read_text_tolerant(target_path): return [] migrated = _parse_prior_recommendations(legacy_text) cleaned = _strip_marker_block(legacy_text) gitignore_hint = f" Ensure {target_path.name} is in your .gitignore so it stays personal." if cleaned: # CLAUDE.md has hand-written content too — keep it, drop only the block. result.add(legacy_path, cleaned) result.warnings.append( f"Moved Headroom learnings out of {legacy_path} into {target_path}: " f"CLAUDE.md is team-shared, so personal learnings now live in " f"{target_path.name}. Review the diff before committing.{gitignore_hint}" ) if not dry_run: legacy_path.write_text(cleaned, encoding="utf-8") else: # CLAUDE.md held nothing but the Headroom block — remove the husk. result.warnings.append( f"Removed {legacy_path} (it contained only Headroom learnings) and " f"moved them into {target_path}.{gitignore_hint}" ) if not dry_run: legacy_path.unlink() return migrated def _resolve_memory_path(self, project: ProjectInfo) -> Path: if project.memory_file: return project.memory_file return project.data_path / "memory" / "MEMORY.md" # ============================================================================= # Codex Writer (OpenAI Codex CLI) # ============================================================================= class CodexWriter(ContextWriter): """Writes learned patterns to AGENTS.md and instructions.md for Codex CLI.""" def write( self, recommendations: list[Recommendation], project: ProjectInfo, dry_run: bool = True, ) -> WriteResult: result = WriteResult() result.dry_run = dry_run context_recs = [r for r in recommendations if r.target == RecommendationTarget.CONTEXT_FILE] memory_recs = [r for r in recommendations if r.target == RecommendationTarget.MEMORY_FILE] if context_recs: agents_md = project.context_file or (project.project_path / "AGENTS.md") full_content = _merge_into_file(agents_md, context_recs) result.add(agents_md, full_content) if not dry_run: agents_md.parent.mkdir(parents=True, exist_ok=True) agents_md.write_text(full_content, encoding="utf-8") if memory_recs: instructions_md = project.memory_file or (project.data_path.parent / "instructions.md") full_content = _merge_into_file(instructions_md, memory_recs) result.add(instructions_md, full_content) if not dry_run: instructions_md.parent.mkdir(parents=True, exist_ok=True) instructions_md.write_text(full_content, encoding="utf-8") return result # ============================================================================= # Gemini Writer (Google Gemini CLI) # ============================================================================= class GeminiWriter(ContextWriter): """Writes learned patterns to GEMINI.md for Gemini CLI.""" def write( self, recommendations: list[Recommendation], project: ProjectInfo, dry_run: bool = True, ) -> WriteResult: result = WriteResult() result.dry_run = dry_run if not recommendations: return result gemini_md = project.context_file or (project.project_path / "GEMINI.md") full_content = _merge_into_file(gemini_md, recommendations) result.add(gemini_md, full_content) if not dry_run: gemini_md.parent.mkdir(parents=True, exist_ok=True) gemini_md.write_text(full_content, encoding="utf-8") return result