lovart3/hooks/use-pending-generation.ts
2025-12-19 01:54:20 +08:00

114 lines
3.4 KiB
TypeScript

"use client";
import { useEffect, useState, useCallback } from "react";
/**
* 持久化的生成任务状态
*/
export interface PendingGeneration {
chatId: string;
modelId: string;
modelType: "gemini" | "volcano-vision" | "volcano-video" | "qianwen-video" | "other";
timestamp: number;
userMessageId?: string;
}
const STORAGE_KEY = "ai-chatbot-pending-generations";
const EXPIRY_TIME = 5 * 60 * 1000; // 5 分钟超时
/**
* 获取所有未完成的生成任务
*/
function getPendingGenerations(): PendingGeneration[] {
if (typeof window === "undefined") return [];
try {
const data = localStorage.getItem(STORAGE_KEY);
if (!data) return [];
const generations: PendingGeneration[] = JSON.parse(data);
// 过滤掉过期的任务
const now = Date.now();
return generations.filter((g) => now - g.timestamp < EXPIRY_TIME);
} catch {
return [];
}
}
/**
* 保存未完成的生成任务
*/
function savePendingGenerations(generations: PendingGeneration[]) {
if (typeof window === "undefined") return;
try {
localStorage.setItem(STORAGE_KEY, JSON.stringify(generations));
} catch (e) {
console.error("Failed to save pending generations:", e);
}
}
/**
* Hook: 管理持久化的生成任务状态
* 解决用户离开页面再返回时加载状态丢失的问题
*/
export function usePendingGeneration(chatId: string) {
const [pendingGeneration, setPendingGeneration] = useState<PendingGeneration | null>(null);
// 初始化:检查是否有当前 chatId 的未完成任务
useEffect(() => {
const generations = getPendingGenerations();
console.log("[PendingGeneration] 初始化检查, chatId:", chatId);
console.log("[PendingGeneration] localStorage 中的所有任务:", generations);
const found = generations.find((g) => g.chatId === chatId);
if (found) {
setPendingGeneration(found);
console.log("[PendingGeneration] ✅ 恢复处理中状态:", found);
} else {
console.log("[PendingGeneration] ❌ 未找到匹配的任务");
}
}, [chatId]);
// 开始生成:保存状态
const startGeneration = useCallback(
(modelId: string, modelType: PendingGeneration["modelType"], userMessageId?: string) => {
const newGeneration: PendingGeneration = {
chatId,
modelId,
modelType,
timestamp: Date.now(),
userMessageId,
};
// 更新 localStorage
const generations = getPendingGenerations();
// 移除同一 chatId 的旧记录
const filtered = generations.filter((g) => g.chatId !== chatId);
savePendingGenerations([...filtered, newGeneration]);
setPendingGeneration(newGeneration);
console.log("[PendingGeneration] 开始生成:", newGeneration);
},
[chatId]
);
// 完成生成:清除状态
const completeGeneration = useCallback(() => {
const generations = getPendingGenerations();
const filtered = generations.filter((g) => g.chatId !== chatId);
savePendingGenerations(filtered);
setPendingGeneration(null);
console.log("[PendingGeneration] 生成完成,清除状态");
}, [chatId]);
// 检查是否有恢复的任务(用户离开再返回)
const hasRestoredPending = pendingGeneration !== null && Date.now() - pendingGeneration.timestamp > 1000;
return {
pendingGeneration,
hasRestoredPending,
startGeneration,
completeGeneration,
};
}