fix(copilot): refresh custom block metadata after hydration

Author: j15z

apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/agent-group/tool-call-item.test.tsx

@@ -1,17 +1,24 @@
 /**
- * @vitest-environment node
+ * @vitest-environment jsdom
  */
-import type { ReactNode, SVGProps } from 'react'
+import { act, type ReactNode, type SVGProps } from 'react'
+import { createRoot, type Root } from 'react-dom/client'
 import { renderToStaticMarkup } from 'react-dom/server'
-import { describe, expect, it, vi } from 'vitest'
-import { getBlockByToolName } from '@/blocks/registry'
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+import { notifyBlockOverlayChanged } from '@/blocks/custom/client-overlay'
+import { getBlock, getBlockByToolName } from '@/blocks/registry'
 import { ToolCallItem } from './tool-call-item'
 
 vi.mock('@/components/ui', () => ({
   ShimmerText: ({ children }: { children: ReactNode }) => <span>{children}</span>,
 }))
 
 describe('ToolCallItem', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    ;(globalThis as { IS_REACT_ACT_ENVIRONMENT?: boolean }).IS_REACT_ACT_ENVIRONMENT = true
+  })
+
   it.each(['executing', 'success', 'error', 'cancelled'] as const)(
     'renders the %s tool row without an icon',
     (status) => {
@@ -115,4 +122,34 @@ describe('ToolCallItem', () => {
     expect(markup).toContain('<svg')
     expect(markup).toContain('Read recent emails')
   })
+
+  it('refreshes the read icon when custom blocks hydrate after mount', () => {
+    vi.mocked(getBlock).mockReturnValue(undefined)
+    const container = document.createElement('div')
+    const root: Root = createRoot(container)
+
+    act(() => {
+      root.render(
+        <ToolCallItem
+          toolName='read'
+          displayTitle='Read Custom block invoice parser'
+          status='success'
+          params={{
+            path: 'organization/custom-blocks/custom_block_invoice_parser.json',
+          }}
+        />
+      )
+    })
+    expect(container.querySelector('[data-testid="custom-block-icon"]')).toBeNull()
+
+    vi.mocked(getBlock).mockReturnValue({
+      type: 'custom_block_invoice_parser',
+      name: 'Invoice Parser',
+      icon: (props: SVGProps<SVGSVGElement>) => <svg {...props} data-testid='custom-block-icon' />,
+    } as ReturnType<typeof getBlock>)
+    act(() => notifyBlockOverlayChanged())
+
+    expect(container.querySelector('[data-testid="custom-block-icon"]')).not.toBeNull()
+    act(() => root.unmount())
+  })
 })

apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/agent-group/tool-call-item.tsx

@@ -13,6 +13,7 @@ import { RETIRED_BROWSER_REQUEST_TAKEOVER_ID } from '@/lib/copilot/tools/retired
 import { extractStreamingStringArgument } from '@/lib/copilot/tools/streaming-args'
 import { getToolStatusDisplayTitle, getWaitCountdownTitle } from '@/lib/copilot/tools/tool-display'
 import { BrandIcon } from '@/blocks/brand-icon'
+import { useCustomBlockOverlayVersion } from '@/blocks/custom/client-overlay'
 import { getBlockByToolName } from '@/blocks/registry'
 import type { ToolCallData, ToolCallStatus } from '../../../../types'
 import { resolveToolDisplayState } from '../../utils'
@@ -122,11 +123,12 @@ export function ToolCallItem({
   toolCallId,
   startedAt,
 }: ToolCallItemProps) {
-  const readBlock = useMemo(() => {
-    if (toolName !== ReadTool.id) return undefined
-    const path = params?.path
-    return typeof path === 'string' ? getReadTargetBlock(path) : undefined
-  }, [toolName, params])
+  useCustomBlockOverlayVersion()
+  const readPath = params?.path
+  const readBlock =
+    toolName === ReadTool.id && typeof readPath === 'string'
+      ? getReadTargetBlock(readPath)
+      : undefined
 
   // Like read's VFS-target resolution above, the gateway uses its exact
   // discovered toolId only as a deterministic registry lookup. This renders

apps/sim/app/workspace/[workspaceId]/home/components/message-content/message-content-overlay.test.tsx

@@ -0,0 +1,84 @@
+/**
+ * @vitest-environment jsdom
+ */
+import { act } from 'react'
+import { createRoot, type Root } from 'react-dom/client'
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+
+const { mockGetBlock } = vi.hoisted(() => ({
+  mockGetBlock: vi.fn(),
+}))
+
+vi.mock('@/blocks/registry', () => ({
+  getBlock: mockGetBlock,
+  getBlockByToolName: vi.fn(),
+  getLatestBlock: vi.fn(),
+}))
+
+vi.mock('@/lib/auth/auth-client', () => ({
+  useSession: vi.fn(() => ({ data: null, isPending: false })),
+}))
+
+interface MockAgentGroupItem {
+  type: string
+  data?: { id: string; displayTitle: string }
+}
+
+vi.mock('./components', () => ({
+  AgentGroup: ({ items }: { items: MockAgentGroupItem[] }) => (
+    <div>
+      {items.map((item) => item.data && <span key={item.data.id}>{item.data.displayTitle}</span>)}
+    </div>
+  ),
+  ChatContent: () => null,
+  CircleStop: () => null,
+  Options: () => null,
+  PendingTagIndicator: () => null,
+}))
+
+import type { ContentBlock } from '@/app/workspace/[workspaceId]/home/types'
+import { notifyBlockOverlayChanged } from '@/blocks/custom/client-overlay'
+import { MessageContent } from './message-content'
+
+describe('MessageContent custom-block hydration', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    ;(globalThis as { IS_REACT_ACT_ENVIRONMENT?: boolean }).IS_REACT_ACT_ENVIRONMENT = true
+  })
+
+  it('refreshes a read title when the custom-block registry hydrates after mount', () => {
+    mockGetBlock.mockReturnValue(undefined)
+    const blocks: ContentBlock[] = [
+      {
+        type: 'tool_call',
+        toolCall: {
+          id: 'read-custom-block',
+          name: 'read',
+          status: 'success',
+          params: {
+            path: 'organization/custom-blocks/custom_block_invoice_parser.json',
+          },
+        },
+        timestamp: 1,
+      },
+    ]
+    const container = document.createElement('div')
+    const root: Root = createRoot(container)
+
+    act(() => {
+      root.render(<MessageContent blocks={blocks} fallbackContent='' isStreaming={false} />)
+    })
+    expect(container.textContent).toContain('Read Custom block invoice parser')
+
+    mockGetBlock.mockReturnValue({
+      type: 'custom_block_invoice_parser',
+      name: 'Invoice Parser',
+      icon: () => null,
+    })
+    act(() => notifyBlockOverlayChanged())
+
+    expect(container.textContent).toContain('Read Invoice Parser')
+    expect(container.textContent).not.toContain('Read Custom block invoice parser')
+    act(() => root.unmount())
+  })
+})

apps/sim/app/workspace/[workspaceId]/home/components/message-content/message-content.tsx

@@ -22,6 +22,7 @@ import {
 } from '@/lib/copilot/tools/tool-display'
 import { useChatSurface } from '@/app/workspace/[workspaceId]/home/components/chat-surface-context'
 import type { CredentialSubmissionPayload } from '@/app/workspace/[workspaceId]/home/components/message-content/components/special-tags'
+import { useCustomBlockOverlayVersion } from '@/blocks/custom/client-overlay'
 import type { ContentBlock, OptionItem, ToolCallData } from '../../types'
 import { SUBAGENT_LABELS } from '../../types'
 import type { AgentGroupItem } from './components'
@@ -851,7 +852,11 @@ function MessageContentInner({
   actions,
 }: MessageContentProps) {
   const { onWorkspaceResourceSelect } = useChatSurface()
-  const parsed = useMemo(() => (blocks.length > 0 ? parseBlocks(blocks) : []), [blocks])
+  const blockOverlayVersion = useCustomBlockOverlayVersion()
+  const parsed = useMemo(
+    () => (blocks.length > 0 ? parseBlocks(blocks) : []),
+    [blocks, blockOverlayVersion]
+  )
 
   const [trailingRevealing, setTrailingRevealing] = useState(false)
   const handleTrailingRevealChange = useCallback((revealing: boolean) => {