diff --git a/dev-notes/compaction-keeps-newest-user-turn.md b/dev-notes/compaction-keeps-newest-user-turn.md new file mode 100644 index 000000000..fd1899930 --- /dev/null +++ b/dev-notes/compaction-keeps-newest-user-turn.md @@ -0,0 +1,41 @@ +# Automatic compaction always keeps the newest user turn + +## What was added + +`_first_recent_context_index` in `tau_coding.session` now clamps its compaction +boundary to the last user message, so automatic compaction can never drop the +newest user turn or its assistant reply. + +## Why it exists + +The function returns the index of the first row to **keep**; rows before it are +summarized and replaced. Its "keep whole user turns" logic moved the boundary +forward to the next user message, but nothing guaranteed that boundary stayed at +or before the **last** user message. That let compaction silently destroy the +most recent context: + +- `[user, assistant, user, assistant]` returned `3`, keeping only the final + `assistant` reply and dropping the newest user prompt. +- `[user, assistant, toolResult, toolResult]` returned `4`, compacting the entire + context including the newest turn. + +The root cause was conflating "keep recent tokens" with "keep whole user turns" +without a guard that the newest user turn is always retained. + +## Fix + +The function computes a single `boundary` value, then clamps it with +`min(boundary, last_user_index)` where `last_user_index` is the index of the last +`user` message (found by a new `_last_user_message_index` helper). When the +boundary clamps to `0`, the caller (`_recent_preserving_compaction_plan`) returns +`None` and skips compaction rather than compacting everything. + +## How to test + +```bash +uv run pytest tests/test_coding_session.py -k "first_recent_context_index or auto_compact or overflow" +``` + +The pure-function tests cover: keeping the newest user turn and reply, never +compacting past the last user message, keeping a pending (unanswered) newest user +prompt, and a normal recent-suffix case. diff --git a/src/tau_coding/session.py b/src/tau_coding/session.py index 78901defd..bc7a86360 100644 --- a/src/tau_coding/session.py +++ b/src/tau_coding/session.py @@ -2650,18 +2650,35 @@ def _first_recent_context_index( candidate_message = rows[candidate_index][1] if candidate_message.role == "user": if candidate_index > 0: - return candidate_index - next_user_index = _next_user_message_index(rows, start=1) - return next_user_index if next_user_index is not None else 0 + boundary = candidate_index + else: + next_user_index = _next_user_message_index(rows, start=1) + boundary = next_user_index if next_user_index is not None else 0 + else: + next_user_index = _next_user_message_index(rows, start=candidate_index + 1) + if next_user_index is not None: + boundary = next_user_index + else: + for index in range(candidate_index, len(rows)): + if rows[index][1].role != "toolResult": + boundary = index + break + else: + boundary = len(rows) + + last_user_index = _last_user_message_index(rows) + if last_user_index is not None: + boundary = min(boundary, last_user_index) + return boundary - next_user_index = _next_user_message_index(rows, start=candidate_index + 1) - if next_user_index is not None: - return next_user_index - for index in range(candidate_index, len(rows)): - if rows[index][1].role != "toolResult": +def _last_user_message_index( + rows: tuple[tuple[str, AgentMessage], ...], +) -> int | None: + for index in range(len(rows) - 1, -1, -1): + if rows[index][1].role == "user": return index - return len(rows) + return None def _next_user_message_index( diff --git a/tests/test_coding_session.py b/tests/test_coding_session.py index 0d99fa55b..23dae876f 100644 --- a/tests/test_coding_session.py +++ b/tests/test_coding_session.py @@ -58,10 +58,15 @@ save_provider_settings, ) from tau_coding import session as coding_session_module +from tau_coding.context_window import estimate_message_tokens from tau_coding.events import QueueUpdateEvent from tau_coding.prompt_templates import PromptTemplate from tau_coding.provider_config import ProviderModelMetadata -from tau_coding.session import _ordered_tree_entries, parse_terminal_command +from tau_coding.session import ( + _first_recent_context_index, + _ordered_tree_entries, + parse_terminal_command, +) async def _collect_session_events(session_stream: object) -> list[object]: @@ -1021,6 +1026,52 @@ def test_parse_terminal_command_prefixes() -> None: assert parse_terminal_command("hello") is None +def _recent_rows(*roles: str) -> tuple[tuple[str, AgentMessage], ...]: + def message(role: str) -> AgentMessage: + if role == "user": + return UserMessage(content=f"user-{role}") + if role == "assistant": + return AssistantMessage(content=f"assistant-{role}") + return ToolResultMessage( + content=f"result-{role}", + tool_call_id=f"call-{role}", + tool_name="tool", + ) + + return tuple((f"entry-{index}", message(role)) for index, role in enumerate(roles)) + + +def _total_tokens(rows: tuple[tuple[str, AgentMessage], ...]) -> int: + return sum(estimate_message_tokens(message) for _entry_id, message in rows) + + +def test_first_recent_context_index_keeps_newest_user_turn_and_reply() -> None: + rows = _recent_rows("user", "assistant", "user", "assistant") + index = _first_recent_context_index(rows, keep_recent_tokens=_total_tokens(rows)) + assert index == 2 + assert [message.role for _entry_id, message in rows[index:]] == ["user", "assistant"] + + +def test_first_recent_context_index_never_compacts_past_last_user() -> None: + rows = _recent_rows("user", "assistant", "toolResult", "toolResult") + index = _first_recent_context_index(rows, keep_recent_tokens=_total_tokens(rows)) + assert index == 0 + + +def test_first_recent_context_index_keeps_pending_newest_user_prompt() -> None: + rows = _recent_rows("user", "assistant", "toolResult", "user") + index = _first_recent_context_index(rows, keep_recent_tokens=_total_tokens(rows)) + assert index == 3 + assert [message.role for _entry_id, message in rows[index:]] == ["user"] + + +def test_first_recent_context_index_keeps_recent_suffix() -> None: + rows = _recent_rows("user", "assistant", "user", "assistant", "user", "assistant") + index = _first_recent_context_index(rows, keep_recent_tokens=1) + assert index == 4 + assert [message.role for _entry_id, message in rows[index:]] == ["user", "assistant"] + + @pytest.mark.anyio async def test_prompt_queues_steering_while_session_is_running(tmp_path: Path) -> None: storage = JsonlSessionStorage(tmp_path / "session.jsonl")