From 85391801577285b4490989f0205b7c61d5b08491 Mon Sep 17 00:00:00 2001 From: Lavindeep Dhillon Date: Fri, 10 Jul 2026 07:05:22 -0400 Subject: [PATCH 1/5] fix(runtime): count full request tokens before every model call Include tool schemas and tool-call arguments in the live estimate, compact at the soft threshold, refuse when the hard limit still cannot be met, and apply the command prompt cap to run_command results. --- shellpilot/cli/slash.py | 2 +- shellpilot/runtime/conversation.py | 61 +++++++++++++++++++++------- shellpilot/runtime/executor.py | 11 ++++- tests/test_budget.py | 8 +--- tests/test_compaction.py | 18 ++++++--- tests/test_conversation.py | 64 +++++++++++++++++++++++++++--- tests/test_executor.py | 27 +++++++++++++ 7 files changed, 157 insertions(+), 34 deletions(-) diff --git a/shellpilot/cli/slash.py b/shellpilot/cli/slash.py index 1a57445..ceeecc5 100644 --- a/shellpilot/cli/slash.py +++ b/shellpilot/cli/slash.py @@ -1077,7 +1077,7 @@ def _context(self) -> None: "[green]yes[/green]", "", ) - total = snapshot.est_system_tokens + tool_tokens + history_tokens + total = self._runtime.estimated_prompt_tokens() table.add_row( "TOTAL", f"of {budget.model_context_tokens} (compact at {budget.compact_at_tokens})", diff --git a/shellpilot/runtime/conversation.py b/shellpilot/runtime/conversation.py index 08345ff..739b151 100644 --- a/shellpilot/runtime/conversation.py +++ b/shellpilot/runtime/conversation.py @@ -411,18 +411,27 @@ def tool_schema_tokens(self) -> int: for definition in self._registry.definitions_for_profile(profile) ) + def _message_tokens(self, message: Message) -> int: + """Estimate one history message the way it is serialized for the model.""" + total = estimate_tokens(message.content) + total += IMAGE_TOKEN_ESTIMATE * len(message.images) + for call in message.tool_calls: + total += estimate_tokens( + json.dumps({"function": {"name": call.name, "arguments": call.arguments}}) + ) + return total + def history_token_estimate(self) -> tuple[int, int]: - """Estimated history tokens (incl. images) and message count, counted - exactly as estimated_prompt_tokens does for the /context breakdown.""" - total = 0 - for message in self._history: - total += estimate_tokens(message.content) - total += IMAGE_TOKEN_ESTIMATE * len(message.images) + """Estimated history tokens (content, images, tool-call args) and count.""" + total = sum(self._message_tokens(message) for message in self._history) return total, len(self._history) def estimated_prompt_tokens(self) -> int: + """Complete request estimate: system + tool schemas + history.""" history_tokens, _ = self.history_token_estimate() - return self._context_snapshot().est_system_tokens + history_tokens + return ( + self._context_snapshot().est_system_tokens + self.tool_schema_tokens() + history_tokens + ) def status(self) -> RuntimeStatus: return RuntimeStatus( @@ -448,12 +457,9 @@ def compact_now(self) -> int: metadata live outside history and are never touched. """ changed = 0 - # System context is invariant during compaction (only self._history - # mutates), so estimate it once and re-test only the changing history sum. - system_tokens = self._context_snapshot().est_system_tokens def over() -> bool: - return system_tokens + self.history_token_estimate()[0] > self.budget.compact_at_tokens + return self.estimated_prompt_tokens() > self.budget.compact_at_tokens # Digestion may reach everything except the in-flight exchange (last 2 # messages); snapshot staleness checks still force a fresh read before @@ -490,6 +496,28 @@ def over() -> bool: changed += 1 return changed + def _ensure_under_hard_limit(self) -> bool: + """Compact at the soft threshold when enabled; refuse past the hard limit.""" + if ( + self._settings.runtime.auto_compact + and self.estimated_prompt_tokens() > self.budget.compact_at_tokens + ): + adjusted = self.compact_now() + if adjusted: + self._ui.show_status(f"Compacted context: adjusted {adjusted} messages.") + if self.estimated_prompt_tokens() <= self.budget.hard_limit_tokens: + return True + self._ui.show_status( + "Context is over the hard limit" + + ( + " even after compaction. Run /clear, or shorten the request." + if self._settings.runtime.auto_compact + else " and automatic compaction is off. " + "Run /compact (or /clear), or turn it back on with /compact auto on." + ) + ) + return False + def run_turn( self, text: str, @@ -541,10 +569,8 @@ def run_turn( audit_kwargs["images"] = len(images) self._audit.write("user_turn", **audit_kwargs) self._record(user(text, images=tuple(images))) - if self._settings.runtime.auto_compact: - adjusted = self.compact_now() - if adjusted: - self._ui.show_status(f"Compacted context: adjusted {adjusted} messages.") + if not self._ensure_under_hard_limit(): + return "" content = self._tool_loop().content self._ui.turn_finished(self._turn_stats(time.monotonic() - started)) return content @@ -608,6 +634,7 @@ def _tool_loop(self) -> Message: profile=self._settings.runtime.security_profile, max_result_tokens=self.budget.max_tool_prompt_tokens, max_total_tokens=self.budget.max_total_tool_prompt_tokens, + max_command_prompt_tokens=self.budget.max_command_prompt_tokens, max_capture_chars=self.budget.max_command_capture_chars, command_timeout_seconds=self._settings.runtime.command_timeout_seconds, ask_approval=self._ui.ask_approval, @@ -622,8 +649,11 @@ def _tool_loop(self) -> Message: nudges_used = 0 empty_nudges_used = 0 consecutive_malformed = 0 + last_reply = Message(role="assistant", content="") while True: + if not self._ensure_under_hard_limit(): + return last_reply messages = [ Message(role="system", content=self._system_message_text()), *self._history, @@ -663,6 +693,7 @@ def _tool_loop(self) -> Message: finally: self._ui.end_response() self._turn_output_tokens += reply.output_tokens + last_reply = reply # History length BEFORE this model step is recorded, so a mid-tool # cancel (below) can roll the step back out and leave no orphaned # tool_call behind (§31.15). diff --git a/shellpilot/runtime/executor.py b/shellpilot/runtime/executor.py index 9fcd825..31c5fde 100644 --- a/shellpilot/runtime/executor.py +++ b/shellpilot/runtime/executor.py @@ -60,6 +60,7 @@ def __init__( profile: str, max_result_tokens: int, max_total_tokens: int, + max_command_prompt_tokens: int | None = None, max_capture_chars: int = 200_000, command_timeout_seconds: int = 600, ask_approval: ApprovalAsker | None = None, @@ -76,6 +77,9 @@ def __init__( self._profile = profile self._max_result_tokens = max_result_tokens self._max_total_tokens = max_total_tokens + self._max_command_prompt_tokens = ( + max_result_tokens if max_command_prompt_tokens is None else max_command_prompt_tokens + ) self._max_capture_chars = max_capture_chars self._command_timeout_seconds = command_timeout_seconds self._ask_approval = ask_approval @@ -284,7 +288,12 @@ def _render(self, name: str, result: ToolResult) -> str: if remaining <= 0: content = "[omitted: total tool-output budget for this turn is spent]" else: - content, _ = truncate_to_tokens(content, min(remaining, self._max_result_tokens)) + per_tool_cap = ( + self._max_command_prompt_tokens + if name == "run_command" + else self._max_result_tokens + ) + content, _ = truncate_to_tokens(content, min(remaining, per_tool_cap)) text = f"{header}\n---\n{content}" if content else header if result.truncated: text += "\n(note: output truncated)" diff --git a/tests/test_budget.py b/tests/test_budget.py index 9ac07f9..2187d21 100644 --- a/tests/test_budget.py +++ b/tests/test_budget.py @@ -27,18 +27,14 @@ def test_large_context_clamps_reservations() -> None: assert budget.max_user_message_tokens == 4096 # min(4096, 25%) assert budget.max_tool_prompt_tokens == 2000 assert budget.max_total_tool_prompt_tokens == 8000 - - -def test_explicit_context_setting_is_not_capped() -> None: - explicit = ContextSettings(model_context_tokens=131_072) - budget = resolve_budget(explicit, detected_context_tokens=None) - assert budget.model_context_tokens == 131_072 + assert budget.max_command_prompt_tokens == 2000 def test_small_context_scales_down_tool_budgets() -> None: budget = resolve_budget(ContextSettings(), detected_context_tokens=4096) assert budget.reserved_response_tokens == 1024 # clamped low assert budget.max_tool_prompt_tokens == 409 # 10% beats the 2000 cap + assert budget.max_command_prompt_tokens == 409 assert budget.max_user_message_tokens == 1024 diff --git a/tests/test_compaction.py b/tests/test_compaction.py index 65a01be..d9aeba2 100644 --- a/tests/test_compaction.py +++ b/tests/test_compaction.py @@ -36,8 +36,10 @@ def make_runtime( def test_old_tool_results_are_digested_before_anything_drops(tmp_path: Path) -> None: - runtime, _, _ = make_runtime(tmp_path, context_tokens=2048) - big_output = "line of tool output\n" * 300 + # Schemas alone are ~1.4k tokens; 4096 leaves room under the hard limit while + # still forcing compaction once a large tool result is present. + runtime, _, _ = make_runtime(tmp_path, context_tokens=4096) + big_output = "line of tool output\n" * 800 runtime._history.extend( [ Message(role="user", content="please inspect the project"), @@ -52,6 +54,7 @@ def test_old_tool_results_are_digested_before_anything_drops(tmp_path: Path) -> Message(role="assistant", content="working on it"), ] ) + assert runtime.estimated_prompt_tokens() > runtime.budget.compact_at_tokens runtime.compact_now() roles = [message.role for message in runtime._history] assert roles.count("user") == 2 # no user message was dropped @@ -62,7 +65,7 @@ def test_old_tool_results_are_digested_before_anything_drops(tmp_path: Path) -> def test_assistant_tool_call_drops_together_with_results(tmp_path: Path) -> None: - runtime, _, _ = make_runtime(tmp_path, context_tokens=256) + runtime, _, _ = make_runtime(tmp_path, context_tokens=2048) runtime._history.extend( [ Message(role="user", content="inspect " + "pad " * 40), @@ -90,7 +93,7 @@ def test_assistant_tool_call_drops_together_with_results(tmp_path: Path) -> None def test_user_messages_drop_last_and_newest_is_kept(tmp_path: Path) -> None: - runtime, _, _ = make_runtime(tmp_path, context_tokens=96) + runtime, _, _ = make_runtime(tmp_path, context_tokens=2048) for index in range(6): runtime._history.append(Message(role="user", content=f"instruction {index} " + "pad " * 30)) runtime.compact_now() @@ -101,9 +104,12 @@ def test_user_messages_drop_last_and_newest_is_kept(tmp_path: Path) -> None: def test_auto_compact_off_refuses_past_hard_limit(tmp_path: Path) -> None: runtime, ui, fake = make_runtime( - tmp_path, context_tokens=256, auto_compact=False, script=[answer("hi")] + tmp_path, context_tokens=4096, auto_compact=False, script=[answer("hi")] + ) + runtime._history.extend(Message(role="user", content="pad " * 600) for _ in range(4)) + assert ( + runtime.estimated_prompt_tokens() + 50 > runtime.budget.hard_limit_tokens ) - runtime._history.extend(Message(role="user", content="pad " * 100) for _ in range(4)) reply = runtime.run_turn("over the limit now") assert reply == "" assert fake.calls == [] # the model was never called diff --git a/tests/test_conversation.py b/tests/test_conversation.py index db82d8e..2cf80b2 100644 --- a/tests/test_conversation.py +++ b/tests/test_conversation.py @@ -621,17 +621,18 @@ def test_oversized_user_message_is_refused(tmp_path: Path) -> None: def test_compaction_drops_oldest_turns(tmp_path: Path) -> None: - # Tiny explicit context so a few turns cross the threshold. - settings = Settings(context=ContextSettings(model_context_tokens=512)) - script = [answer(f"reply {i} " + "pad " * 30) for i in range(6)] + # Context must leave room for tool schemas (~1.4k) while still compacting. + settings = Settings(context=ContextSettings(model_context_tokens=4096)) + script = [answer(f"reply {i} " + "pad " * 40) for i in range(8)] fake = FakeLLM(script=script) ui = FakeUI() runtime = make_runtime(fake, ui, tmp_path, settings) - for i in range(6): - runtime.run_turn(f"question {i} " + "pad " * 30) + for i in range(8): + runtime.run_turn(f"question {i} " + "pad " * 120) assert any("Compacted" in status for status in ui.statuses) + assert fake.calls final_call = fake.calls[-1] contents = [message.content for message in final_call.messages] assert not any("question 0" in content for content in contents) @@ -1021,6 +1022,59 @@ def test_image_token_estimate_counted(tmp_path: Path) -> None: assert estimate_with_image >= estimate_no_image + IMAGE_TOKEN_ESTIMATE +def test_estimated_prompt_tokens_includes_schemas_and_tool_call_args(tmp_path: Path) -> None: + """The live request estimate must count tool schemas and tool-call arguments.""" + import json + + from shellpilot.llm.messages import Message, ToolCall + from shellpilot.runtime.budget import estimate_tokens + + runtime = make_runtime(FakeLLM(script=[]), FakeUI(), tmp_path) + baseline = runtime.estimated_prompt_tokens() + assert baseline == ( + runtime.context_snapshot().est_system_tokens + + runtime.tool_schema_tokens() + + runtime.history_token_estimate()[0] + ) + assert runtime.tool_schema_tokens() > 0 + + huge_args = {"path": "x" * 4000} + runtime.restore_history( + [ + Message( + role="assistant", + content="", + tool_calls=(ToolCall(name="read_file", arguments=huge_args),), + ) + ] + ) + with_args = runtime.estimated_prompt_tokens() + arg_tokens = estimate_tokens( + json.dumps({"function": {"name": "read_file", "arguments": huge_args}}) + ) + assert with_args >= baseline + arg_tokens + + +def test_ensure_under_hard_limit_blocks_when_compaction_cannot_recover( + tmp_path: Path, +) -> None: + """The shared hard-limit gate refuses when kept context still exceeds the limit.""" + from shellpilot.config.model import ContextSettings, RuntimeSettings + from shellpilot.llm.messages import Message + + settings = Settings( + context=ContextSettings(model_context_tokens=2048), + runtime=RuntimeSettings(auto_compact=True), + ) + ui = FakeUI() + runtime = make_runtime(FakeLLM(script=[]), ui, tmp_path, settings=settings) + # Keep a single huge user message so compaction cannot drop below the hard limit. + runtime.restore_history([Message(role="user", content="x" * 20_000)]) + assert runtime.estimated_prompt_tokens() > runtime.budget.hard_limit_tokens + assert not runtime._ensure_under_hard_limit() + assert any("hard limit" in status.lower() for status in ui.statuses) + + def test_set_workspace_rebuilds_project_memory(tmp_path: Path) -> None: """Changing workspace (/cwd) rebuilds the project memory store for the new path, so the previous workspace's facts stop injecting (design section 16); diff --git a/tests/test_executor.py b/tests/test_executor.py index b2a16a1..978d1a1 100644 --- a/tests/test_executor.py +++ b/tests/test_executor.py @@ -856,3 +856,30 @@ def test_approval_display_marks_path_escaping_workspace(tmp_path: Path) -> None: ) assert escape not in request.display assert OUTSIDE_WORKSPACE_DISPLAY in request.display + + +def test_run_command_uses_command_prompt_cap(tmp_path: Path) -> None: + """run_command model text is truncated by max_command_prompt_tokens, not max_result_tokens.""" + from shellpilot.tools.base import ToolResult + from shellpilot.tools.command import RUN_COMMAND + from shellpilot.tools.registry import ToolRegistry + + registry = ToolRegistry() + registry.register(RUN_COMMAND) + executor = ToolExecutor( + registry=registry, + workspace=tmp_path, + profile="balanced", + max_result_tokens=2000, + max_total_tokens=10_000, + max_command_prompt_tokens=20, + ) + rendered = executor._render( + "run_command", + ToolResult(success=True, summary="ok", content="word " * 500), + ) + from shellpilot.runtime.budget import estimate_tokens + + # Header + body should stay near the command cap, far below the general tool cap. + assert estimate_tokens(rendered) < 200 + assert "word" in rendered From 52dbed123aee358a7b63c06b0fdd8a896fe401f7 Mon Sep 17 00:00:00 2001 From: Lavindeep Dhillon Date: Fri, 10 Jul 2026 07:12:25 -0400 Subject: [PATCH 2/5] fix(runtime): roll back turns that hit the hard context limit Refuse oversized turns before they stick in history/session, force-digest in-flight tool results, and discard incomplete assistant+tool exchanges so the next turn is not permanently stuck over budget. --- shellpilot/persistence/sessions.py | 11 ++++ shellpilot/runtime/conversation.py | 96 ++++++++++++++++++++++++------ tests/test_budget.py | 6 ++ tests/test_conversation.py | 83 ++++++++++++++++++++++++++ 4 files changed, 179 insertions(+), 17 deletions(-) diff --git a/shellpilot/persistence/sessions.py b/shellpilot/persistence/sessions.py index efb02e3..8e20cfd 100644 --- a/shellpilot/persistence/sessions.py +++ b/shellpilot/persistence/sessions.py @@ -145,6 +145,14 @@ def truncate_last_turn(self) -> None: """ self._append({"type": "truncate_last_turn"}) + def discard_last_message(self) -> None: + """Record that the trailing transcript message is discarded. + + Used when a user turn is recorded and then refused (e.g. hard context + limit) so --resume does not keep a stuck user message with no reply. + """ + self._append({"type": "discard_last_message"}) + def _message_record(self, kind: str, message: Message) -> dict[str, Any]: content = redact_secrets(message.content) if self._redact else message.content record: dict[str, Any] = { @@ -215,6 +223,9 @@ def load(path: Path) -> LoadedSession: if messages[i].role == "assistant": del messages[i:] break + elif kind == "discard_last_message": + if messages: + messages.pop() elif kind in ("message", "replace_last_message"): role = record.get("role") if not role: diff --git a/shellpilot/runtime/conversation.py b/shellpilot/runtime/conversation.py index 739b151..8450143 100644 --- a/shellpilot/runtime/conversation.py +++ b/shellpilot/runtime/conversation.py @@ -507,16 +507,56 @@ def _ensure_under_hard_limit(self) -> bool: self._ui.show_status(f"Compacted context: adjusted {adjusted} messages.") if self.estimated_prompt_tokens() <= self.budget.hard_limit_tokens: return True - self._ui.show_status( - "Context is over the hard limit" - + ( - " even after compaction. Run /clear, or shorten the request." - if self._settings.runtime.auto_compact - else " and automatic compaction is off. " - "Run /compact (or /clear), or turn it back on with /compact auto on." + self._ui.show_status(self._hard_limit_status()) + return False + + def _hard_limit_status(self) -> str: + """Status text when the hard context limit blocks a model call.""" + if ( + not self._history + and self.estimated_prompt_tokens() > self.budget.hard_limit_tokens + ): + return ( + "System prompt and tool schemas alone exceed the hard limit. " + "Raise context.model_context_tokens, or trim tools/AGENTS.md." ) + if self._settings.runtime.auto_compact: + return ( + "Context is over the hard limit even after compaction. " + "Run /clear, or shorten the request." + ) + return ( + "Context is over the hard limit and automatic compaction is off. " + "Run /compact (or /clear), or turn it back on with /compact auto on." ) - return False + + def _discard_last_user_message(self) -> None: + """Undo a user message that was recorded then refused by the hard limit.""" + if self._history and self._history[-1].role == "user": + self._history.pop() + if self._session is not None: + self._session.discard_last_message() + + def _force_digest_all_tools(self) -> int: + """Digest every tool result, including the normally protected tail.""" + changed = 0 + for index, message in enumerate(self._history): + if message.role != "tool": + continue + digest = _digest_text(message.content) + if digest != message.content: + self._history[index] = Message(role="tool", content=digest) + changed += 1 + return changed + + def _rollback_in_flight_turn(self) -> None: + """Drop from the last assistant message to the end (overflow / cancel).""" + for index in range(len(self._history) - 1, -1, -1): + if self._history[index].role == "assistant": + del self._history[index:] + if self._session is not None: + self._session.truncate_last_turn() + return def run_turn( self, @@ -549,16 +589,22 @@ def run_turn( ) return "" - if not self._settings.runtime.auto_compact and ( + # Compact existing history first when enabled, then preflight the + # incoming turn so a refused request never sticks in history/session. + if ( + self._settings.runtime.auto_compact + and self.estimated_prompt_tokens() > self.budget.compact_at_tokens + ): + adjusted = self.compact_now() + if adjusted: + self._ui.show_status(f"Compacted context: adjusted {adjusted} messages.") + projected = ( self.estimated_prompt_tokens() + estimate_tokens(text) + IMAGE_TOKEN_ESTIMATE * len(images) - > self.budget.hard_limit_tokens - ): - self._ui.show_status( - "Context is over the hard limit and automatic compaction is off. " - "Run /compact (or /clear), or turn it back on with /compact auto on." - ) + ) + if projected > self.budget.hard_limit_tokens: + self._ui.show_status(self._hard_limit_status()) return "" started = time.monotonic() @@ -570,6 +616,7 @@ def run_turn( self._audit.write("user_turn", **audit_kwargs) self._record(user(text, images=tuple(images))) if not self._ensure_under_hard_limit(): + self._discard_last_user_message() return "" content = self._tool_loop().content self._ui.turn_finished(self._turn_stats(time.monotonic() - started)) @@ -652,8 +699,23 @@ def _tool_loop(self) -> Message: last_reply = Message(role="assistant", content="") while True: - if not self._ensure_under_hard_limit(): - return last_reply + if ( + self._settings.runtime.auto_compact + and self.estimated_prompt_tokens() > self.budget.compact_at_tokens + ): + adjusted = self.compact_now() + if adjusted: + self._ui.show_status(f"Compacted context: adjusted {adjusted} messages.") + if self.estimated_prompt_tokens() > self.budget.hard_limit_tokens: + # In-flight tool results sit in the protected compaction window; + # force-digest them once, then retry. If still over, roll the + # incomplete assistant+tool exchange back so the next turn is not stuck. + if self._force_digest_all_tools() and self._settings.runtime.auto_compact: + self.compact_now() + if self.estimated_prompt_tokens() > self.budget.hard_limit_tokens: + self._ui.show_status(self._hard_limit_status()) + self._rollback_in_flight_turn() + return Message(role="assistant", content="") messages = [ Message(role="system", content=self._system_message_text()), *self._history, diff --git a/tests/test_budget.py b/tests/test_budget.py index 2187d21..0e89b21 100644 --- a/tests/test_budget.py +++ b/tests/test_budget.py @@ -45,6 +45,12 @@ def test_explicit_settings_win_over_detection() -> None: assert budget.reserved_response_tokens == 2222 +def test_explicit_context_setting_is_not_capped() -> None: + explicit = ContextSettings(model_context_tokens=131_072) + budget = resolve_budget(explicit, detected_context_tokens=None) + assert budget.model_context_tokens == 131_072 + + def test_estimate_tokens_rounds_up() -> None: assert estimate_tokens("") == 0 assert estimate_tokens("abcd") == 1 diff --git a/tests/test_conversation.py b/tests/test_conversation.py index 2cf80b2..165ea04 100644 --- a/tests/test_conversation.py +++ b/tests/test_conversation.py @@ -544,6 +544,18 @@ def test_load_truncate_last_turn_with_no_assistant_is_ignored(tmp_path: Path) -> assert loaded.messages[0].content == "hi" +def test_load_discard_last_message_pops_trailing(tmp_path: Path) -> None: + store = SessionStore(tmp_path / "sessions", "discard-unit") + store.record_message(Message(role="user", content="keep")) + store.record_message(Message(role="user", content="drop")) + store.discard_last_message() + + loaded = SessionStore.load(store.path) + + assert [m.role for m in loaded.messages] == ["user"] + assert loaded.messages[0].content == "keep" + + def test_plain_decline_after_step_skip_does_not_show_plan_pause_message(tmp_path: Path) -> None: registry, ran = _side_effect_registry() fake = FakeLLM( @@ -1075,6 +1087,77 @@ def test_ensure_under_hard_limit_blocks_when_compaction_cannot_recover( assert any("hard limit" in status.lower() for status in ui.statuses) +def test_hard_limit_refusal_discards_recorded_user_message(tmp_path: Path) -> None: + """A refused turn must not leave a stuck user message in history or on disk.""" + from shellpilot.config.model import RuntimeSettings + + settings = Settings( + context=ContextSettings(model_context_tokens=4096), + runtime=RuntimeSettings(auto_compact=True), + ) + session = SessionStore(tmp_path / "sessions", "hard-limit-user") + ui = FakeUI() + runtime = ConversationRuntime( + llm=FakeLLM(script=[answer("should not run")]), + settings=settings, + workspace=tmp_path, + behavior=BehaviorInstructions(global_text=None, project_text=None), + ui=ui, + session=session, + ) + # Sole user message sits just under the hard limit; any new turn tips it over + # and compaction cannot drop the only user message. + pad = "x" * 6600 + runtime.restore_history([Message(role="user", content=pad)]) + session.record_message(Message(role="user", content=pad)) + assert runtime.estimated_prompt_tokens() <= runtime.budget.hard_limit_tokens + before = list(runtime._history) + reply = runtime.run_turn("tip " * 50) + assert reply == "" + assert runtime._history == before + assert [m.content for m in SessionStore.load(session.path).messages] == [pad] + assert any("hard limit" in status.lower() for status in ui.statuses) + + +def test_mid_tool_loop_hard_limit_rolls_back_in_flight_turn(tmp_path: Path) -> None: + """Oversized in-flight assistant+tool exchanges must not stick past the hard limit.""" + from shellpilot.config.model import RuntimeSettings + from shellpilot.llm.messages import assistant + + (tmp_path / "blob.txt").write_text("ok") + settings = Settings( + context=ContextSettings(model_context_tokens=4096), + runtime=RuntimeSettings(auto_compact=True), + ) + session = SessionStore(tmp_path / "sessions", "hard-limit-tools") + # Huge assistant text is not force-digested (only tool results are), so the + # hard-limit gate must roll the incomplete exchange back out of history/session. + runtime = ConversationRuntime( + llm=FakeLLM( + script=[ + assistant( + "z" * 20_000, + tool_calls=(ToolCall(name="read_file", arguments={"path": "blob.txt"}),), + ), + answer("unreachable"), + ] + ), + settings=settings, + workspace=tmp_path, + behavior=BehaviorInstructions(global_text=None, project_text=None), + ui=FakeUI(), + session=session, + ) + + reply = runtime.run_turn("go") + assert reply == "" + assert [m.role for m in runtime._history] == ["user"] + assert runtime._history[0].content == "go" + loaded = SessionStore.load(session.path) + assert [m.role for m in loaded.messages] == ["user"] + assert loaded.messages[0].content == "go" + + def test_set_workspace_rebuilds_project_memory(tmp_path: Path) -> None: """Changing workspace (/cwd) rebuilds the project memory store for the new path, so the previous workspace's facts stop injecting (design section 16); From c794aa14d967f3547e4406f63e1eb0f520f36544 Mon Sep 17 00:00:00 2001 From: Lavindeep Dhillon Date: Fri, 10 Jul 2026 07:17:02 -0400 Subject: [PATCH 3/5] fix(runtime): keep hard-limit recovery behind auto_compact Force-digesting protected tool results is compaction; with auto_compact off, refuse and roll back the in-flight turn without rewriting earlier tool text. --- shellpilot/runtime/conversation.py | 17 +++++++---- tests/test_conversation.py | 49 ++++++++++++++++++++++++++++++ 2 files changed, 60 insertions(+), 6 deletions(-) diff --git a/shellpilot/runtime/conversation.py b/shellpilot/runtime/conversation.py index 8450143..bc7dd11 100644 --- a/shellpilot/runtime/conversation.py +++ b/shellpilot/runtime/conversation.py @@ -707,12 +707,17 @@ def _tool_loop(self) -> Message: if adjusted: self._ui.show_status(f"Compacted context: adjusted {adjusted} messages.") if self.estimated_prompt_tokens() > self.budget.hard_limit_tokens: - # In-flight tool results sit in the protected compaction window; - # force-digest them once, then retry. If still over, roll the - # incomplete assistant+tool exchange back so the next turn is not stuck. - if self._force_digest_all_tools() and self._settings.runtime.auto_compact: - self.compact_now() - if self.estimated_prompt_tokens() > self.budget.hard_limit_tokens: + # In-flight tool results sit in the protected compaction window. + # Only force-digest them when automatic compaction is on; with it + # off, refuse and roll back without silently rewriting history. + recovered = False + if self._settings.runtime.auto_compact: + if self._force_digest_all_tools(): + self.compact_now() + recovered = ( + self.estimated_prompt_tokens() <= self.budget.hard_limit_tokens + ) + if not recovered: self._ui.show_status(self._hard_limit_status()) self._rollback_in_flight_turn() return Message(role="assistant", content="") diff --git a/tests/test_conversation.py b/tests/test_conversation.py index 165ea04..e9cccb3 100644 --- a/tests/test_conversation.py +++ b/tests/test_conversation.py @@ -1158,6 +1158,55 @@ def test_mid_tool_loop_hard_limit_rolls_back_in_flight_turn(tmp_path: Path) -> N assert loaded.messages[0].content == "go" +def test_mid_tool_loop_hard_limit_with_auto_compact_off_does_not_digest( + tmp_path: Path, +) -> None: + """With auto_compact off, hard-limit refusal must not silently digest tool history.""" + from shellpilot.config.model import RuntimeSettings + + (tmp_path / "blob.txt").write_text("y" * 50_000) + settings = Settings( + context=ContextSettings(model_context_tokens=4096), + runtime=RuntimeSettings(auto_compact=False), + ) + ui = FakeUI() + # ~3639 tokens: under the hard limit, but any truncated tool result tips over. + prior_tool = ("line of prior tool output\n" * 250).rstrip("\n") + runtime = ConversationRuntime( + llm=FakeLLM( + script=[ + tool_call("read_file", path="blob.txt"), + answer("unreachable"), + ] + ), + settings=settings, + workspace=tmp_path, + behavior=BehaviorInstructions(global_text=None, project_text=None), + ui=ui, + ) + runtime.restore_history( + [ + Message(role="user", content="earlier"), + Message( + role="assistant", + content="", + tool_calls=(ToolCall(name="read_file", arguments={"path": "old.txt"}),), + ), + Message(role="tool", content=prior_tool), + ] + ) + assert runtime.estimated_prompt_tokens() <= runtime.budget.hard_limit_tokens + prior_tool_before = runtime._history[2].content + + reply = runtime.run_turn("go") + assert reply == "" + assert len(runtime._llm.calls) == 1 + assert [m.role for m in runtime._history] == ["user", "assistant", "tool", "user"] + assert runtime._history[2].content == prior_tool_before + assert "compacted" not in runtime._history[2].content + assert any("hard limit" in status.lower() for status in ui.statuses) + + def test_set_workspace_rebuilds_project_memory(tmp_path: Path) -> None: """Changing workspace (/cwd) rebuilds the project memory store for the new path, so the previous workspace's facts stop injecting (design section 16); From 941738e93264dceccd341bc5846e18536a1a5847 Mon Sep 17 00:00:00 2001 From: Lavindeep Dhillon Date: Fri, 10 Jul 2026 08:20:35 -0400 Subject: [PATCH 4/5] fix(runtime): drop the unused last_reply remnant and reformat The hard-limit stop returns an empty assistant Message directly, so the last_reply capture in the tool loop was never read (ruff F841). Remove it and apply the formatter to the hard-limit condition and the recovered check. --- shellpilot/runtime/conversation.py | 11 ++--------- tests/test_compaction.py | 4 +--- 2 files changed, 3 insertions(+), 12 deletions(-) diff --git a/shellpilot/runtime/conversation.py b/shellpilot/runtime/conversation.py index bc7dd11..9707a16 100644 --- a/shellpilot/runtime/conversation.py +++ b/shellpilot/runtime/conversation.py @@ -512,10 +512,7 @@ def _ensure_under_hard_limit(self) -> bool: def _hard_limit_status(self) -> str: """Status text when the hard context limit blocks a model call.""" - if ( - not self._history - and self.estimated_prompt_tokens() > self.budget.hard_limit_tokens - ): + if not self._history and self.estimated_prompt_tokens() > self.budget.hard_limit_tokens: return ( "System prompt and tool schemas alone exceed the hard limit. " "Raise context.model_context_tokens, or trim tools/AGENTS.md." @@ -696,7 +693,6 @@ def _tool_loop(self) -> Message: nudges_used = 0 empty_nudges_used = 0 consecutive_malformed = 0 - last_reply = Message(role="assistant", content="") while True: if ( @@ -714,9 +710,7 @@ def _tool_loop(self) -> Message: if self._settings.runtime.auto_compact: if self._force_digest_all_tools(): self.compact_now() - recovered = ( - self.estimated_prompt_tokens() <= self.budget.hard_limit_tokens - ) + recovered = self.estimated_prompt_tokens() <= self.budget.hard_limit_tokens if not recovered: self._ui.show_status(self._hard_limit_status()) self._rollback_in_flight_turn() @@ -760,7 +754,6 @@ def _tool_loop(self) -> Message: finally: self._ui.end_response() self._turn_output_tokens += reply.output_tokens - last_reply = reply # History length BEFORE this model step is recorded, so a mid-tool # cancel (below) can roll the step back out and leave no orphaned # tool_call behind (§31.15). diff --git a/tests/test_compaction.py b/tests/test_compaction.py index d9aeba2..485a685 100644 --- a/tests/test_compaction.py +++ b/tests/test_compaction.py @@ -107,9 +107,7 @@ def test_auto_compact_off_refuses_past_hard_limit(tmp_path: Path) -> None: tmp_path, context_tokens=4096, auto_compact=False, script=[answer("hi")] ) runtime._history.extend(Message(role="user", content="pad " * 600) for _ in range(4)) - assert ( - runtime.estimated_prompt_tokens() + 50 > runtime.budget.hard_limit_tokens - ) + assert runtime.estimated_prompt_tokens() + 50 > runtime.budget.hard_limit_tokens reply = runtime.run_turn("over the limit now") assert reply == "" assert fake.calls == [] # the model was never called From d37d4e866237e6d7b93ee20a9f32daa863527d0c Mon Sep 17 00:00:00 2001 From: Lavindeep Dhillon Date: Fri, 10 Jul 2026 08:20:42 -0400 Subject: [PATCH 5/5] docs: sync DESIGN with tool-schema budgeting and hard-limit recovery Document that the pre-turn and in-loop gates now count tool schemas and refuse (with compaction on or off), the bounded force-digest exception to the protected-tail compaction invariant, and the discard_last_message reconciliation record. --- docs/DESIGN.md | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/docs/DESIGN.md b/docs/DESIGN.md index 3be4310..34cae3d 100644 --- a/docs/DESIGN.md +++ b/docs/DESIGN.md @@ -515,7 +515,7 @@ hard_limit_tokens = floor(model_context_tokens * 0.90) `clamp(min, max, value)` means use `value` but never lower than `min` or higher than `max`. -When automatic compaction is off, a turn is refused if it would cross `hard_limit_tokens`. That pre-turn gate counts the estimated prompt (system prompt + history, including history images) plus this turn's incoming text **and** its incoming images (`IMAGE_TOKEN_ESTIMATE` per image), so an image-heavy turn cannot slip past a limit that text-only turns respect. The term is zero for text-only turns, leaving them unchanged. +A turn is refused if it would cross `hard_limit_tokens`. That pre-turn gate counts the full estimated request — the system prompt, the live profile's encoded tool schemas (`tool_schema_tokens`), and history (including history images) — plus this turn's incoming text **and** its incoming images (`IMAGE_TOKEN_ESTIMATE` per image), so an image-heavy turn cannot slip past a limit that text-only turns respect. The image term is zero for text-only turns, leaving them unchanged. With automatic compaction on, the runtime compacts first and refuses only when compaction cannot bring the projected prompt under the limit; with it off, it refuses directly. Counting the tool schemas closes a blind spot where a request that fit under the limit on paper still overran once the schemas were serialized to the model. Note the floor case: at the 8192-token fallback, after the system prompt, tool schemas, and behavior instructions, the working prompt budget is roughly 3-4k tokens. Small-context operation is a first-class mode, not a degraded one: shorter tool results, more aggressive truncation, and no long conversational tails. @@ -2075,7 +2075,9 @@ V1 compaction was deliberately simple: oldest-first truncation. Selective token- 2. Drop the oldest non-user messages outside the recent window. An assistant tool call takes its tool-result messages with it so no orphans confuse the model. 3. Last resort: drop the oldest user messages, always keeping the newest one. -No model call is involved — compaction is deterministic by design, matching the policy-first philosophy. Model-written summaries of dropped turns were considered and deliberately omitted. `/compact auto on|off` toggles automatic compaction (`[runtime] auto_compact`, default on); with it off, a turn that would exceed the hard limit is refused with guidance instead. +**Bounded exception — hard-limit recovery.** The protection of the in-flight exchange in pass 1 has one deliberate, narrow exception. When a request assembled *inside* the tool loop still exceeds `hard_limit_tokens` after ordinary compaction — a single in-flight tool result large enough to overrun on its own, sitting in the protected recent window — the runtime makes one last-resort recovery attempt before refusing the turn: it force-digests **every** tool result, including the protected tail (`_force_digest_all_tools`), then compacts again and proceeds only if that brought the request under the limit. This is safe for the same reason pass 1 is: a digest is a head/tail excerpt with an omission marker, and snapshot staleness checks force a fresh read before any write, so no correctness depends on the exact digested text surviving. It is gated behind `auto_compact` — it only runs when automatic compaction is on. With compaction off, the runtime never rewrites the protected tail; it refuses and rolls the turn back instead (below). + +No model call is involved — compaction is deterministic by design, matching the policy-first philosophy. Model-written summaries of dropped turns were considered and deliberately omitted. `/compact auto on|off` toggles automatic compaction (`[runtime] auto_compact`, default on). A turn that would exceed the hard limit is refused with guidance either way — the difference is the recovery attempt above: with compaction on the runtime tries the force-digest recovery first (and, on the in-loop gate, rolls the in-flight turn back before refusing); with it off it refuses directly. The in-loop gate mirrors the pre-turn gate: it re-checks the full request before every model call, so a turn that only overruns after several tool results accumulate is caught mid-loop, the in-flight assistant turn is rolled back (from the last assistant message to the end of history), and the turn ends cleanly rather than sending an over-limit request. Even simple truncation must preserve: @@ -2458,7 +2460,7 @@ The rebuild should stay light. The goal is a reliable local harness, not a frame | Memory system | Behavior/project memory, proposals, and optimization move to v2. V1 only reads `AGENTS.md`. Scheduled for v0.3.0 (settled 2026-06-11). | | Token-budget compaction | V1 uses oldest-first truncation; selective compaction is v2. Scheduled for v0.3.0 (settled 2026-06-11). | | `trusted-local` profile | Deferred from v1, and deferred again at the 2026-06-11 v2 scoping. Revisit for v3. | -| Session resume | Shipped in v0.3.0 (settled 2026-06-11): append-only JSONL transcripts at `.shellpilot/sessions/.jsonl`, written incrementally with secrets redacted; compaction trims memory, never the transcript. `shellpilot --resume [id]` restores the latest (or named) session's history; snapshots are never restored, so read-before-write forces fresh reads. `/export` renders the transcript to markdown. Tool-call arguments are redacted recursively (matching the audit log's `_redact_value` logic, now unified in `redact_structure` in `shellpilot/memory/redaction.py`) before they reach the JSONL transcript; `/export` inherits redaction by re-reading the transcript from disk. Fixed in v0.5.2. `session_markdown` re-applies redaction at export time so transcripts written before v0.5.2 (which may contain raw secrets on disk) cannot leak through `/export`; on-disk history is deliberately left untouched. Fixed in v0.5.2 review wave. Plan state now also restores on `--resume` (v0.6.0): an `active_plan` pointer in the transcript is read at boot; if the referenced plan sidecar is live (`proposed`/`active`/`blocked`), `PlanManager.restore` reinstates it (section 11.3). **Read-side traversal guard (v0.10.1):** `SessionStore.find` now rejects any session id whose resolved parent differs from the sessions directory, closing the `--resume ../../../../etc/x` path-traversal vector; the write path was already safe via `path.stem`. **Reconciliation records:** the transcript stays append-only, so mid-turn corrections are records rather than rewrites — on load, `replace_last_message` replaces the last *assistant* record (a mid-batch decline truncates the reply's remaining tool calls, section 14.6) and `truncate_last_turn` deletes from the last assistant record to the end (a mid-tool cancel, section 31.15); a record of either kind with no assistant message present, or an unknown record kind, is ignored. | +| Session resume | Shipped in v0.3.0 (settled 2026-06-11): append-only JSONL transcripts at `.shellpilot/sessions/.jsonl`, written incrementally with secrets redacted; compaction trims memory, never the transcript. `shellpilot --resume [id]` restores the latest (or named) session's history; snapshots are never restored, so read-before-write forces fresh reads. `/export` renders the transcript to markdown. Tool-call arguments are redacted recursively (matching the audit log's `_redact_value` logic, now unified in `redact_structure` in `shellpilot/memory/redaction.py`) before they reach the JSONL transcript; `/export` inherits redaction by re-reading the transcript from disk. Fixed in v0.5.2. `session_markdown` re-applies redaction at export time so transcripts written before v0.5.2 (which may contain raw secrets on disk) cannot leak through `/export`; on-disk history is deliberately left untouched. Fixed in v0.5.2 review wave. Plan state now also restores on `--resume` (v0.6.0): an `active_plan` pointer in the transcript is read at boot; if the referenced plan sidecar is live (`proposed`/`active`/`blocked`), `PlanManager.restore` reinstates it (section 11.3). **Read-side traversal guard (v0.10.1):** `SessionStore.find` now rejects any session id whose resolved parent differs from the sessions directory, closing the `--resume ../../../../etc/x` path-traversal vector; the write path was already safe via `path.stem`. **Reconciliation records:** the transcript stays append-only, so mid-turn corrections are records rather than rewrites — on load, `replace_last_message` replaces the last *assistant* record (a mid-batch decline truncates the reply's remaining tool calls, section 14.6), `truncate_last_turn` deletes from the last assistant record to the end (a mid-tool cancel, section 31.15), and `discard_last_message` pops the single trailing record (a user message that was written to the transcript and then refused by the hard context-limit gate, so `--resume` does not restore a stuck user turn with no reply). A `replace_last_message`/`truncate_last_turn` record with no assistant message present, a `discard_last_message` with an empty transcript, or an unknown record kind, is ignored. | | Agent raw shell | Do not expose `raw_shell` as an agent tool in v1. Keep Manual Shell for direct user-controlled `shell=True`. | | Capability packs (Skills v2) | v0.6.0 shipped instruction-only SKILL.md discovery; v0.7.0 extends it with deterministic trigger selection, four markdown-only builtins, read-only references/templates, script manifest discovery without execution, and enriched `/skills` + `/context` visibility (section 23). | | Capability packs (heavier: tools/handlers/permissions) | Design later after core tools are stable. v3 candidate (2026-06-11). |