@tduyng/pi-context 0.0.0-stage → 0.0.2

package/src/context.ts

@@ -0,0 +1,313 @@
+import {
+    buildSessionContext,
+    DynamicBorder,
+    estimateTokens,
+    type ExtensionAPI,
+    type ThemeColor
+} from '@earendil-works/pi-coding-agent'
+import { Container, HStack, matchesKey, Spacer, Text, truncateToWidth } from '@earendil-works/pi-tui'
+
+const CHARS_PER_TOKEN = 4
+const MAP_COLS = 24
+const MIN_MAP_ROWS = 6
+const MAX_CHILDREN_SHOWN = 30
+const PANEL_SHARE = 0.7
+const CHROME_ROWS = 6
+const MIN_RIGHT_WIDTH = 30
+const PALETTE: ThemeColor[] = [
+    'accent',
+    'success',
+    'warning',
+    'error',
+    'syntaxString',
+    'syntaxFunction',
+    'syntaxKeyword',
+    'syntaxNumber'
+]
+
+function textTokens(text: string): number {
+    return Math.ceil(text.length / CHARS_PER_TOKEN)
+}
+
+function formatTokens(tokens: number): string {
+    return tokens >= 1000 ? `${(tokens / 1000).toFixed(1)}k` : String(tokens)
+}
+
+interface Row {
+    label: string
+    tokens: number
+    children?: Map<string, number>
+}
+
+function addTo(totals: Map<string, number>, label: string, tokens: number): void {
+    totals.set(label, (totals.get(label) ?? 0) + tokens)
+}
+
+function buildMap(
+    rows: Row[],
+    windowTokens: number,
+    mapRows: number,
+    fg: (color: ThemeColor, text: string) => string
+): string[] {
+    const totalCells = MAP_COLS * mapRows
+    const cells: string[] = []
+    let cumTokens = 0
+    let filledCells = 0
+    rows.forEach((row, index) => {
+        cumTokens += row.tokens
+        const target = Math.min(totalCells, Math.round((cumTokens / windowTokens) * totalCells))
+        const count = Math.max(0, target - filledCells)
+        filledCells = target
+        const glyph = fg(PALETTE[index % PALETTE.length] as ThemeColor, '■')
+        for (let i = 0; i < count; i++) cells.push(glyph)
+    })
+    const freeGlyph = fg('muted', '▫')
+    while (cells.length < totalCells) cells.push(freeGlyph)
+
+    const lines: string[] = []
+    for (let r = 0; r < mapRows; r++) lines.push(cells.slice(r * MAP_COLS, (r + 1) * MAP_COLS).join(' '))
+    return lines
+}
+
+function leaderLine(label: string, value: string, leaderWidth: number, indent: string): string {
+    const dots = '.'.repeat(Math.max(1, leaderWidth - label.length))
+    return `${indent}${label} ${dots} ${value}`
+}
+
+function childLines(
+    children: Map<string, number>,
+    leaderWidth: number,
+    fg: (c: ThemeColor, t: string) => string
+): string[] {
+    const sorted = [...children.entries()].toSorted((a, b) => b[1] - a[1])
+    const shown = sorted.slice(0, MAX_CHILDREN_SHOWN)
+    const rest = sorted.slice(MAX_CHILDREN_SHOWN)
+    const lines = shown.map(([name, tokens]) =>
+        fg('dim', leaderLine(name, formatTokens(tokens).padStart(7), leaderWidth, '    ↳ '))
+    )
+    if (rest.length > 0) {
+        const restTokens = rest.reduce((sum, [, tokens]) => sum + tokens, 0)
+        lines.push(
+            fg('dim', leaderLine(`+${rest.length} more`, formatTokens(restTokens).padStart(7), leaderWidth, '    ↳ '))
+        )
+    }
+    return lines
+}
+
+export default function (pi: ExtensionAPI) {
+    pi.registerCommand('context', {
+        description: 'Show context window usage as a colored map and category list',
+        handler: async (_args, ctx) => {
+            if (ctx.mode !== 'tui') {
+                ctx.ui.notify('/context requires TUI mode.', 'warning')
+                return
+            }
+
+            const rowsByLabel = new Map<string, Row>()
+            const rowFor = (label: string): Row => {
+                let row = rowsByLabel.get(label)
+                if (!row) {
+                    row = { label, tokens: 0 }
+                    rowsByLabel.set(label, row)
+                }
+                return row
+            }
+            const add = (label: string, tokens: number): void => {
+                rowFor(label).tokens += tokens
+            }
+            const addChild = (label: string, childName: string, tokens: number): void => {
+                const row = rowFor(label)
+                row.tokens += tokens
+                row.children ??= new Map()
+                addTo(row.children, childName, tokens)
+            }
+
+            add('System prompt', textTokens(ctx.getSystemPrompt()))
+
+            const activeTools = new Set(pi.getActiveTools())
+            for (const tool of pi.getAllTools()) {
+                if (!activeTools.has(tool.name)) continue
+                const tokens = textTokens(tool.name + (tool.description ?? '') + JSON.stringify(tool.parameters ?? {}))
+                addChild(`Tools (${activeTools.size})`, tool.name, tokens)
+            }
+
+            const { messages } = buildSessionContext(ctx.sessionManager.getEntries(), ctx.sessionManager.getLeafId())
+            for (const message of messages) {
+                switch (message.role) {
+                    case 'user':
+                        add('User messages', estimateTokens(message))
+                        break
+                    case 'assistant':
+                        for (const block of message.content) {
+                            if (block.type === 'text') add('Assistant text', textTokens(block.text))
+                            else if (block.type === 'thinking') add('Assistant thinking', textTokens(block.thinking))
+                            else if (block.type === 'toolCall') {
+                                add('Assistant tool calls', textTokens(block.name + JSON.stringify(block.arguments)))
+                            }
+                        }
+                        break
+                    case 'toolResult':
+                        addChild('Tool results', message.toolName, estimateTokens(message))
+                        break
+                    case 'bashExecution':
+                        add('Bash executions', estimateTokens(message))
+                        break
+                    case 'custom':
+                        add(`Extension: ${message.customType}`, estimateTokens(message))
+                        break
+                    case 'branchSummary':
+                    case 'compactionSummary':
+                        add('Compacted history', estimateTokens(message))
+                        break
+                }
+            }
+
+            const rows = [...rowsByLabel.values()]
+                .filter((row) => row.tokens > 0)
+                .toSorted((a, b) => b.tokens - a.tokens)
+            const total = rows.reduce((sum, row) => sum + row.tokens, 0)
+            const reported = ctx.getContextUsage()
+            const windowTokens = reported?.contextWindow ?? (total || 1)
+            const maxParentLabel = Math.max(...rows.map((row) => row.label.length), 'Free space'.length)
+            const maxChildLabel = Math.max(
+                0,
+                ...rows.flatMap((row) => (row.children ? [...row.children.keys()] : [])).map((k) => k.length)
+            )
+            const childLeaderWidth = maxChildLabel + 1
+            const parentLeaderWidth = Math.max(maxParentLabel + 1, childLeaderWidth + 4)
+
+            const reportedPercent = reported?.percent != null ? `${Math.round(reported.percent)}%` : '?%'
+            const headerLine = [
+                ctx.model ? `${ctx.model.provider}/${ctx.model.id}` : undefined,
+                reported?.tokens != null
+                    ? `${formatTokens(reported.tokens)}/${formatTokens(reported.contextWindow)} (${reportedPercent})`
+                    : undefined
+            ]
+                .filter(Boolean)
+                .join('  ·  ')
+
+            await ctx.ui.custom((tui, theme, _kb, done) => {
+                const fg = (color: ThemeColor, text: string) => theme.fg(color, text)
+                const groups = rows.map((row, index) => {
+                    const percent = Math.round((row.tokens / windowTokens) * 100)
+                    const square = fg(PALETTE[index % PALETTE.length] as ThemeColor, '■')
+                    const value = `${formatTokens(row.tokens).padStart(7)}  ${String(percent).padStart(3)}%`
+                    return {
+                        parent: `${square} ${leaderLine(row.label, value, parentLeaderWidth, '')}`,
+                        children: row.children ? childLines(row.children, childLeaderWidth, fg) : []
+                    }
+                })
+                const freeTokens = Math.max(0, windowTokens - total)
+                const freeLine =
+                    reported?.tokens != null
+                        ? fg(
+                              'muted',
+                              `${fg('muted', '▫')} ${leaderLine('Free space', `${formatTokens(freeTokens).padStart(7)}  ${String(Math.round((freeTokens / windowTokens) * 100)).padStart(3)}%`, parentLeaderWidth, '')}`
+                          )
+                        : undefined
+                const titleText = [
+                    fg('accent', theme.bold('Context usage')),
+                    headerLine ? fg('dim', headerLine) : undefined
+                ]
+                    .filter((part): part is string => part !== undefined)
+                    .join('  ')
+                const MAP_TEXT_PAD_X = 1
+                const mapWidth = MAP_COLS * 2 - 1 + MAP_TEXT_PAD_X * 2
+
+                let cached: { width: number; rows: number; lines: string[] } | undefined
+
+                const compose = (width: number): string[] => {
+                    const terminalRows = tui.terminal.rows || 24
+                    const panelRows = Math.max(4, Math.min(terminalRows - 2, Math.floor(terminalRows * PANEL_SHARE)))
+                    const contentRows = Math.max(1, panelRows - CHROME_ROWS)
+                    const showMap = width - mapWidth - 3 >= MIN_RIGHT_WIDTH && contentRows >= MIN_MAP_ROWS
+                    const rightWidth = Math.max(16, (showMap ? width - mapWidth - 3 : width) - 2)
+                    // Parents first: a short window keeps every category it can before any
+                    // child detail, because the categories are the summary and the children
+                    // are the breakdown.
+                    const clipped: string[] = []
+                    if (groups.length > contentRows) {
+                        if (contentRows === 1) {
+                            clipped.push(groups[0]?.parent ?? fg('dim', 'No context in use'))
+                        } else {
+                            const hidden = groups.length - (contentRows - 1)
+                            clipped.push(...groups.slice(0, contentRows - 1).map((group) => group.parent))
+                            clipped.push(
+                                fg(
+                                    'dim',
+                                    truncateToWidth(
+                                        `${hidden} more ${hidden === 1 ? 'category' : 'categories'}`,
+                                        rightWidth,
+                                        '...'
+                                    )
+                                )
+                            )
+                        }
+                    } else {
+                        const reserveFree = freeLine === undefined ? 0 : 1
+                        let childRows = Math.max(0, contentRows - groups.length - reserveFree)
+                        for (const group of groups) {
+                            clipped.push(group.parent)
+                            for (const child of group.children) {
+                                if (childRows <= 0) break
+                                clipped.push(child)
+                                childRows -= 1
+                            }
+                        }
+                        if (freeLine !== undefined && clipped.length < contentRows) clipped.push(freeLine)
+                    }
+                    const lines = clipped.map((line) => truncateToWidth(line, rightWidth, '...'))
+                    const mapRows = Math.min(
+                        contentRows,
+                        Math.max(Math.min(MIN_MAP_ROWS, contentRows), lines.length)
+                    )
+
+                    const container = new Container()
+                    const border = new DynamicBorder((s: string) => theme.fg('accent', s))
+                    container.addChild(border)
+                    container.addChild(new Text(truncateToWidth(titleText, width - 2, '...'), 1, 0))
+                    container.addChild(new Spacer(1))
+                    if (showMap) {
+                        const mapLines = buildMap(rows, windowTokens, mapRows, fg)
+                        container.addChild(
+                            new HStack(
+                                [
+                                    {
+                                        component: new Text(mapLines.join('\n'), MAP_TEXT_PAD_X, 0),
+                                        basis: mapWidth,
+                                        grow: 0,
+                                        shrink: 0
+                                    },
+                                    { component: new Text(lines.join('\n'), 0, 0), grow: 1 }
+                                ],
+                                { gap: 3 }
+                            )
+                        )
+                    } else {
+                        container.addChild(new Text(lines.join('\n'), 0, 0))
+                    }
+                    container.addChild(new Spacer(1))
+                    container.addChild(new Text(fg('dim', 'Press Enter or Esc to close'), 1, 0))
+                    container.addChild(border)
+                    return container.render(width)
+                }
+
+                return {
+                    render: (width: number) => {
+                        const rowsNow = tui.terminal.rows || 24
+                        if (cached?.width !== width || cached.rows !== rowsNow) {
+                            cached = { width, rows: rowsNow, lines: compose(width) }
+                        }
+                        return cached.lines
+                    },
+                    invalidate: () => {
+                        cached = undefined
+                    },
+                    handleInput: (data: string) => {
+                        if (matchesKey(data, 'enter') || matchesKey(data, 'escape')) done(undefined)
+                    }
+                }
+            })
+        }
+    })
+}

package/src/history.ts

@@ -0,0 +1,241 @@
+import { createReadStream, existsSync, readdirSync, statSync } from 'node:fs'
+import { join, relative } from 'node:path'
+import { createInterface } from 'node:readline'
+import type { ExtensionAPI, SessionEntry } from '@earendil-works/pi-coding-agent'
+import type { Infer } from 'ata-validator'
+import { t } from 'ata-validator/t'
+
+const HISTORY_TOOL_NAME = 'history'
+const MAX_READ_CHARS = 20_000
+const MIN_READ_CHARS = 2_000
+const CHARS_PER_TOKEN = 4
+const PAGE_MARGIN_TOKENS = 1_000
+const EXCERPT_LEAD = 100
+const EXCERPT_CHARS = 400
+
+interface FlatEntry {
+    id: string
+    timestamp: string
+    text: string
+    echo: boolean
+}
+
+interface MessageLike {
+    role?: string
+    toolName?: string
+    command?: string
+    output?: string
+    content?: string | Array<{ type: string; text?: string }>
+}
+
+function contentText(content: string | Array<{ type: string; text?: string }> | undefined): string {
+    if (!Array.isArray(content)) return content ?? ''
+    return content
+        .filter((block) => block.type === 'text')
+        .map((block) => block.text ?? '')
+        .join('\n')
+}
+
+export function isEcho(entry: SessionEntry): boolean {
+    if (entry.type === 'compaction' || entry.type === 'branch_summary') return true
+    if (entry.type === 'custom_message') return entry.customType === 'compaction-input-record'
+    const message = entry.type === 'message' ? (entry.message as MessageLike) : undefined
+    return message?.role === 'toolResult' && message.toolName === 'history'
+}
+
+export function flatten(entry: SessionEntry): FlatEntry | undefined {
+    let text: string | undefined
+    if (entry.type === 'message') {
+        const message = entry.message as MessageLike
+        const body =
+            message.role === 'bashExecution'
+                ? `$ ${message.command ?? ''}\n${message.output ?? ''}`
+                : contentText(message.content)
+        text = `[${message.role ?? 'message'}${message.toolName ? `:${message.toolName}` : ''}] ${body}`
+    } else if (entry.type === 'compaction' || entry.type === 'branch_summary') {
+        text = `[${entry.type}] ${entry.summary}`
+    } else if (entry.type === 'custom_message') {
+        text = `[custom:${entry.customType}] ${contentText(entry.content)}`
+    }
+    if (!text?.trim()) return undefined
+    return { id: entry.id, timestamp: entry.timestamp, text, echo: isEcho(entry) }
+}
+
+export function pageChars(usage: { tokens: number | null; contextWindow: number } | undefined): number {
+    if (!usage || usage.tokens == null) return MAX_READ_CHARS
+    const fits = Math.max(0, usage.contextWindow - usage.tokens - PAGE_MARGIN_TOKENS) * CHARS_PER_TOKEN
+    return Math.max(MIN_READ_CHARS, Math.min(MAX_READ_CHARS, fits))
+}
+
+export function queryTerms(query: string): string[] {
+    return query.toLowerCase().split(/\s+/).filter(Boolean)
+}
+
+export function scoreText(text: string, terms: string[]): { score: number; term: string } | undefined {
+    const lower = text.toLowerCase()
+    const phrase = terms.join(' ')
+    if (terms.length > 1 && lower.includes(phrase)) return { score: terms.length + 1, term: phrase }
+    let score = 0
+    let earliest: string | undefined
+    let earliestAt = Number.POSITIVE_INFINITY
+    for (const term of terms) {
+        const at = lower.indexOf(term)
+        if (at === -1) continue
+        score++
+        if (at < earliestAt) {
+            earliestAt = at
+            earliest = term
+        }
+    }
+    return earliest ? { score, term: earliest } : undefined
+}
+
+export function excerptAround(text: string, query: string): string | undefined {
+    const at = text.toLowerCase().indexOf(query)
+    if (at === -1) return undefined
+    const start = Math.max(0, at - EXCERPT_LEAD)
+    const end = start + EXCERPT_CHARS
+    return `${start > 0 ? '…' : ''}${text.slice(start, end)}${end < text.length ? '…' : ''}`
+}
+
+export function byOriginalFirst(a: { score: number; echo: boolean }, b: { score: number; echo: boolean }): number {
+    return Number(a.echo) - Number(b.echo) || b.score - a.score
+}
+
+function formatHit(entry: FlatEntry, excerpt: string, source: string): string {
+    return `${source ? `${source} ` : ''}${entry.timestamp} [${entry.id}] ${entry.echo ? '[echo] ' : ''}${excerpt}`
+}
+
+function sessionFiles(dir: string): string[] {
+    try {
+        if (!existsSync(dir)) return []
+        return readdirSync(dir, { recursive: true, withFileTypes: true })
+            .filter((entry) => entry.isFile() && entry.name.endsWith('.jsonl'))
+            .map((entry) => {
+                const file = join(entry.parentPath, entry.name)
+                return { file, mtimeMs: statSync(file).mtimeMs }
+            })
+            .toSorted((a, b) => b.mtimeMs - a.mtimeMs)
+            .map((entry) => entry.file)
+    } catch {
+        return []
+    }
+}
+
+async function* fileEntries(file: string, signal?: AbortSignal): AsyncGenerator<FlatEntry> {
+    const stream = createReadStream(file, { encoding: 'utf8', signal })
+    const lines = createInterface({ input: stream, crlfDelay: Number.POSITIVE_INFINITY })
+    try {
+        for await (const line of lines) {
+            let entry: FlatEntry | undefined
+            try {
+                entry = flatten(JSON.parse(line) as SessionEntry)
+            } catch {
+                continue
+            }
+            if (entry) yield entry
+        }
+    } finally {
+        lines.close()
+        stream.destroy()
+    }
+}
+
+function textResult(text: string) {
+    return { content: [{ type: 'text' as const, text }], details: undefined }
+}
+
+const historyParams = t.object({
+    op: t.union([t.literal('search'), t.literal('read')]),
+    query: t.optional(t.string()),
+    id: t.optional(t.string()),
+    all: t.optional(t.boolean()),
+    limit: t.optional(t.integer({ minimum: 1, maximum: 50 })),
+    offset: t.optional(t.integer({ minimum: 0 }))
+})
+
+export default function history(pi: ExtensionAPI): void {
+    pi.registerTool({
+        name: HISTORY_TOOL_NAME,
+        label: 'History',
+        description:
+            'Search or read earlier conversation, including pruned or compacted detail. Search the current branch or all project sessions, newest first. Read an entry by id with offset.',
+        promptSnippet: 'Recover earlier conversation',
+        exposure: 'deferred',
+        namespace: {
+            name: 'context',
+            description: 'Earlier conversation: search or read pruned and compacted turns.'
+        },
+        annotations: { readOnlyHint: true, destructiveHint: false, idempotentHint: true, openWorldHint: false },
+        parameters: historyParams,
+
+        async execute(_toolCallId, params: Infer<typeof historyParams>, signal, _onUpdate, ctx) {
+            const manager = ctx.sessionManager
+            const dir = manager.getSessionDir()
+
+            if (params.op === 'search') {
+                const query = params.query?.trim()
+                if (!query) return textResult('query is required for op "search".')
+                const terms = queryTerms(query)
+                const limit = params.limit ?? 10
+                const maxCandidates = limit * 10
+                const candidates: Array<{ score: number; echo: boolean; line: string }> = []
+                const candidate = (entry: FlatEntry, source: string) => {
+                    const hit = scoreText(entry.text, terms)
+                    if (!hit) return undefined
+                    return {
+                        score: hit.score,
+                        echo: entry.echo,
+                        line: formatHit(entry, excerptAround(entry.text, hit.term)!, source)
+                    }
+                }
+
+                if (params.all) {
+                    for (const file of sessionFiles(dir)) {
+                        const recent: Array<{ score: number; echo: boolean; line: string }> = []
+                        for await (const entry of fileEntries(file, signal)) {
+                            const hit = candidate(entry, relative(dir, file))
+                            if (!hit) continue
+                            recent.push(hit)
+                            if (recent.length > maxCandidates) recent.shift()
+                        }
+                        candidates.push(...recent)
+                        if (candidates.length >= maxCandidates) break
+                    }
+                } else {
+                    const branch = manager.getBranch() as SessionEntry[]
+                    for (let i = branch.length - 1; i >= 0 && candidates.length < maxCandidates; i--) {
+                        const entry = flatten(branch[i]!)
+                        const hit = entry && candidate(entry, '')
+                        if (hit) candidates.push(hit)
+                    }
+                }
+                const ranked = candidates.toSorted(byOriginalFirst).slice(0, limit)
+                return textResult(ranked.map((hit) => hit.line).join('\n') || `No history matches "${query}".`)
+            }
+
+            const id = params.id?.trim()
+            if (!id) return textResult('id is required for op "read".')
+            const maxChars = pageChars(ctx.getContextUsage())
+            const slice = (entry: FlatEntry, source: string) => {
+                const offset = Math.min(params.offset ?? 0, entry.text.length)
+                const end = Math.min(entry.text.length, offset + maxChars)
+                const more = end < entry.text.length ? `\nMore remains; read id "${id}" again with offset ${end}.` : ''
+                return textResult(
+                    `${source ? `${source} ` : ''}${entry.timestamp} [${id}] [chars ${offset}-${end} of ${entry.text.length}]\n${entry.text.slice(offset, end)}${more}`
+                )
+            }
+
+            for (const raw of manager.getBranch() as SessionEntry[]) {
+                const entry = raw.id === id ? flatten(raw) : undefined
+                if (entry) return slice(entry, '')
+            }
+            for (const file of sessionFiles(dir)) {
+                for await (const entry of fileEntries(file, signal)) {
+                    if (entry.id === id) return slice(entry, relative(dir, file))
+                }
+            }
+            return textResult(`No history entry with id "${id}".`)
+        }
+    })
+}

package/src/summary.ts

@@ -0,0 +1,55 @@
+import { Validator } from 'ata-validator'
+import { t } from 'ata-validator/t'
+
+import type { CompactionEntry, FileOperations } from '@earendil-works/pi-coding-agent'
+
+const FileListsSchema = t.object({
+    readFiles: t.optional(t.array(t.string())),
+    modifiedFiles: t.optional(t.array(t.string()))
+})
+const isFileListsSchema = new Validator(FileListsSchema).isValidObject
+
+const FileBlock = /^<(read-files|modified-files)>[\s\S]*?<\/\1>\n?/gm
+
+interface FileLists {
+    read: string[]
+    modified: string[]
+}
+
+function listsFrom(details: CompactionEntry['details']): FileLists {
+    if (!isFileListsSchema(details)) return { read: [], modified: [] }
+    return { read: details.readFiles ?? [], modified: details.modifiedFiles ?? [] }
+}
+
+function merge(...lists: string[][]): string[] {
+    return [...new Set(lists.flat())].toSorted()
+}
+
+/**
+ * The summary the model wrote, with pi's own file lists appended. The model may have copied
+ * a previous compaction's lists into the text, so those blocks are stripped first. Lists
+ * from the previous compaction ride along, because pi carries them forward only from its
+ * own compactions.
+ */
+export interface Checkpoint {
+    summary: string
+    details: { readFiles: string[]; modifiedFiles: string[] }
+}
+
+export function checkpointSummary(
+    raw: string,
+    fileOps: FileOperations,
+    previousDetails: CompactionEntry['details']
+): Checkpoint {
+    const previous = listsFrom(previousDetails)
+    const read = merge(previous.read, [...fileOps.read])
+    const modified = merge(previous.modified, [...fileOps.written, ...fileOps.edited])
+    const body = raw.replace(FileBlock, '').trimEnd()
+    return {
+        summary:
+            `${body}\n\n` +
+            `<read-files>\n${read.join('\n')}\n</read-files>\n` +
+            `<modified-files>\n${modified.join('\n')}\n</modified-files>`,
+        details: { readFiles: read, modifiedFiles: modified }
+    }
+}