messages.ts

/**
 * Custom message types and transformers for the coding agent.
 *
 * Extends the base AgentMessage type with coding-agent specific message types,
 * and provides a transformer to convert them to LLM-compatible messages.
 */

import type { AgentMessage } from "@earendil-works/pi-agent-core";
import {
	copyContextProvenance,
	dropFailedAssistantTurns,
	type ImageContent,
	type Message,
	type TextContent,
} from "@earendil-works/pi-ai";
import { ENVIRONMENT_CONTEXT_MESSAGE_TYPE, foldEnvironmentContextIntoNextUserMessage } from "./environment-context.ts";
import { replaceElidedImages } from "./image-elision-placeholder.ts";
import { omitProviderRejectedImages } from "./provider-rejected-images.ts";
import { appendToolResultUserWords } from "./tool-result-user-words.ts";

export const COMPACTION_SUMMARY_PREFIX = `The conversation history before this point was compacted into the following summary:

<summary>
`;

export const COMPACTION_SUMMARY_SUFFIX = `
</summary>`;

export const BRANCH_SUMMARY_PREFIX = `The following is a summary of a branch that this conversation came back from:

<summary>
`;

export const BRANCH_SUMMARY_SUFFIX = `</summary>`;

export const GOAL_CONTINUATION_MESSAGE_TYPE = "goal-continuation";

/**
 * Per-type exclusion must remain false: compaction and branch summarization
 * inspect entries independently, where excluding this type would hide the live
 * goal-continuation message.
 */
export function isContextExcludedCustomMessage(_customType: string): boolean {
	return false;
}

/**
 * Whole-context filtering is deliberately a no-op. Goal continuations are
 * append-only: dropping an already-sent continuation would rewrite a
 * provider-visible turn, so request N would stop being a prefix of request N+1
 * and every cached token ahead of the edit would be invalidated (full re-read,
 * cost spike, 429 storms in team mode). Continuation history is bounded by
 * normal compaction instead of by per-request deletion.
 */
export function filterContextExcludedMessages(messages: AgentMessage[]): AgentMessage[] {
	return messages;
}

/**
 * Message type for bash executions via the ! command.
 */
export interface BashExecutionMessage {
	role: "bashExecution";
	command: string;
	output: string;
	exitCode: number | undefined;
	cancelled: boolean;
	truncated: boolean;
	fullOutputPath?: string;
	timestamp: number;
	/** If true, this message is excluded from LLM context (!! prefix) */
	excludeFromContext?: boolean;
}

/**
 * Message type for extension-injected messages via sendMessage().
 * These are custom messages that extensions can inject into the conversation.
 */
export interface CustomMessage<T = unknown> {
	role: "custom";
	customType: string;
	content: string | (TextContent | ImageContent)[];
	display: boolean;
	details?: T;
	timestamp: number;
}

export interface BranchSummaryMessage {
	role: "branchSummary";
	summary: string;
	fromId: string | null;
	timestamp: number;
}

export interface CompactionSummaryMessage {
	role: "compactionSummary";
	summary: string;
	tokensBefore: number;
	details?: unknown;
	timestamp: number;
}

export interface ConfigurationUpdateMessage {
	role: "configurationUpdate";
	content: (TextContent | ImageContent)[];
	effort: string;
	timestamp: number;
}

// Extend CustomAgentMessages via declaration merging
declare module "@earendil-works/pi-agent-core" {
	interface CustomAgentMessages {
		bashExecution: BashExecutionMessage;
		custom: CustomMessage;
		branchSummary: BranchSummaryMessage;
		configurationUpdate: ConfigurationUpdateMessage;
	}
}

/**
 * Convert a BashExecutionMessage to user message text for LLM context.
 */
export function bashExecutionToText(msg: BashExecutionMessage): string {
	let text = `Ran \`${msg.command}\`\n`;
	if (msg.output) {
		text += `\`\`\`\n${msg.output}\n\`\`\``;
	} else {
		text += "(no output)";
	}
	if (msg.cancelled) {
		text += "\n\n(command cancelled)";
	} else if (msg.exitCode !== null && msg.exitCode !== undefined && msg.exitCode !== 0) {
		text += `\n\nCommand exited with code ${msg.exitCode}`;
	}
	if (msg.truncated && msg.fullOutputPath) {
		text += `\n\n[Output truncated. Full output: ${msg.fullOutputPath}]`;
	}
	return text;
}

export function createBranchSummaryMessage(summary: string, fromId: string, timestamp: string): BranchSummaryMessage {
	return {
		role: "branchSummary",
		summary,
		fromId,
		timestamp: new Date(timestamp).getTime(),
	};
}

export function createCompactionSummaryMessage(
	summary: string,
	tokensBefore: number,
	timestamp: string,
	details?: unknown,
): CompactionSummaryMessage {
	return {
		role: "compactionSummary",
		summary: summary,
		tokensBefore,
		details,
		timestamp: new Date(timestamp).getTime(),
	};
}

/** Convert CustomMessageEntry to AgentMessage format */
export function createCustomMessage(
	customType: string,
	content: string | (TextContent | ImageContent)[],
	display: boolean,
	details: unknown | undefined,
	timestamp: string,
): CustomMessage {
	return {
		role: "custom",
		customType,
		content,
		display,
		details,
		timestamp: new Date(timestamp).getTime(),
	};
}

/**
 * Transform AgentMessages (including custom types) to LLM-compatible Messages.
 *
 * This is used by:
 * - Agent's transormToLlm option (for prompt calls and queued messages)
 * - Compaction's generateSummary (for summarization)
 * - Custom extensions and tools
 */
export function convertToLlm(messages: AgentMessage[]): Message[] {
	const withContextProvenance = <T extends Message>(source: AgentMessage, target: T): T =>
		copyContextProvenance(source, target);
	// Continuations are append-only here too: the transport array must extend the
	// previous request verbatim to keep the provider's cache prefix valid.
	const environmentMessages = new Set<Message>();
	const converted = messages
		.map((m): Message | undefined => {
			switch (m.role) {
				case "bashExecution":
					// Skip messages excluded from context (!! prefix)
					if (m.excludeFromContext) {
						return undefined;
					}
					return withContextProvenance(m, {
						role: "user",
						content: [{ type: "text", text: bashExecutionToText(m) }],
						timestamp: m.timestamp,
					});
				case "custom": {
					if (isContextExcludedCustomMessage(m.customType)) {
						return undefined;
					}

					const content = typeof m.content === "string" ? [{ type: "text" as const, text: m.content }] : m.content;
					const converted = withContextProvenance(m, {
						role: "user" as const,
						content,
						timestamp: m.timestamp,
					});
					if (m.customType === ENVIRONMENT_CONTEXT_MESSAGE_TYPE) environmentMessages.add(converted);
					return converted;
				}
				case "branchSummary":
					return withContextProvenance(m, {
						role: "user",
						content: [{ type: "text" as const, text: BRANCH_SUMMARY_PREFIX + m.summary + BRANCH_SUMMARY_SUFFIX }],
						timestamp: m.timestamp,
					});
				case "compactionSummary":
					return withContextProvenance(m, {
						role: "user",
						content: [
							{ type: "text" as const, text: COMPACTION_SUMMARY_PREFIX + m.summary + COMPACTION_SUMMARY_SUFFIX },
						],
						timestamp: m.timestamp,
					});
				case "configurationUpdate":
					return m;
				case "system":
				case "user":
				case "assistant":
				case "toolResult":
					return m;
				default: {
					const _exhaustiveCheck: never = m;
					return undefined;
				}
			}
		})
		.filter((m) => m !== undefined);
	// Failed provider turns (stopReason error/aborted) must never reach an LLM
	// request: lanes that build requests straight from this output (claude-sdk
	// prompt bridge, cursor turns) would otherwise replay their partial text and
	// unexecuted tool calls, and token estimation would count them. Dropping is
	// deterministic per session state, so the cache-prefix guarantee above holds.
	// senpi#2118: the environment context then becomes the leading block of the
	// user message it precedes, so strict-alternation chat templates never see
	// two consecutive user messages. senpi#2170: images a failed turn proves the
	// provider rejected are replaced first, while the failed turn is still visible.
	// senpi#2920: words a user typed into a tool's answer follow that batch's tool
	// results as a user turn, rebuilt here from the persisted result on every request.
	return foldEnvironmentContextIntoNextUserMessage(
		appendToolResultUserWords(dropFailedAssistantTurns(omitProviderRejectedImages(converted))),
		environmentMessages,
	);
}

// ============================================================================
// Transport image budget
// ============================================================================

/**
 * Cap on inline image base64 (characters approximately equal wire bytes) per
 * provider request. Anthropic documents a 32 MB request limit; reserving the
 * remainder for text and request overhead keeps the transport below that wall.
 */
export const TRANSPORT_IMAGE_BUDGET_BYTES = 24 * 1024 * 1024;

export { IMAGE_ELISION_PLACEHOLDER } from "./image-elision-placeholder.ts";

export const BLOCKED_IMAGE_PLACEHOLDER = "Image reading is disabled.";

export interface ElideOldImagesOptions {
	/** Cumulative inline image base64 budget. Defaults to TRANSPORT_IMAGE_BUDGET_BYTES. */
	budgetBytes?: number;
	/** Newest image blocks always kept regardless of budget. They still consume it. Defaults to 1. */
	alwaysKeepNewest?: number;
	/** Keep at most this many images from turns completed by a later assistant response. */
	maxHistoricalImages?: number;
}

export interface TransportConvertOptions extends ElideOldImagesOptions {
	/** Replace every image when the images.blockImages setting is enabled. */
	blockImages: boolean;
}

/** Drop consecutive duplicates of a placeholder produced by adjacent image replacements. */
export function dedupeConsecutivePlaceholder(
	content: (TextContent | ImageContent)[],
	placeholder: string,
): (TextContent | ImageContent)[] {
	return content.filter(
		(block, index, blocks) =>
			!(
				block.type === "text" &&
				block.text === placeholder &&
				index > 0 &&
				blocks[index - 1].type === "text" &&
				(blocks[index - 1] as TextContent).text === placeholder
			),
	);
}

/**
 * Bound inline image payload at request-build time. The oldest images are
 * replaced by a text placeholder in the message that held them, so a dropped
 * image never moves to another turn.
 *
 * The cutoff moves in fixed steps rather than on every request (senpi#2768).
 * History only grows by appending, so each image's offset from the oldest one
 * never changes; images are dropped once they end below a threshold that only
 * rises when the total crosses the next half-budget step (likewise half the
 * historical-image limit). A request that just adds an image therefore sends
 * every earlier image exactly as before, and a provider cache or a sent-history
 * hash keyed on that history stays valid until the next step. Crossing a step drops up to
 * half the budget at once, so the request shrinks noticeably then grows back image by image.
 *
 * The newest protected blocks, the current turn (when a historical limit is
 * set) and the newest historical images under that limit are always kept. The
 * input and persisted session remain untouched. If every image fits, the
 * original array reference is returned.
 */
export function elideOldImages(messages: Message[], options?: ElideOldImagesOptions): Message[] {
	const budgetBytes = options?.budgetBytes ?? TRANSPORT_IMAGE_BUDGET_BYTES;
	const alwaysKeepNewest = options?.alwaysKeepNewest ?? 1;
	const maxHistoricalImages = options?.maxHistoricalImages;
	const lastAssistantIndex = messages.findLastIndex((message) => message.role === "assistant");

	const images: { key: string; bytes: number; currentTurn: boolean }[] = [];
	messages.forEach((message, messageIndex) => {
		if ((message.role !== "user" && message.role !== "toolResult") || !Array.isArray(message.content)) return;
		message.content.forEach((block, blockIndex) => {
			if (block.type !== "image") return;
			images.push({
				key: `${messageIndex}:${blockIndex}`,
				bytes: block.data.length,
				currentTurn: maxHistoricalImages !== undefined && messageIndex > lastAssistantIndex,
			});
		});
	});

	const imagesToElide = new Set<string>();
	if (maxHistoricalImages !== undefined) {
		const historical = images.filter((image) => !image.currentTurn);
		const step = Math.max(1, Math.floor(maxHistoricalImages / 2));
		const excess = historical.length - maxHistoricalImages;
		const cutoff = excess > 0 ? Math.ceil(excess / step) * step : 0;
		for (const image of historical.slice(0, cutoff)) imagesToElide.add(image.key);
	} else {
		const protectedKeys = new Set(
			images.slice(Math.max(0, images.length - alwaysKeepNewest)).map((image) => image.key),
		);
		const total = images.reduce((sum, image) => sum + image.bytes, 0);
		if (total > budgetBytes) {
			const step = Math.max(1, Math.floor(budgetBytes / 2));
			const threshold = Math.ceil((total - budgetBytes) / step) * step;
			let offset = 0;
			let elidedBytes = 0;
			let straddler: (typeof images)[number] | undefined;
			for (const image of images) {
				const start = offset;
				offset += image.bytes;
				if (protectedKeys.has(image.key)) continue;
				if (offset <= threshold) {
					imagesToElide.add(image.key);
					elidedBytes += image.bytes;
				} else if (start < threshold && straddler === undefined) {
					straddler = image;
				}
			}
			if (straddler !== undefined && total - elidedBytes > budgetBytes) imagesToElide.add(straddler.key);
		}
	}

	if (imagesToElide.size === 0) return messages;

	return messages.map((message, messageIndex) => {
		if ((message.role !== "user" && message.role !== "toolResult") || !Array.isArray(message.content)) {
			return message;
		}
		if (!message.content.some((_block, blockIndex) => imagesToElide.has(`${messageIndex}:${blockIndex}`))) {
			return message;
		}

		return {
			...message,
			content: replaceElidedImages(message.content, (blockIndex) =>
				imagesToElide.has(`${messageIndex}:${blockIndex}`),
			) as typeof message.content,
		};
	});
}

/** Convert messages for the main provider transport and apply its image policy. */
export function convertToLlmForTransport(messages: AgentMessage[], options: TransportConvertOptions): Message[] {
	const converted = convertToLlm(messages);
	if (!options.blockImages) return elideOldImages(converted, options);

	return converted.map((message) => {
		if ((message.role !== "user" && message.role !== "toolResult") || !Array.isArray(message.content)) {
			return message;
		}
		if (!message.content.some((block) => block.type === "image")) return message;

		const replaced = message.content.map((block): TextContent | ImageContent =>
			block.type === "image" ? { type: "text", text: BLOCKED_IMAGE_PLACEHOLDER } : block,
		);
		return {
			...message,
			content: dedupeConsecutivePlaceholder(replaced, BLOCKED_IMAGE_PLACEHOLDER) as typeof message.content,
		};
	});
}