feat: frontend - QOL; backend - QOL

.gitignore

@@ -49,3 +49,5 @@ frontend/.vercel
 # typescript
 frontend/*.tsbuildinfo
 frontend/next-env.d.ts
+
+frontend/.copilotignore

backend/app.py

@@ -1,5 +1,8 @@
+import asyncio
 import json
 import logging
+import os
+import time
 from typing import Dict
 
 import socketio
@@ -17,12 +20,7 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 app = FastAPI()
-CORS_ALLOWED_ORIGINS = [
-    "*",
-    "http://localhost:3000",
-    "http://127.0.0.1:3000",
-    "https://knowledge-net.vercel.app",
-]
+CORS_ALLOWED_ORIGINS = os.getenv("ALLOWED_ORIGINS", ",").split(",")
 app.add_middleware(
     CORSMiddleware,
     allow_origins=CORS_ALLOWED_ORIGINS,
@@ -43,6 +41,7 @@ app.mount("/", socketio.ASGIApp(sio))
 class SessionManager:
     def __init__(self):
         self.sessions: Dict[str, tuple[KNet, CrawlForAIScraper]] = {}
+        self.tasks: Dict[str, asyncio.Task] = {}  # Track research tasks for each session
 
     async def get_or_create_session(self, sid: str) -> tuple[KNet, CrawlForAIScraper]:
         if sid not in self.sessions:
@@ -53,11 +52,27 @@ class SessionManager:
         return self.sessions[sid]
 
     async def cleanup_session(self, sid: str):
+        # Cancel running task if it exists
+        if sid in self.tasks and not self.tasks[sid].done():
+            self.tasks[sid].cancel()
+            try:
+                await self.tasks[sid]
+            except asyncio.CancelledError:
+                logger.info(f"Research task for session {sid} was cancelled")
+            except Exception as e:
+                logger.error(f"Error while cancelling task for {sid}: {str(e)}")
+            finally:
+                del self.tasks[sid]
+
+        # Clean up session resources
         if sid in self.sessions:
             _, scraper = self.sessions[sid]
             await scraper.close()
             del self.sessions[sid]
 
+    def register_task(self, sid: str, task: asyncio.Task):
+        self.tasks[sid] = task
+
 
 session_manager = SessionManager()
 
@@ -96,7 +111,13 @@ async def start_research(sid, data):
         async def progress_callback(status: dict):
             await sio.emit("status", status, room=session_id)
 
-        research_results = await knet.conduct_research(topic, progress_callback, max_depth, num_sites_per_query)
+        task = asyncio.create_task(knet.conduct_research(topic, progress_callback, max_depth, num_sites_per_query))
+        session_manager.register_task(sid, task)
+        research_results = await task
+
+        if not research_results:
+            sio.emit("research_aborted", room=session_id)
+
         logger.info(f"Research completed for topic: {topic}")
         await sio.emit("research_complete", research_results, room=session_id)
 
@@ -105,14 +126,40 @@ async def start_research(sid, data):
         await sio.emit("error", {"message": str(e)}, room=session_id)
 
 
+@sio.event
+async def abort_research(sid):
+    logger.info(f"Aborting research for client {sid}")
+    await session_manager.cleanup_session(sid)
+
+
 @sio.event
 async def test(sid, data):
-    knet, _ = await session_manager.get_or_create_session(sid)
-    print("Testing...")
     data = json.loads(data) if type(data) is not dict else data
-    res = await knet.scraper._scrape_page(data["url"])
-    print(json.dumps(res, indent=2))
-    await sio.emit("test", res, room=sid)
+    topic = data.get("topic").strip().replace("\n", "")
+    logger.info(json.dumps(data, indent=2))
+
+    knet, _ = await session_manager.get_or_create_session(sid)
+    time.sleep(1)
+
+    async def progress_callback(status: dict):
+        await sio.emit("status", status, room=sid)
+
+    # Create a task and register it for proper cancellation
+    task = asyncio.create_task(knet.test(topic, progress_callback))
+    session_manager.register_task(sid, task)
+
+    try:
+        await task
+
+        with open("output.log.json", "r") as f:
+            data = json.load(f)
+        await sio.emit("research_complete", data, room=sid)
+    except asyncio.CancelledError:
+        logger.info(f"Test task for '{topic}' was cancelled")
+        await sio.emit("research_aborted", room=sid)
+    except Exception as e:
+        logger.error(f"Test error: {str(e)}")
+        await sio.emit("error", {"message": str(e)}, room=sid)
 
 
 if __name__ == "__main__":

backend/knet.py

@@ -1,6 +1,8 @@
+import asyncio
 import json
 import logging
 import os
+import time
 from collections import deque
 from datetime import datetime
 from textwrap import dedent
@@ -13,9 +15,11 @@ from google.genai import types
 from research_node import ResearchNode
 from scraper import CrawlForAIScraper
 
-# Load environment variables
 load_dotenv()
 
+# Today's Date | Format 15th Dec, 2025
+DATE = datetime.now().strftime("%d %b, %Y")
+
 
 class Prompt:
     def __init__(self) -> None:
@@ -51,7 +55,7 @@ class Prompt:
 
         Return only decision: true/false""")
 
-        self.search_query = dedent("""Based on the following findings on topic {vertical}, suggest new research directions.
+        self.search_query = dedent("""Based on the following findings on topic {vertical}, create google search queries
         Global Research Plan:
         {research_plan}
 
@@ -67,6 +71,8 @@ class Prompt:
         - Explores different aspects
         - Goes deeper into important details
 
+        - Do not do quote searches
+        - Keep the queries short and to the point
         Return as JSON array of objects with properties:
         - query (string)""")
 
@@ -102,6 +108,9 @@ class Prompt:
         Do not include the heading in the content.
         """)
 
+        for prompt in [self.research_plan, self.site_summary, self.continue_branch, self.search_query]:
+            prompt += f"\n\nFYI Date {DATE}"
+
 
 class Schema:
     def __init__(self) -> None:
@@ -177,7 +186,7 @@ class KNet:
         self.ctx_manager: list[str] = []
         self.token_count: int = 0
 
-    async def conduct_research(self, topic: str, progress_callback, max_depth: int, num_sites_per_query: int) -> dict:
+    async def conduct_research(self, topic: str, progress_callback, max_depth: int, num_sites_per_query: int) -> dict | bool:
         # Local Runtime State
         self.progress = ResearchProgress(progress_callback)
         self.max_depth = max_depth
@@ -193,6 +202,8 @@ class KNet:
         try:
             # Generate research plan
             await self.progress.update(0, "Generating research plan...")
+            self._check_cancelled()
+
             self.research_plan = self.generate_content(self.prompt.research_plan.format(topic=topic), schema=self.schema.research_plan, temp=1.5)[
                 "steps"
             ]
@@ -204,6 +215,8 @@ class KNet:
 
             # Iterate on research plan
             for self.idx_research_plan, _ in enumerate(self.research_plan):
+                self._check_cancelled()
+
                 # Generate initial search query
                 query = self.generate_content(
                     self.prompt.search_query.format(
@@ -221,6 +234,8 @@ class KNet:
                 await self.progress.update(100 / (len(self.research_plan) + 1), f"{self.research_plan[self.idx_research_plan]}")
 
                 while to_explore:
+                    self._check_cancelled()
+
                     current_node, current_depth = to_explore.popleft()
                     if current_depth > self.max_depth:
                         continue
@@ -242,6 +257,8 @@ class KNet:
                             for branch in new_branches:
                                 to_explore.appendleft((branch, current_depth + 1))
 
+            self._check_cancelled()
+
             # Generate final report
             await self.progress.update(100 / (len(self.research_plan) + 1), "Generating final report...")
             final_report = await self._generate_final_report(master_node, topic)
@@ -253,24 +270,38 @@ class KNet:
                 json.dump(final_report, f, indent=2)
             return final_report
 
+        except asyncio.CancelledError:
+            self.logger.info(f"Research task for topic '{topic}' was cancelled")
+            return {"status": False}
         except Exception:
             self.logger.error("Research failed", exc_info=True)
             raise
 
+    def _check_cancelled(self):
+        """Check if the current task has been cancelled and raise CancelledError if so"""
+        if asyncio.current_task() and asyncio.current_task().cancelled():
+            raise asyncio.CancelledError("Research task was cancelled")
+
     async def _generate_final_report(self, root_node: ResearchNode, topic: str, retry_count: int = 1) -> Dict[str, Any]:
         try:
+            self._check_cancelled()
+
             await self.progress.setter(0, "Generating report...")
             findings = "\n\n------\n\n".join(self.ctx_manager)
             with open("ctx_manager.log.txt", "w", encoding="utf-8") as f:
                 f.write(findings)
 
             # Generate report outline
+            self._check_cancelled()
             outline = self.generate_content(self.prompt.report_outline.format(topic=topic, ctx_manager=findings), schema=self.schema.report_outline)
             self.logger.info(f"Report outline:\n{json.dumps(outline, indent=2)}")
             report = []
             raster_report = f"# {outline['title']}\n\n"
+
             # Fill in report outline
             for i, heading in enumerate(outline["headings"]):
+                self._check_cancelled()
+
                 await self.progress.update(100 / (len(outline["headings"]) + 1), "Generating report...")
                 content = self.generate_content(
                     self.prompt.report_fillin.format(
@@ -290,7 +321,7 @@ class KNet:
                 raster_report += f"\n\n## {heading}\n\n{content}"
 
             # Collate multimedia content
-            media_content = {"images": [], "videos": [], "links": [], "references": []}
+            media_content = {"images": [], "videos": [], "links": []}
             all_sources_data = root_node.get_all_data()
             for data in all_sources_data:
                 if data.get("images"):
@@ -331,12 +362,14 @@ class KNet:
                 },
             }
 
+        except asyncio.CancelledError:
+            raise
         except Exception as e:
             if e in ["GEMINI_RECITATION", "NO_RESPONSE"]:
                 self.logger.error("GEMINI_RECITATION or NO_RESPONSE")
             if retry_count < 3:
                 self.logger.error(f"Retrying final report:C:{retry_count} / 3", exc_info=True)
-                return await self._generate_final_report(root_node, retry_count + 1)
+                return await self._generate_final_report(root_node, topic, retry_count + 1)
             self.logger.error("Error generating final report", exc_info=True)
             raise
 
@@ -408,7 +441,7 @@ class KNet:
             self.logger.error("Branch decision failed:", exc_info=True)
             raise
 
-    def generate_content(self, prompt: str, schema: Dict[str, Any] = {}, temp: float = 1, _retry_count: int = 1) -> Dict[str, Any] | str:
+    def generate_content(self, prompt: str, schema: Dict[str, Any] = {}, temp: float = 1) -> Dict[str, Any] | str:
         safe = [
             types.SafetySetting(category=types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, threshold=types.HarmBlockThreshold.BLOCK_NONE),
             types.SafetySetting(category=types.HarmCategory.HARM_CATEGORY_HARASSMENT, threshold=types.HarmBlockThreshold.BLOCK_NONE),
@@ -435,3 +468,27 @@ class KNet:
             if response.candidates[0].finish_reason == types.FinishReason.RECITATION:
                 raise Exception("GEMINI_RECITATION")
             raise
+
+    async def test(self, topic: str, progress_callback):
+        self.progress = ResearchProgress(progress_callback)
+        try:
+            for i in range(5):
+                self._check_cancelled()
+
+                await self.progress.setter(i * 10, f"Researching {topic} {i * 10}%")
+                time.sleep(1)
+                for j in range(5):
+                    self._check_cancelled()
+
+                    await self.progress.setter(i * 10, f"s_ example google search {str(j)}")
+                    time.sleep(1)
+
+            for i in range(10):
+                self._check_cancelled()
+
+                await self.progress.setter(i * 10, "Generating report...")
+                time.sleep(1)
+
+        except asyncio.CancelledError:
+            self.logger.info(f"Test task for '{topic}' was cancelled")
+            raise

backend/scraper.py

@@ -158,9 +158,10 @@ class WebScraper:
 class CrawlForAIScraper:
     def __init__(self) -> None:
         self.logger = logging.getLogger(__name__)
+        self.session = requests.Session()
         self.base_browser = BrowserConfig(
             browser_type="chromium",
-            headless=True,
+            headless=False,
             viewport_width=1920,
             viewport_height=1080,
             accept_downloads=False,
@@ -216,7 +217,7 @@ class CrawlForAIScraper:
                 scan_full_page=True,
             )
 
-            soup = BeautifulSoup(result.cleaned_html, "html.parser")
+            soup = BeautifulSoup(result.html, "html.parser")
             search_results = []
 
             for link in list(soup.select("div > span > a"))[2:]:
@@ -227,16 +228,47 @@ class CrawlForAIScraper:
                     continue
                 search_results.append(url)
 
+            if not search_results:
+                self.logger.warning("No search results found.")
+                self.logger.info("Performing DuckDuckGo search as fallback...")
+                search_results = self._duckduckgo_search(query)
+
             self.logger.info(f"Found {len(search_results)} results")
             return search_results
 
-        except requests.exceptions.RequestException as e:
-            self.logger.error(f"Google search error: {str(e)}", exc_info=True)
-            raise
         except Exception as e:
             self.logger.error(f"Google search error: {str(e)}", exc_info=True)
             raise
 
+    def _duckduckgo_search(self, query: str) -> List[str]:
+        self.logger.info("Performing DuckDuckGo search...")
+        try:
+            encoded_query = quote_plus(query)
+            url = f"https://html.duckduckgo.com/html/?q={encoded_query}"
+
+            response = self.session.get(url, headers=self.headers, timeout=self.timeout)
+            response.raise_for_status()
+
+            soup = BeautifulSoup(response.text, "html.parser")
+            search_results = []
+
+            # DuckDuckGo search results are in elements with class 'result__url'
+            for result in soup.select(".result__url"):
+                url = result.get("href").replace(" ", "").replace("\\n", "")
+                if not url.startswith(("http://", "https://")):
+                    url = "https://" + url
+                search_results.append(url)
+
+            self.logger.info(f"Found {len(search_results)} URLs")
+            return search_results
+
+        except requests.exceptions.RequestException as e:  # Catch network errors specifically
+            self.logger.error(f"DuckDuckGo search error: {str(e)}")
+            return []
+        except Exception as e:  # Catch any other errors
+            self.logger.error(f"DuckDuckGo search error: {str(e)}")
+            return []
+
     async def _scrape_pages(self, urls: str, max_sites: int) -> Dict[str, Any]:
         await self.start()
 
@@ -289,7 +321,7 @@ class CrawlForAIScraper:
                     }
                     scraped_sites.append(data)
                     self.logger.info(f"  - {result.url[:80]}...")
-            return scraped_sites[: max_sites]
+            return scraped_sites[:max_sites]
 
         except Exception as e:
             self.logger.error(f"Scraping error while {urls}: {str(e)}")

frontend/bun.lock

@@ -12,7 +12,7 @@
         "@radix-ui/react-scroll-area": "^1.2.3",
         "@radix-ui/react-select": "^2.1.6",
         "@radix-ui/react-separator": "^1.1.2",
-        "@radix-ui/react-slot": "^1.1.2",
+        "@radix-ui/react-slot": "^1.2.0",
         "@radix-ui/react-tabs": "^1.1.3",
         "@radix-ui/react-tooltip": "^1.1.8",
         "@types/react-syntax-highlighter": "^15.5.13",
@@ -296,7 +296,7 @@
 
     "@radix-ui/react-separator": ["@radix-ui/react-separator@1.1.2", "", { "dependencies": { "@radix-ui/react-primitive": "2.0.2" }, "peerDependencies": { "@types/react": "*", "@types/react-dom": "*", "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc", "react-dom": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" }, "optionalPeers": ["@types/react", "@types/react-dom"] }, "sha512-oZfHcaAp2Y6KFBX6I5P1u7CQoy4lheCGiYj+pGFrHy8E/VNRb5E39TkTr3JrV520csPBTZjkuKFdEsjS5EUNKQ=="],
 
-    "@radix-ui/react-slot": ["@radix-ui/react-slot@1.1.2", "", { "dependencies": { "@radix-ui/react-compose-refs": "1.1.1" }, "peerDependencies": { "@types/react": "*", "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" }, "optionalPeers": ["@types/react"] }, "sha512-YAKxaiGsSQJ38VzKH86/BPRC4rh+b1Jpa+JneA5LRE7skmLPNAyeG8kPJj/oo4STLvlrs8vkf/iYyc3A5stYCQ=="],
+    "@radix-ui/react-slot": ["@radix-ui/react-slot@1.2.0", "", { "dependencies": { "@radix-ui/react-compose-refs": "1.1.2" }, "peerDependencies": { "@types/react": "*", "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" }, "optionalPeers": ["@types/react"] }, "sha512-ujc+V6r0HNDviYqIK3rW4ffgYiZ8g5DEHrGJVk4x7kTlLXRDILnKX9vAUYeIsLOoDpDJ0ujpqMkjH4w2ofuo6w=="],
 
     "@radix-ui/react-tabs": ["@radix-ui/react-tabs@1.1.3", "", { "dependencies": { "@radix-ui/primitive": "1.1.1", "@radix-ui/react-context": "1.1.1", "@radix-ui/react-direction": "1.1.0", "@radix-ui/react-id": "1.1.0", "@radix-ui/react-presence": "1.1.2", "@radix-ui/react-primitive": "2.0.2", "@radix-ui/react-roving-focus": "1.1.2", "@radix-ui/react-use-controllable-state": "1.1.0" }, "peerDependencies": { "@types/react": "*", "@types/react-dom": "*", "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc", "react-dom": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" }, "optionalPeers": ["@types/react", "@types/react-dom"] }, "sha512-9mFyI30cuRDImbmFF6O2KUJdgEOsGh9Vmx9x/Dh9tOhL7BngmQPQfwW4aejKm5OHpfWIdmeV6ySyuxoOGjtNng=="],
 
@@ -1476,6 +1476,20 @@
 
     "@next/eslint-plugin-next/fast-glob": ["fast-glob@3.3.1", "", { "dependencies": { "@nodelib/fs.stat": "^2.0.2", "@nodelib/fs.walk": "^1.2.3", "glob-parent": "^5.1.2", "merge2": "^1.3.0", "micromatch": "^4.0.4" } }, "sha512-kNFPyjhh5cKjrUltxs+wFx+ZkbRaxxmZ+X0ZU31SOsxCEtP9VPgtq2teZw1DebupL5GmDaNQ6yKMMVcM41iqDg=="],
 
+    "@radix-ui/react-collection/@radix-ui/react-slot": ["@radix-ui/react-slot@1.1.2", "", { "dependencies": { "@radix-ui/react-compose-refs": "1.1.1" }, "peerDependencies": { "@types/react": "*", "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" }, "optionalPeers": ["@types/react"] }, "sha512-YAKxaiGsSQJ38VzKH86/BPRC4rh+b1Jpa+JneA5LRE7skmLPNAyeG8kPJj/oo4STLvlrs8vkf/iYyc3A5stYCQ=="],
+
+    "@radix-ui/react-dialog/@radix-ui/react-slot": ["@radix-ui/react-slot@1.1.2", "", { "dependencies": { "@radix-ui/react-compose-refs": "1.1.1" }, "peerDependencies": { "@types/react": "*", "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" }, "optionalPeers": ["@types/react"] }, "sha512-YAKxaiGsSQJ38VzKH86/BPRC4rh+b1Jpa+JneA5LRE7skmLPNAyeG8kPJj/oo4STLvlrs8vkf/iYyc3A5stYCQ=="],
+
+    "@radix-ui/react-menu/@radix-ui/react-slot": ["@radix-ui/react-slot@1.1.2", "", { "dependencies": { "@radix-ui/react-compose-refs": "1.1.1" }, "peerDependencies": { "@types/react": "*", "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" }, "optionalPeers": ["@types/react"] }, "sha512-YAKxaiGsSQJ38VzKH86/BPRC4rh+b1Jpa+JneA5LRE7skmLPNAyeG8kPJj/oo4STLvlrs8vkf/iYyc3A5stYCQ=="],
+
+    "@radix-ui/react-primitive/@radix-ui/react-slot": ["@radix-ui/react-slot@1.1.2", "", { "dependencies": { "@radix-ui/react-compose-refs": "1.1.1" }, "peerDependencies": { "@types/react": "*", "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" }, "optionalPeers": ["@types/react"] }, "sha512-YAKxaiGsSQJ38VzKH86/BPRC4rh+b1Jpa+JneA5LRE7skmLPNAyeG8kPJj/oo4STLvlrs8vkf/iYyc3A5stYCQ=="],
+
+    "@radix-ui/react-select/@radix-ui/react-slot": ["@radix-ui/react-slot@1.1.2", "", { "dependencies": { "@radix-ui/react-compose-refs": "1.1.1" }, "peerDependencies": { "@types/react": "*", "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" }, "optionalPeers": ["@types/react"] }, "sha512-YAKxaiGsSQJ38VzKH86/BPRC4rh+b1Jpa+JneA5LRE7skmLPNAyeG8kPJj/oo4STLvlrs8vkf/iYyc3A5stYCQ=="],
+
+    "@radix-ui/react-slot/@radix-ui/react-compose-refs": ["@radix-ui/react-compose-refs@1.1.2", "", { "peerDependencies": { "@types/react": "*", "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" }, "optionalPeers": ["@types/react"] }, "sha512-z4eqJvfiNnFMHIIvXP3CY57y2WJs5g2v3X0zm9mEJkrkNv4rDxu+sg9Jh8EkXyeqBkB7SOcboo9dMVqhyrACIg=="],
+
+    "@radix-ui/react-tooltip/@radix-ui/react-slot": ["@radix-ui/react-slot@1.1.2", "", { "dependencies": { "@radix-ui/react-compose-refs": "1.1.1" }, "peerDependencies": { "@types/react": "*", "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" }, "optionalPeers": ["@types/react"] }, "sha512-YAKxaiGsSQJ38VzKH86/BPRC4rh+b1Jpa+JneA5LRE7skmLPNAyeG8kPJj/oo4STLvlrs8vkf/iYyc3A5stYCQ=="],
+
     "@ts-morph/common/minimatch": ["minimatch@7.4.6", "", { "dependencies": { "brace-expansion": "^2.0.1" } }, "sha512-sBz8G/YjVniEz6lKPNpKxXwazJe4c19fEfV2GDMX6AjFz+MX9uDWIZW8XreVhkFW3fkIdTv/gxWr/Kks5FFAVw=="],
 
     "@typescript-eslint/parser/debug": ["debug@4.4.0", "", { "dependencies": { "ms": "^2.1.3" } }, "sha512-6WTZ/IxCY/T6BALoZHaE4ctp9xm+Z5kY/pzYaCHRFeyVhojxlrm+46y68HA6hr0TcwEssoxNiDEUJQjfPZ/RYA=="],

frontend/package.json

@@ -17,7 +17,7 @@
     "@radix-ui/react-scroll-area": "^1.2.3",
     "@radix-ui/react-select": "^2.1.6",
     "@radix-ui/react-separator": "^1.1.2",
-    "@radix-ui/react-slot": "^1.1.2",
+    "@radix-ui/react-slot": "^1.2.0",
     "@radix-ui/react-tabs": "^1.1.3",
     "@radix-ui/react-tooltip": "^1.1.8",
     "@types/react-syntax-highlighter": "^15.5.13",

frontend/src/app/layout.tsx

@@ -1,4 +1,5 @@
 import { ThemeProvider } from "@/components/theme-provider";
+import { ChatProvider } from "@/lib/store/ChatContext";
 import type { Metadata } from "next";
 import "./globals.css";
 
@@ -12,7 +13,7 @@ export default function RootLayout({ children }: { children: React.ReactNode })
     <html lang="en" suppressHydrationWarning>
       <body>
         <ThemeProvider attribute="class" defaultTheme="system" enableSystem disableTransitionOnChange>
-          {children}
+          <ChatProvider>{children}</ChatProvider>
         </ThemeProvider>
       </body>
     </html>

frontend/src/app/page.tsx

@@ -1,4 +1,4 @@
-import ChatInterface from '@/components/ChatInterface';
+import ChatInterface from "@/components/ChatInterface";
 
 export default function Home() {
   return <ChatInterface />;

frontend/src/components/ChatHistory.tsx

@@ -1,8 +1,9 @@
 "use client";
 import { ScrollArea } from "@/components/ui/scroll-area";
+import { useChatContext } from "@/lib/store/ChatContext";
 import { Message as MessageType } from "@/lib/types";
-import { Loader2 } from "lucide-react";
-import React, { useEffect, useRef } from "react";
+import { v4 as uuidv4 } from "uuid";
+import React, { useEffect, useRef, useState } from "react";
 import Message from "./Message";
 
 interface ChatHistoryProps {
@@ -10,54 +11,39 @@ interface ChatHistoryProps {
   isLoading: boolean;
 }
 
+// Traditional prop-based component
 const ChatHistory: React.FC<ChatHistoryProps> = ({ messages, isLoading }) => {
   const messagesEndRef = useRef<HTMLDivElement>(null);
+  const [lastProgressMessage, setLastProgressMessage] = useState<string>("");
 
   useEffect(() => {
     messagesEndRef.current?.scrollIntoView({ behavior: "smooth" });
   }, [messages]);
 
   return (
-    <ScrollArea tabIndex={-1}>
-      <div className="pb-24 focus-visible:outline-0" tabIndex={-1}>
-        {messages.length === 0 ? (
-          <div className="h-full flex flex-col items-center justify-center p-8 text-center">
-            <div className="max-w-md space-y-2">
-              <h2 className="text-2xl font-bold">Welcome to the KNet demo</h2>
-              <p className="text-muted-foreground">Ask any research question to get started. The assistant will provide detailed answers backed by research and sources.</p>
-            </div>
+    <div className="pb-24 overflow-auto focus-visible:outline-0" tabIndex={-1}>
+      {messages.length === 0 ? (
+        <div className="h-full flex flex-col items-center justify-center p-8 text-center">
+          <div className="max-w-md space-y-2">
+            <h2 className="text-2xl font-bold">Welcome to the KNet demo</h2>
+            <p className="text-muted-foreground">Ask any research question to get started. The assistant will provide detailed answers backed by research and sources.</p>
           </div>
-        ) : (
-          messages.map((message) => <Message key={message.id} message={message} />)
-        )}
+        </div>
+      ) : (
+        messages.map((message) => {
+          return <Message key={message.id} message={message} lastProgressMessage={lastProgressMessage} setLastProgressMessage={setLastProgressMessage} />;
+        })
+      )}
 
-        {isLoading && (
-          <div className="my-2 mx-4">
-            <div className="max-w-2xl mx-auto flex gap-4 relative">
-              <div className="h-8 w-8 rounded-full shrink-0 bg-muted flex items-center justify-center absolute -left-12 top-0">
-                <Loader2 className="h-5 w-5 animate-spin" />
-              </div>
-              <div className="flex-1">
-                <div className="flex items-center gap-2 mb-1">
-                  <div className="font-medium">KNet</div>
-                  <div className="text-xs text-muted-foreground">Just now</div>
-                </div>
-                <div className="mt-1 bg-muted/50 p-3 rounded-2xl rounded-tl-sm">
-                  <div className="flex space-x-2">
-                    <div className="w-2 h-2 rounded-full bg-muted-foreground/30 animate-pulse"></div>
-                    <div className="w-2 h-2 rounded-full bg-muted-foreground/30 animate-pulse" style={{ animationDelay: "300ms" }}></div>
-                    <div className="w-2 h-2 rounded-full bg-muted-foreground/30 animate-pulse" style={{ animationDelay: "600ms" }}></div>
-                  </div>
-                </div>
-              </div>
-            </div>
-          </div>
-        )}
-
-        <div ref={messagesEndRef} />
-      </div>
-    </ScrollArea>
+      <div ref={messagesEndRef} />
+    </div>
   );
 };
 
+// Context-based component
+export const ChatHistoryWithContext: React.FC = () => {
+  const { chatState } = useChatContext();
+  return <ChatHistory messages={chatState.messages} isLoading={chatState.isLoading} />;
+};
+
 export default ChatHistory;

frontend/src/components/ChatInterface.tsx

@@ -1,317 +1,17 @@
 "use client";
 import { Alert, AlertDescription, AlertTitle } from "@/components/ui/alert";
-import { disconnectSocket, getSocket, initializeSocket } from "@/lib/socket";
-import { ChatData, ChatState, Conversation, Message, ResearchOptions, ResearchResults, StatusUpdate } from "@/lib/types";
-import { useEffect, useRef, useState } from "react";
-import { v4 as uuidv4 } from "uuid";
+import { useChatContext } from "@/lib/store/ChatContext";
+import React from "react";
 import ChatHistory from "./ChatHistory";
 import MessageInput from "./MessageInput";
 import ResearchControls from "./ResearchControls";
 import ChatLayout from "./ui/ChatLayout";
 import ConversationList from "./ui/ConversationList";
 
-const saveToStorage = (data: ChatData) => {
-  if (typeof window !== "undefined") {
-    localStorage.setItem("chatData", JSON.stringify(data));
-  }
-};
-
-const loadFromStorage = (): ChatData => {
-  if (typeof window === "undefined") {
-    return { conversations: [], currentConversationId: null };
-  }
-  const data = localStorage.getItem("chatData");
-  if (!data) {
-    return { conversations: [], currentConversationId: null };
-  }
-  try {
-    const parsed = JSON.parse(data);
-    return {
-      conversations: Array.isArray(parsed.conversations)
-        ? parsed.conversations.map((conv: Conversation) => ({
-            ...conv,
-            messages: Array.isArray(conv.messages)
-              ? conv.messages.map((msg: Message) => ({
-                  ...msg,
-                  // Ensure media property is preserved if it exists
-                  media: msg.media || undefined,
-                }))
-              : [],
-          }))
-        : [],
-      currentConversationId: parsed.currentConversationId,
-    };
-  } catch (e) {
-    return { conversations: [], currentConversationId: null };
-  }
-};
-
 const ChatInterface = () => {
-  const [chatState, setChatState] = useState<ChatState>({ messages: [], isLoading: false, error: null });
-  const [conversations, setConversations] = useState<Conversation[]>([]);
-  const [currentConversationId, setCurrentConversationId] = useState<string | null>(null);
-
-  const [researchOptions, setResearchOptions] = useState<ResearchOptions>({
-    depth: "basic",
-    sources: true,
-    citations: false,
-    max_depth: 1,
-    num_sites_per_query: 3,
-  });
-
-  const userInputRef = useRef<HTMLTextAreaElement>(null);
-
-  // Add this effect for focus management
-  useEffect(() => {
-    const focusInput = () => {
-      setTimeout(() => {
-        userInputRef.current?.focus();
-      }, 100);
-    };
-
-    focusInput();
-  }, [currentConversationId]);
-
-  // Initialize socket once
-  useEffect(() => {
-    const socket = initializeSocket();
-
-    socket.on("connect", () => {
-      console.log("Connected to research server");
-    });
-
-    socket.on("disconnect", () => {
-      console.log("Disconnected from research server");
-    });
-
-    socket.on("status", (data: StatusUpdate) => {
-      setChatState((prevState) => {
-        const messages = [...prevState.messages];
-        const progressText = `(${data.progress}%) ${data.message}`;
-
-        // Find the last assistant message that is a progress update
-        const lastProgressIndex = messages.findLastIndex((msg) => msg.role === "assistant" && msg.content.includes("%)"));
-
-        if (lastProgressIndex !== -1) {
-          // Update existing progress message
-          messages[lastProgressIndex] = {
-            ...messages[lastProgressIndex],
-            content: progressText,
-          };
-        } else {
-          // Add new progress message
-          messages.push({
-            id: uuidv4(),
-            content: progressText,
-            role: "assistant",
-            timestamp: new Date(),
-          });
-        }
-
-        return {
-          ...prevState,
-          messages,
-          isLoading: true,
-        };
-      });
-    });
-
-    socket.on("research_complete", (results: ResearchResults) => {
-      setChatState((prevState) => {
-        const messages = [...prevState.messages];
-
-        // Remove the last progress message if it exists
-        const lastProgressIndex = messages.findLastIndex((msg) => msg.role === "assistant" && msg.content.includes("%)"));
-        if (lastProgressIndex !== -1) {
-          messages.splice(lastProgressIndex, 1);
-        }
-
-        // Format research stats and response
-        const stats = [`Total Queries: ${results.metadata.total_queries}`, `Sources Used: ${results.metadata.total_sources}`, `Search Depth: ${results.metadata.max_depth_reached}`].join(" | ");
-
-        // Format images in a way the Message component can extract
-        const imageMarkdown = results.media?.images?.length ? `\n\n**Relevant Images:**\n${results.media.images.map((img) => `![Image](${img})`).join("\n")}` : "";
-
-        const formattedResponse = [results.content, `\n\n---\n**Research Stats:**\n${stats}`, imageMarkdown].join("");
-
-        const newMessages = [
-          ...messages,
-          {
-            id: uuidv4(),
-            content: formattedResponse,
-            role: "assistant" as const,
-            timestamp: new Date(results.timestamp),
-            // Store the media object directly with the message
-            media: results.media,
-          },
-        ];
-
-        // Save updated messages to localStorage
-        const updatedState = {
-          ...prevState,
-          isLoading: false,
-          messages: newMessages,
-        };
-
-        // Update localStorage with the new messages
-        const updatedData: ChatData = {
-          conversations: conversations.map((conv) => ({
-            ...conv,
-            messages: conv.id === currentConversationId ? newMessages : conv.messages || [],
-            lastUpdated: conv.id === currentConversationId ? new Date().toISOString() : conv.lastUpdated,
-          })),
-          currentConversationId,
-        };
-        saveToStorage(updatedData);
-
-        return updatedState;
-      });
-    });
-
-    socket.on("error", (error: { message: string }) => {
-      setChatState((prevState) => ({
-        ...prevState,
-        error: error.message,
-        isLoading: false,
-      }));
-    });
-
-    return () => {
-      disconnectSocket();
-    };
-  }, []); // Empty dependency array
-
-  // Load initial data
-  useEffect(() => {
-    const data = loadFromStorage();
-    setConversations(data.conversations);
-    setCurrentConversationId(data.currentConversationId);
-
-    if (data.currentConversationId) {
-      const conversation = data.conversations.find((c) => c.id === data.currentConversationId);
-      if (conversation) {
-        setChatState((prev) => ({ ...prev, messages: conversation.messages }));
-      }
-    }
-  }, []);
-
-  useEffect(() => {
-    const currentConv = conversations.find((c) => c.id === currentConversationId);
-
-    const data: ChatData = {
-      conversations: conversations.map((conv) => ({
-        ...conv,
-        messages: conv.id === currentConversationId ? chatState.messages : conv.messages || [],
-      })),
-      currentConversationId,
-    };
-
-    saveToStorage(data);
-  }, [conversations, currentConversationId, chatState.messages]);
-
-  const handleSendMessage = (content: string) => {
-    if (!content.trim()) return;
-
-    let conversationId = currentConversationId;
-    const newMessage: Message = {
-      id: uuidv4(),
-      content,
-      role: "user",
-      timestamp: new Date(),
-    };
-
-    // Create a new conversation if none exists
-    if (!conversationId) {
-      conversationId = uuidv4();
-      setCurrentConversationId(conversationId);
-      setConversations((prev) => [
-        {
-          id: conversationId as string,
-          title: content.length > 30 ? `${content.substring(0, 30)}...` : content,
-          lastUpdated: new Date().toISOString(),
-          messages: [newMessage],
-          active: true,
-        },
-        ...prev.map((c) => ({ ...c, active: false })),
-      ]);
-    } else {
-      // Update the existing conversation
-      setConversations((prev) =>
-        prev.map((conv) => ({
-          ...conv,
-          lastUpdated: conv.id === conversationId ? new Date().toISOString() : conv.lastUpdated,
-          active: conv.id === conversationId,
-          messages: conv.id === conversationId ? [...(conv.messages || []), newMessage] : conv.messages || [],
-        }))
-      );
-    }
-
-    setChatState((prevState) => ({
-      ...prevState,
-      messages: [...prevState.messages, newMessage],
-      isLoading: true,
-      error: null,
-    }));
-
-    // Send message to server via socket
-    try {
-      const socket = getSocket();
-      socket.emit("start_research", {
-        topic: content,
-        max_depth: researchOptions.max_depth,
-        num_sites_per_query: researchOptions.num_sites_per_query,
-      });
-    } catch (error) {
-      setChatState((prevState) => ({
-        ...prevState,
-        error: "Failed to connect to research server",
-        isLoading: false,
-      }));
-    }
-  };
-
-  const handleNewConversation = () => {
-    userInputRef.current?.focus();
-    setCurrentConversationId(null);
-    setChatState(() => ({
-      messages: [],
-      isLoading: false,
-      error: null,
-    }));
-  };
-
-  const handleSelectConversation = (id: string) => {
-    const data = loadFromStorage();
-    const conversation = data.conversations.find((c) => c.id === id);
-
-    setCurrentConversationId(id);
-    setChatState((prev) => ({
-      ...prev,
-      messages: conversation?.messages || [],
-      isLoading: false,
-      error: null,
-    }));
-    setConversations((prev) =>
-      prev.map((conv) => ({
-        ...conv,
-        active: conv.id === id,
-      }))
-    );
-  };
-
-  const handleDeleteConversation = (id: string) => {
-    setConversations((prev) => prev.filter((conv) => conv.id !== id));
-    if (currentConversationId === id) {
-      handleNewConversation();
-    }
-  };
-
-  const handleDeleteAllConversations = () => {
-    setConversations([]);
-    handleNewConversation();
-  };
+  const { chatState, conversations, abortResearch, userInputRef, researchOptions, sendMessage, newConversation, selectConversation, deleteConversation, deleteAllConversations, setResearchOptions } = useChatContext();
 
-  const sidebar = <ConversationList conversations={conversations} onNewConversation={handleNewConversation} onSelectConversation={handleSelectConversation} onDeleteConversation={handleDeleteConversation} onDeleteAllConversations={handleDeleteAllConversations} />;
+  const sidebar = <ConversationList conversations={conversations} onNewConversation={newConversation} onSelectConversation={selectConversation} onDeleteConversation={deleteConversation} onDeleteAllConversations={deleteAllConversations} />;
 
   const mainContent = (
     <div className="flex flex-col w-full h-full relative" tabIndex={-1}>
@@ -324,7 +24,7 @@ const ChatInterface = () => {
         </Alert>
       )}
 
-      <MessageInput onSendMessage={handleSendMessage} isLoading={chatState.isLoading} userInputRef={userInputRef} />
+      <MessageInput onSendMessage={sendMessage} isLoading={chatState.isLoading} userInputRef={userInputRef} onCancel={abortResearch} />
     </div>
   );
 

frontend/src/components/Message.tsx

@@ -3,13 +3,125 @@ import { ScrollArea } from "@/components/ui/scroll-area";
 import { Button } from "@/components/ui/button";
 import { DropdownMenu, DropdownMenuContent, DropdownMenuItem, DropdownMenuTrigger } from "@/components/ui/dropdown-menu";
 import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip";
-import { Message as MessageType } from "@/lib/types";
-import { Bot, Copy, MoreHorizontal, User2 } from "lucide-react";
-import React, { useState, useEffect } from "react";
+import { Message as MessageType, ResearchTree } from "@/lib/types";
+import { Bot, Copy, MoreHorizontal, User2, ExternalLink, Loader2, SearchIcon } from "lucide-react";
+import React, { useState, useEffect, useMemo, useRef } from "react";
 import ReactMarkdown from "react-markdown";
 import remarkGfm from "remark-gfm";
 import { Prism as SyntaxHighlighter } from "react-syntax-highlighter";
 import { oneDark } from "react-syntax-highlighter/dist/cjs/styles/prism";
+import { Badge } from "@/components/ui/badge";
+
+// Function to extract all source URLs from the research tree
+const extractAllSources = (tree: ResearchTree | undefined): Array<{ text: string; url: string }> => {
+  if (!tree) return [];
+
+  // Start with an empty set to avoid duplicates
+  const uniqueSources = new Set<string>();
+
+  // Recursive function to gather all sources
+  const collectSources = (node: ResearchTree) => {
+    // Add all sources from the current node
+    if (node.sources && Array.isArray(node.sources)) {
+      node.sources.forEach((url) => uniqueSources.add(url));
+    }
+
+    // Process all children recursively
+    if (node.children && Array.isArray(node.children)) {
+      node.children.forEach((child) => collectSources(child));
+    }
+  };
+
+  // Start the collection process
+  collectSources(tree);
+
+  // Convert the set to an array of objects with text and url properties
+  return Array.from(uniqueSources).map((url) => {
+    // Try to extract a readable title from the URL
+    let text = "";
+    try {
+      const urlObj = new URL(url);
+      // Remove 'www.' if present and take the hostname
+      text = urlObj.hostname.replace(/^www\./, "");
+
+      // Add the pathname if it's not just "/"
+      if (urlObj.pathname && urlObj.pathname !== "/") {
+        // Format the pathname - keep it short and clean
+        const path = urlObj.pathname.split("/").filter(Boolean);
+        if (path.length > 0) {
+          const lastPathSegment = path[path.length - 1]
+            .replace(/[-_]/g, " ") // Replace dashes and underscores with spaces
+            .replace(/\.html$|\.pdf$|\.php$/, ""); // Remove common extensions
+
+          text = `${text} - ${lastPathSegment}`;
+        }
+      }
+    } catch (e) {
+      // If URL parsing fails, use the URL as is
+      text = url;
+    }
+
+    return { text, url };
+  });
+};
+
+// SourceLinks component for displaying research sources in a scrollable container
+const SourceLinks = ({ links }: { links: Array<{ text: string; url: string }> }) => {
+  if (!links || links.length === 0) return null;
+
+  return (
+    <div className="mt-4 mb-4">
+      <h3 className="text-md font-semibold mb-2">Research Sources:</h3>
+      <ScrollArea className="w-full h-[300px] overflow-auto border border-slate-200 dark:border-slate-700 rounded-xl shadow-sm p-1" type="always">
+        <div className="space-y-2">
+          {links.map((link, index) => {
+            // Extract domain for display
+            let domain = "";
+            try {
+              const urlObj = new URL(link.url);
+              domain = urlObj.hostname.replace(/^www\./, "");
+            } catch (e) {
+              domain = "Unknown source";
+            }
+
+            return (
+              <div key={index} className="flex items-start gap-2 group hover:bg-muted/50 p-2 rounded-md transition-colors">
+                <ExternalLink className="h-4 w-4 text-muted-foreground flex-shrink-0 mt-1" />
+                <div className="flex-1">
+                  <a href={link.url} className="text-primary hover:underline text-sm block" target="_blank" rel="noopener noreferrer">
+                    {domain}
+                  </a>
+                  <a href={link.url} className="text-xs text-muted-foreground hover:underline block truncate" target="_blank" rel="noopener noreferrer" title={link.url}>
+                    {link.url}
+                  </a>
+                </div>
+                <div className="opacity-0 group-hover:opacity-100 transition-opacity">
+                  <TooltipProvider>
+                    <Tooltip>
+                      <TooltipTrigger asChild>
+                        <Button
+                          size="sm"
+                          variant="ghost"
+                          className="h-6 w-6 p-0"
+                          onClick={(e) => {
+                            e.preventDefault();
+                            navigator.clipboard.writeText(link.url);
+                          }}>
+                          <Copy className="h-3 w-3" />
+                        </Button>
+                      </TooltipTrigger>
+                      <TooltipContent>Copy link</TooltipContent>
+                    </Tooltip>
+                  </TooltipProvider>
+                </div>
+              </div>
+            );
+          })}
+        </div>
+      </ScrollArea>
+    </div>
+  );
+};
 
 // ImageGallery component for handling images in a scrollable container
 const ImageGallery = ({ imageUrls }: { imageUrls: string[] }) => {
@@ -45,12 +157,12 @@ const ImageGallery = ({ imageUrls }: { imageUrls: string[] }) => {
   return (
     <div className="mt-4 mb-4">
       <h3 className="text-md font-semibold mb-2">Relevant Images:</h3>
-      <ScrollArea className="w-full max-h-72 rounded-md border">
-        <div className="p-2 space-y-3">
+      <ScrollArea className="w-full h-[300px] overflow-auto border border-slate-200 dark:border-slate-700 rounded-xl shadow-sm p-1" type="always">
+        <div className="p-2 grid grid-cols-2 md:grid-cols-4 gap-2">
           {imageUrls.map((url, index) => (
-            <div key={index} className="image-container">
+            <div key={index} className="image-container h-[150px]">
               <img
-                className="lazy-image rounded-md max-w-full"
+                className="lazy-image rounded-md w-full h-full object-cover shadow-sm border border-slate-100 dark:border-slate-800"
                 src="data:image/svg+xml,%3Csvg xmlns='http://www.w3.org/2000/svg' width='100' height='100' viewBox='0 0 100 100'%3E%3Crect width='100' height='100' fill='%23f1f1f1'/%3E%3C/svg%3E"
                 data-src={url}
                 alt={`Research image ${index + 1}`}
@@ -112,65 +224,90 @@ const MarkdownComponents: Record<string, React.ComponentType<any>> = {
   ),
   blockquote: ({ children }) => <blockquote className="border-l-4 border-border pl-4 italic my-4">{children}</blockquote>,
   table: ({ children }) => (
-    <div className="overflow-x-auto border rounded-2xl my-4">
-      <table className="w-full border-collapse rounded-2xl overflow-hidden shadow-sm">{children}</table>
+    <div className="overflow-x-scroll border rounded-2xl my-4">
+      <table className="w-max border-collapse rounded-2xl overflow-hidden shadow-sm">{children}</table>
     </div>
   ),
   thead: ({ children }) => <thead className="bg-muted">{children}</thead>,
   tbody: ({ children }) => <tbody>{children}</tbody>,
-  th: ({ children }) => <th className="border-r last:border-r-0 border-slate-900 px-4 py-2 text-left font-semibold">{children}</th>,
+  th: ({ children }) => <th className="border-r last:border-r-0 border-slate-900 px-2 py-1 text-left font-semibold">{children}</th>,
   tr: ({ children }) => <tr className="border-b last:border-b-0 border-border">{children}</tr>,
-  td: ({ children }) => <td className="border-r last:border-r-0 border-border px-4 py-2">{children}</td>,
+  td: ({ children }) => <td className="border-r last:border-r-0 border-border px-2 py-1">{children}</td>,
 };
 
-const Message = ({ message }: { message: MessageType }) => {
-  const isUser = message.role === "user";
-  const isProgress = message.content.includes("%)") && message.role === "assistant";
+interface MessageProps {
+  message: MessageType;
+  isLoading?: boolean;
+  lastProgressMessage?: string;
+  setLastProgressMessage?: React.Dispatch<React.SetStateAction<string>>;
+}
+
+const Message = ({ message, isLoading, lastProgressMessage, setLastProgressMessage }: MessageProps) => {
+  const isUser = message?.role === "user";
   const [imageUrls, setImageUrls] = useState<string[]>([]);
+  const progressPercentage = message.progress || 0;
+  const isProgressMessage = message.isProgress === true;
+  const [isSearchMessage, setIsSearchMessage] = useState(message.content.startsWith("s_"));
+
+  // Use useMemo to extract sources only once and only when they change
+  const sourceLinks = useMemo(() => {
+    // If this is the loading component, return empty sources
+    if (isLoading) return [];
+
+    // First, extract sources from the research_tree if available
+    const researchSources = extractAllSources(message.research_tree);
+
+    // If research_tree sources exist, use those
+    if (researchSources.length > 0) {
+      return researchSources;
+    }
+
+    // Otherwise, fall back to any links in the media object
+    return message.media?.links || [];
+  }, [message.research_tree, message.media?.links, isLoading]);
 
   // Extract image URLs from the message content or use the media object
   useEffect(() => {
-    if (!isUser && !isProgress) {
+    if (isLoading) {
+      setImageUrls([]);
+      return;
+    }
+
+    if (!isUser) {
+      // Handle image URLs
       let urls: string[] = [];
 
       // First, check if there's a media object with images
       if (message.media?.images && message.media.images.length > 0) {
         urls = message.media.images;
-      } else {
-        // Fallback to extracting from markdown content
-        const imgRegex = /!\[.*?\]\((.*?)\)/g;
-        let match;
-
-        while ((match = imgRegex.exec(message.content)) !== null) {
-          if (match[1]) {
-            urls.push(match[1]);
-          }
-        }
       }
-
       setImageUrls(urls);
     }
-  }, [message.content, message.media, isUser, isProgress]);
 
-  // Prepare message content without image markdown
-  const cleanContent = message.content.replace(/!\[.*?\]\(.*?\)\n?/g, "").replace(/\*\*Relevant Images:\*\*\n/g, "");
+    setIsSearchMessage(message.content.startsWith("s_"));
+    if ((isLoading || isProgressMessage) && !message.content.startsWith("s_")) {
+      setLastProgressMessage!(message.content);
+    }
+  }, [message.content, message.media, isUser, isLoading, isProgressMessage, setLastProgressMessage]);
 
   const copyToClipboard = () => {
-    navigator.clipboard.writeText(message.content);
+    if (!isLoading) {
+      navigator.clipboard.writeText(message.content);
+    }
   };
 
   return (
-    <div className="py-2 px-4">
-      <div className={`max-w-2xl mx-auto flex gap-4 relative ${isUser ? "flex-row-reverse" : ""}`}>
-        <Avatar className={`h-8 w-8 shrink-0 absolute justify-center item-center ${isUser ? "-right-12" : "-left-12"} top-0`}>{isUser ? <User2 className="h-5 w-5" /> : <Bot className="h-5 w-5" />}</Avatar>
+    <div className="py-2 px-2 sm:px-4 sm:mx-12">
+      <div className={`w-full flex gap-4 relative ${isUser ? "justify-end" : "justify-start"}`}>
+        <Avatar className={`h-8 w-8 rounded-full bg-muted flex justify-center item-center absolute ${isUser ? "right-0 sm:-right-12" : "left-0 sm:-left-12"} top-0 hidden sm:flex`}>{isLoading || isProgressMessage ? <Loader2 className="h-full w-6 animate-spin" /> : isUser ? <User2 className="h-full w-6" /> : <Bot className="h-full w-6" />}</Avatar>
 
-        <div className={`max-w-2xl flex-1 ${isUser ? "items-end" : "items-start"}`}>
+        <div className={`max-w-full ${isLoading || isProgressMessage ? "w-[80%]" : ""} ${isUser ? "items-end ml-auto" : "items-start mr-auto"}`}>
           <div className={`flex items-center gap-2 mb-1 ${isUser ? "justify-end" : "justify-start"}`}>
-            {isUser && <div className="text-xs text-muted-foreground">{new Date(message.timestamp).toLocaleTimeString()}</div>}
             <div className="font-medium">{isUser ? "You" : "KNet"}</div>
-            {!isUser && <div className="text-xs text-muted-foreground">{new Date(message.timestamp).toLocaleTimeString()}</div>}
+            {!isUser && !isLoading && <div className="text-xs text-muted-foreground">{new Date(message.timestamp).toLocaleTimeString()}</div>}
+            {isLoading && <div className="text-xs text-muted-foreground">Just now</div>}
 
-            {!isUser && (
+            {!isUser && !isLoading && !isProgressMessage && (
               <div className="ml-auto flex items-center gap-2">
                 <TooltipProvider>
                   <Tooltip>
@@ -199,12 +336,37 @@ const Message = ({ message }: { message: MessageType }) => {
             )}
           </div>
 
-          <div className={`mt-1 max-w-none ${isUser ? "bg-slate-300 dark:bg-slate-200 dark:text-background text-foreground" : isProgress ? "bg-muted/30" : "bg-muted/50"} p-3 rounded-2xl ${isUser ? "rounded-tr-sm" : "rounded-tl-sm"}`} style={{ overflowWrap: "anywhere" }}>
-            <ReactMarkdown remarkPlugins={[remarkGfm]} components={MarkdownComponents}>
-              {cleanContent}
-            </ReactMarkdown>
+          <div className={`mt-1 w-full ${isUser ? "bg-slate-300 dark:bg-slate-200 dark:text-background text-foreground" : "bg-muted/50"} p-3 rounded-2xl ${isUser ? "rounded-tr-sm" : "rounded-tl-sm"}`} style={{ overflowWrap: "anywhere" }}>
+            {isLoading || isProgressMessage ? (
+              <div className="space-y-2">
+                <div>{lastProgressMessage}</div>
+                <div className={`flex items-center ${isSearchMessage ? "justify-between" : "justify-end"} text-sm w-full`}>
+                  {isSearchMessage ? (
+                    <Badge variant="outline" className="bg-primary/20 text-primary rounded-full">
+                      <SearchIcon className="h-3 w-3 mr-1" />
+                      {message.content.slice(2)}
+                    </Badge>
+                  ) : null}
+                  <Badge variant="outline" className="ml-2 bg-primary/20 text-primary">
+                    {progressPercentage}%
+                  </Badge>
+                </div>
+                <div className="w-full bg-muted/30 h-2.5 rounded-full overflow-hidden">
+                  <div className="h-full bg-primary transition-all duration-500 rounded-full flex items-center justify-end" style={{ width: `${progressPercentage}%` }}>
+                    <div className="h-2 w-2 rounded-full bg-primary-foreground mr-0.5 animate-pulse"></div>
+                  </div>
+                </div>
+              </div>
+            ) : (
+              <>
+                <ReactMarkdown remarkPlugins={[remarkGfm]} components={MarkdownComponents}>
+                  {message.content}
+                </ReactMarkdown>
 
-            {imageUrls.length > 0 && <ImageGallery imageUrls={imageUrls} />}
+                {sourceLinks.length > 0 && <SourceLinks links={sourceLinks} />}
+                {imageUrls.length > 0 && <ImageGallery imageUrls={imageUrls} />}
+              </>
+            )}
           </div>
         </div>
       </div>

frontend/src/components/MessageInput.tsx

@@ -2,16 +2,19 @@
 import { Button } from "@/components/ui/button";
 import { AutosizeTextarea } from "@/components/ui/textarea";
 import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip";
-import { Mic, Paperclip, Send } from "lucide-react";
+import { useChatContext } from "@/lib/store/ChatContext";
+import { Mic, Paperclip, Send, Square, StopCircle } from "lucide-react";
 import React, { useState } from "react";
 
 interface MessageInputProps {
   onSendMessage: (content: string) => void;
   isLoading: boolean;
   userInputRef: React.LegacyRef<any>;
+  onCancel: () => void;
 }
 
-const MessageInput: React.FC<MessageInputProps> = ({ onSendMessage, isLoading, userInputRef }) => {
+// Traditional component that takes props
+const MessageInput: React.FC<MessageInputProps> = ({ onSendMessage, isLoading, userInputRef, onCancel }) => {
   const [message, setMessage] = useState("");
 
   const handleMessageChange = (e: React.ChangeEvent<HTMLTextAreaElement>) => {
@@ -37,7 +40,7 @@ const MessageInput: React.FC<MessageInputProps> = ({ onSendMessage, isLoading, u
     <div className="p-4 absolute bottom-0 left-0 right-0 bg-transparent mb-4">
       <form onSubmit={handleSubmit} className="max-w-4xl mx-auto">
         <div className="relative flex items-center bg-background shadow-lg rounded-[2rem] border overflow-hidden h-full">
-          <AutosizeTextarea placeholder="Ask a research question..." maxHeight={500} minHeight={52} className="pr-36 pl-6 py-4 font-medium border-none h-auto resize-none" value={message} onChange={handleMessageChange} onKeyDown={handleKeyDown} disabled={isLoading} rows={1} autoFocus ref={userInputRef}/>
+          <AutosizeTextarea placeholder="Ask a research question..." maxHeight={500} minHeight={52} className="pr-36 pl-6 py-4 font-medium border-none h-auto resize-none" value={message} onChange={handleMessageChange} onKeyDown={handleKeyDown} disabled={isLoading} rows={1} autoFocus ref={userInputRef} />
 
           <div className="absolute right-3 flex items-center gap-2 h-full">
             <TooltipProvider>
@@ -62,10 +65,17 @@ const MessageInput: React.FC<MessageInputProps> = ({ onSendMessage, isLoading, u
               </Tooltip>
             </TooltipProvider>
 
-            <Button type="submit" size="icon" className="h-9 w-9 rounded-full" disabled={isLoading || !message.trim()}>
-              <Send className="h-4 w-4" />
-              <span className="sr-only">Send message</span>
-            </Button>
+            {isLoading && onCancel ? (
+              <Button type="button" size="icon" variant="ghost" className="h-9 w-9 rounded-full" onClick={onCancel}>
+                <Square className="h-4 w-4" />
+                <span className="sr-only">Cancel</span>
+              </Button>
+            ) : (
+              <Button type="submit" size="icon" className="h-9 w-9 rounded-full" disabled={isLoading || !message.trim()}>
+                <Send className="h-4 w-4" />
+                <span className="sr-only">Send message</span>
+              </Button>
+            )}
           </div>
         </div>
       </form>

frontend/src/components/ResearchControls.tsx

@@ -2,15 +2,17 @@ import { Checkbox } from "@/components/ui/checkbox";
 import { Label } from "@/components/ui/label";
 import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
 import { Separator } from "@/components/ui/separator";
+import { useChatContext } from "@/lib/store/ChatContext";
 import { ResearchOptions } from "@/lib/types";
 import React from "react";
-import { Input } from "@/components/ui/input";  // Make sure you have an Input component
+import { Input } from "@/components/ui/input"; // Make sure you have an Input component
 
 interface ResearchControlsProps {
   options: ResearchOptions;
   onOptionChange: (options: ResearchOptions) => void;
 }
 
+// Traditional prop-based component
 const ResearchControls: React.FC<ResearchControlsProps> = ({ options, onOptionChange }) => {
   return (
     <div className="space-y-6">
@@ -56,24 +58,12 @@ const ResearchControls: React.FC<ResearchControlsProps> = ({ options, onOptionCh
 
         <div className="space-y-2">
           <Label htmlFor="max-depth">Max Depth</Label>
-          <Input
-            type="number"
-            id="max-depth"
-            value={options.max_depth}
-            onChange={(e) => onOptionChange({ ...options, max_depth: parseInt(e.target.value, 10) })}
-            className="w-full"
-          />
+          <Input type="number" id="max-depth" value={options.max_depth} onChange={(e) => onOptionChange({ ...options, max_depth: parseInt(e.target.value, 10) })} className="w-full" />
         </div>
 
         <div className="space-y-2">
           <Label htmlFor="num-sites-per-query">Number of Sites per Query</Label>
-          <Input
-            type="number"
-            id="num-sites-per-query"
-            value={options.num_sites_per_query}
-            onChange={(e) => onOptionChange({ ...options, num_sites_per_query: parseInt(e.target.value, 10) })}
-            className="w-full"
-          />
+          <Input type="number" id="num-sites-per-query" value={options.num_sites_per_query} onChange={(e) => onOptionChange({ ...options, num_sites_per_query: parseInt(e.target.value, 10) })} className="w-full" />
         </div>
       </div>
 
@@ -99,4 +89,11 @@ const ResearchControls: React.FC<ResearchControlsProps> = ({ options, onOptionCh
   );
 };
 
-export default ResearchControls;
+// Context-based component
+export const ResearchControlsWithContext: React.FC = () => {
+  const { researchOptions, setResearchOptions } = useChatContext();
+
+  return <ResearchControls options={researchOptions} onOptionChange={setResearchOptions} />;
+};
+
+export default ResearchControls;

frontend/src/components/ui/ChatLayout.tsx

@@ -29,9 +29,7 @@ const ChatLayout: React.FC<ChatLayoutProps> = ({ sidebar, mainContent, settingsP
           </SheetTrigger>
           <SheetContent side="left" className="w-[80%] sm:w-[350px] p-0">
             <SheetTitle className="sr-only">Mobile Navigation</SheetTitle>
-            <SheetDescription className="sr-only">
-              Sidebar navigation for mobile devices
-            </SheetDescription>
+            <SheetDescription className="sr-only">Sidebar navigation for mobile devices</SheetDescription>
             <div className="border-b p-4">
               <h2 className="text-lg font-semibold">Conversations</h2>
             </div>
@@ -63,7 +61,7 @@ const ChatLayout: React.FC<ChatLayoutProps> = ({ sidebar, mainContent, settingsP
 
       <div className="flex-1 overflow-hidden">
         <ResizablePanelGroup direction="horizontal">
-          <ResizablePanel defaultSize={25} minSize={17} maxSize={30} className="hidden md:block">
+          <ResizablePanel defaultSize={15} className="hidden md:block min-w-[15rem] max-w-[35rem]">
             <Card className="h-full rounded-none border-r border-t-0 border-l-0 border-b-0">
               <ScrollArea className="h-full">{sidebar}</ScrollArea>
             </Card>
@@ -71,7 +69,7 @@ const ChatLayout: React.FC<ChatLayoutProps> = ({ sidebar, mainContent, settingsP
 
           <ResizableHandle withHandle className="hidden md:flex" />
 
-          <ResizablePanel defaultSize={75} className="w-full md:w-auto">
+          <ResizablePanel defaultSize={85} className="w-full md:w-auto">
             <Tabs defaultValue="chat" className="h-full flex flex-col">
               <div className="p-4">
                 <TabsList className="">
@@ -86,7 +84,7 @@ const ChatLayout: React.FC<ChatLayoutProps> = ({ sidebar, mainContent, settingsP
                 </TabsList>
               </div>
 
-              <TabsContent value="chat" className="overflow-auto flex flex-1" tabIndex={-1}>
+              <TabsContent value="chat" className="overflow-auto flex flex-1 !mt-0" tabIndex={-1}>
                 {mainContent}
               </TabsContent>
 

frontend/src/components/ui/ConversationList.tsx

@@ -1,9 +1,18 @@
 "use client";
 import { Button } from "@/components/ui/button";
-import { ConversationListProps } from "@/lib/types";
+import { useChatContext } from "@/lib/store/ChatContext";
 import { MessageSquare, PlusCircle, Trash2, XCircle } from "lucide-react";
 import React from "react";
 
+interface ConversationListProps {
+  conversations: any[];
+  onNewConversation: () => void;
+  onSelectConversation: (id: string) => void;
+  onDeleteConversation: (id: string) => void;
+  onDeleteAllConversations: () => void;
+}
+
+// A component that accepts props for backward compatibility
 const ConversationList: React.FC<ConversationListProps> = ({ conversations, onNewConversation, onSelectConversation, onDeleteConversation, onDeleteAllConversations }) => {
   const handleSelectConversation = (id: string) => {
     onSelectConversation(id);
@@ -22,36 +31,42 @@ const ConversationList: React.FC<ConversationListProps> = ({ conversations, onNe
           </Button>
         )}
       </div>
-
       <div className="px-2 py-2">
         <h2 className="text-sm font-semibold px-2 mb-2">Recent Research</h2>
         <div className="space-y-1">
           {conversations.length === 0 ? (
             <p className="text-sm text-muted-foreground px-2">No conversations yet</p>
           ) : (
-              conversations.map((conversation) => (
-                <div key={conversation.id} className="flex items-center gap-1 px-1 group">
-                  <Button variant={conversation.active ? "secondary" : "ghost"} className={`flex-1 pr-0.5 justify-start text-left truncate ${conversation.active ? "bg-accent" : ""}`} onClick={() => handleSelectConversation(conversation.id)}>
-                    <MessageSquare className="mr-2 h-4 w-4 shrink-0" />
-                    <span className="truncate">{conversation.title}</span>
-                  </Button>
-                  <Button
-                    variant="ghost"
-                    size="icon"
-                    className="h-8 w-8 ml-auto hidden group-hover:flex hover:bg-destructive"
-                    onClick={(e) => {
-                      e.stopPropagation();
-                      onDeleteConversation(conversation.id);
-                    }}>
-                    <XCircle className="h-4 w-4" />
-                  </Button>
-                </div>
-              ))
-            )}
+            conversations.map((conversation) => (
+              <div key={conversation.id} className="flex items-center gap-1 px-1 group">
+                <Button variant={conversation.active ? "secondary" : "ghost"} className={`flex-1 pr-0.5 justify-start text-left truncate ${conversation.active ? "bg-accent" : ""}`} onClick={() => handleSelectConversation(conversation.id)}>
+                  <MessageSquare className="mr-2 h-4 w-4 shrink-0" />
+                  <span className="truncate">{conversation.title}</span>
+                </Button>
+                <Button
+                  variant="ghost"
+                  size="icon"
+                  className="h-8 w-8 ml-auto hidden group-hover:flex hover:bg-destructive"
+                  onClick={(e) => {
+                    e.stopPropagation();
+                    onDeleteConversation(conversation.id);
+                  }}>
+                  <XCircle className="h-4 w-4" />
+                </Button>
+              </div>
+            ))
+          )}
         </div>
       </div>
     </div>
   );
 };
 
+// A component that uses context directly
+export const ConversationListWithContext: React.FC = () => {
+  const { conversations, newConversation, selectConversation, deleteConversation, deleteAllConversations } = useChatContext();
+
+  return <ConversationList conversations={conversations} onNewConversation={newConversation} onSelectConversation={selectConversation} onDeleteConversation={deleteConversation} onDeleteAllConversations={deleteAllConversations} />;
+};
+
 export default ConversationList;

frontend/src/components/ui/badge.tsx

@@ -0,0 +1,36 @@
+import * as React from "react"
+import { cva, type VariantProps } from "class-variance-authority"
+
+import { cn } from "@/lib/utils"
+
+const badgeVariants = cva(
+  "inline-flex items-center rounded-md border px-2.5 py-0.5 text-xs font-semibold transition-colors focus:outline-none focus:ring-2 focus:ring-ring focus:ring-offset-2",
+  {
+    variants: {
+      variant: {
+        default:
+          "border-transparent bg-primary text-primary-foreground shadow hover:bg-primary/80",
+        secondary:
+          "border-transparent bg-secondary text-secondary-foreground hover:bg-secondary/80",
+        destructive:
+          "border-transparent bg-destructive text-destructive-foreground shadow hover:bg-destructive/80",
+        outline: "text-foreground",
+      },
+    },
+    defaultVariants: {
+      variant: "default",
+    },
+  }
+)
+
+export interface BadgeProps
+  extends React.HTMLAttributes<HTMLDivElement>,
+    VariantProps<typeof badgeVariants> {}
+
+function Badge({ className, variant, ...props }: BadgeProps) {
+  return (
+    <div className={cn(badgeVariants({ variant }), className)} {...props} />
+  )
+}
+
+export { Badge, badgeVariants }

frontend/src/lib/store/ChatContext.tsx

@@ -0,0 +1,454 @@
+"use client";
+
+import { ChatData, ChatState, Conversation, Message, ResearchOptions, ResearchResults, StatusUpdate } from "@/lib/types";
+import { ReactNode, createContext, useCallback, useContext, useEffect, useRef, useState } from "react";
+import { v4 as uuidv4 } from "uuid";
+import { disconnectSocket, getSocket, initializeSocket } from "@/lib/socket";
+
+// Utility functions for local storage
+const saveToStorage = (data: ChatData) => {
+  if (typeof window !== "undefined") {
+    localStorage.setItem("chatData", JSON.stringify(data));
+  }
+};
+
+const loadFromStorage = (): ChatData => {
+  if (typeof window === "undefined") {
+    return { conversations: [], currentConversationId: null };
+  }
+  const data = localStorage.getItem("chatData");
+  if (!data) {
+    return { conversations: [], currentConversationId: null };
+  }
+  try {
+    const parsed = JSON.parse(data);
+    return {
+      conversations: Array.isArray(parsed.conversations)
+        ? parsed.conversations.map((conv: Conversation) => ({
+            ...conv,
+            messages: Array.isArray(conv.messages)
+              ? conv.messages.map((msg: Message) => ({
+                  ...msg,
+                  // Ensure media property is preserved if it exists
+                  media: msg.media || undefined,
+                }))
+              : [],
+          }))
+        : [],
+      currentConversationId: parsed.currentConversationId,
+    };
+  } catch (e) {
+    return { conversations: [], currentConversationId: null };
+  }
+};
+
+// Define the context type
+interface ChatContextType {
+  // State
+  chatState: ChatState;
+  conversations: Conversation[];
+  currentConversationId: string | null;
+  researchOptions: ResearchOptions;
+  userInputRef: React.RefObject<HTMLTextAreaElement>;
+
+  // Actions
+  setResearchOptions: (options: ResearchOptions) => void;
+  sendMessage: (content: string) => void;
+  newConversation: () => void;
+  selectConversation: (id: string) => void;
+  deleteConversation: (id: string) => void;
+  deleteAllConversations: () => void;
+  abortResearch: () => void; // New function to abort research
+}
+
+// Create the context with a default value
+const ChatContext = createContext<ChatContextType | undefined>(undefined);
+
+// Provider component
+export const ChatProvider = ({ children }: { children: ReactNode }) => {
+  const [chatState, setChatState] = useState<ChatState>({ messages: [], isLoading: false, error: null });
+  const [conversations, setConversations] = useState<Conversation[]>([]);
+  const [currentConversationId, setCurrentConversationId] = useState<string | null>(null);
+  const [researchOptions, setResearchOptions] = useState<ResearchOptions>({
+    depth: "basic",
+    sources: true,
+    citations: false,
+    max_depth: 1,
+    num_sites_per_query: 3,
+  });
+
+  const userInputRef = useRef<HTMLTextAreaElement>(null);
+
+  // Focus management
+  useEffect(() => {
+    const focusInput = () => {
+      setTimeout(() => {
+        userInputRef.current?.focus();
+      }, 100);
+    };
+
+    focusInput();
+  }, [currentConversationId]);
+
+  // Load initial data
+  useEffect(() => {
+    const data = loadFromStorage();
+    setConversations(data.conversations);
+    setCurrentConversationId(data.currentConversationId);
+
+    if (data.currentConversationId) {
+      const conversation = data.conversations.find((c) => c.id === data.currentConversationId);
+      if (conversation) {
+        // Check if any loaded message has isProgress true
+        let messages = conversation.messages;
+        if (messages.some(msg => msg.isProgress === true)) {
+          // Convert any progress messages to error messages
+          messages = messages.map(msg =>
+            msg.isProgress === true
+              ? { ...msg, content: "Connection error", isProgress: false }
+              : msg
+          );
+        }
+
+        setChatState((prev) => ({ ...prev, messages, isLoading: false }));
+      }
+    }
+  }, []);
+
+  // Socket initialization
+  useEffect(() => {
+    const socket = initializeSocket();
+
+    socket.on("connect", () => {
+      console.log("Connected to research server");
+    });
+
+    socket.on("disconnect", () => {
+      console.log("Disconnected from research server");
+
+      // When socket disconnects, update any progress messages to show connection error
+      setChatState((prevState) => {
+        const updatedMessages = prevState.messages.map(msg =>
+          msg.isProgress === true
+            ? { ...msg, content: "Connection error", isProgress: false }
+            : msg
+        );
+
+        return {
+          ...prevState,
+          messages: updatedMessages,
+          isLoading: false,
+          error: "Lost connection to research server"
+        };
+      });
+    });
+
+    socket.on("status", (data: StatusUpdate) => {
+      setChatState((prevState) => {
+        const messages = [...prevState.messages];
+        const progressText = data.message;
+        const progress = data.progress;
+
+        // Find the last assistant message that is a progress update
+        const lastProgressIndex = messages.findLastIndex((msg) => msg.role === "assistant" && msg.isProgress === true);
+
+        if (lastProgressIndex !== -1) {
+          // Update existing progress message
+          messages[lastProgressIndex] = {
+            ...messages[lastProgressIndex],
+            content: progressText,
+            progress: progress,
+            timestamp: new Date(),
+          };
+        } else {
+          // Add new progress message
+          messages.push({
+            id: uuidv4(),
+            content: progressText,
+            role: "assistant",
+            timestamp: new Date(),
+            progress: progress,
+            isProgress: true,
+          });
+        }
+
+        return {
+          ...prevState,
+          messages,
+          isLoading: true,
+        };
+      });
+    });
+
+    socket.on("research_complete", (results: ResearchResults) => {
+      setChatState((prevState) => {
+        const messages = [...prevState.messages];
+
+        // Remove the last progress message if it exists
+        const lastProgressIndex = messages.findLastIndex((msg) => msg.role === "assistant" && msg.isProgress === true);
+
+        if (lastProgressIndex !== -1) {
+          messages.splice(lastProgressIndex, 1);
+        }
+
+        const newMessages = [
+          ...messages,
+          {
+            id: uuidv4(),
+            content: results.content || "Error: No content available",
+            role: "assistant" as const,
+            timestamp: new Date(results.timestamp),
+            media: results.media,
+            research_tree: results.research_tree,
+          },
+        ];
+
+        // Save updated messages to localStorage
+        const updatedState = {
+          ...prevState,
+          isLoading: false,
+          messages: newMessages,
+        };
+
+        // Update localStorage with the new messages
+        const updatedData: ChatData = {
+          conversations: conversations.map((conv) => ({
+            ...conv,
+            messages: conv.id === currentConversationId ? newMessages : conv.messages || [],
+            lastUpdated: conv.id === currentConversationId ? new Date().toISOString() : conv.lastUpdated,
+          })),
+          currentConversationId,
+        };
+
+        saveToStorage(updatedData);
+
+        return updatedState;
+      });
+    });
+
+    socket.on("research_aborted", () => {
+      setChatState((prevState) => {
+        const messages = [...prevState.messages];
+        const lastProgressIndex = messages.findLastIndex((msg) => msg.role === "assistant" && msg.isProgress === true);
+
+        if (lastProgressIndex !== -1) {
+          messages.splice(lastProgressIndex, 1);
+        }
+
+        // Add a message indicating the research was canceled
+        messages.push({
+          id: uuidv4(),
+          content: "Research has been canceled.",
+          role: "assistant",
+          timestamp: new Date(),
+        });
+
+        return {
+          ...prevState,
+          isLoading: false,
+          messages,
+        };
+      });
+    });
+
+    socket.on("error", (error: { message: string }) => {
+      setChatState((prevState) => ({
+        ...prevState,
+        error: error.message,
+        isLoading: false,
+      }));
+    });
+
+    return () => {
+      disconnectSocket();
+    };
+  }, []);
+
+  // Save data whenever conversations or messages change
+  useEffect(() => {
+    const data: ChatData = {
+      conversations: conversations.map((conv) => ({
+        ...conv,
+        messages: conv.id === currentConversationId ? chatState.messages : conv.messages || [],
+      })),
+      currentConversationId,
+    };
+
+    saveToStorage(data);
+  }, [conversations, currentConversationId, chatState.messages]);
+
+  // Action handlers
+  const sendMessage = useCallback(
+    (content: string) => {
+      if (!content.trim()) return;
+
+      let conversationId = currentConversationId;
+      const newMessage: Message = {
+        id: uuidv4(),
+        content,
+        role: "user",
+        timestamp: new Date(),
+      };
+      const newLoadingMessage: Message = {
+        id: uuidv4(),
+        content: "Loading...",
+        role: "assistant",
+        timestamp: new Date(),
+        isProgress: true,
+      };
+
+      // Create a new conversation if none exists
+      if (!conversationId) {
+        conversationId = uuidv4();
+        setCurrentConversationId(conversationId);
+        setConversations((prev) => [
+          {
+            id: conversationId as string,
+            title: content.length > 30 ? `${content.substring(0, 30)}...` : content,
+            lastUpdated: new Date().toISOString(),
+            messages: [newMessage],
+            active: true,
+          },
+          ...prev.map((c) => ({ ...c, active: false })),
+        ]);
+      } else {
+        // Update the existing conversation
+        setConversations((prev) =>
+          prev.map((conv) => ({
+            ...conv,
+            lastUpdated: conv.id === conversationId ? new Date().toISOString() : conv.lastUpdated,
+            active: conv.id === conversationId,
+            messages: conv.id === conversationId ? [...(conv.messages || []), newMessage] : conv.messages || [],
+          }))
+        );
+      }
+
+      setChatState((prevState) => ({
+        ...prevState,
+        messages: [...prevState.messages, newMessage, newLoadingMessage],
+        isLoading: true,
+        error: null,
+      }));
+
+      // Send message to server via socket
+      try {
+        const socket = getSocket();
+        socket.emit("start_research", {
+          topic: content,
+          max_depth: researchOptions.max_depth,
+          num_sites_per_query: researchOptions.num_sites_per_query,
+        });
+      } catch (error) {
+        setChatState((prevState) => ({
+          ...prevState,
+          error: "Failed to connect to research server",
+          isLoading: false,
+        }));
+      }
+    },
+    [currentConversationId, conversations, researchOptions]
+  );
+
+  const newConversation = useCallback(() => {
+    userInputRef.current?.focus();
+    setCurrentConversationId(null);
+    setChatState(() => ({
+      messages: [],
+      isLoading: false,
+      error: null,
+    }));
+  }, []);
+
+  const selectConversation = useCallback((id: string) => {
+    const data = loadFromStorage();
+    const conversation = data.conversations.find((c) => c.id === id);
+
+    setCurrentConversationId(id);
+    setChatState((prev) => ({
+      ...prev,
+      messages: conversation?.messages || [],
+      isLoading: false,
+      error: null,
+    }));
+    setConversations((prev) =>
+      prev.map((conv) => ({
+        ...conv,
+        active: conv.id === id,
+      }))
+    );
+  }, []);
+
+  const deleteConversation = useCallback(
+    (id: string) => {
+      setConversations((prev) => prev.filter((conv) => conv.id !== id));
+      if (currentConversationId === id) {
+        newConversation();
+      }
+    },
+    [currentConversationId, newConversation]
+  );
+
+  const deleteAllConversations = useCallback(() => {
+    setConversations([]);
+    newConversation();
+  }, [newConversation]);
+
+  const abortResearch = useCallback(() => {
+    try {
+      const socket = getSocket();
+      socket.emit("abort_research");
+
+      setChatState((prevState) => ({
+        ...prevState,
+        isLoading: true,
+      }));
+    } catch (error) {
+      console.error("Failed to abort research:", error);
+      setChatState((prevState) => ({
+        ...prevState,
+        error: "Failed to abort research",
+      }));
+    }
+  }, []);
+
+  // Keyboard shortcuts | Ctrl + I to new chat
+  useEffect(() => {
+    const handleKeyDown = (event: KeyboardEvent) => {
+      if (event.ctrlKey && event.key === "i") {
+        event.preventDefault();
+        newConversation();
+      }
+    };
+
+    window.addEventListener("keydown", handleKeyDown);
+    return () => {
+      window.removeEventListener("keydown", handleKeyDown);
+    };
+  }, []);
+
+  const value = {
+    chatState,
+    conversations,
+    currentConversationId,
+    researchOptions,
+    userInputRef,
+    setResearchOptions,
+    sendMessage,
+    newConversation,
+    selectConversation,
+    deleteConversation,
+    deleteAllConversations,
+    abortResearch,
+  };
+
+  return <ChatContext.Provider value={value}>{children}</ChatContext.Provider>;
+};
+
+// Custom hook for using the chat context
+export const useChatContext = () => {
+  const context = useContext(ChatContext);
+  if (context === undefined) {
+    throw new Error("useChatContext must be used within a ChatProvider");
+  }
+  return context;
+};

frontend/src/lib/types.ts

@@ -1,17 +1,15 @@
 export interface Message {
   id: string;
   content: string;
-  role: "user" | "assistant" | "system";
+  role: "user" | "assistant";
   timestamp: Date;
   media?: {
+    links?: Array<{ text: string; url: string }>;
     images?: string[];
-    videos?: string[];
-    links?: Array<{
-      text: string;
-      url: string;
-    }>;
-    references?: any[];
   };
+  research_tree?: ResearchTree;
+  progress?: number;
+  isProgress?: boolean;
 }
 
 export interface ChatState {
@@ -46,28 +44,28 @@ export interface StatusUpdate {
   progress: number;
 }
 
-export interface ResearchNode {
-  query: string;
-  children?: ResearchNode[];
-}
-
-export interface ResearchMetadata {
-  total_queries: number;
-  total_sources: number;
-  max_depth_reached: number;
-  total_tokens: number;
-}
-
+// Simplified research results based on actual server output
 export interface ResearchResults {
   topic: string;
   timestamp: string;
-  content: string;
+  // Optional fields not present in basic server response
+  content?: string;
   media?: {
     images?: string[];
     videos?: string[];
+    links?: Array<{
+      text: string;
+      url: string;
+    }>;
   };
-  research_tree: ResearchNode;
-  metadata: ResearchMetadata;
+  research_tree?: ResearchTree;
+}
+
+export interface ResearchTree {
+  query: string;
+  depth: number;
+  sources: string[];
+  children: ResearchTree[];
 }
 
 export interface ConversationListProps {