diff --git a/reme/components/file_chunker/json_file_chunker.py b/reme/components/file_chunker/json_file_chunker.py index e8cce166..f1249260 100644 --- a/reme/components/file_chunker/json_file_chunker.py +++ b/reme/components/file_chunker/json_file_chunker.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +from io import StringIO from pathlib import Path from typing import Any @@ -137,10 +138,10 @@ class JsonFileChunker(BaseFileChunker): [chunk], ) else: - lines = raw_text.split("\n", keepends=True) + lines = StringIO(raw_text) chunked_lines = [{"text": "", "start_line": 1, "end_line": 1}] for idx, line in enumerate(lines): - if len(chunked_lines[-1]["text"]) + len(line) > self.chunk_chars: + if chunked_lines[-1]["text"] and len(chunked_lines[-1]["text"]) + len(line) > self.chunk_chars: chunked_lines.append({"text": "", "start_line": idx + 1, "end_line": idx + 1}) chunked_lines[-1]["text"] += line chunked_lines[-1]["end_line"] = idx + 1 diff --git a/tests/unit/test_background_steps.py b/tests/unit/test_background_steps.py index 87867a5b..3e5ae579 100644 --- a/tests/unit/test_background_steps.py +++ b/tests/unit/test_background_steps.py @@ -25,7 +25,7 @@ from watchfiles import Change from reme.components import R from reme.components.agent_wrapper import BaseAgentWrapper -from reme.components.file_chunker import DefaultFileChunker +from reme.components.file_chunker import DefaultFileChunker, JsonFileChunker from reme.components.file_catalog import LocalFileCatalog from reme.components.file_store import LocalFileStore from reme.components.runtime_context import RuntimeContext @@ -885,6 +885,58 @@ def test_update_index_modified_file_reuses_unchanged_embedding(): asyncio.run(run()) +def test_update_index_replaces_large_malformed_json_without_changing_source(): + """Malformed JSON uses the fallback and replaces its old chunks.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + cwd = Path.cwd() + source = cwd / "daily" / "broken.json" + sibling = cwd / "daily" / "healthy.json" + original = ("{" + "x" * 2000).encode() + replacement = ("{" + "y" * 2000).encode() + source.parent.mkdir(parents=True) + source.write_bytes(original) + sibling.write_text('{"healthy": true}', encoding="utf-8") + fs = LocalFileStore(name="default", embedding_store="") + chunker = JsonFileChunker(supported_extensions=["json"], chunk_chars=256) + await fs.start() + await chunker.start() + try: + app_ctx = _make_app_context(cwd) + app_ctx.components = {ComponentEnum.FILE_CHUNKER: {"json": chunker}} + step = UpdateIndexStep(file_store=fs, persist=False, app_context=app_ctx) + response = await step( + RuntimeContext( + changes=[ + {"change": "added", "path": str(source)}, + {"change": "added", "path": str(sibling)}, + ], + ), + ) + + assert response.success is True + assert source.read_bytes() == original + nodes = {node.path: node for node in await fs.get_nodes()} + old_chunk_ids = nodes["daily/broken.json"].chunk_ids + assert nodes["daily/healthy.json"].chunk_ids + + source.write_bytes(replacement) + response = await step(RuntimeContext(changes=[{"change": "modified", "path": str(source)}])) + + assert response.success is True + assert source.read_bytes() == replacement + nodes = {node.path: node for node in await fs.get_nodes()} + assert nodes["daily/broken.json"].chunk_ids != old_chunk_ids + assert all(chunk_id not in fs.file_chunks for chunk_id in old_chunk_ids) + assert nodes["daily/healthy.json"].chunk_ids + finally: + await chunker.close() + await fs.close() + + asyncio.run(run()) + + def test_update_index_skips_oversized_file_and_clears_stale_index(): """Oversized content is not read and any previous index entry is removed.""" diff --git a/tests/unit/test_json_file_chunker.py b/tests/unit/test_json_file_chunker.py index 4d8af06a..597da053 100644 --- a/tests/unit/test_json_file_chunker.py +++ b/tests/unit/test_json_file_chunker.py @@ -255,6 +255,35 @@ class TestEdgeCases: assert len(chunks) == 1 assert chunks[0].start_line == 1 + @pytest.mark.parametrize("length", [2000, 2001]) + def test_large_malformed_fallback_preserves_a_single_long_line(self, make_json, length): + text = "{" + "x" * (length - 1) + path = make_json(raw=text) + + node, chunks = _run(JsonFileChunker().chunk(path)) + + assert "".join(chunk.text for chunk in chunks) == text + assert all(chunk.text for chunk in chunks) + assert node.chunk_ids == [chunk.id for chunk in chunks] + assert [(chunk.start_line, chunk.end_line) for chunk in chunks] == [(1, 1)] + + @pytest.mark.parametrize("ending", ["", "\n", "\r\n"]) + def test_large_malformed_fallback_preserves_lf_lines_and_endings(self, make_json, ending): + text = "first\r\n" + "é\u2028middle\u2029" + "x" * 300 + ending + path = make_json(raw="") + with open(path, "wb") as file: + file.write(text.encode()) + + node, chunks = _run(JsonFileChunker(chunk_chars=256).chunk(path)) + + with open(path, "rb") as file: + assert file.read() == text.encode() + with open(path, encoding="utf-8") as file: + assert "".join(chunk.text for chunk in chunks) == file.read() + assert all(chunk.text for chunk in chunks) + assert node.chunk_ids == [chunk.id for chunk in chunks] + assert [(chunk.start_line, chunk.end_line) for chunk in chunks] == [(1, 1), (2, 2)] + def test_chunk_chars_floor(self): assert JsonFileChunker(chunk_chars=10).chunk_chars == 256