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
read decoded the whole file and split it into lines before applying offset/limit, so large files ran out of memory or hit V8's string and fs size limits. It now streams through the new optional ReadOperations.createReadStream, stores only the lines truncateHead can show and counts the rest; output is unchanged. Abort stops the stream. Custom operations without createReadStream, and fractional or negative ranges, keep the whole-file path. fixes #9654
This commit is contained in:
Petr Houška
parent
5ff8c02708
commit
421da5e070
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