import { useCallback } from "react"; import { useMutation, useQueryClient, type QueryClient, type QueryKey, } from "@tanstack/react-query"; import { useTranslation } from "react-i18next"; import type { CreateScheduleOptions, DaemonClient, UpdateScheduleOptions, } from "@getpaseo/client/internal/daemon-client"; import type { ScheduleSummary } from "@getpaseo/protocol/schedule/types"; import { schedulesQueryBaseKey } from "@/hooks/use-schedules"; import type { FetchAggregatedSchedulesResult } from "@/schedules/aggregated-schedules"; import { useSessionStore } from "@/stores/session-store"; export type CreateScheduleInput = Omit; export type UpdateScheduleInput = Omit; export interface UseScheduleMutationsResult { createSchedule: (input: CreateScheduleInput) => Promise; updateSchedule: (input: UpdateScheduleInput) => Promise; pauseSchedule: (id: string) => Promise; resumeSchedule: (id: string) => Promise; deleteSchedule: (id: string) => Promise; runScheduleNow: (id: string) => Promise; isCreating: boolean; isUpdating: boolean; isPausing: boolean; isResuming: boolean; isDeleting: boolean; isRunningNow: boolean; } interface ScheduleListSnapshot { previous: Array<[QueryKey, FetchAggregatedSchedulesResult | undefined]>; } function requireClient(serverId: string, unavailableMessage: string): DaemonClient { const client = useSessionStore.getState().sessions[serverId]?.client ?? null; if (!client) { throw new Error(unavailableMessage); } return client; } function snapshotSchedules(queryClient: QueryClient): ScheduleListSnapshot { return { previous: queryClient.getQueriesData({ queryKey: schedulesQueryBaseKey, }), }; } function restoreSchedules(queryClient: QueryClient, snapshot: ScheduleListSnapshot): void { for (const [queryKey, previous] of snapshot.previous) { queryClient.setQueryData(queryKey, previous); } } function updateScheduleSections( queryClient: QueryClient, updateSection: ( section: FetchAggregatedSchedulesResult["sections"][number], ) => FetchAggregatedSchedulesResult["sections"][number], ): void { queryClient.setQueriesData( { queryKey: schedulesQueryBaseKey }, (current) => { if (!current) { return current; } return { sections: current.sections.map(updateSection) }; }, ); } function optimisticallySetStatus( queryClient: QueryClient, serverId: string, id: string, status: ScheduleSummary["status"], ): void { const pausedAt = status === "paused" ? new Date().toISOString() : null; updateScheduleSections(queryClient, (section) => section.serverId === serverId ? { ...section, schedules: section.schedules.map((schedule) => schedule.id === id ? { ...schedule, status, pausedAt } : schedule, ), } : section, ); } function optimisticallyRemove(queryClient: QueryClient, serverId: string, id: string): void { updateScheduleSections(queryClient, (section) => section.serverId === serverId ? { ...section, schedules: section.schedules.filter((schedule) => schedule.id !== id), } : section, ); } export function useScheduleMutations({ serverId, }: { serverId: string; }): UseScheduleMutationsResult { const queryClient = useQueryClient(); const { t } = useTranslation(); const invalidate = useCallback(() => { void queryClient.invalidateQueries({ queryKey: schedulesQueryBaseKey }); }, [queryClient]); const createMutation = useMutation({ mutationFn: async (input: CreateScheduleInput): Promise => { const client = requireClient(serverId, t("common.errors.daemonClientUnavailable")); const payload = await client.scheduleCreate(input); if (payload.error) { throw new Error(payload.error); } }, onSettled: invalidate, }); const updateMutation = useMutation({ mutationFn: async (input: UpdateScheduleInput): Promise => { const client = requireClient(serverId, t("common.errors.daemonClientUnavailable")); const payload = await client.scheduleUpdate(input); if (payload.error) { throw new Error(payload.error); } }, onSettled: invalidate, }); const pauseMutation = useMutation({ mutationFn: async (id: string): Promise => { const client = requireClient(serverId, t("common.errors.daemonClientUnavailable")); const payload = await client.schedulePause({ id }); if (payload.error) { throw new Error(payload.error); } }, onMutate: async (id): Promise => { await queryClient.cancelQueries({ queryKey: schedulesQueryBaseKey }); const snapshot = snapshotSchedules(queryClient); optimisticallySetStatus(queryClient, serverId, id, "paused"); return snapshot; }, onError: (_error, _id, context) => { if (context) { restoreSchedules(queryClient, context); } }, onSettled: invalidate, }); const resumeMutation = useMutation({ mutationFn: async (id: string): Promise => { const client = requireClient(serverId, t("common.errors.daemonClientUnavailable")); const payload = await client.scheduleResume({ id }); if (payload.error) { throw new Error(payload.error); } }, onMutate: async (id): Promise => { await queryClient.cancelQueries({ queryKey: schedulesQueryBaseKey }); const snapshot = snapshotSchedules(queryClient); optimisticallySetStatus(queryClient, serverId, id, "active"); return snapshot; }, onError: (_error, _id, context) => { if (context) { restoreSchedules(queryClient, context); } }, onSettled: invalidate, }); const deleteMutation = useMutation({ mutationFn: async (id: string): Promise => { const client = requireClient(serverId, t("common.errors.daemonClientUnavailable")); const payload = await client.scheduleDelete({ id }); if (payload.error) { throw new Error(payload.error); } }, onMutate: async (id): Promise => { await queryClient.cancelQueries({ queryKey: schedulesQueryBaseKey }); const snapshot = snapshotSchedules(queryClient); optimisticallyRemove(queryClient, serverId, id); return snapshot; }, onError: (_error, _id, context) => { if (context) { restoreSchedules(queryClient, context); } }, onSettled: invalidate, }); const runNowMutation = useMutation({ mutationFn: async (id: string): Promise => { const client = requireClient(serverId, t("common.errors.daemonClientUnavailable")); const payload = await client.scheduleRunOnce({ id }); if (payload.error) { throw new Error(payload.error); } }, onSettled: invalidate, }); const createSchedule = useCallback( async (input: CreateScheduleInput): Promise => { await createMutation.mutateAsync(input); }, [createMutation], ); const updateSchedule = useCallback( async (input: UpdateScheduleInput): Promise => { await updateMutation.mutateAsync(input); }, [updateMutation], ); const pauseSchedule = useCallback( async (id: string): Promise => { await pauseMutation.mutateAsync(id); }, [pauseMutation], ); const resumeSchedule = useCallback( async (id: string): Promise => { await resumeMutation.mutateAsync(id); }, [resumeMutation], ); const deleteSchedule = useCallback( async (id: string): Promise => { await deleteMutation.mutateAsync(id); }, [deleteMutation], ); const runScheduleNow = useCallback( async (id: string): Promise => { await runNowMutation.mutateAsync(id); }, [runNowMutation], ); return { createSchedule, updateSchedule, pauseSchedule, resumeSchedule, deleteSchedule, runScheduleNow, isCreating: createMutation.isPending, isUpdating: updateMutation.isPending, isPausing: pauseMutation.isPending, isResuming: resumeMutation.isPending, isDeleting: deleteMutation.isPending, isRunningNow: runNowMutation.isPending, }; }