agent-ui/frontend/src/services/toolRegistry.ts

/**
 * Tool Registry - Single source of truth for MCP tool management
 *
 * All tool registration/unregistration MUST go through this module.
 * Other modules should NOT directly use webmcp registration functions.
 */

import {
  initWebMCP,
  getRegisteredTools as getWebMCPTools,
  clearAllTools as clearWebMCPTools,
  getWebMCP
} from './webmcp'
import { useToolsStore } from '../stores/tools'
import { useCanvasStore } from '../stores/canvas'
import {
  createGlobalHandlers,
  createCanvasHandlers,
  createComponentHandlers,
  createThemeHandlers,
  createDatabaseHandlers,
  createProjectCanvasHandlers,
  createSourceCodeHandlers,
  createTerminalHandlers,
  createResponseHandlers,
  createGitHandlers,
  createTorchHandlers,
  type ToolConfig
} from './tools/handlers'
import { setRouter } from './tools/handlers/globalHandlers'
import { setGiteaCredentials, clearGiteaCredentials } from './tools/handlers/sourceCodeHandlers'
import { ALL_TOOL_METAS, getAllToolNames, type ToolCategory } from './tools/toolDefinitions'

export type PageName = 'home' | 'canvas' | 'components' | 'themes' | 'projects' | 'project-canvas' | 'database' | 'source' | 'terminal' | 'tools' | 'git'

// Internal webmcp functions (not exported for external use)
let webmcpInstance: any = null
const registeredToolsSet = new Set<string>()

async function internalRegisterTool(config: ToolConfig, force: boolean = false): Promise<boolean> {
  if (!webmcpInstance) {
    webmcpInstance = await initWebMCP()
  }

  // Always re-register if force is true (needed after reconnection)
  if (!force && registeredToolsSet.has(config.name)) {
    return false // Already registered
  }

  // Unregister first if exists (to clear any stale handlers from storage)
  if (registeredToolsSet.has(config.name)) {
    webmcpInstance.unregisterTool(config.name)
  }

  webmcpInstance.registerTool(config.name, config.description, config.schema, config.handler)
  registeredToolsSet.add(config.name)
  console.log(`[ToolRegistry] Registered: ${config.name}`)
  return true
}

function internalUnregisterTool(name: string): boolean {
  if (!webmcpInstance || !registeredToolsSet.has(name)) {
    return false
  }

  webmcpInstance.unregisterTool(name)
  registeredToolsSet.delete(name)
  console.log(`[ToolRegistry] Unregistered: ${name}`)
  return true
}

function internalClearAllTools() {
  if (!webmcpInstance) return

  for (const name of registeredToolsSet) {
    webmcpInstance.unregisterTool(name)
  }
  console.log(`[ToolRegistry] Cleared ${registeredToolsSet.size} tools`)
  registeredToolsSet.clear()
}

// Tool configurations cache
let toolConfigsCache: Map<string, ToolConfig> | null = null

function getToolConfigs(): Map<string, ToolConfig> {
  if (toolConfigsCache) return toolConfigsCache

  toolConfigsCache = new Map()

  // Create callbacks for global handlers
  const toolManagementCallbacks = {
    getRegisteredTools: () => Array.from(registeredToolsSet),
    getAllToolNames: () => getAllToolNames(),
    activateTool: async (name: string) => {
      const config = toolConfigsCache?.get(name)
      if (!config) return false
      const result = await internalRegisterTool(config)
      syncStoreWithActiveTools()
      return result
    },
    deactivateTool: (name: string) => {
      const toolsStore = useToolsStore()
      if (toolsStore.isToolPinned(name)) return false
      const result = internalUnregisterTool(name)
      syncStoreWithActiveTools()
      return result
    },
    togglePin: (name: string) => {
      const toolsStore = useToolsStore()
      toolsStore.togglePin(name)
    },
    isToolPinned: (name: string) => {
      const toolsStore = useToolsStore()
      return toolsStore.isToolPinned(name)
    }
  }

  // Create all handlers
  const allHandlers = [
    ...createGlobalHandlers(toolManagementCallbacks),
    ...createCanvasHandlers(),
    ...createComponentHandlers(),
    ...createThemeHandlers(),
    ...createDatabaseHandlers(),
    ...createProjectCanvasHandlers(),
    ...createSourceCodeHandlers(),
    ...createTerminalHandlers(),
    ...createResponseHandlers(),
    ...createGitHandlers(),
    ...createTorchHandlers()
  ]

  for (const config of allHandlers) {
    toolConfigsCache.set(config.name, config)
  }

  return toolConfigsCache
}

// Category to tool names mapping
const categoryTools: Record<ToolCategory, string[]> = {
  global: ['get_current_page', 'navigate_to', 'list_available_tools', 'activate_tool', 'deactivate_tool', 'pin_tool', 'page_refresh'],
  canvas: ['render_html', 'render_vue_component', 'move_window', 'resize_window', 'close_window', 'list_windows', 'inspect_window', 'get_canvas', 'edit_canvas', 'canvas_css', 'canvas_js', 'get_canvas_css'],
  component: ['save_vue_component', 'load_vue_component', 'list_vue_components', 'delete_vue_component'],
  theme: ['get_design_tokens', 'get_active_theme', 'set_theme_variable', 'save_theme', 'list_themes', 'switch_theme', 'reset_theme'],
  database: ['list_tables', 'get_table_schema', 'get_table_data', 'get_database_stats', 'execute_query'],
  source: ['get_repo_info', 'list_repo_files', 'read_repo_file', 'search_repo_code'],
  project: ['list_canvases', 'create_canvas', 'get_canvas', 'update_canvas', 'delete_canvas', 'clone_canvas', 'add_component_to_canvas', 'remove_component_from_canvas', 'get_canvas_components'],
  terminal: ['terminal_open', 'terminal_close', 'terminal_toggle', 'terminal_move', 'terminal_resize', 'bubbleResponse'],
  git: ['get_git_status', 'get_git_diff', 'compare_commits', 'git_log', 'get_git_branches'],
  torch: ['list_torch_clients', 'get_torch_status', 'transfer_torch', 'request_torch', 'release_torch']
}

// Page to categories mapping
const pageCategories: Record<PageName, ToolCategory[]> = {
  home: ['global', 'torch', 'canvas', 'component', 'project', 'terminal'],
  canvas: ['global', 'torch', 'canvas', 'component', 'terminal'],
  'project-canvas': ['global', 'torch', 'canvas', 'component', 'project', 'terminal'],
  projects: ['global', 'torch', 'project', 'terminal'],
  components: ['global', 'torch', 'component', 'terminal'],
  themes: ['global', 'torch', 'theme', 'terminal'],
  database: ['global', 'torch', 'database', 'terminal'],
  source: ['global', 'torch', 'source', 'terminal'],
  terminal: ['global', 'torch', 'terminal'],
  tools: ['global', 'torch', 'terminal'],
  git: ['global', 'torch', 'git', 'terminal']
}

let currentPage: PageName | null = null
let isInitialized = false

/**
 * Check if connected to MCP
 */
function isConnectedToMCP(): boolean {
  const canvasStore = useCanvasStore()
  return canvasStore.isConnected
}

/**
 * Initialize the tool registry with Vue router
 */
export function initToolRegistry(router: any) {
  setRouter(router)
  isInitialized = true
}

/**
 * Set Gitea credentials for source code tools
 */
export function setSourceCodeCredentials(creds: any) {
  setGiteaCredentials(creds)
}

/**
 * Clear Gitea credentials
 */
export function clearSourceCodeCredentials() {
  clearGiteaCredentials()
}

/**
 * Activate tools for a specific page
 * Only works when connected to MCP
 */
export async function activatePageTools(pageName: PageName, force: boolean = false) {
  if (!isInitialized) {
    console.warn('[ToolRegistry] Not initialized')
    return
  }

  // Only register tools if connected to MCP
  if (!isConnectedToMCP()) {
    console.log(`[ToolRegistry] Not connected to MCP, storing page "${pageName}" for later`)
    currentPage = pageName
    return
  }

  if (!force && currentPage === pageName) {
    console.log(`[ToolRegistry] Already on "${pageName}", skipping`)
    return
  }

  console.log(`[ToolRegistry] Switching from "${currentPage}" to "${pageName}"`)

  const toolsStore = useToolsStore()
  const pinnedTools = toolsStore.getPinnedToolNames()
  const configs = getToolConfigs()

  // Get tools for old and new page
  const oldCategories = currentPage ? pageCategories[currentPage] : []
  const newCategories = pageCategories[pageName]

  const oldTools = new Set(oldCategories.flatMap(cat => categoryTools[cat]))
  const newTools = new Set(newCategories.flatMap(cat => categoryTools[cat]))

  // Unregister old tools (except pinned)
  for (const tool of oldTools) {
    if (!newTools.has(tool) && !pinnedTools.includes(tool)) {
      internalUnregisterTool(tool)
    }
  }

  // Register new tools
  for (const toolName of newTools) {
    const config = configs.get(toolName)
    if (config) {
      await internalRegisterTool(config)
    }
  }

  // Ensure pinned tools are registered
  for (const toolName of pinnedTools) {
    const config = configs.get(toolName)
    if (config && !registeredToolsSet.has(toolName)) {
      await internalRegisterTool(config)
    }
  }

  currentPage = pageName
  syncStoreWithActiveTools()

  console.log(`[ToolRegistry] Page "${pageName}" tools activated`)
}

/**
 * Initialize tools on page refresh
 * Just stores the current page, actual registration happens when connected
 */
export async function initToolsOnRefresh(pageName: PageName) {
  if (!isInitialized) {
    console.warn('[ToolRegistry] Not initialized')
    return
  }

  console.log(`[ToolRegistry] Storing current page "${pageName}" for when connected`)
  currentPage = pageName

  // Only activate if already connected
  if (isConnectedToMCP()) {
    internalClearAllTools()
    await activatePageTools(pageName, true)
  }
}

/**
 * Called when torch is obtained and MCP is connected
 * Activates tools for the current page
 */
export async function onTorchConnected() {
  if (!isInitialized) {
    console.warn('[ToolRegistry] Not initialized')
    return
  }

  if (!currentPage) {
    console.warn('[ToolRegistry] No current page set')
    return
  }

  console.log(`[ToolRegistry] Torch connected, activating tools for "${currentPage}"`)

  // Clear any stale tools from WebMCP storage first
  clearWebMCPTools()
  registeredToolsSet.clear()

  // Force re-register all tools for current page
  await activatePageToolsForced(currentPage)
}

/**
 * Force activate tools for a page (used after torch connection)
 */
async function activatePageToolsForced(pageName: PageName) {
  const toolsStore = useToolsStore()
  const pinnedTools = toolsStore.getPinnedToolNames()
  const configs = getToolConfigs()
  const categories = pageCategories[pageName]
  const tools = new Set(categories.flatMap(cat => categoryTools[cat]))

  // Register all tools with force flag
  for (const toolName of tools) {
    const config = configs.get(toolName)
    if (config) {
      await internalRegisterTool(config, true)
    }
  }

  // Also register pinned tools
  for (const toolName of pinnedTools) {
    if (!tools.has(toolName)) {
      const config = configs.get(toolName)
      if (config) {
        await internalRegisterTool(config, true)
      }
    }
  }

  currentPage = pageName
  syncStoreWithActiveTools()
  console.log(`[ToolRegistry] Force activated ${registeredToolsSet.size} tools for "${pageName}"`)
}

/**
 * Called when torch is lost and MCP is disconnected
 * Clears all tools
 */
export function onTorchDisconnected() {
  console.log('[ToolRegistry] Torch disconnected, clearing tools')
  internalClearAllTools()
  syncStoreWithActiveTools()
}

/**
 * Activate a single tool by name
 */
export async function activateTool(toolName: string): Promise<boolean> {
  const configs = getToolConfigs()
  const config = configs.get(toolName)

  if (!config) {
    console.warn(`[ToolRegistry] Tool "${toolName}" not found`)
    return false
  }

  const result = await internalRegisterTool(config)
  syncStoreWithActiveTools()
  return result
}

/**
 * Deactivate a single tool by name
 */
export function deactivateTool(toolName: string): boolean {
  const toolsStore = useToolsStore()

  if (toolsStore.isToolPinned(toolName)) {
    console.warn(`[ToolRegistry] Cannot deactivate pinned tool "${toolName}"`)
    return false
  }

  const result = internalUnregisterTool(toolName)
  syncStoreWithActiveTools()
  return result
}

/**
 * Activate all tools in a category
 */
export async function activateCategory(category: ToolCategory) {
  const configs = getToolConfigs()
  const tools = categoryTools[category] || []

  for (const toolName of tools) {
    const config = configs.get(toolName)
    if (config) {
      await internalRegisterTool(config)
    }
  }

  syncStoreWithActiveTools()
}

/**
 * Deactivate all tools in a category (respecting pinned)
 */
export function deactivateCategory(category: ToolCategory) {
  const toolsStore = useToolsStore()
  const tools = categoryTools[category] || []

  for (const toolName of tools) {
    if (!toolsStore.isToolPinned(toolName)) {
      internalUnregisterTool(toolName)
    }
  }

  syncStoreWithActiveTools()
}

/**
 * Sync the store with currently active tools
 */
export function syncStoreWithActiveTools() {
  const toolsStore = useToolsStore()
  toolsStore.setActiveTools(Array.from(registeredToolsSet))
}

/**
 * Get current page name
 */
export function getCurrentPage(): PageName | null {
  return currentPage
}

/**
 * Get tool names for a page
 */
export function getPageToolNames(pageName: PageName): string[] {
  const categories = pageCategories[pageName] || []
  return categories.flatMap(cat => categoryTools[cat])
}

/**
 * Check if registry is initialized
 */
export function isRegistryInitialized(): boolean {
  return isInitialized
}

/**
 * Get all tool metadata
 */
export function getAllToolMetas() {
  return ALL_TOOL_METAS
}

/**
 * Get registered tools (for internal use)
 */
export function getRegisteredTools(): string[] {
  return Array.from(registeredToolsSet)
}