zk-data-agent/frontend/app/lib/use-training-mode.ts

"use client";

import { useCallback, useEffect, useState } from "react";

export const SESSION_UPDATED_EVENT = "claw-session-updated";
const APPLIED_STORAGE_PREFIX = "claw.training.applied:";

const MODEL_ITERATION_SKILL_NAMES = new Set(["model-iteration"]);

const SKILL_POLL_MS = 8000;

export function useModelIterationEnabled(): {
	enabled: boolean;
	loading: boolean;
} {
	const [enabled, setEnabled] = useState(false);
	const [loading, setLoading] = useState(true);

	useEffect(() => {
		let cancelled = false;
		const tick = async () => {
			try {
				const res = await fetch("/api/claw/skills", { cache: "no-store" });
				if (!res.ok) {
					if (!cancelled) {
						setEnabled(false);
						setLoading(false);
					}
					return;
				}
				const skills = await res.json();
				if (!Array.isArray(skills)) {
					if (!cancelled) {
						setEnabled(false);
						setLoading(false);
					}
					return;
				}
				const isOn = skills.some(
					(s: { name?: string; enabled?: boolean }) =>
						MODEL_ITERATION_SKILL_NAMES.has(String(s?.name ?? "")) &&
						s?.enabled !== false,
				);
				if (!cancelled) {
					setEnabled(isOn);
					setLoading(false);
				}
			} catch {
				if (!cancelled) {
					setEnabled(false);
					setLoading(false);
				}
			}
		};
		tick();
		const id = window.setInterval(tick, SKILL_POLL_MS);
		const onFocus = () => tick();
		window.addEventListener("focus", onFocus);
		return () => {
			cancelled = true;
			window.clearInterval(id);
			window.removeEventListener("focus", onFocus);
		};
	}, []);

	return { enabled, loading };
}

// applied 状态用 localStorage（按 session id）持久化：
// - 关 tab / 浏览器重启都不丢
// - 刷新页面也保持
// 不走后端 PATCH 是为了避开 __LOCALID_* 临时 session id 导致 404
export function useAppliedTraining(sessionId: string | null) {
	// 不再从 localStorage 恢复 applied 状态——面板只通过用户本次会话中
	// 显式点击 skill 按钮打开，避免页面加载时闪烁。
	const [applied, setApplied] = useState(false);

	const setAppliedPersist = useCallback(
		(next: boolean) => {
			setApplied(next);
			if (typeof window === "undefined") return;
			if (!sessionId) return;
			const key = APPLIED_STORAGE_PREFIX + sessionId;
			if (next) {
				window.localStorage.setItem(key, "1");
			} else {
				window.localStorage.removeItem(key);
			}
		},
		[sessionId],
	);

	return { applied, setApplied: setAppliedPersist };
}

export type ProgramStatus = {
	available: boolean;
	path?: string;
	last_modified_ms?: number | null;
	error?: string;
};

const PROGRAM_STATUS_POLL_MS = 5000;

export function useProgramStatus(): ProgramStatus & { loading: boolean } {
	const [status, setStatus] = useState<ProgramStatus>({ available: false });
	const [loading, setLoading] = useState(true);

	useEffect(() => {
		let cancelled = false;
		const tick = async () => {
			try {
				const res = await fetch("/api/claw/training/program-status", {
					cache: "no-store",
				});
				if (!res.ok) {
					if (!cancelled) {
						setStatus({ available: false });
						setLoading(false);
					}
					return;
				}
				const payload = (await res.json()) as ProgramStatus;
				if (!cancelled) {
					setStatus(payload);
					setLoading(false);
				}
			} catch {
				if (!cancelled) {
					setStatus({ available: false });
					setLoading(false);
				}
			}
		};
		tick();
		const id = window.setInterval(tick, PROGRAM_STATUS_POLL_MS);
		const onFocus = () => tick();
		window.addEventListener("focus", onFocus);
		return () => {
			cancelled = true;
			window.clearInterval(id);
			window.removeEventListener("focus", onFocus);
		};
	}, []);

	return { ...status, loading };
}

type SessionUpdatedDetail = {
	sessionId: string;
	is_training?: boolean;
};

export function dispatchSessionUpdated(detail: SessionUpdatedDetail) {
	if (typeof window === "undefined") return;
	window.dispatchEvent(
		new CustomEvent<SessionUpdatedDetail>(SESSION_UPDATED_EVENT, { detail }),
	);
}

export function useTrainingMode(sessionId: string | null) {
	const [isTraining, setIsTraining] = useState(false);
	const [loaded, setLoaded] = useState(false);

	useEffect(() => {
		setLoaded(false);
		if (!sessionId) {
			setIsTraining(false);
			return;
		}
		let cancelled = false;
		fetch(`/api/claw/sessions/${encodeURIComponent(sessionId)}`, {
			cache: "no-store",
		})
			.then(async (res) => {
				if (!res.ok) return null;
				return res.json();
			})
			.then((payload) => {
				if (cancelled) return;
				const next = Boolean(payload?.is_training);
				setIsTraining(next);
				setLoaded(true);
			})
			.catch(() => {
				if (!cancelled) setLoaded(true);
			});
		return () => {
			cancelled = true;
		};
	}, [sessionId]);

	const setTraining = useCallback(
		async (next: boolean) => {
			if (!sessionId) return;
			const previous = isTraining;
			setIsTraining(next);
			try {
				const res = await fetch(
					`/api/claw/sessions/${encodeURIComponent(sessionId)}`,
					{
						method: "PATCH",
						headers: { "content-type": "application/json" },
						body: JSON.stringify({ is_training: next }),
					},
				);
				if (!res.ok) {
					setIsTraining(previous);
					return;
				}
				dispatchSessionUpdated({ sessionId, is_training: next });
			} catch {
				setIsTraining(previous);
			}
		},
		[sessionId, isTraining],
	);

	return { isTraining, setTraining, loaded };
}