set context var on tool call and use to set org id

Author: fcmsilva

src/api/common.py

@@ -8,6 +8,7 @@
 from src.config.config import get_session_request, get_settings
 from src.logger import get_logger
 from src.utils.async_to_sync import async_to_sync
+from src.api.context import get_session_settings
 
 # Set up logger for this module
 logger = get_logger()
@@ -340,13 +341,17 @@ def get_org_id() -> str | None:
         str or None: The organization ID, or None if using API key authentication
     """
     settings = get_settings()
+    session_settings = get_session_settings()
 
     # If using API key authentication, no org_id is needed
     if not settings.is_remote and settings.api_key:
         logger.debug("Using API key authentication, no organization ID needed")
         return None
 
-    org_id = settings.org_id
+    if settings.is_remote and session_settings is not None:
+        org_id = session_settings.get("org_id", None)
+    else:
+        org_id = settings.org_id
 
     if not org_id:
         logger.debug(

src/api/context.py

@@ -0,0 +1,60 @@
+"""Context variable for the API."""
+
+from contextvars import ContextVar
+from mcp.server.fastmcp import Context
+from functools import wraps
+import inspect
+from typing import Callable, Dict, Any
+
+# Context variable to hold the current context object of a given tool call
+session_context = ContextVar("session_context")
+
+
+def get_session_context() -> Context | None:
+    return session_context.get(None)
+
+
+# Using session's context to store session-specific settings like org_id
+def get_session_settings() -> Dict[str, Any] | None:
+    current_context = get_session_context()
+    if current_context is not None:
+        return current_context.request_context.lifespan_context
+
+    raise Exception(
+        "No session context available for this tool call, please try again."
+    )
+
+
+def tool_wrapper(func: Callable) -> Callable:
+    """
+    This function wraps tool functions to set the context variable.
+    """
+
+    # Check if the underlying function accepts a 'ctx' parameter, or generic **kwargs
+    func_signature = inspect.signature(func)
+    accepts_ctx = "ctx" in func_signature.parameters or any(
+        p.kind == inspect.Parameter.VAR_KEYWORD
+        for p in func_signature.parameters.values()
+    )
+
+    @wraps(func)
+    async def wrapper(ctx, *args, **kwargs):
+        # Ensures that ctx is passed to the original function if it accepts it
+        kwargs_cpy = dict(kwargs)
+        if accepts_ctx:
+            kwargs_cpy["ctx"] = ctx
+
+        current_context = get_session_context()
+        if ctx is not None and current_context is None:
+            session_context.set(ctx.request_context)
+
+        if inspect.iscoroutinefunction(func):
+            return await func(*args, **kwargs_cpy)
+        else:
+            return func(*args, **kwargs_cpy)
+
+    # Ensures the mcp library knows that the wrapper expects a 'ctx' parameter
+    if not accepts_ctx:
+        wrapper.__annotations__["ctx"] = Context
+
+    return wrapper

src/api/tools/organization/organization.py

@@ -9,6 +9,7 @@
 import src.api.tools.organization.utils as utils
 from src.config import config
 from src.api.common import query_graphql_organizations
+from src.api.context import get_session_settings
 from src.utils.elicitation import try_elicitation, ElicitationError
 from src.logger import get_logger
 
@@ -72,6 +73,8 @@ async def choose_organization(ctx: Context) -> dict:
 
     settings = config.get_settings()
     user_id = config.get_user_id()
+    session_settings = get_session_settings()
+
     # Track tool call event
     settings.analytics_manager.track_event(
         user_id, "tool_calling", {"name": "choose_organization"}
@@ -154,7 +157,10 @@ class OrganizationChoice(BaseModel):
 
         # Set the selected organization in settings
         if selected_org:
-            settings.org_id = selected_org["orgID"]
+            if settings.is_remote and session_settings:
+                session_settings["org_id"] = selected_org["orgID"]
+            else:
+                settings.org_id = selected_org["orgID"]
 
             return {
                 "status": "success",
@@ -211,6 +217,7 @@ async def set_organization(ctx: Context, organization_id: str) -> dict:
     3. Call set_organization with the chosen ID
     """
     settings = config.get_settings()
+    session_settings = get_session_settings()
     user_id = config.get_user_id()
     # Track tool call event
     settings.analytics_manager.track_event(
@@ -243,10 +250,13 @@ async def set_organization(ctx: Context, organization_id: str) -> dict:
             }
 
         # Set the selected organization in settings
-        if hasattr(settings, "org_id"):
-            settings.org_id = selected_org["orgID"]
+        if settings.is_remote and session_settings:
+            session_settings["org_id"] = selected_org["orgID"]
         else:
-            setattr(settings, "org_id", selected_org["orgID"])
+            if hasattr(settings, "org_id"):
+                settings.org_id = selected_org["orgID"]
+            else:
+                setattr(settings, "org_id", selected_org["orgID"])
 
         await ctx.info(
             f"Organization set to: {selected_org['name']} (ID: {selected_org['orgID']})"

src/api/tools/tools.py

@@ -29,24 +29,24 @@
 
 # Define the tools with their metadata
 tools_definition = [
-    {"func": get_user_info},
-    {"func": organization_info},
-    {"func": choose_organization},
-    {"func": set_organization},
-    {"func": workspace_groups_info},
-    {"func": workspaces_info},
-    {"func": resume_workspace},
-    {"func": list_starter_workspaces},
-    {"func": create_starter_workspace},
-    {"func": terminate_starter_workspace},
-    {"func": list_regions},
-    {"func": list_sharedtier_regions},
-    {"func": run_sql},
-    {"func": create_notebook_file},
-    {"func": upload_notebook_file},
-    {"func": create_job_from_notebook},
-    {"func": get_job},
-    {"func": delete_job},
+    {"tool": get_user_info},
+    {"tool": organization_info},
+    {"tool": choose_organization},
+    {"tool": set_organization},
+    {"tool": workspace_groups_info},
+    {"tool": workspaces_info},
+    {"tool": resume_workspace},
+    {"tool": list_starter_workspaces},
+    {"tool": create_starter_workspace},
+    {"tool": terminate_starter_workspace},
+    {"tool": list_regions},
+    {"tool": list_sharedtier_regions},
+    {"tool": run_sql},
+    {"tool": create_notebook_file},
+    {"tool": upload_notebook_file},
+    {"tool": create_job_from_notebook},
+    {"tool": get_job},
+    {"tool": delete_job},
 ]
 
 # Export the tools

src/api/types.py

@@ -2,6 +2,8 @@
 from enum import Flag, auto
 from typing import Set
 
+from src.api.context import tool_wrapper
+
 
 # Define all possible flags here - ADD NEW FLAGS TO THIS LIST ONLY!
 AVAILABLE_FLAGS = ["deprecated", "internal"]
@@ -83,6 +85,10 @@ def create_from_dict(cls, concept_def: dict):
                 flag_enum = getattr(MCPConceptFlags, flag_name.upper())
                 flags |= flag_enum
 
+        if "tool" in concept_attrs:
+            concept_attrs["func"] = tool_wrapper(concept_attrs["tool"])
+            del concept_attrs["tool"]
+
         # Set title if not explicitly provided and we have a function
         if "title" not in concept_attrs and "func" in concept_attrs:
             concept_attrs["title"] = getattr(concept_attrs["func"], "__name__", "")