Skip to content

Commit 738ee67

Browse files
committed
feat: dual mode should also support mock
Signed-off-by: vsoch <vsoch@users.noreply.github.com>
1 parent 3673ed7 commit 738ee67

3 files changed

Lines changed: 20 additions & 11 deletions

File tree

mcpserver/core/base.py

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,11 @@
22
import json
33
import time
44

5+
from resource_secretary.providers import discover_providers
6+
from resource_secretary.providers.mock import discover_mock_providers
7+
58
import mcpserver.utils as utils
9+
from mcpserver.logger import logger
610

711

812
class WorkerBase:
@@ -12,6 +16,17 @@ class WorkerBase:
1216
its dual mode (acting as worker AND hub.)
1317
"""
1418

19+
def init_providers(self, mock=False):
20+
"""
21+
Probe the local system on startup. E.g., "we found spack, flux, etc."
22+
These can be faux (mock) or real discovered providers
23+
"""
24+
logger.info("📡 Probing local system for resource providers...")
25+
if mock:
26+
self.catalog = discover_mock_providers(self.worker_id, choice=mock)
27+
else:
28+
self.catalog = discover_providers()
29+
1530
def register_agent_tools(self):
1631
"""
1732
Registers the core negotiation tools with the FastMCP instance.

mcpserver/core/hub.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,14 +33,17 @@ def __init__(
3333
dual=False,
3434
hub_id=None,
3535
path="/mcp",
36+
mock=False,
3637
):
38+
# Probably can simplify some of this between worker and hub
3739
self.mcp = mcp
3840
self.host = host
3941
self.port = port
4042
self.path = path
4143
self.secret = secret or secrets.token_urlsafe(32)
4244
self.workers: Dict[str, Dict[str, Any]] = {}
4345
self.hub_id = hub_id or socket.gethostname()
46+
self.mock = mock
4447

4548
# Make requests to hub in batches, in serial, or in parallel
4649
self.set_running_mode(batch, serial, dual)
@@ -103,6 +106,7 @@ def from_args(cls, mcp, args) -> Optional["HubManager"]:
103106
batch=args.batch,
104107
serial=args.serial,
105108
dual=args.dual,
109+
mock=args.mock,
106110
# server path
107111
path=args.path,
108112
)
@@ -397,6 +401,7 @@ def __init__(self, *args, **kwargs):
397401
# Calls super on the HubManager. WorkerBase has no init
398402
super().__init__(*args, **kwargs)
399403
self.setup_dual()
404+
self.init_providers(kwargs.get("mock", False))
400405

401406
def setup_dual(self):
402407
"""

mcpserver/core/worker.py

Lines changed: 0 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -47,17 +47,6 @@ def __init__(
4747
# Register MCP Tools automatically
4848
self.register_agent_tools()
4949

50-
def init_providers(self, mock=False):
51-
"""
52-
Probe the local system on startup. E.g., "we found spack, flux, etc."
53-
These can be faux (mock) or real discovered providers
54-
"""
55-
logger.info("📡 Probing local system for resource providers...")
56-
if mock:
57-
self.catalog = discover_mock_providers(self.worker_id, choice=mock)
58-
else:
59-
self.catalog = discover_providers()
60-
6150
def show(self):
6251
"""
6352
Show providers installed and verbosity.

0 commit comments

Comments
 (0)