diff --git a/surfsense_backend/alembic/versions/124_add_ai_file_sort_enabled.py b/surfsense_backend/alembic/versions/124_add_ai_file_sort_enabled.py deleted file mode 100644 index b77eb9337..000000000 --- a/surfsense_backend/alembic/versions/124_add_ai_file_sort_enabled.py +++ /dev/null @@ -1,44 +0,0 @@ -"""124_add_ai_file_sort_enabled - -Revision ID: 124 -Revises: 123 -Create Date: 2026-04-14 - -Adds ai_file_sort_enabled boolean column to searchspaces. -Defaults to False so AI file sorting is opt-in per search space. -""" - -from __future__ import annotations - -from collections.abc import Sequence - -import sqlalchemy as sa - -from alembic import op - -revision: str = "124" -down_revision: str | None = "123" -branch_labels: str | Sequence[str] | None = None -depends_on: str | Sequence[str] | None = None - - -def upgrade() -> None: - conn = op.get_bind() - existing_columns = [ - col["name"] for col in sa.inspect(conn).get_columns("searchspaces") - ] - - if "ai_file_sort_enabled" not in existing_columns: - op.add_column( - "searchspaces", - sa.Column( - "ai_file_sort_enabled", - sa.Boolean(), - nullable=False, - server_default=sa.text("false"), - ), - ) - - -def downgrade() -> None: - op.drop_column("searchspaces", "ai_file_sort_enabled") diff --git a/surfsense_backend/app/agents/new_chat/chat_deepagent.py b/surfsense_backend/app/agents/new_chat/chat_deepagent.py index 8dc2fa98f..9bf38cad6 100644 --- a/surfsense_backend/app/agents/new_chat/chat_deepagent.py +++ b/surfsense_backend/app/agents/new_chat/chat_deepagent.py @@ -472,7 +472,7 @@ async def create_surfsense_deep_agent( SubAgentMiddleware(backend=StateBackend, subagents=[general_purpose_spec]), create_summarization_middleware(llm, StateBackend), PatchToolCallsMiddleware(), - DedupHITLToolCallsMiddleware(agent_tools=tools), + DedupHITLToolCallsMiddleware(), AnthropicPromptCachingMiddleware(unsupported_model_behavior="ignore"), ] diff --git a/surfsense_backend/app/agents/new_chat/middleware/dedup_tool_calls.py b/surfsense_backend/app/agents/new_chat/middleware/dedup_tool_calls.py index 61494ff1a..f5e8f1235 100644 --- a/surfsense_backend/app/agents/new_chat/middleware/dedup_tool_calls.py +++ b/surfsense_backend/app/agents/new_chat/middleware/dedup_tool_calls.py @@ -20,39 +20,19 @@ from langgraph.runtime import Runtime logger = logging.getLogger(__name__) -_NATIVE_HITL_TOOL_DEDUP_KEYS: dict[str, str] = { - # Gmail - "send_gmail_email": "subject", - "create_gmail_draft": "subject", - "update_gmail_draft": "draft_subject_or_id", - "trash_gmail_email": "email_subject_or_id", - # Google Calendar - "create_calendar_event": "title", - "update_calendar_event": "event_title_or_id", +_HITL_TOOL_DEDUP_KEYS: dict[str, str] = { "delete_calendar_event": "event_title_or_id", - # Google Drive - "create_google_drive_file": "file_name", + "update_calendar_event": "event_title_or_id", + "trash_gmail_email": "email_subject_or_id", + "update_gmail_draft": "draft_subject_or_id", "delete_google_drive_file": "file_name", - # OneDrive - "create_onedrive_file": "file_name", "delete_onedrive_file": "file_name", - # Dropbox - "create_dropbox_file": "file_name", - "delete_dropbox_file": "file_name", - # Notion - "create_notion_page": "title", - "update_notion_page": "page_title", "delete_notion_page": "page_title", - # Linear - "create_linear_issue": "title", - "update_linear_issue": "issue_ref", + "update_notion_page": "page_title", "delete_linear_issue": "issue_ref", - # Jira - "create_jira_issue": "summary", + "update_linear_issue": "issue_ref", "update_jira_issue": "issue_title_or_key", "delete_jira_issue": "issue_title_or_key", - # Confluence - "create_confluence_page": "title", "update_confluence_page": "page_title_or_id", "delete_confluence_page": "page_title_or_id", } @@ -63,39 +43,22 @@ class DedupHITLToolCallsMiddleware(AgentMiddleware): # type: ignore[type-arg] Only the **first** occurrence of each (tool-name, primary-arg-value) pair is kept; subsequent duplicates are silently dropped. - - The dedup map is built from two sources: - - 1. A comprehensive list of native HITL tools (hardcoded above). - 2. Any ``StructuredTool`` instances passed via *agent_tools* whose - ``metadata`` contains ``{"hitl": True, "hitl_dedup_key": "..."}``. - This is how MCP tools automatically get dedup support. """ tools = () - def __init__(self, *, agent_tools: list[Any] | None = None) -> None: - self._dedup_keys: dict[str, str] = dict(_NATIVE_HITL_TOOL_DEDUP_KEYS) - for t in agent_tools or []: - meta = getattr(t, "metadata", None) or {} - if meta.get("hitl") and meta.get("hitl_dedup_key"): - self._dedup_keys[t.name] = meta["hitl_dedup_key"] - def after_model( self, state: AgentState, runtime: Runtime[Any] ) -> dict[str, Any] | None: - return self._dedup(state, self._dedup_keys) + return self._dedup(state) async def aafter_model( self, state: AgentState, runtime: Runtime[Any] ) -> dict[str, Any] | None: - return self._dedup(state, self._dedup_keys) + return self._dedup(state) @staticmethod - def _dedup( - state: AgentState, - dedup_keys: dict[str, str], # type: ignore[type-arg] - ) -> dict[str, Any] | None: + def _dedup(state: AgentState) -> dict[str, Any] | None: # type: ignore[type-arg] messages = state.get("messages") if not messages: return None @@ -110,7 +73,7 @@ class DedupHITLToolCallsMiddleware(AgentMiddleware): # type: ignore[type-arg] for tc in tool_calls: name = tc.get("name", "") - dedup_key_arg = dedup_keys.get(name) + dedup_key_arg = _HITL_TOOL_DEDUP_KEYS.get(name) if dedup_key_arg is not None: arg_val = str(tc.get("args", {}).get(dedup_key_arg, "")).lower() key = (name, arg_val) diff --git a/surfsense_backend/app/agents/new_chat/middleware/filesystem.py b/surfsense_backend/app/agents/new_chat/middleware/filesystem.py index bcd544d61..af5a6925b 100644 --- a/surfsense_backend/app/agents/new_chat/middleware/filesystem.py +++ b/surfsense_backend/app/agents/new_chat/middleware/filesystem.py @@ -9,7 +9,6 @@ from __future__ import annotations import asyncio import logging import re -import secrets from datetime import UTC, datetime from typing import Annotated, Any @@ -28,7 +27,6 @@ from sqlalchemy import delete, select from app.agents.new_chat.sandbox import ( _evict_sandbox_cache, - delete_sandbox, get_or_create_sandbox, is_sandbox_enabled, ) @@ -554,8 +552,7 @@ class SurfSenseFilesystemMiddleware(FilesystemMiddleware): @staticmethod def _wrap_as_python(code: str) -> str: """Wrap Python code in a shell invocation for the sandbox.""" - sentinel = f"_PYEOF_{secrets.token_hex(8)}" - return f"python3 << '{sentinel}'\n{code}\n{sentinel}" + return f"python3 << 'PYEOF'\n{code}\nPYEOF" async def _execute_in_sandbox( self, @@ -575,10 +572,7 @@ class SurfSenseFilesystemMiddleware(FilesystemMiddleware): self._thread_id, first_err, ) - try: - await delete_sandbox(self._thread_id) - except Exception: - _evict_sandbox_cache(self._thread_id) + _evict_sandbox_cache(self._thread_id) try: return await self._try_sandbox_execute(command, runtime, timeout) except Exception: @@ -593,14 +587,7 @@ class SurfSenseFilesystemMiddleware(FilesystemMiddleware): runtime: ToolRuntime[None, FilesystemState], timeout: int | None, ) -> str: - sandbox, _is_new = await get_or_create_sandbox(self._thread_id) - # NOTE: sync_files_to_sandbox is intentionally disabled. - # The virtual FS contains XML-wrapped KB documents whose paths - # would double-nest under SANDBOX_DOCUMENTS_ROOT (e.g. - # /home/daytona/documents/documents/Report.xml) and uploading - # all KB docs on the first execute_code call adds significant - # latency. Re-enable once path mapping is fixed and upload is - # limited to user-created scratch files. + sandbox, is_new = await get_or_create_sandbox(self._thread_id) # files = runtime.state.get("files") or {} # await sync_files_to_sandbox(self._thread_id, files, sandbox, is_new) result = await sandbox.aexecute(command, timeout=timeout) diff --git a/surfsense_backend/app/agents/new_chat/middleware/knowledge_search.py b/surfsense_backend/app/agents/new_chat/middleware/knowledge_search.py index 0460da11d..06ed4ad80 100644 --- a/surfsense_backend/app/agents/new_chat/middleware/knowledge_search.py +++ b/surfsense_backend/app/agents/new_chat/middleware/knowledge_search.py @@ -58,14 +58,6 @@ class KBSearchPlan(BaseModel): default=None, description="Optional ISO end date or datetime for KB search filtering.", ) - is_recency_query: bool = Field( - default=False, - description=( - "True when the user's intent is primarily about recency or temporal " - "ordering (e.g. 'latest', 'newest', 'most recent', 'last uploaded') " - "rather than topical relevance." - ), - ) def _extract_text_from_message(message: BaseMessage) -> str: @@ -253,7 +245,7 @@ def _build_kb_planner_prompt( return ( "You optimize internal knowledge-base search inputs for document retrieval.\n" "Return JSON only with this exact shape:\n" - '{"optimized_query":"string","start_date":"ISO string or null","end_date":"ISO string or null","is_recency_query":bool}\n\n' + '{"optimized_query":"string","start_date":"ISO string or null","end_date":"ISO string or null"}\n\n' "Rules:\n" "- Preserve the user's intent.\n" "- Rewrite the query to improve retrieval using concrete entities, acronyms, projects, tools, people, and document-specific terms when helpful.\n" @@ -261,11 +253,6 @@ def _build_kb_planner_prompt( "- Only use date filters when the latest user request or recent dialogue clearly implies a time range.\n" "- If you use date filters, prefer returning both bounds.\n" "- If no date filter is useful, return null for both dates.\n" - '- Set "is_recency_query" to true ONLY when the user\'s primary intent is about ' - "recency or temporal ordering rather than topical relevance. Examples: " - '"latest file", "newest upload", "most recent document", "what did I save last", ' - '"show me files from today", "last thing I added". ' - "When true, results will be sorted by date instead of relevance.\n" "- Do not include markdown, prose, or explanations.\n\n" f"Today's UTC date: {today}\n\n" f"Recent conversation:\n{recent_conversation or '(none)'}\n\n" @@ -519,135 +506,6 @@ def _resolve_search_types( return list(expanded) if expanded else None -_RECENCY_MAX_CHUNKS_PER_DOC = 5 - - -async def browse_recent_documents( - *, - search_space_id: int, - document_type: list[str] | None = None, - top_k: int = 10, - start_date: datetime | None = None, - end_date: datetime | None = None, -) -> list[dict[str, Any]]: - """Return documents ordered by recency (newest first), no relevance ranking. - - Used when the user's intent is temporal ("latest file", "most recent upload") - and hybrid search would produce poor results because the query has no - meaningful topical signal. - """ - from sqlalchemy import func, select - - from app.db import DocumentType - - async with shielded_async_session() as session: - base_conditions = [ - Document.search_space_id == search_space_id, - func.coalesce(Document.status["state"].astext, "ready") != "deleting", - ] - - if document_type is not None: - import contextlib - - doc_type_enums = [] - for dt in document_type: - if isinstance(dt, str): - with contextlib.suppress(KeyError): - doc_type_enums.append(DocumentType[dt]) - else: - doc_type_enums.append(dt) - if doc_type_enums: - if len(doc_type_enums) == 1: - base_conditions.append(Document.document_type == doc_type_enums[0]) - else: - base_conditions.append(Document.document_type.in_(doc_type_enums)) - - if start_date is not None: - base_conditions.append(Document.updated_at >= start_date) - if end_date is not None: - base_conditions.append(Document.updated_at <= end_date) - - doc_query = ( - select(Document) - .where(*base_conditions) - .order_by(Document.updated_at.desc()) - .limit(top_k) - ) - result = await session.execute(doc_query) - documents = result.scalars().unique().all() - - if not documents: - return [] - - doc_ids = [d.id for d in documents] - - numbered = ( - select( - Chunk.id.label("chunk_id"), - Chunk.document_id, - Chunk.content, - func.row_number() - .over(partition_by=Chunk.document_id, order_by=Chunk.id) - .label("rn"), - ) - .where(Chunk.document_id.in_(doc_ids)) - .subquery("numbered") - ) - - chunk_query = ( - select(numbered.c.chunk_id, numbered.c.document_id, numbered.c.content) - .where(numbered.c.rn <= _RECENCY_MAX_CHUNKS_PER_DOC) - .order_by(numbered.c.document_id, numbered.c.chunk_id) - ) - chunk_result = await session.execute(chunk_query) - fetched_chunks = chunk_result.all() - - doc_chunks: dict[int, list[dict[str, Any]]] = {d.id: [] for d in documents} - for row in fetched_chunks: - if row.document_id in doc_chunks: - doc_chunks[row.document_id].append( - {"chunk_id": row.chunk_id, "content": row.content} - ) - - results: list[dict[str, Any]] = [] - for doc in documents: - chunks_list = doc_chunks.get(doc.id, []) - metadata = doc.document_metadata or {} - results.append( - { - "document_id": doc.id, - "content": "\n\n".join( - c["content"] for c in chunks_list if c.get("content") - ), - "score": 0.0, - "chunks": chunks_list, - "matched_chunk_ids": [], - "document": { - "id": doc.id, - "title": doc.title, - "document_type": ( - doc.document_type.value - if getattr(doc, "document_type", None) - else None - ), - "metadata": metadata, - }, - "source": ( - doc.document_type.value - if getattr(doc, "document_type", None) - else None - ), - } - ) - - logger.info( - "browse_recent_documents: %d docs returned for space=%d", - len(results), - search_space_id, - ) - return results - - async def search_knowledge_base( *, query: str, @@ -846,13 +704,10 @@ class KnowledgeBaseSearchMiddleware(AgentMiddleware): # type: ignore[type-arg] *, messages: Sequence[BaseMessage], user_text: str, - ) -> tuple[str, datetime | None, datetime | None, bool]: - """Rewrite the KB query and infer optional date filters with the LLM. - - Returns (optimized_query, start_date, end_date, is_recency_query). - """ + ) -> tuple[str, datetime | None, datetime | None]: + """Rewrite the KB query and infer optional date filters with the LLM.""" if self.llm is None: - return user_text, None, None, False + return user_text, None, None recent_conversation = _render_recent_conversation( messages, @@ -879,18 +734,15 @@ class KnowledgeBaseSearchMiddleware(AgentMiddleware): # type: ignore[type-arg] plan.start_date, plan.end_date, ) - is_recency = plan.is_recency_query _perf_log.info( - "[kb_fs_middleware] planner in %.3fs query=%r optimized=%r " - "start=%s end=%s recency=%s", + "[kb_fs_middleware] planner in %.3fs query=%r optimized=%r start=%s end=%s", loop.time() - t0, user_text[:80], optimized_query[:120], start_date.isoformat() if start_date else None, end_date.isoformat() if end_date else None, - is_recency, ) - return optimized_query, start_date, end_date, is_recency + return optimized_query, start_date, end_date except (json.JSONDecodeError, ValidationError, ValueError) as exc: logger.warning( "KB planner returned invalid output, using raw query: %s", exc @@ -898,7 +750,7 @@ class KnowledgeBaseSearchMiddleware(AgentMiddleware): # type: ignore[type-arg] except Exception as exc: # pragma: no cover - defensive fallback logger.warning("KB planner failed, using raw query: %s", exc) - return user_text, None, None, False + return user_text, None, None def before_agent( # type: ignore[override] self, @@ -937,12 +789,7 @@ class KnowledgeBaseSearchMiddleware(AgentMiddleware): # type: ignore[type-arg] t0 = _perf_log and asyncio.get_event_loop().time() existing_files = state.get("files") - ( - planned_query, - start_date, - end_date, - is_recency, - ) = await self._plan_search_inputs( + planned_query, start_date, end_date = await self._plan_search_inputs( messages=messages, user_text=user_text, ) @@ -958,28 +805,16 @@ class KnowledgeBaseSearchMiddleware(AgentMiddleware): # type: ignore[type-arg] # messages within the same agent instance. self.mentioned_document_ids = [] - # --- 2. Run KB search (recency browse or hybrid) --- - if is_recency: - doc_types = _resolve_search_types( - self.available_connectors, self.available_document_types - ) - search_results = await browse_recent_documents( - search_space_id=self.search_space_id, - document_type=doc_types, - top_k=self.top_k, - start_date=start_date, - end_date=end_date, - ) - else: - search_results = await search_knowledge_base( - query=planned_query, - search_space_id=self.search_space_id, - available_connectors=self.available_connectors, - available_document_types=self.available_document_types, - top_k=self.top_k, - start_date=start_date, - end_date=end_date, - ) + # --- 2. Run KB hybrid search --- + search_results = await search_knowledge_base( + query=planned_query, + search_space_id=self.search_space_id, + available_connectors=self.available_connectors, + available_document_types=self.available_document_types, + top_k=self.top_k, + start_date=start_date, + end_date=end_date, + ) # --- 3. Merge: mentioned first, then search (dedup by doc id) --- seen_doc_ids: set[int] = set() diff --git a/surfsense_backend/app/agents/new_chat/sandbox.py b/surfsense_backend/app/agents/new_chat/sandbox.py index efac7aae8..614a1b1b9 100644 --- a/surfsense_backend/app/agents/new_chat/sandbox.py +++ b/surfsense_backend/app/agents/new_chat/sandbox.py @@ -16,7 +16,6 @@ import contextlib import logging import os import shutil -import threading from pathlib import Path from daytona import ( @@ -56,16 +55,9 @@ class _TimeoutAwareSandbox(DaytonaSandbox): ) -> ExecuteResponse: # type: ignore[override] return await asyncio.to_thread(self.execute, command, timeout=timeout) - def download_file(self, path: str) -> bytes: - """Download a file from the sandbox filesystem.""" - return self._sandbox.fs.download_file(path) - _daytona_client: Daytona | None = None -_client_lock = threading.Lock() _sandbox_cache: dict[str, _TimeoutAwareSandbox] = {} -_sandbox_locks: dict[str, asyncio.Lock] = {} -_sandbox_locks_mu = asyncio.Lock() _seeded_files: dict[str, dict[str, str]] = {} _SANDBOX_CACHE_MAX_SIZE = 20 THREAD_LABEL_KEY = "surfsense_thread" @@ -78,15 +70,14 @@ def is_sandbox_enabled() -> bool: def _get_client() -> Daytona: global _daytona_client - with _client_lock: - if _daytona_client is None: - config = DaytonaConfig( - api_key=os.environ.get("DAYTONA_API_KEY", ""), - api_url=os.environ.get("DAYTONA_API_URL", "https://app.daytona.io/api"), - target=os.environ.get("DAYTONA_TARGET", "us"), - ) - _daytona_client = Daytona(config) - return _daytona_client + if _daytona_client is None: + config = DaytonaConfig( + api_key=os.environ.get("DAYTONA_API_KEY", ""), + api_url=os.environ.get("DAYTONA_API_URL", "https://app.daytona.io/api"), + target=os.environ.get("DAYTONA_TARGET", "us"), + ) + _daytona_client = Daytona(config) + return _daytona_client def _sandbox_create_params( @@ -138,16 +129,14 @@ def _find_or_create(thread_id: str) -> tuple[_TimeoutAwareSandbox, bool]: try: client.delete(sandbox) except Exception: - logger.debug( - "Could not delete broken sandbox %s", sandbox.id, exc_info=True - ) + logger.debug("Could not delete broken sandbox %s", sandbox.id, exc_info=True) sandbox = client.create(_sandbox_create_params(labels)) is_new = True logger.info("Created replacement sandbox: %s", sandbox.id) elif sandbox.state != SandboxState.STARTED: sandbox.wait_for_sandbox_start(timeout=60) - except DaytonaError: + except Exception: logger.info("No existing sandbox for thread %s — creating one", thread_id) sandbox = client.create(_sandbox_create_params(labels)) is_new = True @@ -156,16 +145,6 @@ def _find_or_create(thread_id: str) -> tuple[_TimeoutAwareSandbox, bool]: return _TimeoutAwareSandbox(sandbox=sandbox), is_new -async def _get_thread_lock(key: str) -> asyncio.Lock: - """Return a per-thread asyncio lock, creating one if needed.""" - async with _sandbox_locks_mu: - lock = _sandbox_locks.get(key) - if lock is None: - lock = asyncio.Lock() - _sandbox_locks[key] = lock - return lock - - async def get_or_create_sandbox( thread_id: int | str, ) -> tuple[_TimeoutAwareSandbox, bool]: @@ -173,52 +152,25 @@ async def get_or_create_sandbox( Uses an in-process cache keyed by thread_id so subsequent messages in the same conversation reuse the sandbox object without an API call. - A per-thread async lock prevents duplicate sandbox creation from - concurrent requests. Returns: Tuple of (sandbox, is_new). *is_new* is True when a fresh sandbox was created, signalling that file tracking should be reset. """ key = str(thread_id) - lock = await _get_thread_lock(key) + cached = _sandbox_cache.get(key) + if cached is not None: + logger.info("Reusing cached sandbox for thread %s", key) + return cached, False + sandbox, is_new = await asyncio.to_thread(_find_or_create, key) + _sandbox_cache[key] = sandbox - async with lock: - cached = _sandbox_cache.get(key) - if cached is not None: - logger.info("Reusing cached sandbox for thread %s", key) - return cached, False - sandbox, is_new = await asyncio.to_thread(_find_or_create, key) - _sandbox_cache[key] = sandbox + if len(_sandbox_cache) > _SANDBOX_CACHE_MAX_SIZE: + oldest_key = next(iter(_sandbox_cache)) + _sandbox_cache.pop(oldest_key, None) + logger.debug("Evicted oldest sandbox cache entry: %s", oldest_key) - if len(_sandbox_cache) > _SANDBOX_CACHE_MAX_SIZE: - oldest_key = next(iter(_sandbox_cache)) - if oldest_key != key: - evicted = _sandbox_cache.pop(oldest_key, None) - _seeded_files.pop(oldest_key, None) - logger.debug("Evicted sandbox cache entry: %s", oldest_key) - if evicted is not None: - _schedule_sandbox_delete(evicted) - - return sandbox, is_new - - -def _schedule_sandbox_delete(sandbox: _TimeoutAwareSandbox) -> None: - """Best-effort background deletion of an evicted sandbox.""" - - def _delete() -> None: - try: - client = _get_client() - client.delete(sandbox._sandbox) - logger.info("Deleted evicted sandbox: %s", sandbox._sandbox.id) - except Exception: - logger.debug("Could not delete evicted sandbox", exc_info=True) - - try: - loop = asyncio.get_running_loop() - loop.run_in_executor(None, _delete) - except RuntimeError: - pass + return sandbox, is_new async def sync_files_to_sandbox( diff --git a/surfsense_backend/app/agents/new_chat/tools/confluence/create_page.py b/surfsense_backend/app/agents/new_chat/tools/confluence/create_page.py index 095413bdb..b4d532b76 100644 --- a/surfsense_backend/app/agents/new_chat/tools/confluence/create_page.py +++ b/surfsense_backend/app/agents/new_chat/tools/confluence/create_page.py @@ -2,10 +2,10 @@ import logging from typing import Any from langchain_core.tools import tool +from langgraph.types import interrupt from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm.attributes import flag_modified -from app.agents.new_chat.tools.hitl import request_approval from app.connectors.confluence_history import ConfluenceHistoryConnector from app.services.confluence import ConfluenceToolMetadataService @@ -65,28 +65,54 @@ def create_create_confluence_page_tool( "connector_type": "confluence", } - result = request_approval( - action_type="confluence_page_creation", - tool_name="create_confluence_page", - params={ - "title": title, - "content": content, - "space_id": space_id, - "connector_id": connector_id, - }, - context=context, + approval = interrupt( + { + "type": "confluence_page_creation", + "action": { + "tool": "create_confluence_page", + "params": { + "title": title, + "content": content, + "space_id": space_id, + "connector_id": connector_id, + }, + }, + "context": context, + } ) - if result.rejected: + decisions_raw = ( + approval.get("decisions", []) if isinstance(approval, dict) else [] + ) + decisions = ( + decisions_raw if isinstance(decisions_raw, list) else [decisions_raw] + ) + decisions = [d for d in decisions if isinstance(d, dict)] + if not decisions: + return {"status": "error", "message": "No approval decision received"} + + decision = decisions[0] + decision_type = decision.get("type") or decision.get("decision_type") + + if decision_type == "reject": return { "status": "rejected", - "message": "User declined. Do not retry or suggest alternatives.", + "message": "User declined. The page was not created.", } - final_title = result.params.get("title", title) - final_content = result.params.get("content", content) or "" - final_space_id = result.params.get("space_id", space_id) - final_connector_id = result.params.get("connector_id", connector_id) + final_params: dict[str, Any] = {} + edited_action = decision.get("edited_action") + if isinstance(edited_action, dict): + edited_args = edited_action.get("args") + if isinstance(edited_args, dict): + final_params = edited_args + elif isinstance(decision.get("args"), dict): + final_params = decision["args"] + + final_title = final_params.get("title", title) + final_content = final_params.get("content", content) or "" + final_space_id = final_params.get("space_id", space_id) + final_connector_id = final_params.get("connector_id", connector_id) if not final_title or not final_title.strip(): return {"status": "error", "message": "Page title cannot be empty."} diff --git a/surfsense_backend/app/agents/new_chat/tools/confluence/delete_page.py b/surfsense_backend/app/agents/new_chat/tools/confluence/delete_page.py index 7c03c2760..ba1dae653 100644 --- a/surfsense_backend/app/agents/new_chat/tools/confluence/delete_page.py +++ b/surfsense_backend/app/agents/new_chat/tools/confluence/delete_page.py @@ -2,10 +2,10 @@ import logging from typing import Any from langchain_core.tools import tool +from langgraph.types import interrupt from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm.attributes import flag_modified -from app.agents.new_chat.tools.hitl import request_approval from app.connectors.confluence_history import ConfluenceHistoryConnector from app.services.confluence import ConfluenceToolMetadataService @@ -74,28 +74,54 @@ def create_delete_confluence_page_tool( document_id = page_data["document_id"] connector_id_from_context = context.get("account", {}).get("id") - result = request_approval( - action_type="confluence_page_deletion", - tool_name="delete_confluence_page", - params={ - "page_id": page_id, - "connector_id": connector_id_from_context, - "delete_from_kb": delete_from_kb, - }, - context=context, + approval = interrupt( + { + "type": "confluence_page_deletion", + "action": { + "tool": "delete_confluence_page", + "params": { + "page_id": page_id, + "connector_id": connector_id_from_context, + "delete_from_kb": delete_from_kb, + }, + }, + "context": context, + } ) - if result.rejected: + decisions_raw = ( + approval.get("decisions", []) if isinstance(approval, dict) else [] + ) + decisions = ( + decisions_raw if isinstance(decisions_raw, list) else [decisions_raw] + ) + decisions = [d for d in decisions if isinstance(d, dict)] + if not decisions: + return {"status": "error", "message": "No approval decision received"} + + decision = decisions[0] + decision_type = decision.get("type") or decision.get("decision_type") + + if decision_type == "reject": return { "status": "rejected", - "message": "User declined. Do not retry or suggest alternatives.", + "message": "User declined. The page was not deleted.", } - final_page_id = result.params.get("page_id", page_id) - final_connector_id = result.params.get( + final_params: dict[str, Any] = {} + edited_action = decision.get("edited_action") + if isinstance(edited_action, dict): + edited_args = edited_action.get("args") + if isinstance(edited_args, dict): + final_params = edited_args + elif isinstance(decision.get("args"), dict): + final_params = decision["args"] + + final_page_id = final_params.get("page_id", page_id) + final_connector_id = final_params.get( "connector_id", connector_id_from_context ) - final_delete_from_kb = result.params.get("delete_from_kb", delete_from_kb) + final_delete_from_kb = final_params.get("delete_from_kb", delete_from_kb) from sqlalchemy.future import select diff --git a/surfsense_backend/app/agents/new_chat/tools/confluence/update_page.py b/surfsense_backend/app/agents/new_chat/tools/confluence/update_page.py index 791d0d8c5..913896f83 100644 --- a/surfsense_backend/app/agents/new_chat/tools/confluence/update_page.py +++ b/surfsense_backend/app/agents/new_chat/tools/confluence/update_page.py @@ -2,10 +2,10 @@ import logging from typing import Any from langchain_core.tools import tool +from langgraph.types import interrupt from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm.attributes import flag_modified -from app.agents.new_chat.tools.hitl import request_approval from app.connectors.confluence_history import ConfluenceHistoryConnector from app.services.confluence import ConfluenceToolMetadataService @@ -78,36 +78,62 @@ def create_update_confluence_page_tool( document_id = page_data.get("document_id") connector_id_from_context = context.get("account", {}).get("id") - result = request_approval( - action_type="confluence_page_update", - tool_name="update_confluence_page", - params={ - "page_id": page_id, - "document_id": document_id, - "new_title": new_title, - "new_content": new_content, - "version": current_version, - "connector_id": connector_id_from_context, - }, - context=context, + approval = interrupt( + { + "type": "confluence_page_update", + "action": { + "tool": "update_confluence_page", + "params": { + "page_id": page_id, + "document_id": document_id, + "new_title": new_title, + "new_content": new_content, + "version": current_version, + "connector_id": connector_id_from_context, + }, + }, + "context": context, + } ) - if result.rejected: + decisions_raw = ( + approval.get("decisions", []) if isinstance(approval, dict) else [] + ) + decisions = ( + decisions_raw if isinstance(decisions_raw, list) else [decisions_raw] + ) + decisions = [d for d in decisions if isinstance(d, dict)] + if not decisions: + return {"status": "error", "message": "No approval decision received"} + + decision = decisions[0] + decision_type = decision.get("type") or decision.get("decision_type") + + if decision_type == "reject": return { "status": "rejected", - "message": "User declined. Do not retry or suggest alternatives.", + "message": "User declined. The page was not updated.", } - final_page_id = result.params.get("page_id", page_id) - final_title = result.params.get("new_title", new_title) or current_title - final_content = result.params.get("new_content", new_content) + final_params: dict[str, Any] = {} + edited_action = decision.get("edited_action") + if isinstance(edited_action, dict): + edited_args = edited_action.get("args") + if isinstance(edited_args, dict): + final_params = edited_args + elif isinstance(decision.get("args"), dict): + final_params = decision["args"] + + final_page_id = final_params.get("page_id", page_id) + final_title = final_params.get("new_title", new_title) or current_title + final_content = final_params.get("new_content", new_content) if final_content is None: final_content = current_body - final_version = result.params.get("version", current_version) - final_connector_id = result.params.get( + final_version = final_params.get("version", current_version) + final_connector_id = final_params.get( "connector_id", connector_id_from_context ) - final_document_id = result.params.get("document_id", document_id) + final_document_id = final_params.get("document_id", document_id) from sqlalchemy.future import select diff --git a/surfsense_backend/app/agents/new_chat/tools/dropbox/create_file.py b/surfsense_backend/app/agents/new_chat/tools/dropbox/create_file.py index 22d8a8a27..ed8034861 100644 --- a/surfsense_backend/app/agents/new_chat/tools/dropbox/create_file.py +++ b/surfsense_backend/app/agents/new_chat/tools/dropbox/create_file.py @@ -5,10 +5,10 @@ from pathlib import Path from typing import Any, Literal from langchain_core.tools import tool +from langgraph.types import interrupt from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.future import select -from app.agents.new_chat.tools.hitl import request_approval from app.connectors.dropbox.client import DropboxClient from app.db import SearchSourceConnector, SearchSourceConnectorType @@ -159,30 +159,56 @@ def create_create_dropbox_file_tool( "supported_types": _SUPPORTED_TYPES, } - result = request_approval( - action_type="dropbox_file_creation", - tool_name="create_dropbox_file", - params={ - "name": name, - "file_type": file_type, - "content": content, - "connector_id": None, - "parent_folder_path": None, - }, - context=context, + approval = interrupt( + { + "type": "dropbox_file_creation", + "action": { + "tool": "create_dropbox_file", + "params": { + "name": name, + "file_type": file_type, + "content": content, + "connector_id": None, + "parent_folder_path": None, + }, + }, + "context": context, + } ) - if result.rejected: + decisions_raw = ( + approval.get("decisions", []) if isinstance(approval, dict) else [] + ) + decisions = ( + decisions_raw if isinstance(decisions_raw, list) else [decisions_raw] + ) + decisions = [d for d in decisions if isinstance(d, dict)] + if not decisions: + return {"status": "error", "message": "No approval decision received"} + + decision = decisions[0] + decision_type = decision.get("type") or decision.get("decision_type") + + if decision_type == "reject": return { "status": "rejected", - "message": "User declined. Do not retry or suggest alternatives.", + "message": "User declined. The file was not created.", } - final_name = result.params.get("name", name) - final_file_type = result.params.get("file_type", file_type) - final_content = result.params.get("content", content) - final_connector_id = result.params.get("connector_id") - final_parent_folder_path = result.params.get("parent_folder_path") + final_params: dict[str, Any] = {} + edited_action = decision.get("edited_action") + if isinstance(edited_action, dict): + edited_args = edited_action.get("args") + if isinstance(edited_args, dict): + final_params = edited_args + elif isinstance(decision.get("args"), dict): + final_params = decision["args"] + + final_name = final_params.get("name", name) + final_file_type = final_params.get("file_type", file_type) + final_content = final_params.get("content", content) + final_connector_id = final_params.get("connector_id") + final_parent_folder_path = final_params.get("parent_folder_path") if not final_name or not final_name.strip(): return {"status": "error", "message": "File name cannot be empty."} diff --git a/surfsense_backend/app/agents/new_chat/tools/dropbox/trash_file.py b/surfsense_backend/app/agents/new_chat/tools/dropbox/trash_file.py index 12559b57a..e15dc3092 100644 --- a/surfsense_backend/app/agents/new_chat/tools/dropbox/trash_file.py +++ b/surfsense_backend/app/agents/new_chat/tools/dropbox/trash_file.py @@ -2,11 +2,11 @@ import logging from typing import Any from langchain_core.tools import tool +from langgraph.types import interrupt from sqlalchemy import String, and_, cast, func from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.future import select -from app.agents.new_chat.tools.hitl import request_approval from app.connectors.dropbox.client import DropboxClient from app.db import ( Document, @@ -174,26 +174,53 @@ def create_delete_dropbox_file_tool( }, } - result = request_approval( - action_type="dropbox_file_trash", - tool_name="delete_dropbox_file", - params={ - "file_path": file_path, - "connector_id": connector.id, - "delete_from_kb": delete_from_kb, - }, - context=context, + approval = interrupt( + { + "type": "dropbox_file_trash", + "action": { + "tool": "delete_dropbox_file", + "params": { + "file_path": file_path, + "connector_id": connector.id, + "delete_from_kb": delete_from_kb, + }, + }, + "context": context, + } ) - if result.rejected: + decisions_raw = ( + approval.get("decisions", []) if isinstance(approval, dict) else [] + ) + decisions = ( + decisions_raw if isinstance(decisions_raw, list) else [decisions_raw] + ) + decisions = [d for d in decisions if isinstance(d, dict)] + if not decisions: + return {"status": "error", "message": "No approval decision received"} + + decision = decisions[0] + decision_type = decision.get("type") or decision.get("decision_type") + logger.info(f"User decision: {decision_type}") + + if decision_type == "reject": return { "status": "rejected", - "message": "User declined. Do not retry or suggest alternatives.", + "message": "User declined. The file was not deleted. Do not ask again or suggest alternatives.", } - final_file_path = result.params.get("file_path", file_path) - final_connector_id = result.params.get("connector_id", connector.id) - final_delete_from_kb = result.params.get("delete_from_kb", delete_from_kb) + final_params: dict[str, Any] = {} + edited_action = decision.get("edited_action") + if isinstance(edited_action, dict): + edited_args = edited_action.get("args") + if isinstance(edited_args, dict): + final_params = edited_args + elif isinstance(decision.get("args"), dict): + final_params = decision["args"] + + final_file_path = final_params.get("file_path", file_path) + final_connector_id = final_params.get("connector_id", connector.id) + final_delete_from_kb = final_params.get("delete_from_kb", delete_from_kb) if final_connector_id != connector.id: result = await db_session.execute( diff --git a/surfsense_backend/app/agents/new_chat/tools/gmail/create_draft.py b/surfsense_backend/app/agents/new_chat/tools/gmail/create_draft.py index 0bd044695..a812f621a 100644 --- a/surfsense_backend/app/agents/new_chat/tools/gmail/create_draft.py +++ b/surfsense_backend/app/agents/new_chat/tools/gmail/create_draft.py @@ -6,9 +6,9 @@ from email.mime.text import MIMEText from typing import Any from langchain_core.tools import tool +from langgraph.types import interrupt from sqlalchemy.ext.asyncio import AsyncSession -from app.agents.new_chat.tools.hitl import request_approval from app.services.gmail import GmailToolMetadataService logger = logging.getLogger(__name__) @@ -85,32 +85,60 @@ def create_create_gmail_draft_tool( logger.info( f"Requesting approval for creating Gmail draft: to='{to}', subject='{subject}'" ) - result = request_approval( - action_type="gmail_draft_creation", - tool_name="create_gmail_draft", - params={ - "to": to, - "subject": subject, - "body": body, - "cc": cc, - "bcc": bcc, - "connector_id": None, - }, - context=context, + approval = interrupt( + { + "type": "gmail_draft_creation", + "action": { + "tool": "create_gmail_draft", + "params": { + "to": to, + "subject": subject, + "body": body, + "cc": cc, + "bcc": bcc, + "connector_id": None, + }, + }, + "context": context, + } ) - if result.rejected: + decisions_raw = ( + approval.get("decisions", []) if isinstance(approval, dict) else [] + ) + decisions = ( + decisions_raw if isinstance(decisions_raw, list) else [decisions_raw] + ) + decisions = [d for d in decisions if isinstance(d, dict)] + if not decisions: + logger.warning("No approval decision received") + return {"status": "error", "message": "No approval decision received"} + + decision = decisions[0] + decision_type = decision.get("type") or decision.get("decision_type") + logger.info(f"User decision: {decision_type}") + + if decision_type == "reject": return { "status": "rejected", "message": "User declined. The draft was not created. Do not ask again or suggest alternatives.", } - final_to = result.params.get("to", to) - final_subject = result.params.get("subject", subject) - final_body = result.params.get("body", body) - final_cc = result.params.get("cc", cc) - final_bcc = result.params.get("bcc", bcc) - final_connector_id = result.params.get("connector_id") + final_params: dict[str, Any] = {} + edited_action = decision.get("edited_action") + if isinstance(edited_action, dict): + edited_args = edited_action.get("args") + if isinstance(edited_args, dict): + final_params = edited_args + elif isinstance(decision.get("args"), dict): + final_params = decision["args"] + + final_to = final_params.get("to", to) + final_subject = final_params.get("subject", subject) + final_body = final_params.get("body", body) + final_cc = final_params.get("cc", cc) + final_bcc = final_params.get("bcc", bcc) + final_connector_id = final_params.get("connector_id") from sqlalchemy.future import select diff --git a/surfsense_backend/app/agents/new_chat/tools/gmail/send_email.py b/surfsense_backend/app/agents/new_chat/tools/gmail/send_email.py index c3f0999f4..2599578bd 100644 --- a/surfsense_backend/app/agents/new_chat/tools/gmail/send_email.py +++ b/surfsense_backend/app/agents/new_chat/tools/gmail/send_email.py @@ -6,9 +6,9 @@ from email.mime.text import MIMEText from typing import Any from langchain_core.tools import tool +from langgraph.types import interrupt from sqlalchemy.ext.asyncio import AsyncSession -from app.agents.new_chat.tools.hitl import request_approval from app.services.gmail import GmailToolMetadataService logger = logging.getLogger(__name__) @@ -86,32 +86,60 @@ def create_send_gmail_email_tool( logger.info( f"Requesting approval for sending Gmail email: to='{to}', subject='{subject}'" ) - result = request_approval( - action_type="gmail_email_send", - tool_name="send_gmail_email", - params={ - "to": to, - "subject": subject, - "body": body, - "cc": cc, - "bcc": bcc, - "connector_id": None, - }, - context=context, + approval = interrupt( + { + "type": "gmail_email_send", + "action": { + "tool": "send_gmail_email", + "params": { + "to": to, + "subject": subject, + "body": body, + "cc": cc, + "bcc": bcc, + "connector_id": None, + }, + }, + "context": context, + } ) - if result.rejected: + decisions_raw = ( + approval.get("decisions", []) if isinstance(approval, dict) else [] + ) + decisions = ( + decisions_raw if isinstance(decisions_raw, list) else [decisions_raw] + ) + decisions = [d for d in decisions if isinstance(d, dict)] + if not decisions: + logger.warning("No approval decision received") + return {"status": "error", "message": "No approval decision received"} + + decision = decisions[0] + decision_type = decision.get("type") or decision.get("decision_type") + logger.info(f"User decision: {decision_type}") + + if decision_type == "reject": return { "status": "rejected", "message": "User declined. The email was not sent. Do not ask again or suggest alternatives.", } - final_to = result.params.get("to", to) - final_subject = result.params.get("subject", subject) - final_body = result.params.get("body", body) - final_cc = result.params.get("cc", cc) - final_bcc = result.params.get("bcc", bcc) - final_connector_id = result.params.get("connector_id") + final_params: dict[str, Any] = {} + edited_action = decision.get("edited_action") + if isinstance(edited_action, dict): + edited_args = edited_action.get("args") + if isinstance(edited_args, dict): + final_params = edited_args + elif isinstance(decision.get("args"), dict): + final_params = decision["args"] + + final_to = final_params.get("to", to) + final_subject = final_params.get("subject", subject) + final_body = final_params.get("body", body) + final_cc = final_params.get("cc", cc) + final_bcc = final_params.get("bcc", bcc) + final_connector_id = final_params.get("connector_id") from sqlalchemy.future import select diff --git a/surfsense_backend/app/agents/new_chat/tools/gmail/trash_email.py b/surfsense_backend/app/agents/new_chat/tools/gmail/trash_email.py index 1f1f6227a..146020845 100644 --- a/surfsense_backend/app/agents/new_chat/tools/gmail/trash_email.py +++ b/surfsense_backend/app/agents/new_chat/tools/gmail/trash_email.py @@ -4,9 +4,9 @@ from datetime import datetime from typing import Any from langchain_core.tools import tool +from langgraph.types import interrupt from sqlalchemy.ext.asyncio import AsyncSession -from app.agents.new_chat.tools.hitl import request_approval from app.services.gmail import GmailToolMetadataService logger = logging.getLogger(__name__) @@ -101,28 +101,56 @@ def create_trash_gmail_email_tool( logger.info( f"Requesting approval for trashing Gmail email: '{email_subject_or_id}' (message_id={message_id}, delete_from_kb={delete_from_kb})" ) - result = request_approval( - action_type="gmail_email_trash", - tool_name="trash_gmail_email", - params={ - "message_id": message_id, - "connector_id": connector_id_from_context, - "delete_from_kb": delete_from_kb, - }, - context=context, + approval = interrupt( + { + "type": "gmail_email_trash", + "action": { + "tool": "trash_gmail_email", + "params": { + "message_id": message_id, + "connector_id": connector_id_from_context, + "delete_from_kb": delete_from_kb, + }, + }, + "context": context, + } ) - if result.rejected: + decisions_raw = ( + approval.get("decisions", []) if isinstance(approval, dict) else [] + ) + decisions = ( + decisions_raw if isinstance(decisions_raw, list) else [decisions_raw] + ) + decisions = [d for d in decisions if isinstance(d, dict)] + if not decisions: + logger.warning("No approval decision received") + return {"status": "error", "message": "No approval decision received"} + + decision = decisions[0] + decision_type = decision.get("type") or decision.get("decision_type") + logger.info(f"User decision: {decision_type}") + + if decision_type == "reject": return { "status": "rejected", "message": "User declined. The email was not trashed. Do not ask again or suggest alternatives.", } - final_message_id = result.params.get("message_id", message_id) - final_connector_id = result.params.get( + edited_action = decision.get("edited_action") + final_params: dict[str, Any] = {} + if isinstance(edited_action, dict): + edited_args = edited_action.get("args") + if isinstance(edited_args, dict): + final_params = edited_args + elif isinstance(decision.get("args"), dict): + final_params = decision["args"] + + final_message_id = final_params.get("message_id", message_id) + final_connector_id = final_params.get( "connector_id", connector_id_from_context ) - final_delete_from_kb = result.params.get("delete_from_kb", delete_from_kb) + final_delete_from_kb = final_params.get("delete_from_kb", delete_from_kb) if not final_connector_id: return { diff --git a/surfsense_backend/app/agents/new_chat/tools/gmail/update_draft.py b/surfsense_backend/app/agents/new_chat/tools/gmail/update_draft.py index 91178cd21..28deec2b4 100644 --- a/surfsense_backend/app/agents/new_chat/tools/gmail/update_draft.py +++ b/surfsense_backend/app/agents/new_chat/tools/gmail/update_draft.py @@ -6,9 +6,9 @@ from email.mime.text import MIMEText from typing import Any from langchain_core.tools import tool +from langgraph.types import interrupt from sqlalchemy.ext.asyncio import AsyncSession -from app.agents.new_chat.tools.hitl import request_approval from app.services.gmail import GmailToolMetadataService logger = logging.getLogger(__name__) @@ -122,37 +122,65 @@ def create_update_gmail_draft_tool( f"Requesting approval for updating Gmail draft: '{original_subject}' " f"(message_id={message_id}, draft_id={draft_id_from_context})" ) - result = request_approval( - action_type="gmail_draft_update", - tool_name="update_gmail_draft", - params={ - "message_id": message_id, - "draft_id": draft_id_from_context, - "to": final_to_default, - "subject": final_subject_default, - "body": body, - "cc": cc, - "bcc": bcc, - "connector_id": connector_id_from_context, - }, - context=context, + approval = interrupt( + { + "type": "gmail_draft_update", + "action": { + "tool": "update_gmail_draft", + "params": { + "message_id": message_id, + "draft_id": draft_id_from_context, + "to": final_to_default, + "subject": final_subject_default, + "body": body, + "cc": cc, + "bcc": bcc, + "connector_id": connector_id_from_context, + }, + }, + "context": context, + } ) - if result.rejected: + decisions_raw = ( + approval.get("decisions", []) if isinstance(approval, dict) else [] + ) + decisions = ( + decisions_raw if isinstance(decisions_raw, list) else [decisions_raw] + ) + decisions = [d for d in decisions if isinstance(d, dict)] + if not decisions: + logger.warning("No approval decision received") + return {"status": "error", "message": "No approval decision received"} + + decision = decisions[0] + decision_type = decision.get("type") or decision.get("decision_type") + logger.info(f"User decision: {decision_type}") + + if decision_type == "reject": return { "status": "rejected", "message": "User declined. The draft was not updated. Do not ask again or suggest alternatives.", } - final_to = result.params.get("to", final_to_default) - final_subject = result.params.get("subject", final_subject_default) - final_body = result.params.get("body", body) - final_cc = result.params.get("cc", cc) - final_bcc = result.params.get("bcc", bcc) - final_connector_id = result.params.get( + final_params: dict[str, Any] = {} + edited_action = decision.get("edited_action") + if isinstance(edited_action, dict): + edited_args = edited_action.get("args") + if isinstance(edited_args, dict): + final_params = edited_args + elif isinstance(decision.get("args"), dict): + final_params = decision["args"] + + final_to = final_params.get("to", final_to_default) + final_subject = final_params.get("subject", final_subject_default) + final_body = final_params.get("body", body) + final_cc = final_params.get("cc", cc) + final_bcc = final_params.get("bcc", bcc) + final_connector_id = final_params.get( "connector_id", connector_id_from_context ) - final_draft_id = result.params.get("draft_id", draft_id_from_context) + final_draft_id = final_params.get("draft_id", draft_id_from_context) if not final_connector_id: return { diff --git a/surfsense_backend/app/agents/new_chat/tools/google_calendar/create_event.py b/surfsense_backend/app/agents/new_chat/tools/google_calendar/create_event.py index 37bcf083e..592ced5ec 100644 --- a/surfsense_backend/app/agents/new_chat/tools/google_calendar/create_event.py +++ b/surfsense_backend/app/agents/new_chat/tools/google_calendar/create_event.py @@ -6,9 +6,9 @@ from typing import Any from google.oauth2.credentials import Credentials from googleapiclient.discovery import build from langchain_core.tools import tool +from langgraph.types import interrupt from sqlalchemy.ext.asyncio import AsyncSession -from app.agents.new_chat.tools.hitl import request_approval from app.services.google_calendar import GoogleCalendarToolMetadataService logger = logging.getLogger(__name__) @@ -90,35 +90,63 @@ def create_create_calendar_event_tool( logger.info( f"Requesting approval for creating calendar event: summary='{summary}'" ) - result = request_approval( - action_type="google_calendar_event_creation", - tool_name="create_calendar_event", - params={ - "summary": summary, - "start_datetime": start_datetime, - "end_datetime": end_datetime, - "description": description, - "location": location, - "attendees": attendees, - "timezone": context.get("timezone"), - "connector_id": None, - }, - context=context, + approval = interrupt( + { + "type": "google_calendar_event_creation", + "action": { + "tool": "create_calendar_event", + "params": { + "summary": summary, + "start_datetime": start_datetime, + "end_datetime": end_datetime, + "description": description, + "location": location, + "attendees": attendees, + "timezone": context.get("timezone"), + "connector_id": None, + }, + }, + "context": context, + } ) - if result.rejected: + decisions_raw = ( + approval.get("decisions", []) if isinstance(approval, dict) else [] + ) + decisions = ( + decisions_raw if isinstance(decisions_raw, list) else [decisions_raw] + ) + decisions = [d for d in decisions if isinstance(d, dict)] + if not decisions: + logger.warning("No approval decision received") + return {"status": "error", "message": "No approval decision received"} + + decision = decisions[0] + decision_type = decision.get("type") or decision.get("decision_type") + logger.info(f"User decision: {decision_type}") + + if decision_type == "reject": return { "status": "rejected", "message": "User declined. The event was not created. Do not ask again or suggest alternatives.", } - final_summary = result.params.get("summary", summary) - final_start_datetime = result.params.get("start_datetime", start_datetime) - final_end_datetime = result.params.get("end_datetime", end_datetime) - final_description = result.params.get("description", description) - final_location = result.params.get("location", location) - final_attendees = result.params.get("attendees", attendees) - final_connector_id = result.params.get("connector_id") + final_params: dict[str, Any] = {} + edited_action = decision.get("edited_action") + if isinstance(edited_action, dict): + edited_args = edited_action.get("args") + if isinstance(edited_args, dict): + final_params = edited_args + elif isinstance(decision.get("args"), dict): + final_params = decision["args"] + + final_summary = final_params.get("summary", summary) + final_start_datetime = final_params.get("start_datetime", start_datetime) + final_end_datetime = final_params.get("end_datetime", end_datetime) + final_description = final_params.get("description", description) + final_location = final_params.get("location", location) + final_attendees = final_params.get("attendees", attendees) + final_connector_id = final_params.get("connector_id") if not final_summary or not final_summary.strip(): return {"status": "error", "message": "Event summary cannot be empty."} diff --git a/surfsense_backend/app/agents/new_chat/tools/google_calendar/delete_event.py b/surfsense_backend/app/agents/new_chat/tools/google_calendar/delete_event.py index 4d9d69b4b..8b088487c 100644 --- a/surfsense_backend/app/agents/new_chat/tools/google_calendar/delete_event.py +++ b/surfsense_backend/app/agents/new_chat/tools/google_calendar/delete_event.py @@ -6,9 +6,9 @@ from typing import Any from google.oauth2.credentials import Credentials from googleapiclient.discovery import build from langchain_core.tools import tool +from langgraph.types import interrupt from sqlalchemy.ext.asyncio import AsyncSession -from app.agents.new_chat.tools.hitl import request_approval from app.services.google_calendar import GoogleCalendarToolMetadataService logger = logging.getLogger(__name__) @@ -100,28 +100,56 @@ def create_delete_calendar_event_tool( logger.info( f"Requesting approval for deleting calendar event: '{event_title_or_id}' (event_id={event_id}, delete_from_kb={delete_from_kb})" ) - result = request_approval( - action_type="google_calendar_event_deletion", - tool_name="delete_calendar_event", - params={ - "event_id": event_id, - "connector_id": connector_id_from_context, - "delete_from_kb": delete_from_kb, - }, - context=context, + approval = interrupt( + { + "type": "google_calendar_event_deletion", + "action": { + "tool": "delete_calendar_event", + "params": { + "event_id": event_id, + "connector_id": connector_id_from_context, + "delete_from_kb": delete_from_kb, + }, + }, + "context": context, + } ) - if result.rejected: + decisions_raw = ( + approval.get("decisions", []) if isinstance(approval, dict) else [] + ) + decisions = ( + decisions_raw if isinstance(decisions_raw, list) else [decisions_raw] + ) + decisions = [d for d in decisions if isinstance(d, dict)] + if not decisions: + logger.warning("No approval decision received") + return {"status": "error", "message": "No approval decision received"} + + decision = decisions[0] + decision_type = decision.get("type") or decision.get("decision_type") + logger.info(f"User decision: {decision_type}") + + if decision_type == "reject": return { "status": "rejected", "message": "User declined. The event was not deleted. Do not ask again or suggest alternatives.", } - final_event_id = result.params.get("event_id", event_id) - final_connector_id = result.params.get( + edited_action = decision.get("edited_action") + final_params: dict[str, Any] = {} + if isinstance(edited_action, dict): + edited_args = edited_action.get("args") + if isinstance(edited_args, dict): + final_params = edited_args + elif isinstance(decision.get("args"), dict): + final_params = decision["args"] + + final_event_id = final_params.get("event_id", event_id) + final_connector_id = final_params.get( "connector_id", connector_id_from_context ) - final_delete_from_kb = result.params.get("delete_from_kb", delete_from_kb) + final_delete_from_kb = final_params.get("delete_from_kb", delete_from_kb) if not final_connector_id: return { diff --git a/surfsense_backend/app/agents/new_chat/tools/google_calendar/update_event.py b/surfsense_backend/app/agents/new_chat/tools/google_calendar/update_event.py index 259f52bba..ed826f1b8 100644 --- a/surfsense_backend/app/agents/new_chat/tools/google_calendar/update_event.py +++ b/surfsense_backend/app/agents/new_chat/tools/google_calendar/update_event.py @@ -6,9 +6,9 @@ from typing import Any from google.oauth2.credentials import Credentials from googleapiclient.discovery import build from langchain_core.tools import tool +from langgraph.types import interrupt from sqlalchemy.ext.asyncio import AsyncSession -from app.agents.new_chat.tools.hitl import request_approval from app.services.google_calendar import GoogleCalendarToolMetadataService logger = logging.getLogger(__name__) @@ -116,45 +116,71 @@ def create_update_calendar_event_tool( logger.info( f"Requesting approval for updating calendar event: '{event_title_or_id}' (event_id={event_id})" ) - result = request_approval( - action_type="google_calendar_event_update", - tool_name="update_calendar_event", - params={ - "event_id": event_id, - "document_id": document_id, - "connector_id": connector_id_from_context, - "new_summary": new_summary, - "new_start_datetime": new_start_datetime, - "new_end_datetime": new_end_datetime, - "new_description": new_description, - "new_location": new_location, - "new_attendees": new_attendees, - }, - context=context, + approval = interrupt( + { + "type": "google_calendar_event_update", + "action": { + "tool": "update_calendar_event", + "params": { + "event_id": event_id, + "document_id": document_id, + "connector_id": connector_id_from_context, + "new_summary": new_summary, + "new_start_datetime": new_start_datetime, + "new_end_datetime": new_end_datetime, + "new_description": new_description, + "new_location": new_location, + "new_attendees": new_attendees, + }, + }, + "context": context, + } ) - if result.rejected: + decisions_raw = ( + approval.get("decisions", []) if isinstance(approval, dict) else [] + ) + decisions = ( + decisions_raw if isinstance(decisions_raw, list) else [decisions_raw] + ) + decisions = [d for d in decisions if isinstance(d, dict)] + if not decisions: + logger.warning("No approval decision received") + return {"status": "error", "message": "No approval decision received"} + + decision = decisions[0] + decision_type = decision.get("type") or decision.get("decision_type") + logger.info(f"User decision: {decision_type}") + + if decision_type == "reject": return { "status": "rejected", "message": "User declined. The event was not updated. Do not ask again or suggest alternatives.", } - final_event_id = result.params.get("event_id", event_id) - final_connector_id = result.params.get( + edited_action = decision.get("edited_action") + final_params: dict[str, Any] = {} + if isinstance(edited_action, dict): + edited_args = edited_action.get("args") + if isinstance(edited_args, dict): + final_params = edited_args + elif isinstance(decision.get("args"), dict): + final_params = decision["args"] + + final_event_id = final_params.get("event_id", event_id) + final_connector_id = final_params.get( "connector_id", connector_id_from_context ) - final_new_summary = result.params.get("new_summary", new_summary) - final_new_start_datetime = result.params.get( + final_new_summary = final_params.get("new_summary", new_summary) + final_new_start_datetime = final_params.get( "new_start_datetime", new_start_datetime ) - final_new_end_datetime = result.params.get( + final_new_end_datetime = final_params.get( "new_end_datetime", new_end_datetime ) - final_new_description = result.params.get( - "new_description", new_description - ) - final_new_location = result.params.get("new_location", new_location) - final_new_attendees = result.params.get("new_attendees", new_attendees) + final_new_description = final_params.get("new_description", new_description) + final_new_location = final_params.get("new_location", new_location) + final_new_attendees = final_params.get("new_attendees", new_attendees) if not final_connector_id: return { diff --git a/surfsense_backend/app/agents/new_chat/tools/google_drive/create_file.py b/surfsense_backend/app/agents/new_chat/tools/google_drive/create_file.py index f36db8f3f..a4fee0965 100644 --- a/surfsense_backend/app/agents/new_chat/tools/google_drive/create_file.py +++ b/surfsense_backend/app/agents/new_chat/tools/google_drive/create_file.py @@ -3,9 +3,9 @@ from typing import Any, Literal from googleapiclient.errors import HttpError from langchain_core.tools import tool +from langgraph.types import interrupt from sqlalchemy.ext.asyncio import AsyncSession -from app.agents.new_chat.tools.hitl import request_approval from app.connectors.google_drive.client import GoogleDriveClient from app.connectors.google_drive.file_types import GOOGLE_DOC, GOOGLE_SHEET from app.services.google_drive import GoogleDriveToolMetadataService @@ -99,30 +99,58 @@ def create_create_google_drive_file_tool( logger.info( f"Requesting approval for creating Google Drive file: name='{name}', type='{file_type}'" ) - result = request_approval( - action_type="google_drive_file_creation", - tool_name="create_google_drive_file", - params={ - "name": name, - "file_type": file_type, - "content": content, - "connector_id": None, - "parent_folder_id": None, - }, - context=context, + approval = interrupt( + { + "type": "google_drive_file_creation", + "action": { + "tool": "create_google_drive_file", + "params": { + "name": name, + "file_type": file_type, + "content": content, + "connector_id": None, + "parent_folder_id": None, + }, + }, + "context": context, + } ) - if result.rejected: + decisions_raw = ( + approval.get("decisions", []) if isinstance(approval, dict) else [] + ) + decisions = ( + decisions_raw if isinstance(decisions_raw, list) else [decisions_raw] + ) + decisions = [d for d in decisions if isinstance(d, dict)] + if not decisions: + logger.warning("No approval decision received") + return {"status": "error", "message": "No approval decision received"} + + decision = decisions[0] + decision_type = decision.get("type") or decision.get("decision_type") + logger.info(f"User decision: {decision_type}") + + if decision_type == "reject": return { "status": "rejected", "message": "User declined. The file was not created. Do not ask again or suggest alternatives.", } - final_name = result.params.get("name", name) - final_file_type = result.params.get("file_type", file_type) - final_content = result.params.get("content", content) - final_connector_id = result.params.get("connector_id") - final_parent_folder_id = result.params.get("parent_folder_id") + final_params: dict[str, Any] = {} + edited_action = decision.get("edited_action") + if isinstance(edited_action, dict): + edited_args = edited_action.get("args") + if isinstance(edited_args, dict): + final_params = edited_args + elif isinstance(decision.get("args"), dict): + final_params = decision["args"] + + final_name = final_params.get("name", name) + final_file_type = final_params.get("file_type", file_type) + final_content = final_params.get("content", content) + final_connector_id = final_params.get("connector_id") + final_parent_folder_id = final_params.get("parent_folder_id") if not final_name or not final_name.strip(): return {"status": "error", "message": "File name cannot be empty."} diff --git a/surfsense_backend/app/agents/new_chat/tools/google_drive/trash_file.py b/surfsense_backend/app/agents/new_chat/tools/google_drive/trash_file.py index 832afff0d..fdf7f9cd3 100644 --- a/surfsense_backend/app/agents/new_chat/tools/google_drive/trash_file.py +++ b/surfsense_backend/app/agents/new_chat/tools/google_drive/trash_file.py @@ -3,9 +3,9 @@ from typing import Any from googleapiclient.errors import HttpError from langchain_core.tools import tool +from langgraph.types import interrupt from sqlalchemy.ext.asyncio import AsyncSession -from app.agents.new_chat.tools.hitl import request_approval from app.connectors.google_drive.client import GoogleDriveClient from app.services.google_drive import GoogleDriveToolMetadataService @@ -101,28 +101,56 @@ def create_delete_google_drive_file_tool( logger.info( f"Requesting approval for deleting Google Drive file: '{file_name}' (file_id={file_id}, delete_from_kb={delete_from_kb})" ) - result = request_approval( - action_type="google_drive_file_trash", - tool_name="delete_google_drive_file", - params={ - "file_id": file_id, - "connector_id": connector_id_from_context, - "delete_from_kb": delete_from_kb, - }, - context=context, + approval = interrupt( + { + "type": "google_drive_file_trash", + "action": { + "tool": "delete_google_drive_file", + "params": { + "file_id": file_id, + "connector_id": connector_id_from_context, + "delete_from_kb": delete_from_kb, + }, + }, + "context": context, + } ) - if result.rejected: + decisions_raw = ( + approval.get("decisions", []) if isinstance(approval, dict) else [] + ) + decisions = ( + decisions_raw if isinstance(decisions_raw, list) else [decisions_raw] + ) + decisions = [d for d in decisions if isinstance(d, dict)] + if not decisions: + logger.warning("No approval decision received") + return {"status": "error", "message": "No approval decision received"} + + decision = decisions[0] + decision_type = decision.get("type") or decision.get("decision_type") + logger.info(f"User decision: {decision_type}") + + if decision_type == "reject": return { "status": "rejected", "message": "User declined. The file was not trashed. Do not ask again or suggest alternatives.", } - final_file_id = result.params.get("file_id", file_id) - final_connector_id = result.params.get( + edited_action = decision.get("edited_action") + final_params: dict[str, Any] = {} + if isinstance(edited_action, dict): + edited_args = edited_action.get("args") + if isinstance(edited_args, dict): + final_params = edited_args + elif isinstance(decision.get("args"), dict): + final_params = decision["args"] + + final_file_id = final_params.get("file_id", file_id) + final_connector_id = final_params.get( "connector_id", connector_id_from_context ) - final_delete_from_kb = result.params.get("delete_from_kb", delete_from_kb) + final_delete_from_kb = final_params.get("delete_from_kb", delete_from_kb) if not final_connector_id: return { diff --git a/surfsense_backend/app/agents/new_chat/tools/hitl.py b/surfsense_backend/app/agents/new_chat/tools/hitl.py deleted file mode 100644 index 64ace547c..000000000 --- a/surfsense_backend/app/agents/new_chat/tools/hitl.py +++ /dev/null @@ -1,142 +0,0 @@ -"""Unified HITL (Human-in-the-Loop) approval utility. - -Provides a single ``request_approval()`` function that encapsulates the -interrupt payload creation, decision parsing, and parameter merging logic -shared by every sensitive tool (native connectors and MCP tools alike). - -Usage inside a tool:: - - from app.agents.new_chat.tools.hitl import request_approval - - result = request_approval( - action_type="gmail_email_send", - tool_name="send_gmail_email", - params={"to": to, "subject": subject, "body": body}, - context=context, - ) - if result.rejected: - return {"status": "rejected", "message": "User declined."} - # result.params contains the final (possibly edited) parameters -""" - -from __future__ import annotations - -import logging -from dataclasses import dataclass, field -from typing import Any - -from langgraph.types import interrupt - -logger = logging.getLogger(__name__) - - -@dataclass(frozen=True, slots=True) -class HITLResult: - """Outcome of a human-in-the-loop approval request.""" - - rejected: bool - decision_type: str - params: dict[str, Any] = field(default_factory=dict) - - -def _parse_decision(approval: Any) -> tuple[str, dict[str, Any]]: - """Extract the first valid decision and its edited parameters. - - Returns: - (decision_type, edited_params) where *decision_type* is one of - ``"approve"``, ``"edit"``, or ``"reject"`` and *edited_params* is - the dict of user-modified arguments (empty when there are none). - - Raises: - ValueError: when no usable decision dict can be found. - """ - decisions_raw = approval.get("decisions", []) if isinstance(approval, dict) else [] - decisions = decisions_raw if isinstance(decisions_raw, list) else [decisions_raw] - decisions = [d for d in decisions if isinstance(d, dict)] - - if not decisions: - raise ValueError("No approval decision received") - - decision = decisions[0] - decision_type: str = ( - decision.get("type") or decision.get("decision_type") or "approve" - ) - - edited_params: dict[str, Any] = {} - edited_action = decision.get("edited_action") - if isinstance(edited_action, dict): - edited_args = edited_action.get("args") - if isinstance(edited_args, dict): - edited_params = edited_args - elif isinstance(decision.get("args"), dict): - edited_params = decision["args"] - - return decision_type, edited_params - - -def request_approval( - *, - action_type: str, - tool_name: str, - params: dict[str, Any], - context: dict[str, Any] | None = None, - trusted_tools: list[str] | None = None, -) -> HITLResult: - """Pause the graph for user approval and return the decision. - - This is a **synchronous** helper (not ``async``) because - ``langgraph.types.interrupt`` is itself synchronous — it raises a - ``GraphInterrupt`` exception that the LangGraph runtime catches. - - Parameters - ---------- - action_type: - A label that the frontend uses to select the correct approval card - (e.g. ``"gmail_email_send"``, ``"mcp_tool_call"``). - tool_name: - The registered LangChain tool name (e.g. ``"send_gmail_email"``). - params: - The original tool arguments. These are shown in the approval card - and used as defaults when the user does not edit anything. - context: - Rich metadata from a ``*ToolMetadataService`` (accounts, folders, - labels, etc.). For MCP tools this can hold the server name and - tool description. - trusted_tools: - An allow-list of tool names the user has previously marked as - "Always Allow". If *tool_name* appears in this list, HITL is - skipped and the tool executes immediately. - - Returns - ------- - HITLResult - ``result.rejected`` is ``True`` when the user chose to deny the - action. Otherwise ``result.params`` contains the final parameter - dict — either the originals or the user-edited version merged on - top. - """ - if trusted_tools and tool_name in trusted_tools: - logger.info("Tool '%s' is user-trusted — skipping HITL", tool_name) - return HITLResult(rejected=False, decision_type="trusted", params=dict(params)) - - approval = interrupt( - { - "type": action_type, - "action": {"tool": tool_name, "params": params}, - "context": context or {}, - } - ) - - try: - decision_type, edited_params = _parse_decision(approval) - except ValueError: - logger.warning("No approval decision received for %s", tool_name) - return HITLResult(rejected=False, decision_type="error", params=params) - - logger.info("User decision for %s: %s", tool_name, decision_type) - - if decision_type == "reject": - return HITLResult(rejected=True, decision_type="reject", params=params) - - final_params = {**params, **edited_params} if edited_params else dict(params) - return HITLResult(rejected=False, decision_type=decision_type, params=final_params) diff --git a/surfsense_backend/app/agents/new_chat/tools/jira/create_issue.py b/surfsense_backend/app/agents/new_chat/tools/jira/create_issue.py index 8b40dde65..d441c49f3 100644 --- a/surfsense_backend/app/agents/new_chat/tools/jira/create_issue.py +++ b/surfsense_backend/app/agents/new_chat/tools/jira/create_issue.py @@ -3,10 +3,10 @@ import logging from typing import Any from langchain_core.tools import tool +from langgraph.types import interrupt from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm.attributes import flag_modified -from app.agents.new_chat.tools.hitl import request_approval from app.connectors.jira_history import JiraHistoryConnector from app.services.jira import JiraToolMetadataService @@ -69,32 +69,58 @@ def create_create_jira_issue_tool( "connector_type": "jira", } - result = request_approval( - action_type="jira_issue_creation", - tool_name="create_jira_issue", - params={ - "project_key": project_key, - "summary": summary, - "issue_type": issue_type, - "description": description, - "priority": priority, - "connector_id": connector_id, - }, - context=context, + approval = interrupt( + { + "type": "jira_issue_creation", + "action": { + "tool": "create_jira_issue", + "params": { + "project_key": project_key, + "summary": summary, + "issue_type": issue_type, + "description": description, + "priority": priority, + "connector_id": connector_id, + }, + }, + "context": context, + } ) - if result.rejected: + decisions_raw = ( + approval.get("decisions", []) if isinstance(approval, dict) else [] + ) + decisions = ( + decisions_raw if isinstance(decisions_raw, list) else [decisions_raw] + ) + decisions = [d for d in decisions if isinstance(d, dict)] + if not decisions: + return {"status": "error", "message": "No approval decision received"} + + decision = decisions[0] + decision_type = decision.get("type") or decision.get("decision_type") + + if decision_type == "reject": return { "status": "rejected", - "message": "User declined. Do not retry or suggest alternatives.", + "message": "User declined. The issue was not created.", } - final_project_key = result.params.get("project_key", project_key) - final_summary = result.params.get("summary", summary) - final_issue_type = result.params.get("issue_type", issue_type) - final_description = result.params.get("description", description) - final_priority = result.params.get("priority", priority) - final_connector_id = result.params.get("connector_id", connector_id) + final_params: dict[str, Any] = {} + edited_action = decision.get("edited_action") + if isinstance(edited_action, dict): + edited_args = edited_action.get("args") + if isinstance(edited_args, dict): + final_params = edited_args + elif isinstance(decision.get("args"), dict): + final_params = decision["args"] + + final_project_key = final_params.get("project_key", project_key) + final_summary = final_params.get("summary", summary) + final_issue_type = final_params.get("issue_type", issue_type) + final_description = final_params.get("description", description) + final_priority = final_params.get("priority", priority) + final_connector_id = final_params.get("connector_id", connector_id) if not final_summary or not final_summary.strip(): return {"status": "error", "message": "Issue summary cannot be empty."} diff --git a/surfsense_backend/app/agents/new_chat/tools/jira/delete_issue.py b/surfsense_backend/app/agents/new_chat/tools/jira/delete_issue.py index 6466c80ea..2f8c370ad 100644 --- a/surfsense_backend/app/agents/new_chat/tools/jira/delete_issue.py +++ b/surfsense_backend/app/agents/new_chat/tools/jira/delete_issue.py @@ -3,10 +3,10 @@ import logging from typing import Any from langchain_core.tools import tool +from langgraph.types import interrupt from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm.attributes import flag_modified -from app.agents.new_chat.tools.hitl import request_approval from app.connectors.jira_history import JiraHistoryConnector from app.services.jira import JiraToolMetadataService @@ -71,28 +71,54 @@ def create_delete_jira_issue_tool( document_id = issue_data["document_id"] connector_id_from_context = context.get("account", {}).get("id") - result = request_approval( - action_type="jira_issue_deletion", - tool_name="delete_jira_issue", - params={ - "issue_key": issue_key, - "connector_id": connector_id_from_context, - "delete_from_kb": delete_from_kb, - }, - context=context, + approval = interrupt( + { + "type": "jira_issue_deletion", + "action": { + "tool": "delete_jira_issue", + "params": { + "issue_key": issue_key, + "connector_id": connector_id_from_context, + "delete_from_kb": delete_from_kb, + }, + }, + "context": context, + } ) - if result.rejected: + decisions_raw = ( + approval.get("decisions", []) if isinstance(approval, dict) else [] + ) + decisions = ( + decisions_raw if isinstance(decisions_raw, list) else [decisions_raw] + ) + decisions = [d for d in decisions if isinstance(d, dict)] + if not decisions: + return {"status": "error", "message": "No approval decision received"} + + decision = decisions[0] + decision_type = decision.get("type") or decision.get("decision_type") + + if decision_type == "reject": return { "status": "rejected", - "message": "User declined. Do not retry or suggest alternatives.", + "message": "User declined. The issue was not deleted.", } - final_issue_key = result.params.get("issue_key", issue_key) - final_connector_id = result.params.get( + final_params: dict[str, Any] = {} + edited_action = decision.get("edited_action") + if isinstance(edited_action, dict): + edited_args = edited_action.get("args") + if isinstance(edited_args, dict): + final_params = edited_args + elif isinstance(decision.get("args"), dict): + final_params = decision["args"] + + final_issue_key = final_params.get("issue_key", issue_key) + final_connector_id = final_params.get( "connector_id", connector_id_from_context ) - final_delete_from_kb = result.params.get("delete_from_kb", delete_from_kb) + final_delete_from_kb = final_params.get("delete_from_kb", delete_from_kb) from sqlalchemy.future import select diff --git a/surfsense_backend/app/agents/new_chat/tools/jira/update_issue.py b/surfsense_backend/app/agents/new_chat/tools/jira/update_issue.py index f6e586a2e..c2b948ae3 100644 --- a/surfsense_backend/app/agents/new_chat/tools/jira/update_issue.py +++ b/surfsense_backend/app/agents/new_chat/tools/jira/update_issue.py @@ -3,10 +3,10 @@ import logging from typing import Any from langchain_core.tools import tool +from langgraph.types import interrupt from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm.attributes import flag_modified -from app.agents.new_chat.tools.hitl import request_approval from app.connectors.jira_history import JiraHistoryConnector from app.services.jira import JiraToolMetadataService @@ -75,34 +75,60 @@ def create_update_jira_issue_tool( document_id = issue_data.get("document_id") connector_id_from_context = context.get("account", {}).get("id") - result = request_approval( - action_type="jira_issue_update", - tool_name="update_jira_issue", - params={ - "issue_key": issue_key, - "document_id": document_id, - "new_summary": new_summary, - "new_description": new_description, - "new_priority": new_priority, - "connector_id": connector_id_from_context, - }, - context=context, + approval = interrupt( + { + "type": "jira_issue_update", + "action": { + "tool": "update_jira_issue", + "params": { + "issue_key": issue_key, + "document_id": document_id, + "new_summary": new_summary, + "new_description": new_description, + "new_priority": new_priority, + "connector_id": connector_id_from_context, + }, + }, + "context": context, + } ) - if result.rejected: + decisions_raw = ( + approval.get("decisions", []) if isinstance(approval, dict) else [] + ) + decisions = ( + decisions_raw if isinstance(decisions_raw, list) else [decisions_raw] + ) + decisions = [d for d in decisions if isinstance(d, dict)] + if not decisions: + return {"status": "error", "message": "No approval decision received"} + + decision = decisions[0] + decision_type = decision.get("type") or decision.get("decision_type") + + if decision_type == "reject": return { "status": "rejected", - "message": "User declined. Do not retry or suggest alternatives.", + "message": "User declined. The issue was not updated.", } - final_issue_key = result.params.get("issue_key", issue_key) - final_summary = result.params.get("new_summary", new_summary) - final_description = result.params.get("new_description", new_description) - final_priority = result.params.get("new_priority", new_priority) - final_connector_id = result.params.get( + final_params: dict[str, Any] = {} + edited_action = decision.get("edited_action") + if isinstance(edited_action, dict): + edited_args = edited_action.get("args") + if isinstance(edited_args, dict): + final_params = edited_args + elif isinstance(decision.get("args"), dict): + final_params = decision["args"] + + final_issue_key = final_params.get("issue_key", issue_key) + final_summary = final_params.get("new_summary", new_summary) + final_description = final_params.get("new_description", new_description) + final_priority = final_params.get("new_priority", new_priority) + final_connector_id = final_params.get( "connector_id", connector_id_from_context ) - final_document_id = result.params.get("document_id", document_id) + final_document_id = final_params.get("document_id", document_id) from sqlalchemy.future import select diff --git a/surfsense_backend/app/agents/new_chat/tools/linear/create_issue.py b/surfsense_backend/app/agents/new_chat/tools/linear/create_issue.py index ff254e133..2b5d37903 100644 --- a/surfsense_backend/app/agents/new_chat/tools/linear/create_issue.py +++ b/surfsense_backend/app/agents/new_chat/tools/linear/create_issue.py @@ -2,9 +2,9 @@ import logging from typing import Any from langchain_core.tools import tool +from langgraph.types import interrupt from sqlalchemy.ext.asyncio import AsyncSession -from app.agents.new_chat.tools.hitl import request_approval from app.connectors.linear_connector import LinearAPIError, LinearConnector from app.services.linear import LinearToolMetadataService @@ -94,37 +94,65 @@ def create_create_linear_issue_tool( } logger.info(f"Requesting approval for creating Linear issue: '{title}'") - result = request_approval( - action_type="linear_issue_creation", - tool_name="create_linear_issue", - params={ - "title": title, - "description": description, - "team_id": None, - "state_id": None, - "assignee_id": None, - "priority": None, - "label_ids": [], - "connector_id": connector_id, - }, - context=context, + approval = interrupt( + { + "type": "linear_issue_creation", + "action": { + "tool": "create_linear_issue", + "params": { + "title": title, + "description": description, + "team_id": None, + "state_id": None, + "assignee_id": None, + "priority": None, + "label_ids": [], + "connector_id": connector_id, + }, + }, + "context": context, + } ) - if result.rejected: + decisions_raw = ( + approval.get("decisions", []) if isinstance(approval, dict) else [] + ) + decisions = ( + decisions_raw if isinstance(decisions_raw, list) else [decisions_raw] + ) + decisions = [d for d in decisions if isinstance(d, dict)] + if not decisions: + logger.warning("No approval decision received") + return {"status": "error", "message": "No approval decision received"} + + decision = decisions[0] + decision_type = decision.get("type") or decision.get("decision_type") + logger.info(f"User decision: {decision_type}") + + if decision_type == "reject": logger.info("Linear issue creation rejected by user") return { "status": "rejected", - "message": "User declined. Do not retry or suggest alternatives.", + "message": "User declined. The issue was not created. Do not ask again or suggest alternatives.", } - final_title = result.params.get("title", title) - final_description = result.params.get("description", description) - final_team_id = result.params.get("team_id") - final_state_id = result.params.get("state_id") - final_assignee_id = result.params.get("assignee_id") - final_priority = result.params.get("priority") - final_label_ids = result.params.get("label_ids") or [] - final_connector_id = result.params.get("connector_id", connector_id) + final_params: dict[str, Any] = {} + edited_action = decision.get("edited_action") + if isinstance(edited_action, dict): + edited_args = edited_action.get("args") + if isinstance(edited_args, dict): + final_params = edited_args + elif isinstance(decision.get("args"), dict): + final_params = decision["args"] + + final_title = final_params.get("title", title) + final_description = final_params.get("description", description) + final_team_id = final_params.get("team_id") + final_state_id = final_params.get("state_id") + final_assignee_id = final_params.get("assignee_id") + final_priority = final_params.get("priority") + final_label_ids = final_params.get("label_ids") or [] + final_connector_id = final_params.get("connector_id", connector_id) if not final_title or not final_title.strip(): logger.error("Title is empty or contains only whitespace") diff --git a/surfsense_backend/app/agents/new_chat/tools/linear/delete_issue.py b/surfsense_backend/app/agents/new_chat/tools/linear/delete_issue.py index 29ef0cdf2..9f4a60953 100644 --- a/surfsense_backend/app/agents/new_chat/tools/linear/delete_issue.py +++ b/surfsense_backend/app/agents/new_chat/tools/linear/delete_issue.py @@ -2,9 +2,9 @@ import logging from typing import Any from langchain_core.tools import tool +from langgraph.types import interrupt from sqlalchemy.ext.asyncio import AsyncSession -from app.agents.new_chat.tools.hitl import request_approval from app.connectors.linear_connector import LinearAPIError, LinearConnector from app.services.linear import LinearToolMetadataService @@ -114,29 +114,57 @@ def create_delete_linear_issue_tool( f"Requesting approval for deleting Linear issue: '{issue_ref}' " f"(id={issue_id}, delete_from_kb={delete_from_kb})" ) - result = request_approval( - action_type="linear_issue_deletion", - tool_name="delete_linear_issue", - params={ - "issue_id": issue_id, - "connector_id": connector_id_from_context, - "delete_from_kb": delete_from_kb, - }, - context=context, + approval = interrupt( + { + "type": "linear_issue_deletion", + "action": { + "tool": "delete_linear_issue", + "params": { + "issue_id": issue_id, + "connector_id": connector_id_from_context, + "delete_from_kb": delete_from_kb, + }, + }, + "context": context, + } ) - if result.rejected: + decisions_raw = ( + approval.get("decisions", []) if isinstance(approval, dict) else [] + ) + decisions = ( + decisions_raw if isinstance(decisions_raw, list) else [decisions_raw] + ) + decisions = [d for d in decisions if isinstance(d, dict)] + if not decisions: + logger.warning("No approval decision received") + return {"status": "error", "message": "No approval decision received"} + + decision = decisions[0] + decision_type = decision.get("type") or decision.get("decision_type") + logger.info(f"User decision: {decision_type}") + + if decision_type == "reject": logger.info("Linear issue deletion rejected by user") return { "status": "rejected", - "message": "User declined. Do not retry or suggest alternatives.", + "message": "User declined. The issue was not deleted. Do not ask again or suggest alternatives.", } - final_issue_id = result.params.get("issue_id", issue_id) - final_connector_id = result.params.get( + edited_action = decision.get("edited_action") + final_params: dict[str, Any] = {} + if isinstance(edited_action, dict): + edited_args = edited_action.get("args") + if isinstance(edited_args, dict): + final_params = edited_args + elif isinstance(decision.get("args"), dict): + final_params = decision["args"] + + final_issue_id = final_params.get("issue_id", issue_id) + final_connector_id = final_params.get( "connector_id", connector_id_from_context ) - final_delete_from_kb = result.params.get("delete_from_kb", delete_from_kb) + final_delete_from_kb = final_params.get("delete_from_kb", delete_from_kb) logger.info( f"Deleting Linear issue with final params: issue_id={final_issue_id}, " diff --git a/surfsense_backend/app/agents/new_chat/tools/linear/update_issue.py b/surfsense_backend/app/agents/new_chat/tools/linear/update_issue.py index f35d0dddd..19af851c1 100644 --- a/surfsense_backend/app/agents/new_chat/tools/linear/update_issue.py +++ b/surfsense_backend/app/agents/new_chat/tools/linear/update_issue.py @@ -2,9 +2,9 @@ import logging from typing import Any from langchain_core.tools import tool +from langgraph.types import interrupt from sqlalchemy.ext.asyncio import AsyncSession -from app.agents.new_chat.tools.hitl import request_approval from app.connectors.linear_connector import LinearAPIError, LinearConnector from app.services.linear import LinearKBSyncService, LinearToolMetadataService @@ -130,45 +130,69 @@ def create_update_linear_issue_tool( logger.info( f"Requesting approval for updating Linear issue: '{issue_ref}' (id={issue_id})" ) - result = request_approval( - action_type="linear_issue_update", - tool_name="update_linear_issue", - params={ - "issue_id": issue_id, - "document_id": document_id, - "new_title": new_title, - "new_description": new_description, - "new_state_id": new_state_id, - "new_assignee_id": new_assignee_id, - "new_priority": new_priority, - "new_label_ids": new_label_ids, - "connector_id": connector_id_from_context, - }, - context=context, + approval = interrupt( + { + "type": "linear_issue_update", + "action": { + "tool": "update_linear_issue", + "params": { + "issue_id": issue_id, + "document_id": document_id, + "new_title": new_title, + "new_description": new_description, + "new_state_id": new_state_id, + "new_assignee_id": new_assignee_id, + "new_priority": new_priority, + "new_label_ids": new_label_ids, + "connector_id": connector_id_from_context, + }, + }, + "context": context, + } ) - if result.rejected: + decisions_raw = ( + approval.get("decisions", []) if isinstance(approval, dict) else [] + ) + decisions = ( + decisions_raw if isinstance(decisions_raw, list) else [decisions_raw] + ) + decisions = [d for d in decisions if isinstance(d, dict)] + if not decisions: + logger.warning("No approval decision received") + return {"status": "error", "message": "No approval decision received"} + + decision = decisions[0] + decision_type = decision.get("type") or decision.get("decision_type") + logger.info(f"User decision: {decision_type}") + + if decision_type == "reject": logger.info("Linear issue update rejected by user") return { "status": "rejected", - "message": "User declined. Do not retry or suggest alternatives.", + "message": "User declined. The issue was not updated. Do not ask again or suggest alternatives.", } - final_issue_id = result.params.get("issue_id", issue_id) - final_document_id = result.params.get("document_id", document_id) - final_new_title = result.params.get("new_title", new_title) - final_new_description = result.params.get( - "new_description", new_description - ) - final_new_state_id = result.params.get("new_state_id", new_state_id) - final_new_assignee_id = result.params.get( - "new_assignee_id", new_assignee_id - ) - final_new_priority = result.params.get("new_priority", new_priority) - final_new_label_ids: list[str] | None = result.params.get( + edited_action = decision.get("edited_action") + final_params: dict[str, Any] = {} + if isinstance(edited_action, dict): + edited_args = edited_action.get("args") + if isinstance(edited_args, dict): + final_params = edited_args + elif isinstance(decision.get("args"), dict): + final_params = decision["args"] + + final_issue_id = final_params.get("issue_id", issue_id) + final_document_id = final_params.get("document_id", document_id) + final_new_title = final_params.get("new_title", new_title) + final_new_description = final_params.get("new_description", new_description) + final_new_state_id = final_params.get("new_state_id", new_state_id) + final_new_assignee_id = final_params.get("new_assignee_id", new_assignee_id) + final_new_priority = final_params.get("new_priority", new_priority) + final_new_label_ids: list[str] | None = final_params.get( "new_label_ids", new_label_ids ) - final_connector_id = result.params.get( + final_connector_id = final_params.get( "connector_id", connector_id_from_context ) diff --git a/surfsense_backend/app/agents/new_chat/tools/mcp_tool.py b/surfsense_backend/app/agents/new_chat/tools/mcp_tool.py index 9743d049d..2fb7ffb06 100644 --- a/surfsense_backend/app/agents/new_chat/tools/mcp_tool.py +++ b/surfsense_backend/app/agents/new_chat/tools/mcp_tool.py @@ -7,11 +7,7 @@ Supports both transport types: - stdio: Local process-based MCP servers (command, args, env) - streamable-http/http/sse: Remote HTTP-based MCP servers (url, headers) -All MCP tools are unconditionally gated by HITL (Human-in-the-Loop) approval. -Per the MCP spec: "Clients MUST consider tool annotations to be untrusted unless -they come from trusted servers." Users can bypass HITL for specific tools by -clicking "Always Allow", which adds the tool name to the connector's -``config.trusted_tools`` allow-list. +This implements real MCP protocol support similar to Cursor's implementation. """ import logging @@ -25,7 +21,6 @@ from pydantic import BaseModel, create_model from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession -from app.agents.new_chat.tools.hitl import request_approval from app.agents.new_chat.tools.mcp_client import MCPClient from app.db import SearchSourceConnector, SearchSourceConnectorType @@ -54,15 +49,27 @@ def _create_dynamic_input_model_from_schema( tool_name: str, input_schema: dict[str, Any], ) -> type[BaseModel]: - """Create a Pydantic model from MCP tool's JSON schema.""" + """Create a Pydantic model from MCP tool's JSON schema. + + Args: + tool_name: Name of the tool (used for model class name) + input_schema: JSON schema from MCP server + + Returns: + Pydantic model class for tool input validation + + """ properties = input_schema.get("properties", {}) required_fields = input_schema.get("required", []) + # Build Pydantic field definitions field_definitions = {} for param_name, param_schema in properties.items(): param_description = param_schema.get("description", "") is_required = param_name in required_fields + # Use Any type for complex schemas to preserve structure + # This allows the MCP server to do its own validation from typing import Any as AnyType from pydantic import Field @@ -78,6 +85,7 @@ def _create_dynamic_input_model_from_schema( Field(None, description=param_description), ) + # Create dynamic model model_name = f"{tool_name.replace(' ', '').replace('-', '_')}Input" return create_model(model_name, **field_definitions) @@ -85,70 +93,55 @@ def _create_dynamic_input_model_from_schema( async def _create_mcp_tool_from_definition_stdio( tool_def: dict[str, Any], mcp_client: MCPClient, - *, - connector_name: str = "", - connector_id: int | None = None, - trusted_tools: list[str] | None = None, ) -> StructuredTool: """Create a LangChain tool from an MCP tool definition (stdio transport). - All MCP tools are unconditionally wrapped with HITL approval. - ``request_approval()`` is called OUTSIDE the try/except so that - ``GraphInterrupt`` propagates cleanly to LangGraph. + Args: + tool_def: Tool definition from MCP server with name, description, input_schema + mcp_client: MCP client instance for calling the tool + + Returns: + LangChain StructuredTool instance + """ tool_name = tool_def.get("name", "unnamed_tool") tool_description = tool_def.get("description", "No description provided") input_schema = tool_def.get("input_schema", {"type": "object", "properties": {}}) + # Log the actual schema for debugging logger.info(f"MCP tool '{tool_name}' input schema: {input_schema}") + # Create dynamic input model from schema input_model = _create_dynamic_input_model_from_schema(tool_name, input_schema) async def mcp_tool_call(**kwargs) -> str: """Execute the MCP tool call via the client with retry support.""" logger.info(f"MCP tool '{tool_name}' called with params: {kwargs}") - # HITL — OUTSIDE try/except so GraphInterrupt propagates to LangGraph - hitl_result = request_approval( - action_type="mcp_tool_call", - tool_name=tool_name, - params=kwargs, - context={ - "mcp_server": connector_name, - "tool_description": tool_description, - "mcp_transport": "stdio", - "mcp_connector_id": connector_id, - }, - trusted_tools=trusted_tools, - ) - if hitl_result.rejected: - return "Tool call rejected by user." - call_kwargs = hitl_result.params - try: + # Connect to server and call tool (connect has built-in retry logic) async with mcp_client.connect(): - result = await mcp_client.call_tool(tool_name, call_kwargs) + result = await mcp_client.call_tool(tool_name, kwargs) return str(result) except RuntimeError as e: + # Connection failures after all retries error_msg = f"MCP tool '{tool_name}' connection failed after retries: {e!s}" logger.error(error_msg) return f"Error: {error_msg}" except Exception as e: + # Tool execution or other errors error_msg = f"MCP tool '{tool_name}' execution failed: {e!s}" logger.exception(error_msg) return f"Error: {error_msg}" + # Create StructuredTool with response_format to preserve exact schema tool = StructuredTool( name=tool_name, description=tool_description, coroutine=mcp_tool_call, args_schema=input_model, - metadata={ - "mcp_input_schema": input_schema, - "mcp_transport": "stdio", - "hitl": True, - "hitl_dedup_key": next(iter(input_schema.get("required", [])), None), - }, + # Store the original MCP schema as metadata so we can access it later + metadata={"mcp_input_schema": input_schema, "mcp_transport": "stdio"}, ) logger.info(f"Created MCP tool (stdio): '{tool_name}'") @@ -159,54 +152,43 @@ async def _create_mcp_tool_from_definition_http( tool_def: dict[str, Any], url: str, headers: dict[str, str], - *, - connector_name: str = "", - connector_id: int | None = None, - trusted_tools: list[str] | None = None, ) -> StructuredTool: """Create a LangChain tool from an MCP tool definition (HTTP transport). - All MCP tools are unconditionally wrapped with HITL approval. - ``request_approval()`` is called OUTSIDE the try/except so that - ``GraphInterrupt`` propagates cleanly to LangGraph. + Args: + tool_def: Tool definition from MCP server with name, description, input_schema + url: URL of the MCP server + headers: HTTP headers for authentication + + Returns: + LangChain StructuredTool instance + """ tool_name = tool_def.get("name", "unnamed_tool") tool_description = tool_def.get("description", "No description provided") input_schema = tool_def.get("input_schema", {"type": "object", "properties": {}}) + # Log the actual schema for debugging logger.info(f"MCP HTTP tool '{tool_name}' input schema: {input_schema}") + # Create dynamic input model from schema input_model = _create_dynamic_input_model_from_schema(tool_name, input_schema) async def mcp_http_tool_call(**kwargs) -> str: """Execute the MCP tool call via HTTP transport.""" logger.info(f"MCP HTTP tool '{tool_name}' called with params: {kwargs}") - # HITL — OUTSIDE try/except so GraphInterrupt propagates to LangGraph - hitl_result = request_approval( - action_type="mcp_tool_call", - tool_name=tool_name, - params=kwargs, - context={ - "mcp_server": connector_name, - "tool_description": tool_description, - "mcp_transport": "http", - "mcp_connector_id": connector_id, - }, - trusted_tools=trusted_tools, - ) - if hitl_result.rejected: - return "Tool call rejected by user." - call_kwargs = hitl_result.params - try: async with ( streamablehttp_client(url, headers=headers) as (read, write, _), ClientSession(read, write) as session, ): await session.initialize() - response = await session.call_tool(tool_name, arguments=call_kwargs) + # Call the tool + response = await session.call_tool(tool_name, arguments=kwargs) + + # Extract content from response result = [] for content in response.content: if hasattr(content, "text"): @@ -227,6 +209,7 @@ async def _create_mcp_tool_from_definition_http( logger.exception(error_msg) return f"Error: {error_msg}" + # Create StructuredTool tool = StructuredTool( name=tool_name, description=tool_description, @@ -236,8 +219,6 @@ async def _create_mcp_tool_from_definition_http( "mcp_input_schema": input_schema, "mcp_transport": "http", "mcp_url": url, - "hitl": True, - "hitl_dedup_key": next(iter(input_schema.get("required", [])), None), }, ) @@ -249,11 +230,20 @@ async def _load_stdio_mcp_tools( connector_id: int, connector_name: str, server_config: dict[str, Any], - trusted_tools: list[str] | None = None, ) -> list[StructuredTool]: - """Load tools from a stdio-based MCP server.""" + """Load tools from a stdio-based MCP server. + + Args: + connector_id: Connector ID for logging + connector_name: Connector name for logging + server_config: Server configuration with command, args, env + + Returns: + List of tools from the MCP server + """ tools: list[StructuredTool] = [] + # Validate required command field command = server_config.get("command") if not command or not isinstance(command, str): logger.warning( @@ -261,6 +251,7 @@ async def _load_stdio_mcp_tools( ) return tools + # Validate args field (must be list if present) args = server_config.get("args", []) if not isinstance(args, list): logger.warning( @@ -268,6 +259,7 @@ async def _load_stdio_mcp_tools( ) return tools + # Validate env field (must be dict if present) env = server_config.get("env", {}) if not isinstance(env, dict): logger.warning( @@ -275,8 +267,10 @@ async def _load_stdio_mcp_tools( ) return tools + # Create MCP client mcp_client = MCPClient(command, args, env) + # Connect and discover tools async with mcp_client.connect(): tool_definitions = await mcp_client.list_tools() @@ -285,15 +279,10 @@ async def _load_stdio_mcp_tools( f"'{command}' (connector {connector_id})" ) + # Create LangChain tools from definitions for tool_def in tool_definitions: try: - tool = await _create_mcp_tool_from_definition_stdio( - tool_def, - mcp_client, - connector_name=connector_name, - connector_id=connector_id, - trusted_tools=trusted_tools, - ) + tool = await _create_mcp_tool_from_definition_stdio(tool_def, mcp_client) tools.append(tool) except Exception as e: logger.exception( @@ -308,11 +297,20 @@ async def _load_http_mcp_tools( connector_id: int, connector_name: str, server_config: dict[str, Any], - trusted_tools: list[str] | None = None, ) -> list[StructuredTool]: - """Load tools from an HTTP-based MCP server.""" + """Load tools from an HTTP-based MCP server. + + Args: + connector_id: Connector ID for logging + connector_name: Connector name for logging + server_config: Server configuration with url, headers + + Returns: + List of tools from the MCP server + """ tools: list[StructuredTool] = [] + # Validate required url field url = server_config.get("url") if not url or not isinstance(url, str): logger.warning( @@ -320,6 +318,7 @@ async def _load_http_mcp_tools( ) return tools + # Validate headers field (must be dict if present) headers = server_config.get("headers", {}) if not isinstance(headers, dict): logger.warning( @@ -327,6 +326,7 @@ async def _load_http_mcp_tools( ) return tools + # Connect and discover tools via HTTP try: async with ( streamablehttp_client(url, headers=headers) as (read, write, _), @@ -334,6 +334,7 @@ async def _load_http_mcp_tools( ): await session.initialize() + # List available tools response = await session.list_tools() tool_definitions = [] for tool in response.tools: @@ -352,15 +353,11 @@ async def _load_http_mcp_tools( f"'{url}' (connector {connector_id})" ) + # Create LangChain tools from definitions for tool_def in tool_definitions: try: tool = await _create_mcp_tool_from_definition_http( - tool_def, - url, - headers, - connector_name=connector_name, - connector_id=connector_id, - trusted_tools=trusted_tools, + tool_def, url, headers ) tools.append(tool) except Exception as e: @@ -401,6 +398,14 @@ async def load_mcp_tools( Results are cached per search space for up to 5 minutes to avoid re-spawning MCP server processes on every chat message. + + Args: + session: Database session + search_space_id: User's search space ID + + Returns: + List of LangChain StructuredTool instances + """ _evict_expired_mcp_cache() @@ -431,7 +436,6 @@ async def load_mcp_tools( try: config = connector.config or {} server_config = config.get("server_config", {}) - trusted_tools = config.get("trusted_tools", []) if not server_config or not isinstance(server_config, dict): logger.warning( @@ -443,17 +447,11 @@ async def load_mcp_tools( if transport in ("streamable-http", "http", "sse"): connector_tools = await _load_http_mcp_tools( - connector.id, - connector.name, - server_config, - trusted_tools=trusted_tools, + connector.id, connector.name, server_config ) else: connector_tools = await _load_stdio_mcp_tools( - connector.id, - connector.name, - server_config, - trusted_tools=trusted_tools, + connector.id, connector.name, server_config ) tools.extend(connector_tools) diff --git a/surfsense_backend/app/agents/new_chat/tools/notion/create_page.py b/surfsense_backend/app/agents/new_chat/tools/notion/create_page.py index 6efffe960..5bb0c52d1 100644 --- a/surfsense_backend/app/agents/new_chat/tools/notion/create_page.py +++ b/surfsense_backend/app/agents/new_chat/tools/notion/create_page.py @@ -2,9 +2,9 @@ import logging from typing import Any from langchain_core.tools import tool +from langgraph.types import interrupt from sqlalchemy.ext.asyncio import AsyncSession -from app.agents.new_chat.tools.hitl import request_approval from app.connectors.notion_history import NotionAPIError, NotionHistoryConnector from app.services.notion import NotionToolMetadataService @@ -99,29 +99,61 @@ def create_create_notion_page_tool( } logger.info(f"Requesting approval for creating Notion page: '{title}'") - result = request_approval( - action_type="notion_page_creation", - tool_name="create_notion_page", - params={ - "title": title, - "content": content, - "parent_page_id": None, - "connector_id": connector_id, - }, - context=context, + approval = interrupt( + { + "type": "notion_page_creation", + "action": { + "tool": "create_notion_page", + "params": { + "title": title, + "content": content, + "parent_page_id": None, + "connector_id": connector_id, + }, + }, + "context": context, + } ) - if result.rejected: + decisions_raw = ( + approval.get("decisions", []) if isinstance(approval, dict) else [] + ) + decisions = ( + decisions_raw if isinstance(decisions_raw, list) else [decisions_raw] + ) + decisions = [d for d in decisions if isinstance(d, dict)] + if not decisions: + logger.warning("No approval decision received") + return { + "status": "error", + "message": "No approval decision received", + } + + decision = decisions[0] + decision_type = decision.get("type") or decision.get("decision_type") + logger.info(f"User decision: {decision_type}") + + if decision_type == "reject": logger.info("Notion page creation rejected by user") return { "status": "rejected", - "message": "User declined. Do not retry or suggest alternatives.", + "message": "User declined. The page was not created. Do not ask again or suggest alternatives.", } - final_title = result.params.get("title", title) - final_content = result.params.get("content", content) - final_parent_page_id = result.params.get("parent_page_id") - final_connector_id = result.params.get("connector_id", connector_id) + edited_action = decision.get("edited_action") + final_params: dict[str, Any] = {} + if isinstance(edited_action, dict): + edited_args = edited_action.get("args") + if isinstance(edited_args, dict): + final_params = edited_args + elif isinstance(decision.get("args"), dict): + # Some interrupt payloads place args directly on the decision. + final_params = decision["args"] + + final_title = final_params.get("title", title) + final_content = final_params.get("content", content) + final_parent_page_id = final_params.get("parent_page_id") + final_connector_id = final_params.get("connector_id", connector_id) if not final_title or not final_title.strip(): logger.error("Title is empty or contains only whitespace") diff --git a/surfsense_backend/app/agents/new_chat/tools/notion/delete_page.py b/surfsense_backend/app/agents/new_chat/tools/notion/delete_page.py index 07f7583d2..fbb7c5004 100644 --- a/surfsense_backend/app/agents/new_chat/tools/notion/delete_page.py +++ b/surfsense_backend/app/agents/new_chat/tools/notion/delete_page.py @@ -2,9 +2,9 @@ import logging from typing import Any from langchain_core.tools import tool +from langgraph.types import interrupt from sqlalchemy.ext.asyncio import AsyncSession -from app.agents.new_chat.tools.hitl import request_approval from app.connectors.notion_history import NotionAPIError, NotionHistoryConnector from app.services.notion.tool_metadata_service import NotionToolMetadataService @@ -114,29 +114,63 @@ def create_delete_notion_page_tool( f"Requesting approval for deleting Notion page: '{page_title}' (page_id={page_id}, delete_from_kb={delete_from_kb})" ) - result = request_approval( - action_type="notion_page_deletion", - tool_name="delete_notion_page", - params={ - "page_id": page_id, - "connector_id": connector_id_from_context, - "delete_from_kb": delete_from_kb, - }, - context=context, + # Request approval before deleting + approval = interrupt( + { + "type": "notion_page_deletion", + "action": { + "tool": "delete_notion_page", + "params": { + "page_id": page_id, + "connector_id": connector_id_from_context, + "delete_from_kb": delete_from_kb, + }, + }, + "context": context, + } ) - if result.rejected: + decisions_raw = ( + approval.get("decisions", []) if isinstance(approval, dict) else [] + ) + decisions = ( + decisions_raw if isinstance(decisions_raw, list) else [decisions_raw] + ) + decisions = [d for d in decisions if isinstance(d, dict)] + if not decisions: + logger.warning("No approval decision received") + return { + "status": "error", + "message": "No approval decision received", + } + + decision = decisions[0] + decision_type = decision.get("type") or decision.get("decision_type") + logger.info(f"User decision: {decision_type}") + + if decision_type == "reject": logger.info("Notion page deletion rejected by user") return { "status": "rejected", - "message": "User declined. Do not retry or suggest alternatives.", + "message": "User declined. The page was not deleted. Do not ask again or suggest alternatives.", } - final_page_id = result.params.get("page_id", page_id) - final_connector_id = result.params.get( + # Extract edited action arguments (if user modified the checkbox) + edited_action = decision.get("edited_action") + final_params: dict[str, Any] = {} + if isinstance(edited_action, dict): + edited_args = edited_action.get("args") + if isinstance(edited_args, dict): + final_params = edited_args + elif isinstance(decision.get("args"), dict): + # Some interrupt payloads place args directly on the decision. + final_params = decision["args"] + + final_page_id = final_params.get("page_id", page_id) + final_connector_id = final_params.get( "connector_id", connector_id_from_context ) - final_delete_from_kb = result.params.get("delete_from_kb", delete_from_kb) + final_delete_from_kb = final_params.get("delete_from_kb", delete_from_kb) logger.info( f"Deleting Notion page with final params: page_id={final_page_id}, connector_id={final_connector_id}, delete_from_kb={final_delete_from_kb}" diff --git a/surfsense_backend/app/agents/new_chat/tools/notion/update_page.py b/surfsense_backend/app/agents/new_chat/tools/notion/update_page.py index 85c08177c..25f2b9918 100644 --- a/surfsense_backend/app/agents/new_chat/tools/notion/update_page.py +++ b/surfsense_backend/app/agents/new_chat/tools/notion/update_page.py @@ -2,9 +2,9 @@ import logging from typing import Any from langchain_core.tools import tool +from langgraph.types import interrupt from sqlalchemy.ext.asyncio import AsyncSession -from app.agents.new_chat.tools.hitl import request_approval from app.connectors.notion_history import NotionAPIError, NotionHistoryConnector from app.services.notion import NotionToolMetadataService @@ -127,27 +127,59 @@ def create_update_notion_page_tool( logger.info( f"Requesting approval for updating Notion page: '{page_title}' (page_id={page_id})" ) - result = request_approval( - action_type="notion_page_update", - tool_name="update_notion_page", - params={ - "page_id": page_id, - "content": content, - "connector_id": connector_id_from_context, - }, - context=context, + approval = interrupt( + { + "type": "notion_page_update", + "action": { + "tool": "update_notion_page", + "params": { + "page_id": page_id, + "content": content, + "connector_id": connector_id_from_context, + }, + }, + "context": context, + } ) - if result.rejected: + decisions_raw = ( + approval.get("decisions", []) if isinstance(approval, dict) else [] + ) + decisions = ( + decisions_raw if isinstance(decisions_raw, list) else [decisions_raw] + ) + decisions = [d for d in decisions if isinstance(d, dict)] + if not decisions: + logger.warning("No approval decision received") + return { + "status": "error", + "message": "No approval decision received", + } + + decision = decisions[0] + decision_type = decision.get("type") or decision.get("decision_type") + logger.info(f"User decision: {decision_type}") + + if decision_type == "reject": logger.info("Notion page update rejected by user") return { "status": "rejected", - "message": "User declined. Do not retry or suggest alternatives.", + "message": "User declined. The page was not updated. Do not ask again or suggest alternatives.", } - final_page_id = result.params.get("page_id", page_id) - final_content = result.params.get("content", content) - final_connector_id = result.params.get( + edited_action = decision.get("edited_action") + final_params: dict[str, Any] = {} + if isinstance(edited_action, dict): + edited_args = edited_action.get("args") + if isinstance(edited_args, dict): + final_params = edited_args + elif isinstance(decision.get("args"), dict): + # Some interrupt payloads place args directly on the decision. + final_params = decision["args"] + + final_page_id = final_params.get("page_id", page_id) + final_content = final_params.get("content", content) + final_connector_id = final_params.get( "connector_id", connector_id_from_context ) diff --git a/surfsense_backend/app/agents/new_chat/tools/onedrive/create_file.py b/surfsense_backend/app/agents/new_chat/tools/onedrive/create_file.py index 21272e01d..8dffb18dd 100644 --- a/surfsense_backend/app/agents/new_chat/tools/onedrive/create_file.py +++ b/surfsense_backend/app/agents/new_chat/tools/onedrive/create_file.py @@ -5,10 +5,10 @@ from pathlib import Path from typing import Any from langchain_core.tools import tool +from langgraph.types import interrupt from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.future import select -from app.agents.new_chat.tools.hitl import request_approval from app.connectors.onedrive.client import OneDriveClient from app.db import SearchSourceConnector, SearchSourceConnectorType @@ -145,28 +145,54 @@ def create_create_onedrive_file_tool( "parent_folders": parent_folders, } - result = request_approval( - action_type="onedrive_file_creation", - tool_name="create_onedrive_file", - params={ - "name": name, - "content": content, - "connector_id": None, - "parent_folder_id": None, - }, - context=context, + approval = interrupt( + { + "type": "onedrive_file_creation", + "action": { + "tool": "create_onedrive_file", + "params": { + "name": name, + "content": content, + "connector_id": None, + "parent_folder_id": None, + }, + }, + "context": context, + } ) - if result.rejected: + decisions_raw = ( + approval.get("decisions", []) if isinstance(approval, dict) else [] + ) + decisions = ( + decisions_raw if isinstance(decisions_raw, list) else [decisions_raw] + ) + decisions = [d for d in decisions if isinstance(d, dict)] + if not decisions: + return {"status": "error", "message": "No approval decision received"} + + decision = decisions[0] + decision_type = decision.get("type") or decision.get("decision_type") + + if decision_type == "reject": return { "status": "rejected", - "message": "User declined. Do not retry or suggest alternatives.", + "message": "User declined. The file was not created.", } - final_name = result.params.get("name", name) - final_content = result.params.get("content", content) - final_connector_id = result.params.get("connector_id") - final_parent_folder_id = result.params.get("parent_folder_id") + final_params: dict[str, Any] = {} + edited_action = decision.get("edited_action") + if isinstance(edited_action, dict): + edited_args = edited_action.get("args") + if isinstance(edited_args, dict): + final_params = edited_args + elif isinstance(decision.get("args"), dict): + final_params = decision["args"] + + final_name = final_params.get("name", name) + final_content = final_params.get("content", content) + final_connector_id = final_params.get("connector_id") + final_parent_folder_id = final_params.get("parent_folder_id") if not final_name or not final_name.strip(): return {"status": "error", "message": "File name cannot be empty."} diff --git a/surfsense_backend/app/agents/new_chat/tools/onedrive/trash_file.py b/surfsense_backend/app/agents/new_chat/tools/onedrive/trash_file.py index a7f13b5df..79d8222fd 100644 --- a/surfsense_backend/app/agents/new_chat/tools/onedrive/trash_file.py +++ b/surfsense_backend/app/agents/new_chat/tools/onedrive/trash_file.py @@ -2,11 +2,11 @@ import logging from typing import Any from langchain_core.tools import tool +from langgraph.types import interrupt from sqlalchemy import String, and_, cast, func from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.future import select -from app.agents.new_chat.tools.hitl import request_approval from app.connectors.onedrive.client import OneDriveClient from app.db import ( Document, @@ -174,26 +174,53 @@ def create_delete_onedrive_file_tool( }, } - result = request_approval( - action_type="onedrive_file_trash", - tool_name="delete_onedrive_file", - params={ - "file_id": file_id, - "connector_id": connector.id, - "delete_from_kb": delete_from_kb, - }, - context=context, + approval = interrupt( + { + "type": "onedrive_file_trash", + "action": { + "tool": "delete_onedrive_file", + "params": { + "file_id": file_id, + "connector_id": connector.id, + "delete_from_kb": delete_from_kb, + }, + }, + "context": context, + } ) - if result.rejected: + decisions_raw = ( + approval.get("decisions", []) if isinstance(approval, dict) else [] + ) + decisions = ( + decisions_raw if isinstance(decisions_raw, list) else [decisions_raw] + ) + decisions = [d for d in decisions if isinstance(d, dict)] + if not decisions: + return {"status": "error", "message": "No approval decision received"} + + decision = decisions[0] + decision_type = decision.get("type") or decision.get("decision_type") + logger.info(f"User decision: {decision_type}") + + if decision_type == "reject": return { "status": "rejected", - "message": "User declined. Do not retry or suggest alternatives.", + "message": "User declined. The file was not trashed. Do not ask again or suggest alternatives.", } - final_file_id = result.params.get("file_id", file_id) - final_connector_id = result.params.get("connector_id", connector.id) - final_delete_from_kb = result.params.get("delete_from_kb", delete_from_kb) + final_params: dict[str, Any] = {} + edited_action = decision.get("edited_action") + if isinstance(edited_action, dict): + edited_args = edited_action.get("args") + if isinstance(edited_args, dict): + final_params = edited_args + elif isinstance(decision.get("args"), dict): + final_params = decision["args"] + + final_file_id = final_params.get("file_id", file_id) + final_connector_id = final_params.get("connector_id", connector.id) + final_delete_from_kb = final_params.get("delete_from_kb", delete_from_kb) if final_connector_id != connector.id: result = await db_session.execute( diff --git a/surfsense_backend/app/db.py b/surfsense_backend/app/db.py index 82d77f847..e69d28ac2 100644 --- a/surfsense_backend/app/db.py +++ b/surfsense_backend/app/db.py @@ -1321,10 +1321,6 @@ class SearchSpace(BaseModel, TimestampMixin): Integer, nullable=True, default=0 ) # For vision/screenshot analysis, defaults to Auto mode - ai_file_sort_enabled = Column( - Boolean, nullable=False, default=False, server_default="false" - ) - user_id = Column( UUID(as_uuid=True), ForeignKey("user.id", ondelete="CASCADE"), nullable=False ) diff --git a/surfsense_backend/app/indexing_pipeline/indexing_pipeline_service.py b/surfsense_backend/app/indexing_pipeline/indexing_pipeline_service.py index e6b2458f3..22c552e5c 100644 --- a/surfsense_backend/app/indexing_pipeline/indexing_pipeline_service.py +++ b/surfsense_backend/app/indexing_pipeline/indexing_pipeline_service.py @@ -422,8 +422,6 @@ class IndexingPipelineService: ) log_index_success(ctx, chunk_count=len(chunks)) - await self._enqueue_ai_sort_if_enabled(document) - except RETRYABLE_LLM_ERRORS as e: log_retryable_llm_error(ctx, e) await rollback_and_persist_failure( @@ -459,29 +457,6 @@ class IndexingPipelineService: return document - async def _enqueue_ai_sort_if_enabled(self, document: Document) -> None: - """Fire-and-forget: enqueue incremental AI sort if the search space has it enabled.""" - try: - from app.db import SearchSpace - - result = await self.session.execute( - select(SearchSpace.ai_file_sort_enabled).where( - SearchSpace.id == document.search_space_id - ) - ) - enabled = result.scalar() - if not enabled: - return - - from app.tasks.celery_tasks.document_tasks import ai_sort_document_task - - user_id = str(document.created_by_id) if document.created_by_id else "" - ai_sort_document_task.delay(document.search_space_id, user_id, document.id) - except Exception: - logging.getLogger(__name__).warning( - "Failed to enqueue AI sort for document %s", document.id, exc_info=True - ) - async def index_batch_parallel( self, connector_docs: list[ConnectorDocument], diff --git a/surfsense_backend/app/routes/editor_routes.py b/surfsense_backend/app/routes/editor_routes.py index 0f986e416..829b2cf69 100644 --- a/surfsense_backend/app/routes/editor_routes.py +++ b/surfsense_backend/app/routes/editor_routes.py @@ -139,19 +139,9 @@ async def get_editor_content( status_code=409, detail="This document is still being processed. Please wait a moment and try again.", ) - if state == "failed": - reason = ( - doc_status.get("reason", "Unknown error") - if isinstance(doc_status, dict) - else "Unknown error" - ) - raise HTTPException( - status_code=422, - detail=f"Processing failed: {reason}. You can delete this document and re-upload it.", - ) raise HTTPException( status_code=400, - detail="This document has no content. It may not have been processed correctly. Try deleting and re-uploading it.", + detail="This document has no viewable content yet. It may still be syncing. Try again in a few seconds, or re-upload if the issue persists.", ) markdown_content = "\n\n".join(chunk_contents) diff --git a/surfsense_backend/app/routes/export_routes.py b/surfsense_backend/app/routes/export_routes.py index 4f2b545a3..641c7fedb 100644 --- a/surfsense_backend/app/routes/export_routes.py +++ b/surfsense_backend/app/routes/export_routes.py @@ -20,9 +20,7 @@ router = APIRouter() @router.get("/search-spaces/{search_space_id}/export") async def export_knowledge_base( search_space_id: int, - folder_id: int | None = Query( - None, description="Export only this folder's subtree" - ), + folder_id: int | None = Query(None, description="Export only this folder's subtree"), session: AsyncSession = Depends(get_async_session), user: User = Depends(current_active_user), ): diff --git a/surfsense_backend/app/routes/sandbox_routes.py b/surfsense_backend/app/routes/sandbox_routes.py index f656e8d76..2c12c3a1e 100644 --- a/surfsense_backend/app/routes/sandbox_routes.py +++ b/surfsense_backend/app/routes/sandbox_routes.py @@ -86,8 +86,9 @@ async def download_sandbox_file( # Fall back to live sandbox download try: - sandbox, _ = await get_or_create_sandbox(thread_id) - content: bytes = await asyncio.to_thread(sandbox.download_file, path) + sandbox = await get_or_create_sandbox(thread_id) + raw_sandbox = sandbox._sandbox + content: bytes = await asyncio.to_thread(raw_sandbox.fs.download_file, path) except Exception as exc: logger.warning("Sandbox file download failed for %s: %s", path, exc) raise HTTPException( diff --git a/surfsense_backend/app/routes/search_source_connectors_routes.py b/surfsense_backend/app/routes/search_source_connectors_routes.py index b87ce28c9..bb20da65d 100644 --- a/surfsense_backend/app/routes/search_source_connectors_routes.py +++ b/surfsense_backend/app/routes/search_source_connectors_routes.py @@ -636,16 +636,9 @@ async def delete_search_source_connector( ) # Delete the connector record - search_space_id = db_connector.search_space_id - is_mcp = db_connector.connector_type == SearchSourceConnectorType.MCP_CONNECTOR await session.delete(db_connector) await session.commit() - if is_mcp: - from app.agents.new_chat.tools.mcp_tool import invalidate_mcp_tools_cache - - invalidate_mcp_tools_cache(search_space_id) - logger.info( f"Connector {connector_id} ({connector_name}) deleted successfully. " f"Total documents deleted: {total_deleted}" @@ -3631,114 +3624,3 @@ async def get_drive_picker_token( status_code=500, detail="Failed to retrieve access token. Check server logs for details.", ) from e - - -# ============================================================================= -# MCP Tool Trust (Allow-List) Routes -# ============================================================================= - - -class MCPTrustToolRequest(BaseModel): - tool_name: str - - -@router.post("/connectors/mcp/{connector_id}/trust-tool") -async def trust_mcp_tool( - connector_id: int, - body: MCPTrustToolRequest, - session: AsyncSession = Depends(get_async_session), - user: User = Depends(current_active_user), -): - """Add a tool to the MCP connector's trusted (always-allow) list. - - Once trusted, the tool executes without HITL approval on subsequent calls. - """ - try: - result = await session.execute( - select(SearchSourceConnector).filter( - SearchSourceConnector.id == connector_id, - SearchSourceConnector.connector_type - == SearchSourceConnectorType.MCP_CONNECTOR, - ) - ) - connector = result.scalars().first() - if not connector: - raise HTTPException(status_code=404, detail="MCP connector not found") - - config = dict(connector.config or {}) - trusted: list[str] = list(config.get("trusted_tools", [])) - if body.tool_name not in trusted: - trusted.append(body.tool_name) - config["trusted_tools"] = trusted - connector.config = config - - from sqlalchemy.orm.attributes import flag_modified - - flag_modified(connector, "config") - await session.commit() - - from app.agents.new_chat.tools.mcp_tool import invalidate_mcp_tools_cache - - invalidate_mcp_tools_cache(connector.search_space_id) - - return {"status": "ok", "trusted_tools": trusted} - - except HTTPException: - raise - except Exception as e: - logger.error(f"Failed to trust MCP tool: {e!s}", exc_info=True) - await session.rollback() - raise HTTPException( - status_code=500, detail=f"Failed to trust tool: {e!s}" - ) from e - - -@router.post("/connectors/mcp/{connector_id}/untrust-tool") -async def untrust_mcp_tool( - connector_id: int, - body: MCPTrustToolRequest, - session: AsyncSession = Depends(get_async_session), - user: User = Depends(current_active_user), -): - """Remove a tool from the MCP connector's trusted list. - - The tool will require HITL approval again on subsequent calls. - """ - try: - result = await session.execute( - select(SearchSourceConnector).filter( - SearchSourceConnector.id == connector_id, - SearchSourceConnector.connector_type - == SearchSourceConnectorType.MCP_CONNECTOR, - ) - ) - connector = result.scalars().first() - if not connector: - raise HTTPException(status_code=404, detail="MCP connector not found") - - config = dict(connector.config or {}) - trusted: list[str] = list(config.get("trusted_tools", [])) - if body.tool_name in trusted: - trusted.remove(body.tool_name) - config["trusted_tools"] = trusted - connector.config = config - - from sqlalchemy.orm.attributes import flag_modified - - flag_modified(connector, "config") - await session.commit() - - from app.agents.new_chat.tools.mcp_tool import invalidate_mcp_tools_cache - - invalidate_mcp_tools_cache(connector.search_space_id) - - return {"status": "ok", "trusted_tools": trusted} - - except HTTPException: - raise - except Exception as e: - logger.error(f"Failed to untrust MCP tool: {e!s}", exc_info=True) - await session.rollback() - raise HTTPException( - status_code=500, detail=f"Failed to untrust tool: {e!s}" - ) from e diff --git a/surfsense_backend/app/routes/search_spaces_routes.py b/surfsense_backend/app/routes/search_spaces_routes.py index 828137518..7fa33ba1c 100644 --- a/surfsense_backend/app/routes/search_spaces_routes.py +++ b/surfsense_backend/app/routes/search_spaces_routes.py @@ -216,7 +216,6 @@ async def read_search_spaces( user_id=space.user_id, citations_enabled=space.citations_enabled, qna_custom_instructions=space.qna_custom_instructions, - ai_file_sort_enabled=space.ai_file_sort_enabled, member_count=member_count, is_owner=is_owner, ) @@ -385,42 +384,6 @@ async def edit_team_memory( return db_search_space -@router.post("/searchspaces/{search_space_id}/ai-sort") -async def trigger_ai_sort( - search_space_id: int, - session: AsyncSession = Depends(get_async_session), - user: User = Depends(current_active_user), -): - """Trigger a full AI file sort for all documents in the search space.""" - try: - await check_permission( - session, - user, - search_space_id, - Permission.SETTINGS_UPDATE.value, - "You don't have permission to trigger AI sort on this search space", - ) - - result = await session.execute( - select(SearchSpace).filter(SearchSpace.id == search_space_id) - ) - db_search_space = result.scalars().first() - if not db_search_space: - raise HTTPException(status_code=404, detail="Search space not found") - - from app.tasks.celery_tasks.document_tasks import ai_sort_search_space_task - - ai_sort_search_space_task.delay(search_space_id, str(user.id)) - return {"message": "AI sort started"} - except HTTPException: - raise - except Exception as e: - logger.error(f"Failed to trigger AI sort: {e!s}", exc_info=True) - raise HTTPException( - status_code=500, detail=f"Failed to trigger AI sort: {e!s}" - ) from e - - @router.delete("/searchspaces/{search_space_id}", response_model=dict) async def delete_search_space( search_space_id: int, diff --git a/surfsense_backend/app/schemas/search_space.py b/surfsense_backend/app/schemas/search_space.py index 77e34ea4b..e3b50be65 100644 --- a/surfsense_backend/app/schemas/search_space.py +++ b/surfsense_backend/app/schemas/search_space.py @@ -22,7 +22,6 @@ class SearchSpaceUpdate(BaseModel): citations_enabled: bool | None = None qna_custom_instructions: str | None = None shared_memory_md: str | None = None - ai_file_sort_enabled: bool | None = None class SearchSpaceRead(SearchSpaceBase, IDModel, TimestampModel): @@ -32,7 +31,6 @@ class SearchSpaceRead(SearchSpaceBase, IDModel, TimestampModel): citations_enabled: bool qna_custom_instructions: str | None = None shared_memory_md: str | None = None - ai_file_sort_enabled: bool = False model_config = ConfigDict(from_attributes=True) diff --git a/surfsense_backend/app/services/ai_file_sort_service.py b/surfsense_backend/app/services/ai_file_sort_service.py deleted file mode 100644 index 2f04131a6..000000000 --- a/surfsense_backend/app/services/ai_file_sort_service.py +++ /dev/null @@ -1,329 +0,0 @@ -"""AI File Sort Service: builds connector-type/date/category/subcategory folder paths.""" - -from __future__ import annotations - -import json -import logging -import re -from datetime import UTC, datetime - -from langchain_core.messages import HumanMessage -from sqlalchemy.ext.asyncio import AsyncSession -from sqlalchemy.future import select -from sqlalchemy.orm import selectinload - -from app.db import ( - Chunk, - Document, - DocumentType, - SearchSourceConnector, - SearchSourceConnectorType, -) -from app.services.folder_service import ensure_folder_hierarchy_with_depth_validation - -logger = logging.getLogger(__name__) - -_DOCTYPE_TO_CONNECTOR_LABEL: dict[str, str] = { - DocumentType.EXTENSION: "Browser Extension", - DocumentType.CRAWLED_URL: "Web Crawl", - DocumentType.FILE: "File Upload", - DocumentType.SLACK_CONNECTOR: "Slack", - DocumentType.TEAMS_CONNECTOR: "Teams", - DocumentType.ONEDRIVE_FILE: "OneDrive", - DocumentType.NOTION_CONNECTOR: "Notion", - DocumentType.YOUTUBE_VIDEO: "YouTube", - DocumentType.GITHUB_CONNECTOR: "GitHub", - DocumentType.LINEAR_CONNECTOR: "Linear", - DocumentType.DISCORD_CONNECTOR: "Discord", - DocumentType.JIRA_CONNECTOR: "Jira", - DocumentType.CONFLUENCE_CONNECTOR: "Confluence", - DocumentType.CLICKUP_CONNECTOR: "ClickUp", - DocumentType.GOOGLE_CALENDAR_CONNECTOR: "Google Calendar", - DocumentType.GOOGLE_GMAIL_CONNECTOR: "Gmail", - DocumentType.GOOGLE_DRIVE_FILE: "Google Drive", - DocumentType.AIRTABLE_CONNECTOR: "Airtable", - DocumentType.LUMA_CONNECTOR: "Luma", - DocumentType.ELASTICSEARCH_CONNECTOR: "Elasticsearch", - DocumentType.BOOKSTACK_CONNECTOR: "BookStack", - DocumentType.CIRCLEBACK: "Circleback", - DocumentType.OBSIDIAN_CONNECTOR: "Obsidian", - DocumentType.NOTE: "Notes", - DocumentType.DROPBOX_FILE: "Dropbox", - DocumentType.COMPOSIO_GOOGLE_DRIVE_CONNECTOR: "Google Drive (Composio)", - DocumentType.COMPOSIO_GMAIL_CONNECTOR: "Gmail (Composio)", - DocumentType.COMPOSIO_GOOGLE_CALENDAR_CONNECTOR: "Google Calendar (Composio)", - DocumentType.LOCAL_FOLDER_FILE: "Local Folder", -} - -_CONNECTOR_TYPE_LABEL: dict[str, str] = { - SearchSourceConnectorType.SERPER_API: "Serper Search", - SearchSourceConnectorType.TAVILY_API: "Tavily Search", - SearchSourceConnectorType.SEARXNG_API: "SearXNG Search", - SearchSourceConnectorType.LINKUP_API: "Linkup Search", - SearchSourceConnectorType.BAIDU_SEARCH_API: "Baidu Search", - SearchSourceConnectorType.SLACK_CONNECTOR: "Slack", - SearchSourceConnectorType.TEAMS_CONNECTOR: "Teams", - SearchSourceConnectorType.ONEDRIVE_CONNECTOR: "OneDrive", - SearchSourceConnectorType.NOTION_CONNECTOR: "Notion", - SearchSourceConnectorType.GITHUB_CONNECTOR: "GitHub", - SearchSourceConnectorType.LINEAR_CONNECTOR: "Linear", - SearchSourceConnectorType.DISCORD_CONNECTOR: "Discord", - SearchSourceConnectorType.JIRA_CONNECTOR: "Jira", - SearchSourceConnectorType.CONFLUENCE_CONNECTOR: "Confluence", - SearchSourceConnectorType.CLICKUP_CONNECTOR: "ClickUp", - SearchSourceConnectorType.GOOGLE_CALENDAR_CONNECTOR: "Google Calendar", - SearchSourceConnectorType.GOOGLE_GMAIL_CONNECTOR: "Gmail", - SearchSourceConnectorType.GOOGLE_DRIVE_CONNECTOR: "Google Drive", - SearchSourceConnectorType.AIRTABLE_CONNECTOR: "Airtable", - SearchSourceConnectorType.LUMA_CONNECTOR: "Luma", - SearchSourceConnectorType.ELASTICSEARCH_CONNECTOR: "Elasticsearch", - SearchSourceConnectorType.WEBCRAWLER_CONNECTOR: "Web Crawl", - SearchSourceConnectorType.BOOKSTACK_CONNECTOR: "BookStack", - SearchSourceConnectorType.CIRCLEBACK_CONNECTOR: "Circleback", - SearchSourceConnectorType.OBSIDIAN_CONNECTOR: "Obsidian", - SearchSourceConnectorType.MCP_CONNECTOR: "MCP", - SearchSourceConnectorType.DROPBOX_CONNECTOR: "Dropbox", - SearchSourceConnectorType.COMPOSIO_GOOGLE_DRIVE_CONNECTOR: "Google Drive (Composio)", - SearchSourceConnectorType.COMPOSIO_GMAIL_CONNECTOR: "Gmail (Composio)", - SearchSourceConnectorType.COMPOSIO_GOOGLE_CALENDAR_CONNECTOR: "Google Calendar (Composio)", -} - -_MAX_CONTENT_CHARS = 4000 -_MAX_CHUNKS_FOR_CONTEXT = 5 - -_CATEGORY_PROMPT = ( - "Based on the document information below, classify it into a broad category " - "and a more specific subcategory.\n\n" - "Rules:\n" - "- category: 1-2 word broad theme (e.g. Science, Finance, Engineering, Communication, Media)\n" - "- subcategory: 1-2 word specific topic within the category " - "(e.g. Physics, Tax Reports, Backend, Team Updates)\n" - "- Use nouns only. Do not include generic terms like 'General' or 'Miscellaneous'.\n\n" - "Title: {title}\n\n" - "Content: {summary}\n\n" - 'Respond with ONLY a JSON object: {{"category": "...", "subcategory": "..."}}' -) - -_SAFE_NAME_RE = re.compile(r"[^a-zA-Z0-9 _\-()]") -_FALLBACK_CATEGORY = "Uncategorized" -_FALLBACK_SUBCATEGORY = "General" - - -def resolve_root_folder_label( - document: Document, connector: SearchSourceConnector | None -) -> str: - if connector is not None: - return _CONNECTOR_TYPE_LABEL.get( - connector.connector_type, str(connector.connector_type) - ) - return _DOCTYPE_TO_CONNECTOR_LABEL.get( - document.document_type, str(document.document_type) - ) - - -def resolve_date_folder(document: Document) -> str: - ts = document.updated_at or document.created_at - if ts is None: - ts = datetime.now(UTC) - return ts.strftime("%Y-%m-%d") - - -def sanitize_category_folder_name( - value: str | None, fallback: str = _FALLBACK_CATEGORY -) -> str: - if not value or not value.strip(): - return fallback - cleaned = _SAFE_NAME_RE.sub("", value.strip()) - cleaned = " ".join(cleaned.split()) - if not cleaned: - return fallback - return cleaned[:50] - - -async def _resolve_document_text( - session: AsyncSession, - document: Document, -) -> str: - """Build the best available text representation for taxonomy generation. - - Prefers ``document.content``; falls back to joining the first few chunks - when content is empty or too short to be useful. - """ - text = (document.content or "").strip() - if len(text) >= 100: - return text[:_MAX_CONTENT_CHARS] - - stmt = ( - select(Chunk.content) - .where(Chunk.document_id == document.id) - .order_by(Chunk.id) - .limit(_MAX_CHUNKS_FOR_CONTEXT) - ) - result = await session.execute(stmt) - chunk_texts = [row[0] for row in result.all() if row[0]] - if chunk_texts: - combined = "\n\n".join(chunk_texts) - return combined[:_MAX_CONTENT_CHARS] - - return text[:_MAX_CONTENT_CHARS] - - -def _get_cached_taxonomy(document: Document) -> tuple[str, str] | None: - """Return (category, subcategory) from document metadata cache, or None.""" - meta = document.document_metadata - if not isinstance(meta, dict): - return None - cat = meta.get("ai_sort_category") - subcat = meta.get("ai_sort_subcategory") - if cat and subcat and isinstance(cat, str) and isinstance(subcat, str): - return cat, subcat - return None - - -def _set_cached_taxonomy(document: Document, category: str, subcategory: str) -> None: - """Persist the AI taxonomy on document metadata for deterministic re-sorts.""" - meta = dict(document.document_metadata or {}) - meta["ai_sort_category"] = category - meta["ai_sort_subcategory"] = subcategory - document.document_metadata = meta - - -async def generate_ai_taxonomy( - title: str, - summary_or_content: str, - llm, -) -> tuple[str, str]: - """Return (category, subcategory) using a single structured LLM call.""" - text = (summary_or_content or "").strip() - if not text: - return _FALLBACK_CATEGORY, _FALLBACK_SUBCATEGORY - - if len(text) > _MAX_CONTENT_CHARS: - text = text[:_MAX_CONTENT_CHARS] - - prompt = _CATEGORY_PROMPT.format(title=title or "Untitled", summary=text) - try: - result = await llm.ainvoke([HumanMessage(content=prompt)]) - raw = result.content.strip() - if raw.startswith("```"): - raw = raw.split("\n", 1)[-1].rsplit("```", 1)[0].strip() - parsed = json.loads(raw) - category = sanitize_category_folder_name( - parsed.get("category"), _FALLBACK_CATEGORY - ) - subcategory = sanitize_category_folder_name( - parsed.get("subcategory"), _FALLBACK_SUBCATEGORY - ) - return category, subcategory - except Exception: - logger.warning("AI taxonomy generation failed, using fallback", exc_info=True) - return _FALLBACK_CATEGORY, _FALLBACK_SUBCATEGORY - - -def _build_path_segments( - root_label: str, - date_label: str, - category: str, - subcategory: str, -) -> list[dict]: - return [ - {"name": root_label, "metadata": {"ai_sort": True, "ai_sort_level": 1}}, - {"name": date_label, "metadata": {"ai_sort": True, "ai_sort_level": 2}}, - {"name": category, "metadata": {"ai_sort": True, "ai_sort_level": 3}}, - {"name": subcategory, "metadata": {"ai_sort": True, "ai_sort_level": 4}}, - ] - - -async def _resolve_taxonomy( - session: AsyncSession, - document: Document, - llm, -) -> tuple[str, str]: - """Return (category, subcategory), reusing cached values when available.""" - cached = _get_cached_taxonomy(document) - if cached is not None: - return cached - - content_text = await _resolve_document_text(session, document) - category, subcategory = await generate_ai_taxonomy( - document.title, content_text, llm - ) - _set_cached_taxonomy(document, category, subcategory) - return category, subcategory - - -async def ai_sort_document( - session: AsyncSession, - document: Document, - llm, -) -> Document: - """Sort a single document into the 4-level AI folder hierarchy.""" - connector: SearchSourceConnector | None = None - if document.connector_id is not None: - connector = await session.get(SearchSourceConnector, document.connector_id) - - root_label = resolve_root_folder_label(document, connector) - date_label = resolve_date_folder(document) - - category, subcategory = await _resolve_taxonomy(session, document, llm) - - segments = _build_path_segments(root_label, date_label, category, subcategory) - - leaf_folder = await ensure_folder_hierarchy_with_depth_validation( - session, - document.search_space_id, - segments, - ) - - document.folder_id = leaf_folder.id - await session.flush() - return document - - -async def ai_sort_all_documents( - session: AsyncSession, - search_space_id: int, - llm, -) -> tuple[int, int]: - """Sort all documents in a search space. Returns (sorted_count, failed_count).""" - stmt = ( - select(Document) - .where(Document.search_space_id == search_space_id) - .options(selectinload(Document.connector)) - ) - result = await session.execute(stmt) - documents = list(result.scalars().all()) - - sorted_count = 0 - failed_count = 0 - - for doc in documents: - try: - connector = doc.connector - root_label = resolve_root_folder_label(doc, connector) - date_label = resolve_date_folder(doc) - - category, subcategory = await _resolve_taxonomy(session, doc, llm) - segments = _build_path_segments( - root_label, date_label, category, subcategory - ) - - leaf_folder = await ensure_folder_hierarchy_with_depth_validation( - session, - search_space_id, - segments, - ) - doc.folder_id = leaf_folder.id - sorted_count += 1 - except Exception: - logger.error("Failed to AI-sort document %s", doc.id, exc_info=True) - failed_count += 1 - - await session.commit() - logger.info( - "AI sort complete for search_space=%d: sorted=%d, failed=%d", - search_space_id, - sorted_count, - failed_count, - ) - return sorted_count, failed_count diff --git a/surfsense_backend/app/services/folder_service.py b/surfsense_backend/app/services/folder_service.py index f5b608600..fc1fb8a75 100644 --- a/surfsense_backend/app/services/folder_service.py +++ b/surfsense_backend/app/services/folder_service.py @@ -142,58 +142,6 @@ async def generate_folder_position( return generate_key_between(last_position, None) -async def ensure_folder_hierarchy_with_depth_validation( - session: AsyncSession, - search_space_id: int, - path_segments: list[dict], -) -> Folder: - """Create or return a nested folder chain, validating depth at each step. - - Each item in ``path_segments`` is a dict with: - - ``name`` (str): folder display name - - ``metadata`` (dict | None): optional ``folder_metadata`` JSONB payload - - Returns the deepest (leaf) Folder in the chain. - """ - parent_id: int | None = None - current_folder: Folder | None = None - - for segment in path_segments: - name = segment["name"] - metadata = segment.get("metadata") - - stmt = select(Folder).where( - Folder.search_space_id == search_space_id, - Folder.name == name, - Folder.parent_id == parent_id - if parent_id is not None - else Folder.parent_id.is_(None), - ) - result = await session.execute(stmt) - folder = result.scalar_one_or_none() - - if folder is None: - await validate_folder_depth(session, parent_id, subtree_depth=0) - position = await generate_folder_position( - session, search_space_id, parent_id - ) - folder = Folder( - name=name, - search_space_id=search_space_id, - parent_id=parent_id, - position=position, - folder_metadata=metadata, - ) - session.add(folder) - await session.flush() - - current_folder = folder - parent_id = folder.id - - assert current_folder is not None, "path_segments must not be empty" - return current_folder - - async def get_folder_subtree_ids(session: AsyncSession, folder_id: int) -> list[int]: """Return all folder IDs in the subtree rooted at folder_id (inclusive).""" result = await session.execute( diff --git a/surfsense_backend/app/tasks/celery_tasks/document_tasks.py b/surfsense_backend/app/tasks/celery_tasks/document_tasks.py index 719cfb940..fc946b4bc 100644 --- a/surfsense_backend/app/tasks/celery_tasks/document_tasks.py +++ b/surfsense_backend/app/tasks/celery_tasks/document_tasks.py @@ -4,7 +4,6 @@ import asyncio import contextlib import logging import os -import time from uuid import UUID from app.celery_app import celery_app @@ -1552,121 +1551,3 @@ async def _index_uploaded_folder_files_async( heartbeat_task.cancel() if notification_id is not None: _stop_heartbeat(notification_id) - - -# ===== AI File Sort tasks ===== - -AI_SORT_LOCK_TTL_SECONDS = 600 # 10 minutes -_ai_sort_redis = None - - -def _get_ai_sort_redis(): - import redis - - global _ai_sort_redis - if _ai_sort_redis is None: - _ai_sort_redis = redis.from_url(config.REDIS_APP_URL, decode_responses=True) - return _ai_sort_redis - - -def _ai_sort_lock_key(search_space_id: int) -> str: - return f"ai_sort:search_space:{search_space_id}:lock" - - -@celery_app.task(name="ai_sort_search_space", bind=True, max_retries=1) -def ai_sort_search_space_task(self, search_space_id: int, user_id: str): - """Full AI sort for all documents in a search space.""" - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - try: - loop.run_until_complete(_ai_sort_search_space_async(search_space_id, user_id)) - finally: - loop.close() - - -async def _ai_sort_search_space_async(search_space_id: int, user_id: str): - r = _get_ai_sort_redis() - lock_key = _ai_sort_lock_key(search_space_id) - - if not r.set(lock_key, "running", nx=True, ex=AI_SORT_LOCK_TTL_SECONDS): - logger.info( - "AI sort already running for search_space=%d, skipping", - search_space_id, - ) - return - - t_start = time.perf_counter() - try: - from app.services.ai_file_sort_service import ai_sort_all_documents - from app.services.llm_service import get_document_summary_llm - - async with get_celery_session_maker()() as session: - llm = await get_document_summary_llm( - session, search_space_id, disable_streaming=True - ) - if llm is None: - logger.warning( - "No LLM configured for search_space=%d, skipping AI sort", - search_space_id, - ) - return - - sorted_count, failed_count = await ai_sort_all_documents( - session, search_space_id, llm - ) - elapsed = time.perf_counter() - t_start - logger.info( - "AI sort search_space=%d done in %.1fs: sorted=%d failed=%d", - search_space_id, - elapsed, - sorted_count, - failed_count, - ) - finally: - r.delete(lock_key) - - -@celery_app.task( - name="ai_sort_document", bind=True, max_retries=2, default_retry_delay=10 -) -def ai_sort_document_task(self, search_space_id: int, user_id: str, document_id: int): - """Incremental AI sort for a single document after indexing.""" - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - try: - loop.run_until_complete( - _ai_sort_document_async(search_space_id, user_id, document_id) - ) - finally: - loop.close() - - -async def _ai_sort_document_async(search_space_id: int, user_id: str, document_id: int): - from app.db import Document - from app.services.ai_file_sort_service import ai_sort_document - from app.services.llm_service import get_document_summary_llm - - async with get_celery_session_maker()() as session: - document = await session.get(Document, document_id) - if document is None: - logger.warning("Document %d not found, skipping AI sort", document_id) - return - - llm = await get_document_summary_llm( - session, search_space_id, disable_streaming=True - ) - if llm is None: - logger.warning( - "No LLM for search_space=%d, skipping AI sort of doc=%d", - search_space_id, - document_id, - ) - return - - await ai_sort_document(session, document, llm) - await session.commit() - logger.info( - "AI sorted document=%d into search_space=%d", - document_id, - search_space_id, - ) diff --git a/surfsense_backend/app/tasks/chat/stream_new_chat.py b/surfsense_backend/app/tasks/chat/stream_new_chat.py index 47a270568..fd118528e 100644 --- a/surfsense_backend/app/tasks/chat/stream_new_chat.py +++ b/surfsense_backend/app/tasks/chat/stream_new_chat.py @@ -61,7 +61,6 @@ from app.services.new_streaming_service import VercelStreamingService from app.utils.content_utils import bootstrap_history_from_db from app.utils.perf import get_perf_logger, log_system_snapshot, trim_native_heap -_background_tasks: set[asyncio.Task] = set() _perf_log = get_perf_logger() @@ -143,7 +142,7 @@ class StreamResult: accumulated_text: str = "" is_interrupted: bool = False interrupt_value: dict[str, Any] | None = None - sandbox_files: list[str] = field(default_factory=list) + sandbox_files: list[str] = field(default_factory=list) # unused, kept for compat agent_called_update_memory: bool = False @@ -441,7 +440,7 @@ async def _stream_agent_events( status="in_progress", items=last_active_step_items, ) - elif tool_name in ("execute", "execute_code"): + elif tool_name == "execute": cmd = ( tool_input.get("command", "") if isinstance(tool_input, dict) @@ -739,7 +738,7 @@ async def _stream_agent_events( status="completed", items=completed_items, ) - elif tool_name in ("execute", "execute_code"): + elif tool_name == "execute": raw_text = ( tool_output.get("result", "") if isinstance(tool_output, dict) @@ -986,7 +985,7 @@ async def _stream_agent_events( if isinstance(tool_output, dict) else {"result": tool_output}, ) - elif tool_name in ("execute", "execute_code"): + elif tool_name == "execute": raw_text = ( tool_output.get("result", "") if isinstance(tool_output, dict) @@ -1553,7 +1552,7 @@ async def stream_new_chat( # Shared threads write to team memory; private threads write to user memory. if not stream_result.agent_called_update_memory: if visibility == ChatVisibility.SEARCH_SPACE: - task = asyncio.create_task( + asyncio.create_task( extract_and_save_team_memory( user_message=user_query, search_space_id=search_space_id, @@ -1561,18 +1560,14 @@ async def stream_new_chat( author_display_name=current_user_display_name, ) ) - _background_tasks.add(task) - task.add_done_callback(_background_tasks.discard) elif user_id: - task = asyncio.create_task( + asyncio.create_task( extract_and_save_memory( user_message=user_query, user_id=user_id, llm=llm, ) ) - _background_tasks.add(task) - task.add_done_callback(_background_tasks.discard) # Finish the step and message yield streaming_service.format_finish_step() @@ -1622,21 +1617,6 @@ async def stream_new_chat( with contextlib.suppress(Exception): await session.close() - # Persist any sandbox-produced files to local storage so they - # remain downloadable after the Daytona sandbox auto-deletes. - if stream_result and stream_result.sandbox_files: - with contextlib.suppress(Exception): - from app.agents.new_chat.sandbox import ( - is_sandbox_enabled, - persist_and_delete_sandbox, - ) - - if is_sandbox_enabled(): - with anyio.CancelScope(shield=True): - await persist_and_delete_sandbox( - chat_id, stream_result.sandbox_files - ) - # Break circular refs held by the agent graph, tools, and LLM # wrappers so the GC can reclaim them in a single pass. agent = llm = connector_service = None diff --git a/surfsense_backend/app/tasks/connector_indexers/google_drive_indexer.py b/surfsense_backend/app/tasks/connector_indexers/google_drive_indexer.py index 21cdbd29f..d8f95da63 100644 --- a/surfsense_backend/app/tasks/connector_indexers/google_drive_indexer.py +++ b/surfsense_backend/app/tasks/connector_indexers/google_drive_indexer.py @@ -961,7 +961,6 @@ async def index_google_drive_files( vision_llm = None if connector_enable_vision_llm: from app.services.llm_service import get_vision_llm - vision_llm = await get_vision_llm(session, search_space_id) drive_client = GoogleDriveClient( session, connector_id, credentials=pre_built_credentials @@ -1169,7 +1168,6 @@ async def index_google_drive_single_file( vision_llm = None if connector_enable_vision_llm: from app.services.llm_service import get_vision_llm - vision_llm = await get_vision_llm(session, search_space_id) drive_client = GoogleDriveClient( session, connector_id, credentials=pre_built_credentials @@ -1308,7 +1306,6 @@ async def index_google_drive_selected_files( vision_llm = None if connector_enable_vision_llm: from app.services.llm_service import get_vision_llm - vision_llm = await get_vision_llm(session, search_space_id) drive_client = GoogleDriveClient( session, connector_id, credentials=pre_built_credentials diff --git a/surfsense_backend/app/tasks/connector_indexers/local_folder_indexer.py b/surfsense_backend/app/tasks/connector_indexers/local_folder_indexer.py index b6797f77a..2d5f9648d 100644 --- a/surfsense_backend/app/tasks/connector_indexers/local_folder_indexer.py +++ b/surfsense_backend/app/tasks/connector_indexers/local_folder_indexer.py @@ -1360,9 +1360,7 @@ async def index_uploaded_files( try: content, content_hash = await _compute_file_content_hash( - temp_path, - filename, - search_space_id, + temp_path, filename, search_space_id, vision_llm=vision_llm_instance, ) except Exception as e: diff --git a/surfsense_backend/app/tasks/connector_indexers/onedrive_indexer.py b/surfsense_backend/app/tasks/connector_indexers/onedrive_indexer.py index 2def799f3..aa654a9a9 100644 --- a/surfsense_backend/app/tasks/connector_indexers/onedrive_indexer.py +++ b/surfsense_backend/app/tasks/connector_indexers/onedrive_indexer.py @@ -656,7 +656,6 @@ async def index_onedrive_files( vision_llm = None if connector_enable_vision_llm: from app.services.llm_service import get_vision_llm - vision_llm = await get_vision_llm(session, search_space_id) onedrive_client = OneDriveClient(session, connector_id) diff --git a/surfsense_backend/scripts/create_sandbox_snapshot.py b/surfsense_backend/scripts/create_sandbox_snapshot.py index 82d452174..97ed6dfe8 100644 --- a/surfsense_backend/scripts/create_sandbox_snapshot.py +++ b/surfsense_backend/scripts/create_sandbox_snapshot.py @@ -21,11 +21,7 @@ from pathlib import Path from dotenv import load_dotenv _here = Path(__file__).parent -for candidate in [ - _here / "../surfsense_backend/.env", - _here / ".env", - _here / "../.env", -]: +for candidate in [_here / "../surfsense_backend/.env", _here / ".env", _here / "../.env"]: if candidate.exists(): load_dotenv(candidate) break @@ -61,10 +57,7 @@ def main() -> None: api_key = os.environ.get("DAYTONA_API_KEY") if not api_key: print("ERROR: DAYTONA_API_KEY is not set.", file=sys.stderr) - print( - "Add it to surfsense_backend/.env or export it in your shell.", - file=sys.stderr, - ) + print("Add it to surfsense_backend/.env or export it in your shell.", file=sys.stderr) sys.exit(1) daytona = Daytona() @@ -74,7 +67,7 @@ def main() -> None: print(f"Deleting existing snapshot '{SNAPSHOT_NAME}' …") daytona.snapshot.delete(existing) print(f"Deleted '{SNAPSHOT_NAME}'. Waiting for removal to propagate …") - for _attempt in range(30): + for attempt in range(30): time.sleep(2) try: daytona.snapshot.get(SNAPSHOT_NAME) @@ -82,9 +75,7 @@ def main() -> None: print(f"Confirmed '{SNAPSHOT_NAME}' is gone.\n") break else: - print( - f"WARNING: '{SNAPSHOT_NAME}' may still exist after 60s. Proceeding anyway.\n" - ) + print(f"WARNING: '{SNAPSHOT_NAME}' may still exist after 60s. Proceeding anyway.\n") except Exception: pass diff --git a/surfsense_backend/tests/unit/etl_pipeline/test_etl_pipeline_service.py b/surfsense_backend/tests/unit/etl_pipeline/test_etl_pipeline_service.py index bb8d7ef83..1a94d4263 100644 --- a/surfsense_backend/tests/unit/etl_pipeline/test_etl_pipeline_service.py +++ b/surfsense_backend/tests/unit/etl_pipeline/test_etl_pipeline_service.py @@ -431,9 +431,7 @@ async def test_llamacloud_heif_accepted_only_with_azure_di(tmp_path, mocker): mocker.patch("app.config.config.AZURE_DI_ENDPOINT", None, create=True) mocker.patch("app.config.config.AZURE_DI_KEY", None, create=True) - with pytest.raises( - EtlUnsupportedFileError, match="document parser does not support this format" - ): + with pytest.raises(EtlUnsupportedFileError, match="document parser does not support this format"): await EtlPipelineService().extract( EtlRequest(file_path=str(heif_file), filename="photo.heif") ) diff --git a/surfsense_backend/tests/unit/middleware/test_knowledge_search.py b/surfsense_backend/tests/unit/middleware/test_knowledge_search.py index 1aaf5d127..a8cf5c93b 100644 --- a/surfsense_backend/tests/unit/middleware/test_knowledge_search.py +++ b/surfsense_backend/tests/unit/middleware/test_knowledge_search.py @@ -6,7 +6,6 @@ import pytest from langchain_core.messages import AIMessage, HumanMessage from app.agents.new_chat.middleware.knowledge_search import ( - KBSearchPlan, KnowledgeBaseSearchMiddleware, _build_document_xml, _normalize_optional_date_range, @@ -367,146 +366,3 @@ class TestKnowledgeBaseSearchMiddlewarePlanner: assert captured["query"] == "deel founders guide summary" assert captured["start_date"] is None assert captured["end_date"] is None - - async def test_middleware_routes_to_recency_browse_when_flagged( - self, - monkeypatch, - ): - """When the planner sets is_recency_query=true, browse_recent_documents - is called instead of search_knowledge_base.""" - browse_captured: dict = {} - search_called = False - - async def fake_browse_recent_documents(**kwargs): - browse_captured.update(kwargs) - return [] - - async def fake_search_knowledge_base(**kwargs): - nonlocal search_called - search_called = True - return [] - - async def fake_build_scoped_filesystem(**kwargs): - return {}, {} - - monkeypatch.setattr( - "app.agents.new_chat.middleware.knowledge_search.browse_recent_documents", - fake_browse_recent_documents, - ) - monkeypatch.setattr( - "app.agents.new_chat.middleware.knowledge_search.search_knowledge_base", - fake_search_knowledge_base, - ) - monkeypatch.setattr( - "app.agents.new_chat.middleware.knowledge_search.build_scoped_filesystem", - fake_build_scoped_filesystem, - ) - - llm = FakeLLM( - json.dumps( - { - "optimized_query": "latest uploaded file", - "start_date": None, - "end_date": None, - "is_recency_query": True, - } - ) - ) - middleware = KnowledgeBaseSearchMiddleware(llm=llm, search_space_id=42) - - result = await middleware.abefore_agent( - {"messages": [HumanMessage(content="what's my latest file?")]}, - runtime=None, - ) - - assert result is not None - assert browse_captured["search_space_id"] == 42 - assert not search_called - - async def test_middleware_uses_hybrid_search_when_not_recency( - self, - monkeypatch, - ): - """When is_recency_query is false (default), hybrid search is used.""" - search_captured: dict = {} - browse_called = False - - async def fake_browse_recent_documents(**kwargs): - nonlocal browse_called - browse_called = True - return [] - - async def fake_search_knowledge_base(**kwargs): - search_captured.update(kwargs) - return [] - - async def fake_build_scoped_filesystem(**kwargs): - return {}, {} - - monkeypatch.setattr( - "app.agents.new_chat.middleware.knowledge_search.browse_recent_documents", - fake_browse_recent_documents, - ) - monkeypatch.setattr( - "app.agents.new_chat.middleware.knowledge_search.search_knowledge_base", - fake_search_knowledge_base, - ) - monkeypatch.setattr( - "app.agents.new_chat.middleware.knowledge_search.build_scoped_filesystem", - fake_build_scoped_filesystem, - ) - - llm = FakeLLM( - json.dumps( - { - "optimized_query": "quarterly revenue report analysis", - "start_date": None, - "end_date": None, - "is_recency_query": False, - } - ) - ) - middleware = KnowledgeBaseSearchMiddleware(llm=llm, search_space_id=42) - - await middleware.abefore_agent( - {"messages": [HumanMessage(content="find the quarterly revenue report")]}, - runtime=None, - ) - - assert search_captured["query"] == "quarterly revenue report analysis" - assert not browse_called - - -# ── KBSearchPlan schema ──────────────────────────────────────────────── - - -class TestKBSearchPlanSchema: - def test_is_recency_query_defaults_to_false(self): - plan = KBSearchPlan(optimized_query="test query") - assert plan.is_recency_query is False - - def test_is_recency_query_parses_true(self): - plan = _parse_kb_search_plan_response( - json.dumps( - { - "optimized_query": "latest uploaded file", - "start_date": None, - "end_date": None, - "is_recency_query": True, - } - ) - ) - assert plan.is_recency_query is True - assert plan.optimized_query == "latest uploaded file" - - def test_missing_is_recency_query_defaults_to_false(self): - plan = _parse_kb_search_plan_response( - json.dumps( - { - "optimized_query": "meeting notes", - "start_date": None, - "end_date": None, - } - ) - ) - assert plan.is_recency_query is False diff --git a/surfsense_backend/tests/unit/services/test_ai_file_sort_service.py b/surfsense_backend/tests/unit/services/test_ai_file_sort_service.py deleted file mode 100644 index 860c2ffa2..000000000 --- a/surfsense_backend/tests/unit/services/test_ai_file_sort_service.py +++ /dev/null @@ -1,275 +0,0 @@ -"""Unit tests for AI file sort service: folder label resolution, date extraction, category sanitization.""" - -from datetime import UTC, datetime -from unittest.mock import AsyncMock, MagicMock - -import pytest - -pytestmark = pytest.mark.unit - - -# ── resolve_root_folder_label ── - - -def _make_document(document_type: str, connector_id=None): - doc = MagicMock() - doc.document_type = document_type - doc.connector_id = connector_id - return doc - - -def _make_connector(connector_type: str): - conn = MagicMock() - conn.connector_type = connector_type - return conn - - -def test_root_label_uses_connector_type_when_available(): - from app.services.ai_file_sort_service import resolve_root_folder_label - - doc = _make_document("FILE", connector_id=1) - conn = _make_connector("GOOGLE_DRIVE_CONNECTOR") - assert resolve_root_folder_label(doc, conn) == "Google Drive" - - -def test_root_label_falls_back_to_document_type(): - from app.services.ai_file_sort_service import resolve_root_folder_label - - doc = _make_document("SLACK_CONNECTOR") - assert resolve_root_folder_label(doc, None) == "Slack" - - -def test_root_label_unknown_doctype_returns_raw_value(): - from app.services.ai_file_sort_service import resolve_root_folder_label - - doc = _make_document("UNKNOWN_TYPE") - assert resolve_root_folder_label(doc, None) == "UNKNOWN_TYPE" - - -# ── resolve_date_folder ── - - -def test_date_folder_from_updated_at(): - from app.services.ai_file_sort_service import resolve_date_folder - - doc = MagicMock() - doc.updated_at = datetime(2025, 3, 15, 10, 30, 0, tzinfo=UTC) - doc.created_at = datetime(2025, 1, 1, 0, 0, 0, tzinfo=UTC) - assert resolve_date_folder(doc) == "2025-03-15" - - -def test_date_folder_falls_back_to_created_at(): - from app.services.ai_file_sort_service import resolve_date_folder - - doc = MagicMock() - doc.updated_at = None - doc.created_at = datetime(2024, 12, 25, 23, 59, 0, tzinfo=UTC) - assert resolve_date_folder(doc) == "2024-12-25" - - -def test_date_folder_both_none_uses_today(): - from app.services.ai_file_sort_service import resolve_date_folder - - doc = MagicMock() - doc.updated_at = None - doc.created_at = None - result = resolve_date_folder(doc) - today = datetime.now(UTC).strftime("%Y-%m-%d") - assert result == today - - -# ── sanitize_category_folder_name ── - - -def test_sanitize_normal_value(): - from app.services.ai_file_sort_service import sanitize_category_folder_name - - assert sanitize_category_folder_name("Machine Learning") == "Machine Learning" - - -def test_sanitize_strips_special_chars(): - from app.services.ai_file_sort_service import sanitize_category_folder_name - - assert sanitize_category_folder_name("Tax/Reports!") == "TaxReports" - - -def test_sanitize_empty_returns_fallback(): - from app.services.ai_file_sort_service import sanitize_category_folder_name - - assert sanitize_category_folder_name("") == "Uncategorized" - assert sanitize_category_folder_name(None) == "Uncategorized" - - -def test_sanitize_truncates_long_names(): - from app.services.ai_file_sort_service import sanitize_category_folder_name - - long_name = "A" * 100 - result = sanitize_category_folder_name(long_name) - assert len(result) <= 50 - - -# ── generate_ai_taxonomy ── - - -@pytest.mark.asyncio -async def test_generate_ai_taxonomy_parses_json(): - from app.services.ai_file_sort_service import generate_ai_taxonomy - - mock_llm = AsyncMock() - mock_result = MagicMock() - mock_result.content = '{"category": "Science", "subcategory": "Physics"}' - mock_llm.ainvoke.return_value = mock_result - - cat, sub = await generate_ai_taxonomy( - "Physics Paper", "Some science document about physics", mock_llm - ) - assert cat == "Science" - assert sub == "Physics" - - -@pytest.mark.asyncio -async def test_generate_ai_taxonomy_handles_markdown_code_block(): - from app.services.ai_file_sort_service import generate_ai_taxonomy - - mock_llm = AsyncMock() - mock_result = MagicMock() - mock_result.content = ( - '```json\n{"category": "Finance", "subcategory": "Tax Reports"}\n```' - ) - mock_llm.ainvoke.return_value = mock_result - - cat, sub = await generate_ai_taxonomy("Tax Doc", "A tax report document", mock_llm) - assert cat == "Finance" - assert sub == "Tax Reports" - - -@pytest.mark.asyncio -async def test_generate_ai_taxonomy_includes_title_in_prompt(): - from app.services.ai_file_sort_service import generate_ai_taxonomy - - mock_llm = AsyncMock() - mock_result = MagicMock() - mock_result.content = '{"category": "Engineering", "subcategory": "Backend"}' - mock_llm.ainvoke.return_value = mock_result - - await generate_ai_taxonomy("API Design Guide", "content about REST APIs", mock_llm) - - prompt_text = mock_llm.ainvoke.call_args[0][0][0].content - assert "API Design Guide" in prompt_text - assert "content about REST APIs" in prompt_text - - -@pytest.mark.asyncio -async def test_generate_ai_taxonomy_fallback_on_error(): - from app.services.ai_file_sort_service import generate_ai_taxonomy - - mock_llm = AsyncMock() - mock_llm.ainvoke.side_effect = RuntimeError("LLM down") - - cat, sub = await generate_ai_taxonomy("Title", "some content", mock_llm) - assert cat == "Uncategorized" - assert sub == "General" - - -@pytest.mark.asyncio -async def test_generate_ai_taxonomy_fallback_on_empty_content(): - from app.services.ai_file_sort_service import generate_ai_taxonomy - - mock_llm = AsyncMock() - cat, sub = await generate_ai_taxonomy("Title", "", mock_llm) - assert cat == "Uncategorized" - assert sub == "General" - mock_llm.ainvoke.assert_not_called() - - -@pytest.mark.asyncio -async def test_generate_ai_taxonomy_fallback_on_invalid_json(): - from app.services.ai_file_sort_service import generate_ai_taxonomy - - mock_llm = AsyncMock() - mock_result = MagicMock() - mock_result.content = "not valid json at all" - mock_llm.ainvoke.return_value = mock_result - - cat, sub = await generate_ai_taxonomy("Title", "some content", mock_llm) - assert cat == "Uncategorized" - assert sub == "General" - - -# ── taxonomy caching ── - - -def test_get_cached_taxonomy_returns_none_when_no_metadata(): - from app.services.ai_file_sort_service import _get_cached_taxonomy - - doc = MagicMock() - doc.document_metadata = None - assert _get_cached_taxonomy(doc) is None - - -def test_get_cached_taxonomy_returns_none_when_keys_missing(): - from app.services.ai_file_sort_service import _get_cached_taxonomy - - doc = MagicMock() - doc.document_metadata = {"some_other_key": "value"} - assert _get_cached_taxonomy(doc) is None - - -def test_get_cached_taxonomy_returns_cached_values(): - from app.services.ai_file_sort_service import _get_cached_taxonomy - - doc = MagicMock() - doc.document_metadata = { - "ai_sort_category": "Finance", - "ai_sort_subcategory": "Tax Reports", - } - assert _get_cached_taxonomy(doc) == ("Finance", "Tax Reports") - - -def test_set_cached_taxonomy_persists_on_metadata(): - from app.services.ai_file_sort_service import _set_cached_taxonomy - - doc = MagicMock() - doc.document_metadata = {"existing_key": "keep_me"} - _set_cached_taxonomy(doc, "Science", "Physics") - assert doc.document_metadata["ai_sort_category"] == "Science" - assert doc.document_metadata["ai_sort_subcategory"] == "Physics" - assert doc.document_metadata["existing_key"] == "keep_me" - - -def test_set_cached_taxonomy_creates_metadata_when_none(): - from app.services.ai_file_sort_service import _set_cached_taxonomy - - doc = MagicMock() - doc.document_metadata = None - _set_cached_taxonomy(doc, "Engineering", "Backend") - assert doc.document_metadata == { - "ai_sort_category": "Engineering", - "ai_sort_subcategory": "Backend", - } - - -# ── _build_path_segments ── - - -def test_build_path_segments_structure(): - from app.services.ai_file_sort_service import _build_path_segments - - segments = _build_path_segments("Google Drive", "2025-03-15", "Science", "Physics") - assert len(segments) == 4 - assert segments[0] == { - "name": "Google Drive", - "metadata": {"ai_sort": True, "ai_sort_level": 1}, - } - assert segments[1] == { - "name": "2025-03-15", - "metadata": {"ai_sort": True, "ai_sort_level": 2}, - } - assert segments[2] == { - "name": "Science", - "metadata": {"ai_sort": True, "ai_sort_level": 3}, - } - assert segments[3] == { - "name": "Physics", - "metadata": {"ai_sort": True, "ai_sort_level": 4}, - } diff --git a/surfsense_backend/tests/unit/services/test_ai_sort_task_dedupe.py b/surfsense_backend/tests/unit/services/test_ai_sort_task_dedupe.py deleted file mode 100644 index fd9018514..000000000 --- a/surfsense_backend/tests/unit/services/test_ai_sort_task_dedupe.py +++ /dev/null @@ -1,43 +0,0 @@ -"""Unit tests for AI sort task Redis deduplication lock.""" - -from unittest.mock import MagicMock, patch - -import pytest - -pytestmark = pytest.mark.unit - - -def test_lock_key_format(): - from app.tasks.celery_tasks.document_tasks import _ai_sort_lock_key - - key = _ai_sort_lock_key(42) - assert key == "ai_sort:search_space:42:lock" - - -def test_lock_prevents_duplicate_run(): - """When the Redis lock already exists, the task should skip execution.""" - - mock_redis = MagicMock() - mock_redis.set.return_value = False # Lock already held - - with ( - patch( - "app.tasks.celery_tasks.document_tasks._get_ai_sort_redis", - return_value=mock_redis, - ), - patch( - "app.tasks.celery_tasks.document_tasks.get_celery_session_maker" - ) as mock_session_maker, - ): - import asyncio - - from app.tasks.celery_tasks.document_tasks import _ai_sort_search_space_async - - loop = asyncio.new_event_loop() - try: - loop.run_until_complete(_ai_sort_search_space_async(1, "user-123")) - finally: - loop.close() - - # Session maker should never be called since lock was not acquired - mock_session_maker.assert_not_called() diff --git a/surfsense_backend/tests/unit/services/test_folder_hierarchy.py b/surfsense_backend/tests/unit/services/test_folder_hierarchy.py deleted file mode 100644 index 9077f6b0e..000000000 --- a/surfsense_backend/tests/unit/services/test_folder_hierarchy.py +++ /dev/null @@ -1,87 +0,0 @@ -"""Unit tests for ensure_folder_hierarchy_with_depth_validation.""" - -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest - -pytestmark = pytest.mark.unit - - -@pytest.mark.asyncio -async def test_creates_missing_folders_in_chain(): - """Should create all folders when none exist.""" - from app.services.folder_service import ( - ensure_folder_hierarchy_with_depth_validation, - ) - - session = AsyncMock() - # All lookups return None (no existing folders) - mock_result = MagicMock() - mock_result.scalar_one_or_none.return_value = None - session.execute.return_value = mock_result - - folder_instances = [] - - def track_add(obj): - folder_instances.append(obj) - - session.add = track_add - - with ( - patch( - "app.services.folder_service.validate_folder_depth", new_callable=AsyncMock - ), - patch( - "app.services.folder_service.generate_folder_position", - new_callable=AsyncMock, - return_value="a0", - ), - ): - # Mock flush to assign IDs - call_count = 0 - - async def mock_flush(): - nonlocal call_count - call_count += 1 - if folder_instances: - folder_instances[-1].id = call_count - - session.flush = mock_flush - - segments = [ - {"name": "Slack", "metadata": {"ai_sort": True, "ai_sort_level": 1}}, - {"name": "2025-03-15", "metadata": {"ai_sort": True, "ai_sort_level": 2}}, - ] - - result = await ensure_folder_hierarchy_with_depth_validation( - session, 1, segments - ) - - assert len(folder_instances) == 2 - assert folder_instances[0].name == "Slack" - assert folder_instances[1].name == "2025-03-15" - assert result is folder_instances[-1] - - -@pytest.mark.asyncio -async def test_reuses_existing_folder(): - """When a folder already exists, it should be reused, not created.""" - from app.services.folder_service import ( - ensure_folder_hierarchy_with_depth_validation, - ) - - session = AsyncMock() - - existing_folder = MagicMock() - existing_folder.id = 42 - - mock_result = MagicMock() - mock_result.scalar_one_or_none.return_value = existing_folder - session.execute.return_value = mock_result - - segments = [{"name": "Existing", "metadata": None}] - - result = await ensure_folder_hierarchy_with_depth_validation(session, 1, segments) - - assert result is existing_folder - session.add.assert_not_called() diff --git a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx index 58eb58f4b..0b1369340 100644 --- a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx +++ b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx @@ -798,7 +798,7 @@ export default function NewChatPage() { }); } else { const tcId = `interrupt-${action.name}`; - addToolCall(contentPartsState, TOOLS_WITH_UI, tcId, action.name, action.args, true); + addToolCall(contentPartsState, TOOLS_WITH_UI, tcId, action.name, action.args); updateToolCall(contentPartsState, tcId, { result: { __interrupt__: true, ...interruptData }, }); @@ -1125,7 +1125,7 @@ export default function NewChatPage() { }); } else { const tcId = `interrupt-${action.name}`; - addToolCall(contentPartsState, TOOLS_WITH_UI, tcId, action.name, action.args, true); + addToolCall(contentPartsState, TOOLS_WITH_UI, tcId, action.name, action.args); updateToolCall(contentPartsState, tcId, { result: { __interrupt__: true, diff --git a/surfsense_web/app/dashboard/[search_space_id]/user-settings/components/ProfileContent.tsx b/surfsense_web/app/dashboard/[search_space_id]/user-settings/components/ProfileContent.tsx index 32377194a..c1e81283b 100644 --- a/surfsense_web/app/dashboard/[search_space_id]/user-settings/components/ProfileContent.tsx +++ b/surfsense_web/app/dashboard/[search_space_id]/user-settings/components/ProfileContent.tsx @@ -78,7 +78,7 @@ export function ProfileContent() { ) : (
-
+
diff --git a/surfsense_web/app/dashboard/error.tsx b/surfsense_web/app/dashboard/error.tsx deleted file mode 100644 index 4d872a69f..000000000 --- a/surfsense_web/app/dashboard/error.tsx +++ /dev/null @@ -1,44 +0,0 @@ -"use client"; - -import Link from "next/link"; -import { useEffect } from "react"; - -export default function DashboardError({ - error, - reset, -}: { - error: globalThis.Error & { digest?: string }; - reset: () => void; -}) { - useEffect(() => { - import("posthog-js") - .then(({ default: posthog }) => { - posthog.captureException(error); - }) - .catch(() => {}); - }, [error]); - - return ( -
-

Something went wrong

-

- An error occurred in this section. Your dashboard is still available. -

-
- - - Go to dashboard home - -
-
- ); -} diff --git a/surfsense_web/components/assistant-ui/assistant-message.tsx b/surfsense_web/components/assistant-ui/assistant-message.tsx index d0d7616d1..764acabba 100644 --- a/surfsense_web/components/assistant-ui/assistant-message.tsx +++ b/surfsense_web/components/assistant-ui/assistant-message.tsx @@ -43,7 +43,6 @@ import { useComments } from "@/hooks/use-comments"; import { useMediaQuery } from "@/hooks/use-media-query"; import { useElectronAPI } from "@/hooks/use-platform"; import { cn } from "@/lib/utils"; -import { openSafeNavigationHref, resolveSafeNavigationHref } from "@/components/tool-ui/shared/media"; // Captured once at module load — survives client-side navigations that strip the query param. const IS_QUICK_ASSIST_WINDOW = @@ -385,7 +384,6 @@ const AssistantMessageInner: FC = () => { generate_image: GenerateImageToolUI, update_memory: UpdateMemoryToolUI, execute: SandboxExecuteToolUI, - execute_code: SandboxExecuteToolUI, create_notion_page: CreateNotionPageToolUI, update_notion_page: UpdateNotionPageToolUI, delete_notion_page: DeleteNotionPageToolUI, diff --git a/surfsense_web/components/assistant-ui/connector-popup/tabs/active-connectors-tab.tsx b/surfsense_web/components/assistant-ui/connector-popup/tabs/active-connectors-tab.tsx index eb161f485..c08871445 100644 --- a/surfsense_web/components/assistant-ui/connector-popup/tabs/active-connectors-tab.tsx +++ b/surfsense_web/components/assistant-ui/connector-popup/tabs/active-connectors-tab.tsx @@ -2,7 +2,7 @@ import { Search, Unplug } from "lucide-react"; import type { FC } from "react"; -import { getDocumentTypeLabel } from "@/lib/documents/document-type-labels"; +import { getDocumentTypeLabel } from "@/components/documents/DocumentTypeIcon"; import { Button } from "@/components/ui/button"; import { Spinner } from "@/components/ui/spinner"; import { TabsContent } from "@/components/ui/tabs"; diff --git a/surfsense_web/components/assistant-ui/inline-mention-editor.tsx b/surfsense_web/components/assistant-ui/inline-mention-editor.tsx index f5de0be7d..2d55f4d20 100644 --- a/surfsense_web/components/assistant-ui/inline-mention-editor.tsx +++ b/surfsense_web/components/assistant-ui/inline-mention-editor.tsx @@ -499,14 +499,10 @@ export const InlineMentionEditor = forwardRef 0) { const range = selection.getRangeAt(0); @@ -516,41 +512,63 @@ export const InlineMentionEditor = forwardRef= 0; i--) { + if (textContent[i] === "@") { + atIndex = i; + break; + } + // Stop if we hit a space (@ must be at word boundary) if (textContent[i] === " " || textContent[i] === "\n") { - wordStart = i + 1; break; } } - let triggerChar: "@" | "/" | null = null; - let triggerIndex = -1; - for (let i = wordStart; i < cursorPos; i++) { - if (textContent[i] === "@" || textContent[i] === "/") { - triggerChar = textContent[i] as "@" | "/"; - triggerIndex = i; - break; - } - } - - if (triggerChar === "@" && triggerIndex !== -1) { - const query = textContent.slice(triggerIndex + 1, cursorPos); + if (atIndex !== -1) { + const query = textContent.slice(atIndex + 1, cursorPos); + // Only trigger if query doesn't start with space if (!query.startsWith(" ")) { shouldTriggerMention = true; mentionQuery = query; } - } else if (triggerChar === "/" && triggerIndex !== -1) { - if ( - triggerIndex === 0 || - textContent[triggerIndex - 1] === " " || - textContent[triggerIndex - 1] === "\n" - ) { - const query = textContent.slice(triggerIndex + 1, cursorPos); - if (!query.startsWith(" ")) { - shouldTriggerAction = true; - actionQuery = query; - } + } + } + } + + // Check for / actions (same pattern as @) + let shouldTriggerAction = false; + let actionQuery = ""; + + if (!shouldTriggerMention && selection && selection.rangeCount > 0) { + const range = selection.getRangeAt(0); + const textNode = range.startContainer; + + if (textNode.nodeType === Node.TEXT_NODE) { + const textContent = textNode.textContent || ""; + const cursorPos = range.startOffset; + + let slashIndex = -1; + for (let i = cursorPos - 1; i >= 0; i--) { + if (textContent[i] === "/") { + slashIndex = i; + break; + } + if (textContent[i] === " " || textContent[i] === "\n") { + break; + } + } + + if ( + slashIndex !== -1 && + (slashIndex === 0 || + textContent[slashIndex - 1] === " " || + textContent[slashIndex - 1] === "\n") + ) { + const query = textContent.slice(slashIndex + 1, cursorPos); + if (!query.startsWith(" ")) { + shouldTriggerAction = true; + actionQuery = query; } } } diff --git a/surfsense_web/components/assistant-ui/thread.tsx b/surfsense_web/components/assistant-ui/thread.tsx index dcb1e3e9e..7fb01401f 100644 --- a/surfsense_web/components/assistant-ui/thread.tsx +++ b/surfsense_web/components/assistant-ui/thread.tsx @@ -28,7 +28,8 @@ import { import { AnimatePresence, motion } from "motion/react"; import Image from "next/image"; import { useParams } from "next/navigation"; -import { type FC, useCallback, useEffect, useMemo, useRef, useState } from "react"; +import { type FC, useCallback, useEffect, useLayoutEffect, useMemo, useRef, useState } from "react"; +import { createPortal } from "react-dom"; import { agentToolsAtom, disabledToolsAtom, @@ -60,7 +61,7 @@ import { } from "@/components/assistant-ui/inline-mention-editor"; import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; import { UserMessage } from "@/components/assistant-ui/user-message"; -import { SLIDEOUT_PANEL_OPENED_EVENT } from "@/lib/layout-events"; +import { SLIDEOUT_PANEL_OPENED_EVENT } from "@/components/layout/ui/sidebar/SidebarSlideOutPanel"; import { DocumentMentionPicker, type DocumentMentionPickerRef, @@ -123,18 +124,16 @@ const ThreadContent: FC = () => { }} /> - !thread.isEmpty}> -
- - - !thread.isEmpty}> - - + !thread.isEmpty}> +
+ +
+ @@ -340,7 +339,10 @@ const Composer: FC = () => { const [showPromptPicker, setShowPromptPicker] = useState(false); const [mentionQuery, setMentionQuery] = useState(""); const [actionQuery, setActionQuery] = useState(""); + const [containerPos, setContainerPos] = useState({ bottom: "200px", left: "50%", top: "auto" }); const editorRef = useRef(null); + const editorContainerRef = useRef(null); + const composerBoxRef = useRef(null); const documentPickerRef = useRef(null); const promptPickerRef = useRef(null); const viewportRef = useRef(null); @@ -361,13 +363,38 @@ const Composer: FC = () => { viewportRef.current = document.querySelector(".aui-thread-viewport"); }, []); + // Compute picker positions using ResizeObserver to avoid layout reads during render + useLayoutEffect(() => { + if (!editorContainerRef.current) return; + + const updatePosition = () => { + if (!editorContainerRef.current) return; + const rect = editorContainerRef.current.getBoundingClientRect(); + const composerRect = composerBoxRef.current?.getBoundingClientRect(); + setContainerPos({ + bottom: `${window.innerHeight - rect.top + 8}px`, + left: `${rect.left}px`, + top: composerRect ? `${composerRect.bottom + 8}px` : "auto", + }); + }; + + updatePosition(); + const ro = new ResizeObserver(updatePosition); + ro.observe(editorContainerRef.current); + if (composerBoxRef.current) { + ro.observe(composerBoxRef.current); + } + + return () => ro.disconnect(); + }, []); + const electronAPI = useElectronAPI(); const [clipboardInitialText, setClipboardInitialText] = useState(); const clipboardLoadedRef = useRef(false); useEffect(() => { if (!electronAPI || clipboardLoadedRef.current) return; clipboardLoadedRef.current = true; - electronAPI.getQuickAskText().then((text: string) => { + electronAPI.getQuickAskText().then((text) => { if (text) { setClipboardInitialText(text); } @@ -560,15 +587,23 @@ const Composer: FC = () => { // Submit message (blocked during streaming, document picker open, or AI responding to another user) const handleSubmit = useCallback(() => { - if (isThreadRunning || isBlockedByOtherUser) return; - if (showDocumentPopover || showPromptPicker) return; - - if (clipboardInitialText) { - const userText = editorRef.current?.getText() ?? ""; - const combined = userText ? `${userText}\n\n${clipboardInitialText}` : clipboardInitialText; - aui.composer().setText(combined); - setClipboardInitialText(undefined); + if (isThreadRunning || isBlockedByOtherUser) { + return; } + if (!showDocumentPopover && !showPromptPicker) { + if (clipboardInitialText) { + const userText = editorRef.current?.getText() ?? ""; + const combined = userText ? `${userText}\n\n${clipboardInitialText}` : clipboardInitialText; + aui.composer().setText(combined); + setClipboardInitialText(undefined); + } + aui.composer().send(); + editorRef.current?.clear(); + setMentionedDocuments([]); + setSidebarDocs([]); + } + if (isThreadRunning || isBlockedByOtherUser) return; + if (showDocumentPopover) return; const viewportEl = viewportRef.current; const heightBefore = viewportEl?.scrollHeight ?? 0; @@ -582,14 +617,18 @@ const Composer: FC = () => { // assistant message so that scrolling-to-bottom actually positions the // user message at the TOP of the viewport. That slack height is // calculated asynchronously (ResizeObserver → style → layout). - // Poll via rAF for ~500ms, re-scrolling whenever scrollHeight changes. + // + // We poll via rAF for ~2 s, re-scrolling whenever scrollHeight changes + // (user msg render → assistant placeholder → ViewportSlack min-height → + // first streamed content). Backup setTimeout calls cover cases where + // the batcher's 50 ms throttle delays the DOM update past the rAF. const scrollToBottom = () => threadViewportStore.getState().scrollToBottom({ behavior: "instant" }); let lastHeight = heightBefore; let frames = 0; let cancelled = false; - const POLL_FRAMES = 30; + const POLL_FRAMES = 120; const pollAndScroll = () => { if (cancelled) return; @@ -609,11 +648,16 @@ const Composer: FC = () => { const t1 = setTimeout(scrollToBottom, 100); const t2 = setTimeout(scrollToBottom, 300); + const t3 = setTimeout(scrollToBottom, 600); + // Cleanup if component unmounts during the polling window. The ref is + // checked inside pollAndScroll; timeouts are cleared in the return below. + // Store cleanup fn so it can be called from a useEffect cleanup if needed. submitCleanupRef.current = () => { cancelled = true; clearTimeout(t1); clearTimeout(t2); + clearTimeout(t3); }; }, [ showDocumentPopover, @@ -661,54 +705,28 @@ const Composer: FC = () => { ); return ( - + - {showDocumentPopover && ( -
- { - setShowDocumentPopover(false); - setMentionQuery(""); - }} - initialSelectedDocuments={mentionedDocuments} - externalSearch={mentionQuery} - /> -
- )} - {showPromptPicker && ( -
- { - setShowPromptPicker(false); - setActionQuery(""); - }} - externalSearch={actionQuery} - /> -
- )} -
+
{clipboardInitialText && ( setClipboardInitialText(undefined)} /> )} -
+ {/* Inline editor with @mention support */} +
{ className="min-h-[24px]" />
+ {/* Document picker popover (portal to body for proper z-index stacking) */} + {showDocumentPopover && + typeof document !== "undefined" && + createPortal( + { + setShowDocumentPopover(false); + setMentionQuery(""); + }} + initialSelectedDocuments={mentionedDocuments} + externalSearch={mentionQuery} + containerStyle={{ + bottom: containerPos.bottom, + left: containerPos.left, + }} + />, + document.body + )} + {showPromptPicker && + typeof document !== "undefined" && + createPortal( + { + setShowPromptPicker(false); + setActionQuery(""); + }} + externalSearch={actionQuery} + containerStyle={{ + position: "fixed", + ...(clipboardInitialText + ? { top: containerPos.top } + : { bottom: containerPos.bottom }), + left: containerPos.left, + zIndex: 50, + }} + />, + document.body + )} diff --git a/surfsense_web/components/assistant-ui/tool-fallback.tsx b/surfsense_web/components/assistant-ui/tool-fallback.tsx index d9833b387..b658dba6d 100644 --- a/surfsense_web/components/assistant-ui/tool-fallback.tsx +++ b/surfsense_web/components/assistant-ui/tool-fallback.tsx @@ -1,16 +1,14 @@ import type { ToolCallMessagePartComponent } from "@assistant-ui/react"; import { CheckIcon, ChevronDownIcon, ChevronUpIcon, XCircleIcon } from "lucide-react"; import { useMemo, useState } from "react"; -import { GenericHitlApprovalToolUI } from "@/components/tool-ui/generic-hitl-approval"; import { getToolIcon } from "@/contracts/enums/toolIcons"; -import { isInterruptResult } from "@/lib/hitl"; import { cn } from "@/lib/utils"; function formatToolName(name: string): string { return name.replace(/_/g, " ").replace(/\b\w/g, (c) => c.toUpperCase()); } -const DefaultToolFallbackInner: ToolCallMessagePartComponent = ({ +export const ToolFallback: ToolCallMessagePartComponent = ({ toolName, argsText, result, @@ -147,10 +145,3 @@ const DefaultToolFallbackInner: ToolCallMessagePartComponent = ({
); }; - -export const ToolFallback: ToolCallMessagePartComponent = (props) => { - if (isInterruptResult(props.result)) { - return ; - } - return ; -}; diff --git a/surfsense_web/components/chat-comments/comment-item/comment-item.tsx b/surfsense_web/components/chat-comments/comment-item/comment-item.tsx index 359439d07..eb374ba49 100644 --- a/surfsense_web/components/chat-comments/comment-item/comment-item.tsx +++ b/surfsense_web/components/chat-comments/comment-item/comment-item.tsx @@ -9,7 +9,6 @@ import { Button } from "@/components/ui/button"; import { cn } from "@/lib/utils"; import { CommentComposer } from "../comment-composer/comment-composer"; import { CommentActions } from "./comment-actions"; -import { convertRenderedToDisplay } from "@/lib/comments/utils"; import type { CommentItemProps } from "./types"; function getInitials(name: string | null, email: string): string { @@ -70,6 +69,10 @@ function formatTimestamp(dateString: string): string { ); } +export function convertRenderedToDisplay(contentRendered: string): string { + // Convert @{DisplayName} format to @DisplayName for editing + return contentRendered.replace(/@\{([^}]+)\}/g, "@$1"); +} function renderMentions(content: string): React.ReactNode { // Match @{DisplayName} format from backend diff --git a/surfsense_web/components/documents/DocumentNode.tsx b/surfsense_web/components/documents/DocumentNode.tsx index 478037520..04cec5f89 100644 --- a/surfsense_web/components/documents/DocumentNode.tsx +++ b/surfsense_web/components/documents/DocumentNode.tsx @@ -82,12 +82,11 @@ export const DocumentNode = React.memo(function DocumentNode({ onContextMenuOpenChange, }: DocumentNodeProps) { const statusState = doc.status?.state ?? "ready"; - const isFailed = statusState === "failed"; - const isProcessing = statusState === "pending" || statusState === "processing"; - const isUnavailable = isProcessing || isFailed; - const isSelectable = !isUnavailable; + const isSelectable = statusState !== "pending" && statusState !== "processing"; const isEditable = - EDITABLE_DOCUMENT_TYPES.has(doc.document_type) && !isUnavailable; + EDITABLE_DOCUMENT_TYPES.has(doc.document_type) && + statusState !== "pending" && + statusState !== "processing"; const handleCheckChange = useCallback(() => { if (isSelectable) { @@ -104,6 +103,7 @@ export const DocumentNode = React.memo(function DocumentNode({ [doc.id] ); + const isProcessing = statusState === "pending" || statusState === "processing"; const [dropdownOpen, setDropdownOpen] = useState(false); const [exporting, setExporting] = useState(null); const [titleTooltipOpen, setTitleTooltipOpen] = useState(false); @@ -261,38 +261,38 @@ export const DocumentNode = React.memo(function DocumentNode({ className="w-40" onClick={(e) => e.stopPropagation()} > - onPreview(doc)} disabled={isUnavailable}> - - Open - - {isEditable && ( - onEdit(doc)}> - - Edit + onPreview(doc)} disabled={isProcessing}> + + Open - )} - onMove(doc)}> - - Move to... - - {onExport && ( - - - - Export - - - - - - )} - {onVersionHistory && isVersionableType(doc.document_type) && ( - onVersionHistory(doc)}> - - Versions + {isEditable && ( + onEdit(doc)}> + + Edit + + )} + onMove(doc)}> + + Move to... - )} - onDelete(doc)}> + {onExport && ( + + + + Export + + + + + + )} + {onVersionHistory && isVersionableType(doc.document_type) && ( + onVersionHistory(doc)}> + + Versions + + )} + onDelete(doc)}> Delete @@ -304,38 +304,38 @@ export const DocumentNode = React.memo(function DocumentNode({ {contextMenuOpen && ( e.stopPropagation()}> - onPreview(doc)} disabled={isUnavailable}> - - Open - - {isEditable && ( - onEdit(doc)}> - - Edit + onPreview(doc)} disabled={isProcessing}> + + Open - )} - onMove(doc)}> - - Move to... - - {onExport && ( - - - - Export - - - - - - )} - {onVersionHistory && isVersionableType(doc.document_type) && ( - onVersionHistory(doc)}> - - Versions + {isEditable && ( + onEdit(doc)}> + + Edit + + )} + onMove(doc)}> + + Move to... - )} - onDelete(doc)}> + {onExport && ( + + + + Export + + + + + + )} + {onVersionHistory && isVersionableType(doc.document_type) && ( + onVersionHistory(doc)}> + + Versions + + )} + onDelete(doc)}> Delete diff --git a/surfsense_web/components/documents/DocumentTypeIcon.tsx b/surfsense_web/components/documents/DocumentTypeIcon.tsx index 4c6507081..5c03d96fa 100644 --- a/surfsense_web/components/documents/DocumentTypeIcon.tsx +++ b/surfsense_web/components/documents/DocumentTypeIcon.tsx @@ -4,12 +4,52 @@ import type React from "react"; import { useEffect, useRef, useState } from "react"; import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; import { getConnectorIcon } from "@/contracts/enums/connectorIcons"; -import { getDocumentTypeLabel } from "@/lib/documents/document-type-labels"; export function getDocumentTypeIcon(type: string, className?: string): React.ReactNode { return getConnectorIcon(type, className); } +export function getDocumentTypeLabel(type: string): string { + const labelMap: Record = { + EXTENSION: "Extension", + CRAWLED_URL: "Web Page", + FILE: "File", + SLACK_CONNECTOR: "Slack", + TEAMS_CONNECTOR: "Microsoft Teams", + ONEDRIVE_FILE: "OneDrive", + DROPBOX_FILE: "Dropbox", + NOTION_CONNECTOR: "Notion", + YOUTUBE_VIDEO: "YouTube Video", + GITHUB_CONNECTOR: "GitHub", + LINEAR_CONNECTOR: "Linear", + DISCORD_CONNECTOR: "Discord", + JIRA_CONNECTOR: "Jira", + CONFLUENCE_CONNECTOR: "Confluence", + CLICKUP_CONNECTOR: "ClickUp", + GOOGLE_CALENDAR_CONNECTOR: "Google Calendar", + GOOGLE_GMAIL_CONNECTOR: "Gmail", + GOOGLE_DRIVE_FILE: "Google Drive", + AIRTABLE_CONNECTOR: "Airtable", + LUMA_CONNECTOR: "Luma", + ELASTICSEARCH_CONNECTOR: "Elasticsearch", + BOOKSTACK_CONNECTOR: "BookStack", + CIRCLEBACK: "Circleback", + OBSIDIAN_CONNECTOR: "Obsidian", + LOCAL_FOLDER_FILE: "Local Folder", + SURFSENSE_DOCS: "SurfSense Docs", + NOTE: "Note", + COMPOSIO_GOOGLE_DRIVE_CONNECTOR: "Composio Google Drive", + COMPOSIO_GMAIL_CONNECTOR: "Composio Gmail", + COMPOSIO_GOOGLE_CALENDAR_CONNECTOR: "Composio Google Calendar", + }; + return ( + labelMap[type] || + type + .split("_") + .map((word) => word.charAt(0) + word.slice(1).toLowerCase()) + .join(" ") + ); +} export function DocumentTypeChip({ type, className }: { type: string; className?: string }) { const icon = getDocumentTypeIcon(type, "h-4 w-4"); diff --git a/surfsense_web/components/documents/DocumentsFilters.tsx b/surfsense_web/components/documents/DocumentsFilters.tsx index 2b7cf0f10..703c9c3b4 100644 --- a/surfsense_web/components/documents/DocumentsFilters.tsx +++ b/surfsense_web/components/documents/DocumentsFilters.tsx @@ -1,8 +1,6 @@ "use client"; -import { IconBinaryTree, IconBinaryTreeFilled } from "@tabler/icons-react"; -import { FolderPlus, ListFilter, Search, Upload, X } from "lucide-react"; -import { AnimatePresence, motion } from "motion/react"; +import { Download, FolderPlus, ListFilter, Loader2, Search, Upload, X } from "lucide-react"; import { useTranslations } from "next-intl"; import React, { useCallback, useMemo, useRef, useState } from "react"; import { useDocumentUploadDialog } from "@/components/assistant-ui/document-upload-popup"; @@ -12,10 +10,8 @@ import { Input } from "@/components/ui/input"; import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover"; import { ToggleGroup, ToggleGroupItem } from "@/components/ui/toggle-group"; import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; -import { cn } from "@/lib/utils"; import type { DocumentTypeEnum } from "@/contracts/types/document.types"; -import { getDocumentTypeIcon } from "./DocumentTypeIcon"; -import { getDocumentTypeLabel } from "@/lib/documents/document-type-labels"; +import { getDocumentTypeIcon, getDocumentTypeLabel } from "./DocumentTypeIcon"; export function DocumentsFilters({ typeCounts: typeCountsRecord, @@ -24,9 +20,8 @@ export function DocumentsFilters({ onToggleType, activeTypes, onCreateFolder, - aiSortEnabled = false, - aiSortBusy = false, - onToggleAiSort, + onExportKB, + isExporting, }: { typeCounts: Partial>; onSearch: (v: string) => void; @@ -34,9 +29,8 @@ export function DocumentsFilters({ onToggleType: (type: DocumentTypeEnum, checked: boolean) => void; activeTypes: DocumentTypeEnum[]; onCreateFolder?: () => void; - aiSortEnabled?: boolean; - aiSortBusy?: boolean; - onToggleAiSort?: () => void; + onExportKB?: () => void; + isExporting?: boolean; }) { const t = useTranslations("documents"); const id = React.useId(); @@ -74,7 +68,7 @@ export function DocumentsFilters({ return (
- {/* New Folder + Filter Toggle Group */} + {/* Filter + New Folder Toggle Group */} {onCreateFolder && ( @@ -91,8 +85,33 @@ export function DocumentsFilters({ New folder - - )} + + )} + + {onExportKB && ( + + + { + e.preventDefault(); + onExportKB(); + }} + > + {isExporting ? ( + + ) : ( + + )} + + + + {isExporting ? "Exporting…" : "Export knowledge base"} + + + )} @@ -182,70 +201,6 @@ export function DocumentsFilters({ - {/* AI Sort Toggle */} - {onToggleAiSort && ( - - - - - - {aiSortBusy - ? "AI sort in progress..." - : aiSortEnabled - ? "AI sort active — click to disable" - : "Enable AI sort"} - - - )} - {/* Search Input */}
diff --git a/surfsense_web/components/documents/FolderTreeView.tsx b/surfsense_web/components/documents/FolderTreeView.tsx index ca833949c..4988e87e7 100644 --- a/surfsense_web/components/documents/FolderTreeView.tsx +++ b/surfsense_web/components/documents/FolderTreeView.tsx @@ -90,8 +90,6 @@ export function FolderTreeView({ const [openContextMenuId, setOpenContextMenuId] = useState(null); - const [manuallyCollapsedAiIds, setManuallyCollapsedAiIds] = useState>(new Set()); - // Single subscription for rename state — derived boolean passed to each FolderNode const [renamingFolderId, setRenamingFolderId] = useAtom(renamingFolderIdAtom); const handleStartRename = useCallback( @@ -100,38 +98,6 @@ export function FolderTreeView({ ); const handleCancelRename = useCallback(() => setRenamingFolderId(null), [setRenamingFolderId]); - const aiSortFolderLevels = useMemo(() => { - const map = new Map(); - for (const f of folders) { - if (f.metadata?.ai_sort === true && typeof f.metadata?.ai_sort_level === "number") { - map.set(f.id, f.metadata.ai_sort_level as number); - } - } - return map; - }, [folders]); - - const handleToggleExpand = useCallback( - (folderId: number) => { - const aiLevel = aiSortFolderLevels.get(folderId); - if (aiLevel !== undefined && aiLevel < 4) { - // AI-auto-expanded folder: only toggle the manual-collapse set. - // Calling onToggleExpand would add it to expandedIds and fight auto-expand. - setManuallyCollapsedAiIds((prev) => { - const next = new Set(prev); - if (next.has(folderId)) { - next.delete(folderId); - } else { - next.add(folderId); - } - return next; - }); - return; - } - onToggleExpand(folderId); - }, - [onToggleExpand, aiSortFolderLevels] - ); - const effectiveActiveTypes = useMemo(() => { if ( activeTypes.includes("FILE" as DocumentTypeEnum) && @@ -246,16 +212,9 @@ export function FolderTreeView({ function renderLevel(parentId: number | null, depth: number): React.ReactNode[] { const key = parentId ?? "root"; - const childFolders = (foldersByParent[key] ?? []).slice().sort((a, b) => { - const aIsDate = - a.metadata?.ai_sort === true && a.metadata?.ai_sort_level === 2; - const bIsDate = - b.metadata?.ai_sort === true && b.metadata?.ai_sort_level === 2; - if (aIsDate && bIsDate) { - return b.name.localeCompare(a.name); - } - return a.position.localeCompare(b.position); - }); + const childFolders = (foldersByParent[key] ?? []) + .slice() + .sort((a, b) => a.position.localeCompare(b.position)); const visibleFolders = hasDescendantMatch ? childFolders.filter((f) => hasDescendantMatch[f.id]) : childFolders; @@ -267,32 +226,6 @@ export function FolderTreeView({ const nodes: React.ReactNode[] = []; - if (parentId === null) { - const processingDocs = childDocs.filter((d) => { - const state = d.status?.state; - return state === "pending" || state === "processing"; - }); - for (const d of processingDocs) { - nodes.push( - setOpenContextMenuId(open ? `doc-${d.id}` : null)} - /> - ); - } - } - for (let i = 0; i < visibleFolders.length; i++) { const f = visibleFolders[i]; const siblingPositions = { @@ -300,15 +233,8 @@ export function FolderTreeView({ after: i < visibleFolders.length - 1 ? visibleFolders[i + 1].position : null, }; - const isSearchAutoExpanded = !!searchQuery && !!hasDescendantMatch?.[f.id]; - const isAiAutoExpandCandidate = - f.metadata?.ai_sort === true && - typeof f.metadata?.ai_sort_level === "number" && - (f.metadata.ai_sort_level as number) < 4; - const isManuallyCollapsed = manuallyCollapsedAiIds.has(f.id); - const isExpanded = isManuallyCollapsed - ? isSearchAutoExpanded - : expandedIds.has(f.id) || isSearchAutoExpanded || isAiAutoExpandCandidate; + const isAutoExpanded = !!searchQuery && !!hasDescendantMatch?.[f.id]; + const isExpanded = expandedIds.has(f.id) || isAutoExpanded; nodes.push( { - const state = d.status?.state; - return state !== "pending" && state !== "processing"; - }) - : childDocs; - - for (const d of remainingDocs) { + for (const d of childDocs) { nodes.push( ). Arbitrary markdown from document conversions -// (e.g. PDF-to-markdown via Azure/DocIntel) can contain constructs that -// break the MDX parser. This module sanitises them before deserialization. +// remarkMdx treats { } as JSX expression delimiters. Arbitrary markdown +// (e.g. AI-generated reports) can contain curly braces that are NOT valid JS +// expressions, which makes acorn throw "Could not parse expression". +// We escape unescaped { and } *outside* of fenced code blocks and inline code +// so remarkMdx treats them as literal characters while still parsing +// , , , etc. tags correctly. // --------------------------------------------------------------------------- const FENCED_OR_INLINE_CODE = /(```[\s\S]*?```|`[^`\n]+`)/g; -// Strip HTML comments that MDX cannot parse. -// PDF converters emit , , etc. -// MDX uses JSX-style comments and chokes on HTML comments, causing the -// parser to stop at the first occurrence. -// - becomes a thematic break (---) -// - All other HTML comments are removed -function stripHtmlComments(md: string): string { - return md - .replace(//gi, "\n---\n") - .replace(//g, ""); -} - -// Convert
...
blocks to plain text blockquotes. -//
with arbitrary text content is not valid JSX, causing the MDX -// parser to fail. -function convertFigureBlocks(md: string): string { - return md.replace(/]*>([\s\S]*?)<\/figure>/gi, (_match, inner: string) => { - const trimmed = (inner as string).trim(); - if (!trimmed) return ""; - const quoted = trimmed - .split("\n") - .map((line) => `> ${line}`) - .join("\n"); - return `\n${quoted}\n`; - }); -} - -// Escape unescaped { and } outside of fenced/inline code so remarkMdx -// treats them as literal characters rather than JSX expression delimiters. -function escapeCurlyBraces(md: string): string { +export function escapeMdxExpressions(md: string): string { const parts = md.split(FENCED_OR_INLINE_CODE); return parts .map((part, i) => { + // Odd indices are code blocks / inline code – leave untouched if (i % 2 === 1) return part; + // Escape { and } that are NOT already escaped (no preceding \) return part.replace(/(? k + 1); - // Sync Next.js router internals so useParams/usePathname stay correct going forward - router.replace(`/dashboard/${searchSpaceId}/new-chat`); + router.replace(`/dashboard/${searchSpaceId}/new-chat/${currentThreadState.id}`); + // Allow router to sync, then navigate to fresh new-chat + setTimeout(() => { + router.push(`/dashboard/${searchSpaceId}/new-chat`); + }, 0); } else { + // Normal navigation - router is in sync router.push(`/dashboard/${searchSpaceId}/new-chat`); } }, [router, searchSpaceId, currentThreadState.id, params?.chat_id, resetCurrentThread]); diff --git a/surfsense_web/components/layout/ui/sidebar/DocumentsSidebar.tsx b/surfsense_web/components/layout/ui/sidebar/DocumentsSidebar.tsx index ee0826595..20b25a2d2 100644 --- a/surfsense_web/components/layout/ui/sidebar/DocumentsSidebar.tsx +++ b/surfsense_web/components/layout/ui/sidebar/DocumentsSidebar.tsx @@ -15,7 +15,6 @@ import { expandedFolderIdsAtom } from "@/atoms/documents/folder.atoms"; import { agentCreatedDocumentsAtom } from "@/atoms/documents/ui.atoms"; import { openEditorPanelAtom } from "@/atoms/editor/editor-panel.atom"; import { rightPanelCollapsedAtom } from "@/atoms/layout/right-panel.atom"; -import { searchSpacesAtom } from "@/atoms/search-spaces/search-space-query.atoms"; import { CreateFolderDialog } from "@/components/documents/CreateFolderDialog"; import type { DocumentNodeDoc } from "@/components/documents/DocumentNode"; import { DocumentsFilters } from "@/components/documents/DocumentsFilters"; @@ -50,7 +49,6 @@ import { useMediaQuery } from "@/hooks/use-media-query"; import { useElectronAPI } from "@/hooks/use-platform"; import { documentsApiService } from "@/lib/apis/documents-api.service"; import { foldersApiService } from "@/lib/apis/folders-api.service"; -import { searchSpacesApiService } from "@/lib/apis/search-spaces-api.service"; import { authenticatedFetch } from "@/lib/auth-utils"; import { uploadFolderScan } from "@/lib/folder-sync-upload"; import { getSupportedExtensionsSet } from "@/lib/supported-extensions"; @@ -110,47 +108,6 @@ export function DocumentsSidebar({ const [watchInitialFolder, setWatchInitialFolder] = useState(null); const isElectron = typeof window !== "undefined" && !!window.electronAPI; - // AI File Sort state - const { data: searchSpaces, refetch: refetchSearchSpaces } = useAtomValue(searchSpacesAtom); - const activeSearchSpace = useMemo( - () => searchSpaces?.find((s) => s.id === searchSpaceId), - [searchSpaces, searchSpaceId] - ); - const aiSortEnabled = activeSearchSpace?.ai_file_sort_enabled ?? false; - const [aiSortBusy, setAiSortBusy] = useState(false); - const [aiSortConfirmOpen, setAiSortConfirmOpen] = useState(false); - - const handleToggleAiSort = useCallback(() => { - if (aiSortEnabled) { - // Disable: just update the setting, no confirmation needed - setAiSortBusy(true); - searchSpacesApiService - .updateSearchSpace({ id: searchSpaceId, data: { ai_file_sort_enabled: false } }) - .then(() => { - refetchSearchSpaces(); - toast.success("AI file sorting disabled"); - }) - .catch(() => toast.error("Failed to disable AI file sorting")) - .finally(() => setAiSortBusy(false)); - } else { - setAiSortConfirmOpen(true); - } - }, [aiSortEnabled, searchSpaceId, refetchSearchSpaces]); - - const handleConfirmEnableAiSort = useCallback(() => { - setAiSortConfirmOpen(false); - setAiSortBusy(true); - searchSpacesApiService - .updateSearchSpace({ id: searchSpaceId, data: { ai_file_sort_enabled: true } }) - .then(() => searchSpacesApiService.triggerAiSort(searchSpaceId)) - .then(() => { - refetchSearchSpaces(); - toast.success("AI file sorting enabled — organizing your documents in the background"); - }) - .catch(() => toast.error("Failed to enable AI file sorting")) - .finally(() => setAiSortBusy(false)); - }, [searchSpaceId, refetchSearchSpaces]); - const handleWatchLocalFolder = useCallback(async () => { const api = window.electronAPI; if (!api?.selectFolder) return; @@ -449,13 +406,22 @@ export function DocumentsSidebar({ setFolderPickerOpen(true); }, []); - const [, setIsExportingKB] = useState(false); + const [isExportingKB, setIsExportingKB] = useState(false); const [exportWarningOpen, setExportWarningOpen] = useState(false); const [exportWarningContext, setExportWarningContext] = useState<{ - folder: FolderDisplay; + type: "kb" | "folder"; + folder?: FolderDisplay; pendingCount: number; } | null>(null); + const pendingDocuments = useMemo( + () => + treeDocuments.filter( + (d) => d.status?.state === "pending" || d.status?.state === "processing" + ), + [treeDocuments] + ); + const doExport = useCallback(async (url: string, downloadName: string) => { const response = await authenticatedFetch(url, { method: "GET" }); if (!response.ok) { @@ -474,29 +440,69 @@ export function DocumentsSidebar({ URL.revokeObjectURL(blobUrl); }, []); - const handleExportWarningConfirm = useCallback(async () => { - setExportWarningOpen(false); - const ctx = exportWarningContext; - if (!ctx?.folder) return; + const handleExportKB = useCallback(async () => { + if (isExportingKB) return; + + if (pendingDocuments.length > 0) { + setExportWarningContext({ type: "kb", pendingCount: pendingDocuments.length }); + setExportWarningOpen(true); + return; + } setIsExportingKB(true); try { - const safeName = - ctx.folder.name - .replace(/[^a-zA-Z0-9 _-]/g, "_") - .trim() - .slice(0, 80) || "folder"; await doExport( - `${process.env.NEXT_PUBLIC_FASTAPI_BACKEND_URL}/api/v1/search-spaces/${searchSpaceId}/export?folder_id=${ctx.folder.id}`, - `${safeName}.zip` + `${process.env.NEXT_PUBLIC_FASTAPI_BACKEND_URL}/api/v1/search-spaces/${searchSpaceId}/export`, + "knowledge-base.zip" ); - toast.success(`Folder "${ctx.folder.name}" exported`); + toast.success("Knowledge base exported"); } catch (err) { - console.error("Folder export failed:", err); + console.error("KB export failed:", err); toast.error(err instanceof Error ? err.message : "Export failed"); } finally { setIsExportingKB(false); } + }, [searchSpaceId, isExportingKB, pendingDocuments.length, doExport]); + + const handleExportWarningConfirm = useCallback(async () => { + setExportWarningOpen(false); + const ctx = exportWarningContext; + if (!ctx) return; + + if (ctx.type === "kb") { + setIsExportingKB(true); + try { + await doExport( + `${process.env.NEXT_PUBLIC_FASTAPI_BACKEND_URL}/api/v1/search-spaces/${searchSpaceId}/export`, + "knowledge-base.zip" + ); + toast.success("Knowledge base exported"); + } catch (err) { + console.error("KB export failed:", err); + toast.error(err instanceof Error ? err.message : "Export failed"); + } finally { + setIsExportingKB(false); + } + } else if (ctx.type === "folder" && ctx.folder) { + setIsExportingKB(true); + try { + const safeName = + ctx.folder.name + .replace(/[^a-zA-Z0-9 _-]/g, "_") + .trim() + .slice(0, 80) || "folder"; + await doExport( + `${process.env.NEXT_PUBLIC_FASTAPI_BACKEND_URL}/api/v1/search-spaces/${searchSpaceId}/export?folder_id=${ctx.folder.id}`, + `${safeName}.zip` + ); + toast.success(`Folder "${ctx.folder.name}" exported`); + } catch (err) { + console.error("Folder export failed:", err); + toast.error(err instanceof Error ? err.message : "Export failed"); + } finally { + setIsExportingKB(false); + } + } setExportWarningContext(null); }, [exportWarningContext, searchSpaceId, doExport]); @@ -524,6 +530,7 @@ export function DocumentsSidebar({ const folderPendingCount = getPendingCountInSubtree(folder.id); if (folderPendingCount > 0) { setExportWarningContext({ + type: "folder", folder, pendingCount: folderPendingCount, }); @@ -670,10 +677,9 @@ export function DocumentsSidebar({ function collectSubtreeDocs(parentId: number): DocumentNodeDoc[] { const directDocs = (treeDocuments ?? []).filter( (d) => - d.folderId === parentId && - d.status?.state !== "pending" && - d.status?.state !== "processing" && - d.status?.state !== "failed" + d.folderId === parentId && + d.status?.state !== "pending" && + d.status?.state !== "processing" ); const childFolders = foldersByParent[String(parentId)] ?? []; const descendantDocs = childFolders.flatMap((cf) => collectSubtreeDocs(cf.id)); @@ -948,9 +954,8 @@ export function DocumentsSidebar({ onToggleType={onToggleType} activeTypes={activeTypes} onCreateFolder={() => handleCreateFolder(null)} - aiSortEnabled={aiSortEnabled} - aiSortBusy={aiSortBusy} - onToggleAiSort={handleToggleAiSort} + onExportKB={handleExportKB} + isExporting={isExportingKB} />
@@ -1112,25 +1117,6 @@ export function DocumentsSidebar({ - - - - - Enable AI File Sorting? - - All documents in this search space will be organized into folders by - connector type, date, and AI-generated categories. New documents will - also be sorted automatically. You can disable this at any time. - - - - Cancel - - Enable - - - - ); diff --git a/surfsense_web/components/layout/ui/sidebar/InboxSidebar.tsx b/surfsense_web/components/layout/ui/sidebar/InboxSidebar.tsx index 5d8b530c0..371f3dc6d 100644 --- a/surfsense_web/components/layout/ui/sidebar/InboxSidebar.tsx +++ b/surfsense_web/components/layout/ui/sidebar/InboxSidebar.tsx @@ -22,8 +22,8 @@ import { useParams, useRouter } from "next/navigation"; import { useTranslations } from "next-intl"; import { useCallback, useDeferredValue, useEffect, useMemo, useRef, useState } from "react"; import { setTargetCommentIdAtom } from "@/atoms/chat/current-thread.atom"; -import { convertRenderedToDisplay } from "@/lib/comments/utils"; -import { getDocumentTypeLabel } from "@/lib/documents/document-type-labels"; +import { convertRenderedToDisplay } from "@/components/chat-comments/comment-item/comment-item"; +import { getDocumentTypeLabel } from "@/components/documents/DocumentTypeIcon"; import { Tabs, TabsList, TabsTrigger } from "@/components/ui/animated-tabs"; import { Avatar, AvatarFallback, AvatarImage } from "@/components/ui/avatar"; import { Button } from "@/components/ui/button"; diff --git a/surfsense_web/components/layout/ui/sidebar/SidebarSlideOutPanel.tsx b/surfsense_web/components/layout/ui/sidebar/SidebarSlideOutPanel.tsx index 661f76ed3..5195082cd 100644 --- a/surfsense_web/components/layout/ui/sidebar/SidebarSlideOutPanel.tsx +++ b/surfsense_web/components/layout/ui/sidebar/SidebarSlideOutPanel.tsx @@ -3,8 +3,8 @@ import { AnimatePresence, motion } from "motion/react"; import { useCallback, useEffect } from "react"; import { useMediaQuery } from "@/hooks/use-media-query"; -import { SLIDEOUT_PANEL_OPENED_EVENT } from "@/lib/layout-events"; +export const SLIDEOUT_PANEL_OPENED_EVENT = "slideout-panel-opened"; interface SidebarSlideOutPanelProps { open: boolean; diff --git a/surfsense_web/components/new-chat/chat-header.tsx b/surfsense_web/components/new-chat/chat-header.tsx index 4716418ee..0c5253c6c 100644 --- a/surfsense_web/components/new-chat/chat-header.tsx +++ b/surfsense_web/components/new-chat/chat-header.tsx @@ -44,28 +44,21 @@ export function ChatHeader({ searchSpaceId, className }: ChatHeaderProps) { const [isVisionGlobal, setIsVisionGlobal] = useState(false); const [visionDialogMode, setVisionDialogMode] = useState<"create" | "edit" | "view">("view"); - // Default provider for create dialogs - const [defaultLLMProvider, setDefaultLLMProvider] = useState(); - const [defaultImageProvider, setDefaultImageProvider] = useState(); - const [defaultVisionProvider, setDefaultVisionProvider] = useState(); - // LLM handlers const handleEditLLMConfig = useCallback( (config: NewLLMConfigPublic | GlobalNewLLMConfig, global: boolean) => { setSelectedConfig(config); setIsGlobal(global); setDialogMode(global ? "view" : "edit"); - setDefaultLLMProvider(undefined); setDialogOpen(true); }, [] ); - const handleAddNewLLM = useCallback((provider?: string) => { + const handleAddNewLLM = useCallback(() => { setSelectedConfig(null); setIsGlobal(false); setDialogMode("create"); - setDefaultLLMProvider(provider); setDialogOpen(true); }, []); @@ -75,11 +68,10 @@ export function ChatHeader({ searchSpaceId, className }: ChatHeaderProps) { }, []); // Image model handlers - const handleAddImageModel = useCallback((provider?: string) => { + const handleAddImageModel = useCallback(() => { setSelectedImageConfig(null); setIsImageGlobal(false); setImageDialogMode("create"); - setDefaultImageProvider(provider); setImageDialogOpen(true); }, []); @@ -88,7 +80,6 @@ export function ChatHeader({ searchSpaceId, className }: ChatHeaderProps) { setSelectedImageConfig(config); setIsImageGlobal(global); setImageDialogMode(global ? "view" : "edit"); - setDefaultImageProvider(undefined); setImageDialogOpen(true); }, [] @@ -100,11 +91,10 @@ export function ChatHeader({ searchSpaceId, className }: ChatHeaderProps) { }, []); // Vision model handlers - const handleAddVisionModel = useCallback((provider?: string) => { + const handleAddVisionModel = useCallback(() => { setSelectedVisionConfig(null); setIsVisionGlobal(false); setVisionDialogMode("create"); - setDefaultVisionProvider(provider); setVisionDialogOpen(true); }, []); @@ -113,7 +103,6 @@ export function ChatHeader({ searchSpaceId, className }: ChatHeaderProps) { setSelectedVisionConfig(config); setIsVisionGlobal(global); setVisionDialogMode(global ? "view" : "edit"); - setDefaultVisionProvider(undefined); setVisionDialogOpen(true); }, [] @@ -142,7 +131,6 @@ export function ChatHeader({ searchSpaceId, className }: ChatHeaderProps) { isGlobal={isGlobal} searchSpaceId={searchSpaceId} mode={dialogMode} - defaultProvider={defaultLLMProvider} />
); diff --git a/surfsense_web/components/new-chat/document-mention-picker.tsx b/surfsense_web/components/new-chat/document-mention-picker.tsx index f2985278d..9c6521f31 100644 --- a/surfsense_web/components/new-chat/document-mention-picker.tsx +++ b/surfsense_web/components/new-chat/document-mention-picker.tsx @@ -29,6 +29,8 @@ interface DocumentMentionPickerProps { onDone: () => void; initialSelectedDocuments?: Pick[]; externalSearch?: string; + /** Positioning styles for the container */ + containerStyle?: React.CSSProperties; } const PAGE_SIZE = 20; @@ -73,6 +75,7 @@ export const DocumentMentionPicker = forwardRef< onDone, initialSelectedDocuments = [], externalSearch = "", + containerStyle, }, ref ) { @@ -391,9 +394,19 @@ export const DocumentMentionPicker = forwardRef< [selectableDocuments, highlightedIndex, handleSelectDocument, onDone] ); + // Hide popup when there are no documents to display (regardless of fetch state) + // Search continues in background; popup reappears when results arrive + if (!actualLoading && actualDocuments.length === 0) { + return null; + } + return (
)}
- ) : ( -
-

No matching documents

-
- )} + ) : null}
); diff --git a/surfsense_web/components/new-chat/model-selector.tsx b/surfsense_web/components/new-chat/model-selector.tsx index 26937e18b..ec5bf6760 100644 --- a/surfsense_web/components/new-chat/model-selector.tsx +++ b/surfsense_web/components/new-chat/model-selector.tsx @@ -1,20 +1,8 @@ "use client"; -import type React from "react"; import { useAtomValue } from "jotai"; -import { - Bot, - Check, - ChevronDown, - Edit3, - Eye, - ImageIcon, - Layers, - Plus, - Search, - Zap, -} from "lucide-react"; -import { Fragment, useCallback, useEffect, useMemo, useRef, useState } from "react"; +import { Bot, Check, ChevronDown, Edit3, Eye, ImageIcon, Plus, Search, Zap } from "lucide-react"; +import { type UIEvent, useCallback, useMemo, useState } from "react"; import { toast } from "sonner"; import { globalImageGenConfigsAtom, @@ -34,16 +22,17 @@ import { import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; import { - Drawer, - DrawerContent, - DrawerHandle, - DrawerHeader, - DrawerTitle, - DrawerTrigger, -} from "@/components/ui/drawer"; + Command, + CommandEmpty, + CommandGroup, + CommandInput, + CommandItem, + CommandList, + CommandSeparator, +} from "@/components/ui/command"; import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover"; import { Spinner } from "@/components/ui/spinner"; -import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; +import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import type { GlobalImageGenConfig, GlobalNewLLMConfig, @@ -52,209 +41,16 @@ import type { NewLLMConfigPublic, VisionLLMConfig, } from "@/contracts/types/new-llm-config.types"; -import { useIsMobile } from "@/hooks/use-mobile"; import { getProviderIcon } from "@/lib/provider-icons"; import { cn } from "@/lib/utils"; -// ─── Helpers ──────────────────────────────────────────────────────── - -const PROVIDER_NAMES: Record = { - OPENAI: "OpenAI", - ANTHROPIC: "Anthropic", - GOOGLE: "Google", - AZURE: "Azure", - AZURE_OPENAI: "Azure OpenAI", - AWS_BEDROCK: "AWS Bedrock", - BEDROCK: "Bedrock", - DEEPSEEK: "DeepSeek", - MISTRAL: "Mistral", - COHERE: "Cohere", - GROQ: "Groq", - OLLAMA: "Ollama", - TOGETHER_AI: "Together AI", - FIREWORKS_AI: "Fireworks AI", - REPLICATE: "Replicate", - HUGGINGFACE: "HuggingFace", - PERPLEXITY: "Perplexity", - XAI: "xAI", - OPENROUTER: "OpenRouter", - CEREBRAS: "Cerebras", - SAMBANOVA: "SambaNova", - VERTEX_AI: "Vertex AI", - MINIMAX: "MiniMax", - MOONSHOT: "Moonshot", - ZHIPU: "Zhipu", - DEEPINFRA: "DeepInfra", - CLOUDFLARE: "Cloudflare", - DATABRICKS: "Databricks", - NSCALE: "NScale", - RECRAFT: "Recraft", - XINFERENCE: "XInference", - CUSTOM: "Custom", - AI21: "AI21", - ALIBABA_QWEN: "Qwen", - ANYSCALE: "Anyscale", - COMETAPI: "CometAPI", -}; - -// Provider keys valid per model type, matching backend enums -// (LiteLLMProvider, ImageGenProvider, VisionProvider in db.py) -const LLM_PROVIDER_KEYS: string[] = [ - "OPENAI", - "ANTHROPIC", - "GOOGLE", - "AZURE_OPENAI", - "BEDROCK", - "VERTEX_AI", - "GROQ", - "DEEPSEEK", - "XAI", - "MISTRAL", - "COHERE", - "OPENROUTER", - "TOGETHER_AI", - "FIREWORKS_AI", - "REPLICATE", - "PERPLEXITY", - "OLLAMA", - "CEREBRAS", - "SAMBANOVA", - "DEEPINFRA", - "AI21", - "ALIBABA_QWEN", - "MOONSHOT", - "ZHIPU", - "MINIMAX", - "HUGGINGFACE", - "CLOUDFLARE", - "DATABRICKS", - "ANYSCALE", - "COMETAPI", - "GITHUB_MODELS", - "CUSTOM", -]; - -const IMAGE_PROVIDER_KEYS: string[] = [ - "OPENAI", - "AZURE_OPENAI", - "GOOGLE", - "VERTEX_AI", - "BEDROCK", - "RECRAFT", - "OPENROUTER", - "XINFERENCE", - "NSCALE", -]; - -const VISION_PROVIDER_KEYS: string[] = [ - "OPENAI", - "ANTHROPIC", - "GOOGLE", - "AZURE_OPENAI", - "VERTEX_AI", - "BEDROCK", - "XAI", - "OPENROUTER", - "OLLAMA", - "GROQ", - "TOGETHER_AI", - "FIREWORKS_AI", - "DEEPSEEK", - "MISTRAL", - "CUSTOM", -]; - -const PROVIDER_KEYS_BY_TAB: Record = { - llm: LLM_PROVIDER_KEYS, - image: IMAGE_PROVIDER_KEYS, - vision: VISION_PROVIDER_KEYS, -}; - -function formatProviderName(provider: string): string { - const key = provider.toUpperCase(); - return ( - PROVIDER_NAMES[key] ?? - provider.charAt(0).toUpperCase() + - provider.slice(1).toLowerCase().replace(/_/g, " ") - ); -} - -function normalizeText(input: string): string { - return input - .normalize("NFD") - .replace(/\p{Diacritic}/gu, "") - .toLowerCase() - .replace(/[^a-z0-9]+/g, " ") - .trim(); -} - -interface ConfigBase { - id: number; - name: string; - model_name: string; - provider: string; -} - -function filterAndScore( - configs: T[], - selectedProvider: string, - searchQuery: string, -): T[] { - let result = configs; - - if (selectedProvider !== "all") { - result = result.filter( - (c) => c.provider.toUpperCase() === selectedProvider, - ); - } - - if (!searchQuery.trim()) return result; - - const normalized = normalizeText(searchQuery); - const tokens = normalized.split(/\s+/).filter(Boolean); - - const scored = result.map((c) => { - const aggregate = normalizeText( - [c.name, c.model_name, c.provider].join(" "), - ); - let score = 0; - if (aggregate.includes(normalized)) score += 5; - for (const token of tokens) { - if (aggregate.includes(token)) score += 1; - } - return { config: c, score }; - }); - - return scored - .filter((s) => s.score > 0) - .sort((a, b) => b.score - a.score) - .map((s) => s.config); -} - -interface DisplayItem { - config: ConfigBase & Record; - isGlobal: boolean; - isAutoMode: boolean; -} - -// ─── Component ────────────────────────────────────────────────────── - interface ModelSelectorProps { - onEditLLM: ( - config: NewLLMConfigPublic | GlobalNewLLMConfig, - isGlobal: boolean, - ) => void; - onAddNewLLM: (provider?: string) => void; - onEditImage?: ( - config: ImageGenerationConfig | GlobalImageGenConfig, - isGlobal: boolean, - ) => void; - onAddNewImage?: (provider?: string) => void; - onEditVision?: ( - config: VisionLLMConfig | GlobalVisionLLMConfig, - isGlobal: boolean, - ) => void; - onAddNewVision?: (provider?: string) => void; + onEditLLM: (config: NewLLMConfigPublic | GlobalNewLLMConfig, isGlobal: boolean) => void; + onAddNewLLM: () => void; + onEditImage?: (config: ImageGenerationConfig | GlobalImageGenConfig, isGlobal: boolean) => void; + onAddNewImage?: () => void; + onEditVision?: (config: VisionLLMConfig | GlobalVisionLLMConfig, isGlobal: boolean) => void; + onAddNewVision?: () => void; className?: string; } @@ -268,69 +64,40 @@ export function ModelSelector({ className, }: ModelSelectorProps) { const [open, setOpen] = useState(false); - const [activeTab, setActiveTab] = useState<"llm" | "image" | "vision">( - "llm", + const [activeTab, setActiveTab] = useState<"llm" | "image" | "vision">("llm"); + const [llmSearchQuery, setLlmSearchQuery] = useState(""); + const [imageSearchQuery, setImageSearchQuery] = useState(""); + const [visionSearchQuery, setVisionSearchQuery] = useState(""); + const [llmScrollPos, setLlmScrollPos] = useState<"top" | "middle" | "bottom">("top"); + const [imageScrollPos, setImageScrollPos] = useState<"top" | "middle" | "bottom">("top"); + const [visionScrollPos, setVisionScrollPos] = useState<"top" | "middle" | "bottom">("top"); + const handleListScroll = useCallback( + (setter: typeof setLlmScrollPos) => (e: UIEvent) => { + const el = e.currentTarget; + const atTop = el.scrollTop <= 2; + const atBottom = el.scrollHeight - el.scrollTop - el.clientHeight <= 2; + setter(atTop ? "top" : atBottom ? "bottom" : "middle"); + }, + [] ); - const [searchQuery, setSearchQuery] = useState(""); - const [selectedProvider, setSelectedProvider] = useState("all"); - const [focusedIndex, setFocusedIndex] = useState(-1); - const [showScrollIndicator, setShowScrollIndicator] = useState(true); - const providerSidebarRef = useRef(null); - const modelListRef = useRef(null); - const searchInputRef = useRef(null); - const isMobile = useIsMobile(); - // Reset search + provider when tab changes - useEffect(() => { - setSelectedProvider("all"); - setSearchQuery(""); - setFocusedIndex(-1); - }, [activeTab]); - - // Reset on open - useEffect(() => { - if (open) { - setSearchQuery(""); - setSelectedProvider("all"); - } - }, [open]); - - // Cmd/Ctrl+M shortcut - useEffect(() => { - const handler = (e: KeyboardEvent) => { - if ((e.metaKey || e.ctrlKey) && e.key === "m") { - e.preventDefault(); - setOpen((prev) => !prev); - } - }; - document.addEventListener("keydown", handler); - return () => document.removeEventListener("keydown", handler); - }, []); - - // Focus search input on open - useEffect(() => { - if (open && !isMobile) { - requestAnimationFrame(() => searchInputRef.current?.focus()); - } - }, [open, isMobile, activeTab]); - - // ─── Data ─── - const { data: llmUserConfigs, isLoading: llmUserLoading } = - useAtomValue(newLLMConfigsAtom); + // LLM data + const { data: llmUserConfigs, isLoading: llmUserLoading } = useAtomValue(newLLMConfigsAtom); const { data: llmGlobalConfigs, isLoading: llmGlobalLoading } = useAtomValue(globalNewLLMConfigsAtom); - const { data: preferences, isLoading: prefsLoading } = - useAtomValue(llmPreferencesAtom); + const { data: preferences, isLoading: prefsLoading } = useAtomValue(llmPreferencesAtom); const searchSpaceId = useAtomValue(activeSearchSpaceIdAtom); - const { mutateAsync: updatePreferences } = useAtomValue( - updateLLMPreferencesMutationAtom, - ); + const { mutateAsync: updatePreferences } = useAtomValue(updateLLMPreferencesMutationAtom); + + // Image data const { data: imageGlobalConfigs, isLoading: imageGlobalLoading } = useAtomValue(globalImageGenConfigsAtom); - const { data: imageUserConfigs, isLoading: imageUserLoading } = - useAtomValue(imageGenConfigsAtom); - const { data: visionGlobalConfigs, isLoading: visionGlobalLoading } = - useAtomValue(globalVisionLLMConfigsAtom); + const { data: imageUserConfigs, isLoading: imageUserLoading } = useAtomValue(imageGenConfigsAtom); + + // Vision data + const { data: visionGlobalConfigs, isLoading: visionGlobalLoading } = useAtomValue( + globalVisionLLMConfigsAtom + ); const { data: visionUserConfigs, isLoading: visionUserLoading } = useAtomValue(visionLLMConfigsAtom); @@ -343,220 +110,133 @@ export function ModelSelector({ visionGlobalLoading || visionUserLoading; - // ─── Current selected configs ─── + // ─── LLM current config ─── const currentLLMConfig = useMemo(() => { if (!preferences) return null; - const id = preferences.agent_llm_id; - if (id === null || id === undefined) return null; - if (id <= 0) return llmGlobalConfigs?.find((c) => c.id === id) ?? null; - return llmUserConfigs?.find((c) => c.id === id) ?? null; + const agentLlmId = preferences.agent_llm_id; + if (agentLlmId === null || agentLlmId === undefined) return null; + if (agentLlmId <= 0) { + return llmGlobalConfigs?.find((c) => c.id === agentLlmId) ?? null; + } + return llmUserConfigs?.find((c) => c.id === agentLlmId) ?? null; }, [preferences, llmGlobalConfigs, llmUserConfigs]); const isLLMAutoMode = - currentLLMConfig && - "is_auto_mode" in currentLLMConfig && - currentLLMConfig.is_auto_mode; + currentLLMConfig && "is_auto_mode" in currentLLMConfig && currentLLMConfig.is_auto_mode; + // ─── Image current config ─── const currentImageConfig = useMemo(() => { if (!preferences) return null; const id = preferences.image_generation_config_id; if (id === null || id === undefined) return null; - return ( - imageGlobalConfigs?.find((c) => c.id === id) ?? - imageUserConfigs?.find((c) => c.id === id) ?? - null - ); + const globalMatch = imageGlobalConfigs?.find((c) => c.id === id); + if (globalMatch) return globalMatch; + return imageUserConfigs?.find((c) => c.id === id) ?? null; }, [preferences, imageGlobalConfigs, imageUserConfigs]); const isImageAutoMode = - currentImageConfig && - "is_auto_mode" in currentImageConfig && - currentImageConfig.is_auto_mode; + currentImageConfig && "is_auto_mode" in currentImageConfig && currentImageConfig.is_auto_mode; + // ─── Vision current config ─── const currentVisionConfig = useMemo(() => { if (!preferences) return null; const id = preferences.vision_llm_config_id; if (id === null || id === undefined) return null; - return ( - visionGlobalConfigs?.find((c) => c.id === id) ?? - visionUserConfigs?.find((c) => c.id === id) ?? - null - ); + const globalMatch = visionGlobalConfigs?.find((c) => c.id === id); + if (globalMatch) return globalMatch; + return visionUserConfigs?.find((c) => c.id === id) ?? null; }, [preferences, visionGlobalConfigs, visionUserConfigs]); - const isVisionAutoMode = - currentVisionConfig && - "is_auto_mode" in currentVisionConfig && - currentVisionConfig.is_auto_mode; - - // ─── Filtered configs (separate global / user for section headers) ─── - const filteredLLMGlobal = useMemo( - () => - filterAndScore(llmGlobalConfigs ?? [], selectedProvider, searchQuery), - [llmGlobalConfigs, selectedProvider, searchQuery], - ); - const filteredLLMUser = useMemo( - () => - filterAndScore(llmUserConfigs ?? [], selectedProvider, searchQuery), - [llmUserConfigs, selectedProvider, searchQuery], - ); - const filteredImageGlobal = useMemo( - () => - filterAndScore( - imageGlobalConfigs ?? [], - selectedProvider, - searchQuery, - ), - [imageGlobalConfigs, selectedProvider, searchQuery], - ); - const filteredImageUser = useMemo( - () => - filterAndScore( - imageUserConfigs ?? [], - selectedProvider, - searchQuery, - ), - [imageUserConfigs, selectedProvider, searchQuery], - ); - const filteredVisionGlobal = useMemo( - () => - filterAndScore( - visionGlobalConfigs ?? [], - selectedProvider, - searchQuery, - ), - [visionGlobalConfigs, selectedProvider, searchQuery], - ); - const filteredVisionUser = useMemo( - () => - filterAndScore( - visionUserConfigs ?? [], - selectedProvider, - searchQuery, - ), - [visionUserConfigs, selectedProvider, searchQuery], - ); - - // Combined display list for keyboard navigation - const currentDisplayItems: DisplayItem[] = useMemo(() => { - const toItems = ( - configs: ConfigBase[], - isGlobal: boolean, - ): DisplayItem[] => - configs.map((c) => ({ - config: c as ConfigBase & Record, - isGlobal, - isAutoMode: - isGlobal && - "is_auto_mode" in c && - !!(c as Record).is_auto_mode, - })); - - switch (activeTab) { - case "llm": - return [ - ...toItems(filteredLLMGlobal, true), - ...toItems(filteredLLMUser, false), - ]; - case "image": - return [ - ...toItems(filteredImageGlobal, true), - ...toItems(filteredImageUser, false), - ]; - case "vision": - return [ - ...toItems(filteredVisionGlobal, true), - ...toItems(filteredVisionUser, false), - ]; - } - }, [ - activeTab, - filteredLLMGlobal, - filteredLLMUser, - filteredImageGlobal, - filteredImageUser, - filteredVisionGlobal, - filteredVisionUser, - ]); - - // ─── Provider sidebar data ─── - // Collect which providers actually have configured models for the active tab - const configuredProviderSet = useMemo(() => { - const configs = - activeTab === "llm" - ? [ - ...(llmGlobalConfigs ?? []), - ...(llmUserConfigs ?? []), - ] - : activeTab === "image" - ? [ - ...(imageGlobalConfigs ?? []), - ...(imageUserConfigs ?? []), - ] - : [ - ...(visionGlobalConfigs ?? []), - ...(visionUserConfigs ?? []), - ]; - const set = new Set(); - for (const c of configs) { - if (c.provider) set.add(c.provider.toUpperCase()); - } - return set; - }, [ - activeTab, - llmGlobalConfigs, - llmUserConfigs, - imageGlobalConfigs, - imageUserConfigs, - visionGlobalConfigs, - visionUserConfigs, - ]); - - // Show only providers valid for the active tab; configured ones first - const activeProviders = useMemo(() => { - const tabKeys = PROVIDER_KEYS_BY_TAB[activeTab] ?? LLM_PROVIDER_KEYS; - const configured = tabKeys.filter((p) => - configuredProviderSet.has(p), + const isVisionAutoMode = useMemo(() => { + return ( + currentVisionConfig && + "is_auto_mode" in currentVisionConfig && + currentVisionConfig.is_auto_mode ); - const unconfigured = tabKeys.filter( - (p) => !configuredProviderSet.has(p), + }, [currentVisionConfig]); + + // ─── LLM filtering ─── + const filteredLLMGlobal = useMemo(() => { + if (!llmGlobalConfigs) return []; + if (!llmSearchQuery) return llmGlobalConfigs; + const q = llmSearchQuery.toLowerCase(); + return llmGlobalConfigs.filter( + (c) => + c.name.toLowerCase().includes(q) || + c.model_name.toLowerCase().includes(q) || + c.provider.toLowerCase().includes(q) ); - return ["all", ...configured, ...unconfigured]; - }, [activeTab, configuredProviderSet]); + }, [llmGlobalConfigs, llmSearchQuery]); - const providerModelCounts = useMemo(() => { - const allConfigs = - activeTab === "llm" - ? [ - ...(llmGlobalConfigs ?? []), - ...(llmUserConfigs ?? []), - ] - : activeTab === "image" - ? [ - ...(imageGlobalConfigs ?? []), - ...(imageUserConfigs ?? []), - ] - : [ - ...(visionGlobalConfigs ?? []), - ...(visionUserConfigs ?? []), - ]; - const counts: Record = { all: allConfigs.length }; - for (const c of allConfigs) { - const p = c.provider.toUpperCase(); - counts[p] = (counts[p] || 0) + 1; - } - return counts; - }, [ - activeTab, - llmGlobalConfigs, - llmUserConfigs, - imageGlobalConfigs, - imageUserConfigs, - visionGlobalConfigs, - visionUserConfigs, - ]); + const filteredLLMUser = useMemo(() => { + if (!llmUserConfigs) return []; + if (!llmSearchQuery) return llmUserConfigs; + const q = llmSearchQuery.toLowerCase(); + return llmUserConfigs.filter( + (c) => + c.name.toLowerCase().includes(q) || + c.model_name.toLowerCase().includes(q) || + c.provider.toLowerCase().includes(q) + ); + }, [llmUserConfigs, llmSearchQuery]); - // ─── Selection handlers ─── + const totalLLMModels = (llmGlobalConfigs?.length ?? 0) + (llmUserConfigs?.length ?? 0); + + // ─── Image filtering ─── + const filteredImageGlobal = useMemo(() => { + if (!imageGlobalConfigs) return []; + if (!imageSearchQuery) return imageGlobalConfigs; + const q = imageSearchQuery.toLowerCase(); + return imageGlobalConfigs.filter( + (c) => + c.name.toLowerCase().includes(q) || + c.model_name.toLowerCase().includes(q) || + c.provider.toLowerCase().includes(q) + ); + }, [imageGlobalConfigs, imageSearchQuery]); + + const filteredImageUser = useMemo(() => { + if (!imageUserConfigs) return []; + if (!imageSearchQuery) return imageUserConfigs; + const q = imageSearchQuery.toLowerCase(); + return imageUserConfigs.filter( + (c) => + c.name.toLowerCase().includes(q) || + c.model_name.toLowerCase().includes(q) || + c.provider.toLowerCase().includes(q) + ); + }, [imageUserConfigs, imageSearchQuery]); + + const totalImageModels = (imageGlobalConfigs?.length ?? 0) + (imageUserConfigs?.length ?? 0); + + // ─── Vision filtering ─── + const filteredVisionGlobal = useMemo(() => { + if (!visionGlobalConfigs) return []; + if (!visionSearchQuery) return visionGlobalConfigs; + const q = visionSearchQuery.toLowerCase(); + return visionGlobalConfigs.filter( + (c) => + c.name.toLowerCase().includes(q) || + c.model_name.toLowerCase().includes(q) || + c.provider.toLowerCase().includes(q) + ); + }, [visionGlobalConfigs, visionSearchQuery]); + + const filteredVisionUser = useMemo(() => { + if (!visionUserConfigs) return []; + if (!visionSearchQuery) return visionUserConfigs; + const q = visionSearchQuery.toLowerCase(); + return visionUserConfigs.filter( + (c) => + c.name.toLowerCase().includes(q) || + c.model_name.toLowerCase().includes(q) || + c.provider.toLowerCase().includes(q) + ); + }, [visionUserConfigs, visionSearchQuery]); + + const totalVisionModels = (visionGlobalConfigs?.length ?? 0) + (visionUserConfigs?.length ?? 0); + + // ─── Handlers ─── const handleSelectLLM = useCallback( async (config: NewLLMConfigPublic | GlobalNewLLMConfig) => { if (currentLLMConfig?.id === config.id) { @@ -574,11 +254,21 @@ export function ModelSelector({ }); toast.success(`Switched to ${config.name}`); setOpen(false); - } catch { + } catch (error) { + console.error("Failed to switch model:", error); toast.error("Failed to switch model"); } }, - [currentLLMConfig, searchSpaceId, updatePreferences], + [currentLLMConfig, searchSpaceId, updatePreferences] + ); + + const handleEditLLMConfig = useCallback( + (e: React.MouseEvent, config: NewLLMConfigPublic | GlobalNewLLMConfig, isGlobal: boolean) => { + e.stopPropagation(); + onEditLLM(config, isGlobal); + setOpen(false); + }, + [onEditLLM] ); const handleSelectImage = useCallback( @@ -602,7 +292,7 @@ export function ModelSelector({ toast.error("Failed to switch image model"); } }, - [currentImageConfig, searchSpaceId, updatePreferences], + [currentImageConfig, searchSpaceId, updatePreferences] ); const handleSelectVision = useCallback( @@ -626,687 +316,649 @@ export function ModelSelector({ toast.error("Failed to switch vision model"); } }, - [currentVisionConfig, searchSpaceId, updatePreferences], + [currentVisionConfig, searchSpaceId, updatePreferences] ); - const handleSelectItem = useCallback( - (item: DisplayItem) => { - switch (activeTab) { - case "llm": - handleSelectLLM( - item.config as NewLLMConfigPublic | GlobalNewLLMConfig, - ); - break; - case "image": - handleSelectImage(item.config.id); - break; - case "vision": - handleSelectVision(item.config.id); - break; - } - }, - [activeTab, handleSelectLLM, handleSelectImage, handleSelectVision], - ); - - const handleEditItem = useCallback( - (e: React.MouseEvent, item: DisplayItem) => { - e.stopPropagation(); - setOpen(false); - switch (activeTab) { - case "llm": - onEditLLM( - item.config as NewLLMConfigPublic | GlobalNewLLMConfig, - item.isGlobal, - ); - break; - case "image": - onEditImage?.( - item.config as ImageGenerationConfig | GlobalImageGenConfig, - item.isGlobal, - ); - break; - case "vision": - onEditVision?.( - item.config as VisionLLMConfig | GlobalVisionLLMConfig, - item.isGlobal, - ); - break; - } - }, - [activeTab, onEditLLM, onEditImage, onEditVision], - ); - - // ─── Keyboard navigation ─── - useEffect(() => { - setFocusedIndex(-1); - }, [searchQuery, selectedProvider]); - - useEffect(() => { - if (focusedIndex < 0 || !modelListRef.current) return; - const items = - modelListRef.current.querySelectorAll("[data-model-index]"); - items[focusedIndex]?.scrollIntoView({ - block: "nearest", - behavior: "smooth", - }); - }, [focusedIndex]); - - const handleKeyDown = useCallback( - (e: React.KeyboardEvent) => { - const count = currentDisplayItems.length; - - // Arrow Left/Right cycle provider filters - if (e.key === "ArrowLeft" || e.key === "ArrowRight") { - e.preventDefault(); - const providers = activeProviders; - const idx = providers.indexOf(selectedProvider); - let next: number; - if (e.key === "ArrowLeft") { - next = idx > 0 ? idx - 1 : providers.length - 1; - } else { - next = - idx < providers.length - 1 ? idx + 1 : 0; - } - setSelectedProvider(providers[next]); - if (providerSidebarRef.current) { - const buttons = - providerSidebarRef.current.querySelectorAll("button"); - buttons[next]?.scrollIntoView({ - block: "nearest", - inline: "nearest", - behavior: "smooth", - }); - } - return; - } - - if (count === 0) return; - - switch (e.key) { - case "ArrowDown": - e.preventDefault(); - setFocusedIndex((prev) => - prev < count - 1 ? prev + 1 : 0, - ); - break; - case "ArrowUp": - e.preventDefault(); - setFocusedIndex((prev) => - prev > 0 ? prev - 1 : count - 1, - ); - break; - case "Enter": - e.preventDefault(); - if (focusedIndex >= 0 && focusedIndex < count) { - handleSelectItem(currentDisplayItems[focusedIndex]); - } - break; - case "Home": - e.preventDefault(); - setFocusedIndex(0); - break; - case "End": - e.preventDefault(); - setFocusedIndex(count - 1); - break; - } - }, - [ - currentDisplayItems, - focusedIndex, - activeProviders, - selectedProvider, - handleSelectItem, - ], - ); - - // ─── Render: Provider sidebar ─── - const renderProviderSidebar = () => { - const configuredCount = configuredProviderSet.size; - - return ( -
-
{ - const t = e.currentTarget; - setShowScrollIndicator( - t.scrollHeight - t.scrollTop > - t.clientHeight + 10, - ); - }} - className={cn( - isMobile - ? "flex flex-row gap-0.5 px-2 py-1.5 overflow-x-auto border-b border-border/40" - : "flex flex-col gap-0.5 p-1 overflow-y-auto flex-1 [&::-webkit-scrollbar]:w-0 [&::-webkit-scrollbar-track]:bg-transparent", - )} - > - {activeProviders.map((provider, idx) => { - const isAll = provider === "all"; - const isActive = selectedProvider === provider; - const count = providerModelCounts[provider] || 0; - const isConfigured = - isAll || configuredProviderSet.has(provider); - - // Separator between configured and unconfigured providers - // idx 0 is "all", configured run from 1..configuredCount, unconfigured start at configuredCount+1 - const showSeparator = - !isAll && - idx === configuredCount + 1 && - configuredCount > 0; - - return ( - - {showSeparator && - (isMobile ? ( -
- ) : ( -
- ))} - - - - - - {isAll - ? "All Models" - : formatProviderName( - provider, - )} - {isConfigured - ? ` (${count})` - : " — not configured"} - - - - ); - })} -
- {!isMobile && showScrollIndicator && ( -
- -
- )} -
- ); - }; - - // ─── Render: Model card ─── - const getSelectedId = () => { - switch (activeTab) { - case "llm": - return currentLLMConfig?.id; - case "image": - return currentImageConfig?.id; - case "vision": - return currentVisionConfig?.id; - } - }; - - const renderModelCard = (item: DisplayItem, index: number) => { - const { config, isAutoMode } = item; - const isSelected = getSelectedId() === config.id; - const isFocused = focusedIndex === index; - const hasCitations = - "citations_enabled" in config && !!config.citations_enabled; - - return ( -
handleSelectItem(item)} - onMouseEnter={() => setFocusedIndex(index)} - className={cn( - "group flex items-start gap-2.5 px-2.5 py-2 rounded-lg cursor-pointer", - "transition-all duration-150 mx-1", - "hover:bg-accent/40 active:scale-[0.99]", - isSelected && "bg-primary/6 dark:bg-primary/8", - isFocused && "bg-accent/50 ring-1 ring-primary/20", - )} - > - {/* Provider icon */} -
- {getProviderIcon(config.provider as string, { - isAutoMode, - className: "size-5", - })} -
- - {/* Model info */} -
-
- - {config.name} - - {isAutoMode && ( - - Recommended - - )} -
-
- - {isAutoMode - ? "Auto Mode" - : (config.model_name as string)} - - {!isAutoMode && hasCitations && ( - - Citations - - )} -
-
- - {/* Actions */} -
- {!isAutoMode && ( - - )} - {isSelected && ( - - )} -
-
- ); - }; - - // ─── Render: Full content ─── - const renderContent = () => { - const globalItems = currentDisplayItems.filter((i) => i.isGlobal); - const userItems = currentDisplayItems.filter((i) => !i.isGlobal); - const globalStartIdx = 0; - const userStartIdx = globalItems.length; - - const addHandler = - activeTab === "llm" - ? onAddNewLLM - : activeTab === "image" - ? onAddNewImage - : onAddNewVision; - const addLabel = - activeTab === "llm" - ? "Add Model" - : activeTab === "image" - ? "Add Image Model" - : "Add Vision Model"; - - return ( -
- {/* Tab header */} -
-
- {( - [ - { - value: "llm" as const, - icon: Zap, - label: "LLM", - }, - { - value: "image" as const, - icon: ImageIcon, - label: "Image", - }, - { - value: "vision" as const, - icon: Eye, - label: "Vision", - }, - ] as const - ).map(({ value, icon: Icon, label }) => ( - - ))} -
-
- - {/* Two-pane layout */} -
- {/* Provider sidebar */} - {renderProviderSidebar()} - - {/* Main content */} -
- {/* Search */} -
- - - setSearchQuery(e.target.value) - } - onKeyDown={handleKeyDown} - autoFocus={!isMobile} - role="combobox" - aria-expanded={true} - aria-controls="model-selector-list" - className={cn( - "w-full pl-8 pr-3 py-1.5 text-xs rounded-lg", - "bg-secondary/30 border border-border/40", - "focus:outline-none focus:ring-2 focus:ring-primary/20 focus:border-primary/40", - "placeholder:text-muted-foreground/50", - "transition-[box-shadow,border-color] duration-200", - )} - /> -
- - {/* Provider header when filtered */} - {selectedProvider !== "all" && ( -
- {getProviderIcon(selectedProvider, { - className: "size-4", - })} - - {formatProviderName(selectedProvider)} - - - {configuredProviderSet.has( - selectedProvider, - ) - ? `${providerModelCounts[selectedProvider] || 0} models` - : "Not configured"} - -
- )} - - {/* Model list */} -
- {currentDisplayItems.length === 0 ? ( -
- {selectedProvider !== "all" && - !configuredProviderSet.has( - selectedProvider, - ) ? ( - <> -
- {getProviderIcon( - selectedProvider, - { - className: - "size-10", - }, - )} -
-

- No{" "} - {formatProviderName( - selectedProvider, - )}{" "} - models configured -

-

- Add a model with this - provider to get started -

- {addHandler && ( - - )} - - ) : ( - <> - -

- No models found -

-

- Try a different search - term -

- - )} -
- ) : ( - <> - {globalItems.length > 0 && ( - <> -
- Global Models -
- {globalItems.map((item, i) => - renderModelCard( - item, - globalStartIdx + i, - ), - )} - - )} - {globalItems.length > 0 && - userItems.length > 0 && ( -
- )} - {userItems.length > 0 && ( - <> -
- Your Configurations -
- {userItems.map((item, i) => - renderModelCard( - item, - userStartIdx + i, - ), - )} - - )} - - )} -
- - {/* Add model button */} - {addHandler && ( -
- -
- )} -
-
-
- ); - }; - - // ─── Trigger button ─── - const triggerButton = ( - - ); - - // ─── Shell: Drawer on mobile, Popover on desktop ─── - if (isMobile) { - return ( - - {triggerButton} - - - - Select Model - -
- {renderContent()} -
- - - ); - } - return ( - {triggerButton} + + + + e.preventDefault()} > - {renderContent()} + setActiveTab(v as "llm" | "image" | "vision")} + className="w-full" + > +
+ + + + LLM + + + + Image + + + + Vision + + +
+ + {/* ─── LLM Tab ─── */} + + + {totalLLMModels > 3 && ( +
+ +
+ )} + + + +
+ +

No models found

+

Try a different search term

+
+ + + {/* Global LLM Configs */} + {filteredLLMGlobal.length > 0 && ( + +
+ Global Models +
+ {filteredLLMGlobal.map((config) => { + const isSelected = currentLLMConfig?.id === config.id; + const isAutoMode = "is_auto_mode" in config && config.is_auto_mode; + return ( + handleSelectLLM(config)} + className={cn( + "mx-2 rounded-lg mb-1 cursor-pointer group transition-all", + "hover:bg-accent/50 dark:hover:bg-white/[0.06]", + isSelected && "bg-accent/80 dark:bg-white/[0.06]", + isAutoMode && "" + )} + > +
+
+
+ {getProviderIcon(config.provider, { isAutoMode })} +
+
+
+ {config.name} + {isAutoMode && ( + + Recommended + + )} + {isSelected && ( + + )} +
+
+ + {isAutoMode ? "Auto Mode" : config.model_name} + + {!isAutoMode && config.citations_enabled && ( + + Citations + + )} +
+
+
+ {!isAutoMode && ( + + )} +
+ + ); + })} + + )} + + {filteredLLMGlobal.length > 0 && filteredLLMUser.length > 0 && ( + + )} + + {/* User LLM Configs */} + {filteredLLMUser.length > 0 && ( + +
+ Your Configurations +
+ {filteredLLMUser.map((config) => { + const isSelected = currentLLMConfig?.id === config.id; + return ( + handleSelectLLM(config)} + className={cn( + "mx-2 rounded-lg mb-1 cursor-pointer group transition-all", + "hover:bg-accent/50 dark:hover:bg-white/[0.06]", + isSelected && "bg-accent/80 dark:bg-white/[0.06]" + )} + > +
+
+
{getProviderIcon(config.provider)}
+
+
+ {config.name} + {isSelected && ( + + )} +
+
+ + {config.model_name} + + {config.citations_enabled && ( + + Citations + + )} +
+
+
+ +
+ + ); + })} + + )} + + {/* Add New LLM Config */} +
+ +
+ + + + + {/* ─── Image Tab ─── */} + + + {totalImageModels > 3 && ( +
+ +
+ )} + + +
+ +

No image models found

+

Try a different search term

+
+ + + {/* Global Image Configs */} + {filteredImageGlobal.length > 0 && ( + +
+ Global Image Models +
+ {filteredImageGlobal.map((config) => { + const isSelected = currentImageConfig?.id === config.id; + const isAuto = "is_auto_mode" in config && config.is_auto_mode; + return ( + handleSelectImage(config.id)} + className={cn( + "mx-2 rounded-lg mb-1 cursor-pointer group transition-all hover:bg-accent/50 dark:hover:bg-white/[0.06]", + isSelected && "bg-accent/80 dark:bg-white/[0.06]", + isAuto && "" + )} + > +
+
+ {getProviderIcon(config.provider, { isAutoMode: isAuto })} +
+
+
+ {config.name} + {isAuto && ( + + Recommended + + )} + {isSelected && } +
+ + {isAuto ? "Auto Mode" : config.model_name} + +
+ {onEditImage && !isAuto && ( + + )} +
+ + ); + })} + + )} + + {/* User Image Configs */} + {filteredImageUser.length > 0 && ( + <> + {filteredImageGlobal.length > 0 && ( + + )} + +
+ Your Image Models +
+ {filteredImageUser.map((config) => { + const isSelected = currentImageConfig?.id === config.id; + return ( + handleSelectImage(config.id)} + className={cn( + "mx-2 rounded-lg mb-1 cursor-pointer group transition-all hover:bg-accent/50 dark:hover:bg-white/[0.06]", + isSelected && "bg-accent/80 dark:bg-white/[0.06]" + )} + > +
+
{getProviderIcon(config.provider)}
+
+
+ {config.name} + {isSelected && ( + + )} +
+ + {config.model_name} + +
+ {onEditImage && ( + + )} +
+ + ); + })} + + + )} + + {/* Add New Image Config */} + {onAddNewImage && ( +
+ +
+ )} + + + + + {/* ─── Vision Tab ─── */} + + + {totalVisionModels > 3 && ( +
+ +
+ )} + + +
+ +

No vision models found

+

Try a different search term

+
+ + + {filteredVisionGlobal.length > 0 && ( + +
+ Global Vision Models +
+ {filteredVisionGlobal.map((config) => { + const isSelected = currentVisionConfig?.id === config.id; + const isAuto = "is_auto_mode" in config && config.is_auto_mode; + return ( + handleSelectVision(config.id)} + className={cn( + "mx-2 rounded-lg mb-1 cursor-pointer group transition-all hover:bg-accent/50 dark:hover:bg-white/[0.06]", + isSelected && "bg-accent/80 dark:bg-white/[0.06]" + )} + > +
+
+ {getProviderIcon(config.provider, { isAutoMode: isAuto })} +
+
+
+ {config.name} + {isAuto && ( + + Recommended + + )} + {isSelected && } +
+ + {isAuto ? "Auto Mode" : config.model_name} + +
+ {onEditVision && !isAuto && ( + + )} +
+ + ); + })} + + )} + + {filteredVisionUser.length > 0 && ( + <> + {filteredVisionGlobal.length > 0 && ( + + )} + +
+ Your Vision Models +
+ {filteredVisionUser.map((config) => { + const isSelected = currentVisionConfig?.id === config.id; + return ( + handleSelectVision(config.id)} + className={cn( + "mx-2 rounded-lg mb-1 cursor-pointer group transition-all hover:bg-accent/50 dark:hover:bg-white/[0.06]", + isSelected && "bg-accent/80 dark:bg-white/[0.06]" + )} + > +
+
{getProviderIcon(config.provider)}
+
+
+ {config.name} + {isSelected && ( + + )} +
+ + {config.model_name} + +
+ {onEditVision && ( + + )} +
+ + ); + })} + + + )} + + {onAddNewVision && ( +
+ +
+ )} + + + + ); diff --git a/surfsense_web/components/new-chat/prompt-picker.tsx b/surfsense_web/components/new-chat/prompt-picker.tsx index b1723f0be..3e6457b8c 100644 --- a/surfsense_web/components/new-chat/prompt-picker.tsx +++ b/surfsense_web/components/new-chat/prompt-picker.tsx @@ -15,7 +15,7 @@ import { import { promptsAtom } from "@/atoms/prompts/prompts-query.atoms"; import { userSettingsDialogAtom } from "@/atoms/settings/settings-dialog.atoms"; -import { Skeleton } from "@/components/ui/skeleton"; +import { Spinner } from "@/components/ui/spinner"; import { cn } from "@/lib/utils"; export interface PromptPickerRef { @@ -28,10 +28,11 @@ interface PromptPickerProps { onSelect: (action: { name: string; prompt: string; mode: "transform" | "explore" }) => void; onDone: () => void; externalSearch?: string; + containerStyle?: React.CSSProperties; } export const PromptPicker = forwardRef(function PromptPicker( - { onSelect, onDone, externalSearch = "" }, + { onSelect, onDone, externalSearch = "", containerStyle }, ref ) { const setUserSettingsDialog = useSetAtom(userSettingsDialogAtom); @@ -59,21 +60,13 @@ export const PromptPicker = forwardRef(funct } } - const createPromptIndex = filtered.length; - const totalItems = filtered.length + 1; - const handleSelect = useCallback( (index: number) => { - if (index === createPromptIndex) { - onDone(); - setUserSettingsDialog({ open: true, initialTab: "prompts" }); - return; - } const action = filtered[index]; if (!action) return; onSelect({ name: action.name, prompt: action.prompt, mode: action.mode }); }, - [filtered, onSelect, createPromptIndex, onDone, setUserSettingsDialog] + [filtered, onSelect] ); useEffect(() => { @@ -100,98 +93,69 @@ export const PromptPicker = forwardRef(funct () => ({ selectHighlighted: () => handleSelect(highlightedIndex), moveUp: () => { + if (filtered.length === 0) return; shouldScrollRef.current = true; - setHighlightedIndex((prev) => (prev > 0 ? prev - 1 : totalItems - 1)); + setHighlightedIndex((prev) => (prev > 0 ? prev - 1 : filtered.length - 1)); }, moveDown: () => { + if (filtered.length === 0) return; shouldScrollRef.current = true; - setHighlightedIndex((prev) => (prev < totalItems - 1 ? prev + 1 : 0)); + setHighlightedIndex((prev) => (prev < filtered.length - 1 ? prev + 1 : 0)); }, }), - [totalItems, highlightedIndex, handleSelect] + [filtered.length, highlightedIndex, handleSelect] ); return ( -
-
+
+
{isLoading ? ( -
-
- -
- {["a", "b", "c", "d", "e"].map((id, i) => ( -
= 3 && "hidden sm:flex" - )} - > - - - - - - -
- ))} +
+
) : isError ? ( -
-

Failed to load prompts

-
+

Failed to load prompts

) : filtered.length === 0 ? ( -
-

No matching prompts

-
+

No matching prompts

) : ( -
-
- Saved Prompts -
- {filtered.map((action, index) => ( - - ))} - -
+ filtered.map((action, index) => ( -
+ )) )} + +
+
); diff --git a/surfsense_web/components/pricing/pricing-section.tsx b/surfsense_web/components/pricing/pricing-section.tsx index 9d75e9a1d..43177e383 100644 --- a/surfsense_web/components/pricing/pricing-section.tsx +++ b/surfsense_web/components/pricing/pricing-section.tsx @@ -1,9 +1,4 @@ -"use client"; -import React, { useRef, useEffect, useState } from "react"; -import { AnimatePresence, motion } from "motion/react"; -import { IconPlus } from "@tabler/icons-react"; import { Pricing } from "@/components/pricing"; -import { cn } from "@/lib/utils"; const demoPlans = [ { @@ -64,280 +59,13 @@ const demoPlans = [ }, ]; -interface FAQItem { - question: string; - answer: string; -} - -interface FAQSection { - title: string; - items: FAQItem[]; -} - -const faqData: FAQSection[] = [ - { - title: "Pages & Billing", - items: [ - { - question: "What exactly is a \"page\" in SurfSense?", - answer: - "A page is a simple billing unit that measures how much content you add to your knowledge base. For PDFs, one page equals one real PDF page. For other document types like Word, PowerPoint, and Excel files, pages are automatically estimated based on the file. Every file uses at least 1 page.", - }, - { - question: "How does the Pay As You Go plan work?", - answer: - "There's no monthly subscription. When you need more pages, simply purchase 1,000-page packs at $1 each. Purchased pages are added to your account immediately so you can keep indexing right away. You only pay when you actually need more.", - }, - { - question: "What happens if I run out of pages?", - answer: - "SurfSense checks your remaining pages before processing each file. If you don't have enough, the upload is paused and you'll be notified. You can purchase additional page packs at any time to continue. For cloud connector syncs, a small overage may be allowed so your sync doesn't partially fail.", - }, - { - question: "If I delete a document, do I get my pages back?", - answer: - "No. Deleting a document removes it from your knowledge base, but the pages it used are not refunded. Pages track your total usage over time, not how much is currently stored. So be mindful of what you index. Once pages are spent, they're spent even if you later remove the document.", - }, - ], - }, - { - title: "File Types & Connectors", - items: [ - { - question: "Which file types count toward my page limit?", - answer: - "Page limits only apply to document files that need processing, including PDFs, Word documents (DOC, DOCX, ODT, RTF), presentations (PPT, PPTX, ODP), spreadsheets (XLS, XLSX, ODS), ebooks (EPUB), and images (JPG, PNG, TIFF, WebP, BMP). Plain text files, code files, Markdown, CSV, TSV, HTML, audio, and video files do not consume any pages.", - }, - { - question: "How are pages consumed?", - answer: - "Pages are deducted whenever a document file is successfully indexed into your knowledge base, whether through direct uploads or file-based connector syncs (Google Drive, OneDrive, Dropbox, Local Folder). SurfSense checks your remaining pages before processing and only charges you after the file is indexed. Duplicate documents are automatically detected and won't cost you extra pages.", - }, - { - question: "Do connectors like Slack, Notion, or Gmail use pages?", - answer: - "No. Connectors that work with structured text data like Slack, Discord, Notion, Confluence, Jira, Linear, ClickUp, GitHub, Gmail, Google Calendar, Microsoft Teams, Airtable, Elasticsearch, Web Crawler, BookStack, Obsidian, and Luma do not use pages at all. Page limits only apply to file-based connectors that need document processing, such as Google Drive, OneDrive, Dropbox, and Local Folder syncs.", - }, - ], - }, - { - title: "Self-Hosting", - items: [ - { - question: "Can I self-host SurfSense with unlimited pages?", - answer: - "Yes! When self-hosting, you have full control over your page limits. The default self-hosted setup gives you effectively unlimited pages, so you can index as much data as your infrastructure supports.", - }, - ], - }, -]; - -const GridLineHorizontal = ({ - className, - offset, -}: { - className?: string; - offset?: string; -}) => { - return ( -
- ); -}; - -const GridLineVertical = ({ - className, - offset, -}: { - className?: string; - offset?: string; -}) => { - return ( -
- ); -}; - -function PricingFAQ() { - const [activeId, setActiveId] = useState(null); - const containerRef = useRef(null); - - useEffect(() => { - function handleClickOutside(event: MouseEvent) { - if ( - containerRef.current && - !containerRef.current.contains(event.target as Node) - ) { - setActiveId(null); - } - } - - document.addEventListener("mousedown", handleClickOutside); - return () => document.removeEventListener("mousedown", handleClickOutside); - }, []); - - const toggleQuestion = (id: string) => { - setActiveId(activeId === id ? null : id); - }; - - return ( -
-
-

- Frequently Asked Questions -

-

- Everything you need to know about SurfSense pages and billing. - Can't find what you need? Reach out at{" "} - - rohan@surfsense.com - -

-
- -
- {faqData.map((section) => ( -
-

- {section.title} -

-
- {section.items.map((item, index) => { - const id = `${section.title}-${index}`; - const isActive = activeId === id; - - return ( -
- {isActive && ( -
- - - - -
- )} - - - {isActive && ( - -

- {item.answer} -

- - )} - -
- ); - })} -
-
- ))} -
-
- ); -} - function PricingBasic() { return ( - <> - - - + ); } diff --git a/surfsense_web/components/settings/general-settings-manager.tsx b/surfsense_web/components/settings/general-settings-manager.tsx index 15d44906b..e83525982 100644 --- a/surfsense_web/components/settings/general-settings-manager.tsx +++ b/surfsense_web/components/settings/general-settings-manager.tsx @@ -2,18 +2,18 @@ import { useQuery } from "@tanstack/react-query"; import { useAtomValue } from "jotai"; -import { FolderArchive, Info } from "lucide-react"; +import { Info } from "lucide-react"; import { useTranslations } from "next-intl"; -import { useCallback, useEffect, useState } from "react"; +import { useEffect, useState } from "react"; import { toast } from "sonner"; import { updateSearchSpaceMutationAtom } from "@/atoms/search-spaces/search-space-mutation.atoms"; import { Alert, AlertDescription } from "@/components/ui/alert"; import { Button } from "@/components/ui/button"; +import { Card, CardContent, CardDescription, CardHeader, CardTitle } from "@/components/ui/card"; import { Input } from "@/components/ui/input"; import { Label } from "@/components/ui/label"; import { Skeleton } from "@/components/ui/skeleton"; import { searchSpacesApiService } from "@/lib/apis/search-spaces-api.service"; -import { authenticatedFetch } from "@/lib/auth-utils"; import { cacheKeys } from "@/lib/query-client/cache-keys"; import { Spinner } from "../ui/spinner"; @@ -40,37 +40,6 @@ export function GeneralSettingsManager({ searchSpaceId }: GeneralSettingsManager const [name, setName] = useState(""); const [description, setDescription] = useState(""); const [saving, setSaving] = useState(false); - const [isExporting, setIsExporting] = useState(false); - - const handleExportKB = useCallback(async () => { - if (isExporting) return; - setIsExporting(true); - try { - const response = await authenticatedFetch( - `${process.env.NEXT_PUBLIC_FASTAPI_BACKEND_URL}/api/v1/search-spaces/${searchSpaceId}/export`, - { method: "GET" } - ); - if (!response.ok) { - const errorData = await response.json().catch(() => ({ detail: "Export failed" })); - throw new Error(errorData.detail || "Export failed"); - } - const blob = await response.blob(); - const url = URL.createObjectURL(blob); - const a = document.createElement("a"); - a.href = url; - a.download = "knowledge-base.zip"; - document.body.appendChild(a); - a.click(); - document.body.removeChild(a); - URL.revokeObjectURL(url); - toast.success("Knowledge base exported"); - } catch (err) { - console.error("KB export failed:", err); - toast.error(err instanceof Error ? err.message : "Export failed"); - } finally { - setIsExporting(false); - } - }, [searchSpaceId, isExporting]); // Initialize state from fetched search space useEffect(() => { @@ -114,10 +83,16 @@ export function GeneralSettingsManager({ searchSpaceId }: GeneralSettingsManager if (loading) { return (
-
- - -
+ + + + + + + + + +
); } @@ -138,45 +113,61 @@ export function GeneralSettingsManager({ searchSpaceId }: GeneralSettingsManager - Update your search space name and description. + Update your search space name and description. These details help identify and organize + your workspace. - -
-
- - setName(e.target.value)} - /> -

- {t("general_name_description")} -

-
+ {/* Search Space Details Card */} + + + + Search Space Details + + Manage the basic information for this search space. + + + +
+ + setName(e.target.value)} + className="text-sm md:text-base h-9 md:h-10" + /> +

+ {t("general_name_description")} +

+
-
- - setDescription(e.target.value)} - /> -

- {t("general_description_description")} -

-
-
+
+ + setDescription(e.target.value)} + className="text-sm md:text-base h-9 md:h-10" + /> +

+ {t("general_description_description")} +

+
+ + -
+ {/* Action Buttons */} +
- -
-
- -

- Download all documents in this search space as a ZIP of markdown files. -

-
- -
); } diff --git a/surfsense_web/components/settings/image-model-manager.tsx b/surfsense_web/components/settings/image-model-manager.tsx index fb28e5b1c..831519aa2 100644 --- a/surfsense_web/components/settings/image-model-manager.tsx +++ b/surfsense_web/components/settings/image-model-manager.tsx @@ -1,7 +1,7 @@ "use client"; import { useAtomValue } from "jotai"; -import { AlertCircle, Dot, Edit3, Info, RefreshCw, Trash2 } from "lucide-react"; +import { AlertCircle, Dot, Edit3, Info, RefreshCw, Trash2, Wand2 } from "lucide-react"; import { useMemo, useState } from "react"; import { deleteImageGenConfigMutationAtom } from "@/atoms/image-gen-config/image-gen-config-mutation.atoms"; import { @@ -209,20 +209,20 @@ export function ImageModelManager({ searchSpaceId }: ImageModelManagerProps) { {["skeleton-a", "skeleton-b", "skeleton-c"].map((key) => ( -
- +
-
- - -
- - -
+
+ + +
+
+ + +
@@ -255,25 +255,20 @@ export function ImageModelManager({ searchSpaceId }: ImageModelManagerProps) {
- {/* Header: Icon + Name + Actions */} -
-
-
- {getProviderIcon(config.provider, { className: "size-4" })} -
-
-

- {config.name} -

- {config.description && ( -

- {config.description} -

- )} -
+ {/* Header: Name + Actions */} +
+
+

+ {config.name} +

+ {config.description && ( +

+ {config.description} +

+ )}
{(canUpdate || canDelete) && ( -
+
{canUpdate && ( @@ -282,7 +277,7 @@ export function ImageModelManager({ searchSpaceId }: ImageModelManagerProps) { variant="ghost" size="icon" onClick={() => openEditDialog(config)} - className="h-7 w-7 rounded-lg text-muted-foreground hover:text-foreground" + className="h-7 w-7 text-muted-foreground hover:text-foreground" > @@ -299,7 +294,7 @@ export function ImageModelManager({ searchSpaceId }: ImageModelManagerProps) { variant="ghost" size="icon" onClick={() => setConfigToDelete(config)} - className="h-7 w-7 rounded-lg text-muted-foreground hover:text-destructive" + className="h-7 w-7 text-muted-foreground hover:text-destructive" > @@ -312,9 +307,17 @@ export function ImageModelManager({ searchSpaceId }: ImageModelManagerProps) { )}
+ {/* Provider + Model */} +
+ {getProviderIcon(config.provider, { className: "size-3.5 shrink-0" })} + + {config.model_name} + +
+ {/* Footer: Date + Creator */} -
- +
+ {new Date(config.created_at).toLocaleDateString(undefined, { year: "numeric", month: "short", @@ -323,11 +326,11 @@ export function ImageModelManager({ searchSpaceId }: ImageModelManagerProps) { {member && ( <> - + -
+
{member.avatarUrl && ( @@ -336,7 +339,7 @@ export function ImageModelManager({ searchSpaceId }: ImageModelManagerProps) { {getInitials(member.name)} - + {member.name}
diff --git a/surfsense_web/components/settings/llm-role-manager.tsx b/surfsense_web/components/settings/llm-role-manager.tsx index d6eb7c64d..e280db493 100644 --- a/surfsense_web/components/settings/llm-role-manager.tsx +++ b/surfsense_web/components/settings/llm-role-manager.tsx @@ -6,7 +6,7 @@ import { Bot, CircleCheck, CircleDashed, - ScanEye, + Eye, FileText, ImageIcon, RefreshCw, @@ -74,7 +74,7 @@ const ROLE_DESCRIPTIONS = { configType: "image" as const, }, vision: { - icon: ScanEye, + icon: Eye, title: "Vision LLM", description: "Vision-capable model for screenshot analysis and context extraction", color: "text-muted-foreground", diff --git a/surfsense_web/components/settings/agent-model-manager.tsx b/surfsense_web/components/settings/model-config-manager.tsx similarity index 83% rename from surfsense_web/components/settings/agent-model-manager.tsx rename to surfsense_web/components/settings/model-config-manager.tsx index f7a2fb824..f83251426 100644 --- a/surfsense_web/components/settings/agent-model-manager.tsx +++ b/surfsense_web/components/settings/model-config-manager.tsx @@ -10,6 +10,7 @@ import { MessageSquareQuote, RefreshCw, Trash2, + Wand2, } from "lucide-react"; import { useMemo, useState } from "react"; import { membersAtom, myAccessAtom } from "@/atoms/members/members-query.atoms"; @@ -42,7 +43,7 @@ import { useMediaQuery } from "@/hooks/use-media-query"; import { getProviderIcon } from "@/lib/provider-icons"; import { cn } from "@/lib/utils"; -interface AgentModelManagerProps { +interface ModelConfigManagerProps { searchSpaceId: number; } @@ -54,7 +55,7 @@ function getInitials(name: string): string { return name.slice(0, 2).toUpperCase(); } -export function AgentModelManager({ searchSpaceId }: AgentModelManagerProps) { +export function ModelConfigManager({ searchSpaceId }: ModelConfigManagerProps) { const isDesktop = useMediaQuery("(min-width: 768px)"); // Mutations const { mutateAsync: deleteConfig, isPending: isDeleting } = useAtomValue( @@ -207,26 +208,28 @@ export function AgentModelManager({ searchSpaceId }: AgentModelManagerProps) { {["skeleton-a", "skeleton-b", "skeleton-c"].map((key) => ( - {/* Header: Icon + Name */} -
- + {/* Header */} +
+ {/* Provider + Model */} +
+ + +
{/* Feature badges */}
+
{/* Footer */} -
- - -
- - -
+
+ + +
@@ -259,25 +262,20 @@ export function AgentModelManager({ searchSpaceId }: AgentModelManagerProps) {
- {/* Header: Icon + Name + Actions */} -
-
-
- {getProviderIcon(config.provider, { className: "size-4" })} -
-
-

- {config.name} -

- {config.description && ( -

- {config.description} -

- )} -
+ {/* Header: Name + Actions */} +
+
+

+ {config.name} +

+ {config.description && ( +

+ {config.description} +

+ )}
{(canUpdate || canDelete) && ( -
+
{canUpdate && ( @@ -286,7 +284,7 @@ export function AgentModelManager({ searchSpaceId }: AgentModelManagerProps) { variant="ghost" size="icon" onClick={() => openEditDialog(config)} - className="h-7 w-7 rounded-lg text-muted-foreground hover:text-foreground" + className="h-7 w-7 text-muted-foreground hover:text-foreground" > @@ -303,7 +301,7 @@ export function AgentModelManager({ searchSpaceId }: AgentModelManagerProps) { variant="ghost" size="icon" onClick={() => setConfigToDelete(config)} - className="h-7 w-7 rounded-lg text-muted-foreground hover:text-destructive" + className="h-7 w-7 text-muted-foreground hover:text-destructive" > @@ -316,12 +314,20 @@ export function AgentModelManager({ searchSpaceId }: AgentModelManagerProps) { )}
+ {/* Provider + Model */} +
+ {getProviderIcon(config.provider, { className: "size-3.5 shrink-0" })} + + {config.model_name} + +
+ {/* Feature badges */}
{config.citations_enabled && ( Citations @@ -330,8 +336,8 @@ export function AgentModelManager({ searchSpaceId }: AgentModelManagerProps) { {!config.use_default_system_instructions && config.system_instructions && ( Custom @@ -340,8 +346,8 @@ export function AgentModelManager({ searchSpaceId }: AgentModelManagerProps) {
{/* Footer: Date + Creator */} -
- +
+ {new Date(config.created_at).toLocaleDateString(undefined, { year: "numeric", month: "short", @@ -350,11 +356,11 @@ export function AgentModelManager({ searchSpaceId }: AgentModelManagerProps) { {member && ( <> - + -
+
{member.avatarUrl && ( @@ -363,7 +369,7 @@ export function AgentModelManager({ searchSpaceId }: AgentModelManagerProps) { {getInitials(member.name)} - + {member.name}
diff --git a/surfsense_web/components/settings/roles-manager.tsx b/surfsense_web/components/settings/roles-manager.tsx index b72a53854..5b30b5f60 100644 --- a/surfsense_web/components/settings/roles-manager.tsx +++ b/surfsense_web/components/settings/roles-manager.tsx @@ -4,7 +4,6 @@ import { useQuery } from "@tanstack/react-query"; import { useAtomValue } from "jotai"; import { Bot, - ChevronDown, Edit2, FileText, Globe, @@ -48,6 +47,7 @@ import { DialogDescription, DialogHeader, DialogTitle, + DialogTrigger, } from "@/components/ui/dialog"; import { DropdownMenu, @@ -58,6 +58,7 @@ import { } from "@/components/ui/dropdown-menu"; import { Input } from "@/components/ui/input"; import { Label } from "@/components/ui/label"; +import { ScrollArea } from "@/components/ui/scroll-area"; import { Spinner } from "@/components/ui/spinner"; import type { PermissionInfo } from "@/contracts/types/permissions.types"; import type { @@ -318,6 +319,100 @@ export function RolesManager({ searchSpaceId }: { searchSpaceId: number }) { ); } +// ============ Role Permissions Display ============ + +function RolePermissionsDialog({ + permissions, + roleName, + children, +}: { + permissions: string[]; + roleName: string; + children: React.ReactNode; +}) { + const isFullAccess = permissions.includes("*"); + + const grouped: Record = {}; + if (!isFullAccess) { + for (const perm of permissions) { + const [category, action] = perm.split(":"); + if (!grouped[category]) grouped[category] = []; + grouped[category].push(action); + } + } + + const sortedCategories = Object.keys(grouped).sort((a, b) => { + const orderA = CATEGORY_CONFIG[a]?.order ?? 99; + const orderB = CATEGORY_CONFIG[b]?.order ?? 99; + return orderA - orderB; + }); + + const categoryCount = sortedCategories.length; + + return ( + + {children} + + + {roleName} — Permissions + + {isFullAccess + ? "This role has unrestricted access to all resources" + : `${permissions.length} permissions across ${categoryCount} categories`} + + + {isFullAccess ? ( +
+
+ +
+
+

Full access

+

+ All permissions granted across every category +

+
+
+ ) : ( + +
+ {sortedCategories.map((category) => { + const actions = grouped[category]; + const config = CATEGORY_CONFIG[category] || { + label: category, + icon: FileText, + }; + const IconComponent = config.icon; + return ( +
+
+ + {config.label} +
+
+ {actions.map((action) => ( + + {ACTION_LABELS[action] || action.replace(/_/g, " ")} + + ))} +
+
+ ); + })} +
+ + )} + + + ); +} + function PermissionsBadge({ permissions }: { permissions: string[] }) { if (permissions.includes("*")) { return ( @@ -368,7 +463,6 @@ function RolesContent({ }) { const [showCreateRole, setShowCreateRole] = useState(false); const [editingRoleId, setEditingRoleId] = useState(null); - const [expandedRoleId, setExpandedRoleId] = useState(null); if (loading) { return ( @@ -414,170 +508,91 @@ function RolesContent({ )}
- {roles.map((role) => { - const isExpanded = expandedRoleId === role.id; - const isFullAccess = role.permissions.includes("*"); - - const grouped: Record = {}; - if (!isFullAccess) { - for (const perm of role.permissions) { - const [category, action] = perm.split(":"); - if (!grouped[category]) grouped[category] = []; - grouped[category].push(action); - } - } - const sortedCategories = Object.keys(grouped).sort((a, b) => { - const orderA = CATEGORY_CONFIG[a]?.order ?? 99; - const orderB = CATEGORY_CONFIG[b]?.order ?? 99; - return orderA - orderB; - }); - - return ( -
-
-
- {role.description && ( -

- {role.description} -

- )} - - -
- -
- - {!role.is_system_role && ( -
- - - - - e.preventDefault()}> - {canUpdate && ( - setEditingRoleId(role.id)}> - - Edit Role - - )} - {canDelete && ( - <> - - - - e.preventDefault()}> - - Delete Role - - - - - Delete role? - - This will permanently delete the "{role.name}" role. - Members with this role will lose their permissions. - - - - Cancel - onDeleteRole(role.id)} - className="bg-destructive text-destructive-foreground hover:bg-destructive/90" - > - Delete - - - - - - )} - - -
- )} - - + +
- {isExpanded && ( -
- {isFullAccess ? ( -
- -

- Full access — all permissions granted across every category -

-
- ) : ( -
- {sortedCategories.map((category) => { - const actions = grouped[category]; - const config = CATEGORY_CONFIG[category] || { - label: category, - icon: FileText, - }; - const IconComponent = config.icon; - return ( -
-
- - - {config.label} - -
-
- {actions.map((action) => ( - + +
+ + {!role.is_system_role && ( +
+ + + + + e.preventDefault()}> + {canUpdate && ( + setEditingRoleId(role.id)}> + + Edit Role + + )} + {canDelete && ( + <> + + + + e.preventDefault()}> + + Delete Role + + + + + Delete role? + + This will permanently delete the "{role.name}" role. + Members with this role will lose their permissions. + + + + Cancel + onDeleteRole(role.id)} + className="bg-destructive text-destructive-foreground hover:bg-destructive/90" > - {ACTION_LABELS[action] || action.replace(/_/g, " ")} - - ))} -
-
- ); - })} -
- )} + Delete + + + + + + )} + +
)}
- ); - })} +
+ ))}
); @@ -661,51 +676,46 @@ function PermissionsEditor({ return (
-
- +
onToggleCategory(category)} + onClick={(e) => e.stopPropagation()} aria-label={`Select all ${config.label} permissions`} /> - + Toggle + + +
-
+ {isExpanded && (
@@ -716,29 +726,28 @@ function PermissionsEditor({ const isSelected = selectedPermissions.includes(perm.value); return ( -
onTogglePermission(perm.value)} > - +
onTogglePermission(perm.value)} + onClick={(e) => e.stopPropagation()} className="shrink-0" /> -
+ ); })}
diff --git a/surfsense_web/components/settings/search-space-settings-dialog.tsx b/surfsense_web/components/settings/search-space-settings-dialog.tsx index e021e1e41..56ad0ab8f 100644 --- a/surfsense_web/components/settings/search-space-settings-dialog.tsx +++ b/surfsense_web/components/settings/search-space-settings-dialog.tsx @@ -7,7 +7,7 @@ import { Brain, CircleUser, Earth, - ScanEye, + Eye, ImageIcon, ListChecks, UserKey, @@ -25,10 +25,10 @@ const GeneralSettingsManager = dynamic( })), { ssr: false } ); -const AgentModelManager = dynamic( +const ModelConfigManager = dynamic( () => - import("@/components/settings/agent-model-manager").then((m) => ({ - default: m.AgentModelManager, + import("@/components/settings/model-config-manager").then((m) => ({ + default: m.ModelConfigManager, })), { ssr: false } ); @@ -88,7 +88,7 @@ export function SearchSpaceSettingsDialog({ searchSpaceId }: SearchSpaceSettings const navItems = [ { value: "general", label: t("nav_general"), icon: }, { value: "roles", label: t("nav_role_assignments"), icon: }, - { value: "models", label: t("nav_agent_models"), icon: }, + { value: "models", label: t("nav_agent_configs"), icon: }, { value: "image-models", label: t("nav_image_models"), @@ -97,7 +97,7 @@ export function SearchSpaceSettingsDialog({ searchSpaceId }: SearchSpaceSettings { value: "vision-models", label: t("nav_vision_models"), - icon: , + icon: , }, { value: "team-roles", label: t("nav_team_roles"), icon: }, { @@ -115,7 +115,7 @@ export function SearchSpaceSettingsDialog({ searchSpaceId }: SearchSpaceSettings const content: Record = { general: , - models: , + models: , roles: , "image-models": , "vision-models": , diff --git a/surfsense_web/components/settings/vision-model-manager.tsx b/surfsense_web/components/settings/vision-model-manager.tsx index 81528c86a..57ea8c205 100644 --- a/surfsense_web/components/settings/vision-model-manager.tsx +++ b/surfsense_web/components/settings/vision-model-manager.tsx @@ -208,20 +208,20 @@ export function VisionModelManager({ searchSpaceId }: VisionModelManagerProps) { {["skeleton-a", "skeleton-b", "skeleton-c"].map((key) => ( -
- +
-
- - -
- - -
+
+ + +
+
+ + +
@@ -253,25 +253,19 @@ export function VisionModelManager({ searchSpaceId }: VisionModelManagerProps) {
- {/* Header: Icon + Name + Actions */} -
-
-
- {getProviderIcon(config.provider, { className: "size-4" })} -
-
-

- {config.name} -

- {config.description && ( -

- {config.description} -

- )} -
+
+
+

+ {config.name} +

+ {config.description && ( +

+ {config.description} +

+ )}
{(canUpdate || canDelete) && ( -
+
{canUpdate && ( @@ -280,7 +274,7 @@ export function VisionModelManager({ searchSpaceId }: VisionModelManagerProps) { variant="ghost" size="icon" onClick={() => openEditDialog(config)} - className="h-6 w-6 text-muted-foreground hover:text-foreground" + className="h-7 w-7 text-muted-foreground hover:text-foreground" > @@ -297,7 +291,7 @@ export function VisionModelManager({ searchSpaceId }: VisionModelManagerProps) { variant="ghost" size="icon" onClick={() => setConfigToDelete(config)} - className="h-7 w-7 rounded-lg text-muted-foreground hover:text-destructive" + className="h-7 w-7 text-muted-foreground hover:text-destructive" > @@ -310,9 +304,17 @@ export function VisionModelManager({ searchSpaceId }: VisionModelManagerProps) { )}
- {/* Footer: Date + Creator */} -
- +
+ {getProviderIcon(config.provider, { + className: "size-3.5 shrink-0", + })} + + {config.model_name} + +
+ +
+ {new Date(config.created_at).toLocaleDateString(undefined, { year: "numeric", month: "short", @@ -321,11 +323,11 @@ export function VisionModelManager({ searchSpaceId }: VisionModelManagerProps) { {member && ( <> - + -
+
{member.avatarUrl && ( @@ -334,7 +336,7 @@ export function VisionModelManager({ searchSpaceId }: VisionModelManagerProps) { {getInitials(member.name)} - + {member.name}
diff --git a/surfsense_web/components/shared/image-config-dialog.tsx b/surfsense_web/components/shared/image-config-dialog.tsx index 1b94f0a35..2ae53ccca 100644 --- a/surfsense_web/components/shared/image-config-dialog.tsx +++ b/surfsense_web/components/shared/image-config-dialog.tsx @@ -48,7 +48,6 @@ interface ImageConfigDialogProps { isGlobal: boolean; searchSpaceId: number; mode: "create" | "edit" | "view"; - defaultProvider?: string; } const INITIAL_FORM = { @@ -68,7 +67,6 @@ export function ImageConfigDialog({ isGlobal, searchSpaceId, mode, - defaultProvider, }: ImageConfigDialogProps) { const [isSubmitting, setIsSubmitting] = useState(false); const [formData, setFormData] = useState(INITIAL_FORM); @@ -89,11 +87,11 @@ export function ImageConfigDialog({ api_version: config.api_version || "", }); } else if (mode === "create") { - setFormData({ ...INITIAL_FORM, provider: defaultProvider ?? "" }); + setFormData(INITIAL_FORM); } setScrollPos("top"); } - }, [open, mode, config, isGlobal, defaultProvider]); + }, [open, mode, config, isGlobal]); const { mutateAsync: createConfig } = useAtomValue(createImageGenConfigMutationAtom); const { mutateAsync: updateConfig } = useAtomValue(updateImageGenConfigMutationAtom); diff --git a/surfsense_web/components/shared/model-config-dialog.tsx b/surfsense_web/components/shared/model-config-dialog.tsx index 1a3c8e4a0..4d2373b49 100644 --- a/surfsense_web/components/shared/model-config-dialog.tsx +++ b/surfsense_web/components/shared/model-config-dialog.tsx @@ -28,7 +28,6 @@ interface ModelConfigDialogProps { isGlobal: boolean; searchSpaceId: number; mode: "create" | "edit" | "view"; - defaultProvider?: string; } export function ModelConfigDialog({ @@ -38,7 +37,6 @@ export function ModelConfigDialog({ isGlobal, searchSpaceId, mode, - defaultProvider, }: ModelConfigDialogProps) { const [isSubmitting, setIsSubmitting] = useState(false); const [scrollPos, setScrollPos] = useState<"top" | "middle" | "bottom">("top"); @@ -196,12 +194,10 @@ export function ModelConfigDialog({ {mode === "create" ? ( ) : isGlobal && config ? (
diff --git a/surfsense_web/components/shared/vision-config-dialog.tsx b/surfsense_web/components/shared/vision-config-dialog.tsx index 7f96b0594..7332d3dcd 100644 --- a/surfsense_web/components/shared/vision-config-dialog.tsx +++ b/surfsense_web/components/shared/vision-config-dialog.tsx @@ -49,7 +49,6 @@ interface VisionConfigDialogProps { isGlobal: boolean; searchSpaceId: number; mode: "create" | "edit" | "view"; - defaultProvider?: string; } const INITIAL_FORM = { @@ -69,7 +68,6 @@ export function VisionConfigDialog({ isGlobal, searchSpaceId, mode, - defaultProvider, }: VisionConfigDialogProps) { const [isSubmitting, setIsSubmitting] = useState(false); const [formData, setFormData] = useState(INITIAL_FORM); @@ -89,11 +87,11 @@ export function VisionConfigDialog({ api_version: (config as VisionLLMConfig).api_version || "", }); } else if (mode === "create") { - setFormData({ ...INITIAL_FORM, provider: defaultProvider ?? "" }); + setFormData(INITIAL_FORM); } setScrollPos("top"); } - }, [open, mode, config, isGlobal, defaultProvider]); + }, [open, mode, config, isGlobal]); const { mutateAsync: createConfig } = useAtomValue(createVisionLLMConfigMutationAtom); const { mutateAsync: updateConfig } = useAtomValue(updateVisionLLMConfigMutationAtom); diff --git a/surfsense_web/components/tool-ui/confluence/create-confluence-page.tsx b/surfsense_web/components/tool-ui/confluence/create-confluence-page.tsx index f8cd6ee15..d22f317b8 100644 --- a/surfsense_web/components/tool-ui/confluence/create-confluence-page.tsx +++ b/surfsense_web/components/tool-ui/confluence/create-confluence-page.tsx @@ -15,8 +15,6 @@ import { SelectTrigger, SelectValue, } from "@/components/ui/select"; -import { isInterruptResult, useHitlDecision } from "@/lib/hitl"; -import type { InterruptResult, HitlDecision } from "@/lib/hitl"; import { useHitlPhase } from "@/hooks/use-hitl-phase"; interface ConfluenceAccount { @@ -32,10 +30,24 @@ interface ConfluenceSpace { name: string; } -type CreateConfluencePageInterruptContext = { - accounts?: ConfluenceAccount[]; - spaces?: ConfluenceSpace[]; - error?: string; +interface InterruptResult { + __interrupt__: true; + __decided__?: "approve" | "reject" | "edit"; + __completed__?: boolean; + action_requests: Array<{ + name: string; + args: Record; + }>; + review_configs: Array<{ + action_name: string; + allowed_decisions: Array<"approve" | "edit" | "reject">; + }>; + interrupt_type?: string; + context?: { + accounts?: ConfluenceAccount[]; + spaces?: ConfluenceSpace[]; + error?: string; + }; } interface SuccessResult { @@ -64,12 +76,21 @@ interface InsufficientPermissionsResult { } type CreateConfluencePageResult = - | InterruptResult + | InterruptResult | SuccessResult | ErrorResult | AuthErrorResult | InsufficientPermissionsResult; +function isInterruptResult(result: unknown): result is InterruptResult { + return ( + typeof result === "object" && + result !== null && + "__interrupt__" in result && + (result as InterruptResult).__interrupt__ === true + ); +} + function isErrorResult(result: unknown): result is ErrorResult { return ( typeof result === "object" && @@ -103,8 +124,12 @@ function ApprovalCard({ onDecision, }: { args: { title: string; content?: string; space_id?: string }; - interruptData: InterruptResult; - onDecision: (decision: HitlDecision) => void; + interruptData: InterruptResult; + onDecision: (decision: { + type: "approve" | "reject" | "edit"; + message?: string; + edited_action?: { name: string; args: Record }; + }) => void; }) { const { phase, setProcessing, setRejected } = useHitlPhase(interruptData); const [isPanelOpen, setIsPanelOpen] = useState(false); @@ -439,16 +464,18 @@ export const CreateConfluencePageToolUI = ({ { title: string; content?: string; space_id?: string }, CreateConfluencePageResult >) => { - const { dispatch } = useHitlDecision(); - if (!result) return null; if (isInterruptResult(result)) { return ( } - onDecision={(decision) => dispatch([decision])} + interruptData={result} + onDecision={(decision) => { + window.dispatchEvent( + new CustomEvent("hitl-decision", { detail: { decisions: [decision] } }) + ); + }} /> ); } diff --git a/surfsense_web/components/tool-ui/confluence/delete-confluence-page.tsx b/surfsense_web/components/tool-ui/confluence/delete-confluence-page.tsx index 37d73377e..258f259f0 100644 --- a/surfsense_web/components/tool-ui/confluence/delete-confluence-page.tsx +++ b/surfsense_web/components/tool-ui/confluence/delete-confluence-page.tsx @@ -6,26 +6,38 @@ import { useCallback, useEffect, useState } from "react"; import { TextShimmerLoader } from "@/components/prompt-kit/loader"; import { Button } from "@/components/ui/button"; import { Checkbox } from "@/components/ui/checkbox"; -import { isInterruptResult, useHitlDecision } from "@/lib/hitl"; -import type { InterruptResult, HitlDecision } from "@/lib/hitl"; import { useHitlPhase } from "@/hooks/use-hitl-phase"; -type DeleteConfluencePageInterruptContext = { - account?: { - id: number; +interface InterruptResult { + __interrupt__: true; + __decided__?: "approve" | "reject"; + __completed__?: boolean; + action_requests: Array<{ name: string; - base_url: string; - auth_expired?: boolean; + args: Record; + }>; + review_configs: Array<{ + action_name: string; + allowed_decisions: Array<"approve" | "reject">; + }>; + interrupt_type?: string; + context?: { + account?: { + id: number; + name: string; + base_url: string; + auth_expired?: boolean; + }; + page?: { + page_id: string; + page_title: string; + space_id: string; + connector_id?: number; + document_id?: number; + indexed_at?: string; + }; + error?: string; }; - page?: { - page_id: string; - page_title: string; - space_id: string; - connector_id?: number; - document_id?: number; - indexed_at?: string; - }; - error?: string; } interface SuccessResult { @@ -65,7 +77,7 @@ interface InsufficientPermissionsResult { } type DeleteConfluencePageResult = - | InterruptResult + | InterruptResult | SuccessResult | ErrorResult | NotFoundResult @@ -73,6 +85,15 @@ type DeleteConfluencePageResult = | AuthErrorResult | InsufficientPermissionsResult; +function isInterruptResult(result: unknown): result is InterruptResult { + return ( + typeof result === "object" && + result !== null && + "__interrupt__" in result && + (result as InterruptResult).__interrupt__ === true + ); +} + function isErrorResult(result: unknown): result is ErrorResult { return ( typeof result === "object" && @@ -124,8 +145,12 @@ function ApprovalCard({ interruptData, onDecision, }: { - interruptData: InterruptResult; - onDecision: (decision: HitlDecision) => void; + interruptData: InterruptResult; + onDecision: (decision: { + type: "approve" | "reject"; + message?: string; + edited_action?: { name: string; args: Record }; + }) => void; }) { const { phase, setProcessing, setRejected } = useHitlPhase(interruptData); const [deleteFromKb, setDeleteFromKb] = useState(false); @@ -377,15 +402,18 @@ export const DeleteConfluencePageToolUI = ({ { page_title_or_id: string; delete_from_kb?: boolean }, DeleteConfluencePageResult >) => { - const { dispatch } = useHitlDecision(); - if (!result) return null; if (isInterruptResult(result)) { return ( } - onDecision={(decision) => dispatch([decision])} + interruptData={result} + onDecision={(decision) => { + const event = new CustomEvent("hitl-decision", { + detail: { decisions: [decision] }, + }); + window.dispatchEvent(event); + }} /> ); } diff --git a/surfsense_web/components/tool-ui/confluence/update-confluence-page.tsx b/surfsense_web/components/tool-ui/confluence/update-confluence-page.tsx index 76df10e40..8bc0772a4 100644 --- a/surfsense_web/components/tool-ui/confluence/update-confluence-page.tsx +++ b/surfsense_web/components/tool-ui/confluence/update-confluence-page.tsx @@ -8,27 +8,39 @@ import { openHitlEditPanelAtom } from "@/atoms/chat/hitl-edit-panel.atom"; import { PlateEditor } from "@/components/editor/plate-editor"; import { TextShimmerLoader } from "@/components/prompt-kit/loader"; import { Button } from "@/components/ui/button"; -import { isInterruptResult, useHitlDecision } from "@/lib/hitl"; -import type { InterruptResult, HitlDecision } from "@/lib/hitl"; import { useHitlPhase } from "@/hooks/use-hitl-phase"; -type UpdateConfluencePageInterruptContext = { - account?: { - id: number; +interface InterruptResult { + __interrupt__: true; + __decided__?: "approve" | "reject" | "edit"; + __completed__?: boolean; + action_requests: Array<{ name: string; - base_url: string; - auth_expired?: boolean; + args: Record; + }>; + review_configs: Array<{ + action_name: string; + allowed_decisions: Array<"approve" | "edit" | "reject">; + }>; + interrupt_type?: string; + context?: { + account?: { + id: number; + name: string; + base_url: string; + auth_expired?: boolean; + }; + page?: { + page_id: string; + page_title: string; + space_id: string; + body: string; + version: number; + document_id: number; + indexed_at?: string; + }; + error?: string; }; - page?: { - page_id: string; - page_title: string; - space_id: string; - body: string; - version: number; - document_id: number; - indexed_at?: string; - }; - error?: string; } interface SuccessResult { @@ -62,13 +74,22 @@ interface InsufficientPermissionsResult { } type UpdateConfluencePageResult = - | InterruptResult + | InterruptResult | SuccessResult | ErrorResult | NotFoundResult | AuthErrorResult | InsufficientPermissionsResult; +function isInterruptResult(result: unknown): result is InterruptResult { + return ( + typeof result === "object" && + result !== null && + "__interrupt__" in result && + (result as InterruptResult).__interrupt__ === true + ); +} + function isErrorResult(result: unknown): result is ErrorResult { return ( typeof result === "object" && @@ -115,8 +136,12 @@ function ApprovalCard({ new_title?: string; new_content?: string; }; - interruptData: InterruptResult; - onDecision: (decision: HitlDecision) => void; + interruptData: InterruptResult; + onDecision: (decision: { + type: "approve" | "reject" | "edit"; + message?: string; + edited_action?: { name: string; args: Record }; + }) => void; }) { const { phase, setProcessing, setRejected } = useHitlPhase(interruptData); @@ -477,16 +502,18 @@ export const UpdateConfluencePageToolUI = ({ }, UpdateConfluencePageResult >) => { - const { dispatch } = useHitlDecision(); - if (!result) return null; if (isInterruptResult(result)) { return ( } - onDecision={(decision) => dispatch([decision])} + interruptData={result} + onDecision={(decision) => { + window.dispatchEvent( + new CustomEvent("hitl-decision", { detail: { decisions: [decision] } }) + ); + }} /> ); } diff --git a/surfsense_web/components/tool-ui/dropbox/create-file.tsx b/surfsense_web/components/tool-ui/dropbox/create-file.tsx index 15d454c76..ac45f1f5b 100644 --- a/surfsense_web/components/tool-ui/dropbox/create-file.tsx +++ b/surfsense_web/components/tool-ui/dropbox/create-file.tsx @@ -16,8 +16,6 @@ import { SelectValue, } from "@/components/ui/select"; import { useHitlPhase } from "@/hooks/use-hitl-phase"; -import { isInterruptResult, useHitlDecision } from "@/lib/hitl"; -import type { InterruptResult, HitlDecision } from "@/lib/hitl"; interface DropboxAccount { id: number; @@ -31,11 +29,21 @@ interface SupportedType { label: string; } -type DropboxCreateFileContext = { - accounts?: DropboxAccount[]; - parent_folders?: Record>; - supported_types?: SupportedType[]; - error?: string; +interface InterruptResult { + __interrupt__: true; + __decided__?: "approve" | "reject" | "edit"; + __completed__?: boolean; + action_requests: Array<{ name: string; args: Record }>; + review_configs: Array<{ + action_name: string; + allowed_decisions: Array<"approve" | "edit" | "reject">; + }>; + context?: { + accounts?: DropboxAccount[]; + parent_folders?: Record>; + supported_types?: SupportedType[]; + error?: string; + }; } interface SuccessResult { @@ -57,7 +65,16 @@ interface AuthErrorResult { connector_type?: string; } -type CreateDropboxFileResult = InterruptResult | SuccessResult | ErrorResult | AuthErrorResult; +type CreateDropboxFileResult = InterruptResult | SuccessResult | ErrorResult | AuthErrorResult; + +function isInterruptResult(result: unknown): result is InterruptResult { + return ( + typeof result === "object" && + result !== null && + "__interrupt__" in result && + (result as InterruptResult).__interrupt__ === true + ); +} function isErrorResult(result: unknown): result is ErrorResult { return ( @@ -83,8 +100,12 @@ function ApprovalCard({ onDecision, }: { args: { name: string; file_type?: string; content?: string }; - interruptData: InterruptResult; - onDecision: (decision: HitlDecision) => void; + interruptData: InterruptResult; + onDecision: (decision: { + type: "approve" | "reject" | "edit"; + message?: string; + edited_action?: { name: string; args: Record }; + }) => void; }) { const { phase, setProcessing, setRejected } = useHitlPhase(interruptData); const [isPanelOpen, setIsPanelOpen] = useState(false); @@ -434,14 +455,17 @@ export const CreateDropboxFileToolUI = ({ { name: string; file_type?: string; content?: string }, CreateDropboxFileResult >) => { - const { dispatch } = useHitlDecision(); if (!result) return null; if (isInterruptResult(result)) { return ( } - onDecision={(decision) => dispatch([decision])} + interruptData={result} + onDecision={(decision) => { + window.dispatchEvent( + new CustomEvent("hitl-decision", { detail: { decisions: [decision] } }) + ); + }} /> ); } diff --git a/surfsense_web/components/tool-ui/dropbox/trash-file.tsx b/surfsense_web/components/tool-ui/dropbox/trash-file.tsx index a2fadd20f..0b38777c1 100644 --- a/surfsense_web/components/tool-ui/dropbox/trash-file.tsx +++ b/surfsense_web/components/tool-ui/dropbox/trash-file.tsx @@ -7,8 +7,6 @@ import { TextShimmerLoader } from "@/components/prompt-kit/loader"; import { Button } from "@/components/ui/button"; import { Checkbox } from "@/components/ui/checkbox"; import { useHitlPhase } from "@/hooks/use-hitl-phase"; -import { isInterruptResult, useHitlDecision } from "@/lib/hitl"; -import type { InterruptResult, HitlDecision } from "@/lib/hitl"; interface DropboxAccount { id: number; @@ -24,10 +22,13 @@ interface DropboxFile { document_id?: number; } -type DropboxTrashFileContext = { - account?: DropboxAccount; - file?: DropboxFile; - error?: string; +interface InterruptResult { + __interrupt__: true; + __decided__?: "approve" | "reject"; + __completed__?: boolean; + action_requests: Array<{ name: string; args: Record }>; + review_configs: Array<{ action_name: string; allowed_decisions: Array<"approve" | "reject"> }>; + context?: { account?: DropboxAccount; file?: DropboxFile; error?: string }; } interface SuccessResult { @@ -51,12 +52,20 @@ interface AuthErrorResult { } type DeleteDropboxFileResult = - | InterruptResult + | InterruptResult | SuccessResult | ErrorResult | NotFoundResult | AuthErrorResult; +function isInterruptResult(result: unknown): result is InterruptResult { + return ( + typeof result === "object" && + result !== null && + "__interrupt__" in result && + (result as InterruptResult).__interrupt__ === true + ); +} function isErrorResult(result: unknown): result is ErrorResult { return ( typeof result === "object" && @@ -86,8 +95,12 @@ function ApprovalCard({ interruptData, onDecision, }: { - interruptData: InterruptResult; - onDecision: (decision: HitlDecision) => void; + interruptData: InterruptResult; + onDecision: (decision: { + type: "approve" | "reject"; + message?: string; + edited_action?: { name: string; args: Record }; + }) => void; }) { const { phase, setProcessing, setRejected } = useHitlPhase(interruptData); const [deleteFromKb, setDeleteFromKb] = useState(false); @@ -295,13 +308,16 @@ export const DeleteDropboxFileToolUI = ({ { file_name: string; delete_from_kb?: boolean }, DeleteDropboxFileResult >) => { - const { dispatch } = useHitlDecision(); if (!result) return null; if (isInterruptResult(result)) { return ( } - onDecision={(decision) => dispatch([decision])} + interruptData={result} + onDecision={(decision) => { + window.dispatchEvent( + new CustomEvent("hitl-decision", { detail: { decisions: [decision] } }) + ); + }} /> ); } diff --git a/surfsense_web/components/tool-ui/generic-hitl-approval.tsx b/surfsense_web/components/tool-ui/generic-hitl-approval.tsx deleted file mode 100644 index 2a26b18f7..000000000 --- a/surfsense_web/components/tool-ui/generic-hitl-approval.tsx +++ /dev/null @@ -1,264 +0,0 @@ -"use client"; - -import type { ToolCallMessagePartComponent } from "@assistant-ui/react"; -import { CornerDownLeftIcon, Pen } from "lucide-react"; -import { useCallback, useEffect, useMemo, useState } from "react"; -import { TextShimmerLoader } from "@/components/prompt-kit/loader"; -import { Button } from "@/components/ui/button"; -import { Input } from "@/components/ui/input"; -import { Textarea } from "@/components/ui/textarea"; -import { useHitlPhase } from "@/hooks/use-hitl-phase"; -import { connectorsApiService } from "@/lib/apis/connectors-api.service"; -import { isInterruptResult, useHitlDecision } from "@/lib/hitl"; -import type { HitlDecision, InterruptResult } from "@/lib/hitl"; - -function ParamEditor({ - params, - onChange, - disabled, -}: { - params: Record; - onChange: (updated: Record) => void; - disabled: boolean; -}) { - const entries = Object.entries(params); - if (entries.length === 0) return null; - - return ( -
- {entries.map(([key, value]) => { - const strValue = value == null ? "" : String(value); - const isLong = strValue.length > 120; - const fieldId = `hitl-param-${key}`; - - return ( -
- - {isLong ? ( -