iflow-mcp_xrds76354_sumo-mcp 0.1.0__py3-none-any.whl

mcp_tools/vehicle.py

@@ -0,0 +1,52 @@
+import traci
+from typing import List, Tuple
+from utils.connection import connection_manager
+
+def get_vehicles() -> List[str]:
+    """Get the list of all active vehicle IDs."""
+    if not connection_manager.is_connected():
+        return []
+    return list(traci.vehicle.getIDList())
+
+def get_vehicle_speed(vehicle_id: str) -> float:
+    """Get the speed of a specific vehicle (m/s)."""
+    if not connection_manager.is_connected():
+        raise RuntimeError("Not connected to SUMO.")
+    return float(traci.vehicle.getSpeed(vehicle_id))
+
+def get_vehicle_position(vehicle_id: str) -> Tuple[float, float]:
+    """Get the (x, y) position of a specific vehicle."""
+    if not connection_manager.is_connected():
+        raise RuntimeError("Not connected to SUMO.")
+    x, y = traci.vehicle.getPosition(vehicle_id)
+    return float(x), float(y)
+
+def get_vehicle_acceleration(vehicle_id: str) -> float:
+    """Get the acceleration of a specific vehicle (m/s^2)."""
+    if not connection_manager.is_connected():
+        raise RuntimeError("Not connected to SUMO.")
+    return float(traci.vehicle.getAcceleration(vehicle_id))
+
+def get_vehicle_lane(vehicle_id: str) -> str:
+    """Get the lane ID of a specific vehicle."""
+    if not connection_manager.is_connected():
+        raise RuntimeError("Not connected to SUMO.")
+    return str(traci.vehicle.getLaneID(vehicle_id))
+
+def get_vehicle_route(vehicle_id: str) -> List[str]:
+    """Get the route (list of edge IDs) of a specific vehicle."""
+    if not connection_manager.is_connected():
+        raise RuntimeError("Not connected to SUMO.")
+    return [str(edge) for edge in traci.vehicle.getRoute(vehicle_id)]
+
+def get_simulation_info() -> dict[str, float | int]:
+    """Get general simulation statistics."""
+    if not connection_manager.is_connected():
+        raise RuntimeError("Not connected to SUMO.")
+    return {
+        "time": float(traci.simulation.getTime()),
+        "loaded_vehicles": int(traci.simulation.getLoadedNumber()),
+        "departed_vehicles": int(traci.simulation.getDepartedNumber()),
+        "arrived_vehicles": int(traci.simulation.getArrivedNumber()),
+        "min_expected_vehicles": int(traci.simulation.getMinExpectedNumber()),
+    }

server.py

@@ -0,0 +1,493 @@
+import logging
+import subprocess
+from typing import Any, Dict, Optional
+
+from mcp.server.fastmcp import FastMCP
+
+from utils.traci import ensure_traci_start_stdout_suppressed
+from mcp_tools.simulation import run_simple_simulation
+from mcp_tools.network import netconvert, netgenerate, osm_get
+from mcp_tools.route import random_trips, duarouter, od2trips
+from mcp_tools.signal import tls_cycle_adaptation, tls_coordinator
+from mcp_tools.analysis import analyze_fcd
+from mcp_tools.vehicle import (
+    get_vehicles, get_vehicle_speed, get_vehicle_position, 
+    get_vehicle_acceleration, get_vehicle_lane, get_vehicle_route,
+    get_simulation_info
+)
+from mcp_tools.rl import find_sumo_rl_scenario_files, list_rl_scenarios, run_rl_training
+from utils.connection import connection_manager
+from utils.sumo import find_sumo_binary, find_sumo_home, find_sumo_tools_dir
+from workflows.sim_gen import sim_gen_workflow
+from workflows.signal_opt import signal_opt_workflow
+from workflows.rl_train import rl_train_workflow
+
+# Configure logging to stderr to not interfere with MCP stdio transport
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+# Ensure TraCI never writes to stdout by default (MCP stdio safety).
+ensure_traci_start_stdout_suppressed()
+
+# Initialize MCP Server (official SDK)
+server = FastMCP("SUMO-MCP-Server")
+
+# --- 1. Network Management ---
+@server.tool(description="Manage SUMO network (generate, convert, or download OSM).")
+def manage_network(action: str, output_file: str, params: Optional[Dict[str, Any]] = None) -> str:
+    """
+    actions:
+    - generate: params={'grid': bool, 'grid_number': int}
+    - convert: params={'osm_file': str}
+    - download_osm: output_file is treated as output_dir. params={'bbox': str, 'prefix': str}
+    """
+    params = params or {}
+    options = params.get("options")
+    
+    if action == "generate":
+        spider = bool(params.get("spider", False))
+        grid = bool(params.get("grid", True))
+        grid_number = params.get("grid_number", 3)
+
+        if spider:
+            # Spider network takes precedence over grid settings.
+            grid = False
+            options_list = list(options or [])
+
+            def _strip_flag(flag: str, has_value: bool = False) -> None:
+                while flag in options_list:
+                    idx = options_list.index(flag)
+                    options_list.pop(idx)
+                    if has_value and idx < len(options_list):
+                        options_list.pop(idx)
+
+            def _set_option(flag: str, value: str) -> None:
+                if flag in options_list:
+                    idx = options_list.index(flag)
+                    if idx + 1 < len(options_list):
+                        options_list[idx + 1] = value
+                    else:
+                        options_list.append(value)
+                else:
+                    options_list.extend([flag, value])
+
+            # Enforce Spider/Grid mutual exclusion even when the user provided `options`.
+            _strip_flag("--grid")
+            _strip_flag("--grid.number", has_value=True)
+
+            if "--spider" not in options_list:
+                options_list.insert(0, "--spider")
+
+            arms_raw = params.get("arms", params.get("arm_number"))
+            if arms_raw is not None:
+                try:
+                    arms = int(arms_raw)
+                except (TypeError, ValueError):
+                    return f"Error: arms must be a positive integer, got {arms_raw!r}"
+                if arms <= 0:
+                    return "Error: arms must be > 0"
+                _set_option("--spider.arm-number", str(arms))
+
+            circles_raw = params.get("circles", params.get("circle_number"))
+            if circles_raw is not None:
+                try:
+                    circles = int(circles_raw)
+                except (TypeError, ValueError):
+                    return f"Error: circles must be a positive integer, got {circles_raw!r}"
+                if circles <= 0:
+                    return "Error: circles must be > 0"
+                _set_option("--spider.circle-number", str(circles))
+
+            space_radius_raw = params.get("ring_radius", params.get("space_radius"))
+            if space_radius_raw is not None:
+                try:
+                    space_radius = float(space_radius_raw)
+                except (TypeError, ValueError):
+                    return f"Error: ring_radius must be a number, got {space_radius_raw!r}"
+                if space_radius <= 0:
+                    return "Error: ring_radius must be > 0"
+                _set_option("--spider.space-radius", str(space_radius))
+
+            attach_length_raw = params.get("radial_distance", params.get("attach_length"))
+            if attach_length_raw is not None:
+                try:
+                    attach_length = float(attach_length_raw)
+                except (TypeError, ValueError):
+                    return f"Error: radial_distance must be a number, got {attach_length_raw!r}"
+                if attach_length < 0:
+                    return "Error: radial_distance must be >= 0"
+                _set_option("--spider.attach-length", str(attach_length))
+
+            omit_center_raw = params.get("omit_center")
+            if omit_center_raw:
+                if "--spider.omit-center" not in options_list:
+                    options_list.append("--spider.omit-center")
+
+            options = options_list
+
+        return netgenerate(output_file, grid, grid_number, options)
+        
+    elif action == "convert" or action == "convert_osm":
+        osm_file = params.get("osm_file")
+        if not osm_file: return "Error: osm_file required for convert action"
+        return netconvert(osm_file, output_file, options)
+        
+    elif action == "download_osm":
+        # output_file here acts as output_dir
+        bbox = params.get("bbox")
+        prefix = params.get("prefix", "osm")
+        if not bbox: return "Error: bbox required for download_osm action"
+        return osm_get(bbox, output_file, prefix, options)
+        
+    return f"Unknown action: {action}"
+
+# --- 2. Demand Management ---
+@server.tool(description="Manage traffic demand (random trips, OD matrix, routing).")
+def manage_demand(action: str, net_file: str, output_file: str, params: Optional[Dict[str, Any]] = None) -> str:
+    """
+    actions:
+    - generate_random: params={'end_time': int, 'period': float}
+    - convert_od: params={'od_file': str} (net_file unused but kept for consistency)
+    - compute_routes: params={'route_files': str} (input trips)
+    """
+    params = params or {}
+    options = params.get("options")
+    
+    if action == "generate_random" or action == "random_trips":
+        # Backward/compat aliases: some clients use `end` instead of `end_time`.
+        end_time_raw = params.get("end_time", params.get("end", 3600))
+        period_raw = params.get("period", 1.0)
+        try:
+            end_time = int(end_time_raw)
+        except (TypeError, ValueError):
+            return f"Error: end_time must be an integer, got {end_time_raw!r}"
+        try:
+            period = float(period_raw)
+        except (TypeError, ValueError):
+            return f"Error: period must be a number, got {period_raw!r}"
+        return random_trips(net_file, output_file, end_time, period, options)
+        
+    elif action == "convert_od" or action == "od_matrix":
+        od_file = params.get("od_file")
+        if not od_file: return "Error: od_file required for convert_od"
+        return od2trips(od_file, output_file, options)
+        
+    elif action == "compute_routes" or action == "routing":
+        route_files = params.get("route_files") # Input trips file
+        if not route_files: return "Error: route_files required for compute_routes"
+        return duarouter(net_file, route_files, output_file, options)
+        
+    return f"Unknown action: {action}"
+
+# --- 3. Simulation Control ---
+@server.tool(description="Control SUMO simulation (connect, step, disconnect).")
+def control_simulation(action: str, params: Optional[Dict[str, Any]] = None) -> str:
+    """
+    actions:
+    - connect: params={'config_file': str, 'gui': bool}
+    - step: params={'step': float}
+    - disconnect: no params
+    """
+    params = params or {}
+    
+    try:
+        timeout_s_raw = params.get("timeout_s", params.get("timeout"))
+        timeout_s: Optional[float] = None
+        if timeout_s_raw is not None:
+            try:
+                timeout_s = float(timeout_s_raw)
+            except (TypeError, ValueError):
+                return f"Error: timeout_s must be a number, got {timeout_s_raw!r}"
+
+        if action == "connect":
+            config_file = params.get("config_file")
+            gui = params.get("gui", False)
+            port = params.get("port", 8813)
+            host = params.get("host", "localhost")
+            if timeout_s is None:
+                connection_manager.connect(config_file, gui, port, host)
+            else:
+                connection_manager.connect(config_file, gui, port, host, timeout_s=timeout_s)
+            return "Successfully connected to SUMO."
+            
+        elif action == "step":
+            step = params.get("step", 0)
+            if timeout_s is None:
+                connection_manager.simulation_step(step)
+            else:
+                connection_manager.simulation_step(step, timeout_s=timeout_s)
+            return "Simulation advanced."
+            
+        elif action == "disconnect":
+            if timeout_s is None:
+                connection_manager.disconnect()
+            else:
+                connection_manager.disconnect(timeout_s=timeout_s)
+            return "Successfully disconnected from SUMO."
+            
+    except Exception as e:
+        return f"Error in control_simulation ({action}): {type(e).__name__}: {e}"
+        
+    return f"Unknown action: {action}"
+
+# --- 4. Query State ---
+@server.tool(description="Query simulation state (vehicles, speed, position). Requires active connection.")
+def query_simulation_state(target: str, params: Optional[Dict[str, Any]] = None) -> str:
+    """
+    targets:
+    - vehicle_list: no params
+    - vehicle_variable: params={'vehicle_id': str, 'variable': 'speed'|'position'|'lane'|'acceleration'|'route'}
+    """
+    params = params or {}
+    
+    try:
+        if target == "vehicle_list" or target == "vehicles":
+            vehs = get_vehicles()
+            return f"Active vehicles: {vehs}"
+            
+        elif target == "vehicle_variable":
+            v_id = params.get("vehicle_id")
+            var = params.get("variable")
+            if not v_id or not var: return "Error: vehicle_id and variable required"
+            
+            if var == "speed": return f"Speed: {get_vehicle_speed(v_id)}"
+            if var == "position": return f"Position: {get_vehicle_position(v_id)}"
+            if var == "acceleration": return f"Acceleration: {get_vehicle_acceleration(v_id)}"
+            if var == "lane": return f"Lane: {get_vehicle_lane(v_id)}"
+            if var == "route": return f"Route: {get_vehicle_route(v_id)}"
+            
+            return f"Unknown variable: {var}"
+        
+        elif target == "simulation":
+            info = get_simulation_info()
+            return f"Simulation Info: {info}"
+            
+    except Exception as e:
+        return f"Error querying state: {type(e).__name__}: {e}"
+        
+    return f"Unknown target: {target}"
+
+# --- 5. Optimize Signals ---
+@server.tool(description="Optimize traffic signals.")
+def optimize_traffic_signals(method: str, net_file: str, route_file: str, output_file: str, params: Optional[Dict[str, Any]] = None) -> str:
+    """
+    methods:
+    - cycle_adaptation: adapt TLS cycles
+    - coordination: TLS coordination
+    """
+    params = params or {}
+    options = params.get("options")
+    
+    if method == "cycle_adaptation" or method == "Websters":
+        return tls_cycle_adaptation(net_file, route_file, output_file)
+    elif method == "coordination":
+        return tls_coordinator(net_file, route_file, output_file, options)
+        
+    return f"Unknown method: {method}"
+
+# --- 6. Workflows ---
+@server.tool(
+    description="""Run high-level SUMO workflows. Available workflows:
+
+**sim_gen_eval** - Generate grid network, simulate traffic, analyze results.
+  params:
+  - grid_number (int): Grid size NxN. Default=3. Aliases: grid_size, size
+  - sim_seconds (int): Simulation duration in seconds. Default=100. Aliases: steps, duration, end_time
+  - output_dir (str): Output directory. Default="output"
+  Example: run_workflow("sim_gen_eval", {"grid_number": 3, "sim_seconds": 1000})
+
+**signal_opt** - Optimize traffic signals for existing network.
+  params:
+  - net_file (str): Path to .net.xml file. REQUIRED
+  - route_file (str): Path to .rou.xml file. REQUIRED
+  - sim_seconds (int): Simulation duration. Default=3600. Aliases: steps, duration
+  - use_coordinator (bool): Use tlsCoordinator instead of tlsCycleAdaptation. Default=false
+  - output_dir (str): Output directory. Default="output"
+
+**rl_train** - Train RL agent for traffic signal control.
+  params:
+  - scenario_name (str): Built-in scenario name (use manage_rl_task("list_scenarios") to see options). Aliases: scenario
+  - episodes (int): Number of training episodes. Default=5. Aliases: num_episodes
+  - steps (int): Steps per episode. Default=1000. Aliases: steps_per_episode
+  - output_dir (str): Output directory. Default="output"
+"""
+)
+def run_workflow(workflow_name: str, params: Dict[str, Any]) -> str:
+    """Execute a high-level workflow."""
+
+    # Helper to get param with aliases
+    def get_param(keys: list, default=None):
+        for k in keys:
+            if k in params:
+                return params[k]
+        return default
+
+    if workflow_name in ("sim_gen_eval", "sim_gen_workflow", "sim_gen"):
+        grid_number = get_param(["grid_number", "grid_size", "size"], 3)
+        sim_seconds = get_param(["sim_seconds", "steps", "duration", "end_time"], 100)
+        output_dir = get_param(["output_dir"], "output")
+
+        return sim_gen_workflow(output_dir, int(grid_number), int(sim_seconds))
+
+    elif workflow_name in ("signal_opt", "signal_opt_workflow"):
+        net_file = get_param(["net_file"], "")
+        route_file = get_param(["route_file"], "")
+
+        if not net_file or not route_file:
+            return "Error: signal_opt requires net_file and route_file parameters."
+
+        sim_seconds = get_param(["sim_seconds", "steps", "duration"], 3600)
+        use_coordinator = get_param(["use_coordinator"], False)
+        output_dir = get_param(["output_dir"], "output")
+
+        return signal_opt_workflow(net_file, route_file, output_dir, int(sim_seconds), bool(use_coordinator))
+
+    elif workflow_name == "rl_train":
+        scenario_name = get_param(["scenario_name", "scenario"], "")
+        episodes = get_param(["episodes", "num_episodes"], 5)
+        steps = get_param(["steps", "steps_per_episode"], 1000)
+        output_dir = get_param(["output_dir"], "output")
+
+        return rl_train_workflow(scenario_name, output_dir, int(episodes), int(steps))
+
+    return f"Unknown workflow: {workflow_name}. Available: sim_gen_eval, signal_opt, rl_train"
+
+# --- 7. RL Task Management ---
+@server.tool(description="Manage RL tasks (list scenarios, custom training).")
+def manage_rl_task(action: str, params: Optional[Dict[str, Any]] = None) -> str:
+    """
+    actions:
+    - list_scenarios: no params
+    - train_custom: params={'net_file', 'route_file', 'out_dir', 'episodes', 'steps', 'algorithm', 'reward_type'}
+    """
+    params = params or {}
+    
+    if action == "list_scenarios":
+        return str(list_rl_scenarios())
+        
+    elif action == "train_custom":
+        scenario_name = params.get("scenario") or params.get("scenario_name")
+        net_file = params.get("net_file")
+        route_file = params.get("route_file")
+
+        if scenario_name:
+            net_file, route_file, err = find_sumo_rl_scenario_files(str(scenario_name))
+            if err:
+                return err
+
+        if not net_file or not route_file:
+            return (
+                "Error: train_custom requires either:\n"
+                "  - scenario/scenario_name (built-in sumo-rl scenario), OR\n"
+                "  - net_file + route_file (custom files)\n"
+                "Hint: Use manage_rl_task(list_scenarios) to see available built-in scenarios."
+            )
+
+        out_dir = params.get("out_dir") or params.get("output_dir") or "output"
+
+        episodes_raw = params.get("episodes", params.get("num_episodes", 1))
+        steps_raw = params.get("steps", params.get("steps_per_episode", 1000))
+        try:
+            episodes = int(episodes_raw)
+        except (TypeError, ValueError):
+            return f"Error: episodes must be an integer, got {episodes_raw!r}"
+        try:
+            steps_per_episode = int(steps_raw)
+        except (TypeError, ValueError):
+            return f"Error: steps must be an integer, got {steps_raw!r}"
+
+        if episodes <= 0:
+            return "Error: episodes must be > 0"
+        if steps_per_episode <= 0:
+            return "Error: steps must be > 0"
+
+        algorithm = str(params.get("algorithm", "ql"))
+        reward_type = str(params.get("reward_type", "diff-waiting-time"))
+
+        return run_rl_training(
+            net_file=str(net_file),
+            route_file=str(route_file),
+            out_dir=str(out_dir),
+            episodes=episodes,
+            steps_per_episode=steps_per_episode,
+            algorithm=algorithm,
+            reward_type=reward_type,
+        )
+        
+    return f"Unknown action: {action}"
+
+# --- Legacy/Misc ---
+@server.tool(name="get_sumo_info", description="Get the version and path of the installed SUMO.")
+def get_sumo_info() -> str:
+    try:
+        sumo_binary = find_sumo_binary("sumo")
+        if not sumo_binary:
+            return (
+                "Error: Could not locate SUMO executable. "
+                "Please ensure SUMO is installed and either `sumo` is available in PATH or `SUMO_HOME` is set."
+            )
+
+        result = subprocess.run(
+            [sumo_binary, "--version"],
+            capture_output=True,
+            text=True,
+            check=True,
+            timeout=10,
+        )
+        version_output = (result.stdout.splitlines() or ["Unknown"])[0]
+
+        sumo_home = find_sumo_home()
+        tools_dir = find_sumo_tools_dir()
+        return "\n".join(
+            [
+                f"SUMO Binary: {sumo_binary}",
+                f"SUMO Version: {version_output}",
+                f"SUMO_HOME: {sumo_home or 'Not Set'}",
+                f"SUMO Tools Dir: {tools_dir or 'Not Found'}",
+            ]
+        )
+    except Exception as e:
+        return f"Error checking SUMO: {str(e)}"
+
+@server.tool(name="run_simple_simulation", description="Run a SUMO simulation using a config file.")
+def run_simple_simulation_tool(config_path: str, steps: int = 100) -> str:
+    return run_simple_simulation(config_path, steps)
+
+@server.tool(description="Analyze FCD output.")
+def run_analysis(fcd_file: str) -> str:
+    return analyze_fcd(fcd_file)
+
+if __name__ == "__main__":
+    # NOTE:
+    # MCP stdio transport relies on AnyIO/asyncio to process thread callbacks.
+    # In some environments, a lack of scheduled timers can cause the event loop to
+    # block indefinitely while waiting for stdio worker-thread results. A small
+    # periodic sleep keeps the loop responsive without emitting any stdout output.
+    import anyio
+
+    async def _wakeup_task() -> None:
+        while True:
+            await anyio.sleep(0.1)
+
+    async def _run_stdio_with_wakeup() -> None:
+        async with anyio.create_task_group() as tg:
+            tg.start_soon(_wakeup_task)
+            await server.run_stdio_async()
+
+def main():
+    """Entry point for running the MCP server via uvx."""
+    import anyio
+
+    async def _wakeup_task() -> None:
+        while True:
+            await anyio.sleep(0.1)
+
+    async def _run_stdio_with_wakeup() -> None:
+        async with anyio.create_task_group() as tg:
+            tg.start_soon(_wakeup_task)
+            await server.run_stdio_async()
+
+    anyio.run(_run_stdio_with_wakeup)
+
+if __name__ == "__main__":
+    main()