@tduyng/pi-context 0.0.0-stage → 0.0.2

package/src/cached-request.ts

@@ -0,0 +1,135 @@
+import { Validator } from 'ata-validator'
+import { t } from 'ata-validator/t'
+
+import { convertToLlm, sessionEntryToContextMessages, type SessionEntry } from '@earendil-works/pi-coding-agent'
+import type { AssistantMessage } from '@earendil-works/pi-ai'
+import type { Api, Model } from '@earendil-works/pi-ai'
+const isString = new Validator(t.string()).isValidObject
+
+export type BranchMessage = ReturnType<typeof sessionEntryToContextMessages>[number]
+export type TranscriptMessage = Parameters<typeof convertToLlm>[0][number]
+
+/**
+ * The transcript of the last request, including the leading system message that carries
+ * the prompt and the tool declarations. Replaying it is what makes the summary a cached
+ * continuation rather than a new request.
+ */
+export interface CapturedRequest {
+    transcript: TranscriptMessage[]
+    model: Model<Api>
+    at: number
+}
+
+export interface CaptureStore {
+    captured: CapturedRequest | null
+}
+
+export function forget(store: CaptureStore): void {
+    store.captured = null
+}
+
+const TextBlockSchema = t.object({ type: t.literal('text'), text: t.string() })
+const isTextBlockSchema = new Validator(TextBlockSchema).isValidObject
+
+export function messageText(message: BranchMessage): string {
+    if (!('content' in message)) return ''
+    const content = message.content
+    if (isString(content)) return content
+    if (Array.isArray(content)) {
+        return content.flatMap((block) => (isTextBlockSchema(block) ? [block.text] : [])).join(' ')
+    }
+    return ''
+}
+
+export function assistantText(message: AssistantMessage): string {
+    return message.content.flatMap((block) => (isTextBlockSchema(block) ? [block.text] : [])).join('')
+}
+
+function sameMessage(a: BranchMessage, b: BranchMessage): boolean {
+    if (a.role !== b.role) return false
+    const aTool = 'toolCallId' in a ? a.toolCallId : undefined
+    const bTool = 'toolCallId' in b ? b.toolCallId : undefined
+    if (aTool !== bTool) return false
+    const aTime = 'timestamp' in a ? a.timestamp : undefined
+    const bTime = 'timestamp' in b ? b.timestamp : undefined
+    return aTime === bTime
+}
+
+/**
+ * What the branch gained since the captured request, and proof that the kept region never
+ * moved: pi's own entry projection must contain the captured transcript as a prefix. The
+ * proof would rest on a guess when an entry of an unknown shape cannot be matched, so null
+ * is returned instead.
+ */
+export function walkTrailing(captured: CapturedRequest, branchEntries: SessionEntry[]): BranchMessage[] | null {
+    const projected: BranchMessage[] = []
+    for (const entry of branchEntries) {
+        if (entry.type === 'compaction' || entry.type === 'session_info') continue
+        projected.push(...sessionEntryToContextMessages(entry))
+    }
+    // The capture starts with pi's leading system message; the branch projection holds the
+    // conversation only, so the prefix is compared on the messages both carry.
+    const conversation = convertToLlm(captured.transcript.filter((message) => message.role !== 'system'))
+    if (projected.length < conversation.length) return null
+    for (let index = 0; index < conversation.length; index++) {
+        const expected = conversation[index]
+        const actual = projected[index]
+        if (expected === undefined || actual === undefined || !sameMessage(expected, actual)) return null
+    }
+    return projected.slice(conversation.length)
+}
+
+export interface KeptStart {
+    text: string
+}
+
+const QUOTE_CHARS = 120
+
+/**
+ * The opening of the first message pi keeps verbatim, quoted back to the summarizer. The
+ * quote marks where the kept window begins.
+ */
+export function keptStartQuote(branchEntries: SessionEntry[], firstKeptEntryId: string): KeptStart | null {
+    const start = branchEntries.findIndex((entry) => entry.id === firstKeptEntryId)
+    if (start === -1) return null
+    const kept = branchEntries.slice(start)
+    const quotes = kept
+        .map((entry) => sessionEntryToContextMessages(entry).map(messageText).join(' ').trim())
+        .filter((text) => text !== '')
+    const first = quotes[0]
+    if (first === undefined) return null
+    return { text: first.slice(0, QUOTE_CHARS) }
+}
+
+/**
+ * The summary instruction, appended as the next user turn. It says stop. It says call no
+ * tools. It quotes the opening the summarizer must stop at, and it keeps pi's checkpoint
+ * shape so the file lists and the recent window survive the round trip.
+ */
+export function summaryInstruction(keptStart: KeptStart | null, customFocus: string | undefined): BranchMessage {
+    const boundary =
+        keptStart === null
+            ? 'Summarize the conversation so far.'
+            : `Summarize the conversation only up to the message beginning "${keptStart.text}".\n` +
+              'Do not include anything after the boundary.'
+    const message: BranchMessage = {
+        role: 'user',
+        content:
+            `${boundary}\n\n` +
+            'Stop here. Call no tools.\n\n' +
+            'Write the summary in this exact format:\n' +
+            'Goal: what the session was trying to accomplish.\n' +
+            'Constraints: hard constraints that shaped the approach.\n' +
+            'Done: what has been completed, including files and tests.\n' +
+            'In Progress: what was being worked on when the summary was requested.\n' +
+            'Blocked: blockers or unresolved issues.\n' +
+            'Key Decisions: the important decisions and why.\n' +
+            'Next Steps: what should happen next.\n' +
+            'Critical Context: the critical context to carry forward.\n' +
+            'File Lists: read and modified file lists. If none, write "None".\n' +
+            (customFocus === undefined ? '' : `\n${customFocus}\n`) +
+            '\nWrite no other text before or after the format.',
+        timestamp: Date.now()
+    }
+    return message
+}

package/src/collapse.ts

@@ -0,0 +1,245 @@
+import {
+    estimateTokens,
+    type ExtensionAPI,
+    type SessionBeforeTreeEvent,
+    type SessionEntry
+} from '@earendil-works/pi-coding-agent'
+import type { ImageContent, TextContent, Usage } from '@earendil-works/pi-ai'
+import { Validator } from 'ata-validator'
+import { t } from 'ata-validator/t'
+
+const OUTCOME_CHARS = 2_000
+const FILE_LIMIT = 12
+const FAILURE_LIMIT = 5
+const COMMAND_CHARS = 120
+const FAILURE_CHARS = 200
+
+const VALIDATION_COMMAND =
+    /(^|\s)(just\s+\S+|tsc\b|vitest\b|jest\b|pytest\b|ruff\b|cargo\s+(?:build|check|test)\b|go\s+test\b|swift\s+(?:build|test)\b|nix\s+build\b|npm\s+(?:run\s+)?(?:test|check|lint|typecheck|build)\b|pnpm\s+(?:run\s+)?(?:test|check|lint|typecheck|build)\b|yarn\s+(?:test|lint|typecheck|build)\b|bun\s+(?:run\s+|x\s+)?(?:test|check|lint|typecheck|build|vitest)\b)/
+
+const PathInput = t.object({ path: t.string() })
+const isPathInput = new Validator(PathInput).isValidObject
+const BashInput = t.object({ command: t.string() })
+const isBashInput = new Validator(BashInput).isValidObject
+
+export interface CollapseDetails {
+    entries: number
+    tokens: number
+    commands: number
+    modifiedFiles: string[]
+    readFiles: string[]
+    validationCommands: string[]
+    failures: string[]
+}
+
+export interface Collapsed {
+    outcome: string
+    usage: Usage
+    details: CollapseDetails
+}
+
+export function parseTurns(args: string): number | null {
+    const trimmed = args.trim()
+    if (trimmed === '') return 1
+    if (!/^\d+$/.test(trimmed)) return null
+    const turns = Number(trimmed)
+    return turns >= 1 && turns <= 100 ? turns : null
+}
+
+export function collapseAnchor(
+    entries: readonly SessionEntry[],
+    turns: number
+): { anchorId: string; startIndex: number } | undefined {
+    const starts: number[] = []
+    entries.forEach((entry, index) => {
+        if (entry.type === 'message' && entry.message.role === 'user') starts.push(index)
+    })
+    const first = starts[starts.length - turns]
+    if (first === undefined || first >= entries.length - 1) return undefined
+    const anchorId = entries[first - 1]?.id
+    if (anchorId === undefined) return undefined
+    return { anchorId, startIndex: first }
+}
+
+function addUsage(total: Usage, usage: Usage): Usage {
+    return {
+        input: total.input + usage.input,
+        output: total.output + usage.output,
+        cacheRead: total.cacheRead + usage.cacheRead,
+        cacheWrite: total.cacheWrite + usage.cacheWrite,
+        cacheWrite1h: (total.cacheWrite1h ?? 0) + (usage.cacheWrite1h ?? 0),
+        reasoning: (total.reasoning ?? 0) + (usage.reasoning ?? 0),
+        totalTokens: total.totalTokens + usage.totalTokens,
+        cost: {
+            input: total.cost.input + usage.cost.input,
+            output: total.cost.output + usage.cost.output,
+            cacheRead: total.cost.cacheRead + usage.cost.cacheRead,
+            cacheWrite: total.cost.cacheWrite + usage.cost.cacheWrite,
+            total: total.cost.total + usage.cost.total
+        }
+    }
+}
+
+function oneLine(text: string, limit: number): string {
+    const flat = text.replace(/\s+/g, ' ').trim()
+    return flat.length <= limit ? flat : `${flat.slice(0, limit)}...`
+}
+
+function tail(text: string, limit: number): string {
+    const trimmed = text.trim()
+    return trimmed.length <= limit ? trimmed : `...${trimmed.slice(trimmed.length - limit)}`
+}
+
+function resultText(content: readonly (TextContent | ImageContent)[]): string {
+    return content
+        .filter((block) => block.type === 'text')
+        .map((block) => block.text)
+        .join(' ')
+        .trim()
+}
+
+export function collapseEntries(entries: readonly SessionEntry[]): Collapsed {
+    const modified = new Set<string>()
+    const read = new Set<string>()
+    const validation = new Set<string>()
+    const failures: string[] = []
+    const commands = new Map<string, string>()
+    let outcome = ''
+    let commandCount = 0
+    let tokens = 0
+    let usage: Usage = {
+        input: 0,
+        output: 0,
+        cacheRead: 0,
+        cacheWrite: 0,
+        totalTokens: 0,
+        cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }
+    }
+
+    for (const entry of entries) {
+        if (entry.type !== 'message') continue
+        const message = entry.message
+        tokens += estimateTokens(message)
+        if (message.role === 'assistant') {
+            usage = addUsage(usage, message.usage)
+            for (const block of message.content) {
+                if (block.type === 'text') {
+                    outcome = block.text
+                    continue
+                }
+                if (block.type !== 'toolCall') continue
+                if (isBashInput(block.arguments)) {
+                    commandCount += 1
+                    const command = oneLine(block.arguments.command, COMMAND_CHARS)
+                    commands.set(block.id, command)
+                    if (VALIDATION_COMMAND.test(block.arguments.command)) validation.add(command)
+                    continue
+                }
+                if (!isPathInput(block.arguments)) continue
+                if (block.name === 'read') read.add(block.arguments.path)
+                else if (block.name === 'write' || block.name === 'edit') modified.add(block.arguments.path)
+            }
+            continue
+        }
+        if (message.role !== 'toolResult' || !message.isError) continue
+        const command = commands.get(message.toolCallId)
+        const label = command === undefined ? message.toolName : `${message.toolName} \`${command}\``
+        failures.push(oneLine(`${label}: ${resultText(message.content)}`, FAILURE_CHARS))
+    }
+
+    return {
+        outcome: tail(outcome, OUTCOME_CHARS),
+        usage,
+        details: {
+            entries: entries.length,
+            tokens,
+            commands: commandCount,
+            modifiedFiles: [...modified].toSorted(),
+            readFiles: [...read].filter((path) => !modified.has(path)).toSorted(),
+            validationCommands: [...validation].toSorted(),
+            failures
+        }
+    }
+}
+
+function list(items: readonly string[], limit: number): string {
+    const shown = items.slice(0, limit)
+    const omitted = items.length - shown.length
+    const lines = shown.map((item) => `- ${item}`)
+    if (omitted > 0) lines.push(`- ... ${omitted} more`)
+    return lines.length === 0 ? '- none' : lines.join('\n')
+}
+
+export function renderSummary(details: CollapseDetails, outcome: string): string {
+    const sections = [`Collapsed ${details.entries} entries, about ${details.tokens.toLocaleString('en-US')} tokens.`]
+    if (outcome !== '') sections.push(`Outcome:\n${outcome}`)
+    if (details.modifiedFiles.length > 0)
+        sections.push(`Modified files (${details.modifiedFiles.length}):\n${list(details.modifiedFiles, FILE_LIMIT)}`)
+    if (details.readFiles.length > 0)
+        sections.push(`Read (${details.readFiles.length}):\n${list(details.readFiles, FILE_LIMIT)}`)
+    const validation =
+        details.validationCommands.length === 0
+            ? ''
+            : `; validation: ${list(details.validationCommands, FAILURE_LIMIT)}`
+    sections.push(`Commands: ${details.commands}${validation}`)
+    if (details.failures.length > 0)
+        sections.push(`Failures (${details.failures.length}):\n${list(details.failures, FAILURE_LIMIT)}`)
+    return sections.join('\n\n')
+}
+
+const HANDOFF = [
+    'Earlier turns were collapsed to reclaim context. The summary is below.',
+    '',
+    'Treat it as the state of that work. Do not search session files or logs for the raw turns, and read a project file again only when the next task needs its current content.',
+    '',
+    '<collapsed>'
+].join('\n')
+
+export default function collapseExtension(pi: ExtensionAPI): void {
+    let pending: { anchorId: string; summary?: string } | undefined
+
+    pi.on('session_before_tree', (event: SessionBeforeTreeEvent) => {
+        if (pending === undefined || event.preparation.targetId !== pending.anchorId) return undefined
+        const { outcome, usage, details } = collapseEntries(event.preparation.entriesToSummarize)
+        const summary = renderSummary(details, outcome)
+        pending.summary = summary
+        return { summary: { summary, details, usage } }
+    })
+
+    pi.registerCommand('collapse', {
+        description: 'Replace finished turns with a structural summary. Full history stays in the session tree.',
+        handler: async (args, ctx) => {
+            if (!ctx.isIdle()) {
+                ctx.ui.notify('Agent is busy. Wait for it to finish.', 'warning')
+                return
+            }
+            const turns = parseTurns(args)
+            if (turns === null) {
+                ctx.ui.notify('Usage: /collapse [turns], turns is 1-100.', 'warning')
+                return
+            }
+            const anchor = collapseAnchor(ctx.sessionManager.getBranch(), turns)
+            if (anchor === undefined) {
+                ctx.ui.notify(
+                    `Nothing to collapse: ${turns} finished turn(s) need an entry before them to collapse into.`,
+                    'warning'
+                )
+                return
+            }
+            pending = { anchorId: anchor.anchorId }
+            let summary: string | undefined
+            try {
+                await ctx.navigateTree(anchor.anchorId, { summarize: true, label: 'collapsed' })
+                summary = pending.summary
+            } finally {
+                pending = undefined
+            }
+            if (summary === undefined) return
+            pi.sendMessage(
+                { customType: 'context-collapse', content: `${HANDOFF}\n${summary}\n</collapsed>`, display: false },
+                { deliverAs: 'nextTurn' }
+            )
+            ctx.ui.notify(`Collapsed ${turns} turn(s). Full history stays in /tree.`, 'info')
+        }
+    })
+}

package/src/compaction-handoff.ts

@@ -0,0 +1,147 @@
+import type { CompactionEntry, ExtensionAPI, SessionEntry } from '@earendil-works/pi-coding-agent'
+
+const RECORD_LIMIT_CHARS = 20_000
+const EXCERPT_CHARS = 4_000
+const OVERHEAD_RESERVE_CHARS = 2_000
+
+interface RecordInput {
+    label: string
+    timestamp: string
+    text: string
+}
+
+interface UserMessage {
+    role: string
+    toolName?: string
+    isError?: boolean
+    content?: string | Array<{ type: string; text?: string }>
+}
+
+interface CustomMessageEntryLike {
+    type: 'custom_message'
+    customType: string
+    display: boolean
+    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')
+}
+
+function excerpt(text: string, limit: number): string {
+    if (text.length <= limit) return text
+    const marker = '\n… middle omitted …\n'
+    const head = Math.floor((limit - marker.length) / 2)
+    return `${text.slice(0, head)}${marker}${text.slice(text.length - (limit - marker.length - head))}`
+}
+
+export function recoveryInput(entry: SessionEntry): RecordInput | undefined {
+    if (entry.type === 'message') {
+        const message = entry.message as UserMessage
+        if (message.role === 'user') {
+            return { label: 'user input', timestamp: entry.timestamp, text: contentText(message.content ?? '').trim() }
+        }
+        if (message.role === 'toolResult' && message.toolName === 'ask_question' && message.isError !== true) {
+            return {
+                label: 'user answer via ask_question',
+                timestamp: entry.timestamp,
+                text: contentText(message.content ?? '').trim()
+            }
+        }
+        return undefined
+    }
+    if (entry.type !== 'custom_message') return undefined
+    const custom = entry as CustomMessageEntryLike
+    if (custom.display !== true || custom.customType === 'context-prune-nudge') return undefined
+    return {
+        label: `visible ${custom.customType ?? 'custom'} message (injected, not typed by the user)`,
+        timestamp: entry.timestamp,
+        text: contentText(custom.content).trim()
+    }
+}
+
+export function compactedSpan(branch: SessionEntry[], compaction: CompactionEntry): SessionEntry[] {
+    const indexOf = (id: string | undefined): number => branch.findIndex((entry) => entry.id === id)
+    const boundary = indexOf(compaction.firstKeptEntryId)
+    const end = boundary === -1 ? indexOf(compaction.id) : boundary
+    if (end < 1) return []
+    for (let i = end - 1; i >= 0; i--) {
+        const entry = branch[i]
+        if (entry?.type !== 'compaction' || entry.id === compaction.id) continue
+        const kept = indexOf(entry.firstKeptEntryId)
+        return branch.slice(kept === -1 ? i + 1 : kept, end)
+    }
+    return branch.slice(0, end)
+}
+
+function formatInput(input: RecordInput): string {
+    return `[${input.label} | ${input.timestamp}] ${excerpt(input.text, EXCERPT_CHARS)}`
+}
+
+export function buildRecoveryRecord(span: SessionEntry[]): string | undefined {
+    const inputs = span.map(recoveryInput).filter((input): input is RecordInput => input !== undefined)
+    if (inputs.length === 0) return undefined
+    const preamble = [
+        'Verbatim input record for the span the summary above covers.',
+        'This lists what was asked, not what happened; verify live state before stateful or external work.',
+        'Use the history tool to recover anything else from before the compaction.'
+    ].join('\n')
+
+    const required = new Set<RecordInput>()
+    const formatted = new Map<RecordInput, string>()
+    const format = (input: RecordInput): string => {
+        const existing = formatted.get(input)
+        if (existing !== undefined) return existing
+        const value = formatInput(input)
+        formatted.set(input, value)
+        return value
+    }
+    const firstUser = inputs.find((input) => input.label === 'user input')
+    const latestUser = inputs.findLast((input) => input.label === 'user input')
+    const latest = inputs.at(-1)
+    for (const input of [firstUser, latestUser, latest]) if (input) required.add(input)
+
+    const fixedChars =
+        preamble.length +
+        [...required].reduce((sum, input) => sum + format(input).length + 2, 0) +
+        OVERHEAD_RESERVE_CHARS
+    let budget = Math.max(0, RECORD_LIMIT_CHARS - fixedChars)
+    for (let i = inputs.length - 1; i >= 0; i--) {
+        const input = inputs[i]!
+        if (required.has(input)) continue
+        const roughCost = Math.min(input.text.length, EXCERPT_CHARS) + input.label.length + input.timestamp.length + 32
+        if (roughCost > budget) continue
+        const cost = format(input).length + 2
+        if (cost > budget) continue
+        budget -= cost
+        required.add(input)
+    }
+
+    const omitted = inputs.filter((input) => !required.has(input))
+    const omission =
+        omitted.length > 0
+            ? `Omitted ${omitted.length} input(s) within the record limit (${omitted[0]!.timestamp} through ${omitted.at(-1)!.timestamp}); recover them with the history tool.`
+            : undefined
+    return [preamble, ...inputs.filter((input) => required.has(input)).map((input) => format(input)), omission]
+        .filter(Boolean)
+        .join('\n\n')
+        .slice(0, RECORD_LIMIT_CHARS)
+}
+
+export default function compactionHandoff(pi: ExtensionAPI): void {
+    pi.on('session_compact', (event, ctx) => {
+        if (event.fromExtension) return
+        if (event.reason === 'manual') return
+        const span = compactedSpan(ctx.sessionManager.getBranch() as SessionEntry[], event.compactionEntry)
+        const record = buildRecoveryRecord(span)
+        if (!record) return
+        pi.sendMessage(
+            { customType: 'compaction-input-record', content: record, display: false, details: undefined },
+            { deliverAs: 'steer' }
+        )
+    })
+}

package/src/context-dedupe.ts

@@ -0,0 +1,16 @@
+import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'
+
+export default function dedupeContextFiles(pi: ExtensionAPI): void {
+    pi.on('before_agent_start', (event) => {
+        const files = event.systemPromptOptions.contextFiles ?? []
+        if (!files.length) return
+
+        const seen = new Set<string>()
+        const deduped = files.filter((file) => {
+            if (seen.has(file.content)) return false
+            seen.add(file.content)
+            return true
+        })
+        if (deduped.length !== files.length) event.systemPromptOptions.contextFiles = deduped
+    })
+}