"""Encoding-aware file IO, output truncation, and per-path write locks.""" import asyncio from pathlib import Path from typing import Iterable import aiofiles import aiofiles.os from ...constants import DEFAULT_MAX_BYTES, MAX_FILE_READ_BYTES, TRUNCATION_NOTICE_MARKER from ...schema import FileChunk from ...utils import get_logger logger = get_logger(log_to_file=False) # --------------------------------------------------------------------------- # In-process per-path write lock. # --------------------------------------------------------------------------- _PATH_LOCKS_MAX = 1024 _PATH_LOCKS: dict[str, asyncio.Lock] = {} _PATH_LOCKS_REGISTRY = asyncio.Lock() async def get_path_lock(target: Path) -> asyncio.Lock: """Return the asyncio.Lock for ``target``; created lazily on first request.""" key = str(target) async with _PATH_LOCKS_REGISTRY: lock = _PATH_LOCKS.get(key) if lock is None: if len(_PATH_LOCKS) >= _PATH_LOCKS_MAX: to_remove = [k for k, v in _PATH_LOCKS.items() if not v.locked()] for k in to_remove[: len(_PATH_LOCKS) // 2]: del _PATH_LOCKS[k] lock = asyncio.Lock() _PATH_LOCKS[key] = lock return lock # --------------------------------------------------------------------------- # Encoding detection # --------------------------------------------------------------------------- _STANDARD_TEXT_EXTS = { ".md", ".py", ".js", ".ts", ".json", ".yaml", ".yml", ".html", ".css", ".xml", ".log", ".conf", ".ini", ".txt", ".sh", } _NON_STANDARD_EXTS = {".csv", ".bat", ".cmd", ".reg"} def _try_decode(data: bytes, encodings: Iterable[str]) -> tuple[str, str] | None: """Return ``(text, encoding)`` for the first encoding that decodes ``data`` cleanly.""" for enc in encodings: try: return data.decode(enc), enc except (UnicodeDecodeError, LookupError): continue return None def _decode_known_file(data: bytes, file_extension: str) -> tuple[str, str]: """Decode file bytes using the extension as a hint. Returns ``(text, encoding)``.""" if data.startswith(b"\xef\xbb\xbf"): return data.decode("utf-8-sig"), "utf-8-sig" if data.startswith((b"\xff\xfe", b"\xfe\xff")): try: return data.decode("utf-16"), "utf-16" except UnicodeDecodeError: pass ext = (file_extension or "").lower() if ext in _STANDARD_TEXT_EXTS: try: return data.decode("utf-8-sig"), "utf-8" except UnicodeDecodeError: pass if ext in _NON_STANDARD_EXTS: result = _try_decode(data, ("utf-8-sig", "gbk")) if result is not None: text, enc = result return text, "utf-8" if enc == "utf-8-sig" else enc return data.decode("utf-8", errors="replace"), "utf-8" # --------------------------------------------------------------------------- # File read / write # --------------------------------------------------------------------------- async def read_file_safe(file_path, max_bytes: int = MAX_FILE_READ_BYTES) -> tuple[str, str]: """Read file in byte mode and decode using extension-aware strategy. Returns ``(text, encoding)``. """ stat = await aiofiles.os.stat(str(file_path)) read_size = min(stat.st_size, max_bytes) async with aiofiles.open(str(file_path), "rb") as f: data = await f.read(read_size) return _decode_known_file(data, Path(file_path).suffix) async def read_file_lines_safe( file_path, start_line: int, end_line: int | None, *, max_collect_bytes: int = DEFAULT_MAX_BYTES * 2, ) -> tuple[str, int, str]: """Read a 1-based inclusive line range without loading the full file. Returns ``(text, total_lines, encoding)``. """ encoding = await detect_file_encoding(file_path) lines: list[str] = [] collected_bytes = 0 total = 0 async with aiofiles.open(str(file_path), "r", encoding=encoding, errors="replace") as f: async for line in f: total += 1 if total >= start_line and (end_line is None or total <= end_line): if collected_bytes < max_collect_bytes: cleaned = line.rstrip("\n") lines.append(cleaned) collected_bytes += len(cleaned.encode(encoding, errors="replace")) + 1 return "\n".join(lines), total, encoding async def detect_file_encoding(file_path, sniff_bytes: int = 8192) -> str: """Detect the encoding of an existing file so writes can preserve it.""" try: async with aiofiles.open(str(file_path), "rb") as f: data = await f.read(sniff_bytes) except Exception: return "utf-8" _, enc = _decode_known_file(data, Path(file_path).suffix) return enc async def write_file_safe(file_path: Path, content: str | bytes, encoding: str = "utf-8") -> None: """Write ``content`` to ``file_path`` in binary mode; creates parent dirs.""" file_path.parent.mkdir(parents=True, exist_ok=True) if isinstance(content, str): try: payload = content.encode(encoding) except (UnicodeEncodeError, LookupError): logger.warning( "write_file_safe: %r cannot encode all chars, falling back to utf-8", encoding, ) payload = content.encode("utf-8") else: payload = content async with aiofiles.open(str(file_path), "wb") as f: await f.write(payload) # --------------------------------------------------------------------------- # Output truncation # --------------------------------------------------------------------------- def truncate_text_output( text: str, *, start_line: int = 1, total_lines: int = 0, max_bytes: int = DEFAULT_MAX_BYTES, file_path: str | None = None, encoding: str = "utf-8", ) -> str: """Truncate text by bytes preserving line integrity; append a continuation notice.""" if not text or max_bytes <= 0: return text try: text_bytes = text.encode(encoding) if len(text_bytes) <= max_bytes: return text truncated = text_bytes[:max_bytes] result = truncated.decode(encoding, errors="ignore") newline_count = result.count("\n") next_line = start_line + max(1, newline_count) if next_line <= total_lines: read_from = next_line elif start_line < total_lines: read_from = total_lines else: return result notice = ( TRUNCATION_NOTICE_MARKER + f"\nThe output above was truncated." f"\nThe full content is saved to the file and contains {total_lines} lines in total." f"\nThis excerpt starts at line {start_line} and covers the next {max_bytes} bytes." f"\nIf the current content is not enough, call `read` with file={file_path or ''} " f"start_line={read_from} to read more." ) return result + notice except Exception: logger.warning("truncate_text_output failed, returning original text", exc_info=True) return text def truncate_session_output( text: str, *, start_line: int = 1, total_lines: int = 0, max_bytes: int = DEFAULT_MAX_BYTES, file_path: str | None = None, encoding: str = "utf-8", ) -> str: """Render session jsonl lines then truncate by bytes preserving line integrity. Each line of *text* is expected to be a serialized ``Msg`` JSON from a raw session transcript (``*.jsonl``). Lines are first rendered via :func:`~reme.steps.index._source_format.render_session_chunk_lines` — one message per line as ``[speaker @ created_at] content`` with internal newlines flattened — and then truncated with the same byte-budget logic as :func:`truncate_text_output`. """ if not text or max_bytes <= 0: return text # Render each jsonl line into a human-readable [speaker @ time] content line. # Import locally to avoid a hard dependency from file_io on the index package # at module load time; the import is safe (no circular dependency). from ..index._source_format import render_session_chunk_lines chunk = FileChunk(text=text) rendered_lines = render_session_chunk_lines(chunk) rendered_text = "\n".join(rendered_lines) # Reuse the same byte-budget truncation as truncate_text_output. return truncate_text_output( rendered_text, start_line=start_line, total_lines=total_lines, max_bytes=max_bytes, file_path=file_path, encoding=encoding, )