feat: Implement wait action and enhance action handling in simulation environment

inference.py

@@ -42,33 +42,87 @@ STDOUT FORMAT
     [END] success=true steps=3 score=1.00 rewards=0.00,0.00,1.00
 """
 
-import asyncio
+import inspect
+import json
 import os
 import textwrap
-from typing import List, Optional
+import time
+from typing import Any, Callable, Dict, List, Optional
 
 from openai import OpenAI
-from models import CoenvAction
-from client import CoEnv
-from dotenv import load_dotenv
-load_dotenv()
-IMAGE_NAME = os.getenv("IMAGE_NAME")
-API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY")
-
-API_BASE_URL = os.getenv("API_BASE_URL") or "https://router.huggingface.co/v1"
-MODEL_NAME = os.getenv("MODEL_NAME") or "Qwen/Qwen2.5-72B-Instruct"
-ENV_URL = os.getenv("ENV_URL") or "http://localhost:8000"
-MAX_STEPS = 8
+try:
+    from dotenv import load_dotenv
+except ImportError:
+    load_dotenv = None
+
+try:
+    from models import CoenvAction
+    from client import CoEnv
+except ImportError:
+    from models import CoenvAction
+    from client import CoEnv
+
+from server.graders.grader_pod_recovery import grade as grade_pod_recovery
+from server.graders.grader_autoscaling import grade as grade_autoscaling
+from server.graders.grader_incident import grade as grade_incident
+
+if load_dotenv is not None:
+    load_dotenv()
+
+LLM_BASE_URL = os.getenv("LLM_BASE_URL", "https://router.huggingface.co/v1")
+ENV_URL = os.getenv("API_BASE_URL", "http://localhost:8000")
+API_DELAY = float(os.getenv("API_DELAY", "0"))
+
+MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen3-8B")
+API_KEY = os.getenv("OPENROUTER_API_KEY") or os.getenv("HF_TOKEN")
+
+BENCHMARKS = ["POD_RECOVERY", "AUTOSCALING", "INCIDENT"]
+TASK_NAMES = ["pod_recovery", "autoscaling", "incident"]
+
 TEMPERATURE = 0.7
 MAX_TOKENS = 150
 SUCCESS_SCORE_THRESHOLD = 0.1  # normalized score in [0, 1]
+DEFAULT_MAX_STEPS = 15
+
+SUCCESS_SCORE_THRESHOLD_BY_TASK: Dict[str, float] = {
+    "pod_recovery": 0.9,
+    "autoscaling": 0.9,
+    "incident": 0.8,
+}
 
-_MAX_REWARD_PER_STEP = MAX_TOKENS * 0.1
-MAX_TOTAL_REWARD = MAX_STEPS * _MAX_REWARD_PER_STEP
+MAX_STALL_REPEATS = 4
+REWARD_EPSILON = 1e-9
+
+MAX_STEPS_BY_TASK = {
+    "pod_recovery": 15,
+    "autoscaling": 20,
+    "incident": 30,
+}
+
+GRADERS: Dict[str, Callable[[Dict[str, Any], int, int], float]] = {
+    "pod_recovery": grade_pod_recovery,
+    "autoscaling": grade_autoscaling,
+    "incident": grade_incident,
+}
 
 SYSTEM_PROMPT = textwrap.dedent(
     """
-    You are an agent interacting with an Kubernetes-like simulation environment.
+    You are a Kubernetes incident-response agent.
+    Return ONLY valid JSON for one action with this schema:
+    {
+            "action_type": "scale|delete_pod|patch|rollout_restart|set_hpa|drain_node|describe|wait",
+      "deployment": "... optional ...",
+      "replicas": 1,
+      "pod_name": "...",
+      "resource_type": "deployment|pod|node|service|configmap|hpa",
+      "name": "...",
+      "patch": {},
+      "min_replicas": 1,
+      "max_replicas": 5,
+      "cpu_target_percent": 70,
+      "node_name": "..."
+    }
+    Do not include markdown, prose, or code fences.
     """
 ).strip()
 
@@ -91,22 +145,150 @@ def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> No
     print(f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={rewards_str}", flush=True)
 
 
-def build_user_prompt(step: int, last_echoed: str, last_reward: float, history: List[str]) -> str:
+def _to_dict(obj: Any) -> Dict[str, Any]:
+    if hasattr(obj, "model_dump"):
+        return obj.model_dump()
+    if isinstance(obj, dict):
+        return obj
+    return vars(obj)
+
+
+def _observation_summary(observation: Any) -> str:
+    obs = _to_dict(observation)
+    pods = obs.get("pods", [])
+    deployments = obs.get("deployments", [])
+    events = obs.get("events", [])
+
+    pod_status_counts: Dict[str, int] = {}
+    for pod in pods:
+        status = pod.get("status", "Unknown")
+        pod_status_counts[status] = pod_status_counts.get(status, 0) + 1
+
+    deployment_lines = []
+    for dep in deployments:
+        deployment_lines.append(
+            f"{dep.get('name')}: desired={dep.get('desired_replicas', 0)} available={dep.get('available_replicas', 0)}"
+        )
+
+    recent_events = [
+        f"{e.get('type', 'Normal')}/{e.get('reason', '')}: {e.get('message', '')}"
+        for e in events[-5:]
+    ]
+
+    return textwrap.dedent(
+        f"""
+        Objective: {obs.get('objective', '')}
+        Step: {obs.get('step', 0)}
+        Pod status counts: {pod_status_counts}
+        Deployments:
+        {chr(10).join(deployment_lines) if deployment_lines else 'None'}
+        Recent events:
+        {chr(10).join(recent_events) if recent_events else 'None'}
+        """
+    ).strip()
+
+
+def build_user_prompt(task_name: str, step: int, observation: Any, history: List[str]) -> str:
     history_block = "\n".join(history[-4:]) if history else "None"
     return textwrap.dedent(
         f"""
+        Task: {task_name}
         Step: {step}
-        Last echoed message: {last_echoed!r}
-        Last reward: {last_reward:.2f}
+        Current cluster summary:
+        {_observation_summary(observation)}
         Previous steps:
         {history_block}
-        Send your next message.
+        Return one valid next action as pure JSON.
         """
     ).strip()
 
 
-def get_model_message(client: OpenAI, step: int, last_echoed: str, last_reward: float, history: List[str]) -> str:
-    user_prompt = build_user_prompt(step, last_echoed, last_reward, history)
+def _safe_json_action(text: str) -> Optional[Dict[str, Any]]:
+    try:
+        return json.loads(text)
+    except json.JSONDecodeError:
+        start = text.find("{")
+        end = text.rfind("}")
+        if start != -1 and end != -1 and end > start:
+            try:
+                return json.loads(text[start : end + 1])
+            except json.JSONDecodeError:
+                return None
+    return None
+
+
+def _heuristic_action(task_name: str, observation: Any) -> Dict[str, Any]:
+    obs = _to_dict(observation)
+    pods = obs.get("pods", [])
+
+    if task_name == "pod_recovery":
+        crashloop = [p for p in pods if p.get("deployment") == "frontend" and p.get("status") == "CrashLoopBackOff"]
+        if crashloop:
+            return {"action_type": "rollout_restart", "deployment": "frontend"}
+        return {"action_type": "describe", "resource_type": "deployment", "name": "frontend"}
+
+    if task_name == "autoscaling":
+        return {
+            "action_type": "set_hpa",
+            "deployment": "backend",
+            "min_replicas": 2,
+            "max_replicas": 6,
+            "cpu_target_percent": 70,
+        }
+
+    return {"action_type": "rollout_restart", "deployment": "auth-service"}
+
+
+def _normalize_action(action: Dict[str, Any]) -> Dict[str, Any]:
+    action_type = action.get("action_type", "describe")
+    if isinstance(action_type, str):
+        action_type = {
+            "set_hpas": "set_hpa",
+            "hpa": "set_hpa",
+            "restart_rollout": "rollout_restart",
+            "noop": "wait",
+            "no_op": "wait",
+            "pause": "wait",
+            "sleep": "wait",
+        }.get(action_type.strip().lower(), action_type.strip().lower())
+    else:
+        action_type = "describe"
+    normalized: Dict[str, Any] = {"action_type": action_type}
+
+    allowed_fields = {
+        "deployment",
+        "replicas",
+        "pod_name",
+        "resource_type",
+        "name",
+        "patch",
+        "min_replicas",
+        "max_replicas",
+        "cpu_target_percent",
+        "node_name",
+    }
+    for field in allowed_fields:
+        if field in action and action[field] is not None:
+            normalized[field] = action[field]
+
+    defaults_by_type = {
+        "describe": {"resource_type": "deployment", "name": "frontend"},
+        "scale": {"deployment": "frontend", "replicas": 3},
+        "rollout_restart": {"deployment": "frontend"},
+        "delete_pod": {"pod_name": "frontend-unknown"},
+        "drain_node": {"node_name": "node-1"},
+        "patch": {"resource_type": "deployment", "name": "frontend", "patch": {}},
+        "set_hpa": {"deployment": "backend", "min_replicas": 2, "max_replicas": 6, "cpu_target_percent": 70},
+        "wait": {},
+    }
+    for k, v in defaults_by_type.get(action_type, {}).items():
+        normalized.setdefault(k, v)
+
+    return normalized
+
+
+def get_model_action(client: OpenAI, task_name: str, step: int, observation: Any, history: List[str]) -> Dict[str, Any]:
+    user_prompt = build_user_prompt(task_name, step, observation, history)
     try:
         completion = client.chat.completions.create(
             model=MODEL_NAME,
@@ -119,66 +301,97 @@ def get_model_message(client: OpenAI, step: int, last_echoed: str, last_reward:
             stream=False,
         )
         text = (completion.choices[0].message.content or "").strip()
-        return text if text else "hello"
+        parsed = _safe_json_action(text)
+        if isinstance(parsed, dict):
+            return _normalize_action(parsed)
+        return _heuristic_action(task_name, observation)
     except Exception as exc:
         print(f"[DEBUG] Model request failed: {exc}", flush=True)
-        return "hello"
+        return _heuristic_action(task_name, observation)
 
+def _close_env(env: Any) -> None:
+    maybe = env.close()
+    if inspect.isawaitable(maybe):
+        # CoEnv.sync() should provide sync close(), but support awaitables defensively.
+        try:
+            while True:
+                maybe.send(None)
+        except StopIteration:
+            pass
+
+
+def main() -> None:
+    if not API_KEY:
+        raise RuntimeError("Missing HF_TOKEN/API_KEY for OpenAI client.")
+    for TASK_NAME, BENCHMARK in zip(TASK_NAMES, BENCHMARKS):
+        client = OpenAI(base_url=LLM_BASE_URL, api_key=API_KEY)
+        max_steps = MAX_STEPS_BY_TASK.get(TASK_NAME, DEFAULT_MAX_STEPS)
+        grader = GRADERS.get(TASK_NAME, grade_pod_recovery)
+
+        env = CoEnv(base_url=ENV_URL).sync()
+
+        history: List[str] = []
+        rewards: List[float] = []
+        steps_taken = 0
+        score = 0.0
+        success = False
+        final_obs: Optional[Any] = None
+        episode_done = False
+        stalled = False
+        last_action_str: Optional[str] = None
+        consecutive_same_action = 0
+        last_reward: Optional[float] = None
+
+        log_start(task=TASK_NAME, env=BENCHMARK, model=MODEL_NAME)
 
-async def main() -> None:
-    client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)
-
-    env = await CoEnv.from_docker_image(IMAGE_NAME)
-
-    history: List[str] = []
-    rewards: List[float] = []
-    steps_taken = 0
-    score = 0.0
-    success = False
-
-    log_start(task=TASK_NAME, env=BENCHMARK, model=MODEL_NAME)
-
-    try:
-        result = await env.reset() # OpenENV.reset()
-        last_echoed = result.observation.echoed_message
-        last_reward = 0.0
+        try:
+            result = env.reset(task=TASK_NAME)
+            final_obs = result.observation
 
-        for step in range(1, MAX_STEPS + 1):
-            if result.done:
-                break
+            for step in range(1, max_steps + 1):
+                time.sleep(API_DELAY)
+                if result.done:
+                    break
 
-            message = get_model_message(client, step, last_echoed, last_reward, history)
+                action_payload = get_model_action(client, TASK_NAME, step, final_obs, history)
+                action = CoenvAction(**action_payload)
 
-            result = await env.step(CoenvAction(message=message))
-            obs = result.observation
+                result = env.step(action)
+                obs = result.observation
+                final_obs = obs
 
-            reward = result.reward or 0.0
-            done = result.done
-            error = None
+                reward = result.reward or 0.0
+                done = result.done
+                error = (obs.metadata or {}).get("error") if hasattr(obs, "metadata") else None
 
-            rewards.append(reward)
-            steps_taken = step
-            last_echoed = obs.echoed_message
-            last_reward = reward
+                rewards.append(reward)
+                steps_taken = step
+                episode_done = bool(done)
 
-            log_step(step=step, action=message, reward=reward, done=done, error=error)
+                action_str = json.dumps(action_payload, separators=(",", ":"))
+                log_step(step=step, action=action_str, reward=reward, done=done, error=error)
 
-            history.append(f"Step {step}: {message!r} -> reward {reward:+.2f}")
+                history.append(f"Step {step}: {action_str} -> reward {reward:+.2f}")
 
-            if done:
-                break
+                if done:
+                    break
 
-        score = sum(rewards) / MAX_TOTAL_REWARD if MAX_TOTAL_REWARD > 0 else 0.0
-        score = min(max(score, 0.0), 1.0)  # clamp to [0, 1]
-        success = score >= SUCCESS_SCORE_THRESHOLD
+            world_state = _to_dict(final_obs) if final_obs is not None else {}
+            score = grader(world_state, steps_taken, max_steps)
+            score = min(max(score, 0.0), 1.0)
+            success = (
+                episode_done
+                and not stalled
+                and steps_taken > 0
+            )
 
-    finally:
-        try:
-            await env.close()
-        except Exception as e:
-            print(f"[DEBUG] env.close() error (container cleanup): {e}", flush=True)
-        log_end(success=success, steps=steps_taken, score=score, rewards=rewards)
+        finally:
+            try:
+                _close_env(env)
+            except Exception as e:
+                print(f"[DEBUG] env.close() error (container cleanup): {e}", flush=True)
+            log_end(success=success, steps=steps_taken, score=score, rewards=rewards)
 
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    main()

models.py

@@ -48,6 +48,7 @@ class CoenvAction(Action):
         "set_hpa",
         "drain_node",
         "describe",
+        "wait",
     ] = Field(..., description="Type of action to execute")
 
     deployment: Optional[str] = Field(default=None)

pyproject.toml

@@ -31,6 +31,7 @@ dependencies = [
 [project.optional-dependencies]
 dev = [
     "pytest>=8.0.0",
+    "pytest-asyncio>=0.23.0",
     "pytest-cov>=4.0.0",
 ]
 

server/actions/__init__.py

@@ -5,6 +5,7 @@ from .rollout_action import RolloutRestartAction
 from .hpa_action import SetHPAAction
 from .drain_action import DrainNodeAction
 from .describe_action import DescribeAction
+from .wait_action import WaitAction
 from typing import Union, Any, Dict, Literal
 
 KubeAction = Union[
@@ -14,10 +15,11 @@ KubeAction = Union[
     RolloutRestartAction,
     SetHPAAction,
     DrainNodeAction,
-    DescribeAction
+    DescribeAction,
+    WaitAction,
 ]
 
-ActionType = Literal["scale", "patch", "delete_pod", "rollout_restart", "set_hpa", "drain_node", "describe"]
+ActionType = Literal["scale", "patch", "delete_pod", "rollout_restart", "set_hpa", "drain_node", "describe", "wait"]
 
 
 def parse_action(data: Dict[str, Any]) -> KubeAction:
@@ -36,6 +38,7 @@ def parse_action(data: Dict[str, Any]) -> KubeAction:
         "set_hpa": SetHPAAction,
         "drain_node": DrainNodeAction,
         "describe": DescribeAction,
+        "wait": WaitAction,
     }
     
     action_class = action_map.get(action_type)
@@ -53,6 +56,7 @@ __all__ = [
     "SetHPAAction",
     "DrainNodeAction",
     "DescribeAction",
+    "WaitAction",
     "KubeAction",
     "parse_action",
 ]

server/actions/wait_action.py

@@ -0,0 +1,6 @@
+from pydantic import BaseModel
+from typing import Literal
+
+
+class WaitAction(BaseModel):
+    action_type: Literal["wait"] = "wait"

server/coenv_environment.py

@@ -596,7 +596,7 @@ class World:
         """Reset the world state and optionally inject a failure condition"""
         self.reset_to_healthy()
         if condition:
-            condition.inject(self)
+            condition.inject()
         return self.get_observation()
     
     def get_observation(self, objective: str = "Maintain cluster health"):

server/executor.py

@@ -9,6 +9,7 @@ from server.actions import (
     SetHPAAction,
     DrainNodeAction,
     DescribeAction,
+    WaitAction,
 )
 from server.models import ClusterObservation
 
@@ -35,6 +36,8 @@ def execute(action: KubeAction, world) -> ExecutionResult:
         return _execute_drain_node(action, world)
     elif isinstance(action, DescribeAction):
         return _execute_describe(action, world)
+    elif isinstance(action, WaitAction):
+        return _execute_wait(world)
     else:
         raise ValueError(f"Unknown action type: {type(action)}")
 
@@ -113,3 +116,12 @@ def _execute_describe(action: DescribeAction, world) -> ExecutionResult:
         tick_advanced=False,
         describe_detail=detail
     )
+
+
+def _execute_wait(world) -> ExecutionResult:
+    world.tick()
+    return ExecutionResult(
+        observation=world.get_observation(),
+        action_applied="Waited one simulation tick",
+        tick_advanced=True,
+    )

server/simulation_service.py

@@ -6,7 +6,7 @@ and reward/completion logic so server app wiring stays thin.
 
 from __future__ import annotations
 
-from typing import Dict, Any
+from typing import Dict, Any, Optional
 import json
 import os
 from openenv.core.env_server.interfaces import Environment
@@ -115,29 +115,111 @@ def calculate_reward(world: World, task_id: str) -> float:
     return 0.0
 
 
-def check_task_complete(world: World, task_id: str) -> bool:
-    """Check if task objective is complete."""
+def _collect_task_metrics(world: World) -> Dict[str, Any]:
+    """Collect state metrics used by completion logic."""
+    pods = world.get_pods()
+    deployments = world.get_deployments() if hasattr(world, "get_deployments") else []
+    hpas = world.get_hpas() if hasattr(world, "get_hpas") else []
+
+    def _deployment_running_ratio(name: str) -> float:
+        dep_pods = [p for p in pods if p.deployment == name]
+        if not dep_pods:
+            return 0.0
+        running = [p for p in dep_pods if p.status == "Running"]
+        return len(running) / len(dep_pods)
+
+    def _deployment_unstable_count(name: str, restart_threshold: int = 5) -> int:
+        dep_pods = [p for p in pods if p.deployment == name]
+        unstable = [
+            p for p in dep_pods
+            if p.status != "Running"
+            or p.status == "CrashLoopBackOff"
+            or getattr(p, "restarts", 0) >= restart_threshold
+        ]
+        return len(unstable)
+
+    key_services = ["auth-service", "api-gateway", "frontend"]
+    incident_unhealthy_services = 0
+    for svc in key_services:
+        if _deployment_running_ratio(svc) < 0.8:
+            incident_unhealthy_services += 1
+
+    backend_hpa = next((h for h in hpas if h.name == "backend-hpa"), None)
+    backend_hpa_ok = (
+        backend_hpa is not None
+        and backend_hpa.min_replicas >= 2
+        and backend_hpa.max_replicas >= 6
+        and backend_hpa.cpu_target_percent <= 70
+    )
+
+    backend_dep = next((d for d in deployments if d.name == "backend"), None)
+    backend_available_ratio = 0.0
+    if backend_dep is not None and backend_dep.desired_replicas > 0:
+        backend_available_ratio = backend_dep.available_replicas / backend_dep.desired_replicas
+
+    return {
+        "frontend_unstable": _deployment_unstable_count("frontend"),
+        "frontend_running_ratio": _deployment_running_ratio("frontend"),
+        "backend_unstable": _deployment_unstable_count("backend"),
+        "backend_running_ratio": _deployment_running_ratio("backend"),
+        "backend_hpa_ok": backend_hpa_ok,
+        "backend_available_ratio": backend_available_ratio,
+        "incident_unhealthy_services": incident_unhealthy_services,
+        "incident_key_unstable": sum(_deployment_unstable_count(svc) for svc in key_services),
+    }
+
+
+def check_task_complete(world: World, task_id: str, baseline_metrics: Optional[Dict[str, Any]] = None) -> bool:
+    """Check if task objective is complete via observable state recovery."""
+    metrics = _collect_task_metrics(world)
+    baseline = baseline_metrics or {}
+    has_baseline = bool(baseline)
+
     if task_id == "pod_recovery":
-        pods = world.get_pods()
-        frontend_pods = [p for p in pods if p.deployment == "frontend"]
-        running = [p for p in frontend_pods if p.status == "Running"]
-        return len(frontend_pods) > 0 and len(running) == len(frontend_pods)
+        if not has_baseline:
+            return metrics["frontend_unstable"] == 0 and metrics["frontend_running_ratio"] >= 1.0
+        had_problem = baseline.get("frontend_unstable", 0) > 0
+        recovered = metrics["frontend_unstable"] == 0 and metrics["frontend_running_ratio"] >= 1.0
+        improved = metrics["frontend_unstable"] < baseline.get("frontend_unstable", 0)
+        return had_problem and recovered and improved
 
     if task_id == "autoscaling":
-        pods = world.get_pods()
-        backend_pods = [p for p in pods if p.deployment == "backend"]
-        running = [p for p in backend_pods if p.status == "Running"]
-        return len(backend_pods) >= 2 and len(running) >= 2
+        if not has_baseline:
+            return (
+                metrics["backend_unstable"] == 0
+                and metrics["backend_running_ratio"] >= 1.0
+                and metrics["backend_available_ratio"] >= 1.0
+                and metrics["backend_hpa_ok"]
+            )
+        had_problem = baseline.get("backend_unstable", 0) > 0
+        recovered = (
+            metrics["backend_unstable"] == 0
+            and metrics["backend_running_ratio"] >= 1.0
+            and metrics["backend_available_ratio"] >= 1.0
+        )
+        improved = metrics["backend_unstable"] < baseline.get("backend_unstable", 0)
+        # For autoscaling, both state recovery and effective HPA policy must be visible.
+        return had_problem and recovered and improved and metrics["backend_hpa_ok"]
 
     if task_id == "incident":
-        pods = world.get_pods()
-        key_services = ["auth-service", "api-gateway", "frontend"]
-        for svc in key_services:
-            svc_pods = [p for p in pods if p.deployment == svc]
-            running = [p for p in svc_pods if p.status == "Running"]
-            if svc_pods and len(running) < len(svc_pods) * 0.8:
-                return False
-        return True
+        if not has_baseline:
+            return (
+                metrics["incident_unhealthy_services"] == 0
+                and metrics["incident_key_unstable"] == 0
+            )
+        had_problem = (
+            baseline.get("incident_unhealthy_services", 0) > 0
+            or baseline.get("incident_key_unstable", 0) > 0
+        )
+        recovered = (
+            metrics["incident_unhealthy_services"] == 0
+            and metrics["incident_key_unstable"] == 0
+        )
+        improved = (
+            metrics["incident_unhealthy_services"] < baseline.get("incident_unhealthy_services", 0)
+            or metrics["incident_key_unstable"] < baseline.get("incident_key_unstable", 0)
+        )
+        return had_problem and recovered and improved
 
     return False
 
@@ -151,13 +233,32 @@ class CoenvEnvironment(Environment):
         self.world = World(self.config, seed=self.config.get("seed"))
         self.current_task = "pod_recovery"
         self.current_objective = get_objective_for_task(self.current_task)
+        self._baseline_metrics: Dict[str, Any] = {}
 
     def reset(self, task: str = "pod_recovery", **_: Any) -> CoenvObservation:
         """Reset simulator state for the selected task and return initial observation."""
         self.current_task = task
         self.current_objective = get_objective_for_task(task)
         condition = get_condition_for_task(task, self.world, self.config)
-        self.world.reset(condition)
+
+        # Inject deterministic, task-specific failures so episodes don't start
+        # in an already-solved state.
+        self.world.reset_to_healthy()
+        if condition is not None:
+            if task == "pod_recovery":
+                condition.inject(target_deployment="frontend", failure_rate=0.8)
+            elif task == "autoscaling":
+                condition.inject(target_deployment="backend", failure_rate=0.8)
+            elif task == "incident":
+                condition.inject(root_cause_service="auth-service", failure_probability=0.8)
+                try:
+                    from .conditions.crash_loop import CrashLoopCondition
+                except ImportError:
+                    from conditions.crash_loop import CrashLoopCondition
+                # Ensure cascading impact reaches key downstream services.
+                CrashLoopCondition(self.world, self.config).inject(target_deployment="api-gateway", failure_rate=0.7)
+                CrashLoopCondition(self.world, self.config).inject(target_deployment="frontend", failure_rate=0.5)
+            self._baseline_metrics = _collect_task_metrics(self.world)
         return self._observation(done=False, reward=0.0, info={"task": task})
 
     def step(self, action: CoenvAction, **_: Any) -> CoenvObservation:
@@ -208,6 +309,9 @@ class CoenvEnvironment(Environment):
                 info["described"] = f"{resource_type}/{name}"
                 info["describe_detail"] = self.world.describe(resource_type, name)
 
+            elif action.action_type == "wait":
+                info["waited"] = True
+
             else:
                 info["error"] = f"Unknown action type: {action.action_type}"
 
@@ -218,12 +322,10 @@ class CoenvEnvironment(Environment):
 
         reward = calculate_reward(self.world, self.current_task)
 
-        done = False
+        done = check_task_complete(self.world, self.current_task, self._baseline_metrics)
         max_steps = self.config.get("tasks", {}).get(self.current_task, {}).get("max_steps", 15)
-        if self.world.step_count >= max_steps:
-            done = True
-        if check_task_complete(self.world, self.current_task):
-            done = True
+        if self.world.step_count >= max_steps and not done:
+            info["truncated"] = True
 
         return self._observation(done=done, reward=reward, info=info)
     
@@ -231,7 +333,7 @@ class CoenvEnvironment(Environment):
     def state(self) -> CoenvState:
         """Return current observation without applying an action."""
         reward = calculate_reward(self.world, self.current_task)
-        done = check_task_complete(self.world, self.current_task)
+        done = check_task_complete(self.world, self.current_task, self._baseline_metrics)
         return CoenvState(
             episode_id=self.episode_id,
             step_count=self.world.step_count

server/validator.py

@@ -8,6 +8,7 @@ from server.actions import (
     SetHPAAction,
     DrainNodeAction,
     DescribeAction,
+    WaitAction,
 )
 
 
@@ -26,6 +27,8 @@ def validate(action: KubeAction, world_state: Dict[str, Any]) -> Optional[str]:
         return _validate_drain_node(action, world_state)
     elif isinstance(action, DescribeAction):
         return _validate_describe(action, world_state)
+    elif isinstance(action, WaitAction):
+        return None
     return None
 
 

tests/test_actions.py

@@ -8,6 +8,7 @@ from server.actions import (
     SetHPAAction,
     DrainNodeAction,
     DescribeAction,
+    WaitAction,
     parse_action,
 )
 
@@ -202,6 +203,11 @@ class TestParseAction:
         assert isinstance(action, DescribeAction)
         assert action.name == "frontend"
 
+    def test_parse_wait_action(self):
+        raw = {"action_type": "wait"}
+        action = parse_action(raw)
+        assert isinstance(action, WaitAction)
+
     def test_parse_unknown_action_type(self):
         with pytest.raises(ValueError, match="Unknown action_type"):
             parse_action({"action_type": "unknown_action"})

tests/test_executor.py

@@ -8,6 +8,7 @@ from server.actions import (
     SetHPAAction,
     DrainNodeAction,
     DescribeAction,
+    WaitAction,
 )
 from server.executor import execute
 from server.models import ClusterObservation
@@ -172,3 +173,14 @@ class TestExecutorDescribe:
 
         assert result.describe_detail is not None
         assert result.describe_detail["type"] == "deployment"
+
+
+class TestExecutorWait:
+    def test_wait_ticks_without_other_world_mutations(self):
+        mock_world = MockWorld()
+        action = WaitAction(action_type="wait")
+        result = execute(action, mock_world)
+
+        assert mock_world.tick_called is True
+        assert result.tick_advanced is True
+        assert result.action_applied == "Waited one simulation tick"

tests/test_inference.py

@@ -0,0 +1,19 @@
+from inference import _normalize_action
+
+
+def test_normalize_action_maps_set_hpas_to_set_hpa():
+    action = _normalize_action({"action_type": "set_hpas", "deployment": "backend"})
+
+    assert action["action_type"] == "set_hpa"
+    assert action["deployment"] == "backend"
+    assert action["min_replicas"] == 2
+    assert action["max_replicas"] == 6
+    assert action["cpu_target_percent"] == 70
+
+
+def test_normalize_action_non_string_type_defaults_to_describe():
+    action = _normalize_action({"action_type": ["set_hpa"]})
+
+    assert action["action_type"] == "describe"
+    assert action["resource_type"] == "deployment"
+    assert action["name"] == "frontend"

tests/test_simulation_service.py

@@ -92,6 +92,9 @@ def test_environment_step_scale_and_describe_paths():
     assert "described" in describe_obs.metadata
     assert "describe_detail" in describe_obs.metadata
 
+    wait_obs = env.step(CoenvAction(action_type="wait"))
+    assert wait_obs.metadata.get("waited") is True
+
 
 def test_environment_step_exception_is_captured_in_metadata(monkeypatch):
     env = CoenvEnvironment()