diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index abdaa871c0..af99458349 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -2,7 +2,11 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import type { RememberedLoadSettings } from "@/components/assistant-ui/model-selector/remembered-load-settings"; -import { cancelStagedModelDownload } from "@/features/hub"; +import { + cancelStagedModelDownload, + mirrorHfTokenInto, + useHfTokenStore, +} from "@/features/hub"; import { toast } from "@/lib/toast"; import { create } from "zustand"; import { isExternalModelId, parseExternalModelId } from "../external-providers"; @@ -25,8 +29,6 @@ import { import { useExternalProvidersStore } from "./external-providers-store"; import { PLUS_MENU_PINS_STORAGE_KEY } from "./plus-menu-prefs-store"; -const HF_TOKEN_KEY = "unsloth_hf_token"; -const HF_TOKEN_CHANGED_EVENT = "unsloth:hf-token-changed"; export const CHAT_REASONING_ENABLED_KEY = "unsloth_chat_reasoning_enabled"; export const CHAT_TOOLS_ENABLED_KEY = "unsloth_chat_tools_enabled"; export const CHAT_CODE_TOOLS_ENABLED_KEY = "unsloth_chat_code_tools_enabled"; @@ -495,17 +497,6 @@ export function saveSpeculativeType(value: string | null): void { } } -function notifyHfTokenChanged(value: string): void { - if (!canUseStorage()) return; - try { - window.dispatchEvent( - new CustomEvent(HF_TOKEN_CHANGED_EVENT, { detail: value }), - ); - } catch { - // ignore - } -} - /** A local model staged for a deferred load (see `pendingSelection`). Shape is * a subset of the load hook's `SelectedModelInput`, structurally assignable. */ export type PendingModelSelection = { @@ -1144,7 +1135,7 @@ export const useChatRuntimeStore = create((set, get) => ({ runningByThreadId: {}, cancelByThreadId: {}, autoTitle: false, - hfToken: loadString(HF_TOKEN_KEY, ""), + hfToken: useHfTokenStore.getState().token, modelsError: null, lastModelLoadError: null, activeGgufVariant: null, @@ -1349,11 +1340,7 @@ export const useChatRuntimeStore = create((set, get) => ({ setScalarSettingVersion("autoTitle", autoTitle, state.autoTitle); return { autoTitle }; }), - setHfToken: (hfToken) => { - saveString(HF_TOKEN_KEY, hfToken); - set({ hfToken }); - notifyHfTokenChanged(hfToken); - }, + setHfToken: (hfToken) => useHfTokenStore.getState().setToken(hfToken), setModelsError: (modelsError) => set({ modelsError }), setLastModelLoadError: (lastModelLoadError) => set({ lastModelLoadError }), setCheckpoint: (modelId, ggufVariant) => @@ -1837,6 +1824,12 @@ export const useChatRuntimeStore = create((set, get) => ({ setContextUsage: (contextUsage) => set({ contextUsage }), })); +// Mirror token edits made through the shared store (e.g. Studio's field). +const unsubscribeHfTokenMirror = mirrorHfTokenInto(useChatRuntimeStore); +if (import.meta.hot) { + import.meta.hot.dispose(unsubscribeHfTokenMirror); +} + export function resolveSpeculativeSettingsForLoad({ usePersistedPreference = false, }: { diff --git a/studio/frontend/src/features/hub/index.ts b/studio/frontend/src/features/hub/index.ts index ddcef146b0..3515f6ca76 100644 --- a/studio/frontend/src/features/hub/index.ts +++ b/studio/frontend/src/features/hub/index.ts @@ -2,3 +2,8 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 export { cancelStagedModelDownload } from "./download-manager"; +export { + getHfToken, + mirrorHfTokenInto, + useHfTokenStore, +} from "./stores/hf-token-store"; diff --git a/studio/frontend/src/features/hub/stores/hf-token-store.ts b/studio/frontend/src/features/hub/stores/hf-token-store.ts index b1e2560f02..499ba9f644 100644 --- a/studio/frontend/src/features/hub/stores/hf-token-store.ts +++ b/studio/frontend/src/features/hub/stores/hf-token-store.ts @@ -5,11 +5,9 @@ import { create } from "zustand"; import { bumpInventoryVersion } from "./inventory-events"; const HF_TOKEN_KEY = "unsloth_hf_token"; -const HF_TOKEN_CHANGED_EVENT = "unsloth:hf-token-changed"; const LEGACY_TRAINING_KEY = "unsloth_training_config_v1"; let storageSyncStarted = false; let storageSyncListener: ((event: StorageEvent) => void) | null = null; -let tokenChangedListener: ((event: Event) => void) | null = null; function canUseStorage(): boolean { return typeof window !== "undefined"; @@ -63,10 +61,6 @@ function stopStorageSync(): void { window.removeEventListener("storage", storageSyncListener); storageSyncListener = null; } - if (tokenChangedListener !== null) { - window.removeEventListener(HF_TOKEN_CHANGED_EVENT, tokenChangedListener); - tokenChangedListener = null; - } storageSyncStarted = false; } @@ -100,10 +94,6 @@ export const useHfTokenStore = create((set) => { applyToken(event.newValue ?? "", false); }; window.addEventListener("storage", storageSyncListener); - tokenChangedListener = (event) => { - applyToken((event as CustomEvent).detail ?? "", false); - }; - window.addEventListener(HF_TOKEN_CHANGED_EVENT, tokenChangedListener); } return { @@ -117,6 +107,21 @@ export function getHfToken(): string { return useHfTokenStore.getState().token; } +// Keep a plain zustand store's `hfToken` field in sync with the shared token: +// seed the current value, then mirror later edits. Returns the unsubscribe so +// callers can wire it to HMR disposal. +export function mirrorHfTokenInto(store: { + getState: () => T; + setState: (partial: Partial) => void; +}): () => void { + store.setState({ hfToken: getHfToken() } as Partial); + return useHfTokenStore.subscribe((state) => { + if (store.getState().hfToken !== state.token) { + store.setState({ hfToken: state.token } as Partial); + } + }); +} + // HF's JS client throws on a non-empty token that isn't `hf_...` instead of // browsing anonymously, so treat anything malformed as no token. export function hfApiToken( diff --git a/studio/frontend/src/features/studio/sections/dataset-section.tsx b/studio/frontend/src/features/studio/sections/dataset-section.tsx index bb74b0d061..e7bab1f47f 100644 --- a/studio/frontend/src/features/studio/sections/dataset-section.tsx +++ b/studio/frontend/src/features/studio/sections/dataset-section.tsx @@ -54,7 +54,7 @@ import { // Imported directly from the store module rather than the "@/features/training" // barrel to avoid an import cycle (the barrel re-exports this section's siblings). import { hasSeparateStreamingEvalSplit } from "@/features/training/stores/training-config-store"; -import { useDebouncedValue, useHfTokenValidation } from "@/hooks"; +import { useDebouncedValue } from "@/hooks"; import { translate, useT } from "@/i18n"; import { ChevronDownStandardIcon } from "@/lib/chevron-icons"; import { toast } from "@/lib/toast"; @@ -398,9 +398,6 @@ export function DatasetSection() { enabled: pickerTab === "huggingface", }); - const { error: tokenValidationError, isChecking: isCheckingToken } = - useHfTokenValidation(hfToken); - const hfResultIds = useMemo(() => { const ids = hfResults.map((r) => r.id); if (dataset && !ids.includes(dataset)) { @@ -1005,9 +1002,9 @@ export function DatasetSection() { - {(tokenValidationError ?? hfSearchError) && ( + {hfSearchError && (

- {tokenValidationError ?? hfSearchError} + {hfSearchError} {" — "}

)} - {isCheckingToken && ( -

- {t("studio.dataset.checkingToken")} -

- )} {pickerTab !== activeSourceTab && (

{t("studio.dataset.browsingSource", { diff --git a/studio/frontend/src/features/training/stores/training-config-store.ts b/studio/frontend/src/features/training/stores/training-config-store.ts index 306c02475c..a927f83fd8 100644 --- a/studio/frontend/src/features/training/stores/training-config-store.ts +++ b/studio/frontend/src/features/training/stores/training-config-store.ts @@ -3,6 +3,7 @@ import { CPT_TARGET_MODULES, DEFAULT_HYPERPARAMS, LR_DEFAULT_CPT, LR_DEFAULT_FULL, LR_DEFAULT_LORA, STEPS, TARGET_MODULES } from "@/config/training"; import { authFetch } from "@/features/auth"; +import { getHfToken, mirrorHfTokenInto, useHfTokenStore } from "@/features/hub"; import { isAdapterMethod } from "@/types/training"; import type { DatasetFormat } from "@/types/training"; import type { ModelType, StepNumber, TrainingMethod } from "@/types/training"; @@ -117,7 +118,9 @@ let _datasetFormatAutoForcedByCpt = false; // modelType / isVisionModel / isAudioModel persist so multimodal-only UI // paints right on reload; the model-config fetch still re-derives them. +// hfToken mirrors the shared hf-token-store and is persisted there instead. const NON_PERSISTED_STATE_KEYS: ReadonlySet = new Set([ + "hfToken", "isCheckingVision", "isEmbeddingModel", "isLoadingModelDefaults", @@ -632,8 +635,7 @@ export const useTrainingConfigStore = create()( ), ); }, - setHfToken: (hfToken) => - set({ hfToken: hfToken.trim().replace(/^["']+|["']+$/g, "") }), + setHfToken: (hfToken) => useHfTokenStore.getState().setToken(hfToken), setDatasetSource: (datasetSource) => set({ datasetSource }), selectHfDataset: (dataset) => { _datasetCheckController?.abort(); @@ -923,7 +925,7 @@ export const useTrainingConfigStore = create()( _learningRateManuallySet = false; _yamlLearningRate = undefined; clearCptDatasetFormatTracking(); - set(initialState); + set({ ...initialState, hfToken: getHfToken() }); }, resetToModelDefaults: () => { const { selectedModel } = get(); @@ -947,7 +949,7 @@ export const useTrainingConfigStore = create()( }, { name: "unsloth_training_config_v1", - version: 11, + version: 12, migrate: (persisted, version) => { const s = persisted as Record; if (version < 2 && s.datasetSubset == null && s.datasetConfig != null) { @@ -1000,6 +1002,15 @@ export const useTrainingConfigStore = create()( // own version guard. s.datasetStreaming ??= false; } + if (version < 12) { + // hfToken moved to the shared hf-token-store; seed it once so an + // existing Studio-only token isn't lost. + const legacyToken = typeof s.hfToken === "string" ? s.hfToken.trim() : ""; + if (legacyToken && !getHfToken()) { + useHfTokenStore.getState().setToken(legacyToken); + } + delete s.hfToken; + } return s as unknown as TrainingConfigStore; }, partialize: partializePersistedState, @@ -1022,3 +1033,8 @@ export const useTrainingConfigStore = create()( }, ), ); + +const unsubscribeHfTokenMirror = mirrorHfTokenInto(useTrainingConfigStore); +if (import.meta.hot) { + import.meta.hot.dispose(unsubscribeHfTokenMirror); +}