[fix] Accept calls on tracking server without trailing slashes (#3119)

aimhubio · Mar 14, 2024 · 09c0f93 · 09c0f93
1 parent 9899aea
commit 09c0f93
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 2 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,5 +1,8 @@
 # Changelog
 
+## 3.19.1
+- Accept calls on tracking server without trailing slashes (mihran113)
+
 ## 3.19.0 Mar 13, 2024
 
 ### Enhancements:

diff --git a/aim/ext/transport/router.py b/aim/ext/transport/router.py
@@ -10,10 +10,15 @@ class ClientRouter:
     def __init__(self):
         self.router = APIRouter()
         self.router.add_api_route('/get-version/', self.get_version, methods=['GET'])
+        self.router.add_api_route('/get-version', self.get_version, methods=['GET'], include_in_schema=False)
         self.router.add_api_route('/heartbeat/{client_uri}/', self.heartbeat, methods=['GET'])
+        self.router.add_api_route('/heartbeat/{client_uri}', self.heartbeat, methods=['GET'], include_in_schema=False)
         self.router.add_api_route('/connect/{client_uri}/', self.connect, methods=['GET'])
+        self.router.add_api_route('/connect/{client_uri}', self.connect, methods=['GET'], include_in_schema=False)
         self.router.add_api_route('/reconnect/{client_uri}/', self.reconnect, methods=['GET'])
+        self.router.add_api_route('/reconnect/{client_uri}', self.reconnect, methods=['GET'], include_in_schema=False)
         self.router.add_api_route('/disconnect/{client_uri}/', self.disconnect, methods=['GET'])
+        self.router.add_api_route('/disconnect/{client_uri}', self.disconnect, methods=['GET'], include_in_schema=False)
 
     @classmethod
     def add_client(cls, client_uri):

diff --git a/aim/ext/transport/server.py b/aim/ext/transport/server.py
@@ -64,6 +64,7 @@ def create_app():
     base_path = os.environ.get(AIM_SERVER_BASE_PATH, '')
     app.mount(f'{base_path}/', api_app)
 
+    @api_app.get('/status', include_in_schema=False)
     @api_app.get('/status/')
     async def status():
         return {'status': 'OK'}

diff --git a/aim/ext/transport/tracking.py b/aim/ext/transport/tracking.py
@@ -67,10 +67,20 @@ def __init__(self, resource_registry: ResourceTypeRegistry):
         self.router = APIRouter()
         self.router.add_api_route('/{client_uri}/get-resource/',
                                   self.get_resource, methods=['POST'])
+        self.router.add_api_route('/{client_uri}/get-resource',
+                                  self.get_resource, methods=['POST'], include_in_schema=False)
         self.router.add_api_route('/{client_uri}/release-resource/{resource_handler}/',
                                   self.release_resource, methods=['GET'])
-        self.router.add_api_route('/{client_uri}/read-instruction/', self.run_instruction, methods=['POST'])
-        self.router.add_api_websocket_route('/{client_uri}/write-instruction/', self.run_write_instructions)
+        self.router.add_api_route('/{client_uri}/release-resource/{resource_handler}',
+                                  self.release_resource, methods=['GET'], include_in_schema=False)
+        self.router.add_api_route('/{client_uri}/read-instruction/',
+                                  self.run_instruction, methods=['POST'])
+        self.router.add_api_route('/{client_uri}/read-instruction',
+                                  self.run_instruction, methods=['POST'], include_in_schema=False)
+        self.router.add_api_websocket_route('/{client_uri}/write-instruction/',
+                                            self.run_write_instructions)
+        self.router.add_api_websocket_route('/{client_uri}/write-instruction',
+                                            self.run_write_instructions)
 
     @classmethod
     def cleanup_client_resources(cls, dead_client_uri):