diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index d944c7ab7..983d4f5aa 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -19,6 +19,7 @@ import { toExternalBackendProviderType, } from "../external-providers"; import { pickFriendlyContainerName } from "../lib/friendly-names"; +import { tryAdoptServerActiveModel } from "../lib/apply-inference-status-to-store"; import { clampReasoningEffortToLevels, getExternalMaxOutputTokens, @@ -1129,6 +1130,10 @@ async function autoLoadSmallestModel(): Promise<{ loaded: boolean; blockedByTrustRemoteCode: boolean; }> { + if (await tryAdoptServerActiveModel()) { + return { loaded: true, blockedByTrustRemoteCode: false }; + } + const store = useChatRuntimeStore.getState(); const hfToken = store.hfToken || null; const trustRemoteCode = store.params.trustRemoteCode ?? false; @@ -1472,6 +1477,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { } if (!useChatRuntimeStore.getState().params.checkpoint) { + // Prefer a model already loaded by the CLI/API before auto-loading. let loaded: boolean; let blockedByTrustRemoteCode: boolean; try { diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index 4bd36e00d..23ae1d80a 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -24,12 +24,15 @@ import { } from "../api/chat-api"; import { formatEta, formatRate } from "../utils/format-transfer"; import { - CHAT_REASONING_ENABLED_KEY, - loadOptionalBool, - type ReasoningEffort, resolveToolsEnabledOnLoad, useChatRuntimeStore, } from "../stores/chat-runtime-store"; +import { + applyActiveModelStatusToStore, + clampLocalReasoningEffort, + normalizeSpeculativeType, + resolveInferenceCheckpointId, +} from "../lib/apply-inference-status-to-store"; import { mergeBackendRecommendedInference, resolveLoadMaxSeqLength, @@ -211,39 +214,6 @@ function getTrustRemoteCodeRequiredMessage(modelName: string): string { return `${modelName} needs custom code enabled to load. Turn on "Enable custom code" in Chat Settings, then try again.`; } -// Canonicalises any backend/persisted value onto the Speculative Decoding -// dropdown's modes ("auto"/"mtp"/"ngram"/"mtp+ngram"/"off"/null). Mirrors -// backend _canonicalize_spec_mode so legacy persisted values round-trip. -function normalizeSpeculativeType(v: string | null | undefined): string | null { - if (v == null) return null; - const s = String(v).trim().toLowerCase(); - if (!s) return null; - if (s === "auto" || s === "default") return "auto"; - if (s === "off") return "off"; - if (s === "ngram-simple") return "ngram-simple"; - if (s === "mtp" || s === "draft-mtp") return "mtp"; - if (s === "ngram" || s === "ngram-mod") return "ngram"; - if (s === "mtp+ngram") return "mtp+ngram"; - // Comma-chained legacy values (e.g. from older persisted state). - const parts = s.split(",").map((p) => p.trim()).filter(Boolean); - const hasMtp = parts.some((p) => p === "mtp" || p === "draft-mtp"); - const hasNgram = parts.some((p) => p === "ngram" || p === "ngram-mod"); - if (hasMtp && hasNgram) return "mtp+ngram"; - if (hasMtp) return "mtp"; - if (hasNgram) return "ngram"; - // Unknown -> safe fallback to Auto so the dropdown stays controlled. - return "auto"; -} - -type LocalReasoningEffort = Extract; - -function clampLocalReasoningEffort(value: ReasoningEffort): LocalReasoningEffort { - if (value === "low" || value === "medium" || value === "high") { - return value; - } - return "low"; -} - export function useChatModelRuntime() { const params = useChatRuntimeStore((state) => state.params); const models = useChatRuntimeStore((state) => state.models); @@ -328,132 +298,15 @@ export function useChatModelRuntime() { const selectedCheckpoint = useChatRuntimeStore.getState().params.checkpoint; const isExternalSelectionActive = isExternalModelId(selectedCheckpoint); if (statusRes.active_model && !isExternalSelectionActive) { - setCheckpoint(statusRes.active_model, statusRes.gguf_variant); - - // Apply inference defaults on reconnect (page refresh with model already loaded) - if (statusRes.inference) { - const currentParams = useChatRuntimeStore.getState().params; - setParams( - mergeBackendRecommendedInference({ - current: currentParams, - response: statusRes, - modelId: statusRes.active_model, - presetSource: useChatRuntimeStore.getState().activePresetSource, - }), - ); - } - - // Restore reasoning/tools support flags and context length - const hydratingExistingModel = - selectedCheckpoint !== statusRes.active_model || - useChatRuntimeStore.getState().activeGgufVariant !== - (statusRes.gguf_variant ?? null); - const supportsReasoning = statusRes.supports_reasoning ?? false; - const reasoningAlwaysOn = statusRes.reasoning_always_on ?? false; - const reasoningStyle = statusRes.reasoning_style ?? "enable_thinking"; - const reasoningEffortLevels = - reasoningStyle === "reasoning_effort" - ? (["low", "medium", "high"] as const) - : (["low", "medium", "high"] as const); - const supportsPreserveThinking = statusRes.supports_preserve_thinking ?? false; - const supportsTools = statusRes.supports_tools ?? false; - const storedReasoningEnabled = loadOptionalBool( - CHAT_REASONING_ENABLED_KEY, - ); - const currentGgufContextLength = statusRes.is_gguf - ? (statusRes.context_length ?? null) - : null; - const ggufMaxContextLength = statusRes.is_gguf - ? (statusRes.max_context_length ?? null) - : null; - const ggufNativeContextLength = statusRes.is_gguf - ? (statusRes.native_context_length ?? null) - : null; - const currentSpecType = normalizeSpeculativeType( - statusRes.speculative_type, - ); - // Refresh runs on F5 (needs hydration) and right after a load (store - // already set). For user-configurable params, only hydrate when the - // shadow `loaded*` field is null ("not yet hydrated"); otherwise we'd - // clobber what the load path just applied and revert the user. - const prevState = useChatRuntimeStore.getState(); - const clampedReasoningEffort = clampLocalReasoningEffort( - prevState.reasoningEffort, - ); - const nextDefaultChatTemplate = - statusRes.chat_template === undefined - ? prevState.defaultChatTemplate - : statusRes.chat_template; - useChatRuntimeStore.setState({ - supportsReasoning, - reasoningAlwaysOn, - reasoningStyle, - supportsReasoningOff: reasoningStyle !== "reasoning_effort", - reasoningEffortLevels, - reasoningEffort: clampedReasoningEffort, - supportsPreserveThinking, - supportsTools, - // Reset per-turn reasoning flag so: - // 1. non-reasoning models don't inherit a stale off state, and - // 2. local reasoning-effort models (Off hidden via - // supportsReasoningOff=false) don't carry reasoningEnabled=false - // from an external model where Off was selected -- the composer - // would still show "Think: " but the adapter would omit - // the kwarg, so Harmony falls back to its default effort. - reasoningEnabled: supportsReasoning - ? reasoningStyle === "reasoning_effort" - ? true - : useChatRuntimeStore.getState().reasoningEnabled - : true, - ggufContextLength: currentGgufContextLength, - ggufMaxContextLength, - ggufNativeContextLength, - modelRequiresTrustRemoteCode: - statusRes.requires_trust_remote_code ?? false, - defaultChatTemplate: nextDefaultChatTemplate, - loadedIsMultimodal: isMultimodalResponse(statusRes), - specFallbackReason: statusRes.spec_fallback_reason ?? null, - ...(prevState.loadedSpeculativeType === null && { - speculativeType: currentSpecType, - loadedSpeculativeType: currentSpecType, - }), - ...(statusRes.spec_draft_n_max !== undefined && - prevState.loadedSpecDraftNMax === null && - prevState.specDraftNMax === null && { - specDraftNMax: statusRes.spec_draft_n_max ?? null, - loadedSpecDraftNMax: statusRes.spec_draft_n_max ?? null, - }), - ...(statusRes.cache_type_kv !== undefined && - prevState.loadedKvCacheDtype === null && { - kvCacheDtype: statusRes.cache_type_kv, - loadedKvCacheDtype: statusRes.cache_type_kv, - }), - ...(statusRes.chat_template_override !== undefined && - prevState.loadedChatTemplateOverride === null && - prevState.chatTemplateOverride === null && { - chatTemplateOverride: statusRes.chat_template_override, - loadedChatTemplateOverride: statusRes.chat_template_override, - }), - }); - // setModels(listRes...) above used catalog data, which omits audio - // capability. Re-apply live status so attach gates survive a refresh. - syncModelCapabilities(statusRes.active_model, statusRes); - - // Set reasoning default for Qwen3.5/3.6 small models - if ( - supportsReasoning && - hydratingExistingModel && - storedReasoningEnabled === null - ) { - let reasoningDefault = true; - const mid = statusRes.active_model.toLowerCase(); - if (mid.includes("qwen3.5") || mid.includes("qwen3.6")) { - const sizeMatch = mid.match(/(\d+\.?\d*)\s*b/); - if (sizeMatch && parseFloat(sizeMatch[1]) < 9) { - reasoningDefault = false; - } - } - useChatRuntimeStore.setState({ reasoningEnabled: reasoningDefault }); + const checkpointId = resolveInferenceCheckpointId(statusRes); + if (checkpointId) { + setCheckpoint(checkpointId, statusRes.gguf_variant); + applyActiveModelStatusToStore(statusRes, { + previousCheckpoint: selectedCheckpoint, + }); + // setModels(listRes...) above used catalog data, which omits audio + // capability. Re-apply live status so attach gates survive a refresh. + syncModelCapabilities(checkpointId, statusRes); } } else if (!statusRes.active_model && !isExternalSelectionActive) { useChatRuntimeStore.setState({ diff --git a/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts new file mode 100644 index 000000000..c02d304dc --- /dev/null +++ b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts @@ -0,0 +1,236 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { getInferenceStatus } from "../api/chat-api"; +import { mergeBackendRecommendedInference } from "../presets/preset-policy"; +import { + CHAT_REASONING_ENABLED_KEY, + loadOptionalBool, + type ReasoningEffort, + resolveToolsEnabledOnLoad, + useChatRuntimeStore, +} from "../stores/chat-runtime-store"; +import { isMultimodalResponse, type InferenceStatusResponse } from "../types/api"; +import type { ChatModelSummary } from "../types/runtime"; + +type LocalReasoningEffort = Extract; + +// Canonicalises backend / persisted speculative mode values onto the UI modes. +export function normalizeSpeculativeType( + v: string | null | undefined, +): string | null { + if (v == null) return null; + const s = String(v).trim().toLowerCase(); + if (!s) return null; + if (s === "auto" || s === "default") return "auto"; + if (s === "off") return "off"; + if (s === "ngram-simple") return "ngram-simple"; + if (s === "mtp" || s === "draft-mtp") return "mtp"; + if (s === "ngram" || s === "ngram-mod") return "ngram"; + if (s === "mtp+ngram") return "mtp+ngram"; + const parts = s.split(",").map((p) => p.trim()).filter(Boolean); + const hasMtp = parts.some((p) => p === "mtp" || p === "draft-mtp"); + const hasNgram = parts.some((p) => p === "ngram" || p === "ngram-mod"); + if (hasMtp && hasNgram) return "mtp+ngram"; + if (hasMtp) return "mtp"; + if (hasNgram) return "ngram"; + return "auto"; +} + +export function clampLocalReasoningEffort( + value: ReasoningEffort, +): LocalReasoningEffort { + if (value === "low" || value === "medium" || value === "high") { + return value; + } + return "low"; +} + +export function resolveInferenceCheckpointId( + status: InferenceStatusResponse, +): string | null { + if (!status.active_model) return null; + return status.model_identifier ?? status.active_model; +} + +function ensureActiveModelInStoreList( + status: InferenceStatusResponse, + checkpointId: string, +): void { + const store = useChatRuntimeStore.getState(); + if (store.models.some((model) => model.id === checkpointId)) { + return; + } + const summary: ChatModelSummary = { + id: checkpointId, + name: status.active_model ?? checkpointId, + isVision: status.is_vision ?? false, + isLora: false, + isGguf: status.is_gguf ?? false, + isAudio: status.is_audio ?? false, + audioType: status.audio_type ?? null, + hasAudioInput: status.has_audio_input ?? false, + }; + store.setModels([...store.models, summary]); +} + +export type ApplyInferenceStatusOptions = { + previousCheckpoint?: string; +}; + +/** Mirror refresh() hydration so adopted CLI models get reasoning/tools flags. */ +export function applyActiveModelStatusToStore( + status: InferenceStatusResponse, + options: ApplyInferenceStatusOptions = {}, +): void { + const checkpointId = resolveInferenceCheckpointId(status); + if (!checkpointId) return; + + const store = useChatRuntimeStore.getState(); + const previousCheckpoint = + options.previousCheckpoint ?? store.params.checkpoint; + + if (status.inference) { + store.setParams( + mergeBackendRecommendedInference({ + current: store.params, + response: status, + modelId: checkpointId, + presetSource: store.activePresetSource, + }), + ); + } + + const hydratingExistingModel = + previousCheckpoint !== checkpointId || + store.activeGgufVariant !== (status.gguf_variant ?? null); + const supportsReasoning = status.supports_reasoning ?? false; + const reasoningAlwaysOn = status.reasoning_always_on ?? false; + const reasoningStyle = status.reasoning_style ?? "enable_thinking"; + const reasoningEffortLevels = + reasoningStyle === "reasoning_effort" + ? (["low", "medium", "high"] as const) + : (["low", "medium", "high"] as const); + const supportsPreserveThinking = status.supports_preserve_thinking ?? false; + const supportsTools = status.supports_tools ?? false; + const storedReasoningEnabled = loadOptionalBool(CHAT_REASONING_ENABLED_KEY); + const currentGgufContextLength = status.is_gguf + ? (status.context_length ?? null) + : null; + const ggufMaxContextLength = status.is_gguf + ? (status.max_context_length ?? null) + : null; + const ggufNativeContextLength = status.is_gguf + ? (status.native_context_length ?? null) + : null; + const currentSpecType = normalizeSpeculativeType(status.speculative_type); + const prevState = useChatRuntimeStore.getState(); + const clampedReasoningEffort = clampLocalReasoningEffort( + prevState.reasoningEffort, + ); + const nextDefaultChatTemplate = + status.chat_template === undefined + ? prevState.defaultChatTemplate + : status.chat_template; + + useChatRuntimeStore.setState({ + supportsReasoning, + reasoningAlwaysOn, + reasoningStyle, + supportsReasoningOff: reasoningStyle !== "reasoning_effort", + reasoningEffortLevels, + reasoningEffort: clampedReasoningEffort, + supportsPreserveThinking, + supportsTools, + ...resolveToolsEnabledOnLoad(supportsTools), + reasoningEnabled: supportsReasoning + ? reasoningStyle === "reasoning_effort" + ? true + : useChatRuntimeStore.getState().reasoningEnabled + : true, + ggufContextLength: currentGgufContextLength, + ggufMaxContextLength, + ggufNativeContextLength, + modelRequiresTrustRemoteCode: status.requires_trust_remote_code ?? false, + defaultChatTemplate: nextDefaultChatTemplate, + loadedIsMultimodal: isMultimodalResponse(status), + specFallbackReason: status.spec_fallback_reason ?? null, + ...(prevState.loadedSpeculativeType === null && { + speculativeType: currentSpecType, + loadedSpeculativeType: currentSpecType, + }), + ...(status.spec_draft_n_max !== undefined && + prevState.loadedSpecDraftNMax === null && + prevState.specDraftNMax === null && { + specDraftNMax: status.spec_draft_n_max ?? null, + loadedSpecDraftNMax: status.spec_draft_n_max ?? null, + }), + ...(status.cache_type_kv !== undefined && + prevState.loadedKvCacheDtype === null && { + kvCacheDtype: status.cache_type_kv, + loadedKvCacheDtype: status.cache_type_kv, + }), + ...(status.chat_template_override !== undefined && + prevState.loadedChatTemplateOverride === null && + prevState.chatTemplateOverride === null && { + chatTemplateOverride: status.chat_template_override, + loadedChatTemplateOverride: status.chat_template_override, + }), + }); + + ensureActiveModelInStoreList(status, checkpointId); + + if ( + supportsReasoning && + hydratingExistingModel && + storedReasoningEnabled === null + ) { + let reasoningDefault = true; + const mid = checkpointId.toLowerCase(); + if (mid.includes("qwen3.5") || mid.includes("qwen3.6")) { + const sizeMatch = mid.match(/(\d+\.?\d*)\s*b/); + if (sizeMatch && parseFloat(sizeMatch[1]) < 9) { + reasoningDefault = false; + } + } + useChatRuntimeStore.setState({ reasoningEnabled: reasoningDefault }); + } +} + +/** + * Adopt the model already loaded on the inference server (e.g. via + * ``unsloth studio run -m``) into the chat UI checkpoint without + * triggering a new /api/inference/load. + */ +export async function tryAdoptServerActiveModel(): Promise { + const store = useChatRuntimeStore.getState(); + if (store.params.checkpoint) { + return true; + } + + let status: InferenceStatusResponse; + try { + status = await getInferenceStatus(); + } catch { + // Status endpoint unavailable: fall back to the normal auto-load path. + return false; + } + if (!status.active_model) { + return false; + } + + const checkpointId = resolveInferenceCheckpointId(status); + if (!checkpointId) { + return false; + } + + // Re-check after the await: keep a checkpoint the user picked meanwhile. + const previousCheckpoint = + useChatRuntimeStore.getState().params.checkpoint; + if (previousCheckpoint) { + return true; + } + store.setCheckpoint(checkpointId, status.gguf_variant); + applyActiveModelStatusToStore(status, { previousCheckpoint }); + return true; +}