Repository: keboola/keboola-mcp-server
Pull request #11: KAB-817: add the concept of "session state factory"
Author: login: vita-stejskal
Created: 2025-03-14T10:12:17Z | Merged: 2025-03-19T20:45:39Z

There is quite a bit of refactoring on this branch which can be split in two main parts:

* Extending the FastMCP to allow attaching a custom "state" to the SSE session initiated when an MCP client opens the conversation to an MCP server via the `sse` transport (i.e. the first `GET /sse` request).
* Updating our existing tools to consume the `mcp.server.fastmcp.Context` object and use the custom "state" attached to the `Context`'s `session.state` field. The "state" is created by our own session state factory function which reads the session parameters and creates `KeboolaClient` (Storage API client) and `ConnectionManager` (Snowflake connector) instances. These instances are than available in the tools.

The FastMCP extenstions are in `mcp.py` file and provide the following stuff:

* Definitions of `SessionParams`, `SessionState` and `SessionStateFactory` -- the things that you need in order to implement your own session state factory which is a function that accepts the `SessionParams` and produces `SessionState`.
* `StatefullServerSession` -- the original `mcp.ServerSession` with the extra field for `SessionState` object
* `_KeboolaServer` -- the subclass of `mcp.server.Server` which modifies `run()` function to accept the `SessionState`, creates `StatefullServerSession` and starts the actual SSE server
* `KeboolaMcpServer` -- the subclass of `mcp.server.FastMcp` which modifies `run_stdio_async()` and `run_sse_async()` functions. The `KeboolaMcpServer` accepts `SessionStateFoactory` in its constructor and the `run_*_async()` functions call the factory in order to get the custom "state" that is then passed to `_KeboolaServer` when starting it.
* The `SessionParams` needed for the `SessionStateFactory` are taken either from the process environment (when using `stdio` transport) or from the HTTP query parameters (when using `sse` transport).

Notes:
* I am surprised that the extensions from `mcp.py` are not already part of the original FastMCP implementation in `mcp` python package. I believe that they would be useful for anyone using FastMCP, especially when using the `sse` transport.
* The FastMCP implementation of the "resources" concept from the MCP (protocol) does not seem to support passing the `mcp.server.fastmcp.Context` object to the functions registered as the resources providers. It seems odd and I did not find anything in the MCP (protocol) itself that would justify this. But I could be wrong.
* The testing of the MCP servers implemented using FastMCP seems to be very challenging. There is no support for testing in the `mcp` python package itself. I'd expect something like Starlette's `TestClient` capable of invoking directly the endpoints of a Starlette-base server.
* In general, my feeling is that the main ideas behind FastMCP are great, but its implementation is still immature and hard to use when you need things beyond the examples from its tutorial.


Diff:
diff --git a/src/keboola_mcp_server/cli.py b/src/keboola_mcp_server/cli.py
index 784e6164..93171f90 100644
--- a/src/keboola_mcp_server/cli.py
+++ b/src/keboola_mcp_server/cli.py
@@ -1,7 +1,6 @@
 """Command-line interface for the Keboola MCP server."""
 
 import argparse
-import asyncio
 import logging
 import sys
 from typing import List, Optional
@@ -24,7 +23,7 @@ def parse_args(args: Optional[List[str]] = None) -> argparse.Namespace:
     parser = argparse.ArgumentParser(description="Keboola MCP Server")
     parser.add_argument(
         "--transport",
-        choices=["stdio"],
+        choices=["stdio", "sse"],
         default="stdio",
         help="Transport to use for MCP communication",
     )
@@ -35,7 +34,7 @@ def parse_args(args: Optional[List[str]] = None) -> argparse.Namespace:
         help="Logging level",
     )
     parser.add_argument(
-        "--api-url", help="Keboola Storage API URL (defaults to https://connection.keboola.com)"
+        "--api-url", default="https://connection.keboola.com", help="Keboola Storage API URL"
     )
 
     return parser.parse_args(args)
@@ -49,11 +48,10 @@ def main(args: Optional[List[str]] = None) -> None:
     """
     parsed_args = parse_args(args)
 
-    # Create config from environment, but override with CLI args
-    config = Config.from_env()
-    if parsed_args.api_url:
-        config.storage_api_url = parsed_args.api_url
-    config.log_level = parsed_args.log_level
+    # Create config from the CLI arguments
+    config = Config.from_dict(
+        {"storage_api_url": parsed_args.api_url, "log_level": parsed_args.log_level}
+    )
 
     try:
         # Create and run server
diff --git a/src/keboola_mcp_server/config.py b/src/keboola_mcp_server/config.py
index 7dc24e4d..63f8ce91 100644
--- a/src/keboola_mcp_server/config.py
+++ b/src/keboola_mcp_server/config.py
@@ -1,18 +1,18 @@
 """Configuration handling for the Keboola MCP server."""
 
+import dataclasses
 import logging
-import os
 from dataclasses import dataclass
-from typing import Optional
+from typing import Mapping, Optional
 
 logger = logging.getLogger(__name__)
 
 
-@dataclass
+@dataclass(frozen=True)
 class Config:
     """Server configuration."""
 
-    storage_token: str
+    storage_token: Optional[str] = None
     storage_api_url: str = "https://connection.keboola.com"
     log_level: str = "INFO"
     # Add Snowflake credentials
@@ -24,71 +24,42 @@ class Config:
     snowflake_schema: Optional[str] = None
     snowflake_role: Optional[str] = None
 
-    def __init__(
-        self,
-        storage_token: str,
-        storage_api_url: str = "https://connection.keboola.com",
-        snowflake_account: Optional[str] = None,
-        snowflake_user: Optional[str] = None,
-        snowflake_password: Optional[str] = None,
-        snowflake_warehouse: Optional[str] = None,
-        snowflake_database: Optional[str] = None,
-        snowflake_role: Optional[str] = None,
-        snowflake_schema: Optional[str] = None,
-        log_level: str = "INFO",
-    ):
-        self.storage_token = storage_token
-        self.storage_api_url = storage_api_url
-        self.snowflake_account = snowflake_account
-        self.snowflake_user = snowflake_user
-        self.snowflake_password = snowflake_password
-        self.snowflake_warehouse = snowflake_warehouse
-        self.snowflake_database = snowflake_database
-        self.snowflake_role = snowflake_role
-        self.snowflake_schema = snowflake_schema
-        self.log_level = log_level
+    @classmethod
+    def _read_options(cls, d: Mapping[str, str]) -> Mapping[str, str]:
+        options: dict[str, str] = {}
+        for f in dataclasses.fields(cls):
+            if f.name in d:
+                options[f.name] = d.get(f.name)
+            elif (dict_name := f"KBC_{f.name.upper()}") in d:
+                options[f.name] = d.get(dict_name)
+        return options
 
     @classmethod
-    def from_env(cls) -> "Config":
-        """Create config from environment variables."""
-        # Add debug logging using logger instead of print
-        for env_var in [
-            "KBC_SNOWFLAKE_ACCOUNT",
-            "KBC_SNOWFLAKE_USER",
-            "KBC_SNOWFLAKE_PASSWORD",
-            "KBC_SNOWFLAKE_WAREHOUSE",
-            "KBC_SNOWFLAKE_DATABASE",
-            "KBC_SNOWFLAKE_ROLE",
-            "KBC_SNOWFLAKE_SCHEMA",
-        ]:
-            logger.debug(f"Reading {env_var}: {'set' if os.getenv(env_var) else 'not set'}")
+    def from_dict(cls, d: Mapping[str, str]) -> "Config":
+        """
+        Creates new `Config` instance with values read from the input mapping.
+        The keys in the input mapping can either be the names of the fields in `Config` class
+        or their uppercase variant prefixed with 'KBC_'.
+        """
+        return cls(**cls._read_options(d))
 
-        storage_token = os.getenv("KBC_STORAGE_TOKEN")
-        if not storage_token:
-            raise ValueError("KBC_STORAGE_TOKEN environment variable is required")
+    def replace_by(self, d: Mapping[str, str]) -> "Config":
+        """
+        Creates new `Config` instance from the existing one by replacing the values from the input mapping.
+        The keys in the input mapping can either be the names of the fields in `Config` class
+        or their uppercase variant prefixed with 'KBC_'.
+        """
+        return dataclasses.replace(self, **self._read_options(d))
 
-        return cls(
-            storage_token=storage_token,
-            storage_api_url=os.getenv("KBC_STORAGE_API_URL", "https://connection.keboola.com"),
-            snowflake_account=os.getenv("KBC_SNOWFLAKE_ACCOUNT"),
-            snowflake_user=os.getenv("KBC_SNOWFLAKE_USER"),
-            snowflake_password=os.getenv("KBC_SNOWFLAKE_PASSWORD"),
-            snowflake_warehouse=os.getenv("KBC_SNOWFLAKE_WAREHOUSE"),
-            snowflake_database=os.getenv("KBC_SNOWFLAKE_DATABASE"),
-            snowflake_role=os.getenv("KBC_SNOWFLAKE_ROLE"),
-            snowflake_schema=os.getenv("KBC_SNOWFLAKE_SCHEMA"),
-            log_level=os.getenv("KBC_LOG_LEVEL", "INFO"),
+    def has_storage_config(self) -> bool:
+        """Check if Storage API configuration is complete."""
+        return all(
+            [
+                self.storage_token,
+                self.storage_api_url,
+            ]
         )
 
-    def validate(self) -> None:
-        """Validate the configuration."""
-        if not self.storage_token:
-            raise ValueError("Storage token not configured")
-        if not self.storage_api_url:
-            raise ValueError("Storage API URL is required")
-        if self.log_level not in ["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"]:
-            raise ValueError(f"Invalid log level: {self.log_level}")
-
     def has_snowflake_config(self) -> bool:
         """Check if Snowflake configuration is complete."""
         return all(
@@ -100,3 +71,16 @@ def has_snowflake_config(self) -> bool:
                 self.snowflake_database,
             ]
         )
+
+    def __repr__(self):
+        params: list[str] = []
+        for f in dataclasses.fields(self):
+            value = getattr(self, f.name)
+            if value:
+                if "token" in f.name or "password" in f.name:
+                    params.append(f"{f.name}='****'")
+                else:
+                    params.append(f"{f.name}='{value}'")
+            else:
+                params.append(f"{f.name}=None")
+        return f'Config({", ".join(params)})'
diff --git a/src/keboola_mcp_server/database.py b/src/keboola_mcp_server/database.py
index 283c7516..eda609ec 100644
--- a/src/keboola_mcp_server/database.py
+++ b/src/keboola_mcp_server/database.py
@@ -2,12 +2,10 @@
 
 import logging
 from contextlib import contextmanager
-
-import snowflake.connector
 from dataclasses import dataclass
-from typing import Optional, List, Tuple, Dict, Any
+from typing import Any, Dict, List
 
-from .config import Config
+import snowflake.connector
 
 logger = logging.getLogger(__name__)
 
@@ -179,38 +177,25 @@ def find_working_connection(self) -> str:
             error_msg += f" - {pattern} ({db}): {error}\n"
         raise DatabaseConnectionError(error_msg)
 
+    def create_snowflake_connection(self) -> snowflake.connector.connection:
+        """Create and return a Snowflake connection using configured credentials.
+        Raises:
+            ValueError: If credentials are not fully configured or connection fails
+        """
+        try:
+            database = self.find_working_connection()
+            conn = snowflake.connector.connect(
+                account=self.config.snowflake_account,
+                user=self.config.snowflake_user,
+                password=self.config.snowflake_password,
+                warehouse=self.config.snowflake_warehouse,
+                database=database,
+                schema=self.config.snowflake_schema,
+                role=self.config.snowflake_role,
+            )
 
-def create_snowflake_connection(config: Config) -> snowflake.connector.connection:
-    """Create a return a Snowflake connection using configured credentials.
-
-    Args:
-        config: Configuration object containing Snowflake credentials
-
-    Returns:
-        snowflake.connector.connection: established Snowflake connection
-
-    Raises:
-        ValueError: If credentials are not fully configured or connection fails
-    """
-    if not config.has_snowflake_config():
-        raise ValueError("Snowflake credentials are not fully configured")
-
-    try:
-        connection_manager = ConnectionManager(config)
-        database = connection_manager.find_working_connection()
-
-        conn = snowflake.connector.connect(
-            account=config.snowflake_account,
-            user=config.snowflake_user,
-            password=config.snowflake_password,
-            warehouse=config.snowflake_warehouse,
-            database=database,
-            schema=config.snowflake_schema,
-            role=config.snowflake_role,
-        )
-
-        return conn
-    except DatabaseConnectionError as e:
-        raise ValueError(f"Failed to find working database connection: {str(e)}")
-    except Exception as e:
-        raise ValueError(f"Failed to create Snowflake connection: {str(e)}")
+            return conn
+        except DatabaseConnectionError as e:
+            raise ValueError(f"Failed to find working database connection: {str(e)}")
+        except Exception as e:
+            raise ValueError(f"Failed to create Snowflake connection: {str(e)}")
diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py
new file mode 100644
index 00000000..f70ff478
--- /dev/null
+++ b/src/keboola_mcp_server/mcp.py
@@ -0,0 +1,169 @@
+"""
+This is the extension of mcp.server.FastMCP and mcp.server.Server classes that allows to attach the "state"
+to the SSE session. The state is created by the state factory function that can be plugged in to the MCP server,
+and that creates a state which contains arbitrary objects keyed by string identifiers. The factory is given the
+query parameters from the HTTP request that initiates the SSE connection.
+
+Example:
+def factory(params: HttpRequestParams) -> SessionState:
+    return { 'sapi_client': KeboolaClient(params['storage_token']) }
+
+mcp = KeboolaMcpServer(name="SAPI Connector", session_state_factory=factory)
+
+@mcp.tool()
+def list_all_buckets(ctx: Context):
+    client = ctx.session.state['sapi_client']
+    return client.storage_client.buckets.list()
+
+mcp.run(transport='sse')
+
+Issues:
+  * The current implementation of FastMCP does not support sending `Context` to the registered
+    resources' functions. The parameter is passed only to the registered tools.
+"""
+
+import logging
+import os
+from contextlib import AbstractAsyncContextManager, AsyncExitStack
+from typing import Any, Callable
+
+import anyio
+import mcp.types as types
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from mcp import ServerSession, stdio_server
+from mcp.server import FastMCP, Server
+from mcp.server.lowlevel.server import LifespanResultT
+from mcp.server.models import InitializationOptions
+from mcp.server.sse import SseServerTransport
+
+logger = logging.getLogger(__name__)
+
+SessionParams = dict[str, str]
+SessionState = dict[str, Any]
+SessionStateFactory = Callable[[SessionParams], SessionState]
+
+
+def _default_session_state_factory(params: SessionParams) -> SessionState:
+    return params
+
+
+class StatefullServerSession(ServerSession):
+    def __init__(
+        self,
+        read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception],
+        write_stream: MemoryObjectSendStream[types.JSONRPCMessage],
+        init_options: InitializationOptions,
+        state: SessionState | None = None,
+    ) -> None:
+        super().__init__(read_stream, write_stream, init_options)
+        self._state = state or {}
+
+    @property
+    def state(self) -> SessionState:
+        return self._state
+
+
+class _KeboolaServer(Server):
+    def __init__(
+        self,
+        name: str,
+        version: str | None = None,
+        instructions: str | None = None,
+        lifespan: Callable[["Server"], AbstractAsyncContextManager[LifespanResultT]] | None = None,
+    ) -> None:
+        super().__init__(name, version=version, instructions=instructions, lifespan=lifespan)
+
+    async def run(
+        self,
+        read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception],
+        write_stream: MemoryObjectSendStream[types.JSONRPCMessage],
+        initialization_options: InitializationOptions,
+        # When False, exceptions are returned as messages to the client.
+        # When True, exceptions are raised, which will cause the server to shut down
+        # but also make tracing exceptions much easier during testing and when using
+        # in-process servers.
+        raise_exceptions: bool = False,
+        state: SessionState | None = None,
+    ):
+        async with AsyncExitStack() as stack:
+            lifespan_context = await stack.enter_async_context(self.lifespan(self))
+            session = await stack.enter_async_context(
+                StatefullServerSession(read_stream, write_stream, initialization_options, state)
+            )
+
+            async with anyio.create_task_group() as tg:
+                async for message in session.incoming_messages:
+                    logger.debug(f"Received message: {message}")
+
+                    tg.start_soon(
+                        self._handle_message,
+                        message,
+                        session,
+                        lifespan_context,
+                        raise_exceptions,
+                    )
+
+
+class KeboolaMcpServer(FastMCP):
+    def __init__(
+        self,
+        name: str | None = None,
+        instructions: str | None = None,
+        *,
+        session_state_factory: SessionStateFactory | None = None,
+        **settings: Any,
+    ) -> None:
+        super().__init__(name, instructions, **settings)
+        self._mcp_server = _KeboolaServer(
+            name=self._mcp_server.name,
+            instructions=self._mcp_server.instructions,
+            lifespan=self._mcp_server.lifespan,
+        )
+        self._setup_handlers()
+        self._session_state_factory = session_state_factory or _default_session_state_factory
+
+    async def run_stdio_async(self) -> None:
+        """Run the server using stdio transport."""
+        async with stdio_server() as (read_stream, write_stream):
+            await self._mcp_server.run(
+                read_stream,
+                write_stream,
+                initialization_options=self._mcp_server.create_initialization_options(),
+                state=self._session_state_factory(dict(os.environ)),
+            )
+
+    async def run_sse_async(self) -> None:
+        """Run the server using SSE transport."""
+        from starlette.applications import Starlette
+        from starlette.routing import Mount, Route
+        from starlette.requests import Request
+        import uvicorn
+
+        sse = SseServerTransport("/messages/")
+
+        async def handle_sse(request: Request):
+            async with sse.connect_sse(request.scope, request.receive, request._send) as streams:
+                await self._mcp_server.run(
+                    streams[0],
+                    streams[1],
+                    initialization_options=self._mcp_server.create_initialization_options(),
+                    state=self._session_state_factory(dict(request.query_params)),
+                )
+
+        starlette_app = Starlette(
+            debug=self.settings.debug,
+            routes=[
+                Route("/sse", endpoint=handle_sse),
+                Mount("/messages/", app=sse.handle_post_message),
+                # TODO: add endpoints for health-check and info
+            ],
+        )
+
+        config = uvicorn.Config(
+            starlette_app,
+            host=self.settings.host,
+            port=self.settings.port,
+            log_level=self.settings.log_level.lower(),
+        )
+        server = uvicorn.Server(config)
+        await server.serve()
diff --git a/src/keboola_mcp_server/server.py b/src/keboola_mcp_server/server.py
index 7ece938a..608bb971 100644
--- a/src/keboola_mcp_server/server.py
+++ b/src/keboola_mcp_server/server.py
@@ -3,16 +3,20 @@
 import csv
 import logging
 from io import StringIO
-from typing import Any, Dict, List, Optional, TypedDict, cast
-
+from typing import Any, cast, Dict, List, Optional, TypedDict
 
 import snowflake.connector
-
-from mcp.server.fastmcp import FastMCP
-
+from mcp.server.fastmcp import Context, FastMCP
+
+from keboola_mcp_server.mcp import (
+    KeboolaMcpServer,
+    SessionParams,
+    SessionState,
+    SessionStateFactory,
+)
 from .client import KeboolaClient
 from .config import Config
-from .database import create_snowflake_connection, ConnectionManager, DatabasePathManager
+from .database import ConnectionManager, DatabasePathManager
 
 logger = logging.getLogger(__name__)
 
@@ -44,6 +48,38 @@ class TableDetail(TypedDict):
     db_identifier: str
 
 
+def _create_session_state_factory(config: Optional[Config] = None) -> SessionStateFactory:
+    def _(params: SessionParams) -> SessionState:
+        logger.info(f"Creating SessionState for params: {params.keys()}.")
+
+        if not config:
+            cfg = Config.from_dict(params)
+        else:
+            cfg = config.replace_by(params)
+
+        logger.info(f"Creating SessionState from config: {cfg}.")
+
+        state: SessionState = {}
+        # Create Keboola client instance
+        try:
+            client = KeboolaClient(cfg.storage_token, cfg.storage_api_url)
+            state["sapi_client"] = client
+            logger.info("Successfully initialized Storage API client.")
+        except Exception as e:
+            logger.error(f"Failed to initialize Keboola client: {e}")
+            raise
+
+        connection_manager = ConnectionManager(cfg)
+        db_path_manager = DatabasePathManager(cfg, connection_manager)
+        state["connection_manager"] = connection_manager
+        state["db_path_manager"] = db_path_manager
+        logger.info("Successfully initialized DB connection and path managers.")
+
+        return state
+
+    return _
+
+
 def create_server(config: Optional[Config] = None) -> FastMCP:
     """Create and configure the MCP server.
 
@@ -53,10 +89,6 @@ def create_server(config: Optional[Config] = None) -> FastMCP:
     Returns:
         Configured FastMCP server instance
     """
-    if config is None:
-        config = Config.from_env()
-    config.validate()
-
     # Configure logging
     handler = logging.StreamHandler()
     handler.setFormatter(logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s"))
@@ -64,8 +96,9 @@ def create_server(config: Optional[Config] = None) -> FastMCP:
     logger.setLevel(config.log_level)
 
     # Initialize FastMCP server with system instructions
-    mcp = FastMCP(
+    mcp = KeboolaMcpServer(
         "Keboola Explorer",
+        session_state_factory=_create_session_state_factory(config),
         dependencies=[
             "keboola.storage-api-client",
             "httpx",
@@ -74,67 +107,8 @@ def create_server(config: Optional[Config] = None) -> FastMCP:
         ],
     )
 
-    connection_manager = ConnectionManager(config)
-    db_path_manager = DatabasePathManager(config, connection_manager)
-
-    # Create Keboola client instance
-    try:
-        keboola = KeboolaClient(config.storage_token, config.storage_api_url)
-    except Exception as e:
-        logger.error(f"Failed to initialize Keboola client: {e}")
-        raise
-    logger.info("Successfully initialized Keboola client")
-
-    # Resources
-    @mcp.resource("keboola://buckets")
-    async def list_buckets() -> List[BucketInfo]:
-        """List all available buckets in Keboola project."""
-        buckets = cast(List[BucketInfo], keboola.storage_client.buckets.list())
-        return buckets
-
-    @mcp.resource("keboola://buckets/{bucket_id}/tables")
-    async def list_bucket_tables(bucket_id: str) -> str:
-        """List all tables in a specific bucket."""
-        tables = cast(List[Dict[str, Any]], keboola.storage_client.buckets.list_tables(bucket_id))
-        return "\n".join(
-            f"- {table['id']}: {table['name']} (Rows: {table.get('rowsCount', 'unknown')})"
-            for table in tables
-        )
-
-    @mcp.resource("keboola://components")
-    async def list_components() -> str:
-        """List all available components and their configurations."""
-        components = cast(List[Dict[str, Any]], await keboola.get("components"))
-        return "\n".join(f"- {comp['id']}: {comp['name']}" for comp in components)
-
-    @mcp.resource(
-        "keboola://tables/{table_id}",
-        description="Get detailed information about a Keboola table including its DB identifier and column information",
-    )
-    async def get_table_detail(table_id: str) -> TableDetail:
-        """Get detailed information about a table."""
-        table = cast(Dict[str, Any], keboola.storage_client.tables.detail(table_id))
-
-        # Get column info
-        columns = table.get("columns", [])
-        column_info = [TableColumnInfo(name=col, db_identifier=f'"{col}"') for col in columns]
-
-        return {
-            "id": table["id"],
-            "name": table.get("name", "N/A"),
-            "primary_key": table.get("primaryKey", []),
-            "created": table.get("created", "N/A"),
-            "row_count": table.get("rowsCount", 0),
-            "data_size_bytes": table.get("dataSizeBytes", 0),
-            "columns": columns,
-            "column_identifiers": column_info,
-            "db_identifier": db_path_manager.get_table_db_path(table),
-            "schema_identifier": table["id"].split(".")[0],
-            "table_identifier": table["id"].split(".")[1],
-        }
-
     @mcp.tool()
-    async def query_table(sql_query: str) -> str:
+    async def query_table(sql_query: str, ctx: Context) -> str:
         """
         Execute a SQL query through the proxy service to get data from Storage.
         Before forming the query always check the get_table_metadata tool to get
@@ -145,11 +119,14 @@ async def query_table(sql_query: str) -> str:
         'SELECT * FROM {{db_identifier}}."test_identify"'. Snowflake is case sensitive so always
         wrap the column names in double quotes.
         """
+        connection_manager = ctx.session.state["connection_manager"]
+        assert isinstance(connection_manager, ConnectionManager)
+
         conn = None
         cursor = None
 
         try:
-            conn = create_snowflake_connection(config)
+            conn = connection_manager.create_snowflake_connection()
             cursor = conn.cursor()
             cursor.execute(sql_query)
             result = cursor.fetchall()
@@ -177,9 +154,11 @@ async def query_table(sql_query: str) -> str:
 
     # Tools
     @mcp.tool()
-    async def list_all_buckets() -> str:
+    async def list_all_buckets(ctx: Context) -> str:
         """List all buckets in the project with their basic information."""
-        buckets = await list_buckets()
+        client = ctx.session.state["sapi_client"]
+        assert isinstance(client, KeboolaClient)
+        buckets = cast(List[BucketInfo], client.storage_client.buckets.list())
 
         header = "# Bucket List\n\n"
         header += f"Total Buckets: {len(buckets)}\n\n"
@@ -198,9 +177,11 @@ async def list_all_buckets() -> str:
         return header + "\n" + "\n".join(bucket_details)
 
     @mcp.tool()
-    async def get_bucket_metadata(bucket_id: str) -> str:
+    async def get_bucket_metadata(bucket_id: str, ctx: Context) -> str:
         """Get detailed information about a specific bucket."""
-        bucket = cast(Dict[str, Any], keboola.storage_client.buckets.detail(bucket_id))
+        client = ctx.session.state["sapi_client"]
+        assert isinstance(client, KeboolaClient)
+        bucket = cast(Dict[str, Any], client.storage_client.buckets.detail(bucket_id))
         return (
             f"Bucket Information:\n"
             f"ID: {bucket['id']}\n"
@@ -212,30 +193,47 @@ async def get_bucket_metadata(bucket_id: str) -> str:
         )
 
     @mcp.tool()
-    async def get_table_metadata(table_id: str) -> Dict[str, Any]:
+    async def get_table_metadata(table_id: str, ctx: Context) -> str:
         """Get detailed information about a specific table including its DB identifier and column information."""
-        # Get table details directly from the storage client
-        table = await get_table_detail(table_id)
+        client = ctx.session.state["sapi_client"]
+        assert isinstance(client, KeboolaClient)
+        table = cast(Dict[str, Any], client.storage_client.tables.detail(table_id))
+
+        # Get column info
+        columns = table.get("columns", [])
+        column_info = [TableColumnInfo(name=col, db_identifier=f'"{col}"') for col in columns]
+
+        db_path_manager = ctx.session.state["db_path_manager"]
+        assert isinstance(db_path_manager, DatabasePathManager)
+
         return (
             f"Table Information:\n"
             f"ID: {table['id']}\n"
             f"Name: {table['name']}\n"
-            f"Primary Key: {', '.join(table['primary_key']) if table['primary_key'] else 'None'}\n"
+            f"Primary Key: {', '.join(table['primaryKey']) if table['primaryKey'] else 'None'}\n"
             f"Created: {table['created']}\n"
-            f"Row Count: {table['row_count']}\n"
-            f"Data Size: {table['data_size_bytes']} bytes\n"
-            f"Columns: {', '.join(table['columns'])}\n"
-            f"Database Identifier: {table['db_identifier']}\n"
-            f"Schema: {table['schema_identifier']}\n"
-            f"Table: {table['table_identifier']}"
+            f"Row Count: {table['rowsCount']}\n"
+            f"Data Size: {table['dataSizeBytes']} bytes\n"
+            f"Columns: {', '.join(str(ci) for ci in column_info)}\n"
+            f"Database Identifier: {db_path_manager.get_table_db_path(table)}\n"
+            f"Schema: {table['id'].split('.')[0]}\n"
+            f"Table: {table['id'].split('.')[1]}"
         )
 
     @mcp.tool()
-    async def list_component_configs(component_id: str) -> str:
+    async def list_components(ctx: Context) -> str:
+        """List all available components and their configurations."""
+        client = ctx.session.state["sapi_client"]
+        assert isinstance(client, KeboolaClient)
+        components = cast(List[Dict[str, Any]], await client.get("components"))
+        return "\n".join(f"- {comp['id']}: {comp['name']}" for comp in components)
+
+    @mcp.tool()
+    async def list_component_configs(component_id: str, ctx: Context) -> str:
         """List all configurations for a specific component."""
-        configs = cast(
-            List[Dict[str, Any]], await keboola.get(f"components/{component_id}/configs")
-        )
+        client = ctx.session.state["sapi_client"]
+        assert isinstance(client, KeboolaClient)
+        configs = cast(List[Dict[str, Any]], await client.get(f"components/{component_id}/configs"))
         return "\n".join(
             f"Configuration: {config['id']}\n"
             f"Name: {config['name']}\n"
@@ -246,9 +244,11 @@ async def list_component_configs(component_id: str) -> str:
         )
 
     @mcp.tool()
-    async def list_bucket_tables_tool(bucket_id: str) -> str:
+    async def list_bucket_tables(bucket_id: str, ctx: Context) -> str:
         """List all tables in a specific bucket with their basic information."""
-        tables = cast(List[Dict[str, Any]], keboola.storage_client.buckets.list_tables(bucket_id))
+        client = ctx.session.state["sapi_client"]
+        assert isinstance(client, KeboolaClient)
+        tables = cast(List[Dict[str, Any]], client.storage_client.buckets.list_tables(bucket_id))
         return "\n".join(
             f"Table: {table['id']}\n"
             f"Name: {table.get('name', 'N/A')}\n"
diff --git a/tests/test_config.py b/tests/test_config.py
index 63dd3dc4..cac14509 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -1,59 +1,130 @@
-"""Tests for the config module."""
-
-import os
+from typing import Mapping
 
 import pytest
-from pytest import MonkeyPatch
 
 from keboola_mcp_server.config import Config
 
 
-def test_config_from_env(monkeypatch: MonkeyPatch) -> None:
-    """Test creating config from environment variables."""
-    monkeypatch.setenv("KBC_STORAGE_TOKEN", "test-token")
-    monkeypatch.setenv("KBC_STORAGE_API_URL", "https://test.keboola.com")
-    monkeypatch.setenv("KBC_LOG_LEVEL", "DEBUG")
-
-    config = Config.from_env()
-    assert config.storage_token == "test-token"
-    assert config.storage_api_url == "https://test.keboola.com"
-    assert config.log_level == "DEBUG"
-
-
-def test_config_missing_token(monkeypatch: MonkeyPatch) -> None:
-    """Test error when storage token is missing."""
-    monkeypatch.delenv("KBC_STORAGE_TOKEN", raising=False)
-
-    with pytest.raises(ValueError, match="KBC_STORAGE_TOKEN environment variable is required"):
-        Config.from_env()
-
-
-def test_config_defaults(monkeypatch: MonkeyPatch) -> None:
-    """Test default values are used when optional vars not set."""
-    monkeypatch.setenv("KBC_STORAGE_TOKEN", "test-token")
-    monkeypatch.delenv("KBC_STORAGE_API_URL", raising=False)
-    monkeypatch.delenv("KBC_LOG_LEVEL", raising=False)
-
-    config = Config.from_env()
-    assert config.storage_token == "test-token"
-    assert config.storage_api_url == "https://connection.keboola.com"
-    assert config.log_level == "INFO"
-
-
-def test_config_validate() -> None:
-    """Test config validation."""
-    # Valid config
-    config = Config(
-        storage_token="test-token", storage_api_url="https://test.keboola.com", log_level="INFO"
+class TestConfig:
+    @pytest.mark.parametrize(
+        "d, expected",
+        [
+            (
+                {"storage_token": "foo", "log_level": "DEBUG"},
+                Config(storage_token="foo", log_level="DEBUG"),
+            ),
+            (
+                {"KBC_STORAGE_TOKEN": "foo", "KBC_LOG_LEVEL": "DEBUG"},
+                Config(storage_token="foo", log_level="DEBUG"),
+            ),
+            (
+                {"foo": "bar", "storage_api_url": "http://nowhere"},
+                Config(storage_api_url="http://nowhere"),
+            ),
+        ],
     )
-    config.validate()  # Should not raise
-
-    # Invalid log level
-    config.log_level = "INVALID"
-    with pytest.raises(ValueError, match="Invalid log level"):
-        config.validate()
-
-    # Missing token
-    config = Config(storage_token="", storage_api_url="https://test.keboola.com")
-    with pytest.raises(ValueError, match="Storage token not configured"):
-        config.validate()
+    def test_from_dict(self, d: Mapping[str, str], expected: Config) -> None:
+        assert Config.from_dict(d) == expected
+
+    @pytest.mark.parametrize(
+        "orig, d, expected",
+        [
+            (
+                Config(),
+                {"storage_token": "foo", "log_level": "DEBUG"},
+                Config(storage_token="foo", log_level="DEBUG"),
+            ),
+            (
+                Config(),
+                {"KBC_STORAGE_TOKEN": "foo", "KBC_LOG_LEVEL": "DEBUG"},
+                Config(storage_token="foo", log_level="DEBUG"),
+            ),
+            (
+                Config(storage_token="bar"),
+                {"storage_token": "foo", "log_level": "DEBUG"},
+                Config(storage_token="foo", log_level="DEBUG"),
+            ),
+            (
+                Config(storage_token="bar"),
+                {"storage_token": None, "log_level": "DEBUG"},
+                Config(log_level="DEBUG"),
+            ),
+        ],
+    )
+    def test_replace_by(self, orig: Config, d: Mapping[str, str], expected: Config) -> None:
+        assert orig.replace_by(d) == expected
+
+    def test_defaults(self) -> None:
+        config = Config()
+        assert config.storage_token is None
+        assert config.storage_api_url == "https://connection.keboola.com"
+        assert config.log_level == "INFO"
+        assert config.snowflake_account is None
+        assert config.snowflake_user is None
+        assert config.snowflake_password is None
+        assert config.snowflake_warehouse is None
+        assert config.snowflake_database is None
+        assert config.snowflake_schema is None
+        assert config.snowflake_role is None
+
+    @pytest.mark.parametrize(
+        "d, expected",
+        [
+            ({}, False),
+            ({"storage_token": "foo"}, True),  # relies on the default value of storage_api_url
+            ({"storage_token": "foo", "storage_api_url": ""}, False),
+            ({"storage_token": "foo", "storage_api_url": "bar"}, True),
+        ],
+    )
+    def test_has_storage_config(self, d: Mapping[str, str], expected: bool) -> None:
+        assert Config.from_dict(d).has_storage_config() == expected
+
+    @pytest.mark.parametrize(
+        "d, expected",
+        [
+            ({}, False),
+            ({"snowflake_account": "foo"}, False),
+            ({"snowflake_account": "foo", "snowflake_user": "bar"}, False),
+            (
+                {"snowflake_account": "foo", "snowflake_user": "bar", "snowflake_password": "baz"},
+                False,
+            ),
+            (
+                {
+                    "snowflake_account": "foo",
+                    "snowflake_user": "bar",
+                    "snowflake_password": "baz",
+                    "snowflake_warehouse": "baf",
+                },
+                False,
+            ),
+            (
+                {
+                    "snowflake_account": "foo",
+                    "snowflake_user": "bar",
+                    "snowflake_password": "baz",
+                    "snowflake_warehouse": "baf",
+                    "snowflake_database": "bam",
+                },
+                True,
+            ),
+        ],
+    )
+    def test_has_snowflake_config(self, d: Mapping[str, str], expected: bool) -> None:
+        assert Config.from_dict(d).has_snowflake_config() == expected
+
+    def test_no_token_password_in_repr(self) -> None:
+        config = Config(storage_token="foo", snowflake_password="bar", snowflake_user="baz")
+        assert str(config) == (
+            "Config("
+            "storage_token='****', "
+            "storage_api_url='https://connection.keboola.com', "
+            "log_level='INFO', "
+            "snowflake_account=None, "
+            "snowflake_user='baz', "
+            "snowflake_password='****', "
+            "snowflake_warehouse=None, "
+            "snowflake_database=None, "
+            "snowflake_schema=None, "
+            "snowflake_role=None)"
+        )
diff --git a/tests/test_server.py b/tests/test_server.py
index dcffee8c..cf1152fd 100644
--- a/tests/test_server.py
+++ b/tests/test_server.py
@@ -1,13 +1,11 @@
 """Tests for server functionality."""
 
-from typing import AsyncGenerator, Dict, List, Optional
-from unittest.mock import AsyncMock, MagicMock, patch
+from typing import List, Optional
 
 import pytest
-import snowflake.connector
 
 from keboola_mcp_server.config import Config
-from keboola_mcp_server.server import TableColumnInfo, TableDetail, create_server
+from keboola_mcp_server.server import create_server, TableColumnInfo, TableDetail
 
 
 @pytest.fixture
@@ -45,55 +43,12 @@ def mock_table_detail() -> TableDetail:
     }
 
 
-@pytest.mark.asyncio
-async def test_table_detail_resource(test_config: Config) -> None:
-    mock_table = {
-        "id": "in.c-test.test-table",
-        "name": "test-table",
-        "primaryKey": ["id"],
-        "created": "2024-01-01",
-        "rowsCount": 100,
-        "dataSizeBytes": 1000,
-        "columns": ["id", "name", "value"],
-    }
-
-    with patch("keboola_mcp_server.server.KeboolaClient") as MockKeboolaClient:
-        mock_storage = MagicMock()
-        mock_storage.tables.detail.return_value = mock_table
-        MockKeboolaClient.return_value.storage_client = mock_storage
-
-        server = create_server(test_config)
-
-        old_method = getattr(server, "get_table_detail", None)
-
-        async def mock_get_table_detail(table_id: str):
-            return {
-                "id": mock_table["id"],
-                "name": mock_table["name"],
-                "columns": mock_table["columns"],
-                "column_identifiers": [
-                    TableColumnInfo(name=col, db_identifier=f'"{col}"')
-                    for col in mock_table["columns"]
-                ],
-            }
-
-        setattr(server, "get_table_detail", mock_get_table_detail)
-
-        try:
-            result = await server.get_table_detail("in.c-test.test-table")
-
-            assert result["id"] == "in.c-test.test-table"
-            assert result["name"] == "test-table"
-            assert result["columns"] == ["id", "name", "value"]
-            assert len(result["column_identifiers"]) == 3
-            assert result["column_identifiers"][0]["name"] == "id"
-        finally:
-            if old_method:
-                setattr(server, "get_table_detail", old_method)
-
-
 @pytest.mark.asyncio
 async def test_query_table_data_tool(test_config: Config) -> None:
+    # TODO -- make this test to actually test something; in its current shape it only calls the mock function
+    #   Testing the MCP tools is tricky. Ideally, we would need a test client similar to
+    #   `starlette.testclient.TestClient`, but for the statefull MCP.
+
     mock_table_info = {
         "db_identifier": 'SAPI_10025."in.c-test"."test_table"',
         "column_identifiers": [
@@ -169,108 +124,3 @@ async def query_table_data(
             server.__dict__["query_table"] = original_query_table
         if original_query_table_data:
             server.__dict__["query_table_data"] = original_query_table_data
-
-
-@pytest.mark.asyncio
-async def test_query_table_tool(test_config: Config) -> None:
-    mock_cursor = MagicMock()
-    mock_cursor.fetchall.return_value = [(1, "test", 100)]
-    mock_cursor.description = [("id",), ("name",), ("value",)]
-
-    mock_conn = MagicMock()
-    mock_conn.cursor.return_value = mock_cursor
-
-    with patch("keboola_mcp_server.server.create_snowflake_connection") as mock_create_conn:
-        mock_create_conn.return_value = mock_conn
-
-        server = create_server(test_config)
-
-        # Store original function
-        original_query_table = server.__dict__.get("query_table")
-
-        # Define the actual function we want to test
-        async def query_table(query: str) -> str:
-            try:
-                cursor = mock_conn.cursor()
-                cursor.execute(query)
-                data = cursor.fetchall()
-
-                # Get column names from description
-                columns = [desc[0] for desc in cursor.description]
-
-                # Format as CSV
-                result = [",".join(columns)]
-                result.extend(",".join(str(val) for val in row) for row in data)
-                return "\n".join(result)
-            finally:
-                cursor.close()
-                mock_conn.close()
-
-        # Set our function
-        server.__dict__["query_table"] = query_table
-
-        try:
-            result = await server.query_table('SELECT * FROM SAPI_10025."test_table"')
-
-            assert isinstance(result, str)
-            assert "id,name,value" in result
-            assert "1,test,100" in result
-
-            # Verify proper cleanup
-            mock_cursor.close.assert_called_once()
-            mock_conn.close.assert_called_once()
-        finally:
-            # Restore original if it existed
-            if original_query_table:
-                server.__dict__["query_table"] = original_query_table
-
-
-@pytest.mark.asyncio
-async def test_query_table_error_handling(test_config: Config) -> None:
-    mock_cursor = MagicMock()
-    mock_conn = MagicMock()
-    mock_conn.cursor.return_value = mock_cursor
-
-    # Test Snowflake programming error
-    mock_cursor.execute.side_effect = snowflake.connector.errors.ProgrammingError("Invalid query")
-
-    with patch("keboola_mcp_server.server.create_snowflake_connection") as mock_create_conn:
-        mock_create_conn.return_value = mock_conn
-
-        server = create_server(test_config)
-
-        # Store original function
-        original_query_table = server.__dict__.get("query_table")
-
-        # Define error handling function
-        async def query_table(query: str) -> str:
-            try:
-                cursor = mock_conn.cursor()
-                cursor.execute(query)
-                data = cursor.fetchall()
-                columns = [desc[0] for desc in cursor.description]
-                result = [",".join(columns)]
-                result.extend(",".join(str(val) for val in row) for row in data)
-                return "\n".join(result)
-            except snowflake.connector.errors.ProgrammingError as e:
-                raise ValueError(f"Snowflake query error: {str(e)}")
-            finally:
-                cursor.close()
-                mock_conn.close()
-
-        # Set our function
-        server.__dict__["query_table"] = query_table
-
-        try:
-            with pytest.raises(ValueError) as exc_info:
-                await server.query_table("SELECT * FROM invalid_table")
-
-            assert "Snowflake query error" in str(exc_info.value)
-
-            # Verify cleanup happened even with error
-            mock_cursor.close.assert_called_once()
-            mock_conn.close.assert_called_once()
-        finally:
-            # Restore original if it existed
-            if original_query_table:
-                server.__dict__["query_table"] = original_query_table


Reviews:
user: login: radcliffkey
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T14:30:57Z
user: login: radcliffkey
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T08:51:05Z
user: login: vita-stejskal
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T08:59:38Z
user: login: jordanrburger
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T20:45:27Z

Review comments:
diff_hunk: @@ -0,0 +1,169 @@
+"""
+This is the extension of mcp.server.FastMCP and mcp.server.Server classes that allows to attach the "state"
+to the SSE session. The state is created by the state factory function that can be plugged in to the MCP server,
+and that creates a state which contains arbitrary objects keyed by string identifiers. The factory is given the
+query parameters from the HTTP request that initiates the SSE connection.
+
+Example:
+def factory(params: HttpRequestParams) -> SessionState:
+    return { 'sapi_client': KeboolaClient(params['storage_token']) }
+
+mcp = KeboolaMcpServer(name="SAPI Connector", session_state_factory=factory)
+
+@mcp.tool()
+def list_all_buckets(ctx: Context):
+    client = ctx.session.state['sapi_client']
+    return client.storage_client.buckets.list()
+
+mcp.run(transport='sse')
+
+Issues:
+  * The current implementation of FastMCP does not support sending `Context` to the registered
+    resources' functions. The parameter is passed only to the registered tools.
+"""
+
+import logging
+import os
+from contextlib import AbstractAsyncContextManager, AsyncExitStack
+from typing import Any, Callable
+
+import anyio
+import mcp.types as types
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from mcp import ServerSession, stdio_server
+from mcp.server import FastMCP, Server
+from mcp.server.lowlevel.server import LifespanResultT
+from mcp.server.models import InitializationOptions
+from mcp.server.sse import SseServerTransport
+
+logger = logging.getLogger(__name__)
+
+SessionParams = dict[str, str]
+SessionState = dict[str, Any]
+SessionStateFactory = Callable[[SessionParams], SessionState]
+
+
+def _default_session_state_factory(params: SessionParams) -> SessionState:
+    return params
+
+
+class StatefullServerSession(ServerSession):
+    def __init__(
+        self,
+        read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception],
+        write_stream: MemoryObjectSendStream[types.JSONRPCMessage],
+        init_options: InitializationOptions,
+        state: SessionState | None = None,
+    ) -> None:
+        super().__init__(read_stream, write_stream, init_options)
+        self._state = state or {}
+
+    @property
+    def state(self) -> SessionState:
+        return self._state
+
+
+class _KeboolaServer(Server):
+    def __init__(
+        self,
+        name: str,
+        version: str | None = None,
+        instructions: str | None = None,
+        lifespan: Callable[["Server"], AbstractAsyncContextManager[LifespanResultT]] | None = None,
+    ) -> None:
+        super().__init__(name, version=version, instructions=instructions, lifespan=lifespan)
+
+    async def run(
+        self,
+        read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception],
+        write_stream: MemoryObjectSendStream[types.JSONRPCMessage],
+        initialization_options: InitializationOptions,
+        # When False, exceptions are returned as messages to the client.
+        # When True, exceptions are raised, which will cause the server to shut down
+        # but also make tracing exceptions much easier during testing and when using
+        # in-process servers.
+        raise_exceptions: bool = False,
+        state: SessionState | None = None,
+    ):
+        async with AsyncExitStack() as stack:
+            lifespan_context = await stack.enter_async_context(self.lifespan(self))
+            session = await stack.enter_async_context(
+                StatefullServerSession(read_stream, write_stream, initialization_options, state)
+            )
+
+            async with anyio.create_task_group() as tg:
+                async for message in session.incoming_messages:
+                    logger.debug(f"Received message: {message}")
+
+                    tg.start_soon(
path: src/keboola_mcp_server/mcp.py
user: login: radcliffkey
body: it the server was really stateful, we couldn't really process all messages in parallel, right?
created_at: 2025-03-18T14:30:57Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,169 @@
+"""
+This is the extension of mcp.server.FastMCP and mcp.server.Server classes that allows to attach the "state"
+to the SSE session. The state is created by the state factory function that can be plugged in to the MCP server,
+and that creates a state which contains arbitrary objects keyed by string identifiers. The factory is given the
+query parameters from the HTTP request that initiates the SSE connection.
+
+Example:
+def factory(params: HttpRequestParams) -> SessionState:
+    return { 'sapi_client': KeboolaClient(params['storage_token']) }
+
+mcp = KeboolaMcpServer(name="SAPI Connector", session_state_factory=factory)
+
+@mcp.tool()
+def list_all_buckets(ctx: Context):
+    client = ctx.session.state['sapi_client']
+    return client.storage_client.buckets.list()
+
+mcp.run(transport='sse')
+
+Issues:
+  * The current implementation of FastMCP does not support sending `Context` to the registered
+    resources' functions. The parameter is passed only to the registered tools.
+"""
+
+import logging
+import os
+from contextlib import AbstractAsyncContextManager, AsyncExitStack
+from typing import Any, Callable
+
+import anyio
+import mcp.types as types
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from mcp import ServerSession, stdio_server
+from mcp.server import FastMCP, Server
+from mcp.server.lowlevel.server import LifespanResultT
+from mcp.server.models import InitializationOptions
+from mcp.server.sse import SseServerTransport
+
+logger = logging.getLogger(__name__)
+
+SessionParams = dict[str, str]
+SessionState = dict[str, Any]
+SessionStateFactory = Callable[[SessionParams], SessionState]
+
+
+def _default_session_state_factory(params: SessionParams) -> SessionState:
+    return params
+
+
+class StatefullServerSession(ServerSession):
+    def __init__(
+        self,
+        read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception],
+        write_stream: MemoryObjectSendStream[types.JSONRPCMessage],
+        init_options: InitializationOptions,
+        state: SessionState | None = None,
+    ) -> None:
+        super().__init__(read_stream, write_stream, init_options)
+        self._state = state or {}
+
+    @property
+    def state(self) -> SessionState:
+        return self._state
+
+
+class _KeboolaServer(Server):
+    def __init__(
+        self,
+        name: str,
+        version: str | None = None,
+        instructions: str | None = None,
+        lifespan: Callable[["Server"], AbstractAsyncContextManager[LifespanResultT]] | None = None,
+    ) -> None:
+        super().__init__(name, version=version, instructions=instructions, lifespan=lifespan)
+
+    async def run(
+        self,
+        read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception],
+        write_stream: MemoryObjectSendStream[types.JSONRPCMessage],
+        initialization_options: InitializationOptions,
+        # When False, exceptions are returned as messages to the client.
+        # When True, exceptions are raised, which will cause the server to shut down
+        # but also make tracing exceptions much easier during testing and when using
+        # in-process servers.
+        raise_exceptions: bool = False,
+        state: SessionState | None = None,
+    ):
+        async with AsyncExitStack() as stack:
+            lifespan_context = await stack.enter_async_context(self.lifespan(self))
+            session = await stack.enter_async_context(
+                StatefullServerSession(read_stream, write_stream, initialization_options, state)
+            )
+
+            async with anyio.create_task_group() as tg:
+                async for message in session.incoming_messages:
+                    logger.debug(f"Received message: {message}")
+
+                    tg.start_soon(
path: src/keboola_mcp_server/mcp.py
user: login: vita-stejskal
body: I think you are right. I took this implementation from `FastMCP` and change the bit that creates the `StatefullServerSession`. Our tools are actually stateless and so it should not be a problem if the client sends multiple requests and they are processed asynchronously.
created_at: 2025-03-19T08:59:38Z
author_association: CONTRIBUTOR

PR comments:
user: login: jordanrburger
created_at: 2025-03-17T15:00:35Z
body: From what I remember, the changes in this PR won't actually be compatible with Claude Desktop, right? But it's still important/useful for using the MCP server with Cursor and other services? 
author_association: CONTRIBUTOR
user: login: vita-stejskal
created_at: 2025-03-18T12:32:01Z
body: > From what I remember, the changes in this PR won't actually be compatible with Claude Desktop, right? But it's still important/useful for using the MCP server with Cursor and other services?

The changes on this branch are compatible with Claude Desktop. They have no effect on the "command" transport type which works the same as before -- i.e. you pass the credentials needed by your MCP server via the environment variables.

The Keboola MCP server supports multiple parameters (see [Config](https://github.com/keboola/keboola-mcp-server/blob/KAB-817-statefull-sessions/src/keboola_mcp_server/config.py#L12) class) and they can either be specified via environment variables (use uppercase and prefixed with `KBC_`, e.g. `KBC_STORAGE_TOKEN`) or via the URL parameters (use lowercase, no prefix, e.g. `storage_token`). So, the implementation now supports both types of transport (`command` and `sse`) and they are fully equivalent.
author_association: CONTRIBUTOR
user: login: vita-stejskal
created_at: 2025-03-19T11:02:01Z
body: > From what I remember, the changes in this PR won't actually be compatible with Claude Desktop, right? But it's still important/useful for using the MCP server with Cursor and other services?

AFAIK @mariankrotil successfully tested these changes in Claude Desktop using the `command` transport. So, we should be safe to merge this PR.
author_association: CONTRIBUTOR
user: login: jordanrburger
created_at: 2025-03-19T20:45:19Z
body: 
> > From what I remember, the changes in this PR won't actually be compatible with Claude Desktop, right? But it's still important/useful for using the MCP server with Cursor and other services?
> 
> AFAIK @mariankrotil successfully tested these changes in Claude Desktop using the `command` transport. So, we should be safe to merge this PR.

Awesome! I'll go ahead and merge
author_association: CONTRIBUTOR