Refactor MCP server to model Protocol and a server registration class #2
Description
Issue: Each new simulation model would require custom FastMCP registration logic. This means manually defining each tool, resource, and prompt for every new model, leading to duplicated code and potential inconsistencies. The FastMCP server would need to be directly aware of the simulation's internal methods, coupling them tightly.
@AliHarp It would be good to standardise (or at least minimise) MCP server offerings (i.e. an expected list of tools, resources and prompts) across simulation models. This would make it easier for other modellers to create new servers. Longer term I'm thinking new modellers will be Agents.
I propose that a we make use of a Python Protocol called Simulation for a concrete simulation model to follow i.e. it must implement the methods. We could separate out the FastMCP server code i.e. registration into a simple SimulationMCPServer class.
Initial Design
The end product
The design is more complicated, but will be clean to use! For example, creating the call centre server would look like the below. If you skip to the end of this issue you can see how its re-usable.
from call_centre_simulation import CallCentreSimulation
from simulation_mcp_server import SimulationMCPServer
def main():
# Create the simulation model
call_centre_sim = CallCentreSimulation()
# Create the MCP server
server = SimulationMCPServer(call_centre_sim)
# Run the server
server.run(transport="http", host="127.0.0.1", port=8001, path="/mcp")
if __name__ == "__main__":
main()Model Protocol
The model will follow a Protocol
from typing import Protocol, runtime_checkable, Dict, Any
@runtime_checkable
class SimulationModel(Protocol):
"""Protocol defining the interface for MCP-compatible simulation models."""
def run_simulation(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
"""Execute the simulation with given parameters and return results."""
...
def get_parameter_schema(self) -> Dict[str, Any]:
"""Return JSON schema for valid simulation parameters."""
...
def get_model_description(self) -> str:
"""Return human-readable description of the simulation model."""
...
def validate_parameters(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
"""Validate parameters and return validation results."""
...
@property
def model_name(self) -> str:
"""Return the name/identifier for this simulation model."""
...MCP Server
This would effectively just register the (expected) tools, resources, prompts of the sim model. It also provides all meta-data that the LLM will use (e.g. name and description of the tools available). It delegates all of the work to the concrete implementation of the SimulationModel. Basically its the "thing you run" when you want to startup the server and the "thing an agent interrogates and interacts via" when wanting to use a simulation model.
from fastmcp import FastMCP
from langchain_core.prompts import PromptTemplate
from fastmcp.prompts.prompt import PromptMessage, TextContent
class SimulationMCPServer:
"""MCP server that can work with any SimulationModel implementation."""
def __init__(self, simulation_model: SimulationModel, server_name: str = None):
# Runtime validation (could maybe use @abstractmethod with `SimulationModel` to reduce code.
if not isinstance(simulation_model, SimulationModel):
missing_methods = []
for method in ['run_simulation', 'get_parameter_schema', 'get_model_description',
'validate_parameters']:
if not hasattr(simulation_model, method):
missing_methods.append(method)
if not hasattr(simulation_model, 'model_name'):
missing_methods.append('model_name (property)')
raise TypeError(f"Object {type(simulation_model)} missing required methods: {missing_methods}")
self.model = simulation_model
server_name = server_name or f"{simulation_model.model_name.title()} Simulation MCP Server"
self.mcp = FastMCP(server_name)
self._register_tools()
self._register_resources()
self._register_prompts()
def _register_tools(self):
"""Register MCP tools that delegate to the simulation model."""
@self.mcp.tool(
name=f"run_{self.model.model_name}_simulation",
description=f"""
Runs a discrete-event {self.model.model_name} simulation with specified parameters, returning performance metrics.
Inputs: parameters (dict) — JSON object matching the experiment schema.
Returns: dict with simulation metrics, such as mean wait times and resource utilizations.
Tags: ["simulation", "{self.model.model_name}", "experiment"]
""")
def run_simulation(parameters: dict) -> dict:
return self.model.run_simulation(parameters)
@self.mcp.tool(
name="validate_simulation_parameters",
description="""
Validate a proposed set of simulation parameters (JSON object) against the experiment schema.
Inputs: parameters (dict)
Returns: {"is_valid": bool, "errors": [str, ...]} — status and explanation.
Tags: ["validation", "parameter_check", "pre_run_check", "schema"]
""")
def validate_parameters(parameters: dict) -> dict:
return self.model.validate_parameters(parameters)
def _register_resources(self):
"""Register MCP resources that delegate to the simulation model."""
@self.mcp.resource(
uri="resource://schema/experiment_parameters",
description="""
Returns the JSON schema defining all allowed input parameters, parameter types, and value constraints.
Outputs: dict (JSON schema), sent as a JSON object.
Tags: ["schema", "parameters", "template"]
""")
def get_schema() -> dict:
return self.model.get_parameter_schema()
@self.mcp.resource(
uri="resource://model/description",
description=f"""
Provides a natural language description of the {self.model.model_name} simulation model.
Outputs: str (text description).
Tags: ["model", "description", "documentation"]
""")
def get_description() -> str:
return self.model.get_model_description()
def _register_prompts(self):
"""Register MCP prompts for parameter conversion."""
@self.mcp.prompt(
name="parameter_jsonification_prompt",
description="""
INSTRUCTION TO LLM: Convert a user's freeform simulation request into a JSON object matching a given schema.
Inputs:
- schema (str): JSON Schema as a string
- user_input (str): User's natural language request
Returns: PromptMessage (LLM input) guiding the agent to produce valid JSON parameters.
Tags: ["jsonification", "schema_mapping", "prompt", "parameters"]
""")
def parameter_jsonification_prompt(
schema: str,
user_input: str,
validation_errors: str = ""
) -> PromptMessage:
with open("resources/parameter_prompt.txt", encoding="utf-8") as f:
prompt_template_text = f.read()
prompt = PromptTemplate.from_template(prompt_template_text)
# Handle validation error feedback
if validation_errors and validation_errors.strip():
validation_feedback = (
"**Validation Feedback:**\n"
"Your last attempt did not pass validation for these reasons:\n"
f"{validation_errors}\n\n"
"Please address the issues above and try again."
)
else:
validation_feedback = ""
filled_prompt = prompt.format(
schema=schema,
user_input=user_input,
validation_feedback=validation_feedback
)
return PromptMessage(
role="user",
content=TextContent(type="text", text=filled_prompt)
)
def run(self, **kwargs):
"""Start the MCP server."""
self.mcp.run(**kwargs)Concrete implementation of a call centre model
Basically this is the existing code refactored from functions to a class. But a class that follows our protocol.
import json
from model import run_simulation_from_dict
class CallCentreSimulation:
"""Call centre simulation implementation following the SimulationModel protocol."""
def __init__(self, schema_path: str = "resources/schema.json"):
self.schema_path = schema_path
self._schema = None
@property
def model_name(self) -> str:
return "healthcare_call_centre"
def run_simulation(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
"""Run the discrete-event healthcare call centre simulation."""
return run_simulation_from_dict(parameters)
def get_parameter_schema(self) -> Dict[str, Any]:
"""Load and return the parameter schema."""
if self._schema is None:
with open(self.schema_path) as f:
self._schema = json.load(f)
return self._schema
def get_model_description(self) -> str:
"""Return human-readable description of the call centre model."""
return (
"This is a discrete-event simulation of a healthcare call centre. "
"Patients call in, interact with operators, and a subset may require a nurse callback. "
"Simulation components: SimPy queues and resources. Tracks wait times, utilization, and callback rates. "
"Configurable parameters: number of operators and nurses, call durations and rates, etc. "
"Sample: 'Run with 14 operators and 5% higher demand.'"
)
def validate_parameters(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
"""Validate simulation parameters against the experiment schema."""
schema = self.get_parameter_schema()
errors = []
for key, value in parameters.items():
# Check for unknown parameters
if key not in schema:
errors.append(f"Unknown parameter: {key}")
continue
spec = schema[key]
expected_type = int if spec["type"] == "int" else float
# Type validation
if not isinstance(value, expected_type):
errors.append(f"{key} must be {spec['type']}")
continue
# Range validation
if "minimum" in spec and value < spec["minimum"]:
errors.append(f"{key} below minimum {spec['minimum']}")
if "maximum" in spec and value > spec["maximum"]:
errors.append(f"{key} above maximum {spec['maximum']}")
# Cross-parameter validation
if all(x in parameters for x in ("call_low", "call_mode", "call_high")):
if not (parameters["call_low"] <= parameters["call_mode"] <= parameters["call_high"]):
errors.append("call_low ≤ call_mode ≤ call_high violated")
if all(x in parameters for x in ("nurse_consult_low", "nurse_consult_high")):
if not (parameters["nurse_consult_low"] <= parameters["nurse_consult_high"]):
errors.append("nurse_consult_low ≤ nurse_consult_high violated")
return {"is_valid": len(errors) == 0, "errors": errors}Reusing with another model
I asked Claude to create another model for us to give an example of how the framework is reusable. What you should see is that we simply define a simulation that follows the protocol (has the correct methods) and we can register it as a server and run it.
from protocols import SimulationModel
from typing import Dict, Any
class InventorySimulation:
"""Example of another simulation model using the same protocol."""
@property
def model_name(self) -> str:
return "inventory_management"
def run_simulation(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
# Inventory-specific simulation logic
return {"average_stock_level": 150, "stockout_rate": 0.05}
def get_parameter_schema(self) -> Dict[str, Any]:
return {
"initial_stock": {"type": "int", "minimum": 0, "maximum": 10000},
"reorder_point": {"type": "int", "minimum": 0, "maximum": 1000}
}
def get_model_description(self) -> str:
return "Discrete-event simulation of inventory management system..."
def validate_parameters(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
# Inventory-specific validation logic
return {"is_valid": True, "errors": []}
# Same server class works!
if __name__ == "__main__":
inventory_sim = InventorySimulation()
inventory_server = SimulationMCPServer(inventory_sim)
inventory_server.run(transport="http", host="127.0.0.1", port=8002, path="/mcp")