From 5b9c1fc668c48118c7099460a3c230d8ce11ea7b Mon Sep 17 00:00:00 2001
From: Louis <louis@jan.ai>
Date: Thu, 19 Dec 2024 12:10:14 +0700
Subject: [PATCH] fix: token speed disappear on complete

---
 web/containers/Providers/ModelHandler.tsx | 3 +++
 web/utils/datetime.test.ts                | 8 ++------
 2 files changed, 5 insertions(+), 6 deletions(-)

diff --git a/web/containers/Providers/ModelHandler.tsx b/web/containers/Providers/ModelHandler.tsx
index d72db4e8e..8e5c3ca6b 100644
--- a/web/containers/Providers/ModelHandler.tsx
+++ b/web/containers/Providers/ModelHandler.tsx
@@ -264,6 +264,9 @@ export default function ModelHandler() {
         if (updatedMessage) {
           deleteMessage(message.id)
           addNewMessage(updatedMessage)
+          setTokenSpeed((prev) =>
+            prev ? { ...prev, message: updatedMessage.id } : undefined
+          )
         }
       })()
 
diff --git a/web/utils/datetime.test.ts b/web/utils/datetime.test.ts
index e7e00ada9..605f98d24 100644
--- a/web/utils/datetime.test.ts
+++ b/web/utils/datetime.test.ts
@@ -4,11 +4,7 @@ import { isToday } from './datetime'
 test("should return only time for today's timestamp", () => {
   const today = new Date()
   const timestamp = today.getTime()
-  const expectedTime = `${today.toLocaleDateString(undefined, {
-    day: '2-digit',
-    month: 'short',
-    year: 'numeric',
-  })}, ${today.toLocaleTimeString(undefined, {
+  const expectedTime = `${today.toLocaleTimeString(undefined, {
     hour: '2-digit',
     minute: '2-digit',
     second: '2-digit',
@@ -24,5 +20,5 @@ test('should return N/A for undefined timestamp', () => {
 test("should return true for today's timestamp", () => {
   const today = new Date()
   const timestamp = today.setHours(0, 0, 0, 0)
-  expect(isToday(timestamp)).toBe(true)
+  expect(isToday(timestamp / 1000)).toBe(true)
 })