langchain-ai/deepagents
1
0
Fork 0
mirror of https://github.com/langchain-ai/deepagents.git synced 2026-07-21 17:25:26 -04:00
Code Issues 169 Packages Projects Releases 16 Wiki Activity

feat(code): offer restart after saving Tavily key via /auth (#4560)

Browse Source 
Saving a Tavily API key via `/auth` now offers to restart the agent
server so web search activates without relaunching.

---

`web_search` is bound only when Tavily is configured at server spawn
time (`server_graph._build_tools`), so a Tavily key added to an
already-running server never activates web search until the server
respawns. Previously `/auth` only auto-resumed the server when a saved
key unblocked a credentials-blocked *model* startup, so saving a Tavily
key mid-session silently did nothing.

Now, saving a Tavily key in the `/auth` manager arms an optional restart
prompt that fires once the manager closes (deferrable with Esc),
mirroring the post-install restart offer.
`AuthManagerScreen.CredentialSaved` carries the saved service so the app
can react to spawn-time-gating credentials, `RestartPromptScreen` gained
`verb`/`body` overrides, and the post-install restart offer was
refactored into a shared `_offer_server_restart` reused by both flows.
The stored key is exported to the environment eagerly so a later
`/restart` (offered or manual) picks it up on reload.

Made by [Open
SWE](https://openswe.vercel.app/agents/95be92ba-5015-b080-9309-3b7a9f8bf85f)

---------

Co-authored-by: open-swe[bot] <open-swe@users.noreply.github.com>
This commit is contained in:
Mason Daugherty
2026-07-08 21:13:59 -04:00
committed by GitHub
parent 1110497e17
commit 12df81ad00
8 changed files with 1037 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/code/deepagents_code/app.py
+238 -34
View File
@@ -364,6 +364,7 @@ if TYPE_CHECKING:
from deepagents_code.tui.textual_adapter import TextualUIAdapter
from deepagents_code.tui.widgets.approval import ApprovalMenu
from deepagents_code.tui.widgets.ask_user import AskUserMenu
from deepagents_code.tui.widgets.auth import AuthManagerScreen
from deepagents_code.tui.widgets.goal_review import GoalReviewMenu, GoalReviewResult
from deepagents_code.tui.widgets.model_selector import ModelSelectorScreen
from deepagents_code.tui.widgets.notification_center import (
@@ -2250,6 +2251,22 @@ class DeepAgentsApp(App):
self._pending_mcp_login_reconnect = False
"""Whether a successful MCP login is waiting for reconnect."""
self._pending_web_search_restart = False
"""Whether a Tavily key saved via `/auth` is awaiting an offered restart.
`web_search` is bound only when Tavily is configured at server spawn
time (see `server_graph._build_tools`), so a key added to a running
server takes effect only after a respawn. Set when the saved key gates
a tool the running server lacks; consumed when the `/auth` manager
closes so the restart prompt never stacks over the still-open manager.
"""
self._auth_exported_tavily_original: str | None = None
"""Original shell `TAVILY_API_KEY` value before `/auth` exported Tavily."""
self._auth_exported_tavily = False
"""Whether this app process exported `TAVILY_API_KEY` from `/auth`."""
self._pending_mcp_disable_reconnect_servers: set[str] = set()
"""MCP servers with disable-state changes waiting for reconnect."""
@@ -13479,7 +13496,12 @@ class DeepAgentsApp(App):
self._chat_input.focus_input()
# When the user selected a greyed-out (uninstalled) provider and
# confirmed installing it, install the extra and reopen the manager
# so they can add a key against the now-installed provider.
# so they can add a key against the now-installed provider. A
# pending web-search restart rides along rather than being consumed
# here: `_install_provider_then_reopen_auth` consumes it on every
# one of its exits (reopen defers to the new manager's close; its
# non-reopen paths call the offer directly), so the flag is never
# stranded past this early return.
extra = screen.pending_install_extra
if extra is not None:
from functools import partial
@@ -13494,15 +13516,121 @@ class DeepAgentsApp(App):
return
task = asyncio.create_task(self._resume_server_after_auth_change())
task.add_done_callback(_log_task_exception)
# A saved Tavily key only enables `web_search` after a respawn.
# Offer the restart now that the manager has closed, so the prompt
# never stacks over the still-open manager.
self._maybe_offer_deferred_web_search_restart()
screen = AuthManagerScreen(initial_provider=initial_provider)
self.push_screen(screen, handle_result)
def on_auth_manager_screen_credential_saved(self, event: Message) -> None:
"""Retry credentials-blocked startup immediately after `/auth` saves a key."""
def on_auth_manager_screen_credential_saved(
self, event: AuthManagerScreen.CredentialSaved
) -> None:
"""Retry credentials-blocked startup immediately after `/auth` saves a key.
A saved Tavily key additionally gates the spawn-time `web_search` tool,
so flag that a restart should be offered once the manager closes.
"""
event.stop()
# Kick off the credentials-blocked-startup retry first — it is the
# response the user is actually waiting for. Web-search bookkeeping is
# secondary and runs synchronously after, so a failure there can never
# preempt scheduling the resume.
task = asyncio.create_task(self._resume_server_after_auth_change())
task.add_done_callback(_log_task_exception)
self._note_web_search_restart_if_needed(event.provider)
def on_auth_manager_screen_credential_deleted(
self, event: AuthManagerScreen.CredentialDeleted
) -> None:
"""Clear auth-derived in-memory state after `/auth` deletes a key."""
event.stop()
self._clear_web_search_restart_if_needed(event.provider)
def _note_web_search_restart_if_needed(self, provider: str) -> None:
"""Flag an offered restart when a saved key enables `web_search`.
`web_search` is bound only when Tavily is configured at server spawn
time. A server that already spawned with Tavily has the tool bound, so
only a running server that lacks it needs a respawn. The stored key is
exported to the environment eagerly (as onboarding does) so a later
`/restart` — or the offered one — picks it up on reload. Ownership of
that export is tracked separately from the offer flag so a later delete
can reliably undo it (see `_clear_web_search_restart_if_needed`).
Args:
provider: The `/auth` config key that was just saved.
"""
from deepagents_code.model_config import TAVILY_SERVICE
if provider != TAVILY_SERVICE:
return
from deepagents_code.config import settings
if settings.has_tavily:
return
if self._server_proc is None or self._server_kwargs is None:
return
from deepagents_code.model_config import apply_stored_service_credentials
previous = os.environ.get("TAVILY_API_KEY")
apply_stored_service_credentials()
exported = os.environ.get("TAVILY_API_KEY")
if not exported:
return
if not self._auth_exported_tavily and previous != exported:
self._auth_exported_tavily_original = previous
self._auth_exported_tavily = True
self._pending_web_search_restart = True
def _clear_web_search_restart_if_needed(self, provider: str) -> None:
"""Disarm a Tavily restart offer and undo its env export after a delete.
Only touches `TAVILY_API_KEY` when *this* app exported it (tracked by
`_auth_exported_tavily`, set in `_note_web_search_restart_if_needed`) —
not the offer flag, which is consumed on manager close before a delete
can happen. That distinction is why a shell-provided key survives a
delete (the export flag was never set) while our own export is reverted
to whatever value it shadowed.
Args:
provider: The `/auth` config key that was just deleted.
"""
from deepagents_code.model_config import TAVILY_SERVICE
if provider != TAVILY_SERVICE:
return
if self._auth_exported_tavily:
if self._auth_exported_tavily_original is None:
os.environ.pop("TAVILY_API_KEY", None)
else:
os.environ["TAVILY_API_KEY"] = self._auth_exported_tavily_original
self._auth_exported_tavily = False
self._auth_exported_tavily_original = None
self._pending_web_search_restart = False
def _maybe_offer_deferred_web_search_restart(self) -> None:
"""Offer the deferred Tavily restart once the `/auth` manager has closed.
Consumes `_pending_web_search_restart` so the offer fires exactly once
and never leaks into a later, unrelated manager close. Scheduled via
`call_after_refresh` so the prompt mounts after the dismissing manager
fully unwinds rather than stacking over it.
"""
if not self._pending_web_search_restart:
return
self._pending_web_search_restart = False
self.call_after_refresh(self._launch_web_search_restart_prompt)
def _launch_web_search_restart_prompt(self) -> None:
"""Schedule the deferred web-search restart offer as a background task.
The close-ordering guarantee lives on the caller
(`_maybe_offer_deferred_web_search_restart`, via `call_after_refresh`).
"""
task = asyncio.create_task(self._offer_restart_for_web_search())
task.add_done_callback(_log_task_exception)
async def _resume_server_after_auth_change(self) -> None:
"""Bring the server up after `/auth` if a credential now unblocks it.
@@ -13599,12 +13727,18 @@ class DeepAgentsApp(App):
"Reopen `/auth` to add a key once it has.",
),
)
# We are not reopening the manager, so surface a Tavily restart the
# user armed earlier in this session rather than stranding the flag.
self._maybe_offer_deferred_web_search_restart()
return
# `ready is False`: the extra genuinely didn't install. `_install_extra`
# has already surfaced the reason to the user, so stay in chat rather
# than reopen — but log it so an "install button did nothing" report is
# debuggable without relying on that sibling method's invariant.
logger.debug("Provider extra %r not importable after install attempt", extra)
# No manager reopen on this path either, so consume any armed Tavily
# restart here.
self._maybe_offer_deferred_web_search_restart()
def _switch_agent(self, agent_name: str) -> None:
"""Switch to a different agent and hot-restart the backing server.
@@ -15602,8 +15736,9 @@ class DeepAgentsApp(App):
prompt to run the restart immediately instead of making the user type
`/restart`.
Owns all of its own follow-up messaging so the caller never appends a
redundant hint:
Surfaces its own follow-up messaging so a caller need not append one
(the `/install` extra path relies on this; `/install --package`
deliberately keeps a short hint of its own):
- Owned + idle: show the prompt (its button is the call to action). If
the prompt can't be shown, fall back to a `/restart` hint.
@@ -15615,24 +15750,100 @@ class DeepAgentsApp(App):
Args:
label: Installed extra/package name, surfaced in the prompt title.
"""
await self._offer_server_restart(
label=label,
verb="Installed",
relaunch_hint=f"Relaunch dcode to load '{label}'.",
busy_hint=(
f"Run `/restart` to load '{label}' once the current task finishes."
),
manual_hint=f"Run `/restart` to load '{label}' now.",
fail_hint=(
f"Couldn't restart the server automatically to load "
f"'{label}'. Run `/restart` to load it."
),
)
async def _offer_restart_for_web_search(self) -> None:
"""Offer a restart so a Tavily key saved via `/auth` enables `web_search`.
The app-owned server binds `web_search` only when Tavily is configured
at spawn time (see `server_graph._build_tools`), so a key added
mid-session takes effect only after the server respawns. Mirrors the
post-install offer: same guards, watchdog, and fallback messaging, with
web-search-specific copy.
"""
from deepagents_code.config import settings
if settings.has_tavily:
# A respawn happened between arming the offer and now — e.g. an
# install-on-select in the same `/auth` session auto-restarted the
# server, which reloaded config and rebound `web_search`. The
# restart is no longer needed, so don't offer a redundant one.
return
await self._offer_server_restart(
label="Tavily API key",
verb="Saved",
prompt_body=(
"Restart the server to enable web search, or defer with `/restart`."
),
relaunch_hint="Relaunch dcode to enable web search with your Tavily key.",
busy_hint=(
"Run `/restart` to enable web search once the current task finishes."
),
manual_hint="Run `/restart` to enable web search now.",
fail_hint=(
"Couldn't restart the server automatically to enable web "
"search. Run `/restart` to enable it."
),
)
async def _offer_server_restart(
self,
*,
label: str,
verb: str,
relaunch_hint: str,
busy_hint: str,
manual_hint: str,
fail_hint: str,
prompt_body: str | None = None,
) -> None:
"""Offer a one-keypress owned-server restart with caller-supplied copy.
Shared by the post-install offer and the post-`/auth` web-search offer:
both prompt to respawn the app-owned LangGraph subprocess so a change
that only takes effect at spawn time (a newly installed package, a
newly configured Tavily key) applies without exiting the TUI.
Surfaces its own follow-up messaging so a caller need not append one
(both the post-install and web-search offers rely on this; a caller
that still prints its own hint, like `/install --package`, does so
deliberately):
- Owned + idle: show the prompt (its button is the call to action). If
the prompt can't be shown, fall back to `manual_hint`.
- Owned + busy/connecting: a restart cancels in-flight work, so surface
`busy_hint`.
- No owned subprocess (remote server): surface `relaunch_hint`.
Args:
label: Subject surfaced in the prompt title (e.g. an extra name or
`"Tavily API key"`).
verb: Past-tense action shown before `label` in the prompt title.
relaunch_hint: Message shown when there is no owned subprocess.
busy_hint: Message shown when the server is busy/connecting.
manual_hint: Fallback shown when the prompt can't be displayed.
fail_hint: Message shown when a chosen restart could not run.
prompt_body: Optional override for the prompt's explanatory line.
"""
if self._server_proc is None or self._server_kwargs is None:
await self._mount_message(
AppMessage(f"Relaunch dcode to load '{label}'."),
)
await self._mount_message(AppMessage(relaunch_hint))
return
if self._agent_running or self._connecting:
await self._mount_message(
AppMessage(
f"Run `/restart` to load '{label}' once the current task finishes.",
),
)
await self._mount_message(AppMessage(busy_hint))
return
# Owned + idle. A `/restart` respawns the subprocess (same effect as a
# relaunch, without exiting), so every couldn't-show-the-prompt path
# below degrades to this hint rather than mentioning a relaunch.
manual_hint = f"Run `/restart` to load '{label}' now."
try:
from deepagents_code.tui.widgets.restart_prompt import RestartPromptScreen
except ModuleNotFoundError:
@@ -15645,7 +15856,7 @@ class DeepAgentsApp(App):
# deliberately narrow — a genuine `ImportError` from a broken modal
# still propagates rather than being mistaken for an upgrade race.
logger.warning(
"restart_prompt unavailable after installing %r; falling back "
"restart_prompt unavailable for %r restart; falling back "
"to the manual /restart hint",
label,
exc_info=True,
@@ -15659,12 +15870,14 @@ class DeepAgentsApp(App):
# (compose crash, programmatic teardown that skips the dismiss
# callback).
choice = await asyncio.wait_for(
self._push_screen_wait(RestartPromptScreen(label)),
self._push_screen_wait(
RestartPromptScreen(label, verb=verb, body=prompt_body)
),
timeout=_MODAL_WATCHDOG_TIMEOUT_SECONDS,
)
except TimeoutError:
logger.warning(
"Restart prompt after installing %r timed out; falling back to "
"Restart prompt for %r timed out; falling back to "
"the manual /restart hint",
label,
)
@@ -15673,9 +15886,7 @@ class DeepAgentsApp(App):
except Exception:
# Modal could not be mounted (e.g. another modal hijacked the
# stack). Fall back to the manual `/restart` hint.
logger.exception(
"Failed to mount restart prompt after installing %r", label
)
logger.exception("Failed to mount restart prompt for %r", label)
await self._mount_message(AppMessage(manual_hint))
return
@@ -15685,12 +15896,7 @@ class DeepAgentsApp(App):
# log. Surface a message so an explicit "restart" choice never looks
# like a silent no-op. Mirrors the `auto_restart` path.
if choice == "restart" and not await self._restart_after_install(label):
await self._mount_message(
AppMessage(
f"Couldn't restart the server automatically to load "
f"'{label}'. Run `/restart` to load it.",
),
)
await self._mount_message(AppMessage(fail_hint))
# Otherwise the prompt was shown and the user made an informed choice;
# no further hint is needed.
@@ -15715,13 +15921,11 @@ class DeepAgentsApp(App):
not currently available or fails.
"""
if self._server_proc is None or self._server_kwargs is None:
logger.info(
"Cannot auto-restart after installing %s: no app-owned server", label
)
logger.info("Cannot auto-restart to load %s: no app-owned server", label)
return False
if self._agent_running or self._connecting:
logger.info(
"Cannot auto-restart after installing %s: server is busy",
"Cannot auto-restart to load %s: server is busy",
label,
)
return False
libs/code/deepagents_code/model_config.py
+10 -1
View File
@@ -562,9 +562,18 @@ Storing a key for this service via `/auth` also enables tracing at startup
name, so it gets special handling beyond a plain key copy.
"""
TAVILY_SERVICE = "tavily"
"""Service name for Tavily web search in `SERVICE_API_KEY_ENV`.
Storing a key for this service via `/auth` gates the spawn-time `web_search`
tool (see `server_graph._build_tools`), so a key added to a running server
takes effect only after a respawn — the app offers that restart, and this
constant is the single name its `/auth` handling compares against.
"""
SERVICE_API_KEY_ENV: dict[str, str] = {
LANGSMITH_SERVICE: "LANGSMITH_API_KEY",
"tavily": "TAVILY_API_KEY",
TAVILY_SERVICE: "TAVILY_API_KEY",
}
"""Non-model services configurable via `/auth`, mapped to their API-key env var.
libs/code/deepagents_code/tui/widgets/auth.py
+46 -6
View File
@@ -20,6 +20,7 @@ from __future__ import annotations
import logging
import os
from enum import StrEnum
from functools import partial
from typing import TYPE_CHECKING, ClassVar, NamedTuple
from urllib.parse import urlsplit
@@ -1498,7 +1499,39 @@ class AuthManagerScreen(ModalScreen[None]):
"""
class CredentialSaved(Message):
"""Posted when a key prompt successfully persists credentials."""
"""Posted when a key prompt successfully persists credentials.
Carries the `/auth` config key that was saved so the app can react to
credentials that gate spawn-time behavior — e.g. a Tavily key that
enables the `web_search` tool only after the server respawns.
"""
def __init__(self, provider: str) -> None:
"""Store the saved provider/service identifier.
Args:
provider: The `/auth` config key that was saved (a model
provider name or a service key such as `"tavily"`).
"""
super().__init__()
self.provider = provider
class CredentialDeleted(Message):
"""Posted when a key prompt deletes stored credentials.
Carries the `/auth` config key that was deleted so the app can clear
any in-memory state derived from the now-removed credential.
"""
def __init__(self, provider: str) -> None:
"""Store the deleted provider/service identifier.
Args:
provider: The `/auth` config key that was deleted (a model
provider name or a service key such as `"tavily"`).
"""
super().__init__()
self.provider = provider
BINDINGS: ClassVar[list[BindingType]] = [
Binding("escape", "cancel", "Close", show=False, priority=True),
@@ -1695,13 +1728,13 @@ class AuthManagerScreen(ModalScreen[None]):
# same way as a model-provider key.
self.app.push_screen(
AuthPromptScreen(provider, SERVICE_API_KEY_ENV[provider]),
self._on_prompt_closed,
partial(self._on_prompt_closed, provider),
)
return
env_var = get_credential_env_var(provider)
self.app.push_screen(
AuthPromptScreen(provider, env_var),
self._on_prompt_closed,
partial(self._on_prompt_closed, provider),
)
def _prompt_install_provider(self, provider: str, extra: str) -> None:
@@ -1798,11 +1831,18 @@ class AuthManagerScreen(ModalScreen[None]):
"""Move the option-list cursor up."""
self.query_one("#auth-manager-options", OptionList).action_cursor_up()
def _on_prompt_closed(self, result: AuthResult | None) -> None:
"""Refresh the option list once the prompt dismisses."""
def _on_prompt_closed(self, provider: str, result: AuthResult | None) -> None:
"""Refresh the option list once the prompt dismisses.
Args:
provider: The provider/service whose prompt just closed.
result: Outcome of the prompt interaction.
"""
self._refresh_options()
if result is AuthResult.SAVED:
self.post_message(self.CredentialSaved())
self.post_message(self.CredentialSaved(provider))
elif result is AuthResult.DELETED:
self.post_message(self.CredentialDeleted(provider))
def _refresh_options(self) -> None:
"""Rebuild option labels from current store state."""
libs/code/deepagents_code/tui/widgets/restart_prompt.py
+34 -11
View File
@@ -1,10 +1,12 @@
"""Confirmation modal offered after a restart-capable `/install`.
"""Confirmation modal offered when a change needs an owned-server respawn.
Provider and sandbox extras (and `--package` installs) are imported by the
app-owned LangGraph server subprocess, so a `/restart` loads them without
exiting the TUI. Rather than make the user type `/restart` by hand, this
modal offers to run that restart immediately while leaving deferral one
keypress away.
Some changes take effect only when the app-owned LangGraph server subprocess
spawns: provider/sandbox extras and `--package` installs are imported at spawn
time, and a Tavily key saved via `/auth` binds the `web_search` tool only at
spawn time. A `/restart` respawns the subprocess without exiting the TUI, so
rather than make the user type `/restart` by hand, this modal offers to run
that restart immediately while leaving deferral one keypress away. The title
`verb` and `body` are caller-supplied so one modal serves each flow.
"""
from __future__ import annotations
@@ -28,7 +30,10 @@ RestartChoice = Literal["restart", "later"]
class RestartPromptScreen(ModalScreen[RestartChoice]):
"""Modal asking whether to restart the server after a successful install.
"""Modal asking whether to restart the server for a spawn-time change.
Serves both the post-install offer and the post-`/auth` web-search offer;
the caller supplies the title `verb` and `body` copy.
Dismisses with `"restart"` when the user accepts and `"later"` when the
user defers. Esc is treated as "later" so the user is never forced into a
@@ -75,14 +80,31 @@ class RestartPromptScreen(ModalScreen[RestartChoice]):
}
"""
def __init__(self, label: str) -> None:
_DEFAULT_BODY = "Restart the server to load it now, or defer with `/restart`."
def __init__(
self,
label: str,
*,
verb: str,
body: str | None = None,
) -> None:
"""Initialize the prompt.
Args:
label: Installed extra/package name, surfaced in the title.
label: The subject surfaced in the title (e.g. an installed extra
name, or a saved credential like ``"Tavily API key"``).
verb: Past-tense action shown before `label` in the title — e.g.
`"Installed"` for the post-install flow or `"Saved"` for a
saved credential. Required (no default) so each flow states its
own intent and a future caller can't inherit install-only copy.
body: Optional override for the explanatory line under the title.
Defaults to the generic restart copy.
"""
super().__init__()
self._label = label
self._verb = verb
self._body = body or self._DEFAULT_BODY
def compose(self) -> ComposeResult:
"""Compose the confirmation dialog.
@@ -94,15 +116,16 @@ class RestartPromptScreen(ModalScreen[RestartChoice]):
with Vertical():
yield Static(
Content.from_markup(
"$check Installed [bold]$name[/bold]",
"$check $verb [bold]$name[/bold]",
check=glyphs.checkmark,
verb=self._verb,
name=self._label,
),
classes="restart-prompt-title",
markup=False,
)
yield Static(
"Restart the server to load it now, or defer with `/restart`.",
self._body,
classes="restart-prompt-body",
markup=False,
)
libs/code/tests/unit_tests/test_app.py
+14 -3
View File
@@ -12046,10 +12046,16 @@ class TestInstallExtraAuthContinuation:
async def test_does_not_reopen_auth_when_install_failed(
self, caplog: pytest.LogCaptureFixture
) -> None:
"""A failed install leaves the user in chat and logs the dead-end at DEBUG."""
"""A failed install leaves the user in chat and logs the dead-end at DEBUG.
This non-reopen path must still consume an armed web-search restart so
the flag never strands onto a later, unrelated `/auth` close.
"""
app = DeepAgentsApp()
app._install_extra = AsyncMock(return_value=False) # ty: ignore
app._show_auth_manager = AsyncMock() # ty: ignore
app.call_after_refresh = MagicMock() # ty: ignore
app._pending_web_search_restart = True
with (
patch("deepagents_code.app._extra_is_ready", return_value=False),
@@ -12060,18 +12066,22 @@ class TestInstallExtraAuthContinuation:
app._install_extra.assert_awaited_once_with("baseten", auto_restart=True) # ty: ignore
app._show_auth_manager.assert_not_awaited() # ty: ignore
assert any("baseten" in record.message for record in caplog.records)
assert app._pending_web_search_restart is False
async def test_surfaces_hint_when_install_state_unverifiable(self) -> None:
"""An unknown post-install state points the user back to `/auth`.
When the extra can't be introspected (`_extra_is_ready` returns `None`)
the manager must not reopen, but the flow must not dead-end silently
either — a message tells the user how to finish.
either — a message tells the user how to finish. This path also consumes
an armed web-search restart so the flag is never stranded.
"""
app = DeepAgentsApp()
app._install_extra = AsyncMock(return_value=False) # ty: ignore
app._show_auth_manager = AsyncMock() # ty: ignore
app._mount_message = AsyncMock() # ty: ignore
app.call_after_refresh = MagicMock() # ty: ignore
app._pending_web_search_restart = True
with patch("deepagents_code.app._extra_is_ready", return_value=None):
await app._install_provider_then_reopen_auth("baseten", provider="baseten")
@@ -12080,6 +12090,7 @@ class TestInstallExtraAuthContinuation:
app._mount_message.assert_awaited_once() # ty: ignore
message = app._mount_message.await_args.args[0] # ty: ignore
assert "baseten" in message._content
assert app._pending_web_search_restart is False
class TestExtraIsReady:
@@ -13569,7 +13580,7 @@ class TestDeferredActions:
app._resume_server_after_auth_change = resume # ty: ignore
app.on_auth_manager_screen_credential_saved(
AuthManagerScreen.CredentialSaved()
AuthManagerScreen.CredentialSaved("openai")
)
await asyncio.sleep(0)
libs/code/tests/unit_tests/test_auth_web_search_restart.py
+609
View File
@@ -0,0 +1,609 @@
"""Tests for offering a server restart when a Tavily key is saved via `/auth`.
`web_search` is bound only when Tavily is configured at server spawn time
(see `server_graph._build_tools`), so a key added to an already-running server
takes effect only after a respawn. Saving a Tavily key in the `/auth` manager
should flag — and, once the manager closes, offer — that restart.
"""
from __future__ import annotations
import os
from unittest.mock import AsyncMock, MagicMock
import pytest
from deepagents_code.app import DeepAgentsApp
def _fake_tavily_export(monkeypatch: pytest.MonkeyPatch) -> None:
"""Make `apply_stored_service_credentials` export a Tavily key to the env."""
def _apply() -> None:
monkeypatch.setenv("TAVILY_API_KEY", "tvly-fake")
monkeypatch.setattr(
"deepagents_code.model_config.apply_stored_service_credentials",
_apply,
)
class TestNoteWebSearchRestart:
"""`_note_web_search_restart_if_needed` gating."""
def test_flags_restart_when_running_server_lacks_tavily(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
"""A saved Tavily key on a Tavily-less running server arms the offer."""
from deepagents_code.config import settings
monkeypatch.delenv("TAVILY_API_KEY", raising=False)
monkeypatch.setattr(settings, "tavily_api_key", None)
_fake_tavily_export(monkeypatch)
app = DeepAgentsApp()
app._server_proc = MagicMock()
app._server_kwargs = {"model_name": "anthropic:fake"}
app._note_web_search_restart_if_needed("tavily")
assert app._pending_web_search_restart is True
def test_ignores_non_tavily_service(self, monkeypatch: pytest.MonkeyPatch) -> None:
"""Model-provider keys don't gate a spawn-time tool, so no offer."""
from deepagents_code.config import settings
monkeypatch.setattr(settings, "tavily_api_key", None)
_fake_tavily_export(monkeypatch)
app = DeepAgentsApp()
app._server_proc = MagicMock()
app._server_kwargs = {"model_name": "anthropic:fake"}
app._note_web_search_restart_if_needed("openai")
assert app._pending_web_search_restart is False
def test_skips_when_server_already_has_tavily(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
"""A server spawned with Tavily already bound `web_search`."""
from deepagents_code.config import settings
monkeypatch.setattr(settings, "tavily_api_key", "tvly-existing")
app = DeepAgentsApp()
app._server_proc = MagicMock()
app._server_kwargs = {"model_name": "anthropic:fake"}
app._note_web_search_restart_if_needed("tavily")
assert app._pending_web_search_restart is False
def test_skips_when_no_owned_server(self, monkeypatch: pytest.MonkeyPatch) -> None:
"""With no owned subprocess there is nothing to respawn."""
from deepagents_code.config import settings
monkeypatch.setattr(settings, "tavily_api_key", None)
_fake_tavily_export(monkeypatch)
app = DeepAgentsApp()
app._server_proc = None
app._server_kwargs = None
app._note_web_search_restart_if_needed("tavily")
assert app._pending_web_search_restart is False
def test_skips_when_export_produces_no_key(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
"""If the eager export lands no `TAVILY_API_KEY`, stay disarmed.
A stored entry can be empty/malformed, in which case
`apply_stored_service_credentials` exports nothing; without an env key
a respawn couldn't bind `web_search`, so there is nothing to offer.
"""
from deepagents_code.config import settings
monkeypatch.delenv("TAVILY_API_KEY", raising=False)
monkeypatch.setattr(settings, "tavily_api_key", None)
monkeypatch.setattr(
"deepagents_code.model_config.apply_stored_service_credentials",
lambda: None,
)
app = DeepAgentsApp()
app._server_proc = MagicMock()
app._server_kwargs = {"model_name": "anthropic:fake"}
app._note_web_search_restart_if_needed("tavily")
assert app._pending_web_search_restart is False
def test_deleted_tavily_key_clears_pending_restart_and_env(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Deleting Tavily disarms the deferred restart and removes its env bridge."""
monkeypatch.setenv("TAVILY_API_KEY", "tvly-deleted")
app = DeepAgentsApp()
app._pending_web_search_restart = True
app._auth_exported_tavily = True
app._auth_exported_tavily_original = None
app._clear_web_search_restart_if_needed("tavily")
assert app._pending_web_search_restart is False
assert "TAVILY_API_KEY" not in os.environ
def test_deleted_tavily_key_clears_env_after_offer_was_consumed(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Deleting Tavily still removes `/auth` export after prompt scheduling."""
from deepagents_code.config import settings
monkeypatch.delenv("TAVILY_API_KEY", raising=False)
monkeypatch.setattr(settings, "tavily_api_key", None)
_fake_tavily_export(monkeypatch)
app = DeepAgentsApp()
app._server_proc = MagicMock()
app._server_kwargs = {"model_name": "anthropic:fake"}
app._note_web_search_restart_if_needed("tavily")
app._maybe_offer_deferred_web_search_restart()
app._clear_web_search_restart_if_needed("tavily")
assert app._pending_web_search_restart is False
assert "TAVILY_API_KEY" not in os.environ
def test_deleted_tavily_key_restores_original_env(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Deleting Tavily restores a shell key that `/auth` temporarily replaced."""
from deepagents_code.config import settings
monkeypatch.setenv("TAVILY_API_KEY", "tvly-from-shell")
monkeypatch.setattr(settings, "tavily_api_key", None)
_fake_tavily_export(monkeypatch)
app = DeepAgentsApp()
app._server_proc = MagicMock()
app._server_kwargs = {"model_name": "anthropic:fake"}
app._note_web_search_restart_if_needed("tavily")
app._maybe_offer_deferred_web_search_restart()
app._clear_web_search_restart_if_needed("tavily")
assert app._pending_web_search_restart is False
assert os.environ["TAVILY_API_KEY"] == "tvly-from-shell"
def test_deleted_non_tavily_key_leaves_pending_restart_and_env(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Model-provider deletes do not affect Tavily restart state."""
monkeypatch.setenv("TAVILY_API_KEY", "tvly-still-present")
app = DeepAgentsApp()
app._pending_web_search_restart = True
app._clear_web_search_restart_if_needed("openai")
assert app._pending_web_search_restart is True
assert os.environ["TAVILY_API_KEY"] == "tvly-still-present"
def test_deleted_tavily_key_preserves_env_when_restart_was_not_pending(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Deleting a stored key must not clear an independent shell env key."""
monkeypatch.setenv("TAVILY_API_KEY", "tvly-from-shell")
app = DeepAgentsApp()
app._pending_web_search_restart = False
app._clear_web_search_restart_if_needed("tavily")
assert app._pending_web_search_restart is False
assert os.environ["TAVILY_API_KEY"] == "tvly-from-shell"
def test_deleted_tavily_key_preserves_env_when_only_restart_flag_was_pending(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
"""The one-shot prompt flag alone does not prove `/auth` owns the env var."""
monkeypatch.setenv("TAVILY_API_KEY", "tvly-from-shell")
app = DeepAgentsApp()
app._pending_web_search_restart = True
app._clear_web_search_restart_if_needed("tavily")
assert app._pending_web_search_restart is False
assert os.environ["TAVILY_API_KEY"] == "tvly-from-shell"
class TestOfferRestartForWebSearch:
"""`_offer_restart_for_web_search` messaging and restart dispatch.
Each test pins `settings.tavily_api_key` to `None` so the idempotent
"already configured" guard never short-circuits the offer under test; the
guard itself is covered by `test_skips_offer_when_tavily_already_configured`.
"""
async def test_no_owned_server_recommends_relaunch(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
"""A remote/not-owned server can't be `/restart`ed — recommend relaunch."""
from deepagents_code.config import settings
monkeypatch.setattr(settings, "tavily_api_key", None)
app = DeepAgentsApp()
app._server_proc = None
app._server_kwargs = None
app._mount_message = AsyncMock() # ty: ignore
await app._offer_restart_for_web_search()
contents = " ".join(
str(c.args[0]._content)
for c in app._mount_message.await_args_list # ty: ignore
)
assert "Relaunch dcode" in contents
assert "web search" in contents
assert "/restart" not in contents
async def test_busy_recommends_restart(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
"""An owned-but-busy server points at `/restart`, never a relaunch."""
from deepagents_code.config import settings
monkeypatch.setattr(settings, "tavily_api_key", None)
app = DeepAgentsApp()
app._server_proc = MagicMock()
app._server_kwargs = {"model_name": "anthropic:fake"}
app._agent_running = True
app._connecting = False
app._mount_message = AsyncMock() # ty: ignore
await app._offer_restart_for_web_search()
contents = " ".join(
str(c.args[0]._content)
for c in app._mount_message.await_args_list # ty: ignore
)
assert "/restart" in contents
assert "relaunch" not in contents.lower()
async def test_restart_choice_dispatches_restart(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Choosing restart respawns the owned server via `_restart_after_install`."""
from deepagents_code.config import settings
monkeypatch.setattr(settings, "tavily_api_key", None)
app = DeepAgentsApp()
app._server_proc = MagicMock()
app._server_kwargs = {"model_name": "anthropic:fake"}
app._agent_running = False
app._connecting = False
app._mount_message = AsyncMock() # ty: ignore
app._push_screen_wait = AsyncMock(return_value="restart") # ty: ignore
app._restart_after_install = AsyncMock(return_value=True) # ty: ignore
await app._offer_restart_for_web_search()
app._restart_after_install.assert_awaited_once_with( # ty: ignore
"Tavily API key"
)
async def test_restart_state_flip_surfaces_fallback(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
"""A chosen restart that can't run surfaces a web-search fallback hint."""
from deepagents_code.config import settings
monkeypatch.setattr(settings, "tavily_api_key", None)
app = DeepAgentsApp()
app._server_proc = MagicMock()
app._server_kwargs = {"model_name": "anthropic:fake"}
app._agent_running = False
app._connecting = False
app._mount_message = AsyncMock() # ty: ignore
app._push_screen_wait = AsyncMock(return_value="restart") # ty: ignore
app._restart_after_install = AsyncMock(return_value=False) # ty: ignore
await app._offer_restart_for_web_search()
contents = " ".join(
str(c.args[0]._content)
for c in app._mount_message.await_args_list # ty: ignore
)
assert "Couldn't restart the server automatically to enable web" in contents
async def test_skips_offer_when_tavily_already_configured(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
"""A respawn that already bound `web_search` suppresses a redundant offer.
Guards the install-on-select-in-the-same-session case: the install
auto-restarted the server (reloading config and rebinding `web_search`),
so by the time the reopened manager closes there is nothing left to do.
"""
from deepagents_code.config import settings
monkeypatch.setattr(settings, "tavily_api_key", "tvly-now-configured")
app = DeepAgentsApp()
app._server_proc = MagicMock()
app._server_kwargs = {"model_name": "anthropic:fake"}
app._agent_running = False
app._connecting = False
app._mount_message = AsyncMock() # ty: ignore
app._push_screen_wait = AsyncMock() # ty: ignore
app._restart_after_install = AsyncMock() # ty: ignore
await app._offer_restart_for_web_search()
app._push_screen_wait.assert_not_awaited() # ty: ignore
app._restart_after_install.assert_not_awaited() # ty: ignore
app._mount_message.assert_not_awaited() # ty: ignore
async def test_later_choice_shows_no_followup_hint(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Deferring at the prompt is an informed choice — no extra hint follows.
The prompt's own button was the call to action, so the shared helper
stays silent (neither restarts nor mounts a `/restart` hint) when the
user picks "later".
"""
from deepagents_code.config import settings
monkeypatch.setattr(settings, "tavily_api_key", None)
app = DeepAgentsApp()
app._server_proc = MagicMock()
app._server_kwargs = {"model_name": "anthropic:fake"}
app._agent_running = False
app._connecting = False
app._mount_message = AsyncMock() # ty: ignore
app._push_screen_wait = AsyncMock(return_value="later") # ty: ignore
app._restart_after_install = AsyncMock() # ty: ignore
await app._offer_restart_for_web_search()
app._restart_after_install.assert_not_awaited() # ty: ignore
app._mount_message.assert_not_awaited() # ty: ignore
async def test_watchdog_timeout_falls_back_to_manual_hint(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
"""A prompt that never resolves is bounded by the watchdog → manual hint."""
from deepagents_code.config import settings
monkeypatch.setattr(settings, "tavily_api_key", None)
app = DeepAgentsApp()
app._server_proc = MagicMock()
app._server_kwargs = {"model_name": "anthropic:fake"}
app._agent_running = False
app._connecting = False
app._mount_message = AsyncMock() # ty: ignore
app._push_screen_wait = AsyncMock(side_effect=TimeoutError()) # ty: ignore
app._restart_after_install = AsyncMock() # ty: ignore
await app._offer_restart_for_web_search()
app._restart_after_install.assert_not_awaited() # ty: ignore
contents = " ".join(
str(c.args[0]._content)
for c in app._mount_message.await_args_list # ty: ignore
)
assert "/restart" in contents
assert "web search" in contents
async def test_mount_failure_falls_back_to_manual_hint(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
"""If the modal can't be mounted at all, degrade to the manual hint."""
from deepagents_code.config import settings
monkeypatch.setattr(settings, "tavily_api_key", None)
app = DeepAgentsApp()
app._server_proc = MagicMock()
app._server_kwargs = {"model_name": "anthropic:fake"}
app._agent_running = False
app._connecting = False
app._mount_message = AsyncMock() # ty: ignore
app._push_screen_wait = AsyncMock( # ty: ignore
side_effect=RuntimeError("stack hijacked")
)
app._restart_after_install = AsyncMock() # ty: ignore
await app._offer_restart_for_web_search()
app._restart_after_install.assert_not_awaited() # ty: ignore
contents = " ".join(
str(c.args[0]._content)
for c in app._mount_message.await_args_list # ty: ignore
)
assert "/restart" in contents
assert "web search" in contents
class TestCredentialSavedHandler:
"""The `/auth` credential-saved handler threads the provider through."""
async def test_saved_tavily_key_arms_restart_offer(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Saving Tavily flags the offer without waiting for the manager to close."""
from deepagents_code.config import settings
from deepagents_code.tui.widgets.auth import AuthManagerScreen
monkeypatch.delenv("TAVILY_API_KEY", raising=False)
monkeypatch.setattr(settings, "tavily_api_key", None)
_fake_tavily_export(monkeypatch)
app = DeepAgentsApp()
async with app.run_test() as pilot:
await pilot.pause()
app._server_proc = MagicMock()
app._server_kwargs = {"model_name": "anthropic:fake"}
app._resume_server_after_auth_change = AsyncMock() # ty: ignore
app.on_auth_manager_screen_credential_saved(
AuthManagerScreen.CredentialSaved("tavily")
)
await pilot.pause()
assert app._pending_web_search_restart is True
async def test_deleted_tavily_key_disarms_restart_offer(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Deleting Tavily clears the offer before the manager-close callback runs."""
from deepagents_code.tui.widgets.auth import AuthManagerScreen
monkeypatch.setenv("TAVILY_API_KEY", "tvly-deleted")
app = DeepAgentsApp()
async with app.run_test() as pilot:
await pilot.pause()
app._pending_web_search_restart = True
app._auth_exported_tavily = True
app._auth_exported_tavily_original = None
app.on_auth_manager_screen_credential_deleted(
AuthManagerScreen.CredentialDeleted("tavily")
)
await pilot.pause()
assert app._pending_web_search_restart is False
assert "TAVILY_API_KEY" not in os.environ
async def test_resume_is_scheduled_before_web_search_bookkeeping(self) -> None:
"""The credentials-blocked resume is scheduled even if bookkeeping raises.
The handler kicks off `_resume_server_after_auth_change` (the response
the user is waiting on) before the secondary web-search bookkeeping, so
a failure in the latter can never preempt the resume.
"""
from deepagents_code.tui.widgets.auth import AuthManagerScreen
app = DeepAgentsApp()
async with app.run_test() as pilot:
await pilot.pause()
app._resume_server_after_auth_change = AsyncMock() # ty: ignore
app._note_web_search_restart_if_needed = MagicMock( # ty: ignore
side_effect=RuntimeError("bookkeeping blew up")
)
with pytest.raises(RuntimeError):
app.on_auth_manager_screen_credential_saved(
AuthManagerScreen.CredentialSaved("tavily")
)
await pilot.pause()
app._resume_server_after_auth_change.assert_awaited_once() # ty: ignore
class TestDeferredOfferAfterManagerCloses:
"""The armed offer is consumed only when the `/auth` manager closes.
Spies on `call_after_refresh` rather than asserting the offer coroutine
actually ran: scheduling is synchronous inside the dismiss callback, so it
is deterministic, whereas the post-refresh invocation is timing-dependent.
"""
@staticmethod
def _spy_call_after_refresh(app: DeepAgentsApp) -> list[object]:
"""Record callbacks scheduled via `call_after_refresh`, still delegating."""
scheduled: list[object] = []
real = app.call_after_refresh
def _spy(callback: object, *args: object, **kwargs: object) -> object:
scheduled.append(callback)
return real(callback, *args, **kwargs) # ty: ignore
app.call_after_refresh = _spy # ty: ignore
return scheduled
async def test_offer_deferred_until_manager_closes(self) -> None:
"""An armed offer is scheduled on manager close — never while it is open."""
app = DeepAgentsApp()
async with app.run_test() as pilot:
await pilot.pause()
app._resume_server_after_auth_change = AsyncMock() # ty: ignore
app._launch_web_search_restart_prompt = MagicMock() # ty: ignore
scheduled = self._spy_call_after_refresh(app)
app._pending_web_search_restart = True
await app._show_auth_manager()
await pilot.pause()
# Manager is open: nothing scheduled and the flag is still armed.
assert app._launch_web_search_restart_prompt not in scheduled
assert app._pending_web_search_restart is True
# Closing the manager (no pending install) consumes the flag once
# and schedules the offer.
app.screen.dismiss(None)
await pilot.pause()
assert app._pending_web_search_restart is False
assert app._launch_web_search_restart_prompt in scheduled
async def test_flag_not_reoffered_on_a_later_unrelated_close(self) -> None:
"""Once consumed, a later manager close does not re-schedule the offer."""
app = DeepAgentsApp()
async with app.run_test() as pilot:
await pilot.pause()
app._resume_server_after_auth_change = AsyncMock() # ty: ignore
app._launch_web_search_restart_prompt = MagicMock() # ty: ignore
scheduled = self._spy_call_after_refresh(app)
app._pending_web_search_restart = True
await app._show_auth_manager()
await pilot.pause()
app.screen.dismiss(None)
await pilot.pause()
assert scheduled.count(app._launch_web_search_restart_prompt) == 1
# A second, unrelated open/close must not re-schedule the offer.
await app._show_auth_manager()
await pilot.pause()
app.screen.dismiss(None)
await pilot.pause()
assert scheduled.count(app._launch_web_search_restart_prompt) == 1
async def test_flag_rides_along_when_close_kicks_off_install(self) -> None:
"""A close that starts a provider install must not consume the offer.
When the manager closes with a `pending_install_extra`, the app hands
off to `_install_provider_then_reopen_auth` (which reopens the manager
or surfaces the offer on its non-reopen exits). Consuming the flag here
would strand the restart the user armed, so it must ride along untouched.
"""
app = DeepAgentsApp()
async with app.run_test() as pilot:
await pilot.pause()
app._resume_server_after_auth_change = AsyncMock() # ty: ignore
app._launch_web_search_restart_prompt = MagicMock() # ty: ignore
app._install_provider_then_reopen_auth = AsyncMock() # ty: ignore
scheduled = self._spy_call_after_refresh(app)
app._pending_web_search_restart = True
await app._show_auth_manager()
await pilot.pause()
# Simulate the user having confirmed installing an uninstalled
# provider: the manager records the extra to install on close.
app.screen.pending_install_extra = "baseten" # ty: ignore
app.screen.pending_install_provider = "baseten" # ty: ignore
app.screen.dismiss(None)
await pilot.pause()
# The flag rides along: neither consumed nor scheduled here.
assert app._pending_web_search_restart is True
assert app._launch_web_search_restart_prompt not in scheduled
app._install_provider_then_reopen_auth.assert_awaited_once() # ty: ignore
libs/code/tests/unit_tests/tui/widgets/test_auth_widgets.py
+43 -4
View File
@@ -65,6 +65,9 @@ class _AuthHostApp(App[None]):
self.prompt_result: AuthResult | None = None
self.prompt_dismissed = False
self.credential_saved_count = 0
self.credential_deleted_count = 0
self.last_saved_provider: str | None = None
self.last_deleted_provider: str | None = None
def compose(self) -> ComposeResult:
"""Render a placeholder root."""
@@ -103,10 +106,18 @@ class _AuthHostApp(App[None]):
self.push_screen(AuthManagerScreen(initial_provider=initial_provider))
def on_auth_manager_screen_credential_saved(
self, _event: AuthManagerScreen.CredentialSaved
self, event: AuthManagerScreen.CredentialSaved
) -> None:
"""Record credential-save notifications from the manager."""
self.credential_saved_count += 1
self.last_saved_provider = event.provider
def on_auth_manager_screen_credential_deleted(
self, event: AuthManagerScreen.CredentialDeleted
) -> None:
"""Record credential-delete notifications from the manager."""
self.credential_deleted_count += 1
self.last_deleted_provider = event.provider
class TestCodexAuthScreen:
@@ -1902,10 +1913,11 @@ class TestAuthManagerScreen:
await pilot.pause()
screen = cast("AuthManagerScreen", app.screen)
screen._on_prompt_closed(AuthResult.SAVED)
screen._on_prompt_closed("openai", AuthResult.SAVED)
await pilot.pause()
assert app.credential_saved_count == 1
assert app.last_saved_provider == "openai"
async def test_prompt_cancel_does_not_post_credential_saved_event(self) -> None:
"""Cancelling or deleting credentials should not trigger startup recovery."""
@@ -1915,12 +1927,39 @@ class TestAuthManagerScreen:
await pilot.pause()
screen = cast("AuthManagerScreen", app.screen)
screen._on_prompt_closed(AuthResult.CANCELLED)
screen._on_prompt_closed(AuthResult.DELETED)
screen._on_prompt_closed("openai", AuthResult.CANCELLED)
screen._on_prompt_closed("openai", AuthResult.DELETED)
await pilot.pause()
assert app.credential_saved_count == 0
async def test_prompt_delete_posts_credential_deleted_event(self) -> None:
"""Deleting a key notifies the app before the manager itself closes."""
app = _AuthHostApp()
async with app.run_test() as pilot:
app.show_manager()
await pilot.pause()
screen = cast("AuthManagerScreen", app.screen)
screen._on_prompt_closed("tavily", AuthResult.DELETED)
await pilot.pause()
assert app.credential_deleted_count == 1
assert app.last_deleted_provider == "tavily"
async def test_prompt_cancel_does_not_post_credential_deleted_event(self) -> None:
"""Cancelling a prompt does not clear app-side credential state."""
app = _AuthHostApp()
async with app.run_test() as pilot:
app.show_manager()
await pilot.pause()
screen = cast("AuthManagerScreen", app.screen)
screen._on_prompt_closed("tavily", AuthResult.CANCELLED)
await pilot.pause()
assert app.credential_deleted_count == 0
async def test_lists_known_providers(self) -> None:
"""Every well-known provider appears in the option list."""
app = _AuthHostApp()
libs/code/tests/unit_tests/tui/widgets/test_restart_prompt.py
+43 -4
View File
@@ -25,7 +25,9 @@ class TestRestartPromptScreen:
def on_dismiss(result: str | None) -> None:
outcomes.append(result)
app.push_screen(RestartPromptScreen("fireworks"), on_dismiss)
app.push_screen(
RestartPromptScreen("fireworks", verb="Installed"), on_dismiss
)
await pilot.pause()
await pilot.press("enter")
@@ -42,7 +44,9 @@ class TestRestartPromptScreen:
def on_dismiss(result: str | None) -> None:
outcomes.append(result)
app.push_screen(RestartPromptScreen("fireworks"), on_dismiss)
app.push_screen(
RestartPromptScreen("fireworks", verb="Installed"), on_dismiss
)
await pilot.pause()
await pilot.press("escape")
@@ -67,7 +71,7 @@ class TestRestartPromptScreen:
def on_dismiss(result: str | None) -> None:
outcomes.append(result)
screen = RestartPromptScreen("fireworks")
screen = RestartPromptScreen("fireworks", verb="Installed")
app.push_screen(screen, on_dismiss)
await pilot.pause()
@@ -80,9 +84,44 @@ class TestRestartPromptScreen:
"""The installed extra/package label is surfaced in the modal title."""
app = _RestartTestApp()
async with app.run_test() as pilot:
app.push_screen(RestartPromptScreen("langchain-custom"))
app.push_screen(RestartPromptScreen("langchain-custom", verb="Installed"))
await pilot.pause()
titles = app.screen.query(".restart-prompt-title")
assert len(titles) == 1
assert "langchain-custom" in str(titles.first().render())
async def test_renders_custom_verb(self) -> None:
"""A caller-supplied `verb` replaces the install-flow default in the title."""
app = _RestartTestApp()
async with app.run_test() as pilot:
app.push_screen(RestartPromptScreen("Tavily API key", verb="Saved"))
await pilot.pause()
title = str(app.screen.query_one(".restart-prompt-title").render())
assert "Saved" in title
assert "Tavily API key" in title
assert "Installed" not in title
async def test_renders_default_body(self) -> None:
"""With no `body` override, the generic restart copy is shown."""
app = _RestartTestApp()
async with app.run_test() as pilot:
app.push_screen(RestartPromptScreen("fireworks", verb="Installed"))
await pilot.pause()
body = str(app.screen.query_one(".restart-prompt-body").render())
assert body == RestartPromptScreen._DEFAULT_BODY
async def test_renders_body_override(self) -> None:
"""A caller-supplied `body` replaces the default explanatory line."""
app = _RestartTestApp()
override = "Restart the server to enable web search, or defer with `/restart`."
async with app.run_test() as pilot:
app.push_screen(
RestartPromptScreen("Tavily API key", verb="Saved", body=override)
)
await pilot.pause()
body = str(app.screen.query_one(".restart-prompt-body").render())
assert body == override