|  |  | @@ -1,8 +1,25 @@ | 
		
	
		
			
			|  |  |  | # | 
		
	
		
			
			|  |  |  | #  Copyright 2025 The InfiniFlow Authors. All Rights Reserved. | 
		
	
		
			
			|  |  |  | # | 
		
	
		
			
			|  |  |  | #  Licensed under the Apache License, Version 2.0 (the "License"); | 
		
	
		
			
			|  |  |  | #  you may not use this file except in compliance with the License. | 
		
	
		
			
			|  |  |  | #  You may obtain a copy of the License at | 
		
	
		
			
			|  |  |  | # | 
		
	
		
			
			|  |  |  | #      http://www.apache.org/licenses/LICENSE-2.0 | 
		
	
		
			
			|  |  |  | # | 
		
	
		
			
			|  |  |  | #  Unless required by applicable law or agreed to in writing, software | 
		
	
		
			
			|  |  |  | #  distributed under the License is distributed on an "AS IS" BASIS, | 
		
	
		
			
			|  |  |  | #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | 
		
	
		
			
			|  |  |  | #  See the License for the specific language governing permissions and | 
		
	
		
			
			|  |  |  | #  limitations under the License. | 
		
	
		
			
			|  |  |  | # | 
		
	
		
			
			|  |  |  | 
 | 
		
	
		
			
			|  |  |  | import asyncio | 
		
	
		
			
			|  |  |  | import logging | 
		
	
		
			
			|  |  |  | import threading | 
		
	
		
			
			|  |  |  | import weakref | 
		
	
		
			
			|  |  |  | from concurrent.futures import ThreadPoolExecutor | 
		
	
		
			
			|  |  |  | from concurrent.futures import TimeoutError as FuturesTimeoutError | 
		
	
		
			
			|  |  |  | from string import Template | 
		
	
		
			
			|  |  |  | from typing import Any, Literal | 
		
	
		
			
			|  |  |  | 
 | 
		
	
	
		
			
			|  |  | @@ -101,7 +118,7 @@ class MCPToolCallSession(ToolCallSession): | 
		
	
		
			
			|  |  |  | try: | 
		
	
		
			
			|  |  |  | result: CallToolResult | Exception = await asyncio.wait_for(results.get(), timeout=timeout) | 
		
	
		
			
			|  |  |  | except asyncio.TimeoutError: | 
		
	
		
			
			|  |  |  | raise TimeoutError(f"MCP task '{task_type}' timeout after {timeout}s") | 
		
	
		
			
			|  |  |  | raise asyncio.TimeoutError(f"MCP task '{task_type}' timeout after {timeout}s") | 
		
	
		
			
			|  |  |  | 
 | 
		
	
		
			
			|  |  |  | if isinstance(result, Exception): | 
		
	
		
			
			|  |  |  | raise result | 
		
	
	
		
			
			|  |  | @@ -128,8 +145,8 @@ class MCPToolCallSession(ToolCallSession): | 
		
	
		
			
			|  |  |  | future = asyncio.run_coroutine_threadsafe(self._get_tools_from_mcp_server(), self._event_loop) | 
		
	
		
			
			|  |  |  | try: | 
		
	
		
			
			|  |  |  | return future.result(timeout=timeout) | 
		
	
		
			
			|  |  |  | except TimeoutError: | 
		
	
		
			
			|  |  |  | logging.error(f"Timeout when fetching tools from MCP server: {self._mcp_server.id}") | 
		
	
		
			
			|  |  |  | except FuturesTimeoutError: | 
		
	
		
			
			|  |  |  | logging.error(f"Timeout when fetching tools from MCP server: {self._mcp_server.id} (timeout={timeout})") | 
		
	
		
			
			|  |  |  | return [] | 
		
	
		
			
			|  |  |  | except Exception: | 
		
	
		
			
			|  |  |  | logging.exception(f"Error fetching tools from MCP server: {self._mcp_server.id}") | 
		
	
	
		
			
			|  |  | @@ -140,9 +157,9 @@ class MCPToolCallSession(ToolCallSession): | 
		
	
		
			
			|  |  |  | future = asyncio.run_coroutine_threadsafe(self._call_mcp_tool(name, arguments), self._event_loop) | 
		
	
		
			
			|  |  |  | try: | 
		
	
		
			
			|  |  |  | return future.result(timeout=timeout) | 
		
	
		
			
			|  |  |  | except TimeoutError as te: | 
		
	
		
			
			|  |  |  | logging.error(f"Timeout calling tool '{name}' on MCP server: {self._mcp_server.id}") | 
		
	
		
			
			|  |  |  | return f"Timeout calling tool '{name}': {te}." | 
		
	
		
			
			|  |  |  | except FuturesTimeoutError: | 
		
	
		
			
			|  |  |  | logging.error(f"Timeout calling tool '{name}' on MCP server: {self._mcp_server.id} (timeout={timeout})") | 
		
	
		
			
			|  |  |  | return f"Timeout calling tool '{name}' (timeout={timeout})." | 
		
	
		
			
			|  |  |  | except Exception as e: | 
		
	
		
			
			|  |  |  | logging.exception(f"Error calling tool '{name}' on MCP server: {self._mcp_server.id}") | 
		
	
		
			
			|  |  |  | return f"Error calling tool '{name}': {e}." | 
		
	
	
		
			
			|  |  | @@ -164,8 +181,8 @@ class MCPToolCallSession(ToolCallSession): | 
		
	
		
			
			|  |  |  | future = asyncio.run_coroutine_threadsafe(self.close(), self._event_loop) | 
		
	
		
			
			|  |  |  | try: | 
		
	
		
			
			|  |  |  | future.result(timeout=timeout) | 
		
	
		
			
			|  |  |  | except TimeoutError: | 
		
	
		
			
			|  |  |  | logging.error(f"Timeout while closing session for server {self._mcp_server.id}") | 
		
	
		
			
			|  |  |  | except FuturesTimeoutError: | 
		
	
		
			
			|  |  |  | logging.error(f"Timeout while closing session for server {self._mcp_server.id} (timeout={timeout})") | 
		
	
		
			
			|  |  |  | except Exception: | 
		
	
		
			
			|  |  |  | logging.exception(f"Unexpected error during close_sync for {self._mcp_server.id}") | 
		
	
		
			
			|  |  |  | 
 |