diff --git a/package.json b/package.json
index e5a9abb8..cda5cfd2 100644
--- a/package.json
+++ b/package.json
@@ -1,7 +1,7 @@
 {
   "name": "koala-client",
   "private": true,
-  "version": "2.0.8",
+  "version": "2.1.0",
   "type": "module",
   "homepage": "./",
   "main": "electron/index.cjs",
diff --git a/src/App.tsx b/src/App.tsx
index f73dfe57..3e079d55 100644
--- a/src/App.tsx
+++ b/src/App.tsx
@@ -22,6 +22,10 @@ function App() {
   const setHideSideMenu = useStore((state) => state.setHideSideMenu);
   const hideSideMenu = useStore((state) => state.hideSideMenu);
   const bottomMessageRef = useStore((state) => state.bottomMessageRef);
+  const setApiAuth = useStore((state) => state.setApiAuth);
+  const apiAuth = useStore((state) => state.apiAuth);
+  const apiKey = useStore((state) => state.apiKey);
+  const setApiKey = useStore((state) => state.setApiKey);
 
   const initialiseNewChat = useInitialiseNewChat();
   const addChat = useAddChat();
@@ -30,6 +34,13 @@ function App() {
   const copyCodeBlock = useCopyCodeBlock();
   const { handleSubmit } = useSubmit();
 
+  if (apiKey && !apiAuth[0].apiKey) {
+    const old = apiAuth;
+    old[0].apiKey = apiKey;
+    setApiAuth(old);
+    setApiKey('');
+  }
+
   const handleGenerate = () => {
     if (useStore.getState().generating) return;
     const updatedChats: ChatInterface[] = JSON.parse(
diff --git a/src/api/api.ts b/src/api/api.ts
index 28897c02..e139fbab 100644
--- a/src/api/api.ts
+++ b/src/api/api.ts
@@ -19,7 +19,7 @@ export const getChatCompletion = async (
   if (isAzureEndpoint(endpoint) && apiKey) {
     headers['api-key'] = apiKey;
 
-    const modelName = modelDef.name;
+    const modelName = modelDef.model;
 
     const apiVersion = '2023-03-15-preview';
 
@@ -69,7 +69,7 @@ export const getChatCompletionStream = async (
   if (isAzureEndpoint(endpoint) && apiKey) {
     headers['api-key'] = apiKey;
 
-    const modelName = modelDef.name;
+    const modelName = modelDef.model;
 
     const apiVersion = '2023-03-15-preview';
 
diff --git a/src/components/ApiMenu/ApiMenu.tsx b/src/components/ApiMenu/ApiMenu.tsx
index 0a11cfae..e76d6571 100644
--- a/src/components/ApiMenu/ApiMenu.tsx
+++ b/src/components/ApiMenu/ApiMenu.tsx
@@ -249,10 +249,12 @@ const ApiMenu = ({
                       }}
                     >
                       <span className='inline-block truncate max-w-full'>
-                        {_apiAuth[modelDef.endpoint].endpoint.replace(
-                          /^https?:\/\//,
-                          ''
-                        )}
+                        {_apiAuth[modelDef.endpoint]
+                          ? _apiAuth[modelDef.endpoint].endpoint.replace(
+                              /^https?:\/\//,
+                              ''
+                            )
+                          : 'Endpoint Undefined'}
                       </span>
 
                       <DownChevronArrow className='absolute right-0 mr-1 flex items-center' />
diff --git a/src/components/ConfigMenu/ModelSelect.tsx b/src/components/ConfigMenu/ModelSelect.tsx
index be44e446..fd72e456 100644
--- a/src/components/ConfigMenu/ModelSelect.tsx
+++ b/src/components/ConfigMenu/ModelSelect.tsx
@@ -17,6 +17,10 @@ export const ModelSelect = ({
   const [dropDown, setDropDown, dropDownRef] = useHideOnOutsideClick();
   const modelDefs = useStore((state: StoreState) => state.modelDefs);
 
+  if (typeof _model !== 'number') {
+    _setModel(0);
+  }
+
   return (
     <div className='mb-4'>
       <button
@@ -25,7 +29,7 @@ export const ModelSelect = ({
         onClick={() => setDropDown((prev) => !prev)}
         aria-label='model'
       >
-        {modelDefs[model]?.name}
+        {modelDefs[model]?.name || modelDefs[model].model}
         <DownChevronArrow />
       </button>
       <div
diff --git a/src/hooks/useSubmit.ts b/src/hooks/useSubmit.ts
index 466ea9a6..b9c28b50 100644
--- a/src/hooks/useSubmit.ts
+++ b/src/hooks/useSubmit.ts
@@ -104,9 +104,7 @@ const useSubmit = () => {
       if (!apiKey || apiKey.length === 0) {
         // official endpoint
         if (apiEndpoint === officialAPIEndpoint) {
-          const error = new Error(t('noApiKeyWarning') as string);
-          setError(error.message);
-          throw error;
+          throw new Error(t('noApiKeyWarning') as string);
         }
 
         // other endpoints
@@ -223,8 +221,7 @@ const useSubmit = () => {
 
         // update tokens used for generating title
         if (countTotalTokens) {
-          const model = config.model_selection;
-          updateTotalTokenUsed(model, [message], {
+          updateTotalTokenUsed(0, [message], {
             role: 'assistant',
             content: title,
           });
@@ -232,6 +229,7 @@ const useSubmit = () => {
       }
     } catch (e: unknown) {
       setError((e as Error).message);
+      setGenerating(false);
       throw e;
     }
     setGenerating(false);
diff --git a/src/store/auth-slice.ts b/src/store/auth-slice.ts
index 3eb95a2b..ae159bb7 100644
--- a/src/store/auth-slice.ts
+++ b/src/store/auth-slice.ts
@@ -4,6 +4,7 @@ import { ModelDefinition } from '@type/chat';
 import { StoreSlice } from './store';
 
 export interface AuthSlice {
+  apiKey?: string;
   firstVisit: boolean;
   apiAuth: EndpointAuth[];
   modelDefs: ModelDefinition[];
@@ -44,12 +45,6 @@ export const createAuthSlice: StoreSlice<AuthSlice> = (set) => ({
       apiKey: apiKey,
     }));
   },
-  setApiEndpoint: (apiEndpoint: string) => {
-    set((prev: AuthSlice) => ({
-      ...prev,
-      apiEndpoint: apiEndpoint,
-    }));
-  },
   setFirstVisit: (firstVisit: boolean) => {
     set((prev: AuthSlice) => ({
       ...prev,
diff --git a/src/utils/messageUtils.ts b/src/utils/messageUtils.ts
index 117cfbbe..bc24a154 100644
--- a/src/utils/messageUtils.ts
+++ b/src/utils/messageUtils.ts
@@ -132,7 +132,7 @@ export const useUpdateTotalTokenUsed = () => {
       completionMessage: MessageInterface
     ) => {
       const updatedTotalTokenUsed = JSON.parse(JSON.stringify(totalTokenUsed));
-      const modelName = modelDefs[model].name;
+      const modelName = modelDefs[model].model;
 
       const newPromptTokens = countTokens(promptMessages, modelName);
       const newCompletionTokens = countTokens([completionMessage], modelName);