Implement workspace isolation, active MCP tools, and Jina search/read tools

backend.py

@@ -27,6 +27,10 @@ import collections
 from pathlib import Path
 from typing import AsyncIterator, Optional, List, Dict, Any
 from pydantic import BaseModel
+import contextvars
+import urllib.parse
+
+workspace_var = contextvars.ContextVar("workspace_dir", default="/tmp/workspace")
 
 # --- Ultimate Agent Brain Imports ---
 from second_brain import SecondBrainWrapper
@@ -36,7 +40,8 @@ from helix_state import helix_db
 from context_engine import ContextEngine
 
 # Instantiate singletons for the orchestrator
-brain = SecondBrainWrapper(space_name="space2-cerebrum")
+SPACE_NAME = os.environ.get("SPACE_NAME", "space2-cerebrum")
+brain = SecondBrainWrapper(space_name=SPACE_NAME)
 context_engine = ContextEngine(brain)
 watchdog = SurvivalWatchdog()
 
@@ -244,6 +249,40 @@ TOOLS = [
             }
         }
     },
+    {
+        "type": "function",
+        "function": {
+            "name": "web_search",
+            "description": "Search the web for up-to-date information, news, papers, or documentation.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "query": {
+                        "type": "string",
+                        "description": "The search query (be specific and detailed)"
+                    }
+                },
+                "required": ["query"]
+            }
+        }
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "web_read",
+            "description": "Read the clean markdown content of any webpage URL to get detailed context, articles, documentation, or code.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "url": {
+                        "type": "string",
+                        "description": "The absolute URL of the webpage to read"
+                    }
+                },
+                "required": ["url"]
+            }
+        }
+    },
 ]
 
 # ---------------------------------------------------------------------------
@@ -252,7 +291,7 @@ TOOLS = [
 
 def _safe_path(rel_path: str) -> Path:
     """Resolve a relative path safely within the workspace."""
-    workspace = Path(WORKSPACE_DIR).resolve()
+    workspace = Path(workspace_var.get()).resolve()
     target = (workspace / rel_path).resolve()
     # Prevent path traversal
     if not str(target).startswith(str(workspace)):
@@ -301,7 +340,174 @@ def repair_arguments(func_name: str, args: dict) -> tuple[dict, list[str]]:
     return repaired_args, notes
 
 
-async def execute_tool(name: str, arguments: dict) -> str:
+# ---------------------------------------------------------------------------
+# Search and MCP Helpers
+# ---------------------------------------------------------------------------
+
+def sanitize_function_name(name: str) -> str:
+    sanitized = re.sub(r'[^a-zA-Z0-9_\.\-]', '_', name)
+    if not re.match(r'^[a-zA-Z_]', sanitized):
+        sanitized = '_' + sanitized
+    if len(sanitized) > 124:
+        sanitized = sanitized[:124]
+    return sanitized
+
+def create_mcp_tool_id(server_name: str, tool_name: str) -> str:
+    san_server = sanitize_function_name(server_name)
+    san_tool = sanitize_function_name(tool_name)
+    max_len = 124
+    sep = "_"
+    if len(san_server) + len(san_tool) + len(sep) > max_len:
+        total = len(san_server) + len(san_tool)
+        server_portion = int((len(san_server) / total) * (max_len - len(sep)))
+        tool_portion = max_len - len(sep) - server_portion
+        return f"{san_server[:server_portion]}{sep}{san_tool[:tool_portion]}"
+    return f"{san_server}{sep}{san_tool}"
+
+async def get_active_mcp_tools():
+    """
+    Queries `mcp_server` table and returns:
+      1. List of OpenAI function definitions to append to dynamic tools list.
+      2. Dictionary mapping `mcp_tool_id` to `(server_name, original_tool_name, config)`.
+    """
+    mcp_tools_list = []
+    mcp_mapping = {}
+    if not db_pool:
+        return mcp_tools_list, mcp_mapping
+    try:
+        async with db_pool.acquire() as conn:
+            rows = await conn.fetch("SELECT name, config, tool_info FROM mcp_server WHERE enabled = true")
+            for row in rows:
+                server_name = row["name"]
+                config_raw = row["config"]
+                if isinstance(config_raw, str):
+                    config = json.loads(config_raw)
+                else:
+                    config = config_raw
+                tool_info_raw = row["tool_info"]
+                if not tool_info_raw:
+                    continue
+                if isinstance(tool_info_raw, str):
+                    tool_info = json.loads(tool_info_raw)
+                else:
+                    tool_info = tool_info_raw
+                for tool in tool_info:
+                    tool_name = tool.get("name")
+                    description = tool.get("description", "")
+                    input_schema = tool.get("inputSchema", {})
+                    tool_id = create_mcp_tool_id(server_name, tool_name)
+                    mcp_mapping[tool_id] = (server_name, tool_name, config)
+                    mcp_tools_list.append({
+                        "type": "function",
+                        "function": {
+                            "name": tool_id,
+                            "description": f"[from MCP server: {server_name}] {description}",
+                            "parameters": input_schema
+                        }
+                    })
+    except Exception as e:
+        log_activity(f"[MCP Database Query Warning] Failed to load MCP tools: {e}")
+    return mcp_tools_list, mcp_mapping
+
+async def execute_web_search(query: str) -> str:
+    tavily_key = os.environ.get("TAVILY_API_KEY", "")
+    if tavily_key:
+        try:
+            async with httpx.AsyncClient(timeout=15.0) as client:
+                r = await client.post("https://api.tavily.com/search", json={
+                    "api_key": tavily_key,
+                    "query": query,
+                    "search_depth": "basic",
+                    "max_results": 5
+                })
+                if r.status_code == 200:
+                    data = r.json()
+                    results = []
+                    for item in data.get("results", []):
+                        results.append(f"Title: {item.get('title')}\nURL: {item.get('url')}\nContent: {item.get('content')}\n---")
+                    return "\n".join(results) if results else "No results found."
+        except Exception as e:
+            log_activity(f"[Tavily Error] {e}")
+
+    exa_key = os.environ.get("EXA_API_KEY", "")
+    if exa_key:
+        try:
+            async with httpx.AsyncClient(timeout=15.0) as client:
+                r = await client.post("https://api.exa.ai/search", headers={
+                    "x-api-key": exa_key,
+                    "Content-Type": "application/json"
+                }, json={
+                    "query": query,
+                    "numResults": 5,
+                    "text": True
+                })
+                if r.status_code == 200:
+                    data = r.json()
+                    results = []
+                    for item in data.get("results", []):
+                        results.append(f"Title: {item.get('title')}\nURL: {item.get('url')}\nContent: {item.get('text', '')[:2000]}\n---")
+                    return "\n".join(results) if results else "No results found."
+        except Exception as e:
+            log_activity(f"[Exa Error] {e}")
+
+    try:
+        escaped_query = urllib.parse.quote(query)
+        async with httpx.AsyncClient(headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64)"}, timeout=10.0) as client:
+            r = await client.get(f"https://html.duckduckgo.com/html/?q={escaped_query}")
+            if r.status_code == 200:
+                try:
+                    from bs4 import BeautifulSoup
+                    soup = BeautifulSoup(r.text, 'html.parser')
+                    results = []
+                    for a in soup.find_all('a', class_='result__snippet')[:5]:
+                        title_el = a.find_previous('a', class_='result__url')
+                        title = title_el.text.strip() if title_el else "No Title"
+                        url = title_el['href'] if title_el and 'href' in title_el.attrs else ""
+                        snippet = a.text.strip()
+                        results.append(f"Title: {title}\nURL: {url}\nSnippet: {snippet}\n---")
+                    return "\n".join(results) if results else "No results found."
+                except Exception:
+                    snippets = re.findall(r'<a class="result__snippet"[^>]*>(.*?)</a>', r.text, re.DOTALL)
+                    results = []
+                    for s in snippets[:5]:
+                        clean_s = re.sub(r'<[^>]*>', '', s).strip()
+                        results.append(f"Snippet: {clean_s}\n---")
+                    return "\n".join(results) if results else "No results found."
+    except Exception as e:
+        log_activity(f"[DDG Scrape Error] {e}")
+    return "Error: Web search failed. Setup API keys (TAVILY_API_KEY, EXA_API_KEY) for best results."
+
+async def execute_web_read(url: str) -> str:
+    jina_url = f"https://r.jina.ai/{url}"
+    try:
+        async with httpx.AsyncClient(timeout=20.0) as client:
+            r = await client.get(jina_url, headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64)"})
+            if r.status_code == 200:
+                content = r.text
+                if len(content) > 30000:
+                    content = content[:30000] + "\n\n[Truncated - webpage content is extremely long]"
+                return content
+    except Exception as e:
+        log_activity(f"[Jina Reader Error] {e}")
+
+    try:
+        async with httpx.AsyncClient(timeout=15.0) as client:
+            r = await client.get(url, headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64)"})
+            if r.status_code == 200:
+                html = r.text
+                clean = re.sub(r'<script.*?>.*?</script>', '', html, flags=re.DOTALL | re.IGNORECASE)
+                clean = re.sub(r'<style.*?>.*?</style>', '', clean, flags=re.DOTALL | re.IGNORECASE)
+                clean = re.sub(r'<.*?>', ' ', clean)
+                clean = re.sub(r'\s+', ' ', clean).strip()
+                if len(clean) > 15000:
+                    clean = clean[:15000] + "\n\n[Truncated]"
+                return clean
+    except Exception as e:
+        log_activity(f"[Direct Scrape Error] {e}")
+    return f"Error: Failed to read URL {url}."
+
+
+async def execute_tool(name: str, arguments: dict, mcp_mapping: dict = None) -> str:
     """Execute a tool and return its output as a string asynchronously."""
     try:
         if name == "read_file":
@@ -333,7 +539,7 @@ async def execute_tool(name: str, arguments: dict) -> str:
                 command,
                 stdout=asyncio.subprocess.PIPE,
                 stderr=asyncio.subprocess.PIPE,
-                cwd=WORKSPACE_DIR,
+                cwd=workspace_var.get(),
                 env={**os.environ, "HOME": "/tmp", "PATH": os.environ.get("PATH", "/usr/local/bin:/usr/bin:/bin")},
             )
             
@@ -390,7 +596,7 @@ async def execute_tool(name: str, arguments: dict) -> str:
                 "grep", "-rn", "--include=*", pattern, str(path),
                 stdout=asyncio.subprocess.PIPE,
                 stderr=asyncio.subprocess.PIPE,
-                cwd=WORKSPACE_DIR,
+                cwd=workspace_var.get(),
             )
             try:
                 stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=10)
@@ -407,6 +613,40 @@ async def execute_tool(name: str, arguments: dict) -> str:
                 output = output[:10000] + "\n\n[Truncated]"
             return output
 
+        elif name == "web_search":
+            query = arguments.get("query")
+            return await execute_web_search(query)
+
+        elif name == "web_read":
+            url = arguments.get("url")
+            return await execute_web_read(url)
+
+        elif mcp_mapping and name in mcp_mapping:
+            server_name, original_tool_name, config = mcp_mapping[name]
+            try:
+                # Ensure registered on Space 4
+                await apost_json(f"{SPACE4_URL}/api/mcp/register", {
+                    "name": server_name,
+                    "config": config
+                })
+                # Call tool on Space 4
+                call_res = await apost_json(f"{SPACE4_URL}/api/mcp/call", {
+                    "serverName": server_name,
+                    "toolName": original_tool_name,
+                    "arguments": arguments
+                })
+                if isinstance(call_res, dict) and "error" in call_res:
+                    return f"Error from MCP server {server_name}: {call_res['error']}"
+                if isinstance(call_res, dict) and "content" in call_res:
+                    texts = []
+                    for item in call_res["content"]:
+                        if isinstance(item, dict) and item.get("type") == "text":
+                            texts.append(item.get("text", ""))
+                    return "\n".join(texts)
+                return str(call_res)
+            except Exception as e:
+                return f"Error executing MCP tool {name}: {str(e)}"
+
         else:
             return f"Error: Unknown tool: {name}"
 
@@ -708,12 +948,19 @@ async def chat_completions(request: Request, authorization: str = Header(None)):
     if user_msg:
         await save_message(session_id, "user", user_msg.get("content", ""))
 
+    session_folder = session_id[:8] if session_id else "default"
+    active_workspace = os.path.join(WORKSPACE_DIR, session_folder)
+    os.makedirs(active_workspace, exist_ok=True)
+
     if not stream:
         # Non-streaming: simple completion
+        workspace_token = workspace_var.set(active_workspace)
         try:
+            mcp_tools, mcp_mapping = await get_active_mcp_tools()
+            active_tools = list(TOOLS) + mcp_tools
             kwargs = {"model": requested_model, "messages": final_messages}
             if is_agentic:
-                kwargs["tools"] = TOOLS
+                kwargs["tools"] = active_tools
                 kwargs["tool_choice"] = "auto"
             async with completions_semaphore:
                 response = await client.chat.completions.create(**kwargs)
@@ -731,18 +978,23 @@ async def chat_completions(request: Request, authorization: str = Header(None)):
         except Exception as e:
             ACTIVE_SESSIONS.discard(session_id)
             return JSONResponse({"error": {"message": str(e), "type": "internal_error"}}, status_code=500)
+        finally:
+            workspace_var.reset(workspace_token)
 
     # Streaming + agentic loop
     async def generate() -> AsyncIterator[str]:
         nonlocal final_messages
-        async with completions_semaphore:
-            try:
+        workspace_token = workspace_var.set(active_workspace)
+        try:
+            mcp_tools, mcp_mapping = await get_active_mcp_tools()
+            active_tools = list(TOOLS) + mcp_tools
+            async with completions_semaphore:
                 for round_num in range(MAX_TOOL_ROUNDS + 1):
                     # Perform auto-compaction before calling NIM API
                     final_messages = compact_history(final_messages)
                     kwargs = {"model": requested_model, "messages": final_messages, "stream": True}
                     if is_agentic:
-                        kwargs["tools"] = TOOLS
+                        kwargs["tools"] = active_tools
                         kwargs["tool_choice"] = "auto"
 
                     # Collect streamed response
@@ -836,7 +1088,6 @@ async def chat_completions(request: Request, authorization: str = Header(None)):
                         yield make_chunk(request_id, requested_model, f"\n\n🔧 **{func_name}**")
                         if repair_notes:
                             yield make_chunk(request_id, requested_model, " *(Auto-Repaired)*")
-                    
                         if func_name == "run_bash" and "command" in repaired_args:
                             yield make_chunk(request_id, requested_model, f": `{repaired_args['command']}`\n")
                         elif func_name == "read_file" and "path" in repaired_args:
@@ -847,11 +1098,17 @@ async def chat_completions(request: Request, authorization: str = Header(None)):
                             yield make_chunk(request_id, requested_model, f": `{repaired_args.get('path', '.')}`\n")
                         elif func_name == "grep_search":
                             yield make_chunk(request_id, requested_model, f": `{repaired_args.get('pattern', '')}`\n")
+                        elif func_name == "web_search" and "query" in repaired_args:
+                            yield make_chunk(request_id, requested_model, f": `{repaired_args['query']}`\n")
+                        elif func_name == "web_read" and "url" in repaired_args:
+                            yield make_chunk(request_id, requested_model, f": `{repaired_args['url']}`\n")
+                        elif mcp_mapping and func_name in mcp_mapping:
+                            yield make_chunk(request_id, requested_model, f": calling tool\n")
                         else:
                             yield make_chunk(request_id, requested_model, "\n")
 
                         # Execute the tool
-                        result = await execute_tool(func_name, repaired_args)
+                        result = await execute_tool(func_name, repaired_args, mcp_mapping)
 
                         # Append teaching note if repaired
                         if repair_notes:
@@ -877,11 +1134,13 @@ async def chat_completions(request: Request, authorization: str = Header(None)):
                 yield make_chunk(request_id, requested_model, finish_reason="stop")
                 yield "data: [DONE]\n\n"
 
-            except Exception as e:
-                error_msg = f"\n\n❌ Error: {str(e)}"
-                yield make_chunk(request_id, requested_model, error_msg)
-                yield make_chunk(request_id, requested_model, finish_reason="stop")
-                yield "data: [DONE]\n\n"
+        except Exception as e:
+            error_msg = f"\n\n❌ Error: {str(e)}"
+            yield make_chunk(request_id, requested_model, error_msg)
+            yield make_chunk(request_id, requested_model, finish_reason="stop")
+            yield "data: [DONE]\n\n"
+        finally:
+            workspace_var.reset(workspace_token)
 
     return StreamingResponse(
         generate(),
@@ -954,7 +1213,18 @@ async def get_workspace_tree():
             }
             
     try:
-        w_path = Path(WORKSPACE_DIR).resolve()
+        active_project = ""
+        if db_pool:
+            try:
+                async with db_pool.acquire() as conn:
+                    row = await conn.fetchrow("SELECT project_name FROM eternity_projects WHERE is_active = true ORDER BY id DESC LIMIT 1")
+                    if row:
+                        active_project = row["project_name"].replace(" ", "-").lower()
+            except Exception:
+                pass
+        
+        project_folder = active_project if active_project else "default"
+        w_path = (Path(WORKSPACE_DIR) / project_folder).resolve()
         if not w_path.exists():
             w_path.mkdir(parents=True, exist_ok=True)
         return build_tree(w_path, w_path)
@@ -965,11 +1235,27 @@ async def get_workspace_tree():
 @app.get("/api/workspace/file")
 async def get_workspace_file(path: str):
     try:
-        safe_p = _safe_path(path)
-        if not safe_p.exists() or not safe_p.is_file():
-            raise HTTPException(status_code=404, detail="File not found")
-        content = safe_p.read_text(encoding="utf-8", errors="replace")
-        return {"path": path, "content": content}
+        active_project = ""
+        if db_pool:
+            try:
+                async with db_pool.acquire() as conn:
+                    row = await conn.fetchrow("SELECT project_name FROM eternity_projects WHERE is_active = true ORDER BY id DESC LIMIT 1")
+                    if row:
+                        active_project = row["project_name"].replace(" ", "-").lower()
+            except Exception:
+                pass
+        
+        project_folder = active_project if active_project else "default"
+        w_path = (Path(WORKSPACE_DIR) / project_folder).resolve()
+        token = workspace_var.set(str(w_path))
+        try:
+            safe_p = _safe_path(path)
+            if not safe_p.exists() or not safe_p.is_file():
+                raise HTTPException(status_code=404, detail="File not found")
+            content = safe_p.read_text(encoding="utf-8", errors="replace")
+            return {"path": path, "content": content}
+        finally:
+            workspace_var.reset(token)
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
@@ -977,7 +1263,18 @@ async def get_workspace_file(path: str):
 @app.get("/api/workspace/latest-screenshot")
 async def get_latest_screenshot():
     try:
-        w_path = Path(WORKSPACE_DIR).resolve()
+        active_project = ""
+        if db_pool:
+            try:
+                async with db_pool.acquire() as conn:
+                    row = await conn.fetchrow("SELECT project_name FROM eternity_projects WHERE is_active = true ORDER BY id DESC LIMIT 1")
+                    if row:
+                        active_project = row["project_name"].replace(" ", "-").lower()
+            except Exception:
+                pass
+        
+        project_folder = active_project if active_project else "default"
+        w_path = (Path(WORKSPACE_DIR) / project_folder).resolve()
         png_files = []
         for p in w_path.rglob("*.png"):
             if any(part.startswith(".") for part in p.parts):
@@ -1747,7 +2044,21 @@ class ForgeExecuteRequest(BaseModel):
 async def forge_execute(req: ForgeExecuteRequest, authorization: str = Header(None)):
     auth(authorization)
     log_activity(f"[Forge] Received execution request for task {req.task_id}: '{req.prompt[:80]}'")
+    
+    project_folder = "default"
+    if "_" in req.task_id:
+        parts = req.task_id.split("_")
+        if len(parts) >= 2:
+            project_folder = parts[1].replace(" ", "-").lower()
+    active_workspace = os.path.join(WORKSPACE_DIR, project_folder)
+    os.makedirs(active_workspace, exist_ok=True)
+    workspace_token = workspace_var.set(active_workspace)
+
     try:
+        # Load active MCP tools
+        mcp_tools, mcp_mapping = await get_active_mcp_tools()
+        active_tools = list(TOOLS) + mcp_tools
+
         # Inject project_state.md into system context if provided (Karpathy protocol)
         system_content = AGENTIC_SYSTEM_PROMPT + f"\nContext rules: {req.context_rules}"
         if req.project_state_md:
@@ -1770,7 +2081,7 @@ async def forge_execute(req: ForgeExecuteRequest, authorization: str = Header(No
                 response = await nim_client.chat.completions.create(
                     model=RECOMMENDED_MODEL,
                     messages=final_messages,
-                    tools=TOOLS,
+                    tools=active_tools,
                     tool_choice="auto",
                     timeout=8.0
                 )
@@ -1781,7 +2092,7 @@ async def forge_execute(req: ForgeExecuteRequest, authorization: str = Header(No
                         response = await mistral_client.chat.completions.create(
                             model="mistral-large-latest",
                             messages=final_messages,
-                            tools=TOOLS,
+                            tools=active_tools,
                             tool_choice="auto"
                         )
                     else:
@@ -1831,7 +2142,7 @@ async def forge_execute(req: ForgeExecuteRequest, authorization: str = Header(No
                         func_args = json.loads(repaired_str)
                     except Exception:
                         func_args = {}
-                result = await execute_tool(func_name, func_args)
+                result = await execute_tool(func_name, func_args, mcp_mapping)
                 
                 final_messages.append({
                     "role": "tool",
@@ -1859,11 +2170,8 @@ async def forge_execute(req: ForgeExecuteRequest, authorization: str = Header(No
             "status": "error",
             "error": str(e)
         }
-
-
-# ---------------------------------------------------------------------------
-# Database Retention & Heartbeat loops
-# ---------------------------------------------------------------------------
+    finally:
+        workspace_var.reset(workspace_token)
 
 async def db_heartbeat_loop():
     log_activity("Database Heartbeat task started")
@@ -2513,6 +2821,9 @@ async def db_heartbeat_loop():
 
 
 def run_backup_loop():
+    if os.environ.get("SPACE_NAME", "space2-cerebrum") == "space2-cerebrum":
+        log_activity("Local Git Backup Loop disabled on Cerebrum (Space 2)")
+        return
     log_activity("Local Git Backup Loop started")
     while True:
         # Wait 5 minutes between runs