server/step_middleware.py

"""
Middleware to normalize /step payload to wrapped format that OpenEnv expects.

OpenEnv's auto-generated /step endpoint requires:
  {"action": {"fixed_config": "..."}}

Validator likely sends:
  {"fixed_config": "..."}

This middleware wraps the direct format to match what OpenEnv expects.
"""
from fastapi import Request
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.datastructures import MutableHeaders
import json


class StepPayloadWrapperMiddleware(BaseHTTPMiddleware):
    """Wrap direct format to match OpenEnv's expected action format."""
    
    async def dispatch(self, request, call_next):
        # Log session info
        session_id = request.headers.get("x-session-id", request.cookies.get("session_id", "NO_SESSION_ID"))
        print(f"[SESSION_MIDDLEWARE] /step request, session_id={session_id}")
        
        # Only process /step POST requests
        if request.url.path != "/step" or request.method != "POST":
            return await call_next(request)
        
        try:
            body = await request.body()
            payload = json.loads(body) if body else {}
            
            print(f"[STEP_RAW] Input payload: {json.dumps(payload)}")
            
            # Check if payload needs wrapping
            if isinstance(payload, dict):
                # If already wrapped in "action", pass through
                if "action" in payload and isinstance(payload.get("action"), dict):
                    print(f"[STEP_WRAP] Already wrapped, passing through")
                    return await call_next(request)
                
                # If direct format (has "fixed_config" but no "action"), wrap it
                if "fixed_config" in payload and "action" not in payload:
                    wrapped = {"action": payload}
                    wrapped_body = json.dumps(wrapped).encode()
                    print(f"[STEP_WRAP] Wrapped direct format to {json.dumps(wrapped)}")
                    
                    # Modify the request body
                    request._body = wrapped_body
                    
                    # Update content-length header
                    headers = MutableHeaders(scope=request.scope)
                    headers["content-length"] = str(len(wrapped_body))
                    
            # Clear the body cache to force re-reading
            request._stream_consumed = False
            
        except Exception as e:
            print(f"[STEP_WRAP_ERROR] {type(e).__name__}: {e}")
            # If anything fails, just pass through original request
        
        response = await call_next(request)
        print(f"[SESSION_MIDDLEWARE] /step response completed for session_id={session_id}")
        return response