claude-code/src/hooks/usePromptsFromClaudeInChrome.tsx

import type { ContentBlockParam } from '@anthropic-ai/sdk/resources/messages.mjs'
import { useEffect, useRef } from 'react'
import { logError } from 'src/utils/log.js'
import { z } from 'zod/v4'
import { callIdeRpc } from '../services/mcp/client.js'
import type {
  ConnectedMCPServer,
  MCPServerConnection,
} from '../services/mcp/types.js'
import type { PermissionMode } from '../types/permissions.js'
import {
  CLAUDE_IN_CHROME_MCP_SERVER_NAME,
  isTrackedClaudeInChromeTabId,
} from '../utils/claudeInChrome/common.js'
import { lazySchema } from '../utils/lazySchema.js'
import { enqueuePendingNotification } from '../utils/messageQueueManager.js'

// Schema for the prompt notification from Chrome extension (JSON-RPC 2.0 format)
const ClaudeInChromePromptNotificationSchema = lazySchema(() =>
  z.object({
    method: z.literal('notifications/message'),
    params: z.object({
      prompt: z.string(),
      image: z
        .object({
          type: z.literal('base64'),
          media_type: z.enum([
            'image/jpeg',
            'image/png',
            'image/gif',
            'image/webp',
          ]),
          data: z.string(),
        })
        .optional(),
      tabId: z.number().optional(),
    }),
  }),
)

/**
 * A hook that listens for prompt notifications from the Claude for Chrome extension,
 * enqueues them as user prompts, and syncs permission mode changes to the extension.
 */
export function usePromptsFromClaudeInChrome(
  mcpClients: MCPServerConnection[],
  toolPermissionMode: PermissionMode,
): void {
  const mcpClientRef = useRef<ConnectedMCPServer | undefined>(undefined)

  useEffect(() => {
    if (process.env.USER_TYPE !== 'ant') {
      return
    }

    const mcpClient = findChromeClient(mcpClients)
    if (mcpClientRef.current !== mcpClient) {
      mcpClientRef.current = mcpClient
    }

    if (mcpClient) {
      mcpClient.client.setNotificationHandler(
        ClaudeInChromePromptNotificationSchema(),
        notification => {
          if (mcpClientRef.current !== mcpClient) {
            return
          }
          const { tabId, prompt, image } = notification.params

          // Process notifications from tabs we're tracking since notifications are broadcasted
          if (
            typeof tabId !== 'number' ||
            !isTrackedClaudeInChromeTabId(tabId)
          ) {
            return
          }

          try {
            // Build content blocks if there's an image, otherwise just use the prompt string
            if (image) {
              const contentBlocks: ContentBlockParam[] = [
                { type: 'text', text: prompt },
                {
                  type: 'image',
                  source: {
                    type: image.type,
                    media_type: image.media_type,
                    data: image.data,
                  },
                },
              ]
              enqueuePendingNotification({
                value: contentBlocks,
                mode: 'prompt',
              })
            } else {
              enqueuePendingNotification({ value: prompt, mode: 'prompt' })
            }
          } catch (error) {
            logError(error as Error)
          }
        },
      )
    }
  }, [mcpClients])

  // Sync permission mode with Chrome extension whenever it changes
  useEffect(() => {
    const chromeClient = findChromeClient(mcpClients)
    if (!chromeClient) return

    const chromeMode =
      toolPermissionMode === 'bypassPermissions'
        ? 'skip_all_permission_checks'
        : 'ask'

    void callIdeRpc('set_permission_mode', { mode: chromeMode }, chromeClient)
  }, [mcpClients, toolPermissionMode])
}

function findChromeClient(
  clients: MCPServerConnection[],
): ConnectedMCPServer | undefined {
  return clients.find(
    (client): client is ConnectedMCPServer =>
      client.type === 'connected' &&
      client.name === CLAUDE_IN_CHROME_MCP_SERVER_NAME,
  )
}