Merge pull request #6 from Razee4315/feat/security-quality-and-feature-improvements

feat: security hardening + close_model/server_info tools + dead code cleanup

Author: Razee4315

CHANGELOG.md

@@ -7,6 +7,37 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Security & validation
+
+- `set_parameter` now validates the `name` argument against a NetLogo
+  identifier regex before interpolating it into a `set <name> <value>`
+  command. Closes a command-injection vector where a name like
+  `"x setup -- "` would have appended an arbitrary NetLogo command past
+  the `set`. Legitimate kebab-case + predicate names (`show-energy?`,
+  `initial-number-sheep`) are unaffected.
+- `run_simulation` rejects empty / non-string entries in `reporters`,
+  and `get_patch_data` rejects blank `attribute`s — these formerly
+  produced confusing NetLogo compile errors deep in the call.
+- 25 new tests cover the validation surface (injection-shape names,
+  blank reporters, blank attributes).
+
+### Added — workspace control & introspection
+
+- New `close_model` tool — issues `clear-all` and forgets the current
+  model path so AI clients can drop pending state without bouncing the
+  JVM (which costs 30-60s on the next startup).
+- New `server_info` tool — pure config / filesystem inspection that
+  returns server version, GUI mode, configured paths, and whether the
+  BehaviorSpace headless launcher is reachable. Useful as a pre-flight
+  check before launching long sweeps; does not require the JVM to be
+  warm.
+
+### Removed
+
+- Dead `get_or_create_netlogo` lazy-init helper in `server.py` — the
+  eager `lifespan()` startup is the only path now, and the duplicated
+  init was an attractive nuisance for future contributors.
+
 ### Added — BehaviorSpace integration & NetLogo 7 transition guide
 
 - 3 new tools for running NetLogo BehaviorSpace experiments:

README.md

@@ -60,6 +60,8 @@ By default, a real NetLogo window opens so you can watch your simulations run li
 | `export_view()` | Export current view as PNG image |
 | `save_model(name, code)` | Save model to file |
 | `export_world()` | Export full world state to CSV |
+| `close_model()` | Unload the current model and `clear-all` the workspace |
+| `server_info()` | Inspect server config (paths, GUI mode, headless launcher) |
 | `list_models()` | List model files in models directory |
 | `search_comses(query)` | Search the CoMSES Net model library |
 | `get_comses_model(uuid)` | Fetch metadata + citation text for one COMSES model |

src/netlogo_mcp/server.py

@@ -14,15 +14,13 @@
 import logging
 from collections.abc import AsyncIterator, Callable, Iterator
 from contextlib import asynccontextmanager, contextmanager
-from threading import RLock
 from typing import Any, TypeVar
 
 from fastmcp import FastMCP
 
 from .config import get_gui_mode, get_jvm_path, get_netlogo_home
 
 logger = logging.getLogger("netlogo_mcp")
-_workspace_lock = RLock()
 T = TypeVar("T")
 
 
@@ -45,49 +43,6 @@ def run_with_stdout_protection(
         return func(*args, **kwargs)
 
 
-def get_or_create_netlogo(lifespan_context: dict[str, Any]) -> Any:
-    """Create the shared NetLogo workspace on first use, then reuse it.
-
-    Lazy init: the JVM starts on the first tool call (30-60s), but
-    this ensures it runs on the correct thread with the right Java
-    class loader context. All subsequent calls are instant.
-    """
-    existing = lifespan_context.get("netlogo")
-    if existing is not None:
-        return existing
-
-    with _workspace_lock:
-        existing = lifespan_context.get("netlogo")
-        if existing is not None:
-            return existing
-
-        with protect_stdout():
-            import pynetlogo
-
-            nl_home = get_netlogo_home()
-            jvm_path = get_jvm_path()
-            gui = get_gui_mode()
-            mode_str = "GUI (live window)" if gui else "headless"
-
-            logger.info("Starting Java Virtual Machine...")
-            logger.info(
-                "Initializing NetLogo in %s mode (NETLOGO_HOME=%s)",
-                mode_str,
-                nl_home,
-            )
-
-            nl = pynetlogo.NetLogoLink(
-                netlogo_home=nl_home,
-                gui=gui,
-                thd=False,  # JPype handles Swing EDT; thd=True hangs on Windows
-                jvm_path=jvm_path,
-            )
-
-        lifespan_context["netlogo"] = nl
-        logger.info("NetLogo workspace ready (%s) — all tools available", mode_str)
-        return nl
-
-
 def shutdown_netlogo(lifespan_context: dict[str, Any]) -> None:
     """Close the shared NetLogo workspace if it was created."""
     nl = lifespan_context.pop("netlogo", None)

src/netlogo_mcp/tools.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import json
+import re
 from datetime import datetime
 from pathlib import Path
 from typing import Any
@@ -20,11 +21,19 @@
     get_comses_cache_dir,
     get_comses_max_download_mb,
     get_exports_dir,
+    get_gui_mode,
     get_models_dir,
     get_netlogo_home,
 )
 from .server import mcp
 
+# NetLogo identifiers may contain letters, digits, and the punctuation
+# characters NetLogo permits in variable / breed names: ``- _ . ? !``.
+# We reject anything else to prevent command injection through `set_parameter`
+# (the value is escaped, but the name is interpolated literally into a
+# `set <name> <value>` command).
+_NETLOGO_IDENTIFIER_RE = re.compile(r"^[A-Za-z][A-Za-z0-9_\-.?!]*$")
+
 # ── Helpers ──────────────────────────────────────────────────────────────────
 
 
@@ -216,6 +225,9 @@ async def run_simulation(
         raise ToolError("ticks must be between 1 and 10000.")
     if not reporters:
         raise ToolError("reporters list cannot be empty.")
+    for i, r in enumerate(reporters):
+        if not isinstance(r, str) or not r.strip():
+            raise ToolError(f"reporters[{i}] must be a non-empty string (got {r!r}).")
     if max_rows < 0:
         raise ToolError("max_rows must be >= 0 (0 = no cap).")
 
@@ -285,8 +297,17 @@ async def set_parameter(name: str, value: Any, ctx: Context) -> str:
 
     Args:
         name: Name of the global variable (e.g. 'initial-number-sheep').
+              Must be a valid NetLogo identifier — letters, digits, and any
+              of ``- _ . ? !`` only. Names with whitespace or shell/NetLogo
+              meta-characters are rejected to prevent command injection.
         value: The value to set. Numbers, strings, booleans accepted.
     """
+    if not isinstance(name, str) or not _NETLOGO_IDENTIFIER_RE.match(name):
+        raise ToolError(
+            f"Invalid parameter name: {name!r}. "
+            "Must start with a letter and contain only letters, digits, "
+            "or any of '-', '_', '.', '?', '!' (NetLogo's identifier rules)."
+        )
     nl = _require_model(ctx)
     # Format the value for NetLogo
     if isinstance(value, bool):
@@ -354,6 +375,8 @@ async def get_patch_data(
         JSON 2D array (rows = y descending, cols = x ascending) by default,
         or a compact summary dict if `summary_only=True`.
     """
+    if not isinstance(attribute, str) or not attribute.strip():
+        raise ToolError("attribute must be a non-empty string.")
     nl = _require_model(ctx)
     try:
         data = nl.patch_report(attribute)
@@ -573,6 +596,76 @@ async def export_world(ctx: Context) -> str:
     return f"World exported to {export_path}"
 
 
+@mcp.tool()
+async def close_model(ctx: Context) -> str:
+    """Unload the currently loaded model and reset the workspace.
+
+    Useful when you want to discard pending state (mid-run agents, set
+    parameters, pending plots) and start fresh, or before opening a new
+    model file from disk to make sure cached compilation state isn't carried
+    forward.
+
+    Note: this does NOT shut down the JVM or NetLogo workspace — only the
+    model. The next `open_model` / `create_model` call will reuse the same
+    JVM (no 30-60s warmup).
+    """
+    nl = _nl(ctx)
+    if _current_model_path(ctx) is None:
+        raise ToolError("No model is currently loaded.")
+    try:
+        # NetLogo doesn't expose an explicit "close model" primitive; the
+        # closest stable equivalent is `clear-all`, which wipes turtles,
+        # patches, ticks, plots, and globals. We then forget the path so
+        # subsequent BehaviorSpace / world-state tools refuse to run until
+        # a new model is loaded.
+        nl.command("clear-all")
+    except Exception as e:
+        raise _wrap_netlogo_error(e) from e
+    _set_current_model_path(ctx, None)
+    return "Model closed. World cleared. Open another model to continue."
+
+
+@mcp.tool()
+async def server_info(ctx: Context) -> str:
+    """Return a snapshot of the running NetLogo MCP server's configuration.
+
+    Useful as a no-cost health check for the AI / user — returns the server
+    version, configured paths, GUI mode, currently-loaded model, and whether
+    a NetLogo headless launcher is reachable for BehaviorSpace runs.
+
+    No JVM round-trip; this is a pure config / filesystem inspection so it
+    works even before the workspace is fully initialized.
+    """
+    from . import __version__
+
+    info: dict[str, Any] = {
+        "server_version": __version__,
+        "gui_mode": get_gui_mode(),
+        "models_dir": str(get_models_dir()),
+        "exports_dir": str(get_exports_dir()),
+        "comses_cache_dir": str(get_comses_cache_dir()),
+        "comses_max_download_mb": get_comses_max_download_mb(),
+        "current_model_path": _current_model_path(ctx),
+    }
+    try:
+        info["netlogo_home"] = get_netlogo_home()
+    except OSError as exc:
+        info["netlogo_home"] = None
+        info["netlogo_home_error"] = str(exc)
+
+    if info.get("netlogo_home"):
+        try:
+            launcher = _bspace.locate_headless_launcher(info["netlogo_home"])
+            info["headless_launcher"] = str(launcher)
+        except _bspace.BSpaceError as exc:
+            info["headless_launcher"] = None
+            info["headless_launcher_error"] = str(exc)
+    else:
+        info["headless_launcher"] = None
+
+    return json.dumps(info, indent=2)
+
+
 # ── CoMSES Net integration ──────────────────────────────────────────────────
 #
 # Five tools + one prompt that let an AI client browse, inspect, download, and