ewilliams28/turnstone
1
0
Fork 0
mirror of https://github.com/turnstonelabs/turnstone.git synced 2026-08-12 23:12:23 -06:00
Code Issues Packages Projects Releases Wiki Activity

fix(session): provider-anchored context budget + cooperative compaction

Browse Source 
Unify truncation and compaction on one provider-anchored fullness measure
(_estimated_prompt_tokens), closing the 80-100% dead zone where tool output
was truncated but compaction never fired. Make compaction cooperative: advise
the model to wrap up and record its plan, compact if it continues, auto-resume
after a cooperative stop, and compact-before-truncate (preserving the in-flight
tool-call turn). Floor auto_compact_pct at 0.1 (invalid 0 -> default 0.8).
This commit is contained in:
Patrick Buckley
2026-06-27 12:52:42 -07:00
parent 9f54a97cc6
commit 65d1552ffa
5 changed files with 688 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_cooperative_compaction.py
+457
View File
@@ -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"
turnstone/core/metacognition.py
+17
View File
@@ -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``
turnstone/core/session.py
+204 -52
View File
@@ -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 --------------------------------------------------------
turnstone/core/settings_registry.py
+9 -3
View File
@@ -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 "
turnstone/core/tool_advisory.py
+1
View File
@@ -122,6 +122,7 @@ SYSTEM_TURN_SOURCES: Final = frozenset(
"start",
"tool_error",
"repeat",
"compaction_pending",
"idle_children",
"watch_triggered",
}