langchain-ai/deepagents
1
0
Fork 0
mirror of https://github.com/langchain-ai/deepagents.git synced 2026-07-22 01:35:28 -04:00
Code Issues 169 Packages Projects Releases 16 Wiki Activity

fix(cli): show pending token placeholder during streaming and restore on early exit (#3288)

Browse Source 
The status bar previously cleared the token count entirely during
streaming, leaving an empty gap that looked like a UI glitch. Switch to
showing a `"... tokens"` placeholder so users know a count is coming
back, and ensure the accurate count is reported even when a turn ends
early (e.g. on HITL rejection or cancellation).

## Changes

- **Replaced `hide_tokens` with `show_pending_tokens`** across
`StatusBar`, `TextualUIAdapter`, and `DeepAgentsApp`. The new method
renders `"... tokens"` instead of blanking the display, making the
pending state explicit.
- **Persist token counts on early-exit paths** — added
`_report_and_persist_tokens` calls in `execute_task_textual` before
returning for HITL rejection and `ask_user` cancellation, so the last
known count is restored to the status bar instead of staying in the
pending state.
- **Kept `set_tokens` refresh semantics** — documented and preserved the
forced-re-render behavior when the reactive value is unchanged, which is
required after the placeholder swaps out the widget text.
- **Added/updated unit tests** for placeholder display, early-exit token
restoration, and no-op safety when the widget is not yet mounted.
This commit is contained in:
Mason Daugherty
2026-05-10 16:02:05 -04:00
committed by GitHub
parent c1dd6c8a0f
commit 882232dbc4
5 changed files with 143 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/cli/deepagents_cli/app.py
+4 -4
View File
@@ -1553,7 +1553,7 @@ class DeepAgentsApp(App):
)
# Wire token display callbacks
self._ui_adapter._on_tokens_update = self._on_tokens_update
self._ui_adapter._on_tokens_hide = self._hide_tokens
self._ui_adapter._on_tokens_pending = self._show_pending_tokens
self._ui_adapter._on_tokens_show = self._show_tokens
# Fire-and-forget workers — none of these block the event loop.
@@ -2751,10 +2751,10 @@ class DeepAgentsApp(App):
approximate=self._tokens_approximate,
)
def _hide_tokens(self) -> None:
"""Hide the token display during streaming."""
def _show_pending_tokens(self) -> None:
"""Show the unknown token count placeholder during streaming."""
if self._status_bar:
self._status_bar.hide_tokens()
self._status_bar.show_pending_tokens()
def _check_hydration_needed(self) -> None:
"""Check if we need to hydrate messages from the store.
libs/cli/deepagents_cli/textual_adapter.py
+16 -8
View File
@@ -280,8 +280,8 @@ class TextualUIAdapter:
self._on_tokens_update: _TokensUpdateCallback | None = None
"""Called with total context tokens after each LLM response."""
self._on_tokens_hide: Callable[[], None] | None = None
"""Called to hide the token display during streaming."""
self._on_tokens_pending: Callable[[], None] | None = None
"""Called to show an unknown token count during streaming."""
self._on_tokens_show: _TokensShowCallback | None = None
"""Called to restore the token display with the cached value."""
@@ -478,21 +478,21 @@ async def execute_task_textual(
# should be set together to avoid inconsistent status-bar behavior.
token_cbs = (
adapter._on_tokens_update,
adapter._on_tokens_hide,
adapter._on_tokens_pending,
adapter._on_tokens_show,
)
if any(token_cbs) and not all(token_cbs):
logger.warning(
"Token callbacks partially wired (update=%s, hide=%s, show=%s); "
"Token callbacks partially wired (update=%s, pending=%s, show=%s); "
"token display may behave inconsistently",
adapter._on_tokens_update is not None,
adapter._on_tokens_hide is not None,
adapter._on_tokens_pending is not None,
adapter._on_tokens_show is not None,
)
# Hide token display during streaming (will be shown with accurate count at end)
if adapter._on_tokens_hide:
adapter._on_tokens_hide()
# Show unknown token count during streaming; the accurate count arrives at turn end.
if adapter._on_tokens_pending:
adapter._on_tokens_pending()
file_op_tracker = FileOpTracker(assistant_id=assistant_id, backend=backend)
displayed_tool_ids: set[str] = set()
@@ -1252,6 +1252,14 @@ async def execute_task_textual(
)
await adapter._mount_message(AppMessage(message))
turn_stats.wall_time_seconds = time.monotonic() - start_time
await _report_and_persist_tokens(
adapter,
agent,
config,
captured_input_tokens,
captured_output_tokens,
shield=True,
)
return turn_stats
stream_input = Command(resume=resume_payload)
libs/cli/deepagents_cli/widgets/status.py
+7 -6
View File
@@ -406,9 +406,10 @@ class StatusBar(Horizontal):
def set_tokens(self, count: int, *, approximate: bool = False) -> None:
"""Set the token count.
Forces a display refresh even when the value is unchanged, because
`hide_tokens` clears the widget text without updating the reactive
attribute.
Forces a display refresh even when the value is unchanged. During
streaming, `show_pending_tokens` replaces the widget text without
changing the reactive token value, so a later update with the same
count still needs to re-render the exact count.
Args:
count: Current context token count.
@@ -423,10 +424,10 @@ class StatusBar(Horizontal):
# self._approximate for the suffix.
self.tokens = count
def hide_tokens(self) -> None:
"""Hide the token display (e.g., during streaming)."""
def show_pending_tokens(self) -> None:
"""Show an unknown token count placeholder during streaming."""
try:
self.query_one("#tokens-display", Static).update("")
self.query_one("#tokens-display", Static).update("... tokens")
except NoMatches:
return
libs/cli/tests/unit_tests/test_status.py
+38 -9
View File
@@ -195,27 +195,36 @@ class TestTokenDisplay:
display = pilot.app.query_one("#tokens-display")
assert "5.0K" in str(display.render())
async def test_hide_tokens_clears_display(self) -> None:
async def test_show_pending_tokens_shows_unknown_placeholder(self) -> None:
async with StatusBarApp().run_test() as pilot:
bar = pilot.app.query_one("#status-bar", StatusBar)
bar.set_tokens(5000)
await pilot.pause()
bar.hide_tokens()
bar.show_pending_tokens()
await pilot.pause()
display = pilot.app.query_one("#tokens-display")
assert str(display.render()) == ""
assert str(display.render()) == "... tokens"
async def test_set_tokens_after_hide_restores_display(self) -> None:
async def test_show_pending_tokens_before_count_shows_placeholder(self) -> None:
async with StatusBarApp().run_test() as pilot:
bar = pilot.app.query_one("#status-bar", StatusBar)
bar.show_pending_tokens()
await pilot.pause()
display = pilot.app.query_one("#tokens-display")
assert str(display.render()) == "... tokens"
async def test_set_tokens_after_pending_restores_display(self) -> None:
"""Regression: set_tokens must refresh even when value is unchanged.
hide_tokens clears the widget text without updating the reactive,
so a subsequent set_tokens with the same value must still re-render.
`show_pending_tokens` replaces the widget text without updating the
reactive value, so a subsequent `set_tokens` with the same count must
still re-render.
"""
async with StatusBarApp().run_test() as pilot:
bar = pilot.app.query_one("#status-bar", StatusBar)
bar.set_tokens(5000)
await pilot.pause()
bar.hide_tokens()
bar.show_pending_tokens()
await pilot.pause()
# Same value — previously skipped by reactive dedup
bar.set_tokens(5000)
@@ -223,6 +232,26 @@ class TestTokenDisplay:
display = pilot.app.query_one("#tokens-display")
assert "5.0K" in str(display.render())
async def test_show_pending_tokens_after_count_change_keeps_placeholder(
self,
) -> None:
async with StatusBarApp().run_test() as pilot:
bar = pilot.app.query_one("#status-bar", StatusBar)
bar.set_tokens(5000)
await pilot.pause()
bar.show_pending_tokens()
await pilot.pause()
bar.set_tokens(7500)
await pilot.pause()
bar.show_pending_tokens()
await pilot.pause()
display = pilot.app.query_one("#tokens-display")
assert str(display.render()) == "... tokens"
def test_show_pending_tokens_without_mount_is_noop(self) -> None:
bar = StatusBar()
bar.show_pending_tokens()
async def test_approximate_appends_plus(self) -> None:
"""approximate=True should append '+' to the token count."""
async with StatusBarApp().run_test() as pilot:
@@ -233,13 +262,13 @@ class TestTokenDisplay:
rendered = str(display.render())
assert "5.0K+" in rendered
async def test_approximate_after_hide_restores_with_plus(self) -> None:
async def test_approximate_after_pending_restores_with_plus(self) -> None:
"""Interrupted restore: same value + approximate should show count with '+'."""
async with StatusBarApp().run_test() as pilot:
bar = pilot.app.query_one("#status-bar", StatusBar)
bar.set_tokens(5000)
await pilot.pause()
bar.hide_tokens()
bar.show_pending_tokens()
await pilot.pause()
bar.set_tokens(5000, approximate=True)
await pilot.pause()
libs/cli/tests/unit_tests/test_textual_adapter.py
+78 -4
View File
@@ -93,7 +93,7 @@ class TestTextualUIAdapterInit:
request_approval=_mock_approval,
)
assert adapter._on_tokens_update is None
assert adapter._on_tokens_hide is None
assert adapter._on_tokens_pending is None
assert adapter._on_tokens_show is None
def test_on_tool_complete_defaults_to_none_and_accepts_callback(self) -> None:
@@ -125,17 +125,17 @@ class TestTextualUIAdapterInit:
def update_cb(count: int, *, approximate: bool = False) -> None:
pass
def hide_cb() -> None:
def pending_cb() -> None:
pass
def show_cb(*, approximate: bool = False) -> None:
pass
adapter._on_tokens_update = update_cb
adapter._on_tokens_hide = hide_cb
adapter._on_tokens_pending = pending_cb
adapter._on_tokens_show = show_cb
assert adapter._on_tokens_update is update_cb
assert adapter._on_tokens_hide is hide_cb
assert adapter._on_tokens_pending is pending_cb
assert adapter._on_tokens_show is show_cb
def test_finalize_pending_tools_with_error_marks_and_clears(self) -> None:
@@ -477,6 +477,12 @@ def _ask_user_interrupt_chunk(payload: dict[str, Any]) -> tuple[Any, ...]:
return ((), "updates", {"__interrupt__": [interrupt]})
def _hitl_interrupt_chunk(payload: dict[str, Any]) -> tuple[Any, ...]:
"""Build an updates-stream chunk containing one HITL interrupt."""
interrupt = SimpleNamespace(id="interrupt-1", value=payload)
return ((), "updates", {"__interrupt__": [interrupt]})
class TestExecuteTaskTextualSummarizationFeedback:
"""Tests for summarization spinner and notification feedback."""
@@ -1239,6 +1245,7 @@ class TestExecuteTaskTextualAskUser:
async def test_ask_user_cancelled_marks_row_rejected_and_halts(self) -> None:
"""Cancelled result should reject the row and not resume generation."""
mounted: list[object] = []
token_events: list[str] = []
future: asyncio.Future[object] = asyncio.Future()
future.set_result({"type": "cancelled"})
@@ -1272,6 +1279,10 @@ class TestExecuteTaskTextualAskUser:
request_approval=_mock_approval,
request_ask_user=request_ask_user,
)
adapter._on_tokens_pending = lambda: token_events.append("pending")
adapter._on_tokens_show = lambda *, approximate=False: token_events.append(
f"show:{approximate}"
)
await execute_task_textual(
user_input="hello",
@@ -1286,6 +1297,69 @@ class TestExecuteTaskTextualAskUser:
app_messages = [widget for widget in mounted if isinstance(widget, AppMessage)]
assert len(app_messages) == 1
assert "Question cancelled" in str(app_messages[0]._content)
assert token_events == ["pending", "show:False"]
async def test_hitl_rejection_restores_token_display_before_halt(self) -> None:
"""Rejected approval should restore tokens before returning early."""
mounted: list[object] = []
token_events: list[str] = []
future: asyncio.Future[object] = asyncio.Future()
future.set_result({"type": "reject"})
async def mount_message(widget: object) -> None:
await asyncio.sleep(0)
mounted.append(widget)
async def request_approval(
_action_requests: list[dict[str, Any]],
_assistant_id: str | None,
) -> asyncio.Future[object]:
await asyncio.sleep(0)
return future
agent = _SequencedAgent(
streams_by_call=[
[
_hitl_interrupt_chunk(
{
"action_requests": [
{"name": "read_file", "args": {"path": "notes.txt"}}
],
"review_configs": [
{
"action_name": "read_file",
"allowed_decisions": ["approve", "reject"],
}
],
}
)
],
[],
]
)
adapter = TextualUIAdapter(
mount_message=mount_message,
update_status=_noop_status,
request_approval=request_approval,
)
adapter._on_tokens_pending = lambda: token_events.append("pending")
adapter._on_tokens_show = lambda *, approximate=False: token_events.append(
f"show:{approximate}"
)
await execute_task_textual(
user_input="hello",
agent=agent,
assistant_id="assistant",
session_state=SimpleNamespace(thread_id="thread-1", auto_approve=False),
adapter=adapter,
)
assert len(agent.stream_inputs) == 1
app_messages = [widget for widget in mounted if isinstance(widget, AppMessage)]
assert len(app_messages) == 1
assert "Command rejected" in str(app_messages[0]._content)
assert token_events == ["pending", "show:False"]
async def test_ask_user_invalid_answers_payload_marks_row_error(self) -> None:
"""Non-list answers should mark row as error and pop it."""