feat(selection): penalise uneven context spread across agents

config.py

@@ -94,6 +94,10 @@ def __init__(self):
         # Length penalty: genotype size at which the penalty starts to grow
         self.length_penalty_baseline_chars: int = 8000
         self.length_penalty_lambda: float = 0.05
+        # Context-dispersity penalty: strength of the term subtracted from
+        # qd_score when a workflow concentrates context on one agent instead
+        # of spreading it. 0.0 disables it, which is the default.
+        self.context_dispersity_lambda: float = 0.0
         # Selection pressure / archive settings
         self.min_improvement_threshold: float = 0.01
         self.population_size: int = 20
@@ -295,6 +299,7 @@ def jsonify(
             "novelty_previous_n": self.novelty_previous_n,
             "length_penalty_baseline_chars": self.length_penalty_baseline_chars,
             "length_penalty_lambda": self.length_penalty_lambda,
+            "context_dispersity_lambda": self.context_dispersity_lambda,
             "min_improvement_threshold": self.min_improvement_threshold,
             "population_size": self.population_size,
             "novelty_k_neighbours": self.novelty_k_neighbours,
@@ -362,6 +367,9 @@ def from_json(self, data: dict[str, Any]) -> None:
         self.length_penalty_lambda = float(
             data.get("length_penalty_lambda", self.length_penalty_lambda)
         )
+        self.context_dispersity_lambda = float(
+            data.get("context_dispersity_lambda", self.context_dispersity_lambda)
+        )
         self.min_improvement_threshold = float(
             data.get("min_improvement_threshold", self.min_improvement_threshold)
         )

docs/reference/configuration.md

@@ -57,6 +57,7 @@ engine](../concepts/evolution-engine.md), not day-to-day settings.
 | `novelty_previous_n` | `int` | `15` | Window size when `novelty_comparison` is previous-N based. |
 | `length_penalty_baseline_chars` | `int` | `8000` | Genotype size (chars) at which the length penalty starts to grow. |
 | `length_penalty_lambda` | `float` | `0.05` | Length-penalty growth rate. |
+| `context_dispersity_lambda` | `float` | `0.0` | Strength of the context-dispersity penalty subtracted from `qd_score`. `0.0` disables it. See the note below. |
 | `min_improvement_threshold` | `float` | `0.01` | Minimum relative improvement required for admission in greedy mode. |
 | `population_size` | `int` | `20` | Max individuals kept in the QD archive. |
 | `novelty_k_neighbours` | `int` | `15` | `k` for k-nearest-neighbour novelty. |
@@ -68,6 +69,19 @@ engine](../concepts/evolution-engine.md), not day-to-day settings.
 | `parent_threshold_similarity` | `float` | `0.8` | Cold-start disk scan: minimum goal-embedding cosine similarity. |
 | `parent_threshold_score` | `float` | `0.01` | Cold-start disk scan / parent draw: minimum score to be considered. |
 
+!!! note "`context_dispersity_lambda` penalises unevenness, not size"
+    The term measures how unevenly the final context is spread across a
+    workflow's agents: it is `0` when every agent ends with the same context
+    and `1` when a single agent holds all of it. It is scale free, so four
+    agents each ending at 900k tokens score `0` and are **not** penalised —
+    only the imbalance is. Setting it above `0` also makes the engine read
+    each agent's memory file once per iteration to recover those lengths.
+
+    A first trial value of `0.05`, matching `length_penalty_lambda`, keeps the
+    term a tie-breaker: the most it can subtract is `0.05`, well below
+    `admit_threshold`, so it cannot by itself gate a workflow out of the
+    archive.
+
 ## OpenRouter routing
 
 | Field | Type | Default | Description |

sources/core/evolution_engine.py

@@ -27,6 +27,7 @@
 )
 from sources.evaluators.evaluator import WorkflowEvaluator
 from sources.benchmark_evaluation.scenario_loader import ScenarioLoader
+from sources.utils.agent_context import read_agent_context_lengths
 from sources.utils.notify import PushNotifier
 from sources.utils.pricing import PricingCalculator
 from sources.utils.run_metrics import append_jsonl, write_run_metrics
@@ -103,6 +104,7 @@ def __init__(
             previous_n=getattr(config, "novelty_previous_n", 5),
             length_penalty_baseline_chars=getattr(config, "length_penalty_baseline_chars", 5000),
             length_penalty_lambda=getattr(config, "length_penalty_lambda", 0.05),
+            context_dispersity_lambda=getattr(config, "context_dispersity_lambda", 0.0),
         )
         self.initial_population = getattr(config, "initial_population", 2) # number of initial random workflows before enabling mutation
 
@@ -532,6 +534,11 @@ async def evolve_generation(
         runs[-1].current_uuid = uuid
         runs[-1].answers = wf_info.answers if wf_info else []
         runs[-1].state_result = wf_info.state_result if wf_info else {}
+        # Only read the agent memory when the penalty is active: recovering the
+        # lengths costs a full parse of every agent memory file, and those files
+        # are largest for exactly the high-context runs this term targets.
+        if self.selection.dispersity_lambda > 0:
+            runs[-1].agent_context_lengths = read_agent_context_lengths(self.config.memory_dir, uuid)
         agents_answers = self.extract_agents_behavior(wf_info.state_result) if wf_info else ""
         self.show_answers(agents_answers)
 

sources/core/schema.py

@@ -70,6 +70,10 @@ class IndividualRun:
     # persisted via sources.core.lineage.record_lineage afterwards.
     parent_uuids: list[str] = field(default_factory=list)
     evolution_kind: str = "seed"  # "seed" | "mutation" | "crossover"
+    # Final context length (input tokens) of each agent in the executed
+    # workflow. Read by SelectionPressure to penalise workflows that pile
+    # context onto a single agent instead of distributing it.
+    agent_context_lengths: list[int] = field(default_factory=list)
 
     def __str__(self) -> str:
         """Return a verbose, debug-style summary of the run."""

sources/core/selection.py

@@ -51,7 +51,10 @@ class PopulationMember:
         genotype_chars: Raw character length of the workflow source —
             used for the length-penalty term in ``qd_score``.
         novelty_score: Mean cosine distance to the current comparison set.
-        qd_score: Combined quality-diversity score with length penalty.
+        qd_score: Combined quality-diversity score, net of the length and
+            context-dispersity penalties.
+        context_dispersity: How unevenly context was spread across the
+            workflow's agents, in ``[0, 1]``.
         reward_uncapped: Base reward with no hard-fail cap.
         created_at: Wall-clock timestamp of construction.
     """
@@ -64,6 +67,7 @@ class PopulationMember:
     novelty_score: float = 0.0
     qd_score: float = 0.0
     reward_uncapped: float = 0.0
+    context_dispersity: float = 0.0
     created_at: datetime = field(default_factory=datetime.now)
 
 
@@ -95,6 +99,7 @@ def __init__(
         previous_n: int = 5,
         length_penalty_baseline_chars: int = 5000,
         length_penalty_lambda: float = 0.05,
+        context_dispersity_lambda: float = 0.0,
     ) -> None:
         """Configure thresholds and the active selection strategy.
 
@@ -116,6 +121,9 @@ def __init__(
             length_penalty_lambda: Strength of the length penalty term
                 subtracted from ``qd_score``. Kept conservative so it only
                 breaks near-ties.
+            context_dispersity_lambda: Strength of the context-dispersity term
+                subtracted from ``qd_score``, pressuring workflows to spread
+                context across agents. ``0.0`` disables it.
         """
         self.config = config
         self.logger = logging.getLogger(__name__)
@@ -137,6 +145,9 @@ def __init__(
         self.previous_n = max(1, int(previous_n))
         self.length_baseline_chars = max(1, int(length_penalty_baseline_chars))
         self.length_lambda = float(length_penalty_lambda)
+        # Pressure towards workflows that spread context across agents rather
+        # than piling it onto one. Zero by default: opt in per run.
+        self.dispersity_lambda = float(context_dispersity_lambda)
 
         self._archive: list[PopulationMember] = []
         self._previous_descriptors: list[list[float]] = []
@@ -399,7 +410,8 @@ def _validate_open_ended(
         Returns:
             Validation result dict from :meth:`_build_result`, extended with
             ``novelty_score``, ``qd_score``, ``length_penalty``,
-            ``archive_size``, ``admit_rejected`` and ``admit_rejected_total``.
+            ``context_dispersity``, ``archive_size``, ``admit_rejected`` and
+            ``admit_rejected_total``.
         """
         self._last_evicted_uuid = None
 
@@ -416,8 +428,10 @@ def _validate_open_ended(
         quality_norm = min(max(new_reward_uncapped, 0.0), 1.0)
         novelty_norm = min(novelty / max(novelty_range, 1e-6), 1.0)
         length_penalty = _length_penalty(genotype_chars, self.length_baseline_chars)
+        context_dispersity = _context_dispersity(_agent_context_lengths(best_new))
         qd_score = self._compose_qd_score(
             new_reward_uncapped, novelty, genotype_chars, novelty_range,
+            context_dispersity,
         )
 
         absolute_improvement = new_reward - baseline_reward
@@ -434,6 +448,7 @@ def _validate_open_ended(
             novelty_score=novelty,
             qd_score=qd_score,
             reward_uncapped=new_reward_uncapped,
+            context_dispersity=context_dispersity,
         )
         admit_rejected = not self._try_admit(member, is_valid)
         self._record_previous(descriptor)
@@ -449,6 +464,7 @@ def _validate_open_ended(
         result["quality_norm"] = quality_norm
         result["novelty_norm"] = novelty_norm
         result["length_penalty"] = length_penalty
+        result["context_dispersity"] = context_dispersity
         result["behaviour_descriptor"] = descriptor or []
         result["archive_size"] = len(self._archive)
         result["admit_rejected"] = admit_rejected
@@ -459,7 +475,7 @@ def _validate_open_ended(
         if self.strategy in (SelectionStrategy.NOVELTY, SelectionStrategy.QUALITY_DIVERSITY):
             self.logger.info(
                 f"Open-ended: novelty={novelty:.3f}, qd={qd_score:.3f}, "
-                f"len_pen={length_penalty:.3f}, "
+                f"len_pen={length_penalty:.3f}, ctx_disp={context_dispersity:.3f}, "
                 f"archive={len(self._archive)}/{self.population_size}"
             )
         return result
@@ -604,6 +620,7 @@ def _refresh_member_metrics(self) -> None:
         for m in self._archive:
             m.qd_score = self._compose_qd_score(
                 m.reward_uncapped, m.novelty_score, m.genotype_chars, novelty_range,
+                m.context_dispersity,
             )
 
     def _knn_novelty_against_peers(self, m: PopulationMember) -> float:
@@ -622,15 +639,17 @@ def _compose_qd_score(
         novelty: float,
         genotype_chars: int,
         novelty_range: float,
+        context_dispersity: float,
     ) -> float:
-        """``(1-w)·quality + w·novelty − λ·length_penalty`` in one place."""
+        """``(1-w)·quality + w·novelty − λ_len·length_pen − λ_disp·dispersity``."""
         quality_norm = min(max(reward_uncapped, 0.0), 1.0)
         novelty_norm = min(novelty / max(novelty_range, 1e-6), 1.0)
         length_penalty = _length_penalty(genotype_chars, self.length_baseline_chars)
         return (
             (1 - self.novelty_weight) * quality_norm
             + self.novelty_weight * novelty_norm
             - self.length_lambda * length_penalty
+            - self.dispersity_lambda * context_dispersity
         )
 
     # ------------------------------------------------------------------
@@ -794,6 +813,43 @@ def _genotype_chars(run: Any) -> int:
     return len(code) if isinstance(code, str) else 0
 
 
+def _agent_context_lengths(run: Any) -> list[int]:
+    """Per-agent final context lengths recorded for the run, ``[]`` when missing."""
+    lengths = _safe_attr(run, "agent_context_lengths", None)
+    if not isinstance(lengths, list):
+        return []
+    return [n for n in lengths if isinstance(n, int) and n > 0]
+
+
+def _context_dispersity(lengths: list[int]) -> float:
+    """How unevenly context is spread across agents, in ``[0, 1]``.
+
+    The coefficient of variation is divided by ``sqrt(n - 1)``, its maximum
+    for ``n`` agents (reached when one agent holds all the context). The
+    result therefore reaches ``1`` only at total concentration and keeps a
+    monotone gradient up to it, rather than saturating early.
+
+    Evenly shared context scores ``0`` at any scale: this term measures
+    concentration, not absolute size. Fewer than two agents carries no
+    dispersion information and scores ``0``.
+
+    Args:
+        lengths: Final context length of each agent, in tokens.
+
+    Returns:
+        ``clip((stdev / mean) / sqrt(n - 1), 0, 1)``, or ``0.0`` when undefined.
+    """
+    n = len(lengths)
+    if n < 2:
+        return 0.0
+    mean = sum(lengths) / n
+    if mean <= 0:
+        return 0.0
+    variance = sum((length - mean) ** 2 for length in lengths) / n
+    coefficient_of_variation = math.sqrt(variance) / mean
+    return max(0.0, min(1.0, coefficient_of_variation / math.sqrt(n - 1)))
+
+
 if __name__ == "__main__":
     from types import SimpleNamespace
 

sources/utils/agent_context.py

@@ -0,0 +1,85 @@
+"""Read the final context length of each agent from a run's saved memory.
+
+The SmolAgent factory writes one JSON file per agent into
+``<memory_dir>/<workflow_uuid>/``, holding the ordered steps of that agent.
+Each step records its ``token_usage``; the input tokens of the last step are
+the context the agent was carrying when it finished.
+"""
+
+import json
+from pathlib import Path
+
+# Both factories save agent memory as ``task_{agent_name}.json``; the
+# single-agent run saves itself as ``task_single_agent.json``.
+_AGENT_FILE_PREFIX = "task_"
+
+
+def _final_input_tokens(steps: list) -> int:
+    """Input tokens of the last step that recorded any, ``0`` when none did.
+
+    Args:
+        steps: Ordered agent memory steps, as loaded from the memory file.
+
+    Returns:
+        The final step's ``token_usage.input_tokens``, or ``0``.
+    """
+    for step in reversed(steps):
+        if not isinstance(step, dict):
+            continue
+        usage = step.get("token_usage") or {}
+        tokens = usage.get("input_tokens")
+        if isinstance(tokens, int) and tokens > 0:
+            return tokens
+    return 0
+
+
+def read_agent_context_lengths(memory_dir: str, workflow_uuid: str) -> list[int]:
+    """Final context length of every agent that ran under ``workflow_uuid``.
+
+    Missing directories, unreadable files and agents with no recorded token
+    usage are skipped rather than raised, because this feeds a ranking term
+    and must never abort an evolution iteration.
+
+    Args:
+        memory_dir: Root directory holding per-run agent memory folders.
+        workflow_uuid: UUID naming this run's memory folder.
+
+    Returns:
+        One positive length per agent, ordered by memory filename.
+    """
+    if not memory_dir or not workflow_uuid:
+        return []
+
+    run_memory = Path(memory_dir) / workflow_uuid
+    if not run_memory.is_dir():
+        return []
+
+    lengths = []
+    for path in sorted(run_memory.glob("*.json")):
+        if not path.name.startswith(_AGENT_FILE_PREFIX):
+            continue
+        try:
+            steps = json.loads(path.read_text())
+        except Exception:
+            # A ranking term must never abort an iteration; a memory file that
+            # is missing, truncated, or too large to parse simply does not vote.
+            continue
+        if not isinstance(steps, list):
+            continue
+        tokens = _final_input_tokens(steps)
+        if tokens:
+            lengths.append(tokens)
+    return lengths
+
+
+if __name__ == "__main__":
+    import tempfile
+
+    with tempfile.TemporaryDirectory() as root:
+        run = Path(root) / "uuid-1"
+        run.mkdir()
+        (run / "task_a.json").write_text(json.dumps([{"token_usage": {"input_tokens": 10}}]))
+        (run / "task_b.json").write_text(json.dumps([{"token_usage": {"input_tokens": 90}}]))
+        assert read_agent_context_lengths(root, "uuid-1") == [10, 90]
+        assert read_agent_context_lengths(root, "missing") == []
+    print("agent_context smoke check passed")