Add MK2.1 lossless share support

Author: LennartvdM

cli.py

@@ -14,7 +14,7 @@
 from calendar_gen_v2 import _load_yearly_budget
 from engines.base import ScheduleInput
 from engines.engine_mk1 import EngineMK1
-from engines.engine_mk2 import EngineMK2
+from engines.engine_mk2 import EngineMK2, EngineMK21
 from rigs.simple_rig import SimpleRig
 from rigs.workforce_rig import WorkforceRig
 
@@ -24,6 +24,8 @@ def _build_engine(engine_name: str) -> object:
         return EngineMK1()
     if engine_name == "mk2":
         return EngineMK2()
+    if engine_name == "mk2_1":
+        return EngineMK21()
     raise ValueError(f"Unknown engine '{engine_name}'")
 
 
@@ -96,7 +98,7 @@ def build_parser() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser(description="Unified Wyrd Engine CLI")
     parser.add_argument(
         "--engine",
-        choices=["mk1", "mk2"],
+        choices=["mk1", "mk2", "mk2_1"],
         required=True,
         help="Select which engine implementation to use",
     )
@@ -157,8 +159,10 @@ def main(argv: Optional[Sequence[str]] = None) -> None:
         return
 
     if args.rig in {"calendar", "workforce"}:
-        if args.engine != "mk2":
-            parser.error(f"The {args.rig} rig requires the mk2 engine")
+        if args.engine not in {"mk2", "mk2_1"}:
+            parser.error(
+                f"The {args.rig} rig requires an MK2-series engine (mk2 or mk2_1)"
+            )
         _run_workforce(args)
         return
 

engines/engine_mk2.py

@@ -38,6 +38,7 @@
 
 __all__ = [
     "EngineMK2",
+    "EngineMK21",
     "DayPlan",
     "apply_micro_jitter",
     "apply_seasonal_modifiers",
@@ -89,6 +90,7 @@ def __init__(
             Callable[[PersonProfile, date, "UniqueDay"], Optional[List[Activity]]]
         ] = None,
         validator: Optional[Callable[[Dict[str, List[Activity]]], List[object]]] = None,
+        engine_version: str = "mk2",
     ) -> None:
         self._profile_factory = {
             "office": (create_office_worker, DEFAULT_TEMPLATES),
@@ -103,12 +105,17 @@ def __init__(
             [PersonProfile, date, "UniqueDay"], Optional[List[Activity]]
         ]
         self._validator: Callable[[Dict[str, List[Activity]]], List[object]]
+        self._engine_version = engine_version or "mk2"
         self.set_friction_generator(friction_generator or generate_daily_friction)
         self.set_unique_schedule_generator(
             unique_schedule_generator or generate_unique_day_schedule
         )
         self.set_validator(validator or validate_week)
 
+    @property
+    def engine_version(self) -> str:
+        return self._engine_version
+
     def set_calendar_provider(self, provider: CalendarProvider) -> None:
         """Replace the calendar provider used by the engine."""
 
@@ -689,18 +696,21 @@ def generate_complete_week(
             if debug_trace is not None:
                 debug_trace["weekly_totals_from_events"] = weekly_totals_minutes
 
+        metadata = {
+            "total_events": len(events_payload),
+            "issue_count": len(issues),
+            "summary_hours": summary_hours,
+            "compression": compression_metadata,
+            "day_types": {plan.date.isoformat(): plan.day_type for plan in week_plans},
+        }
+        metadata["engine_version"] = self._engine_version
+
         result: Dict[str, Any] = {
             "person": profile.name,
             "week_start": start_date.isoformat(),
             "events": events_payload,
             "issues": [asdict(issue) for issue in issues],
-            "metadata": {
-                "total_events": len(events_payload),
-                "issue_count": len(issues),
-                "summary_hours": summary_hours,
-                "compression": compression_metadata,
-                "day_types": {plan.date.isoformat(): plan.day_type for plan in week_plans},
-            },
+            "metadata": metadata,
         }
 
         if debug and debug_trace is not None:
@@ -716,6 +726,28 @@ def select_profile(self, archetype: str) -> Tuple[PersonProfile, Dict[str, Activ
         return factory(), templates
 
 
+class EngineMK21(EngineMK2):
+    """MK2.1 variant that enables lossless share aggregation."""
+
+    def __init__(
+        self,
+        calendar_provider: Optional[CalendarProvider] = None,
+        *,
+        friction_generator: Optional[Callable[[int, float, float], float]] = None,
+        unique_schedule_generator: Optional[
+            Callable[[PersonProfile, date, "UniqueDay"], Optional[List[Activity]]]
+        ] = None,
+        validator: Optional[Callable[[Dict[str, List[Activity]]], List[object]]] = None,
+    ) -> None:
+        super().__init__(
+            calendar_provider=calendar_provider,
+            friction_generator=friction_generator,
+            unique_schedule_generator=unique_schedule_generator,
+            validator=validator,
+            engine_version="mk2_1",
+        )
+
+
 def normalize_mk2_events(
     events: Iterable[Mapping[str, object]], *, week_start: Optional[date] = None
 ) -> List[Dict[str, Any]]:

engines/web_adapter.py

@@ -8,7 +8,7 @@
 
 from engines.base import ScheduleInput
 from engines.engine_mk1 import EngineMK1
-from engines.engine_mk2 import EngineMK2
+from engines.engine_mk2 import EngineMK2, EngineMK21
 from modules.unique_events import UniqueDay
 from rigs.simple_rig import SimpleRig
 from rigs.workforce_rig import WorkforceRig
@@ -175,12 +175,16 @@ def _coerce_start_date(value: Any) -> Optional[date]:
     return None
 
 
-def _ensure_schema(payload: MutableMapping[str, Any], *, rig: str, seed: int, archetype: str) -> SchemaPayload:
+def _ensure_schema(
+    payload: MutableMapping[str, Any], *, rig: str, seed: int, archetype: str, engine_version: Optional[str] = None
+) -> SchemaPayload:
     metadata = dict(payload.get("metadata", {}))
     metadata.setdefault("summary_hours", {})
+    resolved_engine = engine_version or ("mk2" if rig in {"calendar", "workforce"} else "mk1")
     metadata.update(
         {
-            "engine": "mk2" if rig in {"calendar", "workforce"} else "mk1",
+            "engine": resolved_engine,
+            "engine_version": resolved_engine,
             "rig": rig,
             "seed": seed,
             "archetype": archetype,
@@ -243,6 +247,8 @@ def _convert_events(events: Iterable[Mapping[str, Any]]) -> Iterable[Dict[str, A
 
 _MK2_ENGINE = EngineMK2()
 _MK2_RIG = WorkforceRig(engine=_MK2_ENGINE)
+_MK2_1_ENGINE = EngineMK21()
+_MK2_1_RIG = WorkforceRig(engine=_MK2_1_ENGINE)
 
 
 def mk1_run_web(archetype: str, week_start: Optional[str], seed: Any) -> SchemaPayload:
@@ -278,49 +284,34 @@ def mk1_run_web(archetype: str, week_start: Optional[str], seed: Any) -> SchemaP
         },
     }
 
-    return _ensure_schema(payload, rig="default", seed=seed_value, archetype=archetype_key or "office")
-
-
-def mk2_run_calendar_web(
-    archetype: str, week_start: Optional[str], seed: Any, debug: bool = False
-) -> SchemaPayload:
-    archetype_key = str(archetype or "office").strip().lower()
-    seed_value = _coerce_seed(seed)
-    start_date = _coerce_start_date(week_start) or date.today()
-
-    profile, templates = _MK2_RIG.select_profile(archetype_key)
-    result = _MK2_RIG.generate_complete_week(
-        profile, start_date, seed_value, templates, None, debug=debug
+    return _ensure_schema(
+        payload,
+        rig="default",
+        seed=seed_value,
+        archetype=archetype_key or "office",
+        engine_version="mk1",
     )
 
-    payload: MutableMapping[str, Any] = dict(result)
-    payload.setdefault("issues", [])
-    payload.setdefault("events", [])
-    payload.setdefault("week_start", start_date.isoformat())
-    payload["person"] = profile.name
-    payload["metadata"] = {
-        **payload.get("metadata", {}),
-        "profile": profile.name,
-    }
-
-    return _ensure_schema(payload, rig="calendar", seed=seed_value, archetype=archetype_key)
 
-
-def mk2_run_workforce_web(
+def _run_mk2_variant(
+    rig_instance: WorkforceRig,
     archetype: str,
     week_start: Optional[str],
     seed: Any,
-    yearly_budget: Optional[Mapping[str, Any]],
+    *,
+    engine_version: str,
+    rig_label: str,
+    yearly_budget: Optional[Mapping[str, Any]] = None,
     debug: bool = False,
 ) -> SchemaPayload:
     archetype_key = str(archetype or "office").strip().lower()
     seed_value = _coerce_seed(seed)
     start_date = _coerce_start_date(week_start) or date.today()
 
-    profile, templates = _MK2_RIG.select_profile(archetype_key)
-    budget = _build_yearly_budget(yearly_budget)
+    profile, templates = rig_instance.select_profile(archetype_key)
+    budget = _build_yearly_budget(yearly_budget) if yearly_budget is not None else None
 
-    result = _MK2_RIG.generate_complete_week(
+    result = rig_instance.generate_complete_week(
         profile, start_date, seed_value, templates, budget, debug=debug
     )
 
@@ -331,6 +322,7 @@ def mk2_run_workforce_web(
     payload["person"] = profile.name
     metadata = dict(payload.get("metadata", {}))
     metadata["profile"] = profile.name
+    metadata["engine_version"] = engine_version
     if budget is not None:
         metadata["yearly_budget"] = {
             "person_id": budget.person_id,
@@ -341,11 +333,87 @@ def mk2_run_workforce_web(
         }
     payload["metadata"] = metadata
 
-    return _ensure_schema(payload, rig="workforce", seed=seed_value, archetype=archetype_key)
+    return _ensure_schema(
+        payload,
+        rig=rig_label,
+        seed=seed_value,
+        archetype=archetype_key,
+        engine_version=engine_version,
+    )
+
+
+def mk2_run_calendar_web(
+    archetype: str, week_start: Optional[str], seed: Any, debug: bool = False
+) -> SchemaPayload:
+    return _run_mk2_variant(
+        _MK2_RIG,
+        archetype,
+        week_start,
+        seed,
+        engine_version="mk2",
+        rig_label="calendar",
+        yearly_budget=None,
+        debug=debug,
+    )
+
+
+def mk2_run_workforce_web(
+    archetype: str,
+    week_start: Optional[str],
+    seed: Any,
+    yearly_budget: Optional[Mapping[str, Any]],
+    debug: bool = False,
+) -> SchemaPayload:
+    return _run_mk2_variant(
+        _MK2_RIG,
+        archetype,
+        week_start,
+        seed,
+        engine_version="mk2",
+        rig_label="workforce",
+        yearly_budget=yearly_budget,
+        debug=debug,
+    )
+
+
+def mk2_1_run_calendar_web(
+    archetype: str, week_start: Optional[str], seed: Any, debug: bool = False
+) -> SchemaPayload:
+    return _run_mk2_variant(
+        _MK2_1_RIG,
+        archetype,
+        week_start,
+        seed,
+        engine_version="mk2_1",
+        rig_label="calendar",
+        yearly_budget=None,
+        debug=debug,
+    )
+
+
+def mk2_1_run_workforce_web(
+    archetype: str,
+    week_start: Optional[str],
+    seed: Any,
+    yearly_budget: Optional[Mapping[str, Any]],
+    debug: bool = False,
+) -> SchemaPayload:
+    return _run_mk2_variant(
+        _MK2_1_RIG,
+        archetype,
+        week_start,
+        seed,
+        engine_version="mk2_1",
+        rig_label="workforce",
+        yearly_budget=yearly_budget,
+        debug=debug,
+    )
 
 
 __all__ = [
     "mk1_run_web",
     "mk2_run_calendar_web",
     "mk2_run_workforce_web",
+    "mk2_1_run_calendar_web",
+    "mk2_1_run_workforce_web",
 ]

tests/test_cli_parity.py

@@ -19,6 +19,16 @@
 ROOT = Path(__file__).resolve().parents[1]
 
 
+def _override_metadata_engine(payload: dict, engine_value: str) -> dict:
+    """Return a deep copy of payload with metadata engine fields normalised."""
+
+    clone = json.loads(json.dumps(payload))
+    metadata = clone.setdefault("metadata", {})
+    if engine_value:
+        metadata["engine_version"] = engine_value
+    return clone
+
+
 @pytest.mark.parametrize(
     "cli_args, legacy_runner",
     [
@@ -48,6 +58,21 @@
             ],
             "run_calendar_gen_v2",
         ),
+        (
+            [
+                "--engine",
+                "mk2_1",
+                "--rig",
+                "workforce",
+                "--archetype",
+                "office",
+                "--seed",
+                "7",
+                "--start-date",
+                "2025-01-06",
+            ],
+            "run_calendar_gen_v2",
+        ),
     ],
 )
 def test_cli_matches_legacy(tmp_path: Path, cli_args, legacy_runner) -> None:
@@ -98,7 +123,15 @@ def run_calendar_gen_v2_legacy() -> str:
 
     cli_data = json.loads(cli_output.read_text())
     legacy_data = json.loads(legacy_output.read_text())
-    assert cli_data == legacy_data
+    engine_choice = None
+    if "--engine" in cli_args:
+        engine_index = cli_args.index("--engine")
+        if engine_index + 1 < len(cli_args):
+            engine_choice = cli_args[engine_index + 1]
+    comparison_target = legacy_data
+    if engine_choice == "mk2_1":
+        comparison_target = _override_metadata_engine(legacy_data, engine_choice)
+    assert cli_data == comparison_target
 
     def normalise_output(text: str, output_path: Path) -> str:
         return text.replace(str(output_path), "<OUTPUT>")