design: common base for worker and hub

vsoch · vsoch · commit 55dc2fa5c78b · 2026-04-11T20:12:31.000-07:00
We should have a common base with worker functions to make it
cleaner for a hub to serve as a worker.

Signed-off-by: vsoch &lt;vsoch@users.noreply.github.com&gt;
diff --git a/mcpserver/cli/start.py b/mcpserver/cli/start.py
@@ -15,7 +15,7 @@
 from mcpserver.app import init_mcp
 from mcpserver.cli.manager import get_manager
 from mcpserver.core.config import MCPConfig
-from mcpserver.core.hub import HubManager
+from mcpserver.core.hub import DualHubManager, HubManager
 from mcpserver.core.worker import WorkerManager
 from mcpserver.logger import logger
 
@@ -64,7 +64,9 @@ def main(args, extra, **kwargs):
     app = FastAPI(title="MCP Server", lifespan=mcp_app.lifespan)
 
     # Setup Hub (parent role)
-    if args.hub:
+    if args.dual:
+        mcp.hub_manager = DualHubManager.from_args(mcp, args)
+    elif args.hub:
         mcp.hub_manager = HubManager.from_args(mcp, args)
 
     # Setup Worker (child role) - triggered by --join. We require join secret.
@@ -86,8 +88,8 @@ async def lifespan(app: FastAPI):
 
     app = FastAPI(title="MCP Server", lifespan=lifespan)
 
-    # Bind the /register endpoint if we are a Hub
-    if args.hub:
+    # Bind the /register endpoint if we are a Hub (or Hub and Worker)
+    if args.hub or args.dual:
         mcp.hub_manager.bind_to_app(app)
 
     # Mount the MCP server. Note from V: we can use mount with antother FastMCP
diff --git a/mcpserver/core/base.py b/mcpserver/core/base.py
@@ -0,0 +1,99 @@
+import collections
+import json
+import time
+
+import mcpserver.utils as utils
+
+
+class WorkerBase:
+    """
+    A WorkerBase provides worker interaction functions, e.g., negotiate, status,
+    ask secretary. We provide it here so that a hub can use it to generate
+    its dual mode (acting as worker AND hub.)
+    """
+
+    def register_agent_tools(self):
+        """
+        Registers the core negotiation tools with the FastMCP instance.
+        """
+
+        @self.mcp.tool(name="get_status")
+        async def get_status() -> dict:
+            """
+            Returns the Level 1 Static Manifest of this cluster.
+            Use this to verify hardware, software providers, and site info.
+            """
+            return {
+                "worker_id": self.worker_id,
+                "timestamp": time.time(),
+                "manifest": self.manifest,
+            }
+
+        @self.mcp.tool(name="ask_secretary")
+        async def ask_secretary(request: str) -> dict:
+            """
+            Wakes up the local Secretary Agent to perform a Level 2 investigation.
+            Use this to ask about specific software availability, queue depth, or node health.
+            """
+            from resource_secretary.agents.secretary import SecretaryAgent
+
+            # Flatten the catalog into a list of active provider instances
+            active_providers = [inst for category in self.catalog.values() for inst in category]
+
+            # Verbose mode returns a second block with CALLS
+            agent = SecretaryAgent(active_providers, verbose=self.verbose)
+            proposal = await agent.negotiate(request)
+            return {"worker_id": self.worker_id, "proposal": proposal}
+
+        @self.mcp.tool(name="submit")
+        async def receive_job(request: str) -> dict:
+            """
+            Receive a job. Accepts a job request, invokes the local Secretary to
+            generate a spec, submit it, and verify the job ID.
+            """
+            from resource_secretary.agents.secretary import SecretaryAgent
+
+            active_providers = [inst for cat in self.catalog.values() for inst in cat]
+
+            agent = SecretaryAgent(active_providers)
+            raw_result = await agent.submit(request)
+            try:
+                receipt = json.loads(utils.extract_code_block(raw_result))
+            except:
+                receipt = {"status": "FAILED", "reasoning": raw_result}
+
+            return {"worker_id": self.worker_id, "receipt": receipt}
+
+        @self.mcp.tool(name="export_provider_metadata")
+        def export_provider_metadata() -> str:
+            """
+            Iterates through all providers and returns their internal 'truth' state.
+            This tool is 'hidden' from the Secretary Agent but used by the Hub.
+            """
+            truth_map = {}
+            tool_registry = collections.defaultdict(list)
+
+            # Self.catalog is a dict: {"software": [MockSpackProvider, ...]}
+            for category, providers in self.catalog.items():
+                truth_map[category] = {}
+                for p in providers:
+                    # We check if the provider has the export_truth method
+                    if hasattr(p, "export_truth"):
+                        truth_map[category][p.name] = p.export_truth()
+                    else:
+                        # Fallback to standard metadata if not a mock
+                        truth_map[category][p.name] = p.metadata
+
+                    # Capture all Secretary Tools for this provider
+                    # We can use this for simulations to assess what the agent
+                    # should have called (vs. what it did)
+                    manifest = p.discover_tools(tool_types=["secretary"])
+                    for tool_name in manifest.keys():
+                        tool_registry[category].append(f"{p.name}.{tool_name}")
+
+            metadata = {"truth": truth_map, "registry": dict(tool_registry)}
+
+            # If we have an archetype (mocking something) save it
+            if hasattr(p, "archetype"):
+                metadata["metadata"] = {"archetype": p.archetype.name}
+            return json.dumps(metadata, indent=2)
diff --git a/mcpserver/core/hub.py b/mcpserver/core/hub.py
@@ -13,6 +13,8 @@
 import mcpserver.utils as utils
 from mcpserver.logger import logger
 
+from .base import WorkerBase
+
 
 class HubManager:
     """
@@ -30,10 +32,12 @@ def __init__(
         serial=False,
         dual=False,
         hub_id=None,
+        path="/mcp",
     ):
         self.mcp = mcp
         self.host = host
         self.port = port
+        self.path = path
         self.secret = secret or secrets.token_urlsafe(32)
         self.workers: Dict[str, Dict[str, Any]] = {}
         self.hub_id = hub_id or socket.gethostname()
@@ -81,21 +85,15 @@ def set_running_mode(self, batch_size=None, serial=False, dual=False):
         logger.info(f"🚦 Hub initialized with Batch Size: {batch_size}")
 
         # If we are also running as a worker, add ourselves to the fleet
-        if not dual:
-            return
-
-        hub_id = self.hub_id or socket.gethostname()
-        self.workers[hub_id] = {
-            "url": self.registration_url,
-            "client": Client(self.registration_url),
-        }
+        self.dual = dual
 
     @classmethod
     def from_args(cls, mcp, args) -> Optional["HubManager"]:
         """
         Create a HubManager from CLI arguments.
         """
-        if not getattr(args, "hub", False):
+        # Running in hub or dual mode?
+        if not getattr(args, "hub", False) and not getattr(args, "dual", False):
             return None
         return cls(
             mcp,
@@ -105,6 +103,8 @@ def from_args(cls, mcp, args) -> Optional["HubManager"]:
             batch=args.batch,
             serial=args.serial,
             dual=args.dual,
+            # server path
+            path=args.path,
         )
 
     def _print_banner(self):
@@ -386,3 +386,25 @@ def _create_proxy(self, worker_id: str, tool: Tool):
 
         except Exception as e:
             logger.error(f"❌ Failed to generate dynamic proxy for {tool.name}: {e}")
+
+
+class DualHubManager(WorkerBase, HubManager):
+    """
+    Combined hub and worker base. Aka, a hub that also serves as a worker
+    """
+
+    def __init__(self, *args, **kwargs):
+        # Calls super on the HubManager. WorkerBase has no init
+        super().__init__(*args, **kwargs)
+        self.setup_dual()
+
+    def setup_dual(self):
+        """
+        Setup dual mode, which means adding ourselves to the fleet.
+        """
+        hub_id = self.hub_id or socket.gethostname()
+        default_url = f"http://{self.host}:{self.port}{self.path}"
+        self.workers[hub_id] = {
+            "url": self.registration_url,
+            "client": Client(default_url),
+        }
diff --git a/mcpserver/core/worker.py b/mcpserver/core/worker.py
@@ -1,20 +1,18 @@
 import asyncio
-import collections
-import json
 import socket
-import time
 from typing import Any, Dict, Optional
 
 import httpx
 from resource_secretary.providers import discover_providers
 from resource_secretary.providers.mock import discover_mock_providers
 from rich import print
 
-import mcpserver.utils as utils
 from mcpserver.logger import logger
 
+from .base import WorkerBase
 
-class WorkerManager:
+
+class WorkerManager(WorkerBase):
     """
     A generic worker mcpserver that discovers its own capabilities
     and context using the resource-secretary library.
@@ -93,92 +91,6 @@ def parse_labels(self, label_list: Optional[list]) -> dict:
                 labels[k.strip()] = v.strip()
         return labels
 
-    def register_agent_tools(self):
-        """
-        Registers the core negotiation tools with the FastMCP instance.
-        """
-
-        @self.mcp.tool(name="get_status")
-        async def get_status() -> dict:
-            """
-            Returns the Level 1 Static Manifest of this cluster.
-            Use this to verify hardware, software providers, and site info.
-            """
-            return {
-                "worker_id": self.worker_id,
-                "timestamp": time.time(),
-                "manifest": self.manifest,
-            }
-
-        @self.mcp.tool(name="ask_secretary")
-        async def ask_secretary(request: str) -> dict:
-            """
-            Wakes up the local Secretary Agent to perform a Level 2 investigation.
-            Use this to ask about specific software availability, queue depth, or node health.
-            """
-            from resource_secretary.agents.secretary import SecretaryAgent
-
-            # Flatten the catalog into a list of active provider instances
-            active_providers = [inst for category in self.catalog.values() for inst in category]
-
-            # Verbose mode returns a second block with CALLS
-            agent = SecretaryAgent(active_providers, verbose=self.verbose)
-            proposal = await agent.negotiate(request)
-            return {"worker_id": self.worker_id, "proposal": proposal}
-
-        @self.mcp.tool(name="submit")
-        async def receive_job(request: str) -> dict:
-            """
-            Receive a job. Accepts a job request, invokes the local Secretary to
-            generate a spec, submit it, and verify the job ID.
-            """
-            from resource_secretary.agents.secretary import SecretaryAgent
-
-            active_providers = [inst for cat in self.catalog.values() for inst in cat]
-
-            agent = SecretaryAgent(active_providers)
-            raw_result = await agent.submit(request)
-            try:
-                receipt = json.loads(utils.extract_code_block(raw_result))
-            except:
-                receipt = {"status": "FAILED", "reasoning": raw_result}
-
-            return {"worker_id": self.worker_id, "receipt": receipt}
-
-        @self.mcp.tool(name="export_provider_metadata")
-        def export_provider_metadata() -> str:
-            """
-            Iterates through all providers and returns their internal 'truth' state.
-            This tool is 'hidden' from the Secretary Agent but used by the Hub.
-            """
-            truth_map = {}
-            tool_registry = collections.defaultdict(list)
-
-            # Self.catalog is a dict: {"software": [MockSpackProvider, ...]}
-            for category, providers in self.catalog.items():
-                truth_map[category] = {}
-                for p in providers:
-                    # We check if the provider has the export_truth method
-                    if hasattr(p, "export_truth"):
-                        truth_map[category][p.name] = p.export_truth()
-                    else:
-                        # Fallback to standard metadata if not a mock
-                        truth_map[category][p.name] = p.metadata
-
-                    # Capture all Secretary Tools for this provider
-                    # We can use this for simulations to assess what the agent
-                    # should have called (vs. what it did)
-                    manifest = p.discover_tools(tool_types=["secretary"])
-                    for tool_name in manifest.keys():
-                        tool_registry[category].append(f"{p.name}.{tool_name}")
-
-            metadata = {"truth": truth_map, "registry": dict(tool_registry)}
-
-            # If we have an archetype (mocking something) save it
-            if hasattr(p, "archetype"):
-                metadata["metadata"] = {"archetype": p.archetype.name}
-            return json.dumps(metadata, indent=2)
-
     async def run_registration(self):
         """
         Registers the worker with the Hub.
