feat: Add batch node_ids support to node_start/stop/reload/suspend

- Each tool accepts either node_id (single) or node_ids (batch)
- Batch mode runs actions in parallel via ThreadPoolExecutor
- Useful for starting/stopping nodes by topology region
This commit is contained in:
YueGuobin 2026-06-14 00:14:34 +08:00
parent 3d3c0eb8db
commit c647805cfb
No known key found for this signature in database
2 changed files with 88 additions and 20 deletions

View File

@ -418,34 +418,58 @@ async def node_get(
@mcp.tool() @mcp.tool()
async def node_start( async def node_start(
project_id: Annotated[str, Field(description="UUID of the project")], project_id: Annotated[str, Field(description="UUID of the project")],
node_id: Annotated[str, Field(description="UUID of the node to start")], node_id: Annotated[str | None, Field(description="Node UUID (single mode)")] = None,
node_ids: Annotated[list[str] | None, Field(description="Batch mode: [\"uuid1\",\"uuid2\"] — start multiple nodes in parallel")] = None,
) -> list[dict[str, Any]]: ) -> list[dict[str, Any]]:
"""Start a node in a project.""" """Start one or more nodes. Provide node_id for single, or node_ids for batch."""
return await asyncio.to_thread(_run_handler_sync, start_node_handler, {"project_id": project_id, "node_id": node_id}) params = {"project_id": project_id}
if node_ids:
params["node_ids"] = node_ids
else:
params["node_id"] = node_id
return await asyncio.to_thread(_run_handler_sync, start_node_handler, params)
@mcp.tool() @mcp.tool()
async def node_stop( async def node_stop(
project_id: Annotated[str, Field(description="UUID of the project")], project_id: Annotated[str, Field(description="UUID of the project")],
node_id: Annotated[str, Field(description="UUID of the node to stop")], node_id: Annotated[str | None, Field(description="Node UUID (single mode)")] = None,
node_ids: Annotated[list[str] | None, Field(description="Batch mode: [\"uuid1\",\"uuid2\"] — stop multiple nodes in parallel")] = None,
) -> list[dict[str, Any]]: ) -> list[dict[str, Any]]:
"""Stop a node in a project.""" """Stop one or more nodes. Provide node_id for single, or node_ids for batch."""
return await asyncio.to_thread(_run_handler_sync, stop_node_handler, {"project_id": project_id, "node_id": node_id}) params = {"project_id": project_id}
if node_ids:
params["node_ids"] = node_ids
else:
params["node_id"] = node_id
return await asyncio.to_thread(_run_handler_sync, stop_node_handler, params)
@mcp.tool() @mcp.tool()
async def node_reload( async def node_reload(
project_id: Annotated[str, Field(description="UUID of the project")], project_id: Annotated[str, Field(description="UUID of the project")],
node_id: Annotated[str, Field(description="UUID of the node to reload")], node_id: Annotated[str | None, Field(description="Node UUID (single mode)")] = None,
node_ids: Annotated[list[str] | None, Field(description="Batch mode: [\"uuid1\",\"uuid2\"] — reload multiple nodes in parallel")] = None,
) -> list[dict[str, Any]]: ) -> list[dict[str, Any]]:
"""Reload (restart) a node in a project.""" """Reload (restart) one or more nodes. Provide node_id for single, or node_ids for batch."""
return await asyncio.to_thread(_run_handler_sync, reload_node_handler, {"project_id": project_id, "node_id": node_id}) params = {"project_id": project_id}
if node_ids:
params["node_ids"] = node_ids
else:
params["node_id"] = node_id
return await asyncio.to_thread(_run_handler_sync, reload_node_handler, params)
@mcp.tool() @mcp.tool()
async def node_suspend( async def node_suspend(
project_id: Annotated[str, Field(description="UUID of the project")], project_id: Annotated[str, Field(description="UUID of the project")],
node_id: Annotated[str, Field(description="UUID of the node to suspend")], node_id: Annotated[str | None, Field(description="Node UUID (single mode)")] = None,
node_ids: Annotated[list[str] | None, Field(description="Batch mode: [\"uuid1\",\"uuid2\"] — suspend multiple nodes in parallel")] = None,
) -> list[dict[str, Any]]: ) -> list[dict[str, Any]]:
"""Suspend a node in a project.""" """Suspend one or more nodes. Provide node_id for single, or node_ids for batch."""
return await asyncio.to_thread(_run_handler_sync, suspend_node_handler, {"project_id": project_id, "node_id": node_id}) params = {"project_id": project_id}
if node_ids:
params["node_ids"] = node_ids
else:
params["node_id"] = node_id
return await asyncio.to_thread(_run_handler_sync, suspend_node_handler, params)
@mcp.tool() @mcp.tool()

View File

@ -109,11 +109,31 @@ def get_node_handler(params: dict[str, Any], gns3_ctx: dict[str, Any]) -> dict[s
return node return node
def _batch_lifecycle(project_id, node_ids, action, conn, action_label):
"""Helper to run a lifecycle action on multiple nodes in parallel."""
def _act(nid):
try:
conn.http_call("post", f"{conn.base_url}/projects/{project_id}/nodes/{nid}/{action}")
return {"node_id": nid, "status": "success", "message": f"Node {nid} {action_label}"}
except Exception as e:
return {"node_id": nid, "status": "error", "error": str(e)}
with ThreadPoolExecutor(max_workers=min(len(node_ids), BATCH_MAX_WORKERS)) as pool:
return list(pool.map(_act, node_ids))
def start_node_handler(params: dict[str, Any], gns3_ctx: dict[str, Any]) -> dict[str, Any]: def start_node_handler(params: dict[str, Any], gns3_ctx: dict[str, Any]) -> dict[str, Any]:
project_id = params.get("project_id") project_id = params.get("project_id")
if not project_id:
return {"error": "project_id is required"}
node_ids = params.get("node_ids")
if node_ids:
if not isinstance(node_ids, list):
return {"error": "node_ids must be a list"}
conn = _get_connector(gns3_ctx)
return _batch_lifecycle(project_id, node_ids, "start", conn, "started")
node_id = params.get("node_id") node_id = params.get("node_id")
if not project_id or not node_id: if not node_id:
return {"error": "project_id and node_id are required"} return {"error": "node_id or node_ids is required"}
conn = _get_connector(gns3_ctx) conn = _get_connector(gns3_ctx)
conn.http_call("post", f"{conn.base_url}/projects/{project_id}/nodes/{node_id}/start", json_data={}) conn.http_call("post", f"{conn.base_url}/projects/{project_id}/nodes/{node_id}/start", json_data={})
return {"message": f"Node {node_id} started", "node_id": node_id} return {"message": f"Node {node_id} started", "node_id": node_id}
@ -121,9 +141,17 @@ def start_node_handler(params: dict[str, Any], gns3_ctx: dict[str, Any]) -> dict
def stop_node_handler(params: dict[str, Any], gns3_ctx: dict[str, Any]) -> dict[str, Any]: def stop_node_handler(params: dict[str, Any], gns3_ctx: dict[str, Any]) -> dict[str, Any]:
project_id = params.get("project_id") project_id = params.get("project_id")
if not project_id:
return {"error": "project_id is required"}
node_ids = params.get("node_ids")
if node_ids:
if not isinstance(node_ids, list):
return {"error": "node_ids must be a list"}
conn = _get_connector(gns3_ctx)
return _batch_lifecycle(project_id, node_ids, "stop", conn, "stopped")
node_id = params.get("node_id") node_id = params.get("node_id")
if not project_id or not node_id: if not node_id:
return {"error": "project_id and node_id are required"} return {"error": "node_id or node_ids is required"}
conn = _get_connector(gns3_ctx) conn = _get_connector(gns3_ctx)
conn.http_call("post", f"{conn.base_url}/projects/{project_id}/nodes/{node_id}/stop", json_data={}) conn.http_call("post", f"{conn.base_url}/projects/{project_id}/nodes/{node_id}/stop", json_data={})
return {"message": f"Node {node_id} stopped", "node_id": node_id} return {"message": f"Node {node_id} stopped", "node_id": node_id}
@ -131,9 +159,17 @@ def stop_node_handler(params: dict[str, Any], gns3_ctx: dict[str, Any]) -> dict[
def reload_node_handler(params: dict[str, Any], gns3_ctx: dict[str, Any]) -> dict[str, Any]: def reload_node_handler(params: dict[str, Any], gns3_ctx: dict[str, Any]) -> dict[str, Any]:
project_id = params.get("project_id") project_id = params.get("project_id")
if not project_id:
return {"error": "project_id is required"}
node_ids = params.get("node_ids")
if node_ids:
if not isinstance(node_ids, list):
return {"error": "node_ids must be a list"}
conn = _get_connector(gns3_ctx)
return _batch_lifecycle(project_id, node_ids, "reload", conn, "reloaded")
node_id = params.get("node_id") node_id = params.get("node_id")
if not project_id or not node_id: if not node_id:
return {"error": "project_id and node_id are required"} return {"error": "node_id or node_ids is required"}
conn = _get_connector(gns3_ctx) conn = _get_connector(gns3_ctx)
conn.http_call("post", f"{conn.base_url}/projects/{project_id}/nodes/{node_id}/reload") conn.http_call("post", f"{conn.base_url}/projects/{project_id}/nodes/{node_id}/reload")
return {"message": f"Node {node_id} reloaded", "node_id": node_id} return {"message": f"Node {node_id} reloaded", "node_id": node_id}
@ -141,9 +177,17 @@ def reload_node_handler(params: dict[str, Any], gns3_ctx: dict[str, Any]) -> dic
def suspend_node_handler(params: dict[str, Any], gns3_ctx: dict[str, Any]) -> dict[str, Any]: def suspend_node_handler(params: dict[str, Any], gns3_ctx: dict[str, Any]) -> dict[str, Any]:
project_id = params.get("project_id") project_id = params.get("project_id")
if not project_id:
return {"error": "project_id is required"}
node_ids = params.get("node_ids")
if node_ids:
if not isinstance(node_ids, list):
return {"error": "node_ids must be a list"}
conn = _get_connector(gns3_ctx)
return _batch_lifecycle(project_id, node_ids, "suspend", conn, "suspended")
node_id = params.get("node_id") node_id = params.get("node_id")
if not project_id or not node_id: if not node_id:
return {"error": "project_id and node_id are required"} return {"error": "node_id or node_ids is required"}
conn = _get_connector(gns3_ctx) conn = _get_connector(gns3_ctx)
conn.http_call("post", f"{conn.base_url}/projects/{project_id}/nodes/{node_id}/suspend") conn.http_call("post", f"{conn.base_url}/projects/{project_id}/nodes/{node_id}/suspend")
return {"message": f"Node {node_id} suspended", "node_id": node_id} return {"message": f"Node {node_id} suspended", "node_id": node_id}