Fix stateful UI session, enrich FHIR cache, fix Dockerfile

Dockerfile

@@ -1,81 +1,50 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
+# Dockerfile for Northflank CI/CD.
 #
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
-
-# Multi-stage build using openenv-base
-# This Dockerfile is flexible and works for both:
-# - In-repo environments (with local OpenEnv sources)
-# - Standalone environments (with openenv from PyPI/Git)
-# The build script (openenv build) handles context detection and sets appropriate build args.
+# Runs the OpenEnv environment server — exposes /reset, /step, /state,
+# /schema, /ws, and the UI at /.  Training connects separately via ENV_URL.
+#
+# No GPU required: the env server only does FHIR lookups against the
+# baked-in cache and runs the FastAPI server.
 
-ARG BASE_IMAGE=ghcr.io/meta-pytorch/openenv-base:latest
-FROM ${BASE_IMAGE} AS builder
+FROM python:3.11-slim
 
-WORKDIR /app
+ENV DEBIAN_FRONTEND=noninteractive
 
-# Ensure git is available (required for installing dependencies from VCS)
-RUN apt-get update && \
-    apt-get install -y --no-install-recommends git && \
+# ── System dependencies ────────────────────────────────────────────────────
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        curl git build-essential && \
     rm -rf /var/lib/apt/lists/*
 
-# Build argument to control whether we're building standalone or in-repo
-ARG BUILD_MODE=in-repo
-ARG ENV_NAME=medagentbench_env
-
-# Copy environment code (always at root of build context)
-COPY . /app/env
-
-# For in-repo builds, openenv is already vendored in the build context
-# For standalone builds, openenv will be installed via pyproject.toml
-WORKDIR /app/env
-
-# Ensure uv is available (for local builds where base image lacks it)
-RUN if ! command -v uv >/dev/null 2>&1; then \
-        curl -LsSf https://astral.sh/uv/install.sh | sh && \
-        mv /root/.local/bin/uv /usr/local/bin/uv && \
-        mv /root/.local/bin/uvx /usr/local/bin/uvx; \
-    fi
-
-# Install dependencies using uv sync
-# If uv.lock exists, use it; otherwise resolve on the fly
-RUN --mount=type=cache,target=/root/.cache/uv \
-    if [ -f uv.lock ]; then \
-        uv sync --frozen --no-install-project --no-editable; \
-    else \
-        uv sync --no-install-project --no-editable; \
-    fi
-
-RUN --mount=type=cache,target=/root/.cache/uv \
-    if [ -f uv.lock ]; then \
-        uv sync --frozen --no-editable; \
-    else \
-        uv sync --no-editable; \
-    fi
-
-# Final runtime stage
-FROM ${BASE_IMAGE}
+# ── uv ────────────────────────────────────────────────────────────────────
+RUN curl -LsSf https://astral.sh/uv/install.sh | sh && \
+    mv /root/.local/bin/uv   /usr/local/bin/uv && \
+    mv /root/.local/bin/uvx  /usr/local/bin/uvx
 
 WORKDIR /app
 
-# Copy the virtual environment from builder
-COPY --from=builder /app/env/.venv /app/.venv
+# ── Source code ───────────────────────────────────────────────────────────
+COPY medagentbench_env/ ./medagentbench_env/
 
-# Copy the environment code
-COPY --from=builder /app/env /app/env
+# ── MedAgentBench eval module (refsol graders) ────────────────────────────
+COPY medagentbenchv2/medagentbench_v2/src/ ./medagentbenchv2/medagentbench_v2/src/
 
-# Set PATH to use the virtual environment
+# ── Python package + deps (env server only, no train extras) ──────────────
+RUN uv venv --python 3.11 /app/.venv && \
+    . /app/.venv/bin/activate && \
+    uv pip install -e "medagentbench_env"
+
+# ── Runtime environment ───────────────────────────────────────────────────
 ENV PATH="/app/.venv/bin:$PATH"
+ENV PYTHONPATH="/app:$PYTHONPATH"
+ENV ENABLE_WEB_INTERFACE=true
 
-# Set PYTHONPATH so imports work correctly
-ENV PYTHONPATH="/app/env:$PYTHONPATH"
+# ── Expose env server port ────────────────────────────────────────────────
+EXPOSE 8000
 
-# Health check
-HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \
+HEALTHCHECK --interval=30s --timeout=3s --start-period=15s --retries=3 \
     CMD curl -f http://localhost:8000/health || exit 1
 
-# Run the FastAPI server
-# The module path is constructed to work with the /app/env structure
-ENV ENABLE_WEB_INTERFACE=true
-CMD ["sh", "-c", "cd /app/env && uvicorn server.app:app --host 0.0.0.0 --port 8000"]
+# Run the OpenEnv environment server.
+# Training connects to this service via the ENV_URL env var.
+CMD ["uvicorn", "medagentbench_env.server.app:app", \
+     "--host", "0.0.0.0", "--port", "8000"]

server/app.py

@@ -20,6 +20,7 @@ Usage:
 
 import json
 from pathlib import Path
+from typing import Optional
 
 try:
     from openenv.core.env_server.http_server import create_app
@@ -36,6 +37,13 @@ from starlette.requests import Request
 from medagentbench_env.models import MedAgentBenchAction, MedAgentBenchObservation
 from .medagentbench_env_environment import MedAgentBenchEnvironment
 
+# ---------------------------------------------------------------------------
+# Stateful UI session — one persistent environment instance shared across
+# /api/reset and /api/step so step_count and task context survive between calls.
+# (The built-in /reset and /step from OpenEnv create a fresh env per request.)
+# ---------------------------------------------------------------------------
+_ui_env: Optional[MedAgentBenchEnvironment] = None
+
 _ROOT = Path(__file__).parent.parent
 _UI_HTML = (_ROOT / "ui" / "index.html").read_text()
 
@@ -81,6 +89,42 @@ async def get_tasks():
     ])
 
 
+@app.post("/api/reset")
+async def api_reset(request: Request):
+    """Stateful reset for the UI — creates a persistent env instance."""
+    global _ui_env
+    body = {}
+    try:
+        body = await request.json()
+    except Exception:
+        pass
+    task_index = body.get("task_index", 0)
+    _ui_env = MedAgentBenchEnvironment()
+    obs = _ui_env.reset(task_index=task_index)
+    obs_dict = obs.model_dump(exclude={"reward", "done", "metadata"})
+    return JSONResponse({"observation": obs_dict, "reward": obs.reward, "done": obs.done})
+
+
+@app.post("/api/step")
+async def api_step(request: Request):
+    """Stateful step for the UI — uses the same env instance across calls."""
+    global _ui_env
+    if _ui_env is None:
+        raise HTTPException(status_code=400, detail="No active session. Call /api/reset first.")
+    body = {}
+    try:
+        body = await request.json()
+    except Exception:
+        pass
+    try:
+        action = MedAgentBenchAction.model_validate(body.get("action", {}))
+    except Exception as e:
+        raise HTTPException(status_code=422, detail=str(e))
+    obs = _ui_env.step(action)
+    obs_dict = obs.model_dump(exclude={"reward", "done", "metadata"})
+    return JSONResponse({"observation": obs_dict, "reward": obs.reward, "done": obs.done})
+
+
 @app.get("/api/baseline-results")
 async def get_baseline_results():
     """Return pre-computed baseline evaluation results."""

ui/index.html

@@ -640,7 +640,7 @@ async function startSession() {
 
   // Call /reset
   try {
-    const r = await fetch('/reset', {
+    const r = await fetch('/api/reset', {
       method: 'POST',
       headers: {'Content-Type':'application/json'},
       body: JSON.stringify({task_index: selectedTask.index})
@@ -761,7 +761,7 @@ async function sendAction() {
   document.getElementById('send-btn').disabled = true;
 
   try {
-    const r = await fetch('/step', {
+    const r = await fetch('/api/step', {
       method: 'POST',
       headers: {'Content-Type':'application/json'},
       body: JSON.stringify({
@@ -807,7 +807,7 @@ function buildQuickButtons() {
 
   // Task-specific GET shortcuts
   if (type === 'task10') {
-    gets.splice(2, 0, { label: '🩸 A1C (4548-4)', path: `Observation?patient=${mrn}&code=4548-4&_sort=-date`, resource: 'Observation' });
+    gets.splice(2, 0, { label: '🩸 A1C (A1C/4548-4)', path: `Observation?patient=${mrn}&code=A1C&_count=5000`, resource: 'Observation' });
   }
   if (type === 'task3') {
     gets.splice(2, 0, { label: '💓 Vital Signs', path: `Observation?patient=${mrn}&category=vital-signs&_sort=-date`, resource: 'Observation' });
@@ -918,19 +918,32 @@ function appendAgentAction(type, url, body, answer, raw) {
 
 function appendFhirResponse(text) {
   const id = `resp-${traceSteps.length}`;
-  let parsed = null, summary = '';
+  let parsed = null, summary = '', displayText = text;
+
+  // The env wraps FHIR JSON as: "Here is the response from the GET request:\n{JSON}. Please call FINISH..."
+  // Extract the embedded JSON so we can parse and pretty-print it.
+  const prefix = 'Here is the response from the GET request:\n';
+  const suffix = '. Please call FINISH';
+  const prefixIdx = text.indexOf(prefix);
+  if (prefixIdx !== -1) {
+    const afterPrefix = text.substring(prefixIdx + prefix.length);
+    const suffixIdx = afterPrefix.lastIndexOf(suffix);
+    displayText = suffixIdx !== -1 ? afterPrefix.substring(0, suffixIdx) : afterPrefix;
+  }
+
   try {
-    parsed = JSON.parse(text);
+    parsed = JSON.parse(displayText);
     const total = parsed?.total ?? parsed?.entry?.length;
     const rtype = parsed?.resourceType;
     if (rtype === 'Bundle') {
-      summary = `Bundle · ${parsed.entry?.length ?? 0} entries${total !== undefined ? ` (total ${total})` : ''}`;
+      const count = parsed.entry?.length ?? 0;
+      summary = `Bundle · ${count} ${count === 1 ? 'entry' : 'entries'}${total !== undefined ? ` (total ${total})` : ''}`;
     } else if (rtype) {
       summary = `${rtype}`;
     }
   } catch {}
 
-  const prettyText = parsed ? JSON.stringify(parsed, null, 2) : text;
+  const prettyText = parsed ? JSON.stringify(parsed, null, 2) : displayText;
   const shortText = prettyText.length > 2000 ? prettyText.substring(0, 2000) + '\n… (truncated)' : prettyText;
 
   const div = document.createElement('div');