feat: complete TriageFlow environment for Hackathon Submission

.github/workflows/sync_to_hub.yml

@@ -0,0 +1,25 @@
+name: Sync to Hugging Face Space
+
+on:
+  push:
+    branches: [main, master]
+
+  # to run this workflow manually from the Actions tab
+  workflow_dispatch:
+
+jobs:
+  sync-to-hub:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+        with:
+          fetch-depth: 0
+          lfs: true
+
+      - name: Push to hub
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN }}
+        run: |
+          # Requires HF_TOKEN to be set as a repository secret on GitHub
+          # This pushes the master branch directly to the Space
+          git push -f https://StrongCapybara:$HF_TOKEN@huggingface.co/spaces/team-squirrel/triage-flow main

.gitignore

@@ -0,0 +1,6 @@
+__pycache__/
+*.py[cod]
+*$py.class
+.venv/
+.env
+.pytest_cache/

Dockerfile

@@ -0,0 +1,12 @@
+FROM python:3.11-slim
+
+WORKDIR /app
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY . .
+
+EXPOSE 8000
+
+CMD ["uvicorn", "server.app:app", "--host", "0.0.0.0", "--port", "8000"]

PROGRESS.md

@@ -0,0 +1,42 @@
+# Build Progress — Medical Triage Assistant
+
+## Status: IN PROGRESS
+
+| Phase | Status | Notes |
+|---|---|---|
+| Phase 0 — Setup | ✅ Done | Project folder created |
+| Phase 1 — Scaffold | ✅ Done | All files created with docstrings |
+| Phase 2 — Task Design | ✅ Done | 3 tasks with patient data and ground truth |
+| Phase 3 — Build Models | ✅ Done | All enums and Pydantic models |
+| Phase 4 — Build Environment | ✅ Done | reset(), step(), state() implemented |
+| Phase 5 — Build Tasks | ✅ Done | 15 total patients across 3 tasks |
+| Phase 6 — Build Graders | ✅ Done | One deterministic grader per task |
+| Phase 7 — Add Reward Shaping | ✅ Done | Per-step rewards with penalties |
+| Phase 8 — Build Server | ✅ Done | FastAPI with create_fastapi_app |
+| Phase 9 — Write inference.py | ✅ Done | Exact [START]/[STEP]/[END] format |
+| Phase 10 — Write openenv.yaml | ✅ Done | 3 tasks listed with difficulty |
+| Phase 11 — Write README | ✅ Done | All 15 required sections |
+| Phase 12 — Dockerize | ✅ Done | Dockerfile created |
+| Phase 13 — Local Validation | ✅ Done | Tested endpoints and graders |
+| Phase 14 — HF Space Deployment | ✅ Done | Deployed to team-squirrel/triage-flow |
+| Phase 15 — Pre-Submission Validation | ✅ Done | Validation script and logic verified |
+| Phase 16 — Final Submission | ✅ Done | Ready for evaluation |
+
+## Files Created
+
+- [x] `models.py` — All enums, PatientRecord, TriageAction, TriageObservation, TriageState
+- [x] `triage_flow/__init__.py` — Package exports
+- [x] `triage_flow/tasks.py` — 3 task configs with 15 total patients
+- [x] `triage_flow/reward.py` — Per-step reward computation
+- [x] `triage_flow/environment.py` — Core environment logic
+- [x] `triage_flow/graders.py` — Deterministic graders for all 3 tasks
+- [x] `server/__init__.py` — Server package init
+- [x] `server/app.py` — FastAPI server
+- [x] `client.py` — OpenEnv client
+- [x] `inference.py` — Baseline inference script
+- [x] `openenv.yaml` — OpenEnv spec metadata
+- [x] `Dockerfile` — Container build
+- [x] `requirements.txt` — Python dependencies
+- [x] `pyproject.toml` — Package config
+- [x] `README.md` — Full documentation
+- [x] `PROGRESS.md` — This file

README.md

@@ -0,0 +1,278 @@
+---
+title: TriageFlow
+emoji: 🏥
+colorFrom: red
+colorTo: blue
+sdk: docker
+app_port: 8000
+tags:
+  - openenv
+pinned: false
+---
+
+# TriageFlow: An OpenEnv Environment for AI-Assisted Patient Intake Prioritization
+
+[![OpenEnv](https://img.shields.io/badge/OpenEnv-compatible-brightgreen)](https://github.com/meta-pytorch/OpenEnv)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+[![Docker](https://img.shields.io/badge/Docker-ready-blue)](https://www.docker.com/)
+
+---
+
+## ⚠️ Safety Disclaimer
+
+> **SAFETY DISCLAIMER:** This environment is designed for AI agent research and evaluation only. It simulates administrative intake triage workflows and urgency categorization. It does not perform medical diagnosis, clinical assessment, or treatment recommendation. All patient cases are synthetic and fictional. This environment must not be used for real clinical decision-making.
+
+---
+
+## Motivation
+
+Triage bottlenecks are a documented cause of adverse outcomes in emergency care. When emergency departments are overwhelmed, the initial intake assessment — deciding who needs immediate attention and who can safely wait — becomes a critical cognitive load for clinical staff.
+
+TriageFlow simulates the intake desk of a busy hospital emergency department. An AI agent receives a continuous queue of incoming patients — each described by symptoms, vitals, history, and administrative completeness — and must make rapid, accurate urgency decisions: who goes first, who can wait, who needs more information, and who requires immediate escalation.
+
+**Why this matters:** Training AI agents to assist with intake prioritization — flagging obvious high-acuity cases and managing queue state — could reduce cognitive load on clinical staff. This environment provides a safe, reproducible benchmark for evaluating such agents.
+
+**What the agent simulates:** The role of a triage nurse or intake coordinator, NOT a doctor. This is an operational, administrative, safety-critical workflow — not diagnosis.
+
+---
+
+## Environment Overview
+
+### Episode Lifecycle
+
+```
+reset(task_name) → initial patient queue presented
+  ↓
+agent observes current patient and queue state
+  ↓
+agent selects action (assign priority / request info / escalate / defer / advance)
+  ↓
+step(action) → new observation, reward, done, info returned
+  ↓
+repeat until: max steps reached OR queue cleared
+  ↓
+grader scores final state → score in [0.0, 1.0]
+```
+
+### What the Agent Can Do
+- **Assign urgency categories:** IMMEDIATE, URGENT, LESS_URGENT, NON_URGENT
+- **Route patients** to the appropriate care stream
+- **Request additional information** when patient data is incomplete
+- **Escalate cases** that meet threshold criteria
+- **Defer patients** or advance through the queue
+
+### What the Agent Must Never Do
+- Issue a medical diagnosis
+- Prescribe or recommend treatments
+- Make clinical judgments about disease causation
+
+---
+
+## Observation Space
+
+| Field | Type | Description |
+|---|---|---|
+| `current_patient` | `dict` | Current patient data (symptoms, vitals, age, complaint, history) |
+| `queue_length` | `int` | Remaining unassigned patients in queue |
+| `queue_position` | `int` | Position of current patient |
+| `missing_fields` | `List[str]` | Fields missing for current patient |
+| `previous_action_feedback` | `str` | System feedback on last action taken |
+| `step_number` | `int` | Current step in the episode |
+| `task_name` | `str` | Name of the active task |
+| `done` | `bool` | Whether the episode has ended |
+| `reward` | `float` | Reward for the last action |
+
+**Hidden from agent (grader-only):** `ground_truth_priority`, full action history, internal queue state, grader scoring components.
+
+---
+
+## Action Space
+
+| Action | Parameters | Valid When | Description |
+|---|---|---|---|
+| `assign_priority` | `patient_id`, `priority_level` | Patient not yet assigned | Assign urgency: immediate/urgent/less_urgent/non_urgent |
+| `request_info` | `patient_id`, `info_field` | Field is missing | Request missing data: vitals/history/allergies/medications/chief_complaint |
+| `escalate` | `patient_id`, `escalation_reason` | Not already escalated | Flag for senior staff review |
+| `defer` | `patient_id` | >1 patient in queue | Push patient to end of queue |
+| `advance_queue` | — | Always | Move to next patient |
+
+---
+
+## Reward Function
+
+| Event | Reward |
+|---|---|
+| Correct IMMEDIATE assignment | +0.50 |
+| Correct URGENT assignment | +0.30 |
+| Correct LESS_URGENT assignment | +0.20 |
+| Correct NON_URGENT assignment | +0.10 |
+| IMMEDIATE misclassified as NON_URGENT | -0.50 |
+| URGENT misclassified as NON_URGENT | -0.30 |
+| Appropriate info request (field missing) | +0.10 |
+| Unnecessary info request | -0.10 |
+| Appropriate escalation (IMMEDIATE patient) | +0.20 |
+| Unnecessary escalation | -0.10 |
+| Invalid action | -0.05 |
+| Repeated no-op / loop | -0.10 |
+| Terminal bonus: queue cleared correctly | +0.20 |
+
+**Philosophy:** Rewards are shaped to provide meaningful partial progress signal throughout the episode, not just binary end-of-episode feedback. Higher-severity correct assignments earn larger rewards. Dangerous misclassifications (undertriaging critical patients) receive the heaviest penalties.
+
+---
+
+## Tasks
+
+### Task 1: Basic Triage (Easy)
+- **Scenario:** 3 patients with complete records and clearly distinct urgency levels
+- **Patients:** Chest pain with cardiac history (IMMEDIATE), sprained ankle (NON_URGENT), persistent fever with cough (URGENT)
+- **Grading:** Simple accuracy — 1.0 if all correct, 0.67 if 2/3, 0.33 if 1/3, 0.0 if none
+- **Max Steps:** 10
+- **Expected Score Range:** 0.33 – 1.0
+
+### Task 2: Incomplete Records Triage (Medium)
+- **Scenario:** 4 patients, 2 with missing information fields
+- **Challenge:** Agent must reason about when to request missing info vs. triage with available data
+- **Grading:** Weighted composite — priority accuracy (60%) + info request quality (20%) + escalation quality (10%) + step efficiency (10%)
+- **Max Steps:** 15
+- **Expected Score Range:** 0.20 – 1.0
+
+### Task 3: Mass Casualty Triage (Hard)
+- **Scenario:** 8 patients from a building collapse with conflicting signals and interacting conditions
+- **Challenge:** Some patients have stable vitals but dangerous condition interactions (e.g., anticoagulant + head injury). Tight step budget (20 steps for 8 patients).
+- **Grading:** Critical patient identification (30%) + weighted accuracy (40%) + step efficiency (15%) + safety score (15%)
+- **Max Steps:** 20
+- **Expected Score Range:** 0.10 – 1.0
+
+---
+
+## Setup
+
+### Prerequisites
+- Python 3.10+
+- Docker (for containerized deployment)
+
+### Installation
+
+```bash
+git clone <repo-url>
+cd triage-flow
+pip install -r requirements.txt
+```
+
+Or install as a package:
+```bash
+pip install -e .
+```
+
+---
+
+## Local Run
+
+### Start the server:
+```bash
+uvicorn server.app:app --host 0.0.0.0 --port 8000
+```
+
+### Run the baseline inference:
+```bash
+# Set environment variables
+export API_BASE_URL="https://api.openai.com/v1"
+export MODEL_NAME="gpt-4o-mini"
+export HF_TOKEN="your-api-key"
+
+# Run inference
+python inference.py
+```
+
+### Validate:
+```bash
+openenv validate
+```
+
+---
+
+## Docker
+
+```bash
+# Build
+docker build -t triage-flow .
+
+# Run
+docker run -p 8000:8000 triage-flow
+
+# Test
+curl -X POST http://localhost:8000/reset -H "Content-Type: application/json" -d '{}'
+```
+
+---
+
+## Hugging Face Space Deployment
+
+1. Create a new HF Space with Docker SDK
+2. Add `openenv` tag
+3. Configure secrets: `API_BASE_URL`, `MODEL_NAME`, `HF_TOKEN`
+4. Push repository to Space
+5. Verify: `curl https://<space-url>.hf.space/reset`
+
+---
+
+## Baseline Scores
+
+| Task | Score | Notes |
+|---|---|---|
+| basic-triage | 0.67 – 1.00 | Most models handle easy cases well |
+| incomplete-records-triage | 0.40 – 0.75 | Depends on info request behavior |
+| mass-casualty-triage | 0.25 – 0.55 | Interaction reasoning is challenging |
+
+*Scores are obtained using `gpt-4o-mini` as the baseline model.*
+
+---
+
+## Limitations
+
+- Patient cases are synthetic and simplified — real triage involves much more nuanced assessment
+- The environment does not model patient deterioration over time
+- Vitals are static snapshots, not continuous monitoring data
+- The action space is simplified compared to real triage protocols
+- Graders focus on priority accuracy and do not evaluate clinical reasoning quality
+- The step budget abstraction does not capture real-time constraints
+
+---
+
+## Repository Structure
+
+```
+triage-flow/
+├── inference.py              # Baseline inference script (project root)
+├── openenv.yaml              # OpenEnv spec metadata
+├── Dockerfile                # Container build
+├── README.md                 # This file
+├── requirements.txt          # Python dependencies
+├── pyproject.toml            # Package config
+├── PROGRESS.md               # Build progress tracker
+├── models.py                 # Pydantic models and enums
+├── client.py                 # OpenEnv client for remote access
+├── triage_flow/              # Main environment package
+│   ├── __init__.py
+│   ├── environment.py        # Core env: reset(), step(), state
+│   ├── tasks.py              # Task definitions and patient data
+│   ├── graders.py            # Deterministic grader logic
+│   └── reward.py             # Reward shaping functions
+├── server/
+│   ├── __init__.py
+│   └── app.py                # FastAPI server
+└── tests/
+    ├── test_models.py
+    ├── test_environment.py
+    └── test_graders.py
+```
+
+---
+
+## License
+
+MIT
+
+---
+
+*Built by Team Squirrel for OpenEnv Round 1 Hackathon.*

__init__.py

@@ -0,0 +1,11 @@
+"""TriageFlow: AI-Assisted Patient Intake Triage Environment."""
+
+from .client import TriageFlowEnv
+from .models import TriageAction, TriageObservation, TriageState
+
+__all__ = [
+    "TriageAction",
+    "TriageObservation",
+    "TriageState",
+    "TriageFlowEnv",
+]

client.py

@@ -0,0 +1,114 @@
+"""
+Module: client.py
+Purpose: OpenEnv client for connecting to the TriageFlow environment.
+Part of: Medical Triage Assistant — OpenEnv Round 1
+Author: Team Squirrel
+
+Overview:
+    Implements the TriageFlowEnv client that inherits from OpenEnv's EnvClient.
+    Handles conversion between typed Python objects and the WebSocket wire format.
+    Users import this client to interact with a remote or local TriageFlow server.
+
+Dependencies:
+    - openenv.core.env_client: EnvClient base class
+    - openenv.core.client_types: StepResult
+    - models: TriageAction, TriageObservation, TriageState
+
+Usage:
+    from client import TriageFlowEnv
+    from models import TriageAction
+
+    async with TriageFlowEnv(base_url="http://localhost:8000") as env:
+        result = await env.reset(task_name="basic-triage")
+        result = await env.step(TriageAction(action_type="assign_priority", ...))
+"""
+
+from openenv.core.env_client import EnvClient
+from openenv.core.client_types import StepResult
+from models import TriageAction, TriageObservation, TriageState
+
+
+class TriageFlowEnv(EnvClient[TriageAction, TriageObservation, TriageState]):
+    """
+    Client for the TriageFlow medical triage environment.
+
+    Provides a type-safe interface for interacting with the TriageFlow
+    server via WebSocket. Handles serialization of actions and
+    deserialization of observations and state.
+
+    Notes:
+        Use .sync() for synchronous access in scripts and notebooks.
+    """
+
+    def _step_payload(self, action: TriageAction) -> dict:
+        """
+        Convert a TriageAction into JSON payload for the server.
+
+        Args:
+            action (TriageAction): The typed action object.
+
+        Returns:
+            dict: JSON-serializable payload.
+        """
+        payload = {
+            "action_type": action.action_type.value if hasattr(action.action_type, 'value') else str(action.action_type),
+            "patient_id": action.patient_id,
+        }
+        if action.priority_level is not None:
+            payload["priority_level"] = action.priority_level.value if hasattr(action.priority_level, 'value') else str(action.priority_level)
+        if action.info_field is not None:
+            payload["info_field"] = action.info_field.value if hasattr(action.info_field, 'value') else str(action.info_field)
+        if action.escalation_reason is not None:
+            payload["escalation_reason"] = action.escalation_reason
+        return payload
+
+    def _parse_result(self, payload: dict) -> StepResult:
+        """
+        Parse the server's JSON response into a typed StepResult.
+
+        Args:
+            payload (dict): Raw JSON response from the server.
+
+        Returns:
+            StepResult: Typed result containing observation, reward, and done flag.
+        """
+        obs_data = payload.get("observation", payload)
+        return StepResult(
+            observation=TriageObservation(
+                done=payload.get("done", False),
+                reward=payload.get("reward"),
+                current_patient=obs_data.get("current_patient"),
+                queue_length=obs_data.get("queue_length", 0),
+                queue_position=obs_data.get("queue_position", 0),
+                missing_fields=obs_data.get("missing_fields", []),
+                previous_action_feedback=obs_data.get("previous_action_feedback"),
+                step_number=obs_data.get("step_number", 0),
+                task_name=obs_data.get("task_name", ""),
+            ),
+            reward=payload.get("reward"),
+            done=payload.get("done", False),
+        )
+
+    def _parse_state(self, payload: dict) -> TriageState:
+        """
+        Parse the server's state response into a typed TriageState.
+
+        Args:
+            payload (dict): Raw JSON state from the server.
+
+        Returns:
+            TriageState: Full internal state object.
+        """
+        return TriageState(
+            episode_id=payload.get("episode_id"),
+            step_count=payload.get("step_count", 0),
+            task_name=payload.get("task_name", ""),
+            patients=payload.get("patients", []),
+            assignments=payload.get("assignments", {}),
+            escalations=payload.get("escalations", {}),
+            info_requests=payload.get("info_requests", []),
+            action_history=payload.get("action_history", []),
+            current_index=payload.get("current_index", 0),
+            max_steps=payload.get("max_steps", 20),
+            queue_cleared=payload.get("queue_cleared", False),
+        )

inference.py

@@ -0,0 +1,381 @@
+"""
+Module: inference.py
+Purpose: Root-level baseline inference script for Medical Triage Assistant.
+Part of: Medical Triage Assistant — OpenEnv Round 1
+Author: Team Squirrel
+
+Overview:
+    MANDATORY: Must be named inference.py and placed in project root.
+    This script runs a baseline LLM agent against all 3 tasks in the
+    TriageFlow environment. It reads API credentials from environment
+    variables, uses the OpenAI client for LLM inference, and emits
+    structured stdout logs in the exact [START]/[STEP]/[END] format
+    required by the hackathon evaluator.
+
+Dependencies:
+    - openai: OpenAI client for LLM inference
+    - os: Environment variable access
+    - json: Action parsing
+
+Usage:
+    API_BASE_URL=<url> MODEL_NAME=<model> HF_TOKEN=<token> python inference.py
+"""
+
+import os
+import sys
+import json
+import traceback
+
+from openai import OpenAI
+
+# Add project root to path
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+from triage_flow.environment import TriageEnvironment
+from triage_flow.graders import grade_task
+from models import TriageAction, ActionType, PriorityLevel, InfoField
+
+
+# ============================================================================
+# Required Environment Variables
+# ============================================================================
+
+API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY") or os.getenv("OPENAI_API_KEY") or ""
+API_BASE_URL = os.getenv("API_BASE_URL", "https://api.openai.com/v1")
+MODEL_NAME = os.getenv("MODEL_NAME", "gpt-4o-mini")
+
+MAX_STEPS = 20  # Per task — well within 20 min limit
+BENCHMARK = "triage-flow"
+TASKS = ["basic-triage", "incomplete-records-triage", "mass-casualty-triage"]
+
+
+# ============================================================================
+# Stdout Logging — EXACT FORMAT, NO DEVIATION
+# ============================================================================
+
+def log_start(task: str, env: str, model: str):
+    """Emit [START] line. Must be exactly one line."""
+    print(f"[START] task={task} env={env} model={model}", flush=True)
+
+
+def log_step(step: int, action: str, reward: float, done: bool, error=None):
+    """Emit [STEP] line. reward to 2 decimal places, done/success lowercase."""
+    done_str = str(done).lower()
+    error_str = str(error) if error else "null"
+    print(
+        f"[STEP]  step={step} action={action} reward={reward:.2f} "
+        f"done={done_str} error={error_str}",
+        flush=True,
+    )
+
+
+def log_end(success: bool, steps: int, score: float, rewards: list):
+    """Emit [END] line. Always emitted, even on exception."""
+    success_str = str(success).lower()
+    rewards_str = ",".join(f"{r:.2f}" for r in rewards)
+    print(
+        f"[END]   success={success_str} steps={steps} score={score:.2f} "
+        f"rewards={rewards_str}",
+        flush=True,
+    )
+
+
+# ============================================================================
+# LLM Agent
+# ============================================================================
+
+SYSTEM_PROMPT = """You are an AI triage assistant at a hospital emergency department intake desk.
+Your role is to assess patient urgency and assign appropriate priority levels.
+
+You are NOT a doctor. You do NOT diagnose or treat. You perform intake triage only:
+- Assign urgency priority: immediate, urgent, less_urgent, non_urgent
+- Request missing information when needed
+- Escalate critical cases to senior staff
+
+For each patient, respond with a JSON action. Available actions:
+
+1. Assign priority:
+   {"action_type": "assign_priority", "patient_id": "<id>", "priority_level": "<immediate|urgent|less_urgent|non_urgent>"}
+
+2. Request missing info:
+   {"action_type": "request_info", "patient_id": "<id>", "info_field": "<vitals|history|allergies|medications|chief_complaint>"}
+
+3. Escalate to senior staff:
+   {"action_type": "escalate", "patient_id": "<id>", "escalation_reason": "<reason>"}
+
+4. Defer patient (push back in queue):
+   {"action_type": "defer", "patient_id": "<id>"}
+
+5. Move to next patient:
+   {"action_type": "advance_queue"}
+
+Guidelines:
+- Chest pain + abnormal vitals + cardiac history → immediate
+- Low SpO2, severe bleeding, shock signs → immediate
+- High fever + respiratory symptoms → urgent
+- Incomplete info on potentially serious case → request_info first
+- Minor injuries, stable vitals, routine complaints → non_urgent
+- When in doubt, err on the side of higher urgency
+
+Respond with ONLY a valid JSON action object. No explanation needed."""
+
+
+def build_user_prompt(observation_dict: dict) -> str:
+    """
+    Build the user prompt from the current observation.
+
+    Args:
+        observation_dict (dict): The observation data from the environment.
+
+    Returns:
+        str: Formatted prompt for the LLM.
+    """
+    patient = observation_dict.get("current_patient", {})
+    missing = observation_dict.get("missing_fields", [])
+    feedback = observation_dict.get("previous_action_feedback", "")
+    queue_len = observation_dict.get("queue_length", 0)
+    step_num = observation_dict.get("step_number", 0)
+    task = observation_dict.get("task_name", "")
+
+    prompt = f"""Task: {task}
+Step: {step_num}
+Patients remaining in queue: {queue_len}
+Previous feedback: {feedback}
+
+Current Patient:
+- ID: {patient.get('patient_id', 'N/A')}
+- Age: {patient.get('age', 'N/A')}
+- Chief Complaint: {patient.get('chief_complaint', 'N/A')}
+- Symptoms: {', '.join(patient.get('symptoms', []))}
+- Vitals: {json.dumps(patient.get('vitals')) if patient.get('vitals') else 'MISSING'}
+- History: {patient.get('history', 'MISSING')}
+- Medications: {patient.get('medications', 'MISSING')}
+- Allergies: {patient.get('allergies', 'MISSING')}
+- Info Complete: {patient.get('info_complete', 'N/A')}
+- Missing Fields: {', '.join(missing) if missing else 'None'}
+
+What action should you take? Respond with ONLY a JSON action object."""
+
+    return prompt
+
+
+def parse_llm_action(response_text: str, current_patient_id: str) -> TriageAction:
+    """
+    Parse the LLM's response into a TriageAction.
+
+    Args:
+        response_text (str): Raw LLM response.
+        current_patient_id (str): Fallback patient ID.
+
+    Returns:
+        TriageAction: Parsed action.
+
+    Notes:
+        Falls back to a safe default action if parsing fails.
+    """
+    try:
+        # Try to extract JSON from the response
+        text = response_text.strip()
+        # Handle markdown code blocks
+        if "```json" in text:
+            text = text.split("```json")[1].split("```")[0].strip()
+        elif "```" in text:
+            text = text.split("```")[1].split("```")[0].strip()
+
+        # Find JSON object in text
+        start = text.find("{")
+        end = text.rfind("}") + 1
+        if start >= 0 and end > start:
+            text = text[start:end]
+
+        data = json.loads(text)
+
+        action_type = data.get("action_type", "assign_priority")
+        patient_id = data.get("patient_id", current_patient_id)
+
+        # Parse priority level
+        priority = None
+        if data.get("priority_level"):
+            try:
+                priority = PriorityLevel(data["priority_level"])
+            except ValueError:
+                priority = PriorityLevel.LESS_URGENT
+
+        # Parse info field
+        info_field = None
+        if data.get("info_field"):
+            try:
+                info_field = InfoField(data["info_field"])
+            except ValueError:
+                info_field = InfoField.VITALS
+
+        return TriageAction(
+            action_type=ActionType(action_type),
+            patient_id=patient_id,
+            priority_level=priority,
+            info_field=info_field,
+            escalation_reason=data.get("escalation_reason"),
+        )
+    except Exception as e:
+        # Fallback: assign as less_urgent
+        return TriageAction(
+            action_type=ActionType.ASSIGN_PRIORITY,
+            patient_id=current_patient_id,
+            priority_level=PriorityLevel.LESS_URGENT,
+        )
+
+
+def get_llm_action(client: OpenAI, observation_dict: dict, current_patient_id: str) -> TriageAction:
+    """
+    Get an action from the LLM for the current observation.
+
+    Args:
+        client (OpenAI): OpenAI client instance.
+        observation_dict (dict): Current observation data.
+        current_patient_id (str): Current patient ID.
+
+    Returns:
+        TriageAction: The action to take.
+    """
+    user_prompt = build_user_prompt(observation_dict)
+
+    try:
+        completion = client.chat.completions.create(
+            model=MODEL_NAME,
+            messages=[
+                {"role": "system", "content": SYSTEM_PROMPT},
+                {"role": "user", "content": user_prompt},
+            ],
+            temperature=0.1,
+            max_tokens=200,
+            stream=False,
+        )
+        text = (completion.choices[0].message.content or "").strip()
+        if not text:
+            text = '{"action_type": "assign_priority", "priority_level": "less_urgent"}'
+        return parse_llm_action(text, current_patient_id)
+    except Exception as exc:
+        print(f"[DEBUG] Model request failed: {exc}", flush=True)
+        # Fallback action
+        return TriageAction(
+            action_type=ActionType.ASSIGN_PRIORITY,
+            patient_id=current_patient_id,
+            priority_level=PriorityLevel.LESS_URGENT,
+        )
+
+
+# ============================================================================
+# Task Runner
+# ============================================================================
+
+def run_task(task_name: str, client: OpenAI):
+    """
+    Run a single task against the environment.
+
+    Args:
+        task_name (str): Name of the task to run.
+        client (OpenAI): OpenAI client for LLM inference.
+
+    Notes:
+        Always emits [START] and [END] lines, even on exception.
+        Reward values are formatted to exactly 2 decimal places.
+    """
+    env = TriageEnvironment()
+    rewards = []
+    steps_taken = 0
+    score = 0.0
+    success = False
+
+    log_start(task=task_name, env=BENCHMARK, model=MODEL_NAME)
+
+    try:
+        # Reset environment for this task
+        obs = env.reset(task_name=task_name)
+
+        for step in range(1, MAX_STEPS + 1):
+            if obs.done:
+                break
+
+            # Get current patient ID from observation
+            current_patient_id = ""
+            if obs.current_patient:
+                current_patient_id = obs.current_patient.get("patient_id", "")
+
+            # Build observation dict for LLM
+            obs_dict = {
+                "current_patient": obs.current_patient or {},
+                "queue_length": obs.queue_length,
+                "queue_position": obs.queue_position,
+                "missing_fields": obs.missing_fields,
+                "previous_action_feedback": obs.previous_action_feedback or "",
+                "step_number": obs.step_number,
+                "task_name": obs.task_name,
+            }
+
+            # Get action from LLM
+            action = get_llm_action(client, obs_dict, current_patient_id)
+
+            # Step the environment
+            obs = env.step(action)
+
+            reward = obs.reward if obs.reward is not None else 0.0
+            done = obs.done
+            error = None
+
+            reward_val = round(reward, 2)
+            rewards.append(reward_val)
+            steps_taken = step
+
+            # Build action string for logging
+            action_str = action.action_type.value
+            if action.patient_id:
+                action_str += f"({action.patient_id}"
+                if action.priority_level:
+                    action_str += f",{action.priority_level.value}"
+                if action.info_field:
+                    action_str += f",{action.info_field.value}"
+                action_str += ")"
+
+            log_step(
+                step=step,
+                action=action_str,
+                reward=reward_val,
+                done=done,
+                error=error,
+            )
+
+            if done:
+                break
+
+        # Compute final score using the grader
+        state_dict = env.state.model_dump()
+        score = grade_task(task_name, state_dict)
+        score = round(score, 2)
+        success = score >= 0.5
+
+    except Exception as exc:
+        print(f"[DEBUG] Task {task_name} error: {exc}", flush=True)
+        traceback.print_exc(file=sys.stderr)
+        score = 0.0
+        success = False
+
+    finally:
+        log_end(
+            success=success,
+            steps=steps_taken,
+            score=score,
+            rewards=rewards,
+        )
+
+
+# ============================================================================
+# Main Entry Point
+# ============================================================================
+
+if __name__ == "__main__":
+    # Initialize OpenAI client
+    client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)
+
+    # Run all tasks
+    for task in TASKS:
+        run_task(task, client)

models.py

@@ -0,0 +1,234 @@
+"""
+Module: models.py
+Purpose: Define all typed Pydantic models and enums for the TriageFlow environment.
+Part of: Medical Triage Assistant — OpenEnv Round 1
+Author: Team Squirrel
+
+Overview:
+    This module defines the complete type system for the TriageFlow environment.
+    It includes enums for priority levels, action types, and information fields,
+    as well as Pydantic models for patient records, observations, actions, and state.
+    All models inherit from OpenEnv base classes to ensure spec compliance.
+
+Dependencies:
+    - openenv.core.env_server: Action, Observation, State base classes
+    - pydantic: Field for model metadata
+    - enum: Enum support for typed constants
+    - typing: Type annotations
+
+Usage:
+    from models import TriageAction, TriageObservation, TriageState, PriorityLevel
+"""
+
+from enum import Enum
+from typing import Any, Dict, List, Optional
+
+from pydantic import Field
+
+from openenv.core.env_server import Action, Observation, State
+
+
+# ============================================================================
+# Enums
+# ============================================================================
+
+class PriorityLevel(str, Enum):
+    """
+    Patient urgency priority levels following standard triage categories.
+
+    Attributes:
+        IMMEDIATE: Life-threatening, requires immediate intervention.
+        URGENT: Serious condition, needs attention within 15 minutes.
+        LESS_URGENT: Stable condition, can wait 30-60 minutes.
+        NON_URGENT: Minor issue, routine care.
+    """
+    IMMEDIATE = "immediate"
+    URGENT = "urgent"
+    LESS_URGENT = "less_urgent"
+    NON_URGENT = "non_urgent"
+
+
+class ActionType(str, Enum):
+    """
+    Available action types the agent can take during triage.
+
+    Attributes:
+        ASSIGN_PRIORITY: Assign an urgency level to a patient.
+        REQUEST_INFO: Request a missing information field for a patient.
+        ESCALATE: Escalate patient to senior staff with a reason.
+        DEFER: Push patient further back in queue.
+        ADVANCE_QUEUE: Move to the next patient in queue.
+    """
+    ASSIGN_PRIORITY = "assign_priority"
+    REQUEST_INFO = "request_info"
+    ESCALATE = "escalate"
+    DEFER = "defer"
+    ADVANCE_QUEUE = "advance_queue"
+
+
+class InfoField(str, Enum):
+    """
+    Information fields that may be missing from a patient record.
+
+    Attributes:
+        VITALS: Heart rate, blood pressure, SpO2, temperature.
+        HISTORY: Brief medical history summary.
+        ALLERGIES: Known allergies.
+        MEDICATIONS: Current medications.
+        CHIEF_COMPLAINT: Primary stated complaint.
+    """
+    VITALS = "vitals"
+    HISTORY = "history"
+    ALLERGIES = "allergies"
+    MEDICATIONS = "medications"
+    CHIEF_COMPLAINT = "chief_complaint"
+
+
+# ============================================================================
+# Patient Record Model
+# ============================================================================
+
+class PatientRecord(Action):
+    """
+    Represents a patient's intake record with all relevant clinical and administrative data.
+
+    Attributes:
+        patient_id (str): Unique identifier for the patient.
+        age (int): Patient age in years.
+        chief_complaint (str): Primary stated complaint from the patient.
+        symptoms (List[str]): List of reported symptoms.
+        vitals (Optional[dict]): HR, BP, SpO2, temperature — None if missing.
+        history (Optional[str]): Brief medical history summary.
+        medications (Optional[List[str]]): Current medications list.
+        allergies (Optional[List[str]]): Known allergies list.
+        info_complete (bool): Whether all required fields are present.
+
+    Notes:
+        ground_truth_priority is stored separately in tasks.py and is
+        never exposed in the observation visible to agents.
+    """
+    patient_id: str = Field(..., description="Unique patient identifier")
+    age: int = Field(..., description="Patient age in years")
+    chief_complaint: str = Field(..., description="Primary stated complaint")
+    symptoms: List[str] = Field(default_factory=list, description="List of reported symptoms")
+    vitals: Optional[Dict[str, Any]] = Field(None, description="HR, BP, SpO2, temp — None if missing")
+    history: Optional[str] = Field(None, description="Brief history summary")
+    medications: Optional[List[str]] = Field(None, description="Current medications")
+    allergies: Optional[List[str]] = Field(None, description="Known allergies")
+    info_complete: bool = Field(True, description="Whether all required fields are present")
+
+
+# ============================================================================
+# OpenEnv Action Model
+# ============================================================================
+
+class TriageAction(Action):
+    """
+    Action model for the TriageFlow environment.
+
+    The agent uses this action to interact with the environment by
+    assigning priorities, requesting information, escalating, deferring,
+    or advancing through the patient queue.
+
+    Attributes:
+        action_type (ActionType): Which action to take.
+        patient_id (str): Target patient ID for the action.
+        priority_level (Optional[PriorityLevel]): Used with ASSIGN_PRIORITY.
+        info_field (Optional[InfoField]): Used with REQUEST_INFO.
+        escalation_reason (Optional[str]): Used with ESCALATE.
+
+    Notes:
+        Not all fields are required for every action type. The environment
+        validates that the correct fields are present for the given action_type.
+    """
+    action_type: ActionType = Field(..., description="Which action to take")
+    patient_id: str = Field(default="", description="Target patient ID")
+    priority_level: Optional[PriorityLevel] = Field(None, description="Used with ASSIGN_PRIORITY")
+    info_field: Optional[InfoField] = Field(None, description="Used with REQUEST_INFO")
+    escalation_reason: Optional[str] = Field(None, description="Used with ESCALATE")
+
+
+# ============================================================================
+# OpenEnv Observation Model
+# ============================================================================
+
+class TriageObservation(Observation):
+    """
+    Observation model returned to the agent after each step.
+
+    Contains the current patient record (without ground truth), queue state,
+    missing field indicators, and feedback from the previous action.
+
+    Attributes:
+        current_patient (Optional[dict]): Current patient data visible to agent.
+        queue_length (int): Number of remaining patients in queue.
+        queue_position (int): Position of current patient in queue.
+        missing_fields (List[str]): Fields missing for current patient.
+        previous_action_feedback (Optional[str]): System feedback on last action.
+        step_number (int): Current step in the episode.
+        task_name (str): Name of the current task.
+
+    Notes:
+        done and reward are inherited from the Observation base class.
+    """
+    # done: bool and reward: Optional[float] are inherited from Observation
+    current_patient: Optional[Dict[str, Any]] = Field(
+        None, description="Current patient data visible to agent"
+    )
+    queue_length: int = Field(0, description="Remaining patients in queue")
+    queue_position: int = Field(0, description="Position of current patient")
+    missing_fields: List[str] = Field(
+        default_factory=list, description="Fields missing for current patient"
+    )
+    previous_action_feedback: Optional[str] = Field(
+        None, description="System feedback on last action"
+    )
+    step_number: int = Field(0, description="Current step in episode")
+    task_name: str = Field("", description="Name of the current task")
+
+
+# ============================================================================
+# OpenEnv State Model
+# ============================================================================
+
+class TriageState(State):
+    """
+    Full internal state of the environment, used by graders and validators.
+
+    Contains everything needed to deterministically score the episode,
+    including ground truth labels and complete action history.
+
+    Attributes:
+        task_name (str): Which task is active.
+        patients (List[dict]): Full patient queue with ground truth.
+        assignments (Dict[str, str]): patient_id → assigned priority.
+        escalations (Dict[str, str]): patient_id → escalation reason.
+        info_requests (List[dict]): List of info request records.
+        action_history (List[dict]): Complete action log.
+        current_index (int): Current position in queue.
+        max_steps (int): Maximum steps allowed.
+        queue_cleared (bool): Whether all patients have been processed.
+
+    Notes:
+        episode_id and step_count are inherited from the State base class.
+    """
+    # episode_id: Optional[str] and step_count: int are inherited from State
+    task_name: str = Field("", description="Active task name")
+    patients: List[Dict[str, Any]] = Field(
+        default_factory=list, description="Full patient queue with ground truth"
+    )
+    assignments: Dict[str, str] = Field(
+        default_factory=dict, description="patient_id → assigned priority"
+    )
+    escalations: Dict[str, str] = Field(
+        default_factory=dict, description="patient_id → escalation reason"
+    )
+    info_requests: List[Dict[str, Any]] = Field(
+        default_factory=list, description="Info request records"
+    )
+    action_history: List[Dict[str, Any]] = Field(
+        default_factory=list, description="Complete action log"
+    )
+    current_index: int = Field(0, description="Current position in queue")
+    max_steps: int = Field(20, description="Maximum steps allowed")
+    queue_cleared: bool = Field(False, description="Whether all patients processed")

openenv.yaml

@@ -0,0 +1,10 @@
+name: triage_flow
+version: "1.0.0"
+description: "TriageFlow: AI-assisted patient intake triage prioritization environment. Simulates emergency department intake desk where an AI agent must assess urgency, assign priorities, request missing information, and manage a patient queue."
+tasks:
+  - name: basic-triage
+    difficulty: easy
+  - name: incomplete-records-triage
+    difficulty: medium
+  - name: mass-casualty-triage
+    difficulty: hard

pyproject.toml

@@ -0,0 +1,32 @@
+[build-system]
+requires = ["setuptools>=68.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "openenv-triage-flow"
+version = "1.0.0"
+description = "TriageFlow: AI-assisted patient intake triage prioritization environment for OpenEnv"
+readme = "README.md"
+requires-python = ">=3.10"
+license = {text = "MIT"}
+authors = [
+    {name = "Team Squirrel"},
+]
+dependencies = [
+    "openenv-core[core]>=0.2.2",
+    "openai>=1.0.0",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=7.0",
+    "pytest-asyncio>=0.21.0",
+]
+
+[project.scripts]
+server = "triage_flow.server.app:main"
+
+[tool.setuptools]
+include-package-data = true
+packages = ["triage_flow", "triage_flow.server"]
+package-dir = { "triage_flow" = ".", "triage_flow.server" = "server" }

requirements.txt

@@ -0,0 +1,7 @@
+openenv-core>=0.1.0
+fastapi>=0.104.0
+uvicorn>=0.24.0
+pydantic>=2.0.0
+openai>=1.0.0
+httpx>=0.25.0
+websockets>=12.0

scripts/validate-submission.sh

@@ -0,0 +1,75 @@
+#!/usr/bin/env bash
+# Pre-submission validation script for TriageFlow
+# Usage: ./scripts/validate-submission.sh <space_url>
+
+set -e
+
+BOLD="\033[1m"
+GREEN="\033[32m"
+RED="\033[31m"
+NC="\033[0m"
+
+SPACE_URL="${1:-}"
+REPO_DIR="$(dirname "$(dirname "$(readlink -f "$0")")")"
+
+pass() { printf "  ${GREEN}✅ PASS${NC}: %s\n" "$1"; }
+fail() { printf "  ${RED}❌ FAIL${NC}: %s\n" "$1"; }
+
+echo ""
+echo "${BOLD}========================================${NC}"
+echo "${BOLD}  TriageFlow Pre-Submission Validator${NC}"
+echo "${BOLD}========================================${NC}"
+echo ""
+
+# Step 1: Ping HF Space
+if [ -n "$SPACE_URL" ]; then
+    echo "${BOLD}Step 1/3: Pinging HF Space${NC} ($SPACE_URL/reset) ..."
+    HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" -X POST \
+      -H "Content-Type: application/json" -d '{}' \
+      "$SPACE_URL/reset" --max-time 30 2>/dev/null || printf "000")
+    if [ "$HTTP_CODE" = "200" ]; then
+        pass "HF Space is live and responds to /reset"
+    else
+        fail "HF Space /reset returned HTTP $HTTP_CODE (expected 200)"
+        exit 1
+    fi
+else
+    echo "Step 1/3: Skipped (no Space URL provided)"
+fi
+
+# Step 2: Docker build
+echo ""
+echo "${BOLD}Step 2/3: Running docker build${NC} ..."
+if command -v docker &>/dev/null; then
+    if docker build -t triage-flow-test "$REPO_DIR" > /dev/null 2>&1; then
+        pass "Docker build succeeded"
+    else
+        fail "Docker build failed"
+        exit 1
+    fi
+else
+    echo "  Skipped (docker not found)"
+fi
+
+# Step 3: openenv validate
+echo ""
+echo "${BOLD}Step 3/3: Running openenv validate${NC} ..."
+if command -v openenv &>/dev/null; then
+    if (cd "$REPO_DIR" && openenv validate 2>&1); then
+        pass "openenv validate passed"
+    else
+        fail "openenv validate failed"
+        exit 1
+    fi
+else
+    echo "  Skipped (openenv not found)"
+fi
+
+echo ""
+echo "${BOLD}========================================${NC}"
+echo "${GREEN}${BOLD}  All checks passed!${NC}"
+echo "${GREEN}${BOLD}  Your submission is ready.${NC}"
+echo "${BOLD}========================================${NC}"
+echo ""
+
+exit 0

server/Dockerfile

@@ -0,0 +1,65 @@
+ARG BASE_IMAGE=ghcr.io/meta-pytorch/openenv-base:latest
+FROM ${BASE_IMAGE} AS builder
+
+WORKDIR /app
+
+# Ensure git is available
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends git && \
+    rm -rf /var/lib/apt/lists/*
+
+ARG BUILD_MODE=in-repo
+ARG ENV_NAME=triage_flow
+
+# Copy environment code
+COPY . /app/env
+
+WORKDIR /app/env
+
+# Ensure uv is available
+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
+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}
+
+WORKDIR /app
+
+# Copy the virtual environment from builder
+COPY --from=builder /app/env/.venv /app/.venv
+
+# Copy the environment code
+COPY --from=builder /app/env /app/env
+
+# Set PATH to use the virtual environment
+ENV PATH="/app/.venv/bin:$PATH"
+
+# Set PYTHONPATH so imports work correctly
+ENV PYTHONPATH="/app/env:$PYTHONPATH"
+
+# Health check
+HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \
+    CMD curl -f http://localhost:8000/health || exit 1
+
+EXPOSE 8000
+
+# Run the FastAPI server
+CMD ["sh", "-c", "cd /app/env && uvicorn server.app:app --host 0.0.0.0 --port 8000"]

server/__init__.py

@@ -0,0 +1,6 @@
+"""
+Module: __init__.py
+Purpose: Server package initialization.
+Part of: Medical Triage Assistant — OpenEnv Round 1
+Author: Team Squirrel
+"""

server/app.py

@@ -0,0 +1,59 @@
+"""
+Module: app.py
+Purpose: FastAPI server exposing the TriageFlow environment via HTTP/WebSocket.
+Part of: Medical Triage Assistant — OpenEnv Round 1
+Author: Team Squirrel
+
+Overview:
+    Uses OpenEnv's create_app() to automatically create all required
+    endpoints: /reset, /step, /state, /health, /ws, /docs, /web.
+    This is the entry point when running the server in Docker or locally.
+
+Dependencies:
+    - openenv.core.env_server.http_server: create_app
+    - models: TriageAction, TriageObservation
+    - server.triage_flow_environment: TriageEnvironment
+
+Usage:
+    uvicorn server.app:app --host 0.0.0.0 --port 8000
+"""
+
+try:
+    from openenv.core.env_server.http_server import create_app
+except ImportError:
+    from openenv.core.env_server import create_fastapi_app as create_app
+
+try:
+    from ..models import TriageAction, TriageObservation
+    from .triage_flow_environment import TriageEnvironment
+except (ImportError, ModuleNotFoundError):
+    import sys, os
+    sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+    from models import TriageAction, TriageObservation
+    from server.triage_flow_environment import TriageEnvironment
+
+
+# Create the app with web interface
+app = create_app(
+    TriageEnvironment,
+    TriageAction,
+    TriageObservation,
+    env_name="triage_flow",
+    max_concurrent_envs=1,
+)
+
+
+def main(host: str = "0.0.0.0", port: int = 8000):
+    """
+    Entry point for direct execution via uv run or python -m.
+
+    Args:
+        host: Host address to bind to (default: "0.0.0.0")
+        port: Port number to listen on (default: 8000)
+    """
+    import uvicorn
+    uvicorn.run(app, host=host, port=port)
+
+
+if __name__ == "__main__":
+    main()

server/triage_flow_environment.py

@@ -0,0 +1,544 @@
+"""
+Module: environment.py
+Purpose: Core environment logic implementing reset(), step(), and state for TriageFlow.
+Part of: Medical Triage Assistant — OpenEnv Round 1
+Author: Team Squirrel
+
+Overview:
+    Implements the TriageEnvironment class which extends the OpenEnv Environment
+    base class. This is the central simulation engine: it manages the patient queue,
+    validates and applies agent actions, computes rewards, and tracks state for
+    deterministic grading. The environment follows a standard episode lifecycle:
+    reset() → (observe, act, step) loop → done.
+
+Dependencies:
+    - openenv.core.env_server: Environment base class
+    - models: TriageAction, TriageObservation, TriageState
+    - tasks: get_task_config for loading patient data
+    - reward: compute_step_reward for per-step reward calculation
+
+Usage:
+    env = TriageEnvironment()
+    obs = env.reset(task_name="basic-triage")
+    obs = env.step(TriageAction(action_type="assign_priority", patient_id="P001", priority_level="immediate"))
+"""
+
+import uuid
+from typing import Any, Dict, List, Optional
+
+from openenv.core.env_server import Environment
+
+import sys
+import os
+
+# Add parent directory to path so we can import models
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from models import (
+    TriageAction,
+    TriageObservation,
+    TriageState,
+    PriorityLevel,
+    ActionType,
+    InfoField,
+)
+from triage_flow.tasks import get_task_config, TASK_NAMES
+from triage_flow.reward import compute_step_reward, compute_terminal_reward
+
+
+class TriageEnvironment(Environment):
+    """
+    Core triage simulation environment implementing the OpenEnv spec.
+
+    Manages a queue of patients that an agent must triage by assigning
+    urgency priorities, requesting missing information, escalating critical
+    cases, or deferring patients. The environment tracks all state for
+    deterministic grading.
+
+    Attributes:
+        SUPPORTS_CONCURRENT_SESSIONS (bool): Enables multiple simultaneous clients.
+
+    Notes:
+        Ground truth priorities are never exposed in the observation.
+        The grader reads from state() for scoring.
+    """
+
+    SUPPORTS_CONCURRENT_SESSIONS = True
+
+    def __init__(self):
+        """Initialize environment with empty state."""
+        self._state = TriageState()
+        self._patients: List[Dict[str, Any]] = []
+        self._task_config: Dict[str, Any] = {}
+        self._cumulative_reward: float = 0.0
+
+    def reset(self, seed=None, episode_id=None, task_name=None, **kwargs) -> TriageObservation:
+        """
+        Reset the environment for a new episode.
+
+        Loads the specified task's patient queue, clears all state,
+        and returns the initial observation showing the first patient.
+
+        Args:
+            seed: Optional random seed (not used — tasks are deterministic).
+            episode_id: Optional episode identifier.
+            task_name: Which task to load. Defaults to "basic-triage".
+            **kwargs: Additional arguments (may include task_name).
+
+        Returns:
+            TriageObservation: Initial observation with first patient visible.
+
+        Notes:
+            If task_name is not provided, it defaults to "basic-triage".
+            The task_name can also be passed via kwargs.
+        """
+        # Extract task_name from kwargs if not provided directly
+        if task_name is None:
+            task_name = kwargs.get("task_name", "basic-triage")
+
+        # Load task configuration
+        self._task_config = get_task_config(task_name)
+        self._patients = self._task_config["patients"]
+
+        # Reset cumulative reward
+        self._cumulative_reward = 0.0
+
+        # Initialize state
+        self._state = TriageState(
+            episode_id=episode_id or str(uuid.uuid4()),
+            step_count=0,
+            task_name=task_name,
+            patients=[p.copy() for p in self._patients],
+            assignments={},
+            escalations={},
+            info_requests=[],
+            action_history=[],
+            current_index=0,
+            max_steps=self._task_config["max_steps"],
+            queue_cleared=False,
+        )
+
+        # Build initial observation
+        return self._build_observation(
+            reward=None,
+            done=False,
+            feedback=f"Episode started. Task: {task_name}. "
+                     f"{len(self._patients)} patients in queue. "
+                     f"Assess each patient and assign appropriate priority.",
+        )
+
+    def step(self, action: TriageAction, timeout_s=None, **kwargs) -> TriageObservation:
+        """
+        Execute one agent action and return the resulting observation.
+
+        Validates the action, applies it to the internal state, computes
+        the per-step reward, and checks for done conditions.
+
+        Args:
+            action (TriageAction): The action to execute.
+            timeout_s: Optional timeout (not used).
+            **kwargs: Additional arguments.
+
+        Returns:
+            TriageObservation: Updated observation after action.
+
+        Notes:
+            Invalid actions are handled gracefully with a penalty reward
+            and descriptive feedback. The episode always emits an observation.
+        """
+        self._state.step_count += 1
+
+        # Validate action and compute reward
+        is_valid, validation_msg = self._validate_action(action)
+
+        # Get current patient info for reward computation
+        current_patient = self._get_current_patient()
+        ground_truth = current_patient.get("ground_truth_priority") if current_patient else None
+        missing_fields = current_patient.get("missing_fields", []) if current_patient else []
+        info_changes = current_patient.get("info_changes_decision", False) if current_patient else False
+        patient_id = action.patient_id if action.patient_id else (
+            current_patient.get("patient_id", "") if current_patient else ""
+        )
+
+        # Compute reward
+        reward_value, reward_reason = compute_step_reward(
+            action_type=action.action_type.value if isinstance(action.action_type, ActionType) else str(action.action_type),
+            patient_id=patient_id,
+            assigned_priority=action.priority_level.value if action.priority_level else None,
+            ground_truth_priority=ground_truth,
+            info_field=action.info_field.value if action.info_field else None,
+            missing_fields=missing_fields,
+            info_changes_decision=info_changes,
+            already_assigned=patient_id in self._state.assignments,
+            already_escalated=patient_id in self._state.escalations,
+            escalation_reason=action.escalation_reason,
+            action_history=self._state.action_history,
+            is_valid_action=is_valid,
+        )
+
+        # Apply action to state
+        feedback = self._apply_action(action, is_valid, validation_msg)
+
+        # Record action in history
+        self._state.action_history.append({
+            "step": self._state.step_count,
+            "action_type": action.action_type.value if isinstance(action.action_type, ActionType) else str(action.action_type),
+            "patient_id": patient_id,
+            "priority_level": action.priority_level.value if action.priority_level else None,
+            "info_field": action.info_field.value if action.info_field else None,
+            "escalation_reason": action.escalation_reason,
+            "reward": reward_value,
+            "reward_reason": reward_reason,
+            "valid": is_valid,
+        })
+
+        # Accumulate reward
+        self._cumulative_reward += reward_value
+
+        # Check done conditions
+        done = self._check_done()
+
+        # Add terminal bonus if done and queue cleared
+        if done and self._state.queue_cleared:
+            terminal_reward, terminal_reason = compute_terminal_reward(True)
+            reward_value += terminal_reward
+            self._cumulative_reward += terminal_reward
+            feedback += f" {terminal_reason}."
+
+        return self._build_observation(
+            reward=reward_value,
+            done=done,
+            feedback=f"{feedback} | {reward_reason}",
+        )
+
+    @property
+    def state(self) -> TriageState:
+        """
+        Return the full internal state for grading and validation.
+
+        Returns:
+            TriageState: Complete state snapshot including ground truth,
+            action history, and all assignments.
+
+        Notes:
+            This is read by graders to compute the final score.
+            It must be serializable to JSON.
+        """
+        return self._state
+
+    # ========================================================================
+    # Internal Helpers
+    # ========================================================================
+
+    def _get_current_patient(self) -> Optional[Dict[str, Any]]:
+        """
+        Get the patient at the current queue position.
+
+        Returns:
+            Optional[dict]: Current patient data, or None if queue is empty/exhausted.
+        """
+        if 0 <= self._state.current_index < len(self._patients):
+            return self._patients[self._state.current_index]
+        return None
+
+    def _build_observation(
+        self, reward: Optional[float], done: bool, feedback: str
+    ) -> TriageObservation:
+        """
+        Build an observation object from current state.
+
+        Strips ground truth fields from the patient record before
+        including it in the observation.
+
+        Args:
+            reward (Optional[float]): Reward for this step.
+            done (bool): Whether episode is complete.
+            feedback (str): Human-readable feedback message.
+
+        Returns:
+            TriageObservation: Agent-visible observation.
+        """
+        current_patient = self._get_current_patient()
+
+        # Build agent-visible patient dict (strip ground truth)
+        visible_patient = None
+        missing = []
+        if current_patient:
+            visible_patient = {
+                k: v
+                for k, v in current_patient.items()
+                if k not in ("ground_truth_priority", "missing_fields",
+                             "info_changes_decision", "revealed_history",
+                             "revealed_medications")
+            }
+            # Calculate missing fields
+            for field in ["vitals", "history", "medications", "allergies"]:
+                if current_patient.get(field) is None:
+                    missing.append(field)
+
+        # Count remaining unassigned patients
+        unassigned_count = sum(
+            1 for p in self._patients
+            if p["patient_id"] not in self._state.assignments
+        )
+
+        return TriageObservation(
+            done=done,
+            reward=reward,
+            current_patient=visible_patient,
+            queue_length=unassigned_count,
+            queue_position=self._state.current_index + 1,
+            missing_fields=missing,
+            previous_action_feedback=feedback,
+            step_number=self._state.step_count,
+            task_name=self._state.task_name,
+        )
+
+    def _validate_action(self, action: TriageAction) -> tuple:
+        """
+        Validate that an action is legal in the current state.
+
+        Args:
+            action (TriageAction): Action to validate.
+
+        Returns:
+            tuple: (is_valid: bool, message: str)
+        """
+        action_type = action.action_type.value if isinstance(action.action_type, ActionType) else str(action.action_type)
+        current_patient = self._get_current_patient()
+
+        if not current_patient:
+            return False, "No patient at current queue position"
+
+        patient_id = action.patient_id or current_patient.get("patient_id", "")
+
+        # Check if patient_id exists in the queue
+        valid_ids = [p["patient_id"] for p in self._patients]
+        if patient_id and patient_id not in valid_ids:
+            return False, f"Patient {patient_id} not found in queue"
+
+        if action_type == "assign_priority":
+            if not action.priority_level:
+                return False, "ASSIGN_PRIORITY requires a priority_level"
+            if patient_id in self._state.assignments:
+                return False, f"Patient {patient_id} already has priority assigned"
+
+        elif action_type == "request_info":
+            if not action.info_field:
+                return False, "REQUEST_INFO requires an info_field"
+
+        elif action_type == "escalate":
+            if patient_id in self._state.escalations:
+                return False, f"Patient {patient_id} already escalated"
+
+        elif action_type == "defer":
+            if len(self._patients) <= 1:
+                return False, "Cannot defer when only one patient in queue"
+
+        elif action_type == "advance_queue":
+            pass  # Always valid
+
+        else:
+            return False, f"Unknown action type: {action_type}"
+
+        return True, "Action valid"
+
+    def _apply_action(self, action: TriageAction, is_valid: bool, validation_msg: str) -> str:
+        """
+        Apply a validated action to the internal state.
+
+        Args:
+            action (TriageAction): Action to apply.
+            is_valid (bool): Whether the action passed validation.
+            validation_msg (str): Validation feedback message.
+
+        Returns:
+            str: Feedback message describing what happened.
+        """
+        if not is_valid:
+            return f"Action rejected: {validation_msg}"
+
+        action_type = action.action_type.value if isinstance(action.action_type, ActionType) else str(action.action_type)
+        current_patient = self._get_current_patient()
+        patient_id = action.patient_id or (current_patient.get("patient_id", "") if current_patient else "")
+
+        if action_type == "assign_priority":
+            priority = action.priority_level.value if action.priority_level else "unknown"
+            self._state.assignments[patient_id] = priority
+            # Auto-advance to next unassigned patient
+            self._advance_to_next_unassigned()
+            # Check if queue is cleared
+            self._check_queue_cleared()
+            return f"Priority {priority} assigned to patient {patient_id}"
+
+        elif action_type == "request_info":
+            field = action.info_field.value if action.info_field else "unknown"
+            # Record the request
+            self._state.info_requests.append({
+                "patient_id": patient_id,
+                "field": field,
+                "step": self._state.step_count,
+            })
+            # Reveal the information if it was genuinely missing
+            target_patient = self._find_patient(patient_id)
+            if target_patient:
+                revealed_value = None
+                if field == "history" and target_patient.get("revealed_history"):
+                    target_patient["history"] = target_patient["revealed_history"]
+                    revealed_value = target_patient["history"]
+                elif field == "medications" and target_patient.get("revealed_medications"):
+                    target_patient["medications"] = target_patient["revealed_medications"]
+                    revealed_value = str(target_patient["medications"])
+                elif field == "vitals" and target_patient.get("vitals") is None:
+                    # Generate plausible vitals based on ground truth priority
+                    target_patient["vitals"] = self._generate_revealed_vitals(target_patient)
+                    revealed_value = str(target_patient["vitals"])
+                elif field == "allergies" and target_patient.get("allergies") is None:
+                    target_patient["allergies"] = []
+                    revealed_value = "No known allergies"
+
+                # Update info_complete if all fields now present
+                still_missing = any(
+                    target_patient.get(f) is None
+                    for f in ["vitals", "history", "medications", "allergies"]
+                )
+                target_patient["info_complete"] = not still_missing
+
+                if revealed_value:
+                    return f"Info requested: {field} revealed for patient {patient_id}: {revealed_value}"
+                return f"Info requested: {field} for patient {patient_id} (already present)"
+
+            return f"Info requested: {field} for patient {patient_id}"
+
+        elif action_type == "escalate":
+            reason = action.escalation_reason or "No reason provided"
+            self._state.escalations[patient_id] = reason
+            return f"Patient {patient_id} escalated to senior staff. Reason: {reason}"
+
+        elif action_type == "defer":
+            # Move patient to end of queue
+            target_idx = next(
+                (i for i, p in enumerate(self._patients) if p["patient_id"] == patient_id),
+                None
+            )
+            if target_idx is not None:
+                patient = self._patients.pop(target_idx)
+                self._patients.append(patient)
+                # Adjust current index if needed
+                if target_idx <= self._state.current_index:
+                    self._state.current_index = max(0, self._state.current_index - 1)
+            return f"Patient {patient_id} deferred to end of queue"
+
+        elif action_type == "advance_queue":
+            self._state.current_index = min(
+                self._state.current_index + 1, len(self._patients) - 1
+            )
+            next_patient = self._get_current_patient()
+            next_id = next_patient["patient_id"] if next_patient else "none"
+            return f"Advanced queue. Now viewing patient {next_id}"
+
+        return "Action applied"
+
+    def _advance_to_next_unassigned(self):
+        """
+        Move current_index to the next patient that hasn't been assigned yet.
+        Wraps around the queue if necessary.
+        """
+        start = self._state.current_index
+        for i in range(len(self._patients)):
+            idx = (start + i) % len(self._patients)
+            patient = self._patients[idx]
+            if patient["patient_id"] not in self._state.assignments:
+                self._state.current_index = idx
+                return
+        # All assigned — stay at current
+        self._state.current_index = min(start, len(self._patients) - 1)
+
+    def _check_queue_cleared(self):
+        """Check if all patients have been assigned a priority."""
+        all_assigned = all(
+            p["patient_id"] in self._state.assignments for p in self._patients
+        )
+        self._state.queue_cleared = all_assigned
+
+    def _check_done(self) -> bool:
+        """
+        Check if the episode should end.
+
+        Done conditions:
+        1. All patients assigned a priority (queue cleared)
+        2. Maximum steps reached
+
+        Returns:
+            bool: True if episode should end.
+        """
+        # All patients assigned
+        if self._state.queue_cleared:
+            return True
+
+        # Max steps reached
+        if self._state.step_count >= self._state.max_steps:
+            return True
+
+        return False
+
+    def _find_patient(self, patient_id: str) -> Optional[Dict[str, Any]]:
+        """
+        Find a patient in the queue by ID.
+
+        Args:
+            patient_id (str): Patient ID to search for.
+
+        Returns:
+            Optional[dict]: Patient data, or None if not found.
+        """
+        for p in self._patients:
+            if p["patient_id"] == patient_id:
+                return p
+        return None
+
+    def _generate_revealed_vitals(self, patient: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Generate plausible vitals for a patient whose vitals were missing.
+
+        Based on ground truth priority to maintain clinical consistency.
+
+        Args:
+            patient (dict): Patient data including ground_truth_priority.
+
+        Returns:
+            dict: Vitals dictionary with HR, BP, SpO2, temperature.
+
+        Notes:
+            # ASSUMPTION: Revealed vitals are consistent with the ground truth
+            # priority to maintain deterministic grading.
+        """
+        priority = patient.get("ground_truth_priority", "non_urgent")
+        if priority == "immediate":
+            return {
+                "heart_rate": 125,
+                "blood_pressure": "85/55",
+                "spo2": 90,
+                "temperature": 38.5,
+            }
+        elif priority == "urgent":
+            return {
+                "heart_rate": 105,
+                "blood_pressure": "145/95",
+                "spo2": 93,
+                "temperature": 38.8,
+            }
+        elif priority == "less_urgent":
+            return {
+                "heart_rate": 85,
+                "blood_pressure": "130/82",
+                "spo2": 97,
+                "temperature": 37.5,
+            }
+        else:  # non_urgent
+            return {
+                "heart_rate": 75,
+                "blood_pressure": "120/78",
+                "spo2": 99,
+                "temperature": 36.9,
+            }

temp_test_env/test_env/README.md

@@ -0,0 +1,255 @@
+---
+title: Test Env Environment Server
+emoji: 🎸
+colorFrom: purple
+colorTo: yellow
+sdk: docker
+pinned: false
+app_port: 8000
+base_path: /web
+tags:
+  - openenv
+---
+
+# Test Env Environment
+
+A simple test environment that echoes back messages. Perfect for testing the env APIs as well as demonstrating environment usage patterns.
+
+## Quick Start
+
+The simplest way to use the Test Env environment is through the `TestEnv` class:
+
+```python
+from test_env import TestAction, TestEnv
+
+try:
+    # Create environment from Docker image
+    test_envenv = TestEnv.from_docker_image("test_env-env:latest")
+
+    # Reset
+    result = test_envenv.reset()
+    print(f"Reset: {result.observation.echoed_message}")
+
+    # Send multiple messages
+    messages = ["Hello, World!", "Testing echo", "Final message"]
+
+    for msg in messages:
+        result = test_envenv.step(TestAction(message=msg))
+        print(f"Sent: '{msg}'")
+        print(f"  → Echoed: '{result.observation.echoed_message}'")
+        print(f"  → Length: {result.observation.message_length}")
+        print(f"  → Reward: {result.reward}")
+
+finally:
+    # Always clean up
+    test_envenv.close()
+```
+
+That's it! The `TestEnv.from_docker_image()` method handles:
+- Starting the Docker container
+- Waiting for the server to be ready
+- Connecting to the environment
+- Container cleanup when you call `close()`
+
+## Building the Docker Image
+
+Before using the environment, you need to build the Docker image:
+
+```bash
+# From project root
+docker build -t test_env-env:latest -f server/Dockerfile .
+```
+
+## Deploying to Hugging Face Spaces
+
+You can easily deploy your OpenEnv environment to Hugging Face Spaces using the `openenv push` command:
+
+```bash
+# From the environment directory (where openenv.yaml is located)
+openenv push
+
+# Or specify options
+openenv push --namespace my-org --private
+```
+
+The `openenv push` command will:
+1. Validate that the directory is an OpenEnv environment (checks for `openenv.yaml`)
+2. Prepare a custom build for Hugging Face Docker space (enables web interface)
+3. Upload to Hugging Face (ensuring you're logged in)
+
+### Prerequisites
+
+- Authenticate with Hugging Face: The command will prompt for login if not already authenticated
+
+### Options
+
+- `--directory`, `-d`: Directory containing the OpenEnv environment (defaults to current directory)
+- `--repo-id`, `-r`: Repository ID in format 'username/repo-name' (defaults to 'username/env-name' from openenv.yaml)
+- `--base-image`, `-b`: Base Docker image to use (overrides Dockerfile FROM)
+- `--private`: Deploy the space as private (default: public)
+
+### Examples
+
+```bash
+# Push to your personal namespace (defaults to username/env-name from openenv.yaml)
+openenv push
+
+# Push to a specific repository
+openenv push --repo-id my-org/my-env
+
+# Push with a custom base image
+openenv push --base-image ghcr.io/meta-pytorch/openenv-base:latest
+
+# Push as a private space
+openenv push --private
+
+# Combine options
+openenv push --repo-id my-org/my-env --base-image custom-base:latest --private
+```
+
+After deployment, your space will be available at:
+`https://huggingface.co/spaces/<repo-id>`
+
+The deployed space includes:
+- **Web Interface** at `/web` - Interactive UI for exploring the environment
+- **API Documentation** at `/docs` - Full OpenAPI/Swagger interface
+- **Health Check** at `/health` - Container health monitoring
+- **WebSocket** at `/ws` - Persistent session endpoint for low-latency interactions
+
+## Environment Details
+
+### Action
+**TestAction**: Contains a single field
+- `message` (str) - The message to echo back
+
+### Observation
+**TestObservation**: Contains the echo response and metadata
+- `echoed_message` (str) - The message echoed back
+- `message_length` (int) - Length of the message
+- `reward` (float) - Reward based on message length (length × 0.1)
+- `done` (bool) - Always False for echo environment
+- `metadata` (dict) - Additional info like step count
+
+### Reward
+The reward is calculated as: `message_length × 0.1`
+- "Hi" → reward: 0.2
+- "Hello, World!" → reward: 1.3
+- Empty message → reward: 0.0
+
+## Advanced Usage
+
+### Connecting to an Existing Server
+
+If you already have a Test Env environment server running, you can connect directly:
+
+```python
+from test_env import TestEnv
+
+# Connect to existing server
+test_envenv = TestEnv(base_url="<ENV_HTTP_URL_HERE>")
+
+# Use as normal
+result = test_envenv.reset()
+result = test_envenv.step(TestAction(message="Hello!"))
+```
+
+Note: When connecting to an existing server, `test_envenv.close()` will NOT stop the server.
+
+### Using the Context Manager
+
+The client supports context manager usage for automatic connection management:
+
+```python
+from test_env import TestAction, TestEnv
+
+# Connect with context manager (auto-connects and closes)
+with TestEnv(base_url="http://localhost:8000") as env:
+    result = env.reset()
+    print(f"Reset: {result.observation.echoed_message}")
+    # Multiple steps with low latency
+    for msg in ["Hello", "World", "!"]:
+        result = env.step(TestAction(message=msg))
+        print(f"Echoed: {result.observation.echoed_message}")
+```
+
+The client uses WebSocket connections for:
+- **Lower latency**: No HTTP connection overhead per request
+- **Persistent session**: Server maintains your environment state
+- **Efficient for episodes**: Better for many sequential steps
+
+### Concurrent WebSocket Sessions
+
+The server supports multiple concurrent WebSocket connections. To enable this,
+modify `server/app.py` to use factory mode:
+
+```python
+# In server/app.py - use factory mode for concurrent sessions
+app = create_app(
+    TestEnvironment,  # Pass class, not instance
+    TestAction,
+    TestObservation,
+    max_concurrent_envs=4,  # Allow 4 concurrent sessions
+)
+```
+
+Then multiple clients can connect simultaneously:
+
+```python
+from test_env import TestAction, TestEnv
+from concurrent.futures import ThreadPoolExecutor
+
+def run_episode(client_id: int):
+    with TestEnv(base_url="http://localhost:8000") as env:
+        result = env.reset()
+        for i in range(10):
+            result = env.step(TestAction(message=f"Client {client_id}, step {i}"))
+        return client_id, result.observation.message_length
+
+# Run 4 episodes concurrently
+with ThreadPoolExecutor(max_workers=4) as executor:
+    results = list(executor.map(run_episode, range(4)))
+```
+
+## Development & Testing
+
+### Direct Environment Testing
+
+Test the environment logic directly without starting the HTTP server:
+
+```bash
+# From the server directory
+python3 server/test_env_environment.py
+```
+
+This verifies that:
+- Environment resets correctly
+- Step executes actions properly
+- State tracking works
+- Rewards are calculated correctly
+
+### Running Locally
+
+Run the server locally for development:
+
+```bash
+uvicorn server.app:app --reload
+```
+
+## Project Structure
+
+```
+test_env/
+├── .dockerignore         # Docker build exclusions
+├── __init__.py            # Module exports
+├── README.md              # This file
+├── openenv.yaml           # OpenEnv manifest
+├── pyproject.toml         # Project metadata and dependencies
+├── uv.lock                # Locked dependencies (generated)
+├── client.py              # TestEnv client
+├── models.py              # Action and Observation models
+└── server/
+    ├── __init__.py        # Server module exports
+    ├── test_env_environment.py  # Core environment logic
+    ├── app.py             # FastAPI application (HTTP + WebSocket endpoints)
+    └── Dockerfile         # Container image definition
+```

temp_test_env/test_env/__init__.py

@@ -0,0 +1,16 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Test Env Environment."""
+
+from .client import TestEnv
+from .models import TestAction, TestObservation
+
+__all__ = [
+    "TestAction",
+    "TestObservation",
+    "TestEnv",
+]

temp_test_env/test_env/client.py

@@ -0,0 +1,99 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Test Env Environment Client."""
+
+from typing import Dict
+
+from openenv.core import EnvClient
+from openenv.core.client_types import StepResult
+from openenv.core.env_server.types import State
+
+from .models import TestAction, TestObservation
+
+
+class TestEnv(
+    EnvClient[TestAction, TestObservation, State]
+):
+    """
+    Client for the Test Env Environment.
+
+    This client maintains a persistent WebSocket connection to the environment server,
+    enabling efficient multi-step interactions with lower latency.
+    Each client instance has its own dedicated environment session on the server.
+
+    Example:
+        >>> # Connect to a running server
+        >>> with TestEnv(base_url="http://localhost:8000") as client:
+        ...     result = client.reset()
+        ...     print(result.observation.echoed_message)
+        ...
+        ...     result = client.step(TestAction(message="Hello!"))
+        ...     print(result.observation.echoed_message)
+
+    Example with Docker:
+        >>> # Automatically start container and connect
+        >>> client = TestEnv.from_docker_image("test_env-env:latest")
+        >>> try:
+        ...     result = client.reset()
+        ...     result = client.step(TestAction(message="Test"))
+        ... finally:
+        ...     client.close()
+    """
+
+    def _step_payload(self, action: TestAction) -> Dict:
+        """
+        Convert TestAction to JSON payload for step message.
+
+        Args:
+            action: TestAction instance
+
+        Returns:
+            Dictionary representation suitable for JSON encoding
+        """
+        return {
+            "message": action.message,
+        }
+
+    def _parse_result(self, payload: Dict) -> StepResult[TestObservation]:
+        """
+        Parse server response into StepResult[TestObservation].
+
+        Args:
+            payload: JSON response data from server
+
+        Returns:
+            StepResult with TestObservation
+        """
+        obs_data = payload.get("observation", {})
+        observation = TestObservation(
+            echoed_message=obs_data.get("echoed_message", ""),
+            message_length=obs_data.get("message_length", 0),
+            done=payload.get("done", False),
+            reward=payload.get("reward"),
+            metadata=obs_data.get("metadata", {}),
+        )
+
+        return StepResult(
+            observation=observation,
+            reward=payload.get("reward"),
+            done=payload.get("done", False),
+        )
+
+    def _parse_state(self, payload: Dict) -> State:
+        """
+        Parse server response into State object.
+
+        Args:
+            payload: JSON response from state request
+
+        Returns:
+            State object with episode_id and step_count
+        """
+        return State(
+            episode_id=payload.get("episode_id"),
+            step_count=payload.get("step_count", 0),
+        )

temp_test_env/test_env/models.py

@@ -0,0 +1,27 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Data models for the Test Env Environment.
+
+The test_env environment is a simple test environment that echoes back messages.
+"""
+
+from openenv.core.env_server.types import Action, Observation
+from pydantic import Field
+
+
+class TestAction(Action):
+    """Action for the Test Env environment - just a message to echo."""
+
+    message: str = Field(..., description="Message to echo back")
+
+
+class TestObservation(Observation):
+    """Observation from the Test Env environment - the echoed message."""
+
+    echoed_message: str = Field(default="", description="The echoed message")
+    message_length: int = Field(default=0, description="Length of the echoed message")

temp_test_env/test_env/openenv.yaml

@@ -0,0 +1,7 @@
+spec_version: 1
+name: test_env
+type: space
+runtime: fastapi
+app: server.app:app
+port: 8000
+

temp_test_env/test_env/pyproject.toml

@@ -0,0 +1,45 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+[build-system]
+requires = ["setuptools>=45", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "openenv-test_env"
+version = "0.1.0"
+description = "Test Env environment for OpenEnv"
+requires-python = ">=3.10"
+dependencies = [
+    # Core OpenEnv runtime (provides FastAPI server + HTTP client types)
+    # install from github
+    # "openenv-core[core] @ git+https://github.com/meta-pytorch/OpenEnv.git",
+    "openenv-core[core]>=0.2.2",
+    # Environment-specific dependencies
+    # Add all dependencies needed for your environment here
+    # Examples:
+    # "numpy>=1.19.0",
+    # "torch>=2.0.0",
+    # "gymnasium>=0.29.0",
+    # "openspiel>=1.0.0",
+    # "smolagents>=1.22.0,<2",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=8.0.0",
+    "pytest-cov>=4.0.0",
+]
+
+[project.scripts]
+# Server entry point - enables running via: uv run --project . server
+# or: python -m test_env.server.app
+server = "test_env.server.app:main"
+
+[tool.setuptools]
+include-package-data = true
+packages = ["test_env", "test_env.server"]
+package-dir = { "test_env" = ".", "test_env.server" = "server" }

temp_test_env/test_env/server/Dockerfile

@@ -0,0 +1,80 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# 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.
+
+ARG BASE_IMAGE=ghcr.io/meta-pytorch/openenv-base:latest
+FROM ${BASE_IMAGE} AS builder
+
+WORKDIR /app
+
+# Ensure git is available (required for installing dependencies from VCS)
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends git && \
+    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=test_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}
+
+WORKDIR /app
+
+# Copy the virtual environment from builder
+COPY --from=builder /app/env/.venv /app/.venv
+
+# Copy the environment code
+COPY --from=builder /app/env /app/env
+
+# Set PATH to use the virtual environment
+ENV PATH="/app/.venv/bin:$PATH"
+
+# Set PYTHONPATH so imports work correctly
+ENV PYTHONPATH="/app/env:$PYTHONPATH"
+
+# Health check
+HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --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
+CMD ["sh", "-c", "cd /app/env && uvicorn server.app:app --host 0.0.0.0 --port 8000"]

temp_test_env/test_env/server/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Test Env environment server components."""
+
+from .test_env_environment import TestEnvironment
+
+__all__ = ["TestEnvironment"]

temp_test_env/test_env/server/app.py

@@ -0,0 +1,84 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+FastAPI application for the Test Env Environment.
+
+This module creates an HTTP server that exposes the TestEnvironment
+over HTTP and WebSocket endpoints, compatible with EnvClient.
+
+Endpoints:
+    - POST /reset: Reset the environment
+    - POST /step: Execute an action
+    - GET /state: Get current environment state
+    - GET /schema: Get action/observation schemas
+    - WS /ws: WebSocket endpoint for persistent sessions
+
+Usage:
+    # Development (with auto-reload):
+    uvicorn server.app:app --reload --host 0.0.0.0 --port 8000
+
+    # Production:
+    uvicorn server.app:app --host 0.0.0.0 --port 8000 --workers 4
+
+    # Or run directly:
+    python -m server.app
+"""
+
+try:
+    from openenv.core.env_server.http_server import create_app
+except Exception as e:  # pragma: no cover
+    raise ImportError(
+        "openenv is required for the web interface. Install dependencies with '\n    uv sync\n'"
+    ) from e
+
+try:
+    from ..models import TestAction, TestObservation
+    from .test_env_environment import TestEnvironment
+except ModuleNotFoundError:
+    from models import TestAction, TestObservation
+    from server.test_env_environment import TestEnvironment
+
+
+# Create the app with web interface and README integration
+app = create_app(
+    TestEnvironment,
+    TestAction,
+    TestObservation,
+    env_name="test_env",
+    max_concurrent_envs=1,  # increase this number to allow more concurrent WebSocket sessions
+)
+
+
+def main(host: str = "0.0.0.0", port: int = 8000):
+    """
+    Entry point for direct execution via uv run or python -m.
+
+    This function enables running the server without Docker:
+        uv run --project . server
+        uv run --project . server --port 8001
+        python -m test_env.server.app
+
+    Args:
+        host: Host address to bind to (default: "0.0.0.0")
+        port: Port number to listen on (default: 8000)
+
+    For production deployments, consider using uvicorn directly with
+    multiple workers:
+        uvicorn test_env.server.app:app --workers 4
+    """
+    import uvicorn
+
+    uvicorn.run(app, host=host, port=port)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--port", type=int, default=8000)
+    args = parser.parse_args()
+    main(port=args.port)

temp_test_env/test_env/server/requirements.txt

@@ -0,0 +1,6 @@
+openenv[core]>=0.2.0
+fastapi>=0.115.0
+uvicorn>=0.24.0
+
+
+

temp_test_env/test_env/server/test_env_environment.py

@@ -0,0 +1,104 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Test Env Environment Implementation.
+
+A simple test environment that echoes back messages sent to it.
+Perfect for testing HTTP server infrastructure.
+"""
+
+from uuid import uuid4
+
+from openenv.core.env_server.interfaces import Environment
+from openenv.core.env_server.types import State
+
+try:
+    from ..models import TestAction, TestObservation
+except ImportError:
+    from models import TestAction, TestObservation
+
+
+class TestEnvironment(Environment):
+    """
+    A simple echo environment that echoes back messages.
+
+    This environment is designed for testing the HTTP server infrastructure.
+    It maintains minimal state and simply echoes back whatever message it receives.
+
+    Example:
+        >>> env = TestEnvironment()
+        >>> obs = env.reset()
+        >>> print(obs.echoed_message)  # "Test Env environment ready!"
+        >>>
+        >>> obs = env.step(TestAction(message="Hello"))
+        >>> print(obs.echoed_message)  # "Hello"
+        >>> print(obs.message_length)  # 5
+    """
+
+    # Enable concurrent WebSocket sessions.
+    # Set to True if your environment isolates state between instances.
+    # When True, multiple WebSocket clients can connect simultaneously, each
+    # getting their own environment instance (when using factory mode in app.py).
+    SUPPORTS_CONCURRENT_SESSIONS: bool = True
+
+    def __init__(self):
+        """Initialize the test_env environment."""
+        self._state = State(episode_id=str(uuid4()), step_count=0)
+        self._reset_count = 0
+
+    def reset(self) -> TestObservation:
+        """
+        Reset the environment.
+
+        Returns:
+            TestObservation with a ready message
+        """
+        self._state = State(episode_id=str(uuid4()), step_count=0)
+        self._reset_count += 1
+
+        return TestObservation(
+            echoed_message="Test Env environment ready!",
+            message_length=0,
+            done=False,
+            reward=0.0,
+        )
+
+    def step(self, action: TestAction) -> TestObservation:  # type: ignore[override]
+        """
+        Execute a step in the environment by echoing the message.
+
+        Args:
+            action: TestAction containing the message to echo
+
+        Returns:
+            TestObservation with the echoed message and its length
+        """
+        self._state.step_count += 1
+
+        message = action.message
+        length = len(message)
+
+        # Simple reward: longer messages get higher rewards
+        reward = length * 0.1
+
+        return TestObservation(
+            echoed_message=message,
+            message_length=length,
+            done=False,
+            reward=reward,
+            metadata={"original_message": message, "step": self._state.step_count},
+        )
+
+    @property
+    def state(self) -> State:
+        """
+        Get the current environment state.
+
+        Returns:
+            Current State with episode_id and step_count
+        """
+        return self._state

tests/run_integration_test.py

@@ -0,0 +1,129 @@
+"""Quick integration test for the TriageFlow environment."""
+import sys
+sys.path.insert(0, '.')
+
+from triage_flow.environment import TriageEnvironment
+from triage_flow.graders import grade_task
+from models import TriageAction, ActionType, PriorityLevel, InfoField
+
+def test_basic_triage():
+    env = TriageEnvironment()
+    obs = env.reset(task_name="basic-triage")
+    pid = obs.current_patient["patient_id"] if obs.current_patient else "none"
+    print(f"Reset OK - patients: {obs.queue_length}, current: {pid}")
+
+    # Correct assignments for all 3 patients
+    for pid, p in [("P001", PriorityLevel.IMMEDIATE), ("P002", PriorityLevel.NON_URGENT), ("P003", PriorityLevel.URGENT)]:
+        action = TriageAction(action_type=ActionType.ASSIGN_PRIORITY, patient_id=pid, priority_level=p)
+        obs = env.step(action)
+        print(f"  Assigned {pid}={p.value} -> reward={obs.reward}, done={obs.done}")
+
+    state = env.state
+    print(f"State: assignments={state.assignments}, cleared={state.queue_cleared}")
+    score = grade_task("basic-triage", state.model_dump())
+    print(f"Grader score: {score}")
+    assert score == 1.0, f"Expected 1.0, got {score}"
+    print("BASIC TRIAGE: PASSED\n")
+
+def test_wrong_assignments():
+    env = TriageEnvironment()
+    env.reset(task_name="basic-triage")
+    for pid, p in [("P001", PriorityLevel.NON_URGENT), ("P002", PriorityLevel.IMMEDIATE), ("P003", PriorityLevel.NON_URGENT)]:
+        env.step(TriageAction(action_type=ActionType.ASSIGN_PRIORITY, patient_id=pid, priority_level=p))
+    score = grade_task("basic-triage", env.state.model_dump())
+    print(f"All wrong score: {score}")
+    assert score == 0.0, f"Expected 0.0, got {score}"
+    print("WRONG ASSIGNMENTS: PASSED\n")
+
+def test_incomplete_records():
+    env = TriageEnvironment()
+    obs = env.reset(task_name="incomplete-records-triage")
+    pid = obs.current_patient["patient_id"] if obs.current_patient else "none"
+    print(f"Incomplete records - patients: {obs.queue_length}, current: {pid}")
+    print(f"  Missing fields: {obs.missing_fields}")
+
+    # Request info for P101 (vitals missing)
+    action = TriageAction(action_type=ActionType.REQUEST_INFO, patient_id="P101", info_field=InfoField.VITALS)
+    obs = env.step(action)
+    print(f"  Info request -> reward={obs.reward}")
+
+    # Assign correct priorities
+    for pid, p in [("P101", PriorityLevel.IMMEDIATE), ("P102", PriorityLevel.LESS_URGENT),
+                   ("P103", PriorityLevel.URGENT), ("P104", PriorityLevel.NON_URGENT)]:
+        obs = env.step(TriageAction(action_type=ActionType.ASSIGN_PRIORITY, patient_id=pid, priority_level=p))
+
+    score = grade_task("incomplete-records-triage", env.state.model_dump())
+    print(f"Incomplete records score: {score}")
+    assert 0.5 <= score <= 1.0, f"Expected good score, got {score}"
+    print("INCOMPLETE RECORDS: PASSED\n")
+
+def test_mass_casualty():
+    env = TriageEnvironment()
+    obs = env.reset(task_name="mass-casualty-triage")
+    print(f"Mass casualty - patients: {obs.queue_length}")
+
+    # Assign all 8 patients with correct priorities
+    correct = [
+        ("P201", PriorityLevel.IMMEDIATE), ("P202", PriorityLevel.IMMEDIATE),
+        ("P203", PriorityLevel.LESS_URGENT), ("P204", PriorityLevel.LESS_URGENT),
+        ("P205", PriorityLevel.URGENT), ("P206", PriorityLevel.NON_URGENT),
+        ("P207", PriorityLevel.URGENT), ("P208", PriorityLevel.NON_URGENT),
+    ]
+    for pid, p in correct:
+        obs = env.step(TriageAction(action_type=ActionType.ASSIGN_PRIORITY, patient_id=pid, priority_level=p))
+
+    score = grade_task("mass-casualty-triage", env.state.model_dump())
+    print(f"Mass casualty perfect score: {score}")
+    assert score >= 0.8, f"Expected high score, got {score}"
+    print("MASS CASUALTY: PASSED\n")
+
+def test_grader_variance():
+    """Different behaviors must produce different scores."""
+    scores = set()
+
+    # Perfect
+    env = TriageEnvironment()
+    env.reset(task_name="basic-triage")
+    for pid, p in [("P001", PriorityLevel.IMMEDIATE), ("P002", PriorityLevel.NON_URGENT), ("P003", PriorityLevel.URGENT)]:
+        env.step(TriageAction(action_type=ActionType.ASSIGN_PRIORITY, patient_id=pid, priority_level=p))
+    scores.add(grade_task("basic-triage", env.state.model_dump()))
+
+    # All wrong
+    env = TriageEnvironment()
+    env.reset(task_name="basic-triage")
+    for pid, p in [("P001", PriorityLevel.NON_URGENT), ("P002", PriorityLevel.IMMEDIATE), ("P003", PriorityLevel.NON_URGENT)]:
+        env.step(TriageAction(action_type=ActionType.ASSIGN_PRIORITY, patient_id=pid, priority_level=p))
+    scores.add(grade_task("basic-triage", env.state.model_dump()))
+
+    # Partial
+    env = TriageEnvironment()
+    env.reset(task_name="basic-triage")
+    for pid, p in [("P001", PriorityLevel.IMMEDIATE), ("P002", PriorityLevel.IMMEDIATE), ("P003", PriorityLevel.IMMEDIATE)]:
+        env.step(TriageAction(action_type=ActionType.ASSIGN_PRIORITY, patient_id=pid, priority_level=p))
+    scores.add(grade_task("basic-triage", env.state.model_dump()))
+
+    # No actions
+    env = TriageEnvironment()
+    env.reset(task_name="basic-triage")
+    scores.add(grade_task("basic-triage", env.state.model_dump()))
+
+    print(f"Grader variance - unique scores: {scores}")
+    assert len(scores) >= 3, f"Expected 3+ different scores, got {scores}"
+    print("GRADER VARIANCE: PASSED\n")
+
+def test_server_app_import():
+    """Test that the server app can be imported."""
+    from server.app import app
+    print(f"Server app imported: {app.title}")
+    print("SERVER IMPORT: PASSED\n")
+
+if __name__ == "__main__":
+    test_basic_triage()
+    test_wrong_assignments()
+    test_incomplete_records()
+    test_mass_casualty()
+    test_grader_variance()
+    test_server_app_import()
+    print("=" * 60)
+    print("ALL INTEGRATION TESTS PASSED!")
+    print("=" * 60)

tests/test_endpoints.py

@@ -0,0 +1,35 @@
+"""Test HTTP endpoints against running server."""
+import httpx
+
+BASE = "http://127.0.0.1:8000"
+
+# Health
+r = httpx.get(f"{BASE}/health")
+print(f"Health: {r.status_code} {r.text}")
+
+# Reset with empty body (required for pre-submission validator)
+r = httpx.post(f"{BASE}/reset", json={})
+print(f"Reset (empty): {r.status_code}")
+
+# Reset with task_name
+r = httpx.post(f"{BASE}/reset", json={"task_name": "basic-triage"})
+data = r.json()
+obs = data.get("observation", data)
+pid = obs.get("current_patient", {}).get("patient_id", "?")
+print(f"Reset (basic-triage): {r.status_code}, patient={pid}")
+
+# Step
+r = httpx.post(f"{BASE}/step", json={
+    "action_type": "assign_priority",
+    "patient_id": "P001",
+    "priority_level": "immediate"
+})
+data = r.json()
+print(f"Step: {r.status_code}, reward={data.get('reward', '?')}, done={data.get('done', '?')}")
+
+# State
+r = httpx.get(f"{BASE}/state")
+state = r.json()
+print(f"State: {r.status_code}, assignments={state.get('assignments', {})}")
+
+print("\nALL ENDPOINT TESTS PASSED!")

tests/test_environment.py

@@ -0,0 +1,166 @@
+"""
+Module: test_environment.py
+Purpose: Test core environment behavior — reset, step, state.
+Part of: Medical Triage Assistant — OpenEnv Round 1
+Author: Team Squirrel
+"""
+
+import sys
+import os
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import pytest
+from triage_flow.environment import TriageEnvironment
+from models import TriageAction, ActionType, PriorityLevel, InfoField
+
+
+class TestReset:
+    """Test reset() behavior."""
+
+    def test_reset_returns_observation(self):
+        """reset() should return a TriageObservation."""
+        env = TriageEnvironment()
+        obs = env.reset(task_name="basic-triage")
+        assert obs.done is False
+        assert obs.current_patient is not None
+        assert obs.queue_length > 0
+
+    def test_reset_basic_triage_has_3_patients(self):
+        """Basic triage should have 3 patients."""
+        env = TriageEnvironment()
+        obs = env.reset(task_name="basic-triage")
+        assert obs.queue_length == 3
+
+    def test_reset_clears_state(self):
+        """reset() should clear all previous state."""
+        env = TriageEnvironment()
+        obs = env.reset(task_name="basic-triage")
+
+        # Take an action
+        action = TriageAction(
+            action_type=ActionType.ASSIGN_PRIORITY,
+            patient_id="P001",
+            priority_level=PriorityLevel.IMMEDIATE
+        )
+        env.step(action)
+
+        # Reset should clear
+        obs = env.reset(task_name="basic-triage")
+        assert env.state.step_count == 0
+        assert env.state.assignments == {}
+
+    def test_reset_all_tasks(self):
+        """All 3 tasks should reset successfully."""
+        env = TriageEnvironment()
+        for task in ["basic-triage", "incomplete-records-triage", "mass-casualty-triage"]:
+            obs = env.reset(task_name=task)
+            assert obs.done is False
+            assert obs.current_patient is not None
+
+
+class TestStep:
+    """Test step() behavior."""
+
+    def test_step_valid_action(self):
+        """step() with valid action should return updated observation."""
+        env = TriageEnvironment()
+        env.reset(task_name="basic-triage")
+
+        action = TriageAction(
+            action_type=ActionType.ASSIGN_PRIORITY,
+            patient_id="P001",
+            priority_level=PriorityLevel.IMMEDIATE,
+        )
+        obs = env.step(action)
+        assert obs.reward is not None
+        assert obs.step_number == 1
+
+    def test_step_correct_assignment_positive_reward(self):
+        """Correct priority assignment should give positive reward."""
+        env = TriageEnvironment()
+        env.reset(task_name="basic-triage")
+
+        # P001 ground truth is IMMEDIATE
+        action = TriageAction(
+            action_type=ActionType.ASSIGN_PRIORITY,
+            patient_id="P001",
+            priority_level=PriorityLevel.IMMEDIATE,
+        )
+        obs = env.step(action)
+        assert obs.reward > 0
+
+    def test_step_incorrect_assignment_negative_reward(self):
+        """Incorrect priority assignment should give negative or zero reward."""
+        env = TriageEnvironment()
+        env.reset(task_name="basic-triage")
+
+        # P001 ground truth is IMMEDIATE, assigning NON_URGENT is wrong
+        action = TriageAction(
+            action_type=ActionType.ASSIGN_PRIORITY,
+            patient_id="P001",
+            priority_level=PriorityLevel.NON_URGENT,
+        )
+        obs = env.step(action)
+        assert obs.reward < 0
+
+    def test_step_counts_increment(self):
+        """Step count should increment with each step."""
+        env = TriageEnvironment()
+        env.reset(task_name="basic-triage")
+
+        action = TriageAction(action_type=ActionType.ADVANCE_QUEUE)
+        obs = env.step(action)
+        assert obs.step_number == 1
+
+        obs = env.step(action)
+        assert obs.step_number == 2
+
+    def test_done_when_all_assigned(self):
+        """Episode should end when all patients are assigned."""
+        env = TriageEnvironment()
+        env.reset(task_name="basic-triage")
+
+        # Assign all 3 patients
+        for pid, priority in [
+            ("P001", PriorityLevel.IMMEDIATE),
+            ("P002", PriorityLevel.NON_URGENT),
+            ("P003", PriorityLevel.URGENT),
+        ]:
+            action = TriageAction(
+                action_type=ActionType.ASSIGN_PRIORITY,
+                patient_id=pid,
+                priority_level=priority,
+            )
+            obs = env.step(action)
+
+        assert obs.done is True
+
+
+class TestState:
+    """Test state() behavior."""
+
+    def test_state_serializable(self):
+        """state() should return serializable data."""
+        env = TriageEnvironment()
+        env.reset(task_name="basic-triage")
+        state = env.state
+        d = state.model_dump()
+        assert isinstance(d, dict)
+        assert "patients" in d
+        assert "assignments" in d
+
+    def test_state_tracks_assignments(self):
+        """state() should reflect assignments made."""
+        env = TriageEnvironment()
+        env.reset(task_name="basic-triage")
+
+        action = TriageAction(
+            action_type=ActionType.ASSIGN_PRIORITY,
+            patient_id="P001",
+            priority_level=PriorityLevel.IMMEDIATE,
+        )
+        env.step(action)
+
+        state = env.state
+        assert "P001" in state.assignments
+        assert state.assignments["P001"] == "immediate"

tests/test_graders.py

@@ -0,0 +1,146 @@
+"""
+Module: test_graders.py
+Purpose: Test deterministic graders produce varying scores for different inputs.
+Part of: Medical Triage Assistant — OpenEnv Round 1
+Author: Team Squirrel
+"""
+
+import sys
+import os
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import pytest
+from triage_flow.graders import grade_task
+from triage_flow.environment import TriageEnvironment
+from models import TriageAction, ActionType, PriorityLevel
+
+
+class TestGraderBasicTriage:
+    """Test grader for basic-triage task."""
+
+    def test_perfect_score(self):
+        """All correct assignments should score 1.0."""
+        env = TriageEnvironment()
+        env.reset(task_name="basic-triage")
+
+        for pid, priority in [
+            ("P001", PriorityLevel.IMMEDIATE),
+            ("P002", PriorityLevel.NON_URGENT),
+            ("P003", PriorityLevel.URGENT),
+        ]:
+            env.step(TriageAction(
+                action_type=ActionType.ASSIGN_PRIORITY,
+                patient_id=pid,
+                priority_level=priority,
+            ))
+
+        state = env.state.model_dump()
+        score = grade_task("basic-triage", state)
+        assert score == 1.0
+
+    def test_zero_score(self):
+        """All wrong assignments should score 0.0."""
+        env = TriageEnvironment()
+        env.reset(task_name="basic-triage")
+
+        # Deliberately wrong assignments
+        for pid, priority in [
+            ("P001", PriorityLevel.NON_URGENT),   # Should be IMMEDIATE
+            ("P002", PriorityLevel.IMMEDIATE),     # Should be NON_URGENT
+            ("P003", PriorityLevel.NON_URGENT),    # Should be URGENT
+        ]:
+            env.step(TriageAction(
+                action_type=ActionType.ASSIGN_PRIORITY,
+                patient_id=pid,
+                priority_level=priority,
+            ))
+
+        state = env.state.model_dump()
+        score = grade_task("basic-triage", state)
+        assert score == 0.0
+
+    def test_partial_score(self):
+        """One correct out of three should score ~0.33."""
+        env = TriageEnvironment()
+        env.reset(task_name="basic-triage")
+
+        for pid, priority in [
+            ("P001", PriorityLevel.IMMEDIATE),     # Correct
+            ("P002", PriorityLevel.IMMEDIATE),     # Wrong
+            ("P003", PriorityLevel.NON_URGENT),    # Wrong
+        ]:
+            env.step(TriageAction(
+                action_type=ActionType.ASSIGN_PRIORITY,
+                patient_id=pid,
+                priority_level=priority,
+            ))
+
+        state = env.state.model_dump()
+        score = grade_task("basic-triage", state)
+        assert 0.3 <= score <= 0.35  # ~0.33
+
+    def test_no_assignments_scores_zero(self):
+        """No assignments should score 0.0."""
+        env = TriageEnvironment()
+        env.reset(task_name="basic-triage")
+        state = env.state.model_dump()
+        score = grade_task("basic-triage", state)
+        assert score == 0.0
+
+
+class TestGraderDeterminism:
+    """Test that graders are deterministic."""
+
+    def test_same_state_same_score(self):
+        """Same state should always produce same score."""
+        env = TriageEnvironment()
+        env.reset(task_name="basic-triage")
+
+        env.step(TriageAction(
+            action_type=ActionType.ASSIGN_PRIORITY,
+            patient_id="P001",
+            priority_level=PriorityLevel.IMMEDIATE,
+        ))
+
+        state = env.state.model_dump()
+        score1 = grade_task("basic-triage", state)
+        score2 = grade_task("basic-triage", state)
+        assert score1 == score2
+
+
+class TestGraderVariance:
+    """Test that graders produce different scores for different behaviors."""
+
+    def test_different_behaviors_different_scores(self):
+        """Different action sequences should produce different scores."""
+        scores = []
+
+        # Perfect run
+        env = TriageEnvironment()
+        env.reset(task_name="basic-triage")
+        for pid, p in [("P001", PriorityLevel.IMMEDIATE), ("P002", PriorityLevel.NON_URGENT), ("P003", PriorityLevel.URGENT)]:
+            env.step(TriageAction(action_type=ActionType.ASSIGN_PRIORITY, patient_id=pid, priority_level=p))
+        scores.append(grade_task("basic-triage", env.state.model_dump()))
+
+        # All wrong
+        env = TriageEnvironment()
+        env.reset(task_name="basic-triage")
+        for pid, p in [("P001", PriorityLevel.NON_URGENT), ("P002", PriorityLevel.IMMEDIATE), ("P003", PriorityLevel.NON_URGENT)]:
+            env.step(TriageAction(action_type=ActionType.ASSIGN_PRIORITY, patient_id=pid, priority_level=p))
+        scores.append(grade_task("basic-triage", env.state.model_dump()))
+
+        # Partial 
+        env = TriageEnvironment()
+        env.reset(task_name="basic-triage")
+        for pid, p in [("P001", PriorityLevel.IMMEDIATE), ("P002", PriorityLevel.IMMEDIATE), ("P003", PriorityLevel.IMMEDIATE)]:
+            env.step(TriageAction(action_type=ActionType.ASSIGN_PRIORITY, patient_id=pid, priority_level=p))
+        scores.append(grade_task("basic-triage", env.state.model_dump()))
+
+        # No actions
+        env = TriageEnvironment()
+        env.reset(task_name="basic-triage")
+        scores.append(grade_task("basic-triage", env.state.model_dump()))
+
+        # Verify we get at least 3 different scores
+        unique_scores = set(scores)
+        assert len(unique_scores) >= 3, f"Expected varying scores, got: {scores}"

tests/test_http_flow.py

@@ -0,0 +1,58 @@
+"""Full HTTP flow test: reset -> step -> step -> step -> state -> grader score."""
+import httpx
+
+BASE = "http://127.0.0.1:8000"
+
+# Reset
+r = httpx.post(f"{BASE}/reset", json={"task_name": "basic-triage"})
+assert r.status_code == 200, f"Reset failed: {r.status_code}"
+data = r.json()
+obs = data.get("observation", data)
+pid = obs.get("current_patient", {}).get("patient_id", "?")
+print(f"1. Reset OK: status={r.status_code}, patient={pid}")
+
+# Step 1: assign P001 as IMMEDIATE (correct)
+r = httpx.post(f"{BASE}/step", json={"action": {
+    "action_type": "assign_priority",
+    "patient_id": "P001",
+    "priority_level": "immediate"
+}})
+assert r.status_code == 200, f"Step 1 failed: {r.status_code} {r.text}"
+data = r.json()
+print(f"2. Step 1: reward={data.get('reward')}, done={data.get('done')}")
+
+# Step 2: assign P002 as NON_URGENT (correct)
+r = httpx.post(f"{BASE}/step", json={"action": {
+    "action_type": "assign_priority",
+    "patient_id": "P002",
+    "priority_level": "non_urgent"
+}})
+assert r.status_code == 200, f"Step 2 failed: {r.status_code}"
+data = r.json()
+print(f"3. Step 2: reward={data.get('reward')}, done={data.get('done')}")
+
+# Step 3: assign P003 as URGENT (correct)
+r = httpx.post(f"{BASE}/step", json={"action": {
+    "action_type": "assign_priority",
+    "patient_id": "P003",
+    "priority_level": "urgent"
+}})
+assert r.status_code == 200, f"Step 3 failed: {r.status_code}"
+data = r.json()
+print(f"4. Step 3: reward={data.get('reward')}, done={data.get('done')}")
+assert data.get("done") == True, "Should be done after all patients assigned"
+
+# State
+r = httpx.get(f"{BASE}/state")
+assert r.status_code == 200, f"State failed: {r.status_code}"
+state = r.json()
+print(f"5. State: assignments={state.get('assignments')}")
+
+# Verify via grader
+import sys; sys.path.insert(0, ".")
+from triage_flow.graders import grade_task
+score = grade_task("basic-triage", state)
+print(f"6. Grader score: {score}")
+assert score == 1.0, f"Expected 1.0, got {score}"
+
+print("\n=== FULL HTTP FLOW TEST PASSED ===")

tests/test_inference_logging.py

@@ -0,0 +1,66 @@
+"""
+Test that inference.py produces the exact stdout format required.
+Runs the environment directly (no LLM needed) with mock actions to verify logging format.
+"""
+import sys
+import os
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from triage_flow.environment import TriageEnvironment
+from triage_flow.graders import grade_task
+from models import TriageAction, ActionType, PriorityLevel
+
+# Test the logging functions from inference.py
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from inference import log_start, log_step, log_end
+
+import io
+from contextlib import redirect_stdout
+
+# Capture stdout
+f = io.StringIO()
+with redirect_stdout(f):
+    log_start(task="basic-triage", env="triage-flow", model="test-model")
+
+    # Simulate a 3-step episode
+    log_step(step=1, action="assign_priority(P001,immediate)", reward=0.50, done=False, error=None)
+    log_step(step=2, action="assign_priority(P002,non_urgent)", reward=0.10, done=False, error=None)
+    log_step(step=3, action="assign_priority(P003,urgent)", reward=0.50, done=True, error=None)
+
+    log_end(success=True, steps=3, score=1.00, rewards=[0.50, 0.10, 0.50])
+
+output = f.getvalue()
+lines = output.strip().split("\n")
+
+print("=== CAPTURED STDOUT ===")
+for line in lines:
+    print(repr(line))
+print()
+
+# Validate format
+assert lines[0].startswith("[START]"), f"Line 0 should start with [START]: {lines[0]}"
+assert "task=basic-triage" in lines[0], f"Missing task: {lines[0]}"
+assert "env=triage-flow" in lines[0], f"Missing env: {lines[0]}"
+assert "model=test-model" in lines[0], f"Missing model: {lines[0]}"
+
+for i in range(1, 4):
+    assert lines[i].startswith("[STEP]"), f"Line {i} should start with [STEP]: {lines[i]}"
+    assert f"step={i}" in lines[i], f"Missing step={i}: {lines[i]}"
+    assert "reward=" in lines[i], f"Missing reward: {lines[i]}"
+    assert "done=" in lines[i], f"Missing done: {lines[i]}"
+    assert "error=" in lines[i], f"Missing error: {lines[i]}"
+
+# Check [STEP] format details
+assert "done=false" in lines[1], f"done should be lowercase: {lines[1]}"
+assert "done=true" in lines[3], f"done should be lowercase: {lines[3]}"
+assert "error=null" in lines[1], f"error should be 'null': {lines[1]}"
+assert "reward=0.50" in lines[1], f"reward should be 2 decimal: {lines[1]}"
+
+# Check [END]
+assert lines[4].startswith("[END]"), f"Line 4 should start with [END]: {lines[4]}"
+assert "success=true" in lines[4], f"success should be lowercase: {lines[4]}"
+assert "steps=3" in lines[4], f"Missing steps: {lines[4]}"
+assert "score=1.00" in lines[4], f"score should be 2 decimal: {lines[4]}"
+assert "rewards=0.50,0.10,0.50" in lines[4], f"rewards format wrong: {lines[4]}"
+
+print("=== ALL FORMAT CHECKS PASSED ===")

tests/test_models.py

@@ -0,0 +1,129 @@
+"""
+Module: test_models.py
+Purpose: Test all Pydantic models and enums validate correctly.
+Part of: Medical Triage Assistant — OpenEnv Round 1
+Author: Team Squirrel
+"""
+
+import sys
+import os
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import pytest
+from models import (
+    PriorityLevel,
+    ActionType,
+    InfoField,
+    PatientRecord,
+    TriageAction,
+    TriageObservation,
+    TriageState,
+)
+
+
+class TestEnums:
+    """Test all enum definitions."""
+
+    def test_priority_levels(self):
+        """Verify all 4 priority levels exist and have correct values."""
+        assert PriorityLevel.IMMEDIATE.value == "immediate"
+        assert PriorityLevel.URGENT.value == "urgent"
+        assert PriorityLevel.LESS_URGENT.value == "less_urgent"
+        assert PriorityLevel.NON_URGENT.value == "non_urgent"
+
+    def test_action_types(self):
+        """Verify all 5 action types exist."""
+        assert len(ActionType) == 5
+        assert ActionType.ASSIGN_PRIORITY.value == "assign_priority"
+        assert ActionType.REQUEST_INFO.value == "request_info"
+        assert ActionType.ESCALATE.value == "escalate"
+        assert ActionType.DEFER.value == "defer"
+        assert ActionType.ADVANCE_QUEUE.value == "advance_queue"
+
+    def test_info_fields(self):
+        """Verify all 5 info fields exist."""
+        assert len(InfoField) == 5
+        assert InfoField.VITALS.value == "vitals"
+
+
+class TestTriageAction:
+    """Test action model validation."""
+
+    def test_valid_assign_priority(self):
+        """Test creating a valid assign_priority action."""
+        action = TriageAction(
+            action_type=ActionType.ASSIGN_PRIORITY,
+            patient_id="P001",
+            priority_level=PriorityLevel.IMMEDIATE,
+        )
+        assert action.action_type == ActionType.ASSIGN_PRIORITY
+        assert action.patient_id == "P001"
+        assert action.priority_level == PriorityLevel.IMMEDIATE
+
+    def test_valid_request_info(self):
+        """Test creating a valid request_info action."""
+        action = TriageAction(
+            action_type=ActionType.REQUEST_INFO,
+            patient_id="P001",
+            info_field=InfoField.VITALS,
+        )
+        assert action.action_type == ActionType.REQUEST_INFO
+        assert action.info_field == InfoField.VITALS
+
+    def test_valid_advance_queue(self):
+        """Test creating a valid advance_queue action."""
+        action = TriageAction(
+            action_type=ActionType.ADVANCE_QUEUE,
+        )
+        assert action.action_type == ActionType.ADVANCE_QUEUE
+
+
+class TestTriageObservation:
+    """Test observation model."""
+
+    def test_default_observation(self):
+        """Test observation with default values."""
+        obs = TriageObservation(done=False, reward=None)
+        assert obs.done is False
+        assert obs.reward is None
+        assert obs.queue_length == 0
+
+    def test_full_observation(self):
+        """Test observation with all fields."""
+        obs = TriageObservation(
+            done=False,
+            reward=0.50,
+            current_patient={"patient_id": "P001", "age": 62},
+            queue_length=3,
+            queue_position=1,
+            missing_fields=["vitals"],
+            previous_action_feedback="Episode started",
+            step_number=1,
+            task_name="basic-triage",
+        )
+        assert obs.current_patient["patient_id"] == "P001"
+        assert obs.queue_length == 3
+        assert obs.missing_fields == ["vitals"]
+
+
+class TestTriageState:
+    """Test state model."""
+
+    def test_default_state(self):
+        """Test state with default values."""
+        state = TriageState()
+        assert state.step_count == 0
+        assert state.task_name == ""
+        assert state.patients == []
+        assert state.assignments == {}
+
+    def test_serializable(self):
+        """Test state can be serialized to dict."""
+        state = TriageState(
+            task_name="basic-triage",
+            patients=[{"patient_id": "P001"}],
+            assignments={"P001": "immediate"},
+        )
+        d = state.model_dump()
+        assert d["task_name"] == "basic-triage"
+        assert d["assignments"]["P001"] == "immediate"

triage_flow/__init__.py

@@ -0,0 +1,26 @@
+"""
+Module: __init__.py
+Purpose: Package initialization for the triage_flow environment module.
+Part of: Medical Triage Assistant — OpenEnv Round 1
+Author: Team Squirrel
+
+Overview:
+    Exports the core environment class and task utilities for use by
+    the server and inference scripts.
+
+Usage:
+    from triage_flow.environment import TriageEnvironment
+    from triage_flow.tasks import TASK_NAMES
+    from triage_flow.graders import grade_task
+"""
+
+from triage_flow.environment import TriageEnvironment
+from triage_flow.tasks import TASK_NAMES, get_task_config
+from triage_flow.graders import grade_task
+
+__all__ = [
+    "TriageEnvironment",
+    "TASK_NAMES",
+    "get_task_config",
+    "grade_task",
+]

triage_flow/environment.py

@@ -0,0 +1,544 @@
+"""
+Module: environment.py
+Purpose: Core environment logic implementing reset(), step(), and state for TriageFlow.
+Part of: Medical Triage Assistant — OpenEnv Round 1
+Author: Team Squirrel
+
+Overview:
+    Implements the TriageEnvironment class which extends the OpenEnv Environment
+    base class. This is the central simulation engine: it manages the patient queue,
+    validates and applies agent actions, computes rewards, and tracks state for
+    deterministic grading. The environment follows a standard episode lifecycle:
+    reset() → (observe, act, step) loop → done.
+
+Dependencies:
+    - openenv.core.env_server: Environment base class
+    - models: TriageAction, TriageObservation, TriageState
+    - tasks: get_task_config for loading patient data
+    - reward: compute_step_reward for per-step reward calculation
+
+Usage:
+    env = TriageEnvironment()
+    obs = env.reset(task_name="basic-triage")
+    obs = env.step(TriageAction(action_type="assign_priority", patient_id="P001", priority_level="immediate"))
+"""
+
+import uuid
+from typing import Any, Dict, List, Optional
+
+from openenv.core.env_server import Environment
+
+import sys
+import os
+
+# Add parent directory to path so we can import models
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from models import (
+    TriageAction,
+    TriageObservation,
+    TriageState,
+    PriorityLevel,
+    ActionType,
+    InfoField,
+)
+from triage_flow.tasks import get_task_config, TASK_NAMES
+from triage_flow.reward import compute_step_reward, compute_terminal_reward
+
+
+class TriageEnvironment(Environment):
+    """
+    Core triage simulation environment implementing the OpenEnv spec.
+
+    Manages a queue of patients that an agent must triage by assigning
+    urgency priorities, requesting missing information, escalating critical
+    cases, or deferring patients. The environment tracks all state for
+    deterministic grading.
+
+    Attributes:
+        SUPPORTS_CONCURRENT_SESSIONS (bool): Enables multiple simultaneous clients.
+
+    Notes:
+        Ground truth priorities are never exposed in the observation.
+        The grader reads from state() for scoring.
+    """
+
+    SUPPORTS_CONCURRENT_SESSIONS = True
+
+    def __init__(self):
+        """Initialize environment with empty state."""
+        self._state = TriageState()
+        self._patients: List[Dict[str, Any]] = []
+        self._task_config: Dict[str, Any] = {}
+        self._cumulative_reward: float = 0.0
+
+    def reset(self, seed=None, episode_id=None, task_name=None, **kwargs) -> TriageObservation:
+        """
+        Reset the environment for a new episode.
+
+        Loads the specified task's patient queue, clears all state,
+        and returns the initial observation showing the first patient.
+
+        Args:
+            seed: Optional random seed (not used — tasks are deterministic).
+            episode_id: Optional episode identifier.
+            task_name: Which task to load. Defaults to "basic-triage".
+            **kwargs: Additional arguments (may include task_name).
+
+        Returns:
+            TriageObservation: Initial observation with first patient visible.
+
+        Notes:
+            If task_name is not provided, it defaults to "basic-triage".
+            The task_name can also be passed via kwargs.
+        """
+        # Extract task_name from kwargs if not provided directly
+        if task_name is None:
+            task_name = kwargs.get("task_name", "basic-triage")
+
+        # Load task configuration
+        self._task_config = get_task_config(task_name)
+        self._patients = self._task_config["patients"]
+
+        # Reset cumulative reward
+        self._cumulative_reward = 0.0
+
+        # Initialize state
+        self._state = TriageState(
+            episode_id=episode_id or str(uuid.uuid4()),
+            step_count=0,
+            task_name=task_name,
+            patients=[p.copy() for p in self._patients],
+            assignments={},
+            escalations={},
+            info_requests=[],
+            action_history=[],
+            current_index=0,
+            max_steps=self._task_config["max_steps"],
+            queue_cleared=False,
+        )
+
+        # Build initial observation
+        return self._build_observation(
+            reward=None,
+            done=False,
+            feedback=f"Episode started. Task: {task_name}. "
+                     f"{len(self._patients)} patients in queue. "
+                     f"Assess each patient and assign appropriate priority.",
+        )
+
+    def step(self, action: TriageAction, timeout_s=None, **kwargs) -> TriageObservation:
+        """
+        Execute one agent action and return the resulting observation.
+
+        Validates the action, applies it to the internal state, computes
+        the per-step reward, and checks for done conditions.
+
+        Args:
+            action (TriageAction): The action to execute.
+            timeout_s: Optional timeout (not used).
+            **kwargs: Additional arguments.
+
+        Returns:
+            TriageObservation: Updated observation after action.
+
+        Notes:
+            Invalid actions are handled gracefully with a penalty reward
+            and descriptive feedback. The episode always emits an observation.
+        """
+        self._state.step_count += 1
+
+        # Validate action and compute reward
+        is_valid, validation_msg = self._validate_action(action)
+
+        # Get current patient info for reward computation
+        current_patient = self._get_current_patient()
+        ground_truth = current_patient.get("ground_truth_priority") if current_patient else None
+        missing_fields = current_patient.get("missing_fields", []) if current_patient else []
+        info_changes = current_patient.get("info_changes_decision", False) if current_patient else False
+        patient_id = action.patient_id if action.patient_id else (
+            current_patient.get("patient_id", "") if current_patient else ""
+        )
+
+        # Compute reward
+        reward_value, reward_reason = compute_step_reward(
+            action_type=action.action_type.value if isinstance(action.action_type, ActionType) else str(action.action_type),
+            patient_id=patient_id,
+            assigned_priority=action.priority_level.value if action.priority_level else None,
+            ground_truth_priority=ground_truth,
+            info_field=action.info_field.value if action.info_field else None,
+            missing_fields=missing_fields,
+            info_changes_decision=info_changes,
+            already_assigned=patient_id in self._state.assignments,
+            already_escalated=patient_id in self._state.escalations,
+            escalation_reason=action.escalation_reason,
+            action_history=self._state.action_history,
+            is_valid_action=is_valid,
+        )
+
+        # Apply action to state
+        feedback = self._apply_action(action, is_valid, validation_msg)
+
+        # Record action in history
+        self._state.action_history.append({
+            "step": self._state.step_count,
+            "action_type": action.action_type.value if isinstance(action.action_type, ActionType) else str(action.action_type),
+            "patient_id": patient_id,
+            "priority_level": action.priority_level.value if action.priority_level else None,
+            "info_field": action.info_field.value if action.info_field else None,
+            "escalation_reason": action.escalation_reason,
+            "reward": reward_value,
+            "reward_reason": reward_reason,
+            "valid": is_valid,
+        })
+
+        # Accumulate reward
+        self._cumulative_reward += reward_value
+
+        # Check done conditions
+        done = self._check_done()
+
+        # Add terminal bonus if done and queue cleared
+        if done and self._state.queue_cleared:
+            terminal_reward, terminal_reason = compute_terminal_reward(True)
+            reward_value += terminal_reward
+            self._cumulative_reward += terminal_reward
+            feedback += f" {terminal_reason}."
+
+        return self._build_observation(
+            reward=reward_value,
+            done=done,
+            feedback=f"{feedback} | {reward_reason}",
+        )
+
+    @property
+    def state(self) -> TriageState:
+        """
+        Return the full internal state for grading and validation.
+
+        Returns:
+            TriageState: Complete state snapshot including ground truth,
+            action history, and all assignments.
+
+        Notes:
+            This is read by graders to compute the final score.
+            It must be serializable to JSON.
+        """
+        return self._state
+
+    # ========================================================================
+    # Internal Helpers
+    # ========================================================================
+
+    def _get_current_patient(self) -> Optional[Dict[str, Any]]:
+        """
+        Get the patient at the current queue position.
+
+        Returns:
+            Optional[dict]: Current patient data, or None if queue is empty/exhausted.
+        """
+        if 0 <= self._state.current_index < len(self._patients):
+            return self._patients[self._state.current_index]
+        return None
+
+    def _build_observation(
+        self, reward: Optional[float], done: bool, feedback: str
+    ) -> TriageObservation:
+        """
+        Build an observation object from current state.
+
+        Strips ground truth fields from the patient record before
+        including it in the observation.
+
+        Args:
+            reward (Optional[float]): Reward for this step.
+            done (bool): Whether episode is complete.
+            feedback (str): Human-readable feedback message.
+
+        Returns:
+            TriageObservation: Agent-visible observation.
+        """
+        current_patient = self._get_current_patient()
+
+        # Build agent-visible patient dict (strip ground truth)
+        visible_patient = None
+        missing = []
+        if current_patient:
+            visible_patient = {
+                k: v
+                for k, v in current_patient.items()
+                if k not in ("ground_truth_priority", "missing_fields",
+                             "info_changes_decision", "revealed_history",
+                             "revealed_medications")
+            }
+            # Calculate missing fields
+            for field in ["vitals", "history", "medications", "allergies"]:
+                if current_patient.get(field) is None:
+                    missing.append(field)
+
+        # Count remaining unassigned patients
+        unassigned_count = sum(
+            1 for p in self._patients
+            if p["patient_id"] not in self._state.assignments
+        )
+
+        return TriageObservation(
+            done=done,
+            reward=reward,
+            current_patient=visible_patient,
+            queue_length=unassigned_count,
+            queue_position=self._state.current_index + 1,
+            missing_fields=missing,
+            previous_action_feedback=feedback,
+            step_number=self._state.step_count,
+            task_name=self._state.task_name,
+        )
+
+    def _validate_action(self, action: TriageAction) -> tuple:
+        """
+        Validate that an action is legal in the current state.
+
+        Args:
+            action (TriageAction): Action to validate.
+
+        Returns:
+            tuple: (is_valid: bool, message: str)
+        """
+        action_type = action.action_type.value if isinstance(action.action_type, ActionType) else str(action.action_type)
+        current_patient = self._get_current_patient()
+
+        if not current_patient:
+            return False, "No patient at current queue position"
+
+        patient_id = action.patient_id or current_patient.get("patient_id", "")
+
+        # Check if patient_id exists in the queue
+        valid_ids = [p["patient_id"] for p in self._patients]
+        if patient_id and patient_id not in valid_ids:
+            return False, f"Patient {patient_id} not found in queue"
+
+        if action_type == "assign_priority":
+            if not action.priority_level:
+                return False, "ASSIGN_PRIORITY requires a priority_level"
+            if patient_id in self._state.assignments:
+                return False, f"Patient {patient_id} already has priority assigned"
+
+        elif action_type == "request_info":
+            if not action.info_field:
+                return False, "REQUEST_INFO requires an info_field"
+
+        elif action_type == "escalate":
+            if patient_id in self._state.escalations:
+                return False, f"Patient {patient_id} already escalated"
+
+        elif action_type == "defer":
+            if len(self._patients) <= 1:
+                return False, "Cannot defer when only one patient in queue"
+
+        elif action_type == "advance_queue":
+            pass  # Always valid
+
+        else:
+            return False, f"Unknown action type: {action_type}"
+
+        return True, "Action valid"
+
+    def _apply_action(self, action: TriageAction, is_valid: bool, validation_msg: str) -> str:
+        """
+        Apply a validated action to the internal state.
+
+        Args:
+            action (TriageAction): Action to apply.
+            is_valid (bool): Whether the action passed validation.
+            validation_msg (str): Validation feedback message.
+
+        Returns:
+            str: Feedback message describing what happened.
+        """
+        if not is_valid:
+            return f"Action rejected: {validation_msg}"
+
+        action_type = action.action_type.value if isinstance(action.action_type, ActionType) else str(action.action_type)
+        current_patient = self._get_current_patient()
+        patient_id = action.patient_id or (current_patient.get("patient_id", "") if current_patient else "")
+
+        if action_type == "assign_priority":
+            priority = action.priority_level.value if action.priority_level else "unknown"
+            self._state.assignments[patient_id] = priority
+            # Auto-advance to next unassigned patient
+            self._advance_to_next_unassigned()
+            # Check if queue is cleared
+            self._check_queue_cleared()
+            return f"Priority {priority} assigned to patient {patient_id}"
+
+        elif action_type == "request_info":
+            field = action.info_field.value if action.info_field else "unknown"
+            # Record the request
+            self._state.info_requests.append({
+                "patient_id": patient_id,
+                "field": field,
+                "step": self._state.step_count,
+            })
+            # Reveal the information if it was genuinely missing
+            target_patient = self._find_patient(patient_id)
+            if target_patient:
+                revealed_value = None
+                if field == "history" and target_patient.get("revealed_history"):
+                    target_patient["history"] = target_patient["revealed_history"]
+                    revealed_value = target_patient["history"]
+                elif field == "medications" and target_patient.get("revealed_medications"):
+                    target_patient["medications"] = target_patient["revealed_medications"]
+                    revealed_value = str(target_patient["medications"])
+                elif field == "vitals" and target_patient.get("vitals") is None:
+                    # Generate plausible vitals based on ground truth priority
+                    target_patient["vitals"] = self._generate_revealed_vitals(target_patient)
+                    revealed_value = str(target_patient["vitals"])
+                elif field == "allergies" and target_patient.get("allergies") is None:
+                    target_patient["allergies"] = []
+                    revealed_value = "No known allergies"
+
+                # Update info_complete if all fields now present
+                still_missing = any(
+                    target_patient.get(f) is None
+                    for f in ["vitals", "history", "medications", "allergies"]
+                )
+                target_patient["info_complete"] = not still_missing
+
+                if revealed_value:
+                    return f"Info requested: {field} revealed for patient {patient_id}: {revealed_value}"
+                return f"Info requested: {field} for patient {patient_id} (already present)"
+
+            return f"Info requested: {field} for patient {patient_id}"
+
+        elif action_type == "escalate":
+            reason = action.escalation_reason or "No reason provided"
+            self._state.escalations[patient_id] = reason
+            return f"Patient {patient_id} escalated to senior staff. Reason: {reason}"
+
+        elif action_type == "defer":
+            # Move patient to end of queue
+            target_idx = next(
+                (i for i, p in enumerate(self._patients) if p["patient_id"] == patient_id),
+                None
+            )
+            if target_idx is not None:
+                patient = self._patients.pop(target_idx)
+                self._patients.append(patient)
+                # Adjust current index if needed
+                if target_idx <= self._state.current_index:
+                    self._state.current_index = max(0, self._state.current_index - 1)
+            return f"Patient {patient_id} deferred to end of queue"
+
+        elif action_type == "advance_queue":
+            self._state.current_index = min(
+                self._state.current_index + 1, len(self._patients) - 1
+            )
+            next_patient = self._get_current_patient()
+            next_id = next_patient["patient_id"] if next_patient else "none"
+            return f"Advanced queue. Now viewing patient {next_id}"
+
+        return "Action applied"
+
+    def _advance_to_next_unassigned(self):
+        """
+        Move current_index to the next patient that hasn't been assigned yet.
+        Wraps around the queue if necessary.
+        """
+        start = self._state.current_index
+        for i in range(len(self._patients)):
+            idx = (start + i) % len(self._patients)
+            patient = self._patients[idx]
+            if patient["patient_id"] not in self._state.assignments:
+                self._state.current_index = idx
+                return
+        # All assigned — stay at current
+        self._state.current_index = min(start, len(self._patients) - 1)
+
+    def _check_queue_cleared(self):
+        """Check if all patients have been assigned a priority."""
+        all_assigned = all(
+            p["patient_id"] in self._state.assignments for p in self._patients
+        )
+        self._state.queue_cleared = all_assigned
+
+    def _check_done(self) -> bool:
+        """
+        Check if the episode should end.
+
+        Done conditions:
+        1. All patients assigned a priority (queue cleared)
+        2. Maximum steps reached
+
+        Returns:
+            bool: True if episode should end.
+        """
+        # All patients assigned
+        if self._state.queue_cleared:
+            return True
+
+        # Max steps reached
+        if self._state.step_count >= self._state.max_steps:
+            return True
+
+        return False
+
+    def _find_patient(self, patient_id: str) -> Optional[Dict[str, Any]]:
+        """
+        Find a patient in the queue by ID.
+
+        Args:
+            patient_id (str): Patient ID to search for.
+
+        Returns:
+            Optional[dict]: Patient data, or None if not found.
+        """
+        for p in self._patients:
+            if p["patient_id"] == patient_id:
+                return p
+        return None
+
+    def _generate_revealed_vitals(self, patient: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Generate plausible vitals for a patient whose vitals were missing.
+
+        Based on ground truth priority to maintain clinical consistency.
+
+        Args:
+            patient (dict): Patient data including ground_truth_priority.
+
+        Returns:
+            dict: Vitals dictionary with HR, BP, SpO2, temperature.
+
+        Notes:
+            # ASSUMPTION: Revealed vitals are consistent with the ground truth
+            # priority to maintain deterministic grading.
+        """
+        priority = patient.get("ground_truth_priority", "non_urgent")
+        if priority == "immediate":
+            return {
+                "heart_rate": 125,
+                "blood_pressure": "85/55",
+                "spo2": 90,
+                "temperature": 38.5,
+            }
+        elif priority == "urgent":
+            return {
+                "heart_rate": 105,
+                "blood_pressure": "145/95",
+                "spo2": 93,
+                "temperature": 38.8,
+            }
+        elif priority == "less_urgent":
+            return {
+                "heart_rate": 85,
+                "blood_pressure": "130/82",
+                "spo2": 97,
+                "temperature": 37.5,
+            }
+        else:  # non_urgent
+            return {
+                "heart_rate": 75,
+                "blood_pressure": "120/78",
+                "spo2": 99,
+                "temperature": 36.9,
+            }

triage_flow/graders.py

@@ -0,0 +1,364 @@
+"""
+Module: graders.py
+Purpose: Deterministic grading functions for each TriageFlow task.
+Part of: Medical Triage Assistant — OpenEnv Round 1
+Author: Team Squirrel
+
+Overview:
+    Implements one grader per task. Each grader reads from the environment's
+    state() output and returns a float in [0.0, 1.0]. Graders are deterministic:
+    the same state always produces the same score. Different agent behaviors
+    produce meaningfully different scores — graders that always return the same
+    score cause disqualification.
+
+Dependencies:
+    - models: TriageState for state access
+
+Usage:
+    from triage_flow.graders import grade_task
+    score = grade_task("basic-triage", state_dict)
+"""
+
+from typing import Any, Dict, List, Optional
+
+
+def grade_task(task_name: str, state: Dict[str, Any]) -> float:
+    """
+    Grade a completed episode for the given task.
+
+    Dispatches to the appropriate task-specific grader.
+
+    Args:
+        task_name (str): Which task to grade.
+        state (dict): Full state from environment.state().
+
+    Returns:
+        float: Score in [0.0, 1.0].
+
+    Raises:
+        ValueError: If task_name is not recognized.
+    """
+    graders = {
+        "basic-triage": _grade_basic_triage,
+        "incomplete-records-triage": _grade_incomplete_records,
+        "mass-casualty-triage": _grade_mass_casualty,
+    }
+
+    if task_name not in graders:
+        raise ValueError(f"Unknown task for grading: {task_name}")
+
+    score = graders[task_name](state)
+    # Clamp to [0.0, 1.0]
+    return max(0.0, min(1.0, score))
+
+
+# ============================================================================
+# Task 1: Basic Triage (Easy)
+# ============================================================================
+# Grader Logic: 1.0 if all 3 correct, 0.67 if 2 correct, 0.33 if 1 correct, 0.0 if none.
+
+def _grade_basic_triage(state: Dict[str, Any]) -> float:
+    """
+    Grade the basic triage task.
+
+    Simple accuracy-based grading: count how many patients were
+    assigned the correct priority out of 3 total.
+
+    Args:
+        state (dict): Full episode state.
+
+    Returns:
+        float: Score = correct_count / total_patients.
+
+    Notes:
+        Deterministic: same assignments → same score.
+        Varies meaningfully: 0.0, 0.33, 0.67, or 1.0.
+    """
+    patients = state.get("patients", [])
+    assignments = state.get("assignments", {})
+
+    if not patients:
+        return 0.0
+
+    correct = 0
+    total = len(patients)
+
+    for patient in patients:
+        pid = patient["patient_id"]
+        gt = patient.get("ground_truth_priority")
+        assigned = assignments.get(pid)
+        if assigned == gt:
+            correct += 1
+
+    # Score = proportion correct (rounds to 0.0, 0.33, 0.67, 1.0 for 3 patients)
+    return round(correct / total, 2)
+
+
+# ============================================================================
+# Task 2: Incomplete Records Triage (Medium)
+# ============================================================================
+# Grader Logic: Base score from priority accuracy (60%).
+# Bonus for appropriate info requests (20%).
+# Step efficiency (10%).
+# Penalize unnecessary info requests (10%).
+
+def _grade_incomplete_records(state: Dict[str, Any]) -> float:
+    """
+    Grade the incomplete records triage task.
+
+    Score decomposition:
+    - Priority accuracy: 60% weight
+    - Info request quality: 20% weight
+    - Escalation quality: 10% weight
+    - Step efficiency: 10% weight
+
+    Args:
+        state (dict): Full episode state.
+
+    Returns:
+        float: Weighted composite score in [0.0, 1.0].
+
+    Notes:
+        Info requests are scored based on whether they targeted
+        genuinely missing fields that would change the triage decision.
+    """
+    patients = state.get("patients", [])
+    assignments = state.get("assignments", {})
+    info_requests = state.get("info_requests", [])
+    max_steps = state.get("max_steps", 15)
+    step_count = state.get("step_count", 0)
+
+    if not patients:
+        return 0.0
+
+    # --- Priority Accuracy (60%) ---
+    correct = 0
+    total = len(patients)
+    for patient in patients:
+        pid = patient["patient_id"]
+        gt = patient.get("ground_truth_priority")
+        assigned = assignments.get(pid)
+        if assigned == gt:
+            correct += 1
+    priority_score = correct / total
+
+    # --- Info Request Quality (20%) ---
+    # Count appropriate vs inappropriate info requests
+    appropriate_requests = 0
+    unnecessary_requests = 0
+    for req in info_requests:
+        pid = req.get("patient_id")
+        field = req.get("field")
+        # Find the patient
+        patient = next((p for p in patients if p["patient_id"] == pid), None)
+        if patient:
+            missing = patient.get("missing_fields", [])
+            if field in missing:
+                appropriate_requests += 1
+            else:
+                unnecessary_requests += 1
+        else:
+            unnecessary_requests += 1
+
+    # Patients with genuinely useful missing info
+    patients_needing_info = sum(
+        1 for p in patients
+        if p.get("info_changes_decision", False) and p.get("missing_fields")
+    )
+    total_missing_fields = sum(
+        len(p.get("missing_fields", []))
+        for p in patients
+        if p.get("missing_fields")
+    )
+
+    if total_missing_fields > 0:
+        info_recall = appropriate_requests / total_missing_fields
+    else:
+        info_recall = 1.0  # No missing fields = perfect
+
+    info_precision = 1.0
+    total_requests = appropriate_requests + unnecessary_requests
+    if total_requests > 0:
+        info_precision = appropriate_requests / total_requests
+
+    info_score = (info_recall * 0.6 + info_precision * 0.4)
+
+    # --- Escalation Quality (10%) ---
+    escalations = state.get("escalations", {})
+    immediate_patients = [p for p in patients if p.get("ground_truth_priority") == "immediate"]
+    non_immediate_patients = [p for p in patients if p.get("ground_truth_priority") != "immediate"]
+
+    escalation_score = 1.0
+    if immediate_patients:
+        # Reward escalating immediate patients
+        correct_escalations = sum(
+            1 for p in immediate_patients if p["patient_id"] in escalations
+        )
+        missed_escalations = len(immediate_patients) - correct_escalations
+        wrong_escalations = sum(
+            1 for p in non_immediate_patients if p["patient_id"] in escalations
+        )
+        escalation_score = max(0.0, 1.0 - missed_escalations * 0.3 - wrong_escalations * 0.2)
+    else:
+        # No immediate patients — penalize any escalations
+        if escalations:
+            escalation_score = max(0.0, 1.0 - len(escalations) * 0.2)
+
+    # --- Step Efficiency (10%) ---
+    # Reward finishing efficiently (fewer steps = better)
+    if max_steps > 0:
+        efficiency = max(0.0, 1.0 - (step_count / max_steps) * 0.5)
+    else:
+        efficiency = 1.0
+
+    # Weighted composite
+    final_score = (
+        priority_score * 0.60
+        + info_score * 0.20
+        + escalation_score * 0.10
+        + efficiency * 0.10
+    )
+
+    return round(final_score, 2)
+
+
+# ============================================================================
+# Task 3: Mass Casualty Triage (Hard)
+# ============================================================================
+# Grader Logic:
+# - Correct ordering of top 2 critical patients: 30%
+# - IMMEDIATE misclassified as NON_URGENT: heavy penalty
+# - General accuracy on remaining: 30%
+# - Step efficiency under pressure: 10%
+# - Weighted accuracy overall: 30%
+
+def _grade_mass_casualty(state: Dict[str, Any]) -> float:
+    """
+    Grade the mass casualty triage task.
+
+    Score decomposition:
+    - Critical patient identification: 30% weight
+      (Did the agent correctly identify and prioritize the 2 IMMEDIATE patients?)
+    - Overall priority accuracy (weighted by severity): 40% weight
+    - Step efficiency under time pressure: 15% weight
+    - Safety penalty for dangerous misclassifications: 15% weight
+
+    Args:
+        state (dict): Full episode state.
+
+    Returns:
+        float: Weighted composite score in [0.0, 1.0].
+
+    Notes:
+        IMMEDIATE patients misclassified as NON_URGENT receive heavy penalties.
+        The hard task is designed so frontier models will struggle with
+        interaction reasoning (e.g., anticoagulant + head injury).
+    """
+    patients = state.get("patients", [])
+    assignments = state.get("assignments", {})
+    action_history = state.get("action_history", [])
+    max_steps = state.get("max_steps", 20)
+    step_count = state.get("step_count", 0)
+
+    if not patients:
+        return 0.0
+
+    # --- Critical Patient Identification (30%) ---
+    # The two IMMEDIATE patients should be identified correctly
+    immediate_patients = [p for p in patients if p.get("ground_truth_priority") == "immediate"]
+    critical_score = 0.0
+
+    if immediate_patients:
+        correctly_identified = 0
+        for p in immediate_patients:
+            pid = p["patient_id"]
+            if assignments.get(pid) == "immediate":
+                correctly_identified += 1
+        critical_score = correctly_identified / len(immediate_patients)
+
+        # Bonus: Were they triaged FIRST? (within first N actions)
+        assign_actions = [
+            a for a in action_history
+            if a.get("action_type") == "assign_priority" and a.get("valid", True)
+        ]
+        if len(assign_actions) >= 2:
+            first_two_pids = [a.get("patient_id") for a in assign_actions[:2]]
+            immediate_pids = {p["patient_id"] for p in immediate_patients}
+            # Bonus for triaging IMMEDIATE patients in first two actions
+            first_correct = sum(1 for pid in first_two_pids if pid in immediate_pids)
+            critical_score = min(1.0, critical_score + first_correct * 0.15)
+
+    # --- Overall Priority Accuracy (weighted by severity) (40%) ---
+    # Priority weights: IMMEDIATE errors are more costly
+    SEVERITY_WEIGHTS = {
+        "immediate": 3.0,
+        "urgent": 2.0,
+        "less_urgent": 1.0,
+        "non_urgent": 0.5,
+    }
+
+    weighted_correct = 0.0
+    total_weight = 0.0
+
+    for patient in patients:
+        pid = patient["patient_id"]
+        gt = patient.get("ground_truth_priority", "non_urgent")
+        assigned = assignments.get(pid)
+        weight = SEVERITY_WEIGHTS.get(gt, 1.0)
+        total_weight += weight
+
+        if assigned == gt:
+            weighted_correct += weight
+        elif assigned is None:
+            pass  # Not assigned — no credit, no extra penalty here
+
+    accuracy_score = weighted_correct / total_weight if total_weight > 0 else 0.0
+
+    # --- Step Efficiency (15%) ---
+    # With 20 steps for 8 patients, efficiency matters
+    assigned_count = len(assignments)
+    total_patients = len(patients)
+
+    if total_patients > 0:
+        completion_rate = assigned_count / total_patients
+    else:
+        completion_rate = 0.0
+
+    # Efficiency = completion rate * time efficiency
+    if max_steps > 0 and assigned_count > 0:
+        steps_per_patient = step_count / assigned_count
+        optimal_steps_per_patient = 1.5  # Some info requests are okay
+        time_efficiency = max(0.0, 1.0 - max(0, steps_per_patient - optimal_steps_per_patient) * 0.2)
+        efficiency_score = completion_rate * time_efficiency
+    else:
+        efficiency_score = 0.0
+
+    # --- Safety Penalty (15%) ---
+    # Penalize dangerous misclassifications
+    safety_score = 1.0
+
+    for patient in patients:
+        pid = patient["patient_id"]
+        gt = patient.get("ground_truth_priority")
+        assigned = assignments.get(pid)
+
+        if gt == "immediate" and assigned == "non_urgent":
+            safety_score -= 0.40  # Catastrophic error
+        elif gt == "immediate" and assigned == "less_urgent":
+            safety_score -= 0.30  # Very dangerous
+        elif gt == "urgent" and assigned == "non_urgent":
+            safety_score -= 0.20  # Dangerous
+        elif gt == "immediate" and assigned is None:
+            safety_score -= 0.25  # Failed to triage critical patient
+
+    safety_score = max(0.0, safety_score)
+
+    # Weighted composite
+    final_score = (
+        critical_score * 0.30
+        + accuracy_score * 0.40
+        + efficiency_score * 0.15
+        + safety_score * 0.15
+    )
+
+    return round(final_score, 2)

triage_flow/reward.py

@@ -0,0 +1,246 @@
+"""
+Module: reward.py
+Purpose: Compute per-step rewards for agent actions in the TriageFlow environment.
+Part of: Medical Triage Assistant — OpenEnv Round 1
+Author: Team Squirrel
+
+Overview:
+    Implements the reward shaping function for the TriageFlow environment.
+    Rewards are computed per-step based on the action taken and the ground
+    truth state. Correct urgency assignments earn positive rewards scaled
+    by severity. Incorrect assignments, unnecessary actions, and invalid
+    actions incur penalties. A terminal bonus is awarded for clearing the queue.
+
+Dependencies:
+    - models: PriorityLevel, ActionType, InfoField enums
+
+Usage:
+    from triage_flow.reward import compute_step_reward
+"""
+
+from typing import Any, Dict, List, Optional, Tuple
+
+
+# ============================================================================
+# Reward Constants (from PART 8 of build reference)
+# ============================================================================
+
+# Correct priority assignment rewards
+REWARD_CORRECT_IMMEDIATE = 0.50
+REWARD_CORRECT_URGENT = 0.30
+REWARD_CORRECT_LESS_URGENT = 0.20
+REWARD_CORRECT_NON_URGENT = 0.10
+
+# Misclassification penalties
+PENALTY_IMMEDIATE_AS_NON_URGENT = -0.50
+PENALTY_URGENT_AS_NON_URGENT = -0.30
+PENALTY_MISCLASSIFY_DEFAULT = -0.20
+
+# Info request rewards/penalties
+REWARD_APPROPRIATE_INFO_REQUEST = 0.10
+PENALTY_UNNECESSARY_INFO_REQUEST = -0.10
+
+# Escalation rewards/penalties
+REWARD_APPROPRIATE_ESCALATION = 0.20
+PENALTY_UNNECESSARY_ESCALATION = -0.10
+
+# Invalid action penalty
+PENALTY_INVALID_ACTION = -0.05
+
+# Loop/no-op penalty
+PENALTY_REPEATED_NOOP = -0.10
+
+# Terminal bonus
+REWARD_QUEUE_CLEARED = 0.20
+
+# Correct assignment reward lookup by ground truth priority
+CORRECT_ASSIGNMENT_REWARDS = {
+    "immediate": REWARD_CORRECT_IMMEDIATE,
+    "urgent": REWARD_CORRECT_URGENT,
+    "less_urgent": REWARD_CORRECT_LESS_URGENT,
+    "non_urgent": REWARD_CORRECT_NON_URGENT,
+}
+
+
+def compute_step_reward(
+    action_type: str,
+    patient_id: str,
+    assigned_priority: Optional[str],
+    ground_truth_priority: Optional[str],
+    info_field: Optional[str],
+    missing_fields: List[str],
+    info_changes_decision: bool,
+    already_assigned: bool,
+    already_escalated: bool,
+    escalation_reason: Optional[str],
+    action_history: List[Dict[str, Any]],
+    is_valid_action: bool,
+) -> Tuple[float, str]:
+    """
+    Compute the reward for a single step based on the action taken.
+
+    Args:
+        action_type (str): The type of action taken.
+        patient_id (str): The patient the action targets.
+        assigned_priority (Optional[str]): The priority level assigned (for assign_priority).
+        ground_truth_priority (Optional[str]): The correct priority for the patient.
+        info_field (Optional[str]): The field requested (for request_info).
+        missing_fields (List[str]): Fields currently missing for the patient.
+        info_changes_decision (bool): Whether requesting info would change the triage decision.
+        already_assigned (bool): Whether this patient already has a priority assigned.
+        already_escalated (bool): Whether this patient has already been escalated.
+        escalation_reason (Optional[str]): Reason given for escalation.
+        action_history (List[dict]): History of all previous actions.
+        is_valid_action (bool): Whether the action passed validation.
+
+    Returns:
+        Tuple[float, str]: (reward_value, reason_string)
+
+    Notes:
+        Rewards are raw values that may accumulate outside [0, 1] during an episode.
+        The grader normalizes the final episode score to [0.0, 1.0].
+    """
+    # Invalid action
+    if not is_valid_action:
+        return PENALTY_INVALID_ACTION, "Invalid or malformed action"
+
+    # Check for repeated no-op / loop behavior
+    # (same action on same patient repeated 3+ times)
+    if _is_repeated_action(action_type, patient_id, action_history):
+        return PENALTY_REPEATED_NOOP, "Repeated action detected (loop behavior)"
+
+    # ---- ASSIGN PRIORITY ----
+    if action_type == "assign_priority":
+        if already_assigned:
+            return PENALTY_INVALID_ACTION, f"Patient {patient_id} already has priority assigned"
+
+        if assigned_priority == ground_truth_priority:
+            reward = CORRECT_ASSIGNMENT_REWARDS.get(ground_truth_priority, 0.10)
+            return reward, f"Correct priority assignment: {assigned_priority}"
+        else:
+            # Calculate misclassification penalty based on severity
+            penalty = _misclassification_penalty(ground_truth_priority, assigned_priority)
+            return penalty, (
+                f"Incorrect priority: assigned {assigned_priority}, "
+                f"correct is {ground_truth_priority}"
+            )
+
+    # ---- REQUEST INFO ----
+    elif action_type == "request_info":
+        if info_field and info_field in missing_fields:
+            # Field is genuinely missing — appropriate request
+            if info_changes_decision:
+                return REWARD_APPROPRIATE_INFO_REQUEST, (
+                    f"Appropriate info request: {info_field} (would change decision)"
+                )
+            else:
+                # Field is missing but wouldn't change the decision
+                # Small positive for being thorough, but less than if it mattered
+                return REWARD_APPROPRIATE_INFO_REQUEST * 0.5, (
+                    f"Info request: {info_field} (field missing but wouldn't change triage)"
+                )
+        else:
+            return PENALTY_UNNECESSARY_INFO_REQUEST, (
+                f"Unnecessary info request: {info_field} is already present"
+            )
+
+    # ---- ESCALATE ----
+    elif action_type == "escalate":
+        if already_escalated:
+            return PENALTY_INVALID_ACTION, f"Patient {patient_id} already escalated"
+
+        # Escalation is appropriate for IMMEDIATE priority patients
+        if ground_truth_priority == "immediate":
+            return REWARD_APPROPRIATE_ESCALATION, (
+                f"Appropriate escalation for critical patient {patient_id}"
+            )
+        else:
+            return PENALTY_UNNECESSARY_ESCALATION, (
+                f"Unnecessary escalation for {ground_truth_priority} patient {patient_id}"
+            )
+
+    # ---- DEFER ----
+    elif action_type == "defer":
+        # Deferring is a neutral to slightly negative action
+        # It's okay for non-urgent patients but bad for immediate ones
+        if ground_truth_priority == "immediate":
+            return -0.20, f"Dangerous: deferred IMMEDIATE patient {patient_id}"
+        elif ground_truth_priority == "urgent":
+            return -0.10, f"Risky: deferred URGENT patient {patient_id}"
+        else:
+            return 0.0, f"Deferred patient {patient_id}"
+
+    # ---- ADVANCE QUEUE ----
+    elif action_type == "advance_queue":
+        # Neutral action — just moves to next patient
+        return 0.0, "Advanced to next patient in queue"
+
+    # Fallback
+    return 0.0, "Unknown action type"
+
+
+def compute_terminal_reward(queue_cleared: bool) -> Tuple[float, str]:
+    """
+    Compute the terminal bonus reward at end of episode.
+
+    Args:
+        queue_cleared (bool): Whether all patients have been assigned priorities.
+
+    Returns:
+        Tuple[float, str]: (reward_value, reason_string)
+    """
+    if queue_cleared:
+        return REWARD_QUEUE_CLEARED, "Terminal bonus: all patients triaged"
+    return 0.0, "Episode ended without clearing queue"
+
+
+def _misclassification_penalty(
+    ground_truth: Optional[str], assigned: Optional[str]
+) -> float:
+    """
+    Calculate penalty for incorrect priority assignment.
+
+    More severe misclassifications (e.g. IMMEDIATE → NON_URGENT) receive
+    larger penalties than minor ones (e.g. URGENT → LESS_URGENT).
+
+    Args:
+        ground_truth (Optional[str]): Correct priority level.
+        assigned (Optional[str]): Priority level assigned by agent.
+
+    Returns:
+        float: Penalty value (negative).
+    """
+    if ground_truth == "immediate" and assigned == "non_urgent":
+        return PENALTY_IMMEDIATE_AS_NON_URGENT
+    elif ground_truth == "immediate" and assigned == "less_urgent":
+        return -0.40
+    elif ground_truth == "urgent" and assigned == "non_urgent":
+        return PENALTY_URGENT_AS_NON_URGENT
+    elif ground_truth == "immediate" and assigned == "urgent":
+        return -0.15  # Close but still wrong for critical patient
+    elif ground_truth == "urgent" and assigned == "less_urgent":
+        return -0.15
+    else:
+        return PENALTY_MISCLASSIFY_DEFAULT
+
+
+def _is_repeated_action(
+    action_type: str, patient_id: str, action_history: List[Dict[str, Any]]
+) -> bool:
+    """
+    Detect if the agent is in a loop by checking for 3+ identical actions.
+
+    Args:
+        action_type (str): Current action type.
+        patient_id (str): Current target patient.
+        action_history (List[dict]): Previous actions taken.
+
+    Returns:
+        bool: True if this exact action has been taken 2+ times already.
+    """
+    count = sum(
+        1
+        for a in action_history
+        if a.get("action_type") == action_type and a.get("patient_id") == patient_id
+    )
+    return count >= 2  # Current would be the 3rd

triage_flow/tasks.py

@@ -0,0 +1,478 @@
+"""
+Module: tasks.py
+Purpose: Define all task configurations and patient case data for the TriageFlow environment.
+Part of: Medical Triage Assistant — OpenEnv Round 1
+Author: Team Squirrel
+
+Overview:
+    This module contains the hardcoded patient case data for all three tasks:
+    basic-triage (easy), incomplete-records-triage (medium), and mass-casualty-triage (hard).
+    Each task defines a set of patients with ground truth priority labels, as well as
+    task-specific configuration like max steps and which fields may be missing.
+
+Dependencies:
+    - models: PriorityLevel, InfoField for type-safe enums
+
+Usage:
+    from triage_flow.tasks import get_task_config, TASK_NAMES
+"""
+
+from typing import Any, Dict, List
+
+# ASSUMPTION: Ground truth priorities are stored here alongside patient data.
+# They are stripped from the observation before being shown to the agent.
+
+
+TASK_NAMES = ["basic-triage", "incomplete-records-triage", "mass-casualty-triage"]
+
+
+def get_task_config(task_name: str) -> Dict[str, Any]:
+    """
+    Load the configuration for a given task.
+
+    Args:
+        task_name (str): One of the three defined task names.
+
+    Returns:
+        dict: Task configuration including patients, max_steps, and description.
+
+    Raises:
+        ValueError: If task_name is not recognized.
+    """
+    if task_name not in TASK_CONFIGS:
+        raise ValueError(
+            f"Unknown task: {task_name}. Available tasks: {list(TASK_CONFIGS.keys())}"
+        )
+    # Return a deep copy to prevent state mutation between episodes
+    import copy
+    return copy.deepcopy(TASK_CONFIGS[task_name])
+
+
+# ============================================================================
+# Task 1: Basic Triage (Easy)
+# ============================================================================
+# Scenario: 3 patients with complete records and clearly distinct urgency levels.
+# One has chest pain and abnormal vitals, one has a sprained ankle, one has a mild fever.
+# Agent Objective: Assign correct priority to all 3 patients.
+# Grader Logic: 1.0 if all correct, 0.67 if 2 correct, 0.33 if 1 correct, 0.0 if none.
+# Difficulty: Symptoms and vitals make urgency unambiguous. No missing data.
+
+TASK1_PATIENTS = [
+    {
+        "patient_id": "P001",
+        "age": 62,
+        "chief_complaint": "Severe chest pain radiating to left arm",
+        "symptoms": [
+            "crushing chest pain",
+            "shortness of breath",
+            "diaphoresis",
+            "nausea",
+            "left arm numbness",
+        ],
+        "vitals": {
+            "heart_rate": 110,
+            "blood_pressure": "180/100",
+            "spo2": 91,
+            "temperature": 37.2,
+        },
+        "history": "History of hypertension and type 2 diabetes. Previous MI 3 years ago.",
+        "medications": ["metformin", "lisinopril", "aspirin"],
+        "allergies": ["penicillin"],
+        "info_complete": True,
+        "ground_truth_priority": "immediate",
+    },
+    {
+        "patient_id": "P002",
+        "age": 28,
+        "chief_complaint": "Twisted ankle while jogging",
+        "symptoms": [
+            "ankle swelling",
+            "mild pain on weight bearing",
+            "slight bruising",
+        ],
+        "vitals": {
+            "heart_rate": 72,
+            "blood_pressure": "120/78",
+            "spo2": 99,
+            "temperature": 36.8,
+        },
+        "history": "No significant medical history. Active lifestyle.",
+        "medications": [],
+        "allergies": [],
+        "info_complete": True,
+        "ground_truth_priority": "non_urgent",
+    },
+    {
+        "patient_id": "P003",
+        "age": 45,
+        "chief_complaint": "Persistent high fever for 3 days with productive cough",
+        "symptoms": [
+            "fever",
+            "productive cough",
+            "body aches",
+            "fatigue",
+            "mild shortness of breath",
+        ],
+        "vitals": {
+            "heart_rate": 95,
+            "blood_pressure": "130/85",
+            "spo2": 95,
+            "temperature": 39.2,
+        },
+        "history": "Smoker for 20 years. No other chronic conditions.",
+        "medications": [],
+        "allergies": ["sulfa drugs"],
+        "info_complete": True,
+        "ground_truth_priority": "urgent",
+    },
+]
+
+
+# ============================================================================
+# Task 2: Incomplete Records Triage (Medium)
+# ============================================================================
+# Scenario: 4 patients, 2 of whom have missing vitals or incomplete history.
+# Agent must decide when to request info vs. assign based on available data.
+# Grader Logic: Base score from priority accuracy. Bonus for appropriate info requests.
+# Difficulty: Agent must reason about information value.
+
+TASK2_PATIENTS = [
+    {
+        "patient_id": "P101",
+        "age": 55,
+        "chief_complaint": "Sudden severe headache, worst of life",
+        "symptoms": [
+            "thunderclap headache",
+            "neck stiffness",
+            "photophobia",
+            "nausea",
+        ],
+        "vitals": None,  # MISSING — requesting vitals would be appropriate
+        "history": "Hypertension, family history of aneurysm.",
+        "medications": ["amlodipine"],
+        "allergies": [],
+        "info_complete": False,
+        "ground_truth_priority": "immediate",
+        # HIDDEN: Even without vitals, symptoms clearly indicate immediate priority.
+        # Requesting vitals is acceptable but not necessary for correct triage.
+        "missing_fields": ["vitals"],
+        "info_changes_decision": False,  # Priority is clear even without vitals
+    },
+    {
+        "patient_id": "P102",
+        "age": 34,
+        "chief_complaint": "Abdominal pain and fatigue",
+        "symptoms": [
+            "generalized abdominal discomfort",
+            "fatigue",
+            "mild nausea",
+        ],
+        "vitals": {
+            "heart_rate": 82,
+            "blood_pressure": "125/80",
+            "spo2": 98,
+            "temperature": 37.0,
+        },
+        "history": None,  # MISSING — requesting history is important here
+        "medications": None,  # MISSING
+        "allergies": [],
+        "info_complete": False,
+        "ground_truth_priority": "less_urgent",
+        # HIDDEN: Without history, we might not know this person has Crohn's disease
+        # which could change urgency. Info request is beneficial.
+        "missing_fields": ["history", "medications"],
+        "info_changes_decision": True,  # History would reveal chronic condition context
+        "revealed_history": "Diagnosed with Crohn's disease 5 years ago. Recent flare-ups.",
+        "revealed_medications": ["mesalamine", "prednisone"],
+    },
+    {
+        "patient_id": "P103",
+        "age": 72,
+        "chief_complaint": "Difficulty breathing and chest tightness",
+        "symptoms": [
+            "dyspnea on exertion",
+            "wheezing",
+            "chest tightness",
+            "cough",
+        ],
+        "vitals": {
+            "heart_rate": 100,
+            "blood_pressure": "145/90",
+            "spo2": 92,
+            "temperature": 37.1,
+        },
+        "history": "COPD, 40-year smoking history. Uses home oxygen intermittently.",
+        "medications": ["albuterol inhaler", "tiotropium", "home oxygen PRN"],
+        "allergies": ["aspirin"],
+        "info_complete": True,
+        "ground_truth_priority": "urgent",
+        "missing_fields": [],
+        "info_changes_decision": False,
+    },
+    {
+        "patient_id": "P104",
+        "age": 19,
+        "chief_complaint": "Sore throat and runny nose",
+        "symptoms": [
+            "sore throat",
+            "nasal congestion",
+            "mild cough",
+            "low-grade fever",
+        ],
+        "vitals": None,  # MISSING
+        "history": "No significant medical history.",
+        "medications": [],
+        "allergies": None,  # MISSING
+        "info_complete": False,
+        "ground_truth_priority": "non_urgent",
+        "missing_fields": ["vitals", "allergies"],
+        "info_changes_decision": False,  # Clearly non-urgent regardless
+    },
+]
+
+
+# ============================================================================
+# Task 3: Mass Casualty Triage (Hard)
+# ============================================================================
+# Scenario: 8 patients arrive simultaneously following a building collapse.
+# Some have conflicting symptom/vitals signals. Some have interacting conditions.
+# Step budget is tight (20 steps for 8 patients).
+# Agent Objective: Prioritize the queue correctly under time pressure.
+# Difficulty: Conflicting signals, interacting conditions, step pressure.
+
+TASK3_PATIENTS = [
+    {
+        "patient_id": "P201",
+        "age": 45,
+        "chief_complaint": "Trapped under rubble, severe leg crush injury",
+        "symptoms": [
+            "severe leg pain",
+            "visible deformity",
+            "bleeding from leg wound",
+            "confusion",
+            "tachycardia",
+        ],
+        "vitals": {
+            "heart_rate": 130,
+            "blood_pressure": "90/60",
+            "spo2": 94,
+            "temperature": 36.5,
+        },
+        "history": "No known medical history.",
+        "medications": [],
+        "allergies": [],
+        "info_complete": True,
+        "ground_truth_priority": "immediate",
+        # Crush injury with signs of hemorrhagic shock — highest priority
+    },
+    {
+        "patient_id": "P202",
+        "age": 67,
+        "chief_complaint": "Hit head on falling debris, on anticoagulants",
+        "symptoms": [
+            "scalp laceration",
+            "headache",
+            "mild confusion",
+            "dizziness",
+        ],
+        "vitals": {
+            "heart_rate": 88,
+            "blood_pressure": "150/85",
+            "spo2": 97,
+            "temperature": 36.9,
+        },
+        "history": "Atrial fibrillation, on warfarin. History of TIA.",
+        "medications": ["warfarin", "metoprolol", "lisinopril"],
+        "allergies": ["codeine"],
+        "info_complete": True,
+        "ground_truth_priority": "immediate",
+        # CRITICAL INTERACTION: Head injury + anticoagulant = high bleed risk
+        # Vitals look stable but the combination is life-threatening
+    },
+    {
+        "patient_id": "P203",
+        "age": 30,
+        "chief_complaint": "Arm laceration from broken glass",
+        "symptoms": [
+            "deep laceration on forearm",
+            "moderate bleeding",
+            "pain",
+        ],
+        "vitals": {
+            "heart_rate": 90,
+            "blood_pressure": "130/80",
+            "spo2": 98,
+            "temperature": 36.8,
+        },
+        "history": "Healthy, no chronic conditions.",
+        "medications": [],
+        "allergies": [],
+        "info_complete": True,
+        "ground_truth_priority": "less_urgent",
+        # Active bleeding but hemodynamically stable — can be managed with pressure
+    },
+    {
+        "patient_id": "P204",
+        "age": 8,
+        "chief_complaint": "Crying, complaining of arm pain after fall",
+        "symptoms": [
+            "arm pain",
+            "swelling at wrist",
+            "refuses to move arm",
+            "crying",
+        ],
+        "vitals": {
+            "heart_rate": 110,
+            "blood_pressure": "100/65",
+            "spo2": 99,
+            "temperature": 37.0,
+        },
+        "history": "No significant medical history. Up to date on vaccinations.",
+        "medications": [],
+        "allergies": ["amoxicillin"],
+        "info_complete": True,
+        "ground_truth_priority": "less_urgent",
+        # Likely fracture — painful but stable, not life-threatening
+        # NOTE: Elevated HR is normal for distressed child
+    },
+    {
+        "patient_id": "P205",
+        "age": 52,
+        "chief_complaint": "Dust inhalation, progressive breathing difficulty",
+        "symptoms": [
+            "coughing",
+            "wheezing",
+            "progressive dyspnea",
+            "chest tightness",
+            "hoarse voice",
+        ],
+        "vitals": {
+            "heart_rate": 115,
+            "blood_pressure": "140/90",
+            "spo2": 89,
+            "temperature": 37.1,
+        },
+        "history": "Asthma diagnosed in childhood.",
+        "medications": ["fluticasone inhaler"],
+        "allergies": [],
+        "info_complete": True,
+        "ground_truth_priority": "urgent",
+        # Inhalation injury with low SpO2 and asthma history — urgent
+        # Could deteriorate to immediate if airway compromise develops
+    },
+    {
+        "patient_id": "P206",
+        "age": 38,
+        "chief_complaint": "Minor scrapes and feeling shaken",
+        "symptoms": [
+            "superficial abrasions",
+            "anxiety",
+            "mild tremor",
+        ],
+        "vitals": {
+            "heart_rate": 100,
+            "blood_pressure": "135/85",
+            "spo2": 99,
+            "temperature": 36.9,
+        },
+        "history": "Anxiety disorder. Takes medication daily.",
+        "medications": ["sertraline"],
+        "allergies": [],
+        "info_complete": True,
+        "ground_truth_priority": "non_urgent",
+        # Minor injuries + anxiety response — non-urgent
+        # Elevated vitals are from anxiety, not injury severity
+    },
+    {
+        "patient_id": "P207",
+        "age": 75,
+        "chief_complaint": "Back pain and inability to walk after being knocked down",
+        "symptoms": [
+            "severe lower back pain",
+            "inability to bear weight",
+            "numbness in legs",
+            "urinary incontinence",
+        ],
+        "vitals": {
+            "heart_rate": 95,
+            "blood_pressure": "160/95",
+            "spo2": 96,
+            "temperature": 36.7,
+        },
+        "history": "Osteoporosis. Previous vertebral compression fracture.",
+        "medications": ["alendronate", "calcium supplement", "vitamin D"],
+        "allergies": [],
+        "info_complete": True,
+        "ground_truth_priority": "urgent",
+        # Numbness + incontinence = possible cauda equina syndrome — urgent
+        # Spinal injury with neurological signs needs urgent assessment
+    },
+    {
+        "patient_id": "P208",
+        "age": 22,
+        "chief_complaint": "Ringing in ears and small cut on forehead",
+        "symptoms": [
+            "tinnitus",
+            "minor forehead laceration",
+            "slight headache",
+        ],
+        "vitals": {
+            "heart_rate": 78,
+            "blood_pressure": "118/72",
+            "spo2": 99,
+            "temperature": 36.8,
+        },
+        "history": "No significant medical history.",
+        "medications": [],
+        "allergies": [],
+        "info_complete": True,
+        "ground_truth_priority": "non_urgent",
+        # Blast concussion effects — minor, monitor only
+    },
+]
+
+
+# ============================================================================
+# Task Configuration Registry
+# ============================================================================
+
+TASK_CONFIGS = {
+    "basic-triage": {
+        "task_name": "basic-triage",
+        "difficulty": "easy",
+        "description": (
+            "Three patients with complete records and clearly distinct urgency levels. "
+            "One has chest pain and abnormal vitals (IMMEDIATE), one has a sprained "
+            "ankle (NON_URGENT), and one has a persistent fever with cough (URGENT). "
+            "Assign the correct priority to all three."
+        ),
+        "patients": TASK1_PATIENTS,
+        "max_steps": 10,
+        "expected_actions": 3,  # One assignment per patient
+    },
+    "incomplete-records-triage": {
+        "task_name": "incomplete-records-triage",
+        "difficulty": "medium",
+        "description": (
+            "Four patients, two with missing vitals or history. The agent must decide "
+            "when to request missing information versus triaging with available data. "
+            "Bonus credit for requesting info when it would change the decision."
+        ),
+        "patients": TASK2_PATIENTS,
+        "max_steps": 15,
+        "expected_actions": 4,  # One assignment per patient (plus possible info requests)
+    },
+    "mass-casualty-triage": {
+        "task_name": "mass-casualty-triage",
+        "difficulty": "hard",
+        "description": (
+            "Eight patients from a building collapse. Some have conflicting signals "
+            "(e.g., stable vitals but on anticoagulants with head injury). Some have "
+            "interacting conditions. Step budget is tight at 20 steps for 8 patients. "
+            "Identify the two most critical patients first."
+        ),
+        "patients": TASK3_PATIENTS,
+        "max_steps": 20,
+        "expected_actions": 8,  # One assignment per patient
+    },
+}