diff --git a/src/chatbox.tsx b/src/chatbox.tsx
index bf33469..27286b3 100644
--- a/src/chatbox.tsx
+++ b/src/chatbox.tsx
@@ -243,8 +243,10 @@ export default function ChatBOX(props: {
     });
 
     // manually calculate token length
-    chatStore.totalTokens += client.calculate_token_length(inputMsg.trim());
-    client.total_tokens += client.calculate_token_length(inputMsg.trim());
+    chatStore.totalTokens +=
+      calculate_token_length(inputMsg.trim()) + calculate_token_length(images);
+    client.total_tokens = chatStore.totalTokens;
+
     setChatStore({ ...chatStore });
     setInputMsg("");
     setImages([]);
@@ -916,7 +918,9 @@ export default function ChatBOX(props: {
               chatStore.history.push({
                 role: "assistant",
                 content: inputMsg,
-                token: calculate_token_length(inputMsg),
+                token:
+                  calculate_token_length(inputMsg) +
+                  calculate_token_length(images),
                 hide: false,
                 example: false,
               });
diff --git a/src/chatgpt.ts b/src/chatgpt.ts
index 74db715..1f604d7 100644
--- a/src/chatgpt.ts
+++ b/src/chatgpt.ts
@@ -71,7 +71,7 @@ export function calculate_token_length(
       tokens += m.image_url?.detail === "high" ? 65 * 4 : 65;
     }
   }
-  return 0;
+  return tokens;
 }
 
 class Chat {
@@ -267,11 +267,11 @@ class Chat {
     return this._fetch(true);
   }
 
-  calculate_token_length(content: string): number {
+  calculate_token_length(content: string | MessageDetail[]): number {
     return calculate_token_length(content);
   }
 
-  user(...messages: string[]) {
+  user(...messages: (string | MessageDetail[])[]) {
     for (const msg of messages) {
       this.messages.push({ role: "user", content: msg });
       this.total_tokens += this.calculate_token_length(msg);
@@ -279,7 +279,7 @@ class Chat {
     }
   }
 
-  assistant(...messages: string[]) {
+  assistant(...messages: (string | MessageDetail[])[]) {
     for (const msg of messages) {
       this.messages.push({ role: "assistant", content: msg });
       this.total_tokens += this.calculate_token_length(msg);
diff --git a/src/message.tsx b/src/message.tsx
index 10866ef..3fdb679 100644
--- a/src/message.tsx
+++ b/src/message.tsx
@@ -57,6 +57,7 @@ function EditMessage(props: EditMessageProps) {
                         if (typeof chat.content === "string") return;
                         chat.content[index].text = event.target.value;
                         chat.token = calculate_token_length(chat.content);
+                        console.log("calculated token length", chat.token);
                         setChatStore({ ...chatStore });
                       }}
                       onKeyPress={(event: any) => {
@@ -129,6 +130,7 @@ function EditMessage(props: EditMessageProps) {
                           const obj = chat.content[index].image_url;
                           if (obj === undefined) return;
                           obj.detail = obj.detail === "high" ? "low" : "high";
+                          chat.token = calculate_token_length(chat.content);
                           setChatStore({ ...chatStore });
                         }}
                       >