haiku.rag/app/frontend/lib/sessionStorage.ts

export interface Citation {
	index: number;
	document_id: string;
	chunk_id: string;
	chunk_ids?: string[];
	document_uri: string;
	document_title: string | null;
	page_numbers: number[];
	headings: string[] | null;
	content: string;
	doc_item_refs?: string[];
}

// Matches RAGState from the backend capability. The fields named here are the
// ones this UI reads; the capability owns the rest of its namespace, including
// the evidence record that compaction builds its capsule from, so the state has
// to round-trip whole rather than be rebuilt from known keys.
export interface RAGState {
	citation_index: Record<string, Citation>;
	citations: string[];
	document_filter: string | null;
	searches: Record<string, unknown[]>;
	[key: string]: unknown;
}

export interface StoredMessage {
	id: string;
	role?: string;
	content?: string;
	[key: string]: unknown;
}

export interface StoredSession {
	id: string;
	title: string;
	messages: StoredMessage[];
	// The whole AG-UI state. The rag namespace is not the only one a capability
	// writes: the citation policy records violations beside it.
	agentState: AgentState;
	// Sessions stored before agentState existed.
	ragState?: RAGState;
	createdAt: string;
	updatedAt: string;
}

export const AGUI_STATE_KEY = "rag";

export type AgentState = Record<string, unknown>;

// Reads the rag namespace out of a stored session, whichever way it was stored.
export function ragStateOf(session?: StoredSession): RAGState {
	const namespaced = session?.agentState?.[AGUI_STATE_KEY] as
		| Partial<RAGState>
		| undefined;
	return normalizeRAGState(namespaced ?? session?.ragState);
}

// The state to seed an agent with when a session is resumed.
export function agentStateOf(session?: StoredSession): AgentState {
	return session?.agentState ?? { [AGUI_STATE_KEY]: ragStateOf(session) };
}

const SESSIONS_KEY = "haiku.rag.sessions";
const ACTIVE_SESSION_KEY = "haiku.rag.activeSession";

export function normalizeRAGState(state?: Partial<RAGState>): RAGState {
	return {
		...state,
		citation_index: state?.citation_index ?? {},
		citations: state?.citations ?? [],
		document_filter: state?.document_filter ?? null,
		searches: state?.searches ?? {},
	};
}

export function getLatestCitations(state: RAGState): Citation[] {
	return state.citations
		.map((id) => state.citation_index[id])
		.filter((c): c is Citation => c !== undefined);
}

export function getAllSessions(): StoredSession[] {
	const raw = localStorage.getItem(SESSIONS_KEY);
	if (!raw) return [];
	try {
		return JSON.parse(raw) as StoredSession[];
	} catch {
		return [];
	}
}

export function getSession(id: string): StoredSession | null {
	return getAllSessions().find((s) => s.id === id) ?? null;
}

export function getActiveSessionId(): string | null {
	return localStorage.getItem(ACTIVE_SESSION_KEY);
}

export function setActiveSessionId(id: string): void {
	localStorage.setItem(ACTIVE_SESSION_KEY, id);
}

export function createSession(): StoredSession {
	const now = new Date().toISOString();
	const session: StoredSession = {
		id: crypto.randomUUID(),
		title: "New Session",
		messages: [],
		agentState: { [AGUI_STATE_KEY]: normalizeRAGState() },
		createdAt: now,
		updatedAt: now,
	};
	const sessions = getAllSessions();
	sessions.unshift(session);
	localStorage.setItem(SESSIONS_KEY, JSON.stringify(sessions));
	setActiveSessionId(session.id);
	return session;
}

export function saveSession(session: StoredSession): void {
	const sessions = getAllSessions();
	const idx = sessions.findIndex((s) => s.id === session.id);
	if (idx >= 0) {
		sessions[idx] = session;
	} else {
		sessions.unshift(session);
	}
	localStorage.setItem(SESSIONS_KEY, JSON.stringify(sessions));
}

export function updateSessionMessages(
	id: string,
	messages: StoredMessage[],
	agentState: AgentState,
): void {
	const sessions = getAllSessions();
	const idx = sessions.findIndex((s) => s.id === id);
	if (idx < 0) return;

	const session = sessions[idx];
	session.messages = messages;
	session.agentState = agentState;
	session.updatedAt = new Date().toISOString();

	// Derive title from first user message
	if (session.title === "New Session") {
		const firstUserMsg = messages.find(
			(m) =>
				m.content &&
				typeof m.role === "string" &&
				m.role.toLowerCase() === "user",
		);
		if (firstUserMsg?.content) {
			session.title =
				firstUserMsg.content.length > 60
					? `${firstUserMsg.content.slice(0, 57)}...`
					: firstUserMsg.content;
		}
	}

	sessions[idx] = session;
	localStorage.setItem(SESSIONS_KEY, JSON.stringify(sessions));
}

export function deleteSession(id: string): void {
	const sessions = getAllSessions().filter((s) => s.id !== id);
	localStorage.setItem(SESSIONS_KEY, JSON.stringify(sessions));
	if (getActiveSessionId() === id) {
		localStorage.removeItem(ACTIVE_SESSION_KEY);
	}
}

export function exportSessionToMarkdown(session: StoredSession): void {
	const lines: string[] = [`# ${session.title}`, ""];
	for (const msg of session.messages) {
		const role = typeof msg.role === "string" ? msg.role.toLowerCase() : "";
		if (role === "user" && msg.content) {
			lines.push(`**User:** ${msg.content}`, "");
		} else if (role === "assistant" && msg.content) {
			lines.push(`**Assistant:** ${msg.content}`, "");
		}
	}
	const blob = new Blob([lines.join("\n")], { type: "text/markdown" });
	const url = URL.createObjectURL(blob);
	const a = document.createElement("a");
	a.href = url;
	a.download = `${session.title.replace(/[^a-zA-Z0-9]/g, "_")}.md`;
	a.click();
	URL.revokeObjectURL(url);
}