Make HTTP sessions stateless for multi-node deployment

server.py

@@ -14,7 +14,12 @@ instances.
 """
 from __future__ import annotations
 
-from uuid import uuid4
+import base64
+import hashlib
+import hmac
+import json
+import os
+import zlib
 
 from fastapi import FastAPI, HTTPException
 from fastapi import Request, Response
@@ -33,14 +38,52 @@ app = FastAPI(
 )
 
 SESSION_COOKIE = "esc_session_id"
-_envs: dict[str, ESCEnv] = {}
+SESSION_SECRET = os.getenv("ESC_SESSION_SECRET", "esc-openenv-dev-secret").encode("utf-8")
+
+
+def _urlsafe_b64encode(data: bytes) -> str:
+    return base64.urlsafe_b64encode(data).decode("ascii")
+
+
+def _urlsafe_b64decode(data: str) -> bytes:
+    padding = "=" * (-len(data) % 4)
+    return base64.urlsafe_b64decode(data + padding)
+
+
+def _sign(payload: str) -> str:
+    return hmac.new(SESSION_SECRET, payload.encode("utf-8"), hashlib.sha256).hexdigest()
+
+
+def _encode_env(env: ESCEnv) -> str:
+    raw = json.dumps(env.export_state(), separators=(",", ":"), ensure_ascii=False).encode("utf-8")
+    payload = _urlsafe_b64encode(zlib.compress(raw, level=9))
+    return f"{payload}.{_sign(payload)}"
+
+
+def _decode_env(token: str) -> ESCEnv:
+    try:
+        payload, signature = token.rsplit(".", 1)
+    except ValueError as exc:
+        raise RuntimeError("Invalid session token") from exc
+
+    expected = _sign(payload)
+    if not hmac.compare_digest(signature, expected):
+        raise RuntimeError("Invalid session signature")
+
+    try:
+        compressed = _urlsafe_b64decode(payload)
+        data = json.loads(zlib.decompress(compressed).decode("utf-8"))
+    except Exception as exc:
+        raise RuntimeError("Invalid session payload") from exc
+
+    return ESCEnv.from_state(data)
 
 
 def _get_env_for_request(request: Request) -> ESCEnv:
-    session_id = request.cookies.get(SESSION_COOKIE)
-    if not session_id or session_id not in _envs:
+    token = request.cookies.get(SESSION_COOKIE)
+    if not token:
         raise RuntimeError("env.step() called before reset()")
-    return _envs[session_id]
+    return _decode_env(token)
 
 
 @app.get("/")
@@ -61,27 +104,40 @@ def list_tasks() -> dict:
 @app.post("/reset")
 def reset(request: Request, response: Response, req: ResetRequest | None = None) -> dict:
     req = req or ResetRequest()
-    session_id = request.cookies.get(SESSION_COOKIE)
-    if not session_id:
-        session_id = uuid4().hex
-    env = _envs.get(session_id)
-    if env is None:
+    token = request.cookies.get(SESSION_COOKIE)
+    if token:
+        try:
+            env = _decode_env(token)
+        except RuntimeError:
+            env = ESCEnv()
+    else:
         env = ESCEnv()
-        _envs[session_id] = env
     try:
         result = env.reset(task_id=req.task_id, seed=req.seed)
     except KeyError as e:
         raise HTTPException(status_code=400, detail=str(e))
-    response.set_cookie(key=SESSION_COOKIE, value=session_id, httponly=True, samesite="lax")
+    response.set_cookie(
+        key=SESSION_COOKIE,
+        value=_encode_env(env),
+        httponly=True,
+        samesite="lax",
+    )
     return result.model_dump()
 
 
 @app.post("/step")
-def step(req: StepRequest, request: Request) -> dict:
+def step(req: StepRequest, request: Request, response: Response) -> dict:
     try:
-        result = _get_env_for_request(request).step(req.action)
+        env = _get_env_for_request(request)
+        result = env.step(req.action)
     except RuntimeError as e:
         raise HTTPException(status_code=409, detail=str(e))
+    response.set_cookie(
+        key=SESSION_COOKIE,
+        value=_encode_env(env),
+        httponly=True,
+        samesite="lax",
+    )
     return result.model_dump()
 
 

src/env.py

@@ -215,3 +215,75 @@ class ESCEnv:
             }
             for t in TASKS.values()
         ]
+
+    # ------------------------------------------------------------- serialization
+
+    def export_state(self) -> Dict[str, Any]:
+        if self._task is None or self._seeker is None:
+            raise RuntimeError("env.export_state() called before reset()")
+
+        seeker_state = {
+            "distress": self._seeker.distress,
+            "trust": self._seeker.trust,
+            "openness": self._seeker.openness,
+            "revealed": self._seeker.revealed,
+            "stage": self._seeker.stage.value,
+            "last_line_idx_by_stage": {
+                stage.value: idx for stage, idx in self._seeker.last_line_idx_by_stage.items()
+            },
+            "turn": self._seeker.turn,
+        }
+
+        return {
+            "task_id": self._task.id,
+            "turn": self._turn,
+            "done": self._done,
+            "cumulative_reward": self._cumulative_reward,
+            "transcript": list(self._transcript),
+            "agent_messages": list(self._agent_messages),
+            "had_safety_reference": self._had_safety_reference,
+            "seeker": seeker_state,
+        }
+
+    @classmethod
+    def from_state(cls, data: Dict[str, Any]) -> "ESCEnv":
+        task = get_task(str(data["task_id"]))
+        seeker_data = data["seeker"]
+
+        env = cls()
+        env._task = task
+        env._turn = int(data["turn"])
+        env._done = bool(data["done"])
+        env._cumulative_reward = float(data["cumulative_reward"])
+        env._transcript = list(data.get("transcript", []))
+        env._agent_messages = list(data.get("agent_messages", []))
+        env._had_safety_reference = bool(data.get("had_safety_reference", False))
+        env._seeker = SeekerState(
+            persona=task.persona,
+            distress=float(seeker_data["distress"]),
+            trust=float(seeker_data["trust"]),
+            openness=float(seeker_data["openness"]),
+            revealed=bool(seeker_data["revealed"]),
+            stage=Stage(str(seeker_data["stage"])),
+            last_line_idx_by_stage={
+                Stage(stage_name): int(idx)
+                for stage_name, idx in seeker_data["last_line_idx_by_stage"].items()
+            },
+            turn=int(seeker_data["turn"]),
+        )
+
+        if env._transcript:
+            last_seeker_text = next(
+                (entry["text"] for entry in reversed(env._transcript) if entry.get("role") == "seeker"),
+                task.persona.surface_concern,
+            )
+            env._last_obs = Observation(
+                seeker_utterance=last_seeker_text,
+                turn=env._turn,
+                remaining_turns=max(0, task.max_turns - env._turn),
+                stage_hint=env._seeker.stage.value,
+                task_id=task.id,
+                scenario_brief=task.persona.scenario_brief,
+            )
+
+        return env