feat(environment): stream ranged reads in ReadFileTool - #7133
Open
chelsealong wants to merge 2 commits into
Open
chelsealong wants to merge 2 commits into
chelsealong wants to merge 2 commits into
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Link to Issue or Description of Change
Problem:
ReadFileTool.run_asyncreads the entire file into memory viaawait self._environment.read_file(path)and then splits lines across thewhole payload before slicing out the requested
start_line/end_linerange. 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
TODOinsrc/google/adk/tools/environment/_read_file_tool.pycalling this out.Solution:
BaseEnvironment.read_file_lines(path, start_line, end_line), aconcrete (non-abstract) method with a backward-compatible default that
delegates to
read_file()and slices in memory, so existingBaseEnvironmentsubclasses (e.g. the e2b and Daytona integrations) keepworking unchanged.
read_file_linesinLocalEnvironmentto stream the file inbounded chunks and only buffer the requested range, instead of reading
the whole file into a single
bytesobject. Line splitting is done by asmall scanner (
_iter_binary_lines) that matchesbytes.splitlines(keepends=True)semantics (\n,\r, and\r\nallend a line) rather than plain binary-mode iteration, which only splits on
\nand would silently disagree with theBaseEnvironmentdefault (andtherefore with e2b/Daytona) on files using bare-CR line endings.
ReadFileTool.run_asyncto callread_file_lines()instead ofread_file(), removing the TODO and the full-filesplitlines()call.Behavior (return values, error messages, line numbering, truncation) is
unchanged, including for
\r-only and mixed line endings; only the memoryprofile of a ranged local read changes.
Testing Plan
Unit Tests:
Added:
tests/unittests/environment/test_local_environment.py::TestReadFileLines— verifies
LocalEnvironment.read_file_linesreturns the correct1-based inclusive range and total line count, and (by monkeypatching the
full-buffer
_sync_readhelper to raise) that a ranged read never goesthrough the whole-file read path.
TestReadFileLines::test_line_splitting_matches_splitlines_semantics(parametrized over lone-
\rand mixed\r/\n/\r\ncontent) —proves
LocalEnvironment.read_file_lines's streamed split agrees withbytes.splitlines(keepends=True), the semanticsBaseEnvironment.read_file_lines's default implementation (used bye2b/Daytona) relies on.
tests/unittests/tools/environment/test_read_file_tool.py— extended thestub environment to implement
read_file_linesdirectly and assertread_fileis never called for a ranged read throughReadFileTool.An earlier revision of this PR iterated the file directly in binary mode
(
for line in f), which only splits on\nand silently disagreed withbytes.splitlines()on bare-\rline endings — a real correctnessregression caught in review. Fixed by replacing that iteration with
_iter_binary_lines, a chunked scanner that reproducessplitlines(keepends=True)boundary handling (\n,\r,\r\n) whilekeeping memory bounded. Verified the new parametrized test fails without
that fix (reverted only
_local_environment.pywithgit stash, keepingthe new test, then restored it):
Also re-verified the original (pre-review) assertions still fail without
the whole
read_file_linesfeature (reverted the three source files withgit stash, keeping all tests, then restored them):With the fix restored:
Full suite:
Also independently re-measured the memory claim with
tracemallocon a200k-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=10read stayed at ~281 KB, versus tens of MB for the oldwhole-file
read_file()+splitlines()path.pre-commit run --files <changed files>passes (ruff,isort,pyink,addlicense,codespell, ADK compliance checks) on every file this PRtouches.
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
Additional context
This PR was prepared with AI assistance (Claude Code).
🤖 Generated with Claude Code