1515from sqlalchemy .ext .asyncio import AsyncConnection
1616
1717from mcp_optimizer .db .config import DatabaseConfig
18- from mcp_optimizer .db .exceptions import DbNotFoundError
18+ from mcp_optimizer .db .exceptions import DbNotFoundError , DuplicateRegistryServersError
1919from mcp_optimizer .db .models import (
2020 McpStatus ,
2121 RegistryServer ,
2727from mcp_optimizer .db .workload_server_ops import WorkloadServerOps
2828from mcp_optimizer .db .workload_tool_ops import WorkloadToolOps
2929from 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+ )
3135from mcp_optimizer .token_counter import TokenCounter
3236from mcp_optimizer .toolhive .api_models .core import Workload
3337from mcp_optimizer .toolhive .api_models .registry import ImageMetadata , Registry , RemoteServerMetadata
34- from mcp_optimizer .toolhive .enums import ToolHiveProxyMode , url_to_toolhive_proxy_mode
3538from mcp_optimizer .toolhive .k8s_client import K8sClient
3639from 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