env/spindleflow_env.py

"""
SpindleFlowEnv — Main RL environment.
Gymnasium-compatible. Wraps SpindleFlow as the execution backend.
LSTM-policy-safe: state representation is complete per-step (no hidden history).

The environment does NOT call SpindleFlow for every episode during training —
that would be too slow and expensive. Instead, for Phase 1/2 training it uses
a simulated specialist execution (fast, free). For evaluation and demo, it
calls real SpindleFlow.
"""

from __future__ import annotations
import time
import numpy as np
import gymnasium as gym
from gymnasium import spaces
from pathlib import Path
from typing import Optional, Any
import yaml

from env.specialist_registry import SpecialistRegistry
from env.delegation_graph import DelegationGraph
from env.scratchpad import SharedScratchpad
from env.state import build_state, EpisodeState
from env.action_space import ActionDecoder, MetaAction, FactoredAction, DelegationMode
from reward.tier_lock import EpisodeTierLock
from reward.tiered_reward import TieredRewardScorer
from reward.latency_reward import LatencySLAConfig, compute_latency_penalty
from reward.failure_reward import (
    SpecialistResult, SpecialistStatus,
    compute_failure_penalty, compute_recovery_bonus,
)
from reward.conflict_reward import detect_conflicts
from reward.consistency_tracker import PathConsistencyTracker
from agents.task_decomposer import TaskDecomposer, EnrichedTask
from agents.conflict_resolver import ConflictResolver
from agents.fallback_chain import FallbackChainResolver
from agents.specialist_memory import SpecialistMemory
from training.spawn_memory import SpawnMemory, SpawnRecord
from training.task_bank import TaskBank


class SpindleFlowEnv(gym.Env):
    """
    RL Environment for SpindleFlow delegation policy training.

    Episode structure:
      1. Reset: Draw task from task bank, embed it, lock tier, set up components
      2. Step loop: Policy chooses action → environment executes → compute reward
      3. Termination: STOP action, max_steps reached, or episode error

    Observation space: Flat vector (see EpisodeState.observation_dim())
    Action space: Box (continuous — decoded by ActionDecoder)
    """

    metadata = {"render_modes": ["human"]}

    def __init__(
        self,
        config_path: str = "configs/training_config.yaml",
        catalog_path: str = "configs/specialist_catalog.yaml",
        use_real_spindleflow: bool = False,
        phase: int = 1,
        render_mode: Optional[str] = None,
        simulate_specialists: bool = False,
    ):
        super().__init__()

        with open(config_path) as f:
            self.config = yaml.safe_load(f)

        env_cfg = self.config["environment"]
        self.max_steps = env_cfg["max_steps_per_episode"]
        self.max_depth = env_cfg["max_delegation_depth"]
        self.max_specialists = env_cfg.get("max_specialists_per_episode", 6)
        self.specialist_timeout_ms = env_cfg["specialist_timeout_ms"]
        self.phase = phase
        self.use_real_spindleflow = use_real_spindleflow
        self.render_mode = render_mode
        # When True: per-step specialist calls use simulation even if OPENAI_API_KEY
        # is set. Episode-level self-learning (finetuner, spawn) still use the key.
        self.simulate_specialists = simulate_specialists

        reward_cfg = self.config["reward"]
        self.latency_sla = LatencySLAConfig(
            budget_ms=10000.0,
            weight=reward_cfg["latency_weight"],
        )

        # Initialize components
        self.registry = SpecialistRegistry(catalog_path)
        self.task_bank = TaskBank(
            phase=phase,
            config_path=config_path,
            catalog_path=catalog_path,
        )
        # Load sector contradiction pairs from catalog (for conflict detection)
        with open(catalog_path) as _f:
            _catalog_meta = yaml.safe_load(_f).get("metadata", {})
        self._contradiction_pairs = [
            tuple(pair) for pair in _catalog_meta.get("contradiction_pairs", [])
        ]

        self.task_decomposer = TaskDecomposer(sector_cfg=self.config.get("sector", {}))
        _resolution_mem_path = self.config.get("agents", {}).get(
            "resolution_memory_path", "data/resolution_memory.jsonl"
        )
        self.conflict_resolver = ConflictResolver(
            config=self.config,
            memory_path=_resolution_mem_path,
        )
        self.fallback_resolver = FallbackChainResolver()
        self.reward_scorer = TieredRewardScorer(registry=self.registry)
        self.consistency_tracker = PathConsistencyTracker(
            specialist_ids=self.registry.list_ids()
        )
        si_cfg = self.config.get("specialist_improvement", {})
        memory_path = si_cfg.get("memory_path", "data/specialist_memory.json")
        self.specialist_memory = SpecialistMemory(path=memory_path)

        spawn_mem_path = env_cfg.get("spawn_memory_path", "data/spawn_memory.jsonl")
        self._spawn_memory = SpawnMemory(
            path=spawn_mem_path,
            max_entries=env_cfg.get("spawn_memory_max_entries", 500),
        )
        self._pending_spawn_records: list[SpawnRecord] = []
        self.action_decoder = ActionDecoder(
            specialist_ids=self.registry.list_ids(),
            max_specialists=self.max_specialists,
        )

        # Spawn config
        self.spawn_threshold: float = env_cfg.get("spawn_threshold", 0.50)
        self.auto_spawn: bool = env_cfg.get("auto_spawn_specialists", True)
        # Max total spawned specialists across the lifetime of this env instance.
        # Caps registry growth so the observation space stays stable during long runs.
        self._spawn_max_total: int = env_cfg.get("spawn_max_total", 8)
        # Minimum episodes between consecutive spawns — prevents burst-spawning on
        # a streak of low-similarity tasks and keeps the action decoder stable.
        self._spawn_cooldown_episodes: int = env_cfg.get("spawn_cooldown_episodes", 20)
        # Lifetime counters (survive across resets)
        self._spawn_total_count: int = 0
        self._last_spawn_episode: int = -999  # episode index of last spawn
        self._episode_index: int = 0

        # Per-episode state
        self.delegation_graph = DelegationGraph(max_depth=self.max_depth)
        self.scratchpad = SharedScratchpad()
        self.current_task: Optional[EnrichedTask] = None
        self.tier_lock: Optional[EpisodeTierLock] = None
        self.specialist_results: list[SpecialistResult] = []
        self.called_ids: list[str] = []
        self.step_count: int = 0
        self.episode_start_ms: float = 0.0
        self.generalist_baseline: str = ""
        self.config_reward = reward_cfg
        self._last_reward_components: dict = {}
        self._last_factored_action: Optional[Any] = None
        # Active roster for this episode (top-K by task similarity, including spawned)
        self.active_specialist_ids: list[str] = self.registry.list_ids()[:self.max_specialists]
        self.spawned_this_episode: list[str] = []
        # Task embedding cached at reset() — constant within an episode, no need to re-embed each step
        self._task_emb: np.ndarray | None = None

        # Spaces
        obs_dim = EpisodeState.observation_dim(self.max_specialists)
        self.observation_space = spaces.Box(
            low=-10.0, high=10.0, shape=(obs_dim,), dtype=np.float32
        )
        self.action_space = spaces.Box(
            low=-1.0, high=1.0,
            shape=(self.action_decoder.get_action_dim(),),
            dtype=np.float32,
        )

    def reset(
        self,
        seed: Optional[int] = None,
        options: Optional[dict] = None,
    ) -> tuple[np.ndarray, dict]:
        super().reset(seed=seed)

        self.delegation_graph.reset()
        self.scratchpad.reset(episode_id=str(time.time()))
        self.specialist_results = []
        self.called_ids = []
        self.step_count = 0
        self.episode_start_ms = time.time() * 1000

        task_desc = self.task_bank.sample()
        self.current_task = self.task_decomposer.decompose(task_desc)

        self.tier_lock = EpisodeTierLock.for_task(
            self.current_task.complexity_class
        )

        self.generalist_baseline = self._generate_generalist_baseline(
            self.current_task.enriched_description
        )

        self.delegation_graph.add_root("orchestrator")
        self._episode_index += 1

        task_desc = self.current_task.enriched_description
        task_emb  = self.registry.embed_query(task_desc)
        assert task_emb is not None and task_emb.shape == (384,), (
            f"Task embedding failed: got shape {getattr(task_emb, 'shape', None)}"
        )
        self._task_emb = task_emb   # cached for entire episode — task doesn't change

        self.spawned_this_episode = []
        self._pending_spawn_records = []
        # Auto-spawn: if no existing specialist covers this task well, create one via LLM.
        if self.auto_spawn:
            self._maybe_spawn_specialist(task_emb, task_desc)

        # ── Build per-episode active roster (top-K by task similarity) ──
        self.active_specialist_ids = self._select_active_specialists(task_emb)

        # ── Rebuild action decoder to reflect the updated roster ──
        self.action_decoder = ActionDecoder(
            specialist_ids=self.active_specialist_ids,
            max_specialists=self.max_specialists,
        )

        state = build_state(
            task_embedding=task_emb,
            registry=self.registry,
            called_ids=[],
            delegation_graph=self.delegation_graph,
            scratchpad=self.scratchpad,
            step_count=0,
            elapsed_ms=0.0,
            sla_budget_ms=self.latency_sla.budget_ms,
            max_specialists=self.max_specialists,
            max_depth=self.max_depth,
            phase=self.phase,
            active_ids=self.active_specialist_ids,
        )

        info = {
            "task":               task_desc,
            "complexity":         self.current_task.complexity_class,
            "tier":               self.tier_lock.locked_tier.name,
            "active_specialists": list(self.active_specialist_ids),
            "spawned_specialists": list(self.spawned_this_episode),
        }

        return state.to_flat_vector(), info

    def step(
        self, action: np.ndarray
    ) -> tuple[np.ndarray, float, bool, bool, dict]:
        """
        Execute one step in the environment.
        Returns: (observation, reward, terminated, truncated, info)
        """
        self.step_count += 1
        elapsed_ms = time.time() * 1000 - self.episode_start_ms

        # Build specialist mask (enforce DAG constraints)
        valid_ids = self.delegation_graph.get_valid_callees(
            "orchestrator", self.active_specialist_ids
        )
        valid_ids = [sid for sid in valid_ids if sid not in self.called_ids]
        mask = self.action_decoder.build_specialist_mask(valid_ids)

        factored: FactoredAction = self.action_decoder.decode(action, mask)

        assert self._task_emb is not None, (
            "step() called before reset() or task embedding failed in reset()"
        )
        task_emb = self._task_emb

        terminated = False
        truncated = False
        step_results = []

        if factored.meta_action == MetaAction.STOP or self.step_count >= self.max_steps:
            terminated = True
        else:
            step_results = self._dispatch_meta_action(factored, elapsed_ms)
            self.specialist_results.extend(step_results)
            _reg = set(self.registry.list_ids())
            self.called_ids.extend(
                r.specialist_id for r in step_results
                if r.specialist_id in _reg
            )

        if self.step_count >= self.max_steps and not terminated:
            truncated = True
        state = build_state(
            task_embedding=task_emb,
            registry=self.registry,
            called_ids=self.called_ids,
            delegation_graph=self.delegation_graph,
            scratchpad=self.scratchpad,
            step_count=self.step_count,
            elapsed_ms=elapsed_ms,
            sla_budget_ms=self.latency_sla.budget_ms,
            max_specialists=self.max_specialists,
            max_depth=self.max_depth,
            phase=self.phase,
            active_ids=self.active_specialist_ids,
        )

        if terminated or truncated:
            reward = self._compute_final_reward(elapsed_ms)
            self._record_episode_to_memory(reward)
        else:
            reward = self._compute_step_reward(
                step_results, task_emb,
                delegation_mode=factored.delegation_mode,
                meta_action=factored.meta_action,
            )

        step_latencies = {r.specialist_id: r.latency_ms for r in step_results}
        info = {
            # Keys expected by the UI / Streamlit dashboard
            "action_name":           factored.meta_action.name,
            "called_specialists":    list(factored.specialist_ids),
            "delegation_mode":       factored.delegation_mode.name,
            "reward_components":     dict(self._last_reward_components),
            "specialist_latencies":  step_latencies,
            "active_specialists":    list(self.active_specialist_ids),
            "spawned_specialists":   list(self.spawned_this_episode),
            # Raw data for debugging / training callbacks
            "action":                factored.to_log_dict(),
            "called_ids":            list(self.called_ids),
            "step_count":            self.step_count,
            "elapsed_ms":            elapsed_ms,
        }

        return state.to_flat_vector(), reward, terminated, truncated, info

    # ── MetaAction dispatch ───────────────────────────────────────────

    def _dispatch_meta_action(
        self, action: FactoredAction, elapsed_ms: float
    ) -> list[SpecialistResult]:
        """Route to the correct handler based on MetaAction."""
        if action.meta_action == MetaAction.CALL_MEDIATOR:
            return self._exec_meta_mediator(action, elapsed_ms)
        if action.meta_action == MetaAction.CLARIFY_TASK:
            return self._exec_meta_clarify(action, elapsed_ms)
        if action.meta_action == MetaAction.DELEGATE_SUBTASK:
            return self._exec_meta_delegate_subtask(action, elapsed_ms)
        if action.meta_action == MetaAction.RETRY_FAILED:
            return self._exec_meta_retry(action, elapsed_ms)
        if action.meta_action == MetaAction.PARALLEL_SPAWN:
            return self._exec_meta_parallel_spawn(action, elapsed_ms)
        if action.meta_action == MetaAction.SPAWN_SPECIALIST:
            return self._exec_meta_spawn_specialist(action, elapsed_ms)
        return self._execute_action(action, elapsed_ms)   # CALL_SPECIALIST default

    # ── DelegationMode dispatch ───────────────────────────────────────

    def _execute_action(
        self, action: FactoredAction, elapsed_ms: float
    ) -> list[SpecialistResult]:
        """Dispatch to the correct DelegationMode handler."""
        handlers = {
            DelegationMode.SEQUENTIAL:     self._exec_sequential,
            DelegationMode.PARALLEL:       self._exec_parallel,
            DelegationMode.FAN_OUT_REDUCE: self._exec_fan_out_reduce,
            DelegationMode.ITERATIVE:      self._exec_iterative,
            DelegationMode.CONDITIONAL:    self._exec_conditional,
            DelegationMode.PRIORITY_QUEUE: self._exec_priority_queue,
            DelegationMode.BROADCAST:      self._exec_broadcast,
        }
        return handlers.get(action.delegation_mode, self._exec_sequential)(action, elapsed_ms)

    # ── Shared helpers ────────────────────────────────────────────────

    def _can_call(self, sid: str, caller_id: str = "orchestrator") -> bool:
        """True when a specialist is registered, not yet called, and DAG-valid."""
        return (
            sid in self.registry.list_ids()
            and sid not in self.called_ids
            and self.delegation_graph.can_delegate(caller_id, sid)
        )

    def _do_call(
        self,
        sid: str,
        task: str,
        elapsed_ms: float,
        mode: str = "SEQUENTIAL",
        context: str | None = None,
        caller_id: str = "orchestrator",
    ) -> list[SpecialistResult]:
        """
        Validate → record in DAG → call specialist → handle fallback → write scratchpad.

        caller_id controls which node in the delegation graph is the caller.
        Defaults to "orchestrator" for top-level calls. Pass a specialist ID
        to record depth-2 delegations (specialist → sub-specialist).
        Returns a list because a fallback may contribute a second result.
        """
        if not self._can_call(sid, caller_id=caller_id):
            return []
        self.delegation_graph.record_delegation(caller_id, sid, mode)
        result = self._call_specialist(sid, task, elapsed_ms, context=context)
        if result.output:
            self.scratchpad.write(
                author_id=sid,
                author_role=self.registry.get(sid).role,
                content=result.output,
            )
        results = [result]
        if self.fallback_resolver.needs_fallback(result):
            fb_id = self.fallback_resolver.get_fallback(sid, self.called_ids)
            if fb_id and self._can_call(fb_id):
                self.delegation_graph.record_delegation("orchestrator", fb_id, mode)
                fb = self._call_specialist(
                    fb_id, self.current_task.enriched_description, elapsed_ms
                )
                fb.fallback_used = True
                if fb.output:
                    self.scratchpad.write(
                        author_id=fb_id,
                        author_role=self.registry.get(fb_id).role,
                        content=fb.output,
                    )
                results.append(fb)
                # Do NOT append fb_id here — step() uniformly extends called_ids
                # from all step_results after _do_call returns, so appending here
                # would cause a double-count (efficiency penalty and DAG mask both
                # use called_ids, making the fallback specialist appear called twice).
        return results

    def _quick_quality_score(self, output: str, task: str) -> float:
        """Fast T1 cosine similarity — used for within-step stopping conditions."""
        try:
            t = self.registry.embed_query(task)
            o = self.registry.embed_query(output[:800])
            return float((self.registry.cosine_similarity(t, o) + 1.0) / 2.0)
        except Exception:
            return 0.5

    def _synthesize_outputs(self, outputs: list[str]) -> str:
        """Merge multiple specialist outputs into one coherent synthesis."""
        import os
        if os.getenv("OPENAI_API_KEY") and len(outputs) >= 2:
            try:
                from openai import OpenAI
                combined = "\n\n---\n\n".join(
                    f"Specialist {i+1}:\n{o[:500]}" for i, o in enumerate(outputs)
                )
                client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
                resp = client.chat.completions.create(
                    model="gpt-4o-mini", max_tokens=600,
                    messages=[
                        {"role": "system", "content":
                            "Synthesize these specialist analyses into one coherent "
                            "recommendation. Resolve contradictions, highlight consensus."},
                        {"role": "user", "content": combined[:2000]},
                    ],
                )
                return resp.choices[0].message.content
            except Exception as exc:
                print(f"[Synthesize] {exc}")
        joined = "\n\n".join(f"[{i+1}] {o[:200]}" for i, o in enumerate(outputs))
        return (
            f"Synthesis of {len(outputs)} specialist outputs:\n{joined}\n"
            "Consensus: structured design, domain best practices, iterative validation."
        )

    # ── DelegationMode handlers ───────────────────────────────────────

    def _exec_sequential(
        self, action: FactoredAction, elapsed_ms: float
    ) -> list[SpecialistResult]:
        """A→B→C: each specialist receives accumulated context from prior outputs.
        Highest quality for dependent sub-problems."""
        results: list[SpecialistResult] = []
        context = ""
        for sid in action.specialist_ids:
            batch = self._do_call(
                sid, self.current_task.enriched_description,
                elapsed_ms, mode="SEQUENTIAL",
                context=context or None,
            )
            results.extend(batch)
            for r in batch:
                if r.output:
                    context += f"\n{r.output[:400]}"
        return results

    def _exec_parallel(
        self, action: FactoredAction, elapsed_ms: float
    ) -> list[SpecialistResult]:
        """All specialists see the same task independently — no context sharing.
        Lower quality than SEQUENTIAL, lower effective latency for independent work."""
        results: list[SpecialistResult] = []
        for sid in action.specialist_ids:
            results.extend(
                self._do_call(
                    sid, self.current_task.enriched_description,
                    elapsed_ms, mode="PARALLEL",
                )
            )
        return results

    def _exec_fan_out_reduce(
        self, action: FactoredAction, elapsed_ms: float
    ) -> list[SpecialistResult]:
        """Fan-out: all specialists run independently; reduce: a synthesis pass
        merges all outputs into one recommendation. Highest quality, highest cost."""
        results = self._exec_parallel(action, elapsed_ms)
        successful_outs = [
            r.output for r in results
            if r.status == SpecialistStatus.SUCCESS and r.output
        ]
        if len(successful_outs) >= 2:
            synthesis = self._synthesize_outputs(successful_outs)
            synth = SpecialistResult(
                specialist_id="synthesizer",
                status=SpecialistStatus.SUCCESS,
                output=synthesis,
                latency_ms=0.0,
            )
            self.scratchpad.write(
                author_id="synthesizer",
                author_role="Synthesis Mediator",
                content=synthesis,
            )
            results.append(synth)
        return results

    def _exec_iterative(
        self, action: FactoredAction, elapsed_ms: float
    ) -> list[SpecialistResult]:
        """Repeatedly call one specialist, feeding its output back as context,
        until quality threshold met or max_rounds exhausted."""
        if not action.specialist_ids:
            return []
        sid        = action.specialist_ids[0]
        max_rounds = int(action.mode_params.get("max_rounds", 3))
        threshold  = float(action.mode_params.get("quality_threshold", 0.70))
        results: list[SpecialistResult] = []
        context = ""
        for _ in range(max(1, max_rounds)):
            batch = self._do_call(
                sid, self.current_task.enriched_description,
                elapsed_ms, mode="ITERATIVE",
                context=context or None,
            )
            results.extend(batch)
            for r in batch:
                if r.output:
                    if self._quick_quality_score(r.output, self.current_task.enriched_description) >= threshold:
                        return results
                    context = r.output
        return results

    def _exec_conditional(
        self, action: FactoredAction, elapsed_ms: float
    ) -> list[SpecialistResult]:
        """Call specialists in order; stop as soon as one meets the quality
        threshold — avoids unnecessary calls when the first is sufficient."""
        threshold = float(action.mode_params.get("condition_threshold", 0.60))
        results: list[SpecialistResult] = []
        for sid in action.specialist_ids:
            batch = self._do_call(
                sid, self.current_task.enriched_description,
                elapsed_ms, mode="CONDITIONAL",
            )
            results.extend(batch)
            for r in batch:
                if r.output and self._quick_quality_score(
                    r.output, self.current_task.enriched_description
                ) >= threshold:
                    return results
        return results

    def _exec_priority_queue(
        self, action: FactoredAction, elapsed_ms: float
    ) -> list[SpecialistResult]:
        """Sort selected specialists by task-similarity, call highest-ranked first,
        stop when output quality meets stop_threshold. Good for SLA-sensitive tasks."""
        threshold = float(action.mode_params.get("stop_threshold", 0.70))
        task_emb  = self.registry.embed_query(self.current_task.enriched_description)
        sorted_sids = sorted(
            [sid for sid in action.specialist_ids if self._can_call(sid)],
            key=lambda s: (
                self.registry.cosine_similarity(
                    task_emb, self.registry.get(s).to_state_vector()
                ) if s in self.registry.list_ids() else 0.0
            ),
            reverse=True,
        )
        results: list[SpecialistResult] = []
        for sid in sorted_sids:
            batch = self._do_call(
                sid, self.current_task.enriched_description,
                elapsed_ms, mode="PRIORITY_QUEUE",
            )
            results.extend(batch)
            for r in batch:
                if r.output and self._quick_quality_score(
                    r.output, self.current_task.enriched_description
                ) >= threshold:
                    return results
        return results

    def _exec_broadcast(
        self, action: FactoredAction, elapsed_ms: float
    ) -> list[SpecialistResult]:
        """Call all specialists independently, return only the single best result.
        Trades extra API calls for a quality ceiling guarantee."""
        results = self._exec_parallel(action, elapsed_ms)
        successful = [
            r for r in results
            if r.status == SpecialistStatus.SUCCESS and r.output
        ]
        if not successful:
            return results
        best = max(
            successful,
            key=lambda r: self._quick_quality_score(
                r.output, self.current_task.enriched_description
            ),
        )
        self.scratchpad.write(
            author_id=best.specialist_id,
            author_role=(
                self.registry.get(best.specialist_id).role
                if best.specialist_id in self.registry.list_ids() else "Specialist"
            ),
            content=f"[BROADCAST WINNER]\n{best.output}",
        )
        return [best]

    # ── MetaAction handlers ───────────────────────────────────────────

    def _exec_meta_mediator(
        self, action: FactoredAction, elapsed_ms: float
    ) -> list[SpecialistResult]:
        """Synthesise all current specialist_results to resolve conflicts.
        Only meaningful after ≥2 specialist outputs exist this episode."""
        outputs = [
            r.output for r in self.specialist_results
            if r.status == SpecialistStatus.SUCCESS and r.output
        ]
        if len(outputs) < 2:
            return []
        synthesis = self._synthesize_outputs(outputs)
        result = SpecialistResult(
            specialist_id="mediator",
            status=SpecialistStatus.SUCCESS,
            output=synthesis,
            latency_ms=0.0,
        )
        self.scratchpad.write(
            author_id="mediator", author_role="Conflict Mediator", content=synthesis
        )
        return [result]

    def _exec_meta_clarify(
        self, action: FactoredAction, elapsed_ms: float
    ) -> list[SpecialistResult]:
        """Enrich the current task description (via LLM when key available).
        All future specialist calls in this episode see the richer description."""
        import os
        original = self.current_task.enriched_description
        if os.getenv("OPENAI_API_KEY"):
            try:
                from openai import OpenAI
                client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
                resp = client.chat.completions.create(
                    model="gpt-4o-mini", max_tokens=250,
                    messages=[
                        {"role": "system", "content":
                            "Expand this task into a more specific, actionable description. "
                            "Add missing technical context. Keep it under 3 sentences."},
                        {"role": "user", "content": original[:500]},
                    ],
                )
                clarified = resp.choices[0].message.content.strip()
            except Exception as exc:
                print(f"[ClarifyTask] {exc}")
                clarified = original + " [Clarified: requires structured design and domain-specific approach]"
        else:
            clarified = (
                original + " [Clarified: requires structured design, "
                "clear acceptance criteria, and a domain-specific technical approach]"
            )
        self.current_task = type(self.current_task)(
            original_description=self.current_task.original_description,
            enriched_description=clarified,
            complexity_class=self.current_task.complexity_class,
            expected_specialists=self.current_task.expected_specialists,
            domain_hints=self.current_task.domain_hints,
            is_ambiguous=False,
            autonomously_enriched=True,
        )
        self.scratchpad.write(
            author_id="orchestrator", author_role="Orchestrator",
            content=f"Task clarified: {clarified[:300]}",
        )
        self._task_emb = self.registry.embed_query(clarified)
        return []  # effect is through improved quality on future specialist calls

    def _exec_meta_delegate_subtask(
        self, action: FactoredAction, elapsed_ms: float
    ) -> list[SpecialistResult]:
        """Decompose the task into 2–3 subtasks and route each to the best-matching
        sub-specialist, with the lead specialist as the DAG caller (depth 1→2).

        This is the only execution path that produces depth > 1 in the delegation
        graph.  The first specialist in action.specialist_ids acts as the delegating
        node; its sub-calls are recorded as specialist → sub-specialist edges so
        self.delegation_graph.depth reaches 2 when max_depth=2 permits it.
        """
        import os, json
        task = self.current_task.enriched_description

        # ── Step 1: call the lead specialist at depth 1 (orchestrator → lead) ──
        lead_id = next(
            (sid for sid in action.specialist_ids if self._can_call(sid, "orchestrator")),
            None,
        )
        results: list[SpecialistResult] = []
        if lead_id:
            results.extend(self._do_call(lead_id, task, elapsed_ms,
                                         mode="DELEGATE_SUBTASK", caller_id="orchestrator"))
        # If no lead could be called, fall through to sequential
        if not lead_id:
            return self._exec_sequential(action, elapsed_ms)

        # ── Step 2: decompose into subtasks ──
        subtasks: list[str] = []
        if os.getenv("OPENAI_API_KEY"):
            try:
                from openai import OpenAI
                client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
                resp = client.chat.completions.create(
                    model="gpt-4o-mini", max_tokens=250,
                    response_format={"type": "json_object"},
                    messages=[
                        {"role": "system", "content":
                            "Break this task into 2-3 distinct subtasks. "
                            "Return JSON: {\"subtasks\": [\"subtask1\", ...]}"},
                        {"role": "user", "content": task[:500]},
                    ],
                )
                subtasks = json.loads(resp.choices[0].message.content).get("subtasks", [])[:3]
            except Exception as exc:
                print(f"[DelegateSubtask] {exc}")
        if not subtasks:
            subtasks = [
                f"{task[:200]} — part 1: design and requirements",
                f"{task[:200]} — part 2: implementation and validation",
            ]

        # ── Step 3: route each subtask from lead_id (depth 1 → 2) ──
        for subtask in subtasks:
            sub_emb = self.registry.embed_query(subtask)
            for sid, _ in self.registry.find_most_similar(sub_emb, top_k=self.max_specialists):
                if self._can_call(sid, caller_id=lead_id):
                    results.extend(self._do_call(sid, subtask, elapsed_ms,
                                                 mode="DELEGATE_SUBTASK", caller_id=lead_id))
                    break
        return results

    def _exec_meta_retry(
        self, action: FactoredAction, elapsed_ms: float
    ) -> list[SpecialistResult]:
        """Retry all failed/timed-out specialist calls using the FallbackChainResolver."""
        failed = [r for r in self.specialist_results if r.status != SpecialistStatus.SUCCESS]
        if not failed:
            return []
        results: list[SpecialistResult] = []
        for fr in failed:
            fb_id = self.fallback_resolver.get_fallback(fr.specialist_id, self.called_ids)
            if fb_id and self._can_call(fb_id):
                batch = self._do_call(
                    fb_id, self.current_task.enriched_description,
                    elapsed_ms, mode="RETRY_FAILED",
                )
                for r in batch:
                    r.fallback_used = True
                results.extend(batch)
        return results

    def _exec_meta_parallel_spawn(
        self, action: FactoredAction, elapsed_ms: float
    ) -> list[SpecialistResult]:
        """Spawn all selected specialists in parallel (delegates to PARALLEL mode)."""
        return self._exec_parallel(action, elapsed_ms)

    # ── Roster management ─────────────────────────────────────────────

    def _select_active_specialists(self, task_emb: np.ndarray) -> list[str]:
        """
        Pick the max_specialists agents most relevant to this task.
        Always ensures any specialist spawned this episode is in the set.
        """
        ranked = self.registry.find_most_similar(
            task_emb, top_k=self.registry.size
        )
        selected = [sid for sid, _ in ranked[: self.max_specialists]]

        # Guarantee newly spawned specialists are in the active window
        for sid in self.spawned_this_episode:
            if sid not in selected:
                selected[-1] = sid  # replace least-relevant

        return selected

    def _exec_meta_spawn_specialist(
        self, action: FactoredAction, elapsed_ms: float
    ) -> list[SpecialistResult]:
        """
        Policy-triggered specialist spawn.
        Guards: OPENAI_API_KEY required, cooldown and total cap enforced.
        After a successful spawn the active roster and action decoder are
        refreshed so the new specialist is immediately selectable.
        """
        import os
        task_desc = self.current_task.enriched_description

        # Guard: no API key
        if not os.getenv("OPENAI_API_KEY"):
            return []

        # Guard: total cap
        if self._spawn_total_count >= self._spawn_max_total:
            return []

        # Guard: cooldown
        episodes_since_last = self._episode_index - self._last_spawn_episode
        if episodes_since_last < self._spawn_cooldown_episodes:
            return []

        # All guards passed — attempt spawn
        prev_count = self._spawn_total_count
        top1 = self.registry.find_most_similar(self._task_emb, top_k=1)
        best_id  = top1[0][0] if top1 else ""
        best_sim = top1[0][1] if top1 else 0.0
        self._spawn_via_llm(task_desc, best_sim=best_sim, best_id=best_id)

        if self._spawn_total_count > prev_count:
            new_id = self.spawned_this_episode[-1]
            # Refresh active roster so the new specialist is immediately reachable
            self.active_specialist_ids = self._select_active_specialists(self._task_emb)
            self.action_decoder = ActionDecoder(
                specialist_ids=self.active_specialist_ids,
                max_specialists=self.max_specialists,
            )
            return [SpecialistResult(
                specialist_id=new_id,
                status=SpecialistStatus.SUCCESS,
                output=f"[SpawnSpecialist] Spawned '{new_id}' successfully.",
                latency_ms=0.0,
            )]
        else:
            return [SpecialistResult(
                specialist_id="spawn_attempt",
                status=SpecialistStatus.ERROR,
                output="[SpawnSpecialist] LLM spawn failed — see logs.",
                latency_ms=0.0,
            )]

    def _maybe_spawn_specialist(
        self, task_emb: np.ndarray, task: str
    ) -> None:
        """
        Auto-spawn a new specialist via LLM when the best existing match
        falls below spawn_threshold.  Skipped when no OPENAI_API_KEY.
        """
        top1 = self.registry.find_most_similar(task_emb, top_k=1)
        if not top1:
            return
        best_id, best_sim = top1[0]
        if best_sim >= self.spawn_threshold:
            return  # roster already covers the task well enough
        self._spawn_via_llm(task, best_sim, best_id)

    def _spawn_via_llm(
        self, task: str, best_sim: float, best_id: str
    ) -> None:
        """
        Ask GPT-4o-mini to design a new specialist for this task,
        then add it to the registry so it enters the active roster.
        Conditions the prompt on past successful spawns for similar tasks.
        """
        import os, json
        existing_roles = [self.registry.get(s).role for s in self.registry.list_ids()]
        best_role = self.registry.get(best_id).role if best_id else "none"

        # Retrieve similar past successful spawns for RAG context
        min_reward = self.config.get("environment", {}).get("spawn_memory_min_reward", 0.0)
        past_spawns = self._spawn_memory.retrieve_similar(
            self._task_emb, top_k=3, min_reward=min_reward
        )
        past_context = ""
        if past_spawns:
            examples = "\n".join(
                f"- Role: {r.specialist_role}  |  "
                f"Desc: {r.specialist_desc[:150]}  |  "
                f"Reward: {r.episode_reward:.2f}"
                for r in past_spawns
            )
            past_context = (
                f"\n\nPast successful spawns for similar tasks:\n{examples}\n"
                "Use these as inspiration but create something distinct if needed."
            )

        try:
            from openai import OpenAI
            client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
            resp = client.chat.completions.create(
                model="gpt-4o-mini",
                max_tokens=350,
                response_format={"type": "json_object"},
                messages=[
                    {
                        "role": "system",
                        "content": (
                            "You design specialist agent definitions for a multi-agent "
                            "delegation system. Return valid JSON only."
                        ),
                    },
                    {
                        "role": "user",
                        "content": (
                            f"Task: {task[:400]}\n\n"
                            f"Existing specialists: {', '.join(existing_roles)}\n"
                            f"Best current match: {best_role} "
                            f"(cosine similarity {best_sim:.2f} — below threshold)."
                            f"{past_context}\n\n"
                            "Define a new specialist better suited to this task. "
                            "Return JSON with keys: id (snake_case), role (title case), "
                            "description (2–3 sentences of domain expertise), "
                            "complexity_affinity (list from [atomic,simple,moderate,complex,enterprise]), "
                            "avg_latency_ms (integer, 2000–8000)."
                        ),
                    },
                ],
            )
            data = json.loads(resp.choices[0].message.content)
            required = {"id", "role", "description", "complexity_affinity", "avg_latency_ms"}
            if not required.issubset(data):
                print(f"[SpawnSpecialist] Incomplete JSON: {data}")
                return
            # Deduplicate ID
            base_id = str(data["id"]).lower().replace(" ", "_")
            uid = base_id
            suffix = 2
            while uid in self.registry.list_ids():
                uid = f"{base_id}_v{suffix}"
                suffix += 1
            data["id"] = uid
            self.registry.add_specialist(data)
            self.spawned_this_episode.append(uid)
            self._spawn_total_count += 1
            self._last_spawn_episode = self._episode_index
            print(
                f"[SpawnSpecialist] Created '{data['role']}' (id={uid}) "
                f"for task (best_sim was {best_sim:.2f}, "
                f"total spawned={self._spawn_total_count}/{self._spawn_max_total})"
            )
            # Stage a pending spawn record — reward filled in at episode end
            self._pending_spawn_records.append(SpawnRecord(
                task_embedding=self._task_emb.tolist(),
                task_description=task,
                specialist_id=uid,
                specialist_role=data["role"],
                specialist_desc=data["description"],
                episode_reward=0.0,    # filled in at episode end
                pre_spawn_sim=best_sim,
                post_spawn_sim=0.0,    # filled after re-ranking
                episode_idx=self._episode_index,
            ))
        except Exception as exc:
            print(f"[SpawnSpecialist] Failed: {exc}")

    # ── Specialist execution ───────────────────────────────────────────

    def _call_specialist(
        self, specialist_id: str, task: str, elapsed_ms: float,
        context: str | None = None,
    ) -> SpecialistResult:
        """
        Call a specialist.
        Priority order:
          1. use_real_spindleflow=True  → TypeScript SpindleFlow subprocess
          2. OPENAI_API_KEY set         → real OpenAI call per specialist
          3. neither                    → fast simulation (training / offline)

        context: optional accumulated output from prior specialists (SEQUENTIAL/ITERATIVE).
        """
        import os
        specialist = self.registry.get(specialist_id)

        if self.use_real_spindleflow:
            output, latency, status = self._call_real_spindleflow(specialist_id, task)
        elif os.getenv("OPENAI_API_KEY") and not self.simulate_specialists:
            output, latency, status = self._call_openai_specialist(specialist_id, task, context=context)
        else:
            output  = self._simulate_specialist_output(specialist_id, task, context=context)
            latency = specialist.avg_latency_ms + np.random.normal(0, 500)
            status  = SpecialistStatus.SUCCESS

        return SpecialistResult(
            specialist_id=specialist_id,
            status=status,
            output=output,
            latency_ms=max(0, latency),
        )

    def _call_openai_specialist(
        self, specialist_id: str, task: str,
        context: str | None = None,
    ) -> tuple[str, float, SpecialistStatus]:
        """Call GPT-4o-mini acting as this specialist. Each gets its own system prompt.

        context: prior specialist output (SEQUENTIAL/ITERATIVE). When present, injected
                 as a user/assistant exchange before the current task so the model builds
                 on accumulated analysis rather than starting fresh.
        """
        import os
        specialist = self.registry.get(specialist_id)
        start = time.time()
        try:
            from openai import OpenAI
            client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
            if specialist.system_prompt:
                system_content = specialist.system_prompt
            else:
                system_content = (
                    f"You are a {specialist.role}. {specialist.description} "
                    f"Give a focused, expert response relevant to your specialty."
                )
            messages = [{"role": "system", "content": system_content}]
            if context:
                messages.append({
                    "role": "user",
                    "content": f"Prior specialist analysis:\n{context[:600]}",
                })
                messages.append({
                    "role": "assistant",
                    "content": "Understood. I'll build on this prior analysis.",
                })
            messages.append({"role": "user", "content": f"Task: {task[:600]}"})
            response = client.chat.completions.create(
                model="gpt-4o-mini",
                max_tokens=600,
                messages=messages,
            )
            latency = (time.time() - start) * 1000
            return response.choices[0].message.content, latency, SpecialistStatus.SUCCESS
        except Exception as exc:
            latency = (time.time() - start) * 1000
            print(f"[OpenAI specialist {specialist_id}] Error: {exc}")
            return "", latency, SpecialistStatus.ERROR

    def _simulate_specialist_output(
        self, specialist_id: str, task: str,
        context: str | None = None,
    ) -> str:
        """
        Simulate specialist output for training (no API key).

        Critically: the task text is NOT embedded in the output.
        Output quality is driven entirely by domain vocabulary from the
        specialist description, which naturally correlates with the task
        embedding when the specialist is a good match. This gives T1
        quality_delta a real signal (specialist–task domain overlap)
        rather than the degenerate case where both sides quote task[:100]
        and collapse quality_delta to noise.

        context: prior specialist output (SEQUENTIAL/ITERATIVE). When present and
                 similarity is high, the output acknowledges and extends prior work.

        Three quality tiers based on specialist-task cosine similarity:
          > 0.45  → rich domain analysis (high T1 score if relevant)
          > 0.25  → partial domain guidance
          ≤ 0.25  → mismatched — minimal domain content (low T1 score)
        """
        specialist = self.registry.get(specialist_id)
        task_emb = self.registry.embed_query(task)
        spec_emb = specialist.to_state_vector()
        similarity = self.registry.cosine_similarity(task_emb, spec_emb)

        context_prefix = ""
        if context and similarity > 0.45:
            context_prefix = (
                f"Building on the prior analysis, I will extend with {specialist.role.lower()} "
                f"expertise.\n"
            )

        if similarity > 0.45:
            return (
                f"{context_prefix}As a {specialist.role}, here is my expert analysis.\n"
                f"{specialist.description}\n"
                f"Key technical considerations from this domain: systematic design, "
                f"stakeholder alignment, iterative validation, and rigorous testing. "
                f"I recommend applying established {specialist.role.lower()} frameworks "
                f"with particular attention to quality gates and domain-specific constraints."
            )
        elif similarity > 0.25:
            return (
                f"As a {specialist.role}, I can provide partial guidance. "
                f"My expertise: {specialist.description[:200]}. "
                f"For aspects outside my specialty, additional expert input is recommended."
            )
        else:
            return (
                f"As a {specialist.role}, this request falls largely outside my primary domain. "
                f"I can offer only general guidance and recommend a more suitable specialist."
            )

    def _call_real_spindleflow(
        self, specialist_id: str, task: str
    ) -> tuple[str, float, SpecialistStatus]:
        """
        Call the real SpindleFlow TypeScript backend via subprocess.
        Returns (output, latency_ms, status).
        """
        import subprocess
        import json
        import os
        import tempfile

        spindleflow_path = os.getenv("SPINDLEFLOW_PATH", "../SpindleFlow")
        specialist = self.registry.get(specialist_id)

        config = {
            "models": {
                "gemini": {
                    "provider": "gemini",
                    "model": "gemini-2.5-flash-lite",
                    "max_tokens": 4096,
                }
            },
            "provider": "gemini",
            "agents": [{
                "id": specialist_id,
                "role": specialist.role,
                "goal": specialist.description,
            }],
            "workflow": {
                "type": "sequential",
                "steps": [{"agent": specialist_id}],
            },
        }

        with tempfile.NamedTemporaryFile(
            mode="w", suffix=".yml", delete=False
        ) as f:
            yaml.dump(config, f)
            config_path = f.name

        start = time.time()
        try:
            result = subprocess.run(
                ["npm", "run", "dev", "--", "run", config_path, "-i", task[:500]],
                cwd=spindleflow_path,
                capture_output=True,
                text=True,
                timeout=self.specialist_timeout_ms / 1000,
            )
            latency = (time.time() - start) * 1000
            if result.returncode == 0:
                output = result.stdout[-2000:]
                return output, latency, SpecialistStatus.SUCCESS
            else:
                return "", latency, SpecialistStatus.ERROR
        except subprocess.TimeoutExpired:
            latency = (time.time() - start) * 1000
            return "", latency, SpecialistStatus.TIMEOUT
        finally:
            try:
                os.unlink(config_path)
            except Exception:
                pass

    def _generate_generalist_baseline(self, task: str) -> str:
        """
        Generate a generalist (non-specialist) response to the task.
        Uses OpenAI when OPENAI_API_KEY is set (regardless of use_real_spindleflow).
        Falls back to a simulated template when no key is available.
        """
        import os
        if getattr(self, "simulate_specialists", False) or not os.getenv("OPENAI_API_KEY"):
            return (
                "General problem-solving approach:\n"
                "1. Gather and clarify requirements\n"
                "2. Research common solution patterns\n"
                "3. Draft a high-level architecture\n"
                "4. Implement in small, testable increments\n"
                "5. Validate against acceptance criteria and deploy\n"
                "No specialist domain expertise applied."
            )
        api_key = os.getenv("OPENAI_API_KEY")
        if api_key:
            try:
                from openai import OpenAI
                client = OpenAI(api_key=api_key)
                response = client.chat.completions.create(
                    model="gpt-4o-mini",
                    max_tokens=500,
                    messages=[{"role": "user", "content": f"Please help with: {task}"}],
                )
                return response.choices[0].message.content
            except Exception as e:
                print(f"[Baseline] OpenAI error: {e}. Using simulated baseline.")
        # Simulation baseline: domain-neutral boilerplate, NO task text.
        # Must embed far from any specific task so quality_delta is positive
        # whenever a matched specialist contributes domain-relevant content.
        return (
            "General problem-solving approach:\n"
            "1. Gather and clarify requirements\n"
            "2. Research common solution patterns\n"
            "3. Draft a high-level architecture\n"
            "4. Implement in small, testable increments\n"
            "5. Validate against acceptance criteria and deploy\n"
            "No specialist domain expertise applied."
        )

    def _compute_step_reward(
        self,
        step_results: list[SpecialistResult],
        task_emb: np.ndarray,
        delegation_mode: "DelegationMode | None" = None,
        meta_action: "MetaAction | None" = None,
    ) -> float:
        """
        Per-step shaping reward for non-terminal steps.

        Base shaping:
          +0.02  per specialist whose cosine-sim with task > 0.35  (good routing)
          -0.01  per specialist below 0.20                          (mismatch)
          -0.01  per failed call

        Mode-specific adjustments (make mode choice matter before terminal reward):

          PARALLEL — specialists ran concurrently; effective wall-clock cost is
            max(latencies) not sum(latencies).  Reward the latency saving when
            ≥2 specialists ran: +0.01 * (1 - max_lat / sum_lat).
            E.g. 3 specialists × 1 s each → sum=3 s, max=1 s → saving=0.67 →
            bonus ≈ +0.0067.  Scales to zero when only one specialist runs.

          SEQUENTIAL — scratchpad-chaining means each specialist built on prior
            output.  Reward the coordination effort: +0.01 per specialist after
            the first one (they had real context to work with), capped at +0.03.

        Scale stays small vs terminal range [-1, 2] so episode quality_delta
        dominates.  Total step shaping over 10 steps tops out at ~0.25.
        """
        if not step_results or not self.current_task:
            self._last_reward_components = {"step_shaping": 0.0}
            return 0.0

        shaped = 0.0
        for result in step_results:
            if result.status != SpecialistStatus.SUCCESS:
                shaped -= 0.01
                continue
            if result.specialist_id not in self.registry.list_ids():
                continue
            spec_emb = self.registry.get(result.specialist_id).to_state_vector()
            sim = self.registry.cosine_similarity(task_emb, spec_emb)
            if sim > 0.35:
                shaped += 0.02
            elif sim < 0.20:
                shaped -= 0.01

        # Mode-specific bonus
        mode_bonus = 0.0
        successful = [r for r in step_results if r.status == SpecialistStatus.SUCCESS]
        if delegation_mode == DelegationMode.PARALLEL and len(successful) >= 2:
            latencies = [r.latency_ms for r in successful]
            sum_lat = sum(latencies)
            if sum_lat > 0:
                saving = 1.0 - max(latencies) / sum_lat
                mode_bonus = round(0.01 * saving, 4)
        elif delegation_mode == DelegationMode.SEQUENTIAL and len(successful) >= 2:
            # Each specialist after the first had chained context
            chained_count = len(successful) - 1
            mode_bonus = min(0.01 * chained_count, 0.03)

        shaped += mode_bonus

        # Spawn quality shaping — only on SPAWN_SPECIALIST steps
        spawn_bonus = 0.0
        if meta_action == MetaAction.SPAWN_SPECIALIST:
            spawn_succeeded = any(
                r.status == SpecialistStatus.SUCCESS
                and r.specialist_id in self.spawned_this_episode
                for r in step_results
            )
            if spawn_succeeded:
                new_id = self.spawned_this_episode[-1]
                try:
                    new_spec_vec = self.registry.get(new_id).to_state_vector()
                    new_sim = float(self.registry.cosine_similarity(task_emb, new_spec_vec))
                    # Reward coverage gap closed above threshold; penalise redundant spawns
                    spawn_bonus = round(0.05 * max(0.0, new_sim - self.spawn_threshold), 4)
                except Exception:
                    spawn_bonus = 0.0
            else:
                # Guard hit or LLM failed — mild penalty to discourage wasteful spawn attempts
                spawn_bonus = -0.02

        shaped += spawn_bonus
        self._last_reward_components = {
            "step_shaping": float(shaped),
            "mode_bonus": float(mode_bonus),
            "spawn_bonus": float(spawn_bonus),
        }
        return float(shaped)

    def _compute_final_reward(self, elapsed_ms: float) -> float:
        """Compute the full reward for a completed episode."""
        _zero = {k: 0.0 for k in [
            "quality_delta", "efficiency_penalty", "failure_penalty",
            "recovery_bonus", "conflict_penalty", "conflict_bonus",
            "consistency_bonus", "latency_penalty", "explanation_bonus",
        ]}
        if not self.specialist_results or not self.current_task:
            self._last_reward_components = {**_zero, "failure_penalty": -0.1}
            return -0.1

        successful_outputs = [
            r.output for r in self.specialist_results
            if r.status == SpecialistStatus.SUCCESS and r.output
        ]

        if not successful_outputs:
            self._last_reward_components = {**_zero, "failure_penalty": -0.2}
            return -0.2

        specialist_output = "\n\n".join(successful_outputs)
        task_desc = self.current_task.enriched_description

        # Delta reward — same tier for both
        specialist_score = self.reward_scorer.score(
            specialist_output, task_desc, self.tier_lock
        )
        baseline_score = self.reward_scorer.score(
            self.generalist_baseline, task_desc, self.tier_lock
        )
        quality_delta = specialist_score - baseline_score

        # Efficiency penalty
        n = len(self.called_ids)
        expected = self.current_task.expected_specialists
        efficiency_penalty = self.config_reward["efficiency_base_penalty"] * \
                             max(0, n - expected)

        # Failure signals
        failure_penalty = compute_failure_penalty(self.specialist_results)
        recovery_bonus = compute_recovery_bonus(
            self.specialist_results, episode_completed=True
        )

        # Conflict signals
        conflicts = detect_conflicts(
            self.specialist_results,
            registry=self.registry,
            contradiction_pairs=self._contradiction_pairs,
            similarity_threshold=self.config_reward.get(
                "conflict_similarity_threshold", 0.25
            ),
        )
        if conflicts:
            self.conflict_resolver.resolve_all(conflicts, self.specialist_results)
        conflict_penalty = self.config_reward["conflict_unresolved_penalty"] * \
                          len([c for c in conflicts if not c.resolved])
        conflict_bonus = self.config_reward["conflict_resolved_bonus"] * \
                        len([c for c in conflicts if c.resolved])

        # Consistency bonus
        path = self.delegation_graph.get_delegation_path()
        consistency = self.consistency_tracker.consistency_score(
            path, self.current_task.complexity_class
        )
        consistency_bonus = self.config_reward["consistency_bonus_weight"] * consistency

        # Latency penalty
        latency_penalty = compute_latency_penalty(elapsed_ms, self.latency_sla)

        # Explanation bonus
        explanation_bonus = (
            self.config_reward["explanation_bonus"]
            if self.delegation_graph.is_auditable()
            else 0.0
        )

        self.consistency_tracker.record_path(
            self.current_task.complexity_class, path
        )

        total_reward = (
            quality_delta
            - efficiency_penalty
            - failure_penalty
            + recovery_bonus
            - conflict_penalty
            + conflict_bonus
            + consistency_bonus
            - latency_penalty
            + explanation_bonus
        )

        self._last_reward_components = {
            "quality_delta":      float(quality_delta),
            "efficiency_penalty": float(-efficiency_penalty),
            "failure_penalty":    float(-failure_penalty),
            "recovery_bonus":     float(recovery_bonus),
            "conflict_penalty":   float(-conflict_penalty),
            "conflict_bonus":     float(conflict_bonus),
            "consistency_bonus":  float(consistency_bonus),
            "latency_penalty":    float(-latency_penalty),
            "explanation_bonus":  float(explanation_bonus),
        }

        total_reward_clipped = float(np.clip(total_reward, -1.0, 2.0))

        # Record conflict resolution outcomes so the bandit can learn
        self.conflict_resolver.record_episode_outcome(
            quality_delta=float(quality_delta),
            episode_idx=self._episode_index,
        )

        # Finalise pending spawn records with the actual episode reward
        if self._pending_spawn_records and self._task_emb is not None:
            top_post = self.registry.find_most_similar(self._task_emb, top_k=1)
            post_sim = top_post[0][1] if top_post else 0.0
            for rec in self._pending_spawn_records:
                rec.episode_reward = total_reward_clipped
                rec.post_spawn_sim = post_sim
                self._spawn_memory.record(rec)
            self._pending_spawn_records = []

        return total_reward_clipped

    def _record_episode_to_memory(self, episode_reward: float) -> None:
        """Record each specialist's output and the episode reward to SpecialistMemory."""
        if not self.current_task:
            return
        task_desc = self.current_task.enriched_description
        for result in self.specialist_results:
            if result.specialist_id in self.spawned_this_episode:
                continue   # skip spawn confirmation messages
            if result.status == SpecialistStatus.SUCCESS and result.output:
                self.specialist_memory.record(
                    specialist_id=result.specialist_id,
                    task=task_desc,
                    output=result.output,
                    reward=episode_reward,
                )

    def render(self) -> None:
        if self.render_mode == "human" and self.current_task:
            print(f"\n[Episode State]")
            print(f"  Task: {self.current_task.enriched_description[:80]}")
            print(f"  Step: {self.step_count}/{self.max_steps}")
            print(f"  Called: {self.called_ids}")
            print(f"  Depth: {self.delegation_graph.depth}")

    def close(self) -> None:
        pass