diff --git a/src/lingtai/kernel/agent_session.py b/src/lingtai/kernel/agent_session.py index d63bb9469..f671d263c 100644 --- a/src/lingtai/kernel/agent_session.py +++ b/src/lingtai/kernel/agent_session.py @@ -244,9 +244,24 @@ def _rebuild_via_sqlite( boundary_source = "boot" if boundary_rows: row = boundary_rows[0] + try: + fields = json.loads(row.get("fields_json")) + except (json.JSONDecodeError, TypeError, ValueError): + return None + boundary_molt_count = ( + fields.get("molt_count") if isinstance(fields, dict) else None + ) + if ( + isinstance(boundary_molt_count, bool) + or not isinstance(boundary_molt_count, int) + or boundary_molt_count != molt_count + ): + return None boundary_ts = float(row.get("ts") or 0.0) boundary_offset = row.get("source_offset") - boundary_source = _boundary_source_from_fields(row.get("fields_json")) + boundary_source = str(fields.get("initiator") or "agent") + elif molt_count != 0: + return None try: agg = query( @@ -430,18 +445,6 @@ def _read_tail_lines(path: Path, max_bytes: int) -> tuple[list[str], bool]: return lines, hit_start -def _boundary_source_from_fields(fields_json: Any) -> str: - if not fields_json: - return "agent" - try: - fields = json.loads(fields_json) if isinstance(fields_json, str) else fields_json - except (json.JSONDecodeError, ValueError): - return "agent" - if isinstance(fields, dict): - return str(fields.get("initiator") or "agent") - return "agent" - - def _iso_from_ts(ts: float) -> str | None: if not ts: return None diff --git a/tests/test_agent_session.py b/tests/test_agent_session.py index c70627379..febc3943b 100644 --- a/tests/test_agent_session.py +++ b/tests/test_agent_session.py @@ -20,6 +20,7 @@ RuntimeSession, _rebuild_via_full_scan, _rebuild_via_reverse_scan, + _rebuild_via_sqlite, new_runtime_session, rebuild_agent_session_from_events, ) @@ -224,6 +225,78 @@ def test_all_three_tiers_agree(tmp_path): assert t1.cached_tokens == 1600 +@pytest.mark.parametrize( + "fields_json", + [ + json.dumps({"molt_count": 1}), + json.dumps({"initiator": "agent"}), + "{malformed", + ], + ids=["stale", "missing", "malformed"], +) +def test_sqlite_rejects_invalid_molt_boundary(fields_json): + def query(sql: str) -> list[dict]: + if "psyche_molt" in sql: + return [{"ts": 2.0, "source_offset": 10, "fields_json": fields_json}] + pytest.fail("invalid boundary must be rejected before aggregation") + + assert _rebuild_via_sqlite(query, molt_count=2) is None + + +def test_sqlite_accepts_matching_molt_boundary(): + def query(sql: str) -> list[dict]: + if "psyche_molt" in sql: + return [ + { + "ts": 2.0, + "source_offset": 10, + "fields_json": json.dumps( + {"molt_count": 2, "initiator": "system"} + ), + } + ] + return [ + { + "n": 1, + "input_tokens": 1000, + "output_tokens": 200, + "cached_tokens": 700, + "thinking_tokens": 10, + } + ] + + session = _rebuild_via_sqlite(query, molt_count=2) + + assert session is not None + assert session.rebuild_tier == "sqlite" + assert session.boundary_source == "system" + assert session.input_tokens == 1000 + + +def test_sqlite_rejects_nonzero_molt_without_boundary(): + def query(sql: str) -> list[dict]: + if "psyche_molt" in sql: + return [] + pytest.fail("missing boundary must be rejected before aggregation") + + assert _rebuild_via_sqlite(query, molt_count=2) is None + + +def test_stale_sqlite_boundary_falls_back_to_authoritative_jsonl(tmp_path): + current_events = [ + _molt_ev(2.0, 2), + _token_ev(3.0, 1000, 200, 700), + ] + _write_events(tmp_path, current_events) + _build_sqlite(tmp_path, [_molt_ev(1.0, 1)]) + + session = rebuild_agent_session_from_events(tmp_path, molt_count=2) + + assert session.rebuild_tier == "reverse_scan" + assert session.api_calls == 1 + assert session.input_tokens == 1000 + + def test_public_entry_falls_back_to_reverse_scan_without_sidecar(tmp_path): events = [ _molt_ev(1.0, 1),