feat: enhance AI Chat with context management, file integration, image support, and improved code handling

- Added context tabs system for managing multiple types of context (files, code snippets, images)
   - Added preview functionality for context items
   - Added ability to expand/collapse context previews
   - Added file selection popup/dropdown
   - Added file search functionality
   - Added image upload button
   - Added image paste support
   - Added image preview in context tabs
   - Added automatic code detection on paste
   - Added line number tracking for code snippets
   - Added source file name preservation
   - Added line range display for code contexts
   - Added model selection dropdown (Claude 3.5 Sonnet/Claude 3)
   - Added Ctrl+Enter for sending with full context
   - Added Backspace to remove last context tab when input is empty
   - Added smart code detection on paste

frontend/components/editor/AIChat/ChatInput.tsx

@@ -1,12 +1,26 @@
-import { Send, StopCircle } from "lucide-react"
+import { Send, StopCircle, AtSign, Image as ImageIcon } from "lucide-react"
 import { Button } from "../../ui/button"
+import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "../../ui/select"
+import { useRef, useEffect, useState } from "react"
+import * as monaco from 'monaco-editor'
+import { TFile, TFolder } from "@/lib/types"
 
 interface ChatInputProps {
   input: string
   setInput: (input: string) => void
   isGenerating: boolean
-  handleSend: () => void
+  handleSend: (useFullContext?: boolean) => void
   handleStopGeneration: () => void
+  onImageUpload: (file: File) => void
+  onFileMention: (fileName: string) => void
+  addContextTab: (type: string, title: string, content: string, lineRange?: { start: number, end: number }) => void
+  activeFileName?: string
+  editorRef: React.MutableRefObject<monaco.editor.IStandaloneCodeEditor | undefined>
+  lastCopiedRangeRef: React.MutableRefObject<{ startLine: number; endLine: number } | null>
+  contextTabs: { id: string; type: string; title: string; content: string; lineRange?: { start: number; end: number } }[]
+  onRemoveTab: (id: string) => void
+  textareaRef: React.RefObject<HTMLTextAreaElement>
+  files: (TFile | TFolder)[]
 }
 
 export default function ChatInput({
@@ -15,17 +29,194 @@ export default function ChatInput({
   isGenerating,
   handleSend,
   handleStopGeneration,
+  onImageUpload,
+  onFileMention,
+  addContextTab,
+  activeFileName,
+  editorRef,
+  lastCopiedRangeRef,
+  contextTabs,
+  onRemoveTab,
+  textareaRef,
+  files,
 }: ChatInputProps) {
+  // Auto-resize textarea as content changes
+  useEffect(() => {
+    if (textareaRef.current) {
+      textareaRef.current.style.height = 'auto'
+      textareaRef.current.style.height = textareaRef.current.scrollHeight + 'px'
+    }
+  }, [input])
+
+  const handleKeyDown = (e: React.KeyboardEvent) => {
+    if (e.key === "Enter") {
+      if (e.ctrlKey) {
+        e.preventDefault()
+        handleSend(true) // Send with full context
+      } else if (!e.shiftKey && !isGenerating) {
+        e.preventDefault()
+        handleSend(false)
+      }
+    } else if (e.key === "Backspace" && input === "" && contextTabs.length > 0) {
+      e.preventDefault()
+      // Remove the last context tab
+      const lastTab = contextTabs[contextTabs.length - 1]
+      onRemoveTab(lastTab.id)
+    }
+  }
+
+  const handlePaste = async (e: React.ClipboardEvent) => {
+    // Handle image paste
+    const items = Array.from(e.clipboardData.items);
+    for (const item of items) {
+      if (item.type.startsWith('image/')) {
+        e.preventDefault();
+        
+        const file = item.getAsFile();
+        if (!file) continue;
+
+        try {
+          const reader = new FileReader();
+          reader.onload = () => {
+            const base64String = reader.result as string;
+            addContextTab(
+              "image",
+              `Image ${new Date().toLocaleTimeString()}`,
+              base64String
+            );
+          };
+          reader.readAsDataURL(file);
+        } catch (error) {
+          console.error('Error processing pasted image:', error);
+        }
+        return;
+      }
+    }
+    
+    const text = e.clipboardData.getData('text');
+    
+    // Helper function to detect if text looks like code
+    const looksLikeCode = (text: string): boolean => {
+      const codeIndicators = [
+        /^import\s+/m,           // import statements
+        /^function\s+/m,         // function declarations
+        /^class\s+/m,           // class declarations
+        /^const\s+/m,           // const declarations
+        /^let\s+/m,             // let declarations
+        /^var\s+/m,             // var declarations
+        /[{}\[\]();]/,          // common code syntax
+        /^\s*\/\//m,            // comments
+        /^\s*\/\*/m,            // multi-line comments
+        /=>/,                   // arrow functions
+        /^export\s+/m,          // export statements
+      ];
+
+      return codeIndicators.some(pattern => pattern.test(text));
+    };
+
+    // If text doesn't contain newlines or doesn't look like code, let it paste normally
+    if (!text || !text.includes('\n') || !looksLikeCode(text)) {
+      return;
+    }
+
+    e.preventDefault();
+    const editor = editorRef.current;
+    const currentSelection = editor?.getSelection();
+    const lines = text.split('\n');
+    
+    // If selection exists in editor, use file name and line numbers
+    if (currentSelection && !currentSelection.isEmpty()) {
+      addContextTab(
+        "code",
+        `${activeFileName} (${currentSelection.startLineNumber}-${currentSelection.endLineNumber})`,
+        text,
+        { start: currentSelection.startLineNumber, end: currentSelection.endLineNumber }
+      );
+      return;
+    }
+    
+    // If we have stored line range from a copy operation in the editor
+    if (lastCopiedRangeRef.current) {
+      const range = lastCopiedRangeRef.current;
+      addContextTab(
+        "code",
+        `${activeFileName} (${range.startLine}-${range.endLine})`,
+        text,
+        { start: range.startLine, end: range.endLine }
+      );
+      return;
+    }
+    
+    // For code pasted from outside the editor
+    addContextTab(
+      "code",
+      `Pasted Code (1-${lines.length})`,
+      text,
+      { start: 1, end: lines.length }
+    );
+  };
+
+  const handleImageUpload = () => {
+    const input = document.createElement('input')
+    input.type = 'file'
+    input.accept = 'image/*'
+    input.onchange = (e) => {
+      const file = (e.target as HTMLInputElement).files?.[0]
+      if (file) onImageUpload(file)
+    }
+    input.click()
+  }
+
+  const handleMentionClick = () => {
+    if (textareaRef.current) {
+      const cursorPosition = textareaRef.current.selectionStart
+      const newValue = input.slice(0, cursorPosition) + '@' + input.slice(cursorPosition)
+      setInput(newValue)
+      // Focus and move cursor after the @
+      textareaRef.current.focus()
+      const newPosition = cursorPosition + 1
+      textareaRef.current.setSelectionRange(newPosition, newPosition)
+    }
+  }
+
+  // Handle @ mentions in input
+  useEffect(() => {
+    const match = input.match(/@(\w+)$/)
+    if (match) {
+      const fileName = match[1]
+      const allFiles = getAllFiles(files)
+      const file = allFiles.find(file => file.name === fileName)
+      if (file) {
+        onFileMention(file.name)
+      }
+    }
+  }, [input, onFileMention, files])
+
+  // Add this helper function to flatten the file tree
+  const getAllFiles = (items: (TFile | TFolder)[]): TFile[] => {
+    return items.reduce((acc: TFile[], item) => {
+      if (item.type === "file") {
+        acc.push(item)
+      } else {
+        acc.push(...getAllFiles(item.children))
+      }
+      return acc
+    }, [])
+  }
+
   return (
+    <div className="space-y-2">
       <div className="flex space-x-2 min-w-0">
-      <input
+        <textarea
-        type="text"
+          ref={textareaRef}
           value={input}
           onChange={(e) => setInput(e.target.value)}
-        onKeyPress={(e) => e.key === "Enter" && !isGenerating && handleSend()}
+          onKeyDown={handleKeyDown}
-        className="flex-grow p-2 border rounded-lg min-w-0 bg-input"
+          onPaste={handlePaste}
+          className="flex-grow p-2 border rounded-lg min-w-0 bg-input resize-none overflow-hidden"
           placeholder="Type your message..."
           disabled={isGenerating}
+          rows={1}
         />
         {isGenerating ? (
           <Button
@@ -38,7 +229,7 @@ export default function ChatInput({
           </Button>
         ) : (
           <Button
-          onClick={handleSend}
+            onClick={() => handleSend(false)}
             disabled={isGenerating}
             size="icon"
             className="h-10 w-10"
@@ -47,5 +238,40 @@ export default function ChatInput({
           </Button>
         )}
       </div>
+      <div className="flex flex-wrap items-center gap-2 w-full">
+        <div className="flex items-center gap-2 px-2 text-sm text-muted-foreground min-w-auto max-w-auto">
+          <Select defaultValue="claude-3.5-sonnet">
+            <SelectTrigger className="h-6 w-full border-none truncate">
+              <SelectValue />
+            </SelectTrigger>
+            <SelectContent>
+              <SelectItem value="claude-3.5-sonnet">claude-3.5-sonnet</SelectItem>
+              <SelectItem value="claude-3">claude-3</SelectItem>
+            </SelectContent>
+          </Select>
+        </div>
+        <div className="flex items-center gap-2 ml-auto">
+          <Button 
+            variant="ghost" 
+            size="sm"
+            className="h-6 px-2 sm:px-3"
+            onClick={handleMentionClick}
+          >
+            <AtSign className="h-3 w-3 sm:mr-1" />
+            <span className="hidden sm:inline">mention</span>
+          </Button>
+          <Button 
+            variant="ghost" 
+            size="sm"
+            className="h-6 px-2 sm:px-3"
+            onClick={handleImageUpload}
+          >
+            <ImageIcon className="h-3 w-3 sm:mr-1" />
+            <span className="hidden sm:inline">Image</span>
+          </Button>
+        </div>
+      </div>
+    </div>
   )
 }
+

frontend/components/editor/AIChat/ChatMessage.tsx

@@ -6,6 +6,8 @@ import { vscDarkPlus } from "react-syntax-highlighter/dist/esm/styles/prism"
 import remarkGfm from "remark-gfm"
 import { Button } from "../../ui/button"
 import { copyToClipboard, stringifyContent } from "./lib/chatUtils"
+import ContextTabs from "./ContextTabs"
+import { Socket } from "socket.io-client"
 
 interface MessageProps {
   message: {
@@ -15,12 +17,14 @@ interface MessageProps {
   }
   setContext: (context: string | null) => void
   setIsContextExpanded: (isExpanded: boolean) => void
+  socket: Socket | null
 }
 
 export default function ChatMessage({
   message,
   setContext,
   setIsContextExpanded,
+  socket,
 }: MessageProps) {
   const [expandedMessageIndex, setExpandedMessageIndex] = useState<
     number | null
@@ -88,34 +92,18 @@ export default function ChatMessage({
             : "bg-transparent text-white"
         } max-w-full`}
       >
-        {message.role === "user" && (
+        {message.role === "user" && message.context && (
-          <div className="absolute top-0 right-0 flex opacity-0 group-hover:opacity-30 transition-opacity">
-            {renderCopyButton(message.content)}
-            <Button
-              onClick={() => askAboutCode(message.content)}
-              size="sm"
-              variant="ghost"
-              className="p-1 h-6"
-            >
-              <CornerUpLeft className="w-4 h-4" />
-            </Button>
-          </div>
-        )}
-        {message.context && (
           <div className="mb-2 bg-input rounded-lg">
-            <div
+            <ContextTabs
-              className="flex justify-between items-center cursor-pointer"
+              socket={socket}
-              onClick={() =>
+              activeFileName=""
-                setExpandedMessageIndex(expandedMessageIndex === 0 ? null : 0)
+              onAddFile={() => {}}
-              }
+              contextTabs={parseContextToTabs(message.context)}
-            >
+              onRemoveTab={() => {}}
-              <span className="text-sm text-gray-300">Context</span>
+              isExpanded={expandedMessageIndex === 0}
-              {expandedMessageIndex === 0 ? (
+              onToggleExpand={() => setExpandedMessageIndex(expandedMessageIndex === 0 ? null : 0)}
-                <ChevronUp size={16} />
+              className="[&_div:first-child>div:first-child>div]:bg-[#0D0D0D] [&_button:first-child]:hidden [&_button:last-child]:hidden"
-              ) : (
+            />
-                <ChevronDown size={16} />
-              )}
-            </div>
             {expandedMessageIndex === 0 && (
               <div className="relative">
                 <div className="absolute top-0 right-0 flex p-1">
@@ -153,6 +141,19 @@ export default function ChatMessage({
             )}
           </div>
         )}
+        {message.role === "user" && (
+          <div className="absolute top-0 right-0 flex opacity-0 group-hover:opacity-30 transition-opacity">
+            {renderCopyButton(message.content)}
+            <Button
+              onClick={() => askAboutCode(message.content)}
+              size="sm"
+              variant="ghost"
+              className="p-1 h-6"
+            >
+              <CornerUpLeft className="w-4 h-4" />
+            </Button>
+          </div>
+        )}
         {message.role === "assistant" ? (
           <ReactMarkdown
             remarkPlugins={[remarkGfm]}
@@ -224,3 +225,25 @@ export default function ChatMessage({
     </div>
   )
 }
+
+function parseContextToTabs(context: string) {
+  const sections = context.split(/(?=File |Code from )/)
+  return sections.map((section, index) => {
+    const lines = section.trim().split('\n')
+    const titleLine = lines[0]
+    let content = lines.slice(1).join('\n').trim()
+    
+    // Remove code block markers for display
+    content = content.replace(/^```[\w-]*\n/, '').replace(/\n```$/, '')
+    
+    const isFile = titleLine.startsWith('File ')
+    const name = titleLine.replace(/^(File |Code from )/, '').replace(':', '')
+    
+    return {
+      id: `context-${index}`,
+      type: isFile ? "file" as const : "code" as const,
+      name: name,
+      content: content
+    }
+  }).filter(tab => tab.content.length > 0)
+}

frontend/components/editor/AIChat/ContextTabs.tsx

@@ -0,0 +1,167 @@
+import { Plus, X, ChevronDown, ChevronUp, Image as ImageIcon, FileText } from "lucide-react"
+import { useState } from "react"
+import { Button } from "../../ui/button"
+import { TFile, TFolder } from "@/lib/types"
+import {
+  Popover,
+  PopoverContent,
+  PopoverTrigger,
+} from "@/components/ui/popover"
+import { Input } from "@/components/ui/input"
+import { Socket } from "socket.io-client"
+
+interface ContextTab {
+  id: string
+  type: "file" | "code" | "image"
+  name: string
+  content: string
+  sourceFile?: string
+  lineRange?: {
+    start: number
+    end: number
+  }
+}
+
+interface ContextTabsProps {
+  activeFileName: string
+  onAddFile: () => void
+  contextTabs: ContextTab[]
+  onRemoveTab: (id: string) => void
+  isExpanded: boolean
+  onToggleExpand: () => void
+  files?: (TFile | TFolder)[]
+  onFileSelect?: (file: TFile) => void
+  socket: Socket | null
+}
+
+export default function ContextTabs({
+  onAddFile,
+  contextTabs,
+  onRemoveTab,
+  className,
+  files = [],
+  onFileSelect,
+}: ContextTabsProps & { className?: string }) {
+  const [previewTab, setPreviewTab] = useState<ContextTab | null>(null)
+  const [searchQuery, setSearchQuery] = useState("")
+
+  const togglePreview = (tab: ContextTab) => {
+    if (previewTab?.id === tab.id) {
+      setPreviewTab(null)
+    } else {
+      setPreviewTab(tab)
+    }
+  }
+
+  const handleRemoveTab = (id: string) => {
+    if (previewTab?.id === id) {
+      setPreviewTab(null)
+    }
+    onRemoveTab(id)
+  }
+
+  const getAllFiles = (items: (TFile | TFolder)[]): TFile[] => {
+    return items.reduce((acc: TFile[], item) => {
+      if (item.type === "file") {
+        acc.push(item)
+      } else {
+        acc.push(...getAllFiles(item.children))
+      }
+      return acc
+    }, [])
+  }
+
+  const allFiles = getAllFiles(files)
+  const filteredFiles = allFiles.filter(file => 
+    file.name.toLowerCase().includes(searchQuery.toLowerCase())
+  )
+
+  return (
+    <div className={`border-none ${className || ''}`}>
+      <div className="flex flex-col">
+        <div className="flex items-center gap-1 overflow-hidden mb-2 flex-wrap">
+          <Popover>
+            <PopoverTrigger asChild>
+              <Button
+                variant="ghost"
+                size="icon"
+                className="h-6 w-6"
+              >
+                <Plus className="h-4 w-4" />
+              </Button>
+            </PopoverTrigger>
+            <PopoverContent className="w-64 p-2">
+              <Input
+                placeholder="Search files..."
+                value={searchQuery}
+                onChange={(e) => setSearchQuery(e.target.value)}
+                className="mb-2"
+              />
+              <div className="max-h-[200px] overflow-y-auto">
+                {filteredFiles.map((file) => (
+                  <Button
+                    key={file.id}
+                    variant="ghost"
+                    className="w-full justify-start text-sm mb-1"
+                    onClick={() => onFileSelect?.(file)}
+                  >
+                    <FileText className="h-4 w-4 mr-2" />
+                    {file.name}
+                  </Button>
+                ))}
+              </div>
+            </PopoverContent>
+          </Popover>
+          {contextTabs.length === 0 && (
+            <div className="flex items-center gap-1 px-2 rounded">
+              <span className="text-sm text-muted-foreground">Add Context</span>
+            </div>
+          )}
+          {contextTabs.map((tab) => (
+            <div
+              key={tab.id}
+              className="flex items-center gap-1 px-2 bg-input rounded text-sm cursor-pointer hover:bg-muted"
+              onClick={() => togglePreview(tab)}
+            >
+              {tab.type === "image" && <ImageIcon className="h-3 w-3" />}
+              <span>{tab.name}</span>
+              <Button
+                variant="ghost"
+                size="icon"
+                className="h-4 w-4"
+                onClick={(e) => {
+                  e.stopPropagation()
+                  handleRemoveTab(tab.id)
+                }}
+              >
+                <X className="h-3 w-3" />
+              </Button>
+            </div>
+          ))}
+        </div>
+
+        {/* Preview Section */}
+        {previewTab && (
+          <div className="p-2 bg-input rounded-md max-h-[200px] overflow-auto mb-2">
+            {previewTab.lineRange && (
+              <div className="text-xs text-muted-foreground mt-1">
+                Lines {previewTab.lineRange.start}-{previewTab.lineRange.end}
+              </div>
+            )}
+            {previewTab.type === "image" ? (
+              <img 
+                src={previewTab.content} 
+                alt={previewTab.name}
+                className="max-w-full h-auto"
+              />
+            ) : (
+              <pre className="text-xs font-mono whitespace-pre-wrap">
+                {previewTab.content}
+              </pre>
+            )}
+          </div>
+        )}
+      </div>
+    </div>
+  )
+}

frontend/components/editor/AIChat/index.tsx

@@ -3,8 +3,12 @@ import { useEffect, useRef, useState } from "react"
 import LoadingDots from "../../ui/LoadingDots"
 import ChatInput from "./ChatInput"
 import ChatMessage from "./ChatMessage"
-import ContextDisplay from "./ContextDisplay"
+import ContextTabs from "./ContextTabs"
 import { handleSend, handleStopGeneration } from "./lib/chatUtils"
+import { nanoid } from 'nanoid'
+import * as monaco from 'monaco-editor'
+import { TFile, TFolder } from "@/lib/types"
+import { useSocket } from "@/context/SocketContext"
 
 interface Message {
   role: "user" | "assistant"
@@ -12,23 +16,41 @@ interface Message {
   context?: string
 }
 
+interface ContextTab {
+  id: string
+  type: "file" | "code" | "image"
+  name: string
+  content: string
+  lineRange?: { start: number; end: number }
+}
+
+interface AIChatProps {
+  activeFileContent: string
+  activeFileName: string
+  onClose: () => void
+  editorRef: React.MutableRefObject<monaco.editor.IStandaloneCodeEditor | undefined>
+  lastCopiedRangeRef: React.MutableRefObject<{ startLine: number; endLine: number } | null>
+  files: (TFile | TFolder)[]
+}
+
 export default function AIChat({
   activeFileContent,
   activeFileName,
   onClose,
-}: {
+  editorRef,
-  activeFileContent: string
+  lastCopiedRangeRef,
-  activeFileName: string
+  files,
-  onClose: () => void
+}: AIChatProps) {
-}) {
+  const { socket } = useSocket()
   const [messages, setMessages] = useState<Message[]>([])
   const [input, setInput] = useState("")
   const [isGenerating, setIsGenerating] = useState(false)
   const chatContainerRef = useRef<HTMLDivElement>(null)
   const abortControllerRef = useRef<AbortController | null>(null)
-  const [context, setContext] = useState<string | null>(null)
+  const [contextTabs, setContextTabs] = useState<ContextTab[]>([])
   const [isContextExpanded, setIsContextExpanded] = useState(false)
   const [isLoading, setIsLoading] = useState(false)
+  const textareaRef = useRef<HTMLTextAreaElement>(null)
 
   useEffect(() => {
     scrollToBottom()
@@ -45,6 +67,109 @@ export default function AIChat({
     }
   }
 
+  const addContextTab = (type: string, name: string, content: string, lineRange?: { start: number; end: number }) => {
+    const newTab = {
+      id: nanoid(),
+      type: type as "file" | "code" | "image",
+      name,
+      content,
+      lineRange
+    }
+    setContextTabs(prev => [...prev, newTab])
+  }
+
+  const removeContextTab = (id: string) => {
+    setContextTabs(prev => prev.filter(tab => tab.id !== id))
+  }
+
+  const handleAddFile = () => {
+    console.log("Add file to context")
+  }
+
+  const formatCodeContent = (content: string) => {
+    // Remove starting and ending code block markers if they exist
+    return content.replace(/^```[\w-]*\n/, '').replace(/\n```$/, '')
+  }
+
+  const getCombinedContext = () => {
+    if (contextTabs.length === 0) return ''
+    
+    return contextTabs.map(tab => {
+      if (tab.type === 'file') {
+        const fileExt = tab.name.split('.').pop() || 'txt'
+        const cleanContent = formatCodeContent(tab.content)
+        return `File ${tab.name}:\n\`\`\`${fileExt}\n${cleanContent}\n\`\`\``
+      } else if (tab.type === 'code') {
+        const cleanContent = formatCodeContent(tab.content)
+        return `Code from ${tab.name}:\n\`\`\`typescript\n${cleanContent}\n\`\`\``
+      }
+      return `${tab.name}:\n${tab.content}`
+    }).join('\n\n')
+  }
+
+  const handleSendWithContext = () => {
+    const combinedContext = getCombinedContext()
+    handleSend(
+      input,
+      combinedContext,
+      messages,
+      setMessages,
+      setInput,
+      setIsContextExpanded,
+      setIsGenerating,
+      setIsLoading,
+      abortControllerRef,
+      activeFileContent
+    )
+    // Clear context tabs after sending
+    setContextTabs([])
+  }
+
+  function setContext(context: string | null, fileName?: string, lineRange?: { start: number; end: number }): void {
+    if (!context) {
+      setContextTabs([])
+      return
+    }
+
+    const existingCodeTab = contextTabs.find(tab => tab.type === 'code')
+    
+    if (existingCodeTab) {
+      setContextTabs(prev => 
+        prev.map(tab => 
+          tab.id === existingCodeTab.id 
+            ? { ...tab, content: context, name: fileName || 'Code Context', lineRange }
+            : tab
+        )
+      )
+    } else {
+      addContextTab('code', fileName || 'Chat Context', context, lineRange)
+    }
+  }
+
+  useEffect(() => {
+    if (editorRef?.current) {
+      const editor = editorRef.current;
+      
+      // Configure editor options for better copy handling
+      editor.updateOptions({
+        copyWithSyntaxHighlighting: true,
+        emptySelectionClipboard: false
+      });
+
+      // Track selection changes
+      const disposable = editor.onDidChangeCursorSelection((e) => {
+        if (!e.selection.isEmpty()) {
+          lastCopiedRangeRef.current = {
+            startLine: e.selection.startLineNumber,
+            endLine: e.selection.endLineNumber
+          };
+        }
+      });
+
+      return () => disposable.dispose();
+    }
+  }, [editorRef?.current]);
+
   return (
     <div className="flex flex-col h-screen w-full">
       <div className="flex justify-between items-center p-2 border-b">
@@ -73,36 +198,60 @@ export default function AIChat({
             message={message}
             setContext={setContext}
             setIsContextExpanded={setIsContextExpanded}
+            socket={socket}
           />
         ))}
         {isLoading && <LoadingDots />}
       </div>
       <div className="p-4 border-t mb-14">
-        <ContextDisplay
+        <ContextTabs
-          context={context}
+          activeFileName={activeFileName}
-          isContextExpanded={isContextExpanded}
+          onAddFile={handleAddFile}
-          setIsContextExpanded={setIsContextExpanded}
+          contextTabs={contextTabs}
-          setContext={setContext}
+          onRemoveTab={removeContextTab}
+          isExpanded={isContextExpanded}
+          onToggleExpand={() => setIsContextExpanded(!isContextExpanded)}
+          files={files}
+          socket={socket}
+          onFileSelect={(file: TFile) => {
+            socket?.emit("getFile", { fileId: file.id }, (response: string) => {
+              const fileExt = file.name.split('.').pop() || 'txt'
+              const formattedContent = `\`\`\`${fileExt}\n${response}\n\`\`\``
+              addContextTab('file', file.name, formattedContent)
+              if (textareaRef.current) {
+                textareaRef.current.focus()
+              }
+            })
+          }}
         />
         <ChatInput
+          textareaRef={textareaRef}
+          files={[]}
+          addContextTab={addContextTab}
+          editorRef={editorRef}
           input={input}
           setInput={setInput}
           isGenerating={isGenerating}
-          handleSend={() =>
+          handleSend={handleSendWithContext}
-            handleSend(
-              input,
-              context,
-              messages,
-              setMessages,
-              setInput,
-              setIsContextExpanded,
-              setIsGenerating,
-              setIsLoading,
-              abortControllerRef,
-              activeFileContent
-            )
-          }
           handleStopGeneration={() => handleStopGeneration(abortControllerRef)}
+          onImageUpload={(file) => {
+            const reader = new FileReader()
+            reader.onload = (e) => {
+              if (e.target?.result) {
+                addContextTab("image", file.name, e.target.result as string)
+              }
+            }
+            reader.readAsDataURL(file)
+          }}
+          onFileMention={(fileName) => {
+          }}
+          lastCopiedRangeRef={lastCopiedRangeRef}
+          activeFileName={activeFileName}
+          contextTabs={contextTabs.map(tab => ({
+            ...tab,
+            title: tab.id // Add missing title property
+          }))}
+          onRemoveTab={removeContextTab}
         />
       </div>
     </div>

frontend/components/editor/index.tsx

@@ -172,6 +172,9 @@ export default function CodeEditor({
   const editorPanelRef = useRef<ImperativePanelHandle>(null)
   const previewWindowRef = useRef<{ refreshIframe: () => void }>(null)
 
+  // Ref to store the last copied range in the editor to be used in the AIChat component
+  const lastCopiedRangeRef = useRef<{ startLine: number; endLine: number } | null>(null);
+
   const debouncedSetIsSelected = useRef(
     debounce((value: boolean) => {
       setIsSelected(value)
@@ -256,6 +259,17 @@ export default function CodeEditor({
           updatedOptions
         )
       }
+
+      // Store the last copied range in the editor to be used in the AIChat component
+      editor.onDidChangeCursorSelection((e) => {
+        const selection = editor.getSelection();
+        if (selection) {
+          lastCopiedRangeRef.current = {
+            startLine: selection.startLineNumber,
+            endLine: selection.endLineNumber
+          };
+        }
+      });
     }
 
     // Call the function with your file structure
@@ -1219,6 +1233,9 @@ export default function CodeEditor({
                     "No file selected"
                   }
                   onClose={toggleAIChat}
+                  editorRef={{ current: editorRef }}
+                  lastCopiedRangeRef={lastCopiedRangeRef}
+                  files={files}
                 />
               </ResizablePanel>
             </>