diff --git a/tests/test_cooperative_compaction.py b/tests/test_cooperative_compaction.py new file mode 100644 index 00000000..d0762eb6 --- /dev/null +++ b/tests/test_cooperative_compaction.py @@ -0,0 +1,457 @@ +"""Tests for provider-anchored context fullness and cooperative compaction. + +Covers the pieces that make compaction agree with tool-output truncation about +how full the context is, and that let the model reach a stopping point before +the harness collapses the transcript: + +- ``_estimated_prompt_tokens`` — the single fullness measure (provider + ``prompt_tokens`` + post-calibration delta, with a local fallback). +- ``_maybe_compact_midturn`` / ``_do_auto_compact`` — the soft-advise / + hard-compact escalation and the shared compaction action. +- the ``_compaction_advised`` latch lifecycle and the ``compaction_pending`` + advisory plumbing. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from tests._session_helpers import make_session +from turnstone.core.trajectory import dicts_from_turns, turns_from_dicts + + +@pytest.fixture +def session(tmp_db, mock_openai_client): + """ChatSession with a small window so thresholds are easy to reason about. + + context_window=10_000, auto_compact_pct default 0.8 → soft=8000, + hard=min(0.95, 0.9)*10_000=9000. Built via the shared ``make_session`` + factory so the session shape stays in lockstep with the sibling + truncation/compaction suites that read the same fullness measure. + """ + return make_session( + client=mock_openai_client, + context_window=10_000, + max_tokens=1_000, + tool_timeout=10, + ) + + +# --------------------------------------------------------------------------- +# _estimated_prompt_tokens — the shared fullness measure +# --------------------------------------------------------------------------- + + +class TestEstimatedPromptTokens: + def test_falls_back_to_local_without_usage(self, session): + """Before the first API call there is no provider anchor.""" + session._last_usage = None + session._system_tokens = 500 + session._msg_tokens = [100, 200] + assert session._estimated_prompt_tokens() == 800 + + def test_anchors_to_provider_usage_plus_delta(self, session): + """Provider prompt_tokens is ground truth; only post-calibration + messages are estimated on top (system + tool-def + cached prefix are + all already inside prompt_tokens).""" + session._last_usage = {"prompt_tokens": 8000} + session._calibrated_msg_count = 2 + session._msg_tokens = [1, 1, 300, 50] # delta = msgs after index 2 + assert session._estimated_prompt_tokens() == 8000 + 350 + + def test_clamps_stale_calibrated_count(self, session): + """A stale calibration index (post-compaction / mutation) must not + over-slice into a negative/garbage delta.""" + session._last_usage = {"prompt_tokens": 5000} + session._calibrated_msg_count = 99 # > len(_msg_tokens) + session._msg_tokens = [10, 20] + assert session._estimated_prompt_tokens() == 5000 + + +# --------------------------------------------------------------------------- +# _maybe_compact_midturn — soft-advise / hard-compact escalation +# --------------------------------------------------------------------------- + + +class TestMidturnCompactionPolicy: + def test_below_soft_threshold_is_noop(self, session): + with ( + patch.object(session, "_estimated_prompt_tokens", return_value=7_000), + patch.object(session, "_do_auto_compact") as compact, + patch.object(session, "_append_system_turn") as advise, + ): + session._maybe_compact_midturn() + compact.assert_not_called() + advise.assert_not_called() + assert session._compaction_advised is False + + def test_first_crossing_advises_not_compacts(self, session): + """Regression for the dead-zone bug: the provider reported ~85% while + the old naive estimate (system + msgs) was far under the 80% soft + threshold, so mid-turn compaction never fired and the model flailed on + truncated output. Now the provider-anchored estimate crosses soft and + the model is advised to wrap up first.""" + # Naive estimate is ~10% of the window... + session._system_tokens = 1_000 + session._msg_tokens = [1, 1] + session._calibrated_msg_count = 2 + assert session._system_tokens + sum(session._msg_tokens) < 8_000 + # ...but the provider counted 8_500 (tool defs + history) = 85%. + session._last_usage = {"prompt_tokens": 8_500} + session._compaction_advised = False + + with ( + patch.object(session, "_do_auto_compact") as compact, + patch.object(session, "_append_system_turn") as advise, + ): + session._maybe_compact_midturn() + + compact.assert_not_called() + advise.assert_called_once() + assert advise.call_args.args[0] == "compaction_pending" + assert session._compaction_advised is True + + def test_continue_after_advisory_compacts(self, session): + """Already advised + still over soft → the model kept working, compact.""" + session._compaction_advised = True + with ( + patch.object(session, "_estimated_prompt_tokens", return_value=8_500), + patch.object(session, "_do_auto_compact") as compact, + patch.object(session, "_append_system_turn") as advise, + ): + session._maybe_compact_midturn() + compact.assert_called_once_with("mid-turn") + advise.assert_not_called() + + def test_hard_ceiling_compacts_without_advisory(self, session): + """Over the hard ceiling → no turn to spare, compact even if never + advised.""" + session._compaction_advised = False + with ( + patch.object(session, "_estimated_prompt_tokens", return_value=9_500), + patch.object(session, "_do_auto_compact") as compact, + patch.object(session, "_append_system_turn") as advise, + ): + session._maybe_compact_midturn() + compact.assert_called_once_with("mid-turn") + advise.assert_not_called() + + def test_do_auto_compact_rounds_percentage(self, session): + """The notice uses round(), not int() — 0.58 must render '58%', not the + float-truncated '57%'.""" + session.auto_compact_pct = 0.58 + with ( + patch.object(session, "_compact_messages") as compact, + patch.object(session, "_print_status_line"), + patch.object(session.ui, "on_info") as on_info, + ): + session._do_auto_compact("mid-turn") + compact.assert_called_once_with(auto=True, preserve_tail=0) + msg = on_info.call_args.args[0] + assert "58%" in msg + assert "mid-turn" in msg + + +# --------------------------------------------------------------------------- +# Latch lifecycle + advisory plumbing +# --------------------------------------------------------------------------- + + +class TestCompactionLatch: + def test_stale_latch_cleared_on_send(self, session): + """A latch left True by a prior abnormal exit (cancel / error / + superseded / resume) must not survive into the next send and trigger an + advisory-skipping compaction. send() entry clears it.""" + session.messages = turns_from_dicts([{"role": "user", "content": "hi"}]) + session._msg_tokens = [1] + session._title_generated = True # don't spawn the auto-title daemon + session._compaction_advised = True # stale latch from a prior turn + + with ( + patch.object(session, "_create_stream_with_retry", return_value=iter([])), + patch.object( + session, "_stream_response", return_value={"role": "assistant", "content": "done"} + ), + patch.object(session, "_full_messages", return_value=[]), + patch.object(session, "_update_token_table"), + patch.object(session, "_print_status_line"), + patch.object(session, "_emit_state"), + patch.object(session, "_compact_messages"), + patch("turnstone.core.session.save_message"), + ): + session.send("hello") + + assert session._compaction_advised is False + + def test_compact_messages_clears_latch_even_when_it_bails(self, session): + """_compact_messages must clear the latch on every attempt — including + its early-return guards — so a bailed forced compaction falls back to + the advisory grace state instead of retry-storming.""" + session._compaction_advised = True + # One message → hits the "Not enough messages to compact" early return. + session.messages = turns_from_dicts([{"role": "user", "content": "hi"}]) + session._compact_messages(auto=True) + assert session._compaction_advised is False + + +class TestEndOfTurnAutoResume: + """End-of-turn: a cooperative stop (model wound down so we could compact) + resumes after compaction; a natural finish goes idle.""" + + def test_advised_stop_resumes_after_compaction(self, session): + """Latch True at the stop → after compaction a user turn re-prompts the + model to continue (the loop does not break to idle).""" + session.messages = turns_from_dicts([{"role": "user", "content": "task"}]) + session._msg_tokens = [1] + session._title_generated = True + + # send() entry resets the latch, so simulate the mid-turn advisory + # firing *during* the first stream (latch True), then the model stops. + calls = {"n": 0} + + def stream(*_a, **_k): + calls["n"] += 1 + if calls["n"] == 1: + session._compaction_advised = True # advisory fired this turn + return {"role": "assistant", "content": "paused; plan recorded"} + return {"role": "assistant", "content": "done"} + + with ( + patch.object(session, "_create_stream_with_retry", return_value=iter([])), + patch.object(session, "_stream_response", side_effect=stream), + patch.object(session, "_full_messages", return_value=[]), + patch.object(session, "_update_token_table"), + patch.object(session, "_print_status_line"), + patch.object(session, "_emit_state"), + patch.object(session, "_estimated_prompt_tokens", return_value=9_999), + patch.object(session, "_do_auto_compact"), + patch.object(session, "_append_user_turn") as resume, + patch("turnstone.core.session.save_message"), + ): + session.send("go") + + # send() also appends the user input ("go") via _append_user_turn, so + # filter for the resume turn specifically (tagged source=compaction_resume). + resume_calls = [ + c for c in resume.call_args_list if c.kwargs.get("source") == "compaction_resume" + ] + assert len(resume_calls) == 1 + + def test_natural_finish_idles_without_resume(self, session): + """Latch False at the stop (task genuinely done) → compact, then idle; + no auto-resume.""" + session.messages = turns_from_dicts([{"role": "user", "content": "task"}]) + session._msg_tokens = [1] + session._title_generated = True + session._compaction_advised = False + + with ( + patch.object(session, "_create_stream_with_retry", return_value=iter([])), + patch.object( + session, "_stream_response", return_value={"role": "assistant", "content": "done"} + ), + patch.object(session, "_full_messages", return_value=[]), + patch.object(session, "_update_token_table"), + patch.object(session, "_print_status_line"), + patch.object(session, "_emit_state") as emit_state, + patch.object(session, "_estimated_prompt_tokens", return_value=9_999), + patch.object(session, "_do_auto_compact") as compact, + patch.object(session, "_append_user_turn") as resume, + patch("turnstone.core.session.save_message"), + ): + session.send("go") + + compact.assert_called_once() + # No resume turn (only the user input "go" was appended). + assert not [ + c for c in resume.call_args_list if c.kwargs.get("source") == "compaction_resume" + ] + emit_state.assert_any_call("idle") + + def test_no_resume_when_compaction_bails(self, session): + """q-1 regression: if compaction bails (returns False — summary error / + too-large / too-few), the resume must NOT fire — there's no summary to + continue from.""" + session.messages = turns_from_dicts([{"role": "user", "content": "task"}]) + session._msg_tokens = [1] + session._title_generated = True + calls = {"n": 0} + + def stream(*_a, **_k): + calls["n"] += 1 + if calls["n"] == 1: + session._compaction_advised = True # advised stop + return {"role": "assistant", "content": "paused"} + return {"role": "assistant", "content": "done"} + + with ( + patch.object(session, "_create_stream_with_retry", return_value=iter([])), + patch.object(session, "_stream_response", side_effect=stream), + patch.object(session, "_full_messages", return_value=[]), + patch.object(session, "_update_token_table"), + patch.object(session, "_print_status_line"), + patch.object(session, "_emit_state"), + patch.object(session, "_estimated_prompt_tokens", return_value=9_999), + patch.object(session, "_do_auto_compact", return_value=False), # bailed + patch.object(session, "_append_user_turn") as resume, + patch("turnstone.core.session.save_message"), + ): + session.send("go") + + assert not [ + c for c in resume.call_args_list if c.kwargs.get("source") == "compaction_resume" + ] + + def test_resume_preserves_alternation(self, session): + """The auto-resume must not produce two consecutive user turns — some + providers require strict user/assistant alternation. Compaction leaves + a trailing assistant (summary) turn, and the resume user turn follows + it. Drives the real compaction + resume end to end.""" + session.messages = turns_from_dicts( + [ + {"role": "user", "content": "do the task"}, + {"role": "assistant", "content": "on it"}, + ] + ) + session._msg_tokens = [5, 5] + session._title_generated = True + session.compact_max_tokens = 100 # positive summary budget at ctx=10k + session._system_tokens = 0 + + summary = SimpleNamespace(content="## Open tasks\nfinish it", finish_reason="stop") + n = {"i": 0} + + def stream(*_a, **_k): + n["i"] += 1 + if n["i"] == 1: + session._compaction_advised = True # advisory fired this turn + return {"role": "assistant", "content": "pausing to compact"} + return {"role": "assistant", "content": "all done"} + + def est(*_a, **_k): + return 9_999 if n["i"] <= 1 else 10 # over threshold only on the stop turn + + with ( + patch.object(session, "_create_stream_with_retry", return_value=iter([])), + patch.object(session, "_stream_response", side_effect=stream), + patch.object(session, "_full_messages", return_value=[]), + patch.object(session, "_update_token_table"), + patch.object(session, "_print_status_line"), + patch.object(session, "_emit_state"), + patch.object(session, "_estimated_prompt_tokens", side_effect=est), + patch.object(session, "_utility_completion", return_value=summary), + patch("turnstone.core.session.save_message"), + ): + session.send("go") + + roles = [m["role"] for m in dicts_from_turns(session.messages)] + assert not any(roles[i] == roles[i + 1] == "user" for i in range(len(roles) - 1)), ( + f"consecutive user turns: {roles}" + ) + # The resume genuinely happened: a user turn sits after the summary. + assert "assistant" in roles and roles[-1] != "user" + + +class TestCompactBeforeTruncate: + """#2: tail-preserving compaction keeps the in-flight tool-call turn so the + fresh tool results aren't orphaned, gated by the shared _compaction_owed.""" + + def test_preserve_tail_keeps_in_flight_tool_call(self, session): + """compact(preserve_tail=1) summarizes the older history but keeps the + last (assistant tool-call) turn verbatim — so a tool result appended + after it still has its matching tool_use.""" + session.compact_max_tokens = 100 # positive summary budget at ctx=10k + session._system_tokens = 0 + tc = {"id": "call_1", "type": "function", "function": {"name": "x", "arguments": "{}"}} + session.messages = turns_from_dicts( + [ + {"role": "user", "content": "do it"}, + {"role": "assistant", "content": "older reply"}, + {"role": "user", "content": "more"}, + {"role": "assistant", "content": "", "tool_calls": [tc]}, # in-flight + ] + ) + session._msg_tokens = [5, 5, 5, 5] + summary = SimpleNamespace(content="dense summary", finish_reason="stop") + + with patch.object(session, "_utility_completion", return_value=summary): + session._compact_messages(auto=True, preserve_tail=1) + + wire = dicts_from_turns(session.messages) + # [summary_user, summary_asst, preserved assistant-tool-call] + assert wire[0]["role"] == "user" and "[Conversation summary]" in wire[0]["content"] + assert wire[1]["role"] == "assistant" + assert wire[-1]["role"] == "assistant" and wire[-1].get("tool_calls") + # The tool_call survived, so a tool result for call_1 won't orphan. + ids = [t["id"] for m in wire if m.get("tool_calls") for t in m["tool_calls"]] + assert "call_1" in ids + + def test_compaction_owed_predicate(self, session): + # over hard ceiling (>9000) → owed regardless of the latch + with patch.object(session, "_estimated_prompt_tokens", return_value=9_500): + session._compaction_advised = False + assert session._compaction_owed() is True + # over soft (>8000) → owed only when advised + with patch.object(session, "_estimated_prompt_tokens", return_value=8_500): + session._compaction_advised = True + assert session._compaction_owed() is True + session._compaction_advised = False + assert session._compaction_owed() is False + # under soft → never owed + with patch.object(session, "_estimated_prompt_tokens", return_value=7_000): + session._compaction_advised = True + assert session._compaction_owed() is False + + def test_owed_compaction_runs_before_truncation_in_tool_path(self, session): + """Wiring: in the tool path, an owed compaction fires with preserve_tail=1 + before the truncation budget is sized.""" + session.messages = turns_from_dicts([{"role": "user", "content": "task"}]) + session._msg_tokens = [1] + session._title_generated = True + tc = {"id": "call_1", "type": "function", "function": {"name": "x", "arguments": "{}"}} + n = {"i": 0} + + def stream(*_a, **_k): + n["i"] += 1 + if n["i"] == 1: + return {"role": "assistant", "content": "", "tool_calls": [tc]} + return {"role": "assistant", "content": "done"} + + with ( + patch.object(session, "_create_stream_with_retry", return_value=iter([])), + patch.object(session, "_stream_response", side_effect=stream), + patch.object(session, "_execute_tools", return_value=([("call_1", "out")], "")), + patch.object(session, "_full_messages", return_value=[]), + patch.object(session, "_update_token_table"), + patch.object(session, "_print_status_line"), + patch.object(session, "_emit_state"), + # Owed on the tool turn (pre-truncation); _estimated_prompt_tokens stays + # small so the end-of-turn path doesn't also compact. + patch.object(session, "_compaction_owed", side_effect=lambda: n["i"] == 1), + patch.object(session, "_maybe_compact_midturn"), # isolate the pre-truncation call + patch.object(session, "_do_auto_compact") as compact, + patch("turnstone.core.session.save_message"), + ): + session.send("go") + + assert any( + c.args == ("mid-turn",) and c.kwargs.get("preserve_tail") == 1 + for c in compact.call_args_list + ), compact.call_args_list + + +def test_compaction_advisory_is_registered(): + """The advisory source and template must be wired across both modules so + ``_append_system_turn('compaction_pending', ...)`` cannot raise.""" + from turnstone.core.metacognition import format_nudge + from turnstone.core.tool_advisory import SYSTEM_TURN_SOURCES, make_system_turn + + assert "compaction_pending" in SYSTEM_TURN_SOURCES + text = format_nudge("compaction_pending") + assert text and "compact" in text.lower() + turn = make_system_turn("compaction_pending", text) + assert turn["role"] == "system" + assert turn["_source"] == "compaction_pending" diff --git a/turnstone/core/metacognition.py b/turnstone/core/metacognition.py index a3db76c9..7a00b38d 100644 --- a/turnstone/core/metacognition.py +++ b/turnstone/core/metacognition.py @@ -110,6 +110,22 @@ NUDGE_REPEAT = ( "tool, different arguments, or ask the user for clarification." ) +NUDGE_COMPACTION = ( + "The conversation is approaching the context limit and will be compacted " + "shortly — older messages will be replaced by a summary. Reach a natural " + "stopping point. Before continuing, record in this turn your current goal, " + "the tasks that remain, and your intended next steps; anything not written " + "down here may be lost when compaction runs. If you can give a final answer " + "now, do so — otherwise state clearly where to resume." +) + +NUDGE_COMPACTION_RESUME = ( + "The conversation was just compacted to free context. If there is remaining " + "work, continue from the summary above — pick up the open tasks and next " + "steps you recorded and keep going without waiting for further instructions. " + "If the task is already complete, give your final answer." +) + _NUDGE_MAP: dict[str, str] = { "correction": NUDGE_CORRECTION, "denial": NUDGE_DENIAL, @@ -118,6 +134,7 @@ _NUDGE_MAP: dict[str, str] = { "start": NUDGE_START, "tool_error": NUDGE_TOOL_ERROR, "repeat": NUDGE_REPEAT, + "compaction_pending": NUDGE_COMPACTION, # idle_children and watch_triggered carry no static body — the # per-fire text comes from a producer (``format_idle_children_nudge`` # for the former, ``format_watch_message`` + ``sanitize_payload`` diff --git a/turnstone/core/session.py b/turnstone/core/session.py index 5b0a8a31..66bf6f9a 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -99,6 +99,7 @@ from turnstone.core.memory_relevance import ( score_memories, ) from turnstone.core.metacognition import ( + NUDGE_COMPACTION_RESUME, RepeatDetector, detect_completion, detect_correction, @@ -110,6 +111,7 @@ from turnstone.core.nudge_queue import TOOL_DRAIN, USER_DRAIN, NudgeQueue from turnstone.core.providers import create_provider from turnstone.core.ratelimit import TokenBucket from turnstone.core.safety import is_command_blocked, sanitize_command +from turnstone.core.settings_registry import DEFAULT_AUTO_COMPACT_PCT from turnstone.core.skill_field_validation import SKILL_RUNTIME_CONFIG_FIELDS from turnstone.core.skill_parser import MAX_SKILL_DESCRIPTION_LEN from turnstone.core.storage._registry import get_storage @@ -967,7 +969,7 @@ class ChatSession: reasoning_effort: str = "medium", context_window: int = 32768, compact_max_tokens: int = 32768, - auto_compact_pct: float = 0.8, + auto_compact_pct: float = DEFAULT_AUTO_COMPACT_PCT, agent_max_turns: int = -1, tool_truncation: int = 0, mcp_client: MCPClientManager | None = None, @@ -1036,7 +1038,20 @@ class ChatSession: self.reasoning_effort = reasoning_effort self.context_window = context_window if context_window > 0 else 32768 self.compact_max_tokens = compact_max_tokens - self.auto_compact_pct = auto_compact_pct + # auto_compact_pct < 0.1 is invalid: 0 used to mean "disabled" (no + # longer supported), and at 0 the "> context_window * pct" checks become + # "> 0" (always true → compact every turn). Coerce to the default, not + # the 0.1 floor — someone who set 0 wanted *less* compaction, so 0.8 is + # closer to intent than near-constant compaction at 10%. Matches the + # settings-registry path, which rejects sub-0.1 stored values and + # reverts to this same default. + self.auto_compact_pct = ( + auto_compact_pct if auto_compact_pct >= 0.1 else DEFAULT_AUTO_COMPACT_PCT + ) + # Cooperative-compaction latch: set once the model has been advised to + # reach a stopping point under context pressure; if it keeps working + # past the advisory the send loop forces a compaction. + self._compaction_advised = False self.agent_max_turns = agent_max_turns self._chars_per_token = 4.0 # calibrated from API usage # Tool output truncation: 0 means auto (50% of context_window in chars) @@ -2437,32 +2452,103 @@ class ChatSession: eid = getattr(self.ui, "_event_id", None) return eid if isinstance(eid, int) else None + def _estimated_prompt_tokens(self) -> int: + """Best estimate of the current prompt size, in tokens. + + Anchors to the provider-reported ``prompt_tokens`` from the last API + call — which already includes tool-definition tokens and the cached + prefix (providers fold cached + non-cached into one count at the API + boundary; see ``_anthropic.py`` ``total_input``) — and adds a local + estimate for only the messages appended since calibration. Falls + back to a pure local estimate before the first API call. + + Single source of truth for "how full is the context": tool-output + truncation (via :meth:`_remaining_token_budget`) and the + auto-compaction triggers all read it, so they cannot disagree about + the fullness of the same state. + """ + if self._last_usage: + # Clamp the index: a stale _calibrated_msg_count must not + # over-slice after compaction or message-list mutations. + start = min(self._calibrated_msg_count, len(self._msg_tokens)) + return self._last_usage["prompt_tokens"] + sum(self._msg_tokens[start:]) + return self._system_tokens + sum(self._msg_tokens) + def _remaining_token_budget(self) -> int: """Estimate how many tokens are available for new content. - When provider-reported usage is available, uses the last API - call's ``prompt_tokens`` as ground truth and only estimates the - delta (messages added since that call). Falls back to pure - local estimates otherwise. - Reserves a response budget (capped at 25% of context window, since - ``max_tokens`` is an upper bound, not guaranteed consumption) plus - a 5% safety margin. Returns at least 0. + ``max_tokens`` is an upper bound, not guaranteed consumption) plus a + 5% safety margin. Returns at least 0. Context fullness comes from + :meth:`_estimated_prompt_tokens` (provider-anchored), shared with the + auto-compaction triggers so truncation and compaction agree. """ - used = self._system_tokens + sum(self._msg_tokens) - if self._last_usage: - # Provider-reported tokens from the last API call - base = self._last_usage["prompt_tokens"] - # Only estimate tokens for messages added AFTER calibration. - # Clamp index to prevent stale _calibrated_msg_count from - # over-slicing after compaction or message list mutations. - start = min(self._calibrated_msg_count, len(self._msg_tokens)) - new_msg_tokens = sum(self._msg_tokens[start:]) - used = base + new_msg_tokens + used = self._estimated_prompt_tokens() response_reserve = min(self.max_tokens, self.context_window // 4) safety_margin = int(self.context_window * 0.05) return max(0, self.context_window - used - response_reserve - safety_margin) + def _maybe_compact_midturn(self) -> None: + """Cooperative mid-turn compaction policy. + + Reads the provider-anchored fullness estimate (the same measure + tool-output truncation uses, so the two cannot disagree about the + same state) and escalates: + + - over the **hard** ceiling — no turn to spare, compact now; + - over the **soft** threshold and already advised — the model kept + working past the wrap-up advisory, compact; + - over the **soft** threshold, first crossing — append a + ``compaction_pending`` advisory asking the model to reach a stopping + point and record its remaining plan, so the summariser preserves it + (a register-spill onto the transcript before the tape is collapsed), + and latch ``_compaction_advised``. + + No-op below the soft threshold. The latch is cleared by + :meth:`_compact_messages` and at end-of-turn. + """ + est = self._estimated_prompt_tokens() + if self._compaction_owed(est): + self._do_auto_compact("mid-turn") + elif est > self.context_window * self.auto_compact_pct: + self._append_system_turn("compaction_pending", format_nudge("compaction_pending")) + self._compaction_advised = True + + def _compaction_owed(self, used: int | None = None) -> bool: + """True when fullness mandates compaction now: over the hard ceiling, or + over the soft threshold after the model worked past a wrap-up advisory. + + Shared by the compact-before-truncate check in the send loop and + :meth:`_maybe_compact_midturn`. Pass a precomputed ``used`` to avoid a + redundant :meth:`_estimated_prompt_tokens` sum. The hard ceiling sits a + band above the soft threshold (capped at 95%); if ``auto_compact_pct`` is + set above that, the band collapses and any over-soft state is "owed". + """ + if used is None: + used = self._estimated_prompt_tokens() + soft = self.context_window * self.auto_compact_pct + hard = self.context_window * min(0.95, self.auto_compact_pct + 0.10) + return used > hard or (used > soft and self._compaction_advised) + + def _do_auto_compact(self, where: str = "", preserve_tail: int = 0) -> bool: + """Emit the auto-compaction notice, compact, and refresh the status + line. Shared by the mid-turn policy (:meth:`_maybe_compact_midturn`) + and the end-of-turn check so the notice wording, the percentage, and + the post-compaction status refresh stay in lockstep. ``where`` is an + optional qualifier for the notice (e.g. ``"mid-turn"``); ``preserve_tail`` + is forwarded to :meth:`_compact_messages` (e.g. to keep an in-flight + tool-call turn during compact-before-truncate). Returns whether a + summary was actually produced (False if compaction bailed) so callers + can avoid acting on a compaction that did not happen.""" + qualifier = f" {where}" if where else "" + pct_display = round(self.auto_compact_pct * 100) + self.ui.on_info( + f"\n[Auto-compacting{qualifier}: prompt exceeds {pct_display}% of context window]" + ) + compacted = self._compact_messages(auto=True, preserve_tail=preserve_tail) + self._print_status_line() + return compacted + def _truncate_output(self, output: str, remaining_budget_tokens: int | None = None) -> str: """Truncate tool output, keeping head + tail. @@ -3943,6 +4029,7 @@ class ChatSession: send_id: str | None = None, *, from_wake: bool = False, + source: str | None = None, ) -> int: """Append a user turn (plain or multipart) and persist it. @@ -4004,6 +4091,11 @@ class ChatSession: # delivered while a wake is in flight, NOT synthetic, so # they must not inherit the wake tag. user_msg["_source"] = self._wake_source_tag + elif source: + # Provenance for other self-prompted turns (e.g. the post-compaction + # auto-resume): marks the turn for audit / replay / UI so it isn't + # mistaken for real user input. Stripped at the sanitize boundary. + user_msg["_source"] = source if attachments: # Sibling metadata so live history replay has the same shape # as reloaded-from-DB (filenames are not recoverable from an @@ -4240,6 +4332,12 @@ class ChatSession: self._budget_exhausted = False self._budget_warned = False self._notify_count = 0 + # Per-send cooperative-compaction latch: each send starts a fresh + # advise→compact cycle, so reset here. This single chokepoint covers + # the cancel / error / superseded / resume / clear / new exits that + # would otherwise leave the latch set on the long-lived session and + # trip a premature, advisory-skipping compaction on the next send. + self._compaction_advised = False self._generation += 1 my_generation = self._generation # Fresh cancel event per generation. The old event object stays @@ -4414,19 +4512,36 @@ class ChatSession: tool_calls = assistant_msg.get("tool_calls") if not tool_calls: - # Auto-compact when prompt exceeds threshold + # Did the model stop because we asked it to wind down for a + # compaction (cooperative), or because the task is actually + # done? Capture before the reset — it gates the auto-resume. + stopped_to_compact = self._compaction_advised + self._compaction_advised = False + # Auto-compact when the context exceeds the threshold, so the + # next turn starts with headroom. Bare-soft check (NOT + # _compaction_owed()): the turn already ended, so there's no + # model cooperation to wait for and the latch was just + # consumed above — compact whenever over soft. None-safe via + # _estimated_prompt_tokens(), so no _last_usage guard. if ( - self._last_usage - and self._last_usage["prompt_tokens"] + self._estimated_prompt_tokens() > self.context_window * self.auto_compact_pct ): - pct_display = int(self.auto_compact_pct * 100) - self.ui.on_info( - f"\n[Auto-compacting: prompt exceeds {pct_display}% of context window]" - ) - self._compact_messages(auto=True) - # Update status bar with post-compaction token counts - self._print_status_line() + compacted = self._do_auto_compact() + if stopped_to_compact and compacted: + # The model paused mid-task to let us compact, not + # because it was finished. Hand the compacted state + # back as a user turn so it resumes instead of being + # stranded at idle — but only when a summary was + # actually produced (else there's nothing to continue + # from). The prompt lets a genuinely-finished model + # give its final answer and stop. + self._append_user_turn( + NUDGE_COMPACTION_RESUME, + (), + source="compaction_resume", + ) + continue # Flush any queued messages that weren't injected # (no tool calls → no advisory seam to inject at). # If anything drained, the model hasn't seen those @@ -4469,6 +4584,19 @@ class ChatSession: # without per-output bookkeeping, N parallel tool results # could each claim the full remaining budget and collectively # overflow the prompt. + # Compact-before-truncate: if a compaction is already owed (over + # the hard ceiling, or the model worked past a wrap-up advisory), + # do it BEFORE sizing the truncation budget — preserving the + # in-flight assistant tool-call turn (preserve_tail=1) so the + # results about to be appended aren't orphaned from their + # tool_use. The freed context then lets the fresh results + # through (largely) untruncated instead of snipping them only to + # summarise them moments later. Generation-guarded so an + # orphaned thread can't replace history under the active one. + pre_attempted_compact = False + if self._generation == my_generation and self._compaction_owed(): + self._do_auto_compact("mid-turn", preserve_tail=1) + pre_attempted_compact = True truncation_budget = self._remaining_token_budget() _truncated: dict[str, str] = {} for tc_id, output in results: @@ -4654,17 +4782,22 @@ class ChatSession: # joined to queued items by ``\n\n``. self._flush_queued_messages(prefix=user_feedback or "") - # Mid-turn compaction: prevent context overflow during long - # tool chains. Uses local estimates since _last_usage reflects - # the previous API call, not the tool results just appended. - estimated_prompt = self._system_tokens + sum(self._msg_tokens) - if estimated_prompt > self.context_window * self.auto_compact_pct: - pct_display = int(self.auto_compact_pct * 100) - self.ui.on_info( - f"\n[Auto-compacting mid-turn: estimated prompt " - f"exceeds {pct_display}% of context window]" - ) - self._compact_messages(auto=True) + # Don't mutate shared history from an orphaned (superseded) + # thread: a force-cancel handoff bumps _generation, and + # _maybe_compact_midturn can replace self.messages out from + # under the active generation. + if self._generation != my_generation: + return + # Cooperative mid-turn compaction — advise once under context + # pressure so the model can reach a stopping point and spill its + # plan, then compact if it keeps working (or immediately over + # the hard ceiling). Skip if a compaction was already attempted + # pre-truncation this iteration: re-running would double the work + # (and retry-storm on a failed summary); truncation already + # bounded the batch, and the next iteration / end-of-turn + # re-checks. + if not pre_attempted_compact: + self._maybe_compact_midturn() except GenerationCancelled: # If a newer send() has started (force cancel), this thread is # orphaned — skip all message mutations and state changes. @@ -5496,18 +5629,28 @@ class ChatSession: parts.append(f"{role}: {content}") return "\n\n".join(parts) - def _compact_messages(self, auto: bool = False) -> None: - """Compact conversation history by summarizing all messages. + def _compact_messages(self, auto: bool = False, preserve_tail: int = 0) -> bool: + """Compact conversation history by summarizing it into a summary turn. - Summarizes the entire conversation via a separate model call, - budget-fitted to 80% of the context window. + Summarizes the conversation via a separate model call, budget-fitted to + ``auto_compact_pct`` of the context window. When auto=True (triggered by context limit), appends a continuation hint with the last user message so the model can resume seamlessly. + + ``preserve_tail`` keeps the last N messages verbatim (re-appended after + the summary) instead of summarizing them — used to keep an in-flight + assistant tool-call turn so the tool results appended after compaction + aren't orphaned from their ``tool_use``. """ + # Clear the cooperative latch on every compaction *attempt*, ahead of + # the early-return guards below — a bailed compaction (too few/large + # messages, summary error) must fall back to the advisory grace state + # next cycle rather than retry-storm on the same over-soft estimate. + self._compaction_advised = False if len(self.messages) < 2: self.ui.on_info("Not enough messages to compact.") - return + return False # Find the last user message for the continuation hint last_user_content = None @@ -5517,7 +5660,12 @@ class ChatSession: last_user_content = m.text or "" break - to_summarize = self.messages + # Optionally keep the last ``preserve_tail`` messages verbatim — e.g. an + # in-flight assistant tool-call whose results are about to be appended — + # so compacting them away can't orphan a tool result. Re-appended after + # the summary, below. + preserved = self.messages[-preserve_tail:] if preserve_tail > 0 else [] + to_summarize = self.messages[:-preserve_tail] if preserve_tail > 0 else self.messages # Budget: fit as many messages as possible into summary request summary_max_tokens = self.compact_max_tokens @@ -5541,7 +5689,7 @@ class ChatSession: if not selected: self.ui.on_info("Messages too large to fit in summary context.") - return + return False # Build summary prompt and call model formatted = self._format_messages_for_summary(dicts_from_turns(selected)) @@ -5615,7 +5763,7 @@ class ChatSession: self.ui.on_info("[Warning: compaction summary was truncated]") except Exception as e: self.ui.on_error(f"Compaction failed: {e}") - return + return False finally: self.ui.on_thinking_stop() @@ -5630,19 +5778,22 @@ class ChatSession: f"Continue assisting from where we left off." ) - # Replace messages + # Replace messages — summary, then any preserved tail verbatim. before_tokens = self._system_tokens + sum(self._msg_tokens) summary_user = {"role": "user", "content": "[Conversation summary]"} summary_asst = {"role": "assistant", "content": summary} - self.messages = turns_from_dicts([summary_user, summary_asst]) + self.messages = turns_from_dicts([summary_user, summary_asst]) + list(preserved) # File contents are gone after compaction — force re-read before edit_file self._read_files.clear() self._repeat_detector.clear() - # Rebuild token table + # Rebuild token table — summary turns + preserved-tail estimates. su_tok = max(1, int(self._msg_char_count(summary_user) / self._chars_per_token)) sa_tok = max(1, int(self._msg_char_count(summary_asst) / self._chars_per_token)) - self._msg_tokens = [su_tok, sa_tok] + tail_toks = [ + max(1, int(self._msg_char_count(m) / self._chars_per_token)) for m in preserved + ] + self._msg_tokens = [su_tok, sa_tok, *tail_toks] self._calibrated_msg_count = len(self.messages) # anchored to compacted state after_tokens = self._system_tokens + sum(self._msg_tokens) @@ -5661,6 +5812,7 @@ class ChatSession: lines.append(f" {line}") lines.append(separator) self.ui.on_info("\n".join(lines)) + return True # -- Intent validation -------------------------------------------------------- diff --git a/turnstone/core/settings_registry.py b/turnstone/core/settings_registry.py index 7d239f33..abea5ce0 100644 --- a/turnstone/core/settings_registry.py +++ b/turnstone/core/settings_registry.py @@ -31,6 +31,12 @@ class SettingDef: reference_url: str = "" # link to arXiv, docs, or provider reference +# Default auto-compaction trigger as a fraction of the context window. Shared +# with the ChatSession constructor (default + sub-0.1 coercion fallback) so the +# "invalid → default" behavior cannot drift between the registry and the engine. +DEFAULT_AUTO_COMPACT_PCT = 0.8 + + def _build_registry() -> dict[str, SettingDef]: """Build the settings registry from declarative definitions.""" defs: list[SettingDef] = [ @@ -138,10 +144,10 @@ def _build_registry() -> dict[str, SettingDef]: SettingDef( "session.auto_compact_pct", "float", - 0.8, - "Auto-compact at this fraction of context window (0 = disabled)", + DEFAULT_AUTO_COMPACT_PCT, + "Auto-compact at this fraction of context window", "session", - min_value=0.0, + min_value=0.1, max_value=1.0, help="Automatically summarize older messages when the conversation fills this percentage " "of the context window. For example, 0.8 means compact when 80% full. This prevents " diff --git a/turnstone/core/tool_advisory.py b/turnstone/core/tool_advisory.py index 11661836..2bf2506a 100644 --- a/turnstone/core/tool_advisory.py +++ b/turnstone/core/tool_advisory.py @@ -122,6 +122,7 @@ SYSTEM_TURN_SOURCES: Final = frozenset( "start", "tool_error", "repeat", + "compaction_pending", "idle_children", "watch_triggered", }