@oh-my-pi/pi-coding-agent 14.0.5 → 14.1.1

package/src/tools/read.ts

@@ -1,3 +1,4 @@
+import { Database } from "bun:sqlite";
 import * as fs from "node:fs/promises";
 import * as path from "node:path";
 import type { AgentTool, AgentToolContext, AgentToolResult, AgentToolUpdateCallback } from "@oh-my-pi/pi-agent-core";
@@ -38,7 +39,6 @@ import { resolveFileDisplayMode } from "../utils/file-display-mode";
 import { ImageInputTooLargeError, loadImageInput, MAX_IMAGE_INPUT_BYTES } from "../utils/image-loading";
 import { convertFileWithMarkit } from "../utils/markit";
 import { type ArchiveReader, openArchive, parseArchivePathCandidates } from "./archive-reader";
-
 import {
 	executeReadUrl,
 	isReadableUrlPath,
@@ -52,6 +52,22 @@ import { applyListLimit } from "./list-limit";
 import { formatFullOutputReference, formatStyledTruncationWarning, type OutputMeta } from "./output-meta";
 import { expandPath, resolveReadPath } from "./path-utils";
 import { formatAge, formatBytes, shortenPath, wrapBrackets } from "./render-utils";
+import {
+	executeReadQuery,
+	getRowByKey,
+	getRowByRowId,
+	getTableSchema,
+	isSqliteFile,
+	listTables,
+	parseSqlitePathCandidates,
+	parseSqliteSelector,
+	queryRows,
+	renderRow,
+	renderSchema,
+	renderTable,
+	renderTableList,
+	resolveTableRowLookup,
+} from "./sqlite-reader";
 import { ToolAbortError, ToolError, throwIfAborted } from "./tool-errors";
 import { toolResult } from "./tool-result";
 
@@ -420,6 +436,29 @@ interface ResolvedArchiveReadPath {
 	suffixResolution?: { from: string; to: string };
 }
 
+interface ResolvedSqliteReadPath {
+	absolutePath: string;
+	sqliteSubPath: string;
+	queryString: string;
+	suffixResolution?: { from: string; to: string };
+}
+
+function parseSqliteSelectorInput(selector: string | undefined): { subPath: string; queryString: string } {
+	if (!selector) {
+		return { subPath: "", queryString: "" };
+	}
+
+	const queryIndex = selector.indexOf("?");
+	if (queryIndex === -1) {
+		return { subPath: selector.replace(/^:+/, ""), queryString: "" };
+	}
+
+	return {
+		subPath: selector.slice(0, queryIndex).replace(/^:+/, ""),
+		queryString: selector.slice(queryIndex + 1),
+	};
+}
+
 /**
  * Read tool implementation.
  *
@@ -502,6 +541,53 @@ export class ReadTool implements AgentTool<typeof readSchema, ReadToolDetails> {
 		return null;
 	}
 
+	async #resolveSqliteReadPath(readPath: string, signal?: AbortSignal): Promise<ResolvedSqliteReadPath | null> {
+		const candidates = parseSqlitePathCandidates(readPath);
+		for (const candidate of candidates) {
+			let absolutePath = resolveReadPath(candidate.sqlitePath, this.session.cwd);
+			let suffixResolution: { from: string; to: string } | undefined;
+
+			try {
+				const stat = await Bun.file(absolutePath).stat();
+				if (stat.isDirectory()) continue;
+				if (!(await isSqliteFile(absolutePath))) continue;
+
+				return {
+					absolutePath,
+					sqliteSubPath: candidate.subPath,
+					queryString: candidate.queryString,
+					suffixResolution,
+				};
+			} catch (error) {
+				if (!isNotFoundError(error) || isRemoteMountPath(absolutePath)) continue;
+
+				const suffixMatch = await findUniqueSuffixMatch(candidate.sqlitePath, this.session.cwd, signal);
+				if (!suffixMatch) continue;
+
+				try {
+					const retryStat = await Bun.file(suffixMatch.absolutePath).stat();
+					if (retryStat.isDirectory()) continue;
+					if (!(await isSqliteFile(suffixMatch.absolutePath))) continue;
+
+					absolutePath = suffixMatch.absolutePath;
+					suffixResolution = { from: candidate.sqlitePath, to: suffixMatch.displayPath };
+					return {
+						absolutePath,
+						sqliteSubPath: candidate.subPath,
+						queryString: candidate.queryString,
+						suffixResolution,
+					};
+				} catch (retryError) {
+					if (!isNotFoundError(retryError)) {
+						throw retryError;
+					}
+				}
+			}
+		}
+
+		return null;
+	}
+
 	#buildInMemoryTextResult(
 		text: string,
 		offset: number | undefined,
@@ -709,6 +795,132 @@ export class ReadTool implements AgentTool<typeof readSchema, ReadToolDetails> {
 		return result;
 	}
 
+	async #readSqlite(
+		sel: string | undefined,
+		resolvedSqlitePath: ResolvedSqliteReadPath,
+		signal?: AbortSignal,
+	): Promise<AgentToolResult<ReadToolDetails>> {
+		throwIfAborted(signal);
+
+		const selectorInput = sel
+			? parseSqliteSelectorInput(sel)
+			: { subPath: resolvedSqlitePath.sqliteSubPath, queryString: resolvedSqlitePath.queryString };
+		const selector = parseSqliteSelector(selectorInput.subPath, selectorInput.queryString);
+		const details: ReadToolDetails = {
+			resolvedPath: resolvedSqlitePath.absolutePath,
+			suffixResolution: resolvedSqlitePath.suffixResolution,
+		};
+
+		let db: Database | null = null;
+		try {
+			db = new Database(resolvedSqlitePath.absolutePath, { readonly: true, strict: true });
+			db.run("PRAGMA busy_timeout = 3000");
+			throwIfAborted(signal);
+
+			switch (selector.kind) {
+				case "list": {
+					const listLimit = applyListLimit(listTables(db), { limit: 500 });
+					const output = prependSuffixResolutionNotice(
+						renderTableList(listLimit.items),
+						resolvedSqlitePath.suffixResolution,
+					);
+					const truncation = truncateHead(output, { maxLines: Number.MAX_SAFE_INTEGER });
+					details.truncation = truncation.truncated ? truncation : undefined;
+					const resultBuilder = toolResult<ReadToolDetails>(details)
+						.text(truncation.content)
+						.sourcePath(resolvedSqlitePath.absolutePath)
+						.limits({ resultLimit: listLimit.meta.resultLimit?.reached });
+					if (truncation.truncated) {
+						resultBuilder.truncation(truncation, { direction: "head" });
+					}
+					return resultBuilder.done();
+				}
+				case "schema": {
+					const sampleRows = queryRows(db, selector.table, { limit: selector.sampleLimit, offset: 0 });
+					let output = renderSchema(getTableSchema(db, selector.table), {
+						columns: sampleRows.columns,
+						rows: sampleRows.rows,
+					});
+					if (sampleRows.rows.length < sampleRows.totalCount) {
+						const remaining = sampleRows.totalCount - sampleRows.rows.length;
+						output += `\n[${remaining} more rows; use sel="${selector.table}?limit=20&offset=${sampleRows.rows.length}" to continue]`;
+					}
+					return toolResult<ReadToolDetails>(details)
+						.text(prependSuffixResolutionNotice(output, resolvedSqlitePath.suffixResolution))
+						.sourcePath(resolvedSqlitePath.absolutePath)
+						.done();
+				}
+				case "row": {
+					const lookup = resolveTableRowLookup(db, selector.table);
+					const row =
+						lookup.kind === "pk"
+							? getRowByKey(db, selector.table, lookup, selector.key)
+							: getRowByRowId(db, selector.table, selector.key);
+					if (!row) {
+						return toolResult<ReadToolDetails>(details)
+							.text(
+								prependSuffixResolutionNotice(
+									`No row found in table '${selector.table}' for key '${selector.key}'.`,
+									resolvedSqlitePath.suffixResolution,
+								),
+							)
+							.sourcePath(resolvedSqlitePath.absolutePath)
+							.done();
+					}
+					return toolResult<ReadToolDetails>(details)
+						.text(prependSuffixResolutionNotice(renderRow(row), resolvedSqlitePath.suffixResolution))
+						.sourcePath(resolvedSqlitePath.absolutePath)
+						.done();
+				}
+				case "query": {
+					const page = queryRows(db, selector.table, selector);
+					return toolResult<ReadToolDetails>(details)
+						.text(
+							prependSuffixResolutionNotice(
+								renderTable(page.columns, page.rows, {
+									totalCount: page.totalCount,
+									offset: selector.offset,
+									limit: selector.limit,
+									table: selector.table,
+									dbPath: resolvedSqlitePath.absolutePath,
+								}),
+								resolvedSqlitePath.suffixResolution,
+							),
+						)
+						.sourcePath(resolvedSqlitePath.absolutePath)
+						.done();
+				}
+				case "raw": {
+					const result = executeReadQuery(db, selector.sql);
+					return toolResult<ReadToolDetails>(details)
+						.text(
+							prependSuffixResolutionNotice(
+								renderTable(result.columns, result.rows, {
+									totalCount: result.rows.length,
+									offset: 0,
+									limit: result.rows.length || DEFAULT_MAX_LINES,
+									table: "query",
+									dbPath: resolvedSqlitePath.absolutePath,
+								}),
+								resolvedSqlitePath.suffixResolution,
+							),
+						)
+						.sourcePath(resolvedSqlitePath.absolutePath)
+						.done();
+				}
+			}
+
+			throw new ToolError("Unsupported SQLite selector");
+		} catch (error) {
+			if (error instanceof ToolError) {
+				throw error;
+			}
+			throw new ToolError(error instanceof Error ? error.message : String(error));
+		} finally {
+			db?.close();
+		}
+	}
+
 	async execute(
 		_toolCallId: string,
 		params: ReadParams,
@@ -769,6 +981,11 @@ export class ReadTool implements AgentTool<typeof readSchema, ReadToolDetails> {
 			return this.#readArchive(readPath, offset, limit, archivePath, signal);
 		}
 
+		const sqlitePath = await this.#resolveSqliteReadPath(readPath, signal);
+		if (sqlitePath) {
+			return this.#readSqlite(sel, sqlitePath, signal);
+		}
+
 		let absolutePath = resolveReadPath(localReadPath, this.session.cwd);
 		let suffixResolution: { from: string; to: string } | undefined;