mirror of
https://github.com/badlogic/pi-mono.git
synced 2026-11-05 23:37:25 +00:00
fix(coding-agent): stream the read tool instead of loading whole files (#11750)
## Problem
`read` loads the whole file before it applies `offset`/`limit`:
`ops.readFile` -> `toString("utf-8")` -> `split("\n")`. A one-line read
still holds the whole file three times: as a Buffer, as a string and as
an array of lines.
Measured on main with a 200 MB log (2.36M lines):
- `offset: 1, limit: 1` peaks at 638 MB RSS. With
`--max-old-space-size=256` it crashes in `Runtime_StringSplit` (OOM).
- Files over about 512 MB fail with `Cannot create a string longer than
0x1fffffe8 characters`. Files over 2 GB fail with
`ERR_FS_FILE_TOO_LARGE`.
- Aborting only rejects the promise; the read keeps running in the
background.
## Solution
- `ReadOperations` gets a new optional `createReadStream(absolutePath, {
signal })`, which returns an `AsyncIterable<Uint8Array>` (a Node
`Readable` works). The default local operations open the file once,
`fstat` the handle, and stream from it. Like `fs.readFile`, they stop a
regular file at its size when opened: text appended during the read is
not counted, and a file that keeps growing still finishes. A file that
reports size 0 (such as `/proc` files) and other file types (FIFOs,
devices) are read until EOF, as `fs.readFile` reads them. No existing
member changed.
- The text path decodes the chunks with a streaming UTF-8 `TextDecoder`
(non-fatal, `ignoreBOM: true`). That matches `Buffer.toString("utf-8")`,
including invalid bytes and a UTF-8 byte order mark kept as U+FEFF.
- It counts every line but stores only the selected lines `truncateHead`
can show. Storing stops once the stored lines are past
`DEFAULT_MAX_LINES` or `DEFAULT_MAX_BYTES`, plus one more line, and the
existing `truncateHead` runs on the stored lines.
`totalLines`/`totalBytes` are then patched from the counts, as grep and
find do on main.
- The extra line is needed because `truncateHead` drops a trailing empty
line. Example: 2001 short lines, then an empty line, then more text. If
storing stopped at the empty line, the stored text would end in `\n`.
`truncateHead` would count 2000 lines and report "not truncated", but
the full text is truncated by lines.
- A line over 50KB is never stored whole; only its size is tracked. A
first line of 1 GB therefore still gives today's `[Line 1 is ...,
exceeds 50.0KB limit ...]` notice.
- Abort: the signal is passed to the stream, and the loop also checks it
between chunks. Leaving the loop closes the stream.
- The whole-file path stays unchanged for custom operations without
`createReadStream`, and for fractional `offset`/`limit` or negative
`limit`. Array slicing defines what those odd ranges select; for
example, `limit: -1` selects all lines but the last.
- Images are unchanged.
Output is the same as main's, byte for byte, including line counts,
notices, details (including key order) and errors. A file appended to
during the read gives the same output too: both read a regular file up
to its size when opened. Out of scope: the trailing-newline off-by-one
in line counts from #7329 is kept as is.
## Testing
- `test/read-streaming.test.ts` keeps a copy of the old text algorithm
as the reference. It compares the streaming path against it on text,
details and error messages, with chunks of 1 and 4096 bytes and offsets
and limits around the boundaries. The cases are: empty file, trailing
newline or not, CRLF, multi-byte characters and a UTF-8 BOM, 2001 lines,
and a first line over 50KB. There are also edge cases where the line or
byte cap falls exactly on an empty line, and the 50KB + trailing newline
case that `truncateHead` reports as `truncatedBy: "lines"`. Further
tests:
- details key order
- multi-byte characters split across chunks, and invalid UTF-8
- the `readFile` fallback, including fractional and negative ranges
- local files through the default stream, including EISDIR
- abort: the stream is closed and no more chunks are pulled
- a local file appended to mid-read (a decode hook in the test appends
after the first chunk): `limit: 1` still says `[1 more lines ...]`. This
test fails with an unbounded `fs.createReadStream`.
- bounded memory: a child process with `--max-old-space-size=64` reads a
200 MB single line, generated in memory as 64 KB chunks from a custom
`createReadStream`, and must print the size notice. It takes about 0.8 s
and uses 17 MB of heap. If over-long lines are stored whole, it runs out
of memory. 64 KB chunks matter here: V8 keeps decoded strings of 1 MB
outside the heap limit, so with 1 MB chunks the test would not catch
that regression.
- With main's `read.ts`, 5 of the 8 tests fail: matrix, key order,
chunked decoding, bounded memory and abort. The fallback, local-file and
append tests pass on main because they check behavior that did not
change. Removing the one-line-past-the-cap rule, or the
trailing-empty-line adjustment of `totalLines`, makes the matrix fail.
- `test/tools.test.ts`, `read-line-range-display`,
`tool-execution-component`, `block-images` and `image-resize-callers`
pass. `npm run check` is clean.
Measured with an ad-hoc script on a 200 MB log. Baseline RSS after
loading modules is about 121 MB.
| Case | main | this PR |
| --- | --- | --- |
| offset 1, limit 1 | 1028 ms, 638 MB peak RSS | 478 ms, 131 MB |
| offset 2000000, limit 10 | 926 ms, 637 MB | 486 ms, 132 MB |
| no offset/limit (counts all lines) | 1561 ms, 1010 MB | 524 ms, 132 MB
|
| offset 1, limit 1, 256 MB heap | OOM crash | 131 MB |
| 1.1 GB single-line file | string length error | 2.2 s, 130 MB,
line-size notice |
Remaining difference: custom operations that only implement `readFile`
keep the whole-file memory use, and their reads continue after an abort,
as before.
This PR is AI-generated.
This commit is contained in:
Petr Houška
GitHub
parent
43fb4e3106
commit
3a7bb6f48c
4 changed files with 459 additions and 24 deletions
|
|
@ -28,6 +28,7 @@
|
|||
- Fixed the grep tool reading every matched file into memory when `context` was set, which ran out of memory on large files; context now comes from ripgrep, which also fixes UTF-16 files and lone CRs in matched lines, and the match limit notice no longer appears when exactly `limit` matches exist ([#9276](https://github.com/earendil-works/pi/issues/9276), [#11423](https://github.com/earendil-works/pi/issues/11423))
|
||||
- Fixed the find tool reporting the result limit when exactly `limit` files match, and `limit: 0` in find and ls returning every result or an empty listing; limits below 1 now mean 1 ([#7121](https://github.com/earendil-works/pi/issues/7121))
|
||||
- Fixed grep match lines, codemode output, compaction summaries of tool results, and MCP error and server summaries cutting an emoji or other non-BMP character in half when truncating, which left an invalid lone surrogate in the text ([#7121](https://github.com/earendil-works/pi/issues/7121))
|
||||
- Fixed the read tool loading whole files into memory, which ran out of memory on large files even with `offset`/`limit`; it now streams the file and keeps only the shown lines. Custom `ReadOperations` can stream too by providing the new optional `createReadStream` ([#9654](https://github.com/earendil-works/pi/issues/9654))
|
||||
|
||||
|
||||
## [1.5.0] - 2026-11-02
|
||||
|
|
|
|||
|
|
@ -2414,6 +2414,8 @@ pi.registerTool({
|
|||
});
|
||||
```
|
||||
|
||||
`ReadOperations.createReadStream` is optional. When it returns the file as an async iterable of chunks (a Node `Readable` works), `read` keeps only the lines it shows in memory, so large files can be read with `offset`/`limit`; without it, `read` loads the whole file through `readFile`.
|
||||
|
||||
The `create*ToolDefinition()` factories, such as `createReadToolDefinition()`, return complete tool definitions, including the built-in prompt snippet and guidelines, for the same purpose. Rendering does not need copying: a tool registered under a built-in name inherits the built-in renderer for each slot it leaves out (see [Overriding Built-in Tools](#overriding-built-in-tools)).
|
||||
|
||||
**Operations interfaces:** `ReadOperations`, `WriteOperations`, `EditOperations`, `BashOperations`, `PowerShellOperations`, `LsOperations`, `GrepOperations`, `FindOperations`
|
||||
|
|
|
|||
|
|
@ -1,7 +1,8 @@
|
|||
import type { Api, ImageContent, Model, ModelImageResizeOptions, TextContent } from "@earendil-works/pi-ai";
|
||||
import { constants } from "fs";
|
||||
import { access as fsAccess, readFile as fsReadFile } from "fs/promises";
|
||||
import { access as fsAccess, open as fsOpen, readFile as fsReadFile } from "fs/promises";
|
||||
import { type Static, Type } from "typebox";
|
||||
import { TextDecoder } from "util";
|
||||
import { processImage } from "../../utils/image-process.ts";
|
||||
import { detectSupportedImageMimeTypeFromFile } from "../../utils/mime.ts";
|
||||
import type { ExtensionContext, ToolDefinition } from "../extensions/types.ts";
|
||||
|
|
@ -51,14 +52,37 @@ export interface ReadOperations {
|
|||
access: (absolutePath: string) => Promise<void>;
|
||||
/** Detect image MIME type, return null or undefined for non-images */
|
||||
detectImageMimeType?: (absolutePath: string) => Promise<string | null | undefined>;
|
||||
/**
|
||||
* Stream file contents in chunks. When present, text files are read through it and only the lines
|
||||
* that will be shown are kept in memory; otherwise readFile loads the whole file. Stop reading when
|
||||
* the signal aborts; the read tool also stops consuming the stream then.
|
||||
*/
|
||||
createReadStream?: (absolutePath: string, options: { signal?: AbortSignal }) => AsyncIterable<Uint8Array>;
|
||||
}
|
||||
|
||||
const defaultReadOperations: ReadOperations = {
|
||||
readFile: (path) => fsReadFile(path),
|
||||
createReadStream: readLocalFileChunks,
|
||||
access: (path) => fsAccess(path, constants.R_OK),
|
||||
detectImageMimeType: detectSupportedImageMimeTypeFromFile,
|
||||
};
|
||||
|
||||
/**
|
||||
* Stream a local file the way fs.readFile reads it: a regular file up to its size when opened, so text
|
||||
* appended during the read is not counted and a growing file still ends; a file reporting size 0 (such
|
||||
* as /proc files) and other file types until EOF.
|
||||
*/
|
||||
async function* readLocalFileChunks(path: string, { signal }: { signal?: AbortSignal }): AsyncIterable<Uint8Array> {
|
||||
const handle = await fsOpen(path, "r");
|
||||
try {
|
||||
const stats = await handle.stat();
|
||||
const end = stats.isFile() && stats.size > 0 ? stats.size - 1 : undefined;
|
||||
yield* handle.createReadStream({ end, signal, autoClose: false });
|
||||
} finally {
|
||||
await handle.close();
|
||||
}
|
||||
}
|
||||
|
||||
export interface ReadToolOptions {
|
||||
/** Whether to auto-resize images. Default: true */
|
||||
autoResizeImages?: boolean;
|
||||
|
|
@ -82,6 +106,125 @@ function getNonVisionImageNote(model: Model<Api> | undefined): string | undefine
|
|||
return "[Current model does not support images. The image will be omitted from this request.]";
|
||||
}
|
||||
|
||||
/** The lines a read shows, with the file's line count; line counts follow `text.split("\n")`. */
|
||||
interface SelectedLines {
|
||||
totalFileLines: number;
|
||||
truncation: TruncationResult;
|
||||
/** Bytes of the first selected line, for the notice when it alone exceeds the byte limit. */
|
||||
firstLineBytes: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Whether a range can be streamed. Other values (fractional, negative limit) keep the whole-file path,
|
||||
* whose array slicing defines what they select, for example limit -1 as all lines but the last.
|
||||
*/
|
||||
function isStreamableRange(offset: number | undefined, limit: number | undefined): boolean {
|
||||
return (
|
||||
(offset === undefined || Number.isInteger(offset)) &&
|
||||
(limit === undefined || (Number.isInteger(limit) && limit >= 0))
|
||||
);
|
||||
}
|
||||
|
||||
/** Select lines from a whole file in memory: for operations without createReadStream and unusual ranges. */
|
||||
function selectLinesFromBuffer(buffer: Buffer, startLine: number, limit: number | undefined): SelectedLines {
|
||||
const allLines = buffer.toString("utf-8").split("\n");
|
||||
const endLine = limit === undefined ? undefined : Math.min(startLine + limit, allLines.length);
|
||||
const truncation = truncateHead(allLines.slice(startLine, endLine).join("\n"));
|
||||
const firstLineBytes = truncation.firstLineExceedsLimit ? Buffer.byteLength(allLines[startLine], "utf-8") : 0;
|
||||
return { totalFileLines: allLines.length, truncation, firstLineBytes };
|
||||
}
|
||||
|
||||
/**
|
||||
* Select lines from a file read in chunks. Only lines truncateHead can show are stored; the rest are
|
||||
* counted, so memory does not grow with the file. Stops reading when the signal aborts.
|
||||
*/
|
||||
async function selectLinesFromStream(
|
||||
chunks: AsyncIterable<Uint8Array>,
|
||||
startLine: number,
|
||||
limit: number | undefined,
|
||||
signal: AbortSignal | undefined,
|
||||
): Promise<SelectedLines> {
|
||||
const endLine = limit === undefined ? Number.POSITIVE_INFINITY : startLine + limit;
|
||||
// Non-fatal UTF-8 that keeps a byte order mark decodes like Buffer.toString("utf-8").
|
||||
const decoder = new TextDecoder("utf-8", { ignoreBOM: true });
|
||||
const storedLines: string[] = [];
|
||||
let storedBytes = 0;
|
||||
let storedPastCaps = false;
|
||||
let line = 0;
|
||||
let selectedLines = 0;
|
||||
let selectedBytes = 0;
|
||||
let firstLineBytes = 0;
|
||||
let lastLineEmpty = true;
|
||||
// The selected line being read: its text, while it is stored, and its size.
|
||||
let pieces: string[] | undefined = [];
|
||||
let lineBytes = 0;
|
||||
let lineHasText = false;
|
||||
|
||||
const appendToLine = (text: string, start: number, end: number) => {
|
||||
lineHasText = true;
|
||||
if (!pieces) return;
|
||||
const piece = text.slice(start, end);
|
||||
// A line over the byte cap is never shown, and a first line over it is only measured: never store
|
||||
// all of a 1 GB line.
|
||||
if (lineBytes <= DEFAULT_MAX_BYTES) pieces.push(piece);
|
||||
lineBytes += Buffer.byteLength(piece, "utf-8");
|
||||
};
|
||||
const finishLine = () => {
|
||||
if (pieces) {
|
||||
if (selectedLines === 0) firstLineBytes = lineBytes;
|
||||
// Storing continues one line past the caps: truncateHead ignores a trailing empty line, so the
|
||||
// line after the cap must be there for it to decide truncation as for the full text.
|
||||
if (storedLines.length > DEFAULT_MAX_LINES || storedBytes > DEFAULT_MAX_BYTES) storedPastCaps = true;
|
||||
const text = pieces.join("");
|
||||
storedBytes += Buffer.byteLength(text, "utf-8") + (storedLines.length > 0 ? 1 : 0);
|
||||
storedLines.push(text);
|
||||
}
|
||||
selectedLines++;
|
||||
lastLineEmpty = !lineHasText;
|
||||
pieces = storedPastCaps ? undefined : [];
|
||||
lineBytes = 0;
|
||||
lineHasText = false;
|
||||
};
|
||||
const select = (text: string) => {
|
||||
let position = 0;
|
||||
// Start of the selected text in this piece, measured as a whole: the joined lines' size.
|
||||
let selectionStart = -1;
|
||||
while (true) {
|
||||
const selected = line >= startLine && line < endLine;
|
||||
if (selected && selectionStart === -1) selectionStart = position;
|
||||
const newline = text.indexOf("\n", position);
|
||||
const end = newline === -1 ? text.length : newline;
|
||||
if (selected && end > position) appendToLine(text, position, end);
|
||||
if (newline === -1) break;
|
||||
if (selected) {
|
||||
finishLine();
|
||||
if (line + 1 === endLine) {
|
||||
selectedBytes += Buffer.byteLength(text.slice(selectionStart, newline), "utf-8");
|
||||
selectionStart = -1;
|
||||
}
|
||||
}
|
||||
line++;
|
||||
position = newline + 1;
|
||||
}
|
||||
if (selectionStart !== -1) selectedBytes += Buffer.byteLength(text.slice(selectionStart), "utf-8");
|
||||
};
|
||||
|
||||
for await (const chunk of chunks) {
|
||||
// Leaving the loop closes the stream.
|
||||
if (signal?.aborted) throw new Error("Operation aborted");
|
||||
select(decoder.decode(chunk, { stream: true }));
|
||||
}
|
||||
select(decoder.decode());
|
||||
if (line >= startLine && line < endLine) finishLine();
|
||||
const truncation: TruncationResult = {
|
||||
...truncateHead(storedLines.join("\n")),
|
||||
// truncateHead's counts for the joined selected lines, which drop a trailing empty line.
|
||||
totalLines: selectedBytes === 0 ? 0 : selectedLines - (lastLineEmpty ? 1 : 0),
|
||||
totalBytes: selectedBytes,
|
||||
};
|
||||
return { totalFileLines: line + 1, truncation, firstLineBytes };
|
||||
}
|
||||
|
||||
export function createReadToolDefinition(
|
||||
cwd: string,
|
||||
options?: ReadToolOptions,
|
||||
|
|
@ -150,34 +293,30 @@ export function createReadToolDefinition(
|
|||
];
|
||||
}
|
||||
} else {
|
||||
// Read text content.
|
||||
const buffer = await ops.readFile(absolutePath);
|
||||
const textContent = buffer.toString("utf-8");
|
||||
const allLines = textContent.split("\n");
|
||||
const totalFileLines = allLines.length;
|
||||
// Apply offset if specified. Convert from 1-indexed input to 0-indexed array access.
|
||||
// Read text content. Convert offset from 1-indexed input to 0-indexed line access.
|
||||
const startLine = offset ? Math.max(0, offset - 1) : 0;
|
||||
const startLineDisplay = startLine + 1;
|
||||
const { totalFileLines, truncation, firstLineBytes } =
|
||||
ops.createReadStream && isStreamableRange(offset, limit)
|
||||
? await selectLinesFromStream(
|
||||
ops.createReadStream(absolutePath, { signal }),
|
||||
startLine,
|
||||
limit,
|
||||
signal,
|
||||
)
|
||||
: selectLinesFromBuffer(await ops.readFile(absolutePath), startLine, limit);
|
||||
if (aborted) return;
|
||||
// Check if offset is out of bounds.
|
||||
if (startLine >= allLines.length) {
|
||||
throw new Error(`Offset ${offset} is beyond end of file (${allLines.length} lines total)`);
|
||||
if (startLine >= totalFileLines) {
|
||||
throw new Error(`Offset ${offset} is beyond end of file (${totalFileLines} lines total)`);
|
||||
}
|
||||
let selectedContent: string;
|
||||
let userLimitedLines: number | undefined;
|
||||
// If limit is specified by the user, honor it first. Otherwise truncateHead decides.
|
||||
if (limit !== undefined) {
|
||||
const endLine = Math.min(startLine + limit, allLines.length);
|
||||
selectedContent = allLines.slice(startLine, endLine).join("\n");
|
||||
userLimitedLines = endLine - startLine;
|
||||
} else {
|
||||
selectedContent = allLines.slice(startLine).join("\n");
|
||||
}
|
||||
// Apply truncation, respecting both line and byte limits.
|
||||
const truncation = truncateHead(selectedContent);
|
||||
// A user-specified limit is honored before truncation.
|
||||
const userLimitedLines =
|
||||
limit === undefined ? undefined : Math.min(startLine + limit, totalFileLines) - startLine;
|
||||
let outputText: string;
|
||||
if (truncation.firstLineExceedsLimit) {
|
||||
// First line alone exceeds the byte limit. Point the model at a bash fallback.
|
||||
const firstLineSize = formatSize(Buffer.byteLength(allLines[startLine], "utf-8"));
|
||||
const firstLineSize = formatSize(firstLineBytes);
|
||||
outputText = `[Line ${startLineDisplay} is ${firstLineSize}, exceeds ${formatSize(DEFAULT_MAX_BYTES)} limit. Use bash: sed -n '${startLineDisplay}p' ${path} | head -c ${DEFAULT_MAX_BYTES}]`;
|
||||
details = { truncation };
|
||||
} else if (truncation.truncated) {
|
||||
|
|
@ -191,9 +330,9 @@ export function createReadToolDefinition(
|
|||
outputText += `\n\n[Showing lines ${startLineDisplay}-${endLineDisplay} of ${totalFileLines} (${formatSize(DEFAULT_MAX_BYTES)} limit). Use offset=${nextOffset} to continue.]`;
|
||||
}
|
||||
details = { truncation };
|
||||
} else if (userLimitedLines !== undefined && startLine + userLimitedLines < allLines.length) {
|
||||
} else if (userLimitedLines !== undefined && startLine + userLimitedLines < totalFileLines) {
|
||||
// User-specified limit stopped early, but the file still has more content.
|
||||
const remaining = allLines.length - (startLine + userLimitedLines);
|
||||
const remaining = totalFileLines - (startLine + userLimitedLines);
|
||||
const nextOffset = startLine + userLimitedLines + 1;
|
||||
outputText = `${truncation.content}\n\n[${remaining} more lines in file. Use offset=${nextOffset} to continue.]`;
|
||||
} else {
|
||||
|
|
|
|||
293
packages/coding-agent/test/read-streaming.test.ts
Normal file
293
packages/coding-agent/test/read-streaming.test.ts
Normal file
|
|
@ -0,0 +1,293 @@
|
|||
import { execFileSync } from "child_process";
|
||||
import { appendFileSync, mkdtempSync, rmSync, writeFileSync } from "fs";
|
||||
import { tmpdir } from "os";
|
||||
import { join, resolve } from "path";
|
||||
import { pathToFileURL } from "url";
|
||||
import type * as util from "util";
|
||||
import { afterAll, describe, expect, it, vi } from "vitest";
|
||||
import { createReadToolDefinition, type ReadOperations } from "../src/core/tools/read.ts";
|
||||
import { DEFAULT_MAX_BYTES, formatSize, type TruncationResult, truncateHead } from "../src/core/tools/truncate.ts";
|
||||
|
||||
// Runs when the read tool decodes a chunk: the file is open and partly read, so a test can change it mid-read.
|
||||
const decodeHook = vi.hoisted(() => ({ run: undefined as (() => void) | undefined }));
|
||||
vi.mock("util", async (importOriginal) => {
|
||||
const actual = await importOriginal<typeof util>();
|
||||
class HookedTextDecoder extends actual.TextDecoder {
|
||||
decode(...args: Parameters<util.TextDecoder["decode"]>): string {
|
||||
decodeHook.run?.();
|
||||
return super.decode(...args);
|
||||
}
|
||||
}
|
||||
return { ...actual, TextDecoder: HookedTextDecoder };
|
||||
});
|
||||
|
||||
// The read tool's text path before it streamed (#9654), kept as the reference the new path must match.
|
||||
function referenceRead(
|
||||
textContent: string,
|
||||
path: string,
|
||||
offset: number | undefined,
|
||||
limit: number | undefined,
|
||||
): { text: string; details: { truncation: TruncationResult } | undefined } {
|
||||
const allLines = textContent.split("\n");
|
||||
const totalFileLines = allLines.length;
|
||||
const startLine = offset ? Math.max(0, offset - 1) : 0;
|
||||
const startLineDisplay = startLine + 1;
|
||||
if (startLine >= allLines.length) {
|
||||
throw new Error(`Offset ${offset} is beyond end of file (${allLines.length} lines total)`);
|
||||
}
|
||||
let selectedContent: string;
|
||||
let userLimitedLines: number | undefined;
|
||||
if (limit !== undefined) {
|
||||
const endLine = Math.min(startLine + limit, allLines.length);
|
||||
selectedContent = allLines.slice(startLine, endLine).join("\n");
|
||||
userLimitedLines = endLine - startLine;
|
||||
} else {
|
||||
selectedContent = allLines.slice(startLine).join("\n");
|
||||
}
|
||||
const truncation = truncateHead(selectedContent);
|
||||
let outputText: string;
|
||||
let details: { truncation: TruncationResult } | undefined;
|
||||
if (truncation.firstLineExceedsLimit) {
|
||||
const firstLineSize = formatSize(Buffer.byteLength(allLines[startLine], "utf-8"));
|
||||
outputText = `[Line ${startLineDisplay} is ${firstLineSize}, exceeds ${formatSize(DEFAULT_MAX_BYTES)} limit. Use bash: sed -n '${startLineDisplay}p' ${path} | head -c ${DEFAULT_MAX_BYTES}]`;
|
||||
details = { truncation };
|
||||
} else if (truncation.truncated) {
|
||||
const endLineDisplay = startLineDisplay + truncation.outputLines - 1;
|
||||
const nextOffset = endLineDisplay + 1;
|
||||
outputText = truncation.content;
|
||||
if (truncation.truncatedBy === "lines") {
|
||||
outputText += `\n\n[Showing lines ${startLineDisplay}-${endLineDisplay} of ${totalFileLines}. Use offset=${nextOffset} to continue.]`;
|
||||
} else {
|
||||
outputText += `\n\n[Showing lines ${startLineDisplay}-${endLineDisplay} of ${totalFileLines} (${formatSize(DEFAULT_MAX_BYTES)} limit). Use offset=${nextOffset} to continue.]`;
|
||||
}
|
||||
details = { truncation };
|
||||
} else if (userLimitedLines !== undefined && startLine + userLimitedLines < allLines.length) {
|
||||
const remaining = allLines.length - (startLine + userLimitedLines);
|
||||
const nextOffset = startLine + userLimitedLines + 1;
|
||||
outputText = `${truncation.content}\n\n[${remaining} more lines in file. Use offset=${nextOffset} to continue.]`;
|
||||
} else {
|
||||
outputText = truncation.content;
|
||||
}
|
||||
return { text: outputText, details };
|
||||
}
|
||||
|
||||
type ReadOutcome = { text: string; details: unknown } | { error: string };
|
||||
|
||||
const FILE_PATH = join(tmpdir(), "pi-read-streaming-test-missing", "file.txt");
|
||||
|
||||
function referenceOutcome(text: string, offset: number | undefined, limit: number | undefined): ReadOutcome {
|
||||
try {
|
||||
return referenceRead(text, FILE_PATH, offset, limit);
|
||||
} catch (error) {
|
||||
return { error: (error as Error).message };
|
||||
}
|
||||
}
|
||||
|
||||
async function outcome(
|
||||
run: () => Promise<{ content: { type: string; text?: string }[]; details: unknown }>,
|
||||
): Promise<ReadOutcome> {
|
||||
try {
|
||||
const result = await run();
|
||||
return { text: result.content[0].text ?? "", details: result.details };
|
||||
} catch (error) {
|
||||
return { error: (error as Error).message };
|
||||
}
|
||||
}
|
||||
|
||||
/** Splits bytes into chunks of `size`, so characters and lines split at chunk ends. */
|
||||
async function* chunked(bytes: Buffer, size: number): AsyncIterable<Uint8Array> {
|
||||
for (let position = 0; position < bytes.length; position += size) {
|
||||
yield bytes.subarray(position, position + size);
|
||||
}
|
||||
}
|
||||
|
||||
function streamingOps(bytes: Buffer, chunkSize: number): ReadOperations {
|
||||
return {
|
||||
readFile: async () => {
|
||||
throw new Error("readFile must not be used when createReadStream is present");
|
||||
},
|
||||
access: async () => {},
|
||||
createReadStream: () => chunked(bytes, chunkSize),
|
||||
};
|
||||
}
|
||||
|
||||
function readWith(ops: ReadOperations, offset: number | undefined, limit: number | undefined): Promise<ReadOutcome> {
|
||||
const tool = createReadToolDefinition(process.cwd(), { operations: ops });
|
||||
return outcome(() =>
|
||||
tool.execute("call", { path: FILE_PATH, offset, limit }, undefined, undefined, undefined as never),
|
||||
);
|
||||
}
|
||||
|
||||
function lines(count: number, make: (i: number) => string): string {
|
||||
return Array.from({ length: count }, (_, i) => make(i)).join("\n");
|
||||
}
|
||||
|
||||
const FIFTY_KB_OF_LINES = `${lines(512, () => "y".repeat(99))}y`;
|
||||
|
||||
// Inputs around the boundaries the notices and truncation depend on.
|
||||
const TEXT_CASES: Record<string, string> = {
|
||||
empty: "",
|
||||
"no trailing newline": "alpha\nbeta\ngamma",
|
||||
"trailing newline": "alpha\nbeta\n",
|
||||
CRLF: "alpha\r\nbeta\r\n\r\ngamma\r\n",
|
||||
"multi-byte characters and UTF-8 byte order mark": "\uFEFFé€😀\nñ\u{10348}字\n😀😀😀",
|
||||
"2001 lines": lines(2001, (i) => `line ${i}`),
|
||||
// truncateHead drops a trailing empty line: the empty line at the line cap must not end the stored lines.
|
||||
"2001 lines, an empty line and more": `${lines(2001, (i) => `line ${i}`)}\n\nmore`,
|
||||
"first line over 50KB": `${"é".repeat(DEFAULT_MAX_BYTES / 2 + 1)}\nsecond`,
|
||||
"line over 50KB after short lines": `a\nb\n${"z".repeat(60 * 1024)}\nc`,
|
||||
// 50KB + 1 byte: truncated, yet every line fits, so truncateHead reports truncatedBy "lines".
|
||||
"exactly 50KB and a trailing newline": `${FIFTY_KB_OF_LINES}\n`,
|
||||
"exactly 50KB, an empty line and more": `${FIFTY_KB_OF_LINES}\n\nmore`,
|
||||
"multi-byte characters over 50KB": lines(700, (i) => `${i} ${"€😀".repeat(13)}`),
|
||||
};
|
||||
|
||||
describe("read tool streaming", () => {
|
||||
const tempDir = mkdtempSync(join(tmpdir(), "pi-read-streaming-"));
|
||||
afterAll(() => rmSync(tempDir, { recursive: true, force: true }));
|
||||
|
||||
it("gives the same output as the whole-file algorithm", async () => {
|
||||
for (const [name, text] of Object.entries(TEXT_CASES)) {
|
||||
const bytes = Buffer.from(text, "utf-8");
|
||||
const total = text.split("\n").length;
|
||||
for (const offset of [undefined, 2, 3, total, total + 1]) {
|
||||
for (const limit of [undefined, 0, 1, total]) {
|
||||
const expected = referenceOutcome(text, offset, limit);
|
||||
for (const chunkSize of [1, 4096]) {
|
||||
const label = `${name}, offset ${offset}, limit ${limit}, chunks of ${chunkSize}`;
|
||||
expect(await readWith(streamingOps(bytes, chunkSize), offset, limit), label).toEqual(expected);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
it("keeps details key order, which session files serialize", async () => {
|
||||
const text = TEXT_CASES["2001 lines"];
|
||||
const streamed = (await readWith(streamingOps(Buffer.from(text), 4096), 1, undefined)) as { details: unknown };
|
||||
const expected = referenceOutcome(text, 1, undefined) as { details: unknown };
|
||||
expect(JSON.stringify(streamed.details)).toBe(JSON.stringify(expected.details));
|
||||
});
|
||||
|
||||
it("decodes multi-byte characters split across chunks and invalid UTF-8 like Buffer.toString", async () => {
|
||||
const bytes = Buffer.concat([
|
||||
Buffer.from("aé€😀\n😀b\u{10348}\n", "utf-8"),
|
||||
Buffer.from([0x61, 0xc3, 0x0a, 0xe2, 0x82, 0x0a, 0xff, 0xfe, 0x80, 0xf0, 0x9f, 0x98, 0x0a, 0xed, 0xa0, 0x80]),
|
||||
]);
|
||||
for (const chunkSize of [1, 2, 3]) {
|
||||
expect(await readWith(streamingOps(bytes, chunkSize), undefined, undefined)).toEqual(
|
||||
referenceOutcome(bytes.toString("utf-8"), undefined, undefined),
|
||||
);
|
||||
}
|
||||
});
|
||||
|
||||
it("uses readFile for operations without createReadStream and for fractional or negative ranges", async () => {
|
||||
const text = "one\ntwo\nthree\nfour";
|
||||
const ops: ReadOperations = { readFile: async () => Buffer.from(text), access: async () => {} };
|
||||
const withStream: ReadOperations = { ...streamingOps(Buffer.from(text), 1), readFile: ops.readFile };
|
||||
for (const [offset, limit] of [
|
||||
[undefined, undefined],
|
||||
[2, 1],
|
||||
[9, undefined],
|
||||
] as const) {
|
||||
expect(await readWith(ops, offset, limit)).toEqual(referenceOutcome(text, offset, limit));
|
||||
}
|
||||
// Array slicing defines these: limit -1 selects all lines but the last.
|
||||
for (const [offset, limit] of [
|
||||
[2.5, 1],
|
||||
[undefined, -1],
|
||||
[1, 1.5],
|
||||
] as const) {
|
||||
expect(await readWith(withStream, offset, limit)).toEqual(referenceOutcome(text, offset, limit));
|
||||
}
|
||||
});
|
||||
|
||||
it("reads local files through the default stream", async () => {
|
||||
const path = join(tempDir, "local.txt");
|
||||
const text = lines(3000, (i) => `${i} ${"w".repeat(i % 40)}`);
|
||||
writeFileSync(path, text);
|
||||
const tool = createReadToolDefinition(tempDir);
|
||||
for (const [offset, limit] of [
|
||||
[undefined, undefined],
|
||||
[1500, 10],
|
||||
] as const) {
|
||||
const result = await outcome(() =>
|
||||
tool.execute("call", { path, offset, limit }, undefined, undefined, undefined as never),
|
||||
);
|
||||
expect(result).toEqual(referenceRead(text, path, offset, limit));
|
||||
}
|
||||
await expect(tool.execute("call", { path: tempDir }, undefined, undefined, undefined as never)).rejects.toThrow(
|
||||
"EISDIR",
|
||||
);
|
||||
});
|
||||
|
||||
it("reads a local file up to its size when opened, like fs.readFile", async () => {
|
||||
const path = join(tempDir, "growing.txt");
|
||||
// Larger than one 64 KB stream chunk, so the append lands before the stream reaches the end.
|
||||
writeFileSync(path, `head\n${"x".repeat(200_000)}`);
|
||||
decodeHook.run = () => {
|
||||
decodeHook.run = undefined;
|
||||
appendFileSync(path, "\nadded");
|
||||
};
|
||||
const tool = createReadToolDefinition(tempDir);
|
||||
const result = await tool.execute("call", { path, limit: 1 }, undefined, undefined, undefined as never);
|
||||
expect(result.content[0]).toEqual({
|
||||
type: "text",
|
||||
text: "head\n\n[1 more lines in file. Use offset=2 to continue.]",
|
||||
});
|
||||
});
|
||||
|
||||
it("reads a 200 MB line with a 64 MB heap", () => {
|
||||
// A child process, so a regression that stores the whole line runs out of memory instead of passing.
|
||||
const sourceResolverUrl = pathToFileURL(resolve(__dirname, "../src/experimental/source-resolver.ts")).href;
|
||||
const readToolUrl = pathToFileURL(resolve(__dirname, "../src/core/tools/read.ts")).href;
|
||||
const script = `
|
||||
import { createReadToolDefinition } from ${JSON.stringify(readToolUrl)};
|
||||
// 64 KB chunks, as fs streams them: larger decoded strings live outside the V8 heap and its limit.
|
||||
const chunk = Buffer.alloc(64 * 1024, "x");
|
||||
async function* createReadStream() { for (let i = 0; i < 3200; i++) yield chunk; }
|
||||
const operations = { readFile: async () => Buffer.alloc(0), access: async () => {}, createReadStream };
|
||||
const tool = createReadToolDefinition(process.cwd(), { operations });
|
||||
const result = await tool.execute("call", { path: "big.log" });
|
||||
process.stdout.write(result.content[0].text);
|
||||
`;
|
||||
const output = execFileSync(
|
||||
process.execPath,
|
||||
["--max-old-space-size=64", "--import", sourceResolverUrl, "--input-type=module", "--eval", script],
|
||||
{ encoding: "utf-8", stdio: ["ignore", "pipe", "pipe"] },
|
||||
);
|
||||
expect(output).toBe(
|
||||
`[Line 1 is 200.0MB, exceeds 50.0KB limit. Use bash: sed -n '1p' big.log | head -c ${DEFAULT_MAX_BYTES}]`,
|
||||
);
|
||||
});
|
||||
|
||||
it("stops consuming the stream when aborted", async () => {
|
||||
const controller = new AbortController();
|
||||
let chunksRead = 0;
|
||||
let closed = false;
|
||||
const ops: ReadOperations = {
|
||||
readFile: async () => Buffer.alloc(0),
|
||||
access: async () => {},
|
||||
async *createReadStream(_path, { signal }) {
|
||||
expect(signal).toBe(controller.signal);
|
||||
try {
|
||||
while (true) {
|
||||
chunksRead++;
|
||||
if (chunksRead === 3) controller.abort();
|
||||
await new Promise((resolve) => setTimeout(resolve, 1));
|
||||
yield Buffer.from("line\n");
|
||||
}
|
||||
} finally {
|
||||
closed = true;
|
||||
}
|
||||
},
|
||||
};
|
||||
const tool = createReadToolDefinition(process.cwd(), { operations: ops });
|
||||
await expect(
|
||||
tool.execute("call", { path: FILE_PATH }, controller.signal, undefined, undefined as never),
|
||||
).rejects.toThrow("Operation aborted");
|
||||
await new Promise((resolve) => setTimeout(resolve, 20));
|
||||
expect(closed).toBe(true);
|
||||
expect(chunksRead).toBe(3);
|
||||
});
|
||||
});
|
||||
Loading…
Add table
Add a link
Reference in a new issue