fix(py): don't clean up a table replaced via server(data_source=)

pkg-py/CHANGELOG.md

@@ -31,7 +31,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
       qc.server(data_source=conn.table("my_table"), client=chat_client)
   ```
 
-  Registering a table this way is no longer blocked by an earlier session having already registered one.
+  Registering a table this way is no longer blocked by an earlier session having already registered one, and no longer tears down a still-in-use data source from an earlier session (replacing it via `server(data_source=)` leaves the replaced source's cleanup to whoever created it).
 
 ### Improvements
 

pkg-py/src/querychat/_querychat_base.py

@@ -505,6 +505,7 @@ def _add_or_replace_table(
         *,
         replace: bool,
         include_in_greeting: bool,
+        cleanup_replaced: bool = True,
     ) -> None:
         """
         Stage a table and rebuild the system prompt/executor cache.
@@ -513,6 +514,16 @@ def _add_or_replace_table(
         directly by ``.server(data_source=...)`` so that each session can
         register (or replace) its own table even after an earlier session's
         ``.server()`` call has already set ``_server_initialized``.
+
+        ``cleanup_replaced=False`` must be used for that per-session
+        replacement: a table replaced here may still be in active use by an
+        earlier, already-running session (e.g. its own
+        ``DataSourceExecutor`` holds a live reference to it), so closing/
+        disposing it here would pull the resource out from under that
+        session. Cleaning it up is then the caller's own responsibility
+        (e.g. via ``session.on_ended()`` in the code that created it). The
+        default (``True``) preserves :meth:`add_table`'s existing behavior,
+        where a config-time replacement has exactly one owner.
         """
         if not isinstance(include_in_greeting, bool):
             raise TypeError(
@@ -549,11 +560,12 @@ def _add_or_replace_table(
 
         old_source = self._data_sources.get(table_name)
         self._data_sources = next_data_sources
-        if old_source is not None and old_source is not normalized:
+        if cleanup_replaced and old_source is not None and old_source is not normalized:
             old_source.cleanup()
         if self._query_executor is not None:
-            with contextlib.suppress(Exception):
-                self._query_executor.cleanup()
+            if cleanup_replaced:
+                with contextlib.suppress(Exception):
+                    self._query_executor.cleanup()
             self._query_executor = None
 
         if include_in_greeting and table_name not in self.greeter.tables:

pkg-py/src/querychat/_shiny.py

@@ -716,6 +716,7 @@ def server(
                 resolved_table_name,
                 replace=True,
                 include_in_greeting=True,
+                cleanup_replaced=False,
             )
 
         self._require_initialized("server")

pkg-py/tests/test_server_data_source.py

@@ -1,6 +1,6 @@
 """Tests for QueryChat.server(data_source=...) parity with R (#300)."""
 
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 import pandas as pd
 import pytest
@@ -129,3 +129,71 @@ def test_add_table_still_blocked_after_server_init(
 
         with pytest.raises(RuntimeError, match="Cannot add tables after server"):
             qc.add_table(other_users_df, "other")
+
+
+class TestServerDataSourceCleanupSafety:
+    def test_second_session_does_not_clean_up_first_sessions_source(
+        self, users_df, other_users_df, captured_mod_server
+    ):
+        """
+        A second session's server(data_source=...) call must not tear down
+        the DataSource object an earlier, still-running session's own
+        DataSourceExecutor holds a live reference to (e.g. closing a DuckDB
+        connection or disposing a SQLAlchemy engine out from under it).
+        """
+        qc = shiny_mod.QueryChat(None, table_name="users")
+
+        qc.server(data_source=users_df)
+        first_source = qc._data_sources["users"]
+
+        with patch.object(first_source, "cleanup") as mock_cleanup:
+            qc.server(data_source=other_users_df)
+            mock_cleanup.assert_not_called()
+
+    def test_public_add_table_replace_still_cleans_up_old_source(
+        self, users_df, other_users_df
+    ):
+        """
+        Config-time add_table(replace=True) (before any session starts) has
+        exactly one owner for the replaced table, so its existing
+        cleanup-on-replace behavior must be unchanged.
+        """
+        qc = shiny_mod.QueryChat(users_df, "users")
+        first_source = qc._data_sources["users"]
+
+        with patch.object(first_source, "cleanup") as mock_cleanup:
+            qc.add_table(other_users_df, "users", replace=True)
+            mock_cleanup.assert_called_once()
+
+    def test_second_session_does_not_clean_up_first_sessions_query_executor(
+        self, users_df, other_users_df, captured_mod_server
+    ):
+        """
+        A second session's server(data_source=...) call must not close the
+        cached QueryExecutor an earlier, still-running session's chat has
+        already captured (e.g. via _create_session_client) and is actively
+        querying through.
+        """
+        qc = shiny_mod.QueryChat(None, table_name="users")
+
+        qc.server(data_source=users_df)
+        first_executor = qc._require_query_executor("test")
+
+        with patch.object(first_executor, "cleanup") as mock_cleanup:
+            qc.server(data_source=other_users_df)
+            mock_cleanup.assert_not_called()
+
+    def test_public_add_table_replace_still_cleans_up_old_query_executor(
+        self, users_df, other_users_df
+    ):
+        """
+        Config-time add_table(replace=True) has exactly one owner, so its
+        existing cleanup-on-replace behavior for the cached executor must be
+        unchanged.
+        """
+        qc = shiny_mod.QueryChat(users_df, "users")
+        first_executor = qc._require_query_executor("test")
+
+        with patch.object(first_executor, "cleanup") as mock_cleanup:
+            qc.add_table(other_users_df, "users", replace=True)
+            mock_cleanup.assert_called_once()