Workerific

2026-09-20 22:08:18 +08:00 · 2023-02-09 02:58:57 -07:00
parent 593becaa19
commit bc5df84ed1
9 changed files with 290 additions and 262 deletions
@@ -62,6 +62,7 @@
        "embla-carousel-react": "^7.0.3",
        "exifr": "^7.1.3",
        "gray-matter": "^4.0.3",
+        "idb-keyval": "^6.2.0",
        "immer": "^9.0.15",
        "lodash": "^4.17.21",
        "masonic": "^3.7.0",
@@ -6791,6 +6792,14 @@
        "entities": "^3.0.1"
      }
    },
+    "node_modules/idb-keyval": {
+      "version": "6.2.0",
+      "resolved": "https://registry.npmjs.org/idb-keyval/-/idb-keyval-6.2.0.tgz",
+      "integrity": "sha512-uw+MIyQn2jl3+hroD7hF8J7PUviBU7BPKWw4f/ISf32D4LoGu98yHjrzWWJDASu9QNrX10tCJqk9YY0ClWm8Ng==",
+      "dependencies": {
+        "safari-14-idb-fix": "^3.0.0"
+      }
+    },
    "node_modules/ieee754": {
      "version": "1.2.1",
      "resolved": "https://registry.npmjs.org/ieee754/-/ieee754-1.2.1.tgz",
@@ -9286,6 +9295,11 @@
        "node": ">=6"
      }
    },
+    "node_modules/safari-14-idb-fix": {
+      "version": "3.0.0",
+      "resolved": "https://registry.npmjs.org/safari-14-idb-fix/-/safari-14-idb-fix-3.0.0.tgz",
+      "integrity": "sha512-eBNFLob4PMq8JA1dGyFn6G97q3/WzNtFK4RnzT1fnLq+9RyrGknzYiM/9B12MnKAxuj1IXr7UKYtTNtjyKMBog=="
+    },
    "node_modules/safe-buffer": {
      "version": "5.1.2",
      "resolved": "https://registry.npmjs.org/safe-buffer/-/safe-buffer-5.1.2.tgz",
@@ -15589,6 +15603,14 @@
        "entities": "^3.0.1"
      }
    },
+    "idb-keyval": {
+      "version": "6.2.0",
+      "resolved": "https://registry.npmjs.org/idb-keyval/-/idb-keyval-6.2.0.tgz",
+      "integrity": "sha512-uw+MIyQn2jl3+hroD7hF8J7PUviBU7BPKWw4f/ISf32D4LoGu98yHjrzWWJDASu9QNrX10tCJqk9YY0ClWm8Ng==",
+      "requires": {
+        "safari-14-idb-fix": "^3.0.0"
+      }
+    },
    "ieee754": {
      "version": "1.2.1",
      "resolved": "https://registry.npmjs.org/ieee754/-/ieee754-1.2.1.tgz",
@@ -17265,6 +17287,11 @@
        "mri": "^1.1.0"
      }
    },
+    "safari-14-idb-fix": {
+      "version": "3.0.0",
+      "resolved": "https://registry.npmjs.org/safari-14-idb-fix/-/safari-14-idb-fix-3.0.0.tgz",
+      "integrity": "sha512-eBNFLob4PMq8JA1dGyFn6G97q3/WzNtFK4RnzT1fnLq+9RyrGknzYiM/9B12MnKAxuj1IXr7UKYtTNtjyKMBog=="
+    },
    "safe-buffer": {
      "version": "5.1.2",
      "resolved": "https://registry.npmjs.org/safe-buffer/-/safe-buffer-5.1.2.tgz",
@@ -79,6 +79,7 @@
    "embla-carousel-react": "^7.0.3",
    "exifr": "^7.1.3",
    "gray-matter": "^4.0.3",
+    "idb-keyval": "^6.2.0",
    "immer": "^9.0.15",
    "lodash": "^4.17.21",
    "masonic": "^3.7.0",
@@ -44,7 +44,7 @@ export function CivitaiLinkPopover() {
 */

 function LinkDropdown() {
-  const { selectedInstance, connected, socketConnected, instances } = useCivitaiLink();
+  const { instance: selectedInstance, connected, socketConnected, instances } = useCivitaiLink();
  return (
    <Paper style={{ overflow: 'hidden' }}>
      <Stack spacing={0}>
@@ -101,7 +101,7 @@ function GetStarted() {
 }

 function ActivityList() {
-  const { selectedInstance } = useCivitaiLink();
+  const { instance: selectedInstance } = useCivitaiLink();
  const ids = useCivitaiLinkStore((state) => state.ids);
  const { classes } = useActivityListStyles();
  return selectedInstance?.connected ? (
@@ -1,9 +1,6 @@
 /* eslint-disable @typescript-eslint/no-empty-function */
-import { createContext, useCallback, useContext, useEffect, useRef, useState } from 'react';
-import {
-  CivitaiLinkInstance,
-  useGetLinkInstances,
-} from '~/components/CivitaiLink/civitai-link-api';
+import { createContext, useContext, useEffect, useRef, useState } from 'react';
+import { CivitaiLinkInstance } from '~/components/CivitaiLink/civitai-link-api';
 import {
  Command,
  ResponseResourcesList,
@@ -17,55 +14,40 @@ import { showNotification } from '@mantine/notifications';
 import { v4 as uuid } from 'uuid';
 import { immer } from 'zustand/middleware/immer';
 import create from 'zustand';
-import { useLocalStorage } from '@mantine/hooks';
-import { useCurrentUser } from '~/hooks/useCurrentUser';
 import isEqual from 'lodash/isEqual';
-
-// #region types
-type Instance = {
-  key: string | null;
-  connected: boolean; // general connection status - aggregate of `clientsConnected` and `sdConnected`
-  clientsConnected: number; // number of people in room, even though it's probably just you
-  sdConnected: boolean; // if the sd instance is available to connect to
-};
-
-type SelectedInstance = CivitaiLinkInstance & Omit<Instance, 'key'>;
-
-type IncomingMessage =
-  | { type: 'ready' }
-  | { type: 'socketConnection'; payload: boolean }
-  | { type: 'error'; msg: string }
-  | { type: 'message'; msg: string }
-  | { type: 'activitiesUpdate'; payload: ActivitiesResponse[] }
-  | { type: 'resourcesUpdate'; payload: ResponseResourcesList['resources'] }
-  | { type: 'commandComplete'; payload: Response }
-  | { type: 'instance'; payload: Instance }
-  | { type: 'socketConnection'; payload: boolean };
-// #endregion
+import {
+  WorkerOutgoingMessage,
+  WorkerIncomingMessage,
+  Instance,
+} from '~/workers/civitai-link-worker-types';

 // #region context
 type CivitaiLinkState = {
  instances: CivitaiLinkInstance[];
-  selectedInstance?: SelectedInstance;
+  instance?: Instance;
  socketConnected: boolean;
  connected: boolean;
  resources: ResponseResourcesList['resources'];
-  fetchInstances: () => Promise<void>;
-  selectInstance: (instance: { key: string }) => Promise<void>;
-  runCommand: (command: CommandRequest) => Promise<unknown>;
+  createInstance: (id?: number) => Promise<void>;
+  deleteInstance: (id: number) => Promise<void>;
+  renameInstance: (id: number, name: string) => Promise<void>;
+  selectInstance: (id: number) => Promise<void>;
  deselectInstance: () => Promise<void>;
+  runCommand: (command: CommandRequest) => Promise<unknown>;
 };

 const CivitaiLinkCtx = createContext<CivitaiLinkState>({
  instances: [],
-  selectedInstance: undefined,
+  instance: undefined,
  connected: false,
  socketConnected: false,
  resources: [],
-  fetchInstances: async () => {},
+  createInstance: async () => {},
+  deleteInstance: async () => {},
+  renameInstance: async () => {},
  selectInstance: async () => {},
-  runCommand: async () => {},
  deselectInstance: async () => {},
+  runCommand: async () => {},
 } as CivitaiLinkState);
 // #endregion

@@ -110,38 +92,16 @@ const commandPromises: Record<

 export const useCivitaiLink = () => useContext(CivitaiLinkCtx);
 export const CivitaiLinkProvider = ({ children }: { children: React.ReactNode }) => {
-  const user = useCurrentUser();
-  const canUseLink = user != null; // TODO: Briant - Check for subscription...
  const workerRef = useRef<SharedWorker>();
  const workerPromise = useRef<Promise<SharedWorker>>();
  const [socketConnected, setSocketConnected] = useState(false);
-  const [selectedInstanceId, setSelectedInstanceId] = useLocalStorage<number | undefined>({
-    key: 'civitai-link-instance-id',
-  });
-  const { data: instances = [], refetch } = useGetLinkInstances();
-  // const [instances, setInstances] = useState<CivitaiLinkInstance[]>([]);
-  const [selectedInstance, setSelectedInstance] = useState<SelectedInstance>();
+  const [instances, setInstances] = useState<CivitaiLinkInstance[]>([]);
+  const [instance, setInstance] = useState<Instance>();
  const [resources, setResources] = useState<ResponseResourcesList['resources']>([]);
  const [connected, setConnected] = useState(false);
  const setActivities = useCivitaiLinkStore((state) => state.setActivities);

-  console.log({ selectedInstance });
-
-  const updateSelectedInstance = useCallback(
-    (instance: Instance) => {
-      const detectedInstance = instances.find((x) => x.key === instance.key);
-      console.log('FIRFIREA', { instance, detectedInstance, instances });
-      if (detectedInstance) {
-        setSelectedInstanceId(detectedInstance?.id);
-        setSelectedInstance({ ...instance, ...detectedInstance });
-      }
-      setConnected(instance.connected);
-    },
-    [instances, setSelectedInstanceId]
-  );
-
  const getWorker = () => {
-    if (!canUseLink) throw Error('User is not logged in');
    if (workerPromise.current) return workerPromise.current;
    if (workerRef.current) return Promise.resolve(workerRef.current);
    const worker = new SharedWorker(
@@ -157,15 +117,9 @@ export const CivitaiLinkProvider = ({ children }: { children: React.ReactNode })
      showNotification({ message: msg });
    };

-    const handleInstance = (instance: Instance) => {
-      // // const detectedInstance = instances.find((x) => x.key === instance.key);
-      // // console.log('FIRFIREA', { instance, detectedInstance, instances });
-      // // if (detectedInstance) {
-      // setSelectedInstanceId(instance.id);
-      // setSelectedInstance({ ...instance });
-      // // }
-      // setConnected(instance.connected);
-      updateSelectedInstance(instance);
+    const handleInstance = (payload: Instance) => {
+      setInstance(payload);
+      setConnected(payload?.connected ?? false);
    };

    const handleActivities = (activities: ActivitiesResponse[]) => {
@@ -175,15 +129,6 @@ export const CivitaiLinkProvider = ({ children }: { children: React.ReactNode })
        return bDate.getTime() - aDate.getTime();
      });

-      // const removed = sorted.filter((x) => x.type === 'resources:remove');
-      // const added = sorted.filter((x) => x.type === 'resources:add');
-
-      // // TODO - determine how to show that an item has been removed while still being able to show the correct status if removing an item fails
-      // const filtered = added.map((activity) => {
-      //   const index = removed.findIndex((x) => x.resource.hash === activity.resource.hash);
-      //   return index > -1 ? removed[index] : activity;
-      // });
-
      setActivities(sorted);
    };

@@ -200,40 +145,39 @@ export const CivitaiLinkProvider = ({ children }: { children: React.ReactNode })
        resolve(worker);
      };

-      worker.port.onmessage = async function ({ data }: { data: IncomingMessage }) {
+      worker.port.onmessage = async function ({ data }: { data: WorkerOutgoingMessage }) {
        if (data.type === 'ready') handleReady();
        else if (data.type === 'error') handleError(data.msg);
        else if (data.type === 'message') handleMessage(data.msg);
        else if (data.type === 'instance') handleInstance(data.payload);
+        else if (data.type === 'instancesUpdate') setInstances(data.payload);
        else if (data.type === 'resourcesUpdate') setResources(data.payload);
        else if (data.type === 'activitiesUpdate') handleActivities(data.payload);
        else if (data.type === 'commandComplete') handleCommandComplete(data.payload);
        else if (data.type === 'socketConnection') setSocketConnected(data.payload);
-        //TODO.Justin
-        // else if (data.type === 'instances.list') setSocketConnected(data.payload);
      };
    });

    return workerPromise.current;
  };

-  const fetchInstances = useCallback(async () => {
-    await refetch();
-  }, [refetch]);
-
-  const selectInstance = async (instance: { key: string }) => {
+  const boot = async () => {
    const worker = await getWorker();
-    worker.port.postMessage({ type: 'join', key: instance.key });
+    return worker;
  };

-  // TODO.Justin
-  // const createInstance = async (instance: { key: string }) => {
-  //   const worker = await getWorker();
-  //   worker.port.postMessage({ type: 'join', key: instance.key });
-  // };
+  const workerReq = async (req: WorkerIncomingMessage) => {
+    const worker = await getWorker();
+    worker.port.postMessage(req);
+  };
+
+  const selectInstance = (id: number) => workerReq({ type: 'join', id });
+  const deselectInstance = () => workerReq({ type: 'leave' });
+  const createInstance = (id?: number) => workerReq({ type: 'create', id });
+  const deleteInstance = (id: number) => workerReq({ type: 'delete', id });
+  const renameInstance = (id: number, name: string) => workerReq({ type: 'rename', id, name });

  const runCommand = async (command: CommandRequest, timeout = 0) => {
-    const worker = await getWorker();
    const payload = command as Command;
    payload.id = uuid();

@@ -248,37 +192,26 @@ export const CivitaiLinkProvider = ({ children }: { children: React.ReactNode })
      }, timeout);
    });

-    worker.port.postMessage({ type: 'command', payload });
+    await workerReq({ type: 'command', payload });

    return promise;
  };

-  const deselectInstance = async () => {
-    if (!selectedInstance) return;
-    const worker = await getWorker();
-    worker.port.postMessage({ type: 'leave' });
-  };
-
  useEffect(() => {
-    if (!canUseLink) return;
-    fetchInstances();
-  }, [fetchInstances, canUseLink]);
-
-  useEffect(() => {
-    if (!canUseLink || !selectedInstanceId || !instances.length || selectedInstance) return;
-    const storedInstance = instances.find((x) => x.id === selectedInstanceId);
-    if (storedInstance) selectInstance(storedInstance);
-  }, [instances, canUseLink]);
+    boot();
+  }, [boot]);

  return (
    <CivitaiLinkCtx.Provider
      value={{
        instances,
-        selectedInstance,
+        instance,
        connected,
        socketConnected,
        resources,
-        fetchInstances,
+        createInstance,
+        deleteInstance,
+        renameInstance,
        selectInstance,
        deselectInstance,
        runCommand,
@@ -288,12 +221,3 @@ export const CivitaiLinkProvider = ({ children }: { children: React.ReactNode })
    </CivitaiLinkCtx.Provider>
  );
 };
-
-// export function ActualProvider({ children }) {
-//   const { civitaiLink } = useFeatureFlags();
-//   return civitaiLink ? <CivitaiLinkProvider>{children}</CivitaiLinkProvider> : children;
-// }
-
-// export function ConditionalProvider({ children, provider, condition }) {
-//   return condition ? provider({ children }) : children;
-// }
@@ -18,10 +18,6 @@ import {
 import { IconCheck, IconChevronRight, IconCopy } from '@tabler/icons';
 import { useState } from 'react';
 import { z } from 'zod';
-import {
-  useCreateLinkInstance,
-  useUpdateLinkInstance,
-} from '~/components/CivitaiLink/civitai-link-api';
 import { useCivitaiLink } from '~/components/CivitaiLink/CivitaiLinkProvider';
 import { createContextModal } from '~/components/Modals/utils/createContextModal';
 import { Form, InputText, useForm } from '~/libs/form';
@@ -37,7 +33,6 @@ const { openModal, Modal } = createContextModal({
  withCloseButton: false,
  closeOnClickOutside: false,
  Element: ({ context, props }) => {
-    const [key, setKey] = useState<string>();
    const [active, setActive] = useState(0);
    const nextStep = () => setActive((current) => (current < 2 ? current + 1 : current));
    const prevStep = () => setActive((current) => (current > 0 ? current - 1 : current));
@@ -46,35 +41,22 @@ const { openModal, Modal } = createContextModal({
      schema,
    });

-    const { connected, selectInstance, selectedInstance } = useCivitaiLink();
-    const { mutate: createLinkInstance, isLoading: isCreatingLinkInstance } =
-      useCreateLinkInstance();
-    const { mutate: updateLinkInstance, isLoading: isUpdatingLinkInstance } =
-      useUpdateLinkInstance();
+    const {
+      connected,
+      instance: selectedInstance,
+      createInstance,
+      renameInstance,
+    } = useCivitaiLink();

    const handleCreateInstance = () => {
      nextStep();
-      if (!key && !isCreatingLinkInstance) {
-        createLinkInstance(undefined, {
-          onSuccess: (result) => {
-            selectInstance({ key: result.key });
-            setKey(result.key);
-            nextStep();
-          },
-        });
-      }
+      createInstance();
    };

    const handleSubmit = (data: z.infer<typeof schema>) => {
-      if (selectedInstance)
-        updateLinkInstance(
-          { ...data, id: selectedInstance.id },
-          {
-            onSuccess: () => {
-              context.close();
-            },
-          }
-        );
+      if (!selectedInstance?.id) return;
+      renameInstance(selectedInstance.id, data.name);
+      context.close();
    };

    return (
@@ -123,7 +105,6 @@ const { openModal, Modal } = createContextModal({
                  <Button
                    onClick={handleCreateInstance}
                    leftIcon={<IconChevronRight />}
-                    loading={isCreatingLinkInstance}
                  >{`Ok, it's installed`}</Button>
                </Group>
              </Stack>
@@ -136,8 +117,8 @@ const { openModal, Modal } = createContextModal({
                </Text>
                <Text> Paste this code into the Civitai Link settings and save.</Text>
                <Center>
-                  {key ? (
-                    <CopyButton value={key}>
+                  {selectedInstance?.key ? (
+                    <CopyButton value={selectedInstance.key}>
                      {({ copied, copy }) => (
                        <Tooltip label="copy">
                          <Button
@@ -145,7 +126,7 @@ const { openModal, Modal } = createContextModal({
                            onClick={copy}
                            rightIcon={copied ? <IconCheck size={16} /> : <IconCopy size={16} />}
                          >
-                            {!copied ? key : 'Copied'}
+                            {!copied ? selectedInstance.key : 'Copied'}
                          </Button>
                        </Tooltip>
                      )}
@@ -184,9 +165,7 @@ const { openModal, Modal } = createContextModal({
                    label="Name your stable diffusion instance"
                    placeholder="name"
                  />
-                  <Button type="submit" loading={isUpdatingLinkInstance}>
-                    Save
-                  </Button>
+                  <Button type="submit">Save</Button>
                </Stack>
              </Form>
            </Stack>
@@ -1,6 +1,4 @@
-import { useMutation, useQuery } from '@tanstack/react-query';
 import { env } from '~/env/client.mjs';
-import { queryClient } from '~/utils/trpc';

 export type CivitaiLinkInstance = {
  id: number;
@@ -11,25 +9,6 @@ export type CivitaiLinkInstance = {
  createdAt: Date;
 };

-const links = [
-  {
-    id: 2,
-    activated: true,
-    key: '327f5d153ddd160e95e97465ae0aaa4893f74e1477fa4a1215ab9e77aac2d51276385398b4e1c94bcaf9e4ab705e0b245f2d1ff05dda9fc6aed8b729072c1d52',
-    origin: null,
-    name: null,
-    createdAt: new Date('2023-02-01T18:50:51.760Z'),
-  },
-  {
-    id: 3,
-    activated: true,
-    key: '920ceef9e51f59ecaa5b790c32dea22a162188eab624929a2cccabd18d40c75c2987b9f92ffaef104b3d9425fd9b9b3794b0975a1ac433f9d7cf55310e7d3d9a',
-    origin: null,
-    name: null,
-    createdAt: new Date('2023-02-02T21:39:08.522Z'),
-  },
-] as CivitaiLinkInstance[];
-
 const clFetch = async (url: string, options: RequestInit = {}) => {
  if (!url.startsWith('/')) url = '/' + url;
  const response = await fetch(env.NEXT_PUBLIC_CIVITAI_LINK + url, {
@@ -44,36 +23,19 @@ const clFetch = async (url: string, options: RequestInit = {}) => {
 };

 export const getLinkInstances = async () => {
-  // return links; // TODO.civitai-link - comment this out
  return (await clFetch('/api/link')) as CivitaiLinkInstance[];
 };

-const queryKey = ['civitai-link-instances'];
-export const useGetLinkInstances = () => {
-  return useQuery({
-    queryKey,
-    queryFn: getLinkInstances,
-  });
-};
-
-export const createLinkInstance = async (data?: { id?: number }) => {
+export const createLinkInstance = async (id?: number) => {
  return (await clFetch(`/api/link`, {
    method: 'POST',
    headers: {
      'Content-Type': 'application/json',
    },
-    body: data ? JSON.stringify(data) : undefined,
+    body: id ? JSON.stringify({ id }) : undefined,
  })) as { id: number; key: string; instanceCount: number; instanceLimit: number; name: string };
 };

-export const useCreateLinkInstance = () => {
-  return useMutation(createLinkInstance, {
-    onSuccess: async (result) => {
-      await queryClient.invalidateQueries({ queryKey });
-    },
-  });
-};
-
 export const updateLinkInstance = async (data: { id: number; name: string }) => {
  if (!data.id) throw new Error('Missing id');

@@ -86,22 +48,8 @@ export const updateLinkInstance = async (data: { id: number; name: string }) =>
  })) as { id: number; name: string };
 };

-export const useUpdateLinkInstance = () =>
-  useMutation(updateLinkInstance, {
-    onSuccess: async (result) => {
-      await queryClient.invalidateQueries({ queryKey });
-    },
-  });
-
 export const deleteLinkInstance = async (id: number) => {
  return (await clFetch(`/api/link?id=${id}`, {
    method: 'DELETE',
  })) as { success: boolean };
 };
-
-export const useDeleteLinkInstance = () =>
-  useMutation(deleteLinkInstance, {
-    onSuccess: async (result) => {
-      await queryClient.invalidateQueries({ queryKey });
-    },
-  });
@@ -3,6 +3,7 @@ import { Socket } from 'socket.io-client';
 export type ClientType = 'client' | 'sd';

 export interface ServerToClientEvents {
+  kicked: () => void;
  roomPresence: (msg: { client: number; sd: number }) => void;
  upgradeKey: (msg: { key: string }) => void;
  error: (msg: { msg: string }) => void;
@@ -0,0 +1,35 @@
+import { CivitaiLinkInstance } from '~/components/CivitaiLink/civitai-link-api';
+import {
+  ActivitiesResponse,
+  Command,
+  ResponseResourcesList,
+  Response,
+} from '~/components/CivitaiLink/shared-types';
+
+export type Instance = {
+  id: number | null;
+  name: string | null;
+  key: string | null;
+  connected: boolean; // general connection status - aggregate of `clientsConnected` and `sdConnected`
+  clientsConnected: number; // number of people in room, even though it's probably just you
+  sdConnected: boolean; // if the sd instance is available to connect to
+};
+
+export type WorkerOutgoingMessage =
+  | { type: 'ready' }
+  | { type: 'socketConnection'; payload: boolean }
+  | { type: 'error'; msg: string }
+  | { type: 'message'; msg: string }
+  | { type: 'activitiesUpdate'; payload: ActivitiesResponse[] }
+  | { type: 'instancesUpdate'; payload: CivitaiLinkInstance[] }
+  | { type: 'resourcesUpdate'; payload: ResponseResourcesList['resources'] }
+  | { type: 'commandComplete'; payload: Response }
+  | { type: 'instance'; payload: Instance };
+
+export type WorkerIncomingMessage =
+  | { type: 'create'; id?: number }
+  | { type: 'delete'; id: number }
+  | { type: 'rename'; id: number; name: string }
+  | { type: 'join'; id: number }
+  | { type: 'leave' }
+  | { type: 'command'; payload: Command };
@@ -6,9 +6,23 @@ import {
  Response,
  ResponseResourcesList,
  ResponseStatus,
+  ActivitiesResponse,
 } from '~/components/CivitaiLink/shared-types';
 import { env } from '~/env/client.mjs';
 import { v4 as uuid } from 'uuid';
+import {
+  CivitaiLinkInstance,
+  createLinkInstance,
+  deleteLinkInstance,
+  getLinkInstances,
+  updateLinkInstance,
+} from '~/components/CivitaiLink/civitai-link-api';
+import {
+  WorkerIncomingMessage,
+  Instance,
+  WorkerOutgoingMessage,
+} from '~/workers/civitai-link-worker-types';
+import { get, set, del } from 'idb-keyval';

 // --------------------------------
 // Types
@@ -18,18 +32,6 @@ interface SharedWorkerGlobalScope {
 }
 const _self: SharedWorkerGlobalScope = self as any;

-type IncomingMessage =
-  | { type: 'join'; key: string }
-  | { type: 'leave' }
-  | { type: 'command'; payload: Command };
-
-type Instance = {
-  key: string | null;
-  connected: boolean;
-  sdConnected: boolean;
-  clientsConnected: number;
-};
-
 // --------------------------------
 // Setup Socket
 // --------------------------------
@@ -50,32 +52,38 @@ const sendCommand = (payload: Omit<Command, 'id' | 'createdAt'>) => {
 // --------------------------------
 // Setup shared state
 // --------------------------------
-let initialized = false;
-const instance: Instance = {
+const defaultInstance: Instance = {
+  id: null,
  key: null,
+  name: null,
  connected: false,
  sdConnected: false,
  clientsConnected: 0,
 };
+let initialized: number | null = null;
+let instance: Instance = { ...defaultInstance };
+let instances: CivitaiLinkInstance[] = [];
 let resources: ResponseResourcesList['resources'] = [];
-let activities: Response[] = [];
+let activities: ActivitiesResponse[] = [];

 // Shared value events
 const sharedCallbacks = {
  resources: [] as (() => void)[],
  activities: [] as (() => void)[],
  instance: [] as (() => void)[],
+  instances: [] as (() => void)[],
  error: [] as ((msg: string) => void)[],
  message: [] as ((msg: string) => void)[],
  completion: [] as ((response: Response) => void)[],
  socketConnection: [] as ((connected: boolean) => void)[],
 };
-const onUpdate = (type: 'resources' | 'activities' | 'instance', cb: () => void) => {
+const onUpdate = (type: UpdateSharedValueProps['type'], cb: () => void) => {
  sharedCallbacks[type].push(cb);
 };
 type UpdateSharedValueProps =
  | { type: 'resources'; value: ResponseResourcesList['resources'] }
-  | { type: 'activities'; value: Response[] }
+  | { type: 'activities'; value: ActivitiesResponse[] }
+  | { type: 'instances'; value: CivitaiLinkInstance[] }
  | { type: 'instance'; value: Partial<Instance> };
 const updateSharedValue = ({ type, value }: UpdateSharedValueProps) => {
  console.log('updateSharedValue', { type, value });
@@ -86,9 +94,11 @@ const updateSharedValue = ({ type, value }: UpdateSharedValueProps) => {
    activities = value;
    sharedCallbacks.activities.forEach((cb) => cb());
  } else if (type === 'instance') {
-    if (value.key) instance.key = value.key;
-    if (value.connected) instance.connected = value.connected;
+    instance = { ...instance, ...value };
    sharedCallbacks.instance.forEach((cb) => cb());
+  } else if (type === 'instances') {
+    instances = value;
+    sharedCallbacks.instances.forEach((cb) => cb());
  }
 };

@@ -128,15 +138,39 @@ const emitMessage = (msg: string) => {
  sharedCallbacks.message.forEach((cb) => cb(msg));
 };

+// Storage
+const storageKey = 'cl-id';
+const storeInstanceId = async () => {
+  if (!instance.id) {
+    console.log(`${storageKey}: clear`);
+    await del(storageKey);
+  } else {
+    console.log(`${storageKey}: ${instance.id}`);
+    await set(storageKey, instance.id.toString());
+  }
+};
+
+const getStoredInstanceId = async () => {
+  const id = await get(storageKey);
+  console.log(`${storageKey}: ${id}`);
+  if (!id) return null;
+  return Number(id);
+};
+
+// --------------------------------
+// Handle Instance Events
+// --------------------------------
+
 // --------------------------------
 // Handle Socket Events
 // --------------------------------
 socket.on('connect', () => {
+  socket.emit('iam', { type: 'client' });
  emitSocketConnection(true);
-  if (instance.key) {
-    // rejoin if key is set
-    handleJoin(instance.key);
-    handleInitialization();
+  if (instance.id) {
+    // rejoin if id is set
+    initialized = null;
+    handleJoin(instance.id);
  }
 });
 socket.on('disconnect', () => {
@@ -152,16 +186,16 @@ socket.on('commandStatus', (payload: Response) => {
    return;
  }

-  let value: Response[] = [];
+  let value: ActivitiesResponse[] = [];
  if (payload.type === 'activities:list' || payload.type === 'activities:clear') {
-    value = payload.activities;
+    value = payload.activities as ActivitiesResponse[];
  } else {
    let found = false;
    for (const activity of activities) {
      if (activity.id !== payload.id) value.push(activity);
      else {
        found = true;
-        value.push(payload);
+        value.push(payload as ActivitiesResponse);

        // emit completion if status changed to a completed status
        const activityCompleted =
@@ -169,50 +203,74 @@ socket.on('commandStatus', (payload: Response) => {
        if (activityCompleted) emitCompletion(payload);
      }
    }
-    if (!found) value.push(payload);
+    if (!found) value.push(payload as ActivitiesResponse);
  }

  updateSharedValue({ type: 'activities', value });
 });

+socket.on('upgradeKey', ({ key }) => {
+  const match = instances.find((x) => x.id === instance.id);
+  if (match) match.key = key;
+
+  updateSharedValue({ type: 'instance', value: { key } });
+});
+
+socket.on('kicked', () => {
+  updateSharedValue({ type: 'instance', value: defaultInstance });
+  storeInstanceId();
+});
+
 socket.on('error', ({ msg }) => {
  emitError(msg);
 });

 socket.on('roomPresence', ({ client, sd }) => {
+  console.log('roomPresence', { client, sd });
  if (!instance.sdConnected && sd > 0) emitMessage('Stable Diffusion service connected');
  else if (instance.sdConnected && sd === 0) emitMessage('Stable Diffusion service disconnected');

+  const connected = sd > 0 && client > 0;
+  if (connected && !instance.connected) handleInitialization();
+  else if (!connected && instance.connected) initialized = null;
  updateSharedValue({
    type: 'instance',
-    value: { sdConnected: sd > 0, clientsConnected: client, connected: sd > 0 && client > 0 },
+    value: { sdConnected: sd > 0, clientsConnected: client, connected },
  });
 });

 // --------------------------------
 // Handle Incoming Messages
 // --------------------------------
-const handleJoin = (key: string) => {
-  if (instance.key === key && instance.connected) return;
+const handleJoin = (id: number) => {
+  if (instance.id === id && instance.connected) return;
+
+  const targetInstance = instances.find((i) => i.id === id);
+  if (!targetInstance) {
+    emitError('Could not find instance');
+    return;
+  }
+  const { key, name } = targetInstance;

  if (!socket.connected) {
    socket.connect();
    socket.emit('iam', { type: 'client' });
  }

-  socket.emit('join', key, ({ success, msg }) => {
-    updateSharedValue({ type: 'instance', value: { key } });
+  socket.emit('join', targetInstance.key, ({ success, msg }) => {
    if (!success && msg) emitError(msg);
+    else updateSharedValue({ type: 'instance', value: { id, key, name } });
  });
 };

 const handleLeave = () => {
-  if (!instance.key) return;
+  if (!instance.id) return;
  socket.emit('leave');
  updateSharedValue({
    type: 'instance',
-    value: { key: null, sdConnected: false, clientsConnected: 0, connected: false },
+    value: defaultInstance,
  });
+  storeInstanceId();
 };

 const handleCommand = (payload: Command) => {
@@ -223,46 +281,101 @@ const handleCommand = (payload: Command) => {
  socket.emit('command', { ...payload, createdAt: new Date() });
 };

+const handleLoadInstances = async () => {
+  try {
+    const result = await getLinkInstances();
+    updateSharedValue({ type: 'instances', value: result });
+  } catch (err: any) {
+    emitError(`Error loading instances: ${err.message}`);
+  }
+};
+
+const handleRename = async (id: number, name: string) => {
+  try {
+    await updateLinkInstance({ id, name });
+    if (instance.id === id) updateSharedValue({ type: 'instance', value: { name } });
+    await handleLoadInstances();
+  } catch (err: any) {
+    emitError(`Error renaming instance: ${err.message}`);
+  }
+};
+
+const handleDelete = async (id: number) => {
+  try {
+    if (instance.id === id) handleLeave();
+    await deleteLinkInstance(id);
+    await handleLoadInstances();
+  } catch (err: any) {
+    emitError(`Error deleting instance: ${err.message}`);
+  }
+};
+
+const handleCreate = async (id?: number) => {
+  try {
+    const result = await createLinkInstance(id);
+    await handleLoadInstances();
+    handleJoin(result.id);
+  } catch (err: any) {
+    emitError(`Error creating instance: ${err.message}`);
+  }
+};
+
 const handleInitialization = () => {
+  if (!instance.id || initialized === instance.id) return;
+
  sendCommand({ type: 'activities:list' });
  sendCommand({ type: 'resources:list' });
-  initialized = true;
+  initialized = instance.id;
+  console.log(`Initialized instance: ${instance.id}`);
+  storeInstanceId();
 };

 // --------------------------------
 // Bootstrap Worker
 // --------------------------------
-const start = (port: MessagePort) => {
+const start = async (port: MessagePort) => {
  if (!port.postMessage) return;

-  onError((msg) => port.postMessage({ type: 'error', msg }));
-  onMessage((msg) => port.postMessage({ type: 'message', msg }));
-  onCompletion((payload) => port.postMessage({ type: 'commandComplete', payload }));
-  port.postMessage({ type: 'instance', payload: instance });
+  const portReq = (req: WorkerOutgoingMessage) => port.postMessage(req);
+
+  onError((msg) => portReq({ type: 'error', msg }));
+  onMessage((msg) => portReq({ type: 'message', msg }));
+  onCompletion((payload) => portReq({ type: 'commandComplete', payload }));
+  portReq({ type: 'instance', payload: instance });
  onUpdate('instance', () => {
-    if (instance.connected && !initialized) handleInitialization();
-    port.postMessage({ type: 'instance', payload: instance });
+    portReq({ type: 'instance', payload: instance });
  });
-  port.postMessage({ type: 'resourcesUpdate', payload: resources });
+  portReq({ type: 'resourcesUpdate', payload: resources });
  onUpdate('resources', () => {
-    port.postMessage({ type: 'resourcesUpdate', payload: resources });
+    portReq({ type: 'resourcesUpdate', payload: resources });
  });
-  port.postMessage({ type: 'activitiesUpdate', payload: activities });
+  portReq({ type: 'activitiesUpdate', payload: activities });
  onUpdate('activities', () => {
-    port.postMessage({ type: 'activitiesUpdate', payload: activities });
+    portReq({ type: 'activitiesUpdate', payload: activities });
  });
-  port.postMessage({ type: 'socketConnection', payload: socket.connected });
+  portReq({ type: 'instancesUpdate', payload: instances });
+  onUpdate('instances', () => {
+    portReq({ type: 'instancesUpdate', payload: instances });
+  });
+  portReq({ type: 'socketConnection', payload: socket.connected });
  onSocketConnection((connected) => {
-    port.postMessage({ type: 'socketConnection', payload: connected });
+    portReq({ type: 'socketConnection', payload: connected });
  });

-  port.onmessage = ({ data }: { data: IncomingMessage }) => {
-    if (data.type === 'join') handleJoin(data.key);
+  port.onmessage = ({ data }: { data: WorkerIncomingMessage }) => {
+    if (data.type === 'join') handleJoin(data.id);
+    else if (data.type === 'create') handleCreate(data.id);
+    else if (data.type === 'delete') handleDelete(data.id);
+    else if (data.type === 'rename') handleRename(data.id, data.name);
    else if (data.type === 'leave') handleLeave();
    else if (data.type === 'command') handleCommand(data.payload);
  };

-  port.postMessage({ type: 'ready' });
+  handleLoadInstances().finally(async () => {
+    portReq({ type: 'ready' });
+    const storedInstanceId = await getStoredInstanceId();
+    if (storedInstanceId) handleJoin(Number(storedInstanceId));
+  });
 };

 _self.onconnect = (e) => {