Merge branch 'dev' into feat/fix_palyground_bug

Author: yuan.wang

src/memos/api/handlers/scheduler_handler.py

@@ -9,20 +9,181 @@
 import time
 import traceback
 
+from collections import Counter
+from datetime import datetime, timezone
 from typing import Any
 
 from fastapi import HTTPException
 from fastapi.responses import StreamingResponse
 
 # Imports for new implementation
-from memos.api.product_models import StatusResponse, StatusResponseItem
+from memos.api.product_models import (
+    AllStatusResponse,
+    AllStatusResponseData,
+    StatusResponse,
+    StatusResponseItem,
+    TaskSummary,
+)
 from memos.log import get_logger
+from memos.mem_scheduler.base_scheduler import BaseScheduler
 from memos.mem_scheduler.utils.status_tracker import TaskStatusTracker
 
 
 logger = get_logger(__name__)
 
 
+def handle_scheduler_allstatus(
+    mem_scheduler: BaseScheduler,
+    status_tracker: TaskStatusTracker,
+) -> AllStatusResponse:
+    """
+    Get aggregated scheduler status metrics (no per-task payload).
+
+    Args:
+        mem_scheduler: The BaseScheduler instance.
+        status_tracker: The TaskStatusTracker instance.
+
+    Returns:
+        AllStatusResponse with aggregated status data.
+    """
+
+    def _summarize_tasks(task_details: list[dict[str, Any]]) -> TaskSummary:
+        """Aggregate counts by status for the provided task details (tracker data)."""
+        counter = Counter()
+        for detail in task_details:
+            status = detail.get("status")
+            if status:
+                counter[status] += 1
+
+        total = sum(counter.values())
+        return TaskSummary(
+            waiting=counter.get("waiting", 0),
+            in_progress=counter.get("in_progress", 0),
+            completed=counter.get("completed", 0),
+            pending=counter.get("pending", counter.get("waiting", 0)),
+            failed=counter.get("failed", 0),
+            cancelled=counter.get("cancelled", 0),
+            total=total,
+        )
+
+    def _aggregate_counts_from_redis(
+        tracker: TaskStatusTracker, max_age_seconds: float = 86400
+    ) -> TaskSummary | None:
+        """Stream status counts directly from Redis to avoid loading all task payloads."""
+        redis_client = getattr(tracker, "redis", None)
+        if not redis_client:
+            return None
+
+        counter = Counter()
+        now = datetime.now(timezone.utc).timestamp()
+
+        # Scan task_meta keys, then hscan each hash in batches
+        cursor: int | str = 0
+        while True:
+            cursor, keys = redis_client.scan(cursor=cursor, match="memos:task_meta:*", count=200)
+            for key in keys:
+                h_cursor: int | str = 0
+                while True:
+                    h_cursor, fields = redis_client.hscan(key, cursor=h_cursor, count=500)
+                    for value in fields.values():
+                        try:
+                            payload = json.loads(
+                                value.decode("utf-8") if isinstance(value, bytes) else value
+                            )
+                            # Skip stale entries to reduce noise and load
+                            ts = payload.get("submitted_at") or payload.get("started_at")
+                            if ts:
+                                try:
+                                    ts_dt = datetime.fromisoformat(ts)
+                                    ts_seconds = ts_dt.timestamp()
+                                except Exception:
+                                    ts_seconds = None
+                                if ts_seconds and (now - ts_seconds) > max_age_seconds:
+                                    continue
+                            status = payload.get("status")
+                            if status:
+                                counter[status] += 1
+                        except Exception:
+                            continue
+                    if h_cursor == 0 or h_cursor == "0":
+                        break
+            if cursor == 0 or cursor == "0":
+                break
+
+        if not counter:
+            return TaskSummary()  # Empty summary if nothing found
+
+        total = sum(counter.values())
+        return TaskSummary(
+            waiting=counter.get("waiting", 0),
+            in_progress=counter.get("in_progress", 0),
+            completed=counter.get("completed", 0),
+            pending=counter.get("pending", counter.get("waiting", 0)),
+            failed=counter.get("failed", 0),
+            cancelled=counter.get("cancelled", 0),
+            total=total,
+        )
+
+    try:
+        # Prefer streaming aggregation to avoid pulling all task payloads
+        all_tasks_summary = _aggregate_counts_from_redis(status_tracker)
+        if all_tasks_summary is None:
+            # Fallback: load all details then aggregate
+            global_tasks = status_tracker.get_all_tasks_global()
+            all_task_details: list[dict[str, Any]] = []
+            for _, tasks in global_tasks.items():
+                all_task_details.extend(tasks.values())
+            all_tasks_summary = _summarize_tasks(all_task_details)
+
+        # Scheduler view: assume tracker contains scheduler tasks; overlay queue monitor for live queue depth
+        sched_waiting = all_tasks_summary.waiting
+        sched_in_progress = all_tasks_summary.in_progress
+        sched_pending = all_tasks_summary.pending
+        sched_completed = all_tasks_summary.completed
+        sched_failed = all_tasks_summary.failed
+        sched_cancelled = all_tasks_summary.cancelled
+
+        # If queue monitor is available, prefer its live waiting/in_progress counts
+        if mem_scheduler.task_schedule_monitor:
+            queue_status_data = mem_scheduler.task_schedule_monitor.get_tasks_status() or {}
+            scheduler_waiting = 0
+            scheduler_in_progress = 0
+            scheduler_pending = 0
+            for key, value in queue_status_data.items():
+                if not key.startswith("scheduler:"):
+                    continue
+                scheduler_in_progress += int(value.get("running", 0) or 0)
+                scheduler_pending += int(value.get("pending", value.get("remaining", 0)) or 0)
+                scheduler_waiting += int(value.get("remaining", 0) or 0)
+            sched_waiting = scheduler_waiting
+            sched_in_progress = scheduler_in_progress
+            sched_pending = scheduler_pending
+
+        scheduler_summary = TaskSummary(
+            waiting=sched_waiting,
+            in_progress=sched_in_progress,
+            pending=sched_pending,
+            completed=sched_completed,
+            failed=sched_failed,
+            cancelled=sched_cancelled,
+            total=sched_waiting
+            + sched_in_progress
+            + sched_completed
+            + sched_failed
+            + sched_cancelled,
+        )
+
+        return AllStatusResponse(
+            data=AllStatusResponseData(
+                scheduler_summary=scheduler_summary,
+                all_tasks_summary=all_tasks_summary,
+            )
+        )
+    except Exception as err:
+        logger.error(f"Failed to get full scheduler status: {traceback.format_exc()}")
+        raise HTTPException(status_code=500, detail="Failed to get full scheduler status") from err
+
+
 def handle_scheduler_status(
     user_id: str, status_tracker: TaskStatusTracker, task_id: str | None = None
 ) -> StatusResponse:

src/memos/api/product_models.py

@@ -882,3 +882,34 @@ class StatusResponse(BaseResponse[list[StatusResponseItem]]):
     """Response model for scheduler status operations."""
 
     message: str = "Memory get status successfully"
+
+
+class TaskSummary(BaseModel):
+    """Aggregated counts of tasks by status."""
+
+    waiting: int = Field(0, description="Number of tasks waiting to run")
+    in_progress: int = Field(0, description="Number of tasks currently running")
+    pending: int = Field(
+        0, description="Number of tasks fetched by workers but not yet acknowledged"
+    )
+    completed: int = Field(0, description="Number of tasks completed")
+    failed: int = Field(0, description="Number of tasks failed")
+    cancelled: int = Field(0, description="Number of tasks cancelled")
+    total: int = Field(0, description="Total number of tasks counted")
+
+
+class AllStatusResponseData(BaseModel):
+    """Aggregated scheduler status metrics."""
+
+    scheduler_summary: TaskSummary = Field(
+        ..., description="Aggregated status for scheduler-managed tasks"
+    )
+    all_tasks_summary: TaskSummary = Field(
+        ..., description="Aggregated status for all tracked tasks"
+    )
+
+
+class AllStatusResponse(BaseResponse[AllStatusResponseData]):
+    """Response model for full scheduler status operations."""
+
+    message: str = "Scheduler status summary retrieved successfully"

src/memos/api/routers/server_router.py

@@ -24,6 +24,7 @@
 from memos.api.handlers.feedback_handler import FeedbackHandler
 from memos.api.handlers.search_handler import SearchHandler
 from memos.api.product_models import (
+    AllStatusResponse,
     APIADDRequest,
     APIChatCompleteRequest,
     APIFeedbackRequest,
@@ -115,6 +116,18 @@ def add_memories(add_req: APIADDRequest):
 # =============================================================================
 
 
+@router.get(  # Changed from post to get
+    "/scheduler/allstatus",
+    summary="Get detailed scheduler status",
+    response_model=AllStatusResponse,
+)
+def scheduler_allstatus():
+    """Get detailed scheduler status including running tasks and queue metrics."""
+    return handlers.scheduler_handler.handle_scheduler_allstatus(
+        mem_scheduler=mem_scheduler, status_tracker=status_tracker
+    )
+
+
 @router.get(  # Changed from post to get
     "/scheduler/status", summary="Get scheduler running status", response_model=StatusResponse
 )

src/memos/graph_dbs/polardb.py

@@ -151,7 +151,7 @@ def __init__(self, config: PolarDBGraphDBConfig):
         # Create connection pool
         self.connection_pool = psycopg2.pool.ThreadedConnectionPool(
             minconn=5,
-            maxconn=2000,
+            maxconn=100,
             host=host,
             port=port,
             user=user,
@@ -1338,6 +1338,7 @@ def get_subgraph(
                 "edges": [...]
             }
         """
+        logger.info(f"[get_subgraph] center_id: {center_id}")
         if not 1 <= depth <= 5:
             raise ValueError("depth must be 1-5")
 
@@ -1375,6 +1376,7 @@ def get_subgraph(
                 $$ ) as (centers agtype, neighbors agtype, rels agtype);
             """
         conn = self._get_connection()
+        logger.info(f"[get_subgraph] Query: {query}")
         try:
             with conn.cursor() as cursor:
                 cursor.execute(query)
@@ -1746,6 +1748,7 @@ def search_by_embedding(
 
         # Build filter conditions using common method
         filter_conditions = self._build_filter_conditions_sql(filter)
+        logger.info(f"[search_by_embedding] filter_conditions: {filter_conditions}")
         where_clauses.extend(filter_conditions)
 
         where_clause = f"WHERE {' AND '.join(where_clauses)}" if where_clauses else ""
@@ -1918,7 +1921,7 @@ def get_by_metadata(
             knowledgebase_ids=knowledgebase_ids,
             default_user_name=self._get_config_value("user_name"),
         )
-        print(f"[111get_by_metadata] user_name_conditions: {user_name_conditions}")
+        logger.info(f"[get_by_metadata] user_name_conditions: {user_name_conditions}")
 
         # Add user_name WHERE clause
         if user_name_conditions:
@@ -1929,6 +1932,7 @@ def get_by_metadata(
 
         # Build filter conditions using common method
         filter_where_clause = self._build_filter_conditions_cypher(filter)
+        logger.info(f"[get_by_metadata] filter_where_clause: {filter_where_clause}")
 
         where_str = " AND ".join(where_conditions) + filter_where_clause
 
@@ -2393,6 +2397,7 @@ def get_all_memory_items(
 
         # Build filter conditions using common method
         filter_where_clause = self._build_filter_conditions_cypher(filter)
+        logger.info(f"[get_all_memory_items] filter_where_clause: {filter_where_clause}")
 
         # Use cypher query to retrieve memory items
         if include_embedding:
@@ -2426,6 +2431,7 @@ def get_all_memory_items(
             nodes = []
             node_ids = set()
             conn = self._get_connection()
+            logger.info(f"[get_all_memory_items] cypher_query: {cypher_query}")
             try:
                 with conn.cursor() as cursor:
                     cursor.execute(cypher_query)
@@ -3456,7 +3462,11 @@ def _convert_graph_edges(self, core_node: dict) -> dict:
         id_map = {}
         core_node = data.get("core_node", {})
         if not core_node:
-            return core_node
+            return {
+                "core_node": None,
+                "neighbors": data.get("neighbors", []),
+                "edges": data.get("edges", []),
+            }
         core_meta = core_node.get("metadata", {})
         if "graph_id" in core_meta and "id" in core_node:
             id_map[core_meta["graph_id"]] = core_node["id"]
@@ -3507,7 +3517,6 @@ def _build_user_name_and_kb_ids_conditions_cypher(
         """
         user_name_conditions = []
         effective_user_name = user_name if user_name else default_user_name
-        print(f"[delete_node_by_prams] effective_user_name: {effective_user_name}")
 
         if effective_user_name:
             escaped_user_name = effective_user_name.replace("'", "''")