diff --git a/libs/code/deepagents_code/app.py b/libs/code/deepagents_code/app.py index be21d63fd4c..56cdc42f413 100644 --- a/libs/code/deepagents_code/app.py +++ b/libs/code/deepagents_code/app.py @@ -354,6 +354,20 @@ def _coerce_provisional_cost_delta_usd(value: object) -> float | None: return delta_usd +def _persisted_context_tokens(state_values: dict[str, Any]) -> int: + """Read a checkpoint's `_context_tokens`, coercing absent or invalid to 0. + + Args: + state_values: Thread state values from the checkpoint. + + Returns: + The persisted context size, or 0 when it is missing or not a + non-negative `int`. + """ + raw = state_values.get("_context_tokens") + return raw if isinstance(raw, int) and raw >= 0 else 0 + + def _warn_discarded_goal_channels(state_values: dict[str, Any]) -> list[str]: """Report persisted goal/rubric channels that are present but malformed. @@ -9337,6 +9351,8 @@ async def _run_session_start_sequence(self) -> None: ), ) return + if start_cause is SessionStartCause.RESUME: + await self._restore_resumed_context_tokens() if not await self._run_session_start_hook(start_cause): self._initial_session_start_stopped = True @@ -9349,6 +9365,9 @@ async def _run_session_start_sequence(self) -> None: self._startup_cmd = None await self._run_startup_command(cmd) + if start_cause is SessionStartCause.RESUME: + await self._maybe_compact_after_resume() + if should_load_history: await self._remount_pending_goal_rubric_review() @@ -16502,10 +16521,7 @@ async def _fetch_thread_history_data(self, thread_id: str) -> _ThreadHistoryPayl context-token count, and the persisted model spec (if any). """ state_values = await self._get_thread_state_values(thread_id) - raw_tokens = state_values.get("_context_tokens") - context_tokens = ( - raw_tokens if isinstance(raw_tokens, int) and raw_tokens >= 0 else 0 - ) + context_tokens = _persisted_context_tokens(state_values) raw_spec = state_values.get("_model_spec") model_spec = raw_spec if isinstance(raw_spec, str) else "" raw_params = state_values.get("_model_params") @@ -20937,6 +20953,9 @@ def _build_agent(url: str) -> Any: # noqa: ANN401 # see docstring ) await self._mount_message(AppMessage(resume_hint)) + if resume_thread_id is not None: + await self._maybe_compact_after_resume() + # Drain any messages the user typed after we cleared the queue # but before the new server was ready. if self._pending_messages and not self._agent_running: @@ -24219,18 +24238,44 @@ async def _handle_threads_command(self, command: str) -> None: if target is None: return - active_agent = self._assistant_id or DEFAULT_ASSISTANT_ID - if target.agent_name == active_agent: - await self._resume_thread(target.thread_id) - return - # This command runs inline on Textual's message pump. Detach before - # awaiting the confirmation modal so Enter/Esc can reach it. + # awaiting the cwd/compaction modals so Enter/Esc can reach them. Both + # continuations mount their own failure messages. + active_agent = self._assistant_id or DEFAULT_ASSISTANT_ID self._schedule_off_message_pump( - self._confirm_then_resume_cross_agent_thread(target), + self._resume_same_agent_thread(target.thread_id) + if target.agent_name == active_agent + else self._confirm_then_resume_cross_agent_thread(target), context=f"threads-resume:{target.thread_id}", ) + async def _resume_same_agent_thread(self, thread_id: str) -> None: + """Resume a thread already owned by the active agent. + + A thin reporting wrapper: `_resume_thread` handles its own failures once + the switch is under way, but its pre-switch phase (the cwd prompt, the + compaction prompt, and their modal setup) sits outside that boundary. + Detached continuations only get logged, so without this a `/threads -r` + would end with no visible outcome. + + Args: + thread_id: The thread to resume on the current agent. + + Raises: + asyncio.CancelledError: If app shutdown cancels the detached flow. + """ + try: + await self._resume_thread(thread_id) + except asyncio.CancelledError: + raise + except Exception: + logger.exception("Same-agent resume failed for thread %s", thread_id) + await self._mount_message( + ErrorMessage( + f"Could not resume thread {thread_id}. Use /threads to try again." + ) + ) + async def _resolve_threads_resume_target( self, requested_id: str | None ) -> _ThreadsResumeTarget | None: @@ -24942,6 +24987,69 @@ async def _preview_project_settings_change(cwd: Path) -> bool: return False return bool(changes) + async def _restore_resumed_context_tokens(self) -> None: + """Seed `_context_tokens` from the checkpoint without loading history. + + A launch that resumes *and* carries an initial prompt skips the history + load, so nothing would populate the context size — and that path is the + one where it matters most, since the prompt is submitted against the + full restored context moments later. + """ + if not self._lc_thread_id or not self._agent: + return + try: + state_values = await self._get_thread_state_values(self._lc_thread_id) + except Exception: + logger.warning( + "Could not read the context size for resumed thread %s; " + "skipping the compaction suggestion", + self._lc_thread_id, + exc_info=True, + ) + return + tokens = _persisted_context_tokens(state_values) + if tokens > 0: + self._on_tokens_update(tokens) + + async def _maybe_compact_after_resume(self) -> None: + """Offer to compact a just-resumed thread with an oversized context. + + Resuming costs nothing on its own, so the prompt runs once the thread is + live and before the user's next turn pays for the restored context. + Declining is free: `/compact` remains available for the rest of the + session. + """ + from deepagents_code.config_manifest import ( + COMPACT_ON_RESUME_THRESHOLD_DEFAULT, + get_option, + load_config_toml, + resolve_scalar, + ) + + threshold = COMPACT_ON_RESUME_THRESHOLD_DEFAULT + option = get_option("threads.compact_on_resume_threshold") + if option is not None: + resolved, _ = resolve_scalar(option, toml_data=load_config_toml()) + if isinstance(resolved, int): + threshold = resolved + if threshold <= 0 or self._context_tokens <= threshold: + return + + from deepagents_code.tui.modals.resume_compact import ResumeCompactPromptScreen + + try: + compact = await self._push_screen_wait( + ResumeCompactPromptScreen( + context_tokens=self._context_tokens, + threshold=threshold, + ) + ) + except Exception: + logger.exception("Failed to show the compact-on-resume prompt") + return + if compact: + await self._handle_offload() + async def _offer_thread_cwd_switch( self, thread_id: str, @@ -25239,6 +25347,12 @@ async def _resume_thread(self, thread_id: str) -> None: # or failed load leaves the hint under something else. await self._mount_previous_thread_hint(prev_session_thread) + # Landing on a new thread re-arms the same-thread toast, so stepping + # back to a thread and re-selecting it announces itself again. + self._last_thread_unchanged = None + if await self._run_session_start_hook(SessionStartCause.RESUME): + await self._maybe_compact_after_resume() + # Deferred above so it lands last, keeping the interactive prompt # at the bottom of the transcript. Guarded because this runs inside # the rollback handler's `try`: failing to restore a review must @@ -25247,12 +25361,6 @@ async def _resume_thread(self, thread_id: str) -> None: await self._remount_pending_goal_rubric_review() except Exception: logger.exception("Failed to restore pending goal review") - - # Landing on a new thread re-arms the same-thread toast, so stepping - # back to a thread and re-selecting it announces itself again. - self._last_thread_unchanged = None - if not await self._run_session_start_hook(SessionStartCause.RESUME): - return except Exception as exc: if prefetched_payload is None: logger.exception("Failed to prefetch history for thread %s", thread_id) diff --git a/libs/code/deepagents_code/config_manifest.py b/libs/code/deepagents_code/config_manifest.py index ddf8c8fb7c2..3822ec49422 100644 --- a/libs/code/deepagents_code/config_manifest.py +++ b/libs/code/deepagents_code/config_manifest.py @@ -108,6 +108,12 @@ and falls through to the next layer / default. """ +COMPACT_ON_RESUME_THRESHOLD_DEFAULT = 400_000 +"""Context size above which a resumed thread is offered compaction. + +Zero or negative disables the suggestion. +""" + LANGSMITH_PROJECT_DEFAULT = "deepagents-code" """Project agent traces fall back to when no project env var is set. @@ -1472,6 +1478,16 @@ def _credential_options() -> tuple[ConfigOption, ...]: settings_field="interpreter_ptc_acknowledge_unsafe", ), # --- Threads (config.toml-only; structured column table excepted) --- + ConfigOption( + key="threads.compact_on_resume_threshold", + group="Threads", + summary=( + "Offer to compact a resumed thread above this context size (0 disables)." + ), + kind=OptionKind.INT, + default=COMPACT_ON_RESUME_THRESHOLD_DEFAULT, + toml_keys=("threads", "compact_on_resume_threshold"), + ), ConfigOption( key="threads.relative_time", group="Threads", diff --git a/libs/code/deepagents_code/tui/modals/resume_compact.py b/libs/code/deepagents_code/tui/modals/resume_compact.py new file mode 100644 index 00000000000..58e7cbce328 --- /dev/null +++ b/libs/code/deepagents_code/tui/modals/resume_compact.py @@ -0,0 +1,121 @@ +"""Prompt for compacting a large resumed thread.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, ClassVar + +from textual.binding import Binding, BindingType +from textual.containers import Vertical +from textual.screen import ModalScreen +from textual.widgets import Static + +from deepagents_code._session_stats import format_token_count +from deepagents_code.config import get_glyphs + +if TYPE_CHECKING: + from textual.app import ComposeResult + + +class ResumeCompactPromptScreen(ModalScreen[bool]): + """Ask whether to compact a just-resumed thread before the next turn.""" + + can_focus = True + + BINDINGS: ClassVar[list[BindingType]] = [ + Binding("enter", "compact", "Compact", show=False, priority=True), + Binding("escape", "skip", "Skip", show=False, priority=True), + Binding( + "ctrl+c", + "quit_or_interrupt", + "Quit/Interrupt", + show=False, + priority=True, + ), + Binding("ctrl+d", "quit_app", "Quit", show=False, priority=True), + ] + + CSS = """ + ResumeCompactPromptScreen { + align: center middle; + } + + ResumeCompactPromptScreen > Vertical { + width: 72; + max-width: 90%; + height: auto; + background: $surface; + border: solid $warning; + padding: 1 2; + } + + ResumeCompactPromptScreen .resume-compact-title { + text-style: bold; + color: $warning; + text-align: center; + margin-bottom: 1; + } + + ResumeCompactPromptScreen .resume-compact-body { + height: auto; + color: $text; + margin-bottom: 1; + } + + ResumeCompactPromptScreen .resume-compact-help { + height: auto; + color: $text-muted; + text-style: italic; + text-align: center; + } + """ + + def __init__(self, *, context_tokens: int, threshold: int) -> None: + """Initialize the prompt. + + Args: + context_tokens: Latest model-reported context size for the thread. + threshold: Configured token count that triggered the suggestion. + """ + super().__init__() + self._context_tokens = context_tokens + self._threshold = threshold + + def compose(self) -> ComposeResult: + """Compose the confirmation dialog. + + Yields: + Title, explanation, and keyboard help. + """ + with Vertical(): + yield Static( + "Compact this thread?", + classes="resume-compact-title", + markup=False, + ) + yield Static( + f"This thread uses {format_token_count(self._context_tokens)} context " + "tokens, above the configured " + f"{format_token_count(self._threshold)} token threshold. Compacting " + "summarizes older messages so later turns cost less.", + classes="resume-compact-body", + markup=False, + ) + yield Static( + f" {get_glyphs().bullet} ".join( + ("Enter: compact now", "Esc: keep full context") + ), + classes="resume-compact-help", + markup=False, + ) + + def on_mount(self) -> None: + """Focus the modal so its bindings receive keyboard input.""" + self.focus() + + def action_compact(self) -> None: + """Compact the thread before the next turn.""" + self.dismiss(True) + + def action_skip(self) -> None: + """Leave the thread's context untouched.""" + self.dismiss(False) diff --git a/libs/code/tests/unit_tests/test_threads_resume.py b/libs/code/tests/unit_tests/test_threads_resume.py index 366c469e737..40ce00e32c3 100644 --- a/libs/code/tests/unit_tests/test_threads_resume.py +++ b/libs/code/tests/unit_tests/test_threads_resume.py @@ -61,7 +61,9 @@ async def test_resume_flag_resolves_and_resumes(self) -> None: app = _make_app() target = _ThreadsResumeTarget("thread-x", "agent") app._resolve_threads_resume_target = AsyncMock(return_value=target) # ty: ignore - await app._handle_threads_command("/threads -r") + with patch.object(app, "_schedule_off_message_pump") as schedule: + await app._handle_threads_command("/threads -r") + await schedule.call_args.args[0] app._resolve_threads_resume_target.assert_awaited_once_with(None) # ty: ignore app._resume_thread.assert_awaited_once_with("thread-x") # ty: ignore app._show_thread_selector.assert_not_awaited() # ty: ignore @@ -70,7 +72,9 @@ async def test_resume_specific_id(self) -> None: app = _make_app() target = _ThreadsResumeTarget("abc", "agent") app._resolve_threads_resume_target = AsyncMock(return_value=target) # ty: ignore - await app._handle_threads_command("/threads -r abc") + with patch.object(app, "_schedule_off_message_pump") as schedule: + await app._handle_threads_command("/threads -r abc") + await schedule.call_args.args[0] app._resolve_threads_resume_target.assert_awaited_once_with("abc") # ty: ignore app._resume_thread.assert_awaited_once_with("abc") # ty: ignore