diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index f335a707..978fb1b8 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -22,9 +22,10 @@ import { type CompactionResult, calculateContextTokens, compact, + generateBranchSummary, prepareCompaction, shouldCompact, -} from "./compaction.js"; +} from "./compaction/index.js"; import type { LoadedCustomTool, SessionEvent as ToolSessionEvent } from "./custom-tools/index.js"; import { exportSessionToHtml } from "./export-html.js"; import type { @@ -1661,10 +1662,17 @@ export class AgentSession { // Run default summarizer if needed let summaryText: string | undefined; if (options.summarize && entriesToSummarize.length > 0 && !hookSummary) { - const result = await this._generateBranchSummary( + const model = this.model!; + const apiKey = await this._modelRegistry.getApiKey(model); + if (!apiKey) { + throw new Error(`No API key for ${model.provider}`); + } + const result = await generateBranchSummary( entriesToSummarize, - options.customInstructions, + model, + apiKey, this._branchSummaryAbortController.signal, + options.customInstructions, ); this._branchSummaryAbortController = undefined; if (result.aborted) { @@ -1738,104 +1746,6 @@ export class AgentSession { return { editorText, cancelled: false, summaryEntry }; } - /** - * Generate a summary of abandoned branch entries. - */ - private async _generateBranchSummary( - entries: SessionEntry[], - customInstructions: string | undefined, - signal: AbortSignal, - ): Promise<{ summary?: string; aborted?: boolean; error?: string }> { - // Convert entries to messages for summarization - const messages: Array<{ role: string; content: string }> = []; - for (const entry of entries) { - if (entry.type === "message") { - const text = this._extractMessageText(entry.message); - if (text) { - messages.push({ role: entry.message.role, content: text }); - } - } else if (entry.type === "custom_message") { - const text = - typeof entry.content === "string" - ? entry.content - : entry.content - .filter((c): c is { type: "text"; text: string } => c.type === "text") - .map((c) => c.text) - .join(""); - if (text) { - messages.push({ role: "user", content: text }); - } - } else if (entry.type === "branch_summary") { - messages.push({ role: "system", content: `[Previous branch summary: ${entry.summary}]` }); - } - } - - if (messages.length === 0) { - return { summary: "No content to summarize" }; - } - - // Build prompt for summarization - const conversationText = messages.map((m) => `${m.role}: ${m.content}`).join("\n\n"); - const instructions = customInstructions - ? `${customInstructions}\n\n` - : "Summarize this conversation branch concisely, capturing key decisions, actions taken, and outcomes.\n\n"; - - const prompt = `${instructions}Conversation:\n${conversationText}`; - - // Get API key for current model (model is checked in navigateTree before calling this) - const model = this.model!; - const apiKey = await this._modelRegistry.getApiKey(model); - if (!apiKey) { - throw new Error(`No API key for ${model.provider}`); - } - - // Call LLM for summarization - const { complete } = await import("@mariozechner/pi-ai"); - const response = await complete( - model, - { - messages: [ - { - role: "user", - content: [{ type: "text", text: prompt }], - timestamp: Date.now(), - }, - ], - }, - { apiKey, signal, maxTokens: 1024 }, - ); - - // Check if aborted or errored - if (response.stopReason === "aborted") { - return { aborted: true }; - } - if (response.stopReason === "error") { - return { error: response.errorMessage || "Summarization failed" }; - } - - const summary = response.content - .filter((c): c is { type: "text"; text: string } => c.type === "text") - .map((c) => c.text) - .join("\n"); - - return { summary: summary || "No summary generated" }; - } - - /** - * Extract text content from any message type. - */ - private _extractMessageText(message: any): string { - if (!message.content) return ""; - if (typeof message.content === "string") return message.content; - if (Array.isArray(message.content)) { - return message.content - .filter((c: any) => c.type === "text") - .map((c: any) => c.text) - .join(""); - } - return ""; - } - /** * Get all user messages from session for branch selector. */ diff --git a/packages/coding-agent/src/core/compaction/branch-summarization.ts b/packages/coding-agent/src/core/compaction/branch-summarization.ts new file mode 100644 index 00000000..7ac097f7 --- /dev/null +++ b/packages/coding-agent/src/core/compaction/branch-summarization.ts @@ -0,0 +1,115 @@ +/** + * Branch summarization for tree navigation. + * + * When navigating to a different point in the session tree, this generates + * a summary of the branch being left so context isn't lost. + */ + +import type { Model } from "@mariozechner/pi-ai"; +import { complete } from "@mariozechner/pi-ai"; +import type { SessionEntry } from "../session-manager.js"; + +const DEFAULT_INSTRUCTIONS = + "Summarize this conversation branch concisely, capturing key decisions, actions taken, and outcomes."; + +export interface BranchSummaryResult { + summary?: string; + aborted?: boolean; + error?: string; +} + +/** + * Extract text content from any message type. + */ +function extractMessageText(message: any): string { + if (!message.content) return ""; + if (typeof message.content === "string") return message.content; + if (Array.isArray(message.content)) { + return message.content + .filter((c: any) => c.type === "text") + .map((c: any) => c.text) + .join(""); + } + return ""; +} + +/** + * Generate a summary of abandoned branch entries. + * + * @param entries - Session entries to summarize + * @param model - Model to use for summarization + * @param apiKey - API key for the model + * @param signal - Abort signal for cancellation + * @param customInstructions - Optional custom instructions for summarization + */ +export async function generateBranchSummary( + entries: SessionEntry[], + model: Model, + apiKey: string, + signal: AbortSignal, + customInstructions?: string, +): Promise { + // Convert entries to messages for summarization + const messages: Array<{ role: string; content: string }> = []; + + for (const entry of entries) { + if (entry.type === "message") { + const text = extractMessageText(entry.message); + if (text) { + messages.push({ role: entry.message.role, content: text }); + } + } else if (entry.type === "custom_message") { + const text = + typeof entry.content === "string" + ? entry.content + : entry.content + .filter((c): c is { type: "text"; text: string } => c.type === "text") + .map((c) => c.text) + .join(""); + if (text) { + messages.push({ role: "user", content: text }); + } + } else if (entry.type === "branch_summary") { + messages.push({ role: "system", content: `[Previous branch summary: ${entry.summary}]` }); + } + } + + if (messages.length === 0) { + return { summary: "No content to summarize" }; + } + + // Build prompt for summarization + const conversationText = messages.map((m) => `${m.role}: ${m.content}`).join("\n\n"); + const instructions = customInstructions ? `${customInstructions}\n\n` : `${DEFAULT_INSTRUCTIONS}\n\n`; + const prompt = `${instructions}Conversation:\n${conversationText}`; + + // Call LLM for summarization + const response = await complete( + model, + { + messages: [ + { + role: "user", + content: [{ type: "text", text: prompt }], + timestamp: Date.now(), + }, + ], + }, + { apiKey, signal, maxTokens: 1024 }, + ); + + // Check if aborted or errored + if (response.stopReason === "aborted") { + return { aborted: true }; + } + if (response.stopReason === "error") { + return { error: response.errorMessage || "Summarization failed" }; + } + + const summary = response.content + .filter((c): c is { type: "text"; text: string } => c.type === "text") + .map((c) => c.text) + .join("\n"); + + return { summary: summary || "No summary generated" }; +} diff --git a/packages/coding-agent/src/core/compaction.ts b/packages/coding-agent/src/core/compaction/compaction.ts similarity index 99% rename from packages/coding-agent/src/core/compaction.ts rename to packages/coding-agent/src/core/compaction/compaction.ts index 35d933a3..61e4b950 100644 --- a/packages/coding-agent/src/core/compaction.ts +++ b/packages/coding-agent/src/core/compaction/compaction.ts @@ -8,8 +8,8 @@ import type { AgentMessage } from "@mariozechner/pi-agent-core"; import type { AssistantMessage, Model, Usage } from "@mariozechner/pi-ai"; import { complete } from "@mariozechner/pi-ai"; -import { convertToLlm, createBranchSummaryMessage, createHookMessage } from "./messages.js"; -import type { CompactionEntry, SessionEntry } from "./session-manager.js"; +import { convertToLlm, createBranchSummaryMessage, createHookMessage } from "../messages.js"; +import type { CompactionEntry, SessionEntry } from "../session-manager.js"; /** * Extract AgentMessage from an entry if it produces one. diff --git a/packages/coding-agent/src/core/compaction/index.ts b/packages/coding-agent/src/core/compaction/index.ts new file mode 100644 index 00000000..4f8ad306 --- /dev/null +++ b/packages/coding-agent/src/core/compaction/index.ts @@ -0,0 +1,6 @@ +/** + * Compaction and summarization utilities. + */ + +export * from "./branch-summarization.js"; +export * from "./compaction.js"; diff --git a/packages/coding-agent/src/core/hooks/types.ts b/packages/coding-agent/src/core/hooks/types.ts index e3d20e31..e8dd689e 100644 --- a/packages/coding-agent/src/core/hooks/types.ts +++ b/packages/coding-agent/src/core/hooks/types.ts @@ -9,7 +9,7 @@ import type { AgentMessage } from "@mariozechner/pi-agent-core"; import type { ImageContent, Message, Model, TextContent, ToolResultMessage } from "@mariozechner/pi-ai"; import type { Component } from "@mariozechner/pi-tui"; import type { Theme } from "../../modes/interactive/theme/theme.js"; -import type { CompactionPreparation, CompactionResult } from "../compaction.js"; +import type { CompactionPreparation, CompactionResult } from "../compaction/index.js"; import type { ExecOptions, ExecResult } from "../exec.js"; import type { HookMessage } from "../messages.js"; import type { ModelRegistry } from "../model-registry.js"; diff --git a/packages/coding-agent/src/core/index.ts b/packages/coding-agent/src/core/index.ts index 4b75e1cd..b25df418 100644 --- a/packages/coding-agent/src/core/index.ts +++ b/packages/coding-agent/src/core/index.ts @@ -12,7 +12,7 @@ export { type SessionStats, } from "./agent-session.js"; export { type BashExecutorOptions, type BashResult, executeBash } from "./bash-executor.js"; -export type { CompactionResult } from "./compaction.js"; +export type { CompactionResult } from "./compaction/index.js"; export { type CustomAgentTool, type CustomToolFactory, diff --git a/packages/coding-agent/src/index.ts b/packages/coding-agent/src/index.ts index 1e563429..1ae08511 100644 --- a/packages/coding-agent/src/index.ts +++ b/packages/coding-agent/src/index.ts @@ -12,6 +12,7 @@ export { export { type ApiKeyCredential, type AuthCredential, AuthStorage, type OAuthCredential } from "./core/auth-storage.js"; // Compaction export { + type BranchSummaryResult, type CompactionResult, type CutPointResult, calculateContextTokens, @@ -20,10 +21,11 @@ export { estimateTokens, findCutPoint, findTurnStartIndex, + generateBranchSummary, generateSummary, getLastAssistantUsage, shouldCompact, -} from "./core/compaction.js"; +} from "./core/compaction/index.js"; // Custom tools export type { AgentToolUpdateCallback, diff --git a/packages/coding-agent/src/modes/rpc/rpc-client.ts b/packages/coding-agent/src/modes/rpc/rpc-client.ts index f8d8f213..0249ca11 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-client.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-client.ts @@ -10,7 +10,7 @@ import type { AgentEvent, AgentMessage, ThinkingLevel } from "@mariozechner/pi-a import type { ImageContent } from "@mariozechner/pi-ai"; import type { SessionStats } from "../../core/agent-session.js"; import type { BashResult } from "../../core/bash-executor.js"; -import type { CompactionResult } from "../../core/compaction.js"; +import type { CompactionResult } from "../../core/compaction/index.js"; import type { RpcCommand, RpcResponse, RpcSessionState } from "./rpc-types.js"; // ============================================================================ diff --git a/packages/coding-agent/src/modes/rpc/rpc-types.ts b/packages/coding-agent/src/modes/rpc/rpc-types.ts index d4f7a73a..5feead90 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-types.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-types.ts @@ -9,7 +9,7 @@ import type { AgentMessage, ThinkingLevel } from "@mariozechner/pi-agent-core"; import type { ImageContent, Model } from "@mariozechner/pi-ai"; import type { SessionStats } from "../../core/agent-session.js"; import type { BashResult } from "../../core/bash-executor.js"; -import type { CompactionResult } from "../../core/compaction.js"; +import type { CompactionResult } from "../../core/compaction/index.js"; // ============================================================================ // RPC Commands (stdin) diff --git a/packages/coding-agent/test/compaction.test.ts b/packages/coding-agent/test/compaction.test.ts index 331c3b1b..0468a582 100644 --- a/packages/coding-agent/test/compaction.test.ts +++ b/packages/coding-agent/test/compaction.test.ts @@ -12,7 +12,7 @@ import { findCutPoint, getLastAssistantUsage, shouldCompact, -} from "../src/core/compaction.js"; +} from "../src/core/compaction/index.js"; import { buildSessionContext, type CompactionEntry,