server.py

# server.py
from dataclasses import dataclass, asdict
import secrets
import logging
import asyncio
import json
from typing import Tuple

from quart import Quart, websocket, request
from quart_schema import QuartSchema, validate_request, validate_response
from quart_cors import cors
from uvicorn.middleware.proxy_headers import ProxyHeadersMiddleware

from broker import SessionBroker, SessionDoesNotExist, ClientRequest, ClientResponse, ClientError

# Configuración
VERSION = "1.0.0"  # Versión de la API
TIMEOUT: int = 40
LOG_LEVEL: int = logging.DEBUG
TRUSTED_HOSTS: list[str] = ["127.0.0.1", "10.16.38.136", "10.16.3.13", "10.16.13.73"]

# 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()

# Modelos de datos
@dataclass
class Status:
    status: str
    version: str

@dataclass
class Session:
    session_id: str

@dataclass
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

# Rutas API
@app.get("/status")
@validate_response(Status)
async def status() -> Status:
    return Status(status="OK", version=VERSION)

@app.websocket('/session')
async def session_handler():
    session_id = secrets.token_hex()
    app.logger.info(f"{websocket.remote_addr} - NEW SESSION - {session_id}")
    await websocket.send_as(Session(session_id=session_id), Session)

    task = None
    try:
        task = asyncio.ensure_future(_receive(session_id))
        async for request in broker.subscribe(session_id):
            app.logger.info(f"{websocket.remote_addr} - REQUEST - {session_id} - {json.dumps(asdict(request))}")
            await websocket.send_as(request, ClientRequest)
    finally:
        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"{websocket.remote_addr} - RESPONSE - {session_id} - {json.dumps(asdict(response))}")
        await broker.receive_response(session_id, response)

@app.post('/command')
@validate_request(Command)
@validate_response(CommandResponse, 200)
@validate_response(ErrorResponse, 500)
async def command(data: Command) -> Tuple[CommandResponse | ErrorResponse, int]:
    try:
        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"{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

# Agregar un endpoint de health check y root
@app.route("/")
async def root():
    return {"message": "Kaio API Server", "version": VERSION}

@app.route("/health")
async def health_check():
    return {"status": "healthy"}

def run():
    app.run(host='0.0.0.0', port=7860)

if __name__ == "__main__":
    run()