Remove duplicate helper for block timing

Author: TheodoreSpeaks

apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts

@@ -11,7 +11,7 @@ import type {
   PersistedFileAttachment,
   PersistedMessage,
 } from '@/lib/copilot/chat/persisted-message'
-import { normalizeMessage } from '@/lib/copilot/chat/persisted-message'
+import { normalizeMessage, withBlockTiming } from '@/lib/copilot/chat/persisted-message'
 import { resolveStreamToolOutcome } from '@/lib/copilot/chat/stream-tool-outcome'
 import { MOTHERSHIP_CHAT_API_PATH, STREAM_STORAGE_KEY } from '@/lib/copilot/constants'
 import type {
@@ -699,15 +699,6 @@ function parseStreamBatchResponse(value: unknown): StreamBatchResponse {
   }
 }
 
-function withBlockTiming(
-  persisted: Record<string, unknown>,
-  src: ContentBlock
-): Record<string, unknown> {
-  if (typeof src.timestamp === 'number') persisted.timestamp = src.timestamp
-  if (typeof src.endedAt === 'number') persisted.endedAt = src.endedAt
-  return persisted
-}
-
 function toRawPersistedContentBlock(block: ContentBlock): Record<string, unknown> | null {
   const persisted = toRawPersistedContentBlockBody(block)
   return persisted ? withBlockTiming(persisted, block) : null

apps/sim/lib/copilot/chat/display-message.ts

@@ -16,6 +16,7 @@ import {
   ToolCallStatus,
 } from '@/app/workspace/[workspaceId]/home/types'
 import type { PersistedContentBlock, PersistedMessage } from './persisted-message'
+import { withBlockTiming } from './persisted-message'
 
 const STATE_TO_STATUS: Record<string, ToolCallStatus> = {
   [MothershipStreamV1ToolOutcome.success]: ToolCallStatus.success,
@@ -43,12 +44,6 @@ function toToolCallInfo(block: PersistedContentBlock): ToolCallInfo | undefined
   }
 }
 
-function withBlockTiming(block: ContentBlock, src: PersistedContentBlock): ContentBlock {
-  if (typeof src.timestamp === 'number') block.timestamp = src.timestamp
-  if (typeof src.endedAt === 'number') block.endedAt = src.endedAt
-  return block
-}
-
 function toDisplayBlock(block: PersistedContentBlock): ContentBlock | undefined {
   const displayed = toDisplayBlockBody(block)
   return displayed ? withBlockTiming(displayed, block) : undefined

apps/sim/lib/copilot/chat/persisted-message.ts

@@ -87,13 +87,20 @@ function resolveToolState(block: ContentBlock): PersistedToolState {
   return tc.status as PersistedToolState
 }
 
-function withBlockTiming<T extends PersistedContentBlock>(
-  persisted: T,
+/**
+ * Copy `timestamp` / `endedAt` from a source object onto a target object.
+ * Shared by every block mapper (persist, display, snapshot) so the timing
+ * metadata that drives the `Thought for Ns` chip survives the full
+ * persist → normalize → display round-trip — and one rule lives in one place.
+ */
+export function withBlockTiming<T>(
+  target: T,
   src: { timestamp?: number; endedAt?: number }
 ): T {
-  if (typeof src.timestamp === 'number') persisted.timestamp = src.timestamp
-  if (typeof src.endedAt === 'number') persisted.endedAt = src.endedAt
-  return persisted
+  const writable = target as { timestamp?: number; endedAt?: number }
+  if (typeof src.timestamp === 'number') writable.timestamp = src.timestamp
+  if (typeof src.endedAt === 'number') writable.endedAt = src.endedAt
+  return target
 }
 
 function mapContentBlock(block: ContentBlock): PersistedContentBlock {