Stop ScaleWorkers dialog overwriting user input on worker-state update

src/components/dashboard/navbar.tsx

@@ -8,17 +8,17 @@ import {
   Card,
   Separator,
   Container,
-  Heading,
   Link,
 } from "@radix-ui/themes";
 import { ChainIdentity } from "../../rakis-core/synthient-chain/db/entities";
 import ChainIdentities from "./chainidentities";
 import ScaleWorkers from "./scaleWorkers";
 import {
-  LLMModelName,
+  AvailableModel,
   LLMWorkerStates,
 } from "../../rakis-core/synthient-chain/llm/types";
 import LiveHelp from "./livehelp";
+import { useMemo } from "react";
 
 const GreenDot = () => (
   <Box
@@ -44,14 +44,23 @@ export default function NavBar({
 }: {
   llmWorkerStates: LLMWorkerStates;
   mySynthientId: string;
-  scaleLLMWorkers: (modelName: LLMModelName, numWorkers: number) => void;
+  scaleLLMWorkers: (modelName: AvailableModel, numWorkers: number) => void;
   chainIdentities: ChainIdentity[];
   addNewChainIdentity: (
     signature: `0x${string}`,
     chain: string,
     signedWithWallet: string
   ) => Promise<void>;
 }) {
+  const workerCount = useMemo(() => {
+    return Object.keys(llmWorkerStates)
+      .map((workerId) => llmWorkerStates[workerId].modelName as AvailableModel)
+      .reduce((acc, cur) => {
+        acc[cur] = (acc[cur] || 0) + 1;
+        return acc;
+      }, {} as Record<AvailableModel, number>);
+  }, [llmWorkerStates]);
+
   return (
     <Flex direction={{ initial: "column", sm: "row" }} justify="center" gap="2">
       <Tooltip content={`${mySynthientId}`}>
@@ -181,13 +190,7 @@ export default function NavBar({
       <Box flexGrow="1"></Box>
       <LiveHelp />
       <ScaleWorkers
-        workerCount={Object.keys(llmWorkerStates)
-          .map((workerId) => llmWorkerStates[workerId].modelName)
-          .reduce((acc, cur) => {
-            acc[cur] ??= 0;
-            acc[cur]++;
-            return acc;
-          }, {} as { [key: string]: number })}
+        workerCount={workerCount}
         scaleLLMWorkers={scaleLLMWorkers}
       />
       <ChainIdentities

src/components/dashboard/scaleWorkers.tsx

@@ -1,7 +1,8 @@
-import { useEffect, useState } from "react";
+import { useEffect, useRef, useState } from "react";
 import {
   availableModels,
   LLMModelName,
+  AvailableModel,
 } from "../../rakis-core/synthient-chain/llm/types";
 import {
   Button as RadixButton,
@@ -27,19 +28,29 @@ export default function ScaleWorkers({
   workerCount,
   scaleLLMWorkers,
 }: {
-  workerCount: { [modelName: string]: number };
-  scaleLLMWorkers: (modelName: LLMModelName, workerCount: number) => void;
+  workerCount: Record<AvailableModel, number>;
+  scaleLLMWorkers: (modelName: AvailableModel, workerCount: number) => void;
 }) {
   const [open, setOpen] = useState(false);
   const [dialogOpen, setDialogOpen] = useState(false);
-  const [selectedModel, setSelectedModel] = useState<LLMModelName | "">(
+  const [selectedModel, setSelectedModel] = useState<AvailableModel | "">(
     "gemma-2b-it-q4f16_1"
   );
   const [scaleCount, setScaleCount] = useState("");
 
+  const workerCountRef = useRef(workerCount);
+
+  useEffect(() => {
+    workerCountRef.current = workerCount;
+  }, [workerCount]);
+
   useEffect(() => {
-    setScaleCount(`${(selectedModel && workerCount[selectedModel] + 1) || 1}`);
-  }, [selectedModel, workerCount]);
+    if (selectedModel) {
+      setScaleCount(`${(workerCountRef.current[selectedModel] ?? 0) + 1}`);
+    } else {
+      setScaleCount("1");
+    }
+  }, [selectedModel]);
 
   function checkScaleWorkers() {
     if (

src/rakis-core/synthient-chain/llm/types.ts

@@ -16,6 +16,8 @@ export const availableModels = [
 
 export type LLMModelName = (typeof availableModels)[number];
 
+export type AvailableModel = typeof availableModels[number];
+
 export type LLMWorkerStates = {
   [workerId: string]: {
     modelName: LLMModelName;