holmesgpt 0.16.2a0__py3-none-any.whl → 0.18.4__py3-none-any.whl

holmes/plugins/toolsets/kubernetes.yaml

@@ -89,12 +89,124 @@ toolsets:
       - name: "kubernetes_jq_query"
         user_description: "Query Kubernetes Resources: kubectl get {{kind}} --all-namespaces -o json | jq -r {{jq_expr}}"
         description: >
-          Use kubectl to get json for all resources of a specific kind pipe the results to jq to filter them. Do not worry about escaping the jq_expr it will be done by the system on an unescaped expression that you give. e.g. give an expression like .items[] | .spec.containers[].image | select(test("^gcr.io/") | not)
-        command: kubectl get {{ kind }} --all-namespaces -o json | jq -r {{ jq_expr }}
+          Use kubectl to get json for all resources of a specific kind and filter with jq.
+          IMPORTANT: The 'kind' parameter must be the plural form of the resource type
+          (e.g., use "pods" not "pod", "services" not "service", "jobs" not "job").
+          Do not worry about escaping the jq_expr - it will be done by the system.
+          Example: .items[] | .spec.containers[].image | select(test("^gcr.io/") | not)
+        script: |
+          #!/bin/bash
+
+          echo "Executing paginated query for {{ kind }} resources..."
+          echo "Expression: {{ jq_expr }}"
+          echo "---"
+
+          # Get the API path for the resource kind using kubectl
+          API_INFO=$(kubectl api-resources --no-headers | grep "^{{ kind }} " | head -1)
+
+          if [ -z "$API_INFO" ]; then
+            echo "Error: Unable to find resource kind '{{ kind }}'" >&2
+            exit 1
+          fi
+
+          # Extract NAMESPACED value 
+          if [[ "$API_INFO" == *" true "* ]]; then
+            NAMESPACED="true"
+            PREFIX=$(echo "$API_INFO" | sed 's/ true .*//')
+          elif [[ "$API_INFO" == *" false "* ]]; then
+            NAMESPACED="false"
+            PREFIX=$(echo "$API_INFO" | sed 's/ false .*//')
+          else
+            echo "Error: Could not find NAMESPACED field (true/false) in API info" >&2
+            exit 1
+          fi
+
+          # Trim trailing spaces from prefix and collapse internal spaces
+          PREFIX=$(echo "$PREFIX" | sed 's/  *$//' | sed 's/  */ /g')
+
+          IFS=' ' read -ra PREFIX_FIELDS <<< "$PREFIX"
+          FIELD_COUNT=0
+          for field in "${PREFIX_FIELDS[@]}"; do
+            ((FIELD_COUNT++))
+          done
+
+          RESOURCE_NAME="${PREFIX_FIELDS[0]}"
+          if [ $FIELD_COUNT -ge 2 ]; then
+            API_VERSION="${PREFIX_FIELDS[$((FIELD_COUNT - 1))]}"
+          else
+            API_VERSION=""
+          fi
+
+          if [ -z "$API_VERSION" ] || [ -z "$RESOURCE_NAME" ]; then
+            echo "Error: Unable to parse API info for resource kind '{{ kind }}'" >&2
+            exit 1
+          fi
+
+          # Build API path
+          if [[ "$API_VERSION" == "v1" ]]; then
+            API_PATH="/api/v1/${RESOURCE_NAME}"
+          else
+            API_PATH="/apis/${API_VERSION}/${RESOURCE_NAME}"
+          fi
+
+          # Process resources in chunks using API pagination
+          LIMIT=500  # Process 500 items at a time
+          CONTINUE=""
+          PROCESSED=0
+          TOTAL_MATCHES=0
+
+          while true; do
+            # Build API query with limit and continue token
+            if [ -z "$CONTINUE" ]; then
+              # First request - get from all namespaces
+              QUERY="${API_PATH}?limit=${LIMIT}"
+            else
+              # Subsequent requests with continue token
+              QUERY="${API_PATH}?limit=${LIMIT}&continue=${CONTINUE}"
+            fi
+
+            OUTPUT=$(kubectl get --raw "$QUERY" 2>&1)
+            exit_code=$?
+
+            if [ $exit_code -ne 0 ]; then
+              echo "Error: $OUTPUT" >&2
+              exit $exit_code
+            fi
+
+            ITEMS_COUNT=$(echo "$OUTPUT" | jq '.items | length')
+
+            MATCHES=$(echo "$OUTPUT" | jq -r {{ jq_expr }} 2>&1)
+            jq_exit=$?
+            if [ $jq_exit -ne 0 ]; then
+              echo "Error: jq expression failed: $MATCHES" >&2
+              exit $jq_exit
+            fi
+
+            if [ "$ITEMS_COUNT" -gt 0 ]; then
+              if [ -n "$MATCHES" ]; then
+                echo "$MATCHES"
+                MATCH_COUNT=$(echo "$MATCHES" | grep -c . || true)
+                TOTAL_MATCHES=$((TOTAL_MATCHES + MATCH_COUNT))
+              fi
+
+              PROCESSED=$((PROCESSED + ITEMS_COUNT))
+
+              echo "Processed $PROCESSED items, found $TOTAL_MATCHES matches so far..." >&2
+            fi
+
+            CONTINUE=$(echo "$OUTPUT" | jq -r '.metadata.continue // empty')
+
+            if [ -z "$CONTINUE" ]; then
+              break
+            fi
+          done
+
+          echo "---" >&2
+          echo "Total items processed: $PROCESSED, matches found: $TOTAL_MATCHES" >&2
         transformers:
           - name: llm_summarize
             config:
-              input_threshold: 1000
+              input_threshold: 10000
               prompt: |
                 Summarize this jq query output focusing on:
                 - Key patterns and commonalities in the data
@@ -106,52 +218,170 @@ toolsets:
                 - Be concise: aim for ≤ 50% of the original text; prioritize aggregates and actionable outliers
                 - Include grep-ready keys/values; avoid repeating entire objects or unchanged defaults
 
+      - name: "kubernetes_tabular_query"
+        user_description: "Tabular output of specific fields: kubectl get {{kind}} --all-namespaces -o custom-columns={{columns}}"
+        description: >
+          Extract specific fields from Kubernetes resources in tabular format with optional filtering.
+          Memory-efficient way to query large clusters - only requested fields are transmitted.
+          Column specification format: HEADER:FIELD_PATH,HEADER2:FIELD_PATH2,...
+
+          Optional filtering parameter:
+          - filter_pattern: Pattern to match in any column (supports grep regex)
+
+          Examples:
+          - Basic fields: NAME:.metadata.name,STATUS:.status.phase,NODE:.spec.nodeName
+          - Filter by status: filter_pattern="Running"
+          - Filter out lines with <none>: filter_pattern="-v '<none>'"
+          - Nested fields: CREATED:.metadata.creationTimestamp,IMAGE:.spec.containers[0].image
+          - Array fields: LABELS:.metadata.labels,PORTS:.spec.ports[*].port
+
+          Note: Output is tabular text with column headers. Filtering works on the entire line.
+          Note: not allowed characters are: ' / ; and newline
+        command: kubectl get {{ kind }} --all-namespaces -o custom-columns='{{ columns }}'{% if filter_pattern %} | (head -n 1; tail -n +2 | grep {{ filter_pattern }}){% endif %}
+        transformers:
+          - name: llm_summarize
+            config:
+              input_threshold: 10000
+              prompt: |
+                Summarize this tabular output focusing on:
+                - Key patterns and trends in the data
+                - Resources that need attention (errors, pending, failures)
+                - Group similar items into aggregate descriptions
+                - Highlight outliers or unusual values
+                - Mention specific resource names only for problematic items
+                - Provide counts and distributions where relevant
+                - Be concise: aim for ≤ 50% of the original size
+                - Keep output actionable and focused on anomalies
+
       - name: "kubernetes_count"
         user_description: "Count Kubernetes Resources: kubectl get {{kind}} --all-namespaces -o json | jq -c -r {{ jq_expr }}"
         description: >
           Use kubectl to get apply a jq filter and then count the results.
           Use this whenever asked to count kubernetes resources.
+          IMPORTANT: The 'kind' parameter must be the plural form of the resource type
+          (e.g., use "pods" not "pod", "services" not "service", "jobs" not "job").
           Use select() to filter objects before extracting properties, e.g. .items[] | select(.metadata.namespace == "test-1") | .metadata.name
           Do not worry about escaping the jq_expr it will be done by the system on an unescaped expression that you give.
           e.g. give an expression like .items[] | select(.spec.containers[].image | test("^gcr.io/") | not) | .metadata.name
         script: |
+          #!/bin/bash
+
           echo "Command executed: kubectl get {{ kind }} --all-namespaces -o json | jq -c -r {{ jq_expr }}"
           echo "---"
 
-          # Execute the command and capture both stdout and stderr separately
-          temp_error=$(mktemp)
-          matches=$(kubectl get {{ kind }} --all-namespaces -o json 2>"$temp_error" | jq -c -r {{ jq_expr }} 2>>"$temp_error")
-          exit_code=$?
-          error_output=$(cat "$temp_error")
-          rm -f "$temp_error"
-
-          if [ $exit_code -ne 0 ]; then
-            echo "Error executing command (exit code: $exit_code):"
-            echo "$error_output"
-            exit $exit_code
+          # Get the API path for the resource kind
+          API_INFO=$(kubectl api-resources --no-headers | grep "^{{ kind }} " | head -1)
+
+          if [ -z "$API_INFO" ]; then
+            echo "Error: Unable to find resource kind '{{ kind }}'" >&2
+            exit 1
+          fi
+
+          if [[ "$API_INFO" == *" true "* ]]; then
+            NAMESPACED="true"
+            PREFIX=$(echo "$API_INFO" | sed 's/ true .*//')
+          elif [[ "$API_INFO" == *" false "* ]]; then
+            NAMESPACED="false"
+            PREFIX=$(echo "$API_INFO" | sed 's/ false .*//')
           else
-            # Show any stderr warnings even if command succeeded
-            if [ -n "$error_output" ]; then
-              echo "Warnings/stderr output:"
-              echo "$error_output"
-              echo "---"
-            fi
+            echo "Error: Could not find NAMESPACED field (true/false) in API info" >&2
+            exit 1
+          fi
+
+          PREFIX=$(echo "$PREFIX" | sed 's/  *$//' | sed 's/  */ /g')
+
+          IFS=' ' read -ra PREFIX_FIELDS <<< "$PREFIX"
+          FIELD_COUNT=0
+          for field in "${PREFIX_FIELDS[@]}"; do
+            ((FIELD_COUNT++))
+          done
+          RESOURCE_NAME="${PREFIX_FIELDS[0]}"
+
+          if [ $FIELD_COUNT -ge 2 ]; then
+            API_VERSION="${PREFIX_FIELDS[$((FIELD_COUNT - 1))]}"
+          else
+            API_VERSION=""
+          fi
+
+          if [ -z "$API_VERSION" ] || [ -z "$RESOURCE_NAME" ]; then
+            echo "Error: Unable to parse API info for resource kind '{{ kind }}'" >&2
+            exit 1
+          fi
 
-            # Filter out empty lines for accurate count
-            filtered_matches=$(echo "$matches" | grep -v '^$' | grep -v '^null$')
-            if [ -z "$filtered_matches" ]; then
-              count=0
+          # Build API path
+          if [[ "$API_VERSION" == "v1" ]]; then
+            API_PATH="/api/v1/${RESOURCE_NAME}"
+          else
+            API_PATH="/apis/${API_VERSION}/${RESOURCE_NAME}"
+          fi
+
+          # Process resources in chunks using API pagination
+          LIMIT=500
+          CONTINUE=""
+          ALL_MATCHES=""
+          BATCH_NUM=0
+          TOTAL_PROCESSED=0
+
+          while true; do
+            BATCH_NUM=$((BATCH_NUM + 1))
+
+            if [ -z "$CONTINUE" ]; then
+              QUERY="${API_PATH}?limit=${LIMIT}"
             else
-              count=$(echo "$filtered_matches" | wc -l)
+              QUERY="${API_PATH}?limit=${LIMIT}&continue=${CONTINUE}"
+            fi
+
+            OUTPUT=$(kubectl get --raw "$QUERY" 2>&1)
+            exit_code=$?
+
+            if [ $exit_code -ne 0 ]; then
+              echo "Error for query $QUERY: $OUTPUT" >&2
+              exit $exit_code
+            fi
+
+            ITEMS_COUNT=$(echo "$OUTPUT" | jq '.items | length')
+            TOTAL_PROCESSED=$((TOTAL_PROCESSED + ITEMS_COUNT))
+
+            BATCH_MATCHES=$(echo "$OUTPUT" | jq -c -r {{ jq_expr }} 2>&1)
+            jq_exit=$?
+            if [ $jq_exit -ne 0 ]; then
+              echo "Error: jq expression failed: $BATCH_MATCHES" >&2
+              exit $jq_exit
             fi
-            preview=$(echo "$filtered_matches" | head -n 10 | cut -c 1-200 | nl)
 
-            echo "$count results"
-            echo "---"
-            echo "A *preview* of results is shown below (up to 10 results, up to 200 chars):"
-            echo "$preview"
+            if [ -n "$BATCH_MATCHES" ]; then
+              if [ -z "$ALL_MATCHES" ]; then
+                ALL_MATCHES="$BATCH_MATCHES"
+              else
+                ALL_MATCHES="$ALL_MATCHES"$'\n'"$BATCH_MATCHES"
+              fi
+            fi
+
+            CONTINUE=$(echo "$OUTPUT" | jq -r '.metadata.continue // empty')
+            if [ -z "$CONTINUE" ]; then
+              break
+            fi
+
+            echo "Processed batch $BATCH_NUM ($TOTAL_PROCESSED items so far)..." >&2
+          done
+
+          # Now process the collected matches
+          filtered_matches=$(echo "$ALL_MATCHES" | grep -v '^$' | grep -v '^null$')
+          if [ -z "$filtered_matches" ]; then
+            count=0
+            preview=""
+          else
+            count=$(echo "$filtered_matches" | wc -l)
+            preview=$(echo "$filtered_matches" | head -n 10 | cut -c 1-200 | nl)
           fi
 
+          echo "$count results"
+          echo "---"
+          echo "A *preview* of results is shown below (up to 10 results, up to 200 chars):"
+          echo "$preview"
+          echo "---"
+          echo "Total items processed: $TOTAL_PROCESSED" >&2
+
     # NOTE: this is only possible for probes with a healthz endpoint - we do this to avoid giving the LLM generic
     # http GET capabilities which are more powerful than we want to expose
     #- name: "check_liveness_probe"

holmes/plugins/toolsets/kubernetes_logs.py

@@ -3,7 +3,8 @@ import re
 import subprocess
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from datetime import datetime, timezone
-from typing import Optional, List, Tuple, Set
+from typing import List, Optional, Set, Tuple
+
 from pydantic import BaseModel
 
 from holmes.common.env_vars import KUBERNETES_LOGS_TIMEOUT_SECONDS
@@ -14,16 +15,15 @@ from holmes.core.tools import (
     ToolsetTag,
 )
 from holmes.plugins.toolsets.logging_utils.logging_api import (
+    DEFAULT_TIME_SPAN_SECONDS,
     BasePodLoggingToolset,
     FetchPodLogsParams,
     LoggingCapability,
     LoggingConfig,
     PodLoggingTool,
-    DEFAULT_TIME_SPAN_SECONDS,
 )
 from holmes.plugins.toolsets.utils import process_timestamps_to_int, to_unix_ms
 
-
 # match ISO 8601 format (YYYY-MM-DDTHH:MM:SS[.fffffffff]Z) or (YYYY-MM-DDTHH:MM:SS[.fffffffff]+/-XX:XX)
 timestamp_pattern = re.compile(
     r"^(?P<ts>\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(?:\.\d+)?(?:Z|[+-]\d{2}:\d{2}))"

holmes/plugins/toolsets/logging_utils/logging_api.py

@@ -1,12 +1,12 @@
-from abc import ABC, abstractmethod
-from datetime import datetime, timedelta
 import logging
+from abc import ABC, abstractmethod
+from datetime import datetime, timedelta, timezone
+from enum import Enum
 from math import ceil
 from typing import Optional, Set
-from enum import Enum
 
 from pydantic import BaseModel, field_validator
-from datetime import timezone
+
 from holmes.core.llm import LLM
 from holmes.core.tools import (
     StructuredToolResult,
@@ -88,9 +88,14 @@ def truncate_logs(
     llm: LLM,
     token_limit: int,
     structured_params: FetchPodLogsParams,
+    tool_call_id: str,
+    tool_name: str,
 ):
     original_token_count = count_tool_response_tokens(
-        llm=llm, structured_tool_result=logging_structured_tool_result
+        llm=llm,
+        structured_tool_result=logging_structured_tool_result,
+        tool_call_id=tool_call_id,
+        tool_name=tool_name,
     )
     token_count = original_token_count
     text = None
@@ -137,7 +142,10 @@ def truncate_logs(
             )
             logging_structured_tool_result.data = text
             token_count = count_tool_response_tokens(
-                llm=llm, structured_tool_result=logging_structured_tool_result
+                llm=llm,
+                structured_tool_result=logging_structured_tool_result,
+                tool_call_id=tool_call_id,
+                tool_name=tool_name,
             )
     if token_count < original_token_count:
         logging.info(
@@ -266,6 +274,8 @@ If you hit the log limit and see lots of repetitive INFO logs, use exclude_filte
             llm=context.llm,
             token_limit=context.max_token_count,
             structured_params=structured_params,
+            tool_call_id=context.tool_call_id,
+            tool_name=context.tool_name,
         )
 
         return result

holmes/plugins/toolsets/mcp/toolset_mcp.py

@@ -1,30 +1,28 @@
+import asyncio
 import json
+import logging
+import threading
+from contextlib import asynccontextmanager
+from enum import Enum
+from typing import Any, Dict, List, Optional, Tuple, Union
 
-from holmes.common.env_vars import SSE_READ_TIMEOUT
-from holmes.core.tools import (
-    ToolInvokeContext,
-    Toolset,
-    Tool,
-    ToolParameter,
-    StructuredToolResult,
-    StructuredToolResultStatus,
-    CallablePrerequisite,
-)
-
-from typing import Dict, Any, List, Optional
 from mcp.client.session import ClientSession
 from mcp.client.sse import sse_client
+from mcp.client.stdio import StdioServerParameters, stdio_client
 from mcp.client.streamable_http import streamablehttp_client
-
 from mcp.types import Tool as MCP_Tool
+from pydantic import AnyUrl, BaseModel, Field, model_validator
 
-import asyncio
-from contextlib import asynccontextmanager
-from pydantic import BaseModel, Field, AnyUrl, model_validator
-from typing import Tuple
-import logging
-from enum import Enum
-import threading
+from holmes.common.env_vars import SSE_READ_TIMEOUT
+from holmes.core.tools import (
+    CallablePrerequisite,
+    StructuredToolResult,
+    StructuredToolResultStatus,
+    Tool,
+    ToolInvokeContext,
+    ToolParameter,
+    Toolset,
+)
 
 # Lock per MCP server URL to serialize calls to the same server
 _server_locks: Dict[str, threading.Lock] = {}
@@ -42,6 +40,7 @@ def get_server_lock(url: str) -> threading.Lock:
 class MCPMode(str, Enum):
     SSE = "sse"
     STREAMABLE_HTTP = "streamable-http"
+    STDIO = "stdio"
 
 
 class MCPConfig(BaseModel):
@@ -49,14 +48,42 @@ class MCPConfig(BaseModel):
     mode: MCPMode = MCPMode.SSE
     headers: Optional[Dict[str, str]] = None
 
+    def get_lock_string(self) -> str:
+        return str(self.url)
+
+
+class StdioMCPConfig(BaseModel):
+    mode: MCPMode = MCPMode.STDIO
+    command: str
+    args: Optional[List[str]] = None
+    env: Optional[Dict[str, str]] = None
+
+    def get_lock_string(self) -> str:
+        return str(self.command)
+
 
 @asynccontextmanager
-async def get_initialized_mcp_session(
-    url: str, headers: Optional[Dict[str, str]], mode: MCPMode
-):
-    if mode == MCPMode.SSE:
+async def get_initialized_mcp_session(toolset: "RemoteMCPToolset"):
+    if toolset._mcp_config is None:
+        raise ValueError("MCP config is not initialized")
+
+    if isinstance(toolset._mcp_config, StdioMCPConfig):
+        server_params = StdioServerParameters(
+            command=toolset._mcp_config.command,
+            args=toolset._mcp_config.args or [],
+            env=toolset._mcp_config.env,
+        )
+        async with stdio_client(server_params) as (
+            read_stream,
+            write_stream,
+        ):
+            async with ClientSession(read_stream, write_stream) as session:
+                _ = await session.initialize()
+                yield session
+    elif toolset._mcp_config.mode == MCPMode.SSE:
+        url = str(toolset._mcp_config.url)
         async with sse_client(
-            url, headers=headers, sse_read_timeout=SSE_READ_TIMEOUT
+            url, toolset._mcp_config.headers, sse_read_timeout=SSE_READ_TIMEOUT
         ) as (
             read_stream,
             write_stream,
@@ -65,8 +92,9 @@ async def get_initialized_mcp_session(
                 _ = await session.initialize()
                 yield session
     else:
+        url = str(toolset._mcp_config.url)
         async with streamablehttp_client(
-            url, headers=headers, sse_read_timeout=SSE_READ_TIMEOUT
+            url, headers=toolset._mcp_config.headers, sse_read_timeout=SSE_READ_TIMEOUT
         ) as (
             read_stream,
             write_stream,
@@ -86,7 +114,8 @@ class RemoteMCPTool(Tool):
             # Different servers can still run in parallel
             if not self.toolset._mcp_config:
                 raise ValueError("MCP config not initialized")
-            lock = get_server_lock(str(self.toolset._mcp_config.url))
+
+            lock = get_server_lock(str(self.toolset._mcp_config.get_lock_string()))
             with lock:
                 return asyncio.run(self._invoke_async(params))
         except Exception as e:
@@ -107,7 +136,7 @@ class RemoteMCPTool(Tool):
             return False
 
     async def _invoke_async(self, params: Dict) -> StructuredToolResult:
-        async with self.toolset.get_initialized_session() as session:
+        async with get_initialized_mcp_session(self.toolset) as session:
             tool_result = await session.call_tool(self.name, params)
 
         merged_text = " ".join(c.text for c in tool_result.content if c.type == "text")
@@ -153,20 +182,23 @@ class RemoteMCPTool(Tool):
         return parameters
 
     def get_parameterized_one_liner(self, params: Dict) -> str:
-        if params:
-            if params.get("cli_command"):  # Return AWS MCP cli command, if available
-                return f"{params.get('cli_command')}"
+        # AWS MCP cli_command
+        if params and params.get("cli_command"):
+            return f"{params.get('cli_command')}"
 
-        url = (
-            str(self.toolset._mcp_config.url) if self.toolset._mcp_config else "unknown"
-        )
-        return f"Call MCP Server ({url} - {self.name})"
+        # gcloud MCP run_gcloud_command
+        if self.name == "run_gcloud_command" and params and "args" in params:
+            args = params.get("args", [])
+            if isinstance(args, list):
+                return f"gcloud {' '.join(str(arg) for arg in args)}"
+
+        return f"{self.toolset.name}: {self.name} {params}"
 
 
 class RemoteMCPToolset(Toolset):
     tools: List[RemoteMCPTool] = Field(default_factory=list)  # type: ignore
     icon_url: str = "https://registry.npmmirror.com/@lobehub/icons-static-png/1.46.0/files/light/mcp.png"
-    _mcp_config: Optional[MCPConfig] = None
+    _mcp_config: Optional[Union[MCPConfig, StdioMCPConfig]] = None
 
     def model_post_init(self, __context: Any) -> None:
         self.prerequisites = [
@@ -211,23 +243,24 @@ class RemoteMCPToolset(Toolset):
             if not config:
                 return (False, f"Config is required for {self.name}")
 
-            if "mode" in config:
-                mode_value = config.get("mode")
-                allowed_modes = [e.value for e in MCPMode]
-                if mode_value not in allowed_modes:
-                    return (
-                        False,
-                        f'Invalid mode "{mode_value}", allowed modes are {", ".join(allowed_modes)}',
-                    )
-
-            self._mcp_config = MCPConfig(**config)
-
-            clean_url_str = str(self._mcp_config.url).rstrip("/")
-
-            if self._mcp_config.mode == MCPMode.SSE and not clean_url_str.endswith(
-                "/sse"
-            ):
-                self._mcp_config.url = AnyUrl(clean_url_str + "/sse")
+            mode_value = config.get("mode", MCPMode.SSE.value)
+            allowed_modes = [e.value for e in MCPMode]
+            if mode_value not in allowed_modes:
+                return (
+                    False,
+                    f'Invalid mode "{mode_value}", allowed modes are {", ".join(allowed_modes)}',
+                )
+
+            if mode_value == MCPMode.STDIO.value:
+                self._mcp_config = StdioMCPConfig(**config)
+            else:
+                self._mcp_config = MCPConfig(**config)
+                clean_url_str = str(self._mcp_config.url).rstrip("/")
+
+                if self._mcp_config.mode == MCPMode.SSE and not clean_url_str.endswith(
+                    "/sse"
+                ):
+                    self._mcp_config.url = AnyUrl(clean_url_str + "/sse")
 
             tools_result = asyncio.run(self._get_server_tools())
 
@@ -242,18 +275,13 @@ class RemoteMCPToolset(Toolset):
         except Exception as e:
             return (
                 False,
-                f"Failed to load mcp server {self.name} {self._mcp_config.url if self._mcp_config else 'unknown'}: {str(e)}",
+                f"Failed to load mcp server {self.name}: {str(e)}",
             )
 
     async def _get_server_tools(self):
-        async with self.get_initialized_session() as session:
+        async with get_initialized_mcp_session(self) as session:
             return await session.list_tools()
 
-    def get_initialized_session(self):
-        return get_initialized_mcp_session(
-            str(self._mcp_config.url), self._mcp_config.headers, self._mcp_config.mode
-        )
-
     def get_example_config(self) -> Dict[str, Any]:
         example_config = MCPConfig(
             url=AnyUrl("http://example.com:8000/mcp/messages"),