115 lines
3.6 KiB
TypeScript
115 lines
3.6 KiB
TypeScript
import { useState, useEffect, useMemo, useCallback } from 'react'
|
|
import type { ACPClient } from '../acp/client'
|
|
import type { ModelInfo, SessionModelState } from '../acp/types'
|
|
|
|
export interface UseModelsResult {
|
|
/** Whether model selection is supported by the current agent */
|
|
supportsModelSelection: boolean
|
|
/** List of available models */
|
|
availableModels: ModelInfo[]
|
|
/** The currently selected model ID */
|
|
currentModelId: string | null
|
|
/** The currently selected model info */
|
|
currentModel: ModelInfo | null
|
|
/** Set the model for the current session */
|
|
setModel: (modelId: string) => Promise<void>
|
|
/** Whether a model change is in progress */
|
|
isLoading: boolean
|
|
}
|
|
|
|
/**
|
|
* Hook to manage model selection state.
|
|
* Reference: Zed's AcpModelSelector reads from state.available_models and state.current_model_id
|
|
*
|
|
* Uses event-driven updates instead of polling:
|
|
* - setModelStateChangedHandler: called on session create/disconnect
|
|
* - setModelChangedHandler: called when model selection changes
|
|
*/
|
|
export function useModels(client: ACPClient): UseModelsResult {
|
|
const [modelState, setModelState] = useState<SessionModelState | null>(
|
|
client.modelState,
|
|
)
|
|
const [isLoading, setIsLoading] = useState(false)
|
|
|
|
// Subscribe to model state changes (session created/destroyed)
|
|
// This replaces the previous 500ms polling approach
|
|
useEffect(() => {
|
|
// Handler for when model state changes (session created or disconnected)
|
|
const handleModelStateChanged = (state: SessionModelState | null) => {
|
|
setModelState(state)
|
|
// Auto-restore previously selected model when a new session is created
|
|
if (state && state.availableModels.length > 0) {
|
|
const saved = localStorage.getItem('acp_model_id')
|
|
if (
|
|
saved &&
|
|
saved !== state.currentModelId &&
|
|
state.availableModels.some(m => m.modelId === saved)
|
|
) {
|
|
client.setSessionModel(saved).catch(() => {})
|
|
}
|
|
}
|
|
}
|
|
|
|
// Handler for when current model changes within a session
|
|
const handleModelChanged = (modelId: string) => {
|
|
setModelState(prev => {
|
|
if (!prev) return null
|
|
return {
|
|
...prev,
|
|
currentModelId: modelId,
|
|
}
|
|
})
|
|
setIsLoading(false)
|
|
}
|
|
|
|
// Register handlers - setModelStateChangedHandler immediately calls with current state
|
|
client.setModelStateChangedHandler(handleModelStateChanged)
|
|
client.setModelChangedHandler(handleModelChanged)
|
|
|
|
return () => {
|
|
// Clear handlers on unmount
|
|
client.setModelStateChangedHandler(() => {})
|
|
client.setModelChangedHandler(() => {})
|
|
}
|
|
}, [client])
|
|
|
|
const availableModels = useMemo(
|
|
() => modelState?.availableModels ?? [],
|
|
[modelState],
|
|
)
|
|
|
|
const currentModelId = modelState?.currentModelId ?? null
|
|
|
|
const currentModel = useMemo(
|
|
() => availableModels.find(m => m.modelId === currentModelId) ?? null,
|
|
[availableModels, currentModelId],
|
|
)
|
|
|
|
const setModel = useCallback(
|
|
async (modelId: string) => {
|
|
if (!modelState) {
|
|
throw new Error('Model selection not supported')
|
|
}
|
|
setIsLoading(true)
|
|
try {
|
|
await client.setSessionModel(modelId)
|
|
localStorage.setItem('acp_model_id', modelId)
|
|
// The model_changed event will update the state
|
|
} catch (error) {
|
|
setIsLoading(false)
|
|
throw error
|
|
}
|
|
},
|
|
[client, modelState],
|
|
)
|
|
|
|
return {
|
|
supportsModelSelection: modelState !== null && availableModels.length > 0,
|
|
availableModels,
|
|
currentModelId,
|
|
currentModel,
|
|
setModel,
|
|
isLoading,
|
|
}
|
|
}
|