yaak-mountain-loop/crates/yaak-models/guest-js/store.ts

import { invoke } from "@tauri-apps/api/core";
import { getCurrentWebviewWindow } from "@tauri-apps/api/webviewWindow";
import { resolvedModelName } from "@yaakapp/yaak-client/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<ModelPayload>("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<string>("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<
  M extends AnyModel["model"],
  T extends ExtractModel<AnyModel, M>,
>(modelType: M | ReadonlyArray<M>): T[] {
  let data = mustStore().get(modelStoreDataAtom);
  const types: ReadonlyArray<M> = Array.isArray(modelType)
    ? modelType
    : [modelType];
  return types.flatMap((t) => Object.values(data[t]) as T[]);
}

export function getModel<
  M extends AnyModel["model"],
  T extends ExtractModel<AnyModel, M>,
>(modelType: M | ReadonlyArray<M>, id: string): T | null {
  let data = mustStore().get(modelStoreDataAtom);
  const types: ReadonlyArray<M> = 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<
  M extends AnyModel["model"],
  T extends ExtractModel<AnyModel, M>,
>(model: M, id: string, patch: Partial<T> | ((prev: T) => T)): Promise<string> {
  let prev = getModel<M, T>(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<
  M extends AnyModel["model"],
  T extends ExtractModel<AnyModel, M>,
>(base: Pick<T, "id" | "model">, patch: Partial<T>): Promise<string> {
  return patchModelById<M, T>(base.model, base.id, patch);
}

export async function updateModel<
  M extends AnyModel["model"],
  T extends ExtractModel<AnyModel, M>,
>(model: T): Promise<string> {
  return invoke<string>("models_upsert", { model });
}

export async function deleteModelById<
  M extends AnyModel["model"],
  T extends ExtractModel<AnyModel, M>,
>(modelType: M | M[], id: string) {
  let model = getModel<M, T>(modelType, id);
  await deleteModel(model);
}

export async function deleteModel<
  M extends AnyModel["model"],
  T extends ExtractModel<AnyModel, M>,
>(model: T | null) {
  if (model == null) {
    throw new Error("Failed to delete null model");
  }
  await invoke<string>("models_delete", { model });
}

export function duplicateModel<
  M extends AnyModel["model"],
  T extends ExtractModel<AnyModel, M>,
>(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( (?<n>\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<string>("models_duplicate", { model: { ...model, name } });
}

export async function createGlobalModel<
  T extends Exclude<AnyModel, { workspaceId: string }>,
>(patch: Partial<T> & Pick<T, "model">): Promise<string> {
  return invoke<string>("models_upsert", { model: patch });
}

export async function createWorkspaceModel<
  T extends Extract<AnyModel, { workspaceId: string }>,
>(patch: Partial<T> & Pick<T, "model" | "workspaceId">): Promise<string> {
  return invoke<string>("models_upsert", { model: patch });
}

export function replaceModelsInStore<
  M extends AnyModel["model"],
  T extends Extract<AnyModel, { model: M }>,
>(model: M, models: T[]) {
  const newModels: Record<string, T> = {};
  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<AnyModel, { model: M }>,
>(model: M, models: T[], filter?: (model: T) => boolean) {
  mustStore().set(modelStoreDataAtom, (prev: ModelStoreData) => {
    const existingModels = { ...prev[model] } as Record<string, T>;

    // 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;
}