@galaxyproject/orbit 0.0.0 → 0.9.0

package/extensions/orbit/session-index/parse.ts

@@ -0,0 +1,254 @@
+import fs from "node:fs";
+
+export interface SessionHeader {
+  id: string;
+  cwd: string;
+  createdAt: string;
+  parentSession: string | null;
+}
+
+export interface EntryRow {
+  entry_id: string;
+  parent_id: string | null;
+  entry_type: string;
+  timestamp: string;
+  role: "user" | "assistant" | null;
+  text_content: string | null;
+  raw_json: string;
+}
+
+export interface ToolCallRow {
+  entry_id: string;
+  tool_use_id: string;
+  tool_name: string;
+  arguments_json: string;
+  result_text: string | null;
+}
+
+export interface ParseResult {
+  header: SessionHeader;
+  entries: EntryRow[];
+  tool_calls: ToolCallRow[];
+  /** Latest customType=galaxy_analyst_plan entry's data.notebookPath, if any. */
+  notebookPath: string | null;
+  /** Latest customType=session_info name, if any. */
+  sessionName: string | null;
+  /**
+   * Byte offset (exclusive) of the last newline-terminated line parsed.
+   * Trailing bytes after the last newline are not consumed -- the next
+   * scan will re-read them, so a torn read while Pi is actively writing
+   * doesn't silently drop the partial line.
+   */
+  endOffset: number;
+}
+
+export interface ParseOptions {
+  /** Start parsing from this byte offset (used for incremental scans). */
+  startOffset?: number;
+  /** Skip emitting the header if we're resuming past it. */
+  skipHeader?: boolean;
+}
+
+/**
+ * Parse a Pi session JSONL file into normalized rows for the index.
+ *
+ * Does NOT throw on malformed lines; invalid JSON lines are skipped and
+ * the caller's endOffset still advances past them (so incremental scans
+ * don't loop on a corrupt line).
+ */
+export function parseSessionFile(filePath: string, opts: ParseOptions = {}): ParseResult {
+  const startOffset = opts.startOffset ?? 0;
+  const buf = fs.readFileSync(filePath);
+  const slice = buf.subarray(startOffset);
+  const text = slice.toString("utf8");
+
+  const entries: EntryRow[] = [];
+  const toolCalls: ToolCallRow[] = [];
+  const pendingResults = new Map<string, string>();
+  let header: SessionHeader | null = null;
+  let notebookPath: string | null = null;
+  let sessionName: string | null = null;
+
+  let lineStart = 0;
+  let lastCompleteLineEnd = 0;
+  for (let i = 0; i < text.length; i++) {
+    if (text.charCodeAt(i) !== 10 /* \n */) continue;
+    const raw = text.slice(lineStart, i);
+    lineStart = i + 1;
+    lastCompleteLineEnd = i + 1;
+    if (raw.length === 0) continue;
+    let obj: Record<string, unknown>;
+    try {
+      obj = JSON.parse(raw);
+    } catch {
+      continue;
+    }
+
+    if (obj.type === "session") {
+      if (!opts.skipHeader) {
+        header = {
+          id: String(obj.id ?? ""),
+          cwd: String(obj.cwd ?? ""),
+          createdAt: String(obj.timestamp ?? ""),
+          parentSession: typeof obj.parentSession === "string" ? obj.parentSession : null,
+        };
+      }
+      continue;
+    }
+
+    const row = entryRowFromObj(obj);
+    if (!row) continue;
+    entries.push(row);
+
+    if (row.entry_type === "message" && row.role === "assistant") {
+      for (const tc of extractToolUses(obj)) {
+        toolCalls.push({ entry_id: row.entry_id, ...tc });
+      }
+    }
+    if (row.entry_type === "message" && row.role === "user") {
+      for (const { tool_use_id, text_content } of extractToolResults(obj)) {
+        pendingResults.set(tool_use_id, text_content);
+      }
+    }
+
+    if (row.entry_type === "custom") {
+      const customType = String((obj as { customType?: unknown }).customType ?? "");
+      if (customType === "galaxy_analyst_plan") {
+        const data = (obj as { data?: { notebookPath?: string } }).data;
+        if (data?.notebookPath) notebookPath = data.notebookPath;
+      }
+    }
+    if (row.entry_type === "session_info") {
+      const name = (obj as { name?: string }).name;
+      if (name) sessionName = name;
+    }
+  }
+
+  // Join tool-call results by tool_use_id (now a first-class field).
+  const byUseId = new Map<string, ToolCallRow>();
+  for (const tc of toolCalls) {
+    byUseId.set(tc.tool_use_id, tc);
+  }
+  for (const [useId, text] of pendingResults) {
+    const target = byUseId.get(useId);
+    if (target) target.result_text = text;
+  }
+
+  if (!header && !opts.skipHeader) {
+    throw new Error(`No session header found in ${filePath}`);
+  }
+
+  return {
+    header: header ?? ({ id: "", cwd: "", createdAt: "", parentSession: null } as SessionHeader),
+    entries,
+    tool_calls: toolCalls,
+    notebookPath,
+    sessionName,
+    endOffset: startOffset + Buffer.byteLength(text.slice(0, lastCompleteLineEnd), "utf8"),
+  };
+}
+
+function entryRowFromObj(obj: Record<string, unknown>): EntryRow | null {
+  const type = typeof obj.type === "string" ? obj.type : "";
+  const entryId = typeof obj.id === "string" ? obj.id : null;
+  if (!type || !entryId) return null;
+
+  const parentId = typeof obj.parentId === "string" ? obj.parentId : null;
+  const timestamp = typeof obj.timestamp === "string" ? obj.timestamp : "";
+
+  let role: "user" | "assistant" | null = null;
+  let text: string | null = null;
+  if (type === "message") {
+    const msg = (obj as { message?: { role?: string; content?: unknown } }).message;
+    const r = msg?.role;
+    if (r === "user" || r === "assistant") role = r;
+    text = flattenTextBlocks(msg?.content);
+  } else if (type === "compaction") {
+    text =
+      typeof (obj as { summary?: string }).summary === "string"
+        ? (obj as { summary: string }).summary
+        : null;
+  } else if (type === "branch_summary") {
+    text =
+      typeof (obj as { summary?: string }).summary === "string"
+        ? (obj as { summary: string }).summary
+        : null;
+  } else if (type === "custom") {
+    text = JSON.stringify((obj as { data?: unknown }).data ?? null);
+  }
+
+  return {
+    entry_id: entryId,
+    parent_id: parentId,
+    entry_type: type,
+    timestamp,
+    role,
+    text_content: text,
+    raw_json: JSON.stringify(obj),
+  };
+}
+
+function flattenTextBlocks(content: unknown): string | null {
+  if (typeof content === "string") return content;
+  if (!Array.isArray(content)) return null;
+  const parts: string[] = [];
+  for (const block of content) {
+    if (typeof block === "string") {
+      parts.push(block);
+      continue;
+    }
+    if (!block || typeof block !== "object") continue;
+    const b = block as { type?: string; text?: string };
+    if (b.type === "text" && typeof b.text === "string") parts.push(b.text);
+  }
+  return parts.length > 0 ? parts.join("\n") : null;
+}
+
+function extractToolUses(
+  obj: Record<string, unknown>,
+): Array<{ tool_use_id: string; tool_name: string; arguments_json: string; result_text: null }> {
+  const content = (obj as { message?: { content?: unknown } }).message?.content;
+  if (!Array.isArray(content)) return [];
+  const out: Array<{
+    tool_use_id: string;
+    tool_name: string;
+    arguments_json: string;
+    result_text: null;
+  }> = [];
+  for (const block of content) {
+    if (!block || typeof block !== "object") continue;
+    const b = block as { type?: string; name?: string; input?: unknown; id?: string };
+    if (b.type !== "tool_use" || typeof b.name !== "string") continue;
+    const tool_use_id =
+      typeof b.id === "string" && b.id.length > 0 ? b.id : `synth-${b.name}-${out.length}`;
+    out.push({
+      tool_use_id,
+      tool_name: b.name,
+      arguments_json: JSON.stringify(b.input ?? {}),
+      result_text: null,
+    });
+  }
+  return out;
+}
+
+function extractToolResults(
+  obj: Record<string, unknown>,
+): Array<{ tool_use_id: string; text_content: string }> {
+  const content = (obj as { message?: { content?: unknown } }).message?.content;
+  if (!Array.isArray(content)) return [];
+  const out: Array<{ tool_use_id: string; text_content: string }> = [];
+  for (const block of content) {
+    if (!block || typeof block !== "object") continue;
+    const b = block as {
+      type?: string;
+      tool_use_id?: string;
+      content?: unknown;
+    };
+    if (b.type !== "tool_result" || typeof b.tool_use_id !== "string") continue;
+    out.push({
+      tool_use_id: b.tool_use_id,
+      text_content: flattenTextBlocks(b.content) ?? "",
+    });
+  }
+  return out;
+}

package/extensions/orbit/session-index/query.ts

@@ -0,0 +1,209 @@
+import type { Database as Db } from "better-sqlite3";
+
+export interface SearchHit {
+  entry_id: string;
+  session_id: string;
+  session_name: string | null;
+  cwd: string;
+  notebook_path: string | null;
+  timestamp: string;
+  role: "user" | "assistant" | null;
+  snippet: string;
+}
+
+export interface ContextRow {
+  entry_id: string;
+  role: "user" | "assistant" | null;
+  entry_type: string;
+  timestamp: string;
+  text: string;
+}
+
+export interface ToolCallHit {
+  entry_id: string;
+  session_id: string;
+  cwd: string;
+  notebook_path: string | null;
+  timestamp: string;
+  arguments: unknown;
+  result_summary: string | null;
+}
+
+export interface SearchParams {
+  query: string;
+  scope?: "all" | "cwd";
+  cwd?: string;
+  limit?: number;
+}
+
+export interface ContextParams {
+  entry_id: string;
+  before?: number;
+  after?: number;
+}
+
+export interface ToolCallParams {
+  tool_name: string;
+  args_contains?: string;
+  scope?: "all" | "cwd";
+  cwd?: string;
+  limit?: number;
+}
+
+const DEFAULT_LIMIT = 20;
+const MAX_LIMIT = 100;
+
+function clampLimit(n?: number): number {
+  if (typeof n !== "number" || !Number.isFinite(n)) return DEFAULT_LIMIT;
+  return Math.max(1, Math.min(MAX_LIMIT, Math.floor(n)));
+}
+
+export function searchChat(db: Db, params: SearchParams): SearchHit[] {
+  const limit = clampLimit(params.limit);
+  const scopeCwd = params.scope === "cwd" && params.cwd;
+  const scopeClause = scopeCwd ? "AND s.cwd = @cwd" : "";
+
+  const rows = db
+    .prepare(
+      `
+      SELECT
+        e.entry_id,
+        e.session_id,
+        s.name          AS session_name,
+        s.cwd,
+        s.notebook_path,
+        e.timestamp,
+        e.role,
+        snippet(entries_fts, 0, '[', ']', '…', 20) AS snippet
+      FROM entries_fts
+      JOIN entries e ON e.rowid = entries_fts.rowid
+      JOIN sessions s ON s.session_id = e.session_id
+      WHERE entries_fts MATCH @query
+        ${scopeClause}
+      ORDER BY e.timestamp DESC
+      LIMIT @limit
+      `,
+    )
+    .all({ query: params.query, cwd: params.cwd ?? "", limit }) as SearchHit[];
+
+  return rows;
+}
+
+export function getSessionContext(db: Db, params: ContextParams): ContextRow[] {
+  const before = Math.max(0, params.before ?? 3);
+  const after = Math.max(0, params.after ?? 3);
+
+  const anchor = db
+    .prepare("SELECT session_id, timestamp FROM entries WHERE entry_id = ?")
+    .get(params.entry_id) as { session_id: string; timestamp: string } | undefined;
+  if (!anchor) return [];
+
+  const beforeRows = db
+    .prepare(
+      `
+      SELECT entry_id, role, entry_type, timestamp, COALESCE(text_content, '') AS text
+      FROM entries
+      WHERE session_id = @sid AND (timestamp, entry_id) < (@ts, @eid)
+      ORDER BY timestamp DESC, entry_id DESC
+      LIMIT @n
+      `,
+    )
+    .all({
+      sid: anchor.session_id,
+      ts: anchor.timestamp,
+      eid: params.entry_id,
+      n: before,
+    }) as ContextRow[];
+
+  const anchorRow = db
+    .prepare(
+      `
+      SELECT entry_id, role, entry_type, timestamp, COALESCE(text_content, '') AS text
+      FROM entries WHERE entry_id = ?
+      `,
+    )
+    .get(params.entry_id) as ContextRow;
+
+  const afterRows = db
+    .prepare(
+      `
+      SELECT entry_id, role, entry_type, timestamp, COALESCE(text_content, '') AS text
+      FROM entries
+      WHERE session_id = @sid AND (timestamp, entry_id) > (@ts, @eid)
+      ORDER BY timestamp ASC, entry_id ASC
+      LIMIT @n
+      `,
+    )
+    .all({
+      sid: anchor.session_id,
+      ts: anchor.timestamp,
+      eid: params.entry_id,
+      n: after,
+    }) as ContextRow[];
+
+  return [...beforeRows.reverse(), anchorRow, ...afterRows];
+}
+
+export function findToolCalls(db: Db, params: ToolCallParams): ToolCallHit[] {
+  const limit = clampLimit(params.limit);
+  const scopeCwd = params.scope === "cwd" && params.cwd;
+  const scopeClause = scopeCwd ? "AND s.cwd = @cwd" : "";
+  const argsClause = params.args_contains ? "AND tc.arguments_json LIKE @args ESCAPE '\\'" : "";
+
+  const bindParams: Record<string, unknown> = { tool: params.tool_name, limit };
+  if (params.args_contains) bindParams.args = `%${escapeLike(params.args_contains)}%`;
+  if (scopeCwd) bindParams.cwd = params.cwd;
+
+  const rows = db
+    .prepare(
+      `
+      SELECT
+        tc.entry_id,
+        tc.session_id,
+        s.cwd,
+        s.notebook_path,
+        e.timestamp,
+        tc.arguments_json,
+        tc.result_text
+      FROM tool_calls tc
+      JOIN entries  e ON e.entry_id  = tc.entry_id
+      JOIN sessions s ON s.session_id = tc.session_id
+      WHERE tc.tool_name = @tool
+        ${argsClause}
+        ${scopeClause}
+      ORDER BY e.timestamp DESC
+      LIMIT @limit
+      `,
+    )
+    .all(bindParams) as Array<{
+    entry_id: string;
+    session_id: string;
+    cwd: string;
+    notebook_path: string | null;
+    timestamp: string;
+    arguments_json: string;
+    result_text: string | null;
+  }>;
+
+  return rows.map((r) => ({
+    entry_id: r.entry_id,
+    session_id: r.session_id,
+    cwd: r.cwd,
+    notebook_path: r.notebook_path,
+    timestamp: r.timestamp,
+    arguments: safeJson(r.arguments_json),
+    result_summary: r.result_text ? r.result_text.slice(0, 500) : null,
+  }));
+}
+
+function safeJson(s: string): unknown {
+  try {
+    return JSON.parse(s);
+  } catch {
+    return s;
+  }
+}
+
+function escapeLike(s: string): string {
+  return s.replace(/[\\%_]/g, "\\$&");
+}

package/extensions/orbit/session-index/schema.ts

@@ -0,0 +1,72 @@
+/** Bump this whenever a migration would be needed. openIndexDb() responds
+ * by dropping the existing DB and rebuilding from scratch -- the source of
+ * truth is the Pi JSONLs, so this is cheap and safe.
+ */
+export const SCHEMA_VERSION = 3;
+
+export const SCHEMA_SQL = `
+CREATE TABLE sessions (
+  session_id          TEXT PRIMARY KEY,
+  file_path           TEXT NOT NULL UNIQUE,
+  cwd                 TEXT NOT NULL,
+  name                TEXT,
+  parent_session      TEXT,
+  notebook_path       TEXT,
+  created_at          TEXT NOT NULL,
+  last_indexed_at     TEXT NOT NULL,
+  last_indexed_offset INTEGER NOT NULL
+);
+CREATE INDEX sessions_cwd ON sessions(cwd);
+
+CREATE TABLE entries (
+  entry_id     TEXT PRIMARY KEY,
+  session_id   TEXT NOT NULL REFERENCES sessions(session_id) ON DELETE CASCADE,
+  parent_id    TEXT,
+  entry_type   TEXT NOT NULL,
+  timestamp    TEXT NOT NULL,
+  role         TEXT,
+  text_content TEXT,
+  raw_json     TEXT NOT NULL
+);
+CREATE INDEX entries_session ON entries(session_id, timestamp);
+CREATE INDEX entries_type    ON entries(entry_type);
+
+CREATE TABLE tool_calls (
+  entry_id       TEXT NOT NULL REFERENCES entries(entry_id) ON DELETE CASCADE,
+  session_id     TEXT NOT NULL REFERENCES sessions(session_id) ON DELETE CASCADE,
+  tool_use_id    TEXT NOT NULL,
+  tool_name      TEXT NOT NULL,
+  arguments_json TEXT NOT NULL,
+  result_text    TEXT,
+  PRIMARY KEY (entry_id, tool_use_id)
+);
+CREATE INDEX tool_calls_name ON tool_calls(tool_name);
+
+CREATE VIRTUAL TABLE entries_fts USING fts5(
+  text_content,
+  content='entries',
+  content_rowid='rowid',
+  tokenize='porter unicode61'
+);
+
+-- Keep FTS in sync with entries
+CREATE TRIGGER entries_ai AFTER INSERT ON entries BEGIN
+  INSERT INTO entries_fts(rowid, text_content)
+  VALUES (new.rowid, coalesce(new.text_content, ''));
+END;
+CREATE TRIGGER entries_ad AFTER DELETE ON entries BEGIN
+  INSERT INTO entries_fts(entries_fts, rowid, text_content)
+  VALUES ('delete', old.rowid, coalesce(old.text_content, ''));
+END;
+CREATE TRIGGER entries_au AFTER UPDATE ON entries BEGIN
+  INSERT INTO entries_fts(entries_fts, rowid, text_content)
+  VALUES ('delete', old.rowid, coalesce(old.text_content, ''));
+  INSERT INTO entries_fts(rowid, text_content)
+  VALUES (new.rowid, coalesce(new.text_content, ''));
+END;
+
+CREATE TABLE meta (
+  key   TEXT PRIMARY KEY,
+  value TEXT NOT NULL
+);
+`;

package/extensions/orbit/session-index/tools.ts

@@ -0,0 +1,152 @@
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { Text } from "@earendil-works/pi-tui";
+import { Type } from "@sinclair/typebox";
+import { openIndexDb, defaultDbPath } from "./db";
+import { scanSessions } from "./indexer";
+import { searchChat, getSessionContext, findToolCalls } from "./query";
+
+export function registerSessionIndexTools(pi: ExtensionAPI): void {
+  // Long-lived DB handle for the lifetime of the extension. The session-index
+  // DB assumes a single writer per user (see db.ts) -- opening once here and
+  // keeping the WAL/shm files live for the process is the intended shape.
+  const db = openIndexDb(defaultDbPath());
+
+  // Dedup refreshes inside a burst of tool calls. Pi routes a model's tool
+  // round-trips through us in quick succession; re-scanning every time would
+  // triple-stat the corpus for no added freshness.
+  const REFRESH_WINDOW_MS = 2_000;
+  let lastScanAt = 0;
+  function refresh(): void {
+    const now = Date.now();
+    if (now - lastScanAt < REFRESH_WINDOW_MS) return;
+    lastScanAt = now;
+    try {
+      scanSessions(db);
+    } catch {
+      // Don't fail a tool call because of a broken scan; return whatever's indexed.
+    }
+  }
+
+  pi.registerTool({
+    name: "chat_search",
+    label: "Search prior sessions",
+    description:
+      "Full-text search across every Pi session you've had in this account. " +
+      "Searches message text only -- for tool-call arguments use chat_find_tool_calls. " +
+      "Default scope is 'all' (every project). Use scope='cwd' to restrict to the " +
+      "current analysis directory. Returns ranked hits with a snippet; follow up " +
+      "with chat_session_context to read surrounding turns.",
+    parameters: Type.Object({
+      query: Type.String({ minLength: 1, description: "FTS5 query (porter-stemmed tokens)." }),
+      scope: Type.Optional(
+        Type.Union([Type.Literal("all"), Type.Literal("cwd")], {
+          description: "Which sessions to search. Default: 'all'.",
+        }),
+      ),
+      limit: Type.Optional(
+        Type.Integer({ minimum: 1, maximum: 100, description: "Max hits. Default 20." }),
+      ),
+    }),
+    async execute(_id, params) {
+      refresh();
+      try {
+        const hits = searchChat(db, {
+          query: params.query,
+          scope: params.scope ?? "all",
+          cwd: process.cwd(),
+          limit: params.limit,
+        });
+        return {
+          content: [{ type: "text" as const, text: JSON.stringify(hits, null, 2) }],
+          details: { count: hits.length, error: undefined as string | undefined },
+        };
+      } catch (err) {
+        const msg = err instanceof Error ? err.message : String(err);
+        return {
+          content: [
+            {
+              type: "text" as const,
+              text: JSON.stringify(
+                { error: msg, hint: "Check FTS5 query syntax (e.g. quoting)." },
+                null,
+                2,
+              ),
+            },
+          ],
+          details: { count: 0, error: msg as string | undefined },
+        };
+      }
+    },
+    renderResult: (result) => {
+      const d = result.details as { count?: number; error?: string } | undefined;
+      if (d?.error) return new Text(`⚠️  ${d.error}`);
+      return new Text(`🔎 ${d?.count ?? 0} hits`);
+    },
+  });
+
+  pi.registerTool({
+    name: "chat_session_context",
+    label: "Read surrounding turns",
+    description:
+      "Given an entry_id (from chat_search or chat_find_tool_calls), return a " +
+      "window of N entries before and after it from the same session. Default " +
+      "window is 3 before + 3 after. Entries above a compaction point are " +
+      "still retrievable (the index keeps everything the JSONL keeps).",
+    parameters: Type.Object({
+      entry_id: Type.String({ minLength: 1 }),
+      before: Type.Optional(Type.Integer({ minimum: 0, maximum: 50, default: 3 })),
+      after: Type.Optional(Type.Integer({ minimum: 0, maximum: 50, default: 3 })),
+    }),
+    async execute(_id, params) {
+      refresh();
+      const rows = getSessionContext(db, {
+        entry_id: params.entry_id,
+        before: params.before,
+        after: params.after,
+      });
+      return {
+        content: [{ type: "text" as const, text: JSON.stringify(rows, null, 2) }],
+        details: { count: rows.length },
+      };
+    },
+    renderResult: (result) => {
+      const d = result.details as { count?: number } | undefined;
+      return new Text(`📜 ${d?.count ?? 0} entries`);
+    },
+  });
+
+  pi.registerTool({
+    name: "chat_find_tool_calls",
+    label: "Find prior tool calls",
+    description:
+      "Structured search over prior tool invocations. Pass tool_name to filter " +
+      "(e.g. 'workflow_set_overrides'); optional args_contains does a substring " +
+      "match on the arguments JSON (underscore and percent are escaped to avoid " +
+      "LIKE wildcard surprises). Returns parsed arguments and a truncated " +
+      "result summary per hit.",
+    parameters: Type.Object({
+      tool_name: Type.String({ minLength: 1 }),
+      args_contains: Type.Optional(Type.String({ minLength: 1 })),
+      scope: Type.Optional(Type.Union([Type.Literal("all"), Type.Literal("cwd")])),
+      limit: Type.Optional(Type.Integer({ minimum: 1, maximum: 100 })),
+    }),
+    async execute(_id, params) {
+      refresh();
+      const calls = findToolCalls(db, {
+        tool_name: params.tool_name,
+        args_contains: params.args_contains,
+        scope: params.scope ?? "all",
+        cwd: process.cwd(),
+        limit: params.limit,
+      });
+      return {
+        content: [{ type: "text" as const, text: JSON.stringify(calls, null, 2) }],
+        details: { count: calls.length, tool: params.tool_name },
+      };
+    },
+    renderResult: (result) => {
+      const d = result.details as { count?: number; tool?: string } | undefined;
+      return new Text(`🛠 ${d?.count ?? 0} ${d?.tool ?? ""} calls`);
+    },
+  });
+}