Skip to content

Commit f12db02

Browse files
Add runtime-aware transport type determination for docker and k8s modes (#214)
- Renamed determine_proxy_mode to determine_transport_type with runtime_mode parameter - Docker mode: checks proxy_mode field (how proxy connects to container) - K8s mode: checks transport_type field (direct connection to pod) - Both modes fall back to URL detection if primary field not set - Added runtime_mode parameter to MCPServerClient constructor (defaults to "docker") - Updated call sites in ingestion.py and server.py to pass runtime_mode - Added explanatory comments about runtime-specific transport determination - Renamed ToolHiveProxyMode to ToolHiveTransportMode for consistency - Updated function name from url_to_toolhive_proxy_mode to url_to_toolhive_transport_mode - Added 7 new tests for k8s mode and cross-mode validation - All 300 tests passing with proper runtime mode separation 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com>
1 parent 30cdd6d commit f12db02

File tree

8 files changed

+385
-236
lines changed

8 files changed

+385
-236
lines changed

helm/mcp-optimizer/Chart.yaml

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,8 @@ apiVersion: v2
22
name: mcp-optimizer
33
description: A Helm chart for deploying MCP Optimizer MCP Server in Kubernetes
44
type: application
5-
version: 0.1.1
6-
appVersion: "0.2.0"
5+
version: 0.2.1
6+
appVersion: "0.2.1"
77
keywords:
88
- mcp
99
- mcp-optimizer
@@ -14,7 +14,3 @@ sources:
1414
- https://github.com/StacklokLabs/mcp-optimizer
1515
maintainers:
1616
- name: MCP Optimizer Team
17-
18-
# Changelog:
19-
# 0.1.1 - Fix groupFiltering.allowedGroups feature with automatic env var merging in podTemplateSpec
20-
# 0.1.0 - Initial release

src/mcp_optimizer/db/models.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
class TransportType(str, Enum):
1919
"""
2020
Enum for transport types.
21-
There is 1:1 relation between ToolHive proxy modes to database transport types.
21+
There is 1:1 relation between ToolHive transport modes to database transport types.
2222
"""
2323

2424
SSE = "sse"

src/mcp_optimizer/ingestion.py

Lines changed: 23 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
from sqlalchemy.ext.asyncio import AsyncConnection
1616

1717
from mcp_optimizer.db.config import DatabaseConfig
18-
from mcp_optimizer.db.exceptions import DbNotFoundError
18+
from mcp_optimizer.db.exceptions import DbNotFoundError, DuplicateRegistryServersError
1919
from mcp_optimizer.db.models import (
2020
McpStatus,
2121
RegistryServer,
@@ -27,11 +27,14 @@
2727
from mcp_optimizer.db.workload_server_ops import WorkloadServerOps
2828
from mcp_optimizer.db.workload_tool_ops import WorkloadToolOps
2929
from mcp_optimizer.embeddings import EmbeddingManager
30-
from mcp_optimizer.mcp_client import MCPServerClient, WorkloadConnectionError
30+
from mcp_optimizer.mcp_client import (
31+
MCPServerClient,
32+
WorkloadConnectionError,
33+
determine_transport_type,
34+
)
3135
from mcp_optimizer.token_counter import TokenCounter
3236
from mcp_optimizer.toolhive.api_models.core import Workload
3337
from mcp_optimizer.toolhive.api_models.registry import ImageMetadata, Registry, RemoteServerMetadata
34-
from mcp_optimizer.toolhive.enums import ToolHiveProxyMode, url_to_toolhive_proxy_mode
3538
from mcp_optimizer.toolhive.k8s_client import K8sClient
3639
from mcp_optimizer.toolhive.toolhive_client import (
3740
ToolhiveClient,
@@ -184,44 +187,6 @@ async def _batch_gather(self, tasks: list[Any], batch_size: int) -> list[Any]:
184187
)
185188
return results
186189

187-
def _map_transport_type(self, workload: Workload) -> TransportType:
188-
"""Map Toolhive transport type to database transport type.
189-
190-
Args:
191-
workload: Workload object with proxy_mode and url
192-
193-
Returns:
194-
Mapped transport type for database storage
195-
196-
Raises:
197-
ValueError: If transport type is not supported
198-
"""
199-
mapping = {
200-
ToolHiveProxyMode.SSE: TransportType.SSE,
201-
ToolHiveProxyMode.STREAMABLE: TransportType.STREAMABLE,
202-
}
203-
204-
# Prefer using the proxy_mode field if available
205-
if workload.proxy_mode:
206-
proxy_mode_str = workload.proxy_mode.lower()
207-
if proxy_mode_str == "sse":
208-
return TransportType.SSE
209-
elif proxy_mode_str == "streamable-http":
210-
return TransportType.STREAMABLE
211-
else:
212-
logger.warning(
213-
f"Unknown proxy_mode '{proxy_mode_str}', falling back to URL detection",
214-
workload=workload.name,
215-
)
216-
217-
# Fallback to URL-based detection for backwards compatibility
218-
if workload.url is None:
219-
raise IngestionError(f"Workload {workload.name} has no URL")
220-
221-
toolhive_proxy_mode = url_to_toolhive_proxy_mode(workload.url)
222-
223-
return mapping[toolhive_proxy_mode]
224-
225190
def _map_workload_status(self, workload_status: str | None) -> McpStatus:
226191
"""Map workload status to McpStatus enum.
227192
@@ -607,8 +572,6 @@ async def _find_and_link_registry_server(
607572
Raises:
608573
DuplicateRegistryServersError: If multiple matching servers found
609574
"""
610-
from mcp_optimizer.db.exceptions import DuplicateRegistryServersError
611-
612575
# Find matching registry servers
613576
matching_servers = await self.registry_server_ops.find_matching_servers(
614577
url=url, package=package, remote=remote, conn=conn
@@ -710,9 +673,8 @@ async def _upsert_workload_server(
710673
ValueError: If workload data is invalid
711674
DuplicateRegistryServersError: If multiple matching registry servers found
712675
"""
713-
from mcp_optimizer.db.exceptions import DuplicateRegistryServersError
714-
715-
transport = self._map_transport_type(workload)
676+
# Cast to TransportType (DB enum) from ToolHiveTransportType
677+
transport = cast(TransportType, determine_transport_type(workload, self.runtime_mode))
716678
status = self._map_workload_status(workload.status)
717679

718680
if not workload.name:
@@ -987,8 +949,6 @@ async def _process_workload(self, workload: Workload, conn: AsyncConnection) ->
987949
Returns:
988950
Processing result with status and counts
989951
"""
990-
from mcp_optimizer.db.exceptions import DuplicateRegistryServersError
991-
992952
result = {
993953
"name": workload.name,
994954
"status": "failed",
@@ -1016,7 +976,9 @@ async def _process_workload(self, workload: Workload, conn: AsyncConnection) ->
1016976
)
1017977

1018978
# Get tools from MCP server
1019-
mcp_client = MCPServerClient(workload, timeout=self.mcp_timeout)
979+
mcp_client = MCPServerClient(
980+
workload, timeout=self.mcp_timeout, runtime_mode=self.runtime_mode
981+
)
1020982
tools_result = await mcp_client.list_tools()
1021983

1022984
# Sync tools with appropriate context
@@ -1027,6 +989,18 @@ async def _process_workload(self, workload: Workload, conn: AsyncConnection) ->
1027989
# Track if anything was updated
1028990
was_updated = server_was_updated or tools_were_updated
1029991

992+
logger.info(
993+
"Processed workload",
994+
server_id=server_id,
995+
workload_name=workload.name,
996+
url=workload.url,
997+
transport_type=workload.transport_type,
998+
group=workload.group,
999+
tools_count=tools_count,
1000+
server_was_updated=server_was_updated,
1001+
tools_were_updated=tools_were_updated,
1002+
)
1003+
10301004
# Calculate autonomous embedding if:
10311005
# 1. Not linked to registry (registry_server_id is None)
10321006
# 2. No server embedding exists

0 commit comments

Comments
 (0)