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(code): run MCP login during a run, queue the restart (#4643)

Browse Source 
`/mcp login` now runs during an active generation instead of waiting for
it to finish; if you accept the reconnect, the server restart is queued
and applied automatically once the current task completes.

---

`/mcp login` previously deferred the *entire* OAuth flow until an active
agent or shell task finished. But only the final server restart
(`_restart_server_for_mcp_refresh` → `_respawn_server`) actually
conflicts with a live run — the OAuth handshake and on-disk token write
never touch the running server. This starts login immediately even
during a generation and queues only the restart, so the user can
authenticate without waiting for the turn to end.

When the user accepts "reconnect" on the post-login prompt while a run
is active, the restart is queued as a new `mcp_reconnect` deferred
action (drained by the existing `_maybe_drain_deferred` on agent/shell
cleanup) instead of tearing down the in-flight generation. The reconnect
is marked pending in the meantime so `/mcp reconnect` and the splash
banner stay accurate, and the queued action re-checks the pending flag
to avoid a redundant restart if a manual `/mcp reconnect` already ran.
When idle, behavior is unchanged (immediate restart).

Made by [Open
SWE](https://openswe.vercel.app/agents/4defa6fb-fe26-d47b-9fa0-10b5a9c29dcd)

---------

Co-authored-by: open-swe[bot] <open-swe@users.noreply.github.com>
This commit is contained in:
Mason Daugherty
2026-07-14 00:12:36 -04:00
committed by GitHub
parent f9915db420
commit 65e1ee802e
2 changed files with 409 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/code/deepagents_code/app.py
+119 -24
View File
@@ -1348,6 +1348,7 @@ DeferredActionKind = Literal[
"chat_output",
"agent_switch",
"mcp_login",
"mcp_reconnect",
"rubric_model_switch",
"rubric_max_iterations_switch",
]
@@ -11413,8 +11414,16 @@ class DeepAgentsApp(App):
Dequeues and processes the next pending message in FIFO order.
Uses the `_processing_pending` flag to prevent reentrant execution.
Leaves the queue untouched while the server is connecting so the
`ServerReady` path can resume draining against the fully initialized
session.
"""
if self._processing_pending or not self._pending_messages or self._exit:
if (
self._processing_pending
or not self._pending_messages
or self._exit
or self._connecting
):
return
self._processing_pending = True
@@ -12645,6 +12654,7 @@ class DeepAgentsApp(App):
self._shell_worker.cancel()
if self._agent_running and self._agent_worker:
self._agent_worker.cancel()
self._warn_dropped_mcp_reconnect()
self._discard_queue()
def _defer_action(self, action: DeferredAction) -> None:
@@ -12693,12 +12703,41 @@ class DeepAgentsApp(App):
exc_info=True,
)
def _cancel_worker(self, worker: Worker[None] | None) -> None:
def _warn_dropped_mcp_reconnect(self) -> None:
"""Warn when an interrupt discards a queued MCP reconnect.
`_start_mcp_login` -> `_prompt_mcp_reconnect` can queue the server
restart while a run is in flight, telling the user it will fire once
the task completes. An interrupt (`Ctrl+C`, `Esc`, `/clear`) drops that
queued action before it drains, so that promise no longer holds. The
token is already on disk and the pending banner / `/mcp reconnect`
state survive the discard, so recovery is one command away — surface
that rather than letting the reconnect lapse silently.
"""
if not any(a.kind == "mcp_reconnect" for a in self._deferred_actions):
return
self.notify(
"Cancelled the queued MCP reconnect. Run `/mcp reconnect` to load "
"the new tools when ready.",
severity="warning",
timeout=8,
markup=False,
)
def _cancel_worker(
self, worker: Worker[None] | None, *, abort_pending_reconnect: bool = True
) -> None:
"""Discard the message queue and cancel an active worker.
Args:
worker: The worker to cancel.
abort_pending_reconnect: When `True` (the interrupt default), warn
if the discarded queue held a promised MCP reconnect. Pass
`False` from paths that fulfill the reconnect another way (a
full server restart) so the notice does not misfire.
"""
if abort_pending_reconnect:
self._warn_dropped_mcp_reconnect()
self._discard_queue()
if worker is not None:
worker.cancel()
@@ -15577,8 +15616,11 @@ class DeepAgentsApp(App):
async def _handle_mcp_reconnect_command(self, *, force: bool = False) -> None:
"""Restart the server to pick up any deferred MCP login tokens.
No-op (with an inline notice) when nothing is pending so the
command is safe to run idempotently. `force=True` bypasses the
Restarts immediately when a login is pending and the session is
idle; when an agent or shell task is running the restart is queued
via `DeferredAction(kind="mcp_reconnect")` and drained once the task
completes. No-op (with an inline notice) when nothing is pending so
the command is safe to run idempotently. `force=True` bypasses the
no-op guard via a confirmation modal — the escape hatch for
stale-cache or externally-edited-config cases where the server
needs a fresh load even though no login is queued in this
@@ -15589,6 +15631,22 @@ class DeepAgentsApp(App):
if no MCP login is queued.
"""
if self._pending_mcp_reconnect:
if self._agent_running or self._shell_running:
self._defer_action(
DeferredAction(
kind="mcp_reconnect",
execute=lambda: self._run_deferred_mcp_reconnect(
"pending login"
),
),
)
self.notify(
"The server will reconnect once the current task completes.",
severity="information",
timeout=8,
markup=False,
)
return
await self._restart_server_for_mcp_refresh("pending login")
return
if not force:
@@ -15973,10 +16031,13 @@ class DeepAgentsApp(App):
Rejects when MCP is disabled, in remote-server mode (no owned server
to restart), or while an agent switch is in progress. When the local
server is still connecting or the session is mid-run, the login is
queued via `_defer_action` and runs once the server is ready and the
user is idle. Config resolution and server-name validation happen
later, in `_run_mcp_login_worker`.
server is still connecting, the login is queued via `_defer_action`
and runs once the server is ready. An active agent or shell run does
*not* defer login: the OAuth handshake and token write never touch the
running server, so they proceed concurrently; only the follow-up
server restart is queued (see `_prompt_mcp_reconnect`). Config
resolution and server-name validation happen later, in
`_run_mcp_login_worker`.
Args:
server_name: MCP server name from `mcpServers`.
@@ -16028,20 +16089,10 @@ class DeepAgentsApp(App):
)
return
if self._agent_running or self._shell_running:
self.notify(
"MCP login will start once the current task completes.",
timeout=5,
markup=False,
)
self._defer_action(
DeferredAction(
kind="mcp_login",
execute=lambda: self._run_mcp_login_worker(server_name),
),
)
return
# An active agent/shell run is intentionally not a defer gate: the
# OAuth handshake and on-disk token write are independent of the
# running server, so login proceeds immediately and only the restart
# is queued once the task finishes (`_prompt_mcp_reconnect`).
self.run_worker(
self._run_mcp_login_worker(server_name),
exclusive=False,
@@ -16232,6 +16283,30 @@ class DeepAgentsApp(App):
choice = "later"
if choice == "reconnect":
if self._agent_running or self._shell_running:
# The restart tears down the server the active run lives on,
# so honor the user's "reconnect now" intent without killing
# the in-flight generation: queue the restart to fire once the
# task finishes (drained via `_maybe_drain_deferred`). The
# token is already on disk, so mark the reconnect pending for
# `/mcp reconnect` and the splash banner in the meantime.
self._pending_mcp_login_reconnect = True
self._sync_pending_mcp_reconnect()
self._apply_optimistic_mcp_login_pending_state(server_name)
self._defer_action(
DeferredAction(
kind="mcp_reconnect",
execute=lambda: self._run_deferred_mcp_reconnect(server_name),
),
)
self.notify(
f"Logged in to {server_name!r}. The server will reconnect "
"once the current task completes.",
severity="information",
timeout=8,
markup=False,
)
return
self._pending_mcp_login_reconnect = False
self._pending_mcp_disable_reconnect_servers.clear()
self._sync_pending_mcp_reconnect()
@@ -16282,6 +16357,24 @@ class DeepAgentsApp(App):
markup=False,
)
async def _run_deferred_mcp_reconnect(self, server_name: str) -> None:
"""Restart for MCP token refresh once the busy state clears.
Queued by `_prompt_mcp_reconnect` when the user accepts the restart
while an agent or shell task is still running — restarting then would
tear down the server the active run depends on. Re-checks the pending
flag so a manual `/mcp reconnect` in the interim isn't followed by a
redundant restart. (The queue is in-memory only, so a relaunch never
reaches this path — the fresh process loads the token on startup.)
Args:
server_name: Server whose login triggered the reconnect.
"""
if not self._pending_mcp_reconnect:
return
self._clear_mcp_login_reconnect_banner_counts(server_name)
await self._restart_server_for_mcp_refresh(server_name)
async def _restart_server_for_mcp_refresh(self, server_name: str) -> None:
"""Restart the app-owned LangGraph server to pick up new MCP tokens.
@@ -16610,9 +16703,11 @@ class DeepAgentsApp(App):
# Sever in-flight work bound to the dying subprocess. `_cancel_worker`
# discards the queued backlog too — those messages would otherwise
# fire against the freshly respawned agent silently.
# fire against the freshly respawned agent silently. This restart *is*
# the reconnect, so suppress the dropped-reconnect warning: the respawn
# below reloads every on-disk MCP token regardless.
if self._agent_running and self._agent_worker:
self._cancel_worker(self._agent_worker)
self._cancel_worker(self._agent_worker, abort_pending_reconnect=False)
self._agent_running = False
else:
self._discard_queue()
libs/code/tests/unit_tests/test_app.py
+290 -11
View File
@@ -3538,6 +3538,47 @@ class TestMessageQueue:
widgets = app.query(QueuedUserMessage)
assert len(widgets) == 1
async def test_deferred_restart_keeps_prompts_queued_until_ready(self) -> None:
"""A deferred restart must not let task cleanup consume queued prompts."""
app = DeepAgentsApp(agent=MagicMock())
async with app.run_test() as pilot:
await pilot.pause()
app._agent_running = True
await app._submit_input("after restart", "normal")
app._agent_running = False
processed = AsyncMock()
app._process_message = processed # ty: ignore
async def begin_restart() -> None:
await asyncio.sleep(0)
app._agent = None
app._connecting = True
app._defer_action(
DeferredAction(kind="mcp_reconnect", execute=begin_restart),
)
# Match task cleanup: drain deferred work, then try queued input.
await app._maybe_drain_deferred()
await app._process_next_from_queue()
processed.assert_not_awaited()
assert [message.text for message in app._pending_messages] == [
"after restart"
]
assert len(app._queued_widgets) == 1
# `ServerReady` clears the connecting state before its startup
# backlog resumes queue draining.
app._agent = MagicMock()
app._connecting = False
await app._process_next_from_queue()
processed.assert_awaited_once_with("after restart", "normal")
assert not app._pending_messages
assert not app._queued_widgets
async def test_message_blocked_while_thread_switching(self) -> None:
"""Submissions should be ignored while thread switching is in-flight."""
app = DeepAgentsApp()
@@ -15035,6 +15076,24 @@ class TestStartMcpLogin:
assert run_worker_kwargs["group"] == "mcp-login-provider"
assert run_worker_kwargs["exclusive"] is False
async def test_starts_worker_immediately_while_agent_running(self) -> None:
"""A live agent run runs login now; only the restart defers later.
The OAuth handshake and token write never touch the running server, so
an active generation must not block login — the follow-up restart is
what gets queued (see `_prompt_mcp_reconnect`).
"""
app = self._make_app()
app._connecting = False
app._server_proc = object() # ty: ignore
app._agent_running = True
app._run_mcp_login_worker = MagicMock() # ty: ignore
app._start_mcp_login("provider")
app.run_worker.assert_called_once() # ty: ignore
app._defer_action.assert_not_called() # ty: ignore
class TestBuildModelSwitchErrorBody:
"""Tests for `_build_model_switch_error_body` link-aware formatting."""
@@ -19738,7 +19797,9 @@ class TestMCPLoginCommand:
- `/mcp login <name>` reaches `_start_mcp_login`.
- Bare `/mcp` still opens the viewer.
- Remote-server mode refuses login and tells the user.
- Busy state defers login via `DeferredAction(kind="mcp_login")`.
- A busy agent/shell run runs login immediately (only the follow-up
reconnect defers); the still-connecting path defers login via
`DeferredAction(kind="mcp_login")`.
- The viewer's dismiss with a server name kicks off login.
"""
@@ -19804,8 +19865,8 @@ class TestMCPLoginCommand:
message = notify.call_args.args[0]
assert "remote server" in message.lower()
async def test_mcp_login_defers_while_agent_running(self) -> None:
"""Busy state queues the login via `DeferredAction(kind='mcp_login')`."""
async def test_mcp_login_runs_immediately_while_agent_running(self) -> None:
"""A live run no longer defers login; the OAuth worker starts at once."""
app = DeepAgentsApp(agent=MagicMock())
async with app.run_test() as pilot:
await pilot.pause()
@@ -19818,10 +19879,13 @@ class TestMCPLoginCommand:
app._server_proc = MagicMock()
app._agent_running = True
try:
with patch.object(app, "run_worker") as run_worker:
with (
patch.object(app, "run_worker") as run_worker,
patch.object(app, "_run_mcp_login_worker", new=MagicMock()),
):
app._start_mcp_login("notion")
run_worker.assert_not_called()
assert any(a.kind == "mcp_login" for a in app._deferred_actions)
run_worker.assert_called_once()
assert not any(a.kind == "mcp_login" for a in app._deferred_actions)
finally:
app._agent_running = False
@@ -20299,8 +20363,8 @@ class TestMCPLoginCommand:
finally:
app._agent_switching = False
async def test_mcp_login_defers_while_shell_running(self) -> None:
"""`_shell_running=True` also defers login via `DeferredAction`."""
async def test_mcp_login_runs_immediately_while_shell_running(self) -> None:
"""`_shell_running=True` no longer defers login; the worker starts now."""
app = DeepAgentsApp(agent=MagicMock())
async with app.run_test() as pilot:
await pilot.pause()
@@ -20313,10 +20377,13 @@ class TestMCPLoginCommand:
app._server_proc = MagicMock()
app._shell_running = True
try:
with patch.object(app, "run_worker") as run_worker:
with (
patch.object(app, "run_worker") as run_worker,
patch.object(app, "_run_mcp_login_worker", new=MagicMock()),
):
app._start_mcp_login("notion")
run_worker.assert_not_called()
assert any(a.kind == "mcp_login" for a in app._deferred_actions)
run_worker.assert_called_once()
assert not any(a.kind == "mcp_login" for a in app._deferred_actions)
finally:
app._shell_running = False
@@ -20613,6 +20680,175 @@ class TestMCPLoginCommand:
restart.assert_awaited_once_with("notion")
assert app._pending_mcp_reconnect is False
async def test_prompt_mcp_reconnect_restart_choice_queues_while_busy(self) -> None:
"""Accepting reconnect during a live run queues the restart, not runs it.
Also pins the two things the queued path must get right: the just-
authenticated server is optimistically flipped to `awaiting_reconnect`
so `/mcp` stops calling it unauthenticated during the busy window, and
draining once idle fires exactly one restart carrying the real server
name.
"""
from deepagents_code.mcp_tools import MCPServerInfo
app = DeepAgentsApp(agent=MagicMock())
async with app.run_test() as pilot:
await pilot.pause()
app._agent_running = True
app._mcp_server_info = [
MCPServerInfo(
name="notion",
transport="http",
status="unauthenticated",
error="Login required.",
),
]
def _push_screen(_screen: object, callback: Any) -> None: # noqa: ANN401 # callback signature matches Textual's variant
callback("reconnect")
try:
with (
patch.object(app, "push_screen", side_effect=_push_screen),
patch.object(
app, "_restart_server_for_mcp_refresh", new=AsyncMock()
) as restart,
patch.object(app, "notify") as notify,
):
await app._prompt_mcp_reconnect("notion")
# Restart is deferred, not run mid-generation.
restart.assert_not_called()
assert app._pending_mcp_reconnect is True
queued = [
a for a in app._deferred_actions if a.kind == "mcp_reconnect"
]
assert len(queued) == 1
message = notify.call_args.args[0]
assert "notion" in message
assert "current task completes" in message
# Optimistic state: notion now reads as awaiting_reconnect,
# not unauthenticated, so `_apply_optimistic_..._state`
# cannot be silently dropped without failing here.
notion = next(s for s in app._mcp_server_info if s.name == "notion")
assert notion.status == "awaiting_reconnect"
assert app._pending_mcp_login_reconnect is True
# Once the run ends, the drain fires one restart with the
# real server name (not the "pending login" sentinel).
app._agent_running = False
await app._drain_deferred_actions()
restart.assert_awaited_once_with("notion")
finally:
app._agent_running = False
async def test_deferred_mcp_reconnect_restarts_when_still_pending(self) -> None:
"""The queued reconnect restarts while the pending flag is still set."""
app = DeepAgentsApp(agent=MagicMock())
async with app.run_test() as pilot:
await pilot.pause()
app._pending_mcp_reconnect = True
with patch.object(
app, "_restart_server_for_mcp_refresh", new=AsyncMock()
) as restart:
await app._run_deferred_mcp_reconnect("notion")
restart.assert_awaited_once_with("notion")
async def test_deferred_mcp_reconnect_skips_when_already_reconnected(self) -> None:
"""A manual `/mcp reconnect` before the drain clears the pending flag.
The queued action must then be a no-op so the server is not restarted
a second time.
"""
app = DeepAgentsApp(agent=MagicMock())
async with app.run_test() as pilot:
await pilot.pause()
app._pending_mcp_reconnect = False
with patch.object(
app, "_restart_server_for_mcp_refresh", new=AsyncMock()
) as restart:
await app._run_deferred_mcp_reconnect("notion")
restart.assert_not_called()
async def test_reconnect_queue_dedupes_across_prompt_and_command(self) -> None:
"""Prompt-queued reconnect + `/mcp reconnect` while busy = one restart.
`_defer_action` is last-write-wins per kind, so the `/mcp reconnect`
subcommand (sentinel name) replaces the prompt's queued action (real
name). Only one survives, and the drain fires a single restart — the
respawn reloads every token regardless of which name it carries.
"""
app = DeepAgentsApp(agent=MagicMock())
async with app.run_test() as pilot:
await pilot.pause()
app._agent_running = True
def _push_screen(_screen: object, callback: Any) -> None: # noqa: ANN401 # callback signature matches Textual's variant
callback("reconnect")
try:
with (
patch.object(app, "push_screen", side_effect=_push_screen),
patch.object(
app, "_restart_server_for_mcp_refresh", new=AsyncMock()
) as restart,
patch.object(app, "notify"),
):
# Accept "reconnect now" mid-run: queues the real name.
await app._prompt_mcp_reconnect("notion")
# Then `/mcp reconnect` mid-run: queues the sentinel,
# replacing the prior action.
await app._handle_mcp_reconnect_command()
queued = [
a for a in app._deferred_actions if a.kind == "mcp_reconnect"
]
assert len(queued) == 1
app._agent_running = False
await app._drain_deferred_actions()
restart.assert_awaited_once()
finally:
app._agent_running = False
async def test_interrupt_warns_when_dropping_queued_reconnect(self) -> None:
"""Interrupting a run that queued a reconnect surfaces a recovery notice.
The token is already on disk, but the "will reconnect once the task
completes" promise no longer holds once the queue is discarded, so the
drop must not be silent.
"""
app = DeepAgentsApp(agent=MagicMock())
async with app.run_test() as pilot:
await pilot.pause()
app._defer_action(
DeferredAction(kind="mcp_reconnect", execute=AsyncMock()),
)
with patch.object(app, "notify") as notify:
app._cancel_worker(None)
assert not app._deferred_actions
messages = [call.args[0] for call in notify.call_args_list if call.args]
assert any("Cancelled the queued MCP reconnect" in m for m in messages)
async def test_restart_path_drops_queued_reconnect_without_warning(self) -> None:
"""A full restart fulfills the reconnect, so it fires no cancel notice."""
app = DeepAgentsApp(agent=MagicMock())
async with app.run_test() as pilot:
await pilot.pause()
app._defer_action(
DeferredAction(kind="mcp_reconnect", execute=AsyncMock()),
)
with patch.object(app, "notify") as notify:
app._cancel_worker(None, abort_pending_reconnect=False)
assert not app._deferred_actions
messages = [call.args[0] for call in notify.call_args_list if call.args]
assert not any("Cancelled the queued MCP reconnect" in m for m in messages)
async def test_prompt_mcp_reconnect_restart_choice_clears_splash_prompts(
self,
) -> None:
@@ -20721,6 +20957,49 @@ class TestMCPLoginCommand:
await pilot.pause()
restart.assert_awaited_once()
@pytest.mark.parametrize(
("agent_running", "shell_running"),
[(True, False), (False, True)],
)
async def test_mcp_reconnect_subcommand_queues_while_busy(
self, *, agent_running: bool, shell_running: bool
) -> None:
"""`/mcp reconnect` waits for an active agent or shell task to finish."""
app = DeepAgentsApp(agent=MagicMock())
async with app.run_test() as pilot:
await pilot.pause()
app._pending_mcp_login_reconnect = True
app._sync_pending_mcp_reconnect()
app._agent_running = agent_running
app._shell_running = shell_running
try:
with (
patch.object(
app, "_restart_server_for_mcp_refresh", new=AsyncMock()
) as restart,
patch.object(app, "notify") as notify,
):
await app._handle_command("/mcp reconnect")
restart.assert_not_called()
queued = [
action
for action in app._deferred_actions
if action.kind == "mcp_reconnect"
]
assert len(queued) == 1
assert "current task completes" in notify.call_args.args[0]
app._agent_running = False
app._shell_running = False
await app._drain_deferred_actions()
restart.assert_awaited_once_with("pending login")
finally:
app._agent_running = False
app._shell_running = False
async def test_mcp_reconnect_subcommand_noop_when_not_pending(self) -> None:
"""`/mcp reconnect` surfaces a notice and does nothing when idle."""
app = DeepAgentsApp(agent=MagicMock())