branch-summarization.ts

/**
 * 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 { AgentMessage, ThinkingLevel } from "@ponythewhite/base-context-agent";
import type { Model, Usage } from "@ponythewhite/base-context-ai";
import { completeInference, InferenceCoordinator } from "../inference-coordinator.js";
import {
	convertToLlm,
	createBranchSummaryMessage,
	createCompactionSummaryMessage,
	createCustomMessage,
} from "../messages.js";
import type { NativeBranchRequestOutputWriter } from "../request-events.js";
import { MODEL_REQUEST_ID_HEADER } from "../semantic-edges.js";
import type { SessionEntry } from "../session-manager.js";
import { estimateTokens } from "./compaction.js";
import {
	computeFileLists,
	createFileOps,
	extractFileOpsFromMessage,
	type FileOperations,
	formatFileOperations,
	SUMMARIZATION_SYSTEM_PROMPT,
	serializeConversation,
} from "./utils.js";
export interface BranchSummaryResult {
	summary?: string;
	readFiles?: string[];
	modifiedFiles?: string[];
	aborted?: boolean;
	error?: string;
	usage?: Usage;
}

// A result must come from this native completion and retain its actual composed text projection.
const nativeBranchResults = new WeakMap<
	BranchSummaryResult,
	{ summary: string; write: NativeBranchRequestOutputWriter }
>();

export function takeNativeBranchSummaryWrite(
	result: BranchSummaryResult | undefined,
	summary: string,
): NativeBranchRequestOutputWriter | undefined {
	const captured = result ? nativeBranchResults.get(result) : undefined;
	if (result) nativeBranchResults.delete(result);
	// Check only an already privately bound projection; never discover a request by matching text.
	return captured?.summary === summary ? captured.write : undefined;
}

/** Details stored in BranchSummaryEntry.details for file tracking */
export interface BranchSummaryDetails {
	readFiles: string[];
	modifiedFiles: string[];
}

export type { FileOperations } from "./utils.js";

export interface BranchPreparation {
	/** Messages extracted for summarization, in chronological order */
	messages: AgentMessage[];
	/** File operations extracted from tool calls */
	fileOps: FileOperations;
	/** Total estimated tokens in messages */
	totalTokens: number;
}

export interface CollectEntriesResult {
	/** Entries to summarize, in chronological order */
	entries: SessionEntry[];
	/** Common ancestor between old and new position, if any */
	commonAncestorId: string | null;
}

export interface GenerateBranchSummaryOptions {
	/** Model to use for summarization */
	model: Model<any>;
	/** Explicit effort only; absent preserves the provider's existing omitted-effort behavior. */
	thinkingLevel?: ThinkingLevel;
	/** API key for the model */
	apiKey: string;
	/** Request headers for the model */
	headers?: Record<string, string>;
	requests?: InferenceCoordinator;
	/** Abort signal for cancellation */
	signal: AbortSignal;
	/** Optional custom instructions for summarization */
	customInstructions?: string;
	/** If true, customInstructions replaces the default prompt instead of being appended */
	replaceInstructions?: boolean;
	/** Tokens reserved for prompt + LLM response (default 16384) */
	reserveTokens?: number;
}
/** Collect the abandoned suffix from two complete, captured chronological paths. */
export function collectEntriesForBranchSummary(
	oldPath: readonly SessionEntry[],
	targetPath: readonly SessionEntry[],
): CollectEntriesResult {
	const oldIds = new Set(oldPath.map((entry) => entry.id));
	let commonAncestorId: string | null = null;
	for (let index = targetPath.length - 1; index >= 0; index--) {
		if (oldIds.has(targetPath[index].id)) {
			commonAncestorId = targetPath[index].id;
			break;
		}
	}
	const first = commonAncestorId === null ? 0 : oldPath.findIndex((entry) => entry.id === commonAncestorId) + 1;
	return { entries: oldPath.slice(first), commonAncestorId };
}
/**
 * Extract AgentMessage from a session entry.
 * Similar to getMessageFromEntry in compaction.ts but also handles compaction entries.
 */
function getMessageFromEntry(entry: SessionEntry): AgentMessage | undefined {
	switch (entry.type) {
		case "message":
			// Tool-result context remains attached to its assistant tool call.
			if (entry.message.role === "toolResult") return undefined;
			return entry.message;

		case "custom_message":
			return createCustomMessage(entry.customType, entry.content, entry.display, entry.details, entry.timestamp);

		case "branch_summary":
			return createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp);

		case "compaction":
			return createCompactionSummaryMessage(
				entry.summary,
				entry.tokensBefore,
				entry.timestamp,
				entry.customInstructions,
			);
		case "thinking_level_change":
		case "model_change":
		case "custom":
		case "label":
		case "session_info":
			return undefined;
	}
}

/**
 * Prepare entries for summarization with token budget.
 *
 * Walks entries from NEWEST to OLDEST, adding messages until we hit the token budget.
 * This ensures we keep the most recent context when the branch is too long.
 *
 * Also collects file operations from:
 * - Tool calls in assistant messages
 * - Existing branch_summary entries' details (for cumulative tracking)
 *
 * @param entries - Entries in chronological order
 * @param tokenBudget - Maximum tokens to include (0 = no limit)
 */
export function prepareBranchEntries(entries: SessionEntry[], tokenBudget: number = 0): BranchPreparation {
	const messages: AgentMessage[] = [];
	const fileOps = createFileOps();
	let totalTokens = 0;

	// First pass: collect file ops from ALL entries (even if they don't fit in token budget)
	// This ensures we capture cumulative file tracking from nested branch summaries
	// Only extract from pi-generated summaries (fromHook !== true), not extension-generated ones
	for (const entry of entries) {
		if (entry.type === "branch_summary" && !entry.fromHook && entry.details) {
			const details = entry.details as BranchSummaryDetails;
			if (Array.isArray(details.readFiles)) {
				for (const f of details.readFiles) fileOps.read.add(f);
			}
			if (Array.isArray(details.modifiedFiles)) {
				for (const f of details.modifiedFiles) {
					fileOps.edited.add(f);
				}
			}
		}
	}
	for (let i = entries.length - 1; i >= 0; i--) {
		const entry = entries[i];
		const message = getMessageFromEntry(entry);
		if (!message) continue;
		extractFileOpsFromMessage(message, fileOps);

		const tokens = estimateTokens(message);
		if (tokenBudget > 0 && totalTokens + tokens > tokenBudget) {
			if (entry.type === "compaction" || entry.type === "branch_summary") {
				if (totalTokens < tokenBudget * 0.9) {
					messages.unshift(message);
					totalTokens += tokens;
				}
			}
			break;
		}

		messages.unshift(message);
		totalTokens += tokens;
	}

	return { messages, fileOps, totalTokens };
}
const BRANCH_SUMMARY_PREAMBLE = `The user explored a different conversation branch before returning here.
Summary of that exploration:

`;

const BRANCH_SUMMARY_PROMPT = `Create a structured summary of this conversation branch for context when returning later.

Use this EXACT format:

## Goal
[What was the user trying to accomplish in this branch?]

## Constraints & Preferences
- [Any constraints, preferences, or requirements mentioned]
- [Or "(none)" if none were mentioned]

## Progress
### Done
- [x] [Completed tasks/changes]

### In Progress
- [ ] [Work that was started but not finished]

### Blocked
- [Issues preventing progress, if any]

## Key Decisions
- **[Decision]**: [Brief rationale]

## Next Steps
1. [What should happen next to continue this work]

Keep each section concise. Preserve exact file paths, function names, and error messages.`;

/**
 * Generate a summary of abandoned branch entries.
 *
 * @param entries - Session entries to summarize (chronological order)
 * @param options - Generation options
 */
export async function generateBranchSummary(
	entries: SessionEntry[],
	options: GenerateBranchSummaryOptions,
): Promise<BranchSummaryResult> {
	const {
		model,
		thinkingLevel,
		apiKey,
		headers,
		requests,
		signal,
		customInstructions,
		replaceInstructions,
		reserveTokens = 16384,
	} = options;
	const contextWindow = model.contextWindow || 128000;
	const tokenBudget = contextWindow - reserveTokens;

	const { messages, fileOps } = prepareBranchEntries(entries, tokenBudget);

	// Nothing model-visible remains after filtering.
	if (messages.length === 0) {
		return { summary: "No content to summarize" };
	}
	// Serialize before the LLM call so it summarizes rather than continues this branch.
	const llmMessages = convertToLlm(messages);
	const conversationText = serializeConversation(llmMessages);
	let instructions: string;
	if (replaceInstructions && customInstructions) {
		instructions = customInstructions;
	} else if (customInstructions) {
		instructions = `${BRANCH_SUMMARY_PROMPT}\n\nAdditional focus: ${customInstructions}`;
	} else {
		instructions = BRANCH_SUMMARY_PROMPT;
	}
	const promptText = `<conversation>\n${conversationText}\n</conversation>\n\n${instructions}`;

	const summarizationMessages = [
		{
			role: "user" as const,
			content: [{ type: "text" as const, text: promptText }],
			timestamp: Date.now(),
		},
	];
	const completion = completeInference(
		requests,
		model,
		{ systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages },
		{
			apiKey,
			headers,
			signal,
			maxTokens: 2048,
			...(thinkingLevel === undefined ? {} : { reasoning: thinkingLevel }),
		},
		{
			purpose: "summary",
			purposeDetail: "branch",
			operationId: headers?.[MODEL_REQUEST_ID_HEADER],
			semanticEdgeId: headers?.[MODEL_REQUEST_ID_HEADER],
		},
	);
	const response = await completion;
	if (response.stopReason === "aborted") {
		return { aborted: true };
	}
	if (response.stopReason === "error") {
		return { error: response.errorMessage || "Summarization failed" };
	}

	let summary = response.content
		.filter((c): c is { type: "text"; text: string } => c.type === "text")
		.map((c) => c.text)
		.join("\n");
	summary = BRANCH_SUMMARY_PREAMBLE + summary;
	const { readFiles, modifiedFiles } = computeFileLists(fileOps);
	summary += formatFileOperations(readFiles, modifiedFiles);

	const result = {
		summary: summary || "No summary generated",
		readFiles,
		modifiedFiles,
		usage: response.usage,
	};
	const write =
		requests instanceof InferenceCoordinator
			? InferenceCoordinator.prototype.takeBranchOutput.call(requests, completion)
			: undefined;
	if (write) nativeBranchResults.set(result, { summary: result.summary, write });
	return result;
}