@khanhicetea/pi-better-tool 0.2.1 → 0.2.2

package/src/read-evidence.ts

@@ -11,29 +11,44 @@ export interface ReadEvidence {
 	/** 0-based, end-exclusive offsets in LF-normalized, BOM-stripped content. */
 	startOffset: number;
 	endOffset: number;
-	/** 1-based inclusive range actually shown to the model. */
+	/** 1-based inclusive range represented by the verified read output. */
 	startLine: number;
 	endLine: number;
 }
 
 interface ReadCall {
+	id: string;
 	path: string;
 	offset?: number;
 	limit?: number;
 }
 
 interface StoredToolResult {
+	role: "toolResult";
 	toolCallId: string;
 	toolName: string;
 	isError: boolean;
 	content: Array<{ type: string; text?: string }>;
 }
 
+interface StoredAssistant {
+	role: "assistant";
+	content: Array<{ type: string; id?: string; name?: string; arguments?: unknown }>;
+}
+
+interface ContextEntryLike {
+	type: string;
+	message?: StoredAssistant | StoredToolResult | { role: string };
+	retainedTail?: Array<StoredAssistant | StoredToolResult | { role: string }>;
+}
+
 /**
- * Find the newest successful read of targetPath still present in the active
- * model context. The returned range is accepted only when the stored tool
- * output exactly matches what the built-in read tool would produce from the
- * current file bytes. This makes stale or custom read output fail closed.
+ * Find the newest read call for targetPath in Pi's compaction-aware stored
+ * context. The call is accepted only when its successful result exactly
+ * matches the built-in read formatting for the current LF-normalized content.
+ *
+ * This establishes stored-context evidence, not proof of the final provider
+ * payload after other extensions' context/provider hooks.
  */
 export async function findLatestReadEvidence(
 	sessionManager: Pick<ExtensionContext["sessionManager"], "buildContextEntries"> | undefined,
@@ -42,49 +57,58 @@ export async function findLatestReadEvidence(
 	resolvePath: (path: string) => Promise<string>,
 ): Promise<ReadEvidence | null> {
 	if (!sessionManager) return null;
-	const entries = sessionManager.buildContextEntries();
-	const calls = new Map<string, ReadCall>();
-
-	for (const entry of entries) {
-		if (entry.type !== "message" || entry.message.role !== "assistant") continue;
-		for (const item of entry.message.content) {
-			if (item.type !== "toolCall" || item.name !== "read") continue;
-			const args = item.arguments as Record<string, unknown>;
-			if (typeof args.path !== "string") continue;
-			calls.set(item.id, {
-				path: args.path,
-				offset: typeof args.offset === "number" ? args.offset : undefined,
-				limit: typeof args.limit === "number" ? args.limit : undefined,
-			});
+	const entries = sessionManager.buildContextEntries() as ContextEntryLike[];
+	const messages = entries.flatMap((entry) => {
+		if (entry.type === "message" && entry.message) return [entry.message];
+		if (entry.type === "compaction" && Array.isArray(entry.retainedTail)) return entry.retainedTail;
+		return [];
+	});
+	const calls: ReadCall[] = [];
+	const results = new Map<string, StoredToolResult>();
+
+	for (const message of messages) {
+		if (message.role === "assistant") {
+			for (const item of (message as StoredAssistant).content) {
+				if (item.type !== "toolCall" || item.name !== "read" || typeof item.id !== "string") continue;
+				const args = item.arguments as Record<string, unknown> | undefined;
+				if (!args || typeof args.path !== "string") continue;
+				calls.push({
+					id: item.id,
+					path: args.path,
+					offset: typeof args.offset === "number" ? args.offset : undefined,
+					limit: typeof args.limit === "number" ? args.limit : undefined,
+				});
+			}
+		} else if (message.role === "toolResult") {
+			const result = message as StoredToolResult;
+			if (result.toolName === "read") results.set(result.toolCallId, result);
 		}
 	}
 
-	for (let i = entries.length - 1; i >= 0; i--) {
-		const entry = entries[i];
-		if (entry.type !== "message" || entry.message.role !== "toolResult") continue;
-		const result = entry.message as StoredToolResult;
-		if (result.toolName !== "read" || result.isError) continue;
-		const call = calls.get(result.toolCallId);
-		if (!call) continue;
-
+	for (let index = calls.length - 1; index >= 0; index--) {
+		const call = calls[index];
 		let readPath: string;
 		try {
 			readPath = await resolvePath(call.path);
 		} catch {
-			continue;
+			// Without a canonical identity we cannot prove that this newer read is
+			// unrelated, so fail closed instead of falling back to older intent.
+			return null;
 		}
 		if (readPath !== targetPath) continue;
 
-		// This is the latest successful read of this file. If it cannot be
-		// verified, do not silently fall back to older, potentially stale intent.
+		// Never fall back to older intent when the newest same-file read is
+		// missing, failed, malformed, or stale.
+		const result = results.get(call.id);
+		if (!result || result.isError) return null;
 		return evidenceFromBuiltinRead(normalizedContent, call, result.content);
 	}
 	return null;
 }
 
-function evidenceFromBuiltinRead(
+export function evidenceFromBuiltinRead(
 	content: string,
-	call: ReadCall,
+	call: Pick<ReadCall, "path" | "offset" | "limit">,
 	blocks: Array<{ type: string; text?: string }>,
 ): ReadEvidence | null {
 	if (blocks.length !== 1 || blocks[0].type !== "text" || typeof blocks[0].text !== "string") return null;

package/src/similarity.ts

@@ -3,16 +3,20 @@
  * failed oldText almost matched.
  */
 
-import { toFuzzyLines } from "./text.ts";
+import { normalizeForFuzzyMatch, normalizeToLF } from "./text.ts";
 
 export interface AlignOp {
-	type: "equal" | "changed" | "file-only" | "old-only";
-	/** 1-based line number in the file (for equal/changed/file-only). */
+	type: "equal" | "similar" | "changed" | "file-only" | "old-only";
+	/** 1-based line number in the file (for equal/similar/changed/file-only). */
 	fileLine?: number;
+	/** Original LF-normalized file text, never fuzzy-normalized. */
 	fileText?: string;
 	/** 1-based line number in the model-provided oldText. */
 	oldLine?: number;
+	/** Original LF-normalized oldText, never fuzzy-normalized. */
 	oldText?: string;
+	/** Heuristic similarity for aligned pairs. */
+	similarity?: number;
 }
 
 export interface ClosestRegion {
@@ -21,10 +25,15 @@ export interface ClosestRegion {
 	/** 0..1 order-preserving, one-to-one line similarity. */
 	score: number;
 	ops: AlignOp[];
+	/** Number of exactly equal original lines. */
 	equalCount: number;
 	totalOldLines: number;
 	/** oldText was truncated for comparison (very large oldText). */
 	truncated: boolean;
+	/** Best distinct, non-overlapping competitor, when retained by the bounded search. */
+	competitor?: { startLine: number; endLine: number; score: number };
+	/** A discarded window was too competitive to prove a safe score margin. */
+	competitionIncomplete: boolean;
 }
 
 interface PreparedLine {
@@ -42,15 +51,28 @@ function prepareLine(text: string): PreparedLine {
 	return { text, grams, gramCount: Math.max(0, text.length - 1) };
 }
 
+class WorkBudget {
+	private remaining: number;
+	constructor(remaining: number) {
+		this.remaining = remaining;
+	}
+	spend(amount: number): boolean {
+		this.remaining -= amount;
+		return this.remaining >= 0;
+	}
+}
+
 /** Dice coefficient over character bigrams; 1.0 for identical strings. */
 export function lineSimilarity(a: string, b: string): number {
 	return preparedSimilarity(prepareLine(a), prepareLine(b));
 }
 
-function preparedSimilarity(a: PreparedLine, b: PreparedLine): number {
+function preparedSimilarity(a: PreparedLine, b: PreparedLine, budget?: WorkBudget): number {
+	if (budget && !budget.spend(1)) throw new Error("similarity-work-budget-exhausted");
 	if (a.text === b.text) return 1;
 	if (a.gramCount === 0 || b.gramCount === 0) return 0;
 	const [small, large] = a.grams.size <= b.grams.size ? [a.grams, b.grams] : [b.grams, a.grams];
+	if (budget && !budget.spend(small.size)) throw new Error("similarity-work-budget-exhausted");
 	let overlap = 0;
 	for (const [gram, count] of small) {
 		const other = large.get(gram);
@@ -61,7 +83,10 @@ function preparedSimilarity(a: PreparedLine, b: PreparedLine): number {
 
 const MAX_COMPARE_LINES = 300;
 const MAX_SEARCH_FILE_LINES = 10_000;
-const MAX_CANDIDATE_WINDOWS = 24;
+const MAX_CANDIDATE_WINDOWS = 48;
+const MAX_QUERY_BYTES = 128 * 1024;
+const MAX_LINE_CHARS = 16 * 1024;
+const MAX_SIMILARITY_WORK = 20_000_000;
 const MIN_SCORE = 0.3;
 const MATCH_THRESHOLD = 0.75;
 
@@ -71,81 +96,122 @@ interface CandidateWindow {
 	positionalScore: number;
 }
 
+interface ScoredWindow extends CandidateWindow {
+	score: number;
+}
+
+function originalLines(text: string): string[] {
+	const lines = normalizeToLF(text).split("\n");
+	// Remove only the synthetic segment after a terminal newline. Do not remove
+	// a real, unterminated whitespace-only line.
+	if (lines.length > 0 && lines[lines.length - 1] === "") lines.pop();
+	return lines;
+}
+
 /**
- * Find the window of file lines most similar to oldText.
- *
- * The first pass ranks every window using cheap positional similarity. Only a
- * small bounded set of candidates receives the more expensive sequence score.
- * The sequence score is order-preserving and one-to-one, so repeated query
- * lines cannot all claim the same file line.
+ * Find the best window and a distinct runner-up under explicit input/work
+ * budgets. Returning null is the safe fallback when a budget is exhausted.
  */
 export function findClosestRegion(content: string, oldText: string): ClosestRegion | null {
-	const allQLines = toFuzzyLines(oldText);
-	const cLines = toFuzzyLines(content);
-	if (allQLines.length === 0 || cLines.length === 0 || cLines.length > MAX_SEARCH_FILE_LINES) return null;
+	if (Buffer.byteLength(oldText, "utf8") > MAX_QUERY_BYTES) return null;
+	const originalQ = originalLines(oldText);
+	const originalContent = originalLines(content);
+	if (originalQ.length === 0 || originalContent.length === 0 || originalContent.length > MAX_SEARCH_FILE_LINES) return null;
+	if ([...originalQ, ...originalContent].some((line) => line.length > MAX_LINE_CHARS)) return null;
 
+	const allQLines = originalQ.map(normalizeForFuzzyMatch);
+	const cLines = originalContent.map(normalizeForFuzzyMatch);
 	const truncated = allQLines.length > MAX_COMPARE_LINES;
 	const q = truncated ? allQLines.slice(0, MAX_COMPARE_LINES) : allQLines;
+	const qOriginal = truncated ? originalQ.slice(0, MAX_COMPARE_LINES) : originalQ;
 	const L = q.length;
 	const sizes = [...new Set([L - 1, L, L + 1].filter((size) => size >= 1 && size <= cLines.length))];
 	const preparedQ = q.map(prepareLine);
 	const preparedContent = cLines.map(prepareLine);
 	const candidates: CandidateWindow[] = [];
+	let discardedMaxPositionalScore = -1;
+	const budget = new WorkBudget(MAX_SIMILARITY_WORK);
 
-	for (const size of sizes) {
-		for (let start = 0; start + size <= cLines.length; start++) {
-			const pairs = Math.min(L, size);
-			let sum = 0;
-			for (let i = 0; i < pairs; i++) {
-				sum += preparedSimilarity(preparedQ[i], preparedContent[start + i]);
+	try {
+		for (const size of sizes) {
+			for (let start = 0; start + size <= cLines.length; start++) {
+				const pairs = Math.min(L, size);
+				let sum = 0;
+				for (let i = 0; i < pairs; i++) {
+					sum += preparedSimilarity(preparedQ[i], preparedContent[start + i], budget);
+				}
+				const discarded = keepCandidate(candidates, { start, size, positionalScore: sum / Math.max(L, size) });
+				if (discarded !== undefined) discardedMaxPositionalScore = Math.max(discardedMaxPositionalScore, discarded);
 			}
-			keepCandidate(candidates, { start, size, positionalScore: sum / Math.max(L, size) });
 		}
-	}
 
-	let bestScore = -1;
-	let best: CandidateWindow | undefined;
-	for (const candidate of candidates) {
-		const window = preparedContent.slice(candidate.start, candidate.start + candidate.size);
-		const score = sequenceSimilarity(preparedQ, window);
-		if (score > bestScore || (score === bestScore && candidate.positionalScore > (best?.positionalScore ?? -1))) {
-			bestScore = score;
-			best = candidate;
-		}
+		const scored: ScoredWindow[] = candidates.map((candidate) => ({
+			...candidate,
+			score: sequenceSimilarity(
+				preparedQ,
+				preparedContent.slice(candidate.start, candidate.start + candidate.size),
+				budget,
+			),
+		}));
+		scored.sort(
+			(a, b) => b.score - a.score || b.positionalScore - a.positionalScore || a.start - b.start || a.size - b.size,
+		);
+		const best = scored[0];
+		if (!best || best.score < MIN_SCORE) return null;
+		const bestEnd = best.start + best.size;
+		const runner = scored.find((candidate) => {
+			const candidateEnd = candidate.start + candidate.size;
+			return candidateEnd <= best.start || candidate.start >= bestEnd;
+		});
+
+		const windowLines = cLines.slice(best.start, bestEnd);
+		const windowOriginal = originalContent.slice(best.start, bestEnd);
+		const ops = alignLines(q, windowLines, qOriginal, windowOriginal, best.start, budget);
+		return {
+			startLine: best.start + 1,
+			endLine: bestEnd,
+			score: best.score,
+			ops,
+			equalCount: ops.filter((op) => op.type === "equal").length,
+			totalOldLines: L,
+			truncated,
+			competitor: runner
+				? { startLine: runner.start + 1, endLine: runner.start + runner.size, score: runner.score }
+				: undefined,
+			competitionIncomplete: discardedMaxPositionalScore >= best.positionalScore - MIN_DIRECT_POSITIONAL_GAP,
+		};
+	} catch (error) {
+		if (error instanceof Error && error.message === "similarity-work-budget-exhausted") return null;
+		throw error;
 	}
-	if (!best || bestScore < MIN_SCORE) return null;
-
-	const windowLines = cLines.slice(best.start, best.start + best.size);
-	const ops = alignLines(q, windowLines, best.start);
-	return {
-		startLine: best.start + 1,
-		endLine: best.start + best.size,
-		score: bestScore,
-		ops,
-		equalCount: ops.filter((op) => op.type === "equal").length,
-		totalOldLines: L,
-		truncated,
-	};
 }
 
-function keepCandidate(candidates: CandidateWindow[], candidate: CandidateWindow): void {
+const MIN_DIRECT_POSITIONAL_GAP = 0.1;
+
+/** Return the positional score of any candidate discarded by the bounded heap. */
+function keepCandidate(candidates: CandidateWindow[], candidate: CandidateWindow): number | undefined {
 	if (candidates.length < MAX_CANDIDATE_WINDOWS) {
 		candidates.push(candidate);
-		candidates.sort((a, b) => a.positionalScore - b.positionalScore);
-		return;
+		candidates.sort((a, b) => a.positionalScore - b.positionalScore || b.start - a.start || b.size - a.size);
+		return undefined;
 	}
-	if (candidate.positionalScore <= candidates[0].positionalScore) return;
+	const worst = candidates[0];
+	if (
+		candidate.positionalScore < worst.positionalScore ||
+		(candidate.positionalScore === worst.positionalScore && candidate.start >= worst.start)
+	) return candidate.positionalScore;
 	candidates[0] = candidate;
-	candidates.sort((a, b) => a.positionalScore - b.positionalScore);
+	candidates.sort((a, b) => a.positionalScore - b.positionalScore || b.start - a.start || b.size - a.size);
+	return worst.positionalScore;
 }
 
 /** Weighted LCS: lines are matched at most once and in source order. */
-function sequenceSimilarity(query: PreparedLine[], window: PreparedLine[]): number {
+function sequenceSimilarity(query: PreparedLine[], window: PreparedLine[], budget: WorkBudget): number {
 	let previous = new Float64Array(window.length + 1);
 	for (let i = 1; i <= query.length; i++) {
 		const current = new Float64Array(window.length + 1);
 		for (let j = 1; j <= window.length; j++) {
-			const matched = previous[j - 1] + preparedSimilarity(query[i - 1], window[j - 1]);
+			const matched = previous[j - 1] + preparedSimilarity(query[i - 1], window[j - 1], budget);
 			current[j] = Math.max(previous[j], current[j - 1], matched);
 		}
 		previous = current;
@@ -153,18 +219,21 @@ function sequenceSimilarity(query: PreparedLine[], window: PreparedLine[]): numb
 	return previous[window.length] / Math.max(query.length, window.length);
 }
 
-/**
- * LCS-style alignment between oldText lines and the best window. Lines with
- * similarity >= MATCH_THRESHOLD count as equal; adjacent insert/delete pairs
- * are merged into changed operations for compact diagnostics.
- */
-function alignLines(q: string[], w: string[], wOffset: number): AlignOp[] {
+/** Align fuzzy lines while attaching original strings to every operation. */
+function alignLines(
+	q: string[],
+	w: string[],
+	qOriginal: string[],
+	wOriginal: string[],
+	wOffset: number,
+	budget: WorkBudget,
+): AlignOp[] {
 	const n = q.length;
 	const m = w.length;
 	const preparedQ = q.map(prepareLine);
 	const preparedW = w.map(prepareLine);
 	const similarities = Array.from({ length: n }, (_, i) =>
-		Array.from({ length: m }, (_, j) => preparedSimilarity(preparedQ[i], preparedW[j])),
+		Array.from({ length: m }, (_, j) => preparedSimilarity(preparedQ[i], preparedW[j], budget)),
 	);
 	const dp: number[][] = Array.from({ length: n + 1 }, () => new Array<number>(m + 1).fill(0));
 	for (let i = n - 1; i >= 0; i--) {
@@ -180,17 +249,25 @@ function alignLines(q: string[], w: string[], wOffset: number): AlignOp[] {
 	let j = 0;
 	while (i < n && j < m) {
 		if (similarities[i][j] >= MATCH_THRESHOLD) {
-			raw.push({ type: "equal", fileLine: wOffset + j + 1, fileText: w[j], oldLine: i + 1, oldText: q[i] });
+			const exact = qOriginal[i] === wOriginal[j];
+			raw.push({
+				type: exact ? "equal" : "similar",
+				fileLine: wOffset + j + 1,
+				fileText: wOriginal[j],
+				oldLine: i + 1,
+				oldText: qOriginal[i],
+				similarity: similarities[i][j],
+			});
 			i++;
 			j++;
 		} else if (dp[i + 1][j] > dp[i][j + 1]) {
-			raw.push({ type: "old-only", oldLine: i + 1, oldText: q[i++] });
+			raw.push({ type: "old-only", oldLine: i + 1, oldText: qOriginal[i++] });
 		} else {
-			raw.push({ type: "file-only", fileLine: wOffset + j + 1, fileText: w[j++] });
+			raw.push({ type: "file-only", fileLine: wOffset + j + 1, fileText: wOriginal[j++] });
 		}
 	}
-	while (i < n) raw.push({ type: "old-only", oldLine: i + 1, oldText: q[i++] });
-	while (j < m) raw.push({ type: "file-only", fileLine: wOffset + j + 1, fileText: w[j++] });
+	while (i < n) raw.push({ type: "old-only", oldLine: i + 1, oldText: qOriginal[i++] });
+	while (j < m) raw.push({ type: "file-only", fileLine: wOffset + j + 1, fileText: wOriginal[j++] });
 
 	const merged: AlignOp[] = [];
 	for (let k = 0; k < raw.length; k++) {

package/src/text.ts

@@ -71,6 +71,17 @@ export function splitLinesWithEndings(content: string): string[] {
 	return content.match(/[^\n]*\n|[^\n]+/g) ?? [];
 }
 
+/**
+ * Split every logical line while retaining endings. Unlike
+ * splitLinesWithEndings(), this includes the zero-length logical line after a
+ * terminal newline. That distinction is needed when fuzzy normalization turns
+ * a whitespace-only final line into an empty line.
+ */
+export function splitLogicalLinesWithEndings(content: string): string[] {
+	const parts = content.split("\n");
+	return parts.map((part, index) => (index < parts.length - 1 ? `${part}\n` : part));
+}
+
 export function getLineSpans(content: string): LineSpan[] {
 	let offset = 0;
 	return splitLinesWithEndings(content).map((line) => {
@@ -80,6 +91,16 @@ export function getLineSpans(content: string): LineSpan[] {
 	});
 }
 
+/** Logical spans, including a possible zero-length final line. */
+export function getLogicalLineSpans(content: string): LineSpan[] {
+	let offset = 0;
+	return splitLogicalLinesWithEndings(content).map((line) => {
+		const span = { start: offset, end: offset + line.length };
+		offset = span.end;
+		return span;
+	});
+}
+
 /** 0-based index of the line containing `offset` (clamped to the last line). */
 export function lineAt(spans: LineSpan[], offset: number): number {
 	let lo = 0;
@@ -102,17 +123,25 @@ export function lineAt(spans: LineSpan[], offset: number): number {
  * same normalization the matching engine uses for its uniqueness check.
  * `fuzzyContent` must already be `normalizeForFuzzyMatch`-ed.
  */
-export function countFuzzyOccurrences(fuzzyContent: string, needle: string): number {
-	if (!needle) return 0;
-	return fuzzyContent.split(normalizeForFuzzyMatch(needle)).length - 1;
+export function countFuzzyOccurrences(fuzzyContent: string, needle: string, stopAfter = Number.POSITIVE_INFINITY): number {
+	const normalizedNeedle = normalizeForFuzzyMatch(needle);
+	if (!normalizedNeedle) return 0;
+	let count = 0;
+	let index = fuzzyContent.indexOf(normalizedNeedle);
+	while (index !== -1) {
+		count++;
+		if (count >= stopAfter) return count;
+		index = fuzzyContent.indexOf(normalizedNeedle, index + normalizedNeedle.length);
+	}
+	return count;
 }
 
-/** All start offsets of `needle` in `haystack` (literal string search). */
-export function findAllOccurrences(haystack: string, needle: string): number[] {
+/** All non-overlapping start offsets of `needle`, optionally bounded. */
+export function findAllOccurrences(haystack: string, needle: string, limit = Number.POSITIVE_INFINITY): number[] {
 	const offsets: number[] = [];
-	if (!needle) return offsets;
+	if (!needle || limit <= 0) return offsets;
 	let idx = haystack.indexOf(needle);
-	while (idx !== -1) {
+	while (idx !== -1 && offsets.length < limit) {
 		offsets.push(idx);
 		idx = haystack.indexOf(needle, idx + needle.length);
 	}