fix(py): let server(data_source=) survive a second session

pkg-py/CHANGELOG.md

@@ -31,6 +31,8 @@ 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.
+
 ### Improvements
 
 * The `"visualize"` tool is now included in the default toolset (`tools=("filter", "query", "visualize")`). If the visualization dependencies are not installed (the `viz` extra), the tool is dropped with a warning instead of raising an `ImportError`.

pkg-py/src/querychat/_querychat_base.py

@@ -491,7 +491,29 @@ def add_table(
                 "Cannot add tables after server initialization. "
                 "Add all tables before calling .server() or .app()."
             )
+        self._add_or_replace_table(
+            data_source,
+            table_name,
+            replace=replace,
+            include_in_greeting=include_in_greeting,
+        )
 
+    def _add_or_replace_table(
+        self,
+        data_source: IntoFrame | sqlalchemy.Engine | BaseBoard,
+        table_name: str,
+        *,
+        replace: bool,
+        include_in_greeting: bool,
+    ) -> None:
+        """
+        Stage a table and rebuild the system prompt/executor cache.
+
+        This is the guard-free core of :meth:`add_table`. It's also called
+        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``.
+        """
         if not isinstance(include_in_greeting, bool):
             raise TypeError(
                 "include_in_greeting must be True or False, got "

pkg-py/src/querychat/_shiny.py

@@ -711,7 +711,7 @@ def server(
                     "or table_name to the QueryChat constructor, or register a "
                     "table first with add_table()."
                 )
-            self.add_table(
+            self._add_or_replace_table(
                 data_source,
                 resolved_table_name,
                 replace=True,

pkg-py/tests/test_server_data_source.py

@@ -100,20 +100,32 @@ def test_no_data_source_leaves_tables_unchanged(
         assert qc.table_names() == ["users"]
 
 
-class TestServerDataSourceCurrentSingleSessionLimitation:
+class TestServerDataSourceSurvivesSecondSession:
     """
-    Known limitation (tracked by a follow-up PR): server(data_source=...)
-    reuses the public add_table(), so it still hits the "no changes after
-    server initialization" guard on a second session -- the same bug R's
-    existing $server(data_source=) has today. Making this survive a second
-    session is a separate, focused change.
+    server(data_source=...) must not be blocked by an earlier session having
+    already registered a table -- unlike the public add_table(), which still
+    guards against changes after server initialization.
     """
 
-    def test_second_session_currently_raises(
+    def test_second_session_does_not_raise(
         self, users_df, other_users_df, captured_mod_server
     ):
         qc = shiny_mod.QueryChat(None, table_name="users")
         qc.server(data_source=users_df)
 
+        qc.server(data_source=other_users_df)  # must not raise
+
+        assert list(captured_mod_server[1]["data_sources"].keys()) == ["users"]
+
+    def test_add_table_still_blocked_after_server_init(
+        self, users_df, other_users_df, captured_mod_server
+    ):
+        """
+        The public add_table() guard must remain intact; only the
+        server(data_source=...) path bypasses it.
+        """
+        qc = shiny_mod.QueryChat(None, table_name="users")
+        qc.server(data_source=users_df)
+
         with pytest.raises(RuntimeError, match="Cannot add tables after server"):
-            qc.server(data_source=other_users_df)
+            qc.add_table(other_users_df, "other")