fix(py): make cleanup() best-effort and prevent duplicate mod_server() retries

pkg-py/CHANGELOG.md

@@ -5,6 +5,14 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [Unreleased]
+
+### Bug fixes
+
+* `.server()` now counts a session as live only after its setup has fully succeeded, so a failed `.server()` call no longer blocks `add_table()`/`remove_table()` or defers replacement cleanup for as long as the session lives. In Shiny Express apps, a session whose setup fails this way no longer risks a duplicate `mod_server()` attempt (and duplicate reactive-effect/bookmark registration) on a later lazy call from `.df()`, `.sql()`, etc.
+
+* `cleanup()` now always flushes resources retired by concurrent sessions, even if cleaning a current data source or the query executor throws, and no longer aborts partway through a failure: remaining data sources and owned chatlas clients still get cleaned up, with a warning for each failure. Retired resources whose cleanup fails are retained and retried (with a warning) on the next flush instead of being dropped permanently or retried silently.
+
 ## [0.8.0] - 2026-09-12
 
 ### New features

pkg-py/src/querychat/_querychat_base.py

@@ -781,6 +781,8 @@ def _mark_server_initialized(self, session) -> None:
         The add/remove_table guards and cleanup-on-replace in
         ``server(data_source=...)`` key off the number of *live* sessions:
         a session that has ended can no longer be using a replaced resource.
+        Call this only after ``mod_server()`` has returned successfully --
+        counting a session whose setup failed would treat it as live.
         """
         self._active_sessions += 1
 
@@ -801,9 +803,14 @@ def _flush_retired_resources(self) -> None:
         retired = self._retired_resources
         self._retired_resources = []
         for resource in retired:
-            # Best-effort: one failing cleanup must not leave the rest open.
-            with contextlib.suppress(Exception):
+            # Best-effort: one failing cleanup must not leave the rest open,
+            # but retain the failed entry so a later flush (or cleanup()) can
+            # retry it.
+            try:
                 resource.cleanup()
+            except Exception as e:  # noqa: PERF203 (teardown of a few resources, not a hot loop)
+                warnings.warn(f"Failed to clean up retired resource: {e}", stacklevel=2)
+                self._retired_resources.append(resource)
 
     def cleanup(self) -> None:
         """
@@ -821,9 +828,15 @@ def cleanup(self) -> None:
         when the app shuts down (e.g., via `atexit`).
         """
         if self._query_executor is not None:
-            self._query_executor.cleanup()
+            try:
+                self._query_executor.cleanup()
+            except Exception as e:
+                warnings.warn(f"Failed to clean up query executor: {e}", stacklevel=2)
         for source in self._data_sources.values():
-            source.cleanup()
+            try:
+                source.cleanup()
+            except Exception as e:  # noqa: PERF203 (teardown of a few resources, not a hot loop)
+                warnings.warn(f"Failed to clean up data source: {e}", stacklevel=2)
         self._flush_retired_resources()
         for client in self._owned_clients:
             # Best-effort: one provider's close() failing must not leave the

pkg-py/src/querychat/_shiny.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import warnings
+import weakref
 from typing import TYPE_CHECKING, Any, Literal, Optional, overload
 
 import chatlas
@@ -407,7 +408,6 @@ def app_ui(request):
             )
 
         def app_server(input: Inputs, output: Outputs, session: Session):
-            self._mark_server_initialized(session)
             if enable_bookmarking:
                 session.bookmark.exclude.extend(["reset_query", "sql_editor"])
             vals = mod_server(
@@ -422,6 +422,7 @@ def app_server(input: Inputs, output: Outputs, session: Session):
                 greeting_base=None,
                 greeting_tables=list(self.greeter.tables),
             )
+            self._mark_server_initialized(session)
 
             @reactive.calc
             def active_table_name() -> str:
@@ -760,8 +761,7 @@ def create_session_client(**kwargs) -> chatlas.Chat:
             )
         )
 
-        self._mark_server_initialized(session)
-        return mod_server(
+        result = mod_server(
             id or self.id,
             data_sources=dict(self._data_sources),
             executor=self._require_query_executor("server"),
@@ -773,6 +773,8 @@ def create_session_client(**kwargs) -> chatlas.Chat:
             greeting_base=resolved_client,
             greeting_tables=list(self.greeter.tables),
         )
+        self._mark_server_initialized(session)
+        return result
 
 
 class QueryChatExpress(QueryChatBase[IntoFrameT]):
@@ -1021,6 +1023,7 @@ def __init__(
 
         self._enable_bookmarking = enable_bookmarking
         self._vals: ServerValues[IntoFrameT] | None = None
+        self._attempted_sessions: weakref.WeakSet[Session] = weakref.WeakSet()
 
     def _ensure_server_started(self) -> None:
         """
@@ -1030,6 +1033,10 @@ def _ensure_server_started(self) -> None:
         module-level add_table() calls (which happen after __init__ but before
         sidebar()/ui()) can complete before server initialization locks the
         table set.
+
+        Each session gets at most one mod_server() attempt: retrying after a
+        failed attempt would re-register its non-idempotent reactive effects
+        and bookmark/history hooks a second time.
         """
         if self._active_sessions > 0:
             return
@@ -1038,8 +1045,10 @@ def _ensure_server_started(self) -> None:
             return
         if not self._data_sources:
             return
+        if session in self._attempted_sessions:
+            return
+        self._attempted_sessions.add(session)
         self._require_initialized("_ensure_server_started")
-        self._mark_server_initialized(session)
         resolved_history: bool | HistoryOptions = (
             self.history
             if self.history is not None
@@ -1061,6 +1070,7 @@ def _ensure_server_started(self) -> None:
             greeting_base=None,
             greeting_tables=list(self.greeter.tables),
         )
+        self._mark_server_initialized(session)
 
     def sidebar(
         self,

pkg-py/tests/test_cleanup.py

@@ -252,6 +252,118 @@ def test_cleanup_cleans_retired_resources(self, sample_df, ended_callbacks):
             executor_cleanup.assert_called_once()
         assert qc._retired_resources == []
 
+    def test_cleanup_flushes_retired_resources_when_a_cleanup_fails(
+        self, sample_df, ended_callbacks
+    ):
+        """
+        The retired-resource flush is the fallback cleanup path; a failing
+        source cleanup must not prevent it from running, nor abort the rest
+        of cleanup().
+        """
+        qc = shiny_mod.QueryChat(sample_df, "users")
+        qc.server()
+        old_source = qc._data_sources["users"]
+        qc.server(data_source=sample_df.copy())  # retires old_source
+
+        with (
+            patch.object(old_source, "cleanup") as retired_cleanup,
+            patch.object(
+                qc._data_sources["users"],
+                "cleanup",
+                side_effect=RuntimeError("boom"),
+            ),
+            pytest.warns(UserWarning, match="Failed to clean up data source"),
+        ):
+            qc.cleanup()
+        retired_cleanup.assert_called_once()
+
+    def test_cleanup_continues_to_other_sources_after_one_fails(
+        self, monkeypatch, sample_df
+    ):
+        """
+        One data source's cleanup() raising must not skip the rest of
+        cleanup(): remaining sources and owned clients still get closed.
+        """
+        monkeypatch.setenv("OPENAI_API_KEY", "sk-dummy-key-for-testing")
+        qc = QueryChatBase(sample_df, "users", client="openai")
+        qc.add_table(sample_df.copy(), "more_users")
+        failing_source = qc._data_sources["users"]
+        other_source = qc._data_sources["more_users"]
+
+        with (
+            patch.object(failing_source, "cleanup", side_effect=RuntimeError("boom")),
+            patch.object(other_source, "cleanup") as other_cleanup,
+            pytest.warns(UserWarning, match="Failed to clean up data source"),
+        ):
+            qc.cleanup()
+        other_cleanup.assert_called_once()
+        assert qc._base_client.provider._client.is_closed()
+
+    def test_flush_retired_resources_warns_when_cleanup_fails(
+        self, sample_df, ended_callbacks
+    ):
+        qc = shiny_mod.QueryChat(sample_df, "users")
+        qc.server()
+        old_source = qc._data_sources["users"]
+        qc.server(data_source=sample_df.copy())  # retires old_source
+
+        def end_sessions() -> None:
+            for cb in ended_callbacks:
+                cb()
+
+        with (
+            patch.object(old_source, "cleanup", side_effect=RuntimeError("boom")),
+            pytest.warns(UserWarning, match="Failed to clean up retired resource"),
+        ):
+            end_sessions()
+        assert old_source in qc._retired_resources
+
+    def test_failed_retired_cleanup_is_retained_and_retried(
+        self, monkeypatch, sample_df
+    ):
+        """A transient cleanup failure must not permanently drop the resource."""
+        sessions = []
+
+        def next_session():
+            session = MagicMock()
+            session._ended_callbacks = []
+            session.on_ended = session._ended_callbacks.append
+            sessions.append(session)
+            return session
+
+        monkeypatch.setattr(shiny_mod, "get_current_session", next_session)
+        monkeypatch.setattr(shiny_mod, "mod_server", lambda *args, **kwargs: None)
+
+        qc = shiny_mod.QueryChat(sample_df, "users")
+        qc.server()
+        old_source = qc._data_sources["users"]
+        qc.server(data_source=sample_df.copy())  # session 2 retires old_source
+
+        attempts = 0
+
+        def fail_once():
+            nonlocal attempts
+            attempts += 1
+            if attempts == 1:
+                raise RuntimeError("transient failure")
+
+        def end_session(index: int) -> None:
+            for cb in sessions[index]._ended_callbacks:
+                cb()
+
+        with patch.object(old_source, "cleanup", side_effect=fail_once):
+            end_session(0)
+            assert attempts == 0  # session 2 still live: no flush yet
+
+            with pytest.warns(UserWarning, match="Failed to clean up retired resource"):
+                end_session(1)
+            assert attempts == 1  # flush ran; cleanup failed transiently
+            assert old_source in qc._retired_resources
+
+            qc.cleanup()  # retries the retained resource
+            assert attempts == 2
+            assert qc._retired_resources == []
+
 
 class TestCleanupDataSources:
     """Existing executor/source cleanup behavior is preserved."""

pkg-py/tests/test_server_data_source.py

@@ -290,6 +290,28 @@ def test_add_table_allowed_once_all_sessions_have_ended(
         qc.add_table(other_users_df, "other")  # must not raise
         assert qc.table_names() == ["users", "other"]
 
+    def test_failed_server_call_does_not_count_as_live_session(
+        self, users_df, other_users_df, monkeypatch
+    ):
+        """
+        A .server() call that fails mid-setup (after registration) must not
+        linger in the live-session count, which would block add_table() and
+        defer replacement cleanup forever.
+        """
+
+        def boom(*args, **kwargs):
+            raise RuntimeError("mod_server failed")
+
+        monkeypatch.setattr(shiny_mod, "mod_server", boom)
+        monkeypatch.setattr(shiny_mod, "get_current_session", lambda: MagicMock())
+
+        qc = shiny_mod.QueryChat(None, table_name="users")
+        with pytest.raises(RuntimeError, match="mod_server failed"):
+            qc.server(data_source=users_df)
+
+        qc.add_table(other_users_df, "other")  # must not raise
+        assert qc.table_names() == ["users", "other"]
+
 
 class TestServerDataSourceGreetingSnapshot:
     def test_server_passes_greeting_tables_snapshot_to_mod_server(

pkg-py/tests/test_shiny.py

@@ -186,6 +186,44 @@ def fake_mod_server(*args, **kwargs):
     assert captured["history"].restore_mode == "bookmark"
 
 
+def test_ensure_server_started_does_not_retry_after_failed_attempt(monkeypatch):
+    """
+    A mod_server() failure must not be retried within the same session by a
+    later lazy call (e.g. from .df()/.sql()/.ui()) -- retrying would
+    re-register mod_server()'s non-idempotent reactive effects and
+    bookmark/history hooks a second time.
+    """
+    from unittest.mock import MagicMock
+
+    import pandas as pd
+    from querychat._shiny import QueryChatExpress
+    from shiny._namespaces import Root
+    from shiny.session import session_context
+
+    calls = 0
+
+    def failing_mod_server(*args, **kwargs):
+        nonlocal calls
+        calls += 1
+        raise RuntimeError("mod_server failed")
+
+    monkeypatch.setattr("querychat._shiny.mod_server", failing_mod_server)
+
+    mock_session = MagicMock()
+    mock_session.ns = Root
+    with session_context(mock_session):
+        qc = QueryChatExpress(pd.DataFrame({"a": [1, 2, 3]}), "a_table")
+
+        with pytest.raises(RuntimeError, match="mod_server failed"):
+            qc._ensure_server_started()
+        assert calls == 1
+
+        # A later lazy call must not retry mod_server() a second time.
+        with pytest.raises(RuntimeError, match="not initialized"):
+            qc._require_vals()
+        assert calls == 1
+
+
 def test_express_explicit_enable_bookmarking_warns():
     from unittest.mock import MagicMock, patch