Skip to content

feat(environment): stream ranged reads in ReadFileTool - #7133

Open
chelsealong wants to merge 2 commits into
google:mainfrom
chelsealong:fix/read-file-tool-stream-lines
Open

chelsealong wants to merge 2 commits into
google:mainfrom
chelsealong:fix/read-file-tool-stream-lines

Conversation

@chelsealong

Copy link
Copy Markdown
Contributor

Link to Issue or Description of Change

Problem:
ReadFileTool.run_async reads the entire file into memory via
await self._environment.read_file(path) and then splits lines across the
whole payload before slicing out the requested start_line/end_line
range. For very large files (hundreds of MB or GB), this risks high memory
usage or OOM crashes even when the caller only requests a small range of
lines. There was an explicit TODO in
src/google/adk/tools/environment/_read_file_tool.py calling this out.

Solution:

  • Added BaseEnvironment.read_file_lines(path, start_line, end_line), a
    concrete (non-abstract) method with a backward-compatible default that
    delegates to read_file() and slices in memory, so existing
    BaseEnvironment subclasses (e.g. the e2b and Daytona integrations) keep
    working unchanged.
  • Overrode read_file_lines in LocalEnvironment to stream the file in
    bounded chunks and only buffer the requested range, instead of reading
    the whole file into a single bytes object. Line splitting is done by a
    small scanner (_iter_binary_lines) that matches
    bytes.splitlines(keepends=True) semantics (\n, \r, and \r\n all
    end a line) rather than plain binary-mode iteration, which only splits on
    \n and would silently disagree with the BaseEnvironment default (and
    therefore with e2b/Daytona) on files using bare-CR line endings.
  • Updated ReadFileTool.run_async to call read_file_lines() instead of
    read_file(), removing the TODO and the full-file splitlines() call.

Behavior (return values, error messages, line numbering, truncation) is
unchanged, including for \r-only and mixed line endings; only the memory
profile of a ranged local read changes.

Testing Plan

Unit Tests:

  • I have added or updated unit tests for my change.
  • All unit tests pass locally.

Added:

  • tests/unittests/environment/test_local_environment.py::TestReadFileLines
    — verifies LocalEnvironment.read_file_lines returns the correct
    1-based inclusive range and total line count, and (by monkeypatching the
    full-buffer _sync_read helper to raise) that a ranged read never goes
    through the whole-file read path.
  • TestReadFileLines::test_line_splitting_matches_splitlines_semantics
    (parametrized over lone-\r and mixed \r/\n/\r\n content) —
    proves LocalEnvironment.read_file_lines's streamed split agrees with
    bytes.splitlines(keepends=True), the semantics
    BaseEnvironment.read_file_lines's default implementation (used by
    e2b/Daytona) relies on.
  • tests/unittests/tools/environment/test_read_file_tool.py — extended the
    stub environment to implement read_file_lines directly and assert
    read_file is never called for a ranged read through ReadFileTool.

An earlier revision of this PR iterated the file directly in binary mode
(for line in f), which only splits on \n and silently disagreed with
bytes.splitlines() on bare-\r line endings — a real correctness
regression caught in review. Fixed by replacing that iteration with
_iter_binary_lines, a chunked scanner that reproduces
splitlines(keepends=True) boundary handling (\n, \r, \r\n) while
keeping memory bounded. Verified the new parametrized test fails without
that fix (reverted only _local_environment.py with git stash, keeping
the new test, then restored it):

$ python -m pytest tests/unittests/environment/test_local_environment.py -q -k test_line_splitting_matches_splitlines_semantics
FAILED ...[lone_cr] - AssertionError: assert [] == [b'two\r', b'three\r']
FAILED ...[mixed] - AssertionError: assert [b'two\rthree\n', b'four\r\n'] == [b'two\r', b'three\n']
2 failed, 15 deselected in 0.15s

Also re-verified the original (pre-review) assertions still fail without
the whole read_file_lines feature (reverted the three source files with
git stash, keeping all tests, then restored them):

$ python -m pytest tests/unittests/tools/environment/test_read_file_tool.py tests/unittests/environment/test_local_environment.py -q
...
FAILED tests/unittests/tools/environment/test_read_file_tool.py::test_read_file_with_line_range_uses_direct_file_read - AssertionError: assert ['notes.txt'] == []
FAILED tests/unittests/environment/test_local_environment.py::TestReadFileLines::test_returns_requested_range_and_total - AttributeError: 'LocalEnvironment' object has no attribute 'read_file_lines'
FAILED tests/unittests/environment/test_local_environment.py::TestReadFileLines::test_does_not_read_the_whole_file_at_once - AttributeError: 'LocalEnvironment' object has no attribute 'read_file_lines'
3 failed, 18 passed in 2.29s

With the fix restored:

$ python -m pytest tests/unittests/tools/environment/ tests/unittests/environment/ tests/unittests/integrations/e2b tests/unittests/integrations/daytona -q
.......................................................                  [100%]
57 passed, 2 warnings in 4.66s

Full suite:

$ python -m pytest tests/unittests -n auto -q
14984 passed, 82 skipped, 27 xfailed, 2 xpassed, 2169 warnings, 28 subtests passed in 175.23s

Also independently re-measured the memory claim with tracemalloc on a
200k-line file (unaffected by this fix, since the chunked scanner keeps
the same bounded-memory property): peak bytes for a start_line=5, end_line=10 read stayed at ~281 KB, versus tens of MB for the old
whole-file read_file() + splitlines() path.

pre-commit run --files <changed files> passes (ruff, isort, pyink,
addlicense, codespell, ADK compliance checks) on every file this PR
touches.

Manual End-to-End (E2E) Tests:

Not applicable — this is an internal memory-usage change to an existing
tool with unchanged external behavior, covered by the unit tests above.

Checklist

  • I have read the CONTRIBUTING.md document.
  • I have performed a self-review of my own code.
  • I have added tests that prove my fix is effective.
  • New and existing unit tests pass locally with my changes.

Additional context

This PR was prepared with AI assistance (Claude Code).

🤖 Generated with Claude Code

ReadFileTool.run_async loaded the whole file into memory via
read_file() before slicing out the requested line range, risking OOM
on large files even when only a small range was requested.

Add read_file_lines() to BaseEnvironment, with a default fallback that
delegates to read_file() for backward compatibility, and override it
in LocalEnvironment to stream the file line-by-line so only the
requested range is buffered. ReadFileTool now calls read_file_lines()
instead of read_file().

Closes: google#7131
_sync_read_lines iterated the file with `for line in f` in binary mode,
which only splits on \n and misses bare \r (old Mac line endings) that
bytes.splitlines() treats as a line boundary. That made LocalEnvironment
report different total_lines/content than BaseEnvironment's default
splitlines()-based implementation (used by e2b/Daytona) for the same
file, contradicting the PR's claim that line numbering is unchanged.

Replace the direct iteration with a small chunked scanner that matches
bytes.splitlines(keepends=True) semantics (\n, \r, \r\n) while keeping
memory bounded to the scan buffer plus the selected range.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

feat(environment): stream lines in ReadFileTool to prevent OOM on large files

2 participants