import { invoke } from '@tauri-apps/api/core'; import { getCurrentWebviewWindow } from '@tauri-apps/api/webviewWindow'; import { resolvedModelName } from '@yaakapp/app/lib/resolvedModelName'; import { AnyModel, ModelPayload } from '../bindings/gen_models'; import { modelStoreDataAtom } from './atoms'; import { ExtractModel, JotaiStore, ModelStoreData } from './types'; import { newStoreData } from './util'; let _store: JotaiStore | null = null; export function initModelStore(store: JotaiStore) { _store = store; getCurrentWebviewWindow() .listen('model_write', ({ payload }) => { if (shouldIgnoreModel(payload)) return; mustStore().set(modelStoreDataAtom, (prev: ModelStoreData) => { if (payload.change.type === 'upsert') { return { ...prev, [payload.model.model]: { ...prev[payload.model.model], [payload.model.id]: payload.model, }, }; } else { const modelData = { ...prev[payload.model.model] }; delete modelData[payload.model.id]; return { ...prev, [payload.model.model]: modelData }; } }); }) .catch(console.error); } function mustStore(): JotaiStore { if (_store == null) { throw new Error('Model store was not initialized'); } return _store; } let _activeWorkspaceId: string | null = null; export async function changeModelStoreWorkspace(workspaceId: string | null) { console.log('Syncing models with new workspace', workspaceId); const workspaceModelsStr = await invoke('models_workspace_models', { workspaceId, // NOTE: if no workspace id provided, it will just fetch global models }); const workspaceModels = JSON.parse(workspaceModelsStr) as AnyModel[]; const data = newStoreData(); for (const model of workspaceModels) { data[model.model][model.id] = model; } mustStore().set(modelStoreDataAtom, data); console.log('Synced model store with workspace', workspaceId, data); _activeWorkspaceId = workspaceId; } export function listModels>( modelType: M | ReadonlyArray, ): T[] { let data = mustStore().get(modelStoreDataAtom); const types: ReadonlyArray = Array.isArray(modelType) ? modelType : [modelType]; return types.flatMap((t) => Object.values(data[t]) as T[]); } export function getModel>( modelType: M | ReadonlyArray, id: string, ): T | null { let data = mustStore().get(modelStoreDataAtom); const types: ReadonlyArray = Array.isArray(modelType) ? modelType : [modelType]; for (const t of types) { let v = data[t][id]; if (v?.model === t) return v as T; } return null; } export function getAnyModel( id: string, ): AnyModel | null { let data = mustStore().get(modelStoreDataAtom); for (const t of Object.keys(data)) { let v = (data as any)[t]?.[id]; if (v?.model === t) return v; } return null; } export function patchModelById>( model: M, id: string, patch: Partial | ((prev: T) => T), ): Promise { let prev = getModel(model, id); if (prev == null) { throw new Error(`Failed to get model to patch id=${id} model=${model}`); } const newModel = typeof patch === 'function' ? patch(prev) : { ...prev, ...patch }; return updateModel(newModel); } export async function patchModel>( base: Pick, patch: Partial, ): Promise { return patchModelById(base.model, base.id, patch); } export async function updateModel>( model: T, ): Promise { return invoke('models_upsert', { model }); } export async function deleteModelById< M extends AnyModel['model'], T extends ExtractModel, >(modelType: M | M[], id: string) { let model = getModel(modelType, id); await deleteModel(model); } export async function deleteModel>( model: T | null, ) { if (model == null) { throw new Error('Failed to delete null model'); } await invoke('models_delete', { model }); } export function duplicateModel>( model: T | null, ) { if (model == null) { throw new Error('Failed to duplicate null model'); } // If the model has a name, try to duplicate it with a name that doesn't conflict let name = 'name' in model ? resolvedModelName(model) : undefined; if (name != null) { const existingModels = listModels(model.model); for (let i = 0; i < 100; i++) { const hasConflict = existingModels.some((m) => { if ('folderId' in m && 'folderId' in model && model.folderId !== m.folderId) { return false; } else if (resolvedModelName(m) !== name) { return false; } return true; }); if (!hasConflict) { break; } // Name conflict. Try another one const m: RegExpMatchArray | null = name.match(/ Copy( (?\d+))?$/); if (m != null && m.groups?.n == null) { name = name.substring(0, m.index) + ' Copy 2'; } else if (m != null && m.groups?.n != null) { name = name.substring(0, m.index) + ` Copy ${parseInt(m.groups.n) + 1}`; } else { name = `${name} Copy`; } } } return invoke('models_duplicate', { model: { ...model, name } }); } export async function createGlobalModel>( patch: Partial & Pick, ): Promise { return invoke('models_upsert', { model: patch }); } export async function createWorkspaceModel>( patch: Partial & Pick, ): Promise { return invoke('models_upsert', { model: patch }); } export function replaceModelsInStore< M extends AnyModel['model'], T extends Extract, >(model: M, models: T[]) { const newModels: Record = {}; for (const model of models) { newModels[model.id] = model; } mustStore().set(modelStoreDataAtom, (prev: ModelStoreData) => { return { ...prev, [model]: newModels, }; }); } export function mergeModelsInStore< M extends AnyModel['model'], T extends Extract, >(model: M, models: T[], filter?: (model: T) => boolean) { mustStore().set(modelStoreDataAtom, (prev: ModelStoreData) => { const existingModels = { ...prev[model] } as Record; // Merge in new models first for (const m of models) { existingModels[m.id] = m; } // Then filter out unwanted models if (filter) { for (const [id, m] of Object.entries(existingModels)) { if (!filter(m)) { delete existingModels[id]; } } } return { ...prev, [model]: existingModels, }; }); } function shouldIgnoreModel({ model, updateSource }: ModelPayload) { // Never ignore updates from non-user sources if (updateSource.type !== 'window') { return false; } // Never ignore same-window updates if (updateSource.label === getCurrentWebviewWindow().label) { return false; } // Only sync models that belong to this workspace, if a workspace ID is present if ('workspaceId' in model && model.workspaceId !== _activeWorkspaceId) { return true; } if (model.model === 'key_value' && model.namespace === 'no_sync') { return true; } return false; }