zk-data-agent/frontend/app/lib/claw-session-cache.ts

"use client";

export type SessionReplaySnapshot<TSession = unknown, TRun = unknown> = {
	sessionId: string;
	session: TSession;
	runStatus: TRun | null;
	signature: string;
	fetchedAt: number;
};

const MAX_REPLAY_CACHE_SIZE = 50;
const replayCache = new Map<string, SessionReplaySnapshot>();

export function readCachedSessionReplay<TSession = unknown, TRun = unknown>(
	sessionId: string,
): SessionReplaySnapshot<TSession, TRun> | null {
	const cached = replayCache.get(sessionId);
	if (!cached) return null;
	// Refresh recency for LRU behavior.
	replayCache.delete(sessionId);
	replayCache.set(sessionId, cached);
	return cached as SessionReplaySnapshot<TSession, TRun>;
}

export function invalidateCachedSessionReplay(sessionId: string) {
	replayCache.delete(sessionId);
}

export async function fetchSessionReplaySnapshot<
	TSession = unknown,
	TRun = unknown,
>(
	sessionId: string,
	options: { signal?: AbortSignal } = {},
): Promise<SessionReplaySnapshot<TSession, TRun> | null> {
	const stateSnapshot = await fetchSessionStateSnapshot<TSession, TRun>(
		sessionId,
		options.signal,
	);
	if (stateSnapshot) return stateSnapshot;

	const [sessionResponse, runStatus] = await Promise.all([
		fetch(`/api/claw/sessions/${encodeURIComponent(sessionId)}`, {
			cache: "no-store",
			signal: options.signal,
		}),
		fetchLatestRunStatusForCache<TRun>(sessionId, options.signal),
	]);
	if (!sessionResponse.ok) return null;
	const session = (await sessionResponse.json()) as TSession;
	return writeCachedSessionReplay(sessionId, session, runStatus);
}

type SessionStateResponse<TRun> = {
	session?: Record<string, unknown>;
	messages?: Array<{ seq?: number; message?: unknown }>;
	run?: TRun | null;
};

async function fetchSessionStateSnapshot<TSession, TRun>(
	sessionId: string,
	signal?: AbortSignal,
): Promise<SessionReplaySnapshot<TSession, TRun> | null> {
	try {
		const response = await fetch(
			`/api/claw/sessions/${encodeURIComponent(sessionId)}/state?after_message_seq=0&after_event_seq=0`,
			{ cache: "no-store", signal },
		);
		if (!response.ok) return null;
		const payload = (await response.json()) as SessionStateResponse<TRun>;
		if (!Array.isArray(payload.messages)) return null;
		const orderedMessages = payload.messages
			.slice()
			.sort((left, right) => (left.seq ?? 0) - (right.seq ?? 0))
			.map((entry) => entry.message)
			.filter((message) => Boolean(message));
		const session = {
			...(payload.session ?? {}),
			session_id: sessionId,
			messages: orderedMessages,
		} as TSession;
		return writeCachedSessionReplay(sessionId, session, payload.run ?? null);
	} catch (error) {
		if (error instanceof DOMException && error.name === "AbortError") {
			throw error;
		}
		return null;
	}
}

function writeCachedSessionReplay<TSession, TRun>(
	sessionId: string,
	session: TSession,
	runStatus: TRun | null,
): SessionReplaySnapshot<TSession, TRun> {
	const snapshot: SessionReplaySnapshot<TSession, TRun> = {
		sessionId,
		session,
		runStatus,
		signature: sessionReplaySignature(session, runStatus),
		fetchedAt: Date.now(),
	};
	replayCache.set(sessionId, snapshot as SessionReplaySnapshot);
	while (replayCache.size > MAX_REPLAY_CACHE_SIZE) {
		const oldest = replayCache.keys().next().value;
		if (!oldest) break;
		replayCache.delete(oldest);
	}
	return snapshot;
}

async function fetchLatestRunStatusForCache<TRun>(
	sessionId: string,
	signal?: AbortSignal,
): Promise<TRun | null> {
	try {
		const url = new URL("/api/claw/runs/latest", window.location.origin);
		url.searchParams.set("session_id", sessionId);
		const response = await fetch(url, { cache: "no-store", signal });
		if (!response.ok) return null;
		return (await response.json()) as TRun;
	} catch (error) {
		if (error instanceof DOMException && error.name === "AbortError") {
			throw error;
		}
		return null;
	}
}

function sessionReplaySignature(session: unknown, runStatusValue: unknown) {
	const sessionObject = isObject(session) ? session : {};
	const runObject = isObject(runStatusValue) ? runStatusValue : {};
	const messages = Array.isArray(sessionObject.messages)
		? sessionObject.messages
		: [];
	const lastMessage = messages.at(-1);
	const lastObject = isObject(lastMessage) ? lastMessage : {};
	const metadata = isObject(lastObject.metadata) ? lastObject.metadata : {};
	const runStatus =
		typeof runObject.status === "string" ? runObject.status : "";
	const runIsActive = runStatus === "queued" || runStatus === "running";
	return [
		sessionObject.session_id,
		sessionObject.turns,
		sessionObject.tool_calls,
		messages.length,
		lastObject.role,
		lastObject.message_id,
		typeof lastObject.content === "string" ? lastObject.content.length : "",
		metadata.elapsed_ms,
		runObject.run_id,
		runStatus,
		runObject.current_stage,
		runIsActive ? "" : runObject.updated_at,
		runObject.finished_at,
		runIsActive ? "" : runObject.elapsed_ms,
	].join("|");
}

function isObject(value: unknown): value is Record<string, unknown> {
	return Boolean(value) && typeof value === "object" && !Array.isArray(value);
}