feat(clp-mcp-server): Introduce session-based MCP capabilities:

components/clp-mcp-server/clp_mcp_server/server/constants.py

@@ -0,0 +1,25 @@
+"""Constants for CLP MCP server."""
+
+EXPIRED_SESSION_SWEEP_INTERVAL_SECONDS = 600  # 10 minutes
+ITEM_PER_PAGE = 10
+MAX_CACHED_RESULTS = 1000
+SESSION_TTL_MINUTES = 60
+
+SERVER_NAME = "clp-mcp-server"
+SYSTEM_PROMPT = (
+    "You are an AI assistant that helps users query a log database using KQL (Kibana Query"
+    " Language). When given a user query, you should generate a KQL query that accurately captures"
+    " the user's intent. The KQL query should be as specific as possible to minimize the number of"
+    " log messages returned.\n"
+    "You should also consider the following guidelines when generating KQL queries:\n"
+    "- Use specific field names and values to narrow down the search.\n"
+    "- Avoid using wildcards (*) unless absolutely necessary, as they can lead to large result"
+    " sets.\n"
+    "- Use logical operators (AND, OR, NOT) to combine multiple conditions.\n"
+    "- Consider the time range of the logs you are searching. If the user specifies a time range,"
+    " include it in the KQL query.\n"
+    "- If the user query is ambiguous or lacks detail, ask clarifying questions to better"
+    " understand their intent before generating the KQL query.\n"
+    "- Always ensure that the generated KQL query is syntactically correct and can be executed"
+    " without errors.\n"
+)

components/clp-mcp-server/clp_mcp_server/server/server.py

@@ -2,25 +2,10 @@
 
 from typing import Any
 
-from fastmcp import FastMCP
+from fastmcp import Context, FastMCP
 
-
-class ProtocolConstant:
-    """Constants for the CLP MCP Server."""
-
-    SERVER_NAME = "clp-mcp-server"
-
-    # Tool names
-    TOOL_HELLO_WORLD = "hello_world"
-    TOOL_GET_SERVER_INFO = "get_server_info"
-
-    @classmethod
-    def get_capabilities(cls) -> list[str]:
-        """
-        Gets the capabilities of the server.
-        :return: A list of tool names supported by the server.
-        """
-        return [cls.TOOL_HELLO_WORLD, cls.TOOL_GET_SERVER_INFO]
+from . import constants
+from .session_manager import SessionManager
 
 
 def create_mcp_server() -> FastMCP:
@@ -31,22 +16,39 @@ def create_mcp_server() -> FastMCP:
     :raise: Propagates `FastMCP.__init__`'s exceptions.
     :raise: Propagates `FastMCP.tool`'s exceptions.
     """
-    mcp = FastMCP(name=ProtocolConstant.SERVER_NAME)
+    mcp = FastMCP(name=constants.SERVER_NAME)
+
+    session_manager = SessionManager(constants.SESSION_TTL_MINUTES)
 
-    @mcp.tool()
-    def get_server_info() -> dict[str, Any]:
+    @mcp.tool
+    def get_instructions(ctx: Context) -> str:
         """
-        Gets the MCP server's information.
+        Gets a pre-defined “system prompt” that guides the LLM behavior.
+        This function must be invoked before any other `FastMCP.tool`.
 
-        :return: The server's information with a list of capabilities.
+        :param ctx: The `FastMCP` context containing the metadata of the underlying MCP session.
+        :return: A string of “system prompt”.
         """
-        return {
-            "name": ProtocolConstant.SERVER_NAME,
-            "capabilities": ProtocolConstant.get_capabilities(),
-            "status": "running",
-        }
+        session = session_manager.get_or_create_session(ctx.session_id)
+        session.is_instructions_retrieved = True
+        return constants.SYSTEM_PROMPT
+
+    @mcp.tool
+    def get_nth_page(page_index: int, ctx: Context) -> dict[str, Any]:
+        """
+        Retrieves the n-th page of a paginated response with the paging metadata from the previous
+        query.
+
+        :param page_index: Zero-based index, e.g., 0 for the first page.
+        :param ctx: The `FastMCP` context containing the metadata of the underlying MCP session.
+        :return: Dictionary containing paged log entries and the paging metadata if the
+        page `page_index` can be retrieved.
+        :return: Dictionary with ``{"Error": "error message describing the failure"}`` if fails to
+        retrieve page `page_index`.
+        """
+        return session_manager.get_nth_page(ctx.session_id, page_index)
 
-    @mcp.tool()
+    @mcp.tool
     def hello_world(name: str = "clp-mcp-server user") -> dict[str, Any]:
         """
         Provides a simple hello world greeting.
@@ -56,7 +58,7 @@ def hello_world(name: str = "clp-mcp-server user") -> dict[str, Any]:
         """
         return {
             "message": f"Hello World, {name.strip()}!",
-            "server": ProtocolConstant.SERVER_NAME,
+            "server": constants.SERVER_NAME,
             "status": "running",
         }
 

components/clp-mcp-server/clp_mcp_server/server/session_manager.py

@@ -0,0 +1,223 @@
+"""Session management for CLP MCP Server."""
+
+import threading
+import time
+from dataclasses import dataclass, field
+from datetime import datetime, timedelta, timezone
+from typing import Any
+
+from paginate import Page
+
+from . import constants
+
+
+class PaginatedQueryResult:
+    """Paginates the cached log entries returned from a query's response."""
+
+    def __init__(self, result_log_entries: list[str], num_items_per_page: int) -> None:
+        """
+        Initializes the PaginatedQueryResult.
+
+        :param result_log_entries: List of cached log entries to paginate.
+        :param num_items_per_page:
+        :raise: ValueError if the number of cached results or num_items_per_page is invalid.
+        """
+        if len(result_log_entries) > constants.MAX_CACHED_RESULTS:
+            err_msg = (
+                f"PaginatedQueryResult exceeds maximum allowed cached results:"
+                f" {len(result_log_entries)} > {constants.MAX_CACHED_RESULTS}."
+            )
+            raise ValueError(err_msg)
+
+        if num_items_per_page <= 0:
+            err_msg = (
+                f"Invalid num_items_per_page: {num_items_per_page}, "
+                "it must be a positive integer. "
+            )
+            raise ValueError(err_msg)
+
+        self.result_log_entries = result_log_entries
+        self.num_items_per_page = num_items_per_page
+
+        self._num_pages = (len(result_log_entries) + num_items_per_page - 1) // num_items_per_page
+
+    def get_page(self, page_index: int) -> Page | None:
+        """
+        Returns a page from the cached query results.
+
+        :param page_index: The number of page to retrieve (zero-based index; 0 is the first page).
+        :return: A `Page` object for the specified page.
+        :return: None if `page_index` is out of bounds.
+        """
+        page_number = page_index + 1  # Convert zero-based to one-based
+        if page_number <= 0 or self._num_pages < page_number:
+            return None
+
+        return Page(
+            self.result_log_entries,
+            page=page_number,
+            items_per_page=self.num_items_per_page,
+        )
+
+
+@dataclass
+class SessionState:
+    """Represents the state of a user session."""
+
+    session_id: str
+    _num_items_per_page: int
+    _session_ttl_minutes: int
+
+    is_instructions_retrieved: bool = False
+    last_accessed: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
+    _cached_query_result: PaginatedQueryResult | None = None
+
+    def cache_query_result(
+        self,
+        results: list[str],
+    ) -> None:
+        """
+        Caches the latest query result of the session.
+
+        :param results: Complete log entries from previous query for caching.
+        """
+        self._cached_query_result = PaginatedQueryResult(
+            result_log_entries=results, num_items_per_page=self._num_items_per_page
+        )
+
+    def get_page_data(self, page_index: int) -> dict[str, Any]:
+        """
+        Gets page data and its metadata in a dictionary format.
+
+        :param page_index: The number of page to retrieve (zero-based index; 0 is the first page).
+        :return: Dictionary containing paged log entries and the paging metadata if the
+        page `page_index` can be retrieved.
+        :return: Dictionary with ``{"Error": "error message describing the failure"}`` if fails to
+        retrieve page `page_index`.
+        """
+        if self._cached_query_result is None:
+            return {"Error": "No previous paginated response in this session."}
+
+        page = self._cached_query_result.get_page(page_index)
+        if page is None:
+            return {"Error": "Page index is out of bounds."}
+
+        return {
+            "items": list(page),
+            "total_pages": page.page_count,
+            "total_items": page.item_count,
+            "num_items_per_page": page.items_per_page,
+            "has_next": page.next_page is not None,
+            "has_previous": page.previous_page is not None,
+        }
+
+    def is_expired(self) -> bool:
+        """:return: Whether the session has expired."""
+        time_diff = datetime.now(timezone.utc) - self.last_accessed
+        return time_diff > timedelta(minutes=self._session_ttl_minutes)
+
+    def update_access_time(self) -> None:
+        """Updates the last accessed timestamp."""
+        self.last_accessed = datetime.now(timezone.utc)
+
+
+class SessionManager:
+    """
+    Session manager for concurrent user sessions.
+    `SessionManager` respects MCP Server Concurrency Model, that is:
+    The server supports multiple concurrent clients, where each client only makes synchronous
+    API calls.
+    This model leads to the following design decisions:
+    sessions is a shared variable as there may be multiple session attached to the MCP server
+    session state is NOT a shared variable because each session is accessed by only one
+    connection at a time because API calls for a single session are synchronous.
+    """
+
+    def __init__(self, session_ttl_minutes: int) -> None:
+        """
+        Initializes the SessionManager and starts background cleanup thread.
+
+        :param session_ttl_minutes: Session time-to-live in minutes.
+        """
+        self.sessions: dict[str, SessionState] = {}
+        self._session_ttl_minutes = session_ttl_minutes
+        self._sessions_lock = threading.Lock()
+        self._cleanup_thread = threading.Thread(target=self._cleanup_loop, daemon=True)
+        self._cleanup_thread.start()
+
+    def _cleanup_loop(self) -> None:
+        """Cleans up all expired sessions periodically in a separate cleanup thread."""
+        while True:
+            time.sleep(constants.EXPIRED_SESSION_SWEEP_INTERVAL_SECONDS)
+            self.cleanup_expired_sessions()
+
+    def cleanup_expired_sessions(self) -> None:
+        """Cleans up all expired sessions."""
+        with self._sessions_lock:
+            expired_sessions = [
+                sid for sid, session in self.sessions.items() if session.is_expired()
+            ]
+
+            for sid in expired_sessions:
+                del self.sessions[sid]
+
+    def get_or_create_session(self, session_id: str) -> SessionState:
+        """
+        Gets an existing session or creates a new one.
+
+        :param session_id: Unique identifier for the session.
+        :return: The SessionState object for the given session_id.
+        """
+        with self._sessions_lock:
+            if session_id in self.sessions and self.sessions[session_id].is_expired():
+                del self.sessions[session_id]
+
+            if session_id not in self.sessions:
+                self.sessions[session_id] = SessionState(
+                    session_id, constants.ITEM_PER_PAGE, self._session_ttl_minutes
+                )
+
+            session = self.sessions[session_id]
+
+            session.update_access_time()
+            return session
+
+    def cache_query_result(self, session_id: str, query_results: list[str]) -> dict[str, Any]:
+        """
+        Caches query results for a session and returns the first page and the paging metadata.
+
+        :param session_id: Unique identifier for the session.
+        :param query_results: Complete log entries from previous query for caching.
+        :return: Dictionary containing the first page log entries and the paging metadata if the
+        first page can be retrieved.
+        :return: Dictionary with ``{"Error": "error message describing the failure"}`` if fails to
+        retrieve the first page.
+        """
+        session = self.get_or_create_session(session_id)
+        if session.is_instructions_retrieved is False:
+            return {
+                "Error": "Please call get_instructions() first "
+                "to understand how to use this MCP server."
+            }
+
+        session.cache_query_result(results=query_results)
+
+        return session.get_page_data(0)
+
+    def get_nth_page(self, session_id: str, page_index: int) -> dict[str, Any]:
+        """
+        Retrieves the n-th page of a paginated response with the paging metadata from the previous
+        query.
+
+        :param session_id: Unique identifier for the session.
+        :param page_index: The number of page to retrieve (zero-based index; 0 is the first page).
+        :return: Forwards `SessionState.get_page_data`'s return values.
+        """
+        session = self.get_or_create_session(session_id)
+        if session.is_instructions_retrieved is False:
+            return {
+                "Error": "Please call get_instructions() first "
+                "to understand how to use this MCP server."
+            }
+
+        return session.get_page_data(page_index)

components/clp-mcp-server/pyproject.toml

@@ -6,11 +6,12 @@ authors = [{name = "YScope Inc.", email = "dev@yscope.com"}]
 readme = "README.md"
 requires-python = ">=3.10"
 dependencies = [
+    "aiomysql>=0.2.0",
     "click>=8.3.0",
     "fastmcp>=2.12.4",
-    "pymongo>=4.15.1",
-    "aiomysql>=0.2.0",
     "msgpack>=1.1.1",
+    "paginate>=0.5.7",
+    "pymongo>=4.15.1",
 ]
 
 [project.scripts]
@@ -23,10 +24,11 @@ build-backend = "hatchling.build"
 [dependency-groups]
 dev = [
     "mypy>=1.16.0",
-    "ruff>=0.11.12",
     "pytest>=8.4.1",
-    "pytest-env>=1.1.5",
     "pytest-asyncio>=1.2.0",
+    "pytest-env>=1.1.5",
+    "pytest-repeat>=0.9.4",
+    "ruff>=0.11.12",
 ]
 
 [tool.hatch.metadata]
@@ -69,12 +71,12 @@ ignore = [
 ]
 isort.order-by-type = false
 
+[tool.ruff.lint.per-file-ignores]
+"tests/**" = [
+    "S101",  # Allow usage of pytest `assert`
+    "TC003",  # Ignore performance overhead of imports only used for type checking
+]
+
 [tool.ruff.format]
 docstring-code-format = true
 docstring-code-line-length = 100
-
-[tool.ruff.lint.per-file-ignores]
-    "tests/test_clp_connector.py" = [
-        "INP001",  # Allow implicit namespace package for tests
-        "S101",  # Allow "assert" in test cases
-    ]

components/clp-mcp-server/tests/__init__.py

@@ -0,0 +1 @@
+"""Test Package for CLP MCP Server."""