Verify session closure in controlled trace materialization

scripts/collectivespec/create_controlled_greedy_window.py

@@ -45,6 +45,11 @@ def numeric(value: Any, *, field: str, row_index: int) -> float:
     return value
 
 
+def stable_session_key(value: Any) -> str:
+    """Render a session identifier without conflating JSON types."""
+    return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
+
+
 def main() -> int:
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument("--windows-path", type=Path, required=True)
@@ -80,12 +85,15 @@ def main() -> int:
     source_trace = resolve_trace_path(args.windows_path, trace_file)
     if not source_trace.is_file():
         raise SystemExit(f"source trace does not exist: {source_trace}")
+    source_sampling_strategy = source_window.get("sampling_strategy")
+    session_coherent_source = source_sampling_strategy == "session_coherent_uniform_score"
 
     args.output_dir.mkdir(parents=True, exist_ok=True)
     output_trace = args.output_dir / f"{args.output_window_id}.jsonl"
     selected = 0
     sampling_values: list[float] = []
     timestamps: list[float] = []
+    session_counts: dict[str, dict[str, int]] = {}
     with source_trace.open("r", encoding="utf-8") as src, output_trace.open(
         "w", encoding="utf-8"
     ) as dest:
@@ -96,6 +104,17 @@ def main() -> int:
             if not isinstance(row, dict):
                 raise SystemExit(f"row {row_index}: expected an object")
             sampling_u = numeric(row.get("sampling_u"), field="sampling_u", row_index=row_index)
+            session_entry: dict[str, int] | None = None
+            if session_coherent_source:
+                if "session_root" not in row:
+                    raise SystemExit(
+                        f"row {row_index}: session-coherent source has no session_root"
+                    )
+                session_key = stable_session_key(row["session_root"])
+                session_entry = session_counts.setdefault(
+                    session_key, {"total_turns": 0, "selected_turns": 0}
+                )
+                session_entry["total_turns"] += 1
             if sampling_u > threshold:
                 continue
             timestamp = numeric(row.get("timestamp"), field="timestamp", row_index=row_index)
@@ -104,9 +123,27 @@ def main() -> int:
             selected += 1
             sampling_values.append(sampling_u)
             timestamps.append(timestamp)
+            if session_entry is not None:
+                session_entry["selected_turns"] += 1
     if selected == 0:
         output_trace.unlink(missing_ok=True)
         raise SystemExit("threshold selected zero requests")
+    split_sessions = [
+        entry
+        for entry in session_counts.values()
+        if 0 < entry["selected_turns"] < entry["total_turns"]
+    ]
+    if split_sessions:
+        output_trace.unlink(missing_ok=True)
+        raise SystemExit(
+            "session-coherent source produced a partial selected session; "
+            f"split_session_count={len(split_sessions)}"
+        )
+    selected_session_turn_counts = [
+        entry["selected_turns"]
+        for entry in session_counts.values()
+        if entry["selected_turns"] > 0
+    ]
 
     output_window = dict(source_window)
     output_window.update(
@@ -114,12 +151,18 @@ def main() -> int:
             "window_id": args.output_window_id,
             "trace_file": str(output_trace.resolve()),
             "num_requests": selected,
-            "sampling_strategy": "fixed_uniform_score_then_greedy_temperature",
+            "sampling_strategy": (
+                "session_coherent_uniform_score_then_greedy_temperature"
+                if session_coherent_source
+                else "fixed_uniform_score_then_greedy_temperature"
+            ),
             "collectivespec_source_window_id": args.window_id,
             "collectivespec_sampling_u_threshold": threshold,
             "collectivespec_temperature_override": 0.0,
             "collectivespec_source_trace_path": str(source_trace),
             "collectivespec_source_trace_sha256": sha256(source_trace),
+            "collectivespec_source_sampling_strategy": source_sampling_strategy,
+            "collectivespec_session_closed_within_window": session_coherent_source,
         }
     )
     output_windows = {
@@ -153,6 +196,24 @@ def main() -> int:
             ),
         },
         "temperature": {"distinct_value_count": 1, "all_zero": True},
+        "session": {
+            "source_is_session_coherent": session_coherent_source,
+            "session_root_count": len(session_counts) if session_coherent_source else None,
+            "selected_session_count": (
+                len(selected_session_turn_counts) if session_coherent_source else None
+            ),
+            "split_session_count": len(split_sessions) if session_coherent_source else None,
+            "all_selected_or_dropped": not split_sessions if session_coherent_source else None,
+            "selected_turns": {
+                "min": min(selected_session_turn_counts)
+                if selected_session_turn_counts
+                else None,
+                "max": max(selected_session_turn_counts)
+                if selected_session_turn_counts
+                else None,
+                "distinct_value_count": len(set(selected_session_turn_counts)),
+            },
+        },
         "source_trace_sha256": sha256(source_trace),
         "output_trace_sha256": sha256(output_trace),
     }