Refactor server.py and update requirements.txt for enhanced functionality and logging

- Simplified server.py by removing unused imports and restructuring configuration settings.
- Introduced new data models for read and write operations, improving API clarity.
- Enhanced logging to include remote addresses for better traceability of requests and responses.
- Updated requirements.txt to include uvicorn for ASGI support, ensuring compatibility with the new server structure.

broker.py

@@ -0,0 +1,73 @@
+import asyncio
+import uuid
+from typing import AsyncGenerator, Dict, Tuple, Any
+from dataclasses import dataclass
+
+class SessionDoesNotExist(Exception):
+    pass
+
+class SessionAlreadyExists(Exception):
+    pass
+
+class ClientError(Exception):
+    def __init__(self, message):
+        super().__init__(message)
+        self.message = message
+
+@dataclass
+class ClientRequest:
+    request_id: str
+    data: Any 
+
+@dataclass
+class ClientResponse:
+    request_id: str
+    error: bool
+    data: Any
+
+class SessionBroker:
+    def __init__(self):
+        self.sessions: Dict[str, asyncio.Queue] = {}
+        self.pending_responses: Dict[Tuple[str, str], asyncio.Future] = {}
+
+    async def send_request(self, session_id: str, data: Any, timeout: int = 60) -> Any:
+        if session_id not in self.sessions:
+            raise SessionDoesNotExist()
+
+        request_id = str(uuid.uuid4())
+        future = asyncio.get_event_loop().create_future()
+        self.pending_responses[(session_id, request_id)] = future
+
+        await self.sessions[session_id].put(ClientRequest(request_id=request_id, data=data))
+
+        try:
+            return await asyncio.wait_for(future, timeout)
+        except asyncio.TimeoutError:
+            raise
+        finally:
+            if (session_id, request_id) in self.pending_responses:
+                del self.pending_responses[(session_id, request_id)]
+            
+    async def receive_response(self, session_id: str, response: ClientResponse) -> None:
+        if (session_id, response.request_id) in self.pending_responses:
+            future = self.pending_responses.pop((session_id, response.request_id))
+            if not future.done():
+                if response.error:
+                    future.set_exception(ClientError(message=response.data))
+                else:
+                    future.set_result(response.data)
+
+    async def subscribe(self, session_id: str) -> AsyncGenerator[ClientRequest, None]:
+        if session_id in self.sessions:
+            raise SessionAlreadyExists()
+
+        queue = asyncio.Queue()
+        self.sessions[session_id] = queue
+
+        try:
+            while True:
+                yield await queue.get()
+        finally:
+            if session_id in self.sessions:
+                del self.sessions[session_id]
+                self.pending_responses = {k: v for k, v in self.pending_responses.items() if k[0] != session_id}

requirements.txt

@@ -7,4 +7,5 @@ websockets==11.0.3
 python-json-logger==2.0.7
 prometheus-client==0.17.1
 pydantic==1.10.13
-python-dotenv==1.0.0 
+python-dotenv==1.0.0
+uvicorn==0.24.0 

server.py

@@ -1,45 +1,32 @@
 # server.py
-import asyncio
-import uuid
-from typing import AsyncGenerator, Dict, Tuple, Any, Optional
-from dataclasses import dataclass
-from quart import Quart, websocket, request, Response
-from quart_schema import QuartSchema, validate_request, validate_response
-from quart_cors import cors
+from dataclasses import dataclass, asdict
 import importlib.metadata
 import secrets
 import logging
-import os
+import asyncio
+import json
+from typing import Tuple
+
+from quart import Quart, websocket, request
+from quart_schema import QuartSchema, validate_request, validate_response
+from uvicorn.middleware.proxy_headers import ProxyHeadersMiddleware
 
 from broker import SessionBroker, SessionDoesNotExist, ClientRequest, ClientResponse, ClientError
 
-# Configuración para Hugging Face Spaces
-PORT = int(os.getenv('PORT', 7860))
-TIMEOUT: int = int(os.getenv('TIMEOUT', 60))
-LOG_LEVEL: int = getattr(logging, os.getenv('LOG_LEVEL', 'INFO'))
-MAX_MESSAGE_SIZE: int = int(os.getenv('MAX_MESSAGE_SIZE', 16 * 1024 * 1024))
-RATE_LIMIT: int = int(os.getenv('RATE_LIMIT', 100))
-SESSION_TIMEOUT: int = int(os.getenv('SESSION_TIMEOUT', 3600))
-
-# Configurar logging
-logging.basicConfig(
-    level=LOG_LEVEL,
-    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
-)
-
-# Crear aplicación con CORS habilitado
+# Configuración
+TIMEOUT: int = 40
+LOG_LEVEL: int = logging.DEBUG
+TRUSTED_HOSTS: list[str] = ["127.0.0.1"]
+
+# Create app
 app = Quart(__name__)
-app = cors(app, 
-    allow_origin=["https://*.hf.space", "https://*.huggingface.co"],
-    allow_methods=["GET", "POST", "OPTIONS"],
-    allow_headers=["Content-Type"],
-    max_age=3600
-)
 QuartSchema(app)
+app.asgi_app = ProxyHeadersMiddleware(app.asgi_app, trusted_hosts=TRUSTED_HOSTS)
+app.logger.setLevel(LOG_LEVEL)
 
 broker = SessionBroker()
 
-# Definición de modelos de datos
+# Modelos de datos
 @dataclass
 class Status:
     status: str
@@ -54,12 +41,31 @@ class Command:
     session_id: str
     command: str
 
+@dataclass
+class Read:
+    session_id: str
+    path: str
+
+@dataclass
+class Write:
+    session_id: str
+    path: str
+    content: str
+
 @dataclass
 class CommandResponse:
     return_code: int
     stdout: str
     stderr: str
 
+@dataclass
+class ReadResponse:
+    content: str
+
+@dataclass
+class WriteResponse:
+    size: int
+
 @dataclass
 class ErrorResponse:
     error: str
@@ -73,21 +79,24 @@ async def status() -> Status:
 @app.websocket('/session')
 async def session_handler():
     session_id = secrets.token_hex()
-    app.logger.info(f"New session: {session_id}")
+    app.logger.info(f"{websocket.remote_addr} - NEW SESSION - {session_id}")
     await websocket.send_as(Session(session_id=session_id), Session)
 
-    task = asyncio.ensure_future(_receive(session_id))
+    task = None
     try:
+        task = asyncio.ensure_future(_receive(session_id))
         async for request in broker.subscribe(session_id):
-            app.logger.info(f"Sending request {request.request_id} to client.")
+            app.logger.info(f"{websocket.remote_addr} - REQUEST - {session_id} - {json.dumps(asdict(request))}")
             await websocket.send_as(request, ClientRequest)
     finally:
-        task.cancel()
+        if task is not None:
+            task.cancel()
+            await task
 
 async def _receive(session_id: str) -> None:
     while True:
         response = await websocket.receive_as(ClientResponse)
-        app.logger.info(f"Received response for session {session_id}: {response}")
+        app.logger.info(f"{websocket.remote_addr} - RESPONSE - {session_id} - {json.dumps(asdict(response))}")
         await broker.receive_response(session_id, response)
 
 @app.post('/command')
@@ -96,33 +105,56 @@ async def _receive(session_id: str) -> None:
 @validate_response(ErrorResponse, 500)
 async def command(data: Command) -> Tuple[CommandResponse | ErrorResponse, int]:
     try:
-        response_data = await broker.send_request(
-            session_id=data.session_id, 
-            data={'action': 'command', 'command': data.command}, 
-            timeout=TIMEOUT
-        )
-        response = CommandResponse(**response_data)
+        response = CommandResponse(**await broker.send_request(
+            data.session_id,
+            {'action': 'command', 'command': data.command},
+            timeout=TIMEOUT))
+        return response, 200
+    except SessionDoesNotExist:
+        app.logger.warning(f"{request.remote_addr} - INVALID SESSION ID - {repr(data.session_id)}")
+        return ErrorResponse('Session does not exist.'), 500
+    except ClientError as e:
+        return ErrorResponse(e.message), 500
+    except asyncio.TimeoutError:
+        return ErrorResponse('Timeout when waiting for client.'), 500
+
+@app.post('/read')
+@validate_request(Read)
+@validate_response(ReadResponse, 200)
+@validate_response(ErrorResponse, 500)
+async def read(data: Read) -> Tuple[ReadResponse | ErrorResponse, int]:
+    try:
+        response = ReadResponse(**await broker.send_request(
+            data.session_id,
+            {'action': 'read', 'path': data.path},
+            timeout=TIMEOUT))
+        return response, 200
+    except SessionDoesNotExist:
+        app.logger.warning(f"{request.remote_addr} - INVALID SESSION ID - {repr(data.session_id)}")
+        return ErrorResponse('Session does not exist.'), 500
+    except ClientError as e:
+        return ErrorResponse(e.message), 500
+    except asyncio.TimeoutError:
+        return ErrorResponse('Timeout when waiting for client.'), 500
+
+@app.post('/write')
+@validate_request(Write)
+@validate_response(WriteResponse, 200)
+@validate_response(ErrorResponse, 500)
+async def write(data: Write) -> Tuple[WriteResponse | ErrorResponse, int]:
+    try:
+        response = WriteResponse(**await broker.send_request(
+            data.session_id,
+            {'action': 'write', 'path': data.path, 'content': data.content},
+            timeout=TIMEOUT))
         return response, 200
     except SessionDoesNotExist:
-        app.logger.warning(f"Invalid session ID: {data.session_id}")
-        return ErrorResponse(error='Session does not exist.'), 500
+        app.logger.warning(f"{request.remote_addr} - INVALID SESSION ID - {repr(data.session_id)}")
+        return ErrorResponse('Session does not exist.'), 500
     except ClientError as e:
-        return ErrorResponse(error=e.message), 500
+        return ErrorResponse(e.message), 500
     except asyncio.TimeoutError:
-        return ErrorResponse(error='Timeout waiting for client response.'), 500
+        return ErrorResponse('Timeout when waiting for client.'), 500
 
-# Ejecutar aplicación
 def run():
-    app.run(
-        host='0.0.0.0',
-        port=PORT,
-        debug=False
-    )
-
-# Agregar un endpoint de health check
-@app.route("/health")
-async def health_check():
-    return {"status": "healthy"}
-
-if __name__ == "__main__":
-    run()
+    app.run(host='0.0.0.0', port=7860)