src/lib/server/agent/agentMessageReducer.ts

import type { Message } from "$lib/types/Message";
import {
	MessageReasoningUpdateType,
	MessageUpdateStatus,
	MessageUpdateType,
	type MessageStreamUpdate,
	type MessageUpdate,
} from "$lib/types/MessageUpdate";
import { mergeFinalAnswerContent } from "$lib/utils/mergeFinalAnswer";

export interface AgentMessageReduction {
	message: Message;
	title?: string;
	sessionCostUsd?: number;
}

/**
 * Pure server-side equivalent of the live MessageUpdate consumer. It is used by
 * a replacement relay, which has no in-memory conversation object after a pod
 * restart. Stream tokens are stored as length markers, matching the normal
 * generation writer, while content/reasoning remain fully materialized.
 */
export function reduceAgentMessageUpdates(
	current: Message,
	updates: readonly MessageUpdate[]
): AgentMessageReduction {
	const message: Message = {
		...current,
		files: current.files ? [...current.files] : undefined,
		updates: current.updates ? [...current.updates] : [],
	};
	let title: string | undefined;
	let sessionCostUsd: number | undefined;

	for (const update of updates) {
		if (update.type === MessageUpdateType.Stream) {
			message.content += update.token;
		} else if (
			update.type === MessageUpdateType.Reasoning &&
			update.subtype === MessageReasoningUpdateType.Stream
		) {
			message.reasoning = (message.reasoning ?? "") + update.token;
		} else if (update.type === MessageUpdateType.FinalAnswer) {
			message.content = mergeFinalAnswerContent({
				existing: message.content,
				finalText: update.text ?? "",
				hadTools: (message.updates ?? []).some((item) => item.type === MessageUpdateType.Tool),
				isInterrupted: update.interrupted === true,
			});
			message.interrupted = update.interrupted;
		} else if (update.type === MessageUpdateType.File) {
			message.files = [
				...(message.files ?? []),
				{ type: "hash", name: update.name, value: update.sha, mime: update.mime },
			];
		} else if (update.type === MessageUpdateType.RouterMetadata) {
			message.routerMetadata = {
				route: update.route,
				model: update.model,
				provider: update.provider,
			};
		} else if (update.type === MessageUpdateType.Title) {
			title = update.title.replace(/<\/?think>/gi, "").trim();
		} else if (update.type === MessageUpdateType.AgentStats && update.scope === "session") {
			sessionCostUsd = update.costUsd;
		}

		if (
			update.type === MessageUpdateType.Status &&
			update.status === MessageUpdateStatus.KeepAlive
		) {
			continue;
		}
		appendPersistedUpdate(message, update);
	}

	message.updatedAt = new Date();
	return { message, title, sessionCostUsd };
}

function appendPersistedUpdate(message: Message, update: MessageUpdate): void {
	message.updates ??= [];
	if (update.type !== MessageUpdateType.Stream) {
		message.updates.push(update);
		return;
	}

	const length = update.len ?? update.token.length;
	const previous = message.updates.at(-1);
	if (previous?.type === MessageUpdateType.Stream && previous.token === "") {
		message.updates[message.updates.length - 1] = {
			...previous,
			len: (previous.len ?? 0) + length,
		};
		return;
	}
	message.updates.push({
		type: MessageUpdateType.Stream,
		token: "",
		len: length,
	} satisfies MessageStreamUpdate);
}