init

inference.py

@@ -1,314 +1,995 @@
 """
-Baseline inference script.
-
-Uses an LLM (via OpenAI-compatible API) to play through all 3 incident
-scenarios. The conversation history acts as a soft belief tracker —
-the LLM accumulates evidence across steps.
-
-stdout format: [START], [STEP], [END] blocks with exact field names
-as required by the OpenEnv automated evaluator.
+inference.py — SRE Incident Response + Code Attribution Agent
+
+Two execution modes:
+  baseline  — flat P1-only loop (comparison baseline, no orchestrator)
+  unified   — orchestrator → ops subagent → code subagent (research mode)
+
+Three LLM backends (set BACKEND env var):
+  local     — load checkpoint from LOCAL_MODEL_PATH using transformers + Unsloth
+  vllm      — serve checkpoint via vLLM (faster, needs vllm installed + server running)
+  api       — OpenAI-compatible HTTP API (baseline comparisons only)
+
+stdout contract (OpenEnv evaluator parses this — do not change field names):
+  [START] task=<n>
+  [STEP]  step=<n> phase=<1|2> action=<json>
+  [END]   task=<n> score=<f> reward=<f> steps=<n>
+
+Environment variables:
+  BACKEND          local | vllm | api          (default: local)
+  LOCAL_MODEL_PATH path to checkpoint dir      (default: ./checkpoint)
+  LOAD_IN_4BIT     1 | 0                       (default: 1)
+  VLLM_BASE_URL    vLLM server URL             (default: http://localhost:8001)
+  API_BASE_URL     OpenAI-compatible API URL   (api backend only)
+  API_KEY          API key                     (api backend only)
+  MODEL_NAME       model name string           (api / vllm backends)
+  ENV_BASE_URL     OpenEnv server URL          (default: http://localhost:8000)
+  MODE             baseline | unified          (default: unified)
+  COLLECT          1 to write trajectory JSON  (default: 0)
+  MAX_NEW_TOKENS   token budget per call       (default: 512)
+  TEMPERATURE      sampling temperature        (default: 0.3)
+  ORCH_TEMPERATURE orchestrator temperature    (default: 0.1)
 """
 
 from __future__ import annotations
 
 import json
 import os
+import re
 import sys
-import time
 import traceback
+from abc import ABC, abstractmethod
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
 from typing import Any, Dict, List, Optional
 
 import requests
-from openai import OpenAI
 
 
-# ------------------------------------------------------------------
+# ══════════════════════════════════════════════════════════════════
 # Config
-# ------------------------------------------------------------------
+# ══════════════════════════════════════════════════════════════════
+
+BACKEND          = os.environ.get("BACKEND",          "local")
+LOCAL_MODEL_PATH = os.environ.get("LOCAL_MODEL_PATH", "./checkpoint")
+VLLM_BASE_URL    = os.environ.get("VLLM_BASE_URL",    "http://localhost:8001")
+API_BASE_URL     = os.environ.get("API_BASE_URL",     "")
+API_KEY          = os.environ.get("API_KEY",          "no-key")
+MODEL_NAME       = os.environ.get("MODEL_NAME",       "checkpoint")
+ENV_BASE_URL     = os.environ.get("ENV_BASE_URL",     "http://localhost:8000")
+MODE             = os.environ.get("MODE",             "unified")
+COLLECT          = os.environ.get("COLLECT",          "0") == "1"
+
+MAX_NEW_TOKENS   = int(os.environ.get("MAX_NEW_TOKENS",    "512"))
+TEMPERATURE      = float(os.environ.get("TEMPERATURE",     "0.3"))
+ORCH_TEMPERATURE = float(os.environ.get("ORCH_TEMPERATURE","0.1"))
 
-ENV_BASE_URL = os.environ.get("ENV_BASE_URL", "http://localhost:8000")
-MODEL_NAME = os.environ.get("MODEL_NAME", "gpt-4o-mini")
-HF_TOKEN = os.environ.get("HF_TOKEN", "")
-API_KEY = os.environ.get("API_KEY", "")
-MAX_STEPS = 20
-TEMPERATURE = 0.3
+MAX_P1_STEPS = 20
+MAX_P2_STEPS = 15
 
 
-# ------------------------------------------------------------------
-# System prompt — Layer 3: the LLM acts as an SRE
-# ------------------------------------------------------------------
+# ══════════════════════════════════════════════════════════════════
+# Prompts
+# ══════════════════════════════════════════════════════════════════
 
-SYSTEM_PROMPT = """You are an expert Site Reliability Engineer (SRE) responding to a production incident.
+OPS_SYSTEM_PROMPT = """You are an expert Site Reliability Engineer (SRE) responding to a production incident.
 
-You are interacting with a simulated microservices infrastructure through an environment API.
-Your goal is to:
-1. DIAGNOSE the root cause of the incident
-2. REMEDIATE the issue (fix it)
-3. DECLARE the root cause when confident
+Your goal:
+  1. DIAGNOSE the root cause from observable symptoms only
+  2. REMEDIATE by acting on the correct service
+  3. DECLARE when you are confident
 
-## Available Actions
-You must respond with a single JSON object containing your chosen action:
+## Action schema — respond with ONE JSON object per turn
 
-DIAGNOSTIC (information gathering):
-- {"action_type": "view_alerts"} — See all firing alerts
-- {"action_type": "query_logs", "target_service": "<name>", "parameters": {"level": "ERROR"}} — Query logs
-- {"action_type": "check_metrics", "target_service": "<name>"} — Get metric timeseries
-- {"action_type": "check_dependencies", "target_service": "<name>"} — View dependency graph
-- {"action_type": "check_deploy_history", "target_service": "<name>"} — Recent deploys
-- {"action_type": "run_health_check", "target_service": "<name>"} — Ping a service
+Diagnostic (no state mutation):
+  {"action_type": "view_alerts"}
+  {"action_type": "query_logs",           "target_service": "<svc>", "parameters": {"level": "ERROR"}}
+  {"action_type": "check_metrics",        "target_service": "<svc>"}
+  {"action_type": "check_dependencies",   "target_service": "<svc>"}
+  {"action_type": "check_deploy_history", "target_service": "<svc>"}
+  {"action_type": "run_health_check",     "target_service": "<svc>"}
 
-REMEDIATION (fix actions):
-- {"action_type": "restart_service", "target_service": "<name>"} — Restart a service
-- {"action_type": "rollback_deploy", "target_service": "<name>"} — Rollback to previous deploy
-- {"action_type": "scale_service", "target_service": "<name>", "parameters": {"replicas": 5}} — Scale replicas
+Remediation (mutates state):
+  {"action_type": "restart_service", "target_service": "<svc>"}
+  {"action_type": "rollback_deploy", "target_service": "<svc>"}
+  {"action_type": "scale_service",   "target_service": "<svc>", "parameters": {"replicas": 5}}
 
-DECLARATION:
-- {"action_type": "declare_root_cause", "parameters": {"root_cause": "<your diagnosis>"}}
+Terminal:
+  {"action_type": "declare_root_cause", "parameters": {"root_cause": "<diagnosis>"}}
 
-## Available services: api_gateway, auth, orders, payment, cache, database, queue
+Services: api_gateway, auth, orders, payment, cache, database, queue
 
 ## Strategy
-1. Start by viewing alerts to understand the scope
-2. Check metrics and logs for the most affected services
-3. Check dependency graphs to trace upstream causes
-4. Check deploy history for recently changed services
-5. Apply remediation to the root cause service FIRST
-6. Declare root cause when confident
-
-IMPORTANT: Respond with ONLY a valid JSON object. No explanation, no markdown, just the JSON action.
+  - view_alerts first to understand scope
+  - check_metrics + query_logs on the highest-severity service
+  - check_dependencies to trace upstream root causes
+  - check_deploy_history before any rollback
+  - remediate the ROOT cause service first
+  - declare when confident — do not delay unnecessarily
+
+IMPORTANT: Output ONLY valid JSON. No markdown, no explanation.
 """
 
+ORCHESTRATOR_PROMPT = """You are the orchestrator of a two-phase SRE incident response system.
+
+After each ops agent action you assess the current belief state and decide whether to
+continue Phase 1 (gather more evidence) or transition to Phase 2 (codebase attribution).
+
+Rules:
+  - transition only when suspected_service is identified with reasonable confidence
+  - do NOT transition just because steps are high — bad evidence is worse than no transition
+  - evidence_gaps must list specific missing checks (e.g. "deploy_history_unchecked")
+  - estimated_p2_cost reflects how broad the codebase search will need to be
+
+Output ONLY this XML block — no other text:
+
+<belief_state>
+  <suspected_service>{service name or "unknown"}</suspected_service>
+  <suspected_fault_class>{memory_leak|config_change|deadlock|resource_exhaustion|cascading|none}</suspected_fault_class>
+  <service_confidence>{0.00 to 1.00}</service_confidence>
+  <fault_confidence>{0.00 to 1.00}</fault_confidence>
+  <evidence_gaps>{comma-separated list or "none"}</evidence_gaps>
+  <estimated_p2_cost>{low|medium|high}</estimated_p2_cost>
+  <decision>{continue|transition}</decision>
+  <reasoning>{one concise sentence}</reasoning>
+</belief_state>
+"""
 
-# ------------------------------------------------------------------
-# Environment client (direct HTTP)
-# ------------------------------------------------------------------
+CODE_AGENT_PROMPT = """You are a senior software engineer performing code attribution for a production incident.
 
-class EnvClient:
-    def __init__(self, base_url: str):
-        self.base_url = base_url.rstrip("/")
-        self.session = requests.Session()
+Runtime diagnosis handed off from SRE phase:
+  Faulty service : {service}
+  Fault class    : {fault_class}
+  Bad deploy SHA : {commit_sha}
+  Confidence     : service={service_confidence}  fault={fault_confidence}
 
-    def reset(self, task_name: str, seed: int = 42) -> Dict[str, Any]:
-        resp = self.session.post(f"{self.base_url}/reset", json={
-            "task_name": task_name, "seed": seed})
-        resp.raise_for_status()
-        return resp.json()
+Your job: explore the codebase snapshot, find the exact change that caused the incident,
+then either propose a patch or declare that no code change is needed.
 
-    def step(self, action: Dict[str, Any]) -> Dict[str, Any]:
-        resp = self.session.post(f"{self.base_url}/step", json=action)
-        resp.raise_for_status()
-        return resp.json()
+## Action schema — respond with ONE JSON object per turn
 
-    def state(self) -> Dict[str, Any]:
-        resp = self.session.get(f"{self.base_url}/state")
-        resp.raise_for_status()
-        return resp.json()
+Exploration:
+  {"action_type": "list_dir",      "parameters": {"path": "."}}
+  {"action_type": "read_file",     "parameters": {"path": "<rel_path>"}}
+  {"action_type": "search_code",   "parameters": {"query": "<text>", "file_pattern": "*.py"}}
+  {"action_type": "get_git_log",   "parameters": {"path": "<rel_path>", "n_commits": 5}}
+  {"action_type": "get_file_diff", "parameters": {"commit_sha": "<sha>", "path": "<rel_path>"}}
 
+Terminal:
+  {"action_type": "propose_patch",     "parameters": {"diff": "<unified diff>", "explanation": "<reason>"}}
+  {"action_type": "declare_no_change", "parameters": {"reason": "<why no code change is needed>"}}
 
-# ------------------------------------------------------------------
-# LLM agent
-# ------------------------------------------------------------------
+## Strategy
+  1. list_dir to understand repo structure
+  2. get_git_log on the bad commit SHA to see which files changed
+  3. read_file on each changed file to understand the bug
+  4. propose_patch with a minimal correct unified diff
+  5. If symptoms are infra-only (config, scaling) with no bad code: declare_no_change
 
-def create_openai_client() -> OpenAI:
-    """Create OpenAI client with appropriate config."""
-    api_key = API_KEY or HF_TOKEN or "no-key"
-    base_url = os.environ.get("API_BASE_URL")
+IMPORTANT: Output ONLY valid JSON. No markdown, no explanation.
+"""
 
-    # If using HF inference endpoint, set base_url
-    if HF_TOKEN and not API_KEY and not base_url:
-        base_url = f"https://api-inference.huggingface.co/models/{MODEL_NAME}/v1"
 
-    return OpenAI(api_key=api_key, base_url=base_url)
+# ══════════════════════════════════════════════════════════════════
+# LLM backend abstraction
+# ══════════════════════════════════════════════════════════════════
 
+Message = Dict[str, str]  # {"role": "system"|"user"|"assistant", "content": str}
 
-def parse_llm_action(response_text: str) -> Dict[str, Any]:
-    """Extract JSON action from LLM response. Handles markdown wrapping."""
-    text = response_text.strip()
 
-    # Strip markdown code fences if present
-    if text.startswith("```"):
-        lines = text.split("\n")
-        lines = [l for l in lines if not l.strip().startswith("```")]
-        text = "\n".join(lines).strip()
+class LLMBackend(ABC):
+    """
+    Uniform interface over local checkpoint, vLLM, and API backends.
+    All call sites use backend.generate(messages, temperature, max_new_tokens).
+    Swapping backends requires only changing the BACKEND env var.
+    """
 
-    # Find JSON object
-    start = text.find("{")
-    end = text.rfind("}") + 1
-    if start >= 0 and end > start:
-        return json.loads(text[start:end])
+    @abstractmethod
+    def generate(
+        self,
+        messages:       List[Message],
+        temperature:    float = TEMPERATURE,
+        max_new_tokens: int   = MAX_NEW_TOKENS,
+    ) -> str:
+        """Return the assistant response text, stripped."""
+
+
+# ── Local checkpoint ──────────────────────────────────────────────
+
+class LocalModelBackend(LLMBackend):
+    """
+    Loads a HuggingFace checkpoint from LOCAL_MODEL_PATH.
+
+    Uses Unsloth when available for 2x faster inference with identical output.
+    Falls back to vanilla transformers if Unsloth is not installed.
+
+    The model loads once at construction and is reused across all episodes.
+    apply_chat_template handles the system/user/assistant turn format for
+    Qwen, Llama, Mistral and other chat models automatically.
+    """
+
+    def __init__(self, model_path: str, load_in_4bit: bool = True):
+        self.model_path   = model_path
+        self.load_in_4bit = load_in_4bit
+        self.model        = None
+        self.tokenizer    = None
+        self._load()
+
+    def _load(self) -> None:
+        _log(f"Loading checkpoint: {self.model_path}")
+
+        try:
+            from unsloth import FastLanguageModel
+            self.model, self.tokenizer = FastLanguageModel.from_pretrained(
+                model_name    = self.model_path,
+                max_seq_length= 8192,
+                load_in_4bit  = self.load_in_4bit,
+                dtype         = None,  # auto — bfloat16 on Ampere+
+            )
+            FastLanguageModel.for_inference(self.model)
+            _log(f"Backend: Unsloth  4bit={self.load_in_4bit}")
+
+        except ImportError:
+            import torch
+            from transformers import AutoModelForCausalLM, AutoTokenizer
+
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                self.model_path, trust_remote_code=True
+            )
+            self.model = AutoModelForCausalLM.from_pretrained(
+                self.model_path,
+                torch_dtype      = "auto",
+                device_map       = "auto",
+                trust_remote_code= True,
+            )
+            self.model.eval()
+            _log("Backend: transformers (Unsloth not found, using vanilla)")
+
+    def generate(
+        self,
+        messages:       List[Message],
+        temperature:    float = TEMPERATURE,
+        max_new_tokens: int   = MAX_NEW_TOKENS,
+    ) -> str:
+        import torch
+
+        text = self.tokenizer.apply_chat_template(
+            messages,
+            tokenize              = False,
+            add_generation_prompt = True,
+        )
+        inputs    = self.tokenizer(text, return_tensors="pt").to(self.model.device)
+        input_len = inputs["input_ids"].shape[-1]
+
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                max_new_tokens = max_new_tokens,
+                temperature    = temperature,
+                do_sample      = temperature > 0,
+                pad_token_id   = self.tokenizer.eos_token_id,
+            )
 
-    raise ValueError(f"Could not parse action from: {response_text[:200]}")
+        # Decode only newly generated tokens — strip the echoed prompt
+        new_tokens = outputs[0][input_len:]
+        return self.tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
 
 
-def summarize_observation(obs: Dict[str, Any]) -> str:
-    """Convert observation dict to a readable string for the LLM context."""
-    parts = []
-    parts.append(f"Incident: {obs.get('incident_summary', 'N/A')}")
-    parts.append(f"Severity: {obs.get('severity', 'N/A')}")
-    parts.append(f"Time: {obs.get('time_elapsed_minutes', 0)}/{obs.get('time_budget_minutes', 30)} min")
-    parts.append(f"Steps: {obs.get('steps_taken', 0)}/{obs.get('max_steps', 20)}")
-    parts.append(f"Reward: {obs.get('current_reward', 0)} (cumulative: {obs.get('cumulative_reward', 0)})")
+# ── vLLM backend ──────────────────────────────────────────────────
+
+class VLLMBackend(LLMBackend):
+    """
+    Calls a locally running vLLM server via its OpenAI-compatible endpoint.
+
+    Start the server with:
+      python -m vllm.entrypoints.openai.api_server \\
+        --model ./checkpoint --port 8001
+
+    No model loading here — the server handles it. Use this when running
+    rapid eval loops and model load time is a bottleneck.
+    """
+
+    def __init__(self, base_url: str, model_name: str):
+        try:
+            from openai import OpenAI as _OpenAI
+        except ImportError:
+            raise ImportError("pip install openai  (required for vllm backend)")
+        self._client = _OpenAI(api_key="vllm-local", base_url=base_url)
+        self._model  = model_name
+        _log(f"Backend: vLLM at {base_url}  model={model_name}")
+
+    def generate(
+        self,
+        messages:       List[Message],
+        temperature:    float = TEMPERATURE,
+        max_new_tokens: int   = MAX_NEW_TOKENS,
+    ) -> str:
+        resp = self._client.chat.completions.create(
+            model       = self._model,
+            messages    = messages,
+            temperature = temperature,
+            max_tokens  = max_new_tokens,
+        )
+        return (resp.choices[0].message.content or "").strip()
+
+
+# ── API backend ───────────────────────────────────────────────────
+
+class APIBackend(LLMBackend):
+    """
+    Wraps any OpenAI-compatible HTTP API.
+    Use only for baseline comparisons — not for checkpoint inference.
+    """
+
+    def __init__(self, api_key: str, base_url: Optional[str], model_name: str):
+        try:
+            from openai import OpenAI as _OpenAI
+        except ImportError:
+            raise ImportError("pip install openai  (required for api backend)")
+        kwargs: Dict[str, Any] = {"api_key": api_key}
+        if base_url:
+            kwargs["base_url"] = base_url
+        self._client = _OpenAI(**kwargs)
+        self._model  = model_name
+        _log(f"Backend: API  model={model_name}  base={base_url or 'openai'}")
+
+    def generate(
+        self,
+        messages:       List[Message],
+        temperature:    float = TEMPERATURE,
+        max_new_tokens: int   = MAX_NEW_TOKENS,
+    ) -> str:
+        resp = self._client.chat.completions.create(
+            model       = self._model,
+            messages    = messages,
+            temperature = temperature,
+            max_tokens  = max_new_tokens,
+        )
+        return (resp.choices[0].message.content or "").strip()
+
+
+# ── Factory ───────────────────────────────────────────────────────
+
+def build_backend() -> LLMBackend:
+    if BACKEND == "local":
+        return LocalModelBackend(
+            model_path   = LOCAL_MODEL_PATH,
+            load_in_4bit = os.environ.get("LOAD_IN_4BIT", "1") == "1",
+        )
+    if BACKEND == "vllm":
+        return VLLMBackend(base_url=VLLM_BASE_URL, model_name=MODEL_NAME)
+    if BACKEND == "api":
+        return APIBackend(
+            api_key    = API_KEY,
+            base_url   = API_BASE_URL or None,
+            model_name = MODEL_NAME,
+        )
+    raise ValueError(f"Unknown BACKEND={BACKEND!r}. Choose: local | vllm | api")
+
+
+# ══════════════════════════════════════════════════════════════════
+# Data containers
+# ══════════════════════════════════════════════════════════════════
+
+@dataclass
+class BeliefState:
+    suspected_service:    str   = "unknown"
+    suspected_fault_class: str  = "none"
+    service_confidence:   float = 0.0
+    fault_confidence:     float = 0.0
+    evidence_gaps:        str   = "none"
+    estimated_p2_cost:    str   = "unknown"
+    decision:             str   = "continue"
+    reasoning:            str   = ""
+
+    def confident_enough(self) -> bool:
+        """
+        Orchestrator stopping criterion.
+        Stage 4 GRPO training trains the model to emit the correct
+        <decision> tag — this method is therefore the learned policy
+        expressed as a single field check.
+        """
+        return self.decision == "transition"
+
+
+@dataclass
+class StepRecord:
+    step_number:  int
+    phase:        int
+    action:       Dict[str, Any]
+    reward:       float
+    obs_summary:  Dict[str, Any]
+    belief:       Optional[Dict[str, Any]] = None  # P1 only
+
+
+@dataclass
+class EpisodeRecord:
+    task_name:           str
+    mode:                str
+    seed:                int
+    p1_trajectory:       List[StepRecord] = field(default_factory=list)
+    p2_trajectory:       List[StepRecord] = field(default_factory=list)
+    belief_history:      List[Dict]       = field(default_factory=list)
+    declared_patch:      Optional[str]    = None
+    declared_no_change:  bool             = False
+    phase_transition_at: Optional[int]    = None
+    score:               float            = 0.0
+    cumulative_reward:   float            = 0.0
+    total_steps:         int              = 0
+
+
+# ══════════════════════════════════════════════════════════════════
+# Environment client
+# ══════════════════════════════════════════════════════════════════
 
+class EnvClient:
+    def __init__(self, base_url: str):
+        self.base_url = base_url.rstrip("/")
+        self.session  = requests.Session()
+
+    def reset(self, task_name: str, seed: int = 42) -> Dict[str, Any]:
+        r = self.session.post(
+            f"{self.base_url}/reset",
+            json={"task_name": task_name, "seed": seed},
+        )
+        r.raise_for_status()
+        return r.json()
+
+    def step(self, action: Dict[str, Any]) -> Dict[str, Any]:
+        r = self.session.post(f"{self.base_url}/step", json=action)
+        r.raise_for_status()
+        return r.json()
+
+    def state(self) -> Dict[str, Any]:
+        r = self.session.get(f"{self.base_url}/state")
+        r.raise_for_status()
+        return r.json()
+
+    def unified_score(
+        self,
+        declared_patch:     Optional[str],
+        declared_no_change: bool,
+        belief_history:     List[Dict],
+    ) -> Dict[str, float]:
+        try:
+            r = self.session.post(
+                f"{self.base_url}/score",
+                json={
+                    "declared_patch":     declared_patch,
+                    "declared_no_change": declared_no_change,
+                    "belief_history":     belief_history,
+                },
+            )
+            r.raise_for_status()
+            return r.json()
+        except Exception:
+            return {"final": 0.01}
+
+
+# ════════════════════════════════════════════════════════���═════════
+# Parsing helpers
+# ══════════════════════════════════════════════════════════════════
+
+def parse_action(text: str) -> Dict[str, Any]:
+    """
+    Extract JSON action from model output.
+    Local models sometimes wrap output in prose or markdown — we
+    defensively extract the first complete JSON object.
+    """
+    text = text.strip()
+    if text.startswith("```"):
+        text = "\n".join(
+            l for l in text.split("\n")
+            if not l.strip().startswith("```")
+        ).strip()
+    start = text.find("{")
+    end   = text.rfind("}") + 1
+    if start >= 0 and end > start:
+        return json.loads(text[start:end])
+    raise ValueError(f"No JSON in model output: {text[:300]}")
+
+
+def parse_belief(xml_text: str) -> BeliefState:
+    def _x(tag: str) -> str:
+        m = re.search(rf"<{tag}>(.*?)</{tag}>", xml_text, re.DOTALL)
+        return m.group(1).strip() if m else ""
+
+    return BeliefState(
+        suspected_service    = _x("suspected_service")     or "unknown",
+        suspected_fault_class= _x("suspected_fault_class") or "none",
+        service_confidence   = _safe_float(_x("service_confidence")),
+        fault_confidence     = _safe_float(_x("fault_confidence")),
+        evidence_gaps        = _x("evidence_gaps")         or "none",
+        estimated_p2_cost    = _x("estimated_p2_cost")     or "unknown",
+        decision             = _x("decision")              or "continue",
+        reasoning            = _x("reasoning")             or "",
+    )
+
+
+def _safe_float(s: str) -> float:
+    try:
+        return float(s)
+    except (ValueError, TypeError):
+        return 0.0
+
+
+def summarise_obs(obs: Dict[str, Any]) -> str:
+    parts = [
+        f"Incident : {obs.get('incident_summary', 'N/A')}",
+        f"Severity : {obs.get('severity', 'N/A')}",
+        f"Time     : {obs.get('time_elapsed_minutes', 0)}/{obs.get('time_budget_minutes', 30)} min",
+        f"Steps    : {obs.get('steps_taken', 0)}/{obs.get('max_steps', 20)}",
+        f"Reward   : {obs.get('current_reward', 0):.3f}  (Σ {obs.get('cumulative_reward', 0):.3f})",
+    ]
     statuses = obs.get("service_statuses", {})
     if statuses:
-        status_str = ", ".join(f"{k}: {v}" for k, v in statuses.items())
-        parts.append(f"Services: {status_str}")
+        parts.append("Services : " + "  ".join(f"{k}={v}" for k, v in statuses.items()))
+    parts.append(f"Alerts   : {obs.get('active_alerts_count', 0)} active")
+    parts.append(f"Result   : {obs.get('action_message', '')}")
+    data = obs.get("action_result", {})
+    if data:
+        blob = json.dumps(data, indent=2, default=str)
+        if len(blob) > 2000:
+            blob = blob[:2000] + "\n… (truncated)"
+        parts.append(f"Data:\n{blob}")
+    return "\n".join(parts)
 
-    parts.append(f"Alerts: {obs.get('active_alerts_count', 0)} active")
-    parts.append(f"Action result: {obs.get('action_message', 'N/A')}")
 
-    # Include action_result details (truncated)
-    action_result = obs.get("action_result", {})
-    if action_result:
-        result_str = json.dumps(action_result, indent=2, default=str)
-        if len(result_str) > 2000:
-            result_str = result_str[:2000] + "\n... (truncated)"
-        parts.append(f"Data:\n{result_str}")
+def obs_summary_dict(obs: Dict[str, Any]) -> Dict[str, Any]:
+    return {
+        "incident_summary":    obs.get("incident_summary", ""),
+        "severity":            obs.get("severity", ""),
+        "service_statuses":    obs.get("service_statuses", {}),
+        "active_alerts_count": obs.get("active_alerts_count", 0),
+        "action_message":      obs.get("action_message", ""),
+        "current_phase":       obs.get("current_phase", 1),
+    }
 
-    return "\n".join(parts)
 
+# ══════════════════════════════════════════════════════════════════
+# Phase 1 — Ops subagent
+# ══════════════════════════════════════════════════════════════════
+
+def run_phase1(
+    env:      EnvClient,
+    backend:  LLMBackend,
+    init_obs: Dict[str, Any],
+    episode:  EpisodeRecord,
+) -> tuple[BeliefState, Dict[str, Any]]:
+    """
+    Ops diagnostic loop.
+
+    The orchestrator is invoked after each ops action using a separate
+    system prompt and lower temperature.  It sees the full ops conversation
+    history so it can reason about cumulative evidence, not just the last step.
+
+    The ops and orchestrator calls are kept as separate generate() calls
+    rather than a single call with combined prompt — this lets them be
+    trained independently in Stage 2 and the orchestrator auxiliary loss
+    in Stage 4 without entangling their gradients.
+    """
+    ops_messages: List[Message] = [
+        {"role": "system", "content": OPS_SYSTEM_PROMPT},
+        {"role": "user",   "content": f"INCIDENT TRIGGERED:\n{summarise_obs(init_obs)}"},
+    ]
 
-def run_episode(
-    env: EnvClient,
-    llm: OpenAI,
-    task_name: str,
-    seed: int = 42,
-) -> Dict[str, Any]:
-    """Run a single episode and return results."""
+    belief   = BeliefState()
+    last_obs = init_obs
 
-    # --- [START] ---
-    print(f"[START] task={task_name}")
+    for p1_step in range(1, MAX_P1_STEPS + 1):
 
-    result = env.reset(task_name, seed)
-    obs = result["observation"]
+        # ── Ops agent selects next action ─────────────────────────
+        try:
+            ops_text = backend.generate(ops_messages, temperature=TEMPERATURE)
+            action   = parse_action(ops_text)
+        except Exception as e:
+            _warn(f"P1 ops error step {p1_step}: {e}")
+            action   = {"action_type": "view_alerts"}
+            ops_text = json.dumps(action)
+
+        # ── Orchestrator evaluates belief after seeing the action ─
+        # Orchestrator gets its own system prompt, then the full ops
+        # conversation up to and including the chosen action.
+        orch_messages: List[Message] = (
+            [{"role": "system", "content": ORCHESTRATOR_PROMPT}]
+            + ops_messages[1:]  # history without the ops system prompt
+            + [
+                {"role": "assistant", "content": ops_text},
+                {"role": "user",
+                 "content": "Based on all evidence gathered so far, output your belief state now."},
+            ]
+        )
+        try:
+            orch_text = backend.generate(
+                orch_messages,
+                temperature    = ORCH_TEMPERATURE,
+                max_new_tokens = 300,
+            )
+            belief = parse_belief(orch_text)
+        except Exception as e:
+            _warn(f"Orchestrator error step {p1_step}: {e}")
 
-    # Conversation history for belief tracking (Layer 3)
-    messages: List[Dict[str, str]] = [
-        {"role": "system", "content": SYSTEM_PROMPT},
-        {"role": "user", "content": f"INCIDENT TRIGGERED:\n{summarize_observation(obs)}"},
+        episode.belief_history.append(asdict(belief))
+
+        print(
+            f"[STEP] step={p1_step} phase=1 "
+            f"action={json.dumps(action)} "
+            f"svc={belief.suspected_service} "
+            f"svc_conf={belief.service_confidence:.2f} "
+            f"decision={belief.decision}"
+        )
+
+        # ── Execute in environment ────────────────────────────────
+        try:
+            step_result = env.step(action)
+        except Exception as e:
+            _warn(f"Env step error P1 step {p1_step}: {e}")
+            break
+
+        last_obs   = step_result.get("observation", {})
+        reward     = step_result.get("reward",  0.0)
+        done       = step_result.get("done",    False)
+        episode.cumulative_reward += reward
+        episode.total_steps        = p1_step
+
+        episode.p1_trajectory.append(StepRecord(
+            step_number = p1_step,
+            phase       = 1,
+            action      = action,
+            reward      = reward,
+            obs_summary = obs_summary_dict(last_obs),
+            belief      = asdict(belief),
+        ))
+
+        ops_messages.append({"role": "assistant", "content": ops_text})
+        ops_messages.append({
+            "role":    "user",
+            "content": f"Step {p1_step} result (reward={reward:.3f}):\n{summarise_obs(last_obs)}",
+        })
+
+        if done:
+            break
+
+        if belief.confident_enough():
+            episode.phase_transition_at = p1_step
+            break
+
+    return belief, last_obs
+
+
+# ══════════════════════════════════════════════════════════════════
+# Phase 2 — Code subagent
+# ══════════════════════════════════════════════════════════════════
+
+def run_phase2(
+    env:     EnvClient,
+    backend: LLMBackend,
+    belief:  BeliefState,
+    episode: EpisodeRecord,
+) -> None:
+    """
+    Triggers environment phase transition then runs code exploration.
+
+    The code agent gets a fresh context window — it does NOT receive
+    the P1 ops conversation history.  The handoff is only the structured
+    belief state fields (service, fault class, commit SHA).
+
+    This is deliberate: it forces the code agent to form its own
+    code-level evidence independently, and means belief state quality
+    directly gates Phase 2 search efficiency (the r_cross mechanism).
+    """
+    try:
+        p2_init = env.step({
+            "action_type": "transition_to_phase2",
+            "parameters":  {"belief": asdict(belief)},
+        })
+        p2_obs = p2_init.get("observation", {})
+    except Exception as e:
+        _warn(f"Phase transition failed: {e}")
+        return
+
+    commit_sha = p2_obs.get("bad_commit_sha", "unknown")
+
+    code_prompt = CODE_AGENT_PROMPT.format(
+        service            = belief.suspected_service,
+        fault_class        = belief.suspected_fault_class,
+        commit_sha         = commit_sha,
+        service_confidence = f"{belief.service_confidence:.2f}",
+        fault_confidence   = f"{belief.fault_confidence:.2f}",
+    )
+
+    messages: List[Message] = [
+        {"role": "system", "content": code_prompt},
+        {"role": "user",   "content": f"Codebase context:\n{summarise_obs(p2_obs)}"},
     ]
 
-    episode_reward = 0.0
-    final_info = {}
+    for p2_step in range(1, MAX_P2_STEPS + 1):
+        global_step = episode.total_steps + p2_step
 
-    for step_num in range(1, MAX_STEPS + 1):
         try:
-            # Get LLM action
-            completion = llm.chat.completions.create(
-                model=MODEL_NAME,
-                messages=messages,
-                temperature=TEMPERATURE,
-                max_tokens=256,
+            resp_text = backend.generate(
+                messages,
+                temperature    = TEMPERATURE,
+                max_new_tokens = MAX_NEW_TOKENS,
             )
-            llm_response = completion.choices[0].message.content or ""
+            action = parse_action(resp_text)
+        except Exception as e:
+            _warn(f"P2 action error step {p2_step}: {e}")
+            action    = {"action_type": "list_dir", "parameters": {"path": "."}}
+            resp_text = json.dumps(action)
 
-            # Parse action
-            action = parse_llm_action(llm_response)
+        a_type = action.get("action_type", "")
+        print(f"[STEP] step={global_step} phase=2 action={json.dumps(action)}")
 
-            # --- [STEP] ---
-            print(f"[STEP] step={step_num} action={json.dumps(action)}")
+        if a_type == "propose_patch":
+            episode.declared_patch = action.get("parameters", {}).get("diff", "")
+        elif a_type == "declare_no_change":
+            episode.declared_no_change = True
 
-            # Execute in environment
+        try:
             step_result = env.step(action)
-            obs = step_result["observation"]
-            reward = step_result.get("reward", 0.0)
-            done = step_result.get("done", False)
-            info = step_result.get("info", {})
-            episode_reward += reward
-
-            # Update conversation history (belief tracker)
-            messages.append({"role": "assistant", "content": llm_response})
-            messages.append({
-                "role": "user",
-                "content": f"Step {step_num} result (reward={reward}):\n{summarize_observation(obs)}"
-            })
-
-            if done:
-                final_info = info
-                break
+        except Exception as e:
+            _warn(f"Env step error P2 step {p2_step}: {e}")
+            break
+
+        step_obs   = step_result.get("observation", {})
+        reward     = step_result.get("reward",  0.0)
+        done       = step_result.get("done",    False)
+        episode.cumulative_reward += reward
+        episode.total_steps        = global_step
+
+        episode.p2_trajectory.append(StepRecord(
+            step_number = global_step,
+            phase       = 2,
+            action      = action,
+            reward      = reward,
+            obs_summary = obs_summary_dict(step_obs),
+        ))
+
+        messages.append({"role": "assistant", "content": resp_text})
+        messages.append({
+            "role":    "user",
+            "content": f"Step result:\n{summarise_obs(step_obs)}",
+        })
+
+        if done or a_type in {"propose_patch", "declare_no_change"}:
+            break
+
+
+# ══════════════════════════════════════════════════════════════════
+# Episode runners
+# ══════════════════════════════════════════════════════════════════
+
+def run_episode_baseline(
+    env:       EnvClient,
+    backend:   LLMBackend,
+    task_name: str,
+    seed:      int = 42,
+) -> EpisodeRecord:
+    """
+    Flat P1-only loop — no orchestrator, no Phase 2.
+    Ablation Claim 1: compare against run_episode_unified to prove
+    the orchestrator adds value beyond a fixed-strategy baseline.
+    """
+    print(f"[START] task={task_name}")
+    episode = EpisodeRecord(task_name=task_name, mode="baseline", seed=seed)
+    result  = env.reset(task_name, seed)
+    obs     = result["observation"]
+
+    messages: List[Message] = [
+        {"role": "system", "content": OPS_SYSTEM_PROMPT},
+        {"role": "user",   "content": f"INCIDENT TRIGGERED:\n{summarise_obs(obs)}"},
+    ]
 
+    final_info: Dict[str, Any] = {}
+
+    for step_num in range(1, MAX_P1_STEPS + 1):
+        try:
+            resp_text = backend.generate(messages, temperature=TEMPERATURE)
+            action    = parse_action(resp_text)
         except Exception as e:
-            print(f"[STEP] step={step_num} error={str(e)}", file=sys.stderr)
-            # Fallback action: view alerts
-            action = {"action_type": "view_alerts"}
-            step_result = env.step(action)
-            obs = step_result["observation"]
-            reward = step_result.get("reward", 0.0)
-            done = step_result.get("done", False)
-            episode_reward += reward
-            if done:
-                final_info = step_result.get("info", {})
-                break
-
-    # Get final state
-    final_state = env.state()
-    score = final_info.get("score", 0.01)
-
-    # --- [END] ---
-    print(f"[END] task={task_name} "
-          f"score={score:.3f} "
-          f"reward={episode_reward:.3f} "
-          f"steps={final_state.get('step_count', 0)}")
+            _warn(f"Baseline action error step {step_num}: {e}")
+            action    = {"action_type": "view_alerts"}
+            resp_text = json.dumps(action)
 
-    return {
-        "task_name": task_name,
-        "score": score,
-        "cumulative_reward": episode_reward,
-        "steps": final_state.get("step_count", 0),
-        "declared_root_cause": final_state.get("declared_root_cause"),
+        print(f"[STEP] step={step_num} phase=1 action={json.dumps(action)}")
+
+        try:
+            step_result = env.step(action)
+        except Exception as e:
+            _warn(f"Baseline env step error step {step_num}: {e}")
+            break
+
+        obs     = step_result.get("observation", {})
+        reward  = step_result.get("reward", 0.0)
+        done    = step_result.get("done",   False)
+        info    = step_result.get("info",   {})
+        episode.cumulative_reward += reward
+        episode.total_steps        = step_num
+
+        episode.p1_trajectory.append(StepRecord(
+            step_number = step_num,
+            phase       = 1,
+            action      = action,
+            reward      = reward,
+            obs_summary = obs_summary_dict(obs),
+        ))
+
+        messages.append({"role": "assistant", "content": resp_text})
+        messages.append({
+            "role":    "user",
+            "content": f"Step {step_num} result (reward={reward:.3f}):\n{summarise_obs(obs)}",
+        })
+
+        if done:
+            final_info = info
+            break
+
+    episode.score = final_info.get("score", 0.01)
+    print(
+        f"[END] task={task_name} score={episode.score:.3f} "
+        f"reward={episode.cumulative_reward:.3f} steps={episode.total_steps}"
+    )
+    return episode
+
+
+def run_episode_unified(
+    env:       EnvClient,
+    backend:   LLMBackend,
+    task_name: str,
+    seed:      int = 42,
+) -> EpisodeRecord:
+    """
+    Full two-phase episode:
+      Phase 1 — ops subagent diagnoses runtime incident
+      Orchestrator — belief state + stopping criterion after each P1 step
+      Phase 2 — code subagent explores codebase and proposes patch
+    """
+    print(f"[START] task={task_name}")
+    episode = EpisodeRecord(task_name=task_name, mode="unified", seed=seed)
+    result  = env.reset(task_name, seed)
+    obs     = result["observation"]
+
+    belief, _ = run_phase1(env, backend, obs, episode)
+
+    if episode.phase_transition_at is not None:
+        run_phase2(env, backend, belief, episode)
+
+    score_breakdown = env.unified_score(
+        declared_patch     = episode.declared_patch,
+        declared_no_change = episode.declared_no_change,
+        belief_history     = episode.belief_history,
+    )
+    episode.score = score_breakdown.get("final", 0.01)
+
+    print(
+        f"[END] task={task_name} score={episode.score:.3f} "
+        f"reward={episode.cumulative_reward:.3f} steps={episode.total_steps} "
+        f"transition_at={episode.phase_transition_at}"
+    )
+    return episode
+
+
+# ══════════════════════════════════════════════════════════════════
+# Trajectory persistence (SFT / GRPO data collection)
+# ══════════════════════════════════════════════════════════════════
+
+def save_trajectory(episode: EpisodeRecord) -> None:
+    """
+    Write episode to trajectories/<task>_<mode>_<n>.json.
+
+    Schema matches training/trajectory_collector.py:
+      p1_reward   — grader p1_rca + p1_efficiency  (filled post-hoc)
+      p2_reward   — grader patch + no_change scores (filled post-hoc)
+      r_cross     — counterfactual cross-phase reward (filled in Stage 4)
+      belief_history — per-step orchestrator beliefs, primary signal for r_cross
+    """
+    out_dir = Path("trajectories")
+    out_dir.mkdir(exist_ok=True)
+
+    idx  = len(list(out_dir.glob(f"{episode.task_name}_{episode.mode}_*.json")))
+    path = out_dir / f"{episode.task_name}_{episode.mode}_{idx:04d}.json"
+
+    record = {
+        "task_name":           episode.task_name,
+        "mode":                episode.mode,
+        "seed":                episode.seed,
+        "backend":             BACKEND,
+        "score":               episode.score,
+        "cumulative_reward":   episode.cumulative_reward,
+        "total_steps":         episode.total_steps,
+        "phase_transition_at": episode.phase_transition_at,
+        "declared_patch":      episode.declared_patch,
+        "declared_no_change":  episode.declared_no_change,
+        "belief_history":      episode.belief_history,
+        "p1_actions": [
+            {"step": r.step_number, "action": r.action,
+             "reward": r.reward, "belief": r.belief}
+            for r in episode.p1_trajectory
+        ],
+        "p2_actions": [
+            {"step": r.step_number, "action": r.action, "reward": r.reward}
+            for r in episode.p2_trajectory
+        ],
+        # Reward components filled post-hoc by grader / trajectory_collector
+        "p1_reward": 0.0,
+        "p2_reward": episode.score,
+        "r_cross":   0.0,
     }
 
+    path.write_text(json.dumps(record, indent=2, default=str))
+    _log(f"Trajectory saved → {path}")
+
+
+# ══════════════════════════════════════════════════════════════════
+# Utilities
+# ══════════════════════════════════════════════════════════════════
+
+def _log(msg: str) -> None:
+    print(f"[INFO] {msg}", flush=True)
+
+
+def _warn(msg: str) -> None:
+    print(f"[WARN] {msg}", file=sys.stderr, flush=True)
 
-# ------------------------------------------------------------------
+
+def _print_summary(results: List[EpisodeRecord]) -> None:
+    print(f"\n{'═' * 64}")
+    print(f"  RESULTS SUMMARY   mode={MODE}  backend={BACKEND}")
+    print(f"{'═' * 64}")
+    for r in results:
+        tr = f"→P2@step{r.phase_transition_at}" if r.phase_transition_at else "P1-only"
+        print(
+            f"  {r.task_name:30s}  score={r.score:.3f}  "
+            f"steps={r.total_steps:2d}  {tr}"
+        )
+    if results:
+        avg = sum(r.score for r in results) / len(results)
+        print(f"\n  {'AVERAGE':30s}  score={avg:.3f}")
+    print(f"{'═' * 64}")
+
+
+# ══════════════════════════════════════════════════════════════════
 # Main
-# ------------------------------------------------------------------
+# ══════════════════════════════════════════════════════════════════
 
-def main():
+def main() -> None:
     tasks = ["memory_leak", "cascading_failure", "distributed_deadlock"]
 
-    print("=" * 60)
-    print("SRE Incident Response — OpenEnv Inference")
-    print(f"Model: {MODEL_NAME}")
-    print(f"Environment: {ENV_BASE_URL}")
-    print("=" * 60)
+    print("═" * 64)
+    print("  SRE Incident Response — OpenEnv Inference")
+    print(f"  Backend : {BACKEND}")
+    print(f"  Model   : {LOCAL_MODEL_PATH if BACKEND == 'local' else MODEL_NAME}")
+    print(f"  Mode    : {MODE}")
+    print(f"  Env     : {ENV_BASE_URL}")
+    print(f"  Collect : {COLLECT}")
+    print("═" * 64)
+
+    env     = EnvClient(ENV_BASE_URL)
+    backend = build_backend()   # model loads once here
 
-    env = EnvClient(ENV_BASE_URL)
-    llm = create_openai_client()
+    run_fn = run_episode_unified if MODE == "unified" else run_episode_baseline
+    results: List[EpisodeRecord] = []
 
-    results = []
     for task in tasks:
         print(f"\n{'─' * 40}")
-        print(f"Task: {task}")
+        print(f"  Task: {task}")
         print(f"{'─' * 40}")
-
         try:
-            result = run_episode(env, llm, task)
-            results.append(result)
+            episode = run_fn(env, backend, task)
+            results.append(episode)
+            if COLLECT:
+                save_trajectory(episode)
         except Exception as e:
-            print(f"[ERROR] Task {task} failed: {e}", file=sys.stderr)
+            _warn(f"Task {task} failed: {e}")
             traceback.print_exc()
-            results.append({
-                "task_name": task,
-                "score": 0.01,
-                "cumulative_reward": 0.0,
-                "steps": 0,
-                "error": str(e),
-            })
-
-    # Summary
-    print(f"\n{'=' * 60}")
-    print("RESULTS SUMMARY")
-    print(f"{'=' * 60}")
-    for r in results:
-        score = r.get("score", 0.01)
-        print(f"  {r['task_name']:30s}  score={score:.3f}  "
-              f"steps={r.get('steps', 0):2d}  "
-              f"root_cause={r.get('declared_root_cause', 'N/A')}")
+            print(f"[END] task={task} score=0.010 reward=0.000 steps=0")
+            results.append(
+                EpisodeRecord(task_name=task, mode=MODE, seed=42, score=0.01)
+            )
 
-    avg_score = sum(r.get("score", 0) for r in results) / len(results)
-    print(f"\n  {'AVERAGE':30s}  score={avg_score:.3f}")
-    print(f"{'=' * 60}")
+    _print_summary(results)
 
 
 if __name__ == "__main__":
-    main()
+    main()

models.py

@@ -18,6 +18,7 @@ from typing import Any, Dict, List, Optional
 
 class ActionType(str, Enum):
     """Level-1 action categories — what kind of operation."""
+    # ---- Phase 1: ops / SRE ---------------------------------------
     VIEW_ALERTS = "view_alerts"
     QUERY_LOGS = "query_logs"
     CHECK_METRICS = "check_metrics"
@@ -29,6 +30,20 @@ class ActionType(str, Enum):
     SCALE_SERVICE = "scale_service"
     DECLARE_ROOT_CAUSE = "declare_root_cause"
 
+    # ---- Cross-phase control -------------------------------------
+    TRANSITION_TO_PHASE2 = "transition_to_phase2"
+
+    # ---- Phase 2: codebase exploration ---------------------------
+    READ_FILE = "read_file"
+    SEARCH_CODE = "search_code"
+    LIST_DIR = "list_dir"
+    GET_GIT_LOG = "get_git_log"
+    GET_FILE_DIFF = "get_file_diff"
+
+    # ---- Phase 2: terminal --------------------------------------
+    PROPOSE_PATCH = "propose_patch"
+    DECLARE_NO_CHANGE = "declare_no_change"
+
 
 # Actions that require a target_service (Level 2 — where to apply)
 TARGETED_ACTIONS = {
@@ -59,6 +74,62 @@ REMEDIATION_ACTIONS = {
     ActionType.SCALE_SERVICE,
 }
 
+# ---- Phase classification ------------------------------------------
+PHASE1_ACTIONS = (
+    DIAGNOSTIC_ACTIONS
+    | REMEDIATION_ACTIONS
+    | {ActionType.DECLARE_ROOT_CAUSE}
+)
+
+PHASE2_DIAGNOSTIC_ACTIONS = {
+    ActionType.LIST_DIR,
+    ActionType.READ_FILE,
+    ActionType.SEARCH_CODE,
+    ActionType.GET_GIT_LOG,
+    ActionType.GET_FILE_DIFF,
+}
+
+PHASE2_TERMINAL_ACTIONS = {
+    ActionType.PROPOSE_PATCH,
+    ActionType.DECLARE_NO_CHANGE,
+}
+
+PHASE2_ACTIONS = PHASE2_DIAGNOSTIC_ACTIONS | PHASE2_TERMINAL_ACTIONS
+
+# Cross-phase control action — only legal when transitioning P1 → P2
+CONTROL_ACTIONS = {ActionType.TRANSITION_TO_PHASE2}
+
+@dataclass
+class BeliefState:
+    """
+    Structured scratchpad the orchestrator emits before each transition decision.
+    Making this explicit (not implicit in hidden states) lets us:
+      - supervise confidence calibration with auxiliary loss
+      - audit stopping criterion decisions
+      - compute consistency losses (e.g. empty gaps + low confidence = incoherent)
+    """
+    suspected_service: Optional[str] = None
+    suspected_fault_class: Optional[str] = None   # "memory_leak" | "config_change" | "deadlock" | "none"
+    service_confidence: float = 0.0               # [0, 1] — calibrated against ground truth in Stage 2
+    fault_confidence: float = 0.0                 # [0, 1]
+    evidence_gaps: List[str] = field(default_factory=list)  # e.g. ["deploy_history_unchecked"]
+    estimated_p2_cost: str = "unknown"            # "low" | "medium" | "high"
+    decision: str = "continue"                    # "continue" | "transition" | "abort"
+    reasoning: str = ""                           # free-text, used for consistency loss
+
+@dataclass
+class CodeContext:
+    """
+    Hidden code-layer state — injected into Phase 2 at transition.
+    Agent never sees this directly; it must be inferred by exploration.
+    """
+    repo_snapshot_path: str                       # path to bundled mini-repo
+    bad_commit_sha: str
+    ground_truth_files: List[str]                 # files touched by real PR
+    ground_truth_diff: str                        # unified diff string
+    is_valid_issue: bool = True                   # False = user confusion, no-change correct
+    expected_p2_steps: int = 8                    # baseline for efficiency normalization
+    null_context_p2_score: float = 0.0            # filled in during Stage 3 (Pool B baseline)
 
 @dataclass
 class IncidentAction:
@@ -71,6 +142,9 @@ class IncidentAction:
     action_type: str                            # ActionType value
     target_service: Optional[str] = None        # Required for TARGETED_ACTIONS
     parameters: Dict[str, Any] = field(default_factory=dict)
+    # Phase tracking
+    current_phase: int = 1                            # 1 = ops, 2 = code
+    belief_state: Optional[BeliefState] = None        # orchestrator scratchpad output
 
     def parsed_type(self) -> ActionType:
         return ActionType(self.action_type)
@@ -197,3 +271,5 @@ class StepRecord:
     observation_summary: Dict[str, Any]     # key fields from observation
     service_statuses_after: Dict[str, str]  # service health after this step
     timestamp_minutes: int                  # simulation time
+    phase: int = 1
+    belief_state_snapshot: Optional[dict] = None     # serialized BeliefState for behavioral analysis

pools.py

@@ -0,0 +1,141 @@
+"""
+Pool registry for the four-stage curriculum.
+
+A Pool is a named subset of TASK_REGISTRY plus an *episode mode* that controls
+how the environment behaves for that pool.  The pool name is passed to
+`/reset` via the `pool` field; the server then samples a task from that pool
+and switches the environment into the matching mode.
+
+Stages and pools (per the brief):
+
+    Stage 2  bootstrap ops agent           → Pool A   mode = "p1_only"
+    Stage 3  bootstrap code agent          → Pool B   mode = "p2_only"
+                                                       (P1 context = ground truth)
+    Stage 4  joint training with r_cross   → Pool C   mode = "joint"
+    Final    held-out generalization eval  → Pool D   mode = "joint"
+
+Pool A and Pool C reuse the same scenarios — only the mode differs.
+Pool B is a *bootstrapping* mode where the orchestrator's belief is *synthesized*
+from the scenario's ground truth, so the code agent never trains on garbage
+Phase-1 context.  This implements the brief's "P2-only with ground-truth
+P1 context injected" semantics exactly.
+
+Pool D consists of *held-out* scenarios that never appear during training,
+used to measure whether the learned stopping criterion generalizes.
+"""
+
+from __future__ import annotations
+
+import random
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Set
+
+from .models import BeliefState
+from .scenarios.base import BaseScenario
+
+
+# ──────────────────────────────────────────────────────────────────────
+# Pool definitions
+# ──────────────────────────────────────────────────────────────────────
+
+
+@dataclass(frozen=True)
+class Pool:
+    """A named pool: training scenarios + episode mode."""
+    name:          str                # "A" | "B" | "C" | "D"
+    description:   str
+    task_names:    List[str]
+    mode:          str                # "p1_only" | "p2_only" | "joint"
+    # Stage-3 hints used when mode == "p2_only" (Pool B): if True, the env
+    # auto-injects ground-truth context at handoff so the code agent never
+    # sees a noisy P1 trajectory.
+    inject_oracle_belief: bool = False
+
+
+# Training scenarios (seen during all four training stages).  Phase-A
+# scenarios + the brief's four research scenarios.
+_TRAIN_TASKS = [
+    "memory_leak",
+    "cascading_failure",
+    "distributed_deadlock",
+    "circuit_breaker_noop",
+    "aliased_fault",
+    "severity_inversion",
+    "confidence_inversion",
+    "info_ordering",
+]
+
+
+# Held-out scenarios — same fault families, but combined in ways the agent
+# never trained on.  Defined in scenarios/heldout.py and registered lazily.
+_HELDOUT_TASKS = [
+    "heldout_aliased_severity",      # aliased + severity-inversion combo
+    "heldout_confidence_ordering",   # confidence-inversion + info-ordering combo
+]
+
+
+POOLS: Dict[str, Pool] = {
+    "A": Pool(
+        name        = "A",
+        description = "Stage-2 ops bootstrap — P1 only, declare_root_cause terminates.",
+        task_names  = _TRAIN_TASKS,
+        mode        = "p1_only",
+    ),
+    "B": Pool(
+        name        = "B",
+        description = "Stage-3 code bootstrap — P2 only with oracle P1 context injected.",
+        task_names  = _TRAIN_TASKS,
+        mode        = "p2_only",
+        inject_oracle_belief = True,
+    ),
+    "C": Pool(
+        name        = "C",
+        description = "Stage-4 joint training — full P1 → P2 with r_cross.",
+        task_names  = _TRAIN_TASKS,
+        mode        = "joint",
+    ),
+    "D": Pool(
+        name        = "D",
+        description = "Held-out generalization — never seen during training.",
+        task_names  = _HELDOUT_TASKS,
+        mode        = "joint",
+    ),
+}
+
+
+def get_pool(name: str) -> Pool:
+    name = (name or "").upper()
+    if name not in POOLS:
+        raise ValueError(f"Unknown pool {name!r}. Available: {list(POOLS)}")
+    return POOLS[name]
+
+
+def sample_task(pool_name: str, rng: Optional[random.Random] = None) -> str:
+    """Sample one task from a pool."""
+    rng = rng or random
+    return rng.choice(get_pool(pool_name).task_names)
+
+
+# ──────────────────────────────────────────────────────────────────────
+# Oracle belief synthesis (for Pool B)
+# ──────────────────────────────────────────────────────────────────────
+
+
+def oracle_belief(scenario: BaseScenario) -> BeliefState:
+    """
+    Synthesize a *ground-truth* belief from the scenario's static config.
+
+    Used by Pool B (Stage 3) so the code agent sees a perfect Phase-1
+    handoff — its training signal is then purely Phase-2 quality, not
+    Phase-1 errors.  This is the cleanest possible code-agent bootstrap.
+    """
+    return BeliefState(
+        suspected_service     = scenario.root_cause_service,
+        suspected_fault_class = scenario.fault_class,
+        service_confidence    = 1.0,
+        fault_confidence      = 1.0,
+        evidence_gaps         = [],
+        estimated_p2_cost     = "low",
+        decision              = "transition",
+        reasoning             = "[oracle] ground-truth belief synthesized for Pool B bootstrap",
+    )

scenarios/aliased_fault.py

@@ -0,0 +1,153 @@
+"""
+Phase-B scenario 1 — Aliased fault patterns.
+
+Two distinct faults produce *identical initial observations*:
+  - Memory leak in `orders`             (real cause)
+  - Cache thrashing in `queue` (saturated upstream that orders depends on)
+
+Both surface as "high memory + degraded orders" on the dashboard.
+
+The agent's default prior is "investigate the loudest service" — it heads
+straight to `check_metrics(orders)` and `check_deploy_history(orders)`,
+finds the recent batch-processing deploy, and rolls it back.  In the
+aliased version, the recent deploy on `orders` is innocuous; the *real*
+cause is in `queue`'s flush worker.
+
+The diagnostic that disambiguates is `check_dependencies(orders)` followed
+by `check_metrics(queue)` — only that ordering reveals queue is saturated.
+
+This scenario is designed to *break* the base model's prior.  An RL agent
+with the right credit assignment learns to interleave a dependency check
+*before* committing to the loudest service.
+"""
+
+from __future__ import annotations
+
+import random
+from typing import Dict, List, Set
+
+from .base import BaseScenario
+from .code_context_builder import ALIASED_FAULT_CODE_CONTEXT
+from ..models import CodeContext
+from ..simulation.infrastructure import Infrastructure
+from ..simulation.service import Deploy
+from ..simulation.metrics import (
+    generate_memory_leak_history,
+    generate_high_latency_history,
+)
+
+
+class AliasedFaultScenario(BaseScenario):
+
+    @property
+    def task_name(self) -> str:
+        return "aliased_fault"
+
+    @property
+    def code_context(self) -> CodeContext:
+        return ALIASED_FAULT_CODE_CONTEXT
+
+    @property
+    def fault_class(self) -> str:
+        return "cache_thrash"
+
+    @property
+    def display_name(self) -> str:
+        return "Aliased Fault — Cache Thrash Disguised as Memory Leak"
+
+    @property
+    def incident_summary(self) -> str:
+        return (
+            "INCIDENT: Orders service showing high memory usage (84%), elevated latency, "
+            "and intermittent OOM-like errors. Recent orders deploy v2.4.0 was just "
+            "rolled out 18 minutes ago. On-call SRE paged."
+        )
+
+    @property
+    def severity(self) -> str:
+        return "SEV2"
+
+    @property
+    def correct_root_cause(self) -> str:
+        return "queue cache thrashing — saturated worker overflows shared cache used by orders"
+
+    @property
+    def root_cause_keywords(self) -> List[str]:
+        return ["queue", "cache", "thrash", "worker"]
+
+    @property
+    def involved_services(self) -> Set[str]:
+        return {"queue", "orders"}
+
+    @property
+    def root_cause_service(self) -> str:
+        return "queue"
+
+    @property
+    def correct_remediation_actions(self) -> List[Dict[str, str]]:
+        return [
+            {"action_type": "rollback_deploy", "target_service": "queue"},
+            {"action_type": "restart_service", "target_service": "orders"},
+        ]
+
+    def inject(self, infra: Infrastructure) -> None:
+        orders = infra.get_service("orders")
+        queue  = infra.get_service("queue")
+        if orders is None or queue is None:
+            return
+
+        # --- Innocuous orders deploy (the red herring) ---
+        orders.deploy_history.append(Deploy(
+            version="v2.4.0", timestamp_minutes=-18, author="alice",
+            commit_hash="b7d291", description="Refactor: extract pricing helper",
+            is_bad=False,
+        ))
+
+        # --- Real bad deploy: queue worker (the hidden root cause) ---
+        queue.deploy_history.append(Deploy(
+            version="v2.1.0", timestamp_minutes=-22, author="dan",
+            commit_hash="e1f4a02", description="Optimization: bulk flush via shared cache",
+            is_bad=True,
+        ))
+
+        # --- Symptoms on orders look IDENTICAL to a memory leak ---
+        orders.memory_percent = 84.0 + random.gauss(0, 2)
+        orders.cpu_percent    = 41.0 + random.gauss(0, 3)
+        orders.error_rate_percent = 14.0 + random.gauss(0, 2)
+        orders.latency_p95_ms = 380.0 + random.gauss(0, 30)
+        orders.latency_p99_ms = 920.0 + random.gauss(0, 50)
+        orders.status = "degraded"
+        orders.metric_history = generate_memory_leak_history(
+            minutes=30, start_minute=0, leak_start_offset=12, rate=1.4)
+
+        orders.logs = [
+            {"timestamp": "2025-01-15T14:18:00Z", "level": "WARN", "service": "orders",
+             "message": "GC pressure: heap usage at 78%, GC pause 410ms", "trace_id": None},
+            {"timestamp": "2025-01-15T14:22:00Z", "level": "ERROR", "service": "orders",
+             "message": "Allocation failure: cache eviction backlog growing",
+             "trace_id": "trace-554301"},
+            {"timestamp": "2025-01-15T14:25:00Z", "level": "ERROR", "service": "orders",
+             "message": "OutOfMemoryError-like behaviour: cache write blocked >2s",
+             "trace_id": "trace-554404"},
+        ]
+
+        # --- Real fault on queue (subtle: only visible if you check it) ---
+        queue.inject_fault("high_latency", p99=4500)
+        queue.metric_history = generate_high_latency_history(
+            minutes=30, start_minute=0, latency_start_offset=10, target_p99=4500)
+        queue.cpu_percent          = 88.0 + random.gauss(0, 3)
+        queue.memory_percent       = 71.0 + random.gauss(0, 2)
+        queue.error_rate_percent   = 6.0  + random.gauss(0, 1)
+        queue.latency_p99_ms       = 4500.0 + random.gauss(0, 100)
+        queue.requests_per_sec     = 90.0  + random.gauss(0, 5)
+        queue.status               = "degraded"
+        queue.logs = [
+            {"timestamp": "2025-01-15T14:14:00Z", "level": "INFO", "service": "queue",
+             "message": "Deploy v2.1.0 complete — bulk-flush worker active", "trace_id": None},
+            {"timestamp": "2025-01-15T14:19:00Z", "level": "WARN", "service": "queue",
+             "message": "Worker backlog: 8400 messages awaiting flush, eviction rate 320/s",
+             "trace_id": None},
+            {"timestamp": "2025-01-15T14:24:00Z", "level": "ERROR", "service": "queue",
+             "message": "Cache eviction storm: orders-cache shard 3 evicted 84% of keys",
+             "trace_id": "trace-771204"},
+        ]

scenarios/base.py

@@ -1,30 +1,42 @@
 """
 Base scenario class.
 
-Each scenario defines:
-  - How to inject faults into the infrastructure
-  - The correct root cause string
+A scenario is *static config*:
+  - How to inject faults into the infrastructure (`inject()`)
+  - The correct root cause string + keywords (for grading)
   - Which services are involved (for reward shaping)
-  - The oracle grader (trajectory-only, no hidden state)
+  - The optional code-attribution context (`code_context` property)
+  - The oracle-independent grader (`grade()`)
+
+Per-episode mutable state (phase, p1/p2 trajectories, declared patch,
+code workspace) lives on `IncidentEnvironment`, NOT here. A scenario
+instance is therefore safe to share across episodes.
 """
 
 from __future__ import annotations
 
+import math
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Set
+from typing import Any, Dict, List, Optional, Set
 
+from ..models import StepRecord, CodeContext, BeliefState
 from ..simulation.infrastructure import Infrastructure
-from ..models import StepRecord
 
 
 class BaseScenario(ABC):
     """
-    Abstract scenario.  Subclasses implement inject() and grade().
+    Abstract scenario.  Subclasses implement `inject()` and the static
+    config properties below.
 
-    inject() mutates the infrastructure to set up the incident.
-    grade() evaluates a complete trajectory WITHOUT access to hidden state.
+    Optional Phase-2 support: subclasses override `code_context` to point
+    at a bundled mini-repo + ground-truth diff. If `code_context` returns
+    `None`, the scenario is Phase-1 only (legacy).
     """
 
+    # ------------------------------------------------------------------
+    # Static config (must be overridden by subclasses)
+    # ------------------------------------------------------------------
+
     @property
     @abstractmethod
     def task_name(self) -> str:
@@ -46,7 +58,7 @@ class BaseScenario(ABC):
     @property
     @abstractmethod
     def severity(self) -> str:
-        """SEV1/SEV2/SEV3."""
+        """SEV1 / SEV2 / SEV3."""
         ...
 
     @property
@@ -86,6 +98,57 @@ class BaseScenario(ABC):
     def max_steps(self) -> int:
         return 20
 
+    # ---- Phase-2 hook (optional) ------------------------------------
+
+    @property
+    def code_context(self) -> Optional[CodeContext]:
+        """
+        Override to enable Phase 2 (code attribution).
+        Default: scenario is P1-only.
+        """
+        return None
+
+    @property
+    def fault_class(self) -> str:
+        """
+        Ground-truth fault class for belief-state aux loss in Stage 2.
+        One of: memory_leak | config_change | deadlock | resource_exhaustion |
+                cascading | none
+        """
+        return "none"
+
+    # ---- P2 handoff: synthetic issue text -----------------------------
+
+    def build_p2_issue(self, belief: Optional[BeliefState] = None) -> str:
+        """
+        Build the synthetic GitHub-issue-style text the code agent reads at
+        handoff. Combines the incident summary with whatever runtime evidence
+        Phase 1 surfaced. The agent uses this to seed its codebase search.
+        """
+        lines = [
+            f"## Incident: {self.display_name}",
+            "",
+            self.incident_summary,
+            "",
+        ]
+        if belief is not None:
+            lines.append("## Phase-1 diagnosis (handed off)")
+            lines.append(f"- Suspected service: **{belief.suspected_service or 'unknown'}**")
+            lines.append(f"- Suspected fault class: **{belief.suspected_fault_class or 'unknown'}**")
+            lines.append(f"- Service confidence: {belief.service_confidence:.2f}")
+            lines.append(f"- Fault confidence:   {belief.fault_confidence:.2f}")
+            if belief.evidence_gaps:
+                gaps = belief.evidence_gaps if isinstance(belief.evidence_gaps, list) \
+                    else [belief.evidence_gaps]
+                lines.append(f"- Outstanding evidence gaps: {', '.join(map(str, gaps))}")
+            if belief.reasoning:
+                lines.append(f"- Reasoning: {belief.reasoning}")
+        return "\n".join(lines)
+
+    # ------------------------------------------------------------------
+    # Fault injection (must be implemented)
+    # ------------------------------------------------------------------
+
     @abstractmethod
     def inject(self, infra: Infrastructure) -> None:
         """
@@ -94,130 +157,111 @@ class BaseScenario(ABC):
         """
         ...
 
-    # ---------------------------------------------------------------
-    # Grading — oracle-independent, trajectory-only (Layer 6)
-    # ---------------------------------------------------------------
+    # ==================================================================
+    # Grading — oracle-independent, trajectory-only
+    # ==================================================================
 
     def grade(self, trajectory: List[StepRecord]) -> float:
         """
-        Grade the complete trajectory.
-        Returns float in [0.01, 0.99].
-
-        This function receives ONLY the step records — no hidden state,
-        no infrastructure reference. This is critical: the evaluation
-        harness must be able to call this on a saved trajectory.
+        P1-only grader (legacy). Returns float in [0.01, 0.99].
+        Component breakdown: 40% RCA + 30% remediation + 20% efficiency + 10% restoration.
         """
-        import math
-        
         score = 0.0
         score += self._grade_root_cause(trajectory)        # 0.00 – 0.40
         score += self._grade_remediation(trajectory)        # 0.00 – 0.30
         score += self._grade_efficiency(trajectory)         # 0.00 – 0.20
         score += self._grade_restoration(trajectory)        # 0.00 – 0.10
-        
-        # Final safety check for NaN/Inf
+
         if not math.isfinite(score):
             score = 0.0
-            
-        # Ensure score is strictly open interval (0, 1) to pass OpenEnv validation via affine transform
-        # We target [0.01, 0.99] to stay safely away from boundaries
-        adjusted_score = 0.01 + (min(max(float(score), 0.0), 1.0) * 0.98)
-        
-        # Explicit boundary enforcement for absolute certainty
-        if adjusted_score <= 0.001:
+
+        # OpenEnv validator requires strict (0, 1)
+        adjusted = 0.01 + (min(max(float(score), 0.0), 1.0) * 0.98)
+        if adjusted <= 0.001:
             return 0.01
-        if adjusted_score >= 0.999:
+        if adjusted >= 0.999:
             return 0.99
-            
-        return float(round(adjusted_score, 4))
+        return float(round(adjusted, 4))
 
+    # ---- Component graders (used directly by unified grader) ---------
+
+    def grade_p1_rca(self, p1_trajectory: List[StepRecord]) -> float:
+        """RCA component in [0, 1] (independent of weight)."""
+        return self._grade_root_cause(p1_trajectory) / 0.40
+
+    def grade_p1_efficiency(self, p1_trajectory: List[StepRecord]) -> float:
+        """Efficiency component in [0, 1]: 1.0 at 0 steps to declare, 0 at max_steps."""
+        declare_step = next(
+            (r.step_number for r in p1_trajectory
+             if r.action.action_type == "declare_root_cause"),
+            self.max_steps,
+        )
+        return max(0.0, 1.0 - (declare_step / max(self.max_steps, 1)))
+
+    # ---- Internal raw component graders -------------------------------
 
     def _grade_root_cause(self, trajectory: List[StepRecord]) -> float:
-        """
-        Did the agent correctly declare the root cause?
-        Full credit (0.40) for correct, partial credit for close.
-        """
+        """0.40 if last declaration matches keywords ≥60%, 0.20 ≥30%, else 0."""
         declarations = [
             s for s in trajectory
             if s.action.action_type == "declare_root_cause"
         ]
         if not declarations:
-            return 0.0  # Never declared — 0 points
+            return 0.0
 
-        # Use the LAST declaration
         declared = declarations[-1].action.parameters.get("root_cause", "").lower()
-
-        # Check keyword match
-        keywords = self.root_cause_keywords
-        if not keywords:
-            keywords = self.correct_root_cause.lower().split()
+        keywords = self.root_cause_keywords or self.correct_root_cause.lower().split()
 
         matched = sum(1 for kw in keywords if kw in declared)
-        match_ratio = matched / len(keywords) if keywords else 0
+        match_ratio = matched / len(keywords) if keywords else 0.0
 
         if match_ratio >= 0.6:
-            return 0.40  # Close enough — full credit
-        elif match_ratio >= 0.3:
-            return 0.20  # Partial credit
-        else:
-            return 0.0
+            return 0.40
+        if match_ratio >= 0.3:
+            return 0.20
+        return 0.0
 
     def _grade_remediation(self, trajectory: List[StepRecord]) -> float:
-        """
-        Did the agent take the correct fix actions?
-        """
-        correct_actions = self.correct_remediation_actions
-        if not correct_actions:
+        """Fraction of correct (action, target) pairs taken, scaled to 0.30."""
+        correct = self.correct_remediation_actions
+        if not correct:
             return 0.0
 
-        taken_remediations = [
+        taken = [
             (s.action.action_type, s.action.target_service)
             for s in trajectory
             if s.action.action_type in ("restart_service", "rollback_deploy", "scale_service")
         ]
 
-        matched = 0
-        for ca in correct_actions:
-            needed = (ca["action_type"], ca["target_service"])
-            if needed in taken_remediations:
-                matched += 1
-
-        ratio = matched / len(correct_actions)
-        return round(ratio * 0.30, 3)
+        matched = sum(
+            1 for ca in correct
+            if (ca["action_type"], ca["target_service"]) in taken
+        )
+        return round((matched / len(correct)) * 0.30, 3)
 
     def _grade_efficiency(self, trajectory: List[StepRecord]) -> float:
-        """
-        Fewer steps to reach correct diagnosis = more points.
-        Optimal path (for the scenario) gets full credit.
-        """
-        total_steps = len(trajectory)
-        if total_steps == 0:
+        """Step-count tier credit, max 0.20."""
+        n = len(trajectory)
+        if n == 0:
             return 0.0
-
-        # Generous: < 8 steps is excellent, 8-12 is good, 13-16 is okay, 17+ is bad
-        if total_steps <= 6:
+        if n <= 6:
             return 0.20
-        elif total_steps <= 10:
+        if n <= 10:
             return 0.15
-        elif total_steps <= 14:
+        if n <= 14:
             return 0.10
-        elif total_steps <= 17:
+        if n <= 17:
             return 0.05
-        else:
-            return 0.02
+        return 0.02
 
     def _grade_restoration(self, trajectory: List[StepRecord]) -> float:
-        """
-        Are all services healthy at the end of the episode?
-        Check the LAST step's service_statuses_after.
-        """
+        """Final-step service health, max 0.10."""
         if not trajectory:
             return 0.0
-
-        final_statuses = trajectory[-1].service_statuses_after
-        if all(s == "healthy" for s in final_statuses.values()):
+        final = trajectory[-1].service_statuses_after
+        if not final:
+            return 0.0
+        if all(s == "healthy" for s in final.values()):
             return 0.10
-        # Partial credit: how many are healthy
-        healthy_count = sum(1 for s in final_statuses.values() if s == "healthy")
-        total = len(final_statuses) if final_statuses else 1
-        return round(0.10 * (healthy_count / total), 3)
+        healthy = sum(1 for s in final.values() if s == "healthy")
+        return round(0.10 * (healthy / len(final)), 3)

scenarios/circuit_breaker_noop.py

@@ -0,0 +1,108 @@
+"""
+Phase-A "no-change" scenario.
+
+Symptoms look real (orders is intermittently slow; user filed an issue
+claiming "the new release broke checkout"), but on inspection the only
+recent deploy is a documentation comment update — the slowness is the
+service's normal weekly batch backup window kicking off.
+
+The correct behaviour:
+  - Phase 1: investigate, find that no fault is attributable
+  - Phase 2: read the snapshot and emit `declare_no_change`
+
+The scenario's `code_context.is_valid_issue == False`, so any proposed
+diff scores 0 and `grade_no_change(True)` scores 1.0.
+"""
+
+from __future__ import annotations
+
+import random
+from typing import Dict, List, Set
+
+from .base import BaseScenario
+from .code_context_builder import CIRCUIT_BREAKER_CODE_CONTEXT
+from ..models import CodeContext
+from ..simulation.infrastructure import Infrastructure
+from ..simulation.service import Deploy
+from ..simulation.metrics import generate_healthy_history
+
+
+class CircuitBreakerNoopScenario(BaseScenario):
+
+    @property
+    def task_name(self) -> str:
+        return "circuit_breaker_noop"
+
+    @property
+    def code_context(self) -> CodeContext:
+        return CIRCUIT_BREAKER_CODE_CONTEXT
+
+    @property
+    def fault_class(self) -> str:
+        return "none"
+
+    @property
+    def display_name(self) -> str:
+        return "Spurious Issue — No Code Change Required"
+
+    @property
+    def incident_summary(self) -> str:
+        return (
+            "INCIDENT: User report — 'orders deploy v3.0.0 broke checkout, latency spiked'. "
+            "On-call paged. The orders service is slightly elevated on latency but no "
+            "alert fired."
+        )
+
+    @property
+    def severity(self) -> str:
+        return "SEV3"
+
+    @property
+    def correct_root_cause(self) -> str:
+        return "no fault — orders is in its normal weekly backup window"
+
+    @property
+    def root_cause_keywords(self) -> List[str]:
+        return ["no", "fault", "backup", "normal"]
+
+    @property
+    def involved_services(self) -> Set[str]:
+        return {"orders"}
+
+    @property
+    def root_cause_service(self) -> str:
+        return "orders"
+
+    @property
+    def correct_remediation_actions(self) -> List[Dict[str, str]]:
+        # No remediation is "correct" — declaring no-change in P2 is the goal.
+        return []
+
+    def inject(self, infra: Infrastructure) -> None:
+        orders = infra.get_service("orders")
+        if orders is None:
+            return
+
+        orders.deploy_history.append(Deploy(
+            version="v3.0.0", timestamp_minutes=-30, author="ed",
+            commit_hash="d2b9c11",
+            description="Docs: clarify checkout API contract in comments",
+            is_bad=False,
+        ))
+
+        orders.metric_history = generate_healthy_history(minutes=30, start_minute=0)
+        orders.cpu_percent          = 18.0 + random.gauss(0, 2)
+        orders.memory_percent       = 41.0 + random.gauss(0, 2)
+        orders.error_rate_percent   = 0.6  + random.gauss(0, 0.2)
+        orders.latency_p95_ms       = 220.0 + random.gauss(0, 20)   # slightly elevated
+        orders.latency_p99_ms       = 380.0 + random.gauss(0, 25)
+        orders.status               = "healthy"
+
+        orders.logs = [
+            {"timestamp": "2025-01-15T14:00:00Z", "level": "INFO", "service": "orders",
+             "message": "Weekly backup window started — minor latency overhead expected",
+             "trace_id": None},
+            {"timestamp": "2025-01-15T14:12:00Z", "level": "INFO", "service": "orders",
+             "message": "Backup snapshot 47% complete (12.4 GB / 26.5 GB)",
+             "trace_id": None},
+        ]

scenarios/code_context_builder.py

@@ -0,0 +1,246 @@
+"""
+CodeContext registry for each Phase-2-enabled scenario.
+
+Snapshots live under <repo_root>/snapshots/<name>/, bundled in the repo —
+no live GitHub API calls.  Each context provides the snapshot path, the bad
+commit SHA, the ground-truth files/diff (used by `grader_p2.grade_patch_quality`)
+and a slot for the Pool-B null-context baseline (filled in by the
+`training/run_pool_b_baseline.py` runner — re-imported on demand).
+
+The `null_context_p2_score` field starts at a hand-tuned prior; it is
+overwritten in-place by the baseline runner once we have measurements.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from ..models import CodeContext
+
+
+# ──────────────────────────────────────────────────────────────────────
+# Snapshot-root resolution
+# ──────────────────────────────────────────────────────────────────────
+
+_PKG_ROOT       = Path(__file__).resolve().parent.parent          # incident_env/
+_REPO_ROOT      = _PKG_ROOT.parent                                # project/
+SNAPSHOTS_ROOT  = _REPO_ROOT / "snapshots"
+
+
+def _snap(name: str) -> str:
+    """Absolute path to a snapshot directory under <repo_root>/snapshots/."""
+    return str(SNAPSHOTS_ROOT / name)
+
+
+# ──────────────────────────────────────────────────────────────────────
+# Memory leak (easy)
+# ──────────────────────────────────────────────────────────────────────
+
+MEMORY_LEAK_CODE_CONTEXT = CodeContext(
+    repo_snapshot_path  = _snap("orders_v231"),
+    bad_commit_sha      = "a3f7c91",
+    ground_truth_files  = ["orders/handlers/batch.py"],
+    ground_truth_diff   = """--- a/orders/handlers/batch.py
++++ b/orders/handlers/batch.py
+@@ -41,6 +41,7 @@ class BatchProcessor:
+         for order in orders:
+             self._cache[order.id] = order
++        self._cache.clear()
+         self._notify(orders)
+""",
+    is_valid_issue          = True,
+    expected_p2_steps       = 5,
+    null_context_p2_score   = 0.21,
+)
+
+# ──────────────────────────────────────────────────────────────────────
+# Cascading failure (medium)
+# ──────────────────────────────────────────────────────────────────────
+
+CASCADING_FAILURE_CODE_CONTEXT = CodeContext(
+    repo_snapshot_path  = _snap("auth_v180"),
+    bad_commit_sha      = "b8e2d44",
+    ground_truth_files  = ["auth/config.py"],
+    ground_truth_diff   = """--- a/auth/config.py
++++ b/auth/config.py
+@@ -10,3 +10,3 @@
+-JWT_SECRET = os.environ.get("JWT_SECRET")
++JWT_SECRET = os.environ.get("JWT_SECRET") or _DEFAULT_DEV_SECRET
+""",
+    is_valid_issue          = True,
+    expected_p2_steps       = 4,
+    null_context_p2_score   = 0.18,
+)
+
+# ──────────────────────────────────────────────────────────────────────
+# Distributed deadlock (hard)
+# ──────────────────────────────────────────────────────────────────────
+
+DISTRIBUTED_DEADLOCK_CODE_CONTEXT = CodeContext(
+    repo_snapshot_path  = _snap("payment_v310"),
+    bad_commit_sha      = "c5a1f77",
+    ground_truth_files  = ["payment/processor.py"],
+    ground_truth_diff   = """--- a/payment/processor.py
++++ b/payment/processor.py
+@@ -85,5 +85,7 @@ class PaymentProcessor:
+     def retry(self, txn):
++        delay = min(2 ** self.retry_count, 30)
++        time.sleep(delay)
+         self._queue.enqueue(txn)
+""",
+    is_valid_issue          = True,
+    expected_p2_steps       = 10,
+    null_context_p2_score   = 0.09,
+)
+
+# ──────────────────────────────────────────────────────────────────────
+# Circuit breaker — no-change scenario (the patch grader should reject any
+# proposed diff and reward `declare_no_change`).
+# ──────────────────────────────────────────────────────────────────────
+
+CIRCUIT_BREAKER_CODE_CONTEXT = CodeContext(
+    repo_snapshot_path  = _snap("orders_v300"),
+    bad_commit_sha      = "d2b9c11",
+    ground_truth_files  = [],
+    ground_truth_diff   = "",
+    is_valid_issue      = False,
+    expected_p2_steps   = 6,
+    null_context_p2_score = 0.15,
+)
+
+
+# ──────────────────────────────────────────────────────────────────────
+# Phase-B scenarios (RL-discoverable, the brief's four types)
+# ──────────────────────────────────────────────────────────────────────
+
+ALIASED_FAULT_CODE_CONTEXT = CodeContext(
+    repo_snapshot_path  = _snap("queue_v210"),
+    bad_commit_sha      = "e1f4a02",
+    ground_truth_files  = ["queue/worker.py"],
+    ground_truth_diff   = """--- a/queue/worker.py
++++ b/queue/worker.py
+@@ -22,4 +22,5 @@ class CacheWriter:
+     def flush(self, batch):
+-        for k, v in batch.items():
++        for k, v in list(batch.items()):
+             self._cache.set(k, v)
++        batch.clear()
+""",
+    is_valid_issue          = True,
+    expected_p2_steps       = 7,
+    null_context_p2_score   = 0.16,
+)
+
+SEVERITY_INVERSION_CODE_CONTEXT = CodeContext(
+    repo_snapshot_path  = _snap("orders_retry_storm"),
+    bad_commit_sha      = "f8c9b13",
+    ground_truth_files  = ["orders/auth_client.py"],
+    ground_truth_diff   = """--- a/orders/auth_client.py
++++ b/orders/auth_client.py
+@@ -15,5 +15,6 @@ class AuthClient:
+     def validate(self, token):
+-        return self._call_with_retries(token, retries=20)
++        return self._call_with_retries(token, retries=2,
++                                        backoff_seconds=0.5)
+""",
+    is_valid_issue          = True,
+    expected_p2_steps       = 8,
+    null_context_p2_score   = 0.12,
+)
+
+CONFIDENCE_INVERSION_CODE_CONTEXT = CodeContext(
+    repo_snapshot_path  = _snap("payment_threadpool"),
+    bad_commit_sha      = "11abf04",
+    ground_truth_files  = ["payment/threadpool.py"],
+    ground_truth_diff   = """--- a/payment/threadpool.py
++++ b/payment/threadpool.py
+@@ -8,4 +8,5 @@ class PoolWorker:
+     def acquire(self):
+-        self._lock_a.acquire()
+-        self._lock_b.acquire()
++        with self._global_order:
++            self._lock_a.acquire()
++            self._lock_b.acquire()
+""",
+    is_valid_issue          = True,
+    expected_p2_steps       = 9,
+    null_context_p2_score   = 0.10,
+)
+
+INFO_ORDERING_CODE_CONTEXT = CodeContext(
+    repo_snapshot_path  = _snap("shared_libs_dep"),
+    bad_commit_sha      = "9d2e7af",
+    ground_truth_files  = ["requirements.txt", "shared/serializer.py"],
+    ground_truth_diff   = """--- a/requirements.txt
++++ b/requirements.txt
+@@ -3,1 +3,1 @@
+-shared-serializer==1.4.2
++shared-serializer==1.3.0
+""",
+    is_valid_issue          = True,
+    expected_p2_steps       = 9,
+    null_context_p2_score   = 0.11,
+)
+
+
+# ──────────────────────────────────────────────────────────────────────
+# Pool-D held-out scenarios
+# ──────────────────────────────────────────────────────────────────────
+
+HELDOUT_ALIASED_SEVERITY_CODE_CONTEXT = CodeContext(
+    repo_snapshot_path  = _snap("orders_retry_storm"),
+    bad_commit_sha      = "d09a4f1",
+    ground_truth_files  = ["orders/auth_client.py"],
+    ground_truth_diff   = """--- a/orders/auth_client.py
++++ b/orders/auth_client.py
+@@ -15,5 +15,6 @@ class AuthClient:
+     def validate(self, token):
+-        return self._call_with_retries(token, retries=25)
++        return self._call_with_retries(token, retries=2,
++                                        backoff_seconds=0.5)
+""",
+    is_valid_issue          = True,
+    expected_p2_steps       = 9,
+    null_context_p2_score   = 0.10,
+)
+
+HELDOUT_CONFIDENCE_ORDERING_CODE_CONTEXT = CodeContext(
+    repo_snapshot_path  = _snap("shared_libs_dep"),
+    bad_commit_sha      = "9d2e7af",
+    ground_truth_files  = ["requirements.txt"],
+    ground_truth_diff   = """--- a/requirements.txt
++++ b/requirements.txt
+@@ -3,1 +3,1 @@
+-shared-serializer==1.4.2
++shared-serializer==1.3.0
+""",
+    is_valid_issue          = True,
+    expected_p2_steps       = 10,
+    null_context_p2_score   = 0.10,
+)
+
+
+# ──────────────────────────────────────────────────────────────────────
+# Lookup helpers (used by Pool-B baseline runner to write baselines back)
+# ──────────────────────────────────────────────────────────────────────
+
+CODE_CONTEXTS = {
+    "memory_leak":                   MEMORY_LEAK_CODE_CONTEXT,
+    "cascading_failure":             CASCADING_FAILURE_CODE_CONTEXT,
+    "distributed_deadlock":          DISTRIBUTED_DEADLOCK_CODE_CONTEXT,
+    "circuit_breaker_noop":          CIRCUIT_BREAKER_CODE_CONTEXT,
+    "aliased_fault":                 ALIASED_FAULT_CODE_CONTEXT,
+    "severity_inversion":            SEVERITY_INVERSION_CODE_CONTEXT,
+    "confidence_inversion":          CONFIDENCE_INVERSION_CODE_CONTEXT,
+    "info_ordering":                 INFO_ORDERING_CODE_CONTEXT,
+    "heldout_aliased_severity":      HELDOUT_ALIASED_SEVERITY_CODE_CONTEXT,
+    "heldout_confidence_ordering":   HELDOUT_CONFIDENCE_ORDERING_CODE_CONTEXT,
+}
+
+
+def update_null_baseline(task_name: str, score: float) -> None:
+    """Mutate the in-process null-context baseline for `task_name`."""
+    ctx = CODE_CONTEXTS.get(task_name)
+    if ctx is None:
+        raise KeyError(f"No code context for task {task_name}")
+    ctx.null_context_p2_score = float(score)

scenarios/confidence_inversion.py

@@ -0,0 +1,131 @@
+"""
+Phase-B scenario 3 — Confidence inversion.
+
+The symptom pattern most strongly resembles a memory leak (smoothly
+climbing memory, OOM-flavoured logs, p99 latency drift) but the
+deploy history shows nothing memory-related and the *real* root cause
+is a distributed deadlock — accumulating threads pinned waiting on a
+lock cycle.  RAM grows because thread stacks accumulate, not because
+of a heap leak.
+
+Why this scenario matters:
+  - The ops agent's belief should be LOW confidence even though the
+    surface evidence is "high confidence memory leak".
+  - The orchestrator must learn to *not* transition with high
+    confidence based on apparent-evidence — it must actively check
+    deploys, then thread metrics, then escalate.
+
+This directly stresses the stopping criterion: keep investigating
+despite clean-looking symptoms.
+"""
+
+from __future__ import annotations
+
+import random
+from typing import Dict, List, Set
+
+from .base import BaseScenario
+from .code_context_builder import CONFIDENCE_INVERSION_CODE_CONTEXT
+from ..models import CodeContext
+from ..simulation.infrastructure import Infrastructure
+from ..simulation.service import Deploy
+from ..simulation.metrics import generate_memory_leak_history
+
+
+class ConfidenceInversionScenario(BaseScenario):
+
+    @property
+    def task_name(self) -> str:
+        return "confidence_inversion"
+
+    @property
+    def code_context(self) -> CodeContext:
+        return CONFIDENCE_INVERSION_CODE_CONTEXT
+
+    @property
+    def fault_class(self) -> str:
+        return "deadlock"
+
+    @property
+    def display_name(self) -> str:
+        return "Confidence Inversion — Deadlock Masquerading as Memory Leak"
+
+    @property
+    def incident_summary(self) -> str:
+        return (
+            "INCIDENT: Payment service memory has climbed from 38% to 81% over the "
+            "last 25 minutes. Latency p99 drifting upward. OOM-flavoured exceptions "
+            "appearing in logs. Looks like a textbook memory leak."
+        )
+
+    @property
+    def severity(self) -> str:
+        return "SEV2"
+
+    @property
+    def correct_root_cause(self) -> str:
+        return "payment threadpool deadlock — accumulating pinned threads, not a heap leak"
+
+    @property
+    def root_cause_keywords(self) -> List[str]:
+        return ["payment", "deadlock", "thread", "lock"]
+
+    @property
+    def involved_services(self) -> Set[str]:
+        return {"payment"}
+
+    @property
+    def root_cause_service(self) -> str:
+        return "payment"
+
+    @property
+    def correct_remediation_actions(self) -> List[Dict[str, str]]:
+        return [
+            {"action_type": "rollback_deploy", "target_service": "payment"},
+            {"action_type": "restart_service", "target_service": "payment"},
+        ]
+
+    def inject(self, infra: Infrastructure) -> None:
+        payment = infra.get_service("payment")
+        if payment is None:
+            return
+
+        # --- The "innocent" recent deploy: a config change unrelated to memory.
+        #     The agent who only checks deploy_history will see this and
+        #     correctly note: nothing memory-shaped here.  That's the trap —
+        #     the absence of a memory deploy is *evidence* the symptom is
+        #     misleading, but the base-model will dismiss it.
+        payment.deploy_history.append(Deploy(
+            version="v3.1.2", timestamp_minutes=-25, author="dave",
+            commit_hash="11abf04",
+            description="Refactor: lock acquisition order in PoolWorker",
+            is_bad=True,
+        ))
+
+        payment.inject_fault("memory_leak", rate=1.7)
+        payment.metric_history = generate_memory_leak_history(
+            minutes=30, start_minute=0, leak_start_offset=8, rate=1.7)
+        payment.memory_percent      = 81.0 + random.gauss(0, 2)
+        payment.cpu_percent         = 22.0 + random.gauss(0, 2)   # ← cpu LOW (deadlock signature)
+        payment.error_rate_percent  = 9.0  + random.gauss(0, 1)
+        payment.latency_p95_ms      = 1800.0 + random.gauss(0, 100)
+        payment.latency_p99_ms      = 3400.0 + random.gauss(0, 200)
+        payment.requests_per_sec    = 110.0 + random.gauss(0, 8)  # ← throughput collapsed
+        payment.status              = "degraded"
+
+        payment.logs = [
+            {"timestamp": "2025-01-15T14:09:00Z", "level": "INFO", "service": "payment",
+             "message": "Deploy v3.1.2 complete — pool worker refactor live", "trace_id": None},
+            {"timestamp": "2025-01-15T14:14:00Z", "level": "WARN", "service": "payment",
+             "message": "GC pressure: heap usage at 64%, GC pause 290ms (mostly old-gen)",
+             "trace_id": None},
+            {"timestamp": "2025-01-15T14:18:00Z", "level": "WARN", "service": "payment",
+             "message": "Thread pool active count: 198/200 (CPU usage low — threads blocked)",
+             "trace_id": None},
+            {"timestamp": "2025-01-15T14:23:00Z", "level": "ERROR", "service": "payment",
+             "message": "Lock contention: PoolWorker.acquire blocked for 8400ms on lock_b",
+             "trace_id": "trace-883011"},
+            {"timestamp": "2025-01-15T14:25:00Z", "level": "WARN", "service": "payment",
+             "message": "OutOfMemoryError-like: thread stack pool exhausted",
+             "trace_id": None},
+        ]

scenarios/easy_memory_leak.py

@@ -12,6 +12,8 @@ import random
 from typing import Dict, List, Set
 
 from .base import BaseScenario
+from .code_context_builder import MEMORY_LEAK_CODE_CONTEXT
+from ..models import CodeContext
 from ..simulation.infrastructure import Infrastructure
 from ..simulation.service import Deploy
 from ..simulation.metrics import generate_memory_leak_history
@@ -23,6 +25,14 @@ class MemoryLeakScenario(BaseScenario):
     def task_name(self) -> str:
         return "memory_leak"
 
+    @property
+    def code_context(self) -> CodeContext:
+        return MEMORY_LEAK_CODE_CONTEXT
+
+    @property
+    def fault_class(self) -> str:
+        return "memory_leak"
+
     @property
     def display_name(self) -> str:
         return "Memory Leak — Orders Service"

scenarios/grader_p2.py

@@ -0,0 +1,103 @@
+"""
+Phase 2 grader — oracle-independent.
+Takes only the P2 trajectory (List[StepRecord]) and the declared patch/no-change.
+Scores patch quality using three tiers.
+"""
+
+import ast
+import difflib
+from typing import List, Optional
+from ..models import StepRecord, CodeContext
+
+
+def grade_patch_quality(proposed_diff: str, ctx: CodeContext) -> float:
+    """
+    Three-tier patch scoring:
+      Tier 1 (40%): file overlap — did they touch the right files?
+      Tier 2 (30%): AST hunk similarity — do the changed functions match?
+      Tier 3 (30%): syntax validity — does the patch parse cleanly?
+    """
+    proposed_files = _extract_files_from_diff(proposed_diff)
+    ground_truth_files = set(ctx.ground_truth_files)
+
+    # Tier 1
+    if not ground_truth_files:
+        file_score = 0.0
+    else:
+        intersection = proposed_files & ground_truth_files
+        union = proposed_files | ground_truth_files
+        file_score = len(intersection) / len(union) if union else 0.0
+
+    # Tier 2
+    hunk_score = _ast_hunk_similarity(proposed_diff, ctx.ground_truth_diff)
+
+    # Tier 3
+    syntax_score = 1.0 if _patch_parses_cleanly(proposed_diff) else 0.0
+
+    return 0.4 * file_score + 0.3 * hunk_score + 0.3 * syntax_score
+
+
+def grade_no_change(declared: bool, ctx: CodeContext) -> float:
+    """1.0 if agent correctly identified spurious issue, 0.0 otherwise."""
+    if not ctx.is_valid_issue and declared:
+        return 1.0
+    if ctx.is_valid_issue and not declared:
+        return 0.0
+    return 0.0   # wrong in either direction
+
+
+def grade_p2_efficiency(p2_steps: int, expected_steps: int) -> float:
+    """
+    Normalized efficiency — doesn't penalize inherently hard bugs.
+    Score = 1.0 at expected_steps, decays to 0 at 2x expected.
+    """
+    ratio = p2_steps / max(expected_steps, 1)
+    return max(0.0, 1.0 - max(0.0, ratio - 1.0))
+
+
+def _extract_files_from_diff(diff: str) -> set:
+    files = set()
+    for line in diff.split("\n"):
+        if line.startswith("+++ b/"):
+            files.add(line[6:].strip())
+    return files
+
+
+def _ast_hunk_similarity(proposed: str, ground_truth: str) -> float:
+    """
+    Extract (file, function_name) pairs from both diffs.
+    Score = Jaccard overlap of those sets.
+    """
+    proposed_fns = _extract_changed_functions(proposed)
+    truth_fns = _extract_changed_functions(ground_truth)
+    if not truth_fns:
+        return 1.0  # trivial patch, full credit if syntax valid
+    intersection = proposed_fns & truth_fns
+    union = proposed_fns | truth_fns
+    return len(intersection) / len(union) if union else 0.0
+
+
+def _extract_changed_functions(diff: str) -> set:
+    """Parse diff hunks, extract function names via simple @@ line parsing."""
+    fns = set()
+    current_file = ""
+    for line in diff.split("\n"):
+        if line.startswith("+++ b/"):
+            current_file = line[6:].strip()
+        elif line.startswith("@@") and "def " in line:
+            # hunk header often contains function context
+            parts = line.split("def ")
+            if len(parts) > 1:
+                fn_name = parts[1].split("(")[0].strip()
+                fns.add(f"{current_file}:{fn_name}")
+    return fns
+
+
+def _patch_parses_cleanly(diff: str) -> bool:
+    """Extract added lines, try to parse as Python."""
+    added = [l[1:] for l in diff.split("\n") if l.startswith("+") and not l.startswith("+++")]
+    try:
+        ast.parse("\n".join(added))
+        return True
+    except SyntaxError:
+        return False

scenarios/hard_distributed_deadlock.py

@@ -18,6 +18,8 @@ import random
 from typing import Dict, List, Set
 
 from .base import BaseScenario
+from .code_context_builder import DISTRIBUTED_DEADLOCK_CODE_CONTEXT
+from ..models import CodeContext
 from ..simulation.infrastructure import Infrastructure
 from ..simulation.service import Deploy
 from ..simulation.metrics import generate_high_latency_history, generate_healthy_history
@@ -29,6 +31,14 @@ class DistributedDeadlockScenario(BaseScenario):
     def task_name(self) -> str:
         return "distributed_deadlock"
 
+    @property
+    def code_context(self) -> CodeContext:
+        return DISTRIBUTED_DEADLOCK_CODE_CONTEXT
+
+    @property
+    def fault_class(self) -> str:
+        return "deadlock"
+
     @property
     def display_name(self) -> str:
         return "Distributed Deadlock — Payment/Orders/Queue Circular Wait"

scenarios/heldout.py

@@ -0,0 +1,249 @@
+"""
+Pool D — held-out scenarios.
+
+These combine fault families the agent saw individually during training
+into novel compound scenarios that test whether the *strategy* generalized
+(rather than just memorising scenario fingerprints).
+
+Each held-out scenario reuses the same code-context infrastructure (real
+mini-repo snapshots, ground-truth diffs) but in a configuration the agent
+has *never* seen during training.
+"""
+
+from __future__ import annotations
+
+import random
+from typing import Dict, List, Set
+
+from .base import BaseScenario
+from .code_context_builder import (
+    HELDOUT_ALIASED_SEVERITY_CODE_CONTEXT,
+    HELDOUT_CONFIDENCE_ORDERING_CODE_CONTEXT,
+)
+from ..models import CodeContext
+from ..simulation.infrastructure import Infrastructure
+from ..simulation.service import Deploy
+from ..simulation.metrics import (
+    generate_memory_leak_history,
+    generate_error_spike_history,
+    generate_high_latency_history,
+)
+
+
+# ──────────────────────────────────────────────────────────────────────
+# 1. aliased + severity_inversion combo
+#
+# A retry storm in `orders` saturates `auth` (severity inversion), AND the
+# saturation causes `auth`'s memory to climb because thread queues build up
+# (aliased symptom — looks like memory leak in auth).
+#
+# Diagnostic strategy must combine: skip the loud service AND check
+# dependencies.  Neither aliased_fault nor severity_inversion alone teach
+# both rules simultaneously.
+# ──────────────────────────────────────────────────────────────────────
+
+
+class HeldoutAliasedSeverityScenario(BaseScenario):
+
+    @property
+    def task_name(self) -> str:
+        return "heldout_aliased_severity"
+
+    @property
+    def code_context(self) -> CodeContext:
+        return HELDOUT_ALIASED_SEVERITY_CODE_CONTEXT
+
+    @property
+    def fault_class(self) -> str:
+        return "retry_storm"
+
+    @property
+    def display_name(self) -> str:
+        return "[Held-out] Aliased + Severity Inversion"
+
+    @property
+    def incident_summary(self) -> str:
+        return (
+            "INCIDENT: Auth service is firing CRITICAL alerts (HighMemoryUsage, "
+            "HighErrorRate). It looks like a memory leak in auth. Orders shows "
+            "only a WARN-level retry-counter alert."
+        )
+
+    @property
+    def severity(self) -> str:
+        return "SEV1"
+
+    @property
+    def correct_root_cause(self) -> str:
+        return "orders auth-client retry storm — auth is the victim, memory growth is queue build-up"
+
+    @property
+    def root_cause_keywords(self) -> List[str]:
+        return ["orders", "retry", "auth-client", "storm"]
+
+    @property
+    def involved_services(self) -> Set[str]:
+        return {"orders", "auth"}
+
+    @property
+    def root_cause_service(self) -> str:
+        return "orders"
+
+    @property
+    def correct_remediation_actions(self) -> List[Dict[str, str]]:
+        return [
+            {"action_type": "rollback_deploy", "target_service": "orders"},
+            {"action_type": "restart_service", "target_service": "auth"},
+        ]
+
+    def inject(self, infra: Infrastructure) -> None:
+        orders = infra.get_service("orders")
+        auth   = infra.get_service("auth")
+        if orders is None or auth is None:
+            return
+
+        orders.deploy_history.append(Deploy(
+            version="v2.6.0", timestamp_minutes=-14, author="frankie",
+            commit_hash="d09a4f1",
+            description="Resilience: bump auth-client retries 3 → 25 with no jitter",
+            is_bad=True,
+        ))
+        # Auth has no recent change
+        auth.deploy_history.append(Deploy(
+            version="v1.7.4", timestamp_minutes=-86400, author="bob",
+            commit_hash="a01122", description="Routine: bump TLS cert", is_bad=False,
+        ))
+
+        # Auth shows BOTH high memory and high error rate (aliased pattern!)
+        auth.inject_fault("memory_leak", rate=0.9)
+        auth.inject_fault("high_error_rate", rate=58.0)
+        auth.metric_history = generate_memory_leak_history(
+            minutes=30, start_minute=0, leak_start_offset=14, rate=0.9)
+        auth.memory_percent      = 79.0 + random.gauss(0, 2)
+        auth.error_rate_percent  = 58.0 + random.gauss(0, 3)
+        auth.latency_p95_ms      = 3500 + random.gauss(0, 200)
+        auth.latency_p99_ms      = 6200 + random.gauss(0, 400)
+        auth.requests_per_sec    = 4100 + random.gauss(0, 200)
+        auth.status              = "down"
+        auth.logs = [
+            {"timestamp": "2025-01-15T14:13:00Z", "level": "WARN", "service": "auth",
+             "message": "Throughput surged 1300 → 4100 RPS in 4 min", "trace_id": None},
+            {"timestamp": "2025-01-15T14:18:00Z", "level": "ERROR", "service": "auth",
+             "message": "OOM-like: token validation queue at 11400 entries",
+             "trace_id": "trace-771233"},
+        ]
+
+        orders.error_rate_percent = 1.8 + random.gauss(0, 0.4)
+        orders.latency_p95_ms     = 130.0 + random.gauss(0, 10)
+        orders.status             = "degraded"
+        orders.logs = [
+            {"timestamp": "2025-01-15T14:14:00Z", "level": "INFO", "service": "orders",
+             "message": "Deploy v2.6.0 complete — auth-client retry policy updated",
+             "trace_id": None},
+            {"timestamp": "2025-01-15T14:18:00Z", "level": "WARN", "service": "orders",
+             "message": "auth-client retry counter elevated: avg 17 retries per validation",
+             "trace_id": None},
+        ]
+
+
+# ──────────────────────────────────────────────────────────────────────
+# 2. confidence_inversion + info_ordering combo
+#
+# Symptoms scream "memory leak in payment" but the cause is a shared
+# dependency downgrade that happened to interact badly with payment's
+# threadpool — so the fix is in `requirements.txt`, not in payment's
+# service code.  The base model with high confidence on memory_leak
+# will rollback payment, fail to fix it, then look at payment's deploy
+# history (which is clean), and run out of time.
+# ──────────────────────────────────────────────────────────────────────
+
+
+class HeldoutConfidenceOrderingScenario(BaseScenario):
+
+    @property
+    def task_name(self) -> str:
+        return "heldout_confidence_ordering"
+
+    @property
+    def code_context(self) -> CodeContext:
+        return HELDOUT_CONFIDENCE_ORDERING_CODE_CONTEXT
+
+    @property
+    def fault_class(self) -> str:
+        return "shared_dependency"
+
+    @property
+    def display_name(self) -> str:
+        return "[Held-out] Confidence + Info-Ordering Inversion"
+
+    @property
+    def incident_summary(self) -> str:
+        return (
+            "INCIDENT: Payment service memory at 83%, climbing. Latency p99 4200ms "
+            "and rising. Looks like a clear memory leak. No payment deploys in 24h."
+        )
+
+    @property
+    def severity(self) -> str:
+        return "SEV2"
+
+    @property
+    def correct_root_cause(self) -> str:
+        return "shared-serializer dependency downgrade interacting with payment threadpool"
+
+    @property
+    def root_cause_keywords(self) -> List[str]:
+        return ["shared", "dependency", "serializer", "thread"]
+
+    @property
+    def involved_services(self) -> Set[str]:
+        return {"payment"}
+
+    @property
+    def root_cause_service(self) -> str:
+        return "payment"
+
+    @property
+    def correct_remediation_actions(self) -> List[Dict[str, str]]:
+        return [
+            {"action_type": "rollback_deploy", "target_service": "payment"},
+        ]
+
+    def inject(self, infra: Infrastructure) -> None:
+        payment = infra.get_service("payment")
+        if payment is None:
+            return
+
+        payment.deploy_history.append(Deploy(
+            version="img-payment-732", timestamp_minutes=-9, author="lib-bot",
+            commit_hash="9d2e7af",
+            description="Image rebuild: pulled latest shared-serializer (downgraded)",
+            is_bad=True,
+        ))
+
+        payment.inject_fault("memory_leak", rate=1.6)
+        payment.inject_fault("high_latency", p99=4200)
+        payment.metric_history = generate_memory_leak_history(
+            minutes=30, start_minute=0, leak_start_offset=8, rate=1.6)
+        payment.memory_percent       = 83.0 + random.gauss(0, 2)
+        payment.cpu_percent          = 24.0 + random.gauss(0, 2)
+        payment.error_rate_percent   = 11.0 + random.gauss(0, 1)
+        payment.latency_p99_ms       = 4200.0 + random.gauss(0, 200)
+        payment.requests_per_sec     = 95.0 + random.gauss(0, 5)
+        payment.status               = "degraded"
+
+        payment.logs = [
+            {"timestamp": "2025-01-15T14:08:00Z", "level": "INFO", "service": "payment",
+             "message": "Image rebuild complete (shared-serializer pulled)",
+             "trace_id": None},
+            {"timestamp": "2025-01-15T14:14:00Z", "level": "WARN", "service": "payment",
+             "message": "Heap usage at 71%, GC pause 380ms",
+             "trace_id": None},
+            {"timestamp": "2025-01-15T14:19:00Z", "level": "ERROR", "service": "payment",
+             "message": "Deserialization failure: UnknownFieldError "
+                        "(check shared-serializer version)",
+             "trace_id": "trace-441312"},
+            {"timestamp": "2025-01-15T14:22:00Z", "level": "WARN", "service": "payment",
+             "message": "Thread pool saturated: 196/200 threads blocked on retry",
+             "trace_id": None},
+        ]

scenarios/info_ordering.py

@@ -0,0 +1,146 @@
+"""
+Phase-B scenario 4 — Information ordering dependency.
+
+Three services (`orders`, `payment`, `queue`) all degrade *simultaneously*
+because a shared dependency (`shared-serializer`) was downgraded in
+`requirements.txt` by an unrelated PR.  None of the three service-local
+deploy histories show anything related — `check_deploy_history(orders)`
+returns clean for each service.
+
+The ONLY way to find the cause is to ask: "what was the most recent
+commit that touched a file *every degraded service depends on*?"
+That requires checking the **shared dependency file's git log**, which
+the base model (following per-service deploy-history priors) won't do.
+
+This is the canonical "the base model has the wrong prior on which
+artifact to inspect first" scenario.  It's the one most likely to
+expose a difference between prompt-engineered baselines and an
+RL-trained agent that has internalized "correlated multi-service
+degradation → look at shared dependency".
+
+In Phase 2, the diff is in `requirements.txt` (not in any service's
+own source tree), so the code agent must learn to inspect dependency
+manifests as well as service-local code.
+"""
+
+from __future__ import annotations
+
+import random
+from typing import Dict, List, Set
+
+from .base import BaseScenario
+from .code_context_builder import INFO_ORDERING_CODE_CONTEXT
+from ..models import CodeContext
+from ..simulation.infrastructure import Infrastructure
+from ..simulation.service import Deploy
+from ..simulation.metrics import (
+    generate_error_spike_history,
+    generate_high_latency_history,
+)
+
+
+class InfoOrderingScenario(BaseScenario):
+
+    @property
+    def task_name(self) -> str:
+        return "info_ordering"
+
+    @property
+    def code_context(self) -> CodeContext:
+        return INFO_ORDERING_CODE_CONTEXT
+
+    @property
+    def fault_class(self) -> str:
+        return "shared_dependency"
+
+    @property
+    def display_name(self) -> str:
+        return "Info Ordering — Shared-Library Downgrade Hits Three Services"
+
+    @property
+    def incident_summary(self) -> str:
+        return (
+            "INCIDENT: Three independent services — orders, payment, queue — all "
+            "started erroring simultaneously 8 minutes ago. Errors are deserialization "
+            "failures (UnknownFieldError, version mismatch). No service was deployed "
+            "in the last 24 hours."
+        )
+
+    @property
+    def severity(self) -> str:
+        return "SEV1"
+
+    @property
+    def correct_root_cause(self) -> str:
+        return "shared-serializer dependency downgrade in requirements.txt — affects all three services"
+
+    @property
+    def root_cause_keywords(self) -> List[str]:
+        return ["shared", "serializer", "dependency", "requirements"]
+
+    @property
+    def involved_services(self) -> Set[str]:
+        return {"orders", "payment", "queue"}
+
+    @property
+    def root_cause_service(self) -> str:
+        # No single service deploy is to blame — but for remediation purposes
+        # we treat `orders` as the canonical target (rolling its image will
+        # pull the correct shared-serializer back in).
+        return "orders"
+
+    @property
+    def correct_remediation_actions(self) -> List[Dict[str, str]]:
+        return [
+            {"action_type": "rollback_deploy", "target_service": "orders"},
+            {"action_type": "rollback_deploy", "target_service": "payment"},
+            {"action_type": "rollback_deploy", "target_service": "queue"},
+        ]
+
+    def inject(self, infra: Infrastructure) -> None:
+        orders  = infra.get_service("orders")
+        payment = infra.get_service("payment")
+        queue   = infra.get_service("queue")
+        if not all([orders, payment, queue]):
+            return
+
+        # --- A *single* shared-deps PR was merged — it touches no service
+        # individually but affects all three image rebuilds.
+        # We attach the same Deploy stamp to each so the agent can see "all
+        # three have a deploy at the same minute" only after exhaustively
+        # checking each one.
+        for svc in (orders, payment, queue):
+            svc.deploy_history.append(Deploy(
+                version=f"img-{svc.name}-{random.randint(100, 999)}",
+                timestamp_minutes=-8, author="lib-bot",
+                commit_hash="9d2e7af",
+                description="Image rebuild: pulled latest shared-serializer (downgraded)",
+                is_bad=True,
+            ))
+
+        # --- Symptoms: all three show high error rate, similar log message
+        for svc, p99 in [(orders, 1500), (payment, 2200), (queue, 1800)]:
+            svc.inject_fault("high_error_rate", rate=24.0)
+            svc.error_rate_percent  = 24.0 + random.gauss(0, 3)
+            svc.latency_p95_ms      = p99 * 0.6 + random.gauss(0, 80)
+            svc.latency_p99_ms      = p99 + random.gauss(0, 100)
+            svc.status              = "degraded"
+            svc.metric_history      = generate_error_spike_history(
+                minutes=30, start_minute=0, spike_start_offset=22,
+                error_rate_target=24.0)
+
+        # --- Each service's logs say the same deserialisation error ---
+        for svc in (orders, payment, queue):
+            svc.logs = [
+                {"timestamp": "2025-01-15T14:21:30Z", "level": "INFO", "service": svc.name,
+                 "message": "Image rebuild complete — container restarted",
+                 "trace_id": None},
+                {"timestamp": "2025-01-15T14:22:30Z", "level": "ERROR", "service": svc.name,
+                 "message": ("Deserialization failed: UnknownFieldError 'event_v2.idempotency_key' "
+                             "(shared-serializer mismatch?)"),
+                 "trace_id": "trace-441100"},
+                {"timestamp": "2025-01-15T14:24:00Z", "level": "ERROR", "service": svc.name,
+                 "message": ("Schema version mismatch: expected 1.4.x, got 1.3.0 "
+                             "(check shared-serializer version)"),
+                 "trace_id": "trace-441101"},
+            ]

scenarios/medium_cascading_failure.py

@@ -14,6 +14,8 @@ import random
 from typing import Dict, List, Set
 
 from .base import BaseScenario
+from .code_context_builder import CASCADING_FAILURE_CODE_CONTEXT
+from ..models import CodeContext
 from ..simulation.infrastructure import Infrastructure
 from ..simulation.service import Deploy
 from ..simulation.metrics import generate_error_spike_history, generate_healthy_history
@@ -25,6 +27,14 @@ class CascadingFailureScenario(BaseScenario):
     def task_name(self) -> str:
         return "cascading_failure"
 
+    @property
+    def code_context(self) -> CodeContext:
+        return CASCADING_FAILURE_CODE_CONTEXT
+
+    @property
+    def fault_class(self) -> str:
+        return "config_change"
+
     @property
     def display_name(self) -> str:
         return "Cascading Failure — Auth Service Configuration"

scenarios/severity_inversion.py

@@ -0,0 +1,147 @@
+"""
+Phase-B scenario 2 — Misleading severity inversion.
+
+The highest-severity alert is on a service that is the *downstream victim*,
+not the cause.  A retry-storm in `orders` (caused by an aggressive auth
+client config) floods `auth` with traffic; `auth` falls over with
+CRITICAL alerts, while `orders` itself only shows WARN-level "elevated
+retry counter" alerts.
+
+A base model prompted to "follow the highest-severity alert" goes straight
+to `auth` and finds nothing — auth's own deploy history is clean.
+
+The RL signal: in retry-storm scenarios, the *quietest degraded service*
+is the culprit and the loudest is the victim.  This is a scenario-class
+specific reasoning rule no fixed prompt can encode.
+"""
+
+from __future__ import annotations
+
+import random
+from typing import Dict, List, Set
+
+from .base import BaseScenario
+from .code_context_builder import SEVERITY_INVERSION_CODE_CONTEXT
+from ..models import CodeContext
+from ..simulation.infrastructure import Infrastructure
+from ..simulation.service import Deploy
+from ..simulation.metrics import (
+    generate_error_spike_history,
+    generate_healthy_history,
+)
+
+
+class SeverityInversionScenario(BaseScenario):
+
+    @property
+    def task_name(self) -> str:
+        return "severity_inversion"
+
+    @property
+    def code_context(self) -> CodeContext:
+        return SEVERITY_INVERSION_CODE_CONTEXT
+
+    @property
+    def fault_class(self) -> str:
+        return "retry_storm"
+
+    @property
+    def display_name(self) -> str:
+        return "Severity Inversion — Auth Drowning, Orders Is the Culprit"
+
+    @property
+    def incident_summary(self) -> str:
+        return (
+            "INCIDENT: Auth service is firing CRITICAL alerts (HighErrorRate, "
+            "ServiceUnreachable, LatencyP99>5000ms). Customer logins are failing. "
+            "Orders service shows a WARN-level alert: 'Elevated retry counter'."
+        )
+
+    @property
+    def severity(self) -> str:
+        return "SEV1"
+
+    @property
+    def correct_root_cause(self) -> str:
+        return "orders auth-client retry storm overwhelming auth — orders deploy is root cause"
+
+    @property
+    def root_cause_keywords(self) -> List[str]:
+        return ["orders", "retry", "storm", "auth-client"]
+
+    @property
+    def involved_services(self) -> Set[str]:
+        return {"orders", "auth"}
+
+    @property
+    def root_cause_service(self) -> str:
+        return "orders"
+
+    @property
+    def correct_remediation_actions(self) -> List[Dict[str, str]]:
+        return [
+            {"action_type": "rollback_deploy", "target_service": "orders"},
+            {"action_type": "restart_service", "target_service": "auth"},
+        ]
+
+    def inject(self, infra: Infrastructure) -> None:
+        orders = infra.get_service("orders")
+        auth   = infra.get_service("auth")
+        if orders is None or auth is None:
+            return
+
+        # --- The actual bad deploy: orders changed auth-client retry policy ---
+        orders.deploy_history.append(Deploy(
+            version="v2.5.1", timestamp_minutes=-12, author="carol",
+            commit_hash="f8c9b13",
+            description="Resilience: increase auth-client retries 3 → 20",
+            is_bad=True,
+        ))
+
+        # --- Auth has a CLEAN recent deploy history (no smoking gun there) ---
+        auth.deploy_history.append(Deploy(
+            version="v1.7.4", timestamp_minutes=-86400, author="bob",
+            commit_hash="a01122", description="Routine: bump TLS cert",
+            is_bad=False,
+        ))
+
+        # --- Auth is the loud victim: high error rate, big latency ---
+        auth.inject_fault("high_error_rate", rate=72.0)
+        auth.error_rate_percent = 72.0 + random.gauss(0, 3)
+        auth.latency_p95_ms     = 4200 + random.gauss(0, 200)
+        auth.latency_p99_ms     = 7100 + random.gauss(0, 400)
+        auth.requests_per_sec   = 4500 + random.gauss(0, 200)   # ← anomalously high RPS!
+        auth.status             = "down"
+        auth.metric_history     = generate_error_spike_history(
+            minutes=30, start_minute=0, spike_start_offset=12, error_rate_target=72.0)
+        auth.logs = [
+            {"timestamp": "2025-01-15T14:14:00Z", "level": "WARN", "service": "auth",
+             "message": "Request throughput tripled in last 4 minutes (1500 → 4400 RPS)",
+             "trace_id": None},
+            {"timestamp": "2025-01-15T14:18:00Z", "level": "ERROR", "service": "auth",
+             "message": "Token validation queue overflow — dropping requests",
+             "trace_id": "trace-991201"},
+            {"timestamp": "2025-01-15T14:21:00Z", "level": "ERROR", "service": "auth",
+             "message": "Health check failed: auth returned HTTP 500 (overload)",
+             "trace_id": None},
+        ]
+
+        # --- Orders looks almost healthy; only subtle clue is the retry counter ---
+        orders.metric_history = generate_healthy_history(minutes=30, start_minute=0)
+        orders.cpu_percent          = 38.0 + random.gauss(0, 3)
+        orders.memory_percent       = 47.0 + random.gauss(0, 2)
+        orders.error_rate_percent   = 1.5  + random.gauss(0, 0.3)
+        orders.latency_p95_ms       = 110.0 + random.gauss(0, 8)
+        orders.latency_p99_ms       = 240.0 + random.gauss(0, 15)
+        orders.status               = "degraded"   # ← quiet degradation, easy to miss
+        orders.logs = [
+            {"timestamp": "2025-01-15T14:13:00Z", "level": "INFO", "service": "orders",
+             "message": "Deploy v2.5.1 complete — auth-client retry policy updated",
+             "trace_id": None},
+            {"timestamp": "2025-01-15T14:17:00Z", "level": "WARN", "service": "orders",
+             "message": "auth-client retry counter elevated: avg 12 retries per validation",
+             "trace_id": None},
+            {"timestamp": "2025-01-15T14:22:00Z", "level": "WARN", "service": "orders",
+             "message": "auth-client circuit-breaker NOT tripped (retries policy ignores it)",
+             "trace_id": None},
+        ]

server/app.py

@@ -1,51 +1,64 @@
 """
 Thin FastAPI server — marshals JSON in/out.
 No simulation logic lives here.
+
+Endpoints:
+  GET  /health                      health check
+  GET  /tasks                       list available scenarios
+  POST /reset    {task_name, seed}  start a new episode
+  POST /step     {action_type, ...} execute one action (phase-aware)
+  GET  /state                       per-episode metadata
+  GET  /trajectory                  full P1+P2 step records
+  POST /score    {declared_patch, declared_no_change, belief_history}
+                                    unified grader breakdown
 """
 
 from __future__ import annotations
 
-from typing import Any, Dict, Optional
+from dataclasses import asdict
+from typing import Any, Dict, List, Optional
 
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, Request, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 
 from .incident_environment import IncidentEnvironment
 
+
 # ------------------------------------------------------------------
 # App
 # ------------------------------------------------------------------
 
 app = FastAPI(
-    title="SRE Incident Response Environment",
-    description="An OpenEnv environment for training AI agents on production incident response.",
-    version="0.1.0",
+    title       = "SRE Incident Response Environment",
+    description = "Two-phase OpenEnv environment (P1 ops + P2 code attribution).",
+    version     = "0.2.0",
 )
 
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],
-    allow_methods=["*"],
-    allow_headers=["*"],
+    allow_origins = ["*"],
+    allow_methods = ["*"],
+    allow_headers = ["*"],
 )
 
 env = IncidentEnvironment()
 
 
 # ------------------------------------------------------------------
-# Request / Response models (thin wrappers)
+# Request models
 # ------------------------------------------------------------------
 
-class ResetRequest(BaseModel):
-    task_name: Optional[str] = None
-    seed: Optional[int] = None
+class StepRequest(BaseModel):
+    action_type:    str
+    target_service: Optional[str]      = None
+    parameters:     Dict[str, Any]     = {}
 
 
-class StepRequest(BaseModel):
-    action_type: str
-    target_service: Optional[str] = None
-    parameters: Dict[str, Any] = {}
+class ScoreRequest(BaseModel):
+    declared_patch:     Optional[str]   = None
+    declared_no_change: bool            = False
+    belief_history:     List[Dict[str, Any]] = []
 
 
 # ------------------------------------------------------------------
@@ -54,76 +67,153 @@ class StepRequest(BaseModel):
 
 @app.get("/health")
 def health() -> Dict[str, str]:
-    """Health check — the validator pings this first."""
     return {"status": "healthy"}
 
 
-from fastapi import Request
-
 @app.post("/reset")
 async def reset(request: Request) -> Dict[str, Any]:
     """
     Initialize a new incident episode.
-    POST /reset {"task_name": "memory_leak", "seed": 42}
+
+    Accepts (all optional):
+        task_name : str       specific scenario, otherwise sampled from pool
+        seed      : int       RNG seed for deterministic replay
+        pool      : "A"|"B"|"C"|"D"   selects training pool (sets default mode)
+        mode      : "p1_only"|"p2_only"|"joint"   force episode mode
     """
     try:
         body = await request.json()
     except Exception:
         body = {}
-        
     if not isinstance(body, dict):
         body = {}
-        
-    result = env.reset(
-        task_name=body.get("task_name"),
-        seed=body.get("seed"),
+    return env.reset(
+        task_name = body.get("task_name"),
+        seed      = body.get("seed"),
+        pool      = body.get("pool"),
+        mode      = body.get("mode"),
     )
-    return result
+
+
+@app.get("/pools")
+def list_pools() -> Dict[str, Any]:
+    """Pool registry — used by training runners to discover task names."""
+    from ..pools import POOLS
+    return {
+        name: {
+            "name":        p.name,
+            "description": p.description,
+            "task_names":  list(p.task_names),
+            "mode":        p.mode,
+            "inject_oracle_belief": p.inject_oracle_belief,
+        }
+        for name, p in POOLS.items()
+    }
 
 
 @app.post("/step")
 def step(request: StepRequest) -> Dict[str, Any]:
-    """
-    Execute one agent action.
-    POST /step {"action_type": "view_alerts"}
-    """
-    action_data = {
-        "action_type": request.action_type,
+    """Execute one agent action — phase-aware dispatch."""
+    return env.step({
+        "action_type":    request.action_type,
         "target_service": request.target_service,
-        "parameters": request.parameters,
-    }
-    result = env.step(action_data)
-    return result
+        "parameters":     request.parameters or {},
+    })
 
 
 @app.get("/state")
 def state() -> Dict[str, Any]:
+    return env.get_state()
+
+
+@app.get("/trajectory")
+def trajectory() -> Dict[str, Any]:
+    """Return the current episode's full P1 + P2 trajectory."""
+    return {
+        "p1": [_serialize_step(r) for r in env.get_p1_trajectory()],
+        "p2": [_serialize_step(r) for r in env.get_p2_trajectory()],
+    }
+
+
+@app.post("/score")
+def score(req: ScoreRequest) -> Dict[str, Any]:
     """
-    Get current episode metadata.
-    GET /state
+    Unified grader breakdown + counterfactual r_cross.
+
+    Returns:
+        final, p1_rca, p1_efficiency, patch_quality, no_change_detection,
+        p2_efficiency, r_cross, null_context_p2_score
     """
-    return env.get_state()
+    from ..tasks import compute_r_cross
+    breakdown = env.score_unified(belief_history=req.belief_history)
+    state = env.get_state()
+    task = state.get("task_name")
+    r_cross = 0.0
+    null_baseline = 0.0
+    if task:
+        try:
+            r_cross = compute_r_cross(
+                task_name          = task,
+                declared_patch     = state.get("declared_patch"),
+                declared_no_change = bool(state.get("declared_no_change")),
+                p2_trajectory      = env.get_p2_trajectory(),
+            )
+            from ..tasks import get_scenario
+            ctx = get_scenario(task).code_context
+            if ctx is not None:
+                null_baseline = float(ctx.null_context_p2_score)
+        except Exception:
+            pass
+    return {
+        **breakdown,
+        "r_cross":               round(r_cross, 4),
+        "null_context_p2_score": round(null_baseline, 4),
+    }
 
 
 @app.get("/tasks")
 def list_tasks() -> Dict[str, Any]:
-    """List available tasks with descriptions."""
     from ..tasks import TASK_REGISTRY
-    tasks = {}
+    out: Dict[str, Any] = {}
     for name, cls in TASK_REGISTRY.items():
         scenario = cls()
-        tasks[name] = {
-            "display_name": scenario.display_name,
-            "severity": scenario.severity,
-            "max_steps": scenario.max_steps,
+        out[name] = {
+            "display_name":        scenario.display_name,
+            "severity":            scenario.severity,
+            "max_steps":           scenario.max_steps,
             "time_budget_minutes": scenario.time_budget_minutes,
+            "has_phase2":          scenario.code_context is not None,
+            "fault_class":         scenario.fault_class,
         }
-    return {"tasks": tasks}
+    return {"tasks": out}
 
-def main():
+
+# ------------------------------------------------------------------
+# Helpers
+# ------------------------------------------------------------------
+
+def _serialize_step(r) -> Dict[str, Any]:
+    """Convert a StepRecord into a JSON-safe dict."""
+    return {
+        "step_number":           r.step_number,
+        "phase":                 r.phase,
+        "action": {
+            "action_type":    r.action.action_type,
+            "target_service": r.action.target_service,
+            "parameters":     r.action.parameters,
+        },
+        "reward":                r.reward,
+        "observation_summary":   r.observation_summary,
+        "service_statuses_after": r.service_statuses_after,
+        "timestamp_minutes":     r.timestamp_minutes,
+        "belief_state_snapshot": r.belief_state_snapshot,
+    }
+
+
+def main() -> None:
     import uvicorn
     uvicorn.run("incident_env.server.app:app", host="0.0.0.0", port=8000, reload=False)
 
+
 if __name__ == "__main__":
     main()
-

server/code_workspace.py

@@ -0,0 +1,308 @@
+"""
+CodeWorkspace — read-only filesystem + fake-git layer for Phase 2 exploration.
+
+Each scenario points `CodeContext.repo_snapshot_path` at a directory under
+`snapshots/`. That directory contains:
+
+  snapshots/<name>/
+      tree/             ← actual source files the agent reads
+          <pkg>/<file>.py
+          ...
+      git_log.json      ← list of commits (sha, author, date, message, files[])
+      diffs/<sha>.patch ← unified diff for that commit (any file path)
+
+This is a "fake git" by design — it's tighter, deterministic, and trivially
+serializable for trajectory replay.  No subprocess, no real .git directory.
+
+CodeWorkspace is constructed at the start of Phase 2 and lives for the rest
+of the episode.  It exposes safe, sandboxed file access (no `..`, no absolute
+paths, no symlinks).
+"""
+
+from __future__ import annotations
+
+import fnmatch
+import json
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+
+@dataclass
+class CommitRecord:
+    sha: str
+    author: str
+    date: str            # ISO-ish string
+    message: str
+    files: List[str]
+
+
+class CodeWorkspaceError(Exception):
+    """Raised on illegal path access or missing files (returned to agent)."""
+
+
+class CodeWorkspace:
+    """
+    Sandboxed read-only view over a snapshot.
+
+    All paths are interpreted as relative to `tree/` inside the snapshot.
+    Access to anything outside `tree/` raises CodeWorkspaceError.
+    """
+
+    MAX_FILE_BYTES   = 64 * 1024     # truncate large files in read_file()
+    MAX_LIST_ENTRIES = 200
+    MAX_SEARCH_HITS  = 50
+    MAX_SEARCH_BYTES = 5 * 1024 * 1024  # don't grep monsters
+
+    def __init__(self, snapshot_root: str, bad_commit_sha: str = ""):
+        root = Path(snapshot_root).resolve()
+        if not root.exists():
+            raise CodeWorkspaceError(f"Snapshot not found: {snapshot_root}")
+        self.root            = root
+        self.tree_root       = (root / "tree").resolve()
+        if not self.tree_root.exists():
+            raise CodeWorkspaceError(
+                f"Snapshot {snapshot_root} missing tree/ subdir")
+        self.bad_commit_sha  = bad_commit_sha
+        self._git_log: Optional[List[CommitRecord]] = None
+        self._diffs_root     = (root / "diffs").resolve()
+
+    # ------------------------------------------------------------------
+    # Public file-system API (1 method per agent action_type)
+    # ------------------------------------------------------------------
+
+    def list_dir(self, path: str = ".") -> Dict[str, Any]:
+        """List files + subdirs at a relative path under tree/."""
+        target = self._resolve_tree(path)
+        if not target.is_dir():
+            raise CodeWorkspaceError(f"Not a directory: {path}")
+
+        entries = []
+        for child in sorted(target.iterdir()):
+            if child.name.startswith("."):
+                continue
+            entries.append({
+                "name":   child.name,
+                "type":   "dir" if child.is_dir() else "file",
+                "size":   child.stat().st_size if child.is_file() else None,
+            })
+            if len(entries) >= self.MAX_LIST_ENTRIES:
+                break
+        return {
+            "path":    self._rel_to_tree(target),
+            "entries": entries,
+            "count":   len(entries),
+        }
+
+    def read_file(self, path: str) -> Dict[str, Any]:
+        """Read a file under tree/. Truncates if larger than MAX_FILE_BYTES."""
+        target = self._resolve_tree(path)
+        if not target.is_file():
+            raise CodeWorkspaceError(f"Not a file: {path}")
+        data = target.read_bytes()
+        truncated = False
+        if len(data) > self.MAX_FILE_BYTES:
+            data = data[: self.MAX_FILE_BYTES]
+            truncated = True
+        try:
+            text = data.decode("utf-8")
+        except UnicodeDecodeError:
+            text = data.decode("utf-8", errors="replace")
+        return {
+            "path":      self._rel_to_tree(target),
+            "content":   text,
+            "size":      target.stat().st_size,
+            "truncated": truncated,
+        }
+
+    def search_code(
+        self,
+        query:        str,
+        file_pattern: str = "*.py",
+        max_hits:     Optional[int] = None,
+    ) -> Dict[str, Any]:
+        """
+        Substring search across files matching `file_pattern` under tree/.
+        Returns up to `max_hits` (or MAX_SEARCH_HITS) hits with line context.
+        """
+        if not query:
+            return {"query": query, "hits": [], "count": 0}
+        cap = min(max_hits or self.MAX_SEARCH_HITS, self.MAX_SEARCH_HITS)
+        hits: List[Dict[str, Any]] = []
+        bytes_scanned = 0
+
+        for fp in self._iter_tree_files(file_pattern):
+            try:
+                text = fp.read_text("utf-8", errors="replace")
+            except OSError:
+                continue
+            bytes_scanned += len(text)
+            if bytes_scanned > self.MAX_SEARCH_BYTES:
+                break
+            for ln, line in enumerate(text.splitlines(), 1):
+                if query in line:
+                    hits.append({
+                        "path":  self._rel_to_tree(fp),
+                        "line":  ln,
+                        "match": line.strip()[:240],
+                    })
+                    if len(hits) >= cap:
+                        return {"query": query, "hits": hits, "count": len(hits),
+                                "truncated": True}
+        return {"query": query, "hits": hits, "count": len(hits),
+                "truncated": False}
+
+    def get_git_log(
+        self,
+        path:      str = "",
+        n_commits: int = 10,
+    ) -> Dict[str, Any]:
+        """
+        Return up to `n_commits` commits from the snapshot's pre-baked git_log.
+        If `path` is provided, filters to commits that touched a file matching
+        that exact path (or that path's directory).
+        """
+        log = self._load_git_log()
+        if path:
+            target = path.strip("/")
+            log = [c for c in log if any(
+                f == target or f.startswith(target.rstrip("/") + "/") for f in c.files
+            )]
+        log = log[: max(1, n_commits)]
+        return {
+            "path":    path or ".",
+            "commits": [
+                {"sha": c.sha, "author": c.author, "date": c.date,
+                 "message": c.message, "files": list(c.files)}
+                for c in log
+            ],
+            "count":   len(log),
+        }
+
+    def get_file_diff(
+        self,
+        commit_sha: str,
+        path:       str = "",
+    ) -> Dict[str, Any]:
+        """
+        Return the unified diff for `commit_sha`, optionally filtered to
+        hunks touching files matching `path`.
+        """
+        diff_path = (self._diffs_root / f"{commit_sha}.patch")
+        try:
+            diff_path = diff_path.resolve()
+            if not str(diff_path).startswith(str(self._diffs_root)):
+                raise CodeWorkspaceError(f"Illegal diff path: {commit_sha}")
+        except OSError:
+            raise CodeWorkspaceError(f"Diff not found for {commit_sha}")
+
+        if not diff_path.exists():
+            raise CodeWorkspaceError(f"Diff not found for {commit_sha}")
+
+        text = diff_path.read_text("utf-8", errors="replace")
+        if path:
+            text = self._filter_diff_by_path(text, path)
+        return {
+            "commit_sha": commit_sha,
+            "path":       path or "*",
+            "diff":       text,
+        }
+
+    # ------------------------------------------------------------------
+    # Lightweight introspection (used to seed the code agent at handoff)
+    # ------------------------------------------------------------------
+
+    def file_tree(self, max_depth: int = 3) -> List[str]:
+        """Flat list of files under tree/, capped to a sane depth."""
+        out: List[str] = []
+        for fp in self._iter_tree_files("*", max_depth=max_depth):
+            out.append(self._rel_to_tree(fp))
+            if len(out) >= self.MAX_LIST_ENTRIES:
+                break
+        return sorted(out)
+
+    def bad_commit_metadata(self) -> Optional[Dict[str, Any]]:
+        """Return commit metadata for `bad_commit_sha` (without the diff)."""
+        if not self.bad_commit_sha:
+            return None
+        for c in self._load_git_log():
+            if c.sha.startswith(self.bad_commit_sha) or self.bad_commit_sha.startswith(c.sha):
+                return {"sha": c.sha, "author": c.author, "date": c.date,
+                        "message": c.message, "files": list(c.files)}
+        return None
+
+    # ------------------------------------------------------------------
+    # Internals
+    # ------------------------------------------------------------------
+
+    def _load_git_log(self) -> List[CommitRecord]:
+        if self._git_log is not None:
+            return self._git_log
+        path = self.root / "git_log.json"
+        if not path.exists():
+            self._git_log = []
+            return self._git_log
+        raw = json.loads(path.read_text("utf-8"))
+        self._git_log = [
+            CommitRecord(
+                sha     = c["sha"],
+                author  = c.get("author", "unknown"),
+                date    = c.get("date", ""),
+                message = c.get("message", ""),
+                files   = list(c.get("files", [])),
+            )
+            for c in raw
+        ]
+        return self._git_log
+
+    def _resolve_tree(self, path: str) -> Path:
+        """Resolve a user-supplied relative path under tree/, blocking escapes."""
+        cleaned = (path or ".").lstrip("/").lstrip(os.sep)
+        if cleaned in ("", "."):
+            return self.tree_root
+        target = (self.tree_root / cleaned).resolve()
+        if not str(target).startswith(str(self.tree_root)):
+            raise CodeWorkspaceError(f"Illegal path (escapes sandbox): {path}")
+        if not target.exists():
+            raise CodeWorkspaceError(f"Path not found: {path}")
+        return target
+
+    def _rel_to_tree(self, p: Path) -> str:
+        try:
+            return str(p.relative_to(self.tree_root)) or "."
+        except ValueError:
+            return str(p)
+
+    def _iter_tree_files(self, pattern: str, max_depth: int = 16):
+        """Yield Paths under tree/ matching pattern (glob-style)."""
+        for dirpath, dirnames, filenames in os.walk(self.tree_root):
+            depth = Path(dirpath).relative_to(self.tree_root).parts
+            if len(depth) > max_depth:
+                dirnames[:] = []
+                continue
+            dirnames[:] = [d for d in dirnames if not d.startswith(".")]
+            for fname in filenames:
+                if fname.startswith("."):
+                    continue
+                if pattern == "*" or fnmatch.fnmatch(fname, pattern):
+                    yield Path(dirpath) / fname
+
+    @staticmethod
+    def _filter_diff_by_path(diff: str, path: str) -> str:
+        """Return only diff hunks where the +++ b/<file> matches `path`."""
+        out: List[str] = []
+        keep = False
+        target = path.strip("/")
+        for line in diff.split("\n"):
+            if line.startswith("diff --git ") or line.startswith("--- a/") or line.startswith("+++ b/"):
+                if line.startswith("+++ b/"):
+                    file_in_hunk = line[6:].strip()
+                    keep = (file_in_hunk == target
+                            or file_in_hunk.startswith(target.rstrip("/") + "/"))
+                if line.startswith("diff --git "):
+                    # carry through; we'll re-evaluate at +++
+                    pass
+            if keep or line.startswith("diff --git "):
+                out.append(line)
+        return "\n".join(out)

server/incident_environment.py

@@ -1,11 +1,23 @@
 """
 Core Environment implementation.
 
-Execution order per step:  validate → mutate → tick → observe → reward.
-
-The environment uses oracle-shaped rewards for training (they peek at hidden
-state to compute whether the agent investigated the right service) but the
-grader used for evaluation is oracle-independent (trajectory-only).
+Per-step execution order:  validate → mutate → tick → observe → reward.
+
+Two-phase architecture:
+  Phase 1 — ops/SRE diagnostic loop (existing behavior).
+  Phase 2 — code attribution loop, sandboxed under a CodeWorkspace.
+
+Mode selection is automatic per scenario:
+  - Scenario with `code_context = None`  → legacy P1-only episode
+                                            (declare_root_cause terminates)
+  - Scenario with `code_context != None` → unified P1 → P2 episode
+                                            (declare_root_cause is silent;
+                                             transition_to_phase2 switches phase;
+                                             propose_patch / declare_no_change
+                                             terminate the episode)
+
+The environment uses oracle-shaped per-step rewards for training.  The
+oracle-INDEPENDENT graders live on `BaseScenario` and `scenarios.grader_p2`.
 """
 
 from __future__ import annotations
@@ -18,16 +30,32 @@ from typing import Any, Dict, List, Optional, Tuple
 from ..models import (
     ActionType,
     IncidentAction,
-    IncidentObservation,
     IncidentState,
     StepRecord,
+    BeliefState,
     DIAGNOSTIC_ACTIONS,
     REMEDIATION_ACTIONS,
     TARGETED_ACTIONS,
+    PHASE1_ACTIONS,
+    PHASE2_ACTIONS,
+    PHASE2_DIAGNOSTIC_ACTIONS,
+    PHASE2_TERMINAL_ACTIONS,
 )
 from ..simulation.infrastructure import Infrastructure, SERVICE_NAMES
 from ..tasks import get_scenario, TASK_NAMES
 from ..scenarios.base import BaseScenario
+from ..pools import POOLS, get_pool, sample_task, oracle_belief
+from .code_workspace import CodeWorkspace, CodeWorkspaceError
+
+
+# Per-step reward constants ------------------------------------------------
+_STEP_PENALTY     = -0.02
+_REPEAT_PENALTY   = -0.05
+_INVALID_PENALTY  = -0.05
+
+# Phase 2 shaping (small — terminal patch quality is graded post-hoc)
+_P2_DIAG_REWARD   = +0.05
+_P2_TERMINAL_BONUS = +0.10
 
 
 class IncidentEnvironment:
@@ -35,416 +63,866 @@ class IncidentEnvironment:
     SRE Incident Response Environment.
 
     Implements the three OpenEnv methods:
-      - reset(task_name) → IncidentObservation
-      - step(action)     → dict with observation, reward, done
-      - state()          → IncidentState
+      - reset(task_name)  → initial observation + info
+      - step(action)      → dict with observation, reward, done, info
+      - state()           → IncidentState for monitoring
+
+    Plus two extras used by the unified evaluator:
+      - get_trajectory()        → P1 + P2 step records
+      - score_unified(...)      → component scores for unified grader
     """
 
     def __init__(self) -> None:
         self._infra: Optional[Infrastructure] = None
         self._scenario: Optional[BaseScenario] = None
         self._state = IncidentState()
-        self._trajectory: List[StepRecord] = []
+
+        # ---- Per-episode mutable state ----
+        self._phase: int = 1
+        self._workspace: Optional[CodeWorkspace] = None
+        self._belief_at_transition: Optional[BeliefState] = None
+        self._p1_trajectory: List[StepRecord] = []
+        self._p2_trajectory: List[StepRecord] = []
+        self._declared_patch: Optional[str] = None
+        self._declared_no_change: bool = False
+        self._declared_root_cause: Optional[str] = None
         self._cumulative_reward: float = 0.0
         self._done: bool = False
-        self._root_cause_declared: bool = False
 
-    # ------------------------------------------------------------------
+        # ---- Pool / mode (set by reset, drives episode semantics) ----
+        # mode in {"joint" (default), "p1_only" (Pool A), "p2_only" (Pool B)}
+        self._pool: Optional[str] = None
+        self._mode: str = "joint"
+        self._inject_oracle_belief: bool = False
+
+        # P2-only tracking (for repeat detection inside P2)
+        self._p2_actions_taken: List[Tuple[str, str]] = []   # (atype, primary_param)
+
+    # ==================================================================
     # reset()
-    # ------------------------------------------------------------------
+    # ==================================================================
 
     def reset(
         self,
         task_name: Optional[str] = None,
-        seed: Optional[int] = None,
-        **kwargs: Any,
+        seed:      Optional[int] = None,
+        pool:      Optional[str] = None,
+        mode:      Optional[str] = None,
+        **kwargs:  Any,
     ) -> Dict[str, Any]:
         """
         Initialize a new incident episode.
 
-        Args:
-            task_name: One of "memory_leak", "cascading_failure", "distributed_deadlock".
-                       If None, picks randomly.
-            seed: Optional random seed for reproducibility.
-
-        Returns:
-            Dict with observation, reward=0.0, done=False.
+        `pool`  selects training pool A/B/C/D (overrides default mode).
+        `mode`  forces episode semantics ("p1_only"|"p2_only"|"joint").
+                Explicit `mode` always wins over pool defaults.
         """
         if seed is not None:
             random.seed(seed)
 
+        # ---- Pool / task selection ----
+        pool_obj = None
+        if pool:
+            pool_obj = get_pool(pool)
+            if task_name is None:
+                task_name = sample_task(pool, rng=random)
+            self._pool = pool_obj.name
+            self._mode = pool_obj.mode
+            self._inject_oracle_belief = pool_obj.inject_oracle_belief
+        else:
+            self._pool = None
+            self._mode = "joint"
+            self._inject_oracle_belief = False
+
+        if mode:
+            self._mode = mode
+            if mode == "p2_only":
+                self._inject_oracle_belief = True
+
         if task_name is None:
             task_name = random.choice(TASK_NAMES)
 
-        # Create fresh infrastructure
-        self._infra = Infrastructure()
+        self._infra    = Infrastructure()
         self._scenario = get_scenario(task_name)
         self._infra.time_budget_minutes = self._scenario.time_budget_minutes
-
-        # Inject scenario faults
         self._scenario.inject(self._infra)
 
-        # Run a few ticks to let cascades propagate
+        # Let cascades propagate a few minutes
         for _ in range(3):
             self._infra.tick()
 
-        # Reset episode state
         self._state = IncidentState(
-            episode_id=str(uuid.uuid4()),
-            task_name=task_name,
-            step_count=0,
-            time_elapsed_minutes=self._infra.current_minute,
-            done=False,
-            cumulative_reward=0.0,
+            episode_id           = str(uuid.uuid4()),
+            task_name            = task_name,
+            step_count           = 0,
+            time_elapsed_minutes = self._infra.current_minute,
+            done                 = False,
+            cumulative_reward    = 0.0,
         )
-        self._trajectory = []
-        self._cumulative_reward = 0.0
-        self._done = False
-        self._root_cause_declared = False
+        self._phase                 = 1
+        self._workspace             = None
+        self._belief_at_transition  = None
+        self._p1_trajectory         = []
+        self._p2_trajectory         = []
+        self._declared_patch        = None
+        self._declared_no_change    = False
+        self._declared_root_cause   = None
+        self._cumulative_reward     = 0.0
+        self._done                  = False
+        self._p2_actions_taken      = []
+
+        # ---- Pool B (p2_only) auto-handoff with oracle belief --------
+        # The agent never sees Phase 1; we synthesise a perfect handoff and
+        # immediately switch the env into Phase 2.
+        if self._mode == "p2_only" and self._scenario.code_context is not None:
+            belief = oracle_belief(self._scenario)
+            self._handle_transition(IncidentAction(
+                action_type    = ActionType.TRANSITION_TO_PHASE2.value,
+                target_service = None,
+                parameters     = {"belief": asdict(belief)},
+            ))
+            # _handle_transition already returned; we just consume its
+            # observation as the reset observation so caller sees Phase 2.
+            obs = self._build_observation(
+                action_result   = {
+                    "message":          "[Pool B] Auto-handoff with oracle Phase-1 belief.",
+                    "issue":            self._scenario.build_p2_issue(belief),
+                    "file_tree":        (self._workspace.file_tree(max_depth=4)
+                                         if self._workspace else []),
+                    "bad_commit_sha":   self._scenario.code_context.bad_commit_sha,
+                    "bad_commit":       (self._workspace.bad_commit_metadata()
+                                         if self._workspace else None),
+                },
+                action_success  = True,
+                action_message  = "Episode started in Pool B (P2-only) mode",
+                reward          = 0.0,
+            )
+            return {
+                "observation": obs,
+                "reward":      0.01,
+                "done":        False,
+                "info":        {"task_name": task_name,
+                                "pool":      self._pool,
+                                "mode":      self._mode,
+                                "has_phase2": True,
+                                "phase":     2},
+            }
 
         obs = self._build_observation(
-            action_result={"message": "Incident triggered. Begin investigation."},
-            action_success=True,
-            action_message="Episode started",
-            reward=0.0,
+            action_result   = {"message": "Incident triggered. Begin investigation."},
+            action_success  = True,
+            action_message  = "Episode started",
+            reward          = 0.0,
         )
-
         return {
             "observation": obs,
-            "reward": 0.01,
-            "done": False,
+            "reward":      0.01,
+            "done":        False,
+            "info":        {"task_name":  task_name,
+                            "pool":       self._pool,
+                            "mode":       self._mode,
+                            "has_phase2": self._scenario.code_context is not None},
         }
 
-    # ------------------------------------------------------------------
-    # step()  —  validate → mutate → tick → observe → reward
-    # ------------------------------------------------------------------
+    # ==================================================================
+    # step()
+    # ==================================================================
 
     def step(self, action_data: Dict[str, Any]) -> Dict[str, Any]:
-        """
-        Execute one agent action.
-
-        Args:
-            action_data: Dict with action_type, target_service, parameters.
-
-        Returns:
-            Dict with observation, reward, done, info.
-        """
+        """Execute one agent action — phase-aware dispatch."""
         if self._done:
-            obs = self._build_observation(
-                action_result={"error": "Episode is already done."},
-                action_success=False,
-                action_message="Episode already finished",
-                reward=0.0,
-            )
-            final_grade = self._scenario.grade(self._trajectory) if self._scenario else 0.01
-            return {"observation": obs, "reward": 0.01, "done": True, "info": {"score": final_grade}}
+            return self._final_step_response()
 
         if self._infra is None or self._scenario is None:
-            obs = self._build_observation(
-                action_result={"error": "Environment not initialized. Call reset() first."},
-                action_success=False,
-                action_message="Not initialized",
-                reward=0.0,
-            )
-            return {"observation": obs, "reward": 0.01, "done": False, "info": {}}
+            return self._not_initialized_response()
 
-        # Parse action
         action = IncidentAction(
-            action_type=action_data.get("action_type", ""),
-            target_service=action_data.get("target_service"),
-            parameters=action_data.get("parameters", {}),
+            action_type    = action_data.get("action_type", ""),
+            target_service = action_data.get("target_service"),
+            parameters     = action_data.get("parameters", {}) or {},
         )
 
-        # ---- VALIDATE ----
-        is_valid, error_msg = self._infra.validate_action(
-            action.action_type, action.target_service)
+        # ---- Type validation ----------------------------------------
+        try:
+            atype = ActionType(action.action_type)
+        except ValueError:
+            return self._invalid_action_response(
+                f"Unknown action type: {action.action_type!r}",
+                action,
+            )
 
-        if not is_valid:
-            reward = -0.05
-            self._cumulative_reward += reward
-            self._state.step_count += 1
-            obs = self._build_observation(
-                action_result={"error": error_msg},
-                action_success=False,
-                action_message=f"Invalid action: {error_msg}",
-                reward=reward,
+        # ---- Phase-aware dispatch -----------------------------------
+        if atype == ActionType.TRANSITION_TO_PHASE2:
+            return self._handle_transition(action)
+
+        if self._phase == 1:
+            if atype not in PHASE1_ACTIONS:
+                return self._invalid_action_response(
+                    f"Action {atype.value!r} not allowed in Phase 1", action,
+                )
+            return self._step_phase1(action, atype)
+
+        # Phase 2
+        if atype not in PHASE2_ACTIONS:
+            return self._invalid_action_response(
+                f"Action {atype.value!r} not allowed in Phase 2", action,
             )
-            self._record_step(action, reward, obs)
-            return {"observation": obs, "reward": reward, "done": False, "info": {"error": error_msg}}
+        return self._step_phase2(action, atype)
 
-        # ---- MUTATE ----
-        action_result, action_msg = self._execute_action(action)
+    # ------------------------------------------------------------------
+    # Phase 1 step
+    # ------------------------------------------------------------------
 
-        # ---- TICK ----
+    def _step_phase1(
+        self,
+        action: IncidentAction,
+        atype:  ActionType,
+    ) -> Dict[str, Any]:
+        # Validate target / preconditions via Infrastructure
+        is_valid, err = self._infra.validate_action(
+            action.action_type, action.target_service)
+        if not is_valid:
+            return self._invalid_action_response(err, action)
+
+        # Mutate
+        action_result, action_msg = self._execute_p1_action(action, atype)
+
+        # Tick simulation
         self._infra.tick()
-        self._state.step_count += 1
+        self._state.step_count           += 1
         self._state.time_elapsed_minutes = self._infra.current_minute
 
-        # ---- REWARD (oracle-shaped for training) ----
-        # Must compute BEFORE recording — so repeat detection doesn't
-        # flag the current action as already taken.
-        reward = self._compute_reward(action)
+        # Reward (compute BEFORE recording so repeat-detection sees prior actions)
+        reward = self._compute_p1_reward(action, atype)
         self._infra.record_action(action.action_type, action.target_service)
+
         self._cumulative_reward += reward
         self._state.cumulative_reward = self._cumulative_reward
 
-        # ---- CHECK DONE ----
-        done = self._check_done(action)
+        # Done check
+        done = self._check_done_p1(atype)
         self._done = done
         self._state.done = done
 
-        # ---- OBSERVE ----
         obs = self._build_observation(
-            action_result=action_result,
-            action_success=True,
-            action_message=action_msg,
-            reward=reward,
+            action_result   = action_result,
+            action_success  = True,
+            action_message  = action_msg,
+            reward          = reward,
         )
 
-        self._record_step(action, reward, obs)
+        record = StepRecord(
+            step_number             = self._state.step_count,
+            action                  = action,
+            reward                  = reward,
+            observation_summary     = {
+                "action_message":      obs.get("action_message", ""),
+                "active_alerts_count": obs.get("active_alerts_count", 0),
+            },
+            service_statuses_after  = dict(obs.get("service_statuses", {})),
+            timestamp_minutes       = self._infra.current_minute,
+            phase                   = 1,
+        )
+        self._p1_trajectory.append(record)
 
         info: Dict[str, Any] = {}
         if done:
-            # Compute final grade (oracle-independent)
-            final_grade = self._scenario.grade(self._trajectory)
-            info["score"] = final_grade
-            info["task_name"] = self._scenario.task_name
-            info["steps_taken"] = self._state.step_count
-            info["trajectory_length"] = len(self._trajectory)
+            info["score"]            = self._scenario.grade(self._p1_trajectory)
+            info["task_name"]        = self._scenario.task_name
+            info["steps_taken"]      = self._state.step_count
+            info["trajectory_length"] = len(self._p1_trajectory)
 
         return {"observation": obs, "reward": reward, "done": done, "info": info}
 
     # ------------------------------------------------------------------
-    # state()
+    # Phase 2 step
     # ------------------------------------------------------------------
 
+    def _step_phase2(
+        self,
+        action: IncidentAction,
+        atype:  ActionType,
+    ) -> Dict[str, Any]:
+        if self._workspace is None:
+            return self._invalid_action_response(
+                "Phase 2 not initialised — must transition_to_phase2 first.",
+                action,
+            )
+
+        params = action.parameters or {}
+
+        # ---- Execute action ----
+        try:
+            if atype == ActionType.LIST_DIR:
+                result = self._workspace.list_dir(params.get("path", "."))
+                msg = f"Listed {result.get('count', 0)} entries in {result.get('path', '.')}"
+            elif atype == ActionType.READ_FILE:
+                result = self._workspace.read_file(params.get("path", ""))
+                msg = f"Read {result.get('path')} ({result.get('size', 0)} bytes)"
+            elif atype == ActionType.SEARCH_CODE:
+                result = self._workspace.search_code(
+                    query        = params.get("query", ""),
+                    file_pattern = params.get("file_pattern", "*.py"),
+                    max_hits     = params.get("max_hits"),
+                )
+                msg = f"Found {result.get('count', 0)} hit(s) for {params.get('query')!r}"
+            elif atype == ActionType.GET_GIT_LOG:
+                result = self._workspace.get_git_log(
+                    path      = params.get("path", ""),
+                    n_commits = int(params.get("n_commits", 10)),
+                )
+                msg = f"Returned {result.get('count', 0)} commit(s)"
+            elif atype == ActionType.GET_FILE_DIFF:
+                result = self._workspace.get_file_diff(
+                    commit_sha = params.get("commit_sha", ""),
+                    path       = params.get("path", ""),
+                )
+                msg = f"Diff for {result.get('commit_sha')[:8]} ({len(result.get('diff', ''))} bytes)"
+            elif atype == ActionType.PROPOSE_PATCH:
+                diff = params.get("diff", "")
+                self._declared_patch = diff
+                result = {"accepted": True, "patch_bytes": len(diff)}
+                msg = "Patch proposal accepted — episode terminating."
+            elif atype == ActionType.DECLARE_NO_CHANGE:
+                self._declared_no_change = True
+                reason = params.get("reason", "")
+                result = {"accepted": True, "reason": reason}
+                msg = "no-change declaration accepted — episode terminating."
+            else:
+                return self._invalid_action_response(
+                    f"Unhandled P2 action type: {atype.value!r}", action,
+                )
+            success = True
+        except CodeWorkspaceError as e:
+            result  = {"error": str(e)}
+            msg     = f"Workspace error: {e}"
+            success = False
+
+        # ---- Tick (simulation time still advances during P2) ----
+        self._infra.tick()
+        self._state.step_count           += 1
+        self._state.time_elapsed_minutes = self._infra.current_minute
+
+        # ---- Reward ----
+        reward = self._compute_p2_reward(action, atype, success)
+        self._cumulative_reward += reward
+        self._state.cumulative_reward = self._cumulative_reward
+
+        # ---- Done ----
+        done = (atype in PHASE2_TERMINAL_ACTIONS) or self._exceeded_step_budget()
+        self._done = done
+        self._state.done = done
+
+        obs = self._build_observation(
+            action_result   = result,
+            action_success  = success,
+            action_message  = msg,
+            reward          = reward,
+        )
+
+        # Record
+        record = StepRecord(
+            step_number             = self._state.step_count,
+            action                  = action,
+            reward                  = reward,
+            observation_summary     = {
+                "action_message": obs.get("action_message", ""),
+                "p2_action":      atype.value,
+            },
+            service_statuses_after  = dict(obs.get("service_statuses", {})),
+            timestamp_minutes       = self._infra.current_minute,
+            phase                   = 2,
+        )
+        self._p2_trajectory.append(record)
+
+        # Track repeats inside P2
+        prim_param = self._p2_primary_param(atype, params)
+        self._p2_actions_taken.append((atype.value, prim_param))
+
+        info: Dict[str, Any] = {}
+        if done:
+            info["score"]            = self._compute_unified_final_score()
+            info["task_name"]        = self._scenario.task_name
+            info["steps_taken"]      = self._state.step_count
+            info["trajectory_length"] = len(self._p1_trajectory) + len(self._p2_trajectory)
+
+        return {"observation": obs, "reward": reward, "done": done, "info": info}
+
+    # ------------------------------------------------------------------
+    # transition_to_phase2 handler
+    # ------------------------------------------------------------------
+
+    def _handle_transition(self, action: IncidentAction) -> Dict[str, Any]:
+        if self._phase != 1:
+            return self._invalid_action_response(
+                "Already in Phase 2 — cannot transition again.", action,
+            )
+        if self._scenario is None or self._scenario.code_context is None:
+            return self._invalid_action_response(
+                "Scenario has no code_context — Phase 2 unavailable.", action,
+            )
+
+        ctx = self._scenario.code_context
+
+        # Construct workspace
+        try:
+            self._workspace = CodeWorkspace(
+                snapshot_root  = ctx.repo_snapshot_path,
+                bad_commit_sha = ctx.bad_commit_sha,
+            )
+        except CodeWorkspaceError as e:
+            return self._invalid_action_response(
+                f"Cannot open snapshot: {e}", action,
+            )
+
+        # Capture handoff belief
+        belief_dict = (action.parameters or {}).get("belief") or {}
+        self._belief_at_transition = self._coerce_belief(belief_dict)
+
+        # Switch phase
+        self._phase                       = 2
+        self._state.step_count           += 1
+        self._infra.tick()
+        self._state.time_elapsed_minutes  = self._infra.current_minute
+
+        # Initial P2 obs
+        issue_text = self._scenario.build_p2_issue(self._belief_at_transition)
+        file_tree  = self._workspace.file_tree(max_depth=4)
+        commit_meta = self._workspace.bad_commit_metadata()
+
+        action_result = {
+            "phase":            2,
+            "issue":            issue_text,
+            "file_tree":        file_tree,
+            "bad_commit_sha":   ctx.bad_commit_sha,
+            "bad_commit":       commit_meta,
+            "snapshot_root":    str(self._workspace.tree_root),
+        }
+
+        # Reward: small handoff bonus only when belief is non-trivial
+        reward = 0.0
+        if self._belief_at_transition.suspected_service:
+            reward += 0.05
+        self._cumulative_reward += reward
+        self._state.cumulative_reward = self._cumulative_reward
+
+        obs = self._build_observation(
+            action_result   = action_result,
+            action_success  = True,
+            action_message  = "Transitioned to Phase 2 (code attribution).",
+            reward          = reward,
+        )
+
+        record = StepRecord(
+            step_number             = self._state.step_count,
+            action                  = action,
+            reward                  = reward,
+            observation_summary     = {
+                "action_message": "transition_to_phase2",
+                "transition":     True,
+            },
+            service_statuses_after  = dict(obs.get("service_statuses", {})),
+            timestamp_minutes       = self._infra.current_minute,
+            phase                   = 2,
+            belief_state_snapshot   = asdict(self._belief_at_transition),
+        )
+        self._p2_trajectory.append(record)
+
+        return {"observation": obs, "reward": reward, "done": False, "info": {}}
+
+    @staticmethod
+    def _coerce_belief(d: Dict[str, Any]) -> BeliefState:
+        """Best-effort: turn an inference-side dict into the canonical BeliefState."""
+        gaps = d.get("evidence_gaps", [])
+        if isinstance(gaps, str):
+            gaps = [g.strip() for g in gaps.split(",") if g.strip() and g.strip() != "none"]
+        return BeliefState(
+            suspected_service     = d.get("suspected_service") or None,
+            suspected_fault_class = d.get("suspected_fault_class") or None,
+            service_confidence    = float(d.get("service_confidence") or 0.0),
+            fault_confidence      = float(d.get("fault_confidence") or 0.0),
+            evidence_gaps         = list(gaps),
+            estimated_p2_cost     = d.get("estimated_p2_cost") or "unknown",
+            decision              = d.get("decision") or "transition",
+            reasoning             = d.get("reasoning") or "",
+        )
+
+    # ==================================================================
+    # state()
+    # ==================================================================
+
     @property
     def state(self) -> IncidentState:
         return self._state
 
     def get_state(self) -> Dict[str, Any]:
         return {
-            "episode_id": self._state.episode_id,
-            "task_name": self._state.task_name,
-            "step_count": self._state.step_count,
+            "episode_id":           self._state.episode_id,
+            "task_name":            self._state.task_name,
+            "step_count":           self._state.step_count,
             "time_elapsed_minutes": self._state.time_elapsed_minutes,
-            "done": self._state.done,
-            "cumulative_reward": round(self._state.cumulative_reward, 3),
-            "declared_root_cause": self._state.declared_root_cause,
+            "done":                 self._state.done,
+            "cumulative_reward":    round(self._state.cumulative_reward, 3),
+            "declared_root_cause":  self._declared_root_cause,
+            "declared_patch":       self._declared_patch,
+            "declared_no_change":   self._declared_no_change,
+            "phase":                self._phase,
+            "phase_transition_at":  next(
+                (r.step_number for r in self._p2_trajectory
+                 if r.action.action_type == ActionType.TRANSITION_TO_PHASE2.value),
+                None,
+            ),
         }
 
-    # ------------------------------------------------------------------
-    # Action execution handlers
-    # ------------------------------------------------------------------
+    # ==================================================================
+    # Phase 1 action execution
+    # ==================================================================
 
-    def _execute_action(
-        self, action: IncidentAction
+    def _execute_p1_action(
+        self,
+        action: IncidentAction,
+        atype:  ActionType,
     ) -> Tuple[Dict[str, Any], str]:
-        """Execute a validated action. Returns (result_dict, message)."""
-        at = action.parsed_type()
         target = action.target_service
+        params = action.parameters or {}
 
-        if at == ActionType.VIEW_ALERTS:
+        if atype == ActionType.VIEW_ALERTS:
             alerts = self._infra.get_alerts()
-            return {"alerts": alerts, "count": len(alerts)}, f"Viewing {len(alerts)} active alerts"
-
-        elif at == ActionType.QUERY_LOGS:
-            level_filter = action.parameters.get("level")
-            keyword = action.parameters.get("keyword")
-            limit = action.parameters.get("limit", 15)
-            logs = self._infra.get_logs_for_service(target, level_filter, keyword, limit)
+            return {"alerts": alerts, "count": len(alerts)}, \
+                   f"Viewing {len(alerts)} active alerts"
+
+        if atype == ActionType.QUERY_LOGS:
+            level   = params.get("level")
+            keyword = params.get("keyword")
+            limit   = params.get("limit", 15)
+            logs = self._infra.get_logs_for_service(target, level, keyword, limit)
             return {"logs": logs, "count": len(logs), "service": target}, \
                    f"Queried {len(logs)} logs from {target}"
 
-        elif at == ActionType.CHECK_METRICS:
+        if atype == ActionType.CHECK_METRICS:
             metrics = self._infra.get_metrics_for_service(target)
-            return {"metrics": metrics, "service": target, "data_points": len(metrics)}, \
-                   f"Retrieved {len(metrics)} metric data points for {target}"
+            return {"metrics": metrics, "service": target,
+                    "data_points": len(metrics)}, \
+                   f"Retrieved {len(metrics)} metric points for {target}"
 
-        elif at == ActionType.CHECK_DEPENDENCIES:
+        if atype == ActionType.CHECK_DEPENDENCIES:
             deps = self._infra.get_dependencies_for_service(target)
             return {"dependencies": deps, "service": target}, \
                    f"Retrieved dependency map for {target}"
 
-        elif at == ActionType.CHECK_DEPLOY_HISTORY:
+        if atype == ActionType.CHECK_DEPLOY_HISTORY:
             deploys = self._infra.get_deploy_history_for_service(target)
-            return {"deploys": deploys, "service": target, "count": len(deploys)}, \
+            return {"deploys": deploys, "service": target,
+                    "count": len(deploys)}, \
                    f"Retrieved {len(deploys)} deploys for {target}"
 
-        elif at == ActionType.RUN_HEALTH_CHECK:
-            health = self._infra.run_health_check(target)
-            return {"health_check": health, "service": target}, \
-                   f"Health check for {target}: {health['status']}"
+        if atype == ActionType.RUN_HEALTH_CHECK:
+            h = self._infra.run_health_check(target)
+            return {"health_check": h, "service": target}, \
+                   f"Health check for {target}: {h['status']}"
 
-        elif at == ActionType.RESTART_SERVICE:
+        if atype == ActionType.RESTART_SERVICE:
             svc = self._infra.get_service(target)
             msg = svc.restart(self._infra.current_minute) if svc else "Service not found"
             return {"result": msg, "service": target}, msg
 
-        elif at == ActionType.ROLLBACK_DEPLOY:
+        if atype == ActionType.ROLLBACK_DEPLOY:
             svc = self._infra.get_service(target)
-            msg = svc.rollback_deploy(self._infra.current_minute) if svc else "Service not found"
+            msg = svc.rollback_deploy(self._infra.current_minute) \
+                  if svc else "Service not found"
             return {"result": msg, "service": target}, msg
 
-        elif at == ActionType.SCALE_SERVICE:
+        if atype == ActionType.SCALE_SERVICE:
             svc = self._infra.get_service(target)
-            new_replicas = action.parameters.get("replicas", 5)
-            msg = svc.scale(new_replicas, self._infra.current_minute) if svc else "Service not found"
+            new_replicas = params.get("replicas", 5)
+            msg = svc.scale(new_replicas, self._infra.current_minute) \
+                  if svc else "Service not found"
             return {"result": msg, "service": target}, msg
 
-        elif at == ActionType.DECLARE_ROOT_CAUSE:
-            root_cause = action.parameters.get("root_cause", "")
-            self._state.declared_root_cause = root_cause
-            self._root_cause_declared = True
+        if atype == ActionType.DECLARE_ROOT_CAUSE:
+            rc = params.get("root_cause", "")
+            self._declared_root_cause = rc
+            self._state.declared_root_cause = rc
             return {
-                "declared": root_cause,
-                "message": "Root cause declaration registered. Episode will end after this step.",
-            }, f"Root cause declared: {root_cause}"
+                "declared": rc,
+                "message":  ("Root cause declared. " +
+                             ("Episode continues — Phase 2 awaits."
+                              if self._scenario.code_context
+                              else "Episode will end after this step.")),
+            }, f"Root cause declared: {rc[:120]}"
 
-        else:
-            return {"error": f"Unhandled action type: {at}"}, "Unknown action"
+        return {"error": f"Unhandled action type: {atype.value}"}, "Unknown action"
 
-    # ------------------------------------------------------------------
-    # Reward computation (oracle-shaped — Layer 6)
-    # ------------------------------------------------------------------
+    # ==================================================================
+    # Reward computation
+    # ==================================================================
 
-    def _compute_reward(self, action: IncidentAction) -> float:
-        """
-        Compute per-step reward using oracle-shaped signal.
-        The training reward has access to hidden state (involved_services,
-        root_cause_service) — this is necessary for learning.
-        The GRADER does NOT use this; it scores trajectory-only.
-        """
-        at = action.parsed_type()
-        target = action.target_service
+    def _compute_p1_reward(
+        self,
+        action: IncidentAction,
+        atype:  ActionType,
+    ) -> float:
         scenario = self._scenario
-        reward = 0.0
-
-        # --- Step penalty (efficiency pressure) ---
-        reward -= 0.02
+        target   = action.target_service
+        reward   = _STEP_PENALTY
 
-        # --- Repeat detection ---
         if self._infra.was_action_taken(action.action_type, target):
-            reward -= 0.05
-            return round(reward, 3)
+            return round(reward + _REPEAT_PENALTY, 3)
 
-        # --- Diagnostic actions ---
-        if at in DIAGNOSTIC_ACTIONS:
+        if atype in DIAGNOSTIC_ACTIONS:
             if target and target in scenario.involved_services:
-                reward += 0.15  # Investigating a relevant service
+                reward += 0.15
             elif target and target not in scenario.involved_services:
-                reward += 0.05  # Exploring — not penalized heavily
-            elif at == ActionType.VIEW_ALERTS:
-                reward += 0.15  # Always good to view alerts
-
-        # --- Remediation actions ---
-        elif at in REMEDIATION_ACTIONS:
+                reward += 0.05
+            elif atype == ActionType.VIEW_ALERTS:
+                reward += 0.15
+        elif atype in REMEDIATION_ACTIONS:
             if target == scenario.root_cause_service:
-                reward += 0.30  # Correct remediation target
+                reward += 0.30
             elif target and target in scenario.involved_services:
-                reward += 0.10  # Helpful but not the root cause
+                reward += 0.10
             else:
-                reward -= 0.15  # Remediating healthy/uninvolved service
-
-        # --- Root cause declaration ---
-        elif at == ActionType.DECLARE_ROOT_CAUSE:
-            declared = action.parameters.get("root_cause", "").lower()
-            keywords = scenario.root_cause_keywords
-            if keywords:
-                matched = sum(1 for kw in keywords if kw in declared)
-                ratio = matched / len(keywords)
+                reward -= 0.15
+        elif atype == ActionType.DECLARE_ROOT_CAUSE:
+            declared = (action.parameters or {}).get("root_cause", "").lower()
+            kws = scenario.root_cause_keywords
+            if kws:
+                ratio = sum(1 for k in kws if k in declared) / len(kws)
                 if ratio >= 0.6:
-                    reward += 0.40  # Correct
+                    reward += 0.40
                 elif ratio >= 0.3:
-                    reward += 0.15  # Partial
+                    reward += 0.15
                 else:
-                    reward -= 0.20  # Wrong
+                    reward -= 0.20
             else:
                 reward -= 0.20
 
-        # --- Episode completion bonus/penalty ---
-        if self._root_cause_declared:
+        # Completion bonus when episode terminates
+        if self._declared_root_cause and not scenario.code_context:
             if self._infra.all_services_healthy():
-                reward += 0.20  # All services restored
+                reward += 0.20
             if self._infra.current_minute > self._infra.time_budget_minutes:
-                reward -= 0.10  # Exceeded time budget
+                reward -= 0.10
 
         return round(reward, 3)
 
-    # ------------------------------------------------------------------
-    # Done check
-    # ------------------------------------------------------------------
+    def _compute_p2_reward(
+        self,
+        action:  IncidentAction,
+        atype:   ActionType,
+        success: bool,
+    ) -> float:
+        params = action.parameters or {}
+        prim   = self._p2_primary_param(atype, params)
+        reward = _STEP_PENALTY
 
-    def _check_done(self, action: IncidentAction) -> bool:
-        """Episode ends when root cause is declared or max steps reached."""
-        if self._root_cause_declared:
-            return True
-        if self._state.step_count >= self._scenario.max_steps:
+        if not success:
+            return round(reward + _INVALID_PENALTY, 3)
+
+        if (atype.value, prim) in self._p2_actions_taken:
+            return round(reward + _REPEAT_PENALTY, 3)
+
+        if atype in PHASE2_DIAGNOSTIC_ACTIONS:
+            reward += _P2_DIAG_REWARD
+        elif atype in PHASE2_TERMINAL_ACTIONS:
+            reward += _P2_TERMINAL_BONUS
+
+        return round(reward, 3)
+
+    @staticmethod
+    def _p2_primary_param(atype: ActionType, params: Dict[str, Any]) -> str:
+        if atype == ActionType.LIST_DIR:
+            return params.get("path", ".")
+        if atype == ActionType.READ_FILE:
+            return params.get("path", "")
+        if atype == ActionType.SEARCH_CODE:
+            return params.get("query", "")
+        if atype == ActionType.GET_GIT_LOG:
+            return params.get("path", "")
+        if atype == ActionType.GET_FILE_DIFF:
+            return f'{params.get("commit_sha", "")}:{params.get("path", "")}'
+        return ""
+
+    # ==================================================================
+    # Done logic
+    # ==================================================================
+
+    def _check_done_p1(self, atype: ActionType) -> bool:
+        # Pool A / explicit p1_only mode: declare_root_cause always terminates,
+        # regardless of whether the scenario could otherwise transition to P2.
+        if atype == ActionType.DECLARE_ROOT_CAUSE:
+            if self._mode == "p1_only" or self._scenario.code_context is None:
+                return True
+        if self._exceeded_step_budget():
             return True
         return False
 
-    # ------------------------------------------------------------------
+    def _exceeded_step_budget(self) -> bool:
+        budget = self._scenario.max_steps if self._scenario else 20
+        # When code_context exists, allow a bit more headroom for P2 exploration
+        if self._scenario and self._scenario.code_context is not None:
+            budget = budget + 15
+        return self._state.step_count >= budget
+
+    # ==================================================================
     # Observation builder
-    # ------------------------------------------------------------------
+    # ==================================================================
 
     def _build_observation(
         self,
-        action_result: Dict[str, Any],
+        action_result:  Dict[str, Any],
         action_success: bool,
         action_message: str,
-        reward: float,
+        reward:         float,
     ) -> Dict[str, Any]:
-        """Build the POMDP observation dict (no hidden state exposed)."""
-        statuses = self._infra.get_all_statuses() if self._infra else {}
-        alerts = self._infra.get_alerts() if self._infra else []
-        valid_actions = self._infra.get_valid_actions() if self._infra else []
+        statuses      = self._infra.get_all_statuses() if self._infra else {}
+        alerts        = self._infra.get_alerts() if self._infra else []
+        valid_actions = self._valid_actions_for_phase()
 
         return {
-            "incident_summary": self._scenario.incident_summary if self._scenario else "",
-            "severity": self._scenario.severity if self._scenario else "SEV3",
+            "incident_summary":     self._scenario.incident_summary if self._scenario else "",
+            "severity":             self._scenario.severity if self._scenario else "SEV3",
             "time_elapsed_minutes": self._infra.current_minute if self._infra else 0,
-            "time_budget_minutes": self._infra.time_budget_minutes if self._infra else 30,
-            "action_result": action_result,
-            "action_success": action_success,
-            "action_message": action_message,
-            "service_statuses": statuses,
-            "active_alerts_count": len(alerts),
-            "valid_actions": valid_actions,
-            "available_services": list(SERVICE_NAMES),
-            "current_reward": reward,
-            "cumulative_reward": round(self._cumulative_reward, 3),
-            "steps_taken": self._state.step_count,
-            "max_steps": self._scenario.max_steps if self._scenario else 20,
-            "done": self._done,
+            "time_budget_minutes":  self._infra.time_budget_minutes if self._infra else 30,
+            "action_result":        action_result,
+            "action_success":       action_success,
+            "action_message":       action_message,
+            "service_statuses":     statuses,
+            "active_alerts_count":  len(alerts),
+            "valid_actions":        valid_actions,
+            "available_services":   list(SERVICE_NAMES),
+            "current_phase":        self._phase,
+            "current_reward":       reward,
+            "cumulative_reward":    round(self._cumulative_reward, 3),
+            "steps_taken":          self._state.step_count,
+            "max_steps":            self._scenario.max_steps if self._scenario else 20,
+            "done":                 self._done,
+            # Convenience field surfaced after transition (so the inference loop
+            # can grab it without re-issuing a step) — only meaningful after
+            # transition_to_phase2 has been called.
+            "bad_commit_sha":       (self._scenario.code_context.bad_commit_sha
+                                     if self._scenario and self._scenario.code_context else None),
         }
 
-    # ------------------------------------------------------------------
-    # Trajectory recording
-    # ------------------------------------------------------------------
+    def _valid_actions_for_phase(self) -> List[str]:
+        if self._phase == 1:
+            base = self._infra.get_valid_actions() if self._infra else []
+            # Filter to only P1 + (optionally) transition_to_phase2
+            valid = [a for a in base
+                     if a.split(":", 1)[0] in {at.value for at in PHASE1_ACTIONS}]
+            if self._scenario and self._scenario.code_context is not None:
+                valid.append(ActionType.TRANSITION_TO_PHASE2.value)
+            return valid
+        # Phase 2
+        return [at.value for at in PHASE2_ACTIONS]
+
+    # ==================================================================
+    # Trajectory access (used by /score endpoint and Pool runners)
+    # ==================================================================
+
+    def get_trajectory(self) -> List[StepRecord]:
+        return list(self._p1_trajectory) + list(self._p2_trajectory)
+
+    def get_p1_trajectory(self) -> List[StepRecord]:
+        return list(self._p1_trajectory)
+
+    def get_p2_trajectory(self) -> List[StepRecord]:
+        return list(self._p2_trajectory)
+
+    def get_belief_at_transition(self) -> Optional[BeliefState]:
+        return self._belief_at_transition
+
+    # ==================================================================
+    # Final unified scoring
+    # ==================================================================
+
+    def _compute_unified_final_score(self) -> float:
+        """Quick wrapper for the in-step `info.score` field."""
+        from ..tasks import grade_trajectory_unified
+        if self._scenario is None:
+            return 0.01
+        breakdown = grade_trajectory_unified(
+            task_name           = self._scenario.task_name,
+            p1_trajectory       = self._p1_trajectory,
+            p2_trajectory       = self._p2_trajectory,
+            declared_patch      = self._declared_patch,
+            declared_no_change  = self._declared_no_change,
+            p1_belief_history   = [],
+        )
+        return float(breakdown.get("final", 0.01))
+
+    def score_unified(
+        self,
+        belief_history: Optional[List[Dict[str, Any]]] = None,
+    ) -> Dict[str, float]:
+        """Public wrapper exposed by the /score endpoint."""
+        from ..tasks import grade_trajectory_unified
+        if self._scenario is None:
+            return {"final": 0.01}
+        return grade_trajectory_unified(
+            task_name           = self._scenario.task_name,
+            p1_trajectory       = self._p1_trajectory,
+            p2_trajectory       = self._p2_trajectory,
+            declared_patch      = self._declared_patch,
+            declared_no_change  = self._declared_no_change,
+            p1_belief_history   = belief_history or [],
+        )
 
-    def _record_step(
+    # ==================================================================
+    # Error / fallback responses
+    # ==================================================================
+
+    def _invalid_action_response(
         self,
+        msg:    str,
         action: IncidentAction,
-        reward: float,
-        observation: Dict[str, Any],
-    ) -> None:
-        """Record step for trajectory-based grading."""
+    ) -> Dict[str, Any]:
+        reward = _INVALID_PENALTY
+        self._cumulative_reward += reward
+        self._state.step_count  += 1
+        obs = self._build_observation(
+            action_result   = {"error": msg},
+            action_success  = False,
+            action_message  = f"Invalid action: {msg}",
+            reward          = reward,
+        )
+        # Still record the failed attempt so trajectory analysis sees it
         record = StepRecord(
-            step_number=self._state.step_count,
-            action=action,
-            reward=reward,
-            observation_summary={
-                "action_message": observation.get("action_message", ""),
-                "active_alerts_count": observation.get("active_alerts_count", 0),
-            },
-            service_statuses_after=dict(observation.get("service_statuses", {})),
-            timestamp_minutes=self._infra.current_minute if self._infra else 0,
+            step_number             = self._state.step_count,
+            action                  = action,
+            reward                  = reward,
+            observation_summary     = {"action_message": f"invalid: {msg}"},
+            service_statuses_after  = dict(obs.get("service_statuses", {})),
+            timestamp_minutes       = self._infra.current_minute if self._infra else 0,
+            phase                   = self._phase,
         )
-        self._trajectory.append(record)
+        if self._phase == 1:
+            self._p1_trajectory.append(record)
+        else:
+            self._p2_trajectory.append(record)
 
-    # ------------------------------------------------------------------
-    # Trajectory access (for external grading)
-    # ------------------------------------------------------------------
+        return {"observation": obs, "reward": reward, "done": False,
+                "info": {"error": msg}}
 
-    def get_trajectory(self) -> List[StepRecord]:
-        return list(self._trajectory)
+    def _final_step_response(self) -> Dict[str, Any]:
+        obs = self._build_observation(
+            action_result   = {"error": "Episode is already done."},
+            action_success  = False,
+            action_message  = "Episode already finished",
+            reward          = 0.0,
+        )
+        score = (self._compute_unified_final_score()
+                 if self._scenario and self._scenario.code_context
+                 else (self._scenario.grade(self._p1_trajectory)
+                       if self._scenario else 0.01))
+        return {"observation": obs, "reward": 0.01, "done": True,
+                "info": {"score": score}}
+
+    def _not_initialized_response(self) -> Dict[str, Any]:
+        obs = self._build_observation(
+            action_result   = {"error": "Environment not initialized. Call reset() first."},
+            action_success  = False,
+            action_message  = "Not initialized",
+            reward          = 0.0,
+        )
+        return {"observation": obs, "reward": 0.01, "done": False, "info": {}}

snapshots/auth_v180/diffs/b8e2d44.patch

@@ -0,0 +1,22 @@
+diff --git a/auth/config.py b/auth/config.py
+--- a/auth/config.py
++++ b/auth/config.py
+@@ -8,6 +8,9 @@
+ import os
+ from dataclasses import dataclass
+ 
++
++_DEFAULT_DEV_SECRET = "dev-only-do-not-use-in-prod-zZ3kf81P"
++
+ 
+ @dataclass(frozen=True)
+ class AuthConfig:
+@@ -16,7 +19,7 @@
+     issuer: str
+ 
+ 
+-JWT_SECRET = os.environ.get("JWT_SECRET")
++JWT_SECRET = os.environ.get("JWT_SECRET") or _DEFAULT_DEV_SECRET
+ 
+ AUTH_CONFIG = AuthConfig(
+     jwt_secret        = JWT_SECRET,

snapshots/auth_v180/git_log.json

@@ -0,0 +1,32 @@
+[
+  {
+    "sha": "b8e2d44",
+    "author": "bob",
+    "date": "2025-01-15T13:55:00Z",
+    "message": "Config update: rotate JWT signing secret + dev fallback (v1.8.0)\n\nAdds a baked-in dev secret as a fallback when JWT_SECRET is unset. The\nintent is to make local dev work without a .env file. Note: the env-check\nthat would have gated this in prod was lost in rebase — see CR-2014.",
+    "files": [
+      "auth/config.py"
+    ]
+  },
+  {
+    "sha": "55cb7a9",
+    "author": "alice",
+    "date": "2025-01-13T11:00:00Z",
+    "message": "Add TokenService.issue() + .validate() roundtrip tests",
+    "files": [
+      "auth/token.py"
+    ]
+  },
+  {
+    "sha": "0a4ee10",
+    "author": "alice",
+    "date": "2025-01-08T09:00:00Z",
+    "message": "Initial auth service skeleton",
+    "files": [
+      "auth/__init__.py",
+      "auth/config.py",
+      "auth/token.py",
+      "auth/server.py"
+    ]
+  }
+]

snapshots/auth_v180/tree/auth/__init__.py

@@ -0,0 +1,7 @@
+"""auth service — token signing, JWT validation, role lookup."""
+
+from .config import AUTH_CONFIG
+from .token import TokenService
+from .server import AuthServer
+
+__all__ = ["AUTH_CONFIG", "TokenService", "AuthServer"]

snapshots/auth_v180/tree/auth/config.py

@@ -0,0 +1,35 @@
+"""
+Auth service configuration.
+
+`JWT_SECRET` is loaded from the environment.  v1.8.0 added a fallback to
+a baked-in dev secret to "make local testing easier" — this is the change
+shipped on commit b8e2d44.
+
+The fallback was *intended* to apply only when `ENV=dev`, but the env
+check was forgotten — so production now happily falls back too if the
+real secret is missing or unreadable.
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+
+
+_DEFAULT_DEV_SECRET = "dev-only-do-not-use-in-prod-zZ3kf81P"
+
+
+@dataclass(frozen=True)
+class AuthConfig:
+    jwt_secret: str
+    token_ttl_seconds: int
+    issuer: str
+
+
+JWT_SECRET = os.environ.get("JWT_SECRET") or _DEFAULT_DEV_SECRET
+
+AUTH_CONFIG = AuthConfig(
+    jwt_secret        = JWT_SECRET,
+    token_ttl_seconds = 3600,
+    issuer            = "auth.example.com",
+)

snapshots/auth_v180/tree/auth/server.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+from .token import TokenService
+
+
+class AuthServer:
+    """HTTP entry-point for the auth service."""
+
+    def __init__(self) -> None:
+        self._tokens = TokenService()
+
+    def login(self, user: str, password: str) -> str:
+        if not self._authenticate(user, password):
+            raise PermissionError("invalid credentials")
+        return self._tokens.issue(user, ["read", "write"])
+
+    def validate(self, token: str) -> bool:
+        return self._tokens.validate(token)
+
+    def _authenticate(self, user: str, password: str) -> bool:
+        return bool(user) and bool(password)

snapshots/auth_v180/tree/auth/token.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+import time
+from typing import Optional
+
+from .config import AUTH_CONFIG
+
+
+class TokenService:
+    """Issues + validates short-lived JWT-style tokens."""
+
+    def __init__(self, secret: Optional[str] = None) -> None:
+        self._secret = secret or AUTH_CONFIG.jwt_secret
+
+    def issue(self, subject: str, scopes: list[str]) -> str:
+        payload = f"{subject}|{','.join(scopes)}|{int(time.time())}"
+        return f"{payload}.{self._sign(payload)}"
+
+    def validate(self, token: str) -> bool:
+        try:
+            payload, sig = token.rsplit(".", 1)
+        except ValueError:
+            return False
+        return sig == self._sign(payload)
+
+    def _sign(self, payload: str) -> str:
+        return f"sig-{abs(hash((payload, self._secret))) & 0xFFFFFFFF:08x}"

snapshots/orders_retry_storm/diffs/d09a4f1.patch

@@ -0,0 +1,8 @@
+diff --git a/orders/auth_client.py b/orders/auth_client.py
+--- a/orders/auth_client.py
++++ b/orders/auth_client.py
+@@ -15,5 +15,6 @@ class AuthClient:
+     def validate(self, token):
+-        return self._call_with_retries(token, retries=25)
++        return self._call_with_retries(token, retries=2,
++                                        backoff_seconds=0.5)

snapshots/orders_retry_storm/diffs/f8c9b13.patch

@@ -0,0 +1,8 @@
+diff --git a/orders/auth_client.py b/orders/auth_client.py
+--- a/orders/auth_client.py
++++ b/orders/auth_client.py
+@@ -15,5 +15,6 @@ class AuthClient:
+     def validate(self, token):
+-        return self._call_with_retries(token, retries=20)
++        return self._call_with_retries(token, retries=2,
++                                        backoff_seconds=0.5)

snapshots/orders_retry_storm/git_log.json

@@ -0,0 +1,23 @@
+[
+  {
+    "sha":     "f8c9b13",
+    "author":  "carol",
+    "date":    "2026-04-25T14:01:00Z",
+    "message": "resilience(orders): bump auth-client retries 3 -> 20",
+    "files":   ["orders/auth_client.py"]
+  },
+  {
+    "sha":     "d09a4f1",
+    "author":  "frankie",
+    "date":    "2026-04-25T14:03:00Z",
+    "message": "resilience(orders): bump auth-client retries 3 -> 25 with no jitter",
+    "files":   ["orders/auth_client.py"]
+  },
+  {
+    "sha":     "7b1aa90",
+    "author":  "carol",
+    "date":    "2026-04-21T09:30:00Z",
+    "message": "Initial auth-client implementation",
+    "files":   ["orders/auth_client.py"]
+  }
+]

snapshots/orders_retry_storm/tree/orders/auth_client.py

@@ -0,0 +1,35 @@
+"""Thin auth-service client used by the orders pipeline.
+
+After v2.5.1 the retry policy was bumped to 20 retries with no backoff,
+so any transient blip on auth amplifies into a flood of validation
+requests — auth's queue overflows and orders' validation times out
+again, triggering more retries.  Classic retry storm.
+"""
+
+from __future__ import annotations
+
+import time
+from typing import Any
+
+
+class _AuthRPCError(Exception):
+    pass
+
+
+class AuthClient:
+    """Validates user tokens against the auth service."""
+
+    def __init__(self, transport):
+        self._transport = transport
+
+    def validate(self, token: str) -> dict:
+        return self._call_with_retries(token, retries=20)
+
+    def _call_with_retries(self, token: str, retries: int) -> dict:
+        last_err: Exception | None = None
+        for _ in range(retries):
+            try:
+                return self._transport.rpc("auth.validate", {"token": token})
+            except _AuthRPCError as e:
+                last_err = e
+        raise _AuthRPCError(f"auth validate failed after {retries} retries: {last_err}")

snapshots/orders_v231/diffs/a3f7c91.patch

@@ -0,0 +1,50 @@
+diff --git a/orders/handlers/batch.py b/orders/handlers/batch.py
+new file mode 100644
+--- /dev/null
++++ b/orders/handlers/batch.py
+@@ -0,0 +1,52 @@
++"""Batched order processor."""
++
++from __future__ import annotations
++
++from typing import Iterable
++
++from ..models import Order
++from ..storage import OrderStore
++from ..notifier import Notifier
++
++
++class BatchProcessor:
++    def __init__(self, store, notifier, batch_size=100):
++        self._store      = store
++        self._notifier   = notifier
++        self._batch_size = batch_size
++        self._cache: dict[str, Order] = {}
++
++    def submit(self, order):
++        self._cache[order.id] = order
++        if len(self._cache) >= self._batch_size:
++            self.flush()
++
++    def flush(self):
++        orders = list(self._cache.values())
++        self._store.persist_many(orders)
++        self._notify(orders)
++
++    def submit_many(self, orders):
++        for order in orders:
++            self._cache[order.id] = order
++        self._notify(orders)
++
++    def _notify(self, orders):
++        for order in orders:
++            self._notifier.send(order.id, "submitted")
+diff --git a/orders/handlers/__init__.py b/orders/handlers/__init__.py
+--- a/orders/handlers/__init__.py
++++ b/orders/handlers/__init__.py
+@@ -1,3 +1,4 @@
+ from .single import SingleOrderHandler
++from .batch import BatchProcessor
+ 
+-__all__ = ["SingleOrderHandler"]
++__all__ = ["SingleOrderHandler", "BatchProcessor"]

snapshots/orders_v231/git_log.json

@@ -0,0 +1,34 @@
+[
+  {
+    "sha": "a3f7c91",
+    "author": "alice",
+    "date": "2025-01-15T13:50:00Z",
+    "message": "Feature: batch order processing with in-memory cache (v2.3.1)\n\nAdds BatchProcessor with submit_many path and an internal _cache to\navoid double-fetches when retrying batches. Reduces upstream queue load.",
+    "files": [
+      "orders/handlers/batch.py",
+      "orders/handlers/__init__.py",
+      "orders/__init__.py"
+    ]
+  },
+  {
+    "sha": "31aa72c",
+    "author": "alice",
+    "date": "2025-01-12T09:14:00Z",
+    "message": "Add Notifier abstraction + dependency injection",
+    "files": [
+      "orders/notifier.py",
+      "orders/handlers/single.py"
+    ]
+  },
+  {
+    "sha": "0e1d8b2",
+    "author": "bob",
+    "date": "2025-01-09T16:02:00Z",
+    "message": "Initial single-order handler",
+    "files": [
+      "orders/handlers/single.py",
+      "orders/storage.py",
+      "orders/models.py"
+    ]
+  }
+]

snapshots/orders_v231/tree/orders/__init__.py

@@ -0,0 +1,7 @@
+"""orders service — order ingestion, batching, and persistence."""
+
+from .handlers.batch import BatchProcessor
+from .handlers.single import SingleOrderHandler
+from .models import Order
+
+__all__ = ["BatchProcessor", "SingleOrderHandler", "Order"]

snapshots/orders_v231/tree/orders/handlers/__init__.py

@@ -0,0 +1,4 @@
+from .batch import BatchProcessor
+from .single import SingleOrderHandler
+
+__all__ = ["BatchProcessor", "SingleOrderHandler"]

snapshots/orders_v231/tree/orders/handlers/batch.py

@@ -0,0 +1,60 @@
+"""
+Batched order processor.
+
+Released in v2.3.1 to "speed up high-volume checkout windows".  The batch
+handler consumes incoming orders, persists them, and emits notifications.
+
+NOTE: The previous version (v1.x) processed orders one-by-one — see
+single.py.  v2.3.1 added an in-memory cache to avoid re-querying upstream.
+"""
+
+from __future__ import annotations
+
+from typing import Iterable, List
+
+from ..models import Order
+from ..storage import OrderStore
+from ..notifier import Notifier
+
+
+class BatchProcessor:
+    """
+    Accumulate `Order` objects in batches and flush to storage.
+
+    The `_cache` field was added in v2.3.1 to avoid double-fetching the
+    same order from the upstream queue if a batch was retried.  See the
+    PR description on commit a3f7c91 for context.
+    """
+
+    def __init__(
+        self,
+        store: OrderStore,
+        notifier: Notifier,
+        batch_size: int = 100,
+    ) -> None:
+        self._store      = store
+        self._notifier   = notifier
+        self._batch_size = batch_size
+        self._cache: dict[str, Order] = {}
+
+    def submit(self, order: Order) -> None:
+        self._cache[order.id] = order
+        if len(self._cache) >= self._batch_size:
+            self.flush()
+
+    def flush(self) -> None:
+        orders = list(self._cache.values())
+        self._store.persist_many(orders)
+        self._notify(orders)
+
+    def submit_many(self, orders: Iterable[Order]) -> None:
+        for order in orders:
+            self._cache[order.id] = order
+        self._notify(orders)
+
+    def _notify(self, orders: Iterable[Order]) -> None:
+        for order in orders:
+            self._notifier.send(order.id, "submitted")
+
+    def cache_size(self) -> int:
+        return len(self._cache)

snapshots/orders_v231/tree/orders/handlers/single.py

@@ -0,0 +1,17 @@
+"""Pre-v2.3.1 single-order handler. Still used as the fallback path."""
+
+from __future__ import annotations
+
+from ..models import Order
+from ..storage import OrderStore
+from ..notifier import Notifier
+
+
+class SingleOrderHandler:
+    def __init__(self, store: OrderStore, notifier: Notifier) -> None:
+        self._store    = store
+        self._notifier = notifier
+
+    def submit(self, order: Order) -> None:
+        self._store.persist(order)
+        self._notifier.send(order.id, "submitted")

snapshots/orders_v231/tree/orders/models.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Dict, Optional
+
+
+@dataclass
+class Order:
+    id: str
+    customer_id: str
+    line_items: Dict[str, int] = field(default_factory=dict)
+    total_cents: int = 0
+    coupon_code: Optional[str] = None
+    metadata: Dict[str, str] = field(default_factory=dict)

snapshots/orders_v231/tree/orders/notifier.py

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+
+class Notifier:
+    """Outbound notification client — pushes status to downstream queue."""
+
+    def send(self, order_id: str, event: str) -> None:
+        ...

snapshots/orders_v231/tree/orders/storage.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+from typing import Iterable, List
+
+from .models import Order
+
+
+class OrderStore:
+    """Postgres-backed order persistence (mock — talks to db.orders)."""
+
+    def persist(self, order: Order) -> None:
+        ...
+
+    def persist_many(self, orders: Iterable[Order]) -> None:
+        for o in orders:
+            self.persist(o)
+
+    def get(self, order_id: str) -> Order:
+        ...

snapshots/orders_v300/diffs/d2b9c11.patch

@@ -0,0 +1,67 @@
+diff --git a/orders/circuit_breaker.py b/orders/circuit_breaker.py
+new file mode 100644
+--- /dev/null
++++ b/orders/circuit_breaker.py
+@@ -0,0 +1,55 @@
++"""Circuit breaker on the payment dependency (intentional)."""
++
++from __future__ import annotations
++
++import time
++from dataclasses import dataclass
++
++
++@dataclass
++class BreakerConfig:
++    failure_threshold:    int   = 10
++    success_threshold:    int   = 3
++    open_duration_seconds: float = 30.0
++    latency_threshold_ms:  float = 1500.0
++
++
++class CircuitBreaker:
++    def __init__(self, config: BreakerConfig | None = None) -> None:
++        self._cfg          = config or BreakerConfig()
++        self._state        = "closed"
++        self._failures     = 0
++        self._successes    = 0
++        self._opened_at: float = 0.0
++
++    def allow(self) -> bool:
++        if self._state == "open":
++            if time.time() - self._opened_at >= self._cfg.open_duration_seconds:
++                self._state = "half_open"
++                self._successes = 0
++                return True
++            return False
++        return True
++
++    def record_success(self, latency_ms: float) -> None:
++        if latency_ms > self._cfg.latency_threshold_ms:
++            self.record_failure()
++            return
++        if self._state == "half_open":
++            self._successes += 1
++            if self._successes >= self._cfg.success_threshold:
++                self._state = "closed"
++                self._failures = 0
++
++    def record_failure(self) -> None:
++        self._failures += 1
++        if self._failures >= self._cfg.failure_threshold:
++            self._state = "open"
++            self._opened_at = time.time()
+diff --git a/orders/handlers/checkout.py b/orders/handlers/checkout.py
+--- a/orders/handlers/checkout.py
++++ b/orders/handlers/checkout.py
+@@ -3,6 +3,8 @@
+ import time
+ 
++from ..circuit_breaker import CircuitBreaker
++
+ 
+ class Checkout:
+-    def __init__(self, payment_client) -> None:
++    def __init__(self, payment_client, breaker: CircuitBreaker | None = None) -> None:
+         self._payment = payment_client
++        self._breaker = breaker or CircuitBreaker()

snapshots/orders_v300/git_log.json

@@ -0,0 +1,22 @@
+[
+  {
+    "sha": "d2b9c11",
+    "author": "alice",
+    "date": "2025-01-15T13:42:00Z",
+    "message": "Add circuit breaker on payment dependency (v3.0.0)\n\nDefensive guard against the retry-storm pattern from incident-2024-Q4-3.\nReturns 503 immediately when payment latency exceeds 1500ms for 10\nconsecutive calls. Resets after a 30-second cooling period.",
+    "files": [
+      "orders/circuit_breaker.py",
+      "orders/handlers/checkout.py",
+      "orders/__init__.py"
+    ]
+  },
+  {
+    "sha": "97cba12",
+    "author": "bob",
+    "date": "2025-01-12T15:00:00Z",
+    "message": "Initial checkout handler",
+    "files": [
+      "orders/handlers/checkout.py"
+    ]
+  }
+]

snapshots/orders_v300/tree/orders/__init__.py

@@ -0,0 +1,6 @@
+"""orders service v3.0 — adds defensive circuit-breaker on payment dependency."""
+
+from .handlers.checkout import Checkout
+from .circuit_breaker import CircuitBreaker
+
+__all__ = ["Checkout", "CircuitBreaker"]

snapshots/orders_v300/tree/orders/circuit_breaker.py

@@ -0,0 +1,65 @@
+"""
+Circuit breaker on the payment dependency.
+
+This component is INTENTIONAL — it was added in v3.0.0 (commit d2b9c11)
+to prevent the well-known retry-storm failure mode that brought down
+auth in incident-2024-Q4-3.
+
+When payment latency exceeds the threshold for too long, the breaker
+trips and returns 503 immediately to upstream callers, draining load.
+This is correct behaviour — *not* a bug — but customers who only see
+"orders is returning 503s" sometimes file tickets thinking it is one.
+
+Do NOT remove this guard.
+"""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+
+
+@dataclass
+class BreakerConfig:
+    failure_threshold:    int   = 10
+    success_threshold:    int   = 3
+    open_duration_seconds: float = 30.0
+    latency_threshold_ms:  float = 1500.0
+
+
+class CircuitBreaker:
+    """Three-state circuit breaker: closed → open → half-open → closed."""
+
+    def __init__(self, config: BreakerConfig | None = None) -> None:
+        self._cfg          = config or BreakerConfig()
+        self._state        = "closed"   # closed | open | half_open
+        self._failures     = 0
+        self._successes    = 0
+        self._opened_at: float = 0.0
+
+    def allow(self) -> bool:
+        if self._state == "open":
+            if time.time() - self._opened_at >= self._cfg.open_duration_seconds:
+                self._state = "half_open"
+                self._successes = 0
+                return True
+            return False
+        return True
+
+    def record_success(self, latency_ms: float) -> None:
+        if latency_ms > self._cfg.latency_threshold_ms:
+            self.record_failure()
+            return
+        if self._state == "half_open":
+            self._successes += 1
+            if self._successes >= self._cfg.success_threshold:
+                self._state = "closed"
+                self._failures = 0
+        elif self._state == "closed":
+            self._failures = max(0, self._failures - 1)
+
+    def record_failure(self) -> None:
+        self._failures += 1
+        if self._failures >= self._cfg.failure_threshold:
+            self._state = "open"
+            self._opened_at = time.time()

snapshots/orders_v300/tree/orders/handlers/__init__.py

@@ -0,0 +1,3 @@
+from .checkout import Checkout
+
+__all__ = ["Checkout"]

snapshots/orders_v300/tree/orders/handlers/checkout.py

@@ -0,0 +1,29 @@
+"""Checkout endpoint — wraps the payment call in the circuit breaker."""
+
+from __future__ import annotations
+
+import time
+
+from ..circuit_breaker import CircuitBreaker
+
+
+class Checkout:
+    def __init__(self, payment_client, breaker: CircuitBreaker | None = None) -> None:
+        self._payment = payment_client
+        self._breaker = breaker or CircuitBreaker()
+
+    def submit(self, order_id: str, amount_cents: int) -> str:
+        if not self._breaker.allow():
+            return "503 — payment temporarily unavailable, please retry"
+        start = time.time()
+        try:
+            ok = self._payment.charge(order_id, amount_cents)
+        except Exception:
+            self._breaker.record_failure()
+            raise
+        latency_ms = (time.time() - start) * 1000
+        if ok:
+            self._breaker.record_success(latency_ms)
+            return "201 ok"
+        self._breaker.record_failure()
+        return "402 declined"

snapshots/payment_threadpool/diffs/11abf04.patch

@@ -0,0 +1,10 @@
+diff --git a/payment/threadpool.py b/payment/threadpool.py
+--- a/payment/threadpool.py
++++ b/payment/threadpool.py
+@@ -8,4 +8,5 @@ class PoolWorker:
+     def acquire(self):
+-        self._lock_a.acquire()
+-        self._lock_b.acquire()
++        with self._global_order:
++            self._lock_a.acquire()
++            self._lock_b.acquire()

snapshots/payment_threadpool/git_log.json

@@ -0,0 +1,16 @@
+[
+  {
+    "sha":     "11abf04",
+    "author":  "dave",
+    "date":    "2026-04-25T13:35:00Z",
+    "message": "refactor(payment): drop global order, hold per-row locks directly",
+    "files":   ["payment/threadpool.py"]
+  },
+  {
+    "sha":     "660ff21",
+    "author":  "dave",
+    "date":    "2026-04-21T14:02:00Z",
+    "message": "Initial PoolWorker implementation",
+    "files":   ["payment/threadpool.py"]
+  }
+]

snapshots/payment_threadpool/tree/payment/threadpool.py

@@ -0,0 +1,53 @@
+"""Bounded thread pool used to acquire two row-locks per payment txn.
+
+After v3.1.2 the lock-acquisition order was changed but the global
+ordering invariant was dropped, so two concurrent workers can each
+hold one of the locks the other needs — classic AB-BA deadlock.
+
+Symptom: threads accumulate (active_count climbs to pool max) but
+CPU stays low.  Memory grows because each blocked thread keeps its
+stack pinned.  From the outside it looks indistinguishable from a
+heap memory leak.
+"""
+
+from __future__ import annotations
+
+import threading
+from typing import Any
+
+
+class _LockOrder:
+    """Sentinel holding the global acquisition order (must hold first)."""
+    def __init__(self):
+        self._lock = threading.Lock()
+
+    def __enter__(self):
+        self._lock.acquire()
+        return self
+
+    def __exit__(self, *exc):
+        self._lock.release()
+
+
+class PoolWorker:
+    """Acquires two per-account row locks before mutating payment state."""
+
+    def __init__(self):
+        self._lock_a = threading.Lock()
+        self._lock_b = threading.Lock()
+        self._global_order = _LockOrder()
+
+    def acquire(self) -> None:
+        self._lock_a.acquire()
+        self._lock_b.acquire()
+
+    def release(self) -> None:
+        self._lock_b.release()
+        self._lock_a.release()
+
+    def transfer(self, amount: int) -> None:
+        self.acquire()
+        try:
+            ...
+        finally:
+            self.release()

snapshots/payment_v310/diffs/c5a1f77.patch

@@ -0,0 +1,29 @@
+diff --git a/payment/processor.py b/payment/processor.py
+--- a/payment/processor.py
++++ b/payment/processor.py
+@@ -1,7 +1,6 @@
+ """Payment processor — orchestrates the gateway → queue → orders flow."""
+ 
+ from __future__ import annotations
+ 
+-import time
+-from typing import Optional
++from typing import Optional
+ 
+@@ -23,7 +22,7 @@ class PaymentProcessor:
+         self,
+         gateway: StripeGateway,
+         queue:   QueueClient,
+-        max_retries: int = 3,
++        max_retries: int = 10,
+     ) -> None:
+         self._gateway      = gateway
+         self._queue        = queue
+@@ -40,9 +39,7 @@ class PaymentProcessor:
+     def retry(self, txn: Transaction) -> None:
+         if self.retry_count >= self._max_retries:
+             raise RuntimeError(f"retries exhausted for {txn.id}")
+         self.retry_count += 1
+-        delay = min(2 ** self.retry_count, 30)
+-        time.sleep(delay)
+         self._queue.enqueue(txn)

snapshots/payment_v310/git_log.json

@@ -0,0 +1,33 @@
+[
+  {
+    "sha": "c5a1f77",
+    "author": "charlie",
+    "date": "2025-01-15T13:48:00Z",
+    "message": "Improve payment reliability: increase retry count, drop backoff (v3.1.0)\n\nUnder high failure rates the previous exponential-backoff was too slow\nto recover. Bumps retry_count 3 -> 10 and removes the time.sleep() so\nretries fire as fast as the queue can accept them.",
+    "files": [
+      "payment/processor.py"
+    ]
+  },
+  {
+    "sha": "44d3e90",
+    "author": "alice",
+    "date": "2025-01-13T08:10:00Z",
+    "message": "Add QueueClient.ack polling stub",
+    "files": [
+      "payment/queue_client.py"
+    ]
+  },
+  {
+    "sha": "2bb01af",
+    "author": "alice",
+    "date": "2025-01-09T14:30:00Z",
+    "message": "Initial payment processor + gateway scaffolding",
+    "files": [
+      "payment/__init__.py",
+      "payment/processor.py",
+      "payment/gateway.py",
+      "payment/queue_client.py",
+      "payment/txn.py"
+    ]
+  }
+]

snapshots/payment_v310/tree/payment/__init__.py

@@ -0,0 +1,6 @@
+"""payment service — txn lifecycle, gateway adapters, retry orchestration."""
+
+from .processor import PaymentProcessor
+from .gateway import StripeGateway
+
+__all__ = ["PaymentProcessor", "StripeGateway"]

snapshots/payment_v310/tree/payment/gateway.py

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+
+class StripeGateway:
+    """Outbound HTTP client to the upstream payment processor."""
+
+    def charge(self, txn_id: str, amount_cents: int, card_token: str) -> bool:
+        ...