fix(shell): reload MCP and model sessions cleanly

Author: elkaix

src/pythinker_code/app.py

@@ -467,6 +467,22 @@ def session(self) -> Session:
         """Get the Session instance."""
         return self._runtime.session
 
+    async def cleanup_runtime_resources(self) -> None:
+        """Cleanup per-CLI-instance resources without stopping persisted background tasks."""
+        # Cancel the startup managed-model refresh task if it is still running
+        # so it does not outlive this CLI instance across reloads.
+        if self._bg_refresh_task is not None and not self._bg_refresh_task.done():
+            self._bg_refresh_task.cancel()
+
+        # Cleanup MCP connections held by this instance's toolset. Background
+        # task workers are persisted elsewhere and intentionally left alone.
+        from pythinker_code.soul.toolset import PythinkerToolset
+
+        toolset = self._soul.agent.toolset
+        if isinstance(toolset, PythinkerToolset):
+            with contextlib.suppress(Exception):
+                await toolset.cleanup()
+
     async def shutdown_background_tasks(self) -> None:
         """Kill active background tasks on exit, unless keep_alive_on_exit is configured.
 
@@ -480,18 +496,7 @@ async def shutdown_background_tasks(self) -> None:
         store corruption must not propagate and replace the real exit code
         with a traceback.
         """
-        # Cancel the startup managed-model refresh task if it is still running
-        # so it does not outlive the CLI process.
-        if self._bg_refresh_task is not None and not self._bg_refresh_task.done():
-            self._bg_refresh_task.cancel()
-
-        # Cleanup MCP connections held by the toolset
-        from pythinker_code.soul.toolset import PythinkerToolset
-
-        toolset = self._soul.agent.toolset
-        if isinstance(toolset, PythinkerToolset):
-            with contextlib.suppress(Exception):
-                await toolset.cleanup()
+        await self.cleanup_runtime_resources()
 
         bg_config = self._runtime.config.background
         if bg_config.keep_alive_on_exit:

src/pythinker_code/cli/__init__.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import json
 import os
 from importlib import import_module
 from pathlib import Path
@@ -173,6 +174,51 @@ class ExitCode:
     RETRYABLE = 75  # EX_TEMPFAIL from sysexits.h
 
 
+def _load_mcp_configs_from_cli_inputs(
+    mcp_config_file: list[Path] | None,
+    mcp_config: list[str] | None,
+) -> list[Any]:
+    """Load MCP config JSON from the current CLI inputs.
+
+    This intentionally re-resolves the default global MCP file on every call so
+    `/reload` observes servers added after the process started.
+    """
+    from .mcp import get_global_mcp_config_file
+
+    file_configs = list(mcp_config_file or [])
+    raw_mcp_config = list(mcp_config or [])
+
+    # Use default MCP config file if no MCP config file is provided. Keep this
+    # lookup live for reloads: the file may be created after process startup.
+    if not file_configs:
+        default_mcp_file = get_global_mcp_config_file()
+        if default_mcp_file.exists():
+            file_configs.append(default_mcp_file)
+
+    configs: list[Any] = []
+    for conf in file_configs:
+        try:
+            configs.append(json.loads(conf.read_text(encoding="utf-8")))
+        except json.JSONDecodeError as e:
+            raise typer.BadParameter(
+                f"Invalid JSON in MCP config file {conf}: {e}",
+                param_hint="--mcp-config-file",
+            ) from e
+        except OSError as e:
+            raise typer.BadParameter(
+                f"Cannot read MCP config file {conf}: {e}",
+                param_hint="--mcp-config-file",
+            ) from e
+
+    for conf in raw_mcp_config:
+        try:
+            configs.append(json.loads(conf))
+        except json.JSONDecodeError as e:
+            raise typer.BadParameter(f"Invalid JSON: {e}", param_hint="--mcp-config") from e
+
+    return configs
+
+
 InputFormat = Literal["text", "stream-json"]
 OutputFormat = Literal["text", "stream-json"]
 
@@ -502,7 +548,6 @@ def pythinker(
     """Pythinker, your next CLI agent."""
     import asyncio
     import contextlib
-    import json
 
     from pythinker_code.utils.proctitle import init_process_name
 
@@ -525,8 +570,6 @@ def pythinker(
     from pythinker_code.ui.shell.startup import ShellStartupProgress
     from pythinker_code.utils.logging import logger, open_original_stderr, redirect_stderr_to_logger
 
-    from .mcp import get_global_mcp_config_file
-
     # Don't redirect stderr during argument parsing. Our stderr redirector
     # replaces fd=2 with a pipe, which would swallow Click/Typer startup errors.
     # Redirection is installed later, right before PythinkerCLI.create(), so that
@@ -648,25 +691,6 @@ def _emit_fatal_error(message: str) -> None:
     elif config_file is not None:
         config = config_file
 
-    file_configs = list(mcp_config_file or [])
-    raw_mcp_config = list(mcp_config or [])
-
-    # Use default MCP config file if no MCP config is provided
-    if not file_configs:
-        default_mcp_file = get_global_mcp_config_file()
-        if default_mcp_file.exists():
-            file_configs.append(default_mcp_file)
-
-    try:
-        mcp_configs = [json.loads(conf.read_text(encoding="utf-8")) for conf in file_configs]
-    except json.JSONDecodeError as e:
-        raise typer.BadParameter(f"Invalid JSON: {e}", param_hint="--mcp-config-file") from e
-
-    try:
-        mcp_configs += [json.loads(conf) for conf in raw_mcp_config]
-    except json.JSONDecodeError as e:
-        raise typer.BadParameter(f"Invalid JSON: {e}", param_hint="--mcp-config") from e
-
     # Honor --no-telemetry by exporting the env var before any subsystem (Sentry,
     # OTel, sink) reads it during PythinkerCLI.create.
     if no_telemetry:
@@ -748,6 +772,8 @@ async def _run(session_id: str | None, prefill_text: str | None = None) -> tuple
                 if changed:
                     session.save_state()
 
+            mcp_configs = _load_mcp_configs_from_cli_inputs(mcp_config_file, mcp_config)
+
             # Redirect stderr *before* PythinkerCLI.create() so that MCP server
             # subprocesses (e.g. mcp-remote OAuth debug logs) write to the log
             # file instead of polluting the user's terminal.  CLI argument
@@ -899,9 +925,11 @@ async def _run(session_id: str | None, prefill_text: str | None = None) -> tuple
                         timeout=5,
                     )
 
-                if not preserve_background_tasks:
+                if preserve_background_tasks:
+                    await instance.cleanup_runtime_resources()
+                else:
                     await instance.shutdown_background_tasks()
-                    await instance.await_bg_tasks_shutdown()
+                await instance.await_bg_tasks_shutdown()
 
             return session, exit_code
         finally:

src/pythinker_code/ui/shell/setup.py

@@ -223,4 +223,4 @@ def reload(app: Shell, args: str):
     """Reload configuration"""
     from pythinker_code.cli import Reload
 
-    raise Reload
+    raise Reload()

src/pythinker_code/ui/shell/slash.py

@@ -355,6 +355,15 @@ async def model(app: Shell, args: str):
     if model_changed and selected_model_cfg.provider == "managed:lm-studio":
         await _preload_lm_studio_model(selected_provider, selected_model_cfg.model)
 
+    if model_changed:
+        current_session = soul.runtime.session
+        session = await Session.create(current_session.work_dir)
+        session.state.additional_dirs = list(current_session.state.additional_dirs)
+        if session.state.additional_dirs:
+            session.save_state()
+        console.print(f"[{_t.success}]Starting fresh session for the new model...[/]")
+        raise Reload(session_id=session.id)
+
     raise Reload(session_id=soul.runtime.session.id)
 
 

tests/core/test_cli_reload.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+import json
+
+from pythinker_code.cli import _load_mcp_configs_from_cli_inputs
+from pythinker_code.cli.mcp import get_global_mcp_config_file
+
+
+def test_load_mcp_configs_rechecks_default_file_between_reloads() -> None:
+    """Reload must see MCP servers added after process startup."""
+    default_mcp_file = get_global_mcp_config_file()
+    assert not default_mcp_file.exists()
+    assert _load_mcp_configs_from_cli_inputs(None, None) == []
+
+    default_mcp_file.parent.mkdir(parents=True, exist_ok=True)
+    expected = {
+        "mcpServers": {"context7": {"url": "https://mcp.example.test", "transport": "http"}}
+    }
+    default_mcp_file.write_text(json.dumps(expected), encoding="utf-8")
+
+    assert _load_mcp_configs_from_cli_inputs(None, None) == [expected]

tests/ui_and_conv/test_shell_slash_commands.py

@@ -9,10 +9,12 @@
 from unittest.mock import Mock
 
 import pytest
+from pydantic import SecretStr
 from pythinker_core.message import Message
 from pythinker_host.path import HostPath
 
 from pythinker_code.cli import Reload
+from pythinker_code.config import Config, LLMModel, LLMProvider
 from pythinker_code.session import Session
 from pythinker_code.subagents.models import AgentTypeDefinition, ToolPolicy
 from pythinker_code.ui.shell.slash import ShellSlashCmdFunc, shell_mode_registry
@@ -89,6 +91,90 @@ def test_blackbox_style_slash_aliases_are_registered() -> None:
         assert command.name == canonical
 
 
+async def test_model_switch_starts_fresh_session(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Changing models should reload into a new session so old context is not reused."""
+    from pythinker_code.soul.pythinkersoul import PythinkerSoul
+    from pythinker_code.ui.shell import model_picker
+    from pythinker_code.ui.shell import slash as shell_slash
+
+    config = Config(
+        is_from_default_location=True,
+        default_model="model-a",
+        providers={
+            "test-provider": LLMProvider(
+                type="pythinker",
+                base_url="https://example.test",
+                api_key=SecretStr("test-key"),
+            )
+        },
+        models={
+            "model-a": LLMModel(
+                provider="test-provider",
+                model="model-a",
+                max_context_size=100_000,
+            ),
+            "model-b": LLMModel(
+                provider="test-provider",
+                model="model-b",
+                max_context_size=100_000,
+            ),
+        },
+    )
+    current_session = SimpleNamespace(
+        id="current-session",
+        work_dir=HostPath("/tmp/work"),
+        state=SimpleNamespace(additional_dirs=["/extra"]),
+    )
+    mock_soul = Mock(spec=PythinkerSoul)
+    mock_soul.runtime = SimpleNamespace(
+        config=config,
+        llm=SimpleNamespace(model_config=config.models["model-a"]),
+        session=current_session,
+    )
+    mock_soul.thinking_effort = "off"
+    mock_soul.thinking = False
+    shell = SimpleNamespace(soul=mock_soul)
+
+    class _ModelPicker:
+        def __init__(self, *args: Any, **kwargs: Any) -> None:
+            pass
+
+        async def run(self) -> str:
+            return "model-b"
+
+    saved_configs: list[Config] = []
+
+    async def fake_refresh_managed_models(_config: Config) -> None:
+        return None
+
+    fresh_session = SimpleNamespace(
+        id="fresh-session",
+        state=SimpleNamespace(additional_dirs=[]),
+        save_state=Mock(),
+    )
+
+    async def fake_create_session(work_dir: HostPath) -> SimpleNamespace:
+        assert work_dir == current_session.work_dir
+        return fresh_session
+
+    monkeypatch.setattr(shell_slash, "refresh_managed_models", fake_refresh_managed_models)
+    monkeypatch.setattr(model_picker, "ModelPickerApp", _ModelPicker)
+    monkeypatch.setattr(shell_slash, "load_config", lambda: config.model_copy(deep=True))
+    monkeypatch.setattr(shell_slash, "save_config", saved_configs.append)
+    monkeypatch.setattr(shell_slash.Session, "create", fake_create_session)
+    monkeypatch.setattr(shell_slash.console, "print", Mock())
+
+    cmd = shell_slash_registry.find_command("model")
+    assert cmd is not None
+    with pytest.raises(Reload) as exc_info:
+        await _invoke_slash_command(cmd, shell)
+
+    assert exc_info.value.session_id == "fresh-session"
+    assert fresh_session.state.additional_dirs == ["/extra"]
+    fresh_session.save_state.assert_called_once_with()
+    assert saved_configs[-1].default_model == "model-b"
+
+
 async def test_mcp_slash_persists_only_final_snapshot(
     monkeypatch: pytest.MonkeyPatch, capsys: Any
 ) -> None: