127 lines
3.5 KiB
TypeScript
127 lines
3.5 KiB
TypeScript
import { useEffect, useMemo, useRef, useCallback } from "react";
|
|
import { useShallow } from "zustand/react/shallow";
|
|
import type { GenerationQueueItem } from "../stores/useGenerationStore";
|
|
import { useGenerationStore } from "../stores/useGenerationStore";
|
|
import {
|
|
startBackgroundPolling,
|
|
subscribeToTaskUpdates,
|
|
} from "../services/backgroundTaskRunner";
|
|
|
|
interface UseGenerationTasksOptions {
|
|
sourceView: string;
|
|
autoResume?: boolean;
|
|
}
|
|
|
|
export function useGenerationTasks(options: UseGenerationTasksOptions) {
|
|
const { sourceView, autoResume = true } = options;
|
|
const {
|
|
queue,
|
|
addTask,
|
|
updateTask: updateStoredTask,
|
|
getRunningTasks,
|
|
} = useGenerationStore(useShallow((s) => ({
|
|
queue: s.queue,
|
|
addTask: s.addTask,
|
|
updateTask: s.updateTask,
|
|
getRunningTasks: s.getRunningTasks,
|
|
})));
|
|
const pollingStartedRef = useRef(false);
|
|
|
|
// ── Auto-resume: re-subscribe to running tasks on mount ────
|
|
useEffect(() => {
|
|
if (!autoResume || pollingStartedRef.current) return;
|
|
pollingStartedRef.current = true;
|
|
|
|
const active = getRunningTasks().filter((t) => t.sourceView === sourceView);
|
|
if (active.length > 0) {
|
|
startBackgroundPolling();
|
|
}
|
|
|
|
return () => {
|
|
pollingStartedRef.current = false;
|
|
};
|
|
}, [autoResume, sourceView, getRunningTasks]);
|
|
|
|
// ── Subscribe to live updates ───────────────────────────
|
|
useEffect(() => {
|
|
return subscribeToTaskUpdates((updated) => {
|
|
updateStoredTask(updated.id, updated);
|
|
});
|
|
}, [updateStoredTask]);
|
|
|
|
// ── View-scoped computed lists ──────────────────────────
|
|
const myTasks = useMemo(
|
|
() => queue.filter((t) => t.sourceView === sourceView),
|
|
[queue, sourceView],
|
|
);
|
|
|
|
const activeTasks = useMemo(
|
|
() => myTasks.filter((t) => t.status === "running" || t.status === "pending"),
|
|
[myTasks],
|
|
);
|
|
|
|
const completedTasks = useMemo(
|
|
() => myTasks.filter((t) => t.status === "completed"),
|
|
[myTasks],
|
|
);
|
|
|
|
const failedTasks = useMemo(
|
|
() => myTasks.filter((t) => t.status === "failed"),
|
|
[myTasks],
|
|
);
|
|
|
|
// ── Actions ─────────────────────────────────────────────
|
|
const submitTask = useCallback(
|
|
(task: Omit<GenerationQueueItem, "id" | "createdAt">) => {
|
|
const id = `gen-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
|
|
addTask({ ...task, id, createdAt: Date.now() });
|
|
return id;
|
|
},
|
|
[addTask],
|
|
);
|
|
|
|
const updateTask = useCallback(
|
|
(id: string, patch: Partial<GenerationQueueItem>) => {
|
|
updateStoredTask(id, patch);
|
|
},
|
|
[updateStoredTask],
|
|
);
|
|
|
|
const markCompleted = useCallback(
|
|
(id: string, resultUrl: string) => {
|
|
updateStoredTask(id, { status: "completed", progress: 100, resultUrl });
|
|
},
|
|
[updateStoredTask],
|
|
);
|
|
|
|
const markFailed = useCallback(
|
|
(id: string, error: string) => {
|
|
updateStoredTask(id, { status: "failed", error });
|
|
},
|
|
[updateStoredTask],
|
|
);
|
|
|
|
const retryTask = useCallback(
|
|
(id: string) => {
|
|
const task = queue.find((t) => t.id === id);
|
|
if (task) {
|
|
updateStoredTask(id, { status: "pending", progress: 0, error: null });
|
|
}
|
|
},
|
|
[queue, updateStoredTask],
|
|
);
|
|
|
|
return {
|
|
tasks: myTasks,
|
|
activeTasks,
|
|
completedTasks,
|
|
failedTasks,
|
|
submitTask,
|
|
updateTask,
|
|
markCompleted,
|
|
markFailed,
|
|
retryTask,
|
|
hasActiveTasks: activeTasks.length > 0,
|
|
};
|
|
}
|