Merge pull request #8 from reward-protocol/derekx/execution-manager

Flatten Rollout Manager and Session Manager into an Execution Manager

Author: xzrderek

examples/rollout_control_plane_demo.py

@@ -22,9 +22,9 @@
 # Add the mcp directory to path
 sys.path.insert(0, str(Path(__file__).parent.parent))
 
-from reward_protocol.mcp.execution.rollout import RolloutManager
-from reward_protocol.mcp.session.manager import GeneralMCPVectorEnv, SessionManager
-from reward_protocol.mcp.types import DatasetRow, MCPSession, MCPToolCall, Trajectory
+from reward_kit.mcp.execution.manager import ExecutionManager
+from reward_kit.mcp.session.manager import GeneralMCPVectorEnv
+from reward_kit.mcp.types import DatasetRow, MCPSession, MCPToolCall, Trajectory
 
 
 class DemoPolicy:
@@ -93,8 +93,7 @@ async def demonstrate_control_plane_rollout():
     ]
 
     # Create rollout manager
-    session_manager = SessionManager()
-    rollout_manager = RolloutManager(session_manager)
+    execution_manager = ExecutionManager()
 
     # Mock the vector environment to simulate control plane separation
     with patch.object(GeneralMCPVectorEnv, "__init__", return_value=None), patch.object(
@@ -274,7 +273,7 @@ def mock_step_side_effect(tool_calls):
         print()
 
         # Execute rollout
-        trajectories = await rollout_manager.execute_rollout(mock_env, policy, steps=10)
+        trajectories = await execution_manager.execute_rollout(mock_env, policy, steps=10)
 
         print()
         print("📋 ROLLOUT RESULTS")
@@ -378,4 +377,4 @@ def mock_step_side_effect(tool_calls):
 
 
 if __name__ == "__main__":
-    asyncio.run(demonstrate_control_plane_rollout())
+    asyncio.run(demonstrate_control_plane_rollout())

reward_kit/mcp/__init__.py

@@ -11,13 +11,13 @@
 
 # New refactored components
 from .client import MCPConnectionManager
-from .execution import FireworksPolicy, LLMBasePolicy, RolloutManager
+from .execution import FireworksPolicy, LLMBasePolicy, ExecutionManager
 from .gym_production_server import GymProductionServer
 
 # North Star MCP-Gym Framework
 from .mcpgym import McpGym
 from .server import MCPEnvironmentServer
-from .session import GeneralMCPVectorEnv, SessionManager
+from .session import GeneralMCPVectorEnv
 from .simulation_server import SimulationServerBase
 from .types import DatasetRow, MCPSession, MCPToolCall, Trajectory
 
@@ -31,13 +31,12 @@
     "MCPConnectionManager",
     "LLMBasePolicy",
     "FireworksPolicy",
-    "RolloutManager",
-    "SessionManager",
+    "ExecutionManager",
     "GeneralMCPVectorEnv",
     "MCPSession",
     "MCPToolCall",
     "DatasetRow",
     "Trajectory",
     # North Star MCP-Gym Framework
     "McpGym",
-]
+]

reward_kit/mcp/execution/__init__.py

@@ -5,10 +5,10 @@
 """
 
 from .policy import FireworksPolicy, LLMBasePolicy
-from .rollout import RolloutManager
+from .manager import ExecutionManager
 
 __all__ = [
     "LLMBasePolicy",
     "FireworksPolicy",
-    "RolloutManager",
-]
+    "ExecutionManager",
+]