Spaces:
Sleeping
Sleeping
File size: 2,751 Bytes
ed3a617 bac0ba4 ed3a617 bac0ba4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | """
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
|