fix(py): scope server(data_source=) cleanup and table guards to live sessions

pkg-py/src/querychat/_querychat_base.py

@@ -100,8 +100,14 @@ def __init__(
         self._data_sources: dict[str, DataSource] = {}
         self._query_executor: QueryExecutor | None = None
 
-        # Track server initialization state for add/remove table validation
-        self._server_initialized = False
+        # Live Shiny session count. Shared-state guards key off this: an
+        # ended session can no longer be using a resource it registered.
+        self._active_sessions = 0
+
+        # Sources/executors replaced while sessions were still live. Their
+        # cleanup is deferred until the last live session ends (or cleanup()
+        # runs) so a still-running session never loses a resource it uses.
+        self._retired_resources: list[DataSource | QueryExecutor] = []
 
         # Name to register at .server(data_source=...) time when constructed
         # with data_source=None (the deferred pattern).
@@ -143,6 +149,15 @@ def __init__(
                     )
             self.add_table(data_source, table_name, include_in_greeting=True)
         else:
+            # Validate now: a bad deferred name would otherwise surface at
+            # .server() registration time, after the module id is built.
+            if table_name is not None and not re.match(
+                r"^[a-zA-Z][a-zA-Z0-9_]*$", table_name
+            ):
+                raise ValueError(
+                    "Table name must begin with a letter and contain only "
+                    "letters, numbers, and underscores"
+                )
             self._deferred_table_name = table_name
 
     def _build_system_prompt(
@@ -487,12 +502,12 @@ def add_table(
         ValueError
             If table_name already exists (and replace=False) or is invalid.
         RuntimeError
-            If called after server() has been invoked.
+            If called while a server session is active.
 
         """
-        if self._server_initialized:
+        if self._active_sessions > 0:
             raise RuntimeError(
-                "Cannot add tables after server initialization. "
+                "Cannot add tables while a server session is active. "
                 "Add all tables before calling .server() or .app()."
             )
         self._add_or_replace_table(
@@ -516,14 +531,14 @@ def _add_or_replace_table(
 
         Guard-free core of :meth:`add_table`, also called directly by
         ``.server(data_source=...)`` so each session can register its own
-        table even after an earlier session's ``.server()`` call has set
-        ``_server_initialized``.
+        table even while earlier sessions are still running.
 
         ``cleanup_replaced=False`` is for that per-session path: the
         replaced table may still be in active use by an earlier,
         still-running session, so cleaning it up here would pull the
-        resource out from under it. Cleanup becomes the caller's
-        responsibility (e.g. via ``session.on_ended()``).
+        resource out from under it. The retired source and cached executor
+        are retained and cleaned up once the last live session ends
+        (see ``_mark_server_initialized``).
         """
         if not isinstance(include_in_greeting, bool):
             raise TypeError(
@@ -560,12 +575,17 @@ def _add_or_replace_table(
 
         old_source = self._data_sources.get(table_name)
         self._data_sources = next_data_sources
-        if cleanup_replaced and old_source is not None and old_source is not normalized:
-            old_source.cleanup()
+        if old_source is not None and old_source is not normalized:
+            if cleanup_replaced:
+                old_source.cleanup()
+            else:
+                self._retired_resources.append(old_source)
         if self._query_executor is not None:
             if cleanup_replaced:
                 with contextlib.suppress(Exception):
                     self._query_executor.cleanup()
+            else:
+                self._retired_resources.append(self._query_executor)
             self._query_executor = None
 
         if include_in_greeting and table_name not in self.greeter.tables:
@@ -611,7 +631,7 @@ def add_tables(  # noqa: PLR0912
             If the resolved table list is empty, any name is invalid, or any
             name already exists (and ``replace=False``).
         RuntimeError
-            If called after :meth:`server` has been invoked.
+            If called while a server session is active.
 
         Examples
         --------
@@ -631,9 +651,9 @@ def add_tables(  # noqa: PLR0912
         >>> qc.add_tables(backend)
 
         """
-        if self._server_initialized:
+        if self._active_sessions > 0:
             raise RuntimeError(
-                "Cannot add tables after server initialization. "
+                "Cannot add tables while a server session is active. "
                 "Add all tables before calling .server() or .app()."
             )
 
@@ -722,12 +742,12 @@ def remove_table(self, table_name: str) -> None:
         ValueError
             If table doesn't exist or is the last remaining table.
         RuntimeError
-            If called after server() has been invoked.
+            If called while a server session is active.
 
         """
-        if self._server_initialized:
+        if self._active_sessions > 0:
             raise RuntimeError(
-                "Cannot remove tables after server initialization. "
+                "Cannot remove tables while a server session is active. "
                 "Configure all tables before calling .server() or .app()."
             )
 
@@ -754,9 +774,36 @@ def remove_table(self, table_name: str) -> None:
             self._query_executor = None
         removed_source.cleanup()
 
-    def _mark_server_initialized(self) -> None:
-        """Mark that the server has been initialized. Prevents add/remove_table."""
-        self._server_initialized = True
+    def _mark_server_initialized(self, session) -> None:
+        """
+        Track a newly started session until it ends.
+
+        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.
+        """
+        self._active_sessions += 1
+
+        def untrack_session() -> None:
+            self._active_sessions -= 1
+            if self._active_sessions == 0:
+                self._flush_retired_resources()
+
+        session.on_ended(untrack_session)
+
+    def _flush_retired_resources(self) -> None:
+        """
+        Clean up resources retired while sessions were still live.
+
+        Retired sources/executors may still be in use by a live session, so
+        this only runs once no sessions remain (or from ``cleanup()``).
+        """
+        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):
+                resource.cleanup()
 
     def cleanup(self) -> None:
         """
@@ -777,6 +824,7 @@ def cleanup(self) -> None:
             self._query_executor.cleanup()
         for source in self._data_sources.values():
             source.cleanup()
+        self._flush_retired_resources()
         for client in self._owned_clients:
             # Best-effort: one provider's close() failing must not leave the
             # remaining owned clients open.

pkg-py/src/querychat/_shiny.py

@@ -407,7 +407,7 @@ def app_ui(request):
             )
 
         def app_server(input: Inputs, output: Outputs, session: Session):
-            self._mark_server_initialized()
+            self._mark_server_initialized(session)
             if enable_bookmarking:
                 session.bookmark.exclude.extend(["reset_query", "sql_editor"])
             vals = mod_server(
@@ -717,7 +717,9 @@ def server(
                 resolved_table_name,
                 replace=True,
                 include_in_greeting=True,
-                cleanup_replaced=False,
+                # A live session may still be using the replaced source,
+                # so defer its cleanup until no sessions are active.
+                cleanup_replaced=self._active_sessions == 0,
             )
 
         self._require_initialized("server")
@@ -758,7 +760,7 @@ def create_session_client(**kwargs) -> chatlas.Chat:
             )
         )
 
-        self._mark_server_initialized()
+        self._mark_server_initialized(session)
         return mod_server(
             id or self.id,
             data_sources=dict(self._data_sources),
@@ -1029,15 +1031,15 @@ def _ensure_server_started(self) -> None:
         sidebar()/ui()) can complete before server initialization locks the
         table set.
         """
-        if self._server_initialized:
+        if self._active_sessions > 0:
             return
         session = get_current_session()
-        if isinstance(session, ExpressStubSession):
+        if session is None or isinstance(session, ExpressStubSession):
             return
         if not self._data_sources:
             return
         self._require_initialized("_ensure_server_started")
-        self._mark_server_initialized()
+        self._mark_server_initialized(session)
         resolved_history: bool | HistoryOptions = (
             self.history
             if self.history is not None

pkg-py/tests/test_base.py

@@ -372,8 +372,8 @@ def test_empty_list_raises(self, multi_table_engine):
 
     def test_after_server_raises(self, multi_table_engine):
         qc = QueryChatBase()
-        qc._server_initialized = True
-        with pytest.raises(RuntimeError, match="Cannot add tables after server"):
+        qc._active_sessions = 1
+        with pytest.raises(RuntimeError, match="Cannot add tables while a server session"):
             qc.add_tables(multi_table_engine)
 
     def test_system_prompt_built_exactly_once(self, multi_table_engine):

pkg-py/tests/test_cleanup.py

@@ -153,21 +153,22 @@ def test_owned_override_closed_on_session_end(
         qc = shiny_mod.QueryChat(sample_df, "users")
         qc.server(client="openai")
 
-        assert len(ended_callbacks) == 1
         (override,) = qc._owned_clients
-        ended_callbacks[0]()
+        for cb in ended_callbacks:
+            cb()
         assert override.provider._client.is_closed()
         assert qc._owned_clients == []
 
-    def test_user_supplied_override_gets_no_on_ended(
+    def test_user_supplied_override_not_closed_on_session_end(
         self, monkeypatch, sample_df, ended_callbacks
     ):
         monkeypatch.setenv("OPENAI_API_KEY", "sk-dummy-key-for-testing")
         qc = shiny_mod.QueryChat(sample_df, "users")
         chat = ChatOpenAI()
         qc.server(client=chat)
 
-        assert ended_callbacks == []
+        for cb in ended_callbacks:
+            cb()
         qc.cleanup()
         assert not chat.provider._client.is_closed()
 
@@ -184,6 +185,74 @@ def test_owned_override_tracked_and_closed_by_cleanup(
         assert override.provider._client.is_closed()
 
 
+class TestRetiredResourceCleanup:
+    """Resources replaced while sessions are live are cleaned once they end."""
+
+    @pytest.fixture
+    def ended_callbacks(self, monkeypatch):
+        callbacks = []
+        fake_session = MagicMock()
+        fake_session.on_ended = callbacks.append
+        monkeypatch.setattr(shiny_mod, "get_current_session", lambda: fake_session)
+        monkeypatch.setattr(shiny_mod, "mod_server", lambda *args, **kwargs: None)
+        return callbacks
+
+    def test_replacement_with_live_session_defers_cleanup(
+        self, sample_df, ended_callbacks
+    ):
+        qc = shiny_mod.QueryChat(sample_df, "users")
+        qc.server()
+        old_source = qc._data_sources["users"]
+        old_executor = qc._query_executor
+
+        # Second session replaces the table while the first is still live
+        replacement = sample_df.copy()
+        qc.server(data_source=replacement)
+
+        assert old_source in qc._retired_resources
+        assert old_executor in qc._retired_resources
+        with (
+            patch.object(old_source, "cleanup") as source_cleanup,
+            patch.object(old_executor, "cleanup") as executor_cleanup,
+        ):
+            for cb in ended_callbacks:
+                cb()
+            source_cleanup.assert_called_once()
+            executor_cleanup.assert_called_once()
+        assert qc._retired_resources == []
+
+    def test_replacement_without_live_session_cleans_immediately(
+        self, sample_df, ended_callbacks
+    ):
+        qc = shiny_mod.QueryChat(sample_df, "users")
+        qc.server()
+        for cb in ended_callbacks:
+            cb()
+        old_source = qc._data_sources["users"]
+
+        with patch.object(old_source, "cleanup") as source_cleanup:
+            qc.server(data_source=sample_df.copy())
+            source_cleanup.assert_called_once()
+        assert qc._retired_resources == []
+
+    def test_cleanup_cleans_retired_resources(self, sample_df, ended_callbacks):
+        qc = shiny_mod.QueryChat(sample_df, "users")
+        qc.server()
+        old_source = qc._data_sources["users"]
+        old_executor = qc._query_executor
+        qc.server(data_source=sample_df.copy())
+
+        # Sessions never end: cleanup() still releases retired resources
+        with (
+            patch.object(old_source, "cleanup") as source_cleanup,
+            patch.object(old_executor, "cleanup") as executor_cleanup,
+        ):
+            qc.cleanup()
+            source_cleanup.assert_called_once()
+            executor_cleanup.assert_called_once()
+        assert qc._retired_resources == []
+
+
 class TestCleanupDataSources:
     """Existing executor/source cleanup behavior is preserved."""
 

pkg-py/tests/test_deferred_shiny.py

@@ -125,8 +125,8 @@ def test_multiple_server_overrides_do_not_leak_into_shared_state(self, sample_df
         with session_context(ExpressStubSession()):
             qc.server(client=first_override)
 
-        # Reset server_initialized for sequential test
-        qc._server_initialized = False
+        # Reset live-session count for sequential test
+        qc._active_sessions = 0
 
         with session_context(ExpressStubSession()):
             qc.server(client=second_override)
@@ -151,17 +151,17 @@ def test_add_table_does_not_raise_after_init_stub_session(
         """add_table() must succeed after __init__ during stub session."""
         with session_context(ExpressStubSession()):
             qc = ExpressQueryChat(orders_df, "orders")
-            # Without the fix, _server_initialized would be True here and
+            # Without the fix, a session would be tracked here and
             # add_table() would raise RuntimeError.
             qc.add_table(customers_df, "customers")
 
         assert qc.table_names() == ["orders", "customers"]
 
     def test_server_not_initialized_after_init_stub_session(self, orders_df):
-        """_server_initialized must remain False after __init__ in stub session."""
+        """No session may be tracked after __init__ in a stub session."""
         with session_context(ExpressStubSession()):
             qc = ExpressQueryChat(orders_df, "orders")
-            assert not qc._server_initialized
+            assert qc._active_sessions == 0
 
     def test_ensure_server_started_noop_during_stub_session(
         self, orders_df, monkeypatch
@@ -196,11 +196,11 @@ def fake_mod_server(*args, **kwargs):
         mock_session.ns = Root
         with session_context(mock_session):
             qc = ExpressQueryChat(orders_df, "orders")
-            assert not qc._server_initialized
+            assert qc._active_sessions == 0
             qc._ensure_server_started()
 
         assert len(called) == 1
-        assert qc._server_initialized
+        assert qc._active_sessions == 1
 
     def test_ensure_server_started_idempotent(self, orders_df, monkeypatch):
         """_ensure_server_started() called twice starts server only once."""