fix(cli): restore read tool streaming

This commit is contained in:
marius-kilocode
2026-06-22 14:26:32 +02:00
parent d378114b8b
commit 15f42d4bec
7 changed files with 299 additions and 73 deletions
+6
View File
@@ -0,0 +1,6 @@
---
"@kilocode/cli": patch
"kilo-code": patch
---
Restore bounded text-file reads and keep zero-limit pagination and Unicode truncation from producing unusable tool output.
+48 -40
View File
@@ -1,5 +1,6 @@
import { createReadStream } from "fs"
import { PassThrough, Readable } from "stream"
import type { AppFileSystem } from "@opencode-ai/core/filesystem"
import { Effect, Stream } from "effect"
import { addAbortSignal, Readable } from "stream"
import * as Encoding from "./encoding"
/**
@@ -17,54 +18,61 @@ export class InvalidUtf8Error extends Error {
}
}
/**
* UTF-8 text Readable for `filepath`. A leading UTF-8 BOM passes through as
* U+FEFF — same as `createReadStream({ encoding: "utf8" })`.
*/
export function openUtf8(filepath: string): Readable {
const out = new PassThrough({ encoding: "utf8" })
const raw = createReadStream(filepath)
const decoder = new TextDecoder("utf-8", { fatal: true })
raw.on("data", (chunk) => {
try {
const text = decoder.decode(chunk as Buffer, { stream: true })
if (text) out.write(text)
} catch {
raw.destroy()
out.destroy(new InvalidUtf8Error())
}
})
raw.on("end", () => {
try {
const tail = decoder.decode()
if (tail) out.write(tail)
out.end()
} catch {
out.destroy(new InvalidUtf8Error())
}
})
raw.on("error", (err) => out.destroy(err))
// Propagate consumer-side teardown so early-exit (line / byte cap, fallback)
// stops pulling chunks from disk instead of running to EOF.
out.on("close", () => raw.destroy())
return out
type FileSystem = Pick<AppFileSystem.Interface, "readFile" | "stream">
function decode(decoder: TextDecoder, bytes?: Uint8Array) {
try {
return decoder.decode(bytes, bytes ? { stream: true } : undefined)
} catch {
throw new InvalidUtf8Error()
}
}
/** Whole-file UTF-8 Readable via {@link Encoding.read}; buffers the entire decoded file. */
export async function openDecoded(filepath: string): Promise<Readable> {
const decoded = await Encoding.read(filepath)
return Readable.from([decoded.text])
async function* chunks(fs: FileSystem, filepath: string) {
const decoder = new TextDecoder("utf-8", { fatal: true })
for await (const bytes of Stream.toAsyncIterable(fs.stream(filepath))) {
const text = decode(decoder, bytes)
if (text) yield text
}
const tail = decode(decoder)
if (tail) yield tail
}
export function abortable(stream: Readable, signal?: AbortSignal) {
return signal ? addAbortSignal(signal, stream) : stream
}
/** UTF-8 text stream backed by the injected filesystem service. */
export function openUtf8(fs: FileSystem, filepath: string, signal?: AbortSignal): Readable {
return abortable(Readable.from(chunks(fs, filepath)), signal)
}
export function safeSlice(text: string, end: number) {
const sliced = text.slice(0, end)
const last = sliced.charCodeAt(sliced.length - 1)
return last >= 0xd800 && last <= 0xdbff ? sliced.slice(0, -1) : sliced
}
/** Whole-file decoded Readable; buffers legacy encodings only after UTF-8 streaming fails. */
export async function openDecoded(fs: FileSystem, filepath: string, signal?: AbortSignal): Promise<Readable> {
const bytes = Buffer.from(await Effect.runPromise(fs.readFile(filepath), { signal }))
return abortable(Readable.from([Encoding.decode(bytes, Encoding.detect(bytes))]), signal)
}
/**
* Run `fn` against an optimistic UTF-8 stream; on {@link InvalidUtf8Error}
* retry once against {@link openDecoded}. Other errors propagate.
*/
export async function withFallback<T>(filepath: string, fn: (input: Readable) => Promise<T>): Promise<T> {
export async function withFallback<T>(
fs: FileSystem,
filepath: string,
fn: (input: Readable) => Promise<T>,
signal?: AbortSignal,
): Promise<T> {
try {
return await fn(openUtf8(filepath))
return await fn(openUtf8(fs, filepath, signal))
} catch (err) {
if (!(err instanceof InvalidUtf8Error)) throw err
}
return fn(await openDecoded(filepath))
return fn(await openDecoded(fs, filepath, signal))
}
+6 -2
View File
@@ -24,6 +24,7 @@ import { ModelID, ProviderID } from "@/provider/schema"
import { SessionNetwork } from "./network" // kilocode_change
import { CodexAuthExpiredError } from "@/kilocode/provider/codex-refresh" // kilocode_change
import { KiloSessionMessageOrder } from "@/kilocode/session/message-order" // kilocode_change
import * as TextStream from "@/kilocode/text-stream" // kilocode_change
import { Effect, Schema, Types } from "effect"
import { NonNegativeInt } from "@opencode-ai/core/schema"
import * as EffectLogger from "@opencode-ai/core/effect/logger"
@@ -283,8 +284,11 @@ export type ToolStateCompleted = Types.DeepMutable<Schema.Schema.Type<typeof Too
function truncateToolOutput(text: string, maxChars?: number) {
if (!maxChars || text.length <= maxChars) return text
const omitted = text.length - maxChars
return `${text.slice(0, maxChars)}\n[Tool output truncated for compaction: omitted ${omitted} chars]`
// kilocode_change start - avoid persisting malformed Unicode in compacted tool output
const sliced = TextStream.safeSlice(text, maxChars)
const omitted = text.length - sliced.length
return `${sliced}\n[Tool output truncated for compaction: omitted ${omitted} chars]`
// kilocode_change end
}
export const ToolStateError = Schema.Struct({
+45 -28
View File
@@ -15,11 +15,14 @@ import { Reference } from "@/reference/reference"
// kilocode_change start
import * as Encoding from "../kilocode/encoding"
import * as Extract from "../kilocode/tool/read-extract"
import * as TextStream from "../kilocode/text-stream"
// kilocode_change end
const DEFAULT_READ_LIMIT = 2000
const MAX_LINE_LENGTH = 2000
const MAX_LINE_SUFFIX = `... (line truncated to ${MAX_LINE_LENGTH} chars)`
// kilocode_change start - report the safe Unicode slice length
const suffix = (length: number) => `... (line truncated to ${length} chars)`
// kilocode_change end
const MAX_BYTES = 50 * 1024
const MAX_BYTES_LABEL = `${MAX_BYTES / 1024} KB`
const SAMPLE_BYTES = 4096
@@ -110,30 +113,33 @@ export const ReadTool = Tool.define(
)
})
const lines = Effect.fn("ReadTool.lines")((filepath: string, opts: { limit: number; offset: number }) =>
// kilocode_change - extracted formats still need their native readers; ordinary text stays on AppFileSystem
Effect.tryPromise({
try: () => Extract.open(filepath),
catch: (err) => (err instanceof Error ? err : new Error(String(err))),
}).pipe(
Effect.flatMap((extracted) =>
extracted
? Effect.tryPromise({
try: () => collect(extracted, opts),
catch: (err) => (err instanceof Error ? err : new Error(String(err))),
})
: fs.readFile(filepath).pipe(
Effect.map((bytes) => Encoding.decode(Buffer.from(bytes), Encoding.detect(Buffer.from(bytes)))),
Effect.flatMap((text) =>
Effect.tryPromise({
try: () => collect(Readable.from([text]), opts),
catch: (err) => (err instanceof Error ? err : new Error(String(err))),
}),
),
),
// kilocode_change start - extracted formats use native readers; ordinary text streams through AppFileSystem
const lines = Effect.fn("ReadTool.lines")(
(filepath: string, opts: { limit: number; offset: number }, abort: AbortSignal) =>
Effect.tryPromise({
try: () => Extract.open(filepath),
catch: (err) => (err instanceof Error ? err : new Error(String(err))),
}).pipe(
Effect.flatMap((extracted) =>
extracted
? Effect.tryPromise({
try: (signal) => collect(TextStream.abortable(extracted, AbortSignal.any([abort, signal])), opts),
catch: (err) => (err instanceof Error ? err : new Error(String(err))),
})
: Effect.tryPromise({
try: (signal) =>
TextStream.withFallback(
fs,
filepath,
(stream) => collect(stream, opts),
AbortSignal.any([abort, signal]),
),
catch: (err) => (err instanceof Error ? err : new Error(String(err))),
}),
),
),
),
)
// kilocode_change end
const isBinaryFile = (filepath: string, bytes: Uint8Array) => {
const ext = path.extname(filepath).toLowerCase()
@@ -196,6 +202,7 @@ export const ReadTool = Tool.define(
filepath: string,
items: string[],
directory: string,
abort: AbortSignal,
) {
const entries = yield* fs.readDirectoryEntries(filepath).pipe(Effect.catch(() => Effect.succeed([])))
const types = new Map(entries.map((entry) => [entry.name, entry.type]))
@@ -209,7 +216,7 @@ export const ReadTool = Tool.define(
Effect.catch(() => Effect.succeed(new Uint8Array())),
)
if (isBinaryFile(child, sample)) return
const file = yield* lines(child, { limit: DEFAULT_READ_LIMIT, offset: 1 }).pipe(
const file = yield* lines(child, { limit: DEFAULT_READ_LIMIT, offset: 1 }, abort).pipe(
Effect.catch(() => Effect.void),
)
if (!file) return
@@ -264,14 +271,14 @@ export const ReadTool = Tool.define(
if (stat.type === "Directory") {
const items = yield* list(filepath)
const limit = params.limit ?? DEFAULT_READ_LIMIT
const limit = Math.max(1, params.limit ?? DEFAULT_READ_LIMIT) // kilocode_change - prevent zero-limit loops
const offset = params.offset || 1
const start = offset - 1
const sliced = items.slice(start, start + limit)
const truncated = start + sliced.length < items.length
// kilocode_change start
const expand = Boolean(ctx.extra?.["includeDirectoryFiles"])
const loaded = expand ? yield* readDirectoryFiles(filepath, sliced, instance.directory) : []
const loaded = expand ? yield* readDirectoryFiles(filepath, sliced, instance.directory, ctx.abort) : []
const content = loaded.map((item) => item.content).join("\n\n")
// kilocode_change end
@@ -332,7 +339,14 @@ export const ReadTool = Tool.define(
return yield* Effect.fail(new Error(`Cannot read binary file: ${filepath}`))
}
const file = yield* lines(filepath, { limit: params.limit ?? DEFAULT_READ_LIMIT, offset: params.offset || 1 })
const file = yield* lines(
filepath,
{
limit: Math.max(1, params.limit ?? DEFAULT_READ_LIMIT),
offset: params.offset || 1,
},
ctx.abort,
)
if (file.count < file.offset && !(file.count === 0 && file.offset === 1)) {
return yield* Effect.fail(
new Error(`Offset ${file.offset} is out of range for this file (${file.count} lines)`),
@@ -399,7 +413,10 @@ async function collect(stream: Readable, opts: { limit: number; offset: number }
more = true
continue
}
const line = text.length > MAX_LINE_LENGTH ? text.substring(0, MAX_LINE_LENGTH) + MAX_LINE_SUFFIX : text
// kilocode_change start - keep truncated output valid Unicode
const sliced = TextStream.safeSlice(text, MAX_LINE_LENGTH)
const line = text.length > MAX_LINE_LENGTH ? sliced + suffix(sliced.length) : text
// kilocode_change end
const size = Buffer.byteLength(line, "utf-8") + (raw.length > 0 ? 1 : 0)
if (bytes + size > MAX_BYTES) {
cut = true
@@ -95,6 +95,20 @@ describe("kilocode directory reads", () => {
}),
)
it.live("clamps a zero entry limit and advances pagination", () =>
Effect.gen(function* () {
const dir = yield* tmpdirScoped()
yield* put(path.join(dir, "folder", "a.txt"), "alpha")
yield* put(path.join(dir, "folder", "b.txt"), "beta")
const result = yield* exec(dir, { filePath: path.join(dir, "folder"), limit: 0 }, baseCtx)
expect(result.output).toContain("a.txt")
expect(result.output).not.toContain("b.txt")
expect(result.output).toContain("beyond entry 2")
}),
)
if (process.platform !== "win32") {
it.live("skips symlinked top-level files", () =>
Effect.gen(function* () {
@@ -8,8 +8,14 @@ import { KiloSessionMessageOrder } from "../../src/kilocode/session/message-orde
import { MessageV2 } from "../../src/session/message-v2"
import { ModelID, ProviderID } from "../../src/provider/schema"
import { MessageID, PartID, SessionID } from "../../src/session/schema"
import type { Provider } from "../../src/provider/provider"
const sessionID = SessionID.make("ses_safety")
const model = {
id: ModelID.make("test"),
providerID: ProviderID.make("test"),
api: { id: "test", npm: "@ai-sdk/openai" },
} as Provider.Model
function userInfo(id: string): MessageV2.User {
return {
@@ -581,3 +587,25 @@ describe("KiloSessionPrompt.maybeStripHistoricalMedia", () => {
expect(result[3].parts[0].type).toBe("text")
})
})
describe("MessageV2 tool output truncation", () => {
test("does not split a surrogate pair during compaction", async () => {
const part = toolPart("msg_a", "completed")
if (part.state.status !== "completed") throw new Error("expected completed tool part")
part.state.output = "x".repeat(1999) + "📁" + "tail"
const result = await MessageV2.toModelMessages(
[user("msg_u", [textPart("msg_u", "read")]), assistant("msg_a", "msg_u", [part])],
model,
{ toolOutputMaxChars: 2000 },
)
const message = result[2]
if (message.role !== "tool") throw new Error("expected tool message")
const item = message.content[0]
if (item.type !== "tool-result" || item.output.type !== "text") throw new Error("expected text tool result")
const isolated = /[\uD800-\uDBFF](?![\uDC00-\uDFFF])|(?<![\uD800-\uDBFF])[\uDC00-\uDFFF]/
expect(isolated.test(item.output.value)).toBe(false)
expect(item.output.value).toContain("omitted 6 chars")
})
})
@@ -4,7 +4,7 @@
// directly so we validate end-to-end behaviour.
import { afterEach, describe, expect } from "bun:test"
import { Effect, Layer } from "effect"
import { Effect, Exit, Layer, Stream } from "effect"
import path from "path"
import fs from "fs/promises"
import iconv from "iconv-lite"
@@ -54,11 +54,11 @@ const it = testEffect(
),
)
const runRead = (args: Tool.InferParameters<typeof ReadTool>) =>
const runRead = (args: Tool.InferParameters<typeof ReadTool>, next: Tool.Context = ctx) =>
Effect.gen(function* () {
const info = yield* ReadTool
const tool = yield* info.init()
return yield* tool.execute(args, ctx)
return yield* tool.execute(args, next)
})
const runWrite = (args: Tool.InferParameters<typeof WriteTool>) =>
@@ -189,6 +189,155 @@ describe("tool encoding preservation", () => {
)
})
describe("ReadTool streaming and pagination", () => {
it.live("streams UTF-8 files and stops after the output cap", () =>
provideTmpdirInstance((dir) =>
Effect.gen(function* () {
const filepath = path.join(dir, "large.txt")
const content = `${"x".repeat(80)}\n`.repeat(50_000)
yield* Effect.promise(() => fs.writeFile(filepath, content))
const base = yield* AppFileSystem.Service
const counter = { bytes: 0 }
const result = yield* runRead({ filePath: filepath }).pipe(
Effect.provideService(
AppFileSystem.Service,
AppFileSystem.Service.of({
...base,
stream: (file, options) =>
base.stream(file, options).pipe(
Stream.tap((chunk) =>
Effect.sync(() => {
counter.bytes += chunk.length
}),
),
),
}),
),
)
expect(result.metadata.truncated).toBe(true)
expect(counter.bytes).toBeGreaterThan(0)
expect(counter.bytes).toBeLessThan(Buffer.byteLength(content, "utf-8") / 2)
}),
),
)
it.live("stops the filesystem stream when the tool is aborted", () =>
provideTmpdirInstance((dir) =>
Effect.gen(function* () {
const filepath = path.join(dir, "abort.txt")
yield* Effect.promise(() => fs.writeFile(filepath, `${"x".repeat(80)}\n`.repeat(50_000)))
const base = yield* AppFileSystem.Service
const controller = new AbortController()
const state = { chunks: 0, closed: false }
const exit = yield* runRead({ filePath: filepath }, { ...ctx, abort: controller.signal }).pipe(
Effect.provideService(
AppFileSystem.Service,
AppFileSystem.Service.of({
...base,
stream: (file, options) =>
base.stream(file, options).pipe(
Stream.tap(() =>
Effect.sync(() => {
state.chunks += 1
controller.abort()
}),
),
Stream.ensuring(
Effect.sync(() => {
state.closed = true
}),
),
),
}),
),
Effect.exit,
)
expect(Exit.isFailure(exit)).toBe(true)
expect(state.chunks).toBeGreaterThan(0)
expect(state.closed).toBe(true)
}),
),
)
it.live("restarts cleanly when invalid UTF-8 appears after streamed lines", () =>
provideTmpdirInstance((dir) =>
Effect.gen(function* () {
const filepath = path.join(dir, "legacy.txt")
const lines = Array.from({ length: 1_000 }, (_, i) => `valid-${i + 1}-${"x".repeat(70)}`)
const content = Buffer.concat([
Buffer.from(lines.join("\n") + "\n"),
iconv.encode(samples.shiftJis, "Shift_JIS"),
Buffer.from("\nlast"),
])
yield* Effect.promise(() => fs.writeFile(filepath, content))
const base = yield* AppFileSystem.Service
const calls = { bytes: 0, reads: 0 }
const result = yield* runRead({ filePath: filepath, offset: 999, limit: 5 }).pipe(
Effect.provideService(
AppFileSystem.Service,
AppFileSystem.Service.of({
...base,
readFile: (file) =>
Effect.sync(() => {
calls.reads += 1
}).pipe(Effect.andThen(base.readFile(file))),
stream: (file, options) =>
base.stream(file, { ...options, chunkSize: 1024 }).pipe(
Stream.tap((chunk) =>
Effect.sync(() => {
calls.bytes += chunk.length
}),
),
),
}),
),
)
expect(calls.bytes).toBeGreaterThan(64 * 1024)
expect(calls.reads).toBe(1)
expect(result.output.match(/999: valid-999-/g)?.length).toBe(1)
expect(result.output).toContain(`1001: ${samples.shiftJis}`)
expect(result.output).toContain("1002: last")
}),
),
)
it.live("clamps a zero line limit and advances pagination", () =>
provideTmpdirInstance((dir) =>
Effect.gen(function* () {
const filepath = path.join(dir, "lines.txt")
yield* Effect.promise(() => fs.writeFile(filepath, "first\nsecond"))
const result = yield* runRead({ filePath: filepath, limit: 0 })
expect(result.output).toContain("1: first")
expect(result.output).not.toContain("2: second")
expect(result.output).toContain("Use offset=2")
}),
),
)
it.live("keeps truncated lines valid when an emoji crosses the boundary", () =>
provideTmpdirInstance((dir) =>
Effect.gen(function* () {
const filepath = path.join(dir, "emoji.txt")
yield* Effect.promise(() => fs.writeFile(filepath, "x".repeat(1999) + "📁" + "tail"))
const result = yield* runRead({ filePath: filepath })
const isolated = /[\uD800-\uDBFF](?![\uDC00-\uDFFF])|(?<![\uD800-\uDBFF])[\uDC00-\uDFFF]/
expect(isolated.test(result.output)).toBe(false)
expect(result.output).toContain("(line truncated to 1999 chars)")
}),
),
)
})
describe("WriteTool preserves existing file encoding when overwriting", () => {
const cases: Array<[string, string, string]> = [
["UTF-8 with BOM", UTF8_BOM, samples.utf8],