From 6613bc0e478af2be9794546cbb584efcbf985fdd Mon Sep 17 00:00:00 2001 From: Carson Date: Fri, 11 Sep 2026 20:25:09 -0500 Subject: [PATCH 1/3] fix(py): greeting reflects the session that registered its table The welcome greeting is generated asynchronously, after .server() already returns, by asking the LLM to describe the registered table(s). It reads the shared, mutable QueryChatGreeter.tables and QueryChat instance's data sources at that later point -- so a later session's own server(data_source=...) call can mutate both before an earlier session's greeting actually runs, and that earlier session's greeting ends up describing the wrong table. .server() now captures greeting_tables=list(self.greeter.tables) synchronously at call time and threads it (alongside the already per-session data_sources snapshot) through mod_server() to a new private QueryChatGreeter._generate_async_snapshot(), so the async greeting no longer reads live, shared state at generation time. Kept off the public generate()/generate_async()/build_client() API -- this is an internal fix for a race the per-session pattern introduces, not a new capability. --- pkg-py/CHANGELOG.md | 2 +- pkg-py/src/querychat/_querychat_base.py | 6 +- pkg-py/src/querychat/_querychat_greeter.py | 29 +++++++++ pkg-py/src/querychat/_shiny.py | 3 + pkg-py/src/querychat/_shiny_module.py | 7 ++- pkg-py/tests/test_querychat.py | 71 +++++++++++++++++++++- pkg-py/tests/test_server_data_source.py | 16 +++++ pkg-py/tests/test_shiny_module.py | 69 +++++++++++++++++++++ 8 files changed, 199 insertions(+), 4 deletions(-) diff --git a/pkg-py/CHANGELOG.md b/pkg-py/CHANGELOG.md index e8be2fdc2..8bca28401 100644 --- a/pkg-py/CHANGELOG.md +++ b/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, 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). + Registering a table this way is no longer blocked by an earlier session having already registered one, 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), and each session's auto-generated greeting reflects its own table even if a later session registers a different one before that greeting is generated. ### Improvements diff --git a/pkg-py/src/querychat/_querychat_base.py b/pkg-py/src/querychat/_querychat_base.py index b0a7fb8ae..851baf12e 100644 --- a/pkg-py/src/querychat/_querychat_base.py +++ b/pkg-py/src/querychat/_querychat_base.py @@ -383,10 +383,14 @@ def client_factory( tables: list[str], prompt: str | Path, base: chatlas.Chat | None = None, + *, + data_sources: dict[str, DataSource] | None = None, ) -> chatlas.Chat: sp = QueryChatSystemPrompt( prompt_template=prompt, - data_sources=self._data_sources, + data_sources=( + self._data_sources if data_sources is None else data_sources + ), data_description=self._data_description, extra_instructions=None, categorical_threshold=self._categorical_threshold, diff --git a/pkg-py/src/querychat/_querychat_greeter.py b/pkg-py/src/querychat/_querychat_greeter.py index ef95551ef..f3c5c6c3f 100644 --- a/pkg-py/src/querychat/_querychat_greeter.py +++ b/pkg-py/src/querychat/_querychat_greeter.py @@ -12,6 +12,8 @@ import chatlas + from ._datasource import DataSource + class QueryChatGreeter: """Controls greeting generation for a QueryChat instance. Access via ``qc.greeter``.""" @@ -68,3 +70,30 @@ async def generate_async(self, *, base: chatlas.Chat | None = None): """Stream a greeting response from the greeting client.""" client = self.build_client(base) return await client.stream_async(GREETING_PROMPT, echo="none") + + async def _generate_async_snapshot( + self, + *, + base: chatlas.Chat | None, + tables: list[str] | None, + data_sources: dict[str, DataSource], + ): + """ + Stream a greeting response from an explicit session snapshot. + + Internal counterpart to :meth:`generate_async`, used by + ``mod_server()`` for Shiny sessions. `tables` and `data_sources` + override the live `self.tables` and the QueryChat instance's data + sources with a point-in-time snapshot captured when the session's + `.server()` call ran, rather than reading that shared, mutable state + whenever shinychat gets around to invoking the (lazily-scheduled, + asynchronous) greeting callback -- by then, a *later* session's own + `.server(data_source=...)` call may have already mutated it. + """ + client = self._client_factory( + self._tables if tables is None else tables, + self._prompt, + base, + data_sources=data_sources, + ) + return await client.stream_async(GREETING_PROMPT, echo="none") diff --git a/pkg-py/src/querychat/_shiny.py b/pkg-py/src/querychat/_shiny.py index 0e2d1b6ed..f394b5acc 100644 --- a/pkg-py/src/querychat/_shiny.py +++ b/pkg-py/src/querychat/_shiny.py @@ -420,6 +420,7 @@ def app_server(input: Inputs, output: Outputs, session: Session): tools=self.tools, greeter=self.greeter, greeting_base=None, + greeting_tables=list(self.greeter.tables), ) @reactive.calc @@ -768,6 +769,7 @@ def create_session_client(**kwargs) -> chatlas.Chat: tools=self.tools, greeter=self.greeter, greeting_base=resolved_client, + greeting_tables=list(self.greeter.tables), ) @@ -1055,6 +1057,7 @@ def _ensure_server_started(self) -> None: tools=self.tools, greeter=self.greeter, greeting_base=None, + greeting_tables=list(self.greeter.tables), ) def sidebar( diff --git a/pkg-py/src/querychat/_shiny_module.py b/pkg-py/src/querychat/_shiny_module.py index 16f2d9630..b238b4191 100644 --- a/pkg-py/src/querychat/_shiny_module.py +++ b/pkg-py/src/querychat/_shiny_module.py @@ -248,6 +248,7 @@ def mod_server( tools: set[str] | None = None, greeter: QueryChatGreeter, greeting_base: chatlas.Chat | None = None, + greeting_tables: list[str] | None = None, ) -> ServerValues[IntoFrameT]: if not callable(client): raise TypeError("mod_server() requires a callable client factory.") @@ -344,7 +345,11 @@ async def _make_greeting(): GreetWarning, stacklevel=1, ) - stream = await greeter.generate_async(base=greeting_base) + stream = await greeter._generate_async_snapshot( + base=greeting_base, + tables=greeting_tables, + data_sources=data_sources, + ) return shinychat.chat_greeting(stream, persistent=True) greeting_arg = ( diff --git a/pkg-py/tests/test_querychat.py b/pkg-py/tests/test_querychat.py index 5dcbdcb8e..3511f3757 100644 --- a/pkg-py/tests/test_querychat.py +++ b/pkg-py/tests/test_querychat.py @@ -1,14 +1,21 @@ +import asyncio import os import tempfile from pathlib import Path from unittest.mock import patch import ibis +import narwhals.stable.v1 as nw import pandas as pd import polars as pl import pytest from querychat import QueryChat -from querychat._datasource import IbisSource, PolarsLazySource +from querychat._datasource import ( + DataFrameSource, + DataSource, + IbisSource, + PolarsLazySource, +) from sqlalchemy import create_engine, text @@ -429,3 +436,65 @@ def test_remove_table_prunes_greeter_tables(sqlite_engine): qc.remove_table("orders") assert "orders" not in qc.greeter.tables assert "customers" in qc.greeter.tables + + +class TestGreeterSnapshotOverrides: + """ + _generate_async_snapshot() (private; used internally by mod_server() for + Shiny sessions) renders a greeting from an explicit tables/data_sources + snapshot instead of the live, shared greeter.tables/ + QueryChat._data_sources -- which can be mutated by a later Shiny + session's server(data_source=...) call before an earlier session's async + greeting generation runs. The public build_client()/generate()/ + generate_async() API is unaffected and keeps reading live state. + """ + + def test_build_client_uses_live_state(self, sample_df): + qc = QueryChat(sample_df, "test_table") + prompt = qc.greeter.build_client().system_prompt + assert prompt is not None + assert "test_table" in prompt + + def test_snapshot_tables_override_ignores_live_greeter_tables(self, sample_df): + qc = QueryChat(sample_df, "test_table") + qc.greeter.tables = [] # live state says "no tables" + seen: dict[str, str | None] = {} + + async def fake_stream_async(self, *args, **kwargs): + seen["system_prompt"] = self.system_prompt + return "stream" + + with patch("chatlas.Chat.stream_async", fake_stream_async): + asyncio.run( + qc.greeter._generate_async_snapshot( + base=None, tables=["test_table"], data_sources=qc._data_sources + ) + ) + + assert seen["system_prompt"] is not None + assert "test_table" in seen["system_prompt"] + + def test_snapshot_data_sources_override_ignores_live_data_sources(self, sample_df): + qc = QueryChat(sample_df, "test_table") + other_df = pd.DataFrame({"z": [1, 2, 3]}) + snapshot: dict[str, DataSource] = { + "other_table": DataFrameSource( + nw.from_native(other_df, eager_only=True), "other_table" + ) + } + seen: dict[str, str | None] = {} + + async def fake_stream_async(self, *args, **kwargs): + seen["system_prompt"] = self.system_prompt + return "stream" + + with patch("chatlas.Chat.stream_async", fake_stream_async): + asyncio.run( + qc.greeter._generate_async_snapshot( + base=None, tables=["other_table"], data_sources=snapshot + ) + ) + + assert seen["system_prompt"] is not None + assert "other_table" in seen["system_prompt"] + assert "test_table" not in seen["system_prompt"] diff --git a/pkg-py/tests/test_server_data_source.py b/pkg-py/tests/test_server_data_source.py index 125eb96a1..911bd6242 100644 --- a/pkg-py/tests/test_server_data_source.py +++ b/pkg-py/tests/test_server_data_source.py @@ -197,3 +197,19 @@ def test_public_add_table_replace_still_cleans_up_old_query_executor( with patch.object(first_executor, "cleanup") as mock_cleanup: qc.add_table(other_users_df, "users", replace=True) mock_cleanup.assert_called_once() + + +class TestServerDataSourceGreetingSnapshot: + def test_server_passes_greeting_tables_snapshot_to_mod_server( + self, users_df, captured_mod_server + ): + """ + .server() must pass a snapshot of greeter.tables captured at call + time, so mod_server's greeting generation doesn't read the live, + mutable greeter.tables from an async task that may run after a + later session has changed it. + """ + qc = shiny_mod.QueryChat(None, table_name="users") + qc.server(data_source=users_df) + + assert captured_mod_server[0]["greeting_tables"] == ["users"] diff --git a/pkg-py/tests/test_shiny_module.py b/pkg-py/tests/test_shiny_module.py index 45b3252de..1c94cfcd2 100644 --- a/pkg-py/tests/test_shiny_module.py +++ b/pkg-py/tests/test_shiny_module.py @@ -156,6 +156,75 @@ def fake_chat_constructor( assert handoff_server_mock.call_args.kwargs["executor"] is fake_executor +def test_mod_server_generates_greeting_from_session_snapshot_not_live_state(): + """ + _make_greeting() must render the greeting from this session's own + data_sources/greeting_tables snapshot, not by reading the shared, + mutable greeter/QueryChat._data_sources live at async-generation time -- + a later Shiny session's server(data_source=...) call can mutate that + shared state before an earlier session's greeting finishes streaming. + """ + import asyncio + from unittest.mock import AsyncMock, MagicMock, patch + + from querychat._shiny_module import mod_server + + captured = {} + + def fake_chat_constructor(id, *, client=None, greeting=None, history=None, **kwargs): + captured["greeting"] = greeting + return MagicMock() + + fake_source = MagicMock() + fake_source.get_data.return_value = [] + fake_executor = MagicMock() + fake_executor.execute_query.return_value = [] + + def client_factory(**kwargs): + return MagicMock(spec=["stream_async"]) + + fake_greeter = MagicMock() + fake_greeter._generate_async_snapshot = AsyncMock(return_value=MagicMock()) + fake_greeting_base = MagicMock() + + inner_fn = _unwrap_module_server(mod_server) + + fake_input = MagicMock() + fake_input.__getitem__ = MagicMock(return_value=MagicMock()) + fake_session = MagicMock() + fake_session.is_stub_session.return_value = False + + with ( + patch( + "querychat._shiny_module.shinychat.Chat", side_effect=fake_chat_constructor + ), + patch("querychat._shiny_module.has_viz_tool", return_value=False), + patch("querychat._shiny_module.shinychat.chat_greeting", return_value=MagicMock()), + ): + inner_fn( + fake_input, + MagicMock(), + fake_session, + data_sources={"t": fake_source}, + executor=fake_executor, + greeting=None, + client=client_factory, + history=True, + tools=None, + greeter=fake_greeter, + greeting_base=fake_greeting_base, + greeting_tables=["t"], + ) + + asyncio.run(captured["greeting"]()) + + fake_greeter._generate_async_snapshot.assert_called_once_with( + base=fake_greeting_base, + tables=["t"], + data_sources={"t": fake_source}, + ) + + def test_mod_server_registers_chat_bookmarking_with_no_auto_trigger_when_history_not_bookmark_mode(): """ Chat.enable_bookmarking() is called when `history` isn't bookmark mode, so From 17bcfdabecd1f3c749b345b915962ca9e1bba0ca Mon Sep 17 00:00:00 2001 From: Carson Date: Sat, 12 Sep 2026 10:20:21 -0500 Subject: [PATCH 2/3] docs(py): trim comments/docstrings to what isn't inferable from the code Also applies ruff format to test_shiny_module.py (two of the three spots were added in this stack; the third is pre-existing). --- pkg-py/src/querychat/_querychat_base.py | 24 ++++++------- pkg-py/src/querychat/_querychat_greeter.py | 11 +++--- pkg-py/tests/test_querychat.py | 11 +++--- pkg-py/tests/test_server_data_source.py | 40 ++++++---------------- pkg-py/tests/test_shiny_module.py | 20 +++++------ 5 files changed, 39 insertions(+), 67 deletions(-) diff --git a/pkg-py/src/querychat/_querychat_base.py b/pkg-py/src/querychat/_querychat_base.py index 851baf12e..a1433f6a1 100644 --- a/pkg-py/src/querychat/_querychat_base.py +++ b/pkg-py/src/querychat/_querychat_base.py @@ -514,20 +514,16 @@ def _add_or_replace_table( """ 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``. - - ``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. + 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``. + + ``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()``). """ if not isinstance(include_in_greeting, bool): raise TypeError( diff --git a/pkg-py/src/querychat/_querychat_greeter.py b/pkg-py/src/querychat/_querychat_greeter.py index f3c5c6c3f..f5d4422db 100644 --- a/pkg-py/src/querychat/_querychat_greeter.py +++ b/pkg-py/src/querychat/_querychat_greeter.py @@ -82,13 +82,10 @@ async def _generate_async_snapshot( Stream a greeting response from an explicit session snapshot. Internal counterpart to :meth:`generate_async`, used by - ``mod_server()`` for Shiny sessions. `tables` and `data_sources` - override the live `self.tables` and the QueryChat instance's data - sources with a point-in-time snapshot captured when the session's - `.server()` call ran, rather than reading that shared, mutable state - whenever shinychat gets around to invoking the (lazily-scheduled, - asynchronous) greeting callback -- by then, a *later* session's own - `.server(data_source=...)` call may have already mutated it. + ``mod_server()``. The snapshot matters because greeting generation + is scheduled lazily: by the time it runs, a later session's + ``.server(data_source=...)`` call may have already mutated the + shared live state. """ client = self._client_factory( self._tables if tables is None else tables, diff --git a/pkg-py/tests/test_querychat.py b/pkg-py/tests/test_querychat.py index 3511f3757..6f3e356c5 100644 --- a/pkg-py/tests/test_querychat.py +++ b/pkg-py/tests/test_querychat.py @@ -440,13 +440,10 @@ def test_remove_table_prunes_greeter_tables(sqlite_engine): class TestGreeterSnapshotOverrides: """ - _generate_async_snapshot() (private; used internally by mod_server() for - Shiny sessions) renders a greeting from an explicit tables/data_sources - snapshot instead of the live, shared greeter.tables/ - QueryChat._data_sources -- which can be mutated by a later Shiny - session's server(data_source=...) call before an earlier session's async - greeting generation runs. The public build_client()/generate()/ - generate_async() API is unaffected and keeps reading live state. + _generate_async_snapshot() renders from an explicit tables/data_sources + snapshot instead of live shared state, which a later Shiny session may + have mutated before an earlier session's async greeting runs. The + public build_client()/generate()/generate_async() API is unaffected. """ def test_build_client_uses_live_state(self, sample_df): diff --git a/pkg-py/tests/test_server_data_source.py b/pkg-py/tests/test_server_data_source.py index 911bd6242..87f069130 100644 --- a/pkg-py/tests/test_server_data_source.py +++ b/pkg-py/tests/test_server_data_source.py @@ -76,11 +76,7 @@ def test_missing_table_name_raises(self, users_df, captured_mod_server): def test_empty_explicit_table_name_raises_instead_of_falling_back( self, users_df, captured_mod_server ): - """ - An explicit but invalid table_name="" must be validated and rejected, - not silently treated as omitted and fall back to the deferred/first - table name. - """ + """An explicit table_name="" must be rejected, not treated as omitted.""" qc = shiny_mod.QueryChat(None, table_name="users") with pytest.raises(ValueError, match="must begin with a letter"): qc.server(data_source=users_df, table_name="") @@ -101,12 +97,6 @@ def test_no_data_source_leaves_tables_unchanged( class TestServerDataSourceSurvivesSecondSession: - """ - 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_does_not_raise( self, users_df, other_users_df, captured_mod_server ): @@ -136,10 +126,8 @@ 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). + An earlier, still-running session's executor holds a live + reference to the source a later session's registration replaces. """ qc = shiny_mod.QueryChat(None, table_name="users") @@ -154,9 +142,8 @@ 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. + Config-time replacement has a single owner, so cleanup-on-replace + is unchanged on the public path. """ qc = shiny_mod.QueryChat(users_df, "users") first_source = qc._data_sources["users"] @@ -169,10 +156,8 @@ 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. + An earlier, still-running session's chat has already captured the + cached executor and may be querying through it. """ qc = shiny_mod.QueryChat(None, table_name="users") @@ -187,9 +172,8 @@ 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. + Config-time replacement has a single owner, so executor cleanup + is unchanged on the public path. """ qc = shiny_mod.QueryChat(users_df, "users") first_executor = qc._require_query_executor("test") @@ -204,10 +188,8 @@ def test_server_passes_greeting_tables_snapshot_to_mod_server( self, users_df, captured_mod_server ): """ - .server() must pass a snapshot of greeter.tables captured at call - time, so mod_server's greeting generation doesn't read the live, - mutable greeter.tables from an async task that may run after a - later session has changed it. + Greeting generation runs lazily, after a later session may have + mutated the live greeter.tables -- hence the call-time snapshot. """ qc = shiny_mod.QueryChat(None, table_name="users") qc.server(data_source=users_df) diff --git a/pkg-py/tests/test_shiny_module.py b/pkg-py/tests/test_shiny_module.py index 1c94cfcd2..8bc884f3f 100644 --- a/pkg-py/tests/test_shiny_module.py +++ b/pkg-py/tests/test_shiny_module.py @@ -158,11 +158,9 @@ def fake_chat_constructor( def test_mod_server_generates_greeting_from_session_snapshot_not_live_state(): """ - _make_greeting() must render the greeting from this session's own - data_sources/greeting_tables snapshot, not by reading the shared, - mutable greeter/QueryChat._data_sources live at async-generation time -- - a later Shiny session's server(data_source=...) call can mutate that - shared state before an earlier session's greeting finishes streaming. + _make_greeting() must render from this session's own snapshot, not the + shared live state, which a later session's server(data_source=...) call + may have mutated before an earlier session's greeting streams. """ import asyncio from unittest.mock import AsyncMock, MagicMock, patch @@ -171,7 +169,9 @@ def test_mod_server_generates_greeting_from_session_snapshot_not_live_state(): captured = {} - def fake_chat_constructor(id, *, client=None, greeting=None, history=None, **kwargs): + def fake_chat_constructor( + id, *, client=None, greeting=None, history=None, **kwargs + ): captured["greeting"] = greeting return MagicMock() @@ -199,7 +199,9 @@ def client_factory(**kwargs): "querychat._shiny_module.shinychat.Chat", side_effect=fake_chat_constructor ), patch("querychat._shiny_module.has_viz_tool", return_value=False), - patch("querychat._shiny_module.shinychat.chat_greeting", return_value=MagicMock()), + patch( + "querychat._shiny_module.shinychat.chat_greeting", return_value=MagicMock() + ), ): inner_fn( fake_input, @@ -417,9 +419,7 @@ def test_shinychat_chat_contract_used_by_mod_server(): mock_session.app = None with session_context(mock_session): - chat = shinychat.Chat( - "chat", client=MagicMock(), greeting=None, history=True - ) + chat = shinychat.Chat("chat", client=MagicMock(), greeting=None, history=True) @chat.history.on_save def _on_save(values): From 8ace1adf94f025c597a15be2f547ac5e4686d643 Mon Sep 17 00:00:00 2001 From: Carson Date: Sat, 12 Sep 2026 11:02:55 -0500 Subject: [PATCH 3/3] test(py): lock in mixed config-time add_table() + per-session server(data_source=) behavior Parity with the R tests in #306: unnamed registration replaces the config-time table without cleaning it up, an explicit table_name adds alongside it, and the registry is shared and cumulative across sessions. --- pkg-py/tests/test_server_data_source.py | 62 +++++++++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/pkg-py/tests/test_server_data_source.py b/pkg-py/tests/test_server_data_source.py index 87f069130..9261d428e 100644 --- a/pkg-py/tests/test_server_data_source.py +++ b/pkg-py/tests/test_server_data_source.py @@ -195,3 +195,65 @@ def test_server_passes_greeting_tables_snapshot_to_mod_server( qc.server(data_source=users_df) assert captured_mod_server[0]["greeting_tables"] == ["users"] + + +class TestServerDataSourceMixedWithConfigTimeAddTable: + def test_unnamed_registration_replaces_config_time_table( + self, users_df, other_users_df, captured_mod_server + ): + qc = shiny_mod.QueryChat() + qc.add_table(users_df, "orders") + + qc.server(data_source=other_users_df) + + # Same table name, but the session's data replaces the config-time data + sources = captured_mod_server[0]["data_sources"] + assert list(sources.keys()) == ["orders"] + assert sources["orders"].get_data()["id"].tolist() == [4, 5] + + def test_replacing_config_time_table_does_not_clean_it_up( + self, users_df, other_users_df, captured_mod_server + ): + """ + Consistent with per-session replacement: the replaced source's + cleanup is left to whoever created it. + """ + qc = shiny_mod.QueryChat() + qc.add_table(users_df, "orders") + config_source = qc._data_sources["orders"] + + with patch.object(config_source, "cleanup") as mock_cleanup: + qc.server(data_source=other_users_df) + mock_cleanup.assert_not_called() + + def test_explicit_table_name_adds_alongside_config_time_table( + self, users_df, other_users_df, captured_mod_server + ): + qc = shiny_mod.QueryChat() + qc.add_table(users_df, "orders") + + qc.server(data_source=other_users_df, table_name="returns") + + sources = captured_mod_server[0]["data_sources"] + assert list(sources.keys()) == ["orders", "returns"] + # The config-time table's own data is untouched + assert sources["orders"].get_data()["id"].tolist() == [1, 2, 3] + assert sources["returns"].get_data()["id"].tolist() == [4, 5] + + def test_later_session_snapshot_includes_earlier_sessions_table( + self, users_df, other_users_df, captured_mod_server + ): + """The registry is shared and cumulative across sessions.""" + qc = shiny_mod.QueryChat() + qc.add_table(users_df, "orders") + + # Session 1 adds its own table alongside the config-time one + qc.server(data_source=other_users_df, table_name="returns") + # Session 2 replaces "orders" only -- but still sees session 1's table + third_df = pd.DataFrame({"id": [7, 8, 9]}) + qc.server(data_source=third_df, table_name="orders") + + sources = captured_mod_server[1]["data_sources"] + assert list(sources.keys()) == ["orders", "returns"] + assert sources["orders"].get_data()["id"].tolist() == [7, 8, 9] + assert sources["returns"].get_data()["id"].tolist() == [4, 5]