From 37628ba524c9c06a2ad2c21445c8c835a42e4418 Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Sat, 28 Mar 2026 18:41:09 +0800 Subject: [PATCH] refactor(truncation): improve file truncation logic (#184) * refactor(file_store): simplify ChromaDB client initialization and improve file truncation logic - Remove shutil import and _create_chroma_client method from chroma_file_store.py - Directly initialize ChromaDB PersistentClient in start method without retry logic - Reduce DEFAULT_MAX_BYTES from 100KB to 50KB in file_utils.py - Update truncation notice format to provide clearer continuation instructions - Add _truncate_fresh and _retruncate functions for better text truncation handling - Replace inline truncation logic with dedicated function calls in file_utils.py - Rename skills_tool_ids to md_file_tool_ids in tool_result_compactor.py - Update file detection logic to identify any .md files instead of only skill.md - Create comprehensive unit tests for truncation functionality in test_truncate_text_output.py * chore(version): bump version to 0.3.1.6 - Update __version__ from 0.3.1.5 to 0.3.1.6 in __init__.py --- reme/__init__.py | 2 +- reme/core/file_store/chroma_file_store.py | 28 +- .../components/tool_result_compactor.py | 12 +- reme/memory/file_based/tools/file_io.py | 11 +- reme/memory/file_based/utils/file_utils.py | 200 ++++++--- tests/light/test_truncate_text_output.py | 381 ++++++++++++++++++ 6 files changed, 538 insertions(+), 96 deletions(-) create mode 100644 tests/light/test_truncate_text_output.py diff --git a/reme/__init__.py b/reme/__init__.py index 1b211a70..f9ade852 100644 --- a/reme/__init__.py +++ b/reme/__init__.py @@ -6,7 +6,7 @@ from . import extension from . import memory from .reme import ReMe -__version__ = "0.3.1.5" +__version__ = "0.3.1.6" __all__ = [ "config", diff --git a/reme/core/file_store/chroma_file_store.py b/reme/core/file_store/chroma_file_store.py index 4deb07d0..d87f5fef 100644 --- a/reme/core/file_store/chroma_file_store.py +++ b/reme/core/file_store/chroma_file_store.py @@ -2,7 +2,6 @@ import json import random -import shutil import time from pathlib import Path @@ -109,9 +108,13 @@ class ChromaFileStore(BaseFileStore): except Exception as e: logger.error(f"Failed to save file metadata to {self._metadata_file}: {e}") - def _create_chroma_client(self): - """Create a ChromaDB PersistentClient instance.""" - return chromadb.PersistentClient( + async def start(self) -> None: + """Initialize ChromaDB client and collection.""" + if self.client is not None: + return + + # Initialize persistent ChromaDB client + self.client = chromadb.PersistentClient( path=str(self.db_path), settings=Settings( anonymized_telemetry=False, @@ -119,23 +122,6 @@ class ChromaFileStore(BaseFileStore): ), ) - async def start(self) -> None: - """Initialize ChromaDB client and collection.""" - if self.client is not None: - return - - # Initialize persistent ChromaDB client, retry once after wiping db_path on failure - try: - self.client = self._create_chroma_client() - except Exception as e: - logger.warning( - f"ChromaDB failed to initialize at {self.db_path} ({e}). " f"Deleting corrupted database and retrying.", - ) - if self.db_path.exists(): - shutil.rmtree(self.db_path) - logger.info(f"Deleted ChromaDB directory: {self.db_path}") - self.client = self._create_chroma_client() - # Get or create the chunks collection # ChromaDB uses cosine distance by default for similarity self.chunks_collection = self.client.get_or_create_collection( diff --git a/reme/memory/file_based/components/tool_result_compactor.py b/reme/memory/file_based/components/tool_result_compactor.py index d808974c..d80b2f61 100644 --- a/reme/memory/file_based/components/tool_result_compactor.py +++ b/reme/memory/file_based/components/tool_result_compactor.py @@ -91,7 +91,7 @@ class ToolResultCompactor(BaseOp): recent_n += 1 split_index = max(0, len(messages) - max(recent_n, self.recent_n)) - skills_tool_ids = set() + md_file_tool_ids = set() try: for msg in messages: if not isinstance(msg.content, list): @@ -105,12 +105,12 @@ class ToolResultCompactor(BaseOp): if ( block.get("name", "").lower() == "read_file" - and "skill.md" in (block.get("raw_input") or "").lower() + and ".md" in (block.get("raw_input") or "").lower() ): - skills_tool_ids.add(tool_id) + md_file_tool_ids.add(tool_id) except Exception as e: - logger.warning("Failed to detect skill tool ids: %s", e) - logger.info(f"skills_tool_ids: {skills_tool_ids}") + logger.warning("Failed to detect md file tool ids: %s", e) + logger.info(f"md_file_tool_ids: {md_file_tool_ids}") for idx, msg in enumerate(messages): if not isinstance(msg.content, list): @@ -120,7 +120,7 @@ class ToolResultCompactor(BaseOp): for block in msg.content: if isinstance(block, dict) and block.get("type") == "tool_result" and block.get("output"): tool_use_id = block.get("id", "") - if tool_use_id in skills_tool_ids: + if tool_use_id in md_file_tool_ids: effective_max_bytes = self.recent_max_bytes else: effective_max_bytes = max_bytes diff --git a/reme/memory/file_based/tools/file_io.py b/reme/memory/file_based/tools/file_io.py index d876407c..3e28addc 100644 --- a/reme/memory/file_based/tools/file_io.py +++ b/reme/memory/file_based/tools/file_io.py @@ -153,10 +153,13 @@ class FileIO: if text == selected_content and e < total: content_bytes = len(text.encode("utf-8")) notice = ( - TRUNCATION_NOTICE_MARKER - + f"\nFile: {file_path}\nStarting at start_line={s}, next {content_bytes} bytes." - f"\nTotal lines: {total}" - f"\nUse start_line={e + 1} to continue." + TRUNCATION_NOTICE_MARKER + f"\nThe output above was truncated." + f"\nThe full content is saved to the file " + f"and contains {total} lines in total." + f"\nThis excerpt starts at line {s} and " + f"covers the next {content_bytes} bytes." + "\nIf the current content is not enough, " + f"call `read_file` with file_path={file_path} start_line={e + 1} to read more." ) text = text + notice diff --git a/reme/memory/file_based/utils/file_utils.py b/reme/memory/file_based/utils/file_utils.py index e7204537..50248554 100644 --- a/reme/memory/file_based/utils/file_utils.py +++ b/reme/memory/file_based/utils/file_utils.py @@ -8,7 +8,7 @@ from ....core.utils import get_logger logger = get_logger() # Default truncation limit -DEFAULT_MAX_BYTES = 100 * 1024 +DEFAULT_MAX_BYTES = 50 * 1024 # Maximum file size to read into memory (1GB) MAX_FILE_READ_BYTES = 1024 * 1024 * 1024 @@ -16,17 +16,136 @@ MAX_FILE_READ_BYTES = 1024 * 1024 * 1024 # Marker prepended to every truncation notice. # Format: # <<>> -# File: -# Starting at start_line=X, next N bytes. -# Total lines: Z -# Use start_line=Y to continue. +# The output above was truncated. +# The full content is saved to the file and contains Z lines in total. +# This excerpt starts at line X and covers the next N bytes. +# If the current content is not enough, call `read_file` with file_path= start_line=Y to read more. # # Split output on this marker to recover the original (untruncated) portion: # original = output.split(TRUNCATION_NOTICE_MARKER)[0] TRUNCATION_NOTICE_MARKER = "<<>>" -# pylint: disable=too-many-return-statements +def _truncate_fresh( + text: str, + start_line: int, + total_lines: int, + max_bytes: int, + file_path: str | None, + encoding: str, +) -> str: + """Truncate fresh text (no prior truncation marker) by bytes with line integrity. + + Slices at the byte boundary and appends a truncation notice with a continuation + hint so callers know which line to read next. + + Returns the original text unchanged when it fits within max_bytes, or when the + last line itself exceeds max_bytes (unhandled edge case). + """ + text_bytes = text.encode(encoding) + + # Under the byte limit — return as-is without any modification. + if len(text_bytes) <= max_bytes: + return text + + # Slice at the byte boundary. + # Assuming every single line is shorter than DEFAULT_MAX_BYTES, this cut always + # lands mid-line, guaranteeing at least one complete line before the boundary. + # Lines that exceed DEFAULT_MAX_BYTES are not handled and may be skipped entirely. + truncated = text_bytes[:max_bytes] + # Decode back to str; errors="ignore" drops any split multi-byte character + # at the cut boundary without raising an exception. + result = truncated.decode(encoding, errors="ignore") + + # Count '\n' characters to determine how many complete lines are included. + # The tail after the final '\n' is a partial line that will be covered by + # the next read starting at next_line. + newline_count = result.count("\n") + + # Compute the first line number not yet fully included in this chunk. + # max(1, ...) prevents next_line from equaling start_line when a single line + # exceeds max_bytes (newline_count == 0), which would make the caller retry + # the same range indefinitely. + next_line = start_line + max(1, newline_count) + + if next_line <= total_lines: + # Truncation fell before the last line — continue reading from next_line. + read_from = next_line + elif start_line < total_lines: + # next_line overshot total_lines, meaning the cut landed inside the last line. + # Re-read from the start of the last line so the caller gets it in full. + read_from = total_lines + else: + # start_line == total_lines: the last line itself exceeds DEFAULT_MAX_BYTES. + # This case is outside our handled range — return without a truncation notice. + 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_file` with file_path={file_path or ''} " + f"start_line={read_from} to read more." + ) + + return result + notice + + +def _retruncate( + text: str, + max_bytes: int, + encoding: str, +) -> str: + """Re-truncate text that was previously truncated (contains TRUNCATION_NOTICE_MARKER). + + Extracts the original content before the marker, applies the new byte limit, and + updates the embedded notice (byte count and continuation line number) via regex. + + Returns the original text unchanged when: + - the content already fits within max_bytes (with a small slack); + - required metadata fields cannot be parsed from the existing notice. + """ + parts = text.split(TRUNCATION_NOTICE_MARKER, 1) + original_content = parts[0] + old_notice = parts[1] + + text_bytes = original_content.encode(encoding) + + # Allow a small slack to avoid unnecessary re-truncation when content is just + # barely over the limit (e.g. due to minor encoding differences). + if len(text_bytes) <= max_bytes + 100: + return text + + # Parse start_line from notice; return text unchanged if not found + start_match = re.search(r"starts at line (\d+)", old_notice) + if not start_match: + return text + start_line_parsed = int(start_match.group(1)) + + # Re-slice to the new byte limit. + # Because every line is assumed to be shorter than DEFAULT_MAX_BYTES, the cut + # always falls somewhere mid-line, so at least one complete line is preserved. + truncated_bytes = text_bytes[:max_bytes] + # errors="ignore" silently drops any incomplete multi-byte character at the cut boundary. + result = truncated_bytes.decode(encoding, errors="ignore") + # Each '\n' in result corresponds to one fully-included line; + # anything after the last '\n' is a partial line that was cut off. + newline_count = result.count("\n") + + # The next read should start at the line immediately after all complete lines. + # max(1, ...) guards against the theoretical zero-newline case + # (impossible when every line is shorter than DEFAULT_MAX_BYTES). + next_line = start_line_parsed + max(1, newline_count) + + if not re.search(r"covers the next \d+ bytes", old_notice): + return text + # _truncate_fresh always includes a continuation hint, so both fields are always present. + new_notice = re.sub(r"covers the next \d+ bytes", f"covers the next {max_bytes} bytes", old_notice) + new_notice = re.sub(r"start_line=\d+ to read more", f"start_line={next_line} to read more", new_notice) + + return result + TRUNCATION_NOTICE_MARKER + new_notice + + def truncate_text_output( text: str, start_line: int = 1, @@ -41,9 +160,9 @@ def truncate_text_output( If over limit, truncate at the last complete line that fits, allowing the next read to start from a fresh line. - If TRUNCATION_NOTICE_MARKER is already in text (previously truncated), - extract the original content, re-truncate it, and update the - truncation notice using regex. + Dispatches to :func:`_truncate_fresh` for text seen for the first time, or to + :func:`_retruncate` when the text already contains a TRUNCATION_NOTICE_MARKER + from a previous pass. Args: text: The output text to truncate. @@ -65,63 +184,16 @@ def truncate_text_output( try: if TRUNCATION_NOTICE_MARKER in text: - parts = text.split(TRUNCATION_NOTICE_MARKER, 1) - original_content = parts[0] - old_notice = parts[1] - - text_bytes = original_content.encode(encoding) - - # Allow a small slack to avoid re-truncating near-limit content - if len(text_bytes) <= max_bytes + 100: - return text - - # Parse start_line and total_lines from notice; return text unchanged if not found - start_match = re.search(r"Starting at start_line=(\d+)", old_notice) - total_match = re.search(r"Total lines: (\d+)", old_notice) - if not start_match or not total_match: - return text - start_line_parsed = int(start_match.group(1)) - total_lines_parsed = int(total_match.group(1)) - - truncated_bytes = text_bytes[:max_bytes] - result = truncated_bytes.decode(encoding, errors="ignore") - newline_count = result.count("\n") - - next_line = start_line_parsed + max(1, newline_count) - - if not re.search(r"next \d+ bytes", old_notice): - return text - has_continuation = bool(re.search(r"Use start_line=\d+", old_notice)) - new_notice = re.sub(r"next \d+ bytes", f"next {max_bytes} bytes", old_notice) - if has_continuation: - new_notice = re.sub(r"Use start_line=\d+", f"Use start_line={next_line}", new_notice) - elif next_line <= total_lines_parsed: - new_notice = re.sub(r"(Total lines: \d+)", f"\\1\nUse start_line={next_line} to continue.", new_notice) - - return result + TRUNCATION_NOTICE_MARKER + new_notice - + return _retruncate(text, max_bytes=max_bytes, encoding=encoding) else: - 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) - - continuation = f"\nUse start_line={next_line} to continue." if next_line <= total_lines else "" - notice = ( - TRUNCATION_NOTICE_MARKER - + f"\nFile: {file_path or ''}\nStarting at start_line={start_line}, next {max_bytes} bytes." - f"\nTotal lines: {total_lines}{continuation}" + return _truncate_fresh( + text, + start_line=start_line, + total_lines=total_lines, + max_bytes=max_bytes, + file_path=file_path, + encoding=encoding, ) - - return result + notice - except Exception: logger.warning("truncate_text_output failed, returning original text", exc_info=True) return text diff --git a/tests/light/test_truncate_text_output.py b/tests/light/test_truncate_text_output.py new file mode 100644 index 00000000..ba25a842 --- /dev/null +++ b/tests/light/test_truncate_text_output.py @@ -0,0 +1,381 @@ +# -*- coding: utf-8 -*- +# pylint: disable=missing-function-docstring +"""Unit tests for _truncate_fresh, _retruncate, and truncate_text_output. + +Assumptions that mirror the production code: +- Every single line is shorter than DEFAULT_MAX_BYTES. +- Lines that exceed DEFAULT_MAX_BYTES are explicitly ignored / not fully read. +- _truncate_fresh always includes a continuation hint in the notice, so + _retruncate can assume the hint is always present. + +Run: + cd tests/light && python test_truncate_text_output.py +""" + +import re + + +from reme.memory.file_based.utils.file_utils import ( + TRUNCATION_NOTICE_MARKER, + _truncate_fresh, + _retruncate, + truncate_text_output, +) + +# ── helpers ────────────────────────────────────────────────────────────────── + +LINE_BYTES = 20 # every test line is exactly this many bytes +ENC = "utf-8" + + +def make_line(i: int) -> str: + """Return a line that is exactly LINE_BYTES bytes in UTF-8. + + Format: "L{i}" padded with underscores, terminated with "\\n". + Works for i up to 999. + """ + prefix = f"L{i}" + return prefix + "_" * (LINE_BYTES - 1 - len(prefix)) + "\n" + + +def make_text(n: int, start: int = 1) -> str: + """Build n lines starting from line number `start`.""" + return "".join(make_line(i) for i in range(start, start + n)) + + +def content_of(result: str) -> str: + """Return the portion before the truncation marker.""" + return result.split(TRUNCATION_NOTICE_MARKER)[0] + + +def notice_of(result: str) -> str: + """Return the portion after the truncation marker, or '' if absent.""" + parts = result.split(TRUNCATION_NOTICE_MARKER, 1) + return parts[1] if len(parts) > 1 else "" + + +def parse_next_line(result: str) -> int: + """Extract 'start_line=X' from the notice, or -1 if absent.""" + m = re.search(r"start_line=(\d+) to read more", notice_of(result)) + return int(m.group(1)) if m else -1 + + +def parse_covers_bytes(result: str) -> int: + """Extract 'covers the next X bytes' from the notice.""" + m = re.search(r"covers the next (\d+) bytes", notice_of(result)) + return int(m.group(1)) if m else -1 + + +def assert_eq(a, b, msg=""): + assert a == b, f"{msg}: expected {b!r}, got {a!r}" + + +def assert_in(needle, haystack, msg=""): + assert needle in haystack, f"{msg}: {needle!r} not found in {haystack!r}" + + +def assert_not_in(needle, haystack, msg=""): + assert needle not in haystack, f"{msg}: {needle!r} unexpectedly found in {haystack!r}" + + +# ── _truncate_fresh tests ───────────────────────────────────────────────────── + + +def test_fresh_no_truncation_when_under_limit(): + text = make_text(3) # 60 bytes + result = _truncate_fresh(text, start_line=1, total_lines=3, max_bytes=100, file_path=None, encoding=ENC) + assert_eq(result, text, "under limit: no change") + assert_not_in(TRUNCATION_NOTICE_MARKER, result) + + +def test_fresh_no_truncation_at_exact_limit(): + text = make_text(3) # 60 bytes + result = _truncate_fresh(text, start_line=1, total_lines=3, max_bytes=60, file_path=None, encoding=ENC) + assert_eq(result, text, "at exact limit: no change") + assert_not_in(TRUNCATION_NOTICE_MARKER, result) + + +def test_fresh_mid_file_correct_next_line(): + # 10 lines × 20 bytes = 200 bytes; max_bytes=95 + # 95 bytes → 4 complete lines (80 bytes) + 15 bytes into line 5 + # newline_count=4, next_line=1+4=5, 5<=10 → read_from=5 + text = make_text(10) + result = _truncate_fresh(text, start_line=1, total_lines=10, max_bytes=95, file_path=None, encoding=ENC) + + assert_in(TRUNCATION_NOTICE_MARKER, result) + assert_eq(parse_next_line(result), 5, "next_line after mid-file cut") + assert_eq(parse_covers_bytes(result), 95) + assert_in("L4", content_of(result), "line 4 fully included") + assert_not_in("L5\n", content_of(result), "line 5 not fully included") + + +def test_fresh_notice_contains_total_lines_and_start_line(): + text = make_text(10, start=3) + result = _truncate_fresh(text, start_line=3, total_lines=12, max_bytes=95, file_path="/tmp/foo.txt", encoding=ENC) + notice = notice_of(result) + assert_in("contains 12 lines in total", notice) + assert_in("starts at line 3", notice) + assert_in("file_path=/tmp/foo.txt", notice) + + +def test_fresh_next_line_equals_total_lines_reads_from_last(): + # 5 lines × 20 = 100 bytes; max_bytes=85 + # 85 bytes → 4 complete lines (80 bytes) + 5 bytes into line 5 + # newline_count=4, next_line=1+4=5 = total_lines → read_from=5 + text = make_text(5) + result = _truncate_fresh(text, start_line=1, total_lines=5, max_bytes=85, file_path=None, encoding=ENC) + + assert_in(TRUNCATION_NOTICE_MARKER, result) + assert_eq(parse_next_line(result), 5, "should re-read the last line") + + +def test_fresh_next_line_overshoots_total_lines(): + # max_bytes=115 → 5 complete lines (100 bytes) + 15 bytes into line 6 + # newline_count=5, next_line=1+5=6 > total_lines(5), start_line(1) < 5 + # → read_from = total_lines = 5 + text = make_text(10) + result = _truncate_fresh(text, start_line=1, total_lines=5, max_bytes=115, file_path=None, encoding=ENC) + + assert_in(TRUNCATION_NOTICE_MARKER, result) + assert_eq(parse_next_line(result), 5, "overshoot: fall back to total_lines") + + +def test_fresh_long_first_line_skips_to_next_line(): + # Line 1 is 200 bytes (> max_bytes=100), no '\n' in truncated result. + # newline_count=0, next_line=1+max(1,0)=2, 2<=5 → read_from=2 + long_line = "A" * 199 + "\n" # 200 bytes + text = long_line + make_text(4, start=2) + result = _truncate_fresh(text, start_line=1, total_lines=5, max_bytes=100, file_path=None, encoding=ENC) + + assert_in(TRUNCATION_NOTICE_MARKER, result) + assert_eq(parse_next_line(result), 2, "skip to line 2 after long line 1") + assert_not_in("A" * 101, content_of(result)) + + +def test_fresh_long_middle_line_skips_to_next_line(): + # Lines 1-2 normal (40 bytes); line 3 is 200 bytes; lines 4-5 normal. + # max_bytes=100: fits lines 1-2 (40 bytes) + 60 bytes of line 3 (no '\n') + # newline_count=2, next_line=1+2=3, 3<=5 → read_from=3 + text = make_text(2, start=1) + "B" * 199 + "\n" + make_text(2, start=4) + result = _truncate_fresh(text, start_line=1, total_lines=5, max_bytes=100, file_path=None, encoding=ENC) + + assert_in(TRUNCATION_NOTICE_MARKER, result) + assert_eq(parse_next_line(result), 3, "restart from long line 3") + + +def test_fresh_long_last_line_start_equals_total_no_notice(): + # start_line == total_lines, single line too long → unhandled case, no notice. + # newline_count=0, next_line=5+1=6 > 5, start_line(5)==total_lines(5) + long_line = "C" * 199 + "\n" # 200 bytes + result = _truncate_fresh(long_line, start_line=5, total_lines=5, max_bytes=100, file_path=None, encoding=ENC) + + assert_not_in(TRUNCATION_NOTICE_MARKER, result) + assert_eq(len(result), 100, "only max_bytes bytes returned") + + +def test_fresh_long_last_line_arrived_from_previous_chunk_no_notice(): + long_line = "D" * 199 + "\n" + result = _truncate_fresh(long_line, start_line=5, total_lines=5, max_bytes=80, file_path=None, encoding=ENC) + + assert_not_in(TRUNCATION_NOTICE_MARKER, result) + + +# ── _retruncate tests ───────────────────────────────────────────────────────── + + +def _first_pass(n_lines: int = 50, max_bytes: int = 490) -> str: + """Produce a first-truncated text via _truncate_fresh.""" + return _truncate_fresh( + make_text(n_lines), + start_line=1, + total_lines=n_lines, + max_bytes=max_bytes, + file_path="/file.txt", + encoding=ENC, + ) + + +def test_retruncate_within_slack_returns_unchanged(): + # content ≈ 490 bytes; re-truncate with max_bytes=400. + # 490 <= 400+100=500 → slack hit, return unchanged. + pass1 = _first_pass(n_lines=50, max_bytes=490) + result = _retruncate(pass1, max_bytes=400, encoding=ENC) + assert_eq(result, pass1, "within slack: unchanged") + + +def test_retruncate_updates_byte_count_in_notice(): + pass1 = _first_pass(n_lines=50, max_bytes=490) + result = _retruncate(pass1, max_bytes=195, encoding=ENC) + + assert_in(TRUNCATION_NOTICE_MARKER, result) + assert_eq(parse_covers_bytes(result), 195, "notice byte count updated") + + +def test_retruncate_updates_next_line(): + # pass1 content ≈ 490 bytes (24 complete lines), start_line=1. + # re-truncate at 195 bytes → 9 complete lines, next_line=1+9=10. + pass1 = _first_pass(n_lines=50, max_bytes=490) + result = _retruncate(pass1, max_bytes=195, encoding=ENC) + + assert_eq(parse_next_line(result), 10, "next_line updated to 10") + + +def test_retruncate_content_is_smaller(): + pass1 = _first_pass(n_lines=50, max_bytes=490) + result = _retruncate(pass1, max_bytes=195, encoding=ENC) + + assert len(content_of(result)) < len(content_of(pass1)), "content should shrink" + + +def test_retruncate_missing_starts_at_line_returns_unchanged(): + # Notice missing "starts at line X" → return unchanged. + # Use large content (600 bytes) to bypass the 100-byte slack. + large_content = make_text(30) # 600 bytes + broken = ( + large_content + TRUNCATION_NOTICE_MARKER + "\nThe output above was truncated." + "\nThe full content is saved to the file and contains 30 lines in total." + "\nThis excerpt covers the next 600 bytes." # no "starts at line X" + "\nIf the current content is not enough, call `read_file` with file_path=/f.txt start_line=5 to read more." + ) + result = _retruncate(broken, max_bytes=10, encoding=ENC) + assert_eq(result, broken, "missing 'starts at line': unchanged") + + +def test_retruncate_missing_covers_next_bytes_returns_unchanged(): + # Notice has malformed "covers ??? bytes" → return unchanged. + large_content = "X" * 500 + "\n" + broken = ( + large_content + TRUNCATION_NOTICE_MARKER + "\nThe output above was truncated." + "\nThe full content is saved to the file and contains 10 lines in total." + "\nThis excerpt starts at line 1 and covers ??? bytes." + "\nIf the current content is not enough, call `read_file` with file_path=/f.txt start_line=5 to read more." + ) + result = _retruncate(broken, max_bytes=10, encoding=ENC) + assert_eq(result, broken, "missing 'covers the next N bytes': unchanged") + + +# ── truncate_text_output dispatch / guard tests ─────────────────────────────── + + +def test_dispatch_empty_string(): + result = truncate_text_output("", start_line=1, total_lines=0, max_bytes=10) + assert_eq(result, "", "empty string bypassed") + + +def test_dispatch_max_bytes_zero(): + text = make_text(5) + result = truncate_text_output(text, max_bytes=0) + assert_eq(result, text, "max_bytes=0 bypassed") + + +def test_dispatch_routes_to_fresh_when_no_marker(): + text = make_text(10) + result = truncate_text_output(text, start_line=1, total_lines=10, max_bytes=95) + assert_in(TRUNCATION_NOTICE_MARKER, result) + assert_eq(parse_next_line(result), 5) + + +def test_dispatch_routes_to_retruncate_when_marker_present(): + pass1 = _first_pass(n_lines=50, max_bytes=490) + pass2 = truncate_text_output(pass1, max_bytes=195) + assert_in(TRUNCATION_NOTICE_MARKER, pass2) + assert_eq(parse_covers_bytes(pass2), 195) + + +# ── multi-pass integration tests ────────────────────────────────────────────── + + +def test_three_pass_decreasing_truncation(): + """Three successive truncations with shrinking max_bytes.""" + n_lines = 50 + text = make_text(n_lines) # 1000 bytes + + pass1 = truncate_text_output(text, start_line=1, total_lines=n_lines, max_bytes=490, file_path="/f.txt") + assert TRUNCATION_NOTICE_MARKER in pass1 + assert parse_covers_bytes(pass1) == 490 + assert parse_next_line(pass1) > 1 + + # 490 > 195+100=295 → re-truncation proceeds + pass2 = truncate_text_output(pass1, max_bytes=195) + assert TRUNCATION_NOTICE_MARKER in pass2 + assert parse_covers_bytes(pass2) == 195 + assert parse_next_line(pass2) < parse_next_line(pass1), "next_line regresses" + assert len(content_of(pass2)) < len(content_of(pass1)) + + # content_of(pass2) ≈ 195 bytes; 195 > 90+100=190 → re-truncation proceeds + pass3 = truncate_text_output(pass2, max_bytes=90) + assert TRUNCATION_NOTICE_MARKER in pass3 + assert parse_covers_bytes(pass3) == 90 + assert parse_next_line(pass3) < parse_next_line(pass2), "next_line regresses further" + assert len(content_of(pass3)) < len(content_of(pass2)) + + +def test_three_pass_next_lines_are_consistent(): + """next_line values should monotonically decrease with each re-truncation.""" + n_lines = 50 + text = make_text(n_lines) + + pass1 = truncate_text_output(text, start_line=1, total_lines=n_lines, max_bytes=490, file_path="/f.txt") + pass2 = truncate_text_output(pass1, max_bytes=195) + pass3 = truncate_text_output(pass2, max_bytes=90) + + n1 = parse_next_line(pass1) + n2 = parse_next_line(pass2) + n3 = parse_next_line(pass3) + + assert n1 > n2 > n3 > 1, f"Expected n1 > n2 > n3 > 1, got {n1} > {n2} > {n3}" + + +# ── runner ──────────────────────────────────────────────────────────────────── + + +def run_all(): + tests = [ + # _truncate_fresh + test_fresh_no_truncation_when_under_limit, + test_fresh_no_truncation_at_exact_limit, + test_fresh_mid_file_correct_next_line, + test_fresh_notice_contains_total_lines_and_start_line, + test_fresh_next_line_equals_total_lines_reads_from_last, + test_fresh_next_line_overshoots_total_lines, + test_fresh_long_first_line_skips_to_next_line, + test_fresh_long_middle_line_skips_to_next_line, + test_fresh_long_last_line_start_equals_total_no_notice, + test_fresh_long_last_line_arrived_from_previous_chunk_no_notice, + # _retruncate + test_retruncate_within_slack_returns_unchanged, + test_retruncate_updates_byte_count_in_notice, + test_retruncate_updates_next_line, + test_retruncate_content_is_smaller, + test_retruncate_missing_starts_at_line_returns_unchanged, + test_retruncate_missing_covers_next_bytes_returns_unchanged, + # truncate_text_output dispatch / guard + test_dispatch_empty_string, + test_dispatch_max_bytes_zero, + test_dispatch_routes_to_fresh_when_no_marker, + test_dispatch_routes_to_retruncate_when_marker_present, + # multi-pass integration + test_three_pass_decreasing_truncation, + test_three_pass_next_lines_are_consistent, + ] + + passed = 0 + failed = 0 + for t in tests: + try: + t() + print(f" PASS {t.__name__}") + passed += 1 + except Exception as e: + print(f" FAIL {t.__name__}: {e}") + failed += 1 + + print(f"\n{passed} passed, {failed} failed out of {len(tests)} tests.") + return failed == 0 + + +if __name__ == "__main__": + import sys + + ok = run_all() + sys.exit(0 if ok else 1)