fix (memory): replaced json laods with extractor

Author: itsskofficial

src/client/app/api/memories/route.js

@@ -45,4 +45,4 @@ export const POST = withAuth(async function POST(request, { authHeader }) {
 		console.error("API Error in /api/memories (POST):", error)
 		return NextResponse.json({ error: error.message }, { status: 500 })
 	}
-})
+})

src/server/main/auth/utils.py

@@ -14,6 +14,7 @@
 from jose.exceptions import JOSEError
 from fastapi import HTTPException, status, Depends, WebSocket, WebSocketDisconnect
 from fastapi.security import OAuth2PasswordBearer
+from json_extractor import JsonExtractor
 
 from main.config import (
     ENVIRONMENT, SELF_HOST_AUTH_SECRET,
@@ -122,9 +123,9 @@ async def get_decoded_payload_with_claims(self, token: str = Depends(oauth2_sche
     async def ws_authenticate_with_data(self, websocket: WebSocket) -> Optional[Dict]:
         try:
             auth_message_str = await websocket.receive_text()
-            auth_message = json.loads(auth_message_str)
+            auth_message = JsonExtractor.extract_valid_json(auth_message_str)
 
-            if auth_message.get("type") != "auth" or not auth_message.get("token"):
+            if not auth_message or not isinstance(auth_message, dict) or auth_message.get("type") != "auth" or not auth_message.get("token"):
                 await websocket.send_json({"type": "auth_failure", "message": "Invalid auth message format."})
                 await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
                 return None
@@ -149,10 +150,6 @@ async def ws_authenticate_with_data(self, websocket: WebSocket) -> Optional[Dict
         except WebSocketDisconnect:
             print(f"[{datetime.datetime.now()}] [WS_AUTH] WebSocket disconnected during auth.")
             return None
-        except json.JSONDecodeError:
-            await websocket.send_json({"type": "auth_failure", "message": "Auth message must be JSON."})
-            await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
-            return None
         except HTTPException as e:
             await websocket.send_json({"type": "auth_failure", "message": f"Token validation failed: {e.detail}"})
             await websocket.close(code=status.WS_1008_POLICY_VIOLATION)

src/server/main/chat/utils.py

@@ -41,7 +41,9 @@ def call(self, params: str, **kwargs) -> str:
             if isinstance(params, dict):
                  parsed_params = params
             else:
-                 parsed_params = json.loads(params)
+                 parsed_params = JsonExtractor.extract_valid_json(params)
+            if not parsed_params:
+                return json.dumps({"status": "failure", "error": "Invalid JSON in params."})
             json_string_to_validate = parsed_params.get('json_string', '')
             if not json_string_to_validate:
                 return json.dumps({"status": "failure", "error": "Input json_string is empty."})
@@ -339,12 +341,16 @@ def worker():
 def msg_to_str(msg: Dict[str, Any]) -> str:
     if msg.get('role') == 'assistant' and msg.get('function_call'):
         args_str = msg['function_call'].get('arguments', '')
-        try: args_pretty = json.dumps(json.loads(args_str), indent=2)
+        try:
+            parsed_args = JsonExtractor.extract_valid_json(args_str)
+            args_pretty = json.dumps(parsed_args, indent=2) if parsed_args else args_str
         except: args_pretty = args_str
         return f"<tool_code name=\"{msg['function_call'].get('name')}\">\n{args_pretty}\n</tool_code>\n"
     elif msg.get('role') == 'function':
         content = msg.get('content', '')
-        try: content_pretty = json.dumps(json.loads(content), indent=2)
+        try:
+            parsed_content = JsonExtractor.extract_valid_json(content)
+            content_pretty = json.dumps(parsed_content, indent=2) if parsed_content else content
         except: content_pretty = content
         return f"<tool_result tool_name=\"{msg.get('name')}\">\n{content_pretty}\n</tool_result>\n"
     elif msg.get('role') == 'assistant' and msg.get('content'):

src/server/main/integrations/utils.py

@@ -4,6 +4,7 @@
 
 from main.db import MongoManager
 from main.auth.utils import aes_encrypt, aes_decrypt
+from json_extractor import JsonExtractor
 
 async def store_encrypted_integration_token(user_id: str, service_name: str, token_data: Dict[str, Any], db_manager: MongoManager) -> bool:
     """Encrypts and stores the entire token object for a service."""
@@ -40,7 +41,7 @@ async def get_decrypted_integration_token(user_id: str, service_name: str, db_ma
 
     try:
         decrypted_token_str = aes_decrypt(integration_data["encrypted_token"])
-        return json.loads(decrypted_token_str)
+        return JsonExtractor.extract_valid_json(decrypted_token_str)
     except Exception:
         # Handle decryption or JSON parsing errors
         return None

src/server/main/memories/routes.py

@@ -24,7 +24,7 @@ async def startup_event():
 
 @router.get("/", summary="Get all memories for a user")
 async def get_all_memories(
-    user_id: str = Depends(PermissionChecker(required_permissions=["read:profile"]))
+    user_id: str = Depends(PermissionChecker(required_permissions=["read:memory"]))
 ):
     pool = None
     try:
@@ -55,7 +55,7 @@ async def get_all_memories(
 
 @router.get("/graph", summary="Get memory graph data for a user")
 async def get_memory_graph(
-    user_id: str = Depends(PermissionChecker(required_permissions=["read:profile"]))
+    user_id: str = Depends(PermissionChecker(required_permissions=["read:memory"]))
 ):
     try:
         graph_data = await utils.create_memory_graph(user_id)
@@ -67,7 +67,7 @@ async def get_memory_graph(
 @router.post("/", summary="Create a new memory for a user")
 async def create_memory(
     request: CreateMemoryRequest,
-    user_id: str = Depends(PermissionChecker(required_permissions=["write:profile"]))
+    user_id: str = Depends(PermissionChecker(required_permissions=["write:memory"]))
 ):
     try:
         result_message = await utils.create_memory(user_id, request.content, request.source)
@@ -80,7 +80,7 @@ async def create_memory(
 async def update_memory(
     memory_id: int,
     request: UpdateMemoryRequest,
-    user_id: str = Depends(PermissionChecker(required_permissions=["write:profile"]))
+    user_id: str = Depends(PermissionChecker(required_permissions=["write:memory"]))
 ):
     try:
         result_message = await utils.update_memory(user_id, memory_id, request.content)
@@ -94,7 +94,7 @@ async def update_memory(
 @router.delete("/{memory_id}", summary="Delete a memory")
 async def delete_memory(
     memory_id: int,
-    user_id: str = Depends(PermissionChecker(required_permissions=["write:profile"]))
+    user_id: str = Depends(PermissionChecker(required_permissions=["write:memory"]))
 ):
     try:
         result_message = await utils.delete_memory(user_id, memory_id)

src/server/main/memories/utils.py

@@ -8,6 +8,7 @@
 
 import google.generativeai as genai
 from pgvector.asyncpg import register_vector
+from json_extractor import JsonExtractor
 
 from . import db, llm
 from .prompts import fact_analysis_user_prompt_template
@@ -110,9 +111,8 @@ async def create_memory(user_id: str, content: str, source: Optional[str] = "man
         prompt = fact_analysis_user_prompt_template.format(text=content)
         analysis_raw = llm.run_agent_with_prompt(agents["fact_analysis"], prompt)
         analysis_cleaned = clean_llm_output(analysis_raw)
-        try:
-            analysis = json.loads(analysis_cleaned)
-        except json.JSONDecodeError:
+        analysis = JsonExtractor.extract_valid_json(analysis_cleaned)
+        if not analysis:
             logger.error(f"Create memory failed due to analysis JSON error. Output: {analysis_cleaned}")
             raise ValueError("Failed to analyze memory content.")
         return await _insert_fact_with_analysis(conn, user_id, content, source, analysis)
@@ -130,9 +130,8 @@ async def update_memory(user_id: str, memory_id: int, new_content: str) -> str:
         prompt = fact_analysis_user_prompt_template.format(text=new_content)
         analysis_raw = llm.run_agent_with_prompt(agents["fact_analysis"], prompt)
         analysis_cleaned = clean_llm_output(analysis_raw)
-        try:
-            analysis = json.loads(analysis_cleaned)
-        except json.JSONDecodeError:
+        analysis = JsonExtractor.extract_valid_json(analysis_cleaned)
+        if not analysis:
             raise ValueError("Failed to analyze updated memory content.")
 
         new_embedding = _get_normalized_embedding(new_content, task_type="RETRIEVAL_DOCUMENT")

src/server/main/tasks/db.py

@@ -2,6 +2,7 @@
 import datetime
 import uuid
 import json
+from json_extractor import JsonExtractor
 from typing import Dict, List, Optional
 import logging
 
@@ -32,10 +33,7 @@ async def create_task(self, user_id: str, task_data: dict) -> dict:
 
         schedule = task_data.get("schedule")
         if isinstance(schedule, str):
-            try:
-                schedule = json.loads(schedule)
-            except json.JSONDecodeError:
-                schedule = None
+            schedule = JsonExtractor.extract_valid_json(schedule)
 
         task_doc = {
             "task_id": task_id,

src/server/mcp_hub/gcal/auth.py

@@ -10,6 +10,7 @@
 from googleapiclient.discovery import build, Resource
 from fastmcp import Context
 from fastmcp.exceptions import ToolError
+from json_extractor import JsonExtractor
 
 from typing import Optional, Dict, Any
 from dotenv import load_dotenv
@@ -94,7 +95,9 @@ async def get_google_creds(user_id: str) -> Credentials:
 
     try:
         decrypted_creds_str = aes_decrypt(gcal_data["credentials"])
-        token_info = json.loads(decrypted_creds_str)
+        token_info = JsonExtractor.extract_valid_json(decrypted_creds_str)
+        if not token_info:
+            raise ToolError("Failed to parse decrypted credentials for Google Calendar.")
         return Credentials.from_authorized_user_info(token_info)
     except Exception as e:
         raise ToolError(f"Failed to decrypt or parse default OAuth token for Google Calendar: {e}")

src/server/mcp_hub/gdocs/auth.py

@@ -10,6 +10,7 @@
 from googleapiclient.discovery import build, Resource
 from fastmcp import Context
 from fastmcp.exceptions import ToolError
+from json_extractor import JsonExtractor
 
 from typing import Optional
 from dotenv import load_dotenv
@@ -65,7 +66,9 @@ async def get_google_creds(user_id: str) -> Credentials:
 
     try:
         decrypted_creds_str = aes_decrypt(gdocs_data["credentials"])
-        token_info = json.loads(decrypted_creds_str)
+        token_info = JsonExtractor.extract_valid_json(decrypted_creds_str)
+        if not token_info:
+            raise ToolError("Failed to parse decrypted credentials for Google Docs.")
         return Credentials.from_authorized_user_info(token_info)
     except Exception as e:
         raise ToolError(f"Failed to decrypt or parse default OAuth token for Google Docs: {e}")

src/server/mcp_hub/gdrive/auth.py

@@ -10,6 +10,7 @@
 from googleapiclient.discovery import build, Resource
 from fastmcp import Context
 from fastmcp.exceptions import ToolError
+from json_extractor import JsonExtractor
 
 from typing import Optional
 from dotenv import load_dotenv
@@ -71,7 +72,9 @@ async def get_google_creds(user_id: str) -> Credentials:
 
     try:
         decrypted_creds_str = aes_decrypt(gdrive_data["credentials"])
-        token_info = json.loads(decrypted_creds_str)
+        token_info = JsonExtractor.extract_valid_json(decrypted_creds_str)
+        if not token_info:
+            raise ToolError("Failed to parse decrypted credentials for Google Drive.")
         return Credentials.from_authorized_user_info(token_info)
     except Exception as e:
         raise ToolError(f"Failed to decrypt or parse default OAuth token for Google Drive: {e}")