Update models.py

models.py

@@ -4,15 +4,16 @@ models.py – PhishGuard-Env Pydantic Models
 
 Typed request / response schemas used by env.py (FastAPI).
 
-  PhishAction   : Body schema for POST /step
-  StepResponse  : Response schema for POST /step (OpenEnv grader compliance)
-  ResetResponse : Response schema for POST /reset
-
-FIXES
------
-  - StepResponse.task_id is Optional[str] — the episode-over guard returns None.
-  - StepResponse.task_group added — env.py now always returns this field;
-    missing from the model caused silent drops or Pydantic errors.
+PhishAction   : Body schema for POST /step
+StepResponse  : Response schema for POST /step  (OpenEnv grader compliance)
+ResetResponse : Response schema for POST /reset
+
+BUG FIX (v1.0.2 → v1.0.3)
+────────────────────────────────────────────────────────────────────────────
+  StepResponse.task_id was typed as `str` but the episode-already-over guard
+  branch in env.py returns task_id=None.  Pydantic would raise a validation
+  error on every post-episode /step call.
+  Fix: task_id is now Optional[str] with a default of None.
 """
 
 from __future__ import annotations
@@ -36,10 +37,9 @@ class PhishAction(BaseModel):
     reasoning : Optional one-sentence technical justification (for logging).
     """
     action: str = Field(
-        description=(
-            "Triage decision. Must be exactly one of: "
-            "MARK_SAFE | MOVE_TO_SPAM | QUARANTINE | BLOCK_DOMAIN"
-        )
+        max_length=64,
+        description="Triage decision. Must be exactly one of: "
+                    "MARK_SAFE | MOVE_TO_SPAM | QUARANTINE | BLOCK_DOMAIN"
     )
     reasoning: Optional[str] = Field(
         default=None,
@@ -56,7 +56,7 @@ class ResetRequest(BaseModel):
 
 
 # ─────────────────────────────────────────────────────────────────────────────
-# RESPONSE MODELS (OpenEnv spec — all fields required by validator)
+# RESPONSE MODELS  (OpenEnv spec — all fields required by validator)
 # ─────────────────────────────────────────────────────────────────────────────
 
 class StepResponse(BaseModel):
@@ -66,8 +66,9 @@ class StepResponse(BaseModel):
     The OpenEnv validator inspects `task_id` and `is_correct` on every step
     to count how many distinct tasks have been graded.
 
-    task_id    is Optional[str] — episode-over guard returns None.
-    task_group is Optional[str] — episode-over guard returns None.
+    task_id is Optional[str] (not str) because the episode-already-over guard
+    branch returns None — a non-optional field would cause a Pydantic
+    ValidationError on every post-episode call.
     """
     observation: Optional[Dict[str, Any]] = Field(
         description="Next email dict, or null when the episode is done"
@@ -78,13 +79,9 @@ class StepResponse(BaseModel):
     done: bool = Field(
         description="True when all scenarios are complete or health reaches 0"
     )
-    task_id: Optional[str] = Field(
+    task_id: Optional[str] = Field(          # BUG FIX: was `str`, must be Optional
         default=None,
-        description="Semantic task ID e.g. 'task_spam' — required by OpenEnv validator",
-    )
-    task_group: Optional[str] = Field(
-        default=None,
-        description="Difficulty level of this scenario: easy | medium | hard",
+        description="Scenario ID e.g. 'lv3' — required by OpenEnv validator"
     )
     is_correct: bool = Field(
         description="True when reward >= R_PERFECT (0.95)"
@@ -100,7 +97,7 @@ class ResetResponse(BaseModel):
         description="First email observation for this episode"
     )
     task_id: str = Field(
-        description="Semantic task ID of the first scenario e.g. 'task_spam'"
+        description="ID of the first scenario in this episode"
     )
     task_group: str = Field(
         description="Difficulty level of the first scenario: easy | medium | hard"
@@ -109,5 +106,5 @@ class ResetResponse(BaseModel):
         description="Active difficulty level for this episode"
     )
     total_tasks: int = Field(
-        description="Total number of scenarios loaded for this episode"
-    )
+        description="Total number of scenarios in this level"
+    )