getcompanion-ai/co-mono
1
0
Fork 0
mirror of https://github.com/getcompanion-ai/co-mono.git synced 2026-04-15 11:02:17 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Refactor: move compaction code to src/core/compaction/

Browse source 
- Move compaction.ts to src/core/compaction/compaction.ts
- Extract branch summarization to src/core/compaction/branch-summarization.ts
- Add index.ts to re-export all compaction utilities
- Update all imports across the codebase
This commit is contained in:
Mario Zechner 2025-12-29 20:59:35 +01:00
parent aee61b1a6b
commit fd13b53b1c
10 changed files with 142 additions and 109 deletions
Download patch file Download diff file Expand all files Collapse all files

112
packages/coding-agent/src/core/agent-session.ts
View file

@ -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.
*/

115
packages/coding-agent/src/core/compaction/branch-summarization.ts Normal file
View file

@ -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<any>,
apiKey: string,
signal: AbortSignal,
customInstructions?: string,
): Promise<BranchSummaryResult> {
// 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" };
}

4
packages/coding-agent/src/core/compaction.ts → packages/coding-agent/src/core/compaction/compaction.ts
View file

@ -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.

6
packages/coding-agent/src/core/compaction/index.ts Normal file
View file

@ -0,0 +1,6 @@
/**
* Compaction and summarization utilities.
*/
export * from "./branch-summarization.js";
export * from "./compaction.js";

2
packages/coding-agent/src/core/hooks/types.ts
View file

@ -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";

2
packages/coding-agent/src/core/index.ts
View file

@ -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,

4
packages/coding-agent/src/index.ts
View file

@ -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,

2
packages/coding-agent/src/modes/rpc/rpc-client.ts
View file

@ -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";
// ============================================================================

2
packages/coding-agent/src/modes/rpc/rpc-types.ts
View file

@ -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)

2
packages/coding-agent/test/compaction.test.ts
View file

@ -12,7 +12,7 @@ import {
findCutPoint,
getLastAssistantUsage,
shouldCompact,
} from "../src/core/compaction.js";
} from "../src/core/compaction/index.js";
import {
buildSessionContext,
type CompactionEntry,