SurfSense/surfsense_web/components/tool-ui/deepagent-thinking.tsx

"use client";

import { makeAssistantToolUI } from "@assistant-ui/react";
import { Brain, CheckCircle2, Loader2, Search, Sparkles } from "lucide-react";
import { useMemo, useState, useEffect, useRef } from "react";
import {
  ChainOfThought,
  ChainOfThoughtContent,
  ChainOfThoughtItem,
  ChainOfThoughtStep,
  ChainOfThoughtTrigger,
} from "@/components/prompt-kit/chain-of-thought";
import { cn } from "@/lib/utils";

/**
 * Types for the deepagent thinking/reasoning tool
 */
interface ThinkingStep {
  id: string;
  title: string;
  items: string[];
  status: "pending" | "in_progress" | "completed";
}

interface DeepAgentThinkingArgs {
  query?: string;
  context?: string;
}

interface DeepAgentThinkingResult {
  steps?: ThinkingStep[];
  status?: "thinking" | "searching" | "synthesizing" | "completed";
  summary?: string;
}

/**
 * Get icon based on step status and type
 */
function getStepIcon(status: "pending" | "in_progress" | "completed", title: string) {
  // Check for specific step types based on title keywords
  const titleLower = title.toLowerCase();
  
  if (status === "in_progress") {
    return <Loader2 className="size-4 animate-spin text-primary" />;
  }
  
  if (status === "completed") {
    return <CheckCircle2 className="size-4 text-emerald-500" />;
  }
  
  // Default icons based on step type
  if (titleLower.includes("search") || titleLower.includes("knowledge")) {
    return <Search className="size-4 text-muted-foreground" />;
  }
  
  if (titleLower.includes("analy") || titleLower.includes("understand")) {
    return <Brain className="size-4 text-muted-foreground" />;
  }
  
  return <Sparkles className="size-4 text-muted-foreground" />;
}

/**
 * Component to display a single thinking step with controlled open state
 */
function ThinkingStepDisplay({ 
  step, 
  isOpen, 
  onToggle 
}: { 
  step: ThinkingStep; 
  isOpen: boolean;
  onToggle: () => void;
}) {
  const icon = useMemo(() => getStepIcon(step.status, step.title), [step.status, step.title]);
  
  return (
    <ChainOfThoughtStep open={isOpen} onOpenChange={onToggle}>
      <ChainOfThoughtTrigger 
        leftIcon={icon}
        swapIconOnHover={step.status !== "in_progress"}
        className={cn(
          step.status === "in_progress" && "text-foreground font-medium",
          step.status === "completed" && "text-muted-foreground"
        )}
      >
        {step.title}
      </ChainOfThoughtTrigger>
      <ChainOfThoughtContent>
        {step.items.map((item, index) => (
          <ChainOfThoughtItem key={`${step.id}-item-${index}`}>
            {item}
          </ChainOfThoughtItem>
        ))}
      </ChainOfThoughtContent>
    </ChainOfThoughtStep>
  );
}

/**
 * Loading state with animated thinking indicator
 */
function ThinkingLoadingState({ status }: { status?: string }) {
  const statusText = useMemo(() => {
    switch (status) {
      case "searching":
        return "Searching knowledge base...";
      case "synthesizing":
        return "Synthesizing response...";
      case "thinking":
      default:
        return "Thinking...";
    }
  }, [status]);
  
  return (
    <div className="my-3 flex items-center gap-2 rounded-lg border border-border/50 bg-muted/30 px-4 py-3">
      <div className="relative">
        <Brain className="size-5 text-primary" />
        <span className="absolute -right-0.5 -top-0.5 flex size-2">
          <span className="absolute inline-flex size-full animate-ping rounded-full bg-primary/60" />
          <span className="relative inline-flex size-2 rounded-full bg-primary" />
        </span>
      </div>
      <span className="text-sm text-muted-foreground">{statusText}</span>
    </div>
  );
}

/**
 * Smart chain of thought renderer with state management
 */
function SmartChainOfThought({ steps }: { steps: ThinkingStep[] }) {
  // Track which steps the user has manually toggled
  const [manualOverrides, setManualOverrides] = useState<Record<string, boolean>>({});
  // Track previous step statuses to detect changes
  const prevStatusesRef = useRef<Record<string, string>>({});
  
  // Check if any step is currently in progress
  const hasInProgressStep = steps.some(step => step.status === "in_progress");
  
  // Find the last completed step index
  const lastCompletedIndex = steps
    .map((s, i) => s.status === "completed" ? i : -1)
    .filter(i => i !== -1)
    .pop();
  
  // Clear manual overrides when a step's status changes
  useEffect(() => {
    const currentStatuses: Record<string, string> = {};
    steps.forEach(step => {
      currentStatuses[step.id] = step.status;
      // If status changed, clear any manual override for this step
      if (prevStatusesRef.current[step.id] && prevStatusesRef.current[step.id] !== step.status) {
        setManualOverrides(prev => {
          const next = { ...prev };
          delete next[step.id];
          return next;
        });
      }
    });
    prevStatusesRef.current = currentStatuses;
  }, [steps]);
  
  const getStepOpenState = (step: ThinkingStep, index: number): boolean => {
    // If user has manually toggled, respect that
    if (manualOverrides[step.id] !== undefined) {
      return manualOverrides[step.id];
    }
    // Auto behavior: open if in progress
    if (step.status === "in_progress") {
      return true;
    }
    // Auto behavior: keep last completed step open if no in-progress step
    if (!hasInProgressStep && index === lastCompletedIndex) {
      return true;
    }
    // Default: collapsed
    return false;
  };
  
  const handleToggle = (stepId: string, currentOpen: boolean) => {
    setManualOverrides(prev => ({
      ...prev,
      [stepId]: !currentOpen,
    }));
  };
  
  return (
    <ChainOfThought>
      {steps.map((step, index) => {
        const isOpen = getStepOpenState(step, index);
        return (
          <ThinkingStepDisplay 
            key={step.id} 
            step={step} 
            isOpen={isOpen}
            onToggle={() => handleToggle(step.id, isOpen)}
          />
        );
      })}
    </ChainOfThought>
  );
}

/**
 * DeepAgent Thinking Tool UI Component
 *
 * This component displays the agent's chain-of-thought reasoning
 * when the deepagent is processing a query. It shows thinking steps
 * in a collapsible, hierarchical format.
 */
export const DeepAgentThinkingToolUI = makeAssistantToolUI<
  DeepAgentThinkingArgs,
  DeepAgentThinkingResult
>({
  toolName: "deepagent_thinking",
  render: function DeepAgentThinkingUI({ result, status }) {
    // Loading state - tool is still running
    if (status.type === "running" || status.type === "requires-action") {
      return <ThinkingLoadingState status={result?.status} />;
    }

    // Incomplete/cancelled state
    if (status.type === "incomplete") {
      if (status.reason === "cancelled") {
        return null; // Don't show anything if cancelled
      }
      if (status.reason === "error") {
        return null; // Don't show error for thinking - it's not critical
      }
    }

    // No result or no steps - don't render anything
    if (!result?.steps || result.steps.length === 0) {
      return null;
    }

    // Render the chain of thought
    return (
      <div className="my-3 w-full">
        <SmartChainOfThought steps={result.steps} />
      </div>
    );
  },
});

/**
 * Inline Thinking Display Component
 * 
 * A simpler version that can be used inline with the message content
 * for displaying reasoning without the full tool UI infrastructure.
 */
export function InlineThinkingDisplay({
  steps,
  isStreaming = false,
  className,
}: {
  steps: ThinkingStep[];
  isStreaming?: boolean;
  className?: string;
}) {
  if (steps.length === 0 && !isStreaming) {
    return null;
  }

  return (
    <div className={cn("my-3 w-full", className)}>
      {isStreaming && steps.length === 0 ? (
        <ThinkingLoadingState />
      ) : (
        <SmartChainOfThought steps={steps} />
      )}
    </div>
  );
}

export type { ThinkingStep, DeepAgentThinkingArgs, DeepAgentThinkingResult };