new-api/web/default/src/components/ai-elements/tool.tsx

'use client'

import { type ComponentProps, isValidElement, type ReactNode } from 'react'
import type { ToolUIPart } from 'ai'
import {
  CheckCircleIcon,
  ChevronDownIcon,
  CircleIcon,
  ClockIcon,
  WrenchIcon,
  XCircleIcon,
} from 'lucide-react'
import { useTranslation } from 'react-i18next'
import { cn } from '@/lib/utils'
import { Badge } from '@/components/ui/badge'
import {
  Collapsible,
  CollapsibleContent,
  CollapsibleTrigger,
} from '@/components/ui/collapsible'
import { CodeBlock } from './code-block'

// Workaround for missing types in 'ai' package
type ExtendedToolState =
  | ToolUIPart['state']
  | 'approval-requested'
  | 'approval-responded'
  | 'output-denied'

export type ToolProps = ComponentProps<typeof Collapsible>

export const Tool = ({ className, ...props }: ToolProps) => (
  <Collapsible
    className={cn('not-prose mb-4 w-full rounded-md border', className)}
    {...props}
  />
)

export type ToolHeaderProps = {
  title?: string
  type: ToolUIPart['type']
  state: ExtendedToolState
  className?: string
}

const getStatusBadge = (status: ExtendedToolState) => {
  const labels: Record<ExtendedToolState, string> = {
    'input-streaming': 'Pending',
    'input-available': 'Running',
    'approval-requested': 'Awaiting Approval',
    'approval-responded': 'Responded',
    'output-available': 'Completed',
    'output-error': 'Error',
    'output-denied': 'Denied',
  }

  const icons: Record<ExtendedToolState, ReactNode> = {
    'input-streaming': <CircleIcon className='size-4' />,
    'input-available': <ClockIcon className='size-4 animate-pulse' />,
    'approval-requested': <ClockIcon className='size-4 text-yellow-600' />,
    'approval-responded': <CheckCircleIcon className='size-4 text-blue-600' />,
    'output-available': <CheckCircleIcon className='size-4 text-green-600' />,
    'output-error': <XCircleIcon className='size-4 text-red-600' />,
    'output-denied': <XCircleIcon className='size-4 text-orange-600' />,
  }

  return (
    <Badge className='gap-1.5 rounded-full text-xs' variant='secondary'>
      {icons[status]}
      {labels[status]}
    </Badge>
  )
}

export const ToolHeader = ({
  className,
  title,
  type,
  state,
  ...props
}: ToolHeaderProps) => (
  <CollapsibleTrigger
    className={cn(
      'group flex w-full items-center justify-between gap-4 p-3',
      className
    )}
    {...props}
  >
    <div className='flex items-center gap-2'>
      <WrenchIcon className='text-muted-foreground size-4' />
      <span className='text-sm font-medium'>
        {title ?? type.split('-').slice(1).join('-')}
      </span>
      {getStatusBadge(state)}
    </div>
    <ChevronDownIcon className='text-muted-foreground size-4 transition-transform group-data-[panel-open]:rotate-180' />
  </CollapsibleTrigger>
)

export type ToolContentProps = ComponentProps<typeof CollapsibleContent>

export const ToolContent = ({ className, ...props }: ToolContentProps) => (
  <CollapsibleContent
    className={cn(
      'data-closed:fade-out-0 data-closed:slide-out-to-top-2 data-open:slide-in-from-top-2 text-popover-foreground data-closed:animate-out data-open:animate-in outline-none',
      className
    )}
    {...props}
  />
)

export type ToolInputProps = ComponentProps<'div'> & {
  input: ToolUIPart['input']
}

export const ToolInput = ({ className, input, ...props }: ToolInputProps) => {
  const { t } = useTranslation()
  return (
    <div className={cn('space-y-2 overflow-hidden p-4', className)} {...props}>
      <h4 className='text-muted-foreground text-xs font-medium tracking-wide uppercase'>
        {t('Parameters')}
      </h4>
      <div className='bg-muted/50 rounded-md'>
        <CodeBlock code={JSON.stringify(input, null, 2)} language='json' />
      </div>
    </div>
  )
}

export type ToolOutputProps = ComponentProps<'div'> & {
  output: ToolUIPart['output']
  errorText: ToolUIPart['errorText']
}

export const ToolOutput = ({
  className,
  output,
  errorText,
  ...props
}: ToolOutputProps) => {
  if (!(output || errorText)) {
    return null
  }

  let Output = <div>{output as ReactNode}</div>

  if (typeof output === 'object' && !isValidElement(output)) {
    Output = (
      <CodeBlock code={JSON.stringify(output, null, 2)} language='json' />
    )
  } else if (typeof output === 'string') {
    Output = <CodeBlock code={output} language='json' />
  }

  return (
    <div className={cn('space-y-2 p-4', className)} {...props}>
      <h4 className='text-muted-foreground text-xs font-medium tracking-wide uppercase'>
        {errorText ? 'Error' : 'Result'}
      </h4>
      <div
        className={cn(
          'overflow-x-auto rounded-md text-xs [&_table]:w-full',
          errorText
            ? 'bg-destructive/10 text-destructive'
            : 'bg-muted/50 text-foreground'
        )}
      >
        {errorText && <div>{errorText}</div>}
        {Output}
      </div>
    </div>
  )
}