fix(compression): in-place compaction accounts for every row behind a merged user turn
A prompt that never got its reply leaves a durable user;user pair. The alternation repair hands the compressor one dict for both rows, and the in-place commit counted it as one original. On a reload without row ids (gateway) the merged dict matched no row, read as an unpersisted turn, and its two rows were re-sequenced after the compacted set: the old prompts landed behind the running turn, and the tail rewind reached one row into the summarized region and hid it. With row ids the absorbed row was covered but not counted, so the oldest carried original stayed compacted=1 beside its live copy. The commit now resolves a merged user dict to the run of active rows it was built from, and the tail rewind covers every row a tail dict absorbed. (cherry picked from commit ba02ecf54487ca5235b99f70e6370db66aba94f4) (cherry picked from commit 09715419ac7de957b757b5446208d05e26e75f3b)
This commit is contained in:
committed by
kshitij
parent
4789ff600a
commit
44cf71f231
@@ -3753,6 +3753,7 @@ def _commit_compaction(
|
||||
# Tail rows tagged by compress() are archived as superseded duplicates, not
|
||||
# compacted=1. Count against the FINAL list — salvage may have dropped rows.
|
||||
tail_count = sum(1 for m in compressed if id(m) in _tail_tagged_ids)
|
||||
_tail_held = messages[max(0, len(messages) - tail_count):]
|
||||
# The rewind takes the newest `tail_count` durable rows as the tail's originals, so a tail row
|
||||
# with none (this turn's user row, which the CLI and gateway persist after preflight; unflushed
|
||||
# scaffolding) would flag a summarized row superseded instead: gone from display and search.
|
||||
@@ -3766,6 +3767,15 @@ def _commit_compaction(
|
||||
1 for m in messages[max(_turn_idx, len(messages) - tail_count):]
|
||||
if isinstance(m, dict) and not m.get(_DB_PERSISTED_MARKER)
|
||||
and not isinstance(m.get("_row_id"), int))
|
||||
# An alternation repair folds a durable user;user pair into one dict, which then stands
|
||||
# for every row it absorbed: counted once, the oldest carried original stays compacted=1
|
||||
# beside its live copy and is recalled twice. Only rows still active: a live list keeps
|
||||
# the ids after an earlier compaction archived them.
|
||||
tail_count += len({
|
||||
r for m in _tail_held if isinstance(m, dict)
|
||||
for r in m.get("_absorbed_row_ids") or ()
|
||||
if type(r) is int and r > 0
|
||||
and agent._session_db.get_message_role(agent.session_id, r) is not None})
|
||||
persisted = compressed
|
||||
if verbatim_tail:
|
||||
# The kept exchanges are durable rows under the watermark, so the archive below covers
|
||||
|
||||
@@ -8,7 +8,7 @@ import json
|
||||
import logging
|
||||
import re
|
||||
import time
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
from typing import Any, Dict, List, Optional, Set, Tuple
|
||||
|
||||
from agent.context_compressor import (
|
||||
_DB_PERSISTED_MARKER as _DB_PERSISTED_MARKER_KEY, MODEL_ONLY_DISPLAY_METADATA_KEY, _is_checkpoint_item,
|
||||
@@ -994,36 +994,86 @@ class SessionMessagesMixin:
|
||||
"SELECT id FROM messages WHERE session_id = ? AND active = 1 AND role = ? AND content = ?",
|
||||
(session_id, message.get("role"), stored)).fetchall()]
|
||||
|
||||
def _merged_user_run(self, conn, session_id: str, message: Dict[str, Any]) -> Optional[List[int]]:
|
||||
"""Active rows an alternation repair merged into *message*: ``[]`` for none, None when ambiguous.
|
||||
|
||||
A reload without row ids turns a durable ``user;user`` pair (a prompt that never got its reply)
|
||||
into one dict equal to neither row. Read as an unpersisted turn, its rows would be re-sequenced
|
||||
after the compacted set like concurrent appends, behind the turn that is running.
|
||||
"""
|
||||
content = message.get("content")
|
||||
if message.get("role") != "user" or not isinstance(content, str) or "\n\n" not in content:
|
||||
return []
|
||||
rows = [
|
||||
(int(row["id"]), row["role"], self._loaded_view_content(row["role"], self._decode_content(row["content"])))
|
||||
for row in conn.execute(
|
||||
"SELECT id, role, content FROM messages WHERE session_id = ? AND active = 1 ORDER BY id",
|
||||
(session_id,)).fetchall()]
|
||||
runs: List[List[int]] = []
|
||||
for start in range(len(rows)):
|
||||
merged = ""
|
||||
for end in range(start, len(rows)):
|
||||
row_id, role, part = rows[end]
|
||||
if role != "user" or not isinstance(part, str):
|
||||
break
|
||||
merged = f"{merged}\n\n{part}" if merged and part else (merged or part)
|
||||
if not content.startswith(merged):
|
||||
break
|
||||
if end > start and merged == content:
|
||||
runs.append([rows[i][0] for i in range(start, end + 1)])
|
||||
break
|
||||
if len(runs) > 1:
|
||||
return None
|
||||
return runs[0] if runs else []
|
||||
|
||||
def _proved_coverage(
|
||||
self, conn, session_id: str, covered_ids: Optional[List[int]],
|
||||
unresolved_held: Optional[List[Dict[str, Any]]],
|
||||
) -> Optional[List[int]]:
|
||||
"""Ids safe to archive as summarized, or None when a durable held row cannot be named.
|
||||
) -> Optional[Tuple[List[int], Set[int]]]:
|
||||
"""``(ids safe to archive as summarized, ids merged into another held dict)``, or None when
|
||||
a durable held row cannot be named.
|
||||
|
||||
An unresolved dict that still carries the persist marker was loaded from the DB.
|
||||
Failing to name it means the watermark path, which archives the rows the compressor
|
||||
saw, including ones whose ids were stripped. A marker-less miss is an unpersisted
|
||||
turn: it names nothing, and it is not a reason to abandon the ids we do have.
|
||||
Several active rows with the same content are ambiguous, so that also abandons.
|
||||
The second set holds only merges the caller could not count: a dict that lists its
|
||||
own ``_absorbed_row_ids`` is already counted in ``tail_count``.
|
||||
"""
|
||||
if covered_ids is None:
|
||||
return None
|
||||
from agent.context_compressor import _DB_PERSISTED_MARKER
|
||||
|
||||
proved = [int(row_id) for row_id in covered_ids if isinstance(row_id, int) and row_id > 0]
|
||||
merged_away: Set[int] = set()
|
||||
for message in unresolved_held or ():
|
||||
if not isinstance(message, dict):
|
||||
continue
|
||||
matches = self._matching_active_ids(conn, session_id, message)
|
||||
if len(matches) > 1 or (message.get(_DB_PERSISTED_MARKER) and len(matches) != 1):
|
||||
run = [] if matches else self._merged_user_run(conn, session_id, message)
|
||||
if run is None or len(matches) > 1 or (
|
||||
message.get(_DB_PERSISTED_MARKER) and len(matches) != 1 and not run):
|
||||
return None
|
||||
proved.extend(matches)
|
||||
return list(dict.fromkeys(proved))
|
||||
proved.extend(matches or run)
|
||||
merged_away.update(set(run[1:]) - set(message.get("_absorbed_row_ids") or ()))
|
||||
return list(dict.fromkeys(proved)), merged_away
|
||||
|
||||
@staticmethod
|
||||
def _tail_originals(covered_active: List[int], tail_count: int, merged_away: Set[int]) -> List[int]:
|
||||
"""Newest rows behind *tail_count* carried dicts; a dict that merged rows stands for each of them."""
|
||||
width = int(tail_count)
|
||||
while True:
|
||||
window = covered_active[-width:]
|
||||
need = int(tail_count) + sum(1 for row_id in window if row_id in merged_away)
|
||||
if need <= width or width >= len(covered_active):
|
||||
return window
|
||||
width = need
|
||||
|
||||
def _archive_named_rows(
|
||||
self, conn, session_id: str, compacted_messages: List[Dict[str, Any]], covered: List[int], *,
|
||||
tail_count: int, carried_messages: Optional[List[Dict[str, Any]]], patched_model_config: Any,
|
||||
patch: bool,
|
||||
patch: bool, merged_away: Set[int],
|
||||
) -> int:
|
||||
"""Archive *covered* as summarized and clone every other active row after the new set.
|
||||
|
||||
@@ -1039,7 +1089,7 @@ class SessionMessagesMixin:
|
||||
covered_active = [row_id for row_id in active_ids if row_id in covered_set]
|
||||
rewind_ids = list(carried_ids)
|
||||
if tail_count > 0:
|
||||
rewind_ids += covered_active[-int(tail_count):]
|
||||
rewind_ids += self._tail_originals(covered_active, tail_count, merged_away)
|
||||
rewind_ids += unseen
|
||||
rewind_ids = list(dict.fromkeys(rewind_ids))
|
||||
if rewind_ids:
|
||||
@@ -1116,8 +1166,9 @@ class SessionMessagesMixin:
|
||||
proved = self._proved_coverage(conn, session_id, covered_ids, unresolved_held)
|
||||
if proved is not None:
|
||||
return self._archive_named_rows(
|
||||
conn, session_id, compacted_messages, proved, tail_count=tail_count,
|
||||
carried_messages=carried_messages, patched_model_config=patched_model_config, patch=patch)
|
||||
conn, session_id, compacted_messages, proved[0], tail_count=tail_count,
|
||||
carried_messages=carried_messages, patched_model_config=patched_model_config, patch=patch,
|
||||
merged_away=proved[1])
|
||||
tail_ids, tail_tool_calls = ([], 0) if watermark is None else self._tail_rows_after_watermark(
|
||||
conn, "SELECT id, tool_calls FROM messages WHERE session_id = ? AND active = 1 AND id > ? ORDER BY id",
|
||||
(session_id, int(watermark)))
|
||||
|
||||
37
tests/agent/test_in_place_merged_user_rows.py
Normal file
37
tests/agent/test_in_place_merged_user_rows.py
Normal file
@@ -0,0 +1,37 @@
|
||||
"""In-place compaction over a prompt that never got its reply.
|
||||
|
||||
A crash between turn start and the reply leaves a durable ``user;user`` pair. The alternation repair
|
||||
hands the compressor one dict for both rows, so the commit has to account for two originals behind it.
|
||||
"""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.agent.test_in_place_preflight_rewind import _replies_displayed, _turn, session # noqa: F401
|
||||
|
||||
UNANSWERED = "U12x this prompt never got a reply"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("surface", ["cli", "gateway"])
|
||||
def test_compaction_over_an_unanswered_prompt_keeps_one_copy_of_every_row(session, surface):
|
||||
db, agent = session
|
||||
cli = SimpleNamespace(conversation_history=[])
|
||||
for n in range(1, 13):
|
||||
_turn(db, agent, cli, surface, n, 5_000)
|
||||
db.append_message("sid", "user", UNANSWERED)
|
||||
if surface == "cli": # a resumed CLI session starts from the repaired reload
|
||||
cli.conversation_history = db.get_messages_as_conversation("sid", repair_alternation=True)
|
||||
_turn(db, agent, cli, surface, 13, 5_000)
|
||||
_turn(db, agent, cli, surface, 14, 200_000) # real usage over the threshold: the next turn compacts first
|
||||
|
||||
_turn(db, agent, cli, surface, 15, 20_000)
|
||||
|
||||
assert getattr(agent, "_last_compaction_in_place", None) is True
|
||||
assert {f"A{n}" for n in range(1, 16)} <= _replies_displayed(db)
|
||||
live = [m["content"] for m in db.get_messages_as_conversation("sid") if isinstance(m.get("content"), str)]
|
||||
assert [c for c in live if UNANSWERED in c] == [f"{UNANSWERED}\n\nU13 please continue with the next step"]
|
||||
recalled = [row["content"] for row in db._conn.execute(
|
||||
"SELECT content FROM messages WHERE session_id = 'sid' AND (active = 1 OR compacted = 1)").fetchall()]
|
||||
carried = live[next(i for i, c in enumerate(live) if "Numbered steps" in c) + 1:]
|
||||
assert [recalled.count(content) for content in carried] == [1] * len(carried)
|
||||
Reference in New Issue
Block a user