diff --git a/pkg-py/src/querychat/_querychat_base.py b/pkg-py/src/querychat/_querychat_base.py index a1433f6a..cda17bc1 100644 --- a/pkg-py/src/querychat/_querychat_base.py +++ b/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. diff --git a/pkg-py/src/querychat/_shiny.py b/pkg-py/src/querychat/_shiny.py index f394b5ac..ebdf2c2c 100644 --- a/pkg-py/src/querychat/_shiny.py +++ b/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 diff --git a/pkg-py/tests/test_base.py b/pkg-py/tests/test_base.py index ad417dc0..2e2d425e 100644 --- a/pkg-py/tests/test_base.py +++ b/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): diff --git a/pkg-py/tests/test_cleanup.py b/pkg-py/tests/test_cleanup.py index 1bb0e9e2..cbb2fb89 100644 --- a/pkg-py/tests/test_cleanup.py +++ b/pkg-py/tests/test_cleanup.py @@ -153,13 +153,13 @@ 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") @@ -167,7 +167,8 @@ def test_user_supplied_override_gets_no_on_ended( 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.""" diff --git a/pkg-py/tests/test_deferred_shiny.py b/pkg-py/tests/test_deferred_shiny.py index bc1b3f54..28228861 100644 --- a/pkg-py/tests/test_deferred_shiny.py +++ b/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.""" diff --git a/pkg-py/tests/test_multi_table.py b/pkg-py/tests/test_multi_table.py index 2e42aed9..a15a068e 100644 --- a/pkg-py/tests/test_multi_table.py +++ b/pkg-py/tests/test_multi_table.py @@ -215,9 +215,9 @@ def test_add_table_invalid_name_raises(self, orders_df, customers_df): def test_add_table_after_server_raises(self, orders_df, customers_df): """Test that adding table after server init raises error.""" qc = QueryChat(orders_df, "orders", greeting="Hello!") - qc._server_initialized = True # Simulate server initialization + qc._active_sessions = 1 # Simulate a live session - with pytest.raises(RuntimeError, match="Cannot add tables after server"): + with pytest.raises(RuntimeError, match="Cannot add tables while a server session"): qc.add_table(customers_df, "customers") @@ -251,9 +251,9 @@ def test_remove_table_after_server_raises(self, orders_df, customers_df): """Test that removing table after server init raises error.""" qc = QueryChat(orders_df, "orders", greeting="Hello!") qc.add_table(customers_df, "customers") - qc._server_initialized = True + qc._active_sessions = 1 - with pytest.raises(RuntimeError, match="Cannot remove tables after server"): + with pytest.raises(RuntimeError, match="Cannot remove tables while a server session"): qc.remove_table("customers") diff --git a/pkg-py/tests/test_server_data_source.py b/pkg-py/tests/test_server_data_source.py index 9261d428..f71a635e 100644 --- a/pkg-py/tests/test_server_data_source.py +++ b/pkg-py/tests/test_server_data_source.py @@ -36,6 +36,38 @@ def fake_mod_server(*args, **kwargs): return calls +class FakeSession(MagicMock): + """A fake Shiny session whose on_ended callbacks can be fired manually.""" + + def __init__(self): + super().__init__() + self._ended_callbacks: list = [] + self.on_ended = self._ended_callbacks.append + + def end(self): + """Simulate the session ending (fires registered on_ended callbacks).""" + for cb in self._ended_callbacks: + cb() + + +@pytest.fixture +def fake_sessions(monkeypatch): + """Patch mod_server/get_current_session; each server() call gets a new session.""" + sessions: list[FakeSession] = [] + + def fake_mod_server(*args, **kwargs): + return MagicMock() + + def next_session(): + session = FakeSession() + sessions.append(session) + return session + + monkeypatch.setattr(shiny_mod, "mod_server", fake_mod_server) + monkeypatch.setattr(shiny_mod, "get_current_session", next_session) + return sessions + + class TestServerDataSourceRegistersDeferredTable: def test_registers_deferred_table_by_constructor_name( self, users_df, captured_mod_server @@ -73,6 +105,11 @@ def test_missing_table_name_raises(self, users_df, captured_mod_server): with pytest.raises(ValueError, match="table_name"): qc.server(data_source=users_df) + def test_invalid_deferred_table_name_raises_at_construction(self): + """A bad deferred name must fail fast, not at .server() registration.""" + with pytest.raises(ValueError, match="must begin with a letter"): + shiny_mod.QueryChat(None, table_name="bad-name") + def test_empty_explicit_table_name_raises_instead_of_falling_back( self, users_df, captured_mod_server ): @@ -117,7 +154,7 @@ def test_add_table_still_blocked_after_server_init( qc = shiny_mod.QueryChat(None, table_name="users") qc.server(data_source=users_df) - with pytest.raises(RuntimeError, match="Cannot add tables after server"): + with pytest.raises(RuntimeError, match="Cannot add tables while a server session"): qc.add_table(other_users_df, "other") @@ -182,6 +219,77 @@ def test_public_add_table_replace_still_cleans_up_old_query_executor( qc.add_table(other_users_df, "users", replace=True) mock_cleanup.assert_called_once() + def test_first_server_call_cleans_up_constructor_registered_source( + self, users_df, other_users_df, captured_mod_server + ): + """No session can still be using it, so cleanup-on-replace holds.""" + qc = shiny_mod.QueryChat(users_df, "users") + constructor_source = qc._data_sources["users"] + + with patch.object(constructor_source, "cleanup") as mock_cleanup: + qc.server(data_source=other_users_df) + mock_cleanup.assert_called_once() + + def test_second_session_skips_cleanup_even_when_first_cleaned_up( + self, users_df, other_users_df, captured_mod_server + ): + qc = shiny_mod.QueryChat(users_df, "users") + + qc.server(data_source=other_users_df) + session1_source = qc._data_sources["users"] + + third_df = pd.DataFrame({"id": [7, 8, 9], "name": ["F", "G", "H"]}) + with patch.object(session1_source, "cleanup") as mock_cleanup: + qc.server(data_source=third_df) + mock_cleanup.assert_not_called() + + +class TestServerDataSourceSessionLifecycle: + """ + The hazard behind cleanup-on-replace and the add/remove_table guards is + *live* sessions, not past ones: a session that has ended can no longer + be using a resource it registered. + """ + + def test_ended_sessions_source_is_cleaned_up_on_replace( + self, users_df, other_users_df, fake_sessions + ): + qc = shiny_mod.QueryChat(None, table_name="users") + qc.server(data_source=users_df) + first_source = qc._data_sources["users"] + + fake_sessions[0].end() + + with patch.object(first_source, "cleanup") as mock_cleanup: + qc.server(data_source=other_users_df) + mock_cleanup.assert_called_once() + + def test_replaced_source_survives_while_any_session_is_live( + self, users_df, other_users_df, fake_sessions + ): + qc = shiny_mod.QueryChat(None, table_name="users") + qc.server(data_source=users_df) # session 1 + qc.server(data_source=other_users_df) # session 2 replaces s1's source + second_source = qc._data_sources["users"] + + fake_sessions[0].end() # s1 ends; s2 still live + + third_df = pd.DataFrame({"id": [7], "name": ["F"]}) + with patch.object(second_source, "cleanup") as mock_cleanup: + qc.server(data_source=third_df) # session 3 replaces s2's source + mock_cleanup.assert_not_called() + + def test_add_table_allowed_once_all_sessions_have_ended( + self, users_df, other_users_df, fake_sessions + ): + qc = shiny_mod.QueryChat(None, table_name="users") + qc.server(data_source=users_df) + + fake_sessions[0].end() + + 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( @@ -211,12 +319,12 @@ def test_unnamed_registration_replaces_config_time_table( 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( + def test_replacing_config_time_table_on_first_server_call_cleans_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. + No session is running yet, so the replaced source has a single owner + and cleanup-on-replace still holds (only later sessions skip it). """ qc = shiny_mod.QueryChat() qc.add_table(users_df, "orders") @@ -224,7 +332,7 @@ def test_replacing_config_time_table_does_not_clean_it_up( with patch.object(config_source, "cleanup") as mock_cleanup: qc.server(data_source=other_users_df) - mock_cleanup.assert_not_called() + mock_cleanup.assert_called_once() def test_explicit_table_name_adds_alongside_config_time_table( self, users_df, other_users_df, captured_mod_server